feat(models): add Phi-4 multimodal audio 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

@@ -7,6 +7,7 @@
 import threading
 from collections.abc import Mapping, Sequence
 from typing import Any
+from urllib.request import urlopen
 
 import torch
 
@@ -119,6 +120,9 @@ def encode_processor_messages(
     """Use a native multimodal processor to load media and expand soft-token slots."""
 
     normalized = _normalize_multimodal_messages(messages)
+    identity = f"{type(processor).__module__}.{type(processor).__name__}".lower()
+    if "phi4mm" in identity:
+        return _encode_phi4mm_messages(processor, normalized, tools=tools)
     kwargs: dict[str, Any] = {
         "tokenize": True,
         "add_generation_prompt": True,
@@ -144,6 +148,94 @@ def encode_processor_messages(
     return tokens, features or None
 
 
+def _encode_phi4mm_messages(
+    processor: Any,
+    messages: list[dict[str, Any]],
+    *,
+    tools: Any = None,
+) -> tuple[list[int], dict[str, Any] | None]:
+    """Bridge structured API messages to the released Phi-4 processor API."""
+
+    images = []
+    audios = []
+    rendered_messages = []
+    for message in messages:
+        rendered = dict(message)
+        content = rendered.get("content")
+        if isinstance(content, list):
+            pieces = []
+            for part in content:
+                if not isinstance(part, dict):
+                    pieces.append(str(part))
+                    continue
+                kind = str(part.get("type", ""))
+                if kind == "text":
+                    pieces.append(str(part.get("text", "")))
+                elif kind == "image":
+                    images.append(_load_phi4mm_image(part.get("url")))
+                    pieces.append(f"<|image_{len(images)}|>")
+                elif kind == "audio":
+                    audios.append(_load_phi4mm_audio(part.get("url")))
+                    pieces.append(f"<|audio_{len(audios)}|>")
+                else:
+                    raise ValueError(f"Phi4MM does not support multimodal content type {kind!r}")
+            rendered["content"] = "".join(pieces)
+        rendered_messages.append(rendered)
+    template_kwargs: dict[str, Any] = {"tokenize": False, "add_generation_prompt": True}
+    if tools:
+        template_kwargs["tools"] = tools
+    prompt = apply_chat_template_with_options(processor.tokenizer, rendered_messages, **template_kwargs)
+    if prompt.endswith("<|endoftext|>"):
+        prompt = prompt.removesuffix("<|endoftext|>")
+    encoded = processor(
+        text=prompt,
+        images=images or None,
+        audios=audios or None,
+        return_tensors=getattr(processor, "_areno_return_tensors", "pt"),
+    )
+    input_ids = encoded["input_ids"]
+    tokens = normalize_token_ids(input_ids[0].tolist())
+    features = {
+        key: value for key, value in encoded.items() if key not in {"input_ids", "attention_mask", "token_type_ids"}
+    }
+    token_ids = modality_token_ids(processor)
+    features["modality_token_ids"] = token_ids
+    features["image_token_id"] = token_ids["image"]
+    features["audio_token_id"] = token_ids["audio"]
+    return tokens, features
+
+
+def _load_phi4mm_image(reference: Any) -> Any:
+    if not isinstance(reference, str) or not reference:
+        raise ValueError("Phi4MM image content requires a URL or data URI")
+    if reference.startswith("data:"):
+        return _load_base64_image(reference)
+    try:
+        from PIL import Image
+    except ImportError as exc:
+        raise ValueError("Phi4MM image input requires Pillow") from exc
+    if reference.startswith(("http://", "https://")):
+        with urlopen(reference, timeout=30) as response:  # noqa: S310
+            return Image.open(io.BytesIO(response.read())).convert("RGB")
+    return Image.open(reference).convert("RGB")
+
+
+def _load_phi4mm_audio(reference: Any) -> tuple[Any, int]:
+    if not isinstance(reference, str) or not reference:
+        raise ValueError("Phi4MM audio content requires a URL or data URI")
+    try:
+        import soundfile
+    except ImportError as exc:
+        raise ValueError("Phi4MM audio input requires soundfile") from exc
+    if reference.startswith("data:"):
+        _, _, payload = reference.partition(",")
+        return soundfile.read(io.BytesIO(base64.b64decode(payload)))
+    if reference.startswith(("http://", "https://")):
+        with urlopen(reference, timeout=30) as response:  # noqa: S310
+            return soundfile.read(io.BytesIO(response.read()))
+    return soundfile.read(reference)
+
+
 def _ensure_gemma4_torchvision_video_fps(processor: Any) -> None:
     """Backfill FPS metadata omitted by torchvision for some browser videos."""
 
@@ -194,6 +286,17 @@ def modality_token_ids(processor: Any) -> dict[str, int]:
         value = getattr(processor, f"{modality}_token_id", None)
         if isinstance(value, int) and value >= 0:
             result[modality] = int(value)
+    tokenizer = getattr(processor, "tokenizer", None)
+    if tokenizer is not None:
+        for modality, token in (("image", "<|endoftext10|>"), ("audio", "<|endoftext11|>")):
+            if modality in result:
+                continue
+            identity = f"{type(processor).__module__}.{type(processor).__name__}".lower()
+            if "phi4mm" not in identity:
+                continue
+            token_id = tokenizer.convert_tokens_to_ids(token)
+            if isinstance(token_id, int) and token_id >= 0:
+                result[modality] = token_id
     return result
 
 
@@ -430,7 +533,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 +546,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: