def attn_tree
import bendtorch@1.3.0.0/main.bend as Bendtorchimport Base
2916# whether the units (a sequence's key/value group each, fslices) run in the tree: from 8 units on, or from 2 when each is
2917# big. With the tiny Llama's sequences (t = 16, w = 32) the tree cost more than it gave at 2 sequences (0.48 against
2918# 0.37 ms on 16 threads) and gained from 8 on (1.32 against 1.52 ms; 1.5x at 32, 1.7x at 128); with SmolLM2's (t = 128,
2919# w = 576, about half a second of a step each) the two sequences one after another kept a single thread busy for a
2920# third of the step (NOTES, 2026-10-21), and its 2 sequences x 3 groups now give 6 units. t w counts in U32 (a Nat
2921# comparison walks its numbers one unit at a time). The answer never changes the result: the laws attn_par_fwd and
2922# attn_par_bwd show the tree gives the units in order, as one after another.
2923def attn_tree(+b: Nat, +t: Nat, +w: Nat) -> Bool:
2924 Bool.or(Nat.is_le(8n, b), Bool.and(Nat.is_le(2n, b), U32.is_le(4096, (U32.from_nat(t) * U32.from_nat(w) : U32))))