batching: scope token-flattened rows, and warn on writes that can't be scoped

[JadenFiotto-Kaufman]

Items 1, 13 and 14.3 from the stress-sweep decisions.

What was broken

Row scoping was a single shape test (tensor.shape[0] == self.total). Anything
else went to every invoke whole, with three silent consequences: a read saw the
whole batch, an in-place edit landed on every invoke, and a replacement was
dropped (_widen_tensor returned full and discarded edited). The layout that
hits this in practice is the token-flattened one — every transformers MoE block
reshapes to (batch*seq, hidden) before its router, so mlp.gate.output,
mlp.experts.inputs and every .source op below the reshape were unscoped. On
hf-internal-testing/tiny-random-Qwen2MoeForCausalLM, two invokes each forcing
"their" last token onto an expert both wrote the same rows, and the third,
unedited invoke's logits moved.

Forward keywords on an invoke were merged batch-wide, last writer wins, with
nothing said: output_hidden_states came back None for the invoke that asked
for it, and a per-invoke attention_mask/max_new_tokens quietly became the
batch's.

docs/usage/trace.md, the batching.py module docstring and NNsight.md all
promised that each invoke "sees only its own rows of every activation".

What changed

• Batcher._narrow_tensor / _widen_tensor take a leading dim that is a whole
  multiple of the batch size: k = rows // total rows per row of batch, narrowed
  and spliced at [start*k, size*k). Fixes the MoE (B*T, D) and per-head
  (B*T*H, D) cases. Accepted risk, as decided: a leading dim that is a multiple
  by coincidence (four experts against two invokes) is now sliced rather than
  passed through.
• A write to a value neither rule matched warns, naming the location, its leading
  dim and the batch size. A replacement is caught in _widen_tensor (and only
  when the block actually handed back something different, so round-tripping the
  non-batched leaves of a tuple output stays quiet). An in-place edit never comes
  back through the batcher at all, so _narrow_tensor records what it served
  whole with torch's _version counter and _report_unscoped — called at the top
  of narrow/widen — warns for anything whose version moved. That makes the
  report one batcher call late (the next value served to any worker); an edit
  after which nothing at all is read is the case it can miss. Reads deliberately
  do not warn: from a shape alone, a causal mask is indistinguishable from a
  vision tower's patches, and warning on reads would fire constantly.
• TransformersModel._batch_forward warns when two invokes pass different values
  for the same forward keyword.
• Docs: the scoping rule, the three failure modes and the warning in
  docs/usage/invoke-and-batching.md; the corrected sentence in
  docs/usage/trace.md, the batching.py module docstring and NNsight.md §4.4 /
  §7.2; the internals in docs/developing/batching-internals.md; the
  forward-keyword rule; and (item 13, docs only) the ordering rule that statements
  outside the invoke blocks run in the collection pass, before the model — with
  the reduce-after-the-trace remedy, verified here.

Tested

PYTHONPATH=<worktree>/src against the stress env's interpreter, CPU only.

• New tests in tests/test_batching.py: TestFlattenedRows (read, in-place edit
  and replacement on a (batch*seq, hidden) activation, pure torch),
  TestUnscopedWrites (in-place warns, replacement warns and is dropped, a read
  is quiet), TestMoERouterScoping (the tiny-random Qwen2-MoE regression: a
  forced route matches the same prompt traced alone, does something, and leaves
  the control invoke equal to its solo run), TestForwardKwargs. All seven
  behavioural ones fail on a git archive HEAD baseline and pass here.
• tests/test_batching.py 55 passed, 1 skipped. Also green: test_interleaving,
  test_saving, test_backward, test_language, test_modeling, test_envoy,
  test_editing, test_source, test_tracing, test_encoder, test_vision,
  test_chat, test_chunked_tasks, test_multiple_wrappers,
  test_construction_routing, test_vlm, test_diffusion, test_fragments,
  test_util, test_memory, test_serialization, test_tensor_parallel_rules,
  test_deprecations, test_config (~950 tests).
• Not run here: tests/vllm/ and tests/tp/ (no vLLM run, no multi-GPU). The
  vLLM batcher inherits _narrow_tensor, where total is the step's token rows;
  the modulo branch only engages for a tensor whose leading dim is a multiple of
  that, and the warning is a warning.
• Checked for false positives: batched gpt2 traces with edits, tracer.cache()
  (with include_inputs=True), batched generate with tracer.iter, and
  attention-tuple edits produce no warnings.

Deliberately left alone

• The VLM/vision-side case (item 7). The modulo rule fixes equal-images-per-invoke
  by luck and does not address the rest; the axis table is the real answer.
• Time-major (seq, batch, hidden) with seq == total, and a list of
  per-sequence tensors: still wrong, still undetectable from a shape. Both are now
  written down in the docs.
• docs/concepts/batching-and-invokers.md:13 and :76 still state the old
  == total rule — that file belongs to another agent in this batch; it needs the
  same correction. Likewise the nnsight:nnsight skill's batching material,
  which lives outside this repo.
• Item 14's other two points (cache path validation, comment truncation) are not
  mine.

🤖 Generated with Claude Code