This document is relevant for: Trn2, Trn3

Attention Block TKG Kernel API Reference#

[Experimental] Implements a fully fused attention block optimized for Token Generation (autoregressive decoding), keeping all intermediate tensors in SBUF to minimize HBM traffic.

The kernel supports:

  • Fused multi-stage computation: pre-normalization, QKV projection, RoPE, post-normalization, attention, KV cache update, and output projection

  • Multiple KV cache layouts: flat (transposed/non-transposed) and block-based

  • Grouped-Query Attention (GQA) with configurable Q/KV head ratios

  • Optional RMSNorm at multiple stages (pre-projection, post-projection per-head)

  • Optional Rotary Position Embedding (RoPE) with configurable layouts

  • Flexible quantization support (FP8, FP16, BF16)

  • FP8 KV cache quantization support

  • Configurable softmax scaling factor

  • Batch processing with per-batch cache indexing

  • Single program multiple data (SPMD) sharding for distributed computation

Background#

The attention_block_tkg kernel combines multiple stages of transformer attention computation into a single fused operation that minimizes data movement between HBM and on-chip memory (SBUF).

Fused Operations:

The kernel fuses the following stages in SBUF to avoid HBM round-trips:

  1. Pre-normalization: Optional RMSNorm on input hidden states

  2. QKV Projection: Linear projection to Query, Key, Value tensors

  3. RoPE: Optional Rotary Position Embedding on Q and K

  4. Post-normalization: Optional per-head RMSNorm on Q and K

  5. Attention Computation: Scaled dot-product attention with KV cache

  6. KV Cache Update: Write new K/V tokens to cache

  7. Output Projection: Linear projection of attention output

Performance Benefits:

By keeping intermediate tensors in SBUF throughout the computation, this kernel achieves:

  • Reduced HBM bandwidth consumption

  • Lower latency for token generation

  • Better hardware utilization through operation fusion

API Reference#

Source code for this kernel API can be found at: attention_block_tkg.py

attention_block_tkg#

nkilib.core.attention_block_tkg.attention_block_tkg.attention_block_tkg(X: nl.NkiTensor, X_hidden_dim_actual: Optional[int], rmsnorm_X_enabled: bool, rmsnorm_X_eps: Optional[float], rmsnorm_X_gamma: Optional[nl.NkiTensor], W_qkv: nl.NkiTensor, bias_qkv: Optional[nl.NkiTensor], quantization_type_qkv: QuantizationType, weight_dequant_scale_qkv: Optional[nl.NkiTensor], input_dequant_scale_qkv: Optional[nl.NkiTensor], rmsnorm_QK_pre_rope_enabled: bool, rmsnorm_QK_pre_rope_eps: float, rmsnorm_QK_pre_rope_W_Q: Optional[nl.NkiTensor], rmsnorm_QK_pre_rope_W_K: Optional[nl.NkiTensor], cos: Optional[nl.NkiTensor], sin: Optional[nl.NkiTensor], rope_contiguous_layout: bool, rmsnorm_QK_post_rope_enabled: bool, rmsnorm_QK_post_rope_eps: float, rmsnorm_QK_post_rope_W_Q: Optional[nl.NkiTensor], rmsnorm_QK_post_rope_W_K: Optional[nl.NkiTensor], K_cache_transposed: bool, active_blocks_table: Optional[nl.NkiTensor], K_cache: nl.NkiTensor, V_cache: nl.NkiTensor, attention_mask: Optional[nl.NkiTensor], sink: Optional[nl.NkiTensor], update_cache: bool, kv_cache_update_idx: Optional[nl.NkiTensor], W_out: Optional[nl.NkiTensor], bias_out: Optional[nl.NkiTensor], quantization_type_out: QuantizationType, weight_dequant_scale_out: Optional[nl.NkiTensor], input_dequant_scale_out: Optional[nl.NkiTensor], transposed_out: bool, out_in_sb: bool, transposed_in: bool = False, softmax_scale: Optional[float] = None, enable_fa_s_prior_tiling: bool = True, fp8_packed: bool = False, k_scale: Optional[nl.NkiTensor] = None, v_scale: Optional[nl.NkiTensor] = None, sbm: Optional[SbufManager] = None, skip_attention: bool = False, is_h_transposed_by_4: bool = False, KVDP: int = 1, KVDP_replica_group: Optional[ReplicaGroup] = None, KVDP_collective_mode: Optional[KVDPCollectiveMode] = None, KVDP_rank: Optional[nl.NkiTensor] = None, pos_ids: Optional[nl.NkiTensor] = None, swa_start_pos_ids: Optional[nl.NkiTensor] = None, S_ctx: Optional[int] = None, max_context_len: Optional[nl.NkiTensor] = None, dtype_mode: DtypeMode = DtypeMode.NON_OCP, CP: int = 1, CP_replica_group: Optional[ReplicaGroup] = None, CP_collective_mode: Optional[CPCollectiveMode] = None, seq_packed_attention_mask: Optional[nl.NkiTensor] = None, seq_packed_active_blocks_table: Optional[nl.NkiTensor] = None, seq_packed_q_index_table: Optional[nl.NkiTensor] = None, seq_packed_seq_id_table: Optional[nl.NkiTensor] = None, seq_packed_accumulator_partial_route_table: Optional[nl.NkiTensor] = None, seq_packed_accumulator_group_state_route_table: Optional[nl.NkiTensor] = None)#

