This document is relevant for: Trn2, Trn3

Attention Const-Max Kernel API Reference#

Constant-max attention: softmax(Q@K.T * softmax_scale - softmax_max) @ V.

Computes attention assuming a pre-known global softmax max.

Background#

The attention_const_max kernel computes softmax(Q @ K.T * softmax_scale - softmax_max) @ V using a caller-supplied global softmax maximum, so the running-max rescaling of online softmax is not needed.

API Reference#

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

attention_const_max#

nkilib.experimental.attention.attention_const_max(q_hbm: nl.NkiTensor, k_hbm: nl.NkiTensor, v_hbm: nl.NkiTensor, softmax_max: Union[float, nl.NkiTensor], softmax_scale: Optional[float] = None, q_layout: AttentionLayout = AttentionLayout.HEAD_MAJOR, k_layout: AttentionLayout = AttentionLayout.HEAD_MAJOR) → nl.NkiTensor#

Constant-max attention: softmax(Q@K.T * softmax_scale - softmax_max) @ V.

Parameters:

  • q_hbm (nl.NkiTensor) – [N, d, Sq] or [N, Sq, d] bf16 @ HBM, Query.

  • k_hbm (nl.NkiTensor) – [N, d, Sk] or [N, Sk, d] bf16 @ HBM, Key.

  • v_hbm (nl.NkiTensor) – [N, Sk, d] bf16 @ HBM, Value

  • softmax_max (Union[float, nl.NkiTensor]) – Upper bound on (scores * scale). Output is mathematically invariant to this value; it only affects numerical precision. A reasonable estimate suffices. May be either: - a Python scalar (compile-time constant, broadcast to all partitions), or - a [128, 1] f32 HBM tensor holding a single broadcast max replicated across the 128 partitions. Use the tensor form to supply a data-dependent max computed at runtime rather than baked in at trace time.

  • softmax_scale (Optional[float]) – Scaling factor applied to Q. Defaults to 1/sqrt(d).

  • q_layout (AttentionLayout) – Q layout. SEQ_MAJOR for [N, Sq, d], HEAD_MAJOR for [N, d, Sq]. Defaults to HEAD_MAJOR.

  • k_layout (AttentionLayout) – K layout. SEQ_MAJOR for [N, Sk, d], HEAD_MAJOR for [N, d, Sk]. Defaults to HEAD_MAJOR.

Returns:

[N, Sq, d] bf16 @ HBM, Normalized attention output

Return type:

nl.NkiTensor

Notes:

  • Requires std(Q) ≤ 1 and std(K) ≤ 1 for bf16 accuracy with softmax_max=4–5. This holds when QK-RMSNorm is applied after projection.

  • A softmax_max that is too low causes exp() overflow (inf/NaN in output). A softmax_max that is too high clips all scores toward zero, losing precision.

  • For inputs with significantly larger std use a standard online-softmax kernel. Choosing softmax_max from QK-RMSNorm weights: When QK-RMSNorm is applied, the max can be derived from the gamma weights: σ = sqrt(sum(γ_q² * γ_k²) / d) softmax_max ≈ 0.88 * σ * sqrt(2 * ln(Sk)) # Gumbel expected row-max As a safety check, compute the hard bound (Cauchy-Schwarz worst case): hard_bound = sqrt(d) * max|γ_q| * max|γ_k| If hard_bound - softmax_max > ~80, the score range is too wide for constant-max and the block should fall back to online-softmax attention.

Dimensions:

  • N: Number of attention heads (batch of heads)

  • d: Head dimension (must be pmax=128)

  • Sq: Query sequence length

This document is relevant for: Trn2, Trn3