Repository navigation
drama: backbone attention runs the reference's 64-key softmax tiles - #985
Closed
plotarmordev wants to merge 62 commits into
Closed
plotarmordev wants to merge 62 commits into
plotarmordev wants to merge 62 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.
…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.
…m_90 cluster guard zig/src/families/drama/tower.zig: the guard beside 926's pub projections and loadModule.
…he one-CTA kernel below it
…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
…sion types in their own files
…pes 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
…nore the speech audio event
…e speech audio event
Native kept a running max per key and per warp and merged the warps at the end. The reference extend kernel keeps one max per 64-key tile, with Triton's ex2.approx exp, div.full and fused multiply-adds. Same-input replays put full attention 10 to 3000 times further from the reference than MoE or GDN, and router near-ties amplified it down the stack. With the reference's tiles and arithmetic, the attention operator matches on 98.1% and 96.2% of its outputs at layers 3 and 39 (was 70.4% and 57.5%). The 182-token text fixture goes from relative L2 0.124 to 0.053, KL 0.024 to 0.0062, worst row 0.80 to 0.92. Tiles are fixed by key position, so chunk 1/7/128 and reset stay bit-exact.
…ore the speech audio event
This was referenced Oct 17, 2026
Open
4 of 8 tasks
Owner
|
Thank you, @plotarmordev, for the layer-level receipts and the improved layer-3 bit agreement. This increment remains out of 1.1.0 and rides 1.1.1, the next patch release. The all-row cosine, last-row cosine and relative-L2 results still miss their unchanged fidelity bars. Please supply the frozen layer/backbone inputs and reference decode tensors so the three failures can be repaired without relaxing precision or tolerances. The final build also needs its launch controls and matched 40-text voice baseline, with cache-resume and concurrent audio comparisons. Posted by Ash's AI Assistant Synthena |
Owner
|
Thank you, @plotarmordev, for the layer-level receipts and the improved layer-3 bit agreement. This increment remains out of 1.1.0 and rides 1.1.1, the next patch release. The all-row cosine, last-row cosine and relative-L2 results still miss their unchanged fidelity bars. Please supply the frozen layer/backbone inputs and reference decode tensors so the three failures can be repaired without relaxing precision or tolerances. The final build also needs its launch controls and matched 40-text voice baseline, with cache-resume and concurrent audio comparisons. No new Music 3 device pass is recorded, and the three fidelity failures, frozen-reference comparisons and matched audio/concurrency checks remain unresolved. Posted by Ash's AI Assistant Synthena |
Contributor
Author
|
Closing this PR because it edits On CUDA, the landed implementation uses cuBLAS to produce fp32 scores, a full-row two-pass softmax, and fp32 P·V for prompt rows, with a 16-way split-K per-key online softmax for decode. This PR instead targeted sglang-omni's Triton rounding schedule, so it is not a direct patch to the landed attention implementation. We can share the voice-free frozen prefill tensors from our neutral text fixture, including the per-layer attention inputs and outputs, and their mapping to the backbone fields in In our co-batching self-control, the reference itself reached all-row cosine 0.998997 and last-row cosine 0.9982 against its frozen solo output, below the 0.999 bars. We think that is relevant context for interpreting those bars, not proof that no implementation can pass them; the acceptance criteria remain yours to set. Our post-change trace already showed residual differences at layer-0 MoE and layers 1–2 DeltaNet, before the first full-attention layer. That does not establish where the remaining error lies in the landed family, which we have not measured against these frozen tensors. A new attention change would need evidence from that family. |
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
Drama 3's backbone full attention (
drama_attention) runs the reference extend kernel's softmax schedule: 64-keytiles from key 0 with one running max per tile, and Triton's arithmetic for it. Before, each warp kept a running max
per key and the warps were merged at the end. This addresses the accuracy point in the reviews of
#914 and #927.
It is testing, not qualified: it narrows the backbone's gap to the reference on the text fixture, but three
backbone bars still fail (below).
Stacked on #927 (
perf/drama3-pdl); its own commit is thedrama:one, followed by a merge of #927's CI fix. Offered for the combined family to take or leave.Which sums change and why, all inside
drama_attentioninzig/kernels/cuda/drama_backbone_attention.cuh:sm_max[2][8]and one__syncthreadsa tile. Warpwstill holds keysw + 8uof each tile, with the next tile's loads in flight.p = tl_exp(s - m)(ex2.approxofx * log2 e),l = fma(l, c, sum p),o = fma(o, c, sum bf16(p) v)withc = tl_exp(m_old - m); before,expfand unfused multiply-adds per key.o / land the gate's1 / (1 + tl_exp(0 - g))usediv.full, as Triton emits them; before,div.rnandexpf.tl_expanddiv_fullare the existingdrama_gdn_chunk.cuhhelpers. The launch is unchanged.kernel. The reference decodes one request with a split-KV kernel (8 splits of 32-key tiles), not reproduced here.
emulation at layer 3). This change reproduces the reference's rounding; it does not make the math more exact.
Receipt
Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision:
391c64cc, on drama: the backbone chain launches with programmatic dependent launch #927's head with its CI fix (d9869145). No MLX or PyTorch.0.17.0, nvcc 13.0.88 and
CUDA_VISIBLE_DEVICESempty.fishaudio/drama-3ate0a9218a; reference sglang-omniab095ef9.on this change's code (
d2734fd9) against its base on a GB10, with the same results; the end-of-speech and speedresults below come from that run on this change's code. This head's
drama_backbone.cucompiles to the samebytes as the serving build's after the change, at sm_80, 86, 89, 120 and 121 (nvcc 13.0.88, the build's flags).
Exactness, from
tools/bench_concurrent.py --alone --serialortoken_shaagainst"draft": false: n/a fortokens. On the 449-row voice prompt, chunks of 1, 7 and 128 rows and a reset give bit-identical backbone outputs
after the change. Against the reference (backbone-attention-tiles.json, the evidence PR):
182-token text fixture, final hidden state and semantic logits, before → after (bar):
The attention operator alone, on the reference's own layer inputs: bit-equal outputs 70.4% → 98.1% at layer 3,
57.5% → 96.2% at layer 39. The MoE operator replay stayed bit-identical.
The 449-row voice prompt did not improve: relative L2 0.0552 → 0.0558, KL 0.0067 → 0.0087 (bar 0.01), worst
text row 0.9648 → 0.9706, worst audio row 0.9797 → 0.9790.
These fidelity numbers were re-measured on this head on a GB10: the final hidden state is byte-identical to the
serving-build trace above, and the base (drama: the backbone chain launches with programmatic dependent launch #927) reproduces the "before" column byte for byte (7/7 controls).
End of speech, on this head against its base, default voice, no flags, the same 40 test texts, 1024-frame cap; a
stall is 7 s or more under -40 dBFS after the first audible frame:
7, p = 1.0; capped 14 vs 17, p = 0.65).
deterministic trajectory, not a rate.
with a silence guard and long-form settings this series does not have) and does not carry over.
Decode speed before and after, from
tools/bench_openai.py: n/a, because it measures tokens. On a 24-item corpusover two passes (48 requests per build), warm real-time factor 0.659 on this head against 0.671 on its base: no
clear speed difference. Both builds are deterministic, and the base's bodies equal drama: the backbone chain launches with programmatic dependent launch #927's earlier gate on 48 of 48.
Prompt speed before and after: the cold first request's
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. Prompt rowsrun the same kernel, so prompt attention changes too.
Tests run, with pass, fail and skip counts (this head, from an empty local cache):
zig build test --summary all: 86/86 steps, 1412 passed, 0 failed, 7 skipped (JPEG cases; libjpeg-turbo 3.1.4.1is not installed on this machine).
CUDA_VISIBLE_DEVICES= zig build test-cuda-host --summary all: 41/41 steps, 888 passed, 0 failed; the LFM2 JPEGpixel cases inside it return early for the same missing library, so they are not exercised.
zig build test-kernel-contract --summary all: 3/3 steps, 18/18 tests.zig build test-golden --summary all: 4/4 steps.zig build fatbins -Dnvcc=nvcc -Dsm=80,86,89,120,121: 100/100 steps.zig fmt --check build.zig zig: clean.python3 tools/zig/lean_check.py: 517 problems, main's 517 (0 new).git diff --check: clean.Not run, and why:
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. Chunked equals whole on the voice prompt; concurrent-equals-solo is not measured.)
change for fidelity to the reference, with no measurable speed change; the sums are listed above.)
SUMMARYline is in the receipt. (No Drama kernel is registered in this series,drama_attentionincluded.)fidelity oracle needs the reference's frozen outputs and a GPU.)