Enable running iOS and macOS authored models on CUDA GPUs (#99)

* authoring and relevant changes for evals

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>

* fix code format

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>

---------

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>

Author: pkmandke

python/src/coreai_models/export/pipeline.py

@@ -62,6 +62,8 @@ class ExportConfig:
     output_name: str | None = None
     num_layers: int | None = None
     overwrite: bool = False
+    # iOS only. When True, embedding table is not quantized to int8.
+    disable_embedding_quantization: bool = False
     # Optional prebuilt coreai-opt config (KMeansPalettizerConfig or
     # QuantizerConfig) loaded from a user-provided YAML. When set, the pipeline
     # uses this directly and ignores `compression` for config resolution
@@ -210,6 +212,7 @@ async def _async_export_model(config: ExportConfig) -> str:
                 max_context_length=max_context_length,
                 target_dtype=target_dtype,
                 num_layers=config.num_layers,
+                disable_embedding_quantization=config.disable_embedding_quantization,
             )
         model = model.eval()
         # ---- 3. Resolve compression preset ----

python/src/coreai_models/llm/export.py

@@ -144,6 +144,15 @@ def build_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="Allow exporting models without a registry preset. Requires --compute-precision.",
     )
+    parser.add_argument(
+        "--disable-embedding-quantization-ios",
+        action="store_true",
+        help=(
+            "iOS only. Skip int8 quantization of the embedding table and keep it in "
+            "float32. Default: False (embedding is quantized). Rejected when "
+            "--platform is macOS."
+        ),
+    )
     return parser
 
 
@@ -314,6 +323,11 @@ def _resolve_export_config(args: argparse.Namespace) -> ExportConfig:
             "different quantization options and kv-cache limits."
         )
 
+    if args.disable_embedding_quantization_ios and variant != "iOS":
+        raise SystemExit(
+            f"--disable-embedding-quantization-ios requires --platform iOS (got '{variant}')."
+        )
+
     if args.compression_config is not None:
         if not args.compression_config.is_file():
             raise SystemExit(f"--compression-config: file not found: {args.compression_config}")
@@ -350,6 +364,7 @@ def _resolve_export_config(args: argparse.Namespace) -> ExportConfig:
         num_layers=args.num_layers,
         overwrite=args.overwrite,
         compression_config_object=compression_config_object,
+        disable_embedding_quantization=args.disable_embedding_quantization_ios,
     )
 
 
@@ -413,6 +428,8 @@ def main() -> None:
         if config.num_layers:
             print(f"  num_layers:         {config.num_layers}")
         print(f"  overwrite:          {config.overwrite}")
+        if config.variant == "iOS":
+            print(f"  disable_embedding_quantization: {config.disable_embedding_quantization}")
         return
 
     result = export_model(config)

python/src/coreai_models/models/base.py

@@ -340,6 +340,7 @@ def from_hf(
         target_dtype: torch.dtype = torch.float16,
         mmap_path: str | None = None,
         num_layers: int | None = None,
+        disable_embedding_quantization: bool = False,
     ) -> T:
         """Load model from HuggingFace model hub.
 
@@ -353,6 +354,9 @@ def from_hf(
             num_layers: Optional number of transformer layers. When set, only layers
                         0..num_layers-1 are loaded and the config is truncated.
                         Useful for fast smoke tests.
+            disable_embedding_quantization: iOS only. When True, the
+                embedding table is not quantized to int8.
+                Ignored for macOS model classes.
 
         Returns:
             Instance of the model class loaded with HuggingFace weights
@@ -370,8 +374,12 @@ def from_hf(
             hf_model.config, max_context_length, num_layers=num_layers
         )
 
-        # Create our model instance and load the state dict
-        model = cls(config, model_device="meta")
+        # Create our model instance and load the state dict.
+        # disable_embedding_quantization is only accepted by the iOS base class.
+        init_kwargs: dict = {"config": config, "model_device": "meta"}
+        if issubclass(cls, BaseForCausalLMForiOS):
+            init_kwargs["disable_embedding_quantization"] = disable_embedding_quantization
+        model = cls(**init_kwargs)
         model.to(dtype=target_dtype)
         state_dict = hf_model.state_dict()
         if not isinstance(state_dict, collections.abc.MutableMapping):
@@ -414,6 +422,7 @@ def from_hf_memory_efficient(
         num_layers: int | None = None,
         hf_config_attr: str | None = None,
         hf_state_dict_prefix: str = "",
+        disable_embedding_quantization: bool = False,
     ) -> T:
         """Load model from HuggingFace with layer-by-layer memory offloading.
 
