feat(models): add Phi-4 Multimodal text-only adapter

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/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/layers/attention.py

@@ -32,7 +32,7 @@ class CausalSelfAttention(nn.Module):
     reduces across ranks to reassemble the full hidden state.
     """
 
-    def __init__(self, config: ModelConfig, layer_idx: int):
+    def __init__(self, config: ModelConfig, layer_idx: int, *, rotary_embedding: nn.Module | None = None):
         super().__init__()
         ctx = get_tp_context()
         self.layer_idx = layer_idx
@@ -58,7 +58,11 @@ def __init__(self, config: ModelConfig, layer_idx: int):
         # Row-parallel output projection: input is already sharded along
         # head dimension, output is all-reduced across ranks.
         self.o_proj = RowParallelLinear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
-        self.rope = RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta)
+        self.rope = (
+            rotary_embedding
+            if rotary_embedding is not None
+            else RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta)
+        )
         # Optional per-head QK normalization (used by some recent models).
         self.q_norm = RMSNorm(config.head_dim, config.rms_norm_eps) if config.qk_norm else None
         self.k_norm = RMSNorm(config.head_dim, config.rms_norm_eps) if config.qk_norm else None
@@ -93,14 +97,27 @@ def forward(
             k = self.k_norm(k)
         # Rotary embedding is applied on the head dim using position-indexed
         # cos/sin tables; positions are broadcast across heads.
-        q, k = self.rope(q, k, position_ids)
+        q, k = self.apply_rotary(q, k, position_ids, train_meta, infer_meta)
 
         # Presence of infer_meta selects the paged KV-cache backend; otherwise
         # we run the training-mode FlashAttention (padded or varlen packed).
         if infer_meta is not None:
             return self.forward_infer(q, k, v, infer_meta)
         return self.forward_train(q, k, v, train_meta)
 
+    def apply_rotary(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        position_ids: torch.Tensor,
+        train_meta: TrainMeta | None,
+        infer_meta: InferMeta | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply the model's rotary embedding, with a model override hook."""
+
+        del train_meta, infer_meta
+        return self.rope(q, k, position_ids)
+
     def forward_train(
         self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, train_meta: TrainMeta | None
     ) -> torch.Tensor:

areno/engine/layers/mlp.py

@@ -10,7 +10,8 @@
 import torch
 from torch import nn
 
-from areno.accel.ops import areno_silu_and_mul, log_once
+from areno.accel.activations import areno_silu_and_mul
+from areno.accel.utils import log_once
 from areno.engine.config import ModelConfig
 from areno.engine.layers.linear import MergedColumnParallelLinear, RowParallelLinear
 

areno/engine/layers/norm.py

@@ -13,7 +13,7 @@
 from torch import nn
 
 from areno.accel import areno_rmsnorm
-from areno.accel.ops import can_use_cuda_kernel, log_once, rms_norm_gate_fwd
+from areno.accel.utils import can_use_cuda_kernel, log_once
 from areno.engine.layers.linear import mark_tensor_parallel_parameter
 
 
@@ -90,8 +90,10 @@ def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
         # Reshape last dim into (groups_per_rank, group_width) for the kernel.
         x = x.view(*shape[:-1], self.groups_per_rank, self.group_width)
         gate = gate.view(*shape[:-1], self.groups_per_rank, self.group_width)
-        if rms_norm_gate_fwd is None or not can_use_cuda_kernel(x, "fused group RMSNorm sigmoid gate kernel"):
+        if not can_use_cuda_kernel(x, "fused group RMSNorm sigmoid gate kernel"):
             raise RuntimeError("ARENO group RMSNorm sigmoid gate requires the fused CUDA kernel")
+        from areno.accel.kernels.group_rmsnorm import rms_norm_gate_fwd
+
         log_once("group_rmsnorm_sigmoid_gate", "using fused group RMSNorm sigmoid gate kernel")
         # Flatten the leading dims into a single batch so the kernel only
         # sees a 3D (B, groups, width) tensor.

areno/engine/runtime/decode_graph.py

@@ -81,6 +81,7 @@ def __init__(
         """Allocate static input buffers and the `InferMeta` baked into capture."""
 
         self.model = model
+        self.decode_cache_length_limit = getattr(model, "decode_cache_length_limit", None)
         self.bucket = bucket
         self.scratch_block = scratch_block
         self.scratch_recurrent_slot = scratch_recurrent_slot
@@ -156,6 +157,7 @@ def replay_tensors(
         actual = int(input_ids.numel())
         if actual > self.bucket:
             raise ValueError(f"decode payload has {actual} tokens, graph bucket is {self.bucket}")
+        _validate_decode_cache_length(cache_seqlens, actual, self.decode_cache_length_limit)
 
         # Copy the live values into the captured-stable buffers. The graph
         # was recorded against these buffer addresses so `copy_` here is what
@@ -186,3 +188,14 @@ def replay_tensors(
         self.graph.replay()
         assert self.logits_shard is not None
         return self.logits_shard
+
+
+def _validate_decode_cache_length(
+    cache_seqlens: torch.Tensor,
+    actual: int,
+    limit: int | None,
+) -> None:
+    if limit is not None and actual and int(cache_seqlens[:actual].max().item()) >= limit:
+        raise ValueError(
+            "cached decode cannot cross the model's rotary-factor boundary; run a full long-context prefill"
+        )

areno/models/__init__.py

@@ -25,6 +25,13 @@ def _register_qwen35() -> None:
     register_adapter(Qwen35Adapter())
 
 
+def _register_phi4mm() -> None:
+    from areno.models.phi4mm import Phi4MMAdapter
+    from areno.models.registry import register_adapter
+
+    register_adapter(Phi4MMAdapter())
+
+
 def _register_bailing() -> None:
     from areno.models.bailing import BailingMoeLinearV2Adapter
     from areno.models.registry import register_adapter
@@ -71,6 +78,7 @@ def _register_olmo2() -> None:
     "llama": _register_llama,
     "qwen3": _register_qwen3,
     "qwen3_5": _register_qwen35,
+    "phi4mm": _register_phi4mm,
     "bailing": _register_bailing,
     "bailing_v3": _register_bailing_v3,
     "gemma4": _register_gemma4,

areno/models/phi4mm/__init__.py

@@ -0,0 +1,21 @@
+"""Phi-4-Multimodal language-backbone adapter."""
+
+from __future__ import annotations
+
+from areno.models.phi4mm.model import (
+    Phi4MMAdapter,
+    Phi4MMAttention,
+    Phi4MMDecoderLayer,
+    Phi4MMForCausalLM,
+    Phi4MMLongRoPEScaledRotaryEmbedding,
+    Phi4MMModel,
+)
+
+__all__ = [
+    "Phi4MMAdapter",
+    "Phi4MMAttention",
+    "Phi4MMDecoderLayer",
+    "Phi4MMForCausalLM",
+    "Phi4MMLongRoPEScaledRotaryEmbedding",
+    "Phi4MMModel",
+]