Fix/output reorder index space

[thc1006]

Overview

This PR fixes two related indexing issues in output_reorder() and closes a lifecycle bug for output permutations that could lead to silent memory corruption.

1. Row domain mismatch: embd_layer_inp and unmasked embd_nextn are token-indexed, but upstream applied output_swaps (an output-row permutation) to them. This scrambles hidden states when n_outputs < n_tokens. This PR introduces token_swaps to correctly align ubatch-written rows into batch order.
2. Mode-flip corruption: cparams describes the layout for the next evaluation. If a caller toggles cparams.embeddings_nextn_masked after a decode but before reading the buffer, the getters and reorder loops apply the new mode's logic to the old layout. This PR snapshots the layout state into llama_context to guarantee the row domain always matches the allocation.
3. Lifecycle of Output Permutations: Previously, encode() did not clear output_swaps. If an encode() followed a decode() that regrouped tokens, llama_get_logits_ith would apply a stale permutation to the newly generated logits. Because buf_output is a single shared allocation (logits | embd | embd_nextn | ...), this stale swap could index past the logical bounds of the new logits and blindly overwrite neighboring regions within the same physical allocation. Both output_swaps and token_swaps are now explicitly cleared during encode().

Additional information

1. Layout, token_swaps, and Allocator Updates
The extraction paths for token-indexed buffers write in ubatch order, while downstream consumers read by batch index. These orders diverge when the splitter regroups. token_swaps correctly aligns the rows.

Buffer	Mode	Row Domain	Permutation Applied
embd_nextn	masked	output rows	output_swaps (unchanged)
embd_nextn	unmasked	token rows	token_swaps (new)
embd_layer_inp	any	token rows	token_swaps (new)
logits, embd, sampling	-	output rows	output_swaps (unchanged)

To support token_swaps, src/llama-batch.{cpp,h} is updated. llama_batch_allocr now tracks and exposes tok_ids (the batch indices for every token, alongside the existing out_ids for outputs). This gives output_reorder the necessary index mapping to correctly permute the token rows.

2. Layout State Snapshot & Assertions

• Added embd_nextn_masked_output to llama_context. It is recorded inside output_reserve and read at the three sites that previously relied on cparams. This prevents silent data corruption if the mode is flipped between decode and read.
• Bounds check in token-row swap loop: Added to prevent overruns during row swapping into neighboring regions of the monolithic allocation, which ASan cannot detect.
• Width check in extract_layer_inputs: Asserts row_floats == model.hparams.n_embd to explicitly enforce the coupling that output_reorder swaps exactly n_embd floats at a time.

3. Regression Tests & Test Suite Cleanups
Added ~500 lines of tests in tests/test-llama-archs.cpp. The core row-domain and mode-flip tests fail on the base commit and pass with these changes:

• test_output_reorder_token_rows: Tests partial-output and all-output token row alignments using pure embeddings for bit-exact comparisons.
• test_output_reorder_nextn_rows: Tests explicit mode flips using nearest-neighbor row matching.
• probe_shapes: Sweeps 160 batch shapes to cover cases where selected outputs are already sorted (making output_swaps empty), but token rows still need regrouping. Re-gating the existing swap cannot fix this.
• Test Cleanups: Fixed 6 false-positive behaviors in the tests (e.g., the row matcher accepting NaN as a valid match, silent skips on failures, asserting n_embd_out width correctly, etc.) and condensed verbose comments to adhere strictly to the project's 1-2 line descriptive standard.

Note on CONTRIBUTING.md:35 (Missing test for encode()): The encode() lifecycle fix (Overview point 3) lacks a fail-before test. Writing one requires an architecture with an encoder that populates these specific token-indexed buffers, which currently does not exist in the test suite. T5 was tested and passes, confirming the fix does not break existing encoder paths, but a strict fail-before test cannot be structurally constructed here.

4. Performance
Measurements via -p 512 -n 128 (a single sequence with no regrouping, so token_swaps is empty) show that when the feature is practically inactive, the unconditional cost (one tok_ids.push_back() per token) falls cleanly within measurement noise [-1.62%, +1.77%], causing no measurable slowdown.

5. Scope Boundary
This indexing mismatch was found while investigating #27572. While this PR strictly fixes the row-domain corruption and mode-flip bugs, it does not resolve the separate concurrent acceptance collapse reported there. That collapse was traced directly to graph-level write-after-read races when the ring buffer is disabled (addressed independently in #27311), not the indexing issues fixed here. This PR does not claim to fix #27572.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES — An AI agent was used to perform the row-order measurements, verify the transposition logic, benchmark CPU performance, format the tests, and draft this PR description. I manually reviewed and verified all AI-assisted work, fully understand the implementation, and personally authored the final code changes and commit messages. I also edited and approved the final PR description.

[ggml-gh-bot]

Hi @thc1006, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 2 open PRs.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.