[AmitMY]

Stacked on #83.

One FlexAttention core attention

WeLT had two core attention classes: PackedAttention for packed short sequences (the encoders and the bytes decoder), and MaskedAttention for the latent transformer's mask. Both ran FlexAttention. They differed only in input layout and in where the block mask came from, and each built or cached its mask inside the layers.

Now there is one FlexCoreAttention, and the caller builds the block mask once for all layers:

• Packed sequences: run_packed_transformer(..., causal=) puts packed_block_mask(...) on the PackedSeqParams.
• The latent transformer: latent() puts latent_block_mask(...) on the attention mask tensor, which Megatron passes through its layers. latent_block_mask now returns the BlockMask itself.

build_transformer loses its attention argument ("arbitrary" / "causal" / "bidirectional"): every transformer gets the same attention.

Shift blocks in the processor

welt/attention.py held only get_shift_blocks, so it moves into welt/processor.py, its main user. Its warnings are shorter. Its tests move from tests/test_attention.py to tests/test_processor.py.

Verification

• Full suite: 96 passed, 1 skipped, including the packed and latent attention tests against SDPA, Tests: labels and the latent mask against oracles of the words #83's oracles, and vLLM matching Megatron.

• Same training, bit for bit: machine translation (10,000 steps, seed 42) on this stack plus utf8-tokenizer's special token ids instead of copies; AutoConfig directly; FLOPs on collated batches #85 matches Experiment: labels and the latent mask from the words, in the model #82's runs exactly:

  Run	Loss at step 10	Val. bits/byte	Exact	chrF
  Experiment: labels and the latent mask from the words, in the model #82	3.499090	0.5403823	8.2%	45.8
  One FlexAttention core attention; shift blocks in the processor #84 + utf8-tokenizer's special token ids instead of copies; AutoConfig directly; FLOPs on collated batches #85	3.499090	0.5403823	8.6%	46.1

  Generation differs slightly with the same weights: vLLM's batched generation is not bit-deterministic. Experiment: labels and the latent mask from the words, in the model #82's two runs, also with identical weights, gave chrF 45.8 and 45.9.

7 files changed, +82 / −138.

🤖 Generated with Claude Code

https://claude.ai/code/session_015Xm62qoeaTunFSU5GYCQ1f

---

Note

Medium Risk
Attention wiring and mask lifecycle changed across all transformers; behavior is covered by SDPA/oracle tests and reported bit-identical training, but this is core model forward path.

Overview
Unifies Megatron core attention behind a single FlexCoreAttention that handles both packed (THD) and batched (S,B) layouts. Callers attach the FlexAttention BlockMask once per forward—packed_block_mask on PackedSeqParams (with explicit causal in run_packed_transformer) or latent_block_mask on the attention-mask tensor—instead of PackedAttention / MaskedAttention building or caching masks inside every layer. build_transformer drops its attention mode switch; latent_block_mask returns a BlockMask directly (no _welt_block_mask carrier).

Moves shift-block parsing from deleted welt/attention.py into welt/processor.py (get_shift_blocks, shorter warning text); tests move from tests/test_attention.py to tests/test_processor.py, and inference imports follow.

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