# 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 @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, ): """One head of context-encoding attention (SBUF-KV path). 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, ) # 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, ): """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.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 _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, )