From 6636e696492729b44a30193b2f8272c3ec180a32 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 18:34:03 +0000 Subject: [PATCH 01/31] cuda: a one-shot RoCE all-gather for two ranks (roce.zig, roce_gather.cu) and the ranks' control link (tp_link.zig): small gathers as RDMA writes over both halves of the port, after b12x's exchange; NCCL carries the rest --- zig/build/cuda.zig | 1 + zig/kernels/cuda/roce_gather.cu | 114 +++++++ zig/src/cuda/kernels.zig | 1 + zig/src/cuda/nccl.zig | 2 + zig/src/cuda/roce.zig | 543 ++++++++++++++++++++++++++++++++ zig/src/cuda/root.zig | 3 + zig/src/cuda/tp_link.zig | 170 ++++++++++ 7 files changed, 834 insertions(+) create mode 100644 zig/kernels/cuda/roce_gather.cu create mode 100644 zig/src/cuda/roce.zig create mode 100644 zig/src/cuda/tp_link.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..4ec1ae3c9b 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -48,6 +48,7 @@ const kernels = [_]Kernel{ .{ .name = "torch_indexing", .src = "torch_ops/indexing", .flags = torch_ops }, .{ .name = "torch_movement", .src = "torch_ops/movement", .flags = torch_ops }, .{ .name = "torch_nemotron_constants", .src = "torch_ops/nemotron_constants", .flags = torch_ops }, + .{ .name = "roce_gather", .src = "roce_gather", .flags = &.{"-O3"} }, // b12x's one-shot RoCE all-gather, two ranks }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. diff --git a/zig/kernels/cuda/roce_gather.cu b/zig/kernels/cuda/roce_gather.cu new file mode 100644 index 0000000000..8d6ff32e01 --- /dev/null +++ b/zig/kernels/cuda/roce_gather.cu @@ -0,0 +1,114 @@ +// One-shot all-gather between two DGX Sparks over RoCE: the GPU side of b12x's "RoCEnante" exchange +// (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks. +// The host side (zig/src/cuda/roce.zig) owns a pinned region both GPUs and both NICs reach: +// +// recv[slot] the peer's shard, written by the peer's NIC (2 slots of slot_bytes) +// flag[slot][hca] the sequence the peer's NIC writes after its stripe on that HCA (128 bytes apart) +// send[slot] this rank's shard, staged here for its NIC (2 slots) +// ctrl {seq, nbytes of slot 0, nbytes of slot 1, error seq, missing hca}: the proxy's doorbell +// +// One launch: (1) every block stages its share of the shard into send[seq & 1]; (2) the last block to stage +// publishes the slot's size and rings ctrl.seq after a system fence, and the proxy thread RDMA-writes the shard and +// then a flag per HCA stripe on the same queue pair (so a flag never lands before its stripe); (3) each block waits +// for every stripe's flag; (4) every block writes both shards in rank order (the layout NCCL's all-gather gives); +// (5) the last block to finish advances the device epoch, so a CUDA graph replays the exchange with new sequences. +// A wait past the spin limit poisons the exchange: later launches return at once and the host reports the error. + +#include + +namespace { + +constexpr int C_SEQ = 0, C_NBYTES = 1, C_ERROR = 3, C_HCA = 4; // ctrl words (C_NBYTES + slot) +constexpr int FLAG_STRIDE = 128; +// counters (device memory, u32): epoch, poison, then stage and tail arrivals for each grid size class +constexpr int K_EPOCH = 0, K_POISON = 1, K_CLASSES = 4; + +__device__ __forceinline__ uint32_t ld_relaxed_gpu(const uint32_t* p) { + uint32_t v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ uint32_t ld_relaxed_sys(const uint32_t* p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ uint32_t ld_acquire_sys(const uint32_t* p) { + uint32_t v; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ uint4 ld_relaxed_sys4(const uint4* p) { + uint4 v; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ void st_relaxed_sys(uint32_t* p, uint32_t v) { + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} +__device__ __forceinline__ void st_release_gpu(uint32_t* p, uint32_t v) { + asm volatile("st.release.gpu.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} +__device__ __forceinline__ uint32_t arrive(uint32_t* p) { + uint32_t old; + asm volatile("atom.relaxed.gpu.global.add.u32 %0, [%1], 1;" : "=r"(old) : "l"(p) : "memory"); + return old; +} +__device__ __forceinline__ void fence_sys() { asm volatile("fence.sc.sys;" ::: "memory"); } +__device__ __forceinline__ void fence_gpu() { asm volatile("fence.sc.gpu;" ::: "memory"); } + +} // namespace + +// in: this rank's shard (packs of 16 bytes); out: [rank 0's shard | rank 1's shard]. klass: log2(gridDim.x). +extern "C" __global__ void __launch_bounds__(512) roce_gather( + const uint4* __restrict__ in, uint4* __restrict__ out, int packs, int rank, char* recv, char* flags, char* send, + uint32_t* ctrl, long long slot_bytes, uint32_t* counters, int klass, uint32_t spin_limit, int n_hca) { + const int tid = threadIdx.x; + uint32_t* stage = counters + 2 + klass; + uint32_t* tail = counters + 2 + K_CLASSES + klass; + if (ld_relaxed_gpu(counters + K_POISON) != 0u) return; + const uint32_t seq = ld_relaxed_gpu(counters + K_EPOCH) + 1u; + const uint32_t slot = seq & 1u; + const int at = blockIdx.x * blockDim.x + tid, stride = gridDim.x * blockDim.x; + + uint4* mine = reinterpret_cast(send + (long long)slot * slot_bytes); // 1. stage + for (int i = at; i < packs; i += stride) mine[i] = in[i]; + __syncthreads(); + if (tid == 0) { // 2. the doorbell, by the last block to stage + fence_sys(); + if ((arrive(stage) + 1u) % gridDim.x == 0u) { + st_relaxed_sys(ctrl + C_NBYTES + slot, (uint32_t)packs * 16u); + fence_sys(); + st_relaxed_sys(ctrl + C_SEQ, seq); + } + } + if (tid < n_hca) { // 3. the peer's stripes + const uint32_t* flag = reinterpret_cast(flags + ((long long)slot * n_hca + tid) * FLAG_STRIDE); + uint32_t polls = 0; + while (ld_acquire_sys(flag) != seq) { + if (++polls >= spin_limit) { + st_relaxed_sys(ctrl + C_HCA, (uint32_t)tid); + st_relaxed_sys(ctrl + C_ERROR, seq); + st_release_gpu(counters + K_POISON, seq); + break; + } + } + } + __syncthreads(); + if (ld_relaxed_gpu(counters + K_POISON) == 0u) { // 4. both shards in rank order + const uint4* theirs = reinterpret_cast(recv + (long long)slot * slot_bytes); + uint4* own_out = out + (long long)rank * packs; + uint4* peer_out = out + (long long)(1 - rank) * packs; + for (int i = at; i < packs; i += stride) { + own_out[i] = in[i]; + peer_out[i] = ld_relaxed_sys4(theirs + i); + } + } + fence_gpu(); // 5. the epoch, by the last block to finish + __syncthreads(); + if (tid == 0 && (arrive(tail) + 1u) % gridDim.x == 0u) { + fence_gpu(); + if (ld_relaxed_sys(ctrl + C_ERROR) == 0u) st_release_gpu(counters + K_EPOCH, seq); + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..75948fe016 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -46,6 +46,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 roce_gather: []const u8 = if (available) &Blob("fatbin_roce_gather").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/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..ec78fd135d 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -21,6 +21,8 @@ pub const Api = struct { ncclCommInitRank: *const fn (*Comm, c_int, UniqueId, c_int) callconv(.c) R, ncclCommInitAll: *const fn ([*]Comm, c_int, ?[*]const c_int) callconv(.c) R, ncclCommDestroy: *const fn (Comm) callconv(.c) R, + ncclCommAbort: *const fn (Comm) callconv(.c) R, + ncclCommGetAsyncError: *const fn (Comm, *R) callconv(.c) R, ncclGetErrorString: *const fn (R) callconv(.c) ?[*:0]const u8, ncclAllReduce: *const fn (D, D, usize, DataType, RedOp, Comm, abi.Stream) callconv(.c) R, ncclAllGather: *const fn (D, D, usize, DataType, Comm, abi.Stream) callconv(.c) R, diff --git a/zig/src/cuda/roce.zig b/zig/src/cuda/roce.zig new file mode 100644 index 0000000000..93fd1e51f5 --- /dev/null +++ b/zig/src/cuda/roce.zig @@ -0,0 +1,543 @@ +//! Small all-gathers between two ranks as one-shot RDMA writes over RoCE: b12x's "RoCEnante" exchange +//! (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks: the +//! kernel (kernels/cuda/roce_gather.cu) stages this rank's shard in pinned host memory and rings a doorbell; a proxy +//! thread here RDMA-writes it to the peer, striped over the HCAs, each stripe followed by its sequence flag on the +//! same queue pair; the kernel waits for the peer's flags and copies both shards out. A decode step's gathers are a +//! few rows of partials, where NCCL spends tens of microseconds a call in fixed costs; this path takes about ten. +//! libibverbs is opened at run time (as the driver and NCCL are): without it, or without an active HCA, the caller +//! keeps NCCL. + +const std = @import("std"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const memory = @import("memory.zig"); +const Module = @import("module.zig").Module; +const Function = @import("module.zig").Function; +const launch = @import("launch.zig"); +const tp_link = @import("tp_link.zig"); +const kernels = @import("kernels.zig"); + +const max_hcas = 2; +const slots = 2; +const flag_stride = 128; +const send_depth = 256; +const port_num: u8 = 1; +const threads = 512; +const max_blocks = 8; // the largest grid; smaller shards launch fewer (a power of two) +const classes = 4; // grid size classes 1, 2, 4, 8 (each its own arrival counters) +const idle_spins: u64 = 20_000_000; // doorbell polls before the proxy naps between them (a model's gaps are short) +const spin_limit: u32 = 100_000_000; // the kernel's flag polls before it gives up and poisons the exchange + +// -- libibverbs (the ABI of rdma-core's verbs.h; offsets checked against its headers) --------------------------------- +const Device = opaque {}; +const Context = opaque {}; +const Pd = opaque {}; +const Cq = extern struct { context: *Context }; +const Mr = extern struct { context: *Context, pd: *Pd, addr: ?*anyopaque, length: usize, handle: u32, lkey: u32, rkey: u32 }; +const Qp = extern struct { context: *Context, qp_context: ?*anyopaque, pd: *Pd, send_cq: *Cq, recv_cq: *Cq, srq: ?*anyopaque, handle: u32, qp_num: u32 }; +const QpCap = extern struct { max_send_wr: u32 = 0, max_recv_wr: u32 = 0, max_send_sge: u32 = 0, max_recv_sge: u32 = 0, max_inline_data: u32 = 0 }; +const QpInitAttr = extern struct { qp_context: ?*anyopaque = null, send_cq: ?*Cq = null, recv_cq: ?*Cq = null, srq: ?*anyopaque = null, cap: QpCap = .{}, qp_type: c_int = 0, sq_sig_all: c_int = 0 }; +const GlobalRoute = extern struct { dgid: [16]u8 = @splat(0), flow_label: u32 = 0, sgid_index: u8 = 0, hop_limit: u8 = 0, traffic_class: u8 = 0 }; +const AhAttr = extern struct { grh: GlobalRoute = .{}, dlid: u16 = 0, sl: u8 = 0, src_path_bits: u8 = 0, static_rate: u8 = 0, is_global: u8 = 0, port_num: u8 = 0 }; +const QpAttr = extern struct { + qp_state: u32 = 0, cur_qp_state: u32 = 0, path_mtu: u32 = 0, path_mig_state: u32 = 0, qkey: u32 = 0, rq_psn: u32 = 0, + sq_psn: u32 = 0, dest_qp_num: u32 = 0, qp_access_flags: u32 = 0, cap: QpCap = .{}, ah_attr: AhAttr = .{}, alt_ah_attr: AhAttr = .{}, + pkey_index: u16 = 0, alt_pkey_index: u16 = 0, en_sqd_async_notify: u8 = 0, sq_draining: u8 = 0, max_rd_atomic: u8 = 0, + max_dest_rd_atomic: u8 = 0, min_rnr_timer: u8 = 0, port_num: u8 = 0, timeout: u8 = 0, retry_cnt: u8 = 0, rnr_retry: u8 = 0, + alt_port_num: u8 = 0, alt_timeout: u8 = 0, rate_limit: u32 = 0, _tail: u32 = 0, +}; +const Sge = extern struct { addr: u64, length: u32, lkey: u32 }; +const SendWr = extern struct { + wr_id: u64 = 0, next: ?*SendWr = null, sg_list: ?*Sge = null, num_sge: c_int = 0, opcode: u32 = 0, send_flags: u32 = 0, + imm_data: u32 = 0, remote_addr: u64 = 0, rkey: u32 = 0, _tail: [128 - 52]u8 = @splat(0), +}; +const Wc = extern struct { wr_id: u64, status: u32, opcode: u32, vendor_err: u32, _rest: [48 - 20]u8 }; +const PortAttr = extern struct { bytes: [56]u8 align(8) }; + +comptime { + std.debug.assert(@sizeOf(QpInitAttr) == 64 and @offsetOf(QpInitAttr, "cap") == 32 and @offsetOf(QpInitAttr, "qp_type") == 52); + std.debug.assert(@sizeOf(AhAttr) == 32 and @offsetOf(AhAttr, "dlid") == 24 and @offsetOf(AhAttr, "port_num") == 30); + std.debug.assert(@sizeOf(QpAttr) == 144 and @offsetOf(QpAttr, "ah_attr") == 56 and @offsetOf(QpAttr, "pkey_index") == 120); + std.debug.assert(@offsetOf(QpAttr, "rnr_retry") == 132 and @offsetOf(QpAttr, "qp_access_flags") == 32); + std.debug.assert(@sizeOf(SendWr) == 128 and @offsetOf(SendWr, "remote_addr") == 40 and @offsetOf(SendWr, "rkey") == 48); + std.debug.assert(@sizeOf(Wc) == 48 and @offsetOf(Mr, "lkey") == 36 and @offsetOf(Qp, "qp_num") == 52); +} + +// struct ibv_context's ops: the two calls verbs.h inlines +const ops_poll_cq = 96; +const ops_post_send = 208; +const PollCq = *const fn (*Cq, c_int, [*]Wc) callconv(.c) c_int; +const PostSend = *const fn (*Qp, *SendWr, *?*SendWr) callconv(.c) c_int; + +const wr_rdma_write = 0; +const send_signaled = 2; +const send_inline = 8; +const qpt_rc = 2; +const qps_init = 1; +const qps_rtr = 2; +const qps_rts = 3; +const port_active = 4; +const access_local_write = 1; +const access_remote_write = 2; +const qp_state = 1; +const qp_access_flags = 8; +const qp_pkey_index = 16; +const qp_port = 32; +const qp_av = 128; +const qp_path_mtu = 256; +const qp_timeout = 512; +const qp_retry_cnt = 1024; +const qp_rnr_retry = 2048; +const qp_rq_psn = 4096; +const qp_max_qp_rd_atomic = 8192; +const qp_min_rnr_timer = 32768; +const qp_sq_psn = 65536; +const qp_max_dest_rd_atomic = 131072; +const qp_dest_qpn = 1048576; + +const Api = struct { + ibv_get_device_list: *const fn (?*c_int) callconv(.c) ?[*]?*Device, + ibv_free_device_list: *const fn ([*]?*Device) callconv(.c) void, + ibv_get_device_name: *const fn (*Device) callconv(.c) ?[*:0]const u8, + ibv_open_device: *const fn (*Device) callconv(.c) ?*Context, + ibv_close_device: *const fn (*Context) callconv(.c) c_int, + ibv_query_port: *const fn (*Context, u8, *PortAttr) callconv(.c) c_int, + ibv_query_gid: *const fn (*Context, u8, c_int, *[16]u8) callconv(.c) c_int, + ibv_alloc_pd: *const fn (*Context) callconv(.c) ?*Pd, + ibv_dealloc_pd: *const fn (*Pd) callconv(.c) c_int, + ibv_reg_mr: *const fn (*Pd, ?*anyopaque, usize, c_int) callconv(.c) ?*Mr, + ibv_dereg_mr: *const fn (*Mr) callconv(.c) c_int, + ibv_create_cq: *const fn (*Context, c_int, ?*anyopaque, ?*anyopaque, c_int) callconv(.c) ?*Cq, + ibv_destroy_cq: *const fn (*Cq) callconv(.c) c_int, + ibv_create_qp: *const fn (*Pd, *QpInitAttr) callconv(.c) ?*Qp, + ibv_destroy_qp: *const fn (*Qp) callconv(.c) c_int, + ibv_modify_qp: *const fn (*Qp, *QpAttr, c_int) callconv(.c) c_int, + ibv_wc_status_str: *const fn (u32) callconv(.c) ?[*:0]const u8, +}; + +fn opsCall(comptime T: type, ctx: *Context, offset: usize) T { + const slot: *const T = @ptrFromInt(@intFromPtr(ctx) + offset); + return slot.*; +} + +pub const Error = error{ VerbsUnavailable, NoHca, HcaFailed, PeerRefused, ProxyFailed, ExchangeTimedOut }; + +// -- the region's layout (the same on both ranks) --------------------------------------------------------------------- +const Layout = struct { + slot_bytes: usize, + recv: usize, + flags: usize, + send: usize, + ctrl: usize, + total: usize, + + fn of(max_bytes: usize) Layout { + const s = std.mem.alignForward(usize, max_bytes, 4096); + const flags = slots * s; + const send = flags + slots * max_hcas * flag_stride; + const ctrl = send + slots * s; + return .{ .slot_bytes = s, .recv = 0, .flags = flags, .send = send, .ctrl = ctrl, .total = ctrl + flag_stride }; + } +}; +// ctrl words +const c_seq = 0; +const c_nbytes = 1; // + slot +const c_error = 3; +const c_hca = 4; + +/// What a rank tells its peer to connect: its region, and per HCA its key, address and queue pair. +const Blob = extern struct { + ok: u32, + n_hca: u32, + slot_bytes: u64, + addr: u64, + rkey: [max_hcas]u32, + qpn: [max_hcas]u32, + mtu: [max_hcas]u32, + lid: [max_hcas]u16, + gid: [max_hcas][16]u8, +}; + +const Hca = struct { + ctx: ?*Context = null, + pd: ?*Pd = null, + mr: ?*Mr = null, + cq: ?*Cq = null, + qp: ?*Qp = null, + gid: [16]u8 = @splat(0), + gid_index: c_int = 0, + lid: u16 = 0, + mtu: u32 = 0, + outstanding: u32 = 0, +}; + +pub const Gather = struct { + gpa: std.mem.Allocator, + d: *const Driver, + lib: std.DynLib, + api: Api, + rank: u8, + hcas: [max_hcas]Hca = @splat(.{}), + n_hca: u32 = 0, + layout: Layout, + region: memory.HostBuffer, + region_dev: abi.DevicePtr, + counters: memory.DeviceBuffer, + module: Module, + f: Function, + max_bytes: usize, + peer_addr: u64 = 0, + peer_rkey: [max_hcas]u32 = @splat(0), + last_seq: u32 = 0, + running: std.atomic.Value(bool) = .init(false), + failed: std.atomic.Value(bool) = .init(false), + thread: ?std.Thread = null, + err: [256]u8 = @splat(0), + gathers: u64 = 0, + + /// Both ranks' exchange, or null on both when either rank cannot (no libibverbs, no active HCA, a failed + /// connection: each rank's verdict travels over ``link``, so the ranks agree). ``max_bytes``: a rank's shard at most. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const Driver, link: tp_link.Link, wanted: bool, max_bytes: usize) !?*Gather { + const g = try gpa.create(Gather); + var mine: Blob = std.mem.zeroes(Blob); + const ok = wanted and blk: { + g.setup(gpa, io, d, link.rank, max_bytes) catch |e| { + std.log.warn("RoCE all-gather: {s} ({s}); NCCL carries every gather", .{ @errorName(e), std.mem.sliceTo(&g.err, 0) }); + break :blk false; + }; + break :blk true; + }; + if (ok) g.blob(&mine); + var theirs: Blob = undefined; + try link.swap(std.mem.asBytes(&mine), std.mem.asBytes(&theirs)); + if (!ok or theirs.ok == 0 or theirs.n_hca != mine.n_hca or theirs.slot_bytes != mine.slot_bytes) { + if (ok) g.teardown(); + gpa.destroy(g); + return null; + } + const connected = blk: { + g.connect(&theirs) catch |e| { + std.log.warn("RoCE all-gather: {s} ({s}); NCCL carries every gather", .{ @errorName(e), std.mem.sliceTo(&g.err, 0) }); + break :blk false; + }; + break :blk true; + }; + var v: [1]u8 = .{@intFromBool(connected)}; + var pv: [1]u8 = undefined; + try link.swap(&v, &pv); + if (!connected or pv[0] == 0) { + g.teardown(); + gpa.destroy(g); + return null; + } + std.log.info("RoCE all-gather: {d} HCA(s), shards up to {d} KiB by RDMA writes", .{ g.n_hca, max_bytes >> 10 }); + return g; + } + + fn fail(g: *Gather, comptime fmt: []const u8, args: anytype) void { + _ = std.fmt.bufPrint(g.err[0 .. g.err.len - 1], fmt, args) catch {}; + } + + fn setup(g: *Gather, gpa: std.mem.Allocator, io: std.Io, d: *const Driver, rank: u8, max_bytes: usize) !void { + g.* = .{ .gpa = gpa, .d = d, .lib = undefined, .api = undefined, .rank = rank, .layout = .of(max_bytes), .region = undefined, + .region_dev = 0, .counters = undefined, .module = undefined, .f = undefined, .max_bytes = max_bytes }; + g.lib = std.DynLib.open("libibverbs.so.1") catch return error.VerbsUnavailable; + errdefer g.lib.close(); + const info = @typeInfo(Api).@"struct"; + inline for (info.field_names, info.field_types) |name, T| { + @field(g.api, name) = g.lib.lookup(T, name) orelse return error.VerbsUnavailable; + } + var names: [max_hcas][64]u8 = undefined; + var lens: [max_hcas]usize = undefined; + g.n_hca = try hcaNames(io, gpa, &names, &lens); + g.region = try memory.HostBuffer.allocMapped(d, g.layout.total); + errdefer g.region.free(); + @memset(g.region.bytes, 0); + g.region_dev = try g.region.device(); + g.counters = try memory.DeviceBuffer.alloc(d, 64 * 4); + errdefer g.counters.free(); + const zeros: [64 * 4]u8 = @splat(0); + try g.counters.upload(0, &zeros); + g.module = try Module.load(d, kernels.roce_gather); + errdefer g.module.unload(); + g.f = try g.module.function("roce_gather"); + errdefer g.closeHcas(); + for (0..g.n_hca) |h| try g.openHca(io, gpa, h, names[h][0..lens[h]]); + } + + fn closeHcas(g: *Gather) void { + for (&g.hcas) |*x| { + if (x.qp) |q| _ = g.api.ibv_destroy_qp(q); + if (x.cq) |q| _ = g.api.ibv_destroy_cq(q); + if (x.mr) |m| _ = g.api.ibv_dereg_mr(m); + if (x.pd) |p| _ = g.api.ibv_dealloc_pd(p); + if (x.ctx) |c| _ = g.api.ibv_close_device(c); + x.* = .{}; + } + } + + fn openHca(g: *Gather, io: std.Io, gpa: std.mem.Allocator, h: usize, name: []const u8) !void { + const a = &g.api; + var num: c_int = 0; + const list = a.ibv_get_device_list(&num) orelse return error.NoHca; + defer a.ibv_free_device_list(list); + var dev: ?*Device = null; + for (0..@intCast(num)) |i| { + const n = a.ibv_get_device_name(list[i].?) orelse continue; + if (std.mem.eql(u8, std.mem.span(n), name)) dev = list[i]; + } + const hca = &g.hcas[h]; + hca.ctx = a.ibv_open_device(dev orelse { + g.fail("RDMA device {s} not found", .{name}); + return error.NoHca; + }) orelse return error.HcaFailed; + var port: PortAttr = std.mem.zeroes(PortAttr); + if (a.ibv_query_port(hca.ctx.?, port_num, &port) != 0) return error.HcaFailed; + if (std.mem.readInt(u32, port.bytes[0..4], .little) != port_active) { + g.fail("{s} port 1 is not active", .{name}); + return error.NoHca; + } + hca.mtu = std.mem.readInt(u32, port.bytes[8..12], .little); + hca.lid = std.mem.readInt(u16, port.bytes[34..36], .little); + hca.gid_index = try gidIndex(io, gpa, name); + if (a.ibv_query_gid(hca.ctx.?, port_num, hca.gid_index, &hca.gid) != 0) return error.HcaFailed; + hca.pd = a.ibv_alloc_pd(hca.ctx.?) orelse return error.HcaFailed; + hca.mr = a.ibv_reg_mr(hca.pd.?, g.region.bytes.ptr, g.region.bytes.len, access_local_write | access_remote_write) orelse { + g.fail("{s}: registering the pinned region failed", .{name}); + return error.HcaFailed; + }; + hca.cq = a.ibv_create_cq(hca.ctx.?, send_depth, null, null, 0) orelse return error.HcaFailed; + var attr: QpInitAttr = .{ .send_cq = hca.cq, .recv_cq = hca.cq, .qp_type = qpt_rc, + .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 } }; + hca.qp = a.ibv_create_qp(hca.pd.?, &attr) orelse return error.HcaFailed; + var init_attr: QpAttr = .{ .qp_state = qps_init, .port_num = port_num, .qp_access_flags = access_remote_write }; + if (a.ibv_modify_qp(hca.qp.?, &init_attr, qp_state | qp_pkey_index | qp_port | qp_access_flags) != 0) return error.HcaFailed; + } + + fn blob(g: *Gather, b: *Blob) void { + b.ok = 1; + b.n_hca = g.n_hca; + b.slot_bytes = g.layout.slot_bytes; + b.addr = @intFromPtr(g.region.bytes.ptr); + for (0..g.n_hca) |h| { + const x = &g.hcas[h]; + b.rkey[h] = x.mr.?.rkey; + b.qpn[h] = x.qp.?.qp_num; + b.mtu[h] = x.mtu; + b.lid[h] = x.lid; + b.gid[h] = x.gid; + } + } + + fn connect(g: *Gather, peer: *const Blob) !void { + g.peer_addr = peer.addr; + for (0..g.n_hca) |h| { + const x = &g.hcas[h]; + g.peer_rkey[h] = peer.rkey[h]; + var rtr: QpAttr = .{ .qp_state = qps_rtr, .path_mtu = @min(peer.mtu[h], x.mtu), .dest_qp_num = peer.qpn[h], .max_dest_rd_atomic = 1, + .min_rnr_timer = 12, .ah_attr = .{ .is_global = 1, .dlid = peer.lid[h], .port_num = port_num, + .grh = .{ .dgid = peer.gid[h], .sgid_index = @intCast(x.gid_index), .hop_limit = 64 } } }; + if (g.api.ibv_modify_qp(x.qp.?, &rtr, qp_state | qp_av | qp_path_mtu | qp_dest_qpn | qp_rq_psn | qp_max_dest_rd_atomic | qp_min_rnr_timer) != 0) { + g.fail("HCA {d}: the queue pair refused RTR", .{h}); + return error.PeerRefused; + } + var rts: QpAttr = .{ .qp_state = qps_rts, .timeout = 14, .retry_cnt = 7, .rnr_retry = 7, .max_rd_atomic = 1 }; + if (g.api.ibv_modify_qp(x.qp.?, &rts, qp_state | qp_timeout | qp_retry_cnt | qp_rnr_retry | qp_sq_psn | qp_max_qp_rd_atomic) != 0) { + g.fail("HCA {d}: the queue pair refused RTS", .{h}); + return error.PeerRefused; + } + } + g.running.store(true, .release); + g.thread = try std.Thread.spawn(.{}, proxy, .{g}); + } + + fn ctrl(g: *Gather) [*]u32 { + return @ptrCast(@alignCast(g.region.bytes.ptr + g.layout.ctrl)); + } + + fn drain(g: *Gather, h: usize) !void { + const x = &g.hcas[h]; + var wc: [32]Wc = undefined; + const n = opsCall(PollCq, x.ctx.?, ops_poll_cq)(x.cq.?, wc.len, &wc); + if (n < 0) return error.ProxyFailed; + for (wc[0..@intCast(n)]) |w| { + if (w.status != 0) { + const why = g.api.ibv_wc_status_str(w.status) orelse "?"; + g.fail("RDMA write on HCA {d} failed: {s} (vendor 0x{x}, seq {d})", .{ h, std.mem.span(why), w.vendor_err, g.last_seq }); + return error.ProxyFailed; + } + x.outstanding -= 1; + } + } + + /// One sequence's shard to the peer: per HCA a stripe, then its flag (inline, signalled) on the same queue pair. + fn post(g: *Gather, seq: u32, nbytes: u32) !void { + if (nbytes == 0 or nbytes % 16 != 0 or nbytes > g.layout.slot_bytes) { + g.fail("bad shard size {d} at seq {d}", .{ nbytes, seq }); + return error.ProxyFailed; + } + const slot: u64 = seq & 1; + const base = @intFromPtr(g.region.bytes.ptr); + const packs = nbytes / 16; + var seq_copy: u32 = seq; + var offset: u64 = 0; + for (0..g.n_hca) |h| { + const x = &g.hcas[h]; + while (x.outstanding >= send_depth / 4) { + try g.drain(h); + if (!g.running.load(.monotonic)) return error.ProxyFailed; + } + var n = packs / g.n_hca; + if (h < packs % g.n_hca) n += 1; + const bytes: u32 = n * 16; + var flag_sge: Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; + var flag: SendWr = .{ .wr_id = h, .sg_list = &flag_sge, .num_sge = 1, .opcode = wr_rdma_write, .send_flags = send_signaled | send_inline, + .remote_addr = g.peer_addr + g.layout.flags + (slot * g.n_hca + h) * flag_stride, .rkey = g.peer_rkey[h] }; + var data_sge: Sge = .{ .addr = base + g.layout.send + slot * g.layout.slot_bytes + offset, .length = bytes, .lkey = x.mr.?.lkey }; + var data: SendWr = .{ .wr_id = h, .next = &flag, .sg_list = &data_sge, .num_sge = 1, .opcode = wr_rdma_write, + .remote_addr = g.peer_addr + g.layout.recv + slot * g.layout.slot_bytes + offset, .rkey = g.peer_rkey[h] }; + var bad: ?*SendWr = null; + const first: *SendWr = if (bytes != 0) &data else &flag; + const rc = opsCall(PostSend, x.ctx.?, ops_post_send)(x.qp.?, first, &bad); + if (rc != 0) { + g.fail("ibv_post_send: {d}", .{rc}); + return error.ProxyFailed; + } + x.outstanding += 1; + offset += bytes; + } + for (0..g.n_hca) |h| try g.drain(h); + } + + /// The proxy: spins on the doorbell and posts each sequence in order (a missed one from its slot's size). + fn proxy(g: *Gather) void { + const c = g.ctrl(); + var idle: u64 = 0; + const nap: std.c.timespec = .{ .sec = 0, .nsec = 20_000 }; + while (g.running.load(.monotonic)) { + const seq = @atomicLoad(u32, &c[c_seq], .acquire); + if (seq == g.last_seq) { + idle += 1; + if (idle % 64 == 0) for (0..g.n_hca) |h| g.drain(h) catch { + g.failed.store(true, .release); + return; + }; + if (idle >= idle_spins) _ = std.c.nanosleep(&nap, null); + continue; + } + idle = 0; + // a rank's next gather can ring before this thread saw the last one; peers stay within a slot's distance + var pending = seq -% g.last_seq; + if (pending > slots) { + g.fail("the doorbell skipped {d} sequences (last {d}, now {d})", .{ pending, g.last_seq, seq }); + g.failed.store(true, .release); + return; + } + var s = g.last_seq +% 1; + while (pending > 0) : ({ + s +%= 1; + pending -= 1; + }) { + const nbytes = @atomicLoad(u32, &c[c_nbytes + (s & 1)], .acquire); + g.post(s, nbytes) catch { + g.failed.store(true, .release); + return; + }; + g.last_seq = s; + } + } + } + + /// Whether a gather of ``bytes`` a rank (from and to 16-byte aligned buffers) goes this way. + pub fn fits(g: *const Gather, bytes: usize, send: abi.DevicePtr, recv: abi.DevicePtr) bool { + return bytes > 0 and bytes <= g.max_bytes and bytes % 16 == 0 and send % 16 == 0 and recv % 16 == 0; + } + + /// Both ranks' ``bytes`` from ``send`` into ``recv`` in rank order (as ncclAllGather), on ``stream``. + pub fn gather(g: *Gather, stream: abi.Stream, send: abi.DevicePtr, recv: abi.DevicePtr, bytes: usize) !void { + const packs = bytes / 16; + const need = @max(1, std.math.divCeil(usize, packs, 2 * threads) catch unreachable); + const grid: u32 = @intCast(@min(std.math.ceilPowerOfTwo(usize, need) catch max_blocks, max_blocks)); + const klass: i32 = @intCast(std.math.log2_int(u32, grid)); + const base = g.region_dev; + var a: launch.Args = .{}; + a.add(send); + a.add(recv); + a.add(@as(i32, @intCast(packs))); + a.add(@as(i32, g.rank)); + a.add(base + g.layout.recv); + a.add(base + g.layout.flags); + a.add(base + g.layout.send); + a.add(base + g.layout.ctrl); + a.add(@as(i64, @intCast(g.layout.slot_bytes))); + a.add(g.counters.ptr); + a.add(klass); + a.add(spin_limit); + a.add(@as(i32, @intCast(g.n_hca))); + try launch.launch(g.f, .{ .grid = .{ .x = grid }, .block = .{ .x = threads } }, .{ .handle = stream, .d = g.d }, &a); + g.gathers += 1; + } + + /// An error if a wait timed out on the GPU or the proxy failed (the host checks after a synchronize). + pub fn check(g: *Gather) !void { + const c = g.ctrl(); + const bad = @atomicLoad(u32, &c[c_error], .acquire); + if (bad == 0 and !g.failed.load(.acquire)) return; + std.log.err("RoCE all-gather failed (sequence {d}, HCA {d}): {s}", .{ bad, c[c_hca], std.mem.sliceTo(&g.err, 0) }); + return error.ProxyFailed; + } + + fn teardown(g: *Gather) void { + if (g.running.swap(false, .acq_rel)) if (g.thread) |t| t.join(); + g.closeHcas(); + g.module.unload(); + g.counters.free(); + g.region.free(); + g.lib.close(); + } + + pub fn deinit(g: *Gather) void { + g.teardown(); + g.gpa.destroy(g); + } +}; + +/// The HCAs to stripe over, in the order both ranks name them: TF_ROCE_HCA, else NCCL_IB_HCA ("=^" prefixes and +/// ":port" suffixes dropped), at most two. +fn hcaNames(io: std.Io, gpa: std.mem.Allocator, names: *[max_hcas][64]u8, lens: *[max_hcas]usize) !u32 { + _ = io; + _ = gpa; + const raw = std.c.getenv("TF_ROCE_HCA") orelse std.c.getenv("NCCL_IB_HCA") orelse return error.NoHca; + var n: u32 = 0; + var it = std.mem.tokenizeScalar(u8, std.mem.span(raw), ','); + while (it.next()) |w| { + if (n == max_hcas) break; + var s = std.mem.trimStart(u8, std.mem.trim(u8, w, " "), "=^"); + if (std.mem.indexOfScalar(u8, s, ':')) |c| s = s[0..c]; + if (s.len == 0 or s.len > 63) continue; + @memcpy(names[n][0..s.len], s); + lens[n] = s.len; + n += 1; + } + if (n == 0) return error.NoHca; + return n; +} + +/// The HCA's RoCE v2 IPv4 GID (TF_ROCE_GID_INDEX overrides): its index differs between hosts' tables. +fn gidIndex(io: std.Io, gpa: std.mem.Allocator, name: []const u8) !c_int { + if (std.c.getenv("TF_ROCE_GID_INDEX")) |v| return std.fmt.parseInt(c_int, std.mem.span(v), 10) catch error.NoHca; + const cwd = std.Io.Dir.cwd(); + var i: c_int = 0; + while (i < 16) : (i += 1) { + var pbuf: [160]u8 = undefined; + const tpath = try std.fmt.bufPrint(&pbuf, "/sys/class/infiniband/{s}/ports/1/gid_attrs/types/{d}", .{ name, i }); + const t = cwd.readFileAlloc(io, tpath, gpa, .limited(64)) catch continue; + defer gpa.free(t); + if (!std.mem.startsWith(u8, t, "RoCE v2")) continue; + const gpath = try std.fmt.bufPrint(&pbuf, "/sys/class/infiniband/{s}/ports/1/gids/{d}", .{ name, i }); + const gid = cwd.readFileAlloc(io, gpath, gpa, .limited(64)) catch continue; + defer gpa.free(gid); + if (std.mem.startsWith(u8, gid, "0000:0000:0000:0000:0000:ffff:")) return i; + } + return error.NoHca; +} diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..40bd1bb6a0 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -21,6 +21,8 @@ pub const Dim3 = launch.Dim3; pub const graph = @import("graph.zig"); pub const cublaslt = @import("cublaslt.zig"); pub const nccl = @import("nccl.zig"); +pub const tp_link = @import("tp_link.zig"); +pub const roce = @import("roce.zig"); pub const triton = @import("triton.zig"); pub const aot = @import("aot.zig"); pub const kernels = @import("kernels.zig"); @@ -54,4 +56,5 @@ test { _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); + _ = @import("tp_link.zig"); } diff --git a/zig/src/cuda/tp_link.zig b/zig/src/cuda/tp_link.zig new file mode 100644 index 0000000000..1d74b64c50 --- /dev/null +++ b/zig/src/cuda/tp_link.zig @@ -0,0 +1,170 @@ +//! Two ranks over the hosts' network: a TCP control link (rank 0 listens, rank 1 connects; framed messages carry +//! requests and per-round decisions) and the NCCL communicator it brings up (rank 0 makes the unique id and sends it). + +const std = @import("std"); +const posix = std.posix; +const net = std.Io.net; +const Threaded = std.Io.Threaded; +const nccl = @import("nccl.zig"); + +pub const Error = error{ SocketFailed, BindFailed, ListenFailed, AcceptFailed, ConnectTimedOut, PeerClosed, IoFailed, FrameTooLarge }; + +pub const max_frame = 64 << 20; + +/// A message: a tag the protocol defines and its bytes (owned by the caller's allocator). +pub const Message = struct { tag: u32, bytes: []u8 }; + +pub const Link = struct { + fd: posix.socket_t, + rank: u8, + + /// Rank 0: listen on ``port`` (all interfaces) and accept rank 1, waiting up to ``timeout_ms``. + pub fn lead(port: u16, timeout_ms: i32) Error!Link { + const address = net.IpAddress.parse("0.0.0.0", port) catch unreachable; + var storage: Threaded.PosixAddress = undefined; + const len = Threaded.addressToPosix(&address, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const lfd: posix.socket_t = @intCast(rc); + defer _ = posix.system.close(lfd); + const one: c_int = 1; + posix.setsockopt(lfd, posix.SOL.SOCKET, posix.SO.REUSEADDR, std.mem.asBytes(&one)) catch return error.SocketFailed; + if (posix.errno(posix.system.bind(lfd, &storage.any, len)) != .SUCCESS) return error.BindFailed; + if (posix.errno(posix.system.listen(lfd, 1)) != .SUCCESS) return error.ListenFailed; + var fds = [_]posix.pollfd{.{ .fd = lfd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, timeout_ms) catch return error.AcceptFailed; + if (ready == 0) return error.ConnectTimedOut; + const a = posix.system.accept(lfd, null, null); + if (posix.errno(a) != .SUCCESS) return error.AcceptFailed; + const fd: posix.socket_t = @intCast(a); + noDelay(fd); + return .{ .fd = fd, .rank = 0 }; + } + + /// Rank 1: connect to rank 0 at ``address``, retrying every 200 ms for up to ``timeout_ms``. + pub fn follow(io: std.Io, address: net.IpAddress, timeout_ms: i64) Error!Link { + var waited: i64 = 0; + while (true) { + var storage: Threaded.PosixAddress = undefined; + const len = Threaded.addressToPosix(&address, &storage); + const family: u32 = if (address == .ip4) posix.AF.INET else posix.AF.INET6; + const rc = posix.system.socket(family, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const fd: posix.socket_t = @intCast(rc); + if (posix.errno(posix.system.connect(fd, &storage.any, len)) == .SUCCESS) { + noDelay(fd); + return .{ .fd = fd, .rank = 1 }; + } + _ = posix.system.close(fd); + if (waited >= timeout_ms) return error.ConnectTimedOut; + std.Io.sleep(io, .fromMilliseconds(200), .awake) catch {}; + waited += 200; + } + } + + pub fn close(self: *Link) void { + _ = posix.system.shutdown(self.fd, posix.SHUT.RDWR); + _ = posix.system.close(self.fd); + self.* = undefined; + } + + /// One frame: tag and length (little-endian u32 each), then the bytes. + pub fn send(self: Link, tag: u32, bytes: []const u8) Error!void { + if (bytes.len > max_frame) return error.FrameTooLarge; + var head: [8]u8 = undefined; + std.mem.writeInt(u32, head[0..4], tag, .little); + std.mem.writeInt(u32, head[4..8], @intCast(bytes.len), .little); + try writeAll(self.fd, &head); + try writeAll(self.fd, bytes); + } + + /// The next frame; its bytes allocated with ``gpa``. + pub fn recv(self: Link, gpa: std.mem.Allocator) (Error || std.mem.Allocator.Error)!Message { + var head: [8]u8 = undefined; + try readAll(self.fd, &head); + const tag = std.mem.readInt(u32, head[0..4], .little); + const n = std.mem.readInt(u32, head[4..8], .little); + if (n > max_frame) return error.FrameTooLarge; + const bytes = try gpa.alloc(u8, n); + errdefer gpa.free(bytes); + try readAll(self.fd, bytes); + return .{ .tag = tag, .bytes = bytes }; + } + + /// Both ranks' ``mine`` exchanged (equal lengths): returns the peer's into ``theirs``. + pub fn swap(self: Link, mine: []const u8, theirs: []u8) Error!void { + if (self.rank == 0) { + try writeAll(self.fd, mine); + try readAll(self.fd, theirs); + } else { + try readAll(self.fd, theirs); + try writeAll(self.fd, mine); + } + } +}; + +fn noDelay(fd: posix.socket_t) void { + const one: c_int = 1; + posix.setsockopt(fd, posix.IPPROTO.TCP, std.posix.TCP.NODELAY, std.mem.asBytes(&one)) catch {}; +} + +fn writeAll(fd: posix.socket_t, bytes: []const u8) Error!void { + var at: usize = 0; + while (at < bytes.len) { + const rc = posix.system.write(fd, bytes[at..].ptr, bytes.len - at); + switch (posix.errno(rc)) { + .SUCCESS => at += @intCast(rc), + .INTR, .AGAIN => continue, + .PIPE, .CONNRESET => return error.PeerClosed, + else => return error.IoFailed, + } + } +} + +fn readAll(fd: posix.socket_t, bytes: []u8) Error!void { + var at: usize = 0; + while (at < bytes.len) { + const rc = posix.system.read(fd, bytes[at..].ptr, bytes.len - at); + switch (posix.errno(rc)) { + .SUCCESS => { + if (rc == 0) return error.PeerClosed; + at += @intCast(rc); + }, + .INTR, .AGAIN => continue, + .CONNRESET => return error.PeerClosed, + else => return error.IoFailed, + } + } +} + +/// The two ranks' NCCL communicator: rank 0's unique id crosses the link, then both join. +pub fn communicator(lib: *nccl.Library, link: Link) !nccl.Comm { + var uid: nccl.UniqueId = undefined; + if (link.rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&uid), "ncclGetUniqueId"); + try link.send(1, &uid.internal); + } else { + var buf: [256]u8 = undefined; + var fba = std.heap.FixedBufferAllocator.init(&buf); + const m = try link.recv(fba.allocator()); + if (m.tag != 1 or m.bytes.len != uid.internal.len) return error.IoFailed; + @memcpy(&uid.internal, m.bytes); + } + var comm: nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, 2, uid, link.rank), "ncclCommInitRank"); + return comm; +} + +test "frames round-trip through a socket pair" { + var fds: [2]posix.socket_t = undefined; + if (std.c.socketpair(posix.AF.UNIX, posix.SOCK.STREAM, 0, &fds) != 0) return error.SkipZigTest; + const a: Link = .{ .fd = fds[0], .rank = 0 }; + const b: Link = .{ .fd = fds[1], .rank = 1 }; + try a.send(7, "hello ranks"); + const m = try b.recv(std.testing.allocator); + defer std.testing.allocator.free(m.bytes); + try std.testing.expectEqual(@as(u32, 7), m.tag); + try std.testing.expectEqualStrings("hello ranks", m.bytes); + _ = posix.system.close(fds[0]); + _ = posix.system.close(fds[1]); +} From 39de0393a2a3d52844d8d0d7714504d0261aedfb Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 18:36:00 +0000 Subject: [PATCH 02/31] 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 cbe08cf7f4b3e917b2515fa819e1ec1a98348b99 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 18:38:18 +0000 Subject: [PATCH 03/31] GLM-5.3-Flash on two CUDA ranks: the native family (zig/src/families/glm_cuda), its kernels (zig/kernels/cuda/glm: the dense matmul at 4/6/8/16 bits, the EXL3 trellis experts with the fused prompt kernel, latent attention, the indexer, KDA, the hyper-connections, the vision tower's helpers, DFlash2), tf-cuda-test glm-batch/glm-prof/glm-tf, -Dglm (cherry picked from commit 611959f15964cda97ecd12d0e459ea4ac0e84f78) --- zig/build/cuda.zig | 60 +- zig/kernels/cuda/glm/dflash.cu | 360 +++++++ zig/kernels/cuda/glm/exl3.cu | 677 ++++++++++++ zig/kernels/cuda/glm/experts_sym.cu | 302 ++++++ zig/kernels/cuda/glm/experts_sym.cuh | 137 +++ zig/kernels/cuda/glm/experts_sym_prefill.cu | 175 +++ zig/kernels/cuda/glm/index.cu | 250 +++++ zig/kernels/cuda/glm/kda.cu | 416 ++++++++ zig/kernels/cuda/glm/latent.cu | 287 +++++ zig/kernels/cuda/glm/mm.cu | 299 ++++++ zig/kernels/cuda/glm/ops.cu | 195 ++++ zig/kernels/cuda/glm/rows.cu | 426 ++++++++ zig/kernels/cuda/glm/vision.cu | 176 ++++ zig/src/cuda/cublaslt.zig | 9 +- zig/src/cuda/kernels.zig | 15 +- zig/src/families/glm_cuda/api.zig | 141 +++ zig/src/families/glm_cuda/blocks.zig | 462 ++++++++ zig/src/families/glm_cuda/calibrate.zig | 132 +++ zig/src/families/glm_cuda/copy.zig | 117 +++ zig/src/families/glm_cuda/decode.zig | 336 ++++++ zig/src/families/glm_cuda/dflash.zig | 827 +++++++++++++++ zig/src/families/glm_cuda/draw.zig | 171 +++ zig/src/families/glm_cuda/forward.zig | 499 +++++++++ zig/src/families/glm_cuda/guard.zig | 136 +++ zig/src/families/glm_cuda/kern.zig | 590 +++++++++++ zig/src/families/glm_cuda/kern_exl3.zig | 137 +++ zig/src/families/glm_cuda/native.zig | 1014 ++++++++++++++++++ zig/src/families/glm_cuda/proto.zig | 64 ++ zig/src/families/glm_cuda/root.zig | 3 + zig/src/families/glm_cuda/seq.zig | 233 ++++ zig/src/families/glm_cuda/vision.zig | 389 +++++++ zig/src/families/glm_cuda/weights.zig | 1051 +++++++++++++++++++ zig/src/native/cuda.zig | 4 +- zig/src/native/no_family.zig | 2 - zig/tests/cuda/glm/gen_copies.py | 24 + zig/tests/cuda/glm_forward.zig | 709 +++++++++++++ zig/tests/cuda/glm_weights.zig | 78 ++ zig/tests/cuda/main.zig | 16 + 38 files changed, 10902 insertions(+), 17 deletions(-) create mode 100644 zig/kernels/cuda/glm/dflash.cu create mode 100644 zig/kernels/cuda/glm/exl3.cu create mode 100644 zig/kernels/cuda/glm/experts_sym.cu create mode 100644 zig/kernels/cuda/glm/experts_sym.cuh create mode 100644 zig/kernels/cuda/glm/experts_sym_prefill.cu create mode 100644 zig/kernels/cuda/glm/index.cu create mode 100644 zig/kernels/cuda/glm/kda.cu create mode 100644 zig/kernels/cuda/glm/latent.cu create mode 100644 zig/kernels/cuda/glm/mm.cu create mode 100644 zig/kernels/cuda/glm/ops.cu create mode 100644 zig/kernels/cuda/glm/rows.cu create mode 100644 zig/kernels/cuda/glm/vision.cu create mode 100644 zig/src/families/glm_cuda/api.zig create mode 100644 zig/src/families/glm_cuda/blocks.zig create mode 100644 zig/src/families/glm_cuda/calibrate.zig create mode 100644 zig/src/families/glm_cuda/copy.zig create mode 100644 zig/src/families/glm_cuda/decode.zig create mode 100644 zig/src/families/glm_cuda/dflash.zig create mode 100644 zig/src/families/glm_cuda/draw.zig create mode 100644 zig/src/families/glm_cuda/forward.zig create mode 100644 zig/src/families/glm_cuda/guard.zig create mode 100644 zig/src/families/glm_cuda/kern.zig create mode 100644 zig/src/families/glm_cuda/kern_exl3.zig create mode 100644 zig/src/families/glm_cuda/native.zig create mode 100644 zig/src/families/glm_cuda/proto.zig create mode 100644 zig/src/families/glm_cuda/root.zig create mode 100644 zig/src/families/glm_cuda/seq.zig create mode 100644 zig/src/families/glm_cuda/vision.zig create mode 100644 zig/src/families/glm_cuda/weights.zig delete mode 100644 zig/src/native/no_family.zig create mode 100644 zig/tests/cuda/glm/gen_copies.py create mode 100644 zig/tests/cuda/glm_forward.zig create mode 100644 zig/tests/cuda/glm_weights.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 2cf45556de..9b4d4348a8 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -43,13 +43,27 @@ const kernels = [_]Kernel{ .{ .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 + // GLM-5.3-Flash (glm5_next): the device copy of the Python engine's KDA extension (zig/tests/cuda/glm/gen_copies.py) + .{ .name = "glm_kda", .src = "glm/kda", .flags = &.{ "-O3", "--fmad=false" } }, // glm5_next/cuda/kda.py, tensorfold_glm_kda_v2 + .{ .name = "glm_ops", .src = "glm/ops", .flags = &.{ "-O3", "--fmad=false" } }, // ours: the token lists, casts and positions torch ops make + .{ .name = "glm_experts", .src = "glm/experts_sym", .flags = &.{"-O3"} }, // ours: fn_experts on symmetric groups (scales alone) + .{ .name = "glm_experts_prefill", .src = "glm/experts_sym_prefill", .flags = &.{"-O3"} }, + .{ .name = "glm_vision", .src = "glm/vision", .flags = &.{"-O3"} }, // ours: the vision tower's row and elementwise steps + .{ .name = "glm_mm", .src = "glm/mm", .flags = &.{"-O3"} }, // ours: the dense linears on tensor cores, any row count + .{ .name = "glm_rows", .src = "glm/rows", .flags = &.{ "-O3", "--fmad=false" } }, // ours: norms, hyper-connections, router, combine + .{ .name = "glm_latent", .src = "glm/latent", .flags = &.{"-O3"} }, // ours: latent MLA attention, per-row caches + .{ .name = "glm_index", .src = "glm/index", .flags = &.{"-O3"} }, // ours: the DSA indexer, per-row rings and pools + .{ .name = "glm_dflash", .src = "glm/dflash", .flags = &.{"-O3"} }, // ours: the DFlash2 block drafter (opt-in) + .{ .name = "roce_gather", .src = "roce_gather", .flags = &.{"-O3"} }, // ours: b12x's one-shot RoCE all-gather, two ranks + .{ .name = "glm_exl3", .src = "glm/exl3", .flags = &.{"-O3"} }, // python-0.6's EXL3 routed experts (mcg, 4 bits), without torch + .{ .name = "glm_experts3", .src = "glm/experts_sym", .flags = &.{ "-O3", "-DGLM_SYM3" } }, // ours: 3-bit codes in planes + .{ .name = "glm_experts3_prefill", .src = "glm/experts_sym_prefill", .flags = &.{ "-O3", "-DGLM_SYM3" } }, .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, .{ .name = "torch_indexing", .src = "torch_ops/indexing", .flags = torch_ops }, .{ .name = "torch_movement", .src = "torch_ops/movement", .flags = torch_ops }, .{ .name = "torch_nemotron_constants", .src = "torch_ops/nemotron_constants", .flags = torch_ops }, - .{ .name = "roce_gather", .src = "roce_gather", .flags = &.{"-O3"} }, // b12x's one-shot RoCE all-gather, two ranks }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. @@ -80,7 +94,15 @@ 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 } { +var glm_option: ?bool = null; + +/// -Dglm, read once. +fn glmOption(b: *std.Build) bool { + if (glm_option == null) glm_option = b.option(bool, "glm", "register the GLM-5.3-Flash CUDA family (two ranks)") orelse false; + return glm_option.?; +} + +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, glm: ?*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); @@ -92,7 +114,14 @@ 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 }; + // GLM-5.3-Flash on two GB10s (zig/src/families/glm_cuda), registered when built with -Dglm + const glm: ?*std.Build.Module = if (glmOption(b)) blk: { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/glm_cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("lanes", lanes); + break :blk m; + } else null; + return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat, .glm = glm }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -124,14 +153,13 @@ 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, null); 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 })); 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); @@ -141,12 +169,18 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen.zig"), .target = target, .optimize = optimize })); 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); + runner.addImport("qwen_image", vision_build.module(b, target, optimize)); + // the GLM weight loader alone (tf-cuda-test glm-weights checks it against the Python engine's hashes) + const glm_weights = b.createModule(.{ .root_source_file = b.path("zig/src/families/glm_cuda/weights.zig"), .target = target, .optimize = optimize, .link_libc = true }); + glm_weights.addImport("cuda", cuda); + glm_weights.addImport("lanes", mods.lanes); + runner.addImport("glm_weights", glm_weights); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - 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.glm, 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, glm: ?*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 }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), @@ -155,6 +189,10 @@ fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron } }, }); + const fn_options = b.addOptions(); + fn_options.addOption(bool, "glm", glm != null); + mod.addOptions("cuda_families", fn_options); + mod.addImport("glm_cuda", glm orelse b.createModule(.{ .root_source_file = b.path("zig/src/native/no_family_glm.zig"), .target = target, .optimize = optimize })); return .{ .api = api, .engines = mod }; } @@ -164,8 +202,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, glm: ?*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, glm); // 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(.{ @@ -193,7 +231,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, null); 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 }); @@ -249,7 +287,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.glm, 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/kernels/cuda/glm/dflash.cu b/zig/kernels/cuda/glm/dflash.cu new file mode 100644 index 0000000000..1292edc6a8 --- /dev/null +++ b/zig/kernels/cuda/glm/dflash.cu @@ -0,0 +1,360 @@ +// DFlash2 block drafter kernels (our own, after z-lab/dflash's reference, MIT): the target's layer taps, the drafter's +// per-head norms with rotary positions, its grouped two-tap dynamic convolutions, SwiGLU, and its attention over a +// stream's context ring (the target's features, a 2048-token window) and the 8-row block (bidirectional). +// Layout: hidden 4096, 32 query heads and 8 key/value heads of 128, a block of 8 rows a stream. + +#include +#include + +namespace { +constexpr int D = 4096; +constexpr int HD = 128; +constexpr int KVH = 8; +constexpr int QH = 32; +constexpr int KVW = KVH * HD; // 1024 +constexpr int WIN = 2048; // the sliding window (and the ring's slots) +constexpr int BLOCK = 8; +constexpr int GROUP = 16; // channels sharing one dynamic tap weight +constexpr int GROUPS = D / GROUP; +constexpr int CH = 16; // key chunks of the chunked attention (df_attn_part: a block a kv head and chunk) + +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ __nv_bfloat16 tobf(float x) { return __float2bfloat16_rn(x); } + +// The ring of one layer: K then V, [WIN][KVW] each, layers one after another. +__device__ __forceinline__ __nv_bfloat16* ringK(uint64_t base, int layer) { + return reinterpret_cast<__nv_bfloat16*>(base) + (long long)layer * 2 * WIN * KVW; +} +__device__ __forceinline__ __nv_bfloat16* ringV(uint64_t base, int layer) { return ringK(base, layer) + (long long)WIN * KVW; } + +// Qwen3's per-head RMSNorm then rotate-half RoPE of one head (64 threads: dims d and d + 64), in place or to out. +__device__ __forceinline__ void normRope(const __nv_bfloat16* x, const __nv_bfloat16* w, float eps, long long pos, float theta, + __nv_bfloat16* out, float* red) { + const int d = threadIdx.x; // 0..63 + const float a = bf(x[d]), b = bf(x[d + 64]); + float s = a * a + b * b; + for (int o = 16; o > 0; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o); + if ((d & 31) == 0) red[d >> 5] = s; + __syncthreads(); + const float rinv = rsqrtf((red[0] + red[1]) / HD + eps); + __syncthreads(); + const float na = bf(tobf(bf(w[d]) * bf(tobf(a * rinv)))); + const float nb = bf(tobf(bf(w[d + 64]) * bf(tobf(b * rinv)))); + const float inv = powf(theta, -(float)(2 * d) / HD); + const float ang = (float)pos * inv; + const float c = cosf(ang), sn = sinf(ang); + out[d] = tobf(na * c - nb * sn); + out[d + 64] = tobf(nb * c + na * sn); +} +} // namespace + +// The completed residual of a target layer, its 4 streams' mean, into column block ``slot`` of out [R][5 D]. +// grid (R, D / 1024), 256 threads. +// mode 1 (development, TF_GLM_DFLASH2_TAP=s0): stream 0 alone. +extern "C" __global__ void df_tap(const __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ out, int slot, int mode) { + const int r = blockIdx.x; + const __nv_bfloat16* xr = x + (long long)r * 4 * D; + __nv_bfloat16* o = out + (long long)r * 5 * D + (long long)slot * D; + for (int d = blockIdx.y * 1024 + threadIdx.x; d < (blockIdx.y + 1) * 1024; d += blockDim.x) { + if (mode == 1) { + o[d] = xr[d]; + continue; + } + const float a = ((bf(xr[d]) + bf(xr[D + d])) + bf(xr[2 * D + d])) + bf(xr[3 * D + d]); + o[d] = tobf(a / 4.0f); + } +} + +// The block's queries [R][QH*HD] (from the fused qkv rows, stride ldq) normed and rotated at pos[r] into q, and its +// keys likewise into kb, values copied into vb ([R][KVW] each). grid (R, QH + KVH), 64 threads. +extern "C" __global__ void df_block_qkv(const __nv_bfloat16* __restrict__ qkv, int ldq, const __nv_bfloat16* __restrict__ qn, + const __nv_bfloat16* __restrict__ kn, const int* __restrict__ pos, float eps, float theta, + __nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ kb, __nv_bfloat16* __restrict__ vb) { + __shared__ float red[2]; + const int r = blockIdx.x, h = blockIdx.y; + const __nv_bfloat16* row = qkv + (long long)r * ldq; + if (h < QH) { + normRope(row + h * HD, qn, eps, pos[r], theta, q + (long long)r * QH * HD + h * HD, red); + } else { + const int kh = h - QH; + normRope(row + QH * HD + kh * HD, kn, eps, pos[r], theta, kb + (long long)r * KVW + kh * HD, red); + const __nv_bfloat16* v = row + QH * HD + KVW + kh * HD; + __nv_bfloat16* vo = vb + (long long)r * KVW + kh * HD; + vo[threadIdx.x] = v[threadIdx.x]; + vo[threadIdx.x + 64] = v[threadIdx.x + 64]; + } +} + +// Context rows' keys and values ([R][2 KVW]: k | v) normed, rotated at pos[r] and written into each row's stream +// ring (rings[r]) at slot pos % WIN of ``layer``. grid (R, KVH), 64 threads. +extern "C" __global__ void df_ctx_kv(const __nv_bfloat16* __restrict__ kv, const __nv_bfloat16* __restrict__ kn, + const long long* __restrict__ pos, const uint64_t* __restrict__ rings, int layer, float eps, + float theta) { + __shared__ float red[2]; + const int r = blockIdx.x, kh = blockIdx.y; + const long long p = pos[r]; + const int slot = (int)(p % WIN); + const __nv_bfloat16* row = kv + (long long)r * 2 * KVW; + normRope(row + kh * HD, kn, eps, p, theta, ringK(rings[r], layer) + (long long)slot * KVW + kh * HD, red); + __nv_bfloat16* vo = ringV(rings[r], layer) + (long long)slot * KVW + kh * HD; + vo[threadIdx.x] = row[KVW + kh * HD + threadIdx.x]; + vo[threadIdx.x + 64] = row[KVW + kh * HD + threadIdx.x + 64]; +} + +// The grouped two-tap dynamic causal convolution along each stream's block of 8 rows: +// y[t, c] = sum over o in {0, 1} of (base[o][c] + dyn[t][which][o][c / 16]) * z[t - o][c] (z[-1] = 0) +// dyn [R][2 (pre, post)][2 (tap)][GROUPS]; base [2 (tap)][D] (the module's pre or post kernel). add: x += y, else +// out = y. grid (R, D / 1024), 256 threads. +extern "C" __global__ void df_conv(const __nv_bfloat16* __restrict__ z, const __nv_bfloat16* __restrict__ dyn, int which, + const __nv_bfloat16* __restrict__ base, __nv_bfloat16* __restrict__ out, int add) { + const int r = blockIdx.x, t = r % BLOCK; + const __nv_bfloat16* dr = dyn + (long long)r * 4 * GROUPS + which * 2 * GROUPS; + for (int c = blockIdx.y * 1024 + threadIdx.x; c < (blockIdx.y + 1) * 1024; c += blockDim.x) { + const int g = c / GROUP; + float y = (bf(base[c]) + bf(dr[g])) * bf(z[(long long)r * D + c]); + if (t > 0) y += (bf(base[D + c]) + bf(dr[GROUPS + g])) * bf(z[(long long)(r - 1) * D + c]); + __nv_bfloat16* o = out + (long long)r * D + c; + *o = add ? tobf(bf(*o) + bf(tobf(y))) : tobf(y); + } +} + +// SwiGLU over [gate | up] rows (width F each): bf16(bf16(silu(g)) * u). grid (R, F / 1024), 256 threads. +extern "C" __global__ void df_swiglu(const __nv_bfloat16* __restrict__ gu, __nv_bfloat16* __restrict__ out, int F) { + const int r = blockIdx.x; + const __nv_bfloat16* g = gu + (long long)r * 2 * F; + for (int c = blockIdx.y * 1024 + threadIdx.x; c < (blockIdx.y + 1) * 1024 && c < F; c += blockDim.x) { + const float x = bf(g[c]); + const float s = bf(tobf(x / (1.0f + expf(-x)))); + out[(long long)r * F + c] = tobf(s * bf(g[F + c])); + } +} + +// Attention of each stream's block (8 rows, 4 query heads a key/value head) over its ring's context keys in the +// window, positions max(from, p - 2047) .. start - 1 for a row at p, and the block's 8 keys (every one visible): +// grid (S, KVH), 256 threads (a warp a block row, its 4 heads; a lane 4 of the 128 dims). +// q [S*8][QH*HD], kb / vb [S*8][KVW], o [S*8][QH*HD]; starts / froms the streams' block and context starts. +extern "C" __global__ void __launch_bounds__(256) df_attn(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kb, + const __nv_bfloat16* __restrict__ vb, const uint64_t* __restrict__ rings, + const long long* __restrict__ starts, const long long* __restrict__ froms, + int layer, float scale, __nv_bfloat16* __restrict__ o) { + const int s = blockIdx.x, kh = blockIdx.y; + const int j = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int r = s * BLOCK + j; + const long long start = starts[s], p = start + j; + long long lo = p - (WIN - 1); + if (lo < froms[s]) lo = froms[s]; + if (lo < 0) lo = 0; + float qv[4][4], m[4], l[4], acc[4][4]; + for (int h = 0; h < 4; ++h) { + const __nv_bfloat16* qh = q + (long long)r * QH * HD + (kh * 4 + h) * HD + lane * 4; + for (int i = 0; i < 4; ++i) qv[h][i] = bf(qh[i]) * scale; + m[h] = -INFINITY; + l[h] = 0.f; + for (int i = 0; i < 4; ++i) acc[h][i] = 0.f; + } + const __nv_bfloat16* K = ringK(rings[s], layer); + const __nv_bfloat16* V = ringV(rings[s], layer); + const long long n_ctx = start > lo ? start - lo : 0; + for (long long n = 0; n < n_ctx + BLOCK; ++n) { + const __nv_bfloat16 *kr, *vr; + if (n < n_ctx) { + const int slot = (int)((lo + n) % WIN); + kr = K + (long long)slot * KVW + kh * HD + lane * 4; + vr = V + (long long)slot * KVW + kh * HD + lane * 4; + } else { + const long long rb = (long long)s * BLOCK + (n - n_ctx); + kr = kb + rb * KVW + kh * HD + lane * 4; + vr = vb + rb * KVW + kh * HD + lane * 4; + } + float kf[4], vf[4]; + for (int i = 0; i < 4; ++i) { + kf[i] = bf(kr[i]); + vf[i] = bf(vr[i]); + } + for (int h = 0; h < 4; ++h) { + float d = qv[h][0] * kf[0] + qv[h][1] * kf[1] + qv[h][2] * kf[2] + qv[h][3] * kf[3]; + for (int off = 16; off > 0; off >>= 1) d += __shfl_xor_sync(0xffffffffu, d, off); + const float mn = fmaxf(m[h], d); + const float corr = __expf(m[h] - mn), pe = __expf(d - mn); + l[h] = l[h] * corr + pe; + for (int i = 0; i < 4; ++i) acc[h][i] = acc[h][i] * corr + pe * vf[i]; + m[h] = mn; + } + } + for (int h = 0; h < 4; ++h) { + __nv_bfloat16* oh = o + (long long)r * QH * HD + (kh * 4 + h) * HD + lane * 4; + const float inv = 1.0f / l[h]; + for (int i = 0; i < 4; ++i) oh[i] = tobf(acc[h][i] * inv); + } +} + +// df_attn in chunks: block (s, kh, c) takes the c-th 1/CH of each row's keys (its context, then the block's 8) and +// writes the row's 4 heads' running max, sum and unnormalized output (pm, pl [R][QH][CH]; pacc [R][QH][CH][HD], fp32); +// df_attn_merge combines the chunks. The serial walk over 2,048 keys made df_attn latency-bound; this one runs +// CH blocks a kv head. grid (S, KVH, CH), 256 threads (a warp a block row, as df_attn). +extern "C" __global__ void __launch_bounds__(256) df_attn_part(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kb, + const __nv_bfloat16* __restrict__ vb, const uint64_t* __restrict__ rings, const long long* __restrict__ starts, + const long long* __restrict__ froms, int layer, float scale, float* __restrict__ pm, float* __restrict__ pl, + float* __restrict__ pacc) { + const int s = blockIdx.x, kh = blockIdx.y, c = blockIdx.z; + const int j = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int r = s * BLOCK + j; + const long long start = starts[s], p = start + j; + long long lo = p - (WIN - 1); + if (lo < froms[s]) lo = froms[s]; + if (lo < 0) lo = 0; + const long long n_ctx = start > lo ? start - lo : 0; + const long long T = n_ctx + BLOCK, L = (T + CH - 1) / CH; + const long long n0 = (long long)c * L, n1 = n0 + L < T ? n0 + L : T; + float qv[4][4], m[4], l[4], acc[4][4]; + for (int h = 0; h < 4; ++h) { + const __nv_bfloat16* qh = q + (long long)r * QH * HD + (kh * 4 + h) * HD + lane * 4; + for (int i = 0; i < 4; ++i) qv[h][i] = bf(qh[i]) * scale; + m[h] = -INFINITY; + l[h] = 0.f; + for (int i = 0; i < 4; ++i) acc[h][i] = 0.f; + } + const __nv_bfloat16* K = ringK(rings[s], layer); + const __nv_bfloat16* V = ringV(rings[s], layer); + for (long long n = n0; n < n1; ++n) { + const __nv_bfloat16 *kr, *vr; + if (n < n_ctx) { + const int slot = (int)((lo + n) % WIN); + kr = K + (long long)slot * KVW + kh * HD + lane * 4; + vr = V + (long long)slot * KVW + kh * HD + lane * 4; + } else { + const long long rb = (long long)s * BLOCK + (n - n_ctx); + kr = kb + rb * KVW + kh * HD + lane * 4; + vr = vb + rb * KVW + kh * HD + lane * 4; + } + float kf[4], vf[4]; + for (int i = 0; i < 4; ++i) { + kf[i] = bf(kr[i]); + vf[i] = bf(vr[i]); + } + for (int h = 0; h < 4; ++h) { + float d = qv[h][0] * kf[0] + qv[h][1] * kf[1] + qv[h][2] * kf[2] + qv[h][3] * kf[3]; + for (int off = 16; off > 0; off >>= 1) d += __shfl_xor_sync(0xffffffffu, d, off); + const float mn = fmaxf(m[h], d); + const float corr = __expf(m[h] - mn), pe = __expf(d - mn); + l[h] = l[h] * corr + pe; + for (int i = 0; i < 4; ++i) acc[h][i] = acc[h][i] * corr + pe * vf[i]; + m[h] = mn; + } + } + for (int h = 0; h < 4; ++h) { + const long long at = ((long long)r * QH + kh * 4 + h) * CH + c; + float* pa = pacc + at * HD + lane * 4; + for (int i = 0; i < 4; ++i) pa[i] = acc[h][i]; + if (lane == 0) { + pm[at] = m[h]; // -inf with l = 0 for a chunk past the row's keys: the merge weighs it 0 + pl[at] = l[h]; + } + } +} + +// The chunks' partials of row r, head h into o[r][h] (bf16). grid (R, QH), 128 threads (a thread a dim). +extern "C" __global__ void df_attn_merge(const float* __restrict__ pm, const float* __restrict__ pl, const float* __restrict__ pacc, + __nv_bfloat16* __restrict__ o) { + const int r = blockIdx.x, h = blockIdx.y, d = threadIdx.x; + const long long base = ((long long)r * QH + h) * CH; + float m = -INFINITY; + for (int c = 0; c < CH; ++c) m = fmaxf(m, pm[base + c]); + float l = 0.f, a = 0.f; + for (int c = 0; c < CH; ++c) { + const float w = __expf(pm[base + c] - m); + l += pl[base + c] * w; + a += pacc[(base + c) * HD + d] * w; + } + o[(long long)r * QH * HD + h * HD + d] = tobf(a / l); +} + +// x += y (bf16 rows of D). grid (R, D / 1024), 256 threads. +extern "C" __global__ void df_add(__nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ y) { + const long long r = blockIdx.x; + for (int c = blockIdx.y * 1024 + threadIdx.x; c < (blockIdx.y + 1) * 1024; c += blockDim.x) + x[r * D + c] = tobf(bf(x[r * D + c]) + bf(y[r * D + c])); +} + +// out = bf16(g rank 0 + g rank 1): both ranks' gathered fp32 partials [2][R][D] of a split linear (as the engine's +// glm_residual_add reads them). grid (R, D / 1024), 256 threads. +extern "C" __global__ void df_sum2(const float* __restrict__ g, __nv_bfloat16* __restrict__ out, int R) { + const long long r = blockIdx.x; + for (int c = blockIdx.y * 1024 + threadIdx.x; c < (blockIdx.y + 1) * 1024; c += blockDim.x) { + const long long i = r * D + c; + out[i] = tobf(g[i] + g[(long long)R * D + i]); + } +} + +// A rank's 16 largest logits of each head row (bf16 [R][cols]), as global ids and fp32 values, in two stages: +// df_top16_part: each 4096-column chunk's 16 (grid (R, chunks), 256 threads); df_top16_merge: the chunks' into +// out [R][32] (16 ids, then 16 values' bits; grid R, 256 threads). Ties keep the lower id. +namespace { +constexpr int CHUNK_COLS = 4096; +constexpr int TOP = 16; + +__device__ __forceinline__ bool better(float a, int ia, float b, int ib) { return a > b || (a == b && ia < ib); } + +// The block's best (value, index) of sv/si [n] (256 threads), the winner's slot cleared to -inf. +__device__ void takeBest(float* sv, int* si, int n, float* rv, int* ri, float* outv, int* outi) { + float bv = -INFINITY; + int bi = 0x7fffffff, bs = -1; + for (int j = threadIdx.x; j < n; j += blockDim.x) + if (better(sv[j], si[j], bv, bi)) { bv = sv[j]; bi = si[j]; bs = j; } + for (int o = 16; o > 0; o >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, o); + const int oi = __shfl_xor_sync(0xffffffffu, bi, o), os = __shfl_xor_sync(0xffffffffu, bs, o); + if (better(ov, oi, bv, bi)) { bv = ov; bi = oi; bs = os; } + } + const int w = threadIdx.x >> 5; + if ((threadIdx.x & 31) == 0) { rv[w] = bv; ri[w] = bi; ri[8 + w] = bs; } + __syncthreads(); + if (threadIdx.x == 0) { + int best = 0; + for (int k = 1; k < (int)(blockDim.x >> 5); ++k) + if (better(rv[k], ri[k], rv[best], ri[best])) best = k; + *outv = rv[best]; + *outi = ri[best]; + if (ri[8 + best] >= 0) sv[ri[8 + best]] = -INFINITY; + } + __syncthreads(); +} +} // namespace + +extern "C" __global__ void __launch_bounds__(256) df_top16_part(const __nv_bfloat16* __restrict__ logits, int cols, int offset, + float* __restrict__ pv, int* __restrict__ pi) { + __shared__ float sv[CHUNK_COLS]; + __shared__ int si[CHUNK_COLS]; + __shared__ float rv[8]; + __shared__ int ri[16]; + const int r = blockIdx.x, c0 = blockIdx.y * CHUNK_COLS; + const int n = min(CHUNK_COLS, cols - c0); + for (int j = threadIdx.x; j < n; j += blockDim.x) { + sv[j] = bf(logits[(long long)r * cols + c0 + j]); + si[j] = offset + c0 + j; + } + __syncthreads(); + const long long base = ((long long)r * gridDim.y + blockIdx.y) * TOP; + for (int t = 0; t < TOP; ++t) takeBest(sv, si, n, rv, ri, pv + base + t, pi + base + t); +} + +extern "C" __global__ void __launch_bounds__(256) df_top16_merge(const float* __restrict__ pv, const int* __restrict__ pi, int chunks, + int* __restrict__ out) { + __shared__ float sv[64 * TOP]; + __shared__ int si[64 * TOP]; + __shared__ float rv[8]; + __shared__ int ri[16]; + __shared__ float ov[TOP]; + __shared__ int oi[TOP]; + const int r = blockIdx.x, n = chunks * TOP; + for (int j = threadIdx.x; j < n; j += blockDim.x) { + sv[j] = pv[(long long)r * n + j]; + si[j] = pi[(long long)r * n + j]; + } + __syncthreads(); + for (int t = 0; t < TOP; ++t) takeBest(sv, si, n, rv, ri, ov + t, oi + t); + if (threadIdx.x < TOP) { + out[r * 2 * TOP + threadIdx.x] = oi[threadIdx.x]; + out[r * 2 * TOP + TOP + threadIdx.x] = __float_as_int(ov[threadIdx.x]); + } +} diff --git a/zig/kernels/cuda/glm/exl3.cu b/zig/kernels/cuda/glm/exl3.cu new file mode 100644 index 0000000000..64a06dc7d7 --- /dev/null +++ b/zig/kernels/cuda/glm/exl3.cu @@ -0,0 +1,677 @@ +// GLM-5.3-Flash's EXL3 routed experts (4-bit "mcg" trellis, ExLlamaV3's format: MIT, Copyright (c) 2025 Turboderp) +// on CUDA: the grouped trellis GEMV/GEMM and the Hadamard rotations around it. From TensorFold's python-0.6 +// src/tensorfold/families/glm5_next/cuda/exl3.cu, without torch, over the routed pairs alone (the shared expert runs +// on the dense linears), with items of 64 pairs for prompt chunks (each expert's tiles read once a chunk). +// +// A 16x16 tile is 32 little-endian 32-bit words; lane L of a warp decodes the tile's values 8L..8L+7 from words L-1 +// and L, and those eight values are exactly the B fragments of two mma.m16n8k16 (columns 0-7 and 8-15 of the tile). +// Every output depends only on its own pair: pairs share an mma tile but mma keeps rows independent, the K range of +// every warp and split is fixed by the shape, and warps and splits are summed in a fixed order, never with atomics. + +#include +#include +#include + +namespace { + +constexpr float HAD_SCALE = 0.08838834764831845f; // 1 / sqrt(128) + +// Two values of the "mcg" codebook from two 16-bit states, as a half2 (first state in .x). +__device__ __forceinline__ uint32_t mcg2(uint32_t s0, uint32_t s1) { + uint32_t x0 = s0 * 0xCBAC1FEDu; + uint32_t x1 = s1 * 0xCBAC1FEDu; + x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u; + x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u; + uint32_t lo = __byte_perm(x0, x1, 0x5410); + uint32_t hi = __byte_perm(x0, x1, 0x7632); + half2 r = __hadd2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + return *reinterpret_cast(&r); +} + +// This lane's eight values of a 4-bit tile (word = tile[lane]) as the B fragments of its two n8 halves. +__device__ __forceinline__ void decode_tile(uint32_t w, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { + uint32_t p = __shfl_sync(0xffffffffu, w, (lane + 31) & 31); + uint32_t s = __funnelshift_r(w, p, 20); + b0[0] = mcg2((s >> 8) & 0xffffu, (s >> 4) & 0xffffu); + b0[1] = mcg2(s & 0xffffu, w >> 16); + b1[0] = mcg2((w >> 12) & 0xffffu, (w >> 8) & 0xffffu); + b1[1] = mcg2((w >> 4) & 0xffffu, w & 0xffffu); +} + +__device__ __forceinline__ void mma16816(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +__device__ __forceinline__ uint32_t load_pair(const half* x, bool ok) { + return ok ? *reinterpret_cast(x) : 0u; +} + +// Grid (item, n block, mat * SK + split) -> Z[mat][split][pair][n]; an item is an expert and up to 16 MT of its pairs +// (the plan's members); warps sum fixed k ranges, added in warp order. +template +__device__ __forceinline__ void grouped( + const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int SK) { + const int item = blockIdx.x; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + const int split = blockIdx.z % SK; + const int mat = blockIdx.z / SK; + const half* X = mat ? X1 : X0; + const uint32_t* T = mat ? T1 : T0; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int KT = K >> 4, NTILES = N >> 4; + constexpr int ROWS = 16 * MT; + + __shared__ int rows_sh[ROWS]; + for (int i = threadIdx.x; i < ROWS; i += W * 32) rows_sh[i] = i < cnt ? members[first + i] : -1; + __syncthreads(); + int r0[MT], r1[MT]; + const half* x0[MT]; + const half* x1[MT]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + r0[m] = rows_sh[m * 16 + g]; + r1[m] = rows_sh[m * 16 + g + 8]; + x0[m] = X + (size_t)(r0[m] < 0 ? 0 : r0[m]) * K + 2 * t; + x1[m] = X + (size_t)(r1[m] < 0 ? 0 : r1[m]) * K + 2 * t; + } + + const int per_split = KT / SK, per_warp = per_split / W; + const int kt0 = split * per_split + warp * per_warp; + const int nt0 = blockIdx.y * NT; + const uint32_t* tile = T + (((size_t)e * KT + kt0) * NTILES + nt0) * 32 + lane; + + float acc[MT][NT][2][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[m][i][h][c] = 0.f; + + for (int kt = kt0; kt < kt0 + per_warp; ++kt) { + uint32_t words[NT]; +#pragma unroll + for (int i = 0; i < NT; ++i) words[i] = __ldg(tile + i * 32); + const int k = kt * 16; + uint32_t a[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + a[m][0] = load_pair(x0[m] + k, r0[m] >= 0); + a[m][1] = load_pair(x1[m] + k, r1[m] >= 0); + a[m][2] = load_pair(x0[m] + k + 8, r0[m] >= 0); + a[m][3] = load_pair(x1[m] + k + 8, r1[m] >= 0); + } +#pragma unroll + for (int i = 0; i < NT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile(words[i], lane, b0, b1); +#pragma unroll + for (int m = 0; m < MT; ++m) { + mma16816(acc[m][i][0], a[m], b0); + mma16816(acc[m][i][1], a[m], b1); + } + } + tile += (size_t)NTILES * 32; + } + + // warps' partial sums through shared memory, added in warp order + __shared__ float red[W][ROWS][NT * 16]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + red[warp][m * 16 + g][col] = acc[m][i][h][0]; + red[warp][m * 16 + g][col + 1] = acc[m][i][h][1]; + red[warp][m * 16 + g + 8][col] = acc[m][i][h][2]; + red[warp][m * 16 + g + 8][col + 1] = acc[m][i][h][3]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < ROWS * NT * 16; idx += W * 32) { + const int row = idx / (NT * 16), col = idx % (NT * 16); + const int r = rows_sh[row]; + if (r < 0) continue; + float s = red[0][row][col]; +#pragma unroll + for (int w = 1; w < W; ++w) s += red[w][row][col]; + Z[(((size_t)mat * SK + split) * P + r) * N + nt0 * 16 + col] = s; + } +} + +// Fast Walsh-Hadamard transform of 128 values held 4 per lane (lane L: values 4L..4L+3), natural order, fixed +// butterfly order: strides 1, 2 in registers, 4..64 across lanes. +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3]; + v[0] = a + c; v[1] = b + d; v[2] = a - c; v[3] = b - d; +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + + +// -- prompt chunks: the staged kernel ------------------------------------------------------------------------------- +// An item is an expert and up to PM of its pairs (the plan's members, tile PM). Block (item, column block of +// NW * NTW n-tiles, matrix): the four warps split the columns, not K, so a warp's sums need no cross-warp reduction +// and every decoded tile feeds PM / 8 mma; the member rows' k-slices and the tiles' words are staged in shared +// memory by cp.async, STAGES deep, and the A fragments come from the staged rows by ldmatrix. The same sums as the +// grouped kernel for every output, in a fixed order (one warp owns each output; K runs in order). +// +// decode_tile's values with fewer integer instructions: the byte-aligned 16-bit states taken with one PRMT each +// and w >> 4 shared by two states (MiaAI-Lab, GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold patch 0020, Apache-2.0). +// The same states, so the same bits. +__device__ __forceinline__ void decode_tile_p(uint32_t w, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { + const uint32_t p = __shfl_sync(0xffffffffu, w, (lane + 31) & 31); + const uint32_t s = __funnelshift_r(w, p, 20); + const uint32_t t = w >> 4; + b0[0] = mcg2(__byte_perm(s, 0u, 0x4421), (s >> 4) & 0xffffu); // (s >> 8) & 0xffff, (s >> 4) & 0xffff + b0[1] = mcg2(s & 0xffffu, w >> 16); + b1[0] = mcg2(__byte_perm(t, 0u, 0x4421), __byte_perm(w, 0u, 0x4421)); // (w >> 12) & 0xffff, (w >> 8) & 0xffff + b1[1] = mcg2(__byte_perm(t, 0u, 0x4410), w & 0xffffu); // (w >> 4) & 0xffff, w & 0xffff +} + +__device__ __forceinline__ void cp_async16(void* dst, const void* src, bool valid) { + const unsigned d = (unsigned)__cvta_generic_to_shared(dst); + const int n = valid ? 16 : 0; // src-size 0: the 16 bytes come as zeros (a padding row) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(d), "l"(src), "r"(n)); +} +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ void ldmatrix_x4(uint32_t (&a)[4], const half* p) { + const unsigned s = (unsigned)__cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) : "r"(s)); +} + +// xshift: the X row a member pair reads is pair >> xshift (3: one rotated row a token, when every expert of the layer +// shares one input rotation for gate and up; 0: a row a pair). +// SK: K splits -> Z[mat][split][pair][n] (the grouped kernel's layout: glm_exl3_gateup adds the splits in order). A +// small chunk, whose items hold few pairs, gets SK times the blocks at a fraction of the K each. +// Grid: x = (column block, matrix of ``mats``, split), y = item: one item's blocks are adjacent in launch order, so +// its staged rows come from L2 after the first block reads them. +template +__device__ __forceinline__ void prompt_grouped( + const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int xshift, int SK, int mats) { + constexpr int T = NW * 32, MT = PM / 16, NTB = NW * NTW, XS = 16 * KS + 8; // XS: a staged row's halves (padded: ldmatrix conflict-free) + constexpr int XCH = 2 * KS; // 16-byte chunks a staged row takes + static_assert(PM * XCH % T == 0 && NTB * KS * 8 % T == 0, "the stage's chunks spread evenly over the threads"); + __shared__ __align__(16) half xs[STAGES][PM][XS]; + __shared__ __align__(16) uint32_t ws[STAGES][NTB][KS][32]; + __shared__ int rows_sh[PM]; + const int item = blockIdx.y; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + const int cb = blockIdx.x / (mats * SK), ms = blockIdx.x % (mats * SK); + const int mat = ms / SK, split = ms % SK; + const half* X = mat ? X1 : X0; + const uint32_t* Tm = mat ? T1 : T0; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int KT = K >> 4, NTILES = N >> 4; + const int per_split = KT / SK, kt_base = split * per_split; // KS divides per_split (K = 4096 or 1024, SK <= 4) + const int nt0 = cb * NTB; + for (int i = threadIdx.x; i < PM; i += T) rows_sh[i] = i < cnt ? members[first + i] : -1; + __syncthreads(); + const int mt_n = (cnt + 15) >> 4; // m-tiles with a member (uniform over the block) + const uint32_t* tbase = Tm + ((size_t)e * KT * NTILES + nt0) * 32; + + auto issue = [&](int stage, int kt) { // the stage's rows (k-tiles kt .. kt + KS - 1) and words +#pragma unroll + for (int j = 0; j < PM * XCH / T; ++j) { + const int c = threadIdx.x + j * T; + const int row = c / XCH, part = c % XCH; + if (row >= mt_n * 16) continue; // an m-tile without a member: never read + const int r = rows_sh[row]; + const half* src = X + (size_t)(r < 0 ? 0 : (r >> xshift)) * K + kt * 16 + part * 8; + cp_async16(&xs[stage][row][part * 8], src, r >= 0); + } +#pragma unroll + for (int j = 0; j < NTB * KS * 8 / T; ++j) { + const int c = threadIdx.x + j * T; + const int tile = c / (KS * 8), kk = (c / 8) % KS, part = c % 8; + const uint32_t* src = tbase + ((size_t)(kt + kk) * NTILES + tile) * 32 + part * 4; + cp_async16(&ws[stage][tile][kk][part * 4], src, true); + } + cp_async_commit(); + }; + + float acc[MT][NTW][2][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < NTW; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[m][i][h][c] = 0.f; + + const int nstages = per_split / KS; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < nstages) issue(s, kt_base + s * KS); + else cp_async_commit(); + } + for (int st = 0; st < nstages; ++st) { + cp_async_wait(); + __syncthreads(); + const int stage = st % STAGES; + { // the stage after the ones in flight: its buffer was consumed one iteration ago + const int nxt = st + STAGES - 1; + if (nxt < nstages) issue(nxt % STAGES, kt_base + nxt * KS); + else cp_async_commit(); + } +#pragma unroll + for (int kk = 0; kk < KS; ++kk) { + uint32_t b[NTW][2][2]; +#pragma unroll + for (int i = 0; i < NTW; ++i) decode_tile_p(ws[stage][warp * NTW + i][kk][lane], lane, b[i][0], b[i][1]); + for (int m = 0; m < mt_n; ++m) { + uint32_t a[4]; + ldmatrix_x4(a, &xs[stage][m * 16 + (lane & 15)][kk * 16 + (lane >> 4) * 8]); +#pragma unroll + for (int i = 0; i < NTW; ++i) { + mma16816(acc[m][i][0], a, b[i][0]); + mma16816(acc[m][i][1], a, b[i][1]); + } + } + } + } + cp_async_wait<0>(); + // this warp's columns, rows g and g + 8 of each m-tile, straight to Z[mat][split][pair][n] + float* Zs = Z + (size_t)(mat * SK + split) * P * N; + for (int m = 0; m < mt_n; ++m) { + const int r0 = rows_sh[m * 16 + g], r1 = rows_sh[m * 16 + g + 8]; +#pragma unroll + for (int i = 0; i < NTW; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = (nt0 + warp * NTW + i) * 16 + h * 8 + 2 * t; + if (r0 >= 0) *reinterpret_cast(Zs + (size_t)r0 * N + col) = make_float2(acc[m][i][h][0], acc[m][i][h][1]); + if (r1 >= 0) *reinterpret_cast(Zs + (size_t)r1 * N + col) = make_float2(acc[m][i][h][2], acc[m][i][h][3]); + } + } +} + +} // namespace + +// prompt chunks: items of 128 pairs, 4 warps x 2 n-tiles (128 columns a block), 2 k-tiles a stage, 3 stages +extern "C" __global__ void __launch_bounds__(128) glm_exl3_prompt128(const half* __restrict__ X0, const half* __restrict__ X1, + const uint32_t* __restrict__ T0, const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int xshift, int SK, int mats) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + prompt_grouped<128, 4, 2, 2, 3>(X0, X1, T0, T1, items, counts, members, Z, K, N, P, xshift, SK, mats); +} + +// the same with 8 warps (256 columns a block): the item's rows staged half as often +extern "C" __global__ void __launch_bounds__(256) glm_exl3_prompt128w(const half* __restrict__ X0, const half* __restrict__ X1, + const uint32_t* __restrict__ T0, const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int xshift, int SK, int mats) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + prompt_grouped<128, 8, 2, 2, 3>(X0, X1, T0, T1, items, counts, members, Z, K, N, P, xshift, SK, mats); +} + +// Program (row, 128-block of K): Xh[row][block] = fp16((x[row] * suh) @ H), one rotation a token for the gate and up +// projections of a layer whose experts all share one suh (then every pair of the row reads Xh[row]: the same bits +// glm_exl3_rot_in gives each pair). One warp a program. +extern "C" __global__ void glm_exl3_rot_rows(const __nv_bfloat16* __restrict__ x, int x_stride, const half* __restrict__ suh, + half* __restrict__ out, int K) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + const int row = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const half* sh = suh + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(sh[j]); + fwht128(v, lane); + half* o = out + (size_t)row * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +namespace { +} // namespace + +// decode windows: items of 16 pairs, 8 n-tiles a block, 4 warps +extern "C" __global__ void __launch_bounds__(128) glm_exl3_grouped(const half* __restrict__ X0, const half* __restrict__ X1, + const uint32_t* __restrict__ T0, const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int SK) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + grouped<1, 8, 4>(X0, X1, T0, T1, items, counts, members, Z, K, N, P, SK); +} + +// prompt chunks: items of 64 pairs, 2 n-tiles a block, 4 warps +extern "C" __global__ void __launch_bounds__(128) glm_exl3_grouped64(const half* __restrict__ X0, const half* __restrict__ X1, + const uint32_t* __restrict__ T0, const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int SK) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + grouped<4, 2, 4>(X0, X1, T0, T1, items, counts, members, Z, K, N, P, SK); +} + +// Program (pair, 128-block of K, matrix): Xh[mat][pair][block] = fp16((x[row] * suh[mat][e]) @ H) for the gate and up +// projections of every routed pair (pair = row * 8 + slot). One warp a program. +extern "C" __global__ void glm_exl3_rot_in(const __nv_bfloat16* __restrict__ x, int x_stride, const int* __restrict__ pick, + const half* __restrict__ suh0, const half* __restrict__ suh1, half* __restrict__ out0, half* __restrict__ out1, int K) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int row = p >> 3; + const int lane = threadIdx.x; + const int e = pick[p]; + const half* suh = (mat ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? out1 : out0) + (size_t)p * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (pair, 128-block of the rank's expert width): gate and up summed over the splits in order, rotated, scaled +// by svh; GLM's limited SwiGLU with the grouped kernels' bf16 roundings; then the down projection's input rotation: +// Xd[pair][block] = fp16((act * suh_d[e]) @ H). +extern "C" __global__ void glm_exl3_gateup(const float* __restrict__ Z, const int* __restrict__ pick, const half* __restrict__ svh_g, + const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, int P, int N, int SK, float limit) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const int e = pick[p]; + const int n = blk * 128 + 4 * lane; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float sg = 0.f, su = 0.f; + for (int s = 0; s < SK; ++s) { + sg += Z[((size_t)(0 * SK + s) * P + p) * N + n + j]; + su += Z[((size_t)(1 * SK + s) * P + p) * N + n + j]; + } + gv[j] = sg; + uv[j] = su; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float gg = fminf(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = fminf(fmaxf(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), -limit), limit); + float act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (pair, 128-block of the model width): the down projection summed over the splits in order, rotated and +// scaled by svh -> y[row][slot][:] (fp32, rows of ``ys`` slots: the shared expert's last). This rank's share. +extern "C" __global__ void glm_exl3_down(const float* __restrict__ Z, const int* __restrict__ pick, const half* __restrict__ svh_d, + float* __restrict__ y, int P, int D, int SK, int ys) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const int e = pick[p]; + const int n = blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int k = 0; k < SK; ++k) s += Z[((size_t)k * P + p) * D + n + j]; + v[j] = s; + } + fwht128(v, lane); + float* o = y + ((size_t)(p >> 3) * ys + (p & 7)) * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); +} + +// -- prompt chunks, the transposed design: Y^T = W^T X^T ------------------------------------------------------------- +// After MiaAI-Lab's prompt kernels (GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold, patches 0004 and 0009, Apache-2.0), +// written for this engine. A decoded trellis tile is, reordered, an m16n8k16 A fragment of W^T (lane (g, t): +// a0 = b0[0], a1 = b1[0], a2 = b0[1], a3 = b1[1]), so each warp decodes its own tiles straight into registers and +// every decoded fragment feeds the whole item's member rows (PM = 64 pairs) as B fragments from shared memory. A +// block covers 16 weight tiles: gate and up's 8 each of one 128-column Hadamard block (the member rows staged once +// for both), or the down projection's 256 columns; 8 warps hold 2 tiles each. Each output is one fp32 chain of mma +// over ascending k tiles (no split K, no partials in memory), so a row's bits depend on its own values only. The +// epilogues (glm_exl3_gateup's and glm_exl3_down's formulas) run on the sums in shared memory: no round trip. +namespace { + +constexpr int TPM = 64, TNW = 8, TMT = 16 / TNW, TNB = TPM / 8, TT = TNW * 32, TPH = TPM / 2; +constexpr int TBK = 32, TKT = TBK / 16, TSTAGES = 3, TXS = TBK + 8; // a stage: 32 k values a row; +16 B: ldmatrix rows on distinct banks +constexpr int TWORDS = 16 * TKT * 32; // a stage's trellis words (16 tiles x 2 k tiles x 32) +constexpr int GU_ES = 128 + 4, DN_ES = 256 + 4; // the epilogue sums' row strides (floats) +template +constexpr int tstage_halves() { return XR * TXS + 2 * TWORDS; } +constexpr int tsmem_bytes() { + const int gu = TSTAGES * tstage_halves<2 * TPM>() * 2; // gate and up rows apart (no shared rotation) + const int ep = 2 * TPH * GU_ES * 4, dn = TPH * DN_ES * 4; + return gu > ep ? (gu > dn ? gu : dn) : (ep > dn ? ep : dn); +} + +// Which (matrix, weight tile) warp-tile q (0..15) of column block cb is. +template +__device__ __forceinline__ void ttile(int q, int cb, int& mat, int& nt) { + mat = MATS == 2 ? q >> 3 : 0; + nt = MATS == 2 ? cb * 8 + (q & 7) : cb * 16 + q; +} + +// The main loop: acc[i][nb] = this warp's tile i (16 output columns) x members nb * 8 .. + 7, over K. XR: the rows a +// stage holds (TPM when both matrices read the same rows: X0 == X1, xshift 3; 2 * TPM otherwise). +template +__device__ __forceinline__ void tmain(const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, const int* rows_sh, int cnt, int e, int K, int NTILES, int cb, int xshift, + half* stage, float (&acc)[TMT][TNB][4]) { + constexpr int SH = tstage_halves(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + int mat, nt_unused; + ttile(warp * TMT, cb, mat, nt_unused); + const int KT = K >> 4, chunks = K / TBK; + const int nblocks = (cnt + 15) >> 4; // 16-member blocks in use +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int j = 0; j < TNB; ++j) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][j][c] = 0.f; + auto load = [&](int chunk) { + half* dst = stage + (size_t)(chunk % TSTAGES) * SH; + const int k0 = chunk * TBK; + for (int idx = threadIdx.x; idx < XR * (TBK / 8); idx += TT) { + const int m = idx / (TPM * (TBK / 8)), rem = idx % (TPM * (TBK / 8)), i = rem / (TBK / 8), j = rem % (TBK / 8); + if (i >= nblocks * 16) continue; // a 16-member block without a member: never read + const int r = rows_sh[i]; + const half* src = (m ? X1 : X0) + (size_t)(r < 0 ? 0 : (r >> xshift)) * K + k0 + j * 8; + cp_async16(dst + ((size_t)m * TPM + i) * TXS + j * 8, src, r >= 0); + } + uint32_t* wdst = reinterpret_cast(dst + (size_t)XR * TXS); + for (int idx = threadIdx.x; idx < 16 * TKT * 8; idx += TT) { // 16-byte pieces of the 128-byte tile words + const int q = idx / (TKT * 8), kk = (idx / 8) % TKT, part = idx % 8; + int wm, nt; + ttile(q, cb, wm, nt); + const uint32_t* src = (wm ? T1 : T0) + (((size_t)e * KT + chunk * TKT + kk) * NTILES + nt) * 32 + part * 4; + cp_async16(wdst + ((size_t)q * TKT + kk) * 32 + part * 4, src, true); + } + }; +#pragma unroll + for (int c = 0; c < TSTAGES - 1; ++c) { + if (c < chunks) load(c); + cp_async_commit(); + } + for (int c = 0; c < chunks; ++c) { + cp_async_wait(); + __syncthreads(); + if (c + TSTAGES - 1 < chunks) load(c + TSTAGES - 1); + cp_async_commit(); + const half* base = stage + (size_t)(c % TSTAGES) * SH; + const half* xs = base + (size_t)(XR == TPM ? 0 : mat) * TPM * TXS; + const uint32_t* ws = reinterpret_cast(base + (size_t)XR * TXS); +#pragma unroll + for (int kk = 0; kk < TKT; ++kk) { + uint32_t a[TMT][4]; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile_p(ws[((size_t)(warp * TMT + i) * TKT + kk) * 32 + lane], lane, b0, b1); + a[i][0] = b0[0]; a[i][1] = b1[0]; a[i][2] = b0[1]; a[i][3] = b1[1]; + } +#pragma unroll + for (int jb = 0; jb < TPM / 16; ++jb) { + if (jb < nblocks) { + uint32_t r[4]; + const int row = jb * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + ldmatrix_x4(r, xs + (size_t)row * TXS + kk * 16 + ((lane >> 3) & 1) * 8); + const uint32_t bl[2] = {r[0], r[1]}, bh[2] = {r[2], r[3]}; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + mma16816(acc[i][2 * jb], a[i], bl); + mma16816(acc[i][2 * jb + 1], a[i], bh); + } + } + } + } + } + cp_async_wait<0>(); + __syncthreads(); // the stages become the epilogue's sums +} + +// acc (this warp's tiles, members TPH * hf .. + TPH - 1) -> E[member - TPH * hf][col] fp32, row stride ES; col0: the +// first tile's first column within the block's E columns. +__device__ __forceinline__ void tspill(float* E, int ES, int col0, const float (&acc)[TMT][TNB][4], int hf) { + const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int q = 0; q < TPM / 16; ++q) { + const int nb = hf * (TPM / 16) + q; + const int tok = q * 8 + 2 * t, col = col0 + i * 16 + g; + E[(size_t)tok * ES + col] = acc[i][nb][0]; + E[(size_t)(tok + 1) * ES + col] = acc[i][nb][1]; + E[(size_t)tok * ES + col + 8] = acc[i][nb][2]; + E[(size_t)(tok + 1) * ES + col + 8] = acc[i][nb][3]; + } +} + +template +__device__ __forceinline__ void tgateup(const half* __restrict__ Xg, const half* __restrict__ Xu, const uint32_t* __restrict__ Tg, + const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, + const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, + int K, int N, int xshift, float limit, unsigned char* smem_raw, int* rows_sh) { + const int cb = blockIdx.x, item = blockIdx.y; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + tmain<2, XR>(Xg, Xu, Tg, Tu, rows_sh, cnt, e, K, N >> 4, cb, xshift, reinterpret_cast(smem_raw), acc); + float* Eg = reinterpret_cast(smem_raw); + int mat, nt; + ttile<2>(warp * TMT, cb, mat, nt); + const int col0 = (nt - cb * 8) * 16; + const int n = cb * 128 + 4 * lane; + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); // the first half's rows are read + tspill(Eg + (size_t)mat * TPH * GU_ES, GU_ES, col0, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + const int p = rows_sh[hf * TPH + i]; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + gv[j] = Eg[(size_t)i * GU_ES + 4 * lane + j]; + uv[j] = Eg[(size_t)(TPH + i) * GU_ES + 4 * lane + j]; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float gg = fminf(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = fminf(fmaxf(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), -limit), limit); + float act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); + } + } +} + +} // namespace + +// Prompt chunks (items of 64 pairs): block (128-column block of the expert width, item), 8 warps: Xd[pair] = the +// gate/up epilogue of the item's pairs. xshift 3: one rotated row a token read by gate and up alike (the rows staged +// once); 0: a row a pair in Xg and Xu. +extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_gu(const half* __restrict__ Xg, const half* __restrict__ Xu, + const uint32_t* __restrict__ Tg, const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, + half* __restrict__ xd, int K, int N, int xshift, float limit) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + __shared__ __align__(16) unsigned char smem_raw[tsmem_bytes()]; // one buffer for either instantiation + __shared__ int rows_sh[TPM]; + if (xshift) tgateup(Xg, Xu, Tg, Tu, items, counts, members, svh_g, svh_u, suh_d, xd, K, N, xshift, limit, smem_raw, rows_sh); + else tgateup<2 * TPM>(Xg, Xu, Tg, Tu, items, counts, members, svh_g, svh_u, suh_d, xd, K, N, 0, limit, smem_raw, rows_sh); +} + +// Block (256-column block of the model width, item): y[row][slot][:] = the down epilogue of the item's pairs (this +// rank's share). +extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_dn(const half* __restrict__ Xd, const uint32_t* __restrict__ Td, + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, const half* __restrict__ svh_d, + float* __restrict__ y, int K, int D, int ys) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + __shared__ __align__(16) unsigned char smem_raw[tsmem_bytes()]; + __shared__ int rows_sh[TPM]; + const int cb = blockIdx.x, item = blockIdx.y; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + tmain<1, TPM>(Xd, Xd, Td, Td, rows_sh, cnt, e, K, D >> 4, cb, 0, reinterpret_cast(smem_raw), acc); + float* Ed = reinterpret_cast(smem_raw); + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); + tspill(Ed, DN_ES, warp * TMT * 16, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + const int p = rows_sh[hf * TPH + i]; +#pragma unroll + for (int blk = 0; blk < 2; ++blk) { + const int n = cb * 256 + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = Ed[(size_t)i * DN_ES + blk * 128 + 4 * lane + j]; + fwht128(v, lane); + float* o = y + ((size_t)(p >> 3) * ys + (p & 7)) * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); + } + } + } +} diff --git a/zig/kernels/cuda/glm/experts_sym.cu b/zig/kernels/cuda/glm/experts_sym.cu new file mode 100644 index 0000000000..80f01eeeaa --- /dev/null +++ b/zig/kernels/cuda/glm/experts_sym.cu @@ -0,0 +1,302 @@ +// GLM-5.3-Flash's symmetric routed experts, decode form: flashnext/experts.cu reading experts_sym.cuh's blocks, each +// group's bias b = -8 s (exact), so acc = fma(xs, b, fma(p, s, acc)) per group in order gives the affine kernel's bits. + +#include +#include + +#include "experts_sym.cuh" + +namespace { + +__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)); } + +// a row's group sum: the lane's inputs in order, then the quad (every lane of the quad gets the same bits) +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 +struct Stage { + uint4 w[M][Geo::WV]; + uint4 sb[M][Geo::SBV]; + uint4 xa[Geo::XV]; // row gq of the tile: the lane's inputs of the group + uint4 xb[Geo::XV]; // row gq + 8 +}; + +template +__device__ __forceinline__ void load_stage(Stage& st, const uint4* blk, int g, int lane, int t, + const __nv_bfloat16* x0, const __nv_bfloat16* x1, bool v0, bool v1) { + using G = Geo; + const uint4* b = blk + (size_t)g * (M * G::BLOCK); +#pragma unroll + for (int m = 0; m < M; ++m) { +#pragma unroll +#ifdef GLM_SYM3 + for (int c = 0; c < G::WV; ++c) { + const uint32_t* b32 = reinterpret_cast(b + m * G::BLOCK) + c * 96 + lane; + uint32_t p0, p1, p2; + asm volatile("ld.global.nc.L1::no_allocate.u32 %0, [%1];\n" : "=r"(p0) : "l"(b32)); + asm volatile("ld.global.nc.L1::no_allocate.u32 %0, [%1];\n" : "=r"(p1) : "l"(b32 + 32)); + asm volatile("ld.global.nc.L1::no_allocate.u32 %0, [%1];\n" : "=r"(p2) : "l"(b32 + 64)); + st.w[m][c] = unpack3(p0, p1, p2); + } +#else + for (int c = 0; c < G::WV; ++c) st.w[m][c] = ld_w(b + m * G::BLOCK + c * 32 + lane); +#endif +#pragma unroll + for (int c = 0; c < G::SBV; ++c) st.sb[m][c] = ld_w(b + m * G::BLOCK + G::WSEC + t * G::SBV + c); + } + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + const int k0 = g * GS; +#pragma unroll + for (int c = 0; c < G::XV; ++c) { + st.xa[c] = v0 ? ld_x(x0 + k0 + 8 * c) : zero; + st.xb[c] = v1 ? ld_x(x1 + k0 + 8 * c) : zero; + } +} + +template +__device__ __forceinline__ void compute_stage(float (&acc)[M][1][NTW][4], const Stage& st, bool hi) { + using G = Geo; + const float xsa = group_sum(st.xa); + const float xsb = group_sum(st.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) { + // k-step ks: the lane's inputs 4 ks .. 4 ks + 3 (pairs at the mma's k positions 2t and 2t + 8) + const uint32_t a0 = comp(st.xa[ks >> 1], 2 * (ks & 1)); + const uint32_t a2 = comp(st.xa[ks >> 1], 2 * (ks & 1) + 1); + const uint32_t a1 = comp(st.xb[ks >> 1], 2 * (ks & 1)); + const uint32_t a3 = comp(st.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(st.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(st.sb[m][j >> 2], j & 3); + const float s0 = lo_f(sp), s1 = hi_f(sp), b0 = -8.f * s0, b1 = -8.f * s1; + float(&a)[4] = acc[m][0][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 (hi) { + a[2] = fmaf(xsb, b0, fmaf(p[j][2], s0, a[2])); + a[3] = fmaf(xsb, b1, fmaf(p[j][3], s1, a[3])); + } + } + } +} + +template +__device__ __forceinline__ void k_loop(float (&acc)[M][1][NTW][4], const uint4* blk, int KG, int lane, int t, + const __nv_bfloat16* x0, const __nv_bfloat16* x1, bool v0, bool v1) { + Stage st[D]; +#pragma unroll + for (int d = 0; d < D; ++d) + if (d < KG) load_stage(st[d], blk, d, lane, t, x0, x1, v0, v1); + for (int g0 = 0; g0 < KG; g0 += D) { +#pragma unroll + for (int d = 0; d < D; ++d) { + const int g = g0 + d; + if (g < KG) { + compute_stage(acc, st[d], v1); + if (g + D < KG) load_stage(st[d], blk, g + D, lane, t, x0, x1, v0, v1); + } + } + } +} + +// Pair p reads X row p / slots when slots > 0 (X holds tokens), else X row p (X holds a row a pair). +template +__global__ void __launch_bounds__(WARPS * 32) + expert_kernel(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) { + 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); + const bool v0 = gq < cnt, v1 = gq + 8 < cnt; + const int pr0 = v0 ? __ldg(members + first + gq) : 0; + const int pr1 = v1 ? __ldg(members + first + gq + 8) : 0; + const int r0 = slots ? pr0 / slots : pr0, r1 = slots ? pr1 / slots : pr1; + const __nv_bfloat16* x0 = X + (size_t)r0 * x_stride + t * (GS / 4); + const __nv_bfloat16* x1 = X + (size_t)r1 * x_stride + t * (GS / 4); + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (M * Geo::BLOCK); + float acc[M][1][NTW][4]; +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[m][0][j][0] = acc[m][0][j][1] = acc[m][0][j][2] = acc[m][0][j][3] = 0.f; + k_loop(acc, blk, KG, lane, t, x0, x1, v0, v1); + epilogue(acc, 0, out, N, cb * COLS + 2 * t, pr0, pr1, v0, v1, limit); + } +} + +constexpr int PLAN_THREADS = 1024; +constexpr int EMAX = 1024; +constexpr int PSMALL = 1024; // pairs the one-block plan takes; wider calls rank in blocks of PLAN_THREADS + +// Whole block: thread e gives expert e's pair count, gets its first member; items of <= T pairs land in expert order. +__device__ int place_items(int c, int E, int T, int* __restrict__ items, int* __restrict__ counts) { + __shared__ int wsum[3][32]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int tiles = (c + T - 1) / T; + int a = c, b = tiles, u = c > 0; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int ua = __shfl_up_sync(0xffffffffu, a, o), ub = __shfl_up_sync(0xffffffffu, b, o); + const int uu = __shfl_up_sync(0xffffffffu, u, o); + if (lane >= o) { + a += ua; + b += ub; + u += uu; + } + } + if (lane == 31) { + wsum[0][warp] = a; + wsum[1][warp] = b; + wsum[2][warp] = u; + } + __syncthreads(); + if (warp == 0) { +#pragma unroll + for (int s = 0; s < 3; ++s) { + int v = wsum[s][lane]; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int w = __shfl_up_sync(0xffffffffu, v, o); + if (lane >= o) v += w; + } + wsum[s][lane] = v; + } + } + __syncthreads(); + const int off = (warp ? wsum[0][warp - 1] : 0) + a - c; + const int ioff = (warp ? wsum[1][warp - 1] : 0) + b - tiles; + if (tid == 0) { + counts[0] = wsum[1][31]; + counts[1] = wsum[2][31]; + } + if (tid < E) + for (int j = 0; j < tiles; ++j) { + int* it = items + 3 * (ioff + j); + it[0] = tid; + it[1] = off + T * j; + it[2] = min(T, c - T * j); + } + return off; +} + +// One block: pairs p = row * slots + slot grouped by expert (pair order within an expert). +__global__ void __launch_bounds__(PLAN_THREADS) + plan_kernel(const int* __restrict__ picks, int P, int E, int T, int* __restrict__ members, + int* __restrict__ items, int* __restrict__ counts) { + __shared__ int pk[PSMALL]; + __shared__ int cnt[EMAX]; + __shared__ int off[EMAX]; + const int tid = threadIdx.x; + for (int e = tid; e < EMAX; e += PLAN_THREADS) cnt[e] = 0; + __syncthreads(); + for (int p = tid; p < P; p += PLAN_THREADS) { + const int e = picks[p]; + pk[p] = e; + if (e < E) atomicAdd(&cnt[e], 1); // a pick past the table (e >= E) is left out of the plan + } + __syncthreads(); + const int o = place_items(tid < E ? cnt[tid] : 0, E, T, items, counts); + if (tid < E) off[tid] = o; + __syncthreads(); + for (int p = tid; p < P; p += PLAN_THREADS) { + const int e = pk[p]; + if (e >= E) continue; + int rank = 0; + for (int q = 0; q < p; ++q) rank += pk[q] == e; + members[off[e] + rank] = p; + } +} + +// Wide calls, pass 1: block b ranks its pairs within each expert in pair order and writes its counts to hist[b][e]. +__global__ void __launch_bounds__(PLAN_THREADS) + plan_rank(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, int* __restrict__ hist) { + __shared__ int cnt[EMAX]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + for (int e = tid; e < E; e += PLAN_THREADS) cnt[e] = 0; + __syncthreads(); + const int p = blockIdx.x * PLAN_THREADS + tid; + const bool ok = p < P && picks[p] < E; // a pick past the table is left out of the plan + const int e = ok ? picks[p] : -1 - lane; + const unsigned same = __match_any_sync(0xffffffffu, e); + const int below = __popc(same & ((1u << lane) - 1u)); + for (int w = 0; w < PLAN_THREADS / 32; ++w) { + if (warp == w) { + const int base = ok ? cnt[e] : 0; + __syncwarp(); + if (ok && below == 0) cnt[e] = base + __popc(same); + if (ok) rank[p] = base + below; + } + __syncthreads(); + } + for (int x = tid; x < E; x += PLAN_THREADS) hist[(size_t)blockIdx.x * E + x] = cnt[x]; +} + +// Pass 2, one block: each block's first member for each expert in place of its count, and the items. +__global__ void __launch_bounds__(PLAN_THREADS) + plan_offsets(int nblk, int E, int T, int* __restrict__ hist, int* __restrict__ items, int* __restrict__ counts) { + const int tid = threadIdx.x; + int c = 0; + if (tid < E) + for (int b = 0; b < nblk; ++b) { + int* h = hist + (size_t)b * E + tid; + const int v = *h; + *h = c; + c += v; + } + const int off = place_items(c, E, T, items, counts); + if (tid < E) + for (int b = 0; b < nblk; ++b) hist[(size_t)b * E + tid] += off; +} + +// Pass 3: every pair to its place. +__global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const int* __restrict__ rank, + const int* __restrict__ hist, int* __restrict__ members) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p < P && picks[p] < E) members[hist[(size_t)(p / PLAN_THREADS) * E + picks[p]] + rank[p]] = p; +} +} // namespace + +// The specializations the Python build instantiated. +namespace { +template __global__ void expert_kernel<(int)64, (int)1, (int)0, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void expert_kernel<(int)64, (int)1, (int)1, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void expert_kernel<(int)64, (int)2, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +} // namespace diff --git a/zig/kernels/cuda/glm/experts_sym.cuh b/zig/kernels/cuda/glm/experts_sym.cuh new file mode 100644 index 0000000000..f350743112 --- /dev/null +++ b/zig/kernels/cuda/glm/experts_sym.cuh @@ -0,0 +1,137 @@ +// GLM-5.3-Flash's symmetric routed experts: flashnext/experts.cuh with a group's block holding its scales alone (the +// bias of a symmetric int4 group is -8 x its scale, exact in bf16): 16 words a 32-row tile instead of 32. +// Shared by experts_sym.cu (decode form) and experts_sym_prefill.cu (prefill form). +#pragma once + +#include +#include + +namespace { + +constexpr int NTW = 4; // n8 tiles a warp: 32 output columns +constexpr int COLS = 8 * NTW; +constexpr uint32_t LOW = 0x000F000Fu; +constexpr uint32_t K128 = 0x43004300u; // the bf16 pair (128, 128) + +template +struct Geo { + static constexpr int KS = GS / 16; // mma k-steps a group + static constexpr int WV = NTW * GS / 128; // weight uint4 a lane a group + static constexpr int XV = GS / 32; // input uint4 a row a lane a group + static constexpr int SBV = NTW / 4; // scale uint4 a quad a group (the bias: -8 x the scale) +#ifdef GLM_SYM3 + // 3-bit codes (4..11 in the 4-bit words): each lane's uint4 c as three 32-bit planes (two of low 2-bit fields, + // one of high bits), 32 apart a plane: 24 uint4 a c instead of 32 + static constexpr int WSEC = 24 * WV; +#else + static constexpr int WSEC = 32 * WV; +#endif + static constexpr int BLOCK = WSEC + 4 * SBV; +}; + +__device__ __forceinline__ void cp16(void* dst, const void* src) { + const uint32_t d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(d), "l"(src)); +} + +__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n" ::); } + +template +__device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__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; +} + +// nibbles at bits sh and sh + 16 -> the exact bf16 pair (q0, q1): (128 + q) - 128 +#ifdef GLM_SYM3 +// 8 2-bit fields (16 bits) to the low 2 bits of 8 nibbles; 8 bits to bit 0 of 8 nibbles +__device__ __forceinline__ uint32_t spread2(uint32_t x) { + x = (x | (x << 8)) & 0x00FF00FFu; + x = (x | (x << 4)) & 0x0F0F0F0Fu; + return (x | (x << 2)) & 0x33333333u; +} +__device__ __forceinline__ uint32_t spread1(uint32_t x) { + x = (x | (x << 12)) & 0x000F000Fu; + x = (x | (x << 6)) & 0x03030303u; + return (x | (x << 3)) & 0x11111111u; +} +// planes p0 (words 0, 1's low fields), p1 (words 2, 3's), p2 (their high bits, 8 a word) -> the 4-bit words (code + 4) +__device__ __forceinline__ uint4 unpack3(uint32_t p0, uint32_t p1, uint32_t p2) { + uint4 r; + r.x = (spread2(p0 & 0xFFFFu) | (spread1(p2 & 0xFFu) << 2)) + 0x44444444u; + r.y = (spread2(p0 >> 16) | (spread1((p2 >> 8) & 0xFFu) << 2)) + 0x44444444u; + r.z = (spread2(p1 & 0xFFFFu) | (spread1((p2 >> 16) & 0xFFu) << 2)) + 0x44444444u; + r.w = (spread2(p1 >> 16) | (spread1(p2 >> 24) << 2)) + 0x44444444u; + return r; +} +#endif + +__device__ __forceinline__ uint32_t nib2(uint32_t w, int sh) { + uint32_t v = ((w >> sh) & LOW) | K128; + const uint32_t k = K128; + __nv_bfloat162 r = __hsub2(*reinterpret_cast<__nv_bfloat162*>(&v), *reinterpret_cast(&k)); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ void mma(float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, + uint32_t b1) { + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ float lo_f(uint32_t v) { return __uint_as_float(v << 16); } +__device__ __forceinline__ float hi_f(uint32_t v) { return __uint_as_float(v & 0xFFFF0000u); } +__device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +// SwiGLU as the families define it: bf16(bf16(silu(g)) * u) on bf16 g and u, clipped first when limit > 0 +__device__ __forceinline__ float swiglu(float g, float u, float limit) { + float gv = bf(g), uv = bf(u); + if (limit > 0.f) { + gv = fminf(gv, limit); + uv = fminf(fmaxf(uv, -limit), limit); + } + return bf(gv / (1.f + expf(-gv))) * uv; +} + +__device__ __forceinline__ float relu2(float a) { + const float u = fmaxf(bf(a), 0.f); + return u * u; +} + +// EPI 0: fp32 out (down); 1: bf16 relu(bf16(acc))^2; 2: bf16 SwiGLU(matrix 0, matrix 1); 3: bf16 out (prefill down) +template +__device__ __forceinline__ void epilogue(const float (&acc)[M][RT][NTW][4], int r, void* out, int N, int col0, int pr0, + int pr1, bool v0, bool v1, float limit) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + if (!(h ? v1 : v0)) continue; + const size_t row = (size_t)(h ? pr1 : pr0) * N; +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const int col = col0 + 8 * j; + const float a0 = acc[0][r][j][2 * h], a1 = acc[0][r][j][2 * h + 1]; + if constexpr (EPI == 0) { + *reinterpret_cast(reinterpret_cast(out) + row + col) = make_float2(a0, a1); + } else { + float o0, o1; + if constexpr (EPI == 1) { + o0 = relu2(a0); + o1 = relu2(a1); + } else if constexpr (EPI == 2) { + o0 = swiglu(a0, acc[M - 1][r][j][2 * h], limit); + o1 = swiglu(a1, acc[M - 1][r][j][2 * h + 1], limit); + } else { + o0 = a0; + o1 = a1; + } + *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<__nv_bfloat16*>(out) + row + col) = + __floats2bfloat162_rn(o0, o1); + } + } + } +} + +} // namespace diff --git a/zig/kernels/cuda/glm/experts_sym_prefill.cu b/zig/kernels/cuda/glm/experts_sym_prefill.cu new file mode 100644 index 0000000000..fa15fb178a --- /dev/null +++ b/zig/kernels/cuda/glm/experts_sym_prefill.cu @@ -0,0 +1,175 @@ +// GLM-5.3-Flash's symmetric routed experts, prefill form: flashnext/experts_prefill.cu reading experts_sym.cuh's +// blocks, each group's bias pair -8 x its scale pair (exact in bf16). +// Prefill: w = bf16(fma(q, s, b)), one fp32 chain over K; chunk-invariant bits, not decode's (replies prefill again). + +#include +#include + +#include "experts_sym.cuh" + +namespace { +// a bf16 pair times -8 (exact: the sign flipped, the exponent raised by 3; the scales are normal and far from overflow) +__device__ __forceinline__ uint32_t neg8(uint32_t s2) { + const __nv_bfloat162 v = *reinterpret_cast(&s2); + const __nv_bfloat162 r = __hmul2(v, __floats2bfloat162_rn(-8.f, -8.f)); + return *reinterpret_cast(&r); +} +} // namespace + +namespace { + +template +struct Pre { + static constexpr int THREADS = WM * WN * 32, BM = 16 * RT * WM, SG = 64 / GS, XC = 8, WB = M * Geo::BLOCK; + static constexpr int XU = BM * XC, WU = WN * SG * WB, SU = XU + WU; // uint4 a stage: X rows, weight blocks + static constexpr int XPT = (XU + THREADS - 1) / THREADS; + static constexpr int STAGE_BYTES = SU * 16; + static constexpr int STAGES = STAGE_BYTES * 4 <= 49152 ? 4 : STAGE_BYTES * 3 <= 49152 ? 3 : 2; + static __device__ __forceinline__ int xslot(int r, int c) { + return r * XC + (GS == 32 ? c ^ ((r & 1) << 2) : c ^ (r & 1)); + } +}; + +__device__ __forceinline__ uint32_t hfma2(uint32_t q, uint32_t s, uint32_t b) { + __nv_bfloat162 r = __hfma2(*reinterpret_cast<__nv_bfloat162*>(&q), *reinterpret_cast<__nv_bfloat162*>(&s), + *reinterpret_cast<__nv_bfloat162*>(&b)); + return *reinterpret_cast(&r); +} + +template +__global__ void __launch_bounds__(WM * WN * 32) + prefill_kernel(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; + using P = Pre; + constexpr int STAGES = P::STAGES; + extern __shared__ uint4 sm[]; + const int nbt = (NB + WN - 1) / WN, it = blockIdx.x / nbt, cbt = blockIdx.x - it * nbt; + if (it >= __ldg(counts)) return; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, wm = warp / WN, wn = warp - wm * WN; + const int t = lane & 3, gq = lane >> 2, cb = cbt * WN + wn, cbs = min(WN, NB - cbt * WN); + const __nv_bfloat16* xsrc[P::XPT]; + int xdst[P::XPT]; +#pragma unroll + for (int i = 0; i < P::XPT; ++i) { + const int q = tid + i * P::THREADS, r = q / P::XC, c = q - r * P::XC; + xdst[i] = P::xslot(r, c); + xsrc[i] = nullptr; + if (q < P::XU && r < cnt) { + const int p = __ldg(members + first + r); + xsrc[i] = X + (size_t)(slots ? p / slots : p) * x_stride + 8 * c; + } + } + const uint4* wsrc = W + ((size_t)e * NB + cbt * WN) * (size_t)KG * P::WB; + auto stage = [&](int s, int g0) { // groups g0 .. g0 + ng - 1 (ng < SG only at an odd end) + const int ng = min(P::SG, KG - g0); + uint4* xs = sm + s * P::SU; +#pragma unroll + for (int i = 0; i < P::XPT; ++i) + if (xsrc[i] && ((tid + i * P::THREADS) % P::XC) < ng * (GS / 8)) cp16(xs + xdst[i], xsrc[i] + (size_t)g0 * GS); + uint4* ws = xs + P::XU; + for (int q = tid; q < cbs * P::SG * P::WB; q += P::THREADS) { + const int j = q / (P::SG * P::WB), o = q - j * P::SG * P::WB; + if (o < ng * P::WB) cp16(ws + q, wsrc + ((size_t)j * KG + g0) * P::WB + o); + } + }; + const int base = 16 * RT * wm; + const int nt = max(0, min(RT, (cnt - base + 15) >> 4)); + const bool active = nt > 0 && cb < NB; + 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 int steps = (KG + P::SG - 1) / P::SG; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < steps) stage(s, s * P::SG); + cp_commit(); + } + const uint32_t pick = (gq & 1) ? 0x3232u : 0x1010u; // this lane's column: the high or low half of a pair + for (int k = 0; k < steps; ++k) { + cp_wait(); + __syncthreads(); + if (k + STAGES - 1 < steps) stage((k + STAGES - 1) % STAGES, (k + STAGES - 1) * P::SG); + cp_commit(); + if (!active) continue; + const uint4* xs = sm + (k % STAGES) * P::SU; +#pragma unroll + for (int gi = 0; gi < P::SG; ++gi) { + if (k * P::SG + gi >= KG) break; + const uint4* ws = xs + P::XU + (wn * P::SG + gi) * P::WB; + uint4 wv[M][G::WV]; + uint32_t s2[M][NTW], b2[M][NTW]; +#pragma unroll + for (int m = 0; m < M; ++m) { +#pragma unroll +#ifdef GLM_SYM3 + for (int c = 0; c < G::WV; ++c) { + const uint32_t* s32 = reinterpret_cast(ws + m * G::BLOCK) + c * 96 + lane; + wv[m][c] = unpack3(s32[0], s32[32], s32[64]); + } +#else + for (int c = 0; c < G::WV; ++c) wv[m][c] = ws[m * G::BLOCK + c * 32 + lane]; +#endif + const uint4 sp = ws[m * G::BLOCK + G::WSEC + (gq >> 1) * G::SBV]; +#pragma unroll + for (int j = 0; j < NTW; ++j) { + s2[m][j] = __byte_perm(comp(sp, j), 0u, pick); + b2[m][j] = neg8(s2[m][j]); + } + } +#pragma unroll + for (int kk = 0; kk < G::KS / 2; ++kk) { + uint4 xa[RT], xb[RT]; +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + const int r0 = base + 16 * r + gq; + xa[r] = xs[P::xslot(r0, gi * (GS / 8) + t * G::XV + kk)]; + xb[r] = xs[P::xslot(r0 + 8, gi * (GS / 8) + t * G::XV + kk)]; + } +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const int wi = j * (GS / 32) + kk; + const uint32_t word = comp(wv[m][wi >> 2], wi & 3); +#pragma unroll + for (int h = 0; h < 2; ++h) { + const uint32_t b0 = hfma2(nib2(word, 8 * h), s2[m][j], b2[m][j]); + const uint32_t b1 = hfma2(nib2(word, 8 * h + 4), s2[m][j], b2[m][j]); +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + mma(acc[m][r][j], comp(xa[r], 2 * h), comp(xb[r], 2 * h), comp(xa[r], 2 * h + 1), + comp(xb[r], 2 * h + 1), b0, b1); + } + } + } + } + } + } + if (!active) return; +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + const int m0 = base + 16 * r + gq, m1 = m0 + 8; + const bool v0 = m0 < cnt, v1 = m1 < cnt; + const int p0 = v0 ? __ldg(members + first + m0) : 0, p1 = v1 ? __ldg(members + first + m1) : 0; + epilogue(acc, r, out, N, cb * COLS + 2 * t, p0, p1, v0, v1, limit); + } +} +} // namespace + +// The specializations the Python build instantiated. +namespace { +template __global__ void prefill_kernel<(int)64, (int)1, (int)0, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)64, (int)1, (int)1, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)64, (int)1, (int)3, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)64, (int)2, (int)2, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +} // namespace diff --git a/zig/kernels/cuda/glm/index.cu b/zig/kernels/cuda/glm/index.cu new file mode 100644 index 0000000000..d647fed2b4 --- /dev/null +++ b/zig/kernels/cuda/glm/index.cu @@ -0,0 +1,250 @@ +// GLM-5.3-Flash on CUDA: the DSA indexer, ours. A row writes its index key (LayerNorm) and gate into its stream's +// rings, the row that completes a pool of 4 writes the pooled key, each row scores its visible pools, keeps its 512 +// best (the lower pool on ties) and lists their tokens ascending. Rows name their own rings, pools and positions, so +// one launch serves any mix of streams, and a row's choice never depends on the rows beside it. + +#include +#include + +#include "../flashnext/qmm_frag.cuh" + +namespace { + +using namespace qmm_frag; + +constexpr int D = 128, IH = 32, POOL = 4, TOPK = 512, SEL_W = TOPK * POOL + POOL - 1; + +__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +// a score as a uint32 whose unsigned order is the scores' order (-0 counted as +0) +__device__ __forceinline__ uint32_t order_key(float s) { + const int b = __float_as_int(s + 0.0f); + return (uint32_t)(b ^ ((b >> 31) | (int)0x80000000)); +} + +} // namespace + +// row r: LayerNorm(k_raw) -> bf16 at its position of its key ring ``ik[r]``, the fp32 gate -> bf16 into ``ig[r]`` +// (ring pointers already offset so that row p of the positions is at p * 128). 128 threads. +extern "C" __global__ void glm_index_write(const __nv_bfloat16* __restrict__ kr, long long k_stride, const float* __restrict__ gr, + const __nv_bfloat16* __restrict__ lnw, const __nv_bfloat16* __restrict__ lnb, const unsigned long long* __restrict__ ik, + const unsigned long long* __restrict__ ig, const int* __restrict__ pos, float eps) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + __shared__ float red[4]; + const int r = blockIdx.x, d = threadIdx.x, lane = d & 31, warp = d >> 5; + const float x = __bfloat162float(kr[r * k_stride + d]); + auto sum = [&](float v) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) v += __shfl_xor_sync(0xffffffffu, v, off); + __syncthreads(); + if (lane == 0) red[warp] = v; + __syncthreads(); + return ((red[0] + red[1]) + red[2]) + red[3]; + }; + const float mean = sum(x) / D; + const float xc = x - mean; + const float var = sum(xc * xc) / D; + const float y = xc / sqrtf(var + eps) * __bfloat162float(lnw[d]) + __bfloat162float(lnb[d]); + const long long p = pos[r]; + reinterpret_cast<__nv_bfloat16*>(ik[r])[p * D + d] = __float2bfloat16_rn(y); + reinterpret_cast<__nv_bfloat16*>(ig[r])[p * D + d] = __float2bfloat16_rn(gr[(long long)r * D + d]); +} + +// the row that completes a pool (position % 4 == 3) writes its pooled key: the 4 keys weighted by the softmax of +// gate + ape, each product rounded to bf16, summed in order -> pk[r] row pos / 4. 128 threads. Run after +// glm_index_write over the same rows (a pool's earlier keys may be the window's or earlier ones still in the ring). +extern "C" __global__ void glm_pool_keys(const unsigned long long* __restrict__ ik, const unsigned long long* __restrict__ ig, + const __nv_bfloat16* __restrict__ ape, const unsigned long long* __restrict__ pk, const int* __restrict__ pos) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int r = blockIdx.x, d = threadIdx.x; + const long long q = pos[r]; + if ((q & 3) != 3) return; + const long long p = q >> 2; + const __nv_bfloat16* K = reinterpret_cast(ik[r]); + const __nv_bfloat16* G = reinterpret_cast(ig[r]); + float l[4], k[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + l[i] = __bfloat162float(G[(4 * p + i) * D + d]) + __bfloat162float(ape[i * D + d]); + k[i] = __bfloat162float(K[(4 * p + i) * D + d]); + } + const float m = fmaxf(fmaxf(l[0], l[1]), fmaxf(l[2], l[3])); + float e[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) e[i] = expf(l[i] - m); + const float s = ((e[0] + e[1]) + e[2]) + e[3]; + float t[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) t[i] = bf(bf(e[i] / s) * k[i]); + reinterpret_cast<__nv_bfloat16*>(pk[r])[p * D + d] = __float2bfloat16_rn(((t[0] + t[1]) + t[2]) + t[3]); +} + +// each row's pool scores: s_p = sum_h w_h relu(scale * qi_h . pk_p) for its visible pools (p < (pos + 1) / 4), -inf +// for the rest of the np columns. qi [R, 32, 128], w the head weights (bf16, row stride w_stride). grid (R, np / 64), +// 128 threads: a block of 64 pools on tensor cores, the heads summed in a fixed order. +extern "C" __global__ void __launch_bounds__(128) glm_index_scores(const __nv_bfloat16* __restrict__ qi, const __nv_bfloat16* __restrict__ w, + long long w_stride, const unsigned long long* __restrict__ pk, float* __restrict__ out, const int* __restrict__ pos, int np, + float scale, float wscale) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + __shared__ __align__(128) unsigned char sq[IH * D * 2]; + __shared__ __align__(128) unsigned char sk[64 * D * 2]; + __shared__ float hw[IH]; + const int r = blockIdx.x, p0 = blockIdx.y * 64, tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int npool = (pos[r] + 1) >> 2; + float* dst = out + (long long)r * np; + if (p0 >= npool) { + if (tid < 64 && p0 + tid < np) dst[p0 + tid] = -INFINITY; + return; + } + const __nv_bfloat16* P = reinterpret_cast(pk[r]); + // rows are 256 bytes, 16 chunks; chunk c of row x at c ^ (x & 7) within its group of 8 + auto sw = [](int x, int c) { return (c & ~7) | ((c ^ x) & 7); }; + for (int c = tid; c < IH * 16; c += 128) { + const int h = c >> 4, ch = c & 15; + cp16(sq + h * 256 + sw(h, ch) * 16, qi + ((long long)r * IH + h) * D + ch * 8); + } + for (int c = tid; c < 64 * 16; c += 128) { + const int pp = c >> 4, ch = c & 15; + const bool ok = p0 + pp < npool; + cp16z(sk + pp * 256 + sw(pp, ch) * 16, P + (long long)(ok ? p0 + pp : 0) * D + ch * 8, ok); + } + commit(); + if (tid < IH) hw[tid] = __bfloat162float(w[r * w_stride + tid]) * wscale; + wait<0>(); + __syncthreads(); + // warp: pools 16 warp .. +16 (two n8 tiles), all 32 heads (two m16 tiles) + float d[2][2][4]; +#pragma unroll + for (int i = 0; i < 2; ++i) +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) d[i][j][e] = 0.0f; +#pragma unroll + for (int ks = 0; ks < D / 16; ++ks) { + uint32_t a[2][4]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int ar = i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + ldmatrix4(a[i], sq + ar * 256 + sw(ar, ks * 2 + (lane >> 4)) * 16); + } + uint32_t b[4]; + const int br = warp * 16 + (lane & 7) + ((lane >> 4) << 3); + ldmatrix4(b, sk + br * 256 + sw(br, ks * 2 + ((lane >> 3) & 1)) * 16); +#pragma unroll + for (int i = 0; i < 2; ++i) { + mma(d[i][0], a[i], b[0], b[1]); + mma(d[i][1], a[i], b[2], b[3]); + } + } + const int g = lane >> 2, tq = lane & 3; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int u = 0; u < 2; ++u) { + // heads g, g + 8, g + 16, g + 24 of this lane, then the 8 lanes of the column in a xor tree + float s = (hw[g] * fmaxf(d[0][j][u] * scale, 0.0f) + hw[g + 8] * fmaxf(d[0][j][2 + u] * scale, 0.0f)) + + (hw[g + 16] * fmaxf(d[1][j][u] * scale, 0.0f) + hw[g + 24] * fmaxf(d[1][j][2 + u] * scale, 0.0f)); +#pragma unroll + for (int off = 4; off < 32; off <<= 1) s += __shfl_xor_sync(0xffffffffu, s, off); + const int p = p0 + warp * 16 + j * 8 + 2 * tq + u; + if (g == 0 && p < np) dst[p] = p < npool ? s : -INFINITY; + } +} + +// each row's 512 best pools of its visible ones (the lower pool on ties), ascending (int64): a radix select (8 bits a +// pass) finds the 512th best score, then one pass in pool order writes the pools above it and the lowest ties. +// scores [R, np]. 1024 threads. +extern "C" __global__ void __launch_bounds__(1024) glm_index_select(const float* __restrict__ sc, long long* __restrict__ out, int np, + const int* __restrict__ pos) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + __shared__ int hist[256]; + __shared__ int warp_tot[32]; + __shared__ uint32_t s_digit; + __shared__ int s_need; + const int r = blockIdx.x, tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const float* row = sc + (long long)r * np; + const int limit = min(np, max(TOPK, (pos[r] + 1) >> 2)); + uint32_t prefix = 0, fixed = 0; + int need = TOPK; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) hist[tid] = 0; + __syncthreads(); + for (int i = tid; i < limit; i += 1024) { + const uint32_t u = order_key(row[i]); + if ((u & fixed) == prefix) atomicAdd(&hist[(u >> shift) & 0xFF], 1); + } + __syncthreads(); + if (tid == 0) { // the highest digit with at least ``need`` pools at or above it + int above = 0, digit = 0; + for (int b = 255; b >= 0; --b) { + if (above + hist[b] >= need) { + digit = b; + break; + } + above += hist[b]; + } + s_digit = (uint32_t)digit; + s_need = need - above; + } + __syncthreads(); + prefix |= s_digit << shift; + fixed |= 0xFFu << shift; + need = s_need; + __syncthreads(); + } + // pools above the threshold, and the first ``need`` equal to it, in pool order + int written = 0, equal_seen = 0; + for (int c = 0; c < limit; c += 1024) { + const int i = c + tid; + const uint32_t u = i < limit ? order_key(row[i]) : 0u; + const int eq = (i < limit && u == prefix) ? 1 : 0; + // the block's exclusive prefix of eq, then of take + auto scan = [&](int v, int& total) { + int x = v; +#pragma unroll + for (int off = 1; off < 32; off <<= 1) { + const int y = __shfl_up_sync(0xffffffffu, x, off); + if (lane >= off) x += y; + } + __syncthreads(); + if (lane == 31) warp_tot[warp] = x; + __syncthreads(); + int before = 0; + for (int w = 0; w < warp; ++w) before += warp_tot[w]; + total = 0; + for (int w = 0; w < 32; ++w) total += warp_tot[w]; + return before + x - v; + }; + int eq_total; + const int eq_before = scan(eq, eq_total); + const int take = (i < limit && (u > prefix || (eq && eq_before + equal_seen < need))) ? 1 : 0; + int take_total; + const int at = scan(take, take_total); + if (take) out[(long long)r * TOPK + written + at] = i; + written += take_total; + equal_seen += eq_total; + } +} + +// each row's attended tokens (its 512 pools' tokens ascending, then the incomplete last pool's visible tokens, -1 past +// the row) and their count: 2048 plus those when the row has more than 512 complete pools, else 0 (a dense row). +extern "C" __global__ void glm_index_tokens(const long long* __restrict__ pools, const int* __restrict__ pos, + int* __restrict__ tokens, int* __restrict__ counts) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int r = blockIdx.x; + const long long q = pos[r]; + const long long npool = (q + 1) / 4; + int* o = tokens + (long long)r * SEL_W; + for (int i = threadIdx.x; i < TOPK * POOL; i += blockDim.x) o[i] = (int)(pools[(long long)r * TOPK + i / 4] * 4 + i % 4); + if (threadIdx.x < 3) { + const long long t = npool * 4 + threadIdx.x; + o[TOPK * POOL + threadIdx.x] = t <= q ? (int)t : -1; + } + if (threadIdx.x == 0) { + int tail = 0; + for (int j = 0; j < 3; ++j) tail += (npool * 4 + j) <= q; + counts[r] = npool > TOPK ? TOPK * POOL + tail : 0; + } +} diff --git a/zig/kernels/cuda/glm/kda.cu b/zig/kernels/cuda/glm/kda.cu new file mode 100644 index 0000000000..577923d301 --- /dev/null +++ b/zig/kernels/cuda/glm/kda.cu @@ -0,0 +1,416 @@ +// From src/tensorfold/families/glm5_next/cuda/kda.cu (python-0.6), with the chain's body a device function and a +// segmented chain (any number of streams' windows in one launch) added. +// GLM-5.3-Flash's KDA (Kimi delta attention) on CUDA: one block of 1024 threads per head, a chain of R rows. +// +// Per row, following the Hugging Face definition: the depthwise conv over [conv state; q|k|v rows] (4 taps, +// fp32) with SiLU and one bf16 rounding; fp32 L2 norms of q and k (eps inside the sum, q times DK^-0.5); the +// per-channel decay g_i = exp(lower * sigmoid(exp(A_log) * (a_i + dt_bias_i))) in fp32; beta = bf16(sigmoid(b)); +// the delta rule in fp32 (decay along the key channel, read with k, correct toward v, read out with q); +// read-out rounded to bf16; then the gated RMSNorm bf16(w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(gate)). +// Warp w owns value rows 4w .. 4w + 3, lane l the key columns 4l .. 4l + 3. The state update is one routine +// (``update``) shared with ``replay``, compiled without FMA contraction, so a replayed prefix of a window gives +// the bits of the serial steps. + +#include +#include + +namespace { + +constexpr int DK = 128, DV = 128, TAPS = 4; + +__device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float warp_sum(float x) { + for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o); + return x; +} + +__device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); } + +// One delta-rule step on this thread's 4 x 4 block of the state: per-channel decay, read (k), correct toward v. +__device__ __forceinline__ void update(float (&s)[4][4], const float (&kk)[4], const float (&gg)[4], + const float* vrow, int warp, 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] * gg[i]; + 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; + } +} + +__device__ __forceinline__ void chain_body(int h, + int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, + const __nv_bfloat16* __restrict__ A, int a_stride, const __nv_bfloat16* __restrict__ G, int g_stride, + 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, float lower, int rows, + __nv_bfloat16* __restrict__ out, float* __restrict__ state_out, + float* __restrict__ k_save, __nv_bfloat16* __restrict__ v_save, float* __restrict__ g_save, + float* __restrict__ b_save) { + const int C = 3 * H * DK; // conv channels: q | k | v + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float qs[DK], ks[DK], vs[DV], ys[DV], gs[DK]; + __shared__ float beta_s, rinv; + int c = -1; + if (t < 3 * DK) c = (t / DK) * H * DK + h * DK + (t % DK); + float s[4][4]; + const size_t sbase = (size_t)h * 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 float decay_rate = expf(a_log[h]); + 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[(size_t)at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * p_stride + c]); + acc = acc + __bfloat162float(cw[(size_t)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; + } else if (t >= 512 && t < 512 + DK) { + const int i = t - 512; + const float a = __bfloat162float(A[(size_t)r * a_stride + h * DK + i]) + dt_bias[h * DK + i]; + gs[i] = expf(lower * sigmoidf_(decay_rate * a)); + } else if (t == 1023) { + beta_s = bf(sigmoidf_(__bfloat162float(P[(size_t)r * p_stride + b_off + h]))); + } + __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); + __syncwarp(); +#pragma unroll + for (int i = 0; i < 4; ++i) { + float y = v4[i] * inv; + if (warp == 0) y = y * (1.0f / sqrtf((float)DK)); + x[lane * 4 + i] = y; + } + } + __syncthreads(); + const float beta = beta_s; + float kk[4], qq[4], gg[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { kk[i] = ks[lane * 4 + i]; qq[i] = qs[lane * 4 + i]; gg[i] = gs[lane * 4 + i]; } + update(s, kk, gg, vs, warp, 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); + } + if (k_save != nullptr) { + const size_t base = ((size_t)r * H + h) * DK; + if (t < DK) { k_save[base + t] = ks[t]; g_save[base + t] = gs[t]; } + if (t >= DK && t < DK + DV) v_save[base + t - DK] = __float2bfloat16_rn(vs[t - DK]); + if (t == 0) b_save[r * H + h] = beta; + } + __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 = ys[t] * rinv; + const float yw = __bfloat162float(norm_w[t]) * yn; + const float gate = __bfloat162float(G[(size_t)r * g_stride + h * DV + t]); + out[(size_t)r * H * DV + h * DV + t] = __float2bfloat16_rn(yw * sigmoidf_(gate)); + } + __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]; + } +} + +__global__ void __launch_bounds__(1024) chain_kernel( + int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, + const __nv_bfloat16* __restrict__ A, int a_stride, const __nv_bfloat16* __restrict__ G, int g_stride, + 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, float lower, int rows, + __nv_bfloat16* __restrict__ out, float* __restrict__ state_out, + float* __restrict__ k_save, __nv_bfloat16* __restrict__ v_save, float* __restrict__ g_save, + float* __restrict__ b_save) { + chain_body(blockIdx.x, H, P, p_stride, b_off, A, a_stride, G, g_stride, cs, cw, state_in, a_log, dt_bias, norm_w, eps, + lower, rows, out, state_out, k_save, v_save, g_save, b_save); +} + +} // namespace + +// One stream's window in a batched step: its rows (from ``first`` of the step's rows), its states (the live parity +// in, the other out), conv windows, replay inputs and window projections, each the base of layer 0. +struct KdaSeg { + int first, rows; + const float* rec_in; + float* rec_out; + const __nv_bfloat16* conv; + float* k_save; + __nv_bfloat16* v_save; + float* g_save; + float* b_save; + __nv_bfloat16* proj; +}; + +// chain_kernel for every stream of a batched step at layer ``li``: grid (H, streams). The segment's projection rows +// (the conv's channels) are kept in its stream's window buffer for the commit's conv shift. +extern "C" __global__ void __launch_bounds__(1024) glm_kda_chain_seg(const KdaSeg* __restrict__ segs, int li, + int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, + const __nv_bfloat16* __restrict__ A, int a_stride, const __nv_bfloat16* __restrict__ G, int g_stride, + const __nv_bfloat16* __restrict__ cw, const float* __restrict__ a_log, const float* __restrict__ dt_bias, + const __nv_bfloat16* __restrict__ norm_w, float eps, float lower, __nv_bfloat16* __restrict__ out, + int max_rows, int proj_w) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const KdaSeg sg = segs[blockIdx.y]; + const int C = 3 * H * DK; + const size_t rec = (size_t)li * H * DV * DK, conv = (size_t)li * 3 * C; + const size_t kv = (size_t)li * max_rows * H * DK, bb = (size_t)li * max_rows * H; + const __nv_bfloat16* p = P + (size_t)sg.first * p_stride; + if (blockIdx.x == 0) { // the window's conv channels, for the commit (one block a stream copies them) + __nv_bfloat16* dst = sg.proj + (size_t)li * max_rows * proj_w; + for (int i = threadIdx.x; i < sg.rows * C; i += blockDim.x) { + const int r = i / C, c = i % C; + dst[(size_t)r * proj_w + c] = p[(size_t)r * p_stride + c]; + } + } + chain_body(blockIdx.x, H, p, p_stride, b_off, A + (size_t)sg.first * a_stride, a_stride, G + (size_t)sg.first * g_stride, + g_stride, sg.conv + conv, cw, sg.rec_in + rec, a_log, dt_bias, norm_w, eps, lower, sg.rows, + out + (size_t)sg.first * H * DV, sg.rec_out + rec, sg.k_save + kv, sg.v_save + kv, sg.g_save + kv, sg.b_save + bb); +} + +namespace { + +__global__ void __launch_bounds__(1024) replay_kernel( + int H, const float* __restrict__ state_in, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, + const float* __restrict__ b_save, int rows, float* __restrict__ state_out) { + const int h = blockIdx.x; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float vs[DV]; + float s[4][4]; + const size_t sbase = (size_t)h * 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]; + for (int r = 0; r < rows; ++r) { + const size_t base = ((size_t)r * H + h) * DK; + if (t < DV) vs[t] = __bfloat162float(v_save[base + t]); + __syncthreads(); + float kk[4], gg[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { kk[i] = k_save[base + lane * 4 + i]; gg[i] = g_save[base + lane * 4 + i]; } + update(s, kk, gg, vs, warp, b_save[r * H + h]); + __syncthreads(); + } +#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]; +} + +// All layers at once: block (layer, head); per-layer strides of the state buffers and saved rows. +__global__ void __launch_bounds__(1024) replay_layers_kernel( + int H, const float* __restrict__ state_in, size_t state_stride, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, + size_t kv_stride, size_t b_stride, int rows, float* __restrict__ state_out) { + const int layer = blockIdx.x / H, h = blockIdx.x % H; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float vs[DV]; + float s[4][4]; + const float* sin = state_in + layer * state_stride; + float* sout = state_out + layer * state_stride; + const float* ks = k_save + layer * kv_stride; + const __nv_bfloat16* vsv = v_save + layer * kv_stride; + const float* gsv = g_save + layer * kv_stride; + const float* bsv = b_save + layer * b_stride; + const size_t sbase = (size_t)h * DV * DK; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = sin[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i]; + for (int r = 0; r < rows; ++r) { + const size_t base = ((size_t)r * H + h) * DK; + if (t < DV) vs[t] = __bfloat162float(vsv[base + t]); + __syncthreads(); + float kk[4], gg[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { kk[i] = ks[base + lane * 4 + i]; gg[i] = gsv[base + lane * 4 + i]; } + update(s, kk, gg, vs, warp, bsv[r * H + h]); + __syncthreads(); + } +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) sout[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; +} + +// Long windows: state-independent work for all rows (prep_kernel, out_kernel), the delta rule alone row by row (step_kernel); chain_kernel's bits. + +// Block (row, head), 512 threads: conv + SiLU of the 3 x 128 q|k|v channels, the decay, beta, then the q/k norms. +__global__ void __launch_bounds__(512) prep_kernel( + int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, + const __nv_bfloat16* __restrict__ A, int a_stride, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, const float* __restrict__ a_log, const float* __restrict__ dt_bias, + float lower, float* __restrict__ q_out, float* __restrict__ k_save, __nv_bfloat16* __restrict__ v_save, + float* __restrict__ g_save, float* __restrict__ b_save) { + const int C = 3 * H * DK; + const int r = blockIdx.x, h = blockIdx.y; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float qs[DK], ks[DK]; + const size_t base = ((size_t)r * H + h) * DK; + if (t < 3 * DK) { + const int c = (t / DK) * H * DK + h * DK + (t % DK); + 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[(size_t)at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * p_stride + c]); + acc = acc + __bfloat162float(cw[(size_t)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 v_save[base + t - 2 * DK] = __float2bfloat16_rn(act); + } else { + const int i = t - 3 * DK; + const float decay_rate = expf(a_log[h]); + const float a = __bfloat162float(A[(size_t)r * a_stride + h * DK + i]) + dt_bias[h * DK + i]; + g_save[base + i] = expf(lower * sigmoidf_(decay_rate * a)); + if (i == 0) b_save[r * H + h] = bf(sigmoidf_(__bfloat162float(P[(size_t)r * p_stride + b_off + h]))); + } + __syncthreads(); + if (warp < 2) { + const 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); + float* dst = warp == 0 ? q_out : k_save; +#pragma unroll + for (int i = 0; i < 4; ++i) { + float y = v4[i] * inv; + if (warp == 0) y = y * (1.0f / sqrtf((float)DK)); + dst[base + lane * 4 + i] = y; + } + } +} + +// Block (head, part): each warp owns 4 value rows; TR rows of saved k, q, decay, v and beta are staged in shared memory at a time. +template +__global__ void __launch_bounds__(WARPS * 32) step_kernel( + int H, const float* __restrict__ state_in, const float* __restrict__ q_in, const float* __restrict__ k_save, + const __nv_bfloat16* __restrict__ v_save, const float* __restrict__ g_save, const float* __restrict__ b_save, + int rows, __nv_bfloat16* __restrict__ y_out, float* __restrict__ state_out) { + constexpr int NT = WARPS * 32, VR = WARPS * 4; + const int h = blockIdx.x, v0 = blockIdx.y * VR; + const int lw = threadIdx.x >> 5, lane = threadIdx.x & 31, warp = blockIdx.y * WARPS + lw; + __shared__ __align__(16) float kt[TR][DK]; + __shared__ __align__(16) float qt[TR][DK]; + __shared__ __align__(16) float gt[TR][DK]; + __shared__ __align__(16) float vt[TR][VR]; + __shared__ float bt[TR]; + float s[4][4]; + const size_t sbase = (size_t)h * 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]; + for (int r0 = 0; r0 < rows; r0 += TR) { + const int n = rows - r0 < TR ? rows - r0 : TR; + __syncthreads(); + for (int x = threadIdx.x; x < n * DK; x += NT) { + const int rr = x / DK, c = x % DK; + const size_t base = ((size_t)(r0 + rr) * H + h) * DK + c; + kt[rr][c] = k_save[base]; + qt[rr][c] = q_in[base]; + gt[rr][c] = g_save[base]; + } + for (int x = threadIdx.x; x < n * VR; x += NT) { + const int rr = x / VR, c = x % VR; + vt[rr][c] = __bfloat162float(v_save[((size_t)(r0 + rr) * H + h) * DV + v0 + c]); + } + if (threadIdx.x < n) bt[threadIdx.x] = b_save[(r0 + threadIdx.x) * H + h]; + __syncthreads(); + for (int rr = 0; rr < n; ++rr) { + const float4 k4 = reinterpret_cast(kt[rr])[lane]; // one 16-byte load a lane: + const float4 q4 = reinterpret_cast(qt[rr])[lane]; // no bank conflicts + const float4 g4 = reinterpret_cast(gt[rr])[lane]; + const float4 v4 = reinterpret_cast(vt[rr])[lw]; + float kk[4] = {k4.x, k4.y, k4.z, k4.w}, qq[4] = {q4.x, q4.y, q4.z, q4.w}; + float gg[4] = {g4.x, g4.y, g4.z, g4.w}, vv[4] = {v4.x, v4.y, v4.z, v4.w}; + update(s, kk, gg, vv, 0, bt[rr]); +#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) y_out[((size_t)(r0 + rr) * H + h) * DV + warp * 4 + j] = __float2bfloat16_rn(o); + } + } + } + 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]; + } +} + +// Block (row, head), 128 threads: the gated RMSNorm of the read-out. +__global__ void __launch_bounds__(128) out_kernel( + int H, const __nv_bfloat16* __restrict__ y_in, const __nv_bfloat16* __restrict__ G, int g_stride, + const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out) { + const int r = blockIdx.x, h = blockIdx.y; + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float rinv; + const size_t base = ((size_t)r * H + h) * DV; + if (warp == 0) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { const float y = __bfloat162float(y_in[base + 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 = __bfloat162float(y_in[base + t]) * rinv; + const float yw = __bfloat162float(norm_w[t]) * yn; + const float gate = __bfloat162float(G[(size_t)r * g_stride + h * DV + t]); + out[(size_t)r * H * DV + h * DV + t] = __float2bfloat16_rn(yw * sigmoidf_(gate)); +} + +} // namespace + +// the Python build's instantiations +namespace { +template __global__ void step_kernel<(int)4, (int)16>(int, const float *, const float *, const float *, const __nv_bfloat16 *, const float *, const float *, int, __nv_bfloat16 *, float *); +} // namespace diff --git a/zig/kernels/cuda/glm/latent.cu b/zig/kernels/cuda/glm/latent.cu new file mode 100644 index 0000000000..5ce154b389 --- /dev/null +++ b/zig/kernels/cuda/glm/latent.cu @@ -0,0 +1,287 @@ +// GLM-5.3-Flash on CUDA: MLA attention on the 512-wide latent cache, ours. A program is one row's 32 heads over one +// chunk of 512 keys (online softmax in 16-key tiles, tensor-core scores and outputs), and the chunks are merged in +// order; nothing a program computes depends on another row, so a row has the same bits alone, in a decode window, in +// sixteen streams' windows or in a prompt chunk. Every row names its own cache (``lc``, the latent rows of its +// stream) and position (``pos``), so one launch serves any mix of streams. + +#include +#include + +#include "../flashnext/qmm_frag.cuh" + +namespace { + +using namespace qmm_frag; + +constexpr int H = 32, LW = 512, CH = 512, KT = 16, THREADS = 256; +constexpr int DENSE_LIMIT = 2051; // index_topk + pool - 1: rows from here attend to their selected tokens alone +constexpr int ROWB = LW * 2; // a latent row's bytes +constexpr int Q_BYTES = H * ROWB; +constexpr int KV_BYTES = KT * ROWB; +constexpr int SMEM = Q_BYTES + 2 * KV_BYTES + 2 * H * KT * 4 + H * KT * 2 + 3 * H * 4; + +// 16-byte chunk c of a latent row r sits at c ^ (r & 7): ldmatrix's 8 rows land in 8 bank groups +__device__ __forceinline__ int lsw(int r, int c) { return (c & ~7) | ((c ^ r) & 7); } + +template +__device__ __forceinline__ void chunks(const __nv_bfloat16* __restrict__ qa, const unsigned long long* __restrict__ lc, + const int* __restrict__ pos, const int* __restrict__ tok, const int* __restrict__ cnt, float* __restrict__ po, + float* __restrict__ pm, float* __restrict__ pl, int R, float scale) { + extern __shared__ __align__(128) unsigned char sm[]; + unsigned char* sq = sm; + unsigned char* skv = sm + Q_BYTES; + float* ssc = reinterpret_cast(skv + 2 * KV_BYTES); // [2 k halves][H][KT] score partials + __nv_bfloat16* sp = reinterpret_cast<__nv_bfloat16*>(ssc + 2 * H * KT); // [H][KT] probabilities + float* salpha = reinterpret_cast(sp + H * KT); + float* sm_m = salpha + H; + float* sm_l = sm_m + H; + const int r = blockIdx.x, c = blockIdx.y, tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const __nv_bfloat16* cache = reinterpret_cast(lc[r]); + // the keys this program reads: [lo, hi) of the chunk's positions (dense) or of the row's token list (sparse) + int lo = c * CH, hi; + if (SPARSE) { + hi = min(lo + CH, cnt[r]); + } else { + const int q = pos[r]; + hi = q >= DENSE_LIMIT ? lo : min(lo + CH, q + 1); + } + if (tid < H) { + sm_m[tid] = -INFINITY; + sm_l[tid] = 0.0f; + } + float o[16][4]; // warp: heads 16 (warp & 1) .. +16, latent 128 (warp >> 1) .. +128 +#pragma unroll + for (int j = 0; j < 16; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) o[j][e] = 0.0f; + const int ntiles = hi > lo ? (hi - lo + KT - 1) / KT : 0; + auto key_of = [&](int idx) -> long long { return SPARSE ? (long long)tok[(long long)r * 2051 + idx] : (long long)idx; }; + auto load_kv = [&](int buf, int t) { + unsigned char* dst = skv + buf * KV_BYTES; + for (int p = tid; p < KT * (ROWB / 16); p += THREADS) { + const int kr = p / (ROWB / 16), ch = p % (ROWB / 16); + const int idx = lo + t * KT + kr; + const bool ok = idx < hi; + const long long key = ok ? key_of(idx) : 0; + cp16z(dst + kr * ROWB + lsw(kr, ch) * 16, cache + key * LW + ch * 8, ok); + } + }; + if (ntiles > 0) { + for (int p = tid; p < H * (ROWB / 16); p += THREADS) { + const int h = p / (ROWB / 16), ch = p % (ROWB / 16); + cp16(sq + h * ROWB + lsw(h, ch) * 16, qa + ((long long)r * H + h) * LW + ch * 8); + } + load_kv(0, 0); + } + commit(); + const int wm = warp & 1, wn = warp >> 1; + for (int t = 0; t < ntiles; ++t) { + if (t + 1 < ntiles) load_kv((t + 1) & 1, t + 1); + commit(); + wait<1>(); + __syncthreads(); + const unsigned char* kv = skv + (t & 1) * KV_BYTES; + { // scores: warp (heads 16 (warp & 1), keys 8 ((warp >> 1) & 1), latent half warp >> 2) + const int mi = warp & 1, ni = (warp >> 1) & 1, kh = warp >> 2; + float d[4] = {0.0f, 0.0f, 0.0f, 0.0f}; +#pragma unroll 4 + for (int ks = 0; ks < 16; ++ks) { + const int kc = kh * 32 + ks * 2; // the k-step's first 16-byte chunk + uint32_t a[4]; + const int ar = mi * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + ldmatrix4(a, sq + ar * ROWB + lsw(ar, kc + (lane >> 4)) * 16); + uint32_t b[2]; + const int br = ni * 8 + (lane & 7); + ldmatrix2(b, kv + br * ROWB + lsw(br, kc + ((lane >> 3) & 1)) * 16); + mma(d, a, b[0], b[1]); + } + const int g = lane >> 2, tq = lane & 3; + float* dst = ssc + kh * H * KT; + dst[(mi * 16 + g) * KT + ni * 8 + 2 * tq] = d[0]; + dst[(mi * 16 + g) * KT + ni * 8 + 2 * tq + 1] = d[1]; + dst[(mi * 16 + g + 8) * KT + ni * 8 + 2 * tq] = d[2]; + dst[(mi * 16 + g + 8) * KT + ni * 8 + 2 * tq + 1] = d[3]; + } + __syncthreads(); + { // the online softmax: 8 threads a head, 2 keys each + const int h = tid >> 3, j = (tid & 7) * 2; + const int base = lo + t * KT; + float s[2]; + bool valid[2]; +#pragma unroll + for (int u = 0; u < 2; ++u) { + valid[u] = base + j + u < hi; + s[u] = valid[u] ? (ssc[h * KT + j + u] + ssc[H * KT + h * KT + j + u]) * scale : -INFINITY; + } + float tm = fmaxf(s[0], s[1]); +#pragma unroll + for (int off = 4; off > 0; off >>= 1) tm = fmaxf(tm, __shfl_xor_sync(0xffffffffu, tm, off)); + const float m = sm_m[h]; + const bool active = tm != -INFINITY; + const float nm = active ? fmaxf(m, tm) : m; + const float alpha = active ? (m == -INFINITY ? 0.0f : expf(m - nm)) : 1.0f; + float p[2], ps = 0.0f; +#pragma unroll + for (int u = 0; u < 2; ++u) { + p[u] = (valid[u] && active) ? expf(s[u] - nm) : 0.0f; + sp[h * KT + j + u] = __float2bfloat16_rn(p[u]); + } + ps = p[0] + p[1]; +#pragma unroll + for (int off = 1; off < 8; off <<= 1) ps += __shfl_xor_sync(0xffffffffu, ps, off); + __syncwarp(); + if ((tid & 7) == 0) { + salpha[h] = alpha; + sm_m[h] = nm; + sm_l[h] = sm_l[h] * alpha + ps; + } + } + __syncthreads(); + { // o = o * alpha + p . kv + const int g = lane >> 2; + const float a0 = salpha[wm * 16 + g], a1 = salpha[wm * 16 + g + 8]; + uint32_t a[4]; + const int ar = wm * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + ldmatrix4(a, reinterpret_cast(sp) + ar * (KT * 2) + (lane >> 4) * 16); +#pragma unroll + for (int j = 0; j < 16; j += 2) { + // two n8 tiles of latent columns: keys 0-7 / 8-15 of each, transposed into B fragments + uint32_t b[4]; + const int kr = (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = wn * 16 + j + (lane >> 4); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b[0]), "=r"(b[1]), "=r"(b[2]), "=r"(b[3]) + : "r"(smem(kv + kr * ROWB + lsw(kr, ch) * 16))); +#pragma unroll + for (int e = 0; e < 4; ++e) { + o[j][e] *= (e < 2 ? a0 : a1); + o[j + 1][e] *= (e < 2 ? a0 : a1); + } + mma(o[j], a, b[0], b[1]); + mma(o[j + 1], a, b[2], b[3]); + } + } + __syncthreads(); + } + wait<0>(); + // the partials: o [H][LW], m, l of (chunk, row) + const long long base = ((long long)c * R + r) * H; + { + const int g = lane >> 2, tq = lane & 3; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const int col = wn * 128 + j * 8 + 2 * tq; + const int h0 = wm * 16 + g; + *reinterpret_cast(po + (base + h0) * LW + col) = make_float2(o[j][0], o[j][1]); + *reinterpret_cast(po + (base + h0 + 8) * LW + col) = make_float2(o[j][2], o[j][3]); + } + } + __syncthreads(); + if (tid < H) { + pm[base + tid] = sm_m[tid]; + pl[base + tid] = sm_l[tid]; + } +} + +} // namespace + +extern "C" __global__ void __launch_bounds__(THREADS) glm_attn_dense(const __nv_bfloat16* __restrict__ qa, + const unsigned long long* __restrict__ lc, const int* __restrict__ pos, float* __restrict__ po, float* __restrict__ pm, + float* __restrict__ pl, int R, float scale) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + chunks(qa, lc, pos, nullptr, nullptr, po, pm, pl, R, scale); +} + +extern "C" __global__ void __launch_bounds__(THREADS) glm_attn_sparse(const __nv_bfloat16* __restrict__ qa, + const unsigned long long* __restrict__ lc, const int* __restrict__ tok, const int* __restrict__ cnt, float* __restrict__ po, + float* __restrict__ pm, float* __restrict__ pl, int R, float scale) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + chunks(qa, lc, nullptr, tok, cnt, po, pm, pl, R, scale); +} + +// the shared size the chunk kernels take (the host sets it as their dynamic shared memory) +extern "C" __global__ void glm_attn_smem(int* out) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + out[0] = SMEM; } + +// the chunk partials of row r, head h in chunk order -> out [R, H, LW] bf16. dense rows past the limit and sparse +// rows with no tokens are left alone (the other path writes them). grid (R, H), 128 threads. +extern "C" __global__ void glm_attn_merge(const float* __restrict__ po, const float* __restrict__ pm, const float* __restrict__ pl, + __nv_bfloat16* __restrict__ out, const int* __restrict__ pos, const int* __restrict__ cnt, int R, int nch, int sparse) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int r = blockIdx.x, h = blockIdx.y; + if (sparse ? cnt[r] == 0 : pos[r] >= DENSE_LIMIT) return; + float m = -INFINITY, l = 0.0f, acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int c = 0; c < nch; ++c) { + const long long base = ((long long)c * R + r) * H + h; + const float cm = pm[base], cl = pl[base]; + const bool active = cl > 0.0f; + const float nm = active ? fmaxf(m, cm) : m; + const float a = active ? (m == -INFINITY ? 0.0f : expf(m - nm)) : 1.0f; + const float b = active ? expf(cm - nm) : 0.0f; + const float4 co = *reinterpret_cast(po + base * LW + threadIdx.x * 4); + acc[0] = acc[0] * a + co.x * b; + acc[1] = acc[1] * a + co.y * b; + acc[2] = acc[2] * a + co.z * b; + acc[3] = acc[3] * a + co.w * b; + l = l * a + cl * b; + m = nm; + } + __nv_bfloat16* dst = out + ((long long)r * H + h) * LW + threadIdx.x * 4; +#pragma unroll + for (int e = 0; e < 4; ++e) dst[e] = __float2bfloat16_rn(acc[e] / l); +} + +// the window rows' latents into their caches: row r's 512 values at slot pos[r] of lc[r] +extern "C" __global__ void glm_lat_write(const __nv_bfloat16* __restrict__ lat, long long lat_stride, + const unsigned long long* __restrict__ lc, const int* __restrict__ pos) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int r = blockIdx.x; + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(lc[r]) + (long long)pos[r] * LW; + const uint4* src = reinterpret_cast(lat + r * lat_stride); + for (int i = threadIdx.x; i < LW / 8; i += blockDim.x) reinterpret_cast(dst)[i] = src[i]; +} + +// absorb: qa[r, h, n] = sum_k q[r, h, k] wk[h, k, n] (fp32 in k order, then bf16); q [R, H, 256], wk [H, 256, 512]. +// grid (H, LW / 64, ceil(R / 16)), 64 threads: the weight block once in shared memory, each row summed alone. +extern "C" __global__ void glm_absorb_q(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ wk, + __nv_bfloat16* __restrict__ qa, int R) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int D = 256; + __shared__ __nv_bfloat16 w[D][64]; + __shared__ float qs[D]; + const int h = blockIdx.x, n0 = blockIdx.y * 64, n = threadIdx.x; + for (int k = 0; k < D; ++k) w[k][n] = wk[((long long)h * D + k) * LW + n0 + n]; + const int r1 = min(R, (int)(blockIdx.z + 1) * 16); + for (int r = blockIdx.z * 16; r < r1; ++r) { + __syncthreads(); + for (int k = n; k < D; k += 64) qs[k] = __bfloat162float(q[((long long)r * H + h) * D + k]); + __syncthreads(); + float acc = 0.0f; + for (int k = 0; k < D; ++k) acc = fmaf(qs[k], __bfloat162float(w[k][n]), acc); + qa[((long long)r * H + h) * LW + n0 + n] = __float2bfloat16_rn(acc); + } +} + +// expand: out[r, h, n] = sum_k ol[r, h, k] wv[h, n, k]; ol [R, H, 512], wv [H, 256, 512] -> out [R, H, 256]. +// grid (H, 256 / 8, ceil(R / 16)), 256 threads: a warp an output, lanes over k in order, then a xor tree. +extern "C" __global__ void glm_expand_v(const __nv_bfloat16* __restrict__ ol, const __nv_bfloat16* __restrict__ wv, + __nv_bfloat16* __restrict__ out, int R) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int DV = 256; + const int h = blockIdx.x, lane = threadIdx.x & 31, n = blockIdx.y * 8 + (threadIdx.x >> 5); + float w[16]; + const __nv_bfloat16* wr = wv + ((long long)h * DV + n) * LW; +#pragma unroll + for (int i = 0; i < 16; ++i) w[i] = __bfloat162float(wr[i * 32 + lane]); + const int r1 = min(R, (int)(blockIdx.z + 1) * 16); + for (int r = blockIdx.z * 16; r < r1; ++r) { + const __nv_bfloat16* x = ol + ((long long)r * H + h) * LW; + float acc = 0.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) acc = fmaf(w[i], __bfloat162float(x[i * 32 + lane]), acc); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) acc += __shfl_xor_sync(0xffffffffu, acc, off); + if (lane == 0) out[((long long)r * H + h) * DV + n] = __float2bfloat16_rn(acc); + } +} diff --git a/zig/kernels/cuda/glm/mm.cu b/zig/kernels/cuda/glm/mm.cu new file mode 100644 index 0000000000..2afb66d54f --- /dev/null +++ b/zig/kernels/cuda/glm/mm.cu @@ -0,0 +1,299 @@ +// GLM-5.3-Flash on CUDA: the dense linears, ours. One tensor-core kernel for every row count (a decode window's few +// rows, sixteen streams' windows at once, a prompt chunk's 2048): each group of 64 inputs is summed from zero by the +// bf16 mma, then acc = fma(P, scale, acc) in group order, and the K slices (fixed by the shape) are added in slice +// order, so a row's bits never depend on the rows beside it. +// +// Weights are tiled at load (glm_mm_tile*): a column tile of 8 outputs and a group of 64 inputs is 32 lanes' chunks of +// 16 values in mma B-fragment order (lane l: column l / 4; value 4 kt + j at input 16 kt + {2t, 2t+1, 2t+8, 2t+9}[j], +// t = l % 4), each chunk BITS / 2 words (ops.cu's 16-value packing; BITS 16: the bf16 values themselves). Scales are +// fp16 [N / 8][K / 64][8]. + +#include +#include +#include + +#include "../flashnext/qmm_frag.cuh" + +namespace { + +using namespace qmm_frag; + +// 16 codes from a packed chunk (ops.cu's layout), stored + 2^(BITS-1) +template +__device__ __forceinline__ void codes16(const uint32_t* __restrict__ p, int (&q)[16]) { + if constexpr (BITS == 4) { +#pragma unroll + for (int i = 0; i < 16; ++i) q[i] = (int)((p[i / 8] >> (4 * (i % 8))) & 15u) - 8; + } else if constexpr (BITS == 8) { +#pragma unroll + for (int i = 0; i < 16; ++i) q[i] = (int)((p[i / 4] >> (8 * (i % 4))) & 255u) - 128; + } else { +#pragma unroll + for (int i = 0; i < 15; ++i) q[i] = (int)((p[i / 5] >> (6 * (i % 5))) & 63u) - 32; + q[15] = (int)((p[0] >> 30) | ((p[1] >> 30) << 2) | ((p[2] >> 30) << 4)) - 32; + } +} + +template +__device__ __forceinline__ void pack16(const int (&q)[16], uint32_t* __restrict__ out) { + constexpr int half = 1 << (BITS - 1); + uint32_t w[BITS / 2]; +#pragma unroll + for (int i = 0; i < BITS / 2; ++i) w[i] = 0u; + if constexpr (BITS == 4) { +#pragma unroll + for (int i = 0; i < 16; ++i) w[i / 8] |= (uint32_t)(q[i] + half) << (4 * (i % 8)); + } else if constexpr (BITS == 8) { +#pragma unroll + for (int i = 0; i < 16; ++i) w[i / 4] |= (uint32_t)(q[i] + half) << (8 * (i % 4)); + } else { +#pragma unroll + for (int i = 0; i < 15; ++i) w[i / 5] |= (uint32_t)(q[i] + half) << (6 * (i % 5)); + const uint32_t last = (uint32_t)(q[15] + half); +#pragma unroll + for (int j = 0; j < 3; ++j) w[j] |= ((last >> (2 * j)) & 3u) << 30; + } +#pragma unroll + for (int i = 0; i < BITS / 2; ++i) out[i] = w[i]; +} + +// input of value v (0..15) of lane ``lane``'s chunk within its group of 64 +__device__ __forceinline__ int frag_k(int lane, int v) { + const int t = lane & 3, kt = v >> 2, j = v & 3; + return kt * 16 + 2 * t + (j & 1) + (j >> 1) * 8; +} + +// the B-fragment bf16 pairs of a chunk: pairs 2 kt and 2 kt + 1 are step kt's b0 and b1 +template +__device__ __forceinline__ void bpairs(const uint32_t* __restrict__ p, uint32_t (&b)[8]) { + if constexpr (BITS == 16) { +#pragma unroll + for (int i = 0; i < 8; ++i) b[i] = p[i]; + } else { + int q[16]; + codes16(p, q); +#pragma unroll + for (int i = 0; i < 8; ++i) { + const __nv_bfloat162 h = __floats2bfloat162_rn((float)q[2 * i], (float)q[2 * i + 1]); // exact integers + b[i] = *reinterpret_cast(&h); + } + } +} + +template +struct Tile { + static constexpr int WN = 4, BN = 64, NT = 2, MT = BM / WM / 16, THREADS = WM * WN * 32; + static constexpr int WPC = BITS / 2; + static constexpr int ROW = 128, CHUNKS = 8; // a row's 64 inputs: 128 bytes, 8 chunks of 16 + static constexpr int X = BM * ROW; + static constexpr int W = 8 * 32 * WPC * 4; // 8 column tiles' chunks + static constexpr int S = 8 * 8 * 2; // their scales (fp16) + static constexpr int STAGE = X + W + S; + static constexpr int SMEM = STAGES * STAGE; +}; + +template +__device__ __forceinline__ void mm_kernel(const __nv_bfloat16* __restrict__ x, long long ldx, + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, + int f32, int M, int N, int K, int SK) { + using T = Tile; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / T::WN, wn = warp % T::WN; + const int KG = K >> 6, per = KG / SK, N8 = N >> 3; + const int n0 = blockIdx.x * T::BN, m0 = blockIdx.y * BM, slice = blockIdx.z, g0 = slice * per; + const int t0 = n0 >> 3; // the block's first column tile + auto stage = [&](int st) { return buf + st * T::STAGE; }; + auto loadX = [&](int st, int g) { + unsigned char* p = stage(st); + 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 + (long long)row * ldx + g * 64 + ch * 8, m0 + r < M); + } + }; + auto loadW = [&](int st, int g) { + unsigned char* pw = stage(st) + T::X; + constexpr int PER_TILE = 32 * T::WPC * 4 / 16; // 16-byte pieces of one column tile's chunks + for (int c = tid; c < 8 * PER_TILE; c += T::THREADS) { + const int tt = c / PER_TILE, off = c % PER_TILE; + if (t0 + tt >= N8) continue; + const long long src = ((long long)(t0 + tt) * KG + g) * 32 * T::WPC * 4 + off * 16; + cp16(pw + c * 16, reinterpret_cast(w) + src); + } + if constexpr (BITS != 16) { + unsigned char* ps = pw + T::W; + for (int tt = tid; tt < 8; tt += T::THREADS) { + if (t0 + tt >= N8) continue; + cp16(ps + tt * 16, s + ((long long)(t0 + tt) * KG + g) * 8); + } + } + }; + 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; + // dependent launch: the first stages' weights stream in while the kernel before still runs; its outputs (this + // one's input, and any buffer this one writes) only after griddepcontrol.wait + grid_launch(); +#pragma unroll + for (int st = 0; st < STAGES - 1; ++st) + if (st < per) loadW(st, g0 + st); + grid_wait(); +#pragma unroll + for (int st = 0; st < STAGES - 1; ++st) { + if (st < per) loadX(st, g0 + st); + commit(); + } + for (int it = 0; it < per; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) { + loadX(next % STAGES, g0 + next); + loadW(next % STAGES, g0 + next); + } + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __half* ps = reinterpret_cast(p + T::X + T::W); + uint32_t b[T::NT][8]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) bpairs(pw + ((wn * T::NT + j) * 32 + lane) * T::WPC, b[j]); + float d[T::MT][T::NT][4]; +#pragma unroll + for (int kt = 0; kt < 4; ++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) +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[i][j], a[i], b[j][0], b[j][1]); + else mma(d[i][j], a[i], b[j][2 * kt], b[j][2 * kt + 1]); + } + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + float s0 = 1.0f, s1 = 1.0f; + if constexpr (BITS != 16) { + const __half2 h = *reinterpret_cast(ps + (wn * T::NT + j) * 8 + (lane & 3) * 2); + s0 = __low2float(h); + s1 = __high2float(h); + } +#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(d[i][j][e], (e & 1) ? s1 : s0, acc[i][j][e]); + } + } + wait<0>(); +#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 * T::NT + j) * 8 + (lane & 3) * 2; + if (col >= N) continue; +#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) { + float* dst = part + ((long long)slice * M + row) * N + col; + dst[0] = v0; + dst[1] = v1; + } else if (f32) { + float* dst = reinterpret_cast(out) + (long long)row * N + col; + dst[0] = v0; + dst[1] = v1; + } else { + *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<__nv_bfloat16*>(out) + (long long)row * N + col) = + __floats2bfloat162_rn(v0, v1); + } + } + } +} + +} // namespace + +// The K slices' partials added in slice order, then rounded once. +extern "C" __global__ void glm_mm_reduce(const float* __restrict__ part, void* __restrict__ out, long long total, int SK, int f32) { + qmm_frag::grid_launch(); + const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + float acc = part[i]; + for (int sl = 1; sl < SK; ++sl) acc = acc + part[sl * total + i]; + if (f32) reinterpret_cast(out)[i] = acc; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); +} + +// one instance a (bits, row bucket): (16, 1 warp row), (32, 1), (64, 2), (128, 2); three stages +#define GLM_MM_ALL(B) \ + extern "C" __global__ void __launch_bounds__(128) glm_mm##B##_16(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(128) glm_mm##B##_32(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(256) glm_mm##B##_64(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(256) glm_mm##B##_128(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } +GLM_MM_ALL(4) +GLM_MM_ALL(6) +GLM_MM_ALL(8) +GLM_MM_ALL(16) + +// -- tiling at load --------------------------------------------------------------------------------------------------- +// packed rows (ops.cu's [N][K / 16][BITS / 2] words, fp16 scales [N][K / 64]) -> the tiled chunks and scales +template +__device__ __forceinline__ void tile_packed(const uint32_t* __restrict__ src, const __half* __restrict__ ss, + uint32_t* __restrict__ dst, __half* __restrict__ ds, int N, int K) { + constexpr int WPC = BITS / 2; + const int KG = K >> 6, N8 = N >> 3; + const long long c = (long long)blockIdx.x * blockDim.x + threadIdx.x; // a destination chunk: (tile, group, lane) + if (c >= (long long)N8 * KG * 32) return; + const int lane = (int)(c & 31); + const long long tg = c >> 5; + const int g = (int)(tg % KG), tt = (int)(tg / KG); + const int n = tt * 8 + (lane >> 2); + const uint32_t* row = src + (long long)n * (K >> 4) * WPC; + int q[16]; +#pragma unroll + for (int v = 0; v < 16; ++v) { + const int k = g * 64 + frag_k(lane, v); + int sq[16]; + codes16(row + (long long)(k >> 4) * WPC, sq); + q[v] = sq[k & 15]; + } + pack16(q, dst + c * WPC); + if (lane < 8) ds[(tg * 8) + lane] = ss[(long long)(tt * 8 + lane) * KG + g]; +} + +extern "C" __global__ void glm_mm_tile4(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<4>(src, ss, dst, ds, N, K); } +extern "C" __global__ void glm_mm_tile6(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<6>(src, ss, dst, ds, N, K); } +extern "C" __global__ void glm_mm_tile8(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<8>(src, ss, dst, ds, N, K); } + +// bf16 rows [N][K] -> the tiled chunks (16 bf16 a chunk, in pair order) +extern "C" __global__ void glm_mm_tile16(const __nv_bfloat16* __restrict__ src, __nv_bfloat16* __restrict__ dst, int N, int K) { + const int KG = K >> 6, N8 = N >> 3; + const long long c = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (c >= (long long)N8 * KG * 32) return; + const int lane = (int)(c & 31); + const long long tg = c >> 5; + const int g = (int)(tg % KG), tt = (int)(tg / KG); + const int n = tt * 8 + (lane >> 2); +#pragma unroll + for (int v = 0; v < 16; ++v) dst[c * 16 + v] = src[(long long)n * K + g * 64 + frag_k(lane, v)]; +} diff --git a/zig/kernels/cuda/glm/ops.cu b/zig/kernels/cuda/glm/ops.cu new file mode 100644 index 0000000000..14038c496e --- /dev/null +++ b/zig/kernels/cuda/glm/ops.cu @@ -0,0 +1,195 @@ +// GLM-5.3-Flash on CUDA: the small steps the Python engine runs as torch ops (integers and casts, so the same values). + +#include +#include +#include + +// sparse.select_tokens' token list: row r (position q = pos[0] + r) attends to its 512 pools' tokens, ascending +// (pools[r] ascending, int64), then the incomplete last pool's visible tokens (-1 past q); count = 2048 + those when +// the row has more than 512 complete pools, else 0 (the row stays dense). +extern "C" __global__ void glm_sparse_tokens(const int64_t* __restrict__ pools, const int* __restrict__ pos, int rows, + int* __restrict__ tokens, int* __restrict__ counts) { + const int r = blockIdx.x; + if (r >= rows) return; + const int64_t q = (int64_t)pos[0] + r; + const int64_t npool = (q + 1) / 4; + int* out = tokens + (int64_t)r * 2051; + for (int i = threadIdx.x; i < 2048; i += blockDim.x) out[i] = (int)(pools[(int64_t)r * 512 + i / 4] * 4 + i % 4); + if (threadIdx.x < 3) { + const int64_t t = npool * 4 + threadIdx.x; + out[2048 + threadIdx.x] = t <= q ? (int)t : -1; + } + if (threadIdx.x == 0) { + int tail = 0; + for (int j = 0; j < 3; ++j) tail += (npool * 4 + j) <= q; + counts[r] = npool > 512 ? 2048 + tail : 0; + } +} + +// dst[r * dst_stride + c] = bf16(src[r * src_stride + c]) (torch's fp32 -> bf16 copy_: round to nearest even) +extern "C" __global__ void glm_cast_bf16(const float* __restrict__ src, long long src_stride, __nv_bfloat16* __restrict__ dst, + long long dst_stride, int cols) { + const int r = blockIdx.y; + for (int c = blockIdx.x * blockDim.x + threadIdx.x; c < cols; c += gridDim.x * blockDim.x) + dst[r * dst_stride + c] = __float2bfloat16_rn(src[r * src_stride + c]); +} + +// out[i] = base[0] + off + i * step (a block's first-row position on the device: Python's pos_dev + r0) +extern "C" __global__ void glm_positions(const int* __restrict__ base, int off, int step, int n, int* __restrict__ out) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[i] = base[0] + off + i * step; +} + +// rows of ``words`` 32-bit words from src (row pitch spitch words) to dst (dpitch words): torch's strided copy_ +extern "C" __global__ void glm_copy_rows(const uint32_t* __restrict__ src, long long spitch, uint32_t* __restrict__ dst, + long long dpitch, int words) { + const long long r = blockIdx.y; + for (int c = blockIdx.x * blockDim.x + threadIdx.x; c < words; c += gridDim.x * blockDim.x) dst[r * dpitch + c] = src[r * spitch + c]; +} + +// -- packed dense linears (the loader's TF_GLM_DENSE_BITS) ------------------------------------------------------------ +// A row is K / 16 chunks of 16 symmetric BITS-bit values (stored + 2^(BITS-1)), each chunk BITS / 2 32-bit words, with +// an fp16 scale per 64 inputs (4 chunks). 4 bits: value i at bits 4 (i % 8) of word i / 8. 8 bits: byte i % 4 of word +// i / 4. 6 bits: values 0-4, 5-9, 10-14 at bits 6 (i % 5) of words 0, 1, 2; value 15's bits 0-1, 2-3, 4-5 at bits +// 30-31 of words 0, 1, 2. + +template +__device__ __forceinline__ void glm_unpack16(const uint32_t* __restrict__ p, float v[16]) { + if constexpr (BITS == 4) { + const uint32_t a = p[0], b = p[1]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + v[i] = (float)((int)((a >> (4 * i)) & 15u) - 8); + v[8 + i] = (float)((int)((b >> (4 * i)) & 15u) - 8); + } + } else if constexpr (BITS == 8) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t w = p[j]; +#pragma unroll + for (int i = 0; i < 4; ++i) v[4 * j + i] = (float)((int)((w >> (8 * i)) & 255u) - 128); + } + } else { + const uint32_t w0 = p[0], w1 = p[1], w2 = p[2]; +#pragma unroll + for (int i = 0; i < 5; ++i) { + v[i] = (float)((int)((w0 >> (6 * i)) & 63u) - 32); + v[5 + i] = (float)((int)((w1 >> (6 * i)) & 63u) - 32); + v[10 + i] = (float)((int)((w2 >> (6 * i)) & 63u) - 32); + } + v[15] = (float)((int)((w0 >> 30) | ((w1 >> 30) << 2) | ((w2 >> 30) << 4)) - 32); + } +} + +template +__device__ __forceinline__ void glm_load_chunk(const uint32_t* __restrict__ p, uint32_t* w) { + if constexpr (BITS == 4) { + const uint2 a = __ldg(reinterpret_cast(p)); + w[0] = a.x; w[1] = a.y; + } else if constexpr (BITS == 8) { + const uint4 a = __ldg(reinterpret_cast(p)); + w[0] = a.x; w[1] = a.y; w[2] = a.z; w[3] = a.w; + } else { + w[0] = __ldg(p); w[1] = __ldg(p + 1); w[2] = __ldg(p + 2); + } +} + +// out[m, n] = sum_k x[m, k] q[n, k] s[n, k / 64] for M <= 16 rows (x bf16, row stride xs elements, 16-byte aligned; +// out [M, N] bf16 or fp32): a warp a row, lanes over the chunks, each chunk's 16 products summed in fp32 then scaled, +// the lanes' sums reduced by xor shuffles. A row's sum depends only on its own inputs, never on M. +template +__device__ __forceinline__ void glm_qgemv(const __nv_bfloat16* __restrict__ x, long long xs, const uint32_t* __restrict__ w, + const __half* __restrict__ s, void* __restrict__ out, int f32, int M, int N, int K) { + constexpr int WPC = BITS / 2; + const int lane = threadIdx.x & 31; + const int n = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); + if (n >= N) return; + const int chunks = K >> 4; + const uint32_t* wr = w + (long long)n * chunks * WPC; + const __half* sr = s + (long long)n * (K >> 6); + float acc[16]; +#pragma unroll + for (int m = 0; m < 16; ++m) acc[m] = 0.f; + for (int c = lane; c < chunks; c += 32) { + uint32_t pw[WPC]; + glm_load_chunk(wr + (long long)c * WPC, pw); + float v[16]; + glm_unpack16(pw, v); + const float sc = __half2float(sr[c >> 2]); +#pragma unroll + for (int m = 0; m < 16; ++m) { + if (m < M) { + const uint4* xp = reinterpret_cast(x + m * xs + (c << 4)); + const uint4 a = __ldg(xp), b = __ldg(xp + 1); + const __nv_bfloat162* ha = reinterpret_cast(&a); + const __nv_bfloat162* hb = reinterpret_cast(&b); + float p = 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 f = __bfloat1622float2(ha[i]); + p = fmaf(v[2 * i], f.x, p); + p = fmaf(v[2 * i + 1], f.y, p); + } +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 f = __bfloat1622float2(hb[i]); + p = fmaf(v[8 + 2 * i], f.x, p); + p = fmaf(v[8 + 2 * i + 1], f.y, p); + } + acc[m] = fmaf(p, sc, acc[m]); + } + } + } +#pragma unroll + for (int m = 0; m < 16; ++m) { + if (m < M) { + float a = acc[m]; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) a += __shfl_xor_sync(0xffffffffu, a, off); + if (lane == 0) { + if (f32) reinterpret_cast(out)[(long long)m * N + n] = a; + else reinterpret_cast<__nv_bfloat16*>(out)[(long long)m * N + n] = __float2bfloat16_rn(a); + } + } + } +} + +// the rows as bf16 [N, K] (rounded to nearest even from value * scale), for prompt-sized calls on the bf16 matmul +template +__device__ __forceinline__ void glm_qdequant(const uint32_t* __restrict__ w, const __half* __restrict__ s, + __nv_bfloat16* __restrict__ out, long long chunks) { + const long long c = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (c >= chunks) return; + uint32_t pw[BITS / 2]; + glm_load_chunk(w + c * (BITS / 2), pw); + float v[16]; + glm_unpack16(pw, v); + const float sc = __half2float(s[c >> 2]); + __align__(16) __nv_bfloat16 o[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) o[i] = __float2bfloat16_rn(v[i] * sc); + uint4* dst = reinterpret_cast(out + c * 16); + dst[0] = reinterpret_cast(o)[0]; + dst[1] = reinterpret_cast(o)[1]; +} + +#define GLM_Q(B) \ + extern "C" __global__ void __launch_bounds__(256) glm_qgemv##B(const __nv_bfloat16* __restrict__ x, long long xs, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, int f32, int M, int N, \ + int K) { glm_qgemv(x, xs, w, s, out, f32, M, N, K); } \ + extern "C" __global__ void glm_qdequant##B(const uint32_t* __restrict__ w, const __half* __restrict__ s, \ + __nv_bfloat16* __restrict__ out, long long chunks) { glm_qdequant(w, s, out, chunks); } +GLM_Q(4) +GLM_Q(6) +GLM_Q(8) + +// x rows 0 .. rows - 1 (each ``copies`` streams of 4096 bf16) take feats rows 0 .. rows - 1 [rows, 4096] in every stream: +// an image's features in place of its placeholder tokens' embeddings +extern "C" __global__ void glm_inject_rows(const __nv_bfloat16* __restrict__ feats, __nv_bfloat16* __restrict__ x, int rows, + int copies) { + const long long n = (long long)rows * copies * 4096; + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) { + const long long r = i / (copies * 4096), c = i % 4096; + x[i] = feats[r * 4096 + c]; + } +} diff --git a/zig/kernels/cuda/glm/rows.cu b/zig/kernels/cuda/glm/rows.cu new file mode 100644 index 0000000000..1a9b2e1d59 --- /dev/null +++ b/zig/kernels/cuda/glm/rows.cu @@ -0,0 +1,426 @@ +// GLM-5.3-Flash on CUDA: the row steps between the linears, ours (embedding, RMSNorm, the hyper-connections, the +// residuals, SwiGLU, the router and its top-8, the experts' combine, the KDA conv window). Each row is one block's, +// with fp32 sums in a fixed order (threads strided over the row, then a fixed tree), so a row's bits depend on its own +// inputs alone, never on how many rows the window holds. + +#include +#include + +namespace { + +__device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } + +// the block's sum of v (blockDim.x a multiple of 32, at most 1024): xor tree in each warp, then the warps in order +__device__ __forceinline__ float block_sum(float v, float* red) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) v += __shfl_xor_sync(0xffffffffu, v, off); + const int warp = threadIdx.x >> 5, nw = blockDim.x >> 5; + __syncthreads(); + if ((threadIdx.x & 31) == 0) red[warp] = v; + __syncthreads(); + float s = red[0]; + for (int i = 1; i < nw; ++i) s += red[i]; + return s; +} + +__device__ __forceinline__ float sigmoid(float x) { return 1.0f / (1.0f + expf(-x)); } + +} // namespace + +// out[r, c * D + d] = table[ids[r], d] for each of ``copies`` streams (bf16 rows, D a multiple of 8) +extern "C" __global__ void glm_embed(const int* __restrict__ ids, const __nv_bfloat16* __restrict__ table, + __nv_bfloat16* __restrict__ out, int D, int copies) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int r = blockIdx.x; + const uint4* src = reinterpret_cast(table + (long long)ids[r] * D); + uint4* dst = reinterpret_cast(out + (long long)r * copies * D); + for (int i = threadIdx.x; i < D / 8; i += blockDim.x) { + const uint4 v = src[i]; + for (int c = 0; c < copies; ++c) dst[c * (D / 8) + i] = v; + } +} + +// RMSNorm (Llama: bf16(w * bf16(x * rinv))) of row r of x (stride xs_) -> out (stride os_), and the 64-groups' fp32 +// sums of the outputs when xs is given. D <= 4096. +extern "C" __global__ void glm_rmsnorm(const __nv_bfloat16* __restrict__ x, long long x_stride, const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ out, long long o_stride, float* __restrict__ xs, float eps, int D) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + __shared__ float red[32]; + __shared__ float ys[4096]; + const int r = blockIdx.x; + const __nv_bfloat16* xr = x + r * x_stride; + float ss = 0.0f; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float v = __bfloat162float(xr[d]); + ss = fmaf(v, v, ss); + } + const float rinv = 1.0f / sqrtf(block_sum(ss, red) / (float)D + eps); + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float y = bf(__bfloat162float(w[d]) * bf(__bfloat162float(xr[d]) * rinv)); + out[r * o_stride + d] = __float2bfloat16_rn(y); + ys[d] = y; + } + if (xs == nullptr) return; + __syncthreads(); + for (int g = threadIdx.x; g < D / 64; g += blockDim.x) { + float s = 0.0f; + for (int i = 0; i < 64; ++i) s += ys[g * 64 + i]; + xs[(long long)r * (D / 64) + g] = s; + } +} + +// hc_pre (hyper-connections, 4 streams of D = 4096), in two launches. The partials: row r's 24 mixing dots of the +// streams [16384] with fn [24, 16384] and their sum of squares over 16 fixed slices of 1024 inputs, one block a +// (row, slice) -> part [R, 16, 25]. 128 threads. +extern "C" __global__ void __launch_bounds__(128) glm_hc_part(const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ fn, + float* __restrict__ part) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int D = 4096, WIDE = 4 * D, M = 24, SL = 1024; + __shared__ float red[32]; + const int r = blockIdx.x, sl = blockIdx.y, tid = threadIdx.x; + const int k = sl * SL + tid * 8; // each thread 8 consecutive inputs of the slice + const uint4 xv = *reinterpret_cast(x + (long long)r * WIDE + k); + const __nv_bfloat162* xh = reinterpret_cast(&xv); + float xf[8]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __bfloat1622float2(xh[j]); + xf[2 * j] = f.x; + xf[2 * j + 1] = f.y; + } + float ss = 0.0f; +#pragma unroll + for (int j = 0; j < 8; ++j) ss = fmaf(xf[j], xf[j], ss); + float* out = part + ((long long)r * 16 + sl) * (M + 1); +#pragma unroll 4 + for (int m = 0; m < M; ++m) { + const uint4 wv = *reinterpret_cast(fn + (long long)m * WIDE + k); + const __nv_bfloat162* wh = reinterpret_cast(&wv); + float p = 0.0f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __bfloat1622float2(wh[j]); + p = fmaf(f.x, xf[2 * j], p); + p = fmaf(f.y, xf[2 * j + 1], p); + } + const float sm = block_sum(p, red); + if (tid == 0) out[m] = sm; + } + const float sm = block_sum(ss, red); + if (tid == 0) out[M] = sm; +} + +// The finish: the slices' partials in slice order, the pre / post / comb weights (comb: ``iters`` Sinkhorn steps), the +// collapsed row RMSNormed by nw -> out [D] and its 64-groups' sums xs; post [4] and comb [4, 4] kept for hc_post. +// 512 threads. +extern "C" __global__ void __launch_bounds__(512) glm_hc_finish(const __nv_bfloat16* __restrict__ x, const float* __restrict__ part, + const float* __restrict__ base, const float* __restrict__ scale, const __nv_bfloat16* __restrict__ nw, + __nv_bfloat16* __restrict__ out, float* __restrict__ xs, float* __restrict__ post, float* __restrict__ comb, + float eps, float hc_eps, int iters) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int D = 4096, WIDE = 4 * D, M = 24; + __shared__ float red[32]; + __shared__ float pre[4]; + __shared__ float ys[D]; + const int r = blockIdx.x, tid = threadIdx.x; + const __nv_bfloat16* xr = x + (long long)r * WIDE; + __shared__ float mixs[M + 1]; + if (tid <= M) { // each sum: the slices' partials in slice order + float a = 0.0f; + for (int sl = 0; sl < 16; ++sl) a += part[((long long)r * 16 + sl) * (M + 1) + tid]; + mixs[tid] = a; + } + __syncthreads(); + if (tid < 32) { // warp 0: lane i * 4 + j holds comb[i][j]; lanes 16 .. 23 the pre and post weights + const int lane = tid; + const float rinv = 1.0f / sqrtf(mixs[M] / (float)WIDE + eps); + const float s_pre = scale[0], s_post = scale[1], s_comb = scale[2]; + if (lane >= 16 && lane < 20) pre[lane - 16] = sigmoid(mixs[lane - 16] * rinv * s_pre + base[lane - 16]) + hc_eps; + if (lane >= 20 && lane < 24) post[r * 4 + lane - 20] = 2.0f * sigmoid(mixs[4 + lane - 20] * rinv * s_post + base[4 + lane - 20]); + const int q = lane & 15; + const float l = mixs[8 + q] * rinv * s_comb + base[8 + q]; + // row softmax over j (lanes of one i: xor 1, 2), then Sinkhorn: columns (xor 4, 8) and rows in turn + float mx = fmaxf(l, __shfl_xor_sync(0xffffffffu, l, 1)); + mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 2)); + const float ex = expf(l - mx); + float tot = ex + __shfl_xor_sync(0xffffffffu, ex, 1); + tot = tot + __shfl_xor_sync(0xffffffffu, tot, 2); + float c = ex / tot + hc_eps; + auto col = [&](float v) { + float t = v + __shfl_xor_sync(0xffffffffu, v, 4); + t = t + __shfl_xor_sync(0xffffffffu, t, 8); + return v / (t + hc_eps); + }; + auto row = [&](float v) { + float t = v + __shfl_xor_sync(0xffffffffu, v, 1); + t = t + __shfl_xor_sync(0xffffffffu, t, 2); + return v / (t + hc_eps); + }; + c = col(c); + for (int it = 0; it < iters - 1; ++it) c = col(row(c)); + if (lane < 16) comb[r * 16 + lane] = c; + } + __syncthreads(); + const float p0 = pre[0], p1 = pre[1], p2 = pre[2], p3 = pre[3]; + float cs = 0.0f; + for (int d = tid; d < D; d += 512) { + const float c = bf(fmaf(p3, __bfloat162float(xr[3 * D + d]), fmaf(p2, __bfloat162float(xr[2 * D + d]), + fmaf(p1, __bfloat162float(xr[D + d]), p0 * __bfloat162float(xr[d]))))); + ys[d] = c; + cs = fmaf(c, c, cs); + } + const float rinv2 = 1.0f / sqrtf(block_sum(cs, red) / (float)D + eps); + for (int d = tid; d < D; d += 512) { + const float y = bf(__bfloat162float(nw[d]) * bf(ys[d] * rinv2)); + out[(long long)r * D + d] = __float2bfloat16_rn(y); + ys[d] = y; + } + __syncthreads(); + for (int g = tid; g < D / 64; g += 512) { + float s = 0.0f; + for (int i = 0; i < 64; ++i) s += ys[g * 64 + i]; + xs[(long long)r * (D / 64) + g] = s; + } +} + +// hc_post: the branch (both ranks' fp32 partials g [2, R, D], rank 0 first, rounded once) into each of the 4 streams: +// x_s = bf16(post_s * branch + sum_i comb[i][s] x_i), x in place. grid (R, D / 1024), 256 threads. +extern "C" __global__ void glm_hc_post(__nv_bfloat16* __restrict__ x, const float* __restrict__ g, const float* __restrict__ post, + const float* __restrict__ comb, int R) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int D = 4096; + const int r = blockIdx.x; + const float* cb = comb + r * 16; + for (int d = blockIdx.y * 1024 + threadIdx.x; d < (blockIdx.y + 1) * 1024; d += blockDim.x) { + const float branch = bf(g[(long long)r * D + d] + g[(long long)R * D + (long long)r * D + d]); + __nv_bfloat16* xr = x + (long long)r * 4 * D; + const float x0 = __bfloat162float(xr[d]), x1 = __bfloat162float(xr[D + d]), x2 = __bfloat162float(xr[2 * D + d]), + x3 = __bfloat162float(xr[3 * D + d]); +#pragma unroll + for (int s = 0; s < 4; ++s) { + const float mixed = fmaf(cb[12 + s], x3, fmaf(cb[8 + s], x2, fmaf(cb[4 + s], x1, cb[s] * x0))); + xr[s * D + d] = __float2bfloat16_rn(fmaf(post[r * 4 + s], branch, mixed)); + } + } +} + +// the plain residual (the MTP block): x = bf16(x + bf16(g rank 0 + g rank 1)), x [R, D] in place +extern "C" __global__ void glm_residual_add(__nv_bfloat16* __restrict__ x, const float* __restrict__ g, int R) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int D = 4096; + const int r = blockIdx.x; + for (int d = blockIdx.y * 1024 + threadIdx.x; d < (blockIdx.y + 1) * 1024; d += blockDim.x) { + const long long i = (long long)r * D + d; + x[i] = __float2bfloat16_rn(__bfloat162float(x[i]) + bf(g[i] + g[(long long)R * D + i])); + } +} + +// the final collapse: bf16((x0 + x1 + x2 + x3) / 4) of the 4 streams +extern "C" __global__ void glm_stream_mean(const __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ out) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int D = 4096; + const int r = blockIdx.x; + for (int d = blockIdx.y * 1024 + threadIdx.x; d < (blockIdx.y + 1) * 1024; d += blockDim.x) { + const __nv_bfloat16* xr = x + (long long)r * 4 * D; + const float a = ((__bfloat162float(xr[d]) + __bfloat162float(xr[D + d])) + __bfloat162float(xr[2 * D + d])) + + __bfloat162float(xr[3 * D + d]); + out[(long long)r * D + d] = __float2bfloat16_rn(a / 4.0f); + } +} + +// SwiGLU over [gate | up] (bf16, width W each): bf16(bf16(silu(min(g, L))) * clip(u, -L, L)) and its 64-groups' sums. +// grid (R, W / 512), 256 threads. +extern "C" __global__ void glm_swiglu(const __nv_bfloat16* __restrict__ gu, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, + float limit, int W) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + __shared__ float ys[512]; + const int r = blockIdx.x, c0 = blockIdx.y * 512; + for (int i = threadIdx.x; i < 512; i += blockDim.x) { + const int d = c0 + i; + const float g = fminf(__bfloat162float(gu[(long long)r * 2 * W + d]), limit); + const float u = fminf(fmaxf(__bfloat162float(gu[(long long)r * 2 * W + W + d]), -limit), limit); + const float a = bf(bf(g / (1.0f + expf(-g))) * u); + out[(long long)r * W + d] = __float2bfloat16_rn(a); + ys[i] = a; + } + __syncthreads(); + if (threadIdx.x < 8) { + float s = 0.0f; + for (int i = 0; i < 64; ++i) s += ys[threadIdx.x * 64 + i]; + xs[(long long)r * (W / 64) + blockIdx.y * 8 + threadIdx.x] = s; + } +} + +// router logits: out[m, e] = fp32 x[m] . w[e] (bf16 [NE, D]). A block takes 2 experts, a warp a quarter of D (D a +// multiple of 1024): its lanes over 8-input pieces in order, a xor tree, then the quarters summed in order: each row's +// sum is its own. grid (NE / 2, ceil(M / 16)), 256 threads. +extern "C" __global__ void __launch_bounds__(256) glm_router(const __nv_bfloat16* __restrict__ x, long long x_stride, + const __nv_bfloat16* __restrict__ w, float* __restrict__ out, int M, int NE, int D) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + __shared__ float part[8][16]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int e = blockIdx.x * 2 + (warp >> 2), q = warp & 3; + const int m0 = blockIdx.y * 16; + const int len = D >> 2; + float acc[16]; +#pragma unroll + for (int m = 0; m < 16; ++m) acc[m] = 0.0f; + if (e < NE) { + const __nv_bfloat16* wr = w + (long long)e * D + q * len; + const __nv_bfloat16* xq = x + q * len; + for (int k = lane * 8; k < len; k += 256) { + const uint4 wv = *reinterpret_cast(wr + k); + const __nv_bfloat162* wh = reinterpret_cast(&wv); + float wf[8]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __bfloat1622float2(wh[j]); + wf[2 * j] = f.x; + wf[2 * j + 1] = f.y; + } +#pragma unroll + for (int m = 0; m < 16; ++m) { + if (m0 + m < M) { + const uint4 xv = *reinterpret_cast(xq + (long long)(m0 + m) * x_stride + k); + const __nv_bfloat162* xh = reinterpret_cast(&xv); + float p = acc[m]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __bfloat1622float2(xh[j]); + p = fmaf(wf[2 * j], f.x, p); + p = fmaf(wf[2 * j + 1], f.y, p); + } + acc[m] = p; + } + } + } + } +#pragma unroll + for (int m = 0; m < 16; ++m) { + float a = acc[m]; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) a += __shfl_xor_sync(0xffffffffu, a, off); + if (lane == 0) part[warp][m] = a; + } + __syncthreads(); + if (threadIdx.x < 32) { + const int m = threadIdx.x & 15, h = threadIdx.x >> 4, ee = blockIdx.x * 2 + h; + if (ee < NE && m0 + m < M) + out[(long long)(m0 + m) * NE + ee] = ((part[4 * h][m] + part[4 * h + 1][m]) + part[4 * h + 2][m]) + part[4 * h + 3][m]; + } +} + +// each row's 8 experts by sigmoid score plus bias (the lower id on ties), weights the picks' scores normalized in pick +// order times ``scale``, then the shared expert (id NE, weight 1) in slot 8. A warp a row (NE <= 320). +extern "C" __global__ void glm_topk(const float* __restrict__ logits, const float* __restrict__ bias, int* __restrict__ pick, + float* __restrict__ wts, float scale, int NE) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int PER = 10, TOPK = 8, SLOTS = 9; + const int r = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); + const int lane = threadIdx.x & 31; + float score[PER], choice[PER]; +#pragma unroll + for (int i = 0; i < PER; ++i) { + const int e = i * 32 + lane; + if (e < NE) { + score[i] = sigmoid(logits[(long long)r * NE + e]); + choice[i] = score[i] + bias[e]; + } else { + score[i] = 0.0f; + choice[i] = -INFINITY; + } + } + float total = 0.0f, wk[TOPK]; + int ids[TOPK]; +#pragma unroll + for (int k = 0; k < TOPK; ++k) { + float best = -INFINITY; + int at = 0x7fffffff; +#pragma unroll + for (int i = 0; i < PER; ++i) + if (choice[i] > best) { + best = choice[i]; + at = i * 32 + lane; + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ob = __shfl_xor_sync(0xffffffffu, best, off); + const int oa = __shfl_xor_sync(0xffffffffu, at, off); + if (ob > best || (ob == best && oa < at)) { + best = ob; + at = oa; + } + } + float sk = 0.0f; +#pragma unroll + for (int i = 0; i < PER; ++i) + if (i * 32 + lane == at) { + sk = score[i]; + choice[i] = -INFINITY; + } + sk = __shfl_sync(0xffffffffu, sk, at & 31); + ids[k] = at; + wk[k] = sk; + total += sk; + } + if (lane == 0) { + for (int k = 0; k < TOPK; ++k) { + pick[r * SLOTS + k] = ids[k]; + wts[r * SLOTS + k] = wk[k] / (total + 1e-20f) * scale; + } + pick[r * SLOTS + TOPK] = NE; + wts[r * SLOTS + TOPK] = 1.0f; + } +} + +// a rank's fp32 MoE share: sum over the 9 slots in order of w_k y_k (y bf16 in a prompt, fp32 in a window) +extern "C" __global__ void glm_combine(const void* __restrict__ y, int y_f32, const float* __restrict__ wts, float* __restrict__ out) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + constexpr int D = 4096, SLOTS = 9; + const int r = blockIdx.x; + for (int d = blockIdx.y * 1024 + threadIdx.x; d < (blockIdx.y + 1) * 1024; d += blockDim.x) { + float acc = 0.0f; + for (int k = 0; k < SLOTS; ++k) { + const long long i = ((long long)r * SLOTS + k) * D + d; + const float v = y_f32 ? reinterpret_cast(y)[i] : __bfloat162float(reinterpret_cast(y)[i]); + acc = fmaf(v, wts[r * SLOTS + k], acc); + } + out[(long long)r * D + d] = acc; + } +} + +// the KDA conv window: conv [layers, 3, C] (in place) takes rows keep .. keep + 2 of [conv; proj rows] +// (proj [layers, R, >= C], strides in elements). grid (layers, ceil(C / 1024)), 256 threads. +extern "C" __global__ void glm_conv_shift(__nv_bfloat16* __restrict__ conv, const __nv_bfloat16* __restrict__ proj, int keep, + long long conv_layer, long long proj_layer, long long proj_row, int C) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const long long l = blockIdx.x; + for (int c = blockIdx.y * 1024 + threadIdx.x; c < min(C, (int)(blockIdx.y + 1) * 1024); c += blockDim.x) { + __nv_bfloat16 v[3]; +#pragma unroll + for (int t = 0; t < 3; ++t) { + const int src = keep + t; + v[t] = src < 3 ? conv[l * conv_layer + src * C + c] : proj[l * proj_layer + (src - 3) * proj_row + c]; + } +#pragma unroll + for (int t = 0; t < 3; ++t) conv[l * conv_layer + t * C + c] = v[t]; + } +} + +// rows whose position is 0 zeroed (the MTP head's first row reads no token before it): x [R, D] bf16 +extern "C" __global__ void glm_zero_first(__nv_bfloat16* __restrict__ x, const int* __restrict__ pos, int D) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int r = blockIdx.x; + if (pos[r] != 0) return; + for (int d = threadIdx.x; d < D; d += blockDim.x) x[(long long)r * D + d] = __float2bfloat16_rn(0.0f); +} + +// dst row i = src row idx[i] (D bf16 each): each stream's last row of a batched step, for its head +extern "C" __global__ void glm_gather_rows(const __nv_bfloat16* __restrict__ src, const int* __restrict__ idx, + __nv_bfloat16* __restrict__ dst, int D) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights + const int i = blockIdx.x; + const uint4* s = reinterpret_cast(src + (long long)idx[i] * D); + uint4* o = reinterpret_cast(dst + (long long)i * D); + for (int d = threadIdx.x; d < D / 8; d += blockDim.x) o[d] = s[d]; +} diff --git a/zig/kernels/cuda/glm/vision.cu b/zig/kernels/cuda/glm/vision.cu new file mode 100644 index 0000000000..b5d5c44a3e --- /dev/null +++ b/zig/kernels/cuda/glm/vision.cu @@ -0,0 +1,176 @@ +// GLM-5.3-Flash's vision tower on CUDA (families/glm_cuda/vision.zig): the elementwise and row steps around its +// cuBLASLt products, as mlx-vlm's glm5_next/vision.py defines them. The residual stream stays in fp32; a step's bf16 +// output is rounded once, to nearest even. + +#include +#include + +namespace { +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ __nv_bfloat16 tobf(float x) { return __float2bfloat16_rn(x); } + +// a block's sum (blockDim.x a multiple of 32, at most 1024); every thread gets it +__device__ float block_sum(float v) { + __shared__ float part[32]; + __shared__ float total; + for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + __syncthreads(); + if (lane == 0) part[warp] = v; + __syncthreads(); + if (warp == 0) { + float s = lane < (int)(blockDim.x >> 5) ? part[lane] : 0.f; + for (int o = 16; o > 0; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o); + if (lane == 0) total = s; + } + __syncthreads(); + return total; +} + +__device__ __forceinline__ float limited_swiglu(float g, float u, float limit) { + g = fminf(g, limit); + u = fminf(fmaxf(u, -limit), limit); + return g / (1.f + expf(-g)) * u; +} +} // namespace + +// dst[i] = bf16(src[i]) (the processor's fp32 patches) +extern "C" __global__ void glm_vis_to_bf16(const float* __restrict__ src, __nv_bfloat16* __restrict__ dst, long long n) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + dst[i] = tobf(src[i]); +} + +// out[r, c] = bf16(acc[r, c] + bias[c]) (bias may be null); with x32: x32[r, c] += that bf16 value (a residual add) +extern "C" __global__ void glm_vis_bias(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, float* __restrict__ x32, long long n, int cols) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) { + const float v = acc[i] + (bias ? bf(bias[i % cols]) : 0.f); + const __nv_bfloat16 h = tobf(v); + if (out) out[i] = h; + if (x32) x32[i] += bf(h); + } +} + +// x32[i] = acc[i] + bias[i % cols] (the patch embedding: the residual stream's first value, in fp32) +extern "C" __global__ void glm_vis_embed32(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ bias, + float* __restrict__ x32, long long n, int cols) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + x32[i] = bf(tobf(acc[i] + bf(bias[i % cols]))); +} + +// RMSNorm of fp32 rows: y = bf16(x rsqrt(mean(x^2) + eps) w), a block a row +extern "C" __global__ void glm_vis_rmsnorm(const float* __restrict__ x, const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ y, int cols, float eps) { + const float* row = x + (long long)blockIdx.x * cols; + float s = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) s += row[c] * row[c]; + const float inv = rsqrtf(block_sum(s) / cols + eps); + for (int c = threadIdx.x; c < cols; c += blockDim.x) y[(long long)blockIdx.x * cols + c] = tobf(row[c] * inv * bf(w[c])); +} + +// The attention inputs: qkv [N, 3, H, 64] bf16 -> q and k each RMSNorm'd over its head (q_norm, k_norm; bf16), then +// rotated (fp32: x cos + rotate_half(x) sin, angles (row, col) x inv_freq[16] for dims 0-15 / 16-31, again for +// 32-63) into qh, kh [N / seg, H, seg, 64]; v into vt [N / seg, H, 64, seg] (each frame group's head contiguous). +// A warp a (row, head): lane l holds dims l and l + 32. +extern "C" __global__ void glm_vis_qkv(const __nv_bfloat16* __restrict__ qkv, const __nv_bfloat16* __restrict__ qn, + const __nv_bfloat16* __restrict__ kn, const int* __restrict__ pos, + const float* __restrict__ inv_freq, __nv_bfloat16* __restrict__ qh, + __nv_bfloat16* __restrict__ kh, __nv_bfloat16* __restrict__ vt, int N, int seg, + int H, float eps) { + const long long item = blockIdx.x * (long long)(blockDim.x >> 5) + (threadIdx.x >> 5); + if (item >= (long long)N * H) return; + const int lane = threadIdx.x & 31; + const int n = (int)(item / H), h = (int)(item % H); + const int g = n / seg, j = n % seg; + const long long blk = (long long)g * H + h; + const __nv_bfloat16* base = qkv + (long long)n * 3 * H * 64 + h * 64; + const int p = lane < 16 ? pos[2 * n] : pos[2 * n + 1]; + const float ang = (float)p * inv_freq[lane & 15]; + const float c = cosf(ang), s = sinf(ang); + for (int which = 0; which < 2; ++which) { + const __nv_bfloat16* src = base + which * H * 64; + const __nv_bfloat16* wn = which ? kn : qn; + float a = bf(src[lane]), b = bf(src[lane + 32]); + float ss = a * a + b * b; + for (int o = 16; o > 0; o >>= 1) ss += __shfl_xor_sync(0xffffffffu, ss, o); + const float inv = rsqrtf(ss / 64.f + eps); + a = bf(tobf(a * inv * bf(wn[lane]))); + b = bf(tobf(b * inv * bf(wn[lane + 32]))); + // rotate_half: dims l < 32 take -x[l + 32], dims l + 32 take x[l] + const float ra = a * c - b * s, rb = b * c + a * s; + __nv_bfloat16* dst = (which ? kh : qh) + (blk * seg + j) * 64; + dst[lane] = tobf(ra); + dst[lane + 32] = tobf(rb); + } + const __nv_bfloat16* v = base + 2 * H * 64; + vt[(blk * 64 + lane) * seg + j] = v[lane]; + vt[(blk * 64 + lane + 32) * seg + j] = v[lane + 32]; +} + +// A scores row's softmax: pr = bf16(exp(scale (s - max))) and its sum (of the bf16 values), a block a row +extern "C" __global__ void glm_vis_softmax(const float* __restrict__ sc, __nv_bfloat16* __restrict__ pr, + float* __restrict__ sums, int n, float scale) { + const float* row = sc + (long long)blockIdx.x * n; + __shared__ float mx_s[32]; + __shared__ float mx; + float m = -INFINITY; + for (int c = threadIdx.x; c < n; c += blockDim.x) m = fmaxf(m, row[c]); + for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); + if ((threadIdx.x & 31) == 0) mx_s[threadIdx.x >> 5] = m; + __syncthreads(); + if (threadIdx.x < 32) { + float v = threadIdx.x < (blockDim.x >> 5) ? mx_s[threadIdx.x] : -INFINITY; + for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if (threadIdx.x == 0) mx = v; + } + __syncthreads(); + float s = 0.f; + for (int c = threadIdx.x; c < n; c += blockDim.x) { + const __nv_bfloat16 e = tobf(expf((row[c] - mx) * scale)); + pr[(long long)blockIdx.x * n + c] = e; + s += bf(e); + } + s = block_sum(s); + if (threadIdx.x == 0) sums[blockIdx.x] = s; +} + +// att[r, h * 64 + d] = bf16(oh[r, d] / sums[r]) (a head's values product, normalized, into its columns) +extern "C" __global__ void glm_vis_head_out(const float* __restrict__ oh, const float* __restrict__ sums, + __nv_bfloat16* __restrict__ att, int rows, int H, int h) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= rows * 64) return; + const int r = i / 64, d = i % 64; + att[(long long)r * H * 64 + h * 64 + d] = tobf(oh[i] / sums[r]); +} + +// out = bf16(limited_swiglu(bf16(g + gb), bf16(u + ub))) from the two fp32 products (biases may be null) +extern "C" __global__ void glm_vis_swiglu(const float* __restrict__ g, const float* __restrict__ u, + const __nv_bfloat16* __restrict__ gb, const __nv_bfloat16* __restrict__ ub, + __nv_bfloat16* __restrict__ out, long long n, int cols, float limit) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) { + const int c = (int)(i % cols); + const float gv = bf(tobf(g[i] + (gb ? bf(gb[c]) : 0.f))); + const float uv = bf(tobf(u[i] + (ub ? bf(ub[c]) : 0.f))); + out[i] = tobf(limited_swiglu(gv, uv, limit)); + } +} + +// The merger's norm: y = bf16(gelu_erf(LayerNorm(bf16(acc)) w + b)), a block a row of ``cols`` +extern "C" __global__ void glm_vis_ln_gelu(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ b, __nv_bfloat16* __restrict__ y, int cols, + float eps) { + const float* row = acc + (long long)blockIdx.x * cols; + float s = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) s += bf(tobf(row[c])); + const float mean = block_sum(s) / cols; + float v = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + const float d = bf(tobf(row[c])) - mean; + v += d * d; + } + const float inv = rsqrtf(block_sum(v) / cols + eps); + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + const float x = bf(tobf((bf(tobf(row[c])) - mean) * inv * bf(w[c]) + bf(b[c]))); + y[(long long)blockIdx.x * cols + c] = tobf(0.5f * x * (1.f + erff(x * 0.70710678118654752f))); + } +} diff --git a/zig/src/cuda/cublaslt.zig b/zig/src/cuda/cublaslt.zig index cc733b8537..cdf3419551 100644 --- a/zig/src/cuda/cublaslt.zig +++ b/zig/src/cuda/cublaslt.zig @@ -108,6 +108,11 @@ pub const Linear = struct { /// Picks the heuristic's first algorithm for this shape once; `workspace_limit` bounds what it may ask for. pub fn init(lt: *const Library, m: u64, n: u64, k: u64, out: DataType, workspace_limit: usize) Error!Linear { + return initTyped(lt, m, n, k, .bf16, out, workspace_limit); + } + + /// ``init`` with ``in`` the type of X and W (bf16, or f32 for an fp32 product). + pub fn initTyped(lt: *const Library, m: u64, n: u64, k: u64, in: DataType, out: DataType, workspace_limit: usize) Error!Linear { var self: Linear = undefined; self.lt = lt; try lt.check(lt.api.cublasLtCreate(&self.handle), "cublasLtCreate"); @@ -119,9 +124,9 @@ pub const Linear = struct { try lt.check(lt.api.cublasLtMatmulDescSetAttribute(self.desc, desc_transa, &ta, @sizeOf(c_int)), "set transa"); try lt.check(lt.api.cublasLtMatmulDescSetAttribute(self.desc, desc_transb, &tb, @sizeOf(c_int)), "set transb"); // column-major view: D^T[n,m] = W[n,k] (stored k x n, transposed) . X^T[k,m] - try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.a, .bf16, k, n, @intCast(k)), "layout W"); + try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.a, in, k, n, @intCast(k)), "layout W"); errdefer _ = lt.api.cublasLtMatrixLayoutDestroy(self.a); - try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.b, .bf16, k, m, @intCast(k)), "layout X"); + try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.b, in, k, m, @intCast(k)), "layout X"); errdefer _ = lt.api.cublasLtMatrixLayoutDestroy(self.b); try lt.check(lt.api.cublasLtMatrixLayoutCreate(&self.c, out, n, m, @intCast(n)), "layout D"); errdefer _ = lt.api.cublasLtMatrixLayoutDestroy(self.c); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 75948fe016..5536e52bf4 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -33,6 +33,20 @@ pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &. pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; +pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; +pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; +pub const glm_experts: []const u8 = if (available) &Blob("fatbin_glm_experts").bytes else &.{}; +pub const glm_experts_prefill: []const u8 = if (available) &Blob("fatbin_glm_experts_prefill").bytes else &.{}; +pub const glm_vision: []const u8 = if (available) &Blob("fatbin_glm_vision").bytes else &.{}; +pub const glm_mm: []const u8 = if (available) &Blob("fatbin_glm_mm").bytes else &.{}; +pub const glm_rows: []const u8 = if (available) &Blob("fatbin_glm_rows").bytes else &.{}; +pub const glm_latent: []const u8 = if (available) &Blob("fatbin_glm_latent").bytes else &.{}; +pub const glm_index: []const u8 = if (available) &Blob("fatbin_glm_index").bytes else &.{}; +pub const glm_dflash: []const u8 = if (available) &Blob("fatbin_glm_dflash").bytes else &.{}; +pub const roce_gather: []const u8 = if (available) &Blob("fatbin_roce_gather").bytes else &.{}; +pub const glm_exl3: []const u8 = if (available) &Blob("fatbin_glm_exl3").bytes else &.{}; +pub const glm_experts3: []const u8 = if (available) &Blob("fatbin_glm_experts3").bytes else &.{}; +pub const glm_experts3_prefill: []const u8 = if (available) &Blob("fatbin_glm_experts3_prefill").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_pointwise").bytes else &.{}; @@ -46,7 +60,6 @@ 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 roce_gather: []const u8 = if (available) &Blob("fatbin_roce_gather").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/glm_cuda/api.zig b/zig/src/families/glm_cuda/api.zig new file mode 100644 index 0000000000..4060a7f6aa --- /dev/null +++ b/zig/src/families/glm_cuda/api.zig @@ -0,0 +1,141 @@ +//! GLM-5.3-Flash (glm5_next) on CUDA, two ranks: the pieces the weight loader, the forward and the native host share. +//! +//! Weights are a name -> device tensor store whose names are the Python engine's prepared-weight names (its weight +//! tree: "layers.3.dsa.q_b.weight", "layers.5.moe.experts.up", ...), so every tensor the checkpoint loader builds can +//! be checked byte for byte against the Python engine's (tools: capture_ref.py's weights.json hashes). +//! Kernels: the fatbins embedded by the build (this family's, zig/kernels/cuda/glm; the shared grouped experts and the +//! greedy candidates, zig/kernels/cuda/flashnext). + +const std = @import("std"); +const cuda = @import("cuda"); + +pub const DType = enum { bf16, f16, f32, i32, i64, u8, i8, i16, f8e4m3, bool_ }; + +pub fn dtypeBytes(t: DType) usize { + return switch (t) { + .bf16, .f16, .i16 => 2, + .f32, .i32 => 4, + .i64 => 8, + .u8, .i8, .f8e4m3, .bool_ => 1, + }; +} + +pub fn parseDType(s: []const u8) ?DType { + const map = std.StaticStringMap(DType).initComptime(.{ .{ "bfloat16", .bf16 }, .{ "float16", .f16 }, + .{ "float32", .f32 }, .{ "int32", .i32 }, .{ "int64", .i64 }, .{ "uint8", .u8 }, .{ "int8", .i8 }, + .{ "int16", .i16 }, .{ "float8_e4m3fn", .f8e4m3 }, .{ "bool", .bool_ } }); + return map.get(s); +} + +/// A device tensor: address, element type, shape (row-major, contiguous unless ``stride`` says otherwise). +pub const Tensor = struct { + ptr: u64, + dtype: DType, + shape: [5]i64 = .{ 1, 1, 1, 1, 1 }, + nd: u8 = 0, + + pub fn numel(t: Tensor) usize { + var n: usize = 1; + for (t.shape[0..t.nd]) |d| n *= @intCast(d); + return n; + } + pub fn bytes(t: Tensor) usize { + return t.numel() * dtypeBytes(t.dtype); + } + pub fn dim(t: Tensor, i: usize) i64 { + return t.shape[i]; + } +}; + +/// Prepared weights by name, all on the device. ``deinit`` frees what ``owned`` holds. +pub const Store = struct { + gpa: std.mem.Allocator, + map: std.StringHashMap(Tensor), + owned: std.ArrayList(cuda.DeviceBuffer) = .empty, + /// Host-side facts the loader records (vocabulary offset, widths), by name. + ints: std.StringHashMap(i64), + /// Host data by name (owned). + host: std.StringHashMap([]u8), + + pub fn init(gpa: std.mem.Allocator) Store { + return .{ .gpa = gpa, .map = .init(gpa), .ints = .init(gpa), .host = .init(gpa) }; + } + + pub fn deinit(self: *Store) void { + for (self.owned.items) |*b| b.free(); + self.owned.deinit(self.gpa); + var it = self.map.keyIterator(); + while (it.next()) |k| self.gpa.free(k.*); + self.map.deinit(); + var it2 = self.ints.keyIterator(); + while (it2.next()) |k| self.gpa.free(k.*); + self.ints.deinit(); + var it3 = self.host.iterator(); + while (it3.next()) |e| { + self.gpa.free(e.key_ptr.*); + self.gpa.free(e.value_ptr.*); + } + self.host.deinit(); + } + + pub fn put(self: *Store, name: []const u8, t: Tensor) !void { + try self.map.put(try self.gpa.dupe(u8, name), t); + } + + pub fn get(self: *const Store, name: []const u8) !Tensor { + return self.map.get(name) orelse { + std.log.err("glm: no weight named {s}", .{name}); + return error.MissingWeight; + }; + } + + pub fn getf(self: *const Store, comptime fmt: []const u8, args: anytype) !Tensor { + var buf: [256]u8 = undefined; + return self.get(try std.fmt.bufPrint(&buf, fmt, args)); + } +}; + +/// The device, its stream and the two ranks' communicator. +pub const Ctx = struct { + d: *const cuda.Driver, + ctx: *const cuda.Context, + stream: cuda.Stream, + nccl: *cuda.nccl.Library, + comm: cuda.nccl.Comm, + rank: u8, + world: u8 = 2, + /// Small all-gathers as RDMA writes over RoCE (cuda.roce), when both ranks brought it up; NCCL otherwise. + roce: ?*cuda.roce.Gather = null, + + /// Both ranks' ``count`` values from ``send`` into ``recv`` in rank order (ncclAllGather's layout and bits). + pub fn allGather(c: *const Ctx, send: u64, recv: u64, count: usize, dt: cuda.nccl.DataType, what: []const u8) !void { + const size: usize = switch (dt) { + .i8, .u8 => 1, + .f16, .bf16 => 2, + .i32, .u32, .f32 => 4, + .i64, .u64, .f64 => 8, + }; + if (c.roce) |g| if (g.fits(count * size, send, recv)) return g.gather(c.stream.handle, send, recv, count * size); + try c.nccl.check(c.nccl.api.ncclAllGather(send, recv, count, dt, c.comm, c.stream.handle), what); + } +}; + +/// The kernels: this family's modules and the shared ones it launches (embedded fatbins). +pub const Kernels = struct { + ext: [ext_names.len]cuda.Module, + + pub const ext_names = [_][]const u8{ "fn_experts", "fn_experts_prefill", "glm_kda", "fn_sample", "glm_ops", "glm_experts", "glm_experts_prefill", "glm_vision", "glm_experts3", "glm_experts3_prefill", "glm_mm", "glm_rows", "glm_latent", "glm_index", "glm_exl3", "glm_dflash" }; + + pub fn load(d: *const cuda.Driver) !Kernels { + var k: Kernels = .{ .ext = undefined }; + const images = [_][]const u8{ cuda.kernels.fn_experts, cuda.kernels.fn_experts_prefill, cuda.kernels.glm_kda, cuda.kernels.fn_sample, cuda.kernels.glm_ops, + cuda.kernels.glm_experts, cuda.kernels.glm_experts_prefill, cuda.kernels.glm_vision, cuda.kernels.glm_experts3, + cuda.kernels.glm_experts3_prefill, cuda.kernels.glm_mm, cuda.kernels.glm_rows, cuda.kernels.glm_latent, cuda.kernels.glm_index, cuda.kernels.glm_exl3, cuda.kernels.glm_dflash }; + for (images, 0..) |img, i| k.ext[i] = try cuda.Module.load(d, img); + return k; + } + + pub fn deinit(self: *Kernels) void { + for (&self.ext) |*m| m.unload(); + } +}; diff --git a/zig/src/families/glm_cuda/blocks.zig b/zig/src/families/glm_cuda/blocks.zig new file mode 100644 index 0000000000..1d2c4fc335 --- /dev/null +++ b/zig/src/families/glm_cuda/blocks.zig @@ -0,0 +1,462 @@ +//! GLM-5.3-Flash on CUDA: the weights' views and the layers of a step (forward.py's blocks): KDA and DSA attention, +//! the dense MLP and the MoE, the hyper-connections around them, the final norm and the head; and the MTP head. +//! Row-wise work covers the step's rows at once; the KDA recurrences, the attention and the indexer read the step's +//! tables (forward.prepare). + +const std = @import("std"); +const cuda = @import("cuda"); +const api = @import("api.zig"); +const dflash = @import("dflash.zig"); +const kern = @import("kern.zig"); +const sq = @import("seq.zig"); +const forward = @import("forward.zig"); + +const Engine = forward.Engine; +const Buffers = forward.Buffers; +const D = kern.D; +const HL = kern.HL; +const LL = kern.LL; +const PROJ_W = sq.PROJ_W; +const FA_OFF: i64 = 3 * LL * 128; +const GA_OFF: i64 = FA_OFF + 128; +const B_OFF: i64 = GA_OFF + 128; +const MOE_W = forward.MOE_W; +const DENSE_W = forward.DENSE_W; +const IDX_KW = forward.IDX_KW; + +// -- weights ------------------------------------------------------------------------------------------------------------ +const HCW = struct { fn_: u64, base: u64, scale: u64 }; +const KdaW = struct { proj: kern.B16, fb: kern.B16, gb: kern.B16, conv: u64, a_log: u64, dt_bias: u64, norm: u64, o: kern.B16 }; +const IndexW = struct { kw: kern.B16, qb: kern.B16, ln_w: u64, ln_b: u64, gate: u64, ape: u64 }; +const DsaW = struct { proj: kern.B16, q_norm: u64, kv_norm: u64, q_b: kern.B16, wk: u64, wv: u64, o: kern.B16, index: IndexW }; +const MlpW = struct { gu: kern.B16, down: kern.B16 }; +const MoeW = struct { router: u64, bias: u64, up: u64 = 0, down: u64 = 0, exl3: ?kern.Exl3 = null, shared: MlpW }; +const LayerW = struct { + attn_hc: HCW = undefined, + ffn_hc: HCW = undefined, + in_norm: u64, + post_norm: u64, + kda: ?KdaW = null, + dsa: ?DsaW = null, + mlp: ?MlpW = null, + moe: ?MoeW = null, + li: usize = 0, // its KDA or DSA index +}; +const MtpW = struct { enorm: u64, hnorm: u64, eh: kern.B16, norm: u64, layer: LayerW }; + +fn ptr(store: *const api.Store, comptime fmt: []const u8, args: anytype) !u64 { + return (try store.getf(fmt, args)).ptr; +} + +fn b16(store: *const api.Store, comptime fmt: []const u8, args: anytype) !kern.B16 { + var buf: [192]u8 = undefined; + const pre = try std.fmt.bufPrint(&buf, fmt, args); + var nb: [224]u8 = undefined; + if (store.map.get(try std.fmt.bufPrint(&nb, "{s}.qw", .{pre}))) |q| { // packed by the loader (weights.quantize) + var nb2: [224]u8 = undefined; + const sc = try store.get(try std.fmt.bufPrint(&nb2, "{s}.qs", .{pre})); + const kk = sc.shape[1] * 64; + return .{ .w = q.ptr, .n = q.shape[0], .k = kk, .s = sc.ptr, .bits = @intCast(@divExact(q.shape[1] * 8, kk)) }; + } + const t = try store.get(try std.fmt.bufPrint(&nb, "{s}.weight", .{pre})); + return .{ .w = t.ptr, .n = t.shape[0], .k = t.shape[1] }; +} + +fn hcw(store: *const api.Store, p: []const u8, site: []const u8) !HCW { + return .{ .fn_ = try ptr(store, "{s}.{s}.fn", .{ p, site }), .base = try ptr(store, "{s}.{s}.base", .{ p, site }), .scale = try ptr(store, "{s}.{s}.scale", .{ p, site }) }; +} + +fn mlpW(store: *const api.Store, p: []const u8) !MlpW { + return .{ .gu = try b16(store, "{s}.gu", .{p}), .down = try b16(store, "{s}.down", .{p}) }; +} + +fn dsaW(store: *const api.Store, p: []const u8) !DsaW { + return .{ .proj = try b16(store, "{s}.dsa.proj", .{p}), .q_norm = try ptr(store, "{s}.dsa.q_norm", .{p}), .kv_norm = try ptr(store, "{s}.dsa.kv_norm", .{p}), + .q_b = try b16(store, "{s}.dsa.q_b", .{p}), .wk = try ptr(store, "{s}.dsa.absorb.wk", .{p}), .wv = try ptr(store, "{s}.dsa.absorb.wv", .{p}), + .o = try b16(store, "{s}.dsa.o", .{p}), .index = .{ .kw = try b16(store, "{s}.dsa.index.kw", .{p}), .qb = try b16(store, "{s}.dsa.index.qb", .{p}), + .ln_w = try ptr(store, "{s}.dsa.index.ln_w", .{p}), .ln_b = try ptr(store, "{s}.dsa.index.ln_b", .{p}), .gate = try ptr(store, "{s}.dsa.index.gate", .{p}), + .ape = try ptr(store, "{s}.dsa.index.ape", .{p}) } }; +} + +fn moeW(store: *const api.Store, p: []const u8) !MoeW { + var b: [160]u8 = undefined; + var m: MoeW = .{ .router = try ptr(store, "{s}.moe.router", .{p}), .bias = try ptr(store, "{s}.moe.bias", .{p}), + .shared = try mlpW(store, try std.fmt.bufPrint(&b, "{s}.moe.shared", .{p})) }; + if ((store.ints.get("exl3") orelse 0) != 0) { + m.exl3 = .{ .gt = try ptr(store, "{s}.moe.experts.gt", .{p}), .ut = try ptr(store, "{s}.moe.experts.ut", .{p}), .dt = try ptr(store, "{s}.moe.experts.dt", .{p}), + .suh_g = try ptr(store, "{s}.moe.experts.suh_g", .{p}), .suh_u = try ptr(store, "{s}.moe.experts.suh_u", .{p}), + .svh_g = try ptr(store, "{s}.moe.experts.svh_g", .{p}), .svh_u = try ptr(store, "{s}.moe.experts.svh_u", .{p}), + .suh_d = try ptr(store, "{s}.moe.experts.suh_d", .{p}), .svh_d = try ptr(store, "{s}.moe.experts.svh_d", .{p}) }; + } else { + m.up = try ptr(store, "{s}.moe.experts.up", .{p}); + m.down = try ptr(store, "{s}.moe.experts.down", .{p}); + } + return m; +} + +fn layerW(store: *const api.Store, p: []const u8, kda: bool, dense: bool, plain: bool, li: usize) !LayerW { + var l: LayerW = .{ .in_norm = try ptr(store, "{s}.in_norm", .{p}), .post_norm = try ptr(store, "{s}.post_norm", .{p}), .li = li }; + if (!plain) { + l.attn_hc = try hcw(store, p, "attn_hc"); + l.ffn_hc = try hcw(store, p, "ffn_hc"); + } + if (kda) { + l.kda = .{ .proj = try b16(store, "{s}.kda.proj", .{p}), .fb = try b16(store, "{s}.kda.fb", .{p}), .gb = try b16(store, "{s}.kda.gb", .{p}), + .conv = try ptr(store, "{s}.kda.conv", .{p}), .a_log = try ptr(store, "{s}.kda.a_log", .{p}), .dt_bias = try ptr(store, "{s}.kda.dt_bias", .{p}), + .norm = try ptr(store, "{s}.kda.norm", .{p}), .o = try b16(store, "{s}.kda.o", .{p}) }; + } else l.dsa = try dsaW(store, p); + var b: [160]u8 = undefined; + if (dense) l.mlp = try mlpW(store, try std.fmt.bufPrint(&b, "{s}.mlp", .{p})) else l.moe = try moeW(store, p); + return l; +} + +pub const Weights = struct { + layers: [46]LayerW = undefined, + nlayers: usize = 0, + kinds: [46]bool = @splat(false), + nkda: usize = 0, + ndsa: usize = 0, + mtp: ?MtpW = null, + embed: u64 = 0, + norm: u64 = 0, + head: kern.B16 = undefined, + draft_head: kern.B16 = undefined, + vocab_offset: i64 = 0, + experts_sym: bool = false, + experts3: bool = false, + exl3: bool = false, // EXL3 routed experts (glm/exl3.cu) + exl3_shx: bool = false, // every layer's experts share one input rotation for gate and up (one rotated row a token in prompts) + + pub fn caches(w: *const Weights) usize { + return w.ndsa + @intFromBool(w.mtp != null); + } + + /// The store's views, and every dense linear tiled in place for glm/mm.cu (once, at load). + pub fn load(k: *kern.K, store: *const api.Store) !Weights { + var w: Weights = .{}; + w.vocab_offset = store.ints.get("vocab_offset").?; + const kinds = store.host.get("kinds").?; + w.nlayers = kinds.len; + const dense: usize = @intCast(store.ints.get("dense").?); + for (kinds, 0..) |kv, i| { + w.kinds[i] = kv != 0; + var nb: [32]u8 = undefined; + const li = if (w.kinds[i]) w.nkda else w.ndsa; + w.layers[i] = try layerW(store, try std.fmt.bufPrint(&nb, "layers.{d}", .{i}), w.kinds[i], i < dense, false, li); + if (w.kinds[i]) w.nkda += 1 else w.ndsa += 1; + } + if (store.ints.get("mtp").? != 0) { + w.mtp = .{ .enorm = try ptr(store, "mtp.enorm", .{}), .hnorm = try ptr(store, "mtp.hnorm", .{}), .eh = try b16(store, "mtp.eh", .{}), + .norm = try ptr(store, "mtp.norm", .{}), .layer = try layerW(store, "mtp.layer", false, false, true, w.ndsa) }; + } + w.embed = try ptr(store, "embed", .{}); + w.norm = try ptr(store, "norm", .{}); + w.experts_sym = (store.ints.get("experts_sym") orelse 0) != 0; + w.experts3 = (store.ints.get("experts3") orelse 0) != 0; + w.exl3 = (store.ints.get("exl3") orelse 0) != 0; + w.exl3_shx = (store.ints.get("exl3_shx") orelse 0) != 0; + w.head = try b16(store, "head", .{}); + // drafts only propose (verification reads ``head``): a smaller copy when the loader made one + w.draft_head = if (store.map.contains("draft_head.qw")) try b16(store, "draft_head", .{}) else w.head; + try w.tile(k); + return w; + } + + pub fn deinit(_: *Weights, _: std.mem.Allocator) void {} + + fn linears(w: *const Weights, out: *std.ArrayList(kern.B16), gpa: std.mem.Allocator) !void { + const add = struct { + fn f(o: *std.ArrayList(kern.B16), a: std.mem.Allocator, x: kern.B16) !void { + for (o.items) |y| if (y.w == x.w) return; + try o.append(a, x); + } + }.f; + const layer = struct { + fn f(o: *std.ArrayList(kern.B16), a: std.mem.Allocator, l: LayerW) !void { + if (l.kda) |x| for ([_]kern.B16{ x.proj, x.fb, x.gb, x.o }) |m| try add(o, a, m); + if (l.dsa) |x| for ([_]kern.B16{ x.proj, x.q_b, x.o, x.index.kw, x.index.qb }) |m| try add(o, a, m); + if (l.mlp) |x| for ([_]kern.B16{ x.gu, x.down }) |m| try add(o, a, m); + if (l.moe) |x| for ([_]kern.B16{ x.shared.gu, x.shared.down }) |m| try add(o, a, m); + } + }.f; + for (w.layers[0..w.nlayers]) |l| try layer(out, gpa, l); + if (w.mtp) |m| { + try add(out, gpa, m.eh); + try layer(out, gpa, m.layer); + } + try add(out, gpa, w.head); + try add(out, gpa, w.draft_head); + } + + fn tile(w: *const Weights, k: *kern.K) !void { + var all: std.ArrayList(kern.B16) = .empty; + defer all.deinit(k.gpa); + try w.linears(&all, k.gpa); + var most: u64 = 0; + for (all.items) |x| most = @max(most, wbytes(x)); + var tmp = try cuda.DeviceBuffer.alloc(k.d, most + (most / 32) + 512); + defer tmp.free(); + for (all.items) |x| { + if (@mod(x.n, 8) != 0 or @mod(x.k, 64) != 0) return error.LinearNotTileable; + const wb = wbytes(x); + try k.copy(tmp.ptr, x.w, wb); + const sb: u64 = if (x.bits != 0) @intCast(@divExact(x.n * x.k, 64) * 2) else 0; + const st = tmp.ptr + std.mem.alignForward(u64, wb, 256); + if (sb > 0) try k.copy(st, x.s, sb); + try kern.mmTile(k, x.bits, tmp.ptr, st, x.w, x.s, x.n, x.k); + try k.stream.synchronize(); + } + } + + fn wbytes(x: kern.B16) u64 { + const bits: i64 = if (x.bits == 0) 16 else x.bits; + return @intCast(@divExact(x.n * x.k * bits, 8)); + } +}; + +// -- the blocks ----------------------------------------------------------------------------------------------------------- +pub fn mm(e: *Engine, b: *Buffers, x: u64, x_stride: i64, w: kern.B16, out: u64, f32_out: bool, R: i64) !void { + if (e.skip.dense and w.n != forward.HEAD_N) return; + try kern.mm(&e.k, x, x_stride, w, out, f32_out, R, b.sk); +} + +fn outProj(e: *Engine, b: *Buffers, x: u64, x_stride: i64, w: kern.B16, R: i64) !u64 { + try mm(e, b, x, x_stride, w, b.part, true, R); + return forward.gather(e, b, R); +} + +/// A KDA layer: a prompt chunk's rows through its one stream's states (a kept point inside the chunk cuts the +/// chain there), or a decode step's windows, each through its own stream's (the segmented chain). +fn kdaBlock(e: *Engine, b: *Buffers, l: LayerW, R: i64) !u64 { + const k = l.kda.?; + const li = l.li; + const p = b.kproj; + try mm(e, b, b.normed, D, k.proj, p, false, R); + try mm(e, b, p + FA_OFF * 2, PROJ_W, k.fb, b.ka, false, R); + try mm(e, b, p + GA_OFF * 2, PROJ_W, k.gb, b.kg, false, R); + if (!b.prefill) { + try kern.kdaChainSeg(&e.k, b.segs, @intCast(e.step.segs.len), li, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, k.conv, k.a_log, k.dt_bias, + k.norm, b.kout, sq.MAX_ROWS, PROJ_W); + return outProj(e, b, b.kout, LL * 128, k.o, R); + } + const s = e.step.segs[0].s; + const cur = s.cur; + const conv = s.convAt(li); + const sc = b.ks; + const n = e.cut_at; + if (n > 0 and n < R) { // a kept point inside the chunk: the state and window at it, then the rest from them + const cut_rec = e.cut_rec + @as(u64, @intCast(li)) * @as(u64, @intCast(sq.REC * 4)); + const cut_conv = e.cut_conv + @as(u64, @intCast(li)) * @as(u64, @intCast(3 * sq.CONV_C * 2)); + try kern.kdaChain(&e.k, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, conv, k.conv, s.recAt(cur, li), k.a_log, k.dt_bias, k.norm, n, sc, + cut_rec, b.q_tmp, b.y_tmp); + try kern.convShift(&e.k, cut_conv, b.kproj, n, 3 * sq.CONV_C, b.rows * PROJ_W, PROJ_W, 1, sq.CONV_C); + const hw: u64 = @intCast(LL * 128); + const nr: u64 = @intCast(n); + const rest: kern.KdaScratch = .{ .out = sc.out + nr * hw * 2, .k = sc.k + nr * hw * 4, .v = sc.v + nr * hw * 2, .g = sc.g + nr * hw * 4, .b = sc.b + nr * @as(u64, @intCast(LL)) * 4 }; + try kern.kdaChain(&e.k, p + nr * @as(u64, @intCast(PROJ_W)) * 2, PROJ_W, B_OFF, b.ka + nr * hw * 2, LL * 128, b.kg + nr * hw * 2, LL * 128, cut_conv, k.conv, cut_rec, + k.a_log, k.dt_bias, k.norm, R - n, rest, s.recAt(1 - cur, li), b.q_tmp, b.y_tmp); + } else { + try kern.kdaChain(&e.k, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, conv, k.conv, s.recAt(cur, li), k.a_log, k.dt_bias, k.norm, R, sc, + s.recAt(1 - cur, li), b.q_tmp, b.y_tmp); + } + // a prompt chunk keeps every row: the layer's conv window commits now (its state flips with the chunk) + try kern.convShift(&e.k, conv, b.kproj, R, 3 * sq.CONV_C, b.rows * PROJ_W, PROJ_W, 1, sq.CONV_C); + return outProj(e, b, sc.out, LL * 128, k.o, R); +} + +/// Each row's attended tokens (ascending, -1 padded) and their count past the dense limit (0: a dense row). +fn selectTokens(e: *Engine, b: *Buffers, ci: usize, R: i64) !void { + const t = b.rowsTab(); + const B = @min(R, kern.SELECT_ROWS); + var a: i64 = 0; + while (a < R) : (a += B) { + const n = @min(B, R - a); + try kern.scores(&e.k, b.qi + @as(u64, @intCast(a * kern.IDX_HEADS * kern.IDX_DIM * 2)), b.ikr + @as(u64, @intCast(a * IDX_KW + kern.IDX_DIM)) * 2, IDX_KW, + t.at(a), ci, b.scores, n, e.step.np); + try kern.selectRows(&e.k, b.scores, b.pools + @as(u64, @intCast(a * kern.TOPK_POOLS * 8)), e.step.np, t.at(a), n); + } + try kern.sparseTokens(&e.k, b.pools, t, R, b.tokens, b.counts); +} + +/// A DSA layer on latent cache ``ci`` (the MTP layer's is the last): the indexer always (its keys and pools are kept +/// for every position), dense attention for the rows below the dense limit, sparse for the rest. +fn dsaBlock(e: *Engine, b: *Buffers, a: DsaW, ci: usize, R: i64) !u64 { + const t = b.rowsTab(); + try mm(e, b, b.normed, D, a.proj, b.dp, false, R); + try kern.rmsnorm(&e.k, b.dp, kern.Q_LORA + kern.KV_LORA, a.q_norm, b.qr, kern.Q_LORA, b.xs_qr, R, kern.Q_LORA); + try kern.rmsnorm(&e.k, b.dp + kern.Q_LORA * 2, kern.Q_LORA + kern.KV_LORA, a.kv_norm, b.lat, kern.KV_LORA, b.xs_lat, R, kern.KV_LORA); + try mm(e, b, b.qr, kern.Q_LORA, a.q_b, b.q, false, R); + try kern.latentWrite(&e.k, b.lat, kern.KV_LORA, t, ci, R); + try mm(e, b, b.normed, D, a.index.kw, b.ikr, false, R); + try kern.router(&e.k, b.normed, D, a.index.gate, b.igr, R, kern.IDX_DIM); + try kern.indexUpdate(&e.k, b.ikr, IDX_KW, b.igr, a.index.ln_w, a.index.ln_b, a.index.ape, t, ci, R); + if (e.skip.attn) return outProj(e, b, b.vn, HL * kern.DV, a.o, R); + try kern.absorbQ(&e.k, b.q, a.wk, b.qa, R); + const att = b.attRows(); + const lat_row: u64 = @intCast(HL * kern.KV_LORA * 2); + if (e.step.nch > 0) { + var r0: i64 = 0; + while (r0 < R) : (r0 += att) { + const n = @min(att, R - r0); + const off = @as(u64, @intCast(r0)) * lat_row; + try kern.denseAttention(&e.k, b.qa + off, t.at(r0), ci, b.po, b.pm, b.pl, b.ol + off, n, e.step.nch); + } + } + if (e.step.sparse) { + try mm(e, b, b.qr, kern.Q_LORA, a.index.qb, b.qi, false, R); + try selectTokens(e, b, ci, R); + var r0: i64 = 0; + while (r0 < R) : (r0 += att) { + const n = @min(att, R - r0); + const off = @as(u64, @intCast(r0)) * lat_row; + try kern.sparseAttention(&e.k, b.qa + off, t.at(r0), ci, b.tokens + @as(u64, @intCast(r0 * kern.SEL_W * 4)), b.counts + @as(u64, @intCast(r0 * 4)), + b.po, b.pm, b.pl, b.ol + off, n); + } + } + try kern.expandV(&e.k, b.ol, a.wv, b.vn, R); + return outProj(e, b, b.vn, HL * kern.DV, a.o, R); +} + +fn mlpBlock(e: *Engine, b: *Buffers, m: MlpW, R: i64) !u64 { + try mm(e, b, b.normed, D, m.gu, b.gu, false, R); + try kern.swiglu(&e.k, b.gu, b.act, b.xs_act, R, DENSE_W); + return outProj(e, b, b.act, DENSE_W, m.down, R); +} + +/// The MoE: the routed slots through the grouped kernels (8 a row), the shared expert (slot 8) through the dense +/// linears, the slots combined by their weights. +fn moeBlock(e: *Engine, b: *Buffers, m: MoeW, R: i64) !u64 { + try kern.router(&e.k, b.normed, D, m.router, b.mlog, R, kern.EXPERTS); + try kern.topk(&e.k, b.mlog, m.bias, b.pick, b.wts, R); + try kern.copyRows(&e.k, b.rpick, @intCast(kern.TOPK * 4), b.pick, @intCast(kern.SLOTS * 4), @intCast(kern.TOPK * 4), @intCast(R)); + const pairs = R * kern.TOPK; + try kern.expertsPlan(&e.k, b.rpick, pairs, b.plan); + const items = kern.maxItems(pairs, kern.EXPERTS, b.plan.tile); + const gs: i64 = 64; + const nbu = @divExact(MOE_W, 32); + const nbd = @divExact(D, 32); + if (m.exl3) |x| { // the routed slots straight into ey (fp32, rows of 9 slots), then the shared expert's slot + if (!e.skip.experts) try kern.exl3Routed(&e.k, b.normed, D, b.rpick, b.plan, x, b.x3, b.ey, kern.SLOTS, R, MOE_W, b.prefill, e.w.exl3_shx); + const dsz: usize = @intCast(D); + try mm(e, b, b.normed, D, m.shared.gu, b.sgu, false, R); + try kern.swiglu(&e.k, b.sgu, b.sact, b.sxs, R, MOE_W); + try mm(e, b, b.sact, MOE_W, m.shared.down, b.sy, true, R); + try kern.copyRows(&e.k, b.ey + @as(u64, @intCast(kern.TOPK * D * 4)), @as(usize, @intCast(kern.SLOTS)) * dsz * 4, b.sy, dsz * 4, dsz * 4, @intCast(R)); + try kern.combine(&e.k, b.ey, true, b.wts, b.part, R); + return forward.gather(e, b, R); + } + if (b.prefill) { + try kern.expertsCall(&e.k, true, 2, b.normed, D, kern.TOPK, m.up, @divExact(D, gs), nbu, b.plan, b.eact, MOE_W, items, kern.LIMIT); + try kern.expertsCall(&e.k, true, 3, b.eact, MOE_W, 0, m.down, @divExact(MOE_W, gs), nbd, b.plan, b.rey, D, items, 0.0); + } else { + try kern.expertsCall(&e.k, false, 2, b.normed, D, kern.TOPK, m.up, @divExact(D, gs), nbu, b.plan, b.eact, MOE_W, items * nbu, kern.LIMIT); + try kern.expertsCall(&e.k, false, 0, b.eact, MOE_W, 0, m.down, @divExact(MOE_W, gs), nbd, b.plan, b.rey, D, items * nbd, 0.0); + } + const elt: usize = if (b.prefill) 2 else 4; + const dsz: usize = @intCast(D); + const slots: usize = @intCast(kern.SLOTS); + const topk: usize = @intCast(kern.TOPK); + try kern.copyRows(&e.k, b.ey, slots * dsz * elt, b.rey, topk * dsz * elt, topk * dsz * elt, @intCast(R)); + try mm(e, b, b.normed, D, m.shared.gu, b.sgu, false, R); + try kern.swiglu(&e.k, b.sgu, b.sact, b.sxs, R, MOE_W); + try mm(e, b, b.sact, MOE_W, m.shared.down, b.sy, true, R); + const slot8: u64 = @intCast(kern.TOPK * D * @as(i64, @intCast(elt))); + if (b.prefill) { + try kern.castBf16(&e.k, b.sy, D, b.ey + slot8, kern.SLOTS * D, R, D); + } else { + try kern.copyRows(&e.k, b.ey + slot8, slots * dsz * 4, b.sy, dsz * 4, dsz * 4, @intCast(R)); + } + try kern.combine(&e.k, b.ey, !b.prefill, b.wts, b.part, R); + return forward.gather(e, b, R); +} + +fn layerForward(e: *Engine, b: *Buffers, l: LayerW, R: i64) !void { + const h = l.attn_hc; + if (e.skip.hc) { + const g = if (l.kda != null) try kdaBlock(e, b, l, R) else try dsaBlock(e, b, l.dsa.?, l.li, R); + _ = g; + _ = if (l.mlp) |m| try mlpBlock(e, b, m, R) else try moeBlock(e, b, l.moe.?, R); + return; + } + try kern.hcPre(&e.k, b.x, h.fn_, h.base, h.scale, l.in_norm, b.normed, b.xs, b.post, b.comb, b.hcpart, R); + const g = if (l.kda != null) try kdaBlock(e, b, l, R) else try dsaBlock(e, b, l.dsa.?, l.li, R); + try kern.hcPost(&e.k, b.x, g, b.post, b.comb, R); + const f = l.ffn_hc; + try kern.hcPre(&e.k, b.x, f.fn_, f.base, f.scale, l.post_norm, b.normed, b.xs, b.post, b.comb, b.hcpart, R); + const g2 = if (l.mlp) |m| try mlpBlock(e, b, m, R) else try moeBlock(e, b, l.moe.?, R); + try kern.hcPost(&e.k, b.x, g2, b.post, b.comb, R); +} + +/// ``x``'s rows (positions first .. first + R, ``copies`` streams each) that fall in an image's run take its features. +fn injectMedia(e: *Engine, s: *const sq.Seq, first: i64, R: i64, x: u64, copies: i64) !void { + const m = s.media orelse return; + var row0: i64 = 0; // the run's first feature row + for (m.spans) |sp| { + const lo = @max(sp[0], first); + const hi = @min(sp[1], first + R); + if (lo < hi) { + const f = m.feats + @as(u64, @intCast((row0 + lo - sp[0]) * D * 2)); + try kern.injectRows(&e.k, f, x + @as(u64, @intCast((lo - first) * copies * D * 2)), hi - lo, copies); + } + row0 += sp[1] - sp[0]; + } +} + +/// The step's rows through every layer, the final norm, and the head (a prompt chunk: its last row only). The ids +/// are staged in b.ids, the tables prepared. +pub fn compute(e: *Engine, b: *Buffers) !void { + const R = e.step.rows; + const w = &e.w; + try kern.embed(&e.k, b.ids, w.embed, b.x, R, kern.S); + if (b.prefill) try injectMedia(e, e.step.segs[0].s, e.step.first_pos, R, b.x, kern.S); + // DFlash2's taps: a verify step's rows and a prompt chunk's (MTP and the drafter's own runs take none) + const taps: u64 = if (e.drafter) |dr| (if (b.prefill) dr.taps_prefill else if (b == &e.buf) dr.taps_step else 0) else 0; + for (w.layers[0..w.nlayers], 0..) |l, li| { + try layerForward(e, b, l, R); + if (taps != 0 and dflash.isTap(e.drafter.?, li)) try dflash.tap(e.drafter.?, &e.k, b.x, taps, li, R); + } + try kern.streamMean(&e.k, b.x, b.hidden, R); + try kern.rmsnorm(&e.k, b.hidden, D, w.norm, b.fnormed, D, b.fxs, R, D); + if (b.prefill) { + try mm(e, b, b.fnormed + @as(u64, @intCast((R - 1) * D * 2)), D, w.head, b.logits, false, 1); + } else { + try mm(e, b, b.fnormed, D, w.head, b.logits, false, R); + } +} + +/// The MTP head over the step's staged rows (ids in b.ids, input hidden rows in b.hin), then each stream's last row's +/// logits (in segment order). +pub fn mtpCompute(e: *Engine, b: *Buffers) !void { + const m = e.w.mtp.?; + const R = e.step.rows; + const t = b.rowsTab(); + try kern.embed(&e.k, b.ids, e.w.embed, b.me, R, 1); + if (b.prefill) try injectMedia(e, e.step.segs[0].s, e.step.first_pos + 1, R, b.me, 1); // MTP row i reads the token at mtp_len + 1 + i + try zeroFirst(e, b.me, t.posAt(), R); + try kern.rmsnorm(&e.k, b.me, D, m.enorm, b.mcat, 2 * D, null, R, D); + try kern.rmsnorm(&e.k, b.hin, D, m.hnorm, b.mcat + D * 2, 2 * D, null, R, D); + try mm(e, b, b.mcat, 2 * D, m.eh, b.mx, false, R); + const l = m.layer; + try kern.rmsnorm(&e.k, b.mx, D, l.in_norm, b.normed, D, b.xs, R, D); + const g = try dsaBlock(e, b, l.dsa.?, e.w.ndsa, R); + try kern.residualAdd(&e.k, b.mx, g, R); + try kern.rmsnorm(&e.k, b.mx, D, l.post_norm, b.normed, D, b.xs, R, D); + const g2 = try moeBlock(e, b, l.moe.?, R); + try kern.residualAdd(&e.k, b.mx, g2, R); + const nseg: i64 = @intCast(e.step.segs.len); + try gatherRows(e, b.mx, b.last, b.hrows, nseg); + try kern.rmsnorm(&e.k, b.hrows, D, m.norm, b.fnormed, D, b.fxs, nseg, D); + try mm(e, b, b.fnormed, D, e.w.draft_head, b.logits, false, nseg); +} + +fn zeroFirst(e: *Engine, x: u64, pos: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(x); a.add(pos); a.add(@as(i32, @intCast(D))); + try e.k.go(try e.k.ext(.rows, "glm_zero_first"), .{ .x = @intCast(R) }, 256, 0, &a); +} + +fn gatherRows(e: *Engine, src: u64, idx: u64, dst: u64, n: i64) !void { + var a: cuda.Args = .{}; + a.add(src); a.add(idx); a.add(dst); a.add(@as(i32, @intCast(D))); + try e.k.go(try e.k.ext(.rows, "glm_gather_rows"), .{ .x = @intCast(n) }, 256, 0, &a); +} diff --git a/zig/src/families/glm_cuda/calibrate.zig b/zig/src/families/glm_cuda/calibrate.zig new file mode 100644 index 0000000000..89030cae0a --- /dev/null +++ b/zig/src/families/glm_cuda/calibrate.zig @@ -0,0 +1,132 @@ +//! GLM-5.3-Flash's costs for the lane core, timed at load on rank 0 (rank 1 replays every call): a lone stream's +//! verify window by width, shared steps by their total rows (up to 16 streams' windows), and a chained head step. + +const std = @import("std"); +const lanes = @import("lanes"); +const forward = @import("forward.zig"); +const dflash = @import("dflash.zig"); +const decode = @import("decode.zig"); +const proto = @import("proto.zig"); + +const max_rows: usize = @intCast(forward.MAX_ROWS); +const streams = forward.MAX_STREAMS; + +fn msSince(io: std.Io, t0: std.Io.Timestamp) f64 { + return @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(io, .awake)).toNanoseconds())) / 1e6; +} + +/// ``self``: native.zig's Owned. +pub fn run(self: anytype, drafts: bool) !void { + var w: proto.Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + var seqs: [streams]*forward.Seq = undefined; + var toks: [streams]u32 = undefined; + var prompt: [64]u32 = undefined; + for (0..streams) |i| { + for (&prompt, 0..) |*t, j| t.* = @intCast(1000 + (j * 37 + i * 101) % 5000); + const uid = self.next_uid; + self.next_uid += 1; + seqs[i] = try self.newSeq(uid, 4096); + try w.int(uid); + try w.int(@as(i64, 4096)); + try w.tokens(&prompt); + try proto.writeSampling(&w, null); + try w.int(0); + try w.int(0); + try self.send(.prefill, &w); + seqs[i].sampling = null; + toks[i] = try decode.prefill(self.e, seqs[i], &prompt, 0, 0); + } + var picks: [forward.STEP_ROWS]decode.Pick = undefined; + var ids: [streams][max_rows]u32 = undefined; + // one step of ``n`` streams' windows of ``width`` rows (distinct tokens: rows routed to their own experts, as a + // real window's are), each kept at its first row; its ms + const Step = struct { + fn go(o: anytype, wr: *proto.Writer, ss: []*forward.Seq, tk: []u32, idv: *[streams][max_rows]u32, pk: []decode.Pick, n: usize, width: usize, rep: usize) !f64 { + var wins: [streams]decode.Win = undefined; + try wr.int(n); + for (0..n) |i| { + for (idv[i][0..width], 0..) |*t, r| t.* = if (r == 0) tk[i] else @intCast(1000 + ((r * 7919 + rep * 104729 + width * 31 + i * 613) % 60000)); + wins[i] = .{ .s = ss[i], .ids = idv[i][0..width] }; + try wr.int(ss[i].uid); + try wr.tokens(idv[i][0..width]); + } + try o.send(.verify, wr); + const t0 = std.Io.Timestamp.now(o.io, .awake); + try decode.verify(o.e, wins[0..n], pk); + const ms = msSince(o.io, t0); + try wr.int(n); + for (0..n) |i| { + try wr.int(ss[i].uid); + try wr.int(width); + try wr.int(1); + } + try o.send(.keep, wr); + for (0..n) |i| { + try decode.keep(o.e, ss[i], @intCast(width), 1); + tk[i] = pk[i * width].tok; + } + return ms; + } + }; + // warm-up (the GPU's clocks rise over the first steps; every shape's graph captured once) + for (0..16) |rep| _ = try Step.go(self, &w, &seqs, &toks, &ids, &picks, 1, 1, rep); + for (1..max_rows + 1) |width| { + var best: f64 = std.math.inf(f64); + for (0..4) |rep| { + const ms = try Step.go(self, &w, &seqs, &toks, &ids, &picks, 1, width, rep); + if (rep > 0) best = @min(best, ms); + } + self.costs[width - 1] = .{ .width = @intCast(width), .ms = best }; + } + self.cost_count = max_rows; + // shared steps: n streams of w rows (total rows n * w) + const shapes = [_][2]usize{ .{ 2, 1 }, .{ 4, 1 }, .{ 8, 1 }, .{ 16, 1 }, .{ 16, 2 }, .{ 16, 4 }, .{ 16, 8 } }; + for (shapes, 0..) |sh, k| { + var best: f64 = std.math.inf(f64); + for (0..4) |rep| { + const ms = try Step.go(self, &w, &seqs, &toks, &ids, &picks, sh[0], sh[1], rep); + if (rep > 0) best = @min(best, ms); + } + self.shared[k] = .{ .width = @intCast(sh[0] * sh[1]), .ms = best }; + } + self.shared_count = shapes.len; + if (drafts) { + var held: [max_rows]u32 = undefined; + var probs: [max_rows]f64 = undefined; + var at: [2]f64 = .{ std.math.inf(f64), std.math.inf(f64) }; + for ([_]usize{ 1, 7 }, 0..) |depth, k| { + for (0..3) |rep| { + const follow = [_]u32{toks[0]}; + try w.int(1); + try w.int(seqs[0].uid); + try w.tokens(&follow); + try w.int(depth); + try w.int(0); // the hidden row: the last verify's + try w.int(0); // no confidence cut: every level timed + try w.int(@intFromBool(self.e.drafter != null)); // DFlash2 drafts when loaded + try self.send(.draft, &w); + var reqs = [_]decode.DraftReq{.{ .s = seqs[0], .hidden = seqs[0].hidden, .next = &follow, .count = depth, .confidence = 0, + .out = held[0..depth], .probs = probs[0..depth] }}; + const t0 = std.Io.Timestamp.now(self.io, .awake); + // rank 1 replays the frame with the drafter it serves with: DFlash2's round when one is loaded + if (self.e.drafter) |dr| { // as rank 1 replays it: the MTP head kept current (while the mix may pick it), then DFlash2's block + reqs[0].count = 1; + if (@import("native.zig").feedsMtp(self)) try decode.draft(self.e, &reqs); + try dflash.draft(dr, self.e, &.{.{ .s = seqs[0], .n = 1, .anchor = toks[0], .out = held[0..max_rows - 1], .probs = probs[0..max_rows - 1] }}); + } else try decode.draft(self.e, &reqs); + if (rep > 0) at[k] = @min(at[k], msSince(self.io, t0)); + } + } + // a DFlash2 round drafts the whole block whatever the depth: its cost spread over the 7 drafts + self.mtp_ms = if (self.e.drafter != null) at[1] / 7 else @max(0, (at[1] - at[0]) / 6); + } + for (seqs) |s| { + try w.int(s.uid); + try w.int(0); + try self.send(.release, &w); + self.dropSeq(s.uid); + } + std.log.info("glm costs: window 1/2/4/8 rows {d:.1}/{d:.1}/{d:.1}/{d:.1} ms; 16 streams x 1/2/4/8 rows {d:.1}/{d:.1}/{d:.1}/{d:.1} ms; head step {d:.2} ms", .{ + self.costs[0].ms, self.costs[1].ms, self.costs[3].ms, self.costs[7].ms, self.shared[3].ms, self.shared[4].ms, self.shared[5].ms, self.shared[6].ms, self.mtp_ms }); +} diff --git a/zig/src/families/glm_cuda/copy.zig b/zig/src/families/glm_cuda/copy.zig new file mode 100644 index 0000000000..de59affb39 --- /dev/null +++ b/zig/src/families/glm_cuda/copy.zig @@ -0,0 +1,117 @@ +//! Copy drafts (prompt lookup): code, JSON and edits repeat what came before, often longer than the MTP head drafts +//! with confidence. A stream keeps its prompt and every committed token, and an index of each 8-token run's latest +//! end; when the history's last 8 tokens ran before, the tokens that followed them are a draft. A run inside the +//! reply itself must match 16 tokens (a reply's own short repeats mislead more often). The drafts are verified like +//! the head's, so a wrong copy costs rows, never the reply. + +const std = @import("std"); + +pub const K: usize = 8; // tokens a match in the prompt takes by default (TF_GLM_COPY_MIN) +pub const K_REPLY: usize = 16; // tokens a match inside the reply takes by default (TF_GLM_COPY_MIN_REPLY) + +pub const Copy = struct { + hist: std.ArrayList(u32) = .empty, + index: std.AutoHashMapUnmanaged(u64, u32) = .empty, + prompt_len: usize = 0, + from: usize = 0, // the first position indexed (a long prompt's head is left out) + last: ?u32 = null, // where the history's last k tokens ended before (an earlier occurrence) + k: usize = K, // the indexed runs' length: the shorter of the two matches + k_prompt: usize = K, + k_reply: usize = K_REPLY, + + /// ``prompt`` kept and its runs indexed, at most ``most`` of its last tokens; a match takes ``k_prompt`` tokens + /// in the prompt and ``k_reply`` inside the reply. + pub fn init(gpa: std.mem.Allocator, prompt: []const u32, most: usize, k_prompt: usize, k_reply: usize) !Copy { + const kp = @max(2, k_prompt); + const kr = @max(2, k_reply); + var c: Copy = .{ .prompt_len = prompt.len, .from = prompt.len -| most, .k = @min(kp, kr), .k_prompt = kp, .k_reply = kr }; + errdefer c.deinit(gpa); + try c.hist.appendSlice(gpa, prompt); + try c.index.ensureTotalCapacity(gpa, @intCast(@min(prompt.len, most) + 1024)); + var i = @max(c.from, c.k - 1); + while (i < c.hist.items.len) : (i += 1) c.last = c.note(i); + return c; + } + + pub fn deinit(c: *Copy, gpa: std.mem.Allocator) void { + c.hist.deinit(gpa); + c.index.deinit(gpa); + } + + /// Committed tokens appended. + pub fn append(c: *Copy, gpa: std.mem.Allocator, toks: []const u32) !void { + try c.index.ensureUnusedCapacity(gpa, @intCast(toks.len)); + for (toks) |t| { + try c.hist.append(gpa, t); + const i = c.hist.items.len - 1; + if (i >= c.k - 1 and i >= c.from) c.last = c.note(i); + } + } + + /// The run ending at ``i`` indexed; returns where it ended before. + fn note(c: *Copy, i: usize) ?u32 { + const gop = c.index.getOrPutAssumeCapacity(hash(c.hist.items[i + 1 - c.k .. i + 1])); + const prev: ?u32 = if (gop.found_existing) gop.value_ptr.* else null; + gop.value_ptr.* = @intCast(i); + return prev; + } + + /// The tokens that followed the last match, into ``out``; how many (0: no match). + pub fn propose(c: *const Copy, out: []u32) usize { + const p: usize = c.last orelse return 0; + const h = c.hist.items; + const end = h.len - 1; + if (p >= end) return 0; + const need: usize = if (p >= c.prompt_len) c.k_reply else c.k_prompt; + if (p + 1 < need) return 0; + if (!std.mem.eql(u32, h[p + 1 - need .. p + 1], h[end + 1 - need .. end + 1])) return 0; // a hash collision, or short of 16 + const n = @min(out.len, end - p); + @memcpy(out[0..n], h[p + 1 .. p + 1 + n]); + return n; + } +}; + +fn hash(run: []const u32) u64 { + return std.hash.Wyhash.hash(0x636f7079, std.mem.sliceAsBytes(run)); +} + +/// Copied drafts' landing rate, learned from the verifies (shared by the streams): starts at 0.8. +pub const Rate = struct { + tried: f64 = 0, + landed: f64 = 0, + + pub fn chance(r: *const Rate) f64 { + return (r.landed + 4) / (r.tried + 5); + } + + /// A window that verified ``drafted`` copied tokens kept ``kept`` of them; older windows fade. + pub fn record(r: *Rate, drafted: u32, kept: u32) void { + r.tried = r.tried * 0.995 + @as(f64, @floatFromInt(drafted)); + r.landed = r.landed * 0.995 + @as(f64, @floatFromInt(@min(kept, drafted))); + } +}; + +test "a prompt's run is copied, a reply's needs 16 tokens" { + const gpa = std.testing.allocator; + const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 50, 51 }; + var c = try Copy.init(gpa, &prompt, 1 << 20, K, K_REPLY); + defer c.deinit(gpa); + var out: [7]u32 = undefined; + try std.testing.expectEqual(@as(usize, 0), c.propose(&out)); + try c.append(gpa, &.{ 1, 2, 3, 4, 5, 6, 7, 8 }); + const n = c.propose(&out); + try std.testing.expectEqual(@as(usize, 7), n); + try std.testing.expectEqualSlices(u32, &.{ 9, 10, 11, 12, 50, 51, 1 }, out[0..n]); + // a run of 8 that last ran inside the reply: not enough + var d = try Copy.init(gpa, &.{ 99, 98 }, 1 << 20, K, K_REPLY); + defer d.deinit(gpa); + try d.append(gpa, &.{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 77, 1, 2, 3, 4, 5, 6, 7, 8 }); + try std.testing.expectEqual(@as(usize, 0), d.propose(&out)); +} + +test "the rate learns" { + var r: Rate = .{}; + try std.testing.expectApproxEqAbs(@as(f64, 0.8), r.chance(), 1e-9); + for (0..50) |_| r.record(7, 1); + try std.testing.expect(r.chance() < 0.3); +} diff --git a/zig/src/families/glm_cuda/decode.zig b/zig/src/families/glm_cuda/decode.zig new file mode 100644 index 0000000000..d066c50d16 --- /dev/null +++ b/zig/src/families/glm_cuda/decode.zig @@ -0,0 +1,336 @@ +//! GLM-5.3-Flash on CUDA: prompts, verify steps, commits and MTP drafts (decode.py), for any set of streams. A prompt +//! runs alone in chunks of 2048 rows; verify steps and draft steps take every stream's rows at once. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const kern = @import("kern.zig"); +const sq = @import("seq.zig"); +const forward = @import("forward.zig"); +const blocks = @import("blocks.zig"); +const dflash = @import("dflash.zig"); +const draw_mod = @import("draw.zig"); + +const Engine = forward.Engine; +const Buffers = forward.Buffers; +const Seq = sq.Seq; +const Seg = forward.Seg; +const D = kern.D; +const HEAD_N = forward.HEAD_N; + +pub const Pick = struct { tok: u32, p: f64 }; + +/// The gathered candidates -> each row's token (the larger logit; the lower id at a tie: rank 0's) and its +/// probability under the full softmax (both ranks' log-sum-exps). +fn readPicks(e: *Engine, b: *Buffers, rows: i64, out: []Pick) !void { + const d = e.ctx.d; + const n: usize = @intCast(2 * rows * 4); + const w: []i32 = @alignCast(std.mem.bytesAsSlice(i32, e.host[0 .. n * 4])); + try d.check(d.api.cuMemcpyDtoHAsync_v2(@ptrCast(w.ptr), b.cgath, n * 4, e.k.stream.handle), "candidates"); + try e.k.stream.synchronize(); + if (e.ctx.roce) |g| try g.check(); + const R: usize = @intCast(rows); + for (0..R) |r| { + const a = w[4 * r ..][0..4]; + const c = w[4 * (R + r) ..][0..4]; + const v0: f32 = @bitCast(a[1]); + const v1: f32 = @bitCast(c[1]); + const second = v1 > v0 or (v1 == v0 and c[0] < a[0]); + const l0: f64 = @as(f32, @bitCast(a[2])); + const l1: f64 = @as(f32, @bitCast(c[2])); + const top = @max(l0, l1); + const total = top + @log(@exp(l0 - top) + @exp(l1 - top)); + const v: f64 = if (second) v1 else v0; + out[r] = .{ .tok = @intCast(if (second) c[0] else a[0]), .p = @exp(v - total) }; + } +} + +/// A row's draw rule and the position its draw is keyed at (null rule: greedy, read from the candidates). +pub const Draw = struct { rule: ?lanes.Sampling, position: u64 }; + +/// The sampled rows' tokens: each rank's 64 candidates of every row and its log-sum-exps at the row's temperature +/// (fn_cand_topk), gathered, drawn on the host (draw.zig); a rule past the candidates draws over the row's gathered +/// logits. Both ranks run it alike, so their gathers pair up. +/// TF_GLM_TRACE: sampled rows' draw time (the candidates' kernel, gather and copy; the host's draws) and full rows. +pub var draw_stats: struct { gpu_ms: f64 = 0, host_ms: f64 = 0, rows: u64 = 0, full: u64 = 0, calls: u64 = 0, mismatch: u64 = 0 } = .{}; + +fn msFrom(e: *Engine, t: std.Io.Timestamp) f64 { + return @as(f64, @floatFromInt(t.durationTo(std.Io.Timestamp.now(e.io, .awake)).toNanoseconds())) / 1e6; +} + +fn drawRows(e: *Engine, b: *Buffers, rows: i64, draws: []const Draw, out: []Pick) !void { + const k = &e.k; + const t_gpu = std.Io.Timestamp.now(e.io, .awake); + var inv: [forward.STEP_ROWS]f32 = undefined; + for (draws, 0..) |dr, r| inv[r] = if (dr.rule) |rule| @floatCast(1.0 / @max(rule.temperature, 1e-6)) else 1.0; + try k.upload(b.invt, std.mem.sliceAsBytes(inv[0..@intCast(rows)])); + const n: usize = @intCast(rows * draw_mod.W); + // the candidates by radix selection (fn_cand_topk_sel); TF_GLM_TOPK_CHECK=1: the per-thread lists' kernel + // (fn_cand_topk) first, its bytes compared with the selection's (dev check: the same 64, the same sums) + const check = topk_check orelse blk: { + topk_check = std.c.getenv("TF_GLM_TOPK_CHECK") != null; + break :blk topk_check.?; + }; + if (check) { + try candTopk(e, b, rows, "fn_cand_topk"); + try k.stream.synchronize(); + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(e.cand_host.ptr), b.tk, n * 4), "candidates (check)"); + @memcpy(e.cand_host[n .. 2 * n], e.cand_host[0..n]); + } + try candTopk(e, b, rows, "fn_cand_topk_sel"); + if (check) { + try k.stream.synchronize(); + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(e.cand_host.ptr), b.tk, n * 4), "candidates (check)"); + if (!std.mem.eql(i32, e.cand_host[0..n], e.cand_host[n .. 2 * n])) { + draw_stats.mismatch += 1; + std.log.warn("glm draws: fn_cand_topk_sel differs from fn_cand_topk ({d} rows)", .{rows}); + } + } + try e.ctx.allGather(b.tk, b.gtk, n, .i32, "candidates"); + try k.stream.synchronize(); + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(e.cand_host.ptr), b.gtk, 2 * n * 4), "candidates"); + draw_stats.gpu_ms += msFrom(e, t_gpu); + const t_host = std.Io.Timestamp.now(e.io, .awake); + for (draws, 0..) |dr, r| { + const rule = dr.rule orelse continue; + const blocks_ = [2][]const i32{ e.cand_host[r * draw_mod.W ..][0..draw_mod.W], e.cand_host[n + r * draw_mod.W ..][0..draw_mod.W] }; + draw_stats.rows += 1; + out[r] = .{ .tok = switch (try draw_mod.draw(e.gpa, blocks_, rule, dr.position)) { + .token => |t| t, + .full => blk: { + draw_stats.full += 1; + break :blk try drawFull(e, b.logits + @as(u64, @intCast(@as(i64, @intCast(r)) * HEAD_N * 2)), rule, dr.position); + }, + }, .p = 0 }; + } + draw_stats.host_ms += msFrom(e, t_host); + draw_stats.calls += 1; +} + +var topk_check: ?bool = null; + +fn candTopk(e: *Engine, b: *Buffers, rows: i64, name: [:0]const u8) !void { + const f = try e.k.ext(.sample, name); + var a: cuda.Args = .{}; + a.add(b.logits); a.add(@as(i32, @intCast(HEAD_N))); a.add(@as(i32, @intCast(HEAD_N))); a.add(@as(i32, @intCast(e.w.vocab_offset))); + a.add(b.invt); a.add(b.tk); + try e.k.go(f, .{ .x = @intCast(rows) }, 256, 0, &a); +} + +fn drawFull(e: *Engine, row: u64, rule: lanes.Sampling, position: u64) !u32 { + try e.ctx.nccl.check(e.ctx.nccl.api.ncclAllGather(row, e.full.ptr, @intCast(HEAD_N), .bf16, e.ctx.comm, e.k.stream.handle), "full row"); + try e.k.stream.synchronize(); + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(e.full_host.ptr), e.full.ptr, e.full_host.len * 2), "full row"); + return draw_mod.drawFullBf16(e.gpa, e.full_host, rule, position); +} + +/// The MTP head's rows of one prompt chunk (or a resume's pending row): ids and their hidden rows at mtp_len; the +/// caller advances mtp_len. +fn mtpForward(e: *Engine, b: *Buffers, s: *Seq, next: []const u32, hidden: u64) !void { + const n: i64 = @intCast(next.len); + try forward.stageIds(e, b, next); + if (hidden != b.hin) try e.k.copy(b.hin, hidden, @intCast(n * D * 2)); + const segs = [_]Seg{.{ .s = s, .rows = n }}; + try forward.prepare(e, b, &segs, .mtp); + try blocks.mtpCompute(e, b); +} + +/// The prompt in chunks of 2048 rows from ``start`` (0, or the kept state's point: resumed), each chunk's rows +/// absorbed into the MTP cache, and the first token drawn. ``keep_at`` (0: none): the point the prompt's state is kept +/// at (the stream's kept state), cut inside its chunk when it falls inside one. +pub fn prefill(e: *Engine, s: *Seq, prompt: []const u32, start: usize, keep_at: usize) !u32 { + defer s.media = null; // attach's images serve this prefill alone + if (prompt.len == 0 or start >= prompt.len) return error.EmptyPrompt; + if (keep_at > prompt.len or (keep_at > 0 and keep_at < @max(start, 1))) return error.BadKeepPoint; + const b = &e.pbuf; + const n_kda: u64 = @intCast(e.w.nkda); + const rec_all = n_kda * @as(u64, @intCast(sq.REC * 4)); + const conv_all = n_kda * @as(u64, @intCast(3 * sq.CONV_C * 2)); + e.k.prompt = true; + defer e.k.prompt = false; + if (start == 0) { + try sq.reset(&e.k, s); + } else { // restore the kept state, then the pending row absorbed with the prompt's token at the point + if (s.snap_at != start) return error.NoKeptState; + try e.k.copy(s.recAt(s.cur, 0), s.snap_rec, rec_all); + try e.k.copy(s.conv, s.snap_conv, conv_all); + s.pos = @intCast(start); + s.mtp_len = s.snap_mtp_len; + s.mtp_drafted = 0; + try sq.ringSnap(&e.k, s, 0, @intCast(start), true); + try sq.ringSnap(&e.k, s, 1, s.snap_mtp_len, true); + if (e.w.mtp != null) { + try mtpForward(e, b, s, prompt[start .. start + 1], s.snap_pending); + s.mtp_len += 1; + } + } + // the live caches change from here: a state kept below is the only one valid after (a resume kept at its own + // point keeps the state it resumed) + if (!(start > 0 and keep_at == start)) s.snap_at = 0; + s.d_from = @intCast(start); // DFlash2's ring holds this prompt's rows from here (a kept state's were overwritten since) + var first: [1]Pick = undefined; + var at = start; + while (at < prompt.len) { + const R: usize = @min(prompt.len - at, @as(usize, @intCast(e.pbuf.rows))); + try forward.stageIds(e, b, prompt[at .. at + R]); + const point: usize = if (keep_at > at and keep_at <= at + R) keep_at - at else 0; + e.cut_at = if (point > 0 and point < R) @intCast(point) else 0; + if (e.cut_at > 0) try e.k.copy(e.cut_conv, s.conv, conv_all); // the cut's windows start from the chunk's + const segs = [_]Seg{.{ .s = s, .rows = @intCast(R) }}; + try forward.prepare(e, b, &segs, .verify); + try blocks.compute(e, b); + e.cut_at = 0; + if (e.drafter) |dr| { // the chunk's rows within the first draft's window into the drafter's context + const lo = @max(at, prompt.len -| @as(usize, @intCast(dflash.WIN - 1))); + if (lo < at + R) try dflash.ingest(dr, e, &.{.{ .s = s, .taps = dr.taps_prefill + @as(u64, @intCast((lo - at) * 5 * D * 2)), .n = at + R - lo, .pos0 = @intCast(lo) }}); + if (start == 0 and at == 0 and R == prompt.len and e.ctx.rank == 0) try dflash.dumpPrompt(dr, e, R); + } + try e.k.copy(s.last_hidden, b.fnormed + @as(u64, @intCast((R - 1) * D * 2)), @intCast(D * 2)); + if (point > 0) { // the snapshot at the kept point (before the MTP head writes its rows) + if (point < R) { + try e.k.copy(s.snap_rec, e.cut_rec, rec_all); + try e.k.copy(s.snap_conv, e.cut_conv, conv_all); + } else { + try e.k.copy(s.snap_rec, s.recAt(1 - s.cur, 0), rec_all); // the chunk's states (the parity flips below) + try e.k.copy(s.snap_conv, s.conv, conv_all); // the windows after the chunk (shifted in every KDA layer) + } + try e.k.copy(s.snap_pending, b.fnormed + @as(u64, @intCast((point - 1) * D * 2)), @intCast(D * 2)); + s.snap_mtp_len = @intCast(keep_at - 1); + } + // the first token now: the MTP absorb below writes the same logits buffer + if (at + R == prompt.len) { + try forward.candidates(e, b, 1); + if (s.sampled()) |rule| try drawRows(e, b, 1, &.{.{ .rule = rule, .position = prompt.len }}, &first); + } + if (e.w.mtp != null) { + const nxt = prompt[at + 1 .. @min(at + R + 1, prompt.len)]; + if (nxt.len > 0) { + try mtpForward(e, b, s, nxt, b.fnormed); + s.mtp_len += @intCast(nxt.len); + } + } + if (point > 0) { // the kept points' incomplete pool rows (the MTP layer's written by the absorb above) + try sq.ringSnap(&e.k, s, 0, @intCast(keep_at), false); + try sq.ringSnap(&e.k, s, 1, s.snap_mtp_len, false); + } + // a prompt chunk commits whole: its KDA layers already wrote their states (the parity flips) and windows + s.cur = 1 - s.cur; + s.pos += @intCast(R); + if (point > 0) s.snap_at = keep_at; + at += R; + } + if (s.sampled() != null) return first[0].tok; + try readPicks(e, b, 1, &first); + return first[0].tok; +} + +/// A stream's verify window: the pending token, then its drafts. +pub const Win = struct { s: *Seq, ids: []const u32 }; + +/// Every window through the model in one step (each at its stream's position); each row's token into ``out`` (rows +/// in window order). The windows stay uncommitted until ``keep``. +pub fn verify(e: *Engine, wins: []const Win, out: []Pick) !void { + var segs: [forward.MAX_STREAMS]Seg = undefined; + var ids: [forward.STEP_ROWS]u32 = undefined; + var draws: [forward.STEP_ROWS]Draw = undefined; + var R: usize = 0; + var sampled = false; + for (wins, 0..) |w, i| { + if (w.ids.len == 0 or w.ids.len > sq.MAX_ROWS) return error.WindowTooWide; + segs[i] = .{ .s = w.s, .rows = @intCast(w.ids.len) }; + w.s.hidden = e.buf.fnormed + R * @as(usize, @intCast(D)) * 2; + for (w.ids, 0..) |t, r| { + ids[R + r] = t; + const rule = w.s.sampled(); + draws[R + r] = .{ .rule = rule, .position = @intCast(w.s.pos + 1 + @as(i64, @intCast(r))) }; + if (rule != null) sampled = true; + } + R += w.ids.len; + } + const b = &e.buf; + try forward.stageIds(e, b, ids[0..R]); + try forward.prepare(e, b, segs[0..wins.len], .verify); + try forward.run(e); + try readPicks(e, b, @intCast(R), out[0..R]); + if (sampled) try drawRows(e, b, @intCast(R), draws[0..R], out[0..R]); +} + +/// Keep the stream's last verify window's first ``kept`` of ``R`` rows (the KDA states replayed when fewer than all). +pub fn keep(e: *Engine, s: *Seq, R: i64, kept: i64) !void { + if (kept < 1 or kept > R) return error.BadKeep; + const n: i64 = @intCast(e.w.nkda); + if (kept < R) try kern.kdaReplayLayers(&e.k, s.recAt(s.cur, 0), s.ks.k, s.ks.v, s.ks.g, s.ks.b, n, sq.MAX_ROWS, kept, s.recAt(1 - s.cur, 0)); + s.cur = 1 - s.cur; + try kern.convShift(&e.k, s.conv, s.proj, kept, 3 * sq.CONV_C, sq.MAX_ROWS * sq.PROJ_W, sq.PROJ_W, n, sq.CONV_C); + s.pos += kept; +} + +/// A stream's draft request: its kept rows' final-normed rows (``hidden``, D bf16 apart) and the tokens after them, +/// then up to ``count`` drafts while their probabilities' product stays at or above ``confidence`` (the first always +/// kept). ``got`` and ``out``/``probs`` receive the drafts. +pub const DraftReq = struct { s: *Seq, hidden: u64, next: []const u32, count: usize, confidence: f64, out: []u32, probs: []f64, got: usize = 0 }; + +/// Every request's rows absorbed into its stream's MTP cache in one step, then the chains, one row a stream a step, +/// each stream leaving at its own confidence cut. +pub fn draft(e: *Engine, reqs: []DraftReq) !void { + if (e.w.mtp == null) return; + const b = &e.mbuf; + var segs: [forward.MAX_STREAMS]Seg = undefined; + var ids: [forward.STEP_ROWS]u32 = undefined; + var active: [forward.MAX_STREAMS]usize = undefined; + var picks: [forward.MAX_STREAMS]Pick = undefined; + var chain: [forward.MAX_STREAMS]f64 = @splat(1.0); + var na: usize = 0; + var R: usize = 0; + for (reqs, 0..) |*q, i| { + q.got = 0; + if (q.count == 0) continue; + const s = q.s; + if (s.mtp_drafted > 0) { + s.mtp_len -= s.mtp_drafted; + s.mtp_drafted = 0; + } + segs[na] = .{ .s = s, .rows = @intCast(q.next.len) }; + @memcpy(ids[R..][0..q.next.len], q.next); + try e.k.copy(b.hin + R * @as(usize, @intCast(D)) * 2, q.hidden, q.next.len * @as(usize, @intCast(D)) * 2); + R += q.next.len; + active[na] = i; + na += 1; + } + if (na == 0) return; + try forward.stageIds(e, b, ids[0..R]); + try forward.prepare(e, b, segs[0..na], .mtp); + try forward.run(e); + for (segs[0..na]) |g| g.s.mtp_len += g.rows; + var j: usize = 0; + while (na > 0) : (j += 1) { + try readPicks(e, b, @intCast(na), picks[0..na]); + var keep_n: usize = 0; + for (0..na) |a| { + const q = &reqs[active[a]]; + const p = picks[a]; + if (j > 0 and chain[a] * p.p < q.confidence) continue; + q.out[q.got] = p.tok; + q.probs[q.got] = p.p; + q.got += 1; + chain[a] *= p.p; + if (chain[a] < q.confidence or j + 1 >= q.count) continue; + // the stream drafts on: its token and the head's row for it (row a of this step's head rows) + ids[keep_n] = p.tok; + try e.k.copy(b.hin + keep_n * @as(usize, @intCast(D)) * 2, b.fnormed + a * @as(usize, @intCast(D)) * 2, @intCast(D * 2)); + segs[keep_n] = .{ .s = q.s, .rows = 1 }; + chain[keep_n] = chain[a]; + active[keep_n] = active[a]; + keep_n += 1; + } + na = keep_n; + if (na == 0) break; + try forward.stageIds(e, b, ids[0..na]); + try forward.prepare(e, b, segs[0..na], .mtp); + try forward.run(e); + for (segs[0..na]) |g| { + g.s.mtp_len += 1; + g.s.mtp_drafted += 1; + } + } +} diff --git a/zig/src/families/glm_cuda/dflash.zig b/zig/src/families/glm_cuda/dflash.zig new file mode 100644 index 0000000000..c7e536a4e6 --- /dev/null +++ b/zig/src/families/glm_cuda/dflash.zig @@ -0,0 +1,827 @@ +//! DFlash2 block drafts (opt-in: TF_GLM_DRAFTER=dflash2, TF_GLM_DFLASH2 the checkpoint's folder). IncoAI's drafter +//! (incoai/GLM-5.3-Flash-DFlash2, CC BY-NC-ND 4.0: non-commercial use; read from its own download, never changed on +//! disk or redistributed) after z-lab/dflash's DFlash2 architecture (MIT); this implementation is our own. +//! +//! The drafter is a 5-layer Qwen3-style decoder (hidden 4096, 32/8 heads of 128, SwiGLU 12288) that drafts a whole +//! block at once: rows [anchor, MASK x 7] at positions start .. start + 7 attend, bidirectionally, to each other and +//! to the stream's context: the target's completed residuals after layers 5, 14, 24, 33 and 42 (each the mean of its +//! 4 mHC streams), concatenated, through fc and a norm, then each layer's own k and v projections (a ring of the +//! last 2048 positions). Two-tap dynamic convolutions along the block wrap both sublayers. The block's last 7 rows +//! go through the target's head; a path selector picks each draft among the row's 16 likeliest tokens by the +//! previous draft's codebook row. Drafts never change a reply: the target verifies them. +//! +//! Both ranks run the drafter on their own copy (the head is each rank's vocabulary half, its candidates gathered); +//! rank 0 alone runs the selector, and the verify frame carries its drafts to rank 1. + +const std = @import("std"); +const cuda = @import("cuda"); +const kern = @import("kern.zig"); +const forward = @import("forward.zig"); +const blocks = @import("blocks.zig"); +const draw_mod = @import("draw.zig"); +const weights_mod = @import("weights.zig"); + +pub const TAPS = [_]usize{ 5, 14, 24, 33, 42 }; +pub const BLOCK: usize = 8; +pub const MASK: u32 = 154856; +pub const WIN: i64 = 2048; +pub const LAYERS: usize = 5; +const D: usize = 4096; +const KVW: usize = 1024; +const QKV: usize = 4096 + 2 * KVW; +const FF: usize = 12288; +const RANK: usize = 256; +const TOPK: usize = 16; +const THETA: f32 = 10000; +const EPS: f32 = 1e-5; +const SCALE: f32 = 0.08838834764831845; // 128^-0.5 +const VOCAB: usize = 154880; +const MAX_M: usize = forward.MAX_STREAMS * BLOCK; // block rows at most +const MAX_CTX: usize = @intCast(forward.PREFILL_ROWS_MAX); // context rows one ingest takes at most +const CW: usize = 2 * TOPK; // a row's gathered ints: 16 ids, then 16 values' bits +const CHUNKS: usize = 19; // 4096-column chunks of a rank's head rows at most +const QH: usize = 32; +const HD: usize = 128; +const CH: usize = 16; // key chunks of the chunked attention (df_attn_part) + +fn cdiv(a: i64, b: i64) i64 { + return @divFloor(a + b - 1, b); +} + +/// Device bytes a stream's context ring takes (K and V of every layer). +pub const ring_bytes: u64 = LAYERS * 2 * @as(u64, @intCast(WIN)) * KVW * 2; + +const B16 = kern.B16; +const LayerW = struct { + in_norm: u64, post_norm: u64, qn: u64, kn: u64, a_base: u64, m_base: u64, + qkv: B16, kv: B16, o: B16, gu: B16, down: B16, a_proj: B16, m_proj: B16, +}; + +/// One stream's draft request: ``n`` rows of the last verify to take into its context first (0: the prompt's, +/// taken at prefill), the anchor (the token the block starts with, at the stream's position); rank 0 gets ``out``. +pub const Req = struct { s: *forward.Seq, n: usize, anchor: u32, out: ?[]u32 = null, probs: ?[]f64 = null, drafts: bool = true }; + +pub const Drafter = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + weights: cuda.DeviceBuffer, + bits: u8, + layers: [LAYERS]LayerW = undefined, + fc: B16 = undefined, + hproj: B16 = undefined, + hnorm: u64 = 0, + norm: u64 = 0, + pred: []u16, // the selector's codebooks, bf16 [VOCAB][RANK] (host: it reads 16 rows a draft) + succ: []u16, + scratch: cuda.DeviceBuffer, + // block buffers [MAX_M][..] + x: u64 = 0, y: u64 = 0, y2: u64 = 0, dyn: u64 = 0, qkv: u64 = 0, q: u64 = 0, kb: u64 = 0, vb: u64 = 0, att: u64 = 0, a: u64 = 0, + gu: u64 = 0, act: u64 = 0, hn: u64 = 0, g: u64 = 0, + // context buffers [MAX_CTX][..] + cin: u64 = 0, c1: u64 = 0, ctx: u64 = 0, ckv: u64 = 0, + // the target's taps: a verify step's rows and a prompt chunk's, [rows][5 D] + taps_step: u64 = 0, + taps_prefill: u64 = 0, + // tables + ids: u64 = 0, bpos: u64 = 0, rings: u64 = 0, starts: u64 = 0, froms: u64 = 0, cpos: u64 = 0, crings: u64 = 0, invt: u64 = 0, tk: u64 = 0, gtk: u64 = 0, + pv: u64 = 0, pi: u64 = 0, + pm: u64 = 0, pl: u64 = 0, pacc: u64 = 0, // the chunked attention's partials + attn_chunked: bool = true, // TF_GLM_DFLASH2_ATTN=serial: df_attn's one block a kv head (the chunked kernel otherwise) + h_ids: []i32, h_bpos: []i32, h_rings: []u64, h_starts: []i64, h_froms: []i64, h_cpos: []i64, h_crings: []u64, + cand: []i32, // both ranks' gathered candidates, [2][MAX_M][W] + gh: []f32, // the selector's hidden projections [MAX_M][RANK] + forward_ms: f64 = 0, + calls: u64 = 0, + lost: u64 = 0, // streams whose last verify rows were gone at their draft (their context restarts) + prof: bool = false, // TF_GLM_DFLASH2_PROF: each phase synchronized and timed (no graphs) + graphs: [forward.MAX_STREAMS + 1]?cuda.graph.Exec = @splat(null), // the block's pass a stream count + dump: ?[*:0]const u8 = null, // TF_GLM_DFLASH2_DUMP (development): a one-chunk prompt's taps and its first round, written there + dump_state: u8 = 0, // 0: waiting for a prompt, 1: its taps written, 2: done + dump_round: u32 = 0, // the dumped stream's rounds so far, and the one written (TF_GLM_DFLASH2_DUMP_ROUND, 10) + dump_at: u32 = 10, + tap_input: bool = false, + tap_s0: bool = false, + full_head: bool = false, // TF_GLM_DFLASH2_HEAD=full: the candidates from the verify head (8-bit) instead of the drafts' (4-bit) + temp: f64 = 1.0, // TF_GLM_DFLASH2_TEMP: the selector's softmax temperature for the drafts' chances (the lane core's depth rule) + phase_ms: [5]f64 = @splat(0), // context, layers, head and gather, selector, tables + prof_fine: bool = false, // TF_GLM_DFLASH2_PROF=2: every kernel kind of the block pass synchronized and timed + split: bool = false, // the MLP linears split across the two ranks (gate and up by columns, down by K; TF_GLM_DFLASH2_SPLIT=0: whole on each) + ff: usize = FF, // the MLP width this rank holds + fine_ms: [FINE.len]f64 = @splat(0), + fine_t: std.Io.Timestamp = undefined, + + pub fn deinit(dr: *Drafter) void { + for (&dr.graphs) |*g| if (g.*) |*ex| ex.deinit(); + dr.scratch.free(); + dr.weights.free(); + const g = dr.gpa; + g.free(dr.pred); + g.free(dr.succ); + g.free(dr.h_ids); + g.free(dr.h_bpos); + g.free(dr.h_rings); + g.free(dr.h_starts); + g.free(dr.h_froms); + g.free(dr.h_cpos); + g.free(dr.h_crings); + g.free(dr.cand); + g.free(dr.gh); + g.destroy(dr); + } +}; + +/// The block pass's kernel kinds (TF_GLM_DFLASH2_PROF=2), in launch order. +pub const FINE = [_][]const u8{ "embed", "rmsnorm", "a_proj", "conv", "qkv", "block_qkv", "attn", "o_proj", "m_proj", "gate_up", "swiglu", "down", "head_mm", "top16", "cand_gather", "hproj", "down_gather", "down_sum" }; + +fn fine(dr: *Drafter, e: *forward.Engine, which: usize) !void { + if (!dr.prof_fine) return; + try e.k.stream.synchronize(); + const now = std.Io.Timestamp.now(e.io, .awake); + dr.fine_ms[which] += @as(f64, @floatFromInt(dr.fine_t.durationTo(now).toNanoseconds())) / 1e6; + dr.fine_t = now; +} + +/// The kernel kinds' mean ms a round, logged (TF_GLM_DFLASH2_PROF=2) and cleared. +pub fn logFine(dr: *Drafter, rounds: u64) void { + if (!dr.prof_fine) return; + var buf: [512]u8 = undefined; + var w: std.Io.Writer = .fixed(&buf); + const n: f64 = @floatFromInt(@max(rounds, 1)); + for (FINE, dr.fine_ms) |name, ms| w.print(" {s} {d:.3}", .{ name, ms / n }) catch break; + std.log.info("glm dflash2 kernels (ms a round):{s}", .{w.buffered()}); + dr.fine_ms = @splat(0); +} + +/// out[m, n] = x[m, k] . w^T through the engine's own matmul (the drafter's linears packed as the dense layers are). +fn lin(e: *forward.Engine, x: u64, w: B16, out: u64, m: usize, f32_out: bool) !void { + try kern.mm(&e.k, x, w.k, w, out, f32_out, @intCast(m), e.pbuf.sk); // the prompt buffers' split-K scratch: no prefill runs meanwhile +} + +// -- loading ------------------------------------------------------------------------------------------------------- + +const Ent = struct { off: u64, n: u64 }; + +/// TF_GLM_DFLASH2_BITS: the linears packed at load to 4 bits (the default: they only propose, and lands as many as +/// 8 bits did on the ruler), 6 or 8 (a scale per 64), 0 keeps bf16. +fn bitsFromEnv() !u8 { + const v = std.c.getenv("TF_GLM_DFLASH2_BITS") orelse return 4; + const b = std.fmt.parseInt(u8, std.mem.span(v), 10) catch return error.BadDFlash2Bits; + return switch (b) { + 0, 4, 6, 8 => b, + else => error.BadDFlash2Bits, + }; +} + +const Loader = struct { + gpa: std.mem.Allocator, + io: std.Io, + file: std.Io.File, + ents: *std.StringHashMap(Ent), + buf: std.ArrayList(u8) = .empty, + words: std.ArrayList(u32) = .empty, + scales: std.ArrayList(u16) = .empty, + sub: std.ArrayList(u16) = .empty, // a slice of a stored tensor (a rank's share of a split linear) + w: *cuda.DeviceBuffer, + tmp: *cuda.DeviceBuffer, + k: *kern.K, + at: u64 = 0, + bits: u8, + + fn take(l: *Loader, n: u64) u64 { + const p = l.at; + l.at += std.mem.alignForward(u64, n, 256); + return l.w.ptr + p; + } + + fn read(l: *Loader, name: []const u8, want: u64) ![]u8 { + const e = l.ents.get(name) orelse { + std.log.err("dflash2: no {s} in the checkpoint", .{name}); + return error.DrafterTensorMissing; + }; + if (e.n != want) { + std.log.err("dflash2: {s} holds {d} bytes, not {d}", .{ name, e.n, want }); + return error.DrafterTensorShape; + } + try l.buf.resize(l.gpa, e.n); + _ = try l.file.readPositionalAll(l.io, l.buf.items, e.off); + return l.buf.items; + } + + /// A small bf16 tensor as stored. + fn raw(l: *Loader, name: []const u8, n: u64) !u64 { + const p = l.take(n); + try l.w.upload(p - l.w.ptr, try l.read(name, n)); + return p; + } + + /// The tensors ``names`` (rows ``rows`` each, ``kk`` columns) stacked into one linear, packed and tiled for kern.mm. + fn linear(l: *Loader, names: []const []const u8, rows: []const usize, kk: usize) !B16 { + var n: usize = 0; + for (rows) |r| n += r; + const bits = l.bits; + const wb: usize = n * kk * (if (bits == 0) 16 else bits) / 8; + const sb: usize = if (bits == 0) 0 else n * kk / 64 * 2; + const dst = l.take(wb); + const ds: u64 = if (sb > 0) l.take(sb) else 0; + var r0: usize = 0; + if (bits == 0) { + for (names, rows) |name, r| { + try l.tmp.upload(r0 * kk * 2, try l.read(name, r * kk * 2)); + r0 += r; + } + } else { + try l.words.resize(l.gpa, n * kk * bits / 32); + try l.scales.resize(l.gpa, n * kk / 64); + for (names, rows) |name, r| { + const src = try l.read(name, r * kk * 2); + weights_mod.quantize(bits, @alignCast(std.mem.bytesAsSlice(u16, src)), r, kk, + l.words.items[r0 * kk * bits / 32 ..][0 .. r * kk * bits / 32], l.scales.items[r0 * kk / 64 ..][0 .. r * kk / 64]); + r0 += r; + } + try l.tmp.upload(0, std.mem.sliceAsBytes(l.words.items)); + try l.tmp.upload(std.mem.alignForward(u64, wb, 256), std.mem.sliceAsBytes(l.scales.items)); + } + try kern.mmTile(l.k, bits, l.tmp.ptr, l.tmp.ptr + std.mem.alignForward(u64, wb, 256), dst, ds, @intCast(n), @intCast(kk)); + try l.k.stream.synchronize(); + return .{ .w = dst, .n = @intCast(n), .k = @intCast(kk), .s = ds, .bits = bits }; + } + + /// Rows ``row0 .. row0 + rn`` and columns ``col0 .. col0 + cn`` of each of the tensors ``names`` (rows ``rows`` + /// each, ``kk`` columns as stored), stacked into one linear of names.len * rn rows and cn columns: a rank's + /// share of a split linear, packed and tiled as ``linear`` packs a whole one. + fn linearPart(l: *Loader, names: []const []const u8, rows: []const usize, kk: usize, row0: usize, rn: usize, col0: usize, cn: usize) !B16 { + const n = names.len * rn; + const bits = l.bits; + const wb: usize = n * cn * (if (bits == 0) 16 else bits) / 8; + const sb: usize = if (bits == 0) 0 else n * cn / 64 * 2; + const dst = l.take(wb); + const ds: u64 = if (sb > 0) l.take(sb) else 0; + try l.sub.resize(l.gpa, rn * cn); + if (bits != 0) { + try l.words.resize(l.gpa, n * cn * bits / 32); + try l.scales.resize(l.gpa, n * cn / 64); + } + var r0: usize = 0; + for (names, rows) |name, r| { + if (row0 + rn > r or col0 + cn > kk) return error.DrafterTensorShape; + const src: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, try l.read(name, r * kk * 2))); + for (0..rn) |i| @memcpy(l.sub.items[i * cn ..][0..cn], src[(row0 + i) * kk + col0 ..][0..cn]); + if (bits == 0) { + try l.tmp.upload(r0 * cn * 2, std.mem.sliceAsBytes(l.sub.items)); + } else { + weights_mod.quantize(bits, l.sub.items, rn, cn, l.words.items[r0 * cn * bits / 32 ..][0 .. rn * cn * bits / 32], l.scales.items[r0 * cn / 64 ..][0 .. rn * cn / 64]); + } + r0 += rn; + } + if (bits != 0) { + try l.tmp.upload(0, std.mem.sliceAsBytes(l.words.items)); + try l.tmp.upload(std.mem.alignForward(u64, wb, 256), std.mem.sliceAsBytes(l.scales.items)); + } + try kern.mmTile(l.k, bits, l.tmp.ptr, l.tmp.ptr + std.mem.alignForward(u64, wb, 256), dst, ds, @intCast(n), @intCast(cn)); + try l.k.stream.synchronize(); + return .{ .w = dst, .n = @intCast(n), .k = @intCast(cn), .s = ds, .bits = bits }; + } +}; + +fn linearBytes(n: usize, kk: usize, bits: u8) u64 { + const wb = std.mem.alignForward(u64, n * kk * (if (bits == 0) 16 else @as(usize, bits)) / 8, 256); + return wb + if (bits == 0) 0 else std.mem.alignForward(u64, n * kk / 64 * 2, 256); +} + +/// The drafter from ``dir`` (config.json and model.safetensors, bf16), its linears packed for the engine's matmul. +pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *forward.Engine, dir: []const u8) !*Drafter { + try checkConfig(gpa, io, dir); + const bits = try bitsFromEnv(); + const d = e.ctx.d; + const path = try std.fs.path.join(gpa, &.{ dir, "model.safetensors" }); + defer gpa.free(path); + const file = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer file.close(io); + var head: [8]u8 = undefined; + _ = try file.readPositionalAll(io, &head, 0); + const hlen = std.mem.readInt(u64, &head, .little); + if (hlen > 1 << 20) return error.BadDrafterHeader; + const text = try gpa.alloc(u8, hlen); + defer gpa.free(text); + _ = try file.readPositionalAll(io, text, 8); + const hdr = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer hdr.deinit(); + var ents = std.StringHashMap(Ent).init(gpa); + defer ents.deinit(); + var it = hdr.value.object.iterator(); + while (it.next()) |x| { + if (std.mem.eql(u8, x.key_ptr.*, "__metadata__")) continue; + if (!std.mem.eql(u8, x.value_ptr.object.get("dtype").?.string, "BF16")) return error.DrafterNotBf16; + const offs = x.value_ptr.object.get("data_offsets").?.array.items; + try ents.put(x.key_ptr.*, .{ .off = 8 + hlen + @as(u64, @intCast(offs[0].integer)), .n = @intCast(offs[1].integer - offs[0].integer) }); + } + // the device's share: every linear packed (q|k|v and gate|up stacked, k|v again for the context), the norms and + // conv kernels as stored; the codebooks stay on the host + var total: u64 = 0; + const split = e.ctx.world == 2 and if (std.c.getenv("TF_GLM_DFLASH2_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; + const ff: usize = if (split) FF / 2 else FF; + const small: u64 = 6 * 4096 * 2 + 2 * 128 * 2 + 2 * 2 * 2 * D * 2; + total += LAYERS * (linearBytes(QKV, D, bits) + linearBytes(2 * KVW, D, bits) + linearBytes(D, D, bits) + linearBytes(2 * ff, D, bits) + + linearBytes(D, ff, bits) + 2 * linearBytes(KVW, D, bits) + small + 8 * 256); + total += linearBytes(D, 5 * D, bits) + linearBytes(RANK, D, bits) + 2 * D * 2 + 1024; + var wbuf = try cuda.DeviceBuffer.alloc(d, total); + errdefer wbuf.free(); + var tmp = try cuda.DeviceBuffer.alloc(d, linearBytes(2 * FF, D, 0) + 4096); // the largest linear as bf16 + defer tmp.free(); + var l: Loader = .{ .gpa = gpa, .io = io, .file = file, .ents = &ents, .w = &wbuf, .tmp = &tmp, .k = &e.k, .bits = bits }; + defer { + l.buf.deinit(gpa); + l.words.deinit(gpa); + l.scales.deinit(gpa); + l.sub.deinit(gpa); + } + const dr = try gpa.create(Drafter); + errdefer gpa.destroy(dr); + dr.* = .{ .gpa = gpa, .d = d, .weights = wbuf, .bits = bits, .split = split, .ff = ff, .pred = &.{}, .succ = &.{}, .scratch = undefined, .h_ids = undefined, + .h_bpos = undefined, .h_rings = undefined, .h_starts = undefined, .h_froms = undefined, .h_cpos = undefined, .h_crings = undefined, + .cand = undefined, .gh = undefined }; + var nb: [8][96]u8 = undefined; + for (&dr.layers, 0..) |*lw, i| { + const N = struct { + fn n(b: *[96]u8, li: usize, s: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "layers.{d}.{s}", .{ li, s }); + } + }.n; + lw.in_norm = try l.raw(try N(&nb[0], i, "input_layernorm.weight"), D * 2); + lw.post_norm = try l.raw(try N(&nb[0], i, "post_attention_layernorm.weight"), D * 2); + lw.qn = try l.raw(try N(&nb[0], i, "self_attn.q_norm.weight"), 128 * 2); + lw.kn = try l.raw(try N(&nb[0], i, "self_attn.k_norm.weight"), 128 * 2); + lw.a_base = try l.raw(try N(&nb[0], i, "attention_conv.base_kernel"), 2 * 2 * D * 2); + lw.m_base = try l.raw(try N(&nb[0], i, "mlp_conv.base_kernel"), 2 * 2 * D * 2); + const q = try N(&nb[1], i, "self_attn.q_proj.weight"); + const kq = try N(&nb[2], i, "self_attn.k_proj.weight"); + const v = try N(&nb[3], i, "self_attn.v_proj.weight"); + lw.qkv = try l.linear(&.{ q, kq, v }, &.{ D, KVW, KVW }, D); + lw.kv = try l.linear(&.{ kq, v }, &.{ KVW, KVW }, D); + lw.o = try l.linear(&.{try N(&nb[1], i, "self_attn.o_proj.weight")}, &.{D}, D); + const gate = try N(&nb[1], i, "mlp.gate_proj.weight"); + const up = try N(&nb[2], i, "mlp.up_proj.weight"); + const down = try N(&nb[3], i, "mlp.down_proj.weight"); + if (split) { // this rank's half of the MLP width: gate and up rows, down columns + lw.gu = try l.linearPart(&.{ gate, up }, &.{ FF, FF }, D, e.ctx.rank * ff, ff, 0, D); + lw.down = try l.linearPart(&.{down}, &.{D}, FF, 0, D, e.ctx.rank * ff, ff); + } else { + lw.gu = try l.linear(&.{ gate, up }, &.{ FF, FF }, D); + lw.down = try l.linear(&.{down}, &.{D}, FF); + } + lw.a_proj = try l.linear(&.{try N(&nb[1], i, "attention_conv.kernel_projection.weight")}, &.{KVW}, D); + lw.m_proj = try l.linear(&.{try N(&nb[1], i, "mlp_conv.kernel_projection.weight")}, &.{KVW}, D); + } + dr.fc = try l.linear(&.{"fc.weight"}, &.{D}, 5 * D); + dr.hproj = try l.linear(&.{"candidate_selector.hidden_projection.weight"}, &.{RANK}, D); + dr.hnorm = try l.raw("hidden_norm.weight", D * 2); + dr.norm = try l.raw("norm.weight", D * 2); + if (l.at > total) return error.DrafterLayout; + dr.pred = try gpa.alloc(u16, VOCAB * RANK); + @memcpy(std.mem.sliceAsBytes(dr.pred), try l.read("candidate_selector.predecessor_codebook", VOCAB * RANK * 2)); + dr.succ = try gpa.alloc(u16, VOCAB * RANK); + @memcpy(std.mem.sliceAsBytes(dr.succ), try l.read("candidate_selector.successor_codebook", VOCAB * RANK * 2)); + try carveScratch(dr, d); + dr.h_ids = try gpa.alloc(i32, MAX_M); + dr.h_bpos = try gpa.alloc(i32, MAX_M); + dr.h_rings = try gpa.alloc(u64, forward.MAX_STREAMS); + dr.h_starts = try gpa.alloc(i64, forward.MAX_STREAMS); + dr.h_froms = try gpa.alloc(i64, forward.MAX_STREAMS); + dr.h_cpos = try gpa.alloc(i64, MAX_CTX); + dr.h_crings = try gpa.alloc(u64, MAX_CTX); + dr.cand = try gpa.alloc(i32, 2 * MAX_M * CW); + dr.gh = try gpa.alloc(f32, MAX_M * RANK); + dr.prof = std.c.getenv("TF_GLM_DFLASH2_PROF") != null; + if (std.c.getenv("TF_GLM_DFLASH2_ATTN")) |v| dr.attn_chunked = !std.mem.eql(u8, std.mem.span(v), "serial"); + dr.prof_fine = if (std.c.getenv("TF_GLM_DFLASH2_PROF")) |v| std.mem.eql(u8, std.mem.span(v), "2") else false; + dr.dump = std.c.getenv("TF_GLM_DFLASH2_DUMP"); + if (std.c.getenv("TF_GLM_DFLASH2_TAP")) |v| { + dr.tap_input = std.mem.eql(u8, std.mem.span(v), "input"); + dr.tap_s0 = std.mem.eql(u8, std.mem.span(v), "s0"); + } + if (std.c.getenv("TF_GLM_DFLASH2_DUMP_ROUND")) |v| dr.dump_at = std.fmt.parseInt(u32, std.mem.span(v), 10) catch 10; + if (std.c.getenv("TF_GLM_DFLASH2_HEAD")) |v| dr.full_head = std.mem.eql(u8, std.mem.span(v), "full"); + if (std.c.getenv("TF_GLM_DFLASH2_TEMP")) |v| dr.temp = std.fmt.parseFloat(f64, std.mem.span(v)) catch 1.0; + const inv: [MAX_M]f32 = @splat(1.0); + try dr.scratch.upload(dr.invt - dr.scratch.ptr, std.mem.sliceAsBytes(&inv)); + std.log.info("dflash2: drafter loaded from {s}, linears at {d} bits ({d:.2} GiB on the device, {d} MiB of codebooks on the host)", + .{ dir, bits, @as(f64, @floatFromInt(l.at)) / (1 << 30), (2 * VOCAB * RANK * 2) >> 20 }); + return dr; +} + +fn carveScratch(dr: *Drafter, d: *const cuda.Driver) !void { + var at: u64 = 0; + const M: u64 = MAX_M; + const C: u64 = MAX_CTX + 16; + const fields = .{ + .{ "x", M * D * 2 }, .{ "y", M * D * 2 }, .{ "y2", M * D * 2 }, .{ "dyn", M * KVW * 2 }, + .{ "qkv", M * QKV * 2 }, .{ "q", M * D * 2 }, .{ "kb", M * KVW * 2 }, .{ "vb", M * KVW * 2 }, + .{ "att", M * D * 2 }, .{ "a", M * D * 2 }, .{ "gu", M * 2 * FF * 2 }, .{ "act", M * FF * 2 }, + .{ "hn", M * D * 2 }, .{ "g", M * RANK * 4 }, .{ "cin", C * 5 * D * 2 }, .{ "c1", C * D * 2 }, + .{ "ctx", C * D * 2 }, .{ "ckv", C * 2 * KVW * 2 }, + .{ "taps_step", @as(u64, @intCast(forward.STEP_ROWS)) * 5 * D * 2 }, .{ "taps_prefill", @as(u64, @intCast(forward.PREFILL_ROWS_MAX)) * 5 * D * 2 }, + .{ "ids", M * 4 }, .{ "bpos", M * 4 }, .{ "rings", forward.MAX_STREAMS * 8 }, .{ "starts", forward.MAX_STREAMS * 8 }, + .{ "froms", forward.MAX_STREAMS * 8 }, .{ "cpos", C * 8 }, .{ "crings", C * 8 }, .{ "invt", M * 4 }, + .{ "tk", M * CW * 4 }, .{ "gtk", 2 * M * CW * 4 }, + .{ "pv", M * CHUNKS * TOPK * 4 }, .{ "pi", M * CHUNKS * TOPK * 4 }, + .{ "pm", M * QH * CH * 4 }, .{ "pl", M * QH * CH * 4 }, .{ "pacc", M * QH * CH * HD * 4 }, + }; + inline for (fields) |f| { + @field(dr, f[0]) = at; + at += std.mem.alignForward(u64, f[1], 256); + } + dr.scratch = try cuda.DeviceBuffer.alloc(d, at); + inline for (fields) |f| @field(dr, f[0]) += dr.scratch.ptr; + try d.check(d.api.cuMemsetD8_v2(dr.scratch.ptr, 0, at), "zero drafter scratch"); +} + +/// The checkpoint is the one this was written for (its shape and DFlash2 settings). +fn checkConfig(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 text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(text); + const p = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer p.deinit(); + const o = p.value.object; + const df = (o.get("dflash_config") orelse return error.NotDFlash2).object; + const ok = o.get("hidden_size").?.integer == D and o.get("num_hidden_layers").?.integer == LAYERS and + o.get("num_attention_heads").?.integer == 32 and o.get("num_key_value_heads").?.integer == 8 and + o.get("intermediate_size").?.integer == FF and o.get("vocab_size").?.integer == VOCAB and + df.get("block_size").?.integer == BLOCK and df.get("mask_token_id").?.integer == MASK and + df.get("selector_top_k").?.integer == TOPK and df.get("selector_rank").?.integer == RANK and + df.get("conv_kernel_size").?.integer == 2 and df.get("conv_group_size").?.integer == 16; + if (!ok) return error.DrafterConfigUnexpected; + const taps = df.get("target_layer_ids").?.array.items; + if (taps.len != TAPS.len) return error.DrafterConfigUnexpected; + for (taps, TAPS) |t, w| if (t.integer != w) return error.DrafterConfigUnexpected; +} + +// -- the target's side --------------------------------------------------------------------------------------------- + +/// After target layer ``li``: its rows' 4-stream means into the taps (a verify step's or a prompt chunk's). +/// TF_GLM_DFLASH2_TAP (development): `mean` (the default), `input` (layer k's input: k - 1's output), `s0` (stream 0). +pub fn tap(dr: *const Drafter, k: *kern.K, x: u64, taps: u64, li: usize, R: i64) !void { + const at = li + @intFromBool(dr.tap_input); + const slot: i32 = for (TAPS, 0..) |t, i| { + if (t == at) break @intCast(i); + } else return; + var a: cuda.Args = .{}; + a.add(x); a.add(taps); a.add(slot); a.add(@as(i32, @intFromBool(dr.tap_s0))); + try k.go(try k.ext(.dflash, "df_tap"), .{ .x = @intCast(R), .y = @intCast(D / 1024) }, 256, 0, &a); +} + +pub fn isTap(dr: *const Drafter, li: usize) bool { + const at = li + @intFromBool(dr.tap_input); + for (TAPS) |t| if (t == at) return true; + return false; +} + +/// Context runs: ``n`` tap rows from ``taps`` (consecutive) at positions pos0 .. pos0 + n, into ``s``'s ring. +pub const Run = struct { s: *forward.Seq, taps: u64, n: usize, pos0: i64 }; + +/// The runs' rows through fc, the hidden norm and every layer's k and v, into their rings (and synchronized: the +/// host tables are reused). +/// TF_GLM_DFLASH2_DUMP: a whole one-chunk prompt's taps ([n][5 D] bf16, positions 0 .. n), once. +pub fn dumpPrompt(dr: *Drafter, e: *forward.Engine, n: usize) !void { + if (dr.dump == null or dr.dump_state == 2 or n > MAX_CTX) return; // the latest prompt until its first round is written + try e.k.stream.synchronize(); + const buf = try dr.gpa.alloc(u8, n * 5 * D * 2); + defer dr.gpa.free(buf); + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(buf.ptr), dr.taps_prefill, buf.len), "dump taps"); + try writeFile(dr, "taps.bin", buf); + try writeFile(dr, "ingest.bin", ""); + dr.dump_state = 1; + dr.dump_round = 0; +} + +/// While dumping: the ingest's rows (their positions, then their taps) appended to ingest.bin. +fn dumpIngest(dr: *Drafter, e: *forward.Engine, N: usize) !void { + try e.k.stream.synchronize(); + const buf = try dr.gpa.alloc(u8, N * 5 * D * 2); + defer dr.gpa.free(buf); + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(buf.ptr), dr.cin, buf.len), "dump ingest"); + var pb: [512]u8 = undefined; + const path = try std.fmt.bufPrintSentinel(&pb, "{s}/ingest.bin", .{std.mem.span(dr.dump.?)}, 0); + const f = std.c.fopen(path, "ab") orelse return error.DumpOpen; + defer _ = std.c.fclose(f); + const hdr = [2]i64{ @intCast(N), 0 }; + _ = std.c.fwrite(std.mem.asBytes(&hdr).ptr, 1, 16, f); + _ = std.c.fwrite(std.mem.sliceAsBytes(dr.h_cpos[0..N]).ptr, 1, N * 8, f); + _ = std.c.fwrite(buf.ptr, 1, buf.len, f); +} + +fn writeFile(dr: *Drafter, name: []const u8, bytes: []const u8) !void { + var pb: [512]u8 = undefined; + const path = try std.fmt.bufPrintSentinel(&pb, "{s}/{s}", .{ std.mem.span(dr.dump.?), name }, 0); + const f = std.c.fopen(path, "wb") orelse return error.DumpOpen; + defer _ = std.c.fclose(f); + if (std.c.fwrite(bytes.ptr, 1, bytes.len, f) != bytes.len) return error.DumpWrite; +} + +pub fn ingest(dr: *Drafter, e: *forward.Engine, runs: []const Run) !void { + const k = &e.k; + var N: usize = 0; + for (runs) |r| N += r.n; + if (N == 0) return; + if (N > MAX_CTX) return error.DrafterContextTooWide; + var row: usize = 0; + for (runs) |r| { + if (r.n == 0) continue; + try k.copy(dr.cin + row * 5 * D * 2, r.taps, r.n * 5 * D * 2); + for (0..r.n) |j| { + dr.h_cpos[row + j] = r.pos0 + @as(i64, @intCast(j)); + dr.h_crings[row + j] = r.s.draft; + } + row += r.n; + } + if (dr.dump != null and dr.dump_state == 1) try dumpIngest(dr, e, N); + try k.upload(dr.cpos, std.mem.sliceAsBytes(dr.h_cpos[0..N])); + try k.upload(dr.crings, std.mem.sliceAsBytes(dr.h_crings[0..N])); + try lin(e, dr.cin, dr.fc, dr.c1, N, false); + try kern.rmsnorm(k, dr.c1, D, dr.hnorm, dr.ctx, D, null, @intCast(N), D); + for (dr.layers, 0..) |l, li| { + try lin(e, dr.ctx, l.kv, dr.ckv, N, false); + var a: cuda.Args = .{}; + a.add(dr.ckv); a.add(l.kn); a.add(dr.cpos); a.add(dr.crings); a.add(@as(i32, @intCast(li))); a.add(EPS); a.add(THETA); + try k.go(try k.ext(.dflash, "df_ctx_kv"), .{ .x = @intCast(N), .y = 8 }, 64, 0, &a); + } + try k.stream.synchronize(); +} + +// -- drafting ------------------------------------------------------------------------------------------------------ + +/// Every request's block drafted on this rank (both ranks, in step: the head's candidates gather); rank 0 (the +/// requests' ``out``) gets 7 drafts and their chances. +/// Every request's kept rows into its context; then the block pass for the ones that draft (``drafts``: the others +/// draft with the MTP head this round and only keep their context current). +pub fn draft(dr: *Drafter, e: *forward.Engine, all: []const Req) !void { + const k = &e.k; + if (all.len == 0) return; + if (all.len > forward.MAX_STREAMS) return error.TooManyStreams; + const t0 = std.Io.Timestamp.now(e.io, .awake); + var tp = t0; + // the last verify's kept rows into each stream's context + var runs: [forward.MAX_STREAMS]Run = undefined; + var nr: usize = 0; + for (all) |q| { + if (q.n == 0) continue; + const s = q.s; + if (s.v_step != e.verifies or q.n > s.v_rows) { // its rows were overwritten (or never verified): restart its context + s.d_from = s.pos; + dr.lost += 1; + continue; + } + runs[nr] = .{ .s = s, .taps = dr.taps_step + @as(u64, @intCast(s.v_row0)) * 5 * D * 2, .n = q.n, .pos0 = s.v_pos0 }; + nr += 1; + } + try ingest(dr, e, runs[0..nr]); + try phase(dr, e, &tp, 0); + var sub: [forward.MAX_STREAMS]Req = undefined; + var S: usize = 0; + for (all) |q| if (q.drafts) { + sub[S] = q; + S += 1; + }; + if (S == 0) return; + const reqs = sub[0..S]; + // the block: [anchor, MASK x 7] a stream, at positions pos .. pos + 7 + const M = S * BLOCK; + if (dr.split and M > e.buf.rows) return error.DrafterRowsExceedBuffers; // the split's partials use the decode buffers + for (reqs, 0..) |q, i| { + dr.h_rings[i] = q.s.draft; + dr.h_starts[i] = q.s.pos; + dr.h_froms[i] = q.s.d_from; + for (0..BLOCK) |t| { + dr.h_ids[i * BLOCK + t] = @intCast(if (t == 0) q.anchor else MASK); + dr.h_bpos[i * BLOCK + t] = @intCast(q.s.pos + @as(i64, @intCast(t))); + } + } + try k.upload(dr.ids, std.mem.sliceAsBytes(dr.h_ids[0..M])); + try k.upload(dr.bpos, std.mem.sliceAsBytes(dr.h_bpos[0..M])); + try k.upload(dr.rings, std.mem.sliceAsBytes(dr.h_rings[0..S])); + try k.upload(dr.starts, std.mem.sliceAsBytes(dr.h_starts[0..S])); + try k.upload(dr.froms, std.mem.sliceAsBytes(dr.h_froms[0..S])); + try phase(dr, e, &tp, 4); + const rank0 = reqs[0].out != null; + const nw = M * CW; + if (dr.prof or !e.use_graphs or e.sync_each) { + try blockBody(dr, e, S, rank0, &tp); + } else if (dr.graphs[S]) |ex| { + try ex.launchOn(k.stream); + } else { // the block's pass for S streams, captured once (its tables at fixed places, uploaded above) + try cuda.graph.beginCapture(k.stream, .thread_local); + blockBody(dr, e, S, rank0, &tp) catch |err| { + if (cuda.graph.endCapture(k.stream)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return err; + }; + var g = try cuda.graph.endCapture(k.stream); + defer g.deinit(); + dr.graphs[S] = try g.instantiate(); + try dr.graphs[S].?.launchOn(k.stream); + } + try k.stream.synchronize(); + try phase(dr, e, &tp, 2); + if (!rank0 and dr.prof_fine and dr.calls % 256 == 255) logFine(dr, 256); // rank 1 logs its own table (rank 0's goes with the trace) + if (rank0) { + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(dr.cand.ptr), dr.gtk, 2 * nw * 4), "dflash2 candidates"); + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(dr.gh.ptr), dr.g, M * RANK * 4), "dflash2 projections"); + for (reqs, 0..) |q, i| select(dr, i, M, q); + if (dr.dump != null and dr.dump_state == 1 and S == 1) dr.dump_round += 1; + if (dr.dump != null and dr.dump_state == 1 and S == 1 and dr.dump_round == dr.dump_at) { + const hn = try dr.gpa.alloc(u8, M * D * 2); + defer dr.gpa.free(hn); + try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(hn.ptr), dr.hn, hn.len), "dump hidden"); + try writeFile(dr, "hidden.bin", hn); + try writeFile(dr, "cand.bin", std.mem.sliceAsBytes(dr.cand[0 .. 2 * nw])); + try writeFile(dr, "proj.bin", std.mem.sliceAsBytes(dr.gh[0 .. M * RANK])); + var meta: [64]u8 = undefined; + const q = reqs[0]; + const text = try std.fmt.bufPrint(&meta, "{d} {d}", .{ q.anchor, q.s.pos }); + try writeFile(dr, "meta.txt", text); + try writeFile(dr, "drafts.bin", std.mem.sliceAsBytes(q.out.?[0 .. BLOCK - 1])); + try writeFile(dr, "probs.bin", std.mem.sliceAsBytes(q.probs.?[0 .. BLOCK - 1])); + dr.dump_state = 2; + std.log.info("dflash2: dumped a prompt's taps and its first round to {s}", .{std.mem.span(dr.dump.?)}); + } + } + try phase(dr, e, &tp, 3); + dr.forward_ms += @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(e.io, .awake)).toNanoseconds())) / 1e6; + dr.calls += 1; +} + +/// The block's pass for S streams over the uploaded tables: embed, the 5 layers, the final norm, the draft head's 16 +/// best a row gathered from both ranks, and (rank 0) the selector's projections. Capturable. +fn blockBody(dr: *Drafter, e: *forward.Engine, S: usize, rank0: bool, tp: *std.Io.Timestamp) !void { + const k = &e.k; + const M = S * BLOCK; + const nw = M * CW; + if (dr.prof_fine) dr.fine_t = std.Io.Timestamp.now(e.io, .awake); + try kern.embed(k, dr.ids, e.w.embed, dr.x, @intCast(M), 1); + try fine(dr, e, 0); + const R: i64 = @intCast(M); + const grid: cuda.Dim3 = .{ .x = @intCast(M), .y = @intCast(D / 1024) }; + for (dr.layers, 0..) |l, li| { + // attention, wrapped in the attention conv + try kern.rmsnorm(k, dr.x, D, l.in_norm, dr.y, D, null, R, D); + try fine(dr, e, 1); + try lin(e, dr.y, l.a_proj, dr.dyn, M, false); + try fine(dr, e, 2); + try conv(k, dr.y, dr.dyn, 0, l.a_base, dr.y2, false, grid); + try fine(dr, e, 3); + try lin(e, dr.y2, l.qkv, dr.qkv, M, false); + try fine(dr, e, 4); + { + var a: cuda.Args = .{}; + a.add(dr.qkv); a.add(@as(i32, @intCast(QKV))); a.add(l.qn); a.add(l.kn); a.add(dr.bpos); a.add(EPS); a.add(THETA); + a.add(dr.q); a.add(dr.kb); a.add(dr.vb); + try k.go(try k.ext(.dflash, "df_block_qkv"), .{ .x = @intCast(M), .y = 32 + 8 }, 64, 0, &a); + } + try fine(dr, e, 5); + if (dr.attn_chunked) { + var a: cuda.Args = .{}; + a.add(dr.q); a.add(dr.kb); a.add(dr.vb); a.add(dr.rings); a.add(dr.starts); a.add(dr.froms); a.add(@as(i32, @intCast(li))); + a.add(SCALE); a.add(dr.pm); a.add(dr.pl); a.add(dr.pacc); + try k.go(try k.ext(.dflash, "df_attn_part"), .{ .x = @intCast(S), .y = 8, .z = CH }, 256, 0, &a); + var b: cuda.Args = .{}; + b.add(dr.pm); b.add(dr.pl); b.add(dr.pacc); b.add(dr.att); + try k.go(try k.ext(.dflash, "df_attn_merge"), .{ .x = @intCast(M), .y = QH }, HD, 0, &b); + } else { + var a: cuda.Args = .{}; + a.add(dr.q); a.add(dr.kb); a.add(dr.vb); a.add(dr.rings); a.add(dr.starts); a.add(dr.froms); a.add(@as(i32, @intCast(li))); + a.add(SCALE); a.add(dr.att); + try k.go(try k.ext(.dflash, "df_attn"), .{ .x = @intCast(S), .y = 8 }, 256, 0, &a); + } + try fine(dr, e, 6); + try lin(e, dr.att, l.o, dr.a, M, false); + try fine(dr, e, 7); + try conv(k, dr.a, dr.dyn, 1, l.a_base + 2 * D * 2, dr.x, true, grid); + try fine(dr, e, 3); + // the MLP, wrapped in the MLP conv + try kern.rmsnorm(k, dr.x, D, l.post_norm, dr.y, D, null, R, D); + try fine(dr, e, 1); + try lin(e, dr.y, l.m_proj, dr.dyn, M, false); + try fine(dr, e, 8); + try conv(k, dr.y, dr.dyn, 0, l.m_base, dr.y2, false, grid); + try fine(dr, e, 3); + try lin(e, dr.y2, l.gu, dr.gu, M, false); + try fine(dr, e, 9); + { + var a: cuda.Args = .{}; + a.add(dr.gu); a.add(dr.act); a.add(@as(i32, @intCast(dr.ff))); + try k.go(try k.ext(.dflash, "df_swiglu"), .{ .x = @intCast(M), .y = @intCast(dr.ff / 1024) }, 256, 0, &a); + } + try fine(dr, e, 10); + if (dr.split) { // each rank's half of down's K as fp32 partials, gathered and summed + const b = &e.buf; + try lin(e, dr.act, l.down, b.part, M, true); + try fine(dr, e, 11); + try e.ctx.allGather(b.part, b.gath, M * D, .f32, "dflash2 down"); + try fine(dr, e, 16); + var a: cuda.Args = .{}; + a.add(b.gath); a.add(dr.a); a.add(@as(i32, @intCast(M))); + try k.go(try k.ext(.dflash, "df_sum2"), grid, 256, 0, &a); + try fine(dr, e, 17); + } else { + try lin(e, dr.act, l.down, dr.a, M, false); + try fine(dr, e, 11); + } + try conv(k, dr.a, dr.dyn, 1, l.m_base + 2 * D * 2, dr.x, true, grid); + try fine(dr, e, 3); + } + try kern.rmsnorm(k, dr.x, D, dr.norm, dr.hn, D, null, R, D); + try fine(dr, e, 1); + try phase(dr, e, tp, 1); + // every row through the head (a block's first row is the anchor's: the host skips it) and the selector's projection + // the drafts' head: the MTP drafts' packed copy (they only propose), each rank's 16 best a row + try blocks.mm(e, &e.mbuf, dr.hn, D, if (dr.full_head) e.w.head else e.w.draft_head, e.mbuf.logits, false, R); + try fine(dr, e, 12); + { + const chunks = cdiv(forward.HEAD_N, 4096); + var a: cuda.Args = .{}; + a.add(e.mbuf.logits); a.add(@as(i32, @intCast(forward.HEAD_N))); a.add(@as(i32, @intCast(e.w.vocab_offset))); a.add(dr.pv); a.add(dr.pi); + try k.go(try k.ext(.dflash, "df_top16_part"), .{ .x = @intCast(M), .y = @intCast(chunks) }, 256, 0, &a); + var b: cuda.Args = .{}; + b.add(dr.pv); b.add(dr.pi); b.add(@as(i32, @intCast(chunks))); b.add(dr.tk); + try k.go(try k.ext(.dflash, "df_top16_merge"), .{ .x = @intCast(M) }, 256, 0, &b); + } + try fine(dr, e, 13); + try e.ctx.allGather(dr.tk, dr.gtk, nw, .i32, "dflash2 candidates"); + try fine(dr, e, 14); + if (rank0) try lin(e, dr.hn, dr.hproj, dr.g, M, true); + try fine(dr, e, 15); +} + +fn phase(dr: *Drafter, e: *forward.Engine, t: *std.Io.Timestamp, which: usize) !void { + if (!dr.prof) return; + try e.k.stream.synchronize(); + const now = std.Io.Timestamp.now(e.io, .awake); + dr.phase_ms[which] += @as(f64, @floatFromInt(t.durationTo(now).toNanoseconds())) / 1e6; + t.* = now; +} + +fn conv(k: *kern.K, z: u64, dyn: u64, which: i32, base: u64, out: u64, add: bool, grid: cuda.Dim3) !void { + var a: cuda.Args = .{}; + a.add(z); a.add(dyn); a.add(which); a.add(base); a.add(out); a.add(@as(i32, @intFromBool(add))); + try k.go(try k.ext(.dflash, "df_conv"), grid, 256, 0, &a); +} + +const Cand = struct { v: f32, id: u32 }; + +fn bf(x: u16) f32 { + return @bitCast(@as(u32, x) << 16); +} + +/// Stream ``i``'s 7 drafts: at each row, the 16 likeliest of both ranks' candidates, scored by the head's logit plus +/// the previous draft's predecessor row, the row's projection and the candidate's successor row; greedy, its +/// softmax chance kept. +fn select(dr: *Drafter, i: usize, M: usize, q: Req) void { + const out = q.out.?; + const probs = q.probs.?; + const W = CW; + const T = TOPK; + var prev: usize = q.anchor; + for (1..BLOCK) |t| { + const r = i * BLOCK + t; + var c: [TOPK]Cand = undefined; + var n: usize = 0; + for (0..2) |rk| { + const blk = dr.cand[rk * M * W + r * W ..][0..W]; + for (0..T) |j| { + if (blk[j] < 0) continue; + const cand: Cand = .{ .v = @bitCast(blk[T + j]), .id = @intCast(blk[j]) }; + if (n < TOPK) { + c[n] = cand; + n += 1; + continue; + } + var lo: usize = 0; + for (1..TOPK) |m| if (c[m].v < c[lo].v) { + lo = m; + }; + if (cand.v > c[lo].v) c[lo] = cand; + } + } + const g = dr.gh[r * RANK ..][0..RANK]; + const P = dr.pred[prev * RANK ..][0..RANK]; + var pg: [RANK]f32 = undefined; + for (&pg, P, g) |*o, p, x| o.* = bf(p) * x; + var best: usize = 0; + var sc: [TOPK]f32 = undefined; + for (0..n) |m| { + const Sc = dr.succ[@as(usize, c[m].id) * RANK ..][0..RANK]; + var dot: f32 = 0; + for (pg, Sc) |a, b| dot += a * bf(b); + sc[m] = c[m].v + dot; + if (sc[m] > sc[best]) best = m; + } + var z: f64 = 0; + for (sc[0..n]) |x| z += @exp(@as(f64, x - sc[best]) / dr.temp); + out[t - 1] = c[best].id; + probs[t - 1] = 1.0 / z; + prev = c[best].id; + } +} diff --git a/zig/src/families/glm_cuda/draw.zig b/zig/src/families/glm_cuda/draw.zig new file mode 100644 index 0000000000..4986ec0a62 --- /dev/null +++ b/zig/src/families/glm_cuda/draw.zig @@ -0,0 +1,171 @@ +//! A sampled row's token from both ranks' gathered candidates (fn_cand_topk): each rank's 64 largest logits, global +//! ids, its log-sum-exp and its log-sum-exp at the row's temperature. The merged 128 hold the vocabulary's top 64 +//! exactly, so a rule within them draws as lanes.sampling.choose over the whole vocabulary would: top_k 1..64 by +//! choose itself (only the top k take part), a wider rule when its nucleus (top_p) or floor (min_p) ends within the +//! top 64, keyed the same way (seed, position, global id). A rule that reaches past them asks for the full row. + +const std = @import("std"); +const lanes = @import("lanes"); +const smp = lanes.sampling; + +pub const TOPK = 64; +pub const W = 2 * TOPK + 4; // a row's ints in the gathered block: ids, values, lse, lse at the temperature, spare + +pub const Outcome = union(enum) { token: u32, full }; + +const Cand = struct { value: f64, id: u64 }; + +fn desc(_: void, a: Cand, b: Cand) bool { + return a.value > b.value or (a.value == b.value and a.id < b.id); +} + +fn logAddExp(a: f64, b: f64) f64 { + const m = @max(a, b); + if (m == -std.math.inf(f64)) return m; + return m + @log(@exp(a - m) + @exp(b - m)); +} + +/// ``blocks``: the two ranks' W-int rows (rank 0 first). ``position``: the draw's key. +pub fn draw(gpa: std.mem.Allocator, blocks: [2][]const i32, s: smp.Sampling, position: u64) !Outcome { + var cands: [2 * TOPK]Cand = undefined; + var n: usize = 0; + var lse_t: f64 = -std.math.inf(f64); + for (blocks) |b| { + for (0..TOPK) |k| { + if (b[k] < 0) continue; + cands[n] = .{ .value = @as(f32, @bitCast(b[TOPK + k])), .id = @intCast(b[k]) }; + n += 1; + } + lse_t = logAddExp(lse_t, @as(f32, @bitCast(b[2 * TOPK + 1]))); + } + std.mem.sort(Cand, cands[0..n], {}, desc); + const exact = @min(n, TOPK); + if (s.top_k >= 1 and s.top_k <= TOPK) { + const k = @min(@as(usize, s.top_k), exact); + var values: [TOPK]f64 = undefined; + var ids: [TOPK]u64 = undefined; + for (cands[0..k], 0..) |c, i| { + values[i] = c.value; + ids[i] = c.id; + } + return .{ .token = @intCast(try smp.choose(gpa, values[0..k], ids[0..k], position, s)) }; + } + // the whole vocabulary's rule (top_k 0 or past the candidates): its cut must end within the exact top 64 + const t = @max(s.temperature, 1e-6); + var keep: usize = std.math.maxInt(usize); + if (s.top_p > 0.0 and s.top_p < 1.0) { + var cum: f64 = 0; + var below: usize = 0; + for (cands[0..exact]) |c| { + cum += @exp(c.value / t - lse_t); + if (cum < s.top_p) below += 1 else break; + } + if (below >= exact) return .full; + keep = below + 1; + } + if (s.min_p > 0.0) { + const floor = cands[0].value / t + s.minLog(); + var m: usize = 0; + for (cands[0..exact]) |c| { + if (c.value / t >= floor) m += 1; + } + if (m >= exact and keep > exact) return .full; + keep = @min(keep, m); + } + if (s.top_k > TOPK) keep = @min(keep, @as(usize, s.top_k)); + if (keep > exact) return .full; + var best: usize = 0; + var best_score = -std.math.inf(f64); + for (cands[0..keep], 0..) |c, j| { + const gumbel = -@log(-@log(smp.uniform(s.seed, position, c.id))); + const score = c.value / t + gumbel; + if (j == 0 or score > best_score) { + best = j; + best_score = score; + } + } + return .{ .token = @intCast(cands[best].id) }; +} + +/// ``drawFull`` over a whole gathered bf16 row (raw bits by global id): the same order by a counting sort over the +/// 65,536 bf16 values (descending, ids ascending within a value; -0 with +0), then the same rule. +pub fn drawFullBf16(gpa: std.mem.Allocator, row: []const u16, s: smp.Sampling, position: u64) !u32 { + const counts = try gpa.alloc(u32, 1 << 16); + defer gpa.free(counts); + @memset(counts, 0); + for (row) |u| counts[0xffff - orderKey(u)] += 1; + var at: u32 = 0; + for (counts) |*c| { // counts by descending value become start offsets + const n = c.*; + c.* = at; + at += n; + } + const order = try gpa.alloc(smp.Candidate, row.len); + defer gpa.free(order); + for (row, 0..) |u, id| { + const slot = &counts[0xffff - orderKey(u)]; + order[slot.*] = .{ .value = @as(f32, @bitCast(@as(u32, u) << 16)), .id = id }; + slot.* += 1; + } + return @intCast(try smp.chooseOrdered(gpa, order, position, s)); +} + +/// bf16 bits to a key that orders as the value does (-0 as +0). +fn orderKey(raw: u16) u16 { + const u: u16 = if (raw == 0x8000) 0 else raw; + return if (u & 0x8000 != 0) ~u else u | 0x8000; +} + +/// The full rule over a whole gathered row (``values``: the vocabulary's logits by global id). +pub fn drawFull(gpa: std.mem.Allocator, values: []const f64, s: smp.Sampling, position: u64) !u32 { + const ids = try gpa.alloc(u64, values.len); + defer gpa.free(ids); + for (ids, 0..) |*x, i| x.* = i; + return @intCast(try smp.choose(gpa, values, ids, position, s)); +} + +test "a top_k rule draws the same token from the merged candidates as choose over a whole row" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(7); + const r = prng.random(); + const V = 2000; + var row: [V]f64 = undefined; + for (&row) |*x| x.* = @floatCast(@as(f32, @floatCast(r.floatNorm(f64) * 3))); + // each rank's half: its top 64 as the kernel writes them + var blocks: [2][W]i32 = undefined; + for (0..2) |rk| { + var c: [V / 2]Cand = undefined; + for (0..V / 2) |i| c[i] = .{ .value = row[rk * V / 2 + i], .id = rk * V / 2 + i }; + std.mem.sort(Cand, &c, {}, desc); + for (0..TOPK) |k| { + blocks[rk][k] = @intCast(c[k].id); + blocks[rk][TOPK + k] = @bitCast(@as(f32, @floatCast(c[k].value))); + } + blocks[rk][2 * TOPK] = 0; + blocks[rk][2 * TOPK + 1] = 0; + } + for ([_]u32{ 1, 5, 20, 64 }) |k| for (0..20) |pos| { + const s: smp.Sampling = .{ .seed = 11, .temperature = 0.8, .top_k = k, .top_p = 0.95 }; + const got = try draw(gpa, .{ &blocks[0], &blocks[1] }, s, pos); + const want = try drawFull(gpa, &row, s, pos); + try std.testing.expectEqual(want, got.token); + }; +} + +test "the bf16 counting sort draws as choose over the whole row" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(3); + const rnd = prng.random(); + var raw: [4096]u16 = undefined; + var row: [4096]f64 = undefined; + for (0..20) |trial| { + for (&raw, &row) |*u, *v| { + const x: f32 = @floatCast(rnd.floatNorm(f64) * 2.0); + u.* = @intCast(@as(u32, @bitCast(x)) >> 16); // truncated: many ties + if (rnd.uintLessThan(u32, 50) == 0) u.* = 0x8000; // -0 among them + v.* = @as(f32, @bitCast(@as(u32, u.*) << 16)); + } + const s: smp.Sampling = .{ .seed = 5 + trial, .temperature = 0.7 + 0.1 * @as(f64, @floatFromInt(trial % 5)), .top_p = 0.9, .min_p = if (trial % 3 == 0) 0.02 else 0.0 }; + for (0..8) |pos| try std.testing.expectEqual(try drawFull(gpa, &row, s, pos), try drawFullBf16(gpa, &raw, s, pos)); + } +} diff --git a/zig/src/families/glm_cuda/forward.zig b/zig/src/families/glm_cuda/forward.zig new file mode 100644 index 0000000000..897645dbee --- /dev/null +++ b/zig/src/families/glm_cuda/forward.zig @@ -0,0 +1,499 @@ +//! GLM-5.3-Flash (glm5_next) on CUDA, one rank of two: the forward on this family's own kernels, for any mix of +//! streams. A step is a list of segments (a stream and its rows, in order): the row-wise work (projections, norms, +//! hyper-connections, router, experts, head) runs over all the step's rows at once, and the attention, the indexer +//! and the KDA recurrences read the step's tables (each row's position and its stream's caches). Every kernel keeps a +//! row's bits its own, so a stream's tokens are the same alone and beside fifteen others. +//! +//! Both ranks make the same calls in the same order. Tokens are drawn from both ranks' vocabulary shards (an +//! all-gather of each rank's best candidates), so both ranks reach the same tokens with no host decision that could +//! differ. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const api = @import("api.zig"); +const kern = @import("kern.zig"); +const sq = @import("seq.zig"); +const blocks = @import("blocks.zig"); + +pub const Seq = sq.Seq; +pub const Media = sq.Media; +const K = kern.K; +const D = kern.D; +const WIDE = kern.WIDE; +const HL = kern.HL; +const LL = kern.LL; + +pub const MAX_STREAMS: usize = 16; +pub const PREFILL_ROWS_MAX: i64 = 4096; // a prompt chunk at most (buffers sized once take the served value) +/// Rows a prompt chunk takes: TF_GLM_PREFILL_ROWS (2048; 4096 halves what the EXL3 experts decode per row and takes +/// about 2 GiB more of the pool). Read once, before the buffers. +pub fn prefillRows() i64 { + const v = std.c.getenv("TF_GLM_PREFILL_ROWS") orelse return 2048; + const n = std.fmt.parseInt(i64, std.mem.span(v), 10) catch return 2048; + return @max(256, @min(PREFILL_ROWS_MAX, @divFloor(n, 256) * 256)); +} +pub const MAX_ROWS: i64 = sq.MAX_ROWS; +pub const STEP_ROWS: i64 = @as(i64, MAX_STREAMS) * 8; // a batched step's rows at most (the lane core keeps a step within them) +pub const HEAD_N: i64 = 77440; // a rank's vocabulary shard +pub const IDX_KW: i64 = kern.IDX_DIM + kern.IDX_HEADS; // the indexer's key rows and head weights +pub const MOE_W: i64 = 1024; // a rank's expert width +pub const DENSE_W: i64 = 6144; // a rank's dense MLP width + +fn cdiv(a: i64, b: i64) i64 { + return @divFloor(a + b - 1, b); +} + +const Carve = struct { + at: u64 = 0, + fn take(self: *Carve, n: i64) u64 { + const p = self.at; + self.at += std.mem.alignForward(u64, @intCast(@max(n, 1)), 512); + return p; + } +}; + +/// One stream's window in a batched step, laid out as glm_kda_chain_seg reads it. +pub const KdaSeg = extern struct { first: i32, rows: i32, rec_in: u64, rec_out: u64, conv: u64, k_save: u64, v_save: u64, g_save: u64, b_save: u64, proj: u64 }; + +/// Scratch for steps of up to ``rows`` rows: a batched decode step's, the MTP head's, a prompt chunk's. +pub const Buffers = struct { + rows: i64, + prefill: bool, + np_max: i64, // the indexer's pool count at the capacity (the scores' columns at most) + mem: cuda.DeviceBuffer = undefined, + ids: u64 = 0, x: u64 = 0, normed: u64 = 0, xs: u64 = 0, post: u64 = 0, comb: u64 = 0, hcpart: u64 = 0, + ka: u64 = 0, kg: u64 = 0, kproj: u64 = 0, kout: u64 = 0, ks: kern.KdaScratch = undefined, q_tmp: u64 = 0, y_tmp: u64 = 0, + dp: u64 = 0, qr: u64 = 0, xs_qr: u64 = 0, lat: u64 = 0, xs_lat: u64 = 0, q: u64 = 0, vn: u64 = 0, + ikr: u64 = 0, igr: u64 = 0, qi: u64 = 0, qa: u64 = 0, ol: u64 = 0, po: u64 = 0, pm: u64 = 0, pl: u64 = 0, + scores: u64 = 0, pools: u64 = 0, tokens: u64 = 0, counts: u64 = 0, + gu: u64 = 0, act: u64 = 0, xs_act: u64 = 0, + mlog: u64 = 0, pick: u64 = 0, wts: u64 = 0, rpick: u64 = 0, eact: u64 = 0, ey: u64 = 0, rey: u64 = 0, plan: kern.Plan = undefined, + x3: kern.Exl3Scratch = .{ .xg = 0, .xu = 0, .xd = 0, .z = 0 }, // the EXL3 experts' scratch (exl3 checkpoints) + sgu: u64 = 0, sact: u64 = 0, sxs: u64 = 0, sy: u64 = 0, + part: u64 = 0, gath: u64 = 0, sk: u64 = 0, hidden: u64 = 0, fnormed: u64 = 0, fxs: u64 = 0, hrows: u64 = 0, logits: u64 = 0, + me: u64 = 0, mcat: u64 = 0, mx: u64 = 0, hin: u64 = 0, cand: u64 = 0, cgath: u64 = 0, invt: u64 = 0, tk: u64 = 0, gtk: u64 = 0, + /// The step's tables: positions [rows] i32, then lc, ik, ig, pk [12][rows] u64, the KDA segments, the head rows. + tab: u64 = 0, + tab_bytes: u64 = 0, + segs: u64 = 0, + last: u64 = 0, + + pub fn rowsTab(b: *const Buffers) kern.Rows { + const r: u64 = @intCast(b.rows); + const t = b.tab; + return .{ .pos = t, .lc = t + r * 4, .ik = t + r * 4 + 12 * r * 8, .ig = t + r * 4 + 24 * r * 8, .pk = t + r * 4 + 36 * r * 8, .cap = b.rows }; + } + + pub fn ey_f32(b: *const Buffers) bool { + return !b.prefill; + } + + /// Rows one attention call takes (its chunk partials' scratch). + pub fn attRows(b: *const Buffers) i64 { + return @min(b.rows, kern.ATT_ROWS); + } + + fn init(d: *const cuda.Driver, rows: i64, is_prefill: bool, np_max: i64, exl3: bool, shx: bool) !Buffers { + var b: Buffers = .{ .rows = rows, .prefill = is_prefill, .np_max = np_max }; + var c: Carve = .{}; + const pairs = rows * kern.TOPK; + const sel_rows = @min(rows, kern.SELECT_ROWS); + const att = b.attRows(); + const nch = @max(kern.DENSE_CHUNKS, kern.SPARSE_CHUNKS); + const head_rows: i64 = if (is_prefill) 1 else rows; + const ey_elt: i64 = if (is_prefill and !exl3) 2 else 4; + const sk_cols: i64 = 16384; // splitK keeps sk * n at or below this for every linear of the model + b.tab_bytes = @intCast(rows * 4 + 48 * rows * 8 + @as(i64, MAX_STREAMS) * @sizeOf(KdaSeg) + @as(i64, MAX_STREAMS) * 4); + inline for (.{ .{ "ids", rows * 4 }, .{ "x", rows * WIDE * 2 }, .{ "normed", rows * D * 2 }, .{ "xs", rows * @divExact(D, 64) * 4 }, + .{ "post", rows * kern.S * 4 }, .{ "comb", rows * kern.S * kern.S * 4 }, .{ "hcpart", rows * 16 * 25 * 4 }, + .{ "ka", rows * LL * 128 * 2 }, .{ "kg", rows * LL * 128 * 2 }, .{ "kproj", rows * sq.PROJ_W * 2 }, .{ "kout", rows * LL * 128 * 2 }, + .{ "dp", rows * (kern.Q_LORA + kern.KV_LORA) * 2 }, .{ "qr", rows * kern.Q_LORA * 2 }, .{ "xs_qr", rows * @divExact(kern.Q_LORA, 64) * 4 }, + .{ "lat", rows * kern.KV_LORA * 2 }, .{ "xs_lat", rows * @divExact(kern.KV_LORA, 64) * 4 }, .{ "q", rows * HL * kern.QK * 2 }, .{ "vn", rows * HL * kern.DV * 2 }, + .{ "ikr", rows * IDX_KW * 2 }, .{ "igr", rows * kern.IDX_DIM * 4 }, .{ "qi", rows * kern.IDX_HEADS * kern.IDX_DIM * 2 }, + .{ "qa", rows * HL * kern.KV_LORA * 2 }, .{ "ol", rows * HL * kern.KV_LORA * 2 }, + .{ "po", nch * att * HL * kern.KV_LORA * 4 }, .{ "pm", nch * att * HL * 4 }, .{ "pl", nch * att * HL * 4 }, + .{ "scores", sel_rows * @max(np_max, 1024) * 4 }, .{ "pools", rows * kern.TOPK_POOLS * 8 }, .{ "tokens", rows * kern.SEL_W * 4 }, + .{ "counts", rows * 4 }, + .{ "gu", rows * 2 * DENSE_W * 2 }, .{ "act", rows * DENSE_W * 2 }, .{ "xs_act", rows * @divExact(DENSE_W, 64) * 4 }, + .{ "mlog", rows * kern.EXPERTS * 4 }, .{ "pick", rows * kern.SLOTS * 4 }, .{ "wts", rows * kern.SLOTS * 4 }, .{ "rpick", rows * kern.TOPK * 4 }, + .{ "eact", rows * kern.SLOTS * MOE_W * 2 }, .{ "ey", rows * kern.SLOTS * D * ey_elt }, .{ "rey", rows * kern.TOPK * D * ey_elt }, + .{ "sgu", rows * 2 * MOE_W * 2 }, .{ "sact", rows * MOE_W * 2 }, .{ "sxs", rows * @divExact(MOE_W, 64) * 4 }, .{ "sy", rows * D * 4 }, + .{ "part", rows * D * 4 }, .{ "gath", 2 * rows * D * 4 }, .{ "sk", rows * sk_cols * 4 }, + .{ "hidden", rows * D * 2 }, .{ "fnormed", rows * D * 2 }, .{ "fxs", rows * @divExact(D, 64) * 4 }, .{ "hrows", @as(i64, MAX_STREAMS) * D * 2 }, + .{ "logits", head_rows * HEAD_N * 2 }, + .{ "me", rows * D * 2 }, .{ "mcat", rows * 2 * D * 2 }, .{ "mx", rows * D * 2 }, .{ "hin", rows * D * 2 }, + .{ "cand", rows * 4 * 4 }, .{ "cgath", 2 * rows * 4 * 4 }, + .{ "invt", head_rows * 4 }, .{ "tk", head_rows * @as(i64, @intCast(@import("draw.zig").W)) * 4 }, .{ "gtk", 2 * head_rows * @as(i64, @intCast(@import("draw.zig").W)) * 4 }, .{ "tab", @as(i64, @intCast(b.tab_bytes)) } }) |f| + @field(b, f[0]) = c.take(f[1]); + const items = kern.maxItems(pairs, kern.EXPERTS, 16); + const wide = pairs > 1024; + b.plan = .{ .members = c.take(pairs * 4), .items = c.take(items * 3 * 4), .counts = c.take(2 * 4), + .rank = c.take(if (wide) pairs * 4 else 4), .hist = c.take(if (wide) cdiv(pairs, 1024) * kern.EXPERTS * 4 else 4), .tile = if (is_prefill) kern.promptTile() else 16 }; + if (exl3) { + const P = rows * kern.TOPK; + const sk: i64 = if (is_prefill) (if (kern.promptOldKernel()) 1 else kern.promptSplit()) else kern.EXL3_SK; // the staged kernel's splits (its partials: 2 * sk * P * MOE_W fp32) + const xrows: i64 = if (is_prefill and shx and !kern.promptOldKernel()) rows else P; // one rotated row a token (the staged prompt kernel) + b.x3 = .{ .xg = c.take(xrows * D * 2), .xu = c.take((if (xrows == rows) 128 else P * D) * 2), .xd = c.take(P * MOE_W * 2), .z = c.take(@max(2 * sk * P * MOE_W, P * D) * 4) }; + } + if (is_prefill) { + b.ks = .{ .out = c.take(rows * LL * 128 * 2), .k = c.take(rows * LL * 128 * 4), .v = c.take(rows * LL * 128 * 2), .g = c.take(rows * LL * 128 * 4), .b = c.take(rows * LL * 4) }; + b.q_tmp = c.take(rows * LL * 128 * 4); + b.y_tmp = c.take(rows * LL * 128 * 2); + } + b.mem = try cuda.DeviceBuffer.alloc(d, c.at); + const base = b.mem.ptr; + const info = @typeInfo(Buffers).@"struct"; + inline for (info.field_names, info.field_types) |name, ty| { + if (ty == u64 and !std.mem.eql(u8, name, "tab_bytes")) @field(b, name) += base; + } + inline for (.{ "members", "items", "counts", "rank", "hist" }) |f| @field(b.plan, f) += base; + if (exl3) inline for (.{ "xg", "xu", "xd", "z" }) |f| { + @field(b.x3, f) += base; + }; + if (is_prefill) inline for (.{ "out", "k", "v", "g", "b" }) |f| { + @field(b.ks, f) += base; + }; + const r: u64 = @intCast(rows); + b.segs = b.tab + r * 4 + 48 * r * 8; + b.last = b.segs + MAX_STREAMS * @sizeOf(KdaSeg); + try d.check(d.api.cuMemsetD8_v2(b.mem.ptr, 0, c.at), "zero buffers"); + return b; + } +}; + +pub const Options = struct { context: i64 }; + +pub const Skip = struct { + gather: bool = false, + experts: bool = false, + dense: bool = false, + attn: bool = false, + hc: bool = false, + + fn fromEnv() Skip { + var s: Skip = .{}; + const v = std.c.getenv("TF_GLM_SKIP") orelse return s; + var it = std.mem.tokenizeScalar(u8, std.mem.span(v), ','); + while (it.next()) |w| inline for (@typeInfo(Skip).@"struct".field_names) |name| { + if (std.mem.eql(u8, w, name)) @field(s, name) = true; + }; + return s; + } +}; + +/// A step's stream and its rows (consecutive in the step, in segment order). +pub const Seg = struct { s: *Seq, rows: i64 }; + +pub const Kind = enum(u8) { verify = 0, mtp = 1 }; + +/// A decode step's shape: what the launches depend on besides the tables (rows, streams, the dense attention's +/// chunks, whether any row is sparse, the sparse scores' pool bucket). +const GKey = struct { kind: Kind, rows: i64, nseg: usize, nch: i64, np: i64, sparse: bool }; +const Cached = struct { ex: cuda.graph.Exec, used: u64 }; + +fn graphCap() usize { + const v = std.c.getenv("TF_GLM_GRAPHS") orelse return 96; + return @max(8, std.fmt.parseInt(usize, std.mem.span(v), 10) catch 96); +} + +pub const Engine = struct { + gpa: std.mem.Allocator, + io: std.Io, + ctx: *api.Ctx, + k: K, + store: *const api.Store, + w: blocks.Weights = undefined, + rank: u8 = 0, + buf: Buffers = undefined, // batched verify steps + mbuf: Buffers = undefined, // the MTP head's steps + pbuf: Buffers = undefined, // prompt chunks + capacity: i64 = 0, // the largest stream's cache slots + host: []align(16) u8 = &.{}, + tab_host: []align(16) u8 = &.{}, // a step's tables before their upload + cand_host: []i32 = &.{}, // both ranks' candidate blocks of a step (sampled draws) + cut: cuda.DeviceBuffer = undefined, // a prompt chunk's cut: every KDA layer's state and window at the kept point + cut_rec: u64 = 0, + cut_conv: u64 = 0, + cut_at: i64 = 0, // the cut's row in the chunk being computed (0: none) + full: cuda.DeviceBuffer = undefined, // both ranks' logits of one row (a draw past the candidates) + full_host: []u16 = &.{}, + /// Development: every launch synchronizes (and names the one that faulted). + sync_each: bool = false, + /// Decode steps as CUDA graphs, one a shape (``GKey``): the same launches, replayed. The shapes grow with the + /// streams' lengths, so at most ``graph_cap`` stay (TF_GLM_GRAPHS), the least recently run dropped first: each + /// holds device memory the stream pool does not count. + use_graphs: bool = true, + graphs: std.AutoHashMap(GKey, Cached) = undefined, + graph_cap: usize = 96, + graph_clock: u64 = 0, + captures: usize = 0, + evictions: usize = 0, + updates: usize = 0, // captures that took an older graph's instantiation (oldestGraph) + capture_ms: f64 = 0, + /// DFlash2 (opt-in): its taps after the target layers it reads, each stream's ring bytes, and the verify steps + /// run (a stream's last verify rows are good for its draft while no later verify ran). + drafter: ?*@import("dflash.zig").Drafter = null, + ring_bytes: u64 = 0, + verifies: u64 = 0, + /// The step being run (its segments and shape), for the blocks. + step: Step = .{}, + /// Development (TF_GLM_SKIP=gather,experts,dense,attn,hc): a class of work left out, its time measured by the + /// steps' difference (the replies are wrong while set). + skip: Skip = .{}, +}; + +/// The step being computed: its segments, their kind, and the host's view of its rows (for the launches' shape). +pub const Step = struct { + segs: []const Seg = &.{}, + kind: Kind = .verify, + rows: i64 = 0, + first_pos: i64 = 0, // a one-segment step's first position (a prompt chunk's) + nch: i64 = 0, // dense attention chunks (0: no dense row) + sparse: bool = false, // some row attends to its selected tokens + np: i64 = 0, // pools the sparse rows score +}; + +pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *api.Ctx, kernels: *api.Kernels, store: *const api.Store, opts: Options) !*Engine { + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = .{ .gpa = gpa, .io = io, .ctx = ctx, .k = try K.init(gpa, ctx.d, ctx.ctx, ctx.stream, &kernels.ext), .store = store }; + e.rank = ctx.rank; + e.skip = Skip.fromEnv(); + kern.readPdl(); + e.w = try blocks.Weights.load(&e.k, store); + e.k.experts_sym = e.w.experts_sym; + e.k.experts3 = e.w.experts3; + e.capacity = opts.context + MAX_ROWS; + const np = @divFloor(e.capacity, 4); + e.buf = try Buffers.init(ctx.d, STEP_ROWS, false, np, e.w.exl3, false); + e.mbuf = try Buffers.init(ctx.d, STEP_ROWS, false, np, e.w.exl3, false); + e.pbuf = try Buffers.init(ctx.d, prefillRows(), true, np, e.w.exl3, e.w.exl3_shx); + e.host = try gpa.alignedAlloc(u8, .@"16", 1 << 20); + e.tab_host = try gpa.alignedAlloc(u8, .@"16", @intCast(e.pbuf.tab_bytes)); + e.graphs = .init(gpa); + e.graph_cap = graphCap(); + e.cand_host = try gpa.alloc(i32, 2 * @as(usize, @intCast(STEP_ROWS)) * @import("draw.zig").W); + e.full = try cuda.DeviceBuffer.alloc(ctx.d, 2 * HEAD_N * 2); + e.full_host = try gpa.alloc(u16, 2 * HEAD_N); + { + const n: u64 = @intCast(e.w.nkda); + const rec_bytes = n * @as(u64, @intCast(sq.REC * 4)); + e.cut = try cuda.DeviceBuffer.alloc(ctx.d, rec_bytes + n * @as(u64, @intCast(3 * sq.CONV_C * 2))); + e.cut_rec = e.cut.ptr; + e.cut_conv = e.cut.ptr + rec_bytes; + } + return e; +} + +pub fn deinit(e: *Engine) void { + e.buf.mem.free(); + e.mbuf.mem.free(); + e.pbuf.mem.free(); + e.gpa.free(e.host); + e.gpa.free(e.tab_host); + e.gpa.free(e.cand_host); + e.gpa.free(e.full_host); + e.full.free(); + e.cut.free(); + var it = e.graphs.valueIterator(); + while (it.next()) |c| c.ex.deinit(); + e.graphs.deinit(); + e.k.cache.deinit(); + e.w.deinit(e.gpa); + e.gpa.destroy(e); +} + +/// Device bytes a stream of ``capacity`` cache slots takes. +pub fn seqBytes(e: *const Engine, capacity: i64) u64 { + return sq.bytes(e.w.nkda, e.w.caches(), capacity, e.ring_bytes); +} + +pub fn newSeq(e: *Engine, capacity: i64) !*Seq { + return sq.alloc(e.gpa, e.ctx.d, e.w.nkda, e.w.ndsa, e.w.caches(), capacity, e.ring_bytes); +} + +pub fn freeSeq(e: *Engine, s: *Seq) void { + sq.free(e.gpa, s); +} + +/// The next prefill's images (null: none); its chunks put each run's features in place of the embeddings (the main +/// layers' and the MTP head's, a row later), and the prefill drops them. +pub fn attach(s: *Seq, m: ?Media) void { + s.media = m; +} + +/// The step's tables for ``segs`` into ``b`` (uploaded and waited for: the host buffer is reused), its shape into +/// e.step. ``kind`` .mtp: the MTP layer's cache and positions (mtp_len), else every DSA layer's and pos. +pub fn prepare(e: *Engine, b: *Buffers, segs: []const Seg, kind: Kind) !void { + var R: i64 = 0; + for (segs) |g| R += g.rows; + if (R == 0 or R > b.rows or segs.len > MAX_STREAMS) return error.WindowTooWide; + const cap: usize = @intCast(b.rows); + const t = e.tab_host[0..@intCast(b.tab_bytes)]; + const pos: []align(1) i32 = std.mem.bytesAsSlice(i32, t[0 .. cap * 4]); + const tabs: []align(1) u64 = std.mem.bytesAsSlice(u64, t[cap * 4 .. cap * 4 + 48 * cap * 8]); + const kseg: []align(1) KdaSeg = std.mem.bytesAsSlice(KdaSeg, t[cap * 4 + 48 * cap * 8 ..][0 .. MAX_STREAMS * @sizeOf(KdaSeg)]); + const last: []align(1) i32 = std.mem.bytesAsSlice(i32, t[cap * 4 + 48 * cap * 8 + MAX_STREAMS * @sizeOf(KdaSeg) ..][0 .. MAX_STREAMS * 4]); + if (kind == .verify and !b.prefill) e.verifies += 1; + var st: Step = .{ .segs = segs, .kind = kind, .rows = R }; + var row: usize = 0; + var last_pos: i64 = 0; + for (segs, 0..) |g, si| { + const s = g.s; + const p0 = if (kind == .mtp) s.mtp_len else s.pos; + if (p0 + g.rows > s.capacity) return error.ContextFull; + try sq.ensureRing(&e.k, s, if (kind == .mtp) 1 else 0, p0, g.rows); + if (si == 0) st.first_pos = p0; + if (kind == .verify and !b.prefill) { + s.v_step = e.verifies; + s.v_row0 = @intCast(row); + s.v_pos0 = p0; + s.v_rows = @intCast(g.rows); + } + kseg[si] = .{ .first = @intCast(row), .rows = @intCast(g.rows), .rec_in = s.recAt(s.cur, 0), .rec_out = s.recAt(1 - s.cur, 0), .conv = s.conv, + .k_save = s.ks.k, .v_save = s.ks.v, .g_save = s.ks.g, .b_save = s.ks.b, .proj = s.proj }; + last[si] = @intCast(row + @as(usize, @intCast(g.rows)) - 1); + for (0..@intCast(g.rows)) |r| { + const q = p0 + @as(i64, @intCast(r)); + pos[row] = @intCast(q); + if (q < kern.DENSE_LIMIT) st.nch = @max(st.nch, cdiv(q + 1, kern.CHUNK)) else st.sparse = true; + last_pos = @max(last_pos, q); + const lo: usize = if (kind == .mtp) e.w.ndsa else 0; + const hi: usize = if (kind == .mtp) e.w.ndsa + 1 else e.w.ndsa; + for (lo..hi) |ci| { + tabs[ci * cap + row] = s.kc[ci]; + tabs[12 * cap + ci * cap + row] = s.ringPtr(s.ik[ci], ci); + tabs[24 * cap + ci * cap + row] = s.ringPtr(s.ig[ci], ci); + tabs[36 * cap + ci * cap + row] = s.pk[ci]; + } + row += 1; + } + } + if (st.sparse) st.np = kern.poolBucket(last_pos, b.np_max); + e.step = st; + try e.k.upload(b.tab, t); + try e.k.stream.synchronize(); +} + +fn keyOf(e: *const Engine) GKey { + const st = e.step; + return .{ .kind = st.kind, .rows = st.rows, .nseg = st.segs.len, .nch = st.nch, .np = st.np, .sparse = st.sparse }; +} + +/// The step's GPU work (capturable): a verify step's forward and every row's candidates, or the MTP head over its +/// staged rows and each stream's last row's candidates. +fn body(e: *Engine) !void { + switch (e.step.kind) { + .verify => { + try blocks.compute(e, &e.buf); + try candidates(e, &e.buf, e.step.rows); + }, + .mtp => { + try blocks.mtpCompute(e, &e.mbuf); + try candidates(e, &e.mbuf, @intCast(e.step.segs.len)); + }, + } +} + +/// ``body`` from its shape's graph (captured the first time the shape runs), or eager in development modes. +pub fn run(e: *Engine) !void { + if (!e.use_graphs or e.sync_each) return body(e); + const key = keyOf(e); + e.graph_clock += 1; + if (e.graphs.getPtr(key)) |c| { + c.used = e.graph_clock; + return c.ex.launchOn(e.k.stream); + } + const t0 = std.Io.Timestamp.now(e.io, .awake); + defer e.capture_ms += @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(e.io, .awake)).toNanoseconds())) / 1e6; + try cuda.graph.beginCapture(e.k.stream, .thread_local); + body(e) catch |err| { + if (cuda.graph.endCapture(e.k.stream)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return err; + }; + var g = try cuda.graph.endCapture(e.k.stream); + defer g.deinit(); + e.captures += 1; + if (e.graphs.count() >= e.graph_cap) { + // full: the least recently run graph of the same layout takes the new shape's parameters in place (no + // instantiation, no memory); a layout the driver will not update to is dropped for a new instantiation + if (oldestGraph(e, key, true)) |k| { + var c = e.graphs.fetchRemove(k).?.value; + const took = if (c.ex.update(g)) |r| r == .success else |_| false; + if (took) { + e.updates += 1; + try e.graphs.put(key, .{ .ex = c.ex, .used = e.graph_clock }); + return c.ex.launchOn(e.k.stream); + } + c.ex.deinit(); + e.evictions += 1; + } else if (oldestGraph(e, key, false)) |k| { + var c = e.graphs.fetchRemove(k).?.value; + c.ex.deinit(); + e.evictions += 1; + } + } + var ex = try g.instantiate(); + errdefer ex.deinit(); + try e.graphs.put(key, .{ .ex = ex, .used = e.graph_clock }); + try ex.launchOn(e.k.stream); +} + +/// The least recently run graph (``alike``: of ``key``'s kind and attention layout, the ones an update can take). +fn oldestGraph(e: *const Engine, key: GKey, alike: bool) ?GKey { + var oldest: ?GKey = null; + var at: u64 = std.math.maxInt(u64); + var it = e.graphs.iterator(); + while (it.next()) |x| { + const k = x.key_ptr.*; + if (alike and (k.kind != key.kind or k.sparse != key.sparse or (k.nch > 0) != (key.nch > 0) or k.nseg != key.nseg)) continue; + if (x.value_ptr.used < at) { + at = x.value_ptr.used; + oldest = k; + } + } + return oldest; +} + +/// One rank's greedy candidate of each logits row -> both ranks' candidates in b.cgath (rank order). +pub fn candidates(e: *Engine, b: *Buffers, rows: i64) !void { + try kern.rowsTop(&e.k, b.logits, HEAD_N, 0, e.w.vocab_offset, b.cand, rows); + try e.ctx.allGather(b.cand, b.cgath, @intCast(4 * rows), .i32, "candidates"); +} + +/// Every rank's fp32 partial b.part[:R] in rank order -> b.gath [2, R, D] (summed rank 0 first by the consumer). +pub fn gather(e: *Engine, b: *Buffers, R: i64) !u64 { + if (e.skip.gather) return b.part; + if (e.k.prof) |m| { + try e.k.stream.synchronize(); + const t0 = std.Io.Timestamp.now(e.io, .awake); + try e.ctx.allGather(b.part, b.gath, @intCast(R * D), .f32, "all-gather"); + try e.k.stream.synchronize(); + const gop = try m.getOrPut(if (e.ctx.roce != null) "(all-gather)" else "(nccl all-gather)"); + if (!gop.found_existing) gop.value_ptr.* = 0; + gop.value_ptr.* += @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(e.io, .awake)).toNanoseconds())) / 1e6; + return b.gath; + } + try e.ctx.allGather(b.part, b.gath, @intCast(R * D), .f32, "all-gather"); + return b.gath; +} + +/// The ids of a step's rows (in segment order) into b.ids. +pub fn stageIds(e: *Engine, b: *Buffers, ids: []const u32) !void { + if (@as(i64, @intCast(ids.len)) > b.rows) return error.WindowTooWide; + const w: []i32 = @alignCast(std.mem.bytesAsSlice(i32, e.host[0 .. ids.len * 4])); + for (ids, w) |x, *y| y.* = @intCast(x); + try e.k.upload(b.ids, std.mem.sliceAsBytes(w)); + // the upload reads the host buffer when the stream gets to it: later stages reuse it, so wait here + try e.k.stream.synchronize(); +} + +pub fn dropGraphs(e: *Engine) void { + var it = e.graphs.valueIterator(); + while (it.next()) |c| c.ex.deinit(); + e.graphs.clearRetainingCapacity(); +} diff --git a/zig/src/families/glm_cuda/guard.zig b/zig/src/families/glm_cuda/guard.zig new file mode 100644 index 0000000000..495d5e8808 --- /dev/null +++ b/zig/src/families/glm_cuda/guard.zig @@ -0,0 +1,136 @@ +//! A rank's watchdog. The GB10's memory is the host's: when a rank's caches, graphs and buffers take the last of it +//! the kernel stalls the whole machine (no OOM kill comes), and the peer then waits in a collective forever. So a +//! thread reads MemAvailable every 100 ms and the rank exits at the floor (TF_MEM_FLOOR_GIB, 2 GiB), which frees +//! everything it holds; below the floor and 3 GiB more it is ``tight`` and rank 0 admits no new request. It also exits +//! when its peer is gone (the control link hung up or timed out, or NCCL reports an error): a rank alone can only +//! hang, and an exited one can be restarted. + +const std = @import("std"); +const posix = std.posix; +const cuda = @import("cuda"); + +const gib: f64 = 1 << 30; +const POLLRDHUP: i16 = 0x2000; + +pub const Guard = struct { + floor: u64, + tight_at: u64, + fd: posix.socket_t, + nccl: *const cuda.nccl.Library, + comm: cuda.nccl.Comm, + running: std.atomic.Value(bool) = .init(true), + tight: std.atomic.Value(bool) = .init(false), + available: std.atomic.Value(u64) = .init(0), + lowest: std.atomic.Value(u64) = .init(std.math.maxInt(u64)), + thread: ?std.Thread = null, + + /// Watches this process's memory, the link's socket ``fd`` and ``comm`` (kept alive by TCP keepalives). + pub fn start(gpa: std.mem.Allocator, fd: posix.socket_t, nccl: *const cuda.nccl.Library, comm: cuda.nccl.Comm) !*Guard { + const floor = floorBytes(); + const g = try gpa.create(Guard); + g.* = .{ .floor = floor, .tight_at = floor + (3 << 30), .fd = fd, .nccl = nccl, .comm = comm }; + keepAlive(fd); + g.thread = try std.Thread.spawn(.{}, loop, .{g}); + std.log.info("memory guard: exit below {d:.1} GiB available, no new requests below {d:.1} GiB", .{ toGib(g.floor), toGib(g.tight_at) }); + return g; + } + + /// Stopped before the link closes (a closing link is no lost peer). + pub fn stop(g: *Guard, gpa: std.mem.Allocator) void { + g.running.store(false, .release); + if (g.thread) |t| t.join(); + gpa.destroy(g); + } + + pub fn isTight(g: *const Guard) bool { + return g.tight.load(.acquire); + } +}; + +fn loop(g: *Guard) void { + const tick: std.c.timespec = .{ .sec = 0, .nsec = 100_000_000 }; + while (g.running.load(.acquire)) { + _ = std.c.nanosleep(&tick, null); + if (available()) |a| { + g.available.store(a, .release); + _ = g.lowest.fetchMin(a, .monotonic); + if (a < g.floor) die(75, "memory guard: {d:.2} GiB available, under the {d:.1} GiB floor: exiting to free this rank's memory before the machine stalls", .{ toGib(a), toGib(g.floor) }); + const was = g.tight.load(.monotonic); + const now = a < g.tight_at; + if (now != was) { + g.tight.store(now, .release); + if (now) std.log.warn("memory guard: {d:.2} GiB available: no new requests until it is back over {d:.1} GiB", .{ toGib(a), toGib(g.tight_at) }) else std.log.info("memory guard: {d:.2} GiB available again: requests admitted", .{toGib(a)}); + } + } + if (peerGone(g.fd)) { + // a stop frame may be in flight before the peer's close: give this rank's own shutdown time to stop us + const grace: std.c.timespec = .{ .sec = 3, .nsec = 0 }; + _ = std.c.nanosleep(&grace, null); + if (!g.running.load(.acquire)) return; + die(76, "memory guard: the peer rank is gone (control link closed): exiting rather than waiting on it", .{}); + } + var r: cuda.nccl.Result = 0; + if (g.nccl.api.ncclCommGetAsyncError(g.comm, &r) == 0 and r != 0 and r != 7) // 7: ncclInProgress + die(77, "memory guard: NCCL reports error {d} ({s}): the peer rank is gone; exiting", .{ r, std.mem.span(g.nccl.api.ncclGetErrorString(r) orelse "?") }); + } +} + +fn die(code: u8, comptime fmt: []const u8, args: anytype) noreturn { + std.log.err(fmt, args); + std.c._exit(code); +} + +fn peerGone(fd: posix.socket_t) bool { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = POLLRDHUP, .revents = 0 }}; + const n = posix.poll(&fds, 0) catch return false; + return n > 0 and fds[0].revents & (POLLRDHUP | posix.POLL.HUP | posix.POLL.ERR) != 0; +} + +/// A peer that stops answering (its machine stalled, no FIN sent) errors the socket within about 11 s. +fn keepAlive(fd: posix.socket_t) void { + const on: c_int = 1; + const idle: c_int = 5; + const intvl: c_int = 2; + const cnt: c_int = 3; + posix.setsockopt(fd, posix.SOL.SOCKET, posix.SO.KEEPALIVE, std.mem.asBytes(&on)) catch {}; + posix.setsockopt(fd, posix.IPPROTO.TCP, 4, std.mem.asBytes(&idle)) catch {}; // TCP_KEEPIDLE + posix.setsockopt(fd, posix.IPPROTO.TCP, 5, std.mem.asBytes(&intvl)) catch {}; // TCP_KEEPINTVL + posix.setsockopt(fd, posix.IPPROTO.TCP, 6, std.mem.asBytes(&cnt)) catch {}; // TCP_KEEPCNT +} + +/// MemAvailable in bytes (null: unreadable). +pub fn available() ?u64 { + const fd = std.c.open("/proc/meminfo", .{ .ACCMODE = .RDONLY }, @as(std.c.mode_t, 0)); + if (fd < 0) return null; + defer _ = std.c.close(fd); + var buf: [4096]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return null; + return parseAvailable(buf[0..@intCast(n)]); +} + +fn parseAvailable(text: []const u8) ?u64 { + var rows = std.mem.tokenizeScalar(u8, text, '\n'); + while (rows.next()) |row| { + if (!std.mem.startsWith(u8, row, "MemAvailable:")) continue; + var words = std.mem.tokenizeAny(u8, row["MemAvailable:".len..], " \t"); + const kib = std.fmt.parseInt(u64, words.next() orelse return null, 10) catch return null; + return kib * 1024; + } + return null; +} + +fn floorBytes() u64 { + const v = std.c.getenv("TF_MEM_FLOOR_GIB") orelse return 2 << 30; + const g = std.fmt.parseFloat(f64, std.mem.span(v)) catch return 2 << 30; + return @intFromFloat(@max(0.5, g) * gib); +} + +fn toGib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / gib; +} + +test "MemAvailable parsed in bytes" { + try std.testing.expectEqual(@as(?u64, 3 * 1024), parseAvailable("MemTotal: 10 kB\nMemFree: 1 kB\nMemAvailable: 3 kB\n")); + try std.testing.expectEqual(@as(?u64, null), parseAvailable("MemTotal: 10 kB\n")); +} diff --git a/zig/src/families/glm_cuda/kern.zig b/zig/src/families/glm_cuda/kern.zig new file mode 100644 index 0000000000..b9df882837 --- /dev/null +++ b/zig/src/families/glm_cuda/kern.zig @@ -0,0 +1,590 @@ +//! GLM-5.3-Flash on CUDA: every launch the forward makes, on this family's own kernels (glm/mm.cu, rows.cu, +//! latent.cu, index.cu, kda.cu, ops.cu, vision.cu) and the shared grouped experts. Each is written so that a row's +//! bits depend on its own inputs alone: a window's rows, sixteen streams' windows and a prompt chunk's rows go through +//! the same arithmetic. The attention and indexer launches read a step's per-row tables (``Rows``): each row's +//! position and its stream's caches, so one launch serves any mix of streams. + +const std = @import("std"); +const cuda = @import("cuda"); + +fn cdiv(a: i64, b: i64) i64 { + return @divFloor(a + b - 1, b); +} + +fn nextPow2(x: u64) u64 { + return std.math.ceilPowerOfTwo(u64, @max(x, 1)) catch unreachable; +} + +pub const Ext = enum(u8) { experts, experts_prefill, kda, sample, ops, experts_sym, experts_sym_prefill, vision, experts_sym3, experts_sym3_prefill, mm, rows, latent, index, exl3, dflash }; + +/// The launch context: the stream and the modules' functions by name. +pub const K = struct { + d: *const cuda.Driver, + stream: cuda.Stream, + names: [16][]const [:0]const u8 = @splat(&.{}), + funcs: [16][]cuda.Function = @splat(&.{}), + sms: i64 = 48, + cache: std.AutoHashMap(u64, cuda.Function) = undefined, + gpa: std.mem.Allocator, + /// Development: synchronize after every launch and name the one that faulted. + sync_each: bool = false, + /// Development: each kernel's time by name (every launch synchronized first and after: proportions only). + prof: ?*std.StringHashMap(f64) = null, + io: ?std.Io = null, + /// The routed experts packed as symmetric groups (scales alone: glm/experts_sym*). + experts_sym: bool = false, + /// ... with 3-bit codes in planes (glm/experts_sym* built with GLM_SYM3). + experts3: bool = false, + /// The routed experts' current product is a prompt's: the prefill kernels (resumed == fresh: a prompt row always + /// takes them, however few rows its chunk has). + prompt: bool = false, + + fn after(self: *K, what: []const u8) !void { + if (!self.sync_each) return; + self.stream.synchronize() catch |e| { + std.log.err("glm: launch {s} faulted", .{what}); + return e; + }; + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, ctx: *const cuda.Context, stream: cuda.Stream, modules: []const cuda.Module) !K { + var k: K = .{ .d = d, .stream = stream, .gpa = gpa }; + k.cache = .init(gpa); + k.sms = try ctx.attribute(.multiprocessor_count); + for (modules, 0..) |m, i| { + var n: c_uint = 0; + try d.check(d.api.cuModuleGetFunctionCount(&n, m.handle), "cuModuleGetFunctionCount"); + const hs = try gpa.alloc(cuda.abi.Function, n); + defer gpa.free(hs); + try d.check(d.api.cuModuleEnumerateFunctions(hs.ptr, n, m.handle), "cuModuleEnumerateFunctions"); + const names = try gpa.alloc([:0]const u8, n); + const fs = try gpa.alloc(cuda.Function, n); + for (hs, 0..) |h, j| { + var s: ?[*:0]const u8 = null; + try d.check(d.api.cuFuncGetName(&s, h), "cuFuncGetName"); + names[j] = try gpa.dupeSentinel(u8, std.mem.span(s.?), 0); + fs[j] = .{ .d = d, .handle = h }; + } + k.names[i] = names; + k.funcs[i] = fs; + } + return k; + } + + /// A kernel of module ``m`` whose (mangled) name contains ``frag``: the one exact match, else the one containing it. + pub fn ext(self: *K, m: Ext, frag: []const u8) !cuda.Function { + const key = std.hash.Wyhash.hash(@intFromEnum(m), frag); + if (self.cache.get(key)) |f| return f; + var found: ?cuda.Function = null; + for (self.names[@intFromEnum(m)], 0..) |n, j| { + if (std.mem.eql(u8, n, frag)) { + found = self.funcs[@intFromEnum(m)][j]; + break; + } + if (std.mem.indexOf(u8, n, frag) != null) { + if (found != null) { + std.log.err("glm: two kernels in {s} match {s}", .{ @tagName(m), frag }); + return error.AmbiguousKernel; + } + found = self.funcs[@intFromEnum(m)][j]; + } + } + const f = found orelse { + std.log.err("glm: no kernel in {s} matches {s}", .{ @tagName(m), frag }); + return error.MissingKernel; + }; + try self.cache.put(key, f); + return f; + } + + pub fn go(self: *K, f: cuda.Function, grid: cuda.Dim3, block: u32, shared: u32, args: *cuda.Args) !void { + return self.goWith(f, grid, block, shared, args, false); + } + + /// A launch the kernel before may overlap (programmatic dependent launch): the kernel waits at + /// griddepcontrol.wait before it reads what that one writes. TF_GLM_PDL=0: plain launches. + pub fn goDependent(self: *K, f: cuda.Function, grid: cuda.Dim3, block: u32, shared: u32, args: *cuda.Args) !void { + return self.goWith(f, grid, block, shared, args, pdl_on); + } + + fn goWith(self: *K, f: cuda.Function, grid: cuda.Dim3, block: u32, shared: u32, args: *cuda.Args, pdl: bool) !void { + if (self.prof) |m| { + try self.stream.synchronize(); + const t0 = std.Io.Timestamp.now(self.io.?, .awake); + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = block }, .shared = shared }, self.stream, args); + try self.stream.synchronize(); + const ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(self.io.?, .awake)).toNanoseconds())) / 1e6; + var nm: ?[*:0]const u8 = null; + _ = self.d.api.cuFuncGetName(&nm, f.handle); + const name = if (nm) |x| std.mem.span(x) else "kernel"; + const gop = try m.getOrPut(name); + if (!gop.found_existing) gop.value_ptr.* = 0; + gop.value_ptr.* += ms; + return; + } + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = block }, .shared = shared, .pdl = pdl }, self.stream, args); + if (self.sync_each) { + var nm: ?[*:0]const u8 = null; + _ = self.d.api.cuFuncGetName(&nm, f.handle); + try self.after(if (nm) |x| std.mem.span(x) else "kernel"); + } + } + + pub fn memset32(self: *K, dst: u64, v: u32, n: usize) !void { + try self.d.check(self.d.api.cuMemsetD32Async(dst, v, n, self.stream.handle), "memset32"); + } + pub fn copy(self: *K, dst: u64, src: u64, n: usize) !void { + try self.d.check(self.d.api.cuMemcpyDtoDAsync_v2(dst, src, n, self.stream.handle), "copy"); + } + pub fn upload(self: *K, dst: u64, bytes: []const u8) !void { + try self.d.check(self.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, self.stream.handle), "upload"); + } +}; + +// -- the model's fixed dimensions (GLM-5.3-Flash, one rank of two) --------------------------------------------------- +pub const D: i64 = 4096; // hidden +pub const S: i64 = 4; // hyper-connection streams +pub const WIDE: i64 = D * S; +pub const EPS: f32 = 1e-5; // rms_norm_eps +pub const HC_EPS: f32 = 1e-6; +pub const HC_ITERS: i64 = 20; +pub const LIMIT: f32 = 10.0; // swiglu_limit +pub const LOWER: f32 = -5.0; // KDA gate lower bound +pub const ROUTED_SCALE: f32 = 2.5; +pub const HL: i64 = 32; // a rank's attention heads (64 / 2) +pub const LL: i64 = 32; // a rank's KDA heads +pub const QK: i64 = 256; // NoPE: the whole head is the nope part +pub const DV: i64 = 256; +pub const Q_LORA: i64 = 1536; +pub const KV_LORA: i64 = 512; +pub const EXPERTS: i64 = 288; +pub const TOPK: i64 = 8; +pub const SLOTS: i64 = TOPK + 1; // the shared expert's slot last +pub const IDX_HEADS: i64 = 32; +pub const IDX_DIM: i64 = 128; +pub const POOL: i64 = 4; +pub const TOPK_POOLS: i64 = 512; +pub const DENSE_LIMIT: i64 = 2051; // index_topk + kpool - 1 +pub const CHUNK: i64 = 512; // keys a latent attention chunk program reads +pub const DENSE_CHUNKS: i64 = 5; // chunks a dense row (position < DENSE_LIMIT) reads at most +pub const SEL_W: i64 = TOPK_POOLS * POOL + POOL - 1; // 2051: a sparse row's attended tokens +pub const SPARSE_CHUNKS: i64 = 5; // cdiv(SEL_W, CHUNK) +pub const ATT_ROWS: i64 = 512; // rows one attention call takes (its partials' scratch) +pub const SELECT_ROWS: i64 = 512; // rows scored at once (the scores' scratch) +pub const SCALE: f32 = 0.0625; // qk_dim ** -0.5 + +/// A step's per-row tables (device arrays): each row's position (i32) and, for each latent cache ``ci``, the row's +/// stream's latent rows, index key and gate rings (offset so that position p is row p) and pooled keys (u64 each, the +/// ``ci``th array at ci * cap). ``at`` moves to a row block. +pub const Rows = struct { + pos: u64, + lc: u64, + ik: u64, + ig: u64, + pk: u64, + cap: i64, + r0: i64 = 0, + + pub fn at(t: Rows, r0: i64) Rows { + var x = t; + x.r0 += r0; + return x; + } + pub fn posAt(t: Rows) u64 { + return t.pos + @as(u64, @intCast(t.r0 * 4)); + } + fn table(t: Rows, base: u64, ci: usize) u64 { + return base + (@as(u64, @intCast(ci)) * @as(u64, @intCast(t.cap)) + @as(u64, @intCast(t.r0))) * 8; + } + pub fn lcAt(t: Rows, ci: usize) u64 { + return t.table(t.lc, ci); + } + pub fn ikAt(t: Rows, ci: usize) u64 { + return t.table(t.ik, ci); + } + pub fn igAt(t: Rows, ci: usize) u64 { + return t.table(t.ig, ci); + } + pub fn pkAt(t: Rows, ci: usize) u64 { + return t.table(t.pk, ci); + } +}; + +// -- rows.cu ------------------------------------------------------------------------------------------------------------- +pub fn embed(k: *K, ids: u64, table: u64, out: u64, R: i64, copies: i64) !void { + var a: cuda.Args = .{}; + a.add(ids); a.add(table); a.add(out); a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(copies))); + try k.go(try k.ext(.rows, "glm_embed"), .{ .x = @intCast(R) }, 256, 0, &a); +} + +/// rmsnorm(x [R, d] (row stride x_stride), w) -> out (row stride o_stride), with xs the 64-groups' sums when given. +pub fn rmsnorm(k: *K, x: u64, x_stride: i64, w: u64, out: u64, o_stride: i64, xs: ?u64, R: i64, d: i64) !void { + var a: cuda.Args = .{}; + a.add(x); a.add(x_stride); a.add(w); a.add(out); a.add(o_stride); a.add(xs orelse @as(u64, 0)); a.add(EPS); a.add(@as(i32, @intCast(d))); + try k.go(try k.ext(.rows, "glm_rmsnorm"), .{ .x = @intCast(R) }, 256, 0, &a); +} + +/// hc_pre: x [R, S*D] -> out [R, D] normed collapsed rows + xs, post [R, S], comb [R, S, S] (part: [R, 16, 25] fp32). +pub fn hcPre(k: *K, x: u64, fn_: u64, base: u64, scale: u64, norm_w: u64, out: u64, xs: u64, post: u64, comb: u64, part: u64, R: i64) !void { + var p: cuda.Args = .{}; + p.add(x); p.add(fn_); p.add(part); + try k.go(try k.ext(.rows, "glm_hc_part"), .{ .x = @intCast(R), .y = 16 }, 128, 0, &p); + var a: cuda.Args = .{}; + a.add(x); a.add(part); a.add(base); a.add(scale); a.add(norm_w); a.add(out); a.add(xs); a.add(post); a.add(comb); + a.add(EPS); a.add(HC_EPS); a.add(@as(i32, @intCast(HC_ITERS))); + try k.go(try k.ext(.rows, "glm_hc_finish"), .{ .x = @intCast(R) }, 512, 0, &a); +} + +/// hc_post: x (in place) from the gathered [2, R, D] fp32 partials. +pub fn hcPost(k: *K, x: u64, gathered: u64, post: u64, comb: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(x); a.add(gathered); a.add(post); a.add(comb); a.add(@as(i32, @intCast(R))); + try k.go(try k.ext(.rows, "glm_hc_post"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); +} + +pub fn residualAdd(k: *K, x: u64, gathered: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(x); a.add(gathered); a.add(@as(i32, @intCast(R))); + try k.go(try k.ext(.rows, "glm_residual_add"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); +} + +pub fn streamMean(k: *K, x: u64, out: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(x); a.add(out); + try k.go(try k.ext(.rows, "glm_stream_mean"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); +} + +pub fn swiglu(k: *K, gu: u64, out: u64, xs: u64, R: i64, w: i64) !void { + var a: cuda.Args = .{}; + a.add(gu); a.add(out); a.add(xs); a.add(LIMIT); a.add(@as(i32, @intCast(w))); + try k.go(try k.ext(.rows, "glm_swiglu"), .{ .x = @intCast(R), .y = @intCast(@divExact(w, 512)) }, 256, 0, &a); +} + +/// router logits: x [m, D] (row stride x_stride) by w [ne, D] -> out [m, ne] fp32. +pub fn router(k: *K, x: u64, x_stride: i64, w: u64, out: u64, m: i64, ne: i64) !void { + var a: cuda.Args = .{}; + a.add(x); a.add(x_stride); a.add(w); a.add(out); a.add(@as(i32, @intCast(m))); a.add(@as(i32, @intCast(ne))); a.add(@as(i32, @intCast(D))); + try k.go(try k.ext(.rows, "glm_router"), .{ .x = @intCast(cdiv(ne, 2)), .y = @intCast(cdiv(m, 16)) }, 256, 0, &a); +} + +/// each row's top-8 experts and weights, then the shared expert (id 288, weight 1) in slot 8. +pub fn topk(k: *K, logits: u64, bias: u64, pick: u64, wts: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(logits); a.add(bias); a.add(pick); a.add(wts); a.add(ROUTED_SCALE); a.add(@as(i32, @intCast(EXPERTS))); + try k.go(try k.ext(.rows, "glm_topk"), .{ .x = @intCast(R) }, 32, 0, &a); +} + +/// combine: y [R, SLOTS, D] bf16 (prompt) or fp32 (windows) by wts -> out [R, D] fp32. +pub fn combine(k: *K, y: u64, y_f32: bool, wts: u64, out: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(y); a.add(@as(i32, @intFromBool(y_f32))); a.add(wts); a.add(out); + try k.go(try k.ext(.rows, "glm_combine"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); +} + +/// conv [n, 3, C] (in place) takes rows keep .. keep + 2 of [conv; proj rows] (proj [n, R, W >= C]); strides in elements. +pub fn convShift(k: *K, conv: u64, proj: u64, keep: i64, conv_layer: i64, proj_layer: i64, proj_row: i64, n: i64, c: i64) !void { + var a: cuda.Args = .{}; + a.add(conv); a.add(proj); a.add(@as(i32, @intCast(keep))); a.add(conv_layer); a.add(proj_layer); a.add(proj_row); a.add(@as(i32, @intCast(c))); + try k.go(try k.ext(.rows, "glm_conv_shift"), .{ .x = @intCast(n), .y = @intCast(cdiv(c, 1024)) }, 256, 0, &a); +} + +// -- mm.cu: the dense linears ------------------------------------------------------------------------------------------ +/// A dense linear [n, k] as glm/mm.cu reads it, tiled at load: ``bits`` 0 (bf16 values) or 4, 6, 8 (symmetric codes +/// with fp16 scales ``s`` per 64 inputs). +pub const B16 = struct { w: u64, n: i64, k: i64, s: u64 = 0, bits: u8 = 0 }; + +pub const SPLIT_TARGET: i64 = 192; + +/// K slices fixed by the shape (column tiles times slices reach SPLIT_TARGET), in groups of 64 inputs. +var pdl_on: bool = true; + +/// TF_GLM_PDL=0 turns the matmuls' dependent launches off (read once at engine start). +pub fn readPdl() void { + if (std.c.getenv("TF_GLM_PDL")) |v| pdl_on = !std.mem.eql(u8, std.mem.span(v), "0"); +} + +pub fn splitK(n: i64, kk: i64) i64 { + const tiles = cdiv(n, 64); + const steps = @divExact(kk, 64); + var sk: i64 = 1; + while (sk < 8 and tiles * sk < SPLIT_TARGET and @mod(steps, sk * 2) == 0 and @divFloor(steps, sk * 2) >= 4) sk *= 2; + return sk; +} + +fn mmName(bits: u8, bm: i64, buf: []u8) ![]const u8 { + return std.fmt.bufPrint(buf, "glm_mm{d}_{d}", .{ if (bits == 0) @as(u8, 16) else bits, bm }); +} + +fn mmSmem(bits: u8, bm: i64) u32 { + const wpc: i64 = if (bits == 0) 8 else @divExact(@as(i64, bits), 2); + return @intCast(3 * (bm * 128 + 8 * 32 * wpc * 4 + 128)); +} + +/// x [m, k] (row stride x_stride) by w -> out [m, n] contiguous, bf16 or unrounded fp32; ``part`` holds sk * m * n fp32 +/// when the shape splits K. Rows go 128 at a time when there are more (the same arithmetic for every row). +pub fn mm(k: *K, x: u64, x_stride: i64, w: B16, out: u64, f32_out: bool, m: i64, part: u64) !void { + const bm: i64 = if (m <= 16) 16 else if (m <= 32) 32 else if (m <= 64) 64 else 128; + const sk = splitK(w.n, w.k); + var nb: [32]u8 = undefined; + const f = try k.ext(.mm, try mmName(w.bits, bm, &nb)); + const smem = mmSmem(w.bits, bm); + try f.allowDynamicShared(smem); + var a: cuda.Args = .{}; + a.add(x); a.add(x_stride); a.add(w.w); a.add(w.s); a.add(out); a.add(part); a.add(@as(i32, @intFromBool(f32_out))); + a.add(@as(i32, @intCast(m))); a.add(@as(i32, @intCast(w.n))); a.add(@as(i32, @intCast(w.k))); a.add(@as(i32, @intCast(sk))); + try k.goDependent(f, .{ .x = @intCast(cdiv(w.n, 64)), .y = @intCast(cdiv(m, bm)), .z = @intCast(sk) }, if (bm <= 32) 128 else 256, smem, &a); + if (sk > 1) { + const total = m * w.n; + var r: cuda.Args = .{}; + r.add(part); r.add(out); r.add(total); r.add(@as(i32, @intCast(sk))); r.add(@as(i32, @intFromBool(f32_out))); + try k.go(try k.ext(.mm, "glm_mm_reduce"), .{ .x = @intCast(cdiv(total, 256)) }, 256, 0, &r); + } +} + +/// The loader's row-major linear (bf16 [n, k], or packed codes and scales) at ``src`` -> its tiled copy at ``dst`` +/// (the same bytes; scales ``ss`` -> ``ds``). +pub fn mmTile(k: *K, bits: u8, src: u64, ss: u64, dst: u64, ds: u64, n: i64, kk: i64) !void { + const chunks = @divExact(n, 8) * @divExact(kk, 64) * 32; + var a: cuda.Args = .{}; + var nb: [32]u8 = undefined; + if (bits == 0) { + a.add(src); a.add(dst); a.add(@as(i32, @intCast(n))); a.add(@as(i32, @intCast(kk))); + try k.go(try k.ext(.mm, "glm_mm_tile16"), .{ .x = @intCast(cdiv(chunks, 256)) }, 256, 0, &a); + return; + } + a.add(src); a.add(ss); a.add(dst); a.add(ds); a.add(@as(i32, @intCast(n))); a.add(@as(i32, @intCast(kk))); + try k.go(try k.ext(.mm, try std.fmt.bufPrint(&nb, "glm_mm_tile{d}", .{bits})), .{ .x = @intCast(cdiv(chunks, 256)) }, 256, 0, &a); +} + +// -- latent.cu ---------------------------------------------------------------------------------------------------------- +pub fn absorbQ(k: *K, q: u64, wk: u64, qa: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(q); a.add(wk); a.add(qa); a.add(@as(i32, @intCast(R))); + try k.go(try k.ext(.latent, "glm_absorb_q"), .{ .x = @intCast(HL), .y = @intCast(@divExact(KV_LORA, 64)), .z = @intCast(cdiv(R, 16)) }, 64, 0, &a); +} + +pub fn expandV(k: *K, ol: u64, wv: u64, out: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(ol); a.add(wv); a.add(out); a.add(@as(i32, @intCast(R))); + try k.go(try k.ext(.latent, "glm_expand_v"), .{ .x = @intCast(HL), .y = @intCast(@divExact(DV, 8)), .z = @intCast(cdiv(R, 16)) }, 256, 0, &a); +} + +/// each row's latent into its stream's cache ``ci`` at its position. +pub fn latentWrite(k: *K, lat: u64, lat_stride: i64, t: Rows, ci: usize, R: i64) !void { + var a: cuda.Args = .{}; + a.add(lat); a.add(lat_stride); a.add(t.lcAt(ci)); a.add(t.posAt()); + try k.go(try k.ext(.latent, "glm_lat_write"), .{ .x = @intCast(R) }, 64, 0, &a); +} + +const ATT_SMEM: u32 = 32 * 1024 + 2 * 16 * 1024 + 2 * 32 * 16 * 4 + 32 * 16 * 2 + 3 * 32 * 4; + +/// Dense causal attention of qa [R, H, 512] (rows below the dense limit) over ``nch`` 512-key chunks -> out. +pub fn denseAttention(k: *K, qa: u64, t: Rows, ci: usize, po: u64, pm: u64, pl: u64, out: u64, R: i64, nch: i64) !void { + const f = try k.ext(.latent, "glm_attn_dense"); + try f.allowDynamicShared(ATT_SMEM); + var a: cuda.Args = .{}; + a.add(qa); a.add(t.lcAt(ci)); a.add(t.posAt()); a.add(po); a.add(pm); a.add(pl); a.add(@as(i32, @intCast(R))); a.add(SCALE); + try k.go(f, .{ .x = @intCast(R), .y = @intCast(nch) }, 256, ATT_SMEM, &a); + try merge(k, po, pm, pl, out, t.posAt(), 0, R, nch, false); +} + +/// Attention of the rows with counts > 0 over their selected tokens, written into out (others untouched). +pub fn sparseAttention(k: *K, qa: u64, t: Rows, ci: usize, tokens: u64, counts: u64, po: u64, pm: u64, pl: u64, out: u64, R: i64) !void { + const f = try k.ext(.latent, "glm_attn_sparse"); + try f.allowDynamicShared(ATT_SMEM); + var a: cuda.Args = .{}; + a.add(qa); a.add(t.lcAt(ci)); a.add(tokens); a.add(counts); a.add(po); a.add(pm); a.add(pl); a.add(@as(i32, @intCast(R))); a.add(SCALE); + try k.go(f, .{ .x = @intCast(R), .y = @intCast(SPARSE_CHUNKS) }, 256, ATT_SMEM, &a); + try merge(k, po, pm, pl, out, 0, counts, R, SPARSE_CHUNKS, true); +} + +fn merge(k: *K, po: u64, pm: u64, pl: u64, out: u64, pos: u64, counts: u64, R: i64, nch: i64, sparse: bool) !void { + var a: cuda.Args = .{}; + a.add(po); a.add(pm); a.add(pl); a.add(out); a.add(pos); a.add(counts); a.add(@as(i32, @intCast(R))); a.add(@as(i32, @intCast(nch))); + a.add(@as(i32, @intFromBool(sparse))); + try k.go(try k.ext(.latent, "glm_attn_merge"), .{ .x = @intCast(R), .y = @intCast(HL) }, 128, 0, &a); +} + +// -- index.cu: the DSA indexer ------------------------------------------------------------------------------------------- +/// The rows' index keys (LayerNorm) and gates into their streams' rings for cache ``ci``, then the pools they complete. +pub fn indexUpdate(k: *K, kr: u64, k_stride: i64, gate: u64, ln_w: u64, ln_b: u64, ape: u64, t: Rows, ci: usize, R: i64) !void { + var a: cuda.Args = .{}; + a.add(kr); a.add(k_stride); a.add(gate); a.add(ln_w); a.add(ln_b); a.add(t.ikAt(ci)); a.add(t.igAt(ci)); a.add(t.posAt()); a.add(@as(f32, 1e-6)); + try k.go(try k.ext(.index, "glm_index_write"), .{ .x = @intCast(R) }, 128, 0, &a); + var b: cuda.Args = .{}; + b.add(t.ikAt(ci)); b.add(t.igAt(ci)); b.add(ape); b.add(t.pkAt(ci)); b.add(t.posAt()); + try k.go(try k.ext(.index, "glm_pool_keys"), .{ .x = @intCast(R) }, 128, 0, &b); +} + +/// Pools to score for rows up to position ``last`` (the visible ones rounded up to a power of two, at least 1024, at +/// most ``np_max``). +pub fn poolBucket(last: i64, np_max: i64) i64 { + const visible = @divFloor(last + 1, POOL) + 1; + return @min(np_max, @max(1024, @as(i64, @intCast(nextPow2(@intCast(visible)))))); +} + +/// Each of n rows' scores over np pools -> scores [n, np] fp32. +pub fn scores(k: *K, qi: u64, w: u64, w_stride: i64, t: Rows, ci: usize, out: u64, n: i64, np: i64) !void { + var a: cuda.Args = .{}; + a.add(qi); a.add(w); a.add(w_stride); a.add(t.pkAt(ci)); a.add(out); a.add(t.posAt()); a.add(@as(i32, @intCast(np))); + a.add(@as(f32, 0.08838834764831845)); a.add(@as(f32, 1.0 / 5.656854249492381)); + try k.go(try k.ext(.index, "glm_index_scores"), .{ .x = @intCast(n), .y = @intCast(cdiv(np, 64)) }, 128, 0, &a); +} + +/// Each row's 512 best visible pools, ascending (int64). +pub fn selectRows(k: *K, sc: u64, out: u64, np: i64, t: Rows, n: i64) !void { + var a: cuda.Args = .{}; + a.add(sc); a.add(out); a.add(@as(i32, @intCast(np))); a.add(t.posAt()); + try k.go(try k.ext(.index, "glm_index_select"), .{ .x = @intCast(n) }, 1024, 0, &a); +} + +/// Each row's attended tokens and their count (0: a dense row). +pub fn sparseTokens(k: *K, pools: u64, t: Rows, rows: i64, tokens: u64, counts: u64) !void { + var a: cuda.Args = .{}; + a.add(pools); a.add(t.posAt()); a.add(tokens); a.add(counts); + try k.go(try k.ext(.index, "glm_index_tokens"), .{ .x = @intCast(rows) }, 256, 0, &a); +} + +/// A batched step's KDA layer ``li``: every stream's window rows through its own states (segs: [n] device KdaSeg). +pub fn kdaChainSeg(k: *K, segs: u64, nseg: i64, li: usize, p: u64, p_stride: i64, b_off: i64, a_: u64, a_stride: i64, g: u64, g_stride: i64, + cw: u64, a_log: u64, dt_bias: u64, norm_w: u64, out: u64, max_rows: i64, proj_w: i64) !void { + var a: cuda.Args = .{}; + a.add(segs); a.add(@as(i32, @intCast(li))); a.add(@as(i32, @intCast(LL))); a.add(p); a.add(@as(i32, @intCast(p_stride))); a.add(@as(i32, @intCast(b_off))); + a.add(a_); a.add(@as(i32, @intCast(a_stride))); a.add(g); a.add(@as(i32, @intCast(g_stride))); a.add(cw); a.add(a_log); a.add(dt_bias); a.add(norm_w); + a.add(EPS); a.add(LOWER); a.add(out); a.add(@as(i32, @intCast(max_rows))); a.add(@as(i32, @intCast(proj_w))); + try k.go(try k.ext(.kda, "glm_kda_chain_seg"), .{ .x = @intCast(LL), .y = @intCast(nseg) }, 1024, 0, &a); +} + +// -- kda.py: the KDA extension (glm/kda.cu) -------------------------------------------------------------------------- +pub const KdaScratch = struct { out: u64, k: u64, v: u64, g: u64, b: u64 }; + +/// kda.chain: projection rows p [q | k | v | f_a | g_a | b at b_off], bf16 gate rows a and g, the conv window and +/// state -> out rows, state_out, and the replay inputs; windows of 64 rows or more take the three-kernel path. +pub fn kdaChain(k: *K, p: u64, p_stride: i64, b_off: i64, a: u64, a_stride: i64, g: u64, g_stride: i64, cs: u64, cw: u64, state_in: u64, + a_log: u64, dt_bias: u64, norm_w: u64, rows: i64, sc: KdaScratch, state_out: u64, q_tmp: u64, y_tmp: u64) !void { + const H: i32 = @intCast(LL); + if (rows >= 64) { + var pa: cuda.Args = .{}; + pa.add(H); pa.add(p); pa.add(@as(i32, @intCast(p_stride))); pa.add(@as(i32, @intCast(b_off))); pa.add(a); pa.add(@as(i32, @intCast(a_stride))); + pa.add(cs); pa.add(cw); pa.add(a_log); pa.add(dt_bias); pa.add(LOWER); pa.add(q_tmp); pa.add(sc.k); pa.add(sc.v); pa.add(sc.g); pa.add(sc.b); + try k.go(try k.ext(.kda, "11prep_kernel"), .{ .x = @intCast(rows), .y = @intCast(H) }, 512, 0, &pa); + var sa: cuda.Args = .{}; + sa.add(H); sa.add(state_in); sa.add(q_tmp); sa.add(sc.k); sa.add(sc.v); sa.add(sc.g); sa.add(sc.b); sa.add(@as(i32, @intCast(rows))); sa.add(y_tmp); sa.add(state_out); + try k.go(try k.ext(.kda, "11step_kernelILi4ELi16E"), .{ .x = @intCast(H), .y = 128 / 4 / 4 }, 4 * 32, 0, &sa); + var oa: cuda.Args = .{}; + oa.add(H); oa.add(y_tmp); oa.add(g); oa.add(@as(i32, @intCast(g_stride))); oa.add(norm_w); oa.add(EPS); oa.add(sc.out); + try k.go(try k.ext(.kda, "10out_kernel"), .{ .x = @intCast(rows), .y = @intCast(H) }, 128, 0, &oa); + return; + } + var ca: cuda.Args = .{}; + ca.add(H); ca.add(p); ca.add(@as(i32, @intCast(p_stride))); ca.add(@as(i32, @intCast(b_off))); ca.add(a); ca.add(@as(i32, @intCast(a_stride))); + ca.add(g); ca.add(@as(i32, @intCast(g_stride))); ca.add(cs); ca.add(cw); ca.add(state_in); ca.add(a_log); ca.add(dt_bias); ca.add(norm_w); + ca.add(EPS); ca.add(LOWER); ca.add(@as(i32, @intCast(rows))); ca.add(sc.out); ca.add(state_out); ca.add(sc.k); ca.add(sc.v); ca.add(sc.g); ca.add(sc.b); + try k.go(try k.ext(.kda, "12chain_kernel"), .{ .x = @intCast(H) }, 1024, 0, &ca); +} + +/// kda.replay_layers: every layer's state after its first ``rows`` rows (states [layers, H, 128, 128]). +pub fn kdaReplayLayers(k: *K, state_in: u64, sk: u64, sv: u64, sg: u64, sb: u64, layers: i64, rows_cap: i64, rows: i64, state_out: u64) !void { + const H: i32 = @intCast(LL); + var a: cuda.Args = .{}; + a.add(H); a.add(state_in); a.add(@as(u64, @intCast(LL * 128 * 128))); a.add(sk); a.add(sv); a.add(sg); a.add(sb); + a.add(@as(u64, @intCast(rows_cap * LL * 128))); a.add(@as(u64, @intCast(rows_cap * LL))); a.add(@as(i32, @intCast(rows))); a.add(state_out); + try k.go(try k.ext(.kda, "20replay_layers_kernel"), .{ .x = @intCast(layers * LL) }, 1024, 0, &a); +} + +// -- the grouped experts (cuda/experts.py, shared with Flash Next's fatbins) ---------------------------------------- +pub const promptTile = @import("kern_exl3.zig").promptTile; +pub const promptOldKernel = @import("kern_exl3.zig").promptOld; +pub const promptSplit = @import("kern_exl3.zig").promptSplit; + +pub fn maxItems(pairs: i64, experts: i64, tile: i64) i64 { + return @min(pairs, experts) + @divFloor(pairs, tile); +} + +pub const Plan = struct { members: u64, items: u64, counts: u64, rank: u64, hist: u64, tile: i64 }; + +/// experts_v7.plan: picks [R, 8] (the routed slots alone) -> the plan over the 288 experts. +pub fn expertsPlan(k: *K, picks: u64, pairs: i64, plan: Plan) !void { + var a: cuda.Args = .{}; + const E: i32 = @intCast(EXPERTS); + if (pairs <= 1024) { + a.add(picks); a.add(@as(i32, @intCast(pairs))); a.add(E); a.add(@as(i32, @intCast(plan.tile))); a.add(plan.members); a.add(plan.items); a.add(plan.counts); + try k.go(try k.ext(.experts, "11plan_kernelE"), .{ .x = 1 }, 1024, 0, &a); + } else { + const nblk = cdiv(pairs, 1024); + a.add(picks); a.add(@as(i32, @intCast(pairs))); a.add(E); a.add(plan.rank); a.add(plan.hist); + try k.go(try k.ext(.experts, "9plan_rankE"), .{ .x = @intCast(nblk) }, 1024, 0, &a); + var b: cuda.Args = .{}; + b.add(@as(i32, @intCast(nblk))); b.add(E); b.add(@as(i32, @intCast(plan.tile))); b.add(plan.hist); b.add(plan.items); b.add(plan.counts); + try k.go(try k.ext(.experts, "12plan_offsetsE"), .{ .x = 1 }, 1024, 0, &b); + var s: cuda.Args = .{}; + s.add(picks); s.add(@as(i32, @intCast(pairs))); s.add(E); s.add(plan.rank); s.add(plan.hist); s.add(plan.members); + try k.go(try k.ext(.experts, "12plan_scatterE"), .{ .x = @intCast(cdiv(pairs, 256)) }, 256, 0, &s); + } +} + +/// experts_v7.prefill / run: gate-up (epi 2) or down (epi 3 bf16 out / 0 fp32 out) over the plan; ``slots``: the +/// pairs a row of x feeds (0 for down: one pair a row). +pub fn expertsCall(k: *K, prefill: bool, epi: i64, x: u64, x_stride: i64, slots: i64, w: u64, kb: i64, nb: i64, plan: Plan, out: u64, width: i64, items: i64, limit: f32) !void { + const gs: i64 = 64; + const m: i64 = if (epi == 2) 2 else 1; + var a: cuda.Args = .{}; + a.add(x); a.add(@as(i32, @intCast(x_stride))); a.add(@as(i32, @intCast(slots))); a.add(w); a.add(@as(i32, @intCast(kb))); + a.add(@as(i32, @intCast(nb))); a.add(plan.items); a.add(plan.counts); a.add(plan.members); a.add(out); a.add(@as(i32, @intCast(width))); + a.add(limit); + var buf: [128]u8 = undefined; + const sbv: i64 = if (k.experts_sym) 1 else 2; // a quad's scale (and bias) uint4s a group + if (prefill) { + const f = try k.ext(if (k.experts3) .experts_sym3_prefill else if (k.experts_sym) .experts_sym_prefill else .experts_prefill, try std.fmt.bufPrint(&buf, "14prefill_kernelILi{d}ELi{d}ELi{d}ELi2ELi2ELi4E", .{ gs, m, epi })); + const block: i64 = @as(i64, if (k.experts3) 24 else 32) * @divExact(4 * gs, 128) + 4 * sbv; + const su: i64 = 64 * 8 + 4 * @divExact(64, gs) * m * block; + const sb = su * 16; + const stages: i64 = if (sb * 4 <= 49152) 4 else if (sb * 3 <= 49152) 3 else 2; + const smem: u32 = @intCast(stages * sb); + try f.allowDynamicShared(smem); + const grid = items * cdiv(nb, 4); + if (grid >= 1) try k.go(f, .{ .x = @intCast(grid) }, 256, smem, &a); + } else { + const f = try k.ext(if (k.experts3) .experts_sym3 else if (k.experts_sym) .experts_sym else .experts, try std.fmt.bufPrint(&buf, "13expert_kernelILi{d}ELi{d}ELi{d}ELi2ELi4E", .{ gs, m, epi })); + const per_sm: i64 = @max(1, try f.occupancy(128, 0)); + const grid = @min(cdiv(items, 4), per_sm * k.sms); + if (grid >= 1) try k.go(f, .{ .x = @intCast(grid) }, 128, 0, &a); + } +} + +// -- our small kernels (glm/ops.cu) --------------------------------------------------------------------------------- +/// x rows (``copies`` streams each) <- feats rows: image features in place of placeholder embeddings. +pub fn injectRows(k: *K, feats: u64, x: u64, rows: i64, copies: i64) !void { + var a: cuda.Args = .{}; + a.add(feats); a.add(x); a.add(@as(i32, @intCast(rows))); a.add(@as(i32, @intCast(copies))); + try k.go(try k.ext(.ops, "glm_inject_rows"), .{ .x = @intCast(@min(4096, cdiv(rows * copies * 4096, 256))) }, 256, 0, &a); +} + +pub fn castBf16(k: *K, src: u64, src_stride: i64, dst: u64, dst_stride: i64, rows: i64, cols: i64) !void { + var a: cuda.Args = .{}; + a.add(src); a.add(src_stride); a.add(dst); a.add(dst_stride); a.add(@as(i32, @intCast(cols))); + try k.go(try k.ext(.ops, "glm_cast_bf16"), .{ .x = @intCast(cdiv(cols, 256)), .y = @intCast(rows) }, 256, 0, &a); +} + +/// Rows of ``bytes`` (a multiple of 4) from src to dst, row pitches in bytes (torch's strided copy_). +pub fn copyRows(k: *K, dst: u64, dpitch: usize, src: u64, spitch: usize, bytes: usize, rows: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(src); a.add(@as(i64, @intCast(spitch / 4))); a.add(dst); a.add(@as(i64, @intCast(dpitch / 4))); a.add(@as(i32, @intCast(bytes / 4))); + try k.go(try k.ext(.ops, "glm_copy_rows"), .{ .x = @intCast(@min(64, cdiv(@intCast(bytes / 4), 256))), .y = @intCast(rows) }, 256, 0, &a); +} + +/// fn_rows_top (Flash Next's sample.cu): each bf16 row's first maximum as a global id, its value and log-sum-exp. +pub fn rowsTop(k: *K, logits: u64, cols: i64, id_map: u64, offset: i64, out: u64, rows: i64) !void { + const f = try k.ext(.sample, "fn_rows_top"); + var a: cuda.Args = .{}; + a.add(logits); a.add(@as(i32, @intCast(cols))); a.add(@as(i32, @intCast(cols))); a.add(id_map); a.add(@as(i32, @intCast(offset))); a.add(out); + try k.go(f, .{ .x = @intCast(rows) }, 1024, 0, &a); +} + +// -- exl3.cu: EXL3 routed experts (kern_exl3.zig) ----------------------------------------------------------------------- +const x3 = @import("kern_exl3.zig"); +pub const Exl3 = x3.Exl3; +pub const Exl3Scratch = x3.Exl3Scratch; +pub const EXL3_SK = x3.EXL3_SK; +pub const exl3Routed = x3.exl3Routed; diff --git a/zig/src/families/glm_cuda/kern_exl3.zig b/zig/src/families/glm_cuda/kern_exl3.zig new file mode 100644 index 0000000000..8e6e30a364 --- /dev/null +++ b/zig/src/families/glm_cuda/kern_exl3.zig @@ -0,0 +1,137 @@ +//! GLM-5.3-Flash on CUDA: the EXL3 routed experts' launches (glm/exl3.cu), re-exported by kern.zig. + +const std = @import("std"); +const cuda = @import("cuda"); +const kern = @import("kern.zig"); +const K = kern.K; +const Plan = kern.Plan; +const D = kern.D; +const TOPK = kern.TOPK; +const EXPERTS = kern.EXPERTS; +const LIMIT = kern.LIMIT; +const maxItems = kern.maxItems; + +/// One layer's EXL3 experts on this rank: trellis words and fp16 scales (weights.expertsExl3). +pub const Exl3 = struct { gt: u64, ut: u64, dt: u64, suh_g: u64, suh_u: u64, svh_g: u64, svh_u: u64, suh_d: u64, svh_d: u64 }; + +/// Scratch for P routed pairs: the rotated inputs (fp16) and the splits' partials. +pub const Exl3Scratch = struct { xg: u64, xu: u64, xd: u64, z: u64 }; + +pub const EXL3_SK: i64 = 4; // gate/up K slices in decode windows (down: 1); prompt chunks take 1 + +/// The EXL3 experts' prompt kernel (TF_GLM_EXL3_PROMPT, read once): +/// fused (the default): the transposed design (exl3.cu, "prompt chunks, the transposed design"): items of 64 pairs, +/// the weights decoded into A fragments, gate and up in one block with the SwiGLU epilogue, the down projection +/// with its epilogue; no partials in memory, one accumulation chain an output (a chunking-independent result). +/// staged: items of 128 pairs, the four warps over columns, rows and words through shared memory, K split +/// TF_GLM_EXL3_PROMPT_SK ways; wide: its 8-warp block (256 columns a block). +/// old: the grouped decode kernel on items of 64 pairs, warps over K (the engine's original prompt path). +pub const PromptKernel = enum { fused, staged, wide, old }; +var prompt_kernel: ?PromptKernel = null; +pub fn promptKernel() PromptKernel { + if (prompt_kernel == null) { + prompt_kernel = .fused; + if (std.c.getenv("TF_GLM_EXL3_PROMPT")) |v| prompt_kernel = std.meta.stringToEnum(PromptKernel, std.mem.span(v)) orelse .fused; + } + return prompt_kernel.?; +} +pub fn promptOld() bool { + return promptKernel() == .old; +} +pub fn promptWide() bool { + return promptKernel() == .wide; +} +pub fn promptFused() bool { + return promptKernel() == .fused; +} +/// A prompt plan's tile: the pairs an item holds. +pub fn promptTile() i64 { + return if (promptOld() or promptFused()) 64 else 128; +} +/// The staged prompt kernel's K splits (TF_GLM_EXL3_PROMPT_SK=1|2|4, 2): a short prompt's few pairs an expert get SK +/// times the blocks at a fraction of the K each. One value for every chunk, so a row's bits do not depend on the +/// chunking (a resumed prompt equals a fresh one); the chunk scratch is sized for it. Read once. +var prompt_sk: ?i64 = null; +pub fn promptSplit() i64 { + if (prompt_sk == null) { + const v: i64 = if (std.c.getenv("TF_GLM_EXL3_PROMPT_SK")) |v| (std.fmt.parseInt(i64, std.mem.span(v), 10) catch 2) else 2; + prompt_sk = if (v == 1 or v == 2 or v == 4) v else 2; + } + return prompt_sk.?; +} + +/// The routed pairs' outputs (pair = row * 8 + slot, picks [R, 8]) -> y [R, ys, D] fp32 at slots 0..7, through +/// the plan's items (16 pairs each in a window, 64 or 128 in a prompt chunk). +/// ``shx``: every expert of the layer shares one input rotation for gate and up (weights' exl3_shx): a prompt chunk +/// rotates each token once (sc.xg holds a row a token) and the staged kernel reads a pair's row as pair >> 3. +pub fn exl3Routed(k: *K, x: u64, x_stride: i64, picks: u64, plan: Plan, w: Exl3, sc: Exl3Scratch, y: u64, ys: i64, R: i64, width: i64, prompt: bool, shx: bool) !void { + const P = R * TOPK; + const items = maxItems(P, EXPERTS, plan.tile); + const fused = prompt and promptFused(); + const staged = prompt and !promptOld() and !fused; + const wide = staged and promptWide(); + const sk: i64 = if (staged) promptSplit() else if (prompt) 1 else EXL3_SK; + const rows_x = (staged or fused) and shx; + if (rows_x) { + var a: cuda.Args = .{}; + a.add(x); a.add(@as(i32, @intCast(x_stride))); a.add(w.suh_g); a.add(sc.xg); a.add(@as(i32, @intCast(D))); + try k.go(try k.ext(.exl3, "glm_exl3_rot_rows"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 128)) }, 32, 0, &a); + } else { + var a: cuda.Args = .{}; + a.add(x); a.add(@as(i32, @intCast(x_stride))); a.add(picks); a.add(w.suh_g); a.add(w.suh_u); a.add(sc.xg); a.add(sc.xu); a.add(@as(i32, @intCast(D))); + try k.go(try k.ext(.exl3, "glm_exl3_rot_in"), .{ .x = @intCast(P), .y = @intCast(@divExact(D, 128)), .z = 2 }, 32, 0, &a); + } + if (fused) { // gate and up with their epilogue -> sc.xd; the down projection with its epilogue -> y + { + var a: cuda.Args = .{}; + a.add(sc.xg); a.add(if (rows_x) sc.xg else sc.xu); a.add(w.gt); a.add(w.ut); a.add(plan.items); a.add(plan.counts); a.add(plan.members); + a.add(w.svh_g); a.add(w.svh_u); a.add(w.suh_d); a.add(sc.xd); + a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(width))); a.add(@as(i32, if (rows_x) 3 else 0)); a.add(LIMIT); + try k.go(try k.ext(.exl3, "glm_exl3_prompt_gu"), .{ .x = @intCast(@divExact(width, 128)), .y = @intCast(items), .z = 1 }, 256, 0, &a); + } + { + var a: cuda.Args = .{}; + a.add(sc.xd); a.add(w.dt); a.add(plan.items); a.add(plan.counts); a.add(plan.members); a.add(w.svh_d); a.add(y); + a.add(@as(i32, @intCast(width))); a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(ys))); + try k.go(try k.ext(.exl3, "glm_exl3_prompt_dn"), .{ .x = @intCast(@divExact(D, 256)), .y = @intCast(items), .z = 1 }, 256, 0, &a); + } + return; + } + const nt: i64 = if (wide) 16 else if (staged) 8 else if (prompt) 2 else 8; + const threads: u32 = if (wide) 256 else 128; + const gname: []const u8 = if (wide) "glm_exl3_prompt128w" else if (staged) "glm_exl3_prompt128" else if (prompt) "glm_exl3_grouped64" else "glm_exl3_grouped"; + // the staged kernel's grid: x = (column block, matrix, split), one item's blocks adjacent (its rows read from L2) + { + var a: cuda.Args = .{}; + a.add(sc.xg); a.add(if (rows_x) sc.xg else sc.xu); a.add(w.gt); a.add(w.ut); a.add(plan.items); a.add(plan.counts); a.add(plan.members); a.add(sc.z); + a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(width))); a.add(@as(i32, @intCast(P))); + if (staged) a.add(@as(i32, if (rows_x) 3 else 0)); + a.add(@as(i32, @intCast(sk))); + if (staged) a.add(@as(i32, 2)); + const ncb: i64 = @divExact(width, 16 * nt); + const grid: cuda.Dim3 = if (staged) .{ .x = @intCast(ncb * 2 * sk), .y = @intCast(items), .z = 1 } else .{ .x = @intCast(items), .y = @intCast(ncb), .z = @intCast(2 * sk) }; + try k.go(try k.ext(.exl3, gname), grid, threads, 0, &a); + } + { + var a: cuda.Args = .{}; + a.add(sc.z); a.add(picks); a.add(w.svh_g); a.add(w.svh_u); a.add(w.suh_d); a.add(sc.xd); a.add(@as(i32, @intCast(P))); a.add(@as(i32, @intCast(width))); + a.add(@as(i32, @intCast(sk))); a.add(LIMIT); + try k.go(try k.ext(.exl3, "glm_exl3_gateup"), .{ .x = @intCast(P), .y = @intCast(@divExact(width, 128)) }, 32, 0, &a); + } + { + var a: cuda.Args = .{}; + a.add(sc.xd); a.add(sc.xd); a.add(w.dt); a.add(w.dt); a.add(plan.items); a.add(plan.counts); a.add(plan.members); a.add(sc.z); + a.add(@as(i32, @intCast(width))); a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(P))); + if (staged) a.add(@as(i32, 0)); + a.add(@as(i32, 1)); + if (staged) a.add(@as(i32, 1)); + const ncb: i64 = @divExact(D, 16 * nt); + const grid: cuda.Dim3 = if (staged) .{ .x = @intCast(ncb), .y = @intCast(items), .z = 1 } else .{ .x = @intCast(items), .y = @intCast(ncb), .z = 1 }; + try k.go(try k.ext(.exl3, gname), grid, threads, 0, &a); + } + { + var a: cuda.Args = .{}; + a.add(sc.z); a.add(picks); a.add(w.svh_d); a.add(y); a.add(@as(i32, @intCast(P))); a.add(@as(i32, @intCast(D))); a.add(@as(i32, 1)); a.add(@as(i32, @intCast(ys))); + try k.go(try k.ext(.exl3, "glm_exl3_down"), .{ .x = @intCast(P), .y = @intCast(@divExact(D, 128)) }, 32, 0, &a); + } +} diff --git a/zig/src/families/glm_cuda/native.zig b/zig/src/families/glm_cuda/native.zig new file mode 100644 index 0000000000..811723e4e9 --- /dev/null +++ b/zig/src/families/glm_cuda/native.zig @@ -0,0 +1,1014 @@ +//! GLM-5.3-Flash (glm5_next) on two GB10s behind the native family interface: rank 0's lane core drives the engine +//! through a lane backend that first sends each call to rank 1, which replays it on its own half of the model +//! (``follow``). Every token is drawn from both ranks' vocabulary shards (an NCCL exchange), so both ranks reach the +//! same tokens and rank 1 needs nothing back. Up to 16 streams decode together, each with its own caches sized to its +//! request out of one memory pool; a released stream's kept prompt state stays for a later prompt that extends it +//! while the pool has room. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const api = @import("api.zig"); +const weights = @import("weights.zig"); +const forward = @import("forward.zig"); +const decode = @import("decode.zig"); +const proto = @import("proto.zig"); +const Guard = @import("guard.zig").Guard; +const copy = @import("copy.zig"); +const dflash = @import("dflash.zig"); +const be = lanes.backend; + +pub const model_type = "glm5_next"; +pub const formats: []const []const u8 = &.{ "glm-affine-experts-v1", "exl3-mcg4-routed-experts" }; +pub const default_context: i64 = 1048576; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = 2048; +pub const two_ranks = true; +/// Its tower takes GLM's image preparation and placeholders (the server's qwen_image.glm). +pub const vision_glm = true; +const vision = weights.vision; + +pub const Options = struct { + context: usize, + drafts: bool, + segments: usize = 1, + tp: u8 = 2, + rank: u8 = 0, + master: ?[]const u8 = null, + master_port: u16 = 29600, + vision: bool = false, // load the vision tower (rank 0) and take image and video prompts +}; + +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; // v1.0.11: `resumed` (the lane host restored the prompt cache entry) + +pub const Loaded = struct { + backend: be.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32 = 0, // every admitted stream takes its own caches, sized to its request + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, + follow: ?*const fn (*anyopaque) anyerror!void = null, +}; + +/// Device bytes this rank's half of the weights takes. +pub fn deviceWeightBytes(io: std.Io, dir: []const u8) u64 { + return weights.deviceBytes(io, dir) catch 92 << 30; +} + +const max_rows: usize = @intCast(forward.MAX_ROWS); // a stream's verify window at most (pending + 15 drafts) +const max_depth = max_rows - 1; +const max_streams = forward.MAX_STREAMS; +/// The head stops at a draft it gives less than this (the chain's probability product), the Python engine's rule; +/// TF_GLM_DRAFT_CONFIDENCE sets another (rank 0's, sent with each draft frame). +const draft_confidence: f64 = 0.7; + +/// TF_GLM_DRAFT_MARGINAL=p (DFlash2, rank 0): a chain stops at the first draft whose own chance (given the drafts +/// before it) is below p, instead of at the chain's product below the confidence cut; 0 (the default) keeps the product rule. +fn marginalFromEnv() f64 { + const v = std.c.getenv("TF_GLM_DRAFT_MARGINAL") orelse return 0; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch 0; +} + +fn confidenceFromEnv() f64 { + const v = std.c.getenv("TF_GLM_DRAFT_CONFIDENCE") orelse return draft_confidence; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch draft_confidence; +} + +/// TF_GLM_TRACE=1 (rank 0): every 256 rounds, the rounds' mean verify and draft times and rows, logged. +const Trace = struct { + on: bool = false, + rounds: u64 = 0, + verify_ms: f64 = 0, + verify_rows: u64 = 0, + verify_streams: u64 = 0, + draft_ms: f64 = 0, + drafts: u64 = 0, + calls: u64 = 0, + copies: u64 = 0, // draft rounds that took a copy over the head's drafts + df_rounds: u64 = 0, // streams' rounds DFlash2 drafted (the drafter mix) + drafts_calls_total: u64 = 0, // streams' draft rounds + copy_tokens: u64 = 0, + desync: u64 = 0, // windows whose pending token was not the copy history's last (a history out of step) + predicted: f64 = 0, // the verified drafts' chances summed (as chained products) and the drafts they landed + landed: u64 = 0, + verified: u64 = 0, + + fn log(t: *Trace, own: *const Owned) void { + const n: f64 = @floatFromInt(@max(t.rounds, 1)); + const c: f64 = @floatFromInt(@max(t.calls, 1)); + std.log.info("glm trace: {d} verifies, {d:.2} ms each, {d:.2} rows ({d:.2} streams); {d} drafts, {d:.2} ms each, {d:.2} drafted a stream", + .{ t.rounds, t.verify_ms / n, @as(f64, @floatFromInt(t.verify_rows)) / n, @as(f64, @floatFromInt(t.verify_streams)) / n, + t.calls, t.draft_ms / c, @as(f64, @floatFromInt(t.drafts)) / @as(f64, @floatFromInt(@max(t.verify_streams, 1))) }); + std.log.info("glm drafts: {d} verified, {d} landed ({d:.2}), their chances said {d:.2}", .{ t.verified, t.landed, + @as(f64, @floatFromInt(t.landed)) / @as(f64, @floatFromInt(@max(t.verified, 1))), t.predicted / @as(f64, @floatFromInt(@max(t.verified, 1))) }); + std.log.info("glm copies: {d} of {d} draft rounds copied ({d:.2} tokens each), copied drafts land {d:.2}, the drafter's {d:.2}; {d} out of step", + .{ t.copies, t.calls, @as(f64, @floatFromInt(t.copy_tokens)) / @as(f64, @floatFromInt(@max(t.copies, 1))), own.copy_rate.chance(), own.draft_rate.chance(), t.desync }); + { + const ds = &decode.draw_stats; + const dc: f64 = @floatFromInt(@max(ds.calls, 1)); + std.log.info("glm draws: {d} sampled verifies, {d:.2} ms gpu + {d:.2} ms host each, {d} rows, {d} full rows, {d} check mismatches", .{ ds.calls, ds.gpu_ms / dc, ds.host_ms / dc, ds.rows, ds.full, ds.mismatch }); + ds.* = .{}; + } + if (own.e.drafter) |dr| { + const nc: f64 = @floatFromInt(@max(dr.calls, 1)); + std.log.info("glm mix: {d} stream rounds by DFlash2, {d} by the MTP head", .{ t.df_rounds, t.drafts_calls_total - t.df_rounds }); + std.log.info("glm dflash2: {d} rounds, {d:.2} ms each (context {d:.2}, tables {d:.2}, layers {d:.2}, head {d:.2}, selector {d:.2}); {d} streams' contexts restarted", + .{ dr.calls, dr.forward_ms / nc, dr.phase_ms[0] / nc, dr.phase_ms[4] / nc, dr.phase_ms[1] / nc, dr.phase_ms[2] / nc, dr.phase_ms[3] / nc, dr.lost }); + dflash.logFine(dr, dr.calls); + dr.calls = 0; + dr.forward_ms = 0; + dr.phase_ms = @splat(0); + } + const g: f64 = 1 << 30; + const avail: u64 = if (own.guard) |gd| gd.available.load(.monotonic) else 0; + const low: u64 = if (own.guard) |gd| gd.lowest.load(.monotonic) else 0; + std.log.info("glm memory: {d:.2} GiB available (lowest {d:.2}); pool {d:.2} of {d:.2} GiB; {d} graphs ({d} captured, {d} updated in place, {d} dropped, {d:.0} ms capturing); device {d:.2} GiB", + .{ @as(f64, @floatFromInt(avail)) / g, @as(f64, @floatFromInt(low)) / g, @as(f64, @floatFromInt(own.used)) / g, @as(f64, @floatFromInt(own.room)) / g, + own.e.graphs.count(), own.e.captures, own.e.updates, own.e.evictions, own.e.capture_ms, @as(f64, @floatFromInt(cuda.usage(false).device)) / g }); + t.* = .{ .on = true }; + } +}; +/// A stream's caches for admission: the slots a small request takes (the host's count of streams); requests take +/// what they need from the pool (``fits``). +const admission_slots: i64 = 32768; +/// Prompt tokens a stream's copy index covers at most (its last; TF_GLM_COPY on): about 24 MiB of host memory. +const copy_index: usize = 1 << 20; + +const Op = proto.Op; +const Writer = proto.Writer; +const Reader = proto.Reader; + +const Lane = struct { + id: u64, + s: *forward.Seq, + held: [max_depth]u32 = undefined, // drafts held for the next window + probs: [max_depth]f64 = undefined, // the head's chance for each (0 past the confidence cut) + nheld: u32 = 0, + rows: u32 = 0, // the last verify's rows, not yet committed (0: nothing pending) + first: u64 = 0, // the prompt's first token (a handle: readFn) + media: bool = false, // the prompt had images (its kept state is its own) + prompt: []u32 = &.{}, // the prompt (rank 0: the kept state's prefix when released) + copy: ?copy.Copy = null, // rank 0: the history copy drafts come from (copy.zig) + copied: u32 = 0, // copied drafts held for the next window + copy_verified: u32 = 0, // of them, the ones the last window verified + draft_verified: u32 = 0, // the last window's verified drafts when the drafter's alone (its landing rate's sample) + // the drafter mix (DFlash2 loaded, TF_GLM_MIX): each source's tokens a round and round time, learned on this stream + use_df: bool = true, + mix_tok: [2]f64 = .{ 2.7, 3.0 }, // [mtp, dflash2] + mix_ms: [2]f64 = .{ 50, 55 }, + mix_round: u32 = 0, + mix_t: ?std.Io.Timestamp = null, + trace_held: u32 = 0, // drafts the last window verified (TF_GLM_TRACE: how many landed is known at the next draft) +}; + +/// A released stream's caches kept at its prompt's kept point (``ids``: that prefix). +const Kept = struct { s: *forward.Seq, ids: []u32 }; + +pub const Owned = struct { + gpa: std.mem.Allocator, + io: std.Io, + rank: u8, + link: cuda.tp_link.Link, + nccl: cuda.nccl.Library, + ctx: api.Ctx, + kernels: api.Kernels, + store: api.Store, + e: *forward.Engine, + lanes_by: std.AutoHashMap(*lanes.Stream, Lane), + seqs: std.AutoHashMap(u64, *forward.Seq), // every live or kept stream's caches, by uid (both ranks) + kept: std.ArrayList(Kept) = .empty, // rank 0: released streams' kept states, oldest first + next_id: u64 = 1, + next_uid: u64 = 1, + room: u64 = 0, // the pool's bytes (setRoom), and what the streams hold of it + used: u64 = 0, + drawn: [256]u32 = undefined, + next: u64 = 0, + costs: [max_rows]lanes.config.Cost = undefined, + cost_count: usize = 0, + shared: [8]lanes.config.Cost = undefined, + shared_count: usize = 0, + mtp_ms: f64 = 0, + confidence: f64 = draft_confidence, + marginal: f64 = 0, // TF_GLM_DRAFT_MARGINAL: the DFlash2 chain's per-draft chance floor (0: the product rule) + trace: Trace = .{}, + tower: ?vision.Tower = null, // rank 0, with vision: the image encoder + image_token: u32 = 0, // the placeholder the features replace (config image_token_id) + // rank 1: the next prefill's images (its image frame's), freed after it + media_spans: std.ArrayList([2]i64) = .empty, + media_feats: ?cuda.DeviceBuffer = null, + guard: ?*Guard = null, // the memory and peer watchdog (guard.zig) + copy_on: bool = true, // TF_GLM_COPY=0: the head's drafts only + mix_on: bool = true, // TF_GLM_MIX=0: with DFlash2 loaded, it drafts every round (else each stream picks a source a round) + copy_rate: copy.Rate = .{}, + draft_rate: copy.Rate = .{}, // the drafter's drafts' landing rate (its own chances run high: DFlash2's say ~0.9 where ~0.67 land) + copy_min: usize = copy.K, // TF_GLM_COPY_MIN / TF_GLM_COPY_MIN_REPLY: the tokens a match takes in the prompt / the reply + copy_min_reply: usize = copy.K_REPLY, + + pub fn take(self: *Owned, t: u32) u64 { + self.drawn[self.next % self.drawn.len] = t; + self.next += 1; + return self.next - 1; + } + + pub fn send(self: *Owned, op: Op, w: *Writer) !void { + try self.link.send(@intFromEnum(op), w.buf.items); + w.buf.clearRetainingCapacity(); + } + + fn bytesOf(self: *const Owned, capacity: i64) u64 { + return forward.seqBytes(self.e, capacity); + } + + /// A new stream's caches of ``capacity`` slots (``uid``), the pool's bytes counted. + pub fn newSeq(self: *Owned, uid: u64, capacity: i64) !*forward.Seq { + const s = try forward.newSeq(self.e, capacity); + s.uid = uid; + try self.seqs.put(uid, s); + self.used += self.bytesOf(capacity); + return s; + } + + pub fn dropSeq(self: *Owned, uid: u64) void { + const s = self.seqs.get(uid) orelse return; + _ = self.seqs.remove(uid); + self.used -|= self.bytesOf(s.capacity); + forward.freeSeq(self.e, s); + } + + /// Rank 0: the oldest kept state dropped on both ranks (false: none is kept). + fn evictOldest(self: *Owned) !bool { + if (self.kept.items.len == 0) return false; + const k = self.kept.orderedRemove(0); + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(k.s.uid); + try self.send(.evict, &w); + self.gpa.free(k.ids); + self.dropSeq(k.s.uid); + return true; + } + + fn keptBytes(self: *const Owned) u64 { + var n: u64 = 0; + for (self.kept.items) |k| n += self.bytesOf(k.s.capacity); + return n; + } +}; + +/// Whether a request of ``prompt`` tokens and ``max_tokens`` more fits the pool beside the running streams (kept +/// states give way): the lane host keeps it waiting, in order, until it does. +pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { + const self = of(p); + if (self.guard) |g| if (g.isTight()) return false; // the machine's memory runs short: running streams finish first + const need = self.bytesOf(@as(i64, prompt) + @as(i64, max_tokens) + forward.MAX_ROWS); + return self.used -| self.keptBytes() + need <= self.room; +} + +/// The host's pool for the streams (the bytes left after the model and its buffers), set once after open. Rank 1 +/// holds the same caches out of its own memory, so the pair takes the smaller of the two ranks' pools (rank 1 sends +/// its own: the frame after its ready), with a 2% margin. +pub fn setRoom(p: *anyopaque, room: u64) !void { + const self = of(p); + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + if (self.rank == 1) { + try w.int(room); + try self.send(.room, &w); + self.room = room - room / 50; + return; + } + const m = try self.link.recv(self.gpa); + defer self.gpa.free(m.bytes); + if (m.tag != @intFromEnum(Op.room)) return error.RankOneNoRoom; + var r: Reader = .{ .bytes = m.bytes }; + const peer = try r.int(); + const both = @min(room, peer); + self.room = both - both / 50; + std.log.info("glm: stream pool {d:.2} GiB (rank 0 {d:.2}, rank 1 {d:.2} GiB)", .{ @as(f64, @floatFromInt(self.room)) / (1 << 30), @as(f64, @floatFromInt(room)) / (1 << 30), @as(f64, @floatFromInt(peer)) / (1 << 30) }); + std.log.info("glm: prompt chunks of {d} rows (TF_GLM_PREFILL_ROWS), the EXL3 experts' prompt kernel {s}", .{ self.e.pbuf.rows, @tagName(@import("kern_exl3.zig").promptKernel()) }); +} + + +/// Commits the lanes' last windows: ``kept`` of their rows (the lane core calls keep only to cut a window short, so a +/// window it accepted whole is committed here, before the lane's next call). Rank 1 commits the same. +fn settle(self: *Owned, ls: []const *Lane, kept: []const u32) !void { + var n: usize = 0; + for (ls) |l| n += @intFromBool(l.rows != 0); + if (n == 0) return; + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(n); + for (ls, kept) |l, kp| if (l.rows != 0) { + try w.int(l.s.uid); + try w.int(l.rows); + try w.int(kp); + }; + try self.send(.keep, &w); + for (ls, kept) |l, kp| if (l.rows != 0) { + try decode.keep(self.e, l.s, l.rows, kp); + l.rows = 0; + }; +} + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, _: ?[]const u8, o: Options) !Loaded { + if (o.tp != 2) return error.TwoRanksOnly; + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.* = undefined; + own.gpa = gpa; + own.io = io; + own.rank = o.rank; + own.kept = .empty; + own.next_id = 1; + own.next_uid = 1; + own.next = 0; + own.room = 0; + own.used = 0; + own.cost_count = 0; + own.shared_count = 0; + own.mtp_ms = 0; + own.confidence = confidenceFromEnv(); + own.marginal = marginalFromEnv(); + own.copy_on = if (std.c.getenv("TF_GLM_COPY")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; + own.copy_rate = .{}; + own.mix_on = if (std.c.getenv("TF_GLM_MIX")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; + own.copy_min = envUsize("TF_GLM_COPY_MIN", copy.K); + own.copy_min_reply = envUsize("TF_GLM_COPY_MIN_REPLY", copy.K_REPLY); + own.trace = .{ .on = o.rank == 0 and std.c.getenv("TF_GLM_TRACE") != null }; + own.lanes_by = .init(gpa); + own.seqs = .init(gpa); + // the link and the communicator first, while memory is fresh (NCCL registers its buffers at its first use) + own.link = if (o.rank == 0) try cuda.tp_link.Link.lead(o.master_port, 600_000) + else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(o.master orelse return error.NoMaster, o.master_port), 600_000); + errdefer own.link.close(); + own.nccl = try cuda.nccl.Library.open(); + errdefer own.nccl.close(); + const comm = try cuda.tp_link.communicator(&own.nccl, own.link); + own.ctx = .{ .d = ctx.d, .ctx = ctx, .stream = try cuda.Stream.init(ctx.d, true), .nccl = &own.nccl, .comm = comm, .rank = o.rank }; + own.guard = try Guard.start(gpa, own.link.fd, &own.nccl, comm); // from here a short machine or a lost peer ends the rank + errdefer own.guard.?.stop(gpa); + { // NCCL connects (and registers its buffers) at its first collective: here, before the weights take the memory + var warm = try cuda.DeviceBuffer.alloc(ctx.d, 4096); + defer warm.free(); + try own.nccl.check(own.nccl.api.ncclAllGather(warm.ptr, warm.ptr + 2048, 256, .i32, comm, own.ctx.stream.handle), "warm gather"); + try own.ctx.stream.synchronize(); + } + // small gathers over RoCE (TF_ROCE=0: NCCL for all; TF_ROCE_MAX_KB: a rank's shard at most, 512 KiB) + own.ctx.roce = try cuda.roce.Gather.init(gpa, io, ctx.d, own.link, roceWanted(), roceMaxBytes()); + errdefer if (own.ctx.roce) |g| g.deinit(); + own.kernels = try api.Kernels.load(ctx.d); + errdefer own.kernels.deinit(); + own.store = try weights.load(gpa, io, &own.ctx, dir, o.rank); + errdefer own.store.deinit(); + own.e = try forward.init(gpa, io, &own.ctx, &own.kernels, &own.store, .{ .context = @intCast(o.context) }); + errdefer forward.deinit(own.e); + own.tower = null; + own.media_spans = .empty; + own.media_feats = null; + own.image_token = try imageToken(gpa, io, dir); + if (o.vision and o.rank == 0) own.tower = try vision.Tower.load(gpa, io, ctx.d, dir); + errdefer if (own.tower) |*t| t.deinit(); + // TF_GLM_DRAFTER=dflash2: IncoAI's DFlash2 drafter (non-commercial licence) from TF_GLM_DFLASH2 on both ranks + if (drafterWanted()) { + const ddir = std.mem.span(std.c.getenv("TF_GLM_DFLASH2") orelse return error.NoDFlash2Dir); + own.e.drafter = try dflash.load(gpa, io, own.e, ddir); + own.e.ring_bytes = dflash.ring_bytes; + } + errdefer if (own.e.drafter) |dr| dr.deinit(); + { // both ranks draft alike or neither starts + const mine = [1]u8{@intFromBool(own.e.drafter != null)}; + var theirs: [1]u8 = undefined; + try own.link.swap(&mine, &theirs); + if (mine[0] != theirs[0]) return error.DrafterOnOneRank; + } + // both halves loaded before rank 0 serves: rank 1 says so once its engine is up + if (o.rank == 1) try own.link.send(@intFromEnum(Op.ready), "") else { + const m = try own.link.recv(gpa); + defer gpa.free(m.bytes); + if (m.tag != @intFromEnum(Op.ready)) return error.RankOneNotReady; + } + if (o.rank == 0) { // rank 1 replays it in its follow loop, then hears it ended + try @import("calibrate.zig").run(own, o.drafts); + try own.link.send(@intFromEnum(Op.ready), ""); + if (own.e.drafter != null) if (std.c.getenv("TF_GLM_DFLASH2_STEP_MS") != null) { + own.mtp_ms = dflashStepMs(); // calibrate timed a round; the variable sets another + }; + } + const window = windowFromEnv(); + return .{ + .backend = .{ .ptr = own, .vtable = &vtable }, + .facts = .{ .exact_width = window, .mtp = o.drafts, .speculate = o.drafts, .speculate_early = false, .drafts = if (o.drafts) window - 1 else 1, + .hidden_rows = true, .max_streams = max_streams, .batch_rows = @intCast(forward.STEP_ROWS), + .window_costs = own.costs[0..own.cost_count], .mtp_step_ms = own.mtp_ms, .shared_costs = own.shared[0..own.shared_count], + .draft_probabilities = o.drafts, .draft_streams = o.drafts }, + .rows = if (o.drafts) window else 1, + .stream_bytes = forward.seqBytes(own.e, admission_slots), + .ctx = own, + .deinit = release, + .follow = if (o.rank == 1) followLoop else null, + }; +} + +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PromptTooLong, error.ContextFull => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.NoRoom => "the prompt and its reply do not fit the memory left beside the running requests", + error.NoVision => "image and video inputs need the server started with --vision", + error.ImageCountMismatch, error.ImageTokensMismatch => "the prompt's image placeholders do not match its images", + else => null, + }; +} + +fn release(p: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(p)); + if (own.guard) |g| g.stop(own.gpa); // the peer's close that follows is no lost peer + if (own.rank == 0) own.link.send(@intFromEnum(Op.stop), "") catch {}; + for (own.kept.items) |k| own.gpa.free(k.ids); + own.kept.deinit(own.gpa); + var it = own.lanes_by.valueIterator(); + while (it.next()) |l| { + own.gpa.free(l.prompt); + if (l.copy) |*h| h.deinit(own.gpa); + } + own.lanes_by.deinit(); + var si = own.seqs.valueIterator(); + while (si.next()) |s| forward.freeSeq(own.e, s.*); + own.seqs.deinit(); + if (own.tower) |*t| t.deinit(); + if (own.e.drafter) |dr| dr.deinit(); + own.media_spans.deinit(own.gpa); + if (own.media_feats) |*f| f.free(); + forward.deinit(own.e); + own.store.deinit(); + own.kernels.deinit(); + if (own.ctx.roce) |g| g.deinit(); + _ = own.nccl.api.ncclCommAbort(own.ctx.comm); // its threads stop before the library unloads + own.link.close(); + own.nccl.close(); + own.gpa.destroy(own); +} + +pub fn of(p: *anyopaque) *Owned { + return @ptrCast(@alignCast(p)); +} + +const vtable: be.Backend.VTable = .{ + .prefill = prefillFn, + .first = firstFn, + .queue = queueFn, + .read = readFn, + .verify = verifyFn, + .keep = keepFn, + .draft = draftFn, + .release = releaseFn, + .probabilities = probabilitiesFn, +}; + +/// Cache slots a kept state is allocated beyond its own request (TF_GLM_KEEP_HEADROOM, 8192; 0: none). Read once. +var keep_headroom: ?i64 = null; +fn keepHeadroom() i64 { + if (keep_headroom == null) keep_headroom = if (std.c.getenv("TF_GLM_KEEP_HEADROOM")) |v| @max(0, std.fmt.parseInt(i64, std.mem.span(v), 10) catch 8192) else 8192; + return keep_headroom.?; +} + +/// The stream's prompt into its caches (a kept state it extends, or new ones from the pool); the first token is +/// drawn here. +fn prefillFn(p: *anyopaque, s: *lanes.Stream) anyerror!void { + const self = of(p); + const ids = s.prompt(); + const cap: i64 = @intCast(ids.len + s.max_new + max_rows); + if (ids.len == 0 or cap > self.e.capacity) return error.PromptTooLong; + const id = self.next_id; + self.next_id += 1; + // images: encoded, their placeholder runs found, rank 1 sent both, before the prompt (a refusal stays here) + var spans: std.ArrayList([2]i64) = .empty; + defer spans.deinit(self.gpa); + var feats: ?cuda.DeviceBuffer = null; + defer if (feats) |*f| f.free(); + if (s.images.len > 0) feats = try encodeImages(self, ids, s.images, id, &spans); + const media = s.images.len > 0; + // the longest kept state this prompt extends (not with images: their features are this prompt's alone) + var best: ?usize = null; + if (!media) for (self.kept.items, 0..) |k, i| { + if (k.ids.len >= ids.len or k.s.capacity < cap or !std.mem.eql(u32, k.ids, ids[0..k.ids.len])) continue; + if (best == null or k.ids.len > self.kept.items[best.?].ids.len) best = i; + }; + var seq: *forward.Seq = undefined; + var start: usize = 0; + var fresh: i64 = 0; + const keep_at: usize = if (s.drafts and !media) @max(1, ids.len - 1) else 0; + if (best) |i| { + const k = self.kept.orderedRemove(i); + seq = k.s; + start = k.ids.len; + self.gpa.free(k.ids); + } else { + // a state that will be kept gets headroom when the pool has it: a later turn of the conversation (its reply + // and the next message more) resumes it instead of a prompt its exact size only + const head: i64 = if (keep_at > 0) keepHeadroom() else 0; + var alloc = cap; + // (kept states give way, as in fits(): a pool full of them is not a full pool) + if (head > 0 and cap + head <= self.e.capacity and self.used -| self.keptBytes() + self.bytesOf(cap + head) <= self.room) alloc = cap + head; + while (self.used + self.bytesOf(alloc) > self.room) if (!try self.evictOldest()) return error.NoRoom; + fresh = alloc; + seq = try self.newSeq(self.next_uid, alloc); + self.next_uid += 1; + } + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(seq.uid); + try w.int(fresh); + try w.tokens(ids); + try proto.writeSampling(&w, s.sampling); + try w.int(start); + try w.int(keep_at); + try self.send(.prefill, &w); + seq.sampling = s.sampling; + s.cached = @intCast(start); + if (feats) |f| forward.attach(seq, .{ .spans = spans.items, .feats = f.ptr }); + const first = try decode.prefill(self.e, seq, ids, start, keep_at); + var history: ?copy.Copy = if (self.copy_on and s.drafts) try copy.Copy.init(self.gpa, ids, copy_index, self.copy_min, self.copy_min_reply) else null; + errdefer if (history) |*h| h.deinit(self.gpa); + try self.lanes_by.put(s, .{ .id = id, .s = seq, .first = self.take(first), .media = media, .prompt = try self.gpa.dupe(u32, ids), .copy = history }); +} + +/// The prompt's images through the tower (each its own call, features in prompt order) and their placeholder runs +/// (one a image, its length the image's tokens); rank 1 gets both in an image frame. +fn encodeImages(self: *Owned, ids: []const u32, images: []const lanes.Image, id: u64, spans: *std.ArrayList([2]i64)) !cuda.DeviceBuffer { + const tower = &(self.tower orelse return error.NoVision); + var i: usize = 0; + while (i < ids.len) : (i += 1) { + if (ids[i] != self.image_token) continue; + var j = i; + while (j < ids.len and ids[j] == self.image_token) j += 1; + try spans.append(self.gpa, .{ @intCast(i), @intCast(j) }); + i = j; + } + if (spans.items.len != images.len) return error.ImageCountMismatch; + var rows: usize = 0; + for (images, spans.items) |im, sp| { + const n: usize = @intCast(@divExact(im.grid[0] * im.grid[1] * im.grid[2], 4)); + if (n != sp[1] - sp[0]) return error.ImageTokensMismatch; + rows += n; + } + var feats = try cuda.DeviceBuffer.alloc(self.ctx.d, rows * 4096 * 2); + errdefer feats.free(); + var at: usize = 0; + for (images, spans.items) |im, sp| { + try tower.encode(&self.e.k, im.patches, im.grid, feats.ptr + at * 4096 * 2); + at += @intCast(sp[1] - sp[0]); + } + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(id); + try w.int(spans.items.len); + for (spans.items) |sp| { + try w.int(sp[0]); + try w.int(sp[1]); + } + try w.int(rows); + const fat = w.buf.items.len; + try w.buf.resize(self.gpa, fat + rows * 4096 * 2); + try self.ctx.d.check(self.ctx.d.api.cuMemcpyDtoH_v2(w.buf.items[fat..].ptr, feats.ptr, rows * 4096 * 2), "image features"); + try self.send(.image, &w); + return feats; +} + +fn imageToken(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !u32 { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(text); + const cfg = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer cfg.deinit(); + const v = cfg.value.object.get("image_token_id") orelse return error.NoImageToken; + return @intCast(v.integer); +} + +fn laneOf(self: *Owned, s: *lanes.Stream) !*Lane { + return self.lanes_by.getPtr(s) orelse error.NoLane; +} + +fn firstFn(p: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + const self = of(p); + if (position != s.prompt_len) return error.PositionMismatch; + return (try laneOf(self, s)).first; +} + +fn queueFn(_: *anyopaque, _: *lanes.Stream, _: be.Feed, _: u64) anyerror!u64 { + return error.NotPipelined; +} + +fn readFn(p: *anyopaque, handle: u64) anyerror!u32 { + const self = of(p); + if (handle >= self.next or self.next - handle > self.drawn.len) return error.NoSuchToken; + return self.drawn[handle % self.drawn.len]; +} + +fn verifyFn(p: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const self = of(p); + if (windows.len > max_streams) return error.TooManyStreams; + var ls: [max_streams]*Lane = undefined; + var kept: [max_streams]u32 = undefined; + for (windows, 0..) |win, i| { + if (win.parents != null) return error.TreesNotBuilt; + ls[i] = try laneOf(self, win.stream); + kept[i] = ls[i].rows; + if (win.rows() > max_rows) return error.WindowTooWide; + } + try settle(self, ls[0..windows.len], kept[0..windows.len]); // the previous windows, accepted whole + var ids: [max_streams][max_rows]u32 = undefined; + var wins: [max_streams]decode.Win = undefined; + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(windows.len); + for (windows, 0..) |win, i| { + const l = ls[i]; + const rows = win.rows(); + // the core keys row r's draw at position pos + 1 + r, as the forward draws it + if (win.positions.len > 0 and win.positions[0] != @as(u64, @intCast(l.s.pos + 1))) return error.PositionMismatch; + ids[i][0] = win.pending; + if (l.copy) |h| if (self.trace.on and h.hist.items[h.hist.items.len - 1] != win.pending) { + self.trace.desync += 1; + }; + l.copy_verified = @min(l.copied, win.held); + l.draft_verified = if (l.copied == 0) win.held else 0; + if (self.trace.on and win.held > 0) { + var chain: f64 = 1; + for (l.probs[0..win.held]) |pr| { + chain *= pr; + self.trace.predicted += chain; + } + self.trace.verified += win.held; + l.trace_held = win.held; + } + if (win.held > 0) @memcpy(ids[i][1..][0..win.held], l.held[0..win.held]) else @memcpy(ids[i][1..][0..win.tokens.len], win.tokens); + wins[i] = .{ .s = l.s, .ids = ids[i][0..rows] }; + try w.int(l.s.uid); + try w.tokens(ids[i][0..rows]); + } + const t0 = std.Io.Timestamp.now(self.io, .awake); + try self.send(.verify, &w); + var picks: [forward.STEP_ROWS]decode.Pick = undefined; + try decode.verify(self.e, wins[0..windows.len], &picks); + if (self.trace.on) { + const t = &self.trace; + t.verify_ms += msSince(self.io, t0); + t.rounds += 1; + t.verify_streams += windows.len; + for (windows) |win| t.verify_rows += win.rows(); + if (t.rounds == 256) t.log(self); + } + var at: usize = 0; + for (windows, 0..) |win, i| { + const rows = win.rows(); + for (0..rows) |r| out[i].sampled[r] = picks[at + r].tok; + @memcpy(out[i].drafts[0 .. rows - 1], ids[i][1..rows]); + ls[i].rows = @intCast(rows); + ls[i].nheld = 0; + at += rows; + } +} + +fn keepFn(p: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const self = of(p); + var ls: [max_streams]*Lane = undefined; + var kept: [max_streams]u32 = undefined; + for (windows, paths, 0..) |win, path, i| { + if (path.len == 0) return error.EmptyPath; + for (path, 0..) |r, j| if (r != j) return error.TreesNotBuilt; + ls[i] = try laneOf(self, win.stream); + kept[i] = @intCast(path.len); + } + try settle(self, ls[0..windows.len], kept[0..windows.len]); +} + +fn draftFn(p: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + const self = of(p); + if (requests.len > max_streams) return error.TooManyStreams; + var ls: [max_streams]*Lane = undefined; + var kept: [max_streams]u32 = undefined; + for (requests, 0..) |r, i| { + if (r.lanes != null) return error.TreesNotBuilt; + ls[i] = try laneOf(self, r.stream); + // the kept rows (a lone round has kept already and settles nothing) + kept[i] = if (r.rows) |rows| @intCast(rows.len) else ls[i].rows; + } + try settle(self, ls[0..requests.len], kept[0..requests.len]); + var follow: [max_streams][max_rows]u32 = undefined; + var reqs: [max_streams]decode.DraftReq = undefined; + var depths: [max_streams]u32 = undefined; + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(requests.len); + for (requests, 0..) |r, i| { + const l = ls[i]; + var n: usize = 1; + var first = false; + if (r.rows) |rows| { + for (rows, 0..) |row, j| if (row != j) return error.TreesNotBuilt; + n = rows.len; + @memcpy(follow[i][0..n], r.follow[0..n]); + } else { + follow[i][0] = switch (r.first orelse return error.NoFirstToken) { + .handle => |h| try readFn(p, h), + .value => |v| v, + }; + first = true; + } + const depth: u32 = @min(r.depth, max_depth); + depths[i] = depth; + if (self.e.drafter != null) chooseDrafter(self, l, if (r.rows != null) n else 0); + const df = self.e.drafter != null and l.use_df; + // a stream DFlash2 drafts for still feeds the MTP head its rows (one step, its draft unused): it stays current + reqs[i] = .{ .s = l.s, .hidden = if (first) l.s.last_hidden else l.s.hidden, .next = follow[i][0..n], .count = if (df) 1 else depth, + .confidence = self.confidence, .out = l.held[0..depth], .probs = l.probs[0..depth] }; + try w.int(l.s.uid); + try w.tokens(follow[i][0..n]); + try w.int(depth); + try w.int(@intFromBool(first)); // 1: the first draft after the prompt (its hidden row: the prompt's last) + try w.int(@as(u64, @bitCast(self.confidence))); // the confidence cut (rank 1 must stop where rank 0 does: the head's steps gather) + try w.int(@intFromBool(df)); // 1: DFlash2 drafts this stream's round + } + const t0 = std.Io.Timestamp.now(self.io, .awake); + try self.send(.draft, &w); + // the MTP head: every stream's rows, the drafts of the streams it drafts for (none while DFlash2 drafts alone: TF_GLM_MIX=0) + if (feedsMtp(self)) try decode.draft(self.e, reqs[0..requests.len]); + if (self.e.drafter) |dr| { + var dq: [max_streams]dflash.Req = undefined; + for (requests, 0..) |r, i| { + const nf = reqs[i].next.len; + dq[i] = .{ .s = ls[i].s, .n = if (r.rows != null) nf else 0, .anchor = follow[i][nf - 1], .out = ls[i].held[0..max_depth], + .probs = ls[i].probs[0..max_depth], .drafts = ls[i].use_df }; + } + try dflash.draft(dr, self.e, dq[0..requests.len]); + for (reqs[0..requests.len], ls[0..requests.len], depths[0..requests.len]) |*q, l, depth| { // the confidence cut, as the head's chains stop + if (!l.use_df) continue; + var chain: f64 = 1; + q.got = 0; + while (q.got < @min(depth, dflash.BLOCK - 1)) : (q.got += 1) { // a block's 7 (a copy may run longer) + chain *= l.probs[q.got]; + if (q.got > 0 and (if (self.marginal > 0) l.probs[q.got] < self.marginal else chain < q.confidence)) break; + } + } + } + if (self.trace.on) { + self.trace.draft_ms += msSince(self.io, t0); + self.trace.calls += 1; + for (reqs[0..requests.len]) |q| self.trace.drafts += q.got; + self.trace.drafts_calls_total += requests.len; + } + for (requests, 0..) |r, i| { + const l = ls[i]; + const depth: u32 = @min(r.depth, max_depth); + const got = reqs[i].got; + for (got..depth) |j| { // past the confidence cut: held with chance 0 (the allocator leaves them out) + l.held[j] = if (got > 0) l.held[got - 1] else follow[i][reqs[i].next.len - 1]; + l.probs[j] = 0; + } + l.nheld = depth; + if (self.trace.on and r.rows != null and l.trace_held > 0) self.trace.landed += @min(l.trace_held, reqs[i].next.len - 1); + l.trace_held = 0; + if (l.copy) |*h| try copyDrafts(self, l, h, follow[i][0..reqs[i].next.len], r.rows != null, depth, got); + } +} + +/// Whether a draft round runs the MTP head: always without DFlash2, and with it only while the mix may pick the head +/// (both ranks read TF_GLM_MIX, so they agree). +pub fn feedsMtp(self: *const Owned) bool { + return self.e.drafter == null or (self.mix_on and self.e.w.mtp != null); +} + +/// The source for this stream's round (DFlash2 loaded): the one whose measured tokens per round over round time is +/// higher, the other one every 8th round (its numbers stay fresh); ``kept``: the last round's committed tokens (0: none). +fn chooseDrafter(self: *Owned, l: *Lane, kept: usize) void { + const now = std.Io.Timestamp.now(self.io, .awake); + if (!self.mix_on or self.e.w.mtp == null) { + l.use_df = true; + return; + } + if (kept > 0) if (l.mix_t) |t| { + const d: usize = @intFromBool(l.use_df); + const ms = @as(f64, @floatFromInt(t.durationTo(now).toNanoseconds())) / 1e6; + l.mix_tok[d] = 0.85 * l.mix_tok[d] + 0.15 * @as(f64, @floatFromInt(kept)); + l.mix_ms[d] = 0.85 * l.mix_ms[d] + 0.15 * ms; + }; + l.mix_t = now; + l.mix_round += 1; + const df_better = l.mix_tok[1] / l.mix_ms[1] >= l.mix_tok[0] / l.mix_ms[0]; + l.use_df = if (l.mix_round % 8 == 0) !df_better else df_better; + if (self.trace.on) self.trace.df_rounds += @intFromBool(l.use_df); +} + +/// A copy over the head's drafts when its expected landed tokens are more (rank 0: the verify frame carries the +/// held tokens to rank 1). ``follow``: the tokens just committed; ``kept_rows``: the last window's rows are known. +fn copyDrafts(self: *Owned, l: *Lane, h: *copy.Copy, follow: []const u32, kept_rows: bool, depth: u32, got: usize) !void { + const landed: u32 = @intCast(follow.len - 1); + if (kept_rows and l.copy_verified > 0) self.copy_rate.record(l.copy_verified, landed); + if (kept_rows and l.draft_verified > 0) self.draft_rate.record(l.draft_verified, landed); + l.copy_verified = 0; + l.draft_verified = 0; + l.copied = 0; + try h.append(self.gpa, follow); + if (depth == 0) return; + var buf: [max_depth]u32 = undefined; + const n = h.propose(buf[0..depth]); + if (n == 0) return; + const q = self.copy_rate.chance(); + // the copy agrees with every draft and runs on: the drafts stay, the copy's tail follows them + var same: usize = 0; + while (same < @min(n, got) and buf[same] == l.held[same]) same += 1; + const extends = same == got and n > got; + if (!extends) { // else the source whose drafts are expected to land more, by their measured rates + const mine = expected(q, n); + const theirs = expected(self.draft_rate.chance(), got); + if (mine <= theirs) return; + } + const keep = if (extends) got else 0; + @memcpy(l.held[keep..n], buf[keep..n]); + for (l.probs[keep..n]) |*pr| pr.* = q; + for (n..depth) |j| { + l.held[j] = buf[n - 1]; + l.probs[j] = 0; + } + l.copied = @intCast(n); + if (self.trace.on) { + self.trace.copies += 1; + self.trace.copy_tokens += n; + } +} + +/// Tokens a chain of ``n`` drafts lands on average when each lands with chance ``p`` given the ones before. +fn expected(p: f64, n: usize) f64 { + var sum: f64 = 0; + var chain: f64 = 1; + for (0..n) |_| { + chain *= p; + sum += chain; + } + return sum; +} + +/// TF_GLM_WINDOW: a stream's verify window at most, 2 to 16 rows (16: long copy drafts verified in one window). +fn windowFromEnv() u32 { + const v = std.c.getenv("TF_GLM_WINDOW") orelse return max_rows; + const n = std.fmt.parseInt(u32, std.mem.span(v), 10) catch return max_rows; + return @max(2, @min(n, max_rows)); +} + +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 drafterWanted() bool { + const v = std.c.getenv("TF_GLM_DRAFTER") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "dflash2"); +} + +/// TF_GLM_DFLASH2_STEP_MS: a DFlash2 round's cost spread over its 7 drafts (1.6 ms: about 11 ms a round). +fn dflashStepMs() f64 { + const v = std.c.getenv("TF_GLM_DFLASH2_STEP_MS") orelse return 1.6; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch 1.6; +} + +fn roceWanted() bool { + const v = std.c.getenv("TF_ROCE") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +fn roceMaxBytes() usize { + const v = std.c.getenv("TF_ROCE_MAX_KB") orelse return 512 << 10; + return (std.fmt.parseInt(usize, std.mem.span(v), 10) catch 512) << 10; +} + +fn msSince(io: std.Io, t0: std.Io.Timestamp) f64 { + return @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(io, .awake)).toNanoseconds())) / 1e6; +} + +fn probabilitiesFn(p: *anyopaque, s: *lanes.Stream, out: []f64) anyerror!bool { + const self = of(p); + const l = self.lanes_by.getPtr(s) orelse return false; + if (out.len > l.nheld) return false; + @memcpy(out, l.probs[0..out.len]); + return true; +} + +/// The stream's lane ends: its caches kept at the prompt's kept point for a later prompt (rank 0 decides, rank 1 hears +/// it), or freed. +fn releaseFn(p: *anyopaque, s: *lanes.Stream) void { + const self = of(p); + var l = self.lanes_by.get(s) orelse return; + _ = self.lanes_by.remove(s); + if (l.copy) |*h| h.deinit(self.gpa); + const seq = l.s; + const keep_it = seq.snap_at > 0 and !l.media and seq.snap_at <= l.prompt.len; + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + w.int(seq.uid) catch {}; + w.int(@intFromBool(keep_it)) catch {}; + self.send(.release, &w) catch {}; + if (keep_it) { + const ids = self.gpa.dupe(u32, l.prompt[0..seq.snap_at]) catch null; + if (ids) |x| self.kept.append(self.gpa, .{ .s = seq, .ids = x }) catch { + self.gpa.free(x); + self.dropSeq(seq.uid); + } else self.dropSeq(seq.uid); + } else self.dropSeq(seq.uid); + self.gpa.free(l.prompt); +} + +/// Rank 1: rank 0's calls, replayed in order on this rank's half of the model, until rank 0 stops. +fn followLoop(p: *anyopaque) anyerror!void { + const self = of(p); + var picks: [forward.STEP_ROWS]decode.Pick = undefined; + var held: [max_streams][max_depth]u32 = undefined; + var probs: [max_streams][max_depth]f64 = undefined; + while (true) { + const m = self.link.recv(self.gpa) catch |e| switch (e) { + error.PeerClosed => return, + else => return e, + }; + defer self.gpa.free(m.bytes); + var r: Reader = .{ .bytes = m.bytes }; + switch (@as(Op, @enumFromInt(m.tag))) { + .prefill => { + const uid = try r.int(); + const fresh: i64 = @intCast(try r.int()); + const toks = try r.tokens(); + const sampling = try proto.readSampling(&r); + const start: usize = @intCast(try r.int()); + const keep_at: usize = @intCast(try r.int()); + const seq = if (fresh > 0) try self.newSeq(uid, fresh) else self.seqs.get(uid) orelse return error.NoSuchSequence; + seq.sampling = sampling; + defer { // the image frame's features served this prefill alone + self.media_spans.clearRetainingCapacity(); + if (self.media_feats) |*f| f.free(); + self.media_feats = null; + } + if (self.media_feats) |f| forward.attach(seq, .{ .spans = self.media_spans.items, .feats = f.ptr }); + _ = try decode.prefill(self.e, seq, toks, start, keep_at); + }, + .image => { + _ = try r.int(); + const n: usize = @intCast(try r.int()); + self.media_spans.clearRetainingCapacity(); + for (0..n) |_| try self.media_spans.append(self.gpa, .{ @intCast(try r.int()), @intCast(try r.int()) }); + const rows: usize = @intCast(try r.int()); + const bytes = try r.raw(rows * 4096 * 2); + if (self.media_feats) |*f| f.free(); + self.media_feats = try cuda.DeviceBuffer.alloc(self.ctx.d, rows * 4096 * 2); + try self.media_feats.?.upload(0, bytes); + }, + .verify => { + const n: usize = @intCast(try r.int()); + var wins: [max_streams]decode.Win = undefined; + for (0..n) |i| { + const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; + wins[i] = .{ .s = s, .ids = try r.tokens() }; + } + try decode.verify(self.e, wins[0..n], &picks); + }, + .keep => { + const n: usize = @intCast(try r.int()); + for (0..n) |_| { + const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; + const rows: i64 = @intCast(try r.int()); + try decode.keep(self.e, s, rows, @intCast(try r.int())); + } + }, + .draft => { + const n: usize = @intCast(try r.int()); + var reqs: [max_streams]decode.DraftReq = undefined; + var dq: [max_streams]dflash.Req = undefined; + for (0..n) |i| { + const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; + const follow = try r.tokens(); + const depth: usize = @intCast(try r.int()); + const first = try r.int() == 1; + const cut: f64 = @bitCast(try r.int()); + const df = try r.int() == 1; + reqs[i] = .{ .s = s, .hidden = if (first) s.last_hidden else s.hidden, .next = follow, .count = if (df) 1 else depth, .confidence = cut, + .out = held[i][0..depth], .probs = probs[i][0..depth] }; + dq[i] = .{ .s = s, .n = if (first) 0 else follow.len, .anchor = follow[follow.len - 1], .drafts = df }; + } + if (feedsMtp(self)) try decode.draft(self.e, reqs[0..n]); // as rank 0: the MTP head first, then DFlash2 + if (self.e.drafter) |dr| try dflash.draft(dr, self.e, dq[0..n]); + }, + .release => { + const uid = try r.int(); + if (try r.int() == 0) self.dropSeq(uid); + }, + .evict => self.dropSeq(try r.int()), + .stop => { // rank 0 closes the link next: no lost peer + if (self.guard) |g| g.stop(self.gpa); + self.guard = null; + return; + }, + .ready, .room => {}, + } + } +} diff --git a/zig/src/families/glm_cuda/proto.zig b/zig/src/families/glm_cuda/proto.zig new file mode 100644 index 0000000000..6b3fafff1f --- /dev/null +++ b/zig/src/families/glm_cuda/proto.zig @@ -0,0 +1,64 @@ +//! GLM-5.3-Flash's two-rank protocol: the frames rank 0 sends rank 1 (each a call rank 1 replays) and their +//! little-endian encoding. + +const std = @import("std"); +const lanes = @import("lanes"); + +pub const Op = enum(u32) { prefill = 1, verify = 2, keep = 3, draft = 4, release = 5, stop = 6, ready = 7, image = 8, evict = 9, room = 10 }; + +pub const Writer = struct { + buf: std.ArrayList(u8) = .empty, + gpa: std.mem.Allocator, + pub fn int(w: *Writer, v: anytype) !void { + const T = @TypeOf(v); + const x: u64 = switch (@typeInfo(T)) { + .comptime_int => @intCast(v), + .int => |i| if (i.signedness == .signed) @bitCast(@as(i64, v)) else @as(u64, v), + .bool => @intFromBool(v), + else => @compileError("Writer.int: an integer"), + }; + try w.buf.appendSlice(w.gpa, std.mem.asBytes(&x)); + } + pub fn tokens(w: *Writer, t: []const u32) !void { + try w.int(t.len); + try w.buf.appendSlice(w.gpa, std.mem.sliceAsBytes(t)); + } +}; + +pub const Reader = struct { + bytes: []const u8, + at: usize = 0, + pub fn int(r: *Reader) !u64 { + if (r.at + 8 > r.bytes.len) return error.ShortFrame; + const v = std.mem.readInt(u64, r.bytes[r.at..][0..8], .little); + r.at += 8; + return v; + } + pub fn raw(r: *Reader, n: usize) ![]const u8 { + if (r.at + n > r.bytes.len) return error.ShortFrame; + r.at += n; + return r.bytes[r.at - n .. r.at]; + } + pub fn tokens(r: *Reader) ![]const u32 { + const n: usize = @intCast(try r.int()); + if (r.at + 4 * n > r.bytes.len) return error.ShortFrame; + const t: []align(1) const u32 = std.mem.bytesAsSlice(u32, r.bytes[r.at..][0 .. 4 * n]); + r.at += 4 * n; + return @alignCast(t); + } +}; + +pub fn writeSampling(w: *Writer, s: ?lanes.Sampling) !void { + const r = s orelse return w.int(0); + try w.int(1); + try w.int(r.seed); + try w.int(@as(u64, @bitCast(r.temperature))); + try w.int(r.top_k); + try w.int(@as(u64, @bitCast(r.top_p))); + try w.int(@as(u64, @bitCast(r.min_p))); +} + +pub fn readSampling(r: *Reader) !?lanes.Sampling { + if (try r.int() == 0) return null; + return .{ .seed = try r.int(), .temperature = @bitCast(try r.int()), .top_k = @intCast(try r.int()), .top_p = @bitCast(try r.int()), .min_p = @bitCast(try r.int()) }; +} diff --git a/zig/src/families/glm_cuda/root.zig b/zig/src/families/glm_cuda/root.zig new file mode 100644 index 0000000000..78fec41bfa --- /dev/null +++ b/zig/src/families/glm_cuda/root.zig @@ -0,0 +1,3 @@ +//! GLM-5.3-Flash (glm5_next) on CUDA, two ranks: the native family module (native.zig) and its parts. +pub const native = @import("native.zig"); +pub const api = @import("api.zig"); diff --git a/zig/src/families/glm_cuda/seq.zig b/zig/src/families/glm_cuda/seq.zig new file mode 100644 index 0000000000..a5d0ffa4a6 --- /dev/null +++ b/zig/src/families/glm_cuda/seq.zig @@ -0,0 +1,233 @@ +//! GLM-5.3-Flash on CUDA: a stream's committed caches (forward.py's State), sized to its request. KDA conv windows and +//! recurrent states (two parities), the main and MTP latent caches, the indexer's keys and gates (rings: only a +//! window's rows and the 3 before it are read) and pools, and the kept prompt state a later prompt can resume. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const kern = @import("kern.zig"); + +const D = kern.D; +const LL = kern.LL; + +pub const MAX_ROWS: i64 = 16; // a stream's verify window: the pending token and up to 15 drafts (long copies) +pub const PROJ_W: i64 = 3 * LL * 128 + 128 + 128 + LL; // q | k | v | f_a | g_a | b +pub const CONV_C: i64 = 3 * LL * 128; // the conv's channels: q | k | v +pub const REC: i64 = LL * 128 * 128; // a layer's recurrent state (fp32 values) + +/// Image features for a prompt: ``spans`` [begin, end) prompt positions of each image's placeholder run, in order, +/// and ``feats`` their rows [sum of runs, 4096] bf16 on the device (the caller's, live through the prefill). +pub const Media = struct { spans: []const [2]i64, feats: u64 }; + +/// Rows of the index rings (TF_GLM_INDEX_RING, default 16384, at least 4096; the capacity when that is smaller). +pub fn ringRows(capacity: i64) i64 { + var rows: i64 = 16384; + if (std.c.getenv("TF_GLM_INDEX_RING")) |v| rows = std.fmt.parseInt(i64, std.mem.span(v), 10) catch rows; + return @min(capacity, @max(4096, rows)); +} + +/// Rows kept before a rebased window's first pool: a rejected draft's window starts up to MAX_ROWS rows back. +const RING_SLACK: i64 = 2 * MAX_ROWS; + +pub const Seq = struct { + uid: u64 = 0, + capacity: i64, + nkda: usize, + ndsa: usize, // the DSA layers' latent caches + caches: usize, // all latent caches: the DSA layers', then the MTP layer's (when the model has one) + pos: i64 = 0, + mtp_len: i64 = 0, + mtp_drafted: i64 = 0, + cur: u1 = 0, + mem: cuda.DeviceBuffer = undefined, + conv: u64 = 0, // [nkda, 3, CONV_C] bf16 + rec: u64 = 0, // [2, nkda, LL, 128, 128] fp32 + proj: u64 = 0, // [nkda, MAX_ROWS, PROJ_W] bf16: a window's conv channels (the commit's conv shift reads them) + ks: kern.KdaScratch = undefined, // [nkda, MAX_ROWS, ...] the window's replay inputs + kc: [12]u64 = @splat(0), // latent caches [capacity, 512] bf16 + ik: [12]u64 = @splat(0), // index keys [ring, 128] bf16: rows ring_base.. of the positions + ig: [12]u64 = @splat(0), // index gates [ring, 128] bf16 + pk: [12]u64 = @splat(0), // pooled keys [capacity / 4 + 2, 128] bf16 + /// The index rows' first position (a multiple of 4, less the slack): the DSA layers' (0) and the MTP layer's (1). + ring_base: [2]i64 = .{ 0, 0 }, + ring: i64 = 0, + /// The prompt's images (attach): their placeholder runs' positions and features, read by the prompt's chunks. + media: ?Media = null, + /// The last prompt chunk's final-normed last row, the first draft's input. + last_hidden: u64 = 0, + /// The row a window's final-normed rows start at in the step that ran it (the next draft's inputs). + hidden: u64 = 0, + /// The reply's draw rule (null or temperature 0: greedy); drafts stay greedy (they only propose). + sampling: ?lanes.Sampling = null, + /// The kept prompt state: KDA states and conv windows at ``snap_at``, the MTP head's pending row and its cache + /// length, each cache's index rows of the point's incomplete pool. The attention rows below the point stay valid + /// in the live caches while the next prompt extends the kept prefix. 0: none. + snap_at: usize = 0, + snap_mtp_len: i64 = 0, + snap_rec: u64 = 0, // [nkda, LL, 128, 128] fp32 + snap_conv: u64 = 0, // [nkda, 3, CONV_C] bf16 + snap_pending: u64 = 0, // [D] bf16 + snap_ik: u64 = 0, // [12, 3, 128] bf16 + snap_ig: u64 = 0, + // DFlash2 (dflash.zig): the context ring, its first valid position, and the last verify's rows (the step it + // ran in, the rows' first row in it, their first position, how many) + draft: u64 = 0, + d_from: i64 = 0, + v_step: u64 = 0, + v_row0: i64 = 0, + v_pos0: i64 = 0, + v_rows: usize = 0, + + pub fn sampled(s: *const Seq) ?lanes.Sampling { + const r = s.sampling orelse return null; + return if (r.temperature > 0) r else null; + } + + pub fn pools(s: *const Seq) i64 { + return @divFloor(s.capacity, 4); + } + + pub fn recAt(s: *const Seq, parity: u1, li: usize) u64 { + return s.rec + (@as(u64, parity) * s.nkda + li) * @as(u64, @intCast(REC * 4)); + } + + pub fn convAt(s: *const Seq, li: usize) u64 { + return s.conv + li * @as(u64, @intCast(3 * CONV_C * 2)); + } + + /// Cache ``ci``'s ring as the kernels address it (position p at row p - ring_base). + pub fn ringPtr(s: *const Seq, ring: u64, ci: usize) u64 { + const which: usize = if (ci < s.ndsa) 0 else 1; + return ring -% @as(u64, @intCast(s.ring_base[which] * kern.IDX_DIM * 2)); + } + + pub fn kdaScratch(s: *const Seq, li: usize) kern.KdaScratch { + const l: u64 = @intCast(li); + const rows: u64 = @intCast(MAX_ROWS); + const hw: u64 = @intCast(LL * 128); + return .{ .out = s.ks.out + l * rows * hw * 2, .k = s.ks.k + l * rows * hw * 4, .v = s.ks.v + l * rows * hw * 2, + .g = s.ks.g + l * rows * hw * 4, .b = s.ks.b + l * rows * @as(u64, @intCast(LL)) * 4 }; + } +}; + +/// Device bytes a stream of ``capacity`` cache slots takes. +pub fn bytes(nkda: usize, caches: usize, capacity: i64, extra: u64) u64 { + const n: i64 = @intCast(nkda); + const c: i64 = @intCast(caches); + const per = n * 3 * CONV_C * 2 + 2 * n * REC * 4 + n * MAX_ROWS * PROJ_W * 2 + n * MAX_ROWS * LL * 128 * (2 + 4 + 2 + 4) + n * MAX_ROWS * LL * 4 + + c * (capacity * kern.KV_LORA * 2 + 2 * ringRows(capacity) * kern.IDX_DIM * 2 + (@divFloor(capacity, 4) + 2) * kern.IDX_DIM * 2) + D * 2 + + n * REC * 4 + n * 3 * CONV_C * 2 + D * 2 + 2 * 12 * 3 * kern.IDX_DIM * 2; // the kept state + return @as(u64, @intCast(per + 64 * 512)) + extra; +} + +const Carve = struct { + at: u64 = 0, + fn take(self: *Carve, n: i64) u64 { + const p = self.at; + self.at += std.mem.alignForward(u64, @intCast(@max(n, 1)), 512); + return p; + } +}; + +/// ``extra``: bytes of a drafter's ring (0: none). +pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, nkda: usize, ndsa: usize, caches: usize, capacity: i64, extra: u64) !*Seq { + const s = try gpa.create(Seq); + errdefer gpa.destroy(s); + s.* = .{ .capacity = capacity, .nkda = nkda, .ndsa = ndsa, .caches = caches, .ring = ringRows(capacity) }; + var c: Carve = .{}; + const n: i64 = @intCast(nkda); + s.conv = c.take(n * 3 * CONV_C * 2); + s.rec = c.take(2 * n * REC * 4); + s.proj = c.take(n * MAX_ROWS * PROJ_W * 2); + s.ks = .{ .out = c.take(n * MAX_ROWS * LL * 128 * 2), .k = c.take(n * MAX_ROWS * LL * 128 * 4), .v = c.take(n * MAX_ROWS * LL * 128 * 2), + .g = c.take(n * MAX_ROWS * LL * 128 * 4), .b = c.take(n * MAX_ROWS * LL * 4) }; + for (0..caches) |i| { + s.kc[i] = c.take(capacity * kern.KV_LORA * 2); + s.ik[i] = c.take(s.ring * kern.IDX_DIM * 2); + s.ig[i] = c.take(s.ring * kern.IDX_DIM * 2); + s.pk[i] = c.take((@divFloor(capacity, 4) + 2) * kern.IDX_DIM * 2); + } + s.last_hidden = c.take(D * 2); + s.snap_rec = c.take(n * REC * 4); + s.snap_conv = c.take(n * 3 * CONV_C * 2); + s.snap_pending = c.take(D * 2); + s.snap_ik = c.take(12 * 3 * kern.IDX_DIM * 2); + s.snap_ig = c.take(12 * 3 * kern.IDX_DIM * 2); + if (extra > 0) s.draft = c.take(@intCast(extra)); + s.mem = try cuda.DeviceBuffer.alloc(d, c.at); + const base = s.mem.ptr; + inline for (.{ "conv", "rec", "proj", "last_hidden", "snap_rec", "snap_conv", "snap_pending", "snap_ik", "snap_ig" }) |f| @field(s, f) += base; + if (extra > 0) s.draft += base; + inline for (.{ "out", "k", "v", "g", "b" }) |f| @field(s.ks, f) += base; + for (0..caches) |i| { + s.kc[i] += base; + s.ik[i] += base; + s.ig[i] += base; + s.pk[i] += base; + } + try d.check(d.api.cuMemsetD8_v2(s.mem.ptr, 0, c.at), "zero sequence"); + return s; +} + +pub fn free(gpa: std.mem.Allocator, s: *Seq) void { + s.mem.free(); + gpa.destroy(s); +} + +/// State.reset: the conv windows and recurrent states zeroed, positions 0 (the caches' stale rows stay: as Python). +pub fn reset(k: *kern.K, s: *Seq) !void { + const d = k.d; + const n: i64 = @intCast(s.nkda); + try d.check(d.api.cuMemsetD8Async(s.conv, 0, @intCast(n * 3 * CONV_C * 2), k.stream.handle), "reset conv"); + try d.check(d.api.cuMemsetD8Async(s.rec, 0, @intCast(2 * n * REC * 4), k.stream.handle), "reset rec"); + s.cur = 0; + s.pos = 0; + s.mtp_len = 0; + s.mtp_drafted = 0; + s.ring_base = .{ 0, 0 }; +} + +fn ringCaches(s: *const Seq, which: u1) []const usize { + const all = [_]usize{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 }; + return if (which == 0) all[0..s.ndsa] else all[s.ndsa..s.caches]; +} + +/// The index ring ``which`` (0: the DSA layers', 1: the MTP layer's) holding rows host_pos - host_pos % 4 .. +/// host_pos + R: moved down (the rows before the window kept) when the window runs past its end. The kernels read +/// the rings through the step's tables, so nothing captured holds the old base. +pub fn ensureRing(k: *kern.K, s: *Seq, which: u1, host_pos: i64, R: i64) !void { + const lo = host_pos - @mod(host_pos, 4); + const base = s.ring_base[which]; + if (lo >= base and host_pos + R - base <= s.ring) return; + const nb = @max(0, lo - RING_SLACK); + if (nb < base or host_pos - base > s.ring or host_pos + R - nb > s.ring) return error.IndexRowsLost; + const kept: u64 = @intCast(host_pos - nb); + const row: u64 = @intCast(kern.IDX_DIM * 2); + if (kept > 0) for (ringCaches(s, which)) |ci| { + const from: u64 = @intCast(nb - base); + try k.copy(s.ik[ci], s.ik[ci] + from * row, kept * row); + try k.copy(s.ig[ci], s.ig[ci] + from * row, kept * row); + }; + s.ring_base[which] = nb; +} + +/// The kept point's incomplete pool rows (positions at - at % 4 .. at) of ring ``which`` -> the snapshot (or back). +pub fn ringSnap(k: *kern.K, s: *Seq, which: u1, at: i64, restore: bool) !void { + const lo = at - @mod(at, 4); + const n: u64 = @intCast(at - lo); + const row: u64 = @intCast(kern.IDX_DIM * 2); + if (restore) s.ring_base[which] = lo; + if (n == 0) return; + const off: u64 = @intCast(lo - s.ring_base[which]); + for (ringCaches(s, which)) |ci| { + const slot = s.snap_ik + ci * 3 * row; + const gslot = s.snap_ig + ci * 3 * row; + if (restore) { + try k.copy(s.ik[ci] + off * row, slot, n * row); + try k.copy(s.ig[ci] + off * row, gslot, n * row); + } else { + try k.copy(slot, s.ik[ci] + off * row, n * row); + try k.copy(gslot, s.ig[ci] + off * row, n * row); + } + } +} diff --git a/zig/src/families/glm_cuda/vision.zig b/zig/src/families/glm_cuda/vision.zig new file mode 100644 index 0000000000..8eb6c15bc9 --- /dev/null +++ b/zig/src/families/glm_cuda/vision.zig @@ -0,0 +1,389 @@ +//! GLM-5.3-Flash's vision tower on CUDA (rank 0 encodes; rank 1 receives the features), as mlx-vlm's +//! glm5_next/vision.py defines it: the patch embedding (a linear over a patch's 2 x 14 x 14 x 3 values), 24 blocks +//! (RMSNorm, attention with RMSNorm'd q and k and a 2-D rotary over each frame's patches, a limited SwiGLU), the +//! post-norm, the 2 x 2 downsample (a linear over each merge block's four patches) and the merger (a projection, +//! LayerNorm + GELU, a limited SwiGLU) to the language model's 4096. The linears are cuBLASLt products (bf16 in, fp32 +//! accumulate, bias added before the one rounding); the residual stream stays fp32; glm/vision.cu does the rest. +//! Checked against the torch reference (tools/vision_ref.py) by tf-cuda-test glm-vision. + +const std = @import("std"); +const cuda = @import("cuda"); +const kern = @import("kern.zig"); + +const H: i64 = 1024; // hidden +const HEADS: i64 = 16; +const HD: i64 = 64; // a head's width +const MLP: i64 = 4096; +const PATCH: i64 = 1176; // 3 channels x 2 frames x 14 x 14 +const MERGE: i64 = 2; +const OUT: i64 = 4096; // the merged rows' width (the language model's) +const PROJ: i64 = 10240; // the merger's SwiGLU width +pub const DEPTH = 24; +const CHUNK: i64 = 2048; // query rows a scores product covers +const ROWS: i64 = 4096; // rows a linear's fp32 product covers (the scratch's bound) +const EPS: f32 = 1e-5; +const LIMIT: f32 = 10.0; + +const Block = struct { n1: u64, n2: u64, qkv_w: u64, qkv_b: u64, qn: u64, kn: u64, proj_w: u64, proj_b: u64, g_w: u64, g_b: u64, u_w: u64, u_b: u64, d_w: u64, d_b: u64 }; + +const LinKey = struct { m: u64, n: u64, k: u64, f32_out: bool }; + +pub const Tower = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + lt: cuda.cublaslt.Library, + mem: cuda.DeviceBuffer, + blocks: [DEPTH]Block = undefined, + patch_w: u64 = 0, + patch_b: u64 = 0, + post: u64 = 0, + down_w: u64 = 0, // [4096, 2, 2, 1024]: the conv's taps in the merge block's (row, col) order, then the channels + down_b: u64 = 0, + m_proj: u64 = 0, m_ln_w: u64 = 0, m_ln_b: u64 = 0, m_gate: u64 = 0, m_up: u64 = 0, m_down: u64 = 0, + inv_freq: cuda.DeviceBuffer, + linears: std.AutoHashMap(LinKey, cuda.cublaslt.Linear), + ws: cuda.DeviceBuffer, + ws_len: usize = 64 << 20, + bytes: u64 = 0, + /// Development: the fp32 stream after the patch embedding and after block 0 (tf-cuda-test glm-vision). + dbg_embed: ?[]f32 = null, + dbg_block0: ?[]f32 = null, + dbg_blocks: ?[]f32 = null, // the stream after the last block + dbg_all: ?[]f32 = null, // after every block + dbg_down: ?[]u16 = null, // the downsample's rows (bf16) + + /// The tower's tensors (model.visual.*, bf16) from the checkpoint's shards in ``dir``, in one allocation. + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8) !Tower { + const ipath = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(ipath); + const itext = try std.Io.Dir.cwd().readFileAlloc(io, ipath, gpa, .limited(64 << 20)); + defer gpa.free(itext); + const index = try std.json.parseFromSlice(std.json.Value, gpa, itext, .{}); + defer index.deinit(); + const wm = index.value.object.get("weight_map").?.object; + // the shards holding the tower, and its tensors' places in them + const Ent = struct { file: []const u8, off: u64, n: u64, shape: [5]i64, nd: usize }; + var ents = std.StringHashMap(Ent).init(gpa); + defer ents.deinit(); + var files = std.StringHashMap(void).init(gpa); + defer files.deinit(); + var it = wm.iterator(); + while (it.next()) |e| if (std.mem.startsWith(u8, e.key_ptr.*, "model.visual.")) try files.put(e.value_ptr.string, {}); + var headers: std.ArrayList(std.json.Parsed(std.json.Value)) = .empty; + defer { + for (headers.items) |*h| h.deinit(); + headers.deinit(gpa); + } + var total: u64 = 0; + var fit = files.keyIterator(); + while (fit.next()) |fname| { + const path = try std.fs.path.join(gpa, &.{ dir, fname.* }); + defer gpa.free(path); + const file = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer file.close(io); + var head: [8]u8 = undefined; + _ = try file.readPositionalAll(io, &head, 0); + const hlen = std.mem.readInt(u64, &head, .little); + const text = try gpa.alloc(u8, hlen); + defer gpa.free(text); + _ = try file.readPositionalAll(io, text, 8); + try headers.append(gpa, try std.json.parseFromSlice(std.json.Value, gpa, text, .{})); + var hit = headers.items[headers.items.len - 1].value.object.iterator(); + while (hit.next()) |e| { + if (!std.mem.startsWith(u8, e.key_ptr.*, "model.visual.")) continue; + if (!std.mem.eql(u8, e.value_ptr.object.get("dtype").?.string, "BF16")) return error.UnsupportedVisionDType; + const offs = e.value_ptr.object.get("data_offsets").?.array.items; + var ent: Ent = .{ .file = fname.*, .off = 8 + hlen + @as(u64, @intCast(offs[0].integer)), .n = @intCast(offs[1].integer - offs[0].integer), .shape = @splat(0), .nd = 0 }; + for (e.value_ptr.object.get("shape").?.array.items, 0..) |x, i| ent.shape[i] = x.integer; + ent.nd = e.value_ptr.object.get("shape").?.array.items.len; + try ents.put(e.key_ptr.*, ent); + total += std.mem.alignForward(u64, ent.n, 256); + } + } + var t: Tower = .{ .gpa = gpa, .d = d, .lt = try cuda.cublaslt.Library.open(), .mem = try cuda.DeviceBuffer.alloc(d, total), + .inv_freq = try cuda.DeviceBuffer.alloc(d, 16 * 4), .linears = .init(gpa), .ws = undefined, .bytes = total }; + errdefer t.deinit(); + t.ws = try cuda.DeviceBuffer.alloc(d, t.ws_len); + var at: u64 = 0; + var names = std.StringHashMap(u64).init(gpa); + defer names.deinit(); + var buf: std.ArrayList(u8) = .empty; + defer buf.deinit(gpa); + var eit = ents.iterator(); + while (eit.next()) |e| { + const ent = e.value_ptr.*; + try buf.resize(gpa, ent.n); + const path = try std.fs.path.join(gpa, &.{ dir, ent.file }); + defer gpa.free(path); + const file = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer file.close(io); + _ = try file.readPositionalAll(io, buf.items, ent.off); + if (std.mem.eql(u8, e.key_ptr.*, "model.visual.downsample.weight")) { // [o, c, 2, 2] -> [o, 2, 2, c] + if (ent.nd != 4 or ent.shape[0] != OUT or ent.shape[1] != H) return error.BadVisionShape; + const src = std.mem.bytesAsSlice(u16, buf.items); + const perm = try gpa.alloc(u16, src.len); + defer gpa.free(perm); + const c_n: usize = @intCast(H); + for (0..@intCast(OUT)) |o| for (0..c_n) |c| for (0..2) |kh| for (0..2) |kw| { + perm[((o * 2 + kh) * 2 + kw) * c_n + c] = src[((o * c_n + c) * 2 + kh) * 2 + kw]; + }; + @memcpy(buf.items, std.mem.sliceAsBytes(perm)); + } + const ptr = t.mem.ptr + at; + try d.check(d.api.cuMemcpyHtoD_v2(ptr, buf.items.ptr, ent.n), "vision weights"); + try names.put(e.key_ptr.*, ptr); + at += std.mem.alignForward(u64, ent.n, 256); + } + const get = struct { + fn f(m: *const std.StringHashMap(u64), comptime fmt: []const u8, args: anytype) !u64 { + var nb: [128]u8 = undefined; + const name = try std.fmt.bufPrint(&nb, "model.visual." ++ fmt, args); + return m.get(name) orelse { + std.log.err("glm vision: no tensor {s}", .{name}); + return error.MissingVisionTensor; + }; + } + }.f; + t.patch_w = try get(&names, "patch_embed.proj.weight", .{}); + t.patch_b = try get(&names, "patch_embed.proj.bias", .{}); + t.post = try get(&names, "post_layernorm.weight", .{}); + t.down_w = try get(&names, "downsample.weight", .{}); + t.down_b = try get(&names, "downsample.bias", .{}); + t.m_proj = try get(&names, "merger.proj.weight", .{}); + t.m_ln_w = try get(&names, "merger.post_projection_norm.weight", .{}); + t.m_ln_b = try get(&names, "merger.post_projection_norm.bias", .{}); + t.m_gate = try get(&names, "merger.gate_proj.weight", .{}); + t.m_up = try get(&names, "merger.up_proj.weight", .{}); + t.m_down = try get(&names, "merger.down_proj.weight", .{}); + for (&t.blocks, 0..) |*bl, i| { + bl.* = .{ + .n1 = try get(&names, "blocks.{d}.norm1.weight", .{i}), .n2 = try get(&names, "blocks.{d}.norm2.weight", .{i}), + .qkv_w = try get(&names, "blocks.{d}.attn.qkv.weight", .{i}), .qkv_b = try get(&names, "blocks.{d}.attn.qkv.bias", .{i}), + .qn = try get(&names, "blocks.{d}.attn.q_norm.weight", .{i}), .kn = try get(&names, "blocks.{d}.attn.k_norm.weight", .{i}), + .proj_w = try get(&names, "blocks.{d}.attn.proj.weight", .{i}), .proj_b = try get(&names, "blocks.{d}.attn.proj.bias", .{i}), + .g_w = try get(&names, "blocks.{d}.mlp.gate_proj.weight", .{i}), .g_b = try get(&names, "blocks.{d}.mlp.gate_proj.bias", .{i}), + .u_w = try get(&names, "blocks.{d}.mlp.up_proj.weight", .{i}), .u_b = try get(&names, "blocks.{d}.mlp.up_proj.bias", .{i}), + .d_w = try get(&names, "blocks.{d}.mlp.down_proj.weight", .{i}), .d_b = try get(&names, "blocks.{d}.mlp.down_proj.bias", .{i}), + }; + } + // the rotary's frequencies: 1 / 10000^(i / 16), i = 0, 2 .. 30 over 32 (head width / 2) + var inv: [16]f32 = undefined; + for (&inv, 0..) |*v, i| v.* = 1.0 / std.math.pow(f32, 10000.0, @as(f32, @floatFromInt(2 * i)) / 32.0); + try t.inv_freq.upload(0, std.mem.sliceAsBytes(&inv)); + return t; + } + + pub fn deinit(t: *Tower) void { + var it = t.linears.valueIterator(); + while (it.next()) |l| l.deinit(); + t.linears.deinit(); + t.mem.free(); + t.inv_freq.free(); + t.ws.free(); + t.lt.close(); + } + + /// D[m, n] = X[m, k] . W[n, k]^T (fp32 or bf16 out), a plan per shape kept for the next image. + fn linear(t: *Tower, k: *kern.K, x: u64, w: u64, d: u64, m: i64, n: i64, kk: i64, f32_out: bool) !void { + const key: LinKey = .{ .m = @intCast(m), .n = @intCast(n), .k = @intCast(kk), .f32_out = f32_out }; + const gop = try t.linears.getOrPut(key); + if (!gop.found_existing) { + gop.value_ptr.* = cuda.cublaslt.Linear.initTyped(&t.lt, key.m, key.n, key.k, .bf16, if (f32_out) .f32 else .bf16, t.ws_len) catch |e| { + _ = t.linears.remove(key); + return e; + }; + } + try gop.value_ptr.run(x, w, d, t.ws.ptr, t.ws_len, k.stream.handle); + } + + fn grid1(n: i64) cuda.Dim3 { + return .{ .x = @intCast(@min(std.math.divCeil(i64, n, 256) catch unreachable, 65535)) }; + } + + /// out = bf16(x . W^T + bias) (bias 0: none), in row chunks through ``acc``; with ``x32``, x32 += that instead. + fn dense(t: *Tower, k: *kern.K, x: u64, w: u64, bias: u64, out: u64, x32: u64, acc: u64, m: i64, n: i64, kk: i64) !void { + var r: i64 = 0; + while (r < m) : (r += ROWS) { + const rows = @min(ROWS, m - r); + try t.linear(k, x + @as(u64, @intCast(r * kk * 2)), w, acc, rows, n, kk, true); + var a: cuda.Args = .{}; + a.add(acc); a.add(bias); a.add(if (out != 0) out + @as(u64, @intCast(r * n * 2)) else 0); + a.add(if (x32 != 0) x32 + @as(u64, @intCast(r * n * 4)) else 0); a.add(@as(i64, rows * n)); a.add(@as(i32, @intCast(n))); + try k.go(try k.ext(.vision, "glm_vis_bias"), grid1(rows * n), 256, 0, &a); + } + } + + /// The limited SwiGLU of two linears (biases 0: none) over x [m, kk] -> out [m, n] bf16, in row chunks. + fn swiglu(t: *Tower, k: *kern.K, x: u64, gw: u64, gb: u64, uw: u64, ub: u64, out: u64, acc: u64, acc2: u64, m: i64, n: i64, kk: i64) !void { + var r: i64 = 0; + const step = @divFloor(ROWS * MLP, n); // the fp32 products' rows within the scratch + while (r < m) : (r += step) { + const rows = @min(step, m - r); + const xr = x + @as(u64, @intCast(r * kk * 2)); + try t.linear(k, xr, gw, acc, rows, n, kk, true); + try t.linear(k, xr, uw, acc2, rows, n, kk, true); + var a: cuda.Args = .{}; + a.add(acc); a.add(acc2); a.add(gb); a.add(ub); a.add(out + @as(u64, @intCast(r * n * 2))); a.add(@as(i64, rows * n)); + a.add(@as(i32, @intCast(n))); a.add(LIMIT); + try k.go(try k.ext(.vision, "glm_vis_swiglu"), grid1(rows * n), 256, 0, &a); + } + } + + fn rmsnorm(k: *kern.K, x32: u64, w: u64, y: u64, rows: i64, cols: i64) !void { + var a: cuda.Args = .{}; + a.add(x32); a.add(w); a.add(y); a.add(@as(i32, @intCast(cols))); a.add(EPS); + try k.go(try k.ext(.vision, "glm_vis_rmsnorm"), .{ .x = @intCast(rows) }, 256, 0, &a); + } + + /// Scratch bytes ``encode`` takes for ``patches`` patches in segments of ``seg``. + pub fn scratchBytes(patches: i64, seg: i64) u64 { + const N = patches; + const q = @min(CHUNK, seg); + return @intCast(N * PATCH * 4 + N * H * 4 + N * H * 2 * 5 + N * MLP * 2 + 2 * ROWS * MLP * 4 + q * seg * 6 + q * HD * 4 + q * 4 + N * 8 + 16 * 1024); + } + + /// Features of one image or video (``grid`` [t, h, w] in patches, t frame groups) from its patches (fp32, + /// [t h w, 1176] in the processor's merge-block order) into ``out`` ([t h w / 4, 4096] bf16). + pub fn encode(t: *Tower, k: *kern.K, patches: []const f32, grid: [3]i64, out: u64) !void { + const d = t.d; + const gt = grid[0]; + const gh = grid[1]; + const gw = grid[2]; + if (gt < 1 or gh < MERGE or gw < MERGE or @mod(gh, MERGE) != 0 or @mod(gw, MERGE) != 0) return error.BadImageGrid; + const seg = gh * gw; + const N = gt * seg; + if (@as(i64, @intCast(patches.len)) != N * PATCH) return error.BadImagePatches; + // each patch's (row, col) in merge-block order, every frame group alike + const pos = try t.gpa.alloc(i32, @intCast(N * 2)); + defer t.gpa.free(pos); + var i: i64 = 0; + while (i < N) : (i += 1) { + const j = @mod(i, seg); + const in_col = @mod(j, MERGE); + const in_row = @mod(@divTrunc(j, MERGE), MERGE); + const block_col = @mod(@divTrunc(j, MERGE * MERGE), @divTrunc(gw, MERGE)); + const block_row = @divTrunc(j, MERGE * MERGE * @divTrunc(gw, MERGE)); + pos[@intCast(i * 2)] = @intCast(block_row * MERGE + in_row); + pos[@intCast(i * 2 + 1)] = @intCast(block_col * MERGE + in_col); + } + var scratch = try cuda.DeviceBuffer.alloc(d, scratchBytes(N, seg)); + defer scratch.free(); + var at: u64 = scratch.ptr; + const take = struct { + fn f(p: *u64, n: i64) u64 { + const r = p.*; + p.* += std.mem.alignForward(u64, @intCast(n), 256); + return r; + } + }.f; + const q_rows = @min(CHUNK, seg); + const p32 = take(&at, N * PATCH * 4); + const x32 = take(&at, N * H * 4); + const xn = take(&at, N * H * 2); + const qh = take(&at, N * H * 2); + const kh = take(&at, N * H * 2); + const vt = take(&at, N * H * 2); + const att = take(&at, N * H * 2); + const big = take(&at, N * MLP * 2); + const acc = take(&at, ROWS * MLP * 4); + const acc2 = take(&at, ROWS * MLP * 4); + const sc = take(&at, q_rows * seg * 4); + const pr = take(&at, q_rows * seg * 2); + const oh = take(&at, q_rows * HD * 4); + const sums = take(&at, q_rows * 4); + const dpos = take(&at, N * 8); + try d.check(d.api.cuMemcpyHtoDAsync_v2(dpos, pos.ptr, pos.len * 4, k.stream.handle), "vision positions"); + try d.check(d.api.cuMemcpyHtoDAsync_v2(p32, patches.ptr, patches.len * 4, k.stream.handle), "vision patches"); + { // the patch embedding: the patches as bf16, then the conv as a linear (with its bias) into the fp32 stream + var a: cuda.Args = .{}; + a.add(p32); a.add(big); a.add(@as(i64, N * PATCH)); + try k.go(try k.ext(.vision, "glm_vis_to_bf16"), grid1(N * PATCH), 256, 0, &a); + var r: i64 = 0; + while (r < N) : (r += ROWS) { + const rows = @min(ROWS, N - r); + try t.linear(k, big + @as(u64, @intCast(r * PATCH * 2)), t.patch_w, acc, rows, H, PATCH, true); + var b: cuda.Args = .{}; + b.add(acc); b.add(t.patch_b); b.add(x32 + @as(u64, @intCast(r * H * 4))); b.add(@as(i64, rows * H)); b.add(@as(i32, @intCast(H))); + try k.go(try k.ext(.vision, "glm_vis_embed32"), grid1(rows * H), 256, 0, &b); + } + } + const scale: f32 = 1.0 / @sqrt(@as(f32, HD)); + if (t.dbg_embed) |dst| { + try k.stream.synchronize(); + try d.check(d.api.cuMemcpyDtoH_v2(@ptrCast(dst.ptr), x32, dst.len * 4), "vision debug"); + } + for (t.blocks, 0..) |bl, bi| { + defer if (t.dbg_all) |dst| { + const n: usize = @intCast(N * H); + k.stream.synchronize() catch {}; + d.check(d.api.cuMemcpyDtoH_v2(@ptrCast(dst[bi * n ..].ptr), x32, n * 4), "vision debug") catch {}; + }; + defer if (bi == 0) if (t.dbg_block0) |dst| { + k.stream.synchronize() catch {}; + d.check(d.api.cuMemcpyDtoH_v2(@ptrCast(dst.ptr), x32, dst.len * 4), "vision debug") catch {}; + }; + try rmsnorm(k, x32, bl.n1, xn, N, H); + try t.dense(k, xn, bl.qkv_w, bl.qkv_b, big, 0, acc, N, 3 * H, H); + { + var a: cuda.Args = .{}; + a.add(big); a.add(bl.qn); a.add(bl.kn); a.add(dpos); a.add(t.inv_freq.ptr); a.add(qh); a.add(kh); a.add(vt); + a.add(@as(i32, @intCast(N))); a.add(@as(i32, @intCast(seg))); a.add(@as(i32, @intCast(HEADS))); a.add(EPS); + try k.go(try k.ext(.vision, "glm_vis_qkv"), .{ .x = @intCast(std.math.divCeil(i64, N * HEADS, 8) catch unreachable) }, 256, 0, &a); + } + // each frame group attends within itself; per head, chunks of queries (q, k: [t, H, seg, 64]; v: [t, H, 64, seg]) + var g: i64 = 0; + while (g < gt) : (g += 1) { + var h: i64 = 0; + while (h < HEADS) : (h += 1) { + const blk: u64 = @intCast((g * HEADS + h) * seg * HD * 2); + var c: i64 = 0; + while (c < seg) : (c += CHUNK) { + const rows = @min(CHUNK, seg - c); + try t.linear(k, qh + blk + @as(u64, @intCast(c * HD * 2)), kh + blk, sc, rows, seg, HD, true); + var a: cuda.Args = .{}; + a.add(sc); a.add(pr); a.add(sums); a.add(@as(i32, @intCast(seg))); a.add(scale); + try k.go(try k.ext(.vision, "glm_vis_softmax"), .{ .x = @intCast(rows) }, 256, 0, &a); + try t.linear(k, pr, vt + blk, oh, rows, HD, seg, true); + var b: cuda.Args = .{}; + b.add(oh); b.add(sums); b.add(att + @as(u64, @intCast((g * seg + c) * H * 2))); b.add(@as(i32, @intCast(rows))); + b.add(@as(i32, @intCast(HEADS))); b.add(@as(i32, @intCast(h))); + try k.go(try k.ext(.vision, "glm_vis_head_out"), grid1(rows * HD), 256, 0, &b); + } + } + } + try t.dense(k, att, bl.proj_w, bl.proj_b, 0, x32, acc, N, H, H); + try rmsnorm(k, x32, bl.n2, xn, N, H); + try t.swiglu(k, xn, bl.g_w, bl.g_b, bl.u_w, bl.u_b, big, acc, acc2, N, MLP, H); + try t.dense(k, big, bl.d_w, bl.d_b, 0, x32, acc, N, H, MLP); + } + if (t.dbg_blocks) |dst| { + try k.stream.synchronize(); + try d.check(d.api.cuMemcpyDtoH_v2(@ptrCast(dst.ptr), x32, dst.len * 4), "vision debug"); + } + // the post-norm, then each merge block's four patches (adjacent rows) through the 2 x 2 conv as one linear + try rmsnorm(k, x32, t.post, xn, N, H); + const M = @divExact(N, MERGE * MERGE); + try t.dense(k, xn, t.down_w, t.down_b, att, 0, acc, M, OUT, 4 * H); + if (t.dbg_down) |dst| { + try k.stream.synchronize(); + try d.check(d.api.cuMemcpyDtoH_v2(@ptrCast(dst.ptr), att, dst.len * 2), "vision debug"); + } + // the merger: projection, LayerNorm + GELU, the limited SwiGLU to 10240 and back + var r: i64 = 0; + while (r < M) : (r += ROWS) { + const rows = @min(ROWS, M - r); + try t.linear(k, att + @as(u64, @intCast(r * OUT * 2)), t.m_proj, acc, rows, OUT, OUT, true); + var a: cuda.Args = .{}; + a.add(acc); a.add(t.m_ln_w); a.add(t.m_ln_b); a.add(xn + @as(u64, @intCast(r * OUT * 2))); a.add(@as(i32, @intCast(OUT))); a.add(EPS); + try k.go(try k.ext(.vision, "glm_vis_ln_gelu"), .{ .x = @intCast(rows) }, 256, 0, &a); + } + r = 0; + const step = @divFloor(ROWS * MLP, PROJ); + while (r < M) : (r += step) { + const rows = @min(step, M - r); + try t.swiglu(k, xn + @as(u64, @intCast(r * OUT * 2)), t.m_gate, 0, t.m_up, 0, big, acc, acc2, rows, PROJ, OUT); + try t.linear(k, big, t.m_down, out + @as(u64, @intCast(r * OUT * 2)), rows, OUT, PROJ, false); + } + try k.stream.synchronize(); // the scratch is freed on return + } +}; diff --git a/zig/src/families/glm_cuda/weights.zig b/zig/src/families/glm_cuda/weights.zig new file mode 100644 index 0000000000..371493428f --- /dev/null +++ b/zig/src/families/glm_cuda/weights.zig @@ -0,0 +1,1051 @@ +//! GLM-5.3-Flash's prepared weights for one rank of two, read from the converted checkpoint (glm-affine-experts-v1: +//! MLX affine 4-bit routed experts in groups of 64 with fp16 scales and biases, everything else BF16 / FP32), as the +//! Python engine prepares them (families/glm5_next/cuda/weights.py, its affine-experts path, split.py's rank rules): +//! every tensor under the name, shape and bytes of the Python engine's prepared weights (its weight tree, the names +//! tools/capture_ref.py hashes), uploaded to this rank's GPU. +//! +//! The derivations are exact: rank rows and columns cut on the host, BF16 rows stacked, the kv_b blocks split per head +//! (the latent path's absorb weights), fp16 expert scales rounded to bf16 (round to nearest even, as torch), and the +//! experts' blocks in experts_pack.cu's index map. The draft head is the verification head itself (no 4-bit copy). + +const std = @import("std"); +const cuda = @import("cuda"); +pub const api = @import("api.zig"); +pub const kern = @import("kern.zig"); +/// The forward and its engine API (tf-cuda-test glm-ref reaches it through this module). +pub const forward = @import("forward.zig"); +pub const decode = @import("decode.zig"); +pub const vision = @import("vision.zig"); + +const BASE = "model.language_model."; +const WORLD = 2; + +// --------------------------------------------------------------------------------------------------------------- +// The checkpoint: its shards' headers, files opened once, tensors read whole or as row ranges. + +const Entry = struct { dtype: []const u8, shape: []const i64, begin: u64, end: u64, file: usize }; + +const Ckpt = struct { + gpa: std.mem.Allocator, + io: std.Io, + arena: std.heap.ArenaAllocator, + files: std.ArrayList(std.Io.File) = .empty, + paths: std.ArrayList([]const u8) = .empty, + entries: std.StringHashMap(Entry), + + /// ``dir``: the served folder. A shard (or the index) that is a symlink to a container path under /cache/tf/ + /// resolves against the folder's parent directory, so the same folder serves inside the container and on the host. + fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Ckpt { + var ck: Ckpt = .{ .gpa = gpa, .io = io, .arena = .init(gpa), .entries = .init(gpa) }; + errdefer ck.deinit(); + const a = ck.arena.allocator(); + const index_path = try resolve(a, io, dir, "model.safetensors.index.json"); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, a, .limited(1 << 26)); + const idx = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); + var shard_of = std.StringHashMap(usize).init(a); + var it = idx.object.get("weight_map").?.object.iterator(); + while (it.next()) |e| { + const shard = e.value_ptr.string; + if (shard_of.contains(shard)) continue; + const path = try resolve(a, io, dir, shard); + try shard_of.put(shard, ck.files.items.len); + try ck.files.append(gpa, try std.Io.Dir.cwd().openFile(io, path, .{})); + try ck.paths.append(gpa, path); + try ck.readHeader(ck.files.items.len - 1); + } + return ck; + } + + fn readHeader(ck: *Ckpt, f: usize) !void { + const a = ck.arena.allocator(); + var head: [8]u8 = undefined; + _ = try ck.files.items[f].readPositionalAll(ck.io, &head, 0); + const n = std.mem.readInt(u64, &head, .little); + const text = try a.alloc(u8, n); + _ = try ck.files.items[f].readPositionalAll(ck.io, text, 8); + const v = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); + var it = v.object.iterator(); + while (it.next()) |e| { + if (std.mem.eql(u8, e.key_ptr.*, "__metadata__")) continue; + const o = e.value_ptr.object; + const shp = o.get("shape").?.array.items; + const shape = try a.alloc(i64, shp.len); + for (shp, 0..) |x, i| shape[i] = x.integer; + const offs = o.get("data_offsets").?.array.items; + try ck.entries.put(e.key_ptr.*, .{ .dtype = o.get("dtype").?.string, .shape = shape, + .begin = 8 + n + @as(u64, @intCast(offs[0].integer)), .end = 8 + n + @as(u64, @intCast(offs[1].integer)), .file = f }); + } + } + + fn deinit(ck: *Ckpt) void { + // the shards are read once: their page cache would starve NCCL's first registration (and the n-gram tables) + for (ck.files.items) |f| { + _ = std.os.linux.fadvise(f.handle, 0, 0, std.os.linux.POSIX_FADV.DONTNEED); + f.close(ck.io); + } + ck.files.deinit(ck.gpa); + ck.paths.deinit(ck.gpa); + ck.entries.deinit(); + ck.arena.deinit(); + } + + fn has(ck: *const Ckpt, name: []const u8) bool { + return ck.entries.contains(name); + } + + fn get(ck: *const Ckpt, name: []const u8) !Entry { + return ck.entries.get(name) orelse { + std.log.err("glm weights: the checkpoint has no tensor {s}", .{name}); + return error.MissingTensor; + }; + } + + /// ``len`` bytes at ``off`` into the tensor's data. + fn readAt(ck: *const Ckpt, e: Entry, off: u64, out: []u8) !void { + if (off + out.len > e.end - e.begin) return error.ShortRead; + const got = try ck.files.items[e.file].readPositionalAll(ck.io, out, e.begin + off); + if (got != out.len) return error.ShortRead; + } + + fn readAll(ck: *const Ckpt, gpa: std.mem.Allocator, name: []const u8) ![]align(16) u8 { + const e = try ck.get(name); + const buf = try gpa.alignedAlloc(u8, .@"16", e.end - e.begin); + errdefer gpa.free(buf); + try ck.readAt(e, 0, buf); + return buf; + } +}; + +fn resolve(a: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) ![]const u8 { + const path = try std.fs.path.join(a, &.{ dir, name }); + if (std.Io.Dir.cwd().openFile(io, path, .{})) |f| { + f.close(io); + return path; + } else |_| {} + var buf: [4096]u8 = undefined; + const n = std.Io.Dir.cwd().readLink(io, path, &buf) catch return path; + const target = buf[0..n]; + const container = "/cache/tf/"; + if (std.mem.startsWith(u8, target, container)) + return std.fs.path.join(a, &.{ std.fs.path.dirname(dir) orelse ".", target[container.len..] }); + if (std.fs.path.isAbsolute(target)) return a.dupe(u8, target); + return std.fs.path.join(a, &.{ dir, target }); +} + +fn bf16ToF32(b: u16) f32 { + return @bitCast(@as(u32, b) << 16); +} + +fn f32ToBf16(x: f32) u16 { + const u: u32 = @bitCast(x); + if (std.math.isNan(x)) return 0x7FC0; + return @intCast((u + 0x7FFF + ((u >> 16) & 1)) >> 16); +} + +fn f16ToBf16(h: u16) u16 { + return f32ToBf16(@floatCast(@as(f16, @bitCast(h)))); +} + +fn numelOf(shape: []const i64) usize { + var n: usize = 1; + for (shape) |d| n *= @intCast(d); + return n; +} + +fn dtypeOf(s: []const u8) !api.DType { + if (std.mem.eql(u8, s, "BF16")) return .bf16; + if (std.mem.eql(u8, s, "F32")) return .f32; + if (std.mem.eql(u8, s, "F16")) return .f16; + return error.UnsupportedDType; +} + +// --------------------------------------------------------------------------------------------------------------- +// split.py's rules: a checkpoint tensor's rank part (rows, the last axis' columns, or the whole tensor). + +const Split = enum { rep, row, col }; + +fn splitOf(name: []const u8) Split { + const ends = std.mem.endsWith; + const has = struct { + fn f(n: []const u8, frag: []const u8) bool { + return std.mem.indexOf(u8, n, frag) != null; + } + }.f; + if (has(name, ".mlp.experts.") or has(name, ".mlp.shared_experts.") or has(name, ".mlp.gate_proj.") or has(name, ".mlp.up_proj.") or has(name, ".mlp.down_proj.")) + return if (has(name, "down_proj.")) .col else .row; + if (has(name, ".self_attn.indexer.")) return .rep; + for ([_][]const u8{ ".self_attn.q_proj.", ".self_attn.k_proj.", ".self_attn.v_proj.", "_conv1d.", ".self_attn.f_b_proj.", ".self_attn.g_b_proj.", + ".self_attn.b_proj.", ".self_attn.q_b_proj.", ".self_attn.kv_b_proj." }) |f| if (has(name, f)) return .row; + if (ends(u8, name, ".self_attn.A_log") or ends(u8, name, ".self_attn.dt_bias")) return .row; + if (has(name, ".self_attn.o_proj.")) return .col; + return .rep; +} + +const Builder = struct { + gpa: std.mem.Allocator, + ck: *const Ckpt, + dry: bool, + rank: usize, + store: ?*api.Store = null, + d: ?*const cuda.Driver = null, + bytes: u64 = 0, + tensors: usize = 0, + /// Dense linears (all but the indexer's) and the head packed to these bits (0: bf16 as stored): see ``quantize``. + dense_bits: u8 = 0, + head_bits: u8 = 0, + /// TF_GLM_KDA_PROJ_BITS: the KDA layers' stacked in-projection (the largest dense linear, 1.3 GB a rank at 6 bits) + /// at its own width (0: dense_bits). + kda_proj_bits: u8 = 0, + /// A packed copy of the head the MTP drafts read (0: they read the head). + draft_head_bits: u8 = 0, + /// The routed experts' groups as scales alone (their biases -8 x the scales, checked): TF_GLM_EXPERTS_SYM=0 keeps + /// the affine blocks. + experts_sym: bool = true, + /// Development (TF_GLM_EXPERT_SIM_BITS=2 or 3): the experts re-quantized to fewer levels in their 4-bit words, to + /// measure a narrower expert format's quality before its kernel exists. + expert_sim_bits: u8 = 0, + /// TF_GLM_EXPERT_BITS=3: the experts re-quantized to 3-bit codes and packed as glm/experts_sym3's planes. + experts3: bool = false, + /// TF_GLM_EXPERT3_PROJ=down|up: 3-bit codes for that projection alone (the other stays 4-bit in the planes: a + /// sensitivity probe; the kernel reads 3-bit planes either way, so "4-bit" here means 8 levels of the 4-bit + /// value re-rounded). TF_GLM_EXPERT_SIM_BITS probes quality on the 4-bit kernels instead: with + /// TF_GLM_EXPERT_SIM_PROJ=down|up it re-quantizes that projection alone. + sim_proj: u8 = 0, // 0: all, 1: down only, 2: gate/up only + /// The checkpoint's routed experts are EXL3 (4-bit "mcg" trellises, every other weight BF16): read as stored. + exl3: bool = false, + exl3_shx: bool = true, // every layer's experts share one input rotation for gate and up (cleared by a layer that does not) + suh_first: [2 * 4096]u8 = undefined, // the first expert's gate suh of the layer being read + + fn put(b: *Builder, name: []const u8, dtype: api.DType, shape: []const i64, data: ?[]const u8) !void { + var t: api.Tensor = .{ .ptr = 0, .dtype = dtype, .nd = @intCast(shape.len) }; + for (shape, 0..) |x, i| t.shape[i] = x; + const n = t.bytes(); + b.bytes += n; + b.tensors += 1; + if (b.dry) return; + const store = b.store.?; + var buf = try cuda.DeviceBuffer.alloc(b.d.?, @max(n, 1)); + errdefer buf.free(); + if (n > 0) try buf.upload(0, data.?[0..n]); + t.ptr = buf.ptr; + try store.owned.append(store.gpa, buf); + try store.put(name, t); + } + + fn putf(b: *Builder, comptime fmt: []const u8, args: anytype, dtype: api.DType, shape: []const i64, data: ?[]const u8) !void { + var nb: [256]u8 = undefined; + try b.put(try std.fmt.bufPrint(&nb, fmt, args), dtype, shape, data); + } + + /// A checkpoint tensor's rank part as stored: its dtype, rank shape and bytes (null in a dry run). + const Part = struct { dtype: api.DType, shape: [4]i64, nd: usize, data: ?[]align(16) u8 }; + + fn part(b: *Builder, full: []const u8, how: Split) !Part { + const e = try b.ck.get(full); + var p: Part = .{ .dtype = try dtypeOf(e.dtype), .shape = .{ 1, 1, 1, 1 }, .nd = e.shape.len, .data = null }; + for (e.shape, 0..) |x, i| p.shape[i] = x; + const elem = api.dtypeBytes(p.dtype); + const rows: usize = @intCast(p.shape[0]); + const inner: usize = numelOf(e.shape) / rows; + switch (how) { + .rep => { + if (!b.dry) p.data = try b.ck.readAll(b.gpa, full); + }, + .row => { + const per = rows / WORLD; + p.shape[0] = @intCast(per); + if (!b.dry) { + p.data = try b.gpa.alignedAlloc(u8, .@"16", per * inner * elem); + try b.ck.readAt(e, b.rank * per * inner * elem, p.data.?); + } + }, + .col => { + if (e.shape.len != 2) return error.ColumnSplitNeeds2D; + const k: usize = @intCast(e.shape[1]); + const ck = k / WORLD; + p.shape[1] = @intCast(ck); + if (!b.dry) { + const all = try b.ck.readAll(b.gpa, full); + defer b.gpa.free(all); + p.data = try b.gpa.alignedAlloc(u8, .@"16", rows * ck * elem); + for (0..rows) |r| @memcpy(p.data.?[r * ck * elem ..][0 .. ck * elem], all[(r * k + b.rank * ck) * elem ..][0 .. ck * elem]); + } + }, + } + return p; + } + + fn free(b: *Builder, p: Part) void { + if (p.data) |x| b.gpa.free(x); + } + + /// One checkpoint tensor (its rank part) under ``out``, cast to fp32 when ``f32`` (bf16 -> fp32 is exact). + fn tensor(b: *Builder, out: []const u8, full: []const u8, f32_out: bool) !void { + const p = try b.part(full, splitOf(full)); + defer b.free(p); + if (f32_out and p.dtype == .bf16) { + const n = numelOf(p.shape[0..p.nd]); + if (b.dry) return b.put(out, .f32, p.shape[0..p.nd], null); + const f = try b.gpa.alloc(f32, n); + defer b.gpa.free(f); + const h = std.mem.bytesAsSlice(u16, p.data.?); + for (f, 0..) |*o, i| o.* = bf16ToF32(h[i]); + return b.put(out, .f32, p.shape[0..p.nd], std.mem.sliceAsBytes(f)); + } + try b.put(out, p.dtype, p.shape[0..p.nd], p.data); + } + + /// BF16 linears of the same K, their rank rows stacked in order -> ``out``.weight [sum n, k] (stack_b16). + fn stack(b: *Builder, out: []const u8, fulls: []const []const u8) !void { + var parts: [8]Part = undefined; + var n: i64 = 0; + var kk: i64 = 0; + for (fulls, 0..) |f, i| { + parts[i] = try b.part(f, splitOf(f)); + if (parts[i].dtype != .bf16 or parts[i].nd != 2) return error.NotBf16Matrix; + n += parts[i].shape[0]; + kk = parts[i].shape[1]; + } + defer for (parts[0..fulls.len]) |p| b.free(p); + var nb: [256]u8 = undefined; + var bits: u8 = if (std.mem.indexOf(u8, out, "index.") != null or @mod(kk, 64) != 0) 0 else b.dense_bits; + if (b.kda_proj_bits != 0 and fulls.len == 6) bits = b.kda_proj_bits; // the KDA in-projection: 6 stacked linears + const name = try std.fmt.bufPrint(&nb, "{s}.weight", .{out}); + if (b.dry and bits == 0) return b.put(name, .bf16, &.{ n, kk }, null); + if (b.dry) return b.putPacked(out, bits, n, kk, null); + const all = try b.gpa.alignedAlloc(u8, .@"16", @intCast(n * kk * 2)); + defer b.gpa.free(all); + var at: usize = 0; + for (parts[0..fulls.len]) |p| { + @memcpy(all[at..][0..p.data.?.len], p.data.?); + at += p.data.?.len; + } + if (bits != 0) return b.putPacked(out, bits, n, kk, all); + try b.put(name, .bf16, &.{ n, kk }, all); + } + + /// bf16 rows [n, kk] -> ``out``.qw (u8 [n, kk * bits / 8]) and ``out``.qs (fp16 [n, kk / 64]). + fn putPacked(b: *Builder, out: []const u8, bits: u8, n: i64, kk: i64, rows: ?[]align(16) const u8) !void { + const nu: usize = @intCast(n); + const ku: usize = @intCast(kk); + var nb: [256]u8 = undefined; + var nb2: [256]u8 = undefined; + const qn = try std.fmt.bufPrint(&nb, "{s}.qw", .{out}); + const sn = try std.fmt.bufPrint(&nb2, "{s}.qs", .{out}); + const wshape = [_]i64{ n, @divExact(kk * bits, 8) }; + const sshape = [_]i64{ n, @divExact(kk, 64) }; + if (rows == null) { + try b.put(qn, .u8, &wshape, null); + return b.put(sn, .f16, &sshape, null); + } + const words = try b.gpa.alloc(u32, nu * ku * bits / 32); + defer b.gpa.free(words); + const scales = try b.gpa.alloc(u16, nu * ku / 64); + defer b.gpa.free(scales); + quantize(bits, std.mem.bytesAsSlice(u16, rows.?), nu, ku, words, scales); + try b.put(qn, .u8, &wshape, std.mem.sliceAsBytes(words)); + try b.put(sn, .f16, &sshape, std.mem.sliceAsBytes(scales)); + } + + fn linear(b: *Builder, out: []const u8, full: []const u8) !void { + try b.stack(out, &.{full}); + } +}; + +fn fullName(buf: []u8, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(buf, BASE ++ fmt, args); +} + +// --------------------------------------------------------------------------------------------------------------- +// The routed experts: per-expert MLX triples, the rank's rows (gate, up) or columns (down), packed as experts_pack.cu. + +const Triple = struct { + e: usize, + n: usize, + k8: usize, + kg: usize, + words: []u32, + scales: []u16, + biases: []u16, + + fn free(t: Triple, gpa: std.mem.Allocator) void { + gpa.free(t.words); + gpa.free(t.scales); + gpa.free(t.biases); + } +}; + +/// Every expert's ``proj`` of layer ``layer``: rows r0..r1 (gate, up) or the rank's columns (down), scales and biases +/// rounded from fp16 to bf16. +fn readExperts(b: *Builder, li: usize, proj: []const u8, count: usize) !Triple { + var nb: [256]u8 = undefined; + const w0 = try b.ck.get(try fullName(&nb, "layers.{d}.mlp.experts.0.{s}.weight", .{ li, proj })); + const n_full: usize = @intCast(w0.shape[0]); + const k8_full: usize = @intCast(w0.shape[1]); + var nb2: [256]u8 = undefined; + const s0 = try b.ck.get(try fullName(&nb2, "layers.{d}.mlp.experts.0.{s}.scales", .{ li, proj })); + const kg_full: usize = @intCast(s0.shape[1]); + const down = std.mem.eql(u8, proj, "down_proj"); + const nr = if (down) n_full else n_full / WORLD; + const r0 = if (down) 0 else b.rank * nr; + const ok8 = if (down) k8_full / WORLD else k8_full; + const okg = if (down) kg_full / WORLD else kg_full; + const c8 = if (down) b.rank * ok8 else 0; + const cg = if (down) b.rank * okg else 0; + var t: Triple = .{ .e = count, .n = nr, .k8 = ok8, .kg = okg, .words = try b.gpa.alloc(u32, count * nr * ok8), + .scales = try b.gpa.alloc(u16, count * nr * okg), .biases = try b.gpa.alloc(u16, count * nr * okg) }; + errdefer t.free(b.gpa); + const rowbuf = try b.gpa.alloc(u8, nr * k8_full * 4); + defer b.gpa.free(rowbuf); + const sbuf = try b.gpa.alloc(u8, nr * kg_full * 2); + defer b.gpa.free(sbuf); + for (0..count) |x| { + var n1: [256]u8 = undefined; + const we = try b.ck.get(try fullName(&n1, "layers.{d}.mlp.experts.{d}.{s}.weight", .{ li, x, proj })); + try b.ck.readAt(we, r0 * k8_full * 4, rowbuf); + const w32 = std.mem.bytesAsSlice(u32, rowbuf); + for (0..nr) |r| @memcpy(t.words[(x * nr + r) * ok8 ..][0..ok8], w32[r * k8_full + c8 ..][0..ok8]); + for ([_][]const u8{ "scales", "biases" }, 0..) |what, which| { + var n2: [256]u8 = undefined; + const ent = try b.ck.get(try fullName(&n2, "layers.{d}.mlp.experts.{d}.{s}.{s}", .{ li, x, proj, what })); + try b.ck.readAt(ent, r0 * kg_full * 2, sbuf); + const s16 = std.mem.bytesAsSlice(u16, sbuf); + const dst = if (which == 0) t.scales else t.biases; + for (0..nr) |r| for (0..okg) |g| { + dst[(x * nr + r) * okg + g] = f16ToBf16(s16[r * kg_full + cg + g]); + }; + } + } + return t; +} + +/// experts_pack.cu's index map: each (group, column block, expert) block holds 32 * 4 * H weight words (nibbles +/// reordered i0 i2 i4 i6 i1 i3 i5 i7) then 32 scale/bias words; written as slot ``m`` of ``mcount`` per block. +/// Every group's bias is -8 x its scale (a symmetric int4 checkpoint as affine groups): the symmetric blocks hold +/// the scales alone. +fn checkSymmetric(t: Triple) !void { + for (t.scales, t.biases) |sc, bi| { + if (bf16Value(bi) != -8 * bf16Value(sc)) return error.ExpertsNotSymmetric; + } +} + +fn packExperts(t: Triple, gs: usize, out: []u32, mcount: usize, m: usize, sym: bool, three: bool) void { + const h = gs / 32; + const Work = struct { + fn run(tt: Triple, hh: usize, o: []u32, mc: usize, mm: usize, e0: usize, e1: usize, sy: bool, th: bool) void { + const wp = 32 * 4 * hh; + const wsec = if (th) wp * 3 / 4 else wp; + const blk = wsec + @as(usize, if (sy) 16 else 32); + var tmp: [512]u32 = undefined; // a 3-bit block's 4-bit words before their planes + const nbb = tt.n / 32; + for (e0..e1) |e| for (0..nbb) |bb| for (0..tt.kg) |g| { + const dst = o[(((e * nbb + bb) * tt.kg + g) * mc + mm) * blk ..][0..blk]; + const wdst = if (th) tmp[0..wp] else dst[0..wp]; + for (0..wp) |i| { + const row = i / (4 * hh); + const kk = i % (4 * hh); + const tq = row / 8; + const r = row % 8; + const q = kk / hh; + const j = kk % hh; + const w = tt.words[(e * tt.n + bb * 32 + row) * tt.k8 + g * 4 * hh + kk]; + const f = ((tq * hh + j) * 8 + r) * 4 + q; + wdst[(f / 128) * 128 + (f % 32) * 4 + (f / 32) % 4] = shuffle(w); + } + if (th) for (0..hh) |c| for (0..32) |lane| { // the lane's uint4 c as three planes (experts_sym.cuh) + var lo: [4]u32 = undefined; + var hi: [4]u32 = undefined; + for (0..4) |k| { + const word = tmp[(c * 32 + lane) * 4 + k]; + lo[k] = 0; + hi[k] = 0; + for (0..8) |i| { + const code = ((word >> @intCast(4 * i)) & 15) -% 4; // 3-bit codes are stored + 4 + lo[k] |= (code & 3) << @intCast(2 * i); + hi[k] |= ((code >> 2) & 1) << @intCast(i); + } + } + dst[c * 96 + lane] = lo[0] | (lo[1] << 16); + dst[c * 96 + 32 + lane] = lo[2] | (lo[3] << 16); + dst[c * 96 + 64 + lane] = hi[0] | (hi[1] << 8) | (hi[2] << 16) | (hi[3] << 24); + }; + if (sy) { // quad p: rows j * 8 + 2p and + 1's scales, j = 0..3 (experts_sym.cuh) + for (0..16) |s| { + const p = s / 4; + const j = s % 4; + const at = (e * tt.n + bb * 32 + j * 8 + p * 2) * tt.kg + g; + dst[wsec + s] = @as(u32, tt.scales[at]) | (@as(u32, tt.scales[at + tt.kg]) << 16); + } + } else for (0..32) |s| { + const p = s / 8; + const tt4 = s % 4; + const src = if ((s / 4) % 2 == 1) tt.biases else tt.scales; + const at = (e * tt.n + bb * 32 + tt4 * 8 + p * 2) * tt.kg + g; + dst[wp + s] = @as(u32, src[at]) | (@as(u32, src[at + tt.kg]) << 16); + } + }; + } + }; + const threads = 16; + var pool: [threads]?std.Thread = @splat(null); + const per = (t.e + threads - 1) / threads; + for (0..threads) |i| { + const e0 = i * per; + const e1 = @min(t.e, e0 + per); + if (e0 >= e1) break; + pool[i] = std.Thread.spawn(.{}, Work.run, .{ t, h, out, mcount, m, e0, e1, sym, three }) catch blk: { + Work.run(t, h, out, mcount, m, e0, e1, sym, three); + break :blk null; + }; + } + for (pool) |th| if (th) |x| x.join(); +} + +fn shuffle(w: u32) u32 { + var even = w & 0x0F0F0F0F; + var odd = (w >> 4) & 0x0F0F0F0F; + even = (even | (even >> 4)) & 0x00FF00FF; + odd = (odd | (odd >> 4)) & 0x00FF00FF; + even = (even | (even >> 8)) & 0x0000FFFF; + odd = (odd | (odd >> 8)) & 0x0000FFFF; + return even | (odd << 16); +} + +/// grouped.make([gate, up], down, 64): up [E, N/32, K/64, 2, block], down [E, D/32, NI/64, 1, block] int32. +fn experts(b: *Builder, out: []const u8, li: usize) !void { + const E: usize = @intCast(kern.EXPERTS); + const gs: usize = 64; + if (b.experts3 and !b.experts_sym) return error.Experts3NeedsSymmetric; + const wsec: usize = 32 * 4 * (gs / 32) * @as(usize, if (b.experts3) 3 else 4) / 4; + const blk: usize = wsec + @as(usize, if (b.experts_sym) 16 else 32); + const width: usize = @intCast(@divExact(2048, WORLD)); // moe_intermediate_size / ranks + const dims: usize = @intCast(kern.D); + const up_shape = [_]i64{ @intCast(E), @intCast(width / 32), @intCast(dims / gs), 2, @intCast(blk) }; + const down_shape = [_]i64{ @intCast(E), @intCast(dims / 32), @intCast(width / gs), 1, @intCast(blk) }; + var nb: [256]u8 = undefined; + var nb2: [256]u8 = undefined; + const up_name = try std.fmt.bufPrint(&nb, "{s}.up", .{out}); + const down_name = try std.fmt.bufPrint(&nb2, "{s}.down", .{out}); + if (b.dry) { + try b.put(up_name, .i32, &up_shape, null); + try b.put(down_name, .i32, &down_shape, null); + return; + } + { + const up = try b.gpa.alloc(u32, numelOf(&up_shape)); + defer b.gpa.free(up); + for ([_][]const u8{ "gate_proj", "up_proj" }, 0..) |proj, m| { + const t = try readExperts(b, li, proj, E); + defer t.free(b.gpa); + if (t.n != width or t.k8 * 8 != dims) return error.ExpertShape; + if (b.experts_sym) try checkSymmetric(t); + if (b.sim_proj != 1) { + if (b.experts3) simulateBits(t, 3) else if (b.expert_sim_bits != 0) simulateBits(t, b.expert_sim_bits); + } + packExperts(t, gs, up, 2, m, b.experts_sym, b.experts3); + } + try b.put(up_name, .i32, &up_shape, std.mem.sliceAsBytes(up)); + } + const dn = try b.gpa.alloc(u32, numelOf(&down_shape)); + defer b.gpa.free(dn); + const t = try readExperts(b, li, "down_proj", E); + defer t.free(b.gpa); + if (t.n != dims or t.k8 * 8 != width) return error.ExpertShape; + if (b.experts_sym) try checkSymmetric(t); + if (b.sim_proj != 2) { + if (b.experts3) simulateBits(t, 3) else if (b.expert_sim_bits != 0) simulateBits(t, b.expert_sim_bits); + } + packExperts(t, gs, dn, 1, 0, b.experts_sym, b.experts3); + try b.put(down_name, .i32, &down_shape, std.mem.sliceAsBytes(dn)); +} + +/// The EXL3 routed experts of layer ``li``, this rank's whole tiles (glm/exl3.cu's layout): gate and up by output tile +/// columns, down by input tile rows, each matrix's experts stacked as trellis words [E, K/16, N/16, 32] (the stored +/// bytes) and fp16 scales [E, K] (suh) and [E, N] (svh), the rank's parts. +fn expertsExl3(b: *Builder, out: []const u8, li: usize) !void { + const E: usize = @intCast(kern.EXPERTS); + const D: usize = @intCast(kern.D); + const width: usize = @intCast(@divExact(2048, WORLD)); + const tile: usize = 128; // a 16x16 tile's bytes at 4 bits + var nb: [256]u8 = undefined; + const Mat = struct { proj: []const u8, kt: usize, nt: usize, by_cols: bool }; + const mats = [_]Mat{ .{ .proj = "gate_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, .{ .proj = "up_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, + .{ .proj = "down_proj", .kt = width / 16, .nt = D / 16, .by_cols = false } }; + const short = [_][]const u8{ "g", "u", "d" }; + for (mats, short) |m, sh| { + const tshape = [_]i64{ @intCast(E), @intCast(m.kt), @intCast(m.nt), 32 }; + const k_all: usize = if (m.by_cols) D else 2048; + const n_all: usize = if (m.by_cols) 2048 else D; + const sus = [_]i64{ @intCast(E), @intCast(m.kt * 16) }; + const svs = [_]i64{ @intCast(E), @intCast(m.nt * 16) }; + const tn = try std.fmt.bufPrint(&nb, "{s}.{s}t", .{ out, sh }); + if (b.dry) { + try b.put(tn, .i32, &tshape, null); + try b.putf("{s}.suh_{s}", .{ out, sh }, .f16, &sus, null); + try b.putf("{s}.svh_{s}", .{ out, sh }, .f16, &svs, null); + continue; + } + const per = m.kt * m.nt * tile; + const words = try b.gpa.alignedAlloc(u8, .@"16", E * per); + defer b.gpa.free(words); + const suh = try b.gpa.alloc(u16, E * m.kt * 16); + defer b.gpa.free(suh); + const svh = try b.gpa.alloc(u16, E * m.nt * 16); + defer b.gpa.free(svh); + var f: [256]u8 = undefined; + for (0..E) |x| { + const tr = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.trellis", .{ li, x, m.proj })); + if (!std.mem.eql(u8, tr.dtype, "I16") or tr.shape.len != 3 or tr.shape[2] != 64 or tr.shape[0] != k_all / 16 or tr.shape[1] != n_all / 16) return error.NotExl3Mcg4; + if (!b.ck.has(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.mcg", .{ li, x, m.proj }))) return error.NotExl3Mcg4; + const dst = words[x * per ..][0..per]; + if (m.by_cols) { // every tile row, this rank's tile columns + const all = try b.ck.readAll(b.gpa, try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.trellis", .{ li, x, m.proj })); + defer b.gpa.free(all); + const row = (n_all / 16) * tile; + for (0..m.kt) |r| @memcpy(dst[r * m.nt * tile ..][0 .. m.nt * tile], all[r * row + b.rank * m.nt * tile ..][0 .. m.nt * tile]); + } else try b.ck.readAt(tr, b.rank * per, dst); // this rank's tile rows: one contiguous part + const se = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.suh", .{ li, x, m.proj })); + const ve = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.svh", .{ li, x, m.proj })); + if (!std.mem.eql(u8, se.dtype, "F16") or !std.mem.eql(u8, ve.dtype, "F16")) return error.NotExl3Mcg4; + const su = std.mem.sliceAsBytes(suh[x * m.kt * 16 ..][0 .. m.kt * 16]); + const sv = std.mem.sliceAsBytes(svh[x * m.nt * 16 ..][0 .. m.nt * 16]); + try b.ck.readAt(se, if (m.by_cols) 0 else b.rank * su.len, su); + try b.ck.readAt(ve, if (m.by_cols) b.rank * sv.len else 0, sv); + if (m.by_cols) { // gate and up: one input rotation shared by every expert of the layer? (exl3_shx) + if (std.mem.eql(u8, sh, "g") and x == 0) @memcpy(b.suh_first[0..su.len], su) + else if (!std.mem.eql(u8, su, b.suh_first[0..su.len])) b.exl3_shx = false; + } + } + try b.put(tn, .i32, &tshape, words); + try b.putf("{s}.suh_{s}", .{ out, sh }, .f16, &sus, std.mem.sliceAsBytes(suh)); + try b.putf("{s}.svh_{s}", .{ out, sh }, .f16, &svs, std.mem.sliceAsBytes(svh)); + } +} + +// --------------------------------------------------------------------------------------------------------------- +// The layers, as weights.py's load() builds them (names: its weight tree). + +fn hc(b: *Builder, out: []const u8, li: usize, site: []const u8) !void { + var o: [128]u8 = undefined; + var f: [256]u8 = undefined; + try b.tensor(try std.fmt.bufPrint(&o, "{s}.fn", .{out}), try fullName(&f, "layers.{d}.hc_{s}_fn", .{ li, site }), false); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.base", .{out}), try fullName(&f, "layers.{d}.hc_{s}_base", .{ li, site }), true); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.scale", .{out}), try fullName(&f, "layers.{d}.hc_{s}_scale", .{ li, site }), true); +} + +fn kda(b: *Builder, out: []const u8, i: usize) !void { + var o: [160]u8 = undefined; + var n: [6][256]u8 = undefined; + var names: [6][]const u8 = undefined; + for ([_][]const u8{ "q_proj", "k_proj", "v_proj", "f_a_proj", "g_a_proj", "b_proj" }, 0..) |p, j| + names[j] = try fullName(&n[j], "layers.{d}.self_attn.{s}.weight", .{ i, p }); + try b.stack(try std.fmt.bufPrint(&o, "{s}.proj", .{out}), &names); + var f: [256]u8 = undefined; + try b.linear(try std.fmt.bufPrint(&o, "{s}.fb", .{out}), try fullName(&f, "layers.{d}.self_attn.f_b_proj.weight", .{i})); + try b.linear(try std.fmt.bufPrint(&o, "{s}.gb", .{out}), try fullName(&f, "layers.{d}.self_attn.g_b_proj.weight", .{i})); + // conv: q | k | v taps, each [rank channels, 1, 4], as [3 HL 128, 4] + { + var parts: [3]Builder.Part = undefined; + for ("qkv", 0..) |x, j| { + var cn: [256]u8 = undefined; + parts[j] = try b.part(try fullName(&cn, "layers.{d}.self_attn.{c}_conv1d.weight", .{ i, x }), .row); + } + defer for (parts) |p| b.free(p); + const ch = parts[0].shape[0] * 3; + const shape = [_]i64{ ch, 4 }; + if (b.dry) { + try b.put(try std.fmt.bufPrint(&o, "{s}.conv", .{out}), .bf16, &shape, null); + } else { + const all = try b.gpa.alloc(u8, @intCast(ch * 4 * 2)); + defer b.gpa.free(all); + var at: usize = 0; + for (parts) |p| { + @memcpy(all[at..][0..p.data.?.len], p.data.?); + at += p.data.?.len; + } + try b.put(try std.fmt.bufPrint(&o, "{s}.conv", .{out}), .bf16, &shape, all); + } + } + try b.tensor(try std.fmt.bufPrint(&o, "{s}.a_log", .{out}), try fullName(&f, "layers.{d}.self_attn.A_log", .{i}), true); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.dt_bias", .{out}), try fullName(&f, "layers.{d}.self_attn.dt_bias", .{i}), true); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.norm", .{out}), try fullName(&f, "layers.{d}.self_attn.o_norm.weight", .{i}), false); + try b.linear(try std.fmt.bufPrint(&o, "{s}.o", .{out}), try fullName(&f, "layers.{d}.self_attn.o_proj.weight", .{i})); +} + +fn dsa(b: *Builder, out: []const u8, i: usize) !void { + var o: [160]u8 = undefined; + var f: [256]u8 = undefined; + var f2: [256]u8 = undefined; + try b.stack(try std.fmt.bufPrint(&o, "{s}.proj", .{out}), &.{ try fullName(&f, "layers.{d}.self_attn.q_a_proj.weight", .{i}), + try fullName(&f2, "layers.{d}.self_attn.kv_a_proj_with_mqa.weight", .{i}) }); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.q_norm", .{out}), try fullName(&f, "layers.{d}.self_attn.q_a_layernorm.weight", .{i}), false); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.kv_norm", .{out}), try fullName(&f, "layers.{d}.self_attn.kv_a_layernorm.weight", .{i}), false); + try b.linear(try std.fmt.bufPrint(&o, "{s}.q_b", .{out}), try fullName(&f, "layers.{d}.self_attn.q_b_proj.weight", .{i})); + // kv_b per head: 256 key rows then 256 value rows -> absorb.wk, absorb.wv [HL, 256, 512] (the latent path's only use) + { + const p = try b.part(try fullName(&f, "layers.{d}.self_attn.kv_b_proj.weight", .{i}), .row); + defer b.free(p); + const HL: usize = @intCast(kern.HL); + const lw: usize = @intCast(p.shape[1]); + const shape = [_]i64{ kern.HL, kern.QK, @intCast(lw) }; + if (b.dry) { + try b.put(try std.fmt.bufPrint(&o, "{s}.absorb.wk", .{out}), .bf16, &shape, null); + try b.put(try std.fmt.bufPrint(&o, "{s}.absorb.wv", .{out}), .bf16, &shape, null); + } else { + const half: usize = @intCast(kern.QK); + const wk = try b.gpa.alloc(u8, HL * half * lw * 2); + defer b.gpa.free(wk); + const wv = try b.gpa.alloc(u8, HL * half * lw * 2); + defer b.gpa.free(wv); + const row = lw * 2; + for (0..HL) |h| { + @memcpy(wk[h * half * row ..][0 .. half * row], p.data.?[(h * 2 * half) * row ..][0 .. half * row]); + @memcpy(wv[h * half * row ..][0 .. half * row], p.data.?[(h * 2 * half + half) * row ..][0 .. half * row]); + } + try b.put(try std.fmt.bufPrint(&o, "{s}.absorb.wk", .{out}), .bf16, &shape, wk); + try b.put(try std.fmt.bufPrint(&o, "{s}.absorb.wv", .{out}), .bf16, &shape, wv); + } + } + try b.linear(try std.fmt.bufPrint(&o, "{s}.o", .{out}), try fullName(&f, "layers.{d}.self_attn.o_proj.weight", .{i})); + try b.stack(try std.fmt.bufPrint(&o, "{s}.index.kw", .{out}), &.{ try fullName(&f, "layers.{d}.self_attn.indexer.wk.weight", .{i}), + try fullName(&f2, "layers.{d}.self_attn.indexer.weights_proj.weight", .{i}) }); + try b.linear(try std.fmt.bufPrint(&o, "{s}.index.qb", .{out}), try fullName(&f, "layers.{d}.self_attn.indexer.wq_b.weight", .{i})); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.index.ln_w", .{out}), try fullName(&f, "layers.{d}.self_attn.indexer.k_norm.weight", .{i}), false); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.index.ln_b", .{out}), try fullName(&f, "layers.{d}.self_attn.indexer.k_norm.bias", .{i}), false); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.index.gate", .{out}), try fullName(&f, "layers.{d}.self_attn.indexer.index_kpool_compress_gate", .{i}), false); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.index.ape", .{out}), try fullName(&f, "layers.{d}.self_attn.indexer.index_kpool_compress_ape", .{i}), false); +} + +fn mlp(b: *Builder, out: []const u8, prefix: []const u8) !void { + var o: [160]u8 = undefined; + var f: [256]u8 = undefined; + var f2: [256]u8 = undefined; + try b.stack(try std.fmt.bufPrint(&o, "{s}.gu", .{out}), &.{ try std.fmt.bufPrint(&f, "{s}gate_proj.weight", .{prefix}), + try std.fmt.bufPrint(&f2, "{s}up_proj.weight", .{prefix}) }); + try b.linear(try std.fmt.bufPrint(&o, "{s}.down", .{out}), try std.fmt.bufPrint(&f, "{s}down_proj.weight", .{prefix})); +} + +fn moe(b: *Builder, out: []const u8, i: usize) !void { + var o: [160]u8 = undefined; + var f: [256]u8 = undefined; + try b.tensor(try std.fmt.bufPrint(&o, "{s}.router", .{out}), try fullName(&f, "layers.{d}.mlp.gate.weight", .{i}), false); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.bias", .{out}), try fullName(&f, "layers.{d}.mlp.gate.e_score_correction_bias", .{i}), true); + if (b.exl3) try expertsExl3(b, try std.fmt.bufPrint(&o, "{s}.experts", .{out}), i) else try experts(b, try std.fmt.bufPrint(&o, "{s}.experts", .{out}), i); + try mlp(b, try std.fmt.bufPrint(&o, "{s}.shared", .{out}), try fullName(&f, "layers.{d}.mlp.shared_experts.", .{i})); +} + +/// One decoder layer (``plain``: the MTP head's, without hyper-connections) under ``out``. +fn layer(b: *Builder, out: []const u8, i: usize, kda_layer: bool, dense: bool, plain: bool) !void { + var o: [160]u8 = undefined; + var f: [256]u8 = undefined; + if (!plain) { + try hc(b, try std.fmt.bufPrint(&o, "{s}.attn_hc", .{out}), i, "attn"); + try hc(b, try std.fmt.bufPrint(&o, "{s}.ffn_hc", .{out}), i, "ffn"); + } + try b.tensor(try std.fmt.bufPrint(&o, "{s}.in_norm", .{out}), try fullName(&f, "layers.{d}.input_layernorm.weight", .{i}), false); + try b.tensor(try std.fmt.bufPrint(&o, "{s}.post_norm", .{out}), try fullName(&f, "layers.{d}.post_attention_layernorm.weight", .{i}), false); + if (kda_layer) try kda(b, try std.fmt.bufPrint(&o, "{s}.kda", .{out}), i) else try dsa(b, try std.fmt.bufPrint(&o, "{s}.dsa", .{out}), i); + if (dense) try mlp(b, try std.fmt.bufPrint(&o, "{s}.mlp", .{out}), try fullName(&f, "layers.{d}.mlp.", .{i})) else try moe(b, try std.fmt.bufPrint(&o, "{s}.moe", .{out}), i); +} + +/// The text config fields the preparation reads, checked against the dimensions the kernels take. +pub const Cfg = struct { + exl3: bool = false, // EXL3 routed experts (else the converted affine ones) + layers: usize, + kinds: [64]bool = @splat(false), // true: a KDA layer + dense: usize, // first_k_dense_replace + mtp: bool, + vocab: usize, + + pub fn read(a: std.mem.Allocator, io: std.Io, dir: []const u8) !Cfg { + const path = try std.fs.path.join(a, &.{ dir, "config.json" }); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 22)); + const v = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); + const q = v.object.get("quantization_config") orelse return error.NotAffineExperts; + var exl3 = false; + if (q.object.get("quant_method")) |m| if (std.mem.eql(u8, m.string, "exl3")) { + const cb = q.object.get("codebook") orelse return error.NotExl3Mcg4; + const bits = q.object.get("bits") orelse return error.NotExl3Mcg4; + if (!std.mem.eql(u8, cb.string, "mcg") or bits.integer != 4) return error.NotExl3Mcg4; + exl3 = true; + }; + if (!exl3) { + const fmt = q.object.get("format") orelse return error.NotAffineExperts; + if (!std.mem.eql(u8, fmt.string, "glm-affine-experts-v1")) return error.NotAffineExperts; + } + const t = v.object.get("text_config").?.object; + const n: usize = @intCast(t.get("num_hidden_layers").?.integer); + var c: Cfg = .{ .exl3 = exl3, .layers = n, .dense = @intCast(t.get("first_k_dense_replace").?.integer), .mtp = (if (t.get("num_nextn_predict_layers")) |x| x.integer else 0) > 0, + .vocab = @intCast(t.get("vocab_size").?.integer) }; + if (n > 64) return error.TooManyLayers; + for (t.get("layer_types").?.array.items, 0..) |x, i| c.kinds[i] = std.mem.eql(u8, x.string, "linear_attention"); + if (t.get("hidden_size").?.integer != kern.D or t.get("n_routed_experts").?.integer != kern.EXPERTS or t.get("num_experts_per_tok").?.integer != kern.TOPK or + t.get("num_attention_heads").?.integer != kern.HL * WORLD or t.get("kv_lora_rank").?.integer != kern.KV_LORA) return error.UnexpectedDimensions; + return c; + } +}; + +fn build(b: *Builder, c: *const Cfg) !void { + var o: [64]u8 = undefined; + for (0..c.layers) |i| try layer(b, try std.fmt.bufPrint(&o, "layers.{d}", .{i}), i, c.kinds[i], i < c.dense, false); + var f: [256]u8 = undefined; + try b.tensor("norm", try fullName(&f, "norm.weight", .{}), false); + try b.tensor("embed", BASE ++ "embed_tokens.weight", false); + // the head: this rank's vocabulary rows (packed to head_bits when set) + { + const e = try b.ck.get("lm_head.weight"); + const rows: usize = @intCast(e.shape[0]); + const kk: usize = @intCast(e.shape[1]); + const vl = rows / WORLD; + const shape = [_]i64{ @intCast(vl), @intCast(kk) }; + if (b.dry) { + if (b.head_bits != 0) try b.putPacked("head", b.head_bits, shape[0], shape[1], null) else try b.put("head.weight", .bf16, &shape, null); + if (b.draft_head_bits != 0) try b.putPacked("draft_head", b.draft_head_bits, shape[0], shape[1], null); + } else { + const buf = try b.gpa.alignedAlloc(u8, .@"16", vl * kk * 2); + defer b.gpa.free(buf); + try b.ck.readAt(e, b.rank * vl * kk * 2, buf); + if (b.head_bits != 0) try b.putPacked("head", b.head_bits, shape[0], shape[1], buf) else try b.put("head.weight", .bf16, &shape, buf); + if (b.draft_head_bits != 0) try b.putPacked("draft_head", b.draft_head_bits, shape[0], shape[1], buf); + } + } + if (c.mtp) { + const m = c.layers; + try b.tensor("mtp.enorm", try fullName(&f, "layers.{d}.enorm.weight", .{m}), false); + try b.tensor("mtp.hnorm", try fullName(&f, "layers.{d}.hnorm.weight", .{m}), false); + try b.linear("mtp.eh", try fullName(&f, "layers.{d}.eh_proj.weight", .{m})); + try b.tensor("mtp.norm", try fullName(&f, "layers.{d}.shared_head.norm.weight", .{m}), false); + try layer(b, "mtp.layer", m, false, false, true); + } +} + +pub fn load(gpa: std.mem.Allocator, io: std.Io, ctx: *const api.Ctx, dir: []const u8, rank: u8) !api.Store { + var ck = try Ckpt.open(gpa, io, dir); + defer ck.deinit(); + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const c = try Cfg.read(arena.allocator(), io, dir); + var store = api.Store.init(gpa); + errdefer store.deinit(); + var b: Builder = .{ .gpa = gpa, .ck = &ck, .dry = false, .rank = rank, .store = &store, .d = ctx.d, + .dense_bits = try envBits("TF_GLM_DENSE_BITS"), .head_bits = try envBits("TF_GLM_HEAD_BITS"), .draft_head_bits = try envBits("TF_GLM_DRAFT_HEAD_BITS"), .kda_proj_bits = try envBits("TF_GLM_KDA_PROJ_BITS"), .experts_sym = envSym(), .expert_sim_bits = envSimBits(), .experts3 = envBits3(), .sim_proj = envSimProj() }; + b.exl3 = c.exl3; + if (c.exl3) b.experts_sym = false; + try build(&b, &c); + const vl = c.vocab / WORLD; + const facts = [_]struct { []const u8, i64 }{ + .{ "rank", rank }, .{ "world", WORLD }, .{ "vocab", @intCast(c.vocab) }, .{ "vocab_shard", @intCast(vl) }, + .{ "vocab_offset", @intCast(rank * vl) }, .{ "layers", @intCast(c.layers) }, .{ "dense", @intCast(c.dense) }, + .{ "mtp", @intFromBool(c.mtp) }, .{ "device_bytes", @intCast(b.bytes) }, .{ "experts_sym", @intFromBool(b.experts_sym) }, .{ "experts3", @intFromBool(b.experts3) }, .{ "exl3", @intFromBool(b.exl3) }, + .{ "exl3_shx", @intFromBool(b.exl3 and b.exl3_shx and !b.dry) }, + }; + for (facts) |fct| try store.ints.put(try gpa.dupe(u8, fct[0]), fct[1]); + var kinds: [64]u8 = undefined; + for (0..c.layers) |i| kinds[i] = @intFromBool(c.kinds[i]); + try store.host.put(try gpa.dupe(u8, "kinds"), try gpa.dupe(u8, kinds[0..c.layers])); + return store; +} + +/// Device bytes this rank's prepared weights take: the same preparation's shapes, nothing read but headers. +pub fn deviceBytes(io: std.Io, dir: []const u8) !u64 { + const gpa = std.heap.page_allocator; + var ck = try Ckpt.open(gpa, io, dir); + defer ck.deinit(); + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const c = try Cfg.read(arena.allocator(), io, dir); + var b: Builder = .{ .gpa = gpa, .ck = &ck, .dry = true, .rank = 0, .dense_bits = try envBits("TF_GLM_DENSE_BITS"), + .head_bits = try envBits("TF_GLM_HEAD_BITS"), .draft_head_bits = try envBits("TF_GLM_DRAFT_HEAD_BITS"), .kda_proj_bits = try envBits("TF_GLM_KDA_PROJ_BITS"), .experts_sym = envSym(), .expert_sim_bits = envSimBits(), .experts3 = envBits3(), .sim_proj = envSimProj() }; + b.exl3 = c.exl3; + try build(&b, &c); + return b.bytes; +} + +// --------------------------------------------------------------------------------------------------------------- +// Packed dense linears (ops.cu's glm_qgemv layout). + +fn envBits3() bool { + const v = std.c.getenv("TF_GLM_EXPERT_BITS") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "3"); +} + +fn envSimProj() u8 { + const v = std.c.getenv("TF_GLM_EXPERT_SIM_PROJ") orelse return 0; + const t = std.mem.span(v); + return if (std.mem.eql(u8, t, "down")) 1 else if (std.mem.eql(u8, t, "up")) 2 else 0; +} + +fn envSimBits() u8 { + const v = std.c.getenv("TF_GLM_EXPERT_SIM_BITS") orelse return 0; + return std.fmt.parseInt(u8, std.mem.span(v), 10) catch 0; +} + +/// Each row's group of 64 re-quantized to ``bits`` symmetric levels (-2^(bits-1) .. 2^(bits-1) - 1, the scale from a +/// search over clip ratios for the least squared error), written back as 4-bit codes + 8 with that scale (bias -8 s). +fn simulateBits(t: Triple, bits: u8) void { + const Work = struct { + fn run(tt: Triple, nb: u8, e0: usize, e1: usize) void { + const lo: f32 = -@as(f32, @floatFromInt(@as(u32, 1) << @intCast(nb - 1))); + const hi: f32 = -lo - 1; + const wpg = 8; // 4-bit words a group of 64 + for (e0..e1) |e| for (0..tt.n) |r| for (0..tt.kg) |g| { + const at = (e * tt.n + r) * tt.kg + g; + const words = tt.words[(e * tt.n + r) * tt.k8 + g * wpg ..][0..wpg]; + const s0 = bf16Value(tt.scales[at]); + var v: [64]f32 = undefined; + var amax: f32 = 0; + for (0..64) |i| { + const q: i32 = @intCast((words[i / 8] >> @intCast(4 * (i % 8))) & 15); + v[i] = @as(f32, @floatFromInt(q - 8)) * s0; + amax = @max(amax, @abs(v[i])); + } + if (amax == 0) continue; + var best_s: f32 = amax / -lo; + var best_e: f32 = std.math.inf(f32); + for ([_]f32{ 1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6 }) |c| { + const sc: f32 = bf16Value(toBf16(c * amax / -lo)); + if (sc == 0) continue; + var err: f32 = 0; + for (v) |x| { + const q = @max(lo, @min(hi, @round(x / sc))); + err += (q * sc - x) * (q * sc - x); + } + if (err < best_e) { + best_e = err; + best_s = sc; + } + } + for (words) |*w| w.* = 0; + for (v, 0..) |x, i| { + const q: i32 = @intFromFloat(@max(lo, @min(hi, @round(x / best_s)))); + words[i / 8] |= @as(u32, @intCast(q + 8)) << @intCast(4 * (i % 8)); + } + tt.scales[at] = toBf16(best_s); + tt.biases[at] = toBf16(-8 * best_s); + }; + } + }; + const threads = 16; + var pool: [threads]?std.Thread = @splat(null); + const per = (t.e + threads - 1) / threads; + for (0..threads) |i| { + const e0 = i * per; + const e1 = @min(t.e, e0 + per); + if (e0 >= e1) break; + pool[i] = std.Thread.spawn(.{}, Work.run, .{ t, bits, e0, e1 }) catch blk: { + Work.run(t, bits, e0, e1); + break :blk null; + }; + } + for (pool) |th| if (th) |x| x.join(); +} + +fn toBf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); // round to nearest even +} + +fn envSym() bool { + const v = std.c.getenv("TF_GLM_EXPERTS_SYM") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// TF_GLM_DENSE_BITS / TF_GLM_HEAD_BITS: 4, 6 or 8 (unset or 0: bf16 as stored). +fn envBits(name: [*:0]const u8) !u8 { + const v = std.c.getenv(name) orelse return 0; + const bits = std.fmt.parseInt(u8, std.mem.span(v), 10) catch return error.BadPackedBits; + if (bits != 0 and bits != 4 and bits != 6 and bits != 8) return error.BadPackedBits; + return bits; +} + +fn bf16Value(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// x rounded to the nearest integer, halves to even. +fn roundEven(x: f32) f32 { + 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; +} + +/// bf16 rows [n, k] -> symmetric ``bits``-bit values: each group of 64 inputs a scale absmax / (2^(bits-1) - 1) +/// rounded to fp16, values round(w / scale) (halves to even) clamped to +-(2^(bits-1) - 1), stored + 2^(bits-1) in +/// 16-value chunks of bits / 2 words (``words`` [n * k * bits / 32]); ``scales`` [n * k / 64] fp16 bits. +pub fn quantize(bits: u8, src: []const u16, n: usize, k: usize, words: []u32, scales: []u16) void { + const half: i32 = @as(i32, 1) << @intCast(bits - 1); + const qmax: f32 = @floatFromInt(half - 1); + const wpc: usize = bits / 2; + const groups = k / 64; + for (0..n) |r| { + const row = src[r * k ..][0..k]; + for (0..groups) |g| { + var amax: f32 = 0; + for (row[g * 64 ..][0..64]) |h| amax = @max(amax, @abs(bf16Value(h))); + const s16: f16 = @floatCast(amax / qmax); + scales[r * groups + g] = @bitCast(s16); + const sc: f32 = s16; + for (0..4) |c| { + var q: [16]u32 = undefined; + for (0..16) |i| { + const v = bf16Value(row[g * 64 + c * 16 + i]); + const t: i32 = if (sc == 0) 0 else @intFromFloat(@max(-qmax, @min(qmax, roundEven(v / sc)))); + q[i] = @intCast(t + half); + } + pack(bits, &q, words[((r * groups + g) * 4 + c) * wpc ..][0..wpc]); + } + } + } +} + +fn pack(bits: u8, q: *const [16]u32, out: []u32) void { + @memset(out, 0); + switch (bits) { + 4 => for (0..16) |i| { + out[i / 8] |= q[i] << @intCast(4 * (i % 8)); + }, + 8 => for (0..16) |i| { + out[i / 4] |= q[i] << @intCast(8 * (i % 4)); + }, + else => { + for (0..15) |i| out[i / 5] |= q[i] << @intCast(6 * (i % 5)); + for (0..3) |j| out[j] |= ((q[15] >> @intCast(2 * j)) & 3) << 30; + }, + } +} + +test "packed rows: values and scales round trip (4, 6, 8 bits)" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(7); + const n = 3; + const k = 128; + var src: [n * k]u16 = undefined; + for (&src) |*h| h.* = @truncate(@as(u32, @bitCast(prng.random().floatNorm(f32) * 0.02)) >> 16); + src[5] = 0; // a zero inside a group + for ([_]u8{ 4, 6, 8 }) |bits| { + const words = try gpa.alloc(u32, n * k * bits / 32); + defer gpa.free(words); + var scales: [n * k / 64]u16 = undefined; + quantize(bits, &src, n, k, words, &scales); + const half: i32 = @as(i32, 1) << @intCast(bits - 1); + const qmax: f32 = @floatFromInt(half - 1); + const wpc: usize = bits / 2; + for (0..n * k / 16) |c| { + const p = words[c * wpc ..][0..wpc]; + for (0..16) |i| { + const raw: u32 = switch (bits) { + 4 => (p[i / 8] >> @intCast(4 * (i % 8))) & 15, + 8 => (p[i / 4] >> @intCast(8 * (i % 4))) & 255, + else => if (i < 15) (p[i / 5] >> @intCast(6 * (i % 5))) & 63 else (p[0] >> 30) | ((p[1] >> 30) << 2) | ((p[2] >> 30) << 4), + }; + const sc: f32 = @as(f16, @bitCast(scales[c / 4])); + const got = @as(f32, @floatFromInt(@as(i32, @intCast(raw)) - half)) * sc; + const want = bf16Value(src[c * 16 + i]); + try std.testing.expect(@abs(got - want) <= sc * 0.5 + 1e-7); + try std.testing.expect(@abs(@as(f32, @floatFromInt(@as(i32, @intCast(raw)) - half))) <= qmax); + } + } + } +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index a34186a4a3..f5d2fb78f5 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -7,13 +7,15 @@ const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const glm = @import("glm_cuda"); +const with_glm = @import("cuda_families").glm; 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 = if (with_glm) .{ nemotron.native, glm.native } else .{nemotron.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 = .{ diff --git a/zig/src/native/no_family.zig b/zig/src/native/no_family.zig deleted file mode 100644 index d81d7567ac..0000000000 --- a/zig/src/native/no_family.zig +++ /dev/null @@ -1,2 +0,0 @@ -//! Stands in for a CUDA family the build leaves out (-Dflashnext=false): nothing is registered. -pub const native = struct {}; diff --git a/zig/tests/cuda/glm/gen_copies.py b/zig/tests/cuda/glm/gen_copies.py new file mode 100644 index 0000000000..ebc386b62e --- /dev/null +++ b/zig/tests/cuda/glm/gen_copies.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python3 +"""The device code of the Python GLM-5.3-Flash engine's own CUDA extension (families/glm5_next/cuda/kda.cu) as a +zig/kernels/cuda/glm copy: Flash Next's gen_copies.py rules (the source up to its first ATen-using function, ATen +includes dropped, the anonymous namespace kept) plus explicit instantiations of every template kernel the Python +build instantiated (demangled from its cubin). The grouped experts and the 4-bit matmuls are the shared extensions +Flash Next's copies already hold (the same symbols). +usage: gen_copies.py """ +import json, sys +from pathlib import Path + +fn: dict = {} +text_fn = (Path(__file__).parent.parent / "flashnext" / "gen_copies.py").read_text().replace("sys.exit(main())", "") +exec(compile(text_fn, "flashnext/gen_copies.py", "exec"), fn) # its rules, without running its main +src, syms, out = Path(sys.argv[1]), json.loads(Path(sys.argv[2]).read_text()), Path(sys.argv[3]) +out.mkdir(parents=True, exist_ok=True) +rel = "families/glm5_next/cuda/kda.cu" +body = fn["device_part"]((src / rel).read_text()) +kda = [s for k, v in syms.items() if "glm_kda" in k for s in v] +inst = fn["instantiations"](kda) +head = f"// Device code of src/tensorfold/{rel} (generated by zig/tests/cuda/glm/gen_copies.py; do not edit).\n" +text = head + body + ("\n// the Python build's instantiations\nnamespace {\n" + "\n".join(inst) + "\n} // namespace\n" + if inst else "") +(out / "kda.cu").write_text(text) +print(f"kda.cu: {len(inst)} instantiations") diff --git a/zig/tests/cuda/glm_forward.zig b/zig/tests/cuda/glm_forward.zig new file mode 100644 index 0000000000..989bd1bc6f --- /dev/null +++ b/zig/tests/cuda/glm_forward.zig @@ -0,0 +1,709 @@ +//! GLM-5.3-Flash on CUDA: the family's kernels against float64 references and the row-invariance contract (a row's +//! bits alone equal its bits in any window, any mix of streams), on one GPU; and the two-rank engine's contract on +//! the real checkpoint: a stream's greedy reply is the same alone, beside 15 other streams, drafted, and resumed from +//! a kept prompt state. +//! usage: tf-cuda-test glm-kcheck +//! tf-cuda-test glm-batch [tokens] [sync] + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const gw = @import("glm_weights"); +const kern = gw.kern; +const fwd = gw.forward; +const dec = gw.decode; +const Gpu = check.Gpu; + +fn bf(x: u16) f32 { + return @bitCast(@as(u32, x) << 16); +} + +fn tobf(v: f32) u16 { // round to nearest even + const u: u32 = @bitCast(v); + return @truncate((u + 0x7FFF + ((u >> 16) & 1)) >> 16); +} + +fn randBf(rnd: std.Random, out: []u16, scale: f32) void { + for (out) |*h| h.* = tobf(rnd.floatNorm(f32) * scale); +} + +fn ms(io: std.Io, t0: std.Io.Timestamp) f64 { + return @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(io, .awake)).toNanoseconds())) / 1e6; +} + +fn up(d: *const cuda.Driver, bytes: []const u8) !cuda.DeviceBuffer { + const b = try cuda.DeviceBuffer.alloc(d, @max(bytes.len, 16)); + try b.upload(0, bytes); + return b; +} + +fn down(gpa: std.mem.Allocator, d: *const cuda.Driver, comptime T: type, b: cuda.DeviceBuffer, n: usize) ![]T { + const out = try gpa.alloc(T, n); + try d.check(d.api.cuMemcpyDtoH_v2(@ptrCast(out.ptr), b.ptr, n * @sizeOf(T)), "download"); + return out; +} + +const Stats = struct { + worst: f64 = 0, + scale: f64 = 0, + fn add(s: *Stats, got: f64, want: f64) void { + s.worst = @max(s.worst, @abs(got - want)); + s.scale = @max(s.scale, @abs(want)); + } + fn rel(s: Stats) f64 { + return s.worst / @max(s.scale, 1e-30); + } +}; + +// -- glm-kcheck --------------------------------------------------------------------------------------------------------- +fn checkMm(gpu: Gpu, k: *kern.K, rnd: std.Random) !void { + const gpa = gpu.gpa; + const d = gpu.d; + const shapes = [_][2]usize{ .{ 1024, 2048 }, .{ 160, 4096 }, .{ 12576, 512 } }; + for (shapes) |sh| { + const n = sh[0]; + const kk = sh[1]; + const M: usize = 128; + const w16 = try gpa.alloc(u16, n * kk); + defer gpa.free(w16); + randBf(rnd, w16, 0.02); + const x16 = try gpa.alloc(u16, M * kk); + defer gpa.free(x16); + randBf(rnd, x16, 1.0); + var xb = try up(d, std.mem.sliceAsBytes(x16)); + defer xb.free(); + var ob = try cuda.DeviceBuffer.alloc(d, M * n * 4); + defer ob.free(); + var part = try cuda.DeviceBuffer.alloc(d, M * 16384 * 4); + defer part.free(); + for ([_]u8{ 0, 8, 6, 4 }) |bits| { + // the loader's row-major form, then tiled on the device as forward's load does + var src: cuda.DeviceBuffer = undefined; + var ss: ?cuda.DeviceBuffer = null; + const deq = try gpa.alloc(f64, n * kk); + defer gpa.free(deq); + if (bits == 0) { + src = try up(d, std.mem.sliceAsBytes(w16)); + for (w16, deq) |h, *v| v.* = bf(h); + } else { + const words = try gpa.alloc(u32, n * kk * bits / 32); + defer gpa.free(words); + const scales = try gpa.alloc(u16, n * kk / 64); + defer gpa.free(scales); + gw.quantize(bits, w16, n, kk, words, scales); + src = try up(d, std.mem.sliceAsBytes(words)); + ss = try up(d, std.mem.sliceAsBytes(scales)); + const half: i32 = @as(i32, 1) << @intCast(bits - 1); + const wpc: usize = bits / 2; + for (0..n * kk / 16) |c| { + const p = words[c * wpc ..][0..wpc]; + for (0..16) |i| { + const raw: u32 = switch (bits) { + 4 => (p[i / 8] >> @intCast(4 * (i % 8))) & 15, + 8 => (p[i / 4] >> @intCast(8 * (i % 4))) & 255, + else => if (i < 15) (p[i / 5] >> @intCast(6 * (i % 5))) & 63 else (p[0] >> 30) | ((p[1] >> 30) << 2) | ((p[2] >> 30) << 4), + }; + const sc: f64 = @as(f16, @bitCast(scales[c / 4])); + deq[c * 16 + i] = @as(f64, @floatFromInt(@as(i32, @intCast(raw)) - half)) * sc; + } + } + } + defer src.free(); + defer if (ss) |*x| x.free(); + const wbytes = n * kk * @as(usize, if (bits == 0) 16 else bits) / 8; + var tw = try cuda.DeviceBuffer.alloc(d, wbytes); + defer tw.free(); + var ts = try cuda.DeviceBuffer.alloc(d, n * kk / 64 * 2 + 16); + defer ts.free(); + try kern.mmTile(k, bits, src.ptr, if (ss) |x| x.ptr else 0, tw.ptr, ts.ptr, @intCast(n), @intCast(kk)); + const w: kern.B16 = .{ .w = tw.ptr, .n = @intCast(n), .k = @intCast(kk), .s = if (bits == 0) 0 else ts.ptr, .bits = bits }; + try kern.mm(k, xb.ptr, @intCast(kk), w, ob.ptr, true, @intCast(M), part.ptr); + try k.stream.synchronize(); + const full = try down(gpa, d, f32, ob, M * n); + defer gpa.free(full); + var st: Stats = .{}; + for (0..M) |r| { + if (r % 17 != 0 and r != M - 1) continue; // a sample of rows against float64 + for (0..n) |c| { + var acc: f64 = 0; + for (0..kk) |i| acc += @as(f64, bf(x16[r * kk + i])) * deq[c * kk + i]; + st.add(full[r * n + c], acc); + } + } + // a row alone and in narrower windows: the same bits as in the 128-row one + var same = true; + for ([_]usize{ 1, 3, 16, 17, 33, 64, 100 }) |m| { + try kern.mm(k, xb.ptr, @intCast(kk), w, ob.ptr, true, @intCast(m), part.ptr); + try k.stream.synchronize(); + const part_rows = try down(gpa, d, f32, ob, m * n); + defer gpa.free(part_rows); + if (!std.mem.eql(u32, @ptrCast(part_rows), @as([]const u32, @ptrCast(full[0 .. m * n])))) same = false; + } + std.debug.print("mm {d}x{d} bits {d}: rel err {e:.2} vs float64; rows alone == in windows: {}\n", .{ n, kk, if (bits == 0) @as(u8, 16) else bits, st.rel(), same }); + try check.expect(st.rel() < 1e-5 and same, "mm {d}x{d} bits {d}", .{ n, kk, bits }); + } + } +} + +/// Latent attention: rows of two streams (each its own cache), dense and sparse, against float64; each row alone +/// against the same row in the mixed window. +fn checkAttention(gpu: Gpu, k: *kern.K, rnd: std.Random) !void { + const gpa = gpu.gpa; + const d = gpu.d; + const H: usize = 32; + const LW: usize = 512; + const L: usize = 9000; // cache rows a stream + const caches = [2][]u16{ try gpa.alloc(u16, L * LW), try gpa.alloc(u16, L * LW) }; + for (caches) |c| randBf(rnd, c, 0.5); + var cb: [2]cuda.DeviceBuffer = .{ try up(d, std.mem.sliceAsBytes(caches[0])), try up(d, std.mem.sliceAsBytes(caches[1])) }; + defer for (&cb) |*b| b.free(); + // rows: stream 0 at positions 5, 600, 2050 (dense), stream 1 at 1999, 2051, 8000 (sparse past the limit) + const pos = [_]i32{ 5, 600, 2050, 1999, 2051, 8000 }; + const which = [_]usize{ 0, 0, 0, 1, 1, 1 }; + const R = pos.len; + const q16 = try gpa.alloc(u16, R * H * LW); + randBf(rnd, q16, 0.2); + // sparse rows' tokens: 2048 distinct ascending tokens <= pos, then the 3 tail slots (-1) + const tok = try gpa.alloc(i32, R * 2051); + var cnt: [R]i32 = @splat(0); + for (0..R) |r| { + @memset(tok[r * 2051 ..][0..2051], -1); + if (pos[r] < 2051) continue; + var t: i32 = 0; + var j: usize = 0; + while (j < 2048) : (t += 1) { + if (rnd.float(f32) < 2048.0 / @as(f32, @floatFromInt(pos[r] + 1)) or @as(i32, @intCast(2048 - j)) > pos[r] - t) { + tok[r * 2051 + j] = t; + j += 1; + } + } + cnt[r] = 2048; + } + var lc: [R]u64 = undefined; + for (0..R) |r| lc[r] = cb[which[r]].ptr; + // the tables as forward.prepare lays them out (positions, then lc for cache 0; one cache here) + var qb = try up(d, std.mem.sliceAsBytes(q16)); + defer qb.free(); + var tb = try up(d, std.mem.sliceAsBytes(tok)); + defer tb.free(); + var nb = try up(d, std.mem.sliceAsBytes(&cnt)); + defer nb.free(); + var po = try cuda.DeviceBuffer.alloc(d, 5 * R * H * LW * 4); + defer po.free(); + var pm = try cuda.DeviceBuffer.alloc(d, 5 * R * H * 4); + defer pm.free(); + var pl = try cuda.DeviceBuffer.alloc(d, 5 * R * H * 4); + defer pl.free(); + var ob = try cuda.DeviceBuffer.alloc(d, R * H * LW * 2); + defer ob.free(); + const run_rows = struct { + fn f(kk: *kern.K, dd: *const cuda.Driver, q: u64, rows: []const usize, posv: []const i32, lcv: []const u64, tokv: u64, cntv: []const i32, b: [3]u64, out: u64) !void { + var p: [8]i32 = undefined; + var l: [8]u64 = undefined; + var c: [8]i32 = undefined; + for (rows, 0..) |r, i| { + p[i] = posv[r]; + l[i] = lcv[r]; + c[i] = cntv[r]; + } + const n = rows.len; + const lo = std.mem.alignForward(usize, n * 4, 16); // the pointer table 8-aligned, as forward.prepare's is + var tab = try cuda.DeviceBuffer.alloc(dd, lo + n * 8); + defer tab.free(); + try tab.upload(0, std.mem.sliceAsBytes(p[0..n])); + try tab.upload(lo, std.mem.sliceAsBytes(l[0..n])); + var cn = try up(dd, std.mem.sliceAsBytes(c[0..n])); + defer cn.free(); + var tk = try cuda.DeviceBuffer.alloc(dd, n * 2051 * 4); + defer tk.free(); + for (rows, 0..) |r, i| try kk.copy(tk.ptr + i * 2051 * 4, tokv + r * 2051 * 4, 2051 * 4); + var qq = try cuda.DeviceBuffer.alloc(dd, n * 32 * 512 * 2); + defer qq.free(); + for (rows, 0..) |r, i| try kk.copy(qq.ptr + i * 32 * 512 * 2, q + r * 32 * 512 * 2, 32 * 512 * 2); + const t: kern.Rows = .{ .pos = tab.ptr, .lc = tab.ptr + lo, .ik = 0, .ig = 0, .pk = 0, .cap = @intCast(n) }; + try kern.denseAttention(kk, qq.ptr, t, 0, b[0], b[1], b[2], out, @intCast(n), kern.DENSE_CHUNKS); + try kern.sparseAttention(kk, qq.ptr, t, 0, tk.ptr, cn.ptr, b[0], b[1], b[2], out, @intCast(n)); + try kk.stream.synchronize(); + } + }.f; + const all = [_]usize{ 0, 1, 2, 3, 4, 5 }; + try run_rows(k, d, qb.ptr, &all, &pos, &lc, tb.ptr, &cnt, .{ po.ptr, pm.ptr, pl.ptr }, ob.ptr); + const mixed = try down(gpa, d, u16, ob, R * H * LW); + var st: Stats = .{}; + for (0..R) |r| { + const cache = caches[which[r]]; + for (0..H) |h| { + const q = q16[(r * H + h) * LW ..][0..LW]; + var keys: std.ArrayList(usize) = .empty; + defer keys.deinit(gpa); + if (pos[r] < 2051) { + for (0..@intCast(pos[r] + 1)) |j| try keys.append(gpa, j); + } else for (tok[r * 2051 ..][0..@intCast(cnt[r])]) |t| try keys.append(gpa, @intCast(t)); + var m: f64 = -std.math.inf(f64); + const sc = try gpa.alloc(f64, keys.items.len); + defer gpa.free(sc); + for (keys.items, 0..) |j, i| { + var a: f64 = 0; + for (0..LW) |c| a += @as(f64, bf(q[c])) * bf(cache[j * LW + c]); + sc[i] = a * 0.0625; + m = @max(m, sc[i]); + } + var l: f64 = 0; + for (sc) |*s| { + s.* = @exp(s.* - m); + l += s.*; + } + for (0..LW) |c| { + var o: f64 = 0; + for (keys.items, 0..) |j, i| o += sc[i] * bf(cache[j * LW + c]); + st.add(bf(mixed[(r * H + h) * LW + c]), o / l); + } + } + } + // each row alone, and the stream-1 rows on their own + var same = true; + for (0..R) |r| { + try run_rows(k, d, qb.ptr, &.{r}, &pos, &lc, tb.ptr, &cnt, .{ po.ptr, pm.ptr, pl.ptr }, ob.ptr); + const one = try down(gpa, d, u16, ob, H * LW); + defer gpa.free(one); + if (!std.mem.eql(u16, one, mixed[r * H * LW ..][0 .. H * LW])) same = false; + } + std.debug.print("latent attention (dense and sparse, two streams): rel err {e:.2} vs float64; rows alone == mixed: {}\n", .{ st.rel(), same }); + try check.expect(st.rel() < 2e-2 and same, "latent attention", .{}); +} + +/// The indexer: pool scores against float64, the 512 best pools (ties to the lower pool) against a host selection +/// over the kernel's own scores, for rows of two streams. +fn checkIndex(gpu: Gpu, k: *kern.K, rnd: std.Random) !void { + const gpa = gpu.gpa; + const d = gpu.d; + const NP: usize = 4096; + const pk16 = try gpa.alloc(u16, NP * 128); + randBf(rnd, pk16, 0.3); + for (1000..1300) |p| @memcpy(pk16[p * 128 ..][0..128], pk16[999 * 128 ..][0..128]); // ties: 301 equal pools + var pkb = try up(d, std.mem.sliceAsBytes(pk16)); + defer pkb.free(); + const pos = [_]i32{ 2051, 5000, 16383, 9999 }; + const R = pos.len; + const qi = try gpa.alloc(u16, R * 32 * 128); + randBf(rnd, qi, 0.5); + const wv = try gpa.alloc(u16, R * 160); + randBf(rnd, wv, 1.0); + var qb = try up(d, std.mem.sliceAsBytes(qi)); + defer qb.free(); + var wb = try up(d, std.mem.sliceAsBytes(wv)); + defer wb.free(); + var tab = try cuda.DeviceBuffer.alloc(d, R * 4 + R * 8 * 4); + defer tab.free(); + try tab.upload(0, std.mem.sliceAsBytes(&pos)); + var pkp: [R]u64 = @splat(pkb.ptr); + try tab.upload(R * 4 + 3 * R * 8, std.mem.sliceAsBytes(&pkp)); // the pk table (lc, ik, ig unused) + const t: kern.Rows = .{ .pos = tab.ptr, .lc = tab.ptr + R * 4, .ik = tab.ptr + R * 4 + R * 8, .ig = tab.ptr + R * 4 + 2 * R * 8, .pk = tab.ptr + R * 4 + 3 * R * 8, .cap = R }; + var sb = try cuda.DeviceBuffer.alloc(d, R * NP * 4); + defer sb.free(); + var ob = try cuda.DeviceBuffer.alloc(d, R * 512 * 8); + defer ob.free(); + try kern.scores(k, qb.ptr, wb.ptr, 160, t, 0, sb.ptr, R, NP); + try kern.selectRows(k, sb.ptr, ob.ptr, NP, t, R); + try k.stream.synchronize(); + const sc = try down(gpa, d, f32, sb, R * NP); + const sel = try down(gpa, d, i64, ob, R * 512); + var st: Stats = .{}; + var picks_ok = true; + for (0..R) |r| { + const npool: usize = @intCast(@divFloor(pos[r] + 1, 4)); + for (0..NP) |p| { + if (p >= npool) { + if (sc[r * NP + p] != -std.math.inf(f32)) picks_ok = false; + continue; + } + var s: f64 = 0; + for (0..32) |h| { + var dot: f64 = 0; + for (0..128) |c| dot += @as(f64, bf(qi[(r * 32 + h) * 128 + c])) * bf(pk16[p * 128 + c]); + s += @as(f64, bf(wv[r * 160 + h])) / 5.656854249492381 * @max(dot * 0.08838834764831845, 0); + } + st.add(sc[r * NP + p], s); + } + // the host's 512: by score descending, the lower pool first on ties, then ascending + const idx = try gpa.alloc(usize, @min(NP, @max(512, npool))); + defer gpa.free(idx); + for (idx, 0..) |*x, i| x.* = i; + const Ctx = struct { + s: []const f32, + fn lt(c: @This(), a: usize, b: usize) bool { + const x = c.s[a] + 0.0; + const y = c.s[b] + 0.0; + return x > y or (x == y and a < b); + } + }; + std.mem.sort(usize, idx, Ctx{ .s = sc[r * NP ..][0..NP] }, Ctx.lt); + std.mem.sort(usize, idx[0..512], {}, std.sort.asc(usize)); + for (0..512) |i| if (sel[r * 512 + i] != idx[i]) { + picks_ok = false; + }; + } + std.debug.print("indexer: scores rel err {e:.2} vs float64; the 512 pools equal the host's selection: {}\n", .{ st.rel(), picks_ok }); + try check.expect(st.rel() < 1e-4 and picks_ok, "indexer", .{}); +} + +pub fn kcheck(gpu: Gpu, _: []const [:0]const u8) !void { + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var kernels = try gw.api.Kernels.load(gpu.d); + defer kernels.deinit(); + var k = try kern.K.init(gpu.gpa, gpu.d, gpu.ctx, stream, &kernels.ext); + var prng = std.Random.DefaultPrng.init(11); + try checkMm(gpu, &k, prng.random()); + try checkAttention(gpu, &k, prng.random()); + try checkIndex(gpu, &k, prng.random()); + std.debug.print("glm-kcheck: all passed\n", .{}); +} + +// -- glm-batch ---------------------------------------------------------------------------------------------------------- +const Stream = struct { s: *fwd.Seq, out: std.ArrayList(u32) = .empty, pending: u32 = 0, held: [7]u32 = undefined, nheld: usize = 0, done: bool = false }; + +/// Greedy decode of every stream to ``n`` tokens, together: windows of the pending token and up to ``depth`` drafts, +/// the longest agreeing prefix kept, the head drafting from the kept rows (the lane core's rule, without its policy). +fn decodeAll(e: *fwd.Engine, ss: []Stream, n: usize, depth: usize) !void { + var picks: [fwd.STEP_ROWS]dec.Pick = undefined; + var wins: [fwd.MAX_STREAMS]dec.Win = undefined; + var ids: [fwd.MAX_STREAMS][8]u32 = undefined; + var map: [fwd.MAX_STREAMS]usize = undefined; + var probs: [fwd.MAX_STREAMS][7]f64 = undefined; + var nexts: [fwd.MAX_STREAMS][8]u32 = undefined; + while (true) { + var nw: usize = 0; + for (ss, 0..) |*x, i| { + if (x.done) continue; + ids[nw][0] = x.pending; + @memcpy(ids[nw][1..][0..x.nheld], x.held[0..x.nheld]); + wins[nw] = .{ .s = x.s, .ids = ids[nw][0 .. 1 + x.nheld] }; + map[nw] = i; + nw += 1; + } + if (nw == 0) return; + try dec.verify(e, wins[0..nw], &picks); + var reqs: [fwd.MAX_STREAMS]dec.DraftReq = undefined; + var nr: usize = 0; + var at: usize = 0; + for (0..nw) |wi| { + const x = &ss[map[wi]]; + const rows = wins[wi].ids.len; + var kept: usize = 1; + while (kept < rows and wins[wi].ids[kept] == picks[at + kept - 1].tok) kept += 1; + try dec.keep(e, x.s, @intCast(rows), @intCast(kept)); + for (0..kept) |r| { + try x.out.append(e.gpa, picks[at + r].tok); + nexts[wi][r] = picks[at + r].tok; + } + x.pending = picks[at + kept - 1].tok; + x.nheld = 0; + if (x.out.items.len >= n) x.done = true else if (depth > 0) { + reqs[nr] = .{ .s = x.s, .hidden = x.s.hidden, .next = nexts[wi][0..kept], .count = depth, .confidence = 0.7, .out = x.held[0..depth], .probs = probs[nr][0..depth] }; + nr += 1; + } + at += rows; + } + if (nr > 0) { + try dec.draft(e, reqs[0..nr]); + var ri: usize = 0; + for (0..nw) |wi| { + const x = &ss[map[wi]]; + if (x.done) continue; + x.nheld = reqs[ri].got; + ri += 1; + } + } + } +} + +pub fn batch(gpu: Gpu, args: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const io = gpu.io; + if (args.len < 4) return error.Usage; + const rank = try std.fmt.parseInt(u8, args[0], 10); + const port = try std.fmt.parseInt(u16, args[2], 10); + const ntok: usize = if (args.len > 4) try std.fmt.parseInt(usize, args[4], 10) else 48; + var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) + else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); + defer link.close(); + var nccl = try cuda.nccl.Library.open(); + defer nccl.close(); + const comm = try cuda.tp_link.communicator(&nccl, link); + // the communicator's threads stop before the library unloads (an unloaded NCCL under live threads crashes the exit) + defer _ = nccl.api.ncclCommAbort(comm); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + { + const warm = try cuda.DeviceBuffer.alloc(gpu.d, 4096); + try nccl.check(nccl.api.ncclAllGather(warm.ptr, warm.ptr + 2048, 256, .i32, comm, stream.handle), "warm gather"); + try stream.synchronize(); + } + const t0 = std.Io.Timestamp.now(io, .awake); + var kernels = try gw.api.Kernels.load(gpu.d); + var ctx: gw.api.Ctx = .{ .d = gpu.d, .ctx = gpu.ctx, .stream = stream, .nccl = &nccl, .comm = comm, .rank = rank }; + var store = try gw.load(gpa, io, &ctx, args[3], rank); + defer store.deinit(); + const e = try fwd.init(gpa, io, &ctx, &kernels, &store, .{ .context = 65536 }); + e.k.sync_each = args.len > 5 and std.mem.eql(u8, args[5], "sync"); + e.sync_each = e.k.sync_each; + std.debug.print("rank {d}: engine up in {d:.1} s\n", .{ rank, ms(io, t0) / 1000 }); + const S = fwd.MAX_STREAMS; + var prompts: [S][]u32 = undefined; + var prng = std.Random.DefaultPrng.init(5); + for (&prompts, 0..) |*p, i| { + p.* = try gpa.alloc(u32, 40 + i * 37); + for (p.*) |*t| t.* = 1000 + prng.random().uintLessThan(u32, 60000); + } + var solo: [S][]u32 = undefined; + // each stream alone, serial (one-row windows, no drafts) + var ts = std.Io.Timestamp.now(io, .awake); + for (0..S) |i| { + var one = [_]Stream{.{ .s = try fwd.newSeq(e, 4096) }}; + one[0].pending = try dec.prefill(e, one[0].s, prompts[i], 0, 0); + try one[0].out.append(gpa, one[0].pending); + try decodeAll(e, &one, ntok, 0); + solo[i] = one[0].out.items; + fwd.freeSeq(e, one[0].s); + } + const solo_ms = ms(io, ts); + std.debug.print("rank {d}: 16 streams one at a time, serial: {d:.1} tok/s\n", .{ rank, @as(f64, @floatFromInt(S * ntok)) / solo_ms * 1000 }); + // all together, serial, then drafted; and one stream alone drafted + for ([_]usize{ 0, 3 }) |depth| { + var all: [S]Stream = undefined; + for (&all, 0..) |*x, i| { + x.* = .{ .s = try fwd.newSeq(e, 4096) }; + x.pending = try dec.prefill(e, x.s, prompts[i], 0, if (depth > 0) prompts[i].len - 1 else 0); + try x.out.append(gpa, x.pending); + } + ts = std.Io.Timestamp.now(io, .awake); + try decodeAll(e, &all, ntok, depth); + const t_ms = ms(io, ts); + var equal: usize = 0; + for (all, 0..) |x, i| equal += @intFromBool(std.mem.eql(u32, x.out.items[0..ntok], solo[i][0..ntok])); + std.debug.print("rank {d}: 16 streams together, {s}: {d} of 16 replies equal to alone, {d:.1} tok/s aggregate\n", .{ rank, + if (depth > 0) "drafted" else "serial", equal, @as(f64, @floatFromInt(S * (ntok - 1))) / t_ms * 1000 }); + try check.expect(equal == S, "16 streams together equal alone", .{}); + if (depth > 0) { // resumed: stream 0's kept state at its prompt's end - 1, a longer prompt extending it + const ext = try std.mem.concat(gpa, u32, &.{ prompts[0], solo[0][0..8] }); + var fresh = [_]Stream{.{ .s = try fwd.newSeq(e, 4096) }}; + fresh[0].pending = try dec.prefill(e, fresh[0].s, ext, 0, 0); + try fresh[0].out.append(gpa, fresh[0].pending); + try decodeAll(e, &fresh, 16, 0); + var res = [_]Stream{.{ .s = all[0].s }}; + res[0].pending = try dec.prefill(e, res[0].s, ext, prompts[0].len - 1, 0); + try res[0].out.append(gpa, res[0].pending); + try decodeAll(e, &res, 16, 0); + const ok = std.mem.eql(u32, res[0].out.items[0..16], fresh[0].out.items[0..16]); + std.debug.print("rank {d}: resumed from the kept state == fresh: {}\n", .{ rank, ok }); + try check.expect(ok, "resumed equals fresh", .{}); + fwd.freeSeq(e, fresh[0].s); + } + for (all) |x| fwd.freeSeq(e, x.s); + } + { + var one = [_]Stream{.{ .s = try fwd.newSeq(e, 4096) }}; + one[0].pending = try dec.prefill(e, one[0].s, prompts[3], 0, prompts[3].len - 1); + try one[0].out.append(gpa, one[0].pending); + ts = std.Io.Timestamp.now(io, .awake); + try decodeAll(e, &one, 4 * ntok, 3); + const t_ms = ms(io, ts); + std.debug.print("rank {d}: one stream drafted: {d:.1} tok/s (graphs {d})\n", .{ rank, @as(f64, @floatFromInt(4 * ntok - 1)) / t_ms * 1000, e.captures }); + fwd.freeSeq(e, one[0].s); + } + std.debug.print("rank {d}: glm-batch passed\n", .{rank}); +} + +// -- glm-tf ------------------------------------------------------------------------------------------------------------- +/// Teacher-forced against the Python engine's greedy reply (capture_ref.py's reference.json: prompt, first, serial): +/// the prompt, then the reply's tokens in windows of 8; how often this engine's top token is Python's next token. +/// usage: tf-cuda-test glm-tf +pub fn teacher(gpu: Gpu, args: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const io = gpu.io; + if (args.len < 5) return error.Usage; + const rank = try std.fmt.parseInt(u8, args[0], 10); + const port = try std.fmt.parseInt(u16, args[2], 10); + var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) + else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); + defer link.close(); + var nccl = try cuda.nccl.Library.open(); + defer nccl.close(); + const comm = try cuda.tp_link.communicator(&nccl, link); + // the communicator's threads stop before the library unloads (an unloaded NCCL under live threads crashes the exit) + defer _ = nccl.api.ncclCommAbort(comm); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + { + const warm = try cuda.DeviceBuffer.alloc(gpu.d, 4096); + try nccl.check(nccl.api.ncclAllGather(warm.ptr, warm.ptr + 2048, 256, .i32, comm, stream.handle), "warm gather"); + try stream.synchronize(); + } + var kernels = try gw.api.Kernels.load(gpu.d); + var ctx: gw.api.Ctx = .{ .d = gpu.d, .ctx = gpu.ctx, .stream = stream, .nccl = &nccl, .comm = comm, .rank = rank }; + var store = try gw.load(gpa, io, &ctx, args[3], rank); + defer store.deinit(); + const e = try fwd.init(gpa, io, &ctx, &kernels, &store, .{ .context = 65536 }); + const text = try std.Io.Dir.cwd().readFileAlloc(io, args[4], gpa, .limited(1 << 26)); + const ref = (try std.json.parseFromSlice(std.json.Value, gpa, text, .{})).value.object; + const ints = struct { + fn f(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const it = v.array.items; + const out = try a.alloc(u32, it.len); + for (it, 0..) |x, i| out[i] = @intCast(x.integer); + return out; + } + }.f; + const prompt = try ints(gpa, ref.get("prompt").?); + const serial = try ints(gpa, ref.get("serial").?); + const s = try fwd.newSeq(e, 8192); + const first = try dec.prefill(e, s, prompt, 0, 0); + var hits: usize = @intFromBool(first == serial[0]); + var picks: [8]dec.Pick = undefined; + var j: usize = 0; + while (j + 1 < serial.len) { + const w = @min(8, serial.len - 1 - j); + try dec.verify(e, &.{.{ .s = s, .ids = serial[j .. j + w] }}, &picks); + try dec.keep(e, s, @intCast(w), @intCast(w)); + for (0..w) |r| hits += @intFromBool(picks[r].tok == serial[j + r + 1]); + j += w; + } + std.debug.print("rank {d}: teacher-forced on {d} prompt tokens: the top token equals Python's next token at {d} of {d} positions\n", .{ rank, prompt.len, hits, serial.len }); +} + +// -- glm-prof ----------------------------------------------------------------------------------------------------------- +/// One stream's verify windows (1 and 4 rows), eager with every launch timed: where a step's time goes (proportions: +/// each launch synchronized), and the graph-replayed step's time for comparison. +/// usage: tf-cuda-test glm-prof [streams] +pub fn prof(gpu: Gpu, args: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const io = gpu.io; + if (args.len < 4) return error.Usage; + const rank = try std.fmt.parseInt(u8, args[0], 10); + const port = try std.fmt.parseInt(u16, args[2], 10); + const nstreams: usize = if (args.len > 4) try std.fmt.parseInt(usize, args[4], 10) else 1; + const graphs_only = args.len > 5 and std.mem.eql(u8, args[5], "graphs"); // for nsys: 50 one-row steps, no eager pass + var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) + else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); + defer link.close(); + var nccl = try cuda.nccl.Library.open(); + defer nccl.close(); + const comm = try cuda.tp_link.communicator(&nccl, link); + // the communicator's threads stop before the library unloads (an unloaded NCCL under live threads crashes the exit) + defer _ = nccl.api.ncclCommAbort(comm); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + { + const warm = try cuda.DeviceBuffer.alloc(gpu.d, 4096); + try nccl.check(nccl.api.ncclAllGather(warm.ptr, warm.ptr + 2048, 256, .i32, comm, stream.handle), "warm gather"); + try stream.synchronize(); + } + var kernels = try gw.api.Kernels.load(gpu.d); + var ctx: gw.api.Ctx = .{ .d = gpu.d, .ctx = gpu.ctx, .stream = stream, .nccl = &nccl, .comm = comm, .rank = rank }; + const no_roce = if (std.c.getenv("TF_ROCE")) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; + ctx.roce = try cuda.roce.Gather.init(gpa, io, gpu.d, link, !no_roce, 512 << 10); + defer if (ctx.roce) |g| g.deinit(); + var store = try gw.load(gpa, io, &ctx, args[3], rank); + defer store.deinit(); + const e = try fwd.init(gpa, io, &ctx, &kernels, &store, .{ .context = 65536 }); + if (std.c.getenv("TF_PROF_PROMPT")) |v| { // a prompt of this many tokens: its wall time, then its kernels (eager chunks) + const n = try std.fmt.parseInt(usize, std.mem.span(v), 10); + const p = try gpa.alloc(u32, n); + defer gpa.free(p); + for (p, 0..) |*t, i| t.* = @intCast(1000 + (i * 7919) % 60000); + const s0 = try fwd.newSeq(e, @intCast(n + 1024)); + const t0 = std.Io.Timestamp.now(io, .awake); + _ = try dec.prefill(e, s0, p, 0, 0); + const wall = ms(io, t0); + if (rank == 0) std.debug.print("prompt of {d} tokens: {d:.0} ms, {d:.0} tok/s\n", .{ n, wall, @as(f64, @floatFromInt(n)) / wall * 1000 }); + var m = std.StringHashMap(f64).init(gpa); + e.k.prof = &m; + e.k.io = io; + const s1 = try fwd.newSeq(e, @intCast(n + 1024)); + _ = try dec.prefill(e, s1, p, 0, 0); + e.k.prof = null; + var total: f64 = 0; + var it = m.iterator(); + const Row = struct { name: []const u8, ms: f64 }; + var rows: std.ArrayList(Row) = .empty; + while (it.next()) |x| { + total += x.value_ptr.*; + try rows.append(gpa, .{ .name = x.key_ptr.*, .ms = x.value_ptr.* }); + } + std.mem.sort(Row, rows.items, {}, struct { + fn lt(_: void, a: Row, b: Row) bool { + return a.ms > b.ms; + } + }.lt); + if (rank == 0) { + std.debug.print("prompt of {d} tokens: eager launches sum {d:.0} ms (synchronized: proportions only)\n", .{ n, total }); + for (rows.items[0..@min(rows.items.len, 30)]) |x| std.debug.print(" {d:8.1} ms {d:5.1}% {s}\n", .{ x.ms, 100 * x.ms / total, x.name[0..@min(x.name.len, 90)] }); + } + return; + } + var seqs: [fwd.MAX_STREAMS]*fwd.Seq = undefined; + var prompt: [512]u32 = undefined; + for (&prompt, 0..) |*t, i| t.* = @intCast(1000 + (i * 7919) % 60000); + var tok: [fwd.MAX_STREAMS]u32 = undefined; + for (0..nstreams) |i| { + seqs[i] = try fwd.newSeq(e, 8192); + prompt[0] = @intCast(2000 + i); + tok[i] = try dec.prefill(e, seqs[i], &prompt, 0, 0); + } + var picks: [fwd.STEP_ROWS]dec.Pick = undefined; + var ids: [fwd.MAX_STREAMS][8]u32 = undefined; + var wins: [fwd.MAX_STREAMS]dec.Win = undefined; + for ([_]usize{ 1, 4, 8 }) |width| { + if (graphs_only and width != 1) break; + const step = struct { + fn f(ee: *fwd.Engine, ss: []*fwd.Seq, tk: []u32, idv: *[fwd.MAX_STREAMS][8]u32, wv: *[fwd.MAX_STREAMS]dec.Win, pk: []dec.Pick, w: usize) !void { + for (ss, 0..) |s, i| { + for (0..w) |r| idv[i][r] = if (r == 0) tk[i] else @intCast(3000 + r * 101 + i); + wv[i] = .{ .s = s, .ids = idv[i][0..w] }; + } + try dec.verify(ee, wv[0..ss.len], pk); + for (ss, 0..) |s, i| { + try dec.keep(ee, s, @intCast(w), 1); + tk[i] = pk[i * w].tok; + } + } + }.f; + for (0..8) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); // graphs captured, clocks up + var best: f64 = std.math.inf(f64); + for (0..10) |_| { + const t0 = std.Io.Timestamp.now(io, .awake); + try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); + best = @min(best, ms(io, t0)); + } + if (graphs_only) { + for (0..40) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); + if (rank == 0) std.debug.print("{d} stream(s) x {d} rows: graph step {d:.2} ms (graphs only)\n", .{ nstreams, width, best }); + continue; + } + var m = std.StringHashMap(f64).init(gpa); + e.k.prof = &m; + e.k.io = io; + e.use_graphs = false; + for (0..4) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); + e.k.prof = null; + e.use_graphs = true; + var total: f64 = 0; + var it = m.iterator(); + const Row = struct { name: []const u8, ms: f64 }; + var rows: std.ArrayList(Row) = .empty; + while (it.next()) |x| { + total += x.value_ptr.* / 4; + try rows.append(gpa, .{ .name = x.key_ptr.*, .ms = x.value_ptr.* / 4 }); + } + std.mem.sort(Row, rows.items, {}, struct { + fn lt(_: void, a: Row, b: Row) bool { + return a.ms > b.ms; + } + }.lt); + if (rank == 0) { + std.debug.print("{d} stream(s) x {d} rows: graph step {d:.2} ms; eager launches sum {d:.2} ms\n", .{ nstreams, width, best, total }); + for (rows.items[0..@min(rows.items.len, 22)]) |x| std.debug.print(" {d:7.3} ms {d:5.1}% {s}\n", .{ x.ms, 100 * x.ms / total, x.name[0..@min(x.name.len, 90)] }); + } + } +} diff --git a/zig/tests/cuda/glm_weights.zig b/zig/tests/cuda/glm_weights.zig new file mode 100644 index 0000000000..ad1c372306 --- /dev/null +++ b/zig/tests/cuda/glm_weights.zig @@ -0,0 +1,78 @@ +//! GLM-5.3-Flash on CUDA: the native weight loader (src/families/glm_cuda/weights.zig) against the Python engine's +//! prepared weights (tools/capture_ref.py's weights.json: each tensor's name, bytes and sha256). +//! usage: tf-cuda-test glm-weights + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +/// Python tensors the native forward does not read (the latent path attends through absorb.wk/wv, the same bytes). +fn unused(name: []const u8) bool { + return std.mem.endsWith(u8, name, ".kv_k.weight") or std.mem.endsWith(u8, name, ".kv_v.weight") or std.mem.startsWith(u8, name, "draft_head."); +} + +pub fn run(gpu: Gpu, dir: []const u8, rank_s: []const u8, ref: []const u8) !void { + const gw = @import("glm_weights"); + const rank = try std.fmt.parseInt(u8, rank_s, 10); + const gpa = gpu.gpa; + const t0 = std.Io.Timestamp.now(gpu.io, .awake); + const planned = try gw.deviceBytes(gpu.io, dir); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + const ctx: gw.api.Ctx = .{ .d = gpu.d, .ctx = gpu.ctx, .stream = stream, .nccl = undefined, .comm = undefined, .rank = rank }; + var store = try gw.load(gpa, gpu.io, &ctx, dir, rank); + defer store.deinit(); + const secs = @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(gpu.io, .awake)).toNanoseconds())) / 1e9; + std.debug.print("loaded {d} tensors, {d:.2} GiB (planned {d:.2} GiB) in {d:.1} s\n", .{ store.map.count(), + @as(f64, @floatFromInt(store.ints.get("device_bytes").?)) / (1 << 30), @as(f64, @floatFromInt(planned)) / (1 << 30), secs }); + const jp = try std.fs.path.join(gpa, &.{ ref, "weights.json" }); + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, jp, gpa, .limited(1 << 27)); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + var equal: usize = 0; + var bad: usize = 0; + var skipped: usize = 0; + var seen = std.StringHashMap(void).init(gpa); + var have: std.ArrayList(u8) = .empty; + defer have.deinit(gpa); + for (parsed.value.object.get("tensors").?.array.items) |tv| { + const t = tv.object; + const sha = t.get("sha256") orelse continue; // config values and host tensors + const name = t.get("name").?.string; + try seen.put(name, {}); + const n: usize = @intCast(t.get("bytes").?.integer); + const got = store.map.get(name) orelse { + if (unused(name)) { + skipped += 1; + continue; + } + bad += 1; + std.debug.print("MISSING {s} ({d} bytes)\n", .{ name, n }); + continue; + }; + if (got.bytes() != n) { + bad += 1; + std.debug.print("SIZE {s}: {d} bytes, Python {d}\n", .{ name, got.bytes(), n }); + continue; + } + try have.resize(gpa, n); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(have.items.ptr, got.ptr, n), "download"); + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(have.items, &digest, .{}); + const hex = std.fmt.bytesToHex(digest, .lower); + if (std.mem.eql(u8, &hex, sha.string)) { + equal += 1; + } else { + bad += 1; + std.debug.print("DIFFERS {s} {s} [{d} bytes]\n", .{ name, @tagName(got.dtype), n }); + } + } + var extra: usize = 0; + var it = store.map.keyIterator(); + while (it.next()) |k| if (!seen.contains(k.*)) { + extra += 1; + std.debug.print("EXTRA {s} (not in Python's tree)\n", .{k.*}); + }; + std.debug.print("{d} equal, {d} not, {d} unused skipped, {d} extra\n", .{ equal, bad, skipped, extra }); + try check.expect(bad == 0 and extra == 0, "{d} tensors equal, {d} not, {d} extra", .{ equal, bad, extra }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aec6b9ec1b..5ca371044f 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -88,6 +88,7 @@ pub fn main(init: std.process.Init) !u8 { run(gpu, cmd, rest) catch |e| { std.debug.print("FAIL {s}: {t}\n", .{ cmd, e }); + if (@errorReturnTrace()) |t| std.debug.dumpStackTrace(t); return 1; }; return 0; @@ -146,6 +147,21 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { return carveout_tests.run(gpu, if (rest.len > 1) rest[1] else cuda.carveout.default_card, mib); } if (std.mem.eql(u8, cmd, "train-pack")) return train_pack.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 24); + if (std.mem.eql(u8, cmd, "glm-kcheck")) return @import("glm_forward.zig").kcheck(gpu, rest); + if (std.mem.eql(u8, cmd, "glm-tf")) return @import("glm_forward.zig").teacher(gpu, rest); + if (std.mem.eql(u8, cmd, "glm-prof")) return @import("glm_forward.zig").prof(gpu, rest); + if (std.mem.eql(u8, cmd, "glm-batch")) return @import("glm_forward.zig").batch(gpu, rest); + if (std.mem.eql(u8, cmd, "glm-weights")) return @import("glm_weights.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); + if (std.mem.eql(u8, cmd, "fn-gen-compiled")) return fn_run.generate(gpu, rest, .compiled); + if (std.mem.eql(u8, cmd, "fn-mtp-profile")) return @import("fn_mtp.zig").generateP(gpu, rest, true, true, true); + if (std.mem.eql(u8, cmd, "fn-mtp-parallel")) { + @import("fn_mtp.zig").parallel = true; + return @import("fn_mtp.zig").generate(gpu, rest, true, true); + } + if (std.mem.eql(u8, cmd, "fn-mtp-host")) return @import("fn_mtp.zig").generate(gpu, rest, true, false); + if (std.mem.eql(u8, cmd, "fn-mtp-eager")) return @import("fn_mtp.zig").generate(gpu, rest, false, true); + if (std.mem.eql(u8, cmd, "fn-gen-profile")) return fn_run.generate(gpu, rest, .profile); + if (std.mem.eql(u8, cmd, "fn-gen-interp")) return fn_run.generate(gpu, rest, .interp); std.debug.print("{s}", .{usage}); return error.UnknownCommand; } From a76f16aef4f6eda400bcd5359b84f61e07349570 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 19:51:25 +0000 Subject: [PATCH 04/31] glm kernels: the tensor-core fragment header beside them (qmm_frag.cuh), no include across families --- zig/kernels/cuda/glm/index.cu | 2 +- zig/kernels/cuda/glm/latent.cu | 2 +- zig/kernels/cuda/glm/mm.cu | 2 +- zig/kernels/cuda/glm/qmm_frag.cuh | 119 ++++++++++++++++++++++++++++++ 4 files changed, 122 insertions(+), 3 deletions(-) create mode 100644 zig/kernels/cuda/glm/qmm_frag.cuh diff --git a/zig/kernels/cuda/glm/index.cu b/zig/kernels/cuda/glm/index.cu index d647fed2b4..31986f8705 100644 --- a/zig/kernels/cuda/glm/index.cu +++ b/zig/kernels/cuda/glm/index.cu @@ -6,7 +6,7 @@ #include #include -#include "../flashnext/qmm_frag.cuh" +#include "qmm_frag.cuh" namespace { diff --git a/zig/kernels/cuda/glm/latent.cu b/zig/kernels/cuda/glm/latent.cu index 5ce154b389..b9d68ef37d 100644 --- a/zig/kernels/cuda/glm/latent.cu +++ b/zig/kernels/cuda/glm/latent.cu @@ -7,7 +7,7 @@ #include #include -#include "../flashnext/qmm_frag.cuh" +#include "qmm_frag.cuh" namespace { diff --git a/zig/kernels/cuda/glm/mm.cu b/zig/kernels/cuda/glm/mm.cu index 2afb66d54f..d3a9f718c5 100644 --- a/zig/kernels/cuda/glm/mm.cu +++ b/zig/kernels/cuda/glm/mm.cu @@ -12,7 +12,7 @@ #include #include -#include "../flashnext/qmm_frag.cuh" +#include "qmm_frag.cuh" namespace { diff --git a/zig/kernels/cuda/glm/qmm_frag.cuh b/zig/kernels/cuda/glm/qmm_frag.cuh new file mode 100644 index 0000000000..f4a4f457a7 --- /dev/null +++ b/zig/kernels/cuda/glm/qmm_frag.cuh @@ -0,0 +1,119 @@ +// Device code of src/tensorfold/cuda/kernels/qmm_frag.cuh (generated by zig/tests/cuda/flashnext/gen_copies.py; do not edit). +// Shared by the 4-bit matmuls: cp.async staging, ldmatrix, the bf16 mma, nibble decoding from B-fragment order. +#pragma once + +#include +#include +#include + +namespace qmm_frag { + +// Lane matmul tile shapes: BM rows by BN columns a block, WM x WN warps, each warp (BM / WM) x (BN / WN). +template +struct LaneTile { + static constexpr int THREADS = WM * WN * 32; + static constexpr int MT = BM / WM / 16; // m16 tiles a warp + static constexpr int NT = BN / WN / 8; // n8 tiles a warp + static constexpr int ROW = GS * 2; // bytes of one input row a group + static constexpr int CHUNKS = ROW / 16; + static constexpr int X = BM * ROW; // stage bytes: inputs, + static constexpr int W = BN * GS / 2; // weights, + static constexpr int S = BN * 2; // scales, biases (bf16), + static constexpr int XS = BM * 4; // and input sums (fp32) + static constexpr int STAGE = X + W + 2 * S + XS; + static constexpr int PARTIALS = MT * NT * 4 * THREADS * 4; // a K slice's partial, parked for the cluster sum + static constexpr int SMEM = STAGES * STAGE > PARTIALS ? STAGES * STAGE : PARTIALS; +}; + +__device__ __forceinline__ uint32_t smem(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void cp16(void* dst, const void* src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(smem(dst)), "l"(src)); +} + +// 16 bytes, or zeros without reading when ``read`` is false (a row past M: computed, never stored) +__device__ __forceinline__ void cp16z(void* dst, const void* src, bool read) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(smem(dst)), "l"(src), "r"(read ? 16 : 0)); +} + +__device__ __forceinline__ void cp8(void* dst, const void* src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" ::"r"(smem(dst)), "l"(src)); +} + +__device__ __forceinline__ void cp4(void* dst, const void* src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"(smem(dst)), "l"(src)); +} + +__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); } + +template +__device__ __forceinline__ void wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ void ldmatrix4(uint32_t (&r)[4], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(smem(p))); +} + +__device__ __forceinline__ void ldmatrix2(uint32_t (&r)[2], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(smem(p))); +} + +__device__ __forceinline__ void mma(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + asm( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +// The same mma from a zero accumulator (a group's first k16 step): no registers to clear. +__device__ __forceinline__ void mma0(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + const float z = 0.0f; + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%10, %10, %10, %10};\n" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1), "f"(z)); +} + +// Nibbles at bits [s, s + 4) and [16 + s, 20 + s) as a bf16 pair: exponent bits make 128 + q, then minus 128. +__device__ __forceinline__ uint32_t pair(uint32_t w, int s) { + const uint32_t t = ((w >> s) & 0x000F000Fu) | 0x43004300u; + uint32_t r; +#if __CUDA_ARCH__ >= 900 + asm("sub.rn.bf16x2 %0, %1, %2;\n" : "=r"(r) : "r"(t), "r"(0x43004300u)); +#else + // no sub.bf16x2 before sm_90: t * 1 - 128 rounds the exact q once too (an exact value, the same bits) + asm("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(r) : "r"(t), "r"(0x3F803F80u), "r"(0xC300C300u)); +#endif + return r; +} + +// Programmatic dependent launch (sm_90+; no-ops before, and when the launch did not ask for it): wait for the +// previous kernel's writes, and let the next kernel's blocks start once every block here has said so. +__device__ __forceinline__ void grid_wait() { +#if __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;\n" ::: "memory"); +#endif +} + +__device__ __forceinline__ void grid_launch() { +#if __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;\n" ::: "memory"); +#endif +} + +// Block b's (first row, first column), row tiles fastest in bands of ``group`` so blocks in flight share L2. +__device__ __forceinline__ int2 tile_of(int b, int M, int N, int BM, int BN, int group) { + const int rows_t = (M + BM - 1) / BM, cols_t = (N + BN - 1) / BN, band = group * cols_t; + const int first = b / band * group, in_band = b % band, height = min(group, rows_t - first); + return make_int2((first + in_band % height) * BM, in_band / height * BN); +} + +// A 16-byte chunk c of row r sits at chunk c ^ (r % CHUNKS) so ldmatrix's eight rows hit eight bank groups. +template +__device__ __forceinline__ int swz(int r, int c) { return c ^ (r % CHUNKS); } + +} // namespace qmm_frag From abb8bbe6efe73a411f75f555d2de590bbbe0dc32 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 19:54:30 +0000 Subject: [PATCH 05/31] glm: the three Flash Next kernels the family loads (experts, experts_prefill, sample) and no Flash Next test commands --- zig/build/cuda.zig | 4 + zig/kernels/cuda/flashnext/experts.cu | 294 ++++++++++++++++++ zig/kernels/cuda/flashnext/experts.cuh | 106 +++++++ zig/kernels/cuda/flashnext/experts_prefill.cu | 162 ++++++++++ zig/kernels/cuda/flashnext/sample.cu | 254 +++++++++++++++ zig/src/cuda/kernels.zig | 3 + zig/tests/cuda/main.zig | 10 - 7 files changed, 823 insertions(+), 10 deletions(-) create mode 100644 zig/kernels/cuda/flashnext/experts.cu create mode 100644 zig/kernels/cuda/flashnext/experts.cuh create mode 100644 zig/kernels/cuda/flashnext/experts_prefill.cu create mode 100644 zig/kernels/cuda/flashnext/sample.cu diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 9b4d4348a8..212dde8885 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -43,6 +43,10 @@ const kernels = [_]Kernel{ .{ .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 + // the routed experts' grouped kernels and the greedy candidates the GLM family shares with Flash Next (zig/kernels/cuda/flashnext) + .{ .name = "fn_experts", .src = "flashnext/experts", .flags = &.{"-O3"} }, + .{ .name = "fn_experts_prefill", .src = "flashnext/experts_prefill", .flags = &.{"-O3"} }, + .{ .name = "fn_sample", .src = "flashnext/sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // GLM-5.3-Flash (glm5_next): the device copy of the Python engine's KDA extension (zig/tests/cuda/glm/gen_copies.py) .{ .name = "glm_kda", .src = "glm/kda", .flags = &.{ "-O3", "--fmad=false" } }, // glm5_next/cuda/kda.py, tensorfold_glm_kda_v2 .{ .name = "glm_ops", .src = "glm/ops", .flags = &.{ "-O3", "--fmad=false" } }, // ours: the token lists, casts and positions torch ops make diff --git a/zig/kernels/cuda/flashnext/experts.cu b/zig/kernels/cuda/flashnext/experts.cu new file mode 100644 index 0000000000..ea680620e4 --- /dev/null +++ b/zig/kernels/cuda/flashnext/experts.cu @@ -0,0 +1,294 @@ +// Device code of src/tensorfold/cuda/experts.cu (generated by zig/tests/cuda/flashnext/gen_copies.py; do not edit). +// Decode form: acc = fma(xs, b, fma(p, s, acc)) per group in order, one pair an mma row, so no pair affects another. + +#include +#include + +#include "experts.cuh" + +namespace { + +__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)); } + +// a row's group sum: the lane's inputs in order, then the quad (every lane of the quad gets the same bits) +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 +struct Stage { + uint4 w[M][Geo::WV]; + uint4 sb[M][Geo::SBV]; + uint4 xa[Geo::XV]; // row gq of the tile: the lane's inputs of the group + uint4 xb[Geo::XV]; // row gq + 8 +}; + +template +__device__ __forceinline__ void load_stage(Stage& st, const uint4* blk, int g, int lane, int t, + const __nv_bfloat16* x0, const __nv_bfloat16* x1, bool v0, bool v1) { + using G = Geo; + const uint4* b = blk + (size_t)g * (M * G::BLOCK); +#pragma unroll + for (int m = 0; m < M; ++m) { +#pragma unroll + for (int c = 0; c < G::WV; ++c) st.w[m][c] = ld_w(b + m * G::BLOCK + c * 32 + lane); +#pragma unroll + for (int c = 0; c < G::SBV; ++c) st.sb[m][c] = ld_w(b + m * G::BLOCK + 32 * G::WV + t * G::SBV + c); + } + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + const int k0 = g * GS; +#pragma unroll + for (int c = 0; c < G::XV; ++c) { + st.xa[c] = v0 ? ld_x(x0 + k0 + 8 * c) : zero; + st.xb[c] = v1 ? ld_x(x1 + k0 + 8 * c) : zero; + } +} + +template +__device__ __forceinline__ void compute_stage(float (&acc)[M][1][NTW][4], const Stage& st, bool hi) { + using G = Geo; + const float xsa = group_sum(st.xa); + const float xsb = group_sum(st.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) { + // k-step ks: the lane's inputs 4 ks .. 4 ks + 3 (pairs at the mma's k positions 2t and 2t + 8) + const uint32_t a0 = comp(st.xa[ks >> 1], 2 * (ks & 1)); + const uint32_t a2 = comp(st.xa[ks >> 1], 2 * (ks & 1) + 1); + const uint32_t a1 = comp(st.xb[ks >> 1], 2 * (ks & 1)); + const uint32_t a3 = comp(st.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(st.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(st.sb[m][j >> 2], j & 3); + const uint32_t bp = comp(st.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][0][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 (hi) { + a[2] = fmaf(xsb, b0, fmaf(p[j][2], s0, a[2])); + a[3] = fmaf(xsb, b1, fmaf(p[j][3], s1, a[3])); + } + } + } +} + +template +__device__ __forceinline__ void k_loop(float (&acc)[M][1][NTW][4], const uint4* blk, int KG, int lane, int t, + const __nv_bfloat16* x0, const __nv_bfloat16* x1, bool v0, bool v1) { + Stage st[D]; +#pragma unroll + for (int d = 0; d < D; ++d) + if (d < KG) load_stage(st[d], blk, d, lane, t, x0, x1, v0, v1); + for (int g0 = 0; g0 < KG; g0 += D) { +#pragma unroll + for (int d = 0; d < D; ++d) { + const int g = g0 + d; + if (g < KG) { + compute_stage(acc, st[d], v1); + if (g + D < KG) load_stage(st[d], blk, g + D, lane, t, x0, x1, v0, v1); + } + } + } +} + +// Pair p reads X row p / slots when slots > 0 (X holds tokens), else X row p (X holds a row a pair). +template +__global__ void __launch_bounds__(WARPS * 32) + expert_kernel(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) { + 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); + const bool v0 = gq < cnt, v1 = gq + 8 < cnt; + const int pr0 = v0 ? __ldg(members + first + gq) : 0; + const int pr1 = v1 ? __ldg(members + first + gq + 8) : 0; + const int r0 = slots ? pr0 / slots : pr0, r1 = slots ? pr1 / slots : pr1; + const __nv_bfloat16* x0 = X + (size_t)r0 * x_stride + t * (GS / 4); + const __nv_bfloat16* x1 = X + (size_t)r1 * x_stride + t * (GS / 4); + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (M * Geo::BLOCK); + float acc[M][1][NTW][4]; +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[m][0][j][0] = acc[m][0][j][1] = acc[m][0][j][2] = acc[m][0][j][3] = 0.f; + k_loop(acc, blk, KG, lane, t, x0, x1, v0, v1); + epilogue(acc, 0, out, N, cb * COLS + 2 * t, pr0, pr1, v0, v1, limit); + } +} + +constexpr int PLAN_THREADS = 1024; +constexpr int EMAX = 1024; +constexpr int PSMALL = 1024; // pairs the one-block plan takes; wider calls rank in blocks of PLAN_THREADS + +// Whole block: thread e gives expert e's pair count, gets its first member; items of <= T pairs land in expert order. +__device__ int place_items(int c, int E, int T, int* __restrict__ items, int* __restrict__ counts) { + __shared__ int wsum[3][32]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int tiles = (c + T - 1) / T; + int a = c, b = tiles, u = c > 0; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int ua = __shfl_up_sync(0xffffffffu, a, o), ub = __shfl_up_sync(0xffffffffu, b, o); + const int uu = __shfl_up_sync(0xffffffffu, u, o); + if (lane >= o) { + a += ua; + b += ub; + u += uu; + } + } + if (lane == 31) { + wsum[0][warp] = a; + wsum[1][warp] = b; + wsum[2][warp] = u; + } + __syncthreads(); + if (warp == 0) { +#pragma unroll + for (int s = 0; s < 3; ++s) { + int v = wsum[s][lane]; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int w = __shfl_up_sync(0xffffffffu, v, o); + if (lane >= o) v += w; + } + wsum[s][lane] = v; + } + } + __syncthreads(); + const int off = (warp ? wsum[0][warp - 1] : 0) + a - c; + const int ioff = (warp ? wsum[1][warp - 1] : 0) + b - tiles; + if (tid == 0) { + counts[0] = wsum[1][31]; + counts[1] = wsum[2][31]; + } + if (tid < E) + for (int j = 0; j < tiles; ++j) { + int* it = items + 3 * (ioff + j); + it[0] = tid; + it[1] = off + T * j; + it[2] = min(T, c - T * j); + } + return off; +} + +// One block: pairs p = row * slots + slot grouped by expert (pair order within an expert). +__global__ void __launch_bounds__(PLAN_THREADS) + plan_kernel(const int* __restrict__ picks, int P, int E, int T, int* __restrict__ members, + int* __restrict__ items, int* __restrict__ counts) { + __shared__ int pk[PSMALL]; + __shared__ int cnt[EMAX]; + __shared__ int off[EMAX]; + const int tid = threadIdx.x; + for (int e = tid; e < EMAX; e += PLAN_THREADS) cnt[e] = 0; + __syncthreads(); + for (int p = tid; p < P; p += PLAN_THREADS) { + const int e = picks[p]; + pk[p] = e; + if (e < E) atomicAdd(&cnt[e], 1); // a pick past the table (e >= E) is left out of the plan + } + __syncthreads(); + const int o = place_items(tid < E ? cnt[tid] : 0, E, T, items, counts); + if (tid < E) off[tid] = o; + __syncthreads(); + for (int p = tid; p < P; p += PLAN_THREADS) { + const int e = pk[p]; + if (e >= E) continue; + int rank = 0; + for (int q = 0; q < p; ++q) rank += pk[q] == e; + members[off[e] + rank] = p; + } +} + +// Wide calls, pass 1: block b ranks its pairs within each expert in pair order and writes its counts to hist[b][e]. +__global__ void __launch_bounds__(PLAN_THREADS) + plan_rank(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, int* __restrict__ hist) { + __shared__ int cnt[EMAX]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + for (int e = tid; e < E; e += PLAN_THREADS) cnt[e] = 0; + __syncthreads(); + const int p = blockIdx.x * PLAN_THREADS + tid; + const bool ok = p < P && picks[p] < E; // a pick past the table is left out of the plan + const int e = ok ? picks[p] : -1 - lane; + const unsigned same = __match_any_sync(0xffffffffu, e); + const int below = __popc(same & ((1u << lane) - 1u)); + for (int w = 0; w < PLAN_THREADS / 32; ++w) { + if (warp == w) { + const int base = ok ? cnt[e] : 0; + __syncwarp(); + if (ok && below == 0) cnt[e] = base + __popc(same); + if (ok) rank[p] = base + below; + } + __syncthreads(); + } + for (int x = tid; x < E; x += PLAN_THREADS) hist[(size_t)blockIdx.x * E + x] = cnt[x]; +} + +// Pass 2, one block: each block's first member for each expert in place of its count, and the items. +__global__ void __launch_bounds__(PLAN_THREADS) + plan_offsets(int nblk, int E, int T, int* __restrict__ hist, int* __restrict__ items, int* __restrict__ counts) { + const int tid = threadIdx.x; + int c = 0; + if (tid < E) + for (int b = 0; b < nblk; ++b) { + int* h = hist + (size_t)b * E + tid; + const int v = *h; + *h = c; + c += v; + } + const int off = place_items(c, E, T, items, counts); + if (tid < E) + for (int b = 0; b < nblk; ++b) hist[(size_t)b * E + tid] += off; +} + +// Pass 3: every pair to its place. +__global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const int* __restrict__ rank, + const int* __restrict__ hist, int* __restrict__ members) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p < P && picks[p] < E) members[hist[(size_t)(p / PLAN_THREADS) * E + picks[p]] + rank[p]] = p; +} +} // namespace + +// The specializations the Python build instantiated. +namespace { +template __global__ void expert_kernel<(int)32, (int)1, (int)0, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void expert_kernel<(int)32, (int)2, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void expert_kernel<(int)64, (int)1, (int)0, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void expert_kernel<(int)64, (int)1, (int)1, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void expert_kernel<(int)64, (int)2, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +} // namespace diff --git a/zig/kernels/cuda/flashnext/experts.cuh b/zig/kernels/cuda/flashnext/experts.cuh new file mode 100644 index 0000000000..f7fac4861d --- /dev/null +++ b/zig/kernels/cuda/flashnext/experts.cuh @@ -0,0 +1,106 @@ +// Device code of src/tensorfold/cuda/experts.cuh (generated by zig/tests/cuda/flashnext/gen_copies.py; do not edit). +// Shared by the grouped-expert kernels (experts.cu: decode form and plan; experts_prefill.cu: prefill form). +#pragma once + +#include +#include + +namespace { + +constexpr int NTW = 4; // n8 tiles a warp: 32 output columns +constexpr int COLS = 8 * NTW; +constexpr uint32_t LOW = 0x000F000Fu; +constexpr uint32_t K128 = 0x43004300u; // the bf16 pair (128, 128) + +template +struct Geo { + static constexpr int KS = GS / 16; // mma k-steps a group + static constexpr int WV = NTW * GS / 128; // weight uint4 a lane a group + static constexpr int XV = GS / 32; // input uint4 a row a lane a group + static constexpr int SBV = NTW / 2; // scale and bias uint4 a quad a group + static constexpr int BLOCK = 32 * WV + 4 * SBV; +}; + +__device__ __forceinline__ void cp16(void* dst, const void* src) { + const uint32_t d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(d), "l"(src)); +} + +__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n" ::); } + +template +__device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__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; +} + +// nibbles at bits sh and sh + 16 -> the exact bf16 pair (q0, q1): (128 + q) - 128 +__device__ __forceinline__ uint32_t nib2(uint32_t w, int sh) { + uint32_t v = ((w >> sh) & LOW) | K128; + const uint32_t k = K128; + __nv_bfloat162 r = __hsub2(*reinterpret_cast<__nv_bfloat162*>(&v), *reinterpret_cast(&k)); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ void mma(float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, + uint32_t b1) { + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ float lo_f(uint32_t v) { return __uint_as_float(v << 16); } +__device__ __forceinline__ float hi_f(uint32_t v) { return __uint_as_float(v & 0xFFFF0000u); } +__device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +// SwiGLU as the families define it: bf16(bf16(silu(g)) * u) on bf16 g and u, clipped first when limit > 0 +__device__ __forceinline__ float swiglu(float g, float u, float limit) { + float gv = bf(g), uv = bf(u); + if (limit > 0.f) { + gv = fminf(gv, limit); + uv = fminf(fmaxf(uv, -limit), limit); + } + return bf(gv / (1.f + expf(-gv))) * uv; +} + +__device__ __forceinline__ float relu2(float a) { + const float u = fmaxf(bf(a), 0.f); + return u * u; +} + +// EPI 0: fp32 out (down); 1: bf16 relu(bf16(acc))^2; 2: bf16 SwiGLU(matrix 0, matrix 1); 3: bf16 out (prefill down) +template +__device__ __forceinline__ void epilogue(const float (&acc)[M][RT][NTW][4], int r, void* out, int N, int col0, int pr0, + int pr1, bool v0, bool v1, float limit) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + if (!(h ? v1 : v0)) continue; + const size_t row = (size_t)(h ? pr1 : pr0) * N; +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const int col = col0 + 8 * j; + const float a0 = acc[0][r][j][2 * h], a1 = acc[0][r][j][2 * h + 1]; + if constexpr (EPI == 0) { + *reinterpret_cast(reinterpret_cast(out) + row + col) = make_float2(a0, a1); + } else { + float o0, o1; + if constexpr (EPI == 1) { + o0 = relu2(a0); + o1 = relu2(a1); + } else if constexpr (EPI == 2) { + o0 = swiglu(a0, acc[M - 1][r][j][2 * h], limit); + o1 = swiglu(a1, acc[M - 1][r][j][2 * h + 1], limit); + } else { + o0 = a0; + o1 = a1; + } + *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<__nv_bfloat16*>(out) + row + col) = + __floats2bfloat162_rn(o0, o1); + } + } + } +} + +} // namespace diff --git a/zig/kernels/cuda/flashnext/experts_prefill.cu b/zig/kernels/cuda/flashnext/experts_prefill.cu new file mode 100644 index 0000000000..6cab13ddf8 --- /dev/null +++ b/zig/kernels/cuda/flashnext/experts_prefill.cu @@ -0,0 +1,162 @@ +// Device code of src/tensorfold/cuda/experts_prefill.cu (generated by zig/tests/cuda/flashnext/gen_copies.py; do not edit). +// Prefill: w = bf16(fma(q, s, b)), one fp32 chain over K; chunk-invariant bits, not decode's (replies prefill again). + +#include +#include + +#include "experts.cuh" + +namespace { + +template +struct Pre { + static constexpr int THREADS = WM * WN * 32, BM = 16 * RT * WM, SG = 64 / GS, XC = 8, WB = M * Geo::BLOCK; + static constexpr int XU = BM * XC, WU = WN * SG * WB, SU = XU + WU; // uint4 a stage: X rows, weight blocks + static constexpr int XPT = (XU + THREADS - 1) / THREADS; + static constexpr int STAGE_BYTES = SU * 16; + static constexpr int STAGES = STAGE_BYTES * 4 <= 49152 ? 4 : STAGE_BYTES * 3 <= 49152 ? 3 : 2; + static __device__ __forceinline__ int xslot(int r, int c) { + return r * XC + (GS == 32 ? c ^ ((r & 1) << 2) : c ^ (r & 1)); + } +}; + +__device__ __forceinline__ uint32_t hfma2(uint32_t q, uint32_t s, uint32_t b) { + __nv_bfloat162 r = __hfma2(*reinterpret_cast<__nv_bfloat162*>(&q), *reinterpret_cast<__nv_bfloat162*>(&s), + *reinterpret_cast<__nv_bfloat162*>(&b)); + return *reinterpret_cast(&r); +} + +template +__global__ void __launch_bounds__(WM * WN * 32) + prefill_kernel(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; + using P = Pre; + constexpr int STAGES = P::STAGES; + extern __shared__ uint4 sm[]; + const int nbt = (NB + WN - 1) / WN, it = blockIdx.x / nbt, cbt = blockIdx.x - it * nbt; + if (it >= __ldg(counts)) return; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, wm = warp / WN, wn = warp - wm * WN; + const int t = lane & 3, gq = lane >> 2, cb = cbt * WN + wn, cbs = min(WN, NB - cbt * WN); + const __nv_bfloat16* xsrc[P::XPT]; + int xdst[P::XPT]; +#pragma unroll + for (int i = 0; i < P::XPT; ++i) { + const int q = tid + i * P::THREADS, r = q / P::XC, c = q - r * P::XC; + xdst[i] = P::xslot(r, c); + xsrc[i] = nullptr; + if (q < P::XU && r < cnt) { + const int p = __ldg(members + first + r); + xsrc[i] = X + (size_t)(slots ? p / slots : p) * x_stride + 8 * c; + } + } + const uint4* wsrc = W + ((size_t)e * NB + cbt * WN) * (size_t)KG * P::WB; + auto stage = [&](int s, int g0) { // groups g0 .. g0 + ng - 1 (ng < SG only at an odd end) + const int ng = min(P::SG, KG - g0); + uint4* xs = sm + s * P::SU; +#pragma unroll + for (int i = 0; i < P::XPT; ++i) + if (xsrc[i] && ((tid + i * P::THREADS) % P::XC) < ng * (GS / 8)) cp16(xs + xdst[i], xsrc[i] + (size_t)g0 * GS); + uint4* ws = xs + P::XU; + for (int q = tid; q < cbs * P::SG * P::WB; q += P::THREADS) { + const int j = q / (P::SG * P::WB), o = q - j * P::SG * P::WB; + if (o < ng * P::WB) cp16(ws + q, wsrc + ((size_t)j * KG + g0) * P::WB + o); + } + }; + const int base = 16 * RT * wm; + const int nt = max(0, min(RT, (cnt - base + 15) >> 4)); + const bool active = nt > 0 && cb < NB; + 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 int steps = (KG + P::SG - 1) / P::SG; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < steps) stage(s, s * P::SG); + cp_commit(); + } + const uint32_t pick = (gq & 1) ? 0x3232u : 0x1010u; // this lane's column: the high or low half of a pair + for (int k = 0; k < steps; ++k) { + cp_wait(); + __syncthreads(); + if (k + STAGES - 1 < steps) stage((k + STAGES - 1) % STAGES, (k + STAGES - 1) * P::SG); + cp_commit(); + if (!active) continue; + const uint4* xs = sm + (k % STAGES) * P::SU; +#pragma unroll + for (int gi = 0; gi < P::SG; ++gi) { + if (k * P::SG + gi >= KG) break; + const uint4* ws = xs + P::XU + (wn * P::SG + gi) * P::WB; + uint4 wv[M][G::WV]; + uint32_t s2[M][NTW], b2[M][NTW]; +#pragma unroll + for (int m = 0; m < M; ++m) { +#pragma unroll + for (int c = 0; c < G::WV; ++c) wv[m][c] = ws[m * G::BLOCK + c * 32 + lane]; + const uint4 sp = ws[m * G::BLOCK + 32 * G::WV + (gq >> 1) * G::SBV]; + const uint4 bp = ws[m * G::BLOCK + 32 * G::WV + (gq >> 1) * G::SBV + 1]; +#pragma unroll + for (int j = 0; j < NTW; ++j) { + s2[m][j] = __byte_perm(comp(sp, j), 0u, pick); + b2[m][j] = __byte_perm(comp(bp, j), 0u, pick); + } + } +#pragma unroll + for (int kk = 0; kk < G::KS / 2; ++kk) { + uint4 xa[RT], xb[RT]; +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + const int r0 = base + 16 * r + gq; + xa[r] = xs[P::xslot(r0, gi * (GS / 8) + t * G::XV + kk)]; + xb[r] = xs[P::xslot(r0 + 8, gi * (GS / 8) + t * G::XV + kk)]; + } +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const int wi = j * (GS / 32) + kk; + const uint32_t word = comp(wv[m][wi >> 2], wi & 3); +#pragma unroll + for (int h = 0; h < 2; ++h) { + const uint32_t b0 = hfma2(nib2(word, 8 * h), s2[m][j], b2[m][j]); + const uint32_t b1 = hfma2(nib2(word, 8 * h + 4), s2[m][j], b2[m][j]); +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + mma(acc[m][r][j], comp(xa[r], 2 * h), comp(xb[r], 2 * h), comp(xa[r], 2 * h + 1), + comp(xb[r], 2 * h + 1), b0, b1); + } + } + } + } + } + } + if (!active) return; +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + const int m0 = base + 16 * r + gq, m1 = m0 + 8; + const bool v0 = m0 < cnt, v1 = m1 < cnt; + const int p0 = v0 ? __ldg(members + first + m0) : 0, p1 = v1 ? __ldg(members + first + m1) : 0; + epilogue(acc, r, out, N, cb * COLS + 2 * t, p0, p1, v0, v1, limit); + } +} +} // namespace + +// The specializations the Python build instantiated. +namespace { +template __global__ void prefill_kernel<(int)32, (int)1, (int)0, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)32, (int)1, (int)3, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)32, (int)2, (int)2, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)64, (int)1, (int)0, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)64, (int)1, (int)1, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)64, (int)1, (int)3, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +template __global__ void prefill_kernel<(int)64, (int)2, (int)2, (int)2, (int)2, (int)4>(const __nv_bfloat16 *, int, int, const uint4 *, int, int, const int *, const int *, const int *, void *, int, float); +} // namespace diff --git a/zig/kernels/cuda/flashnext/sample.cu b/zig/kernels/cuda/flashnext/sample.cu new file mode 100644 index 0000000000..fc23bbb4b4 --- /dev/null +++ b/zig/kernels/cuda/flashnext/sample.cu @@ -0,0 +1,254 @@ +// Flash Next port: device-side greedy candidates. One block per bf16 logits row: the row's first maximum (lowest +// column among equal maxima, as argmax picks it), mapped to a global token id, and the row's log-sum-exp in fp32. +// out[row] = {id, max (fp32 bits), lse (fp32 bits), 0}. Build: nvcc -cubin -arch=sm_121 -O3 -o fn_sample.cubin fn_sample.cu +#include +#include + +struct Top { + float m; // running maximum + int i; // first column holding it + float s; // sum of exp(x - m) +}; + +__device__ __forceinline__ Top merge(Top a, Top b) { + if (b.m > a.m || (b.m == a.m && b.i < a.i)) { Top t = a; a = b; b = t; } + // a holds the larger maximum (or the earlier column at a tie) + float s = a.s + (b.m == -CUDART_INF_F ? 0.f : b.s * __expf(b.m - a.m)); + return {a.m, a.i, s}; +} + +extern "C" __global__ void __launch_bounds__(1024) fn_rows_top(const __nv_bfloat16* __restrict__ logits, int cols, + int stride, const int* __restrict__ id_map, int offset, int* __restrict__ out) { + const int row = blockIdx.x; + const __nv_bfloat16* x = logits + (size_t)row * stride; + Top t = {-CUDART_INF_F, 0x7fffffff, 0.f}; + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + float v = __bfloat162float(x[c]); + if (v > t.m) { t.s = t.s * __expf(t.m - v) + 1.f; t.m = v; t.i = c; } + else t.s += __expf(v - t.m); + } + for (int o = 16; o > 0; o >>= 1) { + Top u = {__shfl_xor_sync(0xffffffff, t.m, o), __shfl_xor_sync(0xffffffff, t.i, o), __shfl_xor_sync(0xffffffff, t.s, o)}; + t = merge(t, u); + } + __shared__ Top warps[32]; + const int lane = threadIdx.x & 31, w = threadIdx.x >> 5; + if (lane == 0) warps[w] = t; + __syncthreads(); + if (w == 0) { + t = lane < (blockDim.x >> 5) ? warps[lane] : Top{-CUDART_INF_F, 0x7fffffff, 0.f}; + for (int o = 16; o > 0; o >>= 1) { + Top u = {__shfl_xor_sync(0xffffffff, t.m, o), __shfl_xor_sync(0xffffffff, t.i, o), __shfl_xor_sync(0xffffffff, t.s, o)}; + t = merge(t, u); + } + if (lane == 0) { + int* o4 = out + 4 * row; + o4[0] = id_map ? id_map[t.i] : t.i + offset; + o4[1] = __float_as_int(t.m); + o4[2] = __float_as_int(t.m + __logf(t.s)); + o4[3] = 0; + } + } +} + +// Sampling candidates: each row's 64 largest logits (by value, then the lower column at a tie), as global token ids +// and fp32 values, the row's log-sum-exp, and its log-sum-exp at the row's temperature (inv_t[row], x * inv_t); both +// ranks' sets are gathered and the keyed draw runs on the host over the merged 128 (whose first 64 are the vocabulary's +// top 64 exactly). Each thread keeps its columns' sorted top 64, then the block merges the threads' lists by repeated +// argmax over their heads (exact: no value of the row's top 64 can be missing from its thread's list). +// out[row * 132]: ids[64], values[64] (fp32 bits), lse, lse at the temperature (fp32 bits), 2 spare. +#define TOPK 64 +#define TOPW (2 * TOPK + 4) +__device__ __forceinline__ bool above(float a, int ia, float b, int ib) { return a > b || (a == b && ia < ib); } + +__device__ __forceinline__ void lse_merge(float& m, float& s, float m2, float s2) { + const float mm = fmaxf(m, m2); + if (mm == -CUDART_INF_F) return; + s = (m == -CUDART_INF_F ? 0.f : s * __expf(m - mm)) + (m2 == -CUDART_INF_F ? 0.f : s2 * __expf(m2 - mm)); + m = mm; +} + +extern "C" __global__ void __launch_bounds__(256) fn_cand_topk(const __nv_bfloat16* __restrict__ logits, int cols, + int stride, int offset, const float* __restrict__ inv_t, int* __restrict__ out) { + const int row = blockIdx.x; + const __nv_bfloat16* x = logits + (size_t)row * stride; + const float it = inv_t ? inv_t[row] : 1.f; + float lv[TOPK]; + int li[TOPK]; + int n = 0; + float m = -CUDART_INF_F, s = 0.f, mt = -CUDART_INF_F, st = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + const float v = __bfloat162float(x[c]); + if (v > m) { s = s * __expf(m - v) + 1.f; m = v; } else s += __expf(v - m); + const float vt = v * it; + if (vt > mt) { st = st * __expf(mt - vt) + 1.f; mt = vt; } else st += __expf(vt - mt); + if (n < TOPK || above(v, c, lv[TOPK - 1], li[TOPK - 1])) { + int j = n < TOPK ? n++ : TOPK - 1; + while (j > 0 && above(v, c, lv[j - 1], li[j - 1])) { lv[j] = lv[j - 1]; li[j] = li[j - 1]; --j; } + lv[j] = v; + li[j] = c; + } + } + __shared__ float sm[8], ss[8], smt[8], sst[8]; + __shared__ float hv[256]; + __shared__ int hi[256], winner; + for (int o = 16; o > 0; o >>= 1) { + lse_merge(m, s, __shfl_xor_sync(0xffffffff, m, o), __shfl_xor_sync(0xffffffff, s, o)); + lse_merge(mt, st, __shfl_xor_sync(0xffffffff, mt, o), __shfl_xor_sync(0xffffffff, st, o)); + } + const int lane = threadIdx.x & 31, w = threadIdx.x >> 5, nw = blockDim.x >> 5; + if (lane == 0) { sm[w] = m; ss[w] = s; smt[w] = mt; sst[w] = st; } + __syncthreads(); + int* o = out + (size_t)row * TOPW; + if (threadIdx.x == 0) { + float M = -CUDART_INF_F, S = 0.f, MT = -CUDART_INF_F, ST = 0.f; + for (int i = 0; i < nw; ++i) { lse_merge(M, S, sm[i], ss[i]); lse_merge(MT, ST, smt[i], sst[i]); } + o[2 * TOPK] = __float_as_int(M + __logf(S)); + o[2 * TOPK + 1] = __float_as_int(MT + __logf(ST)); + } + int head = 0; + for (int k = 0; k < TOPK; ++k) { + hv[threadIdx.x] = head < n ? lv[head] : -CUDART_INF_F; + hi[threadIdx.x] = head < n ? li[head] : 0x7fffffff; + __syncthreads(); + if (threadIdx.x == 0) { + int b = 0; + for (int t = 1; t < (int)blockDim.x; ++t) if (above(hv[t], hi[t], hv[b], hi[b])) b = t; + winner = b; + o[k] = hi[b] == 0x7fffffff ? -1 : hi[b] + offset; + o[TOPK + k] = __float_as_int(hv[b]); + } + __syncthreads(); + if (threadIdx.x == winner) ++head; + __syncthreads(); + } +} + +// Structured output: each listed row's tokens its grammar does not allow to -inf (rows[i]: the logits row, bits: the +// row's allowed bits over this rank's columns, token c: bit c % 32 of word c / 32), before the picks and draws. +extern "C" __global__ void __launch_bounds__(256) fn_grammar_rows(__nv_bfloat16* __restrict__ logits, int stride, int cols, + const unsigned* __restrict__ bits, int words, const int* __restrict__ rows) { + __nv_bfloat16* x = logits + (size_t)rows[blockIdx.x] * stride; + const unsigned* w = bits + (size_t)blockIdx.x * words; + for (int c = threadIdx.x; c < cols; c += blockDim.x) + if (!((w[c >> 5] >> (c & 31)) & 1u)) x[c] = __ushort_as_bfloat16((unsigned short)0xff80); +} + +// fn_cand_topk's output, by radix selection over the bf16 bits instead of per-thread lists: a pass histograms the +// order-preserving keys' high byte (and takes both log-sum-exps), a second the low byte within the bin holding the +// 64th largest, which fixes the threshold key K; a third collects every column above K and, of the columns at K, the +// lowest ones (in column order), and the picks are ranked by value, then the lower column. The same 64 in the same +// order as fn_cand_topk (-0 counts as +0, as the float compare does); every pass reads the row once, coalesced. +__device__ __forceinline__ unsigned bf_key(unsigned short u) { + if (u == 0x8000u) u = 0; // -0: the same value as +0 + return (u & 0x8000u) ? (~u & 0xffffu) : (u | 0x8000u); +} + +__device__ __forceinline__ void hist_add(unsigned* hist, int bin, bool valid) { + const unsigned m = __match_any_sync(0xffffffffu, valid ? bin : -1); + if (valid && (threadIdx.x & 31) == __ffs(m) - 1) atomicAdd(&hist[bin], __popc(m)); +} + +extern "C" __global__ void __launch_bounds__(1024) fn_cand_topk_sel(const __nv_bfloat16* __restrict__ logits, int cols, + int stride, int offset, const float* __restrict__ inv_t, int* __restrict__ out) { + const int row = blockIdx.x; + const unsigned short* xu = reinterpret_cast(logits + (size_t)row * stride); + const float it = inv_t ? inv_t[row] : 1.f; + const int lane = threadIdx.x & 31, w = threadIdx.x >> 5, nw = blockDim.x >> 5; + __shared__ unsigned hist[256]; + __shared__ float sm[32], ss[32], smt[32], sst[32]; + __shared__ int s_bin, s_gt, s_cnt, s_eq_seen; + __shared__ int warp_eq[32]; + __shared__ float pv[TOPK]; + __shared__ int pi[TOPK]; + for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0; + if (threadIdx.x == 0) { s_cnt = 0; s_eq_seen = 0; } + __syncthreads(); + float m = -CUDART_INF_F, s = 0.f, mt = -CUDART_INF_F, st = 0.f; + for (int base = 0; base < cols; base += blockDim.x) { + const int c = base + threadIdx.x; + const bool valid = c < cols; + unsigned key = 0; + if (valid) { + const unsigned short u = xu[c]; + key = bf_key(u); + const float v = __uint_as_float((unsigned)u << 16); + if (v > m) { s = s * __expf(m - v) + 1.f; m = v; } else s += __expf(v - m); + const float vt = v * it; + if (vt > mt) { st = st * __expf(mt - vt) + 1.f; mt = vt; } else st += __expf(vt - mt); + } + hist_add(hist, (int)(key >> 8), valid); + } + for (int o = 16; o > 0; o >>= 1) { + lse_merge(m, s, __shfl_xor_sync(0xffffffff, m, o), __shfl_xor_sync(0xffffffff, s, o)); + lse_merge(mt, st, __shfl_xor_sync(0xffffffff, mt, o), __shfl_xor_sync(0xffffffff, st, o)); + } + if (lane == 0) { sm[w] = m; ss[w] = s; smt[w] = mt; sst[w] = st; } + __syncthreads(); + int* o = out + (size_t)row * TOPW; + const int want = cols < TOPK ? cols : TOPK; + if (threadIdx.x == 0) { + float M = -CUDART_INF_F, S = 0.f, MT = -CUDART_INF_F, ST = 0.f; + for (int i = 0; i < nw; ++i) { lse_merge(M, S, sm[i], ss[i]); lse_merge(MT, ST, smt[i], sst[i]); } + o[2 * TOPK] = __float_as_int(M + __logf(S)); + o[2 * TOPK + 1] = __float_as_int(MT + __logf(ST)); + int cum = 0, b = 255; + for (; b > 0; --b) { if (cum + (int)hist[b] >= want) break; cum += hist[b]; } + s_bin = b; + s_gt = cum; // columns whose key's high byte is above the bin + } + __syncthreads(); + const unsigned hi = (unsigned)s_bin; + for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0; + __syncthreads(); + for (int base = 0; base < cols; base += blockDim.x) { + const int c = base + threadIdx.x; + unsigned key = c < cols ? bf_key(xu[c]) : 0; + const bool in = c < cols && (key >> 8) == hi; + hist_add(hist, (int)(key & 0xffu), in); + } + __syncthreads(); + if (threadIdx.x == 0) { + int cum = s_gt, b = 255; + for (; b > 0; --b) { if (cum + (int)hist[b] >= want) break; cum += hist[b]; } + s_bin = (int)((hi << 8) | (unsigned)b); + s_gt = cum; // columns whose key is above K + } + __syncthreads(); + const unsigned K = (unsigned)s_bin; + const int need_eq = want - s_gt; + for (int base = 0; base < cols; base += blockDim.x) { + const int c = base + threadIdx.x; + unsigned key = c < cols ? bf_key(xu[c]) : 0; + const bool gt = c < cols && key > K; + const bool eq = c < cols && key == K; + if (gt) { const int j = atomicAdd(&s_cnt, 1); pv[j] = __uint_as_float((unsigned)xu[c] << 16); pi[j] = c; } + const unsigned bal = __ballot_sync(0xffffffffu, eq); + if (lane == 0) warp_eq[w] = __popc(bal); + __syncthreads(); + if (eq) { + int before = s_eq_seen + __popc(bal & ((1u << lane) - 1u)); + for (int i = 0; i < w; ++i) before += warp_eq[i]; + if (before < need_eq) { const int j = atomicAdd(&s_cnt, 1); pv[j] = __uint_as_float((unsigned)xu[c] << 16); pi[j] = c; } + } + __syncthreads(); + if (threadIdx.x == 0) { int t = 0; for (int i = 0; i < nw; ++i) t += warp_eq[i]; s_eq_seen += t; } + __syncthreads(); + if (s_eq_seen >= need_eq && s_cnt >= want) break; // uniform: shared values after a barrier + } + __syncthreads(); + if (threadIdx.x < TOPK) { + const int k = threadIdx.x; + if (k < want) { + const float v = pv[k]; + const int c = pi[k]; + int r = 0; + for (int j = 0; j < want; ++j) if (above(pv[j], pi[j], v, c)) ++r; + o[r] = c + offset; + o[TOPK + r] = __float_as_int(v); + } else { + o[k] = -1; + o[TOPK + k] = __float_as_int(-CUDART_INF_F); + } + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 5536e52bf4..8be92a2e20 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -33,6 +33,9 @@ pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &. pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; +pub const fn_experts: []const u8 = if (available) &Blob("fatbin_fn_experts").bytes else &.{}; +pub const fn_experts_prefill: []const u8 = if (available) &Blob("fatbin_fn_experts_prefill").bytes else &.{}; +pub const fn_sample: []const u8 = if (available) &Blob("fatbin_fn_sample").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; pub const glm_experts: []const u8 = if (available) &Blob("fatbin_glm_experts").bytes else &.{}; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 5ca371044f..665bc9d6f1 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -152,16 +152,6 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-prof")) return @import("glm_forward.zig").prof(gpu, rest); if (std.mem.eql(u8, cmd, "glm-batch")) return @import("glm_forward.zig").batch(gpu, rest); if (std.mem.eql(u8, cmd, "glm-weights")) return @import("glm_weights.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); - if (std.mem.eql(u8, cmd, "fn-gen-compiled")) return fn_run.generate(gpu, rest, .compiled); - if (std.mem.eql(u8, cmd, "fn-mtp-profile")) return @import("fn_mtp.zig").generateP(gpu, rest, true, true, true); - if (std.mem.eql(u8, cmd, "fn-mtp-parallel")) { - @import("fn_mtp.zig").parallel = true; - return @import("fn_mtp.zig").generate(gpu, rest, true, true); - } - if (std.mem.eql(u8, cmd, "fn-mtp-host")) return @import("fn_mtp.zig").generate(gpu, rest, true, false); - if (std.mem.eql(u8, cmd, "fn-mtp-eager")) return @import("fn_mtp.zig").generate(gpu, rest, false, true); - if (std.mem.eql(u8, cmd, "fn-gen-profile")) return fn_run.generate(gpu, rest, .profile); - if (std.mem.eql(u8, cmd, "fn-gen-interp")) return fn_run.generate(gpu, rest, .interp); std.debug.print("{s}", .{usage}); return error.UnknownCommand; } From 2903dd75fcc2f09d343cc937510721741c526acc Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 19:56:47 +0000 Subject: [PATCH 06/31] cuda abi: the driver entries the GLM family uses (module function enumeration, kernel parameter info, virtual memory) --- zig/src/cuda/abi.zig | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index 485edaabde..96c5f9cdea 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -3,6 +3,23 @@ pub const Result = c_int; pub const Device = c_int; pub const DevicePtr = u64; +/// CUmemGenericAllocationHandle. +pub const MemHandle = u64; +/// CUmemLocation: type 1 = a device, id its ordinal. +pub const MemLocation = extern struct { type: c_int = 1, id: c_int = 0 }; +/// CUmemAllocationProp: type 1 = pinned device memory, no shareable handle. +pub const MemAllocationProp = extern struct { + type: c_int = 1, + requested_handle_types: c_int = 0, + location: MemLocation = .{}, + win32_handle_meta_data: ?*anyopaque = null, + compression_type: u8 = 0, + gpu_direct_rdma_capable: u8 = 0, + usage: u16 = 0, + reserved: [4]u8 = @splat(0), +}; +/// CUmemAccessDesc: flags 3 = read and write. +pub const MemAccessDesc = extern struct { location: MemLocation = .{}, flags: c_int = 3 }; pub const Context = ?*opaque {}; pub const Module = ?*opaque {}; @@ -192,6 +209,15 @@ pub const Api = struct { cuMemsetD32_v2: *const fn (DevicePtr, c_uint, usize) callconv(.c) R, cuMemsetD8Async: *const fn (DevicePtr, u8, usize, Stream) callconv(.c) R, cuMemsetD32Async: *const fn (DevicePtr, c_uint, usize, Stream) callconv(.c) R, + // virtual memory: address ranges reserved once, physical chunks mapped into them as a cache grows + cuMemAddressReserve: *const fn (*DevicePtr, usize, usize, DevicePtr, u64) callconv(.c) R, + cuMemAddressFree: *const fn (DevicePtr, usize) callconv(.c) R, + cuMemCreate: *const fn (*MemHandle, usize, *const MemAllocationProp, u64) callconv(.c) R, + cuMemRelease: *const fn (MemHandle) callconv(.c) R, + cuMemMap: *const fn (DevicePtr, usize, usize, MemHandle, u64) callconv(.c) R, + cuMemUnmap: *const fn (DevicePtr, usize) callconv(.c) R, + cuMemSetAccess: *const fn (DevicePtr, usize, *const MemAccessDesc, usize) callconv(.c) R, + cuMemGetAllocationGranularity: *const fn (*usize, *const MemAllocationProp, c_int) callconv(.c) R, cuStreamCreate: *const fn (*Stream, c_uint) callconv(.c) R, cuStreamDestroy_v2: *const fn (Stream) callconv(.c) R, cuStreamSynchronize: *const fn (Stream) callconv(.c) R, @@ -230,6 +256,14 @@ pub const Api = struct { cuGraphExecDestroy: *const fn (GraphExec) callconv(.c) R, cuGraphExecUpdate_v2: *const fn (GraphExec, Graph, *ExecUpdateResultInfo) callconv(.c) R, cuGraphExecKernelNodeSetParams_v2: *const fn (GraphExec, GraphNode, *const KernelNodeParams) callconv(.c) R, + cuGraphNodeGetType: *const fn (GraphNode, *c_int) callconv(.c) R, + cuGraphMemcpyNodeGetParams: *const fn (GraphNode, *anyopaque) callconv(.c) R, + cuGraphMemsetNodeGetParams: *const fn (GraphNode, *anyopaque) callconv(.c) R, + cuFuncGetParamInfo: *const fn (Function, usize, *usize, *usize) callconv(.c) R, + cuKernelGetParamInfo: *const fn (Kernel, usize, *usize, *usize) callconv(.c) R, + cuModuleGetFunctionCount: *const fn (*c_uint, Module) callconv(.c) R, + cuModuleEnumerateFunctions: *const fn ([*]Function, c_uint, Module) callconv(.c) R, + cuFuncGetName: *const fn (*?[*:0]const u8, Function) callconv(.c) R, cuGetErrorName: *const fn (R, *?[*:0]const u8) callconv(.c) R, cuGetErrorString: *const fn (R, *?[*:0]const u8) callconv(.c) R, }; From 87ef757490aac2599bfb9a2fa8c94de7d26f07f0 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 22:18:12 +0000 Subject: [PATCH 07/31] tf-cuda-test glm-*: exit right after the pass or summary line (the teardown hung in NCCL's service threads under the library's unload, reported on PR 855); the startup line says a family keeps its own prompt states instead of 'no prompt cache' --- zig/src/native/cuda.zig | 2 +- zig/tests/cuda/glm_forward.zig | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index f5d2fb78f5..5af024486e 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -407,7 +407,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O 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", + if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else if (@hasField(@TypeOf(loaded), "target")) ", no prompt cache" else ", prompt states kept by the family", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); if (@hasField(@TypeOf(loaded), "target") and cache > 0) { diff --git a/zig/tests/cuda/glm_forward.zig b/zig/tests/cuda/glm_forward.zig index 989bd1bc6f..860a47b232 100644 --- a/zig/tests/cuda/glm_forward.zig +++ b/zig/tests/cuda/glm_forward.zig @@ -513,6 +513,7 @@ pub fn batch(gpu: Gpu, args: []const [:0]const u8) !void { fwd.freeSeq(e, one[0].s); } std.debug.print("rank {d}: glm-batch passed\n", .{rank}); + std.process.exit(0); // a test binary: no teardown (NCCL's service threads under the library's unload hang it) } // -- glm-tf ------------------------------------------------------------------------------------------------------------- @@ -570,6 +571,7 @@ pub fn teacher(gpu: Gpu, args: []const [:0]const u8) !void { j += w; } std.debug.print("rank {d}: teacher-forced on {d} prompt tokens: the top token equals Python's next token at {d} of {d} positions\n", .{ rank, prompt.len, hits, serial.len }); + std.process.exit(0); // as glm-batch } // -- glm-prof ----------------------------------------------------------------------------------------------------------- @@ -706,4 +708,5 @@ pub fn prof(gpu: Gpu, args: []const [:0]const u8) !void { for (rows.items[0..@min(rows.items.len, 22)]) |x| std.debug.print(" {d:7.3} ms {d:5.1}% {s}\n", .{ x.ms, 100 * x.ms / total, x.name[0..@min(x.name.len, 90)] }); } } + std.process.exit(0); // as glm-batch } From a2eefa232ebf7dfce4e6d832deda1ea715461a34 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 08/31] glm cuda: zig fmt --- zig/src/families/glm_cuda/api.zig | 8 +- zig/src/families/glm_cuda/blocks.zig | 48 ++-- zig/src/families/glm_cuda/calibrate.zig | 8 +- zig/src/families/glm_cuda/decode.zig | 8 +- zig/src/families/glm_cuda/dflash.zig | 172 ++++++++--- zig/src/families/glm_cuda/forward.zig | 102 ++++--- zig/src/families/glm_cuda/kern.zig | 368 ++++++++++++++++++++---- zig/src/families/glm_cuda/kern_exl3.zig | 92 +++++- zig/src/families/glm_cuda/native.zig | 49 ++-- zig/src/families/glm_cuda/seq.zig | 6 +- zig/src/families/glm_cuda/vision.zig | 100 +++++-- zig/src/families/glm_cuda/weights.zig | 39 +-- zig/src/native/cuda.zig | 2 +- zig/tests/cuda/glm_forward.zig | 12 +- zig/tests/cuda/glm_weights.zig | 3 +- 15 files changed, 730 insertions(+), 287 deletions(-) diff --git a/zig/src/families/glm_cuda/api.zig b/zig/src/families/glm_cuda/api.zig index 4060a7f6aa..52f5098701 100644 --- a/zig/src/families/glm_cuda/api.zig +++ b/zig/src/families/glm_cuda/api.zig @@ -21,9 +21,7 @@ pub fn dtypeBytes(t: DType) usize { } pub fn parseDType(s: []const u8) ?DType { - const map = std.StaticStringMap(DType).initComptime(.{ .{ "bfloat16", .bf16 }, .{ "float16", .f16 }, - .{ "float32", .f32 }, .{ "int32", .i32 }, .{ "int64", .i64 }, .{ "uint8", .u8 }, .{ "int8", .i8 }, - .{ "int16", .i16 }, .{ "float8_e4m3fn", .f8e4m3 }, .{ "bool", .bool_ } }); + const map = std.StaticStringMap(DType).initComptime(.{ .{ "bfloat16", .bf16 }, .{ "float16", .f16 }, .{ "float32", .f32 }, .{ "int32", .i32 }, .{ "int64", .i64 }, .{ "uint8", .u8 }, .{ "int8", .i8 }, .{ "int16", .i16 }, .{ "float8_e4m3fn", .f8e4m3 }, .{ "bool", .bool_ } }); return map.get(s); } @@ -128,9 +126,7 @@ pub const Kernels = struct { pub fn load(d: *const cuda.Driver) !Kernels { var k: Kernels = .{ .ext = undefined }; - const images = [_][]const u8{ cuda.kernels.fn_experts, cuda.kernels.fn_experts_prefill, cuda.kernels.glm_kda, cuda.kernels.fn_sample, cuda.kernels.glm_ops, - cuda.kernels.glm_experts, cuda.kernels.glm_experts_prefill, cuda.kernels.glm_vision, cuda.kernels.glm_experts3, - cuda.kernels.glm_experts3_prefill, cuda.kernels.glm_mm, cuda.kernels.glm_rows, cuda.kernels.glm_latent, cuda.kernels.glm_index, cuda.kernels.glm_exl3, cuda.kernels.glm_dflash }; + const images = [_][]const u8{ cuda.kernels.fn_experts, cuda.kernels.fn_experts_prefill, cuda.kernels.glm_kda, cuda.kernels.fn_sample, cuda.kernels.glm_ops, cuda.kernels.glm_experts, cuda.kernels.glm_experts_prefill, cuda.kernels.glm_vision, cuda.kernels.glm_experts3, cuda.kernels.glm_experts3_prefill, cuda.kernels.glm_mm, cuda.kernels.glm_rows, cuda.kernels.glm_latent, cuda.kernels.glm_index, cuda.kernels.glm_exl3, cuda.kernels.glm_dflash }; for (images, 0..) |img, i| k.ext[i] = try cuda.Module.load(d, img); return k; } diff --git a/zig/src/families/glm_cuda/blocks.zig b/zig/src/families/glm_cuda/blocks.zig index 1d2c4fc335..70049cb2dd 100644 --- a/zig/src/families/glm_cuda/blocks.zig +++ b/zig/src/families/glm_cuda/blocks.zig @@ -71,22 +71,14 @@ fn mlpW(store: *const api.Store, p: []const u8) !MlpW { } fn dsaW(store: *const api.Store, p: []const u8) !DsaW { - return .{ .proj = try b16(store, "{s}.dsa.proj", .{p}), .q_norm = try ptr(store, "{s}.dsa.q_norm", .{p}), .kv_norm = try ptr(store, "{s}.dsa.kv_norm", .{p}), - .q_b = try b16(store, "{s}.dsa.q_b", .{p}), .wk = try ptr(store, "{s}.dsa.absorb.wk", .{p}), .wv = try ptr(store, "{s}.dsa.absorb.wv", .{p}), - .o = try b16(store, "{s}.dsa.o", .{p}), .index = .{ .kw = try b16(store, "{s}.dsa.index.kw", .{p}), .qb = try b16(store, "{s}.dsa.index.qb", .{p}), - .ln_w = try ptr(store, "{s}.dsa.index.ln_w", .{p}), .ln_b = try ptr(store, "{s}.dsa.index.ln_b", .{p}), .gate = try ptr(store, "{s}.dsa.index.gate", .{p}), - .ape = try ptr(store, "{s}.dsa.index.ape", .{p}) } }; + return .{ .proj = try b16(store, "{s}.dsa.proj", .{p}), .q_norm = try ptr(store, "{s}.dsa.q_norm", .{p}), .kv_norm = try ptr(store, "{s}.dsa.kv_norm", .{p}), .q_b = try b16(store, "{s}.dsa.q_b", .{p}), .wk = try ptr(store, "{s}.dsa.absorb.wk", .{p}), .wv = try ptr(store, "{s}.dsa.absorb.wv", .{p}), .o = try b16(store, "{s}.dsa.o", .{p}), .index = .{ .kw = try b16(store, "{s}.dsa.index.kw", .{p}), .qb = try b16(store, "{s}.dsa.index.qb", .{p}), .ln_w = try ptr(store, "{s}.dsa.index.ln_w", .{p}), .ln_b = try ptr(store, "{s}.dsa.index.ln_b", .{p}), .gate = try ptr(store, "{s}.dsa.index.gate", .{p}), .ape = try ptr(store, "{s}.dsa.index.ape", .{p}) } }; } fn moeW(store: *const api.Store, p: []const u8) !MoeW { var b: [160]u8 = undefined; - var m: MoeW = .{ .router = try ptr(store, "{s}.moe.router", .{p}), .bias = try ptr(store, "{s}.moe.bias", .{p}), - .shared = try mlpW(store, try std.fmt.bufPrint(&b, "{s}.moe.shared", .{p})) }; + var m: MoeW = .{ .router = try ptr(store, "{s}.moe.router", .{p}), .bias = try ptr(store, "{s}.moe.bias", .{p}), .shared = try mlpW(store, try std.fmt.bufPrint(&b, "{s}.moe.shared", .{p})) }; if ((store.ints.get("exl3") orelse 0) != 0) { - m.exl3 = .{ .gt = try ptr(store, "{s}.moe.experts.gt", .{p}), .ut = try ptr(store, "{s}.moe.experts.ut", .{p}), .dt = try ptr(store, "{s}.moe.experts.dt", .{p}), - .suh_g = try ptr(store, "{s}.moe.experts.suh_g", .{p}), .suh_u = try ptr(store, "{s}.moe.experts.suh_u", .{p}), - .svh_g = try ptr(store, "{s}.moe.experts.svh_g", .{p}), .svh_u = try ptr(store, "{s}.moe.experts.svh_u", .{p}), - .suh_d = try ptr(store, "{s}.moe.experts.suh_d", .{p}), .svh_d = try ptr(store, "{s}.moe.experts.svh_d", .{p}) }; + m.exl3 = .{ .gt = try ptr(store, "{s}.moe.experts.gt", .{p}), .ut = try ptr(store, "{s}.moe.experts.ut", .{p}), .dt = try ptr(store, "{s}.moe.experts.dt", .{p}), .suh_g = try ptr(store, "{s}.moe.experts.suh_g", .{p}), .suh_u = try ptr(store, "{s}.moe.experts.suh_u", .{p}), .svh_g = try ptr(store, "{s}.moe.experts.svh_g", .{p}), .svh_u = try ptr(store, "{s}.moe.experts.svh_u", .{p}), .suh_d = try ptr(store, "{s}.moe.experts.suh_d", .{p}), .svh_d = try ptr(store, "{s}.moe.experts.svh_d", .{p}) }; } else { m.up = try ptr(store, "{s}.moe.experts.up", .{p}); m.down = try ptr(store, "{s}.moe.experts.down", .{p}); @@ -101,9 +93,7 @@ fn layerW(store: *const api.Store, p: []const u8, kda: bool, dense: bool, plain: l.ffn_hc = try hcw(store, p, "ffn_hc"); } if (kda) { - l.kda = .{ .proj = try b16(store, "{s}.kda.proj", .{p}), .fb = try b16(store, "{s}.kda.fb", .{p}), .gb = try b16(store, "{s}.kda.gb", .{p}), - .conv = try ptr(store, "{s}.kda.conv", .{p}), .a_log = try ptr(store, "{s}.kda.a_log", .{p}), .dt_bias = try ptr(store, "{s}.kda.dt_bias", .{p}), - .norm = try ptr(store, "{s}.kda.norm", .{p}), .o = try b16(store, "{s}.kda.o", .{p}) }; + l.kda = .{ .proj = try b16(store, "{s}.kda.proj", .{p}), .fb = try b16(store, "{s}.kda.fb", .{p}), .gb = try b16(store, "{s}.kda.gb", .{p}), .conv = try ptr(store, "{s}.kda.conv", .{p}), .a_log = try ptr(store, "{s}.kda.a_log", .{p}), .dt_bias = try ptr(store, "{s}.kda.dt_bias", .{p}), .norm = try ptr(store, "{s}.kda.norm", .{p}), .o = try b16(store, "{s}.kda.o", .{p}) }; } else l.dsa = try dsaW(store, p); var b: [160]u8 = undefined; if (dense) l.mlp = try mlpW(store, try std.fmt.bufPrint(&b, "{s}.mlp", .{p})) else l.moe = try moeW(store, p); @@ -146,8 +136,7 @@ pub const Weights = struct { if (w.kinds[i]) w.nkda += 1 else w.ndsa += 1; } if (store.ints.get("mtp").? != 0) { - w.mtp = .{ .enorm = try ptr(store, "mtp.enorm", .{}), .hnorm = try ptr(store, "mtp.hnorm", .{}), .eh = try b16(store, "mtp.eh", .{}), - .norm = try ptr(store, "mtp.norm", .{}), .layer = try layerW(store, "mtp.layer", false, false, true, w.ndsa) }; + w.mtp = .{ .enorm = try ptr(store, "mtp.enorm", .{}), .hnorm = try ptr(store, "mtp.hnorm", .{}), .eh = try b16(store, "mtp.eh", .{}), .norm = try ptr(store, "mtp.norm", .{}), .layer = try layerW(store, "mtp.layer", false, false, true, w.ndsa) }; } w.embed = try ptr(store, "embed", .{}); w.norm = try ptr(store, "norm", .{}); @@ -235,8 +224,7 @@ fn kdaBlock(e: *Engine, b: *Buffers, l: LayerW, R: i64) !u64 { try mm(e, b, p + FA_OFF * 2, PROJ_W, k.fb, b.ka, false, R); try mm(e, b, p + GA_OFF * 2, PROJ_W, k.gb, b.kg, false, R); if (!b.prefill) { - try kern.kdaChainSeg(&e.k, b.segs, @intCast(e.step.segs.len), li, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, k.conv, k.a_log, k.dt_bias, - k.norm, b.kout, sq.MAX_ROWS, PROJ_W); + try kern.kdaChainSeg(&e.k, b.segs, @intCast(e.step.segs.len), li, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, k.conv, k.a_log, k.dt_bias, k.norm, b.kout, sq.MAX_ROWS, PROJ_W); return outProj(e, b, b.kout, LL * 128, k.o, R); } const s = e.step.segs[0].s; @@ -247,17 +235,14 @@ fn kdaBlock(e: *Engine, b: *Buffers, l: LayerW, R: i64) !u64 { if (n > 0 and n < R) { // a kept point inside the chunk: the state and window at it, then the rest from them const cut_rec = e.cut_rec + @as(u64, @intCast(li)) * @as(u64, @intCast(sq.REC * 4)); const cut_conv = e.cut_conv + @as(u64, @intCast(li)) * @as(u64, @intCast(3 * sq.CONV_C * 2)); - try kern.kdaChain(&e.k, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, conv, k.conv, s.recAt(cur, li), k.a_log, k.dt_bias, k.norm, n, sc, - cut_rec, b.q_tmp, b.y_tmp); + try kern.kdaChain(&e.k, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, conv, k.conv, s.recAt(cur, li), k.a_log, k.dt_bias, k.norm, n, sc, cut_rec, b.q_tmp, b.y_tmp); try kern.convShift(&e.k, cut_conv, b.kproj, n, 3 * sq.CONV_C, b.rows * PROJ_W, PROJ_W, 1, sq.CONV_C); const hw: u64 = @intCast(LL * 128); const nr: u64 = @intCast(n); const rest: kern.KdaScratch = .{ .out = sc.out + nr * hw * 2, .k = sc.k + nr * hw * 4, .v = sc.v + nr * hw * 2, .g = sc.g + nr * hw * 4, .b = sc.b + nr * @as(u64, @intCast(LL)) * 4 }; - try kern.kdaChain(&e.k, p + nr * @as(u64, @intCast(PROJ_W)) * 2, PROJ_W, B_OFF, b.ka + nr * hw * 2, LL * 128, b.kg + nr * hw * 2, LL * 128, cut_conv, k.conv, cut_rec, - k.a_log, k.dt_bias, k.norm, R - n, rest, s.recAt(1 - cur, li), b.q_tmp, b.y_tmp); + try kern.kdaChain(&e.k, p + nr * @as(u64, @intCast(PROJ_W)) * 2, PROJ_W, B_OFF, b.ka + nr * hw * 2, LL * 128, b.kg + nr * hw * 2, LL * 128, cut_conv, k.conv, cut_rec, k.a_log, k.dt_bias, k.norm, R - n, rest, s.recAt(1 - cur, li), b.q_tmp, b.y_tmp); } else { - try kern.kdaChain(&e.k, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, conv, k.conv, s.recAt(cur, li), k.a_log, k.dt_bias, k.norm, R, sc, - s.recAt(1 - cur, li), b.q_tmp, b.y_tmp); + try kern.kdaChain(&e.k, p, PROJ_W, B_OFF, b.ka, LL * 128, b.kg, LL * 128, conv, k.conv, s.recAt(cur, li), k.a_log, k.dt_bias, k.norm, R, sc, s.recAt(1 - cur, li), b.q_tmp, b.y_tmp); } // a prompt chunk keeps every row: the layer's conv window commits now (its state flips with the chunk) try kern.convShift(&e.k, conv, b.kproj, R, 3 * sq.CONV_C, b.rows * PROJ_W, PROJ_W, 1, sq.CONV_C); @@ -271,8 +256,7 @@ fn selectTokens(e: *Engine, b: *Buffers, ci: usize, R: i64) !void { var a: i64 = 0; while (a < R) : (a += B) { const n = @min(B, R - a); - try kern.scores(&e.k, b.qi + @as(u64, @intCast(a * kern.IDX_HEADS * kern.IDX_DIM * 2)), b.ikr + @as(u64, @intCast(a * IDX_KW + kern.IDX_DIM)) * 2, IDX_KW, - t.at(a), ci, b.scores, n, e.step.np); + try kern.scores(&e.k, b.qi + @as(u64, @intCast(a * kern.IDX_HEADS * kern.IDX_DIM * 2)), b.ikr + @as(u64, @intCast(a * IDX_KW + kern.IDX_DIM)) * 2, IDX_KW, t.at(a), ci, b.scores, n, e.step.np); try kern.selectRows(&e.k, b.scores, b.pools + @as(u64, @intCast(a * kern.TOPK_POOLS * 8)), e.step.np, t.at(a), n); } try kern.sparseTokens(&e.k, b.pools, t, R, b.tokens, b.counts); @@ -309,8 +293,7 @@ fn dsaBlock(e: *Engine, b: *Buffers, a: DsaW, ci: usize, R: i64) !u64 { while (r0 < R) : (r0 += att) { const n = @min(att, R - r0); const off = @as(u64, @intCast(r0)) * lat_row; - try kern.sparseAttention(&e.k, b.qa + off, t.at(r0), ci, b.tokens + @as(u64, @intCast(r0 * kern.SEL_W * 4)), b.counts + @as(u64, @intCast(r0 * 4)), - b.po, b.pm, b.pl, b.ol + off, n); + try kern.sparseAttention(&e.k, b.qa + off, t.at(r0), ci, b.tokens + @as(u64, @intCast(r0 * kern.SEL_W * 4)), b.counts + @as(u64, @intCast(r0 * 4)), b.po, b.pm, b.pl, b.ol + off, n); } } try kern.expandV(&e.k, b.ol, a.wv, b.vn, R); @@ -451,12 +434,17 @@ pub fn mtpCompute(e: *Engine, b: *Buffers) !void { fn zeroFirst(e: *Engine, x: u64, pos: u64, R: i64) !void { var a: cuda.Args = .{}; - a.add(x); a.add(pos); a.add(@as(i32, @intCast(D))); + a.add(x); + a.add(pos); + a.add(@as(i32, @intCast(D))); try e.k.go(try e.k.ext(.rows, "glm_zero_first"), .{ .x = @intCast(R) }, 256, 0, &a); } fn gatherRows(e: *Engine, src: u64, idx: u64, dst: u64, n: i64) !void { var a: cuda.Args = .{}; - a.add(src); a.add(idx); a.add(dst); a.add(@as(i32, @intCast(D))); + a.add(src); + a.add(idx); + a.add(dst); + a.add(@as(i32, @intCast(D))); try e.k.go(try e.k.ext(.rows, "glm_gather_rows"), .{ .x = @intCast(n) }, 256, 0, &a); } diff --git a/zig/src/families/glm_cuda/calibrate.zig b/zig/src/families/glm_cuda/calibrate.zig index 89030cae0a..d4d6cfc31c 100644 --- a/zig/src/families/glm_cuda/calibrate.zig +++ b/zig/src/families/glm_cuda/calibrate.zig @@ -106,14 +106,13 @@ pub fn run(self: anytype, drafts: bool) !void { try w.int(0); // no confidence cut: every level timed try w.int(@intFromBool(self.e.drafter != null)); // DFlash2 drafts when loaded try self.send(.draft, &w); - var reqs = [_]decode.DraftReq{.{ .s = seqs[0], .hidden = seqs[0].hidden, .next = &follow, .count = depth, .confidence = 0, - .out = held[0..depth], .probs = probs[0..depth] }}; + var reqs = [_]decode.DraftReq{.{ .s = seqs[0], .hidden = seqs[0].hidden, .next = &follow, .count = depth, .confidence = 0, .out = held[0..depth], .probs = probs[0..depth] }}; const t0 = std.Io.Timestamp.now(self.io, .awake); // rank 1 replays the frame with the drafter it serves with: DFlash2's round when one is loaded if (self.e.drafter) |dr| { // as rank 1 replays it: the MTP head kept current (while the mix may pick it), then DFlash2's block reqs[0].count = 1; if (@import("native.zig").feedsMtp(self)) try decode.draft(self.e, &reqs); - try dflash.draft(dr, self.e, &.{.{ .s = seqs[0], .n = 1, .anchor = toks[0], .out = held[0..max_rows - 1], .probs = probs[0..max_rows - 1] }}); + try dflash.draft(dr, self.e, &.{.{ .s = seqs[0], .n = 1, .anchor = toks[0], .out = held[0 .. max_rows - 1], .probs = probs[0 .. max_rows - 1] }}); } else try decode.draft(self.e, &reqs); if (rep > 0) at[k] = @min(at[k], msSince(self.io, t0)); } @@ -127,6 +126,5 @@ pub fn run(self: anytype, drafts: bool) !void { try self.send(.release, &w); self.dropSeq(s.uid); } - std.log.info("glm costs: window 1/2/4/8 rows {d:.1}/{d:.1}/{d:.1}/{d:.1} ms; 16 streams x 1/2/4/8 rows {d:.1}/{d:.1}/{d:.1}/{d:.1} ms; head step {d:.2} ms", .{ - self.costs[0].ms, self.costs[1].ms, self.costs[3].ms, self.costs[7].ms, self.shared[3].ms, self.shared[4].ms, self.shared[5].ms, self.shared[6].ms, self.mtp_ms }); + std.log.info("glm costs: window 1/2/4/8 rows {d:.1}/{d:.1}/{d:.1}/{d:.1} ms; 16 streams x 1/2/4/8 rows {d:.1}/{d:.1}/{d:.1}/{d:.1} ms; head step {d:.2} ms", .{ self.costs[0].ms, self.costs[1].ms, self.costs[3].ms, self.costs[7].ms, self.shared[3].ms, self.shared[4].ms, self.shared[5].ms, self.shared[6].ms, self.mtp_ms }); } diff --git a/zig/src/families/glm_cuda/decode.zig b/zig/src/families/glm_cuda/decode.zig index d066c50d16..23d0656d61 100644 --- a/zig/src/families/glm_cuda/decode.zig +++ b/zig/src/families/glm_cuda/decode.zig @@ -112,8 +112,12 @@ var topk_check: ?bool = null; fn candTopk(e: *Engine, b: *Buffers, rows: i64, name: [:0]const u8) !void { const f = try e.k.ext(.sample, name); var a: cuda.Args = .{}; - a.add(b.logits); a.add(@as(i32, @intCast(HEAD_N))); a.add(@as(i32, @intCast(HEAD_N))); a.add(@as(i32, @intCast(e.w.vocab_offset))); - a.add(b.invt); a.add(b.tk); + a.add(b.logits); + a.add(@as(i32, @intCast(HEAD_N))); + a.add(@as(i32, @intCast(HEAD_N))); + a.add(@as(i32, @intCast(e.w.vocab_offset))); + a.add(b.invt); + a.add(b.tk); try e.k.go(f, .{ .x = @intCast(rows) }, 256, 0, &a); } diff --git a/zig/src/families/glm_cuda/dflash.zig b/zig/src/families/glm_cuda/dflash.zig index c7e536a4e6..018f0a1776 100644 --- a/zig/src/families/glm_cuda/dflash.zig +++ b/zig/src/families/glm_cuda/dflash.zig @@ -53,8 +53,19 @@ pub const ring_bytes: u64 = LAYERS * 2 * @as(u64, @intCast(WIN)) * KVW * 2; const B16 = kern.B16; const LayerW = struct { - in_norm: u64, post_norm: u64, qn: u64, kn: u64, a_base: u64, m_base: u64, - qkv: B16, kv: B16, o: B16, gu: B16, down: B16, a_proj: B16, m_proj: B16, + in_norm: u64, + post_norm: u64, + qn: u64, + kn: u64, + a_base: u64, + m_base: u64, + qkv: B16, + kv: B16, + o: B16, + gu: B16, + down: B16, + a_proj: B16, + m_proj: B16, }; /// One stream's draft request: ``n`` rows of the last verify to take into its context first (0: the prompt's, @@ -75,19 +86,52 @@ pub const Drafter = struct { succ: []u16, scratch: cuda.DeviceBuffer, // block buffers [MAX_M][..] - x: u64 = 0, y: u64 = 0, y2: u64 = 0, dyn: u64 = 0, qkv: u64 = 0, q: u64 = 0, kb: u64 = 0, vb: u64 = 0, att: u64 = 0, a: u64 = 0, - gu: u64 = 0, act: u64 = 0, hn: u64 = 0, g: u64 = 0, + x: u64 = 0, + y: u64 = 0, + y2: u64 = 0, + dyn: u64 = 0, + qkv: u64 = 0, + q: u64 = 0, + kb: u64 = 0, + vb: u64 = 0, + att: u64 = 0, + a: u64 = 0, + gu: u64 = 0, + act: u64 = 0, + hn: u64 = 0, + g: u64 = 0, // context buffers [MAX_CTX][..] - cin: u64 = 0, c1: u64 = 0, ctx: u64 = 0, ckv: u64 = 0, + cin: u64 = 0, + c1: u64 = 0, + ctx: u64 = 0, + ckv: u64 = 0, // the target's taps: a verify step's rows and a prompt chunk's, [rows][5 D] taps_step: u64 = 0, taps_prefill: u64 = 0, // tables - ids: u64 = 0, bpos: u64 = 0, rings: u64 = 0, starts: u64 = 0, froms: u64 = 0, cpos: u64 = 0, crings: u64 = 0, invt: u64 = 0, tk: u64 = 0, gtk: u64 = 0, - pv: u64 = 0, pi: u64 = 0, - pm: u64 = 0, pl: u64 = 0, pacc: u64 = 0, // the chunked attention's partials + ids: u64 = 0, + bpos: u64 = 0, + rings: u64 = 0, + starts: u64 = 0, + froms: u64 = 0, + cpos: u64 = 0, + crings: u64 = 0, + invt: u64 = 0, + tk: u64 = 0, + gtk: u64 = 0, + pv: u64 = 0, + pi: u64 = 0, + pm: u64 = 0, + pl: u64 = 0, + pacc: u64 = 0, // the chunked attention's partials attn_chunked: bool = true, // TF_GLM_DFLASH2_ATTN=serial: df_attn's one block a kv head (the chunked kernel otherwise) - h_ids: []i32, h_bpos: []i32, h_rings: []u64, h_starts: []i64, h_froms: []i64, h_cpos: []i64, h_crings: []u64, + h_ids: []i32, + h_bpos: []i32, + h_rings: []u64, + h_starts: []i64, + h_froms: []i64, + h_cpos: []i64, + h_crings: []u64, cand: []i32, // both ranks' gathered candidates, [2][MAX_M][W] gh: []f32, // the selector's hidden projections [MAX_M][RANK] forward_ms: f64 = 0, @@ -234,8 +278,7 @@ const Loader = struct { try l.scales.resize(l.gpa, n * kk / 64); for (names, rows) |name, r| { const src = try l.read(name, r * kk * 2); - weights_mod.quantize(bits, @alignCast(std.mem.bytesAsSlice(u16, src)), r, kk, - l.words.items[r0 * kk * bits / 32 ..][0 .. r * kk * bits / 32], l.scales.items[r0 * kk / 64 ..][0 .. r * kk / 64]); + weights_mod.quantize(bits, @alignCast(std.mem.bytesAsSlice(u16, src)), r, kk, l.words.items[r0 * kk * bits / 32 ..][0 .. r * kk * bits / 32], l.scales.items[r0 * kk / 64 ..][0 .. r * kk / 64]); r0 += r; } try l.tmp.upload(0, std.mem.sliceAsBytes(l.words.items)); @@ -337,9 +380,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *forward.Engine, dir: []const } const dr = try gpa.create(Drafter); errdefer gpa.destroy(dr); - dr.* = .{ .gpa = gpa, .d = d, .weights = wbuf, .bits = bits, .split = split, .ff = ff, .pred = &.{}, .succ = &.{}, .scratch = undefined, .h_ids = undefined, - .h_bpos = undefined, .h_rings = undefined, .h_starts = undefined, .h_froms = undefined, .h_cpos = undefined, .h_crings = undefined, - .cand = undefined, .gh = undefined }; + dr.* = .{ .gpa = gpa, .d = d, .weights = wbuf, .bits = bits, .split = split, .ff = ff, .pred = &.{}, .succ = &.{}, .scratch = undefined, .h_ids = undefined, .h_bpos = undefined, .h_rings = undefined, .h_starts = undefined, .h_froms = undefined, .h_cpos = undefined, .h_crings = undefined, .cand = undefined, .gh = undefined }; var nb: [8][96]u8 = undefined; for (&dr.layers, 0..) |*lw, i| { const N = struct { @@ -404,8 +445,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *forward.Engine, dir: []const if (std.c.getenv("TF_GLM_DFLASH2_TEMP")) |v| dr.temp = std.fmt.parseFloat(f64, std.mem.span(v)) catch 1.0; const inv: [MAX_M]f32 = @splat(1.0); try dr.scratch.upload(dr.invt - dr.scratch.ptr, std.mem.sliceAsBytes(&inv)); - std.log.info("dflash2: drafter loaded from {s}, linears at {d} bits ({d:.2} GiB on the device, {d} MiB of codebooks on the host)", - .{ dir, bits, @as(f64, @floatFromInt(l.at)) / (1 << 30), (2 * VOCAB * RANK * 2) >> 20 }); + std.log.info("dflash2: drafter loaded from {s}, linears at {d} bits ({d:.2} GiB on the device, {d} MiB of codebooks on the host)", .{ dir, bits, @as(f64, @floatFromInt(l.at)) / (1 << 30), (2 * VOCAB * RANK * 2) >> 20 }); return dr; } @@ -414,17 +454,15 @@ fn carveScratch(dr: *Drafter, d: *const cuda.Driver) !void { const M: u64 = MAX_M; const C: u64 = MAX_CTX + 16; const fields = .{ - .{ "x", M * D * 2 }, .{ "y", M * D * 2 }, .{ "y2", M * D * 2 }, .{ "dyn", M * KVW * 2 }, - .{ "qkv", M * QKV * 2 }, .{ "q", M * D * 2 }, .{ "kb", M * KVW * 2 }, .{ "vb", M * KVW * 2 }, - .{ "att", M * D * 2 }, .{ "a", M * D * 2 }, .{ "gu", M * 2 * FF * 2 }, .{ "act", M * FF * 2 }, - .{ "hn", M * D * 2 }, .{ "g", M * RANK * 4 }, .{ "cin", C * 5 * D * 2 }, .{ "c1", C * D * 2 }, - .{ "ctx", C * D * 2 }, .{ "ckv", C * 2 * KVW * 2 }, - .{ "taps_step", @as(u64, @intCast(forward.STEP_ROWS)) * 5 * D * 2 }, .{ "taps_prefill", @as(u64, @intCast(forward.PREFILL_ROWS_MAX)) * 5 * D * 2 }, - .{ "ids", M * 4 }, .{ "bpos", M * 4 }, .{ "rings", forward.MAX_STREAMS * 8 }, .{ "starts", forward.MAX_STREAMS * 8 }, - .{ "froms", forward.MAX_STREAMS * 8 }, .{ "cpos", C * 8 }, .{ "crings", C * 8 }, .{ "invt", M * 4 }, - .{ "tk", M * CW * 4 }, .{ "gtk", 2 * M * CW * 4 }, - .{ "pv", M * CHUNKS * TOPK * 4 }, .{ "pi", M * CHUNKS * TOPK * 4 }, - .{ "pm", M * QH * CH * 4 }, .{ "pl", M * QH * CH * 4 }, .{ "pacc", M * QH * CH * HD * 4 }, + .{ "x", M * D * 2 }, .{ "y", M * D * 2 }, .{ "y2", M * D * 2 }, .{ "dyn", M * KVW * 2 }, + .{ "qkv", M * QKV * 2 }, .{ "q", M * D * 2 }, .{ "kb", M * KVW * 2 }, .{ "vb", M * KVW * 2 }, + .{ "att", M * D * 2 }, .{ "a", M * D * 2 }, .{ "gu", M * 2 * FF * 2 }, .{ "act", M * FF * 2 }, + .{ "hn", M * D * 2 }, .{ "g", M * RANK * 4 }, .{ "cin", C * 5 * D * 2 }, .{ "c1", C * D * 2 }, + .{ "ctx", C * D * 2 }, .{ "ckv", C * 2 * KVW * 2 }, .{ "taps_step", @as(u64, @intCast(forward.STEP_ROWS)) * 5 * D * 2 }, .{ "taps_prefill", @as(u64, @intCast(forward.PREFILL_ROWS_MAX)) * 5 * D * 2 }, + .{ "ids", M * 4 }, .{ "bpos", M * 4 }, .{ "rings", forward.MAX_STREAMS * 8 }, .{ "starts", forward.MAX_STREAMS * 8 }, + .{ "froms", forward.MAX_STREAMS * 8 }, .{ "cpos", C * 8 }, .{ "crings", C * 8 }, .{ "invt", M * 4 }, + .{ "tk", M * CW * 4 }, .{ "gtk", 2 * M * CW * 4 }, .{ "pv", M * CHUNKS * TOPK * 4 }, .{ "pi", M * CHUNKS * TOPK * 4 }, + .{ "pm", M * QH * CH * 4 }, .{ "pl", M * QH * CH * 4 }, .{ "pacc", M * QH * CH * HD * 4 }, }; inline for (fields) |f| { @field(dr, f[0]) = at; @@ -467,7 +505,10 @@ pub fn tap(dr: *const Drafter, k: *kern.K, x: u64, taps: u64, li: usize, R: i64) if (t == at) break @intCast(i); } else return; var a: cuda.Args = .{}; - a.add(x); a.add(taps); a.add(slot); a.add(@as(i32, @intFromBool(dr.tap_s0))); + a.add(x); + a.add(taps); + a.add(slot); + a.add(@as(i32, @intFromBool(dr.tap_s0))); try k.go(try k.ext(.dflash, "df_tap"), .{ .x = @intCast(R), .y = @intCast(D / 1024) }, 256, 0, &a); } @@ -543,7 +584,13 @@ pub fn ingest(dr: *Drafter, e: *forward.Engine, runs: []const Run) !void { for (dr.layers, 0..) |l, li| { try lin(e, dr.ctx, l.kv, dr.ckv, N, false); var a: cuda.Args = .{}; - a.add(dr.ckv); a.add(l.kn); a.add(dr.cpos); a.add(dr.crings); a.add(@as(i32, @intCast(li))); a.add(EPS); a.add(THETA); + a.add(dr.ckv); + a.add(l.kn); + a.add(dr.cpos); + a.add(dr.crings); + a.add(@as(i32, @intCast(li))); + a.add(EPS); + a.add(THETA); try k.go(try k.ext(.dflash, "df_ctx_kv"), .{ .x = @intCast(N), .y = 8 }, 64, 0, &a); } try k.stream.synchronize(); @@ -676,23 +723,50 @@ fn blockBody(dr: *Drafter, e: *forward.Engine, S: usize, rank0: bool, tp: *std.I try fine(dr, e, 4); { var a: cuda.Args = .{}; - a.add(dr.qkv); a.add(@as(i32, @intCast(QKV))); a.add(l.qn); a.add(l.kn); a.add(dr.bpos); a.add(EPS); a.add(THETA); - a.add(dr.q); a.add(dr.kb); a.add(dr.vb); + a.add(dr.qkv); + a.add(@as(i32, @intCast(QKV))); + a.add(l.qn); + a.add(l.kn); + a.add(dr.bpos); + a.add(EPS); + a.add(THETA); + a.add(dr.q); + a.add(dr.kb); + a.add(dr.vb); try k.go(try k.ext(.dflash, "df_block_qkv"), .{ .x = @intCast(M), .y = 32 + 8 }, 64, 0, &a); } try fine(dr, e, 5); if (dr.attn_chunked) { var a: cuda.Args = .{}; - a.add(dr.q); a.add(dr.kb); a.add(dr.vb); a.add(dr.rings); a.add(dr.starts); a.add(dr.froms); a.add(@as(i32, @intCast(li))); - a.add(SCALE); a.add(dr.pm); a.add(dr.pl); a.add(dr.pacc); + a.add(dr.q); + a.add(dr.kb); + a.add(dr.vb); + a.add(dr.rings); + a.add(dr.starts); + a.add(dr.froms); + a.add(@as(i32, @intCast(li))); + a.add(SCALE); + a.add(dr.pm); + a.add(dr.pl); + a.add(dr.pacc); try k.go(try k.ext(.dflash, "df_attn_part"), .{ .x = @intCast(S), .y = 8, .z = CH }, 256, 0, &a); var b: cuda.Args = .{}; - b.add(dr.pm); b.add(dr.pl); b.add(dr.pacc); b.add(dr.att); + b.add(dr.pm); + b.add(dr.pl); + b.add(dr.pacc); + b.add(dr.att); try k.go(try k.ext(.dflash, "df_attn_merge"), .{ .x = @intCast(M), .y = QH }, HD, 0, &b); } else { var a: cuda.Args = .{}; - a.add(dr.q); a.add(dr.kb); a.add(dr.vb); a.add(dr.rings); a.add(dr.starts); a.add(dr.froms); a.add(@as(i32, @intCast(li))); - a.add(SCALE); a.add(dr.att); + a.add(dr.q); + a.add(dr.kb); + a.add(dr.vb); + a.add(dr.rings); + a.add(dr.starts); + a.add(dr.froms); + a.add(@as(i32, @intCast(li))); + a.add(SCALE); + a.add(dr.att); try k.go(try k.ext(.dflash, "df_attn"), .{ .x = @intCast(S), .y = 8 }, 256, 0, &a); } try fine(dr, e, 6); @@ -711,7 +785,9 @@ fn blockBody(dr: *Drafter, e: *forward.Engine, S: usize, rank0: bool, tp: *std.I try fine(dr, e, 9); { var a: cuda.Args = .{}; - a.add(dr.gu); a.add(dr.act); a.add(@as(i32, @intCast(dr.ff))); + a.add(dr.gu); + a.add(dr.act); + a.add(@as(i32, @intCast(dr.ff))); try k.go(try k.ext(.dflash, "df_swiglu"), .{ .x = @intCast(M), .y = @intCast(dr.ff / 1024) }, 256, 0, &a); } try fine(dr, e, 10); @@ -722,7 +798,9 @@ fn blockBody(dr: *Drafter, e: *forward.Engine, S: usize, rank0: bool, tp: *std.I try e.ctx.allGather(b.part, b.gath, M * D, .f32, "dflash2 down"); try fine(dr, e, 16); var a: cuda.Args = .{}; - a.add(b.gath); a.add(dr.a); a.add(@as(i32, @intCast(M))); + a.add(b.gath); + a.add(dr.a); + a.add(@as(i32, @intCast(M))); try k.go(try k.ext(.dflash, "df_sum2"), grid, 256, 0, &a); try fine(dr, e, 17); } else { @@ -742,10 +820,17 @@ fn blockBody(dr: *Drafter, e: *forward.Engine, S: usize, rank0: bool, tp: *std.I { const chunks = cdiv(forward.HEAD_N, 4096); var a: cuda.Args = .{}; - a.add(e.mbuf.logits); a.add(@as(i32, @intCast(forward.HEAD_N))); a.add(@as(i32, @intCast(e.w.vocab_offset))); a.add(dr.pv); a.add(dr.pi); + a.add(e.mbuf.logits); + a.add(@as(i32, @intCast(forward.HEAD_N))); + a.add(@as(i32, @intCast(e.w.vocab_offset))); + a.add(dr.pv); + a.add(dr.pi); try k.go(try k.ext(.dflash, "df_top16_part"), .{ .x = @intCast(M), .y = @intCast(chunks) }, 256, 0, &a); var b: cuda.Args = .{}; - b.add(dr.pv); b.add(dr.pi); b.add(@as(i32, @intCast(chunks))); b.add(dr.tk); + b.add(dr.pv); + b.add(dr.pi); + b.add(@as(i32, @intCast(chunks))); + b.add(dr.tk); try k.go(try k.ext(.dflash, "df_top16_merge"), .{ .x = @intCast(M) }, 256, 0, &b); } try fine(dr, e, 13); @@ -765,7 +850,12 @@ fn phase(dr: *Drafter, e: *forward.Engine, t: *std.Io.Timestamp, which: usize) ! fn conv(k: *kern.K, z: u64, dyn: u64, which: i32, base: u64, out: u64, add: bool, grid: cuda.Dim3) !void { var a: cuda.Args = .{}; - a.add(z); a.add(dyn); a.add(which); a.add(base); a.add(out); a.add(@as(i32, @intFromBool(add))); + a.add(z); + a.add(dyn); + a.add(which); + a.add(base); + a.add(out); + a.add(@as(i32, @intFromBool(add))); try k.go(try k.ext(.dflash, "df_conv"), grid, 256, 0, &a); } diff --git a/zig/src/families/glm_cuda/forward.zig b/zig/src/families/glm_cuda/forward.zig index 897645dbee..06836e439d 100644 --- a/zig/src/families/glm_cuda/forward.zig +++ b/zig/src/families/glm_cuda/forward.zig @@ -62,17 +62,72 @@ pub const Buffers = struct { prefill: bool, np_max: i64, // the indexer's pool count at the capacity (the scores' columns at most) mem: cuda.DeviceBuffer = undefined, - ids: u64 = 0, x: u64 = 0, normed: u64 = 0, xs: u64 = 0, post: u64 = 0, comb: u64 = 0, hcpart: u64 = 0, - ka: u64 = 0, kg: u64 = 0, kproj: u64 = 0, kout: u64 = 0, ks: kern.KdaScratch = undefined, q_tmp: u64 = 0, y_tmp: u64 = 0, - dp: u64 = 0, qr: u64 = 0, xs_qr: u64 = 0, lat: u64 = 0, xs_lat: u64 = 0, q: u64 = 0, vn: u64 = 0, - ikr: u64 = 0, igr: u64 = 0, qi: u64 = 0, qa: u64 = 0, ol: u64 = 0, po: u64 = 0, pm: u64 = 0, pl: u64 = 0, - scores: u64 = 0, pools: u64 = 0, tokens: u64 = 0, counts: u64 = 0, - gu: u64 = 0, act: u64 = 0, xs_act: u64 = 0, - mlog: u64 = 0, pick: u64 = 0, wts: u64 = 0, rpick: u64 = 0, eact: u64 = 0, ey: u64 = 0, rey: u64 = 0, plan: kern.Plan = undefined, + ids: u64 = 0, + x: u64 = 0, + normed: u64 = 0, + xs: u64 = 0, + post: u64 = 0, + comb: u64 = 0, + hcpart: u64 = 0, + ka: u64 = 0, + kg: u64 = 0, + kproj: u64 = 0, + kout: u64 = 0, + ks: kern.KdaScratch = undefined, + q_tmp: u64 = 0, + y_tmp: u64 = 0, + dp: u64 = 0, + qr: u64 = 0, + xs_qr: u64 = 0, + lat: u64 = 0, + xs_lat: u64 = 0, + q: u64 = 0, + vn: u64 = 0, + ikr: u64 = 0, + igr: u64 = 0, + qi: u64 = 0, + qa: u64 = 0, + ol: u64 = 0, + po: u64 = 0, + pm: u64 = 0, + pl: u64 = 0, + scores: u64 = 0, + pools: u64 = 0, + tokens: u64 = 0, + counts: u64 = 0, + gu: u64 = 0, + act: u64 = 0, + xs_act: u64 = 0, + mlog: u64 = 0, + pick: u64 = 0, + wts: u64 = 0, + rpick: u64 = 0, + eact: u64 = 0, + ey: u64 = 0, + rey: u64 = 0, + plan: kern.Plan = undefined, x3: kern.Exl3Scratch = .{ .xg = 0, .xu = 0, .xd = 0, .z = 0 }, // the EXL3 experts' scratch (exl3 checkpoints) - sgu: u64 = 0, sact: u64 = 0, sxs: u64 = 0, sy: u64 = 0, - part: u64 = 0, gath: u64 = 0, sk: u64 = 0, hidden: u64 = 0, fnormed: u64 = 0, fxs: u64 = 0, hrows: u64 = 0, logits: u64 = 0, - me: u64 = 0, mcat: u64 = 0, mx: u64 = 0, hin: u64 = 0, cand: u64 = 0, cgath: u64 = 0, invt: u64 = 0, tk: u64 = 0, gtk: u64 = 0, + sgu: u64 = 0, + sact: u64 = 0, + sxs: u64 = 0, + sy: u64 = 0, + part: u64 = 0, + gath: u64 = 0, + sk: u64 = 0, + hidden: u64 = 0, + fnormed: u64 = 0, + fxs: u64 = 0, + hrows: u64 = 0, + logits: u64 = 0, + me: u64 = 0, + mcat: u64 = 0, + mx: u64 = 0, + hin: u64 = 0, + cand: u64 = 0, + cgath: u64 = 0, + invt: u64 = 0, + tk: u64 = 0, + gtk: u64 = 0, /// The step's tables: positions [rows] i32, then lc, ik, ig, pk [12][rows] u64, the KDA segments, the head rows. tab: u64 = 0, tab_bytes: u64 = 0, @@ -105,31 +160,11 @@ pub const Buffers = struct { const ey_elt: i64 = if (is_prefill and !exl3) 2 else 4; const sk_cols: i64 = 16384; // splitK keeps sk * n at or below this for every linear of the model b.tab_bytes = @intCast(rows * 4 + 48 * rows * 8 + @as(i64, MAX_STREAMS) * @sizeOf(KdaSeg) + @as(i64, MAX_STREAMS) * 4); - inline for (.{ .{ "ids", rows * 4 }, .{ "x", rows * WIDE * 2 }, .{ "normed", rows * D * 2 }, .{ "xs", rows * @divExact(D, 64) * 4 }, - .{ "post", rows * kern.S * 4 }, .{ "comb", rows * kern.S * kern.S * 4 }, .{ "hcpart", rows * 16 * 25 * 4 }, - .{ "ka", rows * LL * 128 * 2 }, .{ "kg", rows * LL * 128 * 2 }, .{ "kproj", rows * sq.PROJ_W * 2 }, .{ "kout", rows * LL * 128 * 2 }, - .{ "dp", rows * (kern.Q_LORA + kern.KV_LORA) * 2 }, .{ "qr", rows * kern.Q_LORA * 2 }, .{ "xs_qr", rows * @divExact(kern.Q_LORA, 64) * 4 }, - .{ "lat", rows * kern.KV_LORA * 2 }, .{ "xs_lat", rows * @divExact(kern.KV_LORA, 64) * 4 }, .{ "q", rows * HL * kern.QK * 2 }, .{ "vn", rows * HL * kern.DV * 2 }, - .{ "ikr", rows * IDX_KW * 2 }, .{ "igr", rows * kern.IDX_DIM * 4 }, .{ "qi", rows * kern.IDX_HEADS * kern.IDX_DIM * 2 }, - .{ "qa", rows * HL * kern.KV_LORA * 2 }, .{ "ol", rows * HL * kern.KV_LORA * 2 }, - .{ "po", nch * att * HL * kern.KV_LORA * 4 }, .{ "pm", nch * att * HL * 4 }, .{ "pl", nch * att * HL * 4 }, - .{ "scores", sel_rows * @max(np_max, 1024) * 4 }, .{ "pools", rows * kern.TOPK_POOLS * 8 }, .{ "tokens", rows * kern.SEL_W * 4 }, - .{ "counts", rows * 4 }, - .{ "gu", rows * 2 * DENSE_W * 2 }, .{ "act", rows * DENSE_W * 2 }, .{ "xs_act", rows * @divExact(DENSE_W, 64) * 4 }, - .{ "mlog", rows * kern.EXPERTS * 4 }, .{ "pick", rows * kern.SLOTS * 4 }, .{ "wts", rows * kern.SLOTS * 4 }, .{ "rpick", rows * kern.TOPK * 4 }, - .{ "eact", rows * kern.SLOTS * MOE_W * 2 }, .{ "ey", rows * kern.SLOTS * D * ey_elt }, .{ "rey", rows * kern.TOPK * D * ey_elt }, - .{ "sgu", rows * 2 * MOE_W * 2 }, .{ "sact", rows * MOE_W * 2 }, .{ "sxs", rows * @divExact(MOE_W, 64) * 4 }, .{ "sy", rows * D * 4 }, - .{ "part", rows * D * 4 }, .{ "gath", 2 * rows * D * 4 }, .{ "sk", rows * sk_cols * 4 }, - .{ "hidden", rows * D * 2 }, .{ "fnormed", rows * D * 2 }, .{ "fxs", rows * @divExact(D, 64) * 4 }, .{ "hrows", @as(i64, MAX_STREAMS) * D * 2 }, - .{ "logits", head_rows * HEAD_N * 2 }, - .{ "me", rows * D * 2 }, .{ "mcat", rows * 2 * D * 2 }, .{ "mx", rows * D * 2 }, .{ "hin", rows * D * 2 }, - .{ "cand", rows * 4 * 4 }, .{ "cgath", 2 * rows * 4 * 4 }, - .{ "invt", head_rows * 4 }, .{ "tk", head_rows * @as(i64, @intCast(@import("draw.zig").W)) * 4 }, .{ "gtk", 2 * head_rows * @as(i64, @intCast(@import("draw.zig").W)) * 4 }, .{ "tab", @as(i64, @intCast(b.tab_bytes)) } }) |f| + inline for (.{ .{ "ids", rows * 4 }, .{ "x", rows * WIDE * 2 }, .{ "normed", rows * D * 2 }, .{ "xs", rows * @divExact(D, 64) * 4 }, .{ "post", rows * kern.S * 4 }, .{ "comb", rows * kern.S * kern.S * 4 }, .{ "hcpart", rows * 16 * 25 * 4 }, .{ "ka", rows * LL * 128 * 2 }, .{ "kg", rows * LL * 128 * 2 }, .{ "kproj", rows * sq.PROJ_W * 2 }, .{ "kout", rows * LL * 128 * 2 }, .{ "dp", rows * (kern.Q_LORA + kern.KV_LORA) * 2 }, .{ "qr", rows * kern.Q_LORA * 2 }, .{ "xs_qr", rows * @divExact(kern.Q_LORA, 64) * 4 }, .{ "lat", rows * kern.KV_LORA * 2 }, .{ "xs_lat", rows * @divExact(kern.KV_LORA, 64) * 4 }, .{ "q", rows * HL * kern.QK * 2 }, .{ "vn", rows * HL * kern.DV * 2 }, .{ "ikr", rows * IDX_KW * 2 }, .{ "igr", rows * kern.IDX_DIM * 4 }, .{ "qi", rows * kern.IDX_HEADS * kern.IDX_DIM * 2 }, .{ "qa", rows * HL * kern.KV_LORA * 2 }, .{ "ol", rows * HL * kern.KV_LORA * 2 }, .{ "po", nch * att * HL * kern.KV_LORA * 4 }, .{ "pm", nch * att * HL * 4 }, .{ "pl", nch * att * HL * 4 }, .{ "scores", sel_rows * @max(np_max, 1024) * 4 }, .{ "pools", rows * kern.TOPK_POOLS * 8 }, .{ "tokens", rows * kern.SEL_W * 4 }, .{ "counts", rows * 4 }, .{ "gu", rows * 2 * DENSE_W * 2 }, .{ "act", rows * DENSE_W * 2 }, .{ "xs_act", rows * @divExact(DENSE_W, 64) * 4 }, .{ "mlog", rows * kern.EXPERTS * 4 }, .{ "pick", rows * kern.SLOTS * 4 }, .{ "wts", rows * kern.SLOTS * 4 }, .{ "rpick", rows * kern.TOPK * 4 }, .{ "eact", rows * kern.SLOTS * MOE_W * 2 }, .{ "ey", rows * kern.SLOTS * D * ey_elt }, .{ "rey", rows * kern.TOPK * D * ey_elt }, .{ "sgu", rows * 2 * MOE_W * 2 }, .{ "sact", rows * MOE_W * 2 }, .{ "sxs", rows * @divExact(MOE_W, 64) * 4 }, .{ "sy", rows * D * 4 }, .{ "part", rows * D * 4 }, .{ "gath", 2 * rows * D * 4 }, .{ "sk", rows * sk_cols * 4 }, .{ "hidden", rows * D * 2 }, .{ "fnormed", rows * D * 2 }, .{ "fxs", rows * @divExact(D, 64) * 4 }, .{ "hrows", @as(i64, MAX_STREAMS) * D * 2 }, .{ "logits", head_rows * HEAD_N * 2 }, .{ "me", rows * D * 2 }, .{ "mcat", rows * 2 * D * 2 }, .{ "mx", rows * D * 2 }, .{ "hin", rows * D * 2 }, .{ "cand", rows * 4 * 4 }, .{ "cgath", 2 * rows * 4 * 4 }, .{ "invt", head_rows * 4 }, .{ "tk", head_rows * @as(i64, @intCast(@import("draw.zig").W)) * 4 }, .{ "gtk", 2 * head_rows * @as(i64, @intCast(@import("draw.zig").W)) * 4 }, .{ "tab", @as(i64, @intCast(b.tab_bytes)) } }) |f| @field(b, f[0]) = c.take(f[1]); const items = kern.maxItems(pairs, kern.EXPERTS, 16); const wide = pairs > 1024; - b.plan = .{ .members = c.take(pairs * 4), .items = c.take(items * 3 * 4), .counts = c.take(2 * 4), - .rank = c.take(if (wide) pairs * 4 else 4), .hist = c.take(if (wide) cdiv(pairs, 1024) * kern.EXPERTS * 4 else 4), .tile = if (is_prefill) kern.promptTile() else 16 }; + b.plan = .{ .members = c.take(pairs * 4), .items = c.take(items * 3 * 4), .counts = c.take(2 * 4), .rank = c.take(if (wide) pairs * 4 else 4), .hist = c.take(if (wide) cdiv(pairs, 1024) * kern.EXPERTS * 4 else 4), .tile = if (is_prefill) kern.promptTile() else 16 }; if (exl3) { const P = rows * kern.TOPK; const sk: i64 = if (is_prefill) (if (kern.promptOldKernel()) 1 else kern.promptSplit()) else kern.EXL3_SK; // the staged kernel's splits (its partials: 2 * sk * P * MOE_W fp32) @@ -351,8 +386,7 @@ pub fn prepare(e: *Engine, b: *Buffers, segs: []const Seg, kind: Kind) !void { s.v_pos0 = p0; s.v_rows = @intCast(g.rows); } - kseg[si] = .{ .first = @intCast(row), .rows = @intCast(g.rows), .rec_in = s.recAt(s.cur, 0), .rec_out = s.recAt(1 - s.cur, 0), .conv = s.conv, - .k_save = s.ks.k, .v_save = s.ks.v, .g_save = s.ks.g, .b_save = s.ks.b, .proj = s.proj }; + kseg[si] = .{ .first = @intCast(row), .rows = @intCast(g.rows), .rec_in = s.recAt(s.cur, 0), .rec_out = s.recAt(1 - s.cur, 0), .conv = s.conv, .k_save = s.ks.k, .v_save = s.ks.v, .g_save = s.ks.g, .b_save = s.ks.b, .proj = s.proj }; last[si] = @intCast(row + @as(usize, @intCast(g.rows)) - 1); for (0..@intCast(g.rows)) |r| { const q = p0 + @as(i64, @intCast(r)); diff --git a/zig/src/families/glm_cuda/kern.zig b/zig/src/families/glm_cuda/kern.zig index b9df882837..44b0ab0d85 100644 --- a/zig/src/families/glm_cuda/kern.zig +++ b/zig/src/families/glm_cuda/kern.zig @@ -73,12 +73,12 @@ pub const K = struct { /// A kernel of module ``m`` whose (mangled) name contains ``frag``: the one exact match, else the one containing it. pub fn ext(self: *K, m: Ext, frag: []const u8) !cuda.Function { - const key = std.hash.Wyhash.hash(@intFromEnum(m), frag); + const key = std.hash.Wyhash.hash(@backingInt(m), frag); if (self.cache.get(key)) |f| return f; var found: ?cuda.Function = null; - for (self.names[@intFromEnum(m)], 0..) |n, j| { + for (self.names[@backingInt(m)], 0..) |n, j| { if (std.mem.eql(u8, n, frag)) { - found = self.funcs[@intFromEnum(m)][j]; + found = self.funcs[@backingInt(m)][j]; break; } if (std.mem.indexOf(u8, n, frag) != null) { @@ -86,7 +86,7 @@ pub const K = struct { std.log.err("glm: two kernels in {s} match {s}", .{ @tagName(m), frag }); return error.AmbiguousKernel; } - found = self.funcs[@intFromEnum(m)][j]; + found = self.funcs[@backingInt(m)][j]; } } const f = found orelse { @@ -213,78 +213,132 @@ pub const Rows = struct { // -- rows.cu ------------------------------------------------------------------------------------------------------------- pub fn embed(k: *K, ids: u64, table: u64, out: u64, R: i64, copies: i64) !void { var a: cuda.Args = .{}; - a.add(ids); a.add(table); a.add(out); a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(copies))); + a.add(ids); + a.add(table); + a.add(out); + a.add(@as(i32, @intCast(D))); + a.add(@as(i32, @intCast(copies))); try k.go(try k.ext(.rows, "glm_embed"), .{ .x = @intCast(R) }, 256, 0, &a); } /// rmsnorm(x [R, d] (row stride x_stride), w) -> out (row stride o_stride), with xs the 64-groups' sums when given. pub fn rmsnorm(k: *K, x: u64, x_stride: i64, w: u64, out: u64, o_stride: i64, xs: ?u64, R: i64, d: i64) !void { var a: cuda.Args = .{}; - a.add(x); a.add(x_stride); a.add(w); a.add(out); a.add(o_stride); a.add(xs orelse @as(u64, 0)); a.add(EPS); a.add(@as(i32, @intCast(d))); + a.add(x); + a.add(x_stride); + a.add(w); + a.add(out); + a.add(o_stride); + a.add(xs orelse @as(u64, 0)); + a.add(EPS); + a.add(@as(i32, @intCast(d))); try k.go(try k.ext(.rows, "glm_rmsnorm"), .{ .x = @intCast(R) }, 256, 0, &a); } /// hc_pre: x [R, S*D] -> out [R, D] normed collapsed rows + xs, post [R, S], comb [R, S, S] (part: [R, 16, 25] fp32). pub fn hcPre(k: *K, x: u64, fn_: u64, base: u64, scale: u64, norm_w: u64, out: u64, xs: u64, post: u64, comb: u64, part: u64, R: i64) !void { var p: cuda.Args = .{}; - p.add(x); p.add(fn_); p.add(part); + p.add(x); + p.add(fn_); + p.add(part); try k.go(try k.ext(.rows, "glm_hc_part"), .{ .x = @intCast(R), .y = 16 }, 128, 0, &p); var a: cuda.Args = .{}; - a.add(x); a.add(part); a.add(base); a.add(scale); a.add(norm_w); a.add(out); a.add(xs); a.add(post); a.add(comb); - a.add(EPS); a.add(HC_EPS); a.add(@as(i32, @intCast(HC_ITERS))); + a.add(x); + a.add(part); + a.add(base); + a.add(scale); + a.add(norm_w); + a.add(out); + a.add(xs); + a.add(post); + a.add(comb); + a.add(EPS); + a.add(HC_EPS); + a.add(@as(i32, @intCast(HC_ITERS))); try k.go(try k.ext(.rows, "glm_hc_finish"), .{ .x = @intCast(R) }, 512, 0, &a); } /// hc_post: x (in place) from the gathered [2, R, D] fp32 partials. pub fn hcPost(k: *K, x: u64, gathered: u64, post: u64, comb: u64, R: i64) !void { var a: cuda.Args = .{}; - a.add(x); a.add(gathered); a.add(post); a.add(comb); a.add(@as(i32, @intCast(R))); + a.add(x); + a.add(gathered); + a.add(post); + a.add(comb); + a.add(@as(i32, @intCast(R))); try k.go(try k.ext(.rows, "glm_hc_post"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); } pub fn residualAdd(k: *K, x: u64, gathered: u64, R: i64) !void { var a: cuda.Args = .{}; - a.add(x); a.add(gathered); a.add(@as(i32, @intCast(R))); + a.add(x); + a.add(gathered); + a.add(@as(i32, @intCast(R))); try k.go(try k.ext(.rows, "glm_residual_add"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); } pub fn streamMean(k: *K, x: u64, out: u64, R: i64) !void { var a: cuda.Args = .{}; - a.add(x); a.add(out); + a.add(x); + a.add(out); try k.go(try k.ext(.rows, "glm_stream_mean"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); } pub fn swiglu(k: *K, gu: u64, out: u64, xs: u64, R: i64, w: i64) !void { var a: cuda.Args = .{}; - a.add(gu); a.add(out); a.add(xs); a.add(LIMIT); a.add(@as(i32, @intCast(w))); + a.add(gu); + a.add(out); + a.add(xs); + a.add(LIMIT); + a.add(@as(i32, @intCast(w))); try k.go(try k.ext(.rows, "glm_swiglu"), .{ .x = @intCast(R), .y = @intCast(@divExact(w, 512)) }, 256, 0, &a); } /// router logits: x [m, D] (row stride x_stride) by w [ne, D] -> out [m, ne] fp32. pub fn router(k: *K, x: u64, x_stride: i64, w: u64, out: u64, m: i64, ne: i64) !void { var a: cuda.Args = .{}; - a.add(x); a.add(x_stride); a.add(w); a.add(out); a.add(@as(i32, @intCast(m))); a.add(@as(i32, @intCast(ne))); a.add(@as(i32, @intCast(D))); + a.add(x); + a.add(x_stride); + a.add(w); + a.add(out); + a.add(@as(i32, @intCast(m))); + a.add(@as(i32, @intCast(ne))); + a.add(@as(i32, @intCast(D))); try k.go(try k.ext(.rows, "glm_router"), .{ .x = @intCast(cdiv(ne, 2)), .y = @intCast(cdiv(m, 16)) }, 256, 0, &a); } /// each row's top-8 experts and weights, then the shared expert (id 288, weight 1) in slot 8. pub fn topk(k: *K, logits: u64, bias: u64, pick: u64, wts: u64, R: i64) !void { var a: cuda.Args = .{}; - a.add(logits); a.add(bias); a.add(pick); a.add(wts); a.add(ROUTED_SCALE); a.add(@as(i32, @intCast(EXPERTS))); + a.add(logits); + a.add(bias); + a.add(pick); + a.add(wts); + a.add(ROUTED_SCALE); + a.add(@as(i32, @intCast(EXPERTS))); try k.go(try k.ext(.rows, "glm_topk"), .{ .x = @intCast(R) }, 32, 0, &a); } /// combine: y [R, SLOTS, D] bf16 (prompt) or fp32 (windows) by wts -> out [R, D] fp32. pub fn combine(k: *K, y: u64, y_f32: bool, wts: u64, out: u64, R: i64) !void { var a: cuda.Args = .{}; - a.add(y); a.add(@as(i32, @intFromBool(y_f32))); a.add(wts); a.add(out); + a.add(y); + a.add(@as(i32, @intFromBool(y_f32))); + a.add(wts); + a.add(out); try k.go(try k.ext(.rows, "glm_combine"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); } /// conv [n, 3, C] (in place) takes rows keep .. keep + 2 of [conv; proj rows] (proj [n, R, W >= C]); strides in elements. pub fn convShift(k: *K, conv: u64, proj: u64, keep: i64, conv_layer: i64, proj_layer: i64, proj_row: i64, n: i64, c: i64) !void { var a: cuda.Args = .{}; - a.add(conv); a.add(proj); a.add(@as(i32, @intCast(keep))); a.add(conv_layer); a.add(proj_layer); a.add(proj_row); a.add(@as(i32, @intCast(c))); + a.add(conv); + a.add(proj); + a.add(@as(i32, @intCast(keep))); + a.add(conv_layer); + a.add(proj_layer); + a.add(proj_row); + a.add(@as(i32, @intCast(c))); try k.go(try k.ext(.rows, "glm_conv_shift"), .{ .x = @intCast(n), .y = @intCast(cdiv(c, 1024)) }, 256, 0, &a); } @@ -330,13 +384,26 @@ pub fn mm(k: *K, x: u64, x_stride: i64, w: B16, out: u64, f32_out: bool, m: i64, const smem = mmSmem(w.bits, bm); try f.allowDynamicShared(smem); var a: cuda.Args = .{}; - a.add(x); a.add(x_stride); a.add(w.w); a.add(w.s); a.add(out); a.add(part); a.add(@as(i32, @intFromBool(f32_out))); - a.add(@as(i32, @intCast(m))); a.add(@as(i32, @intCast(w.n))); a.add(@as(i32, @intCast(w.k))); a.add(@as(i32, @intCast(sk))); + a.add(x); + a.add(x_stride); + a.add(w.w); + a.add(w.s); + a.add(out); + a.add(part); + a.add(@as(i32, @intFromBool(f32_out))); + a.add(@as(i32, @intCast(m))); + a.add(@as(i32, @intCast(w.n))); + a.add(@as(i32, @intCast(w.k))); + a.add(@as(i32, @intCast(sk))); try k.goDependent(f, .{ .x = @intCast(cdiv(w.n, 64)), .y = @intCast(cdiv(m, bm)), .z = @intCast(sk) }, if (bm <= 32) 128 else 256, smem, &a); if (sk > 1) { const total = m * w.n; var r: cuda.Args = .{}; - r.add(part); r.add(out); r.add(total); r.add(@as(i32, @intCast(sk))); r.add(@as(i32, @intFromBool(f32_out))); + r.add(part); + r.add(out); + r.add(total); + r.add(@as(i32, @intCast(sk))); + r.add(@as(i32, @intFromBool(f32_out))); try k.go(try k.ext(.mm, "glm_mm_reduce"), .{ .x = @intCast(cdiv(total, 256)) }, 256, 0, &r); } } @@ -348,31 +415,48 @@ pub fn mmTile(k: *K, bits: u8, src: u64, ss: u64, dst: u64, ds: u64, n: i64, kk: var a: cuda.Args = .{}; var nb: [32]u8 = undefined; if (bits == 0) { - a.add(src); a.add(dst); a.add(@as(i32, @intCast(n))); a.add(@as(i32, @intCast(kk))); + a.add(src); + a.add(dst); + a.add(@as(i32, @intCast(n))); + a.add(@as(i32, @intCast(kk))); try k.go(try k.ext(.mm, "glm_mm_tile16"), .{ .x = @intCast(cdiv(chunks, 256)) }, 256, 0, &a); return; } - a.add(src); a.add(ss); a.add(dst); a.add(ds); a.add(@as(i32, @intCast(n))); a.add(@as(i32, @intCast(kk))); + a.add(src); + a.add(ss); + a.add(dst); + a.add(ds); + a.add(@as(i32, @intCast(n))); + a.add(@as(i32, @intCast(kk))); try k.go(try k.ext(.mm, try std.fmt.bufPrint(&nb, "glm_mm_tile{d}", .{bits})), .{ .x = @intCast(cdiv(chunks, 256)) }, 256, 0, &a); } // -- latent.cu ---------------------------------------------------------------------------------------------------------- pub fn absorbQ(k: *K, q: u64, wk: u64, qa: u64, R: i64) !void { var a: cuda.Args = .{}; - a.add(q); a.add(wk); a.add(qa); a.add(@as(i32, @intCast(R))); + a.add(q); + a.add(wk); + a.add(qa); + a.add(@as(i32, @intCast(R))); try k.go(try k.ext(.latent, "glm_absorb_q"), .{ .x = @intCast(HL), .y = @intCast(@divExact(KV_LORA, 64)), .z = @intCast(cdiv(R, 16)) }, 64, 0, &a); } pub fn expandV(k: *K, ol: u64, wv: u64, out: u64, R: i64) !void { var a: cuda.Args = .{}; - a.add(ol); a.add(wv); a.add(out); a.add(@as(i32, @intCast(R))); + a.add(ol); + a.add(wv); + a.add(out); + a.add(@as(i32, @intCast(R))); try k.go(try k.ext(.latent, "glm_expand_v"), .{ .x = @intCast(HL), .y = @intCast(@divExact(DV, 8)), .z = @intCast(cdiv(R, 16)) }, 256, 0, &a); } /// each row's latent into its stream's cache ``ci`` at its position. pub fn latentWrite(k: *K, lat: u64, lat_stride: i64, t: Rows, ci: usize, R: i64) !void { var a: cuda.Args = .{}; - a.add(lat); a.add(lat_stride); a.add(t.lcAt(ci)); a.add(t.posAt()); + a.add(lat); + a.add(lat_stride); + a.add(t.lcAt(ci)); + a.add(t.posAt()); try k.go(try k.ext(.latent, "glm_lat_write"), .{ .x = @intCast(R) }, 64, 0, &a); } @@ -383,7 +467,14 @@ pub fn denseAttention(k: *K, qa: u64, t: Rows, ci: usize, po: u64, pm: u64, pl: const f = try k.ext(.latent, "glm_attn_dense"); try f.allowDynamicShared(ATT_SMEM); var a: cuda.Args = .{}; - a.add(qa); a.add(t.lcAt(ci)); a.add(t.posAt()); a.add(po); a.add(pm); a.add(pl); a.add(@as(i32, @intCast(R))); a.add(SCALE); + a.add(qa); + a.add(t.lcAt(ci)); + a.add(t.posAt()); + a.add(po); + a.add(pm); + a.add(pl); + a.add(@as(i32, @intCast(R))); + a.add(SCALE); try k.go(f, .{ .x = @intCast(R), .y = @intCast(nch) }, 256, ATT_SMEM, &a); try merge(k, po, pm, pl, out, t.posAt(), 0, R, nch, false); } @@ -393,14 +484,29 @@ pub fn sparseAttention(k: *K, qa: u64, t: Rows, ci: usize, tokens: u64, counts: const f = try k.ext(.latent, "glm_attn_sparse"); try f.allowDynamicShared(ATT_SMEM); var a: cuda.Args = .{}; - a.add(qa); a.add(t.lcAt(ci)); a.add(tokens); a.add(counts); a.add(po); a.add(pm); a.add(pl); a.add(@as(i32, @intCast(R))); a.add(SCALE); + a.add(qa); + a.add(t.lcAt(ci)); + a.add(tokens); + a.add(counts); + a.add(po); + a.add(pm); + a.add(pl); + a.add(@as(i32, @intCast(R))); + a.add(SCALE); try k.go(f, .{ .x = @intCast(R), .y = @intCast(SPARSE_CHUNKS) }, 256, ATT_SMEM, &a); try merge(k, po, pm, pl, out, 0, counts, R, SPARSE_CHUNKS, true); } fn merge(k: *K, po: u64, pm: u64, pl: u64, out: u64, pos: u64, counts: u64, R: i64, nch: i64, sparse: bool) !void { var a: cuda.Args = .{}; - a.add(po); a.add(pm); a.add(pl); a.add(out); a.add(pos); a.add(counts); a.add(@as(i32, @intCast(R))); a.add(@as(i32, @intCast(nch))); + a.add(po); + a.add(pm); + a.add(pl); + a.add(out); + a.add(pos); + a.add(counts); + a.add(@as(i32, @intCast(R))); + a.add(@as(i32, @intCast(nch))); a.add(@as(i32, @intFromBool(sparse))); try k.go(try k.ext(.latent, "glm_attn_merge"), .{ .x = @intCast(R), .y = @intCast(HL) }, 128, 0, &a); } @@ -409,10 +515,22 @@ fn merge(k: *K, po: u64, pm: u64, pl: u64, out: u64, pos: u64, counts: u64, R: i /// The rows' index keys (LayerNorm) and gates into their streams' rings for cache ``ci``, then the pools they complete. pub fn indexUpdate(k: *K, kr: u64, k_stride: i64, gate: u64, ln_w: u64, ln_b: u64, ape: u64, t: Rows, ci: usize, R: i64) !void { var a: cuda.Args = .{}; - a.add(kr); a.add(k_stride); a.add(gate); a.add(ln_w); a.add(ln_b); a.add(t.ikAt(ci)); a.add(t.igAt(ci)); a.add(t.posAt()); a.add(@as(f32, 1e-6)); + a.add(kr); + a.add(k_stride); + a.add(gate); + a.add(ln_w); + a.add(ln_b); + a.add(t.ikAt(ci)); + a.add(t.igAt(ci)); + a.add(t.posAt()); + a.add(@as(f32, 1e-6)); try k.go(try k.ext(.index, "glm_index_write"), .{ .x = @intCast(R) }, 128, 0, &a); var b: cuda.Args = .{}; - b.add(t.ikAt(ci)); b.add(t.igAt(ci)); b.add(ape); b.add(t.pkAt(ci)); b.add(t.posAt()); + b.add(t.ikAt(ci)); + b.add(t.igAt(ci)); + b.add(ape); + b.add(t.pkAt(ci)); + b.add(t.posAt()); try k.go(try k.ext(.index, "glm_pool_keys"), .{ .x = @intCast(R) }, 128, 0, &b); } @@ -426,32 +544,60 @@ pub fn poolBucket(last: i64, np_max: i64) i64 { /// Each of n rows' scores over np pools -> scores [n, np] fp32. pub fn scores(k: *K, qi: u64, w: u64, w_stride: i64, t: Rows, ci: usize, out: u64, n: i64, np: i64) !void { var a: cuda.Args = .{}; - a.add(qi); a.add(w); a.add(w_stride); a.add(t.pkAt(ci)); a.add(out); a.add(t.posAt()); a.add(@as(i32, @intCast(np))); - a.add(@as(f32, 0.08838834764831845)); a.add(@as(f32, 1.0 / 5.656854249492381)); + a.add(qi); + a.add(w); + a.add(w_stride); + a.add(t.pkAt(ci)); + a.add(out); + a.add(t.posAt()); + a.add(@as(i32, @intCast(np))); + a.add(@as(f32, 0.08838834764831845)); + a.add(@as(f32, 1.0 / 5.656854249492381)); try k.go(try k.ext(.index, "glm_index_scores"), .{ .x = @intCast(n), .y = @intCast(cdiv(np, 64)) }, 128, 0, &a); } /// Each row's 512 best visible pools, ascending (int64). pub fn selectRows(k: *K, sc: u64, out: u64, np: i64, t: Rows, n: i64) !void { var a: cuda.Args = .{}; - a.add(sc); a.add(out); a.add(@as(i32, @intCast(np))); a.add(t.posAt()); + a.add(sc); + a.add(out); + a.add(@as(i32, @intCast(np))); + a.add(t.posAt()); try k.go(try k.ext(.index, "glm_index_select"), .{ .x = @intCast(n) }, 1024, 0, &a); } /// Each row's attended tokens and their count (0: a dense row). pub fn sparseTokens(k: *K, pools: u64, t: Rows, rows: i64, tokens: u64, counts: u64) !void { var a: cuda.Args = .{}; - a.add(pools); a.add(t.posAt()); a.add(tokens); a.add(counts); + a.add(pools); + a.add(t.posAt()); + a.add(tokens); + a.add(counts); try k.go(try k.ext(.index, "glm_index_tokens"), .{ .x = @intCast(rows) }, 256, 0, &a); } /// A batched step's KDA layer ``li``: every stream's window rows through its own states (segs: [n] device KdaSeg). -pub fn kdaChainSeg(k: *K, segs: u64, nseg: i64, li: usize, p: u64, p_stride: i64, b_off: i64, a_: u64, a_stride: i64, g: u64, g_stride: i64, - cw: u64, a_log: u64, dt_bias: u64, norm_w: u64, out: u64, max_rows: i64, proj_w: i64) !void { +pub fn kdaChainSeg(k: *K, segs: u64, nseg: i64, li: usize, p: u64, p_stride: i64, b_off: i64, a_: u64, a_stride: i64, g: u64, g_stride: i64, cw: u64, a_log: u64, dt_bias: u64, norm_w: u64, out: u64, max_rows: i64, proj_w: i64) !void { var a: cuda.Args = .{}; - a.add(segs); a.add(@as(i32, @intCast(li))); a.add(@as(i32, @intCast(LL))); a.add(p); a.add(@as(i32, @intCast(p_stride))); a.add(@as(i32, @intCast(b_off))); - a.add(a_); a.add(@as(i32, @intCast(a_stride))); a.add(g); a.add(@as(i32, @intCast(g_stride))); a.add(cw); a.add(a_log); a.add(dt_bias); a.add(norm_w); - a.add(EPS); a.add(LOWER); a.add(out); a.add(@as(i32, @intCast(max_rows))); a.add(@as(i32, @intCast(proj_w))); + a.add(segs); + a.add(@as(i32, @intCast(li))); + a.add(@as(i32, @intCast(LL))); + a.add(p); + a.add(@as(i32, @intCast(p_stride))); + a.add(@as(i32, @intCast(b_off))); + a.add(a_); + a.add(@as(i32, @intCast(a_stride))); + a.add(g); + a.add(@as(i32, @intCast(g_stride))); + a.add(cw); + a.add(a_log); + a.add(dt_bias); + a.add(norm_w); + a.add(EPS); + a.add(LOWER); + a.add(out); + a.add(@as(i32, @intCast(max_rows))); + a.add(@as(i32, @intCast(proj_w))); try k.go(try k.ext(.kda, "glm_kda_chain_seg"), .{ .x = @intCast(LL), .y = @intCast(nseg) }, 1024, 0, &a); } @@ -460,26 +606,74 @@ pub const KdaScratch = struct { out: u64, k: u64, v: u64, g: u64, b: u64 }; /// kda.chain: projection rows p [q | k | v | f_a | g_a | b at b_off], bf16 gate rows a and g, the conv window and /// state -> out rows, state_out, and the replay inputs; windows of 64 rows or more take the three-kernel path. -pub fn kdaChain(k: *K, p: u64, p_stride: i64, b_off: i64, a: u64, a_stride: i64, g: u64, g_stride: i64, cs: u64, cw: u64, state_in: u64, - a_log: u64, dt_bias: u64, norm_w: u64, rows: i64, sc: KdaScratch, state_out: u64, q_tmp: u64, y_tmp: u64) !void { +pub fn kdaChain(k: *K, p: u64, p_stride: i64, b_off: i64, a: u64, a_stride: i64, g: u64, g_stride: i64, cs: u64, cw: u64, state_in: u64, a_log: u64, dt_bias: u64, norm_w: u64, rows: i64, sc: KdaScratch, state_out: u64, q_tmp: u64, y_tmp: u64) !void { const H: i32 = @intCast(LL); if (rows >= 64) { var pa: cuda.Args = .{}; - pa.add(H); pa.add(p); pa.add(@as(i32, @intCast(p_stride))); pa.add(@as(i32, @intCast(b_off))); pa.add(a); pa.add(@as(i32, @intCast(a_stride))); - pa.add(cs); pa.add(cw); pa.add(a_log); pa.add(dt_bias); pa.add(LOWER); pa.add(q_tmp); pa.add(sc.k); pa.add(sc.v); pa.add(sc.g); pa.add(sc.b); + pa.add(H); + pa.add(p); + pa.add(@as(i32, @intCast(p_stride))); + pa.add(@as(i32, @intCast(b_off))); + pa.add(a); + pa.add(@as(i32, @intCast(a_stride))); + pa.add(cs); + pa.add(cw); + pa.add(a_log); + pa.add(dt_bias); + pa.add(LOWER); + pa.add(q_tmp); + pa.add(sc.k); + pa.add(sc.v); + pa.add(sc.g); + pa.add(sc.b); try k.go(try k.ext(.kda, "11prep_kernel"), .{ .x = @intCast(rows), .y = @intCast(H) }, 512, 0, &pa); var sa: cuda.Args = .{}; - sa.add(H); sa.add(state_in); sa.add(q_tmp); sa.add(sc.k); sa.add(sc.v); sa.add(sc.g); sa.add(sc.b); sa.add(@as(i32, @intCast(rows))); sa.add(y_tmp); sa.add(state_out); + sa.add(H); + sa.add(state_in); + sa.add(q_tmp); + sa.add(sc.k); + sa.add(sc.v); + sa.add(sc.g); + sa.add(sc.b); + sa.add(@as(i32, @intCast(rows))); + sa.add(y_tmp); + sa.add(state_out); try k.go(try k.ext(.kda, "11step_kernelILi4ELi16E"), .{ .x = @intCast(H), .y = 128 / 4 / 4 }, 4 * 32, 0, &sa); var oa: cuda.Args = .{}; - oa.add(H); oa.add(y_tmp); oa.add(g); oa.add(@as(i32, @intCast(g_stride))); oa.add(norm_w); oa.add(EPS); oa.add(sc.out); + oa.add(H); + oa.add(y_tmp); + oa.add(g); + oa.add(@as(i32, @intCast(g_stride))); + oa.add(norm_w); + oa.add(EPS); + oa.add(sc.out); try k.go(try k.ext(.kda, "10out_kernel"), .{ .x = @intCast(rows), .y = @intCast(H) }, 128, 0, &oa); return; } var ca: cuda.Args = .{}; - ca.add(H); ca.add(p); ca.add(@as(i32, @intCast(p_stride))); ca.add(@as(i32, @intCast(b_off))); ca.add(a); ca.add(@as(i32, @intCast(a_stride))); - ca.add(g); ca.add(@as(i32, @intCast(g_stride))); ca.add(cs); ca.add(cw); ca.add(state_in); ca.add(a_log); ca.add(dt_bias); ca.add(norm_w); - ca.add(EPS); ca.add(LOWER); ca.add(@as(i32, @intCast(rows))); ca.add(sc.out); ca.add(state_out); ca.add(sc.k); ca.add(sc.v); ca.add(sc.g); ca.add(sc.b); + ca.add(H); + ca.add(p); + ca.add(@as(i32, @intCast(p_stride))); + ca.add(@as(i32, @intCast(b_off))); + ca.add(a); + ca.add(@as(i32, @intCast(a_stride))); + ca.add(g); + ca.add(@as(i32, @intCast(g_stride))); + ca.add(cs); + ca.add(cw); + ca.add(state_in); + ca.add(a_log); + ca.add(dt_bias); + ca.add(norm_w); + ca.add(EPS); + ca.add(LOWER); + ca.add(@as(i32, @intCast(rows))); + ca.add(sc.out); + ca.add(state_out); + ca.add(sc.k); + ca.add(sc.v); + ca.add(sc.g); + ca.add(sc.b); try k.go(try k.ext(.kda, "12chain_kernel"), .{ .x = @intCast(H) }, 1024, 0, &ca); } @@ -487,8 +681,17 @@ pub fn kdaChain(k: *K, p: u64, p_stride: i64, b_off: i64, a: u64, a_stride: i64, pub fn kdaReplayLayers(k: *K, state_in: u64, sk: u64, sv: u64, sg: u64, sb: u64, layers: i64, rows_cap: i64, rows: i64, state_out: u64) !void { const H: i32 = @intCast(LL); var a: cuda.Args = .{}; - a.add(H); a.add(state_in); a.add(@as(u64, @intCast(LL * 128 * 128))); a.add(sk); a.add(sv); a.add(sg); a.add(sb); - a.add(@as(u64, @intCast(rows_cap * LL * 128))); a.add(@as(u64, @intCast(rows_cap * LL))); a.add(@as(i32, @intCast(rows))); a.add(state_out); + a.add(H); + a.add(state_in); + a.add(@as(u64, @intCast(LL * 128 * 128))); + a.add(sk); + a.add(sv); + a.add(sg); + a.add(sb); + a.add(@as(u64, @intCast(rows_cap * LL * 128))); + a.add(@as(u64, @intCast(rows_cap * LL))); + a.add(@as(i32, @intCast(rows))); + a.add(state_out); try k.go(try k.ext(.kda, "20replay_layers_kernel"), .{ .x = @intCast(layers * LL) }, 1024, 0, &a); } @@ -508,17 +711,37 @@ pub fn expertsPlan(k: *K, picks: u64, pairs: i64, plan: Plan) !void { var a: cuda.Args = .{}; const E: i32 = @intCast(EXPERTS); if (pairs <= 1024) { - a.add(picks); a.add(@as(i32, @intCast(pairs))); a.add(E); a.add(@as(i32, @intCast(plan.tile))); a.add(plan.members); a.add(plan.items); a.add(plan.counts); + a.add(picks); + a.add(@as(i32, @intCast(pairs))); + a.add(E); + a.add(@as(i32, @intCast(plan.tile))); + a.add(plan.members); + a.add(plan.items); + a.add(plan.counts); try k.go(try k.ext(.experts, "11plan_kernelE"), .{ .x = 1 }, 1024, 0, &a); } else { const nblk = cdiv(pairs, 1024); - a.add(picks); a.add(@as(i32, @intCast(pairs))); a.add(E); a.add(plan.rank); a.add(plan.hist); + a.add(picks); + a.add(@as(i32, @intCast(pairs))); + a.add(E); + a.add(plan.rank); + a.add(plan.hist); try k.go(try k.ext(.experts, "9plan_rankE"), .{ .x = @intCast(nblk) }, 1024, 0, &a); var b: cuda.Args = .{}; - b.add(@as(i32, @intCast(nblk))); b.add(E); b.add(@as(i32, @intCast(plan.tile))); b.add(plan.hist); b.add(plan.items); b.add(plan.counts); + b.add(@as(i32, @intCast(nblk))); + b.add(E); + b.add(@as(i32, @intCast(plan.tile))); + b.add(plan.hist); + b.add(plan.items); + b.add(plan.counts); try k.go(try k.ext(.experts, "12plan_offsetsE"), .{ .x = 1 }, 1024, 0, &b); var s: cuda.Args = .{}; - s.add(picks); s.add(@as(i32, @intCast(pairs))); s.add(E); s.add(plan.rank); s.add(plan.hist); s.add(plan.members); + s.add(picks); + s.add(@as(i32, @intCast(pairs))); + s.add(E); + s.add(plan.rank); + s.add(plan.hist); + s.add(plan.members); try k.go(try k.ext(.experts, "12plan_scatterE"), .{ .x = @intCast(cdiv(pairs, 256)) }, 256, 0, &s); } } @@ -529,8 +752,17 @@ pub fn expertsCall(k: *K, prefill: bool, epi: i64, x: u64, x_stride: i64, slots: const gs: i64 = 64; const m: i64 = if (epi == 2) 2 else 1; var a: cuda.Args = .{}; - a.add(x); a.add(@as(i32, @intCast(x_stride))); a.add(@as(i32, @intCast(slots))); a.add(w); a.add(@as(i32, @intCast(kb))); - a.add(@as(i32, @intCast(nb))); a.add(plan.items); a.add(plan.counts); a.add(plan.members); a.add(out); a.add(@as(i32, @intCast(width))); + a.add(x); + a.add(@as(i32, @intCast(x_stride))); + a.add(@as(i32, @intCast(slots))); + a.add(w); + a.add(@as(i32, @intCast(kb))); + a.add(@as(i32, @intCast(nb))); + a.add(plan.items); + a.add(plan.counts); + a.add(plan.members); + a.add(out); + a.add(@as(i32, @intCast(width))); a.add(limit); var buf: [128]u8 = undefined; const sbv: i64 = if (k.experts_sym) 1 else 2; // a quad's scale (and bias) uint4s a group @@ -556,13 +788,20 @@ pub fn expertsCall(k: *K, prefill: bool, epi: i64, x: u64, x_stride: i64, slots: /// x rows (``copies`` streams each) <- feats rows: image features in place of placeholder embeddings. pub fn injectRows(k: *K, feats: u64, x: u64, rows: i64, copies: i64) !void { var a: cuda.Args = .{}; - a.add(feats); a.add(x); a.add(@as(i32, @intCast(rows))); a.add(@as(i32, @intCast(copies))); + a.add(feats); + a.add(x); + a.add(@as(i32, @intCast(rows))); + a.add(@as(i32, @intCast(copies))); try k.go(try k.ext(.ops, "glm_inject_rows"), .{ .x = @intCast(@min(4096, cdiv(rows * copies * 4096, 256))) }, 256, 0, &a); } pub fn castBf16(k: *K, src: u64, src_stride: i64, dst: u64, dst_stride: i64, rows: i64, cols: i64) !void { var a: cuda.Args = .{}; - a.add(src); a.add(src_stride); a.add(dst); a.add(dst_stride); a.add(@as(i32, @intCast(cols))); + a.add(src); + a.add(src_stride); + a.add(dst); + a.add(dst_stride); + a.add(@as(i32, @intCast(cols))); try k.go(try k.ext(.ops, "glm_cast_bf16"), .{ .x = @intCast(cdiv(cols, 256)), .y = @intCast(rows) }, 256, 0, &a); } @@ -570,7 +809,11 @@ pub fn castBf16(k: *K, src: u64, src_stride: i64, dst: u64, dst_stride: i64, row pub fn copyRows(k: *K, dst: u64, dpitch: usize, src: u64, spitch: usize, bytes: usize, rows: usize) !void { if (rows == 0) return; var a: cuda.Args = .{}; - a.add(src); a.add(@as(i64, @intCast(spitch / 4))); a.add(dst); a.add(@as(i64, @intCast(dpitch / 4))); a.add(@as(i32, @intCast(bytes / 4))); + a.add(src); + a.add(@as(i64, @intCast(spitch / 4))); + a.add(dst); + a.add(@as(i64, @intCast(dpitch / 4))); + a.add(@as(i32, @intCast(bytes / 4))); try k.go(try k.ext(.ops, "glm_copy_rows"), .{ .x = @intCast(@min(64, cdiv(@intCast(bytes / 4), 256))), .y = @intCast(rows) }, 256, 0, &a); } @@ -578,7 +821,12 @@ pub fn copyRows(k: *K, dst: u64, dpitch: usize, src: u64, spitch: usize, bytes: pub fn rowsTop(k: *K, logits: u64, cols: i64, id_map: u64, offset: i64, out: u64, rows: i64) !void { const f = try k.ext(.sample, "fn_rows_top"); var a: cuda.Args = .{}; - a.add(logits); a.add(@as(i32, @intCast(cols))); a.add(@as(i32, @intCast(cols))); a.add(id_map); a.add(@as(i32, @intCast(offset))); a.add(out); + a.add(logits); + a.add(@as(i32, @intCast(cols))); + a.add(@as(i32, @intCast(cols))); + a.add(id_map); + a.add(@as(i32, @intCast(offset))); + a.add(out); try k.go(f, .{ .x = @intCast(rows) }, 1024, 0, &a); } diff --git a/zig/src/families/glm_cuda/kern_exl3.zig b/zig/src/families/glm_cuda/kern_exl3.zig index 8e6e30a364..a65d6ce6a1 100644 --- a/zig/src/families/glm_cuda/kern_exl3.zig +++ b/zig/src/families/glm_cuda/kern_exl3.zig @@ -74,25 +74,56 @@ pub fn exl3Routed(k: *K, x: u64, x_stride: i64, picks: u64, plan: Plan, w: Exl3, const rows_x = (staged or fused) and shx; if (rows_x) { var a: cuda.Args = .{}; - a.add(x); a.add(@as(i32, @intCast(x_stride))); a.add(w.suh_g); a.add(sc.xg); a.add(@as(i32, @intCast(D))); + a.add(x); + a.add(@as(i32, @intCast(x_stride))); + a.add(w.suh_g); + a.add(sc.xg); + a.add(@as(i32, @intCast(D))); try k.go(try k.ext(.exl3, "glm_exl3_rot_rows"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 128)) }, 32, 0, &a); } else { var a: cuda.Args = .{}; - a.add(x); a.add(@as(i32, @intCast(x_stride))); a.add(picks); a.add(w.suh_g); a.add(w.suh_u); a.add(sc.xg); a.add(sc.xu); a.add(@as(i32, @intCast(D))); + a.add(x); + a.add(@as(i32, @intCast(x_stride))); + a.add(picks); + a.add(w.suh_g); + a.add(w.suh_u); + a.add(sc.xg); + a.add(sc.xu); + a.add(@as(i32, @intCast(D))); try k.go(try k.ext(.exl3, "glm_exl3_rot_in"), .{ .x = @intCast(P), .y = @intCast(@divExact(D, 128)), .z = 2 }, 32, 0, &a); } if (fused) { // gate and up with their epilogue -> sc.xd; the down projection with its epilogue -> y { var a: cuda.Args = .{}; - a.add(sc.xg); a.add(if (rows_x) sc.xg else sc.xu); a.add(w.gt); a.add(w.ut); a.add(plan.items); a.add(plan.counts); a.add(plan.members); - a.add(w.svh_g); a.add(w.svh_u); a.add(w.suh_d); a.add(sc.xd); - a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(width))); a.add(@as(i32, if (rows_x) 3 else 0)); a.add(LIMIT); + a.add(sc.xg); + a.add(if (rows_x) sc.xg else sc.xu); + a.add(w.gt); + a.add(w.ut); + a.add(plan.items); + a.add(plan.counts); + a.add(plan.members); + a.add(w.svh_g); + a.add(w.svh_u); + a.add(w.suh_d); + a.add(sc.xd); + a.add(@as(i32, @intCast(D))); + a.add(@as(i32, @intCast(width))); + a.add(@as(i32, if (rows_x) 3 else 0)); + a.add(LIMIT); try k.go(try k.ext(.exl3, "glm_exl3_prompt_gu"), .{ .x = @intCast(@divExact(width, 128)), .y = @intCast(items), .z = 1 }, 256, 0, &a); } { var a: cuda.Args = .{}; - a.add(sc.xd); a.add(w.dt); a.add(plan.items); a.add(plan.counts); a.add(plan.members); a.add(w.svh_d); a.add(y); - a.add(@as(i32, @intCast(width))); a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(ys))); + a.add(sc.xd); + a.add(w.dt); + a.add(plan.items); + a.add(plan.counts); + a.add(plan.members); + a.add(w.svh_d); + a.add(y); + a.add(@as(i32, @intCast(width))); + a.add(@as(i32, @intCast(D))); + a.add(@as(i32, @intCast(ys))); try k.go(try k.ext(.exl3, "glm_exl3_prompt_dn"), .{ .x = @intCast(@divExact(D, 256)), .y = @intCast(items), .z = 1 }, 256, 0, &a); } return; @@ -103,8 +134,17 @@ pub fn exl3Routed(k: *K, x: u64, x_stride: i64, picks: u64, plan: Plan, w: Exl3, // the staged kernel's grid: x = (column block, matrix, split), one item's blocks adjacent (its rows read from L2) { var a: cuda.Args = .{}; - a.add(sc.xg); a.add(if (rows_x) sc.xg else sc.xu); a.add(w.gt); a.add(w.ut); a.add(plan.items); a.add(plan.counts); a.add(plan.members); a.add(sc.z); - a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(width))); a.add(@as(i32, @intCast(P))); + a.add(sc.xg); + a.add(if (rows_x) sc.xg else sc.xu); + a.add(w.gt); + a.add(w.ut); + a.add(plan.items); + a.add(plan.counts); + a.add(plan.members); + a.add(sc.z); + a.add(@as(i32, @intCast(D))); + a.add(@as(i32, @intCast(width))); + a.add(@as(i32, @intCast(P))); if (staged) a.add(@as(i32, if (rows_x) 3 else 0)); a.add(@as(i32, @intCast(sk))); if (staged) a.add(@as(i32, 2)); @@ -114,14 +154,31 @@ pub fn exl3Routed(k: *K, x: u64, x_stride: i64, picks: u64, plan: Plan, w: Exl3, } { var a: cuda.Args = .{}; - a.add(sc.z); a.add(picks); a.add(w.svh_g); a.add(w.svh_u); a.add(w.suh_d); a.add(sc.xd); a.add(@as(i32, @intCast(P))); a.add(@as(i32, @intCast(width))); - a.add(@as(i32, @intCast(sk))); a.add(LIMIT); + a.add(sc.z); + a.add(picks); + a.add(w.svh_g); + a.add(w.svh_u); + a.add(w.suh_d); + a.add(sc.xd); + a.add(@as(i32, @intCast(P))); + a.add(@as(i32, @intCast(width))); + a.add(@as(i32, @intCast(sk))); + a.add(LIMIT); try k.go(try k.ext(.exl3, "glm_exl3_gateup"), .{ .x = @intCast(P), .y = @intCast(@divExact(width, 128)) }, 32, 0, &a); } { var a: cuda.Args = .{}; - a.add(sc.xd); a.add(sc.xd); a.add(w.dt); a.add(w.dt); a.add(plan.items); a.add(plan.counts); a.add(plan.members); a.add(sc.z); - a.add(@as(i32, @intCast(width))); a.add(@as(i32, @intCast(D))); a.add(@as(i32, @intCast(P))); + a.add(sc.xd); + a.add(sc.xd); + a.add(w.dt); + a.add(w.dt); + a.add(plan.items); + a.add(plan.counts); + a.add(plan.members); + a.add(sc.z); + a.add(@as(i32, @intCast(width))); + a.add(@as(i32, @intCast(D))); + a.add(@as(i32, @intCast(P))); if (staged) a.add(@as(i32, 0)); a.add(@as(i32, 1)); if (staged) a.add(@as(i32, 1)); @@ -131,7 +188,14 @@ pub fn exl3Routed(k: *K, x: u64, x_stride: i64, picks: u64, plan: Plan, w: Exl3, } { var a: cuda.Args = .{}; - a.add(sc.z); a.add(picks); a.add(w.svh_d); a.add(y); a.add(@as(i32, @intCast(P))); a.add(@as(i32, @intCast(D))); a.add(@as(i32, 1)); a.add(@as(i32, @intCast(ys))); + a.add(sc.z); + a.add(picks); + a.add(w.svh_d); + a.add(y); + a.add(@as(i32, @intCast(P))); + a.add(@as(i32, @intCast(D))); + a.add(@as(i32, 1)); + a.add(@as(i32, @intCast(ys))); try k.go(try k.ext(.exl3, "glm_exl3_down"), .{ .x = @intCast(P), .y = @intCast(@divExact(D, 128)) }, 32, 0, &a); } } diff --git a/zig/src/families/glm_cuda/native.zig b/zig/src/families/glm_cuda/native.zig index 811723e4e9..1ad4abd466 100644 --- a/zig/src/families/glm_cuda/native.zig +++ b/zig/src/families/glm_cuda/native.zig @@ -99,13 +99,9 @@ const Trace = struct { fn log(t: *Trace, own: *const Owned) void { const n: f64 = @floatFromInt(@max(t.rounds, 1)); const c: f64 = @floatFromInt(@max(t.calls, 1)); - std.log.info("glm trace: {d} verifies, {d:.2} ms each, {d:.2} rows ({d:.2} streams); {d} drafts, {d:.2} ms each, {d:.2} drafted a stream", - .{ t.rounds, t.verify_ms / n, @as(f64, @floatFromInt(t.verify_rows)) / n, @as(f64, @floatFromInt(t.verify_streams)) / n, - t.calls, t.draft_ms / c, @as(f64, @floatFromInt(t.drafts)) / @as(f64, @floatFromInt(@max(t.verify_streams, 1))) }); - std.log.info("glm drafts: {d} verified, {d} landed ({d:.2}), their chances said {d:.2}", .{ t.verified, t.landed, - @as(f64, @floatFromInt(t.landed)) / @as(f64, @floatFromInt(@max(t.verified, 1))), t.predicted / @as(f64, @floatFromInt(@max(t.verified, 1))) }); - std.log.info("glm copies: {d} of {d} draft rounds copied ({d:.2} tokens each), copied drafts land {d:.2}, the drafter's {d:.2}; {d} out of step", - .{ t.copies, t.calls, @as(f64, @floatFromInt(t.copy_tokens)) / @as(f64, @floatFromInt(@max(t.copies, 1))), own.copy_rate.chance(), own.draft_rate.chance(), t.desync }); + std.log.info("glm trace: {d} verifies, {d:.2} ms each, {d:.2} rows ({d:.2} streams); {d} drafts, {d:.2} ms each, {d:.2} drafted a stream", .{ t.rounds, t.verify_ms / n, @as(f64, @floatFromInt(t.verify_rows)) / n, @as(f64, @floatFromInt(t.verify_streams)) / n, t.calls, t.draft_ms / c, @as(f64, @floatFromInt(t.drafts)) / @as(f64, @floatFromInt(@max(t.verify_streams, 1))) }); + std.log.info("glm drafts: {d} verified, {d} landed ({d:.2}), their chances said {d:.2}", .{ t.verified, t.landed, @as(f64, @floatFromInt(t.landed)) / @as(f64, @floatFromInt(@max(t.verified, 1))), t.predicted / @as(f64, @floatFromInt(@max(t.verified, 1))) }); + std.log.info("glm copies: {d} of {d} draft rounds copied ({d:.2} tokens each), copied drafts land {d:.2}, the drafter's {d:.2}; {d} out of step", .{ t.copies, t.calls, @as(f64, @floatFromInt(t.copy_tokens)) / @as(f64, @floatFromInt(@max(t.copies, 1))), own.copy_rate.chance(), own.draft_rate.chance(), t.desync }); { const ds = &decode.draw_stats; const dc: f64 = @floatFromInt(@max(ds.calls, 1)); @@ -115,8 +111,7 @@ const Trace = struct { if (own.e.drafter) |dr| { const nc: f64 = @floatFromInt(@max(dr.calls, 1)); std.log.info("glm mix: {d} stream rounds by DFlash2, {d} by the MTP head", .{ t.df_rounds, t.drafts_calls_total - t.df_rounds }); - std.log.info("glm dflash2: {d} rounds, {d:.2} ms each (context {d:.2}, tables {d:.2}, layers {d:.2}, head {d:.2}, selector {d:.2}); {d} streams' contexts restarted", - .{ dr.calls, dr.forward_ms / nc, dr.phase_ms[0] / nc, dr.phase_ms[4] / nc, dr.phase_ms[1] / nc, dr.phase_ms[2] / nc, dr.phase_ms[3] / nc, dr.lost }); + std.log.info("glm dflash2: {d} rounds, {d:.2} ms each (context {d:.2}, tables {d:.2}, layers {d:.2}, head {d:.2}, selector {d:.2}); {d} streams' contexts restarted", .{ dr.calls, dr.forward_ms / nc, dr.phase_ms[0] / nc, dr.phase_ms[4] / nc, dr.phase_ms[1] / nc, dr.phase_ms[2] / nc, dr.phase_ms[3] / nc, dr.lost }); dflash.logFine(dr, dr.calls); dr.calls = 0; dr.forward_ms = 0; @@ -125,9 +120,7 @@ const Trace = struct { const g: f64 = 1 << 30; const avail: u64 = if (own.guard) |gd| gd.available.load(.monotonic) else 0; const low: u64 = if (own.guard) |gd| gd.lowest.load(.monotonic) else 0; - std.log.info("glm memory: {d:.2} GiB available (lowest {d:.2}); pool {d:.2} of {d:.2} GiB; {d} graphs ({d} captured, {d} updated in place, {d} dropped, {d:.0} ms capturing); device {d:.2} GiB", - .{ @as(f64, @floatFromInt(avail)) / g, @as(f64, @floatFromInt(low)) / g, @as(f64, @floatFromInt(own.used)) / g, @as(f64, @floatFromInt(own.room)) / g, - own.e.graphs.count(), own.e.captures, own.e.updates, own.e.evictions, own.e.capture_ms, @as(f64, @floatFromInt(cuda.usage(false).device)) / g }); + std.log.info("glm memory: {d:.2} GiB available (lowest {d:.2}); pool {d:.2} of {d:.2} GiB; {d} graphs ({d} captured, {d} updated in place, {d} dropped, {d:.0} ms capturing); device {d:.2} GiB", .{ @as(f64, @floatFromInt(avail)) / g, @as(f64, @floatFromInt(low)) / g, @as(f64, @floatFromInt(own.used)) / g, @as(f64, @floatFromInt(own.room)) / g, own.e.graphs.count(), own.e.captures, own.e.updates, own.e.evictions, own.e.capture_ms, @as(f64, @floatFromInt(cuda.usage(false).device)) / g }); t.* = .{ .on = true }; } }; @@ -214,7 +207,7 @@ pub const Owned = struct { } pub fn send(self: *Owned, op: Op, w: *Writer) !void { - try self.link.send(@intFromEnum(op), w.buf.items); + try self.link.send(@backingInt(op), w.buf.items); w.buf.clearRetainingCapacity(); } @@ -282,7 +275,7 @@ pub fn setRoom(p: *anyopaque, room: u64) !void { } const m = try self.link.recv(self.gpa); defer self.gpa.free(m.bytes); - if (m.tag != @intFromEnum(Op.room)) return error.RankOneNoRoom; + if (m.tag != @backingInt(Op.room)) return error.RankOneNoRoom; var r: Reader = .{ .bytes = m.bytes }; const peer = try r.int(); const both = @min(room, peer); @@ -291,7 +284,6 @@ pub fn setRoom(p: *anyopaque, room: u64) !void { std.log.info("glm: prompt chunks of {d} rows (TF_GLM_PREFILL_ROWS), the EXL3 experts' prompt kernel {s}", .{ self.e.pbuf.rows, @tagName(@import("kern_exl3.zig").promptKernel()) }); } - /// Commits the lanes' last windows: ``kept`` of their rows (the lane core calls keep only to cut a window short, so a /// window it accepted whole is committed here, before the lane's next call). Rank 1 commits the same. fn settle(self: *Owned, ls: []const *Lane, kept: []const u32) !void { @@ -341,8 +333,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ own.lanes_by = .init(gpa); own.seqs = .init(gpa); // the link and the communicator first, while memory is fresh (NCCL registers its buffers at its first use) - own.link = if (o.rank == 0) try cuda.tp_link.Link.lead(o.master_port, 600_000) - else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(o.master orelse return error.NoMaster, o.master_port), 600_000); + own.link = if (o.rank == 0) try cuda.tp_link.Link.lead(o.master_port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(o.master orelse return error.NoMaster, o.master_port), 600_000); errdefer own.link.close(); own.nccl = try cuda.nccl.Library.open(); errdefer own.nccl.close(); @@ -385,14 +376,14 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ if (mine[0] != theirs[0]) return error.DrafterOnOneRank; } // both halves loaded before rank 0 serves: rank 1 says so once its engine is up - if (o.rank == 1) try own.link.send(@intFromEnum(Op.ready), "") else { + if (o.rank == 1) try own.link.send(@backingInt(Op.ready), "") else { const m = try own.link.recv(gpa); defer gpa.free(m.bytes); - if (m.tag != @intFromEnum(Op.ready)) return error.RankOneNotReady; + if (m.tag != @backingInt(Op.ready)) return error.RankOneNotReady; } if (o.rank == 0) { // rank 1 replays it in its follow loop, then hears it ended try @import("calibrate.zig").run(own, o.drafts); - try own.link.send(@intFromEnum(Op.ready), ""); + try own.link.send(@backingInt(Op.ready), ""); if (own.e.drafter != null) if (std.c.getenv("TF_GLM_DFLASH2_STEP_MS") != null) { own.mtp_ms = dflashStepMs(); // calibrate timed a round; the variable sets another }; @@ -400,10 +391,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ const window = windowFromEnv(); return .{ .backend = .{ .ptr = own, .vtable = &vtable }, - .facts = .{ .exact_width = window, .mtp = o.drafts, .speculate = o.drafts, .speculate_early = false, .drafts = if (o.drafts) window - 1 else 1, - .hidden_rows = true, .max_streams = max_streams, .batch_rows = @intCast(forward.STEP_ROWS), - .window_costs = own.costs[0..own.cost_count], .mtp_step_ms = own.mtp_ms, .shared_costs = own.shared[0..own.shared_count], - .draft_probabilities = o.drafts, .draft_streams = o.drafts }, + .facts = .{ .exact_width = window, .mtp = o.drafts, .speculate = o.drafts, .speculate_early = false, .drafts = if (o.drafts) window - 1 else 1, .hidden_rows = true, .max_streams = max_streams, .batch_rows = @intCast(forward.STEP_ROWS), .window_costs = own.costs[0..own.cost_count], .mtp_step_ms = own.mtp_ms, .shared_costs = own.shared[0..own.shared_count], .draft_probabilities = o.drafts, .draft_streams = o.drafts }, .rows = if (o.drafts) window else 1, .stream_bytes = forward.seqBytes(own.e, admission_slots), .ctx = own, @@ -425,7 +413,7 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { fn release(p: *anyopaque) void { const own: *Owned = @ptrCast(@alignCast(p)); if (own.guard) |g| g.stop(own.gpa); // the peer's close that follows is no lost peer - if (own.rank == 0) own.link.send(@intFromEnum(Op.stop), "") catch {}; + if (own.rank == 0) own.link.send(@backingInt(Op.stop), "") catch {}; for (own.kept.items) |k| own.gpa.free(k.ids); own.kept.deinit(own.gpa); var it = own.lanes_by.valueIterator(); @@ -724,8 +712,7 @@ fn draftFn(p: *anyopaque, requests: []const be.DraftRequest) anyerror!void { if (self.e.drafter != null) chooseDrafter(self, l, if (r.rows != null) n else 0); const df = self.e.drafter != null and l.use_df; // a stream DFlash2 drafts for still feeds the MTP head its rows (one step, its draft unused): it stays current - reqs[i] = .{ .s = l.s, .hidden = if (first) l.s.last_hidden else l.s.hidden, .next = follow[i][0..n], .count = if (df) 1 else depth, - .confidence = self.confidence, .out = l.held[0..depth], .probs = l.probs[0..depth] }; + reqs[i] = .{ .s = l.s, .hidden = if (first) l.s.last_hidden else l.s.hidden, .next = follow[i][0..n], .count = if (df) 1 else depth, .confidence = self.confidence, .out = l.held[0..depth], .probs = l.probs[0..depth] }; try w.int(l.s.uid); try w.tokens(follow[i][0..n]); try w.int(depth); @@ -741,8 +728,7 @@ fn draftFn(p: *anyopaque, requests: []const be.DraftRequest) anyerror!void { var dq: [max_streams]dflash.Req = undefined; for (requests, 0..) |r, i| { const nf = reqs[i].next.len; - dq[i] = .{ .s = ls[i].s, .n = if (r.rows != null) nf else 0, .anchor = follow[i][nf - 1], .out = ls[i].held[0..max_depth], - .probs = ls[i].probs[0..max_depth], .drafts = ls[i].use_df }; + dq[i] = .{ .s = ls[i].s, .n = if (r.rows != null) nf else 0, .anchor = follow[i][nf - 1], .out = ls[i].held[0..max_depth], .probs = ls[i].probs[0..max_depth], .drafts = ls[i].use_df }; } try dflash.draft(dr, self.e, dq[0..requests.len]); for (reqs[0..requests.len], ls[0..requests.len], depths[0..requests.len]) |*q, l, depth| { // the confidence cut, as the head's chains stop @@ -934,7 +920,7 @@ fn followLoop(p: *anyopaque) anyerror!void { }; defer self.gpa.free(m.bytes); var r: Reader = .{ .bytes = m.bytes }; - switch (@as(Op, @enumFromInt(m.tag))) { + switch (@as(Op, @fromBackingInt(@intCast(m.tag)))) { .prefill => { const uid = try r.int(); const fresh: i64 = @intCast(try r.int()); @@ -991,8 +977,7 @@ fn followLoop(p: *anyopaque) anyerror!void { const first = try r.int() == 1; const cut: f64 = @bitCast(try r.int()); const df = try r.int() == 1; - reqs[i] = .{ .s = s, .hidden = if (first) s.last_hidden else s.hidden, .next = follow, .count = if (df) 1 else depth, .confidence = cut, - .out = held[i][0..depth], .probs = probs[i][0..depth] }; + reqs[i] = .{ .s = s, .hidden = if (first) s.last_hidden else s.hidden, .next = follow, .count = if (df) 1 else depth, .confidence = cut, .out = held[i][0..depth], .probs = probs[i][0..depth] }; dq[i] = .{ .s = s, .n = if (first) 0 else follow.len, .anchor = follow[follow.len - 1], .drafts = df }; } if (feedsMtp(self)) try decode.draft(self.e, reqs[0..n]); // as rank 0: the MTP head first, then DFlash2 diff --git a/zig/src/families/glm_cuda/seq.zig b/zig/src/families/glm_cuda/seq.zig index a5d0ffa4a6..d90c986118 100644 --- a/zig/src/families/glm_cuda/seq.zig +++ b/zig/src/families/glm_cuda/seq.zig @@ -105,8 +105,7 @@ pub const Seq = struct { const l: u64 = @intCast(li); const rows: u64 = @intCast(MAX_ROWS); const hw: u64 = @intCast(LL * 128); - return .{ .out = s.ks.out + l * rows * hw * 2, .k = s.ks.k + l * rows * hw * 4, .v = s.ks.v + l * rows * hw * 2, - .g = s.ks.g + l * rows * hw * 4, .b = s.ks.b + l * rows * @as(u64, @intCast(LL)) * 4 }; + return .{ .out = s.ks.out + l * rows * hw * 2, .k = s.ks.k + l * rows * hw * 4, .v = s.ks.v + l * rows * hw * 2, .g = s.ks.g + l * rows * hw * 4, .b = s.ks.b + l * rows * @as(u64, @intCast(LL)) * 4 }; } }; @@ -139,8 +138,7 @@ pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, nkda: usize, ndsa: u s.conv = c.take(n * 3 * CONV_C * 2); s.rec = c.take(2 * n * REC * 4); s.proj = c.take(n * MAX_ROWS * PROJ_W * 2); - s.ks = .{ .out = c.take(n * MAX_ROWS * LL * 128 * 2), .k = c.take(n * MAX_ROWS * LL * 128 * 4), .v = c.take(n * MAX_ROWS * LL * 128 * 2), - .g = c.take(n * MAX_ROWS * LL * 128 * 4), .b = c.take(n * MAX_ROWS * LL * 4) }; + s.ks = .{ .out = c.take(n * MAX_ROWS * LL * 128 * 2), .k = c.take(n * MAX_ROWS * LL * 128 * 4), .v = c.take(n * MAX_ROWS * LL * 128 * 2), .g = c.take(n * MAX_ROWS * LL * 128 * 4), .b = c.take(n * MAX_ROWS * LL * 4) }; for (0..caches) |i| { s.kc[i] = c.take(capacity * kern.KV_LORA * 2); s.ik[i] = c.take(s.ring * kern.IDX_DIM * 2); diff --git a/zig/src/families/glm_cuda/vision.zig b/zig/src/families/glm_cuda/vision.zig index 8eb6c15bc9..9a273d7e31 100644 --- a/zig/src/families/glm_cuda/vision.zig +++ b/zig/src/families/glm_cuda/vision.zig @@ -39,7 +39,12 @@ pub const Tower = struct { post: u64 = 0, down_w: u64 = 0, // [4096, 2, 2, 1024]: the conv's taps in the merge block's (row, col) order, then the channels down_b: u64 = 0, - m_proj: u64 = 0, m_ln_w: u64 = 0, m_ln_b: u64 = 0, m_gate: u64 = 0, m_up: u64 = 0, m_down: u64 = 0, + m_proj: u64 = 0, + m_ln_w: u64 = 0, + m_ln_b: u64 = 0, + m_gate: u64 = 0, + m_up: u64 = 0, + m_down: u64 = 0, inv_freq: cuda.DeviceBuffer, linears: std.AutoHashMap(LinKey, cuda.cublaslt.Linear), ws: cuda.DeviceBuffer, @@ -100,8 +105,7 @@ pub const Tower = struct { total += std.mem.alignForward(u64, ent.n, 256); } } - var t: Tower = .{ .gpa = gpa, .d = d, .lt = try cuda.cublaslt.Library.open(), .mem = try cuda.DeviceBuffer.alloc(d, total), - .inv_freq = try cuda.DeviceBuffer.alloc(d, 16 * 4), .linears = .init(gpa), .ws = undefined, .bytes = total }; + var t: Tower = .{ .gpa = gpa, .d = d, .lt = try cuda.cublaslt.Library.open(), .mem = try cuda.DeviceBuffer.alloc(d, total), .inv_freq = try cuda.DeviceBuffer.alloc(d, 16 * 4), .linears = .init(gpa), .ws = undefined, .bytes = total }; errdefer t.deinit(); t.ws = try cuda.DeviceBuffer.alloc(d, t.ws_len); var at: u64 = 0; @@ -157,13 +161,20 @@ pub const Tower = struct { t.m_down = try get(&names, "merger.down_proj.weight", .{}); for (&t.blocks, 0..) |*bl, i| { bl.* = .{ - .n1 = try get(&names, "blocks.{d}.norm1.weight", .{i}), .n2 = try get(&names, "blocks.{d}.norm2.weight", .{i}), - .qkv_w = try get(&names, "blocks.{d}.attn.qkv.weight", .{i}), .qkv_b = try get(&names, "blocks.{d}.attn.qkv.bias", .{i}), - .qn = try get(&names, "blocks.{d}.attn.q_norm.weight", .{i}), .kn = try get(&names, "blocks.{d}.attn.k_norm.weight", .{i}), - .proj_w = try get(&names, "blocks.{d}.attn.proj.weight", .{i}), .proj_b = try get(&names, "blocks.{d}.attn.proj.bias", .{i}), - .g_w = try get(&names, "blocks.{d}.mlp.gate_proj.weight", .{i}), .g_b = try get(&names, "blocks.{d}.mlp.gate_proj.bias", .{i}), - .u_w = try get(&names, "blocks.{d}.mlp.up_proj.weight", .{i}), .u_b = try get(&names, "blocks.{d}.mlp.up_proj.bias", .{i}), - .d_w = try get(&names, "blocks.{d}.mlp.down_proj.weight", .{i}), .d_b = try get(&names, "blocks.{d}.mlp.down_proj.bias", .{i}), + .n1 = try get(&names, "blocks.{d}.norm1.weight", .{i}), + .n2 = try get(&names, "blocks.{d}.norm2.weight", .{i}), + .qkv_w = try get(&names, "blocks.{d}.attn.qkv.weight", .{i}), + .qkv_b = try get(&names, "blocks.{d}.attn.qkv.bias", .{i}), + .qn = try get(&names, "blocks.{d}.attn.q_norm.weight", .{i}), + .kn = try get(&names, "blocks.{d}.attn.k_norm.weight", .{i}), + .proj_w = try get(&names, "blocks.{d}.attn.proj.weight", .{i}), + .proj_b = try get(&names, "blocks.{d}.attn.proj.bias", .{i}), + .g_w = try get(&names, "blocks.{d}.mlp.gate_proj.weight", .{i}), + .g_b = try get(&names, "blocks.{d}.mlp.gate_proj.bias", .{i}), + .u_w = try get(&names, "blocks.{d}.mlp.up_proj.weight", .{i}), + .u_b = try get(&names, "blocks.{d}.mlp.up_proj.bias", .{i}), + .d_w = try get(&names, "blocks.{d}.mlp.down_proj.weight", .{i}), + .d_b = try get(&names, "blocks.{d}.mlp.down_proj.bias", .{i}), }; } // the rotary's frequencies: 1 / 10000^(i / 16), i = 0, 2 .. 30 over 32 (head width / 2) @@ -207,8 +218,12 @@ pub const Tower = struct { const rows = @min(ROWS, m - r); try t.linear(k, x + @as(u64, @intCast(r * kk * 2)), w, acc, rows, n, kk, true); var a: cuda.Args = .{}; - a.add(acc); a.add(bias); a.add(if (out != 0) out + @as(u64, @intCast(r * n * 2)) else 0); - a.add(if (x32 != 0) x32 + @as(u64, @intCast(r * n * 4)) else 0); a.add(@as(i64, rows * n)); a.add(@as(i32, @intCast(n))); + a.add(acc); + a.add(bias); + a.add(if (out != 0) out + @as(u64, @intCast(r * n * 2)) else 0); + a.add(if (x32 != 0) x32 + @as(u64, @intCast(r * n * 4)) else 0); + a.add(@as(i64, rows * n)); + a.add(@as(i32, @intCast(n))); try k.go(try k.ext(.vision, "glm_vis_bias"), grid1(rows * n), 256, 0, &a); } } @@ -223,15 +238,25 @@ pub const Tower = struct { try t.linear(k, xr, gw, acc, rows, n, kk, true); try t.linear(k, xr, uw, acc2, rows, n, kk, true); var a: cuda.Args = .{}; - a.add(acc); a.add(acc2); a.add(gb); a.add(ub); a.add(out + @as(u64, @intCast(r * n * 2))); a.add(@as(i64, rows * n)); - a.add(@as(i32, @intCast(n))); a.add(LIMIT); + a.add(acc); + a.add(acc2); + a.add(gb); + a.add(ub); + a.add(out + @as(u64, @intCast(r * n * 2))); + a.add(@as(i64, rows * n)); + a.add(@as(i32, @intCast(n))); + a.add(LIMIT); try k.go(try k.ext(.vision, "glm_vis_swiglu"), grid1(rows * n), 256, 0, &a); } } fn rmsnorm(k: *kern.K, x32: u64, w: u64, y: u64, rows: i64, cols: i64) !void { var a: cuda.Args = .{}; - a.add(x32); a.add(w); a.add(y); a.add(@as(i32, @intCast(cols))); a.add(EPS); + a.add(x32); + a.add(w); + a.add(y); + a.add(@as(i32, @intCast(cols))); + a.add(EPS); try k.go(try k.ext(.vision, "glm_vis_rmsnorm"), .{ .x = @intCast(rows) }, 256, 0, &a); } @@ -296,14 +321,20 @@ pub const Tower = struct { try d.check(d.api.cuMemcpyHtoDAsync_v2(p32, patches.ptr, patches.len * 4, k.stream.handle), "vision patches"); { // the patch embedding: the patches as bf16, then the conv as a linear (with its bias) into the fp32 stream var a: cuda.Args = .{}; - a.add(p32); a.add(big); a.add(@as(i64, N * PATCH)); + a.add(p32); + a.add(big); + a.add(@as(i64, N * PATCH)); try k.go(try k.ext(.vision, "glm_vis_to_bf16"), grid1(N * PATCH), 256, 0, &a); var r: i64 = 0; while (r < N) : (r += ROWS) { const rows = @min(ROWS, N - r); try t.linear(k, big + @as(u64, @intCast(r * PATCH * 2)), t.patch_w, acc, rows, H, PATCH, true); var b: cuda.Args = .{}; - b.add(acc); b.add(t.patch_b); b.add(x32 + @as(u64, @intCast(r * H * 4))); b.add(@as(i64, rows * H)); b.add(@as(i32, @intCast(H))); + b.add(acc); + b.add(t.patch_b); + b.add(x32 + @as(u64, @intCast(r * H * 4))); + b.add(@as(i64, rows * H)); + b.add(@as(i32, @intCast(H))); try k.go(try k.ext(.vision, "glm_vis_embed32"), grid1(rows * H), 256, 0, &b); } } @@ -326,8 +357,18 @@ pub const Tower = struct { try t.dense(k, xn, bl.qkv_w, bl.qkv_b, big, 0, acc, N, 3 * H, H); { var a: cuda.Args = .{}; - a.add(big); a.add(bl.qn); a.add(bl.kn); a.add(dpos); a.add(t.inv_freq.ptr); a.add(qh); a.add(kh); a.add(vt); - a.add(@as(i32, @intCast(N))); a.add(@as(i32, @intCast(seg))); a.add(@as(i32, @intCast(HEADS))); a.add(EPS); + a.add(big); + a.add(bl.qn); + a.add(bl.kn); + a.add(dpos); + a.add(t.inv_freq.ptr); + a.add(qh); + a.add(kh); + a.add(vt); + a.add(@as(i32, @intCast(N))); + a.add(@as(i32, @intCast(seg))); + a.add(@as(i32, @intCast(HEADS))); + a.add(EPS); try k.go(try k.ext(.vision, "glm_vis_qkv"), .{ .x = @intCast(std.math.divCeil(i64, N * HEADS, 8) catch unreachable) }, 256, 0, &a); } // each frame group attends within itself; per head, chunks of queries (q, k: [t, H, seg, 64]; v: [t, H, 64, seg]) @@ -341,12 +382,20 @@ pub const Tower = struct { const rows = @min(CHUNK, seg - c); try t.linear(k, qh + blk + @as(u64, @intCast(c * HD * 2)), kh + blk, sc, rows, seg, HD, true); var a: cuda.Args = .{}; - a.add(sc); a.add(pr); a.add(sums); a.add(@as(i32, @intCast(seg))); a.add(scale); + a.add(sc); + a.add(pr); + a.add(sums); + a.add(@as(i32, @intCast(seg))); + a.add(scale); try k.go(try k.ext(.vision, "glm_vis_softmax"), .{ .x = @intCast(rows) }, 256, 0, &a); try t.linear(k, pr, vt + blk, oh, rows, HD, seg, true); var b: cuda.Args = .{}; - b.add(oh); b.add(sums); b.add(att + @as(u64, @intCast((g * seg + c) * H * 2))); b.add(@as(i32, @intCast(rows))); - b.add(@as(i32, @intCast(HEADS))); b.add(@as(i32, @intCast(h))); + b.add(oh); + b.add(sums); + b.add(att + @as(u64, @intCast((g * seg + c) * H * 2))); + b.add(@as(i32, @intCast(rows))); + b.add(@as(i32, @intCast(HEADS))); + b.add(@as(i32, @intCast(h))); try k.go(try k.ext(.vision, "glm_vis_head_out"), grid1(rows * HD), 256, 0, &b); } } @@ -374,7 +423,12 @@ pub const Tower = struct { const rows = @min(ROWS, M - r); try t.linear(k, att + @as(u64, @intCast(r * OUT * 2)), t.m_proj, acc, rows, OUT, OUT, true); var a: cuda.Args = .{}; - a.add(acc); a.add(t.m_ln_w); a.add(t.m_ln_b); a.add(xn + @as(u64, @intCast(r * OUT * 2))); a.add(@as(i32, @intCast(OUT))); a.add(EPS); + a.add(acc); + a.add(t.m_ln_w); + a.add(t.m_ln_b); + a.add(xn + @as(u64, @intCast(r * OUT * 2))); + a.add(@as(i32, @intCast(OUT))); + a.add(EPS); try k.go(try k.ext(.vision, "glm_vis_ln_gelu"), .{ .x = @intCast(rows) }, 256, 0, &a); } r = 0; diff --git a/zig/src/families/glm_cuda/weights.zig b/zig/src/families/glm_cuda/weights.zig index 371493428f..5a1f36207e 100644 --- a/zig/src/families/glm_cuda/weights.zig +++ b/zig/src/families/glm_cuda/weights.zig @@ -72,8 +72,7 @@ const Ckpt = struct { const shape = try a.alloc(i64, shp.len); for (shp, 0..) |x, i| shape[i] = x.integer; const offs = o.get("data_offsets").?.array.items; - try ck.entries.put(e.key_ptr.*, .{ .dtype = o.get("dtype").?.string, .shape = shape, - .begin = 8 + n + @as(u64, @intCast(offs[0].integer)), .end = 8 + n + @as(u64, @intCast(offs[1].integer)), .file = f }); + try ck.entries.put(e.key_ptr.*, .{ .dtype = o.get("dtype").?.string, .shape = shape, .begin = 8 + n + @as(u64, @intCast(offs[0].integer)), .end = 8 + n + @as(u64, @intCast(offs[1].integer)), .file = f }); } } @@ -174,8 +173,7 @@ fn splitOf(name: []const u8) Split { if (has(name, ".mlp.experts.") or has(name, ".mlp.shared_experts.") or has(name, ".mlp.gate_proj.") or has(name, ".mlp.up_proj.") or has(name, ".mlp.down_proj.")) return if (has(name, "down_proj.")) .col else .row; if (has(name, ".self_attn.indexer.")) return .rep; - for ([_][]const u8{ ".self_attn.q_proj.", ".self_attn.k_proj.", ".self_attn.v_proj.", "_conv1d.", ".self_attn.f_b_proj.", ".self_attn.g_b_proj.", - ".self_attn.b_proj.", ".self_attn.q_b_proj.", ".self_attn.kv_b_proj." }) |f| if (has(name, f)) return .row; + for ([_][]const u8{ ".self_attn.q_proj.", ".self_attn.k_proj.", ".self_attn.v_proj.", "_conv1d.", ".self_attn.f_b_proj.", ".self_attn.g_b_proj.", ".self_attn.b_proj.", ".self_attn.q_b_proj.", ".self_attn.kv_b_proj." }) |f| if (has(name, f)) return .row; if (ends(u8, name, ".self_attn.A_log") or ends(u8, name, ".self_attn.dt_bias")) return .row; if (has(name, ".self_attn.o_proj.")) return .col; return .rep; @@ -392,8 +390,7 @@ fn readExperts(b: *Builder, li: usize, proj: []const u8, count: usize) !Triple { const okg = if (down) kg_full / WORLD else kg_full; const c8 = if (down) b.rank * ok8 else 0; const cg = if (down) b.rank * okg else 0; - var t: Triple = .{ .e = count, .n = nr, .k8 = ok8, .kg = okg, .words = try b.gpa.alloc(u32, count * nr * ok8), - .scales = try b.gpa.alloc(u16, count * nr * okg), .biases = try b.gpa.alloc(u16, count * nr * okg) }; + var t: Triple = .{ .e = count, .n = nr, .k8 = ok8, .kg = okg, .words = try b.gpa.alloc(u32, count * nr * ok8), .scales = try b.gpa.alloc(u16, count * nr * okg), .biases = try b.gpa.alloc(u16, count * nr * okg) }; errdefer t.free(b.gpa); const rowbuf = try b.gpa.alloc(u8, nr * k8_full * 4); defer b.gpa.free(rowbuf); @@ -569,8 +566,7 @@ fn expertsExl3(b: *Builder, out: []const u8, li: usize) !void { const tile: usize = 128; // a 16x16 tile's bytes at 4 bits var nb: [256]u8 = undefined; const Mat = struct { proj: []const u8, kt: usize, nt: usize, by_cols: bool }; - const mats = [_]Mat{ .{ .proj = "gate_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, .{ .proj = "up_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, - .{ .proj = "down_proj", .kt = width / 16, .nt = D / 16, .by_cols = false } }; + const mats = [_]Mat{ .{ .proj = "gate_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, .{ .proj = "up_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, .{ .proj = "down_proj", .kt = width / 16, .nt = D / 16, .by_cols = false } }; const short = [_][]const u8{ "g", "u", "d" }; for (mats, short) |m, sh| { const tshape = [_]i64{ @intCast(E), @intCast(m.kt), @intCast(m.nt), 32 }; @@ -612,8 +608,7 @@ fn expertsExl3(b: *Builder, out: []const u8, li: usize) !void { try b.ck.readAt(se, if (m.by_cols) 0 else b.rank * su.len, su); try b.ck.readAt(ve, if (m.by_cols) b.rank * sv.len else 0, sv); if (m.by_cols) { // gate and up: one input rotation shared by every expert of the layer? (exl3_shx) - if (std.mem.eql(u8, sh, "g") and x == 0) @memcpy(b.suh_first[0..su.len], su) - else if (!std.mem.eql(u8, su, b.suh_first[0..su.len])) b.exl3_shx = false; + if (std.mem.eql(u8, sh, "g") and x == 0) @memcpy(b.suh_first[0..su.len], su) else if (!std.mem.eql(u8, su, b.suh_first[0..su.len])) b.exl3_shx = false; } } try b.put(tn, .i32, &tshape, words); @@ -676,8 +671,7 @@ fn dsa(b: *Builder, out: []const u8, i: usize) !void { var o: [160]u8 = undefined; var f: [256]u8 = undefined; var f2: [256]u8 = undefined; - try b.stack(try std.fmt.bufPrint(&o, "{s}.proj", .{out}), &.{ try fullName(&f, "layers.{d}.self_attn.q_a_proj.weight", .{i}), - try fullName(&f2, "layers.{d}.self_attn.kv_a_proj_with_mqa.weight", .{i}) }); + try b.stack(try std.fmt.bufPrint(&o, "{s}.proj", .{out}), &.{ try fullName(&f, "layers.{d}.self_attn.q_a_proj.weight", .{i}), try fullName(&f2, "layers.{d}.self_attn.kv_a_proj_with_mqa.weight", .{i}) }); try b.tensor(try std.fmt.bufPrint(&o, "{s}.q_norm", .{out}), try fullName(&f, "layers.{d}.self_attn.q_a_layernorm.weight", .{i}), false); try b.tensor(try std.fmt.bufPrint(&o, "{s}.kv_norm", .{out}), try fullName(&f, "layers.{d}.self_attn.kv_a_layernorm.weight", .{i}), false); try b.linear(try std.fmt.bufPrint(&o, "{s}.q_b", .{out}), try fullName(&f, "layers.{d}.self_attn.q_b_proj.weight", .{i})); @@ -707,8 +701,7 @@ fn dsa(b: *Builder, out: []const u8, i: usize) !void { } } try b.linear(try std.fmt.bufPrint(&o, "{s}.o", .{out}), try fullName(&f, "layers.{d}.self_attn.o_proj.weight", .{i})); - try b.stack(try std.fmt.bufPrint(&o, "{s}.index.kw", .{out}), &.{ try fullName(&f, "layers.{d}.self_attn.indexer.wk.weight", .{i}), - try fullName(&f2, "layers.{d}.self_attn.indexer.weights_proj.weight", .{i}) }); + try b.stack(try std.fmt.bufPrint(&o, "{s}.index.kw", .{out}), &.{ try fullName(&f, "layers.{d}.self_attn.indexer.wk.weight", .{i}), try fullName(&f2, "layers.{d}.self_attn.indexer.weights_proj.weight", .{i}) }); try b.linear(try std.fmt.bufPrint(&o, "{s}.index.qb", .{out}), try fullName(&f, "layers.{d}.self_attn.indexer.wq_b.weight", .{i})); try b.tensor(try std.fmt.bufPrint(&o, "{s}.index.ln_w", .{out}), try fullName(&f, "layers.{d}.self_attn.indexer.k_norm.weight", .{i}), false); try b.tensor(try std.fmt.bufPrint(&o, "{s}.index.ln_b", .{out}), try fullName(&f, "layers.{d}.self_attn.indexer.k_norm.bias", .{i}), false); @@ -720,8 +713,7 @@ fn mlp(b: *Builder, out: []const u8, prefix: []const u8) !void { var o: [160]u8 = undefined; var f: [256]u8 = undefined; var f2: [256]u8 = undefined; - try b.stack(try std.fmt.bufPrint(&o, "{s}.gu", .{out}), &.{ try std.fmt.bufPrint(&f, "{s}gate_proj.weight", .{prefix}), - try std.fmt.bufPrint(&f2, "{s}up_proj.weight", .{prefix}) }); + try b.stack(try std.fmt.bufPrint(&o, "{s}.gu", .{out}), &.{ try std.fmt.bufPrint(&f, "{s}gate_proj.weight", .{prefix}), try std.fmt.bufPrint(&f2, "{s}up_proj.weight", .{prefix}) }); try b.linear(try std.fmt.bufPrint(&o, "{s}.down", .{out}), try std.fmt.bufPrint(&f, "{s}down_proj.weight", .{prefix})); } @@ -775,8 +767,7 @@ pub const Cfg = struct { } const t = v.object.get("text_config").?.object; const n: usize = @intCast(t.get("num_hidden_layers").?.integer); - var c: Cfg = .{ .exl3 = exl3, .layers = n, .dense = @intCast(t.get("first_k_dense_replace").?.integer), .mtp = (if (t.get("num_nextn_predict_layers")) |x| x.integer else 0) > 0, - .vocab = @intCast(t.get("vocab_size").?.integer) }; + var c: Cfg = .{ .exl3 = exl3, .layers = n, .dense = @intCast(t.get("first_k_dense_replace").?.integer), .mtp = (if (t.get("num_nextn_predict_layers")) |x| x.integer else 0) > 0, .vocab = @intCast(t.get("vocab_size").?.integer) }; if (n > 64) return error.TooManyLayers; for (t.get("layer_types").?.array.items, 0..) |x, i| c.kinds[i] = std.mem.eql(u8, x.string, "linear_attention"); if (t.get("hidden_size").?.integer != kern.D or t.get("n_routed_experts").?.integer != kern.EXPERTS or t.get("num_experts_per_tok").?.integer != kern.TOPK or @@ -827,16 +818,15 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ctx: *const api.Ctx, dir: []cons const c = try Cfg.read(arena.allocator(), io, dir); var store = api.Store.init(gpa); errdefer store.deinit(); - var b: Builder = .{ .gpa = gpa, .ck = &ck, .dry = false, .rank = rank, .store = &store, .d = ctx.d, - .dense_bits = try envBits("TF_GLM_DENSE_BITS"), .head_bits = try envBits("TF_GLM_HEAD_BITS"), .draft_head_bits = try envBits("TF_GLM_DRAFT_HEAD_BITS"), .kda_proj_bits = try envBits("TF_GLM_KDA_PROJ_BITS"), .experts_sym = envSym(), .expert_sim_bits = envSimBits(), .experts3 = envBits3(), .sim_proj = envSimProj() }; + var b: Builder = .{ .gpa = gpa, .ck = &ck, .dry = false, .rank = rank, .store = &store, .d = ctx.d, .dense_bits = try envBits("TF_GLM_DENSE_BITS"), .head_bits = try envBits("TF_GLM_HEAD_BITS"), .draft_head_bits = try envBits("TF_GLM_DRAFT_HEAD_BITS"), .kda_proj_bits = try envBits("TF_GLM_KDA_PROJ_BITS"), .experts_sym = envSym(), .expert_sim_bits = envSimBits(), .experts3 = envBits3(), .sim_proj = envSimProj() }; b.exl3 = c.exl3; if (c.exl3) b.experts_sym = false; try build(&b, &c); const vl = c.vocab / WORLD; const facts = [_]struct { []const u8, i64 }{ - .{ "rank", rank }, .{ "world", WORLD }, .{ "vocab", @intCast(c.vocab) }, .{ "vocab_shard", @intCast(vl) }, - .{ "vocab_offset", @intCast(rank * vl) }, .{ "layers", @intCast(c.layers) }, .{ "dense", @intCast(c.dense) }, - .{ "mtp", @intFromBool(c.mtp) }, .{ "device_bytes", @intCast(b.bytes) }, .{ "experts_sym", @intFromBool(b.experts_sym) }, .{ "experts3", @intFromBool(b.experts3) }, .{ "exl3", @intFromBool(b.exl3) }, + .{ "rank", rank }, .{ "world", WORLD }, .{ "vocab", @intCast(c.vocab) }, .{ "vocab_shard", @intCast(vl) }, + .{ "vocab_offset", @intCast(rank * vl) }, .{ "layers", @intCast(c.layers) }, .{ "dense", @intCast(c.dense) }, .{ "mtp", @intFromBool(c.mtp) }, + .{ "device_bytes", @intCast(b.bytes) }, .{ "experts_sym", @intFromBool(b.experts_sym) }, .{ "experts3", @intFromBool(b.experts3) }, .{ "exl3", @intFromBool(b.exl3) }, .{ "exl3_shx", @intFromBool(b.exl3 and b.exl3_shx and !b.dry) }, }; for (facts) |fct| try store.ints.put(try gpa.dupe(u8, fct[0]), fct[1]); @@ -854,8 +844,7 @@ pub fn deviceBytes(io: std.Io, dir: []const u8) !u64 { var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); const c = try Cfg.read(arena.allocator(), io, dir); - var b: Builder = .{ .gpa = gpa, .ck = &ck, .dry = true, .rank = 0, .dense_bits = try envBits("TF_GLM_DENSE_BITS"), - .head_bits = try envBits("TF_GLM_HEAD_BITS"), .draft_head_bits = try envBits("TF_GLM_DRAFT_HEAD_BITS"), .kda_proj_bits = try envBits("TF_GLM_KDA_PROJ_BITS"), .experts_sym = envSym(), .expert_sim_bits = envSimBits(), .experts3 = envBits3(), .sim_proj = envSimProj() }; + var b: Builder = .{ .gpa = gpa, .ck = &ck, .dry = true, .rank = 0, .dense_bits = try envBits("TF_GLM_DENSE_BITS"), .head_bits = try envBits("TF_GLM_HEAD_BITS"), .draft_head_bits = try envBits("TF_GLM_DRAFT_HEAD_BITS"), .kda_proj_bits = try envBits("TF_GLM_KDA_PROJ_BITS"), .experts_sym = envSym(), .expert_sim_bits = envSimBits(), .experts3 = envBits3(), .sim_proj = envSimProj() }; b.exl3 = c.exl3; try build(&b, &c); return b.bytes; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 93710b3088..40b5ba3101 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -406,7 +406,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O // 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 "", + 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 if (@hasField(@TypeOf(loaded), "target")) ", no prompt cache" else ", prompt states kept by the family", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); diff --git a/zig/tests/cuda/glm_forward.zig b/zig/tests/cuda/glm_forward.zig index 860a47b232..b9e1cafc96 100644 --- a/zig/tests/cuda/glm_forward.zig +++ b/zig/tests/cuda/glm_forward.zig @@ -425,8 +425,7 @@ pub fn batch(gpu: Gpu, args: []const [:0]const u8) !void { const rank = try std.fmt.parseInt(u8, args[0], 10); const port = try std.fmt.parseInt(u16, args[2], 10); const ntok: usize = if (args.len > 4) try std.fmt.parseInt(usize, args[4], 10) else 48; - var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) - else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); + var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); defer link.close(); var nccl = try cuda.nccl.Library.open(); defer nccl.close(); @@ -482,8 +481,7 @@ pub fn batch(gpu: Gpu, args: []const [:0]const u8) !void { const t_ms = ms(io, ts); var equal: usize = 0; for (all, 0..) |x, i| equal += @intFromBool(std.mem.eql(u32, x.out.items[0..ntok], solo[i][0..ntok])); - std.debug.print("rank {d}: 16 streams together, {s}: {d} of 16 replies equal to alone, {d:.1} tok/s aggregate\n", .{ rank, - if (depth > 0) "drafted" else "serial", equal, @as(f64, @floatFromInt(S * (ntok - 1))) / t_ms * 1000 }); + std.debug.print("rank {d}: 16 streams together, {s}: {d} of 16 replies equal to alone, {d:.1} tok/s aggregate\n", .{ rank, if (depth > 0) "drafted" else "serial", equal, @as(f64, @floatFromInt(S * (ntok - 1))) / t_ms * 1000 }); try check.expect(equal == S, "16 streams together equal alone", .{}); if (depth > 0) { // resumed: stream 0's kept state at its prompt's end - 1, a longer prompt extending it const ext = try std.mem.concat(gpa, u32, &.{ prompts[0], solo[0][0..8] }); @@ -526,8 +524,7 @@ pub fn teacher(gpu: Gpu, args: []const [:0]const u8) !void { if (args.len < 5) return error.Usage; const rank = try std.fmt.parseInt(u8, args[0], 10); const port = try std.fmt.parseInt(u16, args[2], 10); - var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) - else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); + var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); defer link.close(); var nccl = try cuda.nccl.Library.open(); defer nccl.close(); @@ -586,8 +583,7 @@ pub fn prof(gpu: Gpu, args: []const [:0]const u8) !void { const port = try std.fmt.parseInt(u16, args[2], 10); const nstreams: usize = if (args.len > 4) try std.fmt.parseInt(usize, args[4], 10) else 1; const graphs_only = args.len > 5 and std.mem.eql(u8, args[5], "graphs"); // for nsys: 50 one-row steps, no eager pass - var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) - else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); + var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); defer link.close(); var nccl = try cuda.nccl.Library.open(); defer nccl.close(); diff --git a/zig/tests/cuda/glm_weights.zig b/zig/tests/cuda/glm_weights.zig index ad1c372306..781ac43132 100644 --- a/zig/tests/cuda/glm_weights.zig +++ b/zig/tests/cuda/glm_weights.zig @@ -24,8 +24,7 @@ pub fn run(gpu: Gpu, dir: []const u8, rank_s: []const u8, ref: []const u8) !void var store = try gw.load(gpa, gpu.io, &ctx, dir, rank); defer store.deinit(); const secs = @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(gpu.io, .awake)).toNanoseconds())) / 1e9; - std.debug.print("loaded {d} tensors, {d:.2} GiB (planned {d:.2} GiB) in {d:.1} s\n", .{ store.map.count(), - @as(f64, @floatFromInt(store.ints.get("device_bytes").?)) / (1 << 30), @as(f64, @floatFromInt(planned)) / (1 << 30), secs }); + std.debug.print("loaded {d} tensors, {d:.2} GiB (planned {d:.2} GiB) in {d:.1} s\n", .{ store.map.count(), @as(f64, @floatFromInt(store.ints.get("device_bytes").?)) / (1 << 30), @as(f64, @floatFromInt(planned)) / (1 << 30), secs }); const jp = try std.fs.path.join(gpa, &.{ ref, "weights.json" }); const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, jp, gpa, .limited(1 << 27)); const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); 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 09/31] 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 a76651ad5ed7be9927d6f0c3f96d49322dcf766e 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 10/31] glm cuda: comments one line each, files under 600 lines, the long headers in docs/glm-cuda-notes.md Splits move code unchanged (only imports, aliases and pub change): kern.zig into kern_attn, kern_kda and kern_moe; weights.zig into ckpt, experts_load and qpack; dflash.zig into dflash_load and dflash_select; native.zig into native_draft, native_prompt, native_follow and native_env; exl3.cu's transposed prompt design into exl3_transposed.cuh; tf-cuda-test glm-tf and glm-prof into glm_forward_run.zig. --- docs/glm-cuda-notes.md | 103 +++ zig/build/cuda.zig | 4 +- zig/kernels/cuda/flashnext/experts.cu | 3 +- zig/kernels/cuda/flashnext/experts.cuh | 3 +- zig/kernels/cuda/flashnext/experts_prefill.cu | 3 +- zig/kernels/cuda/flashnext/sample.cu | 20 +- zig/kernels/cuda/glm/dflash.cu | 37 +- zig/kernels/cuda/glm/exl3.cu | 234 +------ zig/kernels/cuda/glm/exl3_transposed.cuh | 168 +++++ zig/kernels/cuda/glm/experts_sym.cu | 3 +- zig/kernels/cuda/glm/experts_sym.cuh | 7 +- zig/kernels/cuda/glm/experts_sym_prefill.cu | 4 +- zig/kernels/cuda/glm/index.cu | 23 +- zig/kernels/cuda/glm/kda.cu | 23 +- zig/kernels/cuda/glm/latent.cu | 15 +- zig/kernels/cuda/glm/mm.cu | 16 +- zig/kernels/cuda/glm/ops.cu | 17 +- zig/kernels/cuda/glm/qmm_frag.cuh | 6 +- zig/kernels/cuda/glm/rows.cu | 32 +- zig/kernels/cuda/glm/vision.cu | 9 +- zig/src/families/glm_cuda/api.zig | 8 +- zig/src/families/glm_cuda/blocks.zig | 24 +- zig/src/families/glm_cuda/calibrate.zig | 6 +- zig/src/families/glm_cuda/ckpt.zig | 151 +++++ zig/src/families/glm_cuda/copy.zig | 9 +- zig/src/families/glm_cuda/decode.zig | 34 +- zig/src/families/glm_cuda/dflash.zig | 421 +----------- zig/src/families/glm_cuda/dflash_load.zig | 315 +++++++++ zig/src/families/glm_cuda/dflash_select.zig | 63 ++ zig/src/families/glm_cuda/draw.zig | 9 +- zig/src/families/glm_cuda/experts_load.zig | 331 ++++++++++ zig/src/families/glm_cuda/forward.zig | 38 +- zig/src/families/glm_cuda/guard.zig | 7 +- zig/src/families/glm_cuda/kern.zig | 412 ++---------- zig/src/families/glm_cuda/kern_attn.zig | 160 +++++ zig/src/families/glm_cuda/kern_exl3.zig | 17 +- zig/src/families/glm_cuda/kern_kda.zig | 126 ++++ zig/src/families/glm_cuda/kern_moe.zig | 91 +++ zig/src/families/glm_cuda/native.zig | 495 ++------------ zig/src/families/glm_cuda/native_draft.zig | 181 +++++ zig/src/families/glm_cuda/native_env.zig | 54 ++ zig/src/families/glm_cuda/native_follow.zig | 104 +++ zig/src/families/glm_cuda/native_prompt.zig | 135 ++++ zig/src/families/glm_cuda/proto.zig | 3 +- zig/src/families/glm_cuda/qpack.zig | 90 +++ zig/src/families/glm_cuda/seq.zig | 18 +- zig/src/families/glm_cuda/vision.zig | 13 +- zig/src/families/glm_cuda/weights.zig | 620 +----------------- zig/tests/cuda/glm/gen_copies.py | 7 +- zig/tests/cuda/glm_forward.zig | 218 +----- zig/tests/cuda/glm_forward_run.zig | 196 ++++++ zig/tests/cuda/glm_weights.zig | 4 +- 52 files changed, 2536 insertions(+), 2554 deletions(-) create mode 100644 docs/glm-cuda-notes.md create mode 100644 zig/kernels/cuda/glm/exl3_transposed.cuh create mode 100644 zig/src/families/glm_cuda/ckpt.zig create mode 100644 zig/src/families/glm_cuda/dflash_load.zig create mode 100644 zig/src/families/glm_cuda/dflash_select.zig create mode 100644 zig/src/families/glm_cuda/experts_load.zig create mode 100644 zig/src/families/glm_cuda/kern_attn.zig create mode 100644 zig/src/families/glm_cuda/kern_kda.zig create mode 100644 zig/src/families/glm_cuda/kern_moe.zig create mode 100644 zig/src/families/glm_cuda/native_draft.zig create mode 100644 zig/src/families/glm_cuda/native_env.zig create mode 100644 zig/src/families/glm_cuda/native_follow.zig create mode 100644 zig/src/families/glm_cuda/native_prompt.zig create mode 100644 zig/src/families/glm_cuda/qpack.zig create mode 100644 zig/tests/cuda/glm_forward_run.zig diff --git a/docs/glm-cuda-notes.md b/docs/glm-cuda-notes.md new file mode 100644 index 0000000000..cb850c3737 --- /dev/null +++ b/docs/glm-cuda-notes.md @@ -0,0 +1,103 @@ +# GLM-5.3-Flash on CUDA: source notes + +The longer descriptions these files carried as multi-line header comments, kept here as written. + +## `zig/kernels/cuda/flashnext/sample.cu` + +Flash Next port: device-side greedy candidates. One block per bf16 logits row: the row's first maximum (lowest column among equal maxima, as argmax picks it), mapped to a global token id, and the row's log-sum-exp in fp32. out[row] = {id, max (fp32 bits), lse (fp32 bits), 0}. Build: nvcc -cubin -arch=sm_121 -O3 -o fn_sample.cubin fn_sample.cu + +## `zig/kernels/cuda/glm/dflash.cu` + +DFlash2 block drafter kernels (our own, after z-lab/dflash's reference, MIT): the target's layer taps, the drafter's per-head norms with rotary positions, its grouped two-tap dynamic convolutions, SwiGLU, and its attention over a stream's context ring (the target's features, a 2048-token window) and the 8-row block (bidirectional). Layout: hidden 4096, 32 query heads and 8 key/value heads of 128, a block of 8 rows a stream. + +## `zig/kernels/cuda/glm/exl3.cu` + +GLM-5.3-Flash's EXL3 routed experts (4-bit "mcg" trellis, ExLlamaV3's format: MIT, Copyright (c) 2025 Turboderp) on CUDA: the grouped trellis GEMV/GEMM and the Hadamard rotations around it. From TensorFold's python-0.6 src/tensorfold/families/glm5_next/cuda/exl3.cu, without torch, over the routed pairs alone (the shared expert runs on the dense linears), with items of 64 pairs for prompt chunks (each expert's tiles read once a chunk). A 16x16 tile is 32 little-endian 32-bit words; lane L of a warp decodes the tile's values 8L..8L+7 from words L-1 and L, and those eight values are exactly the B fragments of two mma.m16n8k16 (columns 0-7 and 8-15 of the tile). Every output depends only on its own pair: pairs share an mma tile but mma keeps rows independent, the K range of every warp and split is fixed by the shape, and warps and splits are summed in a fixed order, never with atomics. + +## `zig/kernels/cuda/glm/experts_sym.cuh` + +GLM-5.3-Flash's symmetric routed experts: flashnext/experts.cuh with a group's block holding its scales alone (the bias of a symmetric int4 group is -8 x its scale, exact in bf16): 16 words a 32-row tile instead of 32. Shared by experts_sym.cu (decode form) and experts_sym_prefill.cu (prefill form). + +## `zig/kernels/cuda/glm/experts_sym_prefill.cu` + +GLM-5.3-Flash's symmetric routed experts, prefill form: flashnext/experts_prefill.cu reading experts_sym.cuh's blocks, each group's bias pair -8 x its scale pair (exact in bf16). Prefill: w = bf16(fma(q, s, b)), one fp32 chain over K; chunk-invariant bits, not decode's (replies prefill again). + +## `zig/kernels/cuda/glm/index.cu` + +GLM-5.3-Flash on CUDA: the DSA indexer, ours. A row writes its index key (LayerNorm) and gate into its stream's rings, the row that completes a pool of 4 writes the pooled key, each row scores its visible pools, keeps its 512 best (the lower pool on ties) and lists their tokens ascending. Rows name their own rings, pools and positions, so one launch serves any mix of streams, and a row's choice never depends on the rows beside it. + +## `zig/kernels/cuda/glm/kda.cu` + +From src/tensorfold/families/glm5_next/cuda/kda.cu (python-0.6), with the chain's body a device function and a segmented chain (any number of streams' windows in one launch) added. GLM-5.3-Flash's KDA (Kimi delta attention) on CUDA: one block of 1024 threads per head, a chain of R rows. Per row, following the Hugging Face definition: the depthwise conv over [conv state; q|k|v rows] (4 taps, fp32) with SiLU and one bf16 rounding; fp32 L2 norms of q and k (eps inside the sum, q times DK^-0.5); the per-channel decay g_i = exp(lower * sigmoid(exp(A_log) * (a_i + dt_bias_i))) in fp32; beta = bf16(sigmoid(b)); the delta rule in fp32 (decay along the key channel, read with k, correct toward v, read out with q); read-out rounded to bf16; then the gated RMSNorm bf16(w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(gate)). Warp w owns value rows 4w .. 4w + 3, lane l the key columns 4l .. 4l + 3. The state update is one routine (``update``) shared with ``replay``, compiled without FMA contraction, so a replayed prefix of a window gives the bits of the serial steps. + +## `zig/kernels/cuda/glm/latent.cu` + +GLM-5.3-Flash on CUDA: MLA attention on the 512-wide latent cache, ours. A program is one row's 32 heads over one chunk of 512 keys (online softmax in 16-key tiles, tensor-core scores and outputs), and the chunks are merged in order; nothing a program computes depends on another row, so a row has the same bits alone, in a decode window, in sixteen streams' windows or in a prompt chunk. Every row names its own cache (``lc``, the latent rows of its stream) and position (``pos``), so one launch serves any mix of streams. + +## `zig/kernels/cuda/glm/mm.cu` + +GLM-5.3-Flash on CUDA: the dense linears, ours. One tensor-core kernel for every row count (a decode window's few rows, sixteen streams' windows at once, a prompt chunk's 2048): each group of 64 inputs is summed from zero by the bf16 mma, then acc = fma(P, scale, acc) in group order, and the K slices (fixed by the shape) are added in slice order, so a row's bits never depend on the rows beside it. Weights are tiled at load (glm_mm_tile*): a column tile of 8 outputs and a group of 64 inputs is 32 lanes' chunks of 16 values in mma B-fragment order (lane l: column l / 4; value 4 kt + j at input 16 kt + {2t, 2t+1, 2t+8, 2t+9}[j], t = l % 4), each chunk BITS / 2 words (ops.cu's 16-value packing; BITS 16: the bf16 values themselves). Scales are fp16 [N / 8][K / 64][8]. + +## `zig/kernels/cuda/glm/rows.cu` + +GLM-5.3-Flash on CUDA: the row steps between the linears, ours (embedding, RMSNorm, the hyper-connections, the residuals, SwiGLU, the router and its top-8, the experts' combine, the KDA conv window). Each row is one block's, with fp32 sums in a fixed order (threads strided over the row, then a fixed tree), so a row's bits depend on its own inputs alone, never on how many rows the window holds. + +## `zig/kernels/cuda/glm/vision.cu` + +GLM-5.3-Flash's vision tower on CUDA (families/glm_cuda/vision.zig): the elementwise and row steps around its cuBLASLt products, as mlx-vlm's glm5_next/vision.py defines them. The residual stream stays in fp32; a step's bf16 output is rounded once, to nearest even. + +## `zig/src/families/glm_cuda/api.zig` + +GLM-5.3-Flash (glm5_next) on CUDA, two ranks: the pieces the weight loader, the forward and the native host share. Weights are a name -> device tensor store whose names are the Python engine's prepared-weight names (its weight tree: "layers.3.dsa.q_b.weight", "layers.5.moe.experts.up", ...), so every tensor the checkpoint loader builds can be checked byte for byte against the Python engine's (tools: capture_ref.py's weights.json hashes). Kernels: the fatbins embedded by the build (this family's, zig/kernels/cuda/glm; the shared grouped experts and the greedy candidates, zig/kernels/cuda/flashnext). + +## `zig/src/families/glm_cuda/blocks.zig` + +GLM-5.3-Flash on CUDA: the weights' views and the layers of a step (forward.py's blocks): KDA and DSA attention, the dense MLP and the MoE, the hyper-connections around them, the final norm and the head; and the MTP head. Row-wise work covers the step's rows at once; the KDA recurrences, the attention and the indexer read the step's tables (forward.prepare). + +## `zig/src/families/glm_cuda/copy.zig` + +Copy drafts (prompt lookup): code, JSON and edits repeat what came before, often longer than the MTP head drafts with confidence. A stream keeps its prompt and every committed token, and an index of each 8-token run's latest end; when the history's last 8 tokens ran before, the tokens that followed them are a draft. A run inside the reply itself must match 16 tokens (a reply's own short repeats mislead more often). The drafts are verified like the head's, so a wrong copy costs rows, never the reply. + +## `zig/src/families/glm_cuda/dflash.zig` + +DFlash2 block drafts (opt-in: TF_GLM_DRAFTER=dflash2, TF_GLM_DFLASH2 the checkpoint's folder). IncoAI's drafter (incoai/GLM-5.3-Flash-DFlash2, CC BY-NC-ND 4.0: non-commercial use; read from its own download, never changed on disk or redistributed) after z-lab/dflash's DFlash2 architecture (MIT); this implementation is our own. The drafter is a 5-layer Qwen3-style decoder (hidden 4096, 32/8 heads of 128, SwiGLU 12288) that drafts a whole block at once: rows [anchor, MASK x 7] at positions start .. start + 7 attend, bidirectionally, to each other and to the stream's context: the target's completed residuals after layers 5, 14, 24, 33 and 42 (each the mean of its 4 mHC streams), concatenated, through fc and a norm, then each layer's own k and v projections (a ring of the last 2048 positions). Two-tap dynamic convolutions along the block wrap both sublayers. The block's last 7 rows go through the target's head; a path selector picks each draft among the row's 16 likeliest tokens by the previous draft's codebook row. Drafts never change a reply: the target verifies them. Both ranks run the drafter on their own copy (the head is each rank's vocabulary half, its candidates gathered); rank 0 alone runs the selector, and the verify frame carries its drafts to rank 1. + +## `zig/src/families/glm_cuda/draw.zig` + +A sampled row's token from both ranks' gathered candidates (fn_cand_topk): each rank's 64 largest logits, global ids, its log-sum-exp and its log-sum-exp at the row's temperature. The merged 128 hold the vocabulary's top 64 exactly, so a rule within them draws as lanes.sampling.choose over the whole vocabulary would: top_k 1..64 by choose itself (only the top k take part), a wider rule when its nucleus (top_p) or floor (min_p) ends within the top 64, keyed the same way (seed, position, global id). A rule that reaches past them asks for the full row. + +## `zig/src/families/glm_cuda/forward.zig` + +GLM-5.3-Flash (glm5_next) on CUDA, one rank of two: the forward on this family's own kernels, for any mix of streams. A step is a list of segments (a stream and its rows, in order): the row-wise work (projections, norms, hyper-connections, router, experts, head) runs over all the step's rows at once, and the attention, the indexer and the KDA recurrences read the step's tables (each row's position and its stream's caches). Every kernel keeps a row's bits its own, so a stream's tokens are the same alone and beside fifteen others. Both ranks make the same calls in the same order. Tokens are drawn from both ranks' vocabulary shards (an all-gather of each rank's best candidates), so both ranks reach the same tokens with no host decision that could differ. + +## `zig/src/families/glm_cuda/guard.zig` + +A rank's watchdog. The GB10's memory is the host's: when a rank's caches, graphs and buffers take the last of it the kernel stalls the whole machine (no OOM kill comes), and the peer then waits in a collective forever. So a thread reads MemAvailable every 100 ms and the rank exits at the floor (TF_MEM_FLOOR_GIB, 2 GiB), which frees everything it holds; below the floor and 3 GiB more it is ``tight`` and rank 0 admits no new request. It also exits when its peer is gone (the control link hung up or timed out, or NCCL reports an error): a rank alone can only hang, and an exited one can be restarted. + +## `zig/src/families/glm_cuda/kern.zig` + +GLM-5.3-Flash on CUDA: every launch the forward makes, on this family's own kernels (glm/mm.cu, rows.cu, latent.cu, index.cu, kda.cu, ops.cu, vision.cu) and the shared grouped experts. Each is written so that a row's bits depend on its own inputs alone: a window's rows, sixteen streams' windows and a prompt chunk's rows go through the same arithmetic. The attention and indexer launches read a step's per-row tables (``Rows``): each row's position and its stream's caches, so one launch serves any mix of streams. + +## `zig/src/families/glm_cuda/native.zig` + +GLM-5.3-Flash (glm5_next) on two GB10s behind the native family interface: rank 0's lane core drives the engine through a lane backend that first sends each call to rank 1, which replays it on its own half of the model (``follow``). Every token is drawn from both ranks' vocabulary shards (an NCCL exchange), so both ranks reach the same tokens and rank 1 needs nothing back. Up to 16 streams decode together, each with its own caches sized to its request out of one memory pool; a released stream's kept prompt state stays for a later prompt that extends it while the pool has room. + +## `zig/src/families/glm_cuda/seq.zig` + +GLM-5.3-Flash on CUDA: a stream's committed caches (forward.py's State), sized to its request. KDA conv windows and recurrent states (two parities), the main and MTP latent caches, the indexer's keys and gates (rings: only a window's rows and the 3 before it are read) and pools, and the kept prompt state a later prompt can resume. + +## `zig/src/families/glm_cuda/vision.zig` + +GLM-5.3-Flash's vision tower on CUDA (rank 0 encodes; rank 1 receives the features), as mlx-vlm's glm5_next/vision.py defines it: the patch embedding (a linear over a patch's 2 x 14 x 14 x 3 values), 24 blocks (RMSNorm, attention with RMSNorm'd q and k and a 2-D rotary over each frame's patches, a limited SwiGLU), the post-norm, the 2 x 2 downsample (a linear over each merge block's four patches) and the merger (a projection, LayerNorm + GELU, a limited SwiGLU) to the language model's 4096. The linears are cuBLASLt products (bf16 in, fp32 accumulate, bias added before the one rounding); the residual stream stays fp32; glm/vision.cu does the rest. Checked against the torch reference (tools/vision_ref.py) by tf-cuda-test glm-vision. + +## `zig/src/families/glm_cuda/weights.zig` + +GLM-5.3-Flash's prepared weights for one rank of two, read from the converted checkpoint (glm-affine-experts-v1: MLX affine 4-bit routed experts in groups of 64 with fp16 scales and biases, everything else BF16 / FP32), as the Python engine prepares them (families/glm5_next/cuda/weights.py, its affine-experts path, split.py's rank rules): every tensor under the name, shape and bytes of the Python engine's prepared weights (its weight tree, the names tools/capture_ref.py hashes), uploaded to this rank's GPU. The derivations are exact: rank rows and columns cut on the host, BF16 rows stacked, the kv_b blocks split per head (the latent path's absorb weights), fp16 expert scales rounded to bf16 (round to nearest even, as torch), and the experts' blocks in experts_pack.cu's index map. The draft head is the verification head itself (no 4-bit copy). + +## `zig/tests/cuda/glm_forward.zig` + +GLM-5.3-Flash on CUDA: the family's kernels against float64 references and the row-invariance contract (a row's bits alone equal its bits in any window, any mix of streams), on one GPU; and the two-rank engine's contract on the real checkpoint: a stream's greedy reply is the same alone, beside 15 other streams, drafted, and resumed from a kept prompt state. usage: tf-cuda-test glm-kcheck tf-cuda-test glm-batch [tokens] [sync] + +## `zig/tests/cuda/glm_weights.zig` + +GLM-5.3-Flash on CUDA: the native weight loader (src/families/glm_cuda/weights.zig) against the Python engine's prepared weights (tools/capture_ref.py's weights.json: each tensor's name, bytes and sha256). usage: tf-cuda-test glm-weights diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 212dde8885..b3732ee0e0 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -43,11 +43,11 @@ const kernels = [_]Kernel{ .{ .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 - // the routed experts' grouped kernels and the greedy candidates the GLM family shares with Flash Next (zig/kernels/cuda/flashnext) + // the grouped routed experts and greedy candidates GLM shares with Flash Next (zig/kernels/cuda/flashnext) .{ .name = "fn_experts", .src = "flashnext/experts", .flags = &.{"-O3"} }, .{ .name = "fn_experts_prefill", .src = "flashnext/experts_prefill", .flags = &.{"-O3"} }, .{ .name = "fn_sample", .src = "flashnext/sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, - // GLM-5.3-Flash (glm5_next): the device copy of the Python engine's KDA extension (zig/tests/cuda/glm/gen_copies.py) + // GLM-5.3-Flash (glm5_next): the Python engine's KDA extension, copied by zig/tests/cuda/glm/gen_copies.py .{ .name = "glm_kda", .src = "glm/kda", .flags = &.{ "-O3", "--fmad=false" } }, // glm5_next/cuda/kda.py, tensorfold_glm_kda_v2 .{ .name = "glm_ops", .src = "glm/ops", .flags = &.{ "-O3", "--fmad=false" } }, // ours: the token lists, casts and positions torch ops make .{ .name = "glm_experts", .src = "glm/experts_sym", .flags = &.{"-O3"} }, // ours: fn_experts on symmetric groups (scales alone) diff --git a/zig/kernels/cuda/flashnext/experts.cu b/zig/kernels/cuda/flashnext/experts.cu index ea680620e4..ad03d5cc02 100644 --- a/zig/kernels/cuda/flashnext/experts.cu +++ b/zig/kernels/cuda/flashnext/experts.cu @@ -1,5 +1,4 @@ -// Device code of src/tensorfold/cuda/experts.cu (generated by zig/tests/cuda/flashnext/gen_copies.py; do not edit). -// Decode form: acc = fma(xs, b, fma(p, s, acc)) per group in order, one pair an mma row, so no pair affects another. +// Copy of src/tensorfold/cuda/experts.cu, do not edit. Decode: acc = fma(xs, b, fma(p, s, acc)) a group, in order. #include #include diff --git a/zig/kernels/cuda/flashnext/experts.cuh b/zig/kernels/cuda/flashnext/experts.cuh index f7fac4861d..4e00c22a56 100644 --- a/zig/kernels/cuda/flashnext/experts.cuh +++ b/zig/kernels/cuda/flashnext/experts.cuh @@ -1,5 +1,4 @@ -// Device code of src/tensorfold/cuda/experts.cuh (generated by zig/tests/cuda/flashnext/gen_copies.py; do not edit). -// Shared by the grouped-expert kernels (experts.cu: decode form and plan; experts_prefill.cu: prefill form). +// Copy of src/tensorfold/cuda/experts.cuh, do not edit: shared by experts.cu (decode, plan) and experts_prefill.cu. #pragma once #include diff --git a/zig/kernels/cuda/flashnext/experts_prefill.cu b/zig/kernels/cuda/flashnext/experts_prefill.cu index 6cab13ddf8..6b21bfef40 100644 --- a/zig/kernels/cuda/flashnext/experts_prefill.cu +++ b/zig/kernels/cuda/flashnext/experts_prefill.cu @@ -1,5 +1,4 @@ -// Device code of src/tensorfold/cuda/experts_prefill.cu (generated by zig/tests/cuda/flashnext/gen_copies.py; do not edit). -// Prefill: w = bf16(fma(q, s, b)), one fp32 chain over K; chunk-invariant bits, not decode's (replies prefill again). +// Copy of src/tensorfold/cuda/experts_prefill.cu, do not edit. Prefill: w = bf16(fma(q, s, b)), one fp32 chain over K. #include #include diff --git a/zig/kernels/cuda/flashnext/sample.cu b/zig/kernels/cuda/flashnext/sample.cu index fc23bbb4b4..39a794a1d4 100644 --- a/zig/kernels/cuda/flashnext/sample.cu +++ b/zig/kernels/cuda/flashnext/sample.cu @@ -1,6 +1,4 @@ -// Flash Next port: device-side greedy candidates. One block per bf16 logits row: the row's first maximum (lowest -// column among equal maxima, as argmax picks it), mapped to a global token id, and the row's log-sum-exp in fp32. -// out[row] = {id, max (fp32 bits), lse (fp32 bits), 0}. Build: nvcc -cubin -arch=sm_121 -O3 -o fn_sample.cubin fn_sample.cu +// Flash Next greedy candidates, a block a bf16 row: out[row] = {first argmax as a global id, max, fp32 lse, 0}. #include #include @@ -51,12 +49,7 @@ extern "C" __global__ void __launch_bounds__(1024) fn_rows_top(const __nv_bfloat } } -// Sampling candidates: each row's 64 largest logits (by value, then the lower column at a tie), as global token ids -// and fp32 values, the row's log-sum-exp, and its log-sum-exp at the row's temperature (inv_t[row], x * inv_t); both -// ranks' sets are gathered and the keyed draw runs on the host over the merged 128 (whose first 64 are the vocabulary's -// top 64 exactly). Each thread keeps its columns' sorted top 64, then the block merges the threads' lists by repeated -// argmax over their heads (exact: no value of the row's top 64 can be missing from its thread's list). -// out[row * 132]: ids[64], values[64] (fp32 bits), lse, lse at the temperature (fp32 bits), 2 spare. +// Each row's top 64 (lower column at ties) as ids and values, its lse and lse at inv_t; out[row * 132], merged on host. #define TOPK 64 #define TOPW (2 * TOPK + 4) __device__ __forceinline__ bool above(float a, int ia, float b, int ib) { return a > b || (a == b && ia < ib); } @@ -124,8 +117,7 @@ extern "C" __global__ void __launch_bounds__(256) fn_cand_topk(const __nv_bfloat } } -// Structured output: each listed row's tokens its grammar does not allow to -inf (rows[i]: the logits row, bits: the -// row's allowed bits over this rank's columns, token c: bit c % 32 of word c / 32), before the picks and draws. +// Structured output: a listed row's banned tokens to -inf (bit c % 32 of word c / 32), before the picks and draws. extern "C" __global__ void __launch_bounds__(256) fn_grammar_rows(__nv_bfloat16* __restrict__ logits, int stride, int cols, const unsigned* __restrict__ bits, int words, const int* __restrict__ rows) { __nv_bfloat16* x = logits + (size_t)rows[blockIdx.x] * stride; @@ -134,11 +126,7 @@ extern "C" __global__ void __launch_bounds__(256) fn_grammar_rows(__nv_bfloat16* if (!((w[c >> 5] >> (c & 31)) & 1u)) x[c] = __ushort_as_bfloat16((unsigned short)0xff80); } -// fn_cand_topk's output, by radix selection over the bf16 bits instead of per-thread lists: a pass histograms the -// order-preserving keys' high byte (and takes both log-sum-exps), a second the low byte within the bin holding the -// 64th largest, which fixes the threshold key K; a third collects every column above K and, of the columns at K, the -// lowest ones (in column order), and the picks are ranked by value, then the lower column. The same 64 in the same -// order as fn_cand_topk (-0 counts as +0, as the float compare does); every pass reads the row once, coalesced. +// fn_cand_topk's 64 in its order by radix selection over the bf16 keys: high byte, low byte, then collect; 3 reads. __device__ __forceinline__ unsigned bf_key(unsigned short u) { if (u == 0x8000u) u = 0; // -0: the same value as +0 return (u & 0x8000u) ? (~u & 0xffffu) : (u | 0x8000u); diff --git a/zig/kernels/cuda/glm/dflash.cu b/zig/kernels/cuda/glm/dflash.cu index 1292edc6a8..013df80ae1 100644 --- a/zig/kernels/cuda/glm/dflash.cu +++ b/zig/kernels/cuda/glm/dflash.cu @@ -1,7 +1,4 @@ -// DFlash2 block drafter kernels (our own, after z-lab/dflash's reference, MIT): the target's layer taps, the drafter's -// per-head norms with rotary positions, its grouped two-tap dynamic convolutions, SwiGLU, and its attention over a -// stream's context ring (the target's features, a 2048-token window) and the 8-row block (bidirectional). -// Layout: hidden 4096, 32 query heads and 8 key/value heads of 128, a block of 8 rows a stream. +// DFlash2 drafter kernels (ours, after z-lab/dflash, MIT): taps, norms+rope, convs, SwiGLU, ring and block attention. #include #include @@ -48,9 +45,7 @@ __device__ __forceinline__ void normRope(const __nv_bfloat16* x, const __nv_bflo } } // namespace -// The completed residual of a target layer, its 4 streams' mean, into column block ``slot`` of out [R][5 D]. -// grid (R, D / 1024), 256 threads. -// mode 1 (development, TF_GLM_DFLASH2_TAP=s0): stream 0 alone. +// A target layer's residual, its 4 streams' mean (TF_GLM_DFLASH2_TAP=s0: stream 0), into column block ``slot``. extern "C" __global__ void df_tap(const __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ out, int slot, int mode) { const int r = blockIdx.x; const __nv_bfloat16* xr = x + (long long)r * 4 * D; @@ -65,8 +60,7 @@ extern "C" __global__ void df_tap(const __nv_bfloat16* __restrict__ x, __nv_bflo } } -// The block's queries [R][QH*HD] (from the fused qkv rows, stride ldq) normed and rotated at pos[r] into q, and its -// keys likewise into kb, values copied into vb ([R][KVW] each). grid (R, QH + KVH), 64 threads. +// The block's queries (fused qkv, stride ldq) normed and rotated at pos[r] into q, keys into kb, values into vb. extern "C" __global__ void df_block_qkv(const __nv_bfloat16* __restrict__ qkv, int ldq, const __nv_bfloat16* __restrict__ qn, const __nv_bfloat16* __restrict__ kn, const int* __restrict__ pos, float eps, float theta, __nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ kb, __nv_bfloat16* __restrict__ vb) { @@ -85,8 +79,7 @@ extern "C" __global__ void df_block_qkv(const __nv_bfloat16* __restrict__ qkv, i } } -// Context rows' keys and values ([R][2 KVW]: k | v) normed, rotated at pos[r] and written into each row's stream -// ring (rings[r]) at slot pos % WIN of ``layer``. grid (R, KVH), 64 threads. +// Context rows' keys and values (k | v) normed, rotated at pos[r], into rings[r] at slot pos % WIN of ``layer``. extern "C" __global__ void df_ctx_kv(const __nv_bfloat16* __restrict__ kv, const __nv_bfloat16* __restrict__ kn, const long long* __restrict__ pos, const uint64_t* __restrict__ rings, int layer, float eps, float theta) { @@ -101,10 +94,7 @@ extern "C" __global__ void df_ctx_kv(const __nv_bfloat16* __restrict__ kv, const vo[threadIdx.x + 64] = row[KVW + kh * HD + threadIdx.x + 64]; } -// The grouped two-tap dynamic causal convolution along each stream's block of 8 rows: -// y[t, c] = sum over o in {0, 1} of (base[o][c] + dyn[t][which][o][c / 16]) * z[t - o][c] (z[-1] = 0) -// dyn [R][2 (pre, post)][2 (tap)][GROUPS]; base [2 (tap)][D] (the module's pre or post kernel). add: x += y, else -// out = y. grid (R, D / 1024), 256 threads. +// Two-tap dynamic causal conv over a block: y[t] = sum_o (base[o] + dyn[t][o][c / 16]) * z[t - o]; add or store. extern "C" __global__ void df_conv(const __nv_bfloat16* __restrict__ z, const __nv_bfloat16* __restrict__ dyn, int which, const __nv_bfloat16* __restrict__ base, __nv_bfloat16* __restrict__ out, int add) { const int r = blockIdx.x, t = r % BLOCK; @@ -129,10 +119,7 @@ extern "C" __global__ void df_swiglu(const __nv_bfloat16* __restrict__ gu, __nv_ } } -// Attention of each stream's block (8 rows, 4 query heads a key/value head) over its ring's context keys in the -// window, positions max(from, p - 2047) .. start - 1 for a row at p, and the block's 8 keys (every one visible): -// grid (S, KVH), 256 threads (a warp a block row, its 4 heads; a lane 4 of the 128 dims). -// q [S*8][QH*HD], kb / vb [S*8][KVW], o [S*8][QH*HD]; starts / froms the streams' block and context starts. +// Each block's 8 rows attend to their ring's window (max(from, p - 2047) .. start - 1) and all 8 block keys. extern "C" __global__ void __launch_bounds__(256) df_attn(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kb, const __nv_bfloat16* __restrict__ vb, const uint64_t* __restrict__ rings, const long long* __restrict__ starts, const long long* __restrict__ froms, @@ -188,10 +175,7 @@ extern "C" __global__ void __launch_bounds__(256) df_attn(const __nv_bfloat16* _ } } -// df_attn in chunks: block (s, kh, c) takes the c-th 1/CH of each row's keys (its context, then the block's 8) and -// writes the row's 4 heads' running max, sum and unnormalized output (pm, pl [R][QH][CH]; pacc [R][QH][CH][HD], fp32); -// df_attn_merge combines the chunks. The serial walk over 2,048 keys made df_attn latency-bound; this one runs -// CH blocks a kv head. grid (S, KVH, CH), 256 threads (a warp a block row, as df_attn). +// df_attn split over CH key chunks (running max, sum, fp32 output), merged by df_attn_merge: not latency-bound. extern "C" __global__ void __launch_bounds__(256) df_attn_part(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kb, const __nv_bfloat16* __restrict__ vb, const uint64_t* __restrict__ rings, const long long* __restrict__ starts, const long long* __restrict__ froms, int layer, float scale, float* __restrict__ pm, float* __restrict__ pl, @@ -276,8 +260,7 @@ extern "C" __global__ void df_add(__nv_bfloat16* __restrict__ x, const __nv_bflo x[r * D + c] = tobf(bf(x[r * D + c]) + bf(y[r * D + c])); } -// out = bf16(g rank 0 + g rank 1): both ranks' gathered fp32 partials [2][R][D] of a split linear (as the engine's -// glm_residual_add reads them). grid (R, D / 1024), 256 threads. +// out = bf16(g rank 0 + g rank 1): both ranks' gathered fp32 partials [2][R][D] of a split linear. extern "C" __global__ void df_sum2(const float* __restrict__ g, __nv_bfloat16* __restrict__ out, int R) { const long long r = blockIdx.x; for (int c = blockIdx.y * 1024 + threadIdx.x; c < (blockIdx.y + 1) * 1024; c += blockDim.x) { @@ -286,9 +269,7 @@ extern "C" __global__ void df_sum2(const float* __restrict__ g, __nv_bfloat16* _ } } -// A rank's 16 largest logits of each head row (bf16 [R][cols]), as global ids and fp32 values, in two stages: -// df_top16_part: each 4096-column chunk's 16 (grid (R, chunks), 256 threads); df_top16_merge: the chunks' into -// out [R][32] (16 ids, then 16 values' bits; grid R, 256 threads). Ties keep the lower id. +// A rank's 16 largest logits a head row as global ids and fp32 values: per-chunk part, then merge; ties keep low id. namespace { constexpr int CHUNK_COLS = 4096; constexpr int TOP = 16; diff --git a/zig/kernels/cuda/glm/exl3.cu b/zig/kernels/cuda/glm/exl3.cu index 64a06dc7d7..1505fe31a5 100644 --- a/zig/kernels/cuda/glm/exl3.cu +++ b/zig/kernels/cuda/glm/exl3.cu @@ -1,12 +1,4 @@ -// GLM-5.3-Flash's EXL3 routed experts (4-bit "mcg" trellis, ExLlamaV3's format: MIT, Copyright (c) 2025 Turboderp) -// on CUDA: the grouped trellis GEMV/GEMM and the Hadamard rotations around it. From TensorFold's python-0.6 -// src/tensorfold/families/glm5_next/cuda/exl3.cu, without torch, over the routed pairs alone (the shared expert runs -// on the dense linears), with items of 64 pairs for prompt chunks (each expert's tiles read once a chunk). -// -// A 16x16 tile is 32 little-endian 32-bit words; lane L of a warp decodes the tile's values 8L..8L+7 from words L-1 -// and L, and those eight values are exactly the B fragments of two mma.m16n8k16 (columns 0-7 and 8-15 of the tile). -// Every output depends only on its own pair: pairs share an mma tile but mma keeps rows independent, the K range of -// every warp and split is fixed by the shape, and warps and splits are summed in a fixed order, never with atomics. +// GLM EXL3 routed experts (4-bit mcg trellis, ExLlamaV3's format: MIT, (c) 2025 Turboderp), python-0.6's, no torch. #include #include @@ -49,8 +41,7 @@ __device__ __forceinline__ uint32_t load_pair(const half* x, bool ok) { return ok ? *reinterpret_cast(x) : 0u; } -// Grid (item, n block, mat * SK + split) -> Z[mat][split][pair][n]; an item is an expert and up to 16 MT of its pairs -// (the plan's members); warps sum fixed k ranges, added in warp order. +// Grid (item, n block, mat * SK + split) -> Z[mat][split][pair][n]: an expert and <= 16 MT of its pairs; warp order. template __device__ __forceinline__ void grouped( const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, @@ -149,8 +140,7 @@ __device__ __forceinline__ void grouped( } } -// Fast Walsh-Hadamard transform of 128 values held 4 per lane (lane L: values 4L..4L+3), natural order, fixed -// butterfly order: strides 1, 2 in registers, 4..64 across lanes. +// Fast Walsh-Hadamard transform of 128 values, 4 a lane, natural order: strides 1, 2 in registers, 4..64 across lanes. __device__ __forceinline__ void fwht128(float (&v)[4], int lane) { float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3]; v[0] = a + c; v[1] = b + d; v[2] = a - c; v[3] = b - d; @@ -167,16 +157,9 @@ __device__ __forceinline__ void fwht128(float (&v)[4], int lane) { __device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } -// -- prompt chunks: the staged kernel ------------------------------------------------------------------------------- -// An item is an expert and up to PM of its pairs (the plan's members, tile PM). Block (item, column block of -// NW * NTW n-tiles, matrix): the four warps split the columns, not K, so a warp's sums need no cross-warp reduction -// and every decoded tile feeds PM / 8 mma; the member rows' k-slices and the tiles' words are staged in shared -// memory by cp.async, STAGES deep, and the A fragments come from the staged rows by ldmatrix. The same sums as the -// grouped kernel for every output, in a fixed order (one warp owns each output; K runs in order). -// -// decode_tile's values with fewer integer instructions: the byte-aligned 16-bit states taken with one PRMT each -// and w >> 4 shared by two states (MiaAI-Lab, GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold patch 0020, Apache-2.0). -// The same states, so the same bits. +// -- prompt chunks: the staged kernel, warps split columns not K, cp.async STAGES deep; the grouped kernel's sums -- + +// decode_tile via PRMT (MiaAI-Lab GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold patch 0020, Apache-2.0): same bits. __device__ __forceinline__ void decode_tile_p(uint32_t w, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { const uint32_t p = __shfl_sync(0xffffffffu, w, (lane + 31) & 31); const uint32_t s = __funnelshift_r(w, p, 20); @@ -202,12 +185,7 @@ __device__ __forceinline__ void ldmatrix_x4(uint32_t (&a)[4], const half* p) { : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) : "r"(s)); } -// xshift: the X row a member pair reads is pair >> xshift (3: one rotated row a token, when every expert of the layer -// shares one input rotation for gate and up; 0: a row a pair). -// SK: K splits -> Z[mat][split][pair][n] (the grouped kernel's layout: glm_exl3_gateup adds the splits in order). A -// small chunk, whose items hold few pairs, gets SK times the blocks at a fraction of the K each. -// Grid: x = (column block, matrix of ``mats``, split), y = item: one item's blocks are adjacent in launch order, so -// its staged rows come from L2 after the first block reads them. +// pair reads X row pair >> xshift (3: one row a token); SK K splits -> Z[mat][split][pair][n]; item blocks adjacent. template __device__ __forceinline__ void prompt_grouped( const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, @@ -331,9 +309,7 @@ extern "C" __global__ void __launch_bounds__(256) glm_exl3_prompt128w(const half prompt_grouped<128, 8, 2, 2, 3>(X0, X1, T0, T1, items, counts, members, Z, K, N, P, xshift, SK, mats); } -// Program (row, 128-block of K): Xh[row][block] = fp16((x[row] * suh) @ H), one rotation a token for the gate and up -// projections of a layer whose experts all share one suh (then every pair of the row reads Xh[row]: the same bits -// glm_exl3_rot_in gives each pair). One warp a program. +// Program (row, K block): Xh[row][block] = fp16((x[row] * suh) @ H) when a layer's experts share suh; rot_in's bits. extern "C" __global__ void glm_exl3_rot_rows(const __nv_bfloat16* __restrict__ x, int x_stride, const half* __restrict__ suh, half* __restrict__ out, int K) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); @@ -369,8 +345,7 @@ extern "C" __global__ void __launch_bounds__(128) glm_exl3_grouped64(const half* grouped<4, 2, 4>(X0, X1, T0, T1, items, counts, members, Z, K, N, P, SK); } -// Program (pair, 128-block of K, matrix): Xh[mat][pair][block] = fp16((x[row] * suh[mat][e]) @ H) for the gate and up -// projections of every routed pair (pair = row * 8 + slot). One warp a program. +// Program (pair, K block, matrix): Xh[mat][pair][block] = fp16((x[row] * suh[mat][e]) @ H), pair = row * 8 + slot. extern "C" __global__ void glm_exl3_rot_in(const __nv_bfloat16* __restrict__ x, int x_stride, const int* __restrict__ pick, const half* __restrict__ suh0, const half* __restrict__ suh1, half* __restrict__ out0, half* __restrict__ out1, int K) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -389,9 +364,7 @@ extern "C" __global__ void glm_exl3_rot_in(const __nv_bfloat16* __restrict__ x, for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); } -// Program (pair, 128-block of the rank's expert width): gate and up summed over the splits in order, rotated, scaled -// by svh; GLM's limited SwiGLU with the grouped kernels' bf16 roundings; then the down projection's input rotation: -// Xd[pair][block] = fp16((act * suh_d[e]) @ H). +// Program (pair, width block): gate, up over splits in order, rotated, svh, limited SwiGLU, Xd = fp16(act suh_d H). extern "C" __global__ void glm_exl3_gateup(const float* __restrict__ Z, const int* __restrict__ pick, const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, int P, int N, int SK, float limit) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -426,8 +399,7 @@ extern "C" __global__ void glm_exl3_gateup(const float* __restrict__ Z, const in for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); } -// Program (pair, 128-block of the model width): the down projection summed over the splits in order, rotated and -// scaled by svh -> y[row][slot][:] (fp32, rows of ``ys`` slots: the shared expert's last). This rank's share. +// Program (pair, model-width block): down over splits in order, rotated, svh -> y[row][slot] fp32; this rank's share. extern "C" __global__ void glm_exl3_down(const float* __restrict__ Z, const int* __restrict__ pick, const half* __restrict__ svh_d, float* __restrict__ y, int P, int D, int SK, int ys) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -448,185 +420,12 @@ extern "C" __global__ void glm_exl3_down(const float* __restrict__ Z, const int* for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); } -// -- prompt chunks, the transposed design: Y^T = W^T X^T ------------------------------------------------------------- -// After MiaAI-Lab's prompt kernels (GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold, patches 0004 and 0009, Apache-2.0), -// written for this engine. A decoded trellis tile is, reordered, an m16n8k16 A fragment of W^T (lane (g, t): -// a0 = b0[0], a1 = b1[0], a2 = b0[1], a3 = b1[1]), so each warp decodes its own tiles straight into registers and -// every decoded fragment feeds the whole item's member rows (PM = 64 pairs) as B fragments from shared memory. A -// block covers 16 weight tiles: gate and up's 8 each of one 128-column Hadamard block (the member rows staged once -// for both), or the down projection's 256 columns; 8 warps hold 2 tiles each. Each output is one fp32 chain of mma -// over ascending k tiles (no split K, no partials in memory), so a row's bits depend on its own values only. The -// epilogues (glm_exl3_gateup's and glm_exl3_down's formulas) run on the sums in shared memory: no round trip. -namespace { - -constexpr int TPM = 64, TNW = 8, TMT = 16 / TNW, TNB = TPM / 8, TT = TNW * 32, TPH = TPM / 2; -constexpr int TBK = 32, TKT = TBK / 16, TSTAGES = 3, TXS = TBK + 8; // a stage: 32 k values a row; +16 B: ldmatrix rows on distinct banks -constexpr int TWORDS = 16 * TKT * 32; // a stage's trellis words (16 tiles x 2 k tiles x 32) -constexpr int GU_ES = 128 + 4, DN_ES = 256 + 4; // the epilogue sums' row strides (floats) -template -constexpr int tstage_halves() { return XR * TXS + 2 * TWORDS; } -constexpr int tsmem_bytes() { - const int gu = TSTAGES * tstage_halves<2 * TPM>() * 2; // gate and up rows apart (no shared rotation) - const int ep = 2 * TPH * GU_ES * 4, dn = TPH * DN_ES * 4; - return gu > ep ? (gu > dn ? gu : dn) : (ep > dn ? ep : dn); -} - -// Which (matrix, weight tile) warp-tile q (0..15) of column block cb is. -template -__device__ __forceinline__ void ttile(int q, int cb, int& mat, int& nt) { - mat = MATS == 2 ? q >> 3 : 0; - nt = MATS == 2 ? cb * 8 + (q & 7) : cb * 16 + q; -} - -// The main loop: acc[i][nb] = this warp's tile i (16 output columns) x members nb * 8 .. + 7, over K. XR: the rows a -// stage holds (TPM when both matrices read the same rows: X0 == X1, xshift 3; 2 * TPM otherwise). -template -__device__ __forceinline__ void tmain(const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, - const uint32_t* __restrict__ T1, const int* rows_sh, int cnt, int e, int K, int NTILES, int cb, int xshift, - half* stage, float (&acc)[TMT][TNB][4]) { - constexpr int SH = tstage_halves(); - const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; - int mat, nt_unused; - ttile(warp * TMT, cb, mat, nt_unused); - const int KT = K >> 4, chunks = K / TBK; - const int nblocks = (cnt + 15) >> 4; // 16-member blocks in use -#pragma unroll - for (int i = 0; i < TMT; ++i) -#pragma unroll - for (int j = 0; j < TNB; ++j) -#pragma unroll - for (int c = 0; c < 4; ++c) acc[i][j][c] = 0.f; - auto load = [&](int chunk) { - half* dst = stage + (size_t)(chunk % TSTAGES) * SH; - const int k0 = chunk * TBK; - for (int idx = threadIdx.x; idx < XR * (TBK / 8); idx += TT) { - const int m = idx / (TPM * (TBK / 8)), rem = idx % (TPM * (TBK / 8)), i = rem / (TBK / 8), j = rem % (TBK / 8); - if (i >= nblocks * 16) continue; // a 16-member block without a member: never read - const int r = rows_sh[i]; - const half* src = (m ? X1 : X0) + (size_t)(r < 0 ? 0 : (r >> xshift)) * K + k0 + j * 8; - cp_async16(dst + ((size_t)m * TPM + i) * TXS + j * 8, src, r >= 0); - } - uint32_t* wdst = reinterpret_cast(dst + (size_t)XR * TXS); - for (int idx = threadIdx.x; idx < 16 * TKT * 8; idx += TT) { // 16-byte pieces of the 128-byte tile words - const int q = idx / (TKT * 8), kk = (idx / 8) % TKT, part = idx % 8; - int wm, nt; - ttile(q, cb, wm, nt); - const uint32_t* src = (wm ? T1 : T0) + (((size_t)e * KT + chunk * TKT + kk) * NTILES + nt) * 32 + part * 4; - cp_async16(wdst + ((size_t)q * TKT + kk) * 32 + part * 4, src, true); - } - }; -#pragma unroll - for (int c = 0; c < TSTAGES - 1; ++c) { - if (c < chunks) load(c); - cp_async_commit(); - } - for (int c = 0; c < chunks; ++c) { - cp_async_wait(); - __syncthreads(); - if (c + TSTAGES - 1 < chunks) load(c + TSTAGES - 1); - cp_async_commit(); - const half* base = stage + (size_t)(c % TSTAGES) * SH; - const half* xs = base + (size_t)(XR == TPM ? 0 : mat) * TPM * TXS; - const uint32_t* ws = reinterpret_cast(base + (size_t)XR * TXS); -#pragma unroll - for (int kk = 0; kk < TKT; ++kk) { - uint32_t a[TMT][4]; -#pragma unroll - for (int i = 0; i < TMT; ++i) { - uint32_t b0[2], b1[2]; - decode_tile_p(ws[((size_t)(warp * TMT + i) * TKT + kk) * 32 + lane], lane, b0, b1); - a[i][0] = b0[0]; a[i][1] = b1[0]; a[i][2] = b0[1]; a[i][3] = b1[1]; - } -#pragma unroll - for (int jb = 0; jb < TPM / 16; ++jb) { - if (jb < nblocks) { - uint32_t r[4]; - const int row = jb * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; - ldmatrix_x4(r, xs + (size_t)row * TXS + kk * 16 + ((lane >> 3) & 1) * 8); - const uint32_t bl[2] = {r[0], r[1]}, bh[2] = {r[2], r[3]}; -#pragma unroll - for (int i = 0; i < TMT; ++i) { - mma16816(acc[i][2 * jb], a[i], bl); - mma16816(acc[i][2 * jb + 1], a[i], bh); - } - } - } - } - } - cp_async_wait<0>(); - __syncthreads(); // the stages become the epilogue's sums -} +// -- prompt chunks, the transposed design: Y^T = W^T X^T, one fp32 mma chain an output over ascending k tiles ------- -// acc (this warp's tiles, members TPH * hf .. + TPH - 1) -> E[member - TPH * hf][col] fp32, row stride ES; col0: the -// first tile's first column within the block's E columns. -__device__ __forceinline__ void tspill(float* E, int ES, int col0, const float (&acc)[TMT][TNB][4], int hf) { - const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; -#pragma unroll - for (int i = 0; i < TMT; ++i) -#pragma unroll - for (int q = 0; q < TPM / 16; ++q) { - const int nb = hf * (TPM / 16) + q; - const int tok = q * 8 + 2 * t, col = col0 + i * 16 + g; - E[(size_t)tok * ES + col] = acc[i][nb][0]; - E[(size_t)(tok + 1) * ES + col] = acc[i][nb][1]; - E[(size_t)tok * ES + col + 8] = acc[i][nb][2]; - E[(size_t)(tok + 1) * ES + col + 8] = acc[i][nb][3]; - } -} - -template -__device__ __forceinline__ void tgateup(const half* __restrict__ Xg, const half* __restrict__ Xu, const uint32_t* __restrict__ Tg, - const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, - const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, - int K, int N, int xshift, float limit, unsigned char* smem_raw, int* rows_sh) { - const int cb = blockIdx.x, item = blockIdx.y; - if (item >= counts[0]) return; - const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; - if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; - __syncthreads(); - const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; - float acc[TMT][TNB][4]; - tmain<2, XR>(Xg, Xu, Tg, Tu, rows_sh, cnt, e, K, N >> 4, cb, xshift, reinterpret_cast(smem_raw), acc); - float* Eg = reinterpret_cast(smem_raw); - int mat, nt; - ttile<2>(warp * TMT, cb, mat, nt); - const int col0 = (nt - cb * 8) * 16; - const int n = cb * 128 + 4 * lane; - for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { - if (hf) __syncthreads(); // the first half's rows are read - tspill(Eg + (size_t)mat * TPH * GU_ES, GU_ES, col0, acc, hf); - __syncthreads(); - for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { - const int p = rows_sh[hf * TPH + i]; - float gv[4], uv[4]; -#pragma unroll - for (int j = 0; j < 4; ++j) { - gv[j] = Eg[(size_t)i * GU_ES + 4 * lane + j]; - uv[j] = Eg[(size_t)(TPH + i) * GU_ES + 4 * lane + j]; - } - fwht128(gv, lane); - fwht128(uv, lane); - float v[4]; -#pragma unroll - for (int j = 0; j < 4; ++j) { - float gg = fminf(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); - float uu = fminf(fmaxf(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), -limit), limit); - float act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); - v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); - } - fwht128(v, lane); - half* o = xd + (size_t)p * N + n; -#pragma unroll - for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); - } - } -} - -} // namespace +// After MiaAI-Lab's prompt kernels (GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold patches 0004, 0009, Apache-2.0). +#include "exl3_transposed.cuh" -// Prompt chunks (items of 64 pairs): block (128-column block of the expert width, item), 8 warps: Xd[pair] = the -// gate/up epilogue of the item's pairs. xshift 3: one rotated row a token read by gate and up alike (the rows staged -// once); 0: a row a pair in Xg and Xu. +// Prompt items of 64 pairs: block (128 expert columns, item), 8 warps: Xd = gate/up epilogue; xshift 3 or 0 as above. extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_gu(const half* __restrict__ Xg, const half* __restrict__ Xu, const uint32_t* __restrict__ Tg, const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, @@ -638,8 +437,7 @@ extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_gu(const ha else tgateup<2 * TPM>(Xg, Xu, Tg, Tu, items, counts, members, svh_g, svh_u, suh_d, xd, K, N, 0, limit, smem_raw, rows_sh); } -// Block (256-column block of the model width, item): y[row][slot][:] = the down epilogue of the item's pairs (this -// rank's share). +// Block (256 model-width columns, item): y[row][slot][:] = the down epilogue of the item's pairs (this rank's share). extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_dn(const half* __restrict__ Xd, const uint32_t* __restrict__ Td, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, const half* __restrict__ svh_d, float* __restrict__ y, int K, int D, int ys) { diff --git a/zig/kernels/cuda/glm/exl3_transposed.cuh b/zig/kernels/cuda/glm/exl3_transposed.cuh new file mode 100644 index 0000000000..77cbe5547d --- /dev/null +++ b/zig/kernels/cuda/glm/exl3_transposed.cuh @@ -0,0 +1,168 @@ +// Transposed prompt helpers after MiaAI-Lab (GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold 0004, 0009, Apache-2.0). +#pragma once +// Included by exl3.cu after its shared helpers (mcg decode, cp.async, ldmatrix), in the same translation unit. + +namespace { + +constexpr int TPM = 64, TNW = 8, TMT = 16 / TNW, TNB = TPM / 8, TT = TNW * 32, TPH = TPM / 2; +constexpr int TBK = 32, TKT = TBK / 16, TSTAGES = 3, TXS = TBK + 8; // a stage: 32 k values a row; +16 B: ldmatrix rows on distinct banks +constexpr int TWORDS = 16 * TKT * 32; // a stage's trellis words (16 tiles x 2 k tiles x 32) +constexpr int GU_ES = 128 + 4, DN_ES = 256 + 4; // the epilogue sums' row strides (floats) +template +constexpr int tstage_halves() { return XR * TXS + 2 * TWORDS; } +constexpr int tsmem_bytes() { + const int gu = TSTAGES * tstage_halves<2 * TPM>() * 2; // gate and up rows apart (no shared rotation) + const int ep = 2 * TPH * GU_ES * 4, dn = TPH * DN_ES * 4; + return gu > ep ? (gu > dn ? gu : dn) : (ep > dn ? ep : dn); +} + +// Which (matrix, weight tile) warp-tile q (0..15) of column block cb is. +template +__device__ __forceinline__ void ttile(int q, int cb, int& mat, int& nt) { + mat = MATS == 2 ? q >> 3 : 0; + nt = MATS == 2 ? cb * 8 + (q & 7) : cb * 16 + q; +} + +// Main loop: acc[i][nb] = warp tile i x members nb * 8 .. + 7 over K; XR rows a stage (TPM if X0 == X1, else 2 TPM). +template +__device__ __forceinline__ void tmain(const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, const int* rows_sh, int cnt, int e, int K, int NTILES, int cb, int xshift, + half* stage, float (&acc)[TMT][TNB][4]) { + constexpr int SH = tstage_halves(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + int mat, nt_unused; + ttile(warp * TMT, cb, mat, nt_unused); + const int KT = K >> 4, chunks = K / TBK; + const int nblocks = (cnt + 15) >> 4; // 16-member blocks in use +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int j = 0; j < TNB; ++j) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][j][c] = 0.f; + auto load = [&](int chunk) { + half* dst = stage + (size_t)(chunk % TSTAGES) * SH; + const int k0 = chunk * TBK; + for (int idx = threadIdx.x; idx < XR * (TBK / 8); idx += TT) { + const int m = idx / (TPM * (TBK / 8)), rem = idx % (TPM * (TBK / 8)), i = rem / (TBK / 8), j = rem % (TBK / 8); + if (i >= nblocks * 16) continue; // a 16-member block without a member: never read + const int r = rows_sh[i]; + const half* src = (m ? X1 : X0) + (size_t)(r < 0 ? 0 : (r >> xshift)) * K + k0 + j * 8; + cp_async16(dst + ((size_t)m * TPM + i) * TXS + j * 8, src, r >= 0); + } + uint32_t* wdst = reinterpret_cast(dst + (size_t)XR * TXS); + for (int idx = threadIdx.x; idx < 16 * TKT * 8; idx += TT) { // 16-byte pieces of the 128-byte tile words + const int q = idx / (TKT * 8), kk = (idx / 8) % TKT, part = idx % 8; + int wm, nt; + ttile(q, cb, wm, nt); + const uint32_t* src = (wm ? T1 : T0) + (((size_t)e * KT + chunk * TKT + kk) * NTILES + nt) * 32 + part * 4; + cp_async16(wdst + ((size_t)q * TKT + kk) * 32 + part * 4, src, true); + } + }; +#pragma unroll + for (int c = 0; c < TSTAGES - 1; ++c) { + if (c < chunks) load(c); + cp_async_commit(); + } + for (int c = 0; c < chunks; ++c) { + cp_async_wait(); + __syncthreads(); + if (c + TSTAGES - 1 < chunks) load(c + TSTAGES - 1); + cp_async_commit(); + const half* base = stage + (size_t)(c % TSTAGES) * SH; + const half* xs = base + (size_t)(XR == TPM ? 0 : mat) * TPM * TXS; + const uint32_t* ws = reinterpret_cast(base + (size_t)XR * TXS); +#pragma unroll + for (int kk = 0; kk < TKT; ++kk) { + uint32_t a[TMT][4]; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile_p(ws[((size_t)(warp * TMT + i) * TKT + kk) * 32 + lane], lane, b0, b1); + a[i][0] = b0[0]; a[i][1] = b1[0]; a[i][2] = b0[1]; a[i][3] = b1[1]; + } +#pragma unroll + for (int jb = 0; jb < TPM / 16; ++jb) { + if (jb < nblocks) { + uint32_t r[4]; + const int row = jb * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + ldmatrix_x4(r, xs + (size_t)row * TXS + kk * 16 + ((lane >> 3) & 1) * 8); + const uint32_t bl[2] = {r[0], r[1]}, bh[2] = {r[2], r[3]}; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + mma16816(acc[i][2 * jb], a[i], bl); + mma16816(acc[i][2 * jb + 1], a[i], bh); + } + } + } + } + } + cp_async_wait<0>(); + __syncthreads(); // the stages become the epilogue's sums +} + +// acc (members TPH * hf .. + TPH - 1) -> E[member - TPH * hf][col] fp32, row stride ES; col0: first tile's column. +__device__ __forceinline__ void tspill(float* E, int ES, int col0, const float (&acc)[TMT][TNB][4], int hf) { + const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int q = 0; q < TPM / 16; ++q) { + const int nb = hf * (TPM / 16) + q; + const int tok = q * 8 + 2 * t, col = col0 + i * 16 + g; + E[(size_t)tok * ES + col] = acc[i][nb][0]; + E[(size_t)(tok + 1) * ES + col] = acc[i][nb][1]; + E[(size_t)tok * ES + col + 8] = acc[i][nb][2]; + E[(size_t)(tok + 1) * ES + col + 8] = acc[i][nb][3]; + } +} + +template +__device__ __forceinline__ void tgateup(const half* __restrict__ Xg, const half* __restrict__ Xu, const uint32_t* __restrict__ Tg, + const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, + const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, + int K, int N, int xshift, float limit, unsigned char* smem_raw, int* rows_sh) { + const int cb = blockIdx.x, item = blockIdx.y; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + tmain<2, XR>(Xg, Xu, Tg, Tu, rows_sh, cnt, e, K, N >> 4, cb, xshift, reinterpret_cast(smem_raw), acc); + float* Eg = reinterpret_cast(smem_raw); + int mat, nt; + ttile<2>(warp * TMT, cb, mat, nt); + const int col0 = (nt - cb * 8) * 16; + const int n = cb * 128 + 4 * lane; + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); // the first half's rows are read + tspill(Eg + (size_t)mat * TPH * GU_ES, GU_ES, col0, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + const int p = rows_sh[hf * TPH + i]; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + gv[j] = Eg[(size_t)i * GU_ES + 4 * lane + j]; + uv[j] = Eg[(size_t)(TPH + i) * GU_ES + 4 * lane + j]; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float gg = fminf(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = fminf(fmaxf(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), -limit), limit); + float act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); + } + } +} + +} // namespace diff --git a/zig/kernels/cuda/glm/experts_sym.cu b/zig/kernels/cuda/glm/experts_sym.cu index 80f01eeeaa..90af4c071e 100644 --- a/zig/kernels/cuda/glm/experts_sym.cu +++ b/zig/kernels/cuda/glm/experts_sym.cu @@ -1,5 +1,4 @@ -// GLM-5.3-Flash's symmetric routed experts, decode form: flashnext/experts.cu reading experts_sym.cuh's blocks, each -// group's bias b = -8 s (exact), so acc = fma(xs, b, fma(p, s, acc)) per group in order gives the affine kernel's bits. +// GLM symmetric routed experts, decode: bias b = -8 s exact, so fma(xs, b, fma(p, s, acc)) gives the affine bits. #include #include diff --git a/zig/kernels/cuda/glm/experts_sym.cuh b/zig/kernels/cuda/glm/experts_sym.cuh index f350743112..72f0112755 100644 --- a/zig/kernels/cuda/glm/experts_sym.cuh +++ b/zig/kernels/cuda/glm/experts_sym.cuh @@ -1,6 +1,4 @@ -// GLM-5.3-Flash's symmetric routed experts: flashnext/experts.cuh with a group's block holding its scales alone (the -// bias of a symmetric int4 group is -8 x its scale, exact in bf16): 16 words a 32-row tile instead of 32. -// Shared by experts_sym.cu (decode form) and experts_sym_prefill.cu (prefill form). +// GLM symmetric experts' blocks for both forms: scales alone (bias -8 s, exact in bf16), 16 words a 32-row tile. #pragma once #include @@ -20,8 +18,7 @@ struct Geo { static constexpr int XV = GS / 32; // input uint4 a row a lane a group static constexpr int SBV = NTW / 4; // scale uint4 a quad a group (the bias: -8 x the scale) #ifdef GLM_SYM3 - // 3-bit codes (4..11 in the 4-bit words): each lane's uint4 c as three 32-bit planes (two of low 2-bit fields, - // one of high bits), 32 apart a plane: 24 uint4 a c instead of 32 + // 3-bit codes (4..11 in 4-bit words): a lane's uint4 as three 32-bit planes, 32 apart, 24 uint4 a c not 32 static constexpr int WSEC = 24 * WV; #else static constexpr int WSEC = 32 * WV; diff --git a/zig/kernels/cuda/glm/experts_sym_prefill.cu b/zig/kernels/cuda/glm/experts_sym_prefill.cu index fa15fb178a..4d7937876f 100644 --- a/zig/kernels/cuda/glm/experts_sym_prefill.cu +++ b/zig/kernels/cuda/glm/experts_sym_prefill.cu @@ -1,6 +1,4 @@ -// GLM-5.3-Flash's symmetric routed experts, prefill form: flashnext/experts_prefill.cu reading experts_sym.cuh's -// blocks, each group's bias pair -8 x its scale pair (exact in bf16). -// Prefill: w = bf16(fma(q, s, b)), one fp32 chain over K; chunk-invariant bits, not decode's (replies prefill again). +// GLM symmetric experts, prefill: w = bf16(fma(q, s, -8 s)), one fp32 chain over K; chunk-invariant, not decode's bits. #include #include diff --git a/zig/kernels/cuda/glm/index.cu b/zig/kernels/cuda/glm/index.cu index 31986f8705..a4ca0d7912 100644 --- a/zig/kernels/cuda/glm/index.cu +++ b/zig/kernels/cuda/glm/index.cu @@ -1,7 +1,4 @@ -// GLM-5.3-Flash on CUDA: the DSA indexer, ours. A row writes its index key (LayerNorm) and gate into its stream's -// rings, the row that completes a pool of 4 writes the pooled key, each row scores its visible pools, keeps its 512 -// best (the lower pool on ties) and lists their tokens ascending. Rows name their own rings, pools and positions, so -// one launch serves any mix of streams, and a row's choice never depends on the rows beside it. +// GLM DSA indexer, ours: rings, pools of 4, each row's 512 best pools; rows name their streams, never depend on others. #include #include @@ -24,8 +21,7 @@ __device__ __forceinline__ uint32_t order_key(float s) { } // namespace -// row r: LayerNorm(k_raw) -> bf16 at its position of its key ring ``ik[r]``, the fp32 gate -> bf16 into ``ig[r]`` -// (ring pointers already offset so that row p of the positions is at p * 128). 128 threads. +// row r: LayerNorm(k_raw) -> bf16 into key ring ``ik[r]``, fp32 gate -> bf16 into ``ig[r]`` (row p at p * 128). extern "C" __global__ void glm_index_write(const __nv_bfloat16* __restrict__ kr, long long k_stride, const float* __restrict__ gr, const __nv_bfloat16* __restrict__ lnw, const __nv_bfloat16* __restrict__ lnb, const unsigned long long* __restrict__ ik, const unsigned long long* __restrict__ ig, const int* __restrict__ pos, float eps) { @@ -50,9 +46,7 @@ extern "C" __global__ void glm_index_write(const __nv_bfloat16* __restrict__ kr, reinterpret_cast<__nv_bfloat16*>(ig[r])[p * D + d] = __float2bfloat16_rn(gr[(long long)r * D + d]); } -// the row that completes a pool (position % 4 == 3) writes its pooled key: the 4 keys weighted by the softmax of -// gate + ape, each product rounded to bf16, summed in order -> pk[r] row pos / 4. 128 threads. Run after -// glm_index_write over the same rows (a pool's earlier keys may be the window's or earlier ones still in the ring). +// pool-completing row (pos % 4 == 3): 4 keys by softmax(gate + ape), bf16 products summed in order -> pk; after write. extern "C" __global__ void glm_pool_keys(const unsigned long long* __restrict__ ik, const unsigned long long* __restrict__ ig, const __nv_bfloat16* __restrict__ ape, const unsigned long long* __restrict__ pk, const int* __restrict__ pos) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -79,9 +73,7 @@ extern "C" __global__ void glm_pool_keys(const unsigned long long* __restrict__ reinterpret_cast<__nv_bfloat16*>(pk[r])[p * D + d] = __float2bfloat16_rn(((t[0] + t[1]) + t[2]) + t[3]); } -// each row's pool scores: s_p = sum_h w_h relu(scale * qi_h . pk_p) for its visible pools (p < (pos + 1) / 4), -inf -// for the rest of the np columns. qi [R, 32, 128], w the head weights (bf16, row stride w_stride). grid (R, np / 64), -// 128 threads: a block of 64 pools on tensor cores, the heads summed in a fixed order. +// pool scores s_p = sum_h w_h relu(scale qi_h . pk_p), p < (pos + 1) / 4, else -inf; 64 pools a block, heads in order. extern "C" __global__ void __launch_bounds__(128) glm_index_scores(const __nv_bfloat16* __restrict__ qi, const __nv_bfloat16* __restrict__ w, long long w_stride, const unsigned long long* __restrict__ pk, float* __restrict__ out, const int* __restrict__ pos, int np, float scale, float wscale) { @@ -152,9 +144,7 @@ extern "C" __global__ void __launch_bounds__(128) glm_index_scores(const __nv_bf } } -// each row's 512 best pools of its visible ones (the lower pool on ties), ascending (int64): a radix select (8 bits a -// pass) finds the 512th best score, then one pass in pool order writes the pools above it and the lowest ties. -// scores [R, np]. 1024 threads. +// a row's 512 best visible pools ascending (int64, lower pool on ties): 8-bit radix select, one ordered pass. extern "C" __global__ void __launch_bounds__(1024) glm_index_select(const float* __restrict__ sc, long long* __restrict__ out, int np, const int* __restrict__ pos) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -228,8 +218,7 @@ extern "C" __global__ void __launch_bounds__(1024) glm_index_select(const float* } } -// each row's attended tokens (its 512 pools' tokens ascending, then the incomplete last pool's visible tokens, -1 past -// the row) and their count: 2048 plus those when the row has more than 512 complete pools, else 0 (a dense row). +// a row's tokens: its pools' ascending, then the last pool's visible ones (-1 past); count 2048 + those, 0 when dense. extern "C" __global__ void glm_index_tokens(const long long* __restrict__ pools, const int* __restrict__ pos, int* __restrict__ tokens, int* __restrict__ counts) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights diff --git a/zig/kernels/cuda/glm/kda.cu b/zig/kernels/cuda/glm/kda.cu index 577923d301..01f8529e92 100644 --- a/zig/kernels/cuda/glm/kda.cu +++ b/zig/kernels/cuda/glm/kda.cu @@ -1,15 +1,4 @@ -// From src/tensorfold/families/glm5_next/cuda/kda.cu (python-0.6), with the chain's body a device function and a -// segmented chain (any number of streams' windows in one launch) added. -// GLM-5.3-Flash's KDA (Kimi delta attention) on CUDA: one block of 1024 threads per head, a chain of R rows. -// -// Per row, following the Hugging Face definition: the depthwise conv over [conv state; q|k|v rows] (4 taps, -// fp32) with SiLU and one bf16 rounding; fp32 L2 norms of q and k (eps inside the sum, q times DK^-0.5); the -// per-channel decay g_i = exp(lower * sigmoid(exp(A_log) * (a_i + dt_bias_i))) in fp32; beta = bf16(sigmoid(b)); -// the delta rule in fp32 (decay along the key channel, read with k, correct toward v, read out with q); -// read-out rounded to bf16; then the gated RMSNorm bf16(w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(gate)). -// Warp w owns value rows 4w .. 4w + 3, lane l the key columns 4l .. 4l + 3. The state update is one routine -// (``update``) shared with ``replay``, compiled without FMA contraction, so a replayed prefix of a window gives -// the bits of the serial steps. +// GLM KDA from python-0.6 glm5_next/cuda/kda.cu plus a segmented chain; HF definition, update shared with replay. #include #include @@ -164,8 +153,7 @@ __global__ void __launch_bounds__(1024) chain_kernel( } // namespace -// One stream's window in a batched step: its rows (from ``first`` of the step's rows), its states (the live parity -// in, the other out), conv windows, replay inputs and window projections, each the base of layer 0. +// One stream's window in a batched step: rows from ``first``, states, conv, replay inputs and projections (layer 0). struct KdaSeg { int first, rows; const float* rec_in; @@ -178,8 +166,7 @@ struct KdaSeg { __nv_bfloat16* proj; }; -// chain_kernel for every stream of a batched step at layer ``li``: grid (H, streams). The segment's projection rows -// (the conv's channels) are kept in its stream's window buffer for the commit's conv shift. +// chain_kernel for every stream of a batched step at ``li``: grid (H, streams), projection rows kept for the shift. extern "C" __global__ void __launch_bounds__(1024) glm_kda_chain_seg(const KdaSeg* __restrict__ segs, int li, int H, const __nv_bfloat16* __restrict__ P, int p_stride, int b_off, const __nv_bfloat16* __restrict__ A, int a_stride, const __nv_bfloat16* __restrict__ G, int g_stride, @@ -271,7 +258,7 @@ __global__ void __launch_bounds__(1024) replay_layers_kernel( for (int i = 0; i < 4; ++i) sout[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; } -// Long windows: state-independent work for all rows (prep_kernel, out_kernel), the delta rule alone row by row (step_kernel); chain_kernel's bits. +// Long windows: prep_kernel and out_kernel for all rows, the delta rule row by row (step_kernel); chain_kernel's bits. // Block (row, head), 512 threads: conv + SiLU of the 3 x 128 q|k|v channels, the decay, beta, then the q/k norms. __global__ void __launch_bounds__(512) prep_kernel( @@ -324,7 +311,7 @@ __global__ void __launch_bounds__(512) prep_kernel( } } -// Block (head, part): each warp owns 4 value rows; TR rows of saved k, q, decay, v and beta are staged in shared memory at a time. +// Block (head, part): a warp owns 4 value rows; TR rows of saved k, q, decay, v and beta staged in shared at a time. template __global__ void __launch_bounds__(WARPS * 32) step_kernel( int H, const float* __restrict__ state_in, const float* __restrict__ q_in, const float* __restrict__ k_save, diff --git a/zig/kernels/cuda/glm/latent.cu b/zig/kernels/cuda/glm/latent.cu index b9d68ef37d..ba07302437 100644 --- a/zig/kernels/cuda/glm/latent.cu +++ b/zig/kernels/cuda/glm/latent.cu @@ -1,8 +1,4 @@ -// GLM-5.3-Flash on CUDA: MLA attention on the 512-wide latent cache, ours. A program is one row's 32 heads over one -// chunk of 512 keys (online softmax in 16-key tiles, tensor-core scores and outputs), and the chunks are merged in -// order; nothing a program computes depends on another row, so a row has the same bits alone, in a decode window, in -// sixteen streams' windows or in a prompt chunk. Every row names its own cache (``lc``, the latent rows of its -// stream) and position (``pos``), so one launch serves any mix of streams. +// GLM MLA on the 512-wide latent cache, ours: a row's 32 heads a 512-key chunk, merged in order; row-stable bits. #include #include @@ -204,8 +200,7 @@ extern "C" __global__ void glm_attn_smem(int* out) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights out[0] = SMEM; } -// the chunk partials of row r, head h in chunk order -> out [R, H, LW] bf16. dense rows past the limit and sparse -// rows with no tokens are left alone (the other path writes them). grid (R, H), 128 threads. +// row r head h's chunk partials in order -> out [R, H, LW] bf16; skips dense rows past limit and empty sparse rows. extern "C" __global__ void glm_attn_merge(const float* __restrict__ po, const float* __restrict__ pm, const float* __restrict__ pl, __nv_bfloat16* __restrict__ out, const int* __restrict__ pos, const int* __restrict__ cnt, int R, int nch, int sparse) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -242,8 +237,7 @@ extern "C" __global__ void glm_lat_write(const __nv_bfloat16* __restrict__ lat, for (int i = threadIdx.x; i < LW / 8; i += blockDim.x) reinterpret_cast(dst)[i] = src[i]; } -// absorb: qa[r, h, n] = sum_k q[r, h, k] wk[h, k, n] (fp32 in k order, then bf16); q [R, H, 256], wk [H, 256, 512]. -// grid (H, LW / 64, ceil(R / 16)), 64 threads: the weight block once in shared memory, each row summed alone. +// absorb: qa[r, h, n] = sum_k q[r, h, k] wk[h, k, n], fp32 in k order then bf16; weight block in shared, rows alone. extern "C" __global__ void glm_absorb_q(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ wk, __nv_bfloat16* __restrict__ qa, int R) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -263,8 +257,7 @@ extern "C" __global__ void glm_absorb_q(const __nv_bfloat16* __restrict__ q, con } } -// expand: out[r, h, n] = sum_k ol[r, h, k] wv[h, n, k]; ol [R, H, 512], wv [H, 256, 512] -> out [R, H, 256]. -// grid (H, 256 / 8, ceil(R / 16)), 256 threads: a warp an output, lanes over k in order, then a xor tree. +// expand: out[r, h, n] = sum_k ol[r, h, k] wv[h, n, k] -> [R, H, 256]; a warp an output, lanes over k, then a xor tree. extern "C" __global__ void glm_expand_v(const __nv_bfloat16* __restrict__ ol, const __nv_bfloat16* __restrict__ wv, __nv_bfloat16* __restrict__ out, int R) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights diff --git a/zig/kernels/cuda/glm/mm.cu b/zig/kernels/cuda/glm/mm.cu index d3a9f718c5..734246fb54 100644 --- a/zig/kernels/cuda/glm/mm.cu +++ b/zig/kernels/cuda/glm/mm.cu @@ -1,12 +1,4 @@ -// GLM-5.3-Flash on CUDA: the dense linears, ours. One tensor-core kernel for every row count (a decode window's few -// rows, sixteen streams' windows at once, a prompt chunk's 2048): each group of 64 inputs is summed from zero by the -// bf16 mma, then acc = fma(P, scale, acc) in group order, and the K slices (fixed by the shape) are added in slice -// order, so a row's bits never depend on the rows beside it. -// -// Weights are tiled at load (glm_mm_tile*): a column tile of 8 outputs and a group of 64 inputs is 32 lanes' chunks of -// 16 values in mma B-fragment order (lane l: column l / 4; value 4 kt + j at input 16 kt + {2t, 2t+1, 2t+8, 2t+9}[j], -// t = l % 4), each chunk BITS / 2 words (ops.cu's 16-value packing; BITS 16: the bf16 values themselves). Scales are -// fp16 [N / 8][K / 64][8]. +// GLM dense linears, ours: one tensor-core kernel for any row count, 64-groups then K slices in order; row-stable bits. #include #include @@ -136,8 +128,7 @@ __device__ __forceinline__ void mm_kernel(const __nv_bfloat16* __restrict__ x, l for (int j = 0; j < T::NT; ++j) #pragma unroll for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f; - // dependent launch: the first stages' weights stream in while the kernel before still runs; its outputs (this - // one's input, and any buffer this one writes) only after griddepcontrol.wait + // dependent launch: first stages' weights stream in early; inputs and outputs only after griddepcontrol.wait grid_launch(); #pragma unroll for (int st = 0; st < STAGES - 1; ++st) @@ -255,8 +246,7 @@ GLM_MM_ALL(6) GLM_MM_ALL(8) GLM_MM_ALL(16) -// -- tiling at load --------------------------------------------------------------------------------------------------- -// packed rows (ops.cu's [N][K / 16][BITS / 2] words, fp16 scales [N][K / 64]) -> the tiled chunks and scales +// -- tiling at load: ops.cu's packed rows and fp16 scales [N][K / 64] -> the tiled chunks and scales --------------- template __device__ __forceinline__ void tile_packed(const uint32_t* __restrict__ src, const __half* __restrict__ ss, uint32_t* __restrict__ dst, __half* __restrict__ ds, int N, int K) { diff --git a/zig/kernels/cuda/glm/ops.cu b/zig/kernels/cuda/glm/ops.cu index 14038c496e..7daf986b64 100644 --- a/zig/kernels/cuda/glm/ops.cu +++ b/zig/kernels/cuda/glm/ops.cu @@ -4,9 +4,7 @@ #include #include -// sparse.select_tokens' token list: row r (position q = pos[0] + r) attends to its 512 pools' tokens, ascending -// (pools[r] ascending, int64), then the incomplete last pool's visible tokens (-1 past q); count = 2048 + those when -// the row has more than 512 complete pools, else 0 (the row stays dense). +// sparse.select_tokens: row r's 512 pools' tokens ascending, then the last pool's visible ones; 0 count when dense. extern "C" __global__ void glm_sparse_tokens(const int64_t* __restrict__ pools, const int* __restrict__ pos, int rows, int* __restrict__ tokens, int* __restrict__ counts) { const int r = blockIdx.x; @@ -47,11 +45,7 @@ extern "C" __global__ void glm_copy_rows(const uint32_t* __restrict__ src, long for (int c = blockIdx.x * blockDim.x + threadIdx.x; c < words; c += gridDim.x * blockDim.x) dst[r * dpitch + c] = src[r * spitch + c]; } -// -- packed dense linears (the loader's TF_GLM_DENSE_BITS) ------------------------------------------------------------ -// A row is K / 16 chunks of 16 symmetric BITS-bit values (stored + 2^(BITS-1)), each chunk BITS / 2 32-bit words, with -// an fp16 scale per 64 inputs (4 chunks). 4 bits: value i at bits 4 (i % 8) of word i / 8. 8 bits: byte i % 4 of word -// i / 4. 6 bits: values 0-4, 5-9, 10-14 at bits 6 (i % 5) of words 0, 1, 2; value 15's bits 0-1, 2-3, 4-5 at bits -// 30-31 of words 0, 1, 2. +// -- packed dense linears (TF_GLM_DENSE_BITS): K / 16 chunks of 16 offset BITS-bit values, fp16 scale a 64 ---------- template __device__ __forceinline__ void glm_unpack16(const uint32_t* __restrict__ p, float v[16]) { @@ -94,9 +88,7 @@ __device__ __forceinline__ void glm_load_chunk(const uint32_t* __restrict__ p, u } } -// out[m, n] = sum_k x[m, k] q[n, k] s[n, k / 64] for M <= 16 rows (x bf16, row stride xs elements, 16-byte aligned; -// out [M, N] bf16 or fp32): a warp a row, lanes over the chunks, each chunk's 16 products summed in fp32 then scaled, -// the lanes' sums reduced by xor shuffles. A row's sum depends only on its own inputs, never on M. +// out[m, n] = sum_k x[m, k] q[n, k] s[n, k / 64], M <= 16: a warp a row, fp32 chunks, xor reduce; independent of M. template __device__ __forceinline__ void glm_qgemv(const __nv_bfloat16* __restrict__ x, long long xs, const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, int f32, int M, int N, int K) { @@ -183,8 +175,7 @@ GLM_Q(4) GLM_Q(6) GLM_Q(8) -// x rows 0 .. rows - 1 (each ``copies`` streams of 4096 bf16) take feats rows 0 .. rows - 1 [rows, 4096] in every stream: -// an image's features in place of its placeholder tokens' embeddings +// x rows 0 .. rows - 1 (``copies`` streams of 4096 bf16 each) take feats [rows, 4096]: image features for placeholders extern "C" __global__ void glm_inject_rows(const __nv_bfloat16* __restrict__ feats, __nv_bfloat16* __restrict__ x, int rows, int copies) { const long long n = (long long)rows * copies * 4096; diff --git a/zig/kernels/cuda/glm/qmm_frag.cuh b/zig/kernels/cuda/glm/qmm_frag.cuh index f4a4f457a7..77a424cffb 100644 --- a/zig/kernels/cuda/glm/qmm_frag.cuh +++ b/zig/kernels/cuda/glm/qmm_frag.cuh @@ -1,5 +1,4 @@ -// Device code of src/tensorfold/cuda/kernels/qmm_frag.cuh (generated by zig/tests/cuda/flashnext/gen_copies.py; do not edit). -// Shared by the 4-bit matmuls: cp.async staging, ldmatrix, the bf16 mma, nibble decoding from B-fragment order. +// Device code of cuda/kernels/qmm_frag.cuh (flashnext/gen_copies.py): cp.async, ldmatrix, bf16 mma, nibble decoding. #pragma once #include @@ -91,8 +90,7 @@ __device__ __forceinline__ uint32_t pair(uint32_t w, int s) { return r; } -// Programmatic dependent launch (sm_90+; no-ops before, and when the launch did not ask for it): wait for the -// previous kernel's writes, and let the next kernel's blocks start once every block here has said so. +// Programmatic dependent launch (sm_90+, else no-ops): wait for the previous kernel, release the next one's blocks. __device__ __forceinline__ void grid_wait() { #if __CUDA_ARCH__ >= 900 asm volatile("griddepcontrol.wait;\n" ::: "memory"); diff --git a/zig/kernels/cuda/glm/rows.cu b/zig/kernels/cuda/glm/rows.cu index 1a9b2e1d59..e9b505af6a 100644 --- a/zig/kernels/cuda/glm/rows.cu +++ b/zig/kernels/cuda/glm/rows.cu @@ -1,7 +1,4 @@ -// GLM-5.3-Flash on CUDA: the row steps between the linears, ours (embedding, RMSNorm, the hyper-connections, the -// residuals, SwiGLU, the router and its top-8, the experts' combine, the KDA conv window). Each row is one block's, -// with fp32 sums in a fixed order (threads strided over the row, then a fixed tree), so a row's bits depend on its own -// inputs alone, never on how many rows the window holds. +// GLM row steps between the linears, ours: a block a row, fp32 sums strided then a fixed tree; row-stable bits. #include #include @@ -40,8 +37,7 @@ extern "C" __global__ void glm_embed(const int* __restrict__ ids, const __nv_bfl } } -// RMSNorm (Llama: bf16(w * bf16(x * rinv))) of row r of x (stride xs_) -> out (stride os_), and the 64-groups' fp32 -// sums of the outputs when xs is given. D <= 4096. +// RMSNorm (Llama: bf16(w * bf16(x * rinv))) of row r -> out, and the outputs' 64-group fp32 sums when xs. D <= 4096. extern "C" __global__ void glm_rmsnorm(const __nv_bfloat16* __restrict__ x, long long x_stride, const __nv_bfloat16* __restrict__ w, __nv_bfloat16* __restrict__ out, long long o_stride, float* __restrict__ xs, float eps, int D) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -69,9 +65,7 @@ extern "C" __global__ void glm_rmsnorm(const __nv_bfloat16* __restrict__ x, long } } -// hc_pre (hyper-connections, 4 streams of D = 4096), in two launches. The partials: row r's 24 mixing dots of the -// streams [16384] with fn [24, 16384] and their sum of squares over 16 fixed slices of 1024 inputs, one block a -// (row, slice) -> part [R, 16, 25]. 128 threads. +// hc_pre partials: row r's 24 mixing dots and sum of squares over 16 fixed 1024-slices -> part [R, 16, 25]. 128 thr. extern "C" __global__ void __launch_bounds__(128) glm_hc_part(const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ fn, float* __restrict__ part) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -110,9 +104,7 @@ extern "C" __global__ void __launch_bounds__(128) glm_hc_part(const __nv_bfloat1 if (tid == 0) out[M] = sm; } -// The finish: the slices' partials in slice order, the pre / post / comb weights (comb: ``iters`` Sinkhorn steps), the -// collapsed row RMSNormed by nw -> out [D] and its 64-groups' sums xs; post [4] and comb [4, 4] kept for hc_post. -// 512 threads. +// hc_pre finish: partials in slice order, pre/post/comb (``iters`` Sinkhorn), RMSNormed row and xs; post, comb kept. extern "C" __global__ void __launch_bounds__(512) glm_hc_finish(const __nv_bfloat16* __restrict__ x, const float* __restrict__ part, const float* __restrict__ base, const float* __restrict__ scale, const __nv_bfloat16* __restrict__ nw, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, float* __restrict__ post, float* __restrict__ comb, @@ -183,8 +175,7 @@ extern "C" __global__ void __launch_bounds__(512) glm_hc_finish(const __nv_bfloa } } -// hc_post: the branch (both ranks' fp32 partials g [2, R, D], rank 0 first, rounded once) into each of the 4 streams: -// x_s = bf16(post_s * branch + sum_i comb[i][s] x_i), x in place. grid (R, D / 1024), 256 threads. +// hc_post: branch (ranks' fp32 partials, rank 0 first, rounded once): x_s = bf16(post_s b + sum_i comb[i][s] x_i). extern "C" __global__ void glm_hc_post(__nv_bfloat16* __restrict__ x, const float* __restrict__ g, const float* __restrict__ post, const float* __restrict__ comb, int R) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -228,8 +219,7 @@ extern "C" __global__ void glm_stream_mean(const __nv_bfloat16* __restrict__ x, } } -// SwiGLU over [gate | up] (bf16, width W each): bf16(bf16(silu(min(g, L))) * clip(u, -L, L)) and its 64-groups' sums. -// grid (R, W / 512), 256 threads. +// SwiGLU over [gate | up]: bf16(bf16(silu(min(g, L))) * clip(u, -L, L)) and 64-group sums; grid (R, W / 512), 256 thr. extern "C" __global__ void glm_swiglu(const __nv_bfloat16* __restrict__ gu, __nv_bfloat16* __restrict__ out, float* __restrict__ xs, float limit, int W) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -251,9 +241,7 @@ extern "C" __global__ void glm_swiglu(const __nv_bfloat16* __restrict__ gu, __nv } } -// router logits: out[m, e] = fp32 x[m] . w[e] (bf16 [NE, D]). A block takes 2 experts, a warp a quarter of D (D a -// multiple of 1024): its lanes over 8-input pieces in order, a xor tree, then the quarters summed in order: each row's -// sum is its own. grid (NE / 2, ceil(M / 16)), 256 threads. +// router logits fp32 x[m] . w[e]: 2 experts a block, a warp a D quarter, lanes then xor tree then quarters in order. extern "C" __global__ void __launch_bounds__(256) glm_router(const __nv_bfloat16* __restrict__ x, long long x_stride, const __nv_bfloat16* __restrict__ w, float* __restrict__ out, int M, int NE, int D) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -310,8 +298,7 @@ extern "C" __global__ void __launch_bounds__(256) glm_router(const __nv_bfloat16 } } -// each row's 8 experts by sigmoid score plus bias (the lower id on ties), weights the picks' scores normalized in pick -// order times ``scale``, then the shared expert (id NE, weight 1) in slot 8. A warp a row (NE <= 320). +// a row's 8 experts by sigmoid + bias (lower id on ties), weights normalized times ``scale``, shared expert in slot 8. extern "C" __global__ void glm_topk(const float* __restrict__ logits, const float* __restrict__ bias, int* __restrict__ pick, float* __restrict__ wts, float scale, int NE) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights @@ -389,8 +376,7 @@ extern "C" __global__ void glm_combine(const void* __restrict__ y, int y_f32, co } } -// the KDA conv window: conv [layers, 3, C] (in place) takes rows keep .. keep + 2 of [conv; proj rows] -// (proj [layers, R, >= C], strides in elements). grid (layers, ceil(C / 1024)), 256 threads. +// KDA conv window: conv [layers, 3, C] in place takes rows keep .. keep + 2 of [conv; proj]; grid (layers, C / 1024). extern "C" __global__ void glm_conv_shift(__nv_bfloat16* __restrict__ conv, const __nv_bfloat16* __restrict__ proj, int keep, long long conv_layer, long long proj_layer, long long proj_row, int C) { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // a dependent matmul may start on its weights diff --git a/zig/kernels/cuda/glm/vision.cu b/zig/kernels/cuda/glm/vision.cu index b5d5c44a3e..633df9c029 100644 --- a/zig/kernels/cuda/glm/vision.cu +++ b/zig/kernels/cuda/glm/vision.cu @@ -1,6 +1,4 @@ -// GLM-5.3-Flash's vision tower on CUDA (families/glm_cuda/vision.zig): the elementwise and row steps around its -// cuBLASLt products, as mlx-vlm's glm5_next/vision.py defines them. The residual stream stays in fp32; a step's bf16 -// output is rounded once, to nearest even. +// GLM vision tower's steps around cuBLASLt, as mlx-vlm glm5_next/vision.py: fp32 residual, bf16 outputs rounded once. #include #include @@ -68,10 +66,7 @@ extern "C" __global__ void glm_vis_rmsnorm(const float* __restrict__ x, const __ for (int c = threadIdx.x; c < cols; c += blockDim.x) y[(long long)blockIdx.x * cols + c] = tobf(row[c] * inv * bf(w[c])); } -// The attention inputs: qkv [N, 3, H, 64] bf16 -> q and k each RMSNorm'd over its head (q_norm, k_norm; bf16), then -// rotated (fp32: x cos + rotate_half(x) sin, angles (row, col) x inv_freq[16] for dims 0-15 / 16-31, again for -// 32-63) into qh, kh [N / seg, H, seg, 64]; v into vt [N / seg, H, 64, seg] (each frame group's head contiguous). -// A warp a (row, head): lane l holds dims l and l + 32. +// attention inputs: q, k head-RMSNormed and rotated (fp32, 2D angles) into qh, kh; v into vt; a warp a (row, head). extern "C" __global__ void glm_vis_qkv(const __nv_bfloat16* __restrict__ qkv, const __nv_bfloat16* __restrict__ qn, const __nv_bfloat16* __restrict__ kn, const int* __restrict__ pos, const float* __restrict__ inv_freq, __nv_bfloat16* __restrict__ qh, diff --git a/zig/src/families/glm_cuda/api.zig b/zig/src/families/glm_cuda/api.zig index 52f5098701..9f4886f371 100644 --- a/zig/src/families/glm_cuda/api.zig +++ b/zig/src/families/glm_cuda/api.zig @@ -1,10 +1,4 @@ -//! GLM-5.3-Flash (glm5_next) on CUDA, two ranks: the pieces the weight loader, the forward and the native host share. -//! -//! Weights are a name -> device tensor store whose names are the Python engine's prepared-weight names (its weight -//! tree: "layers.3.dsa.q_b.weight", "layers.5.moe.experts.up", ...), so every tensor the checkpoint loader builds can -//! be checked byte for byte against the Python engine's (tools: capture_ref.py's weights.json hashes). -//! Kernels: the fatbins embedded by the build (this family's, zig/kernels/cuda/glm; the shared grouped experts and the -//! greedy candidates, zig/kernels/cuda/flashnext). +//! GLM-5.3-Flash on CUDA, two ranks: the weight store (Python's prepared-weight names) and kernels the family shares. const std = @import("std"); const cuda = @import("cuda"); diff --git a/zig/src/families/glm_cuda/blocks.zig b/zig/src/families/glm_cuda/blocks.zig index 70049cb2dd..392ec59ccf 100644 --- a/zig/src/families/glm_cuda/blocks.zig +++ b/zig/src/families/glm_cuda/blocks.zig @@ -1,7 +1,4 @@ -//! GLM-5.3-Flash on CUDA: the weights' views and the layers of a step (forward.py's blocks): KDA and DSA attention, -//! the dense MLP and the MoE, the hyper-connections around them, the final norm and the head; and the MTP head. -//! Row-wise work covers the step's rows at once; the KDA recurrences, the attention and the indexer read the step's -//! tables (forward.prepare). +//! GLM-5.3-Flash on CUDA: weight views and a step's layers (KDA, DSA, MLP, MoE, hyper-connections, head, MTP head). const std = @import("std"); const cuda = @import("cuda"); @@ -24,7 +21,7 @@ const MOE_W = forward.MOE_W; const DENSE_W = forward.DENSE_W; const IDX_KW = forward.IDX_KW; -// -- weights ------------------------------------------------------------------------------------------------------------ +// -- weights --------------------------------------------------------------------------------------------------- const HCW = struct { fn_: u64, base: u64, scale: u64 }; const KdaW = struct { proj: kern.B16, fb: kern.B16, gb: kern.B16, conv: u64, a_log: u64, dt_bias: u64, norm: u64, o: kern.B16 }; const IndexW = struct { kw: kern.B16, qb: kern.B16, ln_w: u64, ln_b: u64, gate: u64, ape: u64 }; @@ -203,7 +200,7 @@ pub const Weights = struct { } }; -// -- the blocks ----------------------------------------------------------------------------------------------------------- +// -- the blocks ------------------------------------------------------------------------------------------------ pub fn mm(e: *Engine, b: *Buffers, x: u64, x_stride: i64, w: kern.B16, out: u64, f32_out: bool, R: i64) !void { if (e.skip.dense and w.n != forward.HEAD_N) return; try kern.mm(&e.k, x, x_stride, w, out, f32_out, R, b.sk); @@ -214,8 +211,7 @@ fn outProj(e: *Engine, b: *Buffers, x: u64, x_stride: i64, w: kern.B16, R: i64) return forward.gather(e, b, R); } -/// A KDA layer: a prompt chunk's rows through its one stream's states (a kept point inside the chunk cuts the -/// chain there), or a decode step's windows, each through its own stream's (the segmented chain). +/// A KDA layer: a prompt chunk through its stream's states (cut at a kept point), or windows on the segmented chain. fn kdaBlock(e: *Engine, b: *Buffers, l: LayerW, R: i64) !u64 { const k = l.kda.?; const li = l.li; @@ -262,8 +258,7 @@ fn selectTokens(e: *Engine, b: *Buffers, ci: usize, R: i64) !void { try kern.sparseTokens(&e.k, b.pools, t, R, b.tokens, b.counts); } -/// A DSA layer on latent cache ``ci`` (the MTP layer's is the last): the indexer always (its keys and pools are kept -/// for every position), dense attention for the rows below the dense limit, sparse for the rest. +/// A DSA layer on latent cache ``ci`` (MTP's last): indexer always, dense attention below the limit, sparse above. fn dsaBlock(e: *Engine, b: *Buffers, a: DsaW, ci: usize, R: i64) !u64 { const t = b.rowsTab(); try mm(e, b, b.normed, D, a.proj, b.dp, false, R); @@ -306,8 +301,7 @@ fn mlpBlock(e: *Engine, b: *Buffers, m: MlpW, R: i64) !u64 { return outProj(e, b, b.act, DENSE_W, m.down, R); } -/// The MoE: the routed slots through the grouped kernels (8 a row), the shared expert (slot 8) through the dense -/// linears, the slots combined by their weights. +/// The MoE: 8 routed slots on the grouped kernels, the shared expert (slot 8) on dense linears, combined by weight. fn moeBlock(e: *Engine, b: *Buffers, m: MoeW, R: i64) !u64 { try kern.router(&e.k, b.normed, D, m.router, b.mlog, R, kern.EXPERTS); try kern.topk(&e.k, b.mlog, m.bias, b.pick, b.wts, R); @@ -385,8 +379,7 @@ fn injectMedia(e: *Engine, s: *const sq.Seq, first: i64, R: i64, x: u64, copies: } } -/// The step's rows through every layer, the final norm, and the head (a prompt chunk: its last row only). The ids -/// are staged in b.ids, the tables prepared. +/// The step's rows through every layer, norm and head (a prompt chunk: last row only); ids in b.ids, tables ready. pub fn compute(e: *Engine, b: *Buffers) !void { const R = e.step.rows; const w = &e.w; @@ -407,8 +400,7 @@ pub fn compute(e: *Engine, b: *Buffers) !void { } } -/// The MTP head over the step's staged rows (ids in b.ids, input hidden rows in b.hin), then each stream's last row's -/// logits (in segment order). +/// The MTP head over the staged rows (ids in b.ids, hidden rows in b.hin), then each stream's last-row logits. pub fn mtpCompute(e: *Engine, b: *Buffers) !void { const m = e.w.mtp.?; const R = e.step.rows; diff --git a/zig/src/families/glm_cuda/calibrate.zig b/zig/src/families/glm_cuda/calibrate.zig index d4d6cfc31c..598ce4412c 100644 --- a/zig/src/families/glm_cuda/calibrate.zig +++ b/zig/src/families/glm_cuda/calibrate.zig @@ -1,5 +1,4 @@ -//! GLM-5.3-Flash's costs for the lane core, timed at load on rank 0 (rank 1 replays every call): a lone stream's -//! verify window by width, shared steps by their total rows (up to 16 streams' windows), and a chained head step. +//! GLM-5.3-Flash's lane-core costs timed on rank 0 at load: verify widths, shared steps by rows, a chained head step. const std = @import("std"); const lanes = @import("lanes"); @@ -39,8 +38,7 @@ pub fn run(self: anytype, drafts: bool) !void { } var picks: [forward.STEP_ROWS]decode.Pick = undefined; var ids: [streams][max_rows]u32 = undefined; - // one step of ``n`` streams' windows of ``width`` rows (distinct tokens: rows routed to their own experts, as a - // real window's are), each kept at its first row; its ms + // one step of ``n`` windows of ``width`` distinct-token rows (own experts), each kept at its first row; its ms const Step = struct { fn go(o: anytype, wr: *proto.Writer, ss: []*forward.Seq, tk: []u32, idv: *[streams][max_rows]u32, pk: []decode.Pick, n: usize, width: usize, rep: usize) !f64 { var wins: [streams]decode.Win = undefined; diff --git a/zig/src/families/glm_cuda/ckpt.zig b/zig/src/families/glm_cuda/ckpt.zig new file mode 100644 index 0000000000..353b07112b --- /dev/null +++ b/zig/src/families/glm_cuda/ckpt.zig @@ -0,0 +1,151 @@ +//! GLM-5.3-Flash's converted checkpoint: shard headers, files opened once, tensors read whole or as row ranges. + +const std = @import("std"); +const api = @import("api.zig"); + +pub const BASE = "model.language_model."; +pub const WORLD = 2; + +// -- the checkpoint: shard headers, files opened once, tensors read whole or as row ranges ------------------------ + +pub const Entry = struct { dtype: []const u8, shape: []const i64, begin: u64, end: u64, file: usize }; + +pub const Ckpt = struct { + gpa: std.mem.Allocator, + io: std.Io, + arena: std.heap.ArenaAllocator, + files: std.ArrayList(std.Io.File) = .empty, + paths: std.ArrayList([]const u8) = .empty, + entries: std.StringHashMap(Entry), + + /// ``dir``: the served folder; a shard symlinked under /cache/tf/ resolves against its parent (container or host). + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Ckpt { + var ck: Ckpt = .{ .gpa = gpa, .io = io, .arena = .init(gpa), .entries = .init(gpa) }; + errdefer ck.deinit(); + const a = ck.arena.allocator(); + const index_path = try resolve(a, io, dir, "model.safetensors.index.json"); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, a, .limited(1 << 26)); + const idx = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); + var shard_of = std.StringHashMap(usize).init(a); + var it = idx.object.get("weight_map").?.object.iterator(); + while (it.next()) |e| { + const shard = e.value_ptr.string; + if (shard_of.contains(shard)) continue; + const path = try resolve(a, io, dir, shard); + try shard_of.put(shard, ck.files.items.len); + try ck.files.append(gpa, try std.Io.Dir.cwd().openFile(io, path, .{})); + try ck.paths.append(gpa, path); + try ck.readHeader(ck.files.items.len - 1); + } + return ck; + } + + pub fn readHeader(ck: *Ckpt, f: usize) !void { + const a = ck.arena.allocator(); + var head: [8]u8 = undefined; + _ = try ck.files.items[f].readPositionalAll(ck.io, &head, 0); + const n = std.mem.readInt(u64, &head, .little); + const text = try a.alloc(u8, n); + _ = try ck.files.items[f].readPositionalAll(ck.io, text, 8); + const v = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); + var it = v.object.iterator(); + while (it.next()) |e| { + if (std.mem.eql(u8, e.key_ptr.*, "__metadata__")) continue; + const o = e.value_ptr.object; + const shp = o.get("shape").?.array.items; + const shape = try a.alloc(i64, shp.len); + for (shp, 0..) |x, i| shape[i] = x.integer; + const offs = o.get("data_offsets").?.array.items; + try ck.entries.put(e.key_ptr.*, .{ .dtype = o.get("dtype").?.string, .shape = shape, .begin = 8 + n + @as(u64, @intCast(offs[0].integer)), .end = 8 + n + @as(u64, @intCast(offs[1].integer)), .file = f }); + } + } + + pub fn deinit(ck: *Ckpt) void { + // the shards are read once: their page cache would starve NCCL's first registration (and the n-gram tables) + for (ck.files.items) |f| { + _ = std.os.linux.fadvise(f.handle, 0, 0, std.os.linux.POSIX_FADV.DONTNEED); + f.close(ck.io); + } + ck.files.deinit(ck.gpa); + ck.paths.deinit(ck.gpa); + ck.entries.deinit(); + ck.arena.deinit(); + } + + pub fn has(ck: *const Ckpt, name: []const u8) bool { + return ck.entries.contains(name); + } + + pub fn get(ck: *const Ckpt, name: []const u8) !Entry { + return ck.entries.get(name) orelse { + std.log.err("glm weights: the checkpoint has no tensor {s}", .{name}); + return error.MissingTensor; + }; + } + + /// ``len`` bytes at ``off`` into the tensor's data. + pub fn readAt(ck: *const Ckpt, e: Entry, off: u64, out: []u8) !void { + if (off + out.len > e.end - e.begin) return error.ShortRead; + const got = try ck.files.items[e.file].readPositionalAll(ck.io, out, e.begin + off); + if (got != out.len) return error.ShortRead; + } + + pub fn readAll(ck: *const Ckpt, gpa: std.mem.Allocator, name: []const u8) ![]align(16) u8 { + const e = try ck.get(name); + const buf = try gpa.alignedAlloc(u8, .@"16", e.end - e.begin); + errdefer gpa.free(buf); + try ck.readAt(e, 0, buf); + return buf; + } +}; + +pub fn resolve(a: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) ![]const u8 { + const path = try std.fs.path.join(a, &.{ dir, name }); + if (std.Io.Dir.cwd().openFile(io, path, .{})) |f| { + f.close(io); + return path; + } else |_| {} + var buf: [4096]u8 = undefined; + const n = std.Io.Dir.cwd().readLink(io, path, &buf) catch return path; + const target = buf[0..n]; + const container = "/cache/tf/"; + if (std.mem.startsWith(u8, target, container)) + return std.fs.path.join(a, &.{ std.fs.path.dirname(dir) orelse ".", target[container.len..] }); + if (std.fs.path.isAbsolute(target)) return a.dupe(u8, target); + return std.fs.path.join(a, &.{ dir, target }); +} + +pub fn bf16ToF32(b: u16) f32 { + return @bitCast(@as(u32, b) << 16); +} + +pub fn f32ToBf16(x: f32) u16 { + const u: u32 = @bitCast(x); + if (std.math.isNan(x)) return 0x7FC0; + return @intCast((u + 0x7FFF + ((u >> 16) & 1)) >> 16); +} + +pub fn f16ToBf16(h: u16) u16 { + return f32ToBf16(@floatCast(@as(f16, @bitCast(h)))); +} + +pub fn numelOf(shape: []const i64) usize { + var n: usize = 1; + for (shape) |d| n *= @intCast(d); + return n; +} + +pub fn dtypeOf(s: []const u8) !api.DType { + if (std.mem.eql(u8, s, "BF16")) return .bf16; + if (std.mem.eql(u8, s, "F32")) return .f32; + if (std.mem.eql(u8, s, "F16")) return .f16; + return error.UnsupportedDType; +} + +pub fn bf16Value(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +pub fn fullName(buf: []u8, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(buf, BASE ++ fmt, args); +} diff --git a/zig/src/families/glm_cuda/copy.zig b/zig/src/families/glm_cuda/copy.zig index de59affb39..ee58955a22 100644 --- a/zig/src/families/glm_cuda/copy.zig +++ b/zig/src/families/glm_cuda/copy.zig @@ -1,8 +1,4 @@ -//! Copy drafts (prompt lookup): code, JSON and edits repeat what came before, often longer than the MTP head drafts -//! with confidence. A stream keeps its prompt and every committed token, and an index of each 8-token run's latest -//! end; when the history's last 8 tokens ran before, the tokens that followed them are a draft. A run inside the -//! reply itself must match 16 tokens (a reply's own short repeats mislead more often). The drafts are verified like -//! the head's, so a wrong copy costs rows, never the reply. +//! Copy drafts (prompt lookup): a repeated 8-token run (16 in the reply) drafts what followed it, verified as MTP's. const std = @import("std"); @@ -19,8 +15,7 @@ pub const Copy = struct { k_prompt: usize = K, k_reply: usize = K_REPLY, - /// ``prompt`` kept and its runs indexed, at most ``most`` of its last tokens; a match takes ``k_prompt`` tokens - /// in the prompt and ``k_reply`` inside the reply. + /// ``prompt`` kept and indexed, at most ``most`` last tokens; matches take ``k_prompt`` or ``k_reply`` tokens. pub fn init(gpa: std.mem.Allocator, prompt: []const u32, most: usize, k_prompt: usize, k_reply: usize) !Copy { const kp = @max(2, k_prompt); const kr = @max(2, k_reply); diff --git a/zig/src/families/glm_cuda/decode.zig b/zig/src/families/glm_cuda/decode.zig index 23d0656d61..8f4d8d974d 100644 --- a/zig/src/families/glm_cuda/decode.zig +++ b/zig/src/families/glm_cuda/decode.zig @@ -1,5 +1,4 @@ -//! GLM-5.3-Flash on CUDA: prompts, verify steps, commits and MTP drafts (decode.py), for any set of streams. A prompt -//! runs alone in chunks of 2048 rows; verify steps and draft steps take every stream's rows at once. +//! GLM-5.3-Flash on CUDA (decode.py): prompts in 2048-row chunks; verify, commit and MTP draft steps for every stream. const std = @import("std"); const cuda = @import("cuda"); @@ -20,8 +19,7 @@ const HEAD_N = forward.HEAD_N; pub const Pick = struct { tok: u32, p: f64 }; -/// The gathered candidates -> each row's token (the larger logit; the lower id at a tie: rank 0's) and its -/// probability under the full softmax (both ranks' log-sum-exps). +/// The gathered candidates -> each row's token (larger logit, rank 0's at a tie) and its probability over both ranks. fn readPicks(e: *Engine, b: *Buffers, rows: i64, out: []Pick) !void { const d = e.ctx.d; const n: usize = @intCast(2 * rows * 4); @@ -48,10 +46,7 @@ fn readPicks(e: *Engine, b: *Buffers, rows: i64, out: []Pick) !void { /// A row's draw rule and the position its draw is keyed at (null rule: greedy, read from the candidates). pub const Draw = struct { rule: ?lanes.Sampling, position: u64 }; -/// The sampled rows' tokens: each rank's 64 candidates of every row and its log-sum-exps at the row's temperature -/// (fn_cand_topk), gathered, drawn on the host (draw.zig); a rule past the candidates draws over the row's gathered -/// logits. Both ranks run it alike, so their gathers pair up. -/// TF_GLM_TRACE: sampled rows' draw time (the candidates' kernel, gather and copy; the host's draws) and full rows. +/// Sampled rows' tokens: both ranks' 64 candidates and tempered lses, drawn on the host (draw.zig); TF_GLM_TRACE. pub var draw_stats: struct { gpu_ms: f64 = 0, host_ms: f64 = 0, rows: u64 = 0, full: u64 = 0, calls: u64 = 0, mismatch: u64 = 0 } = .{}; fn msFrom(e: *Engine, t: std.Io.Timestamp) f64 { @@ -65,8 +60,7 @@ fn drawRows(e: *Engine, b: *Buffers, rows: i64, draws: []const Draw, out: []Pick for (draws, 0..) |dr, r| inv[r] = if (dr.rule) |rule| @floatCast(1.0 / @max(rule.temperature, 1e-6)) else 1.0; try k.upload(b.invt, std.mem.sliceAsBytes(inv[0..@intCast(rows)])); const n: usize = @intCast(rows * draw_mod.W); - // the candidates by radix selection (fn_cand_topk_sel); TF_GLM_TOPK_CHECK=1: the per-thread lists' kernel - // (fn_cand_topk) first, its bytes compared with the selection's (dev check: the same 64, the same sums) + // radix-selected candidates; TF_GLM_TOPK_CHECK=1 also runs fn_cand_topk and compares their bytes (dev check) const check = topk_check orelse blk: { topk_check = std.c.getenv("TF_GLM_TOPK_CHECK") != null; break :blk topk_check.?; @@ -128,8 +122,7 @@ fn drawFull(e: *Engine, row: u64, rule: lanes.Sampling, position: u64) !u32 { return draw_mod.drawFullBf16(e.gpa, e.full_host, rule, position); } -/// The MTP head's rows of one prompt chunk (or a resume's pending row): ids and their hidden rows at mtp_len; the -/// caller advances mtp_len. +/// The MTP head's rows of a prompt chunk (or a resume's pending row) at mtp_len; the caller advances mtp_len. fn mtpForward(e: *Engine, b: *Buffers, s: *Seq, next: []const u32, hidden: u64) !void { const n: i64 = @intCast(next.len); try forward.stageIds(e, b, next); @@ -139,9 +132,7 @@ fn mtpForward(e: *Engine, b: *Buffers, s: *Seq, next: []const u32, hidden: u64) try blocks.mtpCompute(e, b); } -/// The prompt in chunks of 2048 rows from ``start`` (0, or the kept state's point: resumed), each chunk's rows -/// absorbed into the MTP cache, and the first token drawn. ``keep_at`` (0: none): the point the prompt's state is kept -/// at (the stream's kept state), cut inside its chunk when it falls inside one. +/// The prompt from ``start`` in 2048-row chunks into the MTP cache, then the first token; state kept at ``keep_at``. pub fn prefill(e: *Engine, s: *Seq, prompt: []const u32, start: usize, keep_at: usize) !u32 { defer s.media = null; // attach's images serve this prefill alone if (prompt.len == 0 or start >= prompt.len) return error.EmptyPrompt; @@ -168,8 +159,7 @@ pub fn prefill(e: *Engine, s: *Seq, prompt: []const u32, start: usize, keep_at: s.mtp_len += 1; } } - // the live caches change from here: a state kept below is the only one valid after (a resume kept at its own - // point keeps the state it resumed) + // the live caches change from here: a state kept below is the only valid one after (a resume keeps its own) if (!(start > 0 and keep_at == start)) s.snap_at = 0; s.d_from = @intCast(start); // DFlash2's ring holds this prompt's rows from here (a kept state's were overwritten since) var first: [1]Pick = undefined; @@ -231,8 +221,7 @@ pub fn prefill(e: *Engine, s: *Seq, prompt: []const u32, start: usize, keep_at: /// A stream's verify window: the pending token, then its drafts. pub const Win = struct { s: *Seq, ids: []const u32 }; -/// Every window through the model in one step (each at its stream's position); each row's token into ``out`` (rows -/// in window order). The windows stay uncommitted until ``keep``. +/// Every window in one step at its stream's position, each row's token into ``out``; uncommitted until ``keep``. pub fn verify(e: *Engine, wins: []const Win, out: []Pick) !void { var segs: [forward.MAX_STREAMS]Seg = undefined; var ids: [forward.STEP_ROWS]u32 = undefined; @@ -269,13 +258,10 @@ pub fn keep(e: *Engine, s: *Seq, R: i64, kept: i64) !void { s.pos += kept; } -/// A stream's draft request: its kept rows' final-normed rows (``hidden``, D bf16 apart) and the tokens after them, -/// then up to ``count`` drafts while their probabilities' product stays at or above ``confidence`` (the first always -/// kept). ``got`` and ``out``/``probs`` receive the drafts. +/// A stream's draft request: kept rows' normed ``hidden`` and tokens, up to ``count`` drafts above ``confidence``. pub const DraftReq = struct { s: *Seq, hidden: u64, next: []const u32, count: usize, confidence: f64, out: []u32, probs: []f64, got: usize = 0 }; -/// Every request's rows absorbed into its stream's MTP cache in one step, then the chains, one row a stream a step, -/// each stream leaving at its own confidence cut. +/// Every request's rows absorbed into its MTP cache in one step, then chains a row a step to each confidence cut. pub fn draft(e: *Engine, reqs: []DraftReq) !void { if (e.w.mtp == null) return; const b = &e.mbuf; diff --git a/zig/src/families/glm_cuda/dflash.zig b/zig/src/families/glm_cuda/dflash.zig index 018f0a1776..765e2b950f 100644 --- a/zig/src/families/glm_cuda/dflash.zig +++ b/zig/src/families/glm_cuda/dflash.zig @@ -1,17 +1,4 @@ -//! DFlash2 block drafts (opt-in: TF_GLM_DRAFTER=dflash2, TF_GLM_DFLASH2 the checkpoint's folder). IncoAI's drafter -//! (incoai/GLM-5.3-Flash-DFlash2, CC BY-NC-ND 4.0: non-commercial use; read from its own download, never changed on -//! disk or redistributed) after z-lab/dflash's DFlash2 architecture (MIT); this implementation is our own. -//! -//! The drafter is a 5-layer Qwen3-style decoder (hidden 4096, 32/8 heads of 128, SwiGLU 12288) that drafts a whole -//! block at once: rows [anchor, MASK x 7] at positions start .. start + 7 attend, bidirectionally, to each other and -//! to the stream's context: the target's completed residuals after layers 5, 14, 24, 33 and 42 (each the mean of its -//! 4 mHC streams), concatenated, through fc and a norm, then each layer's own k and v projections (a ring of the -//! last 2048 positions). Two-tap dynamic convolutions along the block wrap both sublayers. The block's last 7 rows -//! go through the target's head; a path selector picks each draft among the row's 16 likeliest tokens by the -//! previous draft's codebook row. Drafts never change a reply: the target verifies them. -//! -//! Both ranks run the drafter on their own copy (the head is each rank's vocabulary half, its candidates gathered); -//! rank 0 alone runs the selector, and the verify frame carries its drafts to rank 1. +//! DFlash2 drafts (TF_GLM_DRAFTER=dflash2): incoai/GLM-5.3-Flash-DFlash2 (CC BY-NC-ND 4.0, unmodified), z-lab (MIT). const std = @import("std"); const cuda = @import("cuda"); @@ -20,29 +7,31 @@ const forward = @import("forward.zig"); const blocks = @import("blocks.zig"); const draw_mod = @import("draw.zig"); const weights_mod = @import("weights.zig"); +const dl = @import("dflash_load.zig"); +const ds = @import("dflash_select.zig"); pub const TAPS = [_]usize{ 5, 14, 24, 33, 42 }; pub const BLOCK: usize = 8; pub const MASK: u32 = 154856; pub const WIN: i64 = 2048; pub const LAYERS: usize = 5; -const D: usize = 4096; -const KVW: usize = 1024; -const QKV: usize = 4096 + 2 * KVW; -const FF: usize = 12288; -const RANK: usize = 256; -const TOPK: usize = 16; +pub const D: usize = 4096; +pub const KVW: usize = 1024; +pub const QKV: usize = 4096 + 2 * KVW; +pub const FF: usize = 12288; +pub const RANK: usize = 256; +pub const TOPK: usize = 16; const THETA: f32 = 10000; const EPS: f32 = 1e-5; const SCALE: f32 = 0.08838834764831845; // 128^-0.5 -const VOCAB: usize = 154880; -const MAX_M: usize = forward.MAX_STREAMS * BLOCK; // block rows at most -const MAX_CTX: usize = @intCast(forward.PREFILL_ROWS_MAX); // context rows one ingest takes at most -const CW: usize = 2 * TOPK; // a row's gathered ints: 16 ids, then 16 values' bits -const CHUNKS: usize = 19; // 4096-column chunks of a rank's head rows at most -const QH: usize = 32; -const HD: usize = 128; -const CH: usize = 16; // key chunks of the chunked attention (df_attn_part) +pub const VOCAB: usize = 154880; +pub const MAX_M: usize = forward.MAX_STREAMS * BLOCK; // block rows at most +pub const MAX_CTX: usize = @intCast(forward.PREFILL_ROWS_MAX); // context rows one ingest takes at most +pub const CW: usize = 2 * TOPK; // a row's gathered ints: 16 ids, then 16 values' bits +pub const CHUNKS: usize = 19; // 4096-column chunks of a rank's head rows at most +pub const QH: usize = 32; +pub const HD: usize = 128; +pub const CH: usize = 16; // key chunks of the chunked attention (df_attn_part) fn cdiv(a: i64, b: i64) i64 { return @divFloor(a + b - 1, b); @@ -51,7 +40,7 @@ fn cdiv(a: i64, b: i64) i64 { /// Device bytes a stream's context ring takes (K and V of every layer). pub const ring_bytes: u64 = LAYERS * 2 * @as(u64, @intCast(WIN)) * KVW * 2; -const B16 = kern.B16; +pub const B16 = kern.B16; const LayerW = struct { in_norm: u64, post_norm: u64, @@ -68,8 +57,7 @@ const LayerW = struct { m_proj: B16, }; -/// One stream's draft request: ``n`` rows of the last verify to take into its context first (0: the prompt's, -/// taken at prefill), the anchor (the token the block starts with, at the stream's position); rank 0 gets ``out``. +/// One stream's draft request: ``n`` last-verify rows into its context first (0: prefill's), the anchor, ``out``. pub const Req = struct { s: *forward.Seq, n: usize, anchor: u32, out: ?[]u32 = null, probs: ?[]f64 = null, drafts: bool = true }; pub const Drafter = struct { @@ -201,304 +189,13 @@ fn lin(e: *forward.Engine, x: u64, w: B16, out: u64, m: usize, f32_out: bool) !v try kern.mm(&e.k, x, w.k, w, out, f32_out, @intCast(m), e.pbuf.sk); // the prompt buffers' split-K scratch: no prefill runs meanwhile } -// -- loading ------------------------------------------------------------------------------------------------------- +// -- loading: dflash_load.zig --------------------------------------------------------------------------------------- -const Ent = struct { off: u64, n: u64 }; - -/// TF_GLM_DFLASH2_BITS: the linears packed at load to 4 bits (the default: they only propose, and lands as many as -/// 8 bits did on the ruler), 6 or 8 (a scale per 64), 0 keeps bf16. -fn bitsFromEnv() !u8 { - const v = std.c.getenv("TF_GLM_DFLASH2_BITS") orelse return 4; - const b = std.fmt.parseInt(u8, std.mem.span(v), 10) catch return error.BadDFlash2Bits; - return switch (b) { - 0, 4, 6, 8 => b, - else => error.BadDFlash2Bits, - }; -} - -const Loader = struct { - gpa: std.mem.Allocator, - io: std.Io, - file: std.Io.File, - ents: *std.StringHashMap(Ent), - buf: std.ArrayList(u8) = .empty, - words: std.ArrayList(u32) = .empty, - scales: std.ArrayList(u16) = .empty, - sub: std.ArrayList(u16) = .empty, // a slice of a stored tensor (a rank's share of a split linear) - w: *cuda.DeviceBuffer, - tmp: *cuda.DeviceBuffer, - k: *kern.K, - at: u64 = 0, - bits: u8, - - fn take(l: *Loader, n: u64) u64 { - const p = l.at; - l.at += std.mem.alignForward(u64, n, 256); - return l.w.ptr + p; - } - - fn read(l: *Loader, name: []const u8, want: u64) ![]u8 { - const e = l.ents.get(name) orelse { - std.log.err("dflash2: no {s} in the checkpoint", .{name}); - return error.DrafterTensorMissing; - }; - if (e.n != want) { - std.log.err("dflash2: {s} holds {d} bytes, not {d}", .{ name, e.n, want }); - return error.DrafterTensorShape; - } - try l.buf.resize(l.gpa, e.n); - _ = try l.file.readPositionalAll(l.io, l.buf.items, e.off); - return l.buf.items; - } - - /// A small bf16 tensor as stored. - fn raw(l: *Loader, name: []const u8, n: u64) !u64 { - const p = l.take(n); - try l.w.upload(p - l.w.ptr, try l.read(name, n)); - return p; - } - - /// The tensors ``names`` (rows ``rows`` each, ``kk`` columns) stacked into one linear, packed and tiled for kern.mm. - fn linear(l: *Loader, names: []const []const u8, rows: []const usize, kk: usize) !B16 { - var n: usize = 0; - for (rows) |r| n += r; - const bits = l.bits; - const wb: usize = n * kk * (if (bits == 0) 16 else bits) / 8; - const sb: usize = if (bits == 0) 0 else n * kk / 64 * 2; - const dst = l.take(wb); - const ds: u64 = if (sb > 0) l.take(sb) else 0; - var r0: usize = 0; - if (bits == 0) { - for (names, rows) |name, r| { - try l.tmp.upload(r0 * kk * 2, try l.read(name, r * kk * 2)); - r0 += r; - } - } else { - try l.words.resize(l.gpa, n * kk * bits / 32); - try l.scales.resize(l.gpa, n * kk / 64); - for (names, rows) |name, r| { - const src = try l.read(name, r * kk * 2); - weights_mod.quantize(bits, @alignCast(std.mem.bytesAsSlice(u16, src)), r, kk, l.words.items[r0 * kk * bits / 32 ..][0 .. r * kk * bits / 32], l.scales.items[r0 * kk / 64 ..][0 .. r * kk / 64]); - r0 += r; - } - try l.tmp.upload(0, std.mem.sliceAsBytes(l.words.items)); - try l.tmp.upload(std.mem.alignForward(u64, wb, 256), std.mem.sliceAsBytes(l.scales.items)); - } - try kern.mmTile(l.k, bits, l.tmp.ptr, l.tmp.ptr + std.mem.alignForward(u64, wb, 256), dst, ds, @intCast(n), @intCast(kk)); - try l.k.stream.synchronize(); - return .{ .w = dst, .n = @intCast(n), .k = @intCast(kk), .s = ds, .bits = bits }; - } - - /// Rows ``row0 .. row0 + rn`` and columns ``col0 .. col0 + cn`` of each of the tensors ``names`` (rows ``rows`` - /// each, ``kk`` columns as stored), stacked into one linear of names.len * rn rows and cn columns: a rank's - /// share of a split linear, packed and tiled as ``linear`` packs a whole one. - fn linearPart(l: *Loader, names: []const []const u8, rows: []const usize, kk: usize, row0: usize, rn: usize, col0: usize, cn: usize) !B16 { - const n = names.len * rn; - const bits = l.bits; - const wb: usize = n * cn * (if (bits == 0) 16 else bits) / 8; - const sb: usize = if (bits == 0) 0 else n * cn / 64 * 2; - const dst = l.take(wb); - const ds: u64 = if (sb > 0) l.take(sb) else 0; - try l.sub.resize(l.gpa, rn * cn); - if (bits != 0) { - try l.words.resize(l.gpa, n * cn * bits / 32); - try l.scales.resize(l.gpa, n * cn / 64); - } - var r0: usize = 0; - for (names, rows) |name, r| { - if (row0 + rn > r or col0 + cn > kk) return error.DrafterTensorShape; - const src: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, try l.read(name, r * kk * 2))); - for (0..rn) |i| @memcpy(l.sub.items[i * cn ..][0..cn], src[(row0 + i) * kk + col0 ..][0..cn]); - if (bits == 0) { - try l.tmp.upload(r0 * cn * 2, std.mem.sliceAsBytes(l.sub.items)); - } else { - weights_mod.quantize(bits, l.sub.items, rn, cn, l.words.items[r0 * cn * bits / 32 ..][0 .. rn * cn * bits / 32], l.scales.items[r0 * cn / 64 ..][0 .. rn * cn / 64]); - } - r0 += rn; - } - if (bits != 0) { - try l.tmp.upload(0, std.mem.sliceAsBytes(l.words.items)); - try l.tmp.upload(std.mem.alignForward(u64, wb, 256), std.mem.sliceAsBytes(l.scales.items)); - } - try kern.mmTile(l.k, bits, l.tmp.ptr, l.tmp.ptr + std.mem.alignForward(u64, wb, 256), dst, ds, @intCast(n), @intCast(cn)); - try l.k.stream.synchronize(); - return .{ .w = dst, .n = @intCast(n), .k = @intCast(cn), .s = ds, .bits = bits }; - } -}; - -fn linearBytes(n: usize, kk: usize, bits: u8) u64 { - const wb = std.mem.alignForward(u64, n * kk * (if (bits == 0) 16 else @as(usize, bits)) / 8, 256); - return wb + if (bits == 0) 0 else std.mem.alignForward(u64, n * kk / 64 * 2, 256); -} - -/// The drafter from ``dir`` (config.json and model.safetensors, bf16), its linears packed for the engine's matmul. -pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *forward.Engine, dir: []const u8) !*Drafter { - try checkConfig(gpa, io, dir); - const bits = try bitsFromEnv(); - const d = e.ctx.d; - const path = try std.fs.path.join(gpa, &.{ dir, "model.safetensors" }); - defer gpa.free(path); - const file = try std.Io.Dir.cwd().openFile(io, path, .{}); - defer file.close(io); - var head: [8]u8 = undefined; - _ = try file.readPositionalAll(io, &head, 0); - const hlen = std.mem.readInt(u64, &head, .little); - if (hlen > 1 << 20) return error.BadDrafterHeader; - const text = try gpa.alloc(u8, hlen); - defer gpa.free(text); - _ = try file.readPositionalAll(io, text, 8); - const hdr = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); - defer hdr.deinit(); - var ents = std.StringHashMap(Ent).init(gpa); - defer ents.deinit(); - var it = hdr.value.object.iterator(); - while (it.next()) |x| { - if (std.mem.eql(u8, x.key_ptr.*, "__metadata__")) continue; - if (!std.mem.eql(u8, x.value_ptr.object.get("dtype").?.string, "BF16")) return error.DrafterNotBf16; - const offs = x.value_ptr.object.get("data_offsets").?.array.items; - try ents.put(x.key_ptr.*, .{ .off = 8 + hlen + @as(u64, @intCast(offs[0].integer)), .n = @intCast(offs[1].integer - offs[0].integer) }); - } - // the device's share: every linear packed (q|k|v and gate|up stacked, k|v again for the context), the norms and - // conv kernels as stored; the codebooks stay on the host - var total: u64 = 0; - const split = e.ctx.world == 2 and if (std.c.getenv("TF_GLM_DFLASH2_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; - const ff: usize = if (split) FF / 2 else FF; - const small: u64 = 6 * 4096 * 2 + 2 * 128 * 2 + 2 * 2 * 2 * D * 2; - total += LAYERS * (linearBytes(QKV, D, bits) + linearBytes(2 * KVW, D, bits) + linearBytes(D, D, bits) + linearBytes(2 * ff, D, bits) + - linearBytes(D, ff, bits) + 2 * linearBytes(KVW, D, bits) + small + 8 * 256); - total += linearBytes(D, 5 * D, bits) + linearBytes(RANK, D, bits) + 2 * D * 2 + 1024; - var wbuf = try cuda.DeviceBuffer.alloc(d, total); - errdefer wbuf.free(); - var tmp = try cuda.DeviceBuffer.alloc(d, linearBytes(2 * FF, D, 0) + 4096); // the largest linear as bf16 - defer tmp.free(); - var l: Loader = .{ .gpa = gpa, .io = io, .file = file, .ents = &ents, .w = &wbuf, .tmp = &tmp, .k = &e.k, .bits = bits }; - defer { - l.buf.deinit(gpa); - l.words.deinit(gpa); - l.scales.deinit(gpa); - l.sub.deinit(gpa); - } - const dr = try gpa.create(Drafter); - errdefer gpa.destroy(dr); - dr.* = .{ .gpa = gpa, .d = d, .weights = wbuf, .bits = bits, .split = split, .ff = ff, .pred = &.{}, .succ = &.{}, .scratch = undefined, .h_ids = undefined, .h_bpos = undefined, .h_rings = undefined, .h_starts = undefined, .h_froms = undefined, .h_cpos = undefined, .h_crings = undefined, .cand = undefined, .gh = undefined }; - var nb: [8][96]u8 = undefined; - for (&dr.layers, 0..) |*lw, i| { - const N = struct { - fn n(b: *[96]u8, li: usize, s: []const u8) ![]const u8 { - return std.fmt.bufPrint(b, "layers.{d}.{s}", .{ li, s }); - } - }.n; - lw.in_norm = try l.raw(try N(&nb[0], i, "input_layernorm.weight"), D * 2); - lw.post_norm = try l.raw(try N(&nb[0], i, "post_attention_layernorm.weight"), D * 2); - lw.qn = try l.raw(try N(&nb[0], i, "self_attn.q_norm.weight"), 128 * 2); - lw.kn = try l.raw(try N(&nb[0], i, "self_attn.k_norm.weight"), 128 * 2); - lw.a_base = try l.raw(try N(&nb[0], i, "attention_conv.base_kernel"), 2 * 2 * D * 2); - lw.m_base = try l.raw(try N(&nb[0], i, "mlp_conv.base_kernel"), 2 * 2 * D * 2); - const q = try N(&nb[1], i, "self_attn.q_proj.weight"); - const kq = try N(&nb[2], i, "self_attn.k_proj.weight"); - const v = try N(&nb[3], i, "self_attn.v_proj.weight"); - lw.qkv = try l.linear(&.{ q, kq, v }, &.{ D, KVW, KVW }, D); - lw.kv = try l.linear(&.{ kq, v }, &.{ KVW, KVW }, D); - lw.o = try l.linear(&.{try N(&nb[1], i, "self_attn.o_proj.weight")}, &.{D}, D); - const gate = try N(&nb[1], i, "mlp.gate_proj.weight"); - const up = try N(&nb[2], i, "mlp.up_proj.weight"); - const down = try N(&nb[3], i, "mlp.down_proj.weight"); - if (split) { // this rank's half of the MLP width: gate and up rows, down columns - lw.gu = try l.linearPart(&.{ gate, up }, &.{ FF, FF }, D, e.ctx.rank * ff, ff, 0, D); - lw.down = try l.linearPart(&.{down}, &.{D}, FF, 0, D, e.ctx.rank * ff, ff); - } else { - lw.gu = try l.linear(&.{ gate, up }, &.{ FF, FF }, D); - lw.down = try l.linear(&.{down}, &.{D}, FF); - } - lw.a_proj = try l.linear(&.{try N(&nb[1], i, "attention_conv.kernel_projection.weight")}, &.{KVW}, D); - lw.m_proj = try l.linear(&.{try N(&nb[1], i, "mlp_conv.kernel_projection.weight")}, &.{KVW}, D); - } - dr.fc = try l.linear(&.{"fc.weight"}, &.{D}, 5 * D); - dr.hproj = try l.linear(&.{"candidate_selector.hidden_projection.weight"}, &.{RANK}, D); - dr.hnorm = try l.raw("hidden_norm.weight", D * 2); - dr.norm = try l.raw("norm.weight", D * 2); - if (l.at > total) return error.DrafterLayout; - dr.pred = try gpa.alloc(u16, VOCAB * RANK); - @memcpy(std.mem.sliceAsBytes(dr.pred), try l.read("candidate_selector.predecessor_codebook", VOCAB * RANK * 2)); - dr.succ = try gpa.alloc(u16, VOCAB * RANK); - @memcpy(std.mem.sliceAsBytes(dr.succ), try l.read("candidate_selector.successor_codebook", VOCAB * RANK * 2)); - try carveScratch(dr, d); - dr.h_ids = try gpa.alloc(i32, MAX_M); - dr.h_bpos = try gpa.alloc(i32, MAX_M); - dr.h_rings = try gpa.alloc(u64, forward.MAX_STREAMS); - dr.h_starts = try gpa.alloc(i64, forward.MAX_STREAMS); - dr.h_froms = try gpa.alloc(i64, forward.MAX_STREAMS); - dr.h_cpos = try gpa.alloc(i64, MAX_CTX); - dr.h_crings = try gpa.alloc(u64, MAX_CTX); - dr.cand = try gpa.alloc(i32, 2 * MAX_M * CW); - dr.gh = try gpa.alloc(f32, MAX_M * RANK); - dr.prof = std.c.getenv("TF_GLM_DFLASH2_PROF") != null; - if (std.c.getenv("TF_GLM_DFLASH2_ATTN")) |v| dr.attn_chunked = !std.mem.eql(u8, std.mem.span(v), "serial"); - dr.prof_fine = if (std.c.getenv("TF_GLM_DFLASH2_PROF")) |v| std.mem.eql(u8, std.mem.span(v), "2") else false; - dr.dump = std.c.getenv("TF_GLM_DFLASH2_DUMP"); - if (std.c.getenv("TF_GLM_DFLASH2_TAP")) |v| { - dr.tap_input = std.mem.eql(u8, std.mem.span(v), "input"); - dr.tap_s0 = std.mem.eql(u8, std.mem.span(v), "s0"); - } - if (std.c.getenv("TF_GLM_DFLASH2_DUMP_ROUND")) |v| dr.dump_at = std.fmt.parseInt(u32, std.mem.span(v), 10) catch 10; - if (std.c.getenv("TF_GLM_DFLASH2_HEAD")) |v| dr.full_head = std.mem.eql(u8, std.mem.span(v), "full"); - if (std.c.getenv("TF_GLM_DFLASH2_TEMP")) |v| dr.temp = std.fmt.parseFloat(f64, std.mem.span(v)) catch 1.0; - const inv: [MAX_M]f32 = @splat(1.0); - try dr.scratch.upload(dr.invt - dr.scratch.ptr, std.mem.sliceAsBytes(&inv)); - std.log.info("dflash2: drafter loaded from {s}, linears at {d} bits ({d:.2} GiB on the device, {d} MiB of codebooks on the host)", .{ dir, bits, @as(f64, @floatFromInt(l.at)) / (1 << 30), (2 * VOCAB * RANK * 2) >> 20 }); - return dr; -} - -fn carveScratch(dr: *Drafter, d: *const cuda.Driver) !void { - var at: u64 = 0; - const M: u64 = MAX_M; - const C: u64 = MAX_CTX + 16; - const fields = .{ - .{ "x", M * D * 2 }, .{ "y", M * D * 2 }, .{ "y2", M * D * 2 }, .{ "dyn", M * KVW * 2 }, - .{ "qkv", M * QKV * 2 }, .{ "q", M * D * 2 }, .{ "kb", M * KVW * 2 }, .{ "vb", M * KVW * 2 }, - .{ "att", M * D * 2 }, .{ "a", M * D * 2 }, .{ "gu", M * 2 * FF * 2 }, .{ "act", M * FF * 2 }, - .{ "hn", M * D * 2 }, .{ "g", M * RANK * 4 }, .{ "cin", C * 5 * D * 2 }, .{ "c1", C * D * 2 }, - .{ "ctx", C * D * 2 }, .{ "ckv", C * 2 * KVW * 2 }, .{ "taps_step", @as(u64, @intCast(forward.STEP_ROWS)) * 5 * D * 2 }, .{ "taps_prefill", @as(u64, @intCast(forward.PREFILL_ROWS_MAX)) * 5 * D * 2 }, - .{ "ids", M * 4 }, .{ "bpos", M * 4 }, .{ "rings", forward.MAX_STREAMS * 8 }, .{ "starts", forward.MAX_STREAMS * 8 }, - .{ "froms", forward.MAX_STREAMS * 8 }, .{ "cpos", C * 8 }, .{ "crings", C * 8 }, .{ "invt", M * 4 }, - .{ "tk", M * CW * 4 }, .{ "gtk", 2 * M * CW * 4 }, .{ "pv", M * CHUNKS * TOPK * 4 }, .{ "pi", M * CHUNKS * TOPK * 4 }, - .{ "pm", M * QH * CH * 4 }, .{ "pl", M * QH * CH * 4 }, .{ "pacc", M * QH * CH * HD * 4 }, - }; - inline for (fields) |f| { - @field(dr, f[0]) = at; - at += std.mem.alignForward(u64, f[1], 256); - } - dr.scratch = try cuda.DeviceBuffer.alloc(d, at); - inline for (fields) |f| @field(dr, f[0]) += dr.scratch.ptr; - try d.check(d.api.cuMemsetD8_v2(dr.scratch.ptr, 0, at), "zero drafter scratch"); -} - -/// The checkpoint is the one this was written for (its shape and DFlash2 settings). -fn checkConfig(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 text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); - defer gpa.free(text); - const p = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); - defer p.deinit(); - const o = p.value.object; - const df = (o.get("dflash_config") orelse return error.NotDFlash2).object; - const ok = o.get("hidden_size").?.integer == D and o.get("num_hidden_layers").?.integer == LAYERS and - o.get("num_attention_heads").?.integer == 32 and o.get("num_key_value_heads").?.integer == 8 and - o.get("intermediate_size").?.integer == FF and o.get("vocab_size").?.integer == VOCAB and - df.get("block_size").?.integer == BLOCK and df.get("mask_token_id").?.integer == MASK and - df.get("selector_top_k").?.integer == TOPK and df.get("selector_rank").?.integer == RANK and - df.get("conv_kernel_size").?.integer == 2 and df.get("conv_group_size").?.integer == 16; - if (!ok) return error.DrafterConfigUnexpected; - const taps = df.get("target_layer_ids").?.array.items; - if (taps.len != TAPS.len) return error.DrafterConfigUnexpected; - for (taps, TAPS) |t, w| if (t.integer != w) return error.DrafterConfigUnexpected; -} +pub const load = dl.load; // -- the target's side --------------------------------------------------------------------------------------------- -/// After target layer ``li``: its rows' 4-stream means into the taps (a verify step's or a prompt chunk's). -/// TF_GLM_DFLASH2_TAP (development): `mean` (the default), `input` (layer k's input: k - 1's output), `s0` (stream 0). +/// After target layer ``li``: its rows' 4-stream means into the taps; TF_GLM_DFLASH2_TAP=mean|input|s0 (development). pub fn tap(dr: *const Drafter, k: *kern.K, x: u64, taps: u64, li: usize, R: i64) !void { const at = li + @intFromBool(dr.tap_input); const slot: i32 = for (TAPS, 0..) |t, i| { @@ -521,9 +218,7 @@ pub fn isTap(dr: *const Drafter, li: usize) bool { /// Context runs: ``n`` tap rows from ``taps`` (consecutive) at positions pos0 .. pos0 + n, into ``s``'s ring. pub const Run = struct { s: *forward.Seq, taps: u64, n: usize, pos0: i64 }; -/// The runs' rows through fc, the hidden norm and every layer's k and v, into their rings (and synchronized: the -/// host tables are reused). -/// TF_GLM_DFLASH2_DUMP: a whole one-chunk prompt's taps ([n][5 D] bf16, positions 0 .. n), once. +/// The runs' rows through fc, the norm and each layer's k, v into their rings; TF_GLM_DFLASH2_DUMP dumps one prompt. pub fn dumpPrompt(dr: *Drafter, e: *forward.Engine, n: usize) !void { if (dr.dump == null or dr.dump_state == 2 or n > MAX_CTX) return; // the latest prompt until its first round is written try e.k.stream.synchronize(); @@ -598,10 +293,7 @@ pub fn ingest(dr: *Drafter, e: *forward.Engine, runs: []const Run) !void { // -- drafting ------------------------------------------------------------------------------------------------------ -/// Every request's block drafted on this rank (both ranks, in step: the head's candidates gather); rank 0 (the -/// requests' ``out``) gets 7 drafts and their chances. -/// Every request's kept rows into its context; then the block pass for the ones that draft (``drafts``: the others -/// draft with the MTP head this round and only keep their context current). +/// Kept rows into each context, then the block pass for drafting streams (both ranks); rank 0 gets 7 drafts each. pub fn draft(dr: *Drafter, e: *forward.Engine, all: []const Req) !void { const k = &e.k; if (all.len == 0) return; @@ -676,7 +368,7 @@ pub fn draft(dr: *Drafter, e: *forward.Engine, all: []const Req) !void { if (rank0) { try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(dr.cand.ptr), dr.gtk, 2 * nw * 4), "dflash2 candidates"); try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(dr.gh.ptr), dr.g, M * RANK * 4), "dflash2 projections"); - for (reqs, 0..) |q, i| select(dr, i, M, q); + for (reqs, 0..) |q, i| ds.select(dr, i, M, q); if (dr.dump != null and dr.dump_state == 1 and S == 1) dr.dump_round += 1; if (dr.dump != null and dr.dump_state == 1 and S == 1 and dr.dump_round == dr.dump_at) { const hn = try dr.gpa.alloc(u8, M * D * 2); @@ -700,8 +392,7 @@ pub fn draft(dr: *Drafter, e: *forward.Engine, all: []const Req) !void { dr.calls += 1; } -/// The block's pass for S streams over the uploaded tables: embed, the 5 layers, the final norm, the draft head's 16 -/// best a row gathered from both ranks, and (rank 0) the selector's projections. Capturable. +/// The block pass for S streams: embed, 5 layers, norm, both ranks' 16 best a row, rank 0's selector. Capturable. fn blockBody(dr: *Drafter, e: *forward.Engine, S: usize, rank0: bool, tp: *std.Io.Timestamp) !void { const k = &e.k; const M = S * BLOCK; @@ -813,8 +504,7 @@ fn blockBody(dr: *Drafter, e: *forward.Engine, S: usize, rank0: bool, tp: *std.I try kern.rmsnorm(k, dr.x, D, dr.norm, dr.hn, D, null, R, D); try fine(dr, e, 1); try phase(dr, e, tp, 1); - // every row through the head (a block's first row is the anchor's: the host skips it) and the selector's projection - // the drafts' head: the MTP drafts' packed copy (they only propose), each rank's 16 best a row + // every row through the drafts' packed head (the anchor row skipped by the host), 16 best a rank, selector proj try blocks.mm(e, &e.mbuf, dr.hn, D, if (dr.full_head) e.w.head else e.w.draft_head, e.mbuf.logits, false, R); try fine(dr, e, 12); { @@ -858,60 +548,3 @@ fn conv(k: *kern.K, z: u64, dyn: u64, which: i32, base: u64, out: u64, add: bool a.add(@as(i32, @intFromBool(add))); try k.go(try k.ext(.dflash, "df_conv"), grid, 256, 0, &a); } - -const Cand = struct { v: f32, id: u32 }; - -fn bf(x: u16) f32 { - return @bitCast(@as(u32, x) << 16); -} - -/// Stream ``i``'s 7 drafts: at each row, the 16 likeliest of both ranks' candidates, scored by the head's logit plus -/// the previous draft's predecessor row, the row's projection and the candidate's successor row; greedy, its -/// softmax chance kept. -fn select(dr: *Drafter, i: usize, M: usize, q: Req) void { - const out = q.out.?; - const probs = q.probs.?; - const W = CW; - const T = TOPK; - var prev: usize = q.anchor; - for (1..BLOCK) |t| { - const r = i * BLOCK + t; - var c: [TOPK]Cand = undefined; - var n: usize = 0; - for (0..2) |rk| { - const blk = dr.cand[rk * M * W + r * W ..][0..W]; - for (0..T) |j| { - if (blk[j] < 0) continue; - const cand: Cand = .{ .v = @bitCast(blk[T + j]), .id = @intCast(blk[j]) }; - if (n < TOPK) { - c[n] = cand; - n += 1; - continue; - } - var lo: usize = 0; - for (1..TOPK) |m| if (c[m].v < c[lo].v) { - lo = m; - }; - if (cand.v > c[lo].v) c[lo] = cand; - } - } - const g = dr.gh[r * RANK ..][0..RANK]; - const P = dr.pred[prev * RANK ..][0..RANK]; - var pg: [RANK]f32 = undefined; - for (&pg, P, g) |*o, p, x| o.* = bf(p) * x; - var best: usize = 0; - var sc: [TOPK]f32 = undefined; - for (0..n) |m| { - const Sc = dr.succ[@as(usize, c[m].id) * RANK ..][0..RANK]; - var dot: f32 = 0; - for (pg, Sc) |a, b| dot += a * bf(b); - sc[m] = c[m].v + dot; - if (sc[m] > sc[best]) best = m; - } - var z: f64 = 0; - for (sc[0..n]) |x| z += @exp(@as(f64, x - sc[best]) / dr.temp); - out[t - 1] = c[best].id; - probs[t - 1] = 1.0 / z; - prev = c[best].id; - } -} diff --git a/zig/src/families/glm_cuda/dflash_load.zig b/zig/src/families/glm_cuda/dflash_load.zig new file mode 100644 index 0000000000..cf2972141f --- /dev/null +++ b/zig/src/families/glm_cuda/dflash_load.zig @@ -0,0 +1,315 @@ +//! DFlash2's checkpoint onto the device: the config checked, linears packed, scratch carved from one arena. +const std = @import("std"); +const cuda = @import("cuda"); +const kern = @import("kern.zig"); +const forward = @import("forward.zig"); +const weights_mod = @import("weights.zig"); +const dflash = @import("dflash.zig"); +const D = dflash.D; +const KVW = dflash.KVW; +const QKV = dflash.QKV; +const FF = dflash.FF; +const VOCAB = dflash.VOCAB; +const MAX_M = dflash.MAX_M; +const MAX_CTX = dflash.MAX_CTX; +const CW = dflash.CW; +const CHUNKS = dflash.CHUNKS; +const QH = dflash.QH; +const HD = dflash.HD; +const CH = dflash.CH; +const BLOCK = dflash.BLOCK; +const LAYERS = dflash.LAYERS; +const Drafter = dflash.Drafter; +const RANK = dflash.RANK; +const TOPK = dflash.TOPK; +const TAPS = dflash.TAPS; +const MASK = dflash.MASK; +const B16 = dflash.B16; + +const Ent = struct { off: u64, n: u64 }; + +/// TF_GLM_DFLASH2_BITS: linears packed at load to 4 bits (default; they only propose), 6 or 8; 0 keeps bf16. +fn bitsFromEnv() !u8 { + const v = std.c.getenv("TF_GLM_DFLASH2_BITS") orelse return 4; + const b = std.fmt.parseInt(u8, std.mem.span(v), 10) catch return error.BadDFlash2Bits; + return switch (b) { + 0, 4, 6, 8 => b, + else => error.BadDFlash2Bits, + }; +} + +const Loader = struct { + gpa: std.mem.Allocator, + io: std.Io, + file: std.Io.File, + ents: *std.StringHashMap(Ent), + buf: std.ArrayList(u8) = .empty, + words: std.ArrayList(u32) = .empty, + scales: std.ArrayList(u16) = .empty, + sub: std.ArrayList(u16) = .empty, // a slice of a stored tensor (a rank's share of a split linear) + w: *cuda.DeviceBuffer, + tmp: *cuda.DeviceBuffer, + k: *kern.K, + at: u64 = 0, + bits: u8, + + fn take(l: *Loader, n: u64) u64 { + const p = l.at; + l.at += std.mem.alignForward(u64, n, 256); + return l.w.ptr + p; + } + + fn read(l: *Loader, name: []const u8, want: u64) ![]u8 { + const e = l.ents.get(name) orelse { + std.log.err("dflash2: no {s} in the checkpoint", .{name}); + return error.DrafterTensorMissing; + }; + if (e.n != want) { + std.log.err("dflash2: {s} holds {d} bytes, not {d}", .{ name, e.n, want }); + return error.DrafterTensorShape; + } + try l.buf.resize(l.gpa, e.n); + _ = try l.file.readPositionalAll(l.io, l.buf.items, e.off); + return l.buf.items; + } + + /// A small bf16 tensor as stored. + fn raw(l: *Loader, name: []const u8, n: u64) !u64 { + const p = l.take(n); + try l.w.upload(p - l.w.ptr, try l.read(name, n)); + return p; + } + + /// The tensors ``names`` (``rows`` rows, ``kk`` columns each) stacked into one linear, packed for kern.mm. + fn linear(l: *Loader, names: []const []const u8, rows: []const usize, kk: usize) !B16 { + var n: usize = 0; + for (rows) |r| n += r; + const bits = l.bits; + const wb: usize = n * kk * (if (bits == 0) 16 else bits) / 8; + const sb: usize = if (bits == 0) 0 else n * kk / 64 * 2; + const dst = l.take(wb); + const ds: u64 = if (sb > 0) l.take(sb) else 0; + var r0: usize = 0; + if (bits == 0) { + for (names, rows) |name, r| { + try l.tmp.upload(r0 * kk * 2, try l.read(name, r * kk * 2)); + r0 += r; + } + } else { + try l.words.resize(l.gpa, n * kk * bits / 32); + try l.scales.resize(l.gpa, n * kk / 64); + for (names, rows) |name, r| { + const src = try l.read(name, r * kk * 2); + weights_mod.quantize(bits, @alignCast(std.mem.bytesAsSlice(u16, src)), r, kk, l.words.items[r0 * kk * bits / 32 ..][0 .. r * kk * bits / 32], l.scales.items[r0 * kk / 64 ..][0 .. r * kk / 64]); + r0 += r; + } + try l.tmp.upload(0, std.mem.sliceAsBytes(l.words.items)); + try l.tmp.upload(std.mem.alignForward(u64, wb, 256), std.mem.sliceAsBytes(l.scales.items)); + } + try kern.mmTile(l.k, bits, l.tmp.ptr, l.tmp.ptr + std.mem.alignForward(u64, wb, 256), dst, ds, @intCast(n), @intCast(kk)); + try l.k.stream.synchronize(); + return .{ .w = dst, .n = @intCast(n), .k = @intCast(kk), .s = ds, .bits = bits }; + } + + /// A rank's share of a split linear: rows row0.. and columns col0.. of each of ``names``, stacked and packed. + fn linearPart(l: *Loader, names: []const []const u8, rows: []const usize, kk: usize, row0: usize, rn: usize, col0: usize, cn: usize) !B16 { + const n = names.len * rn; + const bits = l.bits; + const wb: usize = n * cn * (if (bits == 0) 16 else bits) / 8; + const sb: usize = if (bits == 0) 0 else n * cn / 64 * 2; + const dst = l.take(wb); + const ds: u64 = if (sb > 0) l.take(sb) else 0; + try l.sub.resize(l.gpa, rn * cn); + if (bits != 0) { + try l.words.resize(l.gpa, n * cn * bits / 32); + try l.scales.resize(l.gpa, n * cn / 64); + } + var r0: usize = 0; + for (names, rows) |name, r| { + if (row0 + rn > r or col0 + cn > kk) return error.DrafterTensorShape; + const src: []const u16 = @alignCast(std.mem.bytesAsSlice(u16, try l.read(name, r * kk * 2))); + for (0..rn) |i| @memcpy(l.sub.items[i * cn ..][0..cn], src[(row0 + i) * kk + col0 ..][0..cn]); + if (bits == 0) { + try l.tmp.upload(r0 * cn * 2, std.mem.sliceAsBytes(l.sub.items)); + } else { + weights_mod.quantize(bits, l.sub.items, rn, cn, l.words.items[r0 * cn * bits / 32 ..][0 .. rn * cn * bits / 32], l.scales.items[r0 * cn / 64 ..][0 .. rn * cn / 64]); + } + r0 += rn; + } + if (bits != 0) { + try l.tmp.upload(0, std.mem.sliceAsBytes(l.words.items)); + try l.tmp.upload(std.mem.alignForward(u64, wb, 256), std.mem.sliceAsBytes(l.scales.items)); + } + try kern.mmTile(l.k, bits, l.tmp.ptr, l.tmp.ptr + std.mem.alignForward(u64, wb, 256), dst, ds, @intCast(n), @intCast(cn)); + try l.k.stream.synchronize(); + return .{ .w = dst, .n = @intCast(n), .k = @intCast(cn), .s = ds, .bits = bits }; + } +}; + +fn linearBytes(n: usize, kk: usize, bits: u8) u64 { + const wb = std.mem.alignForward(u64, n * kk * (if (bits == 0) 16 else @as(usize, bits)) / 8, 256); + return wb + if (bits == 0) 0 else std.mem.alignForward(u64, n * kk / 64 * 2, 256); +} + +/// The drafter from ``dir`` (config.json and model.safetensors, bf16), its linears packed for the engine's matmul. +pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *forward.Engine, dir: []const u8) !*Drafter { + try checkConfig(gpa, io, dir); + const bits = try bitsFromEnv(); + const d = e.ctx.d; + const path = try std.fs.path.join(gpa, &.{ dir, "model.safetensors" }); + defer gpa.free(path); + const file = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer file.close(io); + var head: [8]u8 = undefined; + _ = try file.readPositionalAll(io, &head, 0); + const hlen = std.mem.readInt(u64, &head, .little); + if (hlen > 1 << 20) return error.BadDrafterHeader; + const text = try gpa.alloc(u8, hlen); + defer gpa.free(text); + _ = try file.readPositionalAll(io, text, 8); + const hdr = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer hdr.deinit(); + var ents = std.StringHashMap(Ent).init(gpa); + defer ents.deinit(); + var it = hdr.value.object.iterator(); + while (it.next()) |x| { + if (std.mem.eql(u8, x.key_ptr.*, "__metadata__")) continue; + if (!std.mem.eql(u8, x.value_ptr.object.get("dtype").?.string, "BF16")) return error.DrafterNotBf16; + const offs = x.value_ptr.object.get("data_offsets").?.array.items; + try ents.put(x.key_ptr.*, .{ .off = 8 + hlen + @as(u64, @intCast(offs[0].integer)), .n = @intCast(offs[1].integer - offs[0].integer) }); + } + // the device's share: linears packed (q|k|v, gate|up, k|v for context), norms and convs; codebooks stay host + var total: u64 = 0; + const split = e.ctx.world == 2 and if (std.c.getenv("TF_GLM_DFLASH2_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; + const ff: usize = if (split) FF / 2 else FF; + const small: u64 = 6 * 4096 * 2 + 2 * 128 * 2 + 2 * 2 * 2 * D * 2; + total += LAYERS * (linearBytes(QKV, D, bits) + linearBytes(2 * KVW, D, bits) + linearBytes(D, D, bits) + linearBytes(2 * ff, D, bits) + + linearBytes(D, ff, bits) + 2 * linearBytes(KVW, D, bits) + small + 8 * 256); + total += linearBytes(D, 5 * D, bits) + linearBytes(RANK, D, bits) + 2 * D * 2 + 1024; + var wbuf = try cuda.DeviceBuffer.alloc(d, total); + errdefer wbuf.free(); + var tmp = try cuda.DeviceBuffer.alloc(d, linearBytes(2 * FF, D, 0) + 4096); // the largest linear as bf16 + defer tmp.free(); + var l: Loader = .{ .gpa = gpa, .io = io, .file = file, .ents = &ents, .w = &wbuf, .tmp = &tmp, .k = &e.k, .bits = bits }; + defer { + l.buf.deinit(gpa); + l.words.deinit(gpa); + l.scales.deinit(gpa); + l.sub.deinit(gpa); + } + const dr = try gpa.create(Drafter); + errdefer gpa.destroy(dr); + dr.* = .{ .gpa = gpa, .d = d, .weights = wbuf, .bits = bits, .split = split, .ff = ff, .pred = &.{}, .succ = &.{}, .scratch = undefined, .h_ids = undefined, .h_bpos = undefined, .h_rings = undefined, .h_starts = undefined, .h_froms = undefined, .h_cpos = undefined, .h_crings = undefined, .cand = undefined, .gh = undefined }; + var nb: [8][96]u8 = undefined; + for (&dr.layers, 0..) |*lw, i| { + const N = struct { + fn n(b: *[96]u8, li: usize, s: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "layers.{d}.{s}", .{ li, s }); + } + }.n; + lw.in_norm = try l.raw(try N(&nb[0], i, "input_layernorm.weight"), D * 2); + lw.post_norm = try l.raw(try N(&nb[0], i, "post_attention_layernorm.weight"), D * 2); + lw.qn = try l.raw(try N(&nb[0], i, "self_attn.q_norm.weight"), 128 * 2); + lw.kn = try l.raw(try N(&nb[0], i, "self_attn.k_norm.weight"), 128 * 2); + lw.a_base = try l.raw(try N(&nb[0], i, "attention_conv.base_kernel"), 2 * 2 * D * 2); + lw.m_base = try l.raw(try N(&nb[0], i, "mlp_conv.base_kernel"), 2 * 2 * D * 2); + const q = try N(&nb[1], i, "self_attn.q_proj.weight"); + const kq = try N(&nb[2], i, "self_attn.k_proj.weight"); + const v = try N(&nb[3], i, "self_attn.v_proj.weight"); + lw.qkv = try l.linear(&.{ q, kq, v }, &.{ D, KVW, KVW }, D); + lw.kv = try l.linear(&.{ kq, v }, &.{ KVW, KVW }, D); + lw.o = try l.linear(&.{try N(&nb[1], i, "self_attn.o_proj.weight")}, &.{D}, D); + const gate = try N(&nb[1], i, "mlp.gate_proj.weight"); + const up = try N(&nb[2], i, "mlp.up_proj.weight"); + const down = try N(&nb[3], i, "mlp.down_proj.weight"); + if (split) { // this rank's half of the MLP width: gate and up rows, down columns + lw.gu = try l.linearPart(&.{ gate, up }, &.{ FF, FF }, D, e.ctx.rank * ff, ff, 0, D); + lw.down = try l.linearPart(&.{down}, &.{D}, FF, 0, D, e.ctx.rank * ff, ff); + } else { + lw.gu = try l.linear(&.{ gate, up }, &.{ FF, FF }, D); + lw.down = try l.linear(&.{down}, &.{D}, FF); + } + lw.a_proj = try l.linear(&.{try N(&nb[1], i, "attention_conv.kernel_projection.weight")}, &.{KVW}, D); + lw.m_proj = try l.linear(&.{try N(&nb[1], i, "mlp_conv.kernel_projection.weight")}, &.{KVW}, D); + } + dr.fc = try l.linear(&.{"fc.weight"}, &.{D}, 5 * D); + dr.hproj = try l.linear(&.{"candidate_selector.hidden_projection.weight"}, &.{RANK}, D); + dr.hnorm = try l.raw("hidden_norm.weight", D * 2); + dr.norm = try l.raw("norm.weight", D * 2); + if (l.at > total) return error.DrafterLayout; + dr.pred = try gpa.alloc(u16, VOCAB * RANK); + @memcpy(std.mem.sliceAsBytes(dr.pred), try l.read("candidate_selector.predecessor_codebook", VOCAB * RANK * 2)); + dr.succ = try gpa.alloc(u16, VOCAB * RANK); + @memcpy(std.mem.sliceAsBytes(dr.succ), try l.read("candidate_selector.successor_codebook", VOCAB * RANK * 2)); + try carveScratch(dr, d); + dr.h_ids = try gpa.alloc(i32, MAX_M); + dr.h_bpos = try gpa.alloc(i32, MAX_M); + dr.h_rings = try gpa.alloc(u64, forward.MAX_STREAMS); + dr.h_starts = try gpa.alloc(i64, forward.MAX_STREAMS); + dr.h_froms = try gpa.alloc(i64, forward.MAX_STREAMS); + dr.h_cpos = try gpa.alloc(i64, MAX_CTX); + dr.h_crings = try gpa.alloc(u64, MAX_CTX); + dr.cand = try gpa.alloc(i32, 2 * MAX_M * CW); + dr.gh = try gpa.alloc(f32, MAX_M * RANK); + dr.prof = std.c.getenv("TF_GLM_DFLASH2_PROF") != null; + if (std.c.getenv("TF_GLM_DFLASH2_ATTN")) |v| dr.attn_chunked = !std.mem.eql(u8, std.mem.span(v), "serial"); + dr.prof_fine = if (std.c.getenv("TF_GLM_DFLASH2_PROF")) |v| std.mem.eql(u8, std.mem.span(v), "2") else false; + dr.dump = std.c.getenv("TF_GLM_DFLASH2_DUMP"); + if (std.c.getenv("TF_GLM_DFLASH2_TAP")) |v| { + dr.tap_input = std.mem.eql(u8, std.mem.span(v), "input"); + dr.tap_s0 = std.mem.eql(u8, std.mem.span(v), "s0"); + } + if (std.c.getenv("TF_GLM_DFLASH2_DUMP_ROUND")) |v| dr.dump_at = std.fmt.parseInt(u32, std.mem.span(v), 10) catch 10; + if (std.c.getenv("TF_GLM_DFLASH2_HEAD")) |v| dr.full_head = std.mem.eql(u8, std.mem.span(v), "full"); + if (std.c.getenv("TF_GLM_DFLASH2_TEMP")) |v| dr.temp = std.fmt.parseFloat(f64, std.mem.span(v)) catch 1.0; + const inv: [MAX_M]f32 = @splat(1.0); + try dr.scratch.upload(dr.invt - dr.scratch.ptr, std.mem.sliceAsBytes(&inv)); + std.log.info("dflash2: drafter loaded from {s}, linears at {d} bits ({d:.2} GiB on the device, {d} MiB of codebooks on the host)", .{ dir, bits, @as(f64, @floatFromInt(l.at)) / (1 << 30), (2 * VOCAB * RANK * 2) >> 20 }); + return dr; +} + +fn carveScratch(dr: *Drafter, d: *const cuda.Driver) !void { + var at: u64 = 0; + const M: u64 = MAX_M; + const C: u64 = MAX_CTX + 16; + const fields = .{ + .{ "x", M * D * 2 }, .{ "y", M * D * 2 }, .{ "y2", M * D * 2 }, .{ "dyn", M * KVW * 2 }, + .{ "qkv", M * QKV * 2 }, .{ "q", M * D * 2 }, .{ "kb", M * KVW * 2 }, .{ "vb", M * KVW * 2 }, + .{ "att", M * D * 2 }, .{ "a", M * D * 2 }, .{ "gu", M * 2 * FF * 2 }, .{ "act", M * FF * 2 }, + .{ "hn", M * D * 2 }, .{ "g", M * RANK * 4 }, .{ "cin", C * 5 * D * 2 }, .{ "c1", C * D * 2 }, + .{ "ctx", C * D * 2 }, .{ "ckv", C * 2 * KVW * 2 }, .{ "taps_step", @as(u64, @intCast(forward.STEP_ROWS)) * 5 * D * 2 }, .{ "taps_prefill", @as(u64, @intCast(forward.PREFILL_ROWS_MAX)) * 5 * D * 2 }, + .{ "ids", M * 4 }, .{ "bpos", M * 4 }, .{ "rings", forward.MAX_STREAMS * 8 }, .{ "starts", forward.MAX_STREAMS * 8 }, + .{ "froms", forward.MAX_STREAMS * 8 }, .{ "cpos", C * 8 }, .{ "crings", C * 8 }, .{ "invt", M * 4 }, + .{ "tk", M * CW * 4 }, .{ "gtk", 2 * M * CW * 4 }, .{ "pv", M * CHUNKS * TOPK * 4 }, .{ "pi", M * CHUNKS * TOPK * 4 }, + .{ "pm", M * QH * CH * 4 }, .{ "pl", M * QH * CH * 4 }, .{ "pacc", M * QH * CH * HD * 4 }, + }; + inline for (fields) |f| { + @field(dr, f[0]) = at; + at += std.mem.alignForward(u64, f[1], 256); + } + dr.scratch = try cuda.DeviceBuffer.alloc(d, at); + inline for (fields) |f| @field(dr, f[0]) += dr.scratch.ptr; + try d.check(d.api.cuMemsetD8_v2(dr.scratch.ptr, 0, at), "zero drafter scratch"); +} + +/// The checkpoint is the one this was written for (its shape and DFlash2 settings). +fn checkConfig(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 text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(text); + const p = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer p.deinit(); + const o = p.value.object; + const df = (o.get("dflash_config") orelse return error.NotDFlash2).object; + const ok = o.get("hidden_size").?.integer == D and o.get("num_hidden_layers").?.integer == LAYERS and + o.get("num_attention_heads").?.integer == 32 and o.get("num_key_value_heads").?.integer == 8 and + o.get("intermediate_size").?.integer == FF and o.get("vocab_size").?.integer == VOCAB and + df.get("block_size").?.integer == BLOCK and df.get("mask_token_id").?.integer == MASK and + df.get("selector_top_k").?.integer == TOPK and df.get("selector_rank").?.integer == RANK and + df.get("conv_kernel_size").?.integer == 2 and df.get("conv_group_size").?.integer == 16; + if (!ok) return error.DrafterConfigUnexpected; + const taps = df.get("target_layer_ids").?.array.items; + if (taps.len != TAPS.len) return error.DrafterConfigUnexpected; + for (taps, TAPS) |t, w| if (t.integer != w) return error.DrafterConfigUnexpected; +} diff --git a/zig/src/families/glm_cuda/dflash_select.zig b/zig/src/families/glm_cuda/dflash_select.zig new file mode 100644 index 0000000000..be194ea321 --- /dev/null +++ b/zig/src/families/glm_cuda/dflash_select.zig @@ -0,0 +1,63 @@ +//! DFlash2's path selector on rank 0: each draft among a row's 16 likeliest by the previous draft's codebook row. +const dflash = @import("dflash.zig"); +const CW = dflash.CW; +const BLOCK = dflash.BLOCK; +const Drafter = dflash.Drafter; +const RANK = dflash.RANK; +const TOPK = dflash.TOPK; +const Req = dflash.Req; + +const Cand = struct { v: f32, id: u32 }; + +fn bf(x: u16) f32 { + return @bitCast(@as(u32, x) << 16); +} + +/// Stream ``i``'s 7 drafts: each row's best of both ranks' 16 by logit plus codebook scores; greedy, chance kept. +pub fn select(dr: *Drafter, i: usize, M: usize, q: Req) void { + const out = q.out.?; + const probs = q.probs.?; + const W = CW; + const T = TOPK; + var prev: usize = q.anchor; + for (1..BLOCK) |t| { + const r = i * BLOCK + t; + var c: [TOPK]Cand = undefined; + var n: usize = 0; + for (0..2) |rk| { + const blk = dr.cand[rk * M * W + r * W ..][0..W]; + for (0..T) |j| { + if (blk[j] < 0) continue; + const cand: Cand = .{ .v = @bitCast(blk[T + j]), .id = @intCast(blk[j]) }; + if (n < TOPK) { + c[n] = cand; + n += 1; + continue; + } + var lo: usize = 0; + for (1..TOPK) |m| if (c[m].v < c[lo].v) { + lo = m; + }; + if (cand.v > c[lo].v) c[lo] = cand; + } + } + const g = dr.gh[r * RANK ..][0..RANK]; + const P = dr.pred[prev * RANK ..][0..RANK]; + var pg: [RANK]f32 = undefined; + for (&pg, P, g) |*o, p, x| o.* = bf(p) * x; + var best: usize = 0; + var sc: [TOPK]f32 = undefined; + for (0..n) |m| { + const Sc = dr.succ[@as(usize, c[m].id) * RANK ..][0..RANK]; + var dot: f32 = 0; + for (pg, Sc) |a, b| dot += a * bf(b); + sc[m] = c[m].v + dot; + if (sc[m] > sc[best]) best = m; + } + var z: f64 = 0; + for (sc[0..n]) |x| z += @exp(@as(f64, x - sc[best]) / dr.temp); + out[t - 1] = c[best].id; + probs[t - 1] = 1.0 / z; + prev = c[best].id; + } +} diff --git a/zig/src/families/glm_cuda/draw.zig b/zig/src/families/glm_cuda/draw.zig index 4986ec0a62..690469b1e2 100644 --- a/zig/src/families/glm_cuda/draw.zig +++ b/zig/src/families/glm_cuda/draw.zig @@ -1,8 +1,4 @@ -//! A sampled row's token from both ranks' gathered candidates (fn_cand_topk): each rank's 64 largest logits, global -//! ids, its log-sum-exp and its log-sum-exp at the row's temperature. The merged 128 hold the vocabulary's top 64 -//! exactly, so a rule within them draws as lanes.sampling.choose over the whole vocabulary would: top_k 1..64 by -//! choose itself (only the top k take part), a wider rule when its nucleus (top_p) or floor (min_p) ends within the -//! top 64, keyed the same way (seed, position, global id). A rule that reaches past them asks for the full row. +//! A sampled token from both ranks' top 64 (exact top 64 of the vocabulary); a rule reaching past them asks the row. const std = @import("std"); const lanes = @import("lanes"); @@ -87,8 +83,7 @@ pub fn draw(gpa: std.mem.Allocator, blocks: [2][]const i32, s: smp.Sampling, pos return .{ .token = @intCast(cands[best].id) }; } -/// ``drawFull`` over a whole gathered bf16 row (raw bits by global id): the same order by a counting sort over the -/// 65,536 bf16 values (descending, ids ascending within a value; -0 with +0), then the same rule. +/// ``drawFull`` over a whole gathered bf16 row: counting sort over the 65,536 bf16 values, then the same rule. pub fn drawFullBf16(gpa: std.mem.Allocator, row: []const u16, s: smp.Sampling, position: u64) !u32 { const counts = try gpa.alloc(u32, 1 << 16); defer gpa.free(counts); diff --git a/zig/src/families/glm_cuda/experts_load.zig b/zig/src/families/glm_cuda/experts_load.zig new file mode 100644 index 0000000000..495918b784 --- /dev/null +++ b/zig/src/families/glm_cuda/experts_load.zig @@ -0,0 +1,331 @@ +//! GLM-5.3-Flash's routed experts for one rank, read from the checkpoint and packed (weights.zig's Builder). + +const std = @import("std"); +const kern = @import("kern.zig"); +const ckpt = @import("ckpt.zig"); +const Builder = @import("weights.zig").Builder; +const WORLD = ckpt.WORLD; +const fullName = ckpt.fullName; +const numelOf = ckpt.numelOf; +const f16ToBf16 = ckpt.f16ToBf16; +const bf16Value = ckpt.bf16Value; + +// -- routed experts: per-expert MLX triples, rank rows (gate, up) or columns (down), packed as experts_pack.cu ---- + +const Triple = struct { + e: usize, + n: usize, + k8: usize, + kg: usize, + words: []u32, + scales: []u16, + biases: []u16, + + fn free(t: Triple, gpa: std.mem.Allocator) void { + gpa.free(t.words); + gpa.free(t.scales); + gpa.free(t.biases); + } +}; + +/// Every expert's ``proj`` of layer ``li``: rank rows (gate, up) or columns (down), fp16 scales and biases to bf16. +fn readExperts(b: *Builder, li: usize, proj: []const u8, count: usize) !Triple { + var nb: [256]u8 = undefined; + const w0 = try b.ck.get(try fullName(&nb, "layers.{d}.mlp.experts.0.{s}.weight", .{ li, proj })); + const n_full: usize = @intCast(w0.shape[0]); + const k8_full: usize = @intCast(w0.shape[1]); + var nb2: [256]u8 = undefined; + const s0 = try b.ck.get(try fullName(&nb2, "layers.{d}.mlp.experts.0.{s}.scales", .{ li, proj })); + const kg_full: usize = @intCast(s0.shape[1]); + const down = std.mem.eql(u8, proj, "down_proj"); + const nr = if (down) n_full else n_full / WORLD; + const r0 = if (down) 0 else b.rank * nr; + const ok8 = if (down) k8_full / WORLD else k8_full; + const okg = if (down) kg_full / WORLD else kg_full; + const c8 = if (down) b.rank * ok8 else 0; + const cg = if (down) b.rank * okg else 0; + var t: Triple = .{ .e = count, .n = nr, .k8 = ok8, .kg = okg, .words = try b.gpa.alloc(u32, count * nr * ok8), .scales = try b.gpa.alloc(u16, count * nr * okg), .biases = try b.gpa.alloc(u16, count * nr * okg) }; + errdefer t.free(b.gpa); + const rowbuf = try b.gpa.alloc(u8, nr * k8_full * 4); + defer b.gpa.free(rowbuf); + const sbuf = try b.gpa.alloc(u8, nr * kg_full * 2); + defer b.gpa.free(sbuf); + for (0..count) |x| { + var n1: [256]u8 = undefined; + const we = try b.ck.get(try fullName(&n1, "layers.{d}.mlp.experts.{d}.{s}.weight", .{ li, x, proj })); + try b.ck.readAt(we, r0 * k8_full * 4, rowbuf); + const w32 = std.mem.bytesAsSlice(u32, rowbuf); + for (0..nr) |r| @memcpy(t.words[(x * nr + r) * ok8 ..][0..ok8], w32[r * k8_full + c8 ..][0..ok8]); + for ([_][]const u8{ "scales", "biases" }, 0..) |what, which| { + var n2: [256]u8 = undefined; + const ent = try b.ck.get(try fullName(&n2, "layers.{d}.mlp.experts.{d}.{s}.{s}", .{ li, x, proj, what })); + try b.ck.readAt(ent, r0 * kg_full * 2, sbuf); + const s16 = std.mem.bytesAsSlice(u16, sbuf); + const dst = if (which == 0) t.scales else t.biases; + for (0..nr) |r| for (0..okg) |g| { + dst[(x * nr + r) * okg + g] = f16ToBf16(s16[r * kg_full + cg + g]); + }; + } + } + return t; +} + +/// experts_pack.cu's index map: per block 32 * 4 * H nibble-shuffled words, then 32 (16 symmetric) scale words. +fn checkSymmetric(t: Triple) !void { + for (t.scales, t.biases) |sc, bi| { + if (bf16Value(bi) != -8 * bf16Value(sc)) return error.ExpertsNotSymmetric; + } +} + +fn packExperts(t: Triple, gs: usize, out: []u32, mcount: usize, m: usize, sym: bool, three: bool) void { + const h = gs / 32; + const Work = struct { + fn run(tt: Triple, hh: usize, o: []u32, mc: usize, mm: usize, e0: usize, e1: usize, sy: bool, th: bool) void { + const wp = 32 * 4 * hh; + const wsec = if (th) wp * 3 / 4 else wp; + const blk = wsec + @as(usize, if (sy) 16 else 32); + var tmp: [512]u32 = undefined; // a 3-bit block's 4-bit words before their planes + const nbb = tt.n / 32; + for (e0..e1) |e| for (0..nbb) |bb| for (0..tt.kg) |g| { + const dst = o[(((e * nbb + bb) * tt.kg + g) * mc + mm) * blk ..][0..blk]; + const wdst = if (th) tmp[0..wp] else dst[0..wp]; + for (0..wp) |i| { + const row = i / (4 * hh); + const kk = i % (4 * hh); + const tq = row / 8; + const r = row % 8; + const q = kk / hh; + const j = kk % hh; + const w = tt.words[(e * tt.n + bb * 32 + row) * tt.k8 + g * 4 * hh + kk]; + const f = ((tq * hh + j) * 8 + r) * 4 + q; + wdst[(f / 128) * 128 + (f % 32) * 4 + (f / 32) % 4] = shuffle(w); + } + if (th) for (0..hh) |c| for (0..32) |lane| { // the lane's uint4 c as three planes (experts_sym.cuh) + var lo: [4]u32 = undefined; + var hi: [4]u32 = undefined; + for (0..4) |k| { + const word = tmp[(c * 32 + lane) * 4 + k]; + lo[k] = 0; + hi[k] = 0; + for (0..8) |i| { + const code = ((word >> @intCast(4 * i)) & 15) -% 4; // 3-bit codes are stored + 4 + lo[k] |= (code & 3) << @intCast(2 * i); + hi[k] |= ((code >> 2) & 1) << @intCast(i); + } + } + dst[c * 96 + lane] = lo[0] | (lo[1] << 16); + dst[c * 96 + 32 + lane] = lo[2] | (lo[3] << 16); + dst[c * 96 + 64 + lane] = hi[0] | (hi[1] << 8) | (hi[2] << 16) | (hi[3] << 24); + }; + if (sy) { // quad p: rows j * 8 + 2p and + 1's scales, j = 0..3 (experts_sym.cuh) + for (0..16) |s| { + const p = s / 4; + const j = s % 4; + const at = (e * tt.n + bb * 32 + j * 8 + p * 2) * tt.kg + g; + dst[wsec + s] = @as(u32, tt.scales[at]) | (@as(u32, tt.scales[at + tt.kg]) << 16); + } + } else for (0..32) |s| { + const p = s / 8; + const tt4 = s % 4; + const src = if ((s / 4) % 2 == 1) tt.biases else tt.scales; + const at = (e * tt.n + bb * 32 + tt4 * 8 + p * 2) * tt.kg + g; + dst[wp + s] = @as(u32, src[at]) | (@as(u32, src[at + tt.kg]) << 16); + } + }; + } + }; + const threads = 16; + var pool: [threads]?std.Thread = @splat(null); + const per = (t.e + threads - 1) / threads; + for (0..threads) |i| { + const e0 = i * per; + const e1 = @min(t.e, e0 + per); + if (e0 >= e1) break; + pool[i] = std.Thread.spawn(.{}, Work.run, .{ t, h, out, mcount, m, e0, e1, sym, three }) catch blk: { + Work.run(t, h, out, mcount, m, e0, e1, sym, three); + break :blk null; + }; + } + for (pool) |th| if (th) |x| x.join(); +} + +fn shuffle(w: u32) u32 { + var even = w & 0x0F0F0F0F; + var odd = (w >> 4) & 0x0F0F0F0F; + even = (even | (even >> 4)) & 0x00FF00FF; + odd = (odd | (odd >> 4)) & 0x00FF00FF; + even = (even | (even >> 8)) & 0x0000FFFF; + odd = (odd | (odd >> 8)) & 0x0000FFFF; + return even | (odd << 16); +} + +/// grouped.make([gate, up], down, 64): up [E, N/32, K/64, 2, block], down [E, D/32, NI/64, 1, block] int32. +pub fn experts(b: *Builder, out: []const u8, li: usize) !void { + const E: usize = @intCast(kern.EXPERTS); + const gs: usize = 64; + if (b.experts3 and !b.experts_sym) return error.Experts3NeedsSymmetric; + const wsec: usize = 32 * 4 * (gs / 32) * @as(usize, if (b.experts3) 3 else 4) / 4; + const blk: usize = wsec + @as(usize, if (b.experts_sym) 16 else 32); + const width: usize = @intCast(@divExact(2048, WORLD)); // moe_intermediate_size / ranks + const dims: usize = @intCast(kern.D); + const up_shape = [_]i64{ @intCast(E), @intCast(width / 32), @intCast(dims / gs), 2, @intCast(blk) }; + const down_shape = [_]i64{ @intCast(E), @intCast(dims / 32), @intCast(width / gs), 1, @intCast(blk) }; + var nb: [256]u8 = undefined; + var nb2: [256]u8 = undefined; + const up_name = try std.fmt.bufPrint(&nb, "{s}.up", .{out}); + const down_name = try std.fmt.bufPrint(&nb2, "{s}.down", .{out}); + if (b.dry) { + try b.put(up_name, .i32, &up_shape, null); + try b.put(down_name, .i32, &down_shape, null); + return; + } + { + const up = try b.gpa.alloc(u32, numelOf(&up_shape)); + defer b.gpa.free(up); + for ([_][]const u8{ "gate_proj", "up_proj" }, 0..) |proj, m| { + const t = try readExperts(b, li, proj, E); + defer t.free(b.gpa); + if (t.n != width or t.k8 * 8 != dims) return error.ExpertShape; + if (b.experts_sym) try checkSymmetric(t); + if (b.sim_proj != 1) { + if (b.experts3) simulateBits(t, 3) else if (b.expert_sim_bits != 0) simulateBits(t, b.expert_sim_bits); + } + packExperts(t, gs, up, 2, m, b.experts_sym, b.experts3); + } + try b.put(up_name, .i32, &up_shape, std.mem.sliceAsBytes(up)); + } + const dn = try b.gpa.alloc(u32, numelOf(&down_shape)); + defer b.gpa.free(dn); + const t = try readExperts(b, li, "down_proj", E); + defer t.free(b.gpa); + if (t.n != dims or t.k8 * 8 != width) return error.ExpertShape; + if (b.experts_sym) try checkSymmetric(t); + if (b.sim_proj != 2) { + if (b.experts3) simulateBits(t, 3) else if (b.expert_sim_bits != 0) simulateBits(t, b.expert_sim_bits); + } + packExperts(t, gs, dn, 1, 0, b.experts_sym, b.experts3); + try b.put(down_name, .i32, &down_shape, std.mem.sliceAsBytes(dn)); +} + +/// Layer ``li``'s EXL3 experts, rank tiles: trellis words [E, K/16, N/16, 32], fp16 suh [E, K] and svh [E, N]. +pub fn expertsExl3(b: *Builder, out: []const u8, li: usize) !void { + const E: usize = @intCast(kern.EXPERTS); + const D: usize = @intCast(kern.D); + const width: usize = @intCast(@divExact(2048, WORLD)); + const tile: usize = 128; // a 16x16 tile's bytes at 4 bits + var nb: [256]u8 = undefined; + const Mat = struct { proj: []const u8, kt: usize, nt: usize, by_cols: bool }; + const mats = [_]Mat{ .{ .proj = "gate_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, .{ .proj = "up_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, .{ .proj = "down_proj", .kt = width / 16, .nt = D / 16, .by_cols = false } }; + const short = [_][]const u8{ "g", "u", "d" }; + for (mats, short) |m, sh| { + const tshape = [_]i64{ @intCast(E), @intCast(m.kt), @intCast(m.nt), 32 }; + const k_all: usize = if (m.by_cols) D else 2048; + const n_all: usize = if (m.by_cols) 2048 else D; + const sus = [_]i64{ @intCast(E), @intCast(m.kt * 16) }; + const svs = [_]i64{ @intCast(E), @intCast(m.nt * 16) }; + const tn = try std.fmt.bufPrint(&nb, "{s}.{s}t", .{ out, sh }); + if (b.dry) { + try b.put(tn, .i32, &tshape, null); + try b.putf("{s}.suh_{s}", .{ out, sh }, .f16, &sus, null); + try b.putf("{s}.svh_{s}", .{ out, sh }, .f16, &svs, null); + continue; + } + const per = m.kt * m.nt * tile; + const words = try b.gpa.alignedAlloc(u8, .@"16", E * per); + defer b.gpa.free(words); + const suh = try b.gpa.alloc(u16, E * m.kt * 16); + defer b.gpa.free(suh); + const svh = try b.gpa.alloc(u16, E * m.nt * 16); + defer b.gpa.free(svh); + var f: [256]u8 = undefined; + for (0..E) |x| { + const tr = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.trellis", .{ li, x, m.proj })); + if (!std.mem.eql(u8, tr.dtype, "I16") or tr.shape.len != 3 or tr.shape[2] != 64 or tr.shape[0] != k_all / 16 or tr.shape[1] != n_all / 16) return error.NotExl3Mcg4; + if (!b.ck.has(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.mcg", .{ li, x, m.proj }))) return error.NotExl3Mcg4; + const dst = words[x * per ..][0..per]; + if (m.by_cols) { // every tile row, this rank's tile columns + const all = try b.ck.readAll(b.gpa, try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.trellis", .{ li, x, m.proj })); + defer b.gpa.free(all); + const row = (n_all / 16) * tile; + for (0..m.kt) |r| @memcpy(dst[r * m.nt * tile ..][0 .. m.nt * tile], all[r * row + b.rank * m.nt * tile ..][0 .. m.nt * tile]); + } else try b.ck.readAt(tr, b.rank * per, dst); // this rank's tile rows: one contiguous part + const se = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.suh", .{ li, x, m.proj })); + const ve = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.svh", .{ li, x, m.proj })); + if (!std.mem.eql(u8, se.dtype, "F16") or !std.mem.eql(u8, ve.dtype, "F16")) return error.NotExl3Mcg4; + const su = std.mem.sliceAsBytes(suh[x * m.kt * 16 ..][0 .. m.kt * 16]); + const sv = std.mem.sliceAsBytes(svh[x * m.nt * 16 ..][0 .. m.nt * 16]); + try b.ck.readAt(se, if (m.by_cols) 0 else b.rank * su.len, su); + try b.ck.readAt(ve, if (m.by_cols) b.rank * sv.len else 0, sv); + if (m.by_cols) { // gate and up: one input rotation shared by every expert of the layer? (exl3_shx) + if (std.mem.eql(u8, sh, "g") and x == 0) @memcpy(b.suh_first[0..su.len], su) else if (!std.mem.eql(u8, su, b.suh_first[0..su.len])) b.exl3_shx = false; + } + } + try b.put(tn, .i32, &tshape, words); + try b.putf("{s}.suh_{s}", .{ out, sh }, .f16, &sus, std.mem.sliceAsBytes(suh)); + try b.putf("{s}.svh_{s}", .{ out, sh }, .f16, &svs, std.mem.sliceAsBytes(svh)); + } +} + +/// Each group of 64 re-quantized to ``bits`` symmetric levels (least-squares clip search), stored as 4-bit codes + 8. +fn simulateBits(t: Triple, bits: u8) void { + const Work = struct { + fn run(tt: Triple, nb: u8, e0: usize, e1: usize) void { + const lo: f32 = -@as(f32, @floatFromInt(@as(u32, 1) << @intCast(nb - 1))); + const hi: f32 = -lo - 1; + const wpg = 8; // 4-bit words a group of 64 + for (e0..e1) |e| for (0..tt.n) |r| for (0..tt.kg) |g| { + const at = (e * tt.n + r) * tt.kg + g; + const words = tt.words[(e * tt.n + r) * tt.k8 + g * wpg ..][0..wpg]; + const s0 = bf16Value(tt.scales[at]); + var v: [64]f32 = undefined; + var amax: f32 = 0; + for (0..64) |i| { + const q: i32 = @intCast((words[i / 8] >> @intCast(4 * (i % 8))) & 15); + v[i] = @as(f32, @floatFromInt(q - 8)) * s0; + amax = @max(amax, @abs(v[i])); + } + if (amax == 0) continue; + var best_s: f32 = amax / -lo; + var best_e: f32 = std.math.inf(f32); + for ([_]f32{ 1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6 }) |c| { + const sc: f32 = bf16Value(toBf16(c * amax / -lo)); + if (sc == 0) continue; + var err: f32 = 0; + for (v) |x| { + const q = @max(lo, @min(hi, @round(x / sc))); + err += (q * sc - x) * (q * sc - x); + } + if (err < best_e) { + best_e = err; + best_s = sc; + } + } + for (words) |*w| w.* = 0; + for (v, 0..) |x, i| { + const q: i32 = @intFromFloat(@max(lo, @min(hi, @round(x / best_s)))); + words[i / 8] |= @as(u32, @intCast(q + 8)) << @intCast(4 * (i % 8)); + } + tt.scales[at] = toBf16(best_s); + tt.biases[at] = toBf16(-8 * best_s); + }; + } + }; + const threads = 16; + var pool: [threads]?std.Thread = @splat(null); + const per = (t.e + threads - 1) / threads; + for (0..threads) |i| { + const e0 = i * per; + const e1 = @min(t.e, e0 + per); + if (e0 >= e1) break; + pool[i] = std.Thread.spawn(.{}, Work.run, .{ t, bits, e0, e1 }) catch blk: { + Work.run(t, bits, e0, e1); + break :blk null; + }; + } + for (pool) |th| if (th) |x| x.join(); +} + +fn toBf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); // round to nearest even +} diff --git a/zig/src/families/glm_cuda/forward.zig b/zig/src/families/glm_cuda/forward.zig index 06836e439d..bd176682f4 100644 --- a/zig/src/families/glm_cuda/forward.zig +++ b/zig/src/families/glm_cuda/forward.zig @@ -1,12 +1,4 @@ -//! GLM-5.3-Flash (glm5_next) on CUDA, one rank of two: the forward on this family's own kernels, for any mix of -//! streams. A step is a list of segments (a stream and its rows, in order): the row-wise work (projections, norms, -//! hyper-connections, router, experts, head) runs over all the step's rows at once, and the attention, the indexer -//! and the KDA recurrences read the step's tables (each row's position and its stream's caches). Every kernel keeps a -//! row's bits its own, so a stream's tokens are the same alone and beside fifteen others. -//! -//! Both ranks make the same calls in the same order. Tokens are drawn from both ranks' vocabulary shards (an -//! all-gather of each rank's best candidates), so both ranks reach the same tokens with no host decision that could -//! differ. +//! GLM-5.3-Flash on CUDA, one rank of two: each row's bits its own at any mix of streams; both ranks call alike. const std = @import("std"); const cuda = @import("cuda"); @@ -26,8 +18,7 @@ const LL = kern.LL; pub const MAX_STREAMS: usize = 16; pub const PREFILL_ROWS_MAX: i64 = 4096; // a prompt chunk at most (buffers sized once take the served value) -/// Rows a prompt chunk takes: TF_GLM_PREFILL_ROWS (2048; 4096 halves what the EXL3 experts decode per row and takes -/// about 2 GiB more of the pool). Read once, before the buffers. +/// Rows a prompt chunk takes: TF_GLM_PREFILL_ROWS (2048; 4096 halves per-row expert decode, ~2 GiB more). Read once. pub fn prefillRows() i64 { const v = std.c.getenv("TF_GLM_PREFILL_ROWS") orelse return 2048; const n = std.fmt.parseInt(i64, std.mem.span(v), 10) catch return 2048; @@ -222,8 +213,7 @@ pub const Seg = struct { s: *Seq, rows: i64 }; pub const Kind = enum(u8) { verify = 0, mtp = 1 }; -/// A decode step's shape: what the launches depend on besides the tables (rows, streams, the dense attention's -/// chunks, whether any row is sparse, the sparse scores' pool bucket). +/// A decode step's shape beyond its tables: rows, streams, dense attention chunks, any sparse row, the pool bucket. const GKey = struct { kind: Kind, rows: i64, nseg: usize, nch: i64, np: i64, sparse: bool }; const Cached = struct { ex: cuda.graph.Exec, used: u64 }; @@ -255,9 +245,7 @@ pub const Engine = struct { full_host: []u16 = &.{}, /// Development: every launch synchronizes (and names the one that faulted). sync_each: bool = false, - /// Decode steps as CUDA graphs, one a shape (``GKey``): the same launches, replayed. The shapes grow with the - /// streams' lengths, so at most ``graph_cap`` stay (TF_GLM_GRAPHS), the least recently run dropped first: each - /// holds device memory the stream pool does not count. + /// Decode steps as CUDA graphs, one a ``GKey`` shape, at most ``graph_cap`` (TF_GLM_GRAPHS), LRU dropped first. use_graphs: bool = true, graphs: std.AutoHashMap(GKey, Cached) = undefined, graph_cap: usize = 96, @@ -266,15 +254,13 @@ pub const Engine = struct { evictions: usize = 0, updates: usize = 0, // captures that took an older graph's instantiation (oldestGraph) capture_ms: f64 = 0, - /// DFlash2 (opt-in): its taps after the target layers it reads, each stream's ring bytes, and the verify steps - /// run (a stream's last verify rows are good for its draft while no later verify ran). + /// DFlash2 (opt-in): its taps, each stream's ring bytes, and the verify steps run (a draft needs the latest). drafter: ?*@import("dflash.zig").Drafter = null, ring_bytes: u64 = 0, verifies: u64 = 0, /// The step being run (its segments and shape), for the blocks. step: Step = .{}, - /// Development (TF_GLM_SKIP=gather,experts,dense,attn,hc): a class of work left out, its time measured by the - /// steps' difference (the replies are wrong while set). + /// Development (TF_GLM_SKIP=gather,experts,dense,attn,hc): work left out to time it; replies wrong while set. skip: Skip = .{}, }; @@ -352,14 +338,12 @@ pub fn freeSeq(e: *Engine, s: *Seq) void { sq.free(e.gpa, s); } -/// The next prefill's images (null: none); its chunks put each run's features in place of the embeddings (the main -/// layers' and the MTP head's, a row later), and the prefill drops them. +/// The next prefill's images (null: none): their features replace the embeddings (MTP a row later), then dropped. pub fn attach(s: *Seq, m: ?Media) void { s.media = m; } -/// The step's tables for ``segs`` into ``b`` (uploaded and waited for: the host buffer is reused), its shape into -/// e.step. ``kind`` .mtp: the MTP layer's cache and positions (mtp_len), else every DSA layer's and pos. +/// The step's tables for ``segs`` into ``b`` (uploaded, waited for) and shape; ``kind`` .mtp: MTP cache and mtp_len. pub fn prepare(e: *Engine, b: *Buffers, segs: []const Seg, kind: Kind) !void { var R: i64 = 0; for (segs) |g| R += g.rows; @@ -415,8 +399,7 @@ fn keyOf(e: *const Engine) GKey { return .{ .kind = st.kind, .rows = st.rows, .nseg = st.segs.len, .nch = st.nch, .np = st.np, .sparse = st.sparse }; } -/// The step's GPU work (capturable): a verify step's forward and every row's candidates, or the MTP head over its -/// staged rows and each stream's last row's candidates. +/// The step's GPU work (capturable): a verify forward and its candidates, or the MTP head and last rows' candidates. fn body(e: *Engine) !void { switch (e.step.kind) { .verify => { @@ -453,8 +436,7 @@ pub fn run(e: *Engine) !void { defer g.deinit(); e.captures += 1; if (e.graphs.count() >= e.graph_cap) { - // full: the least recently run graph of the same layout takes the new shape's parameters in place (no - // instantiation, no memory); a layout the driver will not update to is dropped for a new instantiation + // full: the LRU graph of the same layout takes the new parameters in place; one that won't update is dropped if (oldestGraph(e, key, true)) |k| { var c = e.graphs.fetchRemove(k).?.value; const took = if (c.ex.update(g)) |r| r == .success else |_| false; diff --git a/zig/src/families/glm_cuda/guard.zig b/zig/src/families/glm_cuda/guard.zig index 495d5e8808..3ae8b6fea9 100644 --- a/zig/src/families/glm_cuda/guard.zig +++ b/zig/src/families/glm_cuda/guard.zig @@ -1,9 +1,4 @@ -//! A rank's watchdog. The GB10's memory is the host's: when a rank's caches, graphs and buffers take the last of it -//! the kernel stalls the whole machine (no OOM kill comes), and the peer then waits in a collective forever. So a -//! thread reads MemAvailable every 100 ms and the rank exits at the floor (TF_MEM_FLOOR_GIB, 2 GiB), which frees -//! everything it holds; below the floor and 3 GiB more it is ``tight`` and rank 0 admits no new request. It also exits -//! when its peer is gone (the control link hung up or timed out, or NCCL reports an error): a rank alone can only -//! hang, and an exited one can be restarted. +//! A rank's watchdog: exits at the MemAvailable floor (TF_MEM_FLOOR_GIB, 2 GiB) or when its peer is gone; ``tight``. const std = @import("std"); const posix = std.posix; diff --git a/zig/src/families/glm_cuda/kern.zig b/zig/src/families/glm_cuda/kern.zig index 44b0ab0d85..2d0558eac4 100644 --- a/zig/src/families/glm_cuda/kern.zig +++ b/zig/src/families/glm_cuda/kern.zig @@ -1,17 +1,13 @@ -//! GLM-5.3-Flash on CUDA: every launch the forward makes, on this family's own kernels (glm/mm.cu, rows.cu, -//! latent.cu, index.cu, kda.cu, ops.cu, vision.cu) and the shared grouped experts. Each is written so that a row's -//! bits depend on its own inputs alone: a window's rows, sixteen streams' windows and a prompt chunk's rows go through -//! the same arithmetic. The attention and indexer launches read a step's per-row tables (``Rows``): each row's -//! position and its stream's caches, so one launch serves any mix of streams. +//! GLM-5.3-Flash on CUDA: every launch the forward makes; a row's bits depend on its own inputs alone (``Rows``). const std = @import("std"); const cuda = @import("cuda"); -fn cdiv(a: i64, b: i64) i64 { +pub fn cdiv(a: i64, b: i64) i64 { return @divFloor(a + b - 1, b); } -fn nextPow2(x: u64) u64 { +pub fn nextPow2(x: u64) u64 { return std.math.ceilPowerOfTwo(u64, @max(x, 1)) catch unreachable; } @@ -35,8 +31,7 @@ pub const K = struct { experts_sym: bool = false, /// ... with 3-bit codes in planes (glm/experts_sym* built with GLM_SYM3). experts3: bool = false, - /// The routed experts' current product is a prompt's: the prefill kernels (resumed == fresh: a prompt row always - /// takes them, however few rows its chunk has). + /// The routed experts' current product is a prompt's: the prefill kernels for every prompt row (resumed == fresh). prompt: bool = false, fn after(self: *K, what: []const u8) !void { @@ -71,7 +66,7 @@ pub const K = struct { return k; } - /// A kernel of module ``m`` whose (mangled) name contains ``frag``: the one exact match, else the one containing it. + /// A kernel of module ``m`` whose mangled name contains ``frag``: the exact match, else the one containing it. pub fn ext(self: *K, m: Ext, frag: []const u8) !cuda.Function { const key = std.hash.Wyhash.hash(@backingInt(m), frag); if (self.cache.get(key)) |f| return f; @@ -101,8 +96,7 @@ pub const K = struct { return self.goWith(f, grid, block, shared, args, false); } - /// A launch the kernel before may overlap (programmatic dependent launch): the kernel waits at - /// griddepcontrol.wait before it reads what that one writes. TF_GLM_PDL=0: plain launches. + /// A launch overlapping the one before (PDL): it waits at griddepcontrol.wait; TF_GLM_PDL=0 launches plainly. pub fn goDependent(self: *K, f: cuda.Function, grid: cuda.Dim3, block: u32, shared: u32, args: *cuda.Args) !void { return self.goWith(f, grid, block, shared, args, pdl_on); } @@ -173,9 +167,7 @@ pub const ATT_ROWS: i64 = 512; // rows one attention call takes (its partials' s pub const SELECT_ROWS: i64 = 512; // rows scored at once (the scores' scratch) pub const SCALE: f32 = 0.0625; // qk_dim ** -0.5 -/// A step's per-row tables (device arrays): each row's position (i32) and, for each latent cache ``ci``, the row's -/// stream's latent rows, index key and gate rings (offset so that position p is row p) and pooled keys (u64 each, the -/// ``ci``th array at ci * cap). ``at`` moves to a row block. +/// A step's per-row tables: positions (i32) and per latent cache ``ci`` the rows' cache pointers; ``at``: a row block. pub const Rows = struct { pos: u64, lc: u64, @@ -210,7 +202,7 @@ pub const Rows = struct { } }; -// -- rows.cu ------------------------------------------------------------------------------------------------------------- +// -- rows.cu --------------------------------------------------------------------------------------------------- pub fn embed(k: *K, ids: u64, table: u64, out: u64, R: i64, copies: i64) !void { var a: cuda.Args = .{}; a.add(ids); @@ -329,7 +321,7 @@ pub fn combine(k: *K, y: u64, y_f32: bool, wts: u64, out: u64, R: i64) !void { try k.go(try k.ext(.rows, "glm_combine"), .{ .x = @intCast(R), .y = @intCast(@divExact(D, 1024)) }, 256, 0, &a); } -/// conv [n, 3, C] (in place) takes rows keep .. keep + 2 of [conv; proj rows] (proj [n, R, W >= C]); strides in elements. +/// conv [n, 3, C] in place takes rows keep .. keep + 2 of [conv; proj rows] (proj [n, R, W >= C]), element strides. pub fn convShift(k: *K, conv: u64, proj: u64, keep: i64, conv_layer: i64, proj_layer: i64, proj_row: i64, n: i64, c: i64) !void { var a: cuda.Args = .{}; a.add(conv); @@ -342,9 +334,9 @@ pub fn convShift(k: *K, conv: u64, proj: u64, keep: i64, conv_layer: i64, proj_l try k.go(try k.ext(.rows, "glm_conv_shift"), .{ .x = @intCast(n), .y = @intCast(cdiv(c, 1024)) }, 256, 0, &a); } -// -- mm.cu: the dense linears ------------------------------------------------------------------------------------------ -/// A dense linear [n, k] as glm/mm.cu reads it, tiled at load: ``bits`` 0 (bf16 values) or 4, 6, 8 (symmetric codes -/// with fp16 scales ``s`` per 64 inputs). +// -- mm.cu: the dense linears ---------------------------------------------------------------------------------- + +/// A dense linear [n, k] tiled at load as glm/mm.cu reads it: ``bits`` 0 (bf16) or 4, 6, 8 (fp16 scales ``s`` / 64). pub const B16 = struct { w: u64, n: i64, k: i64, s: u64 = 0, bits: u8 = 0 }; pub const SPLIT_TARGET: i64 = 192; @@ -374,8 +366,7 @@ fn mmSmem(bits: u8, bm: i64) u32 { return @intCast(3 * (bm * 128 + 8 * 32 * wpc * 4 + 128)); } -/// x [m, k] (row stride x_stride) by w -> out [m, n] contiguous, bf16 or unrounded fp32; ``part`` holds sk * m * n fp32 -/// when the shape splits K. Rows go 128 at a time when there are more (the same arithmetic for every row). +/// x [m, k] (stride x_stride) by w -> out [m, n], bf16 or fp32; ``part``: sk * m * n fp32 when K splits; 128 rows a go. pub fn mm(k: *K, x: u64, x_stride: i64, w: B16, out: u64, f32_out: bool, m: i64, part: u64) !void { const bm: i64 = if (m <= 16) 16 else if (m <= 32) 32 else if (m <= 64) 64 else 128; const sk = splitK(w.n, w.k); @@ -408,8 +399,7 @@ pub fn mm(k: *K, x: u64, x_stride: i64, w: B16, out: u64, f32_out: bool, m: i64, } } -/// The loader's row-major linear (bf16 [n, k], or packed codes and scales) at ``src`` -> its tiled copy at ``dst`` -/// (the same bytes; scales ``ss`` -> ``ds``). +/// The loader's row-major linear (bf16 or codes and scales) at ``src`` -> its tiled copy at ``dst`` (scales ss -> ds). pub fn mmTile(k: *K, bits: u8, src: u64, ss: u64, dst: u64, ds: u64, n: i64, kk: i64) !void { const chunks = @divExact(n, 8) * @divExact(kk, 64) * 32; var a: cuda.Args = .{}; @@ -431,360 +421,38 @@ pub fn mmTile(k: *K, bits: u8, src: u64, ss: u64, dst: u64, ds: u64, n: i64, kk: try k.go(try k.ext(.mm, try std.fmt.bufPrint(&nb, "glm_mm_tile{d}", .{bits})), .{ .x = @intCast(cdiv(chunks, 256)) }, 256, 0, &a); } -// -- latent.cu ---------------------------------------------------------------------------------------------------------- -pub fn absorbQ(k: *K, q: u64, wk: u64, qa: u64, R: i64) !void { - var a: cuda.Args = .{}; - a.add(q); - a.add(wk); - a.add(qa); - a.add(@as(i32, @intCast(R))); - try k.go(try k.ext(.latent, "glm_absorb_q"), .{ .x = @intCast(HL), .y = @intCast(@divExact(KV_LORA, 64)), .z = @intCast(cdiv(R, 16)) }, 64, 0, &a); -} - -pub fn expandV(k: *K, ol: u64, wv: u64, out: u64, R: i64) !void { - var a: cuda.Args = .{}; - a.add(ol); - a.add(wv); - a.add(out); - a.add(@as(i32, @intCast(R))); - try k.go(try k.ext(.latent, "glm_expand_v"), .{ .x = @intCast(HL), .y = @intCast(@divExact(DV, 8)), .z = @intCast(cdiv(R, 16)) }, 256, 0, &a); -} - -/// each row's latent into its stream's cache ``ci`` at its position. -pub fn latentWrite(k: *K, lat: u64, lat_stride: i64, t: Rows, ci: usize, R: i64) !void { - var a: cuda.Args = .{}; - a.add(lat); - a.add(lat_stride); - a.add(t.lcAt(ci)); - a.add(t.posAt()); - try k.go(try k.ext(.latent, "glm_lat_write"), .{ .x = @intCast(R) }, 64, 0, &a); -} - -const ATT_SMEM: u32 = 32 * 1024 + 2 * 16 * 1024 + 2 * 32 * 16 * 4 + 32 * 16 * 2 + 3 * 32 * 4; - -/// Dense causal attention of qa [R, H, 512] (rows below the dense limit) over ``nch`` 512-key chunks -> out. -pub fn denseAttention(k: *K, qa: u64, t: Rows, ci: usize, po: u64, pm: u64, pl: u64, out: u64, R: i64, nch: i64) !void { - const f = try k.ext(.latent, "glm_attn_dense"); - try f.allowDynamicShared(ATT_SMEM); - var a: cuda.Args = .{}; - a.add(qa); - a.add(t.lcAt(ci)); - a.add(t.posAt()); - a.add(po); - a.add(pm); - a.add(pl); - a.add(@as(i32, @intCast(R))); - a.add(SCALE); - try k.go(f, .{ .x = @intCast(R), .y = @intCast(nch) }, 256, ATT_SMEM, &a); - try merge(k, po, pm, pl, out, t.posAt(), 0, R, nch, false); -} - -/// Attention of the rows with counts > 0 over their selected tokens, written into out (others untouched). -pub fn sparseAttention(k: *K, qa: u64, t: Rows, ci: usize, tokens: u64, counts: u64, po: u64, pm: u64, pl: u64, out: u64, R: i64) !void { - const f = try k.ext(.latent, "glm_attn_sparse"); - try f.allowDynamicShared(ATT_SMEM); - var a: cuda.Args = .{}; - a.add(qa); - a.add(t.lcAt(ci)); - a.add(tokens); - a.add(counts); - a.add(po); - a.add(pm); - a.add(pl); - a.add(@as(i32, @intCast(R))); - a.add(SCALE); - try k.go(f, .{ .x = @intCast(R), .y = @intCast(SPARSE_CHUNKS) }, 256, ATT_SMEM, &a); - try merge(k, po, pm, pl, out, 0, counts, R, SPARSE_CHUNKS, true); -} - -fn merge(k: *K, po: u64, pm: u64, pl: u64, out: u64, pos: u64, counts: u64, R: i64, nch: i64, sparse: bool) !void { - var a: cuda.Args = .{}; - a.add(po); - a.add(pm); - a.add(pl); - a.add(out); - a.add(pos); - a.add(counts); - a.add(@as(i32, @intCast(R))); - a.add(@as(i32, @intCast(nch))); - a.add(@as(i32, @intFromBool(sparse))); - try k.go(try k.ext(.latent, "glm_attn_merge"), .{ .x = @intCast(R), .y = @intCast(HL) }, 128, 0, &a); -} - -// -- index.cu: the DSA indexer ------------------------------------------------------------------------------------------- -/// The rows' index keys (LayerNorm) and gates into their streams' rings for cache ``ci``, then the pools they complete. -pub fn indexUpdate(k: *K, kr: u64, k_stride: i64, gate: u64, ln_w: u64, ln_b: u64, ape: u64, t: Rows, ci: usize, R: i64) !void { - var a: cuda.Args = .{}; - a.add(kr); - a.add(k_stride); - a.add(gate); - a.add(ln_w); - a.add(ln_b); - a.add(t.ikAt(ci)); - a.add(t.igAt(ci)); - a.add(t.posAt()); - a.add(@as(f32, 1e-6)); - try k.go(try k.ext(.index, "glm_index_write"), .{ .x = @intCast(R) }, 128, 0, &a); - var b: cuda.Args = .{}; - b.add(t.ikAt(ci)); - b.add(t.igAt(ci)); - b.add(ape); - b.add(t.pkAt(ci)); - b.add(t.posAt()); - try k.go(try k.ext(.index, "glm_pool_keys"), .{ .x = @intCast(R) }, 128, 0, &b); -} - -/// Pools to score for rows up to position ``last`` (the visible ones rounded up to a power of two, at least 1024, at -/// most ``np_max``). -pub fn poolBucket(last: i64, np_max: i64) i64 { - const visible = @divFloor(last + 1, POOL) + 1; - return @min(np_max, @max(1024, @as(i64, @intCast(nextPow2(@intCast(visible)))))); -} - -/// Each of n rows' scores over np pools -> scores [n, np] fp32. -pub fn scores(k: *K, qi: u64, w: u64, w_stride: i64, t: Rows, ci: usize, out: u64, n: i64, np: i64) !void { - var a: cuda.Args = .{}; - a.add(qi); - a.add(w); - a.add(w_stride); - a.add(t.pkAt(ci)); - a.add(out); - a.add(t.posAt()); - a.add(@as(i32, @intCast(np))); - a.add(@as(f32, 0.08838834764831845)); - a.add(@as(f32, 1.0 / 5.656854249492381)); - try k.go(try k.ext(.index, "glm_index_scores"), .{ .x = @intCast(n), .y = @intCast(cdiv(np, 64)) }, 128, 0, &a); -} - -/// Each row's 512 best visible pools, ascending (int64). -pub fn selectRows(k: *K, sc: u64, out: u64, np: i64, t: Rows, n: i64) !void { - var a: cuda.Args = .{}; - a.add(sc); - a.add(out); - a.add(@as(i32, @intCast(np))); - a.add(t.posAt()); - try k.go(try k.ext(.index, "glm_index_select"), .{ .x = @intCast(n) }, 1024, 0, &a); -} - -/// Each row's attended tokens and their count (0: a dense row). -pub fn sparseTokens(k: *K, pools: u64, t: Rows, rows: i64, tokens: u64, counts: u64) !void { - var a: cuda.Args = .{}; - a.add(pools); - a.add(t.posAt()); - a.add(tokens); - a.add(counts); - try k.go(try k.ext(.index, "glm_index_tokens"), .{ .x = @intCast(rows) }, 256, 0, &a); -} - -/// A batched step's KDA layer ``li``: every stream's window rows through its own states (segs: [n] device KdaSeg). -pub fn kdaChainSeg(k: *K, segs: u64, nseg: i64, li: usize, p: u64, p_stride: i64, b_off: i64, a_: u64, a_stride: i64, g: u64, g_stride: i64, cw: u64, a_log: u64, dt_bias: u64, norm_w: u64, out: u64, max_rows: i64, proj_w: i64) !void { - var a: cuda.Args = .{}; - a.add(segs); - a.add(@as(i32, @intCast(li))); - a.add(@as(i32, @intCast(LL))); - a.add(p); - a.add(@as(i32, @intCast(p_stride))); - a.add(@as(i32, @intCast(b_off))); - a.add(a_); - a.add(@as(i32, @intCast(a_stride))); - a.add(g); - a.add(@as(i32, @intCast(g_stride))); - a.add(cw); - a.add(a_log); - a.add(dt_bias); - a.add(norm_w); - a.add(EPS); - a.add(LOWER); - a.add(out); - a.add(@as(i32, @intCast(max_rows))); - a.add(@as(i32, @intCast(proj_w))); - try k.go(try k.ext(.kda, "glm_kda_chain_seg"), .{ .x = @intCast(LL), .y = @intCast(nseg) }, 1024, 0, &a); -} - -// -- kda.py: the KDA extension (glm/kda.cu) -------------------------------------------------------------------------- -pub const KdaScratch = struct { out: u64, k: u64, v: u64, g: u64, b: u64 }; - -/// kda.chain: projection rows p [q | k | v | f_a | g_a | b at b_off], bf16 gate rows a and g, the conv window and -/// state -> out rows, state_out, and the replay inputs; windows of 64 rows or more take the three-kernel path. -pub fn kdaChain(k: *K, p: u64, p_stride: i64, b_off: i64, a: u64, a_stride: i64, g: u64, g_stride: i64, cs: u64, cw: u64, state_in: u64, a_log: u64, dt_bias: u64, norm_w: u64, rows: i64, sc: KdaScratch, state_out: u64, q_tmp: u64, y_tmp: u64) !void { - const H: i32 = @intCast(LL); - if (rows >= 64) { - var pa: cuda.Args = .{}; - pa.add(H); - pa.add(p); - pa.add(@as(i32, @intCast(p_stride))); - pa.add(@as(i32, @intCast(b_off))); - pa.add(a); - pa.add(@as(i32, @intCast(a_stride))); - pa.add(cs); - pa.add(cw); - pa.add(a_log); - pa.add(dt_bias); - pa.add(LOWER); - pa.add(q_tmp); - pa.add(sc.k); - pa.add(sc.v); - pa.add(sc.g); - pa.add(sc.b); - try k.go(try k.ext(.kda, "11prep_kernel"), .{ .x = @intCast(rows), .y = @intCast(H) }, 512, 0, &pa); - var sa: cuda.Args = .{}; - sa.add(H); - sa.add(state_in); - sa.add(q_tmp); - sa.add(sc.k); - sa.add(sc.v); - sa.add(sc.g); - sa.add(sc.b); - sa.add(@as(i32, @intCast(rows))); - sa.add(y_tmp); - sa.add(state_out); - try k.go(try k.ext(.kda, "11step_kernelILi4ELi16E"), .{ .x = @intCast(H), .y = 128 / 4 / 4 }, 4 * 32, 0, &sa); - var oa: cuda.Args = .{}; - oa.add(H); - oa.add(y_tmp); - oa.add(g); - oa.add(@as(i32, @intCast(g_stride))); - oa.add(norm_w); - oa.add(EPS); - oa.add(sc.out); - try k.go(try k.ext(.kda, "10out_kernel"), .{ .x = @intCast(rows), .y = @intCast(H) }, 128, 0, &oa); - return; - } - var ca: cuda.Args = .{}; - ca.add(H); - ca.add(p); - ca.add(@as(i32, @intCast(p_stride))); - ca.add(@as(i32, @intCast(b_off))); - ca.add(a); - ca.add(@as(i32, @intCast(a_stride))); - ca.add(g); - ca.add(@as(i32, @intCast(g_stride))); - ca.add(cs); - ca.add(cw); - ca.add(state_in); - ca.add(a_log); - ca.add(dt_bias); - ca.add(norm_w); - ca.add(EPS); - ca.add(LOWER); - ca.add(@as(i32, @intCast(rows))); - ca.add(sc.out); - ca.add(state_out); - ca.add(sc.k); - ca.add(sc.v); - ca.add(sc.g); - ca.add(sc.b); - try k.go(try k.ext(.kda, "12chain_kernel"), .{ .x = @intCast(H) }, 1024, 0, &ca); -} - -/// kda.replay_layers: every layer's state after its first ``rows`` rows (states [layers, H, 128, 128]). -pub fn kdaReplayLayers(k: *K, state_in: u64, sk: u64, sv: u64, sg: u64, sb: u64, layers: i64, rows_cap: i64, rows: i64, state_out: u64) !void { - const H: i32 = @intCast(LL); - var a: cuda.Args = .{}; - a.add(H); - a.add(state_in); - a.add(@as(u64, @intCast(LL * 128 * 128))); - a.add(sk); - a.add(sv); - a.add(sg); - a.add(sb); - a.add(@as(u64, @intCast(rows_cap * LL * 128))); - a.add(@as(u64, @intCast(rows_cap * LL))); - a.add(@as(i32, @intCast(rows))); - a.add(state_out); - try k.go(try k.ext(.kda, "20replay_layers_kernel"), .{ .x = @intCast(layers * LL) }, 1024, 0, &a); -} +// -- latent.cu and index.cu (kern_attn.zig) --------------------------------------------------------------------- +const attn = @import("kern_attn.zig"); +pub const absorbQ = attn.absorbQ; +pub const expandV = attn.expandV; +pub const latentWrite = attn.latentWrite; +pub const denseAttention = attn.denseAttention; +pub const sparseAttention = attn.sparseAttention; +pub const indexUpdate = attn.indexUpdate; +pub const poolBucket = attn.poolBucket; +pub const scores = attn.scores; +pub const selectRows = attn.selectRows; +pub const sparseTokens = attn.sparseTokens; + +// -- kda.cu (kern_kda.zig) ----------------------------------------------------------------------------------------- +const kda_k = @import("kern_kda.zig"); +pub const kdaChainSeg = kda_k.kdaChainSeg; +pub const KdaScratch = kda_k.KdaScratch; +pub const kdaChain = kda_k.kdaChain; +pub const kdaReplayLayers = kda_k.kdaReplayLayers; // -- the grouped experts (cuda/experts.py, shared with Flash Next's fatbins) ---------------------------------------- pub const promptTile = @import("kern_exl3.zig").promptTile; pub const promptOldKernel = @import("kern_exl3.zig").promptOld; pub const promptSplit = @import("kern_exl3.zig").promptSplit; +const moe_k = @import("kern_moe.zig"); +pub const maxItems = moe_k.maxItems; +pub const Plan = moe_k.Plan; +pub const expertsPlan = moe_k.expertsPlan; +pub const expertsCall = moe_k.expertsCall; -pub fn maxItems(pairs: i64, experts: i64, tile: i64) i64 { - return @min(pairs, experts) + @divFloor(pairs, tile); -} - -pub const Plan = struct { members: u64, items: u64, counts: u64, rank: u64, hist: u64, tile: i64 }; - -/// experts_v7.plan: picks [R, 8] (the routed slots alone) -> the plan over the 288 experts. -pub fn expertsPlan(k: *K, picks: u64, pairs: i64, plan: Plan) !void { - var a: cuda.Args = .{}; - const E: i32 = @intCast(EXPERTS); - if (pairs <= 1024) { - a.add(picks); - a.add(@as(i32, @intCast(pairs))); - a.add(E); - a.add(@as(i32, @intCast(plan.tile))); - a.add(plan.members); - a.add(plan.items); - a.add(plan.counts); - try k.go(try k.ext(.experts, "11plan_kernelE"), .{ .x = 1 }, 1024, 0, &a); - } else { - const nblk = cdiv(pairs, 1024); - a.add(picks); - a.add(@as(i32, @intCast(pairs))); - a.add(E); - a.add(plan.rank); - a.add(plan.hist); - try k.go(try k.ext(.experts, "9plan_rankE"), .{ .x = @intCast(nblk) }, 1024, 0, &a); - var b: cuda.Args = .{}; - b.add(@as(i32, @intCast(nblk))); - b.add(E); - b.add(@as(i32, @intCast(plan.tile))); - b.add(plan.hist); - b.add(plan.items); - b.add(plan.counts); - try k.go(try k.ext(.experts, "12plan_offsetsE"), .{ .x = 1 }, 1024, 0, &b); - var s: cuda.Args = .{}; - s.add(picks); - s.add(@as(i32, @intCast(pairs))); - s.add(E); - s.add(plan.rank); - s.add(plan.hist); - s.add(plan.members); - try k.go(try k.ext(.experts, "12plan_scatterE"), .{ .x = @intCast(cdiv(pairs, 256)) }, 256, 0, &s); - } -} - -/// experts_v7.prefill / run: gate-up (epi 2) or down (epi 3 bf16 out / 0 fp32 out) over the plan; ``slots``: the -/// pairs a row of x feeds (0 for down: one pair a row). -pub fn expertsCall(k: *K, prefill: bool, epi: i64, x: u64, x_stride: i64, slots: i64, w: u64, kb: i64, nb: i64, plan: Plan, out: u64, width: i64, items: i64, limit: f32) !void { - const gs: i64 = 64; - const m: i64 = if (epi == 2) 2 else 1; - var a: cuda.Args = .{}; - a.add(x); - a.add(@as(i32, @intCast(x_stride))); - a.add(@as(i32, @intCast(slots))); - a.add(w); - a.add(@as(i32, @intCast(kb))); - a.add(@as(i32, @intCast(nb))); - a.add(plan.items); - a.add(plan.counts); - a.add(plan.members); - a.add(out); - a.add(@as(i32, @intCast(width))); - a.add(limit); - var buf: [128]u8 = undefined; - const sbv: i64 = if (k.experts_sym) 1 else 2; // a quad's scale (and bias) uint4s a group - if (prefill) { - const f = try k.ext(if (k.experts3) .experts_sym3_prefill else if (k.experts_sym) .experts_sym_prefill else .experts_prefill, try std.fmt.bufPrint(&buf, "14prefill_kernelILi{d}ELi{d}ELi{d}ELi2ELi2ELi4E", .{ gs, m, epi })); - const block: i64 = @as(i64, if (k.experts3) 24 else 32) * @divExact(4 * gs, 128) + 4 * sbv; - const su: i64 = 64 * 8 + 4 * @divExact(64, gs) * m * block; - const sb = su * 16; - const stages: i64 = if (sb * 4 <= 49152) 4 else if (sb * 3 <= 49152) 3 else 2; - const smem: u32 = @intCast(stages * sb); - try f.allowDynamicShared(smem); - const grid = items * cdiv(nb, 4); - if (grid >= 1) try k.go(f, .{ .x = @intCast(grid) }, 256, smem, &a); - } else { - const f = try k.ext(if (k.experts3) .experts_sym3 else if (k.experts_sym) .experts_sym else .experts, try std.fmt.bufPrint(&buf, "13expert_kernelILi{d}ELi{d}ELi{d}ELi2ELi4E", .{ gs, m, epi })); - const per_sm: i64 = @max(1, try f.occupancy(128, 0)); - const grid = @min(cdiv(items, 4), per_sm * k.sms); - if (grid >= 1) try k.go(f, .{ .x = @intCast(grid) }, 128, 0, &a); - } -} +// -- our small kernels (glm/ops.cu) ---------------------------------------------------------------------------- -// -- our small kernels (glm/ops.cu) --------------------------------------------------------------------------------- /// x rows (``copies`` streams each) <- feats rows: image features in place of placeholder embeddings. pub fn injectRows(k: *K, feats: u64, x: u64, rows: i64, copies: i64) !void { var a: cuda.Args = .{}; @@ -830,7 +498,7 @@ pub fn rowsTop(k: *K, logits: u64, cols: i64, id_map: u64, offset: i64, out: u64 try k.go(f, .{ .x = @intCast(rows) }, 1024, 0, &a); } -// -- exl3.cu: EXL3 routed experts (kern_exl3.zig) ----------------------------------------------------------------------- +// -- exl3.cu: EXL3 routed experts (kern_exl3.zig) --------------------------------------------------------------- const x3 = @import("kern_exl3.zig"); pub const Exl3 = x3.Exl3; pub const Exl3Scratch = x3.Exl3Scratch; diff --git a/zig/src/families/glm_cuda/kern_attn.zig b/zig/src/families/glm_cuda/kern_attn.zig new file mode 100644 index 0000000000..e730c2747f --- /dev/null +++ b/zig/src/families/glm_cuda/kern_attn.zig @@ -0,0 +1,160 @@ +//! GLM-5.3-Flash on CUDA: latent attention (glm/latent.cu) and DSA indexer (glm/index.cu) launches, via kern.zig. + +const std = @import("std"); +const cuda = @import("cuda"); +const kern = @import("kern.zig"); +const K = kern.K; +const cdiv = kern.cdiv; +const nextPow2 = kern.nextPow2; +const Rows = kern.Rows; +const HL = kern.HL; +const DV = kern.DV; +const KV_LORA = kern.KV_LORA; +const SCALE = kern.SCALE; +const SPARSE_CHUNKS = kern.SPARSE_CHUNKS; +const POOL = kern.POOL; + +// -- latent.cu ------------------------------------------------------------------------------------------------- +pub fn absorbQ(k: *K, q: u64, wk: u64, qa: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(q); + a.add(wk); + a.add(qa); + a.add(@as(i32, @intCast(R))); + try k.go(try k.ext(.latent, "glm_absorb_q"), .{ .x = @intCast(HL), .y = @intCast(@divExact(KV_LORA, 64)), .z = @intCast(cdiv(R, 16)) }, 64, 0, &a); +} + +pub fn expandV(k: *K, ol: u64, wv: u64, out: u64, R: i64) !void { + var a: cuda.Args = .{}; + a.add(ol); + a.add(wv); + a.add(out); + a.add(@as(i32, @intCast(R))); + try k.go(try k.ext(.latent, "glm_expand_v"), .{ .x = @intCast(HL), .y = @intCast(@divExact(DV, 8)), .z = @intCast(cdiv(R, 16)) }, 256, 0, &a); +} + +/// each row's latent into its stream's cache ``ci`` at its position. +pub fn latentWrite(k: *K, lat: u64, lat_stride: i64, t: Rows, ci: usize, R: i64) !void { + var a: cuda.Args = .{}; + a.add(lat); + a.add(lat_stride); + a.add(t.lcAt(ci)); + a.add(t.posAt()); + try k.go(try k.ext(.latent, "glm_lat_write"), .{ .x = @intCast(R) }, 64, 0, &a); +} + +const ATT_SMEM: u32 = 32 * 1024 + 2 * 16 * 1024 + 2 * 32 * 16 * 4 + 32 * 16 * 2 + 3 * 32 * 4; + +/// Dense causal attention of qa [R, H, 512] (rows below the dense limit) over ``nch`` 512-key chunks -> out. +pub fn denseAttention(k: *K, qa: u64, t: Rows, ci: usize, po: u64, pm: u64, pl: u64, out: u64, R: i64, nch: i64) !void { + const f = try k.ext(.latent, "glm_attn_dense"); + try f.allowDynamicShared(ATT_SMEM); + var a: cuda.Args = .{}; + a.add(qa); + a.add(t.lcAt(ci)); + a.add(t.posAt()); + a.add(po); + a.add(pm); + a.add(pl); + a.add(@as(i32, @intCast(R))); + a.add(SCALE); + try k.go(f, .{ .x = @intCast(R), .y = @intCast(nch) }, 256, ATT_SMEM, &a); + try merge(k, po, pm, pl, out, t.posAt(), 0, R, nch, false); +} + +/// Attention of the rows with counts > 0 over their selected tokens, written into out (others untouched). +pub fn sparseAttention(k: *K, qa: u64, t: Rows, ci: usize, tokens: u64, counts: u64, po: u64, pm: u64, pl: u64, out: u64, R: i64) !void { + const f = try k.ext(.latent, "glm_attn_sparse"); + try f.allowDynamicShared(ATT_SMEM); + var a: cuda.Args = .{}; + a.add(qa); + a.add(t.lcAt(ci)); + a.add(tokens); + a.add(counts); + a.add(po); + a.add(pm); + a.add(pl); + a.add(@as(i32, @intCast(R))); + a.add(SCALE); + try k.go(f, .{ .x = @intCast(R), .y = @intCast(SPARSE_CHUNKS) }, 256, ATT_SMEM, &a); + try merge(k, po, pm, pl, out, 0, counts, R, SPARSE_CHUNKS, true); +} + +fn merge(k: *K, po: u64, pm: u64, pl: u64, out: u64, pos: u64, counts: u64, R: i64, nch: i64, sparse: bool) !void { + var a: cuda.Args = .{}; + a.add(po); + a.add(pm); + a.add(pl); + a.add(out); + a.add(pos); + a.add(counts); + a.add(@as(i32, @intCast(R))); + a.add(@as(i32, @intCast(nch))); + a.add(@as(i32, @intFromBool(sparse))); + try k.go(try k.ext(.latent, "glm_attn_merge"), .{ .x = @intCast(R), .y = @intCast(HL) }, 128, 0, &a); +} + +// -- index.cu: the DSA indexer --------------------------------------------------------------------------------- + +/// The rows' index keys (LayerNorm) and gates into their streams' rings for cache ``ci``, then the pools they complete. +pub fn indexUpdate(k: *K, kr: u64, k_stride: i64, gate: u64, ln_w: u64, ln_b: u64, ape: u64, t: Rows, ci: usize, R: i64) !void { + var a: cuda.Args = .{}; + a.add(kr); + a.add(k_stride); + a.add(gate); + a.add(ln_w); + a.add(ln_b); + a.add(t.ikAt(ci)); + a.add(t.igAt(ci)); + a.add(t.posAt()); + a.add(@as(f32, 1e-6)); + try k.go(try k.ext(.index, "glm_index_write"), .{ .x = @intCast(R) }, 128, 0, &a); + var b: cuda.Args = .{}; + b.add(t.ikAt(ci)); + b.add(t.igAt(ci)); + b.add(ape); + b.add(t.pkAt(ci)); + b.add(t.posAt()); + try k.go(try k.ext(.index, "glm_pool_keys"), .{ .x = @intCast(R) }, 128, 0, &b); +} + +/// Pools to score up to position ``last``: the visible ones up to a power of two, at least 1024, at most ``np_max``. +pub fn poolBucket(last: i64, np_max: i64) i64 { + const visible = @divFloor(last + 1, POOL) + 1; + return @min(np_max, @max(1024, @as(i64, @intCast(nextPow2(@intCast(visible)))))); +} + +/// Each of n rows' scores over np pools -> scores [n, np] fp32. +pub fn scores(k: *K, qi: u64, w: u64, w_stride: i64, t: Rows, ci: usize, out: u64, n: i64, np: i64) !void { + var a: cuda.Args = .{}; + a.add(qi); + a.add(w); + a.add(w_stride); + a.add(t.pkAt(ci)); + a.add(out); + a.add(t.posAt()); + a.add(@as(i32, @intCast(np))); + a.add(@as(f32, 0.08838834764831845)); + a.add(@as(f32, 1.0 / 5.656854249492381)); + try k.go(try k.ext(.index, "glm_index_scores"), .{ .x = @intCast(n), .y = @intCast(cdiv(np, 64)) }, 128, 0, &a); +} + +/// Each row's 512 best visible pools, ascending (int64). +pub fn selectRows(k: *K, sc: u64, out: u64, np: i64, t: Rows, n: i64) !void { + var a: cuda.Args = .{}; + a.add(sc); + a.add(out); + a.add(@as(i32, @intCast(np))); + a.add(t.posAt()); + try k.go(try k.ext(.index, "glm_index_select"), .{ .x = @intCast(n) }, 1024, 0, &a); +} + +/// Each row's attended tokens and their count (0: a dense row). +pub fn sparseTokens(k: *K, pools: u64, t: Rows, rows: i64, tokens: u64, counts: u64) !void { + var a: cuda.Args = .{}; + a.add(pools); + a.add(t.posAt()); + a.add(tokens); + a.add(counts); + try k.go(try k.ext(.index, "glm_index_tokens"), .{ .x = @intCast(rows) }, 256, 0, &a); +} diff --git a/zig/src/families/glm_cuda/kern_exl3.zig b/zig/src/families/glm_cuda/kern_exl3.zig index a65d6ce6a1..eca8ea276b 100644 --- a/zig/src/families/glm_cuda/kern_exl3.zig +++ b/zig/src/families/glm_cuda/kern_exl3.zig @@ -19,13 +19,7 @@ pub const Exl3Scratch = struct { xg: u64, xu: u64, xd: u64, z: u64 }; pub const EXL3_SK: i64 = 4; // gate/up K slices in decode windows (down: 1); prompt chunks take 1 -/// The EXL3 experts' prompt kernel (TF_GLM_EXL3_PROMPT, read once): -/// fused (the default): the transposed design (exl3.cu, "prompt chunks, the transposed design"): items of 64 pairs, -/// the weights decoded into A fragments, gate and up in one block with the SwiGLU epilogue, the down projection -/// with its epilogue; no partials in memory, one accumulation chain an output (a chunking-independent result). -/// staged: items of 128 pairs, the four warps over columns, rows and words through shared memory, K split -/// TF_GLM_EXL3_PROMPT_SK ways; wide: its 8-warp block (256 columns a block). -/// old: the grouped decode kernel on items of 64 pairs, warps over K (the engine's original prompt path). +/// TF_GLM_EXL3_PROMPT (read once): fused (default, transposed, no partials), staged (K split), wide (8 warps) or old. pub const PromptKernel = enum { fused, staged, wide, old }; var prompt_kernel: ?PromptKernel = null; pub fn promptKernel() PromptKernel { @@ -48,9 +42,7 @@ pub fn promptFused() bool { pub fn promptTile() i64 { return if (promptOld() or promptFused()) 64 else 128; } -/// The staged prompt kernel's K splits (TF_GLM_EXL3_PROMPT_SK=1|2|4, 2): a short prompt's few pairs an expert get SK -/// times the blocks at a fraction of the K each. One value for every chunk, so a row's bits do not depend on the -/// chunking (a resumed prompt equals a fresh one); the chunk scratch is sized for it. Read once. +/// The staged prompt kernel's K splits (TF_GLM_EXL3_PROMPT_SK=1|2|4, default 2), one for every chunk; read once. var prompt_sk: ?i64 = null; pub fn promptSplit() i64 { if (prompt_sk == null) { @@ -60,10 +52,7 @@ pub fn promptSplit() i64 { return prompt_sk.?; } -/// The routed pairs' outputs (pair = row * 8 + slot, picks [R, 8]) -> y [R, ys, D] fp32 at slots 0..7, through -/// the plan's items (16 pairs each in a window, 64 or 128 in a prompt chunk). -/// ``shx``: every expert of the layer shares one input rotation for gate and up (weights' exl3_shx): a prompt chunk -/// rotates each token once (sc.xg holds a row a token) and the staged kernel reads a pair's row as pair >> 3. +/// Routed pairs (row * 8 + slot) -> y [R, ys, D] fp32 slots 0..7; ``shx``: gate and up share one input rotation. pub fn exl3Routed(k: *K, x: u64, x_stride: i64, picks: u64, plan: Plan, w: Exl3, sc: Exl3Scratch, y: u64, ys: i64, R: i64, width: i64, prompt: bool, shx: bool) !void { const P = R * TOPK; const items = maxItems(P, EXPERTS, plan.tile); diff --git a/zig/src/families/glm_cuda/kern_kda.zig b/zig/src/families/glm_cuda/kern_kda.zig new file mode 100644 index 0000000000..5d67eccd5f --- /dev/null +++ b/zig/src/families/glm_cuda/kern_kda.zig @@ -0,0 +1,126 @@ +//! GLM-5.3-Flash on CUDA: the KDA layers' launches (glm/kda.cu), re-exported by kern.zig. + +const cuda = @import("cuda"); +const kern = @import("kern.zig"); +const K = kern.K; +const LL = kern.LL; +const EPS = kern.EPS; +const LOWER = kern.LOWER; + +/// A batched step's KDA layer ``li``: every stream's window rows through its own states (segs: [n] device KdaSeg). +pub fn kdaChainSeg(k: *K, segs: u64, nseg: i64, li: usize, p: u64, p_stride: i64, b_off: i64, a_: u64, a_stride: i64, g: u64, g_stride: i64, cw: u64, a_log: u64, dt_bias: u64, norm_w: u64, out: u64, max_rows: i64, proj_w: i64) !void { + var a: cuda.Args = .{}; + a.add(segs); + a.add(@as(i32, @intCast(li))); + a.add(@as(i32, @intCast(LL))); + a.add(p); + a.add(@as(i32, @intCast(p_stride))); + a.add(@as(i32, @intCast(b_off))); + a.add(a_); + a.add(@as(i32, @intCast(a_stride))); + a.add(g); + a.add(@as(i32, @intCast(g_stride))); + a.add(cw); + a.add(a_log); + a.add(dt_bias); + a.add(norm_w); + a.add(EPS); + a.add(LOWER); + a.add(out); + a.add(@as(i32, @intCast(max_rows))); + a.add(@as(i32, @intCast(proj_w))); + try k.go(try k.ext(.kda, "glm_kda_chain_seg"), .{ .x = @intCast(LL), .y = @intCast(nseg) }, 1024, 0, &a); +} + +// -- kda.py: the KDA extension (glm/kda.cu) -------------------------------------------------------------------------- +pub const KdaScratch = struct { out: u64, k: u64, v: u64, g: u64, b: u64 }; + +/// kda.chain: projection, gate rows, conv window and state -> out, state_out, replay inputs (64+ rows: 3 kernels). +pub fn kdaChain(k: *K, p: u64, p_stride: i64, b_off: i64, a: u64, a_stride: i64, g: u64, g_stride: i64, cs: u64, cw: u64, state_in: u64, a_log: u64, dt_bias: u64, norm_w: u64, rows: i64, sc: KdaScratch, state_out: u64, q_tmp: u64, y_tmp: u64) !void { + const H: i32 = @intCast(LL); + if (rows >= 64) { + var pa: cuda.Args = .{}; + pa.add(H); + pa.add(p); + pa.add(@as(i32, @intCast(p_stride))); + pa.add(@as(i32, @intCast(b_off))); + pa.add(a); + pa.add(@as(i32, @intCast(a_stride))); + pa.add(cs); + pa.add(cw); + pa.add(a_log); + pa.add(dt_bias); + pa.add(LOWER); + pa.add(q_tmp); + pa.add(sc.k); + pa.add(sc.v); + pa.add(sc.g); + pa.add(sc.b); + try k.go(try k.ext(.kda, "11prep_kernel"), .{ .x = @intCast(rows), .y = @intCast(H) }, 512, 0, &pa); + var sa: cuda.Args = .{}; + sa.add(H); + sa.add(state_in); + sa.add(q_tmp); + sa.add(sc.k); + sa.add(sc.v); + sa.add(sc.g); + sa.add(sc.b); + sa.add(@as(i32, @intCast(rows))); + sa.add(y_tmp); + sa.add(state_out); + try k.go(try k.ext(.kda, "11step_kernelILi4ELi16E"), .{ .x = @intCast(H), .y = 128 / 4 / 4 }, 4 * 32, 0, &sa); + var oa: cuda.Args = .{}; + oa.add(H); + oa.add(y_tmp); + oa.add(g); + oa.add(@as(i32, @intCast(g_stride))); + oa.add(norm_w); + oa.add(EPS); + oa.add(sc.out); + try k.go(try k.ext(.kda, "10out_kernel"), .{ .x = @intCast(rows), .y = @intCast(H) }, 128, 0, &oa); + return; + } + var ca: cuda.Args = .{}; + ca.add(H); + ca.add(p); + ca.add(@as(i32, @intCast(p_stride))); + ca.add(@as(i32, @intCast(b_off))); + ca.add(a); + ca.add(@as(i32, @intCast(a_stride))); + ca.add(g); + ca.add(@as(i32, @intCast(g_stride))); + ca.add(cs); + ca.add(cw); + ca.add(state_in); + ca.add(a_log); + ca.add(dt_bias); + ca.add(norm_w); + ca.add(EPS); + ca.add(LOWER); + ca.add(@as(i32, @intCast(rows))); + ca.add(sc.out); + ca.add(state_out); + ca.add(sc.k); + ca.add(sc.v); + ca.add(sc.g); + ca.add(sc.b); + try k.go(try k.ext(.kda, "12chain_kernel"), .{ .x = @intCast(H) }, 1024, 0, &ca); +} + +/// kda.replay_layers: every layer's state after its first ``rows`` rows (states [layers, H, 128, 128]). +pub fn kdaReplayLayers(k: *K, state_in: u64, sk: u64, sv: u64, sg: u64, sb: u64, layers: i64, rows_cap: i64, rows: i64, state_out: u64) !void { + const H: i32 = @intCast(LL); + var a: cuda.Args = .{}; + a.add(H); + a.add(state_in); + a.add(@as(u64, @intCast(LL * 128 * 128))); + a.add(sk); + a.add(sv); + a.add(sg); + a.add(sb); + a.add(@as(u64, @intCast(rows_cap * LL * 128))); + a.add(@as(u64, @intCast(rows_cap * LL))); + a.add(@as(i32, @intCast(rows))); + a.add(state_out); + try k.go(try k.ext(.kda, "20replay_layers_kernel"), .{ .x = @intCast(layers * LL) }, 1024, 0, &a); +} diff --git a/zig/src/families/glm_cuda/kern_moe.zig b/zig/src/families/glm_cuda/kern_moe.zig new file mode 100644 index 0000000000..bbac1a9c45 --- /dev/null +++ b/zig/src/families/glm_cuda/kern_moe.zig @@ -0,0 +1,91 @@ +//! GLM-5.3-Flash on CUDA: the grouped routed experts' plan and calls (cuda/experts.py), re-exported by kern.zig. + +const std = @import("std"); +const cuda = @import("cuda"); +const kern = @import("kern.zig"); +const K = kern.K; +const cdiv = kern.cdiv; +const EXPERTS = kern.EXPERTS; + +pub fn maxItems(pairs: i64, experts: i64, tile: i64) i64 { + return @min(pairs, experts) + @divFloor(pairs, tile); +} + +pub const Plan = struct { members: u64, items: u64, counts: u64, rank: u64, hist: u64, tile: i64 }; + +/// experts_v7.plan: picks [R, 8] (the routed slots alone) -> the plan over the 288 experts. +pub fn expertsPlan(k: *K, picks: u64, pairs: i64, plan: Plan) !void { + var a: cuda.Args = .{}; + const E: i32 = @intCast(EXPERTS); + if (pairs <= 1024) { + a.add(picks); + a.add(@as(i32, @intCast(pairs))); + a.add(E); + a.add(@as(i32, @intCast(plan.tile))); + a.add(plan.members); + a.add(plan.items); + a.add(plan.counts); + try k.go(try k.ext(.experts, "11plan_kernelE"), .{ .x = 1 }, 1024, 0, &a); + } else { + const nblk = cdiv(pairs, 1024); + a.add(picks); + a.add(@as(i32, @intCast(pairs))); + a.add(E); + a.add(plan.rank); + a.add(plan.hist); + try k.go(try k.ext(.experts, "9plan_rankE"), .{ .x = @intCast(nblk) }, 1024, 0, &a); + var b: cuda.Args = .{}; + b.add(@as(i32, @intCast(nblk))); + b.add(E); + b.add(@as(i32, @intCast(plan.tile))); + b.add(plan.hist); + b.add(plan.items); + b.add(plan.counts); + try k.go(try k.ext(.experts, "12plan_offsetsE"), .{ .x = 1 }, 1024, 0, &b); + var s: cuda.Args = .{}; + s.add(picks); + s.add(@as(i32, @intCast(pairs))); + s.add(E); + s.add(plan.rank); + s.add(plan.hist); + s.add(plan.members); + try k.go(try k.ext(.experts, "12plan_scatterE"), .{ .x = @intCast(cdiv(pairs, 256)) }, 256, 0, &s); + } +} + +/// experts_v7.prefill / run: gate-up (epi 2) or down (epi 3 bf16 / 0 fp32) over the plan; ``slots``: pairs a row feeds. +pub fn expertsCall(k: *K, prefill: bool, epi: i64, x: u64, x_stride: i64, slots: i64, w: u64, kb: i64, nb: i64, plan: Plan, out: u64, width: i64, items: i64, limit: f32) !void { + const gs: i64 = 64; + const m: i64 = if (epi == 2) 2 else 1; + var a: cuda.Args = .{}; + a.add(x); + a.add(@as(i32, @intCast(x_stride))); + a.add(@as(i32, @intCast(slots))); + a.add(w); + a.add(@as(i32, @intCast(kb))); + a.add(@as(i32, @intCast(nb))); + a.add(plan.items); + a.add(plan.counts); + a.add(plan.members); + a.add(out); + a.add(@as(i32, @intCast(width))); + a.add(limit); + var buf: [128]u8 = undefined; + const sbv: i64 = if (k.experts_sym) 1 else 2; // a quad's scale (and bias) uint4s a group + if (prefill) { + const f = try k.ext(if (k.experts3) .experts_sym3_prefill else if (k.experts_sym) .experts_sym_prefill else .experts_prefill, try std.fmt.bufPrint(&buf, "14prefill_kernelILi{d}ELi{d}ELi{d}ELi2ELi2ELi4E", .{ gs, m, epi })); + const block: i64 = @as(i64, if (k.experts3) 24 else 32) * @divExact(4 * gs, 128) + 4 * sbv; + const su: i64 = 64 * 8 + 4 * @divExact(64, gs) * m * block; + const sb = su * 16; + const stages: i64 = if (sb * 4 <= 49152) 4 else if (sb * 3 <= 49152) 3 else 2; + const smem: u32 = @intCast(stages * sb); + try f.allowDynamicShared(smem); + const grid = items * cdiv(nb, 4); + if (grid >= 1) try k.go(f, .{ .x = @intCast(grid) }, 256, smem, &a); + } else { + const f = try k.ext(if (k.experts3) .experts_sym3 else if (k.experts_sym) .experts_sym else .experts, try std.fmt.bufPrint(&buf, "13expert_kernelILi{d}ELi{d}ELi{d}ELi2ELi4E", .{ gs, m, epi })); + const per_sm: i64 = @max(1, try f.occupancy(128, 0)); + const grid = @min(cdiv(items, 4), per_sm * k.sms); + if (grid >= 1) try k.go(f, .{ .x = @intCast(grid) }, 128, 0, &a); + } +} diff --git a/zig/src/families/glm_cuda/native.zig b/zig/src/families/glm_cuda/native.zig index 1ad4abd466..61ee1afaa1 100644 --- a/zig/src/families/glm_cuda/native.zig +++ b/zig/src/families/glm_cuda/native.zig @@ -1,9 +1,4 @@ -//! GLM-5.3-Flash (glm5_next) on two GB10s behind the native family interface: rank 0's lane core drives the engine -//! through a lane backend that first sends each call to rank 1, which replays it on its own half of the model -//! (``follow``). Every token is drawn from both ranks' vocabulary shards (an NCCL exchange), so both ranks reach the -//! same tokens and rank 1 needs nothing back. Up to 16 streams decode together, each with its own caches sized to its -//! request out of one memory pool; a released stream's kept prompt state stays for a later prompt that extends it -//! while the pool has room. +//! GLM-5.3-Flash on two GB10s: rank 0's lane backend sends each call to rank 1 to replay; tokens from both shards. const std = @import("std"); const cuda = @import("cuda"); @@ -16,6 +11,10 @@ const proto = @import("proto.zig"); const Guard = @import("guard.zig").Guard; const copy = @import("copy.zig"); const dflash = @import("dflash.zig"); +const native_follow = @import("native_follow.zig"); +const native_prompt = @import("native_prompt.zig"); +const native_draft = @import("native_draft.zig"); +const native_env = @import("native_env.zig"); const be = lanes.backend; pub const model_type = "glm5_next"; @@ -58,24 +57,21 @@ pub fn deviceWeightBytes(io: std.Io, dir: []const u8) u64 { return weights.deviceBytes(io, dir) catch 92 << 30; } -const max_rows: usize = @intCast(forward.MAX_ROWS); // a stream's verify window at most (pending + 15 drafts) -const max_depth = max_rows - 1; -const max_streams = forward.MAX_STREAMS; -/// The head stops at a draft it gives less than this (the chain's probability product), the Python engine's rule; -/// TF_GLM_DRAFT_CONFIDENCE sets another (rank 0's, sent with each draft frame). -const draft_confidence: f64 = 0.7; - -/// TF_GLM_DRAFT_MARGINAL=p (DFlash2, rank 0): a chain stops at the first draft whose own chance (given the drafts -/// before it) is below p, instead of at the chain's product below the confidence cut; 0 (the default) keeps the product rule. -fn marginalFromEnv() f64 { - const v = std.c.getenv("TF_GLM_DRAFT_MARGINAL") orelse return 0; - return std.fmt.parseFloat(f64, std.mem.span(v)) catch 0; -} +pub const max_rows: usize = @intCast(forward.MAX_ROWS); // a stream's verify window at most (pending + 15 drafts) +pub const max_depth = max_rows - 1; +pub const max_streams = forward.MAX_STREAMS; +/// The head stops at a draft whose chain probability is below this (Python's rule); TF_GLM_DRAFT_CONFIDENCE. +pub const draft_confidence: f64 = 0.7; -fn confidenceFromEnv() f64 { - const v = std.c.getenv("TF_GLM_DRAFT_CONFIDENCE") orelse return draft_confidence; - return std.fmt.parseFloat(f64, std.mem.span(v)) catch draft_confidence; -} +const marginalFromEnv = native_env.marginalFromEnv; +const confidenceFromEnv = native_env.confidenceFromEnv; +const windowFromEnv = native_env.windowFromEnv; +const envUsize = native_env.envUsize; +const drafterWanted = native_env.drafterWanted; +const dflashStepMs = native_env.dflashStepMs; +const roceWanted = native_env.roceWanted; +const roceMaxBytes = native_env.roceMaxBytes; +pub const msSince = native_env.msSince; /// TF_GLM_TRACE=1 (rank 0): every 256 rounds, the rounds' mean verify and draft times and rows, logged. const Trace = struct { @@ -124,17 +120,16 @@ const Trace = struct { t.* = .{ .on = true }; } }; -/// A stream's caches for admission: the slots a small request takes (the host's count of streams); requests take -/// what they need from the pool (``fits``). +/// A stream's caches for admission: a small request's slots; requests take what they need from the pool (``fits``). const admission_slots: i64 = 32768; /// Prompt tokens a stream's copy index covers at most (its last; TF_GLM_COPY on): about 24 MiB of host memory. -const copy_index: usize = 1 << 20; +pub const copy_index: usize = 1 << 20; const Op = proto.Op; const Writer = proto.Writer; const Reader = proto.Reader; -const Lane = struct { +pub const Lane = struct { id: u64, s: *forward.Seq, held: [max_depth]u32 = undefined, // drafts held for the next window @@ -211,7 +206,7 @@ pub const Owned = struct { w.buf.clearRetainingCapacity(); } - fn bytesOf(self: *const Owned, capacity: i64) u64 { + pub fn bytesOf(self: *const Owned, capacity: i64) u64 { return forward.seqBytes(self.e, capacity); } @@ -232,7 +227,7 @@ pub const Owned = struct { } /// Rank 0: the oldest kept state dropped on both ranks (false: none is kept). - fn evictOldest(self: *Owned) !bool { + pub fn evictOldest(self: *Owned) !bool { if (self.kept.items.len == 0) return false; const k = self.kept.orderedRemove(0); var w: Writer = .{ .gpa = self.gpa }; @@ -244,15 +239,14 @@ pub const Owned = struct { return true; } - fn keptBytes(self: *const Owned) u64 { + pub fn keptBytes(self: *const Owned) u64 { var n: u64 = 0; for (self.kept.items) |k| n += self.bytesOf(k.s.capacity); return n; } }; -/// Whether a request of ``prompt`` tokens and ``max_tokens`` more fits the pool beside the running streams (kept -/// states give way): the lane host keeps it waiting, in order, until it does. +/// Whether ``prompt`` + ``max_tokens`` fits the pool beside running streams (kept states give way); else it waits. pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { const self = of(p); if (self.guard) |g| if (g.isTight()) return false; // the machine's memory runs short: running streams finish first @@ -260,9 +254,7 @@ pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { return self.used -| self.keptBytes() + need <= self.room; } -/// The host's pool for the streams (the bytes left after the model and its buffers), set once after open. Rank 1 -/// holds the same caches out of its own memory, so the pair takes the smaller of the two ranks' pools (rank 1 sends -/// its own: the frame after its ready), with a 2% margin. +/// The streams' pool after the model, set once after open: the smaller of both ranks' pools, less 2%. pub fn setRoom(p: *anyopaque, room: u64) !void { const self = of(p); var w: Writer = .{ .gpa = self.gpa }; @@ -284,9 +276,8 @@ pub fn setRoom(p: *anyopaque, room: u64) !void { std.log.info("glm: prompt chunks of {d} rows (TF_GLM_PREFILL_ROWS), the EXL3 experts' prompt kernel {s}", .{ self.e.pbuf.rows, @tagName(@import("kern_exl3.zig").promptKernel()) }); } -/// Commits the lanes' last windows: ``kept`` of their rows (the lane core calls keep only to cut a window short, so a -/// window it accepted whole is committed here, before the lane's next call). Rank 1 commits the same. -fn settle(self: *Owned, ls: []const *Lane, kept: []const u32) !void { +/// Commits the lanes' last windows' ``kept`` rows before each lane's next call; rank 1 commits the same. +pub fn settle(self: *Owned, ls: []const *Lane, kept: []const u32) !void { var n: usize = 0; for (ls) |l| n += @intFromBool(l.rows != 0); if (n == 0) return; @@ -455,128 +446,13 @@ const vtable: be.Backend.VTable = .{ .probabilities = probabilitiesFn, }; -/// Cache slots a kept state is allocated beyond its own request (TF_GLM_KEEP_HEADROOM, 8192; 0: none). Read once. -var keep_headroom: ?i64 = null; -fn keepHeadroom() i64 { - if (keep_headroom == null) keep_headroom = if (std.c.getenv("TF_GLM_KEEP_HEADROOM")) |v| @max(0, std.fmt.parseInt(i64, std.mem.span(v), 10) catch 8192) else 8192; - return keep_headroom.?; -} - -/// The stream's prompt into its caches (a kept state it extends, or new ones from the pool); the first token is -/// drawn here. -fn prefillFn(p: *anyopaque, s: *lanes.Stream) anyerror!void { - const self = of(p); - const ids = s.prompt(); - const cap: i64 = @intCast(ids.len + s.max_new + max_rows); - if (ids.len == 0 or cap > self.e.capacity) return error.PromptTooLong; - const id = self.next_id; - self.next_id += 1; - // images: encoded, their placeholder runs found, rank 1 sent both, before the prompt (a refusal stays here) - var spans: std.ArrayList([2]i64) = .empty; - defer spans.deinit(self.gpa); - var feats: ?cuda.DeviceBuffer = null; - defer if (feats) |*f| f.free(); - if (s.images.len > 0) feats = try encodeImages(self, ids, s.images, id, &spans); - const media = s.images.len > 0; - // the longest kept state this prompt extends (not with images: their features are this prompt's alone) - var best: ?usize = null; - if (!media) for (self.kept.items, 0..) |k, i| { - if (k.ids.len >= ids.len or k.s.capacity < cap or !std.mem.eql(u32, k.ids, ids[0..k.ids.len])) continue; - if (best == null or k.ids.len > self.kept.items[best.?].ids.len) best = i; - }; - var seq: *forward.Seq = undefined; - var start: usize = 0; - var fresh: i64 = 0; - const keep_at: usize = if (s.drafts and !media) @max(1, ids.len - 1) else 0; - if (best) |i| { - const k = self.kept.orderedRemove(i); - seq = k.s; - start = k.ids.len; - self.gpa.free(k.ids); - } else { - // a state that will be kept gets headroom when the pool has it: a later turn of the conversation (its reply - // and the next message more) resumes it instead of a prompt its exact size only - const head: i64 = if (keep_at > 0) keepHeadroom() else 0; - var alloc = cap; - // (kept states give way, as in fits(): a pool full of them is not a full pool) - if (head > 0 and cap + head <= self.e.capacity and self.used -| self.keptBytes() + self.bytesOf(cap + head) <= self.room) alloc = cap + head; - while (self.used + self.bytesOf(alloc) > self.room) if (!try self.evictOldest()) return error.NoRoom; - fresh = alloc; - seq = try self.newSeq(self.next_uid, alloc); - self.next_uid += 1; - } - var w: Writer = .{ .gpa = self.gpa }; - defer w.buf.deinit(self.gpa); - try w.int(seq.uid); - try w.int(fresh); - try w.tokens(ids); - try proto.writeSampling(&w, s.sampling); - try w.int(start); - try w.int(keep_at); - try self.send(.prefill, &w); - seq.sampling = s.sampling; - s.cached = @intCast(start); - if (feats) |f| forward.attach(seq, .{ .spans = spans.items, .feats = f.ptr }); - const first = try decode.prefill(self.e, seq, ids, start, keep_at); - var history: ?copy.Copy = if (self.copy_on and s.drafts) try copy.Copy.init(self.gpa, ids, copy_index, self.copy_min, self.copy_min_reply) else null; - errdefer if (history) |*h| h.deinit(self.gpa); - try self.lanes_by.put(s, .{ .id = id, .s = seq, .first = self.take(first), .media = media, .prompt = try self.gpa.dupe(u32, ids), .copy = history }); -} +const keep_headroom = native_prompt.keep_headroom; +const keepHeadroom = native_prompt.keepHeadroom; +const prefillFn = native_prompt.prefillFn; +const encodeImages = native_prompt.encodeImages; +const imageToken = native_prompt.imageToken; -/// The prompt's images through the tower (each its own call, features in prompt order) and their placeholder runs -/// (one a image, its length the image's tokens); rank 1 gets both in an image frame. -fn encodeImages(self: *Owned, ids: []const u32, images: []const lanes.Image, id: u64, spans: *std.ArrayList([2]i64)) !cuda.DeviceBuffer { - const tower = &(self.tower orelse return error.NoVision); - var i: usize = 0; - while (i < ids.len) : (i += 1) { - if (ids[i] != self.image_token) continue; - var j = i; - while (j < ids.len and ids[j] == self.image_token) j += 1; - try spans.append(self.gpa, .{ @intCast(i), @intCast(j) }); - i = j; - } - if (spans.items.len != images.len) return error.ImageCountMismatch; - var rows: usize = 0; - for (images, spans.items) |im, sp| { - const n: usize = @intCast(@divExact(im.grid[0] * im.grid[1] * im.grid[2], 4)); - if (n != sp[1] - sp[0]) return error.ImageTokensMismatch; - rows += n; - } - var feats = try cuda.DeviceBuffer.alloc(self.ctx.d, rows * 4096 * 2); - errdefer feats.free(); - var at: usize = 0; - for (images, spans.items) |im, sp| { - try tower.encode(&self.e.k, im.patches, im.grid, feats.ptr + at * 4096 * 2); - at += @intCast(sp[1] - sp[0]); - } - var w: Writer = .{ .gpa = self.gpa }; - defer w.buf.deinit(self.gpa); - try w.int(id); - try w.int(spans.items.len); - for (spans.items) |sp| { - try w.int(sp[0]); - try w.int(sp[1]); - } - try w.int(rows); - const fat = w.buf.items.len; - try w.buf.resize(self.gpa, fat + rows * 4096 * 2); - try self.ctx.d.check(self.ctx.d.api.cuMemcpyDtoH_v2(w.buf.items[fat..].ptr, feats.ptr, rows * 4096 * 2), "image features"); - try self.send(.image, &w); - return feats; -} - -fn imageToken(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !u32 { - const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); - defer gpa.free(path); - const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); - defer gpa.free(text); - const cfg = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); - defer cfg.deinit(); - const v = cfg.value.object.get("image_token_id") orelse return error.NoImageToken; - return @intCast(v.integer); -} - -fn laneOf(self: *Owned, s: *lanes.Stream) !*Lane { +pub fn laneOf(self: *Owned, s: *lanes.Stream) !*Lane { return self.lanes_by.getPtr(s) orelse error.NoLane; } @@ -590,7 +466,7 @@ fn queueFn(_: *anyopaque, _: *lanes.Stream, _: be.Feed, _: u64) anyerror!u64 { return error.NotPipelined; } -fn readFn(p: *anyopaque, handle: u64) anyerror!u32 { +pub fn readFn(p: *anyopaque, handle: u64) anyerror!u32 { const self = of(p); if (handle >= self.next or self.next - handle > self.drawn.len) return error.NoSuchToken; return self.drawn[handle % self.drawn.len]; @@ -674,206 +550,11 @@ fn keepFn(p: *anyopaque, windows: []const be.Window, paths: []const []const u32) try settle(self, ls[0..windows.len], kept[0..windows.len]); } -fn draftFn(p: *anyopaque, requests: []const be.DraftRequest) anyerror!void { - const self = of(p); - if (requests.len > max_streams) return error.TooManyStreams; - var ls: [max_streams]*Lane = undefined; - var kept: [max_streams]u32 = undefined; - for (requests, 0..) |r, i| { - if (r.lanes != null) return error.TreesNotBuilt; - ls[i] = try laneOf(self, r.stream); - // the kept rows (a lone round has kept already and settles nothing) - kept[i] = if (r.rows) |rows| @intCast(rows.len) else ls[i].rows; - } - try settle(self, ls[0..requests.len], kept[0..requests.len]); - var follow: [max_streams][max_rows]u32 = undefined; - var reqs: [max_streams]decode.DraftReq = undefined; - var depths: [max_streams]u32 = undefined; - var w: Writer = .{ .gpa = self.gpa }; - defer w.buf.deinit(self.gpa); - try w.int(requests.len); - for (requests, 0..) |r, i| { - const l = ls[i]; - var n: usize = 1; - var first = false; - if (r.rows) |rows| { - for (rows, 0..) |row, j| if (row != j) return error.TreesNotBuilt; - n = rows.len; - @memcpy(follow[i][0..n], r.follow[0..n]); - } else { - follow[i][0] = switch (r.first orelse return error.NoFirstToken) { - .handle => |h| try readFn(p, h), - .value => |v| v, - }; - first = true; - } - const depth: u32 = @min(r.depth, max_depth); - depths[i] = depth; - if (self.e.drafter != null) chooseDrafter(self, l, if (r.rows != null) n else 0); - const df = self.e.drafter != null and l.use_df; - // a stream DFlash2 drafts for still feeds the MTP head its rows (one step, its draft unused): it stays current - reqs[i] = .{ .s = l.s, .hidden = if (first) l.s.last_hidden else l.s.hidden, .next = follow[i][0..n], .count = if (df) 1 else depth, .confidence = self.confidence, .out = l.held[0..depth], .probs = l.probs[0..depth] }; - try w.int(l.s.uid); - try w.tokens(follow[i][0..n]); - try w.int(depth); - try w.int(@intFromBool(first)); // 1: the first draft after the prompt (its hidden row: the prompt's last) - try w.int(@as(u64, @bitCast(self.confidence))); // the confidence cut (rank 1 must stop where rank 0 does: the head's steps gather) - try w.int(@intFromBool(df)); // 1: DFlash2 drafts this stream's round - } - const t0 = std.Io.Timestamp.now(self.io, .awake); - try self.send(.draft, &w); - // the MTP head: every stream's rows, the drafts of the streams it drafts for (none while DFlash2 drafts alone: TF_GLM_MIX=0) - if (feedsMtp(self)) try decode.draft(self.e, reqs[0..requests.len]); - if (self.e.drafter) |dr| { - var dq: [max_streams]dflash.Req = undefined; - for (requests, 0..) |r, i| { - const nf = reqs[i].next.len; - dq[i] = .{ .s = ls[i].s, .n = if (r.rows != null) nf else 0, .anchor = follow[i][nf - 1], .out = ls[i].held[0..max_depth], .probs = ls[i].probs[0..max_depth], .drafts = ls[i].use_df }; - } - try dflash.draft(dr, self.e, dq[0..requests.len]); - for (reqs[0..requests.len], ls[0..requests.len], depths[0..requests.len]) |*q, l, depth| { // the confidence cut, as the head's chains stop - if (!l.use_df) continue; - var chain: f64 = 1; - q.got = 0; - while (q.got < @min(depth, dflash.BLOCK - 1)) : (q.got += 1) { // a block's 7 (a copy may run longer) - chain *= l.probs[q.got]; - if (q.got > 0 and (if (self.marginal > 0) l.probs[q.got] < self.marginal else chain < q.confidence)) break; - } - } - } - if (self.trace.on) { - self.trace.draft_ms += msSince(self.io, t0); - self.trace.calls += 1; - for (reqs[0..requests.len]) |q| self.trace.drafts += q.got; - self.trace.drafts_calls_total += requests.len; - } - for (requests, 0..) |r, i| { - const l = ls[i]; - const depth: u32 = @min(r.depth, max_depth); - const got = reqs[i].got; - for (got..depth) |j| { // past the confidence cut: held with chance 0 (the allocator leaves them out) - l.held[j] = if (got > 0) l.held[got - 1] else follow[i][reqs[i].next.len - 1]; - l.probs[j] = 0; - } - l.nheld = depth; - if (self.trace.on and r.rows != null and l.trace_held > 0) self.trace.landed += @min(l.trace_held, reqs[i].next.len - 1); - l.trace_held = 0; - if (l.copy) |*h| try copyDrafts(self, l, h, follow[i][0..reqs[i].next.len], r.rows != null, depth, got); - } -} - -/// Whether a draft round runs the MTP head: always without DFlash2, and with it only while the mix may pick the head -/// (both ranks read TF_GLM_MIX, so they agree). -pub fn feedsMtp(self: *const Owned) bool { - return self.e.drafter == null or (self.mix_on and self.e.w.mtp != null); -} - -/// The source for this stream's round (DFlash2 loaded): the one whose measured tokens per round over round time is -/// higher, the other one every 8th round (its numbers stay fresh); ``kept``: the last round's committed tokens (0: none). -fn chooseDrafter(self: *Owned, l: *Lane, kept: usize) void { - const now = std.Io.Timestamp.now(self.io, .awake); - if (!self.mix_on or self.e.w.mtp == null) { - l.use_df = true; - return; - } - if (kept > 0) if (l.mix_t) |t| { - const d: usize = @intFromBool(l.use_df); - const ms = @as(f64, @floatFromInt(t.durationTo(now).toNanoseconds())) / 1e6; - l.mix_tok[d] = 0.85 * l.mix_tok[d] + 0.15 * @as(f64, @floatFromInt(kept)); - l.mix_ms[d] = 0.85 * l.mix_ms[d] + 0.15 * ms; - }; - l.mix_t = now; - l.mix_round += 1; - const df_better = l.mix_tok[1] / l.mix_ms[1] >= l.mix_tok[0] / l.mix_ms[0]; - l.use_df = if (l.mix_round % 8 == 0) !df_better else df_better; - if (self.trace.on) self.trace.df_rounds += @intFromBool(l.use_df); -} - -/// A copy over the head's drafts when its expected landed tokens are more (rank 0: the verify frame carries the -/// held tokens to rank 1). ``follow``: the tokens just committed; ``kept_rows``: the last window's rows are known. -fn copyDrafts(self: *Owned, l: *Lane, h: *copy.Copy, follow: []const u32, kept_rows: bool, depth: u32, got: usize) !void { - const landed: u32 = @intCast(follow.len - 1); - if (kept_rows and l.copy_verified > 0) self.copy_rate.record(l.copy_verified, landed); - if (kept_rows and l.draft_verified > 0) self.draft_rate.record(l.draft_verified, landed); - l.copy_verified = 0; - l.draft_verified = 0; - l.copied = 0; - try h.append(self.gpa, follow); - if (depth == 0) return; - var buf: [max_depth]u32 = undefined; - const n = h.propose(buf[0..depth]); - if (n == 0) return; - const q = self.copy_rate.chance(); - // the copy agrees with every draft and runs on: the drafts stay, the copy's tail follows them - var same: usize = 0; - while (same < @min(n, got) and buf[same] == l.held[same]) same += 1; - const extends = same == got and n > got; - if (!extends) { // else the source whose drafts are expected to land more, by their measured rates - const mine = expected(q, n); - const theirs = expected(self.draft_rate.chance(), got); - if (mine <= theirs) return; - } - const keep = if (extends) got else 0; - @memcpy(l.held[keep..n], buf[keep..n]); - for (l.probs[keep..n]) |*pr| pr.* = q; - for (n..depth) |j| { - l.held[j] = buf[n - 1]; - l.probs[j] = 0; - } - l.copied = @intCast(n); - if (self.trace.on) { - self.trace.copies += 1; - self.trace.copy_tokens += n; - } -} - -/// Tokens a chain of ``n`` drafts lands on average when each lands with chance ``p`` given the ones before. -fn expected(p: f64, n: usize) f64 { - var sum: f64 = 0; - var chain: f64 = 1; - for (0..n) |_| { - chain *= p; - sum += chain; - } - return sum; -} - -/// TF_GLM_WINDOW: a stream's verify window at most, 2 to 16 rows (16: long copy drafts verified in one window). -fn windowFromEnv() u32 { - const v = std.c.getenv("TF_GLM_WINDOW") orelse return max_rows; - const n = std.fmt.parseInt(u32, std.mem.span(v), 10) catch return max_rows; - return @max(2, @min(n, max_rows)); -} - -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 drafterWanted() bool { - const v = std.c.getenv("TF_GLM_DRAFTER") orelse return false; - return std.mem.eql(u8, std.mem.span(v), "dflash2"); -} - -/// TF_GLM_DFLASH2_STEP_MS: a DFlash2 round's cost spread over its 7 drafts (1.6 ms: about 11 ms a round). -fn dflashStepMs() f64 { - const v = std.c.getenv("TF_GLM_DFLASH2_STEP_MS") orelse return 1.6; - return std.fmt.parseFloat(f64, std.mem.span(v)) catch 1.6; -} - -fn roceWanted() bool { - const v = std.c.getenv("TF_ROCE") orelse return true; - return !std.mem.eql(u8, std.mem.span(v), "0"); -} - -fn roceMaxBytes() usize { - const v = std.c.getenv("TF_ROCE_MAX_KB") orelse return 512 << 10; - return (std.fmt.parseInt(usize, std.mem.span(v), 10) catch 512) << 10; -} - -fn msSince(io: std.Io, t0: std.Io.Timestamp) f64 { - return @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(io, .awake)).toNanoseconds())) / 1e6; -} +const draftFn = native_draft.draftFn; +pub const feedsMtp = native_draft.feedsMtp; +const chooseDrafter = native_draft.chooseDrafter; +const copyDrafts = native_draft.copyDrafts; +const expected = native_draft.expected; fn probabilitiesFn(p: *anyopaque, s: *lanes.Stream, out: []f64) anyerror!bool { const self = of(p); @@ -883,8 +564,7 @@ fn probabilitiesFn(p: *anyopaque, s: *lanes.Stream, out: []f64) anyerror!bool { return true; } -/// The stream's lane ends: its caches kept at the prompt's kept point for a later prompt (rank 0 decides, rank 1 hears -/// it), or freed. +/// The stream's lane ends: its caches kept at the prompt's kept point for a later prompt (rank 0 decides), or freed. fn releaseFn(p: *anyopaque, s: *lanes.Stream) void { const self = of(p); var l = self.lanes_by.get(s) orelse return; @@ -907,93 +587,4 @@ fn releaseFn(p: *anyopaque, s: *lanes.Stream) void { self.gpa.free(l.prompt); } -/// Rank 1: rank 0's calls, replayed in order on this rank's half of the model, until rank 0 stops. -fn followLoop(p: *anyopaque) anyerror!void { - const self = of(p); - var picks: [forward.STEP_ROWS]decode.Pick = undefined; - var held: [max_streams][max_depth]u32 = undefined; - var probs: [max_streams][max_depth]f64 = undefined; - while (true) { - const m = self.link.recv(self.gpa) catch |e| switch (e) { - error.PeerClosed => return, - else => return e, - }; - defer self.gpa.free(m.bytes); - var r: Reader = .{ .bytes = m.bytes }; - switch (@as(Op, @fromBackingInt(@intCast(m.tag)))) { - .prefill => { - const uid = try r.int(); - const fresh: i64 = @intCast(try r.int()); - const toks = try r.tokens(); - const sampling = try proto.readSampling(&r); - const start: usize = @intCast(try r.int()); - const keep_at: usize = @intCast(try r.int()); - const seq = if (fresh > 0) try self.newSeq(uid, fresh) else self.seqs.get(uid) orelse return error.NoSuchSequence; - seq.sampling = sampling; - defer { // the image frame's features served this prefill alone - self.media_spans.clearRetainingCapacity(); - if (self.media_feats) |*f| f.free(); - self.media_feats = null; - } - if (self.media_feats) |f| forward.attach(seq, .{ .spans = self.media_spans.items, .feats = f.ptr }); - _ = try decode.prefill(self.e, seq, toks, start, keep_at); - }, - .image => { - _ = try r.int(); - const n: usize = @intCast(try r.int()); - self.media_spans.clearRetainingCapacity(); - for (0..n) |_| try self.media_spans.append(self.gpa, .{ @intCast(try r.int()), @intCast(try r.int()) }); - const rows: usize = @intCast(try r.int()); - const bytes = try r.raw(rows * 4096 * 2); - if (self.media_feats) |*f| f.free(); - self.media_feats = try cuda.DeviceBuffer.alloc(self.ctx.d, rows * 4096 * 2); - try self.media_feats.?.upload(0, bytes); - }, - .verify => { - const n: usize = @intCast(try r.int()); - var wins: [max_streams]decode.Win = undefined; - for (0..n) |i| { - const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; - wins[i] = .{ .s = s, .ids = try r.tokens() }; - } - try decode.verify(self.e, wins[0..n], &picks); - }, - .keep => { - const n: usize = @intCast(try r.int()); - for (0..n) |_| { - const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; - const rows: i64 = @intCast(try r.int()); - try decode.keep(self.e, s, rows, @intCast(try r.int())); - } - }, - .draft => { - const n: usize = @intCast(try r.int()); - var reqs: [max_streams]decode.DraftReq = undefined; - var dq: [max_streams]dflash.Req = undefined; - for (0..n) |i| { - const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; - const follow = try r.tokens(); - const depth: usize = @intCast(try r.int()); - const first = try r.int() == 1; - const cut: f64 = @bitCast(try r.int()); - const df = try r.int() == 1; - reqs[i] = .{ .s = s, .hidden = if (first) s.last_hidden else s.hidden, .next = follow, .count = if (df) 1 else depth, .confidence = cut, .out = held[i][0..depth], .probs = probs[i][0..depth] }; - dq[i] = .{ .s = s, .n = if (first) 0 else follow.len, .anchor = follow[follow.len - 1], .drafts = df }; - } - if (feedsMtp(self)) try decode.draft(self.e, reqs[0..n]); // as rank 0: the MTP head first, then DFlash2 - if (self.e.drafter) |dr| try dflash.draft(dr, self.e, dq[0..n]); - }, - .release => { - const uid = try r.int(); - if (try r.int() == 0) self.dropSeq(uid); - }, - .evict => self.dropSeq(try r.int()), - .stop => { // rank 0 closes the link next: no lost peer - if (self.guard) |g| g.stop(self.gpa); - self.guard = null; - return; - }, - .ready, .room => {}, - } - } -} +const followLoop = native_follow.followLoop; diff --git a/zig/src/families/glm_cuda/native_draft.zig b/zig/src/families/glm_cuda/native_draft.zig new file mode 100644 index 0000000000..76568466f6 --- /dev/null +++ b/zig/src/families/glm_cuda/native_draft.zig @@ -0,0 +1,181 @@ +//! The GLM backend's draft round: the MTP head, DFlash2 or a copy, chosen per stream by measured yield. +const std = @import("std"); +const decode = @import("decode.zig"); +const copy = @import("copy.zig"); +const dflash = @import("dflash.zig"); +const be = lanes.backend; +const Writer = proto.Writer; +const lanes = @import("lanes"); +const proto = @import("proto.zig"); +const native = @import("native.zig"); +const max_rows = native.max_rows; +const max_depth = native.max_depth; +const max_streams = native.max_streams; +const Lane = native.Lane; +const Owned = native.Owned; +const settle = native.settle; +const of = native.of; +const laneOf = native.laneOf; +const readFn = native.readFn; +const msSince = native.msSince; + +pub fn draftFn(p: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + const self = of(p); + if (requests.len > max_streams) return error.TooManyStreams; + var ls: [max_streams]*Lane = undefined; + var kept: [max_streams]u32 = undefined; + for (requests, 0..) |r, i| { + if (r.lanes != null) return error.TreesNotBuilt; + ls[i] = try laneOf(self, r.stream); + // the kept rows (a lone round has kept already and settles nothing) + kept[i] = if (r.rows) |rows| @intCast(rows.len) else ls[i].rows; + } + try settle(self, ls[0..requests.len], kept[0..requests.len]); + var follow: [max_streams][max_rows]u32 = undefined; + var reqs: [max_streams]decode.DraftReq = undefined; + var depths: [max_streams]u32 = undefined; + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(requests.len); + for (requests, 0..) |r, i| { + const l = ls[i]; + var n: usize = 1; + var first = false; + if (r.rows) |rows| { + for (rows, 0..) |row, j| if (row != j) return error.TreesNotBuilt; + n = rows.len; + @memcpy(follow[i][0..n], r.follow[0..n]); + } else { + follow[i][0] = switch (r.first orelse return error.NoFirstToken) { + .handle => |h| try readFn(p, h), + .value => |v| v, + }; + first = true; + } + const depth: u32 = @min(r.depth, max_depth); + depths[i] = depth; + if (self.e.drafter != null) chooseDrafter(self, l, if (r.rows != null) n else 0); + const df = self.e.drafter != null and l.use_df; + // a stream DFlash2 drafts for still feeds the MTP head its rows (one step, its draft unused): it stays current + reqs[i] = .{ .s = l.s, .hidden = if (first) l.s.last_hidden else l.s.hidden, .next = follow[i][0..n], .count = if (df) 1 else depth, .confidence = self.confidence, .out = l.held[0..depth], .probs = l.probs[0..depth] }; + try w.int(l.s.uid); + try w.tokens(follow[i][0..n]); + try w.int(depth); + try w.int(@intFromBool(first)); // 1: the first draft after the prompt (its hidden row: the prompt's last) + try w.int(@as(u64, @bitCast(self.confidence))); // the confidence cut (rank 1 must stop where rank 0 does: the head's steps gather) + try w.int(@intFromBool(df)); // 1: DFlash2 drafts this stream's round + } + const t0 = std.Io.Timestamp.now(self.io, .awake); + try self.send(.draft, &w); + // the MTP head: every stream's rows, drafts for the streams it drafts for (none if DFlash2 alone: TF_GLM_MIX=0) + if (feedsMtp(self)) try decode.draft(self.e, reqs[0..requests.len]); + if (self.e.drafter) |dr| { + var dq: [max_streams]dflash.Req = undefined; + for (requests, 0..) |r, i| { + const nf = reqs[i].next.len; + dq[i] = .{ .s = ls[i].s, .n = if (r.rows != null) nf else 0, .anchor = follow[i][nf - 1], .out = ls[i].held[0..max_depth], .probs = ls[i].probs[0..max_depth], .drafts = ls[i].use_df }; + } + try dflash.draft(dr, self.e, dq[0..requests.len]); + for (reqs[0..requests.len], ls[0..requests.len], depths[0..requests.len]) |*q, l, depth| { // the confidence cut, as the head's chains stop + if (!l.use_df) continue; + var chain: f64 = 1; + q.got = 0; + while (q.got < @min(depth, dflash.BLOCK - 1)) : (q.got += 1) { // a block's 7 (a copy may run longer) + chain *= l.probs[q.got]; + if (q.got > 0 and (if (self.marginal > 0) l.probs[q.got] < self.marginal else chain < q.confidence)) break; + } + } + } + if (self.trace.on) { + self.trace.draft_ms += msSince(self.io, t0); + self.trace.calls += 1; + for (reqs[0..requests.len]) |q| self.trace.drafts += q.got; + self.trace.drafts_calls_total += requests.len; + } + for (requests, 0..) |r, i| { + const l = ls[i]; + const depth: u32 = @min(r.depth, max_depth); + const got = reqs[i].got; + for (got..depth) |j| { // past the confidence cut: held with chance 0 (the allocator leaves them out) + l.held[j] = if (got > 0) l.held[got - 1] else follow[i][reqs[i].next.len - 1]; + l.probs[j] = 0; + } + l.nheld = depth; + if (self.trace.on and r.rows != null and l.trace_held > 0) self.trace.landed += @min(l.trace_held, reqs[i].next.len - 1); + l.trace_held = 0; + if (l.copy) |*h| try copyDrafts(self, l, h, follow[i][0..reqs[i].next.len], r.rows != null, depth, got); + } +} + +/// Whether a draft round runs the MTP head: always without DFlash2, else while the mix may pick it (TF_GLM_MIX). +pub fn feedsMtp(self: *const Owned) bool { + return self.e.drafter == null or (self.mix_on and self.e.w.mtp != null); +} + +/// This stream's draft source: the higher tokens per round time, the other every 8th round; ``kept``: last commit. +pub fn chooseDrafter(self: *Owned, l: *Lane, kept: usize) void { + const now = std.Io.Timestamp.now(self.io, .awake); + if (!self.mix_on or self.e.w.mtp == null) { + l.use_df = true; + return; + } + if (kept > 0) if (l.mix_t) |t| { + const d: usize = @intFromBool(l.use_df); + const ms = @as(f64, @floatFromInt(t.durationTo(now).toNanoseconds())) / 1e6; + l.mix_tok[d] = 0.85 * l.mix_tok[d] + 0.15 * @as(f64, @floatFromInt(kept)); + l.mix_ms[d] = 0.85 * l.mix_ms[d] + 0.15 * ms; + }; + l.mix_t = now; + l.mix_round += 1; + const df_better = l.mix_tok[1] / l.mix_ms[1] >= l.mix_tok[0] / l.mix_ms[0]; + l.use_df = if (l.mix_round % 8 == 0) !df_better else df_better; + if (self.trace.on) self.trace.df_rounds += @intFromBool(l.use_df); +} + +/// A copy over the head's drafts when it expects more landed tokens; ``follow``: just committed; ``kept_rows``. +pub fn copyDrafts(self: *Owned, l: *Lane, h: *copy.Copy, follow: []const u32, kept_rows: bool, depth: u32, got: usize) !void { + const landed: u32 = @intCast(follow.len - 1); + if (kept_rows and l.copy_verified > 0) self.copy_rate.record(l.copy_verified, landed); + if (kept_rows and l.draft_verified > 0) self.draft_rate.record(l.draft_verified, landed); + l.copy_verified = 0; + l.draft_verified = 0; + l.copied = 0; + try h.append(self.gpa, follow); + if (depth == 0) return; + var buf: [max_depth]u32 = undefined; + const n = h.propose(buf[0..depth]); + if (n == 0) return; + const q = self.copy_rate.chance(); + // the copy agrees with every draft and runs on: the drafts stay, the copy's tail follows them + var same: usize = 0; + while (same < @min(n, got) and buf[same] == l.held[same]) same += 1; + const extends = same == got and n > got; + if (!extends) { // else the source whose drafts are expected to land more, by their measured rates + const mine = expected(q, n); + const theirs = expected(self.draft_rate.chance(), got); + if (mine <= theirs) return; + } + const keep = if (extends) got else 0; + @memcpy(l.held[keep..n], buf[keep..n]); + for (l.probs[keep..n]) |*pr| pr.* = q; + for (n..depth) |j| { + l.held[j] = buf[n - 1]; + l.probs[j] = 0; + } + l.copied = @intCast(n); + if (self.trace.on) { + self.trace.copies += 1; + self.trace.copy_tokens += n; + } +} + +/// Tokens a chain of ``n`` drafts lands on average when each lands with chance ``p`` given the ones before. +pub fn expected(p: f64, n: usize) f64 { + var sum: f64 = 0; + var chain: f64 = 1; + for (0..n) |_| { + chain *= p; + sum += chain; + } + return sum; +} diff --git a/zig/src/families/glm_cuda/native_env.zig b/zig/src/families/glm_cuda/native_env.zig new file mode 100644 index 0000000000..0a6c50efb1 --- /dev/null +++ b/zig/src/families/glm_cuda/native_env.zig @@ -0,0 +1,54 @@ +//! The GLM backend's settings from the environment (TF_GLM_*, TF_ROCE*) and its timing helper. +const std = @import("std"); +const copy = @import("copy.zig"); +const native = @import("native.zig"); +const max_rows = native.max_rows; +const draft_confidence = native.draft_confidence; + +/// TF_GLM_DRAFT_MARGINAL=p (DFlash2): stop at the first draft whose own chance is below p; 0 keeps the product rule. +pub fn marginalFromEnv() f64 { + const v = std.c.getenv("TF_GLM_DRAFT_MARGINAL") orelse return 0; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch 0; +} + +pub fn confidenceFromEnv() f64 { + const v = std.c.getenv("TF_GLM_DRAFT_CONFIDENCE") orelse return draft_confidence; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch draft_confidence; +} + +/// TF_GLM_WINDOW: a stream's verify window at most, 2 to 16 rows (16: long copy drafts verified in one window). +pub fn windowFromEnv() u32 { + const v = std.c.getenv("TF_GLM_WINDOW") orelse return max_rows; + const n = std.fmt.parseInt(u32, std.mem.span(v), 10) catch return max_rows; + return @max(2, @min(n, max_rows)); +} + +pub 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 drafterWanted() bool { + const v = std.c.getenv("TF_GLM_DRAFTER") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "dflash2"); +} + +/// TF_GLM_DFLASH2_STEP_MS: a DFlash2 round's cost spread over its 7 drafts (1.6 ms: about 11 ms a round). +pub fn dflashStepMs() f64 { + const v = std.c.getenv("TF_GLM_DFLASH2_STEP_MS") orelse return 1.6; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch 1.6; +} + +pub fn roceWanted() bool { + const v = std.c.getenv("TF_ROCE") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +pub fn roceMaxBytes() usize { + const v = std.c.getenv("TF_ROCE_MAX_KB") orelse return 512 << 10; + return (std.fmt.parseInt(usize, std.mem.span(v), 10) catch 512) << 10; +} + +pub fn msSince(io: std.Io, t0: std.Io.Timestamp) f64 { + return @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(io, .awake)).toNanoseconds())) / 1e6; +} diff --git a/zig/src/families/glm_cuda/native_follow.zig b/zig/src/families/glm_cuda/native_follow.zig new file mode 100644 index 0000000000..0607abb13f --- /dev/null +++ b/zig/src/families/glm_cuda/native_follow.zig @@ -0,0 +1,104 @@ +//! Rank 1's loop: each frame rank 0 sends replayed on its half of the model until rank 0 stops. +const cuda = @import("cuda"); +const forward = @import("forward.zig"); +const decode = @import("decode.zig"); +const proto = @import("proto.zig"); +const dflash = @import("dflash.zig"); +const Op = proto.Op; +const Reader = proto.Reader; +const native = @import("native.zig"); +const max_depth = native.max_depth; +const max_streams = native.max_streams; +const of = native.of; +const feedsMtp = native.feedsMtp; + +/// Rank 1: rank 0's calls, replayed in order on this rank's half of the model, until rank 0 stops. +pub fn followLoop(p: *anyopaque) anyerror!void { + const self = of(p); + var picks: [forward.STEP_ROWS]decode.Pick = undefined; + var held: [max_streams][max_depth]u32 = undefined; + var probs: [max_streams][max_depth]f64 = undefined; + while (true) { + const m = self.link.recv(self.gpa) catch |e| switch (e) { + error.PeerClosed => return, + else => return e, + }; + defer self.gpa.free(m.bytes); + var r: Reader = .{ .bytes = m.bytes }; + switch (@as(Op, @fromBackingInt(@intCast(m.tag)))) { + .prefill => { + const uid = try r.int(); + const fresh: i64 = @intCast(try r.int()); + const toks = try r.tokens(); + const sampling = try proto.readSampling(&r); + const start: usize = @intCast(try r.int()); + const keep_at: usize = @intCast(try r.int()); + const seq = if (fresh > 0) try self.newSeq(uid, fresh) else self.seqs.get(uid) orelse return error.NoSuchSequence; + seq.sampling = sampling; + defer { // the image frame's features served this prefill alone + self.media_spans.clearRetainingCapacity(); + if (self.media_feats) |*f| f.free(); + self.media_feats = null; + } + if (self.media_feats) |f| forward.attach(seq, .{ .spans = self.media_spans.items, .feats = f.ptr }); + _ = try decode.prefill(self.e, seq, toks, start, keep_at); + }, + .image => { + _ = try r.int(); + const n: usize = @intCast(try r.int()); + self.media_spans.clearRetainingCapacity(); + for (0..n) |_| try self.media_spans.append(self.gpa, .{ @intCast(try r.int()), @intCast(try r.int()) }); + const rows: usize = @intCast(try r.int()); + const bytes = try r.raw(rows * 4096 * 2); + if (self.media_feats) |*f| f.free(); + self.media_feats = try cuda.DeviceBuffer.alloc(self.ctx.d, rows * 4096 * 2); + try self.media_feats.?.upload(0, bytes); + }, + .verify => { + const n: usize = @intCast(try r.int()); + var wins: [max_streams]decode.Win = undefined; + for (0..n) |i| { + const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; + wins[i] = .{ .s = s, .ids = try r.tokens() }; + } + try decode.verify(self.e, wins[0..n], &picks); + }, + .keep => { + const n: usize = @intCast(try r.int()); + for (0..n) |_| { + const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; + const rows: i64 = @intCast(try r.int()); + try decode.keep(self.e, s, rows, @intCast(try r.int())); + } + }, + .draft => { + const n: usize = @intCast(try r.int()); + var reqs: [max_streams]decode.DraftReq = undefined; + var dq: [max_streams]dflash.Req = undefined; + for (0..n) |i| { + const s = self.seqs.get(try r.int()) orelse return error.NoSuchSequence; + const follow = try r.tokens(); + const depth: usize = @intCast(try r.int()); + const first = try r.int() == 1; + const cut: f64 = @bitCast(try r.int()); + const df = try r.int() == 1; + reqs[i] = .{ .s = s, .hidden = if (first) s.last_hidden else s.hidden, .next = follow, .count = if (df) 1 else depth, .confidence = cut, .out = held[i][0..depth], .probs = probs[i][0..depth] }; + dq[i] = .{ .s = s, .n = if (first) 0 else follow.len, .anchor = follow[follow.len - 1], .drafts = df }; + } + if (feedsMtp(self)) try decode.draft(self.e, reqs[0..n]); // as rank 0: the MTP head first, then DFlash2 + if (self.e.drafter) |dr| try dflash.draft(dr, self.e, dq[0..n]); + }, + .release => { + const uid = try r.int(); + if (try r.int() == 0) self.dropSeq(uid); + }, + .evict => self.dropSeq(try r.int()), + .stop => { // rank 0 closes the link next: no lost peer + if (self.guard) |g| g.stop(self.gpa); + self.guard = null; + return; + }, + .ready, .room => {}, + } + } +} diff --git a/zig/src/families/glm_cuda/native_prompt.zig b/zig/src/families/glm_cuda/native_prompt.zig new file mode 100644 index 0000000000..d6a5223c2c --- /dev/null +++ b/zig/src/families/glm_cuda/native_prompt.zig @@ -0,0 +1,135 @@ +//! The GLM backend's prompt pass: a stream's caches (kept or new), its images through the tower, its first token. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const forward = @import("forward.zig"); +const decode = @import("decode.zig"); +const proto = @import("proto.zig"); +const copy = @import("copy.zig"); +const be = lanes.backend; +const Writer = proto.Writer; +const native = @import("native.zig"); +const max_rows = native.max_rows; +const copy_index = native.copy_index; +const Owned = native.Owned; +const fits = native.fits; +const of = native.of; + +/// Cache slots a kept state is allocated beyond its own request (TF_GLM_KEEP_HEADROOM, 8192; 0: none). Read once. +pub var keep_headroom: ?i64 = null; + +pub fn keepHeadroom() i64 { + if (keep_headroom == null) keep_headroom = if (std.c.getenv("TF_GLM_KEEP_HEADROOM")) |v| @max(0, std.fmt.parseInt(i64, std.mem.span(v), 10) catch 8192) else 8192; + return keep_headroom.?; +} + +/// The stream's prompt into its caches (a kept state it extends, or new ones from the pool); draws the first token. +pub fn prefillFn(p: *anyopaque, s: *lanes.Stream) anyerror!void { + const self = of(p); + const ids = s.prompt(); + const cap: i64 = @intCast(ids.len + s.max_new + max_rows); + if (ids.len == 0 or cap > self.e.capacity) return error.PromptTooLong; + const id = self.next_id; + self.next_id += 1; + // images: encoded, their placeholder runs found, rank 1 sent both, before the prompt (a refusal stays here) + var spans: std.ArrayList([2]i64) = .empty; + defer spans.deinit(self.gpa); + var feats: ?cuda.DeviceBuffer = null; + defer if (feats) |*f| f.free(); + if (s.images.len > 0) feats = try encodeImages(self, ids, s.images, id, &spans); + const media = s.images.len > 0; + // the longest kept state this prompt extends (not with images: their features are this prompt's alone) + var best: ?usize = null; + if (!media) for (self.kept.items, 0..) |k, i| { + if (k.ids.len >= ids.len or k.s.capacity < cap or !std.mem.eql(u32, k.ids, ids[0..k.ids.len])) continue; + if (best == null or k.ids.len > self.kept.items[best.?].ids.len) best = i; + }; + var seq: *forward.Seq = undefined; + var start: usize = 0; + var fresh: i64 = 0; + const keep_at: usize = if (s.drafts and !media) @max(1, ids.len - 1) else 0; + if (best) |i| { + const k = self.kept.orderedRemove(i); + seq = k.s; + start = k.ids.len; + self.gpa.free(k.ids); + } else { + // a state to be kept gets headroom when the pool has it, so the conversation's next turn resumes it + const head: i64 = if (keep_at > 0) keepHeadroom() else 0; + var alloc = cap; + // (kept states give way, as in fits(): a pool full of them is not a full pool) + if (head > 0 and cap + head <= self.e.capacity and self.used -| self.keptBytes() + self.bytesOf(cap + head) <= self.room) alloc = cap + head; + while (self.used + self.bytesOf(alloc) > self.room) if (!try self.evictOldest()) return error.NoRoom; + fresh = alloc; + seq = try self.newSeq(self.next_uid, alloc); + self.next_uid += 1; + } + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(seq.uid); + try w.int(fresh); + try w.tokens(ids); + try proto.writeSampling(&w, s.sampling); + try w.int(start); + try w.int(keep_at); + try self.send(.prefill, &w); + seq.sampling = s.sampling; + s.cached = @intCast(start); + if (feats) |f| forward.attach(seq, .{ .spans = spans.items, .feats = f.ptr }); + const first = try decode.prefill(self.e, seq, ids, start, keep_at); + var history: ?copy.Copy = if (self.copy_on and s.drafts) try copy.Copy.init(self.gpa, ids, copy_index, self.copy_min, self.copy_min_reply) else null; + errdefer if (history) |*h| h.deinit(self.gpa); + try self.lanes_by.put(s, .{ .id = id, .s = seq, .first = self.take(first), .media = media, .prompt = try self.gpa.dupe(u32, ids), .copy = history }); +} + +/// The prompt's images through the tower in order, and their placeholder runs; rank 1 gets both in an image frame. +pub fn encodeImages(self: *Owned, ids: []const u32, images: []const lanes.Image, id: u64, spans: *std.ArrayList([2]i64)) !cuda.DeviceBuffer { + const tower = &(self.tower orelse return error.NoVision); + var i: usize = 0; + while (i < ids.len) : (i += 1) { + if (ids[i] != self.image_token) continue; + var j = i; + while (j < ids.len and ids[j] == self.image_token) j += 1; + try spans.append(self.gpa, .{ @intCast(i), @intCast(j) }); + i = j; + } + if (spans.items.len != images.len) return error.ImageCountMismatch; + var rows: usize = 0; + for (images, spans.items) |im, sp| { + const n: usize = @intCast(@divExact(im.grid[0] * im.grid[1] * im.grid[2], 4)); + if (n != sp[1] - sp[0]) return error.ImageTokensMismatch; + rows += n; + } + var feats = try cuda.DeviceBuffer.alloc(self.ctx.d, rows * 4096 * 2); + errdefer feats.free(); + var at: usize = 0; + for (images, spans.items) |im, sp| { + try tower.encode(&self.e.k, im.patches, im.grid, feats.ptr + at * 4096 * 2); + at += @intCast(sp[1] - sp[0]); + } + var w: Writer = .{ .gpa = self.gpa }; + defer w.buf.deinit(self.gpa); + try w.int(id); + try w.int(spans.items.len); + for (spans.items) |sp| { + try w.int(sp[0]); + try w.int(sp[1]); + } + try w.int(rows); + const fat = w.buf.items.len; + try w.buf.resize(self.gpa, fat + rows * 4096 * 2); + try self.ctx.d.check(self.ctx.d.api.cuMemcpyDtoH_v2(w.buf.items[fat..].ptr, feats.ptr, rows * 4096 * 2), "image features"); + try self.send(.image, &w); + return feats; +} + +pub fn imageToken(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !u32 { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(text); + const cfg = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer cfg.deinit(); + const v = cfg.value.object.get("image_token_id") orelse return error.NoImageToken; + return @intCast(v.integer); +} diff --git a/zig/src/families/glm_cuda/proto.zig b/zig/src/families/glm_cuda/proto.zig index 6b3fafff1f..9e63479668 100644 --- a/zig/src/families/glm_cuda/proto.zig +++ b/zig/src/families/glm_cuda/proto.zig @@ -1,5 +1,4 @@ -//! GLM-5.3-Flash's two-rank protocol: the frames rank 0 sends rank 1 (each a call rank 1 replays) and their -//! little-endian encoding. +//! GLM-5.3-Flash's two-rank protocol: rank 0's frames to rank 1 (calls it replays), little-endian. const std = @import("std"); const lanes = @import("lanes"); diff --git a/zig/src/families/glm_cuda/qpack.zig b/zig/src/families/glm_cuda/qpack.zig new file mode 100644 index 0000000000..e000380ad7 --- /dev/null +++ b/zig/src/families/glm_cuda/qpack.zig @@ -0,0 +1,90 @@ +//! Packed dense linears (ops.cu's glm_qgemv layout): bf16 rows as symmetric 4, 6 or 8-bit codes with fp16 scales. + +const std = @import("std"); +const bf16Value = @import("ckpt.zig").bf16Value; + +/// x rounded to the nearest integer, halves to even. +fn roundEven(x: f32) f32 { + 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; +} + +/// bf16 rows [n, k] -> symmetric ``bits``-bit codes + 2^(bits-1), absmax fp16 scales per 64 (halves to even). +pub fn quantize(bits: u8, src: []const u16, n: usize, k: usize, words: []u32, scales: []u16) void { + const half: i32 = @as(i32, 1) << @intCast(bits - 1); + const qmax: f32 = @floatFromInt(half - 1); + const wpc: usize = bits / 2; + const groups = k / 64; + for (0..n) |r| { + const row = src[r * k ..][0..k]; + for (0..groups) |g| { + var amax: f32 = 0; + for (row[g * 64 ..][0..64]) |h| amax = @max(amax, @abs(bf16Value(h))); + const s16: f16 = @floatCast(amax / qmax); + scales[r * groups + g] = @bitCast(s16); + const sc: f32 = s16; + for (0..4) |c| { + var q: [16]u32 = undefined; + for (0..16) |i| { + const v = bf16Value(row[g * 64 + c * 16 + i]); + const t: i32 = if (sc == 0) 0 else @intFromFloat(@max(-qmax, @min(qmax, roundEven(v / sc)))); + q[i] = @intCast(t + half); + } + pack(bits, &q, words[((r * groups + g) * 4 + c) * wpc ..][0..wpc]); + } + } + } +} + +fn pack(bits: u8, q: *const [16]u32, out: []u32) void { + @memset(out, 0); + switch (bits) { + 4 => for (0..16) |i| { + out[i / 8] |= q[i] << @intCast(4 * (i % 8)); + }, + 8 => for (0..16) |i| { + out[i / 4] |= q[i] << @intCast(8 * (i % 4)); + }, + else => { + for (0..15) |i| out[i / 5] |= q[i] << @intCast(6 * (i % 5)); + for (0..3) |j| out[j] |= ((q[15] >> @intCast(2 * j)) & 3) << 30; + }, + } +} + +test "packed rows: values and scales round trip (4, 6, 8 bits)" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(7); + const n = 3; + const k = 128; + var src: [n * k]u16 = undefined; + for (&src) |*h| h.* = @truncate(@as(u32, @bitCast(prng.random().floatNorm(f32) * 0.02)) >> 16); + src[5] = 0; // a zero inside a group + for ([_]u8{ 4, 6, 8 }) |bits| { + const words = try gpa.alloc(u32, n * k * bits / 32); + defer gpa.free(words); + var scales: [n * k / 64]u16 = undefined; + quantize(bits, &src, n, k, words, &scales); + const half: i32 = @as(i32, 1) << @intCast(bits - 1); + const qmax: f32 = @floatFromInt(half - 1); + const wpc: usize = bits / 2; + for (0..n * k / 16) |c| { + const p = words[c * wpc ..][0..wpc]; + for (0..16) |i| { + const raw: u32 = switch (bits) { + 4 => (p[i / 8] >> @intCast(4 * (i % 8))) & 15, + 8 => (p[i / 4] >> @intCast(8 * (i % 4))) & 255, + else => if (i < 15) (p[i / 5] >> @intCast(6 * (i % 5))) & 63 else (p[0] >> 30) | ((p[1] >> 30) << 2) | ((p[2] >> 30) << 4), + }; + const sc: f32 = @as(f16, @bitCast(scales[c / 4])); + const got = @as(f32, @floatFromInt(@as(i32, @intCast(raw)) - half)) * sc; + const want = bf16Value(src[c * 16 + i]); + try std.testing.expect(@abs(got - want) <= sc * 0.5 + 1e-7); + try std.testing.expect(@abs(@as(f32, @floatFromInt(@as(i32, @intCast(raw)) - half))) <= qmax); + } + } + } +} diff --git a/zig/src/families/glm_cuda/seq.zig b/zig/src/families/glm_cuda/seq.zig index d90c986118..a533b16bf4 100644 --- a/zig/src/families/glm_cuda/seq.zig +++ b/zig/src/families/glm_cuda/seq.zig @@ -1,6 +1,4 @@ -//! GLM-5.3-Flash on CUDA: a stream's committed caches (forward.py's State), sized to its request. KDA conv windows and -//! recurrent states (two parities), the main and MTP latent caches, the indexer's keys and gates (rings: only a -//! window's rows and the 3 before it are read) and pools, and the kept prompt state a later prompt can resume. +//! GLM-5.3-Flash on CUDA: a stream's committed caches (forward.py's State), sized to its request, and its kept state. const std = @import("std"); const cuda = @import("cuda"); @@ -15,8 +13,7 @@ pub const PROJ_W: i64 = 3 * LL * 128 + 128 + 128 + LL; // q | k | v | f_a | g_a pub const CONV_C: i64 = 3 * LL * 128; // the conv's channels: q | k | v pub const REC: i64 = LL * 128 * 128; // a layer's recurrent state (fp32 values) -/// Image features for a prompt: ``spans`` [begin, end) prompt positions of each image's placeholder run, in order, -/// and ``feats`` their rows [sum of runs, 4096] bf16 on the device (the caller's, live through the prefill). +/// Image features for a prompt: ``spans`` of placeholder positions in order, ``feats`` their [rows, 4096] bf16 rows. pub const Media = struct { spans: []const [2]i64, feats: u64 }; /// Rows of the index rings (TF_GLM_INDEX_RING, default 16384, at least 4096; the capacity when that is smaller). @@ -59,9 +56,7 @@ pub const Seq = struct { hidden: u64 = 0, /// The reply's draw rule (null or temperature 0: greedy); drafts stay greedy (they only propose). sampling: ?lanes.Sampling = null, - /// The kept prompt state: KDA states and conv windows at ``snap_at``, the MTP head's pending row and its cache - /// length, each cache's index rows of the point's incomplete pool. The attention rows below the point stay valid - /// in the live caches while the next prompt extends the kept prefix. 0: none. + /// The kept prompt state at ``snap_at`` (KDA, MTP row and length, index rows); attention rows stay live. 0: none. snap_at: usize = 0, snap_mtp_len: i64 = 0, snap_rec: u64 = 0, // [nkda, LL, 128, 128] fp32 @@ -69,8 +64,7 @@ pub const Seq = struct { snap_pending: u64 = 0, // [D] bf16 snap_ik: u64 = 0, // [12, 3, 128] bf16 snap_ig: u64 = 0, - // DFlash2 (dflash.zig): the context ring, its first valid position, and the last verify's rows (the step it - // ran in, the rows' first row in it, their first position, how many) + // DFlash2 (dflash.zig): context ring, its first valid position, and the last verify's step, rows and positions draft: u64 = 0, d_from: i64 = 0, v_step: u64 = 0, @@ -190,9 +184,7 @@ fn ringCaches(s: *const Seq, which: u1) []const usize { return if (which == 0) all[0..s.ndsa] else all[s.ndsa..s.caches]; } -/// The index ring ``which`` (0: the DSA layers', 1: the MTP layer's) holding rows host_pos - host_pos % 4 .. -/// host_pos + R: moved down (the rows before the window kept) when the window runs past its end. The kernels read -/// the rings through the step's tables, so nothing captured holds the old base. +/// Index ring ``which`` (0: DSA layers, 1: MTP) moved down when the window passes its end; read via step tables. pub fn ensureRing(k: *kern.K, s: *Seq, which: u1, host_pos: i64, R: i64) !void { const lo = host_pos - @mod(host_pos, 4); const base = s.ring_base[which]; diff --git a/zig/src/families/glm_cuda/vision.zig b/zig/src/families/glm_cuda/vision.zig index 9a273d7e31..b7f2896d8f 100644 --- a/zig/src/families/glm_cuda/vision.zig +++ b/zig/src/families/glm_cuda/vision.zig @@ -1,10 +1,4 @@ -//! GLM-5.3-Flash's vision tower on CUDA (rank 0 encodes; rank 1 receives the features), as mlx-vlm's -//! glm5_next/vision.py defines it: the patch embedding (a linear over a patch's 2 x 14 x 14 x 3 values), 24 blocks -//! (RMSNorm, attention with RMSNorm'd q and k and a 2-D rotary over each frame's patches, a limited SwiGLU), the -//! post-norm, the 2 x 2 downsample (a linear over each merge block's four patches) and the merger (a projection, -//! LayerNorm + GELU, a limited SwiGLU) to the language model's 4096. The linears are cuBLASLt products (bf16 in, fp32 -//! accumulate, bias added before the one rounding); the residual stream stays fp32; glm/vision.cu does the rest. -//! Checked against the torch reference (tools/vision_ref.py) by tf-cuda-test glm-vision. +//! GLM-5.3-Flash's vision tower on CUDA as glm5_next/vision.py (rank 0 encodes); checked by tf-cuda-test glm-vision. const std = @import("std"); const cuda = @import("cuda"); @@ -267,8 +261,7 @@ pub const Tower = struct { return @intCast(N * PATCH * 4 + N * H * 4 + N * H * 2 * 5 + N * MLP * 2 + 2 * ROWS * MLP * 4 + q * seg * 6 + q * HD * 4 + q * 4 + N * 8 + 16 * 1024); } - /// Features of one image or video (``grid`` [t, h, w] in patches, t frame groups) from its patches (fp32, - /// [t h w, 1176] in the processor's merge-block order) into ``out`` ([t h w / 4, 4096] bf16). + /// Features of an image or video (``grid`` [t, h, w]): fp32 patches [t h w, 1176] -> ``out`` [t h w / 4, 4096]. pub fn encode(t: *Tower, k: *kern.K, patches: []const f32, grid: [3]i64, out: u64) !void { const d = t.d; const gt = grid[0]; @@ -371,7 +364,7 @@ pub const Tower = struct { a.add(EPS); try k.go(try k.ext(.vision, "glm_vis_qkv"), .{ .x = @intCast(std.math.divCeil(i64, N * HEADS, 8) catch unreachable) }, 256, 0, &a); } - // each frame group attends within itself; per head, chunks of queries (q, k: [t, H, seg, 64]; v: [t, H, 64, seg]) + // each frame group attends within itself, per head, in query chunks (q, k [t, H, seg, 64]; v transposed) var g: i64 = 0; while (g < gt) : (g += 1) { var h: i64 = 0; diff --git a/zig/src/families/glm_cuda/weights.zig b/zig/src/families/glm_cuda/weights.zig index 5a1f36207e..88c4752789 100644 --- a/zig/src/families/glm_cuda/weights.zig +++ b/zig/src/families/glm_cuda/weights.zig @@ -1,12 +1,4 @@ -//! GLM-5.3-Flash's prepared weights for one rank of two, read from the converted checkpoint (glm-affine-experts-v1: -//! MLX affine 4-bit routed experts in groups of 64 with fp16 scales and biases, everything else BF16 / FP32), as the -//! Python engine prepares them (families/glm5_next/cuda/weights.py, its affine-experts path, split.py's rank rules): -//! every tensor under the name, shape and bytes of the Python engine's prepared weights (its weight tree, the names -//! tools/capture_ref.py hashes), uploaded to this rank's GPU. -//! -//! The derivations are exact: rank rows and columns cut on the host, BF16 rows stacked, the kv_b blocks split per head -//! (the latent path's absorb weights), fp16 expert scales rounded to bf16 (round to nearest even, as torch), and the -//! experts' blocks in experts_pack.cu's index map. The draft head is the verification head itself (no 4-bit copy). +//! GLM-5.3-Flash's prepared weights for one rank of two as weights.py prepares them, cut exactly, uploaded to its GPU. const std = @import("std"); const cuda = @import("cuda"); @@ -17,149 +9,21 @@ pub const forward = @import("forward.zig"); pub const decode = @import("decode.zig"); pub const vision = @import("vision.zig"); -const BASE = "model.language_model."; -const WORLD = 2; - -// --------------------------------------------------------------------------------------------------------------- -// The checkpoint: its shards' headers, files opened once, tensors read whole or as row ranges. - -const Entry = struct { dtype: []const u8, shape: []const i64, begin: u64, end: u64, file: usize }; - -const Ckpt = struct { - gpa: std.mem.Allocator, - io: std.Io, - arena: std.heap.ArenaAllocator, - files: std.ArrayList(std.Io.File) = .empty, - paths: std.ArrayList([]const u8) = .empty, - entries: std.StringHashMap(Entry), - - /// ``dir``: the served folder. A shard (or the index) that is a symlink to a container path under /cache/tf/ - /// resolves against the folder's parent directory, so the same folder serves inside the container and on the host. - fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Ckpt { - var ck: Ckpt = .{ .gpa = gpa, .io = io, .arena = .init(gpa), .entries = .init(gpa) }; - errdefer ck.deinit(); - const a = ck.arena.allocator(); - const index_path = try resolve(a, io, dir, "model.safetensors.index.json"); - const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, a, .limited(1 << 26)); - const idx = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); - var shard_of = std.StringHashMap(usize).init(a); - var it = idx.object.get("weight_map").?.object.iterator(); - while (it.next()) |e| { - const shard = e.value_ptr.string; - if (shard_of.contains(shard)) continue; - const path = try resolve(a, io, dir, shard); - try shard_of.put(shard, ck.files.items.len); - try ck.files.append(gpa, try std.Io.Dir.cwd().openFile(io, path, .{})); - try ck.paths.append(gpa, path); - try ck.readHeader(ck.files.items.len - 1); - } - return ck; - } - - fn readHeader(ck: *Ckpt, f: usize) !void { - const a = ck.arena.allocator(); - var head: [8]u8 = undefined; - _ = try ck.files.items[f].readPositionalAll(ck.io, &head, 0); - const n = std.mem.readInt(u64, &head, .little); - const text = try a.alloc(u8, n); - _ = try ck.files.items[f].readPositionalAll(ck.io, text, 8); - const v = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); - var it = v.object.iterator(); - while (it.next()) |e| { - if (std.mem.eql(u8, e.key_ptr.*, "__metadata__")) continue; - const o = e.value_ptr.object; - const shp = o.get("shape").?.array.items; - const shape = try a.alloc(i64, shp.len); - for (shp, 0..) |x, i| shape[i] = x.integer; - const offs = o.get("data_offsets").?.array.items; - try ck.entries.put(e.key_ptr.*, .{ .dtype = o.get("dtype").?.string, .shape = shape, .begin = 8 + n + @as(u64, @intCast(offs[0].integer)), .end = 8 + n + @as(u64, @intCast(offs[1].integer)), .file = f }); - } - } - - fn deinit(ck: *Ckpt) void { - // the shards are read once: their page cache would starve NCCL's first registration (and the n-gram tables) - for (ck.files.items) |f| { - _ = std.os.linux.fadvise(f.handle, 0, 0, std.os.linux.POSIX_FADV.DONTNEED); - f.close(ck.io); - } - ck.files.deinit(ck.gpa); - ck.paths.deinit(ck.gpa); - ck.entries.deinit(); - ck.arena.deinit(); - } - - fn has(ck: *const Ckpt, name: []const u8) bool { - return ck.entries.contains(name); - } - - fn get(ck: *const Ckpt, name: []const u8) !Entry { - return ck.entries.get(name) orelse { - std.log.err("glm weights: the checkpoint has no tensor {s}", .{name}); - return error.MissingTensor; - }; - } - - /// ``len`` bytes at ``off`` into the tensor's data. - fn readAt(ck: *const Ckpt, e: Entry, off: u64, out: []u8) !void { - if (off + out.len > e.end - e.begin) return error.ShortRead; - const got = try ck.files.items[e.file].readPositionalAll(ck.io, out, e.begin + off); - if (got != out.len) return error.ShortRead; - } - - fn readAll(ck: *const Ckpt, gpa: std.mem.Allocator, name: []const u8) ![]align(16) u8 { - const e = try ck.get(name); - const buf = try gpa.alignedAlloc(u8, .@"16", e.end - e.begin); - errdefer gpa.free(buf); - try ck.readAt(e, 0, buf); - return buf; - } -}; - -fn resolve(a: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) ![]const u8 { - const path = try std.fs.path.join(a, &.{ dir, name }); - if (std.Io.Dir.cwd().openFile(io, path, .{})) |f| { - f.close(io); - return path; - } else |_| {} - var buf: [4096]u8 = undefined; - const n = std.Io.Dir.cwd().readLink(io, path, &buf) catch return path; - const target = buf[0..n]; - const container = "/cache/tf/"; - if (std.mem.startsWith(u8, target, container)) - return std.fs.path.join(a, &.{ std.fs.path.dirname(dir) orelse ".", target[container.len..] }); - if (std.fs.path.isAbsolute(target)) return a.dupe(u8, target); - return std.fs.path.join(a, &.{ dir, target }); -} - -fn bf16ToF32(b: u16) f32 { - return @bitCast(@as(u32, b) << 16); -} - -fn f32ToBf16(x: f32) u16 { - const u: u32 = @bitCast(x); - if (std.math.isNan(x)) return 0x7FC0; - return @intCast((u + 0x7FFF + ((u >> 16) & 1)) >> 16); -} - -fn f16ToBf16(h: u16) u16 { - return f32ToBf16(@floatCast(@as(f16, @bitCast(h)))); -} - -fn numelOf(shape: []const i64) usize { - var n: usize = 1; - for (shape) |d| n *= @intCast(d); - return n; -} - -fn dtypeOf(s: []const u8) !api.DType { - if (std.mem.eql(u8, s, "BF16")) return .bf16; - if (std.mem.eql(u8, s, "F32")) return .f32; - if (std.mem.eql(u8, s, "F16")) return .f16; - return error.UnsupportedDType; -} - -// --------------------------------------------------------------------------------------------------------------- -// split.py's rules: a checkpoint tensor's rank part (rows, the last axis' columns, or the whole tensor). +const ckpt = @import("ckpt.zig"); +const ex = @import("experts_load.zig"); +const qpack = @import("qpack.zig"); +const BASE = ckpt.BASE; +const WORLD = ckpt.WORLD; +const Ckpt = ckpt.Ckpt; +const fullName = ckpt.fullName; +const numelOf = ckpt.numelOf; +const dtypeOf = ckpt.dtypeOf; +const bf16ToF32 = ckpt.bf16ToF32; +const experts = ex.experts; +const expertsExl3 = ex.expertsExl3; +pub const quantize = qpack.quantize; + +// -- split.py's rules: a checkpoint tensor's rank part (rows, the last axis' columns, or the whole tensor) ------- const Split = enum { rep, row, col }; @@ -179,7 +43,7 @@ fn splitOf(name: []const u8) Split { return .rep; } -const Builder = struct { +pub const Builder = struct { gpa: std.mem.Allocator, ck: *const Ckpt, dry: bool, @@ -191,30 +55,24 @@ const Builder = struct { /// Dense linears (all but the indexer's) and the head packed to these bits (0: bf16 as stored): see ``quantize``. dense_bits: u8 = 0, head_bits: u8 = 0, - /// TF_GLM_KDA_PROJ_BITS: the KDA layers' stacked in-projection (the largest dense linear, 1.3 GB a rank at 6 bits) - /// at its own width (0: dense_bits). + /// TF_GLM_KDA_PROJ_BITS: the KDA stacked in-projection (1.3 GB a rank at 6 bits) at its own width (0: dense_bits). kda_proj_bits: u8 = 0, /// A packed copy of the head the MTP drafts read (0: they read the head). draft_head_bits: u8 = 0, - /// The routed experts' groups as scales alone (their biases -8 x the scales, checked): TF_GLM_EXPERTS_SYM=0 keeps - /// the affine blocks. + /// The routed experts' groups as scales alone (biases -8 x scales, checked); TF_GLM_EXPERTS_SYM=0 keeps affine. experts_sym: bool = true, - /// Development (TF_GLM_EXPERT_SIM_BITS=2 or 3): the experts re-quantized to fewer levels in their 4-bit words, to - /// measure a narrower expert format's quality before its kernel exists. + /// Development (TF_GLM_EXPERT_SIM_BITS=2 or 3): the experts re-quantized to fewer levels in their 4-bit words. expert_sim_bits: u8 = 0, /// TF_GLM_EXPERT_BITS=3: the experts re-quantized to 3-bit codes and packed as glm/experts_sym3's planes. experts3: bool = false, - /// TF_GLM_EXPERT3_PROJ=down|up: 3-bit codes for that projection alone (the other stays 4-bit in the planes: a - /// sensitivity probe; the kernel reads 3-bit planes either way, so "4-bit" here means 8 levels of the 4-bit - /// value re-rounded). TF_GLM_EXPERT_SIM_BITS probes quality on the 4-bit kernels instead: with - /// TF_GLM_EXPERT_SIM_PROJ=down|up it re-quantizes that projection alone. + /// TF_GLM_EXPERT3_PROJ / TF_GLM_EXPERT_SIM_PROJ=down|up: re-quantize that projection alone (a sensitivity probe). sim_proj: u8 = 0, // 0: all, 1: down only, 2: gate/up only /// The checkpoint's routed experts are EXL3 (4-bit "mcg" trellises, every other weight BF16): read as stored. exl3: bool = false, exl3_shx: bool = true, // every layer's experts share one input rotation for gate and up (cleared by a layer that does not) suh_first: [2 * 4096]u8 = undefined, // the first expert's gate suh of the layer being read - fn put(b: *Builder, name: []const u8, dtype: api.DType, shape: []const i64, data: ?[]const u8) !void { + pub fn put(b: *Builder, name: []const u8, dtype: api.DType, shape: []const i64, data: ?[]const u8) !void { var t: api.Tensor = .{ .ptr = 0, .dtype = dtype, .nd = @intCast(shape.len) }; for (shape, 0..) |x, i| t.shape[i] = x; const n = t.bytes(); @@ -230,7 +88,7 @@ const Builder = struct { try store.put(name, t); } - fn putf(b: *Builder, comptime fmt: []const u8, args: anytype, dtype: api.DType, shape: []const i64, data: ?[]const u8) !void { + pub fn putf(b: *Builder, comptime fmt: []const u8, args: anytype, dtype: api.DType, shape: []const i64, data: ?[]const u8) !void { var nb: [256]u8 = undefined; try b.put(try std.fmt.bufPrint(&nb, fmt, args), dtype, shape, data); } @@ -350,275 +208,7 @@ const Builder = struct { } }; -fn fullName(buf: []u8, comptime fmt: []const u8, args: anytype) ![]const u8 { - return std.fmt.bufPrint(buf, BASE ++ fmt, args); -} - -// --------------------------------------------------------------------------------------------------------------- -// The routed experts: per-expert MLX triples, the rank's rows (gate, up) or columns (down), packed as experts_pack.cu. - -const Triple = struct { - e: usize, - n: usize, - k8: usize, - kg: usize, - words: []u32, - scales: []u16, - biases: []u16, - - fn free(t: Triple, gpa: std.mem.Allocator) void { - gpa.free(t.words); - gpa.free(t.scales); - gpa.free(t.biases); - } -}; - -/// Every expert's ``proj`` of layer ``layer``: rows r0..r1 (gate, up) or the rank's columns (down), scales and biases -/// rounded from fp16 to bf16. -fn readExperts(b: *Builder, li: usize, proj: []const u8, count: usize) !Triple { - var nb: [256]u8 = undefined; - const w0 = try b.ck.get(try fullName(&nb, "layers.{d}.mlp.experts.0.{s}.weight", .{ li, proj })); - const n_full: usize = @intCast(w0.shape[0]); - const k8_full: usize = @intCast(w0.shape[1]); - var nb2: [256]u8 = undefined; - const s0 = try b.ck.get(try fullName(&nb2, "layers.{d}.mlp.experts.0.{s}.scales", .{ li, proj })); - const kg_full: usize = @intCast(s0.shape[1]); - const down = std.mem.eql(u8, proj, "down_proj"); - const nr = if (down) n_full else n_full / WORLD; - const r0 = if (down) 0 else b.rank * nr; - const ok8 = if (down) k8_full / WORLD else k8_full; - const okg = if (down) kg_full / WORLD else kg_full; - const c8 = if (down) b.rank * ok8 else 0; - const cg = if (down) b.rank * okg else 0; - var t: Triple = .{ .e = count, .n = nr, .k8 = ok8, .kg = okg, .words = try b.gpa.alloc(u32, count * nr * ok8), .scales = try b.gpa.alloc(u16, count * nr * okg), .biases = try b.gpa.alloc(u16, count * nr * okg) }; - errdefer t.free(b.gpa); - const rowbuf = try b.gpa.alloc(u8, nr * k8_full * 4); - defer b.gpa.free(rowbuf); - const sbuf = try b.gpa.alloc(u8, nr * kg_full * 2); - defer b.gpa.free(sbuf); - for (0..count) |x| { - var n1: [256]u8 = undefined; - const we = try b.ck.get(try fullName(&n1, "layers.{d}.mlp.experts.{d}.{s}.weight", .{ li, x, proj })); - try b.ck.readAt(we, r0 * k8_full * 4, rowbuf); - const w32 = std.mem.bytesAsSlice(u32, rowbuf); - for (0..nr) |r| @memcpy(t.words[(x * nr + r) * ok8 ..][0..ok8], w32[r * k8_full + c8 ..][0..ok8]); - for ([_][]const u8{ "scales", "biases" }, 0..) |what, which| { - var n2: [256]u8 = undefined; - const ent = try b.ck.get(try fullName(&n2, "layers.{d}.mlp.experts.{d}.{s}.{s}", .{ li, x, proj, what })); - try b.ck.readAt(ent, r0 * kg_full * 2, sbuf); - const s16 = std.mem.bytesAsSlice(u16, sbuf); - const dst = if (which == 0) t.scales else t.biases; - for (0..nr) |r| for (0..okg) |g| { - dst[(x * nr + r) * okg + g] = f16ToBf16(s16[r * kg_full + cg + g]); - }; - } - } - return t; -} - -/// experts_pack.cu's index map: each (group, column block, expert) block holds 32 * 4 * H weight words (nibbles -/// reordered i0 i2 i4 i6 i1 i3 i5 i7) then 32 scale/bias words; written as slot ``m`` of ``mcount`` per block. -/// Every group's bias is -8 x its scale (a symmetric int4 checkpoint as affine groups): the symmetric blocks hold -/// the scales alone. -fn checkSymmetric(t: Triple) !void { - for (t.scales, t.biases) |sc, bi| { - if (bf16Value(bi) != -8 * bf16Value(sc)) return error.ExpertsNotSymmetric; - } -} - -fn packExperts(t: Triple, gs: usize, out: []u32, mcount: usize, m: usize, sym: bool, three: bool) void { - const h = gs / 32; - const Work = struct { - fn run(tt: Triple, hh: usize, o: []u32, mc: usize, mm: usize, e0: usize, e1: usize, sy: bool, th: bool) void { - const wp = 32 * 4 * hh; - const wsec = if (th) wp * 3 / 4 else wp; - const blk = wsec + @as(usize, if (sy) 16 else 32); - var tmp: [512]u32 = undefined; // a 3-bit block's 4-bit words before their planes - const nbb = tt.n / 32; - for (e0..e1) |e| for (0..nbb) |bb| for (0..tt.kg) |g| { - const dst = o[(((e * nbb + bb) * tt.kg + g) * mc + mm) * blk ..][0..blk]; - const wdst = if (th) tmp[0..wp] else dst[0..wp]; - for (0..wp) |i| { - const row = i / (4 * hh); - const kk = i % (4 * hh); - const tq = row / 8; - const r = row % 8; - const q = kk / hh; - const j = kk % hh; - const w = tt.words[(e * tt.n + bb * 32 + row) * tt.k8 + g * 4 * hh + kk]; - const f = ((tq * hh + j) * 8 + r) * 4 + q; - wdst[(f / 128) * 128 + (f % 32) * 4 + (f / 32) % 4] = shuffle(w); - } - if (th) for (0..hh) |c| for (0..32) |lane| { // the lane's uint4 c as three planes (experts_sym.cuh) - var lo: [4]u32 = undefined; - var hi: [4]u32 = undefined; - for (0..4) |k| { - const word = tmp[(c * 32 + lane) * 4 + k]; - lo[k] = 0; - hi[k] = 0; - for (0..8) |i| { - const code = ((word >> @intCast(4 * i)) & 15) -% 4; // 3-bit codes are stored + 4 - lo[k] |= (code & 3) << @intCast(2 * i); - hi[k] |= ((code >> 2) & 1) << @intCast(i); - } - } - dst[c * 96 + lane] = lo[0] | (lo[1] << 16); - dst[c * 96 + 32 + lane] = lo[2] | (lo[3] << 16); - dst[c * 96 + 64 + lane] = hi[0] | (hi[1] << 8) | (hi[2] << 16) | (hi[3] << 24); - }; - if (sy) { // quad p: rows j * 8 + 2p and + 1's scales, j = 0..3 (experts_sym.cuh) - for (0..16) |s| { - const p = s / 4; - const j = s % 4; - const at = (e * tt.n + bb * 32 + j * 8 + p * 2) * tt.kg + g; - dst[wsec + s] = @as(u32, tt.scales[at]) | (@as(u32, tt.scales[at + tt.kg]) << 16); - } - } else for (0..32) |s| { - const p = s / 8; - const tt4 = s % 4; - const src = if ((s / 4) % 2 == 1) tt.biases else tt.scales; - const at = (e * tt.n + bb * 32 + tt4 * 8 + p * 2) * tt.kg + g; - dst[wp + s] = @as(u32, src[at]) | (@as(u32, src[at + tt.kg]) << 16); - } - }; - } - }; - const threads = 16; - var pool: [threads]?std.Thread = @splat(null); - const per = (t.e + threads - 1) / threads; - for (0..threads) |i| { - const e0 = i * per; - const e1 = @min(t.e, e0 + per); - if (e0 >= e1) break; - pool[i] = std.Thread.spawn(.{}, Work.run, .{ t, h, out, mcount, m, e0, e1, sym, three }) catch blk: { - Work.run(t, h, out, mcount, m, e0, e1, sym, three); - break :blk null; - }; - } - for (pool) |th| if (th) |x| x.join(); -} - -fn shuffle(w: u32) u32 { - var even = w & 0x0F0F0F0F; - var odd = (w >> 4) & 0x0F0F0F0F; - even = (even | (even >> 4)) & 0x00FF00FF; - odd = (odd | (odd >> 4)) & 0x00FF00FF; - even = (even | (even >> 8)) & 0x0000FFFF; - odd = (odd | (odd >> 8)) & 0x0000FFFF; - return even | (odd << 16); -} - -/// grouped.make([gate, up], down, 64): up [E, N/32, K/64, 2, block], down [E, D/32, NI/64, 1, block] int32. -fn experts(b: *Builder, out: []const u8, li: usize) !void { - const E: usize = @intCast(kern.EXPERTS); - const gs: usize = 64; - if (b.experts3 and !b.experts_sym) return error.Experts3NeedsSymmetric; - const wsec: usize = 32 * 4 * (gs / 32) * @as(usize, if (b.experts3) 3 else 4) / 4; - const blk: usize = wsec + @as(usize, if (b.experts_sym) 16 else 32); - const width: usize = @intCast(@divExact(2048, WORLD)); // moe_intermediate_size / ranks - const dims: usize = @intCast(kern.D); - const up_shape = [_]i64{ @intCast(E), @intCast(width / 32), @intCast(dims / gs), 2, @intCast(blk) }; - const down_shape = [_]i64{ @intCast(E), @intCast(dims / 32), @intCast(width / gs), 1, @intCast(blk) }; - var nb: [256]u8 = undefined; - var nb2: [256]u8 = undefined; - const up_name = try std.fmt.bufPrint(&nb, "{s}.up", .{out}); - const down_name = try std.fmt.bufPrint(&nb2, "{s}.down", .{out}); - if (b.dry) { - try b.put(up_name, .i32, &up_shape, null); - try b.put(down_name, .i32, &down_shape, null); - return; - } - { - const up = try b.gpa.alloc(u32, numelOf(&up_shape)); - defer b.gpa.free(up); - for ([_][]const u8{ "gate_proj", "up_proj" }, 0..) |proj, m| { - const t = try readExperts(b, li, proj, E); - defer t.free(b.gpa); - if (t.n != width or t.k8 * 8 != dims) return error.ExpertShape; - if (b.experts_sym) try checkSymmetric(t); - if (b.sim_proj != 1) { - if (b.experts3) simulateBits(t, 3) else if (b.expert_sim_bits != 0) simulateBits(t, b.expert_sim_bits); - } - packExperts(t, gs, up, 2, m, b.experts_sym, b.experts3); - } - try b.put(up_name, .i32, &up_shape, std.mem.sliceAsBytes(up)); - } - const dn = try b.gpa.alloc(u32, numelOf(&down_shape)); - defer b.gpa.free(dn); - const t = try readExperts(b, li, "down_proj", E); - defer t.free(b.gpa); - if (t.n != dims or t.k8 * 8 != width) return error.ExpertShape; - if (b.experts_sym) try checkSymmetric(t); - if (b.sim_proj != 2) { - if (b.experts3) simulateBits(t, 3) else if (b.expert_sim_bits != 0) simulateBits(t, b.expert_sim_bits); - } - packExperts(t, gs, dn, 1, 0, b.experts_sym, b.experts3); - try b.put(down_name, .i32, &down_shape, std.mem.sliceAsBytes(dn)); -} - -/// The EXL3 routed experts of layer ``li``, this rank's whole tiles (glm/exl3.cu's layout): gate and up by output tile -/// columns, down by input tile rows, each matrix's experts stacked as trellis words [E, K/16, N/16, 32] (the stored -/// bytes) and fp16 scales [E, K] (suh) and [E, N] (svh), the rank's parts. -fn expertsExl3(b: *Builder, out: []const u8, li: usize) !void { - const E: usize = @intCast(kern.EXPERTS); - const D: usize = @intCast(kern.D); - const width: usize = @intCast(@divExact(2048, WORLD)); - const tile: usize = 128; // a 16x16 tile's bytes at 4 bits - var nb: [256]u8 = undefined; - const Mat = struct { proj: []const u8, kt: usize, nt: usize, by_cols: bool }; - const mats = [_]Mat{ .{ .proj = "gate_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, .{ .proj = "up_proj", .kt = D / 16, .nt = width / 16, .by_cols = true }, .{ .proj = "down_proj", .kt = width / 16, .nt = D / 16, .by_cols = false } }; - const short = [_][]const u8{ "g", "u", "d" }; - for (mats, short) |m, sh| { - const tshape = [_]i64{ @intCast(E), @intCast(m.kt), @intCast(m.nt), 32 }; - const k_all: usize = if (m.by_cols) D else 2048; - const n_all: usize = if (m.by_cols) 2048 else D; - const sus = [_]i64{ @intCast(E), @intCast(m.kt * 16) }; - const svs = [_]i64{ @intCast(E), @intCast(m.nt * 16) }; - const tn = try std.fmt.bufPrint(&nb, "{s}.{s}t", .{ out, sh }); - if (b.dry) { - try b.put(tn, .i32, &tshape, null); - try b.putf("{s}.suh_{s}", .{ out, sh }, .f16, &sus, null); - try b.putf("{s}.svh_{s}", .{ out, sh }, .f16, &svs, null); - continue; - } - const per = m.kt * m.nt * tile; - const words = try b.gpa.alignedAlloc(u8, .@"16", E * per); - defer b.gpa.free(words); - const suh = try b.gpa.alloc(u16, E * m.kt * 16); - defer b.gpa.free(suh); - const svh = try b.gpa.alloc(u16, E * m.nt * 16); - defer b.gpa.free(svh); - var f: [256]u8 = undefined; - for (0..E) |x| { - const tr = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.trellis", .{ li, x, m.proj })); - if (!std.mem.eql(u8, tr.dtype, "I16") or tr.shape.len != 3 or tr.shape[2] != 64 or tr.shape[0] != k_all / 16 or tr.shape[1] != n_all / 16) return error.NotExl3Mcg4; - if (!b.ck.has(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.mcg", .{ li, x, m.proj }))) return error.NotExl3Mcg4; - const dst = words[x * per ..][0..per]; - if (m.by_cols) { // every tile row, this rank's tile columns - const all = try b.ck.readAll(b.gpa, try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.trellis", .{ li, x, m.proj })); - defer b.gpa.free(all); - const row = (n_all / 16) * tile; - for (0..m.kt) |r| @memcpy(dst[r * m.nt * tile ..][0 .. m.nt * tile], all[r * row + b.rank * m.nt * tile ..][0 .. m.nt * tile]); - } else try b.ck.readAt(tr, b.rank * per, dst); // this rank's tile rows: one contiguous part - const se = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.suh", .{ li, x, m.proj })); - const ve = try b.ck.get(try fullName(&f, "layers.{d}.mlp.experts.{d}.{s}.svh", .{ li, x, m.proj })); - if (!std.mem.eql(u8, se.dtype, "F16") or !std.mem.eql(u8, ve.dtype, "F16")) return error.NotExl3Mcg4; - const su = std.mem.sliceAsBytes(suh[x * m.kt * 16 ..][0 .. m.kt * 16]); - const sv = std.mem.sliceAsBytes(svh[x * m.nt * 16 ..][0 .. m.nt * 16]); - try b.ck.readAt(se, if (m.by_cols) 0 else b.rank * su.len, su); - try b.ck.readAt(ve, if (m.by_cols) b.rank * sv.len else 0, sv); - if (m.by_cols) { // gate and up: one input rotation shared by every expert of the layer? (exl3_shx) - if (std.mem.eql(u8, sh, "g") and x == 0) @memcpy(b.suh_first[0..su.len], su) else if (!std.mem.eql(u8, su, b.suh_first[0..su.len])) b.exl3_shx = false; - } - } - try b.put(tn, .i32, &tshape, words); - try b.putf("{s}.suh_{s}", .{ out, sh }, .f16, &sus, std.mem.sliceAsBytes(suh)); - try b.putf("{s}.svh_{s}", .{ out, sh }, .f16, &svs, std.mem.sliceAsBytes(svh)); - } -} - -// --------------------------------------------------------------------------------------------------------------- -// The layers, as weights.py's load() builds them (names: its weight tree). +// -- the layers, as weights.py's load() builds them (names: its weight tree) ------------------------------------ fn hc(b: *Builder, out: []const u8, li: usize, site: []const u8) !void { var o: [128]u8 = undefined; @@ -675,7 +265,7 @@ fn dsa(b: *Builder, out: []const u8, i: usize) !void { try b.tensor(try std.fmt.bufPrint(&o, "{s}.q_norm", .{out}), try fullName(&f, "layers.{d}.self_attn.q_a_layernorm.weight", .{i}), false); try b.tensor(try std.fmt.bufPrint(&o, "{s}.kv_norm", .{out}), try fullName(&f, "layers.{d}.self_attn.kv_a_layernorm.weight", .{i}), false); try b.linear(try std.fmt.bufPrint(&o, "{s}.q_b", .{out}), try fullName(&f, "layers.{d}.self_attn.q_b_proj.weight", .{i})); - // kv_b per head: 256 key rows then 256 value rows -> absorb.wk, absorb.wv [HL, 256, 512] (the latent path's only use) + // kv_b per head: 256 key rows then 256 value rows -> absorb.wk, absorb.wv [HL, 256, 512] (the latent path's) { const p = try b.part(try fullName(&f, "layers.{d}.self_attn.kv_b_proj.weight", .{i}), .row); defer b.free(p); @@ -850,9 +440,6 @@ pub fn deviceBytes(io: std.Io, dir: []const u8) !u64 { return b.bytes; } -// --------------------------------------------------------------------------------------------------------------- -// Packed dense linears (ops.cu's glm_qgemv layout). - fn envBits3() bool { const v = std.c.getenv("TF_GLM_EXPERT_BITS") orelse return false; return std.mem.eql(u8, std.mem.span(v), "3"); @@ -869,71 +456,6 @@ fn envSimBits() u8 { return std.fmt.parseInt(u8, std.mem.span(v), 10) catch 0; } -/// Each row's group of 64 re-quantized to ``bits`` symmetric levels (-2^(bits-1) .. 2^(bits-1) - 1, the scale from a -/// search over clip ratios for the least squared error), written back as 4-bit codes + 8 with that scale (bias -8 s). -fn simulateBits(t: Triple, bits: u8) void { - const Work = struct { - fn run(tt: Triple, nb: u8, e0: usize, e1: usize) void { - const lo: f32 = -@as(f32, @floatFromInt(@as(u32, 1) << @intCast(nb - 1))); - const hi: f32 = -lo - 1; - const wpg = 8; // 4-bit words a group of 64 - for (e0..e1) |e| for (0..tt.n) |r| for (0..tt.kg) |g| { - const at = (e * tt.n + r) * tt.kg + g; - const words = tt.words[(e * tt.n + r) * tt.k8 + g * wpg ..][0..wpg]; - const s0 = bf16Value(tt.scales[at]); - var v: [64]f32 = undefined; - var amax: f32 = 0; - for (0..64) |i| { - const q: i32 = @intCast((words[i / 8] >> @intCast(4 * (i % 8))) & 15); - v[i] = @as(f32, @floatFromInt(q - 8)) * s0; - amax = @max(amax, @abs(v[i])); - } - if (amax == 0) continue; - var best_s: f32 = amax / -lo; - var best_e: f32 = std.math.inf(f32); - for ([_]f32{ 1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6 }) |c| { - const sc: f32 = bf16Value(toBf16(c * amax / -lo)); - if (sc == 0) continue; - var err: f32 = 0; - for (v) |x| { - const q = @max(lo, @min(hi, @round(x / sc))); - err += (q * sc - x) * (q * sc - x); - } - if (err < best_e) { - best_e = err; - best_s = sc; - } - } - for (words) |*w| w.* = 0; - for (v, 0..) |x, i| { - const q: i32 = @intFromFloat(@max(lo, @min(hi, @round(x / best_s)))); - words[i / 8] |= @as(u32, @intCast(q + 8)) << @intCast(4 * (i % 8)); - } - tt.scales[at] = toBf16(best_s); - tt.biases[at] = toBf16(-8 * best_s); - }; - } - }; - const threads = 16; - var pool: [threads]?std.Thread = @splat(null); - const per = (t.e + threads - 1) / threads; - for (0..threads) |i| { - const e0 = i * per; - const e1 = @min(t.e, e0 + per); - if (e0 >= e1) break; - pool[i] = std.Thread.spawn(.{}, Work.run, .{ t, bits, e0, e1 }) catch blk: { - Work.run(t, bits, e0, e1); - break :blk null; - }; - } - for (pool) |th| if (th) |x| x.join(); -} - -fn toBf16(x: f32) u16 { - const u: u32 = @bitCast(x); - return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); // round to nearest even -} - fn envSym() bool { const v = std.c.getenv("TF_GLM_EXPERTS_SYM") orelse return true; return !std.mem.eql(u8, std.mem.span(v), "0"); @@ -946,95 +468,3 @@ fn envBits(name: [*:0]const u8) !u8 { if (bits != 0 and bits != 4 and bits != 6 and bits != 8) return error.BadPackedBits; return bits; } - -fn bf16Value(h: u16) f32 { - return @bitCast(@as(u32, h) << 16); -} - -/// x rounded to the nearest integer, halves to even. -fn roundEven(x: f32) f32 { - 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; -} - -/// bf16 rows [n, k] -> symmetric ``bits``-bit values: each group of 64 inputs a scale absmax / (2^(bits-1) - 1) -/// rounded to fp16, values round(w / scale) (halves to even) clamped to +-(2^(bits-1) - 1), stored + 2^(bits-1) in -/// 16-value chunks of bits / 2 words (``words`` [n * k * bits / 32]); ``scales`` [n * k / 64] fp16 bits. -pub fn quantize(bits: u8, src: []const u16, n: usize, k: usize, words: []u32, scales: []u16) void { - const half: i32 = @as(i32, 1) << @intCast(bits - 1); - const qmax: f32 = @floatFromInt(half - 1); - const wpc: usize = bits / 2; - const groups = k / 64; - for (0..n) |r| { - const row = src[r * k ..][0..k]; - for (0..groups) |g| { - var amax: f32 = 0; - for (row[g * 64 ..][0..64]) |h| amax = @max(amax, @abs(bf16Value(h))); - const s16: f16 = @floatCast(amax / qmax); - scales[r * groups + g] = @bitCast(s16); - const sc: f32 = s16; - for (0..4) |c| { - var q: [16]u32 = undefined; - for (0..16) |i| { - const v = bf16Value(row[g * 64 + c * 16 + i]); - const t: i32 = if (sc == 0) 0 else @intFromFloat(@max(-qmax, @min(qmax, roundEven(v / sc)))); - q[i] = @intCast(t + half); - } - pack(bits, &q, words[((r * groups + g) * 4 + c) * wpc ..][0..wpc]); - } - } - } -} - -fn pack(bits: u8, q: *const [16]u32, out: []u32) void { - @memset(out, 0); - switch (bits) { - 4 => for (0..16) |i| { - out[i / 8] |= q[i] << @intCast(4 * (i % 8)); - }, - 8 => for (0..16) |i| { - out[i / 4] |= q[i] << @intCast(8 * (i % 4)); - }, - else => { - for (0..15) |i| out[i / 5] |= q[i] << @intCast(6 * (i % 5)); - for (0..3) |j| out[j] |= ((q[15] >> @intCast(2 * j)) & 3) << 30; - }, - } -} - -test "packed rows: values and scales round trip (4, 6, 8 bits)" { - const gpa = std.testing.allocator; - var prng = std.Random.DefaultPrng.init(7); - const n = 3; - const k = 128; - var src: [n * k]u16 = undefined; - for (&src) |*h| h.* = @truncate(@as(u32, @bitCast(prng.random().floatNorm(f32) * 0.02)) >> 16); - src[5] = 0; // a zero inside a group - for ([_]u8{ 4, 6, 8 }) |bits| { - const words = try gpa.alloc(u32, n * k * bits / 32); - defer gpa.free(words); - var scales: [n * k / 64]u16 = undefined; - quantize(bits, &src, n, k, words, &scales); - const half: i32 = @as(i32, 1) << @intCast(bits - 1); - const qmax: f32 = @floatFromInt(half - 1); - const wpc: usize = bits / 2; - for (0..n * k / 16) |c| { - const p = words[c * wpc ..][0..wpc]; - for (0..16) |i| { - const raw: u32 = switch (bits) { - 4 => (p[i / 8] >> @intCast(4 * (i % 8))) & 15, - 8 => (p[i / 4] >> @intCast(8 * (i % 4))) & 255, - else => if (i < 15) (p[i / 5] >> @intCast(6 * (i % 5))) & 63 else (p[0] >> 30) | ((p[1] >> 30) << 2) | ((p[2] >> 30) << 4), - }; - const sc: f32 = @as(f16, @bitCast(scales[c / 4])); - const got = @as(f32, @floatFromInt(@as(i32, @intCast(raw)) - half)) * sc; - const want = bf16Value(src[c * 16 + i]); - try std.testing.expect(@abs(got - want) <= sc * 0.5 + 1e-7); - try std.testing.expect(@abs(@as(f32, @floatFromInt(@as(i32, @intCast(raw)) - half))) <= qmax); - } - } - } -} diff --git a/zig/tests/cuda/glm/gen_copies.py b/zig/tests/cuda/glm/gen_copies.py index ebc386b62e..e821af0a34 100644 --- a/zig/tests/cuda/glm/gen_copies.py +++ b/zig/tests/cuda/glm/gen_copies.py @@ -1,10 +1,5 @@ #!/usr/bin/env python3 -"""The device code of the Python GLM-5.3-Flash engine's own CUDA extension (families/glm5_next/cuda/kda.cu) as a -zig/kernels/cuda/glm copy: Flash Next's gen_copies.py rules (the source up to its first ATen-using function, ATen -includes dropped, the anonymous namespace kept) plus explicit instantiations of every template kernel the Python -build instantiated (demangled from its cubin). The grouped experts and the 4-bit matmuls are the shared extensions -Flash Next's copies already hold (the same symbols). -usage: gen_copies.py """ +"""The Python GLM engine's KDA extension device code as a zig/kernels/cuda/glm copy: .""" import json, sys from pathlib import Path diff --git a/zig/tests/cuda/glm_forward.zig b/zig/tests/cuda/glm_forward.zig index b9e1cafc96..2df13389b1 100644 --- a/zig/tests/cuda/glm_forward.zig +++ b/zig/tests/cuda/glm_forward.zig @@ -1,9 +1,4 @@ -//! GLM-5.3-Flash on CUDA: the family's kernels against float64 references and the row-invariance contract (a row's -//! bits alone equal its bits in any window, any mix of streams), on one GPU; and the two-rank engine's contract on -//! the real checkpoint: a stream's greedy reply is the same alone, beside 15 other streams, drafted, and resumed from -//! a kept prompt state. -//! usage: tf-cuda-test glm-kcheck -//! tf-cuda-test glm-batch [tokens] [sync] +//! tf-cuda-test glm-kcheck (kernels vs float64, row invariance) and glm-batch (alone = batched = drafted = resumed). const std = @import("std"); const cuda = @import("cuda"); @@ -13,6 +8,9 @@ const kern = gw.kern; const fwd = gw.forward; const dec = gw.decode; const Gpu = check.Gpu; +const run = @import("glm_forward_run.zig"); +pub const teacher = run.teacher; +pub const prof = run.prof; fn bf(x: u16) f32 { return @bitCast(@as(u32, x) << 16); @@ -27,7 +25,7 @@ fn randBf(rnd: std.Random, out: []u16, scale: f32) void { for (out) |*h| h.* = tobf(rnd.floatNorm(f32) * scale); } -fn ms(io: std.Io, t0: std.Io.Timestamp) f64 { +pub fn ms(io: std.Io, t0: std.Io.Timestamp) f64 { return @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(io, .awake)).toNanoseconds())) / 1e6; } @@ -55,7 +53,7 @@ const Stats = struct { } }; -// -- glm-kcheck --------------------------------------------------------------------------------------------------------- +// -- glm-kcheck ------------------------------------------------------------------------------------------------ fn checkMm(gpu: Gpu, k: *kern.K, rnd: std.Random) !void { const gpa = gpu.gpa; const d = gpu.d; @@ -145,8 +143,7 @@ fn checkMm(gpu: Gpu, k: *kern.K, rnd: std.Random) !void { } } -/// Latent attention: rows of two streams (each its own cache), dense and sparse, against float64; each row alone -/// against the same row in the mixed window. +/// Latent attention: two streams' rows, dense and sparse, against float64; each row alone equals it in the window. fn checkAttention(gpu: Gpu, k: *kern.K, rnd: std.Random) !void { const gpa = gpu.gpa; const d = gpu.d; @@ -272,8 +269,7 @@ fn checkAttention(gpu: Gpu, k: *kern.K, rnd: std.Random) !void { try check.expect(st.rel() < 2e-2 and same, "latent attention", .{}); } -/// The indexer: pool scores against float64, the 512 best pools (ties to the lower pool) against a host selection -/// over the kernel's own scores, for rows of two streams. +/// The indexer: pool scores against float64, the 512 best pools against a host pick over the kernel's scores. fn checkIndex(gpu: Gpu, k: *kern.K, rnd: std.Random) !void { const gpa = gpu.gpa; const d = gpu.d; @@ -360,11 +356,10 @@ pub fn kcheck(gpu: Gpu, _: []const [:0]const u8) !void { std.debug.print("glm-kcheck: all passed\n", .{}); } -// -- glm-batch ---------------------------------------------------------------------------------------------------------- +// -- glm-batch ------------------------------------------------------------------------------------------------- const Stream = struct { s: *fwd.Seq, out: std.ArrayList(u32) = .empty, pending: u32 = 0, held: [7]u32 = undefined, nheld: usize = 0, done: bool = false }; -/// Greedy decode of every stream to ``n`` tokens, together: windows of the pending token and up to ``depth`` drafts, -/// the longest agreeing prefix kept, the head drafting from the kept rows (the lane core's rule, without its policy). +/// Greedy decode of every stream to ``n`` tokens: pending token plus ``depth`` drafts, the agreeing prefix kept. fn decodeAll(e: *fwd.Engine, ss: []Stream, n: usize, depth: usize) !void { var picks: [fwd.STEP_ROWS]dec.Pick = undefined; var wins: [fwd.MAX_STREAMS]dec.Win = undefined; @@ -513,196 +508,3 @@ pub fn batch(gpu: Gpu, args: []const [:0]const u8) !void { std.debug.print("rank {d}: glm-batch passed\n", .{rank}); std.process.exit(0); // a test binary: no teardown (NCCL's service threads under the library's unload hang it) } - -// -- glm-tf ------------------------------------------------------------------------------------------------------------- -/// Teacher-forced against the Python engine's greedy reply (capture_ref.py's reference.json: prompt, first, serial): -/// the prompt, then the reply's tokens in windows of 8; how often this engine's top token is Python's next token. -/// usage: tf-cuda-test glm-tf -pub fn teacher(gpu: Gpu, args: []const [:0]const u8) !void { - const gpa = gpu.gpa; - const io = gpu.io; - if (args.len < 5) return error.Usage; - const rank = try std.fmt.parseInt(u8, args[0], 10); - const port = try std.fmt.parseInt(u16, args[2], 10); - var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); - defer link.close(); - var nccl = try cuda.nccl.Library.open(); - defer nccl.close(); - const comm = try cuda.tp_link.communicator(&nccl, link); - // the communicator's threads stop before the library unloads (an unloaded NCCL under live threads crashes the exit) - defer _ = nccl.api.ncclCommAbort(comm); - var stream = try cuda.Stream.init(gpu.d, true); - defer stream.deinit(); - { - const warm = try cuda.DeviceBuffer.alloc(gpu.d, 4096); - try nccl.check(nccl.api.ncclAllGather(warm.ptr, warm.ptr + 2048, 256, .i32, comm, stream.handle), "warm gather"); - try stream.synchronize(); - } - var kernels = try gw.api.Kernels.load(gpu.d); - var ctx: gw.api.Ctx = .{ .d = gpu.d, .ctx = gpu.ctx, .stream = stream, .nccl = &nccl, .comm = comm, .rank = rank }; - var store = try gw.load(gpa, io, &ctx, args[3], rank); - defer store.deinit(); - const e = try fwd.init(gpa, io, &ctx, &kernels, &store, .{ .context = 65536 }); - const text = try std.Io.Dir.cwd().readFileAlloc(io, args[4], gpa, .limited(1 << 26)); - const ref = (try std.json.parseFromSlice(std.json.Value, gpa, text, .{})).value.object; - const ints = struct { - fn f(a: std.mem.Allocator, v: std.json.Value) ![]u32 { - const it = v.array.items; - const out = try a.alloc(u32, it.len); - for (it, 0..) |x, i| out[i] = @intCast(x.integer); - return out; - } - }.f; - const prompt = try ints(gpa, ref.get("prompt").?); - const serial = try ints(gpa, ref.get("serial").?); - const s = try fwd.newSeq(e, 8192); - const first = try dec.prefill(e, s, prompt, 0, 0); - var hits: usize = @intFromBool(first == serial[0]); - var picks: [8]dec.Pick = undefined; - var j: usize = 0; - while (j + 1 < serial.len) { - const w = @min(8, serial.len - 1 - j); - try dec.verify(e, &.{.{ .s = s, .ids = serial[j .. j + w] }}, &picks); - try dec.keep(e, s, @intCast(w), @intCast(w)); - for (0..w) |r| hits += @intFromBool(picks[r].tok == serial[j + r + 1]); - j += w; - } - std.debug.print("rank {d}: teacher-forced on {d} prompt tokens: the top token equals Python's next token at {d} of {d} positions\n", .{ rank, prompt.len, hits, serial.len }); - std.process.exit(0); // as glm-batch -} - -// -- glm-prof ----------------------------------------------------------------------------------------------------------- -/// One stream's verify windows (1 and 4 rows), eager with every launch timed: where a step's time goes (proportions: -/// each launch synchronized), and the graph-replayed step's time for comparison. -/// usage: tf-cuda-test glm-prof [streams] -pub fn prof(gpu: Gpu, args: []const [:0]const u8) !void { - const gpa = gpu.gpa; - const io = gpu.io; - if (args.len < 4) return error.Usage; - const rank = try std.fmt.parseInt(u8, args[0], 10); - const port = try std.fmt.parseInt(u16, args[2], 10); - const nstreams: usize = if (args.len > 4) try std.fmt.parseInt(usize, args[4], 10) else 1; - const graphs_only = args.len > 5 and std.mem.eql(u8, args[5], "graphs"); // for nsys: 50 one-row steps, no eager pass - var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); - defer link.close(); - var nccl = try cuda.nccl.Library.open(); - defer nccl.close(); - const comm = try cuda.tp_link.communicator(&nccl, link); - // the communicator's threads stop before the library unloads (an unloaded NCCL under live threads crashes the exit) - defer _ = nccl.api.ncclCommAbort(comm); - var stream = try cuda.Stream.init(gpu.d, true); - defer stream.deinit(); - { - const warm = try cuda.DeviceBuffer.alloc(gpu.d, 4096); - try nccl.check(nccl.api.ncclAllGather(warm.ptr, warm.ptr + 2048, 256, .i32, comm, stream.handle), "warm gather"); - try stream.synchronize(); - } - var kernels = try gw.api.Kernels.load(gpu.d); - var ctx: gw.api.Ctx = .{ .d = gpu.d, .ctx = gpu.ctx, .stream = stream, .nccl = &nccl, .comm = comm, .rank = rank }; - const no_roce = if (std.c.getenv("TF_ROCE")) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; - ctx.roce = try cuda.roce.Gather.init(gpa, io, gpu.d, link, !no_roce, 512 << 10); - defer if (ctx.roce) |g| g.deinit(); - var store = try gw.load(gpa, io, &ctx, args[3], rank); - defer store.deinit(); - const e = try fwd.init(gpa, io, &ctx, &kernels, &store, .{ .context = 65536 }); - if (std.c.getenv("TF_PROF_PROMPT")) |v| { // a prompt of this many tokens: its wall time, then its kernels (eager chunks) - const n = try std.fmt.parseInt(usize, std.mem.span(v), 10); - const p = try gpa.alloc(u32, n); - defer gpa.free(p); - for (p, 0..) |*t, i| t.* = @intCast(1000 + (i * 7919) % 60000); - const s0 = try fwd.newSeq(e, @intCast(n + 1024)); - const t0 = std.Io.Timestamp.now(io, .awake); - _ = try dec.prefill(e, s0, p, 0, 0); - const wall = ms(io, t0); - if (rank == 0) std.debug.print("prompt of {d} tokens: {d:.0} ms, {d:.0} tok/s\n", .{ n, wall, @as(f64, @floatFromInt(n)) / wall * 1000 }); - var m = std.StringHashMap(f64).init(gpa); - e.k.prof = &m; - e.k.io = io; - const s1 = try fwd.newSeq(e, @intCast(n + 1024)); - _ = try dec.prefill(e, s1, p, 0, 0); - e.k.prof = null; - var total: f64 = 0; - var it = m.iterator(); - const Row = struct { name: []const u8, ms: f64 }; - var rows: std.ArrayList(Row) = .empty; - while (it.next()) |x| { - total += x.value_ptr.*; - try rows.append(gpa, .{ .name = x.key_ptr.*, .ms = x.value_ptr.* }); - } - std.mem.sort(Row, rows.items, {}, struct { - fn lt(_: void, a: Row, b: Row) bool { - return a.ms > b.ms; - } - }.lt); - if (rank == 0) { - std.debug.print("prompt of {d} tokens: eager launches sum {d:.0} ms (synchronized: proportions only)\n", .{ n, total }); - for (rows.items[0..@min(rows.items.len, 30)]) |x| std.debug.print(" {d:8.1} ms {d:5.1}% {s}\n", .{ x.ms, 100 * x.ms / total, x.name[0..@min(x.name.len, 90)] }); - } - return; - } - var seqs: [fwd.MAX_STREAMS]*fwd.Seq = undefined; - var prompt: [512]u32 = undefined; - for (&prompt, 0..) |*t, i| t.* = @intCast(1000 + (i * 7919) % 60000); - var tok: [fwd.MAX_STREAMS]u32 = undefined; - for (0..nstreams) |i| { - seqs[i] = try fwd.newSeq(e, 8192); - prompt[0] = @intCast(2000 + i); - tok[i] = try dec.prefill(e, seqs[i], &prompt, 0, 0); - } - var picks: [fwd.STEP_ROWS]dec.Pick = undefined; - var ids: [fwd.MAX_STREAMS][8]u32 = undefined; - var wins: [fwd.MAX_STREAMS]dec.Win = undefined; - for ([_]usize{ 1, 4, 8 }) |width| { - if (graphs_only and width != 1) break; - const step = struct { - fn f(ee: *fwd.Engine, ss: []*fwd.Seq, tk: []u32, idv: *[fwd.MAX_STREAMS][8]u32, wv: *[fwd.MAX_STREAMS]dec.Win, pk: []dec.Pick, w: usize) !void { - for (ss, 0..) |s, i| { - for (0..w) |r| idv[i][r] = if (r == 0) tk[i] else @intCast(3000 + r * 101 + i); - wv[i] = .{ .s = s, .ids = idv[i][0..w] }; - } - try dec.verify(ee, wv[0..ss.len], pk); - for (ss, 0..) |s, i| { - try dec.keep(ee, s, @intCast(w), 1); - tk[i] = pk[i * w].tok; - } - } - }.f; - for (0..8) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); // graphs captured, clocks up - var best: f64 = std.math.inf(f64); - for (0..10) |_| { - const t0 = std.Io.Timestamp.now(io, .awake); - try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); - best = @min(best, ms(io, t0)); - } - if (graphs_only) { - for (0..40) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); - if (rank == 0) std.debug.print("{d} stream(s) x {d} rows: graph step {d:.2} ms (graphs only)\n", .{ nstreams, width, best }); - continue; - } - var m = std.StringHashMap(f64).init(gpa); - e.k.prof = &m; - e.k.io = io; - e.use_graphs = false; - for (0..4) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); - e.k.prof = null; - e.use_graphs = true; - var total: f64 = 0; - var it = m.iterator(); - const Row = struct { name: []const u8, ms: f64 }; - var rows: std.ArrayList(Row) = .empty; - while (it.next()) |x| { - total += x.value_ptr.* / 4; - try rows.append(gpa, .{ .name = x.key_ptr.*, .ms = x.value_ptr.* / 4 }); - } - std.mem.sort(Row, rows.items, {}, struct { - fn lt(_: void, a: Row, b: Row) bool { - return a.ms > b.ms; - } - }.lt); - if (rank == 0) { - std.debug.print("{d} stream(s) x {d} rows: graph step {d:.2} ms; eager launches sum {d:.2} ms\n", .{ nstreams, width, best, total }); - for (rows.items[0..@min(rows.items.len, 22)]) |x| std.debug.print(" {d:7.3} ms {d:5.1}% {s}\n", .{ x.ms, 100 * x.ms / total, x.name[0..@min(x.name.len, 90)] }); - } - } - std.process.exit(0); // as glm-batch -} diff --git a/zig/tests/cuda/glm_forward_run.zig b/zig/tests/cuda/glm_forward_run.zig new file mode 100644 index 0000000000..78c2d1ec7d --- /dev/null +++ b/zig/tests/cuda/glm_forward_run.zig @@ -0,0 +1,196 @@ +//! tf-cuda-test glm-tf and glm-prof: teacher-forced agreement with the Python engine, and a step's time by launch. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const gw = @import("glm_weights"); +const fwd = gw.forward; +const dec = gw.decode; +const Gpu = check.Gpu; +const ms = @import("glm_forward.zig").ms; + +/// glm-tf : top-token agreement with Python's greedy reply. +pub fn teacher(gpu: Gpu, args: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const io = gpu.io; + if (args.len < 5) return error.Usage; + const rank = try std.fmt.parseInt(u8, args[0], 10); + const port = try std.fmt.parseInt(u16, args[2], 10); + var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); + defer link.close(); + var nccl = try cuda.nccl.Library.open(); + defer nccl.close(); + const comm = try cuda.tp_link.communicator(&nccl, link); + // the communicator's threads stop before the library unloads (an unloaded NCCL under live threads crashes the exit) + defer _ = nccl.api.ncclCommAbort(comm); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + { + const warm = try cuda.DeviceBuffer.alloc(gpu.d, 4096); + try nccl.check(nccl.api.ncclAllGather(warm.ptr, warm.ptr + 2048, 256, .i32, comm, stream.handle), "warm gather"); + try stream.synchronize(); + } + var kernels = try gw.api.Kernels.load(gpu.d); + var ctx: gw.api.Ctx = .{ .d = gpu.d, .ctx = gpu.ctx, .stream = stream, .nccl = &nccl, .comm = comm, .rank = rank }; + var store = try gw.load(gpa, io, &ctx, args[3], rank); + defer store.deinit(); + const e = try fwd.init(gpa, io, &ctx, &kernels, &store, .{ .context = 65536 }); + const text = try std.Io.Dir.cwd().readFileAlloc(io, args[4], gpa, .limited(1 << 26)); + const ref = (try std.json.parseFromSlice(std.json.Value, gpa, text, .{})).value.object; + const ints = struct { + fn f(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const it = v.array.items; + const out = try a.alloc(u32, it.len); + for (it, 0..) |x, i| out[i] = @intCast(x.integer); + return out; + } + }.f; + const prompt = try ints(gpa, ref.get("prompt").?); + const serial = try ints(gpa, ref.get("serial").?); + const s = try fwd.newSeq(e, 8192); + const first = try dec.prefill(e, s, prompt, 0, 0); + var hits: usize = @intFromBool(first == serial[0]); + var picks: [8]dec.Pick = undefined; + var j: usize = 0; + while (j + 1 < serial.len) { + const w = @min(8, serial.len - 1 - j); + try dec.verify(e, &.{.{ .s = s, .ids = serial[j .. j + w] }}, &picks); + try dec.keep(e, s, @intCast(w), @intCast(w)); + for (0..w) |r| hits += @intFromBool(picks[r].tok == serial[j + r + 1]); + j += w; + } + std.debug.print("rank {d}: teacher-forced on {d} prompt tokens: the top token equals Python's next token at {d} of {d} positions\n", .{ rank, prompt.len, hits, serial.len }); + std.process.exit(0); // as glm-batch +} + +/// glm-prof [streams]: 1 and 4-row verify windows, each launch timed, and a graph. +pub fn prof(gpu: Gpu, args: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const io = gpu.io; + if (args.len < 4) return error.Usage; + const rank = try std.fmt.parseInt(u8, args[0], 10); + const port = try std.fmt.parseInt(u16, args[2], 10); + const nstreams: usize = if (args.len > 4) try std.fmt.parseInt(usize, args[4], 10) else 1; + const graphs_only = args.len > 5 and std.mem.eql(u8, args[5], "graphs"); // for nsys: 50 one-row steps, no eager pass + var link = if (rank == 0) try cuda.tp_link.Link.lead(port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(args[1], port), 600_000); + defer link.close(); + var nccl = try cuda.nccl.Library.open(); + defer nccl.close(); + const comm = try cuda.tp_link.communicator(&nccl, link); + // the communicator's threads stop before the library unloads (an unloaded NCCL under live threads crashes the exit) + defer _ = nccl.api.ncclCommAbort(comm); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + { + const warm = try cuda.DeviceBuffer.alloc(gpu.d, 4096); + try nccl.check(nccl.api.ncclAllGather(warm.ptr, warm.ptr + 2048, 256, .i32, comm, stream.handle), "warm gather"); + try stream.synchronize(); + } + var kernels = try gw.api.Kernels.load(gpu.d); + var ctx: gw.api.Ctx = .{ .d = gpu.d, .ctx = gpu.ctx, .stream = stream, .nccl = &nccl, .comm = comm, .rank = rank }; + const no_roce = if (std.c.getenv("TF_ROCE")) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; + ctx.roce = try cuda.roce.Gather.init(gpa, io, gpu.d, link, !no_roce, 512 << 10); + defer if (ctx.roce) |g| g.deinit(); + var store = try gw.load(gpa, io, &ctx, args[3], rank); + defer store.deinit(); + const e = try fwd.init(gpa, io, &ctx, &kernels, &store, .{ .context = 65536 }); + if (std.c.getenv("TF_PROF_PROMPT")) |v| { // a prompt of this many tokens: its wall time, then its kernels (eager chunks) + const n = try std.fmt.parseInt(usize, std.mem.span(v), 10); + const p = try gpa.alloc(u32, n); + defer gpa.free(p); + for (p, 0..) |*t, i| t.* = @intCast(1000 + (i * 7919) % 60000); + const s0 = try fwd.newSeq(e, @intCast(n + 1024)); + const t0 = std.Io.Timestamp.now(io, .awake); + _ = try dec.prefill(e, s0, p, 0, 0); + const wall = ms(io, t0); + if (rank == 0) std.debug.print("prompt of {d} tokens: {d:.0} ms, {d:.0} tok/s\n", .{ n, wall, @as(f64, @floatFromInt(n)) / wall * 1000 }); + var m = std.StringHashMap(f64).init(gpa); + e.k.prof = &m; + e.k.io = io; + const s1 = try fwd.newSeq(e, @intCast(n + 1024)); + _ = try dec.prefill(e, s1, p, 0, 0); + e.k.prof = null; + var total: f64 = 0; + var it = m.iterator(); + const Row = struct { name: []const u8, ms: f64 }; + var rows: std.ArrayList(Row) = .empty; + while (it.next()) |x| { + total += x.value_ptr.*; + try rows.append(gpa, .{ .name = x.key_ptr.*, .ms = x.value_ptr.* }); + } + std.mem.sort(Row, rows.items, {}, struct { + fn lt(_: void, a: Row, b: Row) bool { + return a.ms > b.ms; + } + }.lt); + if (rank == 0) { + std.debug.print("prompt of {d} tokens: eager launches sum {d:.0} ms (synchronized: proportions only)\n", .{ n, total }); + for (rows.items[0..@min(rows.items.len, 30)]) |x| std.debug.print(" {d:8.1} ms {d:5.1}% {s}\n", .{ x.ms, 100 * x.ms / total, x.name[0..@min(x.name.len, 90)] }); + } + return; + } + var seqs: [fwd.MAX_STREAMS]*fwd.Seq = undefined; + var prompt: [512]u32 = undefined; + for (&prompt, 0..) |*t, i| t.* = @intCast(1000 + (i * 7919) % 60000); + var tok: [fwd.MAX_STREAMS]u32 = undefined; + for (0..nstreams) |i| { + seqs[i] = try fwd.newSeq(e, 8192); + prompt[0] = @intCast(2000 + i); + tok[i] = try dec.prefill(e, seqs[i], &prompt, 0, 0); + } + var picks: [fwd.STEP_ROWS]dec.Pick = undefined; + var ids: [fwd.MAX_STREAMS][8]u32 = undefined; + var wins: [fwd.MAX_STREAMS]dec.Win = undefined; + for ([_]usize{ 1, 4, 8 }) |width| { + if (graphs_only and width != 1) break; + const step = struct { + fn f(ee: *fwd.Engine, ss: []*fwd.Seq, tk: []u32, idv: *[fwd.MAX_STREAMS][8]u32, wv: *[fwd.MAX_STREAMS]dec.Win, pk: []dec.Pick, w: usize) !void { + for (ss, 0..) |s, i| { + for (0..w) |r| idv[i][r] = if (r == 0) tk[i] else @intCast(3000 + r * 101 + i); + wv[i] = .{ .s = s, .ids = idv[i][0..w] }; + } + try dec.verify(ee, wv[0..ss.len], pk); + for (ss, 0..) |s, i| { + try dec.keep(ee, s, @intCast(w), 1); + tk[i] = pk[i * w].tok; + } + } + }.f; + for (0..8) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); // graphs captured, clocks up + var best: f64 = std.math.inf(f64); + for (0..10) |_| { + const t0 = std.Io.Timestamp.now(io, .awake); + try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); + best = @min(best, ms(io, t0)); + } + if (graphs_only) { + for (0..40) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); + if (rank == 0) std.debug.print("{d} stream(s) x {d} rows: graph step {d:.2} ms (graphs only)\n", .{ nstreams, width, best }); + continue; + } + var m = std.StringHashMap(f64).init(gpa); + e.k.prof = &m; + e.k.io = io; + e.use_graphs = false; + for (0..4) |_| try step(e, seqs[0..nstreams], &tok, &ids, &wins, &picks, width); + e.k.prof = null; + e.use_graphs = true; + var total: f64 = 0; + var it = m.iterator(); + const Row = struct { name: []const u8, ms: f64 }; + var rows: std.ArrayList(Row) = .empty; + while (it.next()) |x| { + total += x.value_ptr.* / 4; + try rows.append(gpa, .{ .name = x.key_ptr.*, .ms = x.value_ptr.* / 4 }); + } + std.mem.sort(Row, rows.items, {}, struct { + fn lt(_: void, a: Row, b: Row) bool { + return a.ms > b.ms; + } + }.lt); + if (rank == 0) { + std.debug.print("{d} stream(s) x {d} rows: graph step {d:.2} ms; eager launches sum {d:.2} ms\n", .{ nstreams, width, best, total }); + for (rows.items[0..@min(rows.items.len, 22)]) |x| std.debug.print(" {d:7.3} ms {d:5.1}% {s}\n", .{ x.ms, 100 * x.ms / total, x.name[0..@min(x.name.len, 90)] }); + } + } + std.process.exit(0); // as glm-batch +} diff --git a/zig/tests/cuda/glm_weights.zig b/zig/tests/cuda/glm_weights.zig index 781ac43132..1257a2b4bc 100644 --- a/zig/tests/cuda/glm_weights.zig +++ b/zig/tests/cuda/glm_weights.zig @@ -1,6 +1,4 @@ -//! GLM-5.3-Flash on CUDA: the native weight loader (src/families/glm_cuda/weights.zig) against the Python engine's -//! prepared weights (tools/capture_ref.py's weights.json: each tensor's name, bytes and sha256). -//! usage: tf-cuda-test glm-weights +//! tf-cuda-test glm-weights : the native loader vs capture_ref.py's weights.json. const std = @import("std"); const cuda = @import("cuda"); 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 11/31] 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 0e5049866d40e5bffef2ddb384c11ce2fa3a5141 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 12/31] cuda: RoCE gather comments on one line each, the protocol notes in docs/roce-gather.md, zig fmt --- docs/roce-gather.md | 20 +++++++++ zig/kernels/cuda/roce_gather.cu | 16 +------ zig/src/cuda/roce.zig | 75 ++++++++++++++++++++------------- zig/src/cuda/tp_link.zig | 3 +- 4 files changed, 67 insertions(+), 47 deletions(-) create mode 100644 docs/roce-gather.md diff --git a/docs/roce-gather.md b/docs/roce-gather.md new file mode 100644 index 0000000000..7155542f11 --- /dev/null +++ b/docs/roce-gather.md @@ -0,0 +1,20 @@ +# RoCE all-gather for two ranks + +One-shot all-gather between two DGX Sparks over RoCE: the GPU side of b12x's "RoCEnante" exchange (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks. The host side (zig/src/cuda/roce.zig) owns a pinned region both GPUs and both NICs reach: + +``` + recv[slot] the peer's shard, written by the peer's NIC (2 slots of slot_bytes) + flag[slot][hca] the sequence the peer's NIC writes after its stripe on that HCA (128 bytes apart) + send[slot] this rank's shard, staged here for its NIC (2 slots) + ctrl {seq, nbytes of slot 0, nbytes of slot 1, error seq, missing hca}: the proxy's doorbell +``` + +One launch: (1) every block stages its share of the shard into send[seq & 1]; (2) the last block to stage publishes the slot's size and rings ctrl.seq after a system fence, and the proxy thread RDMA-writes the shard and then a flag per HCA stripe on the same queue pair (so a flag never lands before its stripe); (3) each block waits for every stripe's flag; (4) every block writes both shards in rank order (the layout NCCL's all-gather gives); (5) the last block to finish advances the device epoch, so a CUDA graph replays the exchange with new sequences. A wait past the spin limit poisons the exchange: later launches return at once and the host reports the error. + +## Host side + +Small all-gathers between two ranks as one-shot RDMA writes over RoCE: b12x's "RoCEnante" exchange (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks: the kernel (kernels/cuda/roce_gather.cu) stages this rank's shard in pinned host memory and rings a doorbell; a proxy thread here RDMA-writes it to the peer, striped over the HCAs, each stripe followed by its sequence flag on the same queue pair; the kernel waits for the peer's flags and copies both shards out. A decode step's gathers are a few rows of partials, where NCCL spends tens of microseconds a call in fixed costs; this path takes about ten. libibverbs is opened at run time (as the driver and NCCL are): without it, or without an active HCA, the caller keeps NCCL. + +## Control link + +Two ranks over the hosts' network: a TCP control link (rank 0 listens, rank 1 connects; framed messages carry requests and per-round decisions) and the NCCL communicator it brings up (rank 0 makes the unique id and sends it). diff --git a/zig/kernels/cuda/roce_gather.cu b/zig/kernels/cuda/roce_gather.cu index 8d6ff32e01..832afbcf9b 100644 --- a/zig/kernels/cuda/roce_gather.cu +++ b/zig/kernels/cuda/roce_gather.cu @@ -1,18 +1,4 @@ -// One-shot all-gather between two DGX Sparks over RoCE: the GPU side of b12x's "RoCEnante" exchange -// (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks. -// The host side (zig/src/cuda/roce.zig) owns a pinned region both GPUs and both NICs reach: -// -// recv[slot] the peer's shard, written by the peer's NIC (2 slots of slot_bytes) -// flag[slot][hca] the sequence the peer's NIC writes after its stripe on that HCA (128 bytes apart) -// send[slot] this rank's shard, staged here for its NIC (2 slots) -// ctrl {seq, nbytes of slot 0, nbytes of slot 1, error seq, missing hca}: the proxy's doorbell -// -// One launch: (1) every block stages its share of the shard into send[seq & 1]; (2) the last block to stage -// publishes the slot's size and rings ctrl.seq after a system fence, and the proxy thread RDMA-writes the shard and -// then a flag per HCA stripe on the same queue pair (so a flag never lands before its stripe); (3) each block waits -// for every stripe's flag; (4) every block writes both shards in rank order (the layout NCCL's all-gather gives); -// (5) the last block to finish advances the device epoch, so a CUDA graph replays the exchange with new sequences. -// A wait past the spin limit poisons the exchange: later launches return at once and the host reports the error. +// Two-rank one-shot RoCE all-gather after b12x (github.com/local-inference-lab/b12x, Apache-2.0): docs/roce-gather.md. #include diff --git a/zig/src/cuda/roce.zig b/zig/src/cuda/roce.zig index 93fd1e51f5..d4b1b3cc16 100644 --- a/zig/src/cuda/roce.zig +++ b/zig/src/cuda/roce.zig @@ -1,11 +1,4 @@ -//! Small all-gathers between two ranks as one-shot RDMA writes over RoCE: b12x's "RoCEnante" exchange -//! (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks: the -//! kernel (kernels/cuda/roce_gather.cu) stages this rank's shard in pinned host memory and rings a doorbell; a proxy -//! thread here RDMA-writes it to the peer, striped over the HCAs, each stripe followed by its sequence flag on the -//! same queue pair; the kernel waits for the peer's flags and copies both shards out. A decode step's gathers are a -//! few rows of partials, where NCCL spends tens of microseconds a call in fixed costs; this path takes about ten. -//! libibverbs is opened at run time (as the driver and NCCL are): without it, or without an active HCA, the caller -//! keeps NCCL. +//! Two-rank all-gathers as one-shot RoCE writes after b12x's RoCEnante (local-inference-lab/b12x, Apache-2.0). const std = @import("std"); const abi = @import("abi.zig"); @@ -40,16 +33,46 @@ const QpInitAttr = extern struct { qp_context: ?*anyopaque = null, send_cq: ?*Cq const GlobalRoute = extern struct { dgid: [16]u8 = @splat(0), flow_label: u32 = 0, sgid_index: u8 = 0, hop_limit: u8 = 0, traffic_class: u8 = 0 }; const AhAttr = extern struct { grh: GlobalRoute = .{}, dlid: u16 = 0, sl: u8 = 0, src_path_bits: u8 = 0, static_rate: u8 = 0, is_global: u8 = 0, port_num: u8 = 0 }; const QpAttr = extern struct { - qp_state: u32 = 0, cur_qp_state: u32 = 0, path_mtu: u32 = 0, path_mig_state: u32 = 0, qkey: u32 = 0, rq_psn: u32 = 0, - sq_psn: u32 = 0, dest_qp_num: u32 = 0, qp_access_flags: u32 = 0, cap: QpCap = .{}, ah_attr: AhAttr = .{}, alt_ah_attr: AhAttr = .{}, - pkey_index: u16 = 0, alt_pkey_index: u16 = 0, en_sqd_async_notify: u8 = 0, sq_draining: u8 = 0, max_rd_atomic: u8 = 0, - max_dest_rd_atomic: u8 = 0, min_rnr_timer: u8 = 0, port_num: u8 = 0, timeout: u8 = 0, retry_cnt: u8 = 0, rnr_retry: u8 = 0, - alt_port_num: u8 = 0, alt_timeout: u8 = 0, rate_limit: u32 = 0, _tail: u32 = 0, + qp_state: u32 = 0, + cur_qp_state: u32 = 0, + path_mtu: u32 = 0, + path_mig_state: u32 = 0, + qkey: u32 = 0, + rq_psn: u32 = 0, + sq_psn: u32 = 0, + dest_qp_num: u32 = 0, + qp_access_flags: u32 = 0, + cap: QpCap = .{}, + ah_attr: AhAttr = .{}, + alt_ah_attr: AhAttr = .{}, + pkey_index: u16 = 0, + alt_pkey_index: u16 = 0, + en_sqd_async_notify: u8 = 0, + sq_draining: u8 = 0, + max_rd_atomic: u8 = 0, + max_dest_rd_atomic: u8 = 0, + min_rnr_timer: u8 = 0, + port_num: u8 = 0, + timeout: u8 = 0, + retry_cnt: u8 = 0, + rnr_retry: u8 = 0, + alt_port_num: u8 = 0, + alt_timeout: u8 = 0, + rate_limit: u32 = 0, + _tail: u32 = 0, }; const Sge = extern struct { addr: u64, length: u32, lkey: u32 }; const SendWr = extern struct { - wr_id: u64 = 0, next: ?*SendWr = null, sg_list: ?*Sge = null, num_sge: c_int = 0, opcode: u32 = 0, send_flags: u32 = 0, - imm_data: u32 = 0, remote_addr: u64 = 0, rkey: u32 = 0, _tail: [128 - 52]u8 = @splat(0), + wr_id: u64 = 0, + next: ?*SendWr = null, + sg_list: ?*Sge = null, + num_sge: c_int = 0, + opcode: u32 = 0, + send_flags: u32 = 0, + imm_data: u32 = 0, + remote_addr: u64 = 0, + rkey: u32 = 0, + _tail: [128 - 52]u8 = @splat(0), }; const Wc = extern struct { wr_id: u64, status: u32, opcode: u32, vendor_err: u32, _rest: [48 - 20]u8 }; const PortAttr = extern struct { bytes: [56]u8 align(8) }; @@ -195,8 +218,7 @@ pub const Gather = struct { err: [256]u8 = @splat(0), gathers: u64 = 0, - /// Both ranks' exchange, or null on both when either rank cannot (no libibverbs, no active HCA, a failed - /// connection: each rank's verdict travels over ``link``, so the ranks agree). ``max_bytes``: a rank's shard at most. + /// Both ranks' exchange, or null on both when either cannot; verdicts cross ``link``. ``max_bytes``: a shard. pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const Driver, link: tp_link.Link, wanted: bool, max_bytes: usize) !?*Gather { const g = try gpa.create(Gather); var mine: Blob = std.mem.zeroes(Blob); @@ -239,8 +261,7 @@ pub const Gather = struct { } fn setup(g: *Gather, gpa: std.mem.Allocator, io: std.Io, d: *const Driver, rank: u8, max_bytes: usize) !void { - g.* = .{ .gpa = gpa, .d = d, .lib = undefined, .api = undefined, .rank = rank, .layout = .of(max_bytes), .region = undefined, - .region_dev = 0, .counters = undefined, .module = undefined, .f = undefined, .max_bytes = max_bytes }; + g.* = .{ .gpa = gpa, .d = d, .lib = undefined, .api = undefined, .rank = rank, .layout = .of(max_bytes), .region = undefined, .region_dev = 0, .counters = undefined, .module = undefined, .f = undefined, .max_bytes = max_bytes }; g.lib = std.DynLib.open("libibverbs.so.1") catch return error.VerbsUnavailable; errdefer g.lib.close(); const info = @typeInfo(Api).@"struct"; @@ -307,8 +328,7 @@ pub const Gather = struct { return error.HcaFailed; }; hca.cq = a.ibv_create_cq(hca.ctx.?, send_depth, null, null, 0) orelse return error.HcaFailed; - var attr: QpInitAttr = .{ .send_cq = hca.cq, .recv_cq = hca.cq, .qp_type = qpt_rc, - .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 } }; + var attr: QpInitAttr = .{ .send_cq = hca.cq, .recv_cq = hca.cq, .qp_type = qpt_rc, .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 } }; hca.qp = a.ibv_create_qp(hca.pd.?, &attr) orelse return error.HcaFailed; var init_attr: QpAttr = .{ .qp_state = qps_init, .port_num = port_num, .qp_access_flags = access_remote_write }; if (a.ibv_modify_qp(hca.qp.?, &init_attr, qp_state | qp_pkey_index | qp_port | qp_access_flags) != 0) return error.HcaFailed; @@ -334,9 +354,7 @@ pub const Gather = struct { for (0..g.n_hca) |h| { const x = &g.hcas[h]; g.peer_rkey[h] = peer.rkey[h]; - var rtr: QpAttr = .{ .qp_state = qps_rtr, .path_mtu = @min(peer.mtu[h], x.mtu), .dest_qp_num = peer.qpn[h], .max_dest_rd_atomic = 1, - .min_rnr_timer = 12, .ah_attr = .{ .is_global = 1, .dlid = peer.lid[h], .port_num = port_num, - .grh = .{ .dgid = peer.gid[h], .sgid_index = @intCast(x.gid_index), .hop_limit = 64 } } }; + var rtr: QpAttr = .{ .qp_state = qps_rtr, .path_mtu = @min(peer.mtu[h], x.mtu), .dest_qp_num = peer.qpn[h], .max_dest_rd_atomic = 1, .min_rnr_timer = 12, .ah_attr = .{ .is_global = 1, .dlid = peer.lid[h], .port_num = port_num, .grh = .{ .dgid = peer.gid[h], .sgid_index = @intCast(x.gid_index), .hop_limit = 64 } } }; if (g.api.ibv_modify_qp(x.qp.?, &rtr, qp_state | qp_av | qp_path_mtu | qp_dest_qpn | qp_rq_psn | qp_max_dest_rd_atomic | qp_min_rnr_timer) != 0) { g.fail("HCA {d}: the queue pair refused RTR", .{h}); return error.PeerRefused; @@ -391,11 +409,9 @@ pub const Gather = struct { if (h < packs % g.n_hca) n += 1; const bytes: u32 = n * 16; var flag_sge: Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; - var flag: SendWr = .{ .wr_id = h, .sg_list = &flag_sge, .num_sge = 1, .opcode = wr_rdma_write, .send_flags = send_signaled | send_inline, - .remote_addr = g.peer_addr + g.layout.flags + (slot * g.n_hca + h) * flag_stride, .rkey = g.peer_rkey[h] }; + var flag: SendWr = .{ .wr_id = h, .sg_list = &flag_sge, .num_sge = 1, .opcode = wr_rdma_write, .send_flags = send_signaled | send_inline, .remote_addr = g.peer_addr + g.layout.flags + (slot * g.n_hca + h) * flag_stride, .rkey = g.peer_rkey[h] }; var data_sge: Sge = .{ .addr = base + g.layout.send + slot * g.layout.slot_bytes + offset, .length = bytes, .lkey = x.mr.?.lkey }; - var data: SendWr = .{ .wr_id = h, .next = &flag, .sg_list = &data_sge, .num_sge = 1, .opcode = wr_rdma_write, - .remote_addr = g.peer_addr + g.layout.recv + slot * g.layout.slot_bytes + offset, .rkey = g.peer_rkey[h] }; + var data: SendWr = .{ .wr_id = h, .next = &flag, .sg_list = &data_sge, .num_sge = 1, .opcode = wr_rdma_write, .remote_addr = g.peer_addr + g.layout.recv + slot * g.layout.slot_bytes + offset, .rkey = g.peer_rkey[h] }; var bad: ?*SendWr = null; const first: *SendWr = if (bytes != 0) &data else &flag; const rc = opsCall(PostSend, x.ctx.?, ops_post_send)(x.qp.?, first, &bad); @@ -502,8 +518,7 @@ pub const Gather = struct { } }; -/// The HCAs to stripe over, in the order both ranks name them: TF_ROCE_HCA, else NCCL_IB_HCA ("=^" prefixes and -/// ":port" suffixes dropped), at most two. +/// The HCAs to stripe over, at most two: TF_ROCE_HCA, else NCCL_IB_HCA without "=^" prefixes or ":port" suffixes. fn hcaNames(io: std.Io, gpa: std.mem.Allocator, names: *[max_hcas][64]u8, lens: *[max_hcas]usize) !u32 { _ = io; _ = gpa; diff --git a/zig/src/cuda/tp_link.zig b/zig/src/cuda/tp_link.zig index 1d74b64c50..53c9839737 100644 --- a/zig/src/cuda/tp_link.zig +++ b/zig/src/cuda/tp_link.zig @@ -1,5 +1,4 @@ -//! Two ranks over the hosts' network: a TCP control link (rank 0 listens, rank 1 connects; framed messages carry -//! requests and per-round decisions) and the NCCL communicator it brings up (rank 0 makes the unique id and sends it). +//! Two ranks' TCP control link (rank 0 listens, rank 1 connects) and the NCCL communicator it brings up. const std = @import("std"); const posix = std.posix; 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 13/31] 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 14/31] 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 15/31] 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 d0692774291c99ed4ea1030cbb195ed9e91d2ce5 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 07:40:12 +0000 Subject: [PATCH 16/31] cuda: RoCE gather ranks agree one shard cap (each sends its max_bytes; fits tests the smaller on both), its host test in the cuda tests --- zig/src/cuda/roce.zig | 19 ++++++++++++++++++- zig/src/cuda/root.zig | 1 + 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/zig/src/cuda/roce.zig b/zig/src/cuda/roce.zig index d4b1b3cc16..41dc961de8 100644 --- a/zig/src/cuda/roce.zig +++ b/zig/src/cuda/roce.zig @@ -173,6 +173,7 @@ const Blob = extern struct { ok: u32, n_hca: u32, slot_bytes: u64, + max_bytes: u64, // the largest shard this rank would send this way (both ranks use the smaller) addr: u64, rkey: [max_hcas]u32, qpn: [max_hcas]u32, @@ -237,6 +238,7 @@ pub const Gather = struct { gpa.destroy(g); return null; } + g.max_bytes = agreedCap(mine.max_bytes, theirs.max_bytes); // one cap for fits on both ranks const connected = blk: { g.connect(&theirs) catch |e| { std.log.warn("RoCE all-gather: {s} ({s}); NCCL carries every gather", .{ @errorName(e), std.mem.sliceTo(&g.err, 0) }); @@ -252,7 +254,7 @@ pub const Gather = struct { gpa.destroy(g); return null; } - std.log.info("RoCE all-gather: {d} HCA(s), shards up to {d} KiB by RDMA writes", .{ g.n_hca, max_bytes >> 10 }); + std.log.info("RoCE all-gather: {d} HCA(s), shards up to {d} KiB by RDMA writes", .{ g.n_hca, g.max_bytes >> 10 }); return g; } @@ -338,6 +340,7 @@ pub const Gather = struct { b.ok = 1; b.n_hca = g.n_hca; b.slot_bytes = g.layout.slot_bytes; + b.max_bytes = g.max_bytes; b.addr = @intFromPtr(g.region.bytes.ptr); for (0..g.n_hca) |h| { const x = &g.hcas[h]; @@ -464,6 +467,11 @@ pub const Gather = struct { } } + /// The cap both ranks test in ``fits``: the smaller of the two, which each rank's slot holds. + fn agreedCap(mine: u64, theirs: u64) usize { + return @intCast(@min(mine, theirs)); + } + /// Whether a gather of ``bytes`` a rank (from and to 16-byte aligned buffers) goes this way. pub fn fits(g: *const Gather, bytes: usize, send: abi.DevicePtr, recv: abi.DevicePtr) bool { return bytes > 0 and bytes <= g.max_bytes and bytes % 16 == 0 and send % 16 == 0 and recv % 16 == 0; @@ -556,3 +564,12 @@ fn gidIndex(io: std.Io, gpa: std.mem.Allocator, name: []const u8) !c_int { } return error.NoHca; } + +test "RoCE gather: both ranks agree one cap" { + // 511 KiB and 512 KiB round to the same slot; a 512 KiB shard must take the same path on both ranks + try std.testing.expectEqual(Layout.of(511 << 10).slot_bytes, Layout.of(512 << 10).slot_bytes); + const a = Gather.agreedCap(511 << 10, 512 << 10); + const b = Gather.agreedCap(512 << 10, 511 << 10); + try std.testing.expectEqual(a, b); + try std.testing.expectEqual(@as(usize, 511 << 10), a); +} diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 40bd1bb6a0..6d3ebff3c7 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -47,6 +47,7 @@ test { _ = @import("embed_rows.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); + _ = @import("roce.zig"); _ = launch; _ = abi; _ = aot; From b07b11be88a8691d86b456629a48457ded14bb8b Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 08:08:48 +0000 Subject: [PATCH 17/31] glm cuda: the review's four fixes: a RoCE check after every sync that reads a gather (sampled draws, the full row, DFlash2), the ranks' settings agreed before the first collective, an image frame bound to its prefill by request id, a kept state allocated before rank 1 hears the keep --- zig/src/families/glm_cuda/api.zig | 6 +++ zig/src/families/glm_cuda/calibrate.zig | 1 + zig/src/families/glm_cuda/decode.zig | 4 +- zig/src/families/glm_cuda/dflash.zig | 1 + zig/src/families/glm_cuda/dflash_load.zig | 16 +++++++- zig/src/families/glm_cuda/native.zig | 25 +++++++----- zig/src/families/glm_cuda/native_env.zig | 42 +++++++++++++++++++++ zig/src/families/glm_cuda/native_follow.zig | 9 +++-- zig/src/families/glm_cuda/native_prompt.zig | 1 + 9 files changed, 89 insertions(+), 16 deletions(-) diff --git a/zig/src/families/glm_cuda/api.zig b/zig/src/families/glm_cuda/api.zig index 9f4886f371..4ade08a0d0 100644 --- a/zig/src/families/glm_cuda/api.zig +++ b/zig/src/families/glm_cuda/api.zig @@ -99,6 +99,12 @@ pub const Ctx = struct { /// Small all-gathers as RDMA writes over RoCE (cuda.roce), when both ranks brought it up; NCCL otherwise. roce: ?*cuda.roce.Gather = null, + /// After the sync that precedes a host read of anything a gather fed: a RoCE gather that timed out wrote nothing and + /// returned, so its error is raised here rather than the read taking a previous step's values. + pub fn gathered(c: *const Ctx) !void { + if (c.roce) |g| try g.check(); + } + /// Both ranks' ``count`` values from ``send`` into ``recv`` in rank order (ncclAllGather's layout and bits). pub fn allGather(c: *const Ctx, send: u64, recv: u64, count: usize, dt: cuda.nccl.DataType, what: []const u8) !void { const size: usize = switch (dt) { diff --git a/zig/src/families/glm_cuda/calibrate.zig b/zig/src/families/glm_cuda/calibrate.zig index 598ce4412c..375da3cc58 100644 --- a/zig/src/families/glm_cuda/calibrate.zig +++ b/zig/src/families/glm_cuda/calibrate.zig @@ -32,6 +32,7 @@ pub fn run(self: anytype, drafts: bool) !void { try proto.writeSampling(&w, null); try w.int(0); try w.int(0); + try w.int(0); // no image frame try self.send(.prefill, &w); seqs[i].sampling = null; toks[i] = try decode.prefill(self.e, seqs[i], &prompt, 0, 0); diff --git a/zig/src/families/glm_cuda/decode.zig b/zig/src/families/glm_cuda/decode.zig index 8f4d8d974d..ecd5eeee8e 100644 --- a/zig/src/families/glm_cuda/decode.zig +++ b/zig/src/families/glm_cuda/decode.zig @@ -26,7 +26,7 @@ fn readPicks(e: *Engine, b: *Buffers, rows: i64, out: []Pick) !void { const w: []i32 = @alignCast(std.mem.bytesAsSlice(i32, e.host[0 .. n * 4])); try d.check(d.api.cuMemcpyDtoHAsync_v2(@ptrCast(w.ptr), b.cgath, n * 4, e.k.stream.handle), "candidates"); try e.k.stream.synchronize(); - if (e.ctx.roce) |g| try g.check(); + try e.ctx.gathered(); const R: usize = @intCast(rows); for (0..R) |r| { const a = w[4 * r ..][0..4]; @@ -82,6 +82,7 @@ fn drawRows(e: *Engine, b: *Buffers, rows: i64, draws: []const Draw, out: []Pick } try e.ctx.allGather(b.tk, b.gtk, n, .i32, "candidates"); try k.stream.synchronize(); + try e.ctx.gathered(); try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(e.cand_host.ptr), b.gtk, 2 * n * 4), "candidates"); draw_stats.gpu_ms += msFrom(e, t_gpu); const t_host = std.Io.Timestamp.now(e.io, .awake); @@ -118,6 +119,7 @@ fn candTopk(e: *Engine, b: *Buffers, rows: i64, name: [:0]const u8) !void { fn drawFull(e: *Engine, row: u64, rule: lanes.Sampling, position: u64) !u32 { try e.ctx.nccl.check(e.ctx.nccl.api.ncclAllGather(row, e.full.ptr, @intCast(HEAD_N), .bf16, e.ctx.comm, e.k.stream.handle), "full row"); try e.k.stream.synchronize(); + try e.ctx.gathered(); try e.ctx.d.check(e.ctx.d.api.cuMemcpyDtoH_v2(@ptrCast(e.full_host.ptr), e.full.ptr, e.full_host.len * 2), "full row"); return draw_mod.drawFullBf16(e.gpa, e.full_host, rule, position); } diff --git a/zig/src/families/glm_cuda/dflash.zig b/zig/src/families/glm_cuda/dflash.zig index 765e2b950f..3fbb39b77a 100644 --- a/zig/src/families/glm_cuda/dflash.zig +++ b/zig/src/families/glm_cuda/dflash.zig @@ -363,6 +363,7 @@ pub fn draft(dr: *Drafter, e: *forward.Engine, all: []const Req) !void { try dr.graphs[S].?.launchOn(k.stream); } try k.stream.synchronize(); + try e.ctx.gathered(); // the candidates and projections below come through gathers try phase(dr, e, &tp, 2); if (!rank0 and dr.prof_fine and dr.calls % 256 == 255) logFine(dr, 256); // rank 1 logs its own table (rank 0's goes with the trace) if (rank0) { diff --git a/zig/src/families/glm_cuda/dflash_load.zig b/zig/src/families/glm_cuda/dflash_load.zig index cf2972141f..d88ed14f3f 100644 --- a/zig/src/families/glm_cuda/dflash_load.zig +++ b/zig/src/families/glm_cuda/dflash_load.zig @@ -180,7 +180,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *forward.Engine, dir: []const } // the device's share: linears packed (q|k|v, gate|up, k|v for context), norms and convs; codebooks stay host var total: u64 = 0; - const split = e.ctx.world == 2 and if (std.c.getenv("TF_GLM_DFLASH2_SPLIT")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; + const split = e.ctx.world == 2 and splitWanted(); const ff: usize = if (split) FF / 2 else FF; const small: u64 = 6 * 4096 * 2 + 2 * 128 * 2 + 2 * 2 * 2 * D * 2; total += LAYERS * (linearBytes(QKV, D, bits) + linearBytes(2 * KVW, D, bits) + linearBytes(D, D, bits) + linearBytes(2 * ff, D, bits) + @@ -260,7 +260,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, e: *forward.Engine, dir: []const dr.tap_s0 = std.mem.eql(u8, std.mem.span(v), "s0"); } if (std.c.getenv("TF_GLM_DFLASH2_DUMP_ROUND")) |v| dr.dump_at = std.fmt.parseInt(u32, std.mem.span(v), 10) catch 10; - if (std.c.getenv("TF_GLM_DFLASH2_HEAD")) |v| dr.full_head = std.mem.eql(u8, std.mem.span(v), "full"); + dr.full_head = fullHeadWanted(); if (std.c.getenv("TF_GLM_DFLASH2_TEMP")) |v| dr.temp = std.fmt.parseFloat(f64, std.mem.span(v)) catch 1.0; const inv: [MAX_M]f32 = @splat(1.0); try dr.scratch.upload(dr.invt - dr.scratch.ptr, std.mem.sliceAsBytes(&inv)); @@ -313,3 +313,15 @@ fn checkConfig(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { if (taps.len != TAPS.len) return error.DrafterConfigUnexpected; for (taps, TAPS) |t, w| if (t.integer != w) return error.DrafterConfigUnexpected; } + +/// TF_GLM_DFLASH2_SPLIT=0: each rank holds the drafter's whole MLP (else half each, on two ranks). +pub fn splitWanted() bool { + const v = std.c.getenv("TF_GLM_DFLASH2_SPLIT") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// TF_GLM_DFLASH2_HEAD=full: the drafter's tokens from the target's full head. +pub fn fullHeadWanted() bool { + const v = std.c.getenv("TF_GLM_DFLASH2_HEAD") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "full"); +} diff --git a/zig/src/families/glm_cuda/native.zig b/zig/src/families/glm_cuda/native.zig index 61ee1afaa1..52f32414c3 100644 --- a/zig/src/families/glm_cuda/native.zig +++ b/zig/src/families/glm_cuda/native.zig @@ -187,6 +187,7 @@ pub const Owned = struct { // rank 1: the next prefill's images (its image frame's), freed after it media_spans: std.ArrayList([2]i64) = .empty, media_feats: ?cuda.DeviceBuffer = null, + media_id: u64 = 0, // rank 1: the request the held image frame came with (its prefill names it; 0: none) guard: ?*Guard = null, // the memory and peer watchdog (guard.zig) copy_on: bool = true, // TF_GLM_COPY=0: the head's drafts only mix_on: bool = true, // TF_GLM_MIX=0: with DFlash2 loaded, it drafts every round (else each stream picks a source a round) @@ -317,7 +318,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ own.marginal = marginalFromEnv(); own.copy_on = if (std.c.getenv("TF_GLM_COPY")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; own.copy_rate = .{}; - own.mix_on = if (std.c.getenv("TF_GLM_MIX")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; + own.mix_on = native_env.mixWanted(); own.copy_min = envUsize("TF_GLM_COPY_MIN", copy.K); own.copy_min_reply = envUsize("TF_GLM_COPY_MIN_REPLY", copy.K_REPLY); own.trace = .{ .on = o.rank == 0 and std.c.getenv("TF_GLM_TRACE") != null }; @@ -326,6 +327,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ // the link and the communicator first, while memory is fresh (NCCL registers its buffers at its first use) own.link = if (o.rank == 0) try cuda.tp_link.Link.lead(o.master_port, 600_000) else try cuda.tp_link.Link.follow(io, try std.Io.net.IpAddress.parse(o.master orelse return error.NoMaster, o.master_port), 600_000); errdefer own.link.close(); + try native_env.agree(own.link, .fromEnv()); // before the first collective: a pair that would run apart refuses own.nccl = try cuda.nccl.Library.open(); errdefer own.nccl.close(); const comm = try cuda.tp_link.communicator(&own.nccl, own.link); @@ -350,6 +352,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ own.tower = null; own.media_spans = .empty; own.media_feats = null; + own.media_id = 0; own.image_token = try imageToken(gpa, io, dir); if (o.vision and o.rank == 0) own.tower = try vision.Tower.load(gpa, io, ctx.d, dir); errdefer if (own.tower) |*t| t.deinit(); @@ -571,19 +574,21 @@ fn releaseFn(p: *anyopaque, s: *lanes.Stream) void { _ = self.lanes_by.remove(s); if (l.copy) |*h| h.deinit(self.gpa); const seq = l.s; - const keep_it = seq.snap_at > 0 and !l.media and seq.snap_at <= l.prompt.len; + // the kept entry allocated before rank 1 hears: a keep it was told of is one rank 0 holds too + var ids: ?[]u32 = null; + if (seq.snap_at > 0 and !l.media and seq.snap_at <= l.prompt.len) { + ids = self.gpa.dupe(u32, l.prompt[0..seq.snap_at]) catch null; + if (ids) |x| self.kept.ensureUnusedCapacity(self.gpa, 1) catch { + self.gpa.free(x); + ids = null; + }; + } var w: Writer = .{ .gpa = self.gpa }; defer w.buf.deinit(self.gpa); w.int(seq.uid) catch {}; - w.int(@intFromBool(keep_it)) catch {}; + w.int(@intFromBool(ids != null)) catch {}; self.send(.release, &w) catch {}; - if (keep_it) { - const ids = self.gpa.dupe(u32, l.prompt[0..seq.snap_at]) catch null; - if (ids) |x| self.kept.append(self.gpa, .{ .s = seq, .ids = x }) catch { - self.gpa.free(x); - self.dropSeq(seq.uid); - } else self.dropSeq(seq.uid); - } else self.dropSeq(seq.uid); + if (ids) |x| self.kept.appendAssumeCapacity(.{ .s = seq, .ids = x }) else self.dropSeq(seq.uid); self.gpa.free(l.prompt); } diff --git a/zig/src/families/glm_cuda/native_env.zig b/zig/src/families/glm_cuda/native_env.zig index 0a6c50efb1..ff0d2d13d1 100644 --- a/zig/src/families/glm_cuda/native_env.zig +++ b/zig/src/families/glm_cuda/native_env.zig @@ -52,3 +52,45 @@ pub fn roceMaxBytes() usize { pub fn msSince(io: std.Io, t0: std.Io.Timestamp) f64 { return @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(io, .awake)).toNanoseconds())) / 1e6; } + +/// The settings that shape what both ranks run together: each rank reads its own environment, so they are compared +/// over the link before the first collective and a pair that differs does not start. +pub const Agreed = extern struct { + prefill_rows: i64, // TF_GLM_PREFILL_ROWS: a prompt chunk's rows + index_ring: i64, // TF_GLM_INDEX_RING: the index rings' rows (before the capacity caps them) + mix: u8, // TF_GLM_MIX + split: u8, // TF_GLM_DFLASH2_SPLIT + drafter: u8, // TF_GLM_DRAFTER=dflash2 + full_head: u8, // TF_GLM_DFLASH2_HEAD=full + _pad: [4]u8 = @splat(0), + + comptime { // the swap sends every byte: no padding left unset + std.debug.assert(@sizeOf(Agreed) == 24); + } + + pub fn fromEnv() Agreed { + return .{ + .prefill_rows = @import("forward.zig").prefillRows(), + .index_ring = @import("seq.zig").ringRows(std.math.maxInt(i64)), + .mix = @intFromBool(mixWanted()), + .split = @intFromBool(@import("dflash_load.zig").splitWanted()), + .drafter = @intFromBool(drafterWanted()), + .full_head = @intFromBool(@import("dflash_load.zig").fullHeadWanted()), + }; + } +}; + +/// Both ranks' Agreed over ``link``; error.SettingsDiffer on both when they differ (each logs both sides). +pub fn agree(link: anytype, mine: Agreed) !void { + var theirs: Agreed = undefined; + try link.swap(std.mem.asBytes(&mine), std.mem.asBytes(&theirs)); + if (std.meta.eql(mine, theirs)) return; + std.log.err("glm: the ranks' settings differ (this rank {any}, the other {any}); set TF_GLM_PREFILL_ROWS, TF_GLM_INDEX_RING, TF_GLM_MIX, TF_GLM_DFLASH2_SPLIT, TF_GLM_DRAFTER and TF_GLM_DFLASH2_HEAD alike on both", .{ mine, theirs }); + return error.SettingsDiffer; +} + +/// TF_GLM_MIX=0: with DFlash2 loaded, it drafts every round (else each stream picks a source a round). +pub fn mixWanted() bool { + const v = std.c.getenv("TF_GLM_MIX") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} diff --git a/zig/src/families/glm_cuda/native_follow.zig b/zig/src/families/glm_cuda/native_follow.zig index 0607abb13f..ed033392fd 100644 --- a/zig/src/families/glm_cuda/native_follow.zig +++ b/zig/src/families/glm_cuda/native_follow.zig @@ -33,18 +33,21 @@ pub fn followLoop(p: *anyopaque) anyerror!void { const sampling = try proto.readSampling(&r); const start: usize = @intCast(try r.int()); const keep_at: usize = @intCast(try r.int()); + const media_id = try r.int(); const seq = if (fresh > 0) try self.newSeq(uid, fresh) else self.seqs.get(uid) orelse return error.NoSuchSequence; seq.sampling = sampling; - defer { // the image frame's features served this prefill alone + defer { // a held image frame serves this prefill alone, or is dropped here (its own prefill never came) self.media_spans.clearRetainingCapacity(); if (self.media_feats) |*f| f.free(); self.media_feats = null; + self.media_id = 0; } - if (self.media_feats) |f| forward.attach(seq, .{ .spans = self.media_spans.items, .feats = f.ptr }); + if (media_id != 0 and media_id != self.media_id) return error.ImageFrameMissing; + if (media_id != 0) forward.attach(seq, .{ .spans = self.media_spans.items, .feats = self.media_feats.?.ptr }); _ = try decode.prefill(self.e, seq, toks, start, keep_at); }, .image => { - _ = try r.int(); + self.media_id = try r.int(); const n: usize = @intCast(try r.int()); self.media_spans.clearRetainingCapacity(); for (0..n) |_| try self.media_spans.append(self.gpa, .{ @intCast(try r.int()), @intCast(try r.int()) }); diff --git a/zig/src/families/glm_cuda/native_prompt.zig b/zig/src/families/glm_cuda/native_prompt.zig index d6a5223c2c..d3298eae3e 100644 --- a/zig/src/families/glm_cuda/native_prompt.zig +++ b/zig/src/families/glm_cuda/native_prompt.zig @@ -72,6 +72,7 @@ pub fn prefillFn(p: *anyopaque, s: *lanes.Stream) anyerror!void { try proto.writeSampling(&w, s.sampling); try w.int(start); try w.int(keep_at); + try w.int(if (media) id else 0); // the image frame this prefill takes (rank 1 drops one sent for another request) try self.send(.prefill, &w); seq.sampling = s.sampling; s.cached = @intCast(start); 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 18/31] 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 c525ea99185154d551ea6189ced1e18190a96e13 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 06:17:51 -0500 Subject: [PATCH 19/31] glm cuda: the gather check and the agreed settings' doc comments on one line each --- zig/src/families/glm_cuda/api.zig | 3 +-- zig/src/families/glm_cuda/native_env.zig | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/zig/src/families/glm_cuda/api.zig b/zig/src/families/glm_cuda/api.zig index 4ade08a0d0..79d5ab4005 100644 --- a/zig/src/families/glm_cuda/api.zig +++ b/zig/src/families/glm_cuda/api.zig @@ -99,8 +99,7 @@ pub const Ctx = struct { /// Small all-gathers as RDMA writes over RoCE (cuda.roce), when both ranks brought it up; NCCL otherwise. roce: ?*cuda.roce.Gather = null, - /// After the sync that precedes a host read of anything a gather fed: a RoCE gather that timed out wrote nothing and - /// returned, so its error is raised here rather than the read taking a previous step's values. + /// After a sync, before the host reads what a gather fed: a timed-out RoCE gather fails here, not as stale rows. pub fn gathered(c: *const Ctx) !void { if (c.roce) |g| try g.check(); } diff --git a/zig/src/families/glm_cuda/native_env.zig b/zig/src/families/glm_cuda/native_env.zig index ff0d2d13d1..6538814f4f 100644 --- a/zig/src/families/glm_cuda/native_env.zig +++ b/zig/src/families/glm_cuda/native_env.zig @@ -53,8 +53,7 @@ pub fn msSince(io: std.Io, t0: std.Io.Timestamp) f64 { return @as(f64, @floatFromInt(t0.durationTo(std.Io.Timestamp.now(io, .awake)).toNanoseconds())) / 1e6; } -/// The settings that shape what both ranks run together: each rank reads its own environment, so they are compared -/// over the link before the first collective and a pair that differs does not start. +/// What shapes both ranks' collectives, from each rank's environment: swapped before the first, refused if unequal. pub const Agreed = extern struct { prefill_rows: i64, // TF_GLM_PREFILL_ROWS: a prompt chunk's rows index_ring: i64, // TF_GLM_INDEX_RING: the index rings' rows (before the capacity caps them) From b8d45ded31edaacc80fb74f134159415c18ea847 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 10:27:08 +0000 Subject: [PATCH 20/31] cuda: the RoCE gather joins tp_group: its setup crosses the group's NCCL communicator (no TCP link: NCCL_COMM_ID and tp_watch do that now), Group.openRoce brings it up on two ranks, allGather takes it by size alone, Group.check after a sync that reads a gather; tf-cuda-test roce-gather (every size against NCCL's bytes, graph replays, time a call) --- docs/roce-gather.md | 10 +- zig/src/cuda/roce.zig | 53 ++++++--- zig/src/cuda/root.zig | 2 - zig/src/cuda/tp_group.zig | 18 ++- zig/src/cuda/tp_link.zig | 169 --------------------------- zig/tests/cuda/main.zig | 3 + zig/tests/cuda/roce_gather_tests.zig | 135 +++++++++++++++++++++ 7 files changed, 199 insertions(+), 191 deletions(-) delete mode 100644 zig/src/cuda/tp_link.zig create mode 100644 zig/tests/cuda/roce_gather_tests.zig diff --git a/docs/roce-gather.md b/docs/roce-gather.md index 7155542f11..e0ff5ae82c 100644 --- a/docs/roce-gather.md +++ b/docs/roce-gather.md @@ -13,8 +13,12 @@ One launch: (1) every block stages its share of the shard into send[seq & 1]; (2 ## Host side -Small all-gathers between two ranks as one-shot RDMA writes over RoCE: b12x's "RoCEnante" exchange (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks: the kernel (kernels/cuda/roce_gather.cu) stages this rank's shard in pinned host memory and rings a doorbell; a proxy thread here RDMA-writes it to the peer, striped over the HCAs, each stripe followed by its sequence flag on the same queue pair; the kernel waits for the peer's flags and copies both shards out. A decode step's gathers are a few rows of partials, where NCCL spends tens of microseconds a call in fixed costs; this path takes about ten. libibverbs is opened at run time (as the driver and NCCL are): without it, or without an active HCA, the caller keeps NCCL. +Small all-gathers between two ranks as one-shot RDMA writes over RoCE: b12x's "RoCEnante" exchange (https://github.com/local-inference-lab/b12x, b12x/comm/roce, Apache License 2.0), written for two ranks: the kernel (kernels/cuda/roce_gather.cu) stages this rank's shard in pinned host memory and rings a doorbell; a proxy thread here RDMA-writes it to the peer, striped over the HCAs, each stripe followed by its sequence flag on the same queue pair; the kernel waits for the peer's flags and copies both shards out. A decode step's gathers are a few rows of partials. On two GB10s with one ConnectX-7 cable (NCCL 2.30.7, `tf-cuda-test roce-gather`, 2,000 calls back to back) this path took 10.4 us a call at 4 KiB a rank against NCCL's 12.3, 11.6 against 13.5 at 8 KiB, 17.4 against 20.5 at 32 KiB, and 17.4 against 51.5 at 64 KiB. libibverbs is opened at run time (as the driver and NCCL are): without it, or without an active HCA, the caller keeps NCCL. -## Control link +## Setup and use -Two ranks over the hosts' network: a TCP control link (rank 0 listens, rank 1 connects; framed messages carry requests and per-round decisions) and the NCCL communicator it brings up (rank 0 makes the unique id and sends it). +`tp_group.Group.openRoce` brings the path up on a two-rank group: each rank opens its HCAs, and what each brings (its region, keys, queue pairs and the largest shard it would send) crosses the group's own NCCL communicator, as does each rank's verdict after connecting. Both ranks get the path or neither does (`TF_ROCE=0` on either rank keeps NCCL). From then on `Group.allGather` sends a gather of up to the smaller rank's cap (512 KiB by default) this way and leaves larger ones to NCCL; the size alone decides, so both ranks always take the same path. A rank whose buffers are not 16-byte aligned gets `error.Misaligned` instead of falling back alone. Liveness is `tp_watch`'s, as for every group. + +A failed exchange writes nothing into the gathered rows. Call `Group.check` after any synchronize whose result reads a gather: it returns the error the GPU or the proxy recorded. + +`tf-cuda-test roce-gather ` runs on both ranks (NCCL_COMM_ID set): every size up to the cap and past it against NCCL's all-gather, byte for byte, then each path's time a call. diff --git a/zig/src/cuda/roce.zig b/zig/src/cuda/roce.zig index 41dc961de8..25295d9bcc 100644 --- a/zig/src/cuda/roce.zig +++ b/zig/src/cuda/roce.zig @@ -7,7 +7,8 @@ const memory = @import("memory.zig"); const Module = @import("module.zig").Module; const Function = @import("module.zig").Function; const launch = @import("launch.zig"); -const tp_link = @import("tp_link.zig"); +const tp_group = @import("tp_group.zig"); +const Stream = @import("stream.zig").Stream; const kernels = @import("kernels.zig"); const max_hcas = 2; @@ -143,7 +144,20 @@ fn opsCall(comptime T: type, ctx: *Context, offset: usize) T { return slot.*; } -pub const Error = error{ VerbsUnavailable, NoHca, HcaFailed, PeerRefused, ProxyFailed, ExchangeTimedOut }; +pub const Error = error{ VerbsUnavailable, NoHca, HcaFailed, PeerRefused, ProxyFailed, ExchangeTimedOut, Misaligned }; + +/// TF_ROCE=0 keeps every gather on NCCL; unset or any other value asks for this path (both ranks must ask). +pub fn wanted() bool { + const v = std.c.getenv("TF_ROCE") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// The peer's `mine`: both ranks' values cross the group's communicator, rank order. +fn swap(comptime T: type, group: *const tp_group.Group, d: *const Driver, s: Stream, mine: T) !T { + var both: [2]T = undefined; + try group.gather(T, d, s, mine, &both); + return both[1 - group.rank]; +} // -- the region's layout (the same on both ranks) --------------------------------------------------------------------- const Layout = struct { @@ -219,20 +233,24 @@ pub const Gather = struct { err: [256]u8 = @splat(0), gathers: u64 = 0, - /// Both ranks' exchange, or null on both when either cannot; verdicts cross ``link``. ``max_bytes``: a shard. - pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const Driver, link: tp_link.Link, wanted: bool, max_bytes: usize) !?*Gather { + /// Both ranks' exchange or null on both; setup crosses ``group``'s communicator on ``s``. ``max_bytes``: a shard. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const Driver, group: *const tp_group.Group, s: Stream, want: bool, max_bytes: usize) !?*Gather { + if (group.world != 2) return error.PeerRefused; const g = try gpa.create(Gather); var mine: Blob = std.mem.zeroes(Blob); - const ok = wanted and blk: { - g.setup(gpa, io, d, link.rank, max_bytes) catch |e| { + const ok = want and blk: { + g.setup(gpa, io, d, @intCast(group.rank), max_bytes) catch |e| { std.log.warn("RoCE all-gather: {s} ({s}); NCCL carries every gather", .{ @errorName(e), std.mem.sliceTo(&g.err, 0) }); break :blk false; }; break :blk true; }; if (ok) g.blob(&mine); - var theirs: Blob = undefined; - try link.swap(std.mem.asBytes(&mine), std.mem.asBytes(&theirs)); + const theirs = swap(Blob, group, d, s, mine) catch |e| { + if (ok) g.teardown(); + gpa.destroy(g); + return e; + }; if (!ok or theirs.ok == 0 or theirs.n_hca != mine.n_hca or theirs.slot_bytes != mine.slot_bytes) { if (ok) g.teardown(); gpa.destroy(g); @@ -246,10 +264,12 @@ pub const Gather = struct { }; break :blk true; }; - var v: [1]u8 = .{@intFromBool(connected)}; - var pv: [1]u8 = undefined; - try link.swap(&v, &pv); - if (!connected or pv[0] == 0) { + const peer_connected = swap(u32, group, d, s, @intFromBool(connected)) catch |e| { + g.teardown(); + gpa.destroy(g); + return e; + }; + if (!connected or peer_connected == 0) { g.teardown(); gpa.destroy(g); return null; @@ -472,13 +492,14 @@ pub const Gather = struct { return @intCast(@min(mine, theirs)); } - /// Whether a gather of ``bytes`` a rank (from and to 16-byte aligned buffers) goes this way. - pub fn fits(g: *const Gather, bytes: usize, send: abi.DevicePtr, recv: abi.DevicePtr) bool { - return bytes > 0 and bytes <= g.max_bytes and bytes % 16 == 0 and send % 16 == 0 and recv % 16 == 0; + /// Whether a gather of ``bytes`` a rank goes this way: the size alone decides, so both ranks pick the same path. + pub fn fits(g: *const Gather, bytes: usize) bool { + return bytes > 0 and bytes <= g.max_bytes and bytes % 16 == 0; } - /// Both ranks' ``bytes`` from ``send`` into ``recv`` in rank order (as ncclAllGather), on ``stream``. + /// Both ranks' ``bytes`` from ``send`` into ``recv`` in rank order (as ncclAllGather); misaligned is an error. pub fn gather(g: *Gather, stream: abi.Stream, send: abi.DevicePtr, recv: abi.DevicePtr, bytes: usize) !void { + if (send % 16 != 0 or recv % 16 != 0) return error.Misaligned; const packs = bytes / 16; const need = @max(1, std.math.divCeil(usize, packs, 2 * threads) catch unreachable); const grid: u32 = @intCast(@min(std.math.ceilPowerOfTwo(usize, need) catch max_blocks, max_blocks)); diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 2fce0ac8ce..81adb2529d 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -21,7 +21,6 @@ pub const Dim3 = launch.Dim3; pub const graph = @import("graph.zig"); pub const cublaslt = @import("cublaslt.zig"); pub const nccl = @import("nccl.zig"); -pub const tp_link = @import("tp_link.zig"); pub const roce = @import("roce.zig"); pub const triton = @import("triton.zig"); pub const aot = @import("aot.zig"); @@ -64,7 +63,6 @@ test { _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); - _ = @import("tp_link.zig"); _ = @import("ck.zig"); _ = @import("tp_watch.zig"); } diff --git a/zig/src/cuda/tp_group.zig b/zig/src/cuda/tp_group.zig index 7f26c7c846..de2b39948f 100644 --- a/zig/src/cuda/tp_group.zig +++ b/zig/src/cuda/tp_group.zig @@ -1,9 +1,10 @@ -//! A tensor-parallel group over NCCL: NCCL_COMM_ID gives every rank the same id; rows meet in rank order. +//! A tensor-parallel group over NCCL (NCCL_COMM_ID; rows in rank order); two ranks add RoCE small gathers (roce.zig). const std = @import("std"); const nccl = @import("nccl.zig"); const driver = @import("driver.zig"); const stream_ = @import("stream.zig"); const memory = @import("memory.zig"); +const roce = @import("roce.zig"); const Stream = stream_.Stream; @@ -12,6 +13,7 @@ pub const Group = struct { comm: nccl.Comm, rank: usize, world: usize, + fast: ?*roce.Gather = null, // two ranks' small all-gathers by RoCE writes (openRoce), when both ranks brought it up /// Rank `rank` of `world` (NCCL opened at run time; NCCL_COMM_ID names rank 0's address, every rank the same). pub fn open(rank: usize, world: usize) !Group { @@ -26,13 +28,26 @@ pub const Group = struct { } pub fn close(g: *Group) void { + if (g.fast) |f| f.deinit(); _ = g.lib.api.ncclCommDestroy(g.comm); g.lib.close(); g.* = undefined; } + /// Two ranks: gathers up to `max_bytes` a rank go by RoCE writes if both ranks bring it up; else NCCL keeps them. + pub fn openRoce(g: *Group, gpa: std.mem.Allocator, io: std.Io, d: *const driver.Driver, s: Stream, max_bytes: usize) !void { + if (g.world != 2 or g.fast != null) return; + g.fast = try roce.Gather.init(gpa, io, d, g, s, roce.wanted(), max_bytes); + } + + /// A failed RoCE gather's error (its rows hold nothing valid): call it after a synchronize that reads a gather. + pub fn check(g: *const Group) !void { + if (g.fast) |f| try f.check(); + } + /// `bytes` at `mine` from every rank into `all` in rank order (world * bytes there), on `s`. pub fn allGather(g: *const Group, s: Stream, mine: u64, all: u64, bytes: usize) !void { + if (g.fast) |f| if (f.fits(bytes)) return f.gather(s.handle, mine, all, bytes); try g.lib.check(g.lib.api.ncclAllGather(mine, all, bytes, .u8, g.comm, s.handle), "ncclAllGather"); } @@ -68,6 +83,7 @@ pub const Group = struct { try dev.uploadAsync(g.rank * size, std.mem.asBytes(&mine), s.handle); try g.allGather(s, dev.ptr + g.rank * size, dev.ptr, size); try s.synchronize(); + try g.check(); try dev.download(0, std.mem.sliceAsBytes(out)); } }; diff --git a/zig/src/cuda/tp_link.zig b/zig/src/cuda/tp_link.zig deleted file mode 100644 index 53c9839737..0000000000 --- a/zig/src/cuda/tp_link.zig +++ /dev/null @@ -1,169 +0,0 @@ -//! Two ranks' TCP control link (rank 0 listens, rank 1 connects) and the NCCL communicator it brings up. - -const std = @import("std"); -const posix = std.posix; -const net = std.Io.net; -const Threaded = std.Io.Threaded; -const nccl = @import("nccl.zig"); - -pub const Error = error{ SocketFailed, BindFailed, ListenFailed, AcceptFailed, ConnectTimedOut, PeerClosed, IoFailed, FrameTooLarge }; - -pub const max_frame = 64 << 20; - -/// A message: a tag the protocol defines and its bytes (owned by the caller's allocator). -pub const Message = struct { tag: u32, bytes: []u8 }; - -pub const Link = struct { - fd: posix.socket_t, - rank: u8, - - /// Rank 0: listen on ``port`` (all interfaces) and accept rank 1, waiting up to ``timeout_ms``. - pub fn lead(port: u16, timeout_ms: i32) Error!Link { - const address = net.IpAddress.parse("0.0.0.0", port) catch unreachable; - var storage: Threaded.PosixAddress = undefined; - const len = Threaded.addressToPosix(&address, &storage); - const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); - if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; - const lfd: posix.socket_t = @intCast(rc); - defer _ = posix.system.close(lfd); - const one: c_int = 1; - posix.setsockopt(lfd, posix.SOL.SOCKET, posix.SO.REUSEADDR, std.mem.asBytes(&one)) catch return error.SocketFailed; - if (posix.errno(posix.system.bind(lfd, &storage.any, len)) != .SUCCESS) return error.BindFailed; - if (posix.errno(posix.system.listen(lfd, 1)) != .SUCCESS) return error.ListenFailed; - var fds = [_]posix.pollfd{.{ .fd = lfd, .events = posix.POLL.IN, .revents = 0 }}; - const ready = posix.poll(&fds, timeout_ms) catch return error.AcceptFailed; - if (ready == 0) return error.ConnectTimedOut; - const a = posix.system.accept(lfd, null, null); - if (posix.errno(a) != .SUCCESS) return error.AcceptFailed; - const fd: posix.socket_t = @intCast(a); - noDelay(fd); - return .{ .fd = fd, .rank = 0 }; - } - - /// Rank 1: connect to rank 0 at ``address``, retrying every 200 ms for up to ``timeout_ms``. - pub fn follow(io: std.Io, address: net.IpAddress, timeout_ms: i64) Error!Link { - var waited: i64 = 0; - while (true) { - var storage: Threaded.PosixAddress = undefined; - const len = Threaded.addressToPosix(&address, &storage); - const family: u32 = if (address == .ip4) posix.AF.INET else posix.AF.INET6; - const rc = posix.system.socket(family, posix.SOCK.STREAM, 0); - if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; - const fd: posix.socket_t = @intCast(rc); - if (posix.errno(posix.system.connect(fd, &storage.any, len)) == .SUCCESS) { - noDelay(fd); - return .{ .fd = fd, .rank = 1 }; - } - _ = posix.system.close(fd); - if (waited >= timeout_ms) return error.ConnectTimedOut; - std.Io.sleep(io, .fromMilliseconds(200), .awake) catch {}; - waited += 200; - } - } - - pub fn close(self: *Link) void { - _ = posix.system.shutdown(self.fd, posix.SHUT.RDWR); - _ = posix.system.close(self.fd); - self.* = undefined; - } - - /// One frame: tag and length (little-endian u32 each), then the bytes. - pub fn send(self: Link, tag: u32, bytes: []const u8) Error!void { - if (bytes.len > max_frame) return error.FrameTooLarge; - var head: [8]u8 = undefined; - std.mem.writeInt(u32, head[0..4], tag, .little); - std.mem.writeInt(u32, head[4..8], @intCast(bytes.len), .little); - try writeAll(self.fd, &head); - try writeAll(self.fd, bytes); - } - - /// The next frame; its bytes allocated with ``gpa``. - pub fn recv(self: Link, gpa: std.mem.Allocator) (Error || std.mem.Allocator.Error)!Message { - var head: [8]u8 = undefined; - try readAll(self.fd, &head); - const tag = std.mem.readInt(u32, head[0..4], .little); - const n = std.mem.readInt(u32, head[4..8], .little); - if (n > max_frame) return error.FrameTooLarge; - const bytes = try gpa.alloc(u8, n); - errdefer gpa.free(bytes); - try readAll(self.fd, bytes); - return .{ .tag = tag, .bytes = bytes }; - } - - /// Both ranks' ``mine`` exchanged (equal lengths): returns the peer's into ``theirs``. - pub fn swap(self: Link, mine: []const u8, theirs: []u8) Error!void { - if (self.rank == 0) { - try writeAll(self.fd, mine); - try readAll(self.fd, theirs); - } else { - try readAll(self.fd, theirs); - try writeAll(self.fd, mine); - } - } -}; - -fn noDelay(fd: posix.socket_t) void { - const one: c_int = 1; - posix.setsockopt(fd, posix.IPPROTO.TCP, std.posix.TCP.NODELAY, std.mem.asBytes(&one)) catch {}; -} - -fn writeAll(fd: posix.socket_t, bytes: []const u8) Error!void { - var at: usize = 0; - while (at < bytes.len) { - const rc = posix.system.write(fd, bytes[at..].ptr, bytes.len - at); - switch (posix.errno(rc)) { - .SUCCESS => at += @intCast(rc), - .INTR, .AGAIN => continue, - .PIPE, .CONNRESET => return error.PeerClosed, - else => return error.IoFailed, - } - } -} - -fn readAll(fd: posix.socket_t, bytes: []u8) Error!void { - var at: usize = 0; - while (at < bytes.len) { - const rc = posix.system.read(fd, bytes[at..].ptr, bytes.len - at); - switch (posix.errno(rc)) { - .SUCCESS => { - if (rc == 0) return error.PeerClosed; - at += @intCast(rc); - }, - .INTR, .AGAIN => continue, - .CONNRESET => return error.PeerClosed, - else => return error.IoFailed, - } - } -} - -/// The two ranks' NCCL communicator: rank 0's unique id crosses the link, then both join. -pub fn communicator(lib: *nccl.Library, link: Link) !nccl.Comm { - var uid: nccl.UniqueId = undefined; - if (link.rank == 0) { - try lib.check(lib.api.ncclGetUniqueId(&uid), "ncclGetUniqueId"); - try link.send(1, &uid.internal); - } else { - var buf: [256]u8 = undefined; - var fba = std.heap.FixedBufferAllocator.init(&buf); - const m = try link.recv(fba.allocator()); - if (m.tag != 1 or m.bytes.len != uid.internal.len) return error.IoFailed; - @memcpy(&uid.internal, m.bytes); - } - var comm: nccl.Comm = null; - try lib.check(lib.api.ncclCommInitRank(&comm, 2, uid, link.rank), "ncclCommInitRank"); - return comm; -} - -test "frames round-trip through a socket pair" { - var fds: [2]posix.socket_t = undefined; - if (std.c.socketpair(posix.AF.UNIX, posix.SOCK.STREAM, 0, &fds) != 0) return error.SkipZigTest; - const a: Link = .{ .fd = fds[0], .rank = 0 }; - const b: Link = .{ .fd = fds[1], .rank = 1 }; - try a.send(7, "hello ranks"); - const m = try b.recv(std.testing.allocator); - defer std.testing.allocator.free(m.bytes); - try std.testing.expectEqual(@as(u32, 7), m.tag); - try std.testing.expectEqualStrings("hello ranks", m.bytes); - _ = posix.system.close(fds[0]); - _ = posix.system.close(fds[1]); -} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..a68a7a880d 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -8,6 +8,7 @@ const bench = @import("bench.zig"); const oracle_tests = @import("oracle_tests.zig"); const libs_tests = @import("libs_tests.zig"); const rank_scatter_tests = @import("rank_scatter_tests.zig"); +const roce_gather_tests = @import("roce_gather_tests.zig"); const sample_tests = @import("sample_tests.zig"); const glue_tests = @import("glue_tests.zig"); const window_profile = @import("window_profile.zig"); @@ -52,6 +53,7 @@ const usage = \\ cublaslt bf16 GEMM through cuBLASLt against an fp64 reference \\ nccl one-rank NCCL all-reduce and all-gather \\ rank-scatter [id] rank-order reduce-scatter (NCCL_COMM_ID or an id file) against an all-gather + \\ roce-gather two ranks (NCCL_COMM_ID): the RoCE all-gather against NCCL's, graph replays, time a call \\ rank-sum the reduce-scatter's sum kernel alone, 2 to 5 ranks as each rank (one GPU, no NCCL) \\ nccl-pair one rank of a multi-host communicator: send/recv with every peer, all-reduce, all-gather \\ gdn-replay replay_kernel bits against the Python oracle's fixture @@ -146,6 +148,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { const world = try std.fmt.parseInt(usize, try arg(rest, 1), 10); return rank_scatter_tests.run(gpu, rank, world, if (rest.len > 2) rest[2] else null); } + if (std.mem.eql(u8, cmd, "roce-gather")) return roce_gather_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10)); if (std.mem.eql(u8, cmd, "rank-sum")) return rank_scatter_tests.sum(gpu); if (std.mem.eql(u8, cmd, "nccl-pair")) return libs_tests.ncclPair(gpu, try std.fmt.parseInt(c_int, try arg(rest, 0), 10), try std.fmt.parseInt(c_int, try arg(rest, 1), 10), try arg(rest, 2)); if (std.mem.eql(u8, cmd, "gdn-replay")) return oracle_tests.gdnReplay(gpu, try arg(rest, 0)); diff --git a/zig/tests/cuda/roce_gather_tests.zig b/zig/tests/cuda/roce_gather_tests.zig new file mode 100644 index 0000000000..18809d55e2 --- /dev/null +++ b/zig/tests/cuda/roce_gather_tests.zig @@ -0,0 +1,135 @@ +//! Two ranks' RoCE all-gather (Group.openRoce) against NCCL's byte for byte, graph replays, time a call (NCCL_COMM_ID). +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const cap = 512 << 10; +/// Shards a rank: tiny, a few decode rows, the cap, and past it or not a multiple of 16 (those stay on NCCL). +const sizes = [_]usize{ 16, 48, 4096, 8192, 32768, 65536, 262144, cap, cap + 16, 100 }; +const seeds = [_]u64{ 1, 2, 3 }; +const graph_gathers = 4; +const graph_replays = 64; +const reps = 2000; + +/// Rank `rank`'s shard for a case: bytes that differ by rank, size and seed. +fn shard(out: []u8, seed: u64, rank: usize) void { + var rng = std.Random.DefaultPrng.init(seed *% 1000003 +% rank *% 7919 +% out.len); + rng.random().bytes(out); +} + +pub fn run(gpu: Gpu, rank: usize) !void { + const a = gpu.gpa; + try check.expect(std.c.getenv("NCCL_COMM_ID") != null, "roce-gather: NCCL_COMM_ID names rank 0's address", .{}); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var g = try cuda.tp_group.Group.open(rank, 2); + defer g.close(); + try g.openRoce(a, gpu.io, gpu.d, s, cap); + try check.expect(g.fast != null, "roce-gather: the RoCE path did not come up on both ranks (see the warning above)", .{}); + + const most = sizes[std.mem.indexOfMax(usize, &sizes)]; + var mine = try cuda.DeviceBuffer.alloc(gpu.d, most); + defer mine.free(); + var by_group = try cuda.DeviceBuffer.alloc(gpu.d, 2 * most); + defer by_group.free(); + var by_nccl = try cuda.DeviceBuffer.alloc(gpu.d, 2 * most); + defer by_nccl.free(); + const host = try a.alloc(u8, most); + defer a.free(host); + const peer = try a.alloc(u8, most); + defer a.free(peer); + const got = try a.alloc(u8, 2 * most); + defer a.free(got); + const want = try a.alloc(u8, 2 * most); + defer a.free(want); + + var by_roce: usize = 0; + for (sizes) |n| for (seeds) |seed| { + shard(host[0..n], seed, rank); + shard(peer[0..n], seed, 1 - rank); + try mine.upload(0, host[0..n]); + try g.allGather(s, mine.ptr, by_group.ptr, n); + try g.lib.check(g.lib.api.ncclAllGather(mine.ptr, by_nccl.ptr, n, .u8, g.comm, s.handle), "ncclAllGather"); + try s.synchronize(); + try g.check(); + @memcpy(want[rank * n ..][0..n], host[0..n]); + @memcpy(want[(1 - rank) * n ..][0..n], peer[0..n]); + try by_group.download(0, got[0 .. 2 * n]); + try check.sameBytes("group all-gather against both ranks' shards", got[0 .. 2 * n], want[0 .. 2 * n]); + try by_nccl.download(0, got[0 .. 2 * n]); + try check.sameBytes("NCCL all-gather against both ranks' shards", got[0 .. 2 * n], want[0 .. 2 * n]); + if (g.fast.?.fits(n)) by_roce += 1; + }; + check.pass("roce-gather rank {d}: shards of {any} bytes (seeds 1-3) equal NCCL's all-gather; {d} of {d} cases by RoCE", .{ rank, sizes, by_roce, sizes.len * seeds.len }); + + try replays(gpu, &g, s, rank); + try timing(gpu, &g, s, rank); +} + +/// A graph of `graph_gathers` RoCE gathers replayed with new shards each time: each replay's rows are its own. +fn replays(gpu: Gpu, g: *cuda.tp_group.Group, s: cuda.Stream, rank: usize) !void { + const a = gpu.gpa; + const n = 8192; + var mine = try cuda.DeviceBuffer.alloc(gpu.d, graph_gathers * n); + defer mine.free(); + var all = try cuda.DeviceBuffer.alloc(gpu.d, graph_gathers * 2 * n); + defer all.free(); + try cuda.graph.beginCapture(s, .thread_local); + for (0..graph_gathers) |i| try g.allGather(s, mine.ptr + i * n, all.ptr + i * 2 * n, n); + var captured = try cuda.graph.endCapture(s); + defer captured.deinit(); + var exec = try captured.instantiate(); + defer exec.deinit(); + const host = try a.alloc(u8, graph_gathers * n); + defer a.free(host); + const peer = try a.alloc(u8, graph_gathers * n); + defer a.free(peer); + const got = try a.alloc(u8, graph_gathers * 2 * n); + defer a.free(got); + const want = try a.alloc(u8, graph_gathers * 2 * n); + defer a.free(want); + for (0..graph_replays) |r| { + for (0..graph_gathers) |i| { + shard(host[i * n ..][0..n], 100 + r * graph_gathers + i, rank); + shard(peer[i * n ..][0..n], 100 + r * graph_gathers + i, 1 - rank); + @memcpy(want[i * 2 * n + rank * n ..][0..n], host[i * n ..][0..n]); + @memcpy(want[i * 2 * n + (1 - rank) * n ..][0..n], peer[i * n ..][0..n]); + } + try mine.upload(0, host); + try exec.launchOn(s); + try s.synchronize(); + try g.check(); + try all.download(0, got); + try check.sameBytes("a replayed graph's gathers", got, want); + } + check.pass("roce-gather rank {d}: a captured graph of {d} gathers of {d} bytes replayed {d} times, each replay's rows its own", .{ rank, graph_gathers, n, graph_replays }); +} + +/// Median microseconds a call of `reps` back-to-back gathers, RoCE against NCCL, at decode-sized shards. +fn timing(gpu: Gpu, g: *cuda.tp_group.Group, s: cuda.Stream, rank: usize) !void { + var t0 = try cuda.Event.init(gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(gpu.d, true); + defer t1.deinit(); + var mine = try cuda.DeviceBuffer.alloc(gpu.d, 65536); + defer mine.free(); + var all = try cuda.DeviceBuffer.alloc(gpu.d, 2 * 65536); + defer all.free(); + for ([_]usize{ 4096, 8192, 32768, 65536 }) |n| { + var us: [2]f64 = undefined; + for (0..2) |which| { + for (0..2) |round| { // the first round warms + try t0.record(s); + for (0..reps) |_| { + if (which == 0) try g.allGather(s, mine.ptr, all.ptr, n) else try g.lib.check(g.lib.api.ncclAllGather(mine.ptr, all.ptr, n, .u8, g.comm, s.handle), "ncclAllGather"); + } + try t1.record(s); + try t1.synchronize(); + if (round == 1) us[which] = @as(f64, try cuda.Event.elapsedMs(t0, t1)) * 1000.0 / reps; + } + } + try g.check(); + check.pass("roce-gather rank {d}: {d} bytes a rank, {d} calls: RoCE {d:.1} us a call, NCCL {d:.1} us", .{ rank, n, reps, us[0], us[1] }); + } +} From b416ec8aac0e16f2cf86b7e69516ceb892384a54 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 10:56:58 +0000 Subject: [PATCH 21/31] lane core: several queued prompts in one pass for a backend with prefill_many (Engine.addStreams; the lane host takes, sets up and launches each job as before and prefills the shared-pass ones together); a pass that errs or has no room leaves none prefilled and each stream goes alone; the CUDA and recorder wrappers forward it; the fake's batched backend and four host tests --- zig/src/core/lane_host.zig | 63 ++++++++++++---- zig/src/core/lane_host_admit.zig | 51 +++++++++++++ zig/src/core/lane_host_many_test.zig | 103 +++++++++++++++++++++++++++ zig/src/core/lane_host_test.zig | 1 + zig/src/core/lanes/backend.zig | 2 + zig/src/core/lanes/engine.zig | 7 ++ zig/src/core/lanes/fake.zig | 22 ++++++ zig/src/core/lanes/open.zig | 23 ++++++ zig/src/core/lanes/recorder.zig | 8 ++- zig/src/native/cuda_bound.zig | 6 ++ 10 files changed, 273 insertions(+), 13 deletions(-) create mode 100644 zig/src/core/lane_host_many_test.zig diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4a8f7d3493..c2f27f4ec1 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -274,17 +274,31 @@ 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 a free lane (or several, when the backend batches prompts); false when none can go. fn admitOne(h: *LaneHost) bool { + if (admit.batches(h)) return admit.many(h); + const job = switch (h.takeJob()) { + .none => return false, + .failed => return true, + .job => |j| j, + }; + const began = h.setup(job) orelse return true; + return h.launch(job, began); + } + + pub const Taken = union(enum) { none, failed, job: *Job }; + + /// The next queued request a lane is free for, moved to the admitted list (`failed`: it was finished instead). + pub fn takeJob(h: *LaneHost) Taken { h.lock(); // a stale device takes no reply: `learn.refuse` ends the queued ones instead if (h.stale or 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 .none; } const at = admit.pick(h) orelse { // a shared cache pool's head waits for room h.unlock(); - return false; + return .none; }; const job = h.queued.orderedRemove(at); job.seq = h.admissions; @@ -292,11 +306,16 @@ pub const LaneHost = struct { h.admitted.append(h.gpa, job) catch { h.unlock(); h.finish(job, .failed, "out of memory"); - return true; + return .failed; }; h.unlock(); + return .{ .job = job }; + } + + /// A taken job's stream, set up to start; null when the job ended instead (dropped or cancelled). + pub fn setup(h: *LaneHost, job: *Job) ?i96 { 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)) return h.ended(h.cancel(job)); // other work came first 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.cacheKey(), r.bornLen(), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { @@ -305,7 +324,7 @@ 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 return h.ended(h.drop(job, "the drafter could not start")); const recall = recalls.reader(h, job); job.stream = lanes.Stream.init(h.gpa, .{ .id = "request", @@ -332,7 +351,7 @@ pub const LaneHost = struct { .grammar = if (r.structure) |st| h.constraint(st) catch |e| { job.proposer.deinit(); recalls.release(job); - return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + return h.ended(h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)")); } else null, .decode_spans = h.spans(r), .decide = r.decide, @@ -342,26 +361,46 @@ pub const LaneHost = struct { }) catch { job.proposer.deinit(); recalls.release(job); - return h.drop(job, "out of memory"); + return h.ended(h.drop(job, "out of memory")); }; job.started = true; job.fill_on.store(true, .release); // status readers count the fill from here until the pass reports const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); job.began = began; + return began; + } + + fn ended(_: *LaneHost, _: bool) ?i96 { + return null; + } + + /// Whether a set-up job starts outside a shared prompt pass: a warm pass, or the lone driver's. + pub fn startsAlone(h: *LaneHost, job: *Job) bool { + return (job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job); + } + + /// A set-up job into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. + pub fn launch(h: *LaneHost, job: *Job, began: i96) bool { + const r = job.request; 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()) { h.filling.append(h.gpa, job) catch return h.drop(job, "out of memory"); return true; } - h.core.addStream(&job.stream) catch |e| return switch (e) { + h.core.addStream(&job.stream) catch |e| return h.failedPass(job, e); + h.prefilled(job, began); + _ = h.deliver(job); + return true; + } + + /// A prompt pass that failed: cancelled, no lane memory yet (back to the queue), or the backend's words. + pub fn failedPass(h: *LaneHost, job: *Job, e: anyerror) bool { + return switch (e) { error.Cancelled => h.cancel(job), error.LaneMemory => h.requeue(job), else => h.drop(job, h.words(e)), }; - h.prefilled(job, began); - _ = h.deliver(job); - return true; } /// The next filling job's next prompt chunk: the other streams' rounds run between chunks, not after the prompt. diff --git a/zig/src/core/lane_host_admit.zig b/zig/src/core/lane_host_admit.zig index 016ab4398a..fde33f7010 100644 --- a/zig/src/core/lane_host_admit.zig +++ b/zig/src/core/lane_host_admit.zig @@ -2,6 +2,7 @@ const api = @import("engine_api.zig"); const LaneHost = @import("lane_host.zig").LaneHost; const Job = LaneHost.Job; +const lanes = @import("lanes"); /// 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 +22,53 @@ 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); } + +/// Whether to admit several at once: prompts batch, none fill by chunks, and two or more wait with two lanes free. +pub fn batches(h: *LaneHost) bool { + if (!h.core.batchesPrompts() or h.core.fills()) return false; + h.lock(); + defer h.unlock(); + const lanes_now = @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes); + return h.queued.items.len >= 2 and lanes_now -| h.admitted.items.len >= 2; +} + +/// Every queued request a lane is free for, prompts in one pass (lone and warm ones as before); false: none taken. +pub fn many(h: *LaneHost) bool { + var jobs: [lanes.Engine.max_together]*Job = undefined; + var began: [lanes.Engine.max_together]i96 = undefined; + var n: usize = 0; + var took = false; + while (n < jobs.len) { + const job = switch (h.takeJob()) { + .none => break, + .failed => { + took = true; + continue; + }, + .job => |j| j, + }; + took = true; + const at = h.setup(job) orelse continue; + if (h.startsAlone(job)) { + _ = h.launch(job, at); + continue; + } + jobs[n] = job; + began[n] = at; + n += 1; + } + if (n == 1) return h.launch(jobs[0], began[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], began[0..n], errs[0..n]) |job, at, err| { + if (err) |e| { + _ = h.failedPass(job, e); + continue; + } + h.prefilled(job, at); + _ = h.deliver(job); + } + return took; +} diff --git a/zig/src/core/lane_host_many_test.zig b/zig/src/core/lane_host_many_test.zig new file mode 100644 index 0000000000..25b6993b0c --- /dev/null +++ b/zig/src/core/lane_host_many_test.zig @@ -0,0 +1,103 @@ +//! Several queued requests in one prompt pass (a backend with prefill_many): the replies a pass alone would give. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; +const gpa = std.testing.allocator; + +const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *Box = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *Box) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } +}; + +/// Each reply the fake target's own continuation of its prompt. +fn expectOwn(prompt: []const u32, tokens: []const u32) !void { + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, prompt); + for (tokens) |t| { + try std.testing.expectEqual(lanes.fake.next(history.items, null, history.items.len), t); + try history.append(gpa, t); + } +} + +/// `n` requests queued before the host starts (`sampled`: one the fake refuses), replies checked, batched passes. +fn serve(n: usize, memory_lanes: ?usize, sampled: ?usize) !struct { passes: usize, streams: usize } { + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .memory_lanes = memory_lanes, .refuse_sampled = sampled != null }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.batched(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 4 }); + var prompts: [6][5]u32 = undefined; + var requests: [6]Request = undefined; + var boxes: [6]Box = @splat(.{}); + defer for (&boxes) |*b| b.tokens.deinit(gpa); + const e = host.engine(); + for (0..n) |i| { + prompts[i] = .{ 2, 7, 1, 8, @intCast(10 + i) }; + requests[i] = .{ .prompt = &prompts[i], .max_tokens = 12 }; + if (sampled == i) requests[i].sampling = .{ .seed = 3, .temperature = 0.7, .top_k = 5 }; + try e.submit(@intCast(i + 1), &requests[i], .{ .ctx = &boxes[i], .event = Box.event }); + } + try host.start(); + defer host.stop(); + for (0..n) |i| { + if (sampled == i) { + try std.testing.expectEqual(Reason.failed, boxes[i].wait()); + continue; + } + try std.testing.expectEqual(Reason.length, boxes[i].wait()); + try std.testing.expectEqual(@as(usize, 12), boxes[i].tokens.items.len); + try expectOwn(&prompts[i], boxes[i].tokens.items); + } + return .{ .passes = target.many_passes, .streams = target.many_streams }; +} + +test "queued requests share one prompt pass, each reply the target's own continuation" { + const r = try serve(4, null, null); + try std.testing.expectEqual(@as(usize, 1), r.passes); + try std.testing.expectEqual(@as(usize, 4), r.streams); +} + +test "more requests than lanes: the free lanes' prompts in one pass, the rest after" { + const r = try serve(6, null, null); + try std.testing.expect(r.passes >= 1); + try std.testing.expect(r.streams >= 4 and r.streams <= 6); +} + +test "a backend without room for them all: they go one at a time, a later pass only as many as fit, same replies" { + const r = try serve(4, 2, null); + try std.testing.expect(r.streams <= 2 * r.passes); // the fake refuses a pass past its room: none took more +} + +test "one request the backend refuses fails alone: the pass is retried a stream at a time, the others' replies kept" { + const r = try serve(4, null, 2); + try std.testing.expectEqual(@as(usize, 0), r.passes); +} diff --git a/zig/src/core/lane_host_test.zig b/zig/src/core/lane_host_test.zig index 85a34c0e15..56f8e353cd 100644 --- a/zig/src/core/lane_host_test.zig +++ b/zig/src/core/lane_host_test.zig @@ -570,4 +570,5 @@ test { _ = @import("lane_host_borrow_test.zig"); _ = @import("lane_host_warm_test.zig"); _ = @import("lane_host_rows_test.zig"); + _ = @import("lane_host_many_test.zig"); } diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 8378c44900..31cf0f4786 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -68,6 +68,8 @@ 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, + /// Several prompts in one pass, each as after `prefill`; false or an error leaves none prefilled. + 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. diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index f11d8dbec5..7e47830941 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -79,6 +79,13 @@ pub const Engine = struct { } pub const addStream = open.add; + pub const addStreams = open.addStreams; + pub const max_together = open.max_together; + + /// 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; + } pub const opened = open.opened; pub const fills = open.fills; pub const fillStream = open.fillStream; diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index 7827f286ac..234fe8d63e 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -53,6 +53,8 @@ pub const Fake = struct { answer_cycles: bool = false, prefill_chunks: usize = 0, prefill_count: usize = 0, + many_passes: usize = 0, // prefill_many calls that took their streams (batched()) + many_streams: usize = 0, // streams those calls prefilled prefill_hook: ?*const fn (ctx: *anyopaque, s: *Stream, chunk: usize) void = null, prefill_hook_ctx: ?*anyopaque = null, refuse_sampled: bool = false, // prefill refuses a sampled stream with error.SamplingRefused @@ -79,6 +81,26 @@ pub const Fake = struct { return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .first_masked = firstMasked, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .alternatives = alternatives, .held = heldDrafts, .reports_fill = true } }; } + /// The same target taking several prompts in one pass (`prefill_many`). + pub fn batched(x: *Fake) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_many = prefillMany, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .reports_fill = true } }; + } + + fn prefillMany(ptr: *anyopaque, ss: []const *Stream) anyerror!bool { + const x = self(ptr); + if (x.memory_lanes) |n| if (x.lanes.count() + ss.len > n) return false; // not all at once: one at a time instead + for (ss, 0..) |s, i| prefill(ptr, s) catch |err| { // none left prefilled, as the contract says + for (ss[0..i]) |done| if (x.lanes.fetchRemove(done)) |kv| { + var l = kv.value; + freeLane(x.gpa, &l); + }; + return err; + }; + x.many_passes += 1; + x.many_streams += ss.len; + return true; + } + /// The same target filling its prompt a chunk a call (`prefill_chunks`), as a backend with `prefill_step`. pub fn stepped(x: *Fake) be.Backend { return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow, .reports_fill = true } }; diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..e0b0ba45e8 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -23,6 +23,29 @@ pub fn add(e: *Engine, s: *Stream) !void { try opened(e, s); } +/// Streams a backend prefills in one pass at most. +pub const max_together = 16; + +/// `add` for several streams, prompts in one pass when the backend takes them; `errs[i]`: stream i's error, or null. +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 false; // one stream's refusal is its own: each goes alone + if (!together) return addEach(e, ss, errs); + for (ss, errs) |s, *x| opened(e, s) catch |err| { // a stream cancelled in the pass gets no first token + x.* = err; + }; +} + +fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (ss, errs) |s, *x| add(e, s) catch |err| { + x.* = err; + }; +} + /// After the prompt pass: the first token drawn and committed, the first drafts asked; rounds from the next. pub fn opened(e: *Engine, s: *Stream) !void { if (s.isCancelled()) { // cancelled in its last chunk: no first token diff --git a/zig/src/core/lanes/recorder.zig b/zig/src/core/lanes/recorder.zig index e62eda60fd..841424ddb7 100644 --- a/zig/src/core/lanes/recorder.zig +++ b/zig/src/core/lanes/recorder.zig @@ -229,7 +229,7 @@ pub const Recorded = struct { pub fn backend(x: *Recorded) be.Backend { const v = x.inner.vtable; - x.table = .{ .prefill = prefill, .prefill_step = if (v.prefill_step != null) prefillStep else null, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .unspeculate = if (v.unspeculate != null) unspeculate else null, .probabilities = if (v.probabilities != null) probabilities else null, .tree = if (v.tree != null) tree else null, .alternatives = if (v.alternatives != null) alternatives else null, .prepare_features = if (v.prepare_features != null) prepareFeatures else null, .features = features, .first_row = if (v.first_row != null) firstRow else null, .first_masked = if (v.first_masked != null) firstMasked else null, .held = if (v.held != null) held else null, .prompt_rows = v.prompt_rows, .release = release }; + x.table = .{ .prefill = prefill, .prefill_step = if (v.prefill_step != null) prefillStep else null, .prefill_many = if (v.prefill_many != null) prefillMany else null, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .unspeculate = if (v.unspeculate != null) unspeculate else null, .probabilities = if (v.probabilities != null) probabilities else null, .tree = if (v.tree != null) tree else null, .alternatives = if (v.alternatives != null) alternatives else null, .prepare_features = if (v.prepare_features != null) prepareFeatures else null, .features = features, .first_row = if (v.first_row != null) firstRow else null, .first_masked = if (v.first_masked != null) firstMasked else null, .held = if (v.held != null) held else null, .prompt_rows = v.prompt_rows, .release = release }; return .{ .ptr = x, .vtable = &x.table }; } @@ -247,6 +247,12 @@ pub const Recorded = struct { if (done) x.recorder.prompt(x.inner, s); return done; } + fn prefillMany(p: *anyopaque, ss: []const *Stream) anyerror!bool { + const x = of(p); + if (!try x.inner.vtable.prefill_many.?(x.inner.ptr, ss)) return false; + for (ss) |s| x.recorder.prompt(x.inner, s); // each prompt's rows, as after its own prefill + return true; + } fn first(p: *anyopaque, s: *Stream, position: u64) anyerror!u64 { return of(p).inner.first(s, position); } diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 2fd155bceb..c825d65f4b 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); From aa9d4a7425be5950674112f9f966d93bf7c674e4 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 11:14:38 +0000 Subject: [PATCH 22/31] glm cuda: the ranks' communicator is a tp_group.Group; on two ranks the main stream's all-gathers (each layer's partials, the MTP head's, the greedy pick padded to 16 bytes) take its RoCE writes, side-stream prompt exchanges stay on NCCL; Group.check after every draw that read gathered rows --- zig/src/families/glm/cuda_engine.zig | 28 ++++++++++++++++++------ zig/src/families/glm/cuda_forward.zig | 4 +++- zig/src/families/glm/cuda_native.zig | 31 ++++++++++++++++----------- 3 files changed, 42 insertions(+), 21 deletions(-) diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index ee38e00b3e..65308ad312 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -23,6 +23,8 @@ pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) boo /// The most drafts a chain holds (decode.DepthPolicy's top depth). pub const max_drafts = 3; +/// A greedy pick as it crosses ranks: (value, id) and padding, so the RoCE gather (16-byte packs) can carry it. +const pair_bytes = 16; /// A prompt pass's chunk rows (decode.prefill's). pub const prompt_rows = fw.prefill_rows; @@ -63,6 +65,7 @@ pub const Engine = struct { spans: []shard.Span, // every rank's vocabulary span lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, + group: ?*const cuda.tp_group.Group = null, // the main stream's all-gathers (RoCE on two ranks); null: lib and comm settings: ranks.Settings, m: fw.Modules, stream: cuda.Stream, @@ -147,8 +150,9 @@ pub const Engine = struct { errdefer if (e.bs) |*b| b.deinit(); e.hold = try cuda.DeviceBuffer.alloc(d, c.hidden * 2); errdefer e.hold.free(); - e.pairs = try cuda.DeviceBuffer.alloc(d, 8 * world); + e.pairs = try cuda.DeviceBuffer.alloc(d, pair_bytes * world); errdefer e.pairs.free(); + try d.check(d.api.cuMemsetD8_v2(e.pairs.ptr, 0, pair_bytes * world), "cuMemsetD8"); // each pair's padding stays zero e.smp = try sampler.Sampler.init(d, c.vocab, fw.window_rows); errdefer e.smp.deinit(); e.shared_buf = try cuda.DeviceBuffer.alloc(d, 8); @@ -182,7 +186,7 @@ pub const Engine = struct { } fn exchange(e: *const Engine) fw.Exchange { - return .{ .nccl = .{ .lib = e.lib, .comm = e.comm } }; + return .{ .nccl = .{ .lib = e.lib, .comm = e.comm, .group = e.group } }; } fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { @@ -462,7 +466,10 @@ pub const Engine = struct { /// `rows` logits rows of this rank's vocabulary -> the tokens at positions first, first + 1, ... fn draw(e: *Engine, logits: u64, rows: usize, first: u64, out: []u32) !void { - if (e.cur.sampling) |s| return e.smp.draw(e.stream, e.exchange(), e.spans, &.{.{ .rank = e.share.rank, .logits = logits }}, rows, first, sampler.rule(s), out); + if (e.cur.sampling) |s| { + try e.smp.draw(e.stream, e.exchange(), e.spans, &.{.{ .rank = e.share.rank, .logits = logits }}, rows, first, sampler.rule(s), out); + return e.checkGathers(); // the draw read rows the forward's gathers made + } for (0..rows) |r| out[r] = try e.greedy(logits + r * e.share.vocab.len() * 2); } @@ -473,12 +480,14 @@ pub const Engine = struct { a.add(logits); a.add(@as(i64, @intCast(e.share.vocab.len()))); a.add(@as(u32, @intCast(e.share.vocab.lo))); - a.add(e.pairs.ptr + e.share.rank * 8); + a.add(e.pairs.ptr + e.share.rank * pair_bytes); try cuda.launch.launch(e.m.glue.top1, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, e.stream, &a); - try e.lib.check(e.lib.api.ncclAllGather(e.pairs.ptr + e.share.rank * 8, e.pairs.ptr, 2, .u32, e.comm, e.stream.handle), "ncclAllGather"); + const mine = e.pairs.ptr + e.share.rank * pair_bytes; + if (e.group) |g| try g.allGather(e.stream, mine, e.pairs.ptr, pair_bytes) else try e.lib.check(e.lib.api.ncclAllGather(mine, e.pairs.ptr, pair_bytes, .u8, e.comm, e.stream.handle), "ncclAllGather"); try e.stream.synchronize(); - var all: [ranks.max_world][2]u32 = undefined; // Engine.init refuses a larger world - try e.d.check(e.d.api.cuMemcpyDtoH_v2(&all, e.pairs.ptr, world * 8), "cuMemcpyDtoH"); + try e.checkGathers(); + var all: [ranks.max_world][4]u32 = undefined; // Engine.init refuses a larger world + try e.d.check(e.d.api.cuMemcpyDtoH_v2(&all, e.pairs.ptr, world * pair_bytes), "cuMemcpyDtoH"); var pick = all[0]; for (all[1..world]) |p| { const v: f32 = @bitCast(p[0]); @@ -488,6 +497,11 @@ pub const Engine = struct { return pick[1]; } + /// A failed RoCE gather's error, after a synchronize whose result read gathered rows (they hold nothing then). + fn checkGathers(e: *const Engine) !void { + if (e.group) |g| try g.check(); + } + fn heatRoot() []const u8 { const p = std.c.getenv("TF_HEAT_ROOT") orelse return heat.default_root; const s = std.mem.span(p); diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index 1e57e55a08..c93ab0079c 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/zig/src/families/glm/cuda_forward.zig @@ -134,7 +134,8 @@ pub const Pieces = struct { /// How the ranks' partials meet: every rank in-process copying into the others' slots, or one rank over NCCL. pub const Exchange = union(enum) { local, - nccl: struct { lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm }, + // group: the main stream's all-gathers through it (RoCE writes for small ones on two ranks); side streams stay raw + nccl: struct { lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, group: ?*const cuda.tp_group.Group = null }, }; /// One rank of the forward: its weights, its sequence's caches, the window's buffers. @@ -175,6 +176,7 @@ pub const Forward = struct { }, .nccl => |n| { const r = f.ranks[0]; + if (n.group) |g| return g.allGather(f.stream, r.b.gathered + r.w.share.rank * bytes, r.b.gathered, bytes); const count = rows * c.hidden; try n.lib.check(n.lib.api.ncclAllGather(r.b.gathered + r.w.share.rank * bytes, r.b.gathered, count, .f32, n.comm, f.stream.handle), "ncclAllGather"); }, diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index e85cb25ad7..2f09b68f41 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -85,13 +85,15 @@ pub fn capacity(context: usize) usize { const Owned = struct { gpa: std.mem.Allocator, - lib: cuda.nccl.Library, - comm: cuda.nccl.Comm, + group: cuda.tp_group.Group, // the ranks' communicator; two ranks add RoCE writes for the main stream's small gathers e: *engine.Engine, link: follow.Link, lanes: glm_lanes.Lanes, }; +/// The largest shard a rank sends by RoCE writes: a decode window's partials (8 rows of fp32) fit. +const roce_cap = 512 << 10; + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { _ = kernels; // our own kernels: a captured set is not used const p = try place(); @@ -104,23 +106,27 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ const own = try gpa.create(Owned); errdefer gpa.destroy(own); own.gpa = gpa; - own.lib = try cuda.nccl.Library.open(); - errdefer own.lib.close(); + var lib = try cuda.nccl.Library.open(); + errdefer lib.close(); var id: cuda.nccl.UniqueId = undefined; - try own.lib.check(own.lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); // NCCL_COMM_ID: the same id on every rank - own.comm = null; - try own.lib.check(own.lib.api.ncclCommInitRank(&own.comm, @intCast(p.world), id, @intCast(p.rank)), "ncclCommInitRank"); - errdefer _ = own.lib.api.ncclCommDestroy(own.comm); + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); // NCCL_COMM_ID: the same id on every rank + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(p.world), id, @intCast(p.rank)), "ncclCommInitRank"); + own.group = .{ .lib = lib, .comm = comm, .rank = p.rank, .world = p.world }; + errdefer _ = own.group.lib.api.ncclCommDestroy(own.group.comm); const cap = capacity(o.context); - own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, &own.lib, own.comm, p.rank, p.world, .{ .capacity = cap, .mtp = o.drafts, .settings = set }); + own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, &own.group.lib, own.group.comm, p.rank, p.world, .{ .capacity = cap, .mtp = o.drafts, .settings = set }); errdefer own.e.deinit(); - own.link = try follow.Link.init(gpa, ctx.d, &own.lib, own.comm, own.e.stream, cap); + try own.group.openRoce(gpa, io, ctx.d, own.e.stream, roce_cap); // both ranks bring it up or neither does + errdefer if (own.group.fast) |f| f.deinit(); + own.e.group = &own.group; + own.link = try follow.Link.init(gpa, ctx.d, &own.group.lib, own.group.comm, own.e.stream, cap); errdefer own.link.deinit(); // Rank 0 admits by the smallest room of the ranks; its sizes are the largest (the first ranks take the remainder). const rooms = try gpa.alloc(u64, p.world); defer gpa.free(rooms); const mine: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); - try ranks.gather(u64, ctx.d, &own.lib, own.comm, own.e.stream, p.rank, mine, rooms); + try ranks.gather(u64, ctx.d, &own.group.lib, own.group.comm, own.e.stream, p.rank, mine, rooms); var peer: ?Peer = null; for (rooms[1..], 1..) |r, q| if (peer == null or r < peer.?.bytes) { peer = .{ .bytes = r, .rank = q }; @@ -194,8 +200,7 @@ fn close(ptr: *anyopaque) void { own.lanes.deinit(); own.link.deinit(); own.e.deinit(); - _ = own.lib.api.ncclCommDestroy(own.comm); - own.lib.close(); + own.group.close(); own.gpa.destroy(own); } From 653bd37f53122ba82372bb129c43e9149a76525e Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 12:21:17 +0000 Subject: [PATCH 23/31] glm cuda: EXL3 routed experts (Mia-AiLab's checkpoint: 4-bit mcg trellis experts, the rest bf16): admitted beside ModelOpt NVFP4, loaded by the rank's 128-column cut, run by glm_exl3 (grouped decode, fused prompt) on an 8-slot plan into the MoE block's fp32 slot rows; bf16 dense layers through b16; the SwiGLU limits keep a NaN; tf-cuda-test glm-exl3 against float64 from core's EXL3 decoder --- zig/build/cuda.zig | 1 + zig/kernels/cuda/glm_exl3.cu | 363 +++++++++++++++++++++++ zig/kernels/cuda/glm_exl3_transposed.cuh | 168 +++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/glm/config.zig | 2 + zig/src/families/glm/cuda.zig | 2 + zig/src/families/glm/cuda_engine.zig | 4 +- zig/src/families/glm/cuda_exl3.zig | 200 +++++++++++++ zig/src/families/glm/cuda_forward.zig | 9 +- zig/src/families/glm/cuda_mlp.zig | 29 +- zig/src/families/glm/cuda_moe.zig | 15 +- zig/src/families/glm/cuda_shard.zig | 2 +- zig/src/families/glm/cuda_state.zig | 13 +- zig/src/families/glm/cuda_weights.zig | 67 +++-- zig/tests/cuda/glm_exl3_tests.zig | 238 +++++++++++++++ zig/tests/cuda/main.zig | 3 + 16 files changed, 1078 insertions(+), 39 deletions(-) create mode 100644 zig/kernels/cuda/glm_exl3.cu create mode 100644 zig/kernels/cuda/glm_exl3_transposed.cuh create mode 100644 zig/src/families/glm/cuda_exl3.zig create mode 100644 zig/tests/cuda/glm_exl3_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index f36b6780fe..6d7ff7ee81 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -52,6 +52,7 @@ const kernels = [_]Kernel{ .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family + .{ .name = "glm_exl3", .flags = &.{"-O3"} }, // ours: GLM's EXL3 routed experts (Mia-AiLab's checkpoint), decode and prompt .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, diff --git a/zig/kernels/cuda/glm_exl3.cu b/zig/kernels/cuda/glm_exl3.cu new file mode 100644 index 0000000000..4846e60371 --- /dev/null +++ b/zig/kernels/cuda/glm_exl3.cu @@ -0,0 +1,363 @@ +// GLM-5.3-Flash's EXL3 routed experts (4-bit mcg trellis, ExLlamaV3's format: MIT, (c) 2025 Turboderp): decode, prompt. + +#include +#include +#include + +namespace { + +// Lets a dependent launch start on its weights early (sm_90 and newer; older targets run in plain stream order). +__device__ __forceinline__ void launch_dependents() { +#if __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +} + +constexpr float HAD_SCALE = 0.08838834764831845f; // 1 / sqrt(128) + +// Two values of the "mcg" codebook from two 16-bit states, as a half2 (first state in .x). +__device__ __forceinline__ uint32_t mcg2(uint32_t s0, uint32_t s1) { + uint32_t x0 = s0 * 0xCBAC1FEDu; + uint32_t x1 = s1 * 0xCBAC1FEDu; + x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u; + x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u; + uint32_t lo = __byte_perm(x0, x1, 0x5410); + uint32_t hi = __byte_perm(x0, x1, 0x7632); + half2 r = __hadd2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + return *reinterpret_cast(&r); +} + +// This lane's eight values of a 4-bit tile (word = tile[lane]) as the B fragments of its two n8 halves. +__device__ __forceinline__ void decode_tile(uint32_t w, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { + uint32_t p = __shfl_sync(0xffffffffu, w, (lane + 31) & 31); + uint32_t s = __funnelshift_r(w, p, 20); + b0[0] = mcg2((s >> 8) & 0xffffu, (s >> 4) & 0xffffu); + b0[1] = mcg2(s & 0xffffu, w >> 16); + b1[0] = mcg2((w >> 12) & 0xffffu, (w >> 8) & 0xffffu); + b1[1] = mcg2((w >> 4) & 0xffffu, w & 0xffffu); +} + +__device__ __forceinline__ void mma16816(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +__device__ __forceinline__ uint32_t load_pair(const half* x, bool ok) { + return ok ? *reinterpret_cast(x) : 0u; +} + +// Grid (item, n block, mat * SK + split) -> Z[mat][split][pair][n]: an expert and <= 16 MT of its pairs; warp order. +template +__device__ __forceinline__ void grouped( + const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int SK) { + const int item = blockIdx.x; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + const int split = blockIdx.z % SK; + const int mat = blockIdx.z / SK; + const half* X = mat ? X1 : X0; + const uint32_t* T = mat ? T1 : T0; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int KT = K >> 4, NTILES = N >> 4; + constexpr int ROWS = 16 * MT; + + __shared__ int rows_sh[ROWS]; + for (int i = threadIdx.x; i < ROWS; i += W * 32) rows_sh[i] = i < cnt ? members[first + i] : -1; + __syncthreads(); + int r0[MT], r1[MT]; + const half* x0[MT]; + const half* x1[MT]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + r0[m] = rows_sh[m * 16 + g]; + r1[m] = rows_sh[m * 16 + g + 8]; + x0[m] = X + (size_t)(r0[m] < 0 ? 0 : r0[m]) * K + 2 * t; + x1[m] = X + (size_t)(r1[m] < 0 ? 0 : r1[m]) * K + 2 * t; + } + + const int per_split = KT / SK, per_warp = per_split / W; + const int kt0 = split * per_split + warp * per_warp; + const int nt0 = blockIdx.y * NT; + const uint32_t* tile = T + (((size_t)e * KT + kt0) * NTILES + nt0) * 32 + lane; + + float acc[MT][NT][2][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[m][i][h][c] = 0.f; + + for (int kt = kt0; kt < kt0 + per_warp; ++kt) { + uint32_t words[NT]; +#pragma unroll + for (int i = 0; i < NT; ++i) words[i] = __ldg(tile + i * 32); + const int k = kt * 16; + uint32_t a[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + a[m][0] = load_pair(x0[m] + k, r0[m] >= 0); + a[m][1] = load_pair(x1[m] + k, r1[m] >= 0); + a[m][2] = load_pair(x0[m] + k + 8, r0[m] >= 0); + a[m][3] = load_pair(x1[m] + k + 8, r1[m] >= 0); + } +#pragma unroll + for (int i = 0; i < NT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile(words[i], lane, b0, b1); +#pragma unroll + for (int m = 0; m < MT; ++m) { + mma16816(acc[m][i][0], a[m], b0); + mma16816(acc[m][i][1], a[m], b1); + } + } + tile += (size_t)NTILES * 32; + } + + // warps' partial sums through shared memory, added in warp order + __shared__ float red[W][ROWS][NT * 16]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + red[warp][m * 16 + g][col] = acc[m][i][h][0]; + red[warp][m * 16 + g][col + 1] = acc[m][i][h][1]; + red[warp][m * 16 + g + 8][col] = acc[m][i][h][2]; + red[warp][m * 16 + g + 8][col + 1] = acc[m][i][h][3]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < ROWS * NT * 16; idx += W * 32) { + const int row = idx / (NT * 16), col = idx % (NT * 16); + const int r = rows_sh[row]; + if (r < 0) continue; + float s = red[0][row][col]; +#pragma unroll + for (int w = 1; w < W; ++w) s += red[w][row][col]; + Z[(((size_t)mat * SK + split) * P + r) * N + nt0 * 16 + col] = s; + } +} + +// Fast Walsh-Hadamard transform of 128 values, 4 a lane, natural order: strides 1, 2 in registers, 4..64 across lanes. +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3]; + v[0] = a + c; v[1] = b + d; v[2] = a - c; v[3] = b - d; +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +// The SwiGLU limits as torch.clamp takes them: a NaN stays NaN (fminf and fmaxf would drop it for the bound). +__device__ __forceinline__ float clamp_hi(float x, float hi) { return x > hi ? hi : x; } +__device__ __forceinline__ float clamp_both(float x, float l) { return x < -l ? -l : (x > l ? l : x); } + + +// -- helpers the transposed prompt kernels share (glm_exl3_transposed.cuh) -- + +// decode_tile via PRMT (MiaAI-Lab GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold patch 0020, Apache-2.0): same bits. +__device__ __forceinline__ void decode_tile_p(uint32_t w, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { + const uint32_t p = __shfl_sync(0xffffffffu, w, (lane + 31) & 31); + const uint32_t s = __funnelshift_r(w, p, 20); + const uint32_t t = w >> 4; + b0[0] = mcg2(__byte_perm(s, 0u, 0x4421), (s >> 4) & 0xffffu); // (s >> 8) & 0xffff, (s >> 4) & 0xffff + b0[1] = mcg2(s & 0xffffu, w >> 16); + b1[0] = mcg2(__byte_perm(t, 0u, 0x4421), __byte_perm(w, 0u, 0x4421)); // (w >> 12) & 0xffff, (w >> 8) & 0xffff + b1[1] = mcg2(__byte_perm(t, 0u, 0x4410), w & 0xffffu); // (w >> 4) & 0xffff, w & 0xffff +} + +__device__ __forceinline__ void cp_async16(void* dst, const void* src, bool valid) { + const unsigned d = (unsigned)__cvta_generic_to_shared(dst); + const int n = valid ? 16 : 0; // src-size 0: the 16 bytes come as zeros (a padding row) + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(d), "l"(src), "r"(n)); +} +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ void ldmatrix_x4(uint32_t (&a)[4], const half* p) { + const unsigned s = (unsigned)__cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) : "r"(s)); +} + +} // namespace + +// Program (row, K block): Xh[row][block] = fp16((x[row] * suh) @ H) when a layer's experts share suh; rot_in's bits. +extern "C" __global__ void glm_exl3_rot_rows(const __nv_bfloat16* __restrict__ x, int x_stride, const half* __restrict__ suh, + half* __restrict__ out, int K) { + launch_dependents(); + const int row = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const half* sh = suh + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(sh[j]); + fwht128(v, lane); + half* o = out + (size_t)row * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +namespace { +} // namespace + +// decode windows: items of 16 pairs, 8 n-tiles a block, 4 warps +extern "C" __global__ void __launch_bounds__(128) glm_exl3_grouped(const half* __restrict__ X0, const half* __restrict__ X1, + const uint32_t* __restrict__ T0, const uint32_t* __restrict__ T1, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, float* __restrict__ Z, int K, int N, int P, int SK) { + launch_dependents(); // a dependent matmul may start on its weights + grouped<1, 8, 4>(X0, X1, T0, T1, items, counts, members, Z, K, N, P, SK); +} + +// Program (pair, K block, matrix): Xh[mat][pair][block] = fp16((x[row] * suh[mat][e]) @ H), pair = row * 8 + slot. +extern "C" __global__ void glm_exl3_rot_in(const __nv_bfloat16* __restrict__ x, int x_stride, const int* __restrict__ pick, + const half* __restrict__ suh0, const half* __restrict__ suh1, half* __restrict__ out0, half* __restrict__ out1, int K) { + launch_dependents(); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int row = p >> 3; + const int lane = threadIdx.x; + const int e = pick[p]; + const half* suh = (mat ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? out1 : out0) + (size_t)p * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (pair, width block): gate, up over splits in order, rotated, svh, limited SwiGLU, Xd = fp16(act suh_d H). +extern "C" __global__ void glm_exl3_gateup(const float* __restrict__ Z, const int* __restrict__ pick, const half* __restrict__ svh_g, + const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, int P, int N, int SK, float limit) { + launch_dependents(); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const int e = pick[p]; + const int n = blk * 128 + 4 * lane; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float sg = 0.f, su = 0.f; + for (int s = 0; s < SK; ++s) { + sg += Z[((size_t)(0 * SK + s) * P + p) * N + n + j]; + su += Z[((size_t)(1 * SK + s) * P + p) * N + n + j]; + } + gv[j] = sg; + uv[j] = su; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float gg = clamp_hi(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = clamp_both(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), limit); + float act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (pair, model-width block): down over splits in order, rotated, svh -> y[row][slot] fp32; this rank's share. +extern "C" __global__ void glm_exl3_down(const float* __restrict__ Z, const int* __restrict__ pick, const half* __restrict__ svh_d, + float* __restrict__ y, int P, int D, int SK, int ys) { + launch_dependents(); // a dependent matmul may start on its weights + const int p = blockIdx.x, blk = blockIdx.y; + const int lane = threadIdx.x; + const int e = pick[p]; + const int n = blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int k = 0; k < SK; ++k) s += Z[((size_t)k * P + p) * D + n + j]; + v[j] = s; + } + fwht128(v, lane); + float* o = y + ((size_t)(p >> 3) * ys + (p & 7)) * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); +} + +// -- prompt chunks, the transposed design: Y^T = W^T X^T, one fp32 mma chain an output over ascending k tiles ------- + +// After MiaAI-Lab's prompt kernels (GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold patches 0004, 0009, Apache-2.0). +#include "glm_exl3_transposed.cuh" + +// Prompt items of 64 pairs: block (128 expert columns, item), 8 warps: Xd = gate/up epilogue; xshift 3 or 0 as above. +extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_gu(const half* __restrict__ Xg, const half* __restrict__ Xu, + const uint32_t* __restrict__ Tg, const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, + half* __restrict__ xd, int K, int N, int xshift, float limit) { + launch_dependents(); + __shared__ __align__(16) unsigned char smem_raw[tsmem_bytes()]; // one buffer for either instantiation + __shared__ int rows_sh[TPM]; + if (xshift) tgateup(Xg, Xu, Tg, Tu, items, counts, members, svh_g, svh_u, suh_d, xd, K, N, xshift, limit, smem_raw, rows_sh); + else tgateup<2 * TPM>(Xg, Xu, Tg, Tu, items, counts, members, svh_g, svh_u, suh_d, xd, K, N, 0, limit, smem_raw, rows_sh); +} + +// Block (256 model-width columns, item): y[row][slot][:] = the down epilogue of the item's pairs (this rank's share). +extern "C" __global__ void __launch_bounds__(256, 2) glm_exl3_prompt_dn(const half* __restrict__ Xd, const uint32_t* __restrict__ Td, + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, const half* __restrict__ svh_d, + float* __restrict__ y, int K, int D, int ys) { + launch_dependents(); + __shared__ __align__(16) unsigned char smem_raw[tsmem_bytes()]; + __shared__ int rows_sh[TPM]; + const int cb = blockIdx.x, item = blockIdx.y; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + tmain<1, TPM>(Xd, Xd, Td, Td, rows_sh, cnt, e, K, D >> 4, cb, 0, reinterpret_cast(smem_raw), acc); + float* Ed = reinterpret_cast(smem_raw); + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); + tspill(Ed, DN_ES, warp * TMT * 16, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + const int p = rows_sh[hf * TPH + i]; +#pragma unroll + for (int blk = 0; blk < 2; ++blk) { + const int n = cb * 256 + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = Ed[(size_t)i * DN_ES + blk * 128 + 4 * lane + j]; + fwht128(v, lane); + float* o = y + ((size_t)(p >> 3) * ys + (p & 7)) * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); + } + } + } +} + +// Program (row): the routed picks of a row of `slots` (the shared expert last) as a row of `slots - 1`, for the plan. +extern "C" __global__ void glm_exl3_picks(const int* __restrict__ pick, int* __restrict__ out, int rows, int slots) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= rows * (slots - 1)) return; + out[i] = pick[(i / (slots - 1)) * slots + i % (slots - 1)]; +} diff --git a/zig/kernels/cuda/glm_exl3_transposed.cuh b/zig/kernels/cuda/glm_exl3_transposed.cuh new file mode 100644 index 0000000000..3e7919e84a --- /dev/null +++ b/zig/kernels/cuda/glm_exl3_transposed.cuh @@ -0,0 +1,168 @@ +// Transposed prompt helpers after MiaAI-Lab (GLM-5.3-Flash-EXL3-2x-DGX-Sparks-TensorFold 0004, 0009, Apache-2.0). +#pragma once +// Included by exl3.cu after its shared helpers (mcg decode, cp.async, ldmatrix), in the same translation unit. + +namespace { + +constexpr int TPM = 64, TNW = 8, TMT = 16 / TNW, TNB = TPM / 8, TT = TNW * 32, TPH = TPM / 2; +constexpr int TBK = 32, TKT = TBK / 16, TSTAGES = 3, TXS = TBK + 8; // a stage: 32 k values a row; +16 B: ldmatrix rows on distinct banks +constexpr int TWORDS = 16 * TKT * 32; // a stage's trellis words (16 tiles x 2 k tiles x 32) +constexpr int GU_ES = 128 + 4, DN_ES = 256 + 4; // the epilogue sums' row strides (floats) +template +constexpr int tstage_halves() { return XR * TXS + 2 * TWORDS; } +constexpr int tsmem_bytes() { + const int gu = TSTAGES * tstage_halves<2 * TPM>() * 2; // gate and up rows apart (no shared rotation) + const int ep = 2 * TPH * GU_ES * 4, dn = TPH * DN_ES * 4; + return gu > ep ? (gu > dn ? gu : dn) : (ep > dn ? ep : dn); +} + +// Which (matrix, weight tile) warp-tile q (0..15) of column block cb is. +template +__device__ __forceinline__ void ttile(int q, int cb, int& mat, int& nt) { + mat = MATS == 2 ? q >> 3 : 0; + nt = MATS == 2 ? cb * 8 + (q & 7) : cb * 16 + q; +} + +// Main loop: acc[i][nb] = warp tile i x members nb * 8 .. + 7 over K; XR rows a stage (TPM if X0 == X1, else 2 TPM). +template +__device__ __forceinline__ void tmain(const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, const int* rows_sh, int cnt, int e, int K, int NTILES, int cb, int xshift, + half* stage, float (&acc)[TMT][TNB][4]) { + constexpr int SH = tstage_halves(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + int mat, nt_unused; + ttile(warp * TMT, cb, mat, nt_unused); + const int KT = K >> 4, chunks = K / TBK; + const int nblocks = (cnt + 15) >> 4; // 16-member blocks in use +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int j = 0; j < TNB; ++j) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][j][c] = 0.f; + auto load = [&](int chunk) { + half* dst = stage + (size_t)(chunk % TSTAGES) * SH; + const int k0 = chunk * TBK; + for (int idx = threadIdx.x; idx < XR * (TBK / 8); idx += TT) { + const int m = idx / (TPM * (TBK / 8)), rem = idx % (TPM * (TBK / 8)), i = rem / (TBK / 8), j = rem % (TBK / 8); + if (i >= nblocks * 16) continue; // a 16-member block without a member: never read + const int r = rows_sh[i]; + const half* src = (m ? X1 : X0) + (size_t)(r < 0 ? 0 : (r >> xshift)) * K + k0 + j * 8; + cp_async16(dst + ((size_t)m * TPM + i) * TXS + j * 8, src, r >= 0); + } + uint32_t* wdst = reinterpret_cast(dst + (size_t)XR * TXS); + for (int idx = threadIdx.x; idx < 16 * TKT * 8; idx += TT) { // 16-byte pieces of the 128-byte tile words + const int q = idx / (TKT * 8), kk = (idx / 8) % TKT, part = idx % 8; + int wm, nt; + ttile(q, cb, wm, nt); + const uint32_t* src = (wm ? T1 : T0) + (((size_t)e * KT + chunk * TKT + kk) * NTILES + nt) * 32 + part * 4; + cp_async16(wdst + ((size_t)q * TKT + kk) * 32 + part * 4, src, true); + } + }; +#pragma unroll + for (int c = 0; c < TSTAGES - 1; ++c) { + if (c < chunks) load(c); + cp_async_commit(); + } + for (int c = 0; c < chunks; ++c) { + cp_async_wait(); + __syncthreads(); + if (c + TSTAGES - 1 < chunks) load(c + TSTAGES - 1); + cp_async_commit(); + const half* base = stage + (size_t)(c % TSTAGES) * SH; + const half* xs = base + (size_t)(XR == TPM ? 0 : mat) * TPM * TXS; + const uint32_t* ws = reinterpret_cast(base + (size_t)XR * TXS); +#pragma unroll + for (int kk = 0; kk < TKT; ++kk) { + uint32_t a[TMT][4]; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile_p(ws[((size_t)(warp * TMT + i) * TKT + kk) * 32 + lane], lane, b0, b1); + a[i][0] = b0[0]; a[i][1] = b1[0]; a[i][2] = b0[1]; a[i][3] = b1[1]; + } +#pragma unroll + for (int jb = 0; jb < TPM / 16; ++jb) { + if (jb < nblocks) { + uint32_t r[4]; + const int row = jb * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + ldmatrix_x4(r, xs + (size_t)row * TXS + kk * 16 + ((lane >> 3) & 1) * 8); + const uint32_t bl[2] = {r[0], r[1]}, bh[2] = {r[2], r[3]}; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + mma16816(acc[i][2 * jb], a[i], bl); + mma16816(acc[i][2 * jb + 1], a[i], bh); + } + } + } + } + } + cp_async_wait<0>(); + __syncthreads(); // the stages become the epilogue's sums +} + +// acc (members TPH * hf .. + TPH - 1) -> E[member - TPH * hf][col] fp32, row stride ES; col0: first tile's column. +__device__ __forceinline__ void tspill(float* E, int ES, int col0, const float (&acc)[TMT][TNB][4], int hf) { + const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int q = 0; q < TPM / 16; ++q) { + const int nb = hf * (TPM / 16) + q; + const int tok = q * 8 + 2 * t, col = col0 + i * 16 + g; + E[(size_t)tok * ES + col] = acc[i][nb][0]; + E[(size_t)(tok + 1) * ES + col] = acc[i][nb][1]; + E[(size_t)tok * ES + col + 8] = acc[i][nb][2]; + E[(size_t)(tok + 1) * ES + col + 8] = acc[i][nb][3]; + } +} + +template +__device__ __forceinline__ void tgateup(const half* __restrict__ Xg, const half* __restrict__ Xu, const uint32_t* __restrict__ Tg, + const uint32_t* __restrict__ Tu, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, + const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, + int K, int N, int xshift, float limit, unsigned char* smem_raw, int* rows_sh) { + const int cb = blockIdx.x, item = blockIdx.y; + if (item >= counts[0]) return; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + float acc[TMT][TNB][4]; + tmain<2, XR>(Xg, Xu, Tg, Tu, rows_sh, cnt, e, K, N >> 4, cb, xshift, reinterpret_cast(smem_raw), acc); + float* Eg = reinterpret_cast(smem_raw); + int mat, nt; + ttile<2>(warp * TMT, cb, mat, nt); + const int col0 = (nt - cb * 8) * 16; + const int n = cb * 128 + 4 * lane; + for (int hf = 0; hf < 2 && hf * TPH < cnt; ++hf) { + if (hf) __syncthreads(); // the first half's rows are read + tspill(Eg + (size_t)mat * TPH * GU_ES, GU_ES, col0, acc, hf); + __syncthreads(); + for (int i = warp; i < TPH && hf * TPH + i < cnt; i += TNW) { + const int p = rows_sh[hf * TPH + i]; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + gv[j] = Eg[(size_t)i * GU_ES + 4 * lane + j]; + uv[j] = Eg[(size_t)(TPH + i) * GU_ES + 4 * lane + j]; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float gg = clamp_hi(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = clamp_both(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), limit); + float act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); + } + } +} + +} // namespace diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 4301d0e27f..6368e28757 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -39,6 +39,7 @@ pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes else &.{}; pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; +pub const glm_exl3: []const u8 = if (available) &Blob("fatbin_glm_exl3").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; diff --git a/zig/src/families/glm/config.zig b/zig/src/families/glm/config.zig index 2b23df7bc4..f7e9cd27c3 100644 --- a/zig/src/families/glm/config.zig +++ b/zig/src/families/glm/config.zig @@ -14,6 +14,7 @@ pub const Config = struct { run: u32 = 45, // backbone layers loaded and run: all, or the first few for a check that must fit (GLM_LAYERS) vocab: u32 = 154880, head_unit: u32 = 64, // the vocabulary's TP cut in whole units: 128 for EXL3, whose matrices split at 128 columns + exl3: bool = false, // routed experts in EXL3 (quant_method "exl3", Mia-AiLab's checkpoint), the rest bf16 eps: f32 = 1e-5, dense_layers: u32 = 3, dense_inter: u32 = 12288, // this Mac's dense MLP rows: all, or its half in TP2 @@ -151,6 +152,7 @@ pub fn parse(gpa: std.mem.Allocator, json: []const u8) !Config { c.vocab = @intCast(try int(t, "vocab_size")); if (root.get("quantization_config")) |q| if (q == .object) if (q.object.get("quant_method")) |m| if (m == .string and std.mem.eql(u8, m.string, "exl3")) { c.head_unit = 128; + c.exl3 = true; }; c.eps = @floatCast(try float(t, "rms_norm_eps")); c.dense_layers = @intCast(try int(t, "first_k_dense_replace")); diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index af9254de61..aeefc3ffee 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -1,6 +1,7 @@ //! GLM-5.3-Flash on CUDA: the family's parts, ported from the Python family's CUDA path. pub const experts = @import("cuda_experts.zig"); +pub const exl3 = @import("cuda_exl3.zig"); pub const glue = @import("cuda_glue.zig"); pub const kda = @import("cuda_kda.zig"); pub const b16 = @import("cuda_b16.zig"); @@ -25,6 +26,7 @@ pub const native = @import("cuda_native.zig"); test { _ = experts; + _ = exl3; _ = glue; _ = kda; _ = b16; diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 65308ad312..f42cf02e18 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -11,6 +11,7 @@ const ranks = @import("cuda_ranks.zig"); const sampler = @import("cuda_sampler.zig"); const heat = @import("heat"); const modelopt = @import("core").modelopt; +const exl3 = @import("cuda_exl3.zig"); pub const Options = struct { capacity: usize, // tokens a sequence holds @@ -101,7 +102,8 @@ pub const Engine = struct { const cfg_path = try std.fs.path.join(a, &.{ checkpoint, "config.json" }); const text = try std.Io.Dir.cwd().readFileAlloc(io, cfg_path, a, .limited(1 << 24)); var why: modelopt.Why = .{}; - modelopt.admit(gpa, text, &why) catch |err| { // a quant_algo other than NVFP4 is refused before anything loads + const admitted = if (exl3.isExl3(gpa, text)) exl3.admit(gpa, text, &why) else modelopt.admit(gpa, text, &why); + admitted catch |err| { // anything but ModelOpt NVFP4 or EXL3 (4-bit mcg routed experts) is refused before loading std.log.err("{s}: {s}", .{ cfg_path, why.text() }); return err; }; diff --git a/zig/src/families/glm/cuda_exl3.zig b/zig/src/families/glm/cuda_exl3.zig new file mode 100644 index 0000000000..05e380a7fa --- /dev/null +++ b/zig/src/families/glm/cuda_exl3.zig @@ -0,0 +1,200 @@ +//! GLM-5.3-Flash's EXL3 routed experts on CUDA (Mia-AiLab's checkpoint: 4-bit mcg trellis experts, the rest bf16). + +const std = @import("std"); +const modelopt = @import("core").modelopt; +const Span = @import("cuda_shard.zig").Span; +const cuda = @import("cuda"); +const experts = @import("cuda_experts.zig"); + +/// config.json's quantization_config (top level, else text_config's), when it has one. +fn block(root: std.json.Value) ?std.json.ObjectMap { + if (root != .object) return null; + if (root.object.get("quantization_config")) |q| if (q == .object) return q.object; + if (root.object.get("text_config")) |t| if (t == .object) if (t.object.get("quantization_config")) |q| if (q == .object) return q.object; + return null; +} + +/// Whether config.json names EXL3 (quant_method "exl3"): `admit` then says whether this family reads it. +pub fn isExl3(gpa: std.mem.Allocator, config_json: []const u8) bool { + const parsed = std.json.parseFromSlice(std.json.Value, gpa, config_json, .{}) catch return false; + defer parsed.deinit(); + const q = block(parsed.value) orelse return false; + const m = q.get("quant_method") orelse return false; + return m == .string and std.mem.eql(u8, m.string, "exl3"); +} + +/// The EXL3 checkpoints the kernels read: 4-bit routed experts on the "mcg" codebook (everything else bf16). +pub fn admit(gpa: std.mem.Allocator, config_json: []const u8, why: *modelopt.Why) !void { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, config_json, .{}); + defer parsed.deinit(); + const q = block(parsed.value) orelse { + why.set("no quantization_config", .{}); + return error.UnsupportedQuant; + }; + const bits = q.get("bits") orelse std.json.Value{ .null = {} }; + if (bits != .integer or bits.integer != 4) { + why.set("EXL3 at {f} bits: the CUDA family reads 4-bit EXL3 experts", .{std.json.fmt(bits, .{})}); + return error.UnsupportedQuant; + } + const book = q.get("codebook") orelse std.json.Value{ .null = {} }; + if (book != .string or !std.mem.eql(u8, book.string, "mcg")) { + why.set("EXL3 codebook {f}: the CUDA family reads \"mcg\"", .{std.json.fmt(book, .{})}); + return error.UnsupportedQuant; + } +} + +/// One layer's EXL3 experts on a rank: trellis words [E][K/16][N/16][32] and fp16 scales (suh [E][K], svh [E][N]). +pub const Experts = struct { + gt: u64, + ut: u64, + dt: u64, + suh_g: u64, + suh_u: u64, + svh_g: u64, + svh_u: u64, + suh_d: u64, + svh_d: u64, + shx: bool, // every expert's gate and up share one input rotation: a prompt rotates a row once, not once a pick +}; + +const tile = 128; // a 16 x 16 tile's bytes at 4 bits + +/// Layer `i`'s EXL3 experts through `L` (cuda_weights' loader): gate, up by `cut`'s tile columns, down by its rows. +pub fn load(L: anytype, i: u32, count: usize, dims: usize, inter: usize, cut: Span) !Experts { + if (dims % 128 != 0 or cut.lo % 128 != 0 or cut.hi % 128 != 0) return error.NotExl3Mcg4; + const width = cut.len(); + var out: Experts = undefined; + out.shx = true; + const Mat = struct { k: usize, n: usize, by_cols: bool }; + const mats = [_]Mat{ + .{ .k = dims, .n = inter, .by_cols = true }, + .{ .k = dims, .n = inter, .by_cols = true }, + .{ .k = inter, .n = dims, .by_cols = false }, + }; + var first_suh: ?[]const u8 = null; + inline for (mats, .{ "gate_proj", "up_proj", "down_proj" }, .{ "gt", "ut", "dt" }, .{ "suh_g", "suh_u", "suh_d" }, .{ "svh_g", "svh_u", "svh_d" }) |m, proj, tw, su, sv| { + const kt = (if (m.by_cols) m.k else width) / 16; // the rank's tiles + const nt = (if (m.by_cols) width else m.n) / 16; + for (0..count) |e| { + const t = try L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".trellis", .{ i, e }); + if (t.dtype != .i16 or t.rank != 3 or t.dim(0) != m.k / 16 or t.dim(1) != m.n / 16 or t.dim(2) != 64) return error.NotExl3Mcg4; + _ = L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".mcg", .{ i, e }) catch return error.NotExl3Mcg4; + if (m.by_cols) { // every tile row, the rank's tile columns + const row = (m.n / 16) * tile; + for (0..kt) |r| try L.host.appendSlice(L.gpa, t.bytes[r * row + (cut.lo / 16) * tile ..][0 .. nt * tile]); + } else try L.host.appendSlice(L.gpa, t.bytes[(cut.lo / 16) * nt * tile ..][0 .. kt * nt * tile]); // its tile rows + } + @field(out, tw) = try L.put("layers.{d}.moe.exl3." ++ tw, .{i}); + for (0..count) |e| { // suh: the input's scales (gate, up: every input; down: the rank's) + const t = try L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".suh", .{ i, e }); + if (t.dtype != .f16 or t.numel() != m.k) return error.NotExl3Mcg4; + const mine = if (m.by_cols) t.bytes else t.bytes[cut.lo * 2 .. cut.hi * 2]; + if (m.by_cols) { // one rotation shared by the layer's every gate and up input? + if (first_suh) |f| (if (!std.mem.eql(u8, f, mine)) { + out.shx = false; + }) else first_suh = mine; + } + try L.host.appendSlice(L.gpa, mine); + } + @field(out, su) = try L.put("layers.{d}.moe.exl3." ++ su, .{i}); + for (0..count) |e| { // svh: the output's scales (gate, up: the rank's; down: every output) + const t = try L.get("layers.{d}.mlp.experts.{d}." ++ proj ++ ".svh", .{ i, e }); + if (t.dtype != .f16 or t.numel() != m.n) return error.NotExl3Mcg4; + try L.host.appendSlice(L.gpa, if (m.by_cols) t.bytes[cut.lo * 2 .. cut.hi * 2] else t.bytes); + } + @field(out, sv) = try L.put("layers.{d}.moe.exl3." ++ sv, .{i}); + } + return out; +} + +/// glm_exl3's kernels. +pub const Kernels = struct { + rot_rows: cuda.Function, + rot_in: cuda.Function, + grouped: cuda.Function, + gateup: cuda.Function, + down: cuda.Function, + prompt_gu: cuda.Function, + prompt_dn: cuda.Function, + picks: cuda.Function, + + pub fn load(m: cuda.Module) !Kernels { + var k: Kernels = undefined; + inline for (@typeInfo(Kernels).@"struct".field_names) |name| @field(k, name) = try m.function("glm_exl3_" ++ name); + return k; + } +}; + +/// Scratch beside the MoE block's (its eact holds the down inputs): 8-slot picks, rotated inputs, K-split partials. +pub const Scratch = struct { picks: u64, xg: u64, xu: u64, z: u64 }; + +/// The picks a row holds (the kernels' pair = row * 8 + slot). +pub const top_k = 8; +/// Gate and up K slices in a decode window (down: 1); prompt chunks take the fused kernels, no partials. +pub const decode_splits = 4; + +/// Bytes of each Scratch buffer for windows of `rows` rows: (picks, xg, xu, z); `shx`: one input rotation a layer. +pub fn scratchBytes(rows: usize, dims: usize, width: usize, prefill: bool, shx: bool) [4]usize { + const pairs = rows * top_k; + const one = prefill and shx; // a prompt rotates each row once + return .{ + pairs * 4, + (if (one) rows else pairs) * dims * 2, + if (one) 16 else pairs * dims * 2, + if (prefill) 16 else @max(2 * decode_splits * pairs * width, pairs * dims) * 4, + }; +} + +/// One launch's arguments, then the launch. +fn go(f: cuda.Function, stream: cuda.Stream, grid: cuda.Dim3, threads: u32, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = threads } }, stream, &a); +} + +/// Routed slots of `rows` rows (picks [rows][slots], shared last) into `ey` [rows][slots][dims] fp32, shared untouched. +pub fn run(k: *const Kernels, ek: *const experts.Kernels, stream: cuda.Stream, w: Experts, sc: Scratch, plan_buf: cuda.grouped.Plan, x: u64, x_stride: usize, picks: u64, slots: usize, rows: usize, count: usize, dims: usize, width: usize, limit: f32, prompt: bool, ey: u64, xd: u64) !void { + if (slots != top_k + 1) return error.UnsupportedTopK; + if (dims % 256 != 0 or width % 128 != 0) return error.NotExl3Mcg4; + const P = rows * top_k; + const D: i32 = @intCast(dims); + const N: i32 = @intCast(width); + const ys: i32 = @intCast(slots); + try go(k.picks, stream, .{ .x = @intCast((P + 255) / 256) }, 256, .{ picks, sc.picks, @as(i32, @intCast(rows)), @as(i32, @intCast(slots)) }); + const plan: experts.Plan = .{ .route = plan_buf, .slots = top_k, .experts = count, .prompt = prompt }; + try experts.route(ek, stream, sc.picks, P, plan); + const items: u32 = @intCast(experts.maxItems(P, count, if (prompt) experts.staged_tile else experts.tile)); + const r = plan_buf; + const one = prompt and w.shx; // the prompt kernels read row pair >> 3: one rotated row a token + if (one) { + try go(k.rot_rows, stream, .{ .x = @intCast(rows), .y = @intCast(dims / 128) }, 32, .{ x, @as(i32, @intCast(x_stride)), w.suh_g, sc.xg, D }); + } else { + try go(k.rot_in, stream, .{ .x = @intCast(P), .y = @intCast(dims / 128), .z = 2 }, 32, .{ x, @as(i32, @intCast(x_stride)), sc.picks, w.suh_g, w.suh_u, sc.xg, sc.xu, D }); + } + if (prompt) { // gate and up with their epilogue into xd; down with its epilogue into ey + try go(k.prompt_gu, stream, .{ .x = @intCast(width / 128), .y = items }, 256, .{ sc.xg, if (one) sc.xg else sc.xu, w.gt, w.ut, r.items, r.counts, r.members, w.svh_g, w.svh_u, w.suh_d, xd, D, N, @as(i32, if (one) 3 else 0), limit }); + return go(k.prompt_dn, stream, .{ .x = @intCast(dims / 256), .y = items }, 256, .{ xd, w.dt, r.items, r.counts, r.members, w.svh_d, ey, N, D, ys }); + } + const sk: i32 = decode_splits; + const nt = 8; // n tiles a block + try go(k.grouped, stream, .{ .x = items, .y = @intCast(width / (16 * nt)), .z = @intCast(2 * sk) }, 128, .{ sc.xg, sc.xu, w.gt, w.ut, r.items, r.counts, r.members, sc.z, D, N, @as(i32, @intCast(P)), sk }); + try go(k.gateup, stream, .{ .x = @intCast(P), .y = @intCast(width / 128) }, 32, .{ sc.z, sc.picks, w.svh_g, w.svh_u, w.suh_d, xd, @as(i32, @intCast(P)), N, sk, limit }); + try go(k.grouped, stream, .{ .x = items, .y = @intCast(dims / (16 * nt)), .z = 1 }, 128, .{ xd, xd, w.dt, w.dt, r.items, r.counts, r.members, sc.z, N, D, @as(i32, @intCast(P)), @as(i32, 1) }); + try go(k.down, stream, .{ .x = @intCast(P), .y = @intCast(dims / 128) }, 32, .{ sc.z, sc.picks, w.svh_d, ey, @as(i32, @intCast(P)), D, @as(i32, 1), ys }); +} + +test "EXL3 admission: Mia-AiLab's block passes; other widths and codebooks are refused by name" { + const gpa = std.testing.allocator; + const mia = + \\{"quantization_config": {"quant_method": "exl3", "bits": 4, "codebook": "mcg", "head_bits": 16}} + ; + try std.testing.expect(isExl3(gpa, mia)); + var why: modelopt.Why = .{}; + try admit(gpa, mia, &why); + try std.testing.expect(!isExl3(gpa, "{\"quantization_config\": {\"quant_method\": \"modelopt\"}}")); + try std.testing.expect(!isExl3(gpa, "{}")); + try std.testing.expectError(error.UnsupportedQuant, admit(gpa, "{\"quantization_config\": {\"quant_method\": \"exl3\", \"bits\": 3, \"codebook\": \"mcg\"}}", &why)); + try std.testing.expectEqualStrings("EXL3 at 3 bits: the CUDA family reads 4-bit EXL3 experts", why.text()); + try std.testing.expectError(error.UnsupportedQuant, admit(gpa, "{\"text_config\": {\"quantization_config\": {\"quant_method\": \"exl3\", \"bits\": 4, \"codebook\": \"mul1\"}}}", &why)); + try std.testing.expectEqualStrings("EXL3 codebook \"mul1\": the CUDA family reads \"mcg\"", why.text()); +} diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index c93ab0079c..893677e276 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/zig/src/families/glm/cuda_forward.zig @@ -17,6 +17,7 @@ const moe = @import("cuda_moe.zig"); const quant = @import("cuda_quant.zig"); const q4 = @import("cuda_q4.zig"); const state = @import("cuda_state.zig"); +const exl3 = @import("cuda_exl3.zig"); const shared = @import("cuda_shared.zig"); pub const prefill_rows = state.prefill_rows; @@ -27,11 +28,12 @@ pub const Buffers = state.Buffers; /// The modules a forward launches from: our fatbins. pub const Modules = struct { glue: glue.Kernels, - mods: [7]cuda.Module, + mods: [8]cuda.Module, kda: kda.Kernels, lane: cuda.qmmf.Lane, // the NVFP4 linears' decode rows (qlinear) prompt: cuda.nvfp4.Prompt, // and their prompt rows experts: experts.Kernels, + exl3: exl3.Kernels, // an EXL3 checkpoint's routed experts copy_rows: cuda.Function, put_slot: cuda.Function, pack_dense: cuda.Function, @@ -43,7 +45,7 @@ pub const Modules = struct { var m: Modules = undefined; var n: usize = 0; errdefer for (m.mods[0..n]) |*x| x.unload(); - for (&m.mods, [_][]const u8{ k.glm_kda, k.glm_ops, k.experts, k.affine4_pack, k.glm_quant, k.qmm_group, k.glm_glue }) |*x, image| { + for (&m.mods, [_][]const u8{ k.glm_kda, k.glm_ops, k.experts, k.affine4_pack, k.glm_quant, k.qmm_group, k.glm_glue, k.glm_exl3 }) |*x, image| { x.* = try cuda.Module.load(d, image); n += 1; } @@ -60,6 +62,7 @@ pub const Modules = struct { m.quant = try quant.Kernels.load(m.mods[4]); m.q4 = try q4.Kernels.load(m.mods[5], m.mods[3], ctx); m.glue = try glue.Kernels.load(m.mods[6]); + m.exl3 = try exl3.Kernels.load(m.mods[7]); return m; } @@ -250,7 +253,7 @@ pub const Forward = struct { } fn moeKernels(f: *const Forward) moe.Kernels { - return .{ .set = f.set(), .experts = &f.m.experts, .put_slot = f.m.put_slot }; + return .{ .set = f.set(), .experts = &f.m.experts, .exl3 = &f.m.exl3, .put_slot = f.m.put_slot }; } const Site = enum { mixer, ffn }; diff --git a/zig/src/families/glm/cuda_mlp.zig b/zig/src/families/glm/cuda_mlp.zig index 8e8d5eba11..a9e55254b9 100644 --- a/zig/src/families/glm/cuda_mlp.zig +++ b/zig/src/families/glm/cuda_mlp.zig @@ -3,12 +3,19 @@ const std = @import("std"); const cuda = @import("cuda"); const glue = @import("cuda_glue.zig"); +const b16 = @import("cuda_b16.zig"); -/// One dense layer on a rank: gate and up [width][dims] and down [dims][width], NVFP4 as Fp4Linear holds them. -pub const Layer = struct { gate: cuda.qmmf.Weight, up: cuda.qmmf.Weight, down: cuda.qmmf.Weight, width: usize, dims: usize, limit: f32 }; +/// A dense layer's weights: NVFP4 as Fp4Linear holds them, or the checkpoint's bf16 (an EXL3 checkpoint's). +pub const Weights = union(enum) { + nvfp4: struct { gate: cuda.qmmf.Weight, up: cuda.qmmf.Weight, down: cuda.qmmf.Weight }, + bf16: struct { gu: u64, down: u64 }, // [gate | up] rows [2 width][dims], down [dims][width] +}; -/// The block's scratch for up to `rows` rows: one projection's rows, the [gate | up] rows, the activation and its sums. -pub const Scratch = struct { tmp: u64, gu: u64, act: u64, xs_act: u64 }; +/// One dense layer on a rank: gate and up [width][dims] and down [dims][width]. +pub const Layer = struct { w: Weights, width: usize, dims: usize, limit: f32 }; + +/// The block's scratch for `rows` rows: a projection's rows, [gate | up], the activation, its sums, split-K partials. +pub const Scratch = struct { tmp: u64, gu: u64, act: u64, xs_act: u64, mm_part: u64 }; pub const Kernels = struct { set: glue.Set, linear: cuda.qlinear.Linear, copy_rows: cuda.Function }; @@ -31,13 +38,23 @@ pub fn block(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_st /// forward.out_proj's fill for the dense MLP: down of rows lo..hi (s.act) into out [hi - lo][dims] fp32. pub fn fill(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, lo: usize, hi: usize, prompt: bool, out: u64) !void { - try mm(k, stream, l.down, s.act + lo * l.width * 2, l.width, hi - lo, out, true, prompt); + switch (l.w) { + .nvfp4 => |w| try mm(k, stream, w.down, s.act + lo * l.width * 2, l.width, hi - lo, out, true, prompt), + .bf16 => |w| try b16.matmul(k.set, stream, s.act + lo * l.width * 2, l.width, w.down, out, s.mm_part, hi - lo, l.dims, l.width, true), + } } /// The block up to its down projection: [gate | up] and SwiGLU into s.act. pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool) !void { const w = l.width; - for ([_]cuda.qmmf.Weight{ l.gate, l.up }, 0..) |proj, j| { + const q = switch (l.w) { + .nvfp4 => |q| q, + .bf16 => |b| { // [gate | up] in one matmul, as the shared expert's + try b16.matmul(k.set, stream, x, x_stride, b.gu, s.gu, s.mm_part, rows, 2 * w, l.dims, false); + return glue.swiglu(k.set, stream, s.gu, s.act, s.xs_act, l.limit, rows, w); + }, + }; + for ([_]cuda.qmmf.Weight{ q.gate, q.up }, 0..) |proj, j| { // Fp4Linear(out=a column slice): the rows computed contiguous, then copied into their half of each row try mm(k, stream, proj, x, x_stride, rows, s.tmp, false, prompt); var a: cuda.Args = .{}; diff --git a/zig/src/families/glm/cuda_moe.zig b/zig/src/families/glm/cuda_moe.zig index ec860a188e..d50e8707c1 100644 --- a/zig/src/families/glm/cuda_moe.zig +++ b/zig/src/families/glm/cuda_moe.zig @@ -5,12 +5,14 @@ const cuda = @import("cuda"); const glue = @import("cuda_glue.zig"); const experts = @import("cuda_experts.zig"); const b16 = @import("cuda_b16.zig"); +const exl3 = @import("cuda_exl3.zig"); /// One MoE layer on a rank: router (bf16 [count][dims]) and fp32 bias, routed experts, shared [gate | up] and down. pub const Layer = struct { router: u64, bias: u64, - routed: cuda.experts.Layer, + routed: cuda.experts.Layer, // NVFP4 experts, or only the shape (width, dims, limit) when `exl3` holds them + exl3: ?exl3.Experts = null, count: usize, shared_gu: u64, shared_down: u64, @@ -34,10 +36,11 @@ pub const Scratch = struct { sxs: u64, // fp32 [rows][shared / 64] sy: u64, // fp32 [rows][dims] mm_part: u64, // fp32 split-K partials of the shared expert's matmuls + x3: exl3.Scratch = .{ .picks = 0, .xg = 0, .xu = 0, .z = 0 }, // EXL3 experts (eact holds their down inputs) }; /// The kernels the block runs: the captured Triton set, the grouped NVFP4 experts, the plan, and glm_ops. -pub const Kernels = struct { set: glue.Set, experts: *const experts.Kernels, put_slot: cuda.Function }; +pub const Kernels = struct { set: glue.Set, experts: *const experts.Kernels, exl3: *const exl3.Kernels, put_slot: cuda.Function }; /// forward.moe_block: out [rows][dims] fp32 partial; `prompt`: 64-pair items, the prompt kernels, bf16 slot rows. pub fn block(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool, out: u64) !void { @@ -52,6 +55,7 @@ pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const shared_id: i32 = @intCast(l.count); try glue.router(k.set, stream, x, x_stride, l.router, s.mlog, s.router_part, rows, d, l.count); try glue.select(k.set, stream, s.mlog, l.bias, s.pick, s.wts, rows, l.top_k, l.count, l.scale, l.norm); + if (l.exl3) |w| return exl3.run(k.exl3, k.experts, stream, w, s.x3, s.plan, x, x_stride, s.pick, slots, rows, l.count, d, l.routed.width, l.routed.limit, prompt, s.ey, s.eact); const plan: experts.Plan = .{ .route = s.plan, .slots = slots, .experts = l.count + 1, .prompt = prompt }; try experts.route(k.experts, stream, s.pick, rows * slots, plan); try experts.run(k.experts, stream, .gate_up, l.routed, plan, x, x_stride, rows, s.eact, shared_id); @@ -63,14 +67,15 @@ pub fn fill(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const d = l.routed.dims; const slots = l.top_k + 1; const n = hi - lo; - const eb: usize = if (prompt) 2 else 4; // the slot rows: bf16 in a prompt chunk + const bf16_slots = prompt and l.exl3 == null; // the slot rows: bf16 in an NVFP4 prompt chunk, else fp32 + const eb: usize = if (bf16_slots) 2 else 4; try b16.matmul(k.set, stream, x + lo * x_stride * 2, x_stride, l.shared_gu, s.sgu + lo * 2 * l.shared * 2, s.mm_part, n, 2 * l.shared, d, false); try glue.swiglu(k.set, stream, s.sgu + lo * 2 * l.shared * 2, s.sact + lo * l.shared * 2, s.sxs + lo * (l.shared / 64) * 4, l.routed.limit, n, l.shared); try b16.matmul(k.set, stream, s.sact + lo * l.shared * 2, l.shared, l.shared_down, s.sy + lo * d * 4, s.mm_part, n, d, l.shared, true); var a: cuda.Args = .{}; a.add(s.sy + lo * d * 4); a.add(s.ey + lo * slots * d * eb); - for ([_]usize{ n, d, slots, l.top_k, @intFromBool(prompt) }) |v| a.add(@as(c_int, @intCast(v))); + for ([_]usize{ n, d, slots, l.top_k, @intFromBool(bf16_slots) }) |v| a.add(@as(c_int, @intCast(v))); try cuda.launch.launch(k.put_slot, .{ .grid = .{ .x = @intCast((n * d + 255) / 256) }, .block = .{ .x = 256 } }, stream, &a); - try glue.combine(k.set, stream, s.ey + lo * slots * d * eb, !prompt, s.wts + lo * slots * 4, out, n, d, slots); + try glue.combine(k.set, stream, s.ey + lo * slots * d * eb, !bf16_slots, s.wts + lo * slots * 4, out, n, d, slots); } diff --git a/zig/src/families/glm/cuda_shard.zig b/zig/src/families/glm/cuda_shard.zig index 6b00faac54..812cfb508a 100644 --- a/zig/src/families/glm/cuda_shard.zig +++ b/zig/src/families/glm/cuda_shard.zig @@ -35,7 +35,7 @@ pub const Share = struct { .kda = try span(c.kda_heads, 1, rank, world), .mla = try span(c.mla_heads, 1, rank, world), .dense = try span(c.dense_inter, unit, rank, world), - .moe = try span(c.moe_inter, unit, rank, world), + .moe = try span(c.moe_inter, if (c.exl3) 128 else unit, rank, world), // EXL3 tiles and rotations: 128 .shared = try span(c.shared_inter, unit, rank, world), .vocab = try span(c.vocab, c.head_unit, rank, world), }; diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig index 32559505c3..b90123aa15 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -2,6 +2,7 @@ const std = @import("std"); const cuda = @import("cuda"); +const exl3 = @import("cuda_exl3.zig"); const weights = @import("cuda_weights.zig"); const kda = @import("cuda_kda.zig"); const b16 = @import("cuda_b16.zig"); @@ -46,7 +47,7 @@ fn arena(d: *const cuda.Driver, sizes: []const usize, out: []u64) !Arena { } /// A rank's B16 matmul shapes ([n, k]) whose split-K partials the buffers hold. -fn b16Shapes(w: *const weights.Weights) [10][2]usize { +fn b16Shapes(w: *const weights.Weights) [12][2]usize { const c = &w.config; const s = &w.share; const n = s.kda.len() * c.kda_dim; @@ -56,6 +57,8 @@ fn b16Shapes(w: *const weights.Weights) [10][2]usize { .{ c.q_lora + c.kv_lora, c.hidden }, .{ s.mla.len() * c.nope, c.q_lora }, .{ c.i_dim + c.i_heads, c.hidden }, .{ c.i_heads * c.i_dim, c.q_lora }, .{ c.hidden, s.mla.len() * c.v_dim }, .{ 2 * s.shared.len(), c.hidden }, .{ c.hidden, 2 * c.hidden }, // the MTP head's eh_proj + if (c.exl3) .{ 2 * s.dense.len(), c.hidden } else .{ c.hidden, 2 * c.hidden }, // bf16 dense layers: [gate | up] + if (c.exl3) .{ c.hidden, s.dense.len() } else .{ c.hidden, 2 * c.hidden }, // and down }; } @@ -123,6 +126,7 @@ pub const Buffers = struct { var mm: usize = 4; // windows under rows_from rows keep their K slices apart (a prompt's tail chunk on the prefill buffers too) for (b16Shapes(w)) |sh_| mm = @max(mm, b16.partBytes(@min(rows, b16.rows_from - 1), sh_[0], sh_[1])); + const x3: [4]usize = if (c.exl3) exl3.scratchBytes(rows, D, ew, prefill, w.exl3_shx) else @splat(4); const hist_len = if (pairs > cuda.grouped.small) (pairs + cuda.grouped.small - 1) / cuda.grouped.small * (count + 1) else 1; const sizes = [_]usize{ rows * 4, rows * S * D * 2, rows * D * 2, rows * (D / 64) * 4, @@ -160,10 +164,11 @@ pub const Buffers = struct { rows * (D / 64) * 4, (if (prefill) 1 else rows) * s.vocab.len() * 2, rows * D * 2, rows * D * 2, rows * 2 * D * 2, rows * D * 2, // hin, me, mcat, mx rows * D * 4, // part: this rank's fp32 partial rows a split exchange sends from + x3[0], x3[1], x3[2], x3[3], // EXL3 experts: 8-slot picks, rotated inputs, a decode window's partials }; var p: [sizes.len]u64 = undefined; const a = try arena(d, &sizes, &p); - return .{ + var out: Buffers = .{ .arena = a, .rows = rows, .prefill = prefill, @@ -201,7 +206,7 @@ pub const Buffers = struct { .mm_part = p[42], .at = p[43], }, - .dense = .{ .tmp = p[44], .gu = p[45], .act = p[46], .xs_act = p[47] }, + .dense = .{ .tmp = p[44], .gu = p[45], .act = p[46], .xs_act = p[47], .mm_part = p[42] }, .moe = .{ .mlog = p[48], .router_part = p[26], .pick = p[49], .wts = p[50], .plan = .{ .members = p[51], .items = p[52], .counts = p[53], .rank = p[54], .hist = p[55] }, .eact = p[56], .ey = p[57], .sgu = p[58], .sact = p[59], .sxs = p[60], .sy = p[61], .mm_part = p[42] }, .hidden = p[62], .fnormed = p[63], @@ -213,6 +218,8 @@ pub const Buffers = struct { .mx = p[69], .part = p[70], }; + out.moe.x3 = .{ .picks = p[71], .xg = p[72], .xu = p[73], .z = p[74] }; + return out; } pub fn deinit(b: *Buffers) void { diff --git a/zig/src/families/glm/cuda_weights.zig b/zig/src/families/glm/cuda_weights.zig index 80bc44a8e0..ed9b5b7e32 100644 --- a/zig/src/families/glm/cuda_weights.zig +++ b/zig/src/families/glm/cuda_weights.zig @@ -12,6 +12,7 @@ const mlp = @import("cuda_mlp.zig"); const moe = @import("cuda_moe.zig"); const quant = @import("cuda_quant.zig"); const q4 = @import("cuda_q4.zig"); +const exl3 = @import("cuda_exl3.zig"); const Tensor = core.checkpoint.Tensor; const prefix = "model.language_model."; @@ -42,6 +43,7 @@ pub const Weights = struct { norm: u64 = 0, head: u64 = 0, // bf16 [share.vocab][dims] mtp: ?Mtp = null, + exl3_shx: bool = true, // EXL3: every layer's gate and up share one input rotation (a prompt rotates each row once) buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, named: std.ArrayList(Named) = .empty, @@ -79,13 +81,13 @@ const Loader = struct { host: std.ArrayList(u8) = .empty, name_buf: [192]u8 = undefined, - fn get(L: *Loader, comptime fmt: []const u8, args: anytype) !Tensor { + pub fn get(L: *Loader, comptime fmt: []const u8, args: anytype) !Tensor { const name = try std.fmt.bufPrint(&L.name_buf, prefix ++ fmt, args); return L.ck.get(name); } /// The host bytes gathered so far uploaded as one tensor named `fmt`. - fn put(L: *Loader, comptime fmt: []const u8, args: anytype) !u64 { + pub fn put(L: *Loader, comptime fmt: []const u8, args: anytype) !u64 { const b = try cuda.DeviceBuffer.fromHost(L.d, L.host.items); try L.w.buffers.append(L.gpa, b); try L.w.named.append(L.gpa, .{ .name = try std.fmt.allocPrint(L.gpa, fmt, args), .ptr = b.ptr, .len = b.len }); @@ -258,10 +260,18 @@ const Loader = struct { const c = &L.w.config; const s = L.w.share.dense; const d = c.hidden; - return .{ + const w: mlp.Weights = if (c.exl3) bf16: { // an EXL3 checkpoint keeps its dense layers bf16 + inline for (.{ "gate_proj", "up_proj" }) |x| try L.rows(try L.get("layers.{d}.mlp." ++ x ++ ".weight", .{i}), s.lo, s.hi, c.dense_inter, .bf16); + const gu = try L.put("layers.{d}.mlp.gu", .{i}); + try L.cols(try L.get("layers.{d}.mlp.down_proj.weight", .{i}), s.lo, s.hi, .bf16); + break :bf16 .{ .bf16 = .{ .gu = gu, .down = try L.put("layers.{d}.mlp.down", .{i}) } }; + } else .{ .nvfp4 = .{ .gate = try L.fp4("layers.{d}.mlp.gate_proj", i, s.lo, s.hi, 0, d), .up = try L.fp4("layers.{d}.mlp.up_proj", i, s.lo, s.hi, 0, d), .down = try L.fp4("layers.{d}.mlp.down_proj", i, 0, d, s.lo, s.hi), + } }; + return .{ + .w = w, .width = s.len(), .dims = d, .limit = c.swiglu_limit, @@ -344,6 +354,39 @@ const Loader = struct { try L.whole(try L.get("layers.{d}.mlp.gate.e_score_correction_bias", .{i}), .f32); const bias = try L.put("layers.{d}.moe.bias", .{i}); const count: usize = c.experts; + var routed: cuda.experts.Layer = .{ .format = .nvfp4, .up = 0, .down = 0, .width = ni, .dims = d, .experts = count, .limit = c.swiglu_limit }; + var x3: ?exl3.Experts = null; + if (c.exl3) { // the MTP layer's experts too + x3 = try exl3.load(L, i, count, d, c.moe_inter, cut); + if (!x3.?.shx) L.w.exl3_shx = false; + } else routed = try L.nvfp4Experts(i); + const sh = L.w.share.shared; + inline for (.{ "gate_proj", "up_proj" }) |x| try L.rows(try L.get("layers.{d}.mlp.shared_experts." ++ x ++ ".weight", .{i}), sh.lo, sh.hi, c.shared_inter, .bf16); + const gu = try L.put("layers.{d}.moe.shared.gu", .{i}); + try L.cols(try L.get("layers.{d}.mlp.shared_experts.down_proj.weight", .{i}), sh.lo, sh.hi, .bf16); + const sd = try L.put("layers.{d}.moe.shared.down", .{i}); + return .{ + .router = router, + .bias = bias, + .routed = routed, + .exl3 = x3, + .count = count, + .shared_gu = gu, + .shared_down = sd, + .shared = sh.len(), + .top_k = c.topk, + .scale = c.routed_scale, + .norm = c.norm_topk, + }; + } + + /// Layer `i`'s routed NVFP4 experts packed as Experts4 (the MTP layer's quantized from its BF16 ones). + fn nvfp4Experts(L: *Loader, i: u32) !cuda.experts.Layer { + const c = &L.w.config; + const d: usize = c.hidden; + const cut = L.w.share.moe; + const ni = cut.len(); + const count: usize = c.experts; const up_len = X.bytes(1, ni, d, 2) / 4; const down_len = X.bytes(1, d, ni, 1) / 4; const up = try L.gpa.alloc(u32, count * up_len); @@ -383,23 +426,7 @@ const Loader = struct { try L.host.appendSlice(L.gpa, bytes); p.* = try L.put("layers.{d}.moe.experts.{s}", .{ i, what }); } - const sh = L.w.share.shared; - inline for (.{ "gate_proj", "up_proj" }) |x| try L.rows(try L.get("layers.{d}.mlp.shared_experts." ++ x ++ ".weight", .{i}), sh.lo, sh.hi, c.shared_inter, .bf16); - const gu = try L.put("layers.{d}.moe.shared.gu", .{i}); - try L.cols(try L.get("layers.{d}.mlp.shared_experts.down_proj.weight", .{i}), sh.lo, sh.hi, .bf16); - const sd = try L.put("layers.{d}.moe.shared.down", .{i}); - return .{ - .router = router, - .bias = bias, - .routed = .{ .format = .nvfp4, .up = ptrs[0], .down = ptrs[1], .up_scale = ptrs[2], .down_scale = ptrs[3], .width = ni, .dims = d, .experts = count, .limit = c.swiglu_limit }, - .count = count, - .shared_gu = gu, - .shared_down = sd, - .shared = sh.len(), - .top_k = c.topk, - .scale = c.routed_scale, - .norm = c.norm_topk, - }; + return .{ .format = .nvfp4, .up = ptrs[0], .down = ptrs[1], .up_scale = ptrs[2], .down_scale = ptrs[3], .width = ni, .dims = d, .experts = count, .limit = c.swiglu_limit }; } /// weights.load's MTPW: layer `layers` as plain DSA + MoE, norms, eh_proj, draft head from the rank's head rows. diff --git a/zig/tests/cuda/glm_exl3_tests.zig b/zig/tests/cuda/glm_exl3_tests.zig new file mode 100644 index 0000000000..59b1fc37bd --- /dev/null +++ b/zig/tests/cuda/glm_exl3_tests.zig @@ -0,0 +1,238 @@ +//! GLM's EXL3 routed experts (cuda_exl3.run, decode and prompt) against float64 matrices from core's EXL3 decoder. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const glm = @import("glm_cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const Gpu = check.Gpu; + +const fmt = core.exl3_format; +const h = qr.helpers; +const Rig = qr.RigT; + +const E = 12; // routed experts (the shared one is id E, skipped) +const D = 512; // dims +const N = 256; // the rank's width of each expert +const slots = glm.exl3.top_k + 1; +const limit: f32 = 7.0; + +/// One weight set: every expert's trellis words and scales as the loader lays them out, and their float64 matrices. +const Set = struct { + gt: u64, + ut: u64, + dt: u64, + suh_g: u64, + suh_u: u64, + svh_g: u64, + svh_u: u64, + suh_d: u64, + svh_d: u64, + shx: bool, + wg: []f64, // [E][D][N] + wu: []f64, + wd: []f64, // [E][N][D] + + fn experts(s: Set) glm.exl3.Experts { + return .{ .gt = s.gt, .ut = s.ut, .dt = s.dt, .suh_g = s.suh_g, .suh_u = s.suh_u, .svh_g = s.svh_g, .svh_u = s.svh_u, .suh_d = s.suh_d, .svh_d = s.svh_d, .shx = s.shx }; + } +}; + +fn signs(r: *Rig, n: usize, scale: f16) ![]f16 { + const out = try r.a.alloc(f16, n); + for (out) |*v| v.* = if (r.rng.boolean()) scale else -scale; + return out; +} + +/// `E` random K x N matrices: trellis words, suh (shared by every expert when `one`), svh; device copies and float64. +fn matrices(r: *Rig, k: usize, n: usize, shared_suh: ?[]const f16, out_suh: *u64, out_svh: *u64, w: []f64) !u64 { + const words = k / 16 * (n / 16) * 64; + const tr = try r.a.alloc(i16, E * words); + for (tr) |*v| v.* = r.rng.int(i16); + const suh = try r.a.alloc(f16, E * k); + const svh = try r.a.alloc(f16, E * n); + const scratch = try r.a.alloc(f64, 2 * k * n); + for (0..E) |e| { + if (shared_suh) |s| @memcpy(suh[e * k ..][0..k], s) else @memcpy(suh[e * k ..][0..k], try signs(r, k, 1.0)); + @memcpy(svh[e * n ..][0..n], try signs(r, n, 0.02)); + fmt.dequantize(.{ .halves = 8 }, .mcg, tr[e * words ..][0..words], k / 16, n / 16, suh[e * k ..][0..k], svh[e * n ..][0..n], w[e * k * n ..][0 .. k * n], scratch); + } + out_suh.* = try r.dev(std.mem.sliceAsBytes(suh)); + out_svh.* = try r.dev(std.mem.sliceAsBytes(svh)); + return r.dev(std.mem.sliceAsBytes(tr)); +} + +fn weights(r: *Rig, one: bool) !Set { + var s: Set = undefined; + s.shx = one; + s.wg = try r.a.alloc(f64, E * D * N); + s.wu = try r.a.alloc(f64, E * D * N); + s.wd = try r.a.alloc(f64, E * N * D); + const rot = if (one) try signs(r, D, 1.0) else null; // one input rotation for every gate and up + s.gt = try matrices(r, D, N, rot, &s.suh_g, &s.svh_g, s.wg); + s.ut = try matrices(r, D, N, rot, &s.suh_u, &s.svh_u, s.wu); + s.dt = try matrices(r, N, D, null, &s.suh_d, &s.svh_d, s.wd); + return s; +} + +/// Each row's 8 distinct routed picks, then the shared expert's id. +fn picks(r: *Rig, rows: usize) ![]i32 { + const out = try r.a.alloc(i32, rows * slots); + for (0..rows) |row| { + var seen: [E]bool = @splat(false); + for (0..slots - 1) |j| { + var e = r.rng.uintLessThan(usize, E); + while (seen[e]) e = r.rng.uintLessThan(usize, E); + seen[e] = true; + out[row * slots + j] = @intCast(e); + } + out[row * slots + slots - 1] = E; + } + return out; +} + +const Run = struct { ey: u64, rows: usize }; + +/// The family's dispatch on `x` (bf16 [rows][D]) and `pk`: ey [rows][slots][D] fp32, the shared slot left as filled. +fn run(r: *Rig, k: *const glm.exl3.Kernels, ek: *const glm.experts.Kernels, s: Set, x: []const u16, pk: []const i32, rows: usize, prompt: bool) !Run { + const P = rows * glm.exl3.top_k; + const tile: usize = if (prompt) glm.experts.staged_tile else glm.experts.tile; + const wide = P > cuda.grouped.small; + const plan: cuda.grouped.Plan = .{ + .members = try r.zeros(P * 4), + .items = try r.zeros(cuda.grouped.maxItems(P, E, tile) * 12), + .counts = try r.zeros(8), + .rank = try r.zeros((if (wide) P else 1) * 4), + .hist = try r.zeros((if (wide) (P + 1023) / 1024 * E else 1) * 4), + }; + const sz = glm.exl3.scratchBytes(rows, D, N, prompt, s.shx); + const sc: glm.exl3.Scratch = .{ .picks = try r.zeros(sz[0]), .xg = try r.zeros(sz[1]), .xu = try r.zeros(sz[2]), .z = try r.zeros(sz[3]) }; + const ey = try r.zeros(rows * slots * D * 4); + const xd = try r.zeros(P * N * 2); + try glm.exl3.run(k, ek, r.s, s.experts(), sc, plan, try r.dev(std.mem.sliceAsBytes(x)), D, try r.dev(std.mem.sliceAsBytes(pk)), slots, rows, E, D, N, limit, prompt, ey, xd); + try r.s.synchronize(); + return .{ .ey = ey, .rows = rows }; +} + +/// The kernels' SwiGLU: gate and up rounded to bf16 and limited, silu to bf16, the product to bf16. +fn act(g: f64, u: f64) f64 { + const gg = @min(h.bfOf(h.bf16Of(@floatCast(g))), limit); + const uu = @max(@min(h.bfOf(h.bf16Of(@floatCast(u))), limit), -limit); + return h.bfOf(h.bf16Of(@floatCast(h.bfOf(h.bf16Of(@floatCast(gg / (1 + @exp(-gg))))) * uu))); +} + +/// Float64 slot rows for routed slots (the shared slot's want is NaN: checked apart) and their tolerances. +fn reference(r: *Rig, s: Set, x: []const u16, pk: []const i32, rows: usize, want: []f64, tol: []f64) !void { + const a = try r.a.alloc(f64, N); + for (0..rows) |row| for (0..slots - 1) |j| { + const e: usize = @intCast(pk[row * slots + j]); + for (0..N) |n| { + var g: f64 = 0; + var u: f64 = 0; + for (0..D) |kk| { + const xv = h.bfOf(x[row * D + kk]); + g += xv * s.wg[(e * D + kk) * N + n]; + u += xv * s.wu[(e * D + kk) * N + n]; + } + a[n] = act(g, u); + } + for (0..D) |n| { + var y: f64 = 0; + var mag: f64 = 0; + for (0..N) |kk| { + y += a[kk] * s.wd[(e * N + kk) * D + n]; + mag += @abs(a[kk] * s.wd[(e * N + kk) * D + n]); + } + const at = (row * slots + j) * D + n; + want[at] = y; + tol[at] = mag * (1.0 / 64.0) + 1e-6; // fp16 rotated inputs, bf16 activations, fp16 down inputs + } + }; +} + +/// Routed slots within tolerance; the shared slot untouched (the fill's bytes). +fn expect(r: *Rig, name: []const u8, out: Run, s: Set, x: []const u16, pk: []const i32) !void { + const n = out.rows * slots * D; + const want = try r.a.alloc(f64, n); + const tol = try r.a.alloc(f64, n); + try reference(r, s, x, pk, out.rows, want, tol); + const raw = try r.back(out.ey, n * 4); + var most: f64 = 0; + var at: usize = 0; + for (0..out.rows) |row| for (0..slots) |j| for (0..D) |c| { + const i = (row * slots + j) * D + c; + const bits = std.mem.readInt(u32, raw[i * 4 ..][0..4], .little); + if (j == slots - 1) { // the shared slot: the fill's bytes + try check.expect(bits == 0xffffffff, "{s}: the shared slot left alone", .{name}); + continue; + } + const g: f64 = @as(f32, @bitCast(bits)); + const q = @abs(g - want[i]) / tol[i]; + if (!(q <= most)) { + most = q; + at = i; + } + }; + r.worst = @max(r.worst, most); + std.debug.print("RESULT {s} ({d} rows): largest difference {d:.3} of its tolerance (element {d})\n", .{ name, out.rows, most, at }); + try check.expect(most <= 1, "{s} within its float64 tolerance", .{name}); +} + +pub fn run_(gpu: Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var stream = try cuda.Stream.init(gpu.d, false); + defer stream.deinit(); + var prng = std.Random.DefaultPrng.init(7); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = stream, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var xm = try cuda.Module.load(gpu.d, cuda.kernels.glm_exl3); + defer xm.unload(); + var pm = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer pm.unload(); + const k = try glm.exl3.Kernels.load(xm); + var ek = try glm.experts.Kernels.load(gpu.d, pm, @intCast(try gpu.ctx.attribute(.multiprocessor_count))); + defer ek.unload(); + for ([_]bool{ true, false }) |one| { + const s = try weights(&r, one); + for ([_]usize{ 1, 3, 8 }) |rows| { // decode windows + const x = try r.bfs(rows * D, 1.0); + const pk = try picks(&r, rows); + const got = try run(&r, &k, &ek, s, x, pk, rows, false); + try expect(&r, if (one) "glm-exl3 decode, one rotation" else "glm-exl3 decode, a rotation an expert", got, s, x, pk); + if (rows > 1) { // every row's bits as alone + const all = try r.back(got.ey, rows * slots * D * 4); + for (0..rows) |row| { + const alone = try run(&r, &k, &ek, s, x[row * D ..][0..D], pk[row * slots ..][0..slots], 1, false); + const mine = try r.back(alone.ey, slots * D * 4); + try check.expect(std.mem.eql(u8, mine[0 .. (slots - 1) * D * 4], all[row * slots * D * 4 ..][0 .. (slots - 1) * D * 4]), "glm-exl3 decode: row {d} of {d} equals the row alone", .{ row, rows }); + } + } + } + for ([_]usize{ 1, 40, 200 }) |rows| { // prompt chunks + const x = try r.bfs(rows * D, 1.0); + const pk = try picks(&r, rows); + const got = try run(&r, &k, &ek, s, x, pk, rows, true); + try expect(&r, if (one) "glm-exl3 prompt, one rotation" else "glm-exl3 prompt, a rotation an expert", got, s, x, pk); + } + } + { // a NaN in one row's input: that row's routed slots NaN, every other row finite + const s = try weights(&r, true); + const rows = 4; + const x = try r.bfs(rows * D, 1.0); + x[2 * D + 17] = 0x7fc0; + const pk = try picks(&r, rows); + for ([_]bool{ false, true }) |prompt| { + const got = try run(&r, &k, &ek, s, x, pk, rows, prompt); + const raw = try r.back(got.ey, rows * slots * D * 4); + for (0..rows) |row| for (0..slots - 1) |j| for (0..D) |c| { + const v: f32 = @bitCast(std.mem.readInt(u32, raw[((row * slots + j) * D + c) * 4 ..][0..4], .little)); + try check.expect(if (row == 2) std.math.isNan(v) else std.math.isFinite(v), "glm-exl3 {s}: a NaN stays in its row ({d})", .{ if (prompt) "prompt" else "decode", row }); + }; + } + } + check.pass("glm-exl3: decode and prompt slot rows within {d:.3} of their float64 tolerance, rows as alone, NaN kept to its row", .{r.worst}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index a68a7a880d..8ed3f0fc3d 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -22,6 +22,7 @@ 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"); const glm_ref_tests = @import("glm_ref_tests.zig"); +const glm_exl3_tests = @import("glm_exl3_tests.zig"); const carveout_tests = @import("carveout_tests.zig"); const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); @@ -92,6 +93,7 @@ const usage = \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR + \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) \\ glm-ref GLM projections (matmuls, router, attention, indexer, hc) against float64 (synthetic) \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run @@ -160,6 +162,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-ref")) return glm_ref_tests.run(gpu); + if (std.mem.eql(u8, cmd, "glm-exl3")) return glm_exl3_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); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); From 8278074903d38e6b459a7dc41a74ff19ba81d0bb Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 12:26:25 +0000 Subject: [PATCH 24/31] kernel contract: glm_exl3's eight launches with their sm_80/86/89/120/121 usage (tf-cuda-test kernel-contract-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..8ea6c97123 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -45,6 +45,14 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "glm_exl3 rot_rows", .symbol = "glm_exl3_rot_rows", .image = "glm_exl3", .threads = 32 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_exl3 rot_in", .symbol = "glm_exl3_rot_in", .image = "glm_exl3", .threads = 32 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "glm_exl3 grouped", .symbol = "glm_exl3_grouped", .image = "glm_exl3", .threads = 128 }, .usage = .{ .{ .regs = 120, .static = 32832, .bound = 128 }, .{ .regs = 130, .static = 32832, .bound = 128 }, .{ .regs = 130, .static = 32832, .bound = 128 }, .{ .regs = 109, .static = 32832, .bound = 128 }, .{ .regs = 109, .static = 32832, .bound = 128 } } }, + .{ .launch = .{ .name = "glm_exl3 gateup", .symbol = "glm_exl3_gateup", .image = "glm_exl3", .threads = 32 }, .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 = "glm_exl3 down", .symbol = "glm_exl3_down", .image = "glm_exl3", .threads = 32 }, .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 = "glm_exl3 prompt_gu", .symbol = "glm_exl3_prompt_gu", .image = "glm_exl3", .threads = 256 }, .usage = .{ .{ .regs = 127, .static = 43264, .bound = 256 }, .{ .regs = 128, .static = 43264, .bound = 256 }, .{ .regs = 128, .static = 43264, .bound = 256 }, .{ .regs = 126, .static = 43264, .bound = 256 }, .{ .regs = 126, .static = 43264, .bound = 256 } } }, + .{ .launch = .{ .name = "glm_exl3 prompt_dn", .symbol = "glm_exl3_prompt_dn", .image = "glm_exl3", .threads = 256 }, .usage = .{ .{ .regs = 126, .static = 43264, .bound = 256 }, .{ .regs = 126, .static = 43264, .bound = 256 }, .{ .regs = 126, .static = 43264, .bound = 256 }, .{ .regs = 127, .static = 43264, .bound = 256 }, .{ .regs = 127, .static = 43264, .bound = 256 } } }, + .{ .launch = .{ .name = "glm_exl3 picks", .symbol = "glm_exl3_picks", .image = "glm_exl3", .threads = 256 }, .usage = .{ .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 13, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, }; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. From 084c253740bc10ffe3116ca0c951c4510afb175c Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 13:41:36 +0000 Subject: [PATCH 25/31] glm cuda: the ranks watch each other (tp_watch): a rank that fails or dies stops the others instead of leaving them waiting, a failed follower leaves with _exit; every peer's send/recv and NCCL's rings connect before the weights load, so a full GB10 never registers them (ibv_reg_mr refused the first prompt pass's split exchange) --- zig/src/families/glm/cuda_native.zig | 33 +++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 2f09b68f41..5abb1905cb 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -91,6 +91,28 @@ const Owned = struct { lanes: glm_lanes.Lanes, }; +extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int; + +/// Every peer's send/recv connected now, before the weights fill memory (connected on first use, ibv_reg_mr can fail). +fn warmPeers(d: *const cuda.Driver, g: *const cuda.tp_group.Group) !void { + const bytes = 4 << 20; // a prompt piece's exchange: every channel a large send/recv takes + var buf = try cuda.DeviceBuffer.alloc(d, 2 * bytes * g.world); + defer buf.free(); + var s = try cuda.Stream.init(d, false); + defer s.deinit(); + try g.lib.check(g.lib.api.ncclGroupStart(), "ncclGroupStart"); + const posted: anyerror!void = post: { + for (0..g.world) |peer| if (peer != g.rank) { + g.send(s, buf.ptr + peer * bytes, bytes, peer) catch |err| break :post err; + g.recv(s, buf.ptr + (g.world + peer) * bytes, bytes, peer) catch |err| break :post err; + }; + break :post {}; + }; + try g.lib.check(g.lib.api.ncclGroupEnd(), "ncclGroupEnd"); // the group closes on every path + try posted; + try s.synchronize(); +} + /// The largest shard a rank sends by RoCE writes: a decode window's partials (8 rows of fp32) fit. const roce_cap = 512 << 10; @@ -103,9 +125,12 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ std.log.err("{s}", .{problem}); return err; }; + if (p.world > 1) try cuda.tp_watch.start(io, p.rank, p.world); // rank 0's line before NCCL: a failure stops every rank const own = try gpa.create(Owned); errdefer gpa.destroy(own); own.gpa = gpa; + // every connection registered at init, while memory is free (lazily, a full GB10's ibv_reg_mr can refuse it) + if (std.c.getenv("NCCL_RUNTIME_CONNECT") == null) _ = setenv("NCCL_RUNTIME_CONNECT", "0", 0); var lib = try cuda.nccl.Library.open(); errdefer lib.close(); var id: cuda.nccl.UniqueId = undefined; @@ -114,6 +139,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(p.world), id, @intCast(p.rank)), "ncclCommInitRank"); own.group = .{ .lib = lib, .comm = comm, .rank = p.rank, .world = p.world }; errdefer _ = own.group.lib.api.ncclCommDestroy(own.group.comm); + if (p.world > 1) try warmPeers(ctx.d, &own.group); const cap = capacity(o.context); own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, &own.group.lib, own.group.comm, p.rank, p.world, .{ .capacity = cap, .mtp = o.drafts, .settings = set }); errdefer own.e.deinit(); @@ -131,11 +157,15 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ for (rooms[1..], 1..) |r, q| if (peer == null or r < peer.?.bytes) { peer = .{ .bytes = r, .rank = q }; }; + // every rank loaded and agreed (cuda_ranks) before the first step; the digest names what they agreed on + if (p.world > 1) try cuda.tp_watch.ready(io, ranks.hashBytes(std.fs.path.basename(std.mem.trimEnd(u8, dir, "/")))); if (p.rank != 0) { // a follower never serves: it runs rank 0's steps, then leaves the process follow.follow(own.e, &own.link) catch |err| { std.log.err("rank {d}: following rank 0 stopped ({s})", .{ p.rank, @errorName(err) }); - std.process.exit(1); + cuda.tp_watch.fail(@errorName(err)); // rank 0 stops with it instead of waiting on this rank + std.c._exit(1); // the stream may hold kernels waiting on a gone rank: a normal exit would join them }; + cuda.tp_watch.ended(); std.log.info("rank {d}: rank 0 ended; leaving", .{p.rank}); std.process.exit(0); } @@ -196,6 +226,7 @@ pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { fn close(ptr: *anyopaque) void { const own: *Owned = @ptrCast(@alignCast(ptr)); + cuda.tp_watch.ended(); // the other ranks' lines close on purpose from here own.lanes.x.end() catch {}; own.lanes.deinit(); own.link.deinit(); From b090a7bcbf1c245477b4c818b4852be5aacf6859 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 13:41:36 +0000 Subject: [PATCH 26/31] glm cuda: verify windows to the decode buffers' 8 rows, so the lane core's copy proposals reach 7 tokens; head chains stay 3 deep and a first copy as wide as a head round (EXL3 on two GB10s: a rewrite-this-code reply 34.2 to 41.1 tok/s, others unchanged) --- zig/src/families/glm/cuda_lanes.zig | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig index 0106c383b0..38b6376a68 100644 --- a/zig/src/families/glm/cuda_lanes.zig +++ b/zig/src/families/glm/cuda_lanes.zig @@ -11,8 +11,9 @@ const be = lanes.backend; /// First tokens a handle names. const ring = 1024; -/// The widest verify window: the pending token and a full draft chain (glm-decode-check's, the AOT row variants). -pub const max_rows = 1 + engine.max_drafts; +/// The widest verify window: the pending token and a copy (the lane core proposes them), as the decode buffers hold. +pub const max_rows = fw.window_rows; +const mtp_rows = 1 + engine.max_drafts; /// A stream's place in the backend: its sequence's slot, the rows its last verify left uncommitted, its held drafts. const Lane = struct { @@ -56,12 +57,13 @@ pub const Lanes = struct { break :vt v; }; - /// Round loop facts: windows up to max_rows, drafts with the head loaded, shared windows unless TF_GLM_SHARED=0. + /// Round loop facts: copies up to max_rows, head chains to mtp_rows, shared windows unless TF_GLM_SHARED=0. pub fn facts(self: *const Lanes) lanes.Model { const drafting = self.x.e.mtp; const sharing = drafting and !self.alone; return .{ .exact_width = if (drafting) max_rows else 1, + .first_copy_rows = mtp_rows, // a first copy as wide as a head round; it grows as its tokens land .gpu_tokens = false, .mtp = drafting, .speculate = drafting, From 93f5261a8db9ad270717f353a3382f15f58d3ab8 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 13:57:40 +0000 Subject: [PATCH 27/31] glm cuda: the image tower (model.visual.*, rank 0's) on cuBLASLt and glm_vision's steps, the processor's patches through image_core over core's decoder (vision_config and image_core split from the Metal files, which re-export them); tf-cuda-test glm-vision: patches as vision_ref.py's within 4e-7, features 1.003x the torch tower's own bf16 distance from float32 --- zig/build/cuda.zig | 1 + zig/kernels/cuda/glm_vision.cu | 181 +++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/glm/cuda.zig | 7 + zig/src/families/glm/cuda_vision.zig | 403 +++++++++++++++++++++++++ zig/src/families/glm/image.zig | 122 +------- zig/src/families/glm/image_core.zig | 117 +++++++ zig/src/families/glm/vision.zig | 95 +----- zig/src/families/glm/vision_config.zig | 83 +++++ zig/tests/cuda/glm_vision_tests.zig | 104 +++++++ zig/tests/cuda/main.zig | 3 + 11 files changed, 927 insertions(+), 190 deletions(-) create mode 100644 zig/kernels/cuda/glm_vision.cu create mode 100644 zig/src/families/glm/cuda_vision.zig create mode 100644 zig/src/families/glm/image_core.zig create mode 100644 zig/src/families/glm/vision_config.zig create mode 100644 zig/tests/cuda/glm_vision_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 6d7ff7ee81..f0e191a744 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -53,6 +53,7 @@ const kernels = [_]Kernel{ .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family .{ .name = "glm_exl3", .flags = &.{"-O3"} }, // ours: GLM's EXL3 routed experts (Mia-AiLab's checkpoint), decode and prompt + .{ .name = "glm_vision", .flags = &.{"-O3"} }, // ours: GLM's image tower steps around cuBLASLt (norms, rotary, softmax, merger) .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, diff --git a/zig/kernels/cuda/glm_vision.cu b/zig/kernels/cuda/glm_vision.cu new file mode 100644 index 0000000000..d144113051 --- /dev/null +++ b/zig/kernels/cuda/glm_vision.cu @@ -0,0 +1,181 @@ +// GLM vision tower's steps around cuBLASLt, as mlx-vlm glm5_next/vision.py: fp32 residual, bf16 outputs rounded once. + +#include +#include + +namespace { +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ __nv_bfloat16 tobf(float x) { return __float2bfloat16_rn(x); } + +// a block's sum (blockDim.x a multiple of 32, at most 1024); every thread gets it +__device__ float block_sum(float v) { + __shared__ float part[32]; + __shared__ float total; + for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + __syncthreads(); + if (lane == 0) part[warp] = v; + __syncthreads(); + if (warp == 0) { + float s = lane < (int)(blockDim.x >> 5) ? part[lane] : 0.f; + for (int o = 16; o > 0; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o); + if (lane == 0) total = s; + } + __syncthreads(); + return total; +} + +__device__ __forceinline__ float limited_swiglu(float g, float u, float limit) { + g = g > limit ? limit : g; // as torch.clamp: a NaN stays NaN (fminf and fmaxf would drop it for the bound) + u = u < -limit ? -limit : (u > limit ? limit : u); + return g / (1.f + expf(-g)) * u; +} +} // namespace + +// dst[i] = bf16(src[i]) (the processor's fp32 patches) +extern "C" __global__ void glm_vis_to_bf16(const float* __restrict__ src, __nv_bfloat16* __restrict__ dst, long long n) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + dst[i] = tobf(src[i]); +} + +// out[r, c] = bf16(acc[r, c] + bias[c]) (bias may be null); with x32: x32[r, c] += that bf16 value (a residual add) +extern "C" __global__ void glm_vis_bias(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, float* __restrict__ x32, long long n, int cols) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) { + const float v = acc[i] + (bias ? bf(bias[i % cols]) : 0.f); + const __nv_bfloat16 h = tobf(v); + if (out) out[i] = h; + if (x32) x32[i] += bf(h); + } +} + +// x32[i] = acc[i] + bias[i % cols] (the patch embedding: the residual stream's first value, in fp32) +extern "C" __global__ void glm_vis_embed32(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ bias, + float* __restrict__ x32, long long n, int cols) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + x32[i] = bf(tobf(acc[i] + bf(bias[i % cols]))); +} + +// RMSNorm of fp32 rows: y = bf16(x rsqrt(mean(x^2) + eps) w), a block a row +extern "C" __global__ void glm_vis_rmsnorm(const float* __restrict__ x, const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ y, int cols, float eps) { + const float* row = x + (long long)blockIdx.x * cols; + float s = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) s += row[c] * row[c]; + const float inv = rsqrtf(block_sum(s) / cols + eps); + for (int c = threadIdx.x; c < cols; c += blockDim.x) y[(long long)blockIdx.x * cols + c] = tobf(row[c] * inv * bf(w[c])); +} + +// attention inputs: q, k head-RMSNormed and rotated (fp32, 2D angles) into qh, kh; v into vt; a warp a (row, head). +extern "C" __global__ void glm_vis_qkv(const __nv_bfloat16* __restrict__ qkv, const __nv_bfloat16* __restrict__ qn, + const __nv_bfloat16* __restrict__ kn, const int* __restrict__ pos, + const float* __restrict__ inv_freq, __nv_bfloat16* __restrict__ qh, + __nv_bfloat16* __restrict__ kh, __nv_bfloat16* __restrict__ vt, int N, int seg, + int H, float eps) { + const long long item = blockIdx.x * (long long)(blockDim.x >> 5) + (threadIdx.x >> 5); + if (item >= (long long)N * H) return; + const int lane = threadIdx.x & 31; + const int n = (int)(item / H), h = (int)(item % H); + const int g = n / seg, j = n % seg; + const long long blk = (long long)g * H + h; + const __nv_bfloat16* base = qkv + (long long)n * 3 * H * 64 + h * 64; + const int p = lane < 16 ? pos[2 * n] : pos[2 * n + 1]; + const float ang = (float)p * inv_freq[lane & 15]; + const float c = cosf(ang), s = sinf(ang); + for (int which = 0; which < 2; ++which) { + const __nv_bfloat16* src = base + which * H * 64; + const __nv_bfloat16* wn = which ? kn : qn; + float a = bf(src[lane]), b = bf(src[lane + 32]); + float ss = a * a + b * b; + for (int o = 16; o > 0; o >>= 1) ss += __shfl_xor_sync(0xffffffffu, ss, o); + const float inv = rsqrtf(ss / 64.f + eps); + a = bf(tobf(a * inv * bf(wn[lane]))); + b = bf(tobf(b * inv * bf(wn[lane + 32]))); + // rotate_half: dims l < 32 take -x[l + 32], dims l + 32 take x[l] + const float ra = a * c - b * s, rb = b * c + a * s; + __nv_bfloat16* dst = (which ? kh : qh) + (blk * seg + j) * 64; + dst[lane] = tobf(ra); + dst[lane + 32] = tobf(rb); + } + const __nv_bfloat16* v = base + 2 * H * 64; + vt[(blk * 64 + lane) * seg + j] = v[lane]; + vt[(blk * 64 + lane + 32) * seg + j] = v[lane + 32]; +} + +// A scores row's softmax: pr = bf16(exp(scale (s - max))) and its sum (of the bf16 values), a block a row +extern "C" __global__ void glm_vis_softmax(const float* __restrict__ sc, __nv_bfloat16* __restrict__ pr, + float* __restrict__ sums, int n, float scale) { + const float* row = sc + (long long)blockIdx.x * n; + __shared__ float mx_s[32]; + __shared__ float mx; + float m = -INFINITY; + for (int c = threadIdx.x; c < n; c += blockDim.x) m = fmaxf(m, row[c]); + for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); + if ((threadIdx.x & 31) == 0) mx_s[threadIdx.x >> 5] = m; + __syncthreads(); + if (threadIdx.x < 32) { + float v = threadIdx.x < (blockDim.x >> 5) ? mx_s[threadIdx.x] : -INFINITY; + for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if (threadIdx.x == 0) mx = v; + } + __syncthreads(); + float s = 0.f; + for (int c = threadIdx.x; c < n; c += blockDim.x) { + const __nv_bfloat16 e = tobf(expf((row[c] - mx) * scale)); + pr[(long long)blockIdx.x * n + c] = e; + s += bf(e); + } + s = block_sum(s); + if (threadIdx.x == 0) sums[blockIdx.x] = s; +} + +// att[r, h * 64 + d] = bf16(oh[r, d] / sums[r]) (a head's values product, normalized, into its columns) +extern "C" __global__ void glm_vis_head_out(const float* __restrict__ oh, const float* __restrict__ sums, + __nv_bfloat16* __restrict__ att, int rows, int H, int h) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= rows * 64) return; + const int r = i / 64, d = i % 64; + att[(long long)r * H * 64 + h * 64 + d] = tobf(oh[i] / sums[r]); +} + +// out = bf16(limited_swiglu(bf16(g + gb), bf16(u + ub))) from the two fp32 products (biases may be null) +extern "C" __global__ void glm_vis_swiglu(const float* __restrict__ g, const float* __restrict__ u, + const __nv_bfloat16* __restrict__ gb, const __nv_bfloat16* __restrict__ ub, + __nv_bfloat16* __restrict__ out, long long n, int cols, float limit) { + for (long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) { + const int c = (int)(i % cols); + const float gv = bf(tobf(g[i] + (gb ? bf(gb[c]) : 0.f))); + const float uv = bf(tobf(u[i] + (ub ? bf(ub[c]) : 0.f))); + out[i] = tobf(limited_swiglu(gv, uv, limit)); + } +} + +// The merger's norm: y = bf16(gelu_erf(LayerNorm(bf16(acc)) w + b)), a block a row of ``cols`` +extern "C" __global__ void glm_vis_ln_gelu(const float* __restrict__ acc, const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ b, __nv_bfloat16* __restrict__ y, int cols, + float eps) { + const float* row = acc + (long long)blockIdx.x * cols; + float s = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) s += bf(tobf(row[c])); + const float mean = block_sum(s) / cols; + float v = 0.f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + const float d = bf(tobf(row[c])) - mean; + v += d * d; + } + const float inv = rsqrtf(block_sum(v) / cols + eps); + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + const float x = bf(tobf((bf(tobf(row[c])) - mean) * inv * bf(w[c]) + bf(b[c]))); + y[(long long)blockIdx.x * cols + c] = tobf(0.5f * x * (1.f + erff(x * 0.70710678118654752f))); + } +} + +// Placeholder rows take the tower's rows in every stream: row r takes feature row map[r] (map[r] < 0: kept). +extern "C" __global__ void glm_vis_inject(const __nv_bfloat16* __restrict__ feats, __nv_bfloat16* __restrict__ out, + const int* __restrict__ map, int rows, int d, int copies) { + const int r = blockIdx.x; + if (r >= rows || map[r] < 0) return; + const __nv_bfloat16* src = feats + (long long)map[r] * d; + for (int c = threadIdx.x; c < d; c += blockDim.x) + for (int k = 0; k < copies; ++k) out[((long long)r * copies + k) * d + c] = src[c]; +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 6368e28757..1f97980c52 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -40,6 +40,7 @@ pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes e pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_exl3: []const u8 = if (available) &Blob("fatbin_glm_exl3").bytes else &.{}; +pub const glm_vision: []const u8 = if (available) &Blob("fatbin_glm_vision").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index aeefc3ffee..0655caefbb 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -2,6 +2,11 @@ pub const experts = @import("cuda_experts.zig"); pub const exl3 = @import("cuda_exl3.zig"); +pub const vision = @import("cuda_vision.zig"); +pub const vision_config = @import("vision_config.zig"); +pub const image_core = @import("image_core.zig"); +/// An image's bytes to the tower's patches, as the Metal family prepares them (core's decoder). +pub const image = image_core.For(@import("core").picture); pub const glue = @import("cuda_glue.zig"); pub const kda = @import("cuda_kda.zig"); pub const b16 = @import("cuda_b16.zig"); @@ -27,6 +32,8 @@ pub const native = @import("cuda_native.zig"); test { _ = experts; _ = exl3; + _ = vision; + _ = image_core; _ = glue; _ = kda; _ = b16; diff --git a/zig/src/families/glm/cuda_vision.zig b/zig/src/families/glm/cuda_vision.zig new file mode 100644 index 0000000000..6cf2825221 --- /dev/null +++ b/zig/src/families/glm/cuda_vision.zig @@ -0,0 +1,403 @@ +//! GLM-5.3-Flash's image tower on CUDA as glm5_next/vision.py: patches to the language model's rows (rank 0 encodes). + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const H: i64 = 1024; // hidden +const HEADS: i64 = 16; +const HD: i64 = 64; // a head's width +const MLP: i64 = 4096; +const PATCH: i64 = 1176; // 3 channels x 2 frames x 14 x 14 +const MERGE: i64 = 2; +const OUT: i64 = 4096; // the merged rows' width (the language model's) +const PROJ: i64 = 10240; // the merger's SwiGLU width +pub const DEPTH = 24; +const CHUNK: i64 = 2048; // query rows a scores product covers +const ROWS: i64 = 4096; // rows a linear's fp32 product covers (the scratch's bound) +const EPS: f32 = 1e-5; +const LIMIT: f32 = 10.0; + +const Block = struct { n1: u64, n2: u64, qkv_w: u64, qkv_b: u64, qn: u64, kn: u64, proj_w: u64, proj_b: u64, g_w: u64, g_b: u64, u_w: u64, u_b: u64, d_w: u64, d_b: u64 }; + +const LinKey = struct { m: u64, n: u64, k: u64, f32_out: bool }; + +/// glm_vision's kernels and the stream they run on. +pub const Run = struct { + stream: cuda.Stream, + f: Funcs, +}; + +pub const Funcs = struct { + to_bf16: cuda.Function, + bias: cuda.Function, + embed32: cuda.Function, + rmsnorm: cuda.Function, + qkv: cuda.Function, + softmax: cuda.Function, + head_out: cuda.Function, + swiglu: cuda.Function, + ln_gelu: cuda.Function, + inject: cuda.Function, + + pub fn load(m: cuda.Module) !Funcs { + var f: Funcs = undefined; + inline for (@typeInfo(Funcs).@"struct".field_names) |name| @field(f, name) = try m.function("glm_vis_" ++ name); + return f; + } +}; + +fn launch(r: *const Run, f: cuda.Function, grid: cuda.Dim3, threads: u32, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = grid, .block = .{ .x = threads } }, r.stream, a); +} + +pub const Tower = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + lt: cuda.cublaslt.Library, + mem: cuda.DeviceBuffer, + blocks: [DEPTH]Block = undefined, + patch_w: u64 = 0, + patch_b: u64 = 0, + post: u64 = 0, + down_w: u64 = 0, // [4096, 2, 2, 1024]: the conv's taps in the merge block's (row, col) order, then the channels + down_b: u64 = 0, + m_proj: u64 = 0, + m_ln_w: u64 = 0, + m_ln_b: u64 = 0, + m_gate: u64 = 0, + m_up: u64 = 0, + m_down: u64 = 0, + inv_freq: cuda.DeviceBuffer, + linears: std.AutoHashMap(LinKey, cuda.cublaslt.Linear), + ws: cuda.DeviceBuffer, + ws_len: usize = 64 << 20, + bytes: u64 = 0, + stages: ?[]f32 = null, // checks: the fp32 stream after the patch embedding and after every block, read back + + /// The tower's tensors (model.visual.*, bf16) from the checkpoint's shards in ``dir``, in one allocation. + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8) !Tower { + var ck = try core.Checkpoint.openModelPrefix(gpa, io, dir, "model.visual."); + defer ck.close(); + var total: u64 = 0; + for (ck.files.items) |f| for (f.names.keys()) |name| if (std.mem.startsWith(u8, name, "model.visual.")) { + total += std.mem.alignForward(u64, (try ck.get(name)).bytes.len, 256); + }; + var t: Tower = .{ .gpa = gpa, .d = d, .lt = try cuda.cublaslt.Library.open(), .mem = try cuda.DeviceBuffer.alloc(d, total), .inv_freq = try cuda.DeviceBuffer.alloc(d, 16 * 4), .linears = .init(gpa), .ws = undefined, .bytes = total }; + errdefer t.deinit(); + t.ws = try cuda.DeviceBuffer.alloc(d, t.ws_len); + var at: u64 = 0; + var names = std.StringHashMap(u64).init(gpa); + defer names.deinit(); + for (ck.files.items) |f| for (f.names.keys()) |name| if (std.mem.startsWith(u8, name, "model.visual.")) { + const ten = try ck.get(name); + if (ten.dtype != .bf16) return error.UnsupportedVisionDType; + var bytes = ten.bytes; + var perm: []u16 = &.{}; + defer gpa.free(perm); + if (std.mem.eql(u8, name, "model.visual.downsample.weight")) { // [o, c, 2, 2] -> [o, 2, 2, c] + if (ten.rank != 4 or ten.dim(0) != OUT or ten.dim(1) != H) return error.BadVisionShape; + const src = std.mem.bytesAsSlice(u16, @as([]align(1) const u8, bytes)); + perm = try gpa.alloc(u16, src.len); + const c_n: usize = @intCast(H); + for (0..@intCast(OUT)) |o| for (0..c_n) |c| for (0..2) |kh| for (0..2) |kw| { + perm[((o * 2 + kh) * 2 + kw) * c_n + c] = src[((o * c_n + c) * 2 + kh) * 2 + kw]; + }; + bytes = std.mem.sliceAsBytes(perm); + } + const ptr = t.mem.ptr + at; + try d.check(d.api.cuMemcpyHtoD_v2(ptr, bytes.ptr, bytes.len), "vision weights"); + try names.put(name, ptr); + at += std.mem.alignForward(u64, bytes.len, 256); + }; + const get = struct { + fn f(m: *const std.StringHashMap(u64), comptime fmt: []const u8, args: anytype) !u64 { + var nb: [128]u8 = undefined; + const name = try std.fmt.bufPrint(&nb, "model.visual." ++ fmt, args); + return m.get(name) orelse { + std.log.err("glm vision: no tensor {s}", .{name}); + return error.MissingVisionTensor; + }; + } + }.f; + t.patch_w = try get(&names, "patch_embed.proj.weight", .{}); + t.patch_b = try get(&names, "patch_embed.proj.bias", .{}); + t.post = try get(&names, "post_layernorm.weight", .{}); + t.down_w = try get(&names, "downsample.weight", .{}); + t.down_b = try get(&names, "downsample.bias", .{}); + t.m_proj = try get(&names, "merger.proj.weight", .{}); + t.m_ln_w = try get(&names, "merger.post_projection_norm.weight", .{}); + t.m_ln_b = try get(&names, "merger.post_projection_norm.bias", .{}); + t.m_gate = try get(&names, "merger.gate_proj.weight", .{}); + t.m_up = try get(&names, "merger.up_proj.weight", .{}); + t.m_down = try get(&names, "merger.down_proj.weight", .{}); + for (&t.blocks, 0..) |*bl, i| { + bl.* = .{ + .n1 = try get(&names, "blocks.{d}.norm1.weight", .{i}), + .n2 = try get(&names, "blocks.{d}.norm2.weight", .{i}), + .qkv_w = try get(&names, "blocks.{d}.attn.qkv.weight", .{i}), + .qkv_b = try get(&names, "blocks.{d}.attn.qkv.bias", .{i}), + .qn = try get(&names, "blocks.{d}.attn.q_norm.weight", .{i}), + .kn = try get(&names, "blocks.{d}.attn.k_norm.weight", .{i}), + .proj_w = try get(&names, "blocks.{d}.attn.proj.weight", .{i}), + .proj_b = try get(&names, "blocks.{d}.attn.proj.bias", .{i}), + .g_w = try get(&names, "blocks.{d}.mlp.gate_proj.weight", .{i}), + .g_b = try get(&names, "blocks.{d}.mlp.gate_proj.bias", .{i}), + .u_w = try get(&names, "blocks.{d}.mlp.up_proj.weight", .{i}), + .u_b = try get(&names, "blocks.{d}.mlp.up_proj.bias", .{i}), + .d_w = try get(&names, "blocks.{d}.mlp.down_proj.weight", .{i}), + .d_b = try get(&names, "blocks.{d}.mlp.down_proj.bias", .{i}), + }; + } + // the rotary's frequencies: 1 / 10000^(i / 16), i = 0, 2 .. 30 over 32 (head width / 2) + var inv: [16]f32 = undefined; + for (&inv, 0..) |*v, i| v.* = 1.0 / std.math.pow(f32, 10000.0, @as(f32, @floatFromInt(2 * i)) / 32.0); + try t.inv_freq.upload(0, std.mem.sliceAsBytes(&inv)); + return t; + } + + pub fn deinit(t: *Tower) void { + var it = t.linears.valueIterator(); + while (it.next()) |l| l.deinit(); + t.linears.deinit(); + t.mem.free(); + t.inv_freq.free(); + t.ws.free(); + t.lt.close(); + } + + /// Stage `i`'s fp32 stream (N rows) into `dst` (checks only: it waits for the GPU). + fn stage(t: *Tower, ru: *const Run, dst: []f32, i: usize, x32: u64, N: i64) !void { + const n: usize = @intCast(N * H); + if ((i + 1) * n > dst.len) return; + try ru.stream.synchronize(); + try t.d.check(t.d.api.cuMemcpyDtoH_v2(@ptrCast(dst[i * n ..].ptr), x32, n * 4), "vision stage"); + } + + /// D[m, n] = X[m, k] . W[n, k]^T (fp32 or bf16 out), a plan per shape kept for the next image. + fn linear(t: *Tower, ru: *const Run, x: u64, w: u64, d: u64, m: i64, n: i64, kk: i64, f32_out: bool) !void { + const key: LinKey = .{ .m = @intCast(m), .n = @intCast(n), .k = @intCast(kk), .f32_out = f32_out }; + const gop = try t.linears.getOrPut(key); + if (!gop.found_existing) { + gop.value_ptr.* = cuda.cublaslt.Linear.init(&t.lt, key.m, key.n, key.k, if (f32_out) .f32 else .bf16, t.ws_len) catch |e| { + _ = t.linears.remove(key); + return e; + }; + } + try gop.value_ptr.run(x, w, d, t.ws.ptr, t.ws_len, ru.stream.handle); + } + + fn grid1(n: i64) cuda.Dim3 { + return .{ .x = @intCast(@min(std.math.divCeil(i64, n, 256) catch unreachable, 65535)) }; + } + + /// out = bf16(x . W^T + bias) (bias 0: none), in row chunks through ``acc``; with ``x32``, x32 += that instead. + fn dense(t: *Tower, ru: *const Run, x: u64, w: u64, bias: u64, out: u64, x32: u64, acc: u64, m: i64, n: i64, kk: i64) !void { + var r: i64 = 0; + while (r < m) : (r += ROWS) { + const rows = @min(ROWS, m - r); + try t.linear(ru, x + @as(u64, @intCast(r * kk * 2)), w, acc, rows, n, kk, true); + var a: cuda.Args = .{}; + a.add(acc); + a.add(bias); + a.add(if (out != 0) out + @as(u64, @intCast(r * n * 2)) else 0); + a.add(if (x32 != 0) x32 + @as(u64, @intCast(r * n * 4)) else 0); + a.add(@as(i64, rows * n)); + a.add(@as(i32, @intCast(n))); + try launch(ru, ru.f.bias, grid1(rows * n), 256, &a); + } + } + + /// The limited SwiGLU of two linears (biases 0: none) over x [m, kk] -> out [m, n] bf16, in row chunks. + fn swiglu(t: *Tower, ru: *const Run, x: u64, gw: u64, gb: u64, uw: u64, ub: u64, out: u64, acc: u64, acc2: u64, m: i64, n: i64, kk: i64) !void { + var r: i64 = 0; + const step = @divFloor(ROWS * MLP, n); // the fp32 products' rows within the scratch + while (r < m) : (r += step) { + const rows = @min(step, m - r); + const xr = x + @as(u64, @intCast(r * kk * 2)); + try t.linear(ru, xr, gw, acc, rows, n, kk, true); + try t.linear(ru, xr, uw, acc2, rows, n, kk, true); + var a: cuda.Args = .{}; + a.add(acc); + a.add(acc2); + a.add(gb); + a.add(ub); + a.add(out + @as(u64, @intCast(r * n * 2))); + a.add(@as(i64, rows * n)); + a.add(@as(i32, @intCast(n))); + a.add(LIMIT); + try launch(ru, ru.f.swiglu, grid1(rows * n), 256, &a); + } + } + + fn rmsnorm(ru: *const Run, x32: u64, w: u64, y: u64, rows: i64, cols: i64) !void { + var a: cuda.Args = .{}; + a.add(x32); + a.add(w); + a.add(y); + a.add(@as(i32, @intCast(cols))); + a.add(EPS); + try launch(ru, ru.f.rmsnorm, .{ .x = @intCast(rows) }, 256, &a); + } + + /// Scratch bytes ``encode`` takes for ``patches`` patches in segments of ``seg``. + pub fn scratchBytes(patches: i64, seg: i64) u64 { + const N = patches; + const q = @min(CHUNK, seg); + return @intCast(N * PATCH * 4 + N * H * 4 + N * H * 2 * 5 + N * MLP * 2 + 2 * ROWS * MLP * 4 + q * seg * 6 + q * HD * 4 + q * 4 + N * 8 + 16 * 1024); + } + + /// Features of an image or video (``grid`` [t, h, w]): fp32 patches [t h w, 1176] -> ``out`` [t h w / 4, 4096]. + pub fn encode(t: *Tower, ru: *const Run, patches: []const f32, grid: [3]i64, out: u64) !void { + const d = t.d; + const gt = grid[0]; + const gh = grid[1]; + const gw = grid[2]; + if (gt < 1 or gh < MERGE or gw < MERGE or @mod(gh, MERGE) != 0 or @mod(gw, MERGE) != 0) return error.BadImageGrid; + const seg = gh * gw; + const N = gt * seg; + if (@as(i64, @intCast(patches.len)) != N * PATCH) return error.BadImagePatches; + // each patch's (row, col) in merge-block order, every frame group alike + const pos = try t.gpa.alloc(i32, @intCast(N * 2)); + defer t.gpa.free(pos); + var i: i64 = 0; + while (i < N) : (i += 1) { + const j = @mod(i, seg); + const in_col = @mod(j, MERGE); + const in_row = @mod(@divTrunc(j, MERGE), MERGE); + const block_col = @mod(@divTrunc(j, MERGE * MERGE), @divTrunc(gw, MERGE)); + const block_row = @divTrunc(j, MERGE * MERGE * @divTrunc(gw, MERGE)); + pos[@intCast(i * 2)] = @intCast(block_row * MERGE + in_row); + pos[@intCast(i * 2 + 1)] = @intCast(block_col * MERGE + in_col); + } + var scratch = try cuda.DeviceBuffer.alloc(d, scratchBytes(N, seg)); + defer scratch.free(); + var at: u64 = scratch.ptr; + const take = struct { + fn f(p: *u64, n: i64) u64 { + const r = p.*; + p.* += std.mem.alignForward(u64, @intCast(n), 256); + return r; + } + }.f; + const q_rows = @min(CHUNK, seg); + const p32 = take(&at, N * PATCH * 4); + const x32 = take(&at, N * H * 4); + const xn = take(&at, N * H * 2); + const qh = take(&at, N * H * 2); + const kh = take(&at, N * H * 2); + const vt = take(&at, N * H * 2); + const att = take(&at, N * H * 2); + const big = take(&at, N * MLP * 2); + const acc = take(&at, ROWS * MLP * 4); + const acc2 = take(&at, ROWS * MLP * 4); + const sc = take(&at, q_rows * seg * 4); + const pr = take(&at, q_rows * seg * 2); + const oh = take(&at, q_rows * HD * 4); + const sums = take(&at, q_rows * 4); + const dpos = take(&at, N * 8); + try d.check(d.api.cuMemcpyHtoDAsync_v2(dpos, pos.ptr, pos.len * 4, ru.stream.handle), "vision positions"); + try d.check(d.api.cuMemcpyHtoDAsync_v2(p32, patches.ptr, patches.len * 4, ru.stream.handle), "vision patches"); + { // the patch embedding: the patches as bf16, then the conv as a linear (with its bias) into the fp32 stream + var a: cuda.Args = .{}; + a.add(p32); + a.add(big); + a.add(@as(i64, N * PATCH)); + try launch(ru, ru.f.to_bf16, grid1(N * PATCH), 256, &a); + var r: i64 = 0; + while (r < N) : (r += ROWS) { + const rows = @min(ROWS, N - r); + try t.linear(ru, big + @as(u64, @intCast(r * PATCH * 2)), t.patch_w, acc, rows, H, PATCH, true); + var b: cuda.Args = .{}; + b.add(acc); + b.add(t.patch_b); + b.add(x32 + @as(u64, @intCast(r * H * 4))); + b.add(@as(i64, rows * H)); + b.add(@as(i32, @intCast(H))); + try launch(ru, ru.f.embed32, grid1(rows * H), 256, &b); + } + } + const scale: f32 = 1.0 / @sqrt(@as(f32, HD)); + if (t.stages) |dst| try t.stage(ru, dst, 0, x32, N); + for (t.blocks, 1..) |bl, bi| { + defer if (t.stages) |dst| t.stage(ru, dst, bi, x32, N) catch {}; + try rmsnorm(ru, x32, bl.n1, xn, N, H); + try t.dense(ru, xn, bl.qkv_w, bl.qkv_b, big, 0, acc, N, 3 * H, H); + { + var a: cuda.Args = .{}; + a.add(big); + a.add(bl.qn); + a.add(bl.kn); + a.add(dpos); + a.add(t.inv_freq.ptr); + a.add(qh); + a.add(kh); + a.add(vt); + a.add(@as(i32, @intCast(N))); + a.add(@as(i32, @intCast(seg))); + a.add(@as(i32, @intCast(HEADS))); + a.add(EPS); + try launch(ru, ru.f.qkv, .{ .x = @intCast(std.math.divCeil(i64, N * HEADS, 8) catch unreachable) }, 256, &a); + } + // each frame group attends within itself, per head, in query chunks (q, k [t, H, seg, 64]; v transposed) + var g: i64 = 0; + while (g < gt) : (g += 1) { + var h: i64 = 0; + while (h < HEADS) : (h += 1) { + const blk: u64 = @intCast((g * HEADS + h) * seg * HD * 2); + var c: i64 = 0; + while (c < seg) : (c += CHUNK) { + const rows = @min(CHUNK, seg - c); + try t.linear(ru, qh + blk + @as(u64, @intCast(c * HD * 2)), kh + blk, sc, rows, seg, HD, true); + var a: cuda.Args = .{}; + a.add(sc); + a.add(pr); + a.add(sums); + a.add(@as(i32, @intCast(seg))); + a.add(scale); + try launch(ru, ru.f.softmax, .{ .x = @intCast(rows) }, 256, &a); + try t.linear(ru, pr, vt + blk, oh, rows, HD, seg, true); + var b: cuda.Args = .{}; + b.add(oh); + b.add(sums); + b.add(att + @as(u64, @intCast((g * seg + c) * H * 2))); + b.add(@as(i32, @intCast(rows))); + b.add(@as(i32, @intCast(HEADS))); + b.add(@as(i32, @intCast(h))); + try launch(ru, ru.f.head_out, grid1(rows * HD), 256, &b); + } + } + } + try t.dense(ru, att, bl.proj_w, bl.proj_b, 0, x32, acc, N, H, H); + try rmsnorm(ru, x32, bl.n2, xn, N, H); + try t.swiglu(ru, xn, bl.g_w, bl.g_b, bl.u_w, bl.u_b, big, acc, acc2, N, MLP, H); + try t.dense(ru, big, bl.d_w, bl.d_b, 0, x32, acc, N, H, MLP); + } + // the post-norm, then each merge block's four patches (adjacent rows) through the 2 x 2 conv as one linear + try rmsnorm(ru, x32, t.post, xn, N, H); + const M = @divExact(N, MERGE * MERGE); + try t.dense(ru, xn, t.down_w, t.down_b, att, 0, acc, M, OUT, 4 * H); + // the merger: projection, LayerNorm + GELU, the limited SwiGLU to 10240 and back + var r: i64 = 0; + while (r < M) : (r += ROWS) { + const rows = @min(ROWS, M - r); + try t.linear(ru, att + @as(u64, @intCast(r * OUT * 2)), t.m_proj, acc, rows, OUT, OUT, true); + var a: cuda.Args = .{}; + a.add(acc); + a.add(t.m_ln_w); + a.add(t.m_ln_b); + a.add(xn + @as(u64, @intCast(r * OUT * 2))); + a.add(@as(i32, @intCast(OUT))); + a.add(EPS); + try launch(ru, ru.f.ln_gelu, .{ .x = @intCast(rows) }, 256, &a); + } + r = 0; + const step = @divFloor(ROWS * MLP, PROJ); + while (r < M) : (r += step) { + const rows = @min(step, M - r); + try t.swiglu(ru, xn + @as(u64, @intCast(r * OUT * 2)), t.m_gate, 0, t.m_up, 0, big, acc, acc2, rows, PROJ, OUT); + try t.linear(ru, big, t.m_down, out + @as(u64, @intCast(r * OUT * 2)), rows, OUT, PROJ, false); + } + try ru.stream.synchronize(); // the scratch is freed on return + } +}; diff --git a/zig/src/families/glm/image.zig b/zig/src/families/glm/image.zig index b22d34de11..4d0ee62758 100644 --- a/zig/src/families/glm/image.zig +++ b/zig/src/families/glm/image.zig @@ -1,116 +1,18 @@ //! An image's bytes to GLM5-Next vision patches as TensorFold 0.6.6's --vision made them (Pillow, then mlx-vlm). const std = @import("std"); const picture = @import("../../core/picture/picture.zig"); -const vision = @import("vision.zig"); - -/// One image ready for the tower: its patches, its patch grid and its tokens (gh * gw / 4). -pub const Prepared = struct { - pixels: []f32, // [gh * gw, 1176] - gh: u32, - gw: u32, - tokens: u32, - hash: u64, // its identity (core/picture/identity.zig): bytes and grid, so another budget is another image -}; - -/// The shared decoder's refusals, and a token budget too small for an image. -pub const Error = picture.Error || error{ImageTooSmall}; - -/// The words for a refusal the server cannot word: the JPEG library this machine lacks, or a CMYK JPEG. -pub const decoderProblem = picture.decoderProblem; - -// ---- the processor's geometry ---- - -pub const Geometry = struct { target_h: u32, target_w: u32, content_h: u32, content_w: u32 }; - -fn alignUp(v: u64, f: u64) u64 { - return (v + f - 1) / f * f; -} - -/// smart_resize and _resize_geometry: the padded canvas under `max_tokens` (at least `min_tokens`), the content size. -pub fn geometry(h: u32, w: u32, min_tokens: u32, max_tokens: u32) Error!Geometry { - const factor: u64 = vision.patch * vision.merge; - const frames: u64 = vision.temporal; - const per_token = frames * factor * factor; - const min_px = @as(u64, min_tokens) * per_token; - const max_px = @as(u64, max_tokens) * per_token; - if (min_tokens == 0 or max_tokens < min_tokens) return error.ImageTooSmall; - const H: u64 = h; - const W: u64 = w; - var ah = alignUp(H, factor); - var aw = alignUp(W, factor); - var pixels = frames * ah * aw; - if (pixels < min_px) { - const scale = @sqrt(@as(f64, @floatFromInt(min_px)) / @as(f64, @floatFromInt(frames * H * W))); - ah = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(H)) * scale)))), factor); - aw = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(W)) * scale)))), factor); - pixels = frames * ah * aw; - } - if (pixels > max_px) { // the tallest content that fits, by bisection on its height - if (max_px < frames * factor * factor) return error.ImageTooSmall; - var best: [2]u64 = .{ factor, factor }; - var low: i64 = 1; - var high: i64 = @intCast(H); - while (low <= high) { - const ch: u64 = @intCast(@divFloor(low + high, 2)); - const cw: u64 = @max(1, W * ch / H); // math.floor(width * content_h / height) - const cand: [2]u64 = .{ alignUp(ch, factor), alignUp(cw, factor) }; - if (frames * cand[0] * cand[1] <= max_px) { - best = cand; - low = @as(i64, @intCast(ch)) + 1; - } else high = @as(i64, @intCast(ch)) - 1; - } - ah = best[0]; - aw = best[1]; - } - var scale = @min(@as(f64, @floatFromInt(ah)) / @as(f64, @floatFromInt(H)), @as(f64, @floatFromInt(aw)) / @as(f64, @floatFromInt(W))); - if (frames * H * W >= min_px) scale = @min(1.0, scale); - const chf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(H)) * scale)); - const cwf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(W)) * scale)); - return .{ .target_h = @intCast(ah), .target_w = @intCast(aw), .content_h = @intCast(@max(1, @min(ah, chf))), .content_w = @intCast(@max(1, @min(aw, cwf))) }; -} - -/// The resized content padded with black, scaled, normalized and cut into 2x2 windows of 14x14 patches, row-major. -pub fn patchify(gpa: std.mem.Allocator, content: picture.Rgb, g: Geometry, mean: [3]f32, sd: [3]f32) Error![]f32 { - const P = vision.patch; - const M = vision.merge; - const gh = g.target_h / P; - const gw = g.target_w / P; - const out = try gpa.alloc(f32, @as(usize, gh) * gw * vision.patch_width); - const rescale: f32 = @floatCast(@as(f64, 1.0) / 255.0); - var row: usize = 0; - for (0..gh / M) |wy| for (0..gw / M) |wx| for (0..M) |my| for (0..M) |mx| { - const py0 = (wy * M + my) * P; - const px0 = (wx * M + mx) * P; - const dst = out[row * vision.patch_width ..][0..vision.patch_width]; - for (0..3) |ch| for (0..vision.temporal) |t| for (0..P) |dy| for (0..P) |dx| { - const y = py0 + dy; - const x = px0 + dx; - const v: u8 = if (y < content.height and x < content.width) content.pixels[(y * content.width + x) * 3 + ch] else 0; - const f = @as(f32, @floatFromInt(v)) * rescale; - dst[((ch * vision.temporal + t) * P + dy) * P + dx] = (f - mean[ch]) / sd[ch]; - }; - row += 1; - }; - return out; -} - -/// Bytes to patches: decoded as 0.6.6 decoded them (core's picture), sized for `max_tokens`, resized, padded and cut. -pub fn prepare(io: std.Io, gpa: std.mem.Allocator, bytes: []const u8, c: vision.Config, max_tokens: u32) Error!Prepared { - var img = try picture.decode(io, gpa, bytes); - defer img.deinit(gpa); - return prepareRgb(gpa, img, picture.identity.ofBytes(0x1a6e, bytes), c, max_tokens); -} - -/// RGB to patches (Pillow's bicubic resize, the padded canvas cut); the hash covers the bytes' hash and the grid. -pub fn prepareRgb(gpa: std.mem.Allocator, img: picture.Rgb, bytes_hash: u64, c: vision.Config, max_tokens: u32) Error!Prepared { - const g = try geometry(@intCast(img.height), @intCast(img.width), @min(c.min_tokens, max_tokens), max_tokens); - const resized = try picture.resample.resize(gpa, .pillow, .of(img.pixels, img.width, img.height), g.content_w, g.content_h); - defer gpa.free(resized); - const content: picture.Rgb = .{ .width = g.content_w, .height = g.content_h, .pixels = resized }; - const pixels = try patchify(gpa, content, g, c.mean, c.std); - const grid = [2]u32{ g.target_h / vision.patch, g.target_w / vision.patch }; - return .{ .pixels = pixels, .gh = grid[0], .gw = grid[1], .tokens = grid[0] * grid[1] / vision.patches_per_token, .hash = picture.identity.withGrid(bytes_hash, grid) }; -} +const vision = @import("vision_config.zig"); +const core = @import("image_core.zig"); +const steps = core.For(picture); + +pub const Prepared = core.Prepared; +pub const Geometry = core.Geometry; +pub const geometry = core.geometry; +pub const Error = steps.Error; +pub const decoderProblem = steps.decoderProblem; +pub const patchify = steps.patchify; +pub const prepare = steps.prepare; +pub const prepareRgb = steps.prepareRgb; test "geometry follows smart_resize" { // the reference dumps' grids: 517x389 pads to 532x392 unscaled; 96x96 grows to 112; 20x15 to the 16-token floor diff --git a/zig/src/families/glm/image_core.zig b/zig/src/families/glm/image_core.zig new file mode 100644 index 0000000000..3f0a087177 --- /dev/null +++ b/zig/src/families/glm/image_core.zig @@ -0,0 +1,117 @@ +//! An image's bytes to GLM5-Next vision patches as TensorFold 0.6.6's --vision made them, over the caller's decoder. +const std = @import("std"); +const vision = @import("vision_config.zig"); + +/// One image ready for the tower: its patches, its patch grid and its tokens (gh * gw / 4). +pub const Prepared = struct { + pixels: []f32, // [gh * gw, 1176] + gh: u32, + gw: u32, + tokens: u32, + hash: u64, // its identity (core/picture/identity.zig): bytes and grid, so another budget is another image +}; + +// ---- the processor's geometry ---- + +pub const Geometry = struct { target_h: u32, target_w: u32, content_h: u32, content_w: u32 }; + +fn alignUp(v: u64, f: u64) u64 { + return (v + f - 1) / f * f; +} + +/// smart_resize and _resize_geometry: the padded canvas under `max_tokens` (at least `min_tokens`), the content size. +pub fn geometry(h: u32, w: u32, min_tokens: u32, max_tokens: u32) error{ImageTooSmall}!Geometry { + const factor: u64 = vision.patch * vision.merge; + const frames: u64 = vision.temporal; + const per_token = frames * factor * factor; + const min_px = @as(u64, min_tokens) * per_token; + const max_px = @as(u64, max_tokens) * per_token; + if (min_tokens == 0 or max_tokens < min_tokens) return error.ImageTooSmall; + const H: u64 = h; + const W: u64 = w; + var ah = alignUp(H, factor); + var aw = alignUp(W, factor); + var pixels = frames * ah * aw; + if (pixels < min_px) { + const scale = @sqrt(@as(f64, @floatFromInt(min_px)) / @as(f64, @floatFromInt(frames * H * W))); + ah = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(H)) * scale)))), factor); + aw = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(W)) * scale)))), factor); + pixels = frames * ah * aw; + } + if (pixels > max_px) { // the tallest content that fits, by bisection on its height + if (max_px < frames * factor * factor) return error.ImageTooSmall; + var best: [2]u64 = .{ factor, factor }; + var low: i64 = 1; + var high: i64 = @intCast(H); + while (low <= high) { + const ch: u64 = @intCast(@divFloor(low + high, 2)); + const cw: u64 = @max(1, W * ch / H); // math.floor(width * content_h / height) + const cand: [2]u64 = .{ alignUp(ch, factor), alignUp(cw, factor) }; + if (frames * cand[0] * cand[1] <= max_px) { + best = cand; + low = @as(i64, @intCast(ch)) + 1; + } else high = @as(i64, @intCast(ch)) - 1; + } + ah = best[0]; + aw = best[1]; + } + var scale = @min(@as(f64, @floatFromInt(ah)) / @as(f64, @floatFromInt(H)), @as(f64, @floatFromInt(aw)) / @as(f64, @floatFromInt(W))); + if (frames * H * W >= min_px) scale = @min(1.0, scale); + const chf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(H)) * scale)); + const cwf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(W)) * scale)); + return .{ .target_h = @intCast(ah), .target_w = @intCast(aw), .content_h = @intCast(@max(1, @min(ah, chf))), .content_w = @intCast(@max(1, @min(aw, cwf))) }; +} + +/// The steps that read pixels, over `picture` (core/picture: the Metal build's relative import, a CUDA build's core). +pub fn For(comptime picture: type) type { + return struct { + /// The shared decoder's refusals, and a token budget too small for an image. + pub const Error = picture.Error || error{ImageTooSmall}; + + /// The words for a refusal the server cannot word: the JPEG library this machine lacks, or a CMYK JPEG. + pub const decoderProblem = picture.decoderProblem; + + /// The resized content padded black, scaled, normalized, cut into 2x2 windows of 14x14 patches, row-major. + pub fn patchify(gpa: std.mem.Allocator, content: picture.Rgb, g: Geometry, mean: [3]f32, sd: [3]f32) Error![]f32 { + const P = vision.patch; + const M = vision.merge; + const gh = g.target_h / P; + const gw = g.target_w / P; + const out = try gpa.alloc(f32, @as(usize, gh) * gw * vision.patch_width); + const rescale: f32 = @floatCast(@as(f64, 1.0) / 255.0); + var row: usize = 0; + for (0..gh / M) |wy| for (0..gw / M) |wx| for (0..M) |my| for (0..M) |mx| { + const py0 = (wy * M + my) * P; + const px0 = (wx * M + mx) * P; + const dst = out[row * vision.patch_width ..][0..vision.patch_width]; + for (0..3) |ch| for (0..vision.temporal) |t| for (0..P) |dy| for (0..P) |dx| { + const y = py0 + dy; + const x = px0 + dx; + const v: u8 = if (y < content.height and x < content.width) content.pixels[(y * content.width + x) * 3 + ch] else 0; + const f = @as(f32, @floatFromInt(v)) * rescale; + dst[((ch * vision.temporal + t) * P + dy) * P + dx] = (f - mean[ch]) / sd[ch]; + }; + row += 1; + }; + return out; + } + + /// Bytes to patches: decoded as 0.6.6 decoded them, sized for `max_tokens`, resized, padded and cut. + pub fn prepare(io: std.Io, gpa: std.mem.Allocator, bytes: []const u8, c: vision.Config, max_tokens: u32) Error!Prepared { + var img = try picture.decode(io, gpa, bytes); + defer img.deinit(gpa); + return prepareRgb(gpa, img, picture.identity.ofBytes(0x1a6e, bytes), c, max_tokens); + } + + /// RGB to patches (Pillow's bicubic resize, the padded canvas cut); the hash covers the bytes and the grid. + pub fn prepareRgb(gpa: std.mem.Allocator, img: picture.Rgb, bytes_hash: u64, c: vision.Config, max_tokens: u32) Error!Prepared { + const g = try geometry(@intCast(img.height), @intCast(img.width), @min(c.min_tokens, max_tokens), max_tokens); + const resized = try picture.resample.resize(gpa, .pillow, .of(img.pixels, img.width, img.height), g.content_w, g.content_h); + defer gpa.free(resized); + const content: picture.Rgb = .{ .width = g.content_w, .height = g.content_h, .pixels = resized }; + const pixels = try patchify(gpa, content, g, c.mean, c.std); + const grid = [2]u32{ g.target_h / vision.patch, g.target_w / vision.patch }; + return .{ .pixels = pixels, .gh = grid[0], .gw = grid[1], .tokens = grid[0] * grid[1] / vision.patches_per_token, .hash = picture.identity.withGrid(bytes_hash, grid) }; + } + }; +} diff --git a/zig/src/families/glm/vision.zig b/zig/src/families/glm/vision.zig index 3665d63f2d..6e6c8e9c86 100644 --- a/zig/src/families/glm/vision.zig +++ b/zig/src/families/glm/vision.zig @@ -10,86 +10,21 @@ const readAll = @import("weights_fix.zig").readAll; const Engine = @import("engine.zig").Engine; const Ref = wts.Ref; -// The tower this file reads; `Config.parse` refuses a checkpoint whose vision_config differs. -pub const hidden = 1024; -pub const heads = 16; -pub const head_dim = 64; -pub const depth = 24; -pub const mlp_inter = 4096; -pub const out_hidden = 4096; // the language model's width -pub const merger_inter = 10240; -pub const patch = 14; -pub const temporal = 2; -pub const merge = 2; -pub const patch_width = 3 * temporal * patch * patch; // 1176: a patch's values, channel by frame by row by column -pub const patches_per_token = merge * merge; - -/// --vision's limits: image tokens a prompt may hold in all (scratch and rows are sized for them), and images. -pub const Limits = struct { image_tokens: u32 = 4096, max_images: u32 = 4 }; - -pub const Config = struct { - eps: f32 = 1e-5, - limit: f32 = 10, // swiglu_limit - image_token: u32, - image_start: u32, - image_end: u32, - video: ?[3]u32 = null, // video_start_token_id, video_token_id, video_end_token_id: the checkpoint reads videos - mean: [3]f32, - std: [3]f32, - min_tokens: u32, - max_tokens: u32, - - /// config.json's vision_config and image token ids, processor_config.json's image processor; null: no tower. - pub fn parse(arena: std.mem.Allocator, config_json: []const u8, processor_json: ?[]const u8) !?Config { - const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, config_json, .{}); - const vc = doc.object.get("vision_config") orelse return null; - if (vc != .object) return null; - const want = .{ .{ "depth", depth }, .{ "hidden_size", hidden }, .{ "num_heads", heads }, .{ "intermediate_size", mlp_inter }, .{ "out_hidden_size", out_hidden }, .{ "projection_intermediate_size", merger_inter }, .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "spatial_merge_size", merge }, .{ "in_channels", 3 } }; - inline for (want) |kv| { - const v = vc.object.get(kv[0]) orelse return error.VisionConfig; - if (v != .integer or v.integer != kv[1]) { - std.log.err("glm vision: vision_config.{s} is {f}; this tower reads {d}", .{ kv[0], std.json.fmt(v, .{}), kv[1] }); - return error.VisionConfig; - } - } - const act = vc.object.get("hidden_act") orelse return error.VisionConfig; - if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.VisionConfig; - var c: Config = .{ .image_token = 0, .image_start = 0, .image_end = 0, .mean = .{ 0.48145466, 0.4578275, 0.40821073 }, .std = .{ 0.26862954, 0.26130258, 0.27577711 }, .min_tokens = 16, .max_tokens = 8000 }; - if (vc.object.get("rms_norm_eps")) |v| c.eps = @floatCast(number(v) orelse return error.VisionConfig); - if (vc.object.get("swiglu_limit")) |v| c.limit = @floatCast(number(v) orelse return error.VisionConfig); - inline for (.{ .{ "image_token_id", "image_token" }, .{ "image_start_token_id", "image_start" }, .{ "image_end_token_id", "image_end" } }) |kv| { - const v = doc.object.get(kv[0]) orelse return error.VisionConfig; - if (v != .integer) return error.VisionConfig; - @field(c, kv[1]) = @intCast(v.integer); - } - if (doc.object.get("video_start_token_id")) |s| if (doc.object.get("video_token_id")) |t| if (doc.object.get("video_end_token_id")) |e| { - if (s == .integer and t == .integer and e == .integer) c.video = .{ @intCast(s.integer), @intCast(t.integer), @intCast(e.integer) }; - }; - const pj = processor_json orelse return c; - const pdoc = try std.json.parseFromSliceLeaky(std.json.Value, arena, pj, .{}); - const ip = pdoc.object.get("image_processor") orelse return c; - if (ip != .object) return error.VisionConfig; - inline for (.{ .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "merge_size", merge } }) |kv| { - if (ip.object.get(kv[0])) |v| if (v != .integer or v.integer != kv[1]) return error.VisionConfig; - } - if (ip.object.get("patch_expand_factor")) |v| if (v != .integer or v.integer != 1) return error.VisionConfig; - inline for (.{ "image_mean", "image_std" }, .{ &c.mean, &c.std }) |key, dst| if (ip.object.get(key)) |v| { - if (v != .array or v.array.items.len != 3) return error.VisionConfig; - for (v.array.items, 0..) |x, i| dst[i] = @floatCast(number(x) orelse return error.VisionConfig); - }; - if (ip.object.get("min_image_tokens")) |v| c.min_tokens = @intCast(v.integer); - if (ip.object.get("max_image_tokens")) |v| c.max_tokens = @intCast(v.integer); - return c; - } - - fn number(v: std.json.Value) ?f64 { - return switch (v) { - .float => |f| f, - .integer => |i| @floatFromInt(i), - else => null, - }; - } -}; +const vcfg = @import("vision_config.zig"); +pub const hidden = vcfg.hidden; +pub const heads = vcfg.heads; +pub const head_dim = vcfg.head_dim; +pub const depth = vcfg.depth; +pub const mlp_inter = vcfg.mlp_inter; +pub const out_hidden = vcfg.out_hidden; +pub const merger_inter = vcfg.merger_inter; +pub const patch = vcfg.patch; +pub const temporal = vcfg.temporal; +pub const merge = vcfg.merge; +pub const patch_width = vcfg.patch_width; +pub const patches_per_token = vcfg.patches_per_token; +pub const Limits = vcfg.Limits; +pub const Config = vcfg.Config; const Kernels = struct { lib: mtl.Library, diff --git a/zig/src/families/glm/vision_config.zig b/zig/src/families/glm/vision_config.zig new file mode 100644 index 0000000000..6b9ce215c1 --- /dev/null +++ b/zig/src/families/glm/vision_config.zig @@ -0,0 +1,83 @@ +//! GLM-5.3-Flash's image tower: its shape and config reading, Metal-free (every backend's image.zig reads it). +const std = @import("std"); + +// The tower this file reads; `Config.parse` refuses a checkpoint whose vision_config differs. +pub const hidden = 1024; +pub const heads = 16; +pub const head_dim = 64; +pub const depth = 24; +pub const mlp_inter = 4096; +pub const out_hidden = 4096; // the language model's width +pub const merger_inter = 10240; +pub const patch = 14; +pub const temporal = 2; +pub const merge = 2; +pub const patch_width = 3 * temporal * patch * patch; // 1176: a patch's values, channel by frame by row by column +pub const patches_per_token = merge * merge; + +/// --vision's limits: image tokens a prompt may hold in all (scratch and rows are sized for them), and images. +pub const Limits = struct { image_tokens: u32 = 4096, max_images: u32 = 4 }; + +pub const Config = struct { + eps: f32 = 1e-5, + limit: f32 = 10, // swiglu_limit + image_token: u32, + image_start: u32, + image_end: u32, + video: ?[3]u32 = null, // video_start_token_id, video_token_id, video_end_token_id: the checkpoint reads videos + mean: [3]f32, + std: [3]f32, + min_tokens: u32, + max_tokens: u32, + + /// config.json's vision_config and image token ids, processor_config.json's image processor; null: no tower. + pub fn parse(arena: std.mem.Allocator, config_json: []const u8, processor_json: ?[]const u8) !?Config { + const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, config_json, .{}); + const vc = doc.object.get("vision_config") orelse return null; + if (vc != .object) return null; + const want = .{ .{ "depth", depth }, .{ "hidden_size", hidden }, .{ "num_heads", heads }, .{ "intermediate_size", mlp_inter }, .{ "out_hidden_size", out_hidden }, .{ "projection_intermediate_size", merger_inter }, .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "spatial_merge_size", merge }, .{ "in_channels", 3 } }; + inline for (want) |kv| { + const v = vc.object.get(kv[0]) orelse return error.VisionConfig; + if (v != .integer or v.integer != kv[1]) { + std.log.err("glm vision: vision_config.{s} is {f}; this tower reads {d}", .{ kv[0], std.json.fmt(v, .{}), kv[1] }); + return error.VisionConfig; + } + } + const act = vc.object.get("hidden_act") orelse return error.VisionConfig; + if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.VisionConfig; + var c: Config = .{ .image_token = 0, .image_start = 0, .image_end = 0, .mean = .{ 0.48145466, 0.4578275, 0.40821073 }, .std = .{ 0.26862954, 0.26130258, 0.27577711 }, .min_tokens = 16, .max_tokens = 8000 }; + if (vc.object.get("rms_norm_eps")) |v| c.eps = @floatCast(number(v) orelse return error.VisionConfig); + if (vc.object.get("swiglu_limit")) |v| c.limit = @floatCast(number(v) orelse return error.VisionConfig); + inline for (.{ .{ "image_token_id", "image_token" }, .{ "image_start_token_id", "image_start" }, .{ "image_end_token_id", "image_end" } }) |kv| { + const v = doc.object.get(kv[0]) orelse return error.VisionConfig; + if (v != .integer) return error.VisionConfig; + @field(c, kv[1]) = @intCast(v.integer); + } + if (doc.object.get("video_start_token_id")) |s| if (doc.object.get("video_token_id")) |t| if (doc.object.get("video_end_token_id")) |e| { + if (s == .integer and t == .integer and e == .integer) c.video = .{ @intCast(s.integer), @intCast(t.integer), @intCast(e.integer) }; + }; + const pj = processor_json orelse return c; + const pdoc = try std.json.parseFromSliceLeaky(std.json.Value, arena, pj, .{}); + const ip = pdoc.object.get("image_processor") orelse return c; + if (ip != .object) return error.VisionConfig; + inline for (.{ .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "merge_size", merge } }) |kv| { + if (ip.object.get(kv[0])) |v| if (v != .integer or v.integer != kv[1]) return error.VisionConfig; + } + if (ip.object.get("patch_expand_factor")) |v| if (v != .integer or v.integer != 1) return error.VisionConfig; + inline for (.{ "image_mean", "image_std" }, .{ &c.mean, &c.std }) |key, dst| if (ip.object.get(key)) |v| { + if (v != .array or v.array.items.len != 3) return error.VisionConfig; + for (v.array.items, 0..) |x, i| dst[i] = @floatCast(number(x) orelse return error.VisionConfig); + }; + if (ip.object.get("min_image_tokens")) |v| c.min_tokens = @intCast(v.integer); + if (ip.object.get("max_image_tokens")) |v| c.max_tokens = @intCast(v.integer); + return c; + } + + fn number(v: std.json.Value) ?f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => null, + }; + } +}; diff --git a/zig/tests/cuda/glm_vision_tests.zig b/zig/tests/cuda/glm_vision_tests.zig new file mode 100644 index 0000000000..f7d5968b54 --- /dev/null +++ b/zig/tests/cuda/glm_vision_tests.zig @@ -0,0 +1,104 @@ +//! GLM's image tower on CUDA against vision_ref.py: the patches, and features as near float32 as torch's bf16 run. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm_cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +fn readF32(gpu: Gpu, dir: []const u8, name: []const u8) ![]f32 { + const path = try std.fs.path.join(gpu.gpa, &.{ dir, name }); + defer gpu.gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(1 << 30)); + defer gpu.gpa.free(bytes); + const out = try gpu.gpa.alloc(f32, bytes.len / 4); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * 4]); + return out; +} + +/// Root-mean-square of `want - got` over the rms of `want`, and the largest difference over that rms. +fn errors(want: []const f32, got: []const f32) [2]f64 { + var d2: f64 = 0; + var w2: f64 = 0; + var most: f64 = 0; + for (want, got) |w, g| { + const d: f64 = @as(f64, g) - w; + d2 += d * d; + w2 += @as(f64, w) * w; + most = @max(most, @abs(d)); + } + const rms = @sqrt(w2 / @as(f64, @floatFromInt(want.len))); + return .{ @sqrt(d2 / @as(f64, @floatFromInt(want.len))) / rms, most / rms }; +} + +pub fn run(gpu: Gpu, model: []const u8, case: []const u8, bf16_case: []const u8, image: []const u8) !void { + const a = gpu.gpa; + var arena: std.heap.ArenaAllocator = .init(a); + defer arena.deinit(); + const cwd = std.Io.Dir.cwd(); + const cfg_path = try std.fs.path.join(arena.allocator(), &.{ model, "config.json" }); + const proc_path = try std.fs.path.join(arena.allocator(), &.{ model, "processor_config.json" }); + const cfg = try cwd.readFileAlloc(gpu.io, cfg_path, arena.allocator(), .limited(1 << 24)); + const proc = try cwd.readFileAlloc(gpu.io, proc_path, arena.allocator(), .limited(1 << 24)); + const c = (try glm.vision_config.Config.parse(arena.allocator(), cfg, proc)) orelse return error.NoVisionTower; + const want_patches = try readF32(gpu, case, "patches.f32"); + defer a.free(want_patches); + const want = try readF32(gpu, case, "features.f32"); + defer a.free(want); + const own = try readF32(gpu, bf16_case, "features.f32"); // the torch tower in bf16: its own distance from float32 + defer a.free(own); + const bound = errors(want, own); + + // preparation: core's decoder and resize, the processor's canvas and merge order + const bytes = try cwd.readFileAlloc(gpu.io, image, arena.allocator(), .limited(64 << 20)); + const p = try glm.image.prepare(gpu.io, a, bytes, c, c.max_tokens); + defer a.free(p.pixels); + try check.expect(p.pixels.len == want_patches.len, "glm-vision: {d} patch values, the reference {d}", .{ p.pixels.len, want_patches.len }); + const pe = errors(want_patches, p.pixels); + std.debug.print("RESULT glm-vision patches {d} x {d} ({d} tokens): rms difference {e:.2} of their rms, largest {e:.2}\n", .{ p.gh, p.gw, p.tokens, pe[0], pe[1] }); + try check.expect(pe[1] < 0.02, "glm-vision: the patches as the processor makes them", .{}); + + // the tower on the reference's patches + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_vision); + defer mod.unload(); + const ru: glm.vision.Run = .{ .stream = s, .f = try glm.vision.Funcs.load(mod) }; + var t = try glm.vision.Tower.load(a, gpu.io, gpu.d, model); + defer t.deinit(); + var out = try cuda.DeviceBuffer.alloc(gpu.d, @as(usize, p.tokens) * glm.vision_config.out_hidden * 2); + defer out.free(); + const grid = [3]i64{ 1, p.gh, p.gw }; + { // where the stream parts from the reference's: after the embedding, then each block + const all = readF32(gpu, case, "all.f32") catch null; + const embed = readF32(gpu, case, "embed.f32") catch null; + if (all != null and embed != null) { + defer a.free(all.?); + defer a.free(embed.?); + const n = embed.?.len; + const got = try a.alloc(f32, (glm.vision.DEPTH + 1) * n); + defer a.free(got); + t.stages = got; + try t.encode(&ru, want_patches, grid, out.ptr); + t.stages = null; + const e0 = errors(embed.?, got[0..n]); + std.debug.print("STAGE embed: rms {d:.4}\n", .{e0[0]}); + for (0..glm.vision.DEPTH) |b| { + const e = errors(all.?[b * n ..][0..n], got[(b + 1) * n ..][0..n]); + std.debug.print("STAGE block {d}: rms {d:.4}, largest {d:.3}\n", .{ b, e[0], e[1] }); + } + } + } + for ([_][]const f32{ want_patches, p.pixels }, [_][]const u8{ "the reference's patches", "its own patches" }) |patches, what| { + try t.encode(&ru, patches, grid, out.ptr); + const raw = try a.alloc(u16, @as(usize, p.tokens) * glm.vision_config.out_hidden); + defer a.free(raw); + try out.download(0, std.mem.sliceAsBytes(raw)); + const got = try a.alloc(f32, raw.len); + defer a.free(got); + for (got, raw) |*g, r| g.* = @bitCast(@as(u32, r) << 16); + const fe = errors(want, got); + std.debug.print("RESULT glm-vision features from {s}: rms difference {d:.4} of their rms ({d:.3}x the torch bf16 run's {d:.4}), largest {d:.3}\n", .{ what, fe[0], fe[0] / bound[0], bound[0], fe[1] }); + try check.expect(fe[0] <= 1.1 * bound[0], "glm-vision: features as close to float32 as the torch tower's bf16 run", .{}); + } + check.pass("glm-vision: the processor's patches as the reference's, the features as near float32 as torch's bf16 run", .{}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 8ed3f0fc3d..73a49fda08 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -23,6 +23,7 @@ const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); const glm_ref_tests = @import("glm_ref_tests.zig"); const glm_exl3_tests = @import("glm_exl3_tests.zig"); +const glm_vision_tests = @import("glm_vision_tests.zig"); const carveout_tests = @import("carveout_tests.zig"); const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); @@ -94,6 +95,7 @@ const usage = \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) + \\ glm-vision MODEL F32 BF16 IMG GLM's image tower and preparation against vision_ref.py's torch runs \\ glm-ref GLM projections (matmuls, router, attention, indexer, hc) against float64 (synthetic) \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run @@ -163,6 +165,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-ref")) return glm_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-exl3")) return glm_exl3_tests.run_(gpu); + if (std.mem.eql(u8, cmd, "glm-vision")) return glm_vision_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); From f9d4f9fd23eb1d7c11465c217ad6388614b50f0b Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 14:39:07 +0000 Subject: [PATCH 28/31] glm cuda: images on two ranks (--vision): rank 0's tower encodes a prompt's images, an image step sends their spans and rows to the other ranks, every rank splices them over the placeholders after the embedding (the MTP head's one place on) on that sequence only; the ranks agree on --vision before loading; the server's CUDA image offer takes any family's tower (a template that wraps images itself gets no server wrap) --- zig/src/families/glm/cuda_engine.zig | 15 ++- zig/src/families/glm/cuda_follow.zig | 27 ++++- zig/src/families/glm/cuda_forward.zig | 8 ++ zig/src/families/glm/cuda_images.zig | 166 ++++++++++++++++++++++++++ zig/src/families/glm/cuda_lanes.zig | 11 ++ zig/src/families/glm/cuda_native.zig | 14 ++- zig/src/families/glm/cuda_ranks.zig | 9 +- zig/src/native/cuda_vision_offer.zig | 12 +- 8 files changed, 251 insertions(+), 11 deletions(-) create mode 100644 zig/src/families/glm/cuda_images.zig diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index f42cf02e18..908ed93333 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -12,11 +12,13 @@ const sampler = @import("cuda_sampler.zig"); const heat = @import("heat"); const modelopt = @import("core").modelopt; const exl3 = @import("cuda_exl3.zig"); +const images_mod = @import("cuda_images.zig"); pub const Options = struct { capacity: usize, // tokens a sequence holds mtp: bool, // load the MTP head (drafts) settings: ranks.Settings, + vision: u32 = 0, // --vision's image tokens (0: none): every rank must hold the rows rank 0 sends }; /// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. @@ -67,6 +69,7 @@ pub const Engine = struct { lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, group: ?*const cuda.tp_group.Group = null, // the main stream's all-gathers (RoCE on two ranks); null: lib and comm + images: ?*images_mod.Images = null, // --vision: a prompt pass's image rows (rank 0 encodes, the others receive) settings: ranks.Settings, m: fw.Modules, stream: cuda.Stream, @@ -112,7 +115,7 @@ pub const Engine = struct { errdefer stream.deinit(); const config_hash = try ranks.hashFile(a, io, cfg_path); const program_hash = try ranks.hashFile(a, io, "/proc/self/exe"); - const fp = ranks.Fingerprint.of( + var fp = ranks.Fingerprint.of( world, rank, o.settings, @@ -122,6 +125,7 @@ pub const Engine = struct { program_hash, checkpoint, ); + fp.vision = o.vision; try ranks.agree(d, lib, comm, stream, fp, try a.alloc(ranks.Fingerprint, world)); const share = try shard.Share.of(&c, rank, world); const spans = try gpa.alloc(shard.Span, world); @@ -193,7 +197,8 @@ pub const Engine = struct { fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { out.* = .{.{ .w = &e.w, .st = &e.cur.st, .b = b }}; - return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .kda_wide = e.settings.kda_wide }; + const im: ?fw.ImageRows = if (e.images) |v| .{ .v = v, .slot = e.slot } else null; + return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .kda_wide = e.settings.kda_wide, .images = im }; } /// Binds slot `slot`'s sequence for the steps that follow, making it (a zeroed state) on first use. @@ -234,6 +239,12 @@ pub const Engine = struct { return e.cur.st.pos; } + /// A follower's begin: rank 0's next request; no image rows from an earlier one splice into it. + pub fn beginFollowing(e: *Engine, sampling: ?sampler.Sampling) !void { + if (e.images) |v| v.clear(); + return e.begin(sampling); + } + /// A fresh sequence with `sampling` (null or temperature 0: greedy). pub fn begin(e: *Engine, sampling: ?sampler.Sampling) !void { e.cur.sampling = sampler.check(sampling); diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 8072b7cfe1..1f31bad634 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -3,10 +3,12 @@ const std = @import("std"); const cuda = @import("cuda"); const engine = @import("cuda_engine.zig"); +const images_mod = @import("cuda_images.zig"); +const vcfg = @import("vision_config.zig"); const sampler = @import("cuda_sampler.zig"); const Engine = engine.Engine; -pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared }; +pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image }; /// One step as it crosses: 64 bytes, then `n` u32 ids. pub const Header = extern struct { @@ -88,6 +90,12 @@ pub const Link = struct { } } + /// `bytes` at `ptr` from rank 0 into every rank's `ptr` (image rows: the follower's buffer is the same size). + pub fn rows(l: *Link, ptr: u64, bytes: usize) !void { + if (bytes == 0) return; + try l.lib.check(l.lib.api.ncclBroadcast(ptr, ptr, bytes, .u8, 0, l.comm, l.stream.handle), "ncclBroadcast"); + } + /// Other ranks: the next step and its ids (valid until the next receive). pub fn receive(l: *Link) !struct { h: Header, ids: []const u32 } { try l.bcast(64); @@ -153,6 +161,16 @@ pub const Leader = struct { try x.e.begin(c); } + /// The prompt pass's image rows (rank 0 encoded them for the bound slot) to every rank: spans, then the rows. + pub fn images(x: Leader, v: *const images_mod.Images) !void { + try x.flush(); + var words: [3 * 64]u32 = undefined; + if (v.spans.items.len * 3 > words.len) return error.TooManyImages; + for (v.spans.items, 0..) |sp, i| words[3 * i ..][0..3].* = .{ sp.at, sp.n, sp.row }; + try x.send(.{ .kind = @backingInt(Kind.image) }, words[0 .. 3 * v.spans.items.len]); + try x.link.rows(v.rows.ptr, v.rowCount() * vcfg.out_hidden * 2); + } + /// `cancel` (rank 0's), asked before each chunk, decides for every rank; `from` every rank holds; `marks` rank 0's. pub fn prefill(x: Leader, ids: []const u32, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !u32 { try x.flush(); @@ -220,7 +238,7 @@ pub fn follow(e: *Engine, link: *Link) !void { const kind: Kind = @fromBackingInt(@intCast(h.kind)); if (kind != .end and kind != .drop and kind != .shared) try e.bind(h.slot); switch (kind) { - .begin => try e.begin(if (h.flags & 1 != 0) .{ .seed = h.seed, .temperature = h.temperature, .top_p = h.top_p, .min_p = h.min_p, .top_k = h.top_k } else null), + .begin => try e.beginFollowing(if (h.flags & 1 != 0) .{ .seed = h.seed, .temperature = h.temperature, .top_p = h.top_p, .min_p = h.min_p, .top_k = h.top_k } else null), .prefill => { const ids = try e.gpa.dupe(u32, step.ids); defer e.gpa.free(ids); @@ -262,6 +280,11 @@ pub fn follow(e: *Engine, link: *Link) !void { defer e.gpa.free(ids); try e.draft(h.flags & 1 != 0, ids, h.depth, &drafts); }, + .image => { + const v = e.images orelse return error.NoVision; // the ranks agreed on --vision (cuda_ranks) + try v.take(h.slot, step.ids); + try link.rows(v.rows.ptr, v.rowCount() * vcfg.out_hidden * 2); + }, .drop => for (step.ids) |id| e.dropState(id), .release => e.release(h.slot), .end => return, diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index 893677e276..56e4de8743 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/zig/src/families/glm/cuda_forward.zig @@ -17,6 +17,7 @@ const moe = @import("cuda_moe.zig"); const quant = @import("cuda_quant.zig"); const q4 = @import("cuda_q4.zig"); const state = @import("cuda_state.zig"); +const images_mod = @import("cuda_images.zig"); const exl3 = @import("cuda_exl3.zig"); const shared = @import("cuda_shared.zig"); @@ -150,6 +151,9 @@ pub const Rank = struct { segs: []const Seg = &.{}, }; +/// The engine's image rows and the sequence the forward runs on (they splice into that sequence's rows only). +pub const ImageRows = struct { v: *images_mod.Images, slot: u32 }; + /// One stream's rows of a shared window: rows row0..row0 + rows, on sequence `st` from its position. pub const Seg = struct { st: *State, row0: usize, rows: usize }; @@ -162,6 +166,7 @@ pub const Forward = struct { kda_wide: bool = true, // TF_GLM_KDA_DECODE_WIDE (the Python family's default): decode windows on KDA's three-kernel chain pieces: ?*const Pieces = null, // a prompt chunk's exchanges in row pieces on a second stream (null: whole) split: bool = true, // TF_GLM_PREFILL_REDUCE=split (the default with send/recv): each rank glues its own rows + images: ?ImageRows = null, // a prompt pass's image rows over their placeholders' embeddings (its sequence only) pub fn set(f: *const Forward) glue.Set { return .{ .glue = &f.m.glue }; @@ -456,6 +461,7 @@ pub const Forward = struct { } const c = &f.ranks[0].w.config; for (f.ranks) |r| try glue.embedB16(f.set(), f.stream, r.b.ids, r.w.embed, r.b.x, R, c.hidden, c.hc); + if (f.images) |im| for (f.ranks) |r| try im.v.splice(f.stream, im.slot, r.st.pos, r.st.pos + R, r.b.x, c.hc); var cut: [max_pieces][2]usize = undefined; const pieces = if (f.pieces) |p| (if (f.ranks[0].b.prefill) ranges(R, p.count, &cut) else 0) else 0; // overlap.Overlap.begin: pieced with two pieces or more, and split takes a chunk of one piece too @@ -517,6 +523,8 @@ pub const Forward = struct { const m = r.w.mtp.?; const b = r.b; try glue.embedB16(f.set(), f.stream, b.ids, r.w.embed, b.me, n, D, 1); + // the head reads each row's next token: an image's rows one place on + if (f.images) |im| try im.v.splice(f.stream, im.slot, r.st.mtp_len + 1, r.st.mtp_len + 1 + n, b.me, 1); if (r.st.mtp_len == 0) try f.d.check(f.d.api.cuMemsetD8Async(b.me, 0, D * 2, f.stream.handle), "cuMemsetD8Async"); // position 0's embedding is zeroed try glue.rmsnorm(f.set(), f.stream, b.me, D, m.enorm, b.mcat, 2 * D, null, c.eps, n, D); try glue.rmsnorm(f.set(), f.stream, b.hin, D, m.hnorm, b.mcat + D * 2, 2 * D, null, c.eps, n, D); diff --git a/zig/src/families/glm/cuda_images.zig b/zig/src/families/glm/cuda_images.zig new file mode 100644 index 0000000000..cfeb019864 --- /dev/null +++ b/zig/src/families/glm/cuda_images.zig @@ -0,0 +1,166 @@ +//! A prompt pass's image rows on every rank: rank 0's tower encodes them, the others receive them, all splice them. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const vision = @import("cuda_vision.zig"); +const vcfg = @import("vision_config.zig"); +const image_core = @import("image_core.zig"); +const image = image_core.For(@import("core").picture); + +/// What a request may hold: image tokens in all (the rows buffer), and images. +pub const Limits = struct { image_tokens: u32, max_images: u32 }; + +/// Prompt rows [at, at + n) take image rows [row, row + n). +pub const Span = struct { at: u32, n: u32, row: u32 }; + +pub const Images = struct { + gpa: std.mem.Allocator, + io: std.Io, + d: *const cuda.Driver, + module: cuda.Module, + run: vision.Run, + tower: ?vision.Tower = null, // rank 0's; the other ranks receive the rows + cfg: ?vcfg.Config = null, + rows: cuda.DeviceBuffer, // bf16 [limits.image_tokens][out_hidden] + map: cuda.DeviceBuffer, // i32 [max_rows]: a splice's feature row per prompt row + host_map: []i32, + limits: Limits, + spans: std.ArrayList(Span) = .empty, + slot: u32 = 0, // the sequence the spans belong to: another stream's rows at the same positions take none + marks: struct { image: u32, open: u32, close: u32 } = .{ .image = 0, .open = 0, .close = 0 }, + mark_ids: [3]u32 = @splat(0), // begin, image, end: text that spells one stays text + + /// The chat template writes each image's begin and end itself (the server only expands the image token). + pub const template_wraps = true; + + /// Rank 0 (`tower`) or another rank: the rows buffer, and rank 0's tower; null with `problem` when no tower. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, stream: cuda.Stream, dir: []const u8, limits: Limits, max_rows: usize, tower: bool, problem: *[]const u8) !?*Images { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const cwd = std.Io.Dir.cwd(); + const cfg_text = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "config.json" }), a, .limited(1 << 24)); + const proc = cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "processor_config.json" }), a, .limited(1 << 24)) catch null; + const cfg = (try vcfg.Config.parse(a, cfg_text, proc)) orelse { + problem.* = "this checkpoint has no vision tower: serve it without --vision"; + return null; + }; + const v = try gpa.create(Images); + errdefer gpa.destroy(v); + var module = try cuda.Module.load(d, cuda.kernels.glm_vision); + errdefer module.unload(); + v.* = .{ .gpa = gpa, .io = io, .d = d, .module = module, .run = .{ .stream = stream, .f = try vision.Funcs.load(module) }, .cfg = cfg, .rows = undefined, .map = undefined, .host_map = undefined, .limits = limits }; + v.rows = try cuda.DeviceBuffer.alloc(d, @as(usize, limits.image_tokens) * vcfg.out_hidden * 2); + errdefer v.rows.free(); + v.map = try cuda.DeviceBuffer.alloc(d, max_rows * 4); + errdefer v.map.free(); + v.host_map = try gpa.alloc(i32, max_rows); + errdefer gpa.free(v.host_map); + if (tower) v.tower = try vision.Tower.load(gpa, io, d, dir); + v.marks = .{ .image = cfg.image_token, .open = cfg.image_start, .close = cfg.image_end }; + v.mark_ids = .{ cfg.image_start, cfg.image_token, cfg.image_end }; + return v; + } + + pub fn deinit(v: *Images) void { + v.spans.deinit(v.gpa); + if (v.tower) |*t| t.deinit(); + v.gpa.free(v.host_map); + v.map.free(); + v.rows.free(); + v.module.unload(); + v.gpa.destroy(v); + } + + /// Device bytes the tower and the buffers hold. + pub fn bytes(v: *const Images) u64 { + return (if (v.tower) |t| t.bytes else 0) + v.rows.len + v.map.len; + } + + /// No rows splice into any pass until another begins (each request's prompt pass starts here). + pub fn clear(v: *Images) void { + v.spans.clearRetainingCapacity(); + } + + /// Rank 0: a prompt pass on `slot` from row `from`: the images whose rows end past it, through the tower. + pub fn begin(v: *Images, slot: u32, images: []const lanes.stream.Image, from: usize) !void { + v.clear(); + if (images.len > v.limits.max_images) return error.TooManyImages; + const t = &(v.tower orelse return error.NoVision); + var row: u32 = 0; + for (images) |img| { + if (img.gh * img.gw != img.tokens * vcfg.patches_per_token) return error.ImageOutOfStep; + if (img.at + img.tokens <= from) continue; // in the kept state already + if (row + img.tokens > v.limits.image_tokens) return error.ImageRows; + try t.encode(&v.run, img.pixels, .{ 1, img.gh, img.gw }, v.rows.ptr + @as(u64, row) * vcfg.out_hidden * 2); + try v.spans.append(v.gpa, .{ .at = img.at, .n = img.tokens, .row = row }); + row += img.tokens; + } + v.slot = slot; + } + + /// The other ranks: rank 0's spans (as `words`: at, n, row a span) on `slot`; the rows arrive into `rows` next. + pub fn take(v: *Images, slot: u32, words: []const u32) !void { + v.clear(); + if (words.len % 3 != 0) return error.BadStep; + var i: usize = 0; + while (i < words.len) : (i += 3) { + const sp: Span = .{ .at = words[i], .n = words[i + 1], .row = words[i + 2] }; + if (sp.row + sp.n > v.limits.image_tokens) return error.ImageRows; + try v.spans.append(v.gpa, sp); + } + v.slot = slot; + } + + /// Image rows the spans hold (what crosses to the other ranks). + pub fn rowCount(v: *const Images) usize { + var n: usize = 0; + for (v.spans.items) |sp| n = @max(n, sp.row + sp.n); + return n; + } + + /// Prompt rows [lo, hi) of `slot`, embedded at `emb` in `copies` streams, take the image rows they cross. + pub fn splice(v: *Images, s: cuda.Stream, slot: u32, lo: usize, hi: usize, emb: u64, copies: usize) !void { + if (v.spans.items.len == 0 or slot != v.slot) return; + const n = hi - lo; + if (n > v.host_map.len) return error.BadWindow; + var any = false; + for (v.host_map[0..n], lo..) |*m, p| { + m.* = -1; + for (v.spans.items) |sp| if (p >= sp.at and p < sp.at + sp.n) { + m.* = @intCast(sp.row + (p - sp.at)); + any = true; + }; + } + if (!any) return; + try v.map.upload(0, std.mem.sliceAsBytes(v.host_map[0..n])); // waits: the host map is reused next splice + var a: cuda.Args = .{}; + a.add(v.rows.ptr); + a.add(emb); + a.add(v.map.ptr); + a.add(@as(c_int, @intCast(n))); + a.add(@as(c_int, @intCast(vcfg.out_hidden))); + a.add(@as(c_int, @intCast(copies))); + try cuda.launch.launch(v.run.f.inject, .{ .grid = .{ .x = @intCast(n) }, .block = .{ .x = 256 } }, s, &a); + } + + /// The words for a prepare error the server cannot word (a JPEG library this machine lacks). + pub const explain = image.decoderProblem; + + /// An image's bytes as the tower reads them, at most `max_tokens` tokens (the server's prepare). + pub fn prepare(v: *const Images, a: std.mem.Allocator, data: []const u8, max_tokens: u32) !image_core.Prepared { + return image.prepare(v.io, a, data, v.cfg.?, max_tokens); + } +}; + +test "a splice's map: rows inside a span take its rows, others keep their embedding" { + var host: [8]i32 = undefined; + const spans = [_]Span{ .{ .at = 3, .n = 2, .row = 0 }, .{ .at = 6, .n = 4, .row = 2 } }; + for (&host, 2..) |*m, p| { + m.* = -1; + for (spans) |sp| if (p >= sp.at and p < sp.at + sp.n) { + m.* = @intCast(sp.row + (p - sp.at)); + }; + } + try std.testing.expectEqualSlices(i32, &.{ -1, 0, 1, -1, 2, 3, 4, 5 }, &host); +} diff --git a/zig/src/families/glm/cuda_lanes.zig b/zig/src/families/glm/cuda_lanes.zig index 38b6376a68..c257fca5ab 100644 --- a/zig/src/families/glm/cuda_lanes.zig +++ b/zig/src/families/glm/cuda_lanes.zig @@ -123,6 +123,7 @@ pub const Lanes = struct { fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { const self = of(ptr); const from = try self.opening(s); + try self.images(s, from); _ = self.take(try self.x.prefill(s.prompt(), from, .{ .ptr = s, .check = cancelled }, marksOf(s))); } @@ -131,6 +132,7 @@ pub const Lanes = struct { const self = of(ptr); const going: ?usize = if (self.lanes.getPtr(s)) |l| l.fill else null; const from = if (going == null) try self.opening(s) else null; + if (going == null) try self.images(s, from); const l = self.lanes.getPtr(s).?; if (going != null) try self.x.bind(l.slot); const at = going orelse if (from) |r| r.at else 0; @@ -154,11 +156,20 @@ pub const Lanes = struct { gop.value_ptr.held_n = 0; gop.value_ptr.fill = null; try self.x.bind(gop.value_ptr.slot); + if (self.x.e.images) |v| v.clear(); // no earlier request's rows splice into this one try self.x.begin(s.sampling); s.cached = 0; return self.resumeFrom(s, ids.len); } + /// A prompt pass with images: rank 0's tower encodes those past `from`, then every rank holds their rows. + fn images(self: *Lanes, s: *lanes.Stream, from: ?engine.Resume) !void { + if (s.images.len == 0) return; + const v = self.x.e.images orelse return error.NoVision; + try v.begin(self.x.e.slot, s.images, if (from) |r| r.at else 0); // a refusal here leaves every rank as it was + try self.x.images(v); + } + fn marksOf(s: *lanes.Stream) ?engine.Marks { return if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; } diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 5abb1905cb..a894a5bbb5 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); const follow = @import("cuda_follow.zig"); +const images_mod = @import("cuda_images.zig"); const glm_lanes = @import("cuda_lanes.zig"); const ranks = @import("cuda_ranks.zig"); const fw = @import("cuda_forward.zig"); @@ -31,6 +32,7 @@ pub const Options = struct { segments: usize = 1, slide: bool = false, room: ?Room = null, + vision: ?images_mod.Limits = null, // --vision: rank 0's tower, every rank's image rows }; /// The smallest room among the other ranks and whose it is. pub const Peer = struct { bytes: u64, rank: usize }; @@ -49,6 +51,7 @@ pub const Loaded = struct { lone: ?LoneRun = null, target: *engine.Engine, // the prompt cache keeps this engine's states (snap_*) peer_room: ?Peer = null, // rank 0: the server admits streams and kept states by this when it is smaller + vision: ?*images_mod.Images = null, // rank 0's tower for the server's image offer }; /// This rank's place: TF_GLM_RANK of TF_GLM_WORLD (both, or neither for a refusal that names them). @@ -141,8 +144,15 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer _ = own.group.lib.api.ncclCommDestroy(own.group.comm); if (p.world > 1) try warmPeers(ctx.d, &own.group); const cap = capacity(o.context); - own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, &own.group.lib, own.group.comm, p.rank, p.world, .{ .capacity = cap, .mtp = o.drafts, .settings = set }); + own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, &own.group.lib, own.group.comm, p.rank, p.world, .{ .capacity = cap, .mtp = o.drafts, .settings = set, .vision = if (o.vision) |l| l.image_tokens else 0 }); errdefer own.e.deinit(); + if (o.vision) |lim| { // before the rooms: the tower's bytes count against rank 0's + own.e.images = (try images_mod.Images.init(gpa, io, ctx.d, own.e.stream, dir, lim, fw.prefill_rows + 1, p.rank == 0, &problem)) orelse { + std.log.err("{s}", .{problem}); + return error.NoVisionTower; + }; + } + errdefer if (own.e.images) |v| v.deinit(); try own.group.openRoce(gpa, io, ctx.d, own.e.stream, roce_cap); // both ranks bring it up or neither does errdefer if (own.group.fast) |f| f.deinit(); own.e.group = &own.group; @@ -186,6 +196,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .deinit = close, .target = own.e, .peer_room = peer, + .vision = own.e.images, }; } @@ -230,6 +241,7 @@ fn close(ptr: *anyopaque) void { own.lanes.x.end() catch {}; own.lanes.deinit(); own.link.deinit(); + if (own.e.images) |v| v.deinit(); own.e.deinit(); own.group.close(); own.gpa.destroy(own); diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig index ab5b01b770..5d4d79f0c5 100644 --- a/zig/src/families/glm/cuda_ranks.zig +++ b/zig/src/families/glm/cuda_ranks.zig @@ -60,7 +60,7 @@ pub fn environment(name: [:0]const u8) ?[]const u8 { /// What every rank must agree on, as bytes the communicator carries. pub const Fingerprint = extern struct { - layout: u32 = 5, // this struct's version + layout: u32 = 6, // this struct's version world: u32, rank: u32, split: u32, @@ -68,6 +68,7 @@ pub const Fingerprint = extern struct { pieces: u32, layers: u32, // TF_GLM_LAYERS (0: all) mtp: u32, // 1: the MTP head loads and runs (drafts on): every prompt pass adds its collectives + vision: u32 = 0, // --vision's image tokens (0: off): rank 0 sends image rows every rank must hold capacity: u64, // tokens a sequence holds heat_high: f64, // the heat bands in degrees (0 and 0: off): every rank must wait alike heat_low: f64, @@ -84,6 +85,7 @@ pub const Fingerprint = extern struct { pieces, layers, mtp, + vision, capacity, heat_high, heat_low, @@ -138,6 +140,7 @@ pub fn mismatch(all: []const Fingerprint) ?Mismatch { "pieces", "layers", "mtp", + "vision", "capacity", "heat_high", "heat_low", @@ -213,6 +216,7 @@ fn explain(f: Fingerprint.Field) []const u8 { .pieces => "TF_GLM_OVERLAP_PIECES", .layers => "TF_GLM_LAYERS", .mtp => "--no-drafts given to some ranks and not others", + .vision => "--vision (or its image tokens) given to some ranks and not others", .capacity => "the sequence capacity (the context)", .heat_high => "TF_HEAT_HIGH", .heat_low => "TF_HEAT_LOW", @@ -264,4 +268,7 @@ test "fingerprints: agreeing ranks pass; the first differing field and rank are try std.testing.expectEqual(Fingerprint.Field.checkpoint, mismatch(&all).?.field); all[1] = Fingerprint.of(2, 1, .{}, 8192, false, h0, h0, "/hub/snapshots/423acf37"); // --no-drafts on rank 1 only try std.testing.expectEqual(Mismatch{ .field = .mtp, .rank = 1 }, mismatch(&all).?); + all[1] = Fingerprint.of(2, 1, .{}, 8192, true, h0, h0, "/hub/snapshots/423acf37"); + all[1].vision = 4096; // --vision on rank 1 only + try std.testing.expectEqual(Mismatch{ .field = .vision, .rank = 1 }, mismatch(&all).?); } diff --git a/zig/src/native/cuda_vision_offer.zig b/zig/src/native/cuda_vision_offer.zig index 56b8ac6d21..6b4439e8d4 100644 --- a/zig/src/native/cuda_vision_offer.zig +++ b/zig/src/native/cuda_vision_offer.zig @@ -1,18 +1,20 @@ //! What a CUDA family with an image tower offers the server: its processor, image token, markers and limits. const std = @import("std"); const api = @import("engine_api"); -const nemotron = @import("nemotron"); const Allocator = std.mem.Allocator; -/// Info.vision for a Nemotron Omni engine's tower; null without --vision. -pub fn offer(v: ?*nemotron.native.vision.Vision) ?api.Vision { +/// Info.vision for a family's tower (Nemotron Omni's, GLM-5.3-Flash's); null without --vision. +pub fn offer(v: anytype) ?api.Vision { const x = v orelse return null; + const T = @typeInfo(@TypeOf(x)).pointer.child; const prepare = struct { fn f(ctx: *anyopaque, a: Allocator, bytes: []const u8, max_tokens: u32) anyerror!api.PreparedImage { - const t: *nemotron.native.vision.Vision = @ptrCast(@alignCast(ctx)); + const t: *T = @ptrCast(@alignCast(ctx)); const p = try t.prepare(a, bytes, max_tokens); return .{ .pixels = p.pixels, .gh = p.gh, .gw = p.gw, .tokens = p.tokens, .hash = p.hash }; } }.f; - return .{ .ctx = x, .prepare = prepare, .image_token = x.marks.image, .markers = &x.mark_ids, .wrap = .{ x.marks.open, x.marks.close }, .image_tokens = x.limits.image_tokens, .max_images = x.limits.max_images, .explain = nemotron.native.vision.Vision.explain }; + // a template that writes each image's begin and end itself: the server expands the image token alone + const wrap: ?[2]u32 = if (@hasDecl(T, "template_wraps")) null else .{ x.marks.open, x.marks.close }; + return .{ .ctx = x, .prepare = prepare, .image_token = x.marks.image, .markers = &x.mark_ids, .wrap = wrap, .image_tokens = x.limits.image_tokens, .max_images = x.limits.max_images, .explain = T.explain }; } From 3cc45819be870d61b62f03715c9a61b91bf516f3 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 15:07:30 +0000 Subject: [PATCH 29/31] tools/glm_vision_ref.py: the torch reference tf-cuda-test glm-vision reads (float32, and --bf16 for its bound), from processor_config.json --- tools/glm_vision_ref.py | 213 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 213 insertions(+) create mode 100644 tools/glm_vision_ref.py diff --git a/tools/glm_vision_ref.py b/tools/glm_vision_ref.py new file mode 100644 index 0000000000..25c1c4234a --- /dev/null +++ b/tools/glm_vision_ref.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +"""GLM-5.3-Flash's vision tower as mlx-vlm's glm5_next (vision.py, processing.py) defines it, in torch on the CPU +(fp32 math on the checkpoint's bf16 weights): the reference the native tower is checked against. + +usage: glm_vision_ref.py [--bf16] ... + --bf16: the same tower in bfloat16 (tf-cuda-test glm-vision bounds the native tower's error by this run's) + an image: its patches (the GLM processor's canvas, padding, normalization and merge order) -> features + a folder of frame_*.png: a video's sampled frames (already sampled) -> one [t, h, w] grid +writes /case/{patches.f32, grid.json, features.f32, block0.f32} +""" +import json, math, struct, sys +from pathlib import Path + +import numpy as np +import torch + +DT = torch.float32 +from PIL import Image + +MEAN = np.array([0.48145466, 0.4578275, 0.40821073], np.float32) +STD = np.array([0.26862954, 0.26130258, 0.27577711], np.float32) +PATCH, MERGE, TEMPORAL = 14, 2, 2 + + +def smart_resize(num_frames, height, width, temporal_factor, factor, min_pixels, max_pixels): + ppt = temporal_factor * factor ** 2 + min_pixels *= ppt + max_pixels *= ppt + + def align(v, f): + return math.ceil(v / f) * f + + aligned_frames = max(temporal_factor, round(num_frames / temporal_factor) * temporal_factor) + + def fit(): + low, high, best = 1, height, (factor, factor) + while low <= high: + ch = (low + high) // 2 + cw = max(1, math.floor(width * ch / height)) + cand = (align(ch, factor), align(cw, factor)) + if aligned_frames * cand[0] * cand[1] <= max_pixels: + best, low = cand, ch + 1 + else: + high = ch - 1 + return best + + ah, aw = align(height, factor), align(width, factor) + pixels = aligned_frames * ah * aw + if pixels < min_pixels: + scale = math.sqrt(min_pixels / (num_frames * height * width)) + ah, aw = align(max(1, math.ceil(height * scale)), factor), align(max(1, math.ceil(width * scale)), factor) + pixels = aligned_frames * ah * aw + if pixels > max_pixels: + ah, aw = fit() + return ah, aw + + +def geometry(num_frames, height, width, min_tokens, max_tokens): + factor = PATCH * MERGE + ppt = TEMPORAL * factor ** 2 + th, tw = smart_resize(num_frames, height, width, TEMPORAL, factor, min_tokens, max_tokens) + scale = min(th / height, tw / width) + if num_frames * height * width >= min_tokens * ppt: + scale = min(1.0, scale) + return th, tw, max(1, min(th, math.floor(height * scale))), max(1, min(tw, math.floor(width * scale))) + + +def frames_to_patches(frames, min_tokens, max_tokens): + """frames: list of RGB PIL images (1 for an image) -> (patches [n, 1176] fp32, grid [t, h, w]).""" + n = len(frames) + w0, h0 = frames[0].size + th, tw, ch, cw = geometry(TEMPORAL if n == 1 else n, h0, w0, min_tokens, max_tokens) + arr = [] + for f in frames: + f = f if f.size == (cw, ch) else f.resize((cw, ch), Image.BICUBIC) + a = np.asarray(f, np.uint8).transpose(2, 0, 1) + a = np.pad(a, ((0, 0), (0, th - ch), (0, tw - cw))).astype(np.float32) / np.float32(255.0) + arr.append((a - MEAN[:, None, None]) / STD[:, None, None]) + gh, gw = th // PATCH, tw // PATCH + if n == 1: # an image: its one frame twice (the temporal patch) + a = arr[0] + p = a.reshape(3, gh // MERGE, MERGE, PATCH, gw // MERGE, MERGE, PATCH).transpose(1, 4, 2, 5, 0, 3, 6) + p = np.broadcast_to(p[:, :, :, :, :, None], (*p.shape[:5], TEMPORAL, *p.shape[5:])) + return np.ascontiguousarray(p.reshape(gh * gw, 3 * TEMPORAL * PATCH * PATCH)), [1, gh, gw] + f = np.stack(arr) + if n % TEMPORAL: + f = np.concatenate([f, f[-1:]]) + gt = f.shape[0] // TEMPORAL + p = f.reshape(gt, TEMPORAL, 3, gh // MERGE, MERGE, PATCH, gw // MERGE, MERGE, PATCH).transpose(0, 3, 6, 4, 7, 2, 1, 5, 8) + return np.ascontiguousarray(p.reshape(gt * gh * gw, 3 * TEMPORAL * PATCH * PATCH)), [gt, gh, gw] + + +def load_visual(model_dir): + idx = json.loads((model_dir / "model.safetensors.index.json").read_text())["weight_map"] + files = sorted({f for k, f in idx.items() if k.startswith("model.visual.")}) + out = {} + for name in files: + raw = (model_dir / name).read_bytes() + n = struct.unpack(" {feats.shape[0]} tokens, |features| rms {feats.float().pow(2).mean().sqrt():.4f}", flush=True) + + +if __name__ == "__main__": + main() From edc22852727688cbdc6567852069a181cf2d4184 Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 16:25:24 +0000 Subject: [PATCH 30/31] glm cuda: the bf16 linears on glm_mm (ours, from the glm_cuda family): weights tiled in place at load, decode windows in K slices with dependent launches, prompts in one launch of 128-row blocks sharing weight tiles in L2 (TF_GLM_MM=0 keeps the glue's b16); TF_GLM_PROF=1 times a verify's sections; tf-cuda-test glm-mm-bench (EXL3 on two GB10s: a 2-row round 67.4 to 63.3 ms, one stream 4.5-9% faster) --- zig/build/cuda.zig | 1 + zig/kernels/cuda/glm_mm.cu | 295 ++++++++++++++++++++++++++ zig/kernels/cuda/glm_qmm_frag.cuh | 117 ++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/glm/cuda.zig | 2 + zig/src/families/glm/cuda_b16.zig | 2 + zig/src/families/glm/cuda_engine.zig | 25 ++- zig/src/families/glm/cuda_forward.zig | 21 +- zig/src/families/glm/cuda_glue.zig | 3 +- zig/src/families/glm/cuda_mm.zig | 109 ++++++++++ zig/src/families/glm/cuda_moe.zig | 5 +- zig/src/families/glm/cuda_prof.zig | 76 +++++++ zig/src/families/glm/cuda_weights.zig | 61 ++++++ zig/tests/cuda/glm_mm_bench.zig | 86 ++++++++ zig/tests/cuda/main.zig | 3 + 15 files changed, 800 insertions(+), 7 deletions(-) create mode 100644 zig/kernels/cuda/glm_mm.cu create mode 100644 zig/kernels/cuda/glm_qmm_frag.cuh create mode 100644 zig/src/families/glm/cuda_mm.zig create mode 100644 zig/src/families/glm/cuda_prof.zig create mode 100644 zig/tests/cuda/glm_mm_bench.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index f0e191a744..3530629e25 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -54,6 +54,7 @@ const kernels = [_]Kernel{ .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family .{ .name = "glm_exl3", .flags = &.{"-O3"} }, // ours: GLM's EXL3 routed experts (Mia-AiLab's checkpoint), decode and prompt .{ .name = "glm_vision", .flags = &.{"-O3"} }, // ours: GLM's image tower steps around cuBLASLt (norms, rotary, softmax, merger) + .{ .name = "glm_mm", .flags = &.{"-O3"} }, // ours: GLM's bf16 linears on tensor cores, weights tiled at load, row-stable .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, diff --git a/zig/kernels/cuda/glm_mm.cu b/zig/kernels/cuda/glm_mm.cu new file mode 100644 index 0000000000..2fc8238929 --- /dev/null +++ b/zig/kernels/cuda/glm_mm.cu @@ -0,0 +1,295 @@ +// GLM dense linears, ours: one tensor-core kernel for any row count, 64-groups then K slices in order; row-stable bits. + +#include +#include +#include + +#include "glm_qmm_frag.cuh" + +namespace { + +using namespace qmm_frag; + +// 16 codes from a packed chunk (ops.cu's layout), stored + 2^(BITS-1) +template +__device__ __forceinline__ void codes16(const uint32_t* __restrict__ p, int (&q)[16]) { + if constexpr (BITS == 4) { +#pragma unroll + for (int i = 0; i < 16; ++i) q[i] = (int)((p[i / 8] >> (4 * (i % 8))) & 15u) - 8; + } else if constexpr (BITS == 8) { +#pragma unroll + for (int i = 0; i < 16; ++i) q[i] = (int)((p[i / 4] >> (8 * (i % 4))) & 255u) - 128; + } else { +#pragma unroll + for (int i = 0; i < 15; ++i) q[i] = (int)((p[i / 5] >> (6 * (i % 5))) & 63u) - 32; + q[15] = (int)((p[0] >> 30) | ((p[1] >> 30) << 2) | ((p[2] >> 30) << 4)) - 32; + } +} + +template +__device__ __forceinline__ void pack16(const int (&q)[16], uint32_t* __restrict__ out) { + constexpr int half = 1 << (BITS - 1); + uint32_t w[BITS / 2]; +#pragma unroll + for (int i = 0; i < BITS / 2; ++i) w[i] = 0u; + if constexpr (BITS == 4) { +#pragma unroll + for (int i = 0; i < 16; ++i) w[i / 8] |= (uint32_t)(q[i] + half) << (4 * (i % 8)); + } else if constexpr (BITS == 8) { +#pragma unroll + for (int i = 0; i < 16; ++i) w[i / 4] |= (uint32_t)(q[i] + half) << (8 * (i % 4)); + } else { +#pragma unroll + for (int i = 0; i < 15; ++i) w[i / 5] |= (uint32_t)(q[i] + half) << (6 * (i % 5)); + const uint32_t last = (uint32_t)(q[15] + half); +#pragma unroll + for (int j = 0; j < 3; ++j) w[j] |= ((last >> (2 * j)) & 3u) << 30; + } +#pragma unroll + for (int i = 0; i < BITS / 2; ++i) out[i] = w[i]; +} + +// input of value v (0..15) of lane ``lane``'s chunk within its group of 64 +__device__ __forceinline__ int frag_k(int lane, int v) { + const int t = lane & 3, kt = v >> 2, j = v & 3; + return kt * 16 + 2 * t + (j & 1) + (j >> 1) * 8; +} + +// the B-fragment bf16 pairs of a chunk: pairs 2 kt and 2 kt + 1 are step kt's b0 and b1 +template +__device__ __forceinline__ void bpairs(const uint32_t* __restrict__ p, uint32_t (&b)[8]) { + if constexpr (BITS == 16) { +#pragma unroll + for (int i = 0; i < 8; ++i) b[i] = p[i]; + } else { + int q[16]; + codes16(p, q); +#pragma unroll + for (int i = 0; i < 8; ++i) { + const __nv_bfloat162 h = __floats2bfloat162_rn((float)q[2 * i], (float)q[2 * i + 1]); // exact integers + b[i] = *reinterpret_cast(&h); + } + } +} + +template +struct Tile { + static constexpr int WN = 4, BN = 64, NT = 2, MT = BM / WM / 16, THREADS = WM * WN * 32; + static constexpr int WPC = BITS / 2; + static constexpr int ROW = 128, CHUNKS = 8; // a row's 64 inputs: 128 bytes, 8 chunks of 16 + static constexpr int X = BM * ROW; + static constexpr int W = 8 * 32 * WPC * 4; // 8 column tiles' chunks + static constexpr int S = 8 * 8 * 2; // their scales (fp16) + static constexpr int STAGE = X + W + S; + static constexpr int SMEM = STAGES * STAGE; +}; + +// ROWS_FIRST: row blocks fastest, so the blocks reading one weight tile run together and share it in L2 (prompts) +template +__device__ __forceinline__ void mm_kernel(const __nv_bfloat16* __restrict__ x, long long ldx, + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, + int f32, int M, int N, int K, int SK) { + using T = Tile; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / T::WN, wn = warp % T::WN; + const int KG = K >> 6, per = KG / SK, N8 = N >> 3; + const int n0 = (ROWS_FIRST ? blockIdx.y : blockIdx.x) * T::BN, m0 = (ROWS_FIRST ? blockIdx.x : blockIdx.y) * BM; + const int slice = blockIdx.z, g0 = slice * per; + const int t0 = n0 >> 3; // the block's first column tile + auto stage = [&](int st) { return buf + st * T::STAGE; }; + auto loadX = [&](int st, int g) { + unsigned char* p = stage(st); + 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 + (long long)row * ldx + g * 64 + ch * 8, m0 + r < M); + } + }; + auto loadW = [&](int st, int g) { + unsigned char* pw = stage(st) + T::X; + constexpr int PER_TILE = 32 * T::WPC * 4 / 16; // 16-byte pieces of one column tile's chunks + for (int c = tid; c < 8 * PER_TILE; c += T::THREADS) { + const int tt = c / PER_TILE, off = c % PER_TILE; + if (t0 + tt >= N8) continue; + const long long src = ((long long)(t0 + tt) * KG + g) * 32 * T::WPC * 4 + off * 16; + cp16(pw + c * 16, reinterpret_cast(w) + src); + } + if constexpr (BITS != 16) { + unsigned char* ps = pw + T::W; + for (int tt = tid; tt < 8; tt += T::THREADS) { + if (t0 + tt >= N8) continue; + cp16(ps + tt * 16, s + ((long long)(t0 + tt) * KG + g) * 8); + } + } + }; + 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; + // dependent launch: first stages' weights stream in early; inputs and outputs only after griddepcontrol.wait + grid_launch(); +#pragma unroll + for (int st = 0; st < STAGES - 1; ++st) + if (st < per) loadW(st, g0 + st); + grid_wait(); +#pragma unroll + for (int st = 0; st < STAGES - 1; ++st) { + if (st < per) loadX(st, g0 + st); + commit(); + } + for (int it = 0; it < per; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) { + loadX(next % STAGES, g0 + next); + loadW(next % STAGES, g0 + next); + } + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __half* ps = reinterpret_cast(p + T::X + T::W); + uint32_t b[T::NT][8]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) bpairs(pw + ((wn * T::NT + j) * 32 + lane) * T::WPC, b[j]); + float d[T::MT][T::NT][4]; +#pragma unroll + for (int kt = 0; kt < 4; ++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) +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[i][j], a[i], b[j][0], b[j][1]); + else mma(d[i][j], a[i], b[j][2 * kt], b[j][2 * kt + 1]); + } + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + float s0 = 1.0f, s1 = 1.0f; + if constexpr (BITS != 16) { + const __half2 h = *reinterpret_cast(ps + (wn * T::NT + j) * 8 + (lane & 3) * 2); + s0 = __low2float(h); + s1 = __high2float(h); + } +#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(d[i][j][e], (e & 1) ? s1 : s0, acc[i][j][e]); + } + } + wait<0>(); +#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 * T::NT + j) * 8 + (lane & 3) * 2; + if (col >= N) continue; +#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) { + float* dst = part + ((long long)slice * M + row) * N + col; + dst[0] = v0; + dst[1] = v1; + } else if (f32) { + float* dst = reinterpret_cast(out) + (long long)row * N + col; + dst[0] = v0; + dst[1] = v1; + } else { + *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<__nv_bfloat16*>(out) + (long long)row * N + col) = + __floats2bfloat162_rn(v0, v1); + } + } + } +} + +} // namespace + +// The K slices' partials added in slice order, then rounded once. +extern "C" __global__ void glm_mm_reduce(const float* __restrict__ part, void* __restrict__ out, long long total, int SK, int f32) { + qmm_frag::grid_launch(); + const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + float acc = part[i]; + for (int sl = 1; sl < SK; ++sl) acc = acc + part[sl * total + i]; + if (f32) reinterpret_cast(out)[i] = acc; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); +} + +// one instance a (bits, row bucket): (16, 1 warp row), (32, 1), (64, 2), (128, 2); three stages +#define GLM_MM_ALL(B) \ + extern "C" __global__ void __launch_bounds__(128) glm_mm##B##_16(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(128) glm_mm##B##_32(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(256) glm_mm##B##_64(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } \ + extern "C" __global__ void __launch_bounds__(256) glm_mm##B##_128(const __nv_bfloat16* __restrict__ x, long long ldx, \ + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, \ + int f32, int M, int N, int K, int SK) { mm_kernel(x, ldx, w, s, out, part, f32, M, N, K, SK); } +// prompts: 128-row blocks, row blocks fastest (one slice: a prompt's rows fill the GPU) +extern "C" __global__ void __launch_bounds__(256) glm_mm16_rows(const __nv_bfloat16* __restrict__ x, long long ldx, + const uint32_t* __restrict__ w, const __half* __restrict__ s, void* __restrict__ out, float* __restrict__ part, + int f32, int M, int N, int K, int SK) { mm_kernel<16, 128, 2, 3, true>(x, ldx, w, s, out, part, f32, M, N, K, SK); } +GLM_MM_ALL(4) +GLM_MM_ALL(6) +GLM_MM_ALL(8) +GLM_MM_ALL(16) + +// -- tiling at load: ops.cu's packed rows and fp16 scales [N][K / 64] -> the tiled chunks and scales --------------- +template +__device__ __forceinline__ void tile_packed(const uint32_t* __restrict__ src, const __half* __restrict__ ss, + uint32_t* __restrict__ dst, __half* __restrict__ ds, int N, int K) { + constexpr int WPC = BITS / 2; + const int KG = K >> 6, N8 = N >> 3; + const long long c = (long long)blockIdx.x * blockDim.x + threadIdx.x; // a destination chunk: (tile, group, lane) + if (c >= (long long)N8 * KG * 32) return; + const int lane = (int)(c & 31); + const long long tg = c >> 5; + const int g = (int)(tg % KG), tt = (int)(tg / KG); + const int n = tt * 8 + (lane >> 2); + const uint32_t* row = src + (long long)n * (K >> 4) * WPC; + int q[16]; +#pragma unroll + for (int v = 0; v < 16; ++v) { + const int k = g * 64 + frag_k(lane, v); + int sq[16]; + codes16(row + (long long)(k >> 4) * WPC, sq); + q[v] = sq[k & 15]; + } + pack16(q, dst + c * WPC); + if (lane < 8) ds[(tg * 8) + lane] = ss[(long long)(tt * 8 + lane) * KG + g]; +} + +extern "C" __global__ void glm_mm_tile4(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<4>(src, ss, dst, ds, N, K); } +extern "C" __global__ void glm_mm_tile6(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<6>(src, ss, dst, ds, N, K); } +extern "C" __global__ void glm_mm_tile8(const uint32_t* src, const __half* ss, uint32_t* dst, __half* ds, int N, int K) { tile_packed<8>(src, ss, dst, ds, N, K); } + +// bf16 rows [N][K] -> the tiled chunks (16 bf16 a chunk, in pair order) +extern "C" __global__ void glm_mm_tile16(const __nv_bfloat16* __restrict__ src, __nv_bfloat16* __restrict__ dst, int N, int K) { + const int KG = K >> 6, N8 = N >> 3; + const long long c = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (c >= (long long)N8 * KG * 32) return; + const int lane = (int)(c & 31); + const long long tg = c >> 5; + const int g = (int)(tg % KG), tt = (int)(tg / KG); + const int n = tt * 8 + (lane >> 2); +#pragma unroll + for (int v = 0; v < 16; ++v) dst[c * 16 + v] = src[(long long)n * K + g * 64 + frag_k(lane, v)]; +} diff --git a/zig/kernels/cuda/glm_qmm_frag.cuh b/zig/kernels/cuda/glm_qmm_frag.cuh new file mode 100644 index 0000000000..77a424cffb --- /dev/null +++ b/zig/kernels/cuda/glm_qmm_frag.cuh @@ -0,0 +1,117 @@ +// Device code of cuda/kernels/qmm_frag.cuh (flashnext/gen_copies.py): cp.async, ldmatrix, bf16 mma, nibble decoding. +#pragma once + +#include +#include +#include + +namespace qmm_frag { + +// Lane matmul tile shapes: BM rows by BN columns a block, WM x WN warps, each warp (BM / WM) x (BN / WN). +template +struct LaneTile { + static constexpr int THREADS = WM * WN * 32; + static constexpr int MT = BM / WM / 16; // m16 tiles a warp + static constexpr int NT = BN / WN / 8; // n8 tiles a warp + static constexpr int ROW = GS * 2; // bytes of one input row a group + static constexpr int CHUNKS = ROW / 16; + static constexpr int X = BM * ROW; // stage bytes: inputs, + static constexpr int W = BN * GS / 2; // weights, + static constexpr int S = BN * 2; // scales, biases (bf16), + static constexpr int XS = BM * 4; // and input sums (fp32) + static constexpr int STAGE = X + W + 2 * S + XS; + static constexpr int PARTIALS = MT * NT * 4 * THREADS * 4; // a K slice's partial, parked for the cluster sum + static constexpr int SMEM = STAGES * STAGE > PARTIALS ? STAGES * STAGE : PARTIALS; +}; + +__device__ __forceinline__ uint32_t smem(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void cp16(void* dst, const void* src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(smem(dst)), "l"(src)); +} + +// 16 bytes, or zeros without reading when ``read`` is false (a row past M: computed, never stored) +__device__ __forceinline__ void cp16z(void* dst, const void* src, bool read) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(smem(dst)), "l"(src), "r"(read ? 16 : 0)); +} + +__device__ __forceinline__ void cp8(void* dst, const void* src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" ::"r"(smem(dst)), "l"(src)); +} + +__device__ __forceinline__ void cp4(void* dst, const void* src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" ::"r"(smem(dst)), "l"(src)); +} + +__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); } + +template +__device__ __forceinline__ void wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ void ldmatrix4(uint32_t (&r)[4], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(smem(p))); +} + +__device__ __forceinline__ void ldmatrix2(uint32_t (&r)[2], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(smem(p))); +} + +__device__ __forceinline__ void mma(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + asm( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +// The same mma from a zero accumulator (a group's first k16 step): no registers to clear. +__device__ __forceinline__ void mma0(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + const float z = 0.0f; + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%10, %10, %10, %10};\n" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1), "f"(z)); +} + +// Nibbles at bits [s, s + 4) and [16 + s, 20 + s) as a bf16 pair: exponent bits make 128 + q, then minus 128. +__device__ __forceinline__ uint32_t pair(uint32_t w, int s) { + const uint32_t t = ((w >> s) & 0x000F000Fu) | 0x43004300u; + uint32_t r; +#if __CUDA_ARCH__ >= 900 + asm("sub.rn.bf16x2 %0, %1, %2;\n" : "=r"(r) : "r"(t), "r"(0x43004300u)); +#else + // no sub.bf16x2 before sm_90: t * 1 - 128 rounds the exact q once too (an exact value, the same bits) + asm("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(r) : "r"(t), "r"(0x3F803F80u), "r"(0xC300C300u)); +#endif + return r; +} + +// Programmatic dependent launch (sm_90+, else no-ops): wait for the previous kernel, release the next one's blocks. +__device__ __forceinline__ void grid_wait() { +#if __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;\n" ::: "memory"); +#endif +} + +__device__ __forceinline__ void grid_launch() { +#if __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;\n" ::: "memory"); +#endif +} + +// Block b's (first row, first column), row tiles fastest in bands of ``group`` so blocks in flight share L2. +__device__ __forceinline__ int2 tile_of(int b, int M, int N, int BM, int BN, int group) { + const int rows_t = (M + BM - 1) / BM, cols_t = (N + BN - 1) / BN, band = group * cols_t; + const int first = b / band * group, in_band = b % band, height = min(group, rows_t - first); + return make_int2((first + in_band % height) * BM, in_band / height * BN); +} + +// A 16-byte chunk c of row r sits at chunk c ^ (r % CHUNKS) so ldmatrix's eight rows hit eight bank groups. +template +__device__ __forceinline__ int swz(int r, int c) { return c ^ (r % CHUNKS); } + +} // namespace qmm_frag diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 1f97980c52..096ea30aa5 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -41,6 +41,7 @@ pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes els pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_exl3: []const u8 = if (available) &Blob("fatbin_glm_exl3").bytes else &.{}; pub const glm_vision: []const u8 = if (available) &Blob("fatbin_glm_vision").bytes else &.{}; +pub const glm_mm: []const u8 = if (available) &Blob("fatbin_glm_mm").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index 0655caefbb..0c77e40d21 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -3,6 +3,7 @@ pub const experts = @import("cuda_experts.zig"); pub const exl3 = @import("cuda_exl3.zig"); pub const vision = @import("cuda_vision.zig"); +pub const mm = @import("cuda_mm.zig"); pub const vision_config = @import("vision_config.zig"); pub const image_core = @import("image_core.zig"); /// An image's bytes to the tower's patches, as the Metal family prepares them (core's decoder). @@ -33,6 +34,7 @@ test { _ = experts; _ = exl3; _ = vision; + _ = mm; _ = image_core; _ = glue; _ = kda; diff --git a/zig/src/families/glm/cuda_b16.zig b/zig/src/families/glm/cuda_b16.zig index d935b26a64..0509b459d2 100644 --- a/zig/src/families/glm/cuda_b16.zig +++ b/zig/src/families/glm/cuda_b16.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const glue = @import("cuda_glue.zig"); +const mm = @import("cuda_mm.zig"); /// qmm.SPLIT_TARGET, B16_SPLIT_COLS, B16_BK, geometry.MLA_B16_SLICES and MLA_B16_ROWS_FROM. pub const split_target = 192; @@ -55,6 +56,7 @@ pub fn partBytes(m: usize, n: usize, k: usize) usize { /// qmm.matmul for a B16 weight: out [m][n] bf16 (fp32 with `wide`) = x . w; 128+ rows sum in registers, else `part`. pub fn matmul(set: glue.Set, stream: cuda.Stream, x: u64, x_stride: usize, w: u64, out: u64, part: u64, m: usize, n: usize, k: usize, wide: bool) !void { if (k % bk != 0) return error.KNotMultipleOf64; + if (set.mm) |mk| return mm.matmul(mk, stream, x, x_stride, .{ .w = w, .n = n, .k = k }, out, part, m, wide); // tiled weights const sk = splitK(n, k); var a: cuda.Args = .{}; a.add(x); diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 908ed93333..51b60700f1 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -13,6 +13,7 @@ const heat = @import("heat"); const modelopt = @import("core").modelopt; const exl3 = @import("cuda_exl3.zig"); const images_mod = @import("cuda_images.zig"); +const prof_mod = @import("cuda_prof.zig"); pub const Options = struct { capacity: usize, // tokens a sequence holds @@ -70,6 +71,7 @@ pub const Engine = struct { comm: cuda.nccl.Comm, group: ?*const cuda.tp_group.Group = null, // the main stream's all-gathers (RoCE on two ranks); null: lib and comm images: ?*images_mod.Images = null, // --vision: a prompt pass's image rows (rank 0 encodes, the others receive) + profs: [fw.window_rows]?*prof_mod.Prof = @splat(null), // TF_GLM_PROF=1: verify sections' GPU time, by rows settings: ranks.Settings, m: fw.Modules, stream: cuda.Stream, @@ -140,6 +142,10 @@ pub const Engine = struct { for (layers, 0..) |*l, i| l.* = @intCast(i); e.* = .{ .gpa = gpa, .io = io, .heat_gate = .{ .bands = o.settings.heat, .root = heatRoot() }, .shared_buf = undefined, .d = d, .c = c, .share = share, .spans = spans, .lib = lib, .comm = comm, .settings = o.settings, .m = m, .stream = stream, .pieces = pieces, .w = undefined, .capacity = o.capacity, .bp = undefined, .bd = undefined, .bm = undefined, .hold = undefined, .pairs = undefined, .smp = undefined, .mtp = o.mtp }; e.w = try weights.load(gpa, io, d, stream, e.m.tools(), checkpoint, c, share, layers, o.mtp); + if (tiledLinears()) { // our bf16 linears (glm_mm) on their tiled weights; TF_GLM_MM=0 keeps the glue's + try weights.tileLinears(&e.w, &e.m.mm, stream, d); + e.m.tiled = true; + } errdefer e.w.deinit(); try e.bind(0); errdefer { @@ -160,6 +166,9 @@ pub const Engine = struct { errdefer e.pairs.free(); try d.check(d.api.cuMemsetD8_v2(e.pairs.ptr, 0, pair_bytes * world), "cuMemsetD8"); // each pair's padding stays zero e.smp = try sampler.Sampler.init(d, c.vocab, fw.window_rows); + if (prof_mod.Prof.wanted()) for (&e.profs) |*p| { + p.* = try prof_mod.Prof.init(d, io); + }; errdefer e.smp.deinit(); e.shared_buf = try cuda.DeviceBuffer.alloc(d, 8); try e.stream.synchronize(); @@ -411,11 +420,19 @@ pub const Engine = struct { pub fn verify(e: *Engine, ids: []const u32, out: []u32) !void { if (ids.len == 0 or ids.len > fw.window_rows) return error.WindowTooWide; var one: [1]fw.Rank = undefined; - const f = e.view(&e.bd, &one); + var f = e.view(&e.bd, &one); const at = e.cur.st.pos; + f.prof = e.profs[ids.len - 1]; + if (f.prof) |p| try p.begin(e.stream, verifyLabel(ids.len)); try f.run(@ptrCast(ids)); e.cur.rows_at = e.bd.fnormed; try e.draw(e.bd.logits, ids.len, at + 1, out); + if (f.prof) |p| p.end(); + } + + fn verifyLabel(rows: usize) []const u8 { + const l = [_][]const u8{ "verify 1 row", "verify 2 rows", "verify 3 rows", "verify 4 rows", "verify 5 rows", "verify 6 rows", "verify 7 rows", "verify 8 rows" }; + return l[@min(rows, l.len) - 1]; } /// Several streams' windows in one forward: shared row ops, per-part attention and KDA. @@ -515,6 +532,12 @@ pub const Engine = struct { if (e.group) |g| try g.check(); } + /// TF_GLM_MM=0 keeps the bf16 linears on the glue's b16 kernels (the weights stay as the checkpoint stores them). + fn tiledLinears() bool { + const v = std.c.getenv("TF_GLM_MM") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); + } + fn heatRoot() []const u8 { const p = std.c.getenv("TF_HEAT_ROOT") orelse return heat.default_root; const s = std.mem.span(p); diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index 56e4de8743..bbc5a92942 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/zig/src/families/glm/cuda_forward.zig @@ -17,6 +17,8 @@ const moe = @import("cuda_moe.zig"); const quant = @import("cuda_quant.zig"); const q4 = @import("cuda_q4.zig"); const state = @import("cuda_state.zig"); +const mm_mod = @import("cuda_mm.zig"); +const prof_mod = @import("cuda_prof.zig"); const images_mod = @import("cuda_images.zig"); const exl3 = @import("cuda_exl3.zig"); const shared = @import("cuda_shared.zig"); @@ -29,12 +31,14 @@ pub const Buffers = state.Buffers; /// The modules a forward launches from: our fatbins. pub const Modules = struct { glue: glue.Kernels, - mods: [8]cuda.Module, + mods: [9]cuda.Module, kda: kda.Kernels, lane: cuda.qmmf.Lane, // the NVFP4 linears' decode rows (qlinear) prompt: cuda.nvfp4.Prompt, // and their prompt rows experts: experts.Kernels, exl3: exl3.Kernels, // an EXL3 checkpoint's routed experts + mm: mm_mod.Kernels, // glm_mm: the bf16 linears on tiled weights + tiled: bool = false, // the bf16 linears' weights are tiled (Engine.init, TF_GLM_MM unset or 1): glm_mm runs them copy_rows: cuda.Function, put_slot: cuda.Function, pack_dense: cuda.Function, @@ -46,7 +50,7 @@ pub const Modules = struct { var m: Modules = undefined; var n: usize = 0; errdefer for (m.mods[0..n]) |*x| x.unload(); - for (&m.mods, [_][]const u8{ k.glm_kda, k.glm_ops, k.experts, k.affine4_pack, k.glm_quant, k.qmm_group, k.glm_glue, k.glm_exl3 }) |*x, image| { + for (&m.mods, [_][]const u8{ k.glm_kda, k.glm_ops, k.experts, k.affine4_pack, k.glm_quant, k.qmm_group, k.glm_glue, k.glm_exl3, k.glm_mm }) |*x, image| { x.* = try cuda.Module.load(d, image); n += 1; } @@ -64,6 +68,8 @@ pub const Modules = struct { m.q4 = try q4.Kernels.load(m.mods[5], m.mods[3], ctx); m.glue = try glue.Kernels.load(m.mods[6]); m.exl3 = try exl3.Kernels.load(m.mods[7]); + m.mm = try mm_mod.Kernels.load(m.mods[8]); + m.tiled = false; return m; } @@ -167,9 +173,10 @@ pub const Forward = struct { pieces: ?*const Pieces = null, // a prompt chunk's exchanges in row pieces on a second stream (null: whole) split: bool = true, // TF_GLM_PREFILL_REDUCE=split (the default with send/recv): each rank glues its own rows images: ?ImageRows = null, // a prompt pass's image rows over their placeholders' embeddings (its sequence only) + prof: ?*prof_mod.Prof = null, // TF_GLM_PROF=1: events between the sections pub fn set(f: *const Forward) glue.Set { - return .{ .glue = &f.m.glue }; + return .{ .glue = &f.m.glue, .mm = if (f.m.tiled) &f.m.mm else null }; } /// reduce.sums' gather: every rank's partial of `rows` rows into every rank's gathered buffer, rank order. @@ -258,7 +265,7 @@ pub const Forward = struct { } fn moeKernels(f: *const Forward) moe.Kernels { - return .{ .set = f.set(), .experts = &f.m.experts, .exl3 = &f.m.exl3, .put_slot = f.m.put_slot }; + return .{ .set = f.set(), .experts = &f.m.experts, .exl3 = &f.m.exl3, .put_slot = f.m.put_slot, .prof = f.prof }; } const Site = enum { mixer, ffn }; @@ -462,11 +469,13 @@ pub const Forward = struct { const c = &f.ranks[0].w.config; for (f.ranks) |r| try glue.embedB16(f.set(), f.stream, r.b.ids, r.w.embed, r.b.x, R, c.hidden, c.hc); if (f.images) |im| for (f.ranks) |r| try im.v.splice(f.stream, im.slot, r.st.pos, r.st.pos + R, r.b.x, c.hc); + if (f.prof) |p| p.mark(f.stream, .embed); var cut: [max_pieces][2]usize = undefined; const pieces = if (f.pieces) |p| (if (f.ranks[0].b.prefill) ranges(R, p.count, &cut) else 0) else 0; // overlap.Overlap.begin: pieced with two pieces or more, and split takes a chunk of one piece too if (pieces > 1 or (pieces == 1 and f.splits(R))) try f.pieced(f.pieces.?, R, cut[0..pieces]) else try f.whole(R); try f.finish(R); + if (f.prof) |p| p.mark(f.stream, .head); } /// The unpieced loop over the layers (and the stream mean): each site's partials whole, then the glue. @@ -480,15 +489,19 @@ pub const Forward = struct { try f.pre(r, l.attn_hc, l.in_norm, R); if (r.segs.len > 0) try shared.mixer(f, r, l, li, di, R, slot(r, R)) else try f.mixer(r, l, li, di, R, @min(latent.chunksFor(r.st.pos + R), r.b.nch), slot(r, R)); } + if (f.prof) |p| p.mark(f.stream, if (f.ranks[0].w.layers[j].mixer == .kda) .kda else .mla); try f.gather(R); + if (f.prof) |p| p.mark(f.stream, .gather_a); for (f.ranks) |r| { const l = r.w.layers[j]; try f.post(r, R); try f.pre(r, l.ffn_hc, l.post_norm, R); try f.ffn(r, l, R, slot(r, R)); } + if (f.prof) |p| p.mark(f.stream, if (f.ranks[0].w.layers[j].ffn == .dense) .dense else .shared); // MoE: its fill try f.gather(R); for (f.ranks) |r| try f.post(r, R); + if (f.prof) |p| p.mark(f.stream, .gather_b); switch (f.ranks[0].w.layers[j].mixer) { .kda => li += 1, .mla => di += 1, diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index 733edcdb27..53414396d3 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -2,9 +2,10 @@ const std = @import("std"); const cuda = @import("cuda"); +const mm_mod = @import("cuda_mm.zig"); /// What the forward's kernels launch from: our glue's functions. -pub const Set = struct { glue: *const Kernels }; +pub const Set = struct { glue: *const Kernels, mm: ?*const mm_mod.Kernels = null }; // mm: bf16 linears tiled at load /// Our glm_glue.cu's functions. pub const Kernels = struct { diff --git a/zig/src/families/glm/cuda_mm.zig b/zig/src/families/glm/cuda_mm.zig new file mode 100644 index 0000000000..0abdfd7e39 --- /dev/null +++ b/zig/src/families/glm/cuda_mm.zig @@ -0,0 +1,109 @@ +//! GLM's bf16 linears on tensor cores (glm_mm): weights tiled once at load, any row count, row-stable K slices. +const std = @import("std"); +const cuda = @import("cuda"); + +/// A linear's tiled weights: [n / 8][k / 64][32 lanes][8 words] bf16 pairs (glm_mm_tile16's layout). +pub const Tiled = struct { w: u64, n: usize, k: usize }; + +/// glm_mm's bf16 kernels by row bucket, the K-slice reduce and the tiling. +pub const Kernels = struct { + by_rows: [4]cuda.Function, // 16, 32, 64, 128 rows a block + prompt: cuda.Function, // 128-row blocks, row blocks fastest: a prompt's blocks share each weight tile in L2 + reduce: cuda.Function, + tile: cuda.Function, + pdl: bool, // dependent launches: a matmul's weights stream in while the kernel before it finishes + + pub fn load(m: cuda.Module) !Kernels { + var k: Kernels = .{ .by_rows = undefined, .prompt = try m.function("glm_mm16_rows"), .reduce = try m.function("glm_mm_reduce"), .tile = try m.function("glm_mm_tile16"), .pdl = true }; + inline for (.{ "glm_mm16_16", "glm_mm16_32", "glm_mm16_64", "glm_mm16_128" }, 0..) |name, i| { + k.by_rows[i] = try m.function(name); + try k.by_rows[i].allowDynamicShared(smem(bucket(i))); + } + try k.prompt.allowDynamicShared(smem(128)); + if (std.c.getenv("TF_GLM_PDL")) |v| k.pdl = !std.mem.eql(u8, std.mem.span(v), "0"); + return k; + } +}; + +const split_target = 192; // column tiles times K slices a launch aims at (the GB10's SMs, a few blocks each) + +fn bucket(i: usize) usize { + return @as(usize, 16) << @intCast(i); +} + +fn smem(bm: usize) u32 { + return @intCast(3 * (bm * 128 + 8 * 32 * 8 * 4 + 128)); +} + +/// K slices fixed by the shape: column tiles times slices reach the target, each slice at least 4 groups of 64. +pub fn splitK(n: usize, k: usize) usize { + const tiles = (n + 63) / 64; + const steps = k / 64; + var sk: usize = 1; + while (sk < 8 and tiles * sk < split_target and steps % (sk * 2) == 0 and steps / (sk * 2) >= 4) sk *= 2; + return sk; +} + +/// Bytes of fp32 partials a call of `m` rows needs (0: one slice; windows past 64 rows go 64 at a time). +pub fn partBytes(m: usize, w: Tiled) usize { + const sk = splitK(w.n, w.k); + return if (sk > 1) sk * @min(m, 64) * w.n * 4 else 0; +} + +/// x [m][k] (rows `ldx` apart) by `w` -> out [m][n], bf16 or fp32 (`wide`); `part`: partBytes(m, w). +pub fn matmul(k: *const Kernels, s: cuda.Stream, x: u64, ldx: usize, w: Tiled, out: u64, part: u64, m: usize, wide: bool) !void { + if (w.k % 64 != 0 or w.n % 8 != 0 or x % 16 != 0 or ldx % 8 != 0) return error.BadShape; + var done: usize = 0; + while (done < m) { // windows to 64 rows a launch with K slices; prompts in one launch of 128-row blocks, one slice + const big = m - done >= 128; + const rows = if (big) m - done else @min(m - done, 64); + const i: usize = if (big) 3 else if (rows <= 16) 0 else if (rows <= 32) 1 else 2; + const bm = bucket(i); + const sk = if (big) 1 else splitK(w.n, w.k); // a prompt's blocks share each weight tile in L2: no slices + const o = out + done * w.n * @as(usize, if (wide) 4 else 2); + var a: cuda.Args = .{}; + a.add(x + done * ldx * 2); + a.add(@as(c_longlong, @intCast(ldx))); + a.add(w.w); + a.add(@as(u64, 0)); + a.add(o); + a.add(part); + a.add(@as(c_int, @intFromBool(wide))); + a.add(@as(c_int, @intCast(rows))); + a.add(@as(c_int, @intCast(w.n))); + a.add(@as(c_int, @intCast(w.k))); + a.add(@as(c_int, @intCast(sk))); + const cols: u32 = @intCast((w.n + 63) / 64); + const row_blocks: u32 = @intCast((rows + bm - 1) / bm); + if (big) { + try cuda.launch.launch(k.prompt, .{ .grid = .{ .x = row_blocks, .y = cols }, .block = .{ .x = 256 }, .shared = smem(128), .pdl = k.pdl }, s, &a); + } else try cuda.launch.launch(k.by_rows[i], .{ .grid = .{ .x = cols, .y = row_blocks, .z = @intCast(sk) }, .block = .{ .x = if (bm <= 32) 128 else 256 }, .shared = smem(bm), .pdl = k.pdl }, s, &a); + if (sk > 1) { + const total = rows * w.n; + var r: cuda.Args = .{}; + r.add(part); + r.add(o); + r.add(@as(c_longlong, @intCast(total))); + r.add(@as(c_int, @intCast(sk))); + r.add(@as(c_int, @intFromBool(wide))); + try cuda.launch.launch(k.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256) }, .block = .{ .x = 256 } }, s, &r); + } + done += rows; + } +} + +/// A row-major bf16 [n][k] linear at `src` tiled into `dst` (as many bytes). +pub fn tile(k: *const Kernels, s: cuda.Stream, src: u64, dst: u64, n: usize, kk: usize) !void { + const chunks = n / 8 * (kk / 64) * 32; + var a: cuda.Args = .{}; + a.add(src); + a.add(dst); + a.add(@as(c_int, @intCast(n))); + a.add(@as(c_int, @intCast(kk))); + try cuda.launch.launch(k.tile, .{ .grid = .{ .x = @intCast((chunks + 255) / 256) }, .block = .{ .x = 256 } }, s, &a); +} + +test "K slices as the shapes take them" { + try std.testing.expectEqual(@as(usize, 1), splitK(12288, 4096)); // 192 column tiles already + try std.testing.expectEqual(@as(usize, 4), splitK(4096, 2048)); // 64 tiles, 32 groups: 4 slices of 8 +} diff --git a/zig/src/families/glm/cuda_moe.zig b/zig/src/families/glm/cuda_moe.zig index d50e8707c1..72132399d2 100644 --- a/zig/src/families/glm/cuda_moe.zig +++ b/zig/src/families/glm/cuda_moe.zig @@ -6,6 +6,7 @@ const glue = @import("cuda_glue.zig"); const experts = @import("cuda_experts.zig"); const b16 = @import("cuda_b16.zig"); const exl3 = @import("cuda_exl3.zig"); +const prof_mod = @import("cuda_prof.zig"); /// One MoE layer on a rank: router (bf16 [count][dims]) and fp32 bias, routed experts, shared [gate | up] and down. pub const Layer = struct { @@ -40,7 +41,7 @@ pub const Scratch = struct { }; /// The kernels the block runs: the captured Triton set, the grouped NVFP4 experts, the plan, and glm_ops. -pub const Kernels = struct { set: glue.Set, experts: *const experts.Kernels, exl3: *const exl3.Kernels, put_slot: cuda.Function }; +pub const Kernels = struct { set: glue.Set, experts: *const experts.Kernels, exl3: *const exl3.Kernels, put_slot: cuda.Function, prof: ?*prof_mod.Prof = null }; /// forward.moe_block: out [rows][dims] fp32 partial; `prompt`: 64-pair items, the prompt kernels, bf16 slot rows. pub fn block(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_stride: usize, rows: usize, prompt: bool, out: u64) !void { @@ -55,6 +56,8 @@ pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const shared_id: i32 = @intCast(l.count); try glue.router(k.set, stream, x, x_stride, l.router, s.mlog, s.router_part, rows, d, l.count); try glue.select(k.set, stream, s.mlog, l.bias, s.pick, s.wts, rows, l.top_k, l.count, l.scale, l.norm); + if (k.prof) |p| p.mark(stream, .router); + defer if (k.prof) |p| p.mark(stream, .routed); if (l.exl3) |w| return exl3.run(k.exl3, k.experts, stream, w, s.x3, s.plan, x, x_stride, s.pick, slots, rows, l.count, d, l.routed.width, l.routed.limit, prompt, s.ey, s.eact); const plan: experts.Plan = .{ .route = s.plan, .slots = slots, .experts = l.count + 1, .prompt = prompt }; try experts.route(k.experts, stream, s.pick, rows * slots, plan); diff --git a/zig/src/families/glm/cuda_prof.zig b/zig/src/families/glm/cuda_prof.zig new file mode 100644 index 0000000000..f948fa5792 --- /dev/null +++ b/zig/src/families/glm/cuda_prof.zig @@ -0,0 +1,76 @@ +//! TF_GLM_PROF=1: where a forward's GPU time goes (events between its sections) against the call's wall time. +const std = @import("std"); +const cuda = @import("cuda"); + +pub const Section = enum { embed, kda, mla, gather_a, dense, router, routed, shared, gather_b, head, mtp }; +const names = [_][]const u8{ "embed", "kda", "mla", "gather_a", "dense", "router", "routed", "shared", "gather_b", "head", "mtp" }; +const max_marks = 512; + +pub const Prof = struct { + d: *const cuda.Driver, + events: [max_marks + 1]cuda.Event, + ends: [max_marks]Section = undefined, + n: usize = 0, // marks this call + open: bool = false, + t0: i96 = 0, + io: std.Io, + sections: [names.len]f64 = @splat(0), // ms each section took on the GPU, summed over calls + gpu_ms: f64 = 0, + wall_ms: f64 = 0, + calls: u64 = 0, + label: []const u8 = "", + + pub fn init(d: *const cuda.Driver, io: std.Io) !*Prof { + const p = try std.heap.page_allocator.create(Prof); + p.* = .{ .d = d, .events = undefined, .io = io }; + for (&p.events) |*e| e.* = try cuda.Event.init(d, true); + return p; + } + + /// TF_GLM_PROF=1 asks for it. + pub fn wanted() bool { + const v = std.c.getenv("TF_GLM_PROF") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); + } + + /// A call begins on `s` (its wall clock and its first event). + pub fn begin(p: *Prof, s: cuda.Stream, label: []const u8) !void { + p.n = 0; + p.open = true; + p.label = label; + p.t0 = std.Io.Clock.awake.now(p.io).toNanoseconds(); + try p.events[0].record(s); + } + + /// The work queued on `s` since the last mark belongs to `section`. + pub fn mark(p: *Prof, s: cuda.Stream, section: Section) void { + if (!p.open or p.n >= max_marks) return; + p.events[p.n + 1].record(s) catch return; + p.ends[p.n] = section; + p.n += 1; + } + + /// The call ended (its results read: the stream is idle): sections summed, a report every 64 calls. + pub fn end(p: *Prof) void { + if (!p.open) return; + p.open = false; + const wall: f64 = @as(f64, @floatFromInt(std.Io.Clock.awake.now(p.io).toNanoseconds() - p.t0)) / 1e6; + if (p.n == 0) return; + p.events[p.n].synchronize() catch return; + for (0..p.n) |i| p.sections[@backingInt(p.ends[i])] += cuda.Event.elapsedMs(p.events[i], p.events[i + 1]) catch 0; + p.gpu_ms += cuda.Event.elapsedMs(p.events[0], p.events[p.n]) catch 0; + p.wall_ms += wall; + p.calls += 1; + if (p.calls % 64 != 0) return; + const k: f64 = @floatFromInt(p.calls); + var buf: [512]u8 = undefined; + var w: std.Io.Writer = .fixed(&buf); + w.print("glm prof {s} ({d} calls): wall {d:.2} ms, gpu {d:.2} ms a call;", .{ p.label, p.calls, p.wall_ms / k, p.gpu_ms / k }) catch {}; + for (names, p.sections) |nm, ms| if (ms > 0) w.print(" {s} {d:.2}", .{ nm, ms / k }) catch {}; + std.log.info("{s}", .{w.buffered()}); + p.sections = @splat(0); + p.gpu_ms = 0; + p.wall_ms = 0; + p.calls = 0; + } +}; diff --git a/zig/src/families/glm/cuda_weights.zig b/zig/src/families/glm/cuda_weights.zig index ed9b5b7e32..fb9fbee479 100644 --- a/zig/src/families/glm/cuda_weights.zig +++ b/zig/src/families/glm/cuda_weights.zig @@ -13,6 +13,7 @@ const moe = @import("cuda_moe.zig"); const quant = @import("cuda_quant.zig"); const q4 = @import("cuda_q4.zig"); const exl3 = @import("cuda_exl3.zig"); +const mm_mod = @import("cuda_mm.zig"); const Tensor = core.checkpoint.Tensor; const prefix = "model.language_model."; @@ -524,3 +525,63 @@ test "bf16 rounding of fp32 is to nearest even" { try std.testing.expectEqual(@as(u16, 0x3f82), toBf16(@bitCast(@as(u32, 0x3f818000)))); // a tie to even (up) try std.testing.expectEqual(@as(u16, 0x3f81), toBf16(@bitCast(@as(u32, 0x3f808001)))); } + +/// Every bf16 linear b16.matmul reads, retiled in place for glm_mm (the q4 draft head was made from the head already). +pub fn tileLinears(w: *const Weights, mk: *const mm_mod.Kernels, s: cuda.Stream, d: *const cuda.Driver) !void { + var list: std.ArrayList([3]usize) = .empty; // (pointer, n, k) + defer list.deinit(w.gpa); + const add = struct { + fn f(l: *std.ArrayList([3]usize), gpa: std.mem.Allocator, ptr: u64, n: usize, k: usize) !void { + if (ptr != 0) try l.append(gpa, .{ ptr, n, k }); + } + }.f; + const kd = @import("cuda_kda.zig").dk; + var layers: std.ArrayList(Layer) = .empty; + defer layers.deinit(w.gpa); + try layers.appendSlice(w.gpa, w.layers); + if (w.mtp) |m| { + try layers.append(w.gpa, m.layer); + try add(&list, w.gpa, m.eh, w.config.hidden, 2 * w.config.hidden); + } + for (layers.items) |l| { + switch (l.mixer) { + .kda => |x| { + const n = x.heads * kd; + try add(&list, w.gpa, x.proj, x.width(), x.dims); + try add(&list, w.gpa, x.fb, n, kd); + try add(&list, w.gpa, x.gb, n, kd); + try add(&list, w.gpa, x.o, x.dims, n); + }, + .mla => |x| { + try add(&list, w.gpa, x.proj, x.q_lora + x.lw, x.dims); + try add(&list, w.gpa, x.q_b, x.heads * x.qk, x.q_lora); + try add(&list, w.gpa, x.index.kw, x.index.d + x.index.heads, x.dims); + try add(&list, w.gpa, x.index.qb, x.index.heads * x.index.d, x.q_lora); + try add(&list, w.gpa, x.o, x.dims, x.heads * x.dv); + }, + } + switch (l.ffn) { + .dense => |x| switch (x.w) { + .bf16 => |b| { + try add(&list, w.gpa, b.gu, 2 * x.width, x.dims); + try add(&list, w.gpa, b.down, x.dims, x.width); + }, + .nvfp4 => {}, + }, + .moe => |x| { + try add(&list, w.gpa, x.shared_gu, 2 * x.shared, x.routed.dims); + try add(&list, w.gpa, x.shared_down, x.routed.dims, x.shared); + }, + } + } + try add(&list, w.gpa, w.head, w.share.vocab.len(), w.config.hidden); + var most: usize = 0; + for (list.items) |e| most = @max(most, e[1] * e[2] * 2); + var tmp = try cuda.DeviceBuffer.alloc(d, most); + defer tmp.free(); + for (list.items) |e| { + try mm_mod.tile(mk, s, e[0], tmp.ptr, e[1], e[2]); + try d.check(d.api.cuMemcpyDtoDAsync_v2(e[0], tmp.ptr, e[1] * e[2] * 2, s.handle), "cuMemcpyDtoDAsync"); + } + try s.synchronize(); +} diff --git a/zig/tests/cuda/glm_mm_bench.zig b/zig/tests/cuda/glm_mm_bench.zig new file mode 100644 index 0000000000..b8532a4d8f --- /dev/null +++ b/zig/tests/cuda/glm_mm_bench.zig @@ -0,0 +1,86 @@ +//! GLM's bf16 linears on one GPU: the family's b16.matmul against glm_mm (tiled weights) on a rank's shapes, by rows. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm_cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +/// A rank's linears of GLM-5.3-Flash at two ranks (n, k): KDA projections and output, MLA, shared expert, dense layer. +const shapes = [_][2]usize{ .{ 6416, 4096 }, .{ 4096, 2048 }, .{ 2048, 4096 }, .{ 4096, 1024 }, .{ 12288, 4096 }, .{ 4096, 6144 }, .{ 2560, 4096 }, .{ 4096, 4096 } }; +const reps = 30; + +pub fn run(gpu: Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var gm = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer gm.unload(); + var mm_mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_mm); + defer mm_mod.unload(); + const glue = try glm.glue.Kernels.load(gm); + const set: glm.glue.Set = .{ .glue = &glue }; + const mk = try glm.mm.Kernels.load(mm_mod); + var prng = std.Random.DefaultPrng.init(3); + const rng = prng.random(); + var t0 = try cuda.Event.init(gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(gpu.d, true); + defer t1.deinit(); + var total = [2][7]f64{ @splat(0), @splat(0) }; + for (shapes) |sh| { + const n = sh[0]; + const k = sh[1]; + const host_w = try a.alloc(u16, n * k); + defer a.free(host_w); + for (host_w) |*v| v.* = @truncate(@as(u32, @bitCast((rng.float(f32) * 2 - 1) * 0.02)) >> 16); + var w = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host_w)); + defer w.free(); + var wt = try cuda.DeviceBuffer.alloc(gpu.d, n * k * 2); + defer wt.free(); + try glm.mm.tile(&mk, s, w.ptr, wt.ptr, n, k); + const tw: glm.mm.Tiled = .{ .w = wt.ptr, .n = n, .k = k }; + const host_x = try a.alloc(u16, 512 * k); + defer a.free(host_x); + for (host_x) |*v| v.* = @truncate(@as(u32, @bitCast(rng.float(f32) * 2 - 1)) >> 16); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host_x)); + defer x.free(); + var out_a = try cuda.DeviceBuffer.alloc(gpu.d, 512 * n * 4); + defer out_a.free(); + var out_b = try cuda.DeviceBuffer.alloc(gpu.d, 512 * n * 4); + defer out_b.free(); + var part = try cuda.DeviceBuffer.alloc(gpu.d, @max(16, @max(glm.b16.partBytes(127, n, k), glm.mm.partBytes(128, tw)))); + defer part.free(); + var ms: [2][7]f64 = undefined; + for ([_]usize{ 1, 2, 4, 8, 64, 128, 512 }, 0..) |rows, ri| { + for (0..2) |which| { + for (0..2) |round| { + try t0.record(s); + for (0..reps) |_| { + if (which == 0) try glm.b16.matmul(set, s, x.ptr, k, w.ptr, out_a.ptr, part.ptr, rows, n, k, true) else try glm.mm.matmul(&mk, s, x.ptr, k, tw, out_b.ptr, part.ptr, rows, true); + } + try t1.record(s); + try t1.synchronize(); + if (round == 1) ms[which][ri] = @as(f64, try cuda.Event.elapsedMs(t0, t1)) / reps; + } + total[which][ri] += ms[which][ri]; + } + // the two products agree (different summation orders: a few fp32 ulps of the row's magnitude) + const ga = try a.alloc(f32, rows * n); + defer a.free(ga); + const gb = try a.alloc(f32, rows * n); + defer a.free(gb); + try out_a.download(0, std.mem.sliceAsBytes(ga)); + try out_b.download(0, std.mem.sliceAsBytes(gb)); + var most: f32 = 0; + var big: f32 = 0; + for (ga, gb) |u, v| { + most = @max(most, @abs(u - v)); + big = @max(big, @abs(u)); + } + try check.expect(most <= big * 1e-3, "glm-mm-bench: {d} x {d} at {d} rows: the products agree", .{ n, k, rows }); + } + std.debug.print("RESULT {d} x {d} (rows 1/2/4/8/64/128/512):\n b16 {d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3} ms\n glm_mm {d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3}/{d:.3} ms\n", .{ n, k, ms[0][0], ms[0][1], ms[0][2], ms[0][3], ms[0][4], ms[0][5], ms[0][6], ms[1][0], ms[1][1], ms[1][2], ms[1][3], ms[1][4], ms[1][5], ms[1][6] }); + } + std.debug.print("RESULT all shapes once (rows 1/2/4/8/64/128/512):\n b16 {d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2} ms\n glm_mm {d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2}/{d:.2} ms\n", .{ total[0][0], total[0][1], total[0][2], total[0][3], total[0][4], total[0][5], total[0][6], total[1][0], total[1][1], total[1][2], total[1][3], total[1][4], total[1][5], total[1][6] }); + check.pass("glm-mm-bench: both products agree on every shape and row count", .{}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 73a49fda08..89343d4be7 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -24,6 +24,7 @@ const quant_ref_tests = @import("quant_ref_tests.zig"); const glm_ref_tests = @import("glm_ref_tests.zig"); const glm_exl3_tests = @import("glm_exl3_tests.zig"); const glm_vision_tests = @import("glm_vision_tests.zig"); +const glm_mm_bench = @import("glm_mm_bench.zig"); const carveout_tests = @import("carveout_tests.zig"); const nvfp4_tests = @import("nvfp4_tests.zig"); const modelopt_tests = @import("modelopt_tests.zig"); @@ -96,6 +97,7 @@ const usage = \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) \\ glm-vision MODEL F32 BF16 IMG GLM's image tower and preparation against vision_ref.py's torch runs + \\ glm-mm-bench GLM's bf16 linears: b16.matmul against glm_mm, by rows (one GPU) \\ glm-ref GLM projections (matmuls, router, attention, indexer, hc) against float64 (synthetic) \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run @@ -165,6 +167,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-ref")) return glm_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-exl3")) return glm_exl3_tests.run_(gpu); + if (std.mem.eql(u8, cmd, "glm-mm-bench")) return glm_mm_bench.run(gpu); if (std.mem.eql(u8, cmd, "glm-vision")) return glm_vision_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); From 875b695e0cd015c29d4d1f9a606888e42ae44dbd Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Fri, 16 Oct 2026 16:57:55 +0000 Subject: [PATCH 31/31] glm cuda: TF_GLM_PROF=1 also times prompt chunks by section (gate, forward, draw, the MTP head's pass, commit), by size --- zig/src/families/glm/cuda_engine.zig | 18 +++++++++++++++++- zig/src/families/glm/cuda_prof.zig | 7 ++++--- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 51b60700f1..f9e92556e0 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -72,6 +72,7 @@ pub const Engine = struct { group: ?*const cuda.tp_group.Group = null, // the main stream's all-gathers (RoCE on two ranks); null: lib and comm images: ?*images_mod.Images = null, // --vision: a prompt pass's image rows (rank 0 encodes, the others receive) profs: [fw.window_rows]?*prof_mod.Prof = @splat(null), // TF_GLM_PROF=1: verify sections' GPU time, by rows + prompt_profs: [4]?*prof_mod.Prof = @splat(null), // and prompt chunks': to 64, 256, 1024, 2048 rows settings: ranks.Settings, m: fw.Modules, stream: cuda.Stream, @@ -169,6 +170,10 @@ pub const Engine = struct { if (prof_mod.Prof.wanted()) for (&e.profs) |*p| { p.* = try prof_mod.Prof.init(d, io); }; + if (prof_mod.Prof.wanted()) for (&e.prompt_profs) |*p| { + p.* = try prof_mod.Prof.init(d, io); + p.*.?.every = 4; + }; errdefer e.smp.deinit(); e.shared_buf = try cuda.DeviceBuffer.alloc(d, 8); try e.stream.synchronize(); @@ -285,10 +290,15 @@ pub const Engine = struct { const rows = @min(len - start, fw.prefill_rows); if (part.len != @min(len - start, rows + 1)) return error.BadChunk; var one: [1]fw.Rank = undefined; - const f = e.view(&e.bp, &one); + var f = e.view(&e.bp, &one); const as_i32: []const i32 = @ptrCast(part); if (marks) |m| if (start > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(start)) != null) m.call(m.ctx, @intCast(start)); + f.prof = e.prompt_profs[if (rows <= 64) 0 else if (rows <= 256) 1 else if (rows <= 1024) 2 else 3]; + const labels = [_][]const u8{ "prompt chunk to 64 rows", "prompt chunk to 256 rows", "prompt chunk to 1024 rows", "prompt chunk to 2048 rows" }; + if (f.prof) |p| try p.begin(e.stream, labels[if (rows <= 64) 0 else if (rows <= 256) 1 else if (rows <= 1024) 2 else 3]); + defer if (f.prof) |p| p.end(); try e.gate(cancel); + if (f.prof) |p| p.mark(e.stream, .gate); try f.run(as_i32[0..rows]); const last = start + rows == len; var first: ?u32 = null; @@ -297,15 +307,21 @@ pub const Engine = struct { try e.draw(e.bp.logits, 1, len, &t); first = t[0]; } + if (f.prof) |p| p.mark(e.stream, .draw); if (e.mtp) { if (last) try e.copy(e.hold.ptr, e.bp.fnormed + (rows - 1) * e.c.hidden * 2, e.c.hidden * 2); const nxt = as_i32[1..]; if (nxt.len > 0) { // decode._absorb_rows + const mp = f.prof; + f.prof = null; // the head's pass as one section try f.mtp(nxt, &.{e.bp.fnormed}); + f.prof = mp; try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + nxt.len); } } + if (f.prof) |p| p.mark(e.stream, .mtp); try f.commit(rows); + if (f.prof) |p| p.mark(e.stream, .commit); return first; } diff --git a/zig/src/families/glm/cuda_prof.zig b/zig/src/families/glm/cuda_prof.zig index f948fa5792..a4bb46cdf1 100644 --- a/zig/src/families/glm/cuda_prof.zig +++ b/zig/src/families/glm/cuda_prof.zig @@ -2,8 +2,8 @@ const std = @import("std"); const cuda = @import("cuda"); -pub const Section = enum { embed, kda, mla, gather_a, dense, router, routed, shared, gather_b, head, mtp }; -const names = [_][]const u8{ "embed", "kda", "mla", "gather_a", "dense", "router", "routed", "shared", "gather_b", "head", "mtp" }; +pub const Section = enum { gate, embed, kda, mla, gather_a, dense, router, routed, shared, gather_b, head, draw, mtp, commit }; +const names = [_][]const u8{ "gate", "embed", "kda", "mla", "gather_a", "dense", "router", "routed", "shared", "gather_b", "head", "draw", "mtp", "commit" }; const max_marks = 512; pub const Prof = struct { @@ -19,6 +19,7 @@ pub const Prof = struct { wall_ms: f64 = 0, calls: u64 = 0, label: []const u8 = "", + every: u64 = 64, // calls a report sums pub fn init(d: *const cuda.Driver, io: std.Io) !*Prof { const p = try std.heap.page_allocator.create(Prof); @@ -61,7 +62,7 @@ pub const Prof = struct { p.gpu_ms += cuda.Event.elapsedMs(p.events[0], p.events[p.n]) catch 0; p.wall_ms += wall; p.calls += 1; - if (p.calls % 64 != 0) return; + if (p.calls % p.every != 0) return; const k: f64 = @floatFromInt(p.calls); var buf: [512]u8 = undefined; var w: std.Io.Writer = .fixed(&buf);