diff --git "a/torch_hf_optimizations.py" "b/torch_hf_optimizations.py"
new file mode 100644--- /dev/null
+++ "b/torch_hf_optimizations.py"
@@ -0,0 +1,3595 @@
+from __future__ import annotations
+
+import types
+import time
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.cache_utils import DynamicCache
+from transformers.generation.logits_process import (
+    LogitsProcessorList,
+    RepetitionPenaltyLogitsProcessor,
+    TemperatureLogitsWarper,
+    TopKLogitsWarper,
+    TopPLogitsWarper,
+)
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.qwen3.modeling_qwen3 import eager_attention_forward
+
+try:
+    from moss_tts_local_clipper_checkpoint.inference_utils import find_last_equal_C
+except ImportError:
+    def find_last_equal_C(tensor, C):
+        mask = (tensor == C).int()
+        flipped_mask = mask.flip(dims=[1])
+        flipped_indices = flipped_mask.argmax(dim=1)
+        seq_len = tensor.shape[1]
+        last_indices = (seq_len - 1) - flipped_indices
+        actual_values = tensor[torch.arange(tensor.shape[0]), last_indices]
+        no_match = actual_values != C
+        last_indices[no_match] = -1
+        return last_indices
+
+try:
+    import triton
+    import triton.language as tl
+except Exception:  # pragma: no cover - optional optimization path
+    triton = None
+    tl = None
+
+
+if triton is not None:
+    @triton.jit
+    def _triton_rmsnorm_kernel(x_ptr, weight_ptr, y_ptr, n_cols: tl.constexpr, eps: tl.constexpr, block: tl.constexpr):
+        row = tl.program_id(0)
+        offsets = tl.arange(0, block)
+        mask = offsets < n_cols
+        x = tl.load(x_ptr + row * n_cols + offsets, mask=mask, other=0.0).to(tl.float32)
+        weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        variance = tl.sum(x * x, axis=0) / n_cols
+        y = x * tl.rsqrt(variance + eps) * weight
+        tl.store(y_ptr + row * n_cols + offsets, y, mask=mask)
+
+
+    @triton.jit
+    def _triton_top_p_sample_kernel(
+        scores_ptr,
+        uniform_ptr,
+        out_ptr,
+        n_cols: tl.constexpr,
+        stride_row: tl.constexpr,
+        temperature: tl.constexpr,
+        top_p: tl.constexpr,
+        block: tl.constexpr,
+    ):
+        row = tl.program_id(0)
+        offsets = tl.arange(0, block)
+        mask = offsets < n_cols
+        scores = tl.load(
+            scores_ptr + row * stride_row + offsets,
+            mask=mask,
+            other=-float("inf"),
+        ).to(tl.float32) / temperature
+        scores = tl.where(mask, scores, -float("inf"))
+        max_score = tl.max(scores, axis=0)
+        exp_scores = tl.exp(scores - max_score)
+        exp_scores = tl.where(mask, exp_scores, 0.0)
+        total_mass = tl.sum(exp_scores, axis=0)
+
+        sorted_scores = tl.sort(scores, descending=True)
+        sorted_exp = tl.exp(sorted_scores - max_score)
+        sorted_exp = tl.where(offsets < n_cols, sorted_exp, 0.0)
+        sorted_cumulative = tl.cumsum(sorted_exp, 0)
+        cutoff_target = top_p * total_mass
+        cutoff_positions = tl.where(sorted_cumulative >= cutoff_target, offsets, block)
+        cutoff_pos = tl.min(cutoff_positions, axis=0)
+        cutoff_score = tl.max(tl.where(offsets == cutoff_pos, sorted_scores, -float("inf")), axis=0)
+
+        kept_exp = tl.where((scores >= cutoff_score) & mask, exp_scores, 0.0)
+        kept_mass = tl.sum(kept_exp, axis=0)
+        draw = tl.load(uniform_ptr + row).to(tl.float32)
+        draw = tl.minimum(tl.maximum(draw, 5.960464477539063e-8), 0.9999999403953552)
+        sample_target = draw * kept_mass
+        kept_cumulative = tl.cumsum(kept_exp, 0)
+        sample_positions = tl.where(kept_cumulative >= sample_target, offsets, block)
+        sample_pos = tl.min(sample_positions, axis=0)
+        sample_pos = tl.minimum(sample_pos, n_cols - 1)
+        tl.store(out_ptr + row, sample_pos)
+
+
+    @triton.jit
+    def _triton_single_query_gqa_attention_kernel(
+        q_ptr,
+        k_ptr,
+        v_ptr,
+        out_ptr,
+        seq_len: tl.constexpr,
+        num_kv_groups: tl.constexpr,
+        head_dim: tl.constexpr,
+        q_stride_b: tl.constexpr,
+        q_stride_h: tl.constexpr,
+        q_stride_d: tl.constexpr,
+        k_stride_b: tl.constexpr,
+        k_stride_h: tl.constexpr,
+        k_stride_t: tl.constexpr,
+        k_stride_d: tl.constexpr,
+        v_stride_b: tl.constexpr,
+        v_stride_h: tl.constexpr,
+        v_stride_t: tl.constexpr,
+        v_stride_d: tl.constexpr,
+        o_stride_b: tl.constexpr,
+        o_stride_h: tl.constexpr,
+        o_stride_d: tl.constexpr,
+        scale: tl.constexpr,
+        block_t: tl.constexpr,
+        block_d: tl.constexpr,
+    ):
+        batch = tl.program_id(0)
+        head = tl.program_id(1)
+        kv_head = head // num_kv_groups
+        offs_t = tl.arange(0, block_t)
+        offs_d = tl.arange(0, block_d)
+        t_mask = offs_t < seq_len
+        d_mask = offs_d < head_dim
+
+        q = tl.load(
+            q_ptr + batch * q_stride_b + head * q_stride_h + offs_d * q_stride_d,
+            mask=d_mask,
+            other=0.0,
+        ).to(tl.float32)
+        k = tl.load(
+            k_ptr
+            + batch * k_stride_b
+            + kv_head * k_stride_h
+            + offs_t[:, None] * k_stride_t
+            + offs_d[None, :] * k_stride_d,
+            mask=t_mask[:, None] & d_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        scores = tl.sum(k * q[None, :], axis=1) * scale
+        scores = tl.where(t_mask, scores, -float("inf"))
+        scores = scores - tl.max(scores, axis=0)
+        probs = tl.exp(scores)
+        probs = probs / tl.sum(probs, axis=0)
+        v = tl.load(
+            v_ptr
+            + batch * v_stride_b
+            + kv_head * v_stride_h
+            + offs_t[:, None] * v_stride_t
+            + offs_d[None, :] * v_stride_d,
+            mask=t_mask[:, None] & d_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        out = tl.sum(v * probs[:, None], axis=0)
+        tl.store(
+            out_ptr + batch * o_stride_b + head * o_stride_h + offs_d * o_stride_d,
+            out,
+            mask=d_mask,
+        )
+
+
+    @triton.jit
+    def _triton_packed_qkv_norm_cache_kernel(
+        qkv_ptr,
+        q_norm_weight_ptr,
+        k_norm_weight_ptr,
+        query_out_ptr,
+        key_cache_ptr,
+        value_cache_ptr,
+        cache_index,
+        qkv_stride_b: tl.constexpr,
+        qkv_stride_t: tl.constexpr,
+        qkv_stride_d: tl.constexpr,
+        query_stride_b: tl.constexpr,
+        query_stride_h: tl.constexpr,
+        query_stride_d: tl.constexpr,
+        key_stride_b: tl.constexpr,
+        key_stride_h: tl.constexpr,
+        key_stride_t: tl.constexpr,
+        key_stride_d: tl.constexpr,
+        value_stride_b: tl.constexpr,
+        value_stride_h: tl.constexpr,
+        value_stride_t: tl.constexpr,
+        value_stride_d: tl.constexpr,
+        num_q_heads: tl.constexpr,
+        num_kv_heads: tl.constexpr,
+        head_dim: tl.constexpr,
+        q_size: tl.constexpr,
+        k_size: tl.constexpr,
+        q_eps: tl.constexpr,
+        k_eps: tl.constexpr,
+        block_d: tl.constexpr,
+    ):
+        batch = tl.program_id(0)
+        head = tl.program_id(1)
+        offsets = tl.arange(0, block_d)
+        mask = offsets < head_dim
+
+        if head < num_q_heads:
+            q = tl.load(
+                qkv_ptr
+                + batch * qkv_stride_b
+                + head * head_dim * qkv_stride_d
+                + offsets * qkv_stride_d,
+                mask=mask,
+                other=0.0,
+            ).to(tl.float32)
+            q_weight = tl.load(q_norm_weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+            q_var = tl.sum(q * q, axis=0) / head_dim
+            q_out = q * tl.rsqrt(q_var + q_eps) * q_weight
+            tl.store(
+                query_out_ptr
+                + batch * query_stride_b
+                + head * query_stride_h
+                + offsets * query_stride_d,
+                q_out,
+                mask=mask,
+            )
+
+        if head < num_kv_heads:
+            k = tl.load(
+                qkv_ptr
+                + batch * qkv_stride_b
+                + (q_size + head * head_dim) * qkv_stride_d
+                + offsets * qkv_stride_d,
+                mask=mask,
+                other=0.0,
+            ).to(tl.float32)
+            k_weight = tl.load(k_norm_weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+            k_var = tl.sum(k * k, axis=0) / head_dim
+            k_out = k * tl.rsqrt(k_var + k_eps) * k_weight
+            tl.store(
+                key_cache_ptr
+                + batch * key_stride_b
+                + head * key_stride_h
+                + cache_index * key_stride_t
+                + offsets * key_stride_d,
+                k_out,
+                mask=mask,
+            )
+            v = tl.load(
+                qkv_ptr
+                + batch * qkv_stride_b
+                + (q_size + k_size + head * head_dim) * qkv_stride_d
+                + offsets * qkv_stride_d,
+                mask=mask,
+                other=0.0,
+            )
+            tl.store(
+                value_cache_ptr
+                + batch * value_stride_b
+                + head * value_stride_h
+                + cache_index * value_stride_t
+                + offsets * value_stride_d,
+                v,
+                mask=mask,
+            )
+
+
+    @triton.jit
+    def _triton_fused_audio_lm_head_sample_kernel(
+        hidden_ptr,
+        weight_ptr,
+        uniform_ptr,
+        out_ptr,
+        vocab_size: tl.constexpr,
+        hidden_size: tl.constexpr,
+        hidden_stride: tl.constexpr,
+        weight_stride_v: tl.constexpr,
+        weight_stride_h: tl.constexpr,
+        temperature: tl.constexpr,
+        top_p: tl.constexpr,
+        pad_token_id: tl.constexpr,
+        block_v: tl.constexpr,
+        block_d: tl.constexpr,
+    ):
+        offsets_v = tl.arange(0, block_v)
+        offsets_d = tl.arange(0, block_d)
+        mask_v = offsets_v < vocab_size
+        acc = tl.zeros((block_v,), tl.float32)
+        for start_d in range(0, hidden_size, block_d):
+            d = start_d + offsets_d
+            mask_d = d < hidden_size
+            hidden = tl.load(
+                hidden_ptr + d * hidden_stride,
+                mask=mask_d,
+                other=0.0,
+            ).to(tl.float32)
+            weight = tl.load(
+                weight_ptr
+                + offsets_v[:, None] * weight_stride_v
+                + d[None, :] * weight_stride_h,
+                mask=mask_v[:, None] & mask_d[None, :],
+                other=0.0,
+            ).to(tl.float32)
+            acc += tl.sum(weight * hidden[None, :], axis=1)
+
+        scores = tl.where(mask_v, acc, -float("inf"))
+        if pad_token_id >= 0:
+            scores = tl.where(offsets_v == pad_token_id, -float("inf"), scores)
+        scores = scores / temperature
+
+        max_score = tl.max(scores, axis=0)
+        exp_scores = tl.exp(scores - max_score)
+        exp_scores = tl.where(mask_v, exp_scores, 0.0)
+        total_mass = tl.sum(exp_scores, axis=0)
+
+        sorted_scores = tl.sort(scores, descending=True)
+        sorted_exp = tl.exp(sorted_scores - max_score)
+        sorted_exp = tl.where(offsets_v < vocab_size, sorted_exp, 0.0)
+        sorted_cumulative = tl.cumsum(sorted_exp, 0)
+        cutoff_target = top_p * total_mass
+        cutoff_positions = tl.where(sorted_cumulative >= cutoff_target, offsets_v, block_v)
+        cutoff_pos = tl.min(cutoff_positions, axis=0)
+        cutoff_score = tl.max(tl.where(offsets_v == cutoff_pos, sorted_scores, -float("inf")), axis=0)
+
+        kept_exp = tl.where((scores >= cutoff_score) & mask_v, exp_scores, 0.0)
+        kept_mass = tl.sum(kept_exp, axis=0)
+        draw = tl.load(uniform_ptr).to(tl.float32)
+        draw = tl.minimum(tl.maximum(draw, 5.960464477539063e-8), 0.9999999403953552)
+        sample_target = draw * kept_mass
+        kept_cumulative = tl.cumsum(kept_exp, 0)
+        sample_positions = tl.where(kept_cumulative >= sample_target, offsets_v, block_v)
+        sample_pos = tl.min(sample_positions, axis=0)
+        sample_pos = tl.minimum(sample_pos, vocab_size - 1)
+        tl.store(out_ptr, sample_pos)
+
+
+class _TritonRMSNorm(nn.Module):
+    def __init__(self, norm: nn.Module):
+        super().__init__()
+        self.weight = norm.weight
+        self.variance_epsilon = _rmsnorm_eps_float(norm)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if (
+            triton is None
+            or not hidden_states.is_cuda
+            or hidden_states.dtype not in {torch.float16, torch.bfloat16, torch.float32}
+        ):
+            input_dtype = hidden_states.dtype
+            values = hidden_states.to(torch.float32)
+            variance = values.pow(2).mean(-1, keepdim=True)
+            return self.weight * (values * torch.rsqrt(variance + self.variance_epsilon)).to(input_dtype)
+        x = hidden_states.contiguous()
+        n_cols = int(x.shape[-1])
+        if n_cols > 4096:
+            input_dtype = hidden_states.dtype
+            values = hidden_states.to(torch.float32)
+            variance = values.pow(2).mean(-1, keepdim=True)
+            return self.weight * (values * torch.rsqrt(variance + self.variance_epsilon)).to(input_dtype)
+        y = torch.empty_like(x)
+        rows = int(x.numel() // n_cols)
+        block = 1 << (n_cols - 1).bit_length()
+        _triton_rmsnorm_kernel[(rows,)](
+            x,
+            self.weight,
+            y,
+            n_cols,
+            self.variance_epsilon,
+            block=block,
+            num_warps=8 if block >= 2048 else 4,
+        )
+        return y
+
+
+class _PackedGateUpMLP(nn.Module):
+    def __init__(self, mlp: nn.Module):
+        super().__init__()
+        self.gate_proj = mlp.gate_proj
+        self.up_proj = mlp.up_proj
+        self.down_proj = mlp.down_proj
+        self.norm = getattr(mlp, "norm", None)
+        self.prenorm = bool(getattr(mlp, "prenorm", self.norm is not None))
+        self._packed_gate_up_weight: torch.Tensor | None = None
+        self._packed_gate_up_key: tuple[int, int, int, int] | None = None
+        self._packed_gate_up_static = False
+
+    def _packed_weight(self) -> torch.Tensor:
+        if self._packed_gate_up_static and self._packed_gate_up_weight is not None:
+            return self._packed_gate_up_weight
+        gate_weight = self.gate_proj.weight
+        up_weight = self.up_proj.weight
+        key = (
+            gate_weight.data_ptr(),
+            up_weight.data_ptr(),
+            gate_weight._version,
+            up_weight._version,
+        )
+        if self._packed_gate_up_key != key or self._packed_gate_up_weight is None:
+            self._packed_gate_up_weight = torch.cat(
+                (gate_weight.detach(), up_weight.detach()),
+                dim=0,
+            ).contiguous()
+            self._packed_gate_up_key = key
+        return self._packed_gate_up_weight
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.norm is not None:
+            x = self.norm(x)
+        gate_up = F.linear(x, self._packed_weight())
+        gate, up = gate_up.chunk(2, dim=-1)
+        return self.down_proj(F.silu(gate) * up)
+
+
+class _PackedQKVAttention(nn.Module):
+    def __init__(self, attention: nn.Module):
+        super().__init__()
+        self.config = attention.config
+        self.layer_idx = attention.layer_idx
+        self.q_proj = attention.q_proj
+        self.k_proj = attention.k_proj
+        self.v_proj = attention.v_proj
+        self.o_proj = attention.o_proj
+        self.q_norm = attention.q_norm
+        self.k_norm = attention.k_norm
+        self.head_dim = attention.head_dim
+        self.num_key_value_groups = attention.num_key_value_groups
+        self.attention_dropout = attention.attention_dropout
+        self.scaling = attention.scaling
+        self.sliding_window = attention.sliding_window
+        self._packed_qkv_weight: torch.Tensor | None = None
+        self._packed_qkv_key: tuple[int, int, int, int, int, int] | None = None
+        self._packed_qkv_static = False
+
+    def _packed_weight(self) -> torch.Tensor:
+        if self._packed_qkv_static and self._packed_qkv_weight is not None:
+            return self._packed_qkv_weight
+        q_weight = self.q_proj.weight
+        k_weight = self.k_proj.weight
+        v_weight = self.v_proj.weight
+        key = (
+            q_weight.data_ptr(),
+            k_weight.data_ptr(),
+            v_weight.data_ptr(),
+            q_weight._version,
+            k_weight._version,
+            v_weight._version,
+        )
+        if self._packed_qkv_key != key or self._packed_qkv_weight is None:
+            self._packed_qkv_weight = torch.cat(
+                (q_weight.detach(), k_weight.detach(), v_weight.detach()),
+                dim=0,
+            ).contiguous()
+            self._packed_qkv_key = key
+        return self._packed_qkv_weight
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings=None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_value=None,
+        cache_position: torch.Tensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q_size = self.q_proj.out_features
+        k_size = self.k_proj.out_features
+        qkv = F.linear(hidden_states, self._packed_weight())
+        query_states, key_states, value_states = qkv.split(
+            (q_size, k_size, self.v_proj.out_features),
+            dim=-1,
+        )
+        query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(key_states.view(hidden_shape)).transpose(1, 2)
+        value_states = value_states.view(hidden_shape).transpose(1, 2)
+
+        if past_key_value is not None:
+            cache_kwargs = {"cache_position": cache_position}
+            key_states, value_states = past_key_value.update(
+                key_states,
+                value_states,
+                self.layer_idx,
+                cache_kwargs,
+            )
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            if self.config._attn_implementation == "sdpa" and kwargs.get(
+                "output_attentions",
+                False,
+            ):
+                attention_interface = eager_attention_forward
+            else:
+                attention_interface = ALL_ATTENTION_FUNCTIONS[
+                    self.config._attn_implementation
+                ]
+
+        if past_key_value is None:
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                is_causal=True,
+                attention_mask=None,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                sliding_window=self.sliding_window,
+                **kwargs,
+            )
+        else:
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                sliding_window=self.sliding_window,
+                **kwargs,
+            )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        return self.o_proj(attn_output), attn_weights
+
+
+def _pack_local_transformer_qkv(model: Any) -> None:
+    for layer in model.local_transformer.layers[
+        : model.local_transformer.config.num_hidden_layers
+    ]:
+        if not isinstance(layer.self_attn, _PackedQKVAttention):
+            layer.self_attn = _PackedQKVAttention(layer.self_attn)
+
+
+def _pack_local_transformer_mlps(model: Any) -> None:
+    for layer in model.local_transformer.layers[
+        : model.local_transformer.config.num_hidden_layers
+    ]:
+        if not isinstance(layer.mlp, _PackedGateUpMLP):
+            layer.mlp = _PackedGateUpMLP(layer.mlp)
+
+
+def _pack_adapter_mlps(model: Any, scope: str) -> None:
+    if scope not in {"all", "input", "heads"}:
+        raise ValueError("--packed-adapter-mlp-scope must be one of all, input, heads")
+    if scope in {"all", "input"} and not isinstance(
+        model.speech_embedding_to_local_mlp,
+        _PackedGateUpMLP,
+    ):
+        model.speech_embedding_to_local_mlp = _PackedGateUpMLP(
+            model.speech_embedding_to_local_mlp
+        )
+    if scope in {"all", "heads"}:
+        for index, mlp in enumerate(model.local_to_speech_embedding_mlps):
+            if not isinstance(mlp, _PackedGateUpMLP):
+                model.local_to_speech_embedding_mlps[index] = _PackedGateUpMLP(mlp)
+
+
+def _materialize_packed_weights(model: Any, static: bool) -> None:
+    for module in model.modules():
+        if isinstance(module, _PackedGateUpMLP):
+            module._packed_weight()
+            module._packed_gate_up_static = static
+        elif isinstance(module, _PackedQKVAttention):
+            module._packed_weight()
+            module._packed_qkv_static = static
+
+
+def _wrap_rmsnorm(module: nn.Module) -> nn.Module:
+    if isinstance(module, _TritonRMSNorm):
+        return module
+    if hasattr(module, "weight") and (
+        hasattr(module, "variance_epsilon") or hasattr(module, "eps")
+    ):
+        return _TritonRMSNorm(module)
+    return module
+
+
+def _rmsnorm_eps_float(module: nn.Module) -> float:
+    saved = getattr(module, "_torchopt_rmsnorm_eps_float", None)
+    if saved is not None:
+        return float(saved)
+    value = getattr(module, "variance_epsilon", getattr(module, "eps", 1e-6))
+    if isinstance(value, torch.Tensor):
+        # This helper must not be called on a CUDA tensor inside compiled/captured code.
+        return float(value.detach().cpu().item())
+    return float(value)
+
+
+def _install_local_triton_rmsnorms(model: Any) -> None:
+    for layer in model.local_transformer.layers[
+        : model.local_transformer.config.num_hidden_layers
+    ]:
+        layer.input_layernorm = _wrap_rmsnorm(layer.input_layernorm)
+        layer.post_attention_layernorm = _wrap_rmsnorm(layer.post_attention_layernorm)
+        layer.self_attn.q_norm = _wrap_rmsnorm(layer.self_attn.q_norm)
+        layer.self_attn.k_norm = _wrap_rmsnorm(layer.self_attn.k_norm)
+    model.local_transformer.norm = _wrap_rmsnorm(model.local_transformer.norm)
+    for index, norm in enumerate(model.layer_norm_before_lm_heads):
+        model.layer_norm_before_lm_heads[index] = _wrap_rmsnorm(norm)
+
+
+def tensorize_rmsnorm_eps(model: Any, *, device: torch.device | None = None) -> int:
+    count = 0
+    for module in model.modules():
+        weight = getattr(module, "weight", None)
+        if not isinstance(weight, torch.Tensor):
+            continue
+        target_device = device if device is not None else weight.device
+        for attr in ("variance_epsilon", "eps"):
+            value = getattr(module, attr, None)
+            if isinstance(value, (float, int)):
+                setattr(module, "_torchopt_rmsnorm_eps_float", float(value))
+                setattr(
+                    module,
+                    attr,
+                    torch.tensor(float(value), device=target_device, dtype=torch.float32),
+                )
+                count += 1
+    return count
+
+
+def install_torch_frame_sampler(
+    model: Any,
+    *,
+    mode: str = "fixed-full-cudagraph",
+    compile_mode: str | None = "max-autotune-no-cudagraphs",
+    packed_local_qkv: bool = False,
+    packed_local_mlp: bool = False,
+    packed_adapter_mlp: bool = False,
+    packed_adapter_mlp_scope: str = "heads",
+    static_packed_weights: bool = False,
+    triton_rmsnorm: bool = False,
+    triton_top_p: bool = False,
+    triton_fused_lm_head: bool = False,
+    triton_qkv_cache: bool = False,
+    tensorrt_local: bool = False,
+    fast_prepare_inputs: bool = False,
+    fast_control_head: bool = False,
+    feedback_lookup: bool = False,
+    local_compile_fullgraph: bool = False,
+    top_p_prefilter_size: int = 0,
+    top_p_sampler_cudagraph: bool = False,
+) -> None:
+    if mode == "none":
+        return
+    if mode not in {
+        "fixed-full",
+        "fixed-full-cudagraph",
+        "fixed-full-compile",
+        "fixed-full-compile-cudagraph",
+        "greedy-full-compile-cudagraph",
+        "prefix-full-compile-cudagraph",
+        "static-local-cache",
+        "static-local-cache-compile",
+        "static-local-cache-triton-compile-cudagraph",
+    }:
+        raise ValueError(f"Unsupported torch optimization mode: {mode}")
+
+    if not hasattr(model, "_torchopt_original_sample"):
+        model._torchopt_original_sample = model._sample
+    if packed_local_qkv:
+        _pack_local_transformer_qkv(model)
+    if packed_local_mlp:
+        _pack_local_transformer_mlps(model)
+    if packed_adapter_mlp:
+        _pack_adapter_mlps(model, packed_adapter_mlp_scope)
+    if packed_local_qkv or packed_local_mlp or packed_adapter_mlp:
+        _materialize_packed_weights(model, static=static_packed_weights)
+    if triton_rmsnorm:
+        _install_local_triton_rmsnorms(model)
+    if tensorrt_local:
+        _install_tensorrt_local_transformer(model)
+    model._torchopt_mode = mode
+    model._torchopt_compile_mode = compile_mode
+    model._torchopt_compiled_local = None
+    model._torchopt_compiled_greedy_frame = None
+    model._torchopt_compiled_static_local = None
+    model._torchopt_compiled_top_p_samplers = {}
+    model._torchopt_compiled_stochastic_top_p_frames = {}
+    model._torchopt_top_p_graphs = {}
+    model._torchopt_top_p_prefilter_size = max(0, int(top_p_prefilter_size))
+    model._torchopt_top_p_sampler_cudagraph = bool(top_p_sampler_cudagraph)
+    model._torchopt_triton_top_p = bool(triton_top_p)
+    model._torchopt_triton_fused_lm_head = bool(triton_fused_lm_head)
+    model._torchopt_triton_qkv_cache = bool(triton_qkv_cache)
+    model._torchopt_tensorrt_local_enabled = bool(tensorrt_local)
+    model._torchopt_fast_prepare_inputs = bool(fast_prepare_inputs)
+    model._torchopt_fast_control_head = bool(fast_control_head)
+    model._torchopt_feedback_lookup = bool(feedback_lookup)
+    model._torchopt_feedback_lookup_cache = {}
+    model._torchopt_local_compile_fullgraph = bool(local_compile_fullgraph)
+    model._torchopt_control_token_ids = None
+    model._torchopt_control_lm_head_weight = None
+    if fast_control_head:
+        control_ids = torch.tensor(
+            [
+                int(model.config.audio_assistant_gen_slot_token_id),
+                int(model.config.audio_end_token_id),
+            ],
+            device=model.lm_heads[0].weight.device,
+            dtype=torch.long,
+        )
+        model._torchopt_control_token_ids = control_ids
+        model._torchopt_control_lm_head_weight = model.lm_heads[0].weight.detach().index_select(
+            0,
+            control_ids,
+        ).contiguous()
+    model._torchopt_frame_graphs = {}
+    model._torchopt_channel_graphs = {}
+    model._sample = types.MethodType(_optimized_sample, model)
+
+
+class _LocalTransformerOnly(nn.Module):
+    def __init__(self, local_transformer: nn.Module):
+        super().__init__()
+        self.local_transformer = local_transformer
+
+    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
+        hidden_states = inputs_embeds
+        for decoder_layer in self.local_transformer.layers[
+            : self.local_transformer.config.num_hidden_layers
+        ]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=None,
+                position_ids=None,
+                past_key_value=None,
+                output_attentions=False,
+                use_cache=False,
+                cache_position=None,
+                position_embeddings=None,
+            )
+        return self.local_transformer.norm(hidden_states)
+
+
+def _install_tensorrt_local_transformer(model: Any) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("--tensorrt-local requires CUDA.")
+    try:
+        import torch_tensorrt
+    except Exception as exc:
+        raise RuntimeError("torch_tensorrt is required for --tensorrt-local.") from exc
+
+    first_param = next(model.local_transformer.parameters())
+    dtype = first_param.dtype
+    if dtype != torch.float16:
+        raise RuntimeError("--tensorrt-local currently requires the local transformer to be fp16.")
+    shape = (
+        1,
+        int(model.channels),
+        int(model.local_transformer_config.hidden_size),
+    )
+    wrapper = _LocalTransformerOnly(model.local_transformer).to(
+        device=first_param.device,
+        dtype=dtype,
+    )
+    wrapper.eval()
+    trt_module = torch_tensorrt.compile(
+        wrapper,
+        ir="dynamo",
+        inputs=[torch_tensorrt.Input(shape, dtype=dtype)],
+    )
+    model._torchopt_tensorrt_local = trt_module
+    model._torchopt_tensorrt_local_shape = shape
+
+
+def _run_local_transformer_impl(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
+    trt_module = getattr(self, "_torchopt_tensorrt_local", None)
+    trt_shape = getattr(self, "_torchopt_tensorrt_local_shape", None)
+    if trt_module is not None and tuple(inputs_embeds.shape) == tuple(trt_shape):
+        return trt_module(inputs_embeds)
+
+    hidden_states = inputs_embeds
+    for decoder_layer in self.local_transformer.layers[
+        : self.local_transformer.config.num_hidden_layers
+    ]:
+        hidden_states = decoder_layer(
+            hidden_states,
+            attention_mask=None,
+            position_ids=None,
+            past_key_value=None,
+            output_attentions=False,
+            use_cache=False,
+            cache_position=None,
+            position_embeddings=None,
+        )
+    return self.local_transformer.norm(hidden_states)
+
+
+def _can_use_fast_control_head(self, channel: int, do_sample: bool) -> bool:
+    return (
+        channel == 0
+        and not do_sample
+        and bool(getattr(self, "_torchopt_fast_control_head", False))
+        and getattr(self, "_torchopt_control_token_ids", None) is not None
+        and getattr(self, "_torchopt_control_lm_head_weight", None) is not None
+    )
+
+
+def _sample_fast_control_head(self, lm_hidden: torch.Tensor) -> torch.Tensor:
+    control_logits = F.linear(lm_hidden, self._torchopt_control_lm_head_weight)
+    control_index = torch.argmax(control_logits, dim=-1)
+    return self._torchopt_control_token_ids[control_index]
+
+
+def _run_local_transformer_fixed(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
+    mode = getattr(self, "_torchopt_mode", "fixed-full")
+    use_compile = mode in {"fixed-full-compile", "fixed-full-compile-cudagraph"}
+    if not use_compile:
+        return _run_local_transformer_impl(self, inputs_embeds)
+
+    if getattr(self, "_torchopt_compiled_local", None) is None:
+        compile_mode = getattr(self, "_torchopt_compile_mode", None)
+        if compile_mode == "default":
+            compile_mode = None
+        self._torchopt_compiled_local = torch.compile(
+            types.MethodType(_run_local_transformer_impl, self),
+            dynamic=False,
+            mode=compile_mode,
+        )
+    return self._torchopt_compiled_local(inputs_embeds)
+
+
+def _repeat_kv_for_local(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return hidden_states
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch,
+        num_key_value_heads,
+        n_rep,
+        slen,
+        head_dim,
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def _triton_single_query_gqa_attention(
+    query_states: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    seq_len: int,
+    *,
+    scale: float,
+    num_kv_groups: int,
+) -> torch.Tensor | None:
+    if (
+        triton is None
+        or not query_states.is_cuda
+        or query_states.ndim != 3
+        or key_cache.ndim != 4
+        or value_cache.ndim != 4
+        or query_states.dtype not in {torch.float16, torch.bfloat16, torch.float32}
+    ):
+        return None
+    batch, num_heads, head_dim = query_states.shape
+    if int(head_dim) > 256:
+        return None
+    query_states = query_states.contiguous()
+    out = torch.empty_like(query_states)
+    block_t = 1 << (max(1, int(seq_len)) - 1).bit_length()
+    block_t = max(16, min(64, block_t))
+    block_d = 1 << (int(head_dim) - 1).bit_length()
+    _triton_single_query_gqa_attention_kernel[(int(batch), int(num_heads))](
+        query_states,
+        key_cache,
+        value_cache,
+        out,
+        int(seq_len),
+        int(num_kv_groups),
+        int(head_dim),
+        query_states.stride(0),
+        query_states.stride(1),
+        query_states.stride(2),
+        key_cache.stride(0),
+        key_cache.stride(1),
+        key_cache.stride(2),
+        key_cache.stride(3),
+        value_cache.stride(0),
+        value_cache.stride(1),
+        value_cache.stride(2),
+        value_cache.stride(3),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        float(scale),
+        block_t=block_t,
+        block_d=block_d,
+        num_warps=4,
+    )
+    return out
+
+
+def _triton_packed_qkv_norm_cache(
+    qkv: torch.Tensor,
+    attn: nn.Module,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    cache_index: int,
+) -> torch.Tensor | None:
+    if (
+        triton is None
+        or not qkv.is_cuda
+        or qkv.dim() != 3
+        or qkv.shape[1] != 1
+        or qkv.dtype not in {torch.float16, torch.bfloat16, torch.float32}
+    ):
+        return None
+    head_dim = int(attn.head_dim)
+    if head_dim <= 0 or head_dim > 256:
+        return None
+    q_size = int(attn.q_proj.out_features)
+    k_size = int(attn.k_proj.out_features)
+    v_size = int(attn.v_proj.out_features)
+    if q_size % head_dim != 0 or k_size % head_dim != 0 or v_size % head_dim != 0:
+        return None
+    num_q_heads = q_size // head_dim
+    num_kv_heads = k_size // head_dim
+    if v_size // head_dim != num_kv_heads:
+        return None
+    qkv = qkv.contiguous()
+    query_out = torch.empty(
+        qkv.shape[0],
+        num_q_heads,
+        head_dim,
+        device=qkv.device,
+        dtype=qkv.dtype,
+    )
+    block_d = 1 << (head_dim - 1).bit_length()
+    q_eps = _rmsnorm_eps_float(attn.q_norm)
+    k_eps = _rmsnorm_eps_float(attn.k_norm)
+    _triton_packed_qkv_norm_cache_kernel[(int(qkv.shape[0]), max(num_q_heads, num_kv_heads))](
+        qkv,
+        attn.q_norm.weight,
+        attn.k_norm.weight,
+        query_out,
+        key_cache,
+        value_cache,
+        int(cache_index),
+        qkv.stride(0),
+        qkv.stride(1),
+        qkv.stride(2),
+        query_out.stride(0),
+        query_out.stride(1),
+        query_out.stride(2),
+        key_cache.stride(0),
+        key_cache.stride(1),
+        key_cache.stride(2),
+        key_cache.stride(3),
+        value_cache.stride(0),
+        value_cache.stride(1),
+        value_cache.stride(2),
+        value_cache.stride(3),
+        num_q_heads,
+        num_kv_heads,
+        head_dim,
+        q_size,
+        k_size,
+        q_eps,
+        k_eps,
+        block_d=block_d,
+        num_warps=4,
+    )
+    return query_out
+
+
+def _get_feedback_lookup_tables(
+    self,
+    *,
+    local_num_steps: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> list[torch.Tensor | None] | None:
+    if not bool(getattr(self, "_torchopt_feedback_lookup", False)):
+        return None
+    cache = getattr(self, "_torchopt_feedback_lookup_cache", None)
+    if cache is None:
+        cache = {}
+        self._torchopt_feedback_lookup_cache = cache
+    key = (
+        device.index if device.type == "cuda" else -1,
+        device.type,
+        dtype,
+        int(local_num_steps),
+    )
+    tables = cache.get(key)
+    if tables is not None:
+        return tables
+
+    channels = int(getattr(self, "channels", local_num_steps))
+    tables = [None] * channels
+    embedding_list = self.model.embedding_list
+    # Channel 0 may use the text/control vocabulary, so only precompute the
+    # small audio-codebook feedback channels.
+    last_feedback_channel = max(1, min(int(local_num_steps) - 1, len(embedding_list)))
+    with torch.inference_mode():
+        for channel in range(1, last_feedback_channel):
+            embedding_weight = embedding_list[channel].weight.to(
+                device=device,
+                dtype=dtype,
+            )
+            tables[channel] = self.speech_embedding_to_local_mlp(
+                embedding_weight
+            ).contiguous()
+    cache[key] = tables
+    return tables
+
+
+def _run_local_transformer_static_cache_impl(
+    self,
+    inputs_embeds: torch.Tensor,
+    *,
+    key_caches: list[torch.Tensor],
+    value_caches: list[torch.Tensor],
+    cache_index: int,
+) -> torch.Tensor:
+    hidden_states = inputs_embeds[:, None, :]
+    batch = inputs_embeds.shape[0]
+    for layer_index, decoder_layer in enumerate(
+        self.local_transformer.layers[: self.local_transformer.config.num_hidden_layers]
+    ):
+        residual = hidden_states
+        hidden_states = decoder_layer.input_layernorm(hidden_states)
+        attn = decoder_layer.self_attn
+        hidden_shape = (batch, 1, -1, attn.head_dim)
+        if isinstance(attn, _PackedQKVAttention):
+            q_size = attn.q_proj.out_features
+            k_size = attn.k_proj.out_features
+            qkv = F.linear(hidden_states, attn._packed_weight())
+            query_states, key_states, value_states = qkv.split(
+                (q_size, k_size, attn.v_proj.out_features),
+                dim=-1,
+            )
+            query_states = attn.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+            key_states = attn.k_norm(key_states.view(hidden_shape)).transpose(1, 2)
+            value_states = value_states.view(hidden_shape).transpose(1, 2)
+        else:
+            query_states = attn.q_norm(attn.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+            key_states = attn.k_norm(attn.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+            value_states = attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_caches[layer_index][:, :, cache_index : cache_index + 1, :].copy_(key_states)
+        value_caches[layer_index][:, :, cache_index : cache_index + 1, :].copy_(value_states)
+        key_all = _repeat_kv_for_local(
+            key_caches[layer_index][:, :, : cache_index + 1, :],
+            attn.num_key_value_groups,
+        )
+        value_all = _repeat_kv_for_local(
+            value_caches[layer_index][:, :, : cache_index + 1, :],
+            attn.num_key_value_groups,
+        )
+        attn_output = F.scaled_dot_product_attention(
+            query_states,
+            key_all,
+            value_all,
+            attn_mask=None,
+            dropout_p=0.0,
+            is_causal=False,
+            scale=attn.scaling,
+        )
+        attn_output = attn_output.transpose(1, 2).reshape(batch, 1, -1).contiguous()
+        hidden_states = residual + attn.o_proj(attn_output)
+
+        residual = hidden_states
+        hidden_states = decoder_layer.post_attention_layernorm(hidden_states)
+        hidden_states = residual + decoder_layer.mlp(hidden_states)
+    return self.local_transformer.norm(hidden_states)[:, 0, :]
+
+
+def _run_local_transformer_static_cache_triton_attention_impl(
+    self,
+    inputs_embeds: torch.Tensor,
+    *,
+    key_caches: list[torch.Tensor],
+    value_caches: list[torch.Tensor],
+    cache_index: int,
+) -> torch.Tensor:
+    hidden_states = inputs_embeds[:, None, :]
+    batch = inputs_embeds.shape[0]
+    for layer_index, decoder_layer in enumerate(
+        self.local_transformer.layers[: self.local_transformer.config.num_hidden_layers]
+    ):
+        residual = hidden_states
+        hidden_states = decoder_layer.input_layernorm(hidden_states)
+        attn = decoder_layer.self_attn
+        hidden_shape = (batch, 1, -1, attn.head_dim)
+        if isinstance(attn, _PackedQKVAttention):
+            q_size = attn.q_proj.out_features
+            k_size = attn.k_proj.out_features
+            qkv = F.linear(hidden_states, attn._packed_weight())
+            query_for_attention = None
+            if bool(getattr(self, "_torchopt_triton_qkv_cache", False)):
+                query_for_attention = _triton_packed_qkv_norm_cache(
+                    qkv,
+                    attn,
+                    key_caches[layer_index],
+                    value_caches[layer_index],
+                    cache_index,
+                )
+            if query_for_attention is None:
+                query_states, key_states, value_states = qkv.split(
+                    (q_size, k_size, attn.v_proj.out_features),
+                    dim=-1,
+                )
+                query_states = attn.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+                key_states = attn.k_norm(key_states.view(hidden_shape)).transpose(1, 2)
+                value_states = value_states.view(hidden_shape).transpose(1, 2)
+                key_caches[layer_index][:, :, cache_index : cache_index + 1, :].copy_(key_states)
+                value_caches[layer_index][:, :, cache_index : cache_index + 1, :].copy_(value_states)
+                query_for_attention = query_states[:, :, 0, :]
+        else:
+            query_states = attn.q_norm(attn.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+            key_states = attn.k_norm(attn.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+            value_states = attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            key_caches[layer_index][:, :, cache_index : cache_index + 1, :].copy_(key_states)
+            value_caches[layer_index][:, :, cache_index : cache_index + 1, :].copy_(value_states)
+            query_for_attention = query_states[:, :, 0, :]
+        attn_output = _triton_single_query_gqa_attention(
+            query_for_attention,
+            key_caches[layer_index],
+            value_caches[layer_index],
+            cache_index + 1,
+            scale=float(attn.scaling),
+            num_kv_groups=int(attn.num_key_value_groups),
+        )
+        if attn_output is None:
+            key_all = _repeat_kv_for_local(
+                key_caches[layer_index][:, :, : cache_index + 1, :],
+                attn.num_key_value_groups,
+            )
+            value_all = _repeat_kv_for_local(
+                value_caches[layer_index][:, :, : cache_index + 1, :],
+                attn.num_key_value_groups,
+            )
+            attn_output = F.scaled_dot_product_attention(
+                query_for_attention[:, :, None, :],
+                key_all,
+                value_all,
+                attn_mask=None,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=attn.scaling,
+            )
+            attn_output = attn_output.transpose(1, 2).reshape(batch, 1, -1).contiguous()
+        else:
+            attn_output = attn_output.reshape(batch, 1, -1).contiguous()
+        hidden_states = residual + attn.o_proj(attn_output)
+
+        residual = hidden_states
+        hidden_states = decoder_layer.post_attention_layernorm(hidden_states)
+        hidden_states = residual + decoder_layer.mlp(hidden_states)
+    return self.local_transformer.norm(hidden_states)[:, 0, :]
+
+
+def _run_local_transformer_static_cache(
+    self,
+    inputs_embeds: torch.Tensor,
+    *,
+    key_caches: list[torch.Tensor],
+    value_caches: list[torch.Tensor],
+    cache_index: int,
+) -> torch.Tensor:
+    mode = getattr(self, "_torchopt_mode", "fixed-full")
+    if mode != "static-local-cache-compile":
+        return _run_local_transformer_static_cache_impl(
+            self,
+            inputs_embeds,
+            key_caches=key_caches,
+            value_caches=value_caches,
+            cache_index=cache_index,
+        )
+    compiled = getattr(self, "_torchopt_compiled_static_local", None)
+    if compiled is None:
+        compile_mode = getattr(self, "_torchopt_compile_mode", None)
+        if compile_mode == "default":
+            compile_mode = None
+        compiled = torch.compile(
+            types.MethodType(_run_local_transformer_static_cache_impl, self),
+            dynamic=False,
+            mode=compile_mode,
+        )
+        self._torchopt_compiled_static_local = compiled
+    return compiled(
+        inputs_embeds,
+        key_caches=key_caches,
+        value_caches=value_caches,
+        cache_index=cache_index,
+    )
+
+
+def _sample_frame_fixed_full_impl(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    input_ids: torch.Tensor,
+    realprocessor: list[LogitsProcessorList],
+    do_samples: list[bool],
+    local_num_steps: int,
+    fast_top_p: bool = False,
+    temperatures: list[float] | None = None,
+    top_ps: list[float] | None = None,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    dtype = hidden_states.dtype
+    device = hidden_states.device
+    local_inputs = torch.zeros(
+        batch,
+        local_num_steps,
+        self.local_transformer_config.hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    for channel in range(local_num_steps):
+        local_inputs[:, channel, :] = current_input
+        local_outputs = _run_local_transformer_fixed(self, local_inputs)
+        local_hidden = local_outputs[:, channel, :]
+        lm_hidden = self.layer_norm_before_lm_heads[channel](
+            self.local_to_speech_embedding_mlps[channel](local_hidden)
+        )
+        if _can_use_fast_control_head(self, channel, do_samples[channel]):
+            token = _sample_fast_control_head(self, lm_hidden)
+        else:
+            logits = self.lm_heads[channel](lm_hidden)
+            if channel != 0:
+                logits[:, self.config.audio_pad_code] = -torch.inf
+            if do_samples[channel]:
+                if fast_top_p and temperatures is not None and top_ps is not None:
+                    token = _sample_top_p_only_compiled(
+                        self,
+                        logits,
+                        temperature=temperatures[channel],
+                        top_p=top_ps[channel],
+                    )
+                else:
+                    scores = realprocessor[channel](input_ids[..., channel], logits)
+                    token = torch.multinomial(F.softmax(scores, dim=-1), num_samples=1).squeeze(1)
+            else:
+                token = torch.argmax(logits, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            lookup = (
+                feedback_lookup_tables[channel]
+                if feedback_lookup_tables is not None and channel < len(feedback_lookup_tables)
+                else None
+            )
+            if lookup is not None:
+                current_input = lookup.index_select(0, token.reshape(-1))
+            else:
+                current_input = self.speech_embedding_to_local_mlp(
+                    self.model.embedding_list[channel](token)
+                )
+    return next_tokens
+
+
+def _sample_frame_greedy_full_impl(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    local_num_steps: int,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    dtype = hidden_states.dtype
+    device = hidden_states.device
+    local_inputs = torch.zeros(
+        batch,
+        local_num_steps,
+        self.local_transformer_config.hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    for channel in range(local_num_steps):
+        local_inputs[:, channel, :] = current_input
+        local_outputs = _run_local_transformer_impl(self, local_inputs)
+        local_hidden = local_outputs[:, channel, :]
+        lm_hidden = self.layer_norm_before_lm_heads[channel](
+            self.local_to_speech_embedding_mlps[channel](local_hidden)
+        )
+        if _can_use_fast_control_head(self, channel, False):
+            token = _sample_fast_control_head(self, lm_hidden)
+        else:
+            logits = self.lm_heads[channel](lm_hidden)
+            if channel != 0:
+                logits[:, self.config.audio_pad_code] = -torch.inf
+            token = torch.argmax(logits, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            lookup = (
+                feedback_lookup_tables[channel]
+                if feedback_lookup_tables is not None and channel < len(feedback_lookup_tables)
+                else None
+            )
+            if lookup is not None:
+                current_input = lookup.index_select(0, token.reshape(-1))
+            else:
+                current_input = self.speech_embedding_to_local_mlp(
+                    self.model.embedding_list[channel](token)
+                )
+    return next_tokens
+
+
+def _sample_frame_greedy_full_compiled(self, hidden_states: torch.Tensor, *, local_num_steps: int) -> torch.Tensor:
+    compiled = getattr(self, "_torchopt_compiled_greedy_frame", None)
+    if compiled is None:
+        compile_mode = getattr(self, "_torchopt_compile_mode", None)
+        if compile_mode == "default":
+            compile_mode = None
+        compiled = torch.compile(
+            types.MethodType(_sample_frame_greedy_full_impl, self),
+            dynamic=False,
+            mode=compile_mode,
+        )
+        self._torchopt_compiled_greedy_frame = compiled
+    return compiled(hidden_states, local_num_steps=local_num_steps)
+
+
+def _sample_frame_greedy_full_compile_cudagraph(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    input_ids: torch.Tensor,
+    realprocessor: list[LogitsProcessorList],
+    do_samples: list[bool],
+    fast_top_p: bool,
+    temperatures: list[float],
+    top_ps: list[float],
+    local_num_steps: int,
+) -> torch.Tensor:
+    if hidden_states.device.type == "cuda" and _can_stochastic_top_p_frame_graph(
+        self,
+        do_samples=do_samples,
+        fast_top_p=fast_top_p,
+        top_ps=top_ps,
+        local_num_steps=local_num_steps,
+    ):
+        return _sample_frame_stochastic_top_p_compile_cudagraph(
+            self,
+            hidden_states,
+            do_samples=do_samples,
+            temperatures=temperatures,
+            top_ps=top_ps,
+            local_num_steps=local_num_steps,
+        )
+    if hidden_states.device.type != "cuda" or any(do_samples):
+        return _sample_frame_fixed_full_cudagraph(
+            self,
+            hidden_states,
+            input_ids=input_ids,
+            realprocessor=realprocessor,
+            do_samples=do_samples,
+            fast_top_p=fast_top_p,
+            temperatures=temperatures,
+            top_ps=top_ps,
+            local_num_steps=local_num_steps,
+        )
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        int(local_num_steps),
+        "greedy-full-compile-cudagraph",
+        getattr(self, "_torchopt_compile_mode", None),
+        bool(getattr(self, "_torchopt_fast_control_head", False)),
+    )
+    entry = self._torchopt_frame_graphs.get(key)
+    if entry is None:
+        static_hidden = torch.empty_like(hidden_states)
+        warmup_stream = torch.cuda.Stream(device=hidden_states.device)
+        warmup_stream.wait_stream(torch.cuda.current_stream(hidden_states.device))
+        with torch.cuda.stream(warmup_stream):
+            for _ in range(3):
+                _sample_frame_greedy_full_compiled(
+                    self,
+                    static_hidden,
+                    local_num_steps=local_num_steps,
+                )
+        torch.cuda.current_stream(hidden_states.device).wait_stream(warmup_stream)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            static_tokens = _sample_frame_greedy_full_compiled(
+                self,
+                static_hidden,
+                local_num_steps=local_num_steps,
+            )
+        entry = (graph, static_hidden, static_tokens)
+        self._torchopt_frame_graphs[key] = entry
+
+    graph, static_hidden, static_tokens = entry
+    static_hidden.copy_(hidden_states)
+    graph.replay()
+    return static_tokens
+
+
+def _can_stochastic_top_p_frame_graph(
+    self,
+    *,
+    do_samples: list[bool],
+    fast_top_p: bool,
+    top_ps: list[float],
+    local_num_steps: int,
+) -> bool:
+    if not fast_top_p or int(getattr(self, "_torchopt_top_p_prefilter_size", 0) or 0) != 0:
+        return False
+    if getattr(self, "_torchopt_compile_mode", None) != "max-autotune-no-cudagraphs":
+        return False
+    if not any(do_samples[:local_num_steps]):
+        return False
+    gumbel_vocab_sizes = {
+        int(getattr(self.lm_heads[channel], "out_features"))
+        for channel in range(local_num_steps)
+        if do_samples[channel] and top_ps[channel] >= 1.0
+    }
+    if len(gumbel_vocab_sizes) > 1:
+        return False
+    return True
+
+
+def _sample_frame_stochastic_top_p_impl(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    gumbel_uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    temperatures: tuple[float, ...],
+    top_ps: tuple[float, ...],
+    local_num_steps: int,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    dtype = hidden_states.dtype
+    device = hidden_states.device
+    local_inputs = torch.zeros(
+        batch,
+        local_num_steps,
+        self.local_transformer_config.hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    for channel in range(local_num_steps):
+        local_inputs[:, channel, :] = current_input
+        local_outputs = _run_local_transformer_impl(self, local_inputs)
+        local_hidden = local_outputs[:, channel, :]
+        lm_hidden = self.layer_norm_before_lm_heads[channel](
+            self.local_to_speech_embedding_mlps[channel](local_hidden)
+        )
+        if _can_use_fast_control_head(self, channel, do_samples[channel]):
+            token = _sample_fast_control_head(self, lm_hidden)
+        else:
+            logits = self.lm_heads[channel](lm_hidden)
+            if channel != 0:
+                logits[:, self.config.audio_pad_code] = -torch.inf
+            if do_samples[channel]:
+                if top_ps[channel] >= 1.0:
+                    token = _sample_temperature_gumbel_with_uniform(
+                        logits,
+                        temperature=temperatures[channel],
+                        uniform=gumbel_uniforms[:, channel, :],
+                    )
+                else:
+                    token = None
+                    if bool(getattr(self, "_torchopt_triton_top_p", False)):
+                        token = _sample_top_p_only_with_uniform_triton(
+                            logits.contiguous(),
+                            temperature=temperatures[channel],
+                            top_p=top_ps[channel],
+                            uniform=uniforms[:, channel, :],
+                        )
+                    if token is None:
+                        token = _sample_top_p_only_with_uniform(
+                            logits,
+                            temperature=temperatures[channel],
+                            top_p=top_ps[channel],
+                            uniform=uniforms[:, channel, :],
+                        )
+            else:
+                token = torch.argmax(logits, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            current_input = self.speech_embedding_to_local_mlp(
+                self.model.embedding_list[channel](token)
+            )
+    return next_tokens
+
+
+def _sample_frame_stochastic_top_p_compiled(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    gumbel_uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    temperatures: tuple[float, ...],
+    top_ps: tuple[float, ...],
+    local_num_steps: int,
+) -> torch.Tensor:
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        tuple(uniforms.shape),
+        tuple(gumbel_uniforms.shape),
+        int(local_num_steps),
+        do_samples,
+        temperatures,
+        top_ps,
+        getattr(self, "_torchopt_compile_mode", None),
+        bool(getattr(self, "_torchopt_fast_control_head", False)),
+    )
+    compiled_cache = getattr(self, "_torchopt_compiled_stochastic_top_p_frames", None)
+    if compiled_cache is None:
+        compiled_cache = {}
+        self._torchopt_compiled_stochastic_top_p_frames = compiled_cache
+    compiled = compiled_cache.get(key)
+    if compiled is None:
+        compile_mode = getattr(self, "_torchopt_compile_mode", None)
+        if compile_mode == "default":
+            compile_mode = None
+        compiled = torch.compile(
+            types.MethodType(_sample_frame_stochastic_top_p_impl, self),
+            dynamic=False,
+            mode=compile_mode,
+        )
+        compiled_cache[key] = compiled
+    return compiled(
+        hidden_states,
+        uniforms,
+        gumbel_uniforms,
+        do_samples=do_samples,
+        temperatures=temperatures,
+        top_ps=top_ps,
+        local_num_steps=local_num_steps,
+    )
+
+
+def _sample_frame_stochastic_top_p_compile_cudagraph(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    do_samples: list[bool],
+    temperatures: list[float],
+    top_ps: list[float],
+    local_num_steps: int,
+) -> torch.Tensor:
+    do_samples_tuple = tuple(bool(value) for value in do_samples[:local_num_steps])
+    temperatures_tuple = tuple(float(value) for value in temperatures[:local_num_steps])
+    top_ps_tuple = tuple(float(value) for value in top_ps[:local_num_steps])
+    needs_gumbel = any(
+        do_samples_tuple[channel] and top_ps_tuple[channel] >= 1.0
+        for channel in range(local_num_steps)
+    )
+    gumbel_vocab_sizes = {
+        int(getattr(self.lm_heads[channel], "out_features"))
+        for channel in range(local_num_steps)
+        if do_samples_tuple[channel] and top_ps_tuple[channel] >= 1.0
+    }
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        int(local_num_steps),
+        do_samples_tuple,
+        temperatures_tuple,
+        top_ps_tuple,
+        "stochastic-top-p-full-frame-cudagraph",
+        getattr(self, "_torchopt_compile_mode", None),
+        bool(getattr(self, "_torchopt_fast_control_head", False)),
+    )
+    entry = self._torchopt_frame_graphs.get(key)
+    if entry is None:
+        static_hidden = torch.empty_like(hidden_states)
+        static_uniforms = torch.empty(
+            hidden_states.shape[0],
+            local_num_steps,
+            1,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        gumbel_vocab_size = next(iter(gumbel_vocab_sizes), 1)
+        static_gumbel_uniforms = torch.empty(
+            hidden_states.shape[0],
+            local_num_steps,
+            gumbel_vocab_size if needs_gumbel else 1,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        warmup_stream = torch.cuda.Stream(device=hidden_states.device)
+        warmup_stream.wait_stream(torch.cuda.current_stream(hidden_states.device))
+        with torch.cuda.stream(warmup_stream):
+            for _ in range(3):
+                static_uniforms.uniform_()
+                if needs_gumbel:
+                    static_gumbel_uniforms.uniform_()
+                _sample_frame_stochastic_top_p_compiled(
+                    self,
+                    static_hidden,
+                    static_uniforms,
+                    static_gumbel_uniforms,
+                    do_samples=do_samples_tuple,
+                    temperatures=temperatures_tuple,
+                    top_ps=top_ps_tuple,
+                    local_num_steps=local_num_steps,
+                )
+        torch.cuda.current_stream(hidden_states.device).wait_stream(warmup_stream)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            static_tokens = _sample_frame_stochastic_top_p_compiled(
+                self,
+                static_hidden,
+                static_uniforms,
+                static_gumbel_uniforms,
+                do_samples=do_samples_tuple,
+                temperatures=temperatures_tuple,
+                top_ps=top_ps_tuple,
+                local_num_steps=local_num_steps,
+            )
+        entry = (graph, static_hidden, static_uniforms, static_gumbel_uniforms, static_tokens)
+        self._torchopt_frame_graphs[key] = entry
+
+    graph, static_hidden, static_uniforms, static_gumbel_uniforms, static_tokens = entry
+    static_hidden.copy_(hidden_states)
+    static_uniforms.uniform_()
+    if needs_gumbel:
+        static_gumbel_uniforms.uniform_()
+    graph.replay()
+    return static_tokens
+
+
+def _sample_frame_prefix_top_p_impl(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    gumbel_uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    temperatures: tuple[float, ...],
+    top_ps: tuple[float, ...],
+    local_num_steps: int,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    dtype = hidden_states.dtype
+    device = hidden_states.device
+    local_inputs = torch.zeros(
+        batch,
+        local_num_steps,
+        self.local_transformer_config.hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    for channel in range(local_num_steps):
+        local_inputs[:, channel, :] = current_input
+        local_outputs = _run_local_transformer_impl(self, local_inputs[:, : channel + 1, :])
+        local_hidden = local_outputs[:, -1, :]
+        lm_hidden = self.layer_norm_before_lm_heads[channel](
+            self.local_to_speech_embedding_mlps[channel](local_hidden)
+        )
+        logits = self.lm_heads[channel](lm_hidden)
+        if channel != 0:
+            logits[:, self.config.audio_pad_code] = -torch.inf
+        if do_samples[channel]:
+            if top_ps[channel] >= 1.0:
+                token = _sample_temperature_gumbel_with_uniform(
+                    logits,
+                    temperature=temperatures[channel],
+                    uniform=gumbel_uniforms[:, channel, :],
+                )
+            else:
+                token = _sample_top_p_only_with_uniform(
+                    logits,
+                    temperature=temperatures[channel],
+                    top_p=top_ps[channel],
+                    uniform=uniforms[:, channel, :],
+                )
+        else:
+            token = torch.argmax(logits, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            current_input = self.speech_embedding_to_local_mlp(
+                self.model.embedding_list[channel](token)
+            )
+    return next_tokens
+
+
+def _sample_frame_prefix_top_p_compiled(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    gumbel_uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    temperatures: tuple[float, ...],
+    top_ps: tuple[float, ...],
+    local_num_steps: int,
+) -> torch.Tensor:
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        tuple(uniforms.shape),
+        tuple(gumbel_uniforms.shape),
+        int(local_num_steps),
+        do_samples,
+        temperatures,
+        top_ps,
+        getattr(self, "_torchopt_compile_mode", None),
+        "prefix-full-frame",
+    )
+    compiled_cache = getattr(self, "_torchopt_compiled_prefix_top_p_frames", None)
+    if compiled_cache is None:
+        compiled_cache = {}
+        self._torchopt_compiled_prefix_top_p_frames = compiled_cache
+    compiled = compiled_cache.get(key)
+    if compiled is None:
+        compile_mode = getattr(self, "_torchopt_compile_mode", None)
+        if compile_mode == "default":
+            compile_mode = None
+        compiled = torch.compile(
+            types.MethodType(_sample_frame_prefix_top_p_impl, self),
+            dynamic=False,
+            mode=compile_mode,
+        )
+        compiled_cache[key] = compiled
+    return compiled(
+        hidden_states,
+        uniforms,
+        gumbel_uniforms,
+        do_samples=do_samples,
+        temperatures=temperatures,
+        top_ps=top_ps,
+        local_num_steps=local_num_steps,
+    )
+
+
+def _sample_frame_prefix_top_p_compile_cudagraph(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    do_samples: list[bool],
+    temperatures: list[float],
+    top_ps: list[float],
+    local_num_steps: int,
+) -> torch.Tensor:
+    if hidden_states.device.type != "cuda":
+        return _sample_frame_prefix_top_p_impl(
+            self,
+            hidden_states,
+            torch.rand(hidden_states.shape[0], local_num_steps, 1, device=hidden_states.device, dtype=hidden_states.dtype),
+            torch.rand(hidden_states.shape[0], local_num_steps, 1, device=hidden_states.device, dtype=hidden_states.dtype),
+            do_samples=tuple(bool(value) for value in do_samples[:local_num_steps]),
+            temperatures=tuple(float(value) for value in temperatures[:local_num_steps]),
+            top_ps=tuple(float(value) for value in top_ps[:local_num_steps]),
+            local_num_steps=local_num_steps,
+        )
+    do_samples_tuple = tuple(bool(value) for value in do_samples[:local_num_steps])
+    temperatures_tuple = tuple(float(value) for value in temperatures[:local_num_steps])
+    top_ps_tuple = tuple(float(value) for value in top_ps[:local_num_steps])
+    needs_gumbel = any(
+        do_samples_tuple[channel] and top_ps_tuple[channel] >= 1.0
+        for channel in range(local_num_steps)
+    )
+    gumbel_vocab_sizes = {
+        int(getattr(self.lm_heads[channel], "out_features"))
+        for channel in range(local_num_steps)
+        if do_samples_tuple[channel] and top_ps_tuple[channel] >= 1.0
+    }
+    if len(gumbel_vocab_sizes) > 1:
+        return _sample_frame_fixed_full_impl(
+            self,
+            hidden_states,
+            input_ids=torch.empty(
+                hidden_states.shape[0],
+                1,
+                self.channels,
+                device=hidden_states.device,
+                dtype=torch.long,
+            ),
+            realprocessor=[LogitsProcessorList() for _ in range(self.channels)],
+            do_samples=list(do_samples_tuple),
+            fast_top_p=True,
+            temperatures=list(temperatures_tuple),
+            top_ps=list(top_ps_tuple),
+            local_num_steps=local_num_steps,
+        )
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        int(local_num_steps),
+        do_samples_tuple,
+        temperatures_tuple,
+        top_ps_tuple,
+        "prefix-full-compile-cudagraph",
+        getattr(self, "_torchopt_compile_mode", None),
+    )
+    entry = self._torchopt_frame_graphs.get(key)
+    if entry is None:
+        static_hidden = torch.empty_like(hidden_states)
+        static_uniforms = torch.empty(
+            hidden_states.shape[0],
+            local_num_steps,
+            1,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        gumbel_vocab_size = next(iter(gumbel_vocab_sizes), 1)
+        static_gumbel_uniforms = torch.empty(
+            hidden_states.shape[0],
+            local_num_steps,
+            gumbel_vocab_size if needs_gumbel else 1,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        warmup_stream = torch.cuda.Stream(device=hidden_states.device)
+        warmup_stream.wait_stream(torch.cuda.current_stream(hidden_states.device))
+        with torch.cuda.stream(warmup_stream):
+            for _ in range(3):
+                static_uniforms.uniform_()
+                if needs_gumbel:
+                    static_gumbel_uniforms.uniform_()
+                _sample_frame_prefix_top_p_compiled(
+                    self,
+                    static_hidden,
+                    static_uniforms,
+                    static_gumbel_uniforms,
+                    do_samples=do_samples_tuple,
+                    temperatures=temperatures_tuple,
+                    top_ps=top_ps_tuple,
+                    local_num_steps=local_num_steps,
+                )
+        torch.cuda.current_stream(hidden_states.device).wait_stream(warmup_stream)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            static_tokens = _sample_frame_prefix_top_p_compiled(
+                self,
+                static_hidden,
+                static_uniforms,
+                static_gumbel_uniforms,
+                do_samples=do_samples_tuple,
+                temperatures=temperatures_tuple,
+                top_ps=top_ps_tuple,
+                local_num_steps=local_num_steps,
+            )
+        entry = (graph, static_hidden, static_uniforms, static_gumbel_uniforms, static_tokens)
+        self._torchopt_frame_graphs[key] = entry
+
+    graph, static_hidden, static_uniforms, static_gumbel_uniforms, static_tokens = entry
+    static_hidden.copy_(hidden_states)
+    static_uniforms.uniform_()
+    if needs_gumbel:
+        static_gumbel_uniforms.uniform_()
+    graph.replay()
+    return static_tokens
+
+
+def _sample_frame_static_local_cache_impl(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    input_ids: torch.Tensor,
+    realprocessor: list[LogitsProcessorList],
+    do_samples: list[bool],
+    local_num_steps: int,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    dtype = hidden_states.dtype
+    device = hidden_states.device
+    key_caches: list[torch.Tensor] = []
+    value_caches: list[torch.Tensor] = []
+    for decoder_layer in self.local_transformer.layers[
+        : self.local_transformer.config.num_hidden_layers
+    ]:
+        attn = decoder_layer.self_attn
+        num_key_value_heads = attn.k_proj.out_features // attn.head_dim
+        key_cache = torch.empty(
+            batch,
+            num_key_value_heads,
+            int(local_num_steps),
+            attn.head_dim,
+            device=device,
+            dtype=dtype,
+        )
+        key_caches.append(key_cache)
+        value_caches.append(torch.empty_like(key_cache))
+    return _sample_frame_static_local_cache_with_caches(
+        self,
+        hidden_states,
+        input_ids=input_ids,
+        realprocessor=realprocessor,
+        do_samples=do_samples,
+        local_num_steps=local_num_steps,
+        key_caches=key_caches,
+        value_caches=value_caches,
+    )
+
+
+def _sample_frame_static_local_cache_with_caches(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    input_ids: torch.Tensor,
+    realprocessor: list[LogitsProcessorList],
+    do_samples: list[bool],
+    local_num_steps: int,
+    key_caches: list[torch.Tensor],
+    value_caches: list[torch.Tensor],
+    feedback_lookup_tables: list[torch.Tensor | None] | None = None,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    device = hidden_states.device
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    for channel in range(local_num_steps):
+        local_hidden = _run_local_transformer_static_cache(
+            self,
+            current_input,
+            key_caches=key_caches,
+            value_caches=value_caches,
+            cache_index=channel,
+        )
+        lm_hidden = self.layer_norm_before_lm_heads[channel](
+            self.local_to_speech_embedding_mlps[channel](local_hidden)
+        )
+        logits = self.lm_heads[channel](lm_hidden)
+        if channel != 0:
+            logits[:, self.config.audio_pad_code] = -torch.inf
+        scores = realprocessor[channel](input_ids[..., channel], logits)
+        if do_samples[channel]:
+            token = torch.multinomial(F.softmax(scores, dim=-1), num_samples=1).squeeze(1)
+        else:
+            token = torch.argmax(scores, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            lookup = (
+                feedback_lookup_tables[channel]
+                if feedback_lookup_tables is not None and channel < len(feedback_lookup_tables)
+                else None
+            )
+            if lookup is not None:
+                current_input = lookup.index_select(0, token.reshape(-1))
+            else:
+                current_input = self.speech_embedding_to_local_mlp(
+                    self.model.embedding_list[channel](token)
+                )
+    return next_tokens
+
+
+def _sample_frame_static_local_cache_fast_with_caches(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    gumbel_uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    temperatures: tuple[float, ...],
+    top_ps: tuple[float, ...],
+    local_num_steps: int,
+    key_caches: list[torch.Tensor],
+    value_caches: list[torch.Tensor],
+    feedback_lookup_tables: list[torch.Tensor | None] | None = None,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    device = hidden_states.device
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    for channel in range(local_num_steps):
+        local_hidden = _run_local_transformer_static_cache_impl(
+            self,
+            current_input,
+            key_caches=key_caches,
+            value_caches=value_caches,
+            cache_index=channel,
+        )
+        lm_hidden = self.layer_norm_before_lm_heads[channel](
+            self.local_to_speech_embedding_mlps[channel](local_hidden)
+        )
+        if do_samples[channel]:
+            if top_ps[channel] >= 1.0:
+                logits = self.lm_heads[channel](lm_hidden)
+                if channel != 0:
+                    logits[:, self.config.audio_pad_code] = -torch.inf
+                token = _sample_temperature_gumbel_with_uniform(
+                    logits,
+                    temperature=temperatures[channel],
+                    uniform=gumbel_uniforms[:, channel, :],
+                )
+            else:
+                token = _sample_fused_lm_head_top_p_triton(
+                    self,
+                    lm_hidden,
+                    channel=channel,
+                    temperature=temperatures[channel],
+                    top_p=top_ps[channel],
+                    uniform=uniforms[:, channel, :],
+                )
+                logits = None
+                if bool(getattr(self, "_torchopt_triton_top_p", False)):
+                    if token is None:
+                        logits = self.lm_heads[channel](lm_hidden)
+                        if channel != 0:
+                            logits[:, self.config.audio_pad_code] = -torch.inf
+                        token = _sample_top_p_only_with_uniform_triton(
+                            logits.contiguous(),
+                            temperature=temperatures[channel],
+                            top_p=top_ps[channel],
+                            uniform=uniforms[:, channel, :],
+                        )
+                if token is None:
+                    if logits is None:
+                        logits = self.lm_heads[channel](lm_hidden)
+                        if channel != 0:
+                            logits[:, self.config.audio_pad_code] = -torch.inf
+                    token = _sample_top_p_only_with_uniform(
+                        logits,
+                        temperature=temperatures[channel],
+                        top_p=top_ps[channel],
+                        uniform=uniforms[:, channel, :],
+                    )
+        else:
+            logits = self.lm_heads[channel](lm_hidden)
+            if channel != 0:
+                logits[:, self.config.audio_pad_code] = -torch.inf
+            token = torch.argmax(logits, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            lookup = (
+                feedback_lookup_tables[channel]
+                if feedback_lookup_tables is not None and channel < len(feedback_lookup_tables)
+                else None
+            )
+            if lookup is not None:
+                current_input = lookup.index_select(0, token.reshape(-1))
+            else:
+                current_input = self.speech_embedding_to_local_mlp(
+                    self.model.embedding_list[channel](token)
+                )
+    return next_tokens
+
+
+def _sample_frame_static_local_cache_triton_attention_fast_with_caches(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    gumbel_uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    temperatures: tuple[float, ...],
+    top_ps: tuple[float, ...],
+    local_num_steps: int,
+    key_caches: list[torch.Tensor],
+    value_caches: list[torch.Tensor],
+    feedback_lookup_tables: list[torch.Tensor | None] | None = None,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    device = hidden_states.device
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    for channel in range(local_num_steps):
+        local_hidden = _run_local_transformer_static_cache_triton_attention_impl(
+            self,
+            current_input,
+            key_caches=key_caches,
+            value_caches=value_caches,
+            cache_index=channel,
+        )
+        lm_hidden = self.layer_norm_before_lm_heads[channel](
+            self.local_to_speech_embedding_mlps[channel](local_hidden)
+        )
+        logits = self.lm_heads[channel](lm_hidden)
+        if channel != 0:
+            logits[:, self.config.audio_pad_code] = -torch.inf
+        if do_samples[channel]:
+            if top_ps[channel] >= 1.0:
+                token = _sample_temperature_gumbel_with_uniform(
+                    logits,
+                    temperature=temperatures[channel],
+                    uniform=gumbel_uniforms[:, channel, :],
+                )
+            else:
+                token = None
+                if bool(getattr(self, "_torchopt_triton_top_p", False)):
+                    token = _sample_top_p_only_with_uniform_triton(
+                        logits.contiguous(),
+                        temperature=temperatures[channel],
+                        top_p=top_ps[channel],
+                        uniform=uniforms[:, channel, :],
+                    )
+                if token is None:
+                    token = _sample_top_p_only_with_uniform(
+                        logits,
+                        temperature=temperatures[channel],
+                        top_p=top_ps[channel],
+                        uniform=uniforms[:, channel, :],
+                    )
+        else:
+            token = torch.argmax(logits, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            lookup = (
+                feedback_lookup_tables[channel]
+                if feedback_lookup_tables is not None and channel < len(feedback_lookup_tables)
+                else None
+            )
+            if lookup is not None:
+                current_input = lookup.index_select(0, token.reshape(-1))
+            else:
+                current_input = self.speech_embedding_to_local_mlp(
+                    self.model.embedding_list[channel](token)
+                )
+    return next_tokens
+
+
+def _sample_frame_static_local_cache_triton_attention_fast_with_caches_dynamic_params(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    sampling_params: torch.Tensor,
+    local_num_steps: int,
+    key_caches: list[torch.Tensor],
+    value_caches: list[torch.Tensor],
+    feedback_lookup_tables: list[torch.Tensor | None] | None = None,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    device = hidden_states.device
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    for channel in range(local_num_steps):
+        local_hidden = _run_local_transformer_static_cache_triton_attention_impl(
+            self,
+            current_input,
+            key_caches=key_caches,
+            value_caches=value_caches,
+            cache_index=channel,
+        )
+        lm_hidden = self.layer_norm_before_lm_heads[channel](
+            self.local_to_speech_embedding_mlps[channel](local_hidden)
+        )
+        logits = self.lm_heads[channel](lm_hidden)
+        if channel != 0:
+            logits[:, self.config.audio_pad_code] = -torch.inf
+        if do_samples[channel]:
+            token = _sample_top_p_only_with_uniform_dynamic(
+                logits,
+                temperature=sampling_params[channel, 0],
+                top_p=sampling_params[channel, 1],
+                uniform=uniforms[:, channel, :],
+            )
+        else:
+            token = torch.argmax(logits, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            lookup = (
+                feedback_lookup_tables[channel]
+                if feedback_lookup_tables is not None and channel < len(feedback_lookup_tables)
+                else None
+            )
+            if lookup is not None:
+                current_input = lookup.index_select(0, token.reshape(-1))
+            else:
+                current_input = self.speech_embedding_to_local_mlp(
+                    self.model.embedding_list[channel](token)
+                )
+    return next_tokens
+
+
+def _sample_frame_static_local_cache_triton_attention_dynamic_params_impl(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    sampling_params: torch.Tensor,
+    local_num_steps: int,
+    feedback_lookup_tables: list[torch.Tensor | None] | None = None,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    key_caches: list[torch.Tensor] = []
+    value_caches: list[torch.Tensor] = []
+    for decoder_layer in self.local_transformer.layers[
+        : self.local_transformer.config.num_hidden_layers
+    ]:
+        attn = decoder_layer.self_attn
+        num_key_value_heads = attn.k_proj.out_features // attn.head_dim
+        key_cache = torch.empty(
+            batch,
+            num_key_value_heads,
+            int(local_num_steps),
+            attn.head_dim,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        key_caches.append(key_cache)
+        value_caches.append(torch.empty_like(key_cache))
+    return _sample_frame_static_local_cache_triton_attention_fast_with_caches_dynamic_params(
+        self,
+        hidden_states,
+        uniforms,
+        do_samples=do_samples,
+        sampling_params=sampling_params,
+        local_num_steps=local_num_steps,
+        key_caches=key_caches,
+        value_caches=value_caches,
+        feedback_lookup_tables=feedback_lookup_tables,
+    )
+
+
+def _sample_frame_static_local_cache_triton_attention_impl(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    gumbel_uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    temperatures: tuple[float, ...],
+    top_ps: tuple[float, ...],
+    local_num_steps: int,
+    feedback_lookup_tables: list[torch.Tensor | None] | None = None,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    key_caches: list[torch.Tensor] = []
+    value_caches: list[torch.Tensor] = []
+    for decoder_layer in self.local_transformer.layers[
+        : self.local_transformer.config.num_hidden_layers
+    ]:
+        attn = decoder_layer.self_attn
+        num_key_value_heads = attn.k_proj.out_features // attn.head_dim
+        key_cache = torch.empty(
+            batch,
+            num_key_value_heads,
+            int(local_num_steps),
+            attn.head_dim,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        key_caches.append(key_cache)
+        value_caches.append(torch.empty_like(key_cache))
+    return _sample_frame_static_local_cache_triton_attention_fast_with_caches(
+        self,
+        hidden_states,
+        uniforms,
+        gumbel_uniforms,
+        do_samples=do_samples,
+        temperatures=temperatures,
+        top_ps=top_ps,
+        local_num_steps=local_num_steps,
+        key_caches=key_caches,
+        value_caches=value_caches,
+        feedback_lookup_tables=feedback_lookup_tables,
+    )
+
+
+def _sample_frame_static_local_cache_triton_attention_compiled(
+    self,
+    hidden_states: torch.Tensor,
+    uniforms: torch.Tensor,
+    gumbel_uniforms: torch.Tensor,
+    *,
+    do_samples: tuple[bool, ...],
+    temperatures: tuple[float, ...],
+    top_ps: tuple[float, ...],
+    local_num_steps: int,
+    sampling_params: torch.Tensor | None = None,
+) -> torch.Tensor:
+    feedback_lookup_tables = _get_feedback_lookup_tables(
+        self,
+        local_num_steps=local_num_steps,
+        device=hidden_states.device,
+        dtype=hidden_states.dtype,
+    )
+    dynamic_sampling_params = sampling_params is not None
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        tuple(uniforms.shape),
+        tuple(gumbel_uniforms.shape),
+        int(local_num_steps),
+        do_samples,
+        None if dynamic_sampling_params else temperatures,
+        None if dynamic_sampling_params else top_ps,
+        dynamic_sampling_params,
+        getattr(self, "_torchopt_compile_mode", None),
+        False if dynamic_sampling_params else bool(getattr(self, "_torchopt_triton_top_p", False)),
+        bool(getattr(self, "_torchopt_triton_fused_lm_head", False)),
+        bool(getattr(self, "_torchopt_triton_qkv_cache", False)),
+        bool(feedback_lookup_tables is not None),
+        bool(getattr(self, "_torchopt_local_compile_fullgraph", False)),
+        "static-local-cache-triton-attention",
+    )
+    compiled_cache = getattr(self, "_torchopt_compiled_static_local_triton_attention", None)
+    if compiled_cache is None:
+        compiled_cache = {}
+        self._torchopt_compiled_static_local_triton_attention = compiled_cache
+    compiled = compiled_cache.get(key)
+    if compiled is None:
+        compile_mode = getattr(self, "_torchopt_compile_mode", None)
+        if compile_mode == "default":
+            compile_mode = None
+        compile_target = (
+            _sample_frame_static_local_cache_triton_attention_dynamic_params_impl
+            if dynamic_sampling_params
+            else _sample_frame_static_local_cache_triton_attention_impl
+        )
+        compiled = torch.compile(
+            types.MethodType(compile_target, self),
+            dynamic=False,
+            mode=compile_mode,
+            fullgraph=bool(getattr(self, "_torchopt_local_compile_fullgraph", False)),
+        )
+        compiled_cache[key] = compiled
+    if dynamic_sampling_params:
+        return compiled(
+            hidden_states,
+            uniforms,
+            do_samples=do_samples,
+            sampling_params=sampling_params,
+            local_num_steps=local_num_steps,
+            feedback_lookup_tables=feedback_lookup_tables,
+        )
+    return compiled(
+        hidden_states,
+        uniforms,
+        gumbel_uniforms,
+        do_samples=do_samples,
+        temperatures=temperatures,
+        top_ps=top_ps,
+        local_num_steps=local_num_steps,
+        feedback_lookup_tables=feedback_lookup_tables,
+    )
+
+
+def _sample_frame_static_local_cache_triton_attention_compile_cudagraph(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    do_samples: list[bool],
+    temperatures: list[float],
+    top_ps: list[float],
+    local_num_steps: int,
+) -> torch.Tensor:
+    if hidden_states.device.type != "cuda":
+        return _sample_frame_static_local_cache_triton_attention_impl(
+            self,
+            hidden_states,
+            torch.rand(hidden_states.shape[0], local_num_steps, 1, device=hidden_states.device, dtype=hidden_states.dtype),
+            torch.rand(hidden_states.shape[0], local_num_steps, 1, device=hidden_states.device, dtype=hidden_states.dtype),
+            do_samples=tuple(bool(value) for value in do_samples[:local_num_steps]),
+            temperatures=tuple(float(value) for value in temperatures[:local_num_steps]),
+            top_ps=tuple(float(value) for value in top_ps[:local_num_steps]),
+            local_num_steps=local_num_steps,
+        )
+    do_samples_tuple = tuple(bool(value) for value in do_samples[:local_num_steps])
+    temperatures_tuple = tuple(float(value) for value in temperatures[:local_num_steps])
+    top_ps_tuple = tuple(float(value) for value in top_ps[:local_num_steps])
+    needs_gumbel = False
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        int(local_num_steps),
+        do_samples_tuple,
+        "dynamic-sampling-params",
+        "static-local-cache-triton-attention-cudagraph",
+        getattr(self, "_torchopt_compile_mode", None),
+        False,
+        bool(getattr(self, "_torchopt_triton_fused_lm_head", False)),
+        bool(getattr(self, "_torchopt_triton_qkv_cache", False)),
+        bool(getattr(self, "_torchopt_feedback_lookup", False)),
+        bool(getattr(self, "_torchopt_local_compile_fullgraph", False)),
+    )
+    entry = self._torchopt_frame_graphs.get(key)
+    if entry is None:
+        static_hidden = torch.empty_like(hidden_states)
+        static_uniforms = torch.empty(
+            hidden_states.shape[0],
+            local_num_steps,
+            1,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        static_gumbel_uniforms = torch.empty(
+            hidden_states.shape[0],
+            local_num_steps,
+            1,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        static_sampling_params = torch.empty(
+            local_num_steps,
+            2,
+            device=hidden_states.device,
+            dtype=torch.float32,
+        )
+        static_sampling_params[:, 0].copy_(
+            torch.tensor(temperatures_tuple, device=hidden_states.device, dtype=torch.float32)
+        )
+        static_sampling_params[:, 1].copy_(
+            torch.tensor(top_ps_tuple, device=hidden_states.device, dtype=torch.float32)
+        )
+        warmup_stream = torch.cuda.Stream(device=hidden_states.device)
+        warmup_stream.wait_stream(torch.cuda.current_stream(hidden_states.device))
+        with torch.cuda.stream(warmup_stream):
+            for _ in range(3):
+                static_uniforms.uniform_()
+                _sample_frame_static_local_cache_triton_attention_compiled(
+                    self,
+                    static_hidden,
+                    static_uniforms,
+                    static_gumbel_uniforms,
+                    do_samples=do_samples_tuple,
+                    temperatures=temperatures_tuple,
+                    top_ps=top_ps_tuple,
+                    local_num_steps=local_num_steps,
+                    sampling_params=static_sampling_params,
+                )
+        torch.cuda.current_stream(hidden_states.device).wait_stream(warmup_stream)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            static_tokens = _sample_frame_static_local_cache_triton_attention_compiled(
+                self,
+                static_hidden,
+                static_uniforms,
+                static_gumbel_uniforms,
+                do_samples=do_samples_tuple,
+                temperatures=temperatures_tuple,
+                top_ps=top_ps_tuple,
+                local_num_steps=local_num_steps,
+                sampling_params=static_sampling_params,
+            )
+        entry = (
+            graph,
+            static_hidden,
+            static_uniforms,
+            static_gumbel_uniforms,
+            static_sampling_params,
+            static_tokens,
+        )
+        self._torchopt_frame_graphs[key] = entry
+
+    graph, static_hidden, static_uniforms, static_gumbel_uniforms, static_sampling_params, static_tokens = entry
+    static_hidden.copy_(hidden_states)
+    static_sampling_params[:, 0].copy_(
+        torch.tensor(temperatures_tuple, device=hidden_states.device, dtype=torch.float32)
+    )
+    static_sampling_params[:, 1].copy_(
+        torch.tensor(top_ps_tuple, device=hidden_states.device, dtype=torch.float32)
+    )
+    if any(do_samples_tuple):
+        static_uniforms.uniform_()
+    graph.replay()
+    return static_tokens
+
+
+def _sample_frame_static_local_cache_cudagraph(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    input_ids: torch.Tensor,
+    realprocessor: list[LogitsProcessorList],
+    do_samples: list[bool],
+    fast_top_p: bool = False,
+    temperatures: list[float] | None = None,
+    top_ps: list[float] | None = None,
+    local_num_steps: int,
+) -> torch.Tensor:
+    if hidden_states.device.type != "cuda":
+        return _sample_frame_static_local_cache_impl(
+            self,
+            hidden_states,
+            input_ids=input_ids,
+            realprocessor=realprocessor,
+            do_samples=do_samples,
+            local_num_steps=local_num_steps,
+        )
+    if any(do_samples) and (
+        not fast_top_p
+        or temperatures is None
+        or top_ps is None
+        or int(getattr(self, "_torchopt_top_p_prefilter_size", 0) or 0) != 0
+    ):
+        return _sample_frame_static_local_cache_impl(
+            self,
+            hidden_states,
+            input_ids=input_ids,
+            realprocessor=realprocessor,
+            do_samples=do_samples,
+            local_num_steps=local_num_steps,
+        )
+    do_samples_tuple = tuple(bool(value) for value in do_samples[:local_num_steps])
+    temperatures_tuple = tuple(float(value) for value in (temperatures or [1.0] * local_num_steps)[:local_num_steps])
+    top_ps_tuple = tuple(float(value) for value in (top_ps or [1.0] * local_num_steps)[:local_num_steps])
+    needs_gumbel = any(
+        do_samples_tuple[channel] and top_ps_tuple[channel] >= 1.0
+        for channel in range(local_num_steps)
+    )
+    gumbel_vocab_sizes = {
+        int(getattr(self.lm_heads[channel], "out_features"))
+        for channel in range(local_num_steps)
+        if do_samples_tuple[channel] and top_ps_tuple[channel] >= 1.0
+    }
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        int(local_num_steps),
+        "static-local-cache-cudagraph",
+        do_samples_tuple,
+        temperatures_tuple,
+        top_ps_tuple,
+        getattr(self, "_torchopt_triton_top_p", False),
+    )
+    entry = self._torchopt_frame_graphs.get(key)
+    if entry is None:
+        batch = hidden_states.shape[0]
+        key_caches: list[torch.Tensor] = []
+        value_caches: list[torch.Tensor] = []
+        for decoder_layer in self.local_transformer.layers[
+            : self.local_transformer.config.num_hidden_layers
+        ]:
+            attn = decoder_layer.self_attn
+            num_key_value_heads = attn.k_proj.out_features // attn.head_dim
+            key_cache = torch.empty(
+                batch,
+                num_key_value_heads,
+                int(local_num_steps),
+                attn.head_dim,
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )
+            key_caches.append(key_cache)
+            value_caches.append(torch.empty_like(key_cache))
+        static_hidden = torch.zeros_like(hidden_states)
+        static_uniforms = torch.empty(
+            hidden_states.shape[0],
+            local_num_steps,
+            1,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        gumbel_vocab_size = next(iter(gumbel_vocab_sizes), 1)
+        static_gumbel_uniforms = torch.empty(
+            hidden_states.shape[0],
+            local_num_steps,
+            gumbel_vocab_size if needs_gumbel else 1,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        warmup_stream = torch.cuda.Stream(device=hidden_states.device)
+        warmup_stream.wait_stream(torch.cuda.current_stream(hidden_states.device))
+        with torch.cuda.stream(warmup_stream):
+            for _ in range(3):
+                if any(do_samples_tuple):
+                    static_uniforms.uniform_()
+                    if needs_gumbel:
+                        static_gumbel_uniforms.uniform_()
+                    _sample_frame_static_local_cache_fast_with_caches(
+                        self,
+                        static_hidden,
+                        static_uniforms,
+                        static_gumbel_uniforms,
+                        do_samples=do_samples_tuple,
+                        temperatures=temperatures_tuple,
+                        top_ps=top_ps_tuple,
+                        local_num_steps=local_num_steps,
+                        key_caches=key_caches,
+                        value_caches=value_caches,
+                    )
+                else:
+                    _sample_frame_static_local_cache_with_caches(
+                        self,
+                        static_hidden,
+                        input_ids=input_ids,
+                        realprocessor=realprocessor,
+                        do_samples=do_samples,
+                        local_num_steps=local_num_steps,
+                        key_caches=key_caches,
+                        value_caches=value_caches,
+                    )
+        torch.cuda.current_stream(hidden_states.device).wait_stream(warmup_stream)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            if any(do_samples_tuple):
+                static_tokens = _sample_frame_static_local_cache_fast_with_caches(
+                    self,
+                    static_hidden,
+                    static_uniforms,
+                    static_gumbel_uniforms,
+                    do_samples=do_samples_tuple,
+                    temperatures=temperatures_tuple,
+                    top_ps=top_ps_tuple,
+                    local_num_steps=local_num_steps,
+                    key_caches=key_caches,
+                    value_caches=value_caches,
+                )
+            else:
+                static_tokens = _sample_frame_static_local_cache_with_caches(
+                    self,
+                    static_hidden,
+                    input_ids=input_ids,
+                    realprocessor=realprocessor,
+                    do_samples=do_samples,
+                    local_num_steps=local_num_steps,
+                    key_caches=key_caches,
+                    value_caches=value_caches,
+                )
+        entry = (graph, static_hidden, static_uniforms, static_gumbel_uniforms, static_tokens, needs_gumbel)
+        self._torchopt_frame_graphs[key] = entry
+
+    graph, static_hidden, static_uniforms, static_gumbel_uniforms, static_tokens, needs_gumbel = entry
+    static_hidden.copy_(hidden_states)
+    if any(do_samples_tuple):
+        static_uniforms.uniform_()
+        if needs_gumbel:
+            static_gumbel_uniforms.uniform_()
+    graph.replay()
+    return static_tokens
+
+
+def _can_frame_graph(hidden_states: torch.Tensor, do_samples: list[bool]) -> bool:
+    return hidden_states.device.type == "cuda" and not any(do_samples)
+
+
+def _sample_channel_logits_impl(
+    self,
+    local_inputs: torch.Tensor,
+    *,
+    channel: int,
+) -> torch.Tensor:
+    local_outputs = _run_local_transformer_fixed(self, local_inputs)
+    local_hidden = local_outputs[:, channel, :]
+    lm_hidden = self.layer_norm_before_lm_heads[channel](
+        self.local_to_speech_embedding_mlps[channel](local_hidden)
+    )
+    logits = self.lm_heads[channel](lm_hidden)
+    if channel != 0:
+        logits[:, self.config.audio_pad_code] = -torch.inf
+    return logits
+
+
+def _sample_channel_logits_cudagraph(
+    self,
+    local_inputs: torch.Tensor,
+    *,
+    channel: int,
+) -> torch.Tensor:
+    key = (
+        local_inputs.device.index,
+        local_inputs.dtype,
+        tuple(local_inputs.shape),
+        int(channel),
+        getattr(self, "_torchopt_mode", "fixed-full-cudagraph"),
+        getattr(self, "_torchopt_compile_mode", None),
+    )
+    entry = self._torchopt_channel_graphs.get(key)
+    if entry is None:
+        static_local_inputs = torch.zeros_like(local_inputs)
+        warmup_stream = torch.cuda.Stream(device=local_inputs.device)
+        warmup_stream.wait_stream(torch.cuda.current_stream(local_inputs.device))
+        with torch.cuda.stream(warmup_stream):
+            for _ in range(3):
+                _sample_channel_logits_impl(self, static_local_inputs, channel=channel)
+        torch.cuda.current_stream(local_inputs.device).wait_stream(warmup_stream)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            static_logits = _sample_channel_logits_impl(self, static_local_inputs, channel=channel)
+        entry = (graph, static_local_inputs, static_logits)
+        self._torchopt_channel_graphs[key] = entry
+
+    graph, static_local_inputs, static_logits = entry
+    static_local_inputs.copy_(local_inputs)
+    graph.replay()
+    return static_logits
+
+
+def _sample_frame_hybrid_cudagraph(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    input_ids: torch.Tensor,
+    realprocessor: list[LogitsProcessorList],
+    do_samples: list[bool],
+    fast_top_p: bool,
+    temperatures: list[float],
+    top_ps: list[float],
+    local_num_steps: int,
+) -> torch.Tensor:
+    batch = hidden_states.shape[0]
+    dtype = hidden_states.dtype
+    device = hidden_states.device
+    local_inputs = torch.zeros(
+        batch,
+        local_num_steps,
+        self.local_transformer_config.hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    next_tokens = torch.zeros(
+        batch,
+        self.channels,
+        device=device,
+        dtype=torch.long,
+    )
+
+    current_input = self.speech_embedding_to_local_mlp(hidden_states)
+    top_p_uniforms = None
+    if (
+        fast_top_p
+        and bool(getattr(self, "_torchopt_top_p_sampler_cudagraph", False))
+        and any(do_samples[:local_num_steps])
+        and int(getattr(self, "_torchopt_top_p_prefilter_size", 0) or 0) == 0
+        and any(top_ps[channel] < 1.0 for channel in range(local_num_steps) if do_samples[channel])
+    ):
+        top_p_uniforms = torch.rand(
+            batch,
+            local_num_steps,
+            1,
+            device=device,
+            dtype=dtype,
+        )
+    for channel in range(local_num_steps):
+        local_inputs[:, channel, :] = current_input
+        logits = _sample_channel_logits_cudagraph(self, local_inputs, channel=channel)
+        if do_samples[channel]:
+            if fast_top_p:
+                uniform = None
+                if top_p_uniforms is not None and top_ps[channel] < 1.0:
+                    uniform = top_p_uniforms[:, channel, :]
+                token = _sample_top_p_only_compiled(
+                    self,
+                    logits,
+                    temperature=temperatures[channel],
+                    top_p=top_ps[channel],
+                    uniform=uniform,
+                )
+            else:
+                scores = realprocessor[channel](input_ids[..., channel], logits)
+                token = torch.multinomial(F.softmax(scores, dim=-1), num_samples=1).squeeze(1)
+        else:
+            token = torch.argmax(logits, dim=-1)
+        next_tokens[:, channel] = token
+        if channel != local_num_steps - 1:
+            current_input = self.speech_embedding_to_local_mlp(
+                self.model.embedding_list[channel](token)
+            )
+    return next_tokens
+
+
+def _sample_frame_fixed_full_cudagraph(
+    self,
+    hidden_states: torch.Tensor,
+    *,
+    input_ids: torch.Tensor,
+    realprocessor: list[LogitsProcessorList],
+    do_samples: list[bool],
+    fast_top_p: bool,
+    temperatures: list[float],
+    top_ps: list[float],
+    local_num_steps: int,
+) -> torch.Tensor:
+    if hidden_states.device.type != "cuda":
+        return _sample_frame_fixed_full_impl(
+            self,
+            hidden_states,
+            input_ids=input_ids,
+            realprocessor=realprocessor,
+            do_samples=do_samples,
+            fast_top_p=fast_top_p,
+            temperatures=temperatures,
+            top_ps=top_ps,
+            local_num_steps=local_num_steps,
+        )
+    if any(do_samples):
+        return _sample_frame_hybrid_cudagraph(
+            self,
+            hidden_states,
+            input_ids=input_ids,
+            realprocessor=realprocessor,
+            do_samples=do_samples,
+            fast_top_p=fast_top_p,
+            temperatures=temperatures,
+            top_ps=top_ps,
+            local_num_steps=local_num_steps,
+        )
+
+    key = (
+        hidden_states.device.index,
+        hidden_states.dtype,
+        tuple(hidden_states.shape),
+        int(local_num_steps),
+        getattr(self, "_torchopt_mode", "fixed-full-cudagraph"),
+        getattr(self, "_torchopt_compile_mode", None),
+    )
+    entry = self._torchopt_frame_graphs.get(key)
+    if entry is None:
+        static_hidden = torch.empty_like(hidden_states)
+        warmup_stream = torch.cuda.Stream(device=hidden_states.device)
+        warmup_stream.wait_stream(torch.cuda.current_stream(hidden_states.device))
+        with torch.cuda.stream(warmup_stream):
+            for _ in range(3):
+                _sample_frame_fixed_full_impl(
+                    self,
+                    static_hidden,
+                    input_ids=input_ids,
+                    realprocessor=realprocessor,
+                    do_samples=do_samples,
+                    local_num_steps=local_num_steps,
+                )
+        torch.cuda.current_stream(hidden_states.device).wait_stream(warmup_stream)
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            static_tokens = _sample_frame_fixed_full_impl(
+                self,
+                static_hidden,
+                input_ids=input_ids,
+                realprocessor=realprocessor,
+                do_samples=do_samples,
+                fast_top_p=fast_top_p,
+                temperatures=temperatures,
+                top_ps=top_ps,
+                local_num_steps=local_num_steps,
+            )
+        entry = (graph, static_hidden, static_tokens)
+        self._torchopt_frame_graphs[key] = entry
+
+    graph, static_hidden, static_tokens = entry
+    static_hidden.copy_(hidden_states)
+    graph.replay()
+    return static_tokens
+
+
+def _build_channel_processors(generation_config, channels: int) -> tuple[list[bool], list[LogitsProcessorList]]:
+    if generation_config.do_samples is None:
+        raise RuntimeError("MOSS-TTS generation_config.do_samples is required.")
+
+    do_samples = generation_config.do_samples
+    realprocessor = [LogitsProcessorList() for _ in range(channels)]
+    for index, layer_config in enumerate(generation_config.layers):
+        if not do_samples[index]:
+            continue
+        if layer_config.get("repetition_penalty") is not None and index != 0:
+            realprocessor[index].append(
+                RepetitionPenaltyLogitsProcessor(
+                    penalty=layer_config.get("repetition_penalty")
+                )
+            )
+        if layer_config.get("temperature") is not None:
+            realprocessor[index].append(
+                TemperatureLogitsWarper(temperature=layer_config.get("temperature"))
+            )
+        if layer_config.get("top_k") is not None and int(layer_config.get("top_k")) > 0:
+            realprocessor[index].append(TopKLogitsWarper(top_k=layer_config.get("top_k")))
+        if layer_config.get("top_p") is not None:
+            realprocessor[index].append(TopPLogitsWarper(top_p=layer_config.get("top_p")))
+    return do_samples, realprocessor
+
+
+def _build_fast_top_p_params(generation_config, channels: int) -> tuple[bool, list[float], list[float]]:
+    if generation_config.do_samples is None:
+        return False, [], []
+    temperatures: list[float] = []
+    top_ps: list[float] = []
+    can_use_fast = True
+    for index, layer_config in enumerate(generation_config.layers[:channels]):
+        temperature = layer_config.get("temperature")
+        top_p = layer_config.get("top_p")
+        repetition_penalty = layer_config.get("repetition_penalty")
+        temperatures.append(1.0 if temperature is None else float(temperature))
+        top_ps.append(1.0 if top_p is None else float(top_p))
+        if repetition_penalty is not None and float(repetition_penalty) != 1.0:
+            can_use_fast = False
+    return can_use_fast, temperatures, top_ps
+
+
+def _sample_top_p_only(scores: torch.Tensor, temperature: float, top_p: float) -> torch.Tensor:
+    if temperature != 1.0:
+        scores = scores / temperature
+    if top_p >= 1.0:
+        uniform = torch.rand_like(scores).clamp_(
+            min=torch.finfo(scores.dtype).tiny,
+            max=1.0 - torch.finfo(scores.dtype).eps,
+        )
+        gumbel = -torch.log(-torch.log(uniform))
+        return torch.argmax(scores + gumbel, dim=-1)
+
+    sorted_logits, sorted_indices = torch.sort(scores, descending=False)
+    cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
+    sorted_indices_to_remove = cumulative_probs <= (1.0 - top_p)
+    sorted_indices_to_remove[..., -1:] = False
+    removed_mass = cumulative_probs.masked_fill(~sorted_indices_to_remove, 0.0).amax(
+        dim=-1,
+        keepdim=True,
+    )
+    uniform = torch.rand_like(removed_mass)
+    sample_prob = removed_mass + uniform * (1.0 - removed_mass)
+    sorted_token = torch.searchsorted(cumulative_probs, sample_prob, right=True).clamp_max(scores.shape[-1] - 1)
+    return sorted_indices.gather(1, sorted_token).squeeze(1)
+
+
+def _sample_top_p_only_with_uniform_triton(
+    scores: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    uniform: torch.Tensor,
+) -> torch.Tensor | None:
+    if (
+        triton is None
+        or not scores.is_cuda
+        or scores.ndim != 2
+        or not scores.is_contiguous()
+        or scores.dtype not in {torch.float16, torch.bfloat16, torch.float32}
+        or uniform.numel() != int(scores.shape[0])
+        or top_p >= 1.0
+        or temperature <= 0.0
+    ):
+        return None
+    n_cols = int(scores.shape[-1])
+    if n_cols <= 0 or n_cols > 2048:
+        return None
+    out = torch.empty((int(scores.shape[0]),), device=scores.device, dtype=torch.long)
+    block = 1 << (n_cols - 1).bit_length()
+    _triton_top_p_sample_kernel[(int(scores.shape[0]),)](
+        scores,
+        uniform.contiguous().reshape(-1),
+        out,
+        n_cols,
+        int(scores.stride(0)),
+        float(temperature),
+        float(top_p),
+        block=block,
+        num_warps=8 if block >= 2048 else 4,
+    )
+    return out
+
+
+def _sample_fused_lm_head_top_p_triton(
+    self,
+    lm_hidden: torch.Tensor,
+    *,
+    channel: int,
+    temperature: float,
+    top_p: float,
+    uniform: torch.Tensor,
+) -> torch.Tensor | None:
+    if (
+        triton is None
+        or not bool(getattr(self, "_torchopt_triton_fused_lm_head", False))
+        or not lm_hidden.is_cuda
+        or lm_hidden.ndim != 2
+        or int(lm_hidden.shape[0]) != 1
+        or channel <= 0
+        or top_p >= 1.0
+        or temperature <= 0.0
+    ):
+        return None
+    weight = self.lm_heads[channel].weight
+    if (
+        weight.dtype not in {torch.float16, torch.bfloat16, torch.float32}
+        or lm_hidden.dtype not in {torch.float16, torch.bfloat16, torch.float32}
+        or int(weight.shape[0]) > 2048
+        or int(weight.shape[1]) > 4096
+    ):
+        return None
+    hidden = lm_hidden.reshape(-1).contiguous()
+    weight = weight.contiguous()
+    vocab_size = int(weight.shape[0])
+    hidden_size = int(weight.shape[1])
+    out = torch.empty((1,), device=lm_hidden.device, dtype=torch.long)
+    block_v = 1 << (vocab_size - 1).bit_length()
+    _triton_fused_audio_lm_head_sample_kernel[(1,)](
+        hidden,
+        weight,
+        uniform.contiguous().reshape(-1),
+        out,
+        vocab_size,
+        hidden_size,
+        int(hidden.stride(0)),
+        int(weight.stride(0)),
+        int(weight.stride(1)),
+        float(temperature),
+        float(top_p),
+        int(self.config.audio_pad_code),
+        block_v=block_v,
+        block_d=32,
+        num_warps=8,
+    )
+    return out
+
+
+def _sample_top_p_only_with_uniform(
+    scores: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    uniform: torch.Tensor,
+) -> torch.Tensor:
+    if temperature != 1.0:
+        scores = scores / temperature
+    sorted_logits, sorted_indices = torch.sort(scores, descending=False)
+    cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
+    sorted_indices_to_remove = cumulative_probs <= (1.0 - top_p)
+    sorted_indices_to_remove[..., -1:] = False
+    removed_mass = cumulative_probs.masked_fill(~sorted_indices_to_remove, 0.0).amax(
+        dim=-1,
+        keepdim=True,
+    )
+    sample_prob = removed_mass + uniform * (1.0 - removed_mass)
+    sorted_token = torch.searchsorted(cumulative_probs, sample_prob, right=True).clamp_max(scores.shape[-1] - 1)
+    return sorted_indices.gather(1, sorted_token).squeeze(1)
+
+
+def _sample_top_p_only_with_uniform_dynamic(
+    scores: torch.Tensor,
+    *,
+    temperature: torch.Tensor,
+    top_p: torch.Tensor,
+    uniform: torch.Tensor,
+) -> torch.Tensor:
+    scores = scores / temperature.clamp_min(torch.finfo(scores.dtype).tiny)
+    sorted_logits, sorted_indices = torch.sort(scores, descending=False)
+    cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
+    effective_top_p = top_p.clamp(1e-6, 1.0 - 1e-6)
+    sorted_indices_to_remove = cumulative_probs <= (1.0 - effective_top_p)
+    sorted_indices_to_remove[..., -1:] = False
+    removed_mass = cumulative_probs.masked_fill(~sorted_indices_to_remove, 0.0).amax(
+        dim=-1,
+        keepdim=True,
+    )
+    sample_prob = removed_mass + uniform * (1.0 - removed_mass)
+    sorted_token = torch.searchsorted(cumulative_probs, sample_prob, right=True).clamp_max(scores.shape[-1] - 1)
+    return sorted_indices.gather(1, sorted_token).squeeze(1)
+
+
+def _sample_temperature_gumbel_with_uniform(
+    scores: torch.Tensor,
+    *,
+    temperature: float,
+    uniform: torch.Tensor,
+) -> torch.Tensor:
+    if temperature != 1.0:
+        scores = scores / temperature
+    uniform = uniform[..., : scores.shape[-1]]
+    uniform = uniform.clamp(
+        min=torch.finfo(uniform.dtype).tiny,
+        max=1.0 - torch.finfo(uniform.dtype).eps,
+    )
+    gumbel = -torch.log(-torch.log(uniform))
+    return torch.argmax(scores + gumbel, dim=-1)
+
+
+def _sample_top_m_top_p(
+    scores: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    prefilter_size: int,
+) -> torch.Tensor:
+    if temperature != 1.0:
+        scores = scores / temperature
+    vocab_size = scores.shape[-1]
+    top_m = min(max(1, int(prefilter_size)), vocab_size)
+    top_logits, top_indices = torch.topk(scores, k=top_m, dim=-1, largest=True, sorted=True)
+    probs = F.softmax(top_logits, dim=-1)
+
+    if top_p < 1.0:
+        cumulative_probs = probs.cumsum(dim=-1)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = False
+        probs = probs.masked_fill(sorted_indices_to_remove, 0.0)
+
+    cumulative_probs = probs.cumsum(dim=-1)
+    total_mass = cumulative_probs[..., -1:].clamp_min(torch.finfo(cumulative_probs.dtype).tiny)
+    sample_prob = torch.rand_like(total_mass) * total_mass
+    top_token = torch.searchsorted(cumulative_probs, sample_prob, right=True).clamp_max(top_m - 1)
+    return top_indices.gather(1, top_token).squeeze(1)
+
+
+def _sample_top_m_top_p_with_uniform(
+    scores: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    prefilter_size: int,
+    uniform: torch.Tensor,
+) -> torch.Tensor:
+    if temperature != 1.0:
+        scores = scores / temperature
+    vocab_size = scores.shape[-1]
+    top_m = min(max(1, int(prefilter_size)), vocab_size)
+    top_logits, top_indices = torch.topk(scores, k=top_m, dim=-1, largest=True, sorted=True)
+    probs = F.softmax(top_logits, dim=-1)
+
+    if top_p < 1.0:
+        cumulative_probs = probs.cumsum(dim=-1)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = False
+        probs = probs.masked_fill(sorted_indices_to_remove, 0.0)
+
+    cumulative_probs = probs.cumsum(dim=-1)
+    total_mass = cumulative_probs[..., -1:].clamp_min(torch.finfo(cumulative_probs.dtype).tiny)
+    sample_prob = uniform * total_mass
+    top_token = torch.searchsorted(cumulative_probs, sample_prob, right=True).clamp_max(top_m - 1)
+    return top_indices.gather(1, top_token).squeeze(1)
+
+
+def _sample_top_p_only_compiled(
+    self,
+    scores: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    uniform: torch.Tensor | None = None,
+) -> torch.Tensor:
+    prefilter_size = int(getattr(self, "_torchopt_top_p_prefilter_size", 0) or 0)
+    if scores.device.type != "cuda":
+        if uniform is not None and prefilter_size == 0 and top_p < 1.0:
+            return _sample_top_p_only_with_uniform(
+                scores,
+                temperature=temperature,
+                top_p=top_p,
+                uniform=uniform,
+            )
+        if prefilter_size > 0:
+            return _sample_top_m_top_p(
+                scores,
+                temperature=temperature,
+                top_p=top_p,
+                prefilter_size=prefilter_size,
+            )
+        return _sample_top_p_only(scores, temperature=temperature, top_p=top_p)
+    if uniform is not None and prefilter_size == 0 and top_p < 1.0:
+        return _sample_top_p_only_cudagraph(
+            self,
+            scores,
+            temperature=temperature,
+            top_p=top_p,
+            uniform=uniform,
+        )
+    key = (
+        scores.device.index,
+        scores.dtype,
+        tuple(scores.shape),
+        float(temperature),
+        float(top_p),
+        prefilter_size,
+        getattr(self, "_torchopt_compile_mode", None),
+    )
+    samplers = getattr(self, "_torchopt_compiled_top_p_samplers", None)
+    if samplers is None:
+        samplers = {}
+        self._torchopt_compiled_top_p_samplers = samplers
+    sampler = samplers.get(key)
+    if sampler is None:
+        compile_mode = getattr(self, "_torchopt_compile_mode", None)
+        if compile_mode == "default":
+            compile_mode = None
+
+        def sampler_impl(inner_scores: torch.Tensor) -> torch.Tensor:
+            if prefilter_size > 0:
+                return _sample_top_m_top_p(
+                    inner_scores,
+                    temperature=temperature,
+                    top_p=top_p,
+                    prefilter_size=prefilter_size,
+                )
+            return _sample_top_p_only(inner_scores, temperature=temperature, top_p=top_p)
+
+        sampler = torch.compile(sampler_impl, dynamic=False, mode=compile_mode)
+        samplers[key] = sampler
+    return sampler(scores)
+
+
+def _sample_top_p_only_cudagraph(
+    self,
+    scores: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    uniform: torch.Tensor,
+) -> torch.Tensor:
+    key = (
+        scores.device.index,
+        scores.dtype,
+        tuple(scores.shape),
+        tuple(uniform.shape),
+        float(temperature),
+        float(top_p),
+        getattr(self, "_torchopt_compile_mode", None),
+    )
+    graphs = getattr(self, "_torchopt_top_p_graphs", None)
+    if graphs is None:
+        graphs = {}
+        self._torchopt_top_p_graphs = graphs
+    entry = graphs.get(key)
+    if entry is None:
+        compile_mode = getattr(self, "_torchopt_compile_mode", None)
+        if compile_mode == "default":
+            compile_mode = None
+
+        def sampler_impl(inner_scores: torch.Tensor, inner_uniform: torch.Tensor) -> torch.Tensor:
+            return _sample_top_p_only_with_uniform(
+                inner_scores,
+                temperature=temperature,
+                top_p=top_p,
+                uniform=inner_uniform,
+            )
+
+        sampler = torch.compile(sampler_impl, dynamic=False, mode=compile_mode)
+        static_scores = torch.empty_like(scores)
+        static_uniform = torch.empty_like(uniform)
+        static_scores.copy_(scores)
+        static_uniform.copy_(uniform)
+        for _ in range(3):
+            static_output = sampler(static_scores, static_uniform)
+        torch.cuda.synchronize(scores.device)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            static_output = sampler(static_scores, static_uniform)
+        entry = {
+            "graph": graph,
+            "scores": static_scores,
+            "uniform": static_uniform,
+            "output": static_output,
+        }
+        graphs[key] = entry
+
+    entry["scores"].copy_(scores)
+    entry["uniform"].copy_(uniform)
+    entry["graph"].replay()
+    return entry["output"].clone()
+
+
+def _fast_prepare_inputs_for_generation(
+    self,
+    input_ids: torch.Tensor,
+    model_kwargs: dict[str, Any],
+) -> dict[str, Any] | None:
+    past_key_values = model_kwargs.get("past_key_values", None)
+    cache_position = model_kwargs.get("cache_position", None)
+    if past_key_values is None or cache_position is None:
+        return None
+    attention_mask = model_kwargs.get("attention_mask", None)
+    if attention_mask is None:
+        return None
+    if input_ids.ndim != 3 or cache_position.ndim != 1:
+        return None
+
+    model_inputs: dict[str, Any] = {
+        "input_ids": input_ids[:, -cache_position.shape[0] :, :],
+        "inputs_embeds": None,
+        "past_key_values": past_key_values,
+        "attention_mask": attention_mask,
+        "cache_position": cache_position,
+        "position_ids": cache_position.unsqueeze(0),
+    }
+    cache_implementation = model_kwargs.get("cache_implementation", None)
+    if cache_implementation is not None:
+        model_inputs["cache_implementation"] = cache_implementation
+    for key in ("speaker_ids", "speaker_embeds", "style_features", "style_attention_mask"):
+        value = model_kwargs.get(key, None)
+        if value is not None:
+            model_inputs[key] = value
+    return model_inputs
+
+
+def _optimized_sample(
+    self,
+    input_ids: torch.LongTensor,
+    logits_processor,
+    stopping_criteria,
+    generation_config,
+    synced_gpus: bool,
+    streamer=None,
+    **model_kwargs,
+):
+    speech_pad_idx = self.config.audio_pad_code
+    device = input_ids.device
+    eos_token_id = generation_config.eos_token_id
+    output_attentions = generation_config.output_attentions
+    output_hidden_states = generation_config.output_hidden_states
+    output_scores = generation_config.output_scores
+    output_logits = generation_config.output_logits
+    return_dict_in_generate = generation_config.return_dict_in_generate
+    has_eos_stopping_criteria = any(
+        hasattr(criteria, "eos_token_id") for criteria in stopping_criteria
+    )
+
+    scores = () if (return_dict_in_generate and output_scores) else None
+    raw_logits = () if (return_dict_in_generate and output_logits) else None
+    decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+    decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+    batch_size, cur_len, channels = input_ids.shape
+    input_ids_length = cur_len
+    this_peer_finished = False
+    unfinished_sequences = torch.ones(
+        batch_size, dtype=torch.long, device=input_ids.device
+    )
+    speaker_ids = model_kwargs.get("speaker_ids", None)
+    if speaker_ids is not None:
+        speaker_ids = speaker_ids.to(device=input_ids.device, dtype=torch.long)
+        if speaker_ids.ndim != 1 or speaker_ids.shape[0] != batch_size:
+            raise ValueError(
+                f"`speaker_ids` must have shape ({batch_size},), got {tuple(speaker_ids.shape)}."
+            )
+        num_speakers = int(getattr(self.model, "num_speakers", 0) or 0)
+        if num_speakers > 0:
+            min_speaker_id = int(speaker_ids.min().item())
+            max_speaker_id = int(speaker_ids.max().item())
+            if min_speaker_id < 0 or max_speaker_id >= num_speakers:
+                raise ValueError(
+                    f"`speaker_ids` must be in [0, {num_speakers - 1}], "
+                    f"got min={min_speaker_id}, max={max_speaker_id}."
+                )
+        model_kwargs["speaker_ids"] = speaker_ids
+    style_features = model_kwargs.get("style_features", None)
+    if style_features is not None:
+        style_features = style_features.to(device=input_ids.device)
+        if style_features.ndim != 2 or style_features.shape[0] != batch_size:
+            raise ValueError(
+                f"`style_features` must have shape ({batch_size}, feature_dim), got {tuple(style_features.shape)}."
+            )
+        model_kwargs["style_features"] = style_features
+    style_attention_mask = model_kwargs.get("style_attention_mask", None)
+    if style_attention_mask is not None:
+        model_kwargs["style_attention_mask"] = style_attention_mask.to(device=input_ids.device)
+    model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+    do_samples, realprocessor = _build_channel_processors(generation_config, channels)
+    fast_top_p, temperatures, top_ps = _build_fast_top_p_params(generation_config, channels)
+    profile_enabled = bool(getattr(self, "_torchopt_profile_generation", False))
+    profile = None
+    if profile_enabled:
+        profile = {
+            "steps": 0,
+            "prepare_inputs_sec": 0.0,
+            "model_forward_sec": 0.0,
+            "update_cache_sec": 0.0,
+            "local_sample_sec": 0.0,
+            "append_tokens_sec": 0.0,
+            "stopping_sec": 0.0,
+        }
+
+    def profile_start() -> float:
+        if not profile_enabled:
+            return 0.0
+        if input_ids.device.type == "cuda":
+            torch.cuda.synchronize(input_ids.device)
+        return time.perf_counter()
+
+    def profile_add(name: str, start: float) -> None:
+        if not profile_enabled or profile is None:
+            return
+        if input_ids.device.type == "cuda":
+            torch.cuda.synchronize(input_ids.device)
+        profile[name] += time.perf_counter() - start
+
+    while self._has_unfinished_sequences(
+        this_peer_finished, synced_gpus, device=input_ids.device
+    ):
+        if profile is not None:
+            profile["steps"] += 1
+        stage_start = profile_start()
+        model_inputs = None
+        if bool(getattr(self, "_torchopt_fast_prepare_inputs", False)):
+            model_inputs = _fast_prepare_inputs_for_generation(self, input_ids, model_kwargs)
+        if model_inputs is None:
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+        model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+        model_inputs["output_hidden_states"] = True
+        profile_add("prepare_inputs_sec", stage_start)
+
+        stage_start = profile_start()
+        outputs = self(
+            **model_inputs,
+            n_vq_for_inference=generation_config.n_vq_for_inference,
+            return_dict=True,
+        )
+        profile_add("model_forward_sec", stage_start)
+
+        stage_start = profile_start()
+        model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
+        profile_add("update_cache_sec", stage_start)
+
+        if synced_gpus and this_peer_finished:
+            continue
+
+        hidden_source = getattr(outputs, "last_hidden_state", None)
+        if hidden_source is None:
+            hidden_source = outputs.hidden_states[-1]
+        hidden = hidden_source[:, -1, :]
+        local_num_steps = min(channels, 1 + generation_config.n_vq_for_inference)
+        mode = getattr(self, "_torchopt_mode", "fixed-full")
+        stage_start = profile_start()
+        if mode == "greedy-full-compile-cudagraph":
+            next_tokens = _sample_frame_greedy_full_compile_cudagraph(
+                self,
+                hidden,
+                input_ids=input_ids,
+                realprocessor=realprocessor,
+                do_samples=do_samples,
+                fast_top_p=fast_top_p,
+                temperatures=temperatures,
+                top_ps=top_ps,
+                local_num_steps=local_num_steps,
+            )
+        elif mode == "prefix-full-compile-cudagraph":
+            next_tokens = _sample_frame_prefix_top_p_compile_cudagraph(
+                self,
+                hidden,
+                do_samples=do_samples,
+                temperatures=temperatures,
+                top_ps=top_ps,
+                local_num_steps=local_num_steps,
+            )
+        elif mode in {"static-local-cache", "static-local-cache-compile"}:
+            next_tokens = _sample_frame_static_local_cache_cudagraph(
+                self,
+                hidden,
+                input_ids=input_ids,
+                realprocessor=realprocessor,
+                do_samples=do_samples,
+                fast_top_p=fast_top_p,
+                temperatures=temperatures,
+                top_ps=top_ps,
+                local_num_steps=local_num_steps,
+            )
+        elif mode == "static-local-cache-triton-compile-cudagraph":
+            next_tokens = _sample_frame_static_local_cache_triton_attention_compile_cudagraph(
+                self,
+                hidden,
+                do_samples=do_samples,
+                temperatures=temperatures,
+                top_ps=top_ps,
+                local_num_steps=local_num_steps,
+            )
+        elif mode in {"fixed-full-cudagraph", "fixed-full-compile-cudagraph"}:
+            next_tokens = _sample_frame_fixed_full_cudagraph(
+                self,
+                hidden,
+                input_ids=input_ids,
+                realprocessor=realprocessor,
+                do_samples=do_samples,
+                fast_top_p=fast_top_p,
+                temperatures=temperatures,
+                top_ps=top_ps,
+                local_num_steps=local_num_steps,
+            )
+        else:
+            next_tokens = _sample_frame_fixed_full_impl(
+                self,
+                hidden,
+                input_ids=input_ids,
+                realprocessor=realprocessor,
+                do_samples=do_samples,
+                local_num_steps=local_num_steps,
+            )
+        profile_add("local_sample_sec", stage_start)
+
+        if has_eos_stopping_criteria:
+            for index in range(channels):
+                pad = eos_token_id if index == 0 else speech_pad_idx
+                next_tokens[:, index] = (
+                    next_tokens[:, index] * unfinished_sequences
+                    + pad * (1 - unfinished_sequences)
+                )
+
+        stage_start = profile_start()
+        input_ids = torch.cat([input_ids, next_tokens[:, None, :]], dim=1)
+        if streamer is not None:
+            streamer.put(next_tokens[:, 0].cpu())
+        profile_add("append_tokens_sec", stage_start)
+
+        stage_start = profile_start()
+        stopping = stopping_criteria(input_ids[..., 0], scores)
+        unfinished_sequences = unfinished_sequences & ~stopping
+        this_peer_finished = unfinished_sequences.max() == 0
+        profile_add("stopping_sec", stage_start)
+
+        if return_dict_in_generate:
+            if output_scores:
+                raise RuntimeError("output_scores is not supported by torch optimized sampler")
+            if output_logits:
+                raise RuntimeError("output_logits is not supported by torch optimized sampler")
+            if output_attentions:
+                decoder_attentions += (outputs.attentions,)
+            if output_hidden_states:
+                decoder_hidden_states += (outputs.hidden_states,)
+
+        cur_len += 1
+        del outputs
+
+    if streamer is not None:
+        streamer.end()
+
+    if profile is not None:
+        total = sum(value for key, value in profile.items() if key.endswith("_sec"))
+        profile["profiled_sec"] = total
+        profile["sec_per_step"] = total / profile["steps"] if profile["steps"] else 0.0
+        runs = getattr(self, "_torchopt_profile_runs", None)
+        if runs is None:
+            runs = []
+            self._torchopt_profile_runs = runs
+        runs.append(profile)
+
+    if return_dict_in_generate:
+        from moss_tts_local_clipper_checkpoint.modeling_moss_tts import (
+            MossTTSGenerateDecoderOnlyOutput,
+        )
+
+        return MossTTSGenerateDecoderOnlyOutput(
+            sequences=input_ids,
+            scores=scores,
+            logits=raw_logits,
+            attentions=decoder_attentions,
+            hidden_states=decoder_hidden_states,
+            past_key_values=model_kwargs.get("past_key_values"),
+        )
+
+    start_indices = find_last_equal_C(input_ids[..., 0], self.config.audio_start_token_id)
+    start_lengths = input_ids_length - start_indices - 1
+    output = []
+    for start_idx, start_length, cur_generation_ids in zip(
+        start_indices, start_lengths, input_ids
+    ):
+        output.append((start_length, cur_generation_ids[start_idx:]))
+    return output