Issue 381 Fix: Add attention mask to SAM3 export and runtime

models/sam3/README.md

@@ -4,11 +4,13 @@ SAM 3 (Segment Anything Model 3) is a unified vision model from Meta for prompta
 
 This export targets **iOS** by restructuring the model in a BC1S layout (channels-first, `Conv2d(1x1)` projections, fp16-safe primitives, rank-4 window attention) and splitting it into three independently optimizable functions:
 
-| Function       | Compression                                | Inputs                               | Outputs                                                      |
-|----------------|--------------------------------------------|--------------------------------------|--------------------------------------------------------------|
-| `image_encode` | 4-bit k-means palettization (gs=32) + fp16 | `pixel_values`                       | `backbone_features`                                          |
-| `text_encode`  | 6-bit k-means palettization (gs=8) + fp16  | `input_ids`                          | `text_features`                                              |
-| `detect`       | fp16 (no weight compression)               | `backbone_features`, `text_features` | `pred_masks`, `pred_boxes`, `pred_logits`, `presence_logits` |
+| Function       | Compression                                | Inputs                                                 | Outputs                                                                      |
+|----------------|--------------------------------------------|--------------------------------------------------------|------------------------------------------------------------------------------|
+| `image_encode` | 4-bit k-means palettization (gs=32) + fp16 | `pixel_values`                                         | `backbone_features`                                                          |
+| `text_encode`  | 6-bit k-means palettization (gs=8) + fp16  | `input_ids`                                            | `text_features`                                                              |
+| `detect`       | fp16 (no weight compression)               | `backbone_features`, `text_features`, `attention_mask` | `pred_masks`, `pred_boxes`, `pred_logits`, `presence_logits`, `semantic_seg` |
+
+`attention_mask` is `[1, 1, 1, S]` (`S` = `--max-text-seq-len`, default 32), 1 for real prompt tokens and 0 for padding, so `detect` ignores the padding the way HF `Sam3Model` does. `ImageSegmenter` fills it from the token ids.
 
 ## Setup
 
@@ -70,6 +72,8 @@ uv run models/sam3/export.py --full --dtype float16
 
 Full bundles land at `<repo-root>/exports/<model>_<dtype>/` (e.g. `exports/sam3_float32/`), so they sit next to the lite `sam3_lite_336_w4_static/` without colliding.
 
+`main` takes `pixel_values`, `input_ids` `[1, 32]` and `attention_mask` `[1, 32]` (int32, 1 for real prompt tokens and 0 for padding); `ImageSegmenter` fills the mask from the token ids.
+
 ## Running
 
 ### In your iOS and macOS applications

python/src/coreai_models/models/ios/sam3/detr.py

