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.ndarray, X_hidden_dim_actual: Optional[int], rmsnorm_X_enabled: bool, rmsnorm_X_eps: Optional[float], rmsnorm_X_gamma: Optional[nl.ndarray], W_qkv: nl.ndarray, bias_qkv: Optional[nl.ndarray], quantization_type_qkv: QuantizationType, weight_dequant_scale_qkv: Optional[nl.ndarray], input_dequant_scale_qkv: Optional[nl.ndarray], rmsnorm_QK_pre_rope_enabled: bool, rmsnorm_QK_pre_rope_eps: float, cos: Optional[nl.ndarray], sin: Optional[nl.ndarray], rope_contiguous_layout: bool, rmsnorm_QK_post_rope_enabled: bool, rmsnorm_QK_post_rope_eps: float, rmsnorm_QK_post_rope_W_Q: Optional[nl.ndarray], rmsnorm_QK_post_rope_W_K: Optional[nl.ndarray], K_cache_transposed: bool, active_blocks_table: Optional[nl.ndarray], K_cache: nl.ndarray, V_cache: nl.ndarray, attention_mask: nl.ndarray, sink: Optional[nl.ndarray], softmax_scale: Optional[float] = None, update_cache: bool, kv_cache_update_idx: Optional[nl.ndarray], k_scale: Optional[nl.ndarray] = None, v_scale: Optional[nl.ndarray] = None, W_out: Optional[nl.ndarray], bias_out: Optional[nl.ndarray], quantization_type_out: QuantizationType, weight_dequant_scale_out: Optional[nl.ndarray], input_dequant_scale_out: Optional[nl.ndarray], transposed_out: bool, out_in_sb: bool, sbm: Optional[SbufManager] = None, skip_attention: bool = False)#

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.ndarray) – Input hidden states [B, S_tkg, H] @ HBM or [pmax, B*S_tkg, H//pmax] @ SBUF

  • X_hidden_dim_actual (int, optional) – Actual hidden dim if X is padded

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

  • rmsnorm_X_eps (float, optional) – RMSNorm epsilon (default 1e-3)

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

  • W_qkv (nl.ndarray) – QKV projection weights [H, d_head*(q_heads+2)] @ HBM

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

  • quantization_type_qkv (QuantizationType) – Quantization type for QKV projection

  • weight_dequant_scale_qkv (nl.ndarray, optional) – Weight dequantization scale for QKV projection

  • input_dequant_scale_qkv (nl.ndarray, optional) – 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

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

  • sin (nl.ndarray, optional) – RoPE sine embeddings [d_head//2, B, S_tkg] @ HBM (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 (nl.ndarray, optional) – Post-RoPE Q weights [1, d_head] @ HBM

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

  • K_cache_transposed (bool) – K cache layout flag

  • active_blocks_table (nl.ndarray, optional) – Block indices for block KV cache [B, num_blocks] @ HBM

  • K_cache (nl.ndarray) – Key cache @ HBM

  • V_cache (nl.ndarray) – Value cache @ HBM

  • attention_mask (nl.ndarray) – Attention mask [S_ctx, B, q_heads, S_tkg] @ HBM

  • sink (nl.ndarray, optional) – Attention sink tokens [H, 1] @ HBM

  • softmax_scale (float, optional) – Scaling factor for attention scores (Q @ K^T * softmax_scale). If None, defaults to 1.0 / sqrt(d_head).

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

  • kv_cache_update_idx (nl.ndarray, optional) – Cache write positions [B, 1] (uint32_max = skip)

  • k_scale (nl.ndarray, optional) – Key quantization scale for FP8 KV cache. Enables FP8 quantization of K values written to cache.

  • v_scale (nl.ndarray, optional) – Value quantization scale for FP8 KV cache. Enables FP8 quantization of V values written to cache.

  • W_out (nl.ndarray, optional) – Output projection weights [q_heads*d_head, H] @ HBM

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

  • quantization_type_out (QuantizationType) – Quantization type for output projection

  • weight_dequant_scale_out (nl.ndarray, optional) – Weight dequantization scale for output projection

  • input_dequant_scale_out (nl.ndarray, optional) – 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

  • sbm (SbufManager, optional) – SBUF memory manager (otherwise auto-allocated)

  • skip_attention (bool) – Skip attention computation (for testing). Default: False.

Returns:

Tuple of (out, K_out, V_out) - Output tensor, updated K cache or new K tokens, updated V cache or new V tokens

Return type:

tuple

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.