Stabilize Higgs Audio termination and prefill

CMakeLists.txt

@@ -170,6 +170,11 @@ option(FLASHRT_ENABLE_MELBAND_ROFORMER
 # Enable explicitly on a Blackwell (sm_120) build.
 option(FLASHRT_ENABLE_OMNIVOICE
        "Build OmniVoice TTS RTX SM120 fused kernels in flash_rt_omnivoice" OFF)
+# Shared delayed-codebook helper used by discrete-audio paths. Kept behind an
+# explicit option so deployment builds that do not serve audio can remove the
+# TU and its binding in lockstep.
+option(FLASHRT_ENABLE_AUDIO_CODEBOOK
+       "Build shared delayed-codebook helper kernels into flash_rt_kernels" ON)
 option(FLASHRT_ENABLE_SM120_DEV_KERNELS
        "Build SM120 probe/tilesweep tuning kernels and public bindings" OFF)
 # ── Slim build for VLA deployments ──
@@ -1189,6 +1194,15 @@ set_target_properties(flash_rt_kernels PROPERTIES
   LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt
   RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt)
 
+if(FLASHRT_ENABLE_AUDIO_CODEBOOK)
+  target_sources(flash_rt_kernels PRIVATE
+    csrc/kernels/delayed_codebook_kernels.cu)
+  target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_AUDIO_CODEBOOK=1)
+  message(STATUS "Delayed audio-codebook kernels: ENABLED")
+else()
+  message(STATUS "Delayed audio-codebook kernels: DISABLED")
+endif()
+
 # ── Motus VAE FP8 quantize kernels (slim-gated) ──
 # These 5 TUs are used only by the Motus VAE FP8 path (the *_ncdhw* /
 # upsample / per-tensor AWQ-FP8 quant helpers; see flash_rt/models/motus/*).

README.md

@@ -139,10 +139,10 @@ DGX Spark / GB10:
 
 #### Higgs Audio v3
 
