From 502b98f27aa91eaab9d05132a6e2f8aafb86c776 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 13:52:47 +0000 Subject: [PATCH 01/14] Add EasyMagpie RL vLLM Omni core --- .../easymagpie_vllm_omni/__init__.py | 17 + .../easymagpie_vllm_omni/backbone_patches.py | 1006 ++- .../easymagpie_vllm_omni/config.py | 71 +- .../easymagpie_vllm_omni/easymagpie.py | 2038 +++++- .../easymagpie_vllm_omni/local_transformer.py | 415 +- .../easymagpie_vllm_omni/scheduler.py | 253 +- .../easymagpie_vllm_omni/vllm_compat.py | 6224 +++++++++++++++++ .../easymagpie_vllm_omni/tests/conftest.py | 22 +- .../easymagpie_vllm_omni/tests/test_config.py | 28 + .../tests/test_easymagpie_backbone_cache.py | 287 + .../tests/test_import_modes.py | 48 + .../tests/test_local_transformer.py | 162 +- .../tests/test_local_transformer_sampler.py | 115 + .../tests/test_local_transformer_shape.py | 40 + .../tests/test_scheduler.py | 315 + 15 files changed, 10586 insertions(+), 455 deletions(-) create mode 100644 examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py create mode 100644 examples/tts/easymagpie_vllm_omni/tests/test_easymagpie_backbone_cache.py create mode 100644 examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py create mode 100644 examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_sampler.py create mode 100644 examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_shape.py create mode 100644 examples/tts/easymagpie_vllm_omni/tests/test_scheduler.py 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..59976b7002ba 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,830 @@ """ 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_flash_attention_torch_fallback() -> bool: + """Use PyTorch SDPA for EasyMagpie backbone FlashAttention calls. + + The converted EasyMagpie vLLM path currently reaches FlashAttention with + finite QKV tensors at the first Nemotron-H attention layer, then receives + NaN attention outputs. Keep vLLM's scheduler and KV cache contract intact, + but replace only that kernel call for EasyMagpie backbone layers. + """ + + patched = 0 + try: + from vllm.attention.backends.flash_attn import FlashAttentionImpl as LegacyFlashAttentionImpl + + patched += int(_patch_flash_attention_impl(LegacyFlashAttentionImpl, is_v1=False)) + except Exception: + pass + try: + from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl as V1FlashAttentionImpl + + patched += int(_patch_flash_attention_impl(V1FlashAttentionImpl, is_v1=True)) + except Exception: + pass + patched += int(_patch_v1_flash_attn_varlen_func()) + if patched: + logger.info("EasyMagpie FlashAttention torch-SDPA fallback installed on %d implementation(s)", patched) + return bool(patched) + + +def patch_final_rmsnorm_native(backbone) -> bool: + """Run the final Nemotron-H RMSNorm through the fp32 native implementation. + + EasyMagpie feeds the backbone through a persistent ``inputs_embeds`` scratch + buffer. vLLM's fused CUDA add+RMSNorm path is allowed to mutate its inputs + in-place; on this checkpoint it leaves the final hidden state non-finite + even after decoder-layer attention outputs are finite. The native RMSNorm + path performs the residual add and variance in fp32 and returns a fresh + tensor, matching the safer reference math closely enough for inference. + """ + + enabled = os.environ.get("EASYMAGPIE_FINAL_RMSNORM_NATIVE", "1").lower() + if enabled in {"0", "false", "no", "off"}: + return False + + norm_f = getattr(backbone, "norm_f", None) + if norm_f is None or not hasattr(norm_f, "forward_native"): + return False + if getattr(norm_f, "_easymagpie_native_forward_installed", False): + return False + + def forward(x, residual=None): + return norm_f.forward_native(x, residual) + + norm_f.forward = forward + norm_f._easymagpie_native_forward_installed = True + logger.info("EasyMagpie final RMSNorm native fp32 fallback installed") + return True + + +def patch_layer_rmsnorm_native(backbone) -> int: + """Run Nemotron-H decoder-layer RMSNorms through native fp32 math. + + The layer RMSNorm modules update the long-lived residual stream consumed by + every following hybrid layer. If the fused add+RMSNorm kernel writes a + non-finite residual, the final hidden state can be corrupted even when each + mixer's direct output looks finite. The native path preserves vLLM's return + contract while doing the residual add and variance in fp32. + """ + + enabled = os.environ.get("EASYMAGPIE_LAYER_RMSNORM_NATIVE", "1").lower() + if enabled in {"0", "false", "no", "off"}: + return 0 + + patched = 0 + for layer_idx, layer in enumerate(getattr(backbone, "layers", []) or []): + norm = getattr(layer, "norm", None) + if norm is None or not hasattr(norm, "forward_native"): + continue + if getattr(norm, "_easymagpie_native_forward_installed", False): + continue + + def forward(x, residual=None, *, _norm=norm): + return _norm.forward_native(x, residual) + + norm.forward = forward + norm._easymagpie_native_forward_installed = True + norm._easymagpie_native_forward_layer_idx = layer_idx + patched += 1 + + if patched: + logger.info("EasyMagpie layer RMSNorm native fallback installed on %d module(s)", patched) + return patched + + def patch_mamba_streaming_decode() -> None: """Treat 1-token streaming extends as decodes so FULL decode cudagraphs work. @@ -51,15 +866,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 +886,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 +899,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 +918,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..db33d6e628ff 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/.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\d+)\.mixer\.experts\.(?P\d+)\." + r"(?Pup_proj|down_proj)\.weight$" +) +_QKV_WEIGHT_RE = re.compile( + r"^layers\.(?P\d+)\.mixer\.(?Pq_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,30 @@ 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._debug_combined_input_norm = torch.zeros(max_num_tokens, dtype=torch.float32) + self._debug_combined_input_vector = torch.zeros(max_num_tokens, self.embedding_dim, dtype=dtype) + self._debug_text_emb_norm = torch.zeros(max_num_tokens, dtype=torch.float32) + self._debug_phoneme_emb_norm = torch.zeros(max_num_tokens, dtype=torch.float32) + self._debug_audio_emb_norm = torch.zeros(max_num_tokens, dtype=torch.float32) + self._debug_positions = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_input_ids = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_decode_dispatch_index = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_decode_dispatch_count = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_decode_offset = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_mamba_num_prefills = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_mamba_num_decodes = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_mamba_state_indices = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_mamba_cache_state_indices = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_mamba_exec_state_indices = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_attn_num_actual_tokens = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_attn_max_query_len = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_attn_max_seq_len = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_attn_seq_lens = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_attn_slot_mapping = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_audio_feedback_missing = torch.zeros(max_num_tokens, dtype=torch.long) + self._debug_audio_input_code0 = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_phoneme_input_valid = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._debug_phoneme_input_token0 = 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 +371,13 @@ 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._debug_lt_top_ids = torch.full((max_num_tokens, self.num_codebooks, 5), -1, dtype=torch.long) + self._debug_lt_top_values = torch.zeros(max_num_tokens, self.num_codebooks, 5, dtype=torch.float32) + self._debug_outputs_enabled = False + 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 +393,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 +431,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 +615,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 +649,219 @@ 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 + + @staticmethod + def _metadata_int(metadata: Any, *names: str) -> Optional[int]: + for name in names: + value = getattr(metadata, name, None) + if value is None: + continue + if isinstance(value, torch.Tensor): + if value.numel() == 0: + continue + return int(value.detach().reshape(-1)[0].item()) + try: + return int(value) + except Exception: + continue + return None + + @staticmethod + def _metadata_tensor(metadata: Any, *names: str) -> Optional[torch.Tensor]: + for name in names: + value = getattr(metadata, name, None) + if isinstance(value, torch.Tensor) and value.numel() > 0: + return value.detach().reshape(-1).to(dtype=torch.long) + return None + + @staticmethod + def _metadata_list_tensor(metadata: Any, *names: str, device: torch.device) -> Optional[torch.Tensor]: + tensor = EasyMagpieTTSForConditionalGeneration._metadata_tensor(metadata, *names) + if tensor is not None: + return tensor.to(device=device) + for name in names: + value = getattr(metadata, name, None) + if value is None: + continue + try: + items = list(value) + except TypeError: + continue + if items: + return torch.tensor([int(item) for item in items], dtype=torch.long, device=device) + return None + + def _record_state_indices_debug( + self, + target: torch.Tensor, + state_indices: torch.Tensor, + *, + num_tokens: int, + decode_idx: Optional[torch.Tensor] = None, + num_req: int = 0, + ) -> None: + device = target.device + state_indices = state_indices.detach().reshape(-1).to(device=device, dtype=torch.long) + if decode_idx is not None and num_req > 0 and state_indices.numel() >= num_req: + valid = decode_idx[:num_req].detach().to(device=device, dtype=torch.long).reshape(-1) + target[valid].copy_(state_indices[:num_req]) + else: + n = min(num_tokens, int(state_indices.numel())) + target[:n].copy_(state_indices[:n]) + + def _record_mamba_metadata_debug( + self, + num_tokens: int, + *, + decode_idx: Optional[torch.Tensor] = None, + num_req: int = 0, + mamba_cache_params: Any = None, + ) -> None: + try: + attn_metadata = get_forward_context().attn_metadata + except Exception: + attn_metadata = None + + metas = list(self._iter_metadata(attn_metadata)) + mamba_metas = [ + meta for meta in metas if getattr(meta, "state_indices_tensor", None) is not None + ] + counter_metas = mamba_metas or metas + rows = slice(0, num_tokens) + + for metadata in counter_metas: + num_prefills = self._metadata_int(metadata, "num_prefills", "num_prefill_tokens") + if num_prefills is not None: + self._debug_mamba_num_prefills[rows].fill_(num_prefills) + break + for metadata in counter_metas: + num_decodes = self._metadata_int(metadata, "num_decodes", "num_decode_tokens") + if num_decodes is not None: + self._debug_mamba_num_decodes[rows].fill_(num_decodes) + break + + cache_state_indices = None + for name in ("state_indices_tensor_d", "state_indices_tensor", "state_indices_tensor_p"): + value = getattr(mamba_cache_params, name, None) + if isinstance(value, torch.Tensor) and value.numel() > 0: + cache_state_indices = value + break + if cache_state_indices is not None: + self._record_state_indices_debug( + self._debug_mamba_cache_state_indices, + cache_state_indices, + num_tokens=num_tokens, + decode_idx=decode_idx, + num_req=num_req, + ) + + exec_state_indices = None + for metadata in mamba_metas: + exec_state_indices = self._metadata_tensor(metadata, "state_indices_tensor") + if exec_state_indices is not None: + break + if exec_state_indices is None: + for metadata in metas: + block_table = getattr(metadata, "block_table_tensor", None) + if isinstance(block_table, torch.Tensor) and block_table.numel() > 0: + exec_state_indices = block_table[:, 0].detach().reshape(-1).to(dtype=torch.long) + break + + if exec_state_indices is not None: + self._record_state_indices_debug( + self._debug_mamba_exec_state_indices, + exec_state_indices, + num_tokens=num_tokens, + decode_idx=decode_idx, + num_req=num_req, + ) + self._record_state_indices_debug( + self._debug_mamba_state_indices, + exec_state_indices, + num_tokens=num_tokens, + decode_idx=decode_idx, + num_req=num_req, + ) + elif cache_state_indices is not None: + self._record_state_indices_debug( + self._debug_mamba_state_indices, + cache_state_indices, + num_tokens=num_tokens, + decode_idx=decode_idx, + num_req=num_req, + ) + + def _record_debug_forward_metadata( + self, + *, + input_ids: torch.Tensor, + positions: torch.Tensor, + decode_idx: Optional[torch.Tensor], + num_req: int, + mamba_cache_params: Any, + ) -> None: + if not self._debug_outputs_enabled: + return + + num_tokens = int(input_ids.shape[0]) + rows = slice(0, num_tokens) + device = self._debug_positions.device + self._debug_positions[rows].copy_(positions.detach().to(device=device, dtype=torch.long).reshape(-1)[:num_tokens]) + self._debug_input_ids[rows].copy_(input_ids.detach().to(device=device, dtype=torch.long).reshape(-1)[:num_tokens]) + self._debug_decode_dispatch_count[rows].fill_(int(num_req)) + if decode_idx is None: + self._debug_decode_dispatch_index[rows].copy_(torch.arange(num_tokens, dtype=torch.long, device=device)) + elif num_req > 0: + valid = decode_idx[:num_req].detach().to(device=device, dtype=torch.long).reshape(-1) + self._debug_decode_dispatch_index[valid].copy_(valid) + + try: + attn_metadata = get_forward_context().attn_metadata + except Exception: + attn_metadata = None + metas = list(self._iter_metadata(attn_metadata)) + self._record_mamba_metadata_debug( + num_tokens, + decode_idx=decode_idx, + num_req=num_req, + mamba_cache_params=mamba_cache_params, + ) + for metadata in metas: + num_actual = self._metadata_int(metadata, "num_actual_tokens") + if num_actual is not None: + self._debug_attn_num_actual_tokens[rows].fill_(num_actual) + break + for metadata in metas: + max_query_len = self._metadata_int(metadata, "max_query_len") + if max_query_len is not None: + self._debug_attn_max_query_len[rows].fill_(max_query_len) + break + for metadata in metas: + max_seq_len = self._metadata_int(metadata, "max_seq_len", "max_seq_len_q") + if max_seq_len is not None: + self._debug_attn_max_seq_len[rows].fill_(max_seq_len) + break + for metadata in metas: + seq_lens = self._metadata_list_tensor(metadata, "seq_lens", "seq_lens_tensor", device=device) + if seq_lens is not None: + n = min(num_tokens, int(seq_lens.numel())) + self._debug_attn_seq_lens[:n].copy_(seq_lens[:n]) + break + for metadata in metas: + slot_mapping = self._metadata_tensor(metadata, "slot_mapping") + if slot_mapping is not None: + slot_mapping = slot_mapping.to(device=device) + n = min(num_tokens, int(slot_mapping.numel())) + self._debug_attn_slot_mapping[:n].copy_(slot_mapping[:n]) + break + # ------------------------------------------------------------------ # forward # ------------------------------------------------------------------ @@ -454,7 +892,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 +911,43 @@ def forward( valid = decode_idx[:num_req] self._assemble_decode_embeddings(combined, valid) - hidden_states = self.backbone( + if self._debug_outputs_enabled: + debug_rows = slice(0, num_tokens) + self._debug_combined_input_vector[debug_rows].copy_(combined.detach()) + self._debug_combined_input_norm[debug_rows] = combined.detach().float().norm(dim=-1) + self.code_predictor.debug_collect_logits = bool(self._debug_outputs_enabled) + + 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 + + self._record_debug_forward_metadata( input_ids=input_ids, positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=combined, + decode_idx=decode_idx, + num_req=int(num_req), + 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)) + if self._debug_outputs_enabled: + self._debug_lt_top_ids[:num_tokens].copy_(self.code_predictor._debug_top_ids[:num_tokens]) + self._debug_lt_top_values[:num_tokens].copy_(self.code_predictor._debug_top_values[:num_tokens]) self._flag_audio_eos(codes, slice(0, num_tokens)) if self.has_phoneme: self._predict_phonemes(hidden_states, slice(0, num_tokens)) @@ -484,10 +955,18 @@ 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) + if self._debug_outputs_enabled: + self._debug_lt_top_ids[valid] = self.code_predictor._debug_top_ids[:num_req] + self._debug_lt_top_values[valid] = self.code_predictor._debug_top_values[:num_req] self._flag_audio_eos(codes[:num_req], valid) if self.has_phoneme: self._predict_phonemes(hidden_states, valid) @@ -510,26 +989,45 @@ 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._debug_decode_offset[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 + self._debug_audio_emb_norm[idx] = audio_emb.float().norm(dim=-1) + 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 + self._debug_text_emb_norm[idx] = text_emb.float().norm(dim=-1) + 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 + self._debug_phoneme_emb_norm[idx] = phon_emb.float().norm(dim=-1) + assembled += phon_emb + else: + self._debug_phoneme_emb_norm[idx].zero_() + combined[idx] = assembled @torch.no_grad() def _predict_phonemes(self, hidden_states: torch.Tensor, idx) -> None: @@ -598,16 +1096,378 @@ 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() + if self._debug_outputs_enabled: + combined = self._combined_embeddings.index_select( + 0, + row_indices.to(self._combined_embeddings.device), + ).clone() + hidden_debug = hidden[:num_tokens].clone() + multimodal_outputs.update( + { + "debug_text_mask": self._dec_text_mask.index_select( + 0, + row_indices.to(self._dec_text_mask.device), + ).clone(), + "debug_text_tokens": self._dec_text_tokens.index_select( + 0, + row_indices.to(self._dec_text_tokens.device), + ).clone(), + "debug_audio_valid": self._dec_audio_valid.index_select( + 0, + row_indices.to(self._dec_audio_valid.device), + ).clone(), + "debug_text_emb_norm": self._debug_text_emb_norm.index_select( + 0, + row_indices.to(self._debug_text_emb_norm.device), + ).clone(), + "debug_phoneme_emb_norm": self._debug_phoneme_emb_norm.index_select( + 0, + row_indices.to(self._debug_phoneme_emb_norm.device), + ).clone(), + "debug_audio_emb_norm": self._debug_audio_emb_norm.index_select( + 0, + row_indices.to(self._debug_audio_emb_norm.device), + ).clone(), + "debug_positions": self._debug_positions.index_select( + 0, + row_indices.to(self._debug_positions.device), + ).clone(), + "debug_input_ids": self._debug_input_ids.index_select( + 0, + row_indices.to(self._debug_input_ids.device), + ).clone(), + "debug_decode_dispatch_index": self._debug_decode_dispatch_index.index_select( + 0, + row_indices.to(self._debug_decode_dispatch_index.device), + ).clone(), + "debug_decode_dispatch_count": self._debug_decode_dispatch_count.index_select( + 0, + row_indices.to(self._debug_decode_dispatch_count.device), + ).clone(), + "debug_combined_input_norm": self._debug_combined_input_norm.index_select( + 0, + row_indices.to(self._debug_combined_input_norm.device), + ).clone(), + "debug_combined_input_vector": self._debug_combined_input_vector.index_select( + 0, + row_indices.to(self._debug_combined_input_vector.device), + ).clone(), + "debug_combined_norm": combined.float().norm(dim=-1), + "debug_combined_vector": combined, + "debug_hidden_norm": hidden_debug.float().norm(dim=-1), + "debug_hidden_vector": hidden_debug, + "debug_audio_input_code0": self._debug_audio_input_code0.index_select( + 0, + row_indices.to(self._debug_audio_input_code0.device), + ).clone(), + "debug_audio_output_code0": audio_codes[:, 0].clone(), + "debug_decode_offset": self._debug_decode_offset.index_select( + 0, + row_indices.to(self._debug_decode_offset.device), + ).clone(), + "debug_audio_feedback_missing": self._debug_audio_feedback_missing.index_select( + 0, + row_indices.to(self._debug_audio_feedback_missing.device), + ).clone(), + "debug_phoneme_input_valid": self._debug_phoneme_input_valid.index_select( + 0, + row_indices.to(self._debug_phoneme_input_valid.device), + ).clone(), + "debug_phoneme_input_token0": self._debug_phoneme_input_token0.index_select( + 0, + row_indices.to(self._debug_phoneme_input_token0.device), + ).clone(), + "debug_mamba_num_prefills": self._debug_mamba_num_prefills.index_select( + 0, + row_indices.to(self._debug_mamba_num_prefills.device), + ).clone(), + "debug_mamba_num_decodes": self._debug_mamba_num_decodes.index_select( + 0, + row_indices.to(self._debug_mamba_num_decodes.device), + ).clone(), + "debug_mamba_state_indices": self._debug_mamba_state_indices.index_select( + 0, + row_indices.to(self._debug_mamba_state_indices.device), + ).clone(), + "debug_mamba_cache_state_indices": self._debug_mamba_cache_state_indices.index_select( + 0, + row_indices.to(self._debug_mamba_cache_state_indices.device), + ).clone(), + "debug_mamba_exec_state_indices": self._debug_mamba_exec_state_indices.index_select( + 0, + row_indices.to(self._debug_mamba_exec_state_indices.device), + ).clone(), + "debug_attn_num_actual_tokens": self._debug_attn_num_actual_tokens.index_select( + 0, + row_indices.to(self._debug_attn_num_actual_tokens.device), + ).clone(), + "debug_attn_max_query_len": self._debug_attn_max_query_len.index_select( + 0, + row_indices.to(self._debug_attn_max_query_len.device), + ).clone(), + "debug_attn_max_seq_len": self._debug_attn_max_seq_len.index_select( + 0, + row_indices.to(self._debug_attn_max_seq_len.device), + ).clone(), + "debug_attn_seq_lens": self._debug_attn_seq_lens.index_select( + 0, + row_indices.to(self._debug_attn_seq_lens.device), + ).clone(), + "debug_attn_slot_mapping": self._debug_attn_slot_mapping.index_select( + 0, + row_indices.to(self._debug_attn_slot_mapping.device), + ).clone(), + "debug_lt_top_ids": self._debug_lt_top_ids.index_select( + 0, + row_indices.to(self._debug_lt_top_ids.device), + ).clone(), + "debug_lt_top_values": self._debug_lt_top_values.index_select( + 0, + row_indices.to(self._debug_lt_top_values.device), + ).clone(), + } + ) + if self.has_phoneme: + phoneme_indices = row_indices.to(self._dec_phoneme_valid.device) + multimodal_outputs.update( + { + "debug_phoneme_valid": self._dec_phoneme_valid.index_select(0, phoneme_indices).clone(), + "debug_phoneme_input_token0": self._dec_phoneme_tokens.index_select( + 0, + row_indices.to(self._dec_phoneme_tokens.device), + )[:, 0].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", + "_debug_lt_top_ids", + "_debug_lt_top_values", + "_token_stop", + "_sample_stop", + "_debug_combined_input_norm", + "_debug_combined_input_vector", + "_debug_text_emb_norm", + "_debug_phoneme_emb_norm", + "_debug_audio_emb_norm", + "_debug_combined_pre_norm", + "_debug_hidden_norm", + "_debug_backbone_last_layer_norm", + "_debug_backbone_last_residual_norm", + "_debug_final_norm_input_norm", + "_debug_final_norm_residual_norm", + "_debug_final_norm_output_norm", + "_debug_attn_qkv_norm", + "_debug_attn_core_norm", + "_debug_attn_output_norm", + "_debug_audio_feedback_missing", + ) + minus_one_names = ( + "_debug_positions", + "_debug_input_ids", + "_debug_decode_dispatch_index", + "_debug_decode_dispatch_count", + "_debug_backbone_first_bad_layer", + "_debug_backbone_first_bad_residual_layer", + "_debug_mamba_num_prefills", + "_debug_mamba_num_decodes", + "_debug_mamba_state_indices", + "_debug_mamba_cache_state_indices", + "_debug_mamba_exec_state_indices", + "_debug_attn_num_actual_tokens", + "_debug_attn_max_query_len", + "_debug_attn_max_seq_len", + "_debug_attn_seq_lens", + "_debug_attn_slot_mapping", + "_debug_decode_offset", + "_debug_audio_input_code0", + "_debug_phoneme_input_valid", + "_debug_phoneme_input_token0", + ) + 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 +1513,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( @@ -669,24 +1531,26 @@ def _preprocess_prefill( # ``additional_information`` and applied to the code predictor here (once, # at prefill — they are scalars that persist across decode steps). self._maybe_set_lt_sampling_params(info_dict) + self._debug_outputs_enabled = bool(info_dict.get("debug_outputs", False)) prefill_embeds = self._build_prefill_embeds(device, info_dict) 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 +1565,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 +1586,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 ``/speaker_embeddings/.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 +1636,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 +1649,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 +1703,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 +1717,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 +1726,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/.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 +1735,47 @@ 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) + debug_decode_offset = getattr(self, "_debug_decode_offset", None) + if isinstance(debug_decode_offset, torch.Tensor): + debug_decode_offset[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 +1783,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,9 +1816,20 @@ 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 + debug_phoneme_valid = getattr(self, "_debug_phoneme_input_valid", None) + if isinstance(debug_phoneme_valid, torch.Tensor): + debug_phoneme_valid[start] = int(self._dec_phoneme_valid[start].item()) + debug_phoneme_token = getattr(self, "_debug_phoneme_input_token0", None) + if isinstance(debug_phoneme_token, torch.Tensor): + debug_phoneme_token[start] = ( + int(self._dec_phoneme_tokens[start, 0].item()) + if int(self._dec_phoneme_valid[start].item()) and self._dec_phoneme_tokens.ndim == 2 + else -1 + ) # ── Audio channel ── opens at decode step == ``speech_delay`` (seeded with # audio BOS), then feeds back the previous frame's codes. For the leading @@ -1038,20 +1838,33 @@ 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) self._dec_audio_valid[start] = 1 else: last_codes = info_dict.get("last_audio_codes") + missing_feedback = True if isinstance(last_codes, torch.Tensor) and last_codes.numel() > 0: c = last_codes.to(device=device, dtype=torch.long).reshape(-1)[: self.num_codebooks] self._dec_audio_codes[start, : c.shape[0]].copy_(c) self._dec_audio_valid[start] = 1 + missing_feedback = False else: # Fallback (should not happen once audio has started): seed BOS. self._dec_audio_codes[start].fill_(self.arch.audio_bos_id) self._dec_audio_valid[start] = 1 + debug_missing = getattr(self, "_debug_audio_feedback_missing", None) + if isinstance(debug_missing, torch.Tensor): + debug_missing[start] = int(missing_feedback) + debug_audio_input = getattr(self, "_debug_audio_input_code0", None) + if isinstance(debug_audio_input, torch.Tensor): + debug_audio_input[start] = ( + int(self._dec_audio_codes[start, 0].item()) + if int(self._dec_audio_valid[start].item()) + else -1 + ) inputs_embeds_out = torch.zeros((1, self.embedding_dim), device=device, dtype=self._combined_embeddings.dtype) return input_ids, inputs_embeds_out, info_update @@ -1060,22 +1873,68 @@ 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._debug_outputs_enabled = False + 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", + "_debug_text_emb_norm", + "_debug_phoneme_emb_norm", + "_debug_audio_emb_norm", + "_dec_phoneme_tokens", + "_dec_phoneme_valid", + ): + value = getattr(self, name, None) + if isinstance(value, torch.Tensor): + value.zero_() + # 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 +1949,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 +1960,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 +2372,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 +2501,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..1339376b1121 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,52 @@ 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.debug_collect_logits: bool = False + self.debug_logits_top_k: int = 5 + self._debug_top_ids = torch.full((max_num_tokens, self.num_codebooks, 5), -1, dtype=torch.long) + self._debug_top_values = torch.zeros(max_num_tokens, self.num_codebooks, 5, 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 +443,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 +458,64 @@ 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_() + if self.debug_collect_logits: + self._debug_top_ids[:num_tokens].fill_(-1) + self._debug_top_values[:num_tokens].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) + if self.debug_collect_logits: + debug_logits = logits + if forbidden is not None: + debug_logits = debug_logits.masked_fill(forbidden, float("-inf")) + debug_k = min(int(self.debug_logits_top_k or 5), int(self._debug_top_ids.shape[-1]), int(debug_logits.shape[-1])) + vals, ids = torch.topk(debug_logits.detach().float(), k=debug_k, dim=-1) + self._debug_top_ids[:num_tokens, k, :debug_k].copy_(ids) + self._debug_top_values[:num_tokens, k, :debug_k].copy_(vals) + 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..c4d753379a29 --- /dev/null +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -0,0 +1,6224 @@ +# 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. +"""Compatibility shims for running vLLM-Omni against older vLLM builds.""" + +from __future__ import annotations + +import builtins +import contextlib +import dataclasses +import importlib +import inspect +import logging +import os +import queue as stdlib_queue +import random +import sys +import types +from collections import defaultdict +from typing import Any, Literal + +_INSTALLED = False +_IMPORT_PATCH_FLAG = "_easymagpie_vllm_compat_import_patch_installed" +logger = logging.getLogger(__name__) + +try: + import torch as _TORCH_MODULE +except Exception: + _TORCH_MODULE = None + +try: + from vllm.attention.backends.flash_attn import FlashAttentionMetadata as _LEGACY_FLASH_ATTENTION_METADATA_CLS +except Exception: + _LEGACY_FLASH_ATTENTION_METADATA_CLS = None + +try: + from vllm.attention.backends.xformers import XFormersMetadata as _XFORMERS_METADATA_CLS +except Exception: + _XFORMERS_METADATA_CLS = None + +try: + from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata as _V1_FLASH_ATTENTION_METADATA_CLS +except Exception: + _V1_FLASH_ATTENTION_METADATA_CLS = None + +try: + from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata as _V1_TRITON_ATTENTION_METADATA_CLS +except Exception: + _V1_TRITON_ATTENTION_METADATA_CLS = None + +try: + from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata as _V1_MAMBA2_ATTENTION_METADATA_CLS +except Exception: + _V1_MAMBA2_ATTENTION_METADATA_CLS = None + +_SUPPORTED_KWARGS_CACHE: dict[int, set[str] | None] = {} + + +def _as_v1_kv_cache_config(cache_config: Any) -> Any: + if cache_config is None or hasattr(cache_config, "kv_cache_groups"): + return cache_config + cached = getattr(cache_config, "_easymagpie_v1_kv_cache_config", None) + if cached is None: + try: + from vllm.v1.kv_cache_interface import KVCacheConfig + + cached = KVCacheConfig(num_blocks=1, kv_cache_tensors=[], kv_cache_groups=[]) + except Exception: + cached = types.SimpleNamespace(num_blocks=1, kv_cache_tensors=[], kv_cache_groups=[]) + try: + cache_config._easymagpie_v1_kv_cache_config = cached + except Exception: + pass + return cached + + +def _buffer_slice_to_int_list(buffer: Any, length: int) -> list[int] | None: + if buffer is None: + return None + candidates = ( + getattr(buffer, "cpu", None), + getattr(buffer, "np", None), + getattr(buffer, "gpu", None), + buffer, + ) + for candidate in candidates: + if candidate is None: + continue + try: + values = candidate[:length] + if hasattr(values, "detach"): + values = values.detach().cpu().tolist() + elif hasattr(values, "tolist"): + values = values.tolist() + return [int(v) for v in values] + except Exception: + continue + return None + + +def _balanced_profile_lengths(num_tokens: int, num_reqs: int) -> list[int]: + num_reqs = max(1, int(num_reqs)) + base = max(1, int(num_tokens) // num_reqs) + remainder = max(0, int(num_tokens) - base * num_reqs) + lengths = [base] * num_reqs + lengths[-1] += remainder + return lengths + + +def _cache_supported_kwargs(callable_obj: Any) -> set[str] | None: + cache_key = id(callable_obj) + if cache_key in _SUPPORTED_KWARGS_CACHE: + return _SUPPORTED_KWARGS_CACHE[cache_key] + try: + signature = inspect.signature(callable_obj) + except Exception: + _SUPPORTED_KWARGS_CACHE[cache_key] = None + return None + if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values()): + _SUPPORTED_KWARGS_CACHE[cache_key] = None + return None + supported = { + name + for name, parameter in signature.parameters.items() + if parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) + } + _SUPPORTED_KWARGS_CACHE[cache_key] = supported + return supported + + +def _prime_supported_kwargs_cache() -> None: + for callable_obj in ( + _LEGACY_FLASH_ATTENTION_METADATA_CLS, + _XFORMERS_METADATA_CLS, + _V1_FLASH_ATTENTION_METADATA_CLS, + _V1_TRITON_ATTENTION_METADATA_CLS, + _V1_MAMBA2_ATTENTION_METADATA_CLS, + ): + if callable_obj is not None: + _cache_supported_kwargs(callable_obj) + + +def _call_with_supported_kwargs(callable_obj: Any, kwargs: dict[str, Any]) -> Any: + supported = _cache_supported_kwargs(callable_obj) + if supported is None: + return callable_obj(**kwargs) + return callable_obj(**{key: value for key, value in kwargs.items() if key in supported}) + + +_prime_supported_kwargs_cache() + + +def _max_int_sequence(values: Any) -> int: + max_value = 0 + for value in values: + int_value = int(value) + if int_value > max_value: + max_value = int_value + return max_value + + +def _is_torch_dynamo_compiling() -> bool: + try: + torch_module = _TORCH_MODULE + dynamo = getattr(torch_module, "_dynamo", None) + is_compiling = getattr(dynamo, "is_compiling", None) + return bool(is_compiling is not None and is_compiling()) + except Exception: + return False + + +def _attach_profile_attention_metadata_attrs( + metadata: Any, + *, + num_tokens: int, + num_reqs: int, + query_lens: list[int], + seq_lens: list[int], + seq_lens_tensor: Any, + seq_start_loc: Any, + query_start_loc: Any = None, + context_lens_tensor: Any = None, + block_tables: Any = None, + slot_mapping: Any = None, + max_query_len: int | None = None, +) -> Any: + if query_start_loc is None: + query_start_loc = seq_start_loc + resolved_max_query_len = max(int(max_query_len or 0), _max_int_sequence(query_lens)) + for name, value in ( + ("num_prefills", int(num_reqs)), + ("num_prefill_tokens", int(num_tokens)), + ("num_decodes", 0), + ("num_decode_tokens", 0), + ("query_lens", list(query_lens)), + ("seq_lens", list(seq_lens)), + ("seq_lens_tensor", seq_lens_tensor), + ("seq_start_loc", seq_start_loc), + ("query_start_loc", query_start_loc), + ("max_query_len", resolved_max_query_len), + ("max_prefill_seq_len", _max_int_sequence(seq_lens)), + ("max_decode_query_len", 0), + ("max_decode_seq_len", 0), + ("context_lens_tensor", context_lens_tensor), + ("block_tables", block_tables), + ("slot_mapping", slot_mapping), + ("prefill_metadata", metadata), + ("decode_metadata", None), + ): + try: + setattr(metadata, name, value) + except Exception: + pass + return metadata + + +def _is_mamba2_attention_metadata(value: Any) -> bool: + common = all( + hasattr(value, attr) + for attr in ( + "num_prefills", + "num_prefill_tokens", + "num_decodes", + "num_decode_tokens", + "seq_lens", + "chunk_size", + ) + ) + if not common: + return False + legacy_state = all(hasattr(value, attr) for attr in ("query_start_loc", "state_indices_tensor")) + split_state = all( + hasattr(value, attr) + for attr in ("query_start_loc_p", "state_indices_tensor_p", "query_start_loc_d", "state_indices_tensor_d") + ) + return legacy_state or split_state + + +def _first_not_none(*values: Any) -> Any: + for value in values: + if value is not None: + return value + return None + + +def _legacy_mamba2_attention_metadata_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + return { + "num_prefills": kwargs.get("num_prefills", 0), + "num_prefill_tokens": kwargs.get("num_prefill_tokens", 0), + "num_decodes": kwargs.get("num_decodes", 0), + "num_decode_tokens": kwargs.get("num_decode_tokens", 0), + "query_start_loc": _first_not_none( + kwargs.get("query_start_loc"), + kwargs.get("legacy_query_start_loc"), + kwargs.get("query_start_loc_p"), + kwargs.get("query_start_loc_d"), + ), + "seq_lens": kwargs["seq_lens"], + "prep_initial_states": kwargs.get("prep_initial_states", False), + "chunk_size": kwargs.get("chunk_size", 0), + "has_initial_states_p": kwargs.get("has_initial_states_p"), + "seq_idx_p": kwargs.get("seq_idx_p"), + "chunk_indices_p": kwargs.get("chunk_indices_p"), + "chunk_offsets_p": kwargs.get("chunk_offsets_p"), + "state_indices_tensor": _first_not_none( + kwargs.get("state_indices_tensor"), + kwargs.get("legacy_state_indices_tensor"), + kwargs.get("state_indices_tensor_p"), + kwargs.get("state_indices_tensor_d"), + ), + "nums_dict": kwargs.get("nums_dict"), + "cu_seqlen": kwargs.get("cu_seqlen"), + "batch_ptr": kwargs.get("batch_ptr"), + "token_chunk_offset_ptr": kwargs.get("token_chunk_offset_ptr"), + } + + +def _filter_cached_supported_kwargs(metadata_cls: Any, kwargs: dict[str, Any]) -> dict[str, Any]: + supported = _cache_supported_kwargs(metadata_cls) + if supported is None: + return dict(kwargs) + return {key: value for key, value in kwargs.items() if key in supported} + + +def _new_mamba2_attention_metadata(metadata_cls: Any, kwargs: dict[str, Any]) -> Any: + constructor_kwargs = { + key: value + for key, value in kwargs.items() + if key not in {"legacy_query_start_loc", "legacy_state_indices_tensor"} + } + supported = _cache_supported_kwargs(metadata_cls) + if supported is not None and "query_start_loc_p" not in supported and "query_start_loc" in supported: + return metadata_cls(**_filter_cached_supported_kwargs(metadata_cls, _legacy_mamba2_attention_metadata_kwargs(kwargs))) + try: + return metadata_cls(**_filter_cached_supported_kwargs(metadata_cls, constructor_kwargs)) + except TypeError as exc: + message = str(exc) + if "unexpected keyword argument" in message and "query_start_loc_p" in kwargs: + legacy_kwargs = _legacy_mamba2_attention_metadata_kwargs(kwargs) + try: + return metadata_cls(**_filter_cached_supported_kwargs(metadata_cls, legacy_kwargs)) + except TypeError as legacy_exc: + if "num_reqs" not in kwargs or "num_reqs" not in str(legacy_exc): + raise + legacy_kwargs["num_reqs"] = kwargs["num_reqs"] + return metadata_cls(**_filter_cached_supported_kwargs(metadata_cls, legacy_kwargs)) + if "num_reqs" not in kwargs or "num_reqs" not in message: + raise + fallback_kwargs = _filter_cached_supported_kwargs(metadata_cls, constructor_kwargs) + fallback_kwargs.pop("num_reqs", None) + return metadata_cls(**fallback_kwargs) + + +def _resize_mamba_state_indices(state_indices: Any, needed_state_indices: int, device: Any) -> Any: + import torch + + if getattr(state_indices, "numel", lambda: 0)() > needed_state_indices: + return state_indices.to(device=device, dtype=torch.int32)[:needed_state_indices] + + current_numel = int(getattr(state_indices, "numel", lambda: 0)()) + if current_numel > 0: + base = state_indices.to(device=device, dtype=torch.int32).reshape(-1) + repeats = (needed_state_indices + current_numel - 1) // current_numel + return base.repeat(repeats)[:needed_state_indices] + + return torch.arange(needed_state_indices, dtype=torch.int32, device=device) + + +def _replace_mamba2_attention_metadata_fields(value: Any, replacements: dict[str, Any]) -> Any: + if not replacements: + return value + if _is_torch_dynamo_compiling(): + return value + try: + if dataclasses.is_dataclass(value): + supported_fields = getattr(type(value), "__dataclass_fields__", {}) + supported_replacements = { + key: replacement + for key, replacement in replacements.items() + if not supported_fields or key in supported_fields + } + if supported_replacements: + return dataclasses.replace(value, **supported_replacements) + except Exception: + pass + fields = dict(getattr(value, "__dict__", {})) + fields.update(replacements) + return types.SimpleNamespace(**fields) + + +def _balanced_tensor_lengths(total: int, count: int, device: Any) -> Any: + import torch + + count = max(0, int(count)) + if count <= 0: + return torch.empty(0, dtype=torch.int32, device=device) + base, extra = divmod(max(0, int(total)), count) + lengths = torch.full((count,), base, dtype=torch.int32, device=device) + if extra: + lengths[:extra] += 1 + return lengths + + +def _infer_mamba2_prefill_query_start_loc(value: Any, device: Any) -> Any: + import torch + + num_prefills = int(getattr(value, "num_prefills", 0) or 0) + num_prefill_tokens = int(getattr(value, "num_prefill_tokens", 0) or 0) + if num_prefills <= 0: + return torch.zeros(1, dtype=torch.int32, device=device) + + candidates = [] + query_start_loc_p = getattr(value, "query_start_loc_p", None) + if query_start_loc_p is not None and hasattr(query_start_loc_p, "numel"): + candidates.append(query_start_loc_p) + + query_start_loc = getattr(value, "query_start_loc", None) + if query_start_loc is not None and hasattr(query_start_loc, "numel"): + query_start_loc = query_start_loc.to(device=device, dtype=torch.int64) + if int(query_start_loc.numel()) >= num_prefills + 1: + num_decode_tokens = int(getattr(value, "num_decode_tokens", 0) or 0) + candidates.append(query_start_loc[-num_prefills - 1 :] - num_decode_tokens) + candidates.append(query_start_loc[: num_prefills + 1]) + + for candidate in candidates: + try: + candidate = candidate.to(device=device, dtype=torch.int64) + if int(candidate.numel()) != num_prefills + 1: + continue + if _is_torch_dynamo_compiling(): + return candidate.to(dtype=torch.int32) + if int(candidate[0].item()) != 0 or int(candidate[-1].item()) != num_prefill_tokens: + continue + lengths = candidate[1:] - candidate[:-1] + if bool((lengths >= 0).all().item()): + return candidate.to(dtype=torch.int32) + except Exception: + continue + + lengths = _balanced_tensor_lengths(num_prefill_tokens, num_prefills, device) + query_start_loc_p = torch.zeros(num_prefills + 1, dtype=torch.int32, device=device) + if num_prefills > 0: + query_start_loc_p[1:] = torch.cumsum(lengths, dim=0) + return query_start_loc_p + + +def _uses_mamba2_varlen_chunk_metadata(value: Any) -> bool: + return hasattr(value, "cu_chunk_seqlen_p") or hasattr(value, "last_chunk_indices_p") + + +def _compute_mamba2_varlen_chunk_metadata(query_start_loc: Any, chunk_size: int) -> tuple[Any, Any, Any]: + import torch + + try: + module = importlib.import_module("vllm.v1.attention.backends.mamba2_attn") + compute = getattr(module, "compute_varlen_chunk_metadata", None) + if compute is not None: + return compute(query_start_loc, chunk_size) + except Exception: + pass + + query_start_loc = query_start_loc.to(dtype=torch.int64) + device = query_start_loc.device + if int(query_start_loc.numel()) == 0: + zero = torch.tensor([0], dtype=torch.int32, device=device) + return zero, torch.empty(0, dtype=torch.int32, device=device), torch.empty(0, dtype=torch.int32, device=device) + + chunk_size = int(chunk_size or 0) + if chunk_size <= 0: + chunk_size = max(1, int(query_start_loc[-1].item())) + + starts = query_start_loc[:-1].detach().cpu().tolist() + ends = query_start_loc[1:].detach().cpu().tolist() + chunk_lens: list[int] = [] + seq_idx_chunks: list[int] = [] + last_chunk_indices: list[int] = [-1] * len(starts) + + for seq_idx, (start, end) in enumerate(zip(starts, ends)): + pos = int(start) + end = int(end) + while pos < end: + room = chunk_size - (pos % chunk_size) + take = min(room, end - pos) + if take <= 0: + break + chunk_lens.append(int(take)) + seq_idx_chunks.append(seq_idx) + last_chunk_indices[seq_idx] = len(chunk_lens) - 1 + pos += take + + boundaries = [0] + for chunk_len in chunk_lens: + boundaries.append(boundaries[-1] + chunk_len) + cu_chunk_seqlens = torch.tensor(boundaries, dtype=torch.int32, device=device) + last_chunk_indices_t = torch.tensor(last_chunk_indices, dtype=torch.int32, device=device) + seq_idx_chunks_t = torch.tensor(seq_idx_chunks, dtype=torch.int32, device=device) + return cu_chunk_seqlens, last_chunk_indices_t, seq_idx_chunks_t + + +def _repair_mamba2_attention_metadata_prefill_seq_idx(value: Any) -> Any: + if not _is_mamba2_attention_metadata(value): + return value + try: + import torch + except Exception: + return value + + num_prefills = int(getattr(value, "num_prefills", 0) or 0) + num_prefill_tokens = int(getattr(value, "num_prefill_tokens", 0) or 0) + if num_prefills <= 0 or num_prefill_tokens <= 0: + return value + if _is_torch_dynamo_compiling(): + return value + + seq_idx_p = getattr(value, "seq_idx_p", None) + query_start_loc = getattr(value, "query_start_loc", None) + query_start_loc_p = getattr(value, "query_start_loc_p", None) + if seq_idx_p is None and not hasattr(query_start_loc, "numel") and not hasattr(query_start_loc_p, "numel"): + return value + + device = getattr(seq_idx_p, "device", None) + if device is None: + device = getattr(query_start_loc, "device", None) + if device is None: + device = getattr(query_start_loc_p, "device", None) + if device is None: + device = getattr(getattr(value, "state_indices_tensor", None), "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + query_start_loc_p = _infer_mamba2_prefill_query_start_loc(value, device) + lengths = (query_start_loc_p[1:] - query_start_loc_p[:-1]).to(dtype=torch.int32) + if int(lengths.numel()) != num_prefills or int(lengths.sum().item()) != num_prefill_tokens: + lengths = _balanced_tensor_lengths(num_prefill_tokens, num_prefills, device) + query_start_loc_p = torch.zeros(num_prefills + 1, dtype=torch.int32, device=device) + query_start_loc_p[1:] = torch.cumsum(lengths, dim=0) + + if _uses_mamba2_varlen_chunk_metadata(value): + cu_chunk_seqlens, last_chunk_indices, fixed_seq_idx_p = _compute_mamba2_varlen_chunk_metadata( + query_start_loc_p, + int(getattr(value, "chunk_size", 0) or 0), + ) + expected_shape = (int(cu_chunk_seqlens.numel()) - 1,) + current_cu_chunk_seqlens = getattr(value, "cu_chunk_seqlen_p", None) + current_last_chunk_indices = getattr(value, "last_chunk_indices_p", None) + if ( + seq_idx_p is not None + and hasattr(seq_idx_p, "shape") + and tuple(seq_idx_p.shape) == expected_shape + and current_cu_chunk_seqlens is not None + and hasattr(current_cu_chunk_seqlens, "shape") + and tuple(current_cu_chunk_seqlens.shape) == tuple(cu_chunk_seqlens.shape) + and torch.equal(current_cu_chunk_seqlens.to(device=device, dtype=torch.int32), cu_chunk_seqlens) + and current_last_chunk_indices is not None + and hasattr(current_last_chunk_indices, "shape") + and tuple(current_last_chunk_indices.shape) == tuple(last_chunk_indices.shape) + and torch.equal(current_last_chunk_indices.to(device=device, dtype=torch.int32), last_chunk_indices) + and torch.equal(seq_idx_p.to(device=device, dtype=torch.int32), fixed_seq_idx_p) + ): + return value + return _replace_mamba2_attention_metadata_fields( + value, + { + "seq_idx_p": fixed_seq_idx_p, + "cu_chunk_seqlen_p": cu_chunk_seqlens, + "last_chunk_indices_p": last_chunk_indices, + }, + ) + + expected_shape = (1, num_prefill_tokens) + if ( + seq_idx_p is not None + and hasattr(seq_idx_p, "shape") + and tuple(seq_idx_p.shape) == expected_shape + ): + return value + + fixed_seq_idx_p = torch.repeat_interleave( + torch.arange(num_prefills, dtype=torch.int32, device=device), + lengths, + output_size=num_prefill_tokens, + ).unsqueeze(0) + + replacements: dict[str, Any] = {"seq_idx_p": fixed_seq_idx_p} + if bool(getattr(value, "prep_initial_states", False)): + try: + module = importlib.import_module("vllm.v1.attention.backends.mamba2_attn") + converter = getattr(module, "_query_start_loc_to_chunk_indices_offsets") + chunk_indices, chunk_offsets = converter( + query_start_loc_p, + int(getattr(value, "chunk_size", 0) or 0), + num_prefill_tokens, + ) + replacements["chunk_indices_p"] = chunk_indices + replacements["chunk_offsets_p"] = chunk_offsets + except Exception: + pass + return _replace_mamba2_attention_metadata_fields(value, replacements) + + +def _compact_legacy_mamba_seq_idx(seq_idx: Any, device: Any) -> Any: + import torch + + seq_idx_flat = seq_idx.detach().reshape(-1).to(device="cpu", dtype=torch.int64) + valid = seq_idx_flat >= 0 + if not bool(valid.all().item()): + seq_idx_flat = seq_idx_flat[valid] + if seq_idx_flat.numel() == 0: + return seq_idx_flat.to(device=device, dtype=torch.int32) + + _, compact = torch.unique(seq_idx_flat, sorted=True, return_inverse=True) + return compact.to(device=device, dtype=torch.int32) + + +def _mamba2_attention_metadata_from_mapping(value: dict[str, Any]) -> Any | None: + required_fields = ( + "num_prefills", + "num_prefill_tokens", + "num_decodes", + "num_decode_tokens", + "query_start_loc", + "seq_lens", + "prep_initial_states", + "chunk_size", + "has_initial_states_p", + "seq_idx_p", + "chunk_indices_p", + "chunk_offsets_p", + "state_indices_tensor", + ) + if not all(field in value for field in required_fields): + return None + try: + import torch + from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata + + kwargs = {field: value[field] for field in required_fields if field in value} + needed_state_indices = int(value.get("num_prefills") or 0) + int(value.get("num_decode_tokens") or 0) + kwargs["num_reqs"] = int(value.get("num_reqs") or needed_state_indices or 1) + state_indices = value.get("state_indices_tensor") + if (state_indices is None or hasattr(state_indices, "numel")) and ( + getattr(state_indices, "numel", lambda: 0)() != needed_state_indices + ): + device = getattr(state_indices, "device", None) + if device is None: + device = getattr(value.get("query_start_loc"), "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + state_indices = _resize_mamba_state_indices(state_indices, needed_state_indices, device) + kwargs["state_indices_tensor"] = state_indices + return _new_mamba2_attention_metadata(Mamba2AttentionMetadata, kwargs) + except Exception: + return types.SimpleNamespace(**{field: value.get(field) for field in required_fields}) + + +def _repair_mamba2_attention_metadata_state_indices(value: Any) -> Any: + if not _is_mamba2_attention_metadata(value): + return value + try: + import torch + except Exception: + return value + if _is_torch_dynamo_compiling(): + return value + value = _repair_mamba2_attention_metadata_prefill_seq_idx(value) + if hasattr(value, "state_indices_tensor_p") or hasattr(value, "state_indices_tensor_d"): + device = None + state_indices_p = getattr(value, "state_indices_tensor_p", None) + state_indices_d = getattr(value, "state_indices_tensor_d", None) + for state_indices in (state_indices_p, state_indices_d, getattr(value, "query_start_loc_p", None), getattr(value, "query_start_loc_d", None)): + device = getattr(state_indices, "device", None) + if device is not None: + break + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + replacements: dict[str, Any] = {} + needed_prefill_indices = int(getattr(value, "num_prefills", 0) or 0) + if needed_prefill_indices > 0 and ( + state_indices_p is None + or not hasattr(state_indices_p, "numel") + or int(state_indices_p.numel()) != needed_prefill_indices + ): + replacements["state_indices_tensor_p"] = _resize_mamba_state_indices( + state_indices_p, + needed_prefill_indices, + device, + ) + needed_decode_indices = int(getattr(value, "num_decode_tokens", 0) or 0) + if needed_decode_indices > 0 and ( + state_indices_d is None + or not hasattr(state_indices_d, "numel") + or int(state_indices_d.numel()) != needed_decode_indices + ): + replacements["state_indices_tensor_d"] = _resize_mamba_state_indices( + state_indices_d, + needed_decode_indices, + device, + ) + if not replacements: + return value + try: + if dataclasses.is_dataclass(value): + return dataclasses.replace(value, **replacements) + except Exception: + pass + fields = dict(getattr(value, "__dict__", {})) + fields.update(replacements) + return types.SimpleNamespace(**fields) + + needed_state_indices = int(getattr(value, "num_prefills", 0) or 0) + int( + getattr(value, "num_decode_tokens", 0) or 0 + ) + state_indices = getattr(value, "state_indices_tensor", None) + if state_indices is not None and not hasattr(state_indices, "numel"): + return value + if getattr(state_indices, "numel", lambda: 0)() == needed_state_indices: + return value + + device = getattr(state_indices, "device", None) + if device is None: + device = getattr(getattr(value, "query_start_loc", None), "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + fixed_state_indices = _resize_mamba_state_indices(state_indices, needed_state_indices, device) + + try: + if dataclasses.is_dataclass(value): + return dataclasses.replace(value, state_indices_tensor=fixed_state_indices) + except Exception: + pass + fields = dict(getattr(value, "__dict__", {})) + fields["state_indices_tensor"] = fixed_state_indices + return types.SimpleNamespace(**fields) + + +def _repair_mamba_cache_params_state_indices(cache_params: Any, attn_metadata: Any) -> Any: + if cache_params is None or not _is_mamba2_attention_metadata(attn_metadata): + return cache_params + try: + import torch + except Exception: + return cache_params + needed_state_indices = int(getattr(attn_metadata, "num_prefills", 0) or 0) + int( + getattr(attn_metadata, "num_decode_tokens", 0) or 0 + ) + state_indices = getattr(cache_params, "state_indices_tensor", None) + if state_indices is not None and not hasattr(state_indices, "numel"): + return cache_params + + device = getattr(state_indices, "device", None) + if device is None: + device = getattr(getattr(attn_metadata, "query_start_loc", None), "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + metadata_state_indices = getattr(attn_metadata, "state_indices_tensor", None) + if ( + metadata_state_indices is None + or not hasattr(metadata_state_indices, "numel") + or int(metadata_state_indices.numel()) != needed_state_indices + ): + split_indices: list[Any] = [] + prefill_indices = getattr(attn_metadata, "state_indices_tensor_p", None) + decode_indices = getattr(attn_metadata, "state_indices_tensor_d", None) + if prefill_indices is not None and hasattr(prefill_indices, "numel") and int(prefill_indices.numel()) > 0: + split_indices.append(prefill_indices) + if decode_indices is not None and hasattr(decode_indices, "numel") and int(decode_indices.numel()) > 0: + split_indices.append(decode_indices) + if split_indices: + try: + metadata_state_indices = torch.cat( + [item.detach().reshape(-1).to(device=device, dtype=torch.int32) for item in split_indices], + dim=0, + ) + except Exception: + metadata_state_indices = None + + if ( + metadata_state_indices is not None + and hasattr(metadata_state_indices, "numel") + and int(metadata_state_indices.numel()) == needed_state_indices + ): + fixed_state_indices = metadata_state_indices.detach().reshape(-1).to(device=device, dtype=torch.int32) + else: + fixed_state_indices = _resize_mamba_state_indices(state_indices, needed_state_indices, device) + + if getattr(state_indices, "numel", lambda: 0)() == needed_state_indices: + try: + current = state_indices.detach().reshape(-1).to(device=device, dtype=torch.int32) + if torch.equal(current, fixed_state_indices): + return cache_params + except Exception: + return cache_params + + try: + if dataclasses.is_dataclass(cache_params): + return dataclasses.replace(cache_params, state_indices_tensor=fixed_state_indices) + except Exception: + pass + fields = dict(getattr(cache_params, "__dict__", {})) + fields["state_indices_tensor"] = fixed_state_indices + return types.SimpleNamespace(**fields) + + +def _select_mamba2_attention_metadata(value: Any) -> Any | None: + if not isinstance(value, dict): + return _repair_mamba2_attention_metadata_state_indices(value) if _is_mamba2_attention_metadata(value) else None + converted = _mamba2_attention_metadata_from_mapping(value) + if converted is not None: + return _repair_mamba2_attention_metadata_state_indices(converted) + for child in value.values(): + selected = _select_mamba2_attention_metadata(child) + if selected is not None: + return selected + return None + + +def _normalize_mamba2_attention_metadata_groups(attn_metadata: Any) -> Any: + if not isinstance(attn_metadata, dict): + return attn_metadata + normalized = {} + changed = False + for key, value in attn_metadata.items(): + selected = _select_mamba2_attention_metadata(value) + if selected is not None and selected is not value: + normalized[key] = selected + changed = True + continue + normalized_value = _normalize_mamba2_attention_metadata_groups(value) + normalized[key] = normalized_value + changed = changed or normalized_value is not value + return normalized if changed else attn_metadata + + +def _as_v1_mamba2_attention_metadata_dict(attn_metadata: Any, vllm_config: Any) -> Any: + if isinstance(attn_metadata, dict) or not _looks_like_easymagpie_vllm_config(vllm_config): + return attn_metadata + selected = _select_mamba2_attention_metadata(attn_metadata) + if selected is None: + return attn_metadata + selected = _repair_mamba2_attention_metadata_state_indices(selected) + + layer_metadata: dict[str, Any] = {} + static_context = getattr(getattr(vllm_config, "compilation_config", None), "static_forward_context", None) + if isinstance(static_context, dict): + for layer_name in static_context: + if layer_name.endswith(".mixer"): + layer_metadata.setdefault(layer_name, selected) + + pattern = _get_easymagpie_hybrid_pattern(vllm_config) + if pattern: + for layer_idx, layer_type in enumerate(pattern): + if layer_type == "M": + layer_metadata.setdefault(f"backbone.layers.{layer_idx}.mixer", selected) + + return layer_metadata or attn_metadata + + +def _build_flash_attention_metadata_from_lengths( + *, + num_tokens: int, + num_reqs: int, + max_query_len: int, + device: Any = None, + slot_mapping: Any = None, + query_start_values: list[int] | None = None, + seq_lens: list[int] | None = None, +) -> Any | None: + torch = _TORCH_MODULE + FlashAttentionMetadata = _LEGACY_FLASH_ATTENTION_METADATA_CLS + if torch is None or FlashAttentionMetadata is None: + return None + + num_tokens = max(1, int(num_tokens)) + num_reqs = max(1, int(num_reqs)) + if query_start_values is not None and len(query_start_values) == num_reqs + 1: + query_lens = [ + max(0, query_start_values[index + 1] - query_start_values[index]) + for index in range(num_reqs) + ] + if sum(query_lens) != num_tokens: + query_lens = _balanced_profile_lengths(num_tokens, num_reqs) + query_start_values = None + else: + query_lens = _balanced_profile_lengths(num_tokens, num_reqs) + query_start_values = None + + if ( + seq_lens is None + or len(seq_lens) != num_reqs + or any(length <= 0 for length in seq_lens) + or sum(seq_lens) != num_tokens + ): + seq_lens = list(query_lens) + + if query_start_values is None: + query_start_values = [0] + for length in query_lens: + query_start_values.append(query_start_values[-1] + int(length)) + seq_start_values = [0] + for length in seq_lens: + seq_start_values.append(seq_start_values[-1] + int(length)) + + if device is None: + device = getattr(slot_mapping, "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) + seq_start_loc = torch.tensor(seq_start_values, dtype=torch.int32, device=device) + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device) + context_lens_tensor = torch.zeros(num_reqs, dtype=torch.int, device=device) + if slot_mapping is None or not hasattr(slot_mapping, "shape"): + slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=device) + else: + slot_mapping = slot_mapping[:num_tokens].to(device=device, dtype=torch.long) + block_tables = torch.empty((num_reqs, 0), dtype=torch.int, device=device) + + return FlashAttentionMetadata( + num_prefills=num_reqs, + num_prefill_tokens=num_tokens, + num_decode_tokens=0, + slot_mapping=slot_mapping, + multi_modal_placeholder_index_maps={}, + enable_kv_scales_calculation=True, + seq_lens=seq_lens, + seq_lens_tensor=seq_lens_tensor, + max_query_len=max(int(max_query_len), _max_int_sequence(query_lens)), + max_prefill_seq_len=_max_int_sequence(seq_lens), + max_decode_query_len=0, + max_decode_seq_len=0, + query_start_loc=query_start_loc, + seq_start_loc=seq_start_loc, + context_lens_tensor=context_lens_tensor, + block_tables=block_tables, + use_cuda_graph=False, + ) + + +def _requested_attention_backend() -> str: + try: + import os + + return str(os.environ.get("VLLM_ATTENTION_BACKEND") or "").upper() + except Exception: + return "" + + +def _profile_attention_lengths( + *, + num_tokens: int, + num_reqs: int, + query_start_values: list[int] | None = None, + seq_lens: list[int] | None = None, +) -> tuple[list[int], list[int], list[int]]: + num_tokens = max(1, int(num_tokens)) + num_reqs = max(1, int(num_reqs)) + if query_start_values is not None and len(query_start_values) == num_reqs + 1: + query_lens = [ + max(0, query_start_values[index + 1] - query_start_values[index]) + for index in range(num_reqs) + ] + if sum(query_lens) != num_tokens: + query_lens = _balanced_profile_lengths(num_tokens, num_reqs) + query_start_values = None + else: + query_lens = _balanced_profile_lengths(num_tokens, num_reqs) + query_start_values = None + + if ( + seq_lens is None + or len(seq_lens) != num_reqs + or any(length <= 0 for length in seq_lens) + or sum(seq_lens) != num_tokens + ): + seq_lens = list(query_lens) + + if query_start_values is None: + query_start_values = [0] + for length in query_lens: + query_start_values.append(query_start_values[-1] + int(length)) + return query_lens, query_start_values, seq_lens + + +def _build_v1_triton_attention_metadata_from_lengths( + *, + num_tokens: int, + num_reqs: int, + max_query_len: int, + device: Any = None, + slot_mapping: Any = None, + query_start_values: list[int] | None = None, + seq_lens: list[int] | None = None, +) -> Any | None: + torch = _TORCH_MODULE + TritonAttentionMetadata = _V1_TRITON_ATTENTION_METADATA_CLS + if torch is None or TritonAttentionMetadata is None: + return None + + query_lens, query_start_values, seq_lens = _profile_attention_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + if device is None: + device = getattr(slot_mapping, "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + num_tokens = max(1, int(num_tokens)) + num_reqs = max(1, int(num_reqs)) + query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32, device=device) + if slot_mapping is None or not hasattr(slot_mapping, "shape"): + slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=device) + else: + slot_mapping = slot_mapping[:num_tokens].to(device=device, dtype=torch.long) + block_table = torch.empty((num_reqs, 0), dtype=torch.int32, device=device) + empty_float = torch.empty((0,), dtype=torch.float32, device=device) + context_lens_tensor = torch.zeros(num_reqs, dtype=torch.int32, device=device) + resolved_max_query_len = max(int(max_query_len), _max_int_sequence(query_lens)) + + metadata = _call_with_supported_kwargs( + TritonAttentionMetadata, + { + "num_actual_tokens": num_tokens, + "max_query_len": resolved_max_query_len, + "query_start_loc": query_start_loc, + "max_seq_len": _max_int_sequence(seq_lens), + "seq_lens": seq_lens_tensor, + "block_table": block_table, + "slot_mapping": slot_mapping, + "seq_threshold_3D": 0, + "num_par_softmax_segments": 0, + "softmax_segm_output": empty_float, + "softmax_segm_max": empty_float, + "softmax_segm_expsum": empty_float, + "use_cascade": False, + "common_prefix_len": 0, + "cu_prefix_query_lens": None, + "prefix_kv_lens": None, + "suffix_kv_lens": None, + "scheduler_metadata": None, + "prefix_scheduler_metadata": None, + "mm_prefix_range": None, + "mm_prefix_range_tensor": None, + }, + ) + return _attach_profile_attention_metadata_attrs( + metadata, + num_tokens=num_tokens, + num_reqs=num_reqs, + query_lens=query_lens, + seq_lens=seq_lens, + seq_lens_tensor=seq_lens_tensor, + seq_start_loc=query_start_loc, + query_start_loc=query_start_loc, + context_lens_tensor=context_lens_tensor, + block_tables=block_table, + slot_mapping=slot_mapping, + max_query_len=resolved_max_query_len, + ) + + +def _build_v1_flash_attention_metadata_from_lengths( + *, + num_tokens: int, + num_reqs: int, + max_query_len: int, + device: Any = None, + slot_mapping: Any = None, + query_start_values: list[int] | None = None, + seq_lens: list[int] | None = None, +) -> Any | None: + torch = _TORCH_MODULE + FlashAttentionMetadata = _V1_FLASH_ATTENTION_METADATA_CLS + if torch is None or FlashAttentionMetadata is None: + return None + + query_lens, query_start_values, seq_lens = _profile_attention_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + if device is None: + device = getattr(slot_mapping, "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + num_tokens = max(1, int(num_tokens)) + num_reqs = max(1, int(num_reqs)) + query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32, device=device) + if slot_mapping is None or not hasattr(slot_mapping, "shape"): + slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=device) + else: + slot_mapping = slot_mapping[:num_tokens].to(device=device, dtype=torch.long) + block_table = torch.empty((num_reqs, 0), dtype=torch.int32, device=device) + context_lens_tensor = torch.zeros(num_reqs, dtype=torch.int32, device=device) + resolved_max_query_len = max(int(max_query_len), _max_int_sequence(query_lens)) + + metadata = _call_with_supported_kwargs( + FlashAttentionMetadata, + { + "num_actual_tokens": num_tokens, + "max_query_len": resolved_max_query_len, + "query_start_loc": query_start_loc, + "max_seq_len": _max_int_sequence(seq_lens), + "seq_lens": seq_lens_tensor, + "block_table": block_table, + "slot_mapping": slot_mapping, + "use_cascade": False, + "common_prefix_len": 0, + "cu_prefix_query_lens": None, + "prefix_kv_lens": None, + "suffix_kv_lens": None, + "max_dcp_context_kv_len": None, + "dcp_context_kv_lens": None, + "scheduler_metadata": None, + "prefix_scheduler_metadata": None, + "max_num_splits": 0, + "causal": True, + }, + ) + return _attach_profile_attention_metadata_attrs( + metadata, + num_tokens=num_tokens, + num_reqs=num_reqs, + query_lens=query_lens, + seq_lens=seq_lens, + seq_lens_tensor=seq_lens_tensor, + seq_start_loc=query_start_loc, + query_start_loc=query_start_loc, + context_lens_tensor=context_lens_tensor, + block_tables=block_table, + slot_mapping=slot_mapping, + max_query_len=resolved_max_query_len, + ) + + +def _build_vllm_attention_metadata_from_lengths( + *, + num_tokens: int, + num_reqs: int, + max_query_len: int, + device: Any = None, + slot_mapping: Any = None, + query_start_values: list[int] | None = None, + seq_lens: list[int] | None = None, +) -> Any: + requested_backend = _requested_attention_backend() + is_flashinfer_backend = requested_backend.startswith("FLASHINFER") + is_flash_attention_backend = requested_backend.startswith("FLASH_ATTN") or requested_backend in { + "FLASH", + "FLASHATTN", + } + if is_flash_attention_backend: + metadata = _build_v1_flash_attention_metadata_from_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + max_query_len=max_query_len, + device=device, + slot_mapping=slot_mapping, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + if metadata is not None: + return metadata + if not is_flashinfer_backend and ("TRITON" in requested_backend or requested_backend != "XFORMERS"): + metadata = _build_v1_triton_attention_metadata_from_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + max_query_len=max_query_len, + device=device, + slot_mapping=slot_mapping, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + if metadata is not None: + return metadata + metadata = _build_v1_flash_attention_metadata_from_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + max_query_len=max_query_len, + device=device, + slot_mapping=slot_mapping, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + if metadata is not None: + return metadata + + torch = _TORCH_MODULE + XFormersMetadata = _XFORMERS_METADATA_CLS + if torch is None or XFormersMetadata is None: + metadata = _build_v1_flash_attention_metadata_from_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + max_query_len=max_query_len, + device=device, + slot_mapping=slot_mapping, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + if metadata is not None: + return metadata + metadata = _build_v1_triton_attention_metadata_from_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + max_query_len=max_query_len, + device=device, + slot_mapping=slot_mapping, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + if metadata is not None: + return metadata + metadata = _build_flash_attention_metadata_from_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + max_query_len=max_query_len, + device=device, + slot_mapping=slot_mapping, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + if metadata is None: + return None + for attr in ("attn_bias", "encoder_attn_bias", "cross_attn_bias"): + if not hasattr(metadata, attr): + setattr(metadata, attr, None) + return metadata + + num_tokens = max(1, int(num_tokens)) + num_reqs = max(1, int(num_reqs)) + if query_start_values is not None and len(query_start_values) == num_reqs + 1: + query_lens = [ + max(0, query_start_values[index + 1] - query_start_values[index]) + for index in range(num_reqs) + ] + if sum(query_lens) != num_tokens: + query_lens = _balanced_profile_lengths(num_tokens, num_reqs) + query_start_values = None + else: + query_lens = _balanced_profile_lengths(num_tokens, num_reqs) + query_start_values = None + + if ( + seq_lens is None + or len(seq_lens) != num_reqs + or any(length <= 0 for length in seq_lens) + or sum(seq_lens) != num_tokens + ): + seq_lens = list(query_lens) + + if query_start_values is None: + query_start_values = [0] + for length in query_lens: + query_start_values.append(query_start_values[-1] + int(length)) + seq_start_values = [0] + for length in seq_lens: + seq_start_values.append(seq_start_values[-1] + int(length)) + + if device is None: + device = getattr(slot_mapping, "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) + seq_start_loc = torch.tensor(seq_start_values, dtype=torch.int32, device=device) + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device) + context_lens_tensor = torch.zeros(num_reqs, dtype=torch.int, device=device) + if slot_mapping is None or not hasattr(slot_mapping, "shape"): + slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=device) + else: + slot_mapping = slot_mapping[:num_tokens].to(device=device, dtype=torch.long) + block_tables = torch.empty((num_reqs, 0), dtype=torch.int, device=device) + + return XFormersMetadata( + seq_lens_tensor=seq_lens_tensor, + max_decode_seq_len=0, + block_tables=block_tables, + num_prefills=num_reqs, + num_prefill_tokens=num_tokens, + num_decode_tokens=0, + slot_mapping=slot_mapping, + multi_modal_placeholder_index_maps={}, + enable_kv_scales_calculation=True, + max_prefill_seq_len=_max_int_sequence(seq_lens), + use_cuda_graph=False, + seq_lens=seq_lens, + seq_start_loc=seq_start_loc, + context_lens_tensor=context_lens_tensor, + max_query_len=max(int(max_query_len), _max_int_sequence(query_lens)), + max_decode_query_len=0, + query_start_loc=query_start_loc, + ) + + +def _build_profile_flash_attention_metadata(self: Any, num_tokens: int, num_reqs: int, max_query_len: int) -> Any: + query_start_values = _buffer_slice_to_int_list(getattr(self, "query_start_loc", None), int(num_reqs) + 1) + seq_lens = _buffer_slice_to_int_list(getattr(self, "seq_lens", None), int(num_reqs)) + device = getattr(self, "device", None) + if device is None: + query_start_gpu = getattr(getattr(self, "query_start_loc", None), "gpu", None) + device = getattr(query_start_gpu, "device", None) + return _build_vllm_attention_metadata_from_lengths( + num_tokens=num_tokens, + num_reqs=num_reqs, + max_query_len=max_query_len, + device=device, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + + +def _get_easymagpie_hybrid_pattern(vllm_config: Any) -> str | None: + model_config = getattr(vllm_config, "model_config", None) + hf_config = getattr(model_config, "hf_config", None) + for candidate in (hf_config, model_config, vllm_config): + pattern = getattr(candidate, "hybrid_override_pattern", None) + if isinstance(pattern, str) and pattern: + return pattern + return None + + +def _get_easymagpie_chunk_size(vllm_config: Any) -> int: + model_config = getattr(vllm_config, "model_config", None) + get_mamba_chunk_size = getattr(model_config, "get_mamba_chunk_size", None) + if get_mamba_chunk_size is not None: + try: + chunk_size = get_mamba_chunk_size() + if chunk_size is not None: + return int(chunk_size) + except Exception: + pass + hf_config = getattr(model_config, "hf_config", None) + for candidate in (hf_config, model_config, vllm_config): + chunk_size = getattr(candidate, "chunk_size", None) + if chunk_size is not None: + return int(chunk_size) + return 128 + + +def _build_profile_mamba2_attention_metadata(flash_metadata: Any, chunk_size: int) -> Any: + torch = _TORCH_MODULE + Mamba2AttentionMetadata = _V1_MAMBA2_ATTENTION_METADATA_CLS + if torch is None or Mamba2AttentionMetadata is None or flash_metadata is None: + return None + + query_start_loc = getattr(flash_metadata, "query_start_loc", None) + if query_start_loc is None: + device = getattr(getattr(flash_metadata, "slot_mapping", None), "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + num_prefills = int(getattr(flash_metadata, "num_prefills", 1) or 1) + query_start_loc = torch.arange(num_prefills + 1, dtype=torch.int32, device=device) + else: + device = query_start_loc.device + num_prefills = max(0, int(getattr(flash_metadata, "num_prefills", query_start_loc.numel() - 1))) + + num_prefill_tokens_value = getattr(flash_metadata, "num_prefill_tokens", None) + if num_prefill_tokens_value is None: + num_prefill_tokens_value = query_start_loc[-1].item() + num_prefill_tokens = int(num_prefill_tokens_value) + seq_lens = getattr(flash_metadata, "seq_lens_tensor", None) + if seq_lens is None: + seq_lens_values = getattr(flash_metadata, "seq_lens", None) + if seq_lens_values is None: + seq_lens = query_start_loc[1:] - query_start_loc[:-1] + else: + seq_lens = torch.tensor(seq_lens_values, dtype=torch.int, device=device) + else: + seq_lens = seq_lens.to(device=device) + + if num_prefills > 0: + query_lens = _balanced_profile_lengths(num_prefill_tokens, num_prefills) + query_start_values = [0] + for query_len in query_lens: + query_start_values.append(query_start_values[-1] + int(query_len)) + query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) + seq_lens = torch.tensor(query_lens, dtype=torch.int, device=device) + query_start_loc_p = query_start_loc[: num_prefills + 1] + seq_idx_p = torch.repeat_interleave( + torch.arange(num_prefills, dtype=torch.int32, device=device), + query_start_loc_p.diff(), + output_size=num_prefill_tokens, + ) + seq_idx_p.unsqueeze_(0) + has_initial_states_p = torch.zeros(num_prefills, dtype=torch.bool, device=device) + else: + seq_idx_p = None + has_initial_states_p = None + + state_indices_tensor_p = ( + torch.arange(max(1, num_prefills), dtype=torch.int32, device=device)[:num_prefills] + if num_prefills > 0 + else None + ) + num_computed_tokens_p = torch.zeros(num_prefills, dtype=torch.int32, device=device) if num_prefills > 0 else None + kwargs = { + "num_reqs": num_prefills, + "num_prefills": num_prefills, + "num_prefill_tokens": num_prefill_tokens, + "num_decodes": 0, + "num_decode_tokens": 0, + "query_start_loc_p": query_start_loc_p if num_prefills > 0 else None, + "query_start_loc_d": None, + "seq_lens": seq_lens, + "num_computed_tokens_p": num_computed_tokens_p, + "prep_initial_states": False, + "chunk_size": int(chunk_size), + "has_initial_states_p": has_initial_states_p, + "seq_idx_p": seq_idx_p, + "state_indices_tensor_p": state_indices_tensor_p, + "state_indices_tensor_d": None, + "num_accepted_tokens": None, + "block_idx_last_scheduled_token": None, + "block_idx_first_scheduled_token_p": None, + "block_idx_last_computed_token": None, + "cu_chunk_seqlen_p": None, + "last_chunk_indices_p": None, + "nums_dict": None, + "batch_ptr": None, + "token_chunk_offset_ptr": None, + } + metadata = _new_mamba2_attention_metadata(Mamba2AttentionMetadata, kwargs) + return metadata + + +def _build_mamba2_attention_metadata_from_legacy( + legacy_metadata: Any, + mamba_cache_params: Any, + hidden_states: Any, +) -> Any | None: + torch = _TORCH_MODULE + Mamba2AttentionMetadata = _V1_MAMBA2_ATTENTION_METADATA_CLS + if torch is None or Mamba2AttentionMetadata is None: + return None + + if legacy_metadata is None: + return None + + device = getattr(hidden_states, "device", None) + if device is None: + state_indices = getattr(mamba_cache_params, "state_indices_tensor", None) + device = getattr(state_indices, "device", None) + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + hidden_len = int(getattr(hidden_states, "shape", [1])[0] or 1) + legacy_seq_idx_p = getattr(legacy_metadata, "seq_idx", None) + chunk_indices_p = getattr(legacy_metadata, "chunk_indices", None) + chunk_offsets_p = getattr(legacy_metadata, "chunk_offsets", None) + chunk_size = int(getattr(legacy_metadata, "chunk_size", 128) or 128) + + if legacy_seq_idx_p is not None and getattr(legacy_seq_idx_p, "numel", lambda: 0)() > 0: + seq_idx_flat = _compact_legacy_mamba_seq_idx(legacy_seq_idx_p, device) + num_prefill_tokens = int(seq_idx_flat.numel()) + num_prefills = int(seq_idx_flat.max().item()) + 1 if num_prefill_tokens > 0 else 0 + counts = ( + torch.bincount(seq_idx_flat, minlength=num_prefills).to(dtype=torch.int32, device=device) + if num_prefills > 0 + else torch.empty(0, dtype=torch.int32, device=device) + ) + query_start_loc_p = torch.zeros(num_prefills + 1, dtype=torch.int32, device=device) + if num_prefills > 0: + query_start_loc_p[1:] = torch.cumsum(counts, dim=0) + num_decode_tokens = max(0, hidden_len - num_prefill_tokens) + num_decodes = num_decode_tokens + decode_lens = ( + torch.ones(num_decode_tokens, dtype=torch.int32, device=device) + if num_decode_tokens > 0 + else torch.empty(0, dtype=torch.int32, device=device) + ) + seq_lens = torch.cat([counts, decode_lens], dim=0) if num_decode_tokens > 0 else counts + legacy_query_start_loc = torch.zeros(num_prefills + num_decodes + 1, dtype=torch.int32, device=device) + if seq_lens.numel() > 0: + legacy_query_start_loc[1:] = torch.cumsum(seq_lens, dim=0) + has_initial_states_p = getattr(legacy_metadata, "has_initial_states", None) + if has_initial_states_p is None: + has_initial_states_p = torch.zeros(num_prefills, dtype=torch.bool, device=device) + else: + has_initial_states_p = has_initial_states_p.to(device=device, dtype=torch.bool)[:num_prefills] + seq_idx_p = seq_idx_flat.unsqueeze(0) if num_prefill_tokens > 0 else None + else: + num_prefills = 0 + num_prefill_tokens = 0 + num_decodes = max(1, hidden_len) + num_decode_tokens = num_decodes + query_start_loc_p = None + legacy_query_start_loc = torch.arange(num_decodes + 1, dtype=torch.int32, device=device) + seq_lens = torch.ones(num_decodes, dtype=torch.int32, device=device) + has_initial_states_p = None + seq_idx_p = None + + state_indices_tensor = getattr(mamba_cache_params, "state_indices_tensor", None) + state_indices_tensor_p = getattr(mamba_cache_params, "state_indices_tensor_p", None) + if num_prefills > 0 and ( + state_indices_tensor_p is None or getattr(state_indices_tensor_p, "numel", lambda: 0)() != num_prefills + ): + state_indices_tensor_p = _resize_mamba_state_indices( + state_indices_tensor_p if state_indices_tensor_p is not None else state_indices_tensor, + num_prefills, + device, + ) + elif num_prefills <= 0: + state_indices_tensor_p = None + + state_indices_tensor_d = getattr(mamba_cache_params, "state_indices_tensor_d", None) + if num_decode_tokens > 0 and ( + state_indices_tensor_d is None + or getattr(state_indices_tensor_d, "numel", lambda: 0)() != num_decode_tokens + ): + state_indices_tensor_d = _resize_mamba_state_indices( + state_indices_tensor_d if state_indices_tensor_d is not None else state_indices_tensor, + num_decode_tokens, + device, + ) + elif num_decode_tokens <= 0: + state_indices_tensor_d = None + + query_start_loc_d = ( + torch.arange(num_decodes + 1, dtype=torch.int32, device=device) if num_decode_tokens > 0 else None + ) + num_computed_tokens_p = torch.zeros(num_prefills, dtype=torch.int32, device=device) if num_prefills > 0 else None + legacy_state_indices_tensor = _resize_mamba_state_indices( + state_indices_tensor, + num_prefills + num_decode_tokens, + device, + ) + + metadata = _new_mamba2_attention_metadata( + Mamba2AttentionMetadata, + { + "num_reqs": num_prefills + num_decodes, + "num_prefills": num_prefills, + "num_prefill_tokens": num_prefill_tokens, + "num_decodes": num_decodes, + "num_decode_tokens": num_decode_tokens, + "query_start_loc_p": query_start_loc_p, + "query_start_loc_d": query_start_loc_d, + "legacy_query_start_loc": legacy_query_start_loc, + "seq_lens": seq_lens, + "num_computed_tokens_p": num_computed_tokens_p, + "prep_initial_states": bool(getattr(legacy_metadata, "prep_initial_states", False)), + "chunk_size": chunk_size, + "has_initial_states_p": has_initial_states_p, + "seq_idx_p": seq_idx_p, + "state_indices_tensor_p": state_indices_tensor_p, + "state_indices_tensor_d": state_indices_tensor_d, + "legacy_state_indices_tensor": legacy_state_indices_tensor, + "num_accepted_tokens": None, + "block_idx_last_scheduled_token": None, + "block_idx_first_scheduled_token_p": None, + "block_idx_last_computed_token": None, + "cu_chunk_seqlen_p": None, + "last_chunk_indices_p": None, + "nums_dict": None, + "batch_ptr": None, + "token_chunk_offset_ptr": None, + }, + ) + return _repair_mamba2_attention_metadata_prefill_seq_idx(metadata) + + +def _select_legacy_mamba2_metadata(value: Any) -> Any | None: + if not isinstance(value, dict): + if any(hasattr(value, attr) for attr in ("seq_idx", "has_initial_states", "chunk_indices", "chunk_offsets")): + return value + return None + for child in value.values(): + selected = _select_legacy_mamba2_metadata(child) + if selected is not None: + return selected + return None + + +def _build_hybrid_profile_attention_metadata(vllm_config: Any, flash_metadata: Any) -> Any: + layer_metadata: dict[str, Any] = {} + mamba_metadata = _build_profile_mamba2_attention_metadata( + flash_metadata, + _get_easymagpie_chunk_size(vllm_config), + ) + static_context = getattr(getattr(vllm_config, "compilation_config", None), "static_forward_context", None) + if isinstance(static_context, dict): + for layer_name in static_context: + if layer_name.endswith(".attn") and flash_metadata is not None: + layer_metadata.setdefault(layer_name, flash_metadata) + elif layer_name.endswith(".mixer") and mamba_metadata is not None: + layer_metadata.setdefault(layer_name, mamba_metadata) + + pattern = _get_easymagpie_hybrid_pattern(vllm_config) + if pattern: + for layer_idx, layer_type in enumerate(pattern): + if layer_type == "M" and mamba_metadata is not None: + layer_metadata.setdefault(f"backbone.layers.{layer_idx}.mixer", mamba_metadata) + elif layer_type == "*" and flash_metadata is not None: + layer_metadata.setdefault(f"backbone.layers.{layer_idx}.mixer.attn", flash_metadata) + return layer_metadata or flash_metadata + + +def _looks_like_easymagpie_vllm_config(vllm_config: Any) -> bool: + candidates = [vllm_config, getattr(vllm_config, "model_config", None)] + model_config = candidates[-1] + candidates.append(getattr(model_config, "hf_config", None)) + for candidate in candidates: + if candidate is None: + continue + class_name = candidate.__class__.__name__.lower() + if "easymagpie" in class_name or "nemotronh" in class_name or "nemotron_h" in class_name: + return True + pattern = getattr(candidate, "hybrid_override_pattern", None) + if isinstance(pattern, str) and "M" in pattern and hasattr(candidate, "chunk_size"): + return True + for attr_name in ("model", "served_model_name", "model_type"): + value = getattr(candidate, attr_name, None) + if value is not None and ( + "easymagpie" in str(value).lower() + or "nemotronh" in str(value).lower() + or "nemotron_h" in str(value).lower() + ): + return True + architectures = getattr(candidate, "architectures", None) or () + if any( + "easymagpie" in str(architecture).lower() + or "nemotronh" in str(architecture).lower() + or "nemotron_h" in str(architecture).lower() + for architecture in architectures + ): + return True + return False + + +def _is_generic_profile_attention_metadata(attn_metadata: Any) -> bool: + if attn_metadata is None or isinstance(attn_metadata, dict): + return False + if hasattr(attn_metadata, "use_cascade"): + return False + return any( + hasattr(attn_metadata, attr) + for attr in ( + "num_prefills", + "num_prefill_tokens", + "query_start_loc", + "seq_lens", + "seq_lens_tensor", + ) + ) + + +def _is_empty_attention_kv_cache(kv_cache: Any) -> bool: + if kv_cache is None: + return True + shape = getattr(kv_cache, "shape", None) + if shape is not None: + try: + return int(shape[0]) < 2 + except Exception: + pass + try: + return len(kv_cache) < 2 + except Exception: + return False + + +def _filter_supported_kwargs(callable_obj: Any, kwargs: dict[str, Any]) -> dict[str, Any]: + try: + parameters = inspect.signature(callable_obj).parameters + except Exception: + return dict(kwargs) + if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()): + return dict(kwargs) + return {key: value for key, value in kwargs.items() if key in parameters} + + +def _synthetic_profile_attention_metadata_from_context(attn_metadata: Any, vllm_config: Any, kwargs: dict[str, Any]) -> Any: + if not _looks_like_easymagpie_vllm_config(vllm_config): + return attn_metadata + if attn_metadata is not None and not _is_generic_profile_attention_metadata(attn_metadata): + return attn_metadata + + num_tokens = None + num_reqs = None + max_query_len = None + slot_mapping = kwargs.get("slot_mapping") + query_start_values = None + seq_lens = None + if attn_metadata is not None: + num_tokens = getattr(attn_metadata, "num_prefill_tokens", None) + if num_tokens is None: + num_tokens = getattr(attn_metadata, "num_actual_tokens", None) + num_reqs = getattr(attn_metadata, "num_prefills", None) + if num_reqs is None: + num_reqs = getattr(attn_metadata, "num_reqs", None) + max_query_len = getattr(attn_metadata, "max_query_len", None) + metadata_slot_mapping = getattr(attn_metadata, "slot_mapping", None) + if metadata_slot_mapping is not None: + slot_mapping = metadata_slot_mapping + if num_reqs is not None: + query_start_values = _buffer_slice_to_int_list( + getattr(attn_metadata, "query_start_loc", None), + int(num_reqs) + 1, + ) + seq_lens = _buffer_slice_to_int_list( + _first_not_none( + getattr(attn_metadata, "seq_lens_tensor", None), + getattr(attn_metadata, "seq_lens", None), + ), + int(num_reqs), + ) + + if num_tokens is None: + num_tokens = kwargs.get("num_tokens") + if num_tokens is None: + return attn_metadata + if num_reqs is None: + batch_descriptor = kwargs.get("batch_descriptor") + num_reqs = getattr(batch_descriptor, "num_reqs", None) + if num_reqs is None: + num_reqs = min(max(1, int(num_tokens)), 1) + if max_query_len is None: + max_query_len = max(1, int(num_tokens) // max(1, int(num_reqs))) + flash_metadata = _build_vllm_attention_metadata_from_lengths( + num_tokens=int(num_tokens), + num_reqs=int(num_reqs), + max_query_len=int(max_query_len), + slot_mapping=slot_mapping, + query_start_values=query_start_values, + seq_lens=seq_lens, + ) + return _build_hybrid_profile_attention_metadata(vllm_config, flash_metadata) + + +def _install_mamba2_metadata_compat() -> None: + try: + module = importlib.import_module("vllm.model_executor.layers.mamba.mamba2_metadata") + except Exception: + return + original = getattr(module, "prepare_mamba2_metadata", None) + if original is None or getattr(original, "_easymagpie_compat", False): + return + + def _empty_attn_metadata(): + return types.SimpleNamespace( + num_prefills=0, + num_prefill_tokens=0, + context_lens_tensor=None, + query_start_loc=None, + ) + + def _select_attn_metadata(attn_metadata: Any): + if not isinstance(attn_metadata, dict): + return attn_metadata if hasattr(attn_metadata, "num_prefills") else None + for value in attn_metadata.values(): + selected = _select_attn_metadata(value) + if selected is not None: + return selected + return None + + def _balanced_tensor_lengths(total: int, count: int, device: Any): + import torch + + if count <= 0: + return torch.empty(0, dtype=torch.int32, device=device) + base, extra = divmod(max(0, int(total)), int(count)) + lengths = torch.full((count,), base, dtype=torch.int32, device=device) + if extra: + lengths[:extra] += 1 + return lengths + + def _build_safe_legacy_metadata(chunk_size: int, attn_metadata: Any, mamba2_metadata: Any = None): + import torch + + num_prefills = int(getattr(attn_metadata, "num_prefills", 0) or 0) + num_prefill_tokens = int(getattr(attn_metadata, "num_prefill_tokens", 0) or 0) + query_start_loc = getattr(attn_metadata, "query_start_loc", None) + if num_prefills <= 0 or num_prefill_tokens <= 0 or query_start_loc is None: + return None + + query_start_loc = query_start_loc.to(dtype=torch.int64) + if int(query_start_loc.numel()) < 2: + return None + query_lens = query_start_loc[1:] - query_start_loc[:-1] + device = query_start_loc.device + + if int(query_start_loc.numel()) >= num_prefills + 1: + first_prefill_lens = query_start_loc[: num_prefills + 1].diff() + try: + if int(first_prefill_lens.sum().item()) == num_prefill_tokens: + return None + except Exception: + pass + + prefill_lens = query_lens[query_lens > 1] + try: + prefill_lens_ok = ( + int(prefill_lens.numel()) == num_prefills + and int(prefill_lens.sum().item()) == num_prefill_tokens + ) + except Exception: + prefill_lens_ok = False + if not prefill_lens_ok: + prefill_lens = _balanced_tensor_lengths(num_prefill_tokens, num_prefills, device) + else: + prefill_lens = prefill_lens.to(dtype=torch.int32) + + query_start_loc_p = torch.zeros(num_prefills + 1, dtype=torch.int32, device=device) + if num_prefills > 0: + query_start_loc_p[1:] = torch.cumsum(prefill_lens.to(dtype=torch.int32), dim=0) + seq_idx = torch.repeat_interleave( + torch.arange(num_prefills, dtype=torch.int32, device=device), + prefill_lens.to(dtype=torch.int32), + ).unsqueeze(0) + + has_initial_states = None + prep_initial_states = False + context_lens_tensor = getattr(attn_metadata, "context_lens_tensor", None) + if context_lens_tensor is not None: + try: + context_lens_tensor = context_lens_tensor.to(device=device) + if int(context_lens_tensor.numel()) == int(query_lens.numel()) and int( + (query_lens > 1).sum().item() + ) == num_prefills: + context_lens_tensor = context_lens_tensor[query_lens > 1] + else: + context_lens_tensor = context_lens_tensor[:num_prefills] + has_initial_states = context_lens_tensor > 0 + prep_initial_states = bool(torch.any(has_initial_states).item()) + except Exception: + has_initial_states = None + prep_initial_states = False + + chunk_indices, chunk_offsets = None, None + if prep_initial_states: + try: + converter = getattr(module, "_query_start_loc_to_chunk_indices_offsets") + chunk_indices, chunk_offsets = converter(query_start_loc_p, chunk_size, num_prefill_tokens) + except Exception: + chunk_indices, chunk_offsets = None, None + + target = mamba2_metadata + if target is None: + metadata_cls = getattr(module, "Mamba2Metadata", None) + if metadata_cls is None: + return types.SimpleNamespace( + has_initial_states=has_initial_states, + prep_initial_states=prep_initial_states, + chunk_size=chunk_size, + seq_idx=seq_idx, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + cu_seqlen=None, + ) + return metadata_cls( + has_initial_states=has_initial_states, + prep_initial_states=prep_initial_states, + chunk_size=chunk_size, + seq_idx=seq_idx, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + ) + + target.has_initial_states = has_initial_states + target.prep_initial_states = prep_initial_states + target.chunk_size = chunk_size + target.seq_idx = seq_idx + target.chunk_indices = chunk_indices + target.chunk_offsets = chunk_offsets + target.cu_seqlen = None + return target + + def prepare_mamba2_metadata(chunk_size: int, attn_metadata: Any, mamba2_metadata: Any = None): + if attn_metadata is None: + attn_metadata = _empty_attn_metadata() + elif isinstance(attn_metadata, dict): + attn_metadata = _select_attn_metadata(attn_metadata) or _empty_attn_metadata() + if not hasattr(attn_metadata, "num_prefills"): + attn_metadata = types.SimpleNamespace( + num_prefills=0, + num_prefill_tokens=0, + context_lens_tensor=None, + query_start_loc=None, + ) + safe_metadata = _build_safe_legacy_metadata(chunk_size, attn_metadata, mamba2_metadata) + if safe_metadata is not None: + return safe_metadata + return original(chunk_size, attn_metadata, mamba2_metadata) + + prepare_mamba2_metadata._easymagpie_compat = True # type: ignore[attr-defined] + module.prepare_mamba2_metadata = prepare_mamba2_metadata + for module_name in ("vllm.model_executor.models.nemotron_h",): + loaded_module = sys.modules.get(module_name) + if loaded_module is not None and getattr(loaded_module, "prepare_mamba2_metadata", None) is original: + loaded_module.prepare_mamba2_metadata = prepare_mamba2_metadata + + +def _install_triton_attention_profile_metadata_compat() -> None: + try: + module = importlib.import_module("vllm.v1.attention.backends.triton_attn") + except Exception: + return + impl_cls = getattr(module, "TritonAttentionImpl", None) + original = getattr(impl_cls, "forward", None) + if original is None or getattr(original, "_easymagpie_profile_metadata_compat", False): + return + + def forward( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale=None, + output_block_scale=None, + ): + if _is_generic_profile_attention_metadata(attn_metadata): + return original( + self, + layer, + query, + key, + value, + kv_cache, + None, + output, + output_scale, + output_block_scale, + ) + return original( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale, + output_block_scale, + ) + + forward._easymagpie_profile_metadata_compat = True # type: ignore[attr-defined] + impl_cls.forward = forward + + +def _install_flash_attention_profile_metadata_compat() -> None: + try: + module = importlib.import_module("vllm.v1.attention.backends.flash_attn") + except Exception: + return + impl_cls = getattr(module, "FlashAttentionImpl", None) + original = getattr(impl_cls, "forward", None) + if original is None or getattr(original, "_easymagpie_profile_metadata_compat", False): + return + + def forward( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output=None, + output_scale=None, + output_block_scale=None, + **kwargs, + ): + optional_kwargs = _filter_supported_kwargs( + original, + {"output_block_scale": output_block_scale, **kwargs}, + ) + if _is_generic_profile_attention_metadata(attn_metadata) or _is_empty_attention_kv_cache(kv_cache): + return original( + self, + layer, + query, + key, + value, + kv_cache, + None, + output, + output_scale, + **optional_kwargs, + ) + return original( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale, + **optional_kwargs, + ) + + forward._easymagpie_profile_metadata_compat = True # type: ignore[attr-defined] + impl_cls.forward = forward + + +def _install_flashinfer_attention_profile_metadata_compat() -> None: + try: + module = importlib.import_module("vllm.v1.attention.backends.flashinfer") + except Exception: + return + impl_cls = getattr(module, "FlashInferImpl", None) + original = getattr(impl_cls, "forward", None) + if original is None or getattr(original, "_easymagpie_profile_metadata_compat", False): + return + + def forward( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output=None, + output_scale=None, + output_block_scale=None, + **kwargs, + ): + optional_kwargs = _filter_supported_kwargs( + original, + {"output_block_scale": output_block_scale, **kwargs}, + ) + if _is_generic_profile_attention_metadata(attn_metadata) or _is_empty_attention_kv_cache(kv_cache): + return original( + self, + layer, + query, + key, + value, + kv_cache, + None, + output, + output_scale, + **optional_kwargs, + ) + return original( + self, + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + output_scale, + **optional_kwargs, + ) + + forward._easymagpie_profile_metadata_compat = True # type: ignore[attr-defined] + impl_cls.forward = forward + + +def _install_mamba2_profile_no_kv_cache_compat() -> None: + try: + from vllm import envs + from vllm.forward_context import get_forward_context + + module = importlib.import_module("vllm.model_executor.layers.mamba.mamba_mixer2") + except Exception: + return + mixer_cls = getattr(module, "MambaMixer2", None) + original = getattr(mixer_cls, "forward_cuda", None) + if mixer_cls is None: + return + + def _has_usable_mamba_kv_cache_for_conv(self: Any) -> bool: + kv_cache = getattr(self, "kv_cache", None) + try: + if kv_cache is None: + return False + if len(kv_cache) == 1 and isinstance(kv_cache[0], (list, tuple)): + kv_cache = kv_cache[0] + if len(kv_cache) < 2: + return False + conv_cache = kv_cache[0] + ssm_cache = kv_cache[1] + return ( + getattr(conv_cache, "ndim", 0) >= 2 + and getattr(ssm_cache, "ndim", 0) >= 2 + ) + except Exception: + return False + + def _is_profile_only_mamba_metadata(attn_metadata: Any) -> bool: + if attn_metadata is None: + return False + if getattr(attn_metadata, "batch_ptr", None) is not None: + return False + if getattr(attn_metadata, "token_chunk_offset_ptr", None) is not None: + return False + num_decode_tokens = int(getattr(attn_metadata, "num_decode_tokens", 0) or 0) + num_prefill_tokens = int(getattr(attn_metadata, "num_prefill_tokens", 0) or 0) + return num_decode_tokens == 0 and num_prefill_tokens > 0 + + def _with_v1_mamba_metadata_dict( + self, + callback, + *args, + allow_profile_metadata_fallback: bool = False, + **kwargs, + ): + try: + use_v1 = bool(getattr(envs, "VLLM_USE_V1", False)) + except Exception: + use_v1 = False + if not use_v1: + return callback(*args, **kwargs) + + try: + forward_context = get_forward_context() + except Exception: + forward_context = None + attn_metadata = getattr(forward_context, "attn_metadata", None) + prefix = getattr(self, "prefix", None) + selected_mamba_metadata = _select_mamba2_attention_metadata(attn_metadata) + if selected_mamba_metadata is None: + projected_states = kwargs.get("projected_states") + if projected_states is None and args: + projected_states = args[0] + try: + num_tokens = max(1, int(getattr(projected_states, "shape", [1])[0] or 1)) + device = getattr(projected_states, "device", None) + if attn_metadata is not None and hasattr(attn_metadata, "num_prefills"): + flash_metadata = attn_metadata + else: + flash_metadata = _build_vllm_attention_metadata_from_lengths( + num_tokens=num_tokens, + num_reqs=1, + max_query_len=num_tokens, + device=device, + ) + chunk_size = int(getattr(self, "chunk_size", None) or getattr(self, "chunk_size_padded", None) or 256) + selected_mamba_metadata = _build_profile_mamba2_attention_metadata(flash_metadata, chunk_size) + except Exception: + selected_mamba_metadata = None + if prefix is None or selected_mamba_metadata is None or forward_context is None: + return callback(*args, **kwargs) + + if allow_profile_metadata_fallback and ( + not _has_usable_mamba_kv_cache_for_conv(self) + or _is_profile_only_mamba_metadata(selected_mamba_metadata) + ): + previous_attn_metadata = forward_context.attn_metadata + forward_context.attn_metadata = None + try: + return callback(*args, **kwargs) + finally: + forward_context.attn_metadata = previous_attn_metadata + + selected_mamba_metadata = _repair_mamba2_attention_metadata_state_indices(selected_mamba_metadata) + replacement_attn_metadata = None + if isinstance(attn_metadata, dict): + layer_metadata = attn_metadata.get(prefix) + if layer_metadata is not selected_mamba_metadata: + replacement_attn_metadata = dict(attn_metadata) + replacement_attn_metadata[prefix] = selected_mamba_metadata + else: + replacement_attn_metadata = {prefix: selected_mamba_metadata} + + if replacement_attn_metadata is None: + return callback(*args, **kwargs) + + previous_attn_metadata = forward_context.attn_metadata + forward_context.attn_metadata = replacement_attn_metadata + try: + return callback(*args, **kwargs) + finally: + forward_context.attn_metadata = previous_attn_metadata + + original_conv_ssm_forward = getattr(mixer_cls, "conv_ssm_forward", None) + if original_conv_ssm_forward is not None and not getattr( + original_conv_ssm_forward, "_easymagpie_v1_attn_metadata_compat", False + ): + + def conv_ssm_forward(self, *args, **kwargs): + return _with_v1_mamba_metadata_dict( + self, + lambda *call_args, **call_kwargs: original_conv_ssm_forward(self, *call_args, **call_kwargs), + *args, + allow_profile_metadata_fallback=True, + **kwargs, + ) + + conv_ssm_forward._easymagpie_v1_attn_metadata_compat = True # type: ignore[attr-defined] + mixer_cls.conv_ssm_forward = conv_ssm_forward + + original_forward = getattr(mixer_cls, "forward", None) + if original_forward is not None and not getattr(original_forward, "_easymagpie_v1_attn_metadata_compat", False): + + def forward(self, *args, **kwargs): + return _with_v1_mamba_metadata_dict( + self, + lambda *call_args, **call_kwargs: original_forward(self, *call_args, **call_kwargs), + *args, + **kwargs, + ) + + forward._easymagpie_v1_attn_metadata_compat = True # type: ignore[attr-defined] + mixer_cls.forward = forward + + if original is None or getattr(original, "_easymagpie_no_kv_profile_compat", False): + return + + try: + original_signature = inspect.signature(original) + except Exception: + original_signature = None + + def has_explicit_cache_args(self: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: + positional_cache_args = ( + (len(args) > 2 and args[2] is not None) + or (len(args) > 3 and args[3] is not None) + ) + keyword_cache_args = ( + kwargs.get("mamba_cache_params") is not None + or kwargs.get("mamba2_metadata") is not None + ) + if original_signature is None: + return positional_cache_args or keyword_cache_args + try: + bound = original_signature.bind_partial(self, *args, **kwargs) + except Exception: + return positional_cache_args or keyword_cache_args + return ( + bound.arguments.get("mamba_cache_params") is not None + or bound.arguments.get("mamba2_metadata") is not None + or positional_cache_args + or keyword_cache_args + ) + + def bound_forward_argument(self: Any, args: tuple[Any, ...], kwargs: dict[str, Any], name: str) -> Any: + if original_signature is not None: + try: + bound = original_signature.bind_partial(self, *args, **kwargs) + value = bound.arguments.get(name) + if value is not None: + return value + except Exception: + pass + positional_index = { + "hidden_states": 0, + "output": 1, + "mamba_cache_params": 2, + "mamba2_metadata": 3, + "mup_vector": 4, + }.get(name) + if positional_index is not None and len(args) > positional_index: + return args[positional_index] + return kwargs.get(name) + + def replace_forward_argument( + args: tuple[Any, ...], kwargs: dict[str, Any], name: str, value: Any + ) -> tuple[tuple[Any, ...], dict[str, Any]]: + positional_index = { + "hidden_states": 0, + "output": 1, + "mamba_cache_params": 2, + "mamba2_metadata": 3, + "mup_vector": 4, + }.get(name) + if positional_index is not None and len(args) > positional_index: + new_args = list(args) + new_args[positional_index] = value + return tuple(new_args), kwargs + new_kwargs = dict(kwargs) + new_kwargs[name] = value + return args, new_kwargs + + def forward_cuda(self, *args, **kwargs): + try: + use_v1 = bool(getattr(envs, "VLLM_USE_V1", False)) + except Exception: + use_v1 = False + + forward_context = None + original_attn_metadata = None + selected_mamba_metadata = None + restore_attn_metadata = False + try: + forward_context = get_forward_context() + except Exception: + forward_context = None + original_attn_metadata = getattr(forward_context, "attn_metadata", None) + normalized_attn_metadata = _normalize_mamba2_attention_metadata_groups(original_attn_metadata) + prefix = getattr(self, "prefix", None) + force_v0_mamba_forward = False + selected_from_legacy_mapping = False + if isinstance(normalized_attn_metadata, dict) and prefix is not None: + layer_metadata = normalized_attn_metadata.get(prefix) + selected_mamba_metadata = _select_mamba2_attention_metadata(layer_metadata) + if selected_mamba_metadata is None: + legacy_metadata = bound_forward_argument(self, args, kwargs, "mamba2_metadata") + if legacy_metadata is None: + legacy_metadata = _select_legacy_mamba2_metadata(layer_metadata) + if legacy_metadata is None: + legacy_metadata = layer_metadata + selected_mamba_metadata = _build_mamba2_attention_metadata_from_legacy( + legacy_metadata, + bound_forward_argument(self, args, kwargs, "mamba_cache_params"), + bound_forward_argument(self, args, kwargs, "hidden_states"), + ) + force_v0_mamba_forward = selected_mamba_metadata is None and isinstance(layer_metadata, dict) + selected_from_legacy_mapping = selected_mamba_metadata is not None and isinstance(layer_metadata, dict) + if selected_mamba_metadata is not None: + selected_mamba_metadata = _repair_mamba2_attention_metadata_state_indices(selected_mamba_metadata) + if selected_mamba_metadata is not None and selected_mamba_metadata is not layer_metadata: + normalized_attn_metadata = dict(normalized_attn_metadata) + normalized_attn_metadata[prefix] = selected_mamba_metadata + elif _is_mamba2_attention_metadata(normalized_attn_metadata): + selected_mamba_metadata = _repair_mamba2_attention_metadata_state_indices(normalized_attn_metadata) + normalized_attn_metadata = selected_mamba_metadata + if forward_context is not None: + replacement_attn_metadata = normalized_attn_metadata + if use_v1 and selected_mamba_metadata is not None and not isinstance(replacement_attn_metadata, dict): + prefix = getattr(self, "prefix", None) + if prefix is not None: + replacement_attn_metadata = {prefix: selected_mamba_metadata} + elif not use_v1 and selected_mamba_metadata is not None: + replacement_attn_metadata = selected_mamba_metadata + if replacement_attn_metadata is not original_attn_metadata: + forward_context.attn_metadata = replacement_attn_metadata + restore_attn_metadata = True + + call_args = args + call_kwargs = kwargs + if selected_mamba_metadata is not None: + cache_params = bound_forward_argument(self, args, kwargs, "mamba_cache_params") + repaired_cache_params = _repair_mamba_cache_params_state_indices(cache_params, selected_mamba_metadata) + if repaired_cache_params is not cache_params: + call_args, call_kwargs = replace_forward_argument( + args, kwargs, "mamba_cache_params", repaired_cache_params + ) + + explicit_cache_args = has_explicit_cache_args(self, call_args, call_kwargs) + try: + if ( + use_v1 + and forward_context is not None + and selected_mamba_metadata is not None + and not explicit_cache_args + and _is_profile_only_mamba_metadata(selected_mamba_metadata) + ): + previous_attn_metadata = getattr(forward_context, "attn_metadata", None) + forward_context.attn_metadata = None + try: + return original(self, *call_args, **call_kwargs) + finally: + forward_context.attn_metadata = previous_attn_metadata + if ( + use_v1 + and forward_context is not None + and (explicit_cache_args or force_v0_mamba_forward) + and (selected_mamba_metadata is not None or force_v0_mamba_forward) + ): + previous_use_v1 = getattr(envs, "VLLM_USE_V1", None) + previous_attn_metadata = getattr(forward_context, "attn_metadata", None) + envs.VLLM_USE_V1 = False + if selected_mamba_metadata is not None: + if selected_from_legacy_mapping and prefix is not None: + forward_context.attn_metadata = {prefix: selected_mamba_metadata} + else: + forward_context.attn_metadata = selected_mamba_metadata + try: + return original(self, *call_args, **call_kwargs) + finally: + envs.VLLM_USE_V1 = previous_use_v1 + forward_context.attn_metadata = previous_attn_metadata + if ( + selected_mamba_metadata is not None + and forward_context is not None + and hasattr(self, "kv_cache") + and not explicit_cache_args + ): + previous_use_v1 = getattr(envs, "VLLM_USE_V1", None) + previous_attn_metadata = getattr(forward_context, "attn_metadata", None) + prefix = getattr(self, "prefix", None) + if isinstance(normalized_attn_metadata, dict): + v1_attn_metadata = normalized_attn_metadata + elif prefix is not None: + v1_attn_metadata = {prefix: selected_mamba_metadata} + else: + v1_attn_metadata = previous_attn_metadata + envs.VLLM_USE_V1 = True + forward_context.attn_metadata = v1_attn_metadata + try: + return original(self, *call_args, **call_kwargs) + finally: + envs.VLLM_USE_V1 = previous_use_v1 + forward_context.attn_metadata = previous_attn_metadata + return original(self, *call_args, **call_kwargs) + finally: + if restore_attn_metadata and forward_context is not None: + try: + forward_context.attn_metadata = original_attn_metadata + except Exception: + pass + + forward_cuda._easymagpie_no_kv_profile_compat = True # type: ignore[attr-defined] + mixer_cls.forward_cuda = forward_cuda + + +def _install_lora_config_alias() -> None: + try: + from vllm.config import LoRAConfig + except Exception: + return + module = types.ModuleType("vllm.config.lora") + module.LoRAConfig = LoRAConfig + module.MaxLoRARanks = Literal[1, 8, 16, 32, 64, 128, 256, 320, 512] + sys.modules.setdefault("vllm.config.lora", module) + + +def _install_config_decorator_compat() -> None: + try: + import vllm.config.utils as config_utils + except Exception: + return + original = getattr(config_utils, "config", None) + if original is None or getattr(original, "_easymagpie_compat", False): + return + + def compat_config(cls=None, **kwargs): + def _decorate(real_cls): + try: + return original(real_cls, **kwargs) + except TypeError: + if not kwargs: + raise + return original(real_cls) + + if cls is None: + return _decorate + return _decorate(cls) + + compat_config._easymagpie_compat = True # type: ignore[attr-defined] + config_utils.config = compat_config + + +def _install_model_arch_convertor() -> None: + if "vllm.transformers_utils.model_arch_config_convertor" in sys.modules: + return + module = types.ModuleType("vllm.transformers_utils.model_arch_config_convertor") + + class ModelArchConfigConvertorBase: + def __init__(self, hf_config: Any, hf_text_config: Any) -> None: + self.hf_config = hf_config + self.hf_text_config = hf_text_config + + def _normalize_quantization_config(self, config: Any) -> Any: + if config is None: + return None + if isinstance(config, dict): + return config.get("quantization_config") + return getattr(config, "quantization_config", None) + + def get_quantization_config(self) -> Any: + return self._normalize_quantization_config(self.hf_text_config) or self._normalize_quantization_config( + self.hf_config + ) + + def convert(self) -> dict[str, Any]: + quant_config = self.get_quantization_config() + return {"quantization_config": quant_config} if quant_config is not None else {} + + module.ModelArchConfigConvertorBase = ModelArchConfigConvertorBase + sys.modules["vllm.transformers_utils.model_arch_config_convertor"] = module + + +def _install_io_processor_stub() -> None: + if "vllm.plugins.io_processors" in sys.modules: + return + module = types.ModuleType("vllm.plugins.io_processors") + + def get_io_processor(*_args: Any, **_kwargs: Any) -> None: + return None + + module.get_io_processor = get_io_processor + sys.modules["vllm.plugins.io_processors"] = module + + +def _install_tokenizer_aliases() -> None: + try: + import vllm.transformers_utils.tokenizer as tokenizer_module + except Exception: + return + if not hasattr(tokenizer_module, "TokenizerLike"): + tokenizer_module.TokenizerLike = getattr(tokenizer_module, "AnyTokenizer", object) + sys.modules.setdefault("vllm.tokenizers", tokenizer_module) + try: + import vllm.transformers_utils.tokenizers.mistral as mistral_module + except Exception: + return + sys.modules.setdefault("vllm.tokenizers.mistral", mistral_module) + + +def _install_repo_utils_alias() -> None: + if "vllm.transformers_utils.repo_utils" in sys.modules: + return + try: + from vllm.transformers_utils.config import file_or_path_exists + except Exception: + return + module = types.ModuleType("vllm.transformers_utils.repo_utils") + module.file_or_path_exists = file_or_path_exists + sys.modules["vllm.transformers_utils.repo_utils"] = module + + +def _install_config_parser_aliases() -> None: + try: + import vllm.transformers_utils.config as config_module + except Exception: + return + + if not hasattr(config_module, "MistralConfigParser"): + + class MistralConfigParser: + pass + + config_module.MistralConfigParser = MistralConfigParser + + if not hasattr(config_module, "register_config_parser"): + + def register_config_parser(_name: str): + def decorator(cls): + return cls + + return decorator + + config_module.register_config_parser = register_config_parser + + +def _install_nemotron_h_auto_config() -> None: + try: + from transformers import AutoConfig + from vllm.transformers_utils.configs import NemotronHConfig + except Exception: + return + try: + AutoConfig.register("nemotron_h", NemotronHConfig) + except ValueError: + pass + + +def _register_easy_magpie_plugin() -> None: + try: + from vllm_plugin_easymagpie_omni import register + except Exception: + return + register() + + +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_renderer_aliases() -> None: + if "vllm.renderers" not in sys.modules: + renderers = types.ModuleType("vllm.renderers") + renderers.__path__ = [] # type: ignore[attr-defined] + + class BaseRenderer: + pass + + def merge_kwargs(*items: Any) -> dict[str, Any]: + merged: dict[str, Any] = {} + for item in items: + if item: + merged.update(dict(item)) + return merged + + def renderer_from_config(*_args: Any, **_kwargs: Any) -> None: + return None + + renderers.BaseRenderer = BaseRenderer + renderers.merge_kwargs = merge_kwargs + renderers.renderer_from_config = renderer_from_config + sys.modules["vllm.renderers"] = renderers + else: + renderers = sys.modules["vllm.renderers"] + if not hasattr(renderers, "__path__"): + renderers.__path__ = [] # type: ignore[attr-defined] + inputs = sys.modules.get("vllm.renderers.inputs") + if inputs is None: + inputs = types.ModuleType("vllm.renderers.inputs") + sys.modules["vllm.renderers.inputs"] = inputs + inputs.__package__ = "vllm.renderers.inputs" + inputs.__path__ = [] # type: ignore[attr-defined] + try: + from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt + except Exception: + try: + from vllm.inputs.data import EmbedsPrompt, TextPrompt, TokensPrompt + except Exception: + EmbedsPrompt = dict # type: ignore[assignment] + TextPrompt = dict # type: ignore[assignment] + TokensPrompt = dict # type: ignore[assignment] + + try: + from typing import TypedDict + + class EncoderDecoderDictPrompt(TypedDict): + encoder_prompt: Any + decoder_prompt: Any | None + + class EncoderDecoderTokPrompt(TypedDict): + encoder_prompt: Any + decoder_prompt: Any | None + + except Exception: + EncoderDecoderDictPrompt = dict # type: ignore[assignment] + EncoderDecoderTokPrompt = dict # type: ignore[assignment] + + for name, value in { + "EmbedsPrompt": EmbedsPrompt, + "TextPrompt": TextPrompt, + "TokensPrompt": TokensPrompt, + "DecoderDictPrompt": TextPrompt | TokensPrompt, + "DecoderOnlyDictPrompt": TextPrompt | TokensPrompt | EmbedsPrompt, + "DecoderOnlyTokPrompt": TokensPrompt | EmbedsPrompt, + "DecoderTokPrompt": TokensPrompt, + "DictPrompt": dict, + "EncoderDictPrompt": TextPrompt | TokensPrompt, + "EncoderTokPrompt": TokensPrompt, + "SingletonDictPrompt": dict, + "SingletonTokPrompt": dict, + "TokPrompt": dict, + "EncoderDecoderDictPrompt": EncoderDecoderDictPrompt, + "EncoderDecoderTokPrompt": EncoderDecoderTokPrompt, + }.items(): + if not hasattr(inputs, name): + setattr(inputs, name, value) + + preprocess = sys.modules.get("vllm.renderers.inputs.preprocess") + if preprocess is None: + preprocess = types.ModuleType("vllm.renderers.inputs.preprocess") + preprocess.__package__ = "vllm.renderers.inputs" + sys.modules["vllm.renderers.inputs.preprocess"] = preprocess + inputs.preprocess = preprocess + + def _is_list_of(value: object, item_type: type) -> bool: + return isinstance(value, list) and all(isinstance(item, item_type) for item in value) + + def _validate_prompt_dict(prompt: Any) -> None: + if "prompt" not in prompt or "prompt_token_ids" in prompt or "prompt_embeds" in prompt: + return + if not isinstance(prompt["prompt"], str): + raise TypeError("Prompt text should be a string") + + def _parse_singleton_prompt(prompt: object, *, allow_embeds: bool = True) -> Any: + if isinstance(prompt, str): + return TextPrompt(prompt=prompt) + if _is_list_of(prompt, int): + return TokensPrompt(prompt_token_ids=prompt) + if isinstance(prompt, dict): + if "encoder_prompt" in prompt: + raise TypeError("Cannot pass encoder-decoder prompt to a singleton prompt parser") + _validate_prompt_dict(prompt) + if "prompt" in prompt or "prompt_token_ids" in prompt or (allow_embeds and "prompt_embeds" in prompt): + return prompt + expected = "text, tokens, or embeddings" if allow_embeds else "text or tokens" + raise TypeError(f"Prompt dictionary must contain {expected}") + raise TypeError("Prompt should be a string, list of tokens, or dictionary") + + def parse_dec_only_prompt(prompt: object) -> Any: + if isinstance(prompt, dict) and "encoder_prompt" in prompt: + raise TypeError("Cannot pass encoder-decoder prompt to decoder-only models") + return _parse_singleton_prompt(prompt) + + def parse_enc_dec_prompt(prompt: object) -> Any: + if isinstance(prompt, dict) and "encoder_prompt" in prompt: + enc_prompt = prompt["encoder_prompt"] + dec_prompt = prompt.get("decoder_prompt") + else: + enc_prompt = prompt + dec_prompt = None + return EncoderDecoderDictPrompt( + encoder_prompt=_parse_singleton_prompt(enc_prompt, allow_embeds=False), + decoder_prompt=None if dec_prompt is None else _parse_singleton_prompt(dec_prompt, allow_embeds=False), + ) + + for name in ( + "DecoderDictPrompt", + "DecoderOnlyDictPrompt", + "DictPrompt", + "EncoderDictPrompt", + "EncoderDecoderDictPrompt", + "SingletonDictPrompt", + "TextPrompt", + "TokensPrompt", + "EmbedsPrompt", + ): + setattr(preprocess, name, getattr(inputs, name)) + preprocess.parse_dec_only_prompt = parse_dec_only_prompt + preprocess.parse_enc_dec_prompt = parse_enc_dec_prompt + elif not hasattr(inputs, "preprocess"): + inputs.preprocess = preprocess + + tokenize = sys.modules.get("vllm.renderers.inputs.tokenize") + if tokenize is None: + tokenize = types.ModuleType("vllm.renderers.inputs.tokenize") + tokenize.__package__ = "vllm.renderers.inputs" + sys.modules["vllm.renderers.inputs.tokenize"] = tokenize + inputs.tokenize = tokenize + for name in ( + "DecoderOnlyTokPrompt", + "DecoderTokPrompt", + "EncoderTokPrompt", + "EncoderDecoderTokPrompt", + "SingletonTokPrompt", + "TokPrompt", + "TokensPrompt", + "EmbedsPrompt", + ): + setattr(tokenize, name, getattr(inputs, name)) + elif not hasattr(inputs, "tokenize"): + inputs.tokenize = tokenize + + +def _install_input_processor_alias() -> None: + try: + from vllm.inputs.preprocess import InputPreprocessor + from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config + from vllm.v1.engine.processor import Processor + except Exception: + return + + original_init = InputPreprocessor.__init__ + if not getattr(original_init, "_easymagpie_compat", False): + + def compat_init(self, *args: Any, **kwargs: Any) -> None: + renderer = kwargs.pop("renderer", None) + if "vllm_config" in kwargs: + vllm_config = kwargs.pop("vllm_config") + tokenizer = kwargs.pop("tokenizer", None) + if tokenizer is None and not getattr(vllm_config.model_config, "skip_tokenizer_init", False): + tokenizer = cached_tokenizer_from_config(model_config=vllm_config.model_config) + original_init(self, vllm_config.model_config, tokenizer, *args, **kwargs) + self.renderer = renderer + return + original_init(self, *args, **kwargs) + self.renderer = getattr(self, "renderer", renderer) + + compat_init._easymagpie_compat = True # type: ignore[attr-defined] + InputPreprocessor.__init__ = compat_init + + if "vllm.v1.engine.input_processor" in sys.modules: + return + + module = types.ModuleType("vllm.v1.engine.input_processor") + + class InputProcessor(Processor): + def __init__(self, *, vllm_config: Any, tokenizer: Any = None, mm_registry: Any = None, **_kwargs: Any) -> None: + if tokenizer is None and not getattr(vllm_config.model_config, "skip_tokenizer_init", False): + tokenizer = cached_tokenizer_from_config(model_config=vllm_config.model_config) + if mm_registry is None: + super().__init__(vllm_config, tokenizer) + else: + super().__init__(vllm_config, tokenizer, mm_registry) + self.renderer = getattr(self.input_preprocessor, "renderer", None) + + module.InputProcessor = InputProcessor + sys.modules["vllm.v1.engine.input_processor"] = module + + +def _install_processor_process_inputs_compat() -> None: + try: + from vllm.v1.engine.processor import Processor + except Exception: + return + + original_process_inputs = Processor.process_inputs + if getattr(original_process_inputs, "_easymagpie_compat", False): + return + + try: + signature = inspect.signature(original_process_inputs) + parameters = signature.parameters.values() + accepts_supported_tasks = "supported_tasks" in signature.parameters or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters + ) + except Exception: + accepts_supported_tasks = False + + if accepts_supported_tasks: + return + + def compat_process_inputs(self, *args: Any, **kwargs: Any): + has_omni_supported_tasks = "supported_tasks" in kwargs + kwargs.pop("supported_tasks", None) + result = original_process_inputs(self, *args, **kwargs) + if has_omni_supported_tasks and isinstance(result, tuple) and len(result) == 2: + return result[1] + return result + + compat_process_inputs._easymagpie_compat = True # type: ignore[attr-defined] + Processor.process_inputs = compat_process_inputs + + +def _install_omni_input_preprocessor_signature_compat() -> None: + try: + from vllm_omni.inputs.preprocess import OmniInputPreprocessor + except Exception: + return + + original_prompt_to_inputs = OmniInputPreprocessor._prompt_to_llm_inputs + if getattr(original_prompt_to_inputs, "_easymagpie_compat", False): + return + + try: + signature = inspect.signature(original_prompt_to_inputs) + parameters = signature.parameters + accepts_kwargs = any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()) + except Exception: + signature = None + parameters = {} + accepts_kwargs = False + + if accepts_kwargs or {"lora_request", "return_mm_hashes"}.issubset(parameters): + return + + accepted_kwargs = { + name + for name, parameter in parameters.items() + if name != "self" and parameter.kind in {inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY} + } + + def compat_prompt_to_inputs(self, *args: Any, **kwargs: Any): + if signature is not None: + kwargs = {key: value for key, value in kwargs.items() if key in accepted_kwargs} + else: + kwargs.pop("lora_request", None) + kwargs.pop("return_mm_hashes", None) + return original_prompt_to_inputs(self, *args, **kwargs) + + compat_prompt_to_inputs._easymagpie_compat = True # type: ignore[attr-defined] + OmniInputPreprocessor._prompt_to_llm_inputs = compat_prompt_to_inputs + + +def _install_input_preprocessor_truncate_compat() -> None: + try: + from vllm.inputs.preprocess import InputPreprocessor + except Exception: + return + + if hasattr(InputPreprocessor, "_truncate_inputs"): + return + + def _truncate_inputs(self, prompt_token_ids: Any, tokenization_kwargs: dict[str, Any] | None = None): + limit = None + if tokenization_kwargs: + limit = tokenization_kwargs.get("max_length", tokenization_kwargs.get("truncate_prompt_tokens")) + if limit is None: + return prompt_token_ids + try: + limit_int = int(limit) + except Exception: + return prompt_token_ids + if limit_int <= 0: + return prompt_token_ids + return prompt_token_ids[-limit_int:] + + InputPreprocessor._truncate_inputs = _truncate_inputs + + +def _install_omni_engine_request_compat() -> None: + try: + from vllm.v1.engine import EngineCoreRequest + from vllm_omni.engine import OmniEngineCoreRequest + except Exception: + return + + for request_cls in (EngineCoreRequest, OmniEngineCoreRequest): + if not hasattr(request_cls, "prompt_embeds"): + request_cls.prompt_embeds = None + + try: + async_omni_engine = importlib.import_module("vllm_omni.engine.async_omni_engine") + from vllm_omni.engine.serialization import serialize_additional_information + except Exception: + return + + try: + request_signature = inspect.signature(OmniEngineCoreRequest) + except Exception: + request_signature = None + + if not getattr(async_omni_engine._upgrade_to_omni_request, "_easymagpie_compat", False): + + def compat_upgrade_to_omni_request(request: Any, raw_prompt: Any) -> Any: + prompt_embeds = getattr(request, "prompt_embeds", None) + additional_information = None + + if isinstance(raw_prompt, dict): + if prompt_embeds is None: + raw_prompt_embeds = raw_prompt.get("prompt_embeds") + try: + import torch + + if isinstance(raw_prompt_embeds, torch.Tensor): + prompt_embeds = raw_prompt_embeds + except Exception: + pass + additional_information = serialize_additional_information( + raw_prompt.get("additional_information"), + log_prefix="AsyncOmniEngine", + ) + + if prompt_embeds is None and additional_information is None: + return request + + if request_signature is None: + return request + + values: dict[str, Any] = {} + for name, parameter in request_signature.parameters.items(): + if name == "prompt_embeds": + values[name] = prompt_embeds + elif name == "additional_information": + values[name] = additional_information + elif hasattr(request, name): + values[name] = getattr(request, name) + elif parameter.default is inspect.Parameter.empty: + values[name] = None + + return OmniEngineCoreRequest(**values) + + compat_upgrade_to_omni_request._easymagpie_compat = True # type: ignore[attr-defined] + async_omni_engine._upgrade_to_omni_request = compat_upgrade_to_omni_request + + async_engine_cls = getattr(async_omni_engine, "AsyncOmniEngine", None) + original_build = getattr(async_engine_cls, "_build_add_request_message", None) + if original_build is None or getattr(original_build, "_easymagpie_external_req_compat", False): + return + try: + original_build_signature = inspect.signature(original_build) + except Exception: + original_build_signature = None + if original_build_signature is not None: + original_build_params = original_build_signature.parameters + if "prompt_text" in original_build_params and "message_type" in original_build_params: + # Newer vLLM-Omni already handles request construction, prompt text, + # final-stage metadata, and orchestrator admission correctly. The + # patched _upgrade_to_omni_request above is enough for EasyMagpie's + # prompt embeddings/additional-information fields. + return + + engine_core_request_cls = EngineCoreRequest + inject_global_id = getattr(async_omni_engine, "_inject_global_id", None) + + def compat_build_add_request_message( + self: Any, + request_id: str, + prompt: Any, + sampling_params_list: Any = None, + final_stage_id: int = 0, + arrival_time: float | None = None, + ) -> dict[str, Any]: + effective_sampling_params_list = ( + list(sampling_params_list) if sampling_params_list is not None else list(self.default_sampling_params_list) + ) + if not effective_sampling_params_list: + raise ValueError( + f"Missing sampling params for stage 0. Got {len(effective_sampling_params_list)} stage params." + ) + params = effective_sampling_params_list[0] + original_prompt = prompt + + stage_type = self.stage_metadata[0].get("stage_type") + if stage_type != "diffusion" and not isinstance(prompt, engine_core_request_cls): + if inject_global_id is not None: + if isinstance(prompt, dict): + inject_global_id(prompt, request_id) + elif isinstance(prompt, list): + for item in prompt: + inject_global_id(item, request_id) + + request = self.input_processor.process_inputs( + request_id=request_id, + prompt=prompt, + params=params, + supported_tasks=self.supported_tasks, + arrival_time=arrival_time, + ) + request = async_omni_engine._upgrade_to_omni_request(request, prompt) + try: + request.external_req_id = request_id + except AttributeError: + pass + + self.output_processors[0].add_request( + request=request, + prompt=prompt, + parent_req=None, + request_index=0, + queue=None, + ) + prompt = request + + return { + "type": "add_request", + "request_id": request_id, + "prompt": prompt, + "original_prompt": original_prompt, + "sampling_params_list": effective_sampling_params_list, + "final_stage_id": final_stage_id, + } + + compat_build_add_request_message._easymagpie_external_req_compat = True # type: ignore[attr-defined] + async_engine_cls._build_add_request_message = compat_build_add_request_message + + +def _install_omni_request_compat() -> None: + try: + from vllm.v1.request import StructuredOutputRequest + import vllm_omni.request as omni_request_module + except Exception: + return + + OmniRequest = getattr(omni_request_module, "OmniRequest", None) + if OmniRequest is None: + return + try: + BaseRequest = OmniRequest.__mro__[1] + except Exception: + return + + original_init = getattr(OmniRequest, "__init__", None) + if original_init is not None and not getattr(original_init, "_easymagpie_request_init_compat", False): + base_init = BaseRequest.__init__ + try: + base_signature = inspect.signature(base_init) + except Exception: + base_signature = None + + def compat_omni_request_init( + self: Any, + prompt_embeds: Any = None, + external_req_id: str | None = None, + additional_information: Any = None, + *args: Any, + **kwargs: Any, + ) -> None: + prompt_embeds_tensor = self._maybe_decode_prompt_embeds(prompt_embeds) + if base_signature is None: + base_init(self, *args, **kwargs) + else: + accepts_kwargs = any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in base_signature.parameters.values() + ) + base_kwargs = dict(kwargs) + if "prompt_embeds" in base_signature.parameters: + base_kwargs["prompt_embeds"] = prompt_embeds_tensor + if not accepts_kwargs: + base_kwargs = { + key: value + for key, value in base_kwargs.items() + if key in base_signature.parameters and key != "self" + } + base_init(self, *args, **base_kwargs) + + self.prompt_embeds = prompt_embeds_tensor + self.prompt_embeds_payload = prompt_embeds if prompt_embeds is not prompt_embeds_tensor else None + self.external_req_id = external_req_id + self.additional_information = additional_information + + compat_omni_request_init._easymagpie_request_init_compat = True # type: ignore[attr-defined] + OmniRequest.__init__ = compat_omni_request_init + + original_from_request = getattr(OmniRequest, "from_engine_core_request", None) + raw_from_request = getattr(original_from_request, "__func__", original_from_request) + if raw_from_request is None or getattr(raw_from_request, "_easymagpie_request_convert_compat", False): + return + + def compat_from_engine_core_request(cls: Any, request: Any, block_hasher: Any) -> Any: + mm_kwargs = getattr(request, "mm_kwargs", None) + if mm_kwargs is not None: + mm_kwargs = list(mm_kwargs) + elif getattr(request, "mm_features", None) is not None: + mm_kwargs = list(getattr(request, "mm_features")) + + sampling_params = getattr(request, "sampling_params", None) + structured_output_request = None + if sampling_params is not None: + from_sampling_params = getattr(StructuredOutputRequest, "from_sampling_params", None) + if callable(from_sampling_params): + structured_output_request = from_sampling_params(sampling_params) + else: + try: + structured_output_request = StructuredOutputRequest(sampling_params=sampling_params) + except TypeError: + structured_outputs = getattr(sampling_params, "structured_outputs", None) + if structured_outputs and not structured_outputs.all_constraints_none(): + structured_output_request = StructuredOutputRequest(params=structured_outputs) + + return cls( + request_id=request.request_id, + external_req_id=getattr(request, "external_req_id", None) or request.request_id, + client_index=getattr(request, "client_index", 0), + prompt_token_ids=request.prompt_token_ids, + prompt_embeds=getattr(request, "prompt_embeds", None), + multi_modal_kwargs=mm_kwargs, + multi_modal_hashes=getattr(request, "mm_hashes", None), + multi_modal_placeholders=getattr(request, "mm_placeholders", None), + sampling_params=sampling_params, + pooling_params=getattr(request, "pooling_params", None), + eos_token_id=getattr(request, "eos_token_id", None), + arrival_time=getattr(request, "arrival_time", None), + lora_request=getattr(request, "lora_request", None), + structured_output_request=structured_output_request, + cache_salt=getattr(request, "cache_salt", None), + priority=getattr(request, "priority", 0), + block_hasher=block_hasher, + additional_information=getattr(request, "additional_information", None), + ) + + compat_from_engine_core_request._easymagpie_request_convert_compat = True # type: ignore[attr-defined] + OmniRequest.from_engine_core_request = classmethod(compat_from_engine_core_request) + + +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_import_utils_alias() -> None: + if "vllm.utils.import_utils" in sys.modules: + return + module = types.ModuleType("vllm.utils.import_utils") + + class LazyLoader(types.ModuleType): + def __init__(self, local_name: str, parent_globals: dict[str, Any], name: str) -> None: + super().__init__(local_name) + self._local_name = local_name + self._parent_globals = parent_globals + self._module_name = name + self._module = None + + def _load(self): + if self._module is None: + self._module = importlib.import_module(self._module_name) + self._parent_globals[self._local_name] = self._module + return self._module + + def __getattr__(self, item: str) -> Any: + return getattr(self._load(), item) + + def resolve_obj_by_qualname(qualname: str) -> Any: + module_name, _, attr = qualname.replace(":", ".").rpartition(".") + if not module_name: + raise ValueError(f"Invalid qualified name: {qualname}") + obj = importlib.import_module(module_name) + for part in attr.split("."): + obj = getattr(obj, part) + return obj + + try: + from vllm.utils import import_pynvml + except Exception: + + def import_pynvml(): + return importlib.import_module("pynvml") + + module.LazyLoader = LazyLoader + module.resolve_obj_by_qualname = resolve_obj_by_qualname + module.import_pynvml = import_pynvml + sys.modules["vllm.utils.import_utils"] = module + + +def _install_torch_utils_alias() -> None: + if "vllm.utils.torch_utils" in sys.modules: + return + module = types.ModuleType("vllm.utils.torch_utils") + + def set_random_seed(seed: int | None) -> None: + import numpy as np + import torch + + seed = 0 if seed is None else int(seed) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + @contextlib.contextmanager + def set_default_torch_dtype(dtype): + import torch + + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(old_dtype) + + try: + from vllm.utils import supports_xccl + except Exception: + + def supports_xccl() -> bool: + return False + + module.set_random_seed = set_random_seed + module.set_default_torch_dtype = set_default_torch_dtype + module.supports_xccl = supports_xccl + sys.modules["vllm.utils.torch_utils"] = module + + +def _install_math_utils_alias() -> None: + if "vllm.utils.math_utils" in sys.modules: + return + try: + from vllm.utils import cdiv + except Exception: + + def cdiv(a: int, b: int) -> int: + return -(a // -b) + + module = types.ModuleType("vllm.utils.math_utils") + module.cdiv = cdiv + sys.modules["vllm.utils.math_utils"] = module + + +def _install_mem_utils_alias() -> None: + if "vllm.utils.mem_utils" in sys.modules: + return + try: + import vllm.utils as vllm_utils + except Exception: + return + + module = types.ModuleType("vllm.utils.mem_utils") + + class MemorySnapshot(vllm_utils.MemorySnapshot): + def __init__(self, *args, device=None, **kwargs) -> None: + super().__init__(*args, **kwargs) + + def format_gib(num_bytes: int | float) -> str: + return f"{float(num_bytes) / float(vllm_utils.GiB_bytes):.2f}" + + module.MemorySnapshot = MemorySnapshot + module.memory_profiling = vllm_utils.memory_profiling + module.format_gib = format_gib + sys.modules["vllm.utils.mem_utils"] = module + + +def _install_vllm_config_profiler_default() -> None: + try: + from vllm.config import VllmConfig + except Exception: + return + if hasattr(VllmConfig, "profiler_config"): + return + + def get_profiler_config(self): + stored = getattr(self, "_easymagpie_profiler_config", None) + if stored is not None: + return stored + additional_config = getattr(self, "additional_config", None) + if isinstance(additional_config, dict): + return additional_config.get("profiler_config") + return getattr(additional_config, "profiler_config", None) + + def set_profiler_config(self, value) -> None: + object.__setattr__(self, "_easymagpie_profiler_config", value) + + VllmConfig.profiler_config = property(get_profiler_config, set_profiler_config) + + +def _install_parallel_config_defaults() -> None: + try: + from vllm.config import ParallelConfig + except Exception: + return + + missing = object() + + def install_property(name: str, default): + if hasattr(ParallelConfig, name): + return + + storage_name = f"_easymagpie_{name}" + + def getter(self): + stored = getattr(self, storage_name, missing) + if stored is not missing: + return stored + return default(self) if callable(default) else default + + def setter(self, value) -> None: + object.__setattr__(self, storage_name, value) + + setattr(ParallelConfig, name, property(getter, setter)) + + install_property( + "nnodes_within_dp", + lambda self: max( + 1, + int(getattr(self, "data_parallel_size", 1)) // int(getattr(self, "data_parallel_size_local", 1)), + ), + ) + install_property("data_parallel_index", lambda self: getattr(self, "data_parallel_rank", 0)) + install_property( + "local_world_size", + lambda self: ( + int(getattr(self, "pipeline_parallel_size", 1)) + * int(getattr(self, "tensor_parallel_size", 1)) + * int(getattr(self, "data_parallel_size_local", 1)) + ), + ) + install_property("enable_dbo", False) + install_property("num_ubatches", 1) + install_property("use_ubatching", False) + + +def _install_cache_config_defaults() -> None: + try: + from vllm.config import CacheConfig + except Exception: + return + if hasattr(CacheConfig, "kv_cache_memory_bytes"): + return + + missing = object() + storage_name = "_easymagpie_kv_cache_memory_bytes" + + def get_kv_cache_memory_bytes(self): + stored = getattr(self, storage_name, missing) + if stored is not missing: + return stored + return None + + def set_kv_cache_memory_bytes(self, value) -> None: + object.__setattr__(self, storage_name, value) + + CacheConfig.kv_cache_memory_bytes = property(get_kv_cache_memory_bytes, set_kv_cache_memory_bytes) + + +def _install_platform_dtype_check() -> None: + try: + from vllm.platforms import current_platform + except Exception: + return + + dtype_check = getattr(current_platform, "check_if_supports_dtype", None) + if callable(dtype_check): + return + + def check_if_supports_dtype(_dtype) -> None: + return None + + try: + current_platform.check_if_supports_dtype = check_if_supports_dtype + except Exception: + pass + try: + setattr(type(current_platform), "check_if_supports_dtype", staticmethod(check_if_supports_dtype)) + except Exception: + pass + + +def _install_torch_accelerator_compat() -> None: + try: + import torch + except Exception: + return + + accelerator = getattr(torch, "accelerator", None) + if accelerator is None or hasattr(accelerator, "empty_cache"): + return + empty_cache = getattr(getattr(torch, "cuda", None), "empty_cache", None) + if not callable(empty_cache): + return + accelerator.empty_cache = empty_cache + + +def _install_logger_info_once_scope_compat() -> None: + try: + import vllm.logger as vllm_logger + except Exception: + vllm_logger = None + + if vllm_logger is not None: + original_print = getattr(vllm_logger, "_print_info_once", None) + if original_print is not None and not getattr(original_print, "_easymagpie_compat", False): + + def _print_info_once(logger, msg, *args, **kwargs): + kwargs.pop("scope", None) + return original_print(logger, msg, *args, **kwargs) + + _print_info_once._easymagpie_compat = True # type: ignore[attr-defined] + vllm_logger._print_info_once = _print_info_once + methods_to_patch = getattr(vllm_logger, "_METHODS_TO_PATCH", None) + if isinstance(methods_to_patch, dict): + methods_to_patch["info_once"] = _print_info_once + + original = getattr(logging.Logger, "info_once", None) + if original is None or getattr(original, "_easymagpie_compat", False): + return + + def info_once(self, msg, *args, **kwargs): + kwargs.pop("scope", None) + return original(self, msg, *args, **kwargs) + + info_once._easymagpie_compat = True # type: ignore[attr-defined] + logging.Logger.info_once = info_once + + +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 _install_cudagraph_mode_compat() -> None: + try: + from vllm.config import CUDAGraphMode + except Exception: + return + if hasattr(CUDAGraphMode.NONE, "valid_runtime_modes"): + return + + def valid_runtime_modes(self): + return self in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL) + + CUDAGraphMode.valid_runtime_modes = valid_runtime_modes # type: ignore[attr-defined] + + +def _install_cuda_graph_stat_alias() -> None: + try: + import vllm.compilation.cuda_graph as cuda_graph + except Exception: + return + if hasattr(cuda_graph, "CUDAGraphStat"): + return + + class CUDAGraphStat: + pass + + cuda_graph.CUDAGraphStat = CUDAGraphStat + + +def _install_cuda_piecewise_no_sym_shape_compat() -> None: + try: + import vllm.compilation.cuda_piecewise_backend as cuda_piecewise_backend + except Exception: + return + backend_cls = getattr(cuda_piecewise_backend, "PiecewiseBackend", None) + if backend_cls is None: + return + original = getattr(backend_cls, "__call__", None) + if original is None or getattr(original, "_easymagpie_no_sym_shape_compat", False): + return + + def _runtime_shape_from_args(args: tuple[Any, ...]) -> int | None: + for arg in args: + shape = getattr(arg, "shape", None) + if shape is None or len(shape) == 0: + continue + try: + return int(shape[0]) + except Exception: + continue + return None + + def _signature_from_args(args: tuple[Any, ...]) -> tuple[Any, ...]: + signature: list[Any] = [] + for arg in args: + shape = getattr(arg, "shape", None) + stride = getattr(arg, "stride", None) + dtype = getattr(arg, "dtype", None) + device = getattr(arg, "device", None) + if shape is None: + signature.append((type(arg).__qualname__, repr(arg))) + continue + try: + stride_value = tuple(int(x) for x in stride()) if callable(stride) else None + except Exception: + stride_value = None + try: + shape_value = tuple(int(x) for x in shape) + except Exception: + shape_value = tuple(shape) + signature.append( + ( + "tensor", + shape_value, + stride_value, + str(dtype), + str(device), + ) + ) + return tuple(signature) + + def _compile_no_sym_shape(self, args: tuple[Any, ...], runtime_shape: int | None): + compile_fn = getattr(getattr(self, "vllm_backend", None), "compiler_manager", None) + compile_fn = getattr(compile_fn, "compile", None) + if not callable(compile_fn): + return None + return compile_fn( + self.graph, + args, + self.compilation_config.inductor_compile_config, + self.compilation_config, + graph_index=self.piecewise_compile_index, + num_graphs=self.total_piecewise_compiles, + runtime_shape=runtime_shape, + ) + + def __call__(self, *args): + if getattr(self, "sym_shape_indices", None): + return original(self, *args) + signature = _signature_from_args(args) + runtime_shape = _runtime_shape_from_args(args) + runnables = getattr(self, "_easymagpie_no_sym_shape_runnables", None) + if runnables is None: + runnables = {} + self._easymagpie_no_sym_shape_runnables = runnables + if not getattr(self, "first_run_finished", False): + self.first_run_finished = True + runnables[signature] = self.compiled_graph_for_general_shape + check_for_ending_compilation = getattr(self, "check_for_ending_compilation", None) + if callable(check_for_ending_compilation): + check_for_ending_compilation() + self._easymagpie_no_sym_shape_compile_count = 0 + self._easymagpie_no_sym_shape_last_runtime_shape = runtime_shape + self._easymagpie_no_sym_shape_last_signature = signature + return self.compiled_graph_for_general_shape(*args) + runnable = runnables.get(signature) + if runnable is None: + runnable = _compile_no_sym_shape(self, args, runtime_shape) + if runnable is None: + runnable = self.compiled_graph_for_general_shape + runnables[signature] = runnable + self._easymagpie_no_sym_shape_compile_count = int( + getattr(self, "_easymagpie_no_sym_shape_compile_count", 0) + ) + 1 + self._easymagpie_no_sym_shape_last_runtime_shape = runtime_shape + self._easymagpie_no_sym_shape_last_signature = signature + return runnable(*args) + + __call__._easymagpie_no_sym_shape_compat = True # type: ignore[attr-defined] + __call__._easymagpie_original = original # type: ignore[attr-defined] + backend_cls.__call__ = __call__ + + +def _install_kv_connector_stats_alias() -> None: + module_name = "vllm.distributed.kv_transfer.kv_connector.v1.metrics" + if module_name in sys.modules: + return + module = types.ModuleType(module_name) + + class KVConnectorStats: + def aggregate(self, _other): + return self + + module.KVConnectorStats = KVConnectorStats + sys.modules[module_name] = module + + +def _install_perf_stats_alias() -> None: + module_name = "vllm.v1.metrics.perf" + if module_name in sys.modules: + return + module = types.ModuleType(module_name) + + class PerfStats: + pass + + module.PerfStats = PerfStats + sys.modules[module_name] = module + + +def _install_scheduler_make_stats_compat() -> None: + try: + from vllm.v1.core.sched.scheduler import Scheduler + except Exception: + return + original = getattr(Scheduler, "make_stats", None) + if original is None or getattr(original, "_easymagpie_compat", False): + return + + def make_stats(self, spec_decoding_stats=None, *_args, **_kwargs): + return original(self, spec_decoding_stats) + + make_stats._easymagpie_compat = True # type: ignore[attr-defined] + Scheduler.make_stats = make_stats + + +def _install_sched_utils_aliases() -> None: + try: + import vllm.v1.core.sched.utils as sched_utils + except Exception: + return + if not hasattr(sched_utils, "remove_all"): + + def remove_all(items, removed): + removed_set = set(removed) + return [item for item in items if item not in removed_set] + + sched_utils.remove_all = remove_all + + +def _install_sched_interface_aliases() -> None: + try: + import vllm.v1.core.sched.interface as sched_interface + except Exception: + return + if hasattr(sched_interface, "PauseState"): + return + + from enum import Enum + + class PauseState(Enum): + UNPAUSED = 0 + PAUSED_ALL = 1 + + sched_interface.PauseState = PauseState + + +def _install_output_processor_signature_compat() -> None: + try: + from vllm.v1.engine.output_processor import OutputProcessor, RequestState + except Exception: + return + + original_init = getattr(OutputProcessor, "__init__", None) + if original_init is not None and not getattr(original_init, "_easymagpie_compat", False): + + def __init__(self, tokenizer, log_stats: bool, stream_interval: int = 1, tracing_enabled: bool = False): + original_init(self, tokenizer=tokenizer, log_stats=log_stats) + self.stream_interval = int(stream_interval) + self.tracing_enabled = bool(tracing_enabled) + if not hasattr(self, "external_req_ids"): + self.external_req_ids = defaultdict(list) + + __init__._easymagpie_compat = True # type: ignore[attr-defined] + OutputProcessor.__init__ = __init__ + + original_from_new_descriptor = RequestState.__dict__.get("from_new_request") + original_from_new = ( + original_from_new_descriptor.__func__ + if isinstance(original_from_new_descriptor, classmethod) + else original_from_new_descriptor + ) + if original_from_new is None or getattr(original_from_new, "_easymagpie_compat", False): + return + try: + original_from_new_signature = inspect.signature(original_from_new) + except Exception: + original_from_new_signature = None + + def from_new_request(cls, *args, stream_interval: int = 1, **kwargs): + if original_from_new_signature is not None: + param_names = list(original_from_new_signature.parameters) + if "stream_interval" in param_names and "stream_interval" not in kwargs: + stream_interval_arg_index = param_names.index("stream_interval") - 1 + if len(args) <= stream_interval_arg_index: + kwargs["stream_interval"] = stream_interval + state = original_from_new(cls, *args, **kwargs) + if not hasattr(state, "stream_interval"): + state.stream_interval = int(stream_interval) + if not hasattr(state, "sent_tokens_offset"): + state.sent_tokens_offset = 0 + request = kwargs.get("request") + if request is None: + for arg in args: + if hasattr(arg, "request_id") and hasattr(arg, "prompt_token_ids"): + request = arg + break + external_req_id = getattr(request, "external_req_id", None) or getattr(request, "request_id", None) + if external_req_id is not None and not hasattr(state, "external_req_id"): + state.external_req_id = external_req_id + return state + + from_new_request._easymagpie_compat = True # type: ignore[attr-defined] + RequestState.from_new_request = classmethod(from_new_request) + + +def _install_flash_attention_builder_compat() -> None: + try: + import vllm.v1.attention.backends.flash_attn as v1_flash_attn + from vllm.attention.backends.flash_attn import FlashAttentionBackend as LegacyBackend + from vllm.attention.backends.flash_attn import FlashAttentionMetadataBuilder as LegacyBuilder + from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadataBuilder as V1Builder + except Exception: + return + + original_get_sliding_window_configs = getattr(v1_flash_attn, "_get_sliding_window_configs", None) + if original_get_sliding_window_configs is not None and not getattr( + original_get_sliding_window_configs, "_easymagpie_compat", False + ): + + def _get_sliding_window_configs(vllm_config): + sliding_window_configs = set() + layers = v1_flash_attn.get_layers_from_vllm_config(vllm_config, v1_flash_attn.Attention) + for layer in layers.values(): + impl = getattr(layer, "impl", None) + sliding_window_configs.add(getattr(impl, "sliding_window", None)) + return sliding_window_configs + + _get_sliding_window_configs._easymagpie_compat = True # type: ignore[attr-defined] + v1_flash_attn._get_sliding_window_configs = _get_sliding_window_configs + + original_get_builder_cls = getattr(LegacyBackend, "get_builder_cls", None) + if original_get_builder_cls is None or getattr(original_get_builder_cls, "_easymagpie_compat", False): + return + try: + current_builder = LegacyBackend.get_builder_cls() + except Exception: + current_builder = None + if current_builder is not LegacyBuilder: + return + + def get_builder_cls(): + return V1Builder + + get_builder_cls._easymagpie_compat = True # type: ignore[attr-defined] + LegacyBackend.get_builder_cls = staticmethod(get_builder_cls) + + +def _install_worker_utils_compat() -> None: + try: + import vllm.v1.worker.utils as worker_utils + from vllm.utils import GiB_bytes + except Exception: + return + if hasattr(worker_utils, "request_memory"): + request_memory = worker_utils.request_memory + else: + + def request_memory(init_snapshot, cache_config): + requested_memory = init_snapshot.total_memory * cache_config.gpu_memory_utilization + if init_snapshot.free_memory < requested_memory: + gib = lambda value: round(value / GiB_bytes, 2) + raise ValueError( + f"Free memory on device ({gib(init_snapshot.free_memory)}/" + f"{gib(init_snapshot.total_memory)} GiB) on startup is less " + f"than desired GPU memory utilization ({cache_config.gpu_memory_utilization}, " + f"{gib(requested_memory)} GiB). Decrease GPU memory utilization " + "or reduce GPU memory used by other processes." + ) + return requested_memory + + worker_utils.request_memory = request_memory + + if not hasattr(worker_utils, "is_residual_scattered_for_sp"): + + def is_residual_scattered_for_sp(*_args, **_kwargs) -> bool: + return False + + worker_utils.is_residual_scattered_for_sp = is_residual_scattered_for_sp + + +def _install_worker_workspace_alias() -> None: + if "vllm.v1.worker.workspace" in sys.modules: + return + module = types.ModuleType("vllm.v1.worker.workspace") + + def init_workspace_manager(*_args, **_kwargs): + return None + + module.init_workspace_manager = init_workspace_manager + sys.modules["vllm.v1.worker.workspace"] = module + + +def _install_gpu_ar_worker_default_attrs() -> None: + try: + module = importlib.import_module("vllm_omni.worker.gpu_ar_worker") + except Exception: + return + worker_cls = getattr(module, "GPUARWorker", None) + if worker_cls is None or hasattr(worker_cls, "use_v2_model_runner"): + return + # vLLM-Omni still relies on the v1 runner hooks for EasyMagpie. Newer vLLM + # worker bases no longer create this flag, so default the class lookup to + # the same v1 path the worker forces when the flag is present. + worker_cls.use_v2_model_runner = False + + +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 + + +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}", + } + + 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] + 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_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 + + +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_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_easy_magpie_refit_rpc_compat() + _install_v1_serial_utils_dense_tensor_compat() + _install_async_omni_client_compat() + + +def _install_forward_context_kwargs_compat() -> None: + try: + import vllm.forward_context as forward_context + except Exception: + return + original = getattr(forward_context, "set_forward_context", None) + if original is None or getattr(original, "_easymagpie_compat", False): + return + original_signature = inspect.signature(original) + supported_kwargs = set(original_signature.parameters) + + def set_forward_context(*args, **kwargs): + context_kwargs = dict(kwargs) + filtered_kwargs = {key: value for key, value in kwargs.items() if key in supported_kwargs} + bound = original_signature.bind_partial(*args, **filtered_kwargs) + if "attn_metadata" in bound.arguments: + for key, value in bound.arguments.items(): + if key not in ("attn_metadata", "vllm_config"): + context_kwargs.setdefault(key, value) + attn_metadata = _normalize_mamba2_attention_metadata_groups(bound.arguments["attn_metadata"]) + attn_metadata = _synthetic_profile_attention_metadata_from_context( + attn_metadata, + bound.arguments.get("vllm_config"), + context_kwargs, + ) + bound.arguments["attn_metadata"] = _as_v1_mamba2_attention_metadata_dict( + attn_metadata, + bound.arguments.get("vllm_config"), + ) + return original(*bound.args, **bound.kwargs) + + set_forward_context._easymagpie_compat = True # type: ignore[attr-defined] + forward_context.set_forward_context = set_forward_context + for module_name in ( + "vllm_omni.worker.gpu_model_runner", + "vllm_omni.worker.gpu_ar_model_runner", + "vllm_omni.worker.gpu_generation_model_runner", + ): + module = sys.modules.get(module_name) + if module is not None and getattr(module, "set_forward_context", None) is original: + module.set_forward_context = set_forward_context + + +def _install_omni_gpu_model_runner_method_compat() -> None: + try: + from vllm.config import CUDAGraphMode + module = importlib.import_module("vllm_omni.worker.gpu_model_runner") + except Exception: + return + _install_cudagraph_mode_compat() + runner_cls = getattr(module, "OmniGPUModelRunner", None) + if runner_cls is None: + return + if not hasattr(runner_cls, "enable_prompt_embeds"): + runner_cls.enable_prompt_embeds = False + if not hasattr(runner_cls, "uses_xdrope_dim"): + runner_cls.uses_xdrope_dim = 0 + if not hasattr(runner_cls, "kv_cache_config"): + + def kv_cache_config(self): + return _as_v1_kv_cache_config(getattr(self, "_kv_cache_config", getattr(self, "cache_config", None))) + + def set_kv_cache_config(self, value): + self._kv_cache_config = value + + def del_kv_cache_config(self): + if hasattr(self, "_kv_cache_config"): + delattr(self, "_kv_cache_config") + + runner_cls.kv_cache_config = property(kv_cache_config, set_kv_cache_config, del_kv_cache_config) + + if not hasattr(runner_cls, "_determine_batch_execution_and_padding"): + + class _BatchDescriptor: + def __init__(self, *, num_tokens: int, num_reqs: int, uniform_decode: bool = False) -> None: + self.num_tokens = num_tokens + self.num_reqs = num_reqs + self.uniform_decode = uniform_decode + + def _determine_batch_execution_and_padding(self, **kwargs): + num_tokens = int(kwargs["num_tokens"]) + num_reqs = int(kwargs["num_reqs"]) + force_uniform_decode = bool(kwargs.get("force_uniform_decode", False)) + if hasattr(self, "get_dp_padding"): + num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) + num_tokens += int(num_pad) + else: + num_tokens_across_dp = None + batch_desc = _BatchDescriptor( + num_tokens=num_tokens, + num_reqs=num_reqs, + uniform_decode=force_uniform_decode, + ) + return CUDAGraphMode.NONE, batch_desc, False, num_tokens_across_dp, None + + runner_cls._determine_batch_execution_and_padding = _determine_batch_execution_and_padding + + if not hasattr(runner_cls, "synchronize_input_prep"): + + def synchronize_input_prep(self): + return contextlib.nullcontext() + + runner_cls.synchronize_input_prep = synchronize_input_prep + + if not hasattr(runner_cls, "_register_layerwise_nvtx_hooks"): + + def _register_layerwise_nvtx_hooks(self): + return None + + runner_cls._register_layerwise_nvtx_hooks = _register_layerwise_nvtx_hooks + + original_build_attention_metadata = getattr(runner_cls, "_build_attention_metadata", None) + if original_build_attention_metadata is not None and not getattr( + original_build_attention_metadata, "_easymagpie_compat", False + ): + original_build_attention_metadata_signature = inspect.signature(original_build_attention_metadata) + + def _build_attention_metadata(self, *args, **kwargs): + bound = original_build_attention_metadata_signature.bind_partial(self, *args, **kwargs) + result = original_build_attention_metadata(*bound.args, **bound.kwargs) + attn_metadata = result[0] if isinstance(result, tuple) and result else result + if isinstance(attn_metadata, dict) and not attn_metadata: + num_reqs = bound.arguments.get("num_reqs") + num_tokens = bound.arguments.get("num_tokens", 0) + max_query_len = bound.arguments.get("max_query_len", num_tokens) + if num_reqs is not None: + num_reqs = int(num_reqs) + self._easymagpie_mamba_cache_batch_size = int(num_reqs) + attn_metadata = _build_profile_flash_attention_metadata( + self, + num_tokens=int(num_tokens), + num_reqs=num_reqs, + max_query_len=int(max_query_len), + ) + if isinstance(result, tuple): + result = (attn_metadata, *result[1:]) + else: + result = attn_metadata + else: + normalized_attn_metadata = _normalize_mamba2_attention_metadata_groups(attn_metadata) + if normalized_attn_metadata is not attn_metadata: + attn_metadata = normalized_attn_metadata + if isinstance(result, tuple): + result = (attn_metadata, *result[1:]) + else: + result = attn_metadata + num_reqs = bound.arguments.get("num_reqs") + if num_reqs is not None: + self._easymagpie_mamba_cache_batch_size = int(num_reqs) + elif hasattr(self, "_easymagpie_mamba_cache_batch_size"): + delattr(self, "_easymagpie_mamba_cache_batch_size") + return result + + _build_attention_metadata._easymagpie_compat = True # type: ignore[attr-defined] + runner_cls._build_attention_metadata = _build_attention_metadata + + original_init_model_kwargs = getattr(runner_cls, "_init_model_kwargs", None) + if original_init_model_kwargs is not None and not getattr( + original_init_model_kwargs, "_easymagpie_compat", False + ): + try: + signature = inspect.signature(original_init_model_kwargs) + parameters = list(signature.parameters.values()) + accepts_num_tokens = any( + parameter.kind == inspect.Parameter.VAR_POSITIONAL for parameter in parameters + ) or (len(parameters) >= 2) + except Exception: + accepts_num_tokens = True + + def _call_base_init_model_kwargs(self, num_tokens): + for base_cls in getattr(runner_cls, "__mro__", ())[1:]: + base_init_model_kwargs = base_cls.__dict__.get("_init_model_kwargs") + if base_init_model_kwargs is None: + continue + if getattr(base_init_model_kwargs, "_easymagpie_compat", False): + continue + if base_init_model_kwargs is original_init_model_kwargs: + continue + try: + base_signature = inspect.signature(base_init_model_kwargs) + base_parameters = list(base_signature.parameters.values()) + base_accepts_num_tokens = any( + parameter.kind == inspect.Parameter.VAR_POSITIONAL for parameter in base_parameters + ) or (len(base_parameters) >= 2) + except Exception: + base_accepts_num_tokens = True + if base_accepts_num_tokens: + try: + return base_init_model_kwargs(self, num_tokens) + except TypeError: + pass + return base_init_model_kwargs(self) + return None + + def _init_model_kwargs(self, num_tokens=None): + if num_tokens is None: + num_tokens = int(getattr(self, "max_num_tokens", 0) or 0) + if accepts_num_tokens: + try: + model_kwargs = original_init_model_kwargs(self, num_tokens) + except TypeError: + model_kwargs = _call_base_init_model_kwargs(self, num_tokens) + if model_kwargs is None: + model_kwargs = original_init_model_kwargs(self) + else: + try: + model_kwargs = original_init_model_kwargs(self) + except TypeError: + model_kwargs = _call_base_init_model_kwargs(self, num_tokens) + if model_kwargs is None: + raise + cache_batch_size = getattr(self, "_easymagpie_mamba_cache_batch_size", None) + if cache_batch_size is not None: + model_kwargs = dict(model_kwargs) + model_kwargs.setdefault("easymagpie_mamba_cache_batch_size", int(cache_batch_size)) + request_cache_kwargs = _easymagpie_materialize_request_cache_kwargs( + self, + getattr(self, "_easymagpie_mamba_request_cache_state", None) + or getattr(self, "_easymagpie_mamba_request_cache_kwargs", None), + ) + if request_cache_kwargs: + model_kwargs = dict(model_kwargs) + model_kwargs.update(request_cache_kwargs) + return model_kwargs + + _init_model_kwargs._easymagpie_compat = True # type: ignore[attr-defined] + runner_cls._init_model_kwargs = _init_model_kwargs + + original_get_cumsum_and_arange = getattr(runner_cls, "_get_cumsum_and_arange", None) + if original_get_cumsum_and_arange is not None and not getattr( + original_get_cumsum_and_arange, "_easymagpie_compat", False + ): + try: + cumsum_signature = inspect.signature(original_get_cumsum_and_arange) + cumsum_parameters = list(cumsum_signature.parameters.values()) + requires_arange_out = any( + parameter.name == "arange_out" and parameter.default is inspect.Parameter.empty + for parameter in cumsum_parameters + ) + except Exception: + requires_arange_out = False + + def _new_arange_out(self, num_tokens): + import numpy as np + + total_num_tokens = int(np.asarray(num_tokens).sum()) + arange_np = getattr(self, "arange_np", None) + dtype = getattr(arange_np, "dtype", None) or getattr(num_tokens, "dtype", None) or np.int64 + return np.empty(total_num_tokens, dtype=dtype) + + def _get_cumsum_and_arange(self, num_tokens, arange_out=None, *args, **kwargs): + if not requires_arange_out: + if arange_out is None: + return original_get_cumsum_and_arange(self, num_tokens, *args, **kwargs) + return original_get_cumsum_and_arange(self, num_tokens, arange_out, *args, **kwargs) + + if arange_out is not None: + return original_get_cumsum_and_arange(self, num_tokens, arange_out, *args, **kwargs) + + generated_arange_out = _new_arange_out(self, num_tokens) + result = original_get_cumsum_and_arange( + self, + num_tokens, + generated_arange_out, + *args, + **kwargs, + ) + if isinstance(result, tuple): + return result + try: + total_num_tokens = int(result[-1]) + except Exception: + total_num_tokens = int(generated_arange_out.shape[0]) + return result, generated_arange_out[:total_num_tokens] + + _get_cumsum_and_arange._easymagpie_compat = True # type: ignore[attr-defined] + runner_cls._get_cumsum_and_arange = _get_cumsum_and_arange + + def _easymagpie_input_batch_req_ids(self) -> list[str]: + return [ + str(req_id) + for req_id in getattr(getattr(self, "input_batch", None), "req_ids", []) + if req_id is not None + ] + + def _easymagpie_request_ids_to_seq_ids(self, request_ids: list[str]) -> dict[str, list[int]]: + input_batch = getattr(self, "input_batch", None) + req_id_to_index = getattr(input_batch, "req_id_to_index", None) + used_seq_ids: set[int] = set() + next_fallback = 0 + mapped: dict[str, list[int]] = {} + + def fallback_seq_id() -> int: + nonlocal next_fallback + while next_fallback in used_seq_ids: + next_fallback += 1 + seq_idx = next_fallback + used_seq_ids.add(seq_idx) + next_fallback += 1 + return seq_idx + + if isinstance(req_id_to_index, dict): + missing_or_duplicate: list[str] = [] + seq_by_req_id: dict[str, int] = {} + for req_id in request_ids: + raw_idx = req_id_to_index.get(req_id) + if raw_idx is None: + raw_idx = req_id_to_index.get(str(req_id)) + try: + seq_idx = int(raw_idx) + except (TypeError, ValueError): + missing_or_duplicate.append(str(req_id)) + else: + if seq_idx in used_seq_ids: + missing_or_duplicate.append(str(req_id)) + else: + used_seq_ids.add(seq_idx) + seq_by_req_id[str(req_id)] = seq_idx + for req_id in missing_or_duplicate: + seq_by_req_id[req_id] = fallback_seq_id() + for req_id in request_ids: + mapped[str(req_id)] = [seq_by_req_id[str(req_id)]] + return mapped + for req_id in request_ids: + mapped[str(req_id)] = [fallback_seq_id()] + return mapped + + def _easymagpie_request_cache_state(self, scheduler_output: Any) -> dict[str, Any]: + num_scheduled_tokens = getattr(scheduler_output, "num_scheduled_tokens", None) + if not num_scheduled_tokens: + return {} + + scheduled_counts = { + str(req_id): int(count or 0) + for req_id, count in num_scheduled_tokens.items() + if int(count or 0) > 0 + } + if not scheduled_counts: + return {} + + finished_requests_ids = getattr(scheduler_output, "finished_req_ids", None) + if finished_requests_ids is None: + finished_requests_ids = getattr(scheduler_output, "finished_requests_ids", None) + return { + "scheduled_request_ids": list(scheduled_counts), + "scheduled_counts": scheduled_counts, + "finished_requests_ids": ( + None + if finished_requests_ids is None + else sorted(str(req_id) for req_id in finished_requests_ids) + ), + "previous_active_request_ids": sorted( + str(req_id) + for req_id in getattr(self, "_easymagpie_active_mamba_request_ids", set()) + ), + "pre_update_request_ids": _easymagpie_input_batch_req_ids(self), + } + + def _easymagpie_materialize_request_cache_kwargs(self, state: Any) -> dict[str, Any]: + if not state: + return {} + if "request_ids_to_seq_ids" in state: + return state + + scheduled_counts = { + str(req_id): int(count or 0) + for req_id, count in dict(state.get("scheduled_counts", {})).items() + if int(count or 0) > 0 + } + if not scheduled_counts: + return {} + + ordered_request_ids: list[str] = [] + seen: set[str] = set() + for req_id in _easymagpie_input_batch_req_ids(self): + if req_id in scheduled_counts and req_id not in seen: + ordered_request_ids.append(req_id) + seen.add(req_id) + for req_id in state.get("scheduled_request_ids", scheduled_counts): + req_id = str(req_id) + if req_id in scheduled_counts and req_id not in seen: + ordered_request_ids.append(req_id) + seen.add(req_id) + if not ordered_request_ids: + return {} + + current_active = set(str(req_id) for req_id in state.get("pre_update_request_ids", [])) + current_active.update(_easymagpie_input_batch_req_ids(self)) + current_active.update(ordered_request_ids) + previous_active = set(str(req_id) for req_id in state.get("previous_active_request_ids", [])) + implicit_finished = previous_active - current_active + finished_requests_ids = state.get("finished_requests_ids") + if finished_requests_ids is None: + finished_requests_ids = sorted(implicit_finished) + else: + finished_requests_ids = sorted( + set(str(req_id) for req_id in finished_requests_ids) | implicit_finished + ) + self._easymagpie_active_mamba_request_ids = current_active + return { + "request_ids_to_seq_ids": _easymagpie_request_ids_to_seq_ids(self, ordered_request_ids), + "finished_requests_ids": finished_requests_ids, + } + + original_model_forward = getattr(runner_cls, "_model_forward", None) + if original_model_forward is not None and not getattr(original_model_forward, "_easymagpie_compat", False): + + def _model_forward(self, *args, **kwargs): + request_cache_kwargs = _easymagpie_materialize_request_cache_kwargs( + self, + getattr(self, "_easymagpie_mamba_request_cache_state", None) + or getattr(self, "_easymagpie_mamba_request_cache_kwargs", None), + ) + if request_cache_kwargs: + kwargs = dict(kwargs) + for key, value in request_cache_kwargs.items(): + kwargs.setdefault(key, value) + return original_model_forward(self, *args, **kwargs) + + _model_forward._easymagpie_compat = True # type: ignore[attr-defined] + runner_cls._model_forward = _model_forward + + def _install_execute_model_request_cache_compat(target_cls: type[Any]) -> None: + original_execute_model = getattr(target_cls, "execute_model", None) + if original_execute_model is None or getattr(original_execute_model, "_easymagpie_compat", False): + return + original_execute_model_signature = inspect.signature(original_execute_model) + + def execute_model(self, *args, **kwargs): + bound = original_execute_model_signature.bind_partial(self, *args, **kwargs) + scheduler_output = bound.arguments.get("scheduler_output") + request_cache_state = ( + _easymagpie_request_cache_state(self, scheduler_output) + if scheduler_output is not None + else {} + ) + if request_cache_state: + self._easymagpie_mamba_request_cache_state = request_cache_state + elif hasattr(self, "_easymagpie_mamba_request_cache_state"): + delattr(self, "_easymagpie_mamba_request_cache_state") + try: + return original_execute_model(*bound.args, **bound.kwargs) + finally: + for name in ( + "_easymagpie_mamba_request_cache_state", + "_easymagpie_mamba_request_cache_kwargs", + ): + if hasattr(self, name): + delattr(self, name) + + execute_model._easymagpie_compat = True # type: ignore[attr-defined] + target_cls.execute_model = execute_model + + _install_execute_model_request_cache_compat(runner_cls) + for module_name, class_names in ( + ("vllm_omni.worker.gpu_ar_model_runner", ("GPUARModelRunner",)), + ("vllm_omni.worker.gpu_generation_model_runner", ("GPUGenerationModelRunner",)), + ): + extra_module = sys.modules.get(module_name) + if extra_module is None: + continue + for class_name in class_names: + extra_runner_cls = getattr(extra_module, class_name, None) + if extra_runner_cls is not None: + _install_execute_model_request_cache_compat(extra_runner_cls) + + original_lora_context = getattr(runner_cls, "maybe_dummy_run_with_lora", None) + if original_lora_context is not None and not getattr(original_lora_context, "_easymagpie_compat", False): + + def maybe_dummy_run_with_lora(self, lora_config, num_scheduled_tokens, *_args, **_kwargs): + if lora_config is None: + return contextlib.nullcontext() + return original_lora_context(self, lora_config, num_scheduled_tokens) + + maybe_dummy_run_with_lora._easymagpie_compat = True # type: ignore[attr-defined] + runner_cls.maybe_dummy_run_with_lora = maybe_dummy_run_with_lora + + original_randomize_context = getattr(runner_cls, "maybe_randomize_inputs", None) + if original_randomize_context is not None and not getattr( + original_randomize_context, "_easymagpie_compat", False + ): + + def maybe_randomize_inputs(self, input_ids, inputs_embeds=None): + if input_ids is None: + return contextlib.nullcontext() + try: + return original_randomize_context(self, input_ids, inputs_embeds) + except TypeError: + return original_randomize_context(self, input_ids) + + maybe_randomize_inputs._easymagpie_compat = True # type: ignore[attr-defined] + runner_cls.maybe_randomize_inputs = maybe_randomize_inputs + + original_dummy_run = getattr(runner_cls, "_dummy_run", None) + if original_dummy_run is not None and not getattr(original_dummy_run, "_easymagpie_force_attn_compat", False): + original_dummy_run_signature = inspect.signature(original_dummy_run) + original_dummy_run_parameters = original_dummy_run_signature.parameters + original_dummy_run_has_cudagraph_mode = "cudagraph_runtime_mode" in original_dummy_run_parameters + original_dummy_run_var_keyword = next( + ( + name + for name, parameter in original_dummy_run_parameters.items() + if parameter.kind == inspect.Parameter.VAR_KEYWORD + ), + None, + ) + + def get_bound_keyword(bound, name: str, default: Any = None): + if name in bound.arguments: + return bound.arguments[name] + if original_dummy_run_var_keyword is not None: + kwargs = bound.arguments.get(original_dummy_run_var_keyword, {}) + if isinstance(kwargs, dict) and name in kwargs: + return kwargs[name] + return default + + def set_bound_keyword(bound, name: str, value: Any) -> None: + parameter = original_dummy_run_parameters.get(name) + if parameter is not None and parameter.kind != inspect.Parameter.VAR_KEYWORD: + bound.arguments[name] = value + return + if original_dummy_run_var_keyword is not None: + kwargs = dict(bound.arguments.get(original_dummy_run_var_keyword, {})) + kwargs[name] = value + bound.arguments[original_dummy_run_var_keyword] = kwargs + return + bound.arguments[name] = value + + def needs_mamba_attention_metadata(self) -> bool: + candidates = [getattr(self, "model", None)] + seen: set[int] = set() + while candidates: + candidate = candidates.pop() + if candidate is None: + continue + candidate_id = id(candidate) + if candidate_id in seen: + continue + seen.add(candidate_id) + class_name = candidate.__class__.__name__.lower() + config = getattr(candidate, "config", None) + pattern = getattr(config, "hybrid_override_pattern", None) + architectures = getattr(config, "architectures", None) or () + if "easymagpie" in class_name or "easy_magpie" in class_name: + return True + if "nemotronh" in class_name or "nemotron_h" in class_name: + return True + if isinstance(pattern, str) and "M" in pattern and hasattr(config, "chunk_size"): + return True + if any( + "easymagpie" in str(architecture).lower() + or "nemotronh" in str(architecture).lower() + or "nemotron_h" in str(architecture).lower() + for architecture in architectures + ): + return True + for attr_name in ("backbone", "model", "language_model", "llm", "decoder"): + child = getattr(candidate, attr_name, None) + if child is not None: + candidates.append(child) + return False + + def ensure_easy_mamba_no_compile_layers(self) -> None: + vllm_config = getattr(self, "vllm_config", None) + compilation_config = getattr(vllm_config, "compilation_config", None) + static_context = getattr(compilation_config, "static_forward_context", None) + if not isinstance(static_context, dict): + return + backbone = getattr(getattr(self, "model", None), "backbone", None) + layers = getattr(backbone, "layers", None) + if layers is None: + return + for layer_idx, layer in enumerate(layers): + mixer = getattr(layer, "mixer", None) + if mixer is not None: + static_context.setdefault(f"backbone.layers.{layer_idx}.mixer", mixer) + + def _dummy_run(self, *args, **kwargs): + bound = original_dummy_run_signature.bind_partial(self, *args, **kwargs) + force_v1_env = False + if needs_mamba_attention_metadata(self): + ensure_easy_mamba_no_compile_layers(self) + num_tokens = get_bound_keyword(bound, "num_tokens") + max_num_reqs = getattr(getattr(self, "scheduler_config", None), "max_num_seqs", None) + if num_tokens is not None and max_num_reqs is not None: + self._easymagpie_mamba_cache_batch_size = min(max(1, int(num_tokens)), max(1, int(max_num_reqs))) + use_v1_dummy_run = original_dummy_run_has_cudagraph_mode + if not use_v1_dummy_run and original_dummy_run_var_keyword is not None: + try: + from vllm import envs + + use_v1_dummy_run = bool(getattr(envs, "VLLM_USE_V1", False)) + except Exception: + use_v1_dummy_run = False + if use_v1_dummy_run: + set_bound_keyword(bound, "cudagraph_runtime_mode", CUDAGraphMode.NONE) + set_bound_keyword(bound, "force_attention", False) + force_v1_env = True + elif not get_bound_keyword(bound, "force_attention", False): + set_bound_keyword(bound, "force_attention", True) + if not force_v1_env: + return original_dummy_run(*bound.args, **bound.kwargs) + try: + from vllm import envs + except Exception: + return original_dummy_run(*bound.args, **bound.kwargs) + previous_use_v1 = getattr(envs, "VLLM_USE_V1", None) + envs.VLLM_USE_V1 = True + try: + return original_dummy_run(*bound.args, **bound.kwargs) + finally: + envs.VLLM_USE_V1 = previous_use_v1 + + _dummy_run._easymagpie_force_attn_compat = True # type: ignore[attr-defined] + runner_cls._dummy_run = _dummy_run + + +def _install_omni_worker_import_patch() -> None: + if getattr(builtins, _IMPORT_PATCH_FLAG, False): + return + original_import = builtins.__import__ + + def import_with_omni_worker_patch(name, globals=None, locals=None, fromlist=(), level=0): + should_patch_preprocess = ( + name == "vllm_omni.inputs.preprocess" + or name.startswith("vllm_omni.inputs.preprocess.") + or (name == "vllm_omni.inputs" and "preprocess" in fromlist) + ) + if should_patch_preprocess: + _install_vllm_inputs_data_alias() + _install_vllm_multimodal_inputs_alias() + module = original_import(name, globals, locals, fromlist, level) + should_patch = ( + name == "vllm_omni.worker.gpu_model_runner" + or name == "vllm_omni.worker.gpu_ar_model_runner" + or name == "vllm_omni.worker.gpu_generation_model_runner" + or name.startswith("vllm_omni.worker.gpu_model_runner.") + or name.startswith("vllm_omni.worker.gpu_ar_model_runner.") + or name.startswith("vllm_omni.worker.gpu_generation_model_runner.") + or (name == "vllm_omni.worker" and "gpu_model_runner" in fromlist) + or ( + name == "vllm_omni.worker" + and any(item in fromlist for item in ("gpu_ar_model_runner", "gpu_generation_model_runner")) + ) + ) + if should_patch: + _install_forward_context_kwargs_compat() + _install_omni_gpu_model_runner_method_compat() + should_patch_worker = ( + name == "vllm_omni.worker.gpu_ar_worker" + or name == "vllm_omni.worker.gpu_generation_worker" + or name.startswith("vllm_omni.worker.gpu_ar_worker.") + or name.startswith("vllm_omni.worker.gpu_generation_worker.") + or ( + name == "vllm_omni.worker" + and any(item in fromlist for item in ("gpu_ar_worker", "gpu_generation_worker")) + ) + ) + if should_patch_worker: + _install_easy_magpie_refit_rpc_compat() + if should_patch_preprocess: + _install_omni_input_preprocessor_signature_compat() + should_patch_async_engine = ( + name == "vllm_omni.engine.async_omni_engine" + or name.startswith("vllm_omni.engine.async_omni_engine.") + or (name == "vllm_omni.engine" and "async_omni_engine" in fromlist) + ) + if should_patch_async_engine: + _install_omni_engine_request_compat() + should_patch_request = name == "vllm_omni.request" or name.startswith("vllm_omni.request.") + if should_patch_request: + _install_omni_request_compat() + return module + + builtins.__import__ = import_with_omni_worker_patch + setattr(builtins, _IMPORT_PATCH_FLAG, True) + + +def _install_ec_transfer_alias() -> None: + if "vllm.distributed.ec_transfer" in sys.modules: + return + module = types.ModuleType("vllm.distributed.ec_transfer") + + class _DisabledECTransfer: + is_consumer = False + is_producer = False + + def has_ec_transfer() -> bool: + return False + + def get_ec_transfer(): + return _DisabledECTransfer() + + module.has_ec_transfer = has_ec_transfer + module.get_ec_transfer = get_ec_transfer + sys.modules["vllm.distributed.ec_transfer"] = module + + +def _install_routed_experts_capturer_alias() -> None: + module_name = "vllm.model_executor.layers.fused_moe.routed_experts_capturer" + if module_name in sys.modules: + return + module = types.ModuleType(module_name) + + class RoutedExpertsCapturer: + @classmethod + def get_instance(cls): + return None + + def clear_buffer(self) -> None: + return None + + module.RoutedExpertsCapturer = RoutedExpertsCapturer + sys.modules[module_name] = module + + +def _install_structured_output_aliases() -> None: + try: + import vllm.v1.core.sched.output as sched_output + except Exception: + sched_output = None + if sched_output is not None and not hasattr(sched_output, "GrammarOutput"): + sched_output.GrammarOutput = object + + try: + import vllm.v1.structured_output.utils as structured_utils + except Exception: + return + if hasattr(structured_utils, "apply_grammar_bitmask"): + return + + def apply_grammar_bitmask(*_args, **_kwargs): + return None + + structured_utils.apply_grammar_bitmask = apply_grammar_bitmask + + +def _install_outputs_aliases() -> None: + try: + import vllm.v1.outputs as outputs + except Exception: + return + if not hasattr(outputs, "AsyncModelRunnerOutput") and hasattr(outputs, "ModelRunnerOutput"): + outputs.AsyncModelRunnerOutput = outputs.ModelRunnerOutput + if not hasattr(outputs, "make_empty_encoder_model_runner_output"): + + def make_empty_encoder_model_runner_output(*_args, **_kwargs): + return outputs.EMPTY_MODEL_RUNNER_OUTPUT + + outputs.make_empty_encoder_model_runner_output = make_empty_encoder_model_runner_output + + +def _install_spec_decode_aliases() -> None: + class _UnsupportedSpecDecodeProposer: + def __init__(self, *args, **kwargs) -> None: + raise NotImplementedError("This vLLM build does not expose the legacy spec-decode proposer.") + + aliases = { + "vllm.v1.spec_decode.draft_model": ("DraftModelProposer", _UnsupportedSpecDecodeProposer), + "vllm.v1.spec_decode.extract_hidden_states": ( + "ExtractHiddenStatesProposer", + _UnsupportedSpecDecodeProposer, + ), + } + for module_name, (class_name, cls) in aliases.items(): + if module_name in sys.modules: + continue + module = types.ModuleType(module_name) + setattr(module, class_name, cls) + sys.modules[module_name] = module + + +def _install_v1_utils_aliases() -> None: + try: + import vllm.v1.utils as v1_utils + except Exception: + return + if hasattr(v1_utils, "record_function_or_nullcontext"): + return + + def record_function_or_nullcontext(name: str): + try: + import torch + + return torch.profiler.record_function(name) + except Exception: + return contextlib.nullcontext() + + v1_utils.record_function_or_nullcontext = record_function_or_nullcontext + + +def _install_gpu_model_runner_aliases() -> None: + try: + import vllm.v1.worker.gpu_model_runner as gpu_model_runner + except Exception: + return + if not hasattr(gpu_model_runner, "PerLayerAttnMetadata"): + gpu_model_runner.PerLayerAttnMetadata = dict + if hasattr(gpu_model_runner, "AsyncGPUModelRunnerOutput"): + return + + class AsyncGPUModelRunnerOutput: + def __init__( + self, + model_runner_output=None, + sampled_token_ids=None, + logprobs_tensors=None, + invalid_req_indices=None, + async_output_copy_stream=None, + vocab_size=None, + ) -> None: + self.model_runner_output = model_runner_output + self.sampled_token_ids = sampled_token_ids + self.logprobs_tensors = logprobs_tensors + self.invalid_req_indices = invalid_req_indices + self.async_output_copy_stream = async_output_copy_stream + self.vocab_size = vocab_size + self.sampled_token_ids_cpu = sampled_token_ids + self.async_copy_ready_event = None + + gpu_model_runner.AsyncGPUModelRunnerOutput = AsyncGPUModelRunnerOutput + + +def _install_model_interface_aliases() -> None: + try: + import vllm.model_executor.models.interfaces as interfaces + except Exception: + return + if hasattr(interfaces, "supports_mrope"): + return + + def supports_mrope(*_args, **_kwargs) -> bool: + return False + + interfaces.supports_mrope = supports_mrope + + +def _install_ubatch_utils_alias() -> None: + module_name = "vllm.v1.worker.ubatch_utils" + if module_name in sys.modules: + return + module = types.ModuleType(module_name) + + def maybe_create_ubatch_slices(*_args, **_kwargs): + return None, None + + module.maybe_create_ubatch_slices = maybe_create_ubatch_slices + sys.modules[module_name] = module + + +def _install_tracing_instrument_alias() -> None: + try: + import vllm.tracing as tracing + except Exception: + return + if hasattr(tracing, "instrument"): + return + + def instrument(func=None, **_kwargs): + def decorator(inner): + return inner + + if callable(func): + return decorator(func) + return decorator + + tracing.instrument = instrument + + +def _install_executor_alias() -> None: + try: + import vllm.v1.executor as executor_pkg + from vllm.v1.executor.abstract import Executor + except Exception: + return + if not hasattr(executor_pkg, "Executor"): + executor_pkg.Executor = Executor + + +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_omni_model_config_field_compat() -> None: + try: + from vllm_omni.config.model import OmniModelConfig + except Exception: + return + if getattr(OmniModelConfig, "_easymagpie_field_compat", False): + return + + omni_defaults = { + "stage_id": 0, + "async_chunk": False, + "model_stage": "thinker", + "model_arch": None, + "worker_type": None, + "engine_output_type": None, + "hf_config_name": None, + "custom_process_next_stage_input_func": None, + "stage_connector_config": lambda: {"name": "SharedMemoryConnector", "extra": {}}, + "omni_kv_config": None, + "codec_frame_rate_hz": None, + "task_type": None, + "io_processor_plugin": None, + "has_sampling_extra_args": False, + "subtalker_sampling_params": None, + } + + def add_defaults_to_omni_kwargs(cls, omni_kwargs): + for key, default in omni_defaults.items(): + if key not in omni_kwargs: + omni_kwargs[key] = default() if callable(default) else default + + def _validate_omni_fields(cls, **omni_kwargs): + unexpected = set(omni_kwargs) - set(omni_defaults) + if unexpected: + raise ValueError(f"Unexpected omni kwarg: {sorted(unexpected)}") + + OmniModelConfig.add_defaults_to_omni_kwargs = classmethod(add_defaults_to_omni_kwargs) + OmniModelConfig._validate_omni_fields = classmethod(_validate_omni_fields) + if not hasattr(OmniModelConfig, "io_processor_plugin"): + OmniModelConfig.io_processor_plugin = None + if not hasattr(OmniModelConfig, "has_sampling_extra_args"): + OmniModelConfig.has_sampling_extra_args = False + if not hasattr(OmniModelConfig, "subtalker_sampling_params"): + OmniModelConfig.subtalker_sampling_params = None + OmniModelConfig._easymagpie_field_compat = True + + +def install_vllm_omni_compat() -> None: + """Install import/signature shims required by vLLM-Omni on older vLLM.""" + + global _INSTALLED + if _INSTALLED: + _install_ray_objectref_aliases() + _install_ray_placement_group_aliases() + _install_vllm_inputs_data_alias() + _install_vllm_multimodal_inputs_alias() + return + _install_config_decorator_compat() + _install_mamba2_metadata_compat() + _install_triton_attention_profile_metadata_compat() + _install_flash_attention_profile_metadata_compat() + _install_flashinfer_attention_profile_metadata_compat() + _install_mamba2_profile_no_kv_cache_compat() + _install_lora_config_alias() + _install_model_arch_convertor() + _install_io_processor_stub() + _install_tokenizer_aliases() + _install_repo_utils_alias() + _install_config_parser_aliases() + _install_nemotron_h_auto_config() + _install_vllm_inputs_data_alias() + _install_vllm_multimodal_inputs_alias() + _install_renderer_aliases() + _install_input_processor_alias() + _install_input_preprocessor_truncate_compat() + _install_omni_input_preprocessor_signature_compat() + _install_omni_engine_request_compat() + _install_omni_request_compat() + _install_processor_process_inputs_compat() + _install_engine_utils_compat() + _install_import_utils_alias() + _install_torch_utils_alias() + _install_math_utils_alias() + _install_mem_utils_alias() + _install_vllm_config_profiler_default() + _install_parallel_config_defaults() + _install_cache_config_defaults() + _install_platform_dtype_check() + _install_torch_accelerator_compat() + _install_logger_info_once_scope_compat() + _install_v1_serial_utils_dense_tensor_compat() + _install_cudagraph_mode_compat() + _install_cuda_graph_stat_alias() + _install_cuda_piecewise_no_sym_shape_compat() + _install_kv_connector_stats_alias() + _install_perf_stats_alias() + _install_scheduler_make_stats_compat() + _install_sched_utils_aliases() + _install_sched_interface_aliases() + _install_output_processor_signature_compat() + _install_flash_attention_builder_compat() + _install_worker_utils_compat() + _install_worker_workspace_alias() + _install_gpu_ar_worker_default_attrs() + _install_easy_magpie_refit_rpc_compat() + _install_forward_context_kwargs_compat() + _install_omni_gpu_model_runner_method_compat() + _install_omni_worker_import_patch() + _install_ec_transfer_alias() + _install_routed_experts_capturer_alias() + _install_structured_output_aliases() + _install_outputs_aliases() + _install_spec_decode_aliases() + _install_v1_utils_aliases() + _install_model_interface_aliases() + _install_gpu_model_runner_aliases() + _install_ubatch_utils_alias() + # vLLM-Omni's runner import may fail until the late aliases above exist. + _install_forward_context_kwargs_compat() + _install_omni_gpu_model_runner_method_compat() + _install_tracing_instrument_alias() + _install_executor_alias() + _install_ray_objectref_aliases() + _install_ray_placement_group_aliases() + _install_omni_model_config_field_compat() + _install_async_omni_client_compat() + _register_easy_magpie_plugin() + _INSTALLED = True + + +__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..5fa16e967118 --- /dev/null +++ b/examples/tts/easymagpie_vllm_omni/tests/test_easymagpie_backbone_cache.py @@ -0,0 +1,287 @@ +# 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._debug_combined_pre_norm = torch.zeros(max_tokens) + model._debug_hidden_norm = torch.zeros(max_tokens) + model._debug_outputs_enabled = False + model._debug_backbone_layers_enabled = False + model._debug_backbone_active_tokens = 0 + model.code_predictor = SimpleNamespace(debug_collect_logits=False) + 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..1fedb6f6fa54 --- /dev/null +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -0,0 +1,48 @@ +# 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 +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 == [] 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..1d342caf024d --- /dev/null +++ b/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_sampler.py @@ -0,0 +1,115 @@ +# 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.debug_collect_logits = False + 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") From fe05e9ece5c4f81e6975cf97fe56ffe5e93a7f02 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 14:44:07 +0000 Subject: [PATCH 02/14] Trim EasyMagpie RL compat modes --- .../easymagpie_vllm_omni/vllm_compat.py | 5020 +---------------- .../tests/test_import_modes.py | 41 + 2 files changed, 193 insertions(+), 4868 deletions(-) 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 index c4d753379a29..028961762975 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -11,3528 +11,178 @@ # 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. -"""Compatibility shims for running vLLM-Omni against older vLLM builds.""" +"""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 builtins -import contextlib -import dataclasses import importlib import inspect import logging import os import queue as stdlib_queue -import random import sys import types -from collections import defaultdict -from typing import Any, Literal +from typing import Any -_INSTALLED = False -_IMPORT_PATCH_FLAG = "_easymagpie_vllm_compat_import_patch_installed" logger = logging.getLogger(__name__) -try: - import torch as _TORCH_MODULE -except Exception: - _TORCH_MODULE = None - -try: - from vllm.attention.backends.flash_attn import FlashAttentionMetadata as _LEGACY_FLASH_ATTENTION_METADATA_CLS -except Exception: - _LEGACY_FLASH_ATTENTION_METADATA_CLS = None - -try: - from vllm.attention.backends.xformers import XFormersMetadata as _XFORMERS_METADATA_CLS -except Exception: - _XFORMERS_METADATA_CLS = None - -try: - from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata as _V1_FLASH_ATTENTION_METADATA_CLS -except Exception: - _V1_FLASH_ATTENTION_METADATA_CLS = None - -try: - from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata as _V1_TRITON_ATTENTION_METADATA_CLS -except Exception: - _V1_TRITON_ATTENTION_METADATA_CLS = None - -try: - from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata as _V1_MAMBA2_ATTENTION_METADATA_CLS -except Exception: - _V1_MAMBA2_ATTENTION_METADATA_CLS = None - -_SUPPORTED_KWARGS_CACHE: dict[int, set[str] | None] = {} - - -def _as_v1_kv_cache_config(cache_config: Any) -> Any: - if cache_config is None or hasattr(cache_config, "kv_cache_groups"): - return cache_config - cached = getattr(cache_config, "_easymagpie_v1_kv_cache_config", None) - if cached is None: - try: - from vllm.v1.kv_cache_interface import KVCacheConfig - - cached = KVCacheConfig(num_blocks=1, kv_cache_tensors=[], kv_cache_groups=[]) - except Exception: - cached = types.SimpleNamespace(num_blocks=1, kv_cache_tensors=[], kv_cache_groups=[]) - try: - cache_config._easymagpie_v1_kv_cache_config = cached - except Exception: - pass - return cached - - -def _buffer_slice_to_int_list(buffer: Any, length: int) -> list[int] | None: - if buffer is None: - return None - candidates = ( - getattr(buffer, "cpu", None), - getattr(buffer, "np", None), - getattr(buffer, "gpu", None), - buffer, - ) - for candidate in candidates: - if candidate is None: - continue - try: - values = candidate[:length] - if hasattr(values, "detach"): - values = values.detach().cpu().tolist() - elif hasattr(values, "tolist"): - values = values.tolist() - return [int(v) for v in values] - except Exception: - continue - return None - - -def _balanced_profile_lengths(num_tokens: int, num_reqs: int) -> list[int]: - num_reqs = max(1, int(num_reqs)) - base = max(1, int(num_tokens) // num_reqs) - remainder = max(0, int(num_tokens) - base * num_reqs) - lengths = [base] * num_reqs - lengths[-1] += remainder - return lengths - - -def _cache_supported_kwargs(callable_obj: Any) -> set[str] | None: - cache_key = id(callable_obj) - if cache_key in _SUPPORTED_KWARGS_CACHE: - return _SUPPORTED_KWARGS_CACHE[cache_key] - try: - signature = inspect.signature(callable_obj) - except Exception: - _SUPPORTED_KWARGS_CACHE[cache_key] = None - return None - if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values()): - _SUPPORTED_KWARGS_CACHE[cache_key] = None - return None - supported = { - name - for name, parameter in signature.parameters.items() - if parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) - } - _SUPPORTED_KWARGS_CACHE[cache_key] = supported - return supported - - -def _prime_supported_kwargs_cache() -> None: - for callable_obj in ( - _LEGACY_FLASH_ATTENTION_METADATA_CLS, - _XFORMERS_METADATA_CLS, - _V1_FLASH_ATTENTION_METADATA_CLS, - _V1_TRITON_ATTENTION_METADATA_CLS, - _V1_MAMBA2_ATTENTION_METADATA_CLS, - ): - if callable_obj is not None: - _cache_supported_kwargs(callable_obj) - - -def _call_with_supported_kwargs(callable_obj: Any, kwargs: dict[str, Any]) -> Any: - supported = _cache_supported_kwargs(callable_obj) - if supported is None: - return callable_obj(**kwargs) - return callable_obj(**{key: value for key, value in kwargs.items() if key in supported}) - - -_prime_supported_kwargs_cache() - - -def _max_int_sequence(values: Any) -> int: - max_value = 0 - for value in values: - int_value = int(value) - if int_value > max_value: - max_value = int_value - return max_value - - -def _is_torch_dynamo_compiling() -> bool: - try: - torch_module = _TORCH_MODULE - dynamo = getattr(torch_module, "_dynamo", None) - is_compiling = getattr(dynamo, "is_compiling", None) - return bool(is_compiling is not None and is_compiling()) - except Exception: - return False - - -def _attach_profile_attention_metadata_attrs( - metadata: Any, - *, - num_tokens: int, - num_reqs: int, - query_lens: list[int], - seq_lens: list[int], - seq_lens_tensor: Any, - seq_start_loc: Any, - query_start_loc: Any = None, - context_lens_tensor: Any = None, - block_tables: Any = None, - slot_mapping: Any = None, - max_query_len: int | None = None, -) -> Any: - if query_start_loc is None: - query_start_loc = seq_start_loc - resolved_max_query_len = max(int(max_query_len or 0), _max_int_sequence(query_lens)) - for name, value in ( - ("num_prefills", int(num_reqs)), - ("num_prefill_tokens", int(num_tokens)), - ("num_decodes", 0), - ("num_decode_tokens", 0), - ("query_lens", list(query_lens)), - ("seq_lens", list(seq_lens)), - ("seq_lens_tensor", seq_lens_tensor), - ("seq_start_loc", seq_start_loc), - ("query_start_loc", query_start_loc), - ("max_query_len", resolved_max_query_len), - ("max_prefill_seq_len", _max_int_sequence(seq_lens)), - ("max_decode_query_len", 0), - ("max_decode_seq_len", 0), - ("context_lens_tensor", context_lens_tensor), - ("block_tables", block_tables), - ("slot_mapping", slot_mapping), - ("prefill_metadata", metadata), - ("decode_metadata", None), - ): - try: - setattr(metadata, name, value) - except Exception: - pass - return metadata - - -def _is_mamba2_attention_metadata(value: Any) -> bool: - common = all( - hasattr(value, attr) - for attr in ( - "num_prefills", - "num_prefill_tokens", - "num_decodes", - "num_decode_tokens", - "seq_lens", - "chunk_size", - ) - ) - if not common: - return False - legacy_state = all(hasattr(value, attr) for attr in ("query_start_loc", "state_indices_tensor")) - split_state = all( - hasattr(value, attr) - for attr in ("query_start_loc_p", "state_indices_tensor_p", "query_start_loc_d", "state_indices_tensor_d") - ) - return legacy_state or split_state - - -def _first_not_none(*values: Any) -> Any: - for value in values: - if value is not None: - return value - return None - - -def _legacy_mamba2_attention_metadata_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: - return { - "num_prefills": kwargs.get("num_prefills", 0), - "num_prefill_tokens": kwargs.get("num_prefill_tokens", 0), - "num_decodes": kwargs.get("num_decodes", 0), - "num_decode_tokens": kwargs.get("num_decode_tokens", 0), - "query_start_loc": _first_not_none( - kwargs.get("query_start_loc"), - kwargs.get("legacy_query_start_loc"), - kwargs.get("query_start_loc_p"), - kwargs.get("query_start_loc_d"), - ), - "seq_lens": kwargs["seq_lens"], - "prep_initial_states": kwargs.get("prep_initial_states", False), - "chunk_size": kwargs.get("chunk_size", 0), - "has_initial_states_p": kwargs.get("has_initial_states_p"), - "seq_idx_p": kwargs.get("seq_idx_p"), - "chunk_indices_p": kwargs.get("chunk_indices_p"), - "chunk_offsets_p": kwargs.get("chunk_offsets_p"), - "state_indices_tensor": _first_not_none( - kwargs.get("state_indices_tensor"), - kwargs.get("legacy_state_indices_tensor"), - kwargs.get("state_indices_tensor_p"), - kwargs.get("state_indices_tensor_d"), - ), - "nums_dict": kwargs.get("nums_dict"), - "cu_seqlen": kwargs.get("cu_seqlen"), - "batch_ptr": kwargs.get("batch_ptr"), - "token_chunk_offset_ptr": kwargs.get("token_chunk_offset_ptr"), - } - - -def _filter_cached_supported_kwargs(metadata_cls: Any, kwargs: dict[str, Any]) -> dict[str, Any]: - supported = _cache_supported_kwargs(metadata_cls) - if supported is None: - return dict(kwargs) - return {key: value for key, value in kwargs.items() if key in supported} - - -def _new_mamba2_attention_metadata(metadata_cls: Any, kwargs: dict[str, Any]) -> Any: - constructor_kwargs = { - key: value - for key, value in kwargs.items() - if key not in {"legacy_query_start_loc", "legacy_state_indices_tensor"} - } - supported = _cache_supported_kwargs(metadata_cls) - if supported is not None and "query_start_loc_p" not in supported and "query_start_loc" in supported: - return metadata_cls(**_filter_cached_supported_kwargs(metadata_cls, _legacy_mamba2_attention_metadata_kwargs(kwargs))) - try: - return metadata_cls(**_filter_cached_supported_kwargs(metadata_cls, constructor_kwargs)) - except TypeError as exc: - message = str(exc) - if "unexpected keyword argument" in message and "query_start_loc_p" in kwargs: - legacy_kwargs = _legacy_mamba2_attention_metadata_kwargs(kwargs) - try: - return metadata_cls(**_filter_cached_supported_kwargs(metadata_cls, legacy_kwargs)) - except TypeError as legacy_exc: - if "num_reqs" not in kwargs or "num_reqs" not in str(legacy_exc): - raise - legacy_kwargs["num_reqs"] = kwargs["num_reqs"] - return metadata_cls(**_filter_cached_supported_kwargs(metadata_cls, legacy_kwargs)) - if "num_reqs" not in kwargs or "num_reqs" not in message: - raise - fallback_kwargs = _filter_cached_supported_kwargs(metadata_cls, constructor_kwargs) - fallback_kwargs.pop("num_reqs", None) - return metadata_cls(**fallback_kwargs) - - -def _resize_mamba_state_indices(state_indices: Any, needed_state_indices: int, device: Any) -> Any: - import torch - - if getattr(state_indices, "numel", lambda: 0)() > needed_state_indices: - return state_indices.to(device=device, dtype=torch.int32)[:needed_state_indices] - - current_numel = int(getattr(state_indices, "numel", lambda: 0)()) - if current_numel > 0: - base = state_indices.to(device=device, dtype=torch.int32).reshape(-1) - repeats = (needed_state_indices + current_numel - 1) // current_numel - return base.repeat(repeats)[:needed_state_indices] - - return torch.arange(needed_state_indices, dtype=torch.int32, device=device) - - -def _replace_mamba2_attention_metadata_fields(value: Any, replacements: dict[str, Any]) -> Any: - if not replacements: - return value - if _is_torch_dynamo_compiling(): - return value - try: - if dataclasses.is_dataclass(value): - supported_fields = getattr(type(value), "__dataclass_fields__", {}) - supported_replacements = { - key: replacement - for key, replacement in replacements.items() - if not supported_fields or key in supported_fields - } - if supported_replacements: - return dataclasses.replace(value, **supported_replacements) - except Exception: - pass - fields = dict(getattr(value, "__dict__", {})) - fields.update(replacements) - return types.SimpleNamespace(**fields) - - -def _balanced_tensor_lengths(total: int, count: int, device: Any) -> Any: - import torch - - count = max(0, int(count)) - if count <= 0: - return torch.empty(0, dtype=torch.int32, device=device) - base, extra = divmod(max(0, int(total)), count) - lengths = torch.full((count,), base, dtype=torch.int32, device=device) - if extra: - lengths[:extra] += 1 - return lengths - - -def _infer_mamba2_prefill_query_start_loc(value: Any, device: Any) -> Any: - import torch - - num_prefills = int(getattr(value, "num_prefills", 0) or 0) - num_prefill_tokens = int(getattr(value, "num_prefill_tokens", 0) or 0) - if num_prefills <= 0: - return torch.zeros(1, dtype=torch.int32, device=device) - - candidates = [] - query_start_loc_p = getattr(value, "query_start_loc_p", None) - if query_start_loc_p is not None and hasattr(query_start_loc_p, "numel"): - candidates.append(query_start_loc_p) - - query_start_loc = getattr(value, "query_start_loc", None) - if query_start_loc is not None and hasattr(query_start_loc, "numel"): - query_start_loc = query_start_loc.to(device=device, dtype=torch.int64) - if int(query_start_loc.numel()) >= num_prefills + 1: - num_decode_tokens = int(getattr(value, "num_decode_tokens", 0) or 0) - candidates.append(query_start_loc[-num_prefills - 1 :] - num_decode_tokens) - candidates.append(query_start_loc[: num_prefills + 1]) - - for candidate in candidates: - try: - candidate = candidate.to(device=device, dtype=torch.int64) - if int(candidate.numel()) != num_prefills + 1: - continue - if _is_torch_dynamo_compiling(): - return candidate.to(dtype=torch.int32) - if int(candidate[0].item()) != 0 or int(candidate[-1].item()) != num_prefill_tokens: - continue - lengths = candidate[1:] - candidate[:-1] - if bool((lengths >= 0).all().item()): - return candidate.to(dtype=torch.int32) - except Exception: - continue - - lengths = _balanced_tensor_lengths(num_prefill_tokens, num_prefills, device) - query_start_loc_p = torch.zeros(num_prefills + 1, dtype=torch.int32, device=device) - if num_prefills > 0: - query_start_loc_p[1:] = torch.cumsum(lengths, dim=0) - return query_start_loc_p - - -def _uses_mamba2_varlen_chunk_metadata(value: Any) -> bool: - return hasattr(value, "cu_chunk_seqlen_p") or hasattr(value, "last_chunk_indices_p") - - -def _compute_mamba2_varlen_chunk_metadata(query_start_loc: Any, chunk_size: int) -> tuple[Any, Any, Any]: - import torch - - try: - module = importlib.import_module("vllm.v1.attention.backends.mamba2_attn") - compute = getattr(module, "compute_varlen_chunk_metadata", None) - if compute is not None: - return compute(query_start_loc, chunk_size) - except Exception: - pass - - query_start_loc = query_start_loc.to(dtype=torch.int64) - device = query_start_loc.device - if int(query_start_loc.numel()) == 0: - zero = torch.tensor([0], dtype=torch.int32, device=device) - return zero, torch.empty(0, dtype=torch.int32, device=device), torch.empty(0, dtype=torch.int32, device=device) - - chunk_size = int(chunk_size or 0) - if chunk_size <= 0: - chunk_size = max(1, int(query_start_loc[-1].item())) - - starts = query_start_loc[:-1].detach().cpu().tolist() - ends = query_start_loc[1:].detach().cpu().tolist() - chunk_lens: list[int] = [] - seq_idx_chunks: list[int] = [] - last_chunk_indices: list[int] = [-1] * len(starts) - - for seq_idx, (start, end) in enumerate(zip(starts, ends)): - pos = int(start) - end = int(end) - while pos < end: - room = chunk_size - (pos % chunk_size) - take = min(room, end - pos) - if take <= 0: - break - chunk_lens.append(int(take)) - seq_idx_chunks.append(seq_idx) - last_chunk_indices[seq_idx] = len(chunk_lens) - 1 - pos += take - - boundaries = [0] - for chunk_len in chunk_lens: - boundaries.append(boundaries[-1] + chunk_len) - cu_chunk_seqlens = torch.tensor(boundaries, dtype=torch.int32, device=device) - last_chunk_indices_t = torch.tensor(last_chunk_indices, dtype=torch.int32, device=device) - seq_idx_chunks_t = torch.tensor(seq_idx_chunks, dtype=torch.int32, device=device) - return cu_chunk_seqlens, last_chunk_indices_t, seq_idx_chunks_t - - -def _repair_mamba2_attention_metadata_prefill_seq_idx(value: Any) -> Any: - if not _is_mamba2_attention_metadata(value): - return value - try: - import torch - except Exception: - return value - - num_prefills = int(getattr(value, "num_prefills", 0) or 0) - num_prefill_tokens = int(getattr(value, "num_prefill_tokens", 0) or 0) - if num_prefills <= 0 or num_prefill_tokens <= 0: - return value - if _is_torch_dynamo_compiling(): - return value - - seq_idx_p = getattr(value, "seq_idx_p", None) - query_start_loc = getattr(value, "query_start_loc", None) - query_start_loc_p = getattr(value, "query_start_loc_p", None) - if seq_idx_p is None and not hasattr(query_start_loc, "numel") and not hasattr(query_start_loc_p, "numel"): - return value - - device = getattr(seq_idx_p, "device", None) - if device is None: - device = getattr(query_start_loc, "device", None) - if device is None: - device = getattr(query_start_loc_p, "device", None) - if device is None: - device = getattr(getattr(value, "state_indices_tensor", None), "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - query_start_loc_p = _infer_mamba2_prefill_query_start_loc(value, device) - lengths = (query_start_loc_p[1:] - query_start_loc_p[:-1]).to(dtype=torch.int32) - if int(lengths.numel()) != num_prefills or int(lengths.sum().item()) != num_prefill_tokens: - lengths = _balanced_tensor_lengths(num_prefill_tokens, num_prefills, device) - query_start_loc_p = torch.zeros(num_prefills + 1, dtype=torch.int32, device=device) - query_start_loc_p[1:] = torch.cumsum(lengths, dim=0) - - if _uses_mamba2_varlen_chunk_metadata(value): - cu_chunk_seqlens, last_chunk_indices, fixed_seq_idx_p = _compute_mamba2_varlen_chunk_metadata( - query_start_loc_p, - int(getattr(value, "chunk_size", 0) or 0), - ) - expected_shape = (int(cu_chunk_seqlens.numel()) - 1,) - current_cu_chunk_seqlens = getattr(value, "cu_chunk_seqlen_p", None) - current_last_chunk_indices = getattr(value, "last_chunk_indices_p", None) - if ( - seq_idx_p is not None - and hasattr(seq_idx_p, "shape") - and tuple(seq_idx_p.shape) == expected_shape - and current_cu_chunk_seqlens is not None - and hasattr(current_cu_chunk_seqlens, "shape") - and tuple(current_cu_chunk_seqlens.shape) == tuple(cu_chunk_seqlens.shape) - and torch.equal(current_cu_chunk_seqlens.to(device=device, dtype=torch.int32), cu_chunk_seqlens) - and current_last_chunk_indices is not None - and hasattr(current_last_chunk_indices, "shape") - and tuple(current_last_chunk_indices.shape) == tuple(last_chunk_indices.shape) - and torch.equal(current_last_chunk_indices.to(device=device, dtype=torch.int32), last_chunk_indices) - and torch.equal(seq_idx_p.to(device=device, dtype=torch.int32), fixed_seq_idx_p) - ): - return value - return _replace_mamba2_attention_metadata_fields( - value, - { - "seq_idx_p": fixed_seq_idx_p, - "cu_chunk_seqlen_p": cu_chunk_seqlens, - "last_chunk_indices_p": last_chunk_indices, - }, - ) - - expected_shape = (1, num_prefill_tokens) - if ( - seq_idx_p is not None - and hasattr(seq_idx_p, "shape") - and tuple(seq_idx_p.shape) == expected_shape - ): - return value - - fixed_seq_idx_p = torch.repeat_interleave( - torch.arange(num_prefills, dtype=torch.int32, device=device), - lengths, - output_size=num_prefill_tokens, - ).unsqueeze(0) - - replacements: dict[str, Any] = {"seq_idx_p": fixed_seq_idx_p} - if bool(getattr(value, "prep_initial_states", False)): - try: - module = importlib.import_module("vllm.v1.attention.backends.mamba2_attn") - converter = getattr(module, "_query_start_loc_to_chunk_indices_offsets") - chunk_indices, chunk_offsets = converter( - query_start_loc_p, - int(getattr(value, "chunk_size", 0) or 0), - num_prefill_tokens, - ) - replacements["chunk_indices_p"] = chunk_indices - replacements["chunk_offsets_p"] = chunk_offsets - except Exception: - pass - return _replace_mamba2_attention_metadata_fields(value, replacements) - - -def _compact_legacy_mamba_seq_idx(seq_idx: Any, device: Any) -> Any: - import torch - - seq_idx_flat = seq_idx.detach().reshape(-1).to(device="cpu", dtype=torch.int64) - valid = seq_idx_flat >= 0 - if not bool(valid.all().item()): - seq_idx_flat = seq_idx_flat[valid] - if seq_idx_flat.numel() == 0: - return seq_idx_flat.to(device=device, dtype=torch.int32) - - _, compact = torch.unique(seq_idx_flat, sorted=True, return_inverse=True) - return compact.to(device=device, dtype=torch.int32) - - -def _mamba2_attention_metadata_from_mapping(value: dict[str, Any]) -> Any | None: - required_fields = ( - "num_prefills", - "num_prefill_tokens", - "num_decodes", - "num_decode_tokens", - "query_start_loc", - "seq_lens", - "prep_initial_states", - "chunk_size", - "has_initial_states_p", - "seq_idx_p", - "chunk_indices_p", - "chunk_offsets_p", - "state_indices_tensor", - ) - if not all(field in value for field in required_fields): - return None - try: - import torch - from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata - - kwargs = {field: value[field] for field in required_fields if field in value} - needed_state_indices = int(value.get("num_prefills") or 0) + int(value.get("num_decode_tokens") or 0) - kwargs["num_reqs"] = int(value.get("num_reqs") or needed_state_indices or 1) - state_indices = value.get("state_indices_tensor") - if (state_indices is None or hasattr(state_indices, "numel")) and ( - getattr(state_indices, "numel", lambda: 0)() != needed_state_indices - ): - device = getattr(state_indices, "device", None) - if device is None: - device = getattr(value.get("query_start_loc"), "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - state_indices = _resize_mamba_state_indices(state_indices, needed_state_indices, device) - kwargs["state_indices_tensor"] = state_indices - return _new_mamba2_attention_metadata(Mamba2AttentionMetadata, kwargs) - except Exception: - return types.SimpleNamespace(**{field: value.get(field) for field in required_fields}) - - -def _repair_mamba2_attention_metadata_state_indices(value: Any) -> Any: - if not _is_mamba2_attention_metadata(value): - return value - try: - import torch - except Exception: - return value - if _is_torch_dynamo_compiling(): - return value - value = _repair_mamba2_attention_metadata_prefill_seq_idx(value) - if hasattr(value, "state_indices_tensor_p") or hasattr(value, "state_indices_tensor_d"): - device = None - state_indices_p = getattr(value, "state_indices_tensor_p", None) - state_indices_d = getattr(value, "state_indices_tensor_d", None) - for state_indices in (state_indices_p, state_indices_d, getattr(value, "query_start_loc_p", None), getattr(value, "query_start_loc_d", None)): - device = getattr(state_indices, "device", None) - if device is not None: - break - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - replacements: dict[str, Any] = {} - needed_prefill_indices = int(getattr(value, "num_prefills", 0) or 0) - if needed_prefill_indices > 0 and ( - state_indices_p is None - or not hasattr(state_indices_p, "numel") - or int(state_indices_p.numel()) != needed_prefill_indices - ): - replacements["state_indices_tensor_p"] = _resize_mamba_state_indices( - state_indices_p, - needed_prefill_indices, - device, - ) - needed_decode_indices = int(getattr(value, "num_decode_tokens", 0) or 0) - if needed_decode_indices > 0 and ( - state_indices_d is None - or not hasattr(state_indices_d, "numel") - or int(state_indices_d.numel()) != needed_decode_indices - ): - replacements["state_indices_tensor_d"] = _resize_mamba_state_indices( - state_indices_d, - needed_decode_indices, - device, - ) - if not replacements: - return value - try: - if dataclasses.is_dataclass(value): - return dataclasses.replace(value, **replacements) - except Exception: - pass - fields = dict(getattr(value, "__dict__", {})) - fields.update(replacements) - return types.SimpleNamespace(**fields) - - needed_state_indices = int(getattr(value, "num_prefills", 0) or 0) + int( - getattr(value, "num_decode_tokens", 0) or 0 - ) - state_indices = getattr(value, "state_indices_tensor", None) - if state_indices is not None and not hasattr(state_indices, "numel"): - return value - if getattr(state_indices, "numel", lambda: 0)() == needed_state_indices: - return value - - device = getattr(state_indices, "device", None) - if device is None: - device = getattr(getattr(value, "query_start_loc", None), "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - fixed_state_indices = _resize_mamba_state_indices(state_indices, needed_state_indices, device) +def _install_vllm_inputs_data_alias() -> None: + if "vllm.inputs.data" in sys.modules: + return try: - if dataclasses.is_dataclass(value): - return dataclasses.replace(value, state_indices_tensor=fixed_state_indices) - except Exception: - pass - fields = dict(getattr(value, "__dict__", {})) - fields["state_indices_tensor"] = fixed_state_indices - return types.SimpleNamespace(**fields) - - -def _repair_mamba_cache_params_state_indices(cache_params: Any, attn_metadata: Any) -> Any: - if cache_params is None or not _is_mamba2_attention_metadata(attn_metadata): - return cache_params - try: - import torch + import vllm.inputs as inputs except Exception: - return cache_params - needed_state_indices = int(getattr(attn_metadata, "num_prefills", 0) or 0) + int( - getattr(attn_metadata, "num_decode_tokens", 0) or 0 - ) - state_indices = getattr(cache_params, "state_indices_tensor", None) - if state_indices is not None and not hasattr(state_indices, "numel"): - return cache_params - - device = getattr(state_indices, "device", None) - if device is None: - device = getattr(getattr(attn_metadata, "query_start_loc", None), "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - metadata_state_indices = getattr(attn_metadata, "state_indices_tensor", None) - if ( - metadata_state_indices is None - or not hasattr(metadata_state_indices, "numel") - or int(metadata_state_indices.numel()) != needed_state_indices - ): - split_indices: list[Any] = [] - prefill_indices = getattr(attn_metadata, "state_indices_tensor_p", None) - decode_indices = getattr(attn_metadata, "state_indices_tensor_d", None) - if prefill_indices is not None and hasattr(prefill_indices, "numel") and int(prefill_indices.numel()) > 0: - split_indices.append(prefill_indices) - if decode_indices is not None and hasattr(decode_indices, "numel") and int(decode_indices.numel()) > 0: - split_indices.append(decode_indices) - if split_indices: - try: - metadata_state_indices = torch.cat( - [item.detach().reshape(-1).to(device=device, dtype=torch.int32) for item in split_indices], - dim=0, - ) - except Exception: - metadata_state_indices = None - - if ( - metadata_state_indices is not None - and hasattr(metadata_state_indices, "numel") - and int(metadata_state_indices.numel()) == needed_state_indices - ): - fixed_state_indices = metadata_state_indices.detach().reshape(-1).to(device=device, dtype=torch.int32) - else: - fixed_state_indices = _resize_mamba_state_indices(state_indices, needed_state_indices, device) - - if getattr(state_indices, "numel", lambda: 0)() == needed_state_indices: - try: - current = state_indices.detach().reshape(-1).to(device=device, dtype=torch.int32) - if torch.equal(current, fixed_state_indices): - return cache_params - except Exception: - return cache_params - + return try: - if dataclasses.is_dataclass(cache_params): - return dataclasses.replace(cache_params, state_indices_tensor=fixed_state_indices) + if importlib.util.find_spec("vllm.inputs.data") is not None: + return except Exception: pass - fields = dict(getattr(cache_params, "__dict__", {})) - fields["state_indices_tensor"] = fixed_state_indices - return types.SimpleNamespace(**fields) - - -def _select_mamba2_attention_metadata(value: Any) -> Any | None: - if not isinstance(value, dict): - return _repair_mamba2_attention_metadata_state_indices(value) if _is_mamba2_attention_metadata(value) else None - converted = _mamba2_attention_metadata_from_mapping(value) - if converted is not None: - return _repair_mamba2_attention_metadata_state_indices(converted) - for child in value.values(): - selected = _select_mamba2_attention_metadata(child) - if selected is not None: - return selected - return None - - -def _normalize_mamba2_attention_metadata_groups(attn_metadata: Any) -> Any: - if not isinstance(attn_metadata, dict): - return attn_metadata - normalized = {} - changed = False - for key, value in attn_metadata.items(): - selected = _select_mamba2_attention_metadata(value) - if selected is not None and selected is not value: - normalized[key] = selected - changed = True - continue - normalized_value = _normalize_mamba2_attention_metadata_groups(value) - normalized[key] = normalized_value - changed = changed or normalized_value is not value - return normalized if changed else attn_metadata - - -def _as_v1_mamba2_attention_metadata_dict(attn_metadata: Any, vllm_config: Any) -> Any: - if isinstance(attn_metadata, dict) or not _looks_like_easymagpie_vllm_config(vllm_config): - return attn_metadata - selected = _select_mamba2_attention_metadata(attn_metadata) - if selected is None: - return attn_metadata - selected = _repair_mamba2_attention_metadata_state_indices(selected) - - layer_metadata: dict[str, Any] = {} - static_context = getattr(getattr(vllm_config, "compilation_config", None), "static_forward_context", None) - if isinstance(static_context, dict): - for layer_name in static_context: - if layer_name.endswith(".mixer"): - layer_metadata.setdefault(layer_name, selected) - - pattern = _get_easymagpie_hybrid_pattern(vllm_config) - if pattern: - for layer_idx, layer_type in enumerate(pattern): - if layer_type == "M": - layer_metadata.setdefault(f"backbone.layers.{layer_idx}.mixer", selected) - - return layer_metadata or attn_metadata - - -def _build_flash_attention_metadata_from_lengths( - *, - num_tokens: int, - num_reqs: int, - max_query_len: int, - device: Any = None, - slot_mapping: Any = None, - query_start_values: list[int] | None = None, - seq_lens: list[int] | None = None, -) -> Any | None: - torch = _TORCH_MODULE - FlashAttentionMetadata = _LEGACY_FLASH_ATTENTION_METADATA_CLS - if torch is None or FlashAttentionMetadata is None: - return None - - num_tokens = max(1, int(num_tokens)) - num_reqs = max(1, int(num_reqs)) - if query_start_values is not None and len(query_start_values) == num_reqs + 1: - query_lens = [ - max(0, query_start_values[index + 1] - query_start_values[index]) - for index in range(num_reqs) - ] - if sum(query_lens) != num_tokens: - query_lens = _balanced_profile_lengths(num_tokens, num_reqs) - query_start_values = None - else: - query_lens = _balanced_profile_lengths(num_tokens, num_reqs) - query_start_values = None - - if ( - seq_lens is None - or len(seq_lens) != num_reqs - or any(length <= 0 for length in seq_lens) - or sum(seq_lens) != num_tokens - ): - seq_lens = list(query_lens) - - if query_start_values is None: - query_start_values = [0] - for length in query_lens: - query_start_values.append(query_start_values[-1] + int(length)) - seq_start_values = [0] - for length in seq_lens: - seq_start_values.append(seq_start_values[-1] + int(length)) - - if device is None: - device = getattr(slot_mapping, "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) - seq_start_loc = torch.tensor(seq_start_values, dtype=torch.int32, device=device) - seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device) - context_lens_tensor = torch.zeros(num_reqs, dtype=torch.int, device=device) - if slot_mapping is None or not hasattr(slot_mapping, "shape"): - slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=device) - else: - slot_mapping = slot_mapping[:num_tokens].to(device=device, dtype=torch.long) - block_tables = torch.empty((num_reqs, 0), dtype=torch.int, device=device) - - return FlashAttentionMetadata( - num_prefills=num_reqs, - num_prefill_tokens=num_tokens, - num_decode_tokens=0, - slot_mapping=slot_mapping, - multi_modal_placeholder_index_maps={}, - enable_kv_scales_calculation=True, - seq_lens=seq_lens, - seq_lens_tensor=seq_lens_tensor, - max_query_len=max(int(max_query_len), _max_int_sequence(query_lens)), - max_prefill_seq_len=_max_int_sequence(seq_lens), - max_decode_query_len=0, - max_decode_seq_len=0, - query_start_loc=query_start_loc, - seq_start_loc=seq_start_loc, - context_lens_tensor=context_lens_tensor, - block_tables=block_tables, - use_cuda_graph=False, - ) - - -def _requested_attention_backend() -> str: - try: - import os - - return str(os.environ.get("VLLM_ATTENTION_BACKEND") or "").upper() - except Exception: - return "" - - -def _profile_attention_lengths( - *, - num_tokens: int, - num_reqs: int, - query_start_values: list[int] | None = None, - seq_lens: list[int] | None = None, -) -> tuple[list[int], list[int], list[int]]: - num_tokens = max(1, int(num_tokens)) - num_reqs = max(1, int(num_reqs)) - if query_start_values is not None and len(query_start_values) == num_reqs + 1: - query_lens = [ - max(0, query_start_values[index + 1] - query_start_values[index]) - for index in range(num_reqs) - ] - if sum(query_lens) != num_tokens: - query_lens = _balanced_profile_lengths(num_tokens, num_reqs) - query_start_values = None - else: - query_lens = _balanced_profile_lengths(num_tokens, num_reqs) - query_start_values = None - - if ( - seq_lens is None - or len(seq_lens) != num_reqs - or any(length <= 0 for length in seq_lens) - or sum(seq_lens) != num_tokens - ): - seq_lens = list(query_lens) - - if query_start_values is None: - query_start_values = [0] - for length in query_lens: - query_start_values.append(query_start_values[-1] + int(length)) - return query_lens, query_start_values, seq_lens - - -def _build_v1_triton_attention_metadata_from_lengths( - *, - num_tokens: int, - num_reqs: int, - max_query_len: int, - device: Any = None, - slot_mapping: Any = None, - query_start_values: list[int] | None = None, - seq_lens: list[int] | None = None, -) -> Any | None: - torch = _TORCH_MODULE - TritonAttentionMetadata = _V1_TRITON_ATTENTION_METADATA_CLS - if torch is None or TritonAttentionMetadata is None: - return None - - query_lens, query_start_values, seq_lens = _profile_attention_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - if device is None: - device = getattr(slot_mapping, "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - num_tokens = max(1, int(num_tokens)) - num_reqs = max(1, int(num_reqs)) - query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) - seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32, device=device) - if slot_mapping is None or not hasattr(slot_mapping, "shape"): - slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=device) - else: - slot_mapping = slot_mapping[:num_tokens].to(device=device, dtype=torch.long) - block_table = torch.empty((num_reqs, 0), dtype=torch.int32, device=device) - empty_float = torch.empty((0,), dtype=torch.float32, device=device) - context_lens_tensor = torch.zeros(num_reqs, dtype=torch.int32, device=device) - resolved_max_query_len = max(int(max_query_len), _max_int_sequence(query_lens)) - - metadata = _call_with_supported_kwargs( - TritonAttentionMetadata, - { - "num_actual_tokens": num_tokens, - "max_query_len": resolved_max_query_len, - "query_start_loc": query_start_loc, - "max_seq_len": _max_int_sequence(seq_lens), - "seq_lens": seq_lens_tensor, - "block_table": block_table, - "slot_mapping": slot_mapping, - "seq_threshold_3D": 0, - "num_par_softmax_segments": 0, - "softmax_segm_output": empty_float, - "softmax_segm_max": empty_float, - "softmax_segm_expsum": empty_float, - "use_cascade": False, - "common_prefix_len": 0, - "cu_prefix_query_lens": None, - "prefix_kv_lens": None, - "suffix_kv_lens": None, - "scheduler_metadata": None, - "prefix_scheduler_metadata": None, - "mm_prefix_range": None, - "mm_prefix_range_tensor": None, - }, - ) - return _attach_profile_attention_metadata_attrs( - metadata, - num_tokens=num_tokens, - num_reqs=num_reqs, - query_lens=query_lens, - seq_lens=seq_lens, - seq_lens_tensor=seq_lens_tensor, - seq_start_loc=query_start_loc, - query_start_loc=query_start_loc, - context_lens_tensor=context_lens_tensor, - block_tables=block_table, - slot_mapping=slot_mapping, - max_query_len=resolved_max_query_len, - ) - - -def _build_v1_flash_attention_metadata_from_lengths( - *, - num_tokens: int, - num_reqs: int, - max_query_len: int, - device: Any = None, - slot_mapping: Any = None, - query_start_values: list[int] | None = None, - seq_lens: list[int] | None = None, -) -> Any | None: - torch = _TORCH_MODULE - FlashAttentionMetadata = _V1_FLASH_ATTENTION_METADATA_CLS - if torch is None or FlashAttentionMetadata is None: - return None - - query_lens, query_start_values, seq_lens = _profile_attention_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - if device is None: - device = getattr(slot_mapping, "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - num_tokens = max(1, int(num_tokens)) - num_reqs = max(1, int(num_reqs)) - query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) - seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32, device=device) - if slot_mapping is None or not hasattr(slot_mapping, "shape"): - slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=device) - else: - slot_mapping = slot_mapping[:num_tokens].to(device=device, dtype=torch.long) - block_table = torch.empty((num_reqs, 0), dtype=torch.int32, device=device) - context_lens_tensor = torch.zeros(num_reqs, dtype=torch.int32, device=device) - resolved_max_query_len = max(int(max_query_len), _max_int_sequence(query_lens)) - - metadata = _call_with_supported_kwargs( - FlashAttentionMetadata, - { - "num_actual_tokens": num_tokens, - "max_query_len": resolved_max_query_len, - "query_start_loc": query_start_loc, - "max_seq_len": _max_int_sequence(seq_lens), - "seq_lens": seq_lens_tensor, - "block_table": block_table, - "slot_mapping": slot_mapping, - "use_cascade": False, - "common_prefix_len": 0, - "cu_prefix_query_lens": None, - "prefix_kv_lens": None, - "suffix_kv_lens": None, - "max_dcp_context_kv_len": None, - "dcp_context_kv_lens": None, - "scheduler_metadata": None, - "prefix_scheduler_metadata": None, - "max_num_splits": 0, - "causal": True, - }, - ) - return _attach_profile_attention_metadata_attrs( - metadata, - num_tokens=num_tokens, - num_reqs=num_reqs, - query_lens=query_lens, - seq_lens=seq_lens, - seq_lens_tensor=seq_lens_tensor, - seq_start_loc=query_start_loc, - query_start_loc=query_start_loc, - context_lens_tensor=context_lens_tensor, - block_tables=block_table, - slot_mapping=slot_mapping, - max_query_len=resolved_max_query_len, - ) - - -def _build_vllm_attention_metadata_from_lengths( - *, - num_tokens: int, - num_reqs: int, - max_query_len: int, - device: Any = None, - slot_mapping: Any = None, - query_start_values: list[int] | None = None, - seq_lens: list[int] | None = None, -) -> Any: - requested_backend = _requested_attention_backend() - is_flashinfer_backend = requested_backend.startswith("FLASHINFER") - is_flash_attention_backend = requested_backend.startswith("FLASH_ATTN") or requested_backend in { - "FLASH", - "FLASHATTN", - } - if is_flash_attention_backend: - metadata = _build_v1_flash_attention_metadata_from_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - max_query_len=max_query_len, - device=device, - slot_mapping=slot_mapping, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - if metadata is not None: - return metadata - if not is_flashinfer_backend and ("TRITON" in requested_backend or requested_backend != "XFORMERS"): - metadata = _build_v1_triton_attention_metadata_from_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - max_query_len=max_query_len, - device=device, - slot_mapping=slot_mapping, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - if metadata is not None: - return metadata - metadata = _build_v1_flash_attention_metadata_from_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - max_query_len=max_query_len, - device=device, - slot_mapping=slot_mapping, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - if metadata is not None: - return metadata - - torch = _TORCH_MODULE - XFormersMetadata = _XFORMERS_METADATA_CLS - if torch is None or XFormersMetadata is None: - metadata = _build_v1_flash_attention_metadata_from_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - max_query_len=max_query_len, - device=device, - slot_mapping=slot_mapping, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - if metadata is not None: - return metadata - metadata = _build_v1_triton_attention_metadata_from_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - max_query_len=max_query_len, - device=device, - slot_mapping=slot_mapping, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - if metadata is not None: - return metadata - metadata = _build_flash_attention_metadata_from_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - max_query_len=max_query_len, - device=device, - slot_mapping=slot_mapping, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - if metadata is None: - return None - for attr in ("attn_bias", "encoder_attn_bias", "cross_attn_bias"): - if not hasattr(metadata, attr): - setattr(metadata, attr, None) - return metadata - - num_tokens = max(1, int(num_tokens)) - num_reqs = max(1, int(num_reqs)) - if query_start_values is not None and len(query_start_values) == num_reqs + 1: - query_lens = [ - max(0, query_start_values[index + 1] - query_start_values[index]) - for index in range(num_reqs) - ] - if sum(query_lens) != num_tokens: - query_lens = _balanced_profile_lengths(num_tokens, num_reqs) - query_start_values = None - else: - query_lens = _balanced_profile_lengths(num_tokens, num_reqs) - query_start_values = None - - if ( - seq_lens is None - or len(seq_lens) != num_reqs - or any(length <= 0 for length in seq_lens) - or sum(seq_lens) != num_tokens - ): - seq_lens = list(query_lens) - - if query_start_values is None: - query_start_values = [0] - for length in query_lens: - query_start_values.append(query_start_values[-1] + int(length)) - seq_start_values = [0] - for length in seq_lens: - seq_start_values.append(seq_start_values[-1] + int(length)) - - if device is None: - device = getattr(slot_mapping, "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) - seq_start_loc = torch.tensor(seq_start_values, dtype=torch.int32, device=device) - seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device) - context_lens_tensor = torch.zeros(num_reqs, dtype=torch.int, device=device) - if slot_mapping is None or not hasattr(slot_mapping, "shape"): - slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=device) - else: - slot_mapping = slot_mapping[:num_tokens].to(device=device, dtype=torch.long) - block_tables = torch.empty((num_reqs, 0), dtype=torch.int, device=device) - - return XFormersMetadata( - seq_lens_tensor=seq_lens_tensor, - max_decode_seq_len=0, - block_tables=block_tables, - num_prefills=num_reqs, - num_prefill_tokens=num_tokens, - num_decode_tokens=0, - slot_mapping=slot_mapping, - multi_modal_placeholder_index_maps={}, - enable_kv_scales_calculation=True, - max_prefill_seq_len=_max_int_sequence(seq_lens), - use_cuda_graph=False, - seq_lens=seq_lens, - seq_start_loc=seq_start_loc, - context_lens_tensor=context_lens_tensor, - max_query_len=max(int(max_query_len), _max_int_sequence(query_lens)), - max_decode_query_len=0, - query_start_loc=query_start_loc, - ) - - -def _build_profile_flash_attention_metadata(self: Any, num_tokens: int, num_reqs: int, max_query_len: int) -> Any: - query_start_values = _buffer_slice_to_int_list(getattr(self, "query_start_loc", None), int(num_reqs) + 1) - seq_lens = _buffer_slice_to_int_list(getattr(self, "seq_lens", None), int(num_reqs)) - device = getattr(self, "device", None) - if device is None: - query_start_gpu = getattr(getattr(self, "query_start_loc", None), "gpu", None) - device = getattr(query_start_gpu, "device", None) - return _build_vllm_attention_metadata_from_lengths( - num_tokens=num_tokens, - num_reqs=num_reqs, - max_query_len=max_query_len, - device=device, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - - -def _get_easymagpie_hybrid_pattern(vllm_config: Any) -> str | None: - model_config = getattr(vllm_config, "model_config", None) - hf_config = getattr(model_config, "hf_config", None) - for candidate in (hf_config, model_config, vllm_config): - pattern = getattr(candidate, "hybrid_override_pattern", None) - if isinstance(pattern, str) and pattern: - return pattern - return None - - -def _get_easymagpie_chunk_size(vllm_config: Any) -> int: - model_config = getattr(vllm_config, "model_config", None) - get_mamba_chunk_size = getattr(model_config, "get_mamba_chunk_size", None) - if get_mamba_chunk_size is not None: - try: - chunk_size = get_mamba_chunk_size() - if chunk_size is not None: - return int(chunk_size) - except Exception: - pass - hf_config = getattr(model_config, "hf_config", None) - for candidate in (hf_config, model_config, vllm_config): - chunk_size = getattr(candidate, "chunk_size", None) - if chunk_size is not None: - return int(chunk_size) - return 128 - - -def _build_profile_mamba2_attention_metadata(flash_metadata: Any, chunk_size: int) -> Any: - torch = _TORCH_MODULE - Mamba2AttentionMetadata = _V1_MAMBA2_ATTENTION_METADATA_CLS - if torch is None or Mamba2AttentionMetadata is None or flash_metadata is None: - return None - - query_start_loc = getattr(flash_metadata, "query_start_loc", None) - if query_start_loc is None: - device = getattr(getattr(flash_metadata, "slot_mapping", None), "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - num_prefills = int(getattr(flash_metadata, "num_prefills", 1) or 1) - query_start_loc = torch.arange(num_prefills + 1, dtype=torch.int32, device=device) - else: - device = query_start_loc.device - num_prefills = max(0, int(getattr(flash_metadata, "num_prefills", query_start_loc.numel() - 1))) - - num_prefill_tokens_value = getattr(flash_metadata, "num_prefill_tokens", None) - if num_prefill_tokens_value is None: - num_prefill_tokens_value = query_start_loc[-1].item() - num_prefill_tokens = int(num_prefill_tokens_value) - seq_lens = getattr(flash_metadata, "seq_lens_tensor", None) - if seq_lens is None: - seq_lens_values = getattr(flash_metadata, "seq_lens", None) - if seq_lens_values is None: - seq_lens = query_start_loc[1:] - query_start_loc[:-1] - else: - seq_lens = torch.tensor(seq_lens_values, dtype=torch.int, device=device) - else: - seq_lens = seq_lens.to(device=device) - - if num_prefills > 0: - query_lens = _balanced_profile_lengths(num_prefill_tokens, num_prefills) - query_start_values = [0] - for query_len in query_lens: - query_start_values.append(query_start_values[-1] + int(query_len)) - query_start_loc = torch.tensor(query_start_values, dtype=torch.int32, device=device) - seq_lens = torch.tensor(query_lens, dtype=torch.int, device=device) - query_start_loc_p = query_start_loc[: num_prefills + 1] - seq_idx_p = torch.repeat_interleave( - torch.arange(num_prefills, dtype=torch.int32, device=device), - query_start_loc_p.diff(), - output_size=num_prefill_tokens, - ) - seq_idx_p.unsqueeze_(0) - has_initial_states_p = torch.zeros(num_prefills, dtype=torch.bool, device=device) - else: - seq_idx_p = None - has_initial_states_p = None - - state_indices_tensor_p = ( - torch.arange(max(1, num_prefills), dtype=torch.int32, device=device)[:num_prefills] - if num_prefills > 0 - else None - ) - num_computed_tokens_p = torch.zeros(num_prefills, dtype=torch.int32, device=device) if num_prefills > 0 else None - kwargs = { - "num_reqs": num_prefills, - "num_prefills": num_prefills, - "num_prefill_tokens": num_prefill_tokens, - "num_decodes": 0, - "num_decode_tokens": 0, - "query_start_loc_p": query_start_loc_p if num_prefills > 0 else None, - "query_start_loc_d": None, - "seq_lens": seq_lens, - "num_computed_tokens_p": num_computed_tokens_p, - "prep_initial_states": False, - "chunk_size": int(chunk_size), - "has_initial_states_p": has_initial_states_p, - "seq_idx_p": seq_idx_p, - "state_indices_tensor_p": state_indices_tensor_p, - "state_indices_tensor_d": None, - "num_accepted_tokens": None, - "block_idx_last_scheduled_token": None, - "block_idx_first_scheduled_token_p": None, - "block_idx_last_computed_token": None, - "cu_chunk_seqlen_p": None, - "last_chunk_indices_p": None, - "nums_dict": None, - "batch_ptr": None, - "token_chunk_offset_ptr": None, - } - metadata = _new_mamba2_attention_metadata(Mamba2AttentionMetadata, kwargs) - return metadata - - -def _build_mamba2_attention_metadata_from_legacy( - legacy_metadata: Any, - mamba_cache_params: Any, - hidden_states: Any, -) -> Any | None: - torch = _TORCH_MODULE - Mamba2AttentionMetadata = _V1_MAMBA2_ATTENTION_METADATA_CLS - if torch is None or Mamba2AttentionMetadata is None: - return None - - if legacy_metadata is None: - return None - - device = getattr(hidden_states, "device", None) - if device is None: - state_indices = getattr(mamba_cache_params, "state_indices_tensor", None) - device = getattr(state_indices, "device", None) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - hidden_len = int(getattr(hidden_states, "shape", [1])[0] or 1) - legacy_seq_idx_p = getattr(legacy_metadata, "seq_idx", None) - chunk_indices_p = getattr(legacy_metadata, "chunk_indices", None) - chunk_offsets_p = getattr(legacy_metadata, "chunk_offsets", None) - chunk_size = int(getattr(legacy_metadata, "chunk_size", 128) or 128) - - if legacy_seq_idx_p is not None and getattr(legacy_seq_idx_p, "numel", lambda: 0)() > 0: - seq_idx_flat = _compact_legacy_mamba_seq_idx(legacy_seq_idx_p, device) - num_prefill_tokens = int(seq_idx_flat.numel()) - num_prefills = int(seq_idx_flat.max().item()) + 1 if num_prefill_tokens > 0 else 0 - counts = ( - torch.bincount(seq_idx_flat, minlength=num_prefills).to(dtype=torch.int32, device=device) - if num_prefills > 0 - else torch.empty(0, dtype=torch.int32, device=device) - ) - query_start_loc_p = torch.zeros(num_prefills + 1, dtype=torch.int32, device=device) - if num_prefills > 0: - query_start_loc_p[1:] = torch.cumsum(counts, dim=0) - num_decode_tokens = max(0, hidden_len - num_prefill_tokens) - num_decodes = num_decode_tokens - decode_lens = ( - torch.ones(num_decode_tokens, dtype=torch.int32, device=device) - if num_decode_tokens > 0 - else torch.empty(0, dtype=torch.int32, device=device) - ) - seq_lens = torch.cat([counts, decode_lens], dim=0) if num_decode_tokens > 0 else counts - legacy_query_start_loc = torch.zeros(num_prefills + num_decodes + 1, dtype=torch.int32, device=device) - if seq_lens.numel() > 0: - legacy_query_start_loc[1:] = torch.cumsum(seq_lens, dim=0) - has_initial_states_p = getattr(legacy_metadata, "has_initial_states", None) - if has_initial_states_p is None: - has_initial_states_p = torch.zeros(num_prefills, dtype=torch.bool, device=device) - else: - has_initial_states_p = has_initial_states_p.to(device=device, dtype=torch.bool)[:num_prefills] - seq_idx_p = seq_idx_flat.unsqueeze(0) if num_prefill_tokens > 0 else None - else: - num_prefills = 0 - num_prefill_tokens = 0 - num_decodes = max(1, hidden_len) - num_decode_tokens = num_decodes - query_start_loc_p = None - legacy_query_start_loc = torch.arange(num_decodes + 1, dtype=torch.int32, device=device) - seq_lens = torch.ones(num_decodes, dtype=torch.int32, device=device) - has_initial_states_p = None - seq_idx_p = None - - state_indices_tensor = getattr(mamba_cache_params, "state_indices_tensor", None) - state_indices_tensor_p = getattr(mamba_cache_params, "state_indices_tensor_p", None) - if num_prefills > 0 and ( - state_indices_tensor_p is None or getattr(state_indices_tensor_p, "numel", lambda: 0)() != num_prefills - ): - state_indices_tensor_p = _resize_mamba_state_indices( - state_indices_tensor_p if state_indices_tensor_p is not None else state_indices_tensor, - num_prefills, - device, - ) - elif num_prefills <= 0: - state_indices_tensor_p = None - - state_indices_tensor_d = getattr(mamba_cache_params, "state_indices_tensor_d", None) - if num_decode_tokens > 0 and ( - state_indices_tensor_d is None - or getattr(state_indices_tensor_d, "numel", lambda: 0)() != num_decode_tokens - ): - state_indices_tensor_d = _resize_mamba_state_indices( - state_indices_tensor_d if state_indices_tensor_d is not None else state_indices_tensor, - num_decode_tokens, - device, - ) - elif num_decode_tokens <= 0: - state_indices_tensor_d = None - - query_start_loc_d = ( - torch.arange(num_decodes + 1, dtype=torch.int32, device=device) if num_decode_tokens > 0 else None - ) - num_computed_tokens_p = torch.zeros(num_prefills, dtype=torch.int32, device=device) if num_prefills > 0 else None - legacy_state_indices_tensor = _resize_mamba_state_indices( - state_indices_tensor, - num_prefills + num_decode_tokens, - device, - ) - - metadata = _new_mamba2_attention_metadata( - Mamba2AttentionMetadata, - { - "num_reqs": num_prefills + num_decodes, - "num_prefills": num_prefills, - "num_prefill_tokens": num_prefill_tokens, - "num_decodes": num_decodes, - "num_decode_tokens": num_decode_tokens, - "query_start_loc_p": query_start_loc_p, - "query_start_loc_d": query_start_loc_d, - "legacy_query_start_loc": legacy_query_start_loc, - "seq_lens": seq_lens, - "num_computed_tokens_p": num_computed_tokens_p, - "prep_initial_states": bool(getattr(legacy_metadata, "prep_initial_states", False)), - "chunk_size": chunk_size, - "has_initial_states_p": has_initial_states_p, - "seq_idx_p": seq_idx_p, - "state_indices_tensor_p": state_indices_tensor_p, - "state_indices_tensor_d": state_indices_tensor_d, - "legacy_state_indices_tensor": legacy_state_indices_tensor, - "num_accepted_tokens": None, - "block_idx_last_scheduled_token": None, - "block_idx_first_scheduled_token_p": None, - "block_idx_last_computed_token": None, - "cu_chunk_seqlen_p": None, - "last_chunk_indices_p": None, - "nums_dict": None, - "batch_ptr": None, - "token_chunk_offset_ptr": None, - }, - ) - return _repair_mamba2_attention_metadata_prefill_seq_idx(metadata) - - -def _select_legacy_mamba2_metadata(value: Any) -> Any | None: - if not isinstance(value, dict): - if any(hasattr(value, attr) for attr in ("seq_idx", "has_initial_states", "chunk_indices", "chunk_offsets")): - return value - return None - for child in value.values(): - selected = _select_legacy_mamba2_metadata(child) - if selected is not None: - return selected - return None - - -def _build_hybrid_profile_attention_metadata(vllm_config: Any, flash_metadata: Any) -> Any: - layer_metadata: dict[str, Any] = {} - mamba_metadata = _build_profile_mamba2_attention_metadata( - flash_metadata, - _get_easymagpie_chunk_size(vllm_config), - ) - static_context = getattr(getattr(vllm_config, "compilation_config", None), "static_forward_context", None) - if isinstance(static_context, dict): - for layer_name in static_context: - if layer_name.endswith(".attn") and flash_metadata is not None: - layer_metadata.setdefault(layer_name, flash_metadata) - elif layer_name.endswith(".mixer") and mamba_metadata is not None: - layer_metadata.setdefault(layer_name, mamba_metadata) - - pattern = _get_easymagpie_hybrid_pattern(vllm_config) - if pattern: - for layer_idx, layer_type in enumerate(pattern): - if layer_type == "M" and mamba_metadata is not None: - layer_metadata.setdefault(f"backbone.layers.{layer_idx}.mixer", mamba_metadata) - elif layer_type == "*" and flash_metadata is not None: - layer_metadata.setdefault(f"backbone.layers.{layer_idx}.mixer.attn", flash_metadata) - return layer_metadata or flash_metadata - - -def _looks_like_easymagpie_vllm_config(vllm_config: Any) -> bool: - candidates = [vllm_config, getattr(vllm_config, "model_config", None)] - model_config = candidates[-1] - candidates.append(getattr(model_config, "hf_config", None)) - for candidate in candidates: - if candidate is None: - continue - class_name = candidate.__class__.__name__.lower() - if "easymagpie" in class_name or "nemotronh" in class_name or "nemotron_h" in class_name: - return True - pattern = getattr(candidate, "hybrid_override_pattern", None) - if isinstance(pattern, str) and "M" in pattern and hasattr(candidate, "chunk_size"): - return True - for attr_name in ("model", "served_model_name", "model_type"): - value = getattr(candidate, attr_name, None) - if value is not None and ( - "easymagpie" in str(value).lower() - or "nemotronh" in str(value).lower() - or "nemotron_h" in str(value).lower() - ): - return True - architectures = getattr(candidate, "architectures", None) or () - if any( - "easymagpie" in str(architecture).lower() - or "nemotronh" in str(architecture).lower() - or "nemotron_h" in str(architecture).lower() - for architecture in architectures - ): - return True - return False - - -def _is_generic_profile_attention_metadata(attn_metadata: Any) -> bool: - if attn_metadata is None or isinstance(attn_metadata, dict): - return False - if hasattr(attn_metadata, "use_cascade"): - return False - return any( - hasattr(attn_metadata, attr) - for attr in ( - "num_prefills", - "num_prefill_tokens", - "query_start_loc", - "seq_lens", - "seq_lens_tensor", - ) - ) - - -def _is_empty_attention_kv_cache(kv_cache: Any) -> bool: - if kv_cache is None: - return True - shape = getattr(kv_cache, "shape", None) - if shape is not None: - try: - return int(shape[0]) < 2 - except Exception: - pass - try: - return len(kv_cache) < 2 - except Exception: - return False + 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)) -def _filter_supported_kwargs(callable_obj: Any, kwargs: dict[str, Any]) -> dict[str, Any]: - try: - parameters = inspect.signature(callable_obj).parameters - except Exception: - return dict(kwargs) - if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()): - return dict(kwargs) - return {key: value for key, value in kwargs.items() if key in parameters} - - -def _synthetic_profile_attention_metadata_from_context(attn_metadata: Any, vllm_config: Any, kwargs: dict[str, Any]) -> Any: - if not _looks_like_easymagpie_vllm_config(vllm_config): - return attn_metadata - if attn_metadata is not None and not _is_generic_profile_attention_metadata(attn_metadata): - return attn_metadata - - num_tokens = None - num_reqs = None - max_query_len = None - slot_mapping = kwargs.get("slot_mapping") - query_start_values = None - seq_lens = None - if attn_metadata is not None: - num_tokens = getattr(attn_metadata, "num_prefill_tokens", None) - if num_tokens is None: - num_tokens = getattr(attn_metadata, "num_actual_tokens", None) - num_reqs = getattr(attn_metadata, "num_prefills", None) - if num_reqs is None: - num_reqs = getattr(attn_metadata, "num_reqs", None) - max_query_len = getattr(attn_metadata, "max_query_len", None) - metadata_slot_mapping = getattr(attn_metadata, "slot_mapping", None) - if metadata_slot_mapping is not None: - slot_mapping = metadata_slot_mapping - if num_reqs is not None: - query_start_values = _buffer_slice_to_int_list( - getattr(attn_metadata, "query_start_loc", None), - int(num_reqs) + 1, - ) - seq_lens = _buffer_slice_to_int_list( - _first_not_none( - getattr(attn_metadata, "seq_lens_tensor", None), - getattr(attn_metadata, "seq_lens", None), - ), - int(num_reqs), - ) - - if num_tokens is None: - num_tokens = kwargs.get("num_tokens") - if num_tokens is None: - return attn_metadata - if num_reqs is None: - batch_descriptor = kwargs.get("batch_descriptor") - num_reqs = getattr(batch_descriptor, "num_reqs", None) - if num_reqs is None: - num_reqs = min(max(1, int(num_tokens)), 1) - if max_query_len is None: - max_query_len = max(1, int(num_tokens) // max(1, int(num_reqs))) - flash_metadata = _build_vllm_attention_metadata_from_lengths( - num_tokens=int(num_tokens), - num_reqs=int(num_reqs), - max_query_len=int(max_query_len), - slot_mapping=slot_mapping, - query_start_values=query_start_values, - seq_lens=seq_lens, - ) - return _build_hybrid_profile_attention_metadata(vllm_config, flash_metadata) - - -def _install_mamba2_metadata_compat() -> None: - try: - module = importlib.import_module("vllm.model_executor.layers.mamba.mamba2_metadata") - except Exception: - return - original = getattr(module, "prepare_mamba2_metadata", None) - if original is None or getattr(original, "_easymagpie_compat", False): - return - - def _empty_attn_metadata(): - return types.SimpleNamespace( - num_prefills=0, - num_prefill_tokens=0, - context_lens_tensor=None, - query_start_loc=None, - ) - - def _select_attn_metadata(attn_metadata: Any): - if not isinstance(attn_metadata, dict): - return attn_metadata if hasattr(attn_metadata, "num_prefills") else None - for value in attn_metadata.values(): - selected = _select_attn_metadata(value) - if selected is not None: - return selected - return None - - def _balanced_tensor_lengths(total: int, count: int, device: Any): - import torch - - if count <= 0: - return torch.empty(0, dtype=torch.int32, device=device) - base, extra = divmod(max(0, int(total)), int(count)) - lengths = torch.full((count,), base, dtype=torch.int32, device=device) - if extra: - lengths[:extra] += 1 - return lengths - - def _build_safe_legacy_metadata(chunk_size: int, attn_metadata: Any, mamba2_metadata: Any = None): - import torch - - num_prefills = int(getattr(attn_metadata, "num_prefills", 0) or 0) - num_prefill_tokens = int(getattr(attn_metadata, "num_prefill_tokens", 0) or 0) - query_start_loc = getattr(attn_metadata, "query_start_loc", None) - if num_prefills <= 0 or num_prefill_tokens <= 0 or query_start_loc is None: - return None - - query_start_loc = query_start_loc.to(dtype=torch.int64) - if int(query_start_loc.numel()) < 2: - return None - query_lens = query_start_loc[1:] - query_start_loc[:-1] - device = query_start_loc.device - - if int(query_start_loc.numel()) >= num_prefills + 1: - first_prefill_lens = query_start_loc[: num_prefills + 1].diff() - try: - if int(first_prefill_lens.sum().item()) == num_prefill_tokens: - return None - except Exception: - pass - - prefill_lens = query_lens[query_lens > 1] - try: - prefill_lens_ok = ( - int(prefill_lens.numel()) == num_prefills - and int(prefill_lens.sum().item()) == num_prefill_tokens - ) - except Exception: - prefill_lens_ok = False - if not prefill_lens_ok: - prefill_lens = _balanced_tensor_lengths(num_prefill_tokens, num_prefills, device) - else: - prefill_lens = prefill_lens.to(dtype=torch.int32) - - query_start_loc_p = torch.zeros(num_prefills + 1, dtype=torch.int32, device=device) - if num_prefills > 0: - query_start_loc_p[1:] = torch.cumsum(prefill_lens.to(dtype=torch.int32), dim=0) - seq_idx = torch.repeat_interleave( - torch.arange(num_prefills, dtype=torch.int32, device=device), - prefill_lens.to(dtype=torch.int32), - ).unsqueeze(0) - - has_initial_states = None - prep_initial_states = False - context_lens_tensor = getattr(attn_metadata, "context_lens_tensor", None) - if context_lens_tensor is not None: - try: - context_lens_tensor = context_lens_tensor.to(device=device) - if int(context_lens_tensor.numel()) == int(query_lens.numel()) and int( - (query_lens > 1).sum().item() - ) == num_prefills: - context_lens_tensor = context_lens_tensor[query_lens > 1] - else: - context_lens_tensor = context_lens_tensor[:num_prefills] - has_initial_states = context_lens_tensor > 0 - prep_initial_states = bool(torch.any(has_initial_states).item()) - except Exception: - has_initial_states = None - prep_initial_states = False - - chunk_indices, chunk_offsets = None, None - if prep_initial_states: - try: - converter = getattr(module, "_query_start_loc_to_chunk_indices_offsets") - chunk_indices, chunk_offsets = converter(query_start_loc_p, chunk_size, num_prefill_tokens) - except Exception: - chunk_indices, chunk_offsets = None, None - - target = mamba2_metadata - if target is None: - metadata_cls = getattr(module, "Mamba2Metadata", None) - if metadata_cls is None: - return types.SimpleNamespace( - has_initial_states=has_initial_states, - prep_initial_states=prep_initial_states, - chunk_size=chunk_size, - seq_idx=seq_idx, - chunk_indices=chunk_indices, - chunk_offsets=chunk_offsets, - cu_seqlen=None, - ) - return metadata_cls( - has_initial_states=has_initial_states, - prep_initial_states=prep_initial_states, - chunk_size=chunk_size, - seq_idx=seq_idx, - chunk_indices=chunk_indices, - chunk_offsets=chunk_offsets, - ) - - target.has_initial_states = has_initial_states - target.prep_initial_states = prep_initial_states - target.chunk_size = chunk_size - target.seq_idx = seq_idx - target.chunk_indices = chunk_indices - target.chunk_offsets = chunk_offsets - target.cu_seqlen = None - return target - - def prepare_mamba2_metadata(chunk_size: int, attn_metadata: Any, mamba2_metadata: Any = None): - if attn_metadata is None: - attn_metadata = _empty_attn_metadata() - elif isinstance(attn_metadata, dict): - attn_metadata = _select_attn_metadata(attn_metadata) or _empty_attn_metadata() - if not hasattr(attn_metadata, "num_prefills"): - attn_metadata = types.SimpleNamespace( - num_prefills=0, - num_prefill_tokens=0, - context_lens_tensor=None, - query_start_loc=None, - ) - safe_metadata = _build_safe_legacy_metadata(chunk_size, attn_metadata, mamba2_metadata) - if safe_metadata is not None: - return safe_metadata - return original(chunk_size, attn_metadata, mamba2_metadata) - - prepare_mamba2_metadata._easymagpie_compat = True # type: ignore[attr-defined] - module.prepare_mamba2_metadata = prepare_mamba2_metadata - for module_name in ("vllm.model_executor.models.nemotron_h",): - loaded_module = sys.modules.get(module_name) - if loaded_module is not None and getattr(loaded_module, "prepare_mamba2_metadata", None) is original: - loaded_module.prepare_mamba2_metadata = prepare_mamba2_metadata - - -def _install_triton_attention_profile_metadata_compat() -> None: - try: - module = importlib.import_module("vllm.v1.attention.backends.triton_attn") - except Exception: - return - impl_cls = getattr(module, "TritonAttentionImpl", None) - original = getattr(impl_cls, "forward", None) - if original is None or getattr(original, "_easymagpie_profile_metadata_compat", False): - return - - def forward( - self, - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale=None, - output_block_scale=None, - ): - if _is_generic_profile_attention_metadata(attn_metadata): - return original( - self, - layer, - query, - key, - value, - kv_cache, - None, - output, - output_scale, - output_block_scale, - ) - return original( - self, - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale, - output_block_scale, - ) - - forward._easymagpie_profile_metadata_compat = True # type: ignore[attr-defined] - impl_cls.forward = forward - - -def _install_flash_attention_profile_metadata_compat() -> None: - try: - module = importlib.import_module("vllm.v1.attention.backends.flash_attn") - except Exception: - return - impl_cls = getattr(module, "FlashAttentionImpl", None) - original = getattr(impl_cls, "forward", None) - if original is None or getattr(original, "_easymagpie_profile_metadata_compat", False): - return - - def forward( - self, - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output=None, - output_scale=None, - output_block_scale=None, - **kwargs, - ): - optional_kwargs = _filter_supported_kwargs( - original, - {"output_block_scale": output_block_scale, **kwargs}, - ) - if _is_generic_profile_attention_metadata(attn_metadata) or _is_empty_attention_kv_cache(kv_cache): - return original( - self, - layer, - query, - key, - value, - kv_cache, - None, - output, - output_scale, - **optional_kwargs, - ) - return original( - self, - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale, - **optional_kwargs, - ) - - forward._easymagpie_profile_metadata_compat = True # type: ignore[attr-defined] - impl_cls.forward = forward - - -def _install_flashinfer_attention_profile_metadata_compat() -> None: - try: - module = importlib.import_module("vllm.v1.attention.backends.flashinfer") - except Exception: - return - impl_cls = getattr(module, "FlashInferImpl", None) - original = getattr(impl_cls, "forward", None) - if original is None or getattr(original, "_easymagpie_profile_metadata_compat", False): - return - - def forward( - self, - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output=None, - output_scale=None, - output_block_scale=None, - **kwargs, - ): - optional_kwargs = _filter_supported_kwargs( - original, - {"output_block_scale": output_block_scale, **kwargs}, - ) - if _is_generic_profile_attention_metadata(attn_metadata) or _is_empty_attention_kv_cache(kv_cache): - return original( - self, - layer, - query, - key, - value, - kv_cache, - None, - output, - output_scale, - **optional_kwargs, - ) - return original( - self, - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - output_scale, - **optional_kwargs, - ) - - forward._easymagpie_profile_metadata_compat = True # type: ignore[attr-defined] - impl_cls.forward = forward - - -def _install_mamba2_profile_no_kv_cache_compat() -> None: - try: - from vllm import envs - from vllm.forward_context import get_forward_context - - module = importlib.import_module("vllm.model_executor.layers.mamba.mamba_mixer2") - except Exception: - return - mixer_cls = getattr(module, "MambaMixer2", None) - original = getattr(mixer_cls, "forward_cuda", None) - if mixer_cls is None: - return - - def _has_usable_mamba_kv_cache_for_conv(self: Any) -> bool: - kv_cache = getattr(self, "kv_cache", None) - try: - if kv_cache is None: - return False - if len(kv_cache) == 1 and isinstance(kv_cache[0], (list, tuple)): - kv_cache = kv_cache[0] - if len(kv_cache) < 2: - return False - conv_cache = kv_cache[0] - ssm_cache = kv_cache[1] - return ( - getattr(conv_cache, "ndim", 0) >= 2 - and getattr(ssm_cache, "ndim", 0) >= 2 - ) - except Exception: - return False - - def _is_profile_only_mamba_metadata(attn_metadata: Any) -> bool: - if attn_metadata is None: - return False - if getattr(attn_metadata, "batch_ptr", None) is not None: - return False - if getattr(attn_metadata, "token_chunk_offset_ptr", None) is not None: - return False - num_decode_tokens = int(getattr(attn_metadata, "num_decode_tokens", 0) or 0) - num_prefill_tokens = int(getattr(attn_metadata, "num_prefill_tokens", 0) or 0) - return num_decode_tokens == 0 and num_prefill_tokens > 0 - - def _with_v1_mamba_metadata_dict( - self, - callback, - *args, - allow_profile_metadata_fallback: bool = False, - **kwargs, - ): - try: - use_v1 = bool(getattr(envs, "VLLM_USE_V1", False)) - except Exception: - use_v1 = False - if not use_v1: - return callback(*args, **kwargs) - - try: - forward_context = get_forward_context() - except Exception: - forward_context = None - attn_metadata = getattr(forward_context, "attn_metadata", None) - prefix = getattr(self, "prefix", None) - selected_mamba_metadata = _select_mamba2_attention_metadata(attn_metadata) - if selected_mamba_metadata is None: - projected_states = kwargs.get("projected_states") - if projected_states is None and args: - projected_states = args[0] - try: - num_tokens = max(1, int(getattr(projected_states, "shape", [1])[0] or 1)) - device = getattr(projected_states, "device", None) - if attn_metadata is not None and hasattr(attn_metadata, "num_prefills"): - flash_metadata = attn_metadata - else: - flash_metadata = _build_vllm_attention_metadata_from_lengths( - num_tokens=num_tokens, - num_reqs=1, - max_query_len=num_tokens, - device=device, - ) - chunk_size = int(getattr(self, "chunk_size", None) or getattr(self, "chunk_size_padded", None) or 256) - selected_mamba_metadata = _build_profile_mamba2_attention_metadata(flash_metadata, chunk_size) - except Exception: - selected_mamba_metadata = None - if prefix is None or selected_mamba_metadata is None or forward_context is None: - return callback(*args, **kwargs) - - if allow_profile_metadata_fallback and ( - not _has_usable_mamba_kv_cache_for_conv(self) - or _is_profile_only_mamba_metadata(selected_mamba_metadata) - ): - previous_attn_metadata = forward_context.attn_metadata - forward_context.attn_metadata = None - try: - return callback(*args, **kwargs) - finally: - forward_context.attn_metadata = previous_attn_metadata - - selected_mamba_metadata = _repair_mamba2_attention_metadata_state_indices(selected_mamba_metadata) - replacement_attn_metadata = None - if isinstance(attn_metadata, dict): - layer_metadata = attn_metadata.get(prefix) - if layer_metadata is not selected_mamba_metadata: - replacement_attn_metadata = dict(attn_metadata) - replacement_attn_metadata[prefix] = selected_mamba_metadata - else: - replacement_attn_metadata = {prefix: selected_mamba_metadata} - - if replacement_attn_metadata is None: - return callback(*args, **kwargs) - - previous_attn_metadata = forward_context.attn_metadata - forward_context.attn_metadata = replacement_attn_metadata - try: - return callback(*args, **kwargs) - finally: - forward_context.attn_metadata = previous_attn_metadata - - original_conv_ssm_forward = getattr(mixer_cls, "conv_ssm_forward", None) - if original_conv_ssm_forward is not None and not getattr( - original_conv_ssm_forward, "_easymagpie_v1_attn_metadata_compat", False - ): - - def conv_ssm_forward(self, *args, **kwargs): - return _with_v1_mamba_metadata_dict( - self, - lambda *call_args, **call_kwargs: original_conv_ssm_forward(self, *call_args, **call_kwargs), - *args, - allow_profile_metadata_fallback=True, - **kwargs, - ) - - conv_ssm_forward._easymagpie_v1_attn_metadata_compat = True # type: ignore[attr-defined] - mixer_cls.conv_ssm_forward = conv_ssm_forward - - original_forward = getattr(mixer_cls, "forward", None) - if original_forward is not None and not getattr(original_forward, "_easymagpie_v1_attn_metadata_compat", False): - - def forward(self, *args, **kwargs): - return _with_v1_mamba_metadata_dict( - self, - lambda *call_args, **call_kwargs: original_forward(self, *call_args, **call_kwargs), - *args, - **kwargs, - ) - - forward._easymagpie_v1_attn_metadata_compat = True # type: ignore[attr-defined] - mixer_cls.forward = forward - - if original is None or getattr(original, "_easymagpie_no_kv_profile_compat", False): - return - - try: - original_signature = inspect.signature(original) - except Exception: - original_signature = None - - def has_explicit_cache_args(self: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: - positional_cache_args = ( - (len(args) > 2 and args[2] is not None) - or (len(args) > 3 and args[3] is not None) - ) - keyword_cache_args = ( - kwargs.get("mamba_cache_params") is not None - or kwargs.get("mamba2_metadata") is not None - ) - if original_signature is None: - return positional_cache_args or keyword_cache_args - try: - bound = original_signature.bind_partial(self, *args, **kwargs) - except Exception: - return positional_cache_args or keyword_cache_args - return ( - bound.arguments.get("mamba_cache_params") is not None - or bound.arguments.get("mamba2_metadata") is not None - or positional_cache_args - or keyword_cache_args - ) - - def bound_forward_argument(self: Any, args: tuple[Any, ...], kwargs: dict[str, Any], name: str) -> Any: - if original_signature is not None: - try: - bound = original_signature.bind_partial(self, *args, **kwargs) - value = bound.arguments.get(name) - if value is not None: - return value - except Exception: - pass - positional_index = { - "hidden_states": 0, - "output": 1, - "mamba_cache_params": 2, - "mamba2_metadata": 3, - "mup_vector": 4, - }.get(name) - if positional_index is not None and len(args) > positional_index: - return args[positional_index] - return kwargs.get(name) - - def replace_forward_argument( - args: tuple[Any, ...], kwargs: dict[str, Any], name: str, value: Any - ) -> tuple[tuple[Any, ...], dict[str, Any]]: - positional_index = { - "hidden_states": 0, - "output": 1, - "mamba_cache_params": 2, - "mamba2_metadata": 3, - "mup_vector": 4, - }.get(name) - if positional_index is not None and len(args) > positional_index: - new_args = list(args) - new_args[positional_index] = value - return tuple(new_args), kwargs - new_kwargs = dict(kwargs) - new_kwargs[name] = value - return args, new_kwargs - - def forward_cuda(self, *args, **kwargs): - try: - use_v1 = bool(getattr(envs, "VLLM_USE_V1", False)) - except Exception: - use_v1 = False - - forward_context = None - original_attn_metadata = None - selected_mamba_metadata = None - restore_attn_metadata = False - try: - forward_context = get_forward_context() - except Exception: - forward_context = None - original_attn_metadata = getattr(forward_context, "attn_metadata", None) - normalized_attn_metadata = _normalize_mamba2_attention_metadata_groups(original_attn_metadata) - prefix = getattr(self, "prefix", None) - force_v0_mamba_forward = False - selected_from_legacy_mapping = False - if isinstance(normalized_attn_metadata, dict) and prefix is not None: - layer_metadata = normalized_attn_metadata.get(prefix) - selected_mamba_metadata = _select_mamba2_attention_metadata(layer_metadata) - if selected_mamba_metadata is None: - legacy_metadata = bound_forward_argument(self, args, kwargs, "mamba2_metadata") - if legacy_metadata is None: - legacy_metadata = _select_legacy_mamba2_metadata(layer_metadata) - if legacy_metadata is None: - legacy_metadata = layer_metadata - selected_mamba_metadata = _build_mamba2_attention_metadata_from_legacy( - legacy_metadata, - bound_forward_argument(self, args, kwargs, "mamba_cache_params"), - bound_forward_argument(self, args, kwargs, "hidden_states"), - ) - force_v0_mamba_forward = selected_mamba_metadata is None and isinstance(layer_metadata, dict) - selected_from_legacy_mapping = selected_mamba_metadata is not None and isinstance(layer_metadata, dict) - if selected_mamba_metadata is not None: - selected_mamba_metadata = _repair_mamba2_attention_metadata_state_indices(selected_mamba_metadata) - if selected_mamba_metadata is not None and selected_mamba_metadata is not layer_metadata: - normalized_attn_metadata = dict(normalized_attn_metadata) - normalized_attn_metadata[prefix] = selected_mamba_metadata - elif _is_mamba2_attention_metadata(normalized_attn_metadata): - selected_mamba_metadata = _repair_mamba2_attention_metadata_state_indices(normalized_attn_metadata) - normalized_attn_metadata = selected_mamba_metadata - if forward_context is not None: - replacement_attn_metadata = normalized_attn_metadata - if use_v1 and selected_mamba_metadata is not None and not isinstance(replacement_attn_metadata, dict): - prefix = getattr(self, "prefix", None) - if prefix is not None: - replacement_attn_metadata = {prefix: selected_mamba_metadata} - elif not use_v1 and selected_mamba_metadata is not None: - replacement_attn_metadata = selected_mamba_metadata - if replacement_attn_metadata is not original_attn_metadata: - forward_context.attn_metadata = replacement_attn_metadata - restore_attn_metadata = True - - call_args = args - call_kwargs = kwargs - if selected_mamba_metadata is not None: - cache_params = bound_forward_argument(self, args, kwargs, "mamba_cache_params") - repaired_cache_params = _repair_mamba_cache_params_state_indices(cache_params, selected_mamba_metadata) - if repaired_cache_params is not cache_params: - call_args, call_kwargs = replace_forward_argument( - args, kwargs, "mamba_cache_params", repaired_cache_params - ) - - explicit_cache_args = has_explicit_cache_args(self, call_args, call_kwargs) - try: - if ( - use_v1 - and forward_context is not None - and selected_mamba_metadata is not None - and not explicit_cache_args - and _is_profile_only_mamba_metadata(selected_mamba_metadata) - ): - previous_attn_metadata = getattr(forward_context, "attn_metadata", None) - forward_context.attn_metadata = None - try: - return original(self, *call_args, **call_kwargs) - finally: - forward_context.attn_metadata = previous_attn_metadata - if ( - use_v1 - and forward_context is not None - and (explicit_cache_args or force_v0_mamba_forward) - and (selected_mamba_metadata is not None or force_v0_mamba_forward) - ): - previous_use_v1 = getattr(envs, "VLLM_USE_V1", None) - previous_attn_metadata = getattr(forward_context, "attn_metadata", None) - envs.VLLM_USE_V1 = False - if selected_mamba_metadata is not None: - if selected_from_legacy_mapping and prefix is not None: - forward_context.attn_metadata = {prefix: selected_mamba_metadata} - else: - forward_context.attn_metadata = selected_mamba_metadata - try: - return original(self, *call_args, **call_kwargs) - finally: - envs.VLLM_USE_V1 = previous_use_v1 - forward_context.attn_metadata = previous_attn_metadata - if ( - selected_mamba_metadata is not None - and forward_context is not None - and hasattr(self, "kv_cache") - and not explicit_cache_args - ): - previous_use_v1 = getattr(envs, "VLLM_USE_V1", None) - previous_attn_metadata = getattr(forward_context, "attn_metadata", None) - prefix = getattr(self, "prefix", None) - if isinstance(normalized_attn_metadata, dict): - v1_attn_metadata = normalized_attn_metadata - elif prefix is not None: - v1_attn_metadata = {prefix: selected_mamba_metadata} - else: - v1_attn_metadata = previous_attn_metadata - envs.VLLM_USE_V1 = True - forward_context.attn_metadata = v1_attn_metadata - try: - return original(self, *call_args, **call_kwargs) - finally: - envs.VLLM_USE_V1 = previous_use_v1 - forward_context.attn_metadata = previous_attn_metadata - return original(self, *call_args, **call_kwargs) - finally: - if restore_attn_metadata and forward_context is not None: - try: - forward_context.attn_metadata = original_attn_metadata - except Exception: - pass - - forward_cuda._easymagpie_no_kv_profile_compat = True # type: ignore[attr-defined] - mixer_cls.forward_cuda = forward_cuda - - -def _install_lora_config_alias() -> None: - try: - from vllm.config import LoRAConfig - except Exception: - return - module = types.ModuleType("vllm.config.lora") - module.LoRAConfig = LoRAConfig - module.MaxLoRARanks = Literal[1, 8, 16, 32, 64, 128, 256, 320, 512] - sys.modules.setdefault("vllm.config.lora", module) - - -def _install_config_decorator_compat() -> None: - try: - import vllm.config.utils as config_utils - except Exception: - return - original = getattr(config_utils, "config", None) - if original is None or getattr(original, "_easymagpie_compat", False): - return - - def compat_config(cls=None, **kwargs): - def _decorate(real_cls): - try: - return original(real_cls, **kwargs) - except TypeError: - if not kwargs: - raise - return original(real_cls) - - if cls is None: - return _decorate - return _decorate(cls) - - compat_config._easymagpie_compat = True # type: ignore[attr-defined] - config_utils.config = compat_config - - -def _install_model_arch_convertor() -> None: - if "vllm.transformers_utils.model_arch_config_convertor" in sys.modules: - return - module = types.ModuleType("vllm.transformers_utils.model_arch_config_convertor") - - class ModelArchConfigConvertorBase: - def __init__(self, hf_config: Any, hf_text_config: Any) -> None: - self.hf_config = hf_config - self.hf_text_config = hf_text_config - - def _normalize_quantization_config(self, config: Any) -> Any: - if config is None: - return None - if isinstance(config, dict): - return config.get("quantization_config") - return getattr(config, "quantization_config", None) - - def get_quantization_config(self) -> Any: - return self._normalize_quantization_config(self.hf_text_config) or self._normalize_quantization_config( - self.hf_config - ) - - def convert(self) -> dict[str, Any]: - quant_config = self.get_quantization_config() - return {"quantization_config": quant_config} if quant_config is not None else {} - - module.ModelArchConfigConvertorBase = ModelArchConfigConvertorBase - sys.modules["vllm.transformers_utils.model_arch_config_convertor"] = module - - -def _install_io_processor_stub() -> None: - if "vllm.plugins.io_processors" in sys.modules: - return - module = types.ModuleType("vllm.plugins.io_processors") - - def get_io_processor(*_args: Any, **_kwargs: Any) -> None: - return None - - module.get_io_processor = get_io_processor - sys.modules["vllm.plugins.io_processors"] = module - - -def _install_tokenizer_aliases() -> None: - try: - import vllm.transformers_utils.tokenizer as tokenizer_module - except Exception: - return - if not hasattr(tokenizer_module, "TokenizerLike"): - tokenizer_module.TokenizerLike = getattr(tokenizer_module, "AnyTokenizer", object) - sys.modules.setdefault("vllm.tokenizers", tokenizer_module) - try: - import vllm.transformers_utils.tokenizers.mistral as mistral_module - except Exception: - return - sys.modules.setdefault("vllm.tokenizers.mistral", mistral_module) - - -def _install_repo_utils_alias() -> None: - if "vllm.transformers_utils.repo_utils" in sys.modules: - return - try: - from vllm.transformers_utils.config import file_or_path_exists - except Exception: - return - module = types.ModuleType("vllm.transformers_utils.repo_utils") - module.file_or_path_exists = file_or_path_exists - sys.modules["vllm.transformers_utils.repo_utils"] = module - - -def _install_config_parser_aliases() -> None: - try: - import vllm.transformers_utils.config as config_module - except Exception: - return - - if not hasattr(config_module, "MistralConfigParser"): - - class MistralConfigParser: - pass - - config_module.MistralConfigParser = MistralConfigParser - - if not hasattr(config_module, "register_config_parser"): - - def register_config_parser(_name: str): - def decorator(cls): - return cls - - return decorator - - config_module.register_config_parser = register_config_parser - - -def _install_nemotron_h_auto_config() -> None: - try: - from transformers import AutoConfig - from vllm.transformers_utils.configs import NemotronHConfig - except Exception: - return - try: - AutoConfig.register("nemotron_h", NemotronHConfig) - except ValueError: - pass - - -def _register_easy_magpie_plugin() -> None: - try: - from vllm_plugin_easymagpie_omni import register - except Exception: - return - register() - - -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_renderer_aliases() -> None: - if "vllm.renderers" not in sys.modules: - renderers = types.ModuleType("vllm.renderers") - renderers.__path__ = [] # type: ignore[attr-defined] - - class BaseRenderer: - pass - - def merge_kwargs(*items: Any) -> dict[str, Any]: - merged: dict[str, Any] = {} - for item in items: - if item: - merged.update(dict(item)) - return merged - - def renderer_from_config(*_args: Any, **_kwargs: Any) -> None: - return None - - renderers.BaseRenderer = BaseRenderer - renderers.merge_kwargs = merge_kwargs - renderers.renderer_from_config = renderer_from_config - sys.modules["vllm.renderers"] = renderers - else: - renderers = sys.modules["vllm.renderers"] - if not hasattr(renderers, "__path__"): - renderers.__path__ = [] # type: ignore[attr-defined] - inputs = sys.modules.get("vllm.renderers.inputs") - if inputs is None: - inputs = types.ModuleType("vllm.renderers.inputs") - sys.modules["vllm.renderers.inputs"] = inputs - inputs.__package__ = "vllm.renderers.inputs" - inputs.__path__ = [] # type: ignore[attr-defined] - try: - from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt - except Exception: - try: - from vllm.inputs.data import EmbedsPrompt, TextPrompt, TokensPrompt - except Exception: - EmbedsPrompt = dict # type: ignore[assignment] - TextPrompt = dict # type: ignore[assignment] - TokensPrompt = dict # type: ignore[assignment] - - try: - from typing import TypedDict - - class EncoderDecoderDictPrompt(TypedDict): - encoder_prompt: Any - decoder_prompt: Any | None - - class EncoderDecoderTokPrompt(TypedDict): - encoder_prompt: Any - decoder_prompt: Any | None - - except Exception: - EncoderDecoderDictPrompt = dict # type: ignore[assignment] - EncoderDecoderTokPrompt = dict # type: ignore[assignment] - - for name, value in { - "EmbedsPrompt": EmbedsPrompt, - "TextPrompt": TextPrompt, - "TokensPrompt": TokensPrompt, - "DecoderDictPrompt": TextPrompt | TokensPrompt, - "DecoderOnlyDictPrompt": TextPrompt | TokensPrompt | EmbedsPrompt, - "DecoderOnlyTokPrompt": TokensPrompt | EmbedsPrompt, - "DecoderTokPrompt": TokensPrompt, - "DictPrompt": dict, - "EncoderDictPrompt": TextPrompt | TokensPrompt, - "EncoderTokPrompt": TokensPrompt, - "SingletonDictPrompt": dict, - "SingletonTokPrompt": dict, - "TokPrompt": dict, - "EncoderDecoderDictPrompt": EncoderDecoderDictPrompt, - "EncoderDecoderTokPrompt": EncoderDecoderTokPrompt, - }.items(): - if not hasattr(inputs, name): - setattr(inputs, name, value) - - preprocess = sys.modules.get("vllm.renderers.inputs.preprocess") - if preprocess is None: - preprocess = types.ModuleType("vllm.renderers.inputs.preprocess") - preprocess.__package__ = "vllm.renderers.inputs" - sys.modules["vllm.renderers.inputs.preprocess"] = preprocess - inputs.preprocess = preprocess - - def _is_list_of(value: object, item_type: type) -> bool: - return isinstance(value, list) and all(isinstance(item, item_type) for item in value) - - def _validate_prompt_dict(prompt: Any) -> None: - if "prompt" not in prompt or "prompt_token_ids" in prompt or "prompt_embeds" in prompt: - return - if not isinstance(prompt["prompt"], str): - raise TypeError("Prompt text should be a string") - - def _parse_singleton_prompt(prompt: object, *, allow_embeds: bool = True) -> Any: - if isinstance(prompt, str): - return TextPrompt(prompt=prompt) - if _is_list_of(prompt, int): - return TokensPrompt(prompt_token_ids=prompt) - if isinstance(prompt, dict): - if "encoder_prompt" in prompt: - raise TypeError("Cannot pass encoder-decoder prompt to a singleton prompt parser") - _validate_prompt_dict(prompt) - if "prompt" in prompt or "prompt_token_ids" in prompt or (allow_embeds and "prompt_embeds" in prompt): - return prompt - expected = "text, tokens, or embeddings" if allow_embeds else "text or tokens" - raise TypeError(f"Prompt dictionary must contain {expected}") - raise TypeError("Prompt should be a string, list of tokens, or dictionary") - - def parse_dec_only_prompt(prompt: object) -> Any: - if isinstance(prompt, dict) and "encoder_prompt" in prompt: - raise TypeError("Cannot pass encoder-decoder prompt to decoder-only models") - return _parse_singleton_prompt(prompt) - - def parse_enc_dec_prompt(prompt: object) -> Any: - if isinstance(prompt, dict) and "encoder_prompt" in prompt: - enc_prompt = prompt["encoder_prompt"] - dec_prompt = prompt.get("decoder_prompt") - else: - enc_prompt = prompt - dec_prompt = None - return EncoderDecoderDictPrompt( - encoder_prompt=_parse_singleton_prompt(enc_prompt, allow_embeds=False), - decoder_prompt=None if dec_prompt is None else _parse_singleton_prompt(dec_prompt, allow_embeds=False), - ) - - for name in ( - "DecoderDictPrompt", - "DecoderOnlyDictPrompt", - "DictPrompt", - "EncoderDictPrompt", - "EncoderDecoderDictPrompt", - "SingletonDictPrompt", - "TextPrompt", - "TokensPrompt", - "EmbedsPrompt", - ): - setattr(preprocess, name, getattr(inputs, name)) - preprocess.parse_dec_only_prompt = parse_dec_only_prompt - preprocess.parse_enc_dec_prompt = parse_enc_dec_prompt - elif not hasattr(inputs, "preprocess"): - inputs.preprocess = preprocess - - tokenize = sys.modules.get("vllm.renderers.inputs.tokenize") - if tokenize is None: - tokenize = types.ModuleType("vllm.renderers.inputs.tokenize") - tokenize.__package__ = "vllm.renderers.inputs" - sys.modules["vllm.renderers.inputs.tokenize"] = tokenize - inputs.tokenize = tokenize - for name in ( - "DecoderOnlyTokPrompt", - "DecoderTokPrompt", - "EncoderTokPrompt", - "EncoderDecoderTokPrompt", - "SingletonTokPrompt", - "TokPrompt", - "TokensPrompt", - "EmbedsPrompt", - ): - setattr(tokenize, name, getattr(inputs, name)) - elif not hasattr(inputs, "tokenize"): - inputs.tokenize = tokenize - - -def _install_input_processor_alias() -> None: - try: - from vllm.inputs.preprocess import InputPreprocessor - from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config - from vllm.v1.engine.processor import Processor - except Exception: - return - - original_init = InputPreprocessor.__init__ - if not getattr(original_init, "_easymagpie_compat", False): - - def compat_init(self, *args: Any, **kwargs: Any) -> None: - renderer = kwargs.pop("renderer", None) - if "vllm_config" in kwargs: - vllm_config = kwargs.pop("vllm_config") - tokenizer = kwargs.pop("tokenizer", None) - if tokenizer is None and not getattr(vllm_config.model_config, "skip_tokenizer_init", False): - tokenizer = cached_tokenizer_from_config(model_config=vllm_config.model_config) - original_init(self, vllm_config.model_config, tokenizer, *args, **kwargs) - self.renderer = renderer - return - original_init(self, *args, **kwargs) - self.renderer = getattr(self, "renderer", renderer) - - compat_init._easymagpie_compat = True # type: ignore[attr-defined] - InputPreprocessor.__init__ = compat_init - - if "vllm.v1.engine.input_processor" in sys.modules: - return - - module = types.ModuleType("vllm.v1.engine.input_processor") - - class InputProcessor(Processor): - def __init__(self, *, vllm_config: Any, tokenizer: Any = None, mm_registry: Any = None, **_kwargs: Any) -> None: - if tokenizer is None and not getattr(vllm_config.model_config, "skip_tokenizer_init", False): - tokenizer = cached_tokenizer_from_config(model_config=vllm_config.model_config) - if mm_registry is None: - super().__init__(vllm_config, tokenizer) - else: - super().__init__(vllm_config, tokenizer, mm_registry) - self.renderer = getattr(self.input_preprocessor, "renderer", None) - - module.InputProcessor = InputProcessor - sys.modules["vllm.v1.engine.input_processor"] = module - - -def _install_processor_process_inputs_compat() -> None: - try: - from vllm.v1.engine.processor import Processor - except Exception: - return - - original_process_inputs = Processor.process_inputs - if getattr(original_process_inputs, "_easymagpie_compat", False): - return - - try: - signature = inspect.signature(original_process_inputs) - parameters = signature.parameters.values() - accepts_supported_tasks = "supported_tasks" in signature.parameters or any( - parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters - ) - except Exception: - accepts_supported_tasks = False - - if accepts_supported_tasks: - return - - def compat_process_inputs(self, *args: Any, **kwargs: Any): - has_omni_supported_tasks = "supported_tasks" in kwargs - kwargs.pop("supported_tasks", None) - result = original_process_inputs(self, *args, **kwargs) - if has_omni_supported_tasks and isinstance(result, tuple) and len(result) == 2: - return result[1] - return result - - compat_process_inputs._easymagpie_compat = True # type: ignore[attr-defined] - Processor.process_inputs = compat_process_inputs - - -def _install_omni_input_preprocessor_signature_compat() -> None: - try: - from vllm_omni.inputs.preprocess import OmniInputPreprocessor - except Exception: - return - - original_prompt_to_inputs = OmniInputPreprocessor._prompt_to_llm_inputs - if getattr(original_prompt_to_inputs, "_easymagpie_compat", False): - return - - try: - signature = inspect.signature(original_prompt_to_inputs) - parameters = signature.parameters - accepts_kwargs = any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()) - except Exception: - signature = None - parameters = {} - accepts_kwargs = False - - if accepts_kwargs or {"lora_request", "return_mm_hashes"}.issubset(parameters): - return - - accepted_kwargs = { - name - for name, parameter in parameters.items() - if name != "self" and parameter.kind in {inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY} - } - - def compat_prompt_to_inputs(self, *args: Any, **kwargs: Any): - if signature is not None: - kwargs = {key: value for key, value in kwargs.items() if key in accepted_kwargs} - else: - kwargs.pop("lora_request", None) - kwargs.pop("return_mm_hashes", None) - return original_prompt_to_inputs(self, *args, **kwargs) - - compat_prompt_to_inputs._easymagpie_compat = True # type: ignore[attr-defined] - OmniInputPreprocessor._prompt_to_llm_inputs = compat_prompt_to_inputs - - -def _install_input_preprocessor_truncate_compat() -> None: - try: - from vllm.inputs.preprocess import InputPreprocessor - except Exception: - return - - if hasattr(InputPreprocessor, "_truncate_inputs"): - return - - def _truncate_inputs(self, prompt_token_ids: Any, tokenization_kwargs: dict[str, Any] | None = None): - limit = None - if tokenization_kwargs: - limit = tokenization_kwargs.get("max_length", tokenization_kwargs.get("truncate_prompt_tokens")) - if limit is None: - return prompt_token_ids - try: - limit_int = int(limit) - except Exception: - return prompt_token_ids - if limit_int <= 0: - return prompt_token_ids - return prompt_token_ids[-limit_int:] - - InputPreprocessor._truncate_inputs = _truncate_inputs - - -def _install_omni_engine_request_compat() -> None: - try: - from vllm.v1.engine import EngineCoreRequest - from vllm_omni.engine import OmniEngineCoreRequest - except Exception: - return - - for request_cls in (EngineCoreRequest, OmniEngineCoreRequest): - if not hasattr(request_cls, "prompt_embeds"): - request_cls.prompt_embeds = None - - try: - async_omni_engine = importlib.import_module("vllm_omni.engine.async_omni_engine") - from vllm_omni.engine.serialization import serialize_additional_information - except Exception: - return - - try: - request_signature = inspect.signature(OmniEngineCoreRequest) - except Exception: - request_signature = None - - if not getattr(async_omni_engine._upgrade_to_omni_request, "_easymagpie_compat", False): - - def compat_upgrade_to_omni_request(request: Any, raw_prompt: Any) -> Any: - prompt_embeds = getattr(request, "prompt_embeds", None) - additional_information = None - - if isinstance(raw_prompt, dict): - if prompt_embeds is None: - raw_prompt_embeds = raw_prompt.get("prompt_embeds") - try: - import torch - - if isinstance(raw_prompt_embeds, torch.Tensor): - prompt_embeds = raw_prompt_embeds - except Exception: - pass - additional_information = serialize_additional_information( - raw_prompt.get("additional_information"), - log_prefix="AsyncOmniEngine", - ) - - if prompt_embeds is None and additional_information is None: - return request - - if request_signature is None: - return request - - values: dict[str, Any] = {} - for name, parameter in request_signature.parameters.items(): - if name == "prompt_embeds": - values[name] = prompt_embeds - elif name == "additional_information": - values[name] = additional_information - elif hasattr(request, name): - values[name] = getattr(request, name) - elif parameter.default is inspect.Parameter.empty: - values[name] = None - - return OmniEngineCoreRequest(**values) - - compat_upgrade_to_omni_request._easymagpie_compat = True # type: ignore[attr-defined] - async_omni_engine._upgrade_to_omni_request = compat_upgrade_to_omni_request - - async_engine_cls = getattr(async_omni_engine, "AsyncOmniEngine", None) - original_build = getattr(async_engine_cls, "_build_add_request_message", None) - if original_build is None or getattr(original_build, "_easymagpie_external_req_compat", False): - return - try: - original_build_signature = inspect.signature(original_build) - except Exception: - original_build_signature = None - if original_build_signature is not None: - original_build_params = original_build_signature.parameters - if "prompt_text" in original_build_params and "message_type" in original_build_params: - # Newer vLLM-Omni already handles request construction, prompt text, - # final-stage metadata, and orchestrator admission correctly. The - # patched _upgrade_to_omni_request above is enough for EasyMagpie's - # prompt embeddings/additional-information fields. - return - - engine_core_request_cls = EngineCoreRequest - inject_global_id = getattr(async_omni_engine, "_inject_global_id", None) - - def compat_build_add_request_message( - self: Any, - request_id: str, - prompt: Any, - sampling_params_list: Any = None, - final_stage_id: int = 0, - arrival_time: float | None = None, - ) -> dict[str, Any]: - effective_sampling_params_list = ( - list(sampling_params_list) if sampling_params_list is not None else list(self.default_sampling_params_list) - ) - if not effective_sampling_params_list: - raise ValueError( - f"Missing sampling params for stage 0. Got {len(effective_sampling_params_list)} stage params." - ) - params = effective_sampling_params_list[0] - original_prompt = prompt - - stage_type = self.stage_metadata[0].get("stage_type") - if stage_type != "diffusion" and not isinstance(prompt, engine_core_request_cls): - if inject_global_id is not None: - if isinstance(prompt, dict): - inject_global_id(prompt, request_id) - elif isinstance(prompt, list): - for item in prompt: - inject_global_id(item, request_id) - - request = self.input_processor.process_inputs( - request_id=request_id, - prompt=prompt, - params=params, - supported_tasks=self.supported_tasks, - arrival_time=arrival_time, - ) - request = async_omni_engine._upgrade_to_omni_request(request, prompt) - try: - request.external_req_id = request_id - except AttributeError: - pass - - self.output_processors[0].add_request( - request=request, - prompt=prompt, - parent_req=None, - request_index=0, - queue=None, - ) - prompt = request - - return { - "type": "add_request", - "request_id": request_id, - "prompt": prompt, - "original_prompt": original_prompt, - "sampling_params_list": effective_sampling_params_list, - "final_stage_id": final_stage_id, - } - - compat_build_add_request_message._easymagpie_external_req_compat = True # type: ignore[attr-defined] - async_engine_cls._build_add_request_message = compat_build_add_request_message - - -def _install_omni_request_compat() -> None: - try: - from vllm.v1.request import StructuredOutputRequest - import vllm_omni.request as omni_request_module - except Exception: - return - - OmniRequest = getattr(omni_request_module, "OmniRequest", None) - if OmniRequest is None: - return - try: - BaseRequest = OmniRequest.__mro__[1] - except Exception: - return - - original_init = getattr(OmniRequest, "__init__", None) - if original_init is not None and not getattr(original_init, "_easymagpie_request_init_compat", False): - base_init = BaseRequest.__init__ - try: - base_signature = inspect.signature(base_init) - except Exception: - base_signature = None - - def compat_omni_request_init( - self: Any, - prompt_embeds: Any = None, - external_req_id: str | None = None, - additional_information: Any = None, - *args: Any, - **kwargs: Any, - ) -> None: - prompt_embeds_tensor = self._maybe_decode_prompt_embeds(prompt_embeds) - if base_signature is None: - base_init(self, *args, **kwargs) - else: - accepts_kwargs = any( - parameter.kind is inspect.Parameter.VAR_KEYWORD - for parameter in base_signature.parameters.values() - ) - base_kwargs = dict(kwargs) - if "prompt_embeds" in base_signature.parameters: - base_kwargs["prompt_embeds"] = prompt_embeds_tensor - if not accepts_kwargs: - base_kwargs = { - key: value - for key, value in base_kwargs.items() - if key in base_signature.parameters and key != "self" - } - base_init(self, *args, **base_kwargs) - - self.prompt_embeds = prompt_embeds_tensor - self.prompt_embeds_payload = prompt_embeds if prompt_embeds is not prompt_embeds_tensor else None - self.external_req_id = external_req_id - self.additional_information = additional_information - - compat_omni_request_init._easymagpie_request_init_compat = True # type: ignore[attr-defined] - OmniRequest.__init__ = compat_omni_request_init - - original_from_request = getattr(OmniRequest, "from_engine_core_request", None) - raw_from_request = getattr(original_from_request, "__func__", original_from_request) - if raw_from_request is None or getattr(raw_from_request, "_easymagpie_request_convert_compat", False): - return - - def compat_from_engine_core_request(cls: Any, request: Any, block_hasher: Any) -> Any: - mm_kwargs = getattr(request, "mm_kwargs", None) - if mm_kwargs is not None: - mm_kwargs = list(mm_kwargs) - elif getattr(request, "mm_features", None) is not None: - mm_kwargs = list(getattr(request, "mm_features")) - - sampling_params = getattr(request, "sampling_params", None) - structured_output_request = None - if sampling_params is not None: - from_sampling_params = getattr(StructuredOutputRequest, "from_sampling_params", None) - if callable(from_sampling_params): - structured_output_request = from_sampling_params(sampling_params) - else: - try: - structured_output_request = StructuredOutputRequest(sampling_params=sampling_params) - except TypeError: - structured_outputs = getattr(sampling_params, "structured_outputs", None) - if structured_outputs and not structured_outputs.all_constraints_none(): - structured_output_request = StructuredOutputRequest(params=structured_outputs) - - return cls( - request_id=request.request_id, - external_req_id=getattr(request, "external_req_id", None) or request.request_id, - client_index=getattr(request, "client_index", 0), - prompt_token_ids=request.prompt_token_ids, - prompt_embeds=getattr(request, "prompt_embeds", None), - multi_modal_kwargs=mm_kwargs, - multi_modal_hashes=getattr(request, "mm_hashes", None), - multi_modal_placeholders=getattr(request, "mm_placeholders", None), - sampling_params=sampling_params, - pooling_params=getattr(request, "pooling_params", None), - eos_token_id=getattr(request, "eos_token_id", None), - arrival_time=getattr(request, "arrival_time", None), - lora_request=getattr(request, "lora_request", None), - structured_output_request=structured_output_request, - cache_salt=getattr(request, "cache_salt", None), - priority=getattr(request, "priority", 0), - block_hasher=block_hasher, - additional_information=getattr(request, "additional_information", None), - ) - - compat_from_engine_core_request._easymagpie_request_convert_compat = True # type: ignore[attr-defined] - OmniRequest.from_engine_core_request = classmethod(compat_from_engine_core_request) - - -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_import_utils_alias() -> None: - if "vllm.utils.import_utils" in sys.modules: - return - module = types.ModuleType("vllm.utils.import_utils") - - class LazyLoader(types.ModuleType): - def __init__(self, local_name: str, parent_globals: dict[str, Any], name: str) -> None: - super().__init__(local_name) - self._local_name = local_name - self._parent_globals = parent_globals - self._module_name = name - self._module = None - - def _load(self): - if self._module is None: - self._module = importlib.import_module(self._module_name) - self._parent_globals[self._local_name] = self._module - return self._module - - def __getattr__(self, item: str) -> Any: - return getattr(self._load(), item) - - def resolve_obj_by_qualname(qualname: str) -> Any: - module_name, _, attr = qualname.replace(":", ".").rpartition(".") - if not module_name: - raise ValueError(f"Invalid qualified name: {qualname}") - obj = importlib.import_module(module_name) - for part in attr.split("."): - obj = getattr(obj, part) - return obj - - try: - from vllm.utils import import_pynvml - except Exception: - - def import_pynvml(): - return importlib.import_module("pynvml") - - module.LazyLoader = LazyLoader - module.resolve_obj_by_qualname = resolve_obj_by_qualname - module.import_pynvml = import_pynvml - sys.modules["vllm.utils.import_utils"] = module - - -def _install_torch_utils_alias() -> None: - if "vllm.utils.torch_utils" in sys.modules: - return - module = types.ModuleType("vllm.utils.torch_utils") - - def set_random_seed(seed: int | None) -> None: - import numpy as np - import torch - - seed = 0 if seed is None else int(seed) - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - - @contextlib.contextmanager - def set_default_torch_dtype(dtype): - import torch - - old_dtype = torch.get_default_dtype() - torch.set_default_dtype(dtype) - try: - yield - finally: - torch.set_default_dtype(old_dtype) - - try: - from vllm.utils import supports_xccl - except Exception: - - def supports_xccl() -> bool: - return False - - module.set_random_seed = set_random_seed - module.set_default_torch_dtype = set_default_torch_dtype - module.supports_xccl = supports_xccl - sys.modules["vllm.utils.torch_utils"] = module - - -def _install_math_utils_alias() -> None: - if "vllm.utils.math_utils" in sys.modules: - return - try: - from vllm.utils import cdiv - except Exception: - - def cdiv(a: int, b: int) -> int: - return -(a // -b) - - module = types.ModuleType("vllm.utils.math_utils") - module.cdiv = cdiv - sys.modules["vllm.utils.math_utils"] = module - - -def _install_mem_utils_alias() -> None: - if "vllm.utils.mem_utils" in sys.modules: - return - try: - import vllm.utils as vllm_utils - except Exception: - return - - module = types.ModuleType("vllm.utils.mem_utils") - - class MemorySnapshot(vllm_utils.MemorySnapshot): - def __init__(self, *args, device=None, **kwargs) -> None: - super().__init__(*args, **kwargs) - - def format_gib(num_bytes: int | float) -> str: - return f"{float(num_bytes) / float(vllm_utils.GiB_bytes):.2f}" - - module.MemorySnapshot = MemorySnapshot - module.memory_profiling = vllm_utils.memory_profiling - module.format_gib = format_gib - sys.modules["vllm.utils.mem_utils"] = module - - -def _install_vllm_config_profiler_default() -> None: - try: - from vllm.config import VllmConfig - except Exception: - return - if hasattr(VllmConfig, "profiler_config"): - return - - def get_profiler_config(self): - stored = getattr(self, "_easymagpie_profiler_config", None) - if stored is not None: - return stored - additional_config = getattr(self, "additional_config", None) - if isinstance(additional_config, dict): - return additional_config.get("profiler_config") - return getattr(additional_config, "profiler_config", None) - - def set_profiler_config(self, value) -> None: - object.__setattr__(self, "_easymagpie_profiler_config", value) - - VllmConfig.profiler_config = property(get_profiler_config, set_profiler_config) - - -def _install_parallel_config_defaults() -> None: - try: - from vllm.config import ParallelConfig - except Exception: - return - - missing = object() - - def install_property(name: str, default): - if hasattr(ParallelConfig, name): - return - - storage_name = f"_easymagpie_{name}" - - def getter(self): - stored = getattr(self, storage_name, missing) - if stored is not missing: - return stored - return default(self) if callable(default) else default - - def setter(self, value) -> None: - object.__setattr__(self, storage_name, value) - - setattr(ParallelConfig, name, property(getter, setter)) - - install_property( - "nnodes_within_dp", - lambda self: max( - 1, - int(getattr(self, "data_parallel_size", 1)) // int(getattr(self, "data_parallel_size_local", 1)), - ), - ) - install_property("data_parallel_index", lambda self: getattr(self, "data_parallel_rank", 0)) - install_property( - "local_world_size", - lambda self: ( - int(getattr(self, "pipeline_parallel_size", 1)) - * int(getattr(self, "tensor_parallel_size", 1)) - * int(getattr(self, "data_parallel_size_local", 1)) - ), - ) - install_property("enable_dbo", False) - install_property("num_ubatches", 1) - install_property("use_ubatching", False) - - -def _install_cache_config_defaults() -> None: - try: - from vllm.config import CacheConfig - except Exception: - return - if hasattr(CacheConfig, "kv_cache_memory_bytes"): - return - - missing = object() - storage_name = "_easymagpie_kv_cache_memory_bytes" - - def get_kv_cache_memory_bytes(self): - stored = getattr(self, storage_name, missing) - if stored is not missing: - return stored - return None - - def set_kv_cache_memory_bytes(self, value) -> None: - object.__setattr__(self, storage_name, value) + 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"] - CacheConfig.kv_cache_memory_bytes = property(get_kv_cache_memory_bytes, set_kv_cache_memory_bytes) +def _install_vllm_multimodal_inputs_alias() -> None: + """Expose vLLM-Omni's legacy plural multimodal input alias on vLLM 0.21+.""" -def _install_platform_dtype_check() -> None: try: - from vllm.platforms import current_platform + multimodal_inputs = importlib.import_module("vllm.multimodal.inputs") except Exception: return - - dtype_check = getattr(current_platform, "check_if_supports_dtype", None) - if callable(dtype_check): + if hasattr(multimodal_inputs, "MultiModalInputs"): return - def check_if_supports_dtype(_dtype) -> None: - return None - - try: - current_platform.check_if_supports_dtype = check_if_supports_dtype - except Exception: - pass - try: - setattr(type(current_platform), "check_if_supports_dtype", staticmethod(check_if_supports_dtype)) - except Exception: - pass - + 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_torch_accelerator_compat() -> None: +def _install_engine_utils_compat() -> None: try: - import torch + import vllm.v1.engine.utils as engine_utils + from vllm.v1.utils import get_engine_client_zmq_addr except Exception: return - accelerator = getattr(torch, "accelerator", None) - if accelerator is None or hasattr(accelerator, "empty_cache"): - return - empty_cache = getattr(getattr(torch, "cuda", None), "empty_cache", None) - if not callable(empty_cache): - return - accelerator.empty_cache = empty_cache + 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)], + ) -def _install_logger_info_once_scope_compat() -> None: - try: - import vllm.logger as vllm_logger - except Exception: - vllm_logger = None + engine_utils.get_engine_zmq_addresses = get_engine_zmq_addresses - if vllm_logger is not None: - original_print = getattr(vllm_logger, "_print_info_once", None) - if original_print is not None and not getattr(original_print, "_easymagpie_compat", False): + 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 + ) - def _print_info_once(logger, msg, *args, **kwargs): - kwargs.pop("scope", None) - return original_print(logger, msg, *args, **kwargs) + class LaunchCoreEnginesContextCompat: + def __init__(self, context_manager: Any) -> None: + self._context_manager = context_manager - _print_info_once._easymagpie_compat = True # type: ignore[attr-defined] - vllm_logger._print_info_once = _print_info_once - methods_to_patch = getattr(vllm_logger, "_METHODS_TO_PATCH", None) - if isinstance(methods_to_patch, dict): - methods_to_patch["info_once"] = _print_info_once + def __enter__(self) -> Any: + value = self._context_manager.__enter__() + if isinstance(value, tuple) and len(value) > 3: + return value[:3] + return value - original = getattr(logging.Logger, "info_once", None) - if original is None or getattr(original, "_easymagpie_compat", False): - return + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: + return self._context_manager.__exit__(exc_type, exc, tb) - def info_once(self, msg, *args, **kwargs): - kwargs.pop("scope", None) - return original(self, msg, *args, **kwargs) + def compat_launch_core_engines(*args: Any, addresses: Any = None, **kwargs: Any): + if not accepts_addresses: + return launch_core_engines(*args, **kwargs) - info_once._easymagpie_compat = True # type: ignore[attr-defined] - logging.Logger.info_once = info_once + 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_v1_serial_utils_dense_tensor_compat() -> None: try: @@ -3639,389 +289,22 @@ def decode_with_easy_magpie_refit_tensors(self: Any, bufs: Any) -> Any: 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 _install_cudagraph_mode_compat() -> None: - try: - from vllm.config import CUDAGraphMode - except Exception: - return - if hasattr(CUDAGraphMode.NONE, "valid_runtime_modes"): - return - - def valid_runtime_modes(self): - return self in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL) - - CUDAGraphMode.valid_runtime_modes = valid_runtime_modes # type: ignore[attr-defined] - - -def _install_cuda_graph_stat_alias() -> None: - try: - import vllm.compilation.cuda_graph as cuda_graph - except Exception: - return - if hasattr(cuda_graph, "CUDAGraphStat"): - return - - class CUDAGraphStat: - pass - - cuda_graph.CUDAGraphStat = CUDAGraphStat - - -def _install_cuda_piecewise_no_sym_shape_compat() -> None: - try: - import vllm.compilation.cuda_piecewise_backend as cuda_piecewise_backend - except Exception: - return - backend_cls = getattr(cuda_piecewise_backend, "PiecewiseBackend", None) - if backend_cls is None: - return - original = getattr(backend_cls, "__call__", None) - if original is None or getattr(original, "_easymagpie_no_sym_shape_compat", False): - return - - def _runtime_shape_from_args(args: tuple[Any, ...]) -> int | None: - for arg in args: - shape = getattr(arg, "shape", None) - if shape is None or len(shape) == 0: - continue - try: - return int(shape[0]) - except Exception: - continue - return None - - def _signature_from_args(args: tuple[Any, ...]) -> tuple[Any, ...]: - signature: list[Any] = [] - for arg in args: - shape = getattr(arg, "shape", None) - stride = getattr(arg, "stride", None) - dtype = getattr(arg, "dtype", None) - device = getattr(arg, "device", None) - if shape is None: - signature.append((type(arg).__qualname__, repr(arg))) - continue - try: - stride_value = tuple(int(x) for x in stride()) if callable(stride) else None - except Exception: - stride_value = None - try: - shape_value = tuple(int(x) for x in shape) - except Exception: - shape_value = tuple(shape) - signature.append( - ( - "tensor", - shape_value, - stride_value, - str(dtype), - str(device), - ) - ) - return tuple(signature) - - def _compile_no_sym_shape(self, args: tuple[Any, ...], runtime_shape: int | None): - compile_fn = getattr(getattr(self, "vllm_backend", None), "compiler_manager", None) - compile_fn = getattr(compile_fn, "compile", None) - if not callable(compile_fn): - return None - return compile_fn( - self.graph, - args, - self.compilation_config.inductor_compile_config, - self.compilation_config, - graph_index=self.piecewise_compile_index, - num_graphs=self.total_piecewise_compiles, - runtime_shape=runtime_shape, - ) - - def __call__(self, *args): - if getattr(self, "sym_shape_indices", None): - return original(self, *args) - signature = _signature_from_args(args) - runtime_shape = _runtime_shape_from_args(args) - runnables = getattr(self, "_easymagpie_no_sym_shape_runnables", None) - if runnables is None: - runnables = {} - self._easymagpie_no_sym_shape_runnables = runnables - if not getattr(self, "first_run_finished", False): - self.first_run_finished = True - runnables[signature] = self.compiled_graph_for_general_shape - check_for_ending_compilation = getattr(self, "check_for_ending_compilation", None) - if callable(check_for_ending_compilation): - check_for_ending_compilation() - self._easymagpie_no_sym_shape_compile_count = 0 - self._easymagpie_no_sym_shape_last_runtime_shape = runtime_shape - self._easymagpie_no_sym_shape_last_signature = signature - return self.compiled_graph_for_general_shape(*args) - runnable = runnables.get(signature) - if runnable is None: - runnable = _compile_no_sym_shape(self, args, runtime_shape) - if runnable is None: - runnable = self.compiled_graph_for_general_shape - runnables[signature] = runnable - self._easymagpie_no_sym_shape_compile_count = int( - getattr(self, "_easymagpie_no_sym_shape_compile_count", 0) - ) + 1 - self._easymagpie_no_sym_shape_last_runtime_shape = runtime_shape - self._easymagpie_no_sym_shape_last_signature = signature - return runnable(*args) - - __call__._easymagpie_no_sym_shape_compat = True # type: ignore[attr-defined] - __call__._easymagpie_original = original # type: ignore[attr-defined] - backend_cls.__call__ = __call__ - - -def _install_kv_connector_stats_alias() -> None: - module_name = "vllm.distributed.kv_transfer.kv_connector.v1.metrics" - if module_name in sys.modules: - return - module = types.ModuleType(module_name) - - class KVConnectorStats: - def aggregate(self, _other): - return self - - module.KVConnectorStats = KVConnectorStats - sys.modules[module_name] = module - - -def _install_perf_stats_alias() -> None: - module_name = "vllm.v1.metrics.perf" - if module_name in sys.modules: - return - module = types.ModuleType(module_name) - - class PerfStats: - pass - - module.PerfStats = PerfStats - sys.modules[module_name] = module - - -def _install_scheduler_make_stats_compat() -> None: - try: - from vllm.v1.core.sched.scheduler import Scheduler - except Exception: - return - original = getattr(Scheduler, "make_stats", None) - if original is None or getattr(original, "_easymagpie_compat", False): - return - - def make_stats(self, spec_decoding_stats=None, *_args, **_kwargs): - return original(self, spec_decoding_stats) - - make_stats._easymagpie_compat = True # type: ignore[attr-defined] - Scheduler.make_stats = make_stats - - -def _install_sched_utils_aliases() -> None: - try: - import vllm.v1.core.sched.utils as sched_utils - except Exception: - return - if not hasattr(sched_utils, "remove_all"): - - def remove_all(items, removed): - removed_set = set(removed) - return [item for item in items if item not in removed_set] - - sched_utils.remove_all = remove_all - - -def _install_sched_interface_aliases() -> None: - try: - import vllm.v1.core.sched.interface as sched_interface - except Exception: - return - if hasattr(sched_interface, "PauseState"): - return - - from enum import Enum - - class PauseState(Enum): - UNPAUSED = 0 - PAUSED_ALL = 1 - - sched_interface.PauseState = PauseState - - -def _install_output_processor_signature_compat() -> None: - try: - from vllm.v1.engine.output_processor import OutputProcessor, RequestState - except Exception: - return - - original_init = getattr(OutputProcessor, "__init__", None) - if original_init is not None and not getattr(original_init, "_easymagpie_compat", False): - - def __init__(self, tokenizer, log_stats: bool, stream_interval: int = 1, tracing_enabled: bool = False): - original_init(self, tokenizer=tokenizer, log_stats=log_stats) - self.stream_interval = int(stream_interval) - self.tracing_enabled = bool(tracing_enabled) - if not hasattr(self, "external_req_ids"): - self.external_req_ids = defaultdict(list) - - __init__._easymagpie_compat = True # type: ignore[attr-defined] - OutputProcessor.__init__ = __init__ - - original_from_new_descriptor = RequestState.__dict__.get("from_new_request") - original_from_new = ( - original_from_new_descriptor.__func__ - if isinstance(original_from_new_descriptor, classmethod) - else original_from_new_descriptor - ) - if original_from_new is None or getattr(original_from_new, "_easymagpie_compat", False): - return - try: - original_from_new_signature = inspect.signature(original_from_new) - except Exception: - original_from_new_signature = None - - def from_new_request(cls, *args, stream_interval: int = 1, **kwargs): - if original_from_new_signature is not None: - param_names = list(original_from_new_signature.parameters) - if "stream_interval" in param_names and "stream_interval" not in kwargs: - stream_interval_arg_index = param_names.index("stream_interval") - 1 - if len(args) <= stream_interval_arg_index: - kwargs["stream_interval"] = stream_interval - state = original_from_new(cls, *args, **kwargs) - if not hasattr(state, "stream_interval"): - state.stream_interval = int(stream_interval) - if not hasattr(state, "sent_tokens_offset"): - state.sent_tokens_offset = 0 - request = kwargs.get("request") - if request is None: - for arg in args: - if hasattr(arg, "request_id") and hasattr(arg, "prompt_token_ids"): - request = arg - break - external_req_id = getattr(request, "external_req_id", None) or getattr(request, "request_id", None) - if external_req_id is not None and not hasattr(state, "external_req_id"): - state.external_req_id = external_req_id - return state - - from_new_request._easymagpie_compat = True # type: ignore[attr-defined] - RequestState.from_new_request = classmethod(from_new_request) - - -def _install_flash_attention_builder_compat() -> None: - try: - import vllm.v1.attention.backends.flash_attn as v1_flash_attn - from vllm.attention.backends.flash_attn import FlashAttentionBackend as LegacyBackend - from vllm.attention.backends.flash_attn import FlashAttentionMetadataBuilder as LegacyBuilder - from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadataBuilder as V1Builder - except Exception: - return - - original_get_sliding_window_configs = getattr(v1_flash_attn, "_get_sliding_window_configs", None) - if original_get_sliding_window_configs is not None and not getattr( - original_get_sliding_window_configs, "_easymagpie_compat", False - ): - - def _get_sliding_window_configs(vllm_config): - sliding_window_configs = set() - layers = v1_flash_attn.get_layers_from_vllm_config(vllm_config, v1_flash_attn.Attention) - for layer in layers.values(): - impl = getattr(layer, "impl", None) - sliding_window_configs.add(getattr(impl, "sliding_window", None)) - return sliding_window_configs - - _get_sliding_window_configs._easymagpie_compat = True # type: ignore[attr-defined] - v1_flash_attn._get_sliding_window_configs = _get_sliding_window_configs - - original_get_builder_cls = getattr(LegacyBackend, "get_builder_cls", None) - if original_get_builder_cls is None or getattr(original_get_builder_cls, "_easymagpie_compat", False): - return - try: - current_builder = LegacyBackend.get_builder_cls() - except Exception: - current_builder = None - if current_builder is not LegacyBuilder: - return - - def get_builder_cls(): - return V1Builder - - get_builder_cls._easymagpie_compat = True # type: ignore[attr-defined] - LegacyBackend.get_builder_cls = staticmethod(get_builder_cls) - - -def _install_worker_utils_compat() -> None: - try: - import vllm.v1.worker.utils as worker_utils - from vllm.utils import GiB_bytes - except Exception: - return - if hasattr(worker_utils, "request_memory"): - request_memory = worker_utils.request_memory - else: - - def request_memory(init_snapshot, cache_config): - requested_memory = init_snapshot.total_memory * cache_config.gpu_memory_utilization - if init_snapshot.free_memory < requested_memory: - gib = lambda value: round(value / GiB_bytes, 2) - raise ValueError( - f"Free memory on device ({gib(init_snapshot.free_memory)}/" - f"{gib(init_snapshot.total_memory)} GiB) on startup is less " - f"than desired GPU memory utilization ({cache_config.gpu_memory_utilization}, " - f"{gib(requested_memory)} GiB). Decrease GPU memory utilization " - "or reduce GPU memory used by other processes." - ) - return requested_memory - - worker_utils.request_memory = request_memory - - if not hasattr(worker_utils, "is_residual_scattered_for_sp"): - - def is_residual_scattered_for_sp(*_args, **_kwargs) -> bool: - return False - - worker_utils.is_residual_scattered_for_sp = is_residual_scattered_for_sp - - -def _install_worker_workspace_alias() -> None: - if "vllm.v1.worker.workspace" in sys.modules: - return - module = types.ModuleType("vllm.v1.worker.workspace") - - def init_workspace_manager(*_args, **_kwargs): - return None - - module.init_workspace_manager = init_workspace_manager - sys.modules["vllm.v1.worker.workspace"] = module - - -def _install_gpu_ar_worker_default_attrs() -> None: - try: - module = importlib.import_module("vllm_omni.worker.gpu_ar_worker") - except Exception: - return - worker_cls = getattr(module, "GPUARWorker", None) - if worker_cls is None or hasattr(worker_cls, "use_v2_model_runner"): - return - # vLLM-Omni still relies on the v1 runner hooks for EasyMagpie. Newer vLLM - # worker bases no longer create this flag, so default the class lookup to - # the same v1 path the worker forces when the flag is present. - worker_cls.use_v2_model_runner = False - + 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, @@ -4915,886 +1198,6 @@ def install_easy_magpie_runtime_compat() -> None: _install_v1_serial_utils_dense_tensor_compat() _install_async_omni_client_compat() - -def _install_forward_context_kwargs_compat() -> None: - try: - import vllm.forward_context as forward_context - except Exception: - return - original = getattr(forward_context, "set_forward_context", None) - if original is None or getattr(original, "_easymagpie_compat", False): - return - original_signature = inspect.signature(original) - supported_kwargs = set(original_signature.parameters) - - def set_forward_context(*args, **kwargs): - context_kwargs = dict(kwargs) - filtered_kwargs = {key: value for key, value in kwargs.items() if key in supported_kwargs} - bound = original_signature.bind_partial(*args, **filtered_kwargs) - if "attn_metadata" in bound.arguments: - for key, value in bound.arguments.items(): - if key not in ("attn_metadata", "vllm_config"): - context_kwargs.setdefault(key, value) - attn_metadata = _normalize_mamba2_attention_metadata_groups(bound.arguments["attn_metadata"]) - attn_metadata = _synthetic_profile_attention_metadata_from_context( - attn_metadata, - bound.arguments.get("vllm_config"), - context_kwargs, - ) - bound.arguments["attn_metadata"] = _as_v1_mamba2_attention_metadata_dict( - attn_metadata, - bound.arguments.get("vllm_config"), - ) - return original(*bound.args, **bound.kwargs) - - set_forward_context._easymagpie_compat = True # type: ignore[attr-defined] - forward_context.set_forward_context = set_forward_context - for module_name in ( - "vllm_omni.worker.gpu_model_runner", - "vllm_omni.worker.gpu_ar_model_runner", - "vllm_omni.worker.gpu_generation_model_runner", - ): - module = sys.modules.get(module_name) - if module is not None and getattr(module, "set_forward_context", None) is original: - module.set_forward_context = set_forward_context - - -def _install_omni_gpu_model_runner_method_compat() -> None: - try: - from vllm.config import CUDAGraphMode - module = importlib.import_module("vllm_omni.worker.gpu_model_runner") - except Exception: - return - _install_cudagraph_mode_compat() - runner_cls = getattr(module, "OmniGPUModelRunner", None) - if runner_cls is None: - return - if not hasattr(runner_cls, "enable_prompt_embeds"): - runner_cls.enable_prompt_embeds = False - if not hasattr(runner_cls, "uses_xdrope_dim"): - runner_cls.uses_xdrope_dim = 0 - if not hasattr(runner_cls, "kv_cache_config"): - - def kv_cache_config(self): - return _as_v1_kv_cache_config(getattr(self, "_kv_cache_config", getattr(self, "cache_config", None))) - - def set_kv_cache_config(self, value): - self._kv_cache_config = value - - def del_kv_cache_config(self): - if hasattr(self, "_kv_cache_config"): - delattr(self, "_kv_cache_config") - - runner_cls.kv_cache_config = property(kv_cache_config, set_kv_cache_config, del_kv_cache_config) - - if not hasattr(runner_cls, "_determine_batch_execution_and_padding"): - - class _BatchDescriptor: - def __init__(self, *, num_tokens: int, num_reqs: int, uniform_decode: bool = False) -> None: - self.num_tokens = num_tokens - self.num_reqs = num_reqs - self.uniform_decode = uniform_decode - - def _determine_batch_execution_and_padding(self, **kwargs): - num_tokens = int(kwargs["num_tokens"]) - num_reqs = int(kwargs["num_reqs"]) - force_uniform_decode = bool(kwargs.get("force_uniform_decode", False)) - if hasattr(self, "get_dp_padding"): - num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) - num_tokens += int(num_pad) - else: - num_tokens_across_dp = None - batch_desc = _BatchDescriptor( - num_tokens=num_tokens, - num_reqs=num_reqs, - uniform_decode=force_uniform_decode, - ) - return CUDAGraphMode.NONE, batch_desc, False, num_tokens_across_dp, None - - runner_cls._determine_batch_execution_and_padding = _determine_batch_execution_and_padding - - if not hasattr(runner_cls, "synchronize_input_prep"): - - def synchronize_input_prep(self): - return contextlib.nullcontext() - - runner_cls.synchronize_input_prep = synchronize_input_prep - - if not hasattr(runner_cls, "_register_layerwise_nvtx_hooks"): - - def _register_layerwise_nvtx_hooks(self): - return None - - runner_cls._register_layerwise_nvtx_hooks = _register_layerwise_nvtx_hooks - - original_build_attention_metadata = getattr(runner_cls, "_build_attention_metadata", None) - if original_build_attention_metadata is not None and not getattr( - original_build_attention_metadata, "_easymagpie_compat", False - ): - original_build_attention_metadata_signature = inspect.signature(original_build_attention_metadata) - - def _build_attention_metadata(self, *args, **kwargs): - bound = original_build_attention_metadata_signature.bind_partial(self, *args, **kwargs) - result = original_build_attention_metadata(*bound.args, **bound.kwargs) - attn_metadata = result[0] if isinstance(result, tuple) and result else result - if isinstance(attn_metadata, dict) and not attn_metadata: - num_reqs = bound.arguments.get("num_reqs") - num_tokens = bound.arguments.get("num_tokens", 0) - max_query_len = bound.arguments.get("max_query_len", num_tokens) - if num_reqs is not None: - num_reqs = int(num_reqs) - self._easymagpie_mamba_cache_batch_size = int(num_reqs) - attn_metadata = _build_profile_flash_attention_metadata( - self, - num_tokens=int(num_tokens), - num_reqs=num_reqs, - max_query_len=int(max_query_len), - ) - if isinstance(result, tuple): - result = (attn_metadata, *result[1:]) - else: - result = attn_metadata - else: - normalized_attn_metadata = _normalize_mamba2_attention_metadata_groups(attn_metadata) - if normalized_attn_metadata is not attn_metadata: - attn_metadata = normalized_attn_metadata - if isinstance(result, tuple): - result = (attn_metadata, *result[1:]) - else: - result = attn_metadata - num_reqs = bound.arguments.get("num_reqs") - if num_reqs is not None: - self._easymagpie_mamba_cache_batch_size = int(num_reqs) - elif hasattr(self, "_easymagpie_mamba_cache_batch_size"): - delattr(self, "_easymagpie_mamba_cache_batch_size") - return result - - _build_attention_metadata._easymagpie_compat = True # type: ignore[attr-defined] - runner_cls._build_attention_metadata = _build_attention_metadata - - original_init_model_kwargs = getattr(runner_cls, "_init_model_kwargs", None) - if original_init_model_kwargs is not None and not getattr( - original_init_model_kwargs, "_easymagpie_compat", False - ): - try: - signature = inspect.signature(original_init_model_kwargs) - parameters = list(signature.parameters.values()) - accepts_num_tokens = any( - parameter.kind == inspect.Parameter.VAR_POSITIONAL for parameter in parameters - ) or (len(parameters) >= 2) - except Exception: - accepts_num_tokens = True - - def _call_base_init_model_kwargs(self, num_tokens): - for base_cls in getattr(runner_cls, "__mro__", ())[1:]: - base_init_model_kwargs = base_cls.__dict__.get("_init_model_kwargs") - if base_init_model_kwargs is None: - continue - if getattr(base_init_model_kwargs, "_easymagpie_compat", False): - continue - if base_init_model_kwargs is original_init_model_kwargs: - continue - try: - base_signature = inspect.signature(base_init_model_kwargs) - base_parameters = list(base_signature.parameters.values()) - base_accepts_num_tokens = any( - parameter.kind == inspect.Parameter.VAR_POSITIONAL for parameter in base_parameters - ) or (len(base_parameters) >= 2) - except Exception: - base_accepts_num_tokens = True - if base_accepts_num_tokens: - try: - return base_init_model_kwargs(self, num_tokens) - except TypeError: - pass - return base_init_model_kwargs(self) - return None - - def _init_model_kwargs(self, num_tokens=None): - if num_tokens is None: - num_tokens = int(getattr(self, "max_num_tokens", 0) or 0) - if accepts_num_tokens: - try: - model_kwargs = original_init_model_kwargs(self, num_tokens) - except TypeError: - model_kwargs = _call_base_init_model_kwargs(self, num_tokens) - if model_kwargs is None: - model_kwargs = original_init_model_kwargs(self) - else: - try: - model_kwargs = original_init_model_kwargs(self) - except TypeError: - model_kwargs = _call_base_init_model_kwargs(self, num_tokens) - if model_kwargs is None: - raise - cache_batch_size = getattr(self, "_easymagpie_mamba_cache_batch_size", None) - if cache_batch_size is not None: - model_kwargs = dict(model_kwargs) - model_kwargs.setdefault("easymagpie_mamba_cache_batch_size", int(cache_batch_size)) - request_cache_kwargs = _easymagpie_materialize_request_cache_kwargs( - self, - getattr(self, "_easymagpie_mamba_request_cache_state", None) - or getattr(self, "_easymagpie_mamba_request_cache_kwargs", None), - ) - if request_cache_kwargs: - model_kwargs = dict(model_kwargs) - model_kwargs.update(request_cache_kwargs) - return model_kwargs - - _init_model_kwargs._easymagpie_compat = True # type: ignore[attr-defined] - runner_cls._init_model_kwargs = _init_model_kwargs - - original_get_cumsum_and_arange = getattr(runner_cls, "_get_cumsum_and_arange", None) - if original_get_cumsum_and_arange is not None and not getattr( - original_get_cumsum_and_arange, "_easymagpie_compat", False - ): - try: - cumsum_signature = inspect.signature(original_get_cumsum_and_arange) - cumsum_parameters = list(cumsum_signature.parameters.values()) - requires_arange_out = any( - parameter.name == "arange_out" and parameter.default is inspect.Parameter.empty - for parameter in cumsum_parameters - ) - except Exception: - requires_arange_out = False - - def _new_arange_out(self, num_tokens): - import numpy as np - - total_num_tokens = int(np.asarray(num_tokens).sum()) - arange_np = getattr(self, "arange_np", None) - dtype = getattr(arange_np, "dtype", None) or getattr(num_tokens, "dtype", None) or np.int64 - return np.empty(total_num_tokens, dtype=dtype) - - def _get_cumsum_and_arange(self, num_tokens, arange_out=None, *args, **kwargs): - if not requires_arange_out: - if arange_out is None: - return original_get_cumsum_and_arange(self, num_tokens, *args, **kwargs) - return original_get_cumsum_and_arange(self, num_tokens, arange_out, *args, **kwargs) - - if arange_out is not None: - return original_get_cumsum_and_arange(self, num_tokens, arange_out, *args, **kwargs) - - generated_arange_out = _new_arange_out(self, num_tokens) - result = original_get_cumsum_and_arange( - self, - num_tokens, - generated_arange_out, - *args, - **kwargs, - ) - if isinstance(result, tuple): - return result - try: - total_num_tokens = int(result[-1]) - except Exception: - total_num_tokens = int(generated_arange_out.shape[0]) - return result, generated_arange_out[:total_num_tokens] - - _get_cumsum_and_arange._easymagpie_compat = True # type: ignore[attr-defined] - runner_cls._get_cumsum_and_arange = _get_cumsum_and_arange - - def _easymagpie_input_batch_req_ids(self) -> list[str]: - return [ - str(req_id) - for req_id in getattr(getattr(self, "input_batch", None), "req_ids", []) - if req_id is not None - ] - - def _easymagpie_request_ids_to_seq_ids(self, request_ids: list[str]) -> dict[str, list[int]]: - input_batch = getattr(self, "input_batch", None) - req_id_to_index = getattr(input_batch, "req_id_to_index", None) - used_seq_ids: set[int] = set() - next_fallback = 0 - mapped: dict[str, list[int]] = {} - - def fallback_seq_id() -> int: - nonlocal next_fallback - while next_fallback in used_seq_ids: - next_fallback += 1 - seq_idx = next_fallback - used_seq_ids.add(seq_idx) - next_fallback += 1 - return seq_idx - - if isinstance(req_id_to_index, dict): - missing_or_duplicate: list[str] = [] - seq_by_req_id: dict[str, int] = {} - for req_id in request_ids: - raw_idx = req_id_to_index.get(req_id) - if raw_idx is None: - raw_idx = req_id_to_index.get(str(req_id)) - try: - seq_idx = int(raw_idx) - except (TypeError, ValueError): - missing_or_duplicate.append(str(req_id)) - else: - if seq_idx in used_seq_ids: - missing_or_duplicate.append(str(req_id)) - else: - used_seq_ids.add(seq_idx) - seq_by_req_id[str(req_id)] = seq_idx - for req_id in missing_or_duplicate: - seq_by_req_id[req_id] = fallback_seq_id() - for req_id in request_ids: - mapped[str(req_id)] = [seq_by_req_id[str(req_id)]] - return mapped - for req_id in request_ids: - mapped[str(req_id)] = [fallback_seq_id()] - return mapped - - def _easymagpie_request_cache_state(self, scheduler_output: Any) -> dict[str, Any]: - num_scheduled_tokens = getattr(scheduler_output, "num_scheduled_tokens", None) - if not num_scheduled_tokens: - return {} - - scheduled_counts = { - str(req_id): int(count or 0) - for req_id, count in num_scheduled_tokens.items() - if int(count or 0) > 0 - } - if not scheduled_counts: - return {} - - finished_requests_ids = getattr(scheduler_output, "finished_req_ids", None) - if finished_requests_ids is None: - finished_requests_ids = getattr(scheduler_output, "finished_requests_ids", None) - return { - "scheduled_request_ids": list(scheduled_counts), - "scheduled_counts": scheduled_counts, - "finished_requests_ids": ( - None - if finished_requests_ids is None - else sorted(str(req_id) for req_id in finished_requests_ids) - ), - "previous_active_request_ids": sorted( - str(req_id) - for req_id in getattr(self, "_easymagpie_active_mamba_request_ids", set()) - ), - "pre_update_request_ids": _easymagpie_input_batch_req_ids(self), - } - - def _easymagpie_materialize_request_cache_kwargs(self, state: Any) -> dict[str, Any]: - if not state: - return {} - if "request_ids_to_seq_ids" in state: - return state - - scheduled_counts = { - str(req_id): int(count or 0) - for req_id, count in dict(state.get("scheduled_counts", {})).items() - if int(count or 0) > 0 - } - if not scheduled_counts: - return {} - - ordered_request_ids: list[str] = [] - seen: set[str] = set() - for req_id in _easymagpie_input_batch_req_ids(self): - if req_id in scheduled_counts and req_id not in seen: - ordered_request_ids.append(req_id) - seen.add(req_id) - for req_id in state.get("scheduled_request_ids", scheduled_counts): - req_id = str(req_id) - if req_id in scheduled_counts and req_id not in seen: - ordered_request_ids.append(req_id) - seen.add(req_id) - if not ordered_request_ids: - return {} - - current_active = set(str(req_id) for req_id in state.get("pre_update_request_ids", [])) - current_active.update(_easymagpie_input_batch_req_ids(self)) - current_active.update(ordered_request_ids) - previous_active = set(str(req_id) for req_id in state.get("previous_active_request_ids", [])) - implicit_finished = previous_active - current_active - finished_requests_ids = state.get("finished_requests_ids") - if finished_requests_ids is None: - finished_requests_ids = sorted(implicit_finished) - else: - finished_requests_ids = sorted( - set(str(req_id) for req_id in finished_requests_ids) | implicit_finished - ) - self._easymagpie_active_mamba_request_ids = current_active - return { - "request_ids_to_seq_ids": _easymagpie_request_ids_to_seq_ids(self, ordered_request_ids), - "finished_requests_ids": finished_requests_ids, - } - - original_model_forward = getattr(runner_cls, "_model_forward", None) - if original_model_forward is not None and not getattr(original_model_forward, "_easymagpie_compat", False): - - def _model_forward(self, *args, **kwargs): - request_cache_kwargs = _easymagpie_materialize_request_cache_kwargs( - self, - getattr(self, "_easymagpie_mamba_request_cache_state", None) - or getattr(self, "_easymagpie_mamba_request_cache_kwargs", None), - ) - if request_cache_kwargs: - kwargs = dict(kwargs) - for key, value in request_cache_kwargs.items(): - kwargs.setdefault(key, value) - return original_model_forward(self, *args, **kwargs) - - _model_forward._easymagpie_compat = True # type: ignore[attr-defined] - runner_cls._model_forward = _model_forward - - def _install_execute_model_request_cache_compat(target_cls: type[Any]) -> None: - original_execute_model = getattr(target_cls, "execute_model", None) - if original_execute_model is None or getattr(original_execute_model, "_easymagpie_compat", False): - return - original_execute_model_signature = inspect.signature(original_execute_model) - - def execute_model(self, *args, **kwargs): - bound = original_execute_model_signature.bind_partial(self, *args, **kwargs) - scheduler_output = bound.arguments.get("scheduler_output") - request_cache_state = ( - _easymagpie_request_cache_state(self, scheduler_output) - if scheduler_output is not None - else {} - ) - if request_cache_state: - self._easymagpie_mamba_request_cache_state = request_cache_state - elif hasattr(self, "_easymagpie_mamba_request_cache_state"): - delattr(self, "_easymagpie_mamba_request_cache_state") - try: - return original_execute_model(*bound.args, **bound.kwargs) - finally: - for name in ( - "_easymagpie_mamba_request_cache_state", - "_easymagpie_mamba_request_cache_kwargs", - ): - if hasattr(self, name): - delattr(self, name) - - execute_model._easymagpie_compat = True # type: ignore[attr-defined] - target_cls.execute_model = execute_model - - _install_execute_model_request_cache_compat(runner_cls) - for module_name, class_names in ( - ("vllm_omni.worker.gpu_ar_model_runner", ("GPUARModelRunner",)), - ("vllm_omni.worker.gpu_generation_model_runner", ("GPUGenerationModelRunner",)), - ): - extra_module = sys.modules.get(module_name) - if extra_module is None: - continue - for class_name in class_names: - extra_runner_cls = getattr(extra_module, class_name, None) - if extra_runner_cls is not None: - _install_execute_model_request_cache_compat(extra_runner_cls) - - original_lora_context = getattr(runner_cls, "maybe_dummy_run_with_lora", None) - if original_lora_context is not None and not getattr(original_lora_context, "_easymagpie_compat", False): - - def maybe_dummy_run_with_lora(self, lora_config, num_scheduled_tokens, *_args, **_kwargs): - if lora_config is None: - return contextlib.nullcontext() - return original_lora_context(self, lora_config, num_scheduled_tokens) - - maybe_dummy_run_with_lora._easymagpie_compat = True # type: ignore[attr-defined] - runner_cls.maybe_dummy_run_with_lora = maybe_dummy_run_with_lora - - original_randomize_context = getattr(runner_cls, "maybe_randomize_inputs", None) - if original_randomize_context is not None and not getattr( - original_randomize_context, "_easymagpie_compat", False - ): - - def maybe_randomize_inputs(self, input_ids, inputs_embeds=None): - if input_ids is None: - return contextlib.nullcontext() - try: - return original_randomize_context(self, input_ids, inputs_embeds) - except TypeError: - return original_randomize_context(self, input_ids) - - maybe_randomize_inputs._easymagpie_compat = True # type: ignore[attr-defined] - runner_cls.maybe_randomize_inputs = maybe_randomize_inputs - - original_dummy_run = getattr(runner_cls, "_dummy_run", None) - if original_dummy_run is not None and not getattr(original_dummy_run, "_easymagpie_force_attn_compat", False): - original_dummy_run_signature = inspect.signature(original_dummy_run) - original_dummy_run_parameters = original_dummy_run_signature.parameters - original_dummy_run_has_cudagraph_mode = "cudagraph_runtime_mode" in original_dummy_run_parameters - original_dummy_run_var_keyword = next( - ( - name - for name, parameter in original_dummy_run_parameters.items() - if parameter.kind == inspect.Parameter.VAR_KEYWORD - ), - None, - ) - - def get_bound_keyword(bound, name: str, default: Any = None): - if name in bound.arguments: - return bound.arguments[name] - if original_dummy_run_var_keyword is not None: - kwargs = bound.arguments.get(original_dummy_run_var_keyword, {}) - if isinstance(kwargs, dict) and name in kwargs: - return kwargs[name] - return default - - def set_bound_keyword(bound, name: str, value: Any) -> None: - parameter = original_dummy_run_parameters.get(name) - if parameter is not None and parameter.kind != inspect.Parameter.VAR_KEYWORD: - bound.arguments[name] = value - return - if original_dummy_run_var_keyword is not None: - kwargs = dict(bound.arguments.get(original_dummy_run_var_keyword, {})) - kwargs[name] = value - bound.arguments[original_dummy_run_var_keyword] = kwargs - return - bound.arguments[name] = value - - def needs_mamba_attention_metadata(self) -> bool: - candidates = [getattr(self, "model", None)] - seen: set[int] = set() - while candidates: - candidate = candidates.pop() - if candidate is None: - continue - candidate_id = id(candidate) - if candidate_id in seen: - continue - seen.add(candidate_id) - class_name = candidate.__class__.__name__.lower() - config = getattr(candidate, "config", None) - pattern = getattr(config, "hybrid_override_pattern", None) - architectures = getattr(config, "architectures", None) or () - if "easymagpie" in class_name or "easy_magpie" in class_name: - return True - if "nemotronh" in class_name or "nemotron_h" in class_name: - return True - if isinstance(pattern, str) and "M" in pattern and hasattr(config, "chunk_size"): - return True - if any( - "easymagpie" in str(architecture).lower() - or "nemotronh" in str(architecture).lower() - or "nemotron_h" in str(architecture).lower() - for architecture in architectures - ): - return True - for attr_name in ("backbone", "model", "language_model", "llm", "decoder"): - child = getattr(candidate, attr_name, None) - if child is not None: - candidates.append(child) - return False - - def ensure_easy_mamba_no_compile_layers(self) -> None: - vllm_config = getattr(self, "vllm_config", None) - compilation_config = getattr(vllm_config, "compilation_config", None) - static_context = getattr(compilation_config, "static_forward_context", None) - if not isinstance(static_context, dict): - return - backbone = getattr(getattr(self, "model", None), "backbone", None) - layers = getattr(backbone, "layers", None) - if layers is None: - return - for layer_idx, layer in enumerate(layers): - mixer = getattr(layer, "mixer", None) - if mixer is not None: - static_context.setdefault(f"backbone.layers.{layer_idx}.mixer", mixer) - - def _dummy_run(self, *args, **kwargs): - bound = original_dummy_run_signature.bind_partial(self, *args, **kwargs) - force_v1_env = False - if needs_mamba_attention_metadata(self): - ensure_easy_mamba_no_compile_layers(self) - num_tokens = get_bound_keyword(bound, "num_tokens") - max_num_reqs = getattr(getattr(self, "scheduler_config", None), "max_num_seqs", None) - if num_tokens is not None and max_num_reqs is not None: - self._easymagpie_mamba_cache_batch_size = min(max(1, int(num_tokens)), max(1, int(max_num_reqs))) - use_v1_dummy_run = original_dummy_run_has_cudagraph_mode - if not use_v1_dummy_run and original_dummy_run_var_keyword is not None: - try: - from vllm import envs - - use_v1_dummy_run = bool(getattr(envs, "VLLM_USE_V1", False)) - except Exception: - use_v1_dummy_run = False - if use_v1_dummy_run: - set_bound_keyword(bound, "cudagraph_runtime_mode", CUDAGraphMode.NONE) - set_bound_keyword(bound, "force_attention", False) - force_v1_env = True - elif not get_bound_keyword(bound, "force_attention", False): - set_bound_keyword(bound, "force_attention", True) - if not force_v1_env: - return original_dummy_run(*bound.args, **bound.kwargs) - try: - from vllm import envs - except Exception: - return original_dummy_run(*bound.args, **bound.kwargs) - previous_use_v1 = getattr(envs, "VLLM_USE_V1", None) - envs.VLLM_USE_V1 = True - try: - return original_dummy_run(*bound.args, **bound.kwargs) - finally: - envs.VLLM_USE_V1 = previous_use_v1 - - _dummy_run._easymagpie_force_attn_compat = True # type: ignore[attr-defined] - runner_cls._dummy_run = _dummy_run - - -def _install_omni_worker_import_patch() -> None: - if getattr(builtins, _IMPORT_PATCH_FLAG, False): - return - original_import = builtins.__import__ - - def import_with_omni_worker_patch(name, globals=None, locals=None, fromlist=(), level=0): - should_patch_preprocess = ( - name == "vllm_omni.inputs.preprocess" - or name.startswith("vllm_omni.inputs.preprocess.") - or (name == "vllm_omni.inputs" and "preprocess" in fromlist) - ) - if should_patch_preprocess: - _install_vllm_inputs_data_alias() - _install_vllm_multimodal_inputs_alias() - module = original_import(name, globals, locals, fromlist, level) - should_patch = ( - name == "vllm_omni.worker.gpu_model_runner" - or name == "vllm_omni.worker.gpu_ar_model_runner" - or name == "vllm_omni.worker.gpu_generation_model_runner" - or name.startswith("vllm_omni.worker.gpu_model_runner.") - or name.startswith("vllm_omni.worker.gpu_ar_model_runner.") - or name.startswith("vllm_omni.worker.gpu_generation_model_runner.") - or (name == "vllm_omni.worker" and "gpu_model_runner" in fromlist) - or ( - name == "vllm_omni.worker" - and any(item in fromlist for item in ("gpu_ar_model_runner", "gpu_generation_model_runner")) - ) - ) - if should_patch: - _install_forward_context_kwargs_compat() - _install_omni_gpu_model_runner_method_compat() - should_patch_worker = ( - name == "vllm_omni.worker.gpu_ar_worker" - or name == "vllm_omni.worker.gpu_generation_worker" - or name.startswith("vllm_omni.worker.gpu_ar_worker.") - or name.startswith("vllm_omni.worker.gpu_generation_worker.") - or ( - name == "vllm_omni.worker" - and any(item in fromlist for item in ("gpu_ar_worker", "gpu_generation_worker")) - ) - ) - if should_patch_worker: - _install_easy_magpie_refit_rpc_compat() - if should_patch_preprocess: - _install_omni_input_preprocessor_signature_compat() - should_patch_async_engine = ( - name == "vllm_omni.engine.async_omni_engine" - or name.startswith("vllm_omni.engine.async_omni_engine.") - or (name == "vllm_omni.engine" and "async_omni_engine" in fromlist) - ) - if should_patch_async_engine: - _install_omni_engine_request_compat() - should_patch_request = name == "vllm_omni.request" or name.startswith("vllm_omni.request.") - if should_patch_request: - _install_omni_request_compat() - return module - - builtins.__import__ = import_with_omni_worker_patch - setattr(builtins, _IMPORT_PATCH_FLAG, True) - - -def _install_ec_transfer_alias() -> None: - if "vllm.distributed.ec_transfer" in sys.modules: - return - module = types.ModuleType("vllm.distributed.ec_transfer") - - class _DisabledECTransfer: - is_consumer = False - is_producer = False - - def has_ec_transfer() -> bool: - return False - - def get_ec_transfer(): - return _DisabledECTransfer() - - module.has_ec_transfer = has_ec_transfer - module.get_ec_transfer = get_ec_transfer - sys.modules["vllm.distributed.ec_transfer"] = module - - -def _install_routed_experts_capturer_alias() -> None: - module_name = "vllm.model_executor.layers.fused_moe.routed_experts_capturer" - if module_name in sys.modules: - return - module = types.ModuleType(module_name) - - class RoutedExpertsCapturer: - @classmethod - def get_instance(cls): - return None - - def clear_buffer(self) -> None: - return None - - module.RoutedExpertsCapturer = RoutedExpertsCapturer - sys.modules[module_name] = module - - -def _install_structured_output_aliases() -> None: - try: - import vllm.v1.core.sched.output as sched_output - except Exception: - sched_output = None - if sched_output is not None and not hasattr(sched_output, "GrammarOutput"): - sched_output.GrammarOutput = object - - try: - import vllm.v1.structured_output.utils as structured_utils - except Exception: - return - if hasattr(structured_utils, "apply_grammar_bitmask"): - return - - def apply_grammar_bitmask(*_args, **_kwargs): - return None - - structured_utils.apply_grammar_bitmask = apply_grammar_bitmask - - -def _install_outputs_aliases() -> None: - try: - import vllm.v1.outputs as outputs - except Exception: - return - if not hasattr(outputs, "AsyncModelRunnerOutput") and hasattr(outputs, "ModelRunnerOutput"): - outputs.AsyncModelRunnerOutput = outputs.ModelRunnerOutput - if not hasattr(outputs, "make_empty_encoder_model_runner_output"): - - def make_empty_encoder_model_runner_output(*_args, **_kwargs): - return outputs.EMPTY_MODEL_RUNNER_OUTPUT - - outputs.make_empty_encoder_model_runner_output = make_empty_encoder_model_runner_output - - -def _install_spec_decode_aliases() -> None: - class _UnsupportedSpecDecodeProposer: - def __init__(self, *args, **kwargs) -> None: - raise NotImplementedError("This vLLM build does not expose the legacy spec-decode proposer.") - - aliases = { - "vllm.v1.spec_decode.draft_model": ("DraftModelProposer", _UnsupportedSpecDecodeProposer), - "vllm.v1.spec_decode.extract_hidden_states": ( - "ExtractHiddenStatesProposer", - _UnsupportedSpecDecodeProposer, - ), - } - for module_name, (class_name, cls) in aliases.items(): - if module_name in sys.modules: - continue - module = types.ModuleType(module_name) - setattr(module, class_name, cls) - sys.modules[module_name] = module - - -def _install_v1_utils_aliases() -> None: - try: - import vllm.v1.utils as v1_utils - except Exception: - return - if hasattr(v1_utils, "record_function_or_nullcontext"): - return - - def record_function_or_nullcontext(name: str): - try: - import torch - - return torch.profiler.record_function(name) - except Exception: - return contextlib.nullcontext() - - v1_utils.record_function_or_nullcontext = record_function_or_nullcontext - - -def _install_gpu_model_runner_aliases() -> None: - try: - import vllm.v1.worker.gpu_model_runner as gpu_model_runner - except Exception: - return - if not hasattr(gpu_model_runner, "PerLayerAttnMetadata"): - gpu_model_runner.PerLayerAttnMetadata = dict - if hasattr(gpu_model_runner, "AsyncGPUModelRunnerOutput"): - return - - class AsyncGPUModelRunnerOutput: - def __init__( - self, - model_runner_output=None, - sampled_token_ids=None, - logprobs_tensors=None, - invalid_req_indices=None, - async_output_copy_stream=None, - vocab_size=None, - ) -> None: - self.model_runner_output = model_runner_output - self.sampled_token_ids = sampled_token_ids - self.logprobs_tensors = logprobs_tensors - self.invalid_req_indices = invalid_req_indices - self.async_output_copy_stream = async_output_copy_stream - self.vocab_size = vocab_size - self.sampled_token_ids_cpu = sampled_token_ids - self.async_copy_ready_event = None - - gpu_model_runner.AsyncGPUModelRunnerOutput = AsyncGPUModelRunnerOutput - - -def _install_model_interface_aliases() -> None: - try: - import vllm.model_executor.models.interfaces as interfaces - except Exception: - return - if hasattr(interfaces, "supports_mrope"): - return - - def supports_mrope(*_args, **_kwargs) -> bool: - return False - - interfaces.supports_mrope = supports_mrope - - -def _install_ubatch_utils_alias() -> None: - module_name = "vllm.v1.worker.ubatch_utils" - if module_name in sys.modules: - return - module = types.ModuleType(module_name) - - def maybe_create_ubatch_slices(*_args, **_kwargs): - return None, None - - module.maybe_create_ubatch_slices = maybe_create_ubatch_slices - sys.modules[module_name] = module - - -def _install_tracing_instrument_alias() -> None: - try: - import vllm.tracing as tracing - except Exception: - return - if hasattr(tracing, "instrument"): - return - - def instrument(func=None, **_kwargs): - def decorator(inner): - return inner - - if callable(func): - return decorator(func) - return decorator - - tracing.instrument = instrument - - -def _install_executor_alias() -> None: - try: - import vllm.v1.executor as executor_pkg - from vllm.v1.executor.abstract import Executor - except Exception: - return - if not hasattr(executor_pkg, "Executor"): - executor_pkg.Executor = Executor - - def _install_async_omni_client_compat() -> None: try: async_omni_module = importlib.import_module("vllm_omni.entrypoints.async_omni") @@ -6087,134 +1490,15 @@ def qsize(self) -> int: pass -def _install_omni_model_config_field_compat() -> None: - try: - from vllm_omni.config.model import OmniModelConfig - except Exception: - return - if getattr(OmniModelConfig, "_easymagpie_field_compat", False): - return - - omni_defaults = { - "stage_id": 0, - "async_chunk": False, - "model_stage": "thinker", - "model_arch": None, - "worker_type": None, - "engine_output_type": None, - "hf_config_name": None, - "custom_process_next_stage_input_func": None, - "stage_connector_config": lambda: {"name": "SharedMemoryConnector", "extra": {}}, - "omni_kv_config": None, - "codec_frame_rate_hz": None, - "task_type": None, - "io_processor_plugin": None, - "has_sampling_extra_args": False, - "subtalker_sampling_params": None, - } - - def add_defaults_to_omni_kwargs(cls, omni_kwargs): - for key, default in omni_defaults.items(): - if key not in omni_kwargs: - omni_kwargs[key] = default() if callable(default) else default - - def _validate_omni_fields(cls, **omni_kwargs): - unexpected = set(omni_kwargs) - set(omni_defaults) - if unexpected: - raise ValueError(f"Unexpected omni kwarg: {sorted(unexpected)}") - - OmniModelConfig.add_defaults_to_omni_kwargs = classmethod(add_defaults_to_omni_kwargs) - OmniModelConfig._validate_omni_fields = classmethod(_validate_omni_fields) - if not hasattr(OmniModelConfig, "io_processor_plugin"): - OmniModelConfig.io_processor_plugin = None - if not hasattr(OmniModelConfig, "has_sampling_extra_args"): - OmniModelConfig.has_sampling_extra_args = False - if not hasattr(OmniModelConfig, "subtalker_sampling_params"): - OmniModelConfig.subtalker_sampling_params = None - OmniModelConfig._easymagpie_field_compat = True - def install_vllm_omni_compat() -> None: - """Install import/signature shims required by vLLM-Omni on older vLLM.""" - - global _INSTALLED - if _INSTALLED: - _install_ray_objectref_aliases() - _install_ray_placement_group_aliases() - _install_vllm_inputs_data_alias() - _install_vllm_multimodal_inputs_alias() - return - _install_config_decorator_compat() - _install_mamba2_metadata_compat() - _install_triton_attention_profile_metadata_compat() - _install_flash_attention_profile_metadata_compat() - _install_flashinfer_attention_profile_metadata_compat() - _install_mamba2_profile_no_kv_cache_compat() - _install_lora_config_alias() - _install_model_arch_convertor() - _install_io_processor_stub() - _install_tokenizer_aliases() - _install_repo_utils_alias() - _install_config_parser_aliases() - _install_nemotron_h_auto_config() - _install_vllm_inputs_data_alias() - _install_vllm_multimodal_inputs_alias() - _install_renderer_aliases() - _install_input_processor_alias() - _install_input_preprocessor_truncate_compat() - _install_omni_input_preprocessor_signature_compat() - _install_omni_engine_request_compat() - _install_omni_request_compat() - _install_processor_process_inputs_compat() - _install_engine_utils_compat() - _install_import_utils_alias() - _install_torch_utils_alias() - _install_math_utils_alias() - _install_mem_utils_alias() - _install_vllm_config_profiler_default() - _install_parallel_config_defaults() - _install_cache_config_defaults() - _install_platform_dtype_check() - _install_torch_accelerator_compat() - _install_logger_info_once_scope_compat() - _install_v1_serial_utils_dense_tensor_compat() - _install_cudagraph_mode_compat() - _install_cuda_graph_stat_alias() - _install_cuda_piecewise_no_sym_shape_compat() - _install_kv_connector_stats_alias() - _install_perf_stats_alias() - _install_scheduler_make_stats_compat() - _install_sched_utils_aliases() - _install_sched_interface_aliases() - _install_output_processor_signature_compat() - _install_flash_attention_builder_compat() - _install_worker_utils_compat() - _install_worker_workspace_alias() - _install_gpu_ar_worker_default_attrs() - _install_easy_magpie_refit_rpc_compat() - _install_forward_context_kwargs_compat() - _install_omni_gpu_model_runner_method_compat() - _install_omni_worker_import_patch() - _install_ec_transfer_alias() - _install_routed_experts_capturer_alias() - _install_structured_output_aliases() - _install_outputs_aliases() - _install_spec_decode_aliases() - _install_v1_utils_aliases() - _install_model_interface_aliases() - _install_gpu_model_runner_aliases() - _install_ubatch_utils_alias() - # vLLM-Omni's runner import may fail until the late aliases above exist. - _install_forward_context_kwargs_compat() - _install_omni_gpu_model_runner_method_compat() - _install_tracing_instrument_alias() - _install_executor_alias() - _install_ray_objectref_aliases() - _install_ray_placement_group_aliases() - _install_omni_model_config_field_compat() - _install_async_omni_client_compat() - _register_easy_magpie_plugin() - _INSTALLED = True + """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__ = [ diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 1fedb6f6fa54..7fa63e81d43e 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -46,3 +46,44 @@ def test_compat_mode_off_does_not_import_vllm() -> None: ) 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 From 1bca834e421440578136dd710dd44328b4e6f627 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 16:40:18 +0000 Subject: [PATCH 03/14] Expose EasyMagpie runtime reset as named RPC --- .../easymagpie_vllm_omni/vllm_compat.py | 42 ++++++++++++++++++ .../tests/test_import_modes.py | 43 +++++++++++++++++++ 2 files changed, 85 insertions(+) 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 index 028961762975..ceae0ec47ff2 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -713,6 +713,7 @@ def _reset_easy_magpie_runner_state_after_refit(model_runner: Any) -> dict[str, _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: @@ -1153,19 +1154,60 @@ def easymagpie_update_text_embedding_rows(self, payload): "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: diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 7fa63e81d43e..2242a9180355 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -19,6 +19,7 @@ import os import subprocess import sys +import types from pathlib import Path @@ -87,3 +88,45 @@ def test_compat_mode_full_is_not_part_of_review_branch() -> None: ) 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}, + ] From 1b72783e7bc86baed3c5d3e668c04a4258ed24ad Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 20:26:00 +0000 Subject: [PATCH 04/14] Trim EasyMagpie optional debug outputs --- .../easymagpie_vllm_omni/easymagpie.py | 474 +----------------- .../easymagpie_vllm_omni/local_transformer.py | 15 - .../tests/test_easymagpie_backbone_cache.py | 7 +- .../tests/test_local_transformer_sampler.py | 1 - 4 files changed, 7 insertions(+), 490 deletions(-) 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 db33d6e628ff..84f47e44c5b9 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/easymagpie.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/easymagpie.py @@ -335,30 +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._debug_combined_input_norm = torch.zeros(max_num_tokens, dtype=torch.float32) - self._debug_combined_input_vector = torch.zeros(max_num_tokens, self.embedding_dim, dtype=dtype) - self._debug_text_emb_norm = torch.zeros(max_num_tokens, dtype=torch.float32) - self._debug_phoneme_emb_norm = torch.zeros(max_num_tokens, dtype=torch.float32) - self._debug_audio_emb_norm = torch.zeros(max_num_tokens, dtype=torch.float32) - self._debug_positions = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_input_ids = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_decode_dispatch_index = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_decode_dispatch_count = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_decode_offset = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_mamba_num_prefills = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_mamba_num_decodes = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_mamba_state_indices = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_mamba_cache_state_indices = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_mamba_exec_state_indices = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_attn_num_actual_tokens = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_attn_max_query_len = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_attn_max_seq_len = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_attn_seq_lens = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_attn_slot_mapping = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_audio_feedback_missing = torch.zeros(max_num_tokens, dtype=torch.long) - self._debug_audio_input_code0 = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_phoneme_input_valid = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._debug_phoneme_input_token0 = torch.full((max_num_tokens,), -1, dtype=torch.long) + 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) @@ -374,9 +351,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: 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._debug_lt_top_ids = torch.full((max_num_tokens, self.num_codebooks, 5), -1, dtype=torch.long) - self._debug_lt_top_values = torch.zeros(max_num_tokens, self.num_codebooks, 5, dtype=torch.float32) - self._debug_outputs_enabled = False self._last_output_row_indices: Optional[torch.Tensor] = None # ── Audio-EOS → engine stop ───────────────────────────────────── @@ -657,211 +631,6 @@ def _iter_metadata(value: Any) -> Iterable[Any]: elif value is not None: yield value - @staticmethod - def _metadata_int(metadata: Any, *names: str) -> Optional[int]: - for name in names: - value = getattr(metadata, name, None) - if value is None: - continue - if isinstance(value, torch.Tensor): - if value.numel() == 0: - continue - return int(value.detach().reshape(-1)[0].item()) - try: - return int(value) - except Exception: - continue - return None - - @staticmethod - def _metadata_tensor(metadata: Any, *names: str) -> Optional[torch.Tensor]: - for name in names: - value = getattr(metadata, name, None) - if isinstance(value, torch.Tensor) and value.numel() > 0: - return value.detach().reshape(-1).to(dtype=torch.long) - return None - - @staticmethod - def _metadata_list_tensor(metadata: Any, *names: str, device: torch.device) -> Optional[torch.Tensor]: - tensor = EasyMagpieTTSForConditionalGeneration._metadata_tensor(metadata, *names) - if tensor is not None: - return tensor.to(device=device) - for name in names: - value = getattr(metadata, name, None) - if value is None: - continue - try: - items = list(value) - except TypeError: - continue - if items: - return torch.tensor([int(item) for item in items], dtype=torch.long, device=device) - return None - - def _record_state_indices_debug( - self, - target: torch.Tensor, - state_indices: torch.Tensor, - *, - num_tokens: int, - decode_idx: Optional[torch.Tensor] = None, - num_req: int = 0, - ) -> None: - device = target.device - state_indices = state_indices.detach().reshape(-1).to(device=device, dtype=torch.long) - if decode_idx is not None and num_req > 0 and state_indices.numel() >= num_req: - valid = decode_idx[:num_req].detach().to(device=device, dtype=torch.long).reshape(-1) - target[valid].copy_(state_indices[:num_req]) - else: - n = min(num_tokens, int(state_indices.numel())) - target[:n].copy_(state_indices[:n]) - - def _record_mamba_metadata_debug( - self, - num_tokens: int, - *, - decode_idx: Optional[torch.Tensor] = None, - num_req: int = 0, - mamba_cache_params: Any = None, - ) -> None: - try: - attn_metadata = get_forward_context().attn_metadata - except Exception: - attn_metadata = None - - metas = list(self._iter_metadata(attn_metadata)) - mamba_metas = [ - meta for meta in metas if getattr(meta, "state_indices_tensor", None) is not None - ] - counter_metas = mamba_metas or metas - rows = slice(0, num_tokens) - - for metadata in counter_metas: - num_prefills = self._metadata_int(metadata, "num_prefills", "num_prefill_tokens") - if num_prefills is not None: - self._debug_mamba_num_prefills[rows].fill_(num_prefills) - break - for metadata in counter_metas: - num_decodes = self._metadata_int(metadata, "num_decodes", "num_decode_tokens") - if num_decodes is not None: - self._debug_mamba_num_decodes[rows].fill_(num_decodes) - break - - cache_state_indices = None - for name in ("state_indices_tensor_d", "state_indices_tensor", "state_indices_tensor_p"): - value = getattr(mamba_cache_params, name, None) - if isinstance(value, torch.Tensor) and value.numel() > 0: - cache_state_indices = value - break - if cache_state_indices is not None: - self._record_state_indices_debug( - self._debug_mamba_cache_state_indices, - cache_state_indices, - num_tokens=num_tokens, - decode_idx=decode_idx, - num_req=num_req, - ) - - exec_state_indices = None - for metadata in mamba_metas: - exec_state_indices = self._metadata_tensor(metadata, "state_indices_tensor") - if exec_state_indices is not None: - break - if exec_state_indices is None: - for metadata in metas: - block_table = getattr(metadata, "block_table_tensor", None) - if isinstance(block_table, torch.Tensor) and block_table.numel() > 0: - exec_state_indices = block_table[:, 0].detach().reshape(-1).to(dtype=torch.long) - break - - if exec_state_indices is not None: - self._record_state_indices_debug( - self._debug_mamba_exec_state_indices, - exec_state_indices, - num_tokens=num_tokens, - decode_idx=decode_idx, - num_req=num_req, - ) - self._record_state_indices_debug( - self._debug_mamba_state_indices, - exec_state_indices, - num_tokens=num_tokens, - decode_idx=decode_idx, - num_req=num_req, - ) - elif cache_state_indices is not None: - self._record_state_indices_debug( - self._debug_mamba_state_indices, - cache_state_indices, - num_tokens=num_tokens, - decode_idx=decode_idx, - num_req=num_req, - ) - - def _record_debug_forward_metadata( - self, - *, - input_ids: torch.Tensor, - positions: torch.Tensor, - decode_idx: Optional[torch.Tensor], - num_req: int, - mamba_cache_params: Any, - ) -> None: - if not self._debug_outputs_enabled: - return - - num_tokens = int(input_ids.shape[0]) - rows = slice(0, num_tokens) - device = self._debug_positions.device - self._debug_positions[rows].copy_(positions.detach().to(device=device, dtype=torch.long).reshape(-1)[:num_tokens]) - self._debug_input_ids[rows].copy_(input_ids.detach().to(device=device, dtype=torch.long).reshape(-1)[:num_tokens]) - self._debug_decode_dispatch_count[rows].fill_(int(num_req)) - if decode_idx is None: - self._debug_decode_dispatch_index[rows].copy_(torch.arange(num_tokens, dtype=torch.long, device=device)) - elif num_req > 0: - valid = decode_idx[:num_req].detach().to(device=device, dtype=torch.long).reshape(-1) - self._debug_decode_dispatch_index[valid].copy_(valid) - - try: - attn_metadata = get_forward_context().attn_metadata - except Exception: - attn_metadata = None - metas = list(self._iter_metadata(attn_metadata)) - self._record_mamba_metadata_debug( - num_tokens, - decode_idx=decode_idx, - num_req=num_req, - mamba_cache_params=mamba_cache_params, - ) - for metadata in metas: - num_actual = self._metadata_int(metadata, "num_actual_tokens") - if num_actual is not None: - self._debug_attn_num_actual_tokens[rows].fill_(num_actual) - break - for metadata in metas: - max_query_len = self._metadata_int(metadata, "max_query_len") - if max_query_len is not None: - self._debug_attn_max_query_len[rows].fill_(max_query_len) - break - for metadata in metas: - max_seq_len = self._metadata_int(metadata, "max_seq_len", "max_seq_len_q") - if max_seq_len is not None: - self._debug_attn_max_seq_len[rows].fill_(max_seq_len) - break - for metadata in metas: - seq_lens = self._metadata_list_tensor(metadata, "seq_lens", "seq_lens_tensor", device=device) - if seq_lens is not None: - n = min(num_tokens, int(seq_lens.numel())) - self._debug_attn_seq_lens[:n].copy_(seq_lens[:n]) - break - for metadata in metas: - slot_mapping = self._metadata_tensor(metadata, "slot_mapping") - if slot_mapping is not None: - slot_mapping = slot_mapping.to(device=device) - n = min(num_tokens, int(slot_mapping.numel())) - self._debug_attn_slot_mapping[:n].copy_(slot_mapping[:n]) - break - # ------------------------------------------------------------------ # forward # ------------------------------------------------------------------ @@ -911,12 +680,6 @@ def forward( valid = decode_idx[:num_req] self._assemble_decode_embeddings(combined, valid) - if self._debug_outputs_enabled: - debug_rows = slice(0, num_tokens) - self._debug_combined_input_vector[debug_rows].copy_(combined.detach()) - self._debug_combined_input_norm[debug_rows] = combined.detach().float().norm(dim=-1) - self.code_predictor.debug_collect_logits = bool(self._debug_outputs_enabled) - backbone_kwargs = { "input_ids": input_ids, "positions": positions, @@ -928,14 +691,6 @@ def forward( mamba_cache_params = self._current_mamba_cache_params(**kwargs) backbone_kwargs["mamba_cache_params"] = mamba_cache_params - self._record_debug_forward_metadata( - input_ids=input_ids, - positions=positions, - decode_idx=decode_idx, - num_req=int(num_req), - mamba_cache_params=mamba_cache_params, - ) - hidden_states = self.backbone(**backbone_kwargs) # Sample codes (local transformer) only where needed. @@ -945,9 +700,6 @@ def forward( 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)) - if self._debug_outputs_enabled: - self._debug_lt_top_ids[:num_tokens].copy_(self.code_predictor._debug_top_ids[:num_tokens]) - self._debug_lt_top_values[:num_tokens].copy_(self.code_predictor._debug_top_values[:num_tokens]) self._flag_audio_eos(codes, slice(0, num_tokens)) if self.has_phoneme: self._predict_phonemes(hidden_states, slice(0, num_tokens)) @@ -964,9 +716,6 @@ def forward( 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) - if self._debug_outputs_enabled: - self._debug_lt_top_ids[valid] = self.code_predictor._debug_top_ids[:num_req] - self._debug_lt_top_values[valid] = self.code_predictor._debug_top_values[:num_req] self._flag_audio_eos(codes[:num_req], valid) if self.has_phoneme: self._predict_phonemes(hidden_states, valid) @@ -991,7 +740,7 @@ def _flag_audio_eos(self, codes: torch.Tensor, idx) -> None: 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._debug_decode_offset[idx].to(device=eos.device) + 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 @@ -1010,23 +759,18 @@ def _assemble_decode_embeddings(self, combined: torch.Tensor, idx) -> None: 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) - self._debug_audio_emb_norm[idx] = audio_emb.float().norm(dim=-1) 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) - self._debug_text_emb_norm[idx] = text_emb.float().norm(dim=-1) 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) - self._debug_phoneme_emb_norm[idx] = phon_emb.float().norm(dim=-1) assembled += phon_emb - else: - self._debug_phoneme_emb_norm[idx].zero_() combined[idx] = assembled @torch.no_grad() @@ -1118,148 +862,6 @@ def make_omni_output(self, model_outputs, **_: Any) -> OmniOutput: 0, row_indices.to(self._dec_phoneme_tokens.device), ).clone() - if self._debug_outputs_enabled: - combined = self._combined_embeddings.index_select( - 0, - row_indices.to(self._combined_embeddings.device), - ).clone() - hidden_debug = hidden[:num_tokens].clone() - multimodal_outputs.update( - { - "debug_text_mask": self._dec_text_mask.index_select( - 0, - row_indices.to(self._dec_text_mask.device), - ).clone(), - "debug_text_tokens": self._dec_text_tokens.index_select( - 0, - row_indices.to(self._dec_text_tokens.device), - ).clone(), - "debug_audio_valid": self._dec_audio_valid.index_select( - 0, - row_indices.to(self._dec_audio_valid.device), - ).clone(), - "debug_text_emb_norm": self._debug_text_emb_norm.index_select( - 0, - row_indices.to(self._debug_text_emb_norm.device), - ).clone(), - "debug_phoneme_emb_norm": self._debug_phoneme_emb_norm.index_select( - 0, - row_indices.to(self._debug_phoneme_emb_norm.device), - ).clone(), - "debug_audio_emb_norm": self._debug_audio_emb_norm.index_select( - 0, - row_indices.to(self._debug_audio_emb_norm.device), - ).clone(), - "debug_positions": self._debug_positions.index_select( - 0, - row_indices.to(self._debug_positions.device), - ).clone(), - "debug_input_ids": self._debug_input_ids.index_select( - 0, - row_indices.to(self._debug_input_ids.device), - ).clone(), - "debug_decode_dispatch_index": self._debug_decode_dispatch_index.index_select( - 0, - row_indices.to(self._debug_decode_dispatch_index.device), - ).clone(), - "debug_decode_dispatch_count": self._debug_decode_dispatch_count.index_select( - 0, - row_indices.to(self._debug_decode_dispatch_count.device), - ).clone(), - "debug_combined_input_norm": self._debug_combined_input_norm.index_select( - 0, - row_indices.to(self._debug_combined_input_norm.device), - ).clone(), - "debug_combined_input_vector": self._debug_combined_input_vector.index_select( - 0, - row_indices.to(self._debug_combined_input_vector.device), - ).clone(), - "debug_combined_norm": combined.float().norm(dim=-1), - "debug_combined_vector": combined, - "debug_hidden_norm": hidden_debug.float().norm(dim=-1), - "debug_hidden_vector": hidden_debug, - "debug_audio_input_code0": self._debug_audio_input_code0.index_select( - 0, - row_indices.to(self._debug_audio_input_code0.device), - ).clone(), - "debug_audio_output_code0": audio_codes[:, 0].clone(), - "debug_decode_offset": self._debug_decode_offset.index_select( - 0, - row_indices.to(self._debug_decode_offset.device), - ).clone(), - "debug_audio_feedback_missing": self._debug_audio_feedback_missing.index_select( - 0, - row_indices.to(self._debug_audio_feedback_missing.device), - ).clone(), - "debug_phoneme_input_valid": self._debug_phoneme_input_valid.index_select( - 0, - row_indices.to(self._debug_phoneme_input_valid.device), - ).clone(), - "debug_phoneme_input_token0": self._debug_phoneme_input_token0.index_select( - 0, - row_indices.to(self._debug_phoneme_input_token0.device), - ).clone(), - "debug_mamba_num_prefills": self._debug_mamba_num_prefills.index_select( - 0, - row_indices.to(self._debug_mamba_num_prefills.device), - ).clone(), - "debug_mamba_num_decodes": self._debug_mamba_num_decodes.index_select( - 0, - row_indices.to(self._debug_mamba_num_decodes.device), - ).clone(), - "debug_mamba_state_indices": self._debug_mamba_state_indices.index_select( - 0, - row_indices.to(self._debug_mamba_state_indices.device), - ).clone(), - "debug_mamba_cache_state_indices": self._debug_mamba_cache_state_indices.index_select( - 0, - row_indices.to(self._debug_mamba_cache_state_indices.device), - ).clone(), - "debug_mamba_exec_state_indices": self._debug_mamba_exec_state_indices.index_select( - 0, - row_indices.to(self._debug_mamba_exec_state_indices.device), - ).clone(), - "debug_attn_num_actual_tokens": self._debug_attn_num_actual_tokens.index_select( - 0, - row_indices.to(self._debug_attn_num_actual_tokens.device), - ).clone(), - "debug_attn_max_query_len": self._debug_attn_max_query_len.index_select( - 0, - row_indices.to(self._debug_attn_max_query_len.device), - ).clone(), - "debug_attn_max_seq_len": self._debug_attn_max_seq_len.index_select( - 0, - row_indices.to(self._debug_attn_max_seq_len.device), - ).clone(), - "debug_attn_seq_lens": self._debug_attn_seq_lens.index_select( - 0, - row_indices.to(self._debug_attn_seq_lens.device), - ).clone(), - "debug_attn_slot_mapping": self._debug_attn_slot_mapping.index_select( - 0, - row_indices.to(self._debug_attn_slot_mapping.device), - ).clone(), - "debug_lt_top_ids": self._debug_lt_top_ids.index_select( - 0, - row_indices.to(self._debug_lt_top_ids.device), - ).clone(), - "debug_lt_top_values": self._debug_lt_top_values.index_select( - 0, - row_indices.to(self._debug_lt_top_values.device), - ).clone(), - } - ) - if self.has_phoneme: - phoneme_indices = row_indices.to(self._dec_phoneme_valid.device) - multimodal_outputs.update( - { - "debug_phoneme_valid": self._dec_phoneme_valid.index_select(0, phoneme_indices).clone(), - "debug_phoneme_input_token0": self._dec_phoneme_tokens.index_select( - 0, - row_indices.to(self._dec_phoneme_tokens.device), - )[:, 0].clone(), - } - ) multimodal_outputs = { name: tensor.contiguous() if isinstance(tensor, torch.Tensor) else tensor for name, tensor in multimodal_outputs.items() @@ -1416,48 +1018,11 @@ def _clear_runtime_rows(self, start: int, end: int) -> None: "_out_code_logprobs", "_out_code_sampling_logprobs", "_out_frame_logprobs", - "_debug_lt_top_ids", - "_debug_lt_top_values", "_token_stop", "_sample_stop", - "_debug_combined_input_norm", - "_debug_combined_input_vector", - "_debug_text_emb_norm", - "_debug_phoneme_emb_norm", - "_debug_audio_emb_norm", - "_debug_combined_pre_norm", - "_debug_hidden_norm", - "_debug_backbone_last_layer_norm", - "_debug_backbone_last_residual_norm", - "_debug_final_norm_input_norm", - "_debug_final_norm_residual_norm", - "_debug_final_norm_output_norm", - "_debug_attn_qkv_norm", - "_debug_attn_core_norm", - "_debug_attn_output_norm", - "_debug_audio_feedback_missing", ) minus_one_names = ( - "_debug_positions", - "_debug_input_ids", - "_debug_decode_dispatch_index", - "_debug_decode_dispatch_count", - "_debug_backbone_first_bad_layer", - "_debug_backbone_first_bad_residual_layer", - "_debug_mamba_num_prefills", - "_debug_mamba_num_decodes", - "_debug_mamba_state_indices", - "_debug_mamba_cache_state_indices", - "_debug_mamba_exec_state_indices", - "_debug_attn_num_actual_tokens", - "_debug_attn_max_query_len", - "_debug_attn_max_seq_len", - "_debug_attn_seq_lens", - "_debug_attn_slot_mapping", - "_debug_decode_offset", - "_debug_audio_input_code0", - "_debug_phoneme_input_valid", - "_debug_phoneme_input_token0", + "_decode_offsets", ) for name in zero_names: value = getattr(self, name, None) @@ -1531,7 +1096,6 @@ def _preprocess_prefill( # ``additional_information`` and applied to the code predictor here (once, # at prefill — they are scalars that persist across decode steps). self._maybe_set_lt_sampling_params(info_dict) - self._debug_outputs_enabled = bool(info_dict.get("debug_outputs", False)) prefill_embeds = self._build_prefill_embeds(device, info_dict) @@ -1736,9 +1300,7 @@ def _preprocess_decode( 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) - debug_decode_offset = getattr(self, "_debug_decode_offset", None) - if isinstance(debug_decode_offset, torch.Tensor): - debug_decode_offset[start] = decode_offset + 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. Two @@ -1820,16 +1382,6 @@ def _preprocess_decode( self._dec_phoneme_valid[start] = 0 if phoneme_ended or feed_eos: info_update["phoneme_ended"] = True - debug_phoneme_valid = getattr(self, "_debug_phoneme_input_valid", None) - if isinstance(debug_phoneme_valid, torch.Tensor): - debug_phoneme_valid[start] = int(self._dec_phoneme_valid[start].item()) - debug_phoneme_token = getattr(self, "_debug_phoneme_input_token0", None) - if isinstance(debug_phoneme_token, torch.Tensor): - debug_phoneme_token[start] = ( - int(self._dec_phoneme_tokens[start, 0].item()) - if int(self._dec_phoneme_valid[start].item()) and self._dec_phoneme_tokens.ndim == 2 - else -1 - ) # ── Audio channel ── opens at decode step == ``speech_delay`` (seeded with # audio BOS), then feeds back the previous frame's codes. For the leading @@ -1845,26 +1397,14 @@ def _preprocess_decode( self._dec_audio_valid[start] = 1 else: last_codes = info_dict.get("last_audio_codes") - missing_feedback = True if isinstance(last_codes, torch.Tensor) and last_codes.numel() > 0: c = last_codes.to(device=device, dtype=torch.long).reshape(-1)[: self.num_codebooks] self._dec_audio_codes[start, : c.shape[0]].copy_(c) self._dec_audio_valid[start] = 1 - missing_feedback = False else: # Fallback (should not happen once audio has started): seed BOS. self._dec_audio_codes[start].fill_(self.arch.audio_bos_id) self._dec_audio_valid[start] = 1 - debug_missing = getattr(self, "_debug_audio_feedback_missing", None) - if isinstance(debug_missing, torch.Tensor): - debug_missing[start] = int(missing_feedback) - debug_audio_input = getattr(self, "_debug_audio_input_code0", None) - if isinstance(debug_audio_input, torch.Tensor): - debug_audio_input[start] = ( - int(self._dec_audio_codes[start, 0].item()) - if int(self._dec_audio_valid[start].item()) - else -1 - ) inputs_embeds_out = torch.zeros((1, self.embedding_dim), device=device, dtype=self._combined_embeddings.dtype) return input_ids, inputs_embeds_out, info_update @@ -1911,7 +1451,6 @@ def _reset_runtime_state_after_weight_load(self) -> None: """Clear mutable generation state after initial load or live refit.""" self.mamba_cache = None - self._debug_outputs_enabled = False self._last_output_row_indices = None for name in ( "_combined_embeddings", @@ -1925,15 +1464,14 @@ def _reset_runtime_state_after_weight_load(self) -> None: "_out_frame_logprobs", "_token_stop", "_sample_stop", - "_debug_text_emb_norm", - "_debug_phoneme_emb_norm", - "_debug_audio_emb_norm", "_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 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 1339376b1121..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 @@ -369,10 +369,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: 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.debug_collect_logits: bool = False - self.debug_logits_top_k: int = 5 - self._debug_top_ids = torch.full((max_num_tokens, self.num_codebooks, 5), -1, dtype=torch.long) - self._debug_top_values = torch.zeros(max_num_tokens, self.num_codebooks, 5, 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) @@ -481,9 +477,6 @@ def generate_codes_with_logprobs( buf_run.zero_() out_logprobs.zero_() out_sampling_logprobs.zero_() - if self.debug_collect_logits: - self._debug_top_ids[:num_tokens].fill_(-1) - self._debug_top_values[:num_tokens].zero_() # Row 0: projected backbone hidden state (the AR "prompt"). buf[:, 0, :] = self.local_transformer_in_projection(dec_hidden) @@ -497,14 +490,6 @@ def generate_codes_with_logprobs( 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) - if self.debug_collect_logits: - debug_logits = logits - if forbidden is not None: - debug_logits = debug_logits.masked_fill(forbidden, float("-inf")) - debug_k = min(int(self.debug_logits_top_k or 5), int(self._debug_top_ids.shape[-1]), int(debug_logits.shape[-1])) - vals, ids = torch.topk(debug_logits.detach().float(), k=debug_k, dim=-1) - self._debug_top_ids[:num_tokens, k, :debug_k].copy_(ids) - self._debug_top_values[:num_tokens, k, :debug_k].copy_(vals) code_k, model_lp, sampling_lp = sample_codebook_with_logprobs( logits, temperature=self.temperature, 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 index 5fa16e967118..03d5ba8e0502 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_easymagpie_backbone_cache.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_easymagpie_backbone_cache.py @@ -116,12 +116,7 @@ def _minimal_model(backbone: torch.nn.Module) -> EasyMagpieTTSForConditionalGene 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._debug_combined_pre_norm = torch.zeros(max_tokens) - model._debug_hidden_norm = torch.zeros(max_tokens) - model._debug_outputs_enabled = False - model._debug_backbone_layers_enabled = False - model._debug_backbone_active_tokens = 0 - model.code_predictor = SimpleNamespace(debug_collect_logits=False) + model.code_predictor = SimpleNamespace() model._get_decode_idxs = lambda: (torch.empty(0, dtype=torch.long), 0) return model 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 index 1d342caf024d..7b42063d59f1 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_sampler.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_local_transformer_sampler.py @@ -74,7 +74,6 @@ def test_code_predictor_pads_local_transformer_to_stable_graph_tokens(monkeypatc 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.debug_collect_logits = False 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() From 0527042199d7f7553b54d32fd018966b42a1e058 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 22:19:32 +0000 Subject: [PATCH 05/14] Compat EasyMagpie Omni model kwargs --- .../easymagpie_vllm_omni/vllm_compat.py | 51 ++++++++++++++++ .../tests/test_import_modes.py | 60 +++++++++++++++++++ 2 files changed, 111 insertions(+) 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 index ceae0ec47ff2..c7cdf2ab7055 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -184,6 +184,55 @@ def compat_launch_core_engines(*args: Any, addresses: Any = None, **kwargs: Any) 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) + for base_cls in type(self).__mro__[1:]: + base_method = base_cls.__dict__.get("_init_model_kwargs") + if base_method is None: + 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_v1_serial_utils_dense_tensor_compat() -> None: try: import msgspec.msgpack as msgpack @@ -1218,6 +1267,7 @@ def install_easy_magpie_refit_rpc_compat() -> None: _install_vllm_inputs_data_alias() _install_vllm_multimodal_inputs_alias() _install_engine_utils_compat() + _install_omni_gpu_model_runner_init_kwargs_compat() _install_easy_magpie_refit_rpc_compat() _install_async_omni_client_compat() @@ -1236,6 +1286,7 @@ def install_easy_magpie_runtime_compat() -> None: _install_vllm_inputs_data_alias() _install_vllm_multimodal_inputs_alias() _install_engine_utils_compat() + _install_omni_gpu_model_runner_init_kwargs_compat() _install_easy_magpie_refit_rpc_compat() _install_v1_serial_utils_dense_tensor_compat() _install_async_omni_client_compat() diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 2242a9180355..7a3261aef068 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -130,3 +130,63 @@ class FakeWorker: {"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) + + 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() + + assert OmniGPUModelRunner()._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) + + 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() + + assert OmniGPUModelRunner()._init_model_kwargs() == {"source": "old-vllm", "num_tokens": 19} + assert OmniGPUModelRunner()._init_model_kwargs(5) == {"source": "old-vllm", "num_tokens": 5} From 107a884ac3e56798d0d37b371cbba8f16cce0622 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 22:40:46 +0000 Subject: [PATCH 06/14] Fix EasyMagpie Omni init kwargs recursion --- .../easymagpie_vllm_omni/vllm_compat.py | 9 ++++++++- .../easymagpie_vllm_omni/tests/test_import_modes.py | 11 +++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) 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 index c7cdf2ab7055..08b299d9243d 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -203,10 +203,17 @@ def _install_omni_gpu_model_runner_init_kwargs_compat() -> None: 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) - for base_cls in type(self).__mro__[1:]: + 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): diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 7a3261aef068..7b372b04610d 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -147,6 +147,9 @@ def _init_model_kwargs(self, num_tokens=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")) @@ -161,8 +164,10 @@ def _init_model_kwargs(self, num_tokens=None): 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: @@ -180,13 +185,19 @@ def _init_model_kwargs(self, num_tokens=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} From 23e4bcd30dfba9aa67b03c3157ff9c72f18702d4 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 23:38:05 +0000 Subject: [PATCH 07/14] Compat Omni cudagraph fallback mode --- .../easymagpie_vllm_omni/vllm_compat.py | 61 +++++++++++++++++ .../tests/test_import_modes.py | 65 +++++++++++++++++++ 2 files changed, 126 insertions(+) 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 index 08b299d9243d..8915d1190b39 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -240,6 +240,65 @@ def _base_init_model_kwargs(self: Any, num_tokens: int | None = None) -> Any: _base_init_model_kwargs._easymagpie_original = current # type: ignore[attr-defined] runner_cls._init_model_kwargs = _base_init_model_kwargs + +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 when the call is not forced eager. + """ + + 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 + + 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] + if bool(kwargs.get("force_eager", False)) or cudagraph_mode != CUDAGraphMode.NONE: + return result + runtime_mode = _runtime_mode_from_config( + self, + force_uniform_decode=bool(kwargs.get("force_uniform_decode", False)), + ) + if runtime_mode == CUDAGraphMode.NONE: + return result + 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_v1_serial_utils_dense_tensor_compat() -> None: try: import msgspec.msgpack as msgpack @@ -1275,6 +1334,7 @@ def install_easy_magpie_refit_rpc_compat() -> None: _install_vllm_multimodal_inputs_alias() _install_engine_utils_compat() _install_omni_gpu_model_runner_init_kwargs_compat() + _install_omni_batch_execution_padding_compat() _install_easy_magpie_refit_rpc_compat() _install_async_omni_client_compat() @@ -1294,6 +1354,7 @@ def install_easy_magpie_runtime_compat() -> None: _install_vllm_multimodal_inputs_alias() _install_engine_utils_compat() _install_omni_gpu_model_runner_init_kwargs_compat() + _install_omni_batch_execution_padding_compat() _install_easy_magpie_refit_rpc_compat() _install_v1_serial_utils_dense_tensor_compat() _install_async_omni_client_compat() diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 7b372b04610d..7eb814374af6 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -201,3 +201,68 @@ class RuntimeRunner(OmniGPUModelRunner): 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_reports_configured_cudagraph_mode_for_omni_fallback(monkeypatch) -> None: + from vllm.config import CUDAGraphMode + + from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat + + class BatchDescriptor: + num_tokens = 8 + num_reqs = 4 + + class OmniGPUModelRunner: + def __init__(self, cudagraph_mode): + self.vllm_config = types.SimpleNamespace( + compilation_config=types.SimpleNamespace(cudagraph_mode=cudagraph_mode) + ) + + def _init_model_kwargs(self, num_tokens=None): + return {} + + def _determine_batch_execution_and_padding(self, **kwargs): + return CUDAGraphMode.NONE, BatchDescriptor(), 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, + ) + + 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 From 613f3cfb50dedc1d3ea5150469b91ae25c8db2a7 Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 25 Jun 2026 00:13:32 +0000 Subject: [PATCH 08/14] Compat Omni structured output requests --- .../easymagpie_vllm_omni/vllm_compat.py | 46 +++++++++++++++++++ .../tests/test_import_modes.py | 36 +++++++++++++++ 2 files changed, 82 insertions(+) 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 index 8915d1190b39..2af20d98f7fd 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -299,6 +299,50 @@ def _determine_batch_execution_and_padding(self: Any, *args: Any, **kwargs: Any) _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_v1_serial_utils_dense_tensor_compat() -> None: try: import msgspec.msgpack as msgpack @@ -1335,6 +1379,7 @@ def install_easy_magpie_refit_rpc_compat() -> None: _install_engine_utils_compat() _install_omni_gpu_model_runner_init_kwargs_compat() _install_omni_batch_execution_padding_compat() + _install_omni_structured_output_request_compat() _install_easy_magpie_refit_rpc_compat() _install_async_omni_client_compat() @@ -1355,6 +1400,7 @@ def install_easy_magpie_runtime_compat() -> None: _install_engine_utils_compat() _install_omni_gpu_model_runner_init_kwargs_compat() _install_omni_batch_execution_padding_compat() + _install_omni_structured_output_request_compat() _install_easy_magpie_refit_rpc_compat() _install_v1_serial_utils_dense_tensor_compat() _install_async_omni_client_compat() diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 7eb814374af6..a0d54002bc7d 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -266,3 +266,39 @@ def _determine_batch_execution_and_padding(self, **kwargs): assert eager_result[0] == CUDAGraphMode.NONE assert full_and_piecewise_mixed[0] == CUDAGraphMode.PIECEWISE assert full_and_piecewise_decode[0] == CUDAGraphMode.FULL + + +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 From 4fb31c9f6121a3c6bb1551feb18a1f322d7330d1 Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 25 Jun 2026 00:31:57 +0000 Subject: [PATCH 09/14] Compat Omni request multimodal kwargs --- .../easymagpie_vllm_omni/vllm_compat.py | 35 +++++++++++++++++++ .../tests/test_import_modes.py | 25 +++++++++++++ 2 files changed, 60 insertions(+) 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 index 2af20d98f7fd..dd8c7ea468b9 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -343,6 +343,39 @@ def structured_output_request_compat(*args: Any, **kwargs: Any) -> Any: 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 @@ -1380,6 +1413,7 @@ def install_easy_magpie_refit_rpc_compat() -> None: _install_omni_gpu_model_runner_init_kwargs_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() @@ -1401,6 +1435,7 @@ def install_easy_magpie_runtime_compat() -> None: _install_omni_gpu_model_runner_init_kwargs_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() diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index a0d54002bc7d..32b663d8c8ed 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -302,3 +302,28 @@ def __init__(self): 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 From 46f7a5a148167f2026918126e6d7eed5e368bc0c Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 25 Jun 2026 01:26:14 +0000 Subject: [PATCH 10/14] Compat vLLM cumsum arange signature --- .../easymagpie_vllm_omni/vllm_compat.py | 59 +++++++++++ .../tests/test_import_modes.py | 98 ++++++++++++++++++- 2 files changed, 156 insertions(+), 1 deletion(-) 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 index dd8c7ea468b9..898455b11423 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -241,6 +241,63 @@ def _base_init_model_kwargs(self: Any, num_tokens: int | None = None) -> Any: 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 + if not supplied_positionally and "arange_out" not in kwargs: + kwargs["arange_out"] = _make_arange_out(self, num_tokens) + return current(self, num_tokens, *args, **kwargs) + + _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. @@ -1411,6 +1468,7 @@ def install_easy_magpie_refit_rpc_compat() -> None: _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() @@ -1433,6 +1491,7 @@ def install_easy_magpie_runtime_compat() -> None: _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() diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 32b663d8c8ed..c4ce75e89c7c 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -203,8 +203,104 @@ class RuntimeRunner(OmniGPUModelRunner): 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_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 vllm.config import CUDAGraphMode + 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 From 794f258c870eddecc7bd324cf2b089c590394729 Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 25 Jun 2026 02:13:59 +0000 Subject: [PATCH 11/14] Normalize Omni cumsum arange return --- .../easymagpie_vllm_omni/vllm_compat.py | 13 ++++++-- .../tests/test_import_modes.py | 32 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) 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 index 898455b11423..b8d752f0ed8b 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -289,9 +289,18 @@ def _make_arange_out(self: Any, num_tokens: Any) -> Any: 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: - kwargs["arange_out"] = _make_arange_out(self, num_tokens) - return current(self, num_tokens, *args, **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] diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index c4ce75e89c7c..52869f8c5a9e 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -243,6 +243,38 @@ class OmniGPUModelRunner(NewParentRunner): 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 3f66802f2927ca64233f25998eb32fd731579a09 Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 25 Jun 2026 02:38:24 +0000 Subject: [PATCH 12/14] Stabilize EasyMagpie cudagraph descriptors --- .../easymagpie_vllm_omni/vllm_compat.py | 68 ++++++++++++++++++- .../tests/test_import_modes.py | 14 ++++ 2 files changed, 80 insertions(+), 2 deletions(-) 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 index b8d752f0ed8b..61521f2a53e2 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -316,6 +316,13 @@ def _install_omni_batch_execution_padding_compat() -> None: explicitly into ``_dummy_run`` and asserts it matches the returned mode. Preserve Omni's simple descriptor while reporting the configured runtime graph mode when the call is not forced eager. + + 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: @@ -343,19 +350,76 @@ def _runtime_mode_from_config(self: Any, *, force_uniform_decode: bool) -> Any: 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 _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: - return result + 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: - return result + patched = (cudagraph_mode, 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:] diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 52869f8c5a9e..42a64ba40788 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -395,6 +395,20 @@ def _determine_batch_execution_and_padding(self, **kwargs): assert full_and_piecewise_mixed[0] == CUDAGraphMode.PIECEWISE assert full_and_piecewise_decode[0] == CUDAGraphMode.FULL + 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 From f1b5dc06db0da0ac051a362fdba195e2e51d202c Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 25 Jun 2026 03:00:17 +0000 Subject: [PATCH 13/14] Trim unused EasyMagpie backbone patch entrypoints --- .../easymagpie_vllm_omni/backbone_patches.py | 93 ------------------- 1 file changed, 93 deletions(-) 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 59976b7002ba..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 @@ -751,99 +751,6 @@ def forward( return True -def patch_flash_attention_torch_fallback() -> bool: - """Use PyTorch SDPA for EasyMagpie backbone FlashAttention calls. - - The converted EasyMagpie vLLM path currently reaches FlashAttention with - finite QKV tensors at the first Nemotron-H attention layer, then receives - NaN attention outputs. Keep vLLM's scheduler and KV cache contract intact, - but replace only that kernel call for EasyMagpie backbone layers. - """ - - patched = 0 - try: - from vllm.attention.backends.flash_attn import FlashAttentionImpl as LegacyFlashAttentionImpl - - patched += int(_patch_flash_attention_impl(LegacyFlashAttentionImpl, is_v1=False)) - except Exception: - pass - try: - from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl as V1FlashAttentionImpl - - patched += int(_patch_flash_attention_impl(V1FlashAttentionImpl, is_v1=True)) - except Exception: - pass - patched += int(_patch_v1_flash_attn_varlen_func()) - if patched: - logger.info("EasyMagpie FlashAttention torch-SDPA fallback installed on %d implementation(s)", patched) - return bool(patched) - - -def patch_final_rmsnorm_native(backbone) -> bool: - """Run the final Nemotron-H RMSNorm through the fp32 native implementation. - - EasyMagpie feeds the backbone through a persistent ``inputs_embeds`` scratch - buffer. vLLM's fused CUDA add+RMSNorm path is allowed to mutate its inputs - in-place; on this checkpoint it leaves the final hidden state non-finite - even after decoder-layer attention outputs are finite. The native RMSNorm - path performs the residual add and variance in fp32 and returns a fresh - tensor, matching the safer reference math closely enough for inference. - """ - - enabled = os.environ.get("EASYMAGPIE_FINAL_RMSNORM_NATIVE", "1").lower() - if enabled in {"0", "false", "no", "off"}: - return False - - norm_f = getattr(backbone, "norm_f", None) - if norm_f is None or not hasattr(norm_f, "forward_native"): - return False - if getattr(norm_f, "_easymagpie_native_forward_installed", False): - return False - - def forward(x, residual=None): - return norm_f.forward_native(x, residual) - - norm_f.forward = forward - norm_f._easymagpie_native_forward_installed = True - logger.info("EasyMagpie final RMSNorm native fp32 fallback installed") - return True - - -def patch_layer_rmsnorm_native(backbone) -> int: - """Run Nemotron-H decoder-layer RMSNorms through native fp32 math. - - The layer RMSNorm modules update the long-lived residual stream consumed by - every following hybrid layer. If the fused add+RMSNorm kernel writes a - non-finite residual, the final hidden state can be corrupted even when each - mixer's direct output looks finite. The native path preserves vLLM's return - contract while doing the residual add and variance in fp32. - """ - - enabled = os.environ.get("EASYMAGPIE_LAYER_RMSNORM_NATIVE", "1").lower() - if enabled in {"0", "false", "no", "off"}: - return 0 - - patched = 0 - for layer_idx, layer in enumerate(getattr(backbone, "layers", []) or []): - norm = getattr(layer, "norm", None) - if norm is None or not hasattr(norm, "forward_native"): - continue - if getattr(norm, "_easymagpie_native_forward_installed", False): - continue - - def forward(x, residual=None, *, _norm=norm): - return _norm.forward_native(x, residual) - - norm.forward = forward - norm._easymagpie_native_forward_installed = True - norm._easymagpie_native_forward_layer_idx = layer_idx - patched += 1 - - if patched: - logger.info("EasyMagpie layer RMSNorm native fallback installed on %d module(s)", patched) - return patched - - def patch_mamba_streaming_decode() -> None: """Treat 1-token streaming extends as decodes so FULL decode cudagraphs work. From 180fff16d2862dd2e8cd59ec2337d30a51e242f7 Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 25 Jun 2026 03:28:57 +0000 Subject: [PATCH 14/14] Avoid lazy cudagraph capture for oversized batches --- .../easymagpie_vllm_omni/vllm_compat.py | 22 +++++++++++++- .../tests/test_import_modes.py | 30 ++++++++++++++++--- 2 files changed, 47 insertions(+), 5 deletions(-) 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 index 61521f2a53e2..c497a290cbb3 100644 --- a/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py +++ b/examples/tts/easymagpie_vllm_omni/easymagpie_vllm_omni/vllm_compat.py @@ -315,7 +315,9 @@ def _install_omni_batch_execution_padding_compat() -> None: ``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 when the call is not forced eager. + 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 @@ -400,6 +402,19 @@ def _normalize_batch_descriptor(batch_desc: Any) -> Any: 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: @@ -420,6 +435,11 @@ def _determine_batch_execution_and_padding(self: Any, *args: Any, **kwargs: Any) 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:] diff --git a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py index 42a64ba40788..c612d3ac84fb 100644 --- a/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py +++ b/examples/tts/easymagpie_vllm_omni/tests/test_import_modes.py @@ -337,20 +337,34 @@ def decode_mode(self): from easymagpie_vllm_omni.vllm_compat import install_easy_magpie_runtime_compat class BatchDescriptor: - num_tokens = 8 - num_reqs = 4 + 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) + 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(), False, None, None + 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 @@ -389,11 +403,19 @@ def _determine_batch_execution_and_padding(self, **kwargs): 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(