# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Affero General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU Affero General Public License for more details. # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . """FlexAttention shared-prefix kernel for PrefixGrouper (GRPO shared-prompt dedup). In GRPO every prompt spawns ``G = num_generations`` completions that share the same prompt prefix. The full-row packed path forwards the identical prefix ``G`` times. PrefixGrouper stores the prefix ONCE and concatenates only the ``G`` suffixes, with an attention layout where each suffix token attends to ``[the single shared prefix] + [causal within its own suffix]``. This kernel expresses that one-prefix -> many-suffix fan-out via a ``torch.nn.attention.flex_attention`` block mask, so the masked-out cross-suffix / cross-group blocks are never computed and the ``P + G*R`` FLOP saving is realised (not merely a masked dense ``O(T^2)``). Mask semantics (identical to the certified SDPA oracle): keep(q_idx, kv_idx) = same_group(q, kv) AND ( is_prefix[kv_idx] # full prefix visibility OR ( suffix_of_kv[kv_idx] == suffix_of_kv[q_idx] # same suffix ... AND kv_idx <= q_idx ) ) # ... causal within it This module is self-contained (no dependency on any temp/ scratch dir) so PrefixGrouper works from the installed source after a fresh compile. It is only imported lazily from ``attention_dispatch.run_attention`` when ``prefix_seg_info`` is present, which itself is only ever set when ``UNSLOTH_GRPO_PREFIX_GROUPER`` is on and grouping succeeded, so the default (off) path never touches this file. Provided entry points: * ``PrefixSegInfo`` : per-flat-token segment metadata + cache signature. * ``build_seg_info_multigroup``: build PrefixSegInfo for many groups packed flat. * ``build_seg_info_from_layout``: build PrefixSegInfo for ONE group (test helper). * ``get_block_mask`` : cached create_block_mask keyed on the signature. * ``flex_shared_prefix_attention(Q, K, V, prefix_seg_info)`` Q/K/V of shape [1, T, n_heads, head_dim]; returns [1, T, n_heads, head_dim], IDENTICAL semantics to the SDPA oracle. """ from __future__ import annotations import os from dataclasses import dataclass from typing import Dict, List, Optional, Tuple import torch from torch.nn.attention.flex_attention import ( BlockMask, create_block_mask, flex_attention, ) # GRPO feeds many distinct segment lengths; at dynamo's default recompile_limit (8) the # compiled kernel silently reuses a mismatched specialisation (wrong results). Raise it. torch._dynamo.config.recompile_limit = max(getattr(torch._dynamo.config, "recompile_limit", 8), 256) torch._dynamo.config.accumulated_recompile_limit = max( getattr(torch._dynamo.config, "accumulated_recompile_limit", 256), 2048 ) # Compiled kernels: torch.compile fuses the sparse mask into one kernel. dynamic=True is # required: T changes almost every GRPO batch and dynamic=False recompiles per T (~14s # each). T is still padded to a multiple of 128 (_pad_len) for the backward kernel. _flex_attention_compiled = torch.compile(flex_attention, dynamic = True) _create_block_mask_compiled = torch.compile(create_block_mask, dynamic = True) # Flash block sizes by Q dtype (env-overridable). The two disjoint key runs (prefix + # own-suffix) stress online-softmax accumulation: fp32 needs 32/32 for a ~1e-6 floor; # bf16 passes parity at 128/64 and is ~5x faster (128/128 OOMs Triton on B200). _FP32_BLOCK_M = int(os.environ.get("PG_FLEX_BLOCK_M", "32")) _FP32_BLOCK_N = int(os.environ.get("PG_FLEX_BLOCK_N", "32")) _BF16_BLOCK_M = int(os.environ.get("PG_FLEX_BF16_BLOCK_M", "128")) _BF16_BLOCK_N = int(os.environ.get("PG_FLEX_BF16_BLOCK_N", "64")) def _kernel_options_for_dtype(dtype): """Pick the numerically-safe flash block sizes for the Q dtype.""" if dtype == torch.bfloat16 or dtype == torch.float16: return {"BLOCK_M": _BF16_BLOCK_M, "BLOCK_N": _BF16_BLOCK_N} return {"BLOCK_M": _FP32_BLOCK_M, "BLOCK_N": _FP32_BLOCK_N} # Backward-compat constant (fp32 default). _FLEX_KERNEL_OPTIONS = {"BLOCK_M": _FP32_BLOCK_M, "BLOCK_N": _FP32_BLOCK_N} # The compiled backward trips an Inductor assertion when T is not a multiple of 128, so # pad the flat sequence. Pad tokens form a group that attends to / is attended by nothing # (all-masked rows return 0, not NaN) and are sliced off the output. _PAD_MULTIPLE = 128 _PAD_GROUP = -99 # sentinel group id / suffix id for pad tokens def _pad_len(T: int) -> int: return ((T + _PAD_MULTIPLE - 1) // _PAD_MULTIPLE) * _PAD_MULTIPLE # --------------------------------------------------------------------------- # Segment metadata # --------------------------------------------------------------------------- @dataclass class PrefixSegInfo: """Per-flat-token segment metadata driving the shared-prefix block mask. The label tensors are 1-D of length ``T_pad`` (>= real ``T``, padded up to a multiple of 128 so the backward kernel compiles). Positions ``[T:T_pad)`` are pad tokens (group/suffix == _PAD_GROUP) that attend to nothing. Attributes ---------- group_of_kv : LongTensor [T_pad] Group id per flat token (0..num_groups-1); _PAD_GROUP for pad tokens. is_prefix : BoolTensor [T_pad] True iff the token is a prefix token of its group (False for pad). suffix_of_kv : LongTensor [T_pad] Suffix id per flat token; -1 for prefix, _PAD_GROUP for pad. Suffix ids are globally unique across groups. signature : hashable Cache key for the block mask (depends only on the labels + T_pad). T : int Real flat sequence length (Q/K/V of this length are padded internally). T_pad : int Padded length (multiple of 128) at which the block mask is built. """ group_of_kv: torch.Tensor is_prefix: torch.Tensor suffix_of_kv: torch.Tensor signature: Tuple T: int T_pad: int def _pad_labels(group_of_kv, is_prefix, suffix_of_kv, device): """Pad the label tensors up to a multiple of 128 with pad-token sentinels.""" T = int(group_of_kv.numel()) T_pad = _pad_len(T) if T_pad == T: return group_of_kv, is_prefix, suffix_of_kv, T, T_pad pad = T_pad - T group_of_kv = torch.cat( [group_of_kv, torch.full((pad,), _PAD_GROUP, dtype = torch.long, device = device)] ) is_prefix = torch.cat([is_prefix, torch.zeros(pad, dtype = torch.bool, device = device)]) suffix_of_kv = torch.cat( [suffix_of_kv, torch.full((pad,), _PAD_GROUP, dtype = torch.long, device = device)] ) return group_of_kv, is_prefix, suffix_of_kv, T, T_pad def build_seg_info_from_layout(layout, device: Optional[torch.device] = None) -> PrefixSegInfo: """Build PrefixSegInfo for ONE group from an object with ``.flat_ids``, ``.P`` and ``.suffix_slices`` (used by the parity test / oracle helpers).""" if device is None: device = layout.flat_ids.device T = int(layout.flat_ids.shape[1]) P = int(layout.P) group_of_kv = torch.zeros(T, dtype = torch.long, device = device) # single group -> 0 is_prefix = torch.zeros(T, dtype = torch.bool, device = device) is_prefix[:P] = True suffix_of_kv = torch.full((T,), -1, dtype = torch.long, device = device) for i, (s, e) in enumerate(layout.suffix_slices): suffix_of_kv[s:e] = i group_of_kv, is_prefix, suffix_of_kv, T, T_pad = _pad_labels( group_of_kv, is_prefix, suffix_of_kv, device ) sig = ("single", T_pad, P, tuple((s, e) for (s, e) in layout.suffix_slices)) return PrefixSegInfo( group_of_kv = group_of_kv, is_prefix = is_prefix, suffix_of_kv = suffix_of_kv, signature = sig, T = T, T_pad = T_pad, ) def build_seg_info_multigroup( group_specs: List[Tuple[int, List[int]]], device: torch.device ) -> Tuple[PrefixSegInfo, List[dict]]: """Build PrefixSegInfo for several shared-prefix groups packed block-diagonally. Parameters ---------- group_specs : list of (P_g, [R_{g,0}, R_{g,1}, ...]) For each group: prefix length and the list of suffix lengths. Returns ------- seg : PrefixSegInfo group_meta : list of dicts with 'base', 'P', 'prefix_last_index', 'suffix_slices' (flat offsets), enough to build the completion index map. """ group_of_list = [] is_prefix_list = [] suffix_of_list = [] group_meta = [] base = 0 suffix_counter = 0 sig_parts = [] for gid, (P, R_list) in enumerate(group_specs): # prefix group_of_list.append(torch.full((P,), gid, dtype = torch.long, device = device)) is_prefix_list.append(torch.ones(P, dtype = torch.bool, device = device)) suffix_of_list.append(torch.full((P,), -1, dtype = torch.long, device = device)) prefix_last_index = base + P - 1 suffix_slices = [] cursor = base + P for r in R_list: group_of_list.append(torch.full((r,), gid, dtype = torch.long, device = device)) is_prefix_list.append(torch.zeros(r, dtype = torch.bool, device = device)) suffix_of_list.append(torch.full((r,), suffix_counter, dtype = torch.long, device = device)) suffix_slices.append((cursor, cursor + r)) cursor += r suffix_counter += 1 group_meta.append( { "base": base, "P": P, "prefix_last_index": prefix_last_index, "suffix_slices": suffix_slices, } ) sig_parts.append((P, tuple(R_list))) base = cursor group_of_kv = torch.cat(group_of_list) is_prefix = torch.cat(is_prefix_list) suffix_of_kv = torch.cat(suffix_of_list) group_of_kv, is_prefix, suffix_of_kv, T, T_pad = _pad_labels( group_of_kv, is_prefix, suffix_of_kv, device ) sig = ("multi", T_pad, tuple(sig_parts)) seg = PrefixSegInfo( group_of_kv = group_of_kv, is_prefix = is_prefix, suffix_of_kv = suffix_of_kv, signature = sig, T = T, T_pad = T_pad, ) return seg, group_meta # --------------------------------------------------------------------------- # Block-mask builder + cache, keyed on (signature, device): the mask depends only on the # per-token labels and T, so it is reused across layers and steps. _BLOCK_MASK_CACHE: Dict[Tuple, BlockMask] = {} def _make_mask_mod(group_of_kv, is_prefix, suffix_of_kv): """Return a mask_mod closure over the (device) label tensors. keep(q, kv) = same_group AND ( is_prefix[kv] AND kv <= q # causal within/ into prefix OR ( suffix_of_kv[kv] == suffix_of_kv[q] # same suffix ... AND (not is_prefix[q]) # q is a suffix token ... AND kv <= q ) ) # ... causal within it The single ``kv <= q`` guard on the is_prefix branch gives BOTH prefix-causal behaviour (a prefix q sees only earlier prefix tokens) AND full-prefix-visibility for suffixes (every prefix index < every suffix index in a group, so kv <= q always holds for a suffix q vs a prefix kv of its group), matching the SDPA oracle exactly. """ def mask_mod(b, h, q_idx, kv_idx): same_group = group_of_kv[q_idx] == group_of_kv[kv_idx] kv_is_prefix = is_prefix[kv_idx] causal = kv_idx <= q_idx same_suffix = (suffix_of_kv[kv_idx] == suffix_of_kv[q_idx]) & (~is_prefix[q_idx]) keep = same_group & ((kv_is_prefix & causal) | (same_suffix & causal)) return keep return mask_mod def get_block_mask( seg: PrefixSegInfo, device: torch.device, compile_mask: bool = True, ) -> BlockMask: """Return a cached BlockMask for the segment signature (built once, reused). CRITICAL: the block mask is cached and shared across BOTH the no-grad old/ref logprob forward (which runs under torch.inference_mode) and the grad training forward. If the mask were first built under inference_mode, its tensors would be INFERENCE tensors that "cannot be saved for backward" when reused in the grad forward. We therefore build the mask with inference mode explicitly DISABLED, so the same cached BlockMask is a normal tensor usable by autograd. (The mask depends only on integer labels; it needs no grad.) """ key = (seg.signature, str(device)) bm = _BLOCK_MASK_CACHE.get(key) if bm is not None: return bm # Move labels to the consumer (Q) device: with a sharded model the seg tensors live on # input_ids.device and would index cross-device. Copies once per (signature, device). # These copies must also run with inference mode DISABLED (same reason as the mask build): # when this entry is first built under the no-grad old/ref forward's inference_mode and # device != seg.device, a .to(device) copy would be an inference tensor that mask_mod # captures, which then cannot be saved for backward when the grad training forward reuses # the cached mask. builder = _create_block_mask_compiled if compile_mask else create_block_mask with torch.inference_mode(False): mask_mod = _make_mask_mod( seg.group_of_kv.to(device), seg.is_prefix.to(device), seg.suffix_of_kv.to(device) ) bm = builder( mask_mod, B = 1, H = None, Q_LEN = seg.T_pad, KV_LEN = seg.T_pad, device = device, ) # FIFO bound: GRPO lengths change nearly every step, so evict the oldest to cap GPU pins. if len(_BLOCK_MASK_CACHE) >= 8: _BLOCK_MASK_CACHE.pop(next(iter(_BLOCK_MASK_CACHE))) _BLOCK_MASK_CACHE[key] = bm return bm def clear_block_mask_cache(): _BLOCK_MASK_CACHE.clear() def _pad_qkv_seq(x: torch.Tensor, T_pad: int) -> torch.Tensor: """Zero-pad a [B, H, T, D] tensor along the sequence dim up to T_pad.""" T = x.shape[2] if T_pad == T: return x pad = torch.zeros(x.shape[0], x.shape[1], T_pad - T, x.shape[3], device = x.device, dtype = x.dtype) return torch.cat([x, pad], dim = 2) def _run_flex(q, k, v, block_mask, enable_gqa, scale, compiled, T, T_pad): """Pad q/k/v to T_pad, run flex, slice the output back to T. q/k/v: [B,H,T,D].""" qp = _pad_qkv_seq(q, T_pad) kp = _pad_qkv_seq(k, T_pad) vp = _pad_qkv_seq(v, T_pad) if compiled: out = _flex_attention_compiled( qp, kp, vp, block_mask = block_mask, enable_gqa = enable_gqa, scale = scale, kernel_options = _kernel_options_for_dtype(qp.dtype), ) else: # eager path (fp64 parity): dense scores, no kernel_options. out = flex_attention( qp, kp, vp, block_mask = block_mask, enable_gqa = enable_gqa, scale = scale, ) return out[:, :, :T, :] # --------------------------------------------------------------------------- # The kernel entry point # --------------------------------------------------------------------------- def flex_shared_prefix_attention( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, prefix_seg_info: PrefixSegInfo, scale: Optional[float] = None, block_mask: Optional[BlockMask] = None, compiled: bool = True, ) -> torch.Tensor: """Shared-prefix attention via FlexAttention. Parameters ---------- Q, K, V : Tensor [1, T, n_heads, head_dim] (Q has n_heads, K/V have n_kv_heads for GQA). prefix_seg_info : PrefixSegInfo scale : optional float, softmax scale (defaults to 1/sqrt(head_dim)). block_mask : optional precomputed BlockMask (else built/cached from seg info). Returns ------- Tensor [1, T, n_heads, head_dim], identical semantics to the SDPA oracle branch. """ assert Q.dim() == 4 and Q.shape[0] == 1, f"expected [1,T,H,D], got {tuple(Q.shape)}" device = Q.device # FlexAttention wants [B, H, T, D]. q = Q.transpose(1, 2) # [1, n_heads, T, D] k = K.transpose(1, 2) # [1, n_kv_heads, T, D] v = V.transpose(1, 2) n_heads = q.shape[1] n_kv = k.shape[1] enable_gqa = n_heads != n_kv T = q.shape[2] T_pad = prefix_seg_info.T_pad assert T == prefix_seg_info.T, f"Q length {T} != seg.T {prefix_seg_info.T}" if block_mask is None: block_mask = get_block_mask(prefix_seg_info, device, compile_mask = compiled) out = _run_flex(q, k, v, block_mask, enable_gqa, scale, compiled, T, T_pad) # back to [1, T, n_heads, D] return out.transpose(1, 2).contiguous() __all__ = [ "PrefixSegInfo", "build_seg_info_multigroup", "build_seg_info_from_layout", "get_block_mask", "clear_block_mask_cache", "flex_shared_prefix_attention", ]