Skip to content

qwen3-tts: one kern manifest with a streamed codec, one graph per decode step - #4

Merged
xiaguan merged 20 commits into
mainfrom
qwen3-tts-kern
Sep 24, 2026
Merged

qwen3-tts: one kern manifest with a streamed codec, one graph per decode step#4
xiaguan merged 20 commits into
mainfrom
qwen3-tts-kern

Conversation

@xiaguan

@xiaguan xiaguan commented Sep 23, 2026

Copy link
Copy Markdown
Contributor

Summary

Qwen3-TTS now runs as one generated kern manifest on one runtime. The talker, code predictor, sampler and a streaming codec decoder all run from it, and each decode step is one CUDA graph launch. omni-cuda and the FlashInfer submodule are gone.

  • Programs: prefill (eager, ragged), first (graph) and decode (graph). A request runs prefill, then first, then decode until it draws the end token.
    • The per-sequence state slot holds the last frame, the repetition bitmap and the codec's streaming state.
    • Talker KV comes from paged leases in kern-pool.
  • Streaming codec: each frame is decoded once. There is no 72-frame window re-decoded per chunk, and --context-frames is gone.
    • The kernels are in kernels/codec.cu: fused and vectorized epilogues, residual GEMMs that accumulate in place, and programmatic dependent launch.
    • One frame for 64 streams takes 1.02 ms of GPU time.
  • Engine: leases each request's KV and slot up front. A step runs start for new requests and step for running ones. A request that can never fit is rejected rather than queued forever.
  • Tools:
    • tools/qwen3_tts/vs_vllm_omni.sh all reruns the comparison with vLLM-Omni and redraws the README chart; its headline ratios are computed from the results.
    • vLLM-Omni now runs its opt-in single-GPU throughput profile under MPS. That is the strongest setup we found for it (c=64: 202 vs 145 audio-s/s for its default deploy).
  • Docs: docs/qwen3-tts.md covers the programs, state, kernels, correctness and performance. docs/bench.md records the codec kernel breakdown and this A/B. docs/qwen3-tts-vs-vllm-omni.md has new results. CLAUDE.md orientation now covers the crate and its tools.

New dependencies:

  • kern-manifest, kern-pool, kern-runtime (git): the manifest schema, KV and state leases, and graph capture and replay.
  • sha2: names the embedded cubins by content.

Removed: omni-cuda, FlashInfer, cudarc, and CI's cuBLAS/cuRAND installs.

Performance

2026-09-23 on one GB300, measured against main with omni-bench, default 2+8 chunks, pinned cores. Values are the mean of two alternating runs.

c TTFP p50 / p99 ms RTF audio-s/s streams without underrun
1 12.2 / 14.7 → 9.9 / 16.4 0.075 → 0.059 13.3 → 16.8 (1.27x) 128/128 → 128/128
8 23.9 / 46.5 → 19.6 / 29.4 0.116 → 0.069 65.4 → 111.5 (1.71x) 128/128 → 128/128
32 42.3 / 147.9 → 21.2 / 29.6 0.252 → 0.084 116.7 → 327.7 (2.81x) 158/256 → 256/256
64 80.4 / 278.8 → 22.5 / 32.6 0.442 → 0.099 133.0 → 563.6 (4.24x) 128/512 → 512/512

Against vLLM-Omni, on its own benchmark with the same 1+25 schedule, pega-omni reaches first audio 3.7-6.6x sooner and serves 1.4-2.5x the audio per second (c=64: 504 vs 202).

Test plan

  • Golden test: runs the real prefill/first/decode path, with the recorded frames forced.
    • Prompt embeddings: worst cosine 0.999993.
    • Talker logits: 0.99990, argmax agreement 98.8%.
    • Predictor logits: 0.99966, argmax agreement 94.1%.
    • Waveform: 36.57 dB SNR for both the streamed and the lagging stream. Two lockstep streams are bit-identical.
  • cargo fmt --check, cargo clippy --all-targets --all-features -D warnings, cargo test --workspace.
  • tools/openai_sdk_check.py against a live server.
  • Whisper spot check in English and Chinese, including concurrent requests in one batch.
  • nsys at c=1: one cuGraphLaunch per decode step.

Follow-ups: #2 (VoiceDesign) and #3 (Base / voice cloning).

🤖 Generated with Claude Code

xiaguan and others added 20 commits September 23, 2026 17:02
@xiaguan @claude
qwen3-tts: the codec decoder streams one frame per step as a kern man…
68c5460 
…ifest

Every running stream's newest frame is decoded in one batched call per step,
its conv, transposed-conv and attention history kept in a per-stream state
slot, instead of re-decoding a 72-frame window per chunk. The call list is a
generated kern manifest over our own kernels (kernels/codec.cu) and
cuBLASLt, run by kern-runtime as one CUDA graph per batch bucket.

New dependencies:
- kern-manifest, kern-pool, kern-runtime (git): manifest schema, per-stream
  state leases and graph capture/replay, so the codec is described as data
  rather than hand-driven launches.
