n150 n300 T3000 p100 p150 p300c Galaxy Sim 40 min Draft

Attention from Scratch

Following along? The code in this lab lives in your ~/tt-scratchpad/llm-from-scratch/ workspace. If you haven't created it yet, use the Create the LLM-from-Scratch Project button in Lab 0.

Embeddings & the Residual Stream gave the residual stream its first value: x = self.tok_emb(idx). That's a running [B, T, n_embd] tensor every block reads from and adds back into. In a modern Llama-3-shaped model, that's the whole value — there is no + pos_emb(pos) term.

Position was left out of the residual stream on purpose. A Llama-3-style model carries it a different way: as a rotation applied to Q and K, called RoPE. This is the lab where that rotation finally happens.

So far, each position in the stream is an island. Token 12's vector knows nothing about token 3's. Attention is the mechanism that lets every position look at every earlier position and pull in what it needs — and it's where RoPE gets applied. It's the reason a transformer can resolve "it" to the noun five words back, and it's the single most compute-heavy thing the model does.

This is the centerpiece lab. You'll read the math twice.

First in plain PyTorch: the GroupedQueryAttention module from the reference model, with RoPE'd Q/K, grouped-query sharing, and causal softmax. Then as a from-scratch TT-Lang kernel that expresses the single-head core of that same softmax(QKᵀ/√d + mask)V as an explicit reader → compute → writer pipeline.

This lab is also where the arc's honesty about a fast-moving stack earns its keep. Unlike the elementwise add in Embeddings & the Residual Stream, this attention kernel is not runnable in the bundled functional simulator yet. We'll be exact about why, how we gain confidence in it anyway, and what still has to happen before it ships.

You are here:

graph LR
    A["Tokenizer & Data"] --> B["Embeddings & RoPE"] --> C["Attention (GQA)"] --> D["Block & Model"] --> E["Train on Blackhole"]

    style C fill:#1B8EB1,stroke:#333,stroke-width:3px

Coming from CUDA: FlashAttention, and the agentic shortcut

If you've optimized attention on a GPU, you know the enemy is memory traffic, not arithmetic. The naive implementation materializes the full T×T score matrix in HBM, reads it back to softmax it, writes it again, reads it a third time to multiply by V — the matrix bounces off high-bandwidth memory several times, and for long sequences that bandwidth, not the matmul, is what you're waiting on. FlashAttention's whole trick is to never let the score matrix touch HBM: it tiles the computation and keeps a running, "online" softmax in on-chip SRAM, fusing QKᵀ, the softmax, and the ·V into one pass so each score tile is produced, consumed, and discarded without a round-trip.

The TT-Lang expression of attention is that same fusion, made explicit. Every intermediate — the scores, the per-row max, the exponentials, the row sums, the normalized weights — lives in an L1 dataflow buffer and is handed between the reader, compute, and writer threads on one Tensix core. Nothing spills to DRAM between QKᵀ and ·V. Where FlashAttention arranges for the fusion behind a hand-tuned CUDA kernel, TT-Lang has you write the fusion out as tile math over L1 buffers — the keep-it-in-fast-memory story is the same, but here it's the literal structure of the source.

The agentic shortcut. There's a deeper reason this matters for how you'll actually build kernels. On CUDA, "you'll need a custom attention kernel" ends a lot of afternoons — and it's exactly the code AI coding agents are worst at. So much of a CUDA kernel's correctness lives in the programmer's head: which tiles are resident, whether the loads coalesce, which warp got to the barrier first, what the occupancy is at this block size. None of that is in the source, so an agent has nothing concrete to generate against, and it hallucinates plausible-looking code that races or reads garbage.

TT-Lang inverts that. The spec is the source: a reader that declares exactly which tiles arrive and from where, a compute body that is pure tile math on those arrivals, and a writer that declares exactly what leaves. Arrivals in → tile math → departures out, nothing implicit. Because the full contract lives in the three functions, you can hand an agent a plain-English reader/compute/writer spec — "read Q, K, V tiles; compute QKᵀ, scale, mask, softmax, then ·V; write the output tile" — and it has something complete to fill in and something concrete to check itself against. That's why "TT-native from line one" is reachable for a kernel as involved as attention: you're not asking an agent to intuit hidden GPU state, you're asking it to write down a pipeline whose shape is already pinned by the DSL.

Attention, the math: Grouped-Query Attention with RoPE'd Q/K (PyTorch reference)

Here is the reference — GroupedQueryAttention, quoted verbatim from ~/tt-scratchpad/llm-from-scratch/reference_gpt.py so the prose and the code can't drift apart:

class GroupedQueryAttention(nn.Module):
    """Causal grouped-query attention with RoPE — the heart of Lab 3.

    GQA: `n_head` query heads share `n_kv_groups` key/value heads (each KV head
    serves `n_head // n_kv_groups` query heads). This shrinks the KV cache and
    KV projections while keeping full query resolution — the modern inference
    win over vanilla multi-head attention.

    Pedagogy note: ttml's GroupedQueryAttention fuses K and V into one
    `kv_linear` (concat_kv_dim) and creates heads with
    `ttml.ops.multi_head_utils.grouped_heads_creation`; here we keep separate
    q/k/v projections for readability. The math is identical.
    """

    def __init__(self, cfg: LlamaConfig):
        super().__init__()
        assert cfg.n_embd % cfg.n_head == 0, "n_embd must be divisible by n_head"
        assert cfg.n_head % cfg.n_kv_groups == 0, (
            "n_head must be divisible by n_kv_groups"
        )
        self.n_head = cfg.n_head
        self.n_kv_groups = cfg.n_kv_groups
        self.head_dim = cfg.n_embd // cfg.n_head
        self.group_size = cfg.n_head // cfg.n_kv_groups  # query heads per KV head

        # Llama linears carry no bias. Q is full width; K/V are group-width.
        self.q_proj = nn.Linear(cfg.n_embd, self.n_head * self.head_dim, bias=False)
        self.k_proj = nn.Linear(cfg.n_embd, self.n_kv_groups * self.head_dim, bias=False)
        self.v_proj = nn.Linear(cfg.n_embd, self.n_kv_groups * self.head_dim, bias=False)
        self.o_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
        self.attn_dropout = nn.Dropout(cfg.dropout)
        self.resid_dropout = nn.Dropout(cfg.dropout)

        # Lower-triangular causal mask (no token attends to the future).
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(cfg.block_size, cfg.block_size)).view(
                1, 1, cfg.block_size, cfg.block_size
            ),
            persistent=False,
        )

    def forward(self, x, cos, sin):
        B, T, C = x.shape
        # Project, then split into heads. Q has n_head heads; K/V have
        # n_kv_groups heads (that is the whole GQA trick).
        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_groups, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_groups, self.head_dim).transpose(1, 2)

        # RoPE rotates Q and K (never V) — position lives in the Q·K product.
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)

        # Share each KV head across its group of query heads. repeat_interleave
        # (not repeat) keeps heads of the same group adjacent, matching how the
        # query heads are laid out. (B, n_kv_groups, T, hd) -> (B, n_head, T, hd)
        k = k.repeat_interleave(self.group_size, dim=1)
        v = v.repeat_interleave(self.group_size, dim=1)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        weights = self.attn_dropout(F.softmax(scores, dim=-1))
        out = weights @ v                        # (B, n_head, T, head_dim)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.o_proj(out))

Strip away the bookkeeping and it's six steps:

