From 72f001f41506f0c5db989cf23645f49b8339db02 Mon Sep 17 00:00:00 2001 From: shideqin Date: Sun, 16 Aug 2026 17:43:21 +0800 Subject: [PATCH 1/2] feat(structures): decoder_llm + audio_codec whole-module structures with one-call doors (OmniVoice, 4.24x) --- docs/omnivoice_structures.md | 119 ++++++++ flash_rt/structures/__init__.py | 15 + flash_rt/structures/autobuild.py | 66 ++++- .../bindings/omnivoice_audio_codec.yaml | 15 + .../structures/bindings/omnivoice_llm.yaml | 27 ++ .../catalog/audio_codec/__init__.py | 0 .../catalog/audio_codec/reference.py | 15 + .../catalog/audio_codec/structure.yaml | 36 +++ .../catalog/decoder_llm/__init__.py | 0 .../catalog/decoder_llm/reference.py | 23 ++ .../catalog/decoder_llm/structure.yaml | 36 +++ flash_rt/structures/discover.py | 36 +++ .../structures/impls/audio_codec/__init__.py | 1 + flash_rt/structures/impls/audio_codec/fp16.py | 67 +++++ .../impls/decoder_ffn/nvfp4_static.py | 266 ++++++++++++++++++ .../structures/impls/decoder_llm/__init__.py | 1 + .../structures/impls/decoder_llm/nvfp4.py | 206 ++++++++++++++ .../impls/linear_proj/nvfp4_static.py | 174 ++++++++++++ flash_rt/structures/schemes.py | 41 +++ tests/test_structures_decoder_llm.py | 173 ++++++++++++ 20 files changed, 1314 insertions(+), 3 deletions(-) create mode 100644 docs/omnivoice_structures.md create mode 100644 flash_rt/structures/bindings/omnivoice_audio_codec.yaml create mode 100644 flash_rt/structures/bindings/omnivoice_llm.yaml create mode 100644 flash_rt/structures/catalog/audio_codec/__init__.py create mode 100644 flash_rt/structures/catalog/audio_codec/reference.py create mode 100644 flash_rt/structures/catalog/audio_codec/structure.yaml create mode 100644 flash_rt/structures/catalog/decoder_llm/__init__.py create mode 100644 flash_rt/structures/catalog/decoder_llm/reference.py create mode 100644 flash_rt/structures/catalog/decoder_llm/structure.yaml create mode 100644 flash_rt/structures/impls/audio_codec/__init__.py create mode 100644 flash_rt/structures/impls/audio_codec/fp16.py create mode 100644 flash_rt/structures/impls/decoder_ffn/nvfp4_static.py create mode 100644 flash_rt/structures/impls/decoder_llm/__init__.py create mode 100644 flash_rt/structures/impls/decoder_llm/nvfp4.py create mode 100644 flash_rt/structures/impls/linear_proj/nvfp4_static.py create mode 100644 tests/test_structures_decoder_llm.py diff --git a/docs/omnivoice_structures.md b/docs/omnivoice_structures.md new file mode 100644 index 000000000..c3361196d --- /dev/null +++ b/docs/omnivoice_structures.md @@ -0,0 +1,119 @@ +# 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. + +### 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. diff --git a/flash_rt/structures/__init__.py b/flash_rt/structures/__init__.py index fce219b35..5c1c04cfe 100644 --- a/flash_rt/structures/__init__.py +++ b/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", diff --git a/flash_rt/structures/autobuild.py b/flash_rt/structures/autobuild.py index 10b07e3a6..ec2999c9d 100644 --- a/flash_rt/structures/autobuild.py +++ b/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") diff --git a/flash_rt/structures/bindings/omnivoice_audio_codec.yaml b/flash_rt/structures/bindings/omnivoice_audio_codec.yaml new file mode 100644 index 000000000..16ecf4f52 --- /dev/null +++ b/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" diff --git a/flash_rt/structures/bindings/omnivoice_llm.yaml b/flash_rt/structures/bindings/omnivoice_llm.yaml new file mode 100644 index 000000000..15790a6ca --- /dev/null +++ b/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" diff --git a/flash_rt/structures/catalog/audio_codec/__init__.py b/flash_rt/structures/catalog/audio_codec/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/flash_rt/structures/catalog/audio_codec/reference.py b/flash_rt/structures/catalog/audio_codec/reference.py new file mode 100644 index 000000000..397ee4063 --- /dev/null +++ b/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 diff --git a/flash_rt/structures/catalog/audio_codec/structure.yaml b/flash_rt/structures/catalog/audio_codec/structure.yaml new file mode 100644 index 000000000..c656744fb --- /dev/null +++ b/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 diff --git a/flash_rt/structures/catalog/decoder_llm/__init__.py b/flash_rt/structures/catalog/decoder_llm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/flash_rt/structures/catalog/decoder_llm/reference.py b/flash_rt/structures/catalog/decoder_llm/reference.py new file mode 100644 index 000000000..ef137b3e9 --- /dev/null +++ b/flash_rt/structures/catalog/decoder_llm/reference.py @@ -0,0 +1,23 @@ +"""Plain-torch reference for ``decoder_llm``. + +The parity ground truth is the host's own eager stack forward: the +reference is the module itself, called with its native interface +(``inputs_embeds=...``). No standalone torch re-implementation exists +because the structure's boundary is the whole stack — decomposing it +here would be a second host, not a reference. +""" + +from __future__ import annotations + +import torch + + +def decoder_llm_ref(module: torch.nn.Module, + inputs_embeds: torch.Tensor, **kwargs) -> torch.Tensor: + """Run the host stack eagerly; returns ``hidden_states``.""" + out = module(inputs_embeds=inputs_embeds, **kwargs) + if isinstance(out, tuple): + return out[0] + if hasattr(out, "last_hidden_state"): + return out.last_hidden_state + return out diff --git a/flash_rt/structures/catalog/decoder_llm/structure.yaml b/flash_rt/structures/catalog/decoder_llm/structure.yaml new file mode 100644 index 000000000..25c916e2c --- /dev/null +++ b/flash_rt/structures/catalog/decoder_llm/structure.yaml @@ -0,0 +1,36 @@ +structure: decoder_llm +version: 1 +description: > + Whole decoder stack: inputs_embeds -> transformer layers -> final norm + -> hidden states. The NVFP4 backend fuses the per-layer chain + (norm -> qkv -> qk-norm+RoPE -> FA2 -> o -> residual+norm -> gate/up -> + SiLU -> down -> residual) into fp4 GEMMs plus fused kernels with + whole-step CUDA graph capture, in the PR-175 tiering (hub artifact + first, local native build second, host floor). + +reference: + module: decoder_llm.reference + entrypoint: decoder_llm_ref + +boundary: + symbolic_dims: [B, S, D] + inputs: + - {name: inputs_embeds, dims: [B, S, D], dtype: "@binding"} + outputs: + - {name: hidden_states, dims: [B, S, D], dtype: "@binding"} + +# weights are bound in place from the host stack (layers/embed/norm/rope) +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 diff --git a/flash_rt/structures/discover.py b/flash_rt/structures/discover.py index 14f6b3224..737b7270f 100644 --- a/flash_rt/structures/discover.py +++ b/flash_rt/structures/discover.py @@ -371,6 +371,42 @@ def discover( "norm": "adaln_rms" if gated else "rms", "ffn_entry": "fp8_static"}, family=family, layer_index=idx)) + if "audio_codec" in structures and all( + hasattr(module, a) + for a in ("decode", "quantizer", "acoustic_decoder", + "decoder_semantic") + ): + # the neural codec decode region: DAC-style acoustic/semantic + # decoders behind a quantizer. HiggsAudioV2TokenizerModel is + # shared by OmniVoice and Higgs-Audio-v3 (recurring family). + seams.append(Seam( + structure="audio_codec", path=path, + parent_path=parent_path, norm_attr=None, + dims={"T": 0, "C": 0}, + variant={"backend": "fp16_codec"}, + family="audio_codec", layer_index=0)) + continue + if "decoder_llm" in structures and all( + hasattr(module, a) for a in ("layers", "embed_tokens", + "norm", "rotary_emb") + ): + # the whole decoder stack: layers + embeddings + final norm + + # rotary. One seam per host (named_modules yields the stack + # itself once; nested module lists are skipped by the slots + # check). The NVFP4 impl's profile envelope refuses hosts + # outside the native engine's v1 contract. + if not any(isinstance(getattr(m, "layers", None), + torch.nn.ModuleList) + for _, m in (("", module),)): + continue + dim = int(module.embed_tokens.weight.shape[-1]) + seams.append(Seam( + structure="decoder_llm", path=path, + parent_path=parent_path, norm_attr="norm", + dims={"D": dim, "L": len(module.layers)}, + variant={"backend": "nvfp4_llm"}, + family="decoder_stack", layer_index=0)) + continue if "modnorm_qkv_chain" in structures: chain_dims = _is_modnorm_qkv_chain(module) if chain_dims is not None: diff --git a/flash_rt/structures/impls/audio_codec/__init__.py b/flash_rt/structures/impls/audio_codec/__init__.py new file mode 100644 index 000000000..33aff44d1 --- /dev/null +++ b/flash_rt/structures/impls/audio_codec/__init__.py @@ -0,0 +1 @@ +from .fp16 import AudioCodecDecodeFp16, bind_codec_decode # noqa: F401 diff --git a/flash_rt/structures/impls/audio_codec/fp16.py b/flash_rt/structures/impls/audio_codec/fp16.py new file mode 100644 index 000000000..da8b185a0 --- /dev/null +++ b/flash_rt/structures/impls/audio_codec/fp16.py @@ -0,0 +1,67 @@ +"""FP16 implementation of the ``audio_codec`` structure. + +The host codec decode is compute-bound on fp32 cuDNN convolutions; +fp16 autocast is the only measured lever (1.22x, waveform cosine +1.00000 vs the fp32 host — torch.compile and CUDA graph both measure +~1.0x here). The impl replaces the codec module's ``decode`` and keeps +the host module for fallback and attribute delegation. + +Boundary: codes [B, C, T] (int64) -> audio_values [B, 1, N] (fp32). +""" + +from __future__ import annotations + +from functools import lru_cache + +import torch + +from ...guard import CAST_OK, PROCEED, GuardedSeam + + +class AudioCodecDecodeFp16(GuardedSeam, torch.nn.Module): + """Codec decode seam: fp16 autocast over the host's eager decode.""" + + _frt_host_attr = "host_codec" + _frt_can_fallback = True + + def __init__(self, host_codec, device): + super().__init__() + self.host_codec = host_codec + self._device = device + self._frt_arm(dtypes=(torch.long,), device=device) + self._frt_guard.notes["backend"] = "fp16_codec" + + def decode(self, codes: torch.Tensor): + admitted = self._frt_admit(codes) + if admitted is not PROCEED: + return admitted + with torch.autocast("cuda", dtype=torch.float16): + return self.host_codec.decode(codes) + + forward = decode + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + if name == "host_codec": + raise + return getattr(super().__getattr__("host_codec"), name) + + +@torch.no_grad() +def bind_codec_decode(host_codec, *, device=None, original=None): + """Bind the codec decode seam to the fp16 path.""" + dev = device or host_codec.device + seam = AudioCodecDecodeFp16(host_codec, dev) + if original is not None: + seam.host_codec = original + # bind-time smoke (AGENTS.md §2.8): one real decode, zero codes is + # not a valid shape — use a tiny all-zero frame set + probe = seam.decode(torch.zeros(1, 1, 4, device=dev, dtype=torch.long)) + audio = probe.audio_values + if not torch.isfinite(audio.float()).all(): + raise ValueError( + f"refused: audio_codec fp16 bind smoke produced non-finite " + f"audio {tuple(audio.shape)}") + return seam diff --git a/flash_rt/structures/impls/decoder_ffn/nvfp4_static.py b/flash_rt/structures/impls/decoder_ffn/nvfp4_static.py new file mode 100644 index 000000000..4d150f165 --- /dev/null +++ b/flash_rt/structures/impls/decoder_ffn/nvfp4_static.py @@ -0,0 +1,266 @@ +"""NVFP4 (W4A4) implementation of the ``decoder_ffn`` structure — native build. + +The local native build's FP4 GEMM tier serves the MLP seam: activation +quantize to NVFP4 (per-16-block swizzled scales), gate/up GEMM, fused +SiLU + quantize, down GEMM. Unlike the ``w8a16``/``w4a16`` hub decode +bands (M in [1, 8]), the native fp4 GEMM covers the MaskGIT-scale M, so +this backend is the prefill/large-M form. + +Kernel resolution follows the PR-175 tiering: hub artifact first, the +local native build second, the retained host module always the floor. +This impl consumes the local native build (``flash_rt.flash_rt_kernels``) +because the hub's fp4 packages do not ship a torch-2.13 variant. + +Boundary: normed activations in, FFN output out (BF16) — the host's +input-layernorm precedes the MLP, so no norm is fused here. Weights are +checkpoint-native ``[out, in]`` (``w_gate``/``w_up``: ``[F, D]``, +``w_down``: ``[D, F]``); no calibration data is required (per-16-block +weight scales at bind time, per-16-block activation scales per call). +""" + +from __future__ import annotations + +from collections.abc import Mapping +from functools import lru_cache + +import torch + +from ...guard import CAST_OK, PROCEED, GuardedSeam + +SUPPORT = { + "D": {"min": 512, "max": 16384, "multiple_of": 64}, + "F": {"min": 1024, "max": 16384, "multiple_of": 64}, + # the native fp4 GEMM serves any M the host throws at it + "M": {"min": 1, "max": 1 << 30}, + "m_classes": ("micro", "small", "medium", "large"), +} + +#: kernels this backend needs from the local native build +_NATIVE_SYMBOLS = ( + "fp4_w4a16_gemm_sm120_bf16out", + "fp4_w4a16_gemm_sm120_bf16out_pingpong", + "quantize_bf16_to_nvfp4_swizzled", + "quantize_bf16_to_nvfp4_swizzled_mse", + "silu_mul_merged_to_nvfp4_swizzled_bf16", +) + + +def _swizzled_sf_bytes(rows: int, cols: int) -> int: + assert cols % 16 == 0 + n_blocks = cols // 16 + n_row_super = (rows + 127) // 128 + n_col_super = (n_blocks + 3) // 4 + return n_row_super * n_col_super * 128 * 64 + + +@lru_cache(maxsize=1) +def _native(): + """The locally built native extension, or None when absent. + + Absence is a bind refusal, not a silent host path: a seam bound + without its kernels would fall back on every call and the ledger + would count a lie. + """ + try: + from flash_rt import flash_rt_kernels as fk + except ImportError: + return None + if any(getattr(fk, s, None) is None for s in _NATIVE_SYMBOLS): + return None + return fk + + +def _check(weights: Mapping[str, torch.Tensor]) -> tuple[int, int]: + w_gate, w_up, w_down = (weights["w_gate"], weights["w_up"], + weights["w_down"]) + dim_f, dim_d = w_gate.shape + if w_up.shape != (dim_f, dim_d) or w_down.shape != (dim_d, dim_f): + raise ValueError( + f"inconsistent weight dims: gate {tuple(w_gate.shape)}, " + f"up {tuple(w_up.shape)}, down {tuple(w_down.shape)}") + for name, dim in (("D", dim_d), ("F", dim_f)): + bounds = SUPPORT[name] + if not bounds["min"] <= dim <= bounds["max"]: + raise ValueError( + f"{name}={dim} outside support envelope " + f"[{bounds['min']}, {bounds['max']}]") + if dim % bounds["multiple_of"]: + raise ValueError( + f"{name}={dim} must be a multiple of " + f"{bounds['multiple_of']}") + return dim_d, dim_f + + +class BoundDecoderFfnNvfp4: + """MLP-seam callable: normed activations in, FFN output out (BF16).""" + + def __init__(self, fk, gate_up_packed, gate_up_sf, + down_packed, down_sf, dim_d, dim_f): + self._fk = fk + self._gate_up_packed = gate_up_packed + self._gate_up_sf = gate_up_sf + self._down_packed = down_packed + self._down_sf = down_sf + self._dim_d = dim_d + self._dim_f = dim_f + self._gu_variant = "pingpong" if 2 * dim_f >= 4096 else "default" + # per-row-count workspace cache: the host (MaskGIT) keeps a + # constant row count per generation, so the first call allocates + # and every later step reuses — no per-call empty()/zeros() + self._ws: dict[int, dict[str, torch.Tensor]] = {} + + def _workspace(self, m: int, d: int): + ws = self._ws.get(m) + if ws is not None: + return ws + dev = self._gate_up_packed.device + ws = { + "inp_packed": torch.empty(m, d // 2, dtype=torch.uint8, + device=dev), + "inp_sf": torch.zeros(_swizzled_sf_bytes(m, d), + dtype=torch.uint8, device=dev), + "dg": torch.empty(m, 2 * self._dim_f, dtype=torch.bfloat16, + device=dev), + "act_packed": torch.empty(m, self._dim_f // 2, + dtype=torch.uint8, device=dev), + "act_sf": torch.zeros(_swizzled_sf_bytes(m, self._dim_f), + dtype=torch.uint8, device=dev), + "out": torch.empty(m, d, dtype=torch.bfloat16, device=dev), + } + self._ws[m] = ws + return ws + + def ffn(self, normed: torch.Tensor) -> torch.Tensor: + fk = self._fk + shape = normed.shape + x = normed.reshape(-1, shape[-1]).to(torch.bfloat16).contiguous() + m = x.shape[0] + d = x.shape[-1] + st = torch.cuda.current_stream().cuda_stream + ws = self._workspace(m, d) + + fk.quantize_bf16_to_nvfp4_swizzled( + x.data_ptr(), ws["inp_packed"].data_ptr(), + ws["inp_sf"].data_ptr(), m, d, st) + + if self._gu_variant == "pingpong": + fk.fp4_w4a16_gemm_sm120_bf16out_pingpong( + ws["inp_packed"].data_ptr(), self._gate_up_packed.data_ptr(), + ws["dg"].data_ptr(), m, 2 * self._dim_f, d, + ws["inp_sf"].data_ptr(), self._gate_up_sf.data_ptr(), + 1.0, st) + else: + fk.fp4_w4a16_gemm_sm120_bf16out( + ws["inp_packed"].data_ptr(), self._gate_up_packed.data_ptr(), + ws["dg"].data_ptr(), m, 2 * self._dim_f, d, + ws["inp_sf"].data_ptr(), self._gate_up_sf.data_ptr(), + 1.0, st) + + fk.silu_mul_merged_to_nvfp4_swizzled_bf16( + ws["dg"].data_ptr(), ws["act_packed"].data_ptr(), + ws["act_sf"].data_ptr(), m, self._dim_f, st) + + fk.fp4_w4a16_gemm_sm120_bf16out( + ws["act_packed"].data_ptr(), self._down_packed.data_ptr(), + ws["out"].data_ptr(), m, d, self._dim_f, + ws["act_sf"].data_ptr(), self._down_sf.data_ptr(), 1.0, st) + return ws["out"].reshape(shape).to(normed.dtype) + + __call__ = ffn + + +class FusedGluMlpNvfp4(GuardedSeam, torch.nn.Module): + """MLP-seam module backed by the native NVFP4 FFN kernels. + + ``original`` is retained whole (host MLP naming varies across model + families), and attribute lookups fall through to it so hosts that + introspect the module they call keep working. + """ + + _frt_host_attr = "host_mlp" + _frt_can_fallback = True + + def __init__(self, bound: BoundDecoderFfnNvfp4, + original: torch.nn.Module | None = None): + super().__init__() + self._bound = bound + if original is not None: + self.host_mlp = original + guard = self._frt_arm(dtypes=CAST_OK, + device=bound._gate_up_packed.device, + k=int(bound._dim_d)) + guard.notes["backend"] = "nvfp4_static" + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + if name == "host_mlp": + raise + return getattr(super().__getattr__("host_mlp"), name) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + admitted = self._frt_admit(hidden) + if admitted is not PROCEED: + return admitted + return self._bound.ffn(hidden) + + +@torch.no_grad() +def bind_mlp_seam( + weights: Mapping[str, torch.Tensor], + *, + variant: Mapping[str, str], + original: torch.nn.Module | None = None, +): + """Bind the MLP-seam slice of ``decoder_ffn`` with native NVFP4. + + ``weights`` uses checkpoint-native ``[out, in]`` projection layout + (``w_gate``/``w_up``: ``[F, D]``, ``w_down``: ``[D, F]``). Weights + are packed at bind time with MSE per-16-block scales; activations + are quantized per call, so no calibration data is required. + """ + if variant.get("activation", "silu") != "silu": + raise ValueError( + f"refused: nvfp4_static serves the SiLU gated FFN only, " + f"got activation {variant.get('activation')!r}") + fk = _native() + if fk is None: + raise ValueError( + "refused: nvfp4_static needs the locally built " + "flash_rt_kernels (flash_rt_kernels + fp4 GEMM symbols); " + "rebuild with -DGPU_ARCH=120") + dim_d, dim_f = _check(weights) + dev = weights["w_gate"].device + + def _pack(w: torch.Tensor): + w = w.to("cuda", torch.bfloat16).contiguous() + n, k = w.shape + packed = torch.empty(n, k // 2, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_swizzled_sf_bytes(n, k), dtype=torch.uint8, + device="cuda") + fk.quantize_bf16_to_nvfp4_swizzled_mse( + w.data_ptr(), packed.data_ptr(), sf.data_ptr(), n, k, + torch.cuda.current_stream().cuda_stream) + return packed, sf + + gate_up = torch.cat([weights["w_gate"], weights["w_up"]], dim=0) + gate_up_packed, gate_up_sf = _pack(gate_up) + down_packed, down_sf = _pack(weights["w_down"]) + torch.cuda.synchronize() + + bound = BoundDecoderFfnNvfp4( + fk, gate_up_packed, gate_up_sf, down_packed, down_sf, + dim_d, dim_f) + + # bind-time smoke: one launch through the real entry point before the + # seam is handed out (AGENTS.md §2.8). The probe carries the host's + # real rank/shape class — 2D, BF16, CUDA. + probe = bound.ffn(torch.zeros(16, dim_d, device=dev, + dtype=torch.bfloat16)) + if probe.shape != (16, dim_d) or not torch.isfinite(probe).all(): + raise ValueError( + f"refused: nvfp4_static bind smoke produced shape " + f"{tuple(probe.shape)}, " + f"finite={bool(torch.isfinite(probe).all())}") + return FusedGluMlpNvfp4(bound, original=original) diff --git a/flash_rt/structures/impls/decoder_llm/__init__.py b/flash_rt/structures/impls/decoder_llm/__init__.py new file mode 100644 index 000000000..39c8a6f93 --- /dev/null +++ b/flash_rt/structures/impls/decoder_llm/__init__.py @@ -0,0 +1 @@ +from .nvfp4 import (DecoderLlmNvfp4, bind_decoder_stack, maskgit_gen) # noqa: F401 diff --git a/flash_rt/structures/impls/decoder_llm/nvfp4.py b/flash_rt/structures/impls/decoder_llm/nvfp4.py new file mode 100644 index 000000000..b0a592574 --- /dev/null +++ b/flash_rt/structures/impls/decoder_llm/nvfp4.py @@ -0,0 +1,206 @@ +"""NVFP4 whole-LLM implementation of ``decoder_llm`` — native build. + +The native omnivoice engines (``flash_rt.models.omnivoice.pipeline_rtx``) +promoted into the structures layer: the whole decoder stack forward is +the fused FP4 engine (fp4 GEMMs + fused qk-norm+RoPE + FA2 + fused +residual/norm/quant), with whole-step CUDA graph capture. Kernel +resolution is the PR-175 tiering: hub artifact first, the local native +build second, the retained host stack always the floor. This impl +consumes the local build. + +Dispatch inside the stack boundary mirrors the native schedule: the +CFG batch (B=2) rides the BF16 engine, single-stream (B=1) rides the +FP4 graph. The MaskGIT two-phase loop is exposed as ``maskgit_gen`` on +this module — the schedule structure for non-text generation. + +The impl's profile envelope is the native engine's v1 contract +(D=1024, L=28, NH=16, NKV=8, HD=128, FFN=3072, RoPE theta 1e6); other +profiles refuse cleanly (the fp8-KV band precedent). +""" + +from __future__ import annotations + +from functools import lru_cache + +import torch + +from ...guard import CAST_OK, PROCEED, GuardedSeam + +#: the native engine's v1 profile (kernel constants are compiled in) +PROFILE = dict(D=1024, L=28, NH=16, NKV=8, HD=128, FFN=3072) + + +@lru_cache(maxsize=1) +def _native(): + """(FlashRTLlm, FlashRTLlmBF16) from the local build, or None.""" + try: + from flash_rt.models.omnivoice.pipeline_rtx import ( # noqa: F401 + FlashRTLlm, + FlashRTLlmBF16, + ) + from flash_rt import flash_rt_kernels # noqa: F401 + from flash_rt import flash_rt_omnivoice # noqa: F401 + from flash_rt import flash_rt_fa2 # noqa: F401 + return FlashRTLlm, FlashRTLlmBF16 + except ImportError: + return None + + +class DecoderLlmNvfp4(GuardedSeam, torch.nn.Module): + """Whole-stack seam: fused FP4 engine with per-batch dispatch.""" + + _frt_host_attr = "host_llm" + _frt_can_fallback = True + + def __init__(self, host_llm, bf16, fp4, device): + super().__init__() + self.host_llm = host_llm + self._bf16 = bf16 + self._fp4 = fp4 + self._device = device + guard = self._frt_arm(dtypes=CAST_OK, device=device) + guard.notes["backend"] = "nvfp4_llm" + + def forward(self, *a, **kw): + e = kw.get("inputs_embeds") + if e is None and len(a) >= 2 and isinstance(a[1], torch.Tensor): + e = a[1] + if e is None: + return self._frt_host()(*a, **kw) + admitted = self._frt_admit(e, *a, **kw) + if admitted is not PROCEED: + return admitted + if e.shape[0] > 1: + # CFG batch: BF16 fused engine (no quantization drift) + return self._bf16.forward(e, attention_mask=kw.get( + "attention_mask")) + return self._fp4.forward_graph(e, attention_mask=kw.get( + "attention_mask")) + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + if name == "host_llm": + raise + return getattr(super().__getattr__("host_llm"), name) + + +def _profile_of(host_llm) -> dict: + cfg = getattr(host_llm, "config", None) + d = int(cfg.hidden_size) if cfg is not None else 0 + l = len(host_llm.layers) if hasattr(host_llm, "layers") else 0 + if cfg is not None: + nh = int(cfg.num_attention_heads) + nkv = int(getattr(cfg, "num_key_value_heads", nh)) + hd = int(getattr(cfg, "head_dim", d // nh)) + f = int(cfg.intermediate_size) + else: + nh = nkv = hd = f = 0 + return dict(D=d, L=l, NH=nh, NKV=nkv, HD=hd, FFN=f) + + +@torch.no_grad() +def bind_decoder_stack(host_llm, *, variant=None, original=None): + """Bind the whole stack to the native NVFP4 engines. + + ``host_llm`` is the decoder stack module (layers / embed_tokens / + norm / rotary_emb). Weights are read in place; the host stack is + retained for fallback, so it stays whole. + """ + eng = _native() + if eng is None: + raise ValueError( + "refused: decoder_llm nvfp4 needs the locally built " + "flash_rt_kernels + flash_rt_omnivoice + flash_rt_fa2; " + "rebuild with -DFLASHRT_ENABLE_OMNIVOICE=ON -DGPU_ARCH=120") + got = _profile_of(host_llm) + if got != PROFILE: + raise ValueError( + f"refused: decoder_llm nvfp4 profile mismatch — the native " + f"engine's v1 contract is {PROFILE}, this host is {got}") + FlashRTLlm, FlashRTLlmBF16 = eng + dev = str(host_llm.embed_tokens.weight.device) + bf16 = FlashRTLlmBF16(host_llm, dev) + fp4 = FlashRTLlm(host_llm, dev) + # calibrate: packed weights + workspaces (native inject's shapes) + c2 = torch.randn(2, 178, PROFILE["D"], device=dev, + dtype=torch.bfloat16) * 0.02 + bf16.calibrate(c2) + fp4.calibrate(c2) + fp4.WL_bf16 = None + bf16.WL_fp4 = None + bf16._fp4_act = None + bf16._alphas = None + c1 = torch.randn(1, 178, PROFILE["D"], device=dev, + dtype=torch.bfloat16) * 0.02 + for _ in range(3): + fp4.forward(c1) + torch.cuda.synchronize() + fp4._capture_graph(c1) + for _ in range(3): + fp4.forward_graph(c1) + torch.cuda.synchronize() + bf16._graph = None + + seam = DecoderLlmNvfp4(host_llm, bf16, fp4, torch.device(dev)) + seam._fp4 = fp4 # maskgit_gen's schedule reads the engine here + if original is not None: + seam.host_llm = original + + # bind-time smoke: one real forward through the engines (AGENTS §2.8) + probe = seam.forward(inputs_embeds=torch.randn( + 1, 16, PROFILE["D"], device=dev, dtype=torch.bfloat16)) + if isinstance(probe, tuple): + probe = probe[0] + if tuple(probe.shape) != (1, 16, PROFILE["D"]) or \ + not torch.isfinite(probe).all(): + raise ValueError( + f"refused: decoder_llm nvfp4 bind smoke produced " + f"{tuple(probe.shape)}, " + f"finite={bool(torch.isfinite(probe).all())}") + return seam + + +class MaskgitLoop: + """The MaskGIT schedule as a serving object (the decode_loop twin). + + ``structures.maskgit_loop(model)`` returns one of these; ``generate`` + runs the two-phase loop (BF16 CFG steps, then FP4 noCFG single-stream + graph replays) over whatever ``decoder_llm`` seam is attached. + """ + + def __init__(self, model, *, cfg_ratio: float = 0.05, + bookend: bool = False): + self._model = model + self._cfg_ratio = cfg_ratio + self._bookend = bookend + + def generate(self, task, gen_config): + return maskgit_gen(self._model, task, gen_config, + cfg_ratio=self._cfg_ratio, + bookend=self._bookend) + + +def maskgit_gen(model, task, gen_config, cfg_ratio=0.05, bookend=False): + """The MaskGIT two-phase schedule over an attached decoder_llm seam. + + Phase 1: BF16 CFG (B=2, cfg_ratio fraction of steps). Phase 2: FP4 + noCFG (B=1, graph replay). Uses the native ``_optimize_maskgit`` + loop; the seam's per-batch dispatch replaces the inject's forward + monkeypatching (B=2 -> BF16 engine, B=1 -> FP4 graph). + """ + from flash_rt.models.omnivoice import pipeline_rtx as prtx + + seam = model.llm + fp4 = getattr(seam, "_fp4", None) + if fp4 is None: + raise ValueError( + "maskgit_gen: no decoder_llm structure attached (auto_swaps " + "with decoder_llm + scheme nvfp4_static)") + orig = model._generate_iterative + prtx._optimize_maskgit(model, fp4, cfg_ratio, bookend) + try: + return model._generate_iterative(task, gen_config) + finally: + model._generate_iterative = orig diff --git a/flash_rt/structures/impls/linear_proj/nvfp4_static.py b/flash_rt/structures/impls/linear_proj/nvfp4_static.py new file mode 100644 index 000000000..564de02c8 --- /dev/null +++ b/flash_rt/structures/impls/linear_proj/nvfp4_static.py @@ -0,0 +1,174 @@ +"""NVFP4 (W4A4) ``linear_proj`` implementation — native build. + +Single projection as one fp4 GEMM: activation quantize to NVFP4 +(per-16-block swizzled scales), ``fp4_w4a16_gemm_sm120_bf16out``, BF16 +output. Serves the attention Q/K/V/O projections of hosts whose forward +M exceeds the hub decode bands; kernel resolution is the PR-175 tiering +(hub first, local native build second, host floor) — this impl consumes +the local build because the hub's fp4 packages ship no torch-2.13 +variant. + +``weights`` is checkpoint-native ``[N, K]`` (out, in). No calibration +data required: per-16-block weight scales at bind time, per-call +activation scales. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from functools import lru_cache + +import torch + +from ...guard import CAST_OK, PROCEED, GuardedSeam + +SUPPORT = { + "K": {"min": 512, "max": 16384, "multiple_of": 16}, + "N": {"min": 128, "max": 65536, "multiple_of": 8}, +} + +_NATIVE_SYMBOLS = ( + "fp4_w4a16_gemm_sm120_bf16out", + "quantize_bf16_to_nvfp4_swizzled", + "quantize_bf16_to_nvfp4_swizzled_mse", +) + + +def _swizzled_sf_bytes(rows: int, cols: int) -> int: + assert cols % 16 == 0 + n_blocks = cols // 16 + n_row_super = (rows + 127) // 128 + n_col_super = (n_blocks + 3) // 4 + return n_row_super * n_col_super * 128 * 64 + + +@lru_cache(maxsize=1) +def _native(): + try: + from flash_rt import flash_rt_kernels as fk + except ImportError: + return None + if any(getattr(fk, s, None) is None for s in _NATIVE_SYMBOLS): + return None + return fk + + +def _check(weights: Mapping[str, torch.Tensor]) -> tuple[int, int]: + w = weights["w"] + if w.dim() != 2: + raise ValueError(f"w must be [N, K], got {tuple(w.shape)}") + n, k = w.shape + for name, dim in (("K", k), ("N", n)): + bounds = SUPPORT[name] + if dim < bounds["min"]: + raise ValueError( + f"{name}={dim} outside support envelope " + f"(min {bounds['min']})") + if bounds.get("multiple_of") and dim % bounds["multiple_of"]: + raise ValueError( + f"{name}={dim} must be a multiple of " + f"{bounds['multiple_of']}") + return n, k + + +class LinearProjNvfp4(GuardedSeam, torch.nn.Module): + """Single projection: fp4 GEMM with runtime activation scales.""" + + _frt_can_fallback = True + + def __init__(self, fk, w_packed, w_sf, bias, n, k): + super().__init__() + self._fk = fk + self._w_packed = w_packed + self._w_sf = w_sf + self._bias = bias + self._n = n + self._k = k + self._ws: dict[int, dict[str, torch.Tensor]] = {} + self._frt_arm(dtypes=CAST_OK, device=w_packed.device, k=int(k)) + + def _workspace(self, m: int): + ws = self._ws.get(m) + if ws is not None: + return ws + dev = self._w_packed.device + ws = { + "a_packed": torch.empty(m, self._k // 2, dtype=torch.uint8, + device=dev), + "a_sf": torch.zeros(_swizzled_sf_bytes(m, self._k), + dtype=torch.uint8, device=dev), + "y": torch.empty(m, self._n, dtype=torch.bfloat16, device=dev), + } + self._ws[m] = ws + return ws + + def forward(self, x: torch.Tensor) -> torch.Tensor: + admitted = self._frt_admit(x) + if admitted is not PROCEED: + return admitted + fk = self._fk + shape = x.shape + flat = x.reshape(-1, shape[-1]).to(torch.bfloat16).contiguous() + m = flat.shape[0] + st = torch.cuda.current_stream().cuda_stream + ws = self._workspace(m) + fk.quantize_bf16_to_nvfp4_swizzled( + flat.data_ptr(), ws["a_packed"].data_ptr(), + ws["a_sf"].data_ptr(), m, self._k, st) + fk.fp4_w4a16_gemm_sm120_bf16out( + ws["a_packed"].data_ptr(), self._w_packed.data_ptr(), + ws["y"].data_ptr(), m, self._n, self._k, + ws["a_sf"].data_ptr(), self._w_sf.data_ptr(), 1.0, st) + y = ws["y"].reshape(*shape[:-1], self._n) + if self._bias is not None: + y = y + self._bias + return y.type_as(x) + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + if name == "host_linear": + raise + return getattr(super().__getattr__("host_linear"), name) + + +@torch.no_grad() +def bind_proj_seam( + weights: Mapping[str, torch.Tensor], + *, + original: torch.nn.Module | None = None, +): + """Bind one projection from a dense ``[N, K]`` weight.""" + fk = _native() + if fk is None: + raise ValueError( + "refused: linear_proj nvfp4_static needs the locally built " + "flash_rt_kernels (fp4 GEMM symbols); rebuild with " + "-DGPU_ARCH=120") + n, k = _check(weights) + w = weights["w"].to("cuda", torch.bfloat16).contiguous() + packed = torch.empty(n, k // 2, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_swizzled_sf_bytes(n, k), dtype=torch.uint8, + device="cuda") + fk.quantize_bf16_to_nvfp4_swizzled_mse( + w.data_ptr(), packed.data_ptr(), sf.data_ptr(), n, k, + torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize() + bias = None + b = weights.get("bias") + if b is not None and b.numel(): + bias = b.to("cuda", torch.bfloat16).contiguous() + bound = LinearProjNvfp4(fk, packed, sf, bias, n, k) + if original is not None: + bound.host_linear = original + + # bind-time smoke (AGENTS.md §2.8) + probe = bound.forward(torch.zeros(16, k, device="cuda", + dtype=torch.bfloat16)) + if probe.shape != (16, n) or not torch.isfinite(probe).all(): + raise ValueError( + f"refused: linear_proj nvfp4_static bind smoke produced " + f"shape {tuple(probe.shape)}, " + f"finite={bool(torch.isfinite(probe).all())}") + return bound diff --git a/flash_rt/structures/schemes.py b/flash_rt/structures/schemes.py index e684b39e0..21c2da71b 100644 --- a/flash_rt/structures/schemes.py +++ b/flash_rt/structures/schemes.py @@ -214,6 +214,46 @@ def decide(self, report: Mapping[str, Mapping[str, float]]) -> Decision: formats=formats) +class Nvfp4Static(QuantScheme): + """Native NVFP4 W4A4 FFN — the large-M sibling of the W4A16 decode + band, served by the local native build's fp4 GEMM tier (no hub + package needed; the impl refuses without the built kernels). + + Weight-only per-16-block quantization at bind time, per-call + activation quantization — needs no calibration data, exactly like + the weight-only decode twins. Routes ``decoder_ffn`` seams to the + ``nvfp4_static`` impl, whose native GEMM covers MaskGIT-scale M + (the hub decode bands stop at M=8). Other structures stay at host + precision. + """ + + name = "nvfp4_static" + _format = "nvfp4_static" + _linear_format = "nvfp4_static" + + def statistics(self, points: Sequence) -> dict[str, PointStat]: + return {f"{p.path}|{p.name}": PointStat(None) for p in points} + + def decide(self, report: Mapping[str, Mapping[str, float]]) -> Decision: + formats, keep = {}, [] + for seam_path, pts in report.items(): + if any(k.endswith("|act_after_mul") for k in pts): + formats[seam_path] = self._format + elif getattr(pts, "structure", None) == "decoder_llm": + formats[seam_path] = self._format + elif getattr(pts, "structure", None) == "audio_codec": + formats[seam_path] = "fp16_codec" + elif (self._linear_format is not None + and getattr(pts, "structure", None) == "linear_proj"): + formats[seam_path] = self._linear_format + else: + keep.append(seam_path) + return Decision(keep_host=tuple(keep), + reasons={p: f"{self.name} binds LLM/MLP/proj " + f"seams only" for p in keep}, + formats=formats) + + class W4A16Decode(W8A16Decode): """Weight-only NVFP4 (E2M1 packed + block scale factors) twin of :class:`W8A16Decode` — same decode band, same M-dispatch, half the @@ -481,6 +521,7 @@ def resolve_auto() -> str: register("fp8_static_keep_outliers", Fp8Static(keep_outliers=20.0)) register("w8a16_decode", W8A16Decode()) register("w4a16_decode", W4A16Decode()) +register("nvfp4_static", Nvfp4Static()) register("w4a4_decode", W4A4Decode()) register("w4a4_decode_release", W4A4Decode(release_host_weights=True)) register("none", NoQuant()) diff --git a/tests/test_structures_decoder_llm.py b/tests/test_structures_decoder_llm.py new file mode 100644 index 000000000..b9dd916f5 --- /dev/null +++ b/tests/test_structures_decoder_llm.py @@ -0,0 +1,173 @@ +"""Contract pins for the ``decoder_llm`` structure and ``nvfp4_static`` +scheme. CPU-safe: nothing here requires kernels, a GPU, or a model — +the native-build tier refuses cleanly without ``flash_rt_kernels``, and +that refusal is itself the contract pin. +""" + +from __future__ import annotations + +import pytest + +from flash_rt.structures import schemes +from flash_rt.structures.binding import load_binding +from flash_rt.structures.registry import load + + +def test_autoplan_exposes_one_call_attach_door(): + """The article's ``plan.attach()``: an AutoPlan carries its root and + commits its swaps atomically (refuses cleanly without a root).""" + from flash_rt.structures.autobuild import AutoPlan + + plan = AutoPlan() + assert callable(getattr(plan, "attach", None)) + with pytest.raises(ValueError, match="no root host"): + plan.attach() + plan.root = object() + with pytest.raises(ValueError, match="no swaps or routed seams"): + plan.attach() + + +def test_maskgit_loop_door_is_exported(): + from flash_rt import structures + + assert callable(structures.maskgit_loop) + assert "maskgit_loop" in structures.__all__ + + +def test_decoder_llm_catalog_entry_loads(): + spec = load("decoder_llm") + assert spec.kind == "region" + assert spec.symbolic_dims == ("B", "S", "D") + # a whole-stack seam binds weights in place: no remapped slots + assert spec.weight_slots == () + assert spec.reference() is not None + + +def test_omnivoice_llm_binding_loads(): + binding = load_binding("omnivoice_llm") + assert binding.name == "omnivoice_llm" + assert binding.structure.name == "decoder_llm" + assert binding.data["dims"]["D"] == 1024 + assert binding.data["dims"]["L"] == 28 + assert not binding.is_pipeline + + +def test_nvfp4_static_scheme_registered(): + assert "nvfp4_static" in schemes.names() + scheme = schemes.get("nvfp4_static") + assert scheme._format == "nvfp4_static" + assert scheme._linear_format == "nvfp4_static" + + +class _FakeStats(dict): + def __init__(self, *args, structure=None, **kwargs): + super().__init__(*args, **kwargs) + self.structure = structure + + +def test_nvfp4_static_routes_decoder_ffn_and_linear_proj(): + scheme = schemes.get("nvfp4_static") + + report = { + "llm": _FakeStats({}, structure="decoder_llm"), + "llm.layers.0.mlp": _FakeStats( + {"llm.layers.0.mlp|act_after_mul": 0.5}, + structure="decoder_ffn"), + "llm.layers.0.self_attn.q_proj": _FakeStats( + {}, structure="linear_proj"), + "llm.layers.0.self_attn.o_proj": _FakeStats( + {}, structure="linear_proj"), + "some.other.seam": _FakeStats( + {"some.other.seam|x": 0.1}, structure="other"), + } + decision = scheme.decide(report) + assert decision.formats == { + "llm": "nvfp4_static", + "llm.layers.0.mlp": "nvfp4_static", + "llm.layers.0.self_attn.q_proj": "nvfp4_static", + "llm.layers.0.self_attn.o_proj": "nvfp4_static", + } + assert decision.keep_host == ("some.other.seam",) + + +def test_decoder_llm_native_refuses_without_local_build(): + from flash_rt.structures.impls.decoder_llm import nvfp4 as llm_impl + + if llm_impl._native() is not None: + pytest.skip("local flash_rt_kernels build present") + with pytest.raises(ValueError, match="refused: decoder_llm nvfp4"): + llm_impl.bind_decoder_stack(None) + + +def test_decoder_ffn_nvfp4_refuses_without_local_build(): + from flash_rt.structures.impls.decoder_ffn import ( + nvfp4_static as ffn_impl) + + if ffn_impl._native() is not None: + pytest.skip("local flash_rt_kernels build present") + with pytest.raises(ValueError, match="refused: nvfp4_static needs"): + ffn_impl.bind_mlp_seam( + {"w_gate": None, "w_up": None, "w_down": None}, + variant={"activation": "silu"}) + + +def test_linear_proj_nvfp4_refuses_without_local_build(): + from flash_rt.structures.impls.linear_proj import ( + nvfp4_static as proj_impl) + + if proj_impl._native() is not None: + pytest.skip("local flash_rt_kernels build present") + with pytest.raises(ValueError, match="refused: linear_proj"): + proj_impl.bind_proj_seam({"w": None}) + + +def test_audio_codec_catalog_entry_loads(): + spec = load("audio_codec") + assert spec.kind == "region" + assert spec.symbolic_dims == ("B", "C", "T", "N") + + +def test_audio_codec_binding_loads(): + binding = load_binding("omnivoice_audio_codec") + assert binding.structure.name == "audio_codec" + assert binding.data["hosts"]["omnivoice"]["module_path"] == \ + "audio_tokenizer" + + +def test_nvfp4_static_routes_audio_codec(): + scheme = schemes.get("nvfp4_static") + report = { + "audio_tokenizer": _FakeStats({}, structure="audio_codec"), + } + decision = scheme.decide(report) + assert decision.formats == {"audio_tokenizer": "fp16_codec"} + + +def test_discovery_finds_decoder_llm_seam_on_stack_slots(): + """The discovery rule keys on slots (layers/embed_tokens/norm/ + rotary_emb), never on model names.""" + import torch + + from flash_rt.structures import discover + + class DummyRotary(torch.nn.Module): + pass + + class Stack(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList() + self.embed_tokens = torch.nn.Embedding(100, 16) + self.norm = torch.nn.LayerNorm(16) + self.rotary_emb = DummyRotary() + + class Host(torch.nn.Module): + def __init__(self): + super().__init__() + self.llm = Stack() + + seams = discover.discover(Host(), structures=("decoder_llm",)) + assert len(seams) == 1 + assert seams[0].structure == "decoder_llm" + assert seams[0].path == "llm" + assert seams[0].dims["D"] == 16 From 29d09151bfa25c591181f57e67049cd98a441876 Mon Sep 17 00:00:00 2001 From: shideqin Date: Sun, 16 Aug 2026 20:57:50 +0800 Subject: [PATCH 2/2] docs: record CFG-on-FP4 rejection with measurements --- docs/omnivoice_structures.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/omnivoice_structures.md b/docs/omnivoice_structures.md index c3361196d..1805a6430 100644 --- a/docs/omnivoice_structures.md +++ b/docs/omnivoice_structures.md @@ -28,6 +28,19 @@ 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