- sha2: names the embedded cubin on disk by content.
- cc (build): locates nvcc to compile kernels/codec.cu to a cubin.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: codec attention applies rope and reads its window one load…
74d48d9 
… deep

RoPE, the KV ring write and attention are one launch per layer; a block
per (row, head) scores a frame per thread and sums V in sixteen groups
instead of a warp walking the 72 frames serially. Codec attention per
step: 222 -> 33 us at batch 1, 231 -> 62 us at batch 64; codec span
772 -> 580 us and 2489 -> 2312 us. Golden SNR 36.89 dB.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: codec elementwise work runs fused and vectorized
0de1de1 
Every elementwise kernel moves eight channels per thread. A residual
unit is im2col (snake of the input plus the carried bias, scattered once
per source row into the dilated taps, history written in the same
launch), GEMM, bias+snake, and a GEMM that accumulates into the residual
stream (cublaslt_bf16_tn_acc), so bias_residual, the separate snake and
ring_write are gone. Conv histories are double-buffered by frame parity
so a launch never overwrites rows its other threads still read. Biases
that only feed linear ops are folded at load (pre_conv into input_proj,
output_proj and ConvNeXt pwconv2 into the next transposed conv's per-tap
bias); the rest are carried to the residual stream's consumers. The
upsamplers' unfold is gone (the GEMM rows already are the output rows),
dwconv+LayerNorm is vectorized with its bias add and history fused,
conv_out snakes each row once into shared memory, norms are a warp per
row.

Codec per step (graph span, codec bench): batch 1 580 -> 487 us, batch
64 2312 -> 1168 us; 199 -> 150 kernels. Golden SNR 36.59 dB.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: codec snake takes the hardware sine after range reduction
9b26049 
The accurate sinf made every SnakeBeta kernel instruction-bound; a
two-step reduction to [-pi, pi] and __sinf keeps the absolute error
near 4e-7. Batch 64 codec span 1168 -> 1077 us (bias_snake 138 -> 91,
im2col_snake 227 -> 196 us per step); batch 1 487 -> 478 us. Golden
SNR 36.57 dB.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: codec kernels launch with programmatic dependent launch
3c59ffd 
Every handwritten kernel waits on griddepcontrol before touching memory
and then triggers its successor, so launch latency overlaps the previous
kernel's tail (cuBLAS launches stay plain). Codec span batch 1
478 -> 422 us, batch 64 1077 -> 1045 us. Golden SNR 36.57 dB.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: codec attention loads its window while rope runs
6d5c748 
The earlier frames' K and V are loaded before the block syncs on the
rotated q, the current frame's K and V are taken from shared memory, so
attention is one memory round trip deep. Attention on the critical path
per step: batch 1 33 -> 25 us, batch 64 74 -> 63 us; codec span 422 ->
415 us and 1045 -> 1037 us. Golden SNR 36.57 dB.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: codec conv_out and dwconv load every tap in one round
2c556f1 
conv_out is a block of 256 per 64 rows: all loads issued up front, each
row snaked once into shared memory, then a thread slides four outputs
over a sixteenth of the channels with the weights in shared memory.
dwconv_ln's seven taps are a compile-time constant, its weights loaded
before the PDL wait and every tap row in one round. Per step: conv_out
batch 1 14.1 -> 4.8 us, batch 64 34.2 -> 23.6 us; dwconv_ln 15 -> 6.7
and 17 -> 8.8 us. Codec span 415 -> 398 us and 1037 -> 1019 us. Golden
SNR 36.57 dB.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: the talker, code predictor, sampler and codec run as one k…
401a3ad 
…ern manifest

One kern runtime executes the whole synthesis path. `model.rs` generates
the manifest: `init` (the pad embedding, once), `prefill` (ragged prompts,
eager), `first` (hidden state to frame to PCM) and `decode` (last frame to
talker to frame to PCM), the last two CUDA graphs per batch bucket, so a
decode step is one graph launch and one read of codes and PCM.

- kernels/talker.cu: our own paged and dense attention, RoPE with QK-norm,
  fused norms, embeddings and a sampler that keeps the Hugging Face order
  (repetition penalty, suppression, temperature, top-k with ties, draw)
  with per-request uniforms; `force` (-1 = draw) lets tests take the real
  path with recorded frames.
- The talker's KV is kern paged state (one state per layer); a sequence's
  slot of the per-sequence state holds its last frame, repetition bitmap
  and codec decoder state, so nothing round-trips through the host.
- manifest.rs: the generator both halves share (was inside codec.rs);
  stack.rs emits a Qwen3 decoder stack for the talker and the predictor.
- The engine leases each request's KV and slot up front and calls
  `start` for new requests and `step` for running ones.
- omni-cuda and the FlashInfer submodule are gone; CI installs nvcc and
  the CUDA headers only.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
docs: qwen3-tts runs as one kern manifest, with its programs, state a…
135b940 
…nd a/b

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
docs: qwen3-tts codec kernels and the streamed codec's a/b against ma…
8ddf16c 
…in are measured

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: merge the codec kernel work onto the one-manifest talker
c08a6c4 
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
tools: vs_vllm_omni.sh all reproduces the readme chart in one run
188fd7e 
`all` sets up the venv, runs vLLM-Omni and then pega-omni alone on the GPU
through the CI points, and renders assets/qwen3-tts-vs-vllm-omni.png with
chart.py, whose headline ratios come from the results. vLLM-Omni now runs its
opt-in single-GPU throughput profile under a private MPS daemon, the
strongest setup we measured for it (c=64: 210-230 vs 181-190 audio-s/s on
the default deploy).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
tools: vs_vllm_omni.sh finds vllm-omni's deploy profiles without impo…
8dafa68 
…rting it and stops mps on exit

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: padded calls fit the page table, unfittable requests are r…
aec88a1 
…ejected, cubins cache per user

- `kv_pages` holds `pages_max + max_seqs` entries: every padding row of a
  graph call names the pad lease's page again, so a nearly full KV plus a
  padded bucket overflowed the input and stopped the engine.
- `Model::open` returns `ExceedsPool` instead of `Busy` when the lease and
  the pad alone exceed the page budget, so such a request is rejected
  instead of sitting at the queue head forever; the log names its prompt
  tokens and frame cap.
- The cubins go to `$XDG_CACHE_HOME` (else `~/.cache`, else the temp dir)
  under `pega-omni/kernels`: a shared `/tmp/pega-omni-kernels` belongs to
  whoever ran first and fails with EACCES for everyone else.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: the kernels share common.cuh and build.rs finds nvcc witho…
7c8e7b3 
…ut the cc crate

- `kernels/common.cuh` holds what `talker.cu` and `codec.cu` both defined:
  the bf16 conversions, 8-wide loads and stores, warp and block reductions,
  the per-sequence slot address and the SiLU-times-up body. Same code, so
  the same results (golden unchanged).
- The norm kernels state their semantics instead of pointing at FlashInfer;
  `codec.cu` names `codec::build`, which generates its calls.
- build.rs used the `cc` crate only to locate nvcc. It now takes `$NVCC`,
  else `$CUDA_PATH/bin/nvcc`, else `/usr/local/cuda/bin/nvcc`, else `nvcc`
  on the PATH, and `cc` is dropped from the workspace dependencies.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
qwen3-tts: the manifest builder is smaller, its state explicit, and t…
a1a874d 
…he crate exposes only what callers use

- `launch` derives a kernel's module (and whether it is PDL-safe) from its
  entry's prefix; `Gen.module` and its assignments are gone.
- `each8` takes rows as a var expression (or a number) and the values per
  row, and replaces every hand-rolled elementwise launch in the talker and
  stack; `each`, `blocks` and `rows_arg` are gone, `i32a` takes any integer
  that fits an i32 and replaces talker's `int`.
- The talker allocates its slot region like every other generator; the
  `region(..) == 0` side-effecting check is gone.
- `Model::load` turns `kv_gib` into KV tokens itself; the engine no longer
  reads the config a second time.
- Only `config`, `engine`, `model`, `prompt` and `weights` are public; `Gen`'s
  fields are private and `bytes_of` (an unsafe generic) is replaced by a
  safe i32 serializer. Unused `Seq::frames`, `Model::limits` and
  `Model::eos` are gone; the engine asks `Model::is_end`.
- Clarity: `Frame` is a struct, `stack::Out` is `Keep`, `Engine::take` is
  `enqueue`, `Stack::forward`'s two attention blocks are one `attention`;
  the codec imports `GROUPS`, sizes its KV ring by `sliding_window` and its
  zero bias by `dim`; the golden test reads the predictor vocabulary from
  the config.
- `--max-step-tokens` says what it counts: prompt tokens prefilled per step.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
ci: the rust job no longer sets LD_LIBRARY_PATH to the CUDA libraries
26cf632 
Nothing links cuBLAS or cuRAND any more and kern-runtime loads CUDA at
runtime, so the tests need only nvcc and the headers to build.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
docs: orientation covers qwen3-tts and its tools, and the codec bench…
76de264 
… method is reproducible

The codec-alone numbers came from a scratch harness against the codec's
own runtime, which no longer exists; bench.md now says so and how to take
the same numbers from the one-manifest model with nsys.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
tools: chart.py takes the vllm-omni version from vs_vllm_omni.sh
d92f52c 
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@xiaguan @claude
docs: the readme chart, the vllm-omni comparison and qwen3-tts perfor…
7e750bf 
…mance are measured on this branch

vLLM-Omni now runs its single-GPU throughput profile under MPS, redrawn by
vs_vllm_omni.sh all; the A/B against main is recorded in bench.md.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: JinYan Su <751080330@qq.com>
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 23, 2026

Copy link
Copy Markdown

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.
To continue using code reviews, you can upgrade your account or add credits to your account and enable them for code reviews in your settings.

@xiaguan
xiaguan merged commit 83b31e1 into main Sep 24, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@xiaguan