@@ -439,6 +448,9 @@ def from_hf_memory_efficient(
                 prefix are loaded. The prefix is stripped before assigning.
                 Use for multimodal checkpoints where text weights live under
                 a prefix (e.g. ``"language_model."``).
+            disable_embedding_quantization: iOS only. When True, the
+                embedding table is not quantized to int8.
+                Ignored for non-iOS model classes.
         """
         model_dir = snapshot_download(
             huggingface_model_id,
@@ -450,7 +462,11 @@ def from_hf_memory_efficient(
 
         config = cls._get_reauthored_config(hf_config, max_context_length, num_layers=num_layers)
 
-        model = cls(config, model_device="meta")
+        # disable_embedding_quantization is only accepted by the iOS base class.
+        init_kwargs: dict = {"config": config, "model_device": "meta"}
+        if issubclass(cls, BaseForCausalLMForiOS):
+            init_kwargs["disable_embedding_quantization"] = disable_embedding_quantization
+        model = cls(**init_kwargs)
         model.to(dtype=target_dtype)
 
         safetensors_files = _resolve_safetensors_files(model_dir)

python/src/coreai_models/models/macos/mixtral.py

@@ -38,10 +38,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         active_experts_scores = torch.softmax(top_logits, dim=-1).to(x.dtype)
 
         y_active_experts = self.switch_mlp(x, active_experts_indices)
-        active_experts_scores = active_experts_scores.unsqueeze(-1)
+        active_experts_scores = active_experts_scores.unsqueeze(-1).to(y_active_experts.device)
         y_active_experts_weighted_by_scores = y_active_experts * active_experts_scores
         y_active_experts_summary = torch.sum(y_active_experts_weighted_by_scores, dim=-2)
-        return y_active_experts_summary.to(x.dtype)
+        return y_active_experts_summary.to(device=x.device, dtype=x.dtype)
 
 
 class Attention(nn.Module):

python/src/coreai_models/primitives/ios/cache.py

@@ -95,15 +95,23 @@ def __init__(self: Self, n_layers: int, hidden_size: int):
     def gen_slice_args(
         self, layer_idx: int, offset: torch.IntTensor, num_token_updates: int
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        layer_index = self._layer_indices[layer_idx]
-        layer_index_end = self._layer_indices_end[layer_idx]
-        begin = torch.cat([layer_index, self._zero, self._zero, self._zero, offset])
+        layer_index = self._layer_indices[layer_idx].to(offset.device)
+        layer_index_end = self._layer_indices_end[layer_idx].to(offset.device)
+        begin = torch.cat(
+            [
+                layer_index,
+                self._zero.to(offset.device),
+                self._zero.to(offset.device),
+                self._zero.to(offset.device),
+                offset,
+            ]
+        )
         end = torch.cat(
             [
                 layer_index_end,
-                self._one,
-                self._hidden_size,
-                self._one,
+                self._one.to(offset.device),
+                self._hidden_size.to(offset.device),
+                self._one.to(offset.device),
                 offset + num_token_updates,
             ]
         )

python/src/coreai_models/primitives/ios/rope.py

@@ -72,17 +72,17 @@ def __init__(
         self._compute_sin_and_cos()
 
     def _apply(self, fn):
-        # the `.to()` function implicitly calls into this function,
-        # and we should recompute the cos / sin rather then just do
+        # The `.to()` function implicitly calls into this function,
+        # and we should recompute the cos / sin rather than just do
         # a simple cast.
         super()._apply(fn)
-        dummy = torch.tensor(0.0)
-        transformed = fn(dummy)
-        self._compute_sin_and_cos(transformed.dtype)
-
-        target_device = transformed.device
-        self.cos_cached = self.cos_cached.to(device=target_device)
-        self.sin_cached = self.sin_cached.to(device=target_device)
+        # Read dtype/device from the buffer post-apply so device-only
+        # and dtype-only .to(...) calls are both honored.
+        new_dtype = self.cos_cached.dtype
+        new_device = self.cos_cached.device
+        self._compute_sin_and_cos(new_dtype)
+        self.cos_cached = self.cos_cached.to(new_device)
+        self.sin_cached = self.sin_cached.to(new_device)
         return self
 
     def _compute_sin_and_cos(self, dtype: torch.dtype = torch.float32) -> None:

python/src/coreai_models/primitives/ios/sdpa.py

@@ -3,8 +3,11 @@
 # Use of this source code is governed by a BSD-3-clause license that can
 # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
+import os
+
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 
 class SDPA(nn.Module):
@@ -31,6 +34,7 @@ def __init__(
                     if isinstance(scale, torch.Tensor)
                     else nn.Buffer(torch.tensor(scale), persistent=False)
                 )
+        self._use_hf_impl = os.environ.get("USE_HF_IMPL", "").lower() == "true"
 
     # Efficient implementation equivalent to the following:
     def forward(
@@ -52,6 +56,47 @@ def forward(
             torch.Tensor: Attention output with shape (batch_size, n_heads*head_dim, 1, seq_len)
         """
 
+        # use FlashAttention to avoid
+        # materializing the full attention score matrix.
+        # Trim K/V from max_pos to seq_len (cache positions beyond seq_len
+        # are zeros masked by -inf) so we can use is_causal=True, which is
+        # required for the FlashAttention kernel.
+        if query.is_cuda and self._use_hf_impl:
+            B, _, _, S = query.shape
+            n_heads = query.shape[1] // self.head_dim
+            n_kv_heads = key.shape[1] // self.head_dim
+
+            # This path is currently prefill-only: it trims K/V to the first S positions
+            # and relies on is_causal=True. In prefill every valid KV position
+            # lies within [0, S)- a valid (unmasked, == 0) entry at KV index
+            # >= S means this is an extend/decode call, which the trim and
+            # is_causal=True below would silently mishandle.
+            assert not (causal_mask[:, S:] == 0).any(), (
+                "CUDA/HF SDPA path is prefill-only. Got a causal_mask with "
+                "valid KV positions beyond query length S (extend/decode not "
+                "supported)."
+            )
+
+            q = query.reshape(B, n_heads, self.head_dim, S).transpose(2, 3).contiguous()
+            k = key[..., :S].reshape(B, n_kv_heads, self.head_dim, S).transpose(2, 3).contiguous()
+            v = value[..., :S].reshape(B, n_kv_heads, self.head_dim, S).transpose(2, 3).contiguous()
+
+            out = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                # is_causal=True is required by the FlashAttention kernel we
+                # target here, so causal_mask is intentionally not passed as
+                # attn_mask. This path is only taken for prefill, where the
+                # mask is guaranteed causal (asserted above), so is_causal=True
+                # and the provided causal_mask are equivalent.
+                is_causal=True,
+                scale=self._scale_factor,
+                enable_gqa=(n_kv_heads != n_heads),
+            )
+
+            return out.transpose(2, 3).reshape(B, n_heads * self.head_dim, 1, S)
+
         # Apply the scale factor before QK^T for numerical stability
         key = key.transpose(-3, -1) * self._scale_factor
         queries = query.split(self.head_dim, dim=1)