Project — asymmetrically. q_proj produces n_head (6) query heads, full width; k_proj/v_proj produce only n_kv_groups (3) heads each — half as many KV heads as query heads. That asymmetry is the entire GQA trick, and it's visible directly in the two projection widths: self.n_head * self.head_dim vs. self.n_kv_groups * self.head_dim.
Rotate. apply_rope(q, cos, sin) and apply_rope(k, cos, sin) — the exact function Embeddings & the Residual Stream handed you, now finally called. This is where position actually enters the model: not in the residual stream, but here, as a rotation of Q and K (never V) by an angle that depends on each token's position.
Share. k.repeat_interleave(self.group_size, dim=1) and the same for v — each of the 3 KV heads gets copied to serve group_size = 2 query heads apiece, so K and V end up looking [B, n_head, T, head_dim] — the same shape attention needs — without ever having computed n_head independent KV projections.
Score. q @ k.transpose(-2, -1) — every (now RoPE'd) query dotted with every (now RoPE'd, now shared) key, giving a [T, T] matrix per head where entry (i, j) is how much token i attends to token j. Divide by √head_dim to keep the scores from growing with dimension (which would saturate the softmax).
Mask + softmax. masked_fill(... == 0, -inf) sets everything above the diagonal to -inf so a token can't attend to the future (this is what makes it causal), then F.softmax(scores, dim=-1) turns each row into a probability distribution over the keys.
Gather. weights @ v — a weighted sum of (shared) value vectors, so each output position is a blend of the value vectors it chose to attend to.

Why GQA, concretely. Every one of those n_head query heads still gets its own attention pattern — full query resolution is preserved. What shrinks is the key/value side: n_kv_groups = 3 instead of n_head = 6 means the KV cache an inference server has to keep resident (one K and one V vector per token, per KV head, for every token generated so far) is half the size, and the k_proj/v_proj matmuls that produce it are correspondingly cheaper. That's the modern efficiency win: MHA is simply the special case n_kv_groups == n_head (no sharing, no cache savings); GQA dials that ratio down without touching query resolution at all. It's why Llama-3 uses GQA throughout, and why Mini-LLM — the from-scratch project this arc credits — and this arc's own nanollama3 hero config (6 heads, 3 KV groups) both build on it rather than plain MHA.

