v1.9: padding masks on the kernel path for NeuronXLMRobertaModel (per-image additive key bias in attention_block_l1_ln_post_kb); existing models bitwise identical to v1.8

Browse files

Files changed (5) hide show

CHANGELOG.md +26 -0
README.md +10 -8
build/torch-neuron/layers.py +23 -8
build/torch-neuron/nki_kernels/attention_block_l1_ln.py +33 -1
build/torch-neuron/nki_kernels/attention_inner.py +32 -1

CHANGELOG.md CHANGED Viewed

@@ -1,5 +1,31 @@
 # Changelog
 ## v1.8 (2026-10)
 * **New: post-LayerNorm encoders (BERT / RoBERTa / XLM-R).** `NeuronXLMRobertaModel` is a drop-in forward for

 # Changelog
+## v1.9 (2026-10)
+* **New: padding masks on the kernel path for `NeuronXLMRobertaModel`.** v1.8 sent any batch with a padding
+  `attention_mask` to the stock forward. The attention kernel now takes an optional per-image additive key bias
+  (`attention_block_l1_ln_post_kb`: 0 for attended keys, -1e9 for masked keys), and the layer turns the mask into that
+  bias, so variable-length batches stay on the kernels. `VIT_ENC_MASK_KERNEL=0` restores the v1.8 fallback.
+  XLM-RoBERTa-Large, LNC=1, 8 processes on one trn2 device, a log-normal length mix padded to the sequence length
+  (mean valid length 53 / 89 / 191), inferences/s:
+  | seq | stock with mask (v1.8 fallback) | v1.9 kernels with mask | speedup | v1.9 kernels, no mask |
+  |---|---|---|---|---|
+  | 128 | 956.2 | **2,923.0** | 3.06x | 3,177.0 |
+  | 256 | 464.5 | **1,899.2** | 4.09x | 1,922.7 |
+  | 512 | 213.0 | **894.4** | 4.20x | 962.9 |
+  Accuracy vs the fp32 CPU model on the same mask (valid positions): rel_l2 0.005-0.055 for right padding, half
+  padding, length-1 sequences and interior holes (the unmasked v1.8 path gives 0.007-0.051 at the same shapes).
+  Masked positions are computed like the stock model does (their keys still attend to valid keys). Every query row
+  is still computed, so padding is not skipped: throughput follows the padded length.
+* **No change for existing models:** 28/28 regression cases (DINOv3 S/S+/B/L/H+ at 224/448, V-JEPA2, LNC=1 and LNC=2)
+  give output bitwise identical to v1.8 at the same speed (within ±0.15%). Without the new argument the kernels emit
+  the v1.8 instruction stream. Verified on the public SDK 2.33 container
+  (`public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314-sdk2.33.0-ubuntu26.04`: torch 2.14.0, torch-neuronx 2.14.3,
+  neuronx-cc 2.28.7631, NKI 0.7.0) with transformers 4.57.6.
 ## v1.8 (2026-10)
 * **New: post-LayerNorm encoders (BERT / RoBERTa / XLM-R).** `NeuronXLMRobertaModel` is a drop-in forward for

README.md CHANGED Viewed

@@ -144,7 +144,7 @@ activation scale with `calibrate_fp8(runner, gelu_amax)`, because trn2's e4m3 sa
 Which LNC to use: **LNC=1** (8 logical cores, one process per core) gives the most throughput per
 instance. **LNC=2** (4 larger cores) gives the lowest latency per request.
-### XLM-RoBERTa / BERT-family text encoders (v1.8)
 ```python
 import os
@@ -155,13 +155,15 @@ model = XLMRobertaModel.from_pretrained("FacebookAI/xlm-roberta-large", dtype=to
                                         add_pooling_layer=False).eval().to("neuron")
 model.forward = types.MethodType(k.NeuronXLMRobertaModel.forward, model)   # instance graft (k = get_kernel(...) as above)
 model = torch.compile(model, backend="neuron")
-ids = tokenizer(texts, padding="max_length", max_length=128, return_tensors="pt").input_ids   # seq % 128 == 0
-out = model(input_ids=ids.to("neuron")).last_hidden_state
 ```
