chore: import upstream snapshot with attribution

python/mlc_llm/model/phi3v/phi3v_image.py

@@ -0,0 +1,309 @@
+"""
+Implementation for Phi architecture.
+"""
+
+from tvm import relax, tirx
+from tvm.relax.frontend import nn
+from tvm.relax.frontend.nn import Module, Tensor, op
+from tvm.script import tirx as T
+
+from mlc_llm.model.vision import CLIPVisionModel
+from mlc_llm.support.config import ConfigBase
+
+
+# mypy: disable-error-code="attr-defined"
+class ImageProjection(Module):
+    def __init__(self, config: ConfigBase):
+        super().__init__()
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size * 4, config.hidden_size, bias=True
+        )
+        self.act = nn.GELU()
+        self.linear_2 = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
+
+    def forward(self, image_features: Tensor) -> Tensor:
+        shape_1 = tirx.Var("shape_1", "int64")
+        image_features = op.wrap_nested(
+            relax.BlockBuilder()
+            .current()
+            .match_cast(
+                image_features._expr,
+                relax.TensorType([shape_1, image_features.shape[1]], image_features.dtype),
+            ),
+            "image_features",
+        )
+
+        hidden_states = self.linear_1(image_features)
+
+        shape_2 = tirx.Var("shape_2", "int64")
+        hidden_states = op.wrap_nested(
+            relax.BlockBuilder()
+            .current()
+            .match_cast(
+                hidden_states._expr,
+                relax.TensorType([shape_2, hidden_states.shape[1]], hidden_states.dtype),
+            ),
+            "hidden_states",
+        )
+
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class Phi3ImageEmbedding(Module):
+    def __init__(self, config: ConfigBase):
+        super().__init__()
+
+        self.img_processor = CLIPVisionModel(config.vision_config)
+        self.image_dim_out = config.img_processor["image_dim_out"]
+
+        self.glb_GN = nn.Parameter((1, 1, self.image_dim_out * 4))
+        self.sub_GN = nn.Parameter((1, 1, 1, self.image_dim_out * 4))
+
+        self.img_projection = ImageProjection(config)
+        self.image_size = config.vision_config.image_size
+
+    def apply_schedule(self, sch, block, bdx=32, tile=[32, 32]):
+        loop_x, loop_y = sch.get_loops(block)[-2:]
+        xo, xi = sch.split(loop_x, factors=[tile[0], None])
+        yo, yi = sch.split(loop_y, factors=[tile[1], None])
+        sch.reorder(xo, yo, xi, yi)
+        t = sch.fuse(xo, yo)
+        ty, tx = sch.split(t, factors=[None, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+    def dyn_repeat_4d_tensor(self, input_tensor, r0, r1, r2, r3) -> Tensor:
+        assert 4 == input_tensor.ndim, "input_tensor should be 4D data tensor"
+
+        def create_dyn_repeat_func(dtype):
+            @T.prim_func(s_tir=True)
+            def dyn_repeat_4d_tensor_func(  # pylint disable=too-many-locals
+                input_tensor: T.handle,
+                output: T.handle,
+                ch0: T.int64(),
+                ch1: T.int64(),
+                ch2: T.int64(),
+                ch3: T.int64(),
+            ):
+                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
+                n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
+                input_tensor_buf = T.match_buffer(input_tensor, (n, c, h, w), dtype=dtype)
+                out_buf = T.match_buffer(output, (n * ch0, c * ch1, h * ch2, w * ch3), dtype=dtype)
+
+                for n_idx in T.thread_binding(n * ch0, thread="blockIdx.x"):
+                    for c_idx in T.thread_binding(c * ch1, thread="blockIdx.y"):
+                        for h_idx, w_idx in T.grid(h * ch2, w * ch3):
+                            with T.sblock("dyn_repeat_4d_tensor"):
+                                T.reads(input_tensor_buf[n_idx, c_idx, h_idx, w_idx])
+                                T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
+                                out_buf[n_idx, c_idx, h_idx, w_idx] = input_tensor_buf[
+                                    n_idx % n, c_idx % c, h_idx % h, w_idx % w
+                                ]
+
+            return dyn_repeat_4d_tensor_func
+
+        n, c, h, w = input_tensor.shape
+        out = op.tensor_ir_op(
+            create_dyn_repeat_func(input_tensor.dtype),
+            "dyn_repeat_4d_tensor",
+            [input_tensor, r0, r1, r2, r3],
+            [Tensor.placeholder([n * r0, c * r1, h * r2, w * r3], input_tensor.dtype)],
+        )
+        return out
+
+    def dyn_concate_dim_2(self, input_1, input_2) -> Tensor:
+        def create_dyn_concate_func(dtype):
+            @T.prim_func(s_tir=True)
+            def dyn_concate_dim_2_func(input_1: T.handle, input_2: T.handle, output: T.handle):
+                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
+                n, c, h1, h2, w = T.int64(), T.int64(), T.int64(), T.int64(), T.int64()
+                input_1_buf = T.match_buffer(input_1, (n, c, h1, w), dtype=dtype)
+                input_2_buf = T.match_buffer(input_2, (n, c, h2, w), dtype=dtype)
+                out_buf = T.match_buffer(output, (n, c, h1 + h2, w), dtype=dtype)
+
+                for n_idx in T.thread_binding(n, thread="blockIdx.x"):
+                    for c_idx in T.thread_binding(c, thread="blockIdx.y"):
+                        for h_idx, w_idx in T.grid(h1 + h2, w):
+                            with T.sblock("dyn_concate_dim_2"):
+                                T.reads(input_1_buf[n_idx, c_idx, h_idx, w_idx])
+                                T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
+                                if h_idx < h1:
+                                    out_buf[n_idx, c_idx, h_idx, w_idx] = input_1_buf[
+                                        n_idx, c_idx, h_idx, w_idx
+                                    ]
+                                else:
+                                    out_buf[n_idx, c_idx, h_idx, w_idx] = input_2_buf[
+                                        n_idx, c_idx, h_idx - h1, w_idx
+                                    ]
+
+            return dyn_concate_dim_2_func
+
+        n1, c1, h1, w1 = input_1.shape
+        n2, c2, h2, w2 = input_2.shape
+        assert n1 == n2 and c1 == c2 and w1 == w2
+
+        out = op.tensor_ir_op(
+            create_dyn_concate_func(input_1.dtype),
+            "dyn_concate_dim_2",
+            [input_1, input_2],
+            [Tensor.placeholder([n1, c1, h1 + h2, w1], input_1.dtype)],
+        )
+        return out
+
+    def dyn_concate_dim_1(self, input_1, input_2) -> Tensor:
+        def create_dyn_concate_func(dtype):
+            @T.prim_func(s_tir=True)
+            def dyn_concate_dim_1_func(input_1: T.handle, input_2: T.handle, output: T.handle):
+                T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
+                c, h1, h2, w = T.int64(), T.int64(), T.int64(), T.int64()
+                input_1_buf = T.match_buffer(input_1, (c, h1, w), dtype=dtype)
+                input_2_buf = T.match_buffer(input_2, (c, h2, w), dtype=dtype)
+                out_buf = T.match_buffer(output, (c, h1 + h2, w), dtype=dtype)
+
+                for c_idx in T.thread_binding(c, thread="blockIdx.y"):
+                    for h_idx, w_idx in T.grid(h1 + h2, w):
+                        with T.sblock("dyn_concate_dim_1"):
+                            T.reads(input_1_buf[c_idx, h_idx, w_idx])
+                            T.writes(out_buf[c_idx, h_idx, w_idx])
+                            if h_idx < h1:
+                                out_buf[c_idx, h_idx, w_idx] = input_1_buf[c_idx, h_idx, w_idx]
+                            else:
+                                out_buf[c_idx, h_idx, w_idx] = input_2_buf[c_idx, h_idx - h1, w_idx]
+
+            return dyn_concate_dim_1_func
+
+        c1, h1, w1 = input_1.shape
+        c2, h2, w2 = input_2.shape
+        assert c1 == c2 and w1 == w2
+
+        out = op.tensor_ir_op(
+            create_dyn_concate_func(input_1.dtype),
+            "dyn_concate",
+            [input_1, input_2],
+            [Tensor.placeholder([c1, h1 + h2, w1], input_1.dtype)],
+        )
+        return out
+
+    def get_img_features(self, img_embeds: Tensor) -> Tensor:
+        img_processor_output = self.img_processor(img_embeds)
+        patch_feature = nn.op.split(img_processor_output, indices_or_sections=[1], axis=1)
+        return patch_feature[1]
+
+    def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
+        N, L, C = image_features.shape
+        num_images = 1
+        H = int(L**0.5)
+        image_features = nn.op.reshape(image_features, ([N, H, H, C]))  # N, 24, 24, 1024
+        image_features = nn.op.reshape(
+            image_features, ([N, H // 2, 2, H // 2, 2, C])
+        )  # N, 12, 2, 12, 2, 1024
+
+        new_s1 = tirx.Var("new_s1", "int64")
+        new_s2 = tirx.Var("new_s2", "int64")
+
+        image_features = op.wrap_nested(
+            relax.BlockBuilder()
+            .current()
+            .match_cast(
+                image_features._expr,
+                relax.TensorType(
+                    [
+                        image_features.shape[0],
+                        new_s1,
+                        image_features.shape[2],
+                        new_s2,
+                        image_features.shape[4],
+                        image_features.shape[5],
+                    ],
+                    image_features.dtype,
+                ),
+            ),
+            "image_features_1",
+        )
+
+        image_features = nn.op.permute_dims(
+            image_features, axes=([0, 1, 3, 2, 4, 5])
+        )  # N, 12, 12, 2, 2, 1024
+        image_features = nn.op.reshape(image_features, ([N, -1, 4 * C]))  # N, 144, 4096
+        image_features = nn.op.reshape(
+            image_features, ([num_images, h_crop, w_crop, H // 2, H // 2, -1])
+        )
+
+        new_s3 = tirx.Var("new_s3", "int64")
+        new_s4 = tirx.Var("new_s4", "int64")
+
+        image_features = op.wrap_nested(
+            relax.BlockBuilder()
+            .current()
+            .match_cast(
+                image_features._expr,
+                relax.TensorType(
+                    [
+                        image_features.shape[0],
+                        new_s3,
+                        image_features.shape[2],
+                        new_s4,
+                        image_features.shape[4],
+                        image_features.shape[5],
+                    ],
+                    image_features.dtype,
+                ),
+            ),
+            "image_features_2",
+        )
+
+        image_features = nn.op.permute_dims(image_features, axes=([0, 1, 3, 2, 4, 5]))
+        image_features_hd = nn.op.reshape(
+            image_features, ([num_images, h_crop * H // 2, w_crop * H // 2, 4 * C])
+        )
+
+        return image_features_hd
+
+    def add_image_newline(self, image_features_hd):
+        """
+        image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
+        output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)
+        """
+        num_images, h, w, hid_dim = image_features_hd.shape
+        temp_sub_GN = self.dyn_repeat_4d_tensor(
+            self.sub_GN, T.int64(1), T.int64(h), T.int64(1), T.int64(1)
+        )
+        image_features_hd_newline = self.dyn_concate_dim_2(image_features_hd, temp_sub_GN)
+        image_features_hd_newline = nn.op.reshape(
+            image_features_hd_newline, ([num_images, -1, hid_dim])
+        )
+        return image_features_hd_newline
+
+    def forward(self, pixel_values: Tensor, h_crop, w_crop) -> Tensor:
+        img_features = self.get_img_features(pixel_values)
+        img_features = nn.op.split(img_features, indices_or_sections=[1], axis=0)
+
+        global_image_features = img_features[0]
+        global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1)
+        global_image_features_hd_newline = self.add_image_newline(global_image_features_hd)
+
+        sub_image_features = img_features[1]
+        sub_image_features_hd = self.reshape_hd_patches_2x2merge(sub_image_features, h_crop, w_crop)
+        sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd)
+
+        global_image_features_hd = nn.op.squeeze(global_image_features_hd_newline, 0)
+
+        combined_image = self.dyn_concate_dim_1(sub_image_features_hd_newline, self.glb_GN)
+        combined_image = self.dyn_concate_dim_1(combined_image, global_image_features_hd_newline)
+        combined_image = nn.op.squeeze(combined_image, 0)
+
+        new_s7 = tirx.Var("new_s7", "int64")
+
+        combined_image = op.wrap_nested(
+            relax.BlockBuilder()
+            .current()
+            .match_cast(
+                combined_image._expr,
+                relax.TensorType([new_s7, combined_image.shape[1]], combined_image.dtype),
+            ),
+            "combined_image",
+        )
+        output_image = self.img_projection(combined_image)
+        return output_image

python/mlc_llm/model/phi3v/phi3v_loader.py

@@ -0,0 +1,132 @@
+"""
+This file specifies how MLC's Phi parameter maps from other formats, for example HuggingFace
+PyTorch, HuggingFace safetensors.
+"""
+
+import functools
+
+from mlc_llm.loader import ExternMapping
+from mlc_llm.quantization import Quantization
+
+from .phi3v_model import Phi3VConfig, Phi3VForCausalLM
+
+
+def huggingface(model_config: Phi3VConfig, quantization: Quantization) -> ExternMapping:
+    """Returns a parameter mapping that maps from the names of MLC LLM parameters to
+    the names of Phi-1/Phi-1.5 HuggingFace PyTorch parameters.
+
+    Parameters
+    ----------
+    model_config : PhiConfig
+        The configuration of the Phi model.
+
+    quantization : Quantization
+        The quantization configuration.
+
+    Returns
+    -------
+    param_map : ExternMapping
+        The parameter mapping from MLC to HuggingFace PyTorch.
+    """
+    model = Phi3VForCausalLM(model_config)
+    if quantization is not None:
+        model.to(quantization.model_dtype)
+    _, _named_params = model.export_tvm(spec=model.get_default_spec())
+    named_parameters = dict(_named_params)
+
+    mapping = ExternMapping()
+
+    def _add(mlc_name, hf_name=None):
+        if None is hf_name:
+            hf_name = mlc_name
+
+        mapping.add_mapping(
+            mlc_name,
+            [hf_name],
+            functools.partial(
+                lambda x, dtype: x.astype(dtype),
+                dtype=named_parameters[mlc_name].dtype,
+            ),
+        )
+
+    def _add_vision(name):
+        _add(name, "model." + name)
+
+    _add("model.embd.weight", "model.embed_tokens.weight")
+
+    prefix = "model.h"
+    hf_prefix = "model.layers"
+    for i in range(model_config.num_hidden_layers):
+        _add(f"{prefix}.{i}.ln.weight", f"{hf_prefix}.{i}.input_layernorm.weight")
+        _add(
+            f"{prefix}.{i}.mlp.down_proj.weight",
+            f"{hf_prefix}.{i}.mlp.down_proj.weight",
+        )
+        _add(
+            f"{prefix}.{i}.mlp.gate_up_proj.weight",
+            f"{hf_prefix}.{i}.mlp.gate_up_proj.weight",
+        )
+        _add(
+            f"{prefix}.{i}.post_attention_layernorm.weight",
+            f"{hf_prefix}.{i}.post_attention_layernorm.weight",
+        )
+        _add(
+            f"{prefix}.{i}.mixer.out_proj.weight",
+            f"{hf_prefix}.{i}.self_attn.o_proj.weight",
+        )
+        _add(
+            f"{prefix}.{i}.mixer.qkv_proj.weight",
+            f"{hf_prefix}.{i}.self_attn.qkv_proj.weight",
+        )
+
+    prefix = "vision_embed_tokens.img_processor.vision_model.encoder.layers"
+    for i in range(model_config.vision_config.num_hidden_layers):
+        _add_vision(f"{prefix}.{i}.layer_norm1.bias")
+        _add_vision(f"{prefix}.{i}.layer_norm1.weight")
+        _add_vision(f"{prefix}.{i}.layer_norm2.bias")
+        _add_vision(f"{prefix}.{i}.layer_norm2.weight")
+        _add_vision(f"{prefix}.{i}.mlp.fc1.bias")
+        _add_vision(f"{prefix}.{i}.mlp.fc1.weight")
+        _add_vision(f"{prefix}.{i}.mlp.fc2.bias")
+        _add_vision(f"{prefix}.{i}.mlp.fc2.weight")
+        _add_vision(f"{prefix}.{i}.self_attn.k_proj.bias")
+        _add_vision(f"{prefix}.{i}.self_attn.k_proj.weight")
+        _add_vision(f"{prefix}.{i}.self_attn.out_proj.bias")
+        _add_vision(f"{prefix}.{i}.self_attn.out_proj.weight")
+        _add_vision(f"{prefix}.{i}.self_attn.q_proj.bias")
+        _add_vision(f"{prefix}.{i}.self_attn.q_proj.weight")
+        _add_vision(f"{prefix}.{i}.self_attn.v_proj.bias")
+        _add_vision(f"{prefix}.{i}.self_attn.v_proj.weight")
+
+    _add_vision("vision_embed_tokens.sub_GN")
+    _add_vision("vision_embed_tokens.glb_GN")
+    _add_vision("vision_embed_tokens.img_processor.vision_model.embeddings.class_embedding")
+    _add_vision("vision_embed_tokens.img_processor.vision_model.embeddings.patch_embedding.weight")
+    _add_vision(
+        "vision_embed_tokens.img_processor.vision_model.embeddings.position_embedding.weight"
+    )
+    _add_vision("vision_embed_tokens.img_processor.vision_model.post_layernorm.bias")
+    _add_vision("vision_embed_tokens.img_processor.vision_model.post_layernorm.weight")
+    _add_vision("vision_embed_tokens.img_processor.vision_model.pre_layrnorm.bias")
+    _add_vision("vision_embed_tokens.img_processor.vision_model.pre_layrnorm.weight")
+
+    prefix = "vision_embed_tokens.img_projection"
+    _add(f"{prefix}.linear_1.bias", f"model.{prefix}.0.bias")
+    _add(f"{prefix}.linear_1.weight", f"model.{prefix}.0.weight")
+    _add(f"{prefix}.linear_2.bias", f"model.{prefix}.2.bias")
+    _add(f"{prefix}.linear_2.weight", f"model.{prefix}.2.weight")
+
+    for mlc_name, mlc_param in named_parameters.items():
+        if mlc_name not in mapping.param_map:
+            mapping.add_mapping(
+                mlc_name,
+                [mlc_name],
+                functools.partial(
+                    lambda x, dtype: x.astype(dtype),
+                    dtype=mlc_param.dtype,
+                ),
+            )
+
+    mapping.add_unused("model.embed_tokens.weight")
+
+    return mapping

python/mlc_llm/model/phi3v/phi3v_model.py

@@ -0,0 +1,390 @@
+"""
+Implementation for Phi architecture.
+"""
+
+import dataclasses
+from typing import Any, Dict, Optional  # noqa: UP035
+
+from tvm import relax, target, tirx
+from tvm.relax.frontend import nn
+from tvm.relax.frontend.nn import Tensor, op
+
+from mlc_llm import op as op_ext
+from mlc_llm.model.model_utils import index_last_token
+from mlc_llm.model.phi3 import Phi3Model
+from mlc_llm.model.vision import CLIPVisionConfig, ImageProcessor
+from mlc_llm.nn import PagedKVCache, RopeMode
+from mlc_llm.support import logging
+from mlc_llm.support.config import ConfigBase
+from mlc_llm.support.style import bold
+
+from .phi3v_image import Phi3ImageEmbedding
+
+logger = logging.getLogger(__name__)
+
+CLIPVISION_DEFAULT_CONFIG = {
+    "hidden_size": 1024,
+    "image_size": 336,
+    "intermediate_size": 4096,
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "patch_size": 14,
+    "projection_dim": 768,
+    "layer_norm_eps": 1e-05,
+    "vocab_size": None,
+}
+
+
+@dataclasses.dataclass
+class Phi3VConfig(ConfigBase):
+    """Configuration of the Phi-3 Vision model."""
+
+    model_type: str
+    hidden_size: int
+    vocab_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    intermediate_size: int
+    rms_norm_eps: float
+    num_key_value_heads: int
+    max_position_embeddings: int
+    vision_config: CLIPVisionConfig = None
+    img_processor: Optional[Dict[str, Any]] = None  # noqa: UP006
+    position_embedding_base: int = 0
+    rope_scaling: Optional[Dict[str, Any]] = None  # noqa: UP006
+    original_max_position_embeddings: int = 0
+    context_window_size: int = 0
+    prefill_chunk_size: int = 0
+    head_dim: int = 0
+    tensor_parallel_shards: int = 1
+    max_batch_size: int = 1
+    kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)  # noqa: UP006
+
+    def __post_init__(self):
+        vision_config_dict: Dict[str, Any]  # noqa: UP006
+        if isinstance(self.vision_config, CLIPVisionConfig):
+            vision_config_dict = dataclasses.asdict(self.vision_config)
+        else:
+            vision_config_dict = dict(CLIPVISION_DEFAULT_CONFIG)
+
+        for k, v in vision_config_dict.pop("kwargs", {}).items():
+            vision_config_dict[k] = v
+
+        self.vision_config = CLIPVisionConfig.from_dict(vision_config_dict)
+
+        if self.position_embedding_base == 0:
+            if "rope_theta" in self.kwargs:
+                self.position_embedding_base = self.kwargs.pop("rope_theta")
+            else:
+                self.position_embedding_base = 10000
+        if self.rope_scaling is not None:
+            if "type" not in self.rope_scaling:
+                self.rope_scaling = None
+            else:
+                if self.rope_scaling["type"] == "su":
+                    self.rope_scaling["type"] = "longrope"
+
+                assert self.rope_scaling["type"] == "longrope", (
+                    f"Unsupported RoPE scaling type {self.rope_scaling['rope_type']} for Phi3"
+                )
+                self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+                (
+                    self.rope_scaling["max_position_embeddings"],
+                    self.rope_scaling["original_max_position_embeddings"],
+                ) = (
+                    self.max_position_embeddings,
+                    self.original_max_position_embeddings,
+                )
+
+        if self.context_window_size == 0:
+            self.context_window_size = self.max_position_embeddings
+
+        if self.prefill_chunk_size == 0:
+            logger.info(
+                "%s defaults to %d",
+                bold("prefill_chunk_size"),
+                min(self.context_window_size, 8192),
+            )
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
+        elif self.prefill_chunk_size > self.context_window_size:
+            logger.info(
+                "Overriding %s from %d to %d",
+                bold("prefill_chunk_size"),
+                self.prefill_chunk_size,
+                min(self.context_window_size, 8192),
+            )
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
+
+        if self.num_key_value_heads == 0 or self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.head_dim == 0:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+        assert self.head_dim * self.num_attention_heads == self.hidden_size
+        assert self.num_attention_heads % self.num_key_value_heads == 0
+
+
+# mypy: disable-error-code="arg-type,annotation-unchecked"
+class Phi3VForCausalLM(nn.Module):
+    def __init__(self, config: Phi3VConfig) -> None:
+        super().__init__()
+
+        self.config = config
+        self.model = Phi3Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False)
+        self.vision_embed_tokens = Phi3ImageEmbedding(config)
+        self.image_processor = ImageProcessor()
+        self.num_hidden_layers = config.num_hidden_layers
+        self.num_attention_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.hidden_size = config.hidden_size
+        self.vocab_size = config.vocab_size
+        self.rope_scaling = config.rope_scaling
+        self.rope_theta = config.position_embedding_base
+        self.rope_ext_factors = (
+            (config.rope_scaling["long_factor"] + config.rope_scaling["short_factor"])
+            if config.rope_scaling is not None
+            else None
+        )
+        self.tensor_parallel_shards = config.tensor_parallel_shards
+        self.dtype = "float32"
+        self.image_dtype = (
+            "uint32"
+            if target.Target.current() and target.Target.current().kind.name == "webgpu"
+            else "uint8"
+        )
+
+    def to(self, dtype: Optional[str] = None):
+        super().to(dtype=dtype)
+        if dtype is not None:
+            self.dtype = dtype
+
+    def batch_forward(
+        self,
+        input_embeds: Tensor,
+        paged_kv_cache: PagedKVCache,
+        logit_positions: Optional[Tensor] = None,
+    ):
+        op_ext.configure()
+
+        hidden_states = self.model(input_embeds, paged_kv_cache)
+        if logit_positions is not None:
+            hidden_states = op.take(hidden_states, logit_positions, axis=1)
+        lm_logits = self.lm_head(hidden_states)
+        if lm_logits.dtype != "float32":
+            lm_logits = lm_logits.astype("float32")
+        return lm_logits
+
+    def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
+        op_ext.configure()
+
+        hidden_states = self.model(input_embed, paged_kv_cache)
+        hidden_states = index_last_token(hidden_states)
+        logits = self.lm_head(hidden_states)
+
+        if logits.dtype != "float32":
+            logits = logits.astype("float32")
+
+        return logits, paged_kv_cache
+
+    def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
+        op_ext.configure()
+
+        hidden_states = self.model(input_embed, paged_kv_cache)
+        logits = self.lm_head(hidden_states)
+        if logits.dtype != "float32":
+            logits = logits.astype("float32")
+        return logits, paged_kv_cache
+
+    def batch_prefill(
+        self,
+        input_embeds: Tensor,
+        logit_positions: Tensor,
+        paged_kv_cache: PagedKVCache,
+    ):
+        if self.tensor_parallel_shards > 1:
+            logit_positions = op.ccl_broadcast_from_worker0(logit_positions)
+        logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions)
+        return logits, paged_kv_cache
+
+    def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache):
+        logits = self.batch_forward(input_embeds, paged_kv_cache)
+        return logits, paged_kv_cache
+
+    def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache):
+        logits = self.batch_forward(input_embeds, paged_kv_cache)
+        return logits, paged_kv_cache
+
+    def embed(self, input_ids: Tensor):
+        if self.tensor_parallel_shards > 1:
+            input_ids = op.ccl_broadcast_from_worker0(input_ids)
+        embeds = self.model.embd(input_ids)
+        return embeds
+
+    def image_preprocess(
+        self, pixel_values: Tensor, resized_height, resized_width, num_crops=16
+    ) -> Tensor:
+        pixel_values = op.permute_dims(pixel_values, axes=(0, 3, 1, 2))  # NHWC -> NCHW
+        pixel_values = self.image_processor.resize(
+            pixel_values, params={"height": resized_height, "width": resized_width}
+        )
+        pixel_values = self.image_processor.pad(pixel_values, dtype=self.image_dtype)
+        pixel_values = self.image_processor.rescale(pixel_values)
+        pixel_values = self.image_processor.normalize(pixel_values)
+        global_image = self.image_processor.resize(
+            pixel_values, params={"height": 336, "width": 336}
+        )
+        global_image = op.wrap_nested(
+            relax.BlockBuilder()
+            .current()
+            .match_cast(
+                global_image._expr,
+                relax.TensorType(
+                    [global_image.shape[0], global_image.shape[1], 336, 336],
+                    global_image.dtype,
+                ),
+            ),
+            "global_image",
+        )
+
+        n, c, h, w = pixel_values.shape
+        assert isinstance(h, tirx.Mul) and isinstance(h.b, tirx.IntImm) and h.b.value == 336
+        pixel_values = op.reshape(pixel_values, shape=(1, 3, h.a, 336, w // 336, 336))
+        pixel_values = op.permute_dims(pixel_values, axes=(0, 2, 4, 1, 3, 5))
+        pixel_values = op.reshape(pixel_values, shape=(-1, 3, 336, 336))
+        combined_image = op.concat([global_image, pixel_values], dim=0)
+
+        # pad to max num crops tensor
+        b, c, h, w = combined_image.shape
+        zeros = op.zeros((num_crops + 1 - b, c, h, w))
+        combined_image = op.concat([combined_image, zeros], dim=0)
+
+        combined_image = op.wrap_nested(
+            relax.BlockBuilder()
+            .current()
+            .match_cast(
+                combined_image._expr,
+                relax.TensorType([num_crops + 1, c, h, w], combined_image.dtype),
+            ),
+            "combined_image",
+        )
+
+        return combined_image
+
+    def image_embed(
+        self,
+        pixel_values: Tensor,
+        resized_height,
+        resized_width,
+        crop_height,
+        crop_width,
+    ) -> Tensor:
+        n, h, w, c = pixel_values.shape
+        pixel_values = self.image_preprocess(pixel_values, resized_height, resized_width)
+        pixel_values = pixel_values.astype(self.dtype)
+        return self.vision_embed_tokens(pixel_values, crop_height, crop_width)
+
+    def create_paged_kv_cache(
+        self,
+        max_batch_size: tirx.Var,
+        max_total_seq_len: tirx.Var,
+        prefill_chunk_size: tirx.Var,
+        page_size: tirx.Var,
+        support_sliding_window: tirx.Var,
+    ) -> PagedKVCache:
+        return PagedKVCache.create_generic(
+            attn_kind="mha",
+            max_batch_size=max_batch_size,
+            max_total_seq_len=max_total_seq_len,
+            prefill_chunk_size=prefill_chunk_size,
+            page_size=page_size,
+            support_sliding_window=support_sliding_window,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards,
+            num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            rope_mode=RopeMode.NORMAL,
+            rope_scaling=self.rope_scaling,
+            rope_scale=1,
+            rope_theta=self.rope_theta,
+            rope_ext_factors=self.rope_ext_factors,
+            dtype=self.dtype,
+        )
+
+    def get_default_spec(self):
+        mod_spec = {
+            "embed": {
+                "input_ids": nn.spec.Tensor(["seq_len"], "int32"),
+                "$": {
+                    "param_mode": "packed",
+                    "effect_mode": "none",
+                },
+            },
+            "image_embed": {
+                "pixel_values": nn.spec.Tensor(
+                    [1, "image_height", "image_width", 3], self.image_dtype
+                ),
+                "resized_height": nn.spec.Int(),
+                "resized_width": nn.spec.Int(),
+                "crop_height": nn.spec.Int(),
+                "crop_width": nn.spec.Int(),
+                "$": {
+                    "param_mode": "packed",
+                    "effect_mode": "none",
+                },
+            },
+            "prefill": {
+                "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype),
+                "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
+                "$": {
+                    "param_mode": "packed",
+                    "effect_mode": "none",
+                },
+            },
+            "decode": {
+                "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype),
+                "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
+                "$": {
+                    "param_mode": "packed",
+                    "effect_mode": "none",
+                },
+            },
+            "batch_prefill": {
+                "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype),
+                "logit_positions": nn.spec.Tensor(["batch_size"], "int32"),
+                "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
+                "$": {
+                    "param_mode": "packed",
+                    "effect_mode": "none",
+                },
+            },
+            "batch_decode": {
+                "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype),
+                "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
+                "$": {
+                    "param_mode": "packed",
+                    "effect_mode": "none",
+                },
+            },
+            "batch_verify": {
+                "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype),
+                "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
+                "$": {
+                    "param_mode": "packed",
+                    "effect_mode": "none",
+                },
+            },
+            "create_paged_kv_cache": {
+                "max_batch_size": int,
+                "max_total_seq_len": int,
+                "prefill_chunk_size": int,
+                "page_size": int,
+                "support_sliding_window": int,
+                "$": {
+                    "param_mode": "none",
+                    "effect_mode": "none",
+                },
+            },
+        }
+        return nn.spec.ModuleSpec.from_raw(mod_spec, self)