diff --git a/docs/api.md b/docs/api.md index 1cdb049853..0662487981 100644 --- a/docs/api.md +++ b/docs/api.md @@ -129,6 +129,8 @@ server reads images for GLM-5.3-Flash, Qwen3.8-27B and Gemma 4 31B on Metal, wit user messages ([video](vision.md#video-on-the-native-server)). See [image input](vision.md) for data URLs, public image URLs, limits and cache behavior. The native binary refuses `--vision` for other models and answers their image parts with HTTP 400. +The native CUDA server reads audio for Nemotron 3.5 Transcribe: `input_audio` and `audio_url` (data URL) parts with +16 kHz mono WAV in user messages ([recipe](recipes/nemotron-3.5-transcribe.md#audio-in-a-request)). Unsupported image input, audio, video and non-text output requests receive HTTP 400. ## Decisions diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b690ec8f9e..40a32b4a75 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -183,6 +183,11 @@ CUDA also checks Dory's activation rows (`act_l2.cuh`): the scaled sigmoid outpu each row then L2-normed, against float64's ranges at their bf16 points and expf's ulps (four controls, one point each). CUDA also checks `dense_rows`' scaled epilogue (Dory's logits, fp32(bf16(x . w)) times a per-column scale), bit-equal to its unscaled tile's outputs rounded and scaled on the host. +CUDA also checks Nemotron 3.5 Transcribe's audio encoders (`nemo_audio.cu`): the ASR encoder's q/k LayerNorm with +RoPE on 64 of 128 channels and the Sortformer's RoPE alone (each output within float64's bound at its three bf16 +points), attention over a window's valid keys (core/attention_range.zig's range, P as bf16), and the Sortformer head's +im2col, padding and relu rows, exact (five controls: a one-pass variance, a key short of the window, a padding row read +or kept, and a relu by `fmaxf`). HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out). RADIO's fused bias, activation and residual (`tf_vit_bias_act`) replaces `tf_vit_bias` then @@ -203,6 +208,9 @@ CUDA also checks the affine-3 experts' split decode pair, bit-equal to the pair one dropping a last round of K groups short of the warps, one chaining the groups in reverse). CUDA also checks Flash Next's NVFP4 experts on prompt tiles with the decode form's sums (SwiGLU gate-up and fp32 down, each bit-equal to the decode instance it replaces; the scaled-code prompt tile is the planted control). +CUDA also checks the bf16 experts (`experts_bf16.cu`, a checkpoint's own bf16 weights: Nemotron 3.5 Transcribe's routed +and shared experts): relu^2 up, fp32 down and bf16 down, in the decode and prompt forms, against float64 (two controls: +the fp32 down launched one 32-input group short, and a down summing in bf16). CUDA also checks Nemotron's pocket (`nemotron_ops.cu`, the learned neurons the living-weights learner adds beside the shared expert): relu^2 up into act rows and down added into the shared expert's slot, a block row a row and tiled, the down to fp32 and bf16, against float64, the neurons past the count NaN so a read of one shows (two controls: the up's diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 8d38798aed..24b7ac81d7 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -9,6 +9,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL. | --- | --- | | Nemotron 3.5 Lightning | [MLX](nemotron-3.5.md), [Intel Arc, experimental](xpu-arc.md) | | Nemotron 3.5 Super VL | [Native CUDA on one or two GB10s, images](nemotron-3.5-super-vl.md) | +| Nemotron 3.5 Transcribe | [Native CUDA on one GB10, the bf16 release, audio input](nemotron-3.5-transcribe.md) | | Qwen3.8-27B | [MLX, quantization and CUDA](qwen3.8-27b.md) | | Qwen3.5-2B | [Native Metal preview, tied affine 4-bit](qwen3.5-2b-native.md) | | Qwen3.8 Flash Next | [MLX prefill and CUDA](qwen3.8-flash-next.md), [native CUDA](flash-next-cuda.md), [CUDA images](flash-next-vision.md), [exact SSD residency on Metal](fluid-subgraph-residency.md) | diff --git a/docs/recipes/nemotron-3.5-transcribe.md b/docs/recipes/nemotron-3.5-transcribe.md new file mode 100644 index 0000000000..6105f88cf3 --- /dev/null +++ b/docs/recipes/nemotron-3.5-transcribe.md @@ -0,0 +1,107 @@ +# Nemotron 3.5 Transcribe + +NVIDIA's Nemotron 3.5 Transcribe 30B-A4B (`nemo_speechlm`) runs on the native CUDA engine from NVIDIA's bf16 release, +on one GB10. It is a speech LLM: an audio encoder turns a recording into rows that stand in the prompt where the +chat template writes `<|audio|>`, and a Nemotron 3.5 Lightning language model (Mamba-2, attention and MoE blocks, +128 routed experts with relu^2, a shared expert, one MTP layer) writes the transcript, translation or answer. The +language model is the Nemotron family (`zig/src/families/nemotron/`, the checkpoint's names in `cuda_salm.zig`); the +encoder is `zig/src/families/nemo_audio/`, with its kernels in `zig/kernels/cuda/nemo_audio.cu`. + +## Checkpoint + +`nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B` at revision `c421d6acdb436edd7ae00565f6414631fb9bd56d`: one +`model.safetensors` of 63 GiB in NeMo's layout, every tensor bf16 except the Mamba layers' fp32 parameters, the +routers' fp32 biases and the CTC timestamp head. The engine reads these tensors as stored; nothing is converted +ahead of time. The routed experts stay in bf16 on a grouped kernel of their own (`zig/kernels/cuda/experts_bf16.cu`). + +Not read: NVIDIA's FP8 and NVFP4 variants, the DFlash2 draft model, and the CTC timestamp head. + +## Run on one GB10 + +```bash +tensorfold serve /models/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B --name transcribe +``` + +The weights and the encoder take 64.7 GiB. The default window is the card's trained 65,536 tokens (config.json's +`llm_config` names the backbone's 262,144, which the speech model was not trained to; `--context` above 65,536 is +refused). At that window the server admits 8 streams, the engine's own at 0.50 GiB and 7 more taking 512-row pages +(3.09 GiB in all) as they need them, and keeps a prompt cache of 64 states (3.0 GiB) and 13.0 GiB of copied rows. + +The chat template turns thinking on by default. The card asks for transcription, translation and language +identification with thinking off and temperature 0, and for audio questions with thinking on, temperature 0.6 and +top_p 0.95. Send `chat_template_kwargs: {"enable_thinking": false}` per request, or start the server with +`--no-thinking`. + +## Audio in a request + +`/v1/chat/completions` (and `/v1/completions` given `messages`) takes audio in user messages as OpenAI's +`input_audio` part or as vLLM's `audio_url` part with a data URL: + +```json +{"messages": [{"role": "user", "content": [ + {"type": "text", "text": "Transcribe the following audio."}, + {"type": "input_audio", "input_audio": {"data": "", "format": "wav"}}]}], + "temperature": 0, "chat_template_kwargs": {"enable_thinking": false}} +``` + +- A clip is a WAV file: 16 kHz, mono, 16-bit PCM or 32-bit float, at least 10 ms. Other rates, channels and formats + get HTTP 400 naming the clip. +- A request holds up to 8 clips and up to an hour of audio in all (12.5 rows a second, 45,000 rows), and no more + than its window leaves after the text. Chat bodies may be 160 MiB on this model, enough for an hour of 16-bit + WAV in base64; other routes keep the 32 MiB limit. +- `/v1/messages` refuses audio: Anthropic's API has no audio block. +- Audio parts in assistant, system or tool messages are refused, as is text that spells `<|audio|>` with no clip for it. +- The card's [prompt guide](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B#prompt-guide) + lists the instructions for written or verbatim form, context terms, speaker tags (`` ...), translation and + language identification; they are plain text beside the clip. + +The encoder runs on each request's prompt pass: log-mel features on the host, the 32-layer Transformer ASR encoder +and the streaming Sortformer speaker encoder on the GPU, in 30-second windows, fused and projected to the language +model's 2,688 columns. A clip's rows are keyed by a hash of its samples in the prompt cache, so a repeated clip (a +second question about the same recording) resumes after it instead of encoding it again. `tensorfold audio-check +MODEL WAV OUT.safetensors` writes every encoder stage's rows for one clip, for comparison with NeMo's. + +## Drafts + +The checkpoint's MTP layer drafts by default. Its head reads the shifted prompt with the clip's audio rows in place, +as NeMo trains it. `--mtp-drafts` and `--mtp-confidence` set it as for the other Nemotron checkpoints; +`"draft": false` decodes serially. Copied continuations stay off for this model. The head adds 2.9 GiB (64.73 GiB against 61.85 with +`--no-drafts`). + +## Exactness and speed + +Measured on one GB10 at the default 65,536-token window, greedy, thinking off, "Transcribe the following audio.": + +- **Against NeMo's model in bf16:** on five LibriSpeech clips (3.5, 5.9, 34 and 130 seconds, and a 19.5-second + clip of two speakers), four replies are the reference's tokens, the 130-second clip's 379 among them. The 34-second + clip writes "at 'em" for the reference's "at em", a near tie that the engine's bf16 logits settle the other way. +- **Encoder stages:** at every stage the engine's rows are as far from NeMo's fp32 rows as NeMo's own bf16 rows are, + on eight clips; the speaker features are identical on all of them. +- **Drafted against `--no-drafts`:** byte-identical replies on 11, 30 and 60 minutes of audio. +- **Prompt cache:** a server with `--prompt-cache-gib 0` gave the same replies as the default one to every check + below: two instructions after the same 11-minute clip, a question after a transcribed clip, seeded sampling, + speaker tags, translation to German and French, language identification, audio questions, thinking on, a two-turn + chat and the Anthropic route with text. A repeated clip resumed 261 of 267 prompt tokens and answered its first + token in 0.067 s instead of 0.381 s. +- **Concurrency:** four long clips at once (11, 11 and 15 minutes and 130 seconds) took 81.8 s together, each reply + identical to its solo run. + +Long recordings, LibriSpeech test-clean read end to end without repeats (WER on the normalized text): + +| Audio | Prompt tokens | Reply tokens | WER | First token | MTP drafts | `--no-drafts` | +| --- | --- | --- | --- | --- | --- | --- | +| 11 min | 8,730 | 2,061 | 1.13% | 13.3 s | 40.4 s, 76.1 tok/s | 73.7 s, 34.1 tok/s | +| 30 min | 22,374 | 5,238 | 1.27% | 33.9 s | 104.2 s, 74.8 tok/s | 190.8 s, 33.4 tok/s | +| 60 min | 44,872 | 11,084 | 1.18% | 70.2 s | 228.9 s, 70.0 tok/s | 411.2 s, 32.4 tok/s | + +The first token waits for the encoder and the whole prompt pass. The drafts kept 72% to 76% of their proposals +(9,748 of 13,142 on the hour). The 130-second clip answered its first token in 2.32 s. The encoder's attention is a +plain per-window kernel, not yet a tiled one. The harness's quality suite was not run on this build. + +## Not served yet + +- Word timestamps and diarization segments (the card's `ctc_timestamps`, from its CTC head): the head stays + unloaded. +- Audio URLs other than data URLs, and formats other than WAV. +- The DFlash2 draft model, and the FP8 and NVFP4 variants. +- More than one GPU: the model fits one GB10, and the two-rank expert split is Super VL's alone. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 9d1e6e33a8..edf42f96d4 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -32,6 +32,8 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin nemotron.addImport("nemotron_draft_ids", draft_ids); const radio = radioModule(b, target, optimize, cuda, core, lanes); nemotron.addImport("radio", radio); + // Nemotron 3.5 Transcribe's audio encoder: the rows its prompt's <|audio|> positions take + nemotron.addImport("nemo_audio", b.createModule(.{ .root_source_file = b.path("zig/src/families/nemo_audio/audio.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core } } })); nemotron.addImport("living", living); const qwen35 = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3_5/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); qwen35.addImport("living", living); diff --git a/zig/build/cuda_family_host.zig b/zig/build/cuda_family_host.zig index 046b2b0c6a..f1844a2359 100644 --- a/zig/build/cuda_family_host.zig +++ b/zig/build/cuda_family_host.zig @@ -27,6 +27,9 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, mods: anytype, step: *std.B b.step("test-qwen35-formats-host", "MLX, ModelOpt and EXL3 admission and actual embedding dispatch without devices").dependOn(&formats_run.step); step.dependOn(&formats_run.step); @import("glm_contract_host.zig").host(b, cuda, mods.core, step); + const nemo_audio = b.addRunArtifact(b.addTest(.{ .root_module = mods.nemotron.import_table.get("nemo_audio").? })); + b.step("test-nemo-audio-host", "Transcribe's audio config, FFT, bf16 rounding and Sortformer FIFO without devices").dependOn(&nemo_audio.step); + step.dependOn(&nemo_audio.step); @import("dory.zig").host(b, @import("dory.zig").module(b, host, .debug, cuda, mods.core, mods.lanes), step); const attention = b.addRunArtifact(b.addTest(.{ .root_module = mods.qwen35, .filters = &.{"Qwen attention path:"} })); b.step("test-qwen35-attention-path-host", "Contiguous and paged CUDA attention selection and graph identity without devices").dependOn(&attention.step); diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 04bded45c9..f6e54c7a29 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -60,6 +60,7 @@ pub const kernels = [_]Kernel{ .{ .name = "exlr8", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 lane projections, the Metal kernels' order .{ .name = "exlr8_dense", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 dense matrices by row, the Metal order .{ .name = "nvfp4_experts_tile", .flags = &.{"-O3"}, .min_sm = 80 }, + .{ .name = "experts_bf16", .flags = &.{"-O3"} }, // ours: grouped bf16 experts on the plan, mma fragment blocks .{ .name = "nvfp4_experts", .flags = &.{"-O3"} }, // nvfp4/experts.cu's device code, the NVFP4 expert kernels .{ .name = "rank_sum", .flags = &.{"-O3"} }, // ours: the rank-order sum of a reduce-scatter (rank_scatter.zig) .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA @@ -84,6 +85,7 @@ pub const kernels = [_]Kernel{ .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle .{ .name = "moss_image", .flags = glue }, // ours: MOSS-Image's modulated norm, q/k rotary, situ, gated add + .{ .name = "nemo_audio", .flags = glue }, // ours: Nemotron 3.5 Transcribe's audio encoders (q/k norm and RoPE, attention, Sortformer glue) .{ .name = "tree_round", .flags = &.{"-O3"} }, // ours: verified tree rounds for every family (cuda/tree_round.zig) .{ .name = "act_rows", .flags = &.{"-O3"} }, // ours: SwiGLU, output gates, group sums (cuda/act_rows.zig) .{ .name = "gdn_prep", .flags = &.{"-O3"} }, // ours: Gated DeltaNet conv, norms, g and beta (cuda/gdn_prep.zig) diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 63ff744b47..6718c0746f 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -352,6 +352,9 @@ TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_us, tf_act_l2::GS | tf_act_l2::PROD) TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_prod, tf_act_l2::GS | tf_act_l2::US) #undef TF_CTL_SWIGLU_L2 +// nemo_audio.cu's controls: tf_na_qk_rope with a one-pass variance (gate 3), its relu as fmaxf (gate 1). +#include "contract_nemo_audio_controls.cuh" + // Gate 1: tf_pocket_in (a warp a neuron, act rows 64 wide) with relu as fmaxf(s, 0), which turns a NaN sum into 0. extern "C" __global__ void __launch_bounds__(128) tf_ctl_pocket_in_fmax(const __nv_bfloat16* y, int d, const float* up, const int* used, float* act, int rows) { @@ -538,3 +541,26 @@ TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_old_2, 2, tf_exl3::PLAIN) TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_skip, 2, tf_exl3::SKIP) TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_stale, 2, tf_exl3::STALE) #undef TF_CTL_EXL3_GROUPED + +// Gate 4: experts_bf16.cu's fp32 down with each pair's sum rounded to bf16 at every input, a thread a column. +extern "C" __global__ void tf_ctl_bf16_experts_bf16_sums(const __nv_bfloat16* X, int x_stride, const uint4* W, int KG, int NB, + const int* items, const int* counts, const int* members, float* out, + int N) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + if (col >= N) return; + const int cb = col / 32, j = (col % 32) / 8; + for (int it = blockIdx.y; it < counts[0]; it += gridDim.y) { + const int e = items[3 * it], first = items[3 * it + 1], cnt = items[3 * it + 2]; + const __nv_bfloat16* blk = reinterpret_cast(W + (static_cast(e) * NB + cb) * KG * 128); + for (int i = 0; i < cnt; ++i) { + const int pr = members[first + i]; + float s = 0.0f; + for (int k = 0; k < KG * 32; ++k) { + const int q = k % 16, lane = 4 * (col % 8) + q / 4; + const float w = __bfloat162float(blk[(static_cast(k / 32) * 128 + 32 * j + lane) * 8 + 4 * (k % 32 / 16) + q % 4]); + s = __bfloat162float(__float2bfloat16_rn(s + tf_ctl_x(X, static_cast(pr) * x_stride + k) * w)); + } + out[static_cast(pr) * N + col] = s; + } + } +} diff --git a/zig/kernels/cuda/contract_nemo_audio_controls.cuh b/zig/kernels/cuda/contract_nemo_audio_controls.cuh new file mode 100644 index 0000000000..394240646f --- /dev/null +++ b/zig/kernels/cuda/contract_nemo_audio_controls.cuh @@ -0,0 +1,56 @@ +// nemo_audio.cu's planted controls (Nemotron 3.5 Transcribe's audio encoder); only the contract controls include this. +#pragma once +#include + +// tf_na_qk_rope with a one-pass variance, E[x^2] - mean^2: rows far from zero lose their digits, so gate 3 goes red. +__device__ __forceinline__ float tf_ctl_na_bf(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float tf_ctl_na_rbf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); } +__device__ __forceinline__ float tf_ctl_na_sum(float v) { + for (int o = 16; o > 0; o >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, o)); + return v; +} + +extern "C" __global__ void __launch_bounds__(32) tf_ctl_na_qk_onepass(__nv_bfloat16* qkv, const __nv_bfloat16* qw, const __nv_bfloat16* qb, + const __nv_bfloat16* kw, const __nv_bfloat16* kb, const __nv_bfloat16* cs, + const __nv_bfloat16* sn, int H, int D, int R, float eps) { + __shared__ float x[128]; + const int which = blockIdx.x & 1; + const int h = (blockIdx.x >> 1) % H; + const int t = (blockIdx.x >> 1) / H; + __nv_bfloat16* row = qkv + ((static_cast(t) * 3 + which) * H + h) * D; + const int lane = threadIdx.x; + for (int i = lane; i < D; i += 32) x[i] = tf_ctl_na_bf(row[i]); + __syncwarp(); + const __nv_bfloat16* w = which ? kw : qw; + const __nv_bfloat16* b = which ? kb : qb; + float acc = 0.0f, sq = 0.0f; + for (int i = lane; i < D; i += 32) { + acc = __fadd_rn(acc, x[i]); + sq = __fadd_rn(sq, __fmul_rn(x[i], x[i])); + } + const float mean = __fdiv_rn(tf_ctl_na_sum(acc), static_cast(D)); + const float var = __fsub_rn(__fdiv_rn(tf_ctl_na_sum(sq), static_cast(D)), __fmul_rn(mean, mean)); + const float rstd = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(fmaxf(var, 0.0f), eps))); + float y[4]; + for (int i = lane, k = 0; i < D; i += 32, ++k) y[k] = tf_ctl_na_rbf(__fadd_rn(__fmul_rn(__fmul_rn(__fsub_rn(x[i], mean), rstd), tf_ctl_na_bf(w[i])), tf_ctl_na_bf(b[i]))); + __syncwarp(); + for (int i = lane, k = 0; i < D; i += 32, ++k) x[i] = y[k]; + __syncwarp(); + const int half = R / 2; + for (int i = lane; i < D; i += 32) { + float v = x[i]; + if (i < R) { + const float c = tf_ctl_na_bf(cs[static_cast(t) * R + i]); + const float s = tf_ctl_na_bf(sn[static_cast(t) * R + i]); + const float r = i < half ? -x[i + half] : x[i - half]; + v = __fadd_rn(tf_ctl_na_rbf(__fmul_rn(v, c)), tf_ctl_na_rbf(__fmul_rn(r, s))); + } + row[i] = __float2bfloat16_rn(v); + } +} + +// Gate 1: nemo_audio's relu as fmaxf(x, 0), which turns a NaN into 0 where tf_na_relu keeps it. +extern "C" __global__ void __launch_bounds__(256) tf_ctl_na_relu_fmax(__nv_bfloat16* x, int64_t count) { + const int64_t at = static_cast(blockIdx.x) * 256 + threadIdx.x; + if (at < count) x[at] = __float2bfloat16_rn(fmaxf(__bfloat162float(x[at]), 0.0f)); +} diff --git a/zig/kernels/cuda/experts_bf16.cu b/zig/kernels/cuda/experts_bf16.cu new file mode 100644 index 0000000000..d4f29bf011 --- /dev/null +++ b/zig/kernels/cuda/experts_bf16.cu @@ -0,0 +1,120 @@ +// Ours: grouped bf16 experts on mma-fragment blocks; a row's sums chain in input order, whatever its tile or form. + +#include +#include + +#include "experts.cuh" + +namespace tf_bf16_experts { + +constexpr int WV = 4; // uint4 a lane a block + +__device__ __forceinline__ uint4 ld_nc(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; +} + +// RT row tiles of 16 pairs share each block's fragments; a unit is one item's 32-column block. +template +__device__ __forceinline__ void body(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, int skip) { + 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); + if (e == skip) continue; + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (32 * WV) + lane; + for (int r0 = 0; r0 < cnt; r0 += 16 * RT) { + bool v[RT][2]; + int pr[RT][2]; + const __nv_bfloat16* x[RT][2]; +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int i = r0 + 16 * r + gq + 8 * h; + v[r][h] = i < cnt; + pr[r][h] = v[r][h] ? __ldg(members + first + i) : 0; + x[r][h] = X + (size_t)(slots ? pr[r][h] / slots : pr[r][h]) * x_stride + 4 * t; + } + float acc[1][RT][NTW][4]; +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[0][r][j][0] = acc[0][r][j][1] = acc[0][r][j][2] = acc[0][r][j][3] = 0.f; + const uint2 zero = make_uint2(0u, 0u); + for (int g = 0; g < KG; ++g) { + uint4 w[NTW]; +#pragma unroll + for (int j = 0; j < NTW; ++j) w[j] = ld_nc(blk + (size_t)g * (32 * WV) + 32 * j); +#pragma unroll + for (int r = 0; r < RT; ++r) { + uint2 xa[2], xb[2]; +#pragma unroll + for (int q = 0; q < 2; ++q) { + xa[q] = v[r][0] ? __ldg(reinterpret_cast(x[r][0] + 32 * g + 16 * q)) : zero; + xb[q] = v[r][1] ? __ldg(reinterpret_cast(x[r][1] + 32 * g + 16 * q)) : zero; + } +#pragma unroll + for (int q = 0; q < 2; ++q) +#pragma unroll + for (int j = 0; j < NTW; ++j) + mma(acc[0][r][j], xa[q].x, xb[q].x, xa[q].y, xb[q].y, q ? w[j].z : w[j].x, q ? w[j].w : w[j].y); + } + } +#pragma unroll + for (int r = 0; r < RT; ++r) + epilogue(acc, r, out, N, cb * COLS + 2 * t, pr[r][0], pr[r][1], v[r][0], v[r][1], 0.f); + } + } +} + +} // namespace tf_bf16_experts + +// Decode (items of up to 16 pairs) and prompt (up to 64) entries per epilogue: the NVFP4 entries' arguments, no scales. +#define TF_BF16_EXPERTS(NAME, EPI, RT, WARPS) \ + extern "C" __global__ void __launch_bounds__(WARPS * 32) \ + NAME(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, int skip) { \ + tf_bf16_experts::body(X, x_stride, slots, W, KG, NB, items, counts, members, out, N, skip); \ + } +TF_BF16_EXPERTS(tf_bf16_experts_up, 1, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_down_f32, 0, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_down_bf16, 3, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_up, 1, 4, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_down_f32, 0, 4, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_down_bf16, 3, 4, 4) + +// The repack, one thread a uint4: matrix e's [k][n] (in_out, x @ W) or [n][k] (nn.Linear) rows into its blocks. +extern "C" __global__ void tf_bf16_experts_pack(const __nv_bfloat16* __restrict__ src, long long e_stride, int ld, + int n0, int k0, int in_out, int n, int k, uint4* __restrict__ dst) { + const int e = blockIdx.y; + const long long per = (long long)(n / 32) * (k / 32) * 32 * tf_bf16_experts::WV; + const long long q = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (q >= per) return; + const int lane = (int)(q % 32), j = (int)(q / 32 % tf_bf16_experts::WV); + const long long b = q / (32 * tf_bf16_experts::WV); + const int kg = k / 32, g = (int)(b % kg), cb = (int)(b / kg); + const int row = n0 + cb * 32 + 8 * j + (lane >> 2), col = k0 + g * 32 + 4 * (lane & 3); + const __nv_bfloat16* s = src + e * e_stride; + uint16_t h[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int c = col + (i & 3) + 16 * (i >> 2); + const __nv_bfloat16 v = in_out ? s[(size_t)c * ld + row] : s[(size_t)row * ld + c]; + h[i] = __bfloat16_as_ushort(v); + } + uint4 o; + o.x = h[0] | (uint32_t)h[1] << 16; + o.y = h[2] | (uint32_t)h[3] << 16; + o.z = h[4] | (uint32_t)h[5] << 16; + o.w = h[6] | (uint32_t)h[7] << 16; + dst[e * per + q] = o; +} diff --git a/zig/kernels/cuda/nemo_audio.cu b/zig/kernels/cuda/nemo_audio.cu new file mode 100644 index 0000000000..99776119f9 --- /dev/null +++ b/zig/kernels/cuda/nemo_audio.cu @@ -0,0 +1,141 @@ +// Nemotron 3.5 Transcribe's audio encoder kernels (families/nemo_audio), rounding as torch's bf16 path does op for op. + +#include +#include +#include + +#include "common/glue_math.cuh" + +using namespace tf_glue; + +namespace { + +constexpr int MAXD = 128; // the widest head: the ASR encoder's 1280 / 10 +constexpr int AT = 128; // attention threads a block + +__device__ __forceinline__ float warp_max(float v) { +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + return v; +} + +// A block's values reduced: warps in order, every thread gets the result. +template +__device__ __forceinline__ float block_reduce(float v, float* red) { + v = MAX ? warp_max(v) : warp_sum(v); + if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = v; + __syncthreads(); + float r = red[0]; + for (int w = 1; w < AT / 32; ++w) r = MAX ? fmaxf(r, red[w]) : __fadd_rn(r, red[w]); + __syncthreads(); + return r; +} + +} // namespace + +// A warp per (row, head, q or k) of qkv [T, 3, H, D], D <= 128: LayerNorm when qw, then RoPE on R channels, as torch. +extern "C" __global__ void __launch_bounds__(32) tf_na_qk_rope(__nv_bfloat16* qkv, const __nv_bfloat16* qw, const __nv_bfloat16* qb, + const __nv_bfloat16* kw, const __nv_bfloat16* kb, const __nv_bfloat16* cs, + const __nv_bfloat16* sn, int H, int D, int R, float eps) { + __shared__ float x[MAXD]; + const int which = blockIdx.x & 1; + const int h = (blockIdx.x >> 1) % H; + const int t = (blockIdx.x >> 1) / H; + __nv_bfloat16* row = qkv + ((static_cast(t) * 3 + which) * H + h) * D; + const int lane = threadIdx.x; + for (int i = lane; i < D; i += 32) x[i] = bf(row[i]); + __syncwarp(); + const __nv_bfloat16* w = which ? kw : qw; + const __nv_bfloat16* b = which ? kb : qb; + if (w) { + float acc = 0.0f; + for (int i = lane; i < D; i += 32) acc = __fadd_rn(acc, x[i]); + const float mean = __fdiv_rn(warp_sum(acc), static_cast(D)); + acc = 0.0f; + for (int i = lane; i < D; i += 32) { + const float d = __fsub_rn(x[i], mean); + acc = __fadd_rn(acc, __fmul_rn(d, d)); + } + const float var = __fdiv_rn(warp_sum(acc), static_cast(D)); + const float rstd = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(var, eps))); + float y[MAXD / 32]; + for (int i = lane, k = 0; i < D; i += 32, ++k) y[k] = rbf(__fadd_rn(__fmul_rn(__fmul_rn(__fsub_rn(x[i], mean), rstd), bf(w[i])), bf(b[i]))); + __syncwarp(); + for (int i = lane, k = 0; i < D; i += 32, ++k) x[i] = y[k]; + __syncwarp(); + } + const int half = R / 2; + for (int i = lane; i < D; i += 32) { + float v = x[i]; + if (i < R) { + const float c = bf(cs[static_cast(t) * R + i]); + const float s = bf(sn[static_cast(t) * R + i]); + const float r = i < half ? -x[i + half] : x[i - half]; + v = __fadd_rn(rbf(__fmul_rn(v, c)), rbf(__fmul_rn(r, s))); + } + row[i] = __float2bfloat16_rn(v); + } +} + +// A block per (query row, head): fp32 scores over keys [0, len), softmax, P as bf16, P V in fp32; T floats shared. +extern "C" __global__ void __launch_bounds__(AT) tf_na_attention(const __nv_bfloat16* __restrict__ qkv, __nv_bfloat16* __restrict__ out, + int H, int D, int len, float scale) { + extern __shared__ float p[]; + __shared__ float q[MAXD]; + __shared__ float red[AT / 32]; + const int t = blockIdx.x; + const int h = blockIdx.y; + const int64_t stride = 3LL * H * D; // one row of qkv + const __nv_bfloat16* qr = qkv + t * stride + static_cast(h) * D; + for (int i = threadIdx.x; i < D; i += AT) q[i] = bf(qr[i]); + __syncthreads(); + const __nv_bfloat16* k0 = qkv + static_cast(H + h) * D; + const __nv_bfloat16* v0 = qkv + static_cast(2 * H + h) * D; + float top = -INFINITY; + for (int j = threadIdx.x; j < len; j += AT) { + const __nv_bfloat16* kr = k0 + j * stride; + float s = 0.0f; + for (int i = 0; i < D; ++i) s = __fmaf_rn(q[i], bf(kr[i]), s); + s = __fmul_rn(s, scale); + p[j] = s; + top = fmaxf(top, s); + } + top = block_reduce(top, red); + float sum = 0.0f; + for (int j = threadIdx.x; j < len; j += AT) { + const float e = expf(__fsub_rn(p[j], top)); + p[j] = e; + sum = __fadd_rn(sum, e); + } + sum = block_reduce(sum, red); + for (int j = threadIdx.x; j < len; j += AT) p[j] = rbf(__fdiv_rn(p[j], sum)); + __syncthreads(); + for (int i = threadIdx.x; i < D; i += AT) { + float acc = 0.0f; + for (int j = 0; j < len; ++j) acc = __fmaf_rn(p[j], bf(v0[j * stride + i]), acc); + out[(static_cast(t) * H + h) * D + i] = __float2bfloat16_rn(acc); + } +} + +// The subpixel conv's input as a GEMM's: y [T, 3 C], y[t, j C + c] = x[t - 1 + j, c], zero outside rows [0, len). +extern "C" __global__ void __launch_bounds__(256) tf_na_im2col3(const __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ y, int T, + int C, int len) { + const int64_t at = static_cast(blockIdx.x) * 256 + threadIdx.x; + if (at >= static_cast(T) * 3 * C) return; + const int t = static_cast(at / (3 * C)); + const int r = static_cast(at % (3 * C)); + const int src = t - 1 + r / C; + y[at] = (src >= 0 && src < len && src < T) ? x[static_cast(src) * C + r % C] : __float2bfloat16_rn(0.0f); +} + +// x = bf16(relu(x)) in place. +extern "C" __global__ void __launch_bounds__(256) tf_na_relu(__nv_bfloat16* x, int64_t count) { + const int64_t at = static_cast(blockIdx.x) * 256 + threadIdx.x; + if (at < count && bf(x[at]) < 0.0f) x[at] = __float2bfloat16_rn(0.0f); +} + +// Rows [len, T) of x [T, C] zeroed in place (torch's masked_fill of a window's padding). +extern "C" __global__ void __launch_bounds__(256) tf_na_zero_rows(__nv_bfloat16* x, int T, int C, int len) { + const int64_t at = static_cast(len) * C + static_cast(blockIdx.x) * 256 + threadIdx.x; + if (at < static_cast(T) * C) x[at] = __float2bfloat16_rn(0.0f); +} diff --git a/zig/src/cli/cuda_audio_check.zig b/zig/src/cli/cuda_audio_check.zig new file mode 100644 index 0000000000..157693b7eb --- /dev/null +++ b/zig/src/cli/cuda_audio_check.zig @@ -0,0 +1,38 @@ +//! audio-check: a WAV clip through Transcribe's audio encoder, every stage to safetensors under NeMo's names. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const nemo_audio = @import("nemotron").nemo_audio; + +pub fn run(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, model: []const u8, rest: []const []const u8) !u8 { + if (rest.len != 2) { + std.debug.print("audio-check: give MODEL WAV OUT.safetensors\n", .{}); + return 2; + } + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, rest[0], gpa, .limited(1 << 31)); + defer gpa.free(bytes); + const wav = try nemo_audio.mel.decodeWav(gpa, bytes); + defer gpa.free(wav); + const e = try nemo_audio.Encoder.init(gpa, io, ctx.d, model); + defer e.deinit(); + var s = try cuda.Stream.init(ctx.d, true); + defer s.deinit(); + const n = nemo_audio.rows(wav.len); + var out = try cuda.DeviceBuffer.alloc(ctx.d, @max(1, n) * e.c.text_hidden * 2); + defer out.free(); + var dump: nemo_audio.Dump = .{ .gpa = gpa, .d = ctx.d }; + defer dump.deinit(); + // one pass without the dump for the time, then the dumped one + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + const got = try e.encode(s, wav, out.ptr, null); + try s.synchronize(); + const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - t0)) / 1e6; + _ = try e.encode(s, wav, out.ptr, &dump); + var tensors: std.ArrayList(core.shard_edit.New) = .empty; + defer tensors.deinit(gpa); + for (dump.items.items) |*it| try tensors.append(gpa, .{ .name = it.name, .dtype = .f32, .shape = &it.shape, .bytes = std.mem.sliceAsBytes(it.data) }); + std.Io.Dir.cwd().deleteFile(io, rest[1]) catch |err| if (err != error.FileNotFound) return err; + try core.shard_edit.create(gpa, rest[1], &.{}, tensors.items); + std.debug.print("RESULT audio-check: {d} samples ({d:.2} s) -> {d} rows in {d:.1} ms; {d} stages to {s}\n", .{ wav.len, @as(f64, @floatFromInt(wav.len)) / 16000.0, got, ms, dump.items.items.len, rest[1] }); + return 0; +} diff --git a/zig/src/cli/cuda_main.zig b/zig/src/cli/cuda_main.zig index 6fdea9e952..ed81b66c40 100644 --- a/zig/src/cli/cuda_main.zig +++ b/zig/src/cli/cuda_main.zig @@ -9,6 +9,7 @@ const checks = @import("cuda_checks.zig"); const kernel_checks = @import("cuda_kernel_checks.zig"); const precision = @import("cuda_precision.zig"); const kl = @import("cuda_kl.zig"); +const audio_check = @import("cuda_audio_check.zig"); const shared_checks = @import("cuda_shared_checks.zig"); const checkpoint_cli = @import("checkpoint_cli"); const lanes_cli = @import("cuda_lanes.zig"); @@ -44,6 +45,7 @@ const usage = \\ (logits rows to FILE, fed --gold then greedy; or FILE's tokens fed: KL, top-1 and NLL against its rows) \\ tensorfold widths MODEL --tokens ID,... [--max-tokens N] [--eager] [sampling as run] \\ (every verify width, 1-32 for a latent MoE else 1-16, against serial decoding, a wrong draft every third window) + \\ tensorfold audio-check MODEL WAV OUT.safetensors (Transcribe's audio encoder, every stage to OUT) \\ tensorfold models \\ tensorfold info MODEL \\ tensorfold pull REPO[@REVISION] @@ -188,6 +190,7 @@ pub fn main(init: std.process.Init) !u8 { return 2; }, }; + if (std.mem.eql(u8, args[1], "audio-check")) return audio_check.run(gpa, init.io, &ctx, opts.model, positional.items); if (std.mem.eql(u8, args[1], "run") and dory_cli.wants(gpa, init.io, opts.model)) { const s: ?lanes.Sampling = if (opts.sampling.temperature > 0) opts.sampling else null; return dory_cli.run(gpa, init.io, &ctx, opts.model, .{ .tokens = opts.tokens, .max_tokens = opts.max_tokens, .loops = opts.loops, .context = opts.context, .drafts = opts.drafts, .graphs = opts.graphs, .stop_eos = opts.stop_eos, .sampling = s }); diff --git a/zig/src/core/contract/experts_bf16_records.zig b/zig/src/core/contract/experts_bf16_records.zig new file mode 100644 index 0000000000..5ed84f9936 --- /dev/null +++ b/zig/src/core/contract/experts_bf16_records.zig @@ -0,0 +1,30 @@ +//! The bf16 experts' CUDA launches (experts_bf16.cu): decode and prompt rows by epilogue, and the repack. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage, the compiler of the release fatbins. +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_up", .symbol = "tf_bf16_experts_up", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_down_f32", .symbol = "tf_bf16_experts_down_f32", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_down_bf16", .symbol = "tf_bf16_experts_down_bf16", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_up", .symbol = "tf_bf16_experts_prompt_up", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_f32", .symbol = "tf_bf16_experts_prompt_down_f32", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_bf16", .symbol = "tf_bf16_experts_prompt_down_bf16", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_pack", .symbol = "tf_bf16_experts_pack", .image = "experts_bf16", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, +}; + +/// nvcc 13.2's prompt forms on sm_120 and sm_121 (164 registers where 13.0 takes 162); the rest compile as 13.0's. +pub const cuda_132 = [_]registry.Override{ + .{ .name = "experts_bf16 tf_bf16_experts_prompt_up", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_f32", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_bf16", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, +}; + +test "bf16 experts CUDA: every launch is in the live registry and fits every SM" { + const t = @import("std").testing; + for (cuda) |r| { + var found: usize = 0; + for (registry.all_cuda) |active| found += @intFromBool(@import("std").mem.eql(u8, active.launch.name, r.launch.name)); + try t.expectEqual(@as(usize, 1), found); + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{r}, true)); + } +} diff --git a/zig/src/core/contract/nemo_audio_cuda_records.zig b/zig/src/core/contract/nemo_audio_cuda_records.zig new file mode 100644 index 0000000000..86341fc15f --- /dev/null +++ b/zig/src/core/contract/nemo_audio_cuda_records.zig @@ -0,0 +1,21 @@ +//! Nemotron 3.5 Transcribe's audio encoder launches (nemo_audio.cu); attention holds up to 2,048 scores shared. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage, the compiler of the release fatbins; nvcc 13.2 compiles each the same. +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "nemo_audio tf_na_qk_rope", .symbol = "tf_na_qk_rope", .image = "nemo_audio", .threads = 32 }, .usage = .{ .{ .regs = 38, .static = 512, .bound = 32 }, .{ .regs = 39, .static = 512, .bound = 32 }, .{ .regs = 39, .static = 512, .bound = 32 }, .{ .regs = 39, .static = 512, .bound = 32 }, .{ .regs = 39, .static = 512, .bound = 32 } } }, + .{ .launch = .{ .name = "nemo_audio tf_na_attention", .symbol = "tf_na_attention", .image = "nemo_audio", .threads = 128, .dynamic = 8192 }, .usage = .{ .{ .regs = 32, .static = 528, .bound = 128 }, .{ .regs = 40, .static = 528, .bound = 128 }, .{ .regs = 40, .static = 528, .bound = 128 }, .{ .regs = 40, .static = 528, .bound = 128 }, .{ .regs = 40, .static = 528, .bound = 128 } } }, + .{ .launch = .{ .name = "nemo_audio tf_na_im2col3", .symbol = "tf_na_im2col3", .image = "nemo_audio", .threads = 256 }, .usage = .{ .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 16, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "nemo_audio tf_na_relu", .symbol = "tf_na_relu", .image = "nemo_audio", .threads = 256 }, .usage = .{ .{ .regs = 8, .static = 0, .bound = 256 }, .{ .regs = 8, .static = 0, .bound = 256 }, .{ .regs = 8, .static = 0, .bound = 256 }, .{ .regs = 8, .static = 0, .bound = 256 }, .{ .regs = 8, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "nemo_audio tf_na_zero_rows", .symbol = "tf_na_zero_rows", .image = "nemo_audio", .threads = 256 }, .usage = .{ .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 10, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } }, +}; + +test "Nemotron 3.5 Transcribe audio CUDA: every launch is in the live registry and fits every SM" { + const t = @import("std").testing; + for (cuda) |r| { + var found: usize = 0; + for (registry.all_cuda) |active| found += @intFromBool(@import("std").mem.eql(u8, active.launch.name, r.launch.name)); + try t.expectEqual(@as(usize, 1), found); + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{r}, true)); + } +} diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 4a8883b56f..b508a44281 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -205,7 +205,7 @@ pub const cuda = @import("flashnext_cuda_records.zig").cuda ++ [_]CudaRecord{ .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, -} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @import("moss_image_cuda_records.zig").cuda; +} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("experts_bf16_records.zig").cuda ++ @import("nemo_audio_cuda_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @import("moss_image_cuda_records.zig").cuda; /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins. pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 }; @@ -315,7 +315,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "latent tf_super_topk_few", .usage = .{ .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 } } }, .{ .name = "ck tf_ck4_up_relu2_few", .usage = .{ .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 } } }, .{ .name = "mamba tf_nemo_scan", .usage = .{ .{ .regs = 138, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 } } }, - } ++ learning.cuda_132 ++ @import("paging_pending.zig").cuda_132 ++ @import("nemotron_prompt_records.zig").cuda_132 ++ @import("glm_cuda_records.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132 ++ @import("nvfp4_fact_records.zig").cuda_132 ++ @import("pocket_records.zig").cuda_132 ++ @import("minimax_cuda_records.zig").cuda_132) }, + } ++ learning.cuda_132 ++ @import("paging_pending.zig").cuda_132 ++ @import("nemotron_prompt_records.zig").cuda_132 ++ @import("glm_cuda_records.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132 ++ @import("nvfp4_fact_records.zig").cuda_132 ++ @import("experts_bf16_records.zig").cuda_132 ++ @import("pocket_records.zig").cuda_132 ++ @import("minimax_cuda_records.zig").cuda_132) }, .{ .release = .{ .major = 13, .minor = 3 }, .usage = &([_]Override{ .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 } } }, } ++ @import("flashnext_tile_records.zig").cuda_133 ++ @import("glm_cuda_records.zig").cuda_133 ++ @import("nvcc_133_records.zig").cuda_133 ++ @import("minimax_cuda_records.zig").cuda_133) }, diff --git a/zig/src/core/contract/registry_sets_test.zig b/zig/src/core/contract/registry_sets_test.zig index 743c4b1a8c..3355ced419 100644 --- a/zig/src/core/contract/registry_sets_test.zig +++ b/zig/src/core/contract/registry_sets_test.zig @@ -54,4 +54,6 @@ test { _ = @import("exl3_grouped_records.zig"); _ = @import("proposal_sampling_records.zig"); _ = @import("pocket_records.zig"); + _ = @import("experts_bf16_records.zig"); + _ = @import("nemo_audio_cuda_records.zig"); } diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index ad35326aba..f0f393ece8 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -17,6 +17,7 @@ pub const load_ledger = @import("load_ledger.zig"); /// One image of a request (``Request.images``): the tower's input and the placeholder rows its output replaces. pub const Image = lanes.stream.Image; +pub const Audio = lanes.stream.Audio; /// A prompt's reproducible seed (``seed_for``), for requests that name none. pub const seedFor = lanes.sampling.seedFor; @@ -24,23 +25,9 @@ pub const seedFor = lanes.sampling.seedFor; /// A committed token's log probability under the target's raw distribution, with the best tokens at its position. pub const LogprobRow = lanes.LogprobRow; -/// The request's stop strings, checked on the engine's thread after each committed token (the server decodes the tail). -pub const Stop = struct { - ctx: *anyopaque, - check: *const fn (ctx: *anyopaque, emitted: []const u32) bool, -}; - -/// tool_choice "required" or a named function: outside a think block the answer opens a call to an offered tool. -pub const CallGate = struct { - /// The token that opens a call, the template text before the name, and the mark that ends the name. - opener: u32, - lead: []const u8 = "", - tail: []const u8 = "", - /// Offered names (empty: any); the gate completes a name its prefix starts. - names: []const []const u8 = &.{}, - think_open: ?u32 = null, - think_end: ?u32 = null, -}; +/// A request's stop strings and tool-call gate (engine_gates.zig). +pub const Stop = @import("engine_gates.zig").Stop; +pub const CallGate = @import("engine_gates.zig").CallGate; /// response_format and the guided_* fields: every token keeps the reply inside this grammar. pub const Structure = struct { @@ -115,6 +102,8 @@ pub const Request = struct { prompt_logprobs: ?u8 = null, /// Images whose rows replace placeholder tokens of ``prompt``, in order; only when ``Info.vision`` is set. images: []const Image = &.{}, + /// Audio clips whose rows replace placeholder tokens of ``prompt``, in order; only when ``Info.audio`` is set. + audios: []const Audio = &.{}, /// The prompt as the prompt cache matches it (each image's placeholders by its content); empty: ``prompt``. cache_key: []const u32 = &.{}, @@ -138,6 +127,8 @@ pub const VideoSource = vision_types.VideoSource; pub const FrameSink = vision_types.FrameSink; pub const PreparedVideo = vision_types.PreparedVideo; pub const VideoOffer = vision_types.VideoOffer; +pub const PreparedAudio = vision_types.PreparedAudio; +pub const AudioOffer = vision_types.AudioOffer; pub const VisionOpen = vision_types.VisionOpen; pub const SpeechRequest = @import("speech_api.zig").SpeechRequest; @@ -240,6 +231,8 @@ pub const Info = struct { pictures: bool = false, /// Image input (--vision); null: text only. vision: ?Vision = null, + /// Audio input (a speech model's encoder); null: none. + audio: ?AudioOffer = null, /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, /// The learner reads conversation prompts at any length through taps. diff --git a/zig/src/core/engine_gates.zig b/zig/src/core/engine_gates.zig new file mode 100644 index 0000000000..a14ec41b77 --- /dev/null +++ b/zig/src/core/engine_gates.zig @@ -0,0 +1,19 @@ +//! A request's stop strings and tool-call gate, checked on the engine's thread (re-exported by engine_api.zig). + +/// The request's stop strings, checked on the engine's thread after each committed token (the server decodes the tail). +pub const Stop = struct { + ctx: *anyopaque, + check: *const fn (ctx: *anyopaque, emitted: []const u32) bool, +}; + +/// tool_choice "required" or a named function: outside a think block the answer opens a call to an offered tool. +pub const CallGate = struct { + /// The token that opens a call, the template text before the name, and the mark that ends the name. + opener: u32, + lead: []const u8 = "", + tail: []const u8 = "", + /// Offered names (empty: any); the gate completes a name its prefix starts. + names: []const []const u8 = &.{}, + think_open: ?u32 = null, + think_end: ?u32 = null, +}; diff --git a/zig/src/core/engine_vision.zig b/zig/src/core/engine_vision.zig index 987b3c593b..3232c198ad 100644 --- a/zig/src/core/engine_vision.zig +++ b/zig/src/core/engine_vision.zig @@ -78,5 +78,22 @@ pub const VideoOffer = struct { max_frames: u32, }; +/// An audio clip prepared for the engine's encoder (``AudioOffer.prepare``): samples, rows and a content hash. +pub const PreparedAudio = struct { samples: []f32, tokens: u32, hash: u64 }; + +/// Audio input an engine serves (``Info.audio``): a clip's bytes to encoder input, and what a request may hold. +pub const AudioOffer = struct { + ctx: *anyopaque, + /// The clip's bytes (a WAV file) at most ``max_tokens`` rows; allocated in ``a``. + prepare: *const fn (ctx: *anyopaque, a: Allocator, bytes: []const u8, max_tokens: u32) anyerror!PreparedAudio, + /// The template's once-a-clip token, expanded to its rows (0 from an engine; the server sets ``locator``'s id). + audio_token: u32 = 0, + /// The text a user message's content starts with once a clip (vLLM's string form, which the template moves). + locator: []const u8, + /// Audio rows one request may hold in all, and clips. + audio_tokens: u32, + max_audios: u32, +}; + /// --vision: image and video input's limits (``Open.vision``); image_tokens_set: the flag given, over a family default. pub const VisionOpen = struct { image_tokens: u32 = 4096, image_tokens_set: bool = false, max_images: u32 = 4, video_tokens: u32 = 16384, max_videos: u32 = 2, max_frames: u32 = 256 }; diff --git a/zig/src/core/lane_host_fetch.zig b/zig/src/core/lane_host_fetch.zig index 147cf42f2b..f372b1d89f 100644 --- a/zig/src/core/lane_host_fetch.zig +++ b/zig/src/core/lane_host_fetch.zig @@ -67,7 +67,7 @@ fn named(h: *LaneHost, name: Name) bool { /// A request the host's lookup keys by its tokens alone (no pictures), as `lane_stream.open` looks it up. fn keyed(r: *const api.Request) bool { - return r.embedding == null and r.prompt_logprobs == null and !r.untaught and r.images.len == 0 and r.cacheKey().len > 0; + return r.embedding == null and r.prompt_logprobs == null and !r.untaught and r.images.len == 0 and r.audios.len == 0 and r.cacheKey().len > 0; } /// A read finished: the engine thread, perhaps idle with every queued request waiting, looks again. diff --git a/zig/src/core/lane_host_lone.zig b/zig/src/core/lane_host_lone.zig index 9bae71a1f5..0261321122 100644 --- a/zig/src/core/lane_host_lone.zig +++ b/zig/src/core/lane_host_lone.zig @@ -15,7 +15,7 @@ fn takes(h: *LaneHost, job: *Job) bool { const r = job.request; const lone = h.lone orelse return false; const think = lone.min_budget > 0 and (r.think_budget == 0 or r.think_budget > lone.min_budget); - return !((r.sampling != null and !lone.sampled) or !r.drafts or ((r.think_budget > 0 or r.loop_guard) and !think) or r.call != null or (r.structure != null and !lone.structures) or (r.images.len > 0 and !lone.images) or (r.logprobs != null and !lone.logprobs) or r.prompt_logprobs != null); + return !((r.sampling != null and !lone.sampled) or !r.drafts or ((r.think_budget > 0 or r.loop_guard) and !think) or r.call != null or (r.structure != null and !lone.structures) or (r.images.len > 0 and !lone.images) or r.audios.len > 0 or (r.logprobs != null and !lone.logprobs) or r.prompt_logprobs != null); } /// After a round: the lane core's last live stream goes back to the driver when nothing else waits or is due. diff --git a/zig/src/core/lane_host_prefix.zig b/zig/src/core/lane_host_prefix.zig index f0b009bbfd..a5053993aa 100644 --- a/zig/src/core/lane_host_prefix.zig +++ b/zig/src/core/lane_host_prefix.zig @@ -61,7 +61,7 @@ pub fn waits(h: *LaneHost, job: *Job) bool { /// A request whose states the cache keys by its tokens alone, as the host's lookup does (no pictures). fn keyed(r: *const api.Request) bool { - return r.embedding == null and r.prompt_logprobs == null and !r.untaught and r.images.len == 0 and r.cacheKey().len > 0; + return r.embedding == null and r.prompt_logprobs == null and !r.untaught and r.images.len == 0 and r.audios.len == 0 and r.cacheKey().len > 0; } /// Waiting requests go on in queue order: their fill stood at the cut, left the fills, or the bound passed (lock held). diff --git a/zig/src/core/lane_stream.zig b/zig/src/core/lane_stream.zig index d14168f185..f00fd40ba7 100644 --- a/zig/src/core/lane_stream.zig +++ b/zig/src/core/lane_stream.zig @@ -69,6 +69,7 @@ pub fn open(h: *LaneHost, job: *Job) bool { .decide = r.decide, .untaught = r.untaught, .images = r.images, + .audios = r.audios, .key = r.cache_key, }) catch { freeProposer(h, job); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index c8b229bd69..f6e3fde468 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -111,6 +111,14 @@ pub const Image = struct { hash: u64 = 0, // core/picture/identity.zig's: the bytes and the grid, as the prompt cache's key names the image }; +/// One audio clip of a prompt: its 16 kHz mono samples for the `tokens` placeholder rows from `at`. +pub const Audio = struct { + samples: []const f32, + at: u32, + tokens: u32, + hash: u64 = 0, // the clip's samples, as the prompt cache's key names it +}; + pub const Spec = struct { id: []const u8, prompt: []const u32, @@ -143,6 +151,7 @@ pub const Spec = struct { 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 = &.{}, // in prompt order; the backend's tower fills their rows + audios: []const Audio = &.{}, // in prompt order; the backend's audio encoder fills their rows key: []const u32 = &.{}, // the prompt as the cache matches it (image rows by their image); empty: the prompt }; @@ -184,6 +193,7 @@ pub const Stream = struct { decide: ?u32 = null, untaught: bool = false, images: []const Image = &.{}, + audios: []const Audio = &.{}, key: []const u32 = &.{}, cached: u32 = 0, // prompt tokens the backend restored from `reuse` (its prompt pass started there) fill_at: std.atomic.Value(u32) = .init(0), // prompt tokens the pass has prefilled; the engine thread writes @@ -271,6 +281,7 @@ pub const Stream = struct { .grammar = spec.grammar, .decode_spans = spec.decode_spans, .images = spec.images, + .audios = spec.audios, .key = spec.key, }; // under a grammar a forced close ends at the think end: the grammar takes the reply from the next token diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 0fe49c1db8..2f9df3dce4 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -14,15 +14,16 @@ const Function = module.Function; const Stream = stream_.Stream; const Plan = grouped.Plan; const Affine3Experts = @import("experts3.zig").Affine3Experts; +pub const Bf16Experts = @import("experts_bf16.zig").Bf16Experts; -pub const Format = enum { affine4, affine3, fp8g, nvfp4, exl3, iq2xxs }; +pub const Format = enum { affine4, affine3, fp8g, nvfp4, exl3, iq2xxs, bf16 }; /// One layer's experts on the device as each format's make() lays them out; a family stacks its shared ones last. pub const Layer = struct { format: Format, exl3: ?@import("exl3_experts.zig").Layer = null, - up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144]; iq2xxs [E][NI/16][D/256][1056] bytes (ggml.zig) - down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144]; iq2xxs [E][D/16][NI/256][1056] + up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144]; iq2xxs [E][NI/16][D/256][1056] bytes (ggml.zig); bf16 [E][NI/32][D/32][4][32] uint4 + down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144]; iq2xxs [E][D/16][NI/256][1056]; bf16 [E][D/32][NI/32][4][32] up_scale: u64 = 0, // fp8g [E][2][NI/128][D/128] fp32; nvfp4 [E][2] global scales; affine4's are in the blocks down_scale: u64 = 0, // fp8g [E][1][D/128][NI/128]; nvfp4 [E][1] width: usize, // NI @@ -50,6 +51,7 @@ pub const Grouped = struct { fp8g: ?*const Fp8Experts = null, nvfp4: ?*const Nvfp4Experts = null, iq2xxs: ?*const Iq2xxsExperts = null, + bf16: ?*const Bf16Experts = null, /// Decode rows: `items` bounds the plan's items (grouped.maxItems at tile 16); `skip` an expert left out, or -1. pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { @@ -73,6 +75,7 @@ pub const Grouped = struct { .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, false, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, false, epi, in, l, p, out, items, skip), + .bf16 => return (g.bf16 orelse return error.FormatNotLoaded).run(s, false, epi, in, l, p, out, items, skip), } } @@ -98,6 +101,7 @@ pub const Grouped = struct { .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, true, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, true, epi, in, l, p, out, items, skip), + .bf16 => return (g.bf16 orelse return error.FormatNotLoaded).run(s, true, epi, in, l, p, out, items, skip), } } }; @@ -105,7 +109,7 @@ pub const Grouped = struct { /// Pairs a prompt item holds, per format (the Python engines' prompt tiles). pub fn promptTile(f: Format) usize { return switch (f) { - .affine4, .affine3, .fp8g, .nvfp4, .exl3, .iq2xxs => 64, + .affine4, .affine3, .fp8g, .nvfp4, .exl3, .iq2xxs, .bf16 => 64, }; } @@ -577,19 +581,8 @@ test "a packed fp8 expert puts each word where fp8/experts.py's pack does" { try std.testing.expectEqual(words[53 * (k / 4) + 23], out[idx]); } -test "a format that is not loaded, or an option it was not built with, is refused before any launch" { - const g: Grouped = .{}; - const l: Layer = .{ .format = .fp8g, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2 }; - const s: Stream = undefined; // never reached: the refusal comes first - const p: Plan = .{ .members = 0, .items = 0, .counts = 0, .rank = 0, .hist = 0 }; - try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, -1)); - const a4: Affine4Experts = undefined; - const h: Grouped = .{ .affine4 = &a4 }; - var la = l; - la.format = .affine4; - try std.testing.expectError(error.EpilogueNotBuilt, h.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1)); - try std.testing.expectError(error.EpilogueNotBuilt, h.prompt(s, .down_f32, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1)); - try std.testing.expectError(error.OptionNotBuilt, h.decode(s, .up, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 3)); +test { + _ = @import("experts_refuse_test.zig"); } test "NVFP4 decode CTA geometry is explicit, bounded and defaults to four warps" { diff --git a/zig/src/cuda/experts_bf16.zig b/zig/src/cuda/experts_bf16.zig new file mode 100644 index 0000000000..76ed554c2f --- /dev/null +++ b/zig/src/cuda/experts_bf16.zig @@ -0,0 +1,94 @@ +//! bf16 experts (experts_bf16.cu): a checkpoint's bf16 weights in mma fragment blocks, relu^2 up, fp32 or bf16 down. +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const grouped = @import("grouped.zig"); +const driver = @import("driver.zig"); +const experts = @import("experts.zig"); + +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; +const Plan = grouped.Plan; +const Epi = experts.Epi; +const Rows = experts.Rows; +const Layer = experts.Layer; + +pub const Bf16Experts = struct { + mod: Module, + fns: [2][3]Function, // decode, prompt; by Epi + resident: [3]usize, // decode blocks the grid may hold: per SM times SMs + pack_fn: Function, + + pub const cols = 32; // output columns a block + pub const warps = 4; + /// Bytes of a (32 columns, 32 inputs) block. + pub const block = 32 * 32 * 2; + + pub const symbols = struct { + pub const decode = [3][:0]const u8{ "tf_bf16_experts_up", "tf_bf16_experts_down_f32", "tf_bf16_experts_down_bf16" }; + pub const prompt = [3][:0]const u8{ "tf_bf16_experts_prompt_up", "tf_bf16_experts_prompt_down_f32", "tf_bf16_experts_prompt_down_bf16" }; + pub const pack = "tf_bf16_experts_pack"; + }; + + pub fn load(d: *const driver.Driver, sms: usize) !Bf16Experts { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.experts_bf16); + errdefer mod.unload(); + var e: Bf16Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .pack_fn = try mod.function(symbols.pack) }; + for (symbols.decode, symbols.prompt, 0..) |dn, pn, i| { + e.fns[0][i] = try mod.function(dn); + e.fns[1][i] = try mod.function(pn); + e.resident[i] = @max(1, try e.fns[0][i].occupancy(warps * 32, 0)) * sms; + } + return e; + } + + pub fn unload(e: *Bf16Experts) void { + e.mod.unload(); + } + + /// One call on a plan of at most `items` items: 16 pairs an item in decode, 64 in a prompt. + pub fn run(e: *const Bf16Experts, s: Stream, prompt: bool, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { + if (!l.relu2 or l.swiglu or l.limit != 0) return error.OptionNotBuilt; + const up = epi == .up; + const kg = (if (up) l.dims else l.width) / 32; + const nb = (if (up) l.width else l.dims) / cols; + const ei = @backingInt(epi); + const units = items * nb; + const grid = if (prompt) (units + warps - 1) / warps else @min((units + warps - 1) / warps, e.resident[ei]); + if (grid < 1) return; + var a: launch_.Args = .{}; + a.add(in.x); + a.add(@as(i32, @intCast(in.stride))); + a.add(@as(i32, @intCast(if (up) in.slots else 0))); + a.add(if (up) l.up else l.down); + a.add(@as(i32, @intCast(kg))); + a.add(@as(i32, @intCast(nb))); + for ([_]u64{ p.items, p.counts, p.members, out }) |v| a.add(v); + a.add(@as(i32, @intCast(if (up) l.width else l.dims))); + a.add(skip); + try launch_.launch(e.fns[@intFromBool(prompt)][ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = warps * 32, .y = 1, .z = 1 } }, s, &a); + } + + /// Bytes of `count` experts' blocks for n outputs from k inputs. + pub fn bytes(count: usize, n: usize, k: usize) usize { + return count * (n / cols) * (k / 32) * block; + } + + /// A source matrix: experts `stride` apart, rows of `ld`, its first output and input, [k][n] (x @ W) or [n][k]. + pub const Source = struct { stride: usize = 0, ld: usize, n0: usize = 0, k0: usize = 0, in_out: bool }; + + /// `count` bf16 matrices of n outputs from k inputs at `src` into their blocks at `dst` (bytes(count, n, k)). + pub fn pack(e: *const Bf16Experts, s: Stream, src: u64, from: Source, dst: u64, count: usize, n: usize, k: usize) !void { + if (n % cols != 0 or k % 32 != 0) return error.UnexpectedShape; + const per = bytes(1, n, k) / 16; + var a: launch_.Args = .{}; + a.add(src); + a.add(@as(i64, @intCast(from.stride))); + for ([_]usize{ from.ld, from.n0, from.k0, @intFromBool(from.in_out), n, k }) |v| a.add(@as(i32, @intCast(v))); + a.add(dst); + try launch_.launch(e.pack_fn, .{ .grid = .{ .x = @intCast((per + 255) / 256), .y = @intCast(count), .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &a); + } +}; diff --git a/zig/src/cuda/experts_refuse_test.zig b/zig/src/cuda/experts_refuse_test.zig new file mode 100644 index 0000000000..9d496f21ca --- /dev/null +++ b/zig/src/cuda/experts_refuse_test.zig @@ -0,0 +1,23 @@ +//! Grouped experts refuse a format that is not loaded, or an option it was not built with, before any launch. +const std = @import("std"); +const experts = @import("experts.zig"); +const Grouped = experts.Grouped; +const Layer = experts.Layer; +const Plan = @import("grouped.zig").Plan; +const Stream = @import("stream.zig").Stream; +const Affine4Experts = experts.Affine4Experts; + +test "a format that is not loaded, or an option it was not built with, is refused before any launch" { + const g: Grouped = .{}; + const l: Layer = .{ .format = .fp8g, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2 }; + const s: Stream = undefined; // never reached: the refusal comes first + const p: Plan = .{ .members = 0, .items = 0, .counts = 0, .rank = 0, .hist = 0 }; + try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, -1)); + const a4: Affine4Experts = undefined; + const h: Grouped = .{ .affine4 = &a4 }; + var la = l; + la.format = .affine4; + try std.testing.expectError(error.EpilogueNotBuilt, h.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1)); + try std.testing.expectError(error.EpilogueNotBuilt, h.prompt(s, .down_f32, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1)); + try std.testing.expectError(error.OptionNotBuilt, h.decode(s, .up, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 3)); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index db13a74f50..633bbc77ca 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -36,6 +36,7 @@ pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_a pub const lfm2: []const u8 = if (available) &Blob("fatbin_lfm2").bytes else &.{}; pub const vit: []const u8 = if (available) &Blob("fatbin_vit").bytes else &.{}; pub const moss_image: []const u8 = if (available) &Blob("fatbin_moss_image").bytes else &.{}; +pub const nemo_audio: []const u8 = if (available) &Blob("fatbin_nemo_audio").bytes else &.{}; pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes else &.{}; pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; pub const lane_gemv: []const u8 = if (available) &Blob("fatbin_lane_gemv").bytes else &.{}; @@ -46,6 +47,7 @@ pub const fp8_experts: []const u8 = if (available) &Blob("fatbin_fp8_experts").b pub const exlr8: []const u8 = if (available) &Blob("fatbin_exlr8").bytes else &.{}; pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").bytes else &.{}; pub const nvfp4_experts_tile: []const u8 = if (available) &Blob("fatbin_nvfp4_experts_tile").bytes else &.{}; +pub const experts_bf16: []const u8 = if (available) &Blob("fatbin_experts_bf16").bytes else &.{}; pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; pub const rank_sum: []const u8 = if (available) &Blob("fatbin_rank_sum").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; diff --git a/zig/src/families/nemo_audio/audio.zig b/zig/src/families/nemo_audio/audio.zig new file mode 100644 index 0000000000..e8171d8eff --- /dev/null +++ b/zig/src/families/nemo_audio/audio.zig @@ -0,0 +1,289 @@ +//! Transcribe's audio encoder on CUDA as NeMo's vLLM SALM path runs it; n samples give ceil((n / 160) / 8) rows. +const std = @import("std"); +const cuda = @import("cuda"); +pub const Config = @import("config.zig"); +pub const mel = @import("mel.zig"); +pub const sortformer = @import("sortformer.zig"); +pub const Dump = @import("dump.zig").Dump; +pub const kernels = @import("kernels.zig"); +const W = @import("weights.zig"); +const enc = @import("encoder.zig"); +const Allocator = std.mem.Allocator; +const S = Config.speakers; + +test { + _ = Config; + _ = mel; + _ = sortformer; +} + +/// The rows a clip of `samples` gives the prompt. +pub fn rows(samples: usize) usize { + return (samples / Config.hop + Config.sub - 1) / Config.sub; +} + +/// One online window over mel frames: its core [s, e) and the context-extended span [cs, ce). +pub const Window = struct { s: usize, e: usize, cs: usize, ce: usize }; + +pub fn windows(gpa: Allocator, total: usize) ![]Window { + const n = @max(1, (total + Config.window_frames - 1) / Config.window_frames); + const out = try gpa.alloc(Window, n); + for (out, 0..) |*w, i| { + const s = i * Config.window_frames; + const e = @min(s + Config.window_frames, total); + w.* = .{ .s = s, .e = e, .cs = s -| Config.context_frames, .ce = @min(e + Config.context_frames, total) }; + } + return out; +} + +fn ceilDiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +pub const Encoder = struct { + gpa: Allocator, + d: *const cuda.Driver, + c: Config, + w: W.Weights, + r: enc.Runner, + kernel: []u16, // the sinusoid speaker kernel [8, asr d] in bf16 + + /// The audio encoder of the checkpoint at `dir` (config.json's perception block and its perception.* tensors). + pub fn init(gpa: Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8) !*Encoder { + 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 << 24)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const c = try Config.parse(parsed.value.object); + const e = try gpa.create(Encoder); + errdefer gpa.destroy(e); + var w = try W.Weights.load(gpa, io, d, dir, c); + errdefer w.deinit(gpa); + var r = try enc.Runner.init(gpa, d, c); + errdefer r.deinit(); + e.* = .{ .gpa = gpa, .d = d, .c = c, .w = w, .r = r, .kernel = try sinusoid(gpa, c.asr.d) }; + return e; + } + + pub fn deinit(e: *Encoder) void { + e.gpa.free(e.kernel); + e.r.deinit(); + e.w.deinit(e.gpa); + e.gpa.destroy(e); + } + + /// Device bytes the weights and scratch hold. + pub fn bytes(e: *const Encoder) u64 { + return e.w.bytes + e.r.scratch.len; + } + + fn upload(e: *const Encoder, s: cuda.Stream, dst: u64, src: []const u8) !void { + try e.d.check(e.d.api.cuMemcpyHtoDAsync_v2(dst, src.ptr, src.len, s.handle), "cuMemcpyHtoDAsync"); + try s.synchronize(); + } + + fn download(e: *const Encoder, s: cuda.Stream, dst: []u8, src: u64) !void { + try s.synchronize(); + try e.d.check(e.d.api.cuMemcpyDtoH_v2(dst.ptr, src, dst.len), "cuMemcpyDtoH"); + } + + fn copy(e: *const Encoder, s: cuda.Stream, dst: u64, src: u64, len: usize) !void { + try e.d.check(e.d.api.cuMemcpyDtoDAsync_v2(dst, src, len, s.handle), "cuMemcpyDtoDAsync"); + } + + /// FeatureStacking's input for frames [cs, ce) of `feats` [frames, 128] at b.in: rows of 8 frames, zero padded. + fn stacked(e: *const Encoder, s: cuda.Stream, feats: []const u16, cs: usize, ce: usize) !usize { + const n = ceilDiv(ce - cs, Config.sub); + const host = try e.gpa.alloc(u16, n * Config.sub * Config.mels); + defer e.gpa.free(host); + @memset(host, 0); + @memcpy(host[0 .. (ce - cs) * Config.mels], feats[cs * Config.mels .. ce * Config.mels]); + try e.upload(s, e.r.b.in, std.mem.sliceAsBytes(host)); + return n; + } + + /// `wav`'s rows (bf16 [rows(wav.len), text hidden]) at `out`; stage tensors into `dump` when given. + pub fn encode(e: *Encoder, s: cuda.Stream, wav: []const f32, out: u64, dump: ?*Dump) !usize { + const gpa = e.gpa; + const c = e.c; + const b = e.r.b; + var f = try mel.features(gpa, wav, e.w.host.window, e.w.host.fb); + defer f.deinit(gpa); + if (dump) |dm| { + try dm.bits("", "mel", f.mel, f.frames, Config.mels); + try dm.bits("", "mel_asr_norm", f.asr, f.frames, Config.mels); + } + const wins = try windows(gpa, f.frames); + defer gpa.free(wins); + const encoded = ceilDiv(f.valid, Config.sub); + var total_out: usize = 0; + const cores = try gpa.alloc([2]usize, wins.len); // each window's (first core row, core rows) + defer gpa.free(cores); + var room: usize = 0; + for (wins) |w| room += ceilDiv(w.ce - w.cs, Config.sub); + const d = c.asr.d; + // device rows for the whole clip: the ASR encoder's core rows, their norm, the speaker infusion, the fusion + var clip = try cuda.DeviceBuffer.alloc(e.d, 4 * room * d * 2); + defer clip.free(); + const asr_rows = clip.ptr; + const normed = asr_rows + room * d * 2; + const infusion = normed + room * d * 2; + const fused = infusion + room * d * 2; + var tag_buf: [32]u8 = undefined; + // the ASR branch, window by window, its core rows kept + for (wins, 0..) |w, wi| { + const tag = try std.fmt.bufPrint(&tag_buf, "asr.w{d}", .{wi}); + const n = try e.stacked(s, f.asr, w.cs, w.ce); + const len = ceilDiv(std.math.clamp(f.valid -| w.cs, 0, w.ce - w.cs), Config.sub); + try e.r.stackIn(s, &e.w.asr, c.asr, n, b.x); + if (dump) |dm| try dm.device(s, tag, ".pre_encode", b.x, n, d); + try e.r.run(s, &e.w.asr, c.asr, 0, n, len, b.a, dump, tag); + if (dump) |dm| try dm.device(s, tag, ".final_norm", b.a, n, d); + const left = (w.s - w.cs) / Config.sub; + const core = @min(ceilDiv(w.e, Config.sub) - ceilDiv(w.s, Config.sub), n - left); + try e.copy(s, asr_rows + total_out * d * 2, b.a + left * d * 2, core * d * 2); + cores[wi] = .{ total_out, core }; + total_out += core; + } + if (dump) |dm| try dm.device(s, "", "asr_encoded", asr_rows, total_out, d); + // the Sortformer, streaming: each window sees the speaker cache and FIFO before its own rows + const diar = try gpa.alloc(f32, total_out * S); + defer gpa.free(diar); + var st: sortformer.State = .{ .gpa = gpa, .width = c.sf.d }; + defer st.deinit(); + for (wins, 0..) |w, wi| { + const tag = try std.fmt.bufPrint(&tag_buf, "sf.w{d}", .{wi}); + try e.diarize(s, &st, &f, w, wi, wi + 1 == wins.len, cores[wi][1], diar[cores[wi][0] * S ..][0 .. cores[wi][1] * S], dump, tag); + } + if (dump) |dm| try dm.floats("", "sortformer_probs", diar, total_out, S); + // fusion: LayerNorm(asr) + LayerNorm_8(speakers above threshold) @ the sinusoid kernel, then the projection + try e.r.vk.layerNorm(s, asr_rows, e.w.asr_norm, e.w.asr_norm_bias, normed, total_out, d, c.asr.eps); + const inf_host = try gpa.alloc(u16, total_out * d); + defer gpa.free(inf_host); + const feats = try gpa.alloc(f32, total_out * S); + defer gpa.free(feats); + for (0..total_out) |t| { + var x: [S]f32 = undefined; + for (0..S) |k| x[k] = if (diar[t * S + k] > Config.speaker_threshold) 1 else 0; + @memcpy(feats[t * S ..][0..S], &x); + const nt = layerNorm8(x, e.w.host.diar_norm, e.w.host.diar_norm_bias, c.asr.eps); + for (0..d) |col| { + var acc: f32 = 0; + for (0..S) |k| acc += nt[k] * mel.widen(e.kernel[k * d + col]); + inf_host[t * d + col] = mel.bf16(acc); + } + } + if (dump) |dm| { + try dm.floats("", "speaker_features", feats, total_out, S); + try dm.device(s, "", "asr_norm_out", normed, total_out, d); + try dm.bits("", "speaker_infusion", inf_host, total_out, d); + } + try e.upload(s, infusion, std.mem.sliceAsBytes(inf_host)); + try e.r.vk.pointwise(s, infusion, normed, fused, total_out * d, .none); + if (dump) |dm| try dm.device(s, "", "fused", fused, total_out, d); + if (encoded > total_out) return error.AudioRows; + try cuda.vit.linear(&e.r.g, &e.r.vk, s, b.acc, fused, e.w.proj, e.w.proj_bias, out, encoded, c.text_hidden, d); + if (dump) |dm| try dm.device(s, "", "audio_embeddings", out, encoded, c.text_hidden); + return encoded; + } + + /// One Sortformer window: encoder and head over cache, FIFO and window, core predictions to `core_out` [core, 8]. + fn diarize(e: *Encoder, s: cuda.Stream, st: *sortformer.State, f: *const mel.Features, w: Window, wi: usize, last: bool, core: usize, core_out: []f32, dump: ?*Dump, tag: []const u8) !void { + _ = wi; + const gpa = e.gpa; + const c = e.c; + const b = e.r.b; + const k = &e.r.k; + const g = &e.r.g; + const vk = &e.r.vk; + const sd = c.sf.d; + const t = Config.tf_d; + const n = try e.stacked(s, f.mel, w.cs, w.ce); + const pre_len = ceilDiv(std.math.clamp(f.valid -| w.cs, 0, w.ce - w.cs), Config.sub); + try e.r.stackIn(s, &e.w.sf, c.sf, n, b.h); + const pre = try gpa.alloc(u16, n * sd); + defer gpa.free(pre); + try e.download(s, std.mem.sliceAsBytes(pre), b.h); + if (dump) |dm| try dm.bits(tag, ".pre_encode", pre, n, sd); + const spk_n = st.spkRows(); + const fifo_n = st.fifoRows(); + const rows_ = spk_n + fifo_n + n; + const len = spk_n + fifo_n + pre_len; + const cat = try gpa.alloc(u16, rows_ * sd); + defer gpa.free(cat); + @memcpy(cat[0 .. spk_n * sd], st.spkcache.items); + @memcpy(cat[spk_n * sd ..][0 .. fifo_n * sd], st.fifo.items); + @memcpy(cat[(spk_n + fifo_n) * sd ..], pre); + if (dump) |dm| try dm.bits(tag, ".encoder_in", cat, rows_, sd); + try e.upload(s, b.x, std.mem.sliceAsBytes(cat)); + try e.r.run(s, &e.w.sf, c.sf, 1, rows_, len, b.a, dump, tag); + if (dump) |dm| try dm.device(s, tag, ".encoder_out", b.a, rows_, sd); + // the head: encoder_proj, padding zeroed, the subpixel conv (8 frames a row), relu, hidden, relu, speakers + try cuda.vit.linear(g, vk, s, b.acc, b.a, e.w.encoder_proj, e.w.encoder_proj_bias, b.h, rows_, t, sd); + try k.zeroRows(s, b.h, rows_, t, len); + try k.im2col3(s, b.h, b.qkv, rows_, t, len); + try cuda.vit.linear(g, vk, s, b.acc, b.qkv, e.w.upsample, e.w.upsample_bias, b.wide, rows_, Config.sub * t, 3 * t); + const hr = rows_ * Config.sub; + try k.reluInPlace(s, b.wide, hr * t); + try cuda.vit.linear(g, vk, s, b.acc, b.wide, e.w.hidden, e.w.hidden_bias, b.qkv, hr, t, t); + try k.reluInPlace(s, b.qkv, hr * t); + try cuda.vit.linear(g, vk, s, b.acc, b.qkv, e.w.spks, e.w.spks_bias, b.h, hr, S, t); + const logits = try gpa.alloc(u16, hr * S); + defer gpa.free(logits); + try e.download(s, std.mem.sliceAsBytes(logits), b.h); + if (dump) |dm| try dm.bits(tag, ".logits_hr", logits, hr, S); + // sigmoid in bf16, frames past the valid length zeroed + const probs = try gpa.alloc(f32, hr * S); + defer gpa.free(probs); + for (probs, logits, 0..) |*p, l, i| { + const x = mel.widen(l); + p.* = if (i / S < len * Config.sub) mel.rbf(1.0 / (1.0 + @exp(-x))) else 0; + } + if (dump) |dm| try dm.floats(tag, ".probs_hr", probs, hr, S); + const lc = (w.s - w.cs) / Config.sub; + const rc = ceilDiv(w.ce - w.e, Config.sub); + const lr = try sortformer.downsample(gpa, probs); + defer gpa.free(lr); + for (len..lr.len / S) |r| @memset(lr[r * S ..][0..S], 0); + const chunk_len = n - lc - rc; + const start = (spk_n + fifo_n + lc) * Config.sub; + const chunk_lr = try sortformer.downsample(gpa, probs[start * S ..][0 .. chunk_len * Config.sub * S]); + defer gpa.free(chunk_lr); + if (!last) try st.update(pre, lr, lc, rc, e.w.host.sil); + if (dump) |dm| { + try dm.bits(tag, ".spkcache_after", st.spkcache.items, st.spkRows(), sd); + try dm.bits(tag, ".fifo_after", st.fifo.items, st.fifoRows(), sd); + } + // _align_diar_frames: the core rows, the last pooled row repeated when the window gives fewer + const have = chunk_lr.len / S; + for (0..core) |r| @memcpy(core_out[r * S ..][0..S], chunk_lr[@min(r, have - 1) * S ..][0..S]); + } +}; + +/// LayerNorm over the eight speaker features (bf16 0 or 1), fp32 inside and rounded once. +fn layerNorm8(x: [S]f32, w: [S]f32, b: [S]f32, eps: f32) [S]f32 { + var mean: f32 = 0; + for (x) |v| mean += v; + mean /= S; + var v2: f32 = 0; + for (x) |v| v2 += (v - mean) * (v - mean); + const rstd = 1.0 / @sqrt(v2 / S + eps); + var out: [S]f32 = undefined; + for (&out, x, w, b) |*o, v, wi, bi| o.* = mel.rbf((v - mean) * rstd * wi + bi); + return out; +} + +/// _build_sinusoid_position_encoding(8, d) in fp32, cast to bf16 with the module. +fn sinusoid(gpa: Allocator, d: usize) ![]u16 { + const out = try gpa.alloc(u16, S * d); + const step: f32 = @floatCast(-(@log(10000.0) / @as(f64, @floatFromInt(d)))); + for (0..S) |p| for (0..d / 2) |j| { + const div: f32 = @floatCast(@exp(@as(f64, @as(f32, @floatFromInt(2 * j)) * step))); + const a: f32 = @as(f32, @floatFromInt(p)) * div; + out[p * d + 2 * j] = mel.bf16(@floatCast(@sin(@as(f64, a)))); + out[p * d + 2 * j + 1] = mel.bf16(@floatCast(@cos(@as(f64, a)))); + }; + return out; +} diff --git a/zig/src/families/nemo_audio/config.zig b/zig/src/families/nemo_audio/config.zig new file mode 100644 index 0000000000..1b8c5fe580 --- /dev/null +++ b/zig/src/families/nemo_audio/config.zig @@ -0,0 +1,111 @@ +//! Transcribe's audio config: the perception config's ASR encoder, NeMo's defaults for the Sortformer and windows. +const std = @import("std"); + +pub const Encoder = struct { + d: usize, + heads: usize, + layers: usize, + rot_frac: f64, // RoPE on this fraction of each head's channels + base: f64 = 10000.0, + qk_norm: bool, + eps: f32 = 1e-5, + + pub fn dim(e: Encoder) usize { + return e.d / e.heads; + } + + pub fn rot(e: Encoder) usize { + return @intFromFloat(@as(f64, @floatFromInt(e.dim())) * e.rot_frac); + } +}; + +// AudioToMelSpectrogramPreprocessor as configured +pub const sample_rate = 16000; +pub const n_fft = 512; +pub const hop = 160; +pub const win = 400; +pub const mels = 128; +pub const bins = n_fft / 2 + 1; +pub const preemph: f32 = 0.97; +pub const log_guard: f32 = 0x1p-24; +pub const norm_eps: f32 = 1e-5; +/// FeatureStacking: mel frames an encoder frame stacks (both encoders) +pub const sub = 8; + +// ParallelExpertEncoder's online inference: core windows of 500 encoder frames, 50 frames of context a side +pub const window_frames = 500 * sub; +pub const context_frames = 50 * sub; + +// Sortformer: its head's width, its speakers, the streaming caches and the speaker-cache compression constants +pub const tf_d = 192; +pub const speakers = 8; +pub const fifo_len = 40; +pub const spkcache_update_period = 300; +pub const spkcache_len = 188; +pub const sil_frames_per_spk = 1; +pub const pred_score_threshold: f32 = 0.25; +pub const scores_boost_latest: f32 = 0.05; +pub const strong_boost_rate = 0.75; +pub const weak_boost_rate = 1.5; +pub const min_pos_scores_rate = 0.5; +/// speaker_feature_mode thresholded: a speaker is active above this probability +pub const speaker_threshold: f32 = 0.5; + +asr: Encoder, +sf: Encoder = .{ .d = 512, .heads = 8, .layers = 31, .rot_frac = 1.0, .qk_norm = false }, +text_hidden: usize, + +/// config.json's perception block: its encoder must be the TransformerEncoder these kernels run. +pub fn parse(root: std.json.ObjectMap) !@This() { + const p = (root.get("perception") orelse return error.NoPerception).object; + const e = (p.get("encoder") orelse return error.NoPerception).object; + const pre = (p.get("preprocessor") orelse return error.NoPerception).object; + if (int(pre, "features") != mels or int(pre, "n_fft") != n_fft or int(pre, "sample_rate") != sample_rate) return error.UnsupportedAudioConfig; + if (int(e, "subsampling_factor") != sub or int(e, "feat_in") != mels) return error.UnsupportedAudioConfig; + if (!flag(e, "qk_norm") or flag(e, "qkv_bias") or !std.mem.eql(u8, str(e, "self_attention_model"), "rope")) return error.UnsupportedAudioConfig; + const c: @This() = .{ + .asr = .{ .d = int(e, "d_model"), .heads = int(e, "n_heads"), .layers = int(e, "n_layers"), .rot_frac = num(e, "rotary_fraction"), .base = num(e, "rope_base"), .qk_norm = true }, + .text_hidden = int(p, "output_dim"), + }; + if (c.asr.heads == 0 or c.asr.d % c.asr.heads != 0 or int(e, "ff_expansion") != 4) return error.UnsupportedAudioConfig; + return c; +} + +fn int(o: std.json.ObjectMap, k: []const u8) usize { + const v = o.get(k) orelse return 0; + return if (v == .integer and v.integer > 0) @intCast(v.integer) else 0; +} + +fn num(o: std.json.ObjectMap, k: []const u8) f64 { + const v = o.get(k) orelse return 0; + return switch (v) { + .float => v.float, + .integer => @floatFromInt(v.integer), + else => 0, + }; +} + +fn flag(o: std.json.ObjectMap, k: []const u8) bool { + const v = o.get(k) orelse return false; + return v == .bool and v.bool; +} + +fn str(o: std.json.ObjectMap, k: []const u8) []const u8 { + const v = o.get(k) orelse return ""; + return if (v == .string) v.string else ""; +} + +test "the published perception block parses" { + const text = + \\{"perception":{"output_dim":2688,"preprocessor":{"sample_rate":16000,"features":128,"n_fft":512}, + \\"encoder":{"feat_in":128,"d_model":1280,"n_heads":10,"n_layers":32,"qkv_bias":false,"qk_norm":true,"ff_expansion":4, + \\"subsampling_factor":8,"self_attention_model":"rope","rope_base":10000.0,"rotary_fraction":0.5}}} + ; + const parsed = try std.json.parseFromSlice(std.json.Value, std.testing.allocator, text, .{}); + defer parsed.deinit(); + const c = try parse(parsed.value.object); + try std.testing.expectEqual(@as(usize, 128), c.asr.dim()); + try std.testing.expectEqual(@as(usize, 64), c.asr.rot()); + try std.testing.expectEqual(@as(usize, 64), c.sf.rot()); + try std.testing.expectEqual(@as(usize, 2688), c.text_hidden); +} diff --git a/zig/src/families/nemo_audio/dump.zig b/zig/src/families/nemo_audio/dump.zig new file mode 100644 index 0000000000..8b0fa94567 --- /dev/null +++ b/zig/src/families/nemo_audio/dump.zig @@ -0,0 +1,49 @@ +//! Stage tensors an encode pass hands back for checks against a reference (float32, named as the reference names them). +const std = @import("std"); +const cuda = @import("cuda"); +const mel = @import("mel.zig"); +const Allocator = std.mem.Allocator; + +pub const Item = struct { name: []u8, shape: [2]usize, data: []f32 }; + +pub const Dump = struct { + gpa: Allocator, + d: *const cuda.Driver, + items: std.ArrayList(Item) = .empty, + + pub fn deinit(dm: *Dump) void { + for (dm.items.items) |it| { + dm.gpa.free(it.name); + dm.gpa.free(it.data); + } + dm.items.deinit(dm.gpa); + } + + fn add(dm: *Dump, tag: []const u8, name: []const u8, rows: usize, cols: usize) ![]f32 { + const n = try std.fmt.allocPrint(dm.gpa, "{s}{s}", .{ tag, name }); + errdefer dm.gpa.free(n); + const data = try dm.gpa.alloc(f32, rows * cols); + errdefer dm.gpa.free(data); + try dm.items.append(dm.gpa, .{ .name = n, .shape = .{ rows, cols }, .data = data }); + return data; + } + + /// bf16 rows [rows, cols] from the device, after the stream's work so far. + pub fn device(dm: *Dump, s: cuda.Stream, tag: []const u8, name: []const u8, ptr: u64, rows: usize, cols: usize) !void { + const raw = try dm.gpa.alloc(u16, rows * cols); + defer dm.gpa.free(raw); + try s.synchronize(); + try dm.d.check(dm.d.api.cuMemcpyDtoH_v2(raw.ptr, ptr, raw.len * 2), "cuMemcpyDtoH"); + for (try dm.add(tag, name, rows, cols), raw) |*o, v| o.* = mel.widen(v); + } + + /// bf16 bits from the host. + pub fn bits(dm: *Dump, tag: []const u8, name: []const u8, v: []const u16, rows: usize, cols: usize) !void { + for (try dm.add(tag, name, rows, cols), v[0 .. rows * cols]) |*o, x| o.* = mel.widen(x); + } + + /// f32 values from the host. + pub fn floats(dm: *Dump, tag: []const u8, name: []const u8, v: []const f32, rows: usize, cols: usize) !void { + @memcpy(try dm.add(tag, name, rows, cols), v[0 .. rows * cols]); + } +}; diff --git a/zig/src/families/nemo_audio/encoder.zig b/zig/src/families/nemo_audio/encoder.zig new file mode 100644 index 0000000000..ef99d0f7e6 --- /dev/null +++ b/zig/src/families/nemo_audio/encoder.zig @@ -0,0 +1,131 @@ +//! NeMo's TransformerEncoder over one window on CUDA: the ASR encoder and the Sortformer's encoder are two of these. +const std = @import("std"); +const cuda = @import("cuda"); +const Config = @import("config.zig"); +const W = @import("weights.zig"); +const mel = @import("mel.zig"); +const Kernels = @import("kernels.zig").Kernels; +const Dump = @import("dump.zig").Dump; +const vit = cuda.vit; +const Allocator = std.mem.Allocator; + +/// Rows a window holds: the ASR encoder's 600 frames, or the Sortformer's cache, FIFO and window (188 + 40 + 600). +pub const max_rows = Config.spkcache_len + Config.fifo_len + (Config.window_frames + 2 * Config.context_frames) / Config.sub; +/// Columns a buffer holds: the widest model width, and the widest Linear (the ASR FFN, 4 d). +const widest = 1280; + +pub const Runner = struct { + k: Kernels, + vk: vit.Kernels, + g: vit.Gemms, + scratch: cuda.DeviceBuffer, + b: Buffers, + tables: [2]Rope, // the ASR encoder's, then the Sortformer's + + /// Scratch rows, bf16 [rows, widest] unless noted. + pub const Buffers = struct { in: u64, x: u64, h: u64, a: u64, qkv: u64, wide: u64, acc: u64 }; + const Rope = struct { cos: u64, sin: u64, rot: usize }; + + pub fn init(gpa: Allocator, d: *const cuda.Driver, c: Config) !Runner { + if (c.asr.d > widest or c.sf.d > widest) return error.UnsupportedAudioConfig; + var k = try Kernels.load(d); + errdefer k.unload(); + var vk = try vit.Kernels.load(d); + errdefer vk.unload(); + var g = try vit.Gemms.init(gpa, d); + errdefer g.deinit(); + const r = max_rows; + const sizes = [_]usize{ + 2 * r * Config.sub * Config.mels, 2 * r * widest, 2 * r * widest, 2 * r * widest, + 2 * r * 3 * widest, 2 * r * 4 * widest, 4 * r * 4 * widest, + }; + const tables = [2]Config.Encoder{ c.asr, c.sf }; + var total: usize = 0; + for (sizes) |sz| total += std.mem.alignForward(usize, sz, 256); + for (tables) |e| total += 2 * std.mem.alignForward(usize, 2 * r * e.rot(), 256); + var scratch = try cuda.DeviceBuffer.alloc(d, total); + errdefer scratch.free(); + var off: usize = 0; + var ptrs: [sizes.len]u64 = undefined; + for (sizes, &ptrs) |sz, *p| { + p.* = try scratch.at(off); + off += std.mem.alignForward(usize, sz, 256); + } + var ropes: [2]Rope = undefined; + for (tables, &ropes) |e, *rope| { + const cs, const sn = try ropeTables(gpa, e, r); + defer gpa.free(cs); + defer gpa.free(sn); + const bytes = std.mem.alignForward(usize, 2 * r * e.rot(), 256); + rope.* = .{ .cos = try scratch.at(off), .sin = try scratch.at(off + bytes), .rot = e.rot() }; + try scratch.upload(off, std.mem.sliceAsBytes(cs)); + try scratch.upload(off + bytes, std.mem.sliceAsBytes(sn)); + off += 2 * bytes; + } + return .{ .k = k, .vk = vk, .g = g, .scratch = scratch, .b = .{ .in = ptrs[0], .x = ptrs[1], .h = ptrs[2], .a = ptrs[3], .qkv = ptrs[4], .wide = ptrs[5], .acc = ptrs[6] }, .tables = ropes }; + } + + pub fn deinit(r: *Runner) void { + r.scratch.free(); + r.g.deinit(); + r.vk.unload(); + r.k.unload(); + } + + /// FeatureStacking's projection: `rows` stacked frames at b.in [rows, 1024] to `out` [rows, d]. + pub fn stackIn(r: *Runner, s: cuda.Stream, st: *const W.Stack, e: Config.Encoder, rows: usize, out: u64) !void { + try r.g.run(s, r.b.in, st.pre, out, rows, e.d, Config.sub * Config.mels, .bf16); + } + + /// The encoder over `rows` rows at b.x, the first `len` valid, into `out`; `which` 0 is the ASR encoder's RoPE. + pub fn run(r: *Runner, s: cuda.Stream, st: *const W.Stack, e: Config.Encoder, which: usize, rows: usize, len: usize, out: u64, dump: ?*Dump, tag: []const u8) !void { + if (rows > max_rows or len > rows) return error.AudioWindow; + const b = r.b; + const d = e.d; + const g = &r.g; + const vk = &r.vk; + const rope = r.tables[which]; + const scale: f32 = @floatCast(1.0 / @sqrt(@as(f64, @floatFromInt(e.dim())))); + try vk.layerNorm(s, b.x, st.embed_norm, st.embed_norm_bias, b.x, rows, d, e.eps); + if (dump) |dm| try dm.device(s, tag, ".embed_norm", b.x, rows, d); + for (st.layers, 0..) |l, i| { + try vk.layerNorm(s, b.x, l.norm1, l.norm1_bias, b.h, rows, d, e.eps); + try g.run(s, b.h, l.qkv, b.qkv, rows, 3 * d, d, .bf16); + const norm: ?[4]u64 = if (e.qk_norm) .{ l.q_norm, l.q_norm_bias, l.k_norm, l.k_norm_bias } else null; + try r.k.qkRope(s, b.qkv, norm, rope.cos, rope.sin, rows, e.heads, e.dim(), rope.rot, e.eps); + try r.k.attend(s, b.qkv, b.a, rows, e.heads, e.dim(), len, scale); + try vit.linearAct(g, vk, s, b.acc, b.a, l.out, l.out_bias, b.h, rows, d, d, b.x, b.x, .none); + try vk.layerNorm(s, b.x, l.norm2, l.norm2_bias, b.h, rows, d, e.eps); + try vit.linearAct(g, vk, s, b.acc, b.h, l.fc1, l.fc1_bias, b.wide, rows, 4 * d, d, null, b.wide, .gelu); + try vit.linearAct(g, vk, s, b.acc, b.wide, l.fc2, l.fc2_bias, b.h, rows, d, 4 * d, b.x, b.x, .none); + if (dump) |dm| if (i == 0) try dm.device(s, tag, ".layer0", b.x, rows, d); + } + try vk.layerNorm(s, b.x, st.final, st.final_bias, out, rows, d, e.eps); + } +}; + +/// RotaryPositionalEncoding.create_pe's tables for positions [0, rows): fp32 frequencies cast to bf16, [rows, rot]. +fn ropeTables(gpa: Allocator, e: Config.Encoder, rows: usize) !struct { []u16, []u16 } { + const rot = e.rot(); + const cs = try gpa.alloc(u16, rows * rot); + errdefer gpa.free(cs); + const sn = try gpa.alloc(u16, rows * rot); + errdefer gpa.free(sn); + const half = rot / 2; + const base: f32 = @floatCast(e.base); + for (0..half) |j| { + // inv_freq = 1 / base^(2j / rot), each op in fp32 + const ex = @as(f32, @floatFromInt(2 * j)) / @as(f32, @floatFromInt(rot)); + const inv: f32 = 1.0 / @as(f32, @floatCast(std.math.pow(f64, base, ex))); + for (0..rows) |p| { + const f: f32 = @as(f32, @floatFromInt(p)) * inv; + const c = mel.bf16(@floatCast(@cos(@as(f64, f)))); + const si = mel.bf16(@floatCast(@sin(@as(f64, f)))); + for ([2]usize{ j, j + half }) |col| { + cs[p * rot + col] = c; + sn[p * rot + col] = si; + } + } + } + return .{ cs, sn }; +} diff --git a/zig/src/families/nemo_audio/kernels.zig b/zig/src/families/nemo_audio/kernels.zig new file mode 100644 index 0000000000..9f6ff9af79 --- /dev/null +++ b/zig/src/families/nemo_audio/kernels.zig @@ -0,0 +1,100 @@ +//! nemo_audio.cu's launches: q/k LayerNorm with partial RoPE, attention over a window's keys, the Sortformer glue. +const std = @import("std"); +const cuda = @import("cuda"); + +const Args = cuda.Args; +const Stream = cuda.Stream; + +/// The widest head (nemo_audio.cu's MAXD) and the attention block's threads. +pub const max_dim = 128; +const attention_threads = 128; +/// Keys one attention block holds scores for: a window plus the speaker cache and FIFO, with room. +pub const max_keys = 2048; + +pub const Kernels = struct { + mod: cuda.Module, + qk_rope: cuda.Function, + attention: cuda.Function, + im2col: cuda.Function, + relu: cuda.Function, + zero_rows: cuda.Function, + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(d, cuda.kernels.nemo_audio); + errdefer mod.unload(); + const k: Kernels = .{ + .mod = mod, + .qk_rope = try mod.function("tf_na_qk_rope"), + .attention = try mod.function("tf_na_attention"), + .im2col = try mod.function("tf_na_im2col3"), + .relu = try mod.function("tf_na_relu"), + .zero_rows = try mod.function("tf_na_zero_rows"), + }; + try k.attention.allowDynamicShared(max_keys * 4); + return k; + } + + pub fn unload(k: *Kernels) void { + k.mod.unload(); + } + + /// qkv [rows, 3, heads, dim] in place: q and k LayerNormed when `norm`, then RoPE on the first `rot` channels. + pub fn qkRope(k: *const Kernels, s: Stream, qkv: u64, norm: ?[4]u64, cs: u64, sn: u64, rows: usize, heads: usize, dim: usize, rot: usize, eps: f32) !void { + if (dim > max_dim or rot > dim or rot % 2 != 0) return error.UnsupportedShape; + if (rows == 0) return; + var a: Args = .{}; + a.add(qkv); + for (norm orelse .{ 0, 0, 0, 0 }) |p| a.add(p); + a.add(cs); + a.add(sn); + for ([_]usize{ heads, dim, rot }) |n| a.add(int(n)); + a.add(eps); + try cuda.launch.launch(k.qk_rope, .{ .grid = .{ .x = @intCast(rows * heads * 2) }, .block = .{ .x = 32 } }, s, &a); + } + + /// out [rows, heads, dim] = softmax(q k^T scale) v over keys [0, len) of qkv [rows, 3, heads, dim]. + pub fn attend(k: *const Kernels, s: Stream, qkv: u64, out: u64, rows: usize, heads: usize, dim: usize, len: usize, scale: f32) !void { + if (dim > max_dim or len > rows or rows > max_keys) return error.UnsupportedShape; + if (rows == 0) return; + var a: Args = .{}; + a.add(qkv); + a.add(out); + for ([_]usize{ heads, dim, len }) |n| a.add(int(n)); + a.add(scale); + try cuda.launch.launch(k.attention, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(heads) }, .block = .{ .x = attention_threads }, .shared = @intCast(rows * 4) }, s, &a); + } + + /// y [rows, 3 c] from x [rows, c]: the three taps of a kernel-3, pad-1 conv, rows from `len` on read as zero. + pub fn im2col3(k: *const Kernels, s: Stream, x: u64, y: u64, rows: usize, c: usize, len: usize) !void { + var a: Args = .{}; + a.add(x); + a.add(y); + for ([_]usize{ rows, c, len }) |n| a.add(int(n)); + try cuda.launch.launch(k.im2col, .{ .grid = .{ .x = @intCast(blocks(rows * 3 * c)) }, .block = .{ .x = 256 } }, s, &a); + } + + pub fn reluInPlace(k: *const Kernels, s: Stream, x: u64, count: usize) !void { + var a: Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(count))); + try cuda.launch.launch(k.relu, .{ .grid = .{ .x = @intCast(blocks(count)) }, .block = .{ .x = 256 } }, s, &a); + } + + /// Rows [len, rows) of x [rows, c] set to zero. + pub fn zeroRows(k: *const Kernels, s: Stream, x: u64, rows: usize, c: usize, len: usize) !void { + if (len >= rows) return; + var a: Args = .{}; + a.add(x); + for ([_]usize{ rows, c, len }) |n| a.add(int(n)); + try cuda.launch.launch(k.zero_rows, .{ .grid = .{ .x = @intCast(blocks((rows - len) * c)) }, .block = .{ .x = 256 } }, s, &a); + } +}; + +fn blocks(n: usize) usize { + return @max(1, (n + 255) / 256); +} + +fn int(x: usize) i32 { + return @intCast(x); +} diff --git a/zig/src/families/nemo_audio/mel.zig b/zig/src/families/nemo_audio/mel.zig new file mode 100644 index 0000000000..2bb175a010 --- /dev/null +++ b/zig/src/families/nemo_audio/mel.zig @@ -0,0 +1,211 @@ +//! NeMo's log-mel frontend on the host (FilterbankFeatures in eval), rounding where torch's bf16 path does. +const std = @import("std"); +const Config = @import("config.zig"); +const Allocator = std.mem.Allocator; + +/// bf16 bits of x, round to nearest even (x finite or infinite). +pub fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + if (std.math.isNan(x)) return @intCast((u >> 16) | 0x40); + return @intCast((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); +} + +pub fn widen(v: u16) f32 { + return @bitCast(@as(u32, v) << 16); +} + +/// x rounded to bf16 and back. +pub fn rbf(x: f32) f32 { + return widen(bf16(x)); +} + +/// A 16 kHz mono WAV's samples in [-1, 1): 16-bit PCM over 32768, as soundfile reads it, or 32-bit float. +pub fn decodeWav(gpa: Allocator, bytes: []const u8) ![]f32 { + if (bytes.len < 12 or !std.mem.eql(u8, bytes[0..4], "RIFF") or !std.mem.eql(u8, bytes[8..12], "WAVE")) return error.NotWav; + var at: usize = 12; + var format: u16 = 0; + var bits: u16 = 0; + while (at + 8 <= bytes.len) { + const id = bytes[at..][0..4]; + const len = std.mem.readInt(u32, bytes[at + 4 ..][0..4], .little); + const body = bytes[at + 8 ..][0..@min(len, bytes.len - at - 8)]; + if (std.mem.eql(u8, id, "fmt ")) { + if (body.len < 16) return error.BadWav; + format = std.mem.readInt(u16, body[0..2], .little); + const channels = std.mem.readInt(u16, body[2..4], .little); + const rate = std.mem.readInt(u32, body[4..8], .little); + bits = std.mem.readInt(u16, body[14..16], .little); + if (format == 0xfffe and body.len >= 26) format = std.mem.readInt(u16, body[24..26], .little); // WAVE_FORMAT_EXTENSIBLE + if (channels != 1) return error.WavNotMono; + if (rate != Config.sample_rate) return error.WavRate; + } else if (std.mem.eql(u8, id, "data")) { + if (format == 1 and bits == 16) { + const out = try gpa.alloc(f32, body.len / 2); + for (out, 0..) |*o, i| o.* = @as(f32, @floatFromInt(std.mem.readInt(i16, body[2 * i ..][0..2], .little))) / 32768.0; + return out; + } + if (format == 3 and bits == 32) { + const out = try gpa.alloc(f32, body.len / 4); + for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, body[4 * i ..][0..4], .little)); + return out; + } + return error.WavFormat; + } + at += 8 + len + (len & 1); + } + return error.BadWav; +} + +pub const Features = struct { + frames: usize, // 1 + samples / hop (center padding), the last past the valid length + valid: usize, // samples / hop: frames from here on are zero + mel: []u16, // [frames, 128] bf16: the Sortformer reads these as they are + asr: []u16, // [frames, 128] bf16: normalized per feature over the valid frames, for the ASR encoder + + pub fn deinit(f: *Features, gpa: Allocator) void { + gpa.free(f.mel); + gpa.free(f.asr); + } +}; + +/// Log-mel features of `wav` with the checkpoint's window [400] and filterbank [128, 257] (bf16 values as f32). +pub fn features(gpa: Allocator, wav: []const f32, window: []const f32, fb: []const f32) !Features { + if (window.len != Config.win or fb.len != Config.mels * Config.bins) return error.BadFilterbank; + if (wav.len < Config.hop) return error.AudioTooShort; + const n = wav.len; + const frames = n / Config.hop + 1; + const valid = n / Config.hop; + // preemphasis in fp32: x[i] - 0.97 x[i - 1] + const x = try gpa.alloc(f32, n); + defer gpa.free(x); + x[0] = wav[0]; + for (1..n) |i| x[i] = wav[i] - Config.preemph * wav[i - 1]; + // the 400-sample window centred in 512, as torch.stft pads a shorter window + var w: [Config.n_fft]f32 = @splat(0); + const left = (Config.n_fft - Config.win) / 2; + @memcpy(w[left..][0..Config.win], window); + var fft = Fft.init(); + const mel = try gpa.alloc(u16, frames * Config.mels); + errdefer gpa.free(mel); + var re: [Config.n_fft]f64 = undefined; + var im: [Config.n_fft]f64 = undefined; + var power: [Config.bins]f32 = undefined; + for (0..frames) |t| { + // frame t covers samples [t hop - 256, t hop + 256) of the zero-padded signal + for (0..Config.n_fft, &re, &im) |i, *r, *m| { + const s = @as(isize, @intCast(t * Config.hop + i)) - Config.n_fft / 2; + const v: f32 = if (s >= 0 and s < n) x[@intCast(s)] else 0; + r.* = @floatCast(v * w[i]); + m.* = 0; + } + fft.run(&re, &im); + for (&power, 0..) |*p, b| { + const fr: f32 = @floatCast(re[b]); + const fi: f32 = @floatCast(im[b]); + const mag = @sqrt(fr * fr + fi * fi); + p.* = mag * mag; + } + for (0..Config.mels) |m| { + var acc: f64 = 0; + for (fb[m * Config.bins ..][0..Config.bins], power) |f, p| acc += @as(f64, f) * p; + const v: f32 = if (t < valid) @log(@as(f32, @floatCast(acc)) + Config.log_guard) else 0; + mel[t * Config.mels + m] = bf16(v); + } + } + const asr = try gpa.alloc(u16, frames * Config.mels); + errdefer gpa.free(asr); + normalize(mel, asr, frames, valid); + return .{ .frames = frames, .valid = valid, .mel = mel, .asr = asr }; +} + +/// normalize_batch(per_feature) on bf16 features as torch rounds it (variance and std in fp32); later frames zero. +fn normalize(mel: []const u16, out: []u16, frames: usize, valid: usize) void { + @memset(out, 0); + if (valid == 0) return; + const den: f32 = @floatFromInt(valid); + for (0..Config.mels) |c| { + const ref = widen(mel[c]); + var sum: f64 = 0; + for (0..valid) |t| sum += rbf(widen(mel[t * Config.mels + c]) - ref); + const mean = rbf(ref + rbf(rbf(@floatCast(sum)) / den)); + var sq: f64 = 0; + for (0..valid) |t| { + const d = rbf(widen(mel[t * Config.mels + c]) - mean); + sq += rbf(d * d); + } + var std_ = @sqrt(rbf(@floatCast(sq)) / (den - 1.0)); + if (std.math.isNan(std_)) std_ = 0; + std_ += Config.norm_eps; + for (0..@min(valid, frames)) |t| out[t * Config.mels + c] = bf16(rbf(widen(mel[t * Config.mels + c]) - mean) / std_); + } +} + +/// A 512-point complex radix-2 FFT in f64. +const Fft = struct { + cos: [Config.n_fft / 2]f64, + sin: [Config.n_fft / 2]f64, + + fn init() Fft { + var f: Fft = undefined; + for (0..Config.n_fft / 2) |k| { + const a = -2.0 * std.math.pi * @as(f64, @floatFromInt(k)) / Config.n_fft; + f.cos[k] = @cos(a); + f.sin[k] = @sin(a); + } + return f; + } + + fn run(f: *const Fft, re: *[Config.n_fft]f64, im: *[Config.n_fft]f64) void { + const n = Config.n_fft; + const bits = std.math.log2_int(usize, n); + for (0..n) |i| { + const j = @bitReverse(@as(u9, @intCast(i))) >> @intCast(9 - bits); + if (j > i) { + std.mem.swap(f64, &re[i], &re[j]); + std.mem.swap(f64, &im[i], &im[j]); + } + } + var len: usize = 2; + while (len <= n) : (len *= 2) { + const step = n / len; + var i: usize = 0; + while (i < n) : (i += len) for (0..len / 2) |k| { + const wr = f.cos[k * step]; + const wi = f.sin[k * step]; + const a = i + k; + const b = a + len / 2; + const tr = re[b] * wr - im[b] * wi; + const ti = re[b] * wi + im[b] * wr; + re[b] = re[a] - tr; + im[b] = im[a] - ti; + re[a] += tr; + im[a] += ti; + }; + } + } +}; + +test "the FFT of an impulse is flat and a cosine lands in its bin" { + var f = Fft.init(); + var re: [Config.n_fft]f64 = @splat(0); + var im: [Config.n_fft]f64 = @splat(0); + re[0] = 1; + f.run(&re, &im); + for (re, im) |r, i| { + try std.testing.expectApproxEqAbs(@as(f64, 1), r, 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 0), i, 1e-12); + } + for (&re, &im, 0..) |*r, *i, t| { + r.* = @cos(2.0 * std.math.pi * 5.0 * @as(f64, @floatFromInt(t)) / Config.n_fft); + i.* = 0; + } + f.run(&re, &im); + try std.testing.expectApproxEqAbs(@as(f64, Config.n_fft / 2), re[5], 1e-9); + try std.testing.expectApproxEqAbs(@as(f64, 0), re[6], 1e-9); +} + +test "bf16 rounds to nearest even" { + try std.testing.expectEqual(@as(u16, 0x3f80), bf16(1.0)); + try std.testing.expectEqual(@as(u16, 0x3f80), bf16(@bitCast(@as(u32, 0x3f808000)))); // a tie stays even + try std.testing.expectEqual(@as(u16, 0x3f82), bf16(@bitCast(@as(u32, 0x3f818000)))); // a tie rounds up to even +} diff --git a/zig/src/families/nemo_audio/sortformer.zig b/zig/src/families/nemo_audio/sortformer.zig new file mode 100644 index 0000000000..f40d62f2c6 --- /dev/null +++ b/zig/src/families/nemo_audio/sortformer.zig @@ -0,0 +1,222 @@ +//! The Sortformer's streaming speaker cache and FIFO on the host, in bf16 as SortformerModules keeps them in eval. +const std = @import("std"); +const Config = @import("config.zig"); +const mel = @import("mel.zig"); +const Allocator = std.mem.Allocator; +const rbf = mel.rbf; +const S = Config.speakers; + +pub const State = struct { + gpa: Allocator, + width: usize, // the Sortformer's width (512) + spkcache: std.ArrayList(u16) = .empty, // [rows, width] bf16 + fifo: std.ArrayList(u16) = .empty, + spkcache_preds: std.ArrayList(f32) = .empty, // [rows, 8], bf16 values + fifo_preds: std.ArrayList(f32) = .empty, + compressed: bool = false, + + pub fn deinit(st: *State) void { + st.spkcache.deinit(st.gpa); + st.fifo.deinit(st.gpa); + st.spkcache_preds.deinit(st.gpa); + st.fifo_preds.deinit(st.gpa); + } + + pub fn spkRows(st: *const State) usize { + return st.spkcache.items.len / st.width; + } + + pub fn fifoRows(st: *const State) usize { + return st.fifo.items.len / st.width; + } + + /// streaming_update with a window's rows `chunk`, its predictions `preds` and its context rows `lc` and `rc`. + pub fn update(st: *State, chunk: []const u16, preds: []const f32, lc: usize, rc: usize, sil: []const u16) !void { + const w = st.width; + const spk_n = st.spkRows(); + const fifo_n = st.fifoRows(); + const n = chunk.len / w; + const chunk_len = n - lc - rc; + try st.fifo_preds.resize(st.gpa, 0); + try st.fifo_preds.appendSlice(st.gpa, preds[spk_n * S ..][0 .. fifo_n * S]); + try st.fifo.appendSlice(st.gpa, chunk[lc * w ..][0 .. chunk_len * w]); + try st.fifo_preds.appendSlice(st.gpa, preds[(spk_n + fifo_n + lc) * S ..][0 .. chunk_len * S]); + if (fifo_n + chunk_len <= Config.fifo_len) return; + var pop: usize = Config.spkcache_update_period; + pop = @max(pop, (chunk_len + fifo_n) -| Config.fifo_len); + pop = @min(pop, fifo_n + chunk_len); + try st.spkcache.appendSlice(st.gpa, st.fifo.items[0 .. pop * w]); + if (st.compressed) { + try st.spkcache_preds.appendSlice(st.gpa, st.fifo_preds.items[0 .. pop * S]); + } else { + try st.spkcache_preds.resize(st.gpa, 0); + try st.spkcache_preds.appendSlice(st.gpa, preds[0 .. spk_n * S]); + try st.spkcache_preds.appendSlice(st.gpa, st.fifo_preds.items[0 .. pop * S]); + } + try dropFront(u16, st.gpa, &st.fifo, pop * w); + try dropFront(f32, st.gpa, &st.fifo_preds, pop * S); + if (st.spkRows() > Config.spkcache_len) { + try st.compress(sil); + st.compressed = true; + } + } + + /// _compress_spkcache: the cache cut to spkcache_len rows, the best-scored frames of each speaker first. + fn compress(st: *State, sil: []const u16) !void { + const gpa = st.gpa; + const w = st.width; + const n = st.spkRows(); + const preds = st.spkcache_preds.items; + const per_spk = Config.spkcache_len / S - Config.sil_frames_per_spk; + const strong: usize = @intFromFloat(@floor(@as(f64, @floatFromInt(per_spk)) * Config.strong_boost_rate)); + const weak: usize = @intFromFloat(@floor(@as(f64, @floatFromInt(per_spk)) * Config.weak_boost_rate)); + const min_pos: usize = @intFromFloat(@floor(@as(f64, @floatFromInt(per_spk)) * Config.min_pos_scores_rate)); + const nf = n + Config.sil_frames_per_spk; + // scores [nf, 8]: log p - log(1 - p) + sum over speakers of log(1 - p) - log 0.5, each op rounded to bf16 + const scores = try gpa.alloc(f32, nf * S); + defer gpa.free(scores); + const log_half: f32 = @floatCast(@log(0.5)); + const inf = std.math.inf(f32); + for (0..n) |f| { + var l1: [S]f32 = undefined; + var sum: f32 = 0; + for (0..S) |s| { + l1[s] = rbf(@log(@max(rbf(1.0 - preds[f * S + s]), Config.pred_score_threshold))); + sum += l1[s]; + } + sum = rbf(sum); + for (0..S) |s| { + const p = preds[f * S + s]; + const lp = rbf(@log(@max(p, Config.pred_score_threshold))); + const v = rbf(rbf(rbf(lp - l1[s]) + sum) - log_half); + scores[f * S + s] = if (p > 0.5) v else -inf; + } + } + // _disable_low_scores: a speaker with min_pos positive frames loses its speech frames that are not positive + for (0..S) |s| { + var pos: usize = 0; + for (0..n) |f| pos += @intFromBool(scores[f * S + s] > 0); + if (pos < min_pos) continue; + for (0..n) |f| if (!(scores[f * S + s] > 0) and preds[f * S + s] > 0.5) { + scores[f * S + s] = -inf; + }; + } + if (Config.scores_boost_latest > 0) for (Config.spkcache_len..n) |f| for (0..S) |s| { + scores[f * S + s] = rbf(scores[f * S + s] + Config.scores_boost_latest); + }; + try boost(gpa, scores, n, strong, -2.0 * log_half); + try boost(gpa, scores, n, weak, -1.0 * log_half); + for (0..S) |s| scores[n * S + s] = inf; // the silence frame a speaker + // _get_topk_indices: the spkcache_len best of every (speaker, frame), sorted speaker-major then by frame + const flat = try gpa.alloc(usize, S * nf); + defer gpa.free(flat); + for (flat, 0..) |*o, i| o.* = i; + const Key = struct { + v: []const f32, + nf: usize, + fn at(k: @This(), i: usize) f32 { + return k.v[(i % k.nf) * S + i / k.nf]; + } + fn more(k: @This(), a: usize, b: usize) bool { + const x = k.at(a); + const y = k.at(b); + return x > y or (x == y and a < b); + } + }; + const key: Key = .{ .v = scores, .nf = nf }; + std.mem.sort(usize, flat, key, Key.more); + const take = flat[0..Config.spkcache_len]; + for (take) |*i| if (key.at(i.*) == -inf) { + i.* = std.math.maxInt(usize); + }; + std.mem.sort(usize, take, {}, std.sort.asc(usize)); + var emb = try gpa.alloc(u16, Config.spkcache_len * w); + defer gpa.free(emb); + var gp = try gpa.alloc(f32, Config.spkcache_len * S); + defer gpa.free(gp); + for (take, 0..) |i, r| { + const frame = if (i == std.math.maxInt(usize)) nf else i % nf; + if (frame >= n) { + @memcpy(emb[r * w ..][0..w], sil); + @memset(gp[r * S ..][0..S], 0); + } else { + @memcpy(emb[r * w ..][0..w], st.spkcache.items[frame * w ..][0..w]); + @memcpy(gp[r * S ..][0..S], preds[frame * S ..][0..S]); + } + } + try st.spkcache.resize(gpa, 0); + try st.spkcache.appendSlice(gpa, emb); + try st.spkcache_preds.resize(gpa, 0); + try st.spkcache_preds.appendSlice(gpa, gp); + } +}; + +/// _boost_topk_scores: each speaker's `k` best frames of [0, n) raised by `by`. +fn boost(gpa: Allocator, scores: []f32, n: usize, k: usize, by: f32) !void { + const idx = try gpa.alloc(usize, n); + defer gpa.free(idx); + for (0..S) |s| { + for (idx, 0..) |*o, i| o.* = i; + const Key = struct { + v: []const f32, + s: usize, + fn more(c: @This(), a: usize, b: usize) bool { + const x = c.v[a * S + c.s]; + const y = c.v[b * S + c.s]; + return x > y or (x == y and a < b); + } + }; + std.mem.sort(usize, idx, Key{ .v = scores, .s = s }, Key.more); + for (idx[0..@min(k, n)]) |f| scores[f * S + s] = rbf(scores[f * S + s] + by); + } +} + +fn dropFront(comptime T: type, gpa: Allocator, list: *std.ArrayList(T), n: usize) !void { + const items = list.items; + std.mem.copyForwards(T, items[0 .. items.len - n], items[n..]); + try list.resize(gpa, items.len - n); +} + +/// avg_pool1d over frames (kernel = stride = 8, ceil mode, padding not counted) of bf16 predictions [frames, 8]. +pub fn downsample(gpa: Allocator, hr: []const f32) ![]f32 { + const frames = hr.len / S; + const out_n = (frames + Config.sub - 1) / Config.sub; + const out = try gpa.alloc(f32, out_n * S); + for (0..out_n) |o| { + const lo = o * Config.sub; + const hi = @min(lo + Config.sub, frames); + for (0..S) |s| { + var acc: f32 = 0; + for (lo..hi) |f| acc += hr[f * S + s]; + out[o * S + s] = rbf(acc / @as(f32, @floatFromInt(hi - lo))); + } + } + return out; +} + +test "a FIFO past its length pops into the cache, and the cache compresses to its length" { + const gpa = std.testing.allocator; + var st: State = .{ .gpa = gpa, .width = 2 }; + defer st.deinit(); + const n = 500; + const chunk = try gpa.alloc(u16, n * 2); + defer gpa.free(chunk); + for (chunk, 0..) |*c, i| c.* = @intCast(i); + const preds = try gpa.alloc(f32, n * S); + defer gpa.free(preds); + for (preds, 0..) |*p, i| p.* = rbf(@as(f32, @floatFromInt(i % 97)) / 97.0); + try st.update(chunk, preds, 0, 0, &.{ 7, 7 }); + try std.testing.expectEqual(@as(usize, Config.fifo_len), st.fifoRows()); + try std.testing.expectEqual(@as(usize, Config.spkcache_len), st.spkRows()); + try std.testing.expect(st.compressed); +} + +test "pooling averages eight frames and the short last one by its own count" { + const gpa = std.testing.allocator; + var hr: [10 * S]f32 = @splat(0); + for (0..10) |f| hr[f * S] = @floatFromInt(f); + const lr = try downsample(gpa, &hr); + defer gpa.free(lr); + try std.testing.expectEqual(@as(f32, 3.5), lr[0]); + try std.testing.expectEqual(@as(f32, 8.5), lr[S]); +} diff --git a/zig/src/families/nemo_audio/weights.zig b/zig/src/families/nemo_audio/weights.zig new file mode 100644 index 0000000000..df1d2038d3 --- /dev/null +++ b/zig/src/families/nemo_audio/weights.zig @@ -0,0 +1,264 @@ +//! Transcribe's perception tensors as stored, in one device allocation; the few the host reads are kept there too. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const Config = @import("config.zig"); +const Allocator = std.mem.Allocator; +const st = core.safetensors; + +const asr = "perception.encoder.asr_encoder."; +const sf = "perception.encoder.diarization_model."; +const heads_ = sf ++ "sortformer_modules."; + +pub const Layer = struct { + norm1: u64, + norm1_bias: u64, + qkv: u64, // [3 d, d], no bias: q rows, then k, then v, each head by head + q_norm: u64 = 0, // [head dim] with its bias, when the stack has qk_norm + q_norm_bias: u64 = 0, + k_norm: u64 = 0, + k_norm_bias: u64 = 0, + out: u64, + out_bias: u64, + norm2: u64, + norm2_bias: u64, + fc1: u64, // [4 d, d] + fc1_bias: u64, + fc2: u64, // [d, 4 d] + fc2_bias: u64, +}; + +/// One TransformerEncoder: the stacking projection, the embed norm, its layers and the final norm. +pub const Stack = struct { + pre: u64, // [d, 8 * 128], no bias + embed_norm: u64, + embed_norm_bias: u64, + final: u64, + final_bias: u64, + layers: []Layer, +}; + +pub const Weights = struct { + memory: cuda.DeviceBuffer, + asr: Stack, + sf: Stack, + encoder_proj: u64, // [192, 512] + encoder_proj_bias: u64, + upsample: u64, // the subpixel conv [1536, 3 * 192]: tap-major (stored [1536, 192, 3]) + upsample_bias: u64, + hidden: u64, // first_hidden_to_hidden [192, 192] + hidden_bias: u64, + spks: u64, // single_hidden_to_spks [8, 192] + spks_bias: u64, + asr_norm: u64, + asr_norm_bias: u64, + proj: u64, // [text hidden, 1280] + proj_bias: u64, + host: Host, + bytes: u64, + + pub fn load(gpa: Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8, c: Config) !Weights { + const paths = try core.checkpoint.shardFiles(gpa, io, dir); + defer core.checkpoint.freeShardFiles(gpa, paths); + var files: std.ArrayList(st.File) = .empty; + defer { + for (files.items) |*f| { + f.release(); + f.close(io); + } + files.deinit(gpa); + } + for (paths) |path| { + var f = try st.File.openPrefix(gpa, io, path, "perception."); + if (f.names.count() == 0) { + f.close(io); + continue; + } + files.append(gpa, f) catch |err| { + f.close(io); + return err; + }; + } + if (files.items.len == 0) return error.NoAudioEncoder; + var p: Planner = .{ .gpa = gpa, .files = files.items }; + defer { + for (p.pieces.items) |piece| if (piece.owned) gpa.free(piece.bytes); + p.pieces.deinit(gpa); + } + var w: Weights = undefined; + w.asr.layers = try gpa.alloc(Layer, c.asr.layers); + errdefer gpa.free(w.asr.layers); + w.sf.layers = try gpa.alloc(Layer, c.sf.layers); + errdefer gpa.free(w.sf.layers); + w.host = try Host.read(gpa, &p, c); + errdefer w.host.deinit(gpa); + try p.stack(&w.asr, asr, c.asr, true); + try p.stack(&w.sf, sf ++ "encoder.", c.sf, false); + const t = Config.tf_d; + w.encoder_proj = try p.take(heads_ ++ "encoder_proj.weight", &.{ t, c.sf.d }); + w.encoder_proj_bias = try p.take(heads_ ++ "encoder_proj.bias", &.{t}); + w.upsample = try p.taps(heads_ ++ "subpixel_upsample.weight", Config.sub * t, t); + w.upsample_bias = try p.take(heads_ ++ "subpixel_upsample.bias", &.{Config.sub * t}); + w.hidden = try p.take(heads_ ++ "first_hidden_to_hidden.weight", &.{ t, t }); + w.hidden_bias = try p.take(heads_ ++ "first_hidden_to_hidden.bias", &.{t}); + w.spks = try p.take(heads_ ++ "single_hidden_to_spks.weight", &.{ Config.speakers, t }); + w.spks_bias = try p.take(heads_ ++ "single_hidden_to_spks.bias", &.{Config.speakers}); + w.asr_norm = try p.take("perception.encoder.asr_norm.weight", &.{c.asr.d}); + w.asr_norm_bias = try p.take("perception.encoder.asr_norm.bias", &.{c.asr.d}); + w.proj = try p.take("perception.proj.weight", &.{ c.text_hidden, c.asr.d }); + w.proj_bias = try p.take("perception.proj.bias", &.{c.text_hidden}); + var memory = try cuda.DeviceBuffer.alloc(d, p.bytes); + errdefer memory.free(); + for (p.pieces.items) |piece| try memory.upload(piece.dst, piece.bytes); + w.memory = memory; + w.bytes = p.bytes; + shift(&w, memory.ptr); + inline for (.{ &w.asr, &w.sf }) |s| { + shift(s, memory.ptr); + for (s.layers) |*l| shift(l, memory.ptr); + } + return w; + } + + pub fn deinit(w: *Weights, gpa: Allocator) void { + w.memory.free(); + gpa.free(w.asr.layers); + gpa.free(w.sf.layers); + w.host.deinit(gpa); + w.* = undefined; + } +}; + +/// The tensors the host side reads, widened to f32 where it computes with them. +pub const Host = struct { + window: []f32, // [400], bf16 values + fb: []f32, // [128, 257], bf16 values + sil: []u16, // learnable_sil_emb [512] bf16 bits + diar_norm: [Config.speakers]f32, + diar_norm_bias: [Config.speakers]f32, + + fn read(gpa: Allocator, p: *Planner, c: Config) !Host { + const window = try widen(gpa, (try p.get("perception.preprocessor.featurizer.window", &.{Config.win})).bytes); + errdefer gpa.free(window); + const fb = try widen(gpa, (try p.get("perception.preprocessor.featurizer.fb", &.{ 1, Config.mels, Config.bins })).bytes); + errdefer gpa.free(fb); + const sil = try gpa.alloc(u16, c.sf.d); + errdefer gpa.free(sil); + for (sil, std.mem.bytesAsSlice(u16, (try p.get(heads_ ++ "learnable_sil_emb", &.{c.sf.d})).bytes)) |*o, v| o.* = v; + var h: Host = .{ .window = window, .fb = fb, .sil = sil, .diar_norm = undefined, .diar_norm_bias = undefined }; + for ([2][]const u8{ "perception.encoder.diar_norm.weight", "perception.encoder.diar_norm.bias" }, [2]*[Config.speakers]f32{ &h.diar_norm, &h.diar_norm_bias }) |name, out| { + const t = try p.get(name, &.{Config.speakers}); + for (out, std.mem.bytesAsSlice(u16, t.bytes)) |*o, v| o.* = bf(v); + } + return h; + } + + fn deinit(h: *Host, gpa: Allocator) void { + gpa.free(h.window); + gpa.free(h.fb); + gpa.free(h.sil); + } +}; + +pub fn bf(v: u16) f32 { + return @bitCast(@as(u32, v) << 16); +} + +fn widen(gpa: Allocator, bytes: []const u8) ![]f32 { + const out = try gpa.alloc(f32, bytes.len / 2); + for (out, std.mem.bytesAsSlice(u16, bytes)) |*o, v| o.* = bf(v); + return out; +} + +/// One tensor's bytes and where they go in the allocation; `owned` bytes were rearranged on the host. +const Piece = struct { bytes: []const u8, dst: u64, owned: bool = false }; + +const Planner = struct { + gpa: Allocator, + files: []const st.File, + pieces: std.ArrayList(Piece) = .empty, + bytes: u64 = 256, // nothing at offset 0: an offset of 0 is a tensor the stack does not have + + fn get(p: *Planner, name: []const u8, shape: []const usize) !st.Tensor { + for (p.files) |*f| if (f.get(name)) |t| { + if (t.is(.bf16, shape)) return t; + std.log.err("{s}: {t} {any}, expected bf16 {any}", .{ name, t.dtype, t.shape[0..t.rank], shape }); + return error.UnexpectedTensor; + }; + std.log.err("checkpoint has no tensor {s}", .{name}); + return error.MissingTensor; + } + + fn place(p: *Planner, bytes: []const u8, owned: bool) !u64 { + const at = std.mem.alignForward(u64, p.bytes, 256); + try p.pieces.append(p.gpa, .{ .bytes = bytes, .dst = at, .owned = owned }); + p.bytes = at + bytes.len; + return at; + } + + /// `name` placed as stored; its offset from the allocation's start. + fn take(p: *Planner, name: []const u8, shape: []const usize) !u64 { + return p.place((try p.get(name, shape)).bytes, false); + } + + /// A conv1d weight [o, c, 3] as a GEMM's [o, 3 c]: tap j of channel c at column j c + c. + fn taps(p: *Planner, name: []const u8, o: usize, c: usize) !u64 { + const src = std.mem.bytesAsSlice(u16, (try p.get(name, &.{ o, c, 3 })).bytes); + const bytes = try p.gpa.alloc(u8, o * 3 * c * 2); // bytes, as the pieces are freed + errdefer p.gpa.free(bytes); + const out = std.mem.bytesAsSlice(u16, bytes); + for (0..o) |r| for (0..c) |ch| for (0..3) |j| { + out[r * 3 * c + j * c + ch] = src[(r * c + ch) * 3 + j]; + }; + return p.place(bytes, true); + } + + fn stack(p: *Planner, s: *Stack, pre: []const u8, c: Config.Encoder, qk_norm: bool) !void { + var buf: [192]u8 = undefined; + const d = c.d; + const dh = d / c.heads; + s.pre = try p.take(try join(&buf, pre, "pre_encode.proj.weight"), &.{ d, Config.sub * Config.mels }); + s.embed_norm = try p.take(try join(&buf, pre, "embed_norm.weight"), &.{d}); + s.embed_norm_bias = try p.take(try join(&buf, pre, "embed_norm.bias"), &.{d}); + s.final = try p.take(try join(&buf, pre, "final_norm.weight"), &.{d}); + s.final_bias = try p.take(try join(&buf, pre, "final_norm.bias"), &.{d}); + for (s.layers, 0..) |*l, i| { + const n = struct { + fn of(b: []u8, prefix: []const u8, li: usize, part: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "{s}layers.{d}.{s}", .{ prefix, li, part }); + } + }.of; + l.* = .{ + .norm1 = try p.take(try n(&buf, pre, i, "norm1.weight"), &.{d}), + .norm1_bias = try p.take(try n(&buf, pre, i, "norm1.bias"), &.{d}), + .qkv = try p.take(try n(&buf, pre, i, "attn.w_qkv.weight"), &.{ 3 * d, d }), + .out = try p.take(try n(&buf, pre, i, "attn.out_proj.weight"), &.{ d, d }), + .out_bias = try p.take(try n(&buf, pre, i, "attn.out_proj.bias"), &.{d}), + .norm2 = try p.take(try n(&buf, pre, i, "norm2.weight"), &.{d}), + .norm2_bias = try p.take(try n(&buf, pre, i, "norm2.bias"), &.{d}), + .fc1 = try p.take(try n(&buf, pre, i, "ffn.net.0.weight"), &.{ 4 * d, d }), + .fc1_bias = try p.take(try n(&buf, pre, i, "ffn.net.0.bias"), &.{4 * d}), + .fc2 = try p.take(try n(&buf, pre, i, "ffn.net.3.weight"), &.{ d, 4 * d }), + .fc2_bias = try p.take(try n(&buf, pre, i, "ffn.net.3.bias"), &.{d}), + }; + if (qk_norm) { + l.q_norm = try p.take(try n(&buf, pre, i, "attn.q_norm.weight"), &.{dh}); + l.q_norm_bias = try p.take(try n(&buf, pre, i, "attn.q_norm.bias"), &.{dh}); + l.k_norm = try p.take(try n(&buf, pre, i, "attn.k_norm.weight"), &.{dh}); + l.k_norm_bias = try p.take(try n(&buf, pre, i, "attn.k_norm.bias"), &.{dh}); + } + } + } +}; + +fn join(buf: []u8, a: []const u8, b: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}{s}", .{ a, b }); +} + +/// Every placed u64 offset of `s` (not its byte count, nor a 0 that marks a tensor it lacks) moved by `base`. +fn shift(s: anytype, base: u64) void { + const info = @typeInfo(@TypeOf(s.*)).@"struct"; + inline for (info.field_names, info.field_types) |name, T| if (T == u64 and comptime !std.mem.eql(u8, name, "bytes")) { + if (@field(s, name) != 0) @field(s, name) += base; + }; +} diff --git a/zig/src/families/nemotron/config.zig b/zig/src/families/nemotron/config.zig index cb1152677e..6ed9775a1b 100644 --- a/zig/src/families/nemotron/config.zig +++ b/zig/src/families/nemotron/config.zig @@ -101,6 +101,7 @@ pub const Config = struct { eos_count: usize = 0, latent: usize = 0, // moe_latent_size: routed experts read and write this width between two bf16 projections omni: bool = false, // nemotron_h_omni: the text model's tensors sit under "language_model." + salm: bool = false, // nemo_speechlm (Nemotron 3.5 Transcribe): bf16 throughout, NeMo's names under "llm." shares: usize = 1, // latent MoE: expert widths cut in this many shares, summed in share order on every TP world held: ?usize = null, // a split checkpoint (tensorfold_tp): the one share its expert tensors hold rank: usize = 0, // this process's tensor-parallel rank of `world`: set at load, never by config.json @@ -147,6 +148,24 @@ pub const Config = struct { return if (self.omni) "language_model." else ""; } + /// The stem of layer i's names after prefix(): HF's "backbone.layers.", NeMo's "model.layers." (salm). + pub fn layerStem(self: Config) []const u8 { + return if (self.salm) "model.layers." else "backbone.layers."; + } + + /// The token table, the final norm and a Mamba layer's fp32 parameters by their names after prefix(). + pub fn embedName(self: Config) []const u8 { + return if (self.salm) "model.embed_tokens.weight" else "backbone.embeddings.weight"; + } + + pub fn normName(self: Config) []const u8 { + return if (self.salm) "model.norm.weight" else "backbone.norm_f.weight"; + } + + pub fn ssmParams(self: Config) []const u8 { + return if (self.salm) "_fp32_params." else ""; + } + /// The expert width shares this rank holds and sums before the ranks meet. pub fn ownShares(self: Config) usize { return self.shares / self.world; @@ -380,9 +399,10 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config var parsed = try std.json.parseFromSlice(std.json.Value, allocator, text, .{}); defer parsed.deinit(); const root = parsed.value.object; - // nemotron_h_omni (Bolt) wraps the language model's config in llm_config, its quantization at the top + // nemotron_h_omni (Bolt, quantization at the top) and nemo_speechlm (Transcribe) wrap theirs in llm_config const wrapped = isType(root, "nemotron_h_omni"); - const o = if (wrapped) (llmConfig(root) orelse return error.BadConfig) else root; + const salm = isType(root, "nemo_speechlm"); + const o = if (wrapped or salm) (llmConfig(root) orelse return error.BadConfig) else root; if (!isType(o, "nemotron_h") and !isType(o, "nemotron_h_puzzle")) return error.NotNemotronH; var c = Config{ .hidden = try int(o, "hidden_size"), @@ -430,7 +450,7 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config c.layers = blocks.len; } // generation_config's [2, 11] is the omni wrapper's eos plus the language model's - if (wrapped) if (root.get("eos_token_id")) |e| eos(&c, e); + if (wrapped or salm) if (root.get("eos_token_id")) |e| eos(&c, e); if (o.get("eos_token_id")) |e| eos(&c, e); if (o.get("time_step_limit")) |v| if (v == .array and v.array.items.len == 2) { c.dt_min = @floatCast(try number(v.array.items[0])); @@ -443,7 +463,9 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config if (!c.bolt) c.shares = 2; try split(&c, root, why); }; - try quantization(&c, root, why); + c.salm = salm; + // the published bf16 weights, read per tensor as the ModelOpt loader reads a bf16 projection + if (salm) c.format = .modelopt else try quantization(&c, root, why); return c; } diff --git a/zig/src/families/nemotron/cuda.zig b/zig/src/families/nemotron/cuda.zig index 428ff9cb4a..d1ca19fea9 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -23,6 +23,8 @@ pub const native = @import("cuda_native.zig"); pub const train = @import("cuda_train.zig"); /// Bolt's image tower (C-RADIO) and projector: one image's patches to the rows its `` placeholders take. pub const vision = @import("cuda_vision.zig"); +pub const audio = @import("cuda_audio.zig"); +pub const nemo_audio = @import("nemo_audio"); pub const sites = @import("cuda_sites.zig"); pub const train_ops = @import("cuda_train_ops.zig"); pub const glue = @import("cuda_glue.zig"); @@ -34,6 +36,7 @@ pub const pages = @import("cuda_pages.zig"); pub const slide = @import("cuda_slide.zig"); pub const worker = @import("cuda_worker.zig"); pub const omni = @import("cuda_native_omni.zig"); +pub const transcribe = @import("cuda_native_transcribe.zig"); pub const identity = @import("core").checkpoint_identity; test { diff --git a/zig/src/families/nemotron/cuda_audio.zig b/zig/src/families/nemotron/cuda_audio.zig new file mode 100644 index 0000000000..2136bf618f --- /dev/null +++ b/zig/src/families/nemotron/cuda_audio.zig @@ -0,0 +1,100 @@ +//! Transcribe's audio on CUDA: the encoder (nemo_audio) and a prompt pass's audio rows, spliced as image rows are. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const radio = @import("radio"); +const nemo_audio = @import("nemo_audio"); + +/// Audio rows a request may hold at most: the card's hour of 16 kHz audio (12.5 rows a second). +pub const hour_rows: u32 = @intCast(nemo_audio.rows(3600 * nemo_audio.Config.sample_rate)); + +/// The rows buffer for a `context`-token window: an hour of audio, or the window when it holds less. +pub fn limitFor(context: usize) u32 { + return @intCast(@min(context, hour_rows)); +} +/// Clips a request may hold. +pub const max_clips = 8; + +pub const Audio = struct { + gpa: std.mem.Allocator, + rows: radio.Rows, + encoder: ?*nemo_audio.Encoder = null, // null: rows set from the host (a check's reference rows) + spans: []radio.Span = &.{}, // the prompt pass in progress: each clip's rows from its first <|audio|> position + + /// A buffer of `limit` rows of the language model's width, and the encoder of the checkpoint at `dir` if given. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: ?[]const u8, limit: u32, hidden: usize) !*Audio { + const a = try gpa.create(Audio); + errdefer gpa.destroy(a); + a.* = .{ .gpa = gpa, .rows = try radio.Rows.init(d, limit, hidden) }; + errdefer a.rows.deinit(); + if (dir) |path| { + a.encoder = try nemo_audio.Encoder.init(gpa, io, d, path); + if (a.encoder.?.c.text_hidden != hidden) return error.AudioRowShape; + } + return a; + } + + pub fn deinit(a: *Audio) void { + a.gpa.free(a.spans); + if (a.encoder) |x| x.deinit(); + a.rows.deinit(); + a.gpa.destroy(a); + } + + /// Device bytes the encoder and the rows buffer hold. + pub fn bytes(a: *const Audio) u64 { + return a.rows.buf.len + if (a.encoder) |x| x.bytes() else 0; + } + + /// One clip's bf16 rows (`rows.hidden` each) from host memory, standing at prompt positions [at, at + n). + pub fn set(a: *Audio, s: cuda.Stream, at: u32, bf16: []const u8) !void { + const w = a.rows.hidden * 2; + if (bf16.len == 0 or bf16.len % w != 0) return error.AudioRows; + const n: u32 = @intCast(bf16.len / w); + if (n > a.rows.limit) return error.AudioRows; + try a.rows.d.check(a.rows.d.api.cuMemcpyHtoDAsync_v2(a.rows.buf.ptr, bf16.ptr, bf16.len, s.handle), "cuMemcpyHtoDAsync"); + a.gpa.free(a.spans); + a.spans = try a.gpa.dupe(radio.Span, &.{.{ .at = at, .n = n, .row = 0, .image = 0 }}); + } + + /// A prompt pass from row `from`: the clips whose rows end past it, through the encoder into the rows buffer. + pub fn begin(a: *Audio, s: cuda.Stream, clips: []const lanes.stream.Audio, from: usize) !void { + a.end(); + const enc = a.encoder orelse return error.NoAudioEncoder; + var spans: std.ArrayList(radio.Span) = .empty; + errdefer spans.deinit(a.gpa); + var row: u32 = 0; + for (clips, 0..) |c, i| { + if (c.at + c.tokens <= from) continue; // a kept prefix holds it already + if (row + c.tokens > a.rows.limit) return error.AudioRows; + const sp: radio.Span = .{ .at = c.at, .n = c.tokens, .row = row, .image = @intCast(i) }; + if (try enc.encode(s, c.samples, a.rows.at(sp), null) != c.tokens) return error.AudioRows; + try spans.append(a.gpa, sp); + row += c.tokens; + } + a.spans = try spans.toOwnedSlice(a.gpa); + } + + /// The pass ended: no rows splice into later passes until another begins. + pub fn end(a: *Audio) void { + a.gpa.free(a.spans); + a.spans = &.{}; + } + + /// Prompt rows [lo, hi) embedded at `emb` take the audio rows they cross. + pub fn splice(a: *const Audio, s: cuda.Stream, lo: usize, hi: usize, emb: u64) !void { + if (a.spans.len == 0) return; + try a.rows.splice(s, a.spans, @intCast(lo), @intCast(hi), emb); + } + + /// A WAV clip as the encoder reads it: its samples, its rows and a hash of its samples (the prompt cache's key). + pub fn prepare(a: *const Audio, alloc: std.mem.Allocator, wav: []const u8, max_rows: u32) !struct { samples: []f32, tokens: u32, hash: u64 } { + _ = a; + const samples = try nemo_audio.mel.decodeWav(alloc, wav); + errdefer alloc.free(samples); + if (samples.len < nemo_audio.Config.hop) return error.AudioTooShort; + const n = nemo_audio.rows(samples.len); + if (n > max_rows) return error.AudioTooLong; + return .{ .samples = samples, .tokens = @intCast(n), .hash = std.hash.Wyhash.hash(0x6175_6469_6f, std.mem.sliceAsBytes(samples)) }; + } +}; diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index 6484c04f9f..fa22b527e9 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -46,6 +46,11 @@ pub const Cancel = prompt_hooks.Cancel; pub const Keep = prompt_hooks.Keep; pub const Taps = prompt_hooks.Taps; +/// The one tap Nemotron gives: past its last layer, the final norm's rows its head and MTP head read. +pub fn finalTap(c: Config) u32 { + return @intCast(c.layers); +} + /// Serial rounds a host keeps queued ahead of the one it reads. pub const lookahead = 4; @@ -116,6 +121,7 @@ pub const Engine = struct { last_chunk: usize = 0, // where the last prompt pass's last chunk started (the prompt's last row: len - 1 - this) tp: ?cuda.tp_group.Group = null, // a latent MoE over ranks (TF_NEMOTRON_RANK of TF_NEMOTRON_WORLD) vision: ?*@import("cuda_vision.zig").Vision = null, // --vision: the tower and a prompt pass's image rows + audio: ?*@import("cuda_audio.zig").Audio = null, // a Transcribe checkpoint's audio rows for the prompt pass snap_host: bool = false, // kept prompt states in pinned host memory (cuda_reuse.zig), set when the host opens the engine load_pending: ?*weights.PendingLoad = null, load_has_weights: bool = false, @@ -153,6 +159,7 @@ pub const Engine = struct { if (opts.slide) return error.SlideNeedsMlx; // the learner reads MLX 4-bit projections try e.k.loadModelopt(ctx); } + if (e.c.salm) try e.k.loadBf16Experts(ctx); if (e.c.latent > 0) try e.k.loadLatent(); if (e.c.ck) try e.k.loadCk(ctx); // the ranks' lines first: a rank that never comes, refuses, fails or dies is named, not waited on @@ -198,6 +205,12 @@ pub const Engine = struct { }) orelse return error.VisionRefused; } errdefer if (!keep_owner) if (e.vision) |v| v.deinit(); + // a Transcribe checkpoint's audio encoder and its rows, loaded before the caches plan their memory + if (e.c.salm) { + const audio = @import("cuda_audio.zig"); + e.audio = try audio.Audio.init(gpa, io, ctx.d, model_dir, audio.limitFor(opts.context orelse default_context), e.c.hidden); + } + errdefer if (!keep_owner) if (e.audio) |a| a.deinit(); e.b = try state.Buffers.init(ctx.d, e.c, e.max_len, e.nch, e.carve); errdefer if (!keep_owner) e.b.deinit(); try e.b.initShared(e.ops(), e.c); @@ -325,6 +338,7 @@ pub const Engine = struct { e.side.join.deinit(); if (e.tp) |*g| g.close(); if (e.vision) |v| v.deinit(); + if (e.audio) |a| a.deinit(); e.stream.deinit(); e.gpa.destroy(e); } @@ -409,6 +423,7 @@ pub const Engine = struct { f.untaught = e.bound.untaught; f.tp = if (e.tp) |*g| g else null; f.vision = e.vision; + f.audio = e.audio; if (dump == null and (e.c.latent > 0 or (e.c.experts <= 128 and e.c.slots() <= 8))) f.side = e.side; // tf_plan_routed's bounds; a latent MoE plans without it return f; } @@ -442,24 +457,7 @@ pub const Engine = struct { return null; } - /// Dev trace (TF_SLIDE_TRACE): the bound sequence's router scores at its decision row and which lessons act on it. - pub fn traceLessons(e: *const Engine) void { - const ad = e.lessons() orelse return; - if (e.decide == lesson_route.none) return; - const n = lesson_route.max_blocks; - var score: [n]f32 = undefined; - var need: [n]f32 = undefined; - var flags: [n]u32 = undefined; - const o = e.ops(); - o.download(std.mem.sliceAsBytes(&score), ad.score + e.slot * n * 4) catch return; - o.download(std.mem.sliceAsBytes(&need), ad.need) catch return; - o.download(std.mem.sliceAsBytes(&flags), ad.latch + e.slot * n * 4) catch return; - e.stream.synchronize() catch return; - var line: [2048]u8 = undefined; - var w: std.Io.Writer = .fixed(&line); - for (0..ad.rank / 16) |b| w.print(" {d:.2}/{d:.2}{s}", .{ score[b], need[b], if (flags[b] != 0) "*" else "" }) catch break; - std.log.info("slide trace: decision row {d}: scores{s} (* acting)", .{ e.decide, w.buffered() }); - } + pub const traceLessons = @import("cuda_lesson_trace.zig").traceLessons; pub fn trainOps(e: *const Engine) train_ops.Ops { return .{ .k = &e.k.train, .s = e.stream }; diff --git a/zig/src/families/nemotron/cuda_forward.zig b/zig/src/families/nemotron/cuda_forward.zig index fdf77ac8cf..6386aa66be 100644 --- a/zig/src/families/nemotron/cuda_forward.zig +++ b/zig/src/families/nemotron/cuda_forward.zig @@ -63,6 +63,7 @@ pub const Forward = struct { captures: []const kept.Mark = &.{}, // marks inside the prompt chunk: their Mamba states are copied out as it runs tp: ?*const cuda.tp_group.Group = null, // a latent MoE's ranks: their expert rows meet every MoE layer vision: ?*const @import("cuda_vision.zig").Vision = null, // a prompt pass's image rows, spliced after its embeds + audio: ?*const @import("cuda_audio.zig").Audio = null, // a prompt pass's audio rows, spliced the same way page_elems: u64 = 0, // elements between a paged sequence's pages (its layers' K and V of state.chunk_keys rows) pub fn init(c: Config, w: *const weights.Weights, b: *const state.Buffers, ops: kern.Ops, max_len: usize, nch: usize, sampled: bool) Forward { @@ -406,6 +407,7 @@ pub const Forward = struct { const b = f.b; try f.embed(b.p_ids, b.emb, w.rows); if (f.vision) |v| try v.splice(f.ops.s, w.pos, w.pos + w.rows, b.emb); + if (f.audio) |a| try a.splice(f.ops.s, w.pos, w.pos + w.rows, b.emb); w.x = b.emb; } diff --git a/zig/src/families/nemotron/cuda_kernels.zig b/zig/src/families/nemotron/cuda_kernels.zig index bf1bc35e73..e0d5b8b857 100644 --- a/zig/src/families/nemotron/cuda_kernels.zig +++ b/zig/src/families/nemotron/cuda_kernels.zig @@ -72,6 +72,7 @@ pub const Kernels = struct { latent: ?@import("cuda_latent.zig").Fns = null, // a latent MoE's glue (loadLatent), else never ck: ?cuda.ck.Ck = null, // the checkpoint's own math (loadCk): FP8 rows and NVFP4 experts, sm_120 on ggml: ?cuda.ggml.Kernels = null, // loaded by loadGgml for a checkpoint with ggml blocks (Bolt's IQ2_XXS and Q8_0) + bf16x: ?cuda.experts.Bf16Experts = null, // loaded by loadBf16Experts for a bf16 checkpoint's experts (Transcribe) gb10: bool, discrete: bool, // the card has its own memory: checkpoint bytes reach it through page-locked slots @@ -115,6 +116,7 @@ pub const Kernels = struct { k.latent = null; k.ck = null; k.ggml = null; + k.bf16x = null; k.triton = if (triton_dir) |dir| try cuda.aot.Set.load(gpa, io, d, ctx.device, dir) else null; errdefer if (k.triton) |*t| t.deinit(); k.train = try train_ops.Kernels.load(d); @@ -143,6 +145,7 @@ pub const Kernels = struct { m.pack.unload(); } if (k.ggml) |*g| g.unload(); + if (k.bf16x) |*x| x.unload(); k.affine.deinit(); k.train.deinit(); if (k.triton) |*t| t.deinit(); @@ -178,6 +181,12 @@ pub const Kernels = struct { k.modelopt = .{ .lane = lane, .gemm = gemm, .experts = experts, .pack = try cuda.modelopt.Packer.load(k.d) }; } + /// The grouped bf16 experts a checkpoint's own bf16 experts run on (Nemotron 3.5 Transcribe). + pub fn loadBf16Experts(k: *Kernels, ctx: *const cuda.Context) !void { + if (k.bf16x != null) return; + k.bf16x = try cuda.experts.Bf16Experts.load(k.d, @intCast(try ctx.attribute(.multiprocessor_count))); + } + /// The ggml block kernels (IQ2_XXS experts, IQ2_XXS and Q8_0 projections, their repack) Bolt's export needs. pub fn loadGgml(k: *Kernels, ctx: *const cuda.Context) !void { if (k.ggml != null) return; @@ -262,7 +271,7 @@ pub const Ops = struct { /// Every grouped-expert format the kernels loaded (cuda/experts.zig). pub fn grouped(o: Ops) cuda.experts.Grouped { - return .{ .affine4 = &o.k.experts, .nvfp4 = if (o.k.modelopt) |*x| &x.experts else null, .iq2xxs = if (o.k.ggml) |*g| g else null }; + return .{ .affine4 = &o.k.experts, .nvfp4 = if (o.k.modelopt) |*x| &x.experts else null, .iq2xxs = if (o.k.ggml) |*g| g else null, .bf16 = if (o.k.bf16x) |*x| x else null }; } /// LatentMoE: z[r] = sum_k wt[r, k] y[r * k_slots + k], rows of `l` bf16 from fp32 pair rows (the shared combine). diff --git a/zig/src/families/nemotron/cuda_lanes.zig b/zig/src/families/nemotron/cuda_lanes.zig index c0606a92f9..8330ff072f 100644 --- a/zig/src/families/nemotron/cuda_lanes.zig +++ b/zig/src/families/nemotron/cuda_lanes.zig @@ -33,11 +33,6 @@ const Lane = struct { seq: *state.Seq, own: bool = false, paged: bool = false, p /// A recorded stream's rows: its prompt's on the host until its first verify, then where its last verify started. const Taps = struct { prompt: std.ArrayList(u8) = .empty, from: usize = 0, verified: ?u64 = null }; -/// The one tap Nemotron gives: past its last layer, the final norm's rows its head and MTP head read. -pub fn finalTap(c: config.Config) u32 { - return @intCast(c.layers); -} - pub const Cuda = struct { gpa: std.mem.Allocator, e: *Engine, @@ -299,6 +294,8 @@ pub const Cuda = struct { const keep: ?engine.Keep = if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; if (s.images.len > 0) try (e.vision orelse return error.NoVision).begin(e.stream, s.images, from); defer if (e.vision) |v| v.end(); + if (s.audios.len > 0) try (e.audio orelse return error.NoAudio).begin(e.stream, s.audios, from); + defer if (e.audio) |a| a.end(); var sink: PromptSink = .{ .gpa = self.gpa, .s = s }; s.prompt_rows.clearRetainingCapacity(); if (s.prompt_logprobs) |k| { @@ -368,14 +365,14 @@ pub const Cuda = struct { /// Recording asks for the final norm's rows alone; their staging buffer is made here. fn prepareFeaturesFn(ptr: *anyopaque, taps: []const u32) anyerror!void { const self = of(ptr); - if (taps.len != 1 or taps[0] != finalTap(self.e.c)) return error.UnsupportedTaps; + if (taps.len != 1 or taps[0] != engine.finalTap(self.e.c)) return error.UnsupportedTaps; if (self.tapped == null) self.tapped = try cuda.HostBuffer.alloc(self.e.ctx.d, state.pass_rows * self.e.c.hidden * 2); } /// The prompt's rows from the host until the first verify; then the last verify's, read from its sequence. fn featuresFn(ptr: *anyopaque, s: *lanes.Stream, taps: []const u32, start: u64, count: u32) anyerror!be.Features { const self = of(ptr); - if (self.tapped == null or taps.len != 1 or taps[0] != finalTap(self.e.c)) return error.UnsupportedTaps; + if (self.tapped == null or taps.len != 1 or taps[0] != engine.finalTap(self.e.c)) return error.UnsupportedTaps; const l = self.lanes.getPtr(s) orelse return error.UnknownStream; const d = self.e.c.hidden * 2; const t = &l.taps; diff --git a/zig/src/families/nemotron/cuda_lesson_trace.zig b/zig/src/families/nemotron/cuda_lesson_trace.zig new file mode 100644 index 0000000000..7d30f817a8 --- /dev/null +++ b/zig/src/families/nemotron/cuda_lesson_trace.zig @@ -0,0 +1,23 @@ +//! The Nemotron engine's slide trace (TF_SLIDE_TRACE) on CUDA. +const std = @import("std"); +const lesson_route = @import("living").lesson_route; +const Engine = @import("cuda_engine.zig").Engine; + +/// Dev trace (TF_SLIDE_TRACE): the bound sequence's router scores at its decision row and which lessons act on it. +pub fn traceLessons(e: *const Engine) void { + const ad = e.lessons() orelse return; + if (e.decide == lesson_route.none) return; + const n = lesson_route.max_blocks; + var score: [n]f32 = undefined; + var need: [n]f32 = undefined; + var flags: [n]u32 = undefined; + const o = e.ops(); + o.download(std.mem.sliceAsBytes(&score), ad.score + e.slot * n * 4) catch return; + o.download(std.mem.sliceAsBytes(&need), ad.need) catch return; + o.download(std.mem.sliceAsBytes(&flags), ad.latch + e.slot * n * 4) catch return; + e.stream.synchronize() catch return; + var line: [2048]u8 = undefined; + var w: std.Io.Writer = .fixed(&line); + for (0..ad.rank / 16) |b| w.print(" {d:.2}/{d:.2}{s}", .{ score[b], need[b], if (flags[b] != 0) "*" else "" }) catch break; + std.log.info("slide trace: decision row {d}: scores{s} (* acting)", .{ e.decide, w.buffered() }); +} diff --git a/zig/src/families/nemotron/cuda_load.zig b/zig/src/families/nemotron/cuda_load.zig index 88c25be5ef..29474c386e 100644 --- a/zig/src/families/nemotron/cuda_load.zig +++ b/zig/src/families/nemotron/cuda_load.zig @@ -8,6 +8,7 @@ const kern = @import("cuda_kernels.zig"); const Source = @import("cuda_source.zig").Source; const latent = @import("cuda_latent.zig"); const modelopt = @import("cuda_modelopt.zig"); +const salm = @import("cuda_salm.zig"); const join = weights.join; const Block = weights.Block; pub const Pending = struct { @@ -76,6 +77,12 @@ fn fill(p: *Pending, io: std.Io, dir: []const u8, c: Config, with_mtp: bool) !vo var ck = try core.Checkpoint.openModel(gpa, io, dir); defer ck.close(); try src.add(&ck); + // a Transcribe checkpoint: the audio encoder and the CTC aligner load apart; TE's _extra_state holds nothing + if (c.salm) { + for ([_][]const u8{ "perception.", "ctc_timestamp." }) |part| try ck.skip(part); + try salm.skipExtraState(&ck); + ck.scope = "llm."; + } const pre = c.prefix(); w.blocks = try gpa.alloc(Block, c.layers); for (c.kinds[0..c.layers], 0..) |kind, i| { @@ -83,13 +90,13 @@ fn fill(p: *Pending, io: std.Io, dir: []const u8, c: Config, with_mtp: bool) !vo var nb: [128]u8 = undefined; var pb: [128]u8 = undefined; const name = try std.fmt.bufPrint(&nm, "blocks.{d}", .{i}); - const norm = try ck.expect(try std.fmt.bufPrint(&pb, "{s}backbone.layers.{d}.norm.weight", .{ pre, i }), .bf16, &.{c.hidden}); + const norm = try ck.expect(try std.fmt.bufPrint(&pb, "{s}{s}{d}.norm.weight", .{ pre, c.layerStem(), i }), .bf16, &.{c.hidden}); var b: Block = .{ .kind = kind, .norm = try L.raw(try join(&nb, name, ".norm"), norm) }; - const mixer = try std.fmt.bufPrint(&pb, "{s}backbone.layers.{d}.mixer.", .{ pre, i }); + const mixer = try std.fmt.bufPrint(&pb, "{s}{s}{d}.mixer.", .{ pre, c.layerStem(), i }); switch (kind) { .mamba => b.mamba = try L.mamba(&ck, try join(&nb, name, ".mamba"), mixer), .attention => b.attn = try L.attention(&ck, try join(&nb, name, ".attn"), mixer), - .moe => b.moe = if (c.bolt) try modelopt.latent(L, &ck, try join(&nb, name, ".moe"), mixer, c, i) else if (c.latent > 0) try latent.load(L, &ck, try join(&nb, name, ".moe"), mixer, c, @intCast(i)) else try L.moe(&ck, try join(&nb, name, ".moe"), mixer, c), + .moe => b.moe = if (c.salm) try salm.moe(L, &ck, try join(&nb, name, ".moe"), mixer, c) else if (c.bolt) try modelopt.latent(L, &ck, try join(&nb, name, ".moe"), mixer, c, i) else if (c.latent > 0) try latent.load(L, &ck, try join(&nb, name, ".moe"), mixer, c, @intCast(i)) else try L.moe(&ck, try join(&nb, name, ".moe"), mixer, c), } w.blocks[i] = b; } @@ -104,10 +111,10 @@ fn fill(p: *Pending, io: std.Io, dir: []const u8, c: Config, with_mtp: bool) !vo }; } var tail: [96]u8 = undefined; - w.norm_f = try L.raw("norm_f", try ck.expect(try std.fmt.bufPrint(&tail, "{s}backbone.norm_f.weight", .{pre}), .bf16, &.{c.hidden})); + w.norm_f = try L.raw("norm_f", try ck.expect(try std.fmt.bufPrint(&tail, "{s}{s}", .{ pre, c.normName() }), .bf16, &.{c.hidden})); w.head = try L.dense(&ck, "head", &.{try std.fmt.bufPrint(&tail, "{s}lm_head", .{pre})}); // a ModelOpt checkpoint keeps its MTP head in its own shards - if (c.format == .modelopt) if (!with_mtp) try ck.skip(try std.fmt.bufPrint(&tail, "{s}mtp.", .{pre})) else if (c.latent > 0) try latent.mtp(L, &ck, c) else try modelopt.mtp(L, &ck, c); + if (c.format == .modelopt) if (!with_mtp) try ck.skip(try std.fmt.bufPrint(&tail, "{s}mtp.", .{pre})) else if (c.salm) try salm.mtp(L, &ck, c) else if (c.latent > 0) try latent.mtp(L, &ck, c) else try modelopt.mtp(L, &ck, c); // the picture tower and its projector load with the image path (an omni checkpoint's other parts) if (c.omni) for ([_][]const u8{ "vision_model.", "mlp1.", "vision_projector." }) |part| try ck.skip(part); if (ck.unused() != 0) return error.UnusedCheckpointTensors; diff --git a/zig/src/families/nemotron/cuda_modelopt.zig b/zig/src/families/nemotron/cuda_modelopt.zig index c3baac3bfd..c5e4d0db4e 100644 --- a/zig/src/families/nemotron/cuda_modelopt.zig +++ b/zig/src/families/nemotron/cuda_modelopt.zig @@ -322,7 +322,7 @@ fn experts(L: *Loader, ck: *core.Checkpoint, label: []const u8, pre: []const u8, /// The bf16 token table as stored. pub fn embed(L: *Loader, ck: *core.Checkpoint, c: Config) !W.Embed { var buf: [64]u8 = undefined; - const t = try ck.expect(try name(&buf, "{s}backbone.embeddings.weight", .{c.prefix()}), .bf16, &.{ c.vocab, c.hidden }); + const t = try ck.expect(try name(&buf, "{s}{s}", .{ c.prefix(), c.embedName() }), .bf16, &.{ c.vocab, c.hidden }); return .{ .w = try L.raw("embed.weight", t), .n = c.vocab, .k = c.hidden, .bf16 = true }; } diff --git a/zig/src/families/nemotron/cuda_mtp.zig b/zig/src/families/nemotron/cuda_mtp.zig index b67e1fd4b3..11cda41f08 100644 --- a/zig/src/families/nemotron/cuda_mtp.zig +++ b/zig/src/families/nemotron/cuda_mtp.zig @@ -459,8 +459,8 @@ pub const Head = struct { return .{ .token = tok[0], .prob = prob[0] }; } - /// MTPHead.absorb_rows: a prompt chunk's rows into the head's cache, keys and values only. - pub fn absorb(h: *Head, hidden: u64, tokens: []const u32) !void { + /// MTPHead.absorb_rows: a prompt chunk's keys and values into the head's cache, `tokens` at [at, at + len). + pub fn absorb(h: *Head, hidden: u64, tokens: []const u32, at: usize) !void { const e = h.e; const c = e.c; const m = e.w.mtp.?; @@ -472,8 +472,13 @@ pub const Head = struct { try h.copied.synchronize(); const host = h.pinned.slice(u32)[host_tokens..][0..tokens.len]; @memcpy(host, tokens); + // a speech prompt's <|audio|> placeholders take the clip's rows, as the trunk's embeds do (cuda_audio.zig) + if (e.audio != null) for (host) |*t| if (t.* >= c.vocab) { + t.* = 0; + }; try o.upload(h.atok, std.mem.sliceAsBytes(host)); try f.embed(h.atok, h.emb, rows); + if (e.audio) |a| try a.splice(o.s, at, at + rows, h.emb); try f.glue.concatNorms(h.emb, hidden, m.enorm, m.hnorm, h.cat, h.cxs, rows, c.hidden, c.eps); try o.prefillDense(h.cat, m.eh_proj, h.x, rows); try f.glue.addRmsnorm(h.x, null, m.attn_norm, h.x, b.y, b.xs, rows, c.hidden, c.eps); diff --git a/zig/src/families/nemotron/cuda_native.zig b/zig/src/families/nemotron/cuda_native.zig index 9f23a7b634..2c6371dadf 100644 --- a/zig/src/families/nemotron/cuda_native.zig +++ b/zig/src/families/nemotron/cuda_native.zig @@ -27,7 +27,7 @@ pub const slide = @import("cuda_slide.zig"); pub const model_type = "nemotron_h"; pub const explains_open = true; // explain() words open refusals too (--vision on a checkpoint without a tower) /// Nemotron 3.5 Bolt's config: the language model under llm_config, beside its vision tower's. -pub const also_model_types: []const []const u8 = &.{"nemotron_h_omni"}; +pub const also_model_types: []const []const u8 = &.{ "nemotron_h_omni", "nemo_speechlm" }; pub const formats: []const []const u8 = &.{ "mlx-q4g64", "modelopt-nvfp4", "modelopt-ggml" }; pub const default_context: i64 = engine.default_context; pub const max_segments: u32 = @import("cuda_segments.zig").MAX; @@ -63,6 +63,7 @@ pub const Loaded = struct { learns: ?*engine.Engine = null, // --slide: the engine its learner trains, whose sites it attached at load learn_bytes: usize = 0, // --slide: what its trainer allocates at the first lesson, kept out of the streams' room vision: ?*vision.Vision = null, // --vision: the tower the server's image requests reach + audio: ?*@import("cuda_audio.zig").Audio = null, // a Transcribe checkpoint's encoder the server's audio requests reach taps: []const u32 = &.{}, // the rows a recorder may ask `backend` for (empty: none) host_snaps: bool = false, // a discrete card: prompt states go to pinned host memory when the cache does not fit /// The other streams' memory as one pool of KV pages (null: each allocates a window at its prefill). @@ -128,13 +129,13 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer gpa.destroy(own); own.* = .{ .gpa = gpa, .e = e, .head = head, .lanes = try Lanes.init(gpa, e, head), .target = .{ .e = e, .head = head }, .carve = carve }; errdefer own.lanes.deinit(); - own.lanes.copy_drafts = o.drafts; + own.lanes.copy_drafts = o.drafts and !e.c.salm; // a copied <|audio|> has no audio row own.lanes.mtp_drafts = o.mtp_drafts; own.lanes.mtp_confidence = o.mtp_confidence; own.living = try slide.loadLiving(gpa, io, e, dir); // before the lanes' costs are measured: they decode through it errdefer if (own.living) |*s| s.deinit(); try own.lanes.measure(io, dir); - own.taps = .{@import("cuda_lanes.zig").finalTap(e.c)}; + own.taps = .{engine.finalTap(e.c)}; return .{ .backend = own.lanes.backend(), .facts = own.lanes.facts(), @@ -147,6 +148,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .learns = if (o.slide) e else null, .learn_bytes = if (o.slide) @import("cuda_train.zig").Trainer.memory(e.c) else 0, .vision = e.vision, + .audio = e.audio, .taps = &own.taps, .host_snaps = e.k.discrete, // a card with its own memory may keep prompt states in the host's (GB10 shares one) .layout = if (o.slide or e.k.triton != null or paging() == .off) null else pages_mod.layout(e), // our glue pages; the learner reads windows diff --git a/zig/src/families/nemotron/cuda_native_transcribe.zig b/zig/src/families/nemotron/cuda_native_transcribe.zig new file mode 100644 index 0000000000..1bb7f80318 --- /dev/null +++ b/zig/src/families/nemotron/cuda_native_transcribe.zig @@ -0,0 +1,28 @@ +//! Nemotron 3.5 Transcribe (nemo_speechlm) on CUDA: the Nemotron engine with its audio encoder (cuda_audio.zig). +const native = @import("cuda_native.zig"); + +pub const model_type = "nemo_speechlm"; +pub const formats = native.formats; +/// The card's trained window: llm_config's 262144 is the backbone's, which this speech model was not trained to. +pub const trained_context: i64 = 65536; +pub const default_context = trained_context; +pub const max_segments = native.max_segments; +pub const prompt_rows = native.prompt_rows; +pub const Options = native.Options; +pub const explains_open = true; +pub const Loaded = native.Loaded; +pub const LoneRun = native.LoneRun; +pub const open = native.open; +pub const explain = native.explain; +pub const warm_quiet = native.warm_quiet; +pub const lone_min_budget = native.lone_min_budget; +pub const snap_bytes = native.snap_bytes; +pub const snap_save = native.snap_save; +pub const snap_restore = native.snap_restore; +pub const snap_drop = native.snap_drop; +pub const snap_slots = native.snap_slots; +pub const snap_place = native.snap_place; +pub const window_rows = native.window_rows; +pub const attach = native.attach; +pub const fits = native.fits; +pub const tensorWeightBytes = native.tensorWeightBytes; diff --git a/zig/src/families/nemotron/cuda_prompt.zig b/zig/src/families/nemotron/cuda_prompt.zig index 9785adae19..e1d7356cf1 100644 --- a/zig/src/families/nemotron/cuda_prompt.zig +++ b/zig/src/families/nemotron/cuda_prompt.zig @@ -39,7 +39,7 @@ pub fn prefillFrom(e: *Engine, prompt: []const u32, from: usize, dump: ?*Dump, h if (head) |h| try h.reset(); } else if (e.pos != from) return error.BadReuse; const waited = e.heat_gate.waited_s; - if (from == 0 and keep == null and e.vision == null and e.segments > 1 and dump == null and e.prompt_rows == null and e.taps == null and prompt.len > state.prefill_rows and e.decide == lesson_route.none and !Cancel.background(cancel)) { + if (from == 0 and keep == null and e.vision == null and e.audio == null and e.segments > 1 and dump == null and e.prompt_rows == null and e.taps == null and prompt.len > state.prefill_rows and e.decide == lesson_route.none and !Cancel.background(cancel)) { try segs.prefill(e, try e.segmentSet(), prompt, head, e.segments, cancel); } else try e.serialChunks(prompt, from, dump, head, cancel, keep); if (e.heat_gate.waited_s > waited) heat.note(e.heat_gate.waited_s - waited); @@ -96,6 +96,10 @@ pub fn serialChunks(e: *Engine, prompt: []const u32, from: usize, dump: ?*Dump, try e.copied.synchronize(); const host = e.promptHost(chunk.len); @memcpy(host, chunk); + // a speech prompt's <|audio|> placeholders lie past the token table: their rows are the clip's (cuda_audio.zig) + if (e.audio != null) for (host) |*t| if (t.* >= e.c.vocab) { + t.* = 0; + }; try e.ops().upload(e.b.p_ids, std.mem.sliceAsBytes(host)); try e.copied.record(e.stream); try f.chunk(@intCast(chunk.len), e.pos); @@ -103,7 +107,7 @@ pub fn serialChunks(e: *Engine, prompt: []const u32, from: usize, dump: ?*Dump, e.pos += @intCast(chunk.len); const known = @min(chunk.len, prompt.len - 1 - s); if (e.prompt_rows) |p| try prompt_rows.run(e, p, prompt[s + 1 ..][0..known]); - if (head) |h| if (known > 0) try h.absorb(e.b.p_hidden, prompt[s + 1 ..][0..known]); + if (head) |h| if (known > 0) try h.absorb(e.b.p_hidden, prompt[s + 1 ..][0..known], s + 1); if (e.taps) |t| try t.call(t.ctx, e.b.p_hidden, s, chunk.len); try reuse.finish(.{ .e = e, .head = head }); s = end; diff --git a/zig/src/families/nemotron/cuda_salm.zig b/zig/src/families/nemotron/cuda_salm.zig new file mode 100644 index 0000000000..1a2d470e8e --- /dev/null +++ b/zig/src/families/nemotron/cuda_salm.zig @@ -0,0 +1,103 @@ +//! Transcribe's language model: NeMo's bf16 tensors as published, routed experts as x @ W, the shared one nn.Linear. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const W = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); +const modelopt = @import("cuda_modelopt.zig"); +const Config = @import("config.zig").Config; +const draft_ids = @import("draft_ids.zig"); + +const Loader = W.Loader; +const Tensor = core.checkpoint.Tensor; +const Bf16 = cuda.experts.Bf16Experts; + +fn name(buf: []u8, comptime fmt: []const u8, args: anytype) ![]const u8 { + return std.fmt.bufPrint(buf, fmt, args); +} + +/// Every TE `_extra_state` (empty bytes beside a projection) marked as used: there is nothing in them to read. +pub fn skipExtraState(ck: *core.Checkpoint) !void { + for (ck.files.items) |*f| for (f.names.keys()) |k| { + if (std.mem.endsWith(u8, k, "._extra_state")) try ck.used.put(ck.gpa, k, {}); + }; +} + +/// A MoE layer: the bf16 router, its fp32 bias, routed experts and the shared expert's halves as experts E and E + 1. +pub fn moe(L: *Loader, ck: *core.Checkpoint, label: []const u8, pre: []const u8, c: Config) !W.MoE { + var a: [128]u8 = undefined; + var b: [192]u8 = undefined; + const r = try modelopt.router(L, try name(&a, "{s}.router", .{label}), try ck.get(try name(&b, "{s}gate.weight", .{pre})), c); + const bias = try ck.expect(try name(&b, "{s}gate.e_score_correction_bias", .{pre}), .f32, &.{c.experts}); + return .{ .router = r, .bias = try L.raw(try name(&a, "{s}.bias", .{label}), bias), .experts = try experts(L, ck, label, pre, c) }; +} + +fn experts(L: *Loader, ck: *core.Checkpoint, label: []const u8, pre: []const u8, c: Config) !kern.Experts { + const k = &(L.ops.k.bf16x orelse return error.Bf16ExpertsNotLoaded); + const E = c.experts; + const w = c.expert_width; + const D = c.hidden; + var layer: kern.Experts = .{ .format = .bf16, .up = 0, .down = 0, .width = w, .dims = D, .experts = E + 2, .relu2 = true }; + var a: [128]u8 = undefined; + var b: [192]u8 = undefined; + for ([2]bool{ true, false }) |up| { + const n = if (up) w else D; + const kk = if (up) D else w; + const routed = try ck.expect(try name(&b, "{s}experts.{s}", .{ pre, if (up) "gate_and_up_projs" else "down_projs" }), .bf16, &.{ E, kk, n }); + const shared = if (up) + try ck.expect(try name(&b, "{s}shared_experts.up_proj.weight", .{pre}), .bf16, &.{ 2 * w, D }) + else + try ck.expect(try name(&b, "{s}shared_experts.down_proj.weight", .{pre}), .bf16, &.{ D, 2 * w }); + const base = try L.tmp(routed.bytes.len + shared.bytes.len); + try L.src.upload(base, routed.bytes); + try L.src.upload(base + routed.bytes.len, shared.bytes); + const out = try L.alloc(try name(&a, "{s}.experts.{s}", .{ label, if (up) "up" else "down" }), Bf16.bytes(E + 2, n, kk)); + try L.src.flush(); + const s = L.ops.s; + const one = Bf16.bytes(1, n, kk); + try k.pack(s, base, .{ .stride = kk * n, .ld = n, .in_out = true }, out, E, n, kk); + const sh = base + routed.bytes.len; + for (0..2) |h| { + const from: Bf16.Source = if (up) .{ .ld = D, .n0 = h * w, .in_out = false } else .{ .ld = 2 * w, .k0 = h * w, .in_out = false }; + try k.pack(s, sh, from, out + (E + h) * one, 1, n, kk); + } + if (up) layer.up = out else layer.down = out; + } + return layer; +} + +/// The MTP head as stored: bf16 attention, input projection and experts; the draft head is lm_head's bf16 rows. +pub fn mtp(L: *Loader, ck: *core.Checkpoint, c: Config) !void { + const p0 = "mtp.layers.0."; + const p1 = "mtp.layers.1.mixer."; + const att = p0 ++ "mixer."; + L.w.mtp = .{ + .enorm = try L.raw("mtp.enorm", try ck.expect(p0 ++ "enorm.weight", .bf16, &.{c.hidden})), + .hnorm = try L.raw("mtp.hnorm", try ck.expect(p0 ++ "hnorm.weight", .bf16, &.{c.hidden})), + .eh_proj = try modelopt.dense(L, ck, "mtp.eh_proj", &.{p0 ++ "eh_proj"}), + .attn_norm = try L.raw("mtp.attn_norm", try ck.expect(p0 ++ "norm.weight", .bf16, &.{c.hidden})), + .attn = .{ .qkv = try modelopt.dense(L, ck, "mtp.attn.qkv", &.{ att ++ "q_proj", att ++ "k_proj", att ++ "v_proj" }), .o = try modelopt.dense(L, ck, "mtp.attn.o", &.{att ++ "o_proj"}) }, + .moe_norm = try L.raw("mtp.moe_norm", try ck.expect("mtp.layers.1.norm.weight", .bf16, &.{c.hidden})), + .moe = try moe(L, ck, "mtp.moe", p1, c), + .final_norm = try L.raw("mtp.final_norm", try ck.expect("mtp.layers.1.final_layernorm.weight", .bf16, &.{c.hidden})), + }; + try draftHead(L, ck, c); +} + +/// head_rows: lm_head's bf16 rows for the draft ids, packed as a head of their own. +fn draftHead(L: *Loader, ck: *core.Checkpoint, c: Config) !void { + const m = if (L.ops.k.modelopt) |*x| x else return error.ModeloptNotLoaded; + const ids = try draft_ids.load(L.gpa, c.vocab); + defer L.gpa.free(ids); + const p = try modelopt.part(ck, "lm_head"); + if (p.kind != .bf16 or p.n != c.vocab or p.k != c.hidden) return error.UnexpectedTensor; + const host = try L.staging(ids.len * p.k * 2); + try W.gatherRows(L, p.w, ids, host); + const base = try L.tmp(host.len); + try L.ops.upload(base, host); + const out = try L.alloc("draft_head.weight", cuda.modelopt.padded(ids.len) * p.k * 2); + try m.pack.bf16Lane(L.ops.s, base, out, ids.len, p.k); + L.w.draft_head = .{ .bf16 = cuda.modelopt.bf16Weight(out, ids.len, p.k) }; + try W.draftIds(L, ids); +} diff --git a/zig/src/families/nemotron/cuda_segments.zig b/zig/src/families/nemotron/cuda_segments.zig index 5957ac5f1f..4d475cca44 100644 --- a/zig/src/families/nemotron/cuda_segments.zig +++ b/zig/src/families/nemotron/cuda_segments.zig @@ -90,6 +90,7 @@ pub fn prefill(e: *Engine, s: *Segments, prompt: []const u32, head: ?*Head, part f.page_elems = e.pageBytes() / 2; f.untaught = e.bound.untaught; f.vision = e.vision; + f.audio = e.audio; seg[k] = .{ .f = f, .walk = .{ .rows = rows, .pos = e.pos + k * R } }; } try e.copied.record(e.stream); @@ -98,7 +99,7 @@ pub fn prefill(e: *Engine, s: *Segments, prompt: []const u32, head: ?*Head, part for (0..c.parts) |k| { // the caller's stream has joined every segment const start = at + k * R; const known = @min(seg[k].walk.rows, prompt.len - 1 - start); - if (head) |h| if (known > 0) try h.absorb(bufs[k].p_hidden, prompt[start + 1 ..][0..known]); + if (head) |h| if (known > 0) try h.absorb(bufs[k].p_hidden, prompt[start + 1 ..][0..known], start + 1); } last = c.parts - 1; last_rows = seg[last].walk.rows; diff --git a/zig/src/families/nemotron/cuda_weights.zig b/zig/src/families/nemotron/cuda_weights.zig index 8137349a03..6de92bd842 100644 --- a/zig/src/families/nemotron/cuda_weights.zig +++ b/zig/src/families/nemotron/cuda_weights.zig @@ -122,6 +122,11 @@ pub const Loader = struct { /// `.float()` of a bf16 tensor into the host staging buffer: exact. pub fn widened(L: *Loader, t: Tensor) ![]u8 { + if (t.dtype == .f32) { + const out = try L.staging(t.bytes.len); + @memcpy(out, try L.src.view(t.bytes)); + return L.host.items; + } if (t.dtype != .bf16) return error.UnexpectedTensor; const n = t.bytes.len / 2; const out = std.mem.bytesAsSlice(u32, try L.staging(n * 4)); @@ -375,13 +380,15 @@ pub const Loader = struct { m.conv_w = try L.alloc(try join(&a, name, ".conv_w"), L.host.items.len); try L.ops.upload(m.conv_w, L.host.items); m.conv_b = try L.widen(try join(&a, name, ".conv_b"), try ck.get(try join(&b, pre, "conv1d.bias"))); - const alog = try L.widened(try ck.get(try join(&b, pre, "A_log"))); + var fp: [128]u8 = undefined; + const ssm = try join(&fp, pre, L.w.config.ssmParams()); // Transcribe keeps these three in fp32 + const alog = try L.widened(try ck.get(try join(&b, ssm, "A_log"))); const src = try L.tmp(alog.len); try L.ops.upload(src, alog); m.a = try L.alloc(try join(&a, name, ".a"), alog.len); try L.ops.torch().mambaA(src, m.a, alog.len / 4); - m.d = try L.widen(try join(&a, name, ".d"), try ck.get(try join(&b, pre, "D"))); - m.dt_bias = try L.widen(try join(&a, name, ".dt_bias"), try ck.get(try join(&b, pre, "dt_bias"))); + m.d = try L.widen(try join(&a, name, ".d"), try ck.get(try join(&b, ssm, "D"))); + m.dt_bias = try L.widen(try join(&a, name, ".dt_bias"), try ck.get(try join(&b, ssm, "dt_bias"))); m.gnorm = try L.raw(try join(&a, name, ".gnorm"), try ck.get(try join(&b, pre, "norm.weight"))); return m; } diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 6bfb5ed1d0..2fac175e01 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -30,7 +30,7 @@ const vision_offer = @import("cuda_vision_offer.zig"); const drafted = @import("cuda_drafted.zig"); const kv_pages = @import("cuda_kv_pool.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native }; +const registry = .{ nemotron.native, nemotron.omni, nemotron.transcribe, glm.native, minimax.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; @@ -215,11 +215,11 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O problem.* = why; return null; } - // a family whose rope can stretch (Qwen's YaRN, from its config or --rope-yarn) names its own ceiling + // the window: a stretched rope's own ceiling (Qwen's YaRN), else a family's trained_context, else config.json's const native = if (comptime @hasDecl(F, "contextCeiling")) F.contextCeiling(a, io, o.dir, o.rope_yarn) catch |e| { problem.* = try std.fmt.allocPrint(a, "{s} ({s}): YaRN takes rope_type yarn, factor, original_max_position_embeddings and their product as max_position_embeddings, and --rope-yarn a factor from 1 matching it", .{ if (o.rope_yarn != null) "--rope-yarn" else "the checkpoint's rope scaling", @errorName(e) }); return null; - } else modelContext(a, io, o.dir); + } else if (comptime @hasDecl(F, "trained_context")) F.trained_context else modelContext(a, io, o.dir); const window = contextWindow(o.context, native, F.default_context) catch |e| { problem.* = switch (e) { error.Negative => try std.fmt.allocPrint(a, "--context {d}: a token count, or 0 for the model's window", .{o.context.?}), @@ -507,6 +507,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O .penalties = served.backend.penalties, .structures = served.backend.vtable.first_masked != null, .vision = if (@hasField(@TypeOf(loaded), "vision")) vision_offer.offer(loaded.vision, o.vision) else null, + .audio = if (@hasField(@TypeOf(loaded), "audio")) vision_offer.audio(loaded.audio) else null, .admit = admit, .slide_conversations = conversations, }, .{ .ctx = loaded.ctx, .text = F.explain }); diff --git a/zig/src/native/cuda_load_gpu.zig b/zig/src/native/cuda_load_gpu.zig index 011531c48b..23f31a93df 100644 --- a/zig/src/native/cuda_load_gpu.zig +++ b/zig/src/native/cuda_load_gpu.zig @@ -190,6 +190,7 @@ test "native Bolt load failure keeps destinations and kernel GPU parents through e.k.d = d; e.k.modelopt = null; e.k.ggml = null; + e.k.bf16x = null; e.k.triton = null; e.k.affine.split = null; e.k.latent = null; @@ -299,6 +300,7 @@ fn partialEngine(g: *Gpu) !*nemotron.Engine { e.k.d = d; e.k.modelopt = null; e.k.ggml = null; + e.k.bf16x = null; e.k.triton = null; e.k.affine.split = null; e.k.latent = null; diff --git a/zig/src/native/cuda_vision_offer.zig b/zig/src/native/cuda_vision_offer.zig index 70dea6b438..5a8c019897 100644 --- a/zig/src/native/cuda_vision_offer.zig +++ b/zig/src/native/cuda_vision_offer.zig @@ -238,3 +238,17 @@ test "actual Qwen offer keeps its image_url type and template-owned markers" { try std.testing.expectEqualSlices(u32, &q.markers, held.markers); try std.testing.expect(held.wrap == null and held.video == null); } + +/// Info.audio for a family's audio encoder (Nemotron 3.5 Transcribe's): its WAV prepare, locator and limits. +pub fn audio(x: ?*nemotron.audio.Audio) ?api.AudioOffer { + const a = x orelse return null; + if (a.encoder == null) return null; + const prepare = struct { + fn f(ctx: *anyopaque, al: Allocator, bytes: []const u8, max_tokens: u32) anyerror!api.PreparedAudio { + const t: *nemotron.audio.Audio = @ptrCast(@alignCast(ctx)); + const p = try t.prepare(al, bytes, max_tokens); + return .{ .samples = p.samples, .tokens = p.tokens, .hash = p.hash }; + } + }.f; + return .{ .ctx = a, .prepare = prepare, .locator = "<|audio|>", .audio_tokens = a.rows.limit, .max_audios = nemotron.audio.max_clips }; +} diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index 34f7e32984..0bc9d24354 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -52,7 +52,7 @@ fn countTokens(srv: *Server, cx: *Cx, chat: Value) errors.Refused!usize { const on = thinking.enable orelse srv.config.enable_thinking; cx.clear_thinking = thinking.clear; const tools = try tool_specs.active(cx, chat.get("tools"), chat.get("tool_choice")); - const msgs = try messages_mod.normalizeWith(cx, chat.get("messages"), "system", srv.needs_user_after_tool, .of(srv.info.vision)); + const msgs = try messages_mod.normalizeWith(cx, chat.get("messages"), "system", srv.needs_user_after_tool, .ofInfo(srv.info)); const ids_ = prompt.renderIds(srv, cx, msgs, tools, on, srv.effortFor(thinking.effort), true) catch |e| switch (e) { error.Refused => return cx.other(cx.message), else => |x| return x, diff --git a/zig/src/server/audio.zig b/zig/src/server/audio.zig new file mode 100644 index 0000000000..da6a2d7c97 --- /dev/null +++ b/zig/src/server/audio.zig @@ -0,0 +1,179 @@ +//! Audio input (Info.audio): input_audio and audio_url parts to the engine's encoder, each clip's locator to its rows. +const std = @import("std"); +const api = @import("engine_api"); +const json = @import("json"); +const errors = @import("errors.zig"); +const images = @import("images.zig"); +const image_key = @import("image_key.zig"); +const http_body = @import("http_body.zig"); +const Value = json.Value; +const Cx = errors.Cx; + +/// A request body at most when the engine takes audio: an hour of 16-bit 16 kHz WAV (115.2 MB) in base64, and its text. +pub const body_limit = 160 << 20; +/// A clip's encoded bytes at most (the body limit bounds them all). +pub const max_audio_bytes = body_limit; + +/// The body limit of a route that reads chat messages: an hour of audio when the engine takes it. +pub fn bodyLimit(info: api.Info) usize { + return if (info.audio != null) body_limit else http_body.limit; +} + +pub const Expanded = struct { + ids: []const u32, + history_len: usize, + audios: []const api.Audio, + cache_key: []const u32, +}; + +/// The part type normalization writes: an audio_url part with a data URL (vLLM's, and the model card's). +pub const part_type = "audio_url"; + +/// An input_audio or data-URL audio_url part as {"type": "audio_url", "audio_url": {"url"}}, as `expand` reads it. +pub fn part(cx: *Cx, p: Value, kind: []const u8) errors.Refused!Value { + var url: []const u8 = ""; + if (std.mem.eql(u8, kind, "input_audio")) { + const ia = p.get("input_audio") orelse return cx.refuse("an input_audio content part must contain input_audio"); + const data = if (ia == .object) ia.get("data") else null; + if (data == null or data.? != .string or data.?.string.len == 0) return cx.refuse("input_audio must contain base64 data"); + const format = ia.get("format"); + if (format) |f| if (f != .string or !std.ascii.eqlIgnoreCase(f.string, "wav")) return cx.refuse("input_audio format must be wav"); + url = try std.mem.concat(cx.a, u8, &.{ "data:audio/wav;base64,", data.?.string }); + } else { + const au = p.get("audio_url") orelse return cx.refuse("an audio_url content part must contain audio_url"); + const u = if (au == .string) au else if (au == .object) au.get("url") orelse Value.null else Value.null; + if (u != .string or u.string.len == 0) return cx.refuse("audio_url must contain a url string"); + if (!std.ascii.startsWithIgnoreCase(u.string, "data:")) return cx.refuse("audio requires a data URL (data:audio/wav;base64,...)"); + url = u.string; + } + for ([_][]const u8{ "image", "images", "image_url", "input_image", "video", "video_url" }) |k| if (json.truthyField(p, k)) return cx.refuse("audio parts cannot contain other media"); + const inner = try json.newObject(cx.a); + try inner.put(cx.a, "url", .{ .string = url }); + const out = try json.newObject(cx.a); + try out.put(cx.a, "type", .{ .string = part_type }); + try out.put(cx.a, "audio_url", .{ .object = inner }); + return .{ .object = out }; +} + +fn isAudio(p: Value) bool { + const kind = p.get("type") orelse return false; + return kind == .string and std.mem.eql(u8, kind.string, part_type); +} + +/// `messages` for the template: a user message with audio becomes vLLM's string form, a locator a clip, then its text. +pub fn forTemplate(a: std.mem.Allocator, messages: Value, locator: []const u8) !Value { + if (messages != .array) return messages; + const out = try a.dupe(Value, messages.array); + for (out) |*m| { + const content = m.get("content") orelse continue; + if (content != .array) continue; + var clips: usize = 0; + var text: std.ArrayList(u8) = .empty; + for (content.array) |p| { + if (isAudio(p)) { + clips += 1; + } else if (p.get("text")) |t| if (t == .string) try text.appendSlice(a, t.string); + } + if (clips == 0) continue; + var s: std.ArrayList(u8) = .empty; + for (0..clips) |_| { + try s.appendSlice(a, locator); + try s.append(a, '\n'); + } + try s.appendSlice(a, text.items); + const copy = try json.copyObject(a, m.object); + try copy.put(a, "content", .{ .string = s.items }); + m.* = .{ .object = copy }; + } + return .{ .array = out }; +} + +/// The prompt with each clip's placeholder expanded to its rows, its clips and its cache ids; null without audio. +pub fn expand(offer: ?api.AudioOffer, cx: *Cx, messages: Value, ids: []const u32, history_len: usize) errors.Refused!?Expanded { + const o = offer orelse return null; + var urls: std.ArrayList([]const u8) = .empty; + if (messages == .array) for (messages.array) |m| { + const content = m.get("content") orelse continue; + if (content != .array) continue; + for (content.array) |p| if (isAudio(p)) try urls.append(cx.a, p.get("audio_url").?.get("url").?.string); + }; + var marks: usize = 0; + for (ids) |t| marks += @intFromBool(t == o.audio_token); + const n = urls.items.len; + if (n == 0 and marks == 0) return null; + // the template writes a locator only for the ones content starts with: text that spells one adds a placeholder + if (marks > n) return cx.refuse("message text cannot spell the audio locator before a message's text"); + if (marks < n) return cx.fail(.server, "the chat template wrote {d} audio placeholders for {d} audio parts", .{ marks, n }); + if (n > o.max_audios) return cx.refuse(try std.fmt.allocPrint(cx.a, "a request may hold at most {d} audio clips; it has {d}", .{ o.max_audios, n })); + const budget: u32 = o.audio_tokens / @as(u32, @intCast(n)); + const prepared = try cx.a.alloc(api.PreparedAudio, n); + for (urls.items, prepared, 0..) |url, *p, i| { + const bytes = try images.dataBytesOf(cx, url, max_audio_bytes, "audio"); + p.* = o.prepare(o.ctx, cx.a, bytes, budget) catch |e| return refusal(cx, i, e, budget); + if (p.tokens == 0 or p.tokens > budget) return refusal(cx, i, error.AudioTooLong, budget); + } + var out_ids: std.ArrayList(u32) = .empty; + var key: std.ArrayList(u32) = .empty; + var clips: std.ArrayList(api.Audio) = .empty; + var hist = history_len; + var k: usize = 0; + for (ids, 0..) |t, src| { + const before = out_ids.items.len; + if (t == o.audio_token) { + const c = prepared[k]; + k += 1; + try clips.append(cx.a, .{ .samples = c.samples, .at = @intCast(out_ids.items.len), .tokens = c.tokens, .hash = c.hash }); + // the first two rows carry the clip's identity whole, the rest a hash of it and the row + for (0..c.tokens) |j| { + try out_ids.append(cx.a, t); + try key.append(cx.a, switch (j) { + 0 => image_key.open(c.hash), + 1 => image_key.close(c.hash), + else => image_key.row(c.hash, j), + }); + } + } else try images.both(cx.a, &out_ids, &key, t); + if (src < history_len) hist += out_ids.items.len - before - 1; + } + return .{ .ids = out_ids.items, .history_len = hist, .audios = clips.items, .cache_key = key.items }; +} + +fn refusal(cx: *Cx, i: usize, err: anyerror, budget: u32) errors.Refused { + return switch (err) { + error.OutOfMemory => error.OutOfMemory, + error.AudioTooLong => cx.refuse(try std.fmt.allocPrint(cx.a, "audio {d} is longer than this request's audio budget of {d} rows ({d:.0} seconds)", .{ i + 1, budget, @as(f64, @floatFromInt(budget)) * 0.08 })), + error.WavRate => cx.refuse(try std.fmt.allocPrint(cx.a, "audio {d} must be sampled at 16 kHz", .{i + 1})), + error.WavNotMono => cx.refuse(try std.fmt.allocPrint(cx.a, "audio {d} must be mono", .{i + 1})), + error.AudioTooShort => cx.refuse(try std.fmt.allocPrint(cx.a, "audio {d} is shorter than 10 ms", .{i + 1})), + else => cx.refuse(try std.fmt.allocPrint(cx.a, "audio {d} could not be decoded; send a WAV file (16-bit PCM or 32-bit float, 16 kHz, mono)", .{i + 1})), + }; +} + +fn fakePrepare(_: *anyopaque, a: std.mem.Allocator, bytes: []const u8, _: u32) anyerror!api.PreparedAudio { + // the test's clips: one row a byte, the first byte as their identity + return .{ .samples = try a.alloc(f32, 0), .tokens = @intCast(bytes.len), .hash = bytes[0] }; +} + +test "a clip's placeholder becomes its rows, keyed by the clip" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var cx: Cx = .{ .a = a }; + var dummy: u8 = 0; + const o: api.AudioOffer = .{ .ctx = &dummy, .prepare = fakePrepare, .audio_token = 9, .locator = "<|audio|>", .audio_tokens = 64, .max_audios = 2 }; + const raw = (try json.parse(a, "[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"Transcribe.\"},{\"type\":\"input_audio\",\"input_audio\":{\"data\":\"AQID\",\"format\":\"wav\"}}]}]")).ok; + const p = try part(&cx, raw.array[0].get("content").?.array[1], "input_audio"); + raw.array[0].get("content").?.array[1] = p; + // the template's view: the locator first + const t = try forTemplate(a, raw, "<|audio|>"); + try std.testing.expectEqualStrings("<|audio|>\nTranscribe.", t.array[0].get("content").?.string); + const x = (try expand(o, &cx, raw, &.{ 1, 2, 9, 3 }, 3)).?; + try std.testing.expectEqualSlices(u32, &.{ 1, 2, 9, 9, 9, 3 }, x.ids); + try std.testing.expectEqual(@as(usize, 1), x.audios.len); + try std.testing.expectEqual(@as(u32, 2), x.audios[0].at); + try std.testing.expectEqual(@as(u32, 3), x.audios[0].tokens); + try std.testing.expectEqual(@as(usize, 5), x.history_len); + try std.testing.expectEqual(image_key.open(1), x.cache_key[2]); + // a placeholder text spells, with no clip for it, is refused + try std.testing.expectError(error.Refused, expand(o, &cx, raw, &.{ 9, 9, 3 }, 0)); +} diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index fa11b0ae8d..6a20111caf 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -69,7 +69,6 @@ pub const Reply = struct { }; pub const Failure = error{ Refused, Cancelled, Failed, OutOfMemory }; - const nowNs = clock.nowNs; const seconds = clock.seconds; @@ -178,6 +177,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa .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, .images = rendered.images, + .audios = rendered.audios, .cache_key = rendered.cache_key, }; try srv.checkFeatures(cx, f, input.tools.len > 0, thinking, rendered.ids, input.tools, &request); diff --git a/zig/src/server/http_body.zig b/zig/src/server/http_body.zig index ac35079fa4..64f3064141 100644 --- a/zig/src/server/http_body.zig +++ b/zig/src/server/http_body.zig @@ -146,7 +146,12 @@ fn trailerName(line: []const u8) bool { /// The body as JSON (an empty body reads as ``{}``): a framing refusal, or the decoder's message as another exception. pub fn readJson(conn: *Conn, cx: *errors.Cx) errors.Refused!json.Value { - const raw = switch (try read(conn, cx.a, limit)) { + return readJsonMax(conn, cx, limit); +} + +/// readJson with its own limit (a route whose requests carry audio). +pub fn readJsonMax(conn: *Conn, cx: *errors.Cx, max: usize) errors.Refused!json.Value { + const raw = switch (try read(conn, cx.a, max)) { .ok => |b| b, .refused => |message| return cx.refuse(message), }; diff --git a/zig/src/server/messages.zig b/zig/src/server/messages.zig index 0f8a6bdc2a..9916f26eeb 100644 --- a/zig/src/server/messages.zig +++ b/zig/src/server/messages.zig @@ -4,6 +4,7 @@ const json = @import("json"); const errors = @import("errors.zig"); const fields = @import("fields.zig"); const api = @import("engine_api"); +const audio = @import("audio.zig"); const Value = json.Value; const Cx = errors.Cx; @@ -27,15 +28,23 @@ pub fn normalize(cx: *Cx, messages: ?Value, late_system: []const u8, needs_user_ return normalizeWith(cx, messages, late_system, needs_user_after_tool, .{}); } -/// The media parts an engine reads: image_url parts (``Info.vision``), and video_url parts (its ``video``). +/// The media parts an engine reads: images and videos (``Info.vision``), and input_audio or audio_url (``Info.audio``). pub const Media = struct { images: bool = false, videos: bool = false, + audio: bool = false, image_type: []const u8 = "image_url", // the part type the template reads (Vision.part_type) pub fn of(v: ?api.Vision) Media { return .{ .images = v != null, .videos = v != null and v.?.video != null, .image_type = if (v) |x| x.part_type else "image_url" }; } + + /// An engine's media: its images and videos, and its audio. + pub fn ofInfo(info: api.Info) Media { + var m = of(info.vision); + m.audio = info.audio != null; + return m; + } }; /// One video part as the template and video path read it: {"type": "video_url", "video_url": {"url"}, "video": null}. @@ -105,17 +114,23 @@ pub fn normalizeWith(cx: *Cx, messages: ?Value, late_system: []const u8, needs_u has_image = true; continue; } + if (media.audio and (std.mem.eql(u8, kind, "input_audio") or std.mem.eql(u8, kind, "audio_url"))) { + if (!std.mem.eql(u8, role.?.string, "user")) return cx.refuse("audio parts are accepted in user messages only"); + try parts.append(cx.a, try audio.part(cx, part, kind)); + has_image = true; + continue; + } if (media.videos and std.mem.eql(u8, kind, "video_url")) { if (!std.mem.eql(u8, role.?.string, "user")) return cx.refuse("video_url parts are supported only in user messages"); try parts.append(cx.a, try videoPart(cx, part)); has_image = true; continue; } - if (!std.mem.eql(u8, kind, "text") or fields.hasMedia(part)) return cx.refuse(if (media.videos) "content parts must be text, image_url or video_url; audio is unsupported" else if (images) "content parts must be text or image_url; audio and video inputs are unsupported" else "this server accepts text parts only; image, audio and video inputs are unsupported"); + if (!std.mem.eql(u8, kind, "text") or fields.hasMedia(part)) return cx.refuse(if (media.audio) "content parts must be text, input_audio or audio_url; image and video inputs are unsupported" else if (media.videos) "content parts must be text, image_url or video_url; audio is unsupported" else if (images) "content parts must be text or image_url; audio and video inputs are unsupported" else "this server accepts text parts only; image, audio and video inputs are unsupported"); const t = part.get("text") orelse return cx.refuse("a text content part must contain a text string"); if (t != .string) return cx.refuse("a text content part must contain a text string"); try text.appendSlice(cx.a, t.string); - if (images) { + if (images or media.audio) { const tp = try json.newObject(cx.a); try tp.put(cx.a, "type", .{ .string = "text" }); try tp.put(cx.a, "text", t); diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index 833defa91d..be4aeb42c5 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -44,7 +44,7 @@ pub const Out = struct { /// POST /v1/chat/completions or /v1/completions, read from the client and answered to it. pub fn post(srv: *Server, conn: *Conn, a: Allocator, is_chat: bool) void { var cx: Cx = .{ .a = a }; - const body = http_body.readJson(conn, &cx) catch { + const body = http_body.readJsonMax(conn, &cx, @import("audio.zig").bodyLimit(srv.info)) catch { const err = errorBody(a, &cx, if (is_chat) "messages" else "prompt") catch return; return routes.sendValue(conn, a, cx.status(), wrapError(a, err) catch return); }; @@ -134,10 +134,10 @@ 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.normalizeWith(cx, body.get("messages"), "system", srv.needs_user_after_tool, .of(srv.info.vision)); + input.messages = try messages.normalizeWith(cx, body.get("messages"), "system", srv.needs_user_after_tool, .ofInfo(srv.info)); 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.normalizeWith(cx, m, "system", srv.needs_user_after_tool, .of(srv.info.vision)); + input.messages = try messages.normalizeWith(cx, m, "system", srv.needs_user_after_tool, .ofInfo(srv.info)); }; if (!is_chat and input.messages.array.len == 0) input.prompt = try legacyPrompt(srv, cx, body.get("prompt")); // ``body.get("max_tokens") or body.get("max_completion_tokens")``: both are ints or None by now diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index d089854b56..baa767ae9c 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -7,12 +7,13 @@ const messages_mod = @import("messages.zig"); const model_text = @import("model_text.zig"); const chat = @import("chat.zig"); const images = @import("images.zig"); +const audio = @import("audio.zig"); const Server = @import("server.zig").Server; const Value = json.Value; const Cx = errors.Cx; /// With image input, `ids` holds each image's tokens, `images` the images and `cache_key` the prompt the cache matches. -pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0, user_end: usize = 0, images: []const api.Image = &.{}, cache_key: []const u32 = &.{} }; +pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0, user_end: usize = 0, images: []const api.Image = &.{}, audios: []const api.Audio = &.{}, cache_key: []const u32 = &.{} }; /// Where lessons decide to act on a reply: the row before the generation prompt's ``, else before its last. pub fn decideOf(srv: *const Server, ids: []const u32) u32 { @@ -30,7 +31,7 @@ pub const isTitle = messages_mod.isTitleRequest; /// ``render_prompt_ids``: messages normalized for this template, rendered, and a dangling ```` closed. pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, thinking: bool, effort: ?[]const u8, generation: bool) errors.Refused![]const u32 { - const normalized = try messages_mod.toolArguments(cx, try messages_mod.normalizeWith(cx, messages, srv.late_system, srv.needs_user_after_tool, .of(srv.info.vision))); + const normalized = try messages_mod.toolArguments(cx, try messages_mod.normalizeWith(cx, messages, srv.late_system, srv.needs_user_after_tool, .ofInfo(srv.info))); var problem: []const u8 = ""; const options: model_text.RenderOptions = .{ .tools = if (tools.len > 0) Value{ .array = @constCast(tools) } else null, @@ -62,7 +63,9 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t /// The template's ids for `messages`; with image input, a marker the text spells is ordinary text (images.zig). fn render(srv: *Server, a: std.mem.Allocator, messages: Value, options: model_text.RenderOptions, problem: *[]const u8) model_text.Error![]u32 { - var rendered = try srv.text.render(a, messages, options, problem); + // a speech model's template reads each clip as a locator before the message's text (audio.zig) + const view = if (srv.info.audio) |o| try audio.forTemplate(a, messages, o.locator) else messages; + var rendered = try srv.text.render(a, view, options, problem); if (srv.info.vision) |v| if (v.video) |o| if (o.video_text.len > 0) { const word = try srv.text.tokenString(a, o.video_token); rendered = try rendered.replaceTemplate(a, o.video_text, word); @@ -122,6 +125,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: chat.Input, thinking: bool, effort: if (history_len > 0) history_len = continuedLen(srv, cx, input.messages, input.tools, thinking, effort, history, history_len); // --vision: placeholders expanded to image tokens before the window, chunks and cache see it; no rewind kept if (try images.expand(srv.info.vision, if (srv.config.vision_urls) images.Fetch.network(srv.io) else null, srv.text, cx, input.messages, prompt, history_len)) |x| return .{ .ids = x.ids, .history_len = x.history_len, .images = x.images, .cache_key = x.cache_key }; + // audio input: each clip's placeholder expanded to its rows, the same way + if (try audio.expand(srv.info.audio, cx, input.messages, prompt, history_len)) |x| return .{ .ids = x.ids, .history_len = x.history_len, .audios = x.audios, .cache_key = x.cache_key }; const count = if (input.messages == .array) input.messages.array.len else 0; const rewind_len = if (srv.info.prompt_cache) rewindLen(srv, cx, input.messages, count, input.tools, prompt, thinking, effort) else 0; // one more render return .{ .ids = prompt, .history_len = history_len, .rewind_len = rewind_len, .user_end = userEnd(input.messages, history_len) }; diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 5dad7808f9..71fe010995 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -81,6 +81,7 @@ test { _ = @import("request_log.zig"); _ = @import("messages.zig"); _ = @import("images.zig"); + _ = @import("audio.zig"); _ = @import("images_video_marker_test.zig"); _ = @import("image_key.zig"); _ = @import("image_pieces.zig"); diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index d4a474bdcc..0e6e115a6a 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -177,6 +177,13 @@ pub const Server = struct { log.line("--vision: marker token {d} is not an added token in this tokenizer; image and video input cannot be served safely", .{id}); return error.InvalidVisionMarkers; }; + // audio input: the locator the template writes is one added token, the one a clip's rows take + if (srv.info.audio) |*o| { + o.audio_token = text.tokenId(o.locator) orelse { + log.line("audio input: the tokenizer has no single token {s}; audio cannot be served", .{o.locator}); + return error.InvalidAudioLocator; + }; + } if (config.voices_dir) |dir| { if (srv.info.speech == null) return error.VoicesRequireSpeech; srv.voices = try @import("drama_voice.zig").load(srv.arena.allocator(), io, dir); diff --git a/zig/src/server/tokens.zig b/zig/src/server/tokens.zig index d7ca9a3728..ad09f3dd8b 100644 --- a/zig/src/server/tokens.zig +++ b/zig/src/server/tokens.zig @@ -30,7 +30,7 @@ fn tokenize(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { var ids: []const u32 = undefined; if (body.has("messages")) { const generation = try flag(cx, body, "add_generation_prompt", true); - const msgs = try messages.normalizeWith(cx, body.get("messages"), "system", srv.needs_user_after_tool, .of(srv.info.vision)); + const msgs = try messages.normalizeWith(cx, body.get("messages"), "system", srv.needs_user_after_tool, .ofInfo(srv.info)); const tools = tool_specs.active(cx, body.get("tools"), body.get("tool_choice")) catch |e| switch (e) { error.Refused => return cx.refuse(cx.message), // a ValueError here is a RequestError else => |x| return x, @@ -44,6 +44,7 @@ fn tokenize(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { }; // media expanded as the chat route expands them: each image's tokens, each video's frame groups if (try images.expand(srv.info.vision, if (srv.config.vision_urls) images.Fetch.network(srv.io) else null, srv.text, cx, msgs, ids, 0)) |x| ids = x.ids; + if (try @import("audio.zig").expand(srv.info.audio, cx, msgs, ids, 0)) |x| ids = x.ids; } else { const text = body.get("prompt"); if (text == null or text.? != .string) return cx.refuse("prompt must be a string (or send messages)"); diff --git a/zig/src/server/warm.zig b/zig/src/server/warm.zig index fdd7a54361..bf307b315f 100644 --- a/zig/src/server/warm.zig +++ b/zig/src/server/warm.zig @@ -47,13 +47,19 @@ fn start(srv: *Server, input: chat.Input, message: Value, thinking: bool, effort // an image conversation's pass holds its images as its turns did (data URLs only: no URL is fetched twice) const pictures = try images.expand(srv.info.vision, null, srv.text, &cx, .{ .array = warm_prompt.messages }, prompt, prompt.len); if (pictures) |x| prompt = x.ids; + const clips = if (pictures == null) try @import("audio.zig").expand(srv.info.audio, &cx, .{ .array = warm_prompt.messages }, prompt, prompt.len) else null; + if (clips) |x| prompt = x.ids; if (prompt.len == 0 or (srv.info.context_window > 0 and prompt.len >= srv.info.context_window)) return; // its states count as the reply's request's: they supersede none of that turn's own until a later turn does - p.request = .{ .prompt = prompt, .max_tokens = 0, .background = true, .born = @intCast(prompt_len), .history_len = @intCast(prompt.len), .rewind_len = if (pictures != null) 0 else @intCast(prompt_mod.rewindLen(srv, &cx, .{ .array = warm_prompt.messages }, input.messages.array.len, input.tools, prompt, thinking, effort)), .chunks = try srv.chunks.starts(a, prompt), .decode_spans = try prompt_mod.replySpans(srv, &cx, prompt) }; + p.request = .{ .prompt = prompt, .max_tokens = 0, .background = true, .born = @intCast(prompt_len), .history_len = @intCast(prompt.len), .rewind_len = if (pictures != null or clips != null) 0 else @intCast(prompt_mod.rewindLen(srv, &cx, .{ .array = warm_prompt.messages }, input.messages.array.len, input.tools, prompt, thinking, effort)), .chunks = try srv.chunks.starts(a, prompt), .decode_spans = try prompt_mod.replySpans(srv, &cx, prompt) }; if (pictures) |x| { p.request.images = x.images; p.request.cache_key = x.cache_key; } + if (clips) |x| { + p.request.audios = x.audios; + p.request.cache_key = x.cache_key; + } const id = srv.next_id.fetchAdd(1, .monotonic); try srv.engine.submit(id, &p.request, .{ .ctx = p, .event = Pass.event }); handed = true; diff --git a/zig/tests/cuda/bf16_experts_tests.zig b/zig/tests/cuda/bf16_experts_tests.zig new file mode 100644 index 0000000000..278aae3e58 --- /dev/null +++ b/zig/tests/cuda/bf16_experts_tests.zig @@ -0,0 +1,201 @@ +//! Grouped bf16 experts against float64: the repack, up and downs, a row's bits the same at any row count and form. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const grouped = cuda.grouped; +const ex_ = cuda.experts; +const Bf16 = ex_.Bf16Experts; + +const E = 6; // routed experts; the shared expert's halves are experts E and E + 1 +const D = 256; +const W = 96; +const K = 3; // routed picks a row; a row has K + 2 pairs + +fn bf(x: f32) u16 { + const b: u32 = @bitCast(x); + return @intCast((b + 0x7FFF + ((b >> 16) & 1)) >> 16); +} + +fn f(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Host = struct { + routed_up: []u16, // [E][D][W]: x @ W, as the stacked checkpoint tensors + routed_down: []u16, // [E][W][D] + shared_up: []u16, // [2W][D]: nn.Linear + shared_down: []u16, // [D][2W] + + /// Expert e's weight from input i to output o, up or down, as the kernels must read it. + fn at(h: Host, up: bool, e: usize, o: usize, i: usize) f64 { + if (e < E) return f(if (up) h.routed_up[(e * D + i) * W + o] else h.routed_down[(e * W + i) * D + o]); + const half = e - E; + return f(if (up) h.shared_up[(half * W + o) * D + i] else h.shared_down[o * 2 * W + half * W + i]); + } +}; + +pub fn experts(gpu: Gpu) !void { + const a = gpu.gpa; + var rng = std.Random.DefaultPrng.init(0x7f16); + const r = rng.random(); + var host: Host = undefined; + inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }, .{ E * D * W, E * W * D, 2 * W * D, D * 2 * W }) |name, n| { + @field(host, name) = try a.alloc(u16, n); + for (@field(host, name)) |*v| v.* = bf((r.float(f32) - 0.5) * 0.25); + } + defer inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }) |name| a.free(@field(host, name)); + + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + var kx = try Bf16.load(gpu.d, sms); + defer kx.unload(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try grouped.Router.resolve(mod); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + + const e = E + 2; + var up = try cuda.DeviceBuffer.alloc(gpu.d, Bf16.bytes(e, W, D)); + defer up.free(); + var down = try cuda.DeviceBuffer.alloc(gpu.d, Bf16.bytes(e, D, W)); + defer down.free(); + var src: [4]cuda.DeviceBuffer = undefined; + inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }, 0..) |name, i| src[i] = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(@field(host, name))); + defer for (&src) |*b| b.free(); + const one_up = Bf16.bytes(1, W, D); + const one_down = Bf16.bytes(1, D, W); + try kx.pack(s, src[0].ptr, .{ .stride = D * W, .ld = W, .in_out = true }, up.ptr, E, W, D); + try kx.pack(s, src[1].ptr, .{ .stride = W * D, .ld = D, .in_out = true }, down.ptr, E, D, W); + for (0..2) |h| { + try kx.pack(s, src[2].ptr, .{ .ld = D, .n0 = h * W, .in_out = false }, up.ptr + (E + h) * one_up, 1, W, D); + try kx.pack(s, src[3].ptr, .{ .ld = 2 * W, .k0 = h * W, .in_out = false }, down.ptr + (E + h) * one_down, 1, D, W); + } + try s.synchronize(); + const layer: ex_.Layer = .{ .format = .bf16, .up = up.ptr, .down = down.ptr, .width = W, .dims = D, .experts = e, .relu2 = true }; + const g: ex_.Grouped = .{ .bf16 = &kx }; + + // 37 rows: three 16-pair decode items for the busiest expert and a partial prompt tile; then each row alone + const rows = 37; + const slots = K + 2; + const x = try a.alloc(u16, rows * D); + defer a.free(x); + for (x) |*v| v.* = bf(r.float(f32) * 2 - 1); + const picks = try a.alloc(i32, rows * slots); + defer a.free(picks); + for (0..rows) |row| { + var chosen: [K]i32 = undefined; + for (0..K) |k| { + var p: i32 = undefined; + while (true) { + p = @intCast(r.uintLessThan(usize, E)); + if (std.mem.indexOfScalar(i32, chosen[0..k], p) == null) break; + } + chosen[k] = p; + picks[row * slots + k] = p; + } + picks[row * slots + K] = E; + picks[row * slots + K + 1] = E + 1; + } + + // the float64 reference: act = bf16(relu(bf16(x . up))^2), y = act . down + const act_ref = try a.alloc(f64, rows * slots * W); + defer a.free(act_ref); + const y_ref = try a.alloc(f64, rows * slots * D); + defer a.free(y_ref); + const mag = try a.alloc(f64, rows * slots * D); + defer a.free(mag); + for (0..rows * slots) |pr| { + const ex: usize = @intCast(picks[pr]); + const row = pr / slots; + for (0..W) |o| { + var acc: f64 = 0; + for (0..D) |i| acc += f(x[row * D + i]) * host.at(true, ex, o, i); + const u = @max(f(bf(@floatCast(acc))), 0); + act_ref[pr * W + o] = f(bf(@floatCast(u * u))); + } + for (0..D) |o| { + var acc: f64 = 0; + var m: f64 = 0; + for (0..W) |i| { + const t = act_ref[pr * W + i] * host.at(false, ex, o, i); + acc += t; + m += @abs(t); + } + y_ref[pr * D + o] = acc; + mag[pr * D + o] = m; + } + } + + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(x)); + defer dx.free(); + var whole: [3][]u8 = undefined; // decode form, all rows: act, y fp32, y bf16 + var got = [3]bool{ false, false, false }; + defer for (whole, got) |w, ok| if (ok) a.free(w); + for ([_]bool{ false, true }) |prompt| { + const run = try oneCase(gpu, router, g, s, layer, dx.ptr, picks, rows, slots, prompt); + defer for (run) |b| a.free(b); + const act_h = std.mem.bytesAsSlice(u16, run[0]); + var worst: f64 = 0; + for (act_h, act_ref) |gv, want| worst = @max(worst, @abs(f(gv) - want) / @max(@abs(want), 1e-3)); + try check.expect(worst <= 1.0 / 64.0, "bf16 experts {s}: up off by {e} (relative)", .{ if (prompt) "prompt" else "decode", worst }); + const y32 = std.mem.bytesAsSlice(f32, run[1]); + var worst_y: f64 = 0; + for (y32, y_ref, mag) |gv, want, m| worst_y = @max(worst_y, @abs(@as(f64, gv) - want) / @max(m, 1e-6)); + try check.expect(worst_y <= 1e-5, "bf16 experts {s}: fp32 down off by {e} of its terms' magnitude", .{ if (prompt) "prompt" else "decode", worst_y }); + if (!prompt) { + for (run, &whole, &got) |b, *w, *ok| { + w.* = try a.dupe(u8, b); + ok.* = true; + } + } else for (run, whole, [_][]const u8{ "prompt up", "prompt fp32 down", "prompt bf16 down" }) |b, w, what| try check.sameBytes(what, b, w); + check.pass("bf16 experts: {d} rows x {d} pairs ({s}) up within {e}, fp32 down within {e} of the float64 reference", .{ rows, slots, if (prompt) "prompt, items of 64" else "decode, items of 16", worst, worst_y }); + } + // each row alone gives the bytes it gave beside the others + for ([_]usize{ 0, 17, rows - 1 }) |row| { + const run = try oneCase(gpu, router, g, s, layer, dx.ptr + row * D * 2, picks[row * slots ..][0..slots], 1, slots, false); + defer for (run) |b| a.free(b); + for (run, whole, [_]usize{ W * 2, D * 4, D * 2 }, [_][]const u8{ "row up", "row fp32 down", "row bf16 down" }) |b, w, per, what| try check.sameBytes(what, b, w[row * slots * per ..][0 .. slots * per]); + } + check.pass("bf16 experts: a row alone, beside 36 others and in the prompt form gives the same bytes", .{}); +} + +/// Plan, up and both downs over `rows` rows; returns act, fp32 y and bf16 y. +fn oneCase(gpu: Gpu, router: grouped.Router, g: ex_.Grouped, s: cuda.Stream, l: ex_.Layer, x: u64, picks: []const i32, rows: usize, slots: usize, prompt: bool) ![3][]u8 { + const pairs = rows * slots; + const tile: usize = if (prompt) ex_.promptTile(.bf16) else 16; + var dpicks = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer dpicks.free(); + const bound = grouped.maxItems(pairs, l.experts, tile); + var bufs: [5]cuda.DeviceBuffer = undefined; + const big = pairs > grouped.small; + const sizes = [5]usize{ pairs * 4, bound * 12, 8, (if (big) pairs else 1) * 4, (if (big) (pairs + 1023) / 1024 * l.experts else 1) * 4 }; + var made: usize = 0; + defer for (bufs[0..made]) |*b| b.free(); + for (sizes, &bufs) |n, *b| { + b.* = try cuda.DeviceBuffer.alloc(gpu.d, n); + made += 1; + } + const plan: grouped.Plan = .{ .members = bufs[0].ptr, .items = bufs[1].ptr, .counts = bufs[2].ptr, .rank = bufs[3].ptr, .hist = bufs[4].ptr }; + try router.route(s, dpicks.ptr, pairs, l.experts, tile, plan); + var act = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.width * 2); + defer act.free(); + var y32 = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.dims * 4); + defer y32.free(); + var y16 = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.dims * 2); + defer y16.free(); + for ([_]*cuda.DeviceBuffer{ &act, &y32, &y16 }) |b| try b.fill8(0xff, s.handle); // a call that writes nothing fails + const call = if (prompt) &ex_.Grouped.prompt else &ex_.Grouped.decode; + try call(g, s, .up, .{ .x = x, .stride = l.dims, .slots = slots }, l, plan, act.ptr, bound, -1); + try call(g, s, .down_f32, .{ .x = act.ptr, .stride = l.width }, l, plan, y32.ptr, bound, -1); + try call(g, s, .down_bf16, .{ .x = act.ptr, .stride = l.width }, l, plan, y16.ptr, bound, -1); + try s.synchronize(); + var out: [3][]u8 = undefined; + out[0] = try check.download(gpu, act); + errdefer gpu.gpa.free(out[0]); + out[1] = try check.download(gpu, y32); + errdefer gpu.gpa.free(out[1]); + out[2] = try check.download(gpu, y16); + return out; +} diff --git a/zig/tests/cuda/contract_experts_bf16.zig b/zig/tests/cuda/contract_experts_bf16.zig new file mode 100644 index 0000000000..fc27d1af07 --- /dev/null +++ b/zig/tests/cuda/contract_experts_bf16.zig @@ -0,0 +1,319 @@ +//! bf16 experts (experts_bf16.cu) under the contract on Transcribe's shapes, decode and prompt forms; two controls. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const grouped = cuda.grouped; +const ex_ = cuda.experts; +const Bf16 = ex_.Bf16Experts; + +pub const Kind = enum { up, down_f32, down_bf16 }; + +/// A planted fault: the fp32 down launched one 32-input group short, or contract_controls.cu's bf16-summing down. +pub const Fault = enum { none, short, bf16_sums }; + +pub const Entry = struct { name: []const u8, kind: Kind, prompt: bool, fault: Fault = .none }; + +pub const entries = [_]Entry{ + .{ .name = "cuda bf16 experts up tf_bf16_experts_up (relu^2)", .kind = .up, .prompt = false }, + .{ .name = "cuda bf16 experts down tf_bf16_experts_down_f32", .kind = .down_f32, .prompt = false }, + .{ .name = "cuda bf16 experts down tf_bf16_experts_down_bf16", .kind = .down_bf16, .prompt = false }, + .{ .name = "cuda bf16 experts prompt up tf_bf16_experts_prompt_up (relu^2)", .kind = .up, .prompt = true }, + .{ .name = "cuda bf16 experts prompt down tf_bf16_experts_prompt_down_f32", .kind = .down_f32, .prompt = true }, + .{ .name = "cuda bf16 experts prompt down tf_bf16_experts_prompt_down_bf16", .kind = .down_bf16, .prompt = true }, + .{ .name = "control: bf16 experts fp32 down launched one 32-input group short", .kind = .down_f32, .prompt = false, .fault = .short }, + .{ .name = "control: bf16 experts fp32 down summing in bf16", .kind = .down_f32, .prompt = false, .fault = .bf16_sums }, +}; + +const experts = 4; +const slots = 8; // pairs a token row: Transcribe's 6 routed picks and its shared expert's two halves + +/// An mma k16 step a link of each row's chain, plus one step's error over its own products (attention_range's 36). +fn depth(s: c.Shape) usize { + return s.k / 16 + core.attention_range.mma_step; +} + +/// Transcribe's routed experts, up 2,688 to 1,856 and back; decode windows of 1 to 32 rows, prompt items of 2 and 9. +const ups = [2][]const c.Shape{ &.{ .{ .rows = 1, .n = 1856, .k = 2688 }, .{ .rows = 32, .n = 1856, .k = 2688 } }, &.{ .{ .rows = 40, .n = 1856, .k = 2688 }, .{ .rows = 72, .n = 1856, .k = 2688 } } }; +const downs = [2][]const c.Shape{ &.{ .{ .rows = 1, .n = 2688, .k = 1856 }, .{ .rows = 32, .n = 2688, .k = 1856 } }, &.{ .{ .rows = 40, .n = 2688, .k = 1856 }, .{ .rows = 72, .n = 2688, .k = 1856 } } }; + +pub fn spec(e: Entry) c.Spec { + const served = (if (e.kind == .up) ups else downs)[@intFromBool(e.prompt)]; + const tiles = [3]c.Tile{ .{ .size = (if (e.prompt) ex_.promptTile(.bf16) else 16) / slots }, .{ .size = 64, .step = 32, .min = 32 }, .{ .size = 64, .step = 32, .min = 32 } }; + return switch (e.kind) { + // bf16(acc)'s rounding enters the square twice + .up => .{ .name = e.name, .accum = .f32, .depth = depth, .ops = 1, .mids = 2, .out = .bf16, .tiles = tiles, .served = served, .sums = .{ .f16 = 30, .bf16 = 30, .why = "a sum that collapses below 0 gives relu^2 0, which misses by want itself, and the bf16 roundings already allow 3 u of it: at most 1 / (3 u) = 85x is provable" } }, + .down_f32 => .{ .name = e.name, .accum = .f32, .depth = depth, .out = .f32, .tiles = tiles, .served = served }, + .down_bf16 => .{ .name = e.name, .accum = .f32, .depth = depth, .out = .bf16, .tiles = tiles, .served = served }, + }; +} + +/// shape.rows token rows, `slots` pairs each over `experts` experts; up reads a token's row, down a pair's. +pub const Experts = struct { + rig: rig_.Rig, + entry: Entry, + kx: *const Bf16, + router: grouped.Router, + ctl: ?cuda.Function = null, + shape: c.Shape = undefined, + x: []u16 = &.{}, + w: []u16 = &.{}, // [experts][n][k], as nn.Linear stores them + picks: []i32 = &.{}, + xd: u64 = 0, + wd: u64 = 0, + plan: grouped.Plan = undefined, + out: u64 = 0, + + pub fn deinit(self: *Experts) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Experts) void { + const a = self.rig.gpu.gpa; + a.free(self.x); + a.free(self.w); + a.free(self.picks); + self.x = &.{}; + self.w = &.{}; + self.picks = &.{}; + } + + pub fn kernel(self: *Experts) g.Kernel { + const control: ?g.Gate = switch (self.entry.fault) { + .none => null, + .short => .shapes, + .bf16_sums => .sums, + }; + return .{ .spec = spec(self.entry), .ctx = self, .vt = &vt, .control = control }; + } + + const vt = g.hooks(Experts); + + fn up(self: *const Experts) bool { + return self.entry.kind == .up; + } + + fn pairs(self: *const Experts) usize { + return self.shape.rows * slots; + } + + fn xRows(self: *const Experts) usize { + return if (self.up()) self.shape.rows else self.pairs(); + } + + fn tile(self: *const Experts) usize { + return if (self.entry.prompt) ex_.promptTile(.bf16) else 16; + } + + fn items(self: *const Experts) usize { + return grouped.maxItems(self.pairs(), experts, self.tile()); + } + + /// Random bf16 rows and weights; non-finite: NaN, +Inf and -Inf in the first rows; stress: mirrored integers. + pub fn prepare(self: *Experts, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 32 != 0 or shape.k == 0 or shape.k % 32 != 0) return false; + self.free(); + self.rig.reset(); + self.shape = shape; + const a = self.rig.gpu.gpa; + const n = shape.n; + const k = shape.k; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.x = try a.alloc(u16, self.xRows() * k); + self.w = try a.alloc(u16, experts * n * k); + self.picks = try a.alloc(i32, self.pairs()); + for (self.x) |*v| v.* = c.toBf16(rng.float(f32) * 2 - 1); + for (self.w) |*v| v.* = c.toBf16((rng.float(f32) - 0.5) * 0.25); + for (self.picks, 0..) |*p, i| p.* = @intCast(if (i < experts) i else rng.uintLessThan(usize, experts)); + switch (fill) { + .random => {}, + .nonfinite => { + // a row holding NaN, one holding +Inf, one holding both infinities + const planted = [_]struct { usize, usize, f32 }{ .{ 0, k / 3, std.math.nan(f32) }, .{ 1, k / 2, std.math.inf(f32) }, .{ 2, 1, std.math.inf(f32) }, .{ 2, k - 1, -std.math.inf(f32) } }; + for (planted) |p| if (p[0] < self.xRows()) { + self.x[p[0] * k + p[1]] = c.toBf16(p[2]); + }; + }, + .stress => self.stress(rng), + } + self.xd = try self.rig.input(std.mem.sliceAsBytes(self.x)); + const src = try self.rig.scratch(self.w.len * 2); + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = src, .len = self.w.len * 2 }).upload(0, std.mem.sliceAsBytes(self.w)); + const blocks = Bf16.bytes(experts, n, k); + const zeros = try a.alloc(u8, @max(blocks, self.pairs() * 4, self.items() * 12)); + defer a.free(zeros); + @memset(zeros, 0); + self.wd = try self.rig.input(zeros[0..blocks]); + try self.kx.pack(self.rig.s, src, .{ .stride = n * k, .ld = k, .in_out = false }, self.wd, experts, n, k); + const ids = try self.rig.input(std.mem.sliceAsBytes(self.picks)); + const hist = (self.pairs() + 1023) / 1024; + self.plan = .{ .members = try self.rig.input(zeros[0 .. self.pairs() * 4]), .items = try self.rig.input(zeros[0 .. self.items() * 12]), .counts = try self.rig.input(zeros[0..8]), .rank = try self.rig.scratch(self.pairs() * 4), .hist = try self.rig.scratch(hist * experts * 4) }; + try self.router.route(self.rig.s, ids, self.pairs(), experts, self.tile(), self.plan); + try self.rig.s.synchronize(); + self.out = try self.rig.output(self.pairs() * n * c.size(spec(self.entry).out)); + return true; + } + + /// Gate 4: mirrored integer weights and inputs, every partial sum an exact integer under 2^20, the row's sum small. + fn stress(self: *Experts, rng: std.Random) void { + const k = self.shape.k; + const h = k / 2; + for (0..experts * self.shape.n) |j| for (0..h) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 2)); + self.w[j * k + i] = c.toBf16(v); + self.w[j * k + h + i] = c.toBf16(v); + }; + for (0..self.xRows()) |r| for (0..h) |i| { + const v: f32 = @floatFromInt(4 * (1 + 2 * rng.uintLessThan(u32, 64))); + self.x[r * k + i] = c.toBf16(v); + self.x[r * k + h + i] = c.toBf16(if (i < 8) 4 - v else -v); + }; + } + + pub fn run(self: *Experts) !void { + const n = self.shape.n; + const k = self.shape.k; + const in: ex_.Rows = .{ .x = self.xd, .stride = k, .slots = if (self.up()) slots else 0 }; + switch (self.entry.fault) { + .none => { + const epi: ex_.Epi = switch (self.entry.kind) { + .up => .up, + .down_f32 => .down_f32, + .down_bf16 => .down_bf16, + }; + const l: ex_.Layer = .{ .format = .bf16, .up = self.wd, .down = self.wd, .width = if (self.up()) n else k, .dims = if (self.up()) k else n, .experts = experts, .relu2 = true }; + try self.kx.run(self.rig.s, self.entry.prompt, epi, in, l, self.plan, self.out, self.items(), -1); + }, + .short => try self.launch(self.kx.fns[0][1], k / 32 - 1, .{ .x = @intCast((self.items() * n / Bf16.cols + Bf16.warps - 1) / Bf16.warps) }, Bf16.warps * 32), + .bf16_sums => try self.launch(self.ctl.?, k / 32, .{ .x = @intCast((n + 127) / 128), .y = @intCast(self.items()) }, 128), + } + try self.rig.s.synchronize(); + } + + /// A down kernel's launch with its own input groups: experts_bf16.cu's arguments, less `slots` for the control. + fn launch(self: *Experts, f: cuda.Function, kg: usize, grid: struct { x: u32, y: u32 = 1 }, threads: u32) !void { + var a: cuda.Args = .{}; + a.add(self.xd); + a.add(@as(i32, @intCast(self.shape.k))); + if (self.entry.fault == .short) a.add(@as(i32, 0)); + a.add(self.wd); + a.add(@as(i32, @intCast(kg))); + a.add(@as(i32, @intCast(self.shape.n / Bf16.cols))); + for ([_]u64{ self.plan.items, self.plan.counts, self.plan.members, self.out }) |v| a.add(v); + a.add(@as(i32, @intCast(self.shape.n))); + if (self.entry.fault == .short) a.add(@as(i32, -1)); + try cuda.launch.launch(f, .{ .grid = .{ .x = grid.x, .y = grid.y, .z = 1 }, .block = .{ .x = threads } }, self.rig.s, &a); + } + + pub fn count(self: *Experts) usize { + return self.pairs() * self.shape.n; + } + + /// Pair p's products with column j of its expert, in input order. + fn terms(self: *const Experts, p: usize, j: usize, t: []f64) void { + const k = self.shape.k; + const xr = (if (self.up()) p / slots else p) * k; + const wr = (@as(usize, @intCast(self.picks[p])) * self.shape.n + j) * k; + for (t, 0..) |*o, i| o.* = c.fromBf16(self.x[xr + i]) * c.fromBf16(self.w[wr + i]); + } + + /// relu^2 for up, as float64 keeps a NaN; down's sum as it is. + fn epilogue(self: *const Experts, sum: f64) f64 { + if (!self.up()) return sum; + const u = if (std.math.isNan(sum)) sum else @max(sum, 0); + return u * u; + } + + pub fn reference(self: *Experts, want: []f64, mag: []f64) !void { + const t = try self.rig.gpu.gpa.alloc(f64, self.shape.k); + defer self.rig.gpu.gpa.free(t); + for (0..self.pairs()) |p| for (0..self.shape.n) |j| { + self.terms(p, j, t); + var sum: f64 = 0; + var abs: f64 = 0; + for (t) |v| { + sum += v; + abs += @abs(v); + } + const o = p * self.shape.n + j; + want[o] = self.epilogue(sum); + // a sum's error reaches relu^2 through its derivative 2 u + mag[o] = if (self.up()) 2 * @sqrt(want[o]) * abs else abs; + }; + } + + pub fn simulate(self: *Experts, acc: c.Accum, order: g.Order, out: []f64) !void { + const t = try self.rig.gpu.gpa.alloc(f64, self.shape.k); + defer self.rig.gpu.gpa.free(t); + for (0..self.pairs()) |p| for (0..self.shape.n) |j| { + self.terms(p, j, t); + out[p * self.shape.n + j] = self.epilogue(contract.rows.sumIn(acc, order, t)); + }; + } +}; + +/// Every registration, on the heap: the gates hold pointers to them. +pub const Set = struct { + gpa: std.mem.Allocator, + kx: Bf16, + mod: cuda.Module, + ctl: cuda.Module, + items: [entries.len]Experts = undefined, + + const controls = struct { + pub const bytes align(16) = @embedFile("fatbin_contract_controls").*; + }; + + /// The set, its kernels appended to `list` (each control after the kernels it breaks). + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + var kx = try Bf16.load(gpu.d, sms); + errdefer kx.unload(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + errdefer mod.unload(); + var ctl = try cuda.Module.load(gpu.d, &controls.bytes); + errdefer ctl.unload(); + set.* = .{ .gpa = gpu.gpa, .kx = kx, .mod = mod, .ctl = ctl }; + const router = try grouped.Router.resolve(set.mod); + const sums = try set.ctl.function("tf_ctl_bf16_experts_bf16_sums"); + for (&set.items, entries) |*x, e| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .entry = e, .kx = &set.kx, .router = router, .ctl = if (e.fault == .bf16_sums) sums else null }; + for (&set.items) |*x| try list.append(a, x.kernel()); + return set; + } + + pub fn deinit(set: *Set) void { + for (&set.items) |*x| x.deinit(); + set.ctl.unload(); + set.mod.unload(); + set.kx.unload(); + set.gpa.destroy(set); + } +}; + +/// The bf16 experts' gates 1-4 and their controls alone. +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(a); + const set = try Set.init(gpu, s, a, &list); + defer set.deinit(); + var t: g.Tally = .{}; + try g.runAll(a, list.items, 0x62663136, &t); + std.debug.print("SUMMARY bf16 experts: {d} kernels, {d} pass gates 1-4, {d} fail, {d} known findings; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.known, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "experts-bf16-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} + +test "every control is named one, and only the controls plant a fault" { + for (entries) |e| try std.testing.expectEqual(e.fault != .none, std.mem.startsWith(u8, e.name, "control: ")); +} diff --git a/zig/tests/cuda/contract_family_sets.zig b/zig/tests/cuda/contract_family_sets.zig index a602175cb3..d412850995 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -1,4 +1,4 @@ -//! The family sets in the full kernel contract: Nemotron's pocket, MiniMax-M3's, MOSS-Image's, GB10's grouped affine-4. +//! Full-contract family sets: Transcribe audio, bf16 experts, Nemotron's pocket, MiniMax-M3, MOSS-Image, GB10 affine-4. const std = @import("std"); const cuda = @import("cuda"); const check = @import("check.zig"); @@ -8,6 +8,8 @@ const Grouped = @import("contract_gb10_many.zig").Registrations; pub const Sets = struct { gpa: std.mem.Allocator, + nemo_audio: *@import("contract_nemo_audio.zig").Set, + bf16_experts: *@import("contract_experts_bf16.zig").Set, pocket: *@import("contract_pocket.zig").Set, minimax: *@import("contract_minimax.zig").Set, moss: *@import("contract_moss_image.zig").Set, @@ -15,6 +17,10 @@ pub const Sets = struct { /// Each set's kernels appended to `list`, in this order. pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel), affine4: *cuda.qlinear.Affine4Kernels, cap: u32) !Sets { + const nemo_audio = try @import("contract_nemo_audio.zig").Set.init(gpu, s, a, list); + errdefer nemo_audio.deinit(); + const bf16_experts = try @import("contract_experts_bf16.zig").Set.init(gpu, s, a, list); + errdefer bf16_experts.deinit(); const pocket = try @import("contract_pocket.zig").Set.init(gpu, s, a, list); errdefer pocket.deinit(); const minimax = try @import("contract_minimax.zig").Set.init(gpu, s, a, list); @@ -26,7 +32,7 @@ pub const Sets = struct { grouped.* = Grouped.init(gpu, s, affine4); errdefer grouped.deinit(); try grouped.append(cap, a, list); - return .{ .gpa = gpu.gpa, .pocket = pocket, .minimax = minimax, .moss = moss, .grouped = grouped }; + return .{ .gpa = gpu.gpa, .nemo_audio = nemo_audio, .bf16_experts = bf16_experts, .pocket = pocket, .minimax = minimax, .moss = moss, .grouped = grouped }; } pub fn deinit(sets: Sets) void { @@ -35,5 +41,7 @@ pub const Sets = struct { sets.moss.deinit(); sets.minimax.deinit(); sets.pocket.deinit(); + sets.bf16_experts.deinit(); + sets.nemo_audio.deinit(); } }; diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index ca9e662d17..e47e9fc9b8 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -26,7 +26,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "experts_bf16", "nemo_audio", "moss_image" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); return error.UnknownContractImage; diff --git a/zig/tests/cuda/contract_nemo_audio.zig b/zig/tests/cuda/contract_nemo_audio.zig new file mode 100644 index 0000000000..b13e79f3b4 --- /dev/null +++ b/zig/tests/cuda/contract_nemo_audio.zig @@ -0,0 +1,507 @@ +//! nemo_audio.cu under the contract (nemo-audio-contract): the audio encoders' kernels, with controls. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const na = @import("nemotron").nemo_audio.kernels; +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const rows_ = core.contract.rows; +const ar = core.attention_range; +const bfr = core.bf16_round; + +const u: f64 = 0x1p-24; +const half_bf16: f64 = 0x1p-8; +const floor: f64 = 0x1p-133; +const eps: f32 = 1e-5; + +pub const controls = struct { + pub const bytes align(16) = @embedFile("fatbin_contract_controls").*; +}; + +/// The ASR encoder's 600-row windows (10 heads of 128, RoPE on 64) and the Sortformer's (to 828 rows, 8 heads of 64). +const asr = .{ .heads = 10, .d = 128, .rot = 64 }; +const sf = .{ .heads = 8, .d = 64, .rot = 64 }; + +pub fn val(h: u16) f64 { + return bfr.value(h); +} + +fn gamma(n: f64) f64 { + return n * u / (1 - n * u); +} + +pub fn randomBf16(rng: std.Random, out: []u16, lo: f32, hi: f32) void { + for (out) |*v| v.* = bfr.bits(lo + rng.float(f32) * (hi - lo)); +} + +pub const planted = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + +// --- tf_na_qk_rope: q and k LayerNormed (one bf16 rounding, fp32 inside), then bf16(bf16(x c) + bf16(rot(x) s)) ---- + +pub const Rope = struct { + rig: rig_.Rig, + k: *const na.Kernels, + name: []const u8, + heads: usize, + rot: usize, + norm: bool, + ctl: ?cuda.Function = null, // the one-pass variance control: tf_na_qk_rope's arguments + control: ?g.Gate = null, + arena: ?std.heap.ArenaAllocator = null, + t: usize = 0, + d: usize = 0, + r: usize = 0, + x: []u16 = &.{}, + w: [2][]u16 = .{ &.{}, &.{} }, + b: [2][]u16 = .{ &.{}, &.{} }, + cs: []u16 = &.{}, + sn: []u16 = &.{}, + /// Each output's allowance, from reference(). + allow: []f64 = &.{}, + dev: struct { qkv: u64 = 0, w: [2]u64 = .{ 0, 0 }, b: [2]u64 = .{ 0, 0 }, cs: u64 = 0, sn: u64 = 0 } = .{}, + + fn depth(s: c.Shape) usize { + return (s.k + 31) / 32 + 5; + } + + pub fn spec(t: *const Rope) c.Spec { + return .{ + .name = t.name, + .accum = if (t.norm) .f32 else .none, + .depth = if (t.norm) depth else noSums, + .ops = 12, + .mids = 3, + .out = .bf16, + .in_place = true, + .tiles = .{ .{}, .{}, .{ .size = 32, .step = 2, .min = 2 } }, + .served = if (t.heads == asr.heads) + &.{ .{ .rows = 600, .n = asr.heads * asr.d, .k = asr.d }, .{ .rows = 137, .n = asr.heads * asr.d, .k = asr.d } } + else + &.{ .{ .rows = 828, .n = sf.heads * sf.d, .k = sf.d }, .{ .rows = 268, .n = sf.heads * sf.d, .k = sf.d } }, + // the outputs round to bf16 at three points after the sums: the bound is RoPE's, not an exact sum's + .sums = .{ .bf16 = 0, .exact = false, .why = "q and k round to bf16 after their LayerNorm, before RoPE: a bf16 sum's error shows only past that rounding" }, + }; + } + + pub fn deinit(t: *Rope) void { + if (t.arena) |*a| a.deinit(); + t.rig.deinit(); + } + + const vt = g.hooks(Rope); + pub fn kernel(t: *Rope) g.Kernel { + return .{ .spec = t.spec(), .ctx = t, .vt = &vt, .control = t.control }; + } + + /// qkv rows offset by 0, 32 or 256 over noise; stress: pairs v, -v of integers to 256, whose squares overflow fp16. + pub fn prepare(t: *Rope, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.k == 0 or shape.k > na.max_dim or shape.k % 2 != 0) return false; + t.rig.reset(); + if (t.arena) |*a| a.deinit(); + t.arena = std.heap.ArenaAllocator.init(t.rig.gpu.gpa); + const a = t.arena.?.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + t.t = shape.rows; + t.d = shape.k; + t.r = @min(t.rot, t.d) & ~@as(usize, 1); + const d = t.d; + const hd = t.heads * d; + t.x = try a.alloc(u16, t.t * 3 * hd); + for (0..t.t * 3 * t.heads) |v| { + const row = t.x[v * d ..][0..d]; + if (fill == .stress) { + for (0..d / 2) |j| { + const m: f32 = if (j == 0) 256 else @floatFromInt(@as(i32, @intCast(rng.uintLessThan(u32, 513))) - 256); + row[2 * j] = bfr.bits(m); + row[2 * j + 1] = bfr.bits(-m); + } + continue; + } + const offset: f32 = switch (v % 3) { + 0 => 0, + 1 => 32, + else => 256, + }; + for (row, 0..) |*y, i| y.* = bfr.bits(offset + (rng.float(f32) * 2 - 1) * (if (i % 16 == 5) @as(f32, 10) else 1)); + } + if (fill == .nonfinite) for (planted, 0..) |p, i| { + // a q, a k and a v head, each one value + const at = ((i % t.t) * 3 + i) * hd + (i % t.heads) * d + (i * 7) % d; + t.x[at] = bfr.bits(p); + }; + for (0..2) |j| { + t.w[j] = try a.alloc(u16, d); + t.b[j] = try a.alloc(u16, d); + randomBf16(rng, t.w[j], -0.5, 2.5); + randomBf16(rng, t.b[j], -1, 1); + } + t.cs = try a.alloc(u16, t.t * t.r); + t.sn = try a.alloc(u16, t.t * t.r); + randomBf16(rng, t.cs, -1, 1); + randomBf16(rng, t.sn, -1, 1); + t.dev.qkv = try t.rig.output(t.x.len * 2); + if (t.norm) for (0..2) |j| { + t.dev.w[j] = try t.rig.input(std.mem.sliceAsBytes(t.w[j])); + t.dev.b[j] = try t.rig.input(std.mem.sliceAsBytes(t.b[j])); + }; + t.dev.cs = try t.rig.input(std.mem.sliceAsBytes(t.cs)); + t.dev.sn = try t.rig.input(std.mem.sliceAsBytes(t.sn)); + return true; + } + + /// qkv goes in fresh each run (the launch writes it in place). + pub fn run(t: *Rope) !void { + const s = t.rig.s; + const buf: cuda.DeviceBuffer = .{ .d = t.rig.gpu.d, .ptr = t.dev.qkv, .len = t.x.len * 2 }; + try buf.upload(0, std.mem.sliceAsBytes(t.x)); + const norm: ?[4]u64 = if (t.norm) .{ t.dev.w[0], t.dev.b[0], t.dev.w[1], t.dev.b[1] } else null; + if (t.ctl) |f| { + var args: cuda.Args = .{}; + args.add(t.dev.qkv); + for (norm.?) |p| args.add(p); + args.add(t.dev.cs); + args.add(t.dev.sn); + for ([_]usize{ t.heads, t.d, t.r }) |n| args.add(@as(c_int, @intCast(n))); + args.add(eps); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(t.t * t.heads * 2) }, .block = .{ .x = 32 } }, s, &args); + } else try t.k.qkRope(s, t.dev.qkv, norm, t.dev.cs, t.dev.sn, t.t, t.heads, t.d, t.r, eps); + try s.synchronize(); + } + + pub fn count(t: *Rope) usize { + return t.x.len; + } + + /// Every output in float64 with each sum as `sum` gives it; `allow` gets each one's LayerNorm and RoPE bound. + fn outputs(t: *Rope, out: []f64, allow: ?[]f64, sum: anytype) !void { + const a = t.arena.?.allocator(); + const d = t.d; + const y = try a.alloc(f64, d); + const e = try a.alloc(f64, d); + const tmp = try a.alloc(f64, d); + for (0..t.t) |row| for (0..3) |which| for (0..t.heads) |h| { + const at = ((row * 3 + which) * t.heads + h) * d; + const x = t.x[at..][0..d]; + if (which == 2) { // v: as it was + for (x, 0..) |v, i| { + out[at + i] = val(v); + if (allow) |l| l[at + i] = 0; + } + continue; + } + if (t.norm) { + var abs: f64 = 0; + for (tmp, x) |*v, bits| { + v.* = val(bits); + abs += @abs(v.*); + } + const n: f64 = @floatFromInt(d); + const mean = sum.of(tmp) / n; + for (tmp, x) |*v, bits| v.* = (val(bits) - mean) * (val(bits) - mean); + const rstd = 1 / @sqrt(sum.of(tmp) / n + eps); + for (0..d) |i| { + const xv = val(x[i]); + const term = (xv - mean) * rstd * val(t.w[which][i]); + y[i] = term + val(t.b[which][i]); + const mag = (2 * @abs(xv - mean) + abs / n) * rstd * @abs(val(t.w[which][i])); + const fp = gamma(@floatFromInt(depth(.{ .rows = 0, .n = 0, .k = d }))) * mag + 12 * u * (@abs(term) + @abs(val(t.b[which][i]))); + e[i] = fp + half_bf16 * (@abs(y[i]) + fp); // and its bf16 rounding + } + } else for (0..d) |i| { + y[i] = val(x[i]); + e[i] = 0; + } + const hr = t.r / 2; + for (0..d) |i| { + if (i >= t.r) { + out[at + i] = y[i]; + if (allow) |l| l[at + i] = e[i] + floor; + continue; + } + const cv = val(t.cs[row * t.r + i]); + const sv = val(t.sn[row * t.r + i]); + const j = if (i < hr) i + hr else i - hr; + const rv = if (i < hr) -y[j] else y[j]; + const p = y[i] * cv; + const q = rv * sv; + out[at + i] = p + q; + if (allow) |l| { + const work = @abs(cv) * e[i] + @abs(sv) * e[j]; + const rounded = work + half_bf16 * (@abs(p) + @abs(q) + work) * (1 + u); + l[at + i] = rounded + half_bf16 * (@abs(p + q) + rounded) + floor; + } + } + }; + } + + const Exact = struct { + fn of(_: Exact, v: []const f64) f64 { + var s: f64 = 0; + for (v) |x| s += x; + return s; + } + }; + const Sim = struct { + acc: c.Accum, + order: g.Order, + fn of(s: Sim, v: []const f64) f64 { + return rows_.sumIn(s.acc, s.order, v); + } + }; + + pub fn reference(t: *Rope, want: []f64, mag: []f64) !void { + t.allow = try t.arena.?.allocator().alloc(f64, want.len); + try t.outputs(want, t.allow, Exact{}); + for (mag, want) |*m, w| m.* = @abs(w); + } + + pub fn bounds(t: *Rope, _: []const f64, _: []const f64, allow: []f64) !void { + @memcpy(allow, t.allow); + } + + pub fn simulate(t: *Rope, acc: c.Accum, order: g.Order, out: []f64) !void { + try t.outputs(out, null, Sim{ .acc = acc, .order = order }); + } +}; + +pub fn noSums(_: c.Shape) usize { + return 0; +} + +// --- tf_na_attention: softmax(q k^T scale) over keys [0, len) in fp32, P as bf16, P V summed in fp32 --------------- + +pub const Attention = struct { + rig: rig_.Rig, + k: *const na.Kernels, + name: []const u8, + heads: usize, + /// The control: one key short of the window. + short: bool = false, + control: ?g.Gate = null, + arena: ?std.heap.ArenaAllocator = null, + t: usize = 0, + len: usize = 0, + d: usize = 0, + qkv: []u16 = &.{}, + half: []f64 = &.{}, + in: u64 = 0, + out: u64 = 0, + + /// The P V chain over the keys, or a score's over d. + fn depth(s: c.Shape) usize { + return @max(s.n + 1, s.k + 1); + } + + pub fn spec(t: *const Attention) c.Spec { + const why = "P rounds to bf16 before P V: its 2^-8 of each term outweighs what a bf16 sum of the values loses"; + return .{ + .name = t.name, + .accum = .f32, + .depth = depth, + .ops = 4, + .mids = 1, + .out = .bf16, + .tiles = .{ .{}, .{ .size = 128 }, .{} }, + .served = if (t.heads == asr.heads) + &.{ .{ .rows = 600, .n = 600, .k = asr.d }, .{ .rows = 600, .n = 137, .k = asr.d } } + else + &.{ .{ .rows = 828, .n = 828, .k = sf.d }, .{ .rows = 828, .n = 791, .k = sf.d }, .{ .rows = 268, .n = 240, .k = sf.d } }, + .sums = .{ .bf16 = 0, .exact = false, .why = why }, + }; + } + + pub fn deinit(t: *Attention) void { + if (t.arena) |*a| a.deinit(); + t.rig.deinit(); + } + + const vt = g.hooks(Attention); + pub fn kernel(t: *Attention) g.Kernel { + return .{ .spec = t.spec(), .ctx = t, .vt = &vt, .control = t.control }; + } + + /// qkv [t, 3, heads, d] in [-1, 1); stress: q zero (every weight equal), each head's values mirrored over the keys. + pub fn prepare(t: *Attention, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + // rows: the window's (a gate's few queries still read every key), n: its valid keys + const rows = @max(shape.rows, shape.n); + if (shape.rows == 0 or shape.n == 0 or rows > na.max_keys or shape.k > na.max_dim) return false; + t.rig.reset(); + if (t.arena) |*a| a.deinit(); + t.arena = std.heap.ArenaAllocator.init(t.rig.gpu.gpa); + const a = t.arena.?.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + t.t = rows; + t.len = shape.n; + t.d = shape.k; + const hd = t.heads * t.d; + t.qkv = try a.alloc(u16, t.t * 3 * hd); + randomBf16(rng, t.qkv, -1, 1); + if (fill == .stress) { + for (0..t.t) |r| @memset(t.qkv[r * 3 * hd ..][0..hd], 0); + const col = try a.alloc(f32, t.len); + for (0..hd) |e| { + mirrored(rng, col); + for (col, 0..) |v, j| t.qkv[(j * 3 + 2) * hd + e] = bfr.bits(v); + } + } + if (fill == .nonfinite) { + // a key of head 0 (every query of that head), a value of head 1, a query of the last head + t.qkv[(0 * 3 + 1) * hd + 3] = bfr.bits(planted[0]); + if (t.heads > 1) t.qkv[(@min(t.len, t.t) / 2 * 3 + 2) * hd + t.d + 5] = bfr.bits(planted[1]); + t.qkv[((t.t - 1) * 3 + 0) * hd + (t.heads - 1) * t.d] = bfr.bits(planted[2]); + } + t.in = try t.rig.input(std.mem.sliceAsBytes(t.qkv)); + t.out = try t.rig.output(t.t * hd * 2); + return true; + } + + fn scale(t: *const Attention) f32 { + return @floatCast(1.0 / @sqrt(@as(f64, @floatFromInt(t.d)))); + } + + pub fn run(t: *Attention) !void { + try t.k.attend(t.rig.s, t.in, t.out, t.t, t.heads, t.d, if (t.short) t.len - 1 else t.len, t.scale()); + try t.rig.s.synchronize(); + } + + pub fn count(t: *Attention) usize { + return t.t * t.heads * t.d; + } + + /// The kernel's arithmetic in attention_range's form: fma scores, expf, P = bf16(e / l), one fma chain for P V. + fn form(t: *const Attention) ar.Form { + const n: f64 = @floatFromInt(t.len); + return .{ .score = @floatFromInt(t.d + 1), .fast_exp = false, .chunk_exp = false, .p_rel = half_bf16 + u, .num = n + 1, .den = @floor(n / 128) + 11 }; + } + + pub fn reference(t: *Attention, want: []f64, mag: []f64) !void { + const a = t.arena.?.allocator(); + const d = t.d; + const hd = t.heads * d; + const n = t.len; + t.half = try a.alloc(f64, want.len); + const q = try a.alloc(f64, d); + const kk = try a.alloc(f64, n * d); + const vv = try a.alloc(f64, n * d); + const w = try a.alloc(f64, n); + const tau = try a.alloc(f64, n); + for (0..t.heads) |h| { + for (0..n) |j| for (0..d) |e| { + kk[j * d + e] = val(t.qkv[(j * 3 + 1) * hd + h * d + e]); + vv[j * d + e] = val(t.qkv[(j * 3 + 2) * hd + h * d + e]); + }; + for (0..t.t) |r| { + for (q, 0..) |*x, e| x.* = val(t.qkv[r * 3 * hd + h * d + e]); + const at = (r * t.heads + h) * d; + ar.head(q, kk, vv, n, d, t.scale(), t.form(), w, tau, want[at..][0..d], t.half[at..][0..d]); + } + } + for (mag, want) |*m, x| m.* = @abs(x); + } + + /// The range's half width, then the output's own rounding at its far end. + pub fn bounds(t: *Attention, want: []const f64, _: []const f64, allow: []f64) !void { + for (allow, want, t.half) |*l, w, h| l.* = h + half_bf16 * (@abs(w) + h) + floor; + } + + /// Gate 4: every weight is 1, so an output is its value column's sum over the keys, over their count. + pub fn simulate(t: *Attention, acc: c.Accum, order: g.Order, out: []f64) !void { + const a = t.arena.?.allocator(); + const hd = t.heads * t.d; + const terms = try a.alloc(f64, t.len); + for (0..hd) |e| { + for (terms, 0..) |*x, j| x.* = val(t.qkv[(j * 3 + 2) * hd + e]); + const o = rows_.sumIn(acc, order, terms) / @as(f64, @floatFromInt(t.len)); + for (0..t.t) |r| out[r * hd + e] = o; + } + } +}; + +/// Gate 4's values: odd integers, every other one times 64 (bf16-exact), the second half cancelling the first. +fn mirrored(rng: std.Random, out: []f32) void { + const h = out.len / 2; + for (0..h) |i| { + const v: f32 = @floatFromInt((1 + 2 * rng.uintLessThan(u32, 128)) * @as(u32, if (i % 2 == 1) 64 else 1)); + out[i] = v; + out[h + i] = (if (i < 8) @as(f32, 1) else 0) - v; + } + if (out.len % 2 == 1) out[out.len - 1] = 0; +} + +// --- the Sortformer head's glue: im2col3, relu and zero_rows, each exact ------------------------------------------- + +pub const Glue = @import("contract_nemo_audio_glue.zig").Glue; + +/// Every registration, on the heap: the gates hold pointers to them. +pub const Set = struct { + gpa: std.mem.Allocator, + k: na.Kernels, + ctl: cuda.Module, + ropes: [3]Rope = undefined, + attns: [3]Attention = undefined, + glue: [6]Glue = undefined, + + /// The set, its kernels appended to `list` (each control after the kernel it breaks). + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + var k = try na.Kernels.load(gpu.d); + errdefer k.unload(); + var ctl = try cuda.Module.load(gpu.d, &controls.bytes); + errdefer ctl.unload(); + set.* = .{ .gpa = gpu.gpa, .k = k, .ctl = ctl }; + const rig: rig_.Rig = .{ .gpu = gpu, .s = s }; + const kp = &set.k; + set.ropes = .{ + .{ .rig = rig, .k = kp, .name = "cuda nemo_audio tf_na_qk_rope, ASR encoder (q/k LayerNorm, RoPE on 64 of 128)", .heads = asr.heads, .rot = asr.rot, .norm = true }, + .{ .rig = rig, .k = kp, .name = "control: nemo_audio q/k LayerNorm with the variance in one pass (E[x^2] - mean^2)", .heads = asr.heads, .rot = asr.rot, .norm = true, .ctl = try ctl.function("tf_ctl_na_qk_onepass"), .control = .shapes }, + .{ .rig = rig, .k = kp, .name = "cuda nemo_audio tf_na_qk_rope, Sortformer encoder (RoPE on 64 of 64)", .heads = sf.heads, .rot = sf.rot, .norm = false }, + }; + set.attns = .{ + .{ .rig = rig, .k = kp, .name = "cuda nemo_audio tf_na_attention, ASR encoder (10 heads of 128)", .heads = asr.heads }, + .{ .rig = rig, .k = kp, .name = "control: nemo_audio tf_na_attention one key short of the window", .heads = asr.heads, .short = true, .control = .shapes }, + .{ .rig = rig, .k = kp, .name = "cuda nemo_audio tf_na_attention, Sortformer encoder (8 heads of 64)", .heads = sf.heads }, + }; + set.glue = .{ + .{ .rig = rig, .k = kp, .name = "cuda nemo_audio tf_na_im2col3 (the subpixel conv's taps)", .op = .im2col3 }, + .{ .rig = rig, .k = kp, .name = "control: nemo_audio tf_na_im2col3 reading a padding row", .op = .im2col3, .late = true, .control = .shapes }, + .{ .rig = rig, .k = kp, .name = "cuda nemo_audio tf_na_zero_rows (padding rows)", .op = .zero_rows }, + .{ .rig = rig, .k = kp, .name = "control: nemo_audio tf_na_zero_rows keeping a padding row", .op = .zero_rows, .late = true, .control = .shapes }, + .{ .rig = rig, .k = kp, .name = "cuda nemo_audio tf_na_relu", .op = .relu }, + .{ .rig = rig, .k = kp, .name = "control: nemo_audio relu by fmaxf, which drops a NaN", .op = .relu, .ctl = try ctl.function("tf_ctl_na_relu_fmax"), .control = .nonfinite }, + }; + for (&set.ropes) |*x| try list.append(a, x.kernel()); + for (&set.attns) |*x| try list.append(a, x.kernel()); + for (&set.glue) |*x| try list.append(a, x.kernel()); + return set; + } + + pub fn deinit(set: *Set) void { + for (&set.ropes) |*x| x.deinit(); + for (&set.attns) |*x| x.deinit(); + for (&set.glue) |*x| x.deinit(); + set.ctl.unload(); + set.k.unload(); + set.gpa.destroy(set); + } +}; + +/// nemo_audio.cu's gates 1-4 and their controls alone. +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(a); + const set = try Set.init(gpu, s, a, &list); + defer set.deinit(); + var t: g.Tally = .{}; + try g.runAll(a, list.items, 0x6b63, &t); // kernel-contract's seed: the same inputs as the full run + std.debug.print("SUMMARY nemo_audio: {d} kernels, {d} pass gates 1-4, {d} fail; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "nemo-audio-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} diff --git a/zig/tests/cuda/contract_nemo_audio_glue.zig b/zig/tests/cuda/contract_nemo_audio_glue.zig new file mode 100644 index 0000000000..3d3c5a8bb7 --- /dev/null +++ b/zig/tests/cuda/contract_nemo_audio_glue.zig @@ -0,0 +1,129 @@ +//! nemo_audio.cu's Sortformer head glue under the contract: im2col, padding and relu rows, exact. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const na = @import("nemotron").nemo_audio.kernels; +const rig_ = @import("contract_rig.zig"); +const base = @import("contract_nemo_audio.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const bfr = core.bf16_round; +const val = base.val; +const randomBf16 = base.randomBf16; +const planted = base.planted; +const controls = base.controls; +const noSums = base.noSums; +/// The Sortformer head's width and its upsampled rows (8 frames a row). +const tf_d = 192; + +pub const Glue = struct { + rig: rig_.Rig, + k: *const na.Kernels, + ctl: ?cuda.Function = null, // relu's fmaxf control + name: []const u8, + op: Op, + /// The im2col and zero_rows controls: one padding row read as valid. + late: bool = false, + control: ?g.Gate = null, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = undefined, + len: usize = 0, + x: []u16 = &.{}, + in: u64 = 0, + out: u64 = 0, + + pub const Op = enum { im2col3, relu, zero_rows }; + + pub fn spec(t: *const Glue) c.Spec { + return .{ + .name = t.name, + .accum = .none, + .depth = noSums, + .out = .bf16, + .in_place = t.op != .im2col3, + .tiles = .{ .{ .size = 256 }, .{}, .{} }, + // rows: the head's rows (cache, FIFO and window; relu: 8 frames each), n: its width + .served = switch (t.op) { + .relu => &.{ .{ .rows = 828 * 8, .n = tf_d, .k = 0 }, .{ .rows = 268 * 8, .n = tf_d, .k = 0 } }, + else => &.{ .{ .rows = 828, .n = tf_d, .k = 0 }, .{ .rows = 268, .n = tf_d, .k = 0 } }, + }, + }; + } + + pub fn deinit(t: *Glue) void { + if (t.arena) |*a| a.deinit(); + t.rig.deinit(); + } + + const vt = g.hooks(Glue); + pub fn kernel(t: *Glue) g.Kernel { + return .{ .spec = t.spec(), .ctx = t, .vt = &vt, .control = t.control }; + } + + /// x [rows, n] in [-4, 4), the valid rows three quarters of them (at least one). + pub fn prepare(t: *Glue, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0) return false; + t.rig.reset(); + if (t.arena) |*a| a.deinit(); + t.arena = std.heap.ArenaAllocator.init(t.rig.gpu.gpa); + const a = t.arena.?.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + t.shape = shape; + t.len = @max(1, shape.rows * 3 / 4); + t.x = try a.alloc(u16, shape.rows * shape.n); + randomBf16(prng.random(), t.x, -4, 4); + if (fill == .nonfinite) for (planted, 0..) |p, i| { + t.x[(i * 131 + 7) % t.x.len] = bfr.bits(p); + }; + if (t.op == .im2col3) { + t.in = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.out = try t.rig.output(shape.rows * 3 * shape.n * 2); + } else t.out = try t.rig.output(t.x.len * 2); + return true; + } + + pub fn run(t: *Glue) !void { + const s = t.rig.s; + const rows = t.shape.rows; + const n = t.shape.n; + const len = if (t.late) t.len + 1 else t.len; + if (t.op != .im2col3) { + const buf: cuda.DeviceBuffer = .{ .d = t.rig.gpu.d, .ptr = t.out, .len = t.x.len * 2 }; + try buf.upload(0, std.mem.sliceAsBytes(t.x)); + } + switch (t.op) { + .im2col3 => try t.k.im2col3(s, t.in, t.out, rows, n, len), + .zero_rows => try t.k.zeroRows(s, t.out, rows, n, len), + .relu => if (t.ctl) |f| { + var args: cuda.Args = .{}; + args.add(t.out); + args.add(@as(i64, @intCast(t.x.len))); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast((t.x.len + 255) / 256) }, .block = .{ .x = 256 } }, s, &args); + } else try t.k.reluInPlace(s, t.out, t.x.len), + } + try s.synchronize(); + } + + pub fn count(t: *Glue) usize { + return if (t.op == .im2col3) 3 * t.x.len else t.x.len; + } + + /// Copies, zeros and relu(x) = x < 0 ? 0 : x (NaN stays): no sums, no magnitude. + pub fn reference(t: *Glue, want: []f64, mag: []f64) !void { + const n = t.shape.n; + switch (t.op) { + .im2col3 => for (0..t.shape.rows) |r| for (0..3) |j| for (0..n) |ch| { + const src = @as(isize, @intCast(r + j)) - 1; + want[(r * 3 + j) * n + ch] = if (src >= 0 and src < t.len) val(t.x[@as(usize, @intCast(src)) * n + ch]) else 0; + }, + .zero_rows => for (want, t.x, 0..) |*w, x, i| { + w.* = if (i / n < t.len) val(x) else 0; + }, + .relu => for (want, t.x) |*w, x| { + w.* = if (val(x) < 0) 0 else val(x); + }, + } + @memset(mag, 0); + } +}; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index a3ad2d2dad..2c0fe71755 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -26,6 +26,7 @@ const mxfp8_tests = @import("mxfp8_tests.zig"); const projection_f32_tests = @import("projection_f32_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); +const bf16_experts_tests = @import("bf16_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); const glm_ref_tests = @import("glm_ref_tests.zig"); const affine8_tests = @import("affine8_tests.zig"); @@ -198,6 +199,7 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes + \\ bf16-experts grouped bf16 experts against float64: repack, up, downs, rows and forms byte for byte \\ projection-f32 FP32 projection outputs against FP64, rounding and row/reduction invariance \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) @@ -233,9 +235,11 @@ const usage = \\ qmm-prefill-tile22-contract the GB10 large-prompt 2x2 tile against the default tile \\ qmm-prefill-many-bench same-process A-B-A screening for the GB10 launch fusion \\ mtp-rows tf_q36m_attn_rows: gates 1-4 (byte-equal to tf_q36m_attn a segment) and float64 alone + \\ experts-bf16-contract the bf16 experts' up and downs through gates 1-4, their controls red (needs -Dnvcc) \\ qwen-pages qwen35.cu's attention cache through a page table: paged == contiguous, float64, red controls \\ rope-yarn YaRN rope in tf_q35_rope, tf_q35_rope8 and tf_q36m_norm_rope against float64, to 1M \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR + \\ nemo-audio-contract Nemotron 3.5 Transcribe's audio encoder kernels through gates 1-4, controls red \\ pocket-contract Nemotron's pocket kernels (learned neurons) through gates 1-4, controls red \\ minimax-contract MiniMax-M3.1-Flash's indexer and block-sparse attention through gates 1-4, controls red \\ moss-contract MOSS-Image's own kernels and its full attention through gates 1-4, controls red @@ -372,6 +376,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "bf16-experts")) return bf16_experts_tests.experts(gpu); if (std.mem.eql(u8, cmd, "projection-f32")) return projection_f32_tests.run(gpu); if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); @@ -459,10 +464,12 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "qmm-prefill-tile22-contract")) return qmm_prefill_tile22_contract.run(gpu); if (std.mem.eql(u8, cmd, "qmm-prefill-many-bench")) return qmm_prefill_many_contract.bench(gpu); if (std.mem.eql(u8, cmd, "mtp-rows")) return @import("contract_qwen_mtp_rows.zig").alone(gpu); + if (std.mem.eql(u8, cmd, "experts-bf16-contract")) return @import("contract_experts_bf16.zig").run(gpu); if (std.mem.eql(u8, cmd, "qwen-pages")) return @import("contract_qwen_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "minimax-contract")) return @import("contract_minimax.zig").run(gpu); if (std.mem.eql(u8, cmd, "rope-yarn")) return @import("contract_qwen_rope_yarn.zig").run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "nemo-audio-contract")) return @import("contract_nemo_audio.zig").run(gpu); if (std.mem.eql(u8, cmd, "pocket-contract")) return @import("contract_pocket.zig").run(gpu); if (std.mem.eql(u8, cmd, "moss-contract")) return @import("contract_moss_image.zig").run(gpu); if (std.mem.eql(u8, cmd, "tree")) return tree_tests.run(gpu);