Kernels:

---

aws-neuron

/

vit-encoder-neuron-kernels

Like 3

Follow

AWS Inferentia and Trainium 194

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#
# MODIFIED from the AWS Neuron SDK's nkilib ``core/attention/attention_cte.py``
# (Apache-2.0, Copyright Amazon.com, Inc. or its affiliates). Modifications:
#   1. Four engine-placement edits (two memsets -> GpSimd, the tp_out eviction
#      tensor_copy -> Scalar, the list-branch write-back tensor_copy -> GpSimd),
#      measured worth ~1.19x end-to-end in the V-JEPA2 encoder.
#   2. Code paths unreachable for this model's configuration removed (causal/SWA
#      masks, prefix caching, flash sectioning at seq<=10240, sinks, fp8,
#      sequence packing, CP) -- the executed instruction stream is unchanged.
# The original supports the general case; use nkilib for anything but this
# non-causal, single-section, SBUF-KV configuration.
"""K2 inner attention: the vendor ``_attention_cte`` context-encoding kernel, FLATTENED to
exactly the branches the 4-kernel V-JEPA2 encoder reaches, with the four h8/h10 engine pins
applied in place. Self-contained: imports only ``nki``, ``nkilib``, ``torch`` (none here) --
NO imports from ``experiment/`` or the repo-root incumbent modules.

WHAT THIS COMPUTES
==================
``Output = softmax(scale * Q @ K^T) @ V`` for ONE head, called 8x per core by the K2 outer
kernel (``attention_block.py``). The pipeline, in order, per 128-token Q group:

    load_q  ->  QK^T (MM1) + row-max  ->  update running max  ->  exp (softmax numerator)
            ->  P@V (MM2)  ->  write-back (reciprocal-normalise, evict to the caller's SBUF O).

Group software-pipelining overlaps three groups at a time (grp_i does PV+writeback, grp_i+1
does exp, grp_i+2 does load-Q + QK/max), which is the vendor's Level-4 loop structure.

PSUM BANK PLAN (this config, tp_out=True, single 2048 section)
--------------------------------------------------------------
  * MM1 (QK^T) tiles: banks {0,1,2,3}   (``address=(0, (k_tile_idx % 4) * PSUM_BANK_SIZE)``).
  * MM2 (P@V) tiles:  banks {4,5,6,7}   (``address=(0, (4 + large_tile_idx % 4) * ...)``).
  * tp reciprocal / correction transposes: bank {3 or 7} (``((grp%2)*4 + 3) * ...``).
The producer (K2 outer) owns banks 5/6/7 during its phase; the attention region runs after the
producer's high-water mark, so there is no live overlap.

RING / MULTI-BUFFER DEPTHS (from the dp/dp2/dp4 fork lineage; kept verbatim)
---------------------------------------------------------------------------
  * ``exp_sb``   group-axis depth 2 (dp2): exp(g+1) can run while transpose(g) still drains.
  * ``exp_tp_sb`` group-axis depth 8 (dp4): P^T transposes issue as soon as their exp output
    exists, spreading DMA traffic out of the P.V collision window.

THE FOUR ENGINE PINS (h8/h10), and WHY each exists (one line each, from PERF_INVENTORY)
---------------------------------------------------------------------------------------
  1. ``_scale_reciprocal_write_back_impl`` tp_out PSUM->SBUF evict ``tensor_copy`` -> SCALAR
     (h8 site 1): the un-pinned copy defaults to Vector, K2's binding engine; Scalar/Vector are
     the only legal engines (PSUM source excludes GpSimd) and Scalar is the exercised majority.
  2. ``_write_back_o_impl`` tp_out list-branch ``tensor_copy`` -> GPSIMD (h8 site 3): SBUF->SBUF
     fp32->bf16 CAST, all three engines legal, GpSimd is the idle engine so the copy sheds off
     Vector.
  3. ``_qk_and_max_impl`` ``memset`` -> GPSIMD (h10 :2876): initialises mm1_partial_max, first
     consumer is several ops downstream, so a slower idle engine is safe and sheds Vector.
  4. ``_fused_qkmax_and_pv_impl`` ``memset`` -> GPSIMD (h10 :3294): same rationale, 2688/core.
The vendor ``_exp_impl`` memset (dp4 :2993) stays UNPINNED (engine.unknown / Vector) on purpose
-- pinning it would fork ``_exp_impl`` and re-home the softmax-exp source strings; it MUST stay
unpinned to reproduce the shipping instruction stream.

DEAD PATHS REMOVED (§4 of PERF_INVENTORY) -- guard that made each unreachable for this call:
``causal_mask=False, use_sbuf_kv=True, seqlen_k=2048<10240, bs=1, tp_out=True, fuse_v_proj=False,
no prior/APC/SWA/CP/seqpack/sink, cache_softmax=False``. See per-function docstrings.
The allocator usage (``nkilib`` ``ModularAllocator``) is kept EXACTLY as the vendor uses it.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, List, Optional

import nki
import nki.isa as nisa
import nki.language as nl
from nki.isa import engine, reduce_cmd
from nki.language.opcode import maximum as _maximum

from nkilib.core.utils.allocator import align_to
from nkilib.core.utils.kernel_assert import assert_shape, kernel_assert
from nkilib.core.utils.kernel_helpers import (
    PSUM_BANK_SIZE,
    div_ceil,
    get_verified_program_sharding_info,
    is_sbuf_tensor,
)
from nkilib.core.utils.logging import get_logger
from nkilib.core.utils.modular_allocator import ModularAllocator
from nkilib.core.utils.stream_shuffle_broadcast import stream_shuffle_broadcast

logger = get_logger("attention_inner")


_FLOAT32_MIN = -3.4028235e38  # used for initialization and masking

# Tested-range constraints (checked in _attention_cte).
_MAX_BS = 512
_MAX_SEQLEN = 131072
_MAX_BS_TIMES_SEQLEN_QK = 32.0 * 36864 * 36864
_MAX_HEAD_DIM = 128
_MIN_GLOBAL_CP_DEGREE = 1
_MAX_GLOBAL_CP_DEGREE = 32

# Tile size / threshold constants used by the LIVE path.
_Q_GRP_SZ = 128
_V_TILE_SZ = 128  # V tile size for loading and MM2 operations
import os as _os
_WB_FUSED = _os.environ.get("VIT_ENC_WB_FUSED", "0") == "1"      # fused normalise+cast write-back: OFF (measured +6.6% at N=2048)
_MM1_WIDE = _os.environ.get("VIT_ENC_MM1_WIDE", "1") == "1"      # set 0 to restore the v1.3 4-bank MM1 layout
_MM1_BANKS = (0, 1, 2, 5, 6)   # MM1 rotation; bank 4 = P.V (single large tile), 3/7 = reciprocal transposes
_MM1_WIDE_MAX_KTILES = 2       # only widen when <= this many key tiles per group: at 4 key tiles (N=2048) the
                               # 5-bank rotation collides with P.V/producer banks and measured +7% slower
_K_TILE_SZ = 512  # K tile size for loading and MM1+masking operations
_EXP_TILE_SZ = 512  # Tile size for exp instructions (must equal _K_TILE_SZ)
_LARGE_TILE_SZ = 2048  # Larger tile size for allocations/pipelining (4 x 512 tiles)
_FLASH_ATTENTION_THRESHOLD = 10 * 1024  # Use flash attention above this K/V length
_FLASH_ATTENTION_SECTION_LENGTH = 8 * 1024  # Section size when using flash attention


# ==========================================================================================
# Buffer / config dataclasses (kept verbatim -- the four required by the reached path).
# ==========================================================================================
@dataclass
class AttnConfig(nl.NKIObject):
    """High-level attention configuration set at kernel entry point."""

    seqlen_q: int = None
    seqlen_k_active: int = None
    seqlen_k_prior: int = None
    d: int = None

    tp_q: bool = None
    tp_k: bool = None
    tp_out: bool = None

    is_prefix_caching: bool = None
    causal_mask: bool = None

    use_swa: bool = None
    sliding_window: int = None

    use_cp: bool = None
    global_cp_deg: int = None
    cp_strided_q_slicing: bool = None
    cp_striped_input: bool = None

    scale: float = None
    cache_softmax: bool = None
    skip_output_normalization: bool = None
    dtype: Any = None
    softmax_dtype: Any = None
    mm_out_dtype: Any = None

    is_sequence_packed: bool = None


@dataclass
class AttnInternalBuffers(nl.NKIObject):
    """Container for all SBUF and PSUM tensor buffers used in attention computation."""

    q_sb = None
    k_sb = None
    v_sb = None
    k_sb_prior = None
    v_sb_prior = None

    k_scale_sb = None

    mm1_psum = None
    mm1_copy_sb = None
    mm1_affine_select_output = None
    mm1_masked = None
    mm1_partial_max = None
    mm1_section_max = None
    mm1_running_max = None
    prev_mm1_running_max = None
    flash_attn_correction_factor = None

    exp_sb = None
    exp_partial_sum = None
    exp_section_sum = None
    exp_tp_sb = None
    exp_running_sum = None
    prev_exp_running_sum = None
    exp_sum_reciprocal = None

    mm2_psum = None
    mm2_sb = None
    mm2_prev_output = None
    mm2_accum_flash_attn = None
    mm2_final = None

    tp_flash_attn_correction_factor_psum = None
    tp_flash_attn_correction_factor_sb = None
    tp_exp_sum_reciprocal_psum = None
    tp_exp_sum_reciprocal_sb = None
    mm2_prev_output_scaled = None

    zero_bias_tensor = None
    sink_sb = None
    range_sel_lbs = None
    range_sel_ubs = None
    range_sel_lbs_prior = None
    range_sel_ubs_prior = None
    k_offset_sb_u32 = None
    kbias_sb = None


@dataclass
class SectionParams(nl.NKIObject):
    section_idx = None
    section_offset = None
    section_offset_active = None
    next_section_offset_active = None
    section_contains_prefix = None
    next_section_contains_prefix = None
    kv_section_idx = None


@dataclass
class AttnTileParams(nl.NKIObject):
    """Tile and buffer sizing parameters computed during implementation."""

    seqlen_k_active_updated: int = None
    seqlen_k_prior_padded: int = None
    is_causal: bool = None

    sb_p: int = None

    num_grps: int = None
    num_q_grps_per_load: int = None
    can_pack_q_load: bool = None

    num_large_tiles_per_section: int = None
    num_k_tiles_per_section: int = None
    num_v_tiles_per_section: int = None

    exp_inst_elems: int = None
    num_exp_insts_per_large_tile: int = None

    num_tps_in_mm2_grp: int = None
    mm2_grp_sz: int = None

    use_swa_optimized_allocation: bool = None

    dynamic_sel_mask: bool = None

    section_len: int = None
    num_sections: int = None


# ==========================================================================================
# Entry: _attention_cte. Dead paths removed: prefix caching / APC / sink / seqpack / CP / SWA /
# fused-v-proj / seqlen-q sharding remainder / cache_softmax. This call ALWAYS takes
# use_sbuf_kv=True, causal_mask=False, tp_out=True, single-shard (num_shard=1), bs=1.
# ==========================================================================================
def _attention_cte(
    q: nl.ndarray,
    k: Optional[nl.ndarray] = None,
    v: Optional[nl.ndarray] = None,
    scale: float = 1.0,
    causal_mask: bool = True,
    tp_q: bool = True,
    tp_k: bool = False,
    tp_out: bool = False,
    cache_softmax: bool = False,
    softmax_dtype=nl.float32,
    mm_out_dtype=nl.float32,
    k_cache_sbuf: Optional[List[nl.ndarray]] = None,
    v_cache_sbuf: Optional[List[nl.ndarray]] = None,
    out_o_hbm=None,
    init_sbuf_addr: int = 0,
    seqlen_k_valid: int = None,
    kbias_sbuf=None,
):
    """One head of context-encoding attention (SBUF-KV path).
    v1.9: ``kbias_sbuf`` (optional, [128, seqlen_k] fp32 SBUF, every partition row identical) is added to the scaled
    scores before softmax: 0 for attended keys, a large negative value for masked (padding) keys. None keeps the
    v1.8 instruction stream exactly.
    EDIT: ``seqlen_k_valid`` overrides the active key length (default len(k_cache)*512)
    so keys at positions >= seqlen_k_valid (token padding) are skipped by every loop bound.

    LIVE contract for the K2 encoder: ``k=v=None`` (K/V handed in via ``k_cache_sbuf`` /
    ``v_cache_sbuf``), ``causal_mask=False``, ``tp_q=False, tp_k=False, tp_out=True``,
    ``cache_softmax=False``, ``out_o_hbm`` is a LIST of SBUF slices into the caller's O tile.

    DEAD (removed): prefix caching (k_prior/v_prior/prior_used_len), sink, sliding_window, CP
    (cp_offset/global_cp_deg/striped/strided), sequence packing (bound_min/bound_max),
    skip_output_normalization, fp8 k_scale, and the fused v-projection (h_hbm/wv_t/bv).
    """
    is_prefix_caching = False

    use_sbuf_kv = (k_cache_sbuf is not None and len(k_cache_sbuf) > 0) or (
        v_cache_sbuf is not None and len(v_cache_sbuf) > 0
    )
    kernel_assert(
        k_cache_sbuf is not None and len(k_cache_sbuf) > 0, "k_cache_sbuf must be provided when using sbuf kv"
    )
    kernel_assert(
        v_cache_sbuf is not None and len(v_cache_sbuf) > 0, "v_cache_sbuf must be provided when using sbuf kv"
    )

    seqlen_q, seqlen_k_active, seqlen_k_prior, d, out_shape, softmax_shape = _check_input_and_return_shape(
        q,
        k,
        v,
        is_prefix_caching,
        None,
        None,
        None,
        tp_q,
        tp_k,
        tp_out,
        cache_softmax,
        k_cache_sbuf,
        v_cache_sbuf,
        None,
        None,
        skip_v_check=False,
    )
    if seqlen_k_valid is not None:
        kernel_assert(seqlen_k_valid <= seqlen_k_active, "seqlen_k_valid exceeds K cache length")
        seqlen_k_active = seqlen_k_valid

    # Use out_o_hbm if provided (LIVE: it is a list of SBUF slices).
    if out_o_hbm is not None:
        result = out_o_hbm
    else:
        result = nl.ndarray(shape=out_shape, dtype=q.dtype, buffer=nl.shared_hbm)

    bs = q.shape[0]
    bs_kv = 1  # sbuf-kv: batch_size_kv is per-batch tiles

    kernel_assert(bs > 0, f"Batch size must be positive, got {bs}")
    kernel_assert(bs <= _MAX_BS, f"attention_cte kernel is not tested for batch size above {_MAX_BS}, got {bs}.")
    kernel_assert(bs_kv > 0, f"Batch size must be positive, got {bs_kv}")
    kernel_assert(bs % bs_kv == 0, f"Q batch size must be a multiple of KV batch size, got {bs=}, {bs_kv=}")

    seqlen_k_total = seqlen_k_active
    kernel_assert(seqlen_q <= _MAX_SEQLEN, f"seqlen above {_MAX_SEQLEN}, got {seqlen_q=}.")
    kernel_assert(seqlen_k_total <= _MAX_SEQLEN, f"seqlen above {_MAX_SEQLEN}, got {seqlen_k_total=}.")
    bs_seqlen_qk_product = float(bs * seqlen_q) * seqlen_k_total
    if bs_seqlen_qk_product > _MAX_BS_TIMES_SEQLEN_QK:
        logger.warn(
            f"attention_cte kernel is not tested for batch size x seqlen_q x seqlen_k above "
            f"{_MAX_BS_TIMES_SEQLEN_QK}, got {bs_seqlen_qk_product=}.",
        )

    kernel_assert(d > 0, f"d must be > 0, got {d=}.")
    kernel_assert(d <= _MAX_HEAD_DIM, f"we do not support head_dim > {_MAX_HEAD_DIM}, got head dim {d}")

    kernel_assert(
        (str(mm_out_dtype) == str(nl.float32))
        or (str(mm_out_dtype) == str(nl.bfloat16) and (nisa.get_nc_version() >= nisa.nc_version.gen4)),
        f"mm_out_dtype (psum) should be in [float32, bfloat16], got {mm_out_dtype} in {nisa.get_nc_version()}.",
    )

    ac = AttnConfig(
        seqlen_q=seqlen_q,
        seqlen_k_active=seqlen_k_active,
        seqlen_k_prior=seqlen_k_prior,
        d=d,
        tp_q=tp_q,
        tp_k=tp_k,
        tp_out=tp_out,
        is_prefix_caching=is_prefix_caching,
        causal_mask=causal_mask,
        use_swa=False,
        sliding_window=0,
        use_cp=False,
        global_cp_deg=None,
        cp_strided_q_slicing=False,
        cp_striped_input=False,
        scale=scale,
        cache_softmax=cache_softmax,
        skip_output_normalization=False,
        dtype=q.dtype,
        softmax_dtype=softmax_dtype,
        mm_out_dtype=mm_out_dtype,
        is_sequence_packed=False,
    )

    # Sharding info. use_sbuf_kv disables internal sharding (caller already sharded), so
    # num_shard=1, shard_id=0, num_bs_per_shard=bs, has_remainder=False.
    get_verified_program_sharding_info("attention_cte", max_sharding=2)
    num_shard = 1
    shard_id = 0
    num_bs_per_shard = bs // num_shard
    bs_offset = shard_id * num_bs_per_shard

    for batch_idx in range(num_bs_per_shard):
        # use_sbuf_kv: bs is LOCAL batch size and b is the LOCAL index.
        batch_id = batch_idx
        kv_batch_id = _q_to_kv_batch_id(batch_id, bs, bs_kv)
        _attention_cte_impl(
            q,
            k,
            v,
            None,
            None,
            None,
            result,
            batch_id,
            kv_batch_id,
            ac,
            k_cache_sbuf=k_cache_sbuf,
            v_cache_sbuf=v_cache_sbuf,
            init_sbuf_addr=init_sbuf_addr,
            kbias_sbuf=kbias_sbuf,
        )

    # has_remainder is False (num_shard==1); the seqlen-q sharding path is DEAD.
    return result


def _attention_cte_impl(
    q,
    k_active,
    v_active,
    k_prior,
    v_prior,
    prior_used_len,
    o,
    batch_id,
    batch_id_kv,
    ac: AttnConfig,
    k_cache_sbuf: Optional[List[nl.ndarray]] = None,
    v_cache_sbuf: Optional[List[nl.ndarray]] = None,
    init_sbuf_addr: int = 0,
    kbias_sbuf=None,
):
    """Single-batch attention body. LIVE path only: single section (num_sections==1), pipelined
    group loop (num_grps=16 > 1), SBUF-KV (no K/V HBM load), tp_out=True.

    DEAD (removed): seqlen-q sharding, sink, SWA memset, sections>1 flash accumulation, K/V HBM
    load (_load_k_tile / _load_v_tile), SBUF-prior storage, cache_softmax intermediate writes,
    the non-pipelined (shard_seqlen_q_length<=1) branch.
    """
    # Not seqlen-sharded: start at group 0, cover all groups.
    atp = _compute_tile_parameters(ac, is_seqlen_sharded=False)
    shard_seqlen_q_start = 0
    shard_seqlen_q_length = atp.num_grps

    allocator = ModularAllocator(initial_address=init_sbuf_addr)
    bufs = AttnInternalBuffers()
    bufs.k_scale_sb = None
    bufs.kbias_sb = kbias_sbuf

    bufs.zero_bias_tensor = allocator.alloc_sbuf_tensor(shape=(atp.sb_p, 1), dtype=nl.float32)
    nisa.memset(bufs.zero_bias_tensor, 0.0)

    # _setup_range_select_bounds is a NO-OP here: dynamic_sel_mask=False and not prefix caching.
    _setup_range_select_bounds(ac, atp, bufs, allocator, None, None, None, None, batch_id=batch_id)

    bufs.mm1_running_max = allocator.alloc_sbuf_tensor(shape=(atp.sb_p, atp.num_grps), dtype=nl.float32)
    bufs.exp_running_sum = allocator.alloc_sbuf_tensor(shape=(atp.sb_p, atp.num_grps), dtype=nl.float32)
    bufs.exp_sum_reciprocal = allocator.alloc_sbuf_tensor(shape=(atp.sb_p, atp.num_grps), dtype=nl.float32)

    sbuf_addr_outer = allocator.get_current_address()

    # Single section (num_sections == 1).
    for section_idx in range(atp.num_sections):
        section_offset = atp.section_len * section_idx
        section_offset_active, section_contains_prefix = _compute_section_offset_active(
            section_offset, ac.is_prefix_caching, atp.seqlen_k_prior_padded
        )
        next_section_offset = atp.section_len * (section_idx + 1)
        next_section_offset_active, next_section_contains_prefix = _compute_section_offset_active(
            next_section_offset, ac.is_prefix_caching, atp.seqlen_k_prior_padded
        )
        sp = SectionParams(
            section_idx=section_idx,
            section_offset=section_offset,
            section_offset_active=section_offset_active,
            next_section_offset_active=next_section_offset_active,
            section_contains_prefix=section_contains_prefix,
            next_section_contains_prefix=next_section_contains_prefix,
        )

        allocator.set_current_address(sbuf_addr_outer)
        _allocate_attention_buffers(allocator, ac, atp, bufs, sink=None,
                                    k_cache_sbuf=k_cache_sbuf, v_cache_sbuf=v_cache_sbuf)
        sbuf_addr = allocator.get_current_address()

        # SBUF-KV: K and V are already resident, so the HBM K/V load block is skipped entirely.

        shard_seqlen_q_end = shard_seqlen_q_start + shard_seqlen_q_length
        if shard_seqlen_q_length == 1:
            # v1.5: one Q group (N <= 128 tokens) -- no pipelining (the pipelined prologue below assumes >= 2
            # groups and indexed a non-existent second Q tile, so N=128 failed to compile in <= v1.4).
            _load_q_impl(shard_seqlen_q_start, ac, atp, sp, bufs, q, batch_id, sbuf_addr)
            _qk_and_max_impl(shard_seqlen_q_start, ac, atp, sp, bufs, batch_id)
            _update_max_impl(shard_seqlen_q_start, ac, atp, sp, bufs, None)
            _exp_impl(shard_seqlen_q_start, ac, atp, sp, bufs, None)
            _pv_impl(shard_seqlen_q_start, ac, atp, sp, bufs)
            _write_back_impl(shard_seqlen_q_start, ac, atp, sp, bufs, o, batch_id)
        else:
            # Software pipelining: >= 2 Q groups.
            _load_q_impl(shard_seqlen_q_start, ac, atp, sp, bufs, q, batch_id, sbuf_addr)
            _qk_and_max_impl(shard_seqlen_q_start, ac, atp, sp, bufs, batch_id)
            _update_max_impl(shard_seqlen_q_start, ac, atp, sp, bufs, None)
            _exp_impl(shard_seqlen_q_start, ac, atp, sp, bufs, None)

            _load_q_impl(shard_seqlen_q_start + 1, ac, atp, sp, bufs, q, batch_id, sbuf_addr)
            _qk_and_max_impl(shard_seqlen_q_start + 1, ac, atp, sp, bufs, batch_id)
            _update_max_impl(shard_seqlen_q_start + 1, ac, atp, sp, bufs, None)

            for grp_i in range(shard_seqlen_q_start, shard_seqlen_q_end - 2):
                # grp_i: PV, write_back ; grp_i+1: EXP ; grp_i+2: load Q, QK+Max.
                _load_q_impl(grp_i + 2, ac, atp, sp, bufs, q, batch_id, sbuf_addr)
                _exp_impl(grp_i + 1, ac, atp, sp, bufs, None)
                _fused_qkmax_and_pv_impl(grp_i, ac, atp, sp, bufs, batch_id)
                _write_back_impl(grp_i, ac, atp, sp, bufs, o, batch_id)
                _update_max_impl(grp_i + 2, ac, atp, sp, bufs, None)

            _pv_impl(shard_seqlen_q_end - 2, ac, atp, sp, bufs)
            _write_back_impl(shard_seqlen_q_end - 2, ac, atp, sp, bufs, o, batch_id)
            _exp_impl(shard_seqlen_q_end - 1, ac, atp, sp, bufs, None)
            _pv_impl(shard_seqlen_q_end - 1, ac, atp, sp, bufs)
            _write_back_impl(shard_seqlen_q_end - 1, ac, atp, sp, bufs, o, batch_id)

    # cache_softmax intermediate writes (out_neg_max / out_sum_recip) are DEAD (cache_softmax=False).


def _compute_tile_parameters(ac: AttnConfig, is_seqlen_sharded: bool) -> AttnTileParams:
    """Tile parameters for the LIVE config. Dead: SWA/prefix/CP scale checks and seqlen updates,
    flash sectioning (total_seqlen_k=2048 <= 10240 -> single section)."""
    atp = AttnTileParams()

    # LIVE: causal_mask=False, no CP/SWA/seqpack -> is_causal=False, dynamic_sel_mask=False.
    atp.is_causal = ac.causal_mask
    atp.dynamic_sel_mask = False
    atp.seqlen_k_active_updated = ac.seqlen_k_active
    atp.use_swa_optimized_allocation = False

    atp.sb_p = nl.tile_size.pmax
    kernel_assert(_Q_GRP_SZ == atp.sb_p, f"expect Q group size == pmax, got {_Q_GRP_SZ=}, {atp.sb_p=}")
    kernel_assert(_V_TILE_SZ == atp.sb_p, f"expect V tile size == pmax, got {_V_TILE_SZ=}, {atp.sb_p=}")

    atp.num_grps = div_ceil(ac.seqlen_q, atp.sb_p)
    atp.can_pack_q_load = not is_seqlen_sharded

    num_q_grps_per_load_dtype = 4 if ac.dtype == nl.float32 else 8
    atp.num_q_grps_per_load = min(num_q_grps_per_load_dtype if atp.can_pack_q_load else 1, atp.num_grps)
    kernel_assert(atp.num_q_grps_per_load > 0, f"num_q_grps_per_load must be positive, got {atp.num_q_grps_per_load}.")

    atp.seqlen_k_prior_padded = None  # not prefix caching
    total_seqlen_k = atp.seqlen_k_active_updated

    # total_seqlen_k = 2048 <= 10240 -> no flash attention, single section.
    use_flash_attn = total_seqlen_k > _FLASH_ATTENTION_THRESHOLD
    if use_flash_attn:
        atp.section_len = min(total_seqlen_k, _FLASH_ATTENTION_SECTION_LENGTH)
    else:
        atp.section_len = total_seqlen_k

    kernel_assert(atp.section_len > 0, f"section_len must be positive, got {atp.section_len}")
    atp.num_sections = div_ceil(total_seqlen_k, atp.section_len)
    if not use_flash_attn:
        kernel_assert(atp.num_sections == 1, "Logic fault, must only have 1 section if not using flash_attn")

    atp.num_large_tiles_per_section = div_ceil(atp.section_len, _LARGE_TILE_SZ)
    atp.num_k_tiles_per_section = div_ceil(atp.section_len, _K_TILE_SZ)
    atp.num_v_tiles_per_section = div_ceil(atp.section_len, _V_TILE_SZ)

    atp.exp_inst_elems = _EXP_TILE_SZ
    atp.num_exp_insts_per_large_tile = _LARGE_TILE_SZ // atp.exp_inst_elems
    atp.num_tps_in_mm2_grp = _K_TILE_SZ // atp.sb_p
    atp.mm2_grp_sz = _K_TILE_SZ

    return atp


def _setup_range_select_bounds(
    ac: AttnConfig,
    atp: AttnTileParams,
    bufs: AttnInternalBuffers,
    allocator: ModularAllocator,
    cp_offset: Any,
    prior_used_len: Any,
    bound_min: Any,
    bound_max: Any,
    batch_id: int = 0,
) -> tuple:
    """NO-OP for this config: only CP / SWA / prefix-caching / sequence-packing populate range
    bounds, and all are off (``dynamic_sel_mask=False``, ``is_prefix_caching=False``). The entire
    vendor body was gated by ``if atp.dynamic_sel_mask:`` / ``if ac.is_prefix_caching:`` / SWA,
    so it emits no instructions here and is removed."""
    return


def _allocate_attention_buffers(
    allocator: ModularAllocator,
    ac: AttnConfig,
    atp: AttnTileParams,
    bufs: AttnInternalBuffers,
    sink: Any,
    k_cache_sbuf: Optional[List[nl.ndarray]] = None,
    v_cache_sbuf: Optional[List[nl.ndarray]] = None,
):
    """Allocate SBUF/PSUM buffers. LIVE: tp_out=True (so all tp_out branches taken), K/V come
    from k_cache_sbuf/v_cache_sbuf, sink=None, dynamic_sel_mask=False (so the static-mask
    mm1_copy_sb / mm1_affine_select_output buffers ARE allocated -- kept). The tp_out=False
    branches and the fused-v-proj branch are DEAD and removed."""
    mm2_p, mm2_n = atp.sb_p, ac.d

    # K/V provided in SBUF by the caller.
    bufs.k_sb = k_cache_sbuf
    bufs.v_sb = v_cache_sbuf

    bufs.q_sb = allocator.alloc_sbuf_tensor(
        shape=(ac.d, atp.sb_p * atp.num_q_grps_per_load),
        dtype=nl.bfloat16,
        block_dim=[div_ceil(atp.num_grps, atp.num_q_grps_per_load)],
        num_free_tiles=[2],
        align_to=32,
    )

    bufs.flash_attn_correction_factor = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    mm1_partial_max_n_elts = atp.num_k_tiles_per_section  # sink is None
    bufs.mm1_partial_max = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, mm1_partial_max_n_elts), dtype=nl.float32,
        block_dim=[atp.num_grps], num_free_tiles=[2], align_to=4,
    )

    bufs.mm1_section_max = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    n_final_reduce_sum_elts = div_ceil(atp.section_len, atp.exp_inst_elems)  # sink is None
    bufs.exp_partial_sum = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, n_final_reduce_sum_elts), dtype=nl.float32,
        block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    bufs.exp_section_sum = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    bufs.prev_mm1_running_max = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    bufs.prev_exp_running_sum = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, 1), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    # tp_out=True everywhere below.
    bufs.mm2_prev_output = allocator.alloc_sbuf_tensor(
        shape=(ac.d, atp.sb_p), dtype=ac.mm_out_dtype, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    bufs.mm2_accum_flash_attn = allocator.alloc_sbuf_tensor(
        shape=(ac.d, atp.sb_p), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    bufs.mm2_prev_output_scaled = allocator.alloc_sbuf_tensor(
        shape=(ac.d, atp.sb_p), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    bufs.tp_flash_attn_correction_factor_psum = []
    bufs.tp_exp_sum_reciprocal_psum = []
    for grp_idx in range(atp.num_grps):
        tp_flash_attn_correction_factor_psum_tile = nl.ndarray(
            (ac.d, atp.sb_p), dtype=nl.float32, buffer=nl.psum,
            address=(0, ((grp_idx % 2) * 4 + 3) * PSUM_BANK_SIZE),
        )
        bufs.tp_flash_attn_correction_factor_psum.append(tp_flash_attn_correction_factor_psum_tile)
        tp_exp_sum_reciprocal_psum_tile = nl.ndarray(
            (ac.d, atp.sb_p), dtype=nl.float32, buffer=nl.psum,
            address=(0, ((grp_idx % 2) * 4 + 3) * PSUM_BANK_SIZE),
        )
        bufs.tp_exp_sum_reciprocal_psum.append(tp_exp_sum_reciprocal_psum_tile)

    bufs.tp_flash_attn_correction_factor_sb = allocator.alloc_sbuf_tensor(
        shape=(ac.d, atp.sb_p), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[4],
    )

    bufs.tp_exp_sum_reciprocal_sb = allocator.alloc_sbuf_tensor(
        shape=(ac.d, atp.sb_p), dtype=nl.float32, block_dim=[atp.num_grps], num_free_tiles=[4],
    )

    bufs.mm2_final = allocator.alloc_sbuf_tensor(
        shape=(ac.d, atp.sb_p), dtype=ac.mm_out_dtype, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    bufs.mm2_sb = allocator.alloc_sbuf_tensor(
        shape=(mm2_n, mm2_p), dtype=ac.mm_out_dtype, block_dim=[atp.num_grps], num_free_tiles=[2],
    )

    bufs.mm1_masked = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, _LARGE_TILE_SZ), dtype=nl.float32,
        block_dim=[atp.num_grps, atp.num_large_tiles_per_section],
        num_free_tiles=[2, atp.num_large_tiles_per_section],
    )

    bufs.exp_sb = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, _LARGE_TILE_SZ), dtype=nl.bfloat16,
        block_dim=[atp.num_grps, atp.num_large_tiles_per_section],
        # DP2: group-axis depth 2 so exp(g+1) runs while transpose(g) still drains exp_sb(g).
        num_free_tiles=[2, atp.num_large_tiles_per_section],
    )

    # MM1 PSUM tiles: banks {0,1,2,3}.
    bufs.mm1_psum = []
    for grp_idx in range(atp.num_grps):
        grp_row = []
        for large_tile_idx in range(atp.num_large_tiles_per_section):
            tile_row = []
            for k_tile_idx in range(4):
                mm1_psum_tile = nl.ndarray(
                    (atp.sb_p, nl.tile_size.psum_fmax), dtype=ac.mm_out_dtype, buffer=nl.psum,
                    # v1.4: rotate MM1 over 5 banks (0,1,2,5,6) across consecutive Q groups instead
                    # of reusing banks 0-3 per group, so a QK matmul does not wait for the Vector row-max evict
                    # of the tile 4 back. Bank 4 = P.V, 3/7 = reciprocal transposes. Banks 5/6 are also the
                    # producer's projection banks; the tile framework orders reuse by address, so overlap
                    # with a later image's producer stays correct (it may serialise).
                    address=(0, (_MM1_BANKS[(grp_idx * 4 + k_tile_idx) % len(_MM1_BANKS)]
                                 if (_MM1_WIDE and atp.num_k_tiles_per_section <= _MM1_WIDE_MAX_KTILES)
                                 else (k_tile_idx % 4)) * PSUM_BANK_SIZE),
                )
                tile_row.append(mm1_psum_tile)
            grp_row.append(tile_row)
        bufs.mm1_psum.append(grp_row)

    # dynamic_sel_mask=False -> the static-mask staging buffers are allocated (no-mask branch
    # does not read them, but the buffers exist in the vendor's non-dynamic allocation).
    bufs.mm1_copy_sb = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, nl.tile_size.psum_fmax), dtype=ac.mm_out_dtype,
        block_dim=[atp.num_grps, atp.num_large_tiles_per_section, 4], num_free_tiles=[1, 1, 2],
    )
    bufs.mm1_affine_select_output = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, nl.tile_size.psum_fmax), dtype=ac.mm_out_dtype,
        block_dim=[atp.num_grps, atp.num_large_tiles_per_section, 4], num_free_tiles=[1, 1, 2],
    )

    bufs.exp_tp_sb = allocator.alloc_sbuf_tensor(
        shape=(atp.sb_p, atp.mm2_grp_sz), dtype=nl.bfloat16,
        block_dim=[atp.num_grps, atp.num_large_tiles_per_section, atp.num_tps_in_mm2_grp],
        # DP4: 8 group-blocks resident so P^T transposes issue as soon as their exp exists.
        num_free_tiles=[8, atp.num_large_tiles_per_section, atp.num_tps_in_mm2_grp],
        align_to=32,
    )

    # MM2 PSUM tiles: banks {4,5,6,7}. tp_out=True shape (mm2_n, mm2_p).
    bufs.mm2_psum = []
    for grp_idx in range(atp.num_grps):
        grp_row = []
        for large_tile_idx in range(atp.num_large_tiles_per_section):
            mm2_psum_tile = nl.ndarray(
                (mm2_n, mm2_p), dtype=ac.mm_out_dtype, buffer=nl.psum,
                address=(0, ((4 + (large_tile_idx % 4)) * PSUM_BANK_SIZE)),
            )
            grp_row.append(mm2_psum_tile)
        bufs.mm2_psum.append(grp_row)


def _q_to_kv_batch_id(batch_id: int, bs: int, bs_kv: int) -> int:
    """Map Q batch id to KV batch id for native GQA support. bs_kv=1 here -> always 0."""
    return batch_id // (bs // bs_kv)


def _check_input_and_return_shape(
    q,
    k,
    v,
    is_prefix_caching,
    k_prior,
    v_prior,
    prior_used_len,
    tp_q,
    tp_k,
    tp_out,
    cache_softmax,
    k_cache_sbuf=None,
    v_cache_sbuf=None,
    k_prior_sbuf=None,
    v_prior_sbuf=None,
    skip_v_check=False,
) -> tuple:
    """Validate shapes and return (seqlen_q, seqlen_k, seqlen_k_prior, d, out_shape,
    cache_softmax_shape). LIVE: use_sbuf_kv=True, tp_q=False, tp_out=True, not prefix caching,
    cache_softmax=False."""
    # SBUF-KV path.
    kernel_assert(k_cache_sbuf is not None and len(k_cache_sbuf) > 0, "k_cache_sbuf must be provided")
    kernel_assert(v_cache_sbuf is not None and len(v_cache_sbuf) > 0, "v_cache_sbuf must be provided")

    if tp_q:
        batch_size, seqlen_q, d = q.shape
    else:
        batch_size, d, seqlen_q = q.shape

    d_from_k = k_cache_sbuf[0].shape[0]
    kernel_assert(d == d_from_k, f"Head dimension mismatch: q has d={d}, but k_cache_sbuf has d={d_from_k}")
    seqlen_k = len(k_cache_sbuf) * _K_TILE_SZ
    batch_size_kv = 1
    seqlen_k_prior = None  # not prefix caching

    out_seqlen = seqlen_q
    if tp_out:
        out_shape = (batch_size, d, out_seqlen)
    else:
        out_shape = (batch_size, out_seqlen, d)

    cache_softmax_tile_size = 128
    padded_seq_grps = div_ceil(out_seqlen, cache_softmax_tile_size)
    cache_softmax_shape = [batch_size, cache_softmax_tile_size, padded_seq_grps]

    return (seqlen_q, seqlen_k, seqlen_k_prior, d, out_shape, cache_softmax_shape)


def _load_q_tile(q, out, tp_q, batch_id, grp_i, seqlen_offset, load_dtype, sbuf_addr, grps_per_load=1) -> int:
    """Load a Q tile HBM->SBUF. LIVE: tp_q=False (q is [bs, d, seqlen]), no transpose; the whole
    ``if tp_q:`` transpose sub-tree (dma_transpose / nc_transpose paths) is DEAD and removed."""
    kernel_assert(str(load_dtype) == str(out[0].dtype), "Conflicting dtype")
    _, d, seqlen = q.shape

    num_f = min(seqlen - seqlen_offset, _Q_GRP_SZ * grps_per_load)
    out_dst_pat = out[grp_i // grps_per_load].ap(pattern=[[_Q_GRP_SZ * grps_per_load, d], [1, num_f]], offset=0)
    q_src_pat = q.ap(pattern=[[seqlen, d], [1, num_f]], offset=batch_id * d * seqlen + seqlen_offset)
    nisa.dma_copy(dst=out_dst_pat, src=q_src_pat)


def _get_kv_tile_apc(
    is_prefix_caching,
    k_active,
    k_prior,
    seqlen_active,
    seqlen_prior,
    seqlen_offset,
    load_offset_active,
) -> tuple:
    """Determine which KV tensor (active or prior) to use. LIVE: is_prefix_caching=False, so this
    returns the trivial active-tile mapping; the prefix-caching branch is DEAD and removed."""
    return k_active, seqlen_active, seqlen_offset, load_offset_active


def _compute_section_offset_active(k_section_offset, is_prefix_caching, seqlen_k_prior_padded):
    """Active-K offset for a section. LIVE: is_prefix_caching=False -> pass-through."""
    section_contains_prefix = False
    k_section_offset_active = k_section_offset
    return k_section_offset_active, section_contains_prefix


def _load_q_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
                 bufs: AttnInternalBuffers, q, batch_id, sbuf_addr):
    """Load a Q group if compute is needed. LIVE: is_causal=False -> has_any_compute is always
    True; k_scale_sb is None (fp8 dequant path DEAD)."""
    if grp_i % atp.num_q_grps_per_load == 0:
        # is_causal=False -> has_any_compute_pred is True.
        q_seqlen_offset = grp_i * _Q_GRP_SZ
        _load_q_tile(
            q, bufs.q_sb, ac.tp_q, batch_id, grp_i, q_seqlen_offset,
            bufs.q_sb[0].dtype, sbuf_addr, grps_per_load=atp.num_q_grps_per_load,
        )
        # k_scale_sb is None -> the delayed-fp8 Q scaling is DEAD.


def _qk_and_max_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
                     bufs: AttnInternalBuffers, batch_id: int = 0):
    """MM1 (QK^T) + row-max for a Q group. LIVE: is_causal=False -> always compute.

    h10 EDIT (site :2876): the ``mm1_partial_max`` init memset is pinned to GpSimd -- it sheds
    Vector (K2's binding engine) and its first consumer is several ops downstream, so a slower
    idle engine is safe."""
    # is_causal=False -> has_any_compute_pred is True.
    nisa.memset(bufs.mm1_partial_max[grp_i], value=_FLOAT32_MIN, engine=engine.gpsimd)

    for large_tile_idx in range(atp.num_large_tiles_per_section):
        _qk_and_max_large_tile_impl(grp_i, large_tile_idx, ac, atp, sp, bufs, batch_id)


def _update_max_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
                     bufs: AttnInternalBuffers, sink):
    """Update running max. LIVE: is_causal=False (always compute), sink=None, num_sections==1 so
    the reduce writes straight into mm1_running_max (max-direct promotion) and the flash
    multi-section correction is DEAD."""
    # is_causal=False -> has_any_compute_pred is True.

    # num_sections==1: the reduce writes straight into mm1_running_max (max-direct); the [128,1]
    # copy the vendor once had here is gone (3,072 Vector instructions removed upstream).
    _max_dst = bufs.mm1_running_max[:, grp_i]
    _n_exec = min(div_ceil(atp.seqlen_k_active_updated, _K_TILE_SZ), _LARGE_TILE_SZ // _K_TILE_SZ,
                  atp.num_k_tiles_per_section)
    if atp.num_large_tiles_per_section == 1 and _n_exec >= 2:
        # v1.3: the chained row-max left the whole row max in the last executed k-tile's column.
        nisa.tensor_scalar(dst=_max_dst, data=bufs.mm1_partial_max[grp_i][:, _n_exec - 1:_n_exec],
                           op0=nl.multiply, operand0=-1.0)
    else:
        nisa.tensor_reduce(_max_dst, nl.maximum, bufs.mm1_partial_max[grp_i], 1, negate=True)
    # num_sections==1: no running-max copy / flash correction factor. sink block DEAD (sink=None).


def _exp_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
              bufs: AttnInternalBuffers, sink):
    """Softmax exp (Scalar ACTIVATE) + P^T dma_transposes. LIVE: is_causal=False, use_swa=False,
    not prefix caching -> every k-tile is emitted (emit_idxs is full), wide_exp=False, sink=None.

    The ``memset`` on ``exp_partial_sum`` stays UNPINNED (engine.unknown / Vector) ON PURPOSE
    (vendor :2993): pinning it would fork this 280-line function and re-home the softmax-exp and
    dma_transpose source strings, which MUST keep their ``attention_inner`` sites stable and
    match the shipping stream. The ACCUM FORK (reset on first tile, accumulate, drain once) and
    the reduce_cmd discipline are preserved verbatim."""
    # is_causal=False -> has_any_compute_pred is True.
    q_seqlen_offset = grp_i * atp.sb_p
    # v1.2: with one large tile only the last emitted column is ever read (see _write_back_impl),
    # and it is written by the chained exp-sum, so the zero-fill is skipped.
    if atp.num_large_tiles_per_section > 1:
        nisa.memset(bufs.exp_partial_sum[grp_i][...], value=0.0)  # UNPINNED on purpose (:2993).

    for large_tile_idx in range(atp.num_large_tiles_per_section):
        kernel_assert(atp.exp_inst_elems == 512, "Internal validation failed.")

        # ACCUM FORK pre-scan: is_causal=False and not prefix caching -> every tile emits, so
        # emit_idxs is [0, 1, 2, 3]. reset on the first, accumulate, drain once.
        emit_idxs = []
        for scan_idx in range(atp.num_exp_insts_per_large_tile):
            scan_prior, scan_seqlen_k, scan_k_start, _ = _get_kv_tile_apc(
                ac.is_prefix_caching,
                False,
                True,
                atp.seqlen_k_active_updated,
                ac.seqlen_k_prior,
                sp.section_offset + large_tile_idx * _LARGE_TILE_SZ + scan_idx * atp.exp_inst_elems,
                None,
            )
            scan_mask = True  # is_causal=False, use_swa=False
            if scan_mask and scan_seqlen_k > scan_k_start:
                emit_idxs.append(scan_idx)
        first_emit = emit_idxs[0] if emit_idxs else -1
        last_emit = emit_idxs[-1] if emit_idxs else -1

        for exp_tile_idx in range(atp.num_exp_insts_per_large_tile):
            is_prior_tile, seqlen_k, k_start_pos, _ = _get_kv_tile_apc(
                ac.is_prefix_caching,
                False,
                True,
                atp.seqlen_k_active_updated,
                ac.seqlen_k_prior,
                sp.section_offset + large_tile_idx * _LARGE_TILE_SZ + exp_tile_idx * atp.exp_inst_elems,
                None,
            )
            num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
            num_f = min(seqlen_k - k_start_pos, atp.exp_inst_elems)

            exp_sel_mask = True  # is_causal=False, use_swa=False

            if exp_sel_mask and seqlen_k > k_start_pos:
                # Step 1: exp. reduce_cmd is ALWAYS explicit (see the vendor note: an idle cmd
                # with reduce_op/reduce_res would write a stale register).
                if exp_tile_idx == first_emit:
                    exp_reduce_cmd = nisa.reduce_cmd.reset_reduce
                else:
                    exp_reduce_cmd = nisa.reduce_cmd.reduce
                if exp_tile_idx == last_emit:
                    exp_reduce_res = bufs.exp_partial_sum[grp_i][
                        :num_p,
                        large_tile_idx * atp.num_exp_insts_per_large_tile + exp_tile_idx,
                    ]
                else:
                    exp_reduce_res = None
                nisa.activation(
                    bufs.exp_sb[grp_i][large_tile_idx][
                        :num_p, nl.ds(exp_tile_idx * atp.exp_inst_elems, num_f)
                    ],
                    op=nl.exp,
                    data=bufs.mm1_masked[grp_i][large_tile_idx][
                        :num_p, nl.ds(exp_tile_idx * atp.exp_inst_elems, num_f)
                    ],
                    reduce_op=nl.add,
                    reduce_cmd=exp_reduce_cmd,
                    reduce_res=exp_reduce_res,
                    bias=bufs.mm1_running_max[:num_p, grp_i],
                )

                # Step 2: P^T dma_transpose into exp_tp_sb ([128, 4, 128]).
                num_f_outer = num_f // atp.sb_p
                num_f_inner = num_f % atp.sb_p

                if num_f_outer >= 1:
                    nisa.dma_transpose(
                        dst=bufs.exp_tp_sb[grp_i][large_tile_idx][exp_tile_idx].ap(
                            [
                                [atp.mm2_grp_sz, atp.sb_p],
                                [1, 1],
                                [atp.sb_p, num_f_outer],
                                [1, num_p],
                            ]
                        ),
                        src=bufs.exp_sb[grp_i][large_tile_idx].ap(
                            [
                                [_LARGE_TILE_SZ, num_p],
                                [1, 1],
                                [atp.sb_p, num_f_outer],
                                [1, atp.sb_p],
                            ],
                            offset=exp_tile_idx * atp.mm2_grp_sz,
                        ),
                    )

                if num_f_inner > 0:
                    nisa.dma_transpose(
                        dst=bufs.exp_tp_sb[grp_i][large_tile_idx][exp_tile_idx].ap(
                            [
                                [atp.mm2_grp_sz, num_f_inner],
                                [1, 1],
                                [atp.sb_p, 1],
                                [1, num_p],
                            ],
                            offset=num_f_outer * atp.sb_p,
                        ),
                        src=bufs.exp_sb[grp_i][large_tile_idx].ap(
                            [
                                [_LARGE_TILE_SZ, num_p],
                                [1, 1],
                                [atp.sb_p, 1],
                                [1, num_f_inner],
                            ],
                            offset=exp_tile_idx * atp.mm2_grp_sz + num_f_outer * atp.sb_p,
                        ),
                    )

    # sink block DEAD (sink=None).


def _pv_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams, bufs: AttnInternalBuffers):
    """MM2 (P@V) for a Q group. LIVE: is_causal=False -> always compute."""
    # is_causal=False -> has_any_compute_pred is True.
    # v1.2: the mm2_sb zero-fill is skipped when there is a single large tile (N <= 2048, the only
    # configuration this fork serves): its first write is the PSUM->SBUF tensor_copy over exactly the
    # [:, :num_f] region that write-back reads, so the zeros were never observed.
    if atp.num_large_tiles_per_section > 1:
        nisa.memset(bufs.mm2_sb[grp_i][...], value=0.0)

    for large_tile_idx in range(atp.num_large_tiles_per_section):
        _pv_large_tile_impl(grp_i, large_tile_idx, ac, atp, sp, bufs)


def _fused_qkmax_and_pv_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
                             bufs: AttnInternalBuffers, batch_id: int = 0):
    """Pipelined MM2 for grp_i while doing QK+max for grp_i+2. LIVE: is_causal=False -> both
    predicates True.

    h10 EDIT (site :3294): the ``mm1_partial_max[qkmax_grp]`` init memset is pinned to GpSimd
    (same rationale as :2876; 2,688/core, the larger of the two)."""
    qkmax_grp = grp_i + 2
    # is_causal=False -> both has_any_compute predicates are True.
    nisa.memset(bufs.mm1_partial_max[qkmax_grp][...], value=_FLOAT32_MIN, engine=engine.gpsimd)

    for large_tile_idx in range(atp.num_large_tiles_per_section):
        _pv_large_tile_impl(grp_i, large_tile_idx, ac, atp, sp, bufs)
        _qk_and_max_large_tile_impl(qkmax_grp, large_tile_idx, ac, atp, sp, bufs, batch_id)


def _write_back_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
                     bufs: AttnInternalBuffers, o, batch_id):
    """Finalise output. LIVE: is_causal=False (always compute), num_sections==1,
    skip_output_normalization=False -> single-section reciprocal + scale_reciprocal_write_back.
    The flash multi-section accumulation and the skip-normalization branches are DEAD."""
    # is_causal=False -> has_any_compute_pred is True.
    q_seqlen_offset = grp_i * atp.sb_p
    # num_sections==1, not skip_output_normalization.
    if atp.num_large_tiles_per_section == 1:
        # v1.2: with one large tile the chained exp-sum leaves the WHOLE row sum in one column
        # (the last emitted exp tile); every other column is the memset 0.0. Read that column
        # directly instead of reducing the zero-padded row (x + 0.0 is exact, so bit-identical).
        # last emitted exp tile = ceil(seqlen_k / exp_inst_elems) - 1 (clamped to the large tile)
        _last_col = min(div_ceil(atp.seqlen_k_active_updated, atp.exp_inst_elems),
                        _LARGE_TILE_SZ // atp.exp_inst_elems) - 1
        nisa.reciprocal(bufs.exp_sum_reciprocal[:, grp_i], bufs.exp_partial_sum[grp_i][:, _last_col:_last_col + 1])
    else:
        nisa.tensor_reduce(bufs.exp_section_sum[grp_i][...], nl.add, bufs.exp_partial_sum[grp_i], axis=1)
        nisa.reciprocal(bufs.exp_sum_reciprocal[:, grp_i], bufs.exp_section_sum[grp_i])

    num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
    num_f = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)

    # Single section: scale by reciprocal and write current output.
    _scale_reciprocal_write_back_impl(bufs.mm2_sb[grp_i], grp_i, ac, atp, bufs, o, batch_id, num_p, num_f)


def _scale_reciprocal_write_back_impl(src_buf, grp_i, ac: AttnConfig, atp: AttnTileParams,
                                      bufs: AttnInternalBuffers, o, batch_id, num_p, num_f):
    """Reciprocal-normalise and write back (final section). LIVE: tp_out=True.

    h8 EDIT (site 1): the tp_out PSUM->SBUF evict ``tensor_copy`` is pinned to Scalar -- the
    un-pinned copy defaults to Vector (K2's binding engine); GpSimd is illegal (PSUM source), so
    Scalar is the only shedding target and is the exercised majority path. The tp_out=False
    ``tensor_scalar(..., engine=engine.vector)`` branch is DEAD (tp_out always True) and removed."""
    # tp_out=True.
    nisa.nc_transpose(
        bufs.tp_exp_sum_reciprocal_psum[grp_i].ap(pattern=[[atp.sb_p, ac.d], [1, atp.sb_p]], offset=0),
        bufs.exp_sum_reciprocal.ap(pattern=[[atp.num_grps, atp.sb_p], [0, ac.d]], offset=grp_i),
    )
    if _WB_FUSED:
        # v1.5: one Vector multiply reads the transposed reciprocal straight from PSUM and writes bf16 into the
        # caller's O slice -- replaces the Scalar PSUM evict, the fp32 multiply into mm2_final and the GpSimd
        # fp32->bf16 cast (same single rounding to bf16 as before: fp32 product, rounded once on write).
        nisa.tensor_tensor(
            o[grp_i][: ac.d, :num_f],
            src_buf[: ac.d, :num_f],
            bufs.tp_exp_sum_reciprocal_psum[grp_i][: ac.d, :num_f],
            nl.multiply,
        )
        return
    # ---- h8 SITE 1: PSUM->SBUF evict pinned to Scalar. ----
    nisa.tensor_copy(
        bufs.tp_exp_sum_reciprocal_sb[grp_i][: ac.d, :num_f],
        bufs.tp_exp_sum_reciprocal_psum[grp_i][: ac.d, :num_f],
        engine=engine.scalar,
    )
    nisa.tensor_tensor(
        bufs.mm2_final[grp_i][: ac.d, :num_f],
        src_buf[: ac.d, :num_f],
        bufs.tp_exp_sum_reciprocal_sb[grp_i][: ac.d, :num_f],
        nl.multiply,
    )

    _write_back_o_impl(bufs.mm2_final[grp_i], grp_i, ac, atp, o, batch_id, num_p, num_f)


def _write_back_o_impl(src_buf, grp_i, ac: AttnConfig, atp: AttnTileParams, o, batch_id, num_p, num_f):
    """Evict finished O. LIVE: ``o`` is a LIST of caller-owned SBUF slices, tp_out=True, so the
    output is copied into the caller's O tile (kept on-chip for the fused out-projection) rather
    than DMA'd to HBM. The HBM-sink branch is DEAD (o is always a list) and removed.

    h8 EDIT (site 3): the tp_out list-branch ``tensor_copy`` (SBUF->SBUF fp32->bf16 CAST) is
    pinned to GpSimd -- all three engines are legal and GpSimd is the idle engine, so it sheds
    off Vector."""
    # o is a list (caller-owned SBUF slices), tp_out=True.
    # ---- h8 SITE 3: SBUF->SBUF CAST pinned to GpSimd. ----
    nisa.tensor_copy(dst=o[grp_i][: ac.d, :num_f], src=src_buf[: ac.d, :num_f], engine=engine.gpsimd)


def _qk_and_max_large_tile_impl(qkmax_grp, large_tile_idx, ac: AttnConfig, atp: AttnTileParams,
                                sp: SectionParams, bufs: AttnInternalBuffers, batch_id: int = 0):
    """MM1 (QK^T) + masking for a Q group and a 2048 K tile. LIVE: is_causal=False,
    dynamic_sel_mask=False, is_prior_tile=False -> the FINAL ``else: # no masking``
    ``tensor_scalar_reduce`` branch. The diagonal (affine_select) and range_select branches are
    DEAD and removed; SBUF-prefix-caching K-tile selection is DEAD (k_sb_prior is None)."""
    q_seqlen_offset = qkmax_grp * atp.sb_p

    num_k_tiles_in_large_tile = _LARGE_TILE_SZ // _K_TILE_SZ
    # v1.3: chain the row-max across the executed k-tiles of this large tile (reset on the first,
    # accumulate on the rest, write reduce_res only on the last) instead of one read-out per k-tile
    # followed by a tensor_reduce. Executed tiles are those with k_start < seqlen_k (non-causal:
    # every tile up to the active key length). max is exact under any grouping -> bit-identical.
    # single large tile (the only case chained): executed k-tiles are 0 .. ceil(seqlen_k/512)-1
    _n_exec = min(div_ceil(atp.seqlen_k_active_updated, _K_TILE_SZ), num_k_tiles_in_large_tile,
                  atp.num_k_tiles_per_section)
    _chain = (atp.num_large_tiles_per_section == 1 and _n_exec >= 2)
    _first_kt = 0
    _last_kt = _n_exec - 1
    for k_tile_idx in range(num_k_tiles_in_large_tile):
        mm1_psum_tile = bufs.mm1_psum[qkmax_grp][large_tile_idx][k_tile_idx]
        # dynamic_sel_mask=False, but the no-mask branch below reads neither mm1_copy_sb nor
        # mm1_affine_select_output, so their per-tile locals are not needed.
        mm1_masked_tile = bufs.mm1_masked[qkmax_grp][large_tile_idx]
        mm1_partial_max_tile = bufs.mm1_partial_max[qkmax_grp]

        k_tile_idx_in_section = large_tile_idx * num_k_tiles_in_large_tile + k_tile_idx
        _kv_sec_idx = sp.kv_section_idx if sp.kv_section_idx is not None else sp.section_idx
        k_tile_idx_global = atp.num_k_tiles_per_section * _kv_sec_idx + k_tile_idx_in_section
        is_prior_tile, seqlen_k, k_start_pos, _ = _get_kv_tile_apc(
            ac.is_prefix_caching,
            False,
            True,
            atp.seqlen_k_active_updated,
            ac.seqlen_k_prior,
            k_tile_idx_global * _K_TILE_SZ,
            None,
        )

        matmul_selection = True  # is_causal=False

        if q_seqlen_offset >= ac.seqlen_q or k_start_pos >= seqlen_k:
            matmul_selection = False

        if matmul_selection and k_tile_idx_in_section < atp.num_k_tiles_per_section:
            num_f = min(seqlen_k - k_start_pos, _K_TILE_SZ)
            num_q_free = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)

            # Step 1: MM1 matmul. Non-APC path -> use k_tile_idx_in_section.
            k_tile_to_use = bufs.k_sb[k_tile_idx_in_section]

            nisa.nc_matmul(
                mm1_psum_tile[:num_q_free, :num_f],
                bufs.q_sb[qkmax_grp // atp.num_q_grps_per_load][
                    : ac.d,
                    nl.ds((qkmax_grp % atp.num_q_grps_per_load) * _Q_GRP_SZ, num_q_free),
                ],
                k_tile_to_use[:, :num_f],
            )

            # Step 2: no masking (is_causal=False, dynamic_sel_mask=False).
            num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
            num_f = min(seqlen_k - k_start_pos, _K_TILE_SZ)

            if bufs.kbias_sb is not None:
                # v1.9 key-padding bias: s = psum * scale + kbias[k], then the row-max of this k-tile.
                _ko = k_start_pos
                nisa.scalar_tensor_tensor(
                    dst=mm1_masked_tile[:num_p, nl.ds(k_tile_idx * _K_TILE_SZ, num_f)],
                    data=mm1_psum_tile[:num_p, :num_f],
                    op0=nl.multiply,
                    operand0=ac.scale,
                    op1=nl.add,
                    operand1=bufs.kbias_sb[:num_p, nl.ds(_ko, num_f)],
                )
                if _chain:
                    # keep the chained-max contract: the running max over executed k-tiles ends in the last
                    # executed tile's column (read by _update_max_impl).
                    _col = mm1_partial_max_tile[:num_p, k_tile_idx_in_section]
                    nisa.tensor_reduce(_col, nl.maximum,
                                       mm1_masked_tile[:num_p, nl.ds(k_tile_idx * _K_TILE_SZ, num_f)], 1)
                    if k_tile_idx > _first_kt:
                        _prev = mm1_partial_max_tile[:num_p, k_tile_idx_in_section - 1]
                        nisa.tensor_tensor(dst=_col, data1=_col, data2=_prev, op=nl.maximum)
                else:
                    nisa.tensor_reduce(mm1_partial_max_tile[:num_p, k_tile_idx_in_section], nl.maximum,
                                       mm1_masked_tile[:num_p, nl.ds(k_tile_idx * _K_TILE_SZ, num_f)], 1)
            elif _chain:
                nisa.tensor_scalar_reduce(
                    mm1_masked_tile[:num_p, nl.ds(k_tile_idx * _K_TILE_SZ, num_f)],
                    data=mm1_psum_tile[:num_p, :num_f],
                    op0=nl.multiply,
                    operand0=ac.scale,
                    reduce_op=nl.maximum,
                    reduce_cmd=(nisa.reduce_cmd.reset_reduce if k_tile_idx == _first_kt else nisa.reduce_cmd.reduce),
                    reduce_res=(mm1_partial_max_tile[:num_p, k_tile_idx_in_section] if k_tile_idx == _last_kt else None),
                )
            else:
                nisa.tensor_scalar_reduce(
                    mm1_masked_tile[:num_p, nl.ds(k_tile_idx * _K_TILE_SZ, num_f)],
                    data=mm1_psum_tile[:num_p, :num_f],
                    op0=nl.multiply,
                    operand0=ac.scale,
                    reduce_op=nl.maximum,
                    reduce_res=mm1_partial_max_tile[:num_p, k_tile_idx_in_section],
                )


def _pv_large_tile_impl(pv_grp, large_tile_idx, ac: AttnConfig, atp: AttnTileParams,
                        sp: SectionParams, bufs: AttnInternalBuffers):
    """MM2 (P@V) for a Q group and 2048 V tile. LIVE: is_causal=False (all tiles selected),
    tp_out=True, single large tile so ``large_tile_idx==0`` -> the ``tensor_copy`` PSUM->SBUF
    reduction (never the accumulate ``tensor_tensor`` branch). SBUF-prefix V-tile selection is
    DEAD (v_sb_prior is None); the tp_out=False MM2 orientation is DEAD."""
    q_seqlen_offset = pv_grp * atp.sb_p
    num_mm2_grps_in_large_tile = _LARGE_TILE_SZ // atp.mm2_grp_sz
    mm2_psum_set = False

    mm2_psum_tile = bufs.mm2_psum[pv_grp][large_tile_idx]

    # Step 1: matmul + accumulate in PSUM per group.
    for mm2_grp_i in range(num_mm2_grps_in_large_tile):
        num_mm2_per_grp = atp.mm2_grp_sz // _V_TILE_SZ
        num_mm2_per_large_tile = num_mm2_per_grp * num_mm2_grps_in_large_tile

        exp_tp_sb_tile = bufs.exp_tp_sb[pv_grp][large_tile_idx][mm2_grp_i]

        is_prior_tile, seqlen_k, k_start_pos_512_tile, _ = _get_kv_tile_apc(
            ac.is_prefix_caching,
            False,
            True,
            atp.seqlen_k_active_updated,
            ac.seqlen_k_prior,
            sp.section_offset + large_tile_idx * _LARGE_TILE_SZ + mm2_grp_i * atp.mm2_grp_sz,
            None,
        )

        for mm2_i in range(num_mm2_per_grp):
            v_tile_idx = large_tile_idx * num_mm2_per_large_tile + mm2_grp_i * num_mm2_per_grp + mm2_i
            k_start_pos = k_start_pos_512_tile + mm2_i * _V_TILE_SZ

            num_p = min(seqlen_k - k_start_pos, _V_TILE_SZ)
            num_f = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)

            mm2_sel_mask = True  # is_causal=False

            if mm2_sel_mask and v_tile_idx < atp.num_v_tiles_per_section and num_p > 0 and num_f > 0:
                mm2_psum_set = True

                # Non-APC path -> use v_tile_idx.
                v_tile_to_use = bufs.v_sb[v_tile_idx]

                # tp_out=True orientation.
                nisa.nc_matmul(
                    mm2_psum_tile[: ac.d, :num_f],
                    v_tile_to_use[:num_p, : ac.d],
                    exp_tp_sb_tile[:num_p, nl.ds(mm2_i * _V_TILE_SZ, num_f)],
                )

    # Step 2: reduce MM2 groups into the large tile. is_prior_tile=False, is_causal=False ->
    # loop_reduce_sel True; single large tile -> large_tile_idx==0 -> tensor_copy (never the
    # accumulate tensor_tensor branch).
    is_prior_tile, seqlen_k, k_start_pos, _ = _get_kv_tile_apc(
        ac.is_prefix_caching,
        False,
        True,
        atp.seqlen_k_active_updated,
        ac.seqlen_k_prior,
        sp.section_offset + large_tile_idx * _LARGE_TILE_SZ,
        None,
    )
    k_seqlen_mask = k_start_pos < seqlen_k

    loop_reduce_sel = True  # is_causal=False

    if loop_reduce_sel and k_seqlen_mask and mm2_psum_set:
        # tp_out=True.
        num_f = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
        if large_tile_idx == 0:
            nisa.tensor_copy(bufs.mm2_sb[pv_grp][:, :num_f], mm2_psum_tile[:, :num_f])
        else:
            nisa.tensor_tensor(
                bufs.mm2_sb[pv_grp][:, :num_f],
                bufs.mm2_sb[pv_grp][:, :num_f],
                mm2_psum_tile[:, :num_f],
                nl.add,
            )