feat(models): add Phi-4 multimodal vision support

areno/accel/ops.py

@@ -12,11 +12,8 @@
 
 from __future__ import annotations
 
-import logging
 from typing import Any
 
-import torch
-
 from areno.accel.activations import areno_gelu_tanh_and_mul, areno_silu_and_mul
 from areno.accel.attention import (
     areno_causal_attention,
@@ -28,46 +25,7 @@
 from areno.accel.kernels.fused_moe import is_available as fused_moe_is_available
 from areno.accel.kernels.group_rmsnorm import rms_norm_gate_fwd
 from areno.accel.kernels.seg_la import SegLaMeta, seg_la_fwd
-
-logger = logging.getLogger(__name__)
-# Process-wide set of message keys already emitted by log_once/warn_once.
-_LOGGED: set[str] = set()
-
-
-def log_once(key: str, message: str, *, level: int = logging.DEBUG) -> None:
-    """Log ``message`` at most once per process for the given ``key``."""
-
-    if key in _LOGGED:
-        return
-    logger.log(level, message)
-    _LOGGED.add(key)
-
-
-def warn_once(key: str, message: str) -> None:
-    """Emit a warning at most once per process for the given ``key``."""
-
-    log_once(key, message, level=logging.WARNING)
-
-
-@torch._dynamo.disable
-def is_cuda_graph_capturing(tensor: torch.Tensor) -> bool:
-    """True if the tensor lives on CUDA and we are inside a graph capture."""
-
-    return tensor.is_cuda and torch.cuda.is_current_stream_capturing()
-
-
-@torch._dynamo.disable
-def can_use_cuda_kernel(tensor: torch.Tensor, name: str, *, allow_sm121: bool = False) -> bool:
-    """Decide whether to take the fused kernel path for ``tensor``.
-
-    Returns False only on non-CUDA tensors. ``name`` and ``allow_sm121`` are
-    kept for compatibility with existing call sites.
-    """
-
-    if not tensor.is_cuda:
-        return False
-    return True
-
+from areno.accel.utils import can_use_cuda_kernel, is_cuda_graph_capturing, log_once, warn_once
 
 __all__ = [
     "Any",

areno/accel/utils.py

@@ -0,0 +1,41 @@
+"""Lightweight acceleration helpers that do not import optional kernels."""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+_LOGGED: set[str] = set()
+
+
+def log_once(key: str, message: str, *, level: int = logging.DEBUG) -> None:
+    """Log ``message`` at most once per process for the given ``key``."""
+
+    if key in _LOGGED:
+        return
+    logger.log(level, message)
+    _LOGGED.add(key)
+
+
+def warn_once(key: str, message: str) -> None:
+    """Emit a warning at most once per process for the given ``key``."""
+
+    log_once(key, message, level=logging.WARNING)
+
+
+@torch._dynamo.disable
+def is_cuda_graph_capturing(tensor: torch.Tensor) -> bool:
+    """True if the tensor lives on CUDA and we are inside a graph capture."""
+
+    return tensor.is_cuda and torch.cuda.is_current_stream_capturing()
+
+
+@torch._dynamo.disable
+def can_use_cuda_kernel(tensor: torch.Tensor, name: str, *, allow_sm121: bool = False) -> bool:
+    """Return whether a fused CUDA kernel can run for ``tensor``."""
+
+    if not tensor.is_cuda:
+        return False
+    return True

areno/api/multimodal.py

@@ -430,7 +430,7 @@ def _image_processor_from_processor(processor: Any):
 
 def _image_token_id(tokenizer: Any, processor: Any) -> int | None:
     for obj in (processor, tokenizer):
-        for attr in ("image_token_id", "image_token_index"):
+        for attr in ("image_token_id", "image_token_index", "special_image_token_id"):
             value = getattr(obj, attr, None)
             if isinstance(value, int):
                 return int(value)
@@ -443,7 +443,7 @@ def _image_token_id(tokenizer: Any, processor: Any) -> int | None:
                     return int(token_id)
     convert = getattr(tokenizer, "convert_tokens_to_ids", None)
     if callable(convert):
-        for token in ("<|image_pad|>", "<|image|>", "<image>"):
+        for token in ("<|image_pad|>", "<|image|>", "<image>", "<|endoftext10|>"):
             token_id = convert(token)
             if isinstance(token_id, int) and token_id >= 0:
                 return int(token_id)

areno/engine/checkpoints/common.py

@@ -98,6 +98,16 @@ class MergedColumnSpec:
     keys: tuple[str, ...]
 
 
+@dataclass(frozen=True, slots=True)
+class PackedSectionColumnSpec:
+    """One HF tensor whose semantic row sections are TP-sharded separately."""
+
+    key: str
+    tensor_attr: str
+    global_sizes_attr: str
+    local_sizes_attr: str
+
+
 @dataclass(frozen=True, slots=True)
 class KSharedQKVColumnSpec:
     """QKV load spec for checkpoints where later layers may share K/V."""
@@ -403,6 +413,8 @@ def save_checkpoint_weights(
     source_path: str | None,
     spec: CheckpointSpec,
     extra_tensors_fn: Callable[[CheckpointTensorStore], None] | None = None,
+    *,
+    copy_passthrough: bool = True,
 ) -> str | None:
     """Save a tensor-parallel model as a HF sharded safetensors checkpoint."""
 
@@ -430,7 +442,7 @@ def save_checkpoint_weights(
             writer.write(tensors, "extra-tensors")
             tensors.clear()
     saved_path = writer.finish()
-    if saved_path is not None and source_path is not None:
+    if copy_passthrough and saved_path is not None and source_path is not None:
         copy_source_passthrough_weights(
             source_path, saved_path, protected_prefix=_protected_prefix_from_top_level(spec.top_level)
         )
@@ -599,6 +611,9 @@ def load_layer_op(
     if isinstance(op, MergedColumnSpec):
         load_merged_column_spec(module, index, prefix, op, rank, world_size)
         return
+    if isinstance(op, PackedSectionColumnSpec):
+        load_packed_section_column_spec(module, index, prefix, op, rank, world_size)
+        return
     if isinstance(op, KSharedQKVColumnSpec):
         load_k_shared_qkv_column_spec(module, index, prefix, op, rank, world_size)
         return
@@ -642,6 +657,9 @@ def save_layer_op(
     if isinstance(op, SplitColumnSpec):
         save_split_column_spec(tensors, module, prefix, op)
         return
+    if isinstance(op, PackedSectionColumnSpec):
+        save_packed_section_column_spec(tensors, module, prefix, op)
+        return
     if isinstance(op, RangedSplitColumnSpec):
         save_ranged_split_column_spec(tensors, module, prefix, op)
         return
@@ -751,6 +769,47 @@ def load_merged_column_spec(
     copy_merged_column_from_index(dst, index, tensor_keys, rank, world_size)
 
 
+def load_packed_section_column_spec(
+    module: nn.Module,
+    index: SafetensorsIndex,
+    prefix: str,
+    spec: PackedSectionColumnSpec,
+    rank: int,
+    world_size: int,
+) -> None:
+    """Shard each row section of one packed HF tensor independently."""
+
+    dst = attr_path(module, spec.tensor_attr)
+    global_sizes = tuple(int(size) for size in attr_path(module, spec.global_sizes_attr))
+    local_sizes = tuple(int(size) for size in attr_path(module, spec.local_sizes_attr))
+    if len(global_sizes) != len(local_sizes):
+        raise ValueError("packed-section global and local size counts differ")
+    ranges = tuple(_shard_range(size, rank, world_size) for size in global_sizes)
+    expected_local_sizes = tuple(end - start for start, end in ranges)
+    if local_sizes != expected_local_sizes:
+        raise ValueError(f"packed-section local sizes {local_sizes} do not match TP shard sizes {expected_local_sizes}")
+    if dst.shape[0] != sum(local_sizes):
+        raise ValueError(f"packed-section destination has {dst.shape[0]} rows, expected {sum(local_sizes)}")
+
+    tensor_key = key(spec.key, prefix)
+    filename = index.weight_map.get(tensor_key)
+    if filename is None:
+        raise KeyError(f"missing HF weight {tensor_key}")
+    with safe_open(index.model_path / filename, framework="pt", device="cpu") as handle:
+        source = handle.get_slice(tensor_key)
+        source_shape = tuple(source.get_shape())
+        expected_shape = (sum(global_sizes), *dst.shape[1:])
+        if source_shape != expected_shape:
+            raise ValueError(f"checkpoint tensor {tensor_key} has shape {source_shape}, expected {expected_shape}")
+        source_offset = 0
+        destination_offset = 0
+        for global_size, local_size, (start, end) in zip(global_sizes, local_sizes, ranges, strict=True):
+            shard = source[source_offset + start : source_offset + end]
+            dst[destination_offset : destination_offset + local_size].copy_(shard.to(dtype=dst.dtype))
+            source_offset += global_size
+            destination_offset += local_size
+
+
 def load_k_shared_qkv_column_spec(
     module: nn.Module, index: SafetensorsIndex, prefix: str, spec: KSharedQKVColumnSpec, rank: int, world_size: int
 ) -> None:
@@ -810,6 +869,28 @@ def save_split_column_spec(
         tensors[key(template, prefix)] = tensor
 
 
+def save_packed_section_column_spec(
+    tensors: dict[str, torch.Tensor | None],
+    module: nn.Module,
+    prefix: str,
+    spec: PackedSectionColumnSpec,
+) -> None:
+    """Gather local packed sections into their original single HF tensor."""
+
+    tensor = attr_path(module, spec.tensor_attr)
+    global_sizes = tuple(int(size) for size in attr_path(module, spec.global_sizes_attr))
+    local_sizes = [int(size) for size in attr_path(module, spec.local_sizes_attr)]
+    world_size = get_tp_context().world_size
+    if any(
+        global_size != local_size * world_size
+        for global_size, local_size in zip(global_sizes, local_sizes, strict=True)
+    ):
+        raise ValueError("packed-section sizes are incompatible with the tensor-parallel world size")
+    if tensor.shape[0] != sum(local_sizes):
+        raise ValueError(f"packed-section source has {tensor.shape[0]} rows, expected {sum(local_sizes)}")
+    tensors[key(spec.key, prefix)] = gather_tensor_parallel_split_column_tensor(tensor, local_sizes)
+
+
 def save_ranged_split_column_spec(
     tensors: dict[str, torch.Tensor | None], module: nn.Module, prefix: str, spec: RangedSplitColumnSpec
 ) -> None:

areno/engine/data/rollout_state.py

@@ -113,6 +113,7 @@ def build_prefill_payload(self) -> dict | None:
         has_mrope_positions = False
         feature_mask: list[bool] = []
         image_features: list[dict] = []
+        image_sequence_modes: list[bool] = []
         cu_seqlens = [0]
         sample_indices: list[int] = []
         block_table: list[list[int]] = []
@@ -156,6 +157,7 @@ def build_prefill_payload(self) -> dict | None:
                             mrope_position_parts if has_mrope_positions else None,
                             feature_mask,
                             image_features,
+                            image_sequence_modes,
                             cu_seqlens,
                             sample_indices,
                             block_table,
@@ -175,6 +177,7 @@ def build_prefill_payload(self) -> dict | None:
                 chunk_len,
             )
             feature_mask.extend(local_mask)
+            image_sequence_modes.append(_prompt_has_image(self.prompt_features[seq_id], prompt))
             if local_features is not None:
                 image_features.append(local_features)
             local_mrope_positions = _slice_prompt_mrope_positions(
@@ -219,6 +222,7 @@ def build_prefill_payload(self) -> dict | None:
             mrope_position_parts if has_mrope_positions else None,
             feature_mask,
             image_features,
+            image_sequence_modes,
             cu_seqlens,
             sample_indices,
             block_table,
@@ -235,6 +239,7 @@ def _prefill_payload(
         mrope_position_parts: list[torch.Tensor] | None,
         feature_mask: list[bool],
         image_features: list[dict],
+        image_sequence_modes: list[bool],
         cu_seqlens: list[int],
         sample_indices: list[int],
         block_table: list[list[int]],
@@ -256,8 +261,13 @@ def _prefill_payload(
             "cache_block_offsets": torch.tensor(cache_block_offsets, dtype=torch.long),
             "recurrent_slots": torch.tensor(recurrent_slots, dtype=torch.long),
         }
-        if any(feature_mask) or image_features or mrope_position_parts is not None:
-            payload["features"] = _prefill_multimodal_features(feature_mask, image_features, mrope_position_parts)
+        if any(feature_mask) or image_features or any(image_sequence_modes) or mrope_position_parts is not None:
+            payload["features"] = _prefill_multimodal_features(
+                feature_mask,
+                image_features,
+                mrope_position_parts,
+                image_sequence_modes,
+            )
         return payload
 
     def ensure_decode_blocks(self, seq_ids: list[int], next_positions: list[int]) -> None:
@@ -327,6 +337,9 @@ def _slice_prompt_image_features(
             key in features
             for key in (
                 "pixel_values",
+                "input_image_embeds",
+                "image_sizes",
+                "image_attention_mask",
                 "image_grid_thw",
                 "target_sizes",
                 "pixel_values_videos",
@@ -366,6 +379,9 @@ def _slice_prompt_image_features(
         )
         for key in (
             "pixel_values",
+            "input_image_embeds",
+            "image_sizes",
+            "image_attention_mask",
             "image_grid_thw",
             "target_sizes",
             "num_patches_per_image",
@@ -412,8 +428,11 @@ def _prefill_multimodal_features(
     feature_mask: list[bool],
     image_features: list[dict],
     mrope_position_parts: list[torch.Tensor] | None = None,
+    image_sequence_modes: list[bool] | None = None,
 ) -> dict:
     features = {}
+    if image_sequence_modes is not None and any(image_sequence_modes):
+        features["image_sequence_mask"] = torch.tensor(image_sequence_modes, dtype=torch.bool)
     if mrope_position_parts is not None:
         features["mrope_position_ids"] = torch.cat(mrope_position_parts, dim=1).to(dtype=torch.long)
     if not image_features:
@@ -456,6 +475,18 @@ def _prompt_image_mask(features: dict, prompt: list[int]) -> list[bool]:
     return [int(token) in values for token in prompt]
 
 
+def _prompt_has_image(features: dict | None, prompt: list[int]) -> bool:
+    if features is None:
+        return False
+    mask = features.get("image_token_mask")
+    if mask is not None:
+        return bool(torch.as_tensor(mask, dtype=torch.bool).any())
+    image_token_id = features.get("image_token_id")
+    if image_token_id is None:
+        image_token_id = (features.get("modality_token_ids") or {}).get("image")
+    return image_token_id is not None and any(int(token) == int(image_token_id) for token in prompt)
+
+
 def payload_to_infer_meta(payload: dict, device: torch.device) -> InferMeta:
     """Move a scheduler payload to device and expose it as model metadata."""