diff --git a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/__init__.py b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/__init__.py
index 074c48463276..855c9a539b10 100644
--- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/__init__.py
+++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/__init__.py
@@ -23,6 +23,23 @@
 vLLM's ``ModelRegistry`` through the ``vllm.general_plugins`` entry point.
 """
 
+import os
+
+_compat_mode = os.environ.get("EASYMAGPIE_VLLM_COMPAT_MODE", "refit").strip().lower()
+if _compat_mode not in {"0", "false", "none", "off"}:
+    from easymagpie_vllm_omni.vllm_compat import (
+        install_easy_magpie_refit_rpc_compat,
+        install_easy_magpie_runtime_compat,
+        install_vllm_omni_compat,
+    )
+
+    if _compat_mode in {"full", "all"}:
+        install_vllm_omni_compat()
+    elif _compat_mode in {"serial", "runtime", "rl"}:
+        install_easy_magpie_runtime_compat()
+    else:
+        install_easy_magpie_refit_rpc_compat()
+
 from easymagpie_vllm_omni.config import EASYMAGPIE_SMALLMAMBA, EasyMagpieOmniArch
 
 __all__ = ["EASYMAGPIE_SMALLMAMBA", "EasyMagpieOmniArch"]
diff --git a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/backbone_patches.py b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/backbone_patches.py
index faceae41c24d..e2cf87161362 100644
--- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/backbone_patches.py
+++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/backbone_patches.py
@@ -20,15 +20,737 @@
 """
 from __future__ import annotations
 
+import inspect
+import os
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import vllm.v1.attention.backends.mamba_attn as _mamba_attn
+try:
+    import vllm.v1.attention.backends.mamba_attn as _mamba_attn
+except ModuleNotFoundError:
+    import vllm.v1.attention.backends.mamba2_attn as _mamba_attn
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
 
