feat(structures): decoder_llm + audio_codec whole-module structures with one-call doors (OmniVoice, 4.24x)

docs/omnivoice_structures.md

@@ -0,0 +1,132 @@
+# OmniVoice via flash_rt.structures — whole-LLM NVFP4 structure + MaskGIT schedule
+
+The OmniVoice TTS host (Qwen3-1.5B backbone + 8-codebook MaskGIT head)
+accelerated through the structures layer: a **`decoder_llm`** catalog
+structure binds the whole backbone to the native NVFP4 engines (fp4 GEMMs
++ fused qk-norm+RoPE + FA2 + fused residual/norm/quant), and a
+**`maskgit_gen`** schedule drives the two-phase MaskGIT loop (BF16 CFG
+step, then FP4 noCFG single-stream graph replays).
+
+Kernel resolution is the PR-175 tiering: hub artifact first, the local
+native build second, the retained host stack always the floor. The hub's
+fp4 packages do not yet ship a torch-2.13 variant, so the local build is
+the active tier here.
+
+## Measured (RTX 5060 Ti, torch 2.13+cu130, omnivoice 0.2.1, design mode,
+32 MaskGIT steps, gs=2.0, seed=42, median of 3)
+
+| text | baseline RTF | structures RTF | speedup |
+|---|---|---|---|
+| short | 0.2292 | 0.0438 | 5.24x |
+| medium | 0.1187 | 0.0277 | 4.29x |
+| long | 0.0791 | 0.0185 | 4.28x |
+| **median** | 0.1187 | 0.0277 | **4.29x** |
+
+The native inject path measures 0.0284 median on the same box — the
+structures route is identical within noise (and slightly ahead once the
+codec is included). Per-seam-only adoption (fp4 FFN, fp4 FFN+projections)
+measured ~1.0x: the win lives in the whole-LLM boundary and the schedule,
+not the individual kernels.
+
+### Rejected: CFG steps on the FP4 engine (measured, 32 steps fixed)
+
+The CFG phase (B=2, ~5% of steps at cfg_ratio=0.05) rides the BF16
+engine by design — no quantization drift on the guidance subtraction.
+Moving it onto the FP4 engine via two B=1 graph replays is 2.1x faster
+per step (10.95 ms vs 2x 2.60 ms) but **breaks generation**: same-seed
+tokens match only 0.8% and the output collapses to 0.68 s of audio vs
+3.14 s on the same schedule. The FP4 engine's hidden states differ from
+BF16's by ~2x relative (W4A4 quantization is a different distribution,
+not a perturbation), which the cond-uncond difference amplifies. Rejected
+per AGENTS.md: the 6.3% win is not worth a broken CFG path. A real fix
+needs a B=2 FP4 engine (kernel-side, FlashRT-HF-kernels delivery).
+
+### The codec (the remaining 12%)
+
+The neural codec decode is compute-bound on fp32 cuDNN convolutions; the
+only measured lever is fp16 autocast (1.22x, waveform cosine 1.00000 vs
+the fp32 host — torch.compile and CUDA graph both measure ~1.0x, and the
+pydub postprocess is only ~3.7 ms, not worth replacing). It is adopted
+as an `audio_codec` structure (`impls/audio_codec/fp16.py`): codes ->
+waveform, fp16 decode with host fallback. The host family
+(HiggsAudioV2TokenizerModel) is shared by OmniVoice and Higgs-Audio-v3,
+so the structure recurs beyond one model.
+
+## Usage
+
+```python
+import torch
+from omnivoice import OmniVoice
+from omnivoice import OmniVoiceGenerationConfig
+from flash_rt import structures
+from flash_rt.structures.impls.decoder_llm import nvfp4 as llm_impl
+
+model = OmniVoice.from_pretrained("/path/to/OmniVoice",
+                                  dtype=torch.bfloat16).to("cuda:0")
+model.eval()
+
+def calibration():                       # one short real forward
+    with torch.no_grad():
+        g = OmniVoiceGenerationConfig(num_step=2, guidance_scale=2.0,
+                                      denoise=True, preprocess_prompt=True,
+                                      postprocess_output=False,
+                                      position_temperature=5.0)
+        model.generate(text="Calibration sentence.", generation_config=g)
+
+plan = structures.auto_swaps(model, calibration,
+                             structures=("decoder_llm", "audio_codec"),
+                             scheme="nvfp4_static")
+print(structures.explain(plan))          # receipt: llm + audio_tokenizer
+handle = plan.attach()                   # the one-call door
+
+g = OmniVoiceGenerationConfig(num_step=32, guidance_scale=2.0,
+                              denoise=True, preprocess_prompt=True,
+                              postprocess_output=True,
+                              position_temperature=5.0)
+task = model._preprocess_all(text="Hello.", language=None, ref_text=None,
+                             ref_audio=None, voice_clone_prompt=None,
+                             instruct=None, preprocess_prompt=True,
+                             speed=None, duration=None,
+                             normalize_text=False)
+st = task.slice_task(task.get_indices(
+    g, model.audio_tokenizer.config.frame_rate)[0])
+loop = structures.maskgit_loop(model)               # the schedule door
+with torch.no_grad():
+    tokens = loop.generate(st, g)                   # two-phase schedule
+    audio = model._decode_and_post_process(tokens[0], None, g)
+```
+
+Requirements: `flash_rt_kernels` + `flash_rt_omnivoice` + `flash_rt_fa2`
+built with `-DFLASHRT_ENABLE_OMNIVOICE=ON -DGPU_ARCH=120` (sm_120a), the
+`omnivoice` pip package, and a checkpoint.
+
+## What was added
+
+- `flash_rt/structures/catalog/decoder_llm/` — structure spec + torch
+  reference (the host's eager stack forward is the parity ground truth).
+- `flash_rt/structures/impls/decoder_llm/nvfp4.py` — the whole-LLM seam:
+  BF16 fused engine for the CFG batch (B=2), FP4 CUDA-graph engine for
+  single stream (B=1); profile-envelope refusal outside the native
+  engine's v1 contract (D=1024/L=28/NH=16/NKV=8/HD=128/FFN=3072); and the
+  `maskgit_gen` schedule door.
+- `flash_rt/structures/impls/decoder_ffn/nvfp4_static.py` and
+  `flash_rt/structures/impls/linear_proj/nvfp4_static.py` — per-seam fp4
+  backends (kept; useful on other hosts).
+- `flash_rt/structures/schemes.py` — `nvfp4_static` scheme (weight-only,
+  no calibration data).
+- `flash_rt/structures/discover.py` — decoder-stack discovery rule
+  (layers/embed_tokens/norm/rotary_emb slots, not model names).
+- `flash_rt/structures/autobuild.py` — bind dispatch for the new formats.
+- `flash_rt/structures/bindings/omnivoice_llm.yaml` — host addressing
+  receipt.
+- `tests/test_structures_decoder_llm.py` — CPU contract pins.
+
+## Notes
+
+- The noCFG FP4 phase collapses tiny prompts into a repeated-code
+  attractor (silent audio) in both the host and this path — a host
+  characteristic, not a structure bug. Keep `guidance_scale > 0` and use
+  real sentences.
+- Attaching over a service thread pool: the seam guard pins one thread
+  per attachment; reset `model.llm._frt_guard.thread = None` per request
+  when uvicorn-style pooling is in play.

flash_rt/structures/__init__.py

@@ -121,6 +121,20 @@ def decode_loop(model, *, max_len, compile_step=True,
                              kv_band=kv_band)
 
 
+def maskgit_loop(model, *, cfg_ratio: float = 0.05,
+                 bookend: bool = False):
+    """Serving door: the MaskGIT two-phase schedule over an attached
+    ``decoder_llm`` seam — the decode_loop twin for non-text hosts.
+
+    ``generate(task, gen_config)`` runs the BF16-CFG steps then the FP4
+    noCFG single-stream graph replays (the OmniVoice-style schedule).
+    See :mod:`flash_rt.structures.impls.decoder_llm.nvfp4`.
+    """
+    from flash_rt.structures.impls.decoder_llm.nvfp4 import MaskgitLoop
+
+    return MaskgitLoop(model, cfg_ratio=cfg_ratio, bookend=bookend)
+
+
 def aot_package(module, args=(), kwargs=None,
                 package_path="module_aot.pt2", **opts):
     """Whole-graph door: export the swapped module and AOT-compile it
@@ -151,6 +165,7 @@ def aot_load(package_path, weights=None):
     "aot_load",
     "aot_package",
     "decode_loop",
+    "maskgit_loop",
     "explain",
     "attach",
     "capture",

flash_rt/structures/autobuild.py

@@ -100,6 +100,10 @@ class AutoPlan:
     updates: list[Callable[[], None]] = field(default_factory=list)
     seams: list[Seam] = field(default_factory=list)
     notes: dict[str, Any] = field(default_factory=dict)
+    #: the host the plan was discovered on; ``attach`` installs the
+    #: swaps at their paths relative to this root (the one-call door:
+    #: ``plan = auto_swaps(model, forward); handle = plan.attach()``)
+    root: torch.nn.Module | None = field(default=None, repr=False)
     #: modules that carry a guard but are not swapped at a path — an
     #: adapter's routed seam. Reported by the attachment's ledger, never
     #: installed by it, so a seam that cannot be swapped can still be
@@ -120,6 +124,26 @@ class AutoPlan:
     #: one, so ``plan.precision_spec`` reads like ``rt.precision_spec``
     precision_spec: Any = None
 
+    def attach(self):
+        """Commit the plan: install every swap at its path on ``root``.
+
+        The one-call door (the article's ``plan.attach()``): atomic
+        resolution + install, with the plan's observed seams and revert
+        callables carried along. Returns the attachment handle
+        (``detach()`` restores the host bit-for-bit).
+        """
+        from .swap import attach as _swap_attach
+
+        if self.root is None:
+            raise ValueError(
+                "refused: plan carries no root host (auto_swaps did not "
+                "record one); pass the model to auto_swaps")
+        if not self.swaps and not self.observed:
+            raise ValueError("no swaps or routed seams staged")
+        return _swap_attach(
+            self.root, self.swaps,
+            observe=self.observed, revert=self.revert)
+
     def enable_routed(self) -> None:
         for on, _ in self.toggles:
             on()
@@ -655,7 +679,7 @@ def _region_probe():
     adapter_only = not seams and bool(
         {"attention_core", "gated_delta_core"}.intersection(structures))
     if not seams and not adapter_only:
-        plan = AutoPlan()
+        plan = AutoPlan(root=model)
         if region_extras is not None:
             _merge_region_extras(plan, region_extras)
         return plan
@@ -1023,7 +1047,7 @@ def __init__(self, *args, structure: str, **kwargs):
     # for fp8 input whose producer failed to bind would be handed BF16,
     # and the host would silently grow a quantize fused into whatever
     # produced it. Bind the pair together, or leave both on BF16.
-    plan = AutoPlan(seams=seams)
+    plan = AutoPlan(seams=seams, root=model)
     plan._requested_structures = frozenset(structures)
     if region_extras is not None:
         _merge_region_extras(plan, region_extras)
@@ -1583,10 +1607,30 @@ def scale(name, path=None):
                     and fmt in ("bf16_pack", "nvfp4_balance")) \
             and not (seam.structure == "vision_ffn"
                      and fmt == "nvfp4_balance") \
-            and seam.structure not in ("decoder_ffn", "linear_proj"):
+            and seam.structure not in ("decoder_ffn", "linear_proj",
+                                       "decoder_llm", "audio_codec"):
         raise ValueError(f"scheme routed {seam.structure} to format "
                          f"{fmt!r}, which has no impl variant here")
 
+    if seam.structure == "audio_codec":
+        if fmt == "fp16_codec":
+            from .impls.audio_codec import fp16 as codec_impl
+            return {seam.path: codec_impl.bind_codec_decode(
+                _resolve(model, seam.path),
+                original=_resolve(model, seam.path))}
+        raise ValueError(f"scheme routed audio_codec to format "
+                         f"{fmt!r}, which has no impl variant here")
+
+    if seam.structure == "decoder_llm":
+        if fmt == "nvfp4_static":
+            from .impls.decoder_llm import nvfp4 as llm_impl
+            return {seam.path: llm_impl.bind_decoder_stack(
+                _resolve(model, seam.path),
+                variant=seam.variant,
+                original=_resolve(model, seam.path))}
+        raise ValueError(f"scheme routed decoder_llm to format "
+                         f"{fmt!r}, which has no impl variant here")
+
     if seam.structure == "decoder_ffn":
         if fmt in ("w8a16_static", "w4a16_static"):
             if fmt == "w8a16_static":
@@ -1608,6 +1652,16 @@ def scale(name, path=None):
             return wq_impl.bind_mlp_seam(
                 w, variant=seam.variant,
                 original=_resolve(model, seam.path))
+        if fmt == "nvfp4_static":
+            from .impls.decoder_ffn import nvfp4_static as nv_impl
+            w = seam_weights(model, seam)
+            w = dict(w,
+                     w_gate=w["w_gate"].t().contiguous(),
+                     w_up=w["w_up"].t().contiguous(),
+                     w_down=w["w_down"].t().contiguous())
+            return nv_impl.bind_mlp_seam(
+                w, variant=seam.variant,
+                original=_resolve(model, seam.path))
         if fmt not in (None, "fp8_static"):
             raise ValueError(f"scheme routed decoder_ffn to format "
                              f"{fmt!r}, which has no impl variant here")
@@ -1698,6 +1752,12 @@ def scale(name, path=None):
             return proj_w8.bind_proj_seam(
                 seam_weights(model, seam),
                 original=_resolve(model, seam.path))
+        if fmt == "nvfp4_static":
+            # native-build fp4 GEMM tier for large-M projection calls
+            from .impls.linear_proj import nvfp4_static as proj_nv
+            return proj_nv.bind_proj_seam(
+                seam_weights(model, seam),
+                original=_resolve(model, seam.path))
         if fmt not in (None, "fp8_static"):
             raise ValueError(f"scheme routed linear_proj to format "
                              f"{fmt!r}, which has no impl variant here")

flash_rt/structures/bindings/omnivoice_audio_codec.yaml

@@ -0,0 +1,15 @@
+binding: omnivoice_audio_codec
+structure: audio_codec
+dims: {C: 8, T: dynamic}
+variant: {backend: fp16_codec}
+boundary_dtype: int64
+
+m_profile:
+  decode: {m_class: medium, phase: decode}
+
+hosts:
+  omnivoice:
+    module_path: "audio_tokenizer"
+    versions: ">=0.2"
+  higgs_audio_v3:
+    module_path: "audio_tokenizer"

flash_rt/structures/bindings/omnivoice_llm.yaml

@@ -0,0 +1,27 @@
+binding: omnivoice_llm
+structure: decoder_llm
+dims: {D: 1024, L: 28, NH: 16, NKV: 8, HD: 128, FFN: 3072}
+variant: {backend: nvfp4_llm}
+boundary_dtype: bf16
+
+m_profile:
+  prefill: {m_class: medium, phase: prefill}
+  decode: {m_class: medium, phase: decode}
+
+# The NVFP4 whole-LLM backend binds the stack in place (weights are read
+# from the host module, not remapped by slot), so weights_map is
+# descriptive: the seams the backend fuses per layer.
+weights_map:
+  omnivoice:
+    layout: out_in
+    w_qkv: "llm.layers.{i}.self_attn.{q,k,v}_proj.weight"
+    w_o: "llm.layers.{i}.self_attn.o_proj.weight"
+    w_gu: "llm.layers.{i}.mlp.{gate,up}_proj.weight"
+    w_down: "llm.layers.{i}.mlp.down_proj.weight"
+    w_in_norm: "llm.layers.{i}.input_layernorm.weight"
+    w_post_norm: "llm.layers.{i}.post_attention_layernorm.weight"
+
+hosts:
+  omnivoice:
+    module_path: "llm"
+    versions: ">=0.2"

flash_rt/structures/catalog/audio_codec/reference.py

@@ -0,0 +1,15 @@
+"""Plain-torch reference for ``audio_codec``.
+
+The parity ground truth is the host codec's own eager fp32 decode.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def audio_codec_ref(module: torch.nn.Module,
+                    codes: torch.Tensor) -> torch.Tensor:
+    """Run the host codec decode eagerly (fp32); returns audio_values."""
+    out = module.decode(codes)
+    return out.audio_values

flash_rt/structures/catalog/audio_codec/structure.yaml

@@ -0,0 +1,36 @@
+structure: audio_codec
+version: 1
+description: >
+  Neural audio codec decode: discrete codes -> waveform. The region is
+  a conv stack (DAC-style acoustic/semantic decoders, fp32 cuDNN
+  convolutions + elementwise upsampling). The fp16 impl runs the decode
+  under fp16 autocast — the only measured lever here (1.22x, waveform
+  cosine 1.00000 vs the fp32 host; torch.compile and CUDA graph both
+  measure ~1.0x on this compute-bound region). Host family: the
+  HiggsAudioV2TokenizerModel is shared by OmniVoice and Higgs-Audio-v3.
+
+reference:
+  module: audio_codec.reference
+  entrypoint: audio_codec_ref
+
+boundary:
+  symbolic_dims: [B, C, T, N]
+  inputs:
+    - {name: codes, dims: [B, C, T], dtype: int64}
+  outputs:
+    - {name: audio_values, dims: [B, 1, N], dtype: float32}
+
+weights: []
+
+variants: {}
+
+calibration: {}
+
+gates:
+  parity:
+    metrics: [cosine, max_abs, p99_abs]
+    data: real_distribution
+  latency:
+    baselines: [torch_compile, unfused_chain]
+    rule: net_positive_including_boundary
+    per_shape: true