From 1eb21dd1af68876a1525ce244d0f2b7cd3c4e359 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 22 Sep 2026 19:58:04 -0700 Subject: [PATCH 1/3] Add Alpamayo Edge exporter support --- tools/hf/exporters/models/__init__.py | 1 + .../hf/exporters/models/alpamayo/__init__.py | 1 + tools/hf/exporters/models/alpamayo/export.py | 65 +++ tools/hf/exporters/models/alpamayo/helpers.py | 107 ++++ tools/hf/exporters/models/alpamayo/patches.py | 185 +++++++ tools/hf/exporters/models/alpamayo/spec.py | 466 ++++++++++++++++++ tools/hf/exporters/models/common/helpers.py | 7 +- .../hf/exporters/tests/test_edge_exporter.py | 55 +++ tools/hf/run_export.py | 11 + 9 files changed, 897 insertions(+), 1 deletion(-) create mode 100644 tools/hf/exporters/models/alpamayo/__init__.py create mode 100644 tools/hf/exporters/models/alpamayo/export.py create mode 100644 tools/hf/exporters/models/alpamayo/helpers.py create mode 100644 tools/hf/exporters/models/alpamayo/patches.py create mode 100644 tools/hf/exporters/models/alpamayo/spec.py diff --git a/tools/hf/exporters/models/__init__.py b/tools/hf/exporters/models/__init__.py index 182bbd3efa..0a8768a9c5 100644 --- a/tools/hf/exporters/models/__init__.py +++ b/tools/hf/exporters/models/__init__.py @@ -1,5 +1,6 @@ """Model families. Importing this package registers EdgeSpecs.""" +from .alpamayo import spec as _alpamayo # noqa: F401 from .groot import spec as _groot # noqa: F401 from .kimi import spec as _kimi # noqa: F401 from .nemotron import spec as _nemotron # noqa: F401 diff --git a/tools/hf/exporters/models/alpamayo/__init__.py b/tools/hf/exporters/models/alpamayo/__init__.py new file mode 100644 index 0000000000..1f3410e866 --- /dev/null +++ b/tools/hf/exporters/models/alpamayo/__init__.py @@ -0,0 +1 @@ +"""Alpamayo 1.5 Edge exporter support.""" diff --git a/tools/hf/exporters/models/alpamayo/export.py b/tools/hf/exporters/models/alpamayo/export.py new file mode 100644 index 0000000000..2d52ec80d8 --- /dev/null +++ b/tools/hf/exporters/models/alpamayo/export.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import argparse + +import torch + +from ...config import EdgeConfig +from ...utils import force_hf_attention +from .helpers import ( + alpamayo_language, + alpamayo_visual, +) + +DEFAULT_CHECKPOINT = "nvidia/Alpamayo-1.5-10B" + + +def prepare_export( + args: argparse.Namespace, + device: torch.device, + dtype: torch.dtype, +): + """Load a base or ModelOpt-quantized Alpamayo checkpoint.""" + try: + import modelopt.torch.opt as mto + except ImportError as exc: + raise ImportError( + "Alpamayo export requires NVIDIA ModelOpt. " + "Install the nvidia-modelopt package." + ) from exc + + # Must be enabled before ``from_pretrained`` restores modelopt_state.pth. + mto.enable_huggingface_checkpointing() + + from alpamayo1_5.models.alpamayo1_5 import Alpamayo1_5 + + model = ( + Alpamayo1_5.from_pretrained( + args.checkpoint or DEFAULT_CHECKPOINT, + dtype=dtype, + attn_implementation="eager", + ) + .to(device=device, dtype=dtype) + .eval() + ) + + force_hf_attention(alpamayo_visual(model), "eager") + force_hf_attention(alpamayo_language(model), "eager") + force_hf_attention(model.expert, "eager") + + export_config = EdgeConfig( + model_type="alpamayo", + engine_dir=args.engine_dir or "/tmp/alpamayo_edge_exporter", + max_seq_len=args.max_seq_len or 4096, + ) + return ( + model, + { + "device": device, + "dtype": dtype, + "clip_id": getattr(args, "clip_id", None), + "t0_us": getattr(args, "t0_us", 5_100_000), + }, + export_config, + "velocity", + ) diff --git a/tools/hf/exporters/models/alpamayo/helpers.py b/tools/hf/exporters/models/alpamayo/helpers.py new file mode 100644 index 0000000000..d3c8956c8d --- /dev/null +++ b/tools/hf/exporters/models/alpamayo/helpers.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn + + +def alpamayo_vlm(model: nn.Module) -> nn.Module: + """Return the Qwen3-VL conditional-generation module.""" + vlm = getattr(model, "vlm", None) + if not isinstance(vlm, nn.Module): + raise AttributeError(f"{type(model).__name__} has no Alpamayo VLM") + return vlm + + +def alpamayo_vlm_core(model: nn.Module) -> nn.Module: + """Return the Qwen3-VL model containing visual and language towers.""" + core = getattr(alpamayo_vlm(model), "model", None) + if not isinstance(core, nn.Module): + raise AttributeError("Alpamayo VLM has no model") + return core + + +def alpamayo_visual(model: nn.Module) -> nn.Module: + visual = getattr(alpamayo_vlm_core(model), "visual", None) + if not isinstance(visual, nn.Module): + raise AttributeError("Alpamayo VLM has no visual tower") + return visual + + +def alpamayo_language(model: nn.Module) -> nn.Module: + language = getattr(alpamayo_vlm_core(model), "language_model", None) + if not isinstance(language, nn.Module): + language = getattr(alpamayo_vlm(model), "language_model", None) + if not isinstance(language, nn.Module): + raise AttributeError("Alpamayo VLM has no language model") + return language + + +def stack_deepstack_features(features: Any) -> torch.Tensor: + """Normalize Qwen3-VL deepstack features to ``[N, tokens, hidden]``.""" + if isinstance(features, torch.Tensor): + return features + if not isinstance(features, (tuple, list)) or not features: + raise ValueError("Alpamayo visual tower returned no deepstack features") + return torch.stack(tuple(features), dim=0) + + +def scatter_visual_tokens( + visual: torch.Tensor, + text_embeds: torch.Tensor, + image_token_mask: torch.Tensor, +) -> torch.Tensor: + """Insert flattened visual features into Qwen image-token positions.""" + hidden = int(text_embeds.shape[-1]) + flat = text_embeds.reshape(-1, hidden).clone() + mask = image_token_mask.reshape(-1) + values = visual.reshape(-1, hidden).to(device=flat.device, dtype=flat.dtype) + count = int(mask.sum().item()) + if count != int(values.shape[0]): + raise ValueError( + "Alpamayo image token count does not match visual features: " + f"{count} tokens vs {values.shape[0]} features" + ) + flat[mask] = values + return flat.reshape_as(text_embeds) + + +def make_deepstack_tensor( + features: torch.Tensor, + image_token_mask: torch.Tensor, + *, + num_layers: int, + batch_size: int, + seq_len: int, + hidden_size: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Expand sparse Qwen deepstack features into the common dense layout.""" + dense = torch.zeros( + num_layers, + batch_size, + seq_len, + hidden_size, + device=device, + dtype=dtype, + ) + zero_text = torch.zeros( + batch_size, + seq_len, + hidden_size, + device=device, + dtype=dtype, + ) + if int(features.shape[0]) > num_layers: + raise ValueError( + f"{features.shape[0]} deepstack stages exceed {num_layers} language layers" + ) + for layer_index in range(int(features.shape[0])): + dense[layer_index] = scatter_visual_tokens( + features[layer_index], + zero_text, + image_token_mask, + ) + return dense diff --git a/tools/hf/exporters/models/alpamayo/patches.py b/tools/hf/exporters/models/alpamayo/patches.py new file mode 100644 index 0000000000..0b160b5842 --- /dev/null +++ b/tools/hf/exporters/models/alpamayo/patches.py @@ -0,0 +1,185 @@ +"""Alpamayo setattr replacements used only while compiling Edge engines.""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch + +from ...plugin.attn_patches import ( + _patch_language_attention, + register_patch, +) +from ...prefix_cache import PrefixKVCache +from ..common.patches import causal_lm_plugin_forward + +ALPAMAYO = "alpamayo" + + +@register_patch( + ALPAMAYO, + "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionAttention.forward", +) +def _patch_qwen3_vl_vision_attention(original: Callable) -> Callable: + """Route Qwen3-VL vision attention through the Edge ViT plugin.""" + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb=None, + position_embeddings=None, + **kwargs: Any, + ) -> torch.Tensor: + del rotary_pos_emb, kwargs + if position_embeddings is None: + return original( + self, + hidden_states, + cu_seqlens, + position_embeddings=position_embeddings, + ) + + from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + apply_rotary_pos_emb_vision, + ) + + seq_len = int(hidden_states.shape[0]) + q, k, v = ( + self.qkv(hidden_states) + .reshape(seq_len, 3, self.num_heads, -1) + .permute(1, 0, 2, 3) + .unbind(0) + ) + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + q = q.to(torch.float16).contiguous() + k = k.to(torch.float16).contiguous() + v = v.to(torch.float16).contiguous() + + # The carrier's length communicates a safe maximum sequence length. + max_seqlen_carrier = torch.zeros( + hidden_states.shape[0], + device=hidden_states.device, + dtype=torch.int32, + ) + output = torch.ops.trt.vit_attention_plugin.default( + q, + k, + v, + cu_seqlens.to(torch.int32), + max_seqlen_carrier, + int(self.num_heads), + int(q.shape[-1]), + ) + output = output.reshape(seq_len, -1).to(hidden_states.dtype) + return self.proj(output) + + return forward + + +@register_patch( + ALPAMAYO, + "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionModel.forward", +) +def _patch_qwen3_vl_vision_model(original: Callable) -> Callable: + """Return a tensor rather than a Python list for deepstack outputs.""" + + def forward(self, hidden_states, grid_thw, **kwargs: Any): + visual, deepstack = original( + self, + hidden_states, + grid_thw, + **kwargs, + ) + return visual, torch.stack(tuple(deepstack), dim=0) + + return forward + + +register_patch( + ALPAMAYO, + "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextAttention.forward", +)(_patch_language_attention) + + +@register_patch( + ALPAMAYO, + "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLForConditionalGeneration.forward", + "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextModel.forward", +) +def _patch_qwen3_vl_language(original: Callable) -> Callable: + """Use flattened AttentionPlugin I/O for language prefill.""" + + def forward( + self, + inputs_embeds=None, + rope_rotary_cos_sin=None, + context_lengths=None, + kvcache_start_index=None, + last_token_ids=None, + ds_stack=None, + *past_key_values, + **kwargs: Any, + ): + if rope_rotary_cos_sin is None: + return original(self, inputs_embeds=inputs_embeds, **kwargs) + + root = getattr(self, "model", None) + decoder = getattr(root, "language_model", None) + if decoder is None: + decoder = self + return causal_lm_plugin_forward( + decoder, + inputs_embeds, + rope_rotary_cos_sin, + context_lengths, + kvcache_start_index, + last_token_ids, + ds_stack, + *past_key_values, + lm_head=getattr(self, "lm_head", None), + ) + + return forward + + +@register_patch( + ALPAMAYO, + "alpamayo1_5.models.alpamayo1_5.Alpamayo1_5.forward", +) +def _patch_alpamayo_action_step(original: Callable) -> Callable: + """Compile one Alpamayo diffusion velocity step with prefix KV tensors.""" + + def forward( + self, + noisy_action, + timestep=None, + prefix_k=None, + prefix_v=None, + position_ids=None, + attention_mask=None, + *args, + **kwargs: Any, + ): + if prefix_k is None or getattr(prefix_k, "ndim", 0) != 5: + return original(self, noisy_action, timestep, *args, **kwargs) + + action_embeds = self.action_in_proj(noisy_action, timestep) + expert_kwargs: dict[str, Any] = {} + if self.config.expert_non_causal_attention: + expert_kwargs["is_causal"] = False + expert = self.expert( + inputs_embeds=action_embeds, + position_ids=position_ids, + past_key_values=PrefixKVCache(prefix_k, prefix_v), + attention_mask=attention_mask, + use_cache=False, + return_dict=True, + **expert_kwargs, + ) + hidden = expert.last_hidden_state + hidden = hidden[:, -int(noisy_action.shape[1]) :] + return self.action_out_proj(hidden).reshape_as(noisy_action) + + return forward diff --git a/tools/hf/exporters/models/alpamayo/spec.py b/tools/hf/exporters/models/alpamayo/spec.py new file mode 100644 index 0000000000..afca679165 --- /dev/null +++ b/tools/hf/exporters/models/alpamayo/spec.py @@ -0,0 +1,466 @@ +from __future__ import annotations + +from collections.abc import Mapping, MutableMapping +from typing import Any + +import torch +import torch.nn as nn + +from ...ops import call_engine, scatter_image_tokens +from ...prefix_cache import PrefixKVCache +from ...spec import ComponentBundle, EdgeSpec, register_edge_spec +from ..common.helpers import ( + causal_lm_flat, + kv_kwargs, + split_flat_to_kwargs, +) +from .helpers import ( + alpamayo_language, + alpamayo_visual, + alpamayo_vlm, + alpamayo_vlm_core, + make_deepstack_tensor, + scatter_visual_tokens, + stack_deepstack_features, +) +from .patches import ALPAMAYO + + +def _export_module(module: nn.Module, device: torch.device, dtype: torch.dtype): + return module.eval().to(device=device, dtype=dtype) + + +def _append_generation_positions( + position_ids: torch.Tensor, + rope_deltas: torch.Tensor, + max_seq_len: int, +) -> torch.Tensor: + """Extend Qwen multimodal positions with ordinary text decode positions.""" + prompt_len = int(position_ids.shape[-1]) + if prompt_len >= max_seq_len: + return position_ids[..., :max_seq_len] + batch_size = int(position_ids.shape[1]) + tail = torch.arange( + prompt_len, + max_seq_len, + device=position_ids.device, + dtype=position_ids.dtype, + ) + tail = tail.unsqueeze(0).expand(batch_size, -1) + tail = tail + rope_deltas.to(device=tail.device, dtype=tail.dtype) + tail = tail.unsqueeze(0).expand(3, -1, -1) + return torch.cat((position_ids, tail), dim=-1) + + +@register_edge_spec("alpamayo", "alpamayo_r1", "alpamayo1_5") +class AlpamayoSpec(EdgeSpec): # type: ignore[misc] + def apply_patches(self, model=None): + del model + from ...plugin.attn_patches import apply_patches + + return apply_patches(ALPAMAYO) + + def prepare_sample_inputs( + self, + model: nn.Module, + raw: Mapping[str, Any], + config: Any, + ) -> MutableMapping[str, Any]: + del config + required = { + "input_ids", + "attention_mask", + "pixel_values", + "image_grid_thw", + } + if required.issubset(raw): + return dict(raw) + + clip_id = raw.get("clip_id") + if not clip_id: + raise ValueError( + "Alpamayo export needs --clip-id, or pre-tokenized input_ids, " + "attention_mask, pixel_values, and image_grid_thw." + ) + + from alpamayo1_5 import helper + from alpamayo1_5.load_physical_aiavdataset import ( + load_physical_aiavdataset, + ) + + device = raw.get( + "device", + torch.device("cuda" if torch.cuda.is_available() else "cpu"), + ) + dtype = raw.get("dtype", torch.float16) + data = raw.get("data") or load_physical_aiavdataset( + str(clip_id), + t0_us=int(raw.get("t0_us", 5_100_000)), + ) + messages = helper.create_message( + data["image_frames"].flatten(0, 1), + camera_indices=data["camera_indices"], + ) + processor = helper.get_processor(model.tokenizer) + tokenized = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=False, + continue_final_message=True, + return_dict=True, + return_tensors="pt", + ) + input_ids = model.fuse_traj_tokens( + tokenized["input_ids"], + { + "ego_history_xyz": data["ego_history_xyz"], + "ego_history_rot": data["ego_history_rot"], + }, + ) + return { + "input_ids": input_ids.to(device=device, dtype=torch.long), + "attention_mask": tokenized["attention_mask"].to( + device=device, + dtype=torch.long, + ), + "pixel_values": tokenized["pixel_values"].to( + device=device, + dtype=dtype, + ), + "image_grid_thw": tokenized["image_grid_thw"].to( + device=device, + dtype=torch.long, + ), + } + + def capture_eager_outputs( + self, + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + bench: dict[str, float] | None = None, + ) -> dict[str, torch.Tensor]: + del config + from ...measure import cuda_ms + + visual = alpamayo_visual(model) + language = alpamayo_language(model) + px = sample["pixel_values"] + grid = sample["image_grid_thw"] + + with torch.no_grad(): + visual_embeds, deepstack = visual(px, grid) + deepstack = stack_deepstack_features(deepstack) + language_out = language( + inputs_embeds=sample["inputs_embeds"], + attention_mask=sample["attention_mask"], + position_ids=sample["position_ids"], + visual_pos_masks=sample["image_token_mask"], + deepstack_visual_embeds=[ + deepstack[i] for i in range(int(deepstack.shape[0])) + ], + use_cache=False, + return_dict=True, + ) + action_embeds = model.action_in_proj( + sample["step_actions"], + sample["step_timestep"], + ) + expert = model.expert( + inputs_embeds=action_embeds, + position_ids=sample["suffix_position_ids"], + past_key_values=PrefixKVCache( + sample["prefix_k"], + sample["prefix_v"], + ), + attention_mask=sample["suffix_attention_mask"], + use_cache=False, + return_dict=True, + is_causal=not bool(model.config.expert_non_causal_attention), + ) + velocity = model.action_out_proj( + expert.last_hidden_state[:, -int(sample["step_actions"].shape[1]) :] + ).reshape_as(sample["step_actions"]) + + if bench is not None: + bench["vision"] = cuda_ms(lambda: visual(px, grid)[0]) + bench["language"] = cuda_ms( + lambda: language( + inputs_embeds=sample["inputs_embeds"], + attention_mask=sample["attention_mask"], + position_ids=sample["position_ids"], + visual_pos_masks=sample["image_token_mask"], + deepstack_visual_embeds=[ + deepstack[i] for i in range(int(deepstack.shape[0])) + ], + use_cache=False, + return_dict=True, + ).last_hidden_state + ) + return { + "vision": visual_embeds, + "language": language_out.last_hidden_state, + "action": velocity, + } + + def prepare( + self, + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + ) -> dict[str, ComponentBundle]: + from ...plugin.attention import ContextAttentionMaskType + + visual = alpamayo_visual(model) + language = alpamayo_language(model) + vlm = alpamayo_vlm(model) + vlm_core = alpamayo_vlm_core(model) + px = sample["pixel_values"] + grid = sample["image_grid_thw"] + device = px.device + dtype = px.dtype + + vision = ComponentBundle( + module=_export_module(visual, device, dtype), + trace_args=(px, grid), + save_args=(px, grid), + input_names=["pixel_values", "image_grid_thw"], + output_names=["visual_embeds", "deepstack_visual_embeds"], + parity_output="visual_embeds", + model_type="qwen3_vl", + engine_file="visual.engine", + trt_settings={ + "disable_tf32": False, + "use_fp32_acc": False, + "use_explicit_typing": False, + "decompose_attention": True, + }, + ) + + with torch.no_grad(): + visual_embeds, deepstack = visual(px, grid) + deepstack = stack_deepstack_features(deepstack) + text_embeds = model.get_input_embeddings()(sample["input_ids"]) + + image_token_id = int(vlm.config.image_token_id) + image_token_mask = sample["input_ids"] == image_token_id + inputs_embeds = scatter_visual_tokens( + visual_embeds, + text_embeds, + image_token_mask, + ).to(device=device, dtype=dtype) + sample["lang_embeds"] = text_embeds.to(device=device, dtype=dtype) + sample["image_token_mask"] = image_token_mask + sample["inputs_embeds"] = inputs_embeds + + position_ids, rope_deltas = vlm_core.get_rope_index( + sample["input_ids"], + sample["image_grid_thw"], + None, + attention_mask=sample["attention_mask"], + ) + sample["position_ids"] = position_ids + sample["rope_deltas"] = rope_deltas + + decoder_layers = language.layers + ds_stack = make_deepstack_tensor( + deepstack, + image_token_mask, + num_layers=len(decoder_layers), + batch_size=int(inputs_embeds.shape[0]), + seq_len=int(inputs_embeds.shape[1]), + hidden_size=int(inputs_embeds.shape[2]), + device=device, + dtype=dtype, + ) + sample["ds_template"] = ds_stack + + prompt_len = int(inputs_embeds.shape[1]) + if int(config.generation_reserve) < 0: + raise ValueError("generation_reserve must be non-negative") + max_seq_len = max( + int(config.max_seq_len), + prompt_len + int(config.generation_reserve), + ) + full_position_ids = _append_generation_positions( + position_ids, + rope_deltas, + max_seq_len, + ) + flat, meta = causal_lm_flat( + language, + inputs_embeds, + max_seq_len=max_seq_len, + device=device, + dtype=dtype, + seq_len=prompt_len, + position_ids=full_position_ids, + ) + flat = (*flat[:5], ds_stack, *flat[6:]) + sample.update(split_flat_to_kwargs(flat, meta["input_names"])) + + language_bundle = ComponentBundle( + module=_export_module(vlm, device, dtype), + trace_args=flat, + save_args=flat, + input_names=meta["input_names"], + output_names=["logits", "lm_hidden_states", "prefix_k", "prefix_v"], + parity_output="lm_hidden_states", + context_attention_mask_type=int(ContextAttentionMaskType.CAUSAL), + model_type="language", + engine_file="language.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + "assume_dynamic_shape_support": True, + "immutable_weights": False, + }, + ) + + action_dims = tuple(model.action_space.get_action_space_dims()) + batch_size = int(inputs_embeds.shape[0]) + step_actions = sample.get( + "step_actions", + torch.randn( + batch_size, + *action_dims, + device=device, + dtype=dtype, + ), + ) + step_timestep = sample.get( + "step_timestep", + torch.full( + (batch_size, 1, 1), + 0.5, + device=device, + dtype=dtype, + ), + ) + sample["step_actions"] = step_actions + sample["step_timestep"] = step_timestep + + language_cfg = language.config + num_kv_heads = int(language_cfg.num_key_value_heads) + head_dim = int( + getattr( + language_cfg, + "head_dim", + language_cfg.hidden_size // language_cfg.num_attention_heads, + ) + ) + prefix_k = torch.zeros( + len(decoder_layers), + batch_size, + num_kv_heads, + prompt_len, + head_dim, + device=device, + dtype=dtype, + ) + prefix_v = torch.zeros_like(prefix_k) + suffix_position_ids, suffix_attention_mask = ( + model._build_expert_pos_ids_and_attn_mask( + offset=torch.full( + (batch_size,), + prompt_len, + device=device, + dtype=torch.long, + ), + rope_deltas=rope_deltas, + kv_cache_seq_len=prompt_len, + n_diffusion_tokens=int(action_dims[0]), + b_star=batch_size, + device=device, + prefix_mask=sample["attention_mask"], + ) + ) + sample["prefix_k"] = prefix_k + sample["prefix_v"] = prefix_v + sample["suffix_position_ids"] = suffix_position_ids + sample["suffix_attention_mask"] = suffix_attention_mask + + action_args = ( + step_actions, + step_timestep, + prefix_k, + prefix_v, + suffix_position_ids, + suffix_attention_mask, + ) + action = ComponentBundle( + module=_export_module(model, device, dtype), + trace_args=action_args, + save_args=action_args, + input_names=[ + "noisy_action", + "timestep", + "prefix_k", + "prefix_v", + "position_ids", + "attention_mask", + ], + output_names=["velocity"], + model_type="action", + engine_file="action.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + "immutable_weights": False, + }, + ) + return {"vision": vision, "language": language_bundle, "action": action} + + def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: + visual, deepstack = call_engine( + engines["vision"], + "vision", + sample["pixel_values"], + sample["image_grid_thw"], + ) + inputs_embeds = scatter_image_tokens( + visual, + sample["lang_embeds"], + sample["image_token_mask"], + ) + + ds_template = sample["ds_template"] + ds_layers = [] + for layer_index in range(int(ds_template.shape[0])): + if layer_index < int(deepstack.shape[0]): + ds_layers.append( + scatter_image_tokens( + deepstack[layer_index], + ds_template[layer_index], + sample["image_token_mask"], + ) + ) + else: + ds_layers.append(ds_template[layer_index]) + ds_stack = torch.stack(ds_layers, dim=0) + + language = call_engine( + engines["language"], + "language", + inputs_embeds, + sample["rope_rotary_cos_sin"], + sample["context_lengths"], + sample["kvcache_start_index"], + sample["last_token_ids"], + ds_stack, + *kv_kwargs(sample), + ) + return call_engine( + engines["action"], + "action", + sample["step_actions"], + sample["step_timestep"], + language[2], + language[3], + sample["suffix_position_ids"], + sample["suffix_attention_mask"], + ) diff --git a/tools/hf/exporters/models/common/helpers.py b/tools/hf/exporters/models/common/helpers.py index 2eb153e7bf..eb6e8911eb 100644 --- a/tools/hf/exporters/models/common/helpers.py +++ b/tools/hf/exporters/models/common/helpers.py @@ -15,6 +15,7 @@ def causal_lm_flat( device: torch.device, dtype: torch.dtype, seq_len: int | None = None, + position_ids: torch.Tensor | None = None, ) -> tuple[tuple[torch.Tensor, ...], dict[str, Any]]: """inputs_embeds, rope, ctx, kv_start, last_token_ids, ds_stack, *kvs. @@ -32,7 +33,11 @@ def causal_lm_flat( from ...rope import make_rope_rotary_cos_sin rope = make_rope_rotary_cos_sin( - cfg, int(max_seq_len), device, language_model=language + cfg, + int(max_seq_len), + device, + language_model=language, + position_ids=position_ids, ) except ImportError: rope = torch.zeros(int(max_seq_len), 2, 1, head_dim, device=device, dtype=dtype) diff --git a/tools/hf/exporters/tests/test_edge_exporter.py b/tools/hf/exporters/tests/test_edge_exporter.py index 590ae602b2..11348f5f94 100644 --- a/tools/hf/exporters/tests/test_edge_exporter.py +++ b/tools/hf/exporters/tests/test_edge_exporter.py @@ -63,6 +63,9 @@ def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]): @pytest.mark.unit def test_builtin_specs_are_registered(): keys = registered_specs() + assert "alpamayo" in keys + assert "alpamayo_r1" in keys + assert "alpamayo1_5" in keys assert "pi05" in keys assert "groot" in keys assert "nemotron_h" in keys @@ -143,6 +146,58 @@ def test_pi05_backend_registers_vision_and_language(): assert any("PI05Pytorch.forward" in p for p in paths) +@pytest.mark.unit +def test_alpamayo_backend_registers_qwen_vision_language_and_action(): + from exporters.models.alpamayo.patches import ALPAMAYO + from exporters.plugin.attn_patches import _PATCHES + + paths = [path for path, _ in _PATCHES[ALPAMAYO]] + assert any("Qwen3VLVisionAttention.forward" in path for path in paths) + assert any("Qwen3VLTextAttention.forward" in path for path in paths) + assert any("Qwen3VLForConditionalGeneration.forward" in path for path in paths) + assert any("Alpamayo1_5.forward" in path for path in paths) + + +@pytest.mark.unit +def test_alpamayo_scatter_visual_tokens(): + from exporters.models.alpamayo.helpers import scatter_visual_tokens + + text = torch.zeros(1, 5, 3) + visual = torch.arange(6, dtype=torch.float32).reshape(2, 3) + mask = torch.tensor([[False, True, False, True, False]]) + result = scatter_visual_tokens(visual, text, mask) + + torch.testing.assert_close(result[0, 1], visual[0]) + torch.testing.assert_close(result[0, 3], visual[1]) + torch.testing.assert_close(result[0, [0, 2, 4]], torch.zeros(3, 3)) + + +@pytest.mark.unit +def test_alpamayo_generation_positions_extend_multimodal_prompt(): + from exporters.models.alpamayo.spec import _append_generation_positions + + positions = torch.tensor( + [ + [[0, 1, 2]], + [[0, 4, 5]], + [[0, 7, 8]], + ], + dtype=torch.long, + ) + result = _append_generation_positions( + positions, + torch.tensor([[10]], dtype=torch.long), + 5, + ) + + assert result.shape == (3, 1, 5) + torch.testing.assert_close(result[:, :, :3], positions) + torch.testing.assert_close( + result[:, :, 3:], + torch.tensor([[[13, 14]], [[13, 14]], [[13, 14]]]), + ) + + @pytest.mark.unit def test_paligemma_image_features_patch_returns_tensor(): from exporters.models.pi05.patches import ( diff --git a/tools/hf/run_export.py b/tools/hf/run_export.py index b62b75caea..c8092657f9 100644 --- a/tools/hf/run_export.py +++ b/tools/hf/run_export.py @@ -21,6 +21,7 @@ ] EXPORT_PREPARERS = { + "alpamayo": "exporters.models.alpamayo.export:prepare_export", "groot": "exporters.models.groot.export:prepare_export", "pi05": "exporters.models.pi05.export:prepare_export", "nemotron": "exporters.models.nemotron.export:prepare_export", @@ -42,6 +43,16 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--prompt", default="Hello, how are you?") parser.add_argument("--engine-dir") parser.add_argument("--max-seq-len", type=int) + parser.add_argument( + "--clip-id", + help="PhysicalAI clip ID used to prepare Alpamayo sample inputs.", + ) + parser.add_argument( + "--t0-us", + type=int, + default=5_100_000, + help="Timestamp within the Alpamayo sample clip (default: 5100000).", + ) parser.add_argument("--device") parser.add_argument( "--dtype", From ba280f34ae0171938678d4c990d0237432997637 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 22 Sep 2026 20:02:13 -0700 Subject: [PATCH 2/3] Use fixed-grid Alpamayo export wrappers --- tools/hf/exporters/models/alpamayo/helpers.py | 104 ++++++++++++++++ tools/hf/exporters/models/alpamayo/patches.py | 112 +++++------------- tools/hf/exporters/models/alpamayo/spec.py | 93 +++++++++------ .../hf/exporters/tests/test_edge_exporter.py | 35 +++++- 4 files changed, 227 insertions(+), 117 deletions(-) diff --git a/tools/hf/exporters/models/alpamayo/helpers.py b/tools/hf/exporters/models/alpamayo/helpers.py index d3c8956c8d..cb16961651 100644 --- a/tools/hf/exporters/models/alpamayo/helpers.py +++ b/tools/hf/exporters/models/alpamayo/helpers.py @@ -4,6 +4,9 @@ import torch import torch.nn as nn +import torch.nn.functional as F + +from ...prefix_cache import PrefixKVCache def alpamayo_vlm(model: nn.Module) -> nn.Module: @@ -105,3 +108,104 @@ def make_deepstack_tensor( image_token_mask, ) return dense + + +class VisualFixedGrid(nn.Module): + """Qwen3-VL visual tower with grid-dependent tensors baked as buffers.""" + + def __init__(self, visual: nn.Module, grid_thw: torch.Tensor): + super().__init__() + self.visual = visual.eval() + with torch.no_grad(): + pos_embeds = visual.fast_pos_embed_interpolate(grid_thw) + rotary = visual.rot_pos_emb(grid_thw) + seq_len = int(pos_embeds.shape[0]) + rotary = rotary.reshape(seq_len, -1) + rotary = torch.cat((rotary, rotary), dim=-1) + lengths = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], + grid_thw[:, 0], + ) + cu_seqlens = F.pad( + lengths.cumsum(dim=0, dtype=torch.int32), + (1, 0), + value=0, + ) + static_lengths = [int(value) for value in lengths.cpu().tolist()] + for block in visual.blocks: + block.attn._static_lengths = static_lengths + + self.register_buffer("pos_embeds", pos_embeds, persistent=False) + self.register_buffer("cos", rotary.cos(), persistent=False) + self.register_buffer("sin", rotary.sin(), persistent=False) + self.register_buffer("cu_seqlens", cu_seqlens, persistent=False) + + def forward(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + hidden = self.visual.patch_embed(pixel_values) + hidden = hidden + self.pos_embeds.to(hidden.dtype) + position_embeddings = ( + self.cos.to(hidden.dtype), + self.sin.to(hidden.dtype), + ) + deepstack = [] + for layer_index, block in enumerate(self.visual.blocks): + hidden = block( + hidden, + cu_seqlens=self.cu_seqlens, + position_embeddings=position_embeddings, + ) + if layer_index in self.visual.deepstack_visual_indexes: + merger_index = self.visual.deepstack_visual_indexes.index(layer_index) + deepstack.append( + self.visual.deepstack_merger_list[merger_index](hidden) + ) + return self.visual.merger(hidden), torch.stack(tuple(deepstack), dim=0) + + +class StaticKVDiffusionStepModule(nn.Module): + """Fused Alpamayo action projection, expert, and output projection.""" + + def __init__( + self, + action_in_proj: nn.Module, + expert: nn.Module, + action_out_proj: nn.Module, + action_space_dims: tuple[int, ...], + ): + super().__init__() + self.action_in_proj = action_in_proj + self.expert = expert + self.action_out_proj = action_out_proj + self.action_space_dims = action_space_dims + self.n_diffusion_tokens = int(action_space_dims[0]) + + def forward( + self, + noisy_action: torch.Tensor, + timestep: torch.Tensor, + prefix_k: torch.Tensor, + prefix_v: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + batch_size = noisy_action.shape[0] + action_embeds = self.action_in_proj(noisy_action, timestep) + if action_embeds.dim() == 2: + action_embeds = action_embeds.view( + batch_size, + self.n_diffusion_tokens, + -1, + ) + expert = self.expert( + inputs_embeds=action_embeds, + position_ids=position_ids, + past_key_values=PrefixKVCache(prefix_k, prefix_v), + attention_mask=attention_mask, + use_cache=False, + return_dict=True, + ) + hidden = expert.last_hidden_state[:, -self.n_diffusion_tokens :] + return self.action_out_proj(hidden).reshape( + batch_size, + *self.action_space_dims, + ) diff --git a/tools/hf/exporters/models/alpamayo/patches.py b/tools/hf/exporters/models/alpamayo/patches.py index 0b160b5842..83fab62638 100644 --- a/tools/hf/exporters/models/alpamayo/patches.py +++ b/tools/hf/exporters/models/alpamayo/patches.py @@ -10,7 +10,6 @@ _patch_language_attention, register_patch, ) -from ...prefix_cache import PrefixKVCache from ..common.patches import causal_lm_plugin_forward ALPAMAYO = "alpamayo" @@ -21,7 +20,7 @@ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionAttention.forward", ) def _patch_qwen3_vl_vision_attention(original: Callable) -> Callable: - """Route Qwen3-VL vision attention through the Edge ViT plugin.""" + """Use static sequence splits baked by ``VisualFixedGrid``.""" def forward( self, @@ -31,17 +30,20 @@ def forward( position_embeddings=None, **kwargs: Any, ) -> torch.Tensor: - del rotary_pos_emb, kwargs - if position_embeddings is None: + static_lengths = getattr(self, "_static_lengths", None) + if static_lengths is None or position_embeddings is None: return original( self, hidden_states, cu_seqlens, + rotary_pos_emb=rotary_pos_emb, position_embeddings=position_embeddings, + **kwargs, ) from transformers.models.qwen3_vl.modeling_qwen3_vl import ( apply_rotary_pos_emb_vision, + eager_attention_forward, ) seq_len = int(hidden_states.shape[0]) @@ -53,50 +55,35 @@ def forward( ) cos, sin = position_embeddings q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) - q = q.to(torch.float16).contiguous() - k = k.to(torch.float16).contiguous() - v = v.to(torch.float16).contiguous() - - # The carrier's length communicates a safe maximum sequence length. - max_seqlen_carrier = torch.zeros( - hidden_states.shape[0], - device=hidden_states.device, - dtype=torch.int32, - ) - output = torch.ops.trt.vit_attention_plugin.default( - q, - k, - v, - cu_seqlens.to(torch.int32), - max_seqlen_carrier, - int(self.num_heads), - int(q.shape[-1]), + q = q.transpose(0, 1).unsqueeze(0) + k = k.transpose(0, 1).unsqueeze(0) + v = v.transpose(0, 1).unsqueeze(0) + splits = [torch.split(tensor, static_lengths, dim=2) for tensor in (q, k, v)] + outputs = [ + eager_attention_forward( + self, + q_part, + k_part, + v_part, + attention_mask=None, + scaling=self.scaling, + dropout=0.0, + is_causal=False, + **kwargs, + )[0] + for q_part, k_part, v_part in zip(*splits) + ] + output = ( + torch.cat(outputs, dim=1) + .reshape(seq_len, -1) + .contiguous() + .to(hidden_states.dtype) ) - output = output.reshape(seq_len, -1).to(hidden_states.dtype) return self.proj(output) return forward -@register_patch( - ALPAMAYO, - "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionModel.forward", -) -def _patch_qwen3_vl_vision_model(original: Callable) -> Callable: - """Return a tensor rather than a Python list for deepstack outputs.""" - - def forward(self, hidden_states, grid_thw, **kwargs: Any): - visual, deepstack = original( - self, - hidden_states, - grid_thw, - **kwargs, - ) - return visual, torch.stack(tuple(deepstack), dim=0) - - return forward - - register_patch( ALPAMAYO, "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextAttention.forward", @@ -142,44 +129,3 @@ def forward( ) return forward - - -@register_patch( - ALPAMAYO, - "alpamayo1_5.models.alpamayo1_5.Alpamayo1_5.forward", -) -def _patch_alpamayo_action_step(original: Callable) -> Callable: - """Compile one Alpamayo diffusion velocity step with prefix KV tensors.""" - - def forward( - self, - noisy_action, - timestep=None, - prefix_k=None, - prefix_v=None, - position_ids=None, - attention_mask=None, - *args, - **kwargs: Any, - ): - if prefix_k is None or getattr(prefix_k, "ndim", 0) != 5: - return original(self, noisy_action, timestep, *args, **kwargs) - - action_embeds = self.action_in_proj(noisy_action, timestep) - expert_kwargs: dict[str, Any] = {} - if self.config.expert_non_causal_attention: - expert_kwargs["is_causal"] = False - expert = self.expert( - inputs_embeds=action_embeds, - position_ids=position_ids, - past_key_values=PrefixKVCache(prefix_k, prefix_v), - attention_mask=attention_mask, - use_cache=False, - return_dict=True, - **expert_kwargs, - ) - hidden = expert.last_hidden_state - hidden = hidden[:, -int(noisy_action.shape[1]) :] - return self.action_out_proj(hidden).reshape_as(noisy_action) - - return forward diff --git a/tools/hf/exporters/models/alpamayo/spec.py b/tools/hf/exporters/models/alpamayo/spec.py index afca679165..cf1e08bb6c 100644 --- a/tools/hf/exporters/models/alpamayo/spec.py +++ b/tools/hf/exporters/models/alpamayo/spec.py @@ -7,7 +7,6 @@ import torch.nn as nn from ...ops import call_engine, scatter_image_tokens -from ...prefix_cache import PrefixKVCache from ...spec import ComponentBundle, EdgeSpec, register_edge_spec from ..common.helpers import ( causal_lm_flat, @@ -15,6 +14,8 @@ split_flat_to_kwargs, ) from .helpers import ( + StaticKVDiffusionStepModule, + VisualFixedGrid, alpamayo_language, alpamayo_visual, alpamayo_vlm, @@ -162,25 +163,14 @@ def capture_eager_outputs( use_cache=False, return_dict=True, ) - action_embeds = model.action_in_proj( + velocity = sample["action_module"]( sample["step_actions"], sample["step_timestep"], + sample["prefix_k"], + sample["prefix_v"], + sample["suffix_position_ids"], + sample["suffix_attention_mask"], ) - expert = model.expert( - inputs_embeds=action_embeds, - position_ids=sample["suffix_position_ids"], - past_key_values=PrefixKVCache( - sample["prefix_k"], - sample["prefix_v"], - ), - attention_mask=sample["suffix_attention_mask"], - use_cache=False, - return_dict=True, - is_causal=not bool(model.config.expert_non_causal_attention), - ) - velocity = model.action_out_proj( - expert.last_hidden_state[:, -int(sample["step_actions"].shape[1]) :] - ).reshape_as(sample["step_actions"]) if bench is not None: bench["vision"] = cuda_ms(lambda: visual(px, grid)[0]) @@ -220,11 +210,17 @@ def prepare( device = px.device dtype = px.dtype + visual.config.attn_implementation = "sdpa" + visual.config._attn_implementation = "sdpa" + fixed_visual = VisualFixedGrid(visual, grid).to( + device=device, + dtype=dtype, + ) vision = ComponentBundle( - module=_export_module(visual, device, dtype), - trace_args=(px, grid), - save_args=(px, grid), - input_names=["pixel_values", "image_grid_thw"], + module=fixed_visual.eval(), + trace_args=(px,), + save_args=(px,), + input_names=["pixel_values"], output_names=["visual_embeds", "deepstack_visual_embeds"], parity_output="visual_embeds", model_type="qwen3_vl", @@ -253,12 +249,24 @@ def prepare( sample["image_token_mask"] = image_token_mask sample["inputs_embeds"] = inputs_embeds - position_ids, rope_deltas = vlm_core.get_rope_index( - sample["input_ids"], - sample["image_grid_thw"], - None, - attention_mask=sample["attention_mask"], - ) + try: + position_ids, rope_deltas = vlm_core.get_rope_index( + sample["input_ids"], + sample["image_grid_thw"], + None, + attention_mask=sample["attention_mask"], + ) + except (TypeError, IndexError): + image_token_types = ( + sample["input_ids"] == int(vlm_core.config.image_token_id) + ).to(torch.int32) + position_ids, rope_deltas = vlm_core.get_rope_index( + sample["input_ids"], + mm_token_type_ids=image_token_types, + image_grid_thw=sample["image_grid_thw"], + video_grid_thw=None, + attention_mask=sample["attention_mask"], + ) sample["position_ids"] = position_ids sample["rope_deltas"] = rope_deltas @@ -343,16 +351,24 @@ def prepare( sample["step_timestep"] = step_timestep language_cfg = language.config - num_kv_heads = int(language_cfg.num_key_value_heads) + expert_cfg = model.expert.config + for field in ("num_hidden_layers", "num_key_value_heads", "head_dim"): + if int(getattr(language_cfg, field)) != int(getattr(expert_cfg, field)): + raise ValueError( + "Alpamayo language/expert KV layouts differ at " + f"{field}: {getattr(language_cfg, field)} vs " + f"{getattr(expert_cfg, field)}" + ) + num_kv_heads = int(expert_cfg.num_key_value_heads) head_dim = int( getattr( - language_cfg, + expert_cfg, "head_dim", - language_cfg.hidden_size // language_cfg.num_attention_heads, + expert_cfg.hidden_size // expert_cfg.num_attention_heads, ) ) prefix_k = torch.zeros( - len(decoder_layers), + int(expert_cfg.num_hidden_layers), batch_size, num_kv_heads, prompt_len, @@ -382,6 +398,18 @@ def prepare( sample["suffix_position_ids"] = suffix_position_ids sample["suffix_attention_mask"] = suffix_attention_mask + model.expert.config._attn_implementation = "sdpa" + action_module = ( + StaticKVDiffusionStepModule( + model.action_in_proj, + model.expert, + model.action_out_proj, + action_dims, + ) + .to(device=device, dtype=dtype) + .eval() + ) + sample["action_module"] = action_module action_args = ( step_actions, step_timestep, @@ -391,7 +419,7 @@ def prepare( suffix_attention_mask, ) action = ComponentBundle( - module=_export_module(model, device, dtype), + module=action_module, trace_args=action_args, save_args=action_args, input_names=[ @@ -420,7 +448,6 @@ def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: engines["vision"], "vision", sample["pixel_values"], - sample["image_grid_thw"], ) inputs_embeds = scatter_image_tokens( visual, diff --git a/tools/hf/exporters/tests/test_edge_exporter.py b/tools/hf/exporters/tests/test_edge_exporter.py index 11348f5f94..b0bfea76c4 100644 --- a/tools/hf/exporters/tests/test_edge_exporter.py +++ b/tools/hf/exporters/tests/test_edge_exporter.py @@ -155,7 +155,40 @@ def test_alpamayo_backend_registers_qwen_vision_language_and_action(): assert any("Qwen3VLVisionAttention.forward" in path for path in paths) assert any("Qwen3VLTextAttention.forward" in path for path in paths) assert any("Qwen3VLForConditionalGeneration.forward" in path for path in paths) - assert any("Alpamayo1_5.forward" in path for path in paths) + + +@pytest.mark.unit +def test_alpamayo_action_module_forward_shape(): + from types import SimpleNamespace + + from exporters.models.alpamayo.helpers import StaticKVDiffusionStepModule + + class ActionIn(nn.Module): + def forward(self, actions, timestep): + return actions + timestep + + class Expert(nn.Module): + def forward(self, inputs_embeds, **kwargs): + del kwargs + return SimpleNamespace(last_hidden_state=inputs_embeds) + + projection = nn.Linear(2, 2, bias=False) + projection.weight.data.copy_(torch.eye(2)) + module = StaticKVDiffusionStepModule( + ActionIn(), + Expert(), + projection, + (4, 2), + ) + output = module( + torch.zeros(1, 4, 2), + torch.ones(1, 1, 1), + torch.zeros(1, 1, 1, 3, 2), + torch.zeros(1, 1, 1, 3, 2), + torch.zeros(3, 1, 4, dtype=torch.long), + torch.zeros(1, 1, 4, 7), + ) + assert output.shape == (1, 4, 2) @pytest.mark.unit From c7df8a9447c074ed1e33dcae7caa1d7f640610a8 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 22 Sep 2026 20:18:49 -0700 Subject: [PATCH 3/3] Use setattr patches for Alpamayo export --- tools/hf/exporters/models/alpamayo/helpers.py | 132 +++++------------- tools/hf/exporters/models/alpamayo/patches.py | 88 +++++++++++- tools/hf/exporters/models/alpamayo/spec.py | 48 +++---- .../hf/exporters/tests/test_edge_exporter.py | 34 +++-- 4 files changed, 166 insertions(+), 136 deletions(-) diff --git a/tools/hf/exporters/models/alpamayo/helpers.py b/tools/hf/exporters/models/alpamayo/helpers.py index cb16961651..676f5f2e9c 100644 --- a/tools/hf/exporters/models/alpamayo/helpers.py +++ b/tools/hf/exporters/models/alpamayo/helpers.py @@ -6,8 +6,6 @@ import torch.nn as nn import torch.nn.functional as F -from ...prefix_cache import PrefixKVCache - def alpamayo_vlm(model: nn.Module) -> nn.Module: """Return the Qwen3-VL conditional-generation module.""" @@ -110,102 +108,38 @@ def make_deepstack_tensor( return dense -class VisualFixedGrid(nn.Module): - """Qwen3-VL visual tower with grid-dependent tensors baked as buffers.""" - - def __init__(self, visual: nn.Module, grid_thw: torch.Tensor): - super().__init__() - self.visual = visual.eval() - with torch.no_grad(): - pos_embeds = visual.fast_pos_embed_interpolate(grid_thw) - rotary = visual.rot_pos_emb(grid_thw) - seq_len = int(pos_embeds.shape[0]) - rotary = rotary.reshape(seq_len, -1) - rotary = torch.cat((rotary, rotary), dim=-1) - lengths = torch.repeat_interleave( - grid_thw[:, 1] * grid_thw[:, 2], - grid_thw[:, 0], - ) - cu_seqlens = F.pad( - lengths.cumsum(dim=0, dtype=torch.int32), - (1, 0), - value=0, - ) - static_lengths = [int(value) for value in lengths.cpu().tolist()] - for block in visual.blocks: - block.attn._static_lengths = static_lengths - - self.register_buffer("pos_embeds", pos_embeds, persistent=False) - self.register_buffer("cos", rotary.cos(), persistent=False) - self.register_buffer("sin", rotary.sin(), persistent=False) - self.register_buffer("cu_seqlens", cu_seqlens, persistent=False) - - def forward(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - hidden = self.visual.patch_embed(pixel_values) - hidden = hidden + self.pos_embeds.to(hidden.dtype) - position_embeddings = ( - self.cos.to(hidden.dtype), - self.sin.to(hidden.dtype), - ) - deepstack = [] - for layer_index, block in enumerate(self.visual.blocks): - hidden = block( - hidden, - cu_seqlens=self.cu_seqlens, - position_embeddings=position_embeddings, - ) - if layer_index in self.visual.deepstack_visual_indexes: - merger_index = self.visual.deepstack_visual_indexes.index(layer_index) - deepstack.append( - self.visual.deepstack_merger_list[merger_index](hidden) - ) - return self.visual.merger(hidden), torch.stack(tuple(deepstack), dim=0) - - -class StaticKVDiffusionStepModule(nn.Module): - """Fused Alpamayo action projection, expert, and output projection.""" - - def __init__( - self, - action_in_proj: nn.Module, - expert: nn.Module, - action_out_proj: nn.Module, - action_space_dims: tuple[int, ...], - ): - super().__init__() - self.action_in_proj = action_in_proj - self.expert = expert - self.action_out_proj = action_out_proj - self.action_space_dims = action_space_dims - self.n_diffusion_tokens = int(action_space_dims[0]) - - def forward( - self, - noisy_action: torch.Tensor, - timestep: torch.Tensor, - prefix_k: torch.Tensor, - prefix_v: torch.Tensor, - position_ids: torch.Tensor, - attention_mask: torch.Tensor, - ) -> torch.Tensor: - batch_size = noisy_action.shape[0] - action_embeds = self.action_in_proj(noisy_action, timestep) - if action_embeds.dim() == 2: - action_embeds = action_embeds.view( - batch_size, - self.n_diffusion_tokens, - -1, - ) - expert = self.expert( - inputs_embeds=action_embeds, - position_ids=position_ids, - past_key_values=PrefixKVCache(prefix_k, prefix_v), - attention_mask=attention_mask, - use_cache=False, - return_dict=True, +def prepare_fixed_grid_vision( + visual: nn.Module, + grid_thw: torch.Tensor, +) -> None: + """Attach fixed-grid tensors consumed by the temporary vision patch.""" + with torch.no_grad(): + pos_embeds = visual.fast_pos_embed_interpolate(grid_thw) + rotary = visual.rot_pos_emb(grid_thw) + seq_len = int(pos_embeds.shape[0]) + rotary = rotary.reshape(seq_len, -1) + rotary = torch.cat((rotary, rotary), dim=-1) + lengths = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], + grid_thw[:, 0], ) - hidden = expert.last_hidden_state[:, -self.n_diffusion_tokens :] - return self.action_out_proj(hidden).reshape( - batch_size, - *self.action_space_dims, + cu_seqlens = F.pad( + lengths.cumsum(dim=0, dtype=torch.int32), + (1, 0), + value=0, ) + static_lengths = [int(value) for value in lengths.cpu().tolist()] + + fixed_buffers = { + "_edge_pos_embeds": pos_embeds, + "_edge_cos": rotary.cos(), + "_edge_sin": rotary.sin(), + "_edge_cu_seqlens": cu_seqlens, + } + for name, value in fixed_buffers.items(): + if name in visual._buffers: + visual._buffers[name] = value + else: + visual.register_buffer(name, value, persistent=False) + for block in visual.blocks: + block.attn._edge_static_lengths = static_lengths diff --git a/tools/hf/exporters/models/alpamayo/patches.py b/tools/hf/exporters/models/alpamayo/patches.py index 83fab62638..fb06df92fb 100644 --- a/tools/hf/exporters/models/alpamayo/patches.py +++ b/tools/hf/exporters/models/alpamayo/patches.py @@ -10,6 +10,7 @@ _patch_language_attention, register_patch, ) +from ...prefix_cache import PrefixKVCache from ..common.patches import causal_lm_plugin_forward ALPAMAYO = "alpamayo" @@ -20,7 +21,7 @@ "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionAttention.forward", ) def _patch_qwen3_vl_vision_attention(original: Callable) -> Callable: - """Use static sequence splits baked by ``VisualFixedGrid``.""" + """Use static sequence splits prepared on the original vision instance.""" def forward( self, @@ -30,7 +31,7 @@ def forward( position_embeddings=None, **kwargs: Any, ) -> torch.Tensor: - static_lengths = getattr(self, "_static_lengths", None) + static_lengths = getattr(self, "_edge_static_lengths", None) if static_lengths is None or position_embeddings is None: return original( self, @@ -84,6 +85,42 @@ def forward( return forward +@register_patch( + ALPAMAYO, + "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionModel.forward", +) +def _patch_qwen3_vl_vision_model(original: Callable) -> Callable: + """Run Qwen3-VL vision with grid-dependent values prepared on the instance.""" + + def forward(self, hidden_states, grid_thw=None, **kwargs: Any): + if not hasattr(self, "_edge_pos_embeds"): + return original(self, hidden_states, grid_thw, **kwargs) + + del grid_thw + hidden_states = self.patch_embed(hidden_states) + hidden_states = hidden_states + self._edge_pos_embeds.to(hidden_states.dtype) + position_embeddings = ( + self._edge_cos.to(hidden_states.dtype), + self._edge_sin.to(hidden_states.dtype), + ) + deepstack = [] + for layer_index, block in enumerate(self.blocks): + hidden_states = block( + hidden_states, + cu_seqlens=self._edge_cu_seqlens, + position_embeddings=position_embeddings, + **kwargs, + ) + if layer_index in self.deepstack_visual_indexes: + merger_index = self.deepstack_visual_indexes.index(layer_index) + deepstack.append( + self.deepstack_merger_list[merger_index](hidden_states) + ) + return self.merger(hidden_states), torch.stack(tuple(deepstack), dim=0) + + return forward + + register_patch( ALPAMAYO, "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextAttention.forward", @@ -129,3 +166,50 @@ def forward( ) return forward + + +@register_patch( + ALPAMAYO, + "alpamayo1_5.models.alpamayo1_5.Alpamayo1_5.forward", +) +def _patch_alpamayo_action_step(original: Callable) -> Callable: + """Run one diffusion step when explicit stacked prefix KV is supplied.""" + + def forward( + self, + noisy_action, + timestep=None, + prefix_k=None, + prefix_v=None, + position_ids=None, + attention_mask=None, + *args, + **kwargs: Any, + ): + if prefix_k is None or getattr(prefix_k, "ndim", 0) != 5: + return original(self, noisy_action, timestep, *args, **kwargs) + + n_diffusion_tokens = int(noisy_action.shape[1]) + action_embeds = self.action_in_proj(noisy_action, timestep) + if action_embeds.dim() == 2: + action_embeds = action_embeds.view( + noisy_action.shape[0], + n_diffusion_tokens, + -1, + ) + expert_kwargs: dict[str, Any] = {} + if self.config.expert_non_causal_attention: + expert_kwargs["is_causal"] = False + expert = self.expert( + inputs_embeds=action_embeds, + position_ids=position_ids, + past_key_values=PrefixKVCache(prefix_k, prefix_v), + attention_mask=attention_mask, + use_cache=False, + return_dict=True, + **expert_kwargs, + ) + hidden = expert.last_hidden_state[:, -n_diffusion_tokens:] + return self.action_out_proj(hidden).reshape_as(noisy_action) + + return forward diff --git a/tools/hf/exporters/models/alpamayo/spec.py b/tools/hf/exporters/models/alpamayo/spec.py index cf1e08bb6c..d8f1b5f8d5 100644 --- a/tools/hf/exporters/models/alpamayo/spec.py +++ b/tools/hf/exporters/models/alpamayo/spec.py @@ -7,6 +7,7 @@ import torch.nn as nn from ...ops import call_engine, scatter_image_tokens +from ...prefix_cache import PrefixKVCache from ...spec import ComponentBundle, EdgeSpec, register_edge_spec from ..common.helpers import ( causal_lm_flat, @@ -14,13 +15,12 @@ split_flat_to_kwargs, ) from .helpers import ( - StaticKVDiffusionStepModule, - VisualFixedGrid, alpamayo_language, alpamayo_visual, alpamayo_vlm, alpamayo_vlm_core, make_deepstack_tensor, + prepare_fixed_grid_vision, scatter_visual_tokens, stack_deepstack_features, ) @@ -163,14 +163,28 @@ def capture_eager_outputs( use_cache=False, return_dict=True, ) - velocity = sample["action_module"]( + action_embeds = model.action_in_proj( sample["step_actions"], sample["step_timestep"], - sample["prefix_k"], - sample["prefix_v"], - sample["suffix_position_ids"], - sample["suffix_attention_mask"], ) + expert_kwargs: dict[str, Any] = {} + if model.config.expert_non_causal_attention: + expert_kwargs["is_causal"] = False + expert = model.expert( + inputs_embeds=action_embeds, + position_ids=sample["suffix_position_ids"], + past_key_values=PrefixKVCache( + sample["prefix_k"], + sample["prefix_v"], + ), + attention_mask=sample["suffix_attention_mask"], + use_cache=False, + return_dict=True, + **expert_kwargs, + ) + velocity = model.action_out_proj( + expert.last_hidden_state[:, -int(sample["step_actions"].shape[1]) :] + ).reshape_as(sample["step_actions"]) if bench is not None: bench["vision"] = cuda_ms(lambda: visual(px, grid)[0]) @@ -212,12 +226,9 @@ def prepare( visual.config.attn_implementation = "sdpa" visual.config._attn_implementation = "sdpa" - fixed_visual = VisualFixedGrid(visual, grid).to( - device=device, - dtype=dtype, - ) + prepare_fixed_grid_vision(visual, grid) vision = ComponentBundle( - module=fixed_visual.eval(), + module=_export_module(visual, device, dtype), trace_args=(px,), save_args=(px,), input_names=["pixel_values"], @@ -399,17 +410,6 @@ def prepare( sample["suffix_attention_mask"] = suffix_attention_mask model.expert.config._attn_implementation = "sdpa" - action_module = ( - StaticKVDiffusionStepModule( - model.action_in_proj, - model.expert, - model.action_out_proj, - action_dims, - ) - .to(device=device, dtype=dtype) - .eval() - ) - sample["action_module"] = action_module action_args = ( step_actions, step_timestep, @@ -419,7 +419,7 @@ def prepare( suffix_attention_mask, ) action = ComponentBundle( - module=action_module, + module=_export_module(model, device, dtype), trace_args=action_args, save_args=action_args, input_names=[ diff --git a/tools/hf/exporters/tests/test_edge_exporter.py b/tools/hf/exporters/tests/test_edge_exporter.py index b0bfea76c4..fba7510dfa 100644 --- a/tools/hf/exporters/tests/test_edge_exporter.py +++ b/tools/hf/exporters/tests/test_edge_exporter.py @@ -153,15 +153,17 @@ def test_alpamayo_backend_registers_qwen_vision_language_and_action(): paths = [path for path, _ in _PATCHES[ALPAMAYO]] assert any("Qwen3VLVisionAttention.forward" in path for path in paths) + assert any("Qwen3VLVisionModel.forward" in path for path in paths) assert any("Qwen3VLTextAttention.forward" in path for path in paths) assert any("Qwen3VLForConditionalGeneration.forward" in path for path in paths) + assert any("Alpamayo1_5.forward" in path for path in paths) @pytest.mark.unit -def test_alpamayo_action_module_forward_shape(): +def test_alpamayo_action_patch_uses_explicit_prefix_kv(): from types import SimpleNamespace - from exporters.models.alpamayo.helpers import StaticKVDiffusionStepModule + from exporters.models.alpamayo.patches import _patch_alpamayo_action_step class ActionIn(nn.Module): def forward(self, actions, timestep): @@ -172,15 +174,25 @@ def forward(self, inputs_embeds, **kwargs): del kwargs return SimpleNamespace(last_hidden_state=inputs_embeds) - projection = nn.Linear(2, 2, bias=False) - projection.weight.data.copy_(torch.eye(2)) - module = StaticKVDiffusionStepModule( - ActionIn(), - Expert(), - projection, - (4, 2), - ) - output = module( + class Alpamayo(nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(expert_non_causal_attention=True) + self.action_in_proj = ActionIn() + self.expert = Expert() + self.action_out_proj = nn.Linear(2, 2, bias=False) + self.action_out_proj.weight.data.copy_(torch.eye(2)) + + def forward(self, noisy_action, timestep=None, *args, **kwargs): + del timestep, args, kwargs + return noisy_action - 1 + + Alpamayo.forward = _patch_alpamayo_action_step(Alpamayo.forward) + model = Alpamayo() + actions = torch.zeros(1, 4, 2) + torch.testing.assert_close(model(actions), actions - 1) + + output = model( torch.zeros(1, 4, 2), torch.ones(1, 1, 1), torch.zeros(1, 1, 1, 3, 2),