+_EASYMAGPIE_MOE_OP_REGISTERED = False
+
+
+def _easymagpie_moe_forward(
+    hidden_states: torch.Tensor,
+    router_logits: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    from vllm.forward_context import get_forward_context
+
+    forward_context = get_forward_context()
+    layer = forward_context.no_compile_layers[layer_name]
+    return layer.forward_impl(hidden_states, router_logits)
+
+
+def _easymagpie_moe_forward_fake(
+    hidden_states: torch.Tensor,
+    router_logits: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    del router_logits, layer_name
+    return torch.empty_like(hidden_states)
+
+
+def _register_easymagpie_moe_custom_op() -> None:
+    global _EASYMAGPIE_MOE_OP_REGISTERED
+    if _EASYMAGPIE_MOE_OP_REGISTERED:
+        return
+    from vllm.platforms import current_platform
+    from vllm.utils import direct_register_custom_op
+
+    try:
+        direct_register_custom_op(
+            op_name="easymagpie_moe_forward",
+            op_func=_easymagpie_moe_forward,
+            mutates_args=[],
+            fake_impl=_easymagpie_moe_forward_fake,
+            dispatch_key=current_platform.dispatch_key,
+        )
+    except RuntimeError as exc:
+        message = str(exc)
+        if "easymagpie_moe_forward" not in message and "already" not in message.lower():
+            raise
+    _EASYMAGPIE_MOE_OP_REGISTERED = True
+
+
+def _has_vllm_forward_context() -> bool:
+    try:
+        import vllm.forward_context as forward_context
+    except Exception:
+        return False
+    return getattr(forward_context, "_forward_context", None) is not None
+
+
+def _env_flag(name: str, default: bool = True) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    return value.lower() not in {"0", "false", "no", "off"}
+
+
+def _repeat_kv_heads_for_gqa(kv: torch.Tensor, num_query_heads: int) -> torch.Tensor:
+    if kv.shape[1] == num_query_heads:
+        return kv
+    repeats = num_query_heads // kv.shape[1]
+    return kv.repeat_interleave(repeats, dim=1)
+
+
+def _causal_sdpa(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    scale: float,
+    is_causal: bool,
+    query_start_pos: int = 0,
+) -> torch.Tensor:
+    out_dtype = q.dtype
+    num_query_heads = q.shape[1]
+    k = _repeat_kv_heads_for_gqa(k, num_query_heads)
+    v = _repeat_kv_heads_for_gqa(v, num_query_heads)
+
+    qf = torch.nan_to_num(q.float())
+    kf = torch.nan_to_num(k.float())
+    vf = torch.nan_to_num(v.float())
+    scores = torch.matmul(qf.transpose(0, 1), kf.transpose(0, 1).transpose(-2, -1)) * float(scale)
+    if is_causal:
+        q_pos = query_start_pos + torch.arange(q.shape[0], device=q.device)
+        k_pos = torch.arange(k.shape[0], device=q.device)
+        attn_mask = k_pos.unsqueeze(0) <= q_pos.unsqueeze(1)
+        scores = scores.masked_fill(~attn_mask.unsqueeze(0), float("-inf"))
+    attn = torch.softmax(scores, dim=-1)
+    attn = torch.nan_to_num(attn)
+    out = torch.matmul(attn, vf.transpose(0, 1))
+    return out.transpose(0, 1).to(out_dtype).contiguous()
+
+
+def _cache_block_size(key_cache: torch.Tensor) -> int:
+    return int(key_cache.shape[1]) if key_cache.ndim >= 4 else 1
+
+
+def _flatten_cache(cache: torch.Tensor) -> torch.Tensor:
+    return cache.reshape(-1, cache.shape[-2], cache.shape[-1])
+
+
+def _gather_cache_slots(cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor:
+    slots = slots.to(device=cache.device, dtype=torch.long)
+    if cache.ndim >= 4:
+        block_size = _cache_block_size(cache)
+        blocks = slots // block_size
+        offsets = slots % block_size
+        return cache[blocks, offsets]
+    return cache[slots]
+
+
+def _cache_slots_for_sequence(
+    block_tables: torch.Tensor,
+    row: int,
+    seq_len: int,
+    *,
+    block_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    positions = torch.arange(seq_len, dtype=torch.long, device=device)
+    if block_tables is None or block_tables.numel() == 0:
+        return positions
+    table = block_tables[row].to(device=device, dtype=torch.long)
+    blocks = table[positions // block_size]
+    return blocks * block_size + (positions % block_size)
+
+
+def _slice_start_loc(start_loc: torch.Tensor | None, fallback_lens: torch.Tensor) -> torch.Tensor:
+    if start_loc is not None:
+        return start_loc.to(dtype=torch.long)
+    out = torch.zeros(fallback_lens.numel() + 1, dtype=torch.long, device=fallback_lens.device)
+    out[1:] = torch.cumsum(fallback_lens.to(dtype=torch.long), dim=0)
+    return out
+
+
+def _torch_flash_attention_forward(
+    impl,
+    layer,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    output: torch.Tensor,
+) -> torch.Tensor:
+    if kv_cache.numel() > 0 and key is not None and value is not None:
+        torch.ops._C_cache_ops.reshape_and_cache_flash(
+            key,
+            value,
+            kv_cache[0],
+            kv_cache[1],
+            attn_metadata.slot_mapping.flatten(),
+            impl.kv_cache_dtype,
+            layer._k_scale,
+            layer._v_scale,
+        )
+
+    num_prefill_tokens = int(getattr(attn_metadata, "num_prefill_tokens", 0) or 0)
+    decode_tokens = int(getattr(attn_metadata, "num_decode_tokens", 0) or 0)
+    output.zero_()
+
+    prefill_meta = getattr(attn_metadata, "prefill_metadata", None)
+    if prefill_meta is not None and num_prefill_tokens > 0:
+        seq_lens_tensor = getattr(prefill_meta, "seq_lens_tensor", None)
+        if seq_lens_tensor is None:
+            seq_lens_tensor = torch.tensor(
+                list(getattr(prefill_meta, "seq_lens", []) or []),
+                dtype=torch.long,
+                device=query.device,
+            )
+        else:
+            seq_lens_tensor = seq_lens_tensor.to(device=query.device, dtype=torch.long)
+        query_start_loc = _slice_start_loc(getattr(prefill_meta, "query_start_loc", None), seq_lens_tensor)
+        seq_start_loc = _slice_start_loc(getattr(prefill_meta, "seq_start_loc", None), seq_lens_tensor)
+        block_tables = getattr(prefill_meta, "block_tables", None)
+        use_cache_for_prefill = (
+            kv_cache.numel() > 0
+            and block_tables is not None
+            and getattr(block_tables, "numel", lambda: 0)() > 0
+        )
+        if use_cache_for_prefill:
+            key_cache = kv_cache[0]
+            value_cache = kv_cache[1]
+            block_size = _cache_block_size(kv_cache[0])
+        for row in range(int(seq_lens_tensor.numel())):
+            qs = int(query_start_loc[row].item())
+            qe = int(query_start_loc[row + 1].item())
+            ks = int(seq_start_loc[row].item())
+            ke = int(seq_start_loc[row + 1].item())
+            if qe <= qs or ke <= ks:
+                continue
+            q_row = query[qs:qe]
+            if use_cache_for_prefill:
+                slots = _cache_slots_for_sequence(
+                    block_tables,
+                    row,
+                    int(seq_lens_tensor[row].item()),
+                    block_size=block_size,
+                    device=query.device,
+                )
+                k_row = _gather_cache_slots(key_cache, slots)
+                v_row = _gather_cache_slots(value_cache, slots)
+            else:
+                k_row = key[ks:ke]
+                v_row = value[ks:ke]
+            query_start_pos = max(0, int(seq_lens_tensor[row].item()) - int(q_row.shape[0]))
+            output[qs:qe].copy_(
+                _causal_sdpa(
+                    q_row,
+                    k_row,
+                    v_row,
+                    scale=float(impl.scale),
+                    is_causal=True,
+                    query_start_pos=query_start_pos,
+                )
+            )
+
+    decode_meta = getattr(attn_metadata, "decode_metadata", None)
+    if decode_meta is not None and decode_tokens > 0:
+        if kv_cache.numel() == 0:
+            start = num_prefill_tokens
+            output[start : start + decode_tokens].copy_(
+                _causal_sdpa(
+                    query[start : start + decode_tokens],
+                    key[start : start + decode_tokens],
+                    value[start : start + decode_tokens],
+                    scale=float(impl.scale),
+                    is_causal=True,
+                )
+            )
+            return output
+        key_cache = kv_cache[0]
+        value_cache = kv_cache[1]
+        block_size = _cache_block_size(kv_cache[0])
+        seq_lens = getattr(decode_meta, "seq_lens_tensor", None)
+        if seq_lens is None:
+            seq_lens = torch.tensor(
+                list(getattr(decode_meta, "seq_lens", []) or []),
+                dtype=torch.long,
+                device=query.device,
+            )
+        else:
+            seq_lens = seq_lens.to(device=query.device, dtype=torch.long)
+        block_tables = getattr(decode_meta, "block_tables", None)
+        query_start_loc = getattr(decode_meta, "query_start_loc", None)
+        if query_start_loc is None:
+            query_lens = torch.ones(seq_lens.numel(), dtype=torch.long, device=query.device)
+            query_start_loc = _slice_start_loc(None, query_lens)
+        else:
+            query_start_loc = query_start_loc.to(device=query.device, dtype=torch.long)
+        decode_query_offset = num_prefill_tokens
+        for row in range(int(seq_lens.numel())):
+            qs = decode_query_offset + int(query_start_loc[row].item())
+            qe = decode_query_offset + int(query_start_loc[row + 1].item())
+            if qe <= qs:
+                continue
+            seq_len = int(seq_lens[row].item())
+            slots = _cache_slots_for_sequence(
+                block_tables,
+                row,
+                seq_len,
+                block_size=block_size,
+                device=query.device,
+            )
+            output[qs:qe].copy_(
+                _causal_sdpa(
+                    query[qs:qe],
+                    _gather_cache_slots(key_cache, slots),
+                    _gather_cache_slots(value_cache, slots),
+                    scale=float(impl.scale),
+                    is_causal=(qe - qs) > 1,
+                    query_start_pos=max(0, seq_len - (qe - qs)),
+                )
+            )
+    return output
+
+
+def _torch_v1_flash_attention_forward(
+    impl,
+    layer,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    output: torch.Tensor,
+) -> torch.Tensor:
+    if attn_metadata is None:
+        return output
+
+    if (
+        kv_cache.numel() > 0
+        and key is not None
+        and value is not None
+        and getattr(impl, "kv_sharing_target_layer_name", None) is None
+    ):
+        slot_mapping = getattr(attn_metadata, "slot_mapping", None)
+        if slot_mapping is None:
+            slot_mapping = torch.arange(int(key.shape[0]), dtype=torch.long, device=key.device)
+        else:
+            slot_mapping = slot_mapping.to(device=key.device, dtype=torch.long).flatten()
+        torch.ops._C_cache_ops.reshape_and_cache_flash(
+            key,
+            value,
+            kv_cache[0],
+            kv_cache[1],
+            slot_mapping,
+            impl.kv_cache_dtype,
+            layer._k_scale,
+            layer._v_scale,
+        )
+
+    num_actual_tokens_value = getattr(attn_metadata, "num_actual_tokens", None)
+    if num_actual_tokens_value is None:
+        inferred_tokens = int(getattr(attn_metadata, "num_prefill_tokens", 0) or 0) + int(
+            getattr(attn_metadata, "num_decode_tokens", 0) or 0
+        )
+        num_actual_tokens = inferred_tokens if inferred_tokens > 0 else int(query.shape[0])
+    else:
+        num_actual_tokens = int(num_actual_tokens_value or 0)
+    num_actual_tokens = min(num_actual_tokens, int(query.shape[0]))
+    if num_actual_tokens <= 0:
+        return output
+
+    output.zero_()
+    query_start_loc = getattr(attn_metadata, "query_start_loc", None)
+    if query_start_loc is None:
+        query_start_loc = torch.tensor([0, num_actual_tokens], dtype=torch.long, device=query.device)
+    else:
+        query_start_loc = query_start_loc.to(device=query.device, dtype=torch.long)
+    seq_lens = getattr(attn_metadata, "seq_lens", None)
+    if seq_lens is None:
+        seq_lens = torch.diff(query_start_loc)
+    else:
+        seq_lens = seq_lens.to(device=query.device, dtype=torch.long)
+
+    block_table = getattr(attn_metadata, "block_table", None)
+    use_cache = (
+        kv_cache.numel() > 0
+        and block_table is not None
+        and getattr(block_table, "numel", lambda: 0)() > 0
+    )
+    if use_cache:
+        key_cache = kv_cache[0]
+        value_cache = kv_cache[1]
+        block_size = _cache_block_size(kv_cache[0])
+    else:
+        key_cache = value_cache = None
+        block_size = 1
+        block_table = None
+
+    causal = bool(getattr(attn_metadata, "causal", True))
+    for row in range(int(seq_lens.numel())):
+        qs = int(query_start_loc[row].item())
+        qe = int(query_start_loc[row + 1].item())
+        q_len = qe - qs
+        seq_len = int(seq_lens[row].item())
+        if q_len <= 0 or seq_len <= 0:
+            continue
+        if use_cache:
+            slots = _cache_slots_for_sequence(
+                block_table,
+                row,
+                seq_len,
+                block_size=block_size,
+                device=query.device,
+            )
+            k_row = _gather_cache_slots(key_cache, slots)
+            v_row = _gather_cache_slots(value_cache, slots)
+        else:
+            k_row = key[qs : qs + seq_len]
+            v_row = value[qs : qs + seq_len]
+        query_start_pos = max(0, seq_len - q_len)
+        output[qs:qe].copy_(
+            _causal_sdpa(
+                query[qs:qe],
+                k_row,
+                v_row,
+                scale=float(impl.scale),
+                is_causal=causal,
+                query_start_pos=query_start_pos,
+            )
+        )
+    return output
+
+
+def _scale_is_effectively_one(value) -> bool:
+    if value is None:
+        return True
+    if not isinstance(value, torch.Tensor):
+        return float(value) == 1.0
+    if value.numel() == 0:
+        return True
+    return bool(torch.all(value.detach().float() == 1.0).item())
+
+
+def _window_is_unrestricted(value) -> bool:
+    if value is None:
+        return True
+    if isinstance(value, torch.Tensor):
+        if value.numel() == 0:
+            return True
+        return bool(torch.all(value.detach().cpu() == -1).item())
+    try:
+        items = list(value)
+    except TypeError:
+        return False
+    return len(items) > 0 and all(int(item) == -1 for item in items)
+
+
+def _softcap_is_zero(value) -> bool:
+    if value is None:
+        return True
+    if isinstance(value, torch.Tensor):
+        if value.numel() == 0:
+            return True
+        return bool(torch.all(value.detach().float().cpu() == 0).item())
+    return float(value) == 0.0
+
+
+def _torch_flash_attn_varlen_forward(
+    *,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    seqused_k: torch.Tensor | None = None,
+    cu_seqlens_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+) -> torch.Tensor:
+    del max_seqlen_q, max_seqlen_k
+    scale = float(softmax_scale) if softmax_scale is not None else float(q.shape[-1] ** -0.5)
+    cu_seqlens_q = cu_seqlens_q.to(device=q.device, dtype=torch.long)
+    num_rows = int(cu_seqlens_q.numel()) - 1
+    if seqused_k is not None:
+        seq_lens = seqused_k.to(device=q.device, dtype=torch.long)
+    elif cu_seqlens_k is not None:
+        seq_lens = torch.diff(cu_seqlens_k.to(device=q.device, dtype=torch.long))
+    else:
+        seq_lens = torch.diff(cu_seqlens_q)
+
+    use_block_cache = block_table is not None and getattr(block_table, "numel", lambda: 0)() > 0 and k.ndim >= 4
+    if use_block_cache:
+        key_cache = k
+        value_cache = v
+        block_size = _cache_block_size(k)
+    else:
+        key_cache = value_cache = None
+        block_size = 1
+
+    out.zero_()
+    for row in range(num_rows):
+        qs = int(cu_seqlens_q[row].item())
+        qe = int(cu_seqlens_q[row + 1].item())
+        q_len = qe - qs
+        seq_len = int(seq_lens[row].item())
+        if q_len <= 0 or seq_len <= 0:
+            continue
+        if use_block_cache:
+            slots = _cache_slots_for_sequence(
+                block_table,
+                row,
+                seq_len,
+                block_size=block_size,
+                device=q.device,
+            )
+            k_row = _gather_cache_slots(key_cache, slots)
+            v_row = _gather_cache_slots(value_cache, slots)
+        else:
+            if cu_seqlens_k is not None:
+                ks = int(cu_seqlens_k[row].item())
+            else:
+                ks = qs
+            k_row = k[ks : ks + seq_len]
+            v_row = v[ks : ks + seq_len]
+        query_start_pos = max(0, seq_len - q_len)
+        out[qs:qe].copy_(
+            _causal_sdpa(
+                q[qs:qe],
+                k_row,
+                v_row,
+                scale=scale,
+                is_causal=causal,
+                query_start_pos=query_start_pos,
+            )
+        )
+    return out
+
+
+def _patch_v1_flash_attn_varlen_func() -> bool:
+    try:
+        import vllm.v1.attention.backends.flash_attn as v1_flash_attn
+    except Exception:
+        return False
+    original = getattr(v1_flash_attn, "flash_attn_varlen_func", None)
+    if original is None or getattr(original, "_easymagpie_torch_fallback", False):
+        return False
+
+    def flash_attn_varlen_func(q, k, v, *args, **kwargs):
+        enabled = os.environ.get("EASYMAGPIE_TORCH_FLASH_ATTN_FALLBACK", "1").lower()
+        unsupported = (
+            args
+            or enabled in {"0", "false", "no", "off"}
+            or kwargs.get("out") is None
+            or kwargs.get("q_v") is not None
+            or float(kwargs.get("dropout_p", 0.0) or 0.0) != 0.0
+            or not _window_is_unrestricted(kwargs.get("window_size"))
+            or kwargs.get("alibi_slopes") is not None
+            or not _softcap_is_zero(kwargs.get("softcap"))
+            or bool(kwargs.get("return_attn_probs", False))
+            or bool(kwargs.get("return_softmax_lse", False))
+        )
+        if unsupported:
+            if not getattr(flash_attn_varlen_func, "_easymagpie_warned_unsupported", False):
+                logger.warning(
+                    "EasyMagpie torch-SDPA varlen fallback bypassed unsupported FlashAttention call: "
+                    "args=%s window_size=%s alibi=%s softcap=%s q_v=%s out=%s",
+                    bool(args),
+                    kwargs.get("window_size"),
+                    kwargs.get("alibi_slopes") is not None,
+                    kwargs.get("softcap"),
+                    kwargs.get("q_v") is not None,
+                    kwargs.get("out") is not None,
+                )
+                flash_attn_varlen_func._easymagpie_warned_unsupported = True  # type: ignore[attr-defined]
+            return original(q, k, v, *args, **kwargs)
+        if not getattr(flash_attn_varlen_func, "_easymagpie_warned_active", False):
+            logger.warning("EasyMagpie torch-SDPA varlen fallback active for vLLM v1 FlashAttention")
+            flash_attn_varlen_func._easymagpie_warned_active = True  # type: ignore[attr-defined]
+        return _torch_flash_attn_varlen_forward(
+            q=q,
+            k=k,
+            v=v,
+            out=kwargs["out"],
+            cu_seqlens_q=kwargs["cu_seqlens_q"],
+            cu_seqlens_k=kwargs.get("cu_seqlens_k"),
+            seqused_k=kwargs.get("seqused_k"),
+            max_seqlen_q=int(kwargs["max_seqlen_q"]),
+            max_seqlen_k=int(kwargs["max_seqlen_k"]),
+            block_table=kwargs.get("block_table"),
+            softmax_scale=kwargs.get("softmax_scale"),
+            causal=bool(kwargs.get("causal", False)),
+        )
+
+    flash_attn_varlen_func._easymagpie_torch_fallback = True  # type: ignore[attr-defined]
+    v1_flash_attn.flash_attn_varlen_func = flash_attn_varlen_func
+    return True
+
+
+def _patch_flash_attention_impl(FlashAttentionImpl, *, is_v1: bool) -> bool:
+    original = getattr(FlashAttentionImpl, "forward", None)
+    if original is None or getattr(FlashAttentionImpl, "_easymagpie_torch_fallback_installed", False):
+        return False
+
+    original_signature = None
+    try:
+        original_signature = inspect.signature(original)
+        original_accepts_kwargs = any(
+            param.kind == inspect.Parameter.VAR_KEYWORD for param in original_signature.parameters.values()
+        )
+        original_accepts_output_block_scale = "output_block_scale" in original_signature.parameters
+    except Exception:
+        original_accepts_kwargs = False
+        original_accepts_output_block_scale = True
+
+    def _call_original(
+        self,
+        layer,
+        query,
+        key,
+        value,
+        kv_cache,
+        attn_metadata,
+        output,
+        output_scale,
+        output_block_scale,
+        extra_kwargs,
+    ):
+        base_args = (self, layer, query, key, value, kv_cache, attn_metadata)
+        if original_accepts_kwargs:
+            return original(
+                *base_args,
+                output=output,
+                output_scale=output_scale,
+                output_block_scale=output_block_scale,
+                **extra_kwargs,
+            )
+        if extra_kwargs:
+            if original_signature is None:
+                return original(*base_args, output, output_scale, output_block_scale)
+            return original(
+                *base_args,
+                output=output,
+                output_scale=output_scale,
+                **{key: value for key, value in extra_kwargs.items() if key in original_signature.parameters},
+            )
+        if original_accepts_output_block_scale:
+            return original(*base_args, output, output_scale, output_block_scale)
+        return original(*base_args, output, output_scale)
+
+    def _metadata_ok_for_fallback(self, key, value, kv_cache, attn_metadata) -> bool:
+        del self, key, value, kv_cache
+        return (not is_v1) or (attn_metadata is not None)
+
+    def forward(
+        self,
+        layer,
+        query,
+        key,
+        value,
+        kv_cache,
+        attn_metadata,
+        output=None,
+        output_scale=None,
+        output_block_scale=None,
+        **kwargs,
+    ):
+        enabled = os.environ.get("EASYMAGPIE_TORCH_FLASH_ATTN_FALLBACK", "1").lower()
+        window_ok = _window_is_unrestricted(getattr(self, "sliding_window", None))
+        softcap_ok = _softcap_is_zero(getattr(self, "logits_soft_cap", None))
+        metadata_ok = _metadata_ok_for_fallback(self, key, value, kv_cache, attn_metadata)
+        should_fallback = (
+            enabled not in {"0", "false", "no", "off"}
+            and output is not None
+            and output_scale is None
+            and output_block_scale is None
+            and not kwargs
+            and metadata_ok
+            and window_ok
+            and getattr(self, "alibi_slopes", None) is None
+            and softcap_ok
+            and not str(getattr(self, "kv_cache_dtype", "auto")).startswith("fp8")
+        )
+        if not should_fallback:
+            if enabled not in {"0", "false", "no", "off"} and not getattr(
+                forward, "_easymagpie_warned_bypassed", False
+            ):
+                logger.warning(
+                    "EasyMagpie torch-SDPA direct fallback bypassed for %s FlashAttentionImpl: "
+                    "output=%s output_scale=%s output_block_scale=%s extra_kwargs=%s "
+                    "metadata_ok=%s window=%s window_ok=%s alibi=%s softcap=%s "
+                    "softcap_ok=%s kv_cache_dtype=%s",
+                    "v1" if is_v1 else "legacy",
+                    output is not None,
+                    output_scale is not None,
+                    output_block_scale is not None,
+                    sorted(kwargs),
+                    metadata_ok,
+                    getattr(self, "sliding_window", None),
+                    window_ok,
+                    getattr(self, "alibi_slopes", None) is not None,
+                    getattr(self, "logits_soft_cap", None),
+                    softcap_ok,
+                    getattr(self, "kv_cache_dtype", None),
+                )
+                forward._easymagpie_warned_bypassed = True  # type: ignore[attr-defined]
+            return _call_original(
+                self,
+                layer,
+                query,
+                key,
+                value,
+                kv_cache,
+                attn_metadata,
+                output,
+                output_scale,
+                output_block_scale,
+                kwargs,
+            )
+        if not getattr(forward, "_easymagpie_warned_active", False):
+            logger.warning(
+                "EasyMagpie torch-SDPA direct fallback active for %s FlashAttentionImpl: "
+                "window=%s softcap=%s kv_cache_dtype=%s",
+                "v1" if is_v1 else "legacy",
+                getattr(self, "sliding_window", None),
+                getattr(self, "logits_soft_cap", None),
+                getattr(self, "kv_cache_dtype", None),
+            )
+            forward._easymagpie_warned_active = True  # type: ignore[attr-defined]
+        if is_v1:
+            use_cascade = bool(getattr(attn_metadata, "use_cascade", False)) if attn_metadata is not None else False
+            if use_cascade:
+                if not getattr(forward, "_easymagpie_warned_cascade", False):
+                    logger.warning("EasyMagpie torch-SDPA direct fallback bypassed v1 cascade attention")
+                    forward._easymagpie_warned_cascade = True  # type: ignore[attr-defined]
+                return _call_original(
+                    self,
+                    layer,
+                    query,
+                    key,
+                    value,
+                    kv_cache,
+                    attn_metadata,
+                    output,
+                    output_scale,
+                    output_block_scale,
+                    kwargs,
+                )
+            return _torch_v1_flash_attention_forward(self, layer, query, key, value, kv_cache, attn_metadata, output)
+        return _torch_flash_attention_forward(self, layer, query, key, value, kv_cache, attn_metadata, output)
+
+    forward._easymagpie_torch_fallback = True  # type: ignore[attr-defined]
+    FlashAttentionImpl.forward = forward
+    FlashAttentionImpl._easymagpie_torch_fallback_installed = True
+    return True
+
+
 def patch_mamba_streaming_decode() -> None:
     """Treat 1-token streaming extends as decodes so FULL decode cudagraphs work.
 
@@ -51,15 +773,19 @@ def patch_mamba_streaming_decode() -> None:
       captured kernel reads the capture-time dummy slot (0) instead of the
       request's real Mamba-cache slot -> garbage hidden states.
 
-    Forcing ``treat_short_extends_as_decodes=True`` makes single-token extends
-    classify as decodes (``num_prefills==0``), which both matches the dispatched
-    FULL decode graph and re-enables the per-step ``state_indices_tensor_d``
-    refresh. Multi-token context prefills (``query_len>1``) still classify as
-    prefills, so this is safe for mixed batches. Advancing Mamba state by one
-    token via the decode kernels is semantically identical to a 1-token prefill
-    chunk (it reads the slot's state and writes the advanced state back in
-    place), so no state update is lost — the only requirement is exactly one new
-    token per streamed step (``SamplingParams(max_tokens=1)``).
+    Forcing ``treat_short_extends_as_decodes=True`` on vLLM versions that expose
+    that flag makes single-token extends classify as decodes
+    (``num_prefills==0``), which both matches the dispatched FULL decode graph
+    and re-enables the per-step ``state_indices_tensor_d`` refresh. Older vLLM
+    builds expose only ``decode_threshold`` and already classify
+    ``query_len <= decode_threshold`` as decode, so this wrapper must be
+    signature-aware and only forward kwargs the installed helper supports.
+    Multi-token context prefills (``query_len>1``) still classify as prefills,
+    so this is safe for mixed batches. Advancing Mamba state by one token via
+    the decode kernels is semantically identical to a 1-token prefill chunk (it
+    reads the slot's state and writes the advanced state back in place), so no
+    state update is lost — the only requirement is exactly one new token per
+    streamed step (``SamplingParams(max_tokens=1)``).
 
     Idempotent and process-global; the EasyMagpie plugin only ever serves this
     model so the global patch is acceptable.
@@ -67,6 +793,12 @@ def patch_mamba_streaming_decode() -> None:
     orig = _mamba_attn.split_decodes_and_prefills
     if getattr(orig, "_easymagpie_patched", False):
         return
+    try:
+        orig_params = inspect.signature(orig).parameters
+    except (TypeError, ValueError):
+        orig_params = {}
+    supports_require_uniform = "require_uniform" in orig_params
+    supports_treat_short_extends = "treat_short_extends_as_decodes" in orig_params
 
     def patched(
         common_attn_metadata,
@@ -74,12 +806,12 @@ def patched(
         require_uniform: bool = False,
         treat_short_extends_as_decodes: bool = True,
     ):
-        return orig(
-            common_attn_metadata,
-            decode_threshold=decode_threshold,
-            require_uniform=require_uniform,
-            treat_short_extends_as_decodes=True,
-        )
+        kwargs = {"decode_threshold": decode_threshold}
+        if supports_require_uniform:
+            kwargs["require_uniform"] = require_uniform
+        if supports_treat_short_extends:
+            kwargs["treat_short_extends_as_decodes"] = True
+        return orig(common_attn_metadata, **kwargs)
 
     patched._easymagpie_patched = True
     _mamba_attn.split_decodes_and_prefills = patched
@@ -93,6 +825,155 @@ def forward(self, x):
         return F.silu(x)
 
 
+def patch_nemotron_h_moe_layer() -> bool:
+    """Register an ``E`` MoE layer for vLLM Nemotron-H builds that omit it.
+
+    Some vLLM versions parse Nemotron-H but only ship Mamba, attention, and dense
+    MLP layer classes. EasyMagpie SmallMamba checkpoints keep routed expert
+    weights under ``layers.*.mixer.experts.*`` and need a real ``E`` layer to
+    load. This lightweight fallback mirrors the checkpoint layout directly
+    (router + two-matrix routed experts + two-matrix shared expert). It is slower
+    than vLLM's fused MoE path, but it preserves the model contract and keeps the
+    rollout backend usable on those builds.
+    """
+    from vllm.model_executor.layers.layernorm import RMSNorm
+    from vllm.model_executor.models import nemotron_h as nh
+
+    existing = nh.ALL_DECODER_LAYER_TYPES.get("E")
+    if existing is not None and getattr(existing, "_easymagpie_moe_layer", False):
+        return False
+    if existing is not None and existing.__name__ != "NemotronHMLPDecoderLayer":
+        return False
+
+    class _NemotronHTwoMatrixMLP(nn.Module):
+        def __init__(self, hidden_size: int, intermediate_size: int) -> None:
+            super().__init__()
+            self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+            self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+            self.act_fn = _SiluActivation()
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.down_proj(self.act_fn(self.up_proj(x)))
+
+    class NemotronHMoE(nn.Module):
+        def __init__(self, config, prefix: str = "") -> None:
+            super().__init__()
+            self.n_routed_experts = int(config.n_routed_experts)
+            self.num_experts_per_tok = int(config.num_experts_per_tok)
+            self.routed_scaling_factor = float(getattr(config, "routed_scaling_factor", 1.0))
+            self.norm_topk_prob = bool(getattr(config, "norm_topk_prob", True))
+            self.route_indices_on_cpu = _env_flag("EASYMAGPIE_MOE_CPU_ROUTING", True)
+            self.layer_name = prefix
+            self.use_custom_op = False
+
+            hidden_size = int(config.hidden_size)
+            routed_intermediate = int(getattr(config, "moe_intermediate_size", config.intermediate_size))
+            shared_intermediate = int(
+                getattr(config, "moe_shared_expert_intermediate_size", routed_intermediate)
+            )
+
+            self.gate = nn.Linear(hidden_size, self.n_routed_experts, bias=False)
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.empty(self.n_routed_experts), requires_grad=False
+            )
+            self.experts = nn.ModuleList(
+                [_NemotronHTwoMatrixMLP(hidden_size, routed_intermediate) for _ in range(self.n_routed_experts)]
+            )
+            self.shared_experts = _NemotronHTwoMatrixMLP(hidden_size, shared_intermediate)
+
+            if prefix:
+                try:
+                    from vllm import envs
+                    from vllm.config import get_current_vllm_config
+
+                    if bool(getattr(envs, "VLLM_USE_V1", False)):
+                        compilation_config = get_current_vllm_config().compilation_config
+                        static_context = compilation_config.static_forward_context
+                        if prefix in static_context and static_context[prefix] is not self:
+                            raise ValueError(f"Duplicate layer name: {prefix}")
+                        static_context[prefix] = self
+                        _register_easymagpie_moe_custom_op()
+                        self.use_custom_op = True
+                except Exception:
+                    logger.debug("Could not register EasyMagpie MoE custom op", exc_info=True)
+
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            original_shape = hidden_states.shape
+            hidden_states = hidden_states.reshape(-1, original_shape[-1])
+            router_logits = self.gate(hidden_states).float()
+            if self.use_custom_op and hidden_states.is_cuda and _has_vllm_forward_context():
+                output = torch.ops.vllm.easymagpie_moe_forward(
+                    hidden_states,
+                    router_logits,
+                    self.layer_name,
+                )
+            else:
+                output = self.forward_impl(hidden_states, router_logits)
+            return output.reshape(original_shape)
+
+        def forward_impl(
+            self,
+            hidden_states: torch.Tensor,
+            router_logits: torch.Tensor,
+        ) -> torch.Tensor:
+            bias = getattr(self.gate, "e_score_correction_bias", None)
+            if bias is not None:
+                router_logits = router_logits + bias.float().unsqueeze(0)
+            scores = torch.softmax(router_logits, dim=-1)
+            topk_weights, topk_ids = torch.topk(scores, k=self.num_experts_per_tok, dim=-1)
+            if self.norm_topk_prob:
+                topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True).clamp_min(1e-20)
+
+            routed = torch.zeros_like(hidden_states)
+            route_topk_ids = topk_ids
+            route_on_cpu = self.route_indices_on_cpu and topk_ids.is_cuda
+            if route_on_cpu:
+                route_topk_ids = topk_ids.detach().to("cpu")
+            for expert_id, expert in enumerate(self.experts):
+                token_ids, slot_ids = torch.where(route_topk_ids == expert_id)
+                if token_ids.numel() == 0:
+                    continue
+                if route_on_cpu:
+                    token_ids = token_ids.to(device=hidden_states.device, dtype=torch.long)
+                    slot_ids = slot_ids.to(device=hidden_states.device, dtype=torch.long)
+                expert_out = expert(hidden_states[token_ids])
+                weights = topk_weights[token_ids, slot_ids].to(expert_out.dtype).unsqueeze(-1)
+                routed[token_ids] += expert_out * weights
+
+            return routed * self.routed_scaling_factor + self.shared_experts(hidden_states)
+
+    class NemotronHMoEDecoderLayer(nn.Module):
+        _easymagpie_moe_layer = True
+
+        def __init__(
+            self,
+            config,
+            layer_idx: int,
+            model_config=None,
+            cache_config=None,
+            quant_config=None,
+            prefix: str = "",
+        ) -> None:
+            super().__init__()
+            del layer_idx, model_config, cache_config, quant_config
+            self.mixer = NemotronHMoE(config, prefix=f"{prefix}.mixer")
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        def forward(self, hidden_states: torch.Tensor, residual, **kwargs):
+            del kwargs
+            if residual is None:
+                residual = hidden_states
+                hidden_states = self.norm(hidden_states)
+            else:
+                hidden_states, residual = self.norm(hidden_states, residual)
+            hidden_states = self.mixer(hidden_states)
+            return hidden_states, residual
+
+    nh.ALL_DECODER_LAYER_TYPES["E"] = NemotronHMoEDecoderLayer
+    logger.info("EasyMagpie Nemotron-H MoE fallback layer registered")
+    return True
+
+
 def patch_silu_shared_experts(backbone) -> int:
     """Replace ``shared_experts.act_fn`` with SiLU on every NemotronHMoE layer.
 
diff --git a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/config.py b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/config.py
index b6b76e97be0f..66916db1c60b 100644
--- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/config.py
+++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/config.py
@@ -25,7 +25,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, Iterable
 
 # Number of trailing special tokens appended to every audio codebook.
 # Matches ``len(SpecialAudioToken)`` in
@@ -42,6 +42,75 @@
 SPECIAL_AUDIO_MASK: int = 4
 
 
+def normalize_nemotron_h_config(hf_config: Any) -> Any:
+    """Fill vLLM Nemotron-H runtime aliases on converted EasyMagpie configs."""
+    if not hasattr(hf_config, "rms_norm_eps") and hasattr(hf_config, "layer_norm_epsilon"):
+        hf_config.rms_norm_eps = hf_config.layer_norm_epsilon
+    return hf_config
+
+
+def derive_nemotron_h_hybrid_pattern_from_weight_keys(
+    keys: Iterable[str],
+    *,
+    layer_prefix: str = "decoder.layers.",
+) -> str | None:
+    """Infer a Nemotron-H hybrid pattern from converted checkpoint key names.
+
+    Older EasyMagpie conversion snippets could carry a stale
+    ``hybrid_override_pattern``. The state dict is the authoritative source for
+    whether a layer is Mamba (``M``), attention (``*``), dense MLP (``-``), or
+    routed MoE (``E``).
+    """
+    layer_kinds: dict[int, set[str]] = {}
+    for key in keys:
+        if not key.startswith(layer_prefix):
+            continue
+        rest = key[len(layer_prefix) :]
+        layer_str, sep, tail = rest.partition(".")
+        if not sep:
+            continue
+        try:
+            layer_idx = int(layer_str)
+        except ValueError:
+            continue
+        if not tail.startswith("mixer."):
+            layer_kinds.setdefault(layer_idx, set())
+            continue
+
+        mixer_key = tail[len("mixer.") :]
+        layer_kind = layer_kinds.setdefault(layer_idx, set())
+        if mixer_key.startswith(("experts.", "shared_experts.", "gate.")):
+            layer_kind.add("E")
+        elif mixer_key.startswith(("A_log", "D", "conv1d.", "dt_bias", "in_proj.", "out_proj.", "norm.")):
+            layer_kind.add("M")
+        elif mixer_key.startswith(("q_proj.", "k_proj.", "v_proj.", "o_proj.", "qkv_proj.")):
+            layer_kind.add("*")
+        elif mixer_key.startswith(("down_proj.", "gate_proj.", "up_proj.")):
+            layer_kind.add("-")
+
+    if not layer_kinds:
+        return None
+
+    max_layer = max(layer_kinds)
+    if any(idx not in layer_kinds for idx in range(max_layer + 1)):
+        return None
+
+    chars: list[str] = []
+    for idx in range(max_layer + 1):
+        kinds = layer_kinds[idx]
+        if "E" in kinds:
+            chars.append("E")
+        elif "M" in kinds:
+            chars.append("M")
+        elif "*" in kinds:
+            chars.append("*")
+        elif "-" in kinds:
+            chars.append("-")
+        else:
+            return None
+    return "".join(chars)
+
+
 @dataclass
 class EasyMagpieOmniArch:
     """Static architecture description for an EasyMagpieTTS checkpoint.
diff --git a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/easymagpie.py b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/easymagpie.py
index fbb589b21dc9..84f47e44c5b9 100644
--- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/easymagpie.py
+++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/easymagpie.py
@@ -44,19 +44,9 @@
 
 Per-request I/O (via ``additional_information``):
 
-* speaker context audio (prefill only) — either of:
-
-  * ``speaker_id`` — a short string naming a *known* speaker whose
-    ``(T_audio, embedding_dim)`` context-audio embedding lives in the checkpoint
-    as ``speaker_embeddings/<id>.pt``. It is loaded from disk on first use and
-    folded into the cached assembled prefill context (see
-    :meth:`_build_prefill_embeds`), so repeat requests carry just the id — no
-    per-request H2D and no serialization of a multi-hundred-KB tensor through the
-    engine. This is the recommended path for serving a fixed speaker set.
-  * ``speaker_embedding`` — a ``(T_audio, embedding_dim)`` tensor for custom /
-    one-off voices, copied H2D on each prefill.
-
-  ``preprocess`` assembles the full prefill context embedding itself as
+* ``speaker_embedding`` (prefill only) — ``(T_audio, embedding_dim)``
+  speaker-encoded context-audio embedding. ``preprocess`` assembles the full
+  prefill context embedding itself as
   ``[task_embedding | speaker_embedding | context_text_embedded]``, so the
   caller only does the speaker-encoder math and passes plain context text (the
   model tokenizes + embeds it and prepends the per-mode service token).
@@ -106,29 +96,66 @@
 from __future__ import annotations
 
 import bisect
+import hashlib
+import inspect
+import os
+import re
 from collections.abc import Callable, Iterable
+from pathlib import Path
 from typing import Any, Optional
 
 import torch
 from torch import nn
+from vllm import envs
 from vllm.compilation.backends import set_model_tag
 from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.forward_context import BatchDescriptor, get_forward_context
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import (
-    HasInnerState,
-    IsHybrid,
-    SupportsMambaPrefixCaching,
-)
+from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid
+
+try:
+    from vllm.model_executor.models.interfaces import SupportsMambaPrefixCaching
+except ImportError:
+
+    class SupportsMambaPrefixCaching:  # type: ignore[no-redef]
+        """Compatibility marker for vLLM builds that do not expose the mixin."""
+
+        pass
 from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM, NemotronHModel
+try:
+    from vllm.model_executor.models.mamba_cache import MambaCacheManager, MambaCacheParams
+
+    _HAS_VLLM_MAMBA_CACHE_MANAGER = True
+except ImportError:
+    _HAS_VLLM_MAMBA_CACHE_MANAGER = False
+
+    class MambaCacheParams:  # type: ignore[no-redef]
+        """Small compatibility container for vLLM builds without mamba_cache."""
+
+        def __init__(self, conv_state: torch.Tensor, ssm_state: torch.Tensor, state_indices_tensor: torch.Tensor):
+            self.conv_state = conv_state
+            self.ssm_state = ssm_state
+            self.conv_states = conv_state
+            self.ssm_states = ssm_state
+            self.state_indices_tensor = state_indices_tensor
+
+    class MambaCacheManager:  # type: ignore[no-redef]
+        pass
 from vllm.model_executor.models.utils import maybe_prefix
 from vllm.sequence import IntermediateTensors
+try:
+    from vllm.utils import LayerBlockType
+except ImportError:
+
+    class LayerBlockType:  # type: ignore[no-redef]
+        mamba = "mamba"
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 
 from easymagpie_vllm_omni.backbone_patches import (
     patch_mamba_streaming_decode,
     patch_moe_routed_scale,
+    patch_nemotron_h_moe_layer,
     patch_silu_shared_experts,
 )
 from easymagpie_vllm_omni.config import EasyMagpieOmniArch
@@ -136,6 +163,14 @@
 
 logger = init_logger(__name__)
 
+_MOE_EXPERT_WEIGHT_RE = re.compile(
+    r"^layers\.(?P<layer>\d+)\.mixer\.experts\.(?P<expert>\d+)\."
+    r"(?P<projection>up_proj|down_proj)\.weight$"
+)
+_QKV_WEIGHT_RE = re.compile(
+    r"^layers\.(?P<layer>\d+)\.mixer\.(?P<projection>q_proj|k_proj|v_proj)\.weight$"
+)
+
 # Placeholder token id stuffed into the per-step ``input_ids`` returned by
 # ``preprocess`` — the model never consumes ``input_ids`` (decode behaviour is
 # driven by the per-token buffers), and ``compute_logits`` returns
@@ -145,6 +180,17 @@
 # Context text used when the request omits ``context_text``
 _DEFAULT_CONTEXT_TEXT = "[EN]"
 
+
+def _apply_optional_nemotron_h_weight_mapper(
+    weights: Iterable[tuple[str, torch.Tensor]],
+) -> Iterable[tuple[str, torch.Tensor]]:
+    mapper = getattr(NemotronHForCausalLM, "hf_to_vllm_mapper", None)
+    apply = getattr(mapper, "apply", None)
+    if apply is None:
+        return weights
+    return apply(weights)
+
+
 # This class is not wrapped in ``@support_torch_compile``: the Nemotron-H
 # backbone and :class:`EasyMagpieCodePredictor` each manage their own
 # ``torch.compile`` / CUDA-graph capture internally, so the outer ``forward``
@@ -167,7 +213,10 @@ class EasyMagpieTTSForConditionalGeneration(
     # expects these as class attributes.
     get_mamba_state_dtype_from_config = NemotronHForCausalLM.get_mamba_state_dtype_from_config
     get_mamba_state_shape_from_config = NemotronHForCausalLM.get_mamba_state_shape_from_config
-    get_mamba_state_copy_func = NemotronHForCausalLM.get_mamba_state_copy_func
+    _get_mamba_state_copy_func = getattr(NemotronHForCausalLM, "get_mamba_state_copy_func", None)
+    if _get_mamba_state_copy_func is not None:
+        get_mamba_state_copy_func = _get_mamba_state_copy_func
+    del _get_mamba_state_copy_func
 
     # Omni runner hooks.
     has_preprocess: bool = True
@@ -186,6 +235,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         hf_config = vllm_config.model_config.hf_config
         self.hf_config = hf_config
         self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
         self.arch = EasyMagpieOmniArch.from_hf_config(hf_config)
         self.model_path = vllm_config.model_config.model
 
@@ -195,10 +245,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.num_codebooks = arch.num_stacked_codebooks
 
         # ── Backbone (reused vLLM Nemotron-H LM; fed via inputs_embeds) ──
+        patch_nemotron_h_moe_layer()
         self.backbone = NemotronHModel(
             vllm_config=vllm_config,
             prefix=maybe_prefix(prefix, "backbone"),
         )
+        self._backbone_accepts_mamba_cache_params = self._detect_backbone_mamba_cache_param()
+        self.mamba_cache: Optional[MambaCacheManager] = None
         # The checkpoint was trained with mlp_hidden_act=silu but vLLM's
         # NemotronHMLP hard-codes ReLU² in shared_experts. Restore SiLU (no-op
         # when the backbone has no MoE layers).
@@ -227,6 +280,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         # at conversion time and fed additively on every decode step.
         text_vocab_size = int(getattr(hf_config, "text_vocab_size", getattr(hf_config, "vocab_size", 0)))
         self.text_embedding = nn.Embedding(text_vocab_size, self.embedding_dim)
+        # PyTorch can use a distinct text-embedding path for context text on
+        # legacy checkpoints. Most current checkpoints use the target-text table
+        # for context too, so keep that memory-neutral unless config.json
+        # explicitly requests a separate baked context table.
+        if bool(getattr(hf_config, "context_text_embedding_distinct", False)):
+            self.context_text_embedding = nn.Embedding(text_vocab_size, self.embedding_dim)
+        else:
+            self.context_text_embedding = self.text_embedding
 
         # Text-stream EOS id — the last-but-one row of the text vocab, matching
         # the reference ``EasyMagpieTTSInferenceModel.eos_id = num_tokens - 2``.
@@ -274,6 +335,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         dtype = vllm_config.model_config.dtype
         # Combined per-token input embedding fed into the backbone.
         self._combined_embeddings = torch.zeros(max_num_tokens, self.embedding_dim, dtype=dtype)
+        self._decode_offsets = torch.full((max_num_tokens,), -1, dtype=torch.long)
         # Per-token decode inputs assembled by ``preprocess``.
         self._dec_text_tokens = torch.zeros(max_num_tokens, dtype=torch.long)
         self._dec_text_mask = torch.zeros(max_num_tokens, dtype=torch.long)
@@ -286,6 +348,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             self._dec_phoneme_valid = torch.zeros(max_num_tokens, dtype=torch.long)
 
         self._out_codes = torch.zeros(max_num_tokens, self.num_codebooks, dtype=torch.long)
+        self._out_code_logprobs = torch.zeros(max_num_tokens, self.num_codebooks, dtype=torch.float32)
+        self._out_code_sampling_logprobs = torch.zeros(max_num_tokens, self.num_codebooks, dtype=torch.float32)
+        self._out_frame_logprobs = torch.zeros(max_num_tokens, dtype=torch.float32)
+        self._last_output_row_indices: Optional[torch.Tensor] = None
 
         # ── Audio-EOS → engine stop ─────────────────────────────────────
         # The model signals end-of-speech inside the audio codebooks.
@@ -301,19 +367,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         # slice of ``token_stop`` based on ``logit_idx`` that can be used in
         # ``compute_logits``
         self._sample_stop = torch.zeros(max_num_tokens, dtype=torch.bool)
-
-        # ── Assembled prefill context embeddings (the only context cache) ──
-        # ``preprocess`` runs on the host, once per request, serially on the
-        # runner's critical path, so per-request speaker-tensor transfer + the
-        # tokenize/embed/cat dominate TTFT under concurrency. Cache the *whole*
-        # assembled context ``[task | speaker | context_text]`` per
-        # ``(task_mode_id, speaker_id, context_text, device)`` (see
-        # :meth:`_build_prefill_embeds`): for a known speaker it is identical on
-        # every request, so the cache subsumes a separate speaker-embedding table
-        # — the speaker ``.pt`` is read from disk only on the (first) cache miss
-        # for that combo (see :meth:`_load_known_speaker_embedding`), then never
-        # again. Custom raw-tensor voices are one-off and skip the cache.
-        self._prefill_cache: dict[tuple, torch.Tensor] = {}
+        self.min_content_audio_frames_before_eos = max(
+            0,
+            int(os.environ.get("EASYMAGPIE_MIN_CONTENT_AUDIO_FRAMES_BEFORE_EOS", "0") or 0),
+        )
 
     # ------------------------------------------------------------------
     # Embedding helpers
@@ -348,6 +405,137 @@ def _embed_phoneme(self, phoneme_tokens: torch.Tensor) -> torch.Tensor:
     # Decode-token dispatch (which positions need the local transformer)
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _detect_backbone_mamba_cache_param() -> bool:
+        try:
+            return "mamba_cache_params" in inspect.signature(NemotronHModel.forward).parameters
+        except Exception:
+            return False
+
+    def _ensure_v0_mamba_cache(self) -> Optional[MambaCacheManager]:
+        """Create the vLLM v0 Mamba cache lazily, matching NemotronHForCausalLM."""
+        if not _HAS_VLLM_MAMBA_CACHE_MANAGER:
+            return None
+        if getattr(self, "mamba_cache", None) is not None:
+            return self.mamba_cache
+        if not hasattr(self, "vllm_config") or not hasattr(self, "model_config"):
+            return None
+
+        num_mamba_layers = self.model_config.get_num_layers_by_block_type(
+            self.vllm_config.parallel_config,
+            LayerBlockType.mamba,
+        )
+        mamba_state_shape = self.get_mamba_state_shape_from_config(
+            self.vllm_config,
+            use_v1=False,
+        )
+        mamba_state_dtype = self.get_mamba_state_dtype_from_config(self.vllm_config)
+        self.mamba_cache = MambaCacheManager(
+            self.vllm_config,
+            num_mamba_layers,
+            *mamba_state_shape,
+            *mamba_state_dtype,
+        )
+        return self.mamba_cache
+
+    @staticmethod
+    def _metadata_token_count(attn_metadata: Any) -> Optional[int]:
+        if attn_metadata is None:
+            return None
+        metas = list(attn_metadata.values()) if isinstance(attn_metadata, dict) else [attn_metadata]
+        for metadata in metas:
+            num_prefills = getattr(metadata, "num_prefills", None)
+            num_decode_tokens = getattr(metadata, "num_decode_tokens", None)
+            if num_prefills is not None and num_decode_tokens is not None:
+                return int(num_prefills or 0) + int(num_decode_tokens or 0)
+
+            num_prefill_tokens = getattr(metadata, "num_prefill_tokens", None)
+            if num_prefill_tokens is not None and num_decode_tokens is not None:
+                return int(num_prefill_tokens or 0) + int(num_decode_tokens or 0)
+
+            query_start_loc = getattr(metadata, "query_start_loc", None)
+            if query_start_loc is not None:
+                try:
+                    if int(query_start_loc.numel()) > 0:
+                        return int(query_start_loc[-1].item())
+                except Exception:
+                    pass
+        return None
+
+    @staticmethod
+    def _mamba_cache_batch_size_from_metadata(attn_metadata: Any) -> Optional[int]:
+        for metadata in EasyMagpieTTSForConditionalGeneration._iter_metadata(attn_metadata):
+            state_indices = getattr(metadata, "state_indices_tensor", None)
+            if isinstance(state_indices, torch.Tensor) and state_indices.numel() > 0:
+                return int(state_indices.numel())
+            block_table = getattr(metadata, "block_table_tensor", None)
+            if isinstance(block_table, torch.Tensor) and block_table.numel() > 0:
+                return int(block_table.shape[0])
+            num_decodes = getattr(metadata, "num_decode_tokens", None)
+            num_prefills = getattr(metadata, "num_prefills", None)
+            if num_decodes is not None and num_prefills is not None:
+                return int(num_decodes or 0) + int(num_prefills or 0)
+        return None
+
+    def _profile_mamba_cache_params_from_forward_context(
+        self,
+        batch_size: Optional[int] = None,
+    ) -> Optional[MambaCacheParams]:
+        mamba_cache = getattr(self, "mamba_cache", None)
+        if mamba_cache is None:
+            return None
+        if batch_size is None:
+            try:
+                batch_size = self._metadata_token_count(get_forward_context().attn_metadata)
+            except Exception:
+                batch_size = None
+        if batch_size is None:
+            return None
+
+        cache_tensors, state_indices_tensor = mamba_cache.get_seqlen_agnostic_capture_inputs(int(batch_size))
+        return MambaCacheParams(
+            cache_tensors[0],
+            cache_tensors[1],
+            state_indices_tensor,
+        )
+
+    def _current_mamba_cache_params(
+        self,
+        *,
+        mamba_cache_params: Any = None,
+        easymagpie_mamba_cache_batch_size: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Any:
+        if mamba_cache_params is not None:
+            return mamba_cache_params
+
+        if kwargs.get("request_ids_to_seq_ids") is not None and kwargs.get("finished_requests_ids") is not None:
+            mamba_cache = self._ensure_v0_mamba_cache()
+            current_run_tensors = getattr(mamba_cache, "current_run_tensors", None) if mamba_cache is not None else None
+            if current_run_tensors is not None:
+                try:
+                    return current_run_tensors(**kwargs)
+                except Exception:
+                    logger.exception("Failed to build EasyMagpie Mamba cache params from request ids")
+
+        if easymagpie_mamba_cache_batch_size is not None:
+            return self._profile_mamba_cache_params_from_forward_context(easymagpie_mamba_cache_batch_size)
+
+        try:
+            use_v1 = bool(getattr(envs, "VLLM_USE_V1", False))
+        except Exception:
+            use_v1 = False
+        if use_v1:
+            return self._profile_mamba_cache_params_from_forward_context(easymagpie_mamba_cache_batch_size)
+
+        mamba_cache = self._ensure_v0_mamba_cache()
+        if mamba_cache is None:
+            return self._profile_mamba_cache_params_from_forward_context(easymagpie_mamba_cache_batch_size)
+        current_run_tensors = getattr(mamba_cache, "current_run_tensors", None)
+        if current_run_tensors is None:
+            return self._profile_mamba_cache_params_from_forward_context(easymagpie_mamba_cache_batch_size)
+        return current_run_tensors(**kwargs)
+
     @staticmethod
     def _select_query_layout(attn_metadata):
         """Return ``(max_query_len, query_start_loc)`` from heterogeneous metadata.
@@ -401,6 +589,17 @@ def _get_decode_idxs(self):
         if attn_metadata is None:
             return None, 0
 
+        metas = list(attn_metadata.values()) if isinstance(attn_metadata, dict) else [attn_metadata]
+        for metadata in metas:
+            num_prefills = getattr(metadata, "num_prefills", None)
+            num_decode_tokens = getattr(metadata, "num_decode_tokens", None)
+            num_decodes = getattr(metadata, "num_decodes", None)
+            if num_prefills is None or (num_decode_tokens is None and num_decodes is None):
+                continue
+            decode_count = num_decode_tokens if num_decode_tokens is not None else num_decodes
+            if int(num_prefills or 0) > 0 and int(decode_count or 0) == 0:
+                return torch.empty(0, dtype=torch.long, device=self._combined_embeddings.device), 0
+
         max_query_len, start_loc = self._select_query_layout(attn_metadata)
 
         # Decode-only batch (or layout unavailable) -> run the LT on every token.
@@ -424,6 +623,14 @@ def _get_decode_idxs(self):
             )
         return decode_token_indices, num_requests
 
+    @staticmethod
+    def _iter_metadata(value: Any) -> Iterable[Any]:
+        if isinstance(value, dict):
+            for child in value.values():
+                yield from EasyMagpieTTSForConditionalGeneration._iter_metadata(child)
+        elif value is not None:
+            yield value
+
     # ------------------------------------------------------------------
     # forward
     # ------------------------------------------------------------------
@@ -454,7 +661,14 @@ def forward(
         # Reset per-token stop flags for this step (so prefill / warm-up rows stay
         # "continue"); decode positions get set below by :meth:`_flag_audio_eos`.
         self._token_stop[:num_tokens].zero_()
+        self._out_code_logprobs[:num_tokens].zero_()
+        self._out_code_sampling_logprobs[:num_tokens].zero_()
+        self._out_frame_logprobs[:num_tokens].zero_()
         logits_index = kwargs.get("logits_index")
+        if isinstance(logits_index, torch.Tensor) and logits_index.numel() > 0:
+            self._last_output_row_indices = logits_index.detach().reshape(-1).to(dtype=torch.long)
+        else:
+            self._last_output_row_indices = None
 
         decode_idx, num_req = self._get_decode_idxs()
 
@@ -466,17 +680,26 @@ def forward(
             valid = decode_idx[:num_req]
             self._assemble_decode_embeddings(combined, valid)
 
-        hidden_states = self.backbone(
-            input_ids=input_ids,
-            positions=positions,
-            intermediate_tensors=intermediate_tensors,
-            inputs_embeds=combined,
-        )
+        backbone_kwargs = {
+            "input_ids": input_ids,
+            "positions": positions,
+            "intermediate_tensors": intermediate_tensors,
+            "inputs_embeds": combined,
+        }
+        mamba_cache_params = None
+        if self._backbone_accepts_mamba_cache_params:
+            mamba_cache_params = self._current_mamba_cache_params(**kwargs)
+            backbone_kwargs["mamba_cache_params"] = mamba_cache_params
+
+        hidden_states = self.backbone(**backbone_kwargs)
 
         # Sample codes (local transformer) only where needed.
         if decode_idx is None:
-            codes = self.code_predictor.generate_codes(hidden_states)
+            codes, code_logprobs, sampling_logprobs = self.code_predictor.generate_codes_with_logprobs(hidden_states)
             self._out_codes[:num_tokens].copy_(codes)
+            self._out_code_logprobs[:num_tokens].copy_(code_logprobs)
+            self._out_code_sampling_logprobs[:num_tokens].copy_(sampling_logprobs)
+            self._out_frame_logprobs[:num_tokens].copy_(code_logprobs.sum(dim=-1))
             self._flag_audio_eos(codes, slice(0, num_tokens))
             if self.has_phoneme:
                 self._predict_phonemes(hidden_states, slice(0, num_tokens))
@@ -484,10 +707,15 @@ def forward(
             ctx = get_forward_context()
             orig_bd = ctx.batch_descriptor
             ctx.batch_descriptor = BatchDescriptor(num_tokens=decode_idx.shape[0])
-            codes = self.code_predictor.generate_codes(hidden_states[decode_idx])
+            codes, code_logprobs, sampling_logprobs = self.code_predictor.generate_codes_with_logprobs(
+                hidden_states[decode_idx]
+            )
             ctx.batch_descriptor = orig_bd
             valid = decode_idx[:num_req]
             self._out_codes[valid] = codes[:num_req]
+            self._out_code_logprobs[valid] = code_logprobs[:num_req]
+            self._out_code_sampling_logprobs[valid] = sampling_logprobs[:num_req]
+            self._out_frame_logprobs[valid] = code_logprobs[:num_req].sum(dim=-1)
             self._flag_audio_eos(codes[:num_req], valid)
             if self.has_phoneme:
                 self._predict_phonemes(hidden_states, valid)
@@ -510,26 +738,40 @@ def _flag_audio_eos(self, codes: torch.Tensor, idx) -> None:
         i.e. checks if EOS is emited without sampling. Skip for now.
         """
         eos = (codes == self.audio_eos_id).any(dim=1) & (self._dec_audio_valid[idx] == 1)
+        min_content_frames = int(getattr(self, "min_content_audio_frames_before_eos", 0) or 0)
+        if min_content_frames > 0:
+            decode_offset = self._decode_offsets[idx].to(device=eos.device)
+            previous_content_frames = decode_offset - int(self.speech_delay)
+            eos = eos & (previous_content_frames >= min_content_frames)
         self._token_stop[idx] = eos
 
     def _assemble_decode_embeddings(self, combined: torch.Tensor, idx) -> None:
         """Add ``text + phoneme + audio`` embeddings into ``combined`` at ``idx``."""
+        # Mixed prefill/decode batches carry real prefill rows in ``inputs_embeds``.
+        # Decode rows must start empty; otherwise a stale context/prefill vector
+        # is added to the streaming text/phoneme/audio components. Build a fresh
+        # tensor and assign it back because ``idx`` is often a LongTensor; in
+        # that case ``combined[idx].zero_()`` would only clear an advanced-index
+        # copy.
+        assembled = torch.zeros_like(combined[idx])
+
         # Audio: previous-frame codes (gated by validity).
         audio_codes = self._dec_audio_codes[idx]
         audio_emb = self.code_predictor.embed_audio_frame(audio_codes)
         audio_emb = audio_emb * self._dec_audio_valid[idx].unsqueeze(-1).to(audio_emb.dtype)
-        combined[idx] += audio_emb
+        assembled += audio_emb
 
         # Text: current subword token (gated by validity).
         text_emb = self.text_embedding(self._dec_text_tokens[idx])
         text_emb = text_emb * self._dec_text_mask[idx].unsqueeze(-1).to(text_emb.dtype)
-        combined[idx] += text_emb
+        assembled += text_emb
 
         # Phoneme: previous predicted phoneme (gated by validity).
         if self.has_phoneme:
             phon_emb = self._embed_phoneme(self._dec_phoneme_tokens[idx])
             phon_emb = phon_emb * self._dec_phoneme_valid[idx].unsqueeze(-1).to(phon_emb.dtype)
-            combined[idx] += phon_emb
+            assembled += phon_emb
+        combined[idx] = assembled
 
     @torch.no_grad()
     def _predict_phonemes(self, hidden_states: torch.Tensor, idx) -> None:
@@ -598,16 +840,199 @@ def make_omni_output(self, model_outputs, **_: Any) -> OmniOutput:
             return model_outputs
         hidden = model_outputs
         num_tokens = int(hidden.shape[0])
-        audio_codes = self._out_codes[:num_tokens].clone()
+        row_indices = self._row_indices_for_hidden(
+            hidden,
+            int(self._out_codes.shape[0]),
+            self._out_codes.device,
+            explicit_row_indices=getattr(self, "_last_output_row_indices", None),
+        )
+        audio_codes = self._out_codes.index_select(0, row_indices).clone()
+        audio_code_logprobs = self._out_code_logprobs.index_select(0, row_indices).clone()
+        audio_code_sampling_logprobs = self._out_code_sampling_logprobs.index_select(0, row_indices).clone()
+        audio_frame_logprobs = self._out_frame_logprobs.index_select(0, row_indices).clone()
+        multimodal_outputs: dict[str, torch.Tensor] = {
+            "audio_codes": audio_codes,
+            "audio_codes_feedback": audio_codes,
+            "audio_code_logprobs": audio_code_logprobs,
+            "audio_code_sampling_logprobs": audio_code_sampling_logprobs,
+            "audio_frame_logprobs": audio_frame_logprobs,
+        }
+        if self.has_phoneme:
+            multimodal_outputs["phoneme_tokens_feedback"] = self._dec_phoneme_tokens.index_select(
+                0,
+                row_indices.to(self._dec_phoneme_tokens.device),
+            ).clone()
+        multimodal_outputs = {
+            name: tensor.contiguous() if isinstance(tensor, torch.Tensor) else tensor
+            for name, tensor in multimodal_outputs.items()
+        }
         return OmniOutput(
             text_hidden_states=hidden,
-            multimodal_outputs={"audio_codes": audio_codes},
+            multimodal_outputs=multimodal_outputs,
         )
 
+    @staticmethod
+    def _row_indices_for_hidden(
+        hidden_states: torch.Tensor,
+        row_count: int,
+        device: torch.device,
+        *,
+        explicit_row_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Return full-batch row indices matching a request-local hidden slice."""
+        num_rows = int(hidden_states.shape[0])
+        if num_rows <= 0 or row_count <= 0:
+            return torch.empty(0, dtype=torch.long, device=device)
+
+        if isinstance(explicit_row_indices, torch.Tensor) and int(explicit_row_indices.numel()) == num_rows:
+            return explicit_row_indices.detach().reshape(-1).to(device=device, dtype=torch.long)
+
+        stride0 = hidden_states.stride(0) or 1
+        start = int(hidden_states.storage_offset()) // int(stride0)
+        if 0 <= start and start + num_rows <= row_count:
+            return torch.arange(start, start + num_rows, dtype=torch.long, device=device)
+
+        return torch.arange(0, min(num_rows, row_count), dtype=torch.long, device=device)
+
+    @staticmethod
+    def _last_request_row_index(
+        hidden_states: torch.Tensor,
+        row_tensor: torch.Tensor,
+        explicit_row_indices: Optional[torch.Tensor] = None,
+    ) -> int:
+        """Map a request slice to either request-local or full-batch row coordinates."""
+        local_last = int(hidden_states.shape[0]) - 1
+        if int(row_tensor.shape[0]) == int(hidden_states.shape[0]):
+            return local_last
+
+        if (
+            isinstance(explicit_row_indices, torch.Tensor)
+            and int(explicit_row_indices.numel()) == int(hidden_states.shape[0])
+        ):
+            explicit_last = int(explicit_row_indices.detach().reshape(-1)[local_last].item())
+            if 0 <= explicit_last < int(row_tensor.shape[0]):
+                return explicit_last
+
+        stride0 = hidden_states.stride(0) or 1
+        storage_last = hidden_states.storage_offset() // stride0 + local_last
+        if 0 <= storage_last < int(row_tensor.shape[0]):
+            return int(storage_last)
+
+        return min(max(local_last, 0), int(row_tensor.shape[0]) - 1)
+
     # ------------------------------------------------------------------
     # preprocess / postprocess
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _first_str(value: Any) -> str:
+        """Return the first element of a list-wrapped scalar, or the scalar itself, as a string."""
+        if isinstance(value, list):
+            return str(value[0]) if value else ""
+        if value is None:
+            return ""
+        return str(value)
+
+    @staticmethod
+    def _coerce_opt_int(value: Any) -> Optional[int]:
+        """Best-effort extract a single int from a scalar / list / tensor / str.
+
+        Used to read a per-step streamed ``text_token`` out of the request's
+        ``additional_information`` (which may wrap the id as a list, a 1-element
+        tensor, or a string depending on how the caller / transport packed it).
+        Returns ``None`` when no usable integer is present.
+        """
+        if value is None:
+            return None
+        if isinstance(value, bool):  # bool is an int subclass — handle explicitly.
+            return int(value)
+        if isinstance(value, int):
+            return value
+        if isinstance(value, float):
+            return int(value)
+        if isinstance(value, torch.Tensor):
+            return int(value.reshape(-1)[0].item()) if value.numel() > 0 else None
+        if isinstance(value, (list, tuple)):
+            return EasyMagpieTTSForConditionalGeneration._coerce_opt_int(value[0]) if value else None
+        if isinstance(value, str):
+            try:
+                return int(value.strip())
+            except ValueError:
+                return None
+        return None
+
+    @staticmethod
+    def _coerce_int_list(value: Any) -> Optional[list[int]]:
+        """Best-effort normalize a caller-provided 1-D token list."""
+        if value is None:
+            return None
+        if isinstance(value, torch.Tensor):
+            return [int(item) for item in value.detach().cpu().reshape(-1).tolist()]
+        if isinstance(value, (list, tuple)):
+            out: list[int] = []
+            for item in value:
+                parsed = EasyMagpieTTSForConditionalGeneration._coerce_opt_int(item)
+                if parsed is not None:
+                    out.append(int(parsed))
+            return out
+        parsed = EasyMagpieTTSForConditionalGeneration._coerce_opt_int(value)
+        return [int(parsed)] if parsed is not None else None
+
+    @staticmethod
+    def _coerce_int_rows(value: Any, row_width: int) -> Optional[list[list[int]]]:
+        """Best-effort normalize caller-provided 2-D token rows."""
+        if value is None:
+            return None
+        width = max(1, int(row_width))
+        if isinstance(value, torch.Tensor):
+            tensor = value.detach().cpu().long()
+            if tensor.ndim <= 1:
+                return [EasyMagpieTTSForConditionalGeneration._coerce_int_list(tensor) or []]
+            return [[int(item) for item in row[:width].reshape(-1).tolist()] for row in tensor.reshape(-1, tensor.shape[-1])]
+        if isinstance(value, (list, tuple)):
+            if not value:
+                return []
+            first = value[0]
+            if isinstance(first, (list, tuple, torch.Tensor)):
+                rows: list[list[int]] = []
+                for row in value:
+                    parsed = EasyMagpieTTSForConditionalGeneration._coerce_int_list(row)
+                    rows.append((parsed or [])[:width])
+                return rows
+            parsed = EasyMagpieTTSForConditionalGeneration._coerce_int_list(value)
+            return [(parsed or [])[:width]]
+        parsed = EasyMagpieTTSForConditionalGeneration._coerce_int_list(value)
+        return [(parsed or [])[:width]] if parsed is not None else None
+
+    def _clear_runtime_rows(self, start: int, end: int) -> None:
+        """Clear mutable per-row state before reusing prefill/decode slots."""
+        row = slice(max(0, int(start)), max(0, int(end)))
+        zero_names = (
+            "_dec_text_tokens",
+            "_dec_text_mask",
+            "_dec_audio_codes",
+            "_dec_audio_valid",
+            "_dec_phoneme_tokens",
+            "_dec_phoneme_valid",
+            "_out_codes",
+            "_out_code_logprobs",
+            "_out_code_sampling_logprobs",
+            "_out_frame_logprobs",
+            "_token_stop",
+            "_sample_stop",
+        )
+        minus_one_names = (
+            "_decode_offsets",
+        )
+        for name in zero_names:
+            value = getattr(self, name, None)
+            if isinstance(value, torch.Tensor):
+                value[row].zero_()
+        for name in minus_one_names:
+            value = getattr(self, name, None)
+            if isinstance(value, torch.Tensor):
+                value[row].fill_(-1)
+
     def preprocess(
         self,
         input_ids: torch.Tensor,
@@ -653,7 +1078,9 @@ def _batch_slot_offset(input_ids_view: torch.Tensor, fallback: int) -> int:
         The runner passes ``input_ids = input_ids_buffer[s:e]``
         """
         if input_ids_view.dim() == 1 and input_ids_view.is_contiguous():
-            return int(input_ids_view.storage_offset())
+            offset = int(input_ids_view.storage_offset())
+            if offset > 0:
+                return offset
         return int(fallback)
 
     def _preprocess_prefill(
@@ -674,19 +1101,20 @@ def _preprocess_prefill(
 
         offset = int(info_dict.get("prefill_offset", 0) or 0)
         total = int(prefill_embeds.shape[0])
-        take = prefill_embeds[offset : offset + span_len]
-        # The prefill chunk must lie fully within the assembled context. Padding
-        # short chunks with zeros / a repeated last row is invalid: the backbone
-        # was never trained on padded context frames, so silently doing so would
-        # corrupt conditioning rather than fail loudly. This holds iff the caller
-        # sized ``prompt_token_ids`` to ``estimate_prompt_len(...)``.
-        assert int(take.shape[0]) == span_len, (
-            f"EasyMagpieTTS prefill chunk [{offset}:{offset + span_len}] is not fully covered by the "
-            f"assembled context embedding (length {total}). The caller must pass "
-            f"prompt_token_ids of length estimate_prompt_len(...) = "
-            f"[task?] + speaker_embedding.shape[0] + len(tokenize(context_text)); "
-            f"zero-padding the backbone context is invalid (the model was not trained on it)."
-        )
+        s = max(0, min(offset, total))
+        e = max(0, min(offset + span_len, total))
+        take = prefill_embeds[s:e]
+        if int(take.shape[0]) < span_len:
+            pad_n = span_len - int(take.shape[0])
+            pad_rows = (
+                take[-1:].expand(pad_n, -1)
+                if take.shape[0] > 0
+                else prefill_embeds.new_zeros(pad_n, prefill_embeds.shape[-1])
+            )
+            take = torch.cat([take, pad_rows], dim=0)
+
+        row_start = self._batch_slot_offset(input_ids, 0)
+        self._clear_runtime_rows(row_start, row_start + span_len)
 
         info_update = {
             "prefill_offset": offset + span_len,
@@ -701,8 +1129,11 @@ def _preprocess_prefill(
         # ``text_tokens`` is provided the request runs in **streaming-text mode**:
         # no list is baked, and :meth:`_preprocess_decode` instead reads one
         # subword id per step from the streamed ``additional_information.text_token``.
-        if not info_dict.get("text_tokens"):
-            text = info_dict.get("text")
+        text_tokens = self._coerce_int_list(info_dict.get("text_tokens"))
+        if text_tokens is not None:
+            info_update["text_tokens"] = text_tokens
+        else:
+            text = self._first_str(info_dict.get("text"))
             if text:
                 info_update["text_tokens"] = self._encode_text_stream(text)
         input_ids_out = torch.full_like(input_ids, _DUMMY_TOKEN_ID)
@@ -719,121 +1150,46 @@ def _build_prefill_embeds(
 
         from the per-request inputs:
 
-        * speaker context audio — either ``speaker_id`` (a known speaker whose
-          embedding is precomputed model state, see
-          :meth:`_resolve_speaker_embedding`) or, for custom / one-off voices, a
-          2-D ``(T_audio, embedding_dim)`` ``speaker_embedding`` tensor.
+        * ``speaker_embedding`` — the speaker-encoded context-audio embedding,
+          required as a 2-D ``(T_audio, embedding_dim)`` tensor.
         * ``context_text`` — a plain string (e.g. ``"[EN]"``); tokenized in-model
-          and embedded through the baked per-subword ``text_embedding`` table.
+          (see :meth:`_encode_context_text`) and embedded through the baked
+          per-subword ``text_embedding`` table.
         * ``task_mode_id`` — selects the per-mode task ("service token")
           embedding row; prepended only when the checkpoint has a task table.
 
-        Returns the full context embedding; the per-chunk slicing is done by
-        :meth:`_preprocess_prefill`.
-
-        For a known ``speaker_id`` the result is a pure function of
-        ``(task_mode_id, speaker_id, context_text)`` and is cached in
-        ``self._prefill_cache``, so the tokenize + embed + cat below run once per
-        distinct combo instead of on every request's prefill. The returned tensor
-        is only ever read (sliced) downstream, never mutated, so sharing the
-        cached instance is safe.
+        Returns the full context embedding; the per-chunk slicing/padding is done
+        by :meth:`_preprocess_prefill`.
         """
-        speaker_id = info_dict.get("speaker_id")
-        context_text = info_dict.get("context_text") or _DEFAULT_CONTEXT_TEXT
-        if self.task_embedding is not None:
-            task_mode_id = int(info_dict.get("task_mode_id", 0) or 0)
-            task_mode_id = max(0, min(task_mode_id, self.num_task_embeddings - 1))
-        else:
-            task_mode_id = 0
+        dtype = self._combined_embeddings.dtype
 
-        # Custom raw-tensor voices (no speaker_id) are one-off, so skip the cache.
-        cache_key = (task_mode_id, speaker_id, context_text, str(device)) if speaker_id else None
-        if cache_key is not None:
-            cached = self._prefill_cache.get(cache_key)
-            if cached is not None:
-                return cached
+        speaker_embedding = info_dict.get("speaker_embedding")
+        assert isinstance(speaker_embedding, torch.Tensor) and speaker_embedding.ndim == 2, (
+            "EasyMagpieTTS preprocess expects additional_information.speaker_embedding to be a 2-D "
+            "(T_audio, embedding_dim) tensor (the speaker-encoded context audio); "
+            f"got {type(speaker_embedding).__name__}"
+            + (f" with ndim={speaker_embedding.ndim}" if isinstance(speaker_embedding, torch.Tensor) else "")
+        )
 
-        dtype = self._combined_embeddings.dtype
         parts: list[torch.Tensor] = []
 
         # Task / "service token" embedding (prepended), when present.
         if self.task_embedding is not None:
+            task_mode_id = int(info_dict.get("task_mode_id", 0) or 0)
+            task_mode_id = max(0, min(task_mode_id, self.num_task_embeddings - 1))
             task_row = self.task_embedding(torch.tensor([task_mode_id], device=device, dtype=torch.long))
             parts.append(task_row.to(dtype))
 
-        # Speaker-encoded context audio (known-speaker state or custom tensor).
-        parts.append(self._resolve_speaker_embedding(device, info_dict))
+        # Speaker-encoded context audio.
+        parts.append(speaker_embedding.to(device=device, dtype=dtype))
 
         # Context text: tokenized in-model and embedded through the baked table.
+        context_text = self._first_str(info_dict.get("context_text")) or _DEFAULT_CONTEXT_TEXT
         ctx_ids = self._encode_context_text(context_text, device)
         if ctx_ids.numel() > 0:
-            parts.append(self.text_embedding(ctx_ids).to(dtype))
-
-        embeds = torch.cat(parts, dim=0)
-        if cache_key is not None:
-            self._prefill_cache[cache_key] = embeds
-        return embeds
-
-    def _resolve_speaker_embedding(self, device: torch.device, info_dict: dict[str, Any]) -> torch.Tensor:
-        """Return the speaker context-audio embedding on ``device`` in model dtype.
-
-        For a known ``speaker_id`` the embedding is read from disk by
-        :meth:`_load_known_speaker_embedding`; this only ever runs on a
-        prefill-cache miss (see :meth:`_build_prefill_embeds`), i.e. once per
-        ``(speaker_id, context_text, task)`` combo, so there is no separate
-        speaker-embedding table — the assembled prefill cache subsumes it. Falls
-        back to a raw ``speaker_embedding`` tensor (custom / one-off voice),
-        copied H2D here. Exactly one of the two must be supplied.
-        """
-        dtype = self._combined_embeddings.dtype
-        speaker_id = info_dict.get("speaker_id")
-        if speaker_id:
-            return self._load_known_speaker_embedding(speaker_id, device, dtype)
+            parts.append(self.context_text_embedding(ctx_ids).to(dtype))
 
-        speaker_embedding = info_dict.get("speaker_embedding")
-        assert isinstance(speaker_embedding, torch.Tensor) and speaker_embedding.ndim == 2, (
-            "EasyMagpieTTS preprocess expects additional_information.speaker_id (a known speaker) or "
-            "speaker_embedding as a 2-D (T_audio, embedding_dim) tensor (the speaker-encoded context "
-            f"audio); got speaker_embedding={type(speaker_embedding).__name__}"
-            + (f" with ndim={speaker_embedding.ndim}" if isinstance(speaker_embedding, torch.Tensor) else "")
-        )
-        return speaker_embedding.to(device=device, dtype=dtype)
-
-    def _load_known_speaker_embedding(
-        self, speaker_id: str, device: torch.device, dtype: torch.dtype
-    ) -> torch.Tensor:
-        """Read one known speaker's embedding from ``<model_path>/speaker_embeddings/<id>.pt``.
-
-        The file holds either a bare ``(T_audio, embedding_dim)`` tensor or a dict
-        with a ``speaker_encoding`` key (the converter/caller layout); it is moved
-        to ``device`` in model dtype. Called only on a prefill-cache miss, so each
-        known speaker is read at most once per ``(context_text, task)`` combo and
-        the result is then baked into ``self._prefill_cache``. Read from disk (not
-        via :meth:`load_weights`) so known speakers work even under
-        ``--load-format dummy``, which skips weight loading.
-        """
-        import glob
-        import os
-
-        spk_dir = os.path.join(self.model_path, "speaker_embeddings")
-        path = os.path.join(spk_dir, f"{speaker_id}.pt")
-        if not os.path.exists(path):
-            known = sorted(
-                os.path.splitext(os.path.basename(p))[0] for p in glob.glob(os.path.join(spk_dir, "*.pt"))
-            )
-            raise AssertionError(
-                f"EasyMagpieTTS preprocess got unknown speaker_id {speaker_id!r}; known speakers: {known}. "
-                "Register it under the checkpoint's speaker_embeddings/ dir, or pass a raw "
-                "speaker_embedding tensor for a custom voice."
-            )
-        loaded = torch.load(path, map_location="cpu")
-        embedding = loaded["speaker_encoding"] if isinstance(loaded, dict) else loaded
-        assert isinstance(embedding, torch.Tensor) and embedding.ndim == 2, (
-            f"EasyMagpieTTS: speaker embedding {path} must be a 2-D (T_audio, embedding_dim) tensor; "
-            f"got {type(embedding).__name__}"
-            + (f" with ndim={embedding.ndim}" if isinstance(embedding, torch.Tensor) else "")
-        )
-        return embedding.to(device=device, dtype=dtype)
+        return torch.cat(parts, dim=0)
 
     def _maybe_set_lt_sampling_params(self, info_dict: dict[str, Any]) -> None:
         """Apply per-request audio sampling params to the local transformer.
@@ -844,10 +1200,10 @@ def _maybe_set_lt_sampling_params(self, info_dict: dict[str, Any]) -> None:
         """
         temperature = info_dict.get("temperature")
         if temperature is not None:
-            self.code_predictor.temperature = float(temperature)
+            self.code_predictor.temperature = float(self._first_str(temperature) or 0.0)
         top_k = info_dict.get("top_k", info_dict.get("topk"))
         if top_k is not None:
-            self.code_predictor.top_k = int(top_k)
+            self.code_predictor.top_k = int(float(self._first_str(top_k) or 0))
 
     def _get_text_tokenizer(self):
         """Lazily load the context-text tokenizer from the model directory.
@@ -857,9 +1213,25 @@ def _get_text_tokenizer(self):
         on first use from ``model_path``.
         """
         if self._text_tokenizer is None:
-            from transformers import AutoTokenizer
+            import json
 
-            self._text_tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
+            from transformers import AutoTokenizer
+            from transformers import PreTrainedTokenizerFast
+
+            try:
+                self._text_tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
+            except ValueError as exc:
+                if "TokenizersBackend" not in str(exc):
+                    raise
+                tokenizer_config = json.loads((Path(self.model_path) / "tokenizer_config.json").read_text())
+                self._text_tokenizer = PreTrainedTokenizerFast(
+                    tokenizer_file=str(Path(self.model_path) / "tokenizer.json"),
+                    bos_token=tokenizer_config.get("bos_token"),
+                    eos_token=tokenizer_config.get("eos_token"),
+                    unk_token=tokenizer_config.get("unk_token"),
+                    clean_up_tokenization_spaces=bool(tokenizer_config.get("clean_up_tokenization_spaces", False)),
+                    model_max_length=int(tokenizer_config.get("model_max_length", int(1e30))),
+                )
         return self._text_tokenizer
 
     def _encode_context_text(self, context_text: str, device: torch.device) -> torch.Tensor:
@@ -895,19 +1267,13 @@ def estimate_prompt_len(
         context_text: str = _DEFAULT_CONTEXT_TEXT,
         has_task_embedding: bool = False,
     ) -> int:
-        """Length-only mirror of :meth:`_build_prefill_embeds` (custom voice).
+        """Length-only mirror of :meth:`_build_prefill_embeds`.
 
         The engine assembles the prefill context as
         ``[task_embedding? | speaker_embedding | context_text_embedded]``, so the
         caller must pass ``prompt_token_ids = [0] * estimate_prompt_len(...)`` for
         the placeholder length to match the assembled embedding length (otherwise
-        vLLM pads / truncates and quality drops). This is a pure function of
-        lengths, so it stays static — callable in the request-building process
-        without an engine instance.
-
-        For a **known speaker** the caller holds only a ``speaker_id`` (not the
-        tensor); use :meth:`get_prompt_len`, which loads the embedding from the
-        checkpoint dir and calls this method.
+        vLLM pads / truncates and quality drops).
 
         Args:
             speaker_embedding: ``(T_audio, embedding_dim)`` speaker-encoded
@@ -915,7 +1281,7 @@ def estimate_prompt_len(
             tokenize: callable turning ``context_text`` into its subword ids
                 (e.g. ``lambda t: tokenizer.encode(t)``) — must match the
                 tokenizer the engine loads from ``model_path``.
-            context_text: conditioning string (default ``"[EN]"``).
+            context_text: conditioning string (default ``"[NO TEXT CONTEXT]"``).
             has_task_embedding: whether the checkpoint prepends a task /
                 "service token" embedding (``num_task_embeddings > 0``).
         """
@@ -924,39 +1290,6 @@ def estimate_prompt_len(
         task_len = 1 if has_task_embedding else 0
         return task_len + t_audio + ctx_len
 
-    @classmethod
-    def get_prompt_len(cls, speaker_id: str, model_path: str, *, tokenize: Callable[[str], Iterable[int]]) -> int:
-        """Known-speaker convenience wrapper around :meth:`estimate_prompt_len`.
-
-        Resolves everything from the checkpoint dir so it cannot disagree with
-        what the engine actually uses: loads the speaker embedding from
-        ``speaker_embeddings/<speaker_id>.pt`` (the same file the engine reads in
-        :meth:`_load_known_speaker_embedding`), reads ``has_task_embedding`` from
-        ``config.json`` (``num_task_embeddings``),
-        and conditions on the fixed :data:`_DEFAULT_CONTEXT_TEXT`. Lets a caller
-        holding only a ``speaker_id`` size ``prompt_token_ids`` without an engine
-        instance (``context_text`` / ``has_task_embedding`` are intentionally not
-        params — they must match the precomputed checkpoint, not be overridden).
-        """
-        import json
-        import os
-
-        path = os.path.join(model_path, "speaker_embeddings", f"{speaker_id}.pt")
-        if not os.path.exists(path):
-            raise FileNotFoundError(f"EasyMagpieTTS: no speaker embedding {path} for speaker_id {speaker_id!r}")
-        loaded = torch.load(path, map_location="cpu")
-        speaker_embedding = loaded["speaker_encoding"] if isinstance(loaded, dict) else loaded
-
-        with open(os.path.join(model_path, "config.json")) as f:
-            num_task_embeddings = int(json.load(f).get("num_task_embeddings", 0))
-
-        return cls.estimate_prompt_len(
-            speaker_embedding,
-            tokenize=tokenize,
-            context_text=_DEFAULT_CONTEXT_TEXT,
-            has_task_embedding=num_task_embeddings > 0,
-        )
-
     def _preprocess_decode(
         self,
         input_ids: torch.Tensor,
@@ -966,47 +1299,45 @@ def _preprocess_decode(
     ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
         decode_offset = int(info_dict.get("decode_offset", 0) or 0)
         info_update: dict[str, Any] = {"decode_offset": decode_offset + 1}
+        self._clear_runtime_rows(start, start + 1)
+        self._decode_offsets[start] = decode_offset
 
         # ── Text channel ── (delay 0: one subword per step from step 0). The text
-        # stream leads the phoneme/audio streams by their respective delays. The
-        # model always consumes exactly one buffered subword id per decode step,
-        # indexed by ``decode_offset`` from a persistent ``text_tokens`` list. That
-        # list is populated by one of two mutually exclusive input modes:
+        # stream leads the phoneme/audio streams by their respective delays. Two
+        # mutually exclusive input modes are supported:
         #
         # * **Whole-text (non-streaming)** — the caller passed ``text`` whole at
         #   prefill; it was tokenized in-model and stashed as the ``text_tokens``
-        #   list (see :meth:`_preprocess_prefill`). No per-step ``text_token``
-        #   arrives, so the buffer never grows here.
-        # * **Streamed** — the caller did *not* pass ``text`` at prefill and instead
-        #   pushes subword ids during decode via ``additional_information`` under
-        #   ``text_token`` (always a ``list[int]``). Each chunk may carry a single id
-        #   (``[id]`` with ``max_tokens == 1``, one frame per chunk) or several ids at
-        #   once (``max_tokens == N``, so the engine free-runs N frames off one chunk —
-        #   fewer round-trips). Those ids are appended to ``text_tokens`` and consumed
-        #   one per step.
-        #
-        # In both modes, once the buffer is exhausted the channel is masked off
-        # (adds nothing) rather than repeating the last token, so the caller can keep
-        # pumping decode steps (passing an empty ``text_token`` list) while the audio
-        # tail finishes.
-        #
-        # Append-once: a streamed chunk's ``text_token`` payload stays identical on
-        # every decode step of that chunk, so we extend the buffer only when it has
-        # been fully consumed (``decode_offset`` caught up to its length). This is
-        # safe because the chunk's segment stops at ``max_tokens`` and ``preprocess``
-        # is not called again until a fresh chunk has replaced ``text_token`` — hence
-        # the caller must size each chunk's ``max_tokens`` to the number of ids it
-        # carries.
-        text_tokens = info_dict.get("text_tokens") or []
-        incoming = info_dict.get("text_token") or []
-        if incoming and decode_offset >= len(text_tokens):
-            text_tokens = text_tokens + [int(t) for t in incoming]
-            info_update["text_tokens"] = text_tokens
-        if decode_offset < len(text_tokens):
-            self._dec_text_tokens[start] = int(text_tokens[decode_offset])
-            self._dec_text_mask[start] = 1
+        #   list (see :meth:`_preprocess_prefill`). Step k consumes
+        #   ``text_tokens[k]`` (the list ends with the text-EOS id); once the
+        #   stream is exhausted the channel is masked off (adds nothing) rather
+        #   than repeating the last token.
+        # * **Streamed** — the caller did *not* pass ``text`` at prefill and
+        #   instead pushes one subword id per decode step via
+        #   ``additional_information`` under ``text_token`` (a single int / 1-elem
+        #   tensor; close the stream by pushing the text-EOS id as the last real
+        #   token). The model embeds that step's id and masks the channel off on
+        #   any step that carries no id (``text_token`` absent or ``< 0``), so the
+        #   caller can keep pumping decode steps after the text ends while the
+        #   audio tail finishes. Because each streamed chunk overwrites the
+        #   previous ``text_token`` in the per-request buffer, every step gets a
+        #   fresh value (or the caller's sentinel ``-1`` to mask).
+        text_tokens = info_dict.get("text_tokens")
+        if isinstance(text_tokens, list):
+            if decode_offset < len(text_tokens):
+                self._dec_text_tokens[start] = int(text_tokens[decode_offset])
+                self._dec_text_mask[start] = 1
+            else:
+                self._dec_text_tokens[start] = 0
+                self._dec_text_mask[start] = 0
         else:
-            self._dec_text_mask[start] = 0
+            streamed_id = self._coerce_opt_int(info_dict.get("text_token"))
+            if streamed_id is not None and streamed_id >= 0:
+                self._dec_text_tokens[start] = streamed_id
+                self._dec_text_mask[start] = 1
+            else:
+                self._dec_text_tokens[start] = 0
+                self._dec_text_mask[start] = 0
 
         # ── Phoneme channel ── opens at decode step == ``phonemes_delay`` (seeded
         # with phoneme BOS), then feeds back the previous step's prediction, and
@@ -1014,8 +1345,28 @@ def _preprocess_decode(
         if self.has_phoneme:
             phoneme_ended = bool(info_dict.get("phoneme_ended", False))
             feed_eos = False
+            gt_phoneme_rows = self._coerce_int_rows(
+                info_dict.get("gt_phoneme_tokens"),
+                int(self.arch.phoneme_stacking_factor),
+            )
+            gt_phoneme_index = decode_offset - self.phonemes_delay
             if phoneme_ended or decode_offset < self.phonemes_delay:
+                self._dec_phoneme_tokens[start].zero_()
                 self._dec_phoneme_valid[start] = 0
+            elif gt_phoneme_rows is not None:
+                if 0 <= gt_phoneme_index < len(gt_phoneme_rows):
+                    row = torch.tensor(
+                        gt_phoneme_rows[gt_phoneme_index],
+                        device=device,
+                        dtype=torch.long,
+                    ).reshape(-1)[: self.arch.phoneme_stacking_factor]
+                    self._dec_phoneme_tokens[start].zero_()
+                    self._dec_phoneme_tokens[start, : row.shape[0]].copy_(row)
+                    self._dec_phoneme_valid[start] = 1
+                    feed_eos = bool((row == self.phoneme_eos_id).any())
+                else:
+                    self._dec_phoneme_tokens[start].zero_()
+                    self._dec_phoneme_valid[start] = 0
             elif decode_offset == self.phonemes_delay:
                 self._dec_phoneme_tokens[start].fill_(self.phoneme_bos_id)
                 self._dec_phoneme_valid[start] = 1
@@ -1027,6 +1378,7 @@ def _preprocess_decode(
                     self._dec_phoneme_valid[start] = 1
                     feed_eos = bool((p == self.phoneme_eos_id).any())
                 else:
+                    self._dec_phoneme_tokens[start].zero_()
                     self._dec_phoneme_valid[start] = 0
             if phoneme_ended or feed_eos:
                 info_update["phoneme_ended"] = True
@@ -1038,6 +1390,7 @@ def _preprocess_decode(
         # stability but its codes for those frames are discarded by the caller and
         # never fed back here.
         if decode_offset < self.speech_delay:
+            self._dec_audio_codes[start].zero_()
             self._dec_audio_valid[start] = 0
         elif decode_offset == self.speech_delay:
             self._dec_audio_codes[start].fill_(self.arch.audio_bos_id)
@@ -1060,22 +1413,66 @@ def postprocess(self, hidden_states: torch.Tensor, multimodal_outputs: Optional[
         """Stash the last frame's codes (and phoneme) for the next decode step."""
         if hidden_states.numel() == 0:
             return {}
-        stride0 = hidden_states.stride(0) or 1
-        req_start = hidden_states.storage_offset() // stride0
-        last = req_start + hidden_states.shape[0] - 1
-
         out: dict[str, Any] = {}
-        audio_codes = (multimodal_outputs or {}).get("audio_codes")
+        multimodal_outputs = multimodal_outputs or {}
+        audio_codes = multimodal_outputs.get("audio_codes_feedback")
+        if not isinstance(audio_codes, torch.Tensor):
+            audio_codes = multimodal_outputs.get("audio_codes")
         if isinstance(audio_codes, torch.Tensor) and audio_codes.numel() > 0:
+            last = self._last_request_row_index(
+                hidden_states,
+                audio_codes,
+                getattr(self, "_last_output_row_indices", None),
+            )
             out["last_audio_codes"] = audio_codes[last : last + 1].detach()
         if self.has_phoneme:
-            out["last_phoneme_token"] = self._dec_phoneme_tokens[last : last + 1].detach().clone()
+            phoneme_tokens = multimodal_outputs.get("phoneme_tokens_feedback")
+            if isinstance(phoneme_tokens, torch.Tensor) and phoneme_tokens.numel() > 0:
+                last = self._last_request_row_index(
+                    hidden_states,
+                    phoneme_tokens,
+                    getattr(self, "_last_output_row_indices", None),
+                )
+                out["last_phoneme_token"] = phoneme_tokens[last : last + 1].detach().clone()
+            else:
+                last = self._last_request_row_index(
+                    hidden_states,
+                    self._dec_phoneme_tokens,
+                    getattr(self, "_last_output_row_indices", None),
+                )
+                out["last_phoneme_token"] = self._dec_phoneme_tokens[last : last + 1].detach().clone()
         return out
 
     # ------------------------------------------------------------------
     # weight loading
     # ------------------------------------------------------------------
 
+    def _reset_runtime_state_after_weight_load(self) -> None:
+        """Clear mutable generation state after initial load or live refit."""
+
+        self.mamba_cache = None
+        self._last_output_row_indices = None
+        for name in (
+            "_combined_embeddings",
+            "_dec_text_tokens",
+            "_dec_text_mask",
+            "_dec_audio_codes",
+            "_dec_audio_valid",
+            "_out_codes",
+            "_out_code_logprobs",
+            "_out_code_sampling_logprobs",
+            "_out_frame_logprobs",
+            "_token_stop",
+            "_sample_stop",
+            "_dec_phoneme_tokens",
+            "_dec_phoneme_valid",
+        ):
+            value = getattr(self, name, None)
+            if isinstance(value, torch.Tensor):
+                value.zero_()
+        if isinstance(getattr(self, "_decode_offsets", None), torch.Tensor):
+            self._decode_offsets.fill_(-1)
+
     # Checkpoint prefixes (EasyMagpieTTS state dict) → in-model paths.
     # ``decoder.*`` is fed to the vLLM backbone loader separately (it understands
     # HF Nemotron-H naming + Mamba/MoE packing). The TTS submodules are copied
@@ -1090,6 +1487,7 @@ def postprocess(self, hidden_states: torch.Tensor, multimodal_outputs: Optional[
         "phoneme_embeddings.": "phoneme_embeddings.",
         "phoneme_final_proj.": "phoneme_final_proj.",
         "text_embedding.": "text_embedding.",
+        "context_text_embedding.": "context_text_embedding.",
         "task_embedding.": "task_embedding.",
     }
 
@@ -1100,6 +1498,408 @@ def _remap_tts_key(self, name: str) -> Optional[str]:
                 return dst + name[len(src) :]
         return None
 
+    @staticmethod
+    def _embedding_table_is_degenerate(table: torch.Tensor) -> bool:
+        table_f = table.detach().float()
+        return (
+            table_f.numel() == 0
+            or (not bool(torch.isfinite(table_f).all().item()))
+            or float(table_f.abs().max().item()) == 0.0
+        )
+
+    @staticmethod
+    def _sample_refit_indices(numel: int, max_values: int = 8) -> list[int]:
+        if numel <= 0 or max_values <= 0:
+            return []
+        if numel <= max_values:
+            return list(range(numel))
+        candidates = [0, 1, numel // 4, numel // 2, (3 * numel) // 4, numel - 2, numel - 1]
+        indices: list[int] = []
+        for idx in candidates:
+            idx = max(0, min(numel - 1, int(idx)))
+            if idx not in indices:
+                indices.append(idx)
+            if len(indices) >= max_values:
+                break
+        return indices
+
+    @staticmethod
+    def _refit_skip_unchanged_enabled() -> bool:
+        raw = os.getenv("EASYMAGPIE_SKIP_UNCHANGED_REFIT_PAYLOAD", "0").strip().lower()
+        return raw in {"1", "true", "yes", "on"}
+
+    def _refit_default_float_dtype(self) -> torch.dtype:
+        for param in self.parameters():
+            if param.is_floating_point():
+                return param.dtype
+        return torch.float32
+
+    def _refit_source_for_fingerprint(
+        self,
+        name: str,
+        tensor: torch.Tensor,
+        *,
+        own_params: dict[str, torch.Tensor],
+        backbone_params: dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        source = tensor.detach()
+        if name.startswith("decoder."):
+            mapped_weights = list(_apply_optional_nemotron_h_weight_mapper([(name[len("decoder.") :], source)]))
+            if mapped_weights:
+                mapped_name, mapped_source = mapped_weights[0]
+            else:
+                mapped_name, mapped_source = name[len("decoder.") :], source
+            target_name, source, _ = self._backbone_refit_target_name_and_tensor(
+                str(mapped_name),
+                mapped_source,
+                input_name=name,
+            )
+            target = backbone_params.get(target_name)
+        else:
+            mapped_name = self._remap_tts_key(name)
+            target = own_params.get(mapped_name) if mapped_name is not None else None
+            source = tensor.detach()
+
+        if source.is_floating_point():
+            dtype = target.dtype if isinstance(target, torch.Tensor) else self._refit_default_float_dtype()
+            if source.dtype != dtype:
+                source = source.to(dtype=dtype)
+        return source
+
+    @staticmethod
+    def _shape_equal(left: torch.Tensor, right: torch.Tensor) -> bool:
+        return tuple(left.shape) == tuple(right.shape)
+
+    @staticmethod
+    def _candidate_if_shape_matches(candidate: torch.Tensor, reference: torch.Tensor) -> list[torch.Tensor]:
+        if EasyMagpieTTSForConditionalGeneration._shape_equal(candidate, reference):
+            return [candidate.detach()]
+        if candidate.ndim == 2 and EasyMagpieTTSForConditionalGeneration._shape_equal(candidate.t(), reference):
+            return [candidate.t().contiguous().detach()]
+        return []
+
+    def _packed_qkv_resident_sources_for_fingerprint(
+        self,
+        mapped_name: str,
+        source: torch.Tensor,
+        *,
+        backbone_params: dict[str, torch.Tensor],
+    ) -> list[torch.Tensor]:
+        match = _QKV_WEIGHT_RE.match(mapped_name)
+        if match is None:
+            return []
+        layer = match.group("layer")
+        projection = match.group("projection")
+        target = backbone_params.get(f"layers.{layer}.mixer.qkv_proj.weight")
+        if not isinstance(target, torch.Tensor) or target.ndim != 2 or source.ndim != 2:
+            return []
+
+        total_rows = int(target.shape[0])
+        q_rows = int(getattr(self, "hidden_dim", 0) or 0)
+        if q_rows <= 0 or total_rows <= q_rows:
+            q_rows = int(source.shape[0]) if projection == "q_proj" else 0
+        if q_rows <= 0:
+            return []
+        kv_rows = (total_rows - q_rows) // 2
+        if kv_rows <= 0 or q_rows + 2 * kv_rows != total_rows:
+            return []
+
+        if projection == "q_proj":
+            candidate = target[:q_rows, :]
+        elif projection == "k_proj":
+            candidate = target[q_rows : q_rows + kv_rows, :]
+        else:
+            candidate = target[q_rows + kv_rows : q_rows + 2 * kv_rows, :]
+        return self._candidate_if_shape_matches(candidate, source)
+
+    def _packed_moe_resident_sources_for_fingerprint(
+        self,
+        mapped_name: str,
+        source: torch.Tensor,
+        *,
+        backbone_params: dict[str, torch.Tensor],
+    ) -> list[torch.Tensor]:
+        match = _MOE_EXPERT_WEIGHT_RE.match(mapped_name)
+        if match is None:
+            return []
+        layer = match.group("layer")
+        expert_idx = int(match.group("expert"))
+        projection = match.group("projection")
+        target_key = "w13_weight" if projection == "up_proj" else "w2_weight"
+        target = backbone_params.get(f"layers.{layer}.mixer.experts.{target_key}")
+        if not isinstance(target, torch.Tensor) or target.ndim < 3:
+            return []
+        if expert_idx < 0 or expert_idx >= int(target.shape[0]):
+            return []
+
+        slab = target[expert_idx]
+        candidates = self._candidate_if_shape_matches(slab, source)
+        if projection == "up_proj" and slab.ndim == source.ndim and slab.shape[0] == source.shape[0] * 2:
+            first = slab[: source.shape[0]]
+            second = slab[source.shape[0] : source.shape[0] * 2]
+            candidates.extend(self._candidate_if_shape_matches(first, source))
+            candidates.extend(self._candidate_if_shape_matches(second, source))
+        return candidates
+
+    def _resident_sources_for_fingerprint(
+        self,
+        name: str,
+        tensor: torch.Tensor,
+        *,
+        own_params: dict[str, torch.Tensor],
+        backbone_params: dict[str, torch.Tensor],
+    ) -> list[torch.Tensor]:
+        source = self._refit_source_for_fingerprint(
+            name,
+            tensor,
+            own_params=own_params,
+            backbone_params=backbone_params,
+        )
+        if name.startswith("decoder."):
+            mapped_name = name[len("decoder.") :]
+            exact_target_name, _, _ = self._backbone_refit_target_name_and_tensor(
+                mapped_name,
+                tensor.detach(),
+                input_name=name,
+            )
+            target = backbone_params.get(exact_target_name)
+            candidates: list[torch.Tensor] = []
+            if isinstance(target, torch.Tensor):
+                candidates.extend(self._candidate_if_shape_matches(target, source))
+            candidates.extend(
+                self._packed_qkv_resident_sources_for_fingerprint(
+                    mapped_name,
+                    source,
+                    backbone_params=backbone_params,
+                )
+            )
+            candidates.extend(
+                self._packed_moe_resident_sources_for_fingerprint(
+                    mapped_name,
+                    source,
+                    backbone_params=backbone_params,
+                )
+            )
+            return candidates
+
+        mapped = self._remap_tts_key(name)
+        if mapped is None:
+            return []
+        target = own_params.get(mapped)
+        if not isinstance(target, torch.Tensor):
+            return []
+        return self._candidate_if_shape_matches(target, source)
+
+    def _refit_loaded_target_names_for_input(
+        self,
+        name: str,
+        tensor: torch.Tensor,
+        *,
+        own_params: dict[str, torch.Tensor],
+        backbone_params: dict[str, torch.Tensor],
+    ) -> set[str]:
+        if name.startswith("decoder."):
+            mapped_name = name[len("decoder.") :]
+            target_names: set[str] = set()
+            exact_target_name, _, _ = self._backbone_refit_target_name_and_tensor(
+                mapped_name,
+                tensor.detach(),
+                input_name=name,
+            )
+            if exact_target_name in backbone_params:
+                target_names.add(f"backbone.{exact_target_name}")
+            qkv_match = _QKV_WEIGHT_RE.match(mapped_name)
+            if qkv_match is not None:
+                qkv_name = f"layers.{qkv_match.group('layer')}.mixer.qkv_proj.weight"
+                if qkv_name in backbone_params:
+                    target_names.add(f"backbone.{qkv_name}")
+            moe_match = _MOE_EXPERT_WEIGHT_RE.match(mapped_name)
+            if moe_match is not None:
+                target_key = "w13_weight" if moe_match.group("projection") == "up_proj" else "w2_weight"
+                moe_name = f"layers.{moe_match.group('layer')}.mixer.experts.{target_key}"
+                if moe_name in backbone_params:
+                    target_names.add(f"backbone.{moe_name}")
+            return target_names
+
+        mapped = self._remap_tts_key(name)
+        if mapped is not None and mapped in own_params:
+            return {mapped}
+        return set()
+
+    @staticmethod
+    @torch.no_grad()
+    def _refit_source_fingerprint(tensor: torch.Tensor) -> str:
+        """Return a content hash for deciding whether a refit tensor is unchanged."""
+
+        source = tensor.detach().contiguous().cpu()
+        as_bytes = source.view(torch.uint8).numpy()
+        digest = hashlib.sha256()
+        digest.update(str(source.dtype).encode("utf-8"))
+        digest.update(b"|")
+        digest.update(",".join(str(int(dim)) for dim in source.shape).encode("utf-8"))
+        digest.update(b"|")
+        digest.update(memoryview(as_bytes))
+        return digest.hexdigest()
+
+    @classmethod
+    @torch.no_grad()
+    def _sample_refit_copy_check(
+        cls,
+        *,
+        input_name: str,
+        target_name: str,
+        source: torch.Tensor,
+        target: torch.Tensor,
+        group: str,
+        note: str | None = None,
+    ) -> dict[str, Any]:
+        result: dict[str, Any] = {
+            "input_name": input_name,
+            "target_name": target_name,
+            "group": group,
+            "source_shape": list(source.shape),
+            "target_shape": list(target.shape),
+            "source_dtype": str(source.dtype).replace("torch.", ""),
+            "target_dtype": str(target.dtype).replace("torch.", ""),
+            "checked": False,
+            "matched": False,
+        }
+        if note:
+            result["note"] = note
+        if tuple(source.shape) != tuple(target.shape):
+            result["reason"] = "shape_mismatch"
+            return result
+
+        numel = int(target.numel())
+        indices = cls._sample_refit_indices(numel)
+        result["numel"] = numel
+        result["sample_indices"] = indices
+        if not indices:
+            result["checked"] = True
+            result["matched"] = True
+            result["max_sample_abs_diff"] = 0.0
+            return result
+
+        src_flat = source.detach().reshape(-1)
+        tgt_flat = target.detach().reshape(-1)
+        src_idx = torch.tensor(indices, device=src_flat.device, dtype=torch.long)
+        tgt_idx = torch.tensor(indices, device=tgt_flat.device, dtype=torch.long)
+        src_sample = src_flat.index_select(0, src_idx).to(device=tgt_flat.device, dtype=target.dtype)
+        tgt_sample = tgt_flat.index_select(0, tgt_idx)
+        if target.is_floating_point():
+            diff = (tgt_sample.float() - src_sample.float()).abs()
+            max_diff = float(diff.max().item()) if diff.numel() else 0.0
+            if target.dtype in (torch.float16, torch.bfloat16):
+                abs_tol = 5e-3
+                rel_tol = 5e-3
+            else:
+                abs_tol = 1e-5
+                rel_tol = 1e-5
+            matched = bool(torch.allclose(tgt_sample.float(), src_sample.float(), atol=abs_tol, rtol=rel_tol))
+            result["sample_abs_tol"] = abs_tol
+            result["sample_rel_tol"] = rel_tol
+        else:
+            mismatch = tgt_sample.ne(src_sample)
+            max_diff = float(mismatch.to(torch.float32).max().item()) if mismatch.numel() else 0.0
+            matched = not bool(mismatch.any().item())
+        result.update(
+            {
+                "checked": True,
+                "matched": matched,
+                "max_sample_abs_diff": max_diff,
+            }
+        )
+        if not matched:
+            result["source_sample"] = [float(x) for x in src_sample.detach().float().cpu().tolist()]
+            result["target_sample"] = [float(x) for x in tgt_sample.detach().float().cpu().tolist()]
+        return result
+
+    @staticmethod
+    def _backbone_refit_target_name_and_tensor(
+        name: str,
+        tensor: torch.Tensor,
+        *,
+        input_name: str | None = None,
+    ) -> tuple[str, torch.Tensor, str | None]:
+        target_name = name
+        source = tensor
+        if "embeddings" in target_name:
+            target_name = target_name.replace("embeddings", "embed_tokens")
+        original_name = input_name or name
+        if "A_log" in target_name or "A_log" in original_name:
+            target_name = target_name.replace("A_log", "A")
+            # Nemotron-H stores the continuous-time SSM state matrix as
+            # ``A = -exp(A_log)`` in the vLLM module.  Match the loader's value
+            # transformation when verifying resident refit tensors.
+            source = -torch.exp(source.to(torch.float32))
+        if "D" in target_name:
+            source = source.to(torch.float32)
+        if "dt_bias" in target_name:
+            source = source.to(torch.float32)
+        if any(proj in target_name for proj in ("q_proj", "k_proj", "v_proj")):
+            weight_name = next(proj for proj in ("q_proj", "k_proj", "v_proj") if proj in target_name)
+            return target_name.replace(weight_name, "qkv_proj"), source, f"packed_{weight_name}"
+        return target_name, source, None
+
+    @staticmethod
+    def _summarize_refit_copy_checks(checks: list[dict[str, Any]]) -> dict[str, Any]:
+        checked = [item for item in checks if bool(item.get("checked"))]
+        failed = [item for item in checked if not bool(item.get("matched"))]
+        skipped = [item for item in checks if not bool(item.get("checked"))]
+        max_diff = 0.0
+        worst: dict[str, Any] | None = None
+        for item in checked:
+            diff = float(item.get("max_sample_abs_diff", 0.0) or 0.0)
+            if diff >= max_diff:
+                max_diff = diff
+                worst = item
+
+        by_group: dict[str, dict[str, Any]] = {}
+        for item in checks:
+            group = str(item.get("group", "unknown"))
+            bucket = by_group.setdefault(group, {"num_total": 0, "num_checked": 0, "num_matched": 0, "num_failed": 0})
+            bucket["num_total"] += 1
+            if bool(item.get("checked")):
+                bucket["num_checked"] += 1
+                if bool(item.get("matched")):
+                    bucket["num_matched"] += 1
+                else:
+                    bucket["num_failed"] += 1
+
+        return {
+            "ok": not failed,
+            "num_total": len(checks),
+            "num_checked": len(checked),
+            "num_matched": len(checked) - len(failed),
+            "num_failed": len(failed),
+            "num_skipped": len(skipped),
+            "max_sample_abs_diff": max_diff,
+            "worst": worst,
+            "failed_head": failed[:8],
+            "skipped_head": skipped[:8],
+            "by_group": by_group,
+        }
+
+    def _repair_context_text_embedding_if_needed(self, input_names: set[str]) -> bool:
+        """Fallback for older artifacts without a valid context-text table."""
+
+        context_embedding = getattr(self, "context_text_embedding", None)
+        text_embedding = getattr(self, "text_embedding", None)
+        context_weight = getattr(context_embedding, "weight", None)
+        text_weight = getattr(text_embedding, "weight", None)
+        if not isinstance(context_weight, torch.Tensor) or not isinstance(text_weight, torch.Tensor):
+            return False
+        if context_weight is text_weight:
+            return False
+        if tuple(context_weight.shape) != tuple(text_weight.shape):
+            return False
+        context_name = "context_text_embedding.weight"
+        if context_name in input_names and not self._embedding_table_is_degenerate(context_weight):
+            return False
+        with torch.no_grad():
+            context_weight.copy_(text_weight.to(dtype=context_weight.dtype, device=context_weight.device))
+        return True
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load backbone (Nemotron-H) + TTS submodule weights from a converted checkpoint.
 
@@ -1110,38 +1910,127 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         :meth:`NemotronHModel.load_weights` (which handles HF naming + Mamba/MoE
         packing); TTS weights are copied directly by name.
         """
+        materialized_weights = list(weights)
+        input_names = [str(name) for name, _ in materialized_weights]
         own_params = dict(self.named_parameters())
+        backbone_params = dict(self.backbone.named_parameters())
         loaded: set[str] = set()
         backbone_weights: list[tuple[str, torch.Tensor]] = []
+        backbone_input_names: list[str] = []
+        forwarded_backbone_input_names: list[str] = []
+        skipped_unchanged_backbone: list[tuple[str, str, torch.Tensor]] = []
+        tts_input_names: list[str] = []
+        skipped_unmapped: list[str] = []
+        skipped_optional: list[str] = []
+        skipped_unchanged: list[str] = []
+        skipped_unchanged_by_source_fingerprint: list[str] = []
+        skipped_unchanged_by_resident_fingerprint: list[str] = []
+        fingerprint_cache_misses: list[str] = []
+        missing_targets: list[str] = []
+        tts_copy_check_items: list[tuple[str, str, torch.Tensor, torch.Tensor]] = []
+        refit_rpc_active = bool(getattr(self, "_easymagpie_refit_rpc_active", False))
+        skip_unchanged_enabled = self._refit_skip_unchanged_enabled() and refit_rpc_active
+        previous_fingerprints = getattr(self, "_easymagpie_last_refit_source_fingerprints", {})
+        if not isinstance(previous_fingerprints, dict):
+            previous_fingerprints = {}
+        current_fingerprints: dict[str, str] = {}
+        if skip_unchanged_enabled:
+            current_fingerprints = {
+                str(name): self._refit_source_fingerprint(
+                    self._refit_source_for_fingerprint(
+                        str(name),
+                        tensor,
+                        own_params=own_params,
+                        backbone_params=backbone_params,
+                    )
+                )
+                for name, tensor in materialized_weights
+            }
 
-        for name, tensor in weights:
+        for name, tensor in materialized_weights:
+            source_fingerprint_match = (
+                skip_unchanged_enabled
+                and previous_fingerprints.get(str(name)) == current_fingerprints.get(str(name))
+            )
+            resident_fingerprint_match = False
+            if skip_unchanged_enabled and not source_fingerprint_match:
+                current_fingerprint = current_fingerprints.get(str(name))
+                if current_fingerprint is not None:
+                    for resident_source in self._resident_sources_for_fingerprint(
+                        str(name),
+                        tensor,
+                        own_params=own_params,
+                        backbone_params=backbone_params,
+                    ):
+                        if self._refit_source_fingerprint(resident_source) == current_fingerprint:
+                            resident_fingerprint_match = True
+                            break
+                if resident_fingerprint_match:
+                    skipped_unchanged_by_resident_fingerprint.append(str(name))
+                else:
+                    fingerprint_cache_misses.append(str(name))
+            elif source_fingerprint_match:
+                skipped_unchanged_by_source_fingerprint.append(str(name))
+            unchanged = bool(source_fingerprint_match or resident_fingerprint_match)
             if name.startswith("decoder."):
-                backbone_weights.append((name[len("decoder.") :], tensor))
+                backbone_input_names.append(name)
+                if unchanged:
+                    skipped_unchanged.append(name)
+                    skipped_unchanged_backbone.append((name, name[len("decoder.") :], tensor))
+                    loaded.update(
+                        self._refit_loaded_target_names_for_input(
+                            str(name),
+                            tensor,
+                            own_params=own_params,
+                            backbone_params=backbone_params,
+                        )
+                    )
+                else:
+                    backbone_weights.append((name[len("decoder.") :], tensor))
+                    forwarded_backbone_input_names.append(name)
                 continue
             mapped = self._remap_tts_key(name)
             if mapped is None:
                 # Unrelated checkpoint section (codec, speaker encoder, CAS, etc.).
+                skipped_unmapped.append(name)
                 continue
             if mapped.startswith("task_embedding.") and self.task_embedding is None:
                 # Single-mode model: checkpoint may still ship an (unused) table.
+                skipped_optional.append(name)
                 continue
             target = own_params.get(mapped)
             if target is None:
+                if mapped == "context_text_embedding.weight" and self.context_text_embedding is self.text_embedding:
+                    skipped_optional.append(name)
+                    continue
                 logger.warning("EasyMagpieTTS: no parameter for checkpoint key %s -> %s", name, mapped)
+                missing_targets.append(name)
                 continue
-            # The local-transformer FFN ships as kernel-1 ``Conv1d`` weights
-            # (``[out, in, 1]``) but now lives as ``nn.Linear`` (``[out, in]``).
-            # Squeeze the trailing singleton conv dim so the dense layer loads 1:1.
-            if tensor.ndim == target.ndim + 1 and tensor.shape[-1] == 1:
-                tensor = tensor.squeeze(-1)
             if target.shape != tensor.shape:
                 raise RuntimeError(
                     f"EasyMagpieTTS weight shape mismatch at {mapped!r}: "
                     f"ckpt {tuple(tensor.shape)} vs model {tuple(target.shape)}"
                 )
-            with torch.no_grad():
-                target.data.copy_(tensor.to(target.dtype))
+            if unchanged:
+                skipped_unchanged.append(name)
+                loaded.update(
+                    self._refit_loaded_target_names_for_input(
+                        str(name),
+                        tensor,
+                        own_params=own_params,
+                        backbone_params=backbone_params,
+                    )
+                )
+            else:
+                with torch.no_grad():
+                    target.data.copy_(tensor.to(target.dtype))
+            tts_input_names.append(name)
             loaded.add(mapped)
+            tts_copy_check_items.append((name, mapped, tensor, target))
+
+        context_text_embedding_repaired = self._repair_context_text_embedding_if_needed(set(input_names))
+        if context_text_embedding_repaired and "context_text_embedding.weight" in own_params:
+            loaded.add("context_text_embedding.weight")
 
         # ``NemotronHModel.load_weights`` (the inner model) does *not* apply the
         # HF->vLLM renaming that lives on the ``NemotronHForCausalLM`` wrapper, so
@@ -1150,12 +2039,195 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         # Apply that mapper here so the converted checkpoint can keep stock HF
         # Nemotron-H names. The wrapper's ``backbone -> model`` prefix rule is a
         # no-op here because we already stripped the ``decoder.`` prefix.
-        backbone_weights = list(NemotronHForCausalLM.hf_to_vllm_mapper.apply(backbone_weights))
-        backbone_loaded = self.backbone.load_weights(backbone_weights)
+        backbone_weights = list(_apply_optional_nemotron_h_weight_mapper(backbone_weights))
+        if backbone_weights:
+            backbone_loaded = self.backbone.load_weights(backbone_weights)
+        else:
+            backbone_loaded = set()
         loaded |= {f"backbone.{n}" for n in backbone_loaded}
+        if skipped_unchanged:
+            previous_loaded_targets = getattr(self, "_easymagpie_last_refit_loaded_targets", set())
+            if isinstance(previous_loaded_targets, set):
+                loaded |= {str(name) for name in previous_loaded_targets}
+        copy_checks: list[dict[str, Any]] = []
+        for input_name, mapped, tensor, target in tts_copy_check_items:
+            copy_checks.append(
+                self._sample_refit_copy_check(
+                    input_name=input_name,
+                    target_name=mapped,
+                    source=tensor,
+                    target=target,
+                    group="tts",
+                )
+            )
+        for original_name, (mapped_name, tensor) in zip(forwarded_backbone_input_names, backbone_weights):
+            target_name, source, skip_reason = self._backbone_refit_target_name_and_tensor(
+                mapped_name,
+                tensor,
+                input_name=original_name,
+            )
+            target = backbone_params.get(target_name)
+            if target is None:
+                copy_checks.append(
+                    {
+                        "input_name": original_name,
+                        "target_name": target_name,
+                        "group": "backbone",
+                        "checked": False,
+                        "matched": False,
+                        "reason": "target_not_found",
+                    }
+                )
+                continue
+            if skip_reason is not None:
+                copy_checks.append(
+                    {
+                        "input_name": original_name,
+                        "target_name": target_name,
+                        "group": "backbone",
+                        "checked": False,
+                        "matched": False,
+                        "reason": skip_reason,
+                    }
+                )
+                continue
+            copy_checks.append(
+                self._sample_refit_copy_check(
+                    input_name=original_name,
+                    target_name=f"backbone.{target_name}",
+                    source=source,
+                    target=target,
+                    group="backbone",
+                )
+            )
+        skipped_backbone_weights = list(
+            _apply_optional_nemotron_h_weight_mapper(
+                [(mapped_name, tensor) for _, mapped_name, tensor in skipped_unchanged_backbone]
+            )
+        )
+        for original_name, (mapped_name, tensor) in zip(
+            [name for name, _, _ in skipped_unchanged_backbone],
+            skipped_backbone_weights,
+        ):
+            target_name, source, skip_reason = self._backbone_refit_target_name_and_tensor(
+                mapped_name,
+                tensor,
+                input_name=original_name,
+            )
+            target = backbone_params.get(target_name)
+            if target is None:
+                copy_checks.append(
+                    {
+                        "input_name": original_name,
+                        "target_name": target_name,
+                        "group": "backbone",
+                        "checked": False,
+                        "matched": False,
+                        "reason": "skipped_unchanged_target_not_found",
+                    }
+                )
+                continue
+            if skip_reason is not None:
+                copy_checks.append(
+                    {
+                        "input_name": original_name,
+                        "target_name": target_name,
+                        "group": "backbone",
+                        "checked": False,
+                        "matched": False,
+                        "reason": f"skipped_unchanged_{skip_reason}",
+                    }
+                )
+                continue
+            copy_checks.append(
+                self._sample_refit_copy_check(
+                    input_name=original_name,
+                    target_name=f"backbone.{target_name}",
+                    source=source,
+                    target=target,
+                    group="backbone",
+                    note="skipped_unchanged",
+                )
+            )
+        refit_copy_check = self._summarize_refit_copy_checks(copy_checks)
+
+        missing_model_targets = sorted(str(name) for name in own_params if str(name) not in loaded)
+        allow_missing_text_tables = bool(getattr(self, "_easymagpie_allow_missing_text_tables_refit", False))
+        allowed_missing_model_targets = set()
+        if allow_missing_text_tables:
+            allowed_missing_model_targets.add("text_embedding.weight")
+            allowed_missing_model_targets.add("context_text_embedding.weight")
+        blocking_missing_model_targets = [
+            name for name in missing_model_targets if name not in allowed_missing_model_targets
+        ]
+        loaded_names = sorted(str(name) for name in loaded)
+        num_backbone_loaded_targets = sum(1 for name in loaded_names if name.startswith("backbone."))
+        context_text_embedding_aliased = self.context_text_embedding is self.text_embedding
+        fingerprint_cache_size_before = len(previous_fingerprints)
+        fingerprint_cache_size_after = (
+            len({**previous_fingerprints, **current_fingerprints}) if skip_unchanged_enabled else 0
+        )
+        self._last_easy_magpie_load_weights_summary = {
+            "ok": (
+                not skipped_unmapped
+                and not missing_targets
+                and not blocking_missing_model_targets
+                and bool(refit_copy_check.get("ok", False))
+            ),
+            "num_input_tensors": len(materialized_weights),
+            "num_routed_input_tensors": len(backbone_input_names) + len(tts_input_names),
+            "num_backbone_input_tensors": len(backbone_input_names),
+            "num_tts_input_tensors": len(tts_input_names),
+            "num_loaded_targets": len(loaded_names),
+            "num_backbone_loaded_targets": num_backbone_loaded_targets,
+            "num_tts_loaded_targets": len(loaded_names) - num_backbone_loaded_targets,
+            "num_actual_backbone_loader_targets": len(backbone_loaded),
+            "refit_rpc_active": bool(refit_rpc_active),
+            "skip_unchanged_enabled": bool(skip_unchanged_enabled),
+            "fingerprint_cache_size_before": fingerprint_cache_size_before,
+            "fingerprint_cache_size_after": fingerprint_cache_size_after,
+            "num_source_fingerprint_matches": len(skipped_unchanged_by_source_fingerprint),
+            "num_resident_fingerprint_matches": len(skipped_unchanged_by_resident_fingerprint),
+            "num_fingerprint_cache_misses": len(fingerprint_cache_misses),
+            "num_skipped_unchanged": len(skipped_unchanged),
+            "skipped_unchanged_head": skipped_unchanged[:32],
+            "skipped_unchanged_tail": skipped_unchanged[-32:],
+            "skipped_unchanged_by_source_fingerprint_head": skipped_unchanged_by_source_fingerprint[:32],
+            "skipped_unchanged_by_source_fingerprint_tail": skipped_unchanged_by_source_fingerprint[-32:],
+            "skipped_unchanged_by_resident_fingerprint_head": skipped_unchanged_by_resident_fingerprint[:32],
+            "skipped_unchanged_by_resident_fingerprint_tail": skipped_unchanged_by_resident_fingerprint[-32:],
+            "fingerprint_cache_miss_head": fingerprint_cache_misses[:32],
+            "fingerprint_cache_miss_tail": fingerprint_cache_misses[-32:],
+            "num_backbone_weights_forwarded_to_loader": len(backbone_weights),
+            "input_head": input_names[:16],
+            "input_tail": input_names[-16:],
+            "loaded_head": loaded_names[:16],
+            "loaded_tail": loaded_names[-16:],
+            "skipped_unmapped": skipped_unmapped[:32],
+            "skipped_optional": skipped_optional[:32],
+            "missing_targets": missing_targets[:32],
+            "missing_model_targets": missing_model_targets[:32],
+            "blocking_missing_model_targets": blocking_missing_model_targets[:32],
+            "allowed_missing_model_targets": sorted(allowed_missing_model_targets)[:32],
+            "num_skipped_unmapped": len(skipped_unmapped),
+            "num_skipped_optional": len(skipped_optional),
+            "num_missing_targets": len(missing_targets),
+            "num_missing_model_targets": len(missing_model_targets),
+            "num_blocking_missing_model_targets": len(blocking_missing_model_targets),
+            "allow_missing_text_tables": bool(allow_missing_text_tables),
+            "context_text_embedding_repaired": bool(context_text_embedding_repaired),
+            "context_text_embedding_aliased": bool(context_text_embedding_aliased),
+            "refit_copy_check": refit_copy_check,
+        }
 
         # Derived runtime state.
         self.code_predictor.init_forbidden_mask()
+        self._reset_runtime_state_after_weight_load()
+        if skip_unchanged_enabled and bool(self._last_easy_magpie_load_weights_summary.get("ok", False)):
+            cached = dict(previous_fingerprints)
+            cached.update(current_fingerprints)
+            self._easymagpie_last_refit_source_fingerprints = cached
+            self._easymagpie_last_refit_loaded_targets = set(loaded_names)
 
         logger.info("Loaded %d weights for EasyMagpieTTSForConditionalGeneration", len(loaded))
         return loaded
diff --git a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/local_transformer.py b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/local_transformer.py
index 3633b9902c18..0d2b2459fce3 100644
--- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/local_transformer.py
+++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/local_transformer.py
@@ -27,7 +27,7 @@
 * :class:`EasyMagpieCodePredictor` owns the persistent, address-stable scratch
   buffers and runs the per-frame autoregressive loop, re-invoking the compiled
   transformer once per codebook over the **same** buffer (replaying one
-  fixed-shape graph N times is faster and simpler than capturing N separate
+  stable-shape graph N times is faster and simpler than capturing N separate
   graphs).
 
 All sampling is CUDA-graph safe (Gumbel-max + ``topk`` + ``masked_fill`` only;
@@ -35,6 +35,8 @@
 """
 from __future__ import annotations
 
+import os
+
 import torch
 from torch import nn
 from vllm.compilation.decorators import support_torch_compile
@@ -43,15 +45,115 @@
 from easymagpie_vllm_omni.config import EasyMagpieOmniArch
 
 
-# Default top-k width for audio-codebook sampling. Because ``torch.topk``'s ``k``
-# shapes tensors inside the captured graph, this becomes a capture-time constant.
-_DEFAULT_TOP_K = 80
+def _gumbel_argmax(logits: torch.Tensor) -> torch.Tensor:
+    """Gumbel-max categorical draw — CUDA-graph safe.
+
+    Equivalent to sampling from ``softmax(logits)`` but uses only
+    ``uniform_`` + ``log`` + ``argmax`` (all legal inside a captured graph)
+    and degrades gracefully on degenerate warmup logits instead of triggering
+    a device-side assert the way ``multinomial`` does.
+    """
+    u = torch.empty_like(logits).uniform_(1e-20, 1.0 - 1e-20)
+    return (logits - torch.log(-torch.log(u))).argmax(dim=-1)
+
+
+def _sanitize_logits(logits: torch.Tensor) -> torch.Tensor:
+    """Replace non-finite logits before CUDA-graph-safe sampling math."""
+
+    finfo = torch.finfo(logits.dtype)
+    return torch.nan_to_num(
+        logits,
+        nan=-finfo.max,
+        posinf=finfo.max,
+        neginf=-finfo.max,
+    )
+
+
+def _selected_logprob(distribution_logits: torch.Tensor, selected: torch.Tensor) -> torch.Tensor:
+    """Gather finite selected-token logprobs from a row-wise logits tensor."""
+
+    logits_f = distribution_logits.float()
+    selected_lp = (
+        logits_f.gather(-1, selected.unsqueeze(-1)).squeeze(-1)
+        - torch.logsumexp(logits_f, dim=-1)
+    )
+    return torch.where(torch.isfinite(selected_lp), selected_lp, torch.zeros_like(selected_lp))
+
+
+def sample_codebook(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+    top_k: int,
+    forbidden_mask: torch.Tensor | None,
+) -> torch.Tensor:
+    """Sample one codebook's tokens from logits (CUDA-graph safe).
+
+    Args:
+        logits: ``[num_tokens, vocab]`` raw codebook logits.
+        temperature: Sampling temperature; ``<= 0`` falls back to argmax.
+        top_k: Top-k truncation width (``<= 0`` disables truncation).
+        forbidden_mask: Optional ``[vocab]`` bool mask; ``True`` entries are
+            set to ``-inf`` before sampling (reserved/special tokens).
+
+    Returns:
+        ``[num_tokens]`` int64 sampled token ids.
+    """
+    logits = _sanitize_logits(logits)
+    if forbidden_mask is not None:
+        logits = logits.masked_fill(forbidden_mask, float("-inf"))
+
+    if temperature <= 0.0:
+        return logits.argmax(dim=-1)
+
+    logits = logits / temperature
 
-# Minimum sampling temperature used inside the compiled graph. The old eager
-# sampler special-cased ``temperature <= 0`` as exact argmax, but a
-# data-dependent branch is illegal inside a captured graph, so we clamp to a tiny
-# value (near-argmax) and always take the Gumbel-top-k path.
-_MIN_SAMPLING_TEMPERATURE = 1e-4
+    if top_k is not None and top_k > 0:
+        vals, idxs = torch.topk(logits, k=min(top_k, logits.size(-1)), dim=-1)
+        sampled_in_k = _gumbel_argmax(vals)
+        return idxs.gather(-1, sampled_in_k.unsqueeze(-1)).squeeze(-1)
+
+    return _gumbel_argmax(logits)
+
+
+def sample_codebook_with_logprobs(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+    top_k: int,
+    forbidden_mask: torch.Tensor | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Sample one codebook and return selected-code log probabilities.
+
+    Sampling follows :func:`sample_codebook` exactly up to the selected token.
+    The extra logprob work happens only after the token is selected so this
+    helper can be used as an observational rollout API.
+    """
+    raw_model_logits = _sanitize_logits(logits)
+    sampling_logits = raw_model_logits
+    if forbidden_mask is not None:
+        sampling_logits = sampling_logits.masked_fill(forbidden_mask, float("-inf"))
+
+    if temperature <= 0.0:
+        sampled = sampling_logits.argmax(dim=-1)
+        selected_sampling_lp = torch.zeros_like(sampled, dtype=torch.float32)
+    else:
+        sampling_logits = sampling_logits / temperature
+        if top_k is not None and top_k > 0:
+            vals, idxs = torch.topk(
+                sampling_logits,
+                k=min(top_k, sampling_logits.size(-1)),
+                dim=-1,
+            )
+            sampled_in_k = _gumbel_argmax(vals)
+            sampled = idxs.gather(-1, sampled_in_k.unsqueeze(-1)).squeeze(-1)
+            selected_sampling_lp = _selected_logprob(vals, sampled_in_k)
+        else:
+            sampled = _gumbel_argmax(sampling_logits)
+            selected_sampling_lp = _selected_logprob(sampling_logits, sampled)
+
+    selected_model_lp = _selected_logprob(raw_model_logits, sampled)
+    return sampled, selected_model_lp, selected_sampling_lp
 
 
 class EasyMagpieLTSelfAttention(nn.Module):
@@ -74,8 +176,8 @@ def __init__(self, d_model: int, n_heads: int) -> None:
         self.o_net = nn.Linear(n_heads * self.d_head, d_model, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        b, t, _ = x.shape
-        qkv = self.qkv_net(x).reshape(b, t, 3, self.n_heads, self.d_head)
+        qkv = self.qkv_net(x)
+        qkv = qkv.reshape(-1, qkv.shape[-2], 3, self.n_heads, self.d_head)
         q, k, v = qkv.unbind(dim=2)  # each [b, t, nh, dh]
         # [b, nh, t, dh]
         q = q.transpose(1, 2)
@@ -84,46 +186,39 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         attn = torch.nn.functional.scaled_dot_product_attention(
             q, k, v, is_causal=True, scale=self.scale
         )
-        attn = attn.transpose(1, 2).contiguous().view(b, t, -1)
+        attn = attn.transpose(1, 2).contiguous().flatten(2)
         return self.o_net(attn)
 
 
 class EasyMagpieLTFeedForward(nn.Module):
     """Positionwise feed-forward network.
 
-    A ``Conv1d(kernel_size=1)`` over the channel dim is mathematically identical
-    to an ``nn.Linear`` applied on the last dim, but the conv form forces a
-    ``[b, t, c] -> [b, c, t]`` transpose on the way in and out (which torch
-    cannot fuse away and which showed up as ``*_transpose_*`` /
-    ``*_convolution_*`` triton kernels in profiling). We therefore use plain
-    bias-free ``nn.Linear`` layers and operate directly on the ``[b, t, c]``
-    layout. The ``conv`` submodule attribute is kept so the kernel-1 conv
-    weights from the training checkpoint (shape ``[out, in, 1]``) still map 1:1;
-    :meth:`EasyMagpieTTS.load_weights` squeezes the trailing singleton dim.
+    Uses ``Conv1d(kernel_size=1)`` layers named ``proj.conv`` and ``o_net.conv``
+    (no bias). A kernel-1 conv is a plain linear over the channel dim, applied
+    with a single transpose and GELU(tanh) in between. The ``Conv1d`` submodule
+    names match the training checkpoint so weights load 1:1.
     """
 
     def __init__(self, d_model: int, d_ffn: int) -> None:
         super().__init__()
-        self.proj = _LinearWrapper(d_model, d_ffn)
-        self.o_net = _LinearWrapper(d_ffn, d_model)
+        self.proj = _Conv1dWrapper(d_model, d_ffn)
+        self.o_net = _Conv1dWrapper(d_ffn, d_model)
         self.act = nn.GELU(approximate="tanh")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: [b, t, c]; no transpose needed for a kernel-1 conv == linear.
-        return self.o_net(self.act(self.proj(x)))
+        # x: [b, t, c] -> conv expects [b, c, t]
+        h = x.transpose(1, 2)
+        h = self.act(self.proj(h))
+        h = self.o_net(h)
+        return h.transpose(1, 2)
 
 
-class _LinearWrapper(nn.Module):
-    """Holds a bias-free ``nn.Linear`` under attribute name ``conv``.
-
-    The attribute is named ``conv`` purely so the parameter path matches the
-    training checkpoint's kernel-1 ``Conv1d`` (``...proj.conv.weight``); the math
-    is a plain dense projection on the channel dim.
-    """
+class _Conv1dWrapper(nn.Module):
+    """Holds a kernel-1 ``Conv1d`` under attribute name ``conv`` (no bias)."""
 
     def __init__(self, in_ch: int, out_ch: int) -> None:
         super().__init__()
-        self.conv = nn.Linear(in_ch, out_ch, bias=False)
+        self.conv = nn.Conv1d(in_ch, out_ch, kernel_size=1, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.conv(x)
@@ -149,14 +244,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
+# NOTE: ``dynamic_arg_dims`` is passed explicitly rather than relying on
+# vLLM's annotation-based inference. This file uses
+# ``from __future__ import annotations`` (PEP 563), so ``forward``'s
+# annotations are stored as strings (``"torch.Tensor"``) and vLLM's
+# ``v.annotation in [torch.Tensor, ...]`` check would never match, raising
+# "No dynamic dimensions found...". ``inputs_embeds`` is
+# ``[num_tokens, num_codebooks, hidden]`` -> dim 0 (num_tokens) is dynamic.
+@support_torch_compile(dynamic_arg_dims={"inputs_embeds": 0})
 class EasyMagpieLocalTransformer(nn.Module):
-    """Causal transformer stack with learnable positional embeddings.
+    """Compiled causal transformer stack with learnable positional embeddings.
 
-    Plain (uncompiled) module: it is invoked from inside
-    :class:`EasyMagpieCodeLoop`'s compiled forward, so it gets *inlined* into
-    that single captured graph rather than being compiled / replayed on its own.
-    Holds learnable ``position_embeddings``, the stacked ``layers.{i}.*`` and a
-    no-op ``norm_out`` (names match the training checkpoint).
+    Decorated with ``@support_torch_compile`` so vLLM captures a single CUDA
+    graph for the fixed ``(num_tokens, num_stacked_codebooks, d_model)`` input
+    shape. Holds learnable ``position_embeddings``, the stacked ``layers.{i}.*``
+    and a no-op ``norm_out``.
     """
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
@@ -170,11 +272,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         max_len = arch.num_stacked_codebooks + 2
 
         self.position_embeddings = nn.Embedding(max_len, d_model)
-        # Cache a constant ``arange`` so we don't re-materialize it (and re-run an
-        # embedding gather) on every autoregressive step. The positional table is
-        # tiny and fixed; gathering once per forward over a cached index avoids
-        # the ``arange + embedding`` triton kernel seen in profiling.
-        self.register_buffer("_positions", torch.arange(max_len), persistent=False)
         self.layers = nn.ModuleList(
             [EasyMagpieLTLayer(d_model, d_ffn, n_heads) for _ in range(n_layers)]
         )
@@ -182,102 +279,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
 
     def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
         seq_len = inputs_embeds.shape[1]
-        pos_emb = self.position_embeddings(self._positions[:seq_len])
-        x = inputs_embeds + pos_emb.unsqueeze(0)
+        positions = torch.arange(seq_len, device=inputs_embeds.device)
+        x = inputs_embeds + self.position_embeddings(positions).unsqueeze(0)
         for layer in self.layers:
             x = layer(x)
         return self.norm_out(x)
 
 
-# NOTE: ``dynamic_arg_dims`` is passed explicitly rather than relying on vLLM's
-# annotation-based inference. This file uses ``from __future__ import
-# annotations`` (PEP 563), so ``forward``'s annotations are stored as strings
-# (``"torch.Tensor"``) and vLLM's ``v.annotation in [torch.Tensor, ...]`` check
-# would never match, raising "No dynamic dimensions found...". Both ``dec_hidden``
-# and ``gumbel_noise`` are ``[num_tokens, ...]`` -> dim 0 (num_tokens) is dynamic.
-@support_torch_compile(dynamic_arg_dims={"dec_hidden": 0, "gumbel_noise": 0})
-class EasyMagpieCodeLoop(nn.Module):
-    """Compiled single-graph autoregressive codebook loop.
-
-    Runs the *entire* per-frame loop — transformer stack, per-codebook projection
-    heads, and (graph-safe) sampling — under one ``@support_torch_compile`` graph,
-    so vLLM captures a single CUDA graph replayed once per frame instead of
-    replaying the transformer ``N`` times with eager projection / sampling in
-    between. (Total FLOPs are unchanged — this removes per-step Python and
-    kernel-launch overhead, which dominates at low concurrency.)
-
-    It owns **no parameters**: the projection / embedding / out-projection modules
-    and the forbidden mask live on the parent :class:`EasyMagpieCodePredictor` (so
-    the checkpoint still loads 1:1) and are reached through a non-registered
-    reference set by :meth:`bind_predictor`.
-
-    Sampling is kept graph-safe by construction:
-
-    * the Gumbel noise is drawn eagerly *outside* the graph and injected as
-      ``gumbel_noise`` — running ``uniform_`` inside the capture would otherwise
-      reuse the captured random numbers on every replay;
-    * ``temperature`` is a runtime tensor, so per-request temperature works
-      without recompiling;
-    * ``top_k`` shapes the ``topk`` / noise tensors and is therefore a
-      **capture-time constant** (per-request ``top_k`` changes are not honored
-      once the graph is captured).
-    """
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        super().__init__()
-        arch = EasyMagpieOmniArch.from_hf_config(vllm_config.model_config.hf_config)
-        self.num_codebooks = arch.num_stacked_codebooks
-        self.lt_hidden = arch.local_transformer_hidden_dim
-        self.top_k = min(_DEFAULT_TOP_K, arch.num_all_tokens_per_codebook)
-        # Set by :meth:`bind_predictor`; held in a tuple so nn.Module does not
-        # register the parent as a submodule (which would duplicate params).
-        self._predictor_ref: tuple = ()
-
-    def bind_predictor(self, predictor: "EasyMagpieCodePredictor") -> None:
-        self._predictor_ref = (predictor,)
-        self.top_k = predictor._sample_top_k
-
-    def forward(
-        self,
-        dec_hidden: torch.Tensor,
-        gumbel_noise: torch.Tensor,
-        temperature: torch.Tensor,
-    ) -> torch.Tensor:
-        """Sample all ``num_codebooks`` codes for every frame in one graph.
-
-        Args:
-            dec_hidden: ``[num_tokens, embedding_dim]`` backbone hidden state.
-            gumbel_noise: ``[num_tokens, num_codebooks, top_k]`` pre-drawn
-                Gumbel noise (``-log(-log(u))``), one slice per codebook.
-            temperature: ``[1]`` sampling temperature (already clamped > 0).
-
-        Returns:
-            ``[num_tokens, num_codebooks]`` int64 sampled codes.
-        """
-        cp = self._predictor_ref[0]
-        num_tokens = dec_hidden.shape[0]
-        n = self.num_codebooks
-
-        buf = dec_hidden.new_zeros(num_tokens, n, self.lt_hidden)
-        buf[:, 0, :] = cp.local_transformer_in_projection(dec_hidden)
-
-        forbidden = cp.forbidden_mask
-        codes: list[torch.Tensor] = []
-        for k in range(n):
-            hidden = cp.local_transformer(buf)
-            row = cp.local_transformer_audio_out_projection(hidden[:, k, :])
-            logits = cp.local_transformer_out_projections[k](row)
-            logits = logits.masked_fill(forbidden, float("-inf")) / temperature
-            vals, idxs = torch.topk(logits, self.top_k, dim=-1)
-            picked = (vals + gumbel_noise[:, k, :]).argmax(dim=-1, keepdim=True)
-            code_k = idxs.gather(-1, picked).squeeze(-1)
-            codes.append(code_k)
-            if k + 1 < n:
-                emb = cp.audio_in_projection(cp.audio_embeddings[k](code_k))
-                buf[:, k + 1, :] = cp.local_transformer_in_projection(emb)
-        return torch.stack(codes, dim=1)
-
-
 class EasyMagpieCodePredictor(nn.Module):
     """Autoregressive intra-frame codebook predictor (the "local transformer").
 
@@ -350,30 +358,48 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             persistent=False,
         )
 
-        # Sampling knobs (overridable from the outer model / request). ``top_k``
-        # is captured into the compiled loop graph (see ``EasyMagpieCodeLoop``),
-        # so per-request ``top_k`` changes are not honored once captured;
-        # per-request ``temperature`` is, since it is fed as a runtime tensor.
+        # Sampling knobs (overridable from the outer model / request).
         self.temperature: float = 0.7
-        self.top_k: int = _DEFAULT_TOP_K
-        self.lt_hidden = lt_hidden
-        self._sample_top_k = min(self.top_k, self.num_tokens_per_codebook)
-
-        # Compiled single-graph autoregressive loop (owns no params; reaches the
-        # projection heads / embeddings / mask on ``self`` via a bound reference).
-        self._code_loop = EasyMagpieCodeLoop(vllm_config=vllm_config, prefix=f"{prefix}.code_loop")
-        self._code_loop.bind_predictor(self)
+        self.top_k: int = 80
 
         # ── Persistent address-stable scratch buffers ──────────────────
-        # (created on the CUDA default device that vLLM sets during model init).
         max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
         dtype = vllm_config.model_config.dtype
-        # Stable-address input for the captured loop graph.
-        self._dec_hidden_buf = torch.zeros(max_num_tokens, self.embedding_dim, dtype=dtype)
-        # Gumbel noise drawn eagerly each frame and injected into the graph; fp32
-        # so the small ``-log(-log(u))`` values don't underflow in fp16.
-        self._gumbel_buf = torch.zeros(max_num_tokens, self.num_codebooks, self._sample_top_k, dtype=torch.float32)
-        self._temperature_buf = torch.zeros(1, dtype=torch.float32)
+        self._buf_inputs = torch.zeros(max_num_tokens, self.num_codebooks, lt_hidden, dtype=dtype)
+        self._out_codes = torch.zeros(max_num_tokens, self.num_codebooks, dtype=torch.long)
+        self._out_code_logprobs = torch.zeros(max_num_tokens, self.num_codebooks, dtype=torch.float32)
+        self._out_code_sampling_logprobs = torch.zeros(max_num_tokens, self.num_codebooks, dtype=torch.float32)
+        self._max_num_tokens = int(max_num_tokens)
+        self._local_transformer_graph_tokens = self._resolve_local_transformer_graph_tokens(max_num_tokens)
+
+    @staticmethod
+    def _resolve_local_transformer_graph_tokens(max_num_tokens: int) -> int:
+        """Return the explicit stable local-transformer batch size.
+
+        Request concurrency must not implicitly enable local-transformer
+        padding: that optimization changes generated audio quality for this
+        model. Keep the audio-safe default unpadded unless the graph-token
+        knob is set deliberately for a perf experiment.
+        """
+
+        raw_value = (os.environ.get("EASYMAGPIE_LOCAL_TRANSFORMER_GRAPH_TOKENS") or "").strip()
+        if not raw_value:
+            return 0
+        try:
+            value = int(raw_value)
+        except ValueError:
+            return 0
+        if value <= 0:
+            return 0
+        return min(int(max_num_tokens), value)
+
+    def _local_transformer_run_tokens(self, num_tokens: int) -> int:
+        """Pad active rows so CUDA graph replay sees a stable leading dimension."""
+
+        graph_tokens = int(self._local_transformer_graph_tokens or 0)
+        if graph_tokens <= 0 or num_tokens >= graph_tokens:
+            return num_tokens
+        return graph_tokens
 
     @torch.no_grad()
     def init_forbidden_mask(self) -> None:
@@ -413,15 +439,14 @@ def embed_audio_frame(self, codes: torch.Tensor) -> torch.Tensor:
         acc = acc / self.num_codebooks
         return self.audio_in_projection(acc)
 
+    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
+        """Run the compiled local transformer over the input buffer."""
+        return self.local_transformer(inputs_embeds)
+
     @torch.no_grad()
     def generate_codes(self, dec_hidden: torch.Tensor) -> torch.Tensor:
         """Autoregressively sample all ``C * S`` codebooks for each frame.
 
-        Draws this frame's Gumbel noise eagerly into a stable buffer (fresh
-        randomness per frame, outside the captured graph) and stages the inputs
-        at fixed addresses, then runs the whole loop as a single captured graph
-        via :class:`EasyMagpieCodeLoop`.
-
         Args:
             dec_hidden: ``[num_tokens, hidden]`` backbone hidden state (one row
                 per frame being decoded).
@@ -429,14 +454,53 @@ def generate_codes(self, dec_hidden: torch.Tensor) -> torch.Tensor:
         Returns:
             ``[num_tokens, num_codebooks]`` int64 sampled codes.
         """
-        num_tokens = dec_hidden.shape[0]
-        in_buf = self._dec_hidden_buf[:num_tokens]
-        in_buf.copy_(dec_hidden)
+        codes, _, _ = self.generate_codes_with_logprobs(dec_hidden)
+        return codes
 
-        # ``-log(-log(u))`` Gumbel noise, computed in place in fp32.
-        noise = self._gumbel_buf[:num_tokens]
-        noise.uniform_(1e-20, 1.0 - 1e-20)
-        noise.log_().neg_().log_().neg_()
+    @torch.no_grad()
+    def generate_codes_with_logprobs(
+        self, dec_hidden: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Autoregressively sample codebooks and keep selected-code logprobs.
 
-        self._temperature_buf.fill_(max(float(self.temperature), _MIN_SAMPLING_TEMPERATURE))
-        return self._code_loop(in_buf, noise, self._temperature_buf)
+        Returns:
+            ``(codes, model_logprobs, sampling_logprobs)``, each shaped
+            ``[num_tokens, num_codebooks]``.
+        """
+        num_tokens = dec_hidden.shape[0]
+        run_tokens = self._local_transformer_run_tokens(int(num_tokens))
+        buf_run = self._buf_inputs[:run_tokens]
+        buf = buf_run[:num_tokens]
+        out = self._out_codes[:num_tokens]
+        out_logprobs = self._out_code_logprobs[:num_tokens]
+        out_sampling_logprobs = self._out_code_sampling_logprobs[:num_tokens]
+        buf_run.zero_()
+        out_logprobs.zero_()
+        out_sampling_logprobs.zero_()
+
+        # Row 0: projected backbone hidden state (the AR "prompt").
+        buf[:, 0, :] = self.local_transformer_in_projection(dec_hidden)
+
+        # Always pass the mask unconditionally. An all-False mask makes
+        # ``masked_fill`` a no-op, so there's no need to guard with
+        # ``forbidden_mask.any()`` — and that guard is a data-dependent
+        # host sync that is illegal during CUDA-graph capture.
+        forbidden = self.forbidden_mask
+        for k in range(self.num_codebooks):
+            hidden = self(buf_run)  # compiled transformer over the stable-shape buffer
+            row = self.local_transformer_audio_out_projection(hidden[:num_tokens, k, :])
+            logits = self.local_transformer_out_projections[k](row)
+            code_k, model_lp, sampling_lp = sample_codebook_with_logprobs(
+                logits,
+                temperature=self.temperature,
+                top_k=self.top_k,
+                forbidden_mask=forbidden,
+            )
+            out[:, k] = code_k
+            out_logprobs[:, k] = model_lp
+            out_sampling_logprobs[:, k] = sampling_lp
+            if k + 1 < self.num_codebooks:
+                emb = self.audio_in_projection(self.audio_embeddings[k](code_k))
+                buf[:, k + 1, :] = self.local_transformer_in_projection(emb)
+
+        return out[:num_tokens], out_logprobs[:num_tokens], out_sampling_logprobs[:num_tokens]
diff --git a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/scheduler.py b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/scheduler.py
index 6f82bb23332e..ec4eee6128dc 100644
--- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/scheduler.py
+++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/scheduler.py
@@ -31,12 +31,34 @@
 """
 from __future__ import annotations
 
-from vllm.v1.request import Request, StreamingUpdate
+import os
+import logging
 
-from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler
+from vllm.v1.core.sched.request_queue import create_request_queue
+from vllm.v1.request import Request
 
+try:
+    from vllm.v1.request import StreamingUpdate
+except ImportError:
+    StreamingUpdate = object
 
-class EasyMagpieARAsyncScheduler(OmniARAsyncScheduler):
+try:
+    from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler as _OmniARBaseScheduler
+except ImportError:
+    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler as _OmniARBaseScheduler
+
+
+logger = logging.getLogger(__name__)
+
+
+def _env_flag(name: str, default: bool) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    return value.strip().lower() not in {"0", "false", "no", "off"}
+
+
+class EasyMagpieARAsyncScheduler(_OmniARBaseScheduler):
     """``OmniARAsyncScheduler`` that forwards per-chunk ``additional_information``.
 
     Replace (not merge) is the correct session-level semantics: the session field
@@ -50,6 +72,231 @@ class EasyMagpieARAsyncScheduler(OmniARAsyncScheduler):
     value; a truly empty chunk leaves the previous payload intact).
     """
 
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.disable_mixed_prefill_decode = _env_flag(
+            "EASYMAGPIE_DISABLE_MIXED_PREFILL_DECODE", True
+        )
+        self.purge_stale_train_before_post_refit = _env_flag(
+            "EASYMAGPIE_PURGE_STALE_TRAIN_BEFORE_POST_REFIT", True
+        )
+
+    @staticmethod
+    def _request_id(request: Request) -> str:
+        return str(getattr(request, "request_id", ""))
+
+    @staticmethod
+    def _request_ids_head(requests, limit: int = 8) -> list[str]:
+        ids: list[str] = []
+        for request in requests or []:
+            ids.append(EasyMagpieARAsyncScheduler._request_id(request))
+            if len(ids) >= limit:
+                break
+        return ids
+
+    @staticmethod
+    def _safe_len(value) -> int | None:
+        try:
+            return len(value)
+        except Exception:
+            return None
+
+    @staticmethod
+    def _is_train_rollout_request(request: Request) -> bool:
+        request_id = EasyMagpieARAsyncScheduler._request_id(request)
+        return request_id.startswith(
+            (
+                "easymagpie-grpo-train-step-",
+                "easymagpie-grpo-train-rank-",
+            )
+        )
+
+    @staticmethod
+    def _is_post_refit_request(request: Request) -> bool:
+        request_id = EasyMagpieARAsyncScheduler._request_id(request)
+        return request_id.startswith(
+            (
+                "easymagpie-grpo-train-post-refit-step-",
+                "easymagpie-grpo-train-post-refit-rank-",
+            )
+        )
+
+    def _has_waiting_post_refit_request(self) -> bool:
+        return any(self._is_post_refit_request(request) for request in self.waiting)
+
+    def _purge_stale_train_running_before_post_refit(self) -> list[str]:
+        if (
+            not getattr(self, "purge_stale_train_before_post_refit", True)
+            or not self.waiting
+            or not self._has_waiting_post_refit_request()
+            or not self.running
+        ):
+            return []
+
+        kept = []
+        purged = []
+        for request in self.running:
+            if self._is_train_rollout_request(request):
+                purged.append(self._request_id(request))
+            else:
+                kept.append(request)
+
+        if not purged:
+            return []
+        try:
+            self.running[:] = kept
+        except TypeError:
+            self.running = kept
+        self._easymagpie_last_post_refit_purge = {
+            "num_purged_running": len(purged),
+            "purged_running_head": purged[:16],
+            "purged_running_tail": purged[-16:],
+        }
+        logger.info(
+            "EasyMagpie post-refit scheduler purged stale train running requests: "
+            "num_purged=%d waiting_head=%s purged_head=%s",
+            len(purged),
+            self._request_ids_head(self.waiting),
+            purged[:8],
+        )
+        return purged
+
+    @staticmethod
+    def _is_active_decode_request(request: Request) -> bool:
+        is_finished = getattr(request, "is_finished", None)
+        if callable(is_finished) and is_finished():
+            return False
+
+        status = getattr(request, "status", None)
+        status_name = getattr(status, "name", None)
+        if status_name is None and status is not None:
+            status_name = str(status)
+        if status_name is not None and status_name != "RUNNING":
+            return False
+
+        pending = EasyMagpieARAsyncScheduler._num_pending_tokens(request)
+        if pending is not None and pending <= 0:
+            return False
+
+        # Older tests/stubs may not carry vLLM RequestStatus. If no status is
+        # present, keep the conservative old behavior and treat it as active.
+        return True
+
+    @staticmethod
+    def _num_pending_tokens(request: Request) -> int | None:
+        required = (
+            "num_tokens_with_spec",
+            "num_output_placeholders",
+            "num_computed_tokens",
+        )
+        if not all(hasattr(request, name) for name in required):
+            return None
+        return (
+            int(getattr(request, "num_tokens_with_spec") or 0)
+            + int(getattr(request, "num_output_placeholders") or 0)
+            - int(getattr(request, "num_computed_tokens") or 0)
+        )
+
+    def _has_active_decode_requests(self) -> bool:
+        return any(self._is_active_decode_request(request) for request in self.running)
+
+    @staticmethod
+    def _scheduled_token_count(schedule_output) -> int | None:
+        num_scheduled_tokens = getattr(schedule_output, "num_scheduled_tokens", None)
+        if num_scheduled_tokens is None:
+            return None
+        if isinstance(num_scheduled_tokens, dict):
+            try:
+                return sum(int(count or 0) for count in num_scheduled_tokens.values())
+            except Exception:
+                return None
+        try:
+            return int(num_scheduled_tokens or 0)
+        except Exception:
+            return None
+
+    @staticmethod
+    def _restore_deferred_waiting(current_waiting, deferred_waiting):
+        if current_waiting:
+            for request in deferred_waiting:
+                current_waiting.add_request(request)
+            return current_waiting
+        return deferred_waiting
+
+    def schedule(self):  # type: ignore[override]
+        post_refit_waiting = self._has_waiting_post_refit_request()
+        if post_refit_waiting:
+            logger.info(
+                "EasyMagpie post-refit scheduler before schedule: waiting=%s running=%s "
+                "waiting_head=%s running_head=%s active_decode=%s",
+                self._safe_len(self.waiting),
+                self._safe_len(self.running),
+                self._request_ids_head(self.waiting),
+                self._request_ids_head(self.running),
+                self._has_active_decode_requests(),
+            )
+        self._purge_stale_train_running_before_post_refit()
+
+        if (
+            not self.disable_mixed_prefill_decode
+            or not self._has_active_decode_requests()
+            or not self.waiting
+        ):
+            result = super().schedule()
+            if post_refit_waiting:
+                logger.info(
+                    "EasyMagpie post-refit scheduler after direct schedule: scheduled_tokens=%s "
+                    "waiting=%s running=%s waiting_head=%s running_head=%s",
+                    self._scheduled_token_count(result),
+                    self._safe_len(self.waiting),
+                    self._safe_len(self.running),
+                    self._request_ids_head(self.waiting),
+                    self._request_ids_head(self.running),
+                )
+            return result
+
+        # The EasyMagpie Mamba path is unstable when vLLM batches cached
+        # one-token decodes together with fresh prompt prefills.
+        deferred_waiting = self.waiting
+        self.waiting = create_request_queue(self.policy)
+        restored_waiting = False
+        try:
+            result = super().schedule()
+            scheduled_tokens = self._scheduled_token_count(result)
+            self._easymagpie_last_mixed_guard = {
+                "deferred_waiting": len(deferred_waiting),
+                "guarded_scheduled_tokens": scheduled_tokens,
+                "fallback_to_waiting_prefill": scheduled_tokens == 0,
+            }
+            if scheduled_tokens == 0:
+                self.waiting = self._restore_deferred_waiting(self.waiting, deferred_waiting)
+                restored_waiting = True
+                result = super().schedule()
+                if post_refit_waiting:
+                    logger.info(
+                        "EasyMagpie post-refit scheduler after guarded fallback: scheduled_tokens=%s "
+                        "waiting=%s running=%s waiting_head=%s running_head=%s",
+                        self._scheduled_token_count(result),
+                        self._safe_len(self.waiting),
+                        self._safe_len(self.running),
+                        self._request_ids_head(self.waiting),
+                        self._request_ids_head(self.running),
+                    )
+                return result
+            if post_refit_waiting:
+                logger.info(
+                    "EasyMagpie post-refit scheduler after guarded decode schedule: scheduled_tokens=%s "
+                    "deferred_waiting=%s waiting=%s running=%s",
+                    scheduled_tokens,
+                    self._safe_len(deferred_waiting),
+                    self._safe_len(self.waiting),
+                    self._safe_len(self.running),
+                )
+            return result
+        finally:
+            if not restored_waiting:
+                self.waiting = self._restore_deferred_waiting(self.waiting, deferred_waiting)
+
     def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None:
         super()._update_request_as_session(session, update)
 
diff --git a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py
new file mode 100644
index 000000000000..c497a290cbb3
--- /dev/null
+++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py
@@ -0,0 +1,1902 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Focused vLLM/vLLM-Omni shims for EasyMagpie NeMo-RL rollouts.
+
+This review branch intentionally keeps only the compatibility surface used by
+EasyMagpie RL serving: worker refit RPCs, dense tensor serialization for refit
+payloads, light input/import aliases, and AsyncOmni abstract-method fill-ins.
+The older broad ``full/all`` vLLM-Omni compatibility layer is not part of the
+reviewable RL patch series.
+"""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+import logging
+import os
+import queue as stdlib_queue
+import sys
+import types
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _install_vllm_inputs_data_alias() -> None:
+    if "vllm.inputs.data" in sys.modules:
+        return
+    try:
+        import vllm.inputs as inputs
+    except Exception:
+        return
+    try:
+        if importlib.util.find_spec("vllm.inputs.data") is not None:
+            return
+    except Exception:
+        pass
+
+    module = types.ModuleType("vllm.inputs.data")
+    module.__doc__ = "Compatibility aliases for vLLM input types exported from vllm.inputs."
+    module.__package__ = "vllm.inputs"
+    for name in dir(inputs):
+        if not name.startswith("__"):
+            setattr(module, name, getattr(inputs, name))
+    for old_name, new_name in {
+        "TokenInputs": "TokensInput",
+        "EmbedsInputs": "EmbedsInput",
+        "SingletonInputs": "SingletonInput",
+    }.items():
+        if not hasattr(module, old_name) and hasattr(inputs, new_name):
+            setattr(module, old_name, getattr(inputs, new_name))
+
+    sys.modules["vllm.inputs.data"] = module
+    sys.modules.setdefault("vllm.inputs.parse", module)
+    if not hasattr(inputs, "data"):
+        inputs.data = module
+    if not hasattr(inputs, "parse"):
+        inputs.parse = sys.modules["vllm.inputs.parse"]
+
+
+def _install_vllm_multimodal_inputs_alias() -> None:
+    """Expose vLLM-Omni's legacy plural multimodal input alias on vLLM 0.21+."""
+
+    try:
+        multimodal_inputs = importlib.import_module("vllm.multimodal.inputs")
+    except Exception:
+        return
+    if hasattr(multimodal_inputs, "MultiModalInputs"):
+        return
+
+    candidate = None
+    for name in ("MultiModalInputsV2", "MultiModalInput"):
+        candidate = getattr(multimodal_inputs, name, None)
+        if candidate is not None:
+            break
+    if candidate is None:
+        try:
+            inputs = importlib.import_module("vllm.inputs")
+            candidate = getattr(inputs, "MultiModalInput", None)
+        except Exception:
+            candidate = None
+    if candidate is None:
+        candidate = dict
+    multimodal_inputs.MultiModalInputs = candidate
+
+def _install_engine_utils_compat() -> None:
+    try:
+        import vllm.v1.engine.utils as engine_utils
+        from vllm.v1.utils import get_engine_client_zmq_addr
+    except Exception:
+        return
+
+    if not hasattr(engine_utils, "get_engine_zmq_addresses"):
+
+        def get_engine_zmq_addresses(vllm_config: Any, num_api_servers: int = 1):
+            parallel_config = vllm_config.parallel_config
+            dp_size = parallel_config.data_parallel_size
+            local_engine_count = parallel_config.data_parallel_size_local
+            local_start_index = parallel_config.data_parallel_rank_local
+            host = parallel_config.data_parallel_master_ip
+            local_only = (
+                local_start_index is not None
+                or parallel_config.data_parallel_hybrid_lb
+                or parallel_config.data_parallel_external_lb
+                or local_engine_count == dp_size
+            )
+            return engine_utils.EngineZmqAddresses(
+                inputs=[get_engine_client_zmq_addr(local_only, host) for _ in range(num_api_servers)],
+                outputs=[get_engine_client_zmq_addr(local_only, host) for _ in range(num_api_servers)],
+            )
+
+        engine_utils.get_engine_zmq_addresses = get_engine_zmq_addresses
+
+    launch_core_engines = getattr(engine_utils, "launch_core_engines", None)
+    if launch_core_engines is not None and not getattr(launch_core_engines, "_easymagpie_compat", False):
+        signature = inspect.signature(launch_core_engines)
+        parameters = signature.parameters
+        accepts_addresses = "addresses" in parameters
+        positional_parameters = [
+            name
+            for name, parameter in parameters.items()
+            if parameter.kind
+            in {
+                inspect.Parameter.POSITIONAL_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            }
+        ]
+        addresses_pos = positional_parameters.index("addresses") if "addresses" in positional_parameters else None
+        num_api_servers_pos = (
+            positional_parameters.index("num_api_servers") if "num_api_servers" in positional_parameters else None
+        )
+
+        class LaunchCoreEnginesContextCompat:
+            def __init__(self, context_manager: Any) -> None:
+                self._context_manager = context_manager
+
+            def __enter__(self) -> Any:
+                value = self._context_manager.__enter__()
+                if isinstance(value, tuple) and len(value) > 3:
+                    return value[:3]
+                return value
+
+            def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any:
+                return self._context_manager.__exit__(exc_type, exc, tb)
+
+        def compat_launch_core_engines(*args: Any, addresses: Any = None, **kwargs: Any):
+            if not accepts_addresses:
+                return launch_core_engines(*args, **kwargs)
+
+            call_kwargs = dict(kwargs)
+            address_supplied_positionally = addresses_pos is not None and len(args) > addresses_pos
+            if not address_supplied_positionally and "addresses" not in call_kwargs:
+                if addresses is None:
+                    vllm_config = call_kwargs.get("vllm_config")
+                    if vllm_config is None and args:
+                        vllm_config = args[0]
+                    num_api_servers = call_kwargs.get("num_api_servers")
+                    if num_api_servers is None and num_api_servers_pos is not None and len(args) > num_api_servers_pos:
+                        num_api_servers = args[num_api_servers_pos]
+                    if num_api_servers is None:
+                        num_api_servers_parameter = parameters.get("num_api_servers")
+                        if num_api_servers_parameter is not None:
+                            num_api_servers = num_api_servers_parameter.default
+                    if num_api_servers is inspect.Parameter.empty:
+                        num_api_servers = 1
+                    addresses = engine_utils.get_engine_zmq_addresses(vllm_config, int(num_api_servers or 1))
+                call_kwargs["addresses"] = addresses
+            launch_context = launch_core_engines(*args, **call_kwargs)
+            if hasattr(launch_context, "__enter__") and hasattr(launch_context, "__exit__"):
+                return LaunchCoreEnginesContextCompat(launch_context)
+            return launch_context
+
+        compat_launch_core_engines._easymagpie_compat = True  # type: ignore[attr-defined]
+        engine_utils.launch_core_engines = compat_launch_core_engines
+
+def _install_omni_gpu_model_runner_init_kwargs_compat() -> None:
+    """Bridge vLLM-Omni's optional ``num_tokens`` arg across vLLM versions."""
+
+    try:
+        module = importlib.import_module("vllm_omni.worker.gpu_model_runner")
+    except Exception:
+        return
+
+    runner_cls = getattr(module, "OmniGPUModelRunner", None)
+    if runner_cls is None:
+        return
+
+    current = getattr(runner_cls, "_init_model_kwargs", None)
+    if current is None or getattr(current, "_easymagpie_init_model_kwargs_compat", False):
+        return
+
+    def _base_init_model_kwargs(self: Any, num_tokens: int | None = None) -> Any:
+        if num_tokens is None:
+            num_tokens = int(getattr(self, "max_num_tokens", 0) or 0)
+        mro = type(self).__mro__
+        try:
+            base_classes = mro[mro.index(runner_cls) + 1 :]
+        except ValueError:
+            base_classes = mro[1:]
+        for base_cls in base_classes:
+            base_method = base_cls.__dict__.get("_init_model_kwargs")
+            if base_method is None:
+                continue
+            if getattr(base_method, "_easymagpie_init_model_kwargs_compat", False):
+                continue
+            try:
+                signature = inspect.signature(base_method)
+            except (TypeError, ValueError):
+                return base_method(self, num_tokens)
+            positional = [
+                parameter
+                for parameter in signature.parameters.values()
+                if parameter.kind
+                in {
+                    inspect.Parameter.POSITIONAL_ONLY,
+                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                }
+            ]
+            accepts_varargs = any(
+                parameter.kind == inspect.Parameter.VAR_POSITIONAL
+                for parameter in signature.parameters.values()
+            )
+            if accepts_varargs or len(positional) >= 2:
+                return base_method(self, num_tokens)
+            return base_method(self)
+        return {}
+
+    _base_init_model_kwargs._easymagpie_init_model_kwargs_compat = True  # type: ignore[attr-defined]
+    _base_init_model_kwargs._easymagpie_original = current  # type: ignore[attr-defined]
+    runner_cls._init_model_kwargs = _base_init_model_kwargs
+
+
+def _install_omni_cumsum_arange_compat() -> None:
+    """Supply vLLM 0.21's required ``arange_out`` buffer when Omni calls old API."""
+
+    try:
+        import numpy as np
+
+        module = importlib.import_module("vllm_omni.worker.gpu_model_runner")
+    except Exception:
+        return
+
+    runner_cls = getattr(module, "OmniGPUModelRunner", None)
+    if runner_cls is None:
+        return
+
+    current = getattr(runner_cls, "_get_cumsum_and_arange", None)
+    if current is None or getattr(current, "_easymagpie_cumsum_arange_compat", False):
+        return
+
+    try:
+        signature = inspect.signature(current)
+    except (TypeError, ValueError):
+        return
+
+    parameters = signature.parameters
+    if "arange_out" not in parameters:
+        return
+
+    positional_parameters = [
+        name
+        for name, parameter in parameters.items()
+        if parameter.kind
+        in {
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        }
+    ]
+    remaining_positional = [name for name in positional_parameters if name not in {"self", "num_tokens"}]
+    arange_out_args_index = remaining_positional.index("arange_out") if "arange_out" in remaining_positional else None
+
+    def _make_arange_out(self: Any, num_tokens: Any) -> Any:
+        num_tokens_np = np.asarray(num_tokens)
+        total_num_tokens = int(num_tokens_np.sum(dtype=np.int64))
+        arange_np = getattr(self, "arange_np", None)
+        dtype = getattr(arange_np, "dtype", np.int64)
+        return np.empty((total_num_tokens,), dtype=dtype)
+
+    def _get_cumsum_and_arange(self: Any, num_tokens: Any, *args: Any, **kwargs: Any) -> Any:
+        supplied_positionally = arange_out_args_index is not None and len(args) > arange_out_args_index
+        supplied_arange_out = None
+        if supplied_positionally:
+            supplied_arange_out = args[arange_out_args_index]
+        elif "arange_out" in kwargs:
+            supplied_arange_out = kwargs["arange_out"]
+        if not supplied_positionally and "arange_out" not in kwargs:
+            supplied_arange_out = _make_arange_out(self, num_tokens)
+            kwargs["arange_out"] = supplied_arange_out
+        result = current(self, num_tokens, *args, **kwargs)
+        if isinstance(result, tuple) and len(result) >= 2:
+            return result
+        return result, supplied_arange_out
+
+    _get_cumsum_and_arange._easymagpie_cumsum_arange_compat = True  # type: ignore[attr-defined]
+    _get_cumsum_and_arange._easymagpie_original = current  # type: ignore[attr-defined]
+    runner_cls._get_cumsum_and_arange = _get_cumsum_and_arange
+
+
+def _install_omni_batch_execution_padding_compat() -> None:
+    """Let vLLM 0.21 CUDA-graph capture use Omni's fallback batch descriptor.
+
+    The vLLM-Omni branch used by EasyMagpie includes a compatibility fallback
+    for ``_determine_batch_execution_and_padding`` that always returns
+    ``CUDAGraphMode.NONE``. Newer vLLM V1 passes the runtime capture mode
+    explicitly into ``_dummy_run`` and asserts it matches the returned mode.
+    Preserve Omni's simple descriptor while reporting the configured runtime
+    graph mode only for descriptors vLLM captured at startup. Larger runtime
+    batches must stay eager; otherwise vLLM attempts a lazy CUDA-graph capture
+    after capture has been disabled.
+
+    Also normalize Omni's fallback descriptor into a value-keyed object. The
+    fallback descriptor class is created inside the method and is identity
+    hashed, so vLLM's piecewise CUDA-graph wrapper cannot match runtime
+    descriptors against the descriptors captured at engine startup. When that
+    happens, vLLM attempts a lazy graph capture after capture has been disabled
+    and raises ``CUDA graph capturing detected at an inappropriate time``.
+    """
+
+    try:
+        module = importlib.import_module("vllm_omni.worker.gpu_model_runner")
+        from vllm.config import CUDAGraphMode
+    except Exception:
+        return
+
+    runner_cls = getattr(module, "OmniGPUModelRunner", None)
+    if runner_cls is None:
+        return
+
+    current = getattr(runner_cls, "_determine_batch_execution_and_padding", None)
+    if current is None or getattr(current, "_easymagpie_batch_execution_padding_compat", False):
+        return
+
+    def _runtime_mode_from_config(self: Any, *, force_uniform_decode: bool) -> Any:
+        compilation_config = getattr(getattr(self, "vllm_config", None), "compilation_config", None)
+        configured_mode = getattr(compilation_config, "cudagraph_mode", None)
+        if configured_mode is None:
+            return CUDAGraphMode.NONE
+        if force_uniform_decode and hasattr(configured_mode, "decode_mode"):
+            return configured_mode.decode_mode()
+        if hasattr(configured_mode, "mixed_mode"):
+            return configured_mode.mixed_mode()
+        return configured_mode
+
+    class _ValueBatchDescriptor:
+        __slots__ = ("num_tokens", "num_reqs", "uniform_decode")
+
+        def __init__(
+            self,
+            *,
+            num_tokens: int,
+            num_reqs: int | None = None,
+            uniform_decode: bool = False,
+        ) -> None:
+            self.num_tokens = int(num_tokens)
+            self.num_reqs = None if num_reqs is None else int(num_reqs)
+            self.uniform_decode = bool(uniform_decode)
+
+        @property
+        def non_uniform(self) -> "_ValueBatchDescriptor":
+            return _ValueBatchDescriptor(
+                num_tokens=self.num_tokens,
+                num_reqs=self.num_reqs,
+                uniform_decode=False,
+            )
+
+        def __eq__(self, other: Any) -> bool:
+            return (
+                hasattr(other, "num_tokens")
+                and int(getattr(other, "num_tokens")) == self.num_tokens
+                and bool(getattr(other, "uniform_decode", False)) == self.uniform_decode
+            )
+
+        def __hash__(self) -> int:
+            return hash((self.num_tokens, self.uniform_decode))
+
+        def __repr__(self) -> str:
+            return (
+                f"EasyMagpieBatchDescriptor(num_tokens={self.num_tokens}, "
+                f"uniform_decode={self.uniform_decode}, num_reqs={self.num_reqs})"
+            )
+
+    def _normalize_batch_descriptor(batch_desc: Any) -> Any:
+        if batch_desc is None or isinstance(batch_desc, _ValueBatchDescriptor):
+            return batch_desc
+        num_tokens = getattr(batch_desc, "num_tokens", None)
+        if num_tokens is None:
+            return batch_desc
+        return _ValueBatchDescriptor(
+            num_tokens=int(num_tokens),
+            num_reqs=getattr(batch_desc, "num_reqs", None),
+            uniform_decode=bool(getattr(batch_desc, "uniform_decode", False)),
+        )
+
+    def _descriptor_has_captured_graph(self: Any, batch_desc: Any) -> bool:
+        """Return whether vLLM should already have captured this graph key."""
+
+        num_tokens = getattr(batch_desc, "num_tokens", None)
+        if num_tokens is None:
+            return False
+        compilation_config = getattr(getattr(self, "vllm_config", None), "compilation_config", None)
+        capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None)
+        if capture_sizes:
+            return int(num_tokens) in {int(size) for size in capture_sizes}
+        max_capture = getattr(compilation_config, "max_cudagraph_capture_size", None)
+        return max_capture is not None and int(num_tokens) <= int(max_capture)
+
+    def _determine_batch_execution_and_padding(self: Any, *args: Any, **kwargs: Any) -> Any:
+        result = current(self, *args, **kwargs)
+        if not isinstance(result, tuple) or len(result) < 5:
+            return result
+        cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, cudagraph_stats = result[:5]
+        batch_desc = _normalize_batch_descriptor(batch_desc)
+        if bool(kwargs.get("force_eager", False)) or cudagraph_mode != CUDAGraphMode.NONE:
+            patched = (cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, cudagraph_stats)
+            if len(result) > 5:
+                patched += result[5:]
+            return patched
+        runtime_mode = _runtime_mode_from_config(
+            self,
+            force_uniform_decode=bool(kwargs.get("force_uniform_decode", False)),
+        )
+        if runtime_mode == CUDAGraphMode.NONE:
+            patched = (cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, cudagraph_stats)
+            if len(result) > 5:
+                patched += result[5:]
+            return patched
+        if not _descriptor_has_captured_graph(self, batch_desc):
+            patched = (CUDAGraphMode.NONE, batch_desc, should_ubatch, num_tokens_across_dp, cudagraph_stats)
+            if len(result) > 5:
+                patched += result[5:]
+            return patched
+        patched = (runtime_mode, batch_desc, should_ubatch, num_tokens_across_dp, cudagraph_stats)
+        if len(result) > 5:
+            patched += result[5:]
+        return patched
+
+    _determine_batch_execution_and_padding._easymagpie_batch_execution_padding_compat = True  # type: ignore[attr-defined]
+    _determine_batch_execution_and_padding._easymagpie_original = current  # type: ignore[attr-defined]
+    runner_cls._determine_batch_execution_and_padding = _determine_batch_execution_and_padding
+
+
+def _install_omni_structured_output_request_compat() -> None:
+    """Bridge vLLM-Omni's legacy ``StructuredOutputRequest`` construction.
+
+    The EasyMagpie vLLM-Omni branch builds structured-output metadata with
+    ``StructuredOutputRequest(sampling_params=...)`` when every request is
+    converted inside the engine core. Some vLLM 0.21 builds accept that keyword
+    while others only accept positional construction or no construction for
+    requests without guided/structured output constraints. Keep EasyMagpie audio
+    requests moving across both variants without changing vLLM-Omni source.
+    """
+
+    try:
+        request_module = importlib.import_module("vllm_omni.request")
+    except Exception:
+        return
+
+    original = getattr(request_module, "StructuredOutputRequest", None)
+    if original is None or getattr(original, "_easymagpie_structured_output_request_compat", False):
+        return
+
+    def structured_output_request_compat(*args: Any, **kwargs: Any) -> Any:
+        if "sampling_params" not in kwargs:
+            return original(*args, **kwargs)
+
+        sampling_params = kwargs["sampling_params"]
+        try:
+            return original(*args, **kwargs)
+        except TypeError as exc:
+            if "sampling_params" not in str(exc) and "unexpected keyword" not in str(exc):
+                raise
+
+        positional_kwargs = dict(kwargs)
+        positional_kwargs.pop("sampling_params", None)
+        try:
+            return original(sampling_params, *args, **positional_kwargs)
+        except TypeError:
+            return None
+
+    structured_output_request_compat._easymagpie_structured_output_request_compat = True  # type: ignore[attr-defined]
+    structured_output_request_compat._easymagpie_original = original  # type: ignore[attr-defined]
+    request_module.StructuredOutputRequest = structured_output_request_compat
+
+
+def _install_omni_request_mm_kwargs_compat() -> None:
+    """Restore the legacy ``mm_kwargs`` attribute expected by OmniRequest.
+
+    Some vLLM 0.21 builds no longer populate ``Request.mm_kwargs`` before
+    vLLM-Omni's ``OmniRequest.__init__`` aliases it into ``mm_features``.
+    EasyMagpie requests already carry multimodal payloads through Omni's prompt
+    object, so an absent field should behave like an empty mapping rather than
+    making request preprocessing fail before generation starts.
+    """
+
+    try:
+        request_module = importlib.import_module("vllm_omni.request")
+    except Exception:
+        return
+
+    request_cls = getattr(request_module, "OmniRequest", None)
+    if request_cls is None:
+        return
+
+    current = getattr(request_cls, "__init__", None)
+    if current is None or getattr(current, "_easymagpie_mm_kwargs_compat", False):
+        return
+
+    def omni_request_init_compat(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if not hasattr(self, "mm_kwargs"):
+            setattr(self, "mm_kwargs", kwargs.get("mm_features", {}))
+        return current(self, *args, **kwargs)
+
+    omni_request_init_compat._easymagpie_mm_kwargs_compat = True  # type: ignore[attr-defined]
+    omni_request_init_compat._easymagpie_original = current  # type: ignore[attr-defined]
+    request_cls.__init__ = omni_request_init_compat
+
+
+def _install_v1_serial_utils_dense_tensor_compat() -> None:
+    try:
+        import msgspec.msgpack as msgpack
+        import torch
+        import vllm.v1.serial_utils as serial_utils
+    except Exception:
+        return
+
+    encoder_cls = getattr(serial_utils, "MsgpackEncoder", None)
+    if encoder_cls is None:
+        return
+    original = getattr(encoder_cls, "_encode_tensor", None)
+    if original is None or getattr(original, "_easymagpie_dense_tensor_compat", False):
+        return
+
+    custom_type_raw_view = getattr(serial_utils, "CUSTOM_TYPE_RAW_VIEW", None)
+    if custom_type_raw_view is None:
+        return
+
+    def _encode_tensor_dense(
+        self: Any,
+        obj: torch.Tensor,
+    ) -> tuple[str, tuple[int, ...], Any]:
+        assert self.aux_buffers is not None
+        tensor = obj.detach()
+        if tensor.device.type != "cpu":
+            tensor = tensor.cpu()
+        flat = tensor.reshape(-1)
+        dense = torch.empty((int(flat.numel()),), dtype=flat.dtype, device="cpu")
+        dense.copy_(flat)
+        arr = dense.view(torch.uint8).numpy()
+        if obj.nbytes < self.size_threshold:
+            data = msgpack.Ext(custom_type_raw_view, arr.data)
+        else:
+            data = len(self.aux_buffers)
+            self.aux_buffers.append(arr.data)
+        dtype = str(obj.dtype).removeprefix("torch.")
+        return dtype, obj.shape, data
+
+    _encode_tensor_dense._easymagpie_dense_tensor_compat = True  # type: ignore[attr-defined]
+    _encode_tensor_dense._easymagpie_original = original  # type: ignore[attr-defined]
+    encoder_cls._encode_tensor = _encode_tensor_dense
+
+    decoder_cls = getattr(serial_utils, "MsgpackDecoder", None)
+    if decoder_cls is None:
+        return
+    original_decode = getattr(decoder_cls, "decode", None)
+    if original_decode is None or getattr(original_decode, "_easymagpie_easy_refit_tensor_decode_compat", False):
+        return
+
+    easy_magpie_refit_methods = {
+        "easymagpie_load_weights",
+        "easymagpie_load_non_text_weights",
+        "easymagpie_update_text_embedding_rows",
+    }
+
+    def _is_encoded_tensor_triple(obj: Any) -> bool:
+        if not isinstance(obj, (list, tuple)) or len(obj) != 3:
+            return False
+        dtype, shape, data = obj
+        if not isinstance(dtype, str):
+            return False
+        torch_dtype = getattr(torch, dtype, None)
+        if not isinstance(torch_dtype, torch.dtype):
+            return False
+        if not isinstance(shape, (list, tuple)) or not all(isinstance(dim, int) and dim >= 0 for dim in shape):
+            return False
+        if isinstance(data, int):
+            return 0 <= data < len(getattr(_decode_easy_magpie_refit_tensors, "_active_aux_buffers", ()) or ())
+        return isinstance(data, (bytes, bytearray, memoryview))
+
+    def _decode_easy_magpie_refit_tensors(self: Any, obj: Any) -> Any:
+        if _is_encoded_tensor_triple(obj):
+            return self._decode_tensor(obj)
+        if isinstance(obj, list):
+            return [_decode_easy_magpie_refit_tensors(self, item) for item in obj]
+        if isinstance(obj, tuple):
+            return tuple(_decode_easy_magpie_refit_tensors(self, item) for item in obj)
+        if isinstance(obj, dict):
+            return {key: _decode_easy_magpie_refit_tensors(self, value) for key, value in obj.items()}
+        return obj
+
+    def _is_easy_magpie_refit_utility_request(obj: Any) -> bool:
+        if not isinstance(obj, (list, tuple)) or len(obj) < 4:
+            return False
+        utility_method = obj[2]
+        utility_args = obj[3]
+        if utility_method in easy_magpie_refit_methods:
+            return True
+        if utility_method != "collective_rpc" or not isinstance(utility_args, (list, tuple)) or not utility_args:
+            return False
+        return utility_args[0] in easy_magpie_refit_methods
+
+    def decode_with_easy_magpie_refit_tensors(self: Any, bufs: Any) -> Any:
+        bytestr_types = (bytes, bytearray, memoryview)
+        zmq_module = getattr(serial_utils, "zmq", None)
+        frame_cls = getattr(zmq_module, "Frame", None)
+        if frame_cls is not None:
+            bytestr_types = (*bytestr_types, frame_cls)
+        if isinstance(bufs, bytestr_types):
+            decoded = self.decoder.decode(bufs)
+            if _is_easy_magpie_refit_utility_request(decoded):
+                decoded = _decode_easy_magpie_refit_tensors(self, decoded)
+            return decoded
+
+        self.aux_buffers = bufs
+        _decode_easy_magpie_refit_tensors._active_aux_buffers = bufs  # type: ignore[attr-defined]
+        try:
+            decoded = self.decoder.decode(bufs[0])
+            if _is_easy_magpie_refit_utility_request(decoded):
+                decoded = _decode_easy_magpie_refit_tensors(self, decoded)
+            return decoded
+        finally:
+            self.aux_buffers = ()
+            _decode_easy_magpie_refit_tensors._active_aux_buffers = ()  # type: ignore[attr-defined]
+
+    decode_with_easy_magpie_refit_tensors._easymagpie_easy_refit_tensor_decode_compat = (  # type: ignore[attr-defined]
+        True
+    )
+    decode_with_easy_magpie_refit_tensors._easymagpie_original = original_decode  # type: ignore[attr-defined]
+    decoder_cls.decode = decode_with_easy_magpie_refit_tensors
+
+def _filter_easy_magpie_refit_weights_for_model(
+    model: Any,
+    weights: list[tuple[str, Any]],
+) -> tuple[list[tuple[str, Any]], list[str]]:
+    """Drop converted optional weights that the live vLLM model does not define."""
+    try:
+        model_param_names = {str(name) for name, _ in model.named_parameters()}
+    except Exception:
+        model_param_names = set()
+    if "context_text_embedding.weight" in model_param_names:
+        return weights, []
+
+    filtered: list[tuple[str, Any]] = []
+    dropped: list[str] = []
+    for name, tensor in weights:
+        if name == "context_text_embedding.weight":
+            dropped.append(name)
+            continue
+        filtered.append((name, tensor))
+    return filtered, dropped
+
+
+def _safe_tensor_cache_key(tensor: Any) -> tuple[Any, ...] | None:
+    try:
+        return (
+            str(getattr(tensor, "device", "")),
+            int(tensor.data_ptr()),
+            tuple(int(dim) for dim in tensor.shape),
+            tuple(int(stride) for stride in tensor.stride()),
+            int(tensor.storage_offset()),
+        )
+    except Exception:
+        return None
+
+
+def _zero_cache_tensors(
+    value: Any,
+    path: str,
+    seen: set[tuple[Any, ...]],
+) -> tuple[int, int, list[dict[str, Any]], list[str]]:
+    try:
+        import torch
+    except Exception as exc:
+        return 0, 0, [], [f"{path}: import torch failed: {type(exc).__name__}: {exc}"]
+
+    if isinstance(value, torch.Tensor):
+        cache_key = _safe_tensor_cache_key(value)
+        if cache_key is not None and cache_key in seen:
+            return 0, 0, [], []
+        if cache_key is not None:
+            seen.add(cache_key)
+        try:
+            numel = int(value.numel())
+            value.zero_()
+            return (
+                1,
+                numel,
+                [
+                    {
+                        "path": path,
+                        "shape": [int(dim) for dim in value.shape],
+                        "dtype": str(value.dtype).replace("torch.", ""),
+                        "device": str(value.device),
+                        "numel": numel,
+                    }
+                ],
+                [],
+            )
+        except Exception as exc:
+            return 0, 0, [], [f"{path}: {type(exc).__name__}: {exc}"]
+
+    if isinstance(value, dict):
+        total_tensors = 0
+        total_numel = 0
+        items: list[dict[str, Any]] = []
+        errors: list[str] = []
+        for key, item in value.items():
+            sub_tensors, sub_numel, sub_items, sub_errors = _zero_cache_tensors(
+                item,
+                f"{path}.{key}",
+                seen,
+            )
+            total_tensors += sub_tensors
+            total_numel += sub_numel
+            items.extend(sub_items)
+            errors.extend(sub_errors)
+        return total_tensors, total_numel, items, errors
+
+    if isinstance(value, (list, tuple)):
+        total_tensors = 0
+        total_numel = 0
+        items = []
+        errors = []
+        for index, item in enumerate(value):
+            sub_tensors, sub_numel, sub_items, sub_errors = _zero_cache_tensors(
+                item,
+                f"{path}[{index}]",
+                seen,
+            )
+            total_tensors += sub_tensors
+            total_numel += sub_numel
+            items.extend(sub_items)
+            errors.extend(sub_errors)
+        return total_tensors, total_numel, items, errors
+
+    return 0, 0, [], []
+
+
+def _reset_easy_magpie_cache_tensors_after_refit(model_runner: Any) -> dict[str, Any]:
+    seen: set[tuple[Any, ...]] = set()
+    reset_attrs: list[str] = []
+    items: list[dict[str, Any]] = []
+    errors: list[str] = []
+    num_tensors = 0
+    numel = 0
+
+    def reset_value(value: Any, path: str) -> None:
+        nonlocal num_tensors, numel
+        sub_tensors, sub_numel, sub_items, sub_errors = _zero_cache_tensors(value, path, seen)
+        if sub_tensors:
+            reset_attrs.append(path)
+        num_tensors += sub_tensors
+        numel += sub_numel
+        items.extend(sub_items)
+        errors.extend(sub_errors)
+
+    for attr_name in ("kv_caches", "kv_cache", "gpu_cache", "cache_tensors"):
+        if hasattr(model_runner, attr_name):
+            reset_value(getattr(model_runner, attr_name), f"model_runner.{attr_name}")
+
+    model = getattr(model_runner, "model", None)
+    if model is not None:
+        model_mamba_cache = getattr(model, "mamba_cache", None)
+        if model_mamba_cache is not None:
+            reset_value(model_mamba_cache, "model.mamba_cache")
+            try:
+                setattr(model, "mamba_cache", None)
+                reset_attrs.append("model.mamba_cache=None")
+            except Exception as exc:
+                errors.append(f"model.mamba_cache=None: {type(exc).__name__}: {exc}")
+
+        modules = getattr(model, "modules", None)
+        if callable(modules):
+            try:
+                for module_index, module in enumerate(modules()):
+                    if hasattr(module, "kv_cache"):
+                        reset_value(getattr(module, "kv_cache"), f"model.modules[{module_index}].kv_cache")
+            except Exception as exc:
+                errors.append(f"model.modules: {type(exc).__name__}: {exc}")
+
+    vllm_config = getattr(model_runner, "vllm_config", None)
+    static_context = getattr(getattr(vllm_config, "compilation_config", None), "static_forward_context", None)
+    if isinstance(static_context, dict):
+        for layer_name, layer in static_context.items():
+            if hasattr(layer, "kv_cache"):
+                reset_value(getattr(layer, "kv_cache"), f"static_forward_context.{layer_name}.kv_cache")
+
+    return {
+        "num_tensors": num_tensors,
+        "numel": numel,
+        "reset_attrs": reset_attrs,
+        "items_head": items[:16],
+        "items_tail": items[-16:],
+        "errors": errors,
+    }
+
+
+def _scrub_empty_easy_magpie_input_batch(input_batch: Any) -> dict[str, Any]:
+    if input_batch is None:
+        return {"present": False}
+
+    remaining_req_ids = [
+        str(req_id)
+        for req_id in list(getattr(input_batch, "req_ids", []) or [])
+        if req_id is not None
+    ]
+    result: dict[str, Any] = {
+        "present": True,
+        "applied": False,
+        "remaining_request_ids": len(remaining_req_ids),
+        "cleared_fields": [],
+        "errors": [],
+    }
+    if remaining_req_ids:
+        result["request_ids_head"] = remaining_req_ids[:8]
+        return result
+
+    result["applied"] = True
+    for attr_name in ("_req_ids", "req_output_token_ids"):
+        value = getattr(input_batch, attr_name, None)
+        if isinstance(value, list):
+            try:
+                size_before = len(value)
+                value.clear()
+                result["cleared_fields"].append({"name": attr_name, "size_before": size_before})
+            except Exception as exc:
+                result["errors"].append(f"{attr_name}: {type(exc).__name__}: {exc}")
+
+    for attr_name in (
+        "req_id_to_index",
+        "generators",
+        "num_logprobs",
+        "num_prompt_logprobs",
+        "in_progress_prompt_logprobs_cpu",
+        "bad_words_token_ids",
+        "pooling_params",
+    ):
+        value = getattr(input_batch, attr_name, None)
+        if isinstance(value, dict):
+            try:
+                size_before = len(value)
+                value.clear()
+                result["cleared_fields"].append({"name": attr_name, "size_before": size_before})
+            except Exception as exc:
+                result["errors"].append(f"{attr_name}: {type(exc).__name__}: {exc}")
+
+    for attr_name in (
+        "greedy_reqs",
+        "random_reqs",
+        "top_p_reqs",
+        "top_k_reqs",
+        "spec_decode_unsupported_reqs",
+        "frequency_penalties_reqs",
+        "presence_penalties_reqs",
+        "repetition_penalties_reqs",
+        "has_allowed_token_ids",
+    ):
+        value = getattr(input_batch, attr_name, None)
+        if isinstance(value, set):
+            try:
+                size_before = len(value)
+                value.clear()
+                result["cleared_fields"].append({"name": attr_name, "size_before": size_before})
+            except Exception as exc:
+                result["errors"].append(f"{attr_name}: {type(exc).__name__}: {exc}")
+
+    for attr_name in (
+        "token_ids_cpu_tensor",
+        "num_computed_tokens_cpu_tensor",
+        "allowed_token_ids_mask_cpu_tensor",
+    ):
+        value = getattr(input_batch, attr_name, None)
+        if hasattr(value, "zero_"):
+            try:
+                value.zero_()
+                result["cleared_fields"].append({"name": attr_name, "zeroed": True})
+            except Exception as exc:
+                result["errors"].append(f"{attr_name}.zero_: {type(exc).__name__}: {exc}")
+
+    for attr_name in (
+        "token_ids_cpu",
+        "num_tokens",
+        "num_tokens_no_spec",
+        "num_prompt_tokens",
+        "num_computed_tokens_cpu",
+        "temperature_cpu",
+        "top_p_cpu",
+        "top_k_cpu",
+        "frequency_penalties_cpu",
+        "presence_penalties_cpu",
+        "repetition_penalties_cpu",
+    ):
+        value = getattr(input_batch, attr_name, None)
+        if hasattr(value, "fill"):
+            try:
+                value.fill(0)
+                result["cleared_fields"].append({"name": attr_name, "filled": 0})
+            except Exception as exc:
+                result["errors"].append(f"{attr_name}.fill: {type(exc).__name__}: {exc}")
+
+    block_table = getattr(input_batch, "block_table", None)
+    clear = getattr(block_table, "clear", None)
+    if callable(clear):
+        try:
+            clear()
+            result["cleared_fields"].append({"name": "block_table", "cleared": True})
+        except Exception as exc:
+            result["errors"].append(f"block_table.clear: {type(exc).__name__}: {exc}")
+
+    batch_update_builder = getattr(input_batch, "batch_update_builder", None)
+    get_and_reset = getattr(batch_update_builder, "get_and_reset", None)
+    if callable(get_and_reset):
+        try:
+            get_and_reset(0)
+            result["cleared_fields"].append({"name": "batch_update_builder", "reset": True})
+        except Exception as exc:
+            result["errors"].append(f"batch_update_builder.get_and_reset: {type(exc).__name__}: {exc}")
+
+    return result
+
+
+def _reset_easy_magpie_runner_state_after_refit(model_runner: Any) -> dict[str, Any]:
+    """Clear runner-side request/cache state after live refit.
+
+    EasyMagpie refit is a between-rollout operation. Finished or aborted
+    request ids left in the model runner can still influence the next Mamba
+    cache update, so remove them before admitting a post-refit request.
+    """
+
+    if model_runner is None:
+        return {
+            "cleared_attrs": [],
+            "missing_attrs": [],
+            "cleared_mappings": [],
+            "input_batch_reset": {"present": False},
+            "input_batch_scrub": {"present": False},
+            "cache_tensors_reset": {"num_tensors": 0, "numel": 0, "reset_attrs": [], "errors": []},
+            "errors": ["model_runner is None"],
+        }
+
+    cleared_attrs: list[str] = []
+    missing_attrs: list[str] = []
+    cleared_mappings: list[dict[str, Any]] = []
+    errors: list[str] = []
+    for name in (
+        "_easymagpie_active_mamba_request_ids",
+        "_easymagpie_mamba_request_cache_kwargs",
+        "_easymagpie_mamba_request_cache_state",
+        "_easymagpie_mamba_cache_batch_size",
+    ):
+        if not hasattr(model_runner, name):
+            missing_attrs.append(name)
+            continue
+        try:
+            delattr(model_runner, name)
+            cleared_attrs.append(name)
+        except Exception as exc:
+            errors.append(f"{name}: {type(exc).__name__}: {exc}")
+
+    for name in ("requests", "encoder_cache"):
+        mapping = getattr(model_runner, name, None)
+        if not isinstance(mapping, dict):
+            continue
+        size_before = len(mapping)
+        try:
+            mapping.clear()
+            cleared_mappings.append({"name": name, "size_before": size_before})
+        except Exception as exc:
+            errors.append(f"{name}.clear: {type(exc).__name__}: {exc}")
+
+    input_batch_reset: dict[str, Any] = {"present": False}
+    input_batch = getattr(model_runner, "input_batch", None)
+    if input_batch is not None:
+        req_ids = [
+            str(req_id)
+            for req_id in list(getattr(input_batch, "req_ids", []) or [])
+            if req_id is not None
+        ]
+        removed: list[str] = []
+        missing: list[str] = []
+        input_batch_errors: list[str] = []
+        remove_request = getattr(input_batch, "remove_request", None)
+        if callable(remove_request):
+            for req_id in req_ids:
+                try:
+                    removed_index = remove_request(req_id)
+                except Exception as exc:
+                    input_batch_errors.append(f"{req_id}: {type(exc).__name__}: {exc}")
+                    continue
+                if removed_index is None:
+                    missing.append(req_id)
+                else:
+                    removed.append(req_id)
+            condense = getattr(input_batch, "condense", None)
+            if callable(condense):
+                try:
+                    condense()
+                except Exception as exc:
+                    input_batch_errors.append(f"condense: {type(exc).__name__}: {exc}")
+        else:
+            raw_req_ids = getattr(input_batch, "_req_ids", None)
+            req_id_to_index = getattr(input_batch, "req_id_to_index", None)
+            if isinstance(raw_req_ids, list) and isinstance(req_id_to_index, dict):
+                raw_req_ids.clear()
+                req_id_to_index.clear()
+                removed = req_ids
+            else:
+                input_batch_errors.append("input_batch has no remove_request method")
+        errors.extend(f"input_batch.{error}" for error in input_batch_errors)
+        input_batch_reset = {
+            "present": True,
+            "num_request_ids_before": len(req_ids),
+            "num_removed": len(removed),
+            "num_missing": len(missing),
+            "request_ids_before_head": req_ids[:8],
+            "request_ids_before_tail": req_ids[-8:],
+            "errors": input_batch_errors,
+        }
+
+    input_batch_scrub = _scrub_empty_easy_magpie_input_batch(input_batch)
+    cache_tensors_reset = _reset_easy_magpie_cache_tensors_after_refit(model_runner)
+    errors.extend(f"input_batch_scrub.{error}" for error in input_batch_scrub.get("errors", []))
+    errors.extend(f"cache_tensors_reset.{error}" for error in cache_tensors_reset.get("errors", []))
+
+    return {
+        "cleared_attrs": cleared_attrs,
+        "missing_attrs": missing_attrs,
+        "cleared_mappings": cleared_mappings,
+        "input_batch_reset": input_batch_reset,
+        "input_batch_scrub": input_batch_scrub,
+        "cache_tensors_reset": cache_tensors_reset,
+        "errors": errors,
+    }
+
+
+_EASYMAGPIE_REFIT_RPC_COMPAT_VERSION = 5
+_EASYMAGPIE_TEXT_ROW_REFIT_RPC_COMPAT_VERSION = 1
+_EASYMAGPIE_RUNTIME_RESET_RPC_COMPAT_VERSION = 1
+
+
+def _install_easy_magpie_refit_rpc_compat() -> None:
+    """Expose a tiny worker RPC for loading EasyMagpie refit weights."""
+
+    def _looks_like_refit_pair(item: Any) -> bool:
+        return isinstance(item, (list, tuple)) and len(item) == 2 and isinstance(item[0], str)
+
+    def _materialize_refit_weights(weights: Any) -> list[tuple[str, Any]]:
+        if isinstance(weights, (str, bytes, os.PathLike)):
+            try:
+                import torch
+
+                loaded = torch.load(os.fspath(weights), map_location="cpu", weights_only=True)
+            except TypeError:
+                import torch
+
+                loaded = torch.load(os.fspath(weights), map_location="cpu")
+            if isinstance(loaded, dict) and "weights" in loaded:
+                loaded = loaded["weights"]
+            return _materialize_refit_weights(loaded)
+
+        if isinstance(weights, dict) and "path" in weights:
+            return _materialize_refit_weights(weights["path"])
+
+        if isinstance(weights, dict):
+            raw_items = list(weights.items())
+        else:
+            raw_items = list(weights)
+
+        # Some collective_rpc paths preserve the outer argument tuple, so the
+        # worker receives ``([("name", tensor), ...],)`` instead of the list.
+        if (
+            len(raw_items) == 1
+            and isinstance(raw_items[0], (list, tuple))
+            and not _looks_like_refit_pair(raw_items[0])
+            and all(_looks_like_refit_pair(item) for item in raw_items[0])
+        ):
+            raw_items = list(raw_items[0])
+
+        materialized: list[tuple[str, Any]] = []
+        for item in raw_items:
+            if not _looks_like_refit_pair(item):
+                raise TypeError(f"invalid EasyMagpie refit weight entry: {type(item).__name__}")
+            name, tensor = item
+            # A few RPC transports normalize tuple payloads to nested lists,
+            # e.g. ["decoder.foo", [tensor]]. Keep the model loader strict but
+            # undo that harmless wrapper here.
+            if isinstance(tensor, (list, tuple)) and len(tensor) == 1 and hasattr(tensor[0], "shape"):
+                tensor = tensor[0]
+            if not hasattr(tensor, "shape"):
+                try:
+                    import torch
+
+                    tensor = torch.as_tensor(tensor)
+                except Exception as exc:
+                    raise TypeError(f"invalid tensor for EasyMagpie refit weight {name!r}") from exc
+            materialized.append((name, tensor))
+        return materialized
+
+    def _materialize_text_row_payload(payload: Any) -> tuple[Any, list[tuple[str, Any]]]:
+        if isinstance(payload, (str, bytes, os.PathLike)):
+            try:
+                import torch
+
+                loaded = torch.load(os.fspath(payload), map_location="cpu", weights_only=True)
+            except TypeError:
+                import torch
+
+                loaded = torch.load(os.fspath(payload), map_location="cpu")
+            return _materialize_text_row_payload(loaded)
+        if isinstance(payload, dict) and "path" in payload:
+            return _materialize_text_row_payload(payload["path"])
+        if not isinstance(payload, dict):
+            raise TypeError(f"invalid EasyMagpie text-row refit payload: {type(payload).__name__}")
+        row_ids = payload.get("row_ids")
+        rows = payload.get("weights", payload.get("rows"))
+        if row_ids is None or rows is None:
+            raise ValueError("EasyMagpie text-row refit payload needs row_ids and weights")
+        return row_ids, _materialize_refit_weights(rows)
+
+    def _sample_text_row_copy_check(
+        *,
+        name: str,
+        row_ids: Any,
+        source_rows: Any,
+        target_weight: Any,
+    ) -> dict[str, Any]:
+        import torch
+
+        ids = torch.as_tensor(row_ids, dtype=torch.long, device=target_weight.device).reshape(-1)
+        src = source_rows.to(device=target_weight.device, dtype=target_weight.dtype)
+        sampled = target_weight.index_select(0, ids)
+        if src.numel() == 0:
+            max_diff = 0.0
+            matched = True
+        else:
+            diff = (sampled.float() - src.float()).abs()
+            max_diff = float(diff.max().item()) if diff.numel() else 0.0
+            matched = bool(torch.allclose(sampled.float(), src.float(), atol=5e-3, rtol=5e-3))
+        return {
+            "name": name,
+            "checked": True,
+            "matched": matched,
+            "num_rows": int(ids.numel()),
+            "row_ids_head": [int(x) for x in ids[:8].detach().cpu().tolist()],
+            "row_ids_tail": [int(x) for x in ids[-8:].detach().cpu().tolist()],
+            "source_shape": list(source_rows.shape),
+            "target_shape": list(target_weight.shape),
+            "max_sample_abs_diff": max_diff,
+        }
+
+    def _apply_text_row_update(model: Any, row_ids: Any, rows: list[tuple[str, Any]]) -> dict[str, Any]:
+        import torch
+
+        ids_cpu = torch.as_tensor(row_ids, dtype=torch.long, device="cpu").reshape(-1)
+        if ids_cpu.numel() <= 0:
+            raise ValueError("EasyMagpie text-row refit needs at least one row")
+        if len(torch.unique(ids_cpu)) != ids_cpu.numel():
+            raise ValueError("EasyMagpie text-row refit row_ids must be unique")
+        own_params = dict(model.named_parameters())
+        updated: list[str] = []
+        dropped: list[str] = []
+        copy_checks: list[dict[str, Any]] = []
+        for name, tensor in rows:
+            name = str(name)
+            if name not in {"text_embedding.weight", "context_text_embedding.weight"}:
+                raise ValueError(f"unsupported EasyMagpie text-row refit weight: {name!r}")
+            target = own_params.get(name)
+            if target is None:
+                if name == "context_text_embedding.weight" and getattr(model, "context_text_embedding", None) is getattr(model, "text_embedding", None):
+                    dropped.append(name)
+                    continue
+                raise KeyError(f"EasyMagpie text-row refit target not found: {name}")
+            values = tensor if hasattr(tensor, "shape") else torch.as_tensor(tensor)
+            if values.ndim != 2:
+                raise ValueError(f"{name} row tensor must be 2-D, got shape {tuple(values.shape)}")
+            if int(values.shape[0]) != int(ids_cpu.numel()) or int(values.shape[1]) != int(target.shape[1]):
+                raise ValueError(
+                    f"{name} row tensor shape {tuple(values.shape)} does not match "
+                    f"row_ids={int(ids_cpu.numel())}, hidden={int(target.shape[1])}"
+                )
+            min_id = int(ids_cpu.min().item())
+            max_id = int(ids_cpu.max().item())
+            if min_id < 0 or max_id >= int(target.shape[0]):
+                raise ValueError(
+                    f"{name} row ids must be in [0, {int(target.shape[0]) - 1}], got [{min_id}, {max_id}]"
+                )
+            ids = ids_cpu.to(device=target.device)
+            with torch.no_grad():
+                target.data.index_copy_(0, ids, values.to(device=target.device, dtype=target.dtype))
+            updated.append(name)
+            copy_checks.append(
+                _sample_text_row_copy_check(
+                    name=name,
+                    row_ids=ids,
+                    source_rows=values,
+                    target_weight=target,
+                )
+            )
+        failed = [item for item in copy_checks if not bool(item.get("matched"))]
+        return {
+            "ok": not failed and bool(updated),
+            "num_rows": int(ids_cpu.numel()),
+            "row_ids_head": [int(x) for x in ids_cpu[:8].tolist()],
+            "row_ids_tail": [int(x) for x in ids_cpu[-8:].tolist()],
+            "updated": updated,
+            "dropped_unsupported": dropped,
+            "copy_check": {
+                "ok": not failed,
+                "num_checked": len(copy_checks),
+                "num_failed": len(failed),
+                "failed_head": failed[:8],
+                "items": copy_checks[:8],
+            },
+        }
+
+    _MISSING_ATTR = object()
+
+    def _iter_easy_magpie_refit_flag_targets(model: Any):
+        seen: set[int] = set()
+        stack = [model]
+        while stack:
+            target = stack.pop()
+            if target is None:
+                continue
+            target_id = id(target)
+            if target_id in seen:
+                continue
+            seen.add(target_id)
+            if target is model or callable(getattr(target, "load_weights", None)) or hasattr(target, "text_embedding"):
+                yield target
+
+            modules = getattr(target, "modules", None)
+            if callable(modules):
+                try:
+                    stack.extend(child for child in modules() if child is not target)
+                except Exception:
+                    pass
+            for attr_name in ("model", "module", "_module"):
+                try:
+                    child = getattr(target, attr_name, None)
+                except Exception:
+                    child = None
+                if child is not None and child is not target:
+                    stack.append(child)
+
+    def _load_easy_magpie_weights_for_refit(
+        model: Any,
+        *,
+        weights: list[tuple[str, Any]],
+        allow_missing_text_tables: bool,
+    ):
+        flag_targets = list(_iter_easy_magpie_refit_flag_targets(model))
+        previous_refit_active: list[tuple[Any, Any]] = []
+        previous_allow_missing: list[tuple[Any, Any]] = []
+        for target in flag_targets:
+            previous_value = getattr(target, "_easymagpie_refit_rpc_active", _MISSING_ATTR)
+            try:
+                target._easymagpie_refit_rpc_active = True
+            except Exception:
+                continue
+            previous_refit_active.append((target, previous_value))
+        if allow_missing_text_tables:
+            for target in flag_targets:
+                previous_value = getattr(target, "_easymagpie_allow_missing_text_tables_refit", _MISSING_ATTR)
+                try:
+                    target._easymagpie_allow_missing_text_tables_refit = True
+                except Exception:
+                    continue
+                previous_allow_missing.append((target, previous_value))
+        try:
+            return model.load_weights(weights=weights)
+        finally:
+            for target, previous_value in reversed(previous_allow_missing):
+                if previous_value is _MISSING_ATTR:
+                    try:
+                        delattr(target, "_easymagpie_allow_missing_text_tables_refit")
+                    except Exception:
+                        pass
+                else:
+                    try:
+                        target._easymagpie_allow_missing_text_tables_refit = previous_value
+                    except Exception:
+                        pass
+            for target, previous_value in reversed(previous_refit_active):
+                if previous_value is _MISSING_ATTR:
+                    try:
+                        delattr(target, "_easymagpie_refit_rpc_active")
+                    except Exception:
+                        pass
+                else:
+                    try:
+                        target._easymagpie_refit_rpc_active = previous_value
+                    except Exception:
+                        pass
+            try:
+                flag_targets.clear()
+            except Exception:
+                pass
+
+    worker_classes = []
+    for module_name, class_name in (
+        ("vllm_omni.worker.gpu_ar_worker", "GPUARWorker"),
+        ("vllm_omni.worker.gpu_generation_worker", "GPUGenerationWorker"),
+    ):
+        try:
+            module = importlib.import_module(module_name)
+        except Exception:
+            continue
+        worker_cls = getattr(module, class_name, None)
+        if worker_cls is not None:
+            worker_classes.append(worker_cls)
+
+    for worker_cls in worker_classes:
+        installed_version = int(getattr(worker_cls, "_easymagpie_refit_rpc_compat_version", 0) or 0)
+        if (
+            getattr(worker_cls, "_easymagpie_refit_rpc_compat", False)
+            and installed_version >= _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION
+        ):
+            continue
+
+        def _easymagpie_load_weights_impl(self, weights, *, allow_missing_text_tables: bool):
+            model_runner = getattr(self, "model_runner", None)
+            model = getattr(model_runner, "model", None)
+            if model is None or not hasattr(model, "load_weights"):
+                return {
+                    "ok": False,
+                    "refit_rpc_compat_version": _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION,
+                    "error": "EasyMagpie refit RPC could not find model_runner.model.load_weights",
+                }
+            materialized = []
+            try:
+                materialized = _materialize_refit_weights(weights)
+                materialized, dropped = _filter_easy_magpie_refit_weights_for_model(model, materialized)
+                loaded = _load_easy_magpie_weights_for_refit(
+                    model,
+                    weights=materialized,
+                    allow_missing_text_tables=allow_missing_text_tables,
+                )
+                runtime_state_reset = _reset_easy_magpie_runner_state_after_refit(model_runner)
+                try:
+                    import torch
+
+                    if torch.cuda.is_available():
+                        torch.cuda.synchronize()
+                except Exception:
+                    pass
+                input_names = [str(name) for name, _ in materialized]
+                loaded_names = sorted(str(item) for item in loaded) if loaded is not None else []
+                num_loaded = len(loaded_names) if loaded is not None else None
+                load_summary = getattr(model, "_last_easy_magpie_load_weights_summary", None)
+                if isinstance(load_summary, dict):
+                    load_ok = bool(load_summary.get("ok", True))
+                    reset_ok = not runtime_state_reset.get("errors")
+                    result = {
+                        "ok": bool(load_ok and reset_ok),
+                        "num_input_tensors": int(load_summary.get("num_input_tensors", len(materialized))),
+                        "num_loaded": int(load_summary.get("num_loaded_targets", num_loaded or 0)),
+                        "num_loaded_targets": int(load_summary.get("num_loaded_targets", num_loaded or 0)),
+                        "input_head": load_summary.get("input_head", input_names[:16]),
+                        "input_tail": load_summary.get("input_tail", input_names[-16:]),
+                        "loaded_head": load_summary.get("loaded_head", loaded_names[:16]),
+                        "loaded_tail": load_summary.get("loaded_tail", loaded_names[-16:]),
+                        "dropped_unsupported": dropped,
+                        "allow_missing_text_tables": bool(allow_missing_text_tables),
+                        "load_summary": load_summary,
+                        "runtime_state_reset": runtime_state_reset,
+                        "refit_rpc_compat_version": _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION,
+                    }
+                    if not load_ok:
+                        result["error"] = "EasyMagpie refit load summary reported incomplete update"
+                    elif not reset_ok:
+                        result["error"] = "EasyMagpie refit loaded weights but failed to reset runner runtime state"
+                    return result
+                if num_loaded is not None and num_loaded != len(materialized):
+                    return {
+                        "ok": False,
+                        "error": (
+                            "EasyMagpie refit loaded tensor count mismatch: "
+                            f"loaded {num_loaded} of {len(materialized)} input tensors"
+                        ),
+                        "num_input_tensors": len(materialized),
+                        "num_loaded": num_loaded,
+                        "input_head": input_names[:16],
+                        "input_tail": input_names[-16:],
+                        "loaded_head": loaded_names[:16],
+                        "loaded_tail": loaded_names[-16:],
+                        "dropped_unsupported": dropped,
+                        "allow_missing_text_tables": bool(allow_missing_text_tables),
+                        "runtime_state_reset": runtime_state_reset,
+                        "refit_rpc_compat_version": _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION,
+                    }
+                reset_ok = not runtime_state_reset.get("errors")
+                return {
+                    "ok": bool(reset_ok),
+                    "num_input_tensors": len(materialized),
+                    "num_loaded": num_loaded,
+                    "input_head": input_names[:16],
+                    "input_tail": input_names[-16:],
+                    "loaded_head": loaded_names[:16],
+                    "loaded_tail": loaded_names[-16:],
+                    "dropped_unsupported": dropped,
+                    "allow_missing_text_tables": bool(allow_missing_text_tables),
+                    "runtime_state_reset": runtime_state_reset,
+                    "refit_rpc_compat_version": _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION,
+                    **(
+                        {"error": "EasyMagpie refit loaded weights but failed to reset runner runtime state"}
+                        if not reset_ok
+                        else {}
+                    ),
+                }
+            except Exception as exc:
+                logger.exception("EasyMagpie refit RPC failed")
+                try:
+                    num_input_tensors = len(weights)
+                except Exception:
+                    num_input_tensors = None
+                return {
+                    "ok": False,
+                    "refit_rpc_compat_version": _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION,
+                    "error": f"{type(exc).__name__}: {exc}",
+                    "num_input_tensors": num_input_tensors,
+                }
+
+        def easymagpie_load_weights(self, weights):
+            return _easymagpie_load_weights_impl(
+                self,
+                weights,
+                allow_missing_text_tables=False,
+            )
+
+        def easymagpie_load_non_text_weights(self, weights):
+            result = _easymagpie_load_weights_impl(
+                self,
+                weights,
+                allow_missing_text_tables=True,
+            )
+            if isinstance(result, dict):
+                result["refit_mode"] = "non_text_weights"
+            return result
+
+        def easymagpie_update_text_embedding_rows(self, payload):
+            model_runner = getattr(self, "model_runner", None)
+            model = getattr(model_runner, "model", None)
+            if model is None or not hasattr(model, "named_parameters"):
+                return {
+                    "ok": False,
+                    "text_row_refit_rpc_compat_version": _EASYMAGPIE_TEXT_ROW_REFIT_RPC_COMPAT_VERSION,
+                    "error": "EasyMagpie text-row refit RPC could not find model_runner.model.named_parameters",
+                }
+            try:
+                row_ids, rows = _materialize_text_row_payload(payload)
+                update_summary = _apply_text_row_update(model, row_ids, rows)
+                runtime_state_reset = _reset_easy_magpie_runner_state_after_refit(model_runner)
+                try:
+                    import torch
+
+                    if torch.cuda.is_available():
+                        torch.cuda.synchronize()
+                except Exception:
+                    pass
+                reset_ok = not runtime_state_reset.get("errors")
+                result = {
+                    **update_summary,
+                    "runtime_state_reset": runtime_state_reset,
+                    "text_row_refit_rpc_compat_version": _EASYMAGPIE_TEXT_ROW_REFIT_RPC_COMPAT_VERSION,
+                }
+                result["ok"] = bool(update_summary.get("ok")) and reset_ok
+                if not result["ok"] and not reset_ok:
+                    result["error"] = "EasyMagpie text-row refit loaded rows but failed to reset runner runtime state"
+                return result
+            except Exception as exc:
+                logger.exception("EasyMagpie text-row refit RPC failed")
+                return {
+                    "ok": False,
+                    "text_row_refit_rpc_compat_version": _EASYMAGPIE_TEXT_ROW_REFIT_RPC_COMPAT_VERSION,
+                    "error": f"{type(exc).__name__}: {exc}",
+                }
+
+        def easymagpie_reset_runtime_state(self):
+            model_runner = getattr(self, "model_runner", None)
+            if model_runner is None:
+                return {
+                    "ok": False,
+                    "runtime_reset_rpc_compat_version": _EASYMAGPIE_RUNTIME_RESET_RPC_COMPAT_VERSION,
+                    "error": "EasyMagpie runtime reset RPC could not find worker.model_runner",
+                }
+            try:
+                runtime_state_reset = _reset_easy_magpie_runner_state_after_refit(model_runner)
+                try:
+                    import torch
+
+                    if torch.cuda.is_available():
+                        torch.cuda.synchronize()
+                except Exception:
+                    pass
+                reset_ok = not runtime_state_reset.get("errors")
+                return {
+                    "ok": bool(reset_ok),
+                    "runtime_state_reset": runtime_state_reset,
+                    "runtime_reset_rpc_compat_version": _EASYMAGPIE_RUNTIME_RESET_RPC_COMPAT_VERSION,
+                    **(
+                        {"error": "EasyMagpie runtime-state reset reported errors"}
+                        if not reset_ok
+                        else {}
+                    ),
+                }
+            except Exception as exc:
+                logger.exception("EasyMagpie runtime-state reset RPC failed")
+                return {
+                    "ok": False,
+                    "runtime_reset_rpc_compat_version": _EASYMAGPIE_RUNTIME_RESET_RPC_COMPAT_VERSION,
+                    "error": f"{type(exc).__name__}: {exc}",
+                }
+
+        easymagpie_load_weights._easymagpie_refit_rpc_compat = True  # type: ignore[attr-defined]
+        easymagpie_load_weights._easymagpie_refit_rpc_compat_version = _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION  # type: ignore[attr-defined]
+        easymagpie_load_non_text_weights._easymagpie_refit_rpc_compat = True  # type: ignore[attr-defined]
+        easymagpie_load_non_text_weights._easymagpie_refit_rpc_compat_version = _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION  # type: ignore[attr-defined]
+        easymagpie_update_text_embedding_rows._easymagpie_text_row_refit_rpc_compat = True  # type: ignore[attr-defined]
+        easymagpie_update_text_embedding_rows._easymagpie_text_row_refit_rpc_compat_version = _EASYMAGPIE_TEXT_ROW_REFIT_RPC_COMPAT_VERSION  # type: ignore[attr-defined]
+        easymagpie_reset_runtime_state._easymagpie_runtime_reset_rpc_compat = True  # type: ignore[attr-defined]
+        easymagpie_reset_runtime_state._easymagpie_runtime_reset_rpc_compat_version = _EASYMAGPIE_RUNTIME_RESET_RPC_COMPAT_VERSION  # type: ignore[attr-defined]
+        worker_cls.easymagpie_load_weights = easymagpie_load_weights
+        worker_cls.easymagpie_load_non_text_weights = easymagpie_load_non_text_weights
+        worker_cls.easymagpie_update_text_embedding_rows = easymagpie_update_text_embedding_rows
+        worker_cls.easymagpie_reset_runtime_state = easymagpie_reset_runtime_state
+        worker_cls._easymagpie_refit_rpc_compat = True
+        worker_cls._easymagpie_refit_rpc_compat_version = _EASYMAGPIE_REFIT_RPC_COMPAT_VERSION
+        worker_cls._easymagpie_text_row_refit_rpc_compat = True
+        worker_cls._easymagpie_text_row_refit_rpc_compat_version = _EASYMAGPIE_TEXT_ROW_REFIT_RPC_COMPAT_VERSION
+        worker_cls._easymagpie_runtime_reset_rpc_compat = True
+        worker_cls._easymagpie_runtime_reset_rpc_compat_version = _EASYMAGPIE_RUNTIME_RESET_RPC_COMPAT_VERSION
+
+
+def install_easy_magpie_refit_rpc_compat() -> None:
+    """Install only the EasyMagpie refit worker RPC, leaving generation untouched."""
+
+    _install_ray_objectref_aliases()
+    _install_ray_placement_group_aliases()
+    _install_vllm_inputs_data_alias()
+    _install_vllm_multimodal_inputs_alias()
+    _install_engine_utils_compat()
+    _install_omni_gpu_model_runner_init_kwargs_compat()
+    _install_omni_cumsum_arange_compat()
+    _install_omni_batch_execution_padding_compat()
+    _install_omni_structured_output_request_compat()
+    _install_omni_request_mm_kwargs_compat()
+    _install_easy_magpie_refit_rpc_compat()
+    _install_async_omni_client_compat()
+
+
+def install_easy_magpie_runtime_compat() -> None:
+    """Install focused shims needed by EasyMagpie RL rollout workers.
+
+    This intentionally avoids the broad vLLM-Omni import/signature shims in
+    ``install_vllm_omni_compat``. The RL path needs the refit RPC plus a dense
+    tensor serializer for vLLM V1 outputs, but the broader compatibility layer
+    can perturb newer vLLM/Pydantic stacks.
+    """
+
+    _install_ray_objectref_aliases()
+    _install_ray_placement_group_aliases()
+    _install_vllm_inputs_data_alias()
+    _install_vllm_multimodal_inputs_alias()
+    _install_engine_utils_compat()
+    _install_omni_gpu_model_runner_init_kwargs_compat()
+    _install_omni_cumsum_arange_compat()
+    _install_omni_batch_execution_padding_compat()
+    _install_omni_structured_output_request_compat()
+    _install_omni_request_mm_kwargs_compat()
+    _install_easy_magpie_refit_rpc_compat()
+    _install_v1_serial_utils_dense_tensor_compat()
+    _install_async_omni_client_compat()
+
+def _install_async_omni_client_compat() -> None:
+    try:
+        async_omni_module = importlib.import_module("vllm_omni.entrypoints.async_omni")
+    except Exception:
+        return
+
+    AsyncOmni = getattr(async_omni_module, "AsyncOmni", None)
+    if AsyncOmni is None:
+        return
+
+    if not hasattr(AsyncOmni, "get_model_config") or getattr(
+        AsyncOmni.get_model_config, "__isabstractmethod__", False
+    ):
+
+        async def get_model_config(self):
+            return self.model_config
+
+        AsyncOmni.get_model_config = get_model_config
+
+    if not hasattr(AsyncOmni, "get_decoding_config") or getattr(
+        AsyncOmni.get_decoding_config, "__isabstractmethod__", False
+    ):
+
+        async def get_decoding_config(self):
+            vllm_config = self.vllm_config
+            decoding_config = getattr(vllm_config, "decoding_config", None)
+            if decoding_config is not None:
+                return decoding_config
+            model_config = self.model_config
+            return getattr(model_config, "decoding_config", None)
+
+        AsyncOmni.get_decoding_config = get_decoding_config
+
+    if not hasattr(AsyncOmni, "notify_kv_transfer_request_rejected") or getattr(
+        AsyncOmni.notify_kv_transfer_request_rejected, "__isabstractmethod__", False
+    ):
+
+        def notify_kv_transfer_request_rejected(self, *args, **kwargs):
+            return None
+
+        AsyncOmni.notify_kv_transfer_request_rejected = notify_kv_transfer_request_rejected
+
+    abstract_methods = getattr(AsyncOmni, "__abstractmethods__", frozenset())
+    if abstract_methods:
+        AsyncOmni.__abstractmethods__ = frozenset(
+            name
+            for name in abstract_methods
+            if name
+            not in {
+                "get_decoding_config",
+                "get_model_config",
+                "notify_kv_transfer_request_rejected",
+            }
+        )
+
+
+def _install_ray_objectref_aliases() -> None:
+    """Restore Ray type aliases expected by NeMo-RL annotation imports."""
+
+    try:
+        import ray
+    except Exception:
+        return
+    try:
+        import ray._raylet as raylet
+    except Exception:
+        raylet = None
+
+    class _EasyMagpieRayObjectRef:
+        pass
+
+    class _EasyMagpieRayObjectRefGenerator:
+        pass
+
+    for name, fallback in (
+        ("ObjectRef", _EasyMagpieRayObjectRef),
+        ("ObjectRefGenerator", _EasyMagpieRayObjectRefGenerator),
+    ):
+        if hasattr(ray, name):
+            continue
+        alias = getattr(raylet, name, None) if raylet is not None else None
+        if alias is None:
+            alias = fallback
+        try:
+            setattr(ray, name, alias)
+        except Exception:
+            pass
+
+
+def _install_ray_placement_group_aliases() -> None:
+    """Restore Ray placement-group imports used by NeMo-RL type imports."""
+
+    try:
+        import ray
+    except Exception:
+        return
+
+    if not hasattr(ray, "is_initialized"):
+        ray.is_initialized = lambda: False
+
+    if not hasattr(ray, "remote"):
+
+        class _UnavailableRemote:
+            def __init__(self, target, options=None):
+                self._target = target
+                self._options = dict(options or {})
+                self.__name__ = getattr(target, "__name__", type(target).__name__)
+                self.__qualname__ = getattr(target, "__qualname__", self.__name__)
+                self.__doc__ = getattr(target, "__doc__", None)
+
+            def options(self, **kwargs):
+                merged = dict(self._options)
+                merged.update(kwargs)
+                return type(self)(self._target, merged)
+
+            def remote(self, *_args, **_kwargs):
+                raise RuntimeError(
+                    "ray.remote is unavailable in this Ray build. The EasyMagpie "
+                    "compatibility shim only restores this symbol so NeMo-RL modules "
+                    "can be imported in non-Ray train loops."
+                )
+
+        def remote(*args, **kwargs):
+            if len(args) == 1 and callable(args[0]) and not kwargs:
+                return _UnavailableRemote(args[0])
+
+            def decorate(target):
+                return _UnavailableRemote(target, kwargs)
+
+            return decorate
+
+        ray.remote = remote
+
+    actor_module_name = "ray.actor"
+    try:
+        actor_module = importlib.import_module(actor_module_name)
+    except Exception:
+        actor_module = types.ModuleType(actor_module_name)
+        sys.modules[actor_module_name] = actor_module
+
+    if not hasattr(actor_module, "ActorHandle"):
+
+        class ActorHandle:
+            pass
+
+        actor_module.ActorHandle = ActorHandle
+
+    if not hasattr(ray, "actor"):
+        try:
+            setattr(ray, "actor", actor_module)
+        except Exception:
+            pass
+
+    try:
+        import ray.util as ray_util
+    except Exception:
+        ray_util = types.ModuleType("ray.util")
+        sys.modules["ray.util"] = ray_util
+        try:
+            setattr(ray, "util", ray_util)
+        except Exception:
+            pass
+
+    if not hasattr(ray_util, "__path__"):
+        try:
+            ray_util.__path__ = []  # type: ignore[attr-defined]
+        except Exception:
+            pass
+
+    module_name = "ray.util.placement_group"
+    try:
+        module = importlib.import_module(module_name)
+    except Exception:
+        module = types.ModuleType(module_name)
+        sys.modules[module_name] = module
+
+    if not hasattr(module, "PlacementGroup"):
+
+        class PlacementGroup:
+            pass
+
+        module.PlacementGroup = PlacementGroup
+
+    if not hasattr(module, "placement_group"):
+
+        def placement_group(*_args, **_kwargs):
+            raise RuntimeError(
+                "ray.util.placement_group.placement_group is unavailable in this Ray build. "
+                "The EasyMagpie compatibility shim only restores this symbol so NeMo-RL "
+                "modules can be imported in non-Ray train loops."
+            )
+
+        module.placement_group = placement_group
+
+    if not hasattr(ray_util, "placement_group"):
+        try:
+            setattr(ray_util, "placement_group", module)
+        except Exception:
+            pass
+
+    if not hasattr(module, "placement_group_table"):
+
+        def placement_group_table(*_args, **_kwargs):
+            return {}
+
+        module.placement_group_table = placement_group_table
+
+    if not hasattr(module, "remove_placement_group"):
+
+        def remove_placement_group(*_args, **_kwargs):
+            return None
+
+        module.remove_placement_group = remove_placement_group
+
+    if not hasattr(ray_util, "placement_group_table"):
+        ray_util.placement_group_table = module.placement_group_table
+
+    strategies_module_name = "ray.util.scheduling_strategies"
+    try:
+        strategies_module = importlib.import_module(strategies_module_name)
+    except Exception:
+        strategies_module = types.ModuleType(strategies_module_name)
+        sys.modules[strategies_module_name] = strategies_module
+
+    if not hasattr(strategies_module, "PlacementGroupSchedulingStrategy"):
+
+        class PlacementGroupSchedulingStrategy:
+            def __init__(
+                self,
+                placement_group=None,
+                placement_group_bundle_index: int | None = None,
+                placement_group_capture_child_tasks: bool | None = None,
+                **kwargs,
+            ):
+                self.placement_group = placement_group
+                self.placement_group_bundle_index = placement_group_bundle_index
+                self.placement_group_capture_child_tasks = placement_group_capture_child_tasks
+                self.kwargs = kwargs
+
+        strategies_module.PlacementGroupSchedulingStrategy = PlacementGroupSchedulingStrategy
+
+    if not hasattr(strategies_module, "NodeAffinitySchedulingStrategy"):
+
+        class NodeAffinitySchedulingStrategy:
+            def __init__(self, node_id=None, soft: bool | None = None, **kwargs):
+                self.node_id = node_id
+                self.soft = soft
+                self.kwargs = kwargs
+
+        strategies_module.NodeAffinitySchedulingStrategy = NodeAffinitySchedulingStrategy
+
+    if not hasattr(ray_util, "scheduling_strategies"):
+        try:
+            setattr(ray_util, "scheduling_strategies", strategies_module)
+        except Exception:
+            pass
+
+    queue_module_name = "ray.util.queue"
+    try:
+        queue_module = importlib.import_module(queue_module_name)
+    except Exception:
+        queue_module = types.ModuleType(queue_module_name)
+        sys.modules[queue_module_name] = queue_module
+
+    if not hasattr(queue_module, "Queue"):
+
+        class Queue:
+            def __init__(self, maxsize: int = 0, *args, **kwargs):
+                del args, kwargs
+                self._queue = stdlib_queue.Queue(maxsize=maxsize)
+
+            def put(self, item, block: bool = True, timeout: float | None = None):
+                return self._queue.put(item, block=block, timeout=timeout)
+
+            def get(self, block: bool = True, timeout: float | None = None):
+                return self._queue.get(block=block, timeout=timeout)
+
+            def empty(self) -> bool:
+                return self._queue.empty()
+
+            def qsize(self) -> int:
+                return self._queue.qsize()
+
+        queue_module.Queue = Queue
+
+    if not hasattr(ray_util, "queue"):
+        try:
+            setattr(ray_util, "queue", queue_module)
+        except Exception:
+            pass
+
+
+
+def install_vllm_omni_compat() -> None:
+    """Reject the legacy broad compatibility mode in the review branch."""
+
+    raise RuntimeError(
+        "EASYMAGPIE_VLLM_COMPAT_MODE=full/all is not included in the "
+        "review-friendly EasyMagpie RL branch. Use refit, runtime, serial, "
+        "or rl mode for the NeMo-RL rollout/refit path."
+    )
+
+
+__all__ = [
+    "install_easy_magpie_refit_rpc_compat",
+    "install_easy_magpie_runtime_compat",
+    "install_vllm_omni_compat",
+]
diff --git a/examples/tts/easymagpie_vllm_omni/tests/conftest.py b/examples/tts/easymagpie_vllm_omni/tests/conftest.py
index 8170bd62b34c..8d8e794cdd16 100644
--- a/examples/tts/easymagpie_vllm_omni/tests/conftest.py
+++ b/examples/tts/easymagpie_vllm_omni/tests/conftest.py
@@ -15,8 +15,8 @@
 
 The model definition (``easymagpie_vllm_omni.local_transformer`` etc.) is plain
 PyTorch: the ``@support_torch_compile`` decorator short-circuits to eager when
-``compilation_config.mode == CompilationMode.NONE``, and the modules only read a
-handful of scalars off the ``VllmConfig``. So the whole stack can be exercised as
+the compilation config disables compilation, and the modules only read a handful
+of scalars off the ``VllmConfig``. So the whole stack can be exercised as
 ordinary PyTorch with a tiny stand-in config — **no model directory, no engine,
 no GPU required** — which is what these fixtures provide.
 
@@ -55,20 +55,32 @@ def build_vllm_config(**arch_overrides):
     * ``model_config.hf_config`` — arch scalars (read via ``from_hf_config``);
     * ``model_config.dtype`` — buffer dtype;
     * ``scheduler_config.max_num_batched_tokens`` — scratch-buffer length;
-    * ``compilation_config.mode`` — ``CompilationMode.NONE`` so the
+    * ``compilation_config`` — disables compilation so the
       ``@support_torch_compile`` wrapper stays in eager mode.
 
     Any keyword overrides are merged into the default tiny arch profile.
     """
     import torch
-    from vllm.config import CompilationMode
+
+    try:
+        from vllm.config import CompilationLevel
+    except ImportError:
+        from vllm.config.compilation import CompilationLevel
+    try:
+        from vllm.config import CompilationMode
+        compilation_mode = CompilationMode.NONE
+    except ImportError:
+        compilation_mode = None
 
     arch = {**_DEFAULT_ARCH, **arch_overrides}
     hf_config = types.SimpleNamespace(**arch)
     return types.SimpleNamespace(
         model_config=types.SimpleNamespace(hf_config=hf_config, dtype=torch.float32),
         scheduler_config=types.SimpleNamespace(max_num_batched_tokens=128),
-        compilation_config=types.SimpleNamespace(mode=CompilationMode.NONE),
+        compilation_config=types.SimpleNamespace(
+            level=CompilationLevel.NO_COMPILATION,
+            mode=compilation_mode,
+        ),
     )
 
 
diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_config.py b/examples/tts/easymagpie_vllm_omni/tests/test_config.py
index e8955d348328..8c9744a650dc 100644
--- a/examples/tts/easymagpie_vllm_omni/tests/test_config.py
+++ b/examples/tts/easymagpie_vllm_omni/tests/test_config.py
@@ -27,6 +27,8 @@
     SPECIAL_AUDIO_EOS,
     SPECIAL_AUDIO_MASK,
     EasyMagpieOmniArch,
+    derive_nemotron_h_hybrid_pattern_from_weight_keys,
+    normalize_nemotron_h_config,
 )
 
 
@@ -86,3 +88,29 @@ def test_from_hf_config_hidden_size_fallback():
     assert arch.hidden_dim == 999
     # embedding_dim defaults to the same backbone width when not given explicitly.
     assert arch.embedding_dim == 999
+
+
+def test_normalize_nemotron_h_config_adds_norm_alias_but_preserves_moe_layers():
+    hf_config = types.SimpleNamespace(layer_norm_epsilon=1e-5, hybrid_override_pattern="MEM*E")
+    out = normalize_nemotron_h_config(hf_config)
+    assert out.rms_norm_eps == 1e-5
+    assert out.hybrid_override_pattern == "MEM*E"
+
+
+def test_derive_nemotron_h_hybrid_pattern_from_weight_keys_prefers_checkpoint_layout():
+    keys = [
+        "decoder.layers.0.mixer.A_log",
+        "decoder.layers.0.mixer.conv1d.weight",
+        "decoder.layers.1.mixer.experts.0.down_proj.weight",
+        "decoder.layers.1.mixer.experts.0.up_proj.weight",
+        "decoder.layers.1.mixer.shared_experts.down_proj.weight",
+        "decoder.layers.1.mixer.gate.weight",
+        "decoder.layers.2.mixer.q_proj.weight",
+        "decoder.layers.2.mixer.k_proj.weight",
+        "decoder.layers.2.mixer.v_proj.weight",
+        "decoder.layers.2.mixer.o_proj.weight",
+        "decoder.layers.3.mixer.experts.0.down_proj.weight",
+        "decoder.layers.3.mixer.experts.0.up_proj.weight",
+    ]
+
+    assert derive_nemotron_h_hybrid_pattern_from_weight_keys(keys) == "ME*E"
diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_easymagpie_backbone_cache.py b/examples/tts/easymagpie_vllm_omni/tests/test_easymagpie_backbone_cache.py
new file mode 100644
index 000000000000..03d5ba8e0502
--- /dev/null
+++ b/examples/tts/easymagpie_vllm_omni/tests/test_easymagpie_backbone_cache.py
@@ -0,0 +1,282 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for EasyMagpie's vLLM Nemotron-H backbone call contract."""
+from __future__ import annotations
+
+import sys
+import types
+from types import SimpleNamespace
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("vllm")
+
+
+def _install_omni_output_stub() -> None:
+    if "vllm_omni.model_executor.models.output_templates" in sys.modules:
+        return
+
+    omni_pkg = types.ModuleType("vllm_omni")
+    omni_pkg.__path__ = []
+    model_executor_pkg = types.ModuleType("vllm_omni.model_executor")
+    model_executor_pkg.__path__ = []
+    models_pkg = types.ModuleType("vllm_omni.model_executor.models")
+    models_pkg.__path__ = []
+    output_templates = types.ModuleType("vllm_omni.model_executor.models.output_templates")
+
+    class OmniOutput:
+        def __init__(self, text_hidden_states=None, multimodal_outputs=None):
+            self.text_hidden_states = text_hidden_states
+            self.multimodal_outputs = multimodal_outputs or {}
+
+    output_templates.OmniOutput = OmniOutput
+    sys.modules.setdefault("vllm_omni", omni_pkg)
+    sys.modules.setdefault("vllm_omni.model_executor", model_executor_pkg)
+    sys.modules.setdefault("vllm_omni.model_executor.models", models_pkg)
+    sys.modules["vllm_omni.model_executor.models.output_templates"] = output_templates
+
+
+_install_omni_output_stub()
+
+import easymagpie_vllm_omni.easymagpie as easymagpie_module  # noqa: E402
+from easymagpie_vllm_omni.easymagpie import EasyMagpieTTSForConditionalGeneration  # noqa: E402
+
+
+class _BackboneRequiringMambaCache(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.seen_mamba_cache_params = None
+
+    def forward(
+        self,
+        *,
+        input_ids,
+        positions,
+        mamba_cache_params,
+        intermediate_tensors=None,
+        inputs_embeds=None,
+    ):
+        self.seen_mamba_cache_params = mamba_cache_params
+        return inputs_embeds
+
+
+class _FakeMambaCache:
+    def __init__(self) -> None:
+        self.capture_batch_size = None
+        self.conv_state = torch.ones(1, 2, 3)
+        self.ssm_state = torch.ones(1, 2, 4)
+        self.state_indices_tensor = torch.tensor([0, 1], dtype=torch.int32)
+
+    def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
+        self.capture_batch_size = batch_size
+        return (self.conv_state, self.ssm_state), self.state_indices_tensor[:batch_size]
+
+
+class _FakeCurrentRunMambaCache(_FakeMambaCache):
+    def __init__(self) -> None:
+        super().__init__()
+        self.current_run_kwargs = None
+
+    def current_run_tensors(self, **kwargs):
+        self.current_run_kwargs = dict(kwargs)
+        return SimpleNamespace(
+            conv_state=self.conv_state,
+            ssm_state=self.ssm_state,
+            state_indices_tensor=self.state_indices_tensor,
+        )
+
+
+def _minimal_model(backbone: torch.nn.Module) -> EasyMagpieTTSForConditionalGeneration:
+    model = EasyMagpieTTSForConditionalGeneration.__new__(EasyMagpieTTSForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.backbone = backbone
+    model._backbone_accepts_mamba_cache_params = True
+    model.has_phoneme = False
+    model.mamba_cache = None
+
+    max_tokens = 4
+    hidden_dim = 3
+    num_codebooks = 2
+    model._combined_embeddings = torch.zeros(max_tokens, hidden_dim)
+    model._token_stop = torch.zeros(max_tokens, dtype=torch.bool)
+    model._sample_stop = torch.zeros(max_tokens, dtype=torch.bool)
+    model._out_codes = torch.zeros(max_tokens, num_codebooks, dtype=torch.long)
+    model._out_code_logprobs = torch.zeros(max_tokens, num_codebooks)
+    model._out_code_sampling_logprobs = torch.zeros(max_tokens, num_codebooks)
+    model._out_frame_logprobs = torch.zeros(max_tokens)
+    model.code_predictor = SimpleNamespace()
+    model._get_decode_idxs = lambda: (torch.empty(0, dtype=torch.long), 0)
+    return model
+
+
+def test_forward_passes_mamba_cache_params_to_vllm_backbone():
+    backbone = _BackboneRequiringMambaCache()
+    model = _minimal_model(backbone)
+    mamba_cache_params = object()
+
+    output = model(
+        input_ids=torch.tensor([1, 2]),
+        positions=torch.tensor([0, 1]),
+        inputs_embeds=torch.ones(2, 3),
+        mamba_cache_params=mamba_cache_params,
+    )
+
+    assert backbone.seen_mamba_cache_params is mamba_cache_params
+    assert torch.equal(output, torch.ones(2, 3))
+
+
+def test_forward_without_v0_cache_metadata_uses_forward_context_cache(monkeypatch):
+    monkeypatch.setattr(easymagpie_module.envs, "VLLM_USE_V1", False)
+    backbone = _BackboneRequiringMambaCache()
+    model = _minimal_model(backbone)
+
+    output = model(
+        input_ids=torch.tensor([1, 2]),
+        positions=torch.tensor([0, 1]),
+        inputs_embeds=torch.ones(2, 3),
+    )
+
+    assert backbone.seen_mamba_cache_params is None
+    assert torch.equal(output, torch.ones(2, 3))
+
+
+def test_missing_vllm_mamba_cache_manager_disables_v0_cache(monkeypatch):
+    monkeypatch.setattr(easymagpie_module, "_HAS_VLLM_MAMBA_CACHE_MANAGER", False)
+    backbone = _BackboneRequiringMambaCache()
+    model = _minimal_model(backbone)
+    model.vllm_config = object()
+    model.model_config = SimpleNamespace(get_num_layers_by_block_type=lambda *_args, **_kwargs: 1)
+
+    assert model._ensure_v0_mamba_cache() is None
+
+
+def test_forward_without_v0_cache_metadata_uses_profile_forward_context_cache(monkeypatch):
+    monkeypatch.setattr(easymagpie_module.envs, "VLLM_USE_V1", False)
+    monkeypatch.setattr(
+        easymagpie_module,
+        "get_forward_context",
+        lambda: SimpleNamespace(
+            attn_metadata=SimpleNamespace(
+                num_prefills=1,
+                num_decode_tokens=1,
+            )
+        ),
+    )
+    backbone = _BackboneRequiringMambaCache()
+    model = _minimal_model(backbone)
+    cache = _FakeMambaCache()
+    model.mamba_cache = cache
+
+    output = model(
+        input_ids=torch.tensor([1, 2]),
+        positions=torch.tensor([0, 1]),
+        inputs_embeds=torch.ones(2, 3),
+    )
+
+    assert cache.capture_batch_size == 2
+    assert backbone.seen_mamba_cache_params is not None
+    assert backbone.seen_mamba_cache_params.conv_state is cache.conv_state
+    assert backbone.seen_mamba_cache_params.ssm_state is cache.ssm_state
+    assert torch.equal(backbone.seen_mamba_cache_params.state_indices_tensor, cache.state_indices_tensor)
+    assert torch.equal(output, torch.ones(2, 3))
+
+
+def test_forward_without_v0_cache_metadata_uses_explicit_profile_cache_batch_size(monkeypatch):
+    monkeypatch.setattr(easymagpie_module.envs, "VLLM_USE_V1", False)
+    backbone = _BackboneRequiringMambaCache()
+    model = _minimal_model(backbone)
+    cache = _FakeMambaCache()
+    model.mamba_cache = cache
+
+    output = model(
+        input_ids=torch.tensor([1, 2]),
+        positions=torch.tensor([0, 1]),
+        inputs_embeds=torch.ones(2, 3),
+        easymagpie_mamba_cache_batch_size=2,
+    )
+
+    assert cache.capture_batch_size == 2
+    assert backbone.seen_mamba_cache_params is not None
+    assert backbone.seen_mamba_cache_params.conv_state is cache.conv_state
+    assert backbone.seen_mamba_cache_params.ssm_state is cache.ssm_state
+    assert torch.equal(backbone.seen_mamba_cache_params.state_indices_tensor, cache.state_indices_tensor)
+    assert torch.equal(output, torch.ones(2, 3))
+
+
+def test_forward_prefers_request_cache_kwargs_over_profile_batch_size(monkeypatch):
+    monkeypatch.setattr(easymagpie_module.envs, "VLLM_USE_V1", True)
+    backbone = _BackboneRequiringMambaCache()
+    model = _minimal_model(backbone)
+    cache = _FakeCurrentRunMambaCache()
+    model.mamba_cache = cache
+
+    output = model(
+        input_ids=torch.tensor([1, 2]),
+        positions=torch.tensor([0, 1]),
+        inputs_embeds=torch.ones(2, 3),
+        easymagpie_mamba_cache_batch_size=2,
+        request_ids_to_seq_ids={"req-a": [0], "req-b": [1]},
+        finished_requests_ids=["old-req"],
+    )
+
+    assert cache.capture_batch_size is None
+    assert cache.current_run_kwargs["request_ids_to_seq_ids"] == {"req-a": [0], "req-b": [1]}
+    assert cache.current_run_kwargs["finished_requests_ids"] == ["old-req"]
+    assert backbone.seen_mamba_cache_params is not None
+    assert backbone.seen_mamba_cache_params.conv_state is cache.conv_state
+    assert torch.equal(output, torch.ones(2, 3))
+
+
+def test_mamba_prefill_only_metadata_returns_empty_decode_indices(monkeypatch):
+    model = EasyMagpieTTSForConditionalGeneration.__new__(EasyMagpieTTSForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model._combined_embeddings = torch.zeros(4, 3)
+    monkeypatch.setattr(
+        easymagpie_module,
+        "get_forward_context",
+        lambda: SimpleNamespace(
+            attn_metadata=SimpleNamespace(
+                num_prefills=1,
+                num_decode_tokens=0,
+            )
+        ),
+    )
+
+    decode_idx, num_req = model._get_decode_idxs()
+
+    assert num_req == 0
+    assert decode_idx is not None
+    assert decode_idx.numel() == 0
+
+
+def test_mamba_decode_only_metadata_keeps_decode_all_path(monkeypatch):
+    model = EasyMagpieTTSForConditionalGeneration.__new__(EasyMagpieTTSForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model._combined_embeddings = torch.zeros(4, 3)
+    monkeypatch.setattr(
+        easymagpie_module,
+        "get_forward_context",
+        lambda: SimpleNamespace(
+            attn_metadata=SimpleNamespace(
+                num_prefills=0,
+                num_decode_tokens=2,
+            )
+        ),
+    )
+
+    decode_idx, num_req = model._get_decode_idxs()
+
+    assert decode_idx is None
+    assert num_req == 0
diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py
new file mode 100644
index 000000000000..c612d3ac84fb
--- /dev/null
+++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py
@@ -0,0 +1,493 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Import-mode tests for EasyMagpie vLLM-Omni startup hooks."""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import types
+from pathlib import Path
+
+
+def test_compat_mode_off_does_not_import_vllm() -> None:
+    root = Path(__file__).resolve().parents[1]
+    env = os.environ.copy()
+    env["EASYMAGPIE_VLLM_COMPAT_MODE"] = "off"
+    env["PYTHONPATH"] = str(root)
+
+    code = """
+import json
+import sys
+import easymagpie_vllm_omni
+loaded = sorted(name for name in sys.modules if name == "vllm" or name.startswith("vllm."))
+print(json.dumps(loaded))
+"""
+    proc = subprocess.run(
+        [sys.executable, "-c", code],
+        check=True,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        timeout=30,
+    )
+    loaded = json.loads(proc.stdout.strip())
+    assert loaded == []
+
+
+def test_compat_mode_serial_imports_review_path() -> None:
+    root = Path(__file__).resolve().parents[1]
+    env = os.environ.copy()
+    env["EASYMAGPIE_VLLM_COMPAT_MODE"] = "serial"
+    env["PYTHONPATH"] = str(root)
+
+    code = """
+import easymagpie_vllm_omni
+print("ok")
+"""
+    proc = subprocess.run(
+        [sys.executable, "-c", code],
+        check=True,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        timeout=30,
+    )
+    assert proc.stdout.strip().splitlines()[-1] == "ok"
+
+
+def test_compat_mode_full_is_not_part_of_review_branch() -> None:
+    root = Path(__file__).resolve().parents[1]
+    env = os.environ.copy()
+    env["EASYMAGPIE_VLLM_COMPAT_MODE"] = "full"
+    env["PYTHONPATH"] = str(root)
+
+    proc = subprocess.run(
+        [sys.executable, "-c", "import easymagpie_vllm_omni"],
+        check=False,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        timeout=30,
+    )
+    assert proc.returncode != 0
+    assert "not included in the review-friendly EasyMagpie RL branch" in proc.stderr
+
+
+def test_refit_compat_installs_named_runtime_reset_rpc(monkeypatch) -> None:
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_refit_rpc_compat
+
+    class FakeWorker:
+        pass
+
+    worker_module = types.ModuleType("vllm_omni.worker.gpu_ar_worker")
+    worker_module.GPUARWorker = FakeWorker
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker", types.ModuleType("vllm_omni.worker"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker.gpu_ar_worker", worker_module)
+    monkeypatch.delitem(sys.modules, "vllm_omni.worker.gpu_generation_worker", raising=False)
+
+    install_easy_magpie_refit_rpc_compat()
+
+    worker = FakeWorker()
+    worker.model_runner = types.SimpleNamespace(
+        _easymagpie_active_mamba_request_ids={"req"},
+        _easymagpie_mamba_request_cache_kwargs={"req": {}},
+        _easymagpie_mamba_request_cache_state={"req": {}},
+        _easymagpie_mamba_cache_batch_size=1,
+        requests={"req": object()},
+        encoder_cache={"req": object()},
+        input_batch=None,
+    )
+
+    result = worker.easymagpie_reset_runtime_state()
+
+    assert result["ok"] is True
+    assert result["runtime_reset_rpc_compat_version"] >= 1
+    assert set(result["runtime_state_reset"]["cleared_attrs"]) == {
+        "_easymagpie_active_mamba_request_ids",
+        "_easymagpie_mamba_request_cache_kwargs",
+        "_easymagpie_mamba_request_cache_state",
+        "_easymagpie_mamba_cache_batch_size",
+    }
+    assert result["runtime_state_reset"]["cleared_mappings"] == [
+        {"name": "requests", "size_before": 1},
+        {"name": "encoder_cache", "size_before": 1},
+    ]
+
+
+def test_runtime_compat_handles_new_vllm_init_model_kwargs_signature(monkeypatch) -> None:
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat
+
+    class NewParentRunner:
+        def _init_model_kwargs(self):
+            return {"source": "new-vllm"}
+
+    class OmniGPUModelRunner(NewParentRunner):
+        max_num_tokens = 17
+
+        def _init_model_kwargs(self, num_tokens=None):
+            if num_tokens is None:
+                num_tokens = int(getattr(self, "max_num_tokens", 0) or 0)
+            return super()._init_model_kwargs(num_tokens)
+
+    class RuntimeRunner(OmniGPUModelRunner):
+        pass
+
+    runner_module = types.ModuleType("vllm_omni.worker.gpu_model_runner")
+    runner_module.OmniGPUModelRunner = OmniGPUModelRunner
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker", types.ModuleType("vllm_omni.worker"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker.gpu_model_runner", runner_module)
+
+    try:
+        OmniGPUModelRunner()._init_model_kwargs()
+    except TypeError as exc:
+        assert "positional argument" in str(exc)
+    else:
+        raise AssertionError("test fixture should reproduce the vLLM 0.21 signature mismatch")
+
+    install_easy_magpie_runtime_compat()
+    install_easy_magpie_runtime_compat()
+
+    assert OmniGPUModelRunner()._init_model_kwargs() == {"source": "new-vllm"}
+    assert RuntimeRunner()._init_model_kwargs() == {"source": "new-vllm"}
+
+
+def test_runtime_compat_preserves_old_vllm_init_model_kwargs_signature(monkeypatch) -> None:
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat
+
+    class OldParentRunner:
+        def _init_model_kwargs(self, num_tokens=None):
+            return {"source": "old-vllm", "num_tokens": num_tokens}
+
+    class OmniGPUModelRunner(OldParentRunner):
+        max_num_tokens = 19
+
+        def _init_model_kwargs(self, num_tokens=None):
+            if num_tokens is None:
+                num_tokens = int(getattr(self, "max_num_tokens", 0) or 0)
+            return super()._init_model_kwargs(num_tokens)
+
+    class RuntimeRunner(OmniGPUModelRunner):
+        pass
+
+    runner_module = types.ModuleType("vllm_omni.worker.gpu_model_runner")
+    runner_module.OmniGPUModelRunner = OmniGPUModelRunner
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker", types.ModuleType("vllm_omni.worker"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker.gpu_model_runner", runner_module)
+
+    install_easy_magpie_runtime_compat()
+    install_easy_magpie_runtime_compat()
+
+    assert OmniGPUModelRunner()._init_model_kwargs() == {"source": "old-vllm", "num_tokens": 19}
+    assert OmniGPUModelRunner()._init_model_kwargs(5) == {"source": "old-vllm", "num_tokens": 5}
+    assert RuntimeRunner()._init_model_kwargs() == {"source": "old-vllm", "num_tokens": 19}
+    assert RuntimeRunner()._init_model_kwargs(5) == {"source": "old-vllm", "num_tokens": 5}
+
+
+def test_runtime_compat_supplies_new_vllm_cumsum_arange_output_buffer(monkeypatch) -> None:
+    import numpy as np
+
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat
+
+    class NewParentRunner:
+        def __init__(self):
+            self.arange_np = np.arange(16, dtype=np.int32)
+
+        def _get_cumsum_and_arange(self, num_tokens, arange_out):
+            cu_num_tokens = np.cumsum(num_tokens, dtype=np.int32)
+            offsets = np.repeat(cu_num_tokens - num_tokens, num_tokens)
+            np.subtract(self.arange_np[: int(cu_num_tokens[-1])], offsets, out=arange_out)
+            return cu_num_tokens, arange_out
+
+    class OmniGPUModelRunner(NewParentRunner):
+        pass
+
+    runner_module = types.ModuleType("vllm_omni.worker.gpu_model_runner")
+    runner_module.OmniGPUModelRunner = OmniGPUModelRunner
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker", types.ModuleType("vllm_omni.worker"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker.gpu_model_runner", runner_module)
+
+    try:
+        OmniGPUModelRunner()._get_cumsum_and_arange(np.array([2, 3], dtype=np.int32))
+    except TypeError as exc:
+        assert "arange_out" in str(exc)
+    else:
+        raise AssertionError("test fixture should reproduce the vLLM arange_out signature mismatch")
+
+    install_easy_magpie_runtime_compat()
+    install_easy_magpie_runtime_compat()
+
+    cu_num_tokens, arange = OmniGPUModelRunner()._get_cumsum_and_arange(np.array([2, 3], dtype=np.int32))
+
+    assert cu_num_tokens.tolist() == [2, 5]
+    assert arange.tolist() == [0, 1, 0, 1, 2]
+
+
+def test_runtime_compat_normalizes_new_vllm_cumsum_single_return(monkeypatch) -> None:
+    import numpy as np
+
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat
+
+    class NewParentRunner:
+        def __init__(self):
+            self.arange_np = np.arange(16, dtype=np.int32)
+
+        def _get_cumsum_and_arange(self, num_tokens, arange_out):
+            cu_num_tokens = np.cumsum(num_tokens, dtype=np.int32)
+            offsets = np.repeat(cu_num_tokens - num_tokens, num_tokens)
+            np.subtract(self.arange_np[: int(cu_num_tokens[-1])], offsets, out=arange_out)
+            return cu_num_tokens
+
+    class OmniGPUModelRunner(NewParentRunner):
+        pass
+
+    runner_module = types.ModuleType("vllm_omni.worker.gpu_model_runner")
+    runner_module.OmniGPUModelRunner = OmniGPUModelRunner
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker", types.ModuleType("vllm_omni.worker"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker.gpu_model_runner", runner_module)
+
+    install_easy_magpie_runtime_compat()
+
+    cu_num_tokens, arange = OmniGPUModelRunner()._get_cumsum_and_arange(np.array([2, 3], dtype=np.int32))
+
+    assert cu_num_tokens.tolist() == [2, 5]
+    assert arange.tolist() == [0, 1, 0, 1, 2]
+
+
+def test_runtime_compat_preserves_old_vllm_cumsum_arange_signature(monkeypatch) -> None:
+    import numpy as np
+
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat
+
+    class OldParentRunner:
+        def __init__(self):
+            self.arange_np = np.arange(16, dtype=np.int32)
+
+        def _get_cumsum_and_arange(self, num_tokens, cumsum_dtype=None):
+            cu_num_tokens = np.cumsum(num_tokens, dtype=cumsum_dtype)
+            offsets = np.repeat(cu_num_tokens - num_tokens, num_tokens)
+            return cu_num_tokens, self.arange_np[: int(cu_num_tokens[-1])] - offsets
+
+    class OmniGPUModelRunner(OldParentRunner):
+        pass
+
+    runner_module = types.ModuleType("vllm_omni.worker.gpu_model_runner")
+    runner_module.OmniGPUModelRunner = OmniGPUModelRunner
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker", types.ModuleType("vllm_omni.worker"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker.gpu_model_runner", runner_module)
+
+    install_easy_magpie_runtime_compat()
+    install_easy_magpie_runtime_compat()
+
+    cu_num_tokens, arange = OmniGPUModelRunner()._get_cumsum_and_arange(
+        np.array([2, 3], dtype=np.int32),
+        cumsum_dtype=np.int64,
+    )
+
+    assert cu_num_tokens.dtype == np.int64
+    assert cu_num_tokens.tolist() == [2, 5]
+    assert arange.tolist() == [0, 1, 0, 1, 2]
+
+
+def test_runtime_compat_reports_configured_cudagraph_mode_for_omni_fallback(monkeypatch) -> None:
+    from enum import Enum
+
+    class CUDAGraphMode(Enum):
+        NONE = 0
+        PIECEWISE = 1
+        FULL = 2
+        FULL_AND_PIECEWISE = 3
+
+        def mixed_mode(self):
+            if self == CUDAGraphMode.FULL_AND_PIECEWISE:
+                return CUDAGraphMode.PIECEWISE
+            return self
+
+        def decode_mode(self):
+            if self == CUDAGraphMode.FULL_AND_PIECEWISE:
+                return CUDAGraphMode.FULL
+            return self
+
+    config_module = types.ModuleType("vllm.config")
+    config_module.CUDAGraphMode = CUDAGraphMode
+    monkeypatch.setitem(sys.modules, "vllm.config", config_module)
+
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat
+
+    class BatchDescriptor:
+        def __init__(self, num_tokens=8, num_reqs=4):
+            self.num_tokens = num_tokens
+            self.num_reqs = num_reqs
+
+    class OmniGPUModelRunner:
+        def __init__(self, cudagraph_mode):
+            self.vllm_config = types.SimpleNamespace(
+                compilation_config=types.SimpleNamespace(
+                    cudagraph_mode=cudagraph_mode,
+                    cudagraph_capture_sizes=[1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64],
+                    max_cudagraph_capture_size=64,
+                )
+            )
+
+        def _init_model_kwargs(self, num_tokens=None):
+            return {}
+
+        def _determine_batch_execution_and_padding(self, **kwargs):
+            return (
+                CUDAGraphMode.NONE,
+                BatchDescriptor(
+                    num_tokens=kwargs.get("num_tokens", 8),
+                    num_reqs=kwargs.get("num_reqs", 4),
+                ),
+                False,
+                None,
+                None,
+            )
+
+    runner_module = types.ModuleType("vllm_omni.worker.gpu_model_runner")
+    runner_module.OmniGPUModelRunner = OmniGPUModelRunner
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker", types.ModuleType("vllm_omni.worker"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.worker.gpu_model_runner", runner_module)
+
+    install_easy_magpie_runtime_compat()
+    install_easy_magpie_runtime_compat()
+
+    mixed_result = OmniGPUModelRunner(CUDAGraphMode.PIECEWISE)._determine_batch_execution_and_padding(
+        num_tokens=8,
+        num_reqs=4,
+        force_eager=False,
+        force_uniform_decode=False,
+    )
+    eager_result = OmniGPUModelRunner(CUDAGraphMode.PIECEWISE)._determine_batch_execution_and_padding(
+        num_tokens=8,
+        num_reqs=4,
+        force_eager=True,
+        force_uniform_decode=False,
+    )
+    full_and_piecewise_mixed = OmniGPUModelRunner(
+        CUDAGraphMode.FULL_AND_PIECEWISE
+    )._determine_batch_execution_and_padding(
+        num_tokens=8,
+        num_reqs=4,
+        force_eager=False,
+        force_uniform_decode=False,
+    )
+    full_and_piecewise_decode = OmniGPUModelRunner(
+        CUDAGraphMode.FULL_AND_PIECEWISE
+    )._determine_batch_execution_and_padding(
+        num_tokens=8,
+        num_reqs=4,
+        force_eager=False,
+        force_uniform_decode=True,
+    )
+    oversized_mixed = OmniGPUModelRunner(CUDAGraphMode.PIECEWISE)._determine_batch_execution_and_padding(
+        num_tokens=80,
+        num_reqs=32,
+        force_eager=False,
+        force_uniform_decode=False,
+    )
+
+    assert mixed_result[0] == CUDAGraphMode.PIECEWISE
+    assert eager_result[0] == CUDAGraphMode.NONE
+    assert full_and_piecewise_mixed[0] == CUDAGraphMode.PIECEWISE
+    assert full_and_piecewise_decode[0] == CUDAGraphMode.FULL
+    assert oversized_mixed[0] == CUDAGraphMode.NONE
+    assert oversized_mixed[1].num_tokens == 80
+
+    capture_desc = mixed_result[1]
+    runtime_desc = OmniGPUModelRunner(CUDAGraphMode.PIECEWISE)._determine_batch_execution_and_padding(
+        num_tokens=8,
+        num_reqs=4,
+        force_eager=False,
+        force_uniform_decode=False,
+    )[1]
+    assert capture_desc is not runtime_desc
+    assert capture_desc == runtime_desc
+    assert hash(capture_desc) == hash(runtime_desc)
+    assert capture_desc.num_tokens == 8
+    assert capture_desc.num_reqs == 4
+    assert capture_desc.uniform_decode is False
+
+
+def test_runtime_compat_handles_structured_output_request_signature_variants(monkeypatch) -> None:
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat
+
+    class KeywordStructuredOutputRequest:
+        def __init__(self, *, sampling_params):
+            self.mode = "keyword"
+            self.sampling_params = sampling_params
+
+    class PositionalStructuredOutputRequest:
+        def __init__(self, sampling_params):
+            self.mode = "positional"
+            self.sampling_params = sampling_params
+
+    class NoSamplingStructuredOutputRequest:
+        def __init__(self):
+            self.mode = "none"
+
+    request_module = types.ModuleType("vllm_omni.request")
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.request", request_module)
+
+    request_module.StructuredOutputRequest = KeywordStructuredOutputRequest
+    install_easy_magpie_runtime_compat()
+    assert request_module.StructuredOutputRequest(sampling_params="sp").mode == "keyword"
+
+    request_module.StructuredOutputRequest = PositionalStructuredOutputRequest
+    install_easy_magpie_runtime_compat()
+    positional_result = request_module.StructuredOutputRequest(sampling_params="sp")
+    assert positional_result.mode == "positional"
+    assert positional_result.sampling_params == "sp"
+
+    request_module.StructuredOutputRequest = NoSamplingStructuredOutputRequest
+    install_easy_magpie_runtime_compat()
+    assert request_module.StructuredOutputRequest(sampling_params="sp") is None
+
+
+def test_runtime_compat_supplies_legacy_omni_request_mm_kwargs(monkeypatch) -> None:
+    from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat
+
+    class OmniRequest:
+        def __init__(self, *, mm_features=None):
+            del mm_features
+            self.mm_features = self.mm_kwargs
+
+    request_module = types.ModuleType("vllm_omni.request")
+    request_module.OmniRequest = OmniRequest
+    monkeypatch.setitem(sys.modules, "vllm_omni", types.ModuleType("vllm_omni"))
+    monkeypatch.setitem(sys.modules, "vllm_omni.request", request_module)
+
+    install_easy_magpie_runtime_compat()
+
+    empty_request = request_module.OmniRequest()
+    assert empty_request.mm_kwargs == {}
+    assert empty_request.mm_features == {}
+
+    mm_features = {"audio": object()}
+    audio_request = request_module.OmniRequest(mm_features=mm_features)
+    assert audio_request.mm_kwargs is mm_features
+    assert audio_request.mm_features is mm_features
diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer.py b/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer.py
index b290aa510ba4..19ba45abe24a 100644
--- a/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer.py
+++ b/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer.py
@@ -38,7 +38,12 @@
 
 from conftest import build_vllm_config  # noqa: E402
 from easymagpie_vllm_omni.config import EasyMagpieOmniArch  # noqa: E402
-from easymagpie_vllm_omni.local_transformer import EasyMagpieCodePredictor  # noqa: E402
+from easymagpie_vllm_omni.local_transformer import (  # noqa: E402
+    EasyMagpieCodePredictor,
+    sample_codebook,
+    sample_codebook_with_logprobs,
+)
+from nemo.collections.tts.modules.magpietts_modules import LocalTransformerHelper  # noqa: E402
 from torch import nn  # noqa: E402
 
 # Two arch profiles: one where all widths are equal (in/out projections are
@@ -134,7 +139,7 @@ def _vllm_teacher_forced_logits(
     """
     num_tokens = dec_hidden.shape[0]
     n = cp.num_codebooks
-    lt_hidden = cp.lt_hidden
+    lt_hidden = cp._buf_inputs.shape[-1]
     buf = torch.zeros(num_tokens, n, lt_hidden, dtype=dec_hidden.dtype, device=dec_hidden.device)
     buf[:, 0, :] = cp.local_transformer_in_projection(dec_hidden)
     for k in range(n - 1):
@@ -154,14 +159,8 @@ def _copy_nemo_into_vllm(nemo: NeMoLocalTransformerStack, cp: EasyMagpieCodePred
     missing = []
     for name, param in cp.named_parameters():
         if name in nemo_sd:
-            src = nemo_sd[name]
-            # The FFN ships as kernel-1 Conv1d (``[out, in, 1]``) in NeMo but is a
-            # plain ``nn.Linear`` (``[out, in]``) here; squeeze the conv dim to
-            # match (mirrors ``EasyMagpieTTS.load_weights``).
-            if src.ndim == param.ndim + 1 and src.shape[-1] == 1:
-                src = src.squeeze(-1)
-            assert param.shape == src.shape, f"shape mismatch {name}"
-            param.data.copy_(src.to(param.dtype))
+            assert param.shape == nemo_sd[name].shape, f"shape mismatch {name}"
+            param.data.copy_(nemo_sd[name].to(param.dtype))
         else:
             missing.append(name)
     assert not missing, f"vLLM params with no NeMo counterpart: {missing}"
@@ -205,6 +204,43 @@ def test_local_transformer_matches_nemo(profile):
     assert argmax_mismatch == 0, f"{argmax_mismatch} argmax mismatches"
 
 
+@pytest.mark.unit
+def test_autoregressive_argmax_matches_nemo_helper():
+    """The fixed-buffer vLLM AR loop must match NeMo's growing-sequence helper."""
+    cp, nemo, arch = _build_pair(ARCH_PROFILES["equal_dims"])
+    cp.temperature = 0.0
+    topk = min(80, arch.num_all_tokens_per_codebook)
+    cp.top_k = topk
+    helper = LocalTransformerHelper(
+        local_transformer=nemo.local_transformer,
+        audio_embeddings=nemo.audio_embeddings,
+        audio_in_projection=nemo.audio_in_projection,
+        local_transformer_in_projection=nemo.local_transformer_in_projection,
+        local_transformer_audio_out_projection=nemo.local_transformer_audio_out_projection,
+        local_transformer_out_projections=nemo.local_transformer_out_projections,
+        num_audio_codebooks=arch.num_audio_codebooks,
+        frame_stacking_factor=arch.frame_stacking_factor,
+        audio_eos_id=arch.audio_eos_id,
+        mask_token_id=arch.mask_token_id,
+        codebook_size=arch.codebook_size,
+    )
+
+    torch.manual_seed(4321)
+    dec_hidden = torch.randn(4, arch.hidden_dim)
+    nemo_codes = helper.sample_autoregressive(
+        dec_output=dec_hidden,
+        temperature=0.0,
+        topk=topk,
+        use_cfg=False,
+        use_kv_cache=False,
+        sanitize_logits=True,
+    )
+    nemo_flat = nemo_codes.permute(0, 2, 1).reshape(dec_hidden.shape[0], -1)
+    vllm_flat = cp.generate_codes(dec_hidden)
+
+    assert torch.equal(vllm_flat, nemo_flat)
+
+
 @pytest.mark.unit
 def test_generate_codes_shape_dtype_and_range():
     """``generate_codes`` returns valid (num_tokens, num_codebooks) int64 codes within vocab."""
@@ -220,6 +256,29 @@ def test_generate_codes_shape_dtype_and_range():
     assert codes.max().item() < arch.num_all_tokens_per_codebook
 
 
+@pytest.mark.unit
+def test_generate_codes_with_logprobs_shape_and_finiteness():
+    """``generate_codes_with_logprobs`` returns finite selected-code logprobs."""
+    cp, _, arch = _build_pair(ARCH_PROFILES["equal_dims"])
+    num_tokens = 5
+
+    torch.manual_seed(0)
+    codes, model_logprobs, sampling_logprobs = cp.generate_codes_with_logprobs(
+        torch.randn(num_tokens, arch.hidden_dim)
+    )
+
+    expected_shape = (num_tokens, arch.num_stacked_codebooks)
+    assert codes.shape == expected_shape
+    assert model_logprobs.shape == expected_shape
+    assert sampling_logprobs.shape == expected_shape
+    assert model_logprobs.dtype == torch.float32
+    assert sampling_logprobs.dtype == torch.float32
+    assert torch.isfinite(model_logprobs).all()
+    assert torch.isfinite(sampling_logprobs).all()
+    assert (model_logprobs <= 0).all()
+    assert (sampling_logprobs <= 0).all()
+
+
 @pytest.mark.unit
 def test_generate_codes_respects_forbidden_mask():
     """With argmax sampling, forbidden special tokens are never emitted (only EOS stays reachable)."""
@@ -234,6 +293,89 @@ def test_generate_codes_respects_forbidden_mask():
     assert allowed.all(), f"sampled forbidden tokens: {sorted(set(codes[~allowed].tolist()))}"
 
 
+@pytest.mark.unit
+def test_sampler_sanitizes_non_finite_logits():
+    """Match the NeMo inference path's logit sanitization before top-k sampling."""
+
+    logits = torch.tensor([[float("nan"), float("inf"), float("-inf"), 5.0]])
+
+    sampled = sample_codebook(logits, temperature=0.0, top_k=0, forbidden_mask=None)
+    sampled_with_lp, model_lp, sampling_lp = sample_codebook_with_logprobs(
+        logits,
+        temperature=0.0,
+        top_k=0,
+        forbidden_mask=None,
+    )
+
+    assert sampled.tolist() == [1]
+    assert sampled_with_lp.tolist() == [1]
+    assert torch.isfinite(model_lp).all()
+    assert torch.isfinite(sampling_lp).all()
+
+
+@pytest.mark.unit
+def test_sampler_sanitizes_non_finite_logits_with_temperature_topk():
+    """Finite logprobs are required for RL rollouts with sampled top-k audio."""
+
+    logits = torch.tensor(
+        [
+            [float("nan"), float("inf"), float("-inf"), 5.0],
+            [float("nan"), float("-inf"), -4.0, 2.0],
+        ]
+    )
+
+    torch.manual_seed(0)
+    sampled, model_lp, sampling_lp = sample_codebook_with_logprobs(
+        logits,
+        temperature=0.7,
+        top_k=3,
+        forbidden_mask=None,
+    )
+
+    assert sampled.shape == (2,)
+    assert torch.isfinite(model_lp).all()
+    assert torch.isfinite(sampling_lp).all()
+
+
+@pytest.mark.unit
+def test_selected_logprobs_split_raw_model_and_sampling_distributions():
+    """Model logprobs match raw logits while sampling logprobs match masked top-k."""
+
+    logits = torch.tensor(
+        [
+            [1.5, -2.0, 0.25, 3.0],
+            [-0.5, 2.5, 0.75, 1.0],
+        ],
+        dtype=torch.float32,
+    )
+    forbidden_mask = torch.tensor(
+        [
+            [False, True, False, False],
+            [False, False, True, False],
+        ],
+        dtype=torch.bool,
+    )
+    temperature = 0.7
+    top_k = 2
+
+    torch.manual_seed(11)
+    sampled, model_lp, sampling_lp = sample_codebook_with_logprobs(
+        logits,
+        temperature=temperature,
+        top_k=top_k,
+        forbidden_mask=forbidden_mask,
+    )
+
+    expected_model_lp = torch.log_softmax(logits, dim=-1).gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
+    masked = logits.masked_fill(forbidden_mask, float("-inf"))
+    vals, idxs = torch.topk(masked / temperature, k=top_k, dim=-1)
+    sampled_in_k = (idxs == sampled.unsqueeze(-1)).nonzero(as_tuple=False)[:, -1]
+    expected_sampling_lp = torch.log_softmax(vals, dim=-1).gather(-1, sampled_in_k.unsqueeze(-1)).squeeze(-1)
+
+    assert torch.allclose(model_lp, expected_model_lp, atol=1e-6)
+    assert torch.allclose(sampling_lp, expected_sampling_lp, atol=1e-6)
+
+
 @pytest.mark.unit
 def test_generate_codes_deterministic_with_seed():
     """Same seed + same input ⇒ identical sampled codes (sampler is RNG-driven, no host state)."""
diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_sampler.py b/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_sampler.py
new file mode 100644
index 000000000000..7b42063d59f1
--- /dev/null
+++ b/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_sampler.py
@@ -0,0 +1,114 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+import torch
+from torch import nn
+
+from easymagpie_vllm_omni.local_transformer import (
+    EasyMagpieCodePredictor,
+    sample_codebook,
+    sample_codebook_with_logprobs,
+)
+
+
+def test_sampler_sanitizes_non_finite_logits_argmax():
+    logits = torch.tensor([[float("nan"), float("inf"), float("-inf"), 5.0]])
+
+    sampled = sample_codebook(logits, temperature=0.0, top_k=0, forbidden_mask=None)
+    sampled_with_lp, model_lp, sampling_lp = sample_codebook_with_logprobs(
+        logits,
+        temperature=0.0,
+        top_k=0,
+        forbidden_mask=None,
+    )
+
+    assert sampled.tolist() == [1]
+    assert sampled_with_lp.tolist() == [1]
+    assert torch.isfinite(model_lp).all()
+    assert torch.isfinite(sampling_lp).all()
+
+
+def test_sampler_sanitizes_non_finite_logits_with_temperature_topk():
+    logits = torch.tensor(
+        [
+            [float("nan"), float("inf"), float("-inf"), 5.0],
+            [float("nan"), float("-inf"), -4.0, 2.0],
+        ]
+    )
+
+    torch.manual_seed(0)
+    sampled, model_lp, sampling_lp = sample_codebook_with_logprobs(
+        logits,
+        temperature=0.7,
+        top_k=3,
+        forbidden_mask=None,
+    )
+
+    assert sampled.shape == (2,)
+    assert torch.isfinite(model_lp).all()
+    assert torch.isfinite(sampling_lp).all()
+
+
+class _RecordingLocalTransformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.shapes = []
+
+    def forward(self, inputs_embeds):
+        self.shapes.append(tuple(inputs_embeds.shape))
+        return inputs_embeds
+
+
+def test_code_predictor_pads_local_transformer_to_stable_graph_tokens(monkeypatch):
+    monkeypatch.setenv("EASYMAGPIE_LOCAL_TRANSFORMER_GRAPH_TOKENS", "8")
+    cp = EasyMagpieCodePredictor.__new__(EasyMagpieCodePredictor)
+    nn.Module.__init__(cp)
+    cp.num_codebooks = 2
+    cp.num_tokens_per_codebook = 8
+    cp._buf_inputs = torch.zeros(8, cp.num_codebooks, 4)
+    cp._out_codes = torch.zeros(8, cp.num_codebooks, dtype=torch.long)
+    cp._out_code_logprobs = torch.zeros(8, cp.num_codebooks, dtype=torch.float32)
+    cp._out_code_sampling_logprobs = torch.zeros(8, cp.num_codebooks, dtype=torch.float32)
+    cp.forbidden_mask = torch.zeros(cp.num_tokens_per_codebook, dtype=torch.bool)
+    cp._local_transformer_graph_tokens = EasyMagpieCodePredictor._resolve_local_transformer_graph_tokens(8)
+    cp.local_transformer_in_projection = nn.Identity()
+    cp.local_transformer_audio_out_projection = nn.Identity()
+    cp.local_transformer_out_projections = nn.ModuleList(
+        [nn.Linear(4, cp.num_tokens_per_codebook), nn.Linear(4, cp.num_tokens_per_codebook)]
+    )
+    cp.audio_embeddings = nn.ModuleList(
+        [nn.Embedding(cp.num_tokens_per_codebook, 4), nn.Embedding(cp.num_tokens_per_codebook, 4)]
+    )
+    cp.audio_in_projection = nn.Identity()
+    cp.temperature = 0.0
+    cp.top_k = 0
+    recorder = _RecordingLocalTransformer()
+    cp.local_transformer = recorder
+
+    codes_6, model_lp_6, sampling_lp_6 = cp.generate_codes_with_logprobs(
+        torch.randn(6, 4)
+    )
+    assert codes_6.shape == (6, cp.num_codebooks)
+    assert model_lp_6.shape == (6, cp.num_codebooks)
+    assert sampling_lp_6.shape == (6, cp.num_codebooks)
+    assert recorder.shapes
+    assert set(recorder.shapes) == {(8, cp.num_codebooks, cp._buf_inputs.shape[-1])}
+
+    recorder.shapes.clear()
+    codes_7, _, _ = cp.generate_codes_with_logprobs(
+        torch.randn(7, 4)
+    )
+    assert codes_7.shape == (7, cp.num_codebooks)
+    assert set(recorder.shapes) == {(8, cp.num_codebooks, cp._buf_inputs.shape[-1])}
+
+
+def test_max_concurrent_requests_does_not_enable_local_transformer_padding(monkeypatch):
+    monkeypatch.delenv("EASYMAGPIE_LOCAL_TRANSFORMER_GRAPH_TOKENS", raising=False)
+    monkeypatch.setenv("EASYMAGPIE_VLLM_MAX_CONCURRENT_REQUESTS", "8")
+
+    assert EasyMagpieCodePredictor._resolve_local_transformer_graph_tokens(8) == 0
diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_shape.py b/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_shape.py
new file mode 100644
index 000000000000..9a060f627fd9
--- /dev/null
+++ b/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_shape.py
@@ -0,0 +1,40 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Shape regression tests for the EasyMagpie local transformer."""
+from __future__ import annotations
+
+import inspect
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("vllm")
+
+from easymagpie_vllm_omni.local_transformer import EasyMagpieLTSelfAttention  # noqa: E402
+
+
+@pytest.mark.unit
+def test_local_self_attention_keeps_batch_dimension_inferred():
+    """The vLLM compiled graph may be replayed with a different decode batch size."""
+    source = inspect.getsource(EasyMagpieLTSelfAttention.forward)
+    assert ".reshape(-1," in source
+    assert ".reshape(b," not in source
+    assert ".view(b," not in source
+
+    attn = EasyMagpieLTSelfAttention(d_model=64, n_heads=4).eval()
+    first = torch.randn(1, 16, 64)
+    later = torch.randn(7, 16, 64)
+    with torch.no_grad():
+        assert attn(first).shape == first.shape
+        assert attn(later).shape == later.shape
diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_scheduler.py b/examples/tts/easymagpie_vllm_omni/tests/test_scheduler.py
new file mode 100644
index 000000000000..dd33c261b9f0
--- /dev/null
+++ b/examples/tts/easymagpie_vllm_omni/tests/test_scheduler.py
@@ -0,0 +1,315 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+
+
+class _FakeQueue(list):
+    def add_request(self, request):
+        self.append(request)
+
+
+class _FakeBaseScheduler:
+    pass
+
+
+def _load_scheduler_with_stubs(monkeypatch):
+    def fake_package(name: str):
+        module = types.ModuleType(name)
+        module.__path__ = []
+        monkeypatch.setitem(sys.modules, name, module)
+        return module
+
+    def fake_module(name: str):
+        module = types.ModuleType(name)
+        monkeypatch.setitem(sys.modules, name, module)
+        return module
+
+    for name in [
+        "vllm",
+        "vllm.v1",
+        "vllm.v1.core",
+        "vllm.v1.core.sched",
+        "vllm_omni",
+        "vllm_omni.core",
+        "vllm_omni.core.sched",
+    ]:
+        fake_package(name)
+
+    request_queue = fake_module("vllm.v1.core.sched.request_queue")
+    request_queue.create_request_queue = lambda policy: _FakeQueue()
+
+    request = fake_module("vllm.v1.request")
+    request.Request = object
+    request.StreamingUpdate = object
+
+    omni_scheduler = fake_module("vllm_omni.core.sched.omni_ar_scheduler")
+    omni_scheduler.OmniARAsyncScheduler = _FakeBaseScheduler
+    omni_scheduler.OmniARScheduler = _FakeBaseScheduler
+
+    module_path = (
+        Path(__file__).parents[1] / "easymagpie_vllm_omni" / "scheduler.py"
+    )
+    spec = importlib.util.spec_from_file_location(
+        "easymagpie_scheduler_under_test", module_path
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, spec.name, module)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_easy_scheduler_defers_waiting_prefills_while_running(monkeypatch):
+    scheduler_mod = _load_scheduler_with_stubs(monkeypatch)
+    scheduler = object.__new__(scheduler_mod.EasyMagpieARAsyncScheduler)
+    scheduler.disable_mixed_prefill_decode = True
+    scheduler.policy = "fcfs"
+    scheduler.running = [SimpleNamespace(request_id="decode-0")]
+    scheduler.waiting = _FakeQueue(
+        [
+            SimpleNamespace(request_id="prefill-0"),
+            SimpleNamespace(request_id="prefill-1"),
+        ]
+    )
+
+    def fake_base_schedule(self):
+        assert list(self.waiting) == []
+        self.waiting.add_request(SimpleNamespace(request_id="preempted-decode"))
+        return "scheduler-output"
+
+    monkeypatch.setattr(
+        scheduler_mod._OmniARBaseScheduler,
+        "schedule",
+        fake_base_schedule,
+        raising=False,
+    )
+
+    assert scheduler.schedule() == "scheduler-output"
+    assert [request.request_id for request in scheduler.waiting] == [
+        "preempted-decode",
+        "prefill-0",
+        "prefill-1",
+    ]
+    assert scheduler._easymagpie_last_mixed_guard == {
+        "deferred_waiting": 2,
+        "guarded_scheduled_tokens": None,
+        "fallback_to_waiting_prefill": False,
+    }
+
+
+def test_easy_scheduler_allows_prefill_when_guarded_decode_schedules_no_tokens(monkeypatch):
+    scheduler_mod = _load_scheduler_with_stubs(monkeypatch)
+    scheduler = object.__new__(scheduler_mod.EasyMagpieARAsyncScheduler)
+    scheduler.disable_mixed_prefill_decode = True
+    scheduler.policy = "fcfs"
+    scheduler.running = [
+        SimpleNamespace(
+            request_id="stale-decode",
+            status=SimpleNamespace(name="RUNNING"),
+            is_finished=lambda: False,
+            num_tokens_with_spec=129,
+            num_output_placeholders=0,
+            num_computed_tokens=128,
+        )
+    ]
+    scheduler.waiting = _FakeQueue([SimpleNamespace(request_id="prefill-0")])
+    calls = []
+
+    def fake_base_schedule(self):
+        calls.append([request.request_id for request in self.waiting])
+        if len(calls) == 1:
+            assert calls[-1] == []
+            return SimpleNamespace(num_scheduled_tokens={})
+        assert calls[-1] == ["prefill-0"]
+        return "prefill-output"
+
+    monkeypatch.setattr(
+        scheduler_mod._OmniARBaseScheduler,
+        "schedule",
+        fake_base_schedule,
+        raising=False,
+    )
+
+    assert scheduler.schedule() == "prefill-output"
+    assert calls == [[], ["prefill-0"]]
+    assert scheduler._easymagpie_last_mixed_guard == {
+        "deferred_waiting": 1,
+        "guarded_scheduled_tokens": 0,
+        "fallback_to_waiting_prefill": True,
+    }
+
+
+def test_easy_scheduler_does_not_defer_prefill_for_paused_running(monkeypatch):
+    scheduler_mod = _load_scheduler_with_stubs(monkeypatch)
+    scheduler = object.__new__(scheduler_mod.EasyMagpieARAsyncScheduler)
+    scheduler.disable_mixed_prefill_decode = True
+    scheduler.running = [
+        SimpleNamespace(
+            request_id="old-stream",
+            status=SimpleNamespace(name="WAITING_FOR_CHUNK"),
+            is_finished=lambda: False,
+        )
+    ]
+    scheduler.waiting = _FakeQueue([SimpleNamespace(request_id="prefill-0")])
+
+    def fake_base_schedule(self):
+        assert [request.request_id for request in self.waiting] == ["prefill-0"]
+        return "scheduler-output"
+
+    monkeypatch.setattr(
+        scheduler_mod._OmniARBaseScheduler,
+        "schedule",
+        fake_base_schedule,
+        raising=False,
+    )
+
+    assert scheduler.schedule() == "scheduler-output"
+
+
+def test_easy_scheduler_does_not_defer_prefill_for_stale_running_without_pending_tokens(monkeypatch):
+    scheduler_mod = _load_scheduler_with_stubs(monkeypatch)
+    scheduler = object.__new__(scheduler_mod.EasyMagpieARAsyncScheduler)
+    scheduler.disable_mixed_prefill_decode = True
+    scheduler.running = [
+        SimpleNamespace(
+            request_id="stale-running",
+            status=SimpleNamespace(name="RUNNING"),
+            is_finished=lambda: False,
+            num_tokens_with_spec=128,
+            num_output_placeholders=0,
+            num_computed_tokens=128,
+        )
+    ]
+    scheduler.waiting = _FakeQueue([SimpleNamespace(request_id="prefill-0")])
+
+    def fake_base_schedule(self):
+        assert [request.request_id for request in self.waiting] == ["prefill-0"]
+        return "scheduler-output"
+
+    monkeypatch.setattr(
+        scheduler_mod._OmniARBaseScheduler,
+        "schedule",
+        fake_base_schedule,
+        raising=False,
+    )
+
+    assert scheduler.schedule() == "scheduler-output"
+
+
+def test_easy_scheduler_purges_stale_train_running_before_post_refit(monkeypatch):
+    scheduler_mod = _load_scheduler_with_stubs(monkeypatch)
+    scheduler = object.__new__(scheduler_mod.EasyMagpieARAsyncScheduler)
+    scheduler.disable_mixed_prefill_decode = True
+    scheduler.purge_stale_train_before_post_refit = True
+    scheduler.running = [
+        SimpleNamespace(
+            request_id="easymagpie-grpo-train-rank-00000-step-000000-0-0",
+            status=SimpleNamespace(name="RUNNING"),
+            is_finished=lambda: False,
+            num_tokens_with_spec=129,
+            num_output_placeholders=0,
+            num_computed_tokens=128,
+        )
+    ]
+    scheduler.waiting = _FakeQueue(
+        [SimpleNamespace(request_id="easymagpie-grpo-train-post-refit-rank-00000-step-000000-0-0")]
+    )
+
+    def fake_base_schedule(self):
+        assert self.running == []
+        assert [request.request_id for request in self.waiting] == [
+            "easymagpie-grpo-train-post-refit-rank-00000-step-000000-0-0"
+        ]
+        return "post-refit-output"
+
+    monkeypatch.setattr(
+        scheduler_mod._OmniARBaseScheduler,
+        "schedule",
+        fake_base_schedule,
+        raising=False,
+    )
+
+    assert scheduler.schedule() == "post-refit-output"
+    assert scheduler._easymagpie_last_post_refit_purge == {
+        "num_purged_running": 1,
+        "purged_running_head": ["easymagpie-grpo-train-rank-00000-step-000000-0-0"],
+        "purged_running_tail": ["easymagpie-grpo-train-rank-00000-step-000000-0-0"],
+    }
+
+
+def test_easy_scheduler_detects_legacy_and_ranked_train_post_refit_requests(monkeypatch):
+    scheduler_mod = _load_scheduler_with_stubs(monkeypatch)
+    scheduler_cls = scheduler_mod.EasyMagpieARAsyncScheduler
+
+    train_ids = [
+        "easymagpie-grpo-train-step-000000-0-0",
+        "easymagpie-grpo-train-rank-00007-step-000000-1-7",
+    ]
+    post_refit_ids = [
+        "easymagpie-grpo-train-post-refit-step-000000-0-0",
+        "easymagpie-grpo-train-post-refit-rank-00007-step-000000-0-7",
+    ]
+    for request_id in train_ids:
+        assert scheduler_cls._is_train_rollout_request(SimpleNamespace(request_id=request_id))
+        assert not scheduler_cls._is_post_refit_request(SimpleNamespace(request_id=request_id))
+    for request_id in post_refit_ids:
+        assert scheduler_cls._is_post_refit_request(SimpleNamespace(request_id=request_id))
+        assert not scheduler_cls._is_train_rollout_request(SimpleNamespace(request_id=request_id))
+
+
+def test_easy_scheduler_keeps_non_train_running_before_post_refit(monkeypatch):
+    scheduler_mod = _load_scheduler_with_stubs(monkeypatch)
+    scheduler = object.__new__(scheduler_mod.EasyMagpieARAsyncScheduler)
+    scheduler.disable_mixed_prefill_decode = True
+    scheduler.purge_stale_train_before_post_refit = True
+    scheduler.policy = "fcfs"
+    scheduler.running = [
+        SimpleNamespace(
+            request_id="interactive-request",
+            status=SimpleNamespace(name="RUNNING"),
+            is_finished=lambda: False,
+            num_tokens_with_spec=129,
+            num_output_placeholders=0,
+            num_computed_tokens=128,
+        )
+    ]
+    scheduler.waiting = _FakeQueue(
+        [SimpleNamespace(request_id="easymagpie-grpo-train-post-refit-step-000000-0-0")]
+    )
+
+    def fake_base_schedule(self):
+        assert [request.request_id for request in self.waiting] == []
+        return SimpleNamespace(num_scheduled_tokens={"interactive-request": 1})
+
+    monkeypatch.setattr(
+        scheduler_mod._OmniARBaseScheduler,
+        "schedule",
+        fake_base_schedule,
+        raising=False,
+    )
+
+    result = scheduler.schedule()
+
+    assert result.num_scheduled_tokens == {"interactive-request": 1}
+    assert [request.request_id for request in scheduler.running] == ["interactive-request"]
+    assert [request.request_id for request in scheduler.waiting] == [
+        "easymagpie-grpo-train-post-refit-step-000000-0-0"
+    ]
+    assert not hasattr(scheduler, "_easymagpie_last_post_refit_purge")