python/src/coreai_models/primitives/macos/cache.py

@@ -108,6 +108,12 @@ def update_and_fetch(
         torch._check_is_size(seq_len)
         device = self._k_cache.device
 
+        compute_device = k.device
+        cross_device = compute_device != device
+        if cross_device:
+            k = k.to(device)
+            v = v.to(device)
+
         layer_index = torch.tensor((layer_idx,), dtype=torch.int32, device=device)
         layer_index_end = torch.tensor((layer_idx + 1,), dtype=torch.int32, device=device)
 
@@ -162,7 +168,11 @@ def update_and_fetch(
         # return the slice k, v
         k = self._k_cache.narrow(0, layer_idx, 1).narrow(-2, 0, seq_len)
         v = self._v_cache.narrow(0, layer_idx, 1).narrow(-2, 0, seq_len)
-        return k.squeeze(0), v.squeeze(0)
+        k_out = k.squeeze(0)
+        v_out = v.squeeze(0)
+        if cross_device:
+            return k_out.to(compute_device), v_out.to(compute_device)
+        return k_out, v_out
 
 
 class SSMState:

python/src/coreai_models/primitives/macos/rope.py

@@ -109,7 +109,7 @@ def forward(
         return self._rope(
             x,
             position_ids=position_ids,
-            freqs=self._freqs,
+            freqs=self._freqs.to(x.device),
             offset=offset,
         )
 

python/src/coreai_models/primitives/macos/switch.py

@@ -78,6 +78,11 @@ def __init__(
         self.up_proj = SwitchLinear(hidden_size, moe_intermediate_size, 1, num_experts, bias=bias)
         self.down_proj = SwitchLinear(moe_intermediate_size, hidden_size, 1, num_experts, bias=bias)
         self._activate = activation if activation is not None else SwiGLU()
+        # Eager-only optimization. When set, tokens are processed in chunks of
+        # this size to bound the peak GatherMM intermediate. Left None for
+        # export/production so the traced
+        # graph carries no data-dependent control flow on the token dimension.
+        self.eager_chunk_size: int | None = None
 
     def forward(
         self: Self,
@@ -90,12 +95,27 @@ def forward(
         x = x.reshape((-1, 1, 1, hidden_size))
         # batch size mul query length x num active experts
         indices = indices.reshape((-1, num_active_experts))
-        # batch size mul query length x num active experts x 1 x moe intermediate size
-        gate = self.gate_proj(x, indices)
-        up = self.up_proj(x, indices)
-        gated_up = self._activate(up, gate)
-        # batch size mul query length x num active experts x 1 x hidden size
-        x = self.down_proj(gated_up, indices)
+        bsql = x.shape[0]
+
+        chunk_size = self.eager_chunk_size
+        if chunk_size is None or bsql <= chunk_size:
+            gate = self.gate_proj(x, indices)
+            up = self.up_proj(x, indices)
+            gated_up = self._activate(up, gate)
+            # nws x bsql x nae x 1 x hidden size
+            x = self.down_proj(gated_up, indices)
+        else:
+            chunks = []
+            for start in range(0, bsql, chunk_size):
+                x_c = x[start : start + chunk_size]
+                idx_c = indices[start : start + chunk_size]
+                gate_c = self.gate_proj(x_c, idx_c)
+                up_c = self.up_proj(x_c, idx_c)
+                gated_c = self._activate(up_c, gate_c)
+                chunks.append(self.down_proj(gated_c, idx_c))
+            # nws x bsql x nae x 1 x hidden size
+            x = torch.cat(chunks, dim=1)
+
         # batch size x query length x num active experts x hidden size
         x = x.reshape((batch_size, query_length, num_active_experts, hidden_size))
         return x