[AmitMY]

Stacked on #81.

Labels and the latent mask from the words, in the model

Before, the processor tokenized every word twice (once as an input, once as the previous word's label) and built a dense words × words attention mask for every example. Now an example carries:

• each word's bytes (input_ids)
• its packed sequence and shift block ids (sequence_ids, block_ids)
• whether it predicts the next word (label_mask: not PAD words, nor the words of a shift block before its ShiftIn)

From these the model derives:

• the labels: next_word_labels, the next word's bytes in the same sequence, or an empty word after the last word of a sequence
• the latent transformer's FlexAttention block mask: latent_block_mask, causal within a sequence and bidirectional within a shift block

Removed: the second tokenization, get_sequence_labels, the PAD-label loop, get_attention_mask_for_packed_sequence, and the input_attention_mask, attention_mask and labels_* batch keys.

Verification so far

• Equivalence: on 55 real packed examples (translation-style shift blocks, SignWriting, plain text), the derived labels, label positions and attention masks are identical to the old processor's. Examples are half the size (235 KB against 462 KB).
• FLOPs: benchmarks/flops.py counts the same FLOPs per step (3.078 / 2.504 TFLOPs), so the decoder sees the same label lengths.
• Tests: all 91 pass, including vLLM generation matching Megatron and the masked-attention test against SDPA.

Training A/B

Machine translation, 10,000 steps, each run twice with the same seed, against the same stack without this change:

Run	Loss at step 10	Val. bits/byte	Exact	chrF	ms / step
Before, run 1	3.500859	0.5268	8.6%	46.3	187.8
Before, run 2	3.499039	0.5361	8.2%	46.9	185.6
This PR, run 1	3.499090	0.5404	8.2%	45.8	183.6
This PR, run 2	3.499090	0.5404	8.2%	45.9	185.8

• Same quality: the gap is the size of the old code's own run-to-run spread (0.009 bits/byte). On the same weights and texts, both versions' per-byte losses agree to 4e-5 on every scored byte.
• Reproducible: this PR gives identical losses on reruns with the same seed. Before, the same seed gave different losses from step 10.
• Same speed, within noise.

9 files changed, +140 / −163.

🤖 Generated with Claude Code

https://claude.ai/code/session_015Xm62qoeaTunFSU5GYCQ1f

---

Note

Medium Risk
Changes the training/inference batch schema and core forward path (labeling + latent attention), though tests and A/B runs report equivalent behavior.

Overview
Processor batches are slimmer: examples now emit word bytes (input_ids), packed sequence / shift-block ids (sequence_ids, block_ids), and label_mask (who predicts the next word). The old second tokenization pass, dense L×L attention_mask, and labels_* / input_attention_mask keys are gone.

The model owns derivation: next_word_labels builds decoder targets from those fields; latent_block_mask builds the latent transformer’s FlexAttention mask (causal per sequence, bidirectional inside shift blocks). WeLTModel.forward, encode_words, and latent take the new tensors; MaskedAttention expects a pre-built block mask carrier instead of materializing from a dense mask.

Benchmarks/tests follow the new contract (benchmarks/flops.py uses next_word_labels for decoder token counts; attention tests focus on get_shift_blocks only).

^{Reviewed by Cursor Bugbot for commit b81d8c7. Bugbot is set up for automated code reviews on this repo. Configure here.}