[rohansanda]

Adds image input to Command A+ (#379).

• Encoder. The checkpoint's SigLIP tower, pixel shuffle and SwiGLU projector (components/vision.py), loaded strictly from the vision weights. It runs as a vision_encoder node on rank 0, as Qwen3.5 does at TP4.
• prefill_image walk. The encoder runs first, then the LLM's prepare_inputs writes its output over the <|IMG_PATCH|> embeddings, so chunked prefill splits both together. Text prompts are unchanged.
• Prompts. Images go through the checkpoint's image processor, and each placeholder widens to its tiles' tokens. A prompt that's too long for max_seq_len is refused at the first image that doesn't fit. Images come from /generate files, in the order sent, or from {"type": "image"} items in messages.

Validation (4x GB200, TP4, full checkpoint)

• Token IDs and pixel values are identical to Cohere2VisionProcessor for one image and for two images (one of them five tiles).
• validate_checkpoint.py --images (new flag) checks three text and three image prompts x 32 steps against the HF reference. The greedy token matches at 191 of 192 steps. The one miss is a near tie on the 1151-token image prompt, where HF's top two logits are 0.25 apart. The minimum cosine similarity is 0.996, and the maximum RMSE is 0.26 for images against 0.20 for text.
• Through /generate, single-image, multi-image and mixed image/text concurrent requests all work. A 32-request soak at concurrency 16 returned no errors. Warm time to first token is 40 ms for text, 60 ms with a one-tile image and 83 ms with a five-tile image.

Not included: an OpenAI adapter, CUDA graphs for the encoder, caching encoder output across turns, and batching images across requests. Without that, image TTFT at 16 concurrent requests is 400-520 ms vs 126 ms for text. Batching is the next follow-up.

[mstar-review]

Automated review

🟡 No findings on 0c7c789, light review

Description

• The description doesn't say whether ruff and the CPU Core suite were run, and doesn't give the result for the new and changed tests (test_vision.py, test_integration.py, test_registration.py, test_config.py).
• The GB200 validation doesn't name the torch/CUDA build, as AGENTS.md's Testing section asks for GPU runs.
• The description doesn't say how the TTFT numbers (40/60/83 ms) and the 32-request soak at concurrency 16 were produced: which harness, what flags or commands, and how many warmup and measured runs.
• One behavior change isn't mentioned: a text-only prompt or pre-tokenized text_inputs that contains the <|IMG_PATCH|> token is now rejected with a placeholder-count error, where before it was accepted.

What the review checked

• The TP LLM's prepare_inputs embeds tokens with no collective: after fuse_for_inference on CUDA, embed gathers from the replicated table (language_model.py:327), so building input_embeds per rank in prepare_inputs can't break lockstep (invariant 11).
• Encoder on rank 0 feeding the TP4/TP8 LLM uses the same node_groups layout as configs/qwen3_5_27b_tp4.yaml, and the edge reaches every LLM rank the same way.
• Building input_embeds with masked_scatter in prepare_inputs follows the existing Qwen3.5 and Cosmos3 pattern.
• Every walk check in command_a_plus outside the tests now accepts prefill_image where it accepted prefill: _check_walk, declare_step prefill tracking, _sample last-hidden selection, postprocess, check_stop and get_partition_forward_pass_args.
• Prefix caching can't key image prompts by placeholder ids alone: Command A+ declares no prefix_key_streams, so _open_prefix_caches opens no keyed walk for this model.
• _expand_images replaces each placeholder with START + tiles x (tokens_per_tile + line break) + END and checks the placeholder count against the image count before it touches the tokenizer or the processor.
• get_initial_forward_pass_args sends pixel_values to vision_encoder and unpersists the tensors of both input edges. The image walk goes to decode the same way the text walk does.
• Vision weight loading is strict: load_state_dict with the prefix remapping, and the shard-availability check uses the vision prefixes when building the encoder node.
• The image API key image_inputs matches the one Qwen3.5 already uses.

_{Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.}

[arun-elr]

looks clean, the chunked-prefill cut and the TP layout follow the Qwen3.5 pattern. a few small things:

• CI only ran label here since the base is #379's branch. I ran the PR's tests on CPU at 6503923: command_a_plus/tests + test_model_kv_layouts.py 90 passed, 17 skipped (CUDA kernels, HF reference), ruff clean. worth a CI rerun once #379 lands.
• CommandAPlusVisionSubmodule keeps the default can_batch() -> False, so concurrent image prompts go through the encoder one at a time. the 40/60/83 ms TTFTs are warm single requests; what does image TTFT look like at c16? fine as a follow-up.
• _expand_images tiles every image before the max_seq_len check, so an over-long multi-image prompt does all the processor work before it's refused. could check the token count per image as it goes.
• nit: tiles are always cast to bf16 (L299), so an fp32 run gets bf16 pixels. probably fine, just flagging for the reference checks.

[rohansanda]

Thanks for taking a look @arun-elr!

Good call on the encoder. I measured it on 4x GB200 (median TTFT in ms, 16 output tokens):

	c1	c4	c16
text	41	53	126
1-tile image	64	150	396
4-tile image	79	182	518

So images do queue up at c16. The tiles are all the same size, so batching them across requests should be pretty simple. I'll do that in a follow-up.

Fixed the length check. It now refuses the prompt as soon as an image doesn't fit, instead of tiling all of them first, and there's a test for it.

The bf16 cast is on purpose. Serving runs in bf16 anyway, and it halves what gets sent to the encoder. The HF parity check casts its pixels the same way, and the fp32 encoder test feeds the tower directly, so the reference checks aren't affected.

CI should run now that this targets main. Locally, CPU Core passes (6782) and the GPU tests pass on GB200 (869).

[arun-elr]

thanks for the numbers. the length check looks right, and cross-request batching as a follow-up sounds good. lgtm

[NSagan271]

LGTM, one nitpick + I have some live tests still running