[rohansanda]

Follow-up to #838 suggested by @arun-elr. The vision node encoded one request's tiles per call, so concurrent image requests waited on each other. Now one call takes the tiles of every request in the step, encodes them in chunks of 64, and splits the output back per request. 64 is about where the per-tile cost stops dropping on GB200, so the cap just bounds activation memory (~1.6 GB).

For that to be safe, a tile has to get the same embeddings no matter what else is in the call, and it didn't:

• cuDNN picks the patch conv's algorithm by batch size, which changed embeddings enough to flip greedy replies. The patch embedding is now the same weights applied as a matmul.
• cuDNN attention built a new plan for every new tile count, about 0.5 s each on GB200. The tower skips it now.

With both changes, batched and per-request embeddings are bitwise equal on GB200.

Median TTFT in ms, 4x GB200 TP4, 16 output tokens, before → after:

request	c1	c4	c16	c32
text	41 → 41	80 → 78	125 → 125	155 → 161
1-tile image	62 → 61	151 → 137	415 → 256	714 → 365
4-tile image	81 → 81	185 → 167	522 → 452	816 → 765

Text throughput is unchanged (bench_http: 287 / 1716 / 3379 / 3849 tok/s at c1 / 8 / 32 / 64).

Testing:

• CPU suite as in CI, plus ruff and test/ci. The new tests check that two requests batched together, with a chunk boundary inside one of them, get the same embeddings as running each alone, and that the matmul patch embedding matches the conv.
• On GB200 with MSTAR_TEST_HF=1: the Command A+ tests, KV layouts and startup graph validation.
• HF checkpoint parity (6 prompts x 32 tokens): everything is within the thresholds. Greedy matches 190/192 against 191/192 before. Both misses are exact bf16 ties on our side, like the existing one.

[mstar-review]

Automated review

🟡 No findings on 1caf65a, light review

Description

• The TTFT table doesn't name the harness or the command that produced it (concurrency sweep, warmup, image sizes, request count), so it can't be reproduced; the text-throughput line names bench_http but gives no flags either.
• The text row at c32 goes from 155 to 161 ms, and the description doesn't say whether that is run-to-run noise or a cost of running the larger batched vision steps on rank 0, which also runs the LLM.

What the review checked

• engine._forward always takes forward_batched now that can_batch returns True; it passes engine_inputs.request_ids in the same order as the inputs list that preprocess concatenates, so the per-request split in forward_batched lines up with the rids.
• The vision node has no accelerator-graph config and runs alone on rank 0 (configs/command_a_plus_tp4.yaml), so there are no padded rows and no lockstep concerns.
• The forward still works for a single request: forward accepts the extra tile_counts key from preprocess through **kwargs.
• The unfold/permute/flatten patch embedding gives the same (c, kh, kw) channel-major order as Conv2d's weight.flatten(1) and the same row-major patch order as conv.flatten(2).transpose(1, 2); the new CPU test checks this against the conv.
• The new tests (test_vision_node_batches_requests..., test_patch_embedding_matches_the_conv) run on the CPU with tiny encoders and need no GPU or weights.
• Splitting into TILES_PER_CALL chunks is mechanical processing of a batch the scheduler already formed, not batch selection, so it does not break invariant 1. The forward makes no host syncs (invariant 3).

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

[Gaurav-Shah05]

LGTM. I've a TP8 run queued on H100s, and I will post the numbers here soon. Please wait for them before merge.

[Gaurav-Shah05]

Run Configuration:
8xH100s at TP8 with configs/command_a_plus_tp8.yaml: three launches per side, alternating, plus one pair with the API server's preprocessing at 4 threads. The vision step and memory rows come from the vision tower alone on one H100.

Numbers:

	main	this PR	Δ
Vision step, 16 requests, 117 tiles	387 ms	288 ms	1.35x ↓
Extra TTFT at an unseen tile count	33 to 45 ms	1 to 18 ms	~31 ms ↓
TTFT p50, c16	4857 to 4910 ms	4938 to 4958 ms	1.0x
TTFT p50, c1 / c4, preprocessing at 4 threads	207 / 373 ms	176 / 313 ms	1.18x ↓ / 1.19x ↓
Vision tower peak memory, 117 tiles	0.30 GiB	1.82 GiB	6.0x ↑
Text output, c8	794 to 830 tok/s	778 to 802 tok/s	1.0x

Batching rarely engages, because the API server tiles images one at a time (26 to 264 ms each) and feeds the encoder slower than it drains them.

Good to merge.