From 415385c7a57dbfa0cb4a3123a8a76396bd165bba Mon Sep 17 00:00:00 2001
From: LiangSu8899 <7thuniversels@gmail.com>
Date: Thu, 9 Jul 2026 04:52:28 -0400
Subject: [PATCH 1/5] fix(higgs): stabilize batched prefill eoc

---
 CMakeLists.txt                                |  1 +
 csrc/bindings.cpp                             | 30 ++++++
 csrc/kernels/higgs_audio_v3_kernels.cu        | 99 +++++++++++++++++++
 csrc/kernels/higgs_audio_v3_kernels.cuh       | 21 ++++
 .../frontends/torch/_higgs_audio_v3_bf16.py   | 67 +++++++------
 .../frontends/torch/_higgs_audio_v3_fp8.py    | 44 +++++----
 .../frontends/torch/higgs_audio_v3_rtx.py     | 43 +++++---
 tests/test_higgs_audio_v3_eoc.py              | 21 ++++
 8 files changed, 262 insertions(+), 64 deletions(-)
 create mode 100644 csrc/kernels/higgs_audio_v3_kernels.cu
 create mode 100644 csrc/kernels/higgs_audio_v3_kernels.cuh
 create mode 100644 tests/test_higgs_audio_v3_eoc.py