-| Hardware | Mode | Latency | Throughput | Source |
-|---|---|---:|---:|---|
-| RTX 5090 | FP8 AR decode | **3.2 ms/frame** | RTF **0.095-0.11** | [Higgs performance](docs/higgs_audio_v3.md#performance) |
-| RTX 5090 | BF16 AR decode | **6.1 ms/frame** | RTF **0.15** | [Higgs performance](docs/higgs_audio_v3.md#performance) |
+| Hardware | Mode | Latency | TTFA | Throughput | Source |
+|---|---|---:|---:|---:|---|
+| RTX 5090 | FP8 AR decode | **3.6 ms/frame** | **~79 ms** | RTF **0.09** | [Higgs performance](docs/higgs_audio_v3.md#performance) |
+| RTX 5090 | BF16 AR decode | **6.0 ms/frame** | **~127 ms** | RTF **0.151** | [Higgs performance](docs/higgs_audio_v3.md#performance) |
 
 #### Motus Stage3
 

csrc/bindings.cpp

@@ -152,6 +152,9 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128(
 #endif
 #include "kernels/silu_mul_qwen36.cuh"
 #include "kernels/embedding_lookup_bf16.cuh"
+#ifdef FLASHRT_HAVE_AUDIO_CODEBOOK
+#include "kernels/delayed_codebook_kernels.cuh"
+#endif
 #ifdef FLASHRT_HAVE_QWEN36_KERNELS
 #include "kernels/qwen36_misc.cuh"
 #endif
@@ -4411,6 +4414,18 @@ PYBIND11_MODULE(flash_rt_kernels, m) {
         py::arg("x"), py::arg("W"), py::arg("out"),
         py::arg("M"), py::arg("N"), py::arg("K"), py::arg("stream") = 0);
 
+    m.def("bf16_matmul_cublaslt_bf16",
+        [](uintptr_t x, uintptr_t W, uintptr_t out,
+           int M, int N, int K, uintptr_t stream) {
+            flash_rt::kernels::bf16_matmul_cublaslt_bf16(
+                reinterpret_cast<const __nv_bfloat16*>(x),
+                reinterpret_cast<const __nv_bfloat16*>(W),
+                reinterpret_cast<__nv_bfloat16*>(out),
+                M, N, K, to_stream(stream));
+        },
+        py::arg("x"), py::arg("W"), py::arg("out"),
+        py::arg("M"), py::arg("N"), py::arg("K"), py::arg("stream") = 0);
+
 #ifdef FLASHRT_HAVE_QWEN36_KERNELS
     m.def("bf16_matmul_qwen36_bf16",
         [](uintptr_t x, uintptr_t W, uintptr_t out,
@@ -4505,6 +4520,25 @@ PYBIND11_MODULE(flash_rt_kernels, m) {
         py::arg("token_ids"), py::arg("embed"), py::arg("out"),
         py::arg("rows"), py::arg("hidden"), py::arg("stream") = 0);
 
+#ifdef FLASHRT_HAVE_AUDIO_CODEBOOK
+    m.def("delayed_codebook_argmax_embed_bf16",
+        [](uintptr_t logits, uintptr_t codebook, uintptr_t codes_out,
+           uintptr_t embed_out, int num_codebooks, int codebook_vocab,
+           int hidden, int delay, int boc, uintptr_t stream) {
+            flash_rt::kernels::delayed_codebook_argmax_embed_bf16(
+                reinterpret_cast<const __nv_bfloat16*>(logits),
+                reinterpret_cast<const __nv_bfloat16*>(codebook),
+                reinterpret_cast<int64_t*>(codes_out),
+                reinterpret_cast<__nv_bfloat16*>(embed_out),
+                num_codebooks, codebook_vocab, hidden, delay, boc,
+                to_stream(stream));
+        },
+        py::arg("logits"), py::arg("codebook"), py::arg("codes_out"),
+        py::arg("embed_out"), py::arg("num_codebooks"),
+        py::arg("codebook_vocab"), py::arg("hidden"), py::arg("delay"),
+        py::arg("boc"), py::arg("stream") = 0);
+#endif
+
 #ifdef FLASHRT_HAVE_QWEN36_KERNELS
     m.def("qwen36_embedding_lookup_bf16",
         [](uintptr_t token_ids, uintptr_t embed, uintptr_t out,

csrc/kernels/delayed_codebook_kernels.cu

@@ -0,0 +1,99 @@
+#include "delayed_codebook_kernels.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace flash_rt::kernels {
+namespace {
+
+__global__ void delayed_codebook_argmax_kernel(
+    const __nv_bfloat16* logits,
+    int64_t* codes,
+    int num_codebooks,
+    int codebook_vocab,
+    int delay,
+    int boc) {
+  const int cb = blockIdx.x;
+  if (cb >= num_codebooks) return;
+  if (delay < num_codebooks && cb > delay) {
+    if (threadIdx.x == 0) codes[cb] = static_cast<int64_t>(boc);
+    return;
+  }
+
+  extern __shared__ unsigned char smem[];
+  float* vals = reinterpret_cast<float*>(smem);
+  int* idxs = reinterpret_cast<int*>(vals + blockDim.x);
+
+  const int tid = threadIdx.x;
+  float best = -3.402823466e38f;
+  int best_i = 0;
+  const __nv_bfloat16* row = logits + cb * codebook_vocab;
+  for (int i = tid; i < codebook_vocab; i += blockDim.x) {
+    const float v = __bfloat162float(row[i]);
+    if (v > best || (v == best && i < best_i)) {
+      best = v;
+      best_i = i;
+    }
+  }
+  vals[tid] = best;
+  idxs[tid] = best_i;
+  __syncthreads();
+
+  for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
+    if (tid < stride) {
+      const float ov = vals[tid + stride];
+      const int oi = idxs[tid + stride];
+      if (ov > vals[tid] || (ov == vals[tid] && oi < idxs[tid])) {
+        vals[tid] = ov;
+        idxs[tid] = oi;
+      }
+    }
+    __syncthreads();
+  }
+  if (tid == 0) codes[cb] = static_cast<int64_t>(idxs[0]);
+}
+
+__global__ void delayed_codebook_embed_sum_kernel(
+    const int64_t* codes,
+    const __nv_bfloat16* codebook,
+    __nv_bfloat16* embed,
+    int num_codebooks,
+    int codebook_vocab,
+    int hidden) {
+  const int h = blockIdx.x * blockDim.x + threadIdx.x;
+  if (h >= hidden) return;
+  float acc = 0.0f;
+  for (int cb = 0; cb < num_codebooks; ++cb) {
+    const int code = static_cast<int>(codes[cb]);
+    const int row = cb * codebook_vocab + code;
+    acc += __bfloat162float(codebook[row * hidden + h]);
+  }
+  embed[h] = __float2bfloat16(acc);
+}
+
+}  // namespace
+
+void delayed_codebook_argmax_embed_bf16(
+    const __nv_bfloat16* logits,
+    const __nv_bfloat16* codebook,
+    int64_t* codes_out,
+    __nv_bfloat16* embed_out,
+    int num_codebooks,
+    int codebook_vocab,
+    int hidden,
+    int delay,
+    int boc,
+    cudaStream_t stream) {
+  if (num_codebooks <= 0 || codebook_vocab <= 0 || hidden <= 0) return;
+  const int arg_threads = 1024;
+  const size_t smem = arg_threads * (sizeof(float) + sizeof(int));
+  delayed_codebook_argmax_kernel<<<num_codebooks, arg_threads, smem, stream>>>(
+      logits, codes_out, num_codebooks, codebook_vocab, delay, boc);
+  const int emb_threads = 256;
+  const int emb_blocks = (hidden + emb_threads - 1) / emb_threads;
+  delayed_codebook_embed_sum_kernel<<<emb_blocks, emb_threads, 0, stream>>>(
+      codes_out, codebook, embed_out, num_codebooks, codebook_vocab, hidden);
+}
+
+}  // namespace flash_rt::kernels

csrc/kernels/delayed_codebook_kernels.cuh

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace flash_rt::kernels {
+
+void delayed_codebook_argmax_embed_bf16(
+    const __nv_bfloat16* logits,
+    const __nv_bfloat16* codebook,
+    int64_t* codes_out,
+    __nv_bfloat16* embed_out,
+    int num_codebooks,
+    int codebook_vocab,
+    int hidden,
+    int delay,
+    int boc,
+    cudaStream_t stream);
+
+}  // namespace flash_rt::kernels

docs/benchmark_comparison.md

@@ -159,15 +159,18 @@ speedup.
 
 | Metric | FlashRT FP8 | FlashRT BF16 | SGLang |
 |---|---:|---:|---:|
-| RTF | **0.095-0.11** | **0.15** | 0.16-0.19 |
-| TTFA | **~94 ms** | **~138 ms** | 0.36-0.63 s |
-| Per-frame | **~3.2 ms** | **~6.1 ms** | ~6.4 ms |
+| RTF | **~0.09** | **0.15** | 0.16-0.19 |
+| TTFA | **~79 ms** | **~127 ms** | 0.36-0.63 s |
+| Per-frame | **~3.6 ms** | **~6.0 ms** | ~6.4 ms |
 | VRAM | **6.6 GB** | **9.6 GB** | 28.3 GB reserved |
 
 | Mode | FlashRT | Unoptimized PyTorch reference | Speedup |
 |---|---:|---:|---:|
-| FP8 AR decode | **3.2 ms/frame** | 10.8 ms/frame | **3.3x** |
-| BF16 AR decode | **6.1 ms/frame** | 10.8 ms/frame | **1.8x** |
+| FP8 AR decode | **3.6 ms/frame** | 10.8 ms/frame | **3.0x** |
+| BF16 AR decode | **6.0 ms/frame** | 10.8 ms/frame | **1.8x** |
+
+BF16 prefill-only latency improves from **8.42 → 6.73 ms** at P=6 and
+**11.74 → 6.82 ms** at P=13 in the same single-stream frontend benchmark.
 
 ## Video
 

docs/higgs_audio_v3.md

@@ -27,10 +27,10 @@ also the auto-selection on non-FP8 builds):
 
 | Metric | **FP8** (default) | **BF16** (`fp8=False`) |
 |---|---|---|
-| Real-time factor (RTF) | **0.095 – 0.11** (≈ 9–10× real time) | **0.15** (≈ 6.5× real time) |
-| Time to first audio (TTFA) | **≈ 94 ms** | **≈ 138 ms** |
-| Autoregressive decode | **≈ 3.2 ms/frame** | **≈ 6.1 ms/frame** (at the BF16 bandwidth wall) |
-| Prompt prefill | **≈ 1.0 ms/token** | **≈ 0.9 ms/token** (cuBLAS, weight-once) |
+| Real-time factor (RTF) | **≈ 0.09** (≈ 11× real time) | **0.15** (≈ 6.5× real time) |
+| Time to first audio (TTFA) | **≈ 79 ms** | **≈ 127 ms** |
+| Autoregressive decode | **≈ 3.6 ms/frame** | **≈ 6.0 ms/frame** (at the BF16 bandwidth wall) |
+| Prompt prefill | **≈ 0.42 ms/token** at P=13 (≈ 5.5 ms total) | **≈ 0.52 ms/token** at P=13 (≈ 6.8 ms total; weight-once) |
 | Peak VRAM | **6.6 GB** | **9.6 GB** |
 | Fidelity | teacher-forced logits **cos 1.0**; codec **cos 0.99993**; streamed == one-shot **cos 1.0** | same (bit-exact vs eager) |
 | Prefix reuse | shared `system` preamble cuts prefill **~64 %**, output bit-identical | same |
@@ -40,6 +40,18 @@ Both precisions run the **same** fully-kernelised, zero-torch decode path
 weight bytes of FP8, so its per-frame is ~1.7× — the bandwidth floor, not
 overhead.
 
+The BF16 batched prefill path is also fully kernelised. In a same-prompt
+prefill-only comparison, the BF16 path improves from **8.42 → 6.73 ms** at P=6
+and **11.74 → 6.82 ms** at P=13. The short-prompt EOC regression case
+`"Four score."` now exits in about **1 s** of audio in BF16 instead of running
+to the 40.68 s generation cap.
+
+FP8 activation calibration is fixed to a built-in short calibration prompt
+rather than the user's first request, so startup order no longer changes the
+scale set used for later prompts. A deterministic repeated-code stop guard also
+terminates pathological no-EOC loops before the generation cap; normal EOC
+termination remains unchanged.
+
 **Hardware-adaptive**: the precision is auto-selected from the GPU
 (`fp8=None`, the default) — FP8 where its kernels are compiled in, else BF16;
 pass `fp8=True` / `fp8=False` (or `--bf16`) to force, and forced FP8 falls back
@@ -55,8 +67,8 @@ eager backbone):
 
 | Stage | PyTorch eager | FlashRT FP8 | |
 |---|---|---|---|
-| Autoregressive decode (no codec) | 10.8 ms/frame | **3.2 ms/frame** | **3.3× faster** |
-| Prompt prefill | 3.7 ms/token | **1.0 ms/token** | **3.7× faster** |
+| Autoregressive decode (no codec) | 10.8 ms/frame | **3.6 ms/frame** | **3.0× faster** |
+| Prompt prefill | 3.7 ms/token | **0.42 ms/token** | **8.8× faster** |
 | Backbone weight VRAM | 7.3 GB (bf16) | **3.6 GB (FP8)** | **2× smaller** |
 
 The decode math path is fully kernelised (RMSNorm + quant, dedicated M=1 GEMV
@@ -114,15 +126,15 @@ Expected output on a 5090 (numbers vary with text length and clocks):
 
 ```
 [FP8 W8A8] 'The quick brown fox jumps over the lazy dog.'
-  -> fox.wav  (~3.0s audio, ~4.5s wall incl 1st-call setup)
-  bench 1: AR decode ~290 ms (~3.8 ms/frame)
-  bench 2: AR decode ~290 ms (~3.8 ms/frame)
+  -> fox.wav  (~4.0s audio, ~3s wall incl 1st-call setup)
+  bench 1: AR decode ~355 ms (~3.6 ms/frame)
+  bench 2: AR decode ~355 ms (~3.6 ms/frame)
   ...
 ```
 
 First call pays a one-time cost: FP8 activation-scale **calibration** (a short
-BF16 free-run) and **codec load**. Subsequent calls are warm. Add `--bf16` to
-run the BF16 backbone instead of FP8.
+BF16 free-run on a built-in calibration prompt) and **codec load**. Subsequent
+calls are warm. Add `--bf16` to run the BF16 backbone instead of FP8.
 
 ---
 
@@ -212,18 +224,35 @@ teacher-forced cosine above, and the codec is bit-faithful on identical codes.
 
 Headline numbers are in [Performance](#performance) above. Methodology:
 
+Public AR decode harness:
+
+```bash
+python examples/higgs_audio_v3_quickstart.py \
+    --checkpoint <checkpoint> \
+    --text "The quick brown fox jumps over the lazy dog." \
+    --benchmark 10
+
+python examples/higgs_audio_v3_quickstart.py \
+    --checkpoint <checkpoint> \
+    --text "The quick brown fox jumps over the lazy dog." \
+    --bf16 --benchmark 10
+```
+
 - **Full pipeline, warm.** RTF / TTFA are end-to-end text→waveform through the
   standardized `generate` / `generate_stream` frontend, FP8 backbone, after a
   warm-up call (lazy FP8 calibration + codec load + CUDA-graph capture happen
   once on the first call and are excluded).
-- **Decode floor 3.2 ms/frame** is the clean single-position CUDA-graph replay;
+- **Decode floor 3.6 ms/frame** is the clean single-position CUDA-graph replay;
   the per-token GEMMs read 3.6 GB of distinct FP8 weights, so this is genuinely
   HBM-bound (micro-benchmarks that reuse weights report L2-cached fiction). The
   full-pipeline per-frame is slightly higher because attention cost grows with KV
   length over a long generation.
-- **Prefill** is one batched M=P forward (≈ 1 ms/token); a shared `system`
+- **Prefill** is one batched M=P forward (FP8 ≈ 5.5 ms total; BF16 ≈ 6.8 ms total
+  for a 13-token prompt in the benchmark sentence); a shared `system`
   preamble reuses its resident KV across requests (only the new text suffix is
-  prefilled — bit-identical to a cold prefill).
+  prefilled — bit-identical to a cold prefill). Prefill-only numbers time
+  `set_prompt(...)` + `prefill()` after one warm-up and report the median of 20
+  iterations.
 - **Codec** runs in fp32 (ConvTranspose is unstable in low precision) as one
   small pass at the end (≤ 50 ms for 40 s of audio); in streaming it is decoded
   in overlapping windows so the streamed waveform matches the one-shot output.
@@ -235,8 +264,8 @@ Headline numbers are in [Performance](#performance) above. Methodology:
 ## 8. Notes & limitations
 
 - **FP8 calibration** is per-tensor static (activation `amax/448`), measured
-  once from a short BF16 free-run of the first prompt and reused. Activation
-  ranges are stable across prompts; re-instantiate the frontend to recalibrate.
+  once from a short BF16 free-run of a built-in calibration prompt and reused.
+  This keeps activation scales independent of the first user request.
 - The BF16 projection weights are freed after calibration (the FP8 backbone is
   the active path); pass `fp8=False` for the BF16 backbone, which keeps them.
 - **GPU support / precision auto-selection.** The build auto-detects the GPU