@@ -544,7 +544,7 @@ def forward(
         text_features: torch.Tensor,
         vision_pos: torch.Tensor,
         spatial_shapes: torch.Tensor | None = None,
-        text_mask: torch.Tensor | None = None,
+        text_cross_attn_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         batch_size = vision_features.shape[0]
 
@@ -557,15 +557,6 @@ def forward(
         hidden_bsc = torch.cat([presence_token, query_embeds], dim=1)
         hidden_states = hidden_bsc.permute(0, 2, 1).unsqueeze(2)
 
-        text_cross_attn_mask = None
-        if text_mask is not None:
-            inv_mask = (~text_mask).to(dtype=vision_features.dtype) * (-40000.0)
-            text_cross_attn_mask = (
-                inv_mask.unsqueeze(1)
-                .unsqueeze(1)
-                .expand(-1, self.num_heads, hidden_states.shape[3], -1)
-            )
-
         normed_hs = None
         new_ref = None
         presence_logits = None

python/src/coreai_models/models/ios/sam3/mask_decoder.py

@@ -186,11 +186,16 @@ def __init__(
 
         self.sdpa = BidirectionalSDPA(num_heads=num_heads, head_dim=self.head_dim)
 
-    def forward(self, query: torch.Tensor, key_value: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        query: torch.Tensor,
+        key_value: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         q = self.q_proj(query)
         k = self.k_proj(key_value)
         v = self.v_proj(key_value)
-        return self.o_proj(self.sdpa(q, k, v))
+        return self.o_proj(self.sdpa(q, k, v, attention_mask=attention_mask))
 
     @classmethod
     def from_hf_attention(cls, hf_attn: nn.Module) -> "CrossAttentionReauthored":
@@ -231,11 +236,14 @@ def forward(
         backbone_features: list[torch.Tensor],
         encoder_hidden_states: torch.Tensor,
         prompt_features: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         if prompt_features is not None:
             residual = encoder_hidden_states
             normed = self.prompt_cross_attn_norm(encoder_hidden_states)
-            attn_out = self.prompt_cross_attn(normed, prompt_features)
+            attn_out = self.prompt_cross_attn(
+                normed, prompt_features, attention_mask=prompt_attention_mask
+            )
             encoder_hidden_states = residual + attn_out
 
         pixel_embed = self._embed_pixels(backbone_features, encoder_hidden_states)
@@ -329,14 +337,21 @@ def forward(
         self,
         decoder_hidden_states: torch.Tensor,
         text_features: torch.Tensor,
+        text_valid: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        """``text_valid`` is ``(B, 1, 1, seq)``, 1 for real tokens and 0 for padding."""
         orig_text = text_features
         h = self.text_mlp_layer1(text_features)
         h = F.relu(h)
         h = self.text_mlp_layer2(h)
         text_features = self.text_mlp_out_norm(h + orig_text)
 
-        pooled_text = text_features.mean(dim=3, keepdim=True)
+        if text_valid is None:
+            pooled_text = text_features.mean(dim=3, keepdim=True)
+        else:
+            one = torch.tensor(1.0, dtype=text_features.dtype, device=text_features.device)
+            num_valid = text_valid.sum(dim=3, keepdim=True).clamp(min=one)
+            pooled_text = (text_features * text_valid).sum(dim=3, keepdim=True) / num_valid
         proj_text = self.text_proj(pooled_text)
 
         scale = self._scale.to(dtype=decoder_hidden_states.dtype)

python/src/coreai_models/models/ios/sam3/sam3_reauthored.py

@@ -12,7 +12,7 @@
 Usage:
     model = Sam3Lite.from_pretrained("facebook/sam3", image_size=336)
     pred_masks, pred_boxes, pred_logits, presence_logits, semantic_seg =
-     model(pixel_values, input_ids)
+     model(pixel_values, input_ids, attention_mask)
 """
 
 import torch
@@ -42,6 +42,19 @@ def _linear_to_conv2d(linear: nn.Linear) -> nn.Conv2d:
     return conv
 
 
+def text_mask_bias(
+    attention_mask: torch.Tensor, dtype: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Turn a ``(B, seq)`` or ``(B, 1, 1, seq)`` text mask (1 = token, 0 = padding) into
+    ``(bias, valid)``, both ``(B, 1, 1, seq)``: an additive attention bias that is
+    ``-40000.0`` on padding, and the 0/1 weights for the scoring head's masked mean.
+    """
+    valid = attention_mask.reshape(attention_mask.shape[0], 1, 1, -1).to(dtype)
+    one = torch.tensor(1.0, dtype=dtype, device=valid.device)
+    neg = torch.tensor(40000.0, dtype=dtype, device=valid.device)
+    return (valid - one) * neg, valid
+
+
 class Sam3Lite(nn.Module):
     """SAM3 Lite model for iOS (palettized).
 
@@ -51,6 +64,9 @@ class Sam3Lite(nn.Module):
     Inputs:
         pixel_values: (B, 3, image_size, image_size) float tensor.
         input_ids: (B, seq_len) int32 token IDs.
+        attention_mask: optional (B, seq_len) mask, 1 for real tokens and 0 for
+            padding, as returned by the tokenizer. Without it the padding slots
+            are attended to, which changes the scores and masks (HF masks them).
 
     Outputs:
         Tuple ``(pred_masks, pred_boxes, pred_logits, presence_logits, semantic_seg)``.
@@ -81,6 +97,7 @@ def forward(
         self,
         pixel_values: torch.Tensor,
         input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         backbone_features = self.image_encoder(pixel_values)
 
@@ -91,6 +108,10 @@ def forward(
         text_hidden_states = self.text_encoder(input_ids)
         text_features = self.text_projection(text_hidden_states)
 
+        text_bias = text_valid = None
+        if attention_mask is not None:
+            text_bias, text_valid = text_mask_bias(attention_mask, text_features.dtype)
+
         vision_level2 = fpn_hidden_states_trimmed[-1]
         B = vision_level2.shape[0]
         vision_bc1s = vision_level2.reshape(B, 256, 1, -1)
@@ -102,19 +123,22 @@ def forward(
             vision_feats=vision_bc1s,
             text_feats=text_features,
             vision_pos=pos_bc1s,
+            text_cross_attn_mask=text_bias,
         )
 
         final_hidden_states, pred_boxes, presence_logits = self.detr_decoder(
             vision_features=encoder_output,
             text_features=text_features,
             vision_pos=pos_bc1s,
             spatial_shapes=self.spatial_shapes,
+            text_cross_attn_mask=text_bias,
         )
 
         pred_logits = (
             self.scoring(
                 decoder_hidden_states=final_hidden_states.unsqueeze(0),
                 text_features=text_features,
+                text_valid=text_valid,
             )
             .squeeze(-1)
             .squeeze(0)
@@ -125,6 +149,7 @@ def forward(
             backbone_features=list(fpn_hidden_states_trimmed),
             encoder_hidden_states=encoder_output,
             prompt_features=text_features,
+            prompt_attention_mask=text_bias,
         )
         pred_masks = mask_outputs["pred_masks"]
         semantic_seg = mask_outputs["semantic_seg"]

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

@@ -9,7 +9,8 @@
 
 Mask convention: ``attention_mask`` (when provided) is in GPU layout
 (B, num_heads, query_seq, key_seq), with ``-40000.0`` for masked
-positions. ``None`` is unmasked attention.
+positions. Size-1 head / query dims broadcast (a key-padding mask is
+``(B, 1, 1, key_seq)``). ``None`` is unmasked attention.
 """
 
 import torch
@@ -60,7 +61,10 @@ def _chunked_forward(
         for start in range(0, query_seq, chunk_size):
             end = min(start + chunk_size, query_seq)
             q_chunk = query[:, :, :, start:end]
-            mask_chunk = attention_mask[:, :, start:end, :] if attention_mask is not None else None
+            # A mask broadcast over queries (query dim of 1) applies to every chunk as is.
+            mask_chunk = attention_mask
+            if attention_mask is not None and attention_mask.shape[2] > 1:
+                mask_chunk = attention_mask[:, :, start:end, :]
             chunks.append(self._standard_forward(q_chunk, key, value, mask_chunk))
         return torch.cat(chunks, dim=3)
 

python/src/coreai_models/segmentation/export.py

@@ -353,6 +353,7 @@ def main() -> None:
                     "hf_model_id",
                     "image_size",
                     "dtype",
+                    "max_text_seq_len",
                     "output_dir",
                     "output_name",
                     "overwrite",