Repository navigation
drama: load the checkpoint with direct I/O in file order and build the codec on its own thread - #925
Closed
plotarmordev wants to merge 48 commits into
Closed
plotarmordev wants to merge 48 commits into
plotarmordev wants to merge 48 commits into
Conversation
setAllocationLimit bounds the bytes of device buffers plus pinned host buffers. An allocation reserves its bytes before the driver call and a failed call refunds them, so concurrent allocations cannot overshoot the cap and a refused one leaves usage and the measured peak unchanged. Host counts now take the same lock as device counts. An engine that admits a model against a memory budget can hold its own allocations to it.
The tower turns a backbone row into a frame's ten codebook values: four shared and nine residual layers whose cache keeps every preceding text and audio row, a semantic head and the reference sampler (temperature, top-p over the whole vocabulary, then top-k). Its projections run on cuBLASLt; drama_tower.cu holds the mixer, norms, RoPE, attention and sampling kernels.
The codec decodes ten codebooks a frame to mono 44.1 kHz PCM, 2048 samples a frame, as the reference's FP16 vocoder does: RVQ, the windowed post-module transformer, two upsample stages and the DAC decoder. codec_pth.zig reads torch.save's zip archive and runs its state-dict pickle on a VM that knows only the opcodes and globals a tensor state dict uses, so nothing in the file runs. Weight norm is folded in FP32 and tensors are cast to FP16 with round-to-nearest-even, as the reference does. codec.zig holds the decode; codec_load.zig holds the codec's geometry, the weight-norm folding, the FP16 casts and the one device weight image.
An engine may now synthesize speech: SpeechRequest carries a prompt, sampling settings and caller-owned PCM, Info.speech describes the audio it produces, and Info.text_generation lets a speech-only engine refuse the text routes. Text submit gains error.Unsupported, which the chat route reports as a request error.
The native CUDA backend opens fish_omni checkpoints with the Drama 3 host: one synthesis at a time on a worker thread that owns the context, stream, backbone, tower and codec, admitted against the same memory pool the text families use. The pool reading moves to cuda_memory.readPool so both paths share it, and the host caps its own allocations at what it was admitted for. Options the host cannot honor are refused at open.
OpenAI's speech request, built into Drama 3's default-voice prompt as the reference's prompt builder builds it and synthesized whole, then sent as WAV or raw PCM. Fields the route would silently drop are refused, and a speech-only engine starts the server without the text-only services.
Drivers that check the native backbone and tower against captured reference tensors, replay single operators against them, and trace each stage for attribution, plus tf-cuda-test drama-conv and drama-qk on synthetic rows. None of them runs as part of the host tests.
The codec can decode frames from a given frame onward at their absolute rows, reading the rows before them that its causal convolutions and attention windows reach, and writing nothing before them. Every chunk boundary gives the same samples as one decode of all frames, so audio can be decoded while it is generated.
Drivers that decode captured codes with the native codec for comparison with the reference decoder, replay its RVQ and weight preparation in isolation, and trace each decoder stage. None of them runs as part of the host tests.
… generates A streamed speech request now receives audio events: each one hands the caller a final, rising prefix of the PCM before the finished event. The Drama host decodes the first frame alone, for the earliest audio, then frames in groups, and publishes each decoded prefix.
…ve host and their tool steps Embeds the backbone and tower fatbins, adds the drama module to the native server and its host tests, and a build step per backbone and tower oracle, replay and trace driver. The codec's kernels and tools follow on their own.
… are decoded With stream: true and response_format pcm, the route sends HTTP/1.1 chunks as the engine publishes audio. A stream that fails after its first chunk ends without a terminator, so a client cannot mistake it for complete audio.
…memory A single-row projection copies each warp's weight chunks into a shared ring with cp.async ahead of use and reads them back with ldmatrix, issuing the same mma sequence on the same fragments in the same K order as the row tiling it replaces, so the results are the same bits.
Embeds the codec fatbin and adds a build step per codec oracle, replay and trace driver. With this the series builds.
Decodes the same frame codes whole and in the host's chunk policy, one frame at a time and in irregular chunks, each over NaN-poisoned scratch, and requires every chunked decode to equal the whole decode.
Some file systems accept an O_DIRECT open and then fail the aligned read with EINVAL. The reader now clears O_DIRECT on that descriptor and reads again through the page cache instead of failing the load; a test injects the refusal.
A marker test: an image can carry a device global that promises something about its kernels, and the host can check for it before relying on the promise. Any driver failure, a missing name included, answers false rather than an error. A test drives it through a fake driver table.
Drama 3's backbone is a Qwen3.5-35B-A3B-style hybrid: 30 Gated DeltaNet and 10 gated full-attention layers, each with a 256-expert top-8 block-FP8 MoE and a gated shared expert. Model.run appends rows and returns every row's final-norm hidden state and the last row's semantic logits. The kernels follow the arithmetic of the sglang-omni reference: the canonical 64-token chunked GDN prefill (drama_gdn_chunk.cuh, after flash-linear-attention's kernels as SGLang carries them), the packed recurrent decode, gated attention with Q/K norm and partial NeoX RoPE, dynamic 1 x 128 FP8 activation quantization, and the BF16 tensor-core GEMMs in drama_backbone_gemm.cu. Weights load as stored; nothing is requantized. THIRD_PARTY_NOTICES.md credits flash-linear-attention (MIT), whose chunked kernels the GDN prefill is adapted from and whose packed decode the recurrence follows.
…ync ring The one-row path keeps the same mma sequence and fragments as the tilings it replaces, now for the semantic head's grouping too, while its weights stream through an unpadded shared ring whose 16-byte pieces are swizzled so each read still hits distinct banks.
…ation The residual add and RMS norm before the MoE also write the normed rows' FP8 groups, in the same order and rounding as the separate quantization launch they replace.
The routed rows are stored as their FP8 quantization and the shared expert's rows as BF16, in one launch instead of two, with the same arithmetic per row.
Each thread of drama_moe_combine now writes 8 contiguous outputs from 16-byte loads; every sum keeps its order.
A warp loads and scores its keys in batches with the next batch's loads issued early; the running softmax update stays in key order, so the result is unchanged.
A cluster of CTAs shares a head: each scores its share of the key tiles, all gather the scores over distributed shared memory and replay the tile updates in order, then each sums its own slice of the head. Every value is computed as the one-CTA kernel computes it; prompt rows keep that kernel.
… order stager.zig collects each tensor's copy as the loaders walk their names, sorts them by file offset and streams each file through two pinned slots, reading one with O_DIRECT while copying from the other, in pieces that never span two device allocations.
…udio A synthesis that reaches max_frames without an end of speech keeps the audio of those frames, as the reference returns its capped replies, instead of failing without audio. A whole reply carries X-Finish-Reason stop or length; a stream declares the trailer and ends with X-Finish-Reason: length only when it was clipped.
…ame launch For a single row, the attention launch also does drama_attn_prep's norm, RoPE and cache append, taking the new key and value from registers, with the same roundings as the separate launch.
The sampler's ranking comes from a stable LSD radix sort under the same strict total order the bitonic network used, so the permutation and every sum over it are unchanged.
…one loads The codec's host-side weight image is CPU work that the backbone and tower reads leave room for; it now runs on its own thread with the context current there, and the host joins it before it finishes opening.
Closed
4 of 8 tasks
Owner
|
Thank you, @plotarmordev, we will carry the direct/parallel loading work with your credit into the combined Drama owner and loader; the landing order and remaining checks are in our reply on #927. Posted by Ash's AI Assistant Synthena |
Conflicts, each kept from both sides:
- zig/build/cuda.zig: family() returns radio and drama; engines() and
nativeServer() take core, heat and drama; main's tf-cuda-test runner,
kernel contract and GLM/LFM2/RADIO helpers stay ahead of the drama tool steps.
- zig/src/core/engine_api.zig: the speech types beside main's vision types;
SubmitError holds Unsupported and Stale; the VTable keeps submit_speech
beside recall, imprint and drain.
- zig/src/native/cuda.zig: families lists lfm2_vl then fish_omni; pool() and
weightBytes() leave for cuda_memory's readPool() and weightBytes(); openDrama
refuses main's new --vision, --precision, --record and --learn-conversations.
- zig/src/server/{chat,root,routes}.zig, zig/tests/cuda/main.zig: both sides'
error arms, imports and test files.
- zig/src/server/speech.zig: imports the named json module (ashhart#860).
…treaming CHANGELOG.md: the streaming line stays under ## Unreleased, above main's ## 1.0.12.
…ackbone-one-row
…ts image holds no cluster kernel drama_attention_decode's __cluster_dims__ and cluster group exist from sm_90 on, so nvcc refused drama_tower.cu for sm_80, sm_86 and sm_89, and the backbone admits sm_89. The kernel now compiles for sm_90+ images only, and the tower resolves and launches it only when the device has clusters; other devices run the one-CTA drama_attention the multi-row calls use. The sm_120 and sm_121 images are byte-identical.
…n types in their own files, Drama's refusals in cuda_refusal Upstream main's engine_api.zig is at 600 lines and native/cuda.zig at 583, so the merge took them to 632 and 602. The speech request, result, offer and submit error move to core/speech_api.zig, and the image and video input types to core/vision_api.zig, re-exported under their names as rows_ask.zig's are. Drama 3's flag refusals join LFM2's in cuda_refusal.zig, with a test. No declaration, behavior or test is removed.
…ion types in their own files zig/src/core/engine_api.zig: the speech types leave as on 920; the streamed request's comment and stream field move with them into core/speech_api.zig.
…d vision types in their own files
…and vision types in their own files
…vision types in their own files
…ypes in their own files
ashhart#921 adds Event.audio; main's newer lane host tests switch over every event, so zig build test stopped compiling them.
…s ignore the speech audio event
…sts ignore the speech audio event
…ignore the speech audio event
…he speech audio event
Owner
|
Thank you, @plotarmordev, for the speech protocol and infrastructure, and @drowzeys for the family, backends and codec. The selected work from #925 is included through #927 ( Posted by Ash's AI Assistant Synthena |
4 of 8 tasks
Owner
|
The consolidated Drama contribution shipped in 1.1.0 as Posted by Ash's AI Assistant Synthena |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this changes
Checkpoint weights stream onto the GPU with direct I/O in file order through two pinned slots (part of #914;
stacked on #924, its own commits are the last two). The codec builds its weight image on its own thread while
the backbone loads.
This stack implements #914 and depends on #915 (allocation cap), #916 (direct-read fallback) and #917
(Module.holds), inherited through #920.
This inherits the stack's overlap with #844, including the recorded CUDA integration and
CHANGELOG.mdconflicts;land #844 first, then rebase the stack.
Receipt
68aeb3e51dd2, on mainc7137bb3. No MLX or PyTorch.0.17.0, nvcc 13.0.88 and
CUDA_VISIBLE_DEVICESempty.fishaudio/drama-3ate0a9218a; reference sglang-omniab095ef9.1e76a70e. Here they are rebased ontoc7137bb3with comments made single-line and files split. Every Drama fatbin of this head is byte-identical tothe measured series' at the same commit. The host Zig is token-identical apart from the file split, compared
with
std.zig.Ast.tools/bench_concurrent.py --alone --serialortoken_shaagainst"draft": false: n/a fortokens. The same bytes land in the same device buffers. The host tests compare the staged image with a
loader-order copy, and the 46 normal replies stayed byte-identical through this stage (speed-stages.json).
Fidelity against the reference is not qualified (qualification.json); preserving native output does
not establish reference equivalence.
tools/bench_openai.py: n/a. Only loading changes, and decode is untouched.ttft_sfromtools/bench_cache.py --require-coldon a long prompt, or "n/a" with the reason when the change cannot reach prompt processing: Prompt processingwas not measured; these commits change loading rather than the prompt kernels. The following numbers measure
loading/startup, not cold-request TTFT (speed-stages.json):
bef432bbb25e(one GB10, under the recipe's guard): all 46 normal-requestaudio bodies byte-identical to the promoted pre-upstream build, the long-05 clipped body identical in both
passes,
bench/stream_check.py9/9, warm RTF 0.654 against the pre-speed 1.087, and all six tower one-rowprojections on
drama_tower_gemv.CUDA_VISIBLE_DEVICES= zig build test-cuda-host --summary allfrom an empty local cache: 29/29 steps, 496 passed, 0 failed, 0 skipped. This includes the stager's two new tests (read planning and landed bytes for random copies).zig build native -Dnvcc=/usr/local/cuda/bin/nvcc --summary all: 46/46 steps.zig build fatbins -Dnvcc=/usr/local/cuda/bin/nvcc: every Drama fatbin is byte-identical to the measuredseries' at this commit.
python3 tools/zig/lean_check.py: 527 problems at every commit, main's 527 (lean_new.py: 0 new).zig fmt --check: clean at every commit.bef432bbb25e; this head ishost-tested only.
tools/bench_openai.py,bench_concurrent.pyandbench_cache.py: they drive text tokens, so speech used therecipe's harness.
Checklist
generates audio frames in its own host loop, as in drama: Drama 3 (fishaudio/drama-3) text-to-speech on CUDA #920. Text families are unchanged.)
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Nodrafts or resume. The measured speed stages preserved the 46 normal corpus audio bodies; concurrent-equals-solo
is not measured.)