The TT-Lang kernel below implements the numerical core of one head — score → mask → softmax → gather — as tile math. (The projections, RoPE rotation, and KV-head sharing are ordinary ttnn matmuls, elementwise ops, and reshapes around it; the kernel's job is the fused score → mask → softmax → gather that FlashAttention fuses on CUDA, run once per query head after GQA has already lined Q, K, and V up to the same [T, head_dim] shape.)

The TT-Lang attention inception kernel

Recall the arc's framing from Pick Your Altitude: descending to TT-Lang for a hot kernel is inception, not conversion. There is no CUDA or Triton attention kernel being ported here. This kernel is attention_kernel, authored TT-native and kept faithful to the vendor original in vendor/tt-lang/examples/test_transformer_block.py — the same reader → compute → writer shape you met in Embeddings & the Residual Stream, scaled up to the full attention pipeline. It lives in ~/tt-scratchpad/llm-from-scratch/kernels/attention.py.

Where GQA fits against this kernel. attention_kernel below takes Q, K, and V that are already [T, head_dim] for a single head — it has no opinion about how many query heads share that K/V, because by the time this kernel runs, GQA's repeat_interleave has already made every query head's K and V look like an ordinary single-head input. GQA is a sharing pattern layered on top of this kernel, not a change to it: run it once per query head (6 times for nanollama3), and the KV-head-sharing that makes GQA cheap already happened one level up, in the reshape/repeat step, not in the tile math below. On the real training path this reshape has a name — ttml.ops.multi_head_utils.grouped_heads_creation — which is what ttml.models.llama's GQA module calls to turn n_kv_groups KV heads into n_head-shaped tensors before attention runs; the PyTorch reference's repeat_interleave is the from-scratch mirror of exactly that call.

It works at single-tile granularity — one 32×32 tile of tokens (SEQ_TILES = 1), one 32×32 tile of head dimension (EMBD_TILES = 1) — so the pipeline is readable without grid-partitioning bookkeeping getting in the way. The signature and the buffer declarations first:

@ttl.operation(grid=(1, 1))
def attention_kernel(q, k, v, scale, causal_mask, scaler, out):
    """Single-head scaled dot-product attention.

    ``scale`` is a tile filled with 1/sqrt(head_dim); ``scaler`` is a tile of
    ones used by the reduce ops; ``causal_mask`` has -inf above the diagonal.
    """
    q_dfb = ttl.make_dataflow_buffer_like(q, shape=(SEQ_TILES, EMBD_TILES), block_count=2)
    k_dfb = ttl.make_dataflow_buffer_like(k, shape=(SEQ_TILES, EMBD_TILES), block_count=2)
    v_dfb = ttl.make_dataflow_buffer_like(v, shape=(SEQ_TILES, EMBD_TILES), block_count=2)
    scale_dfb = ttl.make_dataflow_buffer_like(scale, shape=(1, 1), block_count=1)
    mask_dfb = ttl.make_dataflow_buffer_like(causal_mask, shape=(SEQ_TILES, SEQ_TILES), block_count=1)
    scaler_dfb = ttl.make_dataflow_buffer_like(scaler, shape=(1, 1), block_count=1)
    out_dfb = ttl.make_dataflow_buffer_like(out, shape=(SEQ_TILES, EMBD_TILES), block_count=2)

    k_t_dfb = ttl.make_dataflow_buffer_like(k, shape=(EMBD_TILES, SEQ_TILES), block_count=2)
    snodes_dfb = ttl.make_dataflow_buffer_like(causal_mask, shape=(SEQ_TILES, SEQ_TILES), block_count=2)
    scale_bcast_dfb = ttl.make_dataflow_buffer_like(causal_mask, shape=(SEQ_TILES, SEQ_TILES), block_count=2)
    scaled_masked_dfb = ttl.make_dataflow_buffer_like(causal_mask, shape=(SEQ_TILES, SEQ_TILES), block_count=2)
    max_dfb = ttl.make_dataflow_buffer_like(scaler, shape=(1, 1), block_count=2)
    max_bcast_dfb = ttl.make_dataflow_buffer_like(causal_mask, shape=(SEQ_TILES, SEQ_TILES), block_count=2)
    exp_dfb = ttl.make_dataflow_buffer_like(causal_mask, shape=(SEQ_TILES, SEQ_TILES), block_count=2)
    sum_dfb = ttl.make_dataflow_buffer_like(scaler, shape=(1, 1), block_count=2)
    sum_bcast_dfb = ttl.make_dataflow_buffer_like(causal_mask, shape=(SEQ_TILES, SEQ_TILES), block_count=2)
    softmax_dfb = ttl.make_dataflow_buffer_like(causal_mask, shape=(SEQ_TILES, SEQ_TILES), block_count=2)

That's the FlashAttention "keep it in fast memory" story spelled out as declarations: the top block is the seven I/O buffers (Q, K, V, the scale tile, the causal mask, the reduce scaler, and the output); the bottom block is the ten intermediate buffers — every partial result the fused pipeline produces (the transposed K, the raw scores, the scaled-and-masked scores, the per-row max and its broadcast, the exponentials, the row sums and their broadcast, and the final normalized weights) gets its own L1 buffer instead of a DRAM round-trip.

DFBs and the 32-per-core ceiling

Count the make_dataflow_buffer_like calls and there are 17 dataflow buffers on this one core. That number is worth pausing on, because it's a real hardware budget, not an abstraction. Per the TT-Lang specification's platform-limits table (TTLangSpecification.md, Appendix E), the maximum number of dataflow buffers is 32 per core, on both Wormhole^™ and Blackhole^®. So attention's 17 buffers sit a little over halfway into the ceiling — comfortably within budget, but close enough that you can feel it: a two-headed fused variant, or one that kept the exponentials around instead of recomputing them, would start crowding 32. This is the kind of resource accounting that on CUDA lives in your head as "register pressure" and "shared-memory occupancy"; in TT-Lang it's a countable line-by-line fact in the source, which is again why an agent can reason about it.

The block_count on each buffer is the same double-buffering knob from Embeddings & the Residual Stream: block_count=2 lets the producer fill the next block while the consumer drains the current one (pipelining); the single-slot block_count=1 buffers (scale_dfb, mask_dfb, scaler_dfb) hold read-once constants that never need a second slot.

The compute thread: score → mask → softmax → gather

Here is the whole compute body, verbatim — the score/mask/softmax/gather core of the PyTorch reference (steps 4 through 6 above, for one already-shared head), now as tile operations handed between L1 buffers:

    @ttl.compute()
    def compute():
        # scores = Q @ K^T
        with k_dfb.wait() as kv, k_t_dfb.reserve() as kt:
            kt.store(ttl.transpose(kv, kt))
        with q_dfb.wait() as qv, k_t_dfb.wait() as ktv:
            with snodes_dfb.reserve() as sc:
                sc.store(ttl.math.matmul(qv, ktv, sc))

        # scaled + masked scores
        with (
            snodes_dfb.wait() as scv,
            scale_dfb.wait() as scalev,
            mask_dfb.wait() as maskv,
        ):
            with scale_bcast_dfb.reserve() as sb:
                sb.store(ttl.math.broadcast(scalev, sb, dims=[0, 1]))
            with scale_bcast_dfb.wait() as sbv, scaled_masked_dfb.reserve() as sm:
                sm.store(scv * sbv + maskv)

        # softmax (max-shifted for numerical stability)
        with scaler_dfb.wait() as scaler_v, scaled_masked_dfb.wait() as smv:
            with max_dfb.reserve() as mx:
                mx.store(ttl.math.reduce_max(smv, scaler_v, mx, dims=[0]))
            with max_dfb.wait() as mxv, max_bcast_dfb.reserve() as mxb:
                mxb.store(ttl.math.broadcast(mxv, mxb, dims=[1]))
            with max_bcast_dfb.wait() as mxbv:
                shifted = smv - mxbv
                with exp_dfb.reserve() as ex:
                    ex.store(ttl.math.exp(shifted))
                with exp_dfb.wait() as exv, sum_dfb.reserve() as sm:
                    sm.store(ttl.math.reduce_sum(exv, scaler_v, sm, dims=[0]))
                with sum_dfb.wait() as smv2, sum_bcast_dfb.reserve() as smb:
                    smb.store(ttl.math.broadcast(smv2, smb, dims=[1]))
                with sum_bcast_dfb.wait() as smbv, softmax_dfb.reserve() as sfm:
                    # Recomputing exp(shifted) here (rather than reusing exv from
                    # exp_dfb above) is faithful to the vendor kernel structure —
                    # `attention_kernel` in test_transformer_block.py @ a19aaa8
                    # does the same recompute. Not "optimized" away on purpose.
                    sfm.store(ttl.math.exp(shifted) / smbv)

        # out = softmax @ V
        with softmax_dfb.wait() as sfmv, v_dfb.wait() as vv:
            with out_dfb.reserve() as o:
                o.store(ttl.math.matmul(sfmv, vv, o))

Read it as the four commented stanzas:

scores = Q @ K^T — ttl.transpose flips K into Kᵀ in L1, then ttl.math.matmul produces the raw score tile. (On CUDA you'd fuse the transpose into the load; here it's one explicit tile op, one buffer.)
scaled + masked scores — ttl.math.broadcast fans the single 1/√d scalar out across the whole score tile, then scv * sbv + maskv scales and adds the causal mask (-inf above the diagonal) in one tile expression. This is step 4's / √head_dim and half of step 5's masked_fill combined.
softmax — this is the interesting part, and the reason for most of those intermediate buffers. Softmax over a tile can't be a single op; it's built from primitives: reduce_max down the key dimension for the per-row max, broadcast it back across the row, subtract it (the standard max-shift for numerical stability, so exp never overflows), exp, reduce_sum the exponentials, broadcast that sum back, and divide. Six tile ops and five buffers to express one F.softmax. This softmax-from-reductions is the centerpiece of the centerpiece — and, as the next section explains, it's also exactly the piece the bundled simulator won't run yet.
out = softmax @ V — one more matmul, the weighted gather, into the output buffer.

The recompute of exp(shifted) in the final softmax line is deliberate and called out in the source: it mirrors the vendor kernel's structure rather than hand-optimizing away a buffer, so the drift check against test_transformer_block.py stays clean.

The reader and writer threads

The two data-movement threads are the plainest part — they just stream the I/O buffers between DRAM and L1:

    @ttl.datamovement()
    def dm_read():
        with q_dfb.reserve() as blk:
            tx = ttl.copy(q[0:SEQ_TILES, 0:EMBD_TILES], blk)
            tx.wait()
        with k_dfb.reserve() as blk:
            tx = ttl.copy(k[0:SEQ_TILES, 0:EMBD_TILES], blk)
            tx.wait()
        with v_dfb.reserve() as blk:
            tx = ttl.copy(v[0:SEQ_TILES, 0:EMBD_TILES], blk)
            tx.wait()
        with scale_dfb.reserve() as blk:
            tx = ttl.copy(scale[0, 0], blk)
            tx.wait()
        with mask_dfb.reserve() as blk:
            tx = ttl.copy(causal_mask[0:SEQ_TILES, 0:SEQ_TILES], blk)
            tx.wait()
        with scaler_dfb.reserve() as blk:
            tx = ttl.copy(scaler[0, 0], blk)
            tx.wait()

    @ttl.datamovement()
    def dm_write():
        with out_dfb.wait() as blk:
            tx = ttl.copy(blk, out[0:SEQ_TILES, 0:EMBD_TILES])
            tx.wait()

dm_read .reserve()s each input buffer (producer role) and ttl.copys a tile in from DRAM; dm_write .wait()s on the finished output (consumer role) and copies it back out. Same reader → compute → writer contract as the add in Embeddings & the Residual Stream — just seven inputs and a much busier compute body between them.

Validate: the PyTorch-reference correlation check

Because this kernel doesn't run in the simulator (next section explains why), its main() is built to still give you an honest, reproducible signal. Run it:

python ~/tt-scratchpad/llm-from-scratch/kernels/attention.py

main() does three things, in order:

Runs the PyTorch reference first, always. It builds random bf16 Q, K, V and a causal mask, computes softmax(QKᵀ·scale + mask) @ V with plain PyTorch (_torch_attention), and prints the first few output values. This is the exact math the lab teaches, and it runs on CPU anywhere — so the file always exits 0 with a real reference result, no device or simulator needed.
Attempts the TT-Lang kernel. It converts the tensors with ttnn.from_torch(..., layout=ttnn.TILE_LAYOUT) and calls attention_kernel. At the pinned tt-lang commit this raises inside the simulator (see below); main() catches it and prints TT-Lang kernel: SIM-UNSUPPORTED at this pin (compiler/hardware path only) with the underlying reason, and still exits 0.
Scores it against the reference — when it runs. If the kernel does execute (on the compiler/hardware path, or a future sim build that accepts the broadcast), main() converts the output back with ttnn.to_torch, computes the max absolute error and the correlation against the PyTorch reference, and prints PASSED when correlation exceeds 0.9. That correlation check is how we gain confidence in the kernel's numerics without a clean sim run — the same reference the lab just taught, used as the oracle.

So the validation you can reproduce today is the PyTorch reference plus the honest SIM-UNSUPPORTED report. The correlation gate is wired and ready; it lights up the moment the kernel runs on a path that accepts it.

Honest flag: the simulator is ahead of the hardware compiler

The elementwise add in Embeddings & the Residual Stream ran two ways — in your browser and in the standalone functional simulator — with max abs error vs torch: 0.000000. This attention kernel does not run in the bundled functional simulator, and this lab will not pretend it does. Here is exactly what's true, from the kernel's own header annotation (verified while authoring, not asserted):

Verification status: COMPILER / HARDWARE PATH ONLY — not runnable in the bundled functional simulator at this tt-lang pin (a19aaa8).

The cause is specific and narrow. Softmax needs reduce_max / reduce_sum down the key dimension, then broadcasts those per-row scalars back across the full score tile. In the pinned simulator build, ttl.math.broadcast requires its source block to have a genuine size-1 sub-tile dimension — but the reductions here yield a full 32×32 tile, so the sim raises "Cannot broadcast along dimension N: dimension must have element size 1". That's the softmax broadcast in the compute body above (max_bcast_dfb, sum_bcast_dfb) hitting a validator the simulator applies and the hardware path does not.

The framing to hold onto is the one the runtime matrix in Pick Your Altitude set up: the functional simulator runs ahead of the hardware compiler. The vendor's own test_transformer_block.py — the kernel this one is faithful to — carries TTLANG_HARDWARE_CI: skip-compiler and runs on a real device, precisely because it targets the compiler/hardware path, not the sim's stricter broadcast checker. This is a simulator gap, not a correctness problem with the kernel: the tile math is right, the vendor runs the same structure on silicon, and our numerics are pinned to the PyTorch reference.

So the honest status of this kernel, stated plainly:

Authored TT-native, from scratch — inception, not a port.
Faithful to the vendor attention_kernel (drift-checked against test_transformer_block.py @ a19aaa8).
Confident by construction and by oracle — faithful to the vendor attention_kernel (drift-checked, above), with the verified PyTorch reference standing in as the oracle it will be scored against. The correlation gate in main() is wired and ready, but pending the kernel's first successful execution (sim or hardware) — not executed on the functional simulator, and not yet run on hardware in this arc.
Scheduled for a hardware pass before publish — construction-faithfulness plus the PyTorch oracle is the confidence we have today; a real on-device run — the one that finally lights up the correlation gate — is the confidence we owe you before this leaves draft.

Contrast the add in Embeddings & the Residual Stream and the matmul in The Transformer Block & the Model, which are sim-runnable and pass cleanly. Attention and RMSNorm are the two kernels in this arc that press the edge the simulator hasn't caught up to yet — and saying so, precisely, is the point.

Graduate box

You've now read attention as the same math twice — a PyTorch module and a 17-buffer TT-Lang pipeline — and you know exactly how far each has been verified. The kernel's confidence today is construction-faithfulness to the vendor kernel plus the verified PyTorch reference acting as the oracle — the correlation gate that will score it against that oracle is wired and ready, pending the kernel's first successful execution; the confidence it still needs is a real on-device run.

That on-device run is Train It & Run for Real. There, the arc's hero run — nanollama3_char, the same 6-head, 3-KV-group GQA and RoPE θ=500000 this lab just walked — trains via ttml.models.llama on a Blackhole p300c, with the loss dropping monotonically from 4.69 to 3.23 over 20 steps, on-device, exit 0. That's where the GQA + softmax path this lab authored gets exercised for real on silicon, closing the loop this honest flag opens.

Attention is built — GQA, RoPE'd Q/K, and all — and its honesty is on the table. Next: the rest of the transformer block — a SwiGLU MLP, RMSNorm (which hits the same simulator edge as attention), and the residual wiring that stacks it all into a model you can run.

→ Continue to Lab 4: The Block & the Model