diff --git a/CMakeLists.txt b/CMakeLists.txt
index cb4079f0..391b046e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1113,6 +1113,7 @@ pybind11_add_module(flash_rt_kernels
   csrc/gemm/fp8_block128_gemm.cu
   csrc/kernels/silu_mul_qwen36.cu
   csrc/kernels/embedding_lookup_bf16.cu
+  csrc/kernels/higgs_audio_v3_kernels.cu
   # qwen3_5_moe-family RTX SM120 kernels are NOT listed here -- they are
   # gated into the module below (FLASHRT_ENABLE_QWEN35MOE), so SM89/87/110
   # and non-Nex-N2 builds never compile them. See the gated block after the
diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp
index 9f09c598..4022986f 100644
--- a/csrc/bindings.cpp
+++ b/csrc/bindings.cpp
@@ -146,6 +146,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128(
 #endif
 #include "kernels/silu_mul_qwen36.cuh"
 #include "kernels/embedding_lookup_bf16.cuh"
+#include "kernels/higgs_audio_v3_kernels.cuh"
 #ifdef FLASHRT_HAVE_QWEN36_KERNELS
 #include "kernels/qwen36_misc.cuh"
 #endif
@@ -4246,6 +4247,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,
@@ -4340,6 +4353,23 @@ 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);
 
+    m.def("higgs_audio_v3_argmax_delay_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::higgs_audio_v3_argmax_delay_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);
+
 #ifdef FLASHRT_HAVE_QWEN36_KERNELS
     m.def("qwen36_embedding_lookup_bf16",
         [](uintptr_t token_ids, uintptr_t embed, uintptr_t out,
diff --git a/csrc/kernels/higgs_audio_v3_kernels.cu b/csrc/kernels/higgs_audio_v3_kernels.cu
new file mode 100644
index 00000000..e839b157
--- /dev/null
+++ b/csrc/kernels/higgs_audio_v3_kernels.cu
@@ -0,0 +1,99 @@
+#include "higgs_audio_v3_kernels.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace flash_rt::kernels {
+namespace {
+
+__global__ void higgs_argmax_delay_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 higgs_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 higgs_audio_v3_argmax_delay_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));
+  higgs_argmax_delay_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;
+  higgs_embed_sum_kernel<<<emb_blocks, emb_threads, 0, stream>>>(
+      codes_out, codebook, embed_out, num_codebooks, codebook_vocab, hidden);
+}
+
+}  // namespace flash_rt::kernels
diff --git a/csrc/kernels/higgs_audio_v3_kernels.cuh b/csrc/kernels/higgs_audio_v3_kernels.cuh
new file mode 100644
index 00000000..269e0e6d
--- /dev/null
+++ b/csrc/kernels/higgs_audio_v3_kernels.cuh
@@ -0,0 +1,21 @@
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace flash_rt::kernels {
+
+void higgs_audio_v3_argmax_delay_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
diff --git a/flash_rt/frontends/torch/_higgs_audio_v3_bf16.py b/flash_rt/frontends/torch/_higgs_audio_v3_bf16.py
index 6fcb7ff1..77990ccb 100644
--- a/flash_rt/frontends/torch/_higgs_audio_v3_bf16.py
+++ b/flash_rt/frontends/torch/_higgs_audio_v3_bf16.py
@@ -18,10 +18,8 @@
 """
 from __future__ import annotations
 
-from typing import Any
-
 import torch
-import torch.nn.functional as F
+from typing import Any
 
 BF16 = torch.bfloat16
 
@@ -177,17 +175,16 @@ def prefill_batched(self, ids: list[int], start_pos: int = 0) -> torch.Tensor:
         H, NQ, NKV, HD, INTER = self.H, self.NQ, self.NKV, self.HD, self.INTER
         NQK, KV, EPS = self.NQK, self.KV, self.EPS
         s = torch.cuda.current_stream().cuda_stream
-        # One-time eager prefill GEMM: cuBLAS (torch.matmul) reads each weight
-        # ONCE across the P prompt rows (5-14x the warp-per-row bf16_matmul,
-        # which re-reads W per row -> 11%/5% effective BW at P=13/33). This is
-        # the one-time setup path, not the per-frame decode graph (which stays
-        # the hand-written bf16 GEMV); the norms/qk-norm+RoPE/silu/attention all
-        # stay kernelised. mv is the M=1 head GEMV.
+        # Keep the prompt path in FlashRT-owned kernels. GEMMs use the cuBLASLt
+        # binding for throughput; q/k norm + RoPE + KV write use the batched
+        # prefill kernels so the prompt KV is produced by one consistent path
+        # instead of a Python loop over strided slices.
         mv = fvk.bf16_matvec_qwen36_bf16
+        fast_gemm = fvk.bf16_matmul_cublaslt_bf16
         rc, rs = fe._rope_cos, fe._rope_sin
 
         ids_t = torch.tensor(ids[start_pos:], device=dev)
-        h = F.embedding(ids_t, fe._weights["text_embed"]).contiguous()  # [S,H]
+        h = torch.empty(S, H, device=dev, dtype=BF16)
         xn = torch.empty(S, H, device=dev, dtype=BF16)
         xn2 = torch.empty(S, H, device=dev, dtype=BF16)
         Dq = torch.empty(S, NQK + 2 * KV, device=dev, dtype=BF16)
@@ -195,42 +192,50 @@ def prefill_batched(self, ids: list[int], start_pos: int = 0) -> torch.Tensor:
         act = torch.empty(S, INTER, device=dev, dtype=BF16)
         tmp = torch.empty(S, H, device=dev, dtype=BF16)
 
+        fvk.embedding_lookup_bf16(ids_t.data_ptr(),
+                                  fe._weights["text_embed"].data_ptr(),
+                                  h.data_ptr(), S, H, s)
         fvk.rms_norm(h.data_ptr(), self.WL[0]["in_norm"].data_ptr(),
                      xn.data_ptr(), S, H, EPS, s)
         for L in range(self.NL):
             w = self.WL[L]
-            torch.matmul(xn, w["qkv"].t(), out=Dq)
-            for j in range(S):
-                pos = start_pos + j
-                qj = Dq[j, :NQK].contiguous()
-                kj = Dq[j, NQK:NQK + KV].contiguous()
-                vj = Dq[j, NQK + KV:].contiguous()
-                fvk.qwen3_q_norm_rope_qstage_bf16(
-                    q_pre=qj.data_ptr(), q_norm_w=w["qn"].data_ptr(),
-                    cos=rc[pos].data_ptr(), sin=rs[pos].data_ptr(),
-                    q_buf_dst=be.Q_buf[:, j].data_ptr(), n_q_heads=NQ, eps=EPS,
-                    stream=s)
-                fvk.qwen3_k_norm_rope_kvwrite_bf16(
-                    k_pre=kj.data_ptr(), v_pre=vj.data_ptr(),
-                    k_norm_w=w["kn"].data_ptr(),
-                    cos=rc[pos].data_ptr(), sin=rs[pos].data_ptr(),
-                    k_cache_dst=be.K_cache[L, pos].data_ptr(),
-                    v_cache_dst=be.V_cache[L, pos].data_ptr(),
-                    n_kv_heads=NKV, eps=EPS, stream=s)
+            fast_gemm(xn.data_ptr(), w["qkv"].data_ptr(), Dq.data_ptr(),
+                      S, NQK + 2 * KV, H, s)
+            q = Dq[:, :NQK]
+            k = Dq[:, NQK:NQK + KV]
+            v = Dq[:, NQK + KV:]
+            fvk.qwen3_q_norm_rope_qstage_prefill_bf16(
+                q.data_ptr(), w["qn"].data_ptr(),
+                rc[start_pos:start_pos + S].data_ptr(),
+                rs[start_pos:start_pos + S].data_ptr(),
+                be.Q_buf[:, :S].data_ptr(), NQ, S, int(q.stride(0)),
+                int(be.Q_buf[:, :S].stride(1)), EPS, s)
+            fvk.qwen3_k_norm_rope_kvwrite_prefill_bf16(
+                k.data_ptr(), v.data_ptr(), w["kn"].data_ptr(),
+                rc[start_pos:start_pos + S].data_ptr(),
+                rs[start_pos:start_pos + S].data_ptr(),
+                be.K_cache[L, start_pos:start_pos + S].data_ptr(),
+                be.V_cache[L, start_pos:start_pos + S].data_ptr(),
+                NKV, S, int(k.stride(0)),
+                int(be.K_cache[L, start_pos:start_pos + S].stride(0)),
+                EPS, s)
             be.O_buf[:, :S].zero_()
             be.run("full", L, S, kv_seq=P, causal=True, stream=s,
                    softmax_scale=HD ** -0.5)
             ao = be.O_buf[:, :S].reshape(S, NQK).contiguous()
-            torch.matmul(ao, w["o"].t(), out=tmp)
+            fast_gemm(ao.data_ptr(), w["o"].data_ptr(), tmp.data_ptr(),
+                      S, H, NQK, s)
             fvk.residual_add_rms_norm(h.data_ptr(), tmp.data_ptr(),
                                       w["post_norm"].data_ptr(), xn2.data_ptr(),
                                       S, H, EPS, s)
-            torch.matmul(xn2, w["gu"].t(), out=Dg)
+            fast_gemm(xn2.data_ptr(), w["gu"].data_ptr(), Dg.data_ptr(),
+                      S, 2 * INTER, H, s)
             g = Dg[:, :INTER].contiguous()
             u = Dg[:, INTER:].contiguous()
             fvk.silu_mul_qwen36_bf16(g.data_ptr(), u.data_ptr(), act.data_ptr(),
                                      S * INTER, s)
-            torch.matmul(act, w["down"].t(), out=tmp)
+            fast_gemm(act.data_ptr(), w["down"].data_ptr(), tmp.data_ptr(),
+                      S, H, INTER, s)
             if L < self.NL - 1:
                 fvk.residual_add_rms_norm(
                     h.data_ptr(), tmp.data_ptr(),
diff --git a/flash_rt/frontends/torch/_higgs_audio_v3_fp8.py b/flash_rt/frontends/torch/_higgs_audio_v3_fp8.py
index 4986bb62..eb25fbb5 100644
--- a/flash_rt/frontends/torch/_higgs_audio_v3_fp8.py
+++ b/flash_rt/frontends/torch/_higgs_audio_v3_fp8.py
@@ -268,7 +268,7 @@ def prefill_batched(self, ids: list[int], start_pos: int = 0) -> torch.Tensor:
         rc, rs = fe._rope_cos, fe._rope_sin
 
         ids_t = torch.tensor(ids[start_pos:], device=dev)
-        h = F.embedding(ids_t, fe._weights["text_embed"]).contiguous()  # [S,H]
+        h = torch.empty(S, H, device=dev, dtype=BF16)
         fp8a = torch.empty(S, H, device=dev, dtype=F8)
         Dq = torch.empty(S, NQK + 2 * KV, device=dev, dtype=BF16)
         Dg = torch.empty(S, 2 * INTER, device=dev, dtype=BF16)
@@ -277,6 +277,9 @@ def prefill_batched(self, ids: list[int], start_pos: int = 0) -> torch.Tensor:
         tmp = torch.empty(S, H, device=dev, dtype=BF16)
         act = torch.empty(S, INTER, device=dev, dtype=BF16)
 
+        fvk.embedding_lookup_bf16(ids_t.data_ptr(),
+                                  fe._weights["text_embed"].data_ptr(),
+                                  h.data_ptr(), S, H, s)
         for L in range(self.NL):
             w = self.WL[L]
             fvk.rms_norm_fp8(h.data_ptr(), w["in_norm"].data_ptr(),
@@ -284,23 +287,24 @@ def prefill_batched(self, ids: list[int], start_pos: int = 0) -> torch.Tensor:
                              self.DS[L][0].data_ptr(), s)
             gemm(fp8a.data_ptr(), w["qkvq"].data_ptr(), Dq.data_ptr(),
                  S, NQK + 2 * KV, H, self.ALP[L][0] * w["qkvs"], s)
-            for j in range(S):
-                pos = start_pos + j
-                qj = Dq[j, :NQK].contiguous()
-                kj = Dq[j, NQK:NQK + KV].contiguous()
-                vj = Dq[j, NQK + KV:].contiguous()
-                fvk.qwen3_q_norm_rope_qstage_bf16(
-                    q_pre=qj.data_ptr(), q_norm_w=w["qn"].data_ptr(),
-                    cos=rc[pos].data_ptr(), sin=rs[pos].data_ptr(),
-                    q_buf_dst=be.Q_buf[:, j].data_ptr(), n_q_heads=NQ,
-                    eps=EPS, stream=s)
-                fvk.qwen3_k_norm_rope_kvwrite_bf16(
-                    k_pre=kj.data_ptr(), v_pre=vj.data_ptr(),
-                    k_norm_w=w["kn"].data_ptr(),
-                    cos=rc[pos].data_ptr(), sin=rs[pos].data_ptr(),
-                    k_cache_dst=be.K_cache[L, pos].data_ptr(),
-                    v_cache_dst=be.V_cache[L, pos].data_ptr(),
-                    n_kv_heads=NKV, eps=EPS, stream=s)
+            q = Dq[:, :NQK]
+            k = Dq[:, NQK:NQK + KV]
+            v = Dq[:, NQK + KV:]
+            fvk.qwen3_q_norm_rope_qstage_prefill_bf16(
+                q.data_ptr(), w["qn"].data_ptr(),
+                rc[start_pos:start_pos + S].data_ptr(),
+                rs[start_pos:start_pos + S].data_ptr(),
+                be.Q_buf[:, :S].data_ptr(), NQ, S, int(q.stride(0)),
+                int(be.Q_buf[:, :S].stride(1)), EPS, s)
+            fvk.qwen3_k_norm_rope_kvwrite_prefill_bf16(
+                k.data_ptr(), v.data_ptr(), w["kn"].data_ptr(),
+                rc[start_pos:start_pos + S].data_ptr(),
+                rs[start_pos:start_pos + S].data_ptr(),
+                be.K_cache[L, start_pos:start_pos + S].data_ptr(),
+                be.V_cache[L, start_pos:start_pos + S].data_ptr(),
+                NKV, S, int(k.stride(0)),
+                int(be.K_cache[L, start_pos:start_pos + S].stride(0)),
+                EPS, s)
             # Pre-zero the query slots: the causal FA2 path can leave a partial
             # query-tile slot unwritten, and reading stale O_buf makes the
             # prefill non-deterministic run-to-run (frame-0 near-tie argmax
@@ -314,7 +318,7 @@ def prefill_batched(self, ids: list[int], start_pos: int = 0) -> torch.Tensor:
                                     self.DS[L][1].data_ptr(), S * NQK, s)
             gemm(fp8o.data_ptr(), w["oq"].data_ptr(), tmp.data_ptr(),
                  S, H, NQK, self.ALP[L][1] * w["os"], s)
-            h = (h.float() + tmp.float()).to(BF16)
+            fvk.residual_add(h.data_ptr(), tmp.data_ptr(), S * H, s)
             fvk.rms_norm_fp8(h.data_ptr(), w["post_norm"].data_ptr(),
                              fp8a.data_ptr(), S, H, EPS,
                              self.DS[L][2].data_ptr(), s)
@@ -328,7 +332,7 @@ def prefill_batched(self, ids: list[int], start_pos: int = 0) -> torch.Tensor:
                                     self.DS[L][3].data_ptr(), S * INTER, s)
             gemm(fp8d.data_ptr(), w["dnq"].data_ptr(), tmp.data_ptr(),
                  S, H, INTER, self.ALP[L][3] * w["dns"], s)
-            h = (h.float() + tmp.float()).to(BF16)
+            fvk.residual_add(h.data_ptr(), tmp.data_ptr(), S * H, s)
 
         hlast = h[S - 1:S].contiguous()
         fvk.rms_norm(hlast.data_ptr(), fe._weights["final_norm"].data_ptr(),
diff --git a/flash_rt/frontends/torch/higgs_audio_v3_rtx.py b/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
index ba8a47dd..ed3ff0fd 100644
--- a/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
+++ b/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
@@ -20,13 +20,21 @@
 from typing import Any
 
 import torch
-import torch.nn.functional as F
 
 from flash_rt.models.higgs_audio_v3.pipeline_rtx import HiggsAudioV3Dims
 
 _SPECIALS = ("<|tts|>", "<|text|>", "<|audio|>")
 
 
+def _delayed_eoc_countdown(codes: torch.Tensor, nc: int, eoc: int) -> int | None:
+    """Remaining delayed rows to flush once any codebook emits EOC."""
+    eoc_pos = (codes == eoc).nonzero(as_tuple=False)
+    if not eoc_pos.numel():
+        return None
+    first_eoc_cb = int(eoc_pos[0])
+    return max(1, nc - 1 - first_eoc_cb)
+
+
 class HiggsAudioV3TorchFrontendRtx:
     """Single-stream Higgs Audio v3 TTS inference on RTX SM120 (BF16 backbone)."""
 
@@ -84,6 +92,8 @@ def __init__(self, checkpoint_path: str, *,
         self._dec: Any = None                          # active decode engine
         self._fp8_decoder: Any = None                  # back-compat alias
         self._codec: Any = None
+        self._codes_dev = None
+        self._embed_buf = None
         self.latency_records: list[float] = []
 
         self._load_weights()
@@ -242,6 +252,9 @@ def _alloc_buffers(self) -> None:
         nc = self._cfg["num_codebooks"]
         self._cb_offsets = (
             torch.arange(nc, device=self.device) * self._cfg["codebook_vocab"])
+        self._codes_dev = torch.empty(nc, device=self.device, dtype=torch.long)
+        self._embed_buf = torch.empty(
+            1, self._cfg["hidden"], device=self.device, dtype=torch.bfloat16)
 
     def _build_rope_table(self) -> None:
         hd = self._cfg["head_dim"]
@@ -253,11 +266,6 @@ def _build_rope_table(self) -> None:
         self._rope_cos = f.cos().to(torch.bfloat16).contiguous()
         self._rope_sin = f.sin().to(torch.bfloat16).contiguous()
 
-    def _embed_codes(self, codes):
-        cb = self._weights["codebook"]
-        ids = codes.to(self.device).long() + self._cb_offsets
-        return F.embedding(ids, cb).sum(0, keepdim=True)
-
     # ── Public API ──
 
     def build_prompt(self, text: str, system: str | None = None) -> list[int]:
@@ -384,7 +392,13 @@ def prefill(self, prompt_ids: list[int] | None = None,
         else:
             te = self._weights["text_embed"]
             for t, tok in enumerate(self._prompt_ids):
-                row = F.embedding(torch.tensor([tok], device=self.device), te)
+                tok_t = torch.tensor([tok], device=self.device)
+                row = torch.empty(1, self._cfg["hidden"], device=self.device,
+                                  dtype=torch.bfloat16)
+                fvk.embedding_lookup_bf16(tok_t.data_ptr(), te.data_ptr(),
+                                          row.data_ptr(), 1,
+                                          self._cfg["hidden"],
+                                          torch.cuda.current_stream().cuda_stream)
                 self._gen_logits = self._frame_logits(fvk, row, t)
         self._gen_pos = P
         self._resident_ids = list(self._prompt_ids)
@@ -400,17 +414,20 @@ def decode_stream(self):
         delay, eoc_countdown, done = 0, None, False
         window: list[torch.Tensor] = []
         for j in range(self.max_new_frames):
-            codes = logits.argmax(-1).clone()
+            fvk.higgs_audio_v3_argmax_delay_embed_bf16(
+                logits.data_ptr(), self._weights["codebook"].data_ptr(),
+                self._codes_dev.data_ptr(), self._embed_buf.data_ptr(),
+                nc, self._cfg["codebook_vocab"], self._cfg["hidden"],
+                delay, boc, torch.cuda.current_stream().cuda_stream)
+            codes = self._codes_dev.cpu()
             if delay < nc:
-                if delay + 1 < nc:
-                    codes[delay + 1:] = boc
                 delay += 1
             elif eoc_countdown is not None:
                 eoc_countdown -= 1
                 if eoc_countdown <= 0:
                     done = True
-            elif int(codes[0]) == eoc:
-                eoc_countdown = nc - 2
+            else:
+                eoc_countdown = _delayed_eoc_countdown(codes, nc, eoc)
             if done:
                 break
             window.append(codes.clone())
@@ -418,7 +435,7 @@ def decode_stream(self):
                 base = len(window) - nc
                 yield torch.stack(
                     [window[base + i][i] for i in range(nc)]).cpu()
-            logits = self._decode_logits(fvk, self._embed_codes(codes), P + j)
+            logits = self._decode_logits(fvk, self._embed_buf, P + j)
 
     @torch.no_grad()
     def predict(self, text: str | None = None) -> torch.Tensor:
diff --git a/tests/test_higgs_audio_v3_eoc.py b/tests/test_higgs_audio_v3_eoc.py
new file mode 100644
index 00000000..b0d1f8ac
--- /dev/null
+++ b/tests/test_higgs_audio_v3_eoc.py
@@ -0,0 +1,21 @@
+import torch
+
+from flash_rt.frontends.torch.higgs_audio_v3_rtx import _delayed_eoc_countdown
+
+
+def test_delayed_eoc_countdown_flushes_full_tail():
+    nc = 8
+    eoc = 1025
+    codes = torch.tensor([eoc, 1, 2, 3, 4, 5, 6, 7])
+    assert _delayed_eoc_countdown(codes, nc, eoc) == nc - 1
+
+
+def test_delayed_eoc_countdown_accepts_later_codebook_eoc():
+    nc = 8
+    eoc = 1025
+    codes = torch.tensor([10, 11, 12, eoc, 14, 15, 16, 17])
+    assert _delayed_eoc_countdown(codes, nc, eoc) == 4
+
+
+def test_delayed_eoc_countdown_absent():
+    assert _delayed_eoc_countdown(torch.arange(8), 8, 1025) is None

From 7fffe6058db09e2c86d1cf22dd24e3af41cd850e Mon Sep 17 00:00:00 2001
From: LiangSu8899 <7thuniversels@gmail.com>
Date: Thu, 9 Jul 2026 05:10:38 -0400
Subject: [PATCH 2/5] Document Higgs BF16 prefill performance

---
 README.md                    |  2 +-
 docs/benchmark_comparison.md |  9 ++++++---
 docs/higgs_audio_v3.md       | 33 ++++++++++++++++++++++++++++-----
 3 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index a503c5f4..737757b3 100644
--- a/README.md
+++ b/README.md
@@ -142,7 +142,7 @@ DGX Spark / GB10:
 | 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) |
+| RTX 5090 | BF16 AR decode | **6.0 ms/frame** | RTF **0.151** | [Higgs performance](docs/higgs_audio_v3.md#performance) |
 
 #### Motus Stage3
 
diff --git a/docs/benchmark_comparison.md b/docs/benchmark_comparison.md
index c9efb61d..7774787a 100644
--- a/docs/benchmark_comparison.md
+++ b/docs/benchmark_comparison.md
@@ -160,14 +160,17 @@ 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 |
+| TTFA | **~94 ms** | **~127 ms** | 0.36-0.63 s |
+| Per-frame | **~3.2 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** |
+| BF16 AR decode | **6.0 ms/frame** | 10.8 ms/frame | **1.8x** |
+
+BF16 prefill-only latency improves from **8.42 → 6.79 ms** at P=6 and
+**11.74 → 6.86 ms** at P=13 in the same single-stream frontend benchmark.
 
 ## Video
 
diff --git a/docs/higgs_audio_v3.md b/docs/higgs_audio_v3.md
index 7cc4e0ef..3a0f35b4 100644
--- a/docs/higgs_audio_v3.md
+++ b/docs/higgs_audio_v3.md
@@ -28,9 +28,9 @@ 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) |
+| Time to first audio (TTFA) | **≈ 94 ms** | **≈ 127 ms** |
+| Autoregressive decode | **≈ 3.2 ms/frame** | **≈ 6.0 ms/frame** (at the BF16 bandwidth wall) |
+| Prompt prefill | **≈ 1.0 ms/token** | **≈ 0.53 ms/token** at P=13 (≈ 6.9 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,12 @@ 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.79 ms** at P=6
+and **11.74 → 6.86 ms** at P=13. The short-prompt EOC regression case
+`"Four score."` now exits at **1.20 s** of audio in BF16 instead of running to
+the 40.68 s generation cap.
+
 **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
@@ -212,6 +218,20 @@ 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
@@ -221,9 +241,12 @@ Headline numbers are in [Performance](#performance) above. Methodology:
   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 ≈ 1 ms/token; BF16 ≈ 6.9 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.

From c7f8cdbcbd35b684b3478da3db8999de93f28686 Mon Sep 17 00:00:00 2001
From: LiangSu8899 <7thuniversels@gmail.com>
Date: Thu, 9 Jul 2026 05:25:46 -0400
Subject: [PATCH 3/5] Stabilize Higgs FP8 calibration

---
 README.md                                     |  2 +-
 docs/benchmark_comparison.md                  | 12 ++---
 docs/higgs_audio_v3.md                        | 44 +++++++++++--------
 .../frontends/torch/higgs_audio_v3_rtx.py     | 17 ++++++-
 tests/test_higgs_audio_v3_eoc.py              | 19 +++++++-
 5 files changed, 66 insertions(+), 28 deletions(-)

diff --git a/README.md b/README.md
index 737757b3..8cd2c5f2 100644
--- a/README.md
+++ b/README.md
@@ -141,7 +141,7 @@ DGX Spark / GB10:
 
 | 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 | FP8 AR decode | **3.6 ms/frame** | RTF **0.09** | [Higgs performance](docs/higgs_audio_v3.md#performance) |
 | RTX 5090 | BF16 AR decode | **6.0 ms/frame** | RTF **0.151** | [Higgs performance](docs/higgs_audio_v3.md#performance) |
 
 #### Motus Stage3
diff --git a/docs/benchmark_comparison.md b/docs/benchmark_comparison.md
index 7774787a..2c891f6d 100644
--- a/docs/benchmark_comparison.md
+++ b/docs/benchmark_comparison.md
@@ -159,18 +159,18 @@ speedup.
 
 | Metric | FlashRT FP8 | FlashRT BF16 | SGLang |
 |---|---:|---:|---:|
-| RTF | **0.095-0.11** | **0.15** | 0.16-0.19 |
-| TTFA | **~94 ms** | **~127 ms** | 0.36-0.63 s |
-| Per-frame | **~3.2 ms** | **~6.0 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** |
+| 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.79 ms** at P=6 and
-**11.74 → 6.86 ms** at P=13 in the same single-stream frontend benchmark.
+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
 
diff --git a/docs/higgs_audio_v3.md b/docs/higgs_audio_v3.md
index 3a0f35b4..93e55707 100644
--- a/docs/higgs_audio_v3.md
+++ b/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** | **≈ 127 ms** |
-| Autoregressive decode | **≈ 3.2 ms/frame** | **≈ 6.0 ms/frame** (at the BF16 bandwidth wall) |
-| Prompt prefill | **≈ 1.0 ms/token** | **≈ 0.53 ms/token** at P=13 (≈ 6.9 ms total; 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 |
@@ -41,10 +41,16 @@ 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.79 ms** at P=6
-and **11.74 → 6.86 ms** at P=13. The short-prompt EOC regression case
-`"Four score."` now exits at **1.20 s** of audio in BF16 instead of running to
-the 40.68 s generation cap.
+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;
@@ -61,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
@@ -120,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.
 
 ---
 
@@ -236,12 +242,12 @@ python examples/higgs_audio_v3_quickstart.py \
   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 (FP8 ≈ 1 ms/token; BF16 ≈ 6.9 ms total
+- **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). Prefill-only numbers time
@@ -258,8 +264,8 @@ python examples/higgs_audio_v3_quickstart.py \
 ## 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
diff --git a/flash_rt/frontends/torch/higgs_audio_v3_rtx.py b/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
index ed3ff0fd..67c8b7aa 100644
--- a/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
+++ b/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
@@ -24,6 +24,8 @@
 from flash_rt.models.higgs_audio_v3.pipeline_rtx import HiggsAudioV3Dims
 
 _SPECIALS = ("<|tts|>", "<|text|>", "<|audio|>")
+_FP8_CALIBRATION_TEXT = "Four score."
+_REPEAT_STOP_FRAMES = 32
 
 
 def _delayed_eoc_countdown(codes: torch.Tensor, nc: int, eoc: int) -> int | None:
@@ -35,6 +37,13 @@ def _delayed_eoc_countdown(codes: torch.Tensor, nc: int, eoc: int) -> int | None
     return max(1, nc - 1 - first_eoc_cb)
 
 
+def _repeat_code_run(codes: torch.Tensor,
+                     prev_key: tuple[int, ...] | None,
+                     count: int) -> tuple[tuple[int, ...], int]:
+    key = tuple(int(x) for x in codes.tolist())
+    return key, count + 1 if key == prev_key else 1
+
+
 class HiggsAudioV3TorchFrontendRtx:
     """Single-stream Higgs Audio v3 TTS inference on RTX SM120 (BF16 backbone)."""
 
@@ -303,7 +312,7 @@ def _ensure_decoder(self) -> None:
                 HiggsAudioV3Fp8Decoder,
             )
             dec = HiggsAudioV3Fp8Decoder(self)
-            dec.calibrate(self._prompt_ids)   # static activation scales (once)
+            dec.calibrate(self.build_prompt(_FP8_CALIBRATION_TEXT))
         else:
             from flash_rt.frontends.torch._higgs_audio_v3_bf16 import (
                 HiggsAudioV3Bf16Decoder,
@@ -412,6 +421,7 @@ def decode_stream(self):
         boc, eoc = self.DIMS.boc_id, self.DIMS.eoc_id
         P, logits = self._gen_pos, self._gen_logits
         delay, eoc_countdown, done = 0, None, False
+        repeat_key, repeat_count = None, 0
         window: list[torch.Tensor] = []
         for j in range(self.max_new_frames):
             fvk.higgs_audio_v3_argmax_delay_embed_bf16(
@@ -428,6 +438,11 @@ def decode_stream(self):
                     done = True
             else:
                 eoc_countdown = _delayed_eoc_countdown(codes, nc, eoc)
+                if eoc_countdown is None:
+                    repeat_key, repeat_count = _repeat_code_run(
+                        codes, repeat_key, repeat_count)
+                    if repeat_count >= _REPEAT_STOP_FRAMES:
+                        done = True
             if done:
                 break
             window.append(codes.clone())
diff --git a/tests/test_higgs_audio_v3_eoc.py b/tests/test_higgs_audio_v3_eoc.py
index b0d1f8ac..cb0fb506 100644
--- a/tests/test_higgs_audio_v3_eoc.py
+++ b/tests/test_higgs_audio_v3_eoc.py
@@ -1,6 +1,9 @@
 import torch
 
-from flash_rt.frontends.torch.higgs_audio_v3_rtx import _delayed_eoc_countdown
+from flash_rt.frontends.torch.higgs_audio_v3_rtx import (
+    _delayed_eoc_countdown,
+    _repeat_code_run,
+)
 
 
 def test_delayed_eoc_countdown_flushes_full_tail():
@@ -19,3 +22,17 @@ def test_delayed_eoc_countdown_accepts_later_codebook_eoc():
 
 def test_delayed_eoc_countdown_absent():
     assert _delayed_eoc_countdown(torch.arange(8), 8, 1025) is None
+
+
+def test_repeat_code_run_counts_identical_rows_and_resets():
+    prev, count = _repeat_code_run(torch.tensor([1, 2, 3]), None, 0)
+    assert prev == (1, 2, 3)
+    assert count == 1
+
+    prev, count = _repeat_code_run(torch.tensor([1, 2, 3]), prev, count)
+    assert prev == (1, 2, 3)
+    assert count == 2
+
+    prev, count = _repeat_code_run(torch.tensor([1, 2, 4]), prev, count)
+    assert prev == (1, 2, 4)
+    assert count == 1

From 2653658a47b4839b8bd123f9563b5c4ca55e7e9c Mon Sep 17 00:00:00 2001
From: LiangSu8899 <7thuniversels@gmail.com>
Date: Thu, 9 Jul 2026 05:32:54 -0400
Subject: [PATCH 4/5] Add Higgs TTFA to benchmark summary

---
 README.md | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 8cd2c5f2..a4f8554b 100644
--- a/README.md
+++ b/README.md
@@ -139,10 +139,10 @@ DGX Spark / GB10:
 
 #### Higgs Audio v3
 
-| Hardware | Mode | Latency | Throughput | Source |
-|---|---|---:|---:|---|
-| RTX 5090 | FP8 AR decode | **3.6 ms/frame** | RTF **0.09** | [Higgs performance](docs/higgs_audio_v3.md#performance) |
-| RTX 5090 | BF16 AR decode | **6.0 ms/frame** | RTF **0.151** | [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
 

From 1cd4f9a9a828ce1097f9fa4feba252417831bb9d Mon Sep 17 00:00:00 2001
From: LiangSu8899 <7thuniversels@gmail.com>
Date: Thu, 9 Jul 2026 05:45:56 -0400
Subject: [PATCH 5/5] Update audio codebook kernel inventory

---
 CMakeLists.txt                                | 15 +++++++++++-
 csrc/bindings.cpp                             | 10 +++++---
 ...kernels.cu => delayed_codebook_kernels.cu} | 12 +++++-----
 ...rnels.cuh => delayed_codebook_kernels.cuh} |  2 +-
 .../frontends/torch/higgs_audio_v3_rtx.py     |  6 ++++-
 scripts/build_inventory.py                    |  3 +++
 tests/test_build_inventory.py                 | 23 ++++++++++++++-----
 tests/test_generic_kernel_bindings.py         | 23 +++++++++++++++++--
 8 files changed, 74 insertions(+), 20 deletions(-)
 rename csrc/kernels/{higgs_audio_v3_kernels.cu => delayed_codebook_kernels.cu} (88%)
 rename csrc/kernels/{higgs_audio_v3_kernels.cuh => delayed_codebook_kernels.cuh} (89%)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 391b046e..34a81f96 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -168,6 +168,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)
 # ── Slim build for VLA deployments ──
 # OFF by default: the standard build keeps the broad, compatibility-oriented
 # kernel set it has today. Turn ON (-DFLASHRT_SLIM_BUILD=ON) to drop
@@ -1113,7 +1118,6 @@ pybind11_add_module(flash_rt_kernels
   csrc/gemm/fp8_block128_gemm.cu
   csrc/kernels/silu_mul_qwen36.cu
   csrc/kernels/embedding_lookup_bf16.cu
-  csrc/kernels/higgs_audio_v3_kernels.cu
   # qwen3_5_moe-family RTX SM120 kernels are NOT listed here -- they are
   # gated into the module below (FLASHRT_ENABLE_QWEN35MOE), so SM89/87/110
   # and non-Nex-N2 builds never compile them. See the gated block after the
@@ -1128,6 +1132,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/*).
diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp
index 4022986f..32ee0b3b 100644
--- a/csrc/bindings.cpp
+++ b/csrc/bindings.cpp
@@ -146,7 +146,9 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128(
 #endif
 #include "kernels/silu_mul_qwen36.cuh"
 #include "kernels/embedding_lookup_bf16.cuh"
-#include "kernels/higgs_audio_v3_kernels.cuh"
+#ifdef FLASHRT_HAVE_AUDIO_CODEBOOK
+#include "kernels/delayed_codebook_kernels.cuh"
+#endif
 #ifdef FLASHRT_HAVE_QWEN36_KERNELS
 #include "kernels/qwen36_misc.cuh"
 #endif
@@ -4353,11 +4355,12 @@ 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);
 
-    m.def("higgs_audio_v3_argmax_delay_embed_bf16",
+#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::higgs_audio_v3_argmax_delay_embed_bf16(
+            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),
@@ -4369,6 +4372,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) {
         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",
diff --git a/csrc/kernels/higgs_audio_v3_kernels.cu b/csrc/kernels/delayed_codebook_kernels.cu
similarity index 88%
rename from csrc/kernels/higgs_audio_v3_kernels.cu
rename to csrc/kernels/delayed_codebook_kernels.cu
index e839b157..3646a93b 100644
--- a/csrc/kernels/higgs_audio_v3_kernels.cu
+++ b/csrc/kernels/delayed_codebook_kernels.cu
@@ -1,4 +1,4 @@
-#include "higgs_audio_v3_kernels.cuh"
+#include "delayed_codebook_kernels.cuh"
 
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
@@ -7,7 +7,7 @@
 namespace flash_rt::kernels {
 namespace {
 
-__global__ void higgs_argmax_delay_kernel(
+__global__ void delayed_codebook_argmax_kernel(
     const __nv_bfloat16* logits,
     int64_t* codes,
     int num_codebooks,
@@ -54,7 +54,7 @@ __global__ void higgs_argmax_delay_kernel(
   if (tid == 0) codes[cb] = static_cast<int64_t>(idxs[0]);
 }
 
-__global__ void higgs_embed_sum_kernel(
+__global__ void delayed_codebook_embed_sum_kernel(
     const int64_t* codes,
     const __nv_bfloat16* codebook,
     __nv_bfloat16* embed,
@@ -74,7 +74,7 @@ __global__ void higgs_embed_sum_kernel(
 
 }  // namespace
 
-void higgs_audio_v3_argmax_delay_embed_bf16(
+void delayed_codebook_argmax_embed_bf16(
     const __nv_bfloat16* logits,
     const __nv_bfloat16* codebook,
     int64_t* codes_out,
@@ -88,11 +88,11 @@ void higgs_audio_v3_argmax_delay_embed_bf16(
   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));
-  higgs_argmax_delay_kernel<<<num_codebooks, arg_threads, smem, stream>>>(
+  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;
-  higgs_embed_sum_kernel<<<emb_blocks, emb_threads, 0, stream>>>(
+  delayed_codebook_embed_sum_kernel<<<emb_blocks, emb_threads, 0, stream>>>(
       codes_out, codebook, embed_out, num_codebooks, codebook_vocab, hidden);
 }
 
diff --git a/csrc/kernels/higgs_audio_v3_kernels.cuh b/csrc/kernels/delayed_codebook_kernels.cuh
similarity index 89%
rename from csrc/kernels/higgs_audio_v3_kernels.cuh
rename to csrc/kernels/delayed_codebook_kernels.cuh
index 269e0e6d..4a66e00f 100644
--- a/csrc/kernels/higgs_audio_v3_kernels.cuh
+++ b/csrc/kernels/delayed_codebook_kernels.cuh
@@ -6,7 +6,7 @@
 
 namespace flash_rt::kernels {
 
-void higgs_audio_v3_argmax_delay_embed_bf16(
+void delayed_codebook_argmax_embed_bf16(
     const __nv_bfloat16* logits,
     const __nv_bfloat16* codebook,
     int64_t* codes_out,
diff --git a/flash_rt/frontends/torch/higgs_audio_v3_rtx.py b/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
index 67c8b7aa..648e59ec 100644
--- a/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
+++ b/flash_rt/frontends/torch/higgs_audio_v3_rtx.py
@@ -126,6 +126,10 @@ def _resolve_precision(fp8: bool | None) -> bool:
             get_gpu_sm_version,
         )
         log = logging.getLogger(__name__)
+        if not hasattr(fvk, "delayed_codebook_argmax_embed_bf16"):
+            raise RuntimeError(
+                "Higgs Audio v3: delayed codebook decode helper is not in this "
+                "build; rebuild with FLASHRT_ENABLE_AUDIO_CODEBOOK=ON.")
         fp8_ok = (hasattr(fvk, "ht_gemv_fp8_m1_w8")
                   and hasattr(fvk, "ht_fp8_gemm_16x192x128_w8"))
         bf16_ok = hasattr(fvk, "ht_gemv_bf16_m1_w4")
@@ -424,7 +428,7 @@ def decode_stream(self):
         repeat_key, repeat_count = None, 0
         window: list[torch.Tensor] = []
         for j in range(self.max_new_frames):
-            fvk.higgs_audio_v3_argmax_delay_embed_bf16(
+            fvk.delayed_codebook_argmax_embed_bf16(
                 logits.data_ptr(), self._weights["codebook"].data_ptr(),
                 self._codes_dev.data_ptr(), self._embed_buf.data_ptr(),
                 nc, self._cfg["codebook_vocab"], self._cfg["hidden"],
diff --git a/scripts/build_inventory.py b/scripts/build_inventory.py
index 7542e8b8..e4120e80 100644
--- a/scripts/build_inventory.py
+++ b/scripts/build_inventory.py
@@ -105,6 +105,9 @@
         "csrc/kernels/rope_qwen3.cu",
         "csrc/kernels/qwen3_qkv_post_proc.cu",
     ),
+    "audio_codebook": (
+        "csrc/kernels/delayed_codebook_kernels.cu",
+    ),
 }
 
 _SRC_RE = re.compile(r"csrc/[A-Za-z0-9_/]+\.(?:cu|cpp)")
diff --git a/tests/test_build_inventory.py b/tests/test_build_inventory.py
index 36876d25..afb36bf0 100644
--- a/tests/test_build_inventory.py
+++ b/tests/test_build_inventory.py
@@ -30,15 +30,15 @@
 # build-mode aware: it reads FLASHRT_SLIM_BUILD from the configured build dir's
 # CMakeCache.txt and asserts the matching surface.
 #
-#   compat default (FLASHRT_SLIM_BUILD=OFF): 55 direct TUs, unchanged.
-#   slim SM89   (FLASHRT_SLIM_BUILD=ON):     33 direct TUs (-22):
+#   compat default (FLASHRT_SLIM_BUILD=OFF): 56 direct TUs, unchanged.
+#   slim SM89   (FLASHRT_SLIM_BUILD=ON):     34 direct TUs (-22):
 #       -5 Motus VAE FP8 (Unit 1), -10 Qwen3.6/linear-attn (Unit 2),
 #       -7 SM120/NVFP4-named (Unit 3).
 # The direct-source breakdowns below are SM89-specific; other arches add/remove
 # arch-owned sources (for example SM120 adds nvfp4_sf_reshape_sm120.cu), so the
 # direct-count/category asserts are limited to SM89.
-COMPAT_KERNELS_TU = 55
-SLIM_SM89_KERNELS_TU = 33
+COMPAT_KERNELS_TU = 56
+SLIM_SM89_KERNELS_TU = 34
 
 # Per-group breakdown of flash_rt_kernels (mirrors AGENTS.md "Current Build
 # Layout"). categorize() drops empty groups, so slim omits sm120_nvfp4_named.
@@ -49,6 +49,7 @@
     "motus_video_fp8_history": 7,
     "dit_video": 2,
     "qwen3_family": 2,
+    "audio_codebook": 1,
     "other": 10,
 }
 SLIM_SM89_KERNELS_CATEGORIES = {
@@ -57,6 +58,7 @@
     "motus_video_fp8_history": 2,
     "dit_video": 2,
     "qwen3_family": 2,
+    "audio_codebook": 1,
     "other": 10,
 }
 
@@ -90,12 +92,21 @@ def _is_sm89() -> bool:
     return (_cache_value("GPU_ARCH") or "").strip() == "89"
 
 
+def _audio_codebook_enabled() -> bool:
+    return _cache_bool("FLASHRT_ENABLE_AUDIO_CODEBOOK")
+
+
 def _expected_tu() -> int:
-    return SLIM_SM89_KERNELS_TU if _is_slim() else COMPAT_KERNELS_TU
+    expected = SLIM_SM89_KERNELS_TU if _is_slim() else COMPAT_KERNELS_TU
+    return expected if _audio_codebook_enabled() else expected - 1
 
 
 def _expected_categories() -> dict:
-    return SLIM_SM89_KERNELS_CATEGORIES if _is_slim() else COMPAT_KERNELS_CATEGORIES
+    expected = dict(SLIM_SM89_KERNELS_CATEGORIES if _is_slim()
+                    else COMPAT_KERNELS_CATEGORIES)
+    if not _audio_codebook_enabled():
+        expected.pop("audio_codebook", None)
+    return expected
 
 
 def _load_inventory():
diff --git a/tests/test_generic_kernel_bindings.py b/tests/test_generic_kernel_bindings.py
index 64d09095..9cbd7919 100644
--- a/tests/test_generic_kernel_bindings.py
+++ b/tests/test_generic_kernel_bindings.py
@@ -85,6 +85,10 @@ def _is_slim_build() -> bool:
     return _cache_bool("FLASHRT_SLIM_BUILD")
 
 
+def _audio_codebook_enabled() -> bool:
+    return _cache_bool("FLASHRT_ENABLE_AUDIO_CODEBOOK")
+
+
 def test_legacy_matmul_bindings_exist():
     """Compat build: legacy name present. Slim build: gated out by design."""
     m = _import_kernels()
@@ -129,8 +133,8 @@ def test_legacy_cublaslt_binding_exists_on_vl_module():
     m = _import_vl_kernels()
     if not hasattr(m, "fp8_block128_gemm_blockscaled_sm89_bf16out"):
         pytest.skip(
-            "bf16_matmul_cublaslt_bf16 is part of the SM89 Qwen3-VL module; "
-            "SM120 builds use the SM120 path instead"
+            "Qwen3-VL bf16_matmul_cublaslt_bf16 is present only on the SM89 "
+            "VL module path; SM120 builds use the SM120 path instead"
         )
     assert hasattr(m, "bf16_matmul_cublaslt_bf16")
 
@@ -141,7 +145,22 @@ def test_neutral_matmul_binding_exists():
     assert hasattr(m, "bf16_matmul_bf16")
 
 
+def test_neutral_cublaslt_matmul_binding_exists():
+    """cuBLASLt BF16 GEMM helper is exported by the core kernel module."""
+    m = _import_kernels()
+    assert hasattr(m, "bf16_matmul_cublaslt_bf16")
+
+
 def test_neutral_embedding_binding_exists():
     """Neutral helper must exist in every build mode (never gated)."""
     m = _import_kernels()
     assert hasattr(m, "embedding_lookup_bf16")
+
+
+def test_delayed_codebook_binding_exists():
+    """Delayed codebook argmax/embed helper is exported by the core module."""
+    m = _import_kernels()
+    if _audio_codebook_enabled():
+        assert hasattr(m, "delayed_codebook_argmax_embed_bf16")
+    else:
+        assert not hasattr(m, "delayed_codebook_argmax_embed_bf16")