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