Repository navigation
drama: the tower's one-row projections on drama_tower_gemv, in cuBLASLt's summation order - #926
Open
plotarmordev wants to merge 30 commits into
Open
plotarmordev wants to merge 30 commits into
plotarmordev wants to merge 30 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.
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.
…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.
Embeds the codec fatbin and adds a build step per codec oracle, replay and trace driver. With this the series builds.
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.
… 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.
… 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.
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.
…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.
…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.
…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.
…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.
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.
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.
… 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.
…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.
…Lt's summation order drama_tower_gemv sums in the order of the gemv kernel cuBLASLt picks for the tower's one-row plans. At init each projection is checked against its cuBLASLt plan bit for bit and uses the kernel only if they agree. The qualification and its probe weights live in tower_gemv.zig.
This was referenced Oct 16, 2026
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
1 participant
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
The tower's one-row projections run on
drama_tower_gemv, which reproduces the summation order of cuBLASLt's gemvkernel (part of #914; stacked on #925, its own commit is the last one). It is enabled per projection
only after that projection's plan matched it bit for bit at init.
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
2757724d5dc6, 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. Each speed stage was gated on all 46 normal replies of a 48-request corpus capture: byte-identical to the
build before any speed change. The two capped long-05 replies were identical to the
reviewed clipped capture, and
bench/stream_check.pypassed 9/9 (speed-stages.json). Fidelity against the reference is not qualified(qualification.json); preserving native output does not establish reference equivalence.
bit for bit; this finite probe is not an exhaustive proof over all inputs.
tools/bench_openai.py: n/a, because it measures tokens. Warm real-time factor0.741 → 0.678 for stage 04, this commit (speed-stages.json).
ttft_sfromtools/bench_cache.py --require-coldon a long prompt, or "n/a" with the reason when the change cannot reach prompt processing: Not measured. Multi-row prompt calls retain cuBLASLt plans; the optimizedprojection dispatch is based on row count, so single-row prompt calls can use the qualified GEMV path.
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, 498 passed, 0 failed, 0 skipped. This includes the two new tests: the probe separates the plans' order from five neighbouring orders, and the qualification commits only when every projection was probed.zig build native -Dnvcc=/usr/local/cuda/bin/nvcc --summary all: 47/47 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.cuBLASLt elsewhere.
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.)