Fused Attention Block for Token Generation (TKG).

Performs end-to-end attention block computation optimized for autoregressive decoding: X → [RMSNorm] → QKV Projection → [RMSNorm Q/K] → [RoPE] → [RMSNorm Q/K] → Attention → KV Cache Update → [Output Projection] → Output

All intermediate tensors remain in SBUF to minimize HBM traffic.

Parameters:

  • X (nl.NkiTensor) – Input hidden states

  • X_hidden_dim_actual (Optional[int]) – Actual hidden dim if X is padded

  • rmsnorm_X_enabled (bool) – Apply RMSNorm to X before QKV projection

  • rmsnorm_X_eps (Optional[float]) – RMSNorm epsilon (default 1e-3)

  • rmsnorm_X_gamma (Optional[nl.NkiTensor]) – [1, H] @ HBM, RMSNorm weights

  • W_qkv (nl.NkiTensor) – [H, d_head*(q_heads + 2*kv_heads)] @ HBM, QKV projection weights. Packs Q (q_heads heads), K (kv_heads heads), V (kv_heads heads) per token.

  • bias_qkv (Optional[nl.NkiTensor]) – [1, d_head*(q_heads + 2*kv_heads)] @ HBM, QKV bias.

  • quantization_type_qkv (QuantizationType) – Type of quantization for QKV projection (NONE, STATIC, ROW).

  • weight_dequant_scale_qkv (Optional[nl.NkiTensor]) – Weight dequantization scale for QKV projection.

  • input_dequant_scale_qkv (Optional[nl.NkiTensor]) – Input dequantization scale for QKV projection.

  • rmsnorm_QK_pre_rope_enabled (bool) – Apply RMSNorm to Q/K before RoPE

  • rmsnorm_QK_pre_rope_eps (float) – Pre-RoPE RMSNorm epsilon

  • rmsnorm_QK_pre_rope_W_Q (Optional[nl.NkiTensor]) – [1, d_head] @ HBM, Pre-RoPE Q gamma weights

  • rmsnorm_QK_pre_rope_W_K (Optional[nl.NkiTensor]) – [1, d_head] @ HBM, Pre-RoPE K gamma weights

  • cos (Optional[nl.NkiTensor]) – [d_head//2, B, S_tkg] @ HBM, RoPE cosine embeddings (None = skip RoPE)

  • sin (Optional[nl.NkiTensor]) – [d_head//2, B, S_tkg] @ HBM, RoPE sine embeddings (None = skip RoPE)

  • rope_contiguous_layout (bool) – True for contiguous halves, False for interleaved

  • rmsnorm_QK_post_rope_enabled (bool) – Apply RMSNorm to Q/K after RoPE

  • rmsnorm_QK_post_rope_eps (float) – Post-RoPE RMSNorm epsilon

  • rmsnorm_QK_post_rope_W_Q (Optional[nl.NkiTensor]) – [1, d_head] @ HBM, Post-RoPE Q weights

  • rmsnorm_QK_post_rope_W_K (Optional[nl.NkiTensor]) – [1, d_head] @ HBM, Post-RoPE K weights

  • K_cache_transposed (bool) – Whether K cache is stored transposed in HBM. If True: K cache is [B, d_head, S_max_ctx]. If False: K cache is [B, S_max_ctx, d_head]. Must be False for block KV cache.

  • active_blocks_table (Optional[nl.NkiTensor]) – @ HBM, Block indices for block KV cache. kv_heads == 1: [B, num_blocks]. kv_heads > 1: [B, kv_heads, num_blocks] (per-head block tables; each (batch, head) pair owns its own slot range so writes never collide). Block indices must use num_blocks * kv_heads addressing.

  • K_cache (nl.NkiTensor) – Key cache @ HBM. Flat KV: [B, kv_heads, d_head, S_max_ctx] if K_cache_transposed else [B, kv_heads, S_max_ctx, d_head]. Block KV: kv_heads == 1: [num_blocks, block_len, d_head] or [num_blocks, 1, block_len, d_head]. kv_heads > 1: [num_blocks, kv_heads, block_len, d_head]. Block KV with fp8_packed: kv_heads == 1: [num_blocks, block_len // 2, d_head, 2] or [num_blocks, 1, block_len // 2, d_head, 2]. kv_heads > 1: [num_blocks, kv_heads, block_len // 2, d_head, 2].

  • V_cache (nl.NkiTensor) – Value cache @ HBM. Flat KV: [B, kv_heads, S_max_ctx, d_head]. Block KV: kv_heads == 1: [num_blocks, block_len, d_head] or [num_blocks, 1, block_len, d_head]. kv_heads > 1: [num_blocks, kv_heads, block_len, d_head].

  • attention_mask (Optional[nl.NkiTensor]) – Attention mask @ HBM. Required for ordinary attention. Sequence-packed attention uses seq_packed_attention_mask instead, so this may be None. When pos_ids is provided: [S_tkg, B, q_heads_attn, S_tkg], active-only portion of the mask (the prior mask is generated in-kernel; layout is the same for swap and default). When pos_ids is None (full pre-generated mask): the layout MUST match the QK-swap decision. Determine it with is_qk_swapped(…) first (see attention_tkg / attention_tkg_utils for the full mask-shape guidance), then supply the matching layout: - default (K-stationary): [S_ctx, B, q_heads_attn, S_tkg] - QK-swap (Q-stationary): [B, q_heads_attn, S_tkg, S_ctx] gen_mask_tkg_hbm produces the correct layout when given the same transposed_out flag.

  • sink (Optional[nl.NkiTensor]) – [q_heads_attn, 1] @ HBM, one value per attention (Q) head. q_heads_attn = q_heads * KVDP (= kv_heads * q_per_group).

  • update_cache (bool) – Update KV cache with new tokens

  • kv_cache_update_idx (Optional[nl.NkiTensor]) – Cache write positions (uint32_max = skip). Flat KV: [B, 1] — the shared start position; consecutive tokens are assumed. Block KV: per-token physical slots. kv_heads == 1: [B, S_tkg]. kv_heads > 1: [B, kv_heads, S_tkg].

  • W_out (Optional[nl.NkiTensor]) – [q_heads*d_head, H] @ HBM, Output projection weights. When using FP8 KV cache (k_scale/v_scale provided), the attention output has magnitude scaled by v_scale. The caller must absorb v_scale into W_out or its dequant scale to compensate: W_out = W_out_original / v_scale (for NONE quantization) or weight_dequant_scale_out = scale / v_scale (for ROW/STATIC).

  • bias_out (Optional[nl.NkiTensor]) – [1, H] @ HBM, Output projection bias

  • quantization_type_out (QuantizationType) – Type of quantization for output projection (NONE, STATIC, ROW).

  • weight_dequant_scale_out (Optional[nl.NkiTensor]) – Weight dequantization scale for output projection.

  • input_dequant_scale_out (Optional[nl.NkiTensor]) – Input dequantization scale for output projection.

  • transposed_out (bool) – Transpose output layout (requires W_out)

  • out_in_sb (bool) – Return output in SBUF instead of HBM

  • transposed_in (bool) – When True, X is in transposed HBM layout [H0=pmax, n_prgs, H1_shard, BxS] from a preceding layer’s transposed_out. (Note that using [H0, n_prgs, BxS, H1_shard] instead would require an extra tensor_copy upstream so all-in-all it doesn’t save us from doing the equivalent in attention.) Requires B*S_tkg <= pmax (small batch path). Default: False.

  • softmax_scale (Optional[float]) – Scaling factor for attention scores. If None, defaults to (1/√D) / k_scale. When using FP8 KV cache (k_scale/v_scale provided) and softmax_scale is None, the kernel automatically divides by k_scale to dequantize the KV cache values in the QK matmul, effectively setting softmax_scale = (1/√D) / k_scale. When softmax_scale is explicitly provided, the caller is responsible for incorporating k_scale (e.g., softmax_scale = scaling / k_scale).

  • enable_fa_s_prior_tiling (bool) – bool: Whether to enable flash attention (FA) for attention computation. When enabled, the attention computation is tiled along the context (s_prior) to reduce peak memory usage.

  • fp8_packed (bool) – When True with block KV, K_cache is fp8 [num_blocks, block_len // 2, d_head, 2] where the last dimension packs two consecutive sequence positions. Enables DMA transpose for FP8 block KV, avoiding the slower PE transpose fallback.

  • k_scale (Optional[nl.NkiTensor]) – Scale for K quantization to FP8. Shape (PMAX, 1) or (1, 1) @ HBM. Must contain a single scalar value (replicated or scalar). When provided with v_scale, enables FP8 KV cache quantization. Supported dtypes: float32, float16, bfloat16.

  • v_scale (Optional[nl.NkiTensor]) – Scale for V quantization to FP8. Shape (PMAX, 1) or (1, 1) @ HBM. Must contain a single scalar value (replicated or scalar). When provided with k_scale, enables FP8 KV cache quantization. Supported dtypes: float32, float16, bfloat16.

  • sbm (Optional[SbufManager]) – SBUF memory manager (otherwise auto-allocated)

  • skip_attention (bool) – Skip attention computation (for testing)

  • is_h_transposed_by_4 (bool) – Whether input X and RMSNorm gamma have been pre-shuffled along the H dimension for MXFP quantization. Required to be True when quantization_type_qkv is MX, STATIC_MX, or ROW_MX. The pre-shuffle reorders [B, S, H//512, 128, 4] → [B, S, 4, H//512, 128] so the data matches the layout expected by the hardware quantize_mx instruction. Default: False.

  • KVDP (int) – KV cache data parallelism degree - number of ranks that shard the KV cache across the batch dimension (1 = disabled). Each rank processes B/KVDP batches.

  • KVDP_replica_group (Optional[ReplicaGroup]) – Replica group for collective ops

  • KVDP_rank (Optional[nl.NkiTensor]) – [1] @ HBM, uint32. This rank’s position within its KVDP replica group (0 to KVDP-1). Used for batch/head slicing. Required when KVDP > 1.

  • pos_ids (Optional[nl.NkiTensor]) – [B, S_tkg] @ HBM, Absolute sequence position of each active token. When provided, the kernel generates the prior causal mask on-chip (attending to KV positions < pos_id) instead of loading a full pre-generated mask from HBM. In this mode, attention_mask carries only the active portion.

  • swa_start_pos_ids (Optional[nl.NkiTensor]) – [B, S_tkg] @ HBM, Per-query SWA window start positions. When provided together with pos_ids, generates a banded sliding window attention mask where each token attends to positions in [start_pos, pos_id). When None, standard causal attention mask is generated (token attends to all positions < pos_id).

  • S_ctx (Optional[int]) – Explicit context length for flat KV with pos_ids. Required when pos_ids is provided and block KV is not used, since S_ctx cannot be derived from the attention mask. Must be None otherwise.

  • max_context_len (Optional[nl.NkiTensor]) – [1] @ HBM, int32 scalar for dynamic FA early exit. When provided, the FA loop exits early after processing ceil(max_context_len / tile_size) tiles. Requires block KV cache and pos_ids.

  • dtype_mode (DtypeMode) – Quantization dtype policy forwarded to all FP8-aware subkernels and to the FP8 KV cache when K_cache / V_cache are allocated with the opaque "float8e4" sentinel. Concrete dtypes are honored as-is. - DtypeMode.NON_OCP (default): nl.float8_e4m3 (max=240). - DtypeMode.OCP: nl.float8_e4m3fn (max=448). TRN3 only. - DtypeMode.AUTO: nl.float8_e4m3fn on TRN3, else nl.float8_e4m3.

  • CP (int) – Context parallelism degree - number of ranks that shard the KV cache along the sequence (context) dimension (1 = disabled). Each rank holds an S_ctx/CP slice of the prior KV cache and runs local flash attention; the distributed softmax results are combined across ranks by the CP collectives.

  • CP_replica_group (Optional[ReplicaGroup]) – Replica group for CP collective ops. Required when CP > 1.

  • CP_collective_mode (Optional[CPCollectiveMode]) – Collective mode for CP output redistribution (ALL_TO_ALL or REDUCE_SCATTER). Defaults to ALL_TO_ALL when None. KV Data Parallelism (KVDP > 1): KV-DP partitions the KV cache across ranks along the batch dimension. Each rank holds B/KVDP batches of the KV cache. Before attention: all_gather Q heads, slice Q/K/V batch. After attention: all_gather output batch, slice heads. When KV data parallelism is enabled, input/output shapes change: - B_attn = B / KVDP (batches per rank for attention) - q_heads_attn = q_heads * KVDP (query heads per rank after gather) Input shape changes: - K_cache, V_cache: [B_attn, …] instead of [B, …] - attention_mask: [S_ctx, B_attn, q_heads_attn, S_tkg] - pos_ids: [B_attn, S_tkg] instead of [B, S_tkg] - swa_start_pos_ids: [B_attn, S_tkg] instead of [B, S_tkg] - kv_cache_update_idx: [B_attn, 1] (caller must slice per rank) Output shape changes (when update_cache=False): - K_out: kv_heads == 1: [d_head, B_attn, S_tkg] kv_heads > 1: [d_head, B_attn, kv_heads, S_tkg] - V_out: [B_attn, kv_heads, S_tkg, d_head] Context Parallelism (CP > 1): CP partitions the prior KV cache across ranks along the sequence (context) dimension. Each rank holds an S_ctx/CP slice and runs local flash attention, returning unnormalized output plus per-rank softmax stats (running max/sum). Before attention: all_gather Q heads (q_heads -> q_heads*CP, no K/V slicing since each rank owns its context shard). After

  • seq_packed_attention_mask (Optional[nl.NkiTensor]) – Packed schedule-row mask. When provided, the active-block, Q-index, and sequence-ID tables are required and the kernel dispatches to attention_tkg_seq_packed. The two accumulator route tables are an optional all-or-none pair. This is the only attention mask consumed on the packed path. See attention_tkg_sequence_packing_design_spec.md in core/attention for the schedule contract.

  • seq_packed_active_blocks_table (Optional[nl.NkiTensor]) – Packed physical block routes.

  • seq_packed_q_index_table (Optional[nl.NkiTensor]) – Packed Q tensor-indirection routes.

  • seq_packed_seq_id_table (Optional[nl.NkiTensor]) – Packed absolute sequence IDs.

  • seq_packed_accumulator_partial_route_table (Optional[nl.NkiTensor]) – Tensor Indirection routes that scatter each row’s slot partials into sequence groups. Must be supplied with the group-state route.

  • seq_packed_accumulator_group_state_route_table (Optional[nl.NkiTensor]) – Tensor Indirection routes that gather and scatter each group’s persistent sequence state. Must be supplied with the partial route.

Returns:

Output tensor with shape depending on projection and output location: - Without projection (W_out=None): - out_in_sb=False: [B, q_heads, d_head, S_tkg] @ HBM - out_in_sb=True: [d_head, B*q_heads*S_tkg] @ SBUF - With projection (W_out provided): - transposed_out=False, out_in_sb=False: [B*S_tkg, H] @ HBM - transposed_out=False, out_in_sb=True: [B*S_tkg, H//lnc] @ SBUF - transposed_out=True, out_in_sb=False: [128, lnc, H//lnc//128, B*S_tkg] @ HBM - transposed_out=True, out_in_sb=True: [128, H//lnc//128, B*S_tkg] @ SBUF

Returns:

  • If update_cache=True: Updated K cache (shape matches K_cache input) - If update_cache=False: New K tokens @ HBM with shape: kv_heads == 1: [d_head, B_attn, S_tkg] kv_heads > 1: [d_head, B_attn, kv_heads, S_tkg]

Returns:

  • If update_cache=True: Updated V cache (shape matches V_cache input) - If update_cache=False: New V tokens [B_attn, kv_heads, S_tkg, d_head] @ HBM

Return type:

nl.NkiTensor

Dimensions:

  • B: batch size

  • S_tkg: number of new tokens to generate

  • S_ctx: KV cache sequence length in current bucket

  • S_max_ctx: maximum KV cache capacity of current bucket

  • H: hidden dimension

  • d_head: head dimension (must be even)

  • q_heads: number of query heads

  • kv_heads: 1 (GQA with single KV head)

Supported Data Types:

  • Supports nl.float16 and nl.bfloat16

Constraints:

  • Requires NeuronCore v3+

  • d_head must be even

  • H must be multiple of 128

  • Requires batch * sequence_tkg * q_heads <= pmax (=128)

Implementation Details#

Computation Flow:

The kernel executes the following stages in sequence:

  1. Input Pre-normalization (optional):

    • Apply RMSNorm to input hidden states: X_norm = RMSNorm(X, rmsnorm_pre_W, rmsnorm_pre_eps)

    • Computed in FP32, result cast back to input dtype

  2. QKV Projection:

    • Compute QKV = X_norm @ W_qkv.T using matrix multiplication

    • Result shape: [B, S_tkg, (q_heads + 2) * d_head]

    • Supports FP8 quantization with dequantization scales

  3. Q/K Processing (per head group):

    • Extract Q heads: Q = QKV[:, :, :q_heads * d_head]

    • Extract K head: K = QKV[:, :, q_heads * d_head : (q_heads + 1) * d_head]

    • Apply RoPE if enabled: Q, K = RoPE(Q, K, cos, sin, position_ids)

    • Apply per-head RMSNorm if enabled: Q = RMSNorm(Q, rmsnorm_post_W_Q), K = RMSNorm(K, rmsnorm_post_W_K)

  4. V Processing:

    • Extract V head: V = QKV[:, :, (q_heads + 1) * d_head :]

  5. KV Cache Update:

    • Write new K/V tokens to cache at positions specified by kv_cache_update_idx

    • Supports multiple cache layouts (flat, transposed, block-based)

    • Uses indirect addressing for efficient batch processing

  6. Attention Computation:

    • Compute scaled dot-product attention: Attn = softmax(Q @ K_cache.T / scale) @ V_cache

    • Apply causal masking based on S_ctx (context lengths)

    • Use FP32 accumulation if mixed_precision=True

    • Supports Grouped-Query Attention by replicating KV heads

  7. Output Projection:

    • Reshape attention output: Attn_flat = Attn.reshape([B, S_tkg, q_heads * d_head])

    • Compute out = Attn_flat @ W_o.T

    • Supports FP8 quantization with dequantization scales

Memory Management:

The kernel uses a custom SBUF memory manager (SbufManager) to efficiently allocate and reuse on-chip memory:

  • Stack-based allocation for temporary tensors

  • Automatic memory reuse after tensor lifetime ends

  • Minimizes SBUF fragmentation

Parallelization:

The kernel supports data parallelism across multiple Neuron Cores:

  • Batch dimension (B) can be sharded across cores

  • Each core processes a subset of batch elements independently

  • KV cache updates use per-core indexing

Cache Layout Support:

  1. Flat Cache (is_block_kv=False):

    • K cache: [B, S_max_ctx, d_head] or [B, d_head, S_max_ctx] (transposed)

    • V cache: [B, S_max_ctx, d_head]

    • Direct indexing by batch and sequence position

  2. Block Cache (is_block_kv=True):

    • K/V cache: [num_blocks, block_len, d_head]

    • Indirect indexing via block slot mapping

    • Efficient for variable-length sequences

Quantization Support:

  • FP8 weights: Provide qkv_scale and o_scale for dequantization

  • Mixed precision: FP32 accumulation with FP16/BF16 inputs

  • Automatic dtype handling throughout the pipeline

Key Implementation Notes:

  1. Grouped-Query Attention: The kernel processes Q heads in groups, where each group shares a single K/V head. This reduces KV cache memory by a factor of q_heads / kv_heads.

  2. RoPE Application: Rotary embeddings are applied using position indices derived from S_ctx (current context length). Supports both contiguous and interleaved layouts.

  3. Causal Masking: Attention scores are masked such that token at position i can only attend to positions 0 to i in the context. Implemented by adding -inf to masked positions before softmax.

  4. Cache Update Optimization:

    • For S_tkg=1: Uses batched vector DMA with vector_offset for all batches in one operation

    • For S_tkg>1: Uses per-batch scalar DMA with scalar_offset

    • Block cache uses indirect addressing via block slot indices

  5. Memory Efficiency: All intermediate tensors (QKV, Q, K, V, attention scores, attention output) remain in SBUF. Only input X, weights, caches, and final output out reside in HBM.

This document is relevant for: Trn2, Trn3