-The kernel path needs LNC=1, a sequence length that is a multiple of 128 (up to 2048) and no padding mask. Pass
-`attention_mask=None` or all ones; masked batches fall back to the stock forward. Per instance, XLM-RoBERTa-Large at
-seq 128 / 256 / 512 runs at 3,145 / 1,870 / 935 inferences/s, against 2,093 / 944 / 397 stock (see CHANGELOG v1.8).
 ## Kernels
@@ -169,7 +171,7 @@ seq 128 / 256 / 512 runs at 3,145 / 1,870 / 935 inferences/s, against 2,093 / 94
 |---|---|
 | `attention_block` | LN'd hidden -> Q/K/V projection + RoPE + attention + output projection, one 8-head (D=64) or 4-head (D=128) group per program; any number of images per launch; padded keys skipped |
 | `attention_block_l1` | LNC=1: one program over **all** head groups, one output (no per-group partials) |
-| `attention_block_l1_ln` | `attention_block_l1` plus residual add and LayerNorm in the epilogue |
 | `res_layernorm_fold` | residual add (sum of up to 4 branches) + affine-free LayerNorm |
 | `mlp` | GELU MLP (fc1 -> GELU -> fc2) |
 | `mlp_res_ln` | `mlp` plus residual add and the next block's LayerNorm in the epilogue |
@@ -196,7 +198,7 @@ RoPE pairings via `rope=`:
 * ViT-7B-sized models (hidden 4096) are not covered by the drop-in layer. They need tensor parallelism.
   The streaming SwiGLU kernel is slower than the compiler's FFN at that width.
 * Validated end to end on DINOv3 ViT-B/L/H+, V-JEPA2 ViT-L and (v1.8) XLM-RoBERTa-Large.
-* Post-LN text encoders (v1.8, `NeuronXLMRobertaModel`): LNC=1 only, no padding mask, no token-type / relative
   position variants beyond the standard BERT embeddings, GELU FFN.
 ## Requirements

 Which LNC to use: **LNC=1** (8 logical cores, one process per core) gives the most throughput per
 instance. **LNC=2** (4 larger cores) gives the lowest latency per request.
+### XLM-RoBERTa / BERT-family text encoders (v1.8; padding masks v1.9)
 ```python
 import os
                                         add_pooling_layer=False).eval().to("neuron")
 model.forward = types.MethodType(k.NeuronXLMRobertaModel.forward, model)   # instance graft (k = get_kernel(...) as above)
 model = torch.compile(model, backend="neuron")
+enc = tokenizer(texts, padding="max_length", max_length=128, return_tensors="pt")              # seq % 128 == 0
+out = model(input_ids=enc.input_ids.to("neuron"), attention_mask=enc.attention_mask.to("neuron")).last_hidden_state
 ```
+The kernel path needs LNC=1 and a sequence length that is a multiple of 128 (up to 2048). Since v1.9 a padding
+`attention_mask` stays on the kernel path (an additive key bias in the attention kernel); `attention_mask=None` skips
+the mask entirely and is slightly faster on fully padded batches. Per instance, XLM-RoBERTa-Large at seq 128 / 256 / 512
+runs at 3,145 / 1,870 / 935 inferences/s without a mask (2,093 / 944 / 397 stock), and at 2,923 / 1,899 / 894 with a
+real padding mask (956 / 465 / 213 stock with the same mask). See CHANGELOG v1.8 and v1.9.
 ## Kernels
 |---|---|
 | `attention_block` | LN'd hidden -> Q/K/V projection + RoPE + attention + output projection, one 8-head (D=64) or 4-head (D=128) group per program; any number of images per launch; padded keys skipped |
 | `attention_block_l1` | LNC=1: one program over **all** head groups, one output (no per-group partials) |
+| `attention_block_l1_ln` | `attention_block_l1` plus residual add and LayerNorm in the epilogue (`_post`: post-LN output; `_post_kb`, v1.9: per-image additive key bias for padding masks) |
 | `res_layernorm_fold` | residual add (sum of up to 4 branches) + affine-free LayerNorm |
 | `mlp` | GELU MLP (fc1 -> GELU -> fc2) |
 | `mlp_res_ln` | `mlp` plus residual add and the next block's LayerNorm in the epilogue |
 * ViT-7B-sized models (hidden 4096) are not covered by the drop-in layer. They need tensor parallelism.
   The streaming SwiGLU kernel is slower than the compiler's FFN at that width.
 * Validated end to end on DINOv3 ViT-B/L/H+, V-JEPA2 ViT-L and (v1.8) XLM-RoBERTa-Large.
+* Post-LN text encoders (v1.8, `NeuronXLMRobertaModel`): LNC=1 only (padding masks supported since v1.9), no token-type / relative
   position variants beyond the standard BERT embeddings, GELU FFN.
 ## Requirements

build/torch-neuron/layers.py CHANGED Viewed

@@ -20,6 +20,7 @@ Usage (graft, works on any Transformers version)::
 """
 from __future__ import annotations
 import torch
 from .encoder import LNC, BlockWeights, EncoderRunner, pack_block, rope_tables
@@ -289,9 +290,10 @@ class NeuronXLMRobertaModel(torch.nn.Module):
     Graft on an instance (``model.forward = types.MethodType(NeuronXLMRobertaModel.forward, model)``); do not use it
     as a class-level KernelConfig replacement, because the fallback calls ``type(model).forward``.
-    Returns ``last_hidden_state``. Kernel path: LNC=1, sequence length a multiple of 128 and <= 2048, all positions
-    attended (``attention_mask`` all ones / no padding). Anything else (padding mask, LNC=2, other lengths,
-    off-device, training) falls back to the stock forward.
     """
     can_torch_compile = True
@@ -307,7 +309,12 @@ class NeuronXLMRobertaModel(torch.nn.Module):
         S = input_ids.shape[1] if input_ids is not None else -1
         # The mask check runs outside the compiled graph (one host sync per call); pass attention_mask=None to skip it.
         masked = attention_mask is not None and _bert_masked(attention_mask)
-        if (input_ids is None or input_ids.device.type != "neuron" or self.training or masked or position_ids is not None
                 or kwargs.get("inputs_embeds") is not None or not _bert_ok(self, S)):
             # stock forward of the grafted instance's own class (XLMRobertaModel / RobertaModel / BertModel). The graft
             # sets an INSTANCE attribute (types.MethodType), so the class attribute is still the stock forward.
@@ -333,12 +340,20 @@ class NeuronXLMRobertaModel(torch.nn.Module):
         h = F.layer_norm(x, (H,), g0, b0, eps0)
         nL = len(R.L)
         from .encoder import _PACK2
-        pk2 = int(_PACK2 and B % 2 == 0 and 2 * S <= 512)
         for li, p in enumerate(R.L):
             ga, ba, _ = st["ln_attn"][li]
-            _, z2, h1 = _ab1ln.attention_block_l1_ln_post(
-                z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo_full"], h, p["scale"], D, S, p["eps2"],
-                ga, ba, B, R._res(_ab1ln, S), pk2, 1)
             go, bo, eo = st["ln_in"][li + 1]
             if li + 1 < nL:
                 _, z1, h = nki_mlp_res_ln_post(z2, p["w1t"], p["b1"], p["w2t"], p["b2"], h1, eo, go, bo)

 """
 from __future__ import annotations
+import os as _osm
 import torch
 from .encoder import LNC, BlockWeights, EncoderRunner, pack_block, rope_tables
     Graft on an instance (``model.forward = types.MethodType(NeuronXLMRobertaModel.forward, model)``); do not use it
     as a class-level KernelConfig replacement, because the fallback calls ``type(model).forward``.
+    Returns ``last_hidden_state``. Kernel path: LNC=1, sequence length a multiple of 128 and <= 2048. v1.9: a
+    padding ``attention_mask`` (any 0/1 pattern) stays on the kernel path as an additive key bias
+    (``VIT_ENC_MASK_KERNEL=0`` restores the v1.8 fallback). Anything else (LNC=2, other lengths, off-device,
+    training) falls back to the stock forward.
     """
     can_torch_compile = True
         S = input_ids.shape[1] if input_ids is not None else -1
         # The mask check runs outside the compiled graph (one host sync per call); pass attention_mask=None to skip it.
         masked = attention_mask is not None and _bert_masked(attention_mask)
+        if masked and _osm.environ.get("VIT_ENC_MASK_KERNEL", "1") == "0":
+            kmask = None
+        else:
+            kmask = attention_mask if masked else None
+        if (input_ids is None or input_ids.device.type != "neuron" or self.training or (masked and kmask is None)
+                or position_ids is not None
                 or kwargs.get("inputs_embeds") is not None or not _bert_ok(self, S)):
             # stock forward of the grafted instance's own class (XLMRobertaModel / RobertaModel / BertModel). The graft
             # sets an INSTANCE attribute (types.MethodType), so the class attribute is still the stock forward.
         h = F.layer_norm(x, (H,), g0, b0, eps0)
         nL = len(R.L)
         from .encoder import _PACK2
+        pk2 = int(_PACK2 and B % 2 == 0 and 2 * S <= 512 and kmask is None)
+        # v1.9: padding mask -> additive key bias [B, S] (0 attended, -1e9 masked), same softmax as the stock
+        # extended mask; every query row (padding rows included) is computed, as in the stock forward.
+        kb = None if kmask is None else ((1.0 - kmask.to(torch.float32)) * -1e9).contiguous()
         for li, p in enumerate(R.L):
             ga, ba, _ = st["ln_attn"][li]
+            if kb is None:
+                _, z2, h1 = _ab1ln.attention_block_l1_ln_post(
+                    z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo_full"], h, p["scale"], D, S,
+                    p["eps2"], ga, ba, B, R._res(_ab1ln, S), pk2, 1)
+            else:
+                _, z2, h1 = _ab1ln.attention_block_l1_ln_post_kb(
+                    z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo_full"], h, p["scale"], D, S,
+                    p["eps2"], ga, ba, kb, B, R._res(_ab1ln, S), 1)
             go, bo, eo = st["ln_in"][li + 1]
             if li + 1 < nL:
                 _, z1, h = nki_mlp_res_ln_post(z2, p["w1t"], p["b1"], p["w2t"], p["b2"], h1, eo, go, bo)

build/torch-neuron/nki_kernels/attention_block_l1_ln.py CHANGED Viewed

@@ -88,7 +88,7 @@ def sbuf_estimate(H: int, Fp: int, N: int, D: int, resident: bool) -> int:
 @nki.jit
 def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
                                   wot_hbm, bo_hbm, x_hbm, scale, head_dim, n_valid, eps, nimg=1, resident=0, pack2=0,
-                                  g_hbm=None, bt_hbm=None, norope=0):
     NT, H = h_hbm.shape
     N = NT // nimg
     Fp = wqk_hbm.shape[2]
@@ -118,6 +118,11 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
     _kernel_assert(NT == nimg * N, "rows must be nimg x N")
     _kernel_assert(x_hbm.shape == (NT, H), "x must be [nimg*N, H] fp32")
     post = g_hbm is not None                       # post-LN mode: also emit h = LN(x_new) * g + bt (fp32)
     if post:
         hout_hbm = nl.ndarray((NT, H), dtype=nl.float32, buffer=nl.shared_hbm)
@@ -234,6 +239,12 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
                     nisa.dma_copy(dst=tab[t_i][s][:, 0:_tw],
                                   src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
         img_base = allocator.get_current_address()
         for g in range(G):
@@ -399,6 +410,7 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
                         out_o_hbm=o_l,
                         init_sbuf_addr=attn_base,
                         seqlen_k_valid=n_valid,
                     )
                     allocator.set_current_address(attn_base)
@@ -489,6 +501,7 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
 _ATTN_L1_LN = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
 _ATTN_L1_LN_POST = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
 _ATTN_L1_LN_POST_NR = _ATTN_L1_LN_POST
 @nki_op("vit_enc::attention_block_l1_ln", mutates_args={})
@@ -536,3 +549,22 @@ def attention_block_l1_ln_post(
     if bqk.dim() == 3 and bqk.shape[1] != 128:
         bqk = _pack_qk_bias(bqk)
     return _ATTN_L1_LN_POST(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident, pack2, g, bt, norope)

 @nki.jit
 def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
                                   wot_hbm, bo_hbm, x_hbm, scale, head_dim, n_valid, eps, nimg=1, resident=0, pack2=0,
+                                  g_hbm=None, bt_hbm=None, norope=0, kbias_hbm=None):
     NT, H = h_hbm.shape
     N = NT // nimg
     Fp = wqk_hbm.shape[2]
     _kernel_assert(NT == nimg * N, "rows must be nimg x N")
     _kernel_assert(x_hbm.shape == (NT, H), "x must be [nimg*N, H] fp32")
+    # v1.9: optional per-image additive key bias [nimg, N] fp32 (0 attended, large negative for padding keys).
+    kb = kbias_hbm is not None
+    if kb:
+        _kernel_assert(kbias_hbm.shape == (nimg, N), "kbias must be [nimg, N] fp32")
+        _kernel_assert(not pack2, "kbias is not supported with pack2")
     post = g_hbm is not None                       # post-LN mode: also emit h = LN(x_new) * g + bt (fp32)
     if post:
         hout_hbm = nl.ndarray((NT, H), dtype=nl.float32, buffer=nl.shared_hbm)
                     nisa.dma_copy(dst=tab[t_i][s][:, 0:_tw],
                                   src=tab_hbm[t_i, 0:_P_MAX, s * _PSUM_FMAX:s * _PSUM_FMAX + _tw])
+        # v1.9: this image's key bias, broadcast to every partition (one stride-0 DMA)
+        kb_sb = None
+        if kb:
+            kb_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, N), dtype=nl.float32, align_to=4)
+            nisa.dma_copy(dst=kb_sb, src=kbias_hbm.ap(pattern=[[0, _P_MAX], [1, N]], offset=img * N))
         img_base = allocator.get_current_address()
         for g in range(G):
                         out_o_hbm=o_l,
                         init_sbuf_addr=attn_base,
                         seqlen_k_valid=n_valid,
+                        kbias_sbuf=kb_sb,
                     )
                     allocator.set_current_address(attn_base)
 _ATTN_L1_LN = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
 _ATTN_L1_LN_POST = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
 _ATTN_L1_LN_POST_NR = _ATTN_L1_LN_POST
+_ATTN_L1_LN_POST_KB = wrap_nki_to_torch_hop(_attention_block_l1_ln_kernel[1])
 @nki_op("vit_enc::attention_block_l1_ln", mutates_args={})
     if bqk.dim() == 3 and bqk.shape[1] != 128:
         bqk = _pack_qk_bias(bqk)
     return _ATTN_L1_LN_POST(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident, pack2, g, bt, norope)
+@nki_op("vit_enc::attention_block_l1_ln_post_kb", mutates_args={})
+def attention_block_l1_ln_post_kb(
+    h: torch.Tensor, wqk: torch.Tensor, bqk: torch.Tensor, tab: torch.Tensor, wvt: torch.Tensor, bv: torch.Tensor,
+    wot: torch.Tensor, bo: torch.Tensor, x: torch.Tensor, scale: float, head_dim: int, n_valid: int, eps: float,
+    g: torch.Tensor, bt: torch.Tensor, kbias: torch.Tensor, nimg: int = 1, resident: int = 0, norope: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """v1.9: ``attention_block_l1_ln_post`` with a per-image additive key bias ``kbias`` [nimg, N] fp32 (0 for
+    attended keys, a large negative value such as -1e9 for padding keys), the BERT-style padding mask. pack2 is off."""
+    for t in (h, wqk, bqk, tab, wvt, bv, wot, bo, x, g, bt, kbias):
+        if t is not None and t.requires_grad:
+            raise RuntimeError("attention_block_l1_ln_post_kb is inference-only; use torch.no_grad().")
+    if h.shape[0] // nimg > 2048:
+        raise NotImplementedError("attention_block_l1_ln_post_kb supports N<=2048 only")
+    if bqk.dim() == 3 and bqk.shape[1] != 128:
+        bqk = _pack_qk_bias(bqk)
+    return _ATTN_L1_LN_POST_KB(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident, 0,
+                               g, bt, norope, kbias)

build/torch-neuron/nki_kernels/attention_inner.py CHANGED Viewed

@@ -208,6 +208,7 @@ class AttnInternalBuffers(nl.NKIObject):
     range_sel_lbs_prior = None
     range_sel_ubs_prior = None
     k_offset_sb_u32 = None
 @dataclass
@@ -275,8 +276,12 @@ def _attention_cte(
     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.
@@ -406,6 +411,7 @@ def _attention_cte(
             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.
@@ -426,6 +432,7 @@ def _attention_cte_impl(
     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.
@@ -442,6 +449,7 @@ def _attention_cte_impl(
     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)
@@ -1196,7 +1204,30 @@ def _qk_and_max_large_tile_impl(qkmax_grp, large_tile_idx, ac: AttnConfig, atp:
             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],

     range_sel_lbs_prior = None
     range_sel_ubs_prior = None
     k_offset_sb_u32 = None
+    kbias_sb = None
 @dataclass
     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.
             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.
     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.
     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)
             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],