Command A+: batch image tiles across requests in the vision encoder

mstar/model/command_a_plus/components/vision.py

@@ -8,9 +8,12 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 from mstar.model.command_a_plus.config import CommandAPlusVisionConfig
 
+_ATTENTION_BACKENDS = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
+
 
 class SiglipAttention(nn.Module):
     def __init__(self, config: CommandAPlusVisionConfig) -> None:
@@ -64,8 +67,13 @@ def __init__(self, config: CommandAPlusVisionConfig) -> None:
         self.position_embedding = nn.Embedding((config.image_size // config.patch_size) ** 2, config.hidden_size)
 
     def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
-        patches = self.patch_embedding(pixel_values.to(self.patch_embedding.weight.dtype))
-        return patches.flatten(2).transpose(1, 2) + self.position_embedding.weight
+        # The conv as a matmul: cuDNN picks its algorithm by batch size, which would
+        # make a tile's embeddings depend on the other tiles in the call.
+        weight = self.patch_embedding.weight
+        p = weight.shape[-1]
+        patches = pixel_values.to(weight.dtype).unfold(2, p, p).unfold(3, p, p)
+        patches = patches.permute(0, 2, 3, 1, 4, 5).flatten(3).flatten(1, 2)
+        return F.linear(patches, weight.flatten(1), self.patch_embedding.bias) + self.position_embedding.weight
 
 
 class CommandAPlusProjector(nn.Module):
@@ -113,8 +121,10 @@ def __init__(self, config: CommandAPlusVisionConfig, text_hidden_size: int) -> N
     def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
         """``[tiles, 3, size, size]`` to ``[tiles * tokens_per_tile, text_hidden]``, tile-major."""
         x = self.embeddings(pixel_values)
-        for layer in self.layers:
-            x = layer(x)
+        # cuDNN attention builds a plan per new tile count, ~0.5 s each on GB200
+        with sdpa_kernel(_ATTENTION_BACKENDS):
+            for layer in self.layers:
+                x = layer(x)
         return self.projector(self.post_layernorm(x)).flatten(0, 1)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:

mstar/model/command_a_plus/submodules.py

@@ -38,10 +38,14 @@
 
 
 class CommandAPlusVisionSubmodule(NodeSubmodule):
-    """Turns a request's image tiles into the embeddings for its ``<|IMG_PATCH|>`` tokens."""
+    """Turns each request's image tiles into the embeddings for its ``<|IMG_PATCH|>`` tokens."""
 
     disable_torch_compile = True
 
+    # On GB200 a tile costs ~1.3 ms from 64 tiles up and ~26 MiB of activations,
+    # so larger calls only grow the peak; a step's tile count is otherwise unbounded.
+    TILES_PER_CALL = 64
+
     def __init__(self, encoder: CommandAPlusVisionEncoder):
         super().__init__()
         self.encoder = encoder
@@ -50,9 +54,27 @@ def prepare_inputs(self, graph_walk, fwd_info, inputs, **kwargs) -> NodeInputs:
         tiles = torch.cat(inputs["pixel_values"]).to(self.get_device())
         return NodeInputs(tensor_inputs={"pixel_values": tiles})
 
+    def preprocess(self, graph_walk, engine_inputs, inputs) -> dict[str, Any]:
+        tiles = [inp.tensor_inputs["pixel_values"] for inp in inputs]
+        return {"pixel_values": torch.cat(tiles), "tile_counts": [len(t) for t in tiles]}
+
     def forward(self, graph_walk, engine_inputs, pixel_values, **kwargs) -> NameToTensorList:
         return {"image_embeds": [self.encoder(pixel_values)]}
 
+    def can_batch(self, batch, model_inputs) -> bool:
+        return True
+
+    def forward_batched(
+        self, graph_walk, engine_inputs, pixel_values, tile_counts, **kwargs,
+    ) -> dict[str, NameToTensorList]:
+        tokens = self.encoder.config.tokens_per_tile
+        embeds = torch.cat([self.encoder(chunk) for chunk in pixel_values.split(self.TILES_PER_CALL)])
+        embeds = embeds.split([n * tokens for n in tile_counts])
+        return {
+            rid: {"image_embeds": [e]}
+            for rid, e in zip(engine_inputs.request_ids, embeds, strict=True)
+        }
+
 
 class CommandAPlusLLMSubmodule(ARNodeSubmodule):
     # Decode replays CUDA graphs of the same eager kernels; torch.compile stays

mstar/model/command_a_plus/tests/test_integration.py

@@ -28,7 +28,7 @@
     SAMPLER,
     CommandAPlusConfig,
 )
-from mstar.model.command_a_plus.submodules import CommandAPlusLLMSubmodule
+from mstar.model.command_a_plus.submodules import CommandAPlusLLMSubmodule, CommandAPlusVisionSubmodule
 from mstar.model.command_a_plus.tests.test_backbone import bind_test_resources
 from mstar.model.command_a_plus.tests.test_weight_loading import checkpoint
 from mstar.model.multimodal import PromptPart
@@ -257,6 +257,26 @@ def test_vision_node_loads_from_the_checkpoint_and_encodes_tiles(self):
             outputs = node.forward("prefill_image", None, **node.preprocess("prefill_image", None, [inputs]))
         self.assertEqual(tuple(outputs["image_embeds"][0].shape), (3 * 4, self.model.config.hidden_size))
 
+    def test_vision_node_batches_requests_and_hands_each_its_own_embeddings(self):
+        from mstar.model.command_a_plus.tests.test_vision import tiny_encoder
+
+        torch.manual_seed(0)
+        node = CommandAPlusVisionSubmodule(tiny_encoder()).eval()
+        node.TILES_PER_CALL = 2  # a chunk straddles the two requests
+        tiles = {"b": torch.rand(3, 3, 32, 32), "a": torch.rand(1, 3, 32, 32)}
+        inputs = [node.prepare_inputs("prefill_image", info("prefill_image"), {"pixel_values": [t]})
+                  for t in tiles.values()]
+        engine = ModelInputsFromEngine(list(tiles), {})
+        self.assertTrue(node.can_batch(None, inputs))
+        with torch.inference_mode():
+            batched = node.forward_batched("prefill_image", engine, **node.preprocess("prefill_image", engine, inputs))
+            alone = [node.forward("prefill_image", None, **node.preprocess("prefill_image", None, [inp]))
+                     for inp in inputs]
+        self.assertEqual(list(batched), ["b", "a"])
+        for rid, single in zip(tiles, alone, strict=True):
+            self.assertEqual(len(batched[rid]["image_embeds"][0]), len(tiles[rid]) * 4)
+            torch.testing.assert_close(batched[rid]["image_embeds"][0], single["image_embeds"][0])
+
     def test_load_dtype_and_reject_sequence_parallelism(self):
         save_file(checkpoint(self.model.config), Path(self.directory) / "model.safetensors")
         node = self.model.get_submodule("LLM", autocast_dtype=torch.bfloat16)

mstar/model/command_a_plus/tests/test_vision.py

@@ -41,6 +41,13 @@ def test_loads_checkpoint_names_strictly(self):
         with self.assertRaisesRegex(RuntimeError, "Missing"):
             tiny_encoder().load_weights(list(weights.items())[1:])
 
+    def test_patch_embedding_matches_the_conv(self):
+        torch.manual_seed(0)
+        embeddings = tiny_encoder().embeddings
+        pixels = torch.rand(2, 3, 32, 32)
+        conv = embeddings.patch_embedding(pixels).flatten(2).transpose(1, 2)
+        torch.testing.assert_close(embeddings(pixels), conv + embeddings.position_embedding.weight)
+
 
 @unittest.skipUnless(os.environ.get("MSTAR_TEST_HF") == "1", "set MSTAR_TEST_HF=1")
 class HuggingFaceVisionParityTests(unittest.TestCase):