[Gaurav-Shah05]

Qwen3-TTS 1.7B can now run its CodePredictor's depth loop, sampling included, as one persistent kernel on Hopper.

Each Talker step runs a 5-layer CodePredictor over 16 depth positions to produce a frame's residual codes. On main that loop is 337 projection, 80 attention, 15 sampling and 15 embedding launches per step. On an H200 at c64 it takes 2.6 ms of a 4.8 ms decode step. This PR runs the loop as one kernel. A frame is 1.19x to 1.25x faster at 64 rows in two measurements, and once the GPU is the limit, throughput rises 1.04x to 1.05x at c64. At one request at a time it's 7% slower, since a frame costs about the same at any row count.

The kernel is opt-in (depth_kernel: true in model_kwargs). It's built only for the 1.7B shapes in bf16 on a Hopper GPU with 128 or more SMs, with Triton 3.7, which comes with torch 2.12. Elsewhere, startup logs why and the loop runs. A step the kernel can't take also runs the loop: one with more than 64 rows, or with sampler settings the kernel can't reproduce.

Before #1182, the worker's host time set the pace at c64, so the GPU time the kernel saves didn't reach throughput. #1182 removed that host cost, and the throughput numbers below are measured on top of it.

How it works

• One launch of 132 CTAs, in clusters of two, walks a static task list per CTA built on the host (schedule.py). A frame is 432 phases: qkv, attention, o, gate/up and down for 5 layers at each of 16 positions, plus the input projection, 15 heads, and 16 embedding steps, 15 of which draw a code first.
• A task waits on a counter in global memory, runs one tile and signals a counter. An attention task waits only on the 16 qkv tiles of its kv head, not on the whole phase.
• Weight tiles stream through a ring of 32 shared-memory pages of 4 KB each. The host decides each tile's pages and the task before which its TMA copy goes out, so a tile loads while the CTA waits, as far as the ring allows.
• Projections use the layouts _cp_linear_kernel uses at 64 rows, with gate and up stacked in one 64-row weight tile. Attention uses cp_attend's layouts, and each draw picks the token fused_sample would.
• The o and down projections split each tile's K loop across the two CTAs of a cluster. Rank 1 writes its partial sum into rank 0's shared memory, and rank 0 adds the two, always in the same order. A frame therefore gives the same result every run.
• The sampler resource gains fused_draw_args() and advance_offsets(). A kernel that draws in place reads the step's per-row temperature, top-k, top-p, seed and offset, and after the launch it moves the offsets as that many sample calls would. fused_draw_args() returns None when sample would do more than one plain draw: an eager step, the repetition penalty, min-p, logprobs, a kept last token, or TP. The Talker then runs the loop.
• The kernel isn't built when M* runs it under MPS (CUDA_MPS_PIPE_DIRECTORY set). Its CTAs wait on each other, so all 132 have to be resident at once, and an MPS limit on this process's SMs would keep some of them from ever being scheduled.

Where the depth loop's time goes

I profiled with and without the kernel, at four levels.

• The Talker step. Under nsys at c64 on an H200, a decode step at 59 rows takes 4.36 ms with the kernel. The frame is 2.22 ms of that (51%), the backbone GEMMs 1.04 ms, attention 0.56 ms, and norms, RoPE and SwiGLU 0.4 ms. Without the kernel the loop takes 2.6 ms of the step.
• The floor. Each of a frame's 16 positions reads the 5 layers' 157 MB of weights again, since they don't fit in the 50 MB L2. A frame therefore reads about 2.6 GB with the heads, which takes 0.54 ms at 4.8 TB/s. The kernel runs at 4x that floor, and the loop at 5x.
• Phases. A traced frame on an H200 spends 470 to 480 µs in the handoffs between phases and 1620 to 1700 µs in each phase's slowest task. A task takes about the same time at 1 row as at 64, so the frame costs a fixed latency per phase, 432 times over. That's why a frame takes 2127 µs at 1 row and 2218 µs at 64.
• Inside a task. Eight timestamps per task at 64 rows split the critical path into seven fixed costs, none above 21%:

part	share
input copy while the MMAs run	21%
attention, draw and embed bodies	17%
the phase's last task starting late	17%
weight tile landing after the wait	15%
epilogue (stores and partial sums)	11%
signal	10%
handoff to the next phase	6%

The ring rarely loads a weight tile ahead. One layer needs about 68 pages per CTA against the ring's 32, so each tile waits for the task before it to free its pages.

Decisions and experiments

• K loops rolled up. With every task type in one kernel, the code didn't fit in an SM's instruction cache. Each task ran about 2x slower than the same task alone until its K loops were rolled up.
• Per-head waits. An attention task starts as soon as its own kv head's qkv tiles are done.
• o and down split across a CTA pair. 128 tasks cover the 64 tiles, so 128 of the 132 SMs have work in those phases. The pair's partial sums go through distributed shared memory and are added in a fixed order, with no atomics.
• Fusing phases, rejected. Fusing down into gate/up and o into attention would remove 160 phases. But each layer would then need a sum of partial outputs across tiles, up to 24 MB of fp32 atomics per phase, in an order that changes from run to run. That's as much traffic as the phases it removes, and the frame would no longer give the same result every run.
• Reading only the step's rows, tried and reverted. Copying only the real rows, with the rest of each 64-row tile zero-filled, made 1 and 16 rows slower: 2829 and 2921 µs against 2127 and 2187. So the activation's L2 traffic doesn't set a task's time.
• Two row groups in flight, rejected on analysis. The idea is to run two halves of the rows one phase apart, so each half's handoffs hide behind the other half's work. But a CTA runs its tasks in order on one warpgroup, and a task costs about the same at any row count. Two half-row tasks would cost about two tasks. It only pays with a producer warp that issues the copies ahead, and that needs shared memory the kernel doesn't have (it uses about 213 of 227 KB).
• A 128-row frame, not built yet. Steps over 64 rows run the loop. When I decided this on H100s, c128 was limited by the worker's GIL, so freeing GPU time there wouldn't have raised throughput. With the one-step codec change, the GPU is 90% busy at c128 on an H200, and the kernel gives only 1.01x there because most steps have more than 64 rows.

When the saving reaches throughput

• On Qwen3-TTS: stacked-head Talker, incremental codec, up to 128 rows a step #978 alone, on H100s: frames ran 1.22x faster, but throughput moved only 0.98x to 1.00x at c64, with the GPU 63 to 69% busy. The worker's three Python threads share one GIL, and it was busy nearly all the time.
• On Naomi's host-overhead series (Host overhead 4/N: Rust backend for the micro-scheduler #851 to Host overhead 9/N: Per-step planning: KV page pool and Rust plan columns, one-copy staging, a segmented stager #932): the kernel gave a median 1.04x at c64 over eight pairs, ranging from 0.99x to 1.14x.
• On main on an H200: the GPU idles 18% of the time at c64, because the codec takes one step per chunk width.
• With one codec step for every width (Qwen3-TTS: Decode Codec Chunks of Every Width in One Step #1182): the GPU becomes the limit, and the kernel saves 0.3 to 0.4 ms of the 0.5 ms the frame benchmark predicts at these row counts. A decode step's GPU time falls from 4.78 to 4.49 ms, and a mixed step's from 5.9 to 5.5 ms.

Cost

At small steps the frame is slower than the loop

A frame takes 2.1 to 2.2 ms at any row count, while the loop's time grows with its rows. At 1 row the frame is 0.96x the loop in my timing and 0.89x in Naomi's live test, and by 16 rows it's ahead (1.03x to 1.09x). Her c1 throughput is 0.93x (20.1 against 21.5 audio-s/s). That's acceptable because the kernel is opt-in for deployments that run near capacity, where steps have tens of rows. Running the loop for small steps would remove the cost, once the crossover is measured.

Codes differ from the loop's

The frame's projections use the 64-row tiles at every row count, while cp_linear picks other tiles below 33 rows, and o and down add their two halves in their own order. Given the same prefix at 64 rows, 12.9% of the frame's draws differ from the loop's, and once one does, the rest of that frame's groups follow different codes. The loop's own draws differ by 5.9% and 4.7% between a step at 16 or 32 rows and one at 64, from its own tile choice. That's acceptable because the audio is as good: WER on the c64 audio is the same in both measurements (1.64 to 1.71% with the kernel, 1.67 to 1.73% without, in the live test), and runaway audio is as rare (8 against 7 in 48,000 requests each). The frame's order is fixed, so a seeded request gives the same audio every time on one compile cache.

The frame holds every SM for 2.2 ms

The frame's CTAs wait on each other for the whole frame, so no kernel on another stream can run beside it. A kernel already running delays the frame until it frees its SMs. On main nothing runs beside the frame anyway: the worker's steps run one at a time, and no config runs a node on a side stream.

It needs torch 2.12, where voice cloning doesn't work yet

Triton 3.7 comes with torch 2.12, and torch 2.11 pins Triton 3.6. torchaudio has no build for torch 2.12, and Qwen3TTSModel.load_audio imports it, so on the Base checkpoint every clone request fails in a deployment that can turn the kernel on. That's acceptable because the kernel is off by default and the shipped 1.7B config serves CustomVoice, which never decodes a clip. docs/models.rst says so.

About 720 lines of Gluon and 450 of host-side scheduling, for one model's shapes

The kernel is built for the 1.7B CodePredictor (hidden 1024, 16 q and 8 kv heads of 128, intermediate 3072, 5 layers) and for Hopper's wgmma, TMA and clusters. It's acceptable because it stays inside the Qwen3-TTS package apart from the two sampler methods, and it's off by default.

What's left in the frame, open

Not built yet:

• A 128-row frame. It would cover c128's steps, which now run the loop on a GPU that is 90% busy.
• Per-chunk waits. A projection would start each K chunk as soon as the producer tiles for that chunk finish. That hides most of the late starts and handoffs, about a quarter of the critical path. From the split above I estimate a frame of about 1.8 ms.
• A producer warp with two row groups. A producer warp would issue the next tasks' copies while the two halves of the rows interleave. From the same split I estimate 1.5 to 1.7 ms. The frame is a chain of 432 dependent phases, so even then it stays well above its 0.54 ms weight floor.

Testing

The tests check:

• The task lists: a split tile's two halves run on one cluster, every wait is matched by its phase's signals, no weight tile is written over pages a running task still reads, and a grid that would overwrite the combine buffer is refused.
• The frame against the loop with given codes, at 1, 13 and 64 rows: codes and embedding sums match exactly, the KV cache and logits stay within 2% of the largest KV value and 1.5% of the largest logit, and a second run gives the same result. The loop itself moves by up to 1.1% and 0.8% between a step of the same rows alone and in 64 rows, where cp_linear picks other tiles, and the frame is within 1.14% and 0.96% of the loop over six seeds at 1 to 64 rows.
• The sampled frame: its group 15 is fused_sample's draw from its own logits at the last offset, and its embedding sums match a run given its drawn codes. A captured frame replays what it ran, and the Talker chooses between the frame and the loop as described above.
• The sampler: fifteen fused draws and one advance_offsets give the tokens and offsets of fifteen sample calls, with and without filters. fused_draw_args() returns None for an eager step, the penalty, a step that writes logprobs, a kept last token and a CPU sampler. The Gluon draw picks fused_sample's token bit for bit.

The table's runs used one H200, torch 2.12.1 with Triton 3.7, configs/qwen3tts_1p7b.yaml and the Python frontend. Both arms were on #1182's branch, with the first 800 Seed-TTS English target texts and voice vivian, in closed loop. Each number is the median of 6 passes, with boots alternating:

	#1182	This PR on it
c32 audio-s/s	357.4/359.5	375.4/372.7	1.04x
c64 audio-s/s	537.3/545.1	567.4/565.9	1.05x
c128 audio-s/s	661.1/666.3	676.2/663.3	1.01x
c64 TTFA p50	68/68 ms	65/66 ms	1.04x

• Frame time: captured, on this branch's final commit. The kernel takes 2127, 2192, 2199 and 2218 µs at 1, 16, 32 and 64 rows. The loop takes 2040, 2389, 2594 and 2771 µs, so the ratios are 0.96x, 1.09x, 1.18x and 1.25x. Naomi's live test measured the kernel within 1% of that and the loop about 5% faster (1924, 2273, 2470 and 2647 µs), so 0.89x, 1.03x, 1.12x and 1.19x.
• Naomi's live test (one H200, three interleaved fresh-cache boots per arm, this PR's first head merged with Qwen3-TTS: Decode Codec Chunks of Every Width in One Step #1182): 0.93x, 1.04x, 1.04x and 1.01x audio-s/s at c1, c32, c64 and c128 with the kernel on against off, and the kernel off matches main. Its other checks are in her summary comment.
• WER: measured with whisper-large-v3-turbo on the c64 passes: 1.53% and 1.60% with the kernel, 1.73% and 1.64% without. Seeded audio is identical across the boots of each arm that share a compile cache. A rebuilt cache changes it, with the kernel on or off.
• The table's kernel arm ran under an environment switch, which this PR replaces with depth_kernel. The kernel code is the same. On the final commit, a boot with depth_kernel: true built the frame and ran c64 at 564.0 audio-s/s.
• The torch 2.12 venv has no torchaudio, so test_qwen3_tts_load_audio_rejects_a_long_clip_before_decoding fails there on main too.

[NSagan271]

Mainly looks good (pending some live testing that I have running), left some comments that you should address before merge.

[NSagan271]

Nit / readability: this __init__ is ~70 lines with few spaces and comments. Would add line breaks and one-liner comments between logical sections of this init function, or break it up into sub functions.

[Gaurav-Shah05]

Split into _stack_weights, _allocate_scratch and _compile_sampled in 1a5007f

[NSagan271]

This doesn't refuse a sampler with logprobs on. The description says fused_draw_args() returns None for logprobs (and that the tests check it), but nothing here reads logprob_buf, and test_sampler_fused_draw.py has no logprobs case.

With enable_logprobs on, sample writes logprob_buf, and step_host_rows hands that row back. A kernel taking this path never writes it, so step_host_rows returns whatever the row last held instead of the resource refusing. Qwen3-TTS's code_predictor spec doesn't enable logprobs, so nothing hits this today, but this is the engine API other kernels will use. Adding self.logprob_buf is not None to this condition and a test would make it match the description.

[Gaurav-Shah05]

Thanks, the description was ahead of the branch. af36689 adds self.logprob_buf is not None to the condition.

[NSagan271]

Comment is out of date: an attention task covers 8 rows, not 16. schedule.py says "one q head, 8 rows", and the ty == 4 branch below runs attend4 twice at 4 rows each, from (tile % 8) * 8.

[Gaurav-Shah05]

Fixed in 1a5007f.

[NSagan271]

Nit: nothing in this file uses these. The task loop and _issue dispatch on literals (ty == 0, ty == 4, ...) and base hard-codes MAXT * 8, while schedule.py defines the same names again. Using the names in the dispatch, imported from schedule, would make the loop easier to read and keep one definition of the encoding. If Gluon needs the literals, I'd delete these and say so in a comment.

[Gaurav-Shah05]

The dispatch in _issue and the task loop now uses schedule's names, and base uses ROW, so the encoding is defined once (1a5007f). Gluon reads a module global only when it's a constexpr, so kernel.py wraps schedule's ints in gl.constexpr.

[NSagan271]

This should also say the kernel needs Triton 3.7. torch 2.11 pins triton==3.6.0, so on the torch most installs get, depth_kernel: true only logs a warning and runs the loop. (to be tested live)

There's a second catch for a deployment that moves to torch 2.12 to get it: torchaudio stops at 2.11, and Qwen3TTSModel.load_audio imports torchaudio.functional. The description mentions a test failing for this reason. Does voice cloning on the 1.7B Base checkpoint then fail in the same deployment that turns the kernel on?

[Gaurav-Shah05]

Added in baf6893.

And yes: load_audio imports torchaudio.functional before it reads the clip, and torchaudio stops at 2.11, so every clone request on the Base checkpoint fails in a torch 2.12 deployment.

The docs and the description now say a cloning deployment can't turn the kernel on yet. Our pyproject requires torchaudio, so a plain install resolves torch 2.11 and the kernel stays off. Making cloning work without torchaudio is a change of its own, since eight other modules import it too, api_server/data_worker.py among them.

[NSagan271]

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human. Numbers come from one allocation, so treat small deltas as provisional.

Live test: Qwen3-TTS: Run the 1.7B Depth Loop as One Persistent Kernel on Hopper

Works as an opt-in: the throughput claim holds (1.04x at c64) and WER doesn't move. The kernel costs 7% at c1, frame ratios come in lower than claimed, codes differ from the loop more than "near a tie" suggests, and fused_draw_args() doesn't refuse logprobs.

1x H200, configs/qwen3tts_1p7b.yaml with depth_kernel: true, Python frontend, torch 2.12.1 / Triton 3.7 venv · PR head a8353aee merged with main a9e81ed8 (has #1182), vs main

Claims

Claim	Result
Throughput 1.04x / 1.05x / 1.01x at c32 / c64 / c128	✅ 1.04x / 1.04x / 1.01x (table below)
c64 TTFA p50 1.04x	✅ 65.8 vs 68.0 ms, 1.03x
Frame 2127 / 2192 / 2199 / 2218 µs, loop 2040 / 2389 / 2594 / 2771 µs at 1 / 16 / 32 / 64 rows, 0.96x to 1.25x	⚠️ Kernel within about 1% (2150 / 2204 / 2210 / 2228 µs), but the loop is about 5% faster here (1924 / 2273 / 2470 / 2647 µs), so 0.89x / 1.03x / 1.12x / 1.19x. Captured, real weights, real sampler, 4 processes. The author's timing harness isn't in the PR, so the loop gap is unexplained
Frame matches the loop with given codes	✅ Real weights, 1 / 13 / 64 rows: codes and embedding sums exact, group 15 logits differ by at most 0.375 on a scale of 34
Frame correct against the reference	✅ qwen_tts 0.1.1, 5 sentences: logit max error 0.625 vs the loop's 0.5625, p50 0.0625 for both
Codes differ "near a tie"	⚠️ Same prefix, 64 rows: 12.9% of draws flip vs the loop (283/2199), 88% of rows diverge within a frame's 15 groups. Loop at 16 / 32 rows vs itself at 64: 5.9% / 4.7%. No effect on WER below
WER unchanged at c64	✅ Kernel 1.71 / 1.67 / 1.64%, off 1.73 / 1.67 / 1.69%, main 1.71 / 1.60 / 1.67% (3 boots). Seeds 1-3: 1.74 / 1.83 / 1.80% vs 1.77 / 1.72 / 1.79%
Seeded audio identical across boots	✅ 20/20 files at c1, same boot and across warm-cache boots, on kernel, off and main. A rebuilt compile cache changes it: 0/20 across fresh-cache boots, kernel on and off
fused_draw_args() is None for logprobs	❌ Returns args with logprob_buf set (deterministic probe). Nothing enables logprobs on this resource today
None for min-p; refused for fp32, MPS env, 0.6B shapes, Triton 3.6	✅ min-p, fp32 and MPS by direct probe. 0.6B and torch 2.11 served: logged the reason, 0 frame runs, requests succeeded
Steps over 64 rows run the loop	✅ Frame captured only at buckets up to 64; c128 runs, 0 errors
Author's GPU tests	✅ PR head 21 passed. Merged, adding test_qwen3_tts_model.py: 114 passed, 1 failed (the disclosed torchaudio test)

Risks

What could break	Probe	Result
Replay doesn't run the kernel	Hook on frame.run, worker phases at c64	✅ Captured at buckets 1 to 64. Decode gpu_exec p50 4.35 ms vs 4.77 ms
Frame stalls the plan stream	Worker phases at c64, 2 boots each	✅ pre_plan p50 0.78-0.80 ms vs 0.78-0.82 ms
Mixed sampler settings in one step	Default, greedy, top-k, top-p, low T, 40 in flight	✅ 40/40 served. WER 0%, greedy 1.09% (8 files each). Kernel off, one greedy request ran away (45.65%)
Runaway audio	48,000 requests per arm, c64 and c128, 30 seeds; audio over 3x the sentence's median	✅ 8 vs 7, not separable (Fisher p = 0.5). The loop overruns too
Disconnects	300 dropped mid-stream at 32 in flight	✅ Requests after it served, no tracebacks
Other layouts	Talker on rank 1; two workers on one GPU, c8	✅ Frame captured in both, 0 errors
Startup and memory	2 fresh boots each	✅ 289-313 s vs 316-319 s, RSS 808-819 MB vs 819-858 MB

Performance

tts_speech_bench, first 800 Seed-TTS sentences, vivian, closed loop, seed 1234. Median of 3 interleaved fresh-cache boots, 2 passes each.

Arm	audio-s/s, c1 / c32 / c64 / c128	TTFA p50 ms, c1 / c32 / c64 / c128
main	21.3 / 356.3 / 536.2 / 661.2	26.3 / 52.4 / 68.6 / 120.3
PR, kernel off	21.5 / 356.8 / 537.5 / 658.1	26.1 / 52.2 / 68.0 / 121.5
PR, kernel on	20.1 / 372.8 / 559.4 / 661.9	27.1 / 50.3 / 65.8 / 119.1
On / off	0.93x / 1.04x / 1.04x / 1.01x

• c1 is 7% slower with the kernel. RTF 0.051 vs 0.047, and the boots don't overlap (20.0-20.1 vs 21.4-21.6). This fits the frame being 0.89x at 1 row, though the description's table starts at c32. A deployment that runs mostly at low concurrency loses throughput by turning it on.
• Kernel off matches main. Ranges across boots overlap at every concurrency.
• The step saving is in line with the description. Decode gpu_exec p50 falls 0.42 ms; it has 4.78 to 4.49 ms.

The node is shared and its load varies. The job's pinned CPUs stayed under 5% busy in the logged windows, and arms were interleaved.

Other

• Regressions: none with the kernel off. Other models not run: the sampler gains two methods that only this kernel calls.
• Not tested: voice clone with the kernel (Base doesn't start on main or here on transformers 5.22.1, owned by Fix models that don't start and request inputs that crash the worker #1167; torch 2.12 also has no torchaudio, which load_audio imports), serving under MPS (gate checked by probe only), nsys kernel names.
• Media: available on request.

Client commands

python -m benchmark.tts_speech_bench --engine mstar --url $URL --model qwen3_tts_1p7b --sentences s800.txt \
  --voice vivian --language English --seed 1234 --concurrency {1,32,64,128} --warmup {c} --repeats 2 \
  --num-sentences {40 at c1, else 800} [--save-audio-dir wav]
python -m benchmark.tts_speech_wer --audio-dir wav --sentences s800.txt
MSTAR_PHASE_TIMING=100 mserve qwen3_tts_1p7b --config <config>   # then c64 as above

s800.txt is the first 800 Seed-TTS English target texts. Runaway check: the same bench with --seed varied, 800 sentences.

[Gaurav-Shah05]

@NSagan271 Thanks for the live test. In order:

• c1 at 0.93x: agreed. A frame costs about the same at any row count, and my own timing has it at 0.96x the loop at 1 row. Running the loop below some row count would remove it. I'd measure the crossover between 1 and 16 rows first and do that in a follow-up, since the kernel is opt-in.
• Frame ratios: your kernel times match mine within 1%, and your loop is about 5% faster at every row count, so the difference is in how the loop was timed.
• Code flips: "near a tie" undersold it. The description is updated.