From 08e133cd6b035adc75474985cb8c097c5b1a3f95 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:26:24 -0700 Subject: [PATCH] Add PrefixGrouper for GRPO: dedup the shared prompt across a group's completions (#6871) * GRPO: optional sequence packing for the no-grad old/ref logp path Add an opt-in sequence-packing fast path to _get_per_token_logps_and_entropies, enabled with UNSLOTH_GRPO_SEQ_PACKING=1. When the batch is text-only, the padded [B, Lmax] per-chunk forward is replaced by a single varlen [1, sum L] forward (BlockDiagonalCausalMask via packed_seq_lengths with reset position_ids). Per-token logps use the same float32 chunked_hidden_states_selective_log_softmax as the padded path, so the old and reference logps are bit-for-bit identical. Safety: the packed path is self-verified once against the padded ground truth on a batch that has at least two rows with real completion tokens (self._unsloth_seq_packing_nograd_ok), so cross-sample contamination would actually manifest; a degenerate all-pad / fully tool-masked batch leaves the verdict unset and re-verifies later. If a backend silently ignores packed_seq_lengths (flat batch run under a normal causal mask, samples leaking across boundaries), the packed logps will not match and packing is disabled instead of corrupting logps. It also forces use_cache=False (a populated past_key_value disables varlen packing), skips packing when a sliding window is shorter than the packed stream, runs the same GPT-OSS offload device_synchronize the padded loop uses, and falls back on any exception (UNSLOTH_GRPO_SEQ_PACKING_DEBUG=1 prints the reason). Default off, so existing behavior is unchanged. Pairs with the matching gradient-path change in unsloth_zoo so the full GRPO logp + loss + backward can run packed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: address review feedback - Cache the packed-vs-padded verdict per unwrapped model instead of on the trainer, so a separately forwarded reference model is verified on its own forward path rather than inheriting the policy model's verdict. - Force the padded path when token_type_ids or mm_token_type_ids are present, matching the extra vision kwargs the padded loop forwards. - Require the xformers varlen backend before packing. Without it the packed mask falls back to a dense O(T^2) SDPA mask that can OOM on the flattened batch, so we keep the padded loop in that case. - On any packed-forward failure (missing backend, OOM, unsupported forward) empty the cache on OOM, disable packing for that model, and fall back to the chunked padded loop instead of retrying every step. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: default-on, verify against per-row reference Redesign of the optional sequence-packing fast path for the no-grad old/ref logprob recompute, after establishing that the packed forward is the exact per-row computation and the padded batch forward is the side that mis-positions left-padded rows on long completions. - Default the packing on (UNSLOTH_GRPO_SEQ_PACKING, disable with 0). - Verify the packed logprobs against the per-row clean forward (each row's real tokens alone, reset 0-based positions, no padding), not the padded batch which is itself wrong for left-padding. Cross-sample contamination (a backend ignoring packed_seq_lengths) shows up as a large mismatch and falls back to the padded loop. - Make the trust decision shape and RoPE aware: re-verify whenever the packed total length or the longest segment grows past what was verified, so a later batch crossing a LongRoPE short/long cache boundary is re-checked instead of trusted blindly. - Run lm_head only on completion-prediction positions instead of every packed prompt token, so long-prompt/short-completion batches do not pay for projecting the whole packed prompt. - Drop the hard xformers import so the path also runs in FlashAttention-only environments; the per-row verification guards correctness regardless of backend. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: disable entirely on cross-sample mismatch When the per-row verification fails, distinguish the two failure modes by magnitude instead of by sequence length: - A large mismatch (>= 1.5) is the cross-sample contamination signature: the model's attention does not honor the block-diagonal packed mask (seen on some MoE / custom-attention models, e.g. qwen2_moe). Disable packing entirely for the model so later batches do not pay the verification cost again. - A moderate mismatch is more likely a length-boundary effect (a LongRoPE short/long cache switch): keep marking just that length region unsafe so packing still runs for smaller shapes. Validated: Qwen1.5-MoE falls back after a single verification (grad and no-grad ok flags go False, no re-verify on later steps); dense Llama-3.2 and Qwen3 still verify and engage packing. * GRPO no-grad packing: trim comments to be concise * GRPO no-grad packing: fix per-row completion boundary for left-padded rows The completion-target selection used a single global boundary (col >= L - logits_to_keep). After left-packing, each row's completion starts at (L - logits_to_keep) - left_pad[row], so for left-padded rows the first left_pad completion tokens fall below the global boundary and were dropped, leaving 0 logprobs at real completion positions that the loss mask keeps. Use the per-row boundary so packed coverage matches create_completion_attention_mask exactly, and widen the self-verify mask to the full per-row completion region so it can catch coverage gaps. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: gate verification on real completion rows Count active rows via create_completion_attention_mask (the same mask the loss uses) instead of any non-pad token in the packed window. Prompt-only rows carry prompt-overflow tokens in the window and could otherwise satisfy the >= 2 verification guard, letting a batch with a single real completion row cache a trust decision. This matches the gradient path, which already gates on the completion mask. The same mask is reused for the self-verify comparison. * GRPO no-grad packing: gate debug logging on UNSLOTH_ENABLE_LOGGING Use the shared UNSLOTH_ENABLE_LOGGING global (import_fixes, re-exported by _utils) instead of a bespoke UNSLOTH_GRPO_SEQ_PACKING_DEBUG env var for the packing debug prints, matching the rest of the codebase. * GRPO packing: import UNSLOTH_ENABLE_LOGGING inside the injected logp function _get_per_token_logps_and_entropies is copied verbatim into the generated GRPO trainer via inspect.getsource, and that module never imported UNSLOTH_ENABLE_LOGGING, so the default-on packing verify path raised NameError (and the except handler re-raised it). Import the flag locally, before the try, so the name is defined in the generated module too. Drop it from the now-unused module-level import. * GRPO no-grad packing: harden unsafe-length skip, verify guard, fallback cleanup Three fixes to the no-grad logp packing path, mirroring the grad path: - skip the packed forward for known-unsafe lengths by reading unsafe_T and gating on it before the forward, instead of running the full packed pass and the result build only to discard them (wastes a pass, can OOM at large T) - only widen the verified T/seg envelope when >= 2 completion rows actually exercised cross-sample packing; a < 2 row batch cannot expose leakage, so it must not extend the trusted shape that later multi-row batches skip verify for - drop the packed intermediates (hidden/sel/result/ref) before the padded fallback loop so it does not run with the flattened hidden state still resident * GRPO no-grad packing: cap the flattened forward at one mini-batch budget The packed path built a single [1, sum L] forward over every row before any size check, so a large batch could exceed the memory the padded path bounds per mini-batch. Gate packing on _pk_T <= _pk_cap (B * seq_len, one padded mini-batch's token budget); larger batches fall back to the chunked padded loop. * GRPO no-grad packing: disable unless unsloth_zoo has the masked-column guard The packed path leaves masked prompt/pad logprob columns at 0, which only stays finite if unsloth_zoo grpo_compute_loss zeroes them before exp() (zoo#840). An older unsloth_zoo without that guard would NaN. Detect the guard once (cached on the model) via inspect.getsource and gate packing on it, so #6738 is safe with any unsloth_zoo version and re-enables packing automatically once a guarded zoo is installed, independent of the pinned lower bound. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO packing: hoist env gates and zoo-guard detection to one-time module checks Read UNSLOTH_GRPO_SEQ_PACKING and detect the unsloth_zoo masked-column guard once at import time (module constants plus RL_PRE_ITEMS for the generated trainer cache) instead of per call, and drop the in-function UNSLOTH_ENABLE_LOGGING import for a module-top one. The UNSLOTH_GRPO_SEQ_PACKING_VERIFY force-verify debug knob is commented out, kept in place for hand re-enable; the first-use and envelope-growth self-verify stays active. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO packing: cap the flattened forward by the padded chunk rows B counts chunks at this point, so B * seq_len understated (small runs) or overstated (large runs) the padded mini-batch token budget; use batch_size * seq_len, the rows the padded loop actually forwards per chunk. * Add PrefixGrouper for GRPO: dedup the shared prompt across a group's completions In GRPO every prompt spawns G=num_generations completions that share the prompt prefix, so the trunk logprob forward re-encodes that prefix G times. PrefixGrouper stores the prefix once and concatenates only the G suffixes behind a FlexAttention shared-prefix mask, cutting the forward from G*(P+R) to P+G*R tokens across both the no-grad old/ref forwards and the grad logp forward. Default off behind the UNSLOTH_GRPO_PREFIX_GROUPER env gate, so the gate-unset path is byte-identical to today. A tok_r auto-gate and a first-use self-verify (fall back and mark the shape unsafe on mismatch) keep it from ever shipping wrong logprobs silently. Wired for llama, mistral, qwen3, gemma2, cohere, granite and falcon_h1, plus qwen2 and gemma through the shared LlamaAttention_fast_forward. Stacked on the GRPO sequence-packing PR (#6738); the grad path lands in a companion unsloth-zoo PR. Also fixes a latent UNSLOTH_ENABLE_LOGGING NameError in the seq-packing no-grad verify path by defining the name as a generated-cache pre-item. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: enforce the sliding-window cap, gate softcap models, bound the mask cache Add a max_segment_cap kwarg to build_group_layout so it falls back when a group's span (prefix + longest suffix) exceeds the model's local window, and pass the config sliding_window into the no-grad engage gate the same way the packed _pk guard derives it. Skip PrefixGrouper entirely for attn_logit_softcapping models, since the FlexAttention kernel never applies logit softcapping. Bound _BLOCK_MASK_CACHE to a FIFO of 8 so per-step lengths cannot pin BlockMasks forever, release the PG hidden before the verify forward, and align the UNSLOTH_ENABLE_LOGGING pre-item truthiness with the canonical form. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: vectorize the real-column scan in build_group_layout Replace the per-row O(B*L) Python scan of the keep mask with a GPU-derived contiguous-run fast path (first real column + count per row), keeping the general scan only as a fallback for non-contiguous rows. Works for both call sites: the no-grad layout (left-padded prompt + right-padded completion, run does not start at column 0) and the grad layout (left-packed). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: hoist the gate and kernel imports to one-time module checks, AGPLv3 headers Read UNSLOTH_GRPO_PREFIX_GROUPER and resolve the prefix_grouper imports once at module level (source constants plus an RL_PRE_ITEMS entry for the generated trainer cache) instead of per call, matching the sequence-packing gates. The prefix_grouper env helpers become one-time module reads with unchanged signatures, and attention_dispatch resolves the FlexAttention kernel once behind the same gate (lazy fallback kept). The two new prefix_grouper files move to AGPLv3 headers. * PrefixGrouper: length-envelope trust and hybrid SSM exclusion Verified signatures now record (max T, max segment) and re-verify when either grows, matching the packed path's envelope. Hybrid SSM models (FalconH1 etc.) are excluded at the gate since only attention gets the shared-prefix isolation, and the FalconH1 wiring is removed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: defer the unverified no-grad forward until the packed reference exists Unverified shapes no longer run the whole-batch shared-prefix forward up front; it now runs at the verify site, only when the packed path produced a reference. A declined packed path (budget, window) therefore costs no wasted PG forward per step. Trusted shapes still run it first to skip the full-row forward, with the same fallback. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: disable under vLLM (fast_inference=True) With colocated vLLM generation the rollout dominates the GRPO step, so the shared-prefix training forward saves little end-to-end and its first-use self-verify (which also runs the full-row path) is net overhead. Gate PG on not use_vllm so it only engages on the raw transformers path, where the training forward is on the critical path. Packing is unaffected. * PrefixGrouper: compile the FlexAttention kernel with dynamic shapes GRPO changes the packed length T almost every batch. With dynamic=False the flex forward+backward kernel recompiled on every new T (~14s each on a 4B trunk), which dominated the step and made PG a net loss. dynamic=True compiles once, then reuses the kernel across all lengths recompile-free (a new shape drops from ~14s to ~1.4ms after a two-graph warmup). T is still padded to a multiple of 128 for the backward block assertion. * PrefixGrouper: default on Enable PrefixGrouper by default (UNSLOTH_GRPO_PREFIX_GROUPER defaults to 1; set 0 to disable). Still auto-disabled under vLLM (fast_inference=True) and by the arch/softcap/ SSM/tok_r gates, and the first-use self-verify falls back on any mismatch, so this is a memory-first default on the raw-transformers path with no correctness risk. * GRPO PrefixGrouper: gate on zoo masked-column guard and exclude MoE - Require the zoo masked-column guard (zoo#840) before PrefixGrouper can engage. PG rides the sequence-packing path, so when the first-step self-verify is off the fast path trusts PG output directly; without the guard those masked columns feed NaN into the packed loss. Gate PG on the same UNSLOTH_ZOO_HAS_MASKED_COL_GUARD the packing path already checks. - Exclude MoE configs (num_experts, num_local_experts, n_routed_experts, moe_intermediate_size) alongside the hybrid-SSM markers. Only the threaded attention forwards carry the shared-prefix isolation, so a MoE decoder that does not forward prefix_seg_info would let suffixes leak across completions. - Refresh the stale default-off comments now that UNSLOTH_GRPO_PREFIX_GROUPER is on by default. * GRPO PrefixGrouper: import chunked_hidden_states_selective_log_softmax The shared-prefix forward passes chunked_hidden_states_selective_log_softmax into extract_logps, but the name was only ever provided by the generated trainer cache (rl.py injects grpo_selective_log_softmax_code), never bound in this module. Import it from unsloth_zoo.rl_replacements next to its sibling chunked_selective_log_softmax so the source resolves the name in every scope (the new _pg_run_forward closure included). No runtime change: the cache still defines the function via template injection. * GRPO PrefixGrouper: dropout gate, device-safe layout, Mistral mask skip Addresses three review findings on the shared-prefix path: - Skip PrefixGrouper when the model sets a nonzero attention_dropout. The normal backends apply config.attention_dropout while training (e.g. Granite dense flash/sdpa/xformers), but the FlexAttention shared-prefix path is deterministic, so gate PG off for those configs rather than train on mismatched activations. - Move the shared-prefix mask labels to the consumer (Q) device in get_block_mask and the target index maps to hidden.device in extract_logps, mirroring the packed path moving its metadata to the consumer device. Prevents cross-device indexing when the model is sharded across GPUs. - Do not synthesize a causal attention_mask in the Mistral forward when prefix_seg_info is present. On the no-xFormers path that synthetic mask tripped resolve_prefix_seg_info and forced PG to always fall back to the packed forward. * GRPO sequence packing: tighten comments * GRPO PrefixGrouper: tighten comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO PrefixGrouper: persistent disable on runtime failure; build block-mask labels with inference mode disabled - rl_replacements: on a PG forward exception (FlexAttention/Triton compile failure or OOM), set a model-level _unsloth_prefix_grouper_nograd_disabled flag and consult it in the engage gate, mirroring the seq-packing handler, so a GPU-wide failure is not retried and re-paid every step. - prefix_grouper_kernel: move the .to(device) label copies inside the inference_mode(False) block so a cross-device (model-parallel shard) first build does not capture inference tensors, which otherwise cannot be saved for backward when the grad training forward reuses the cached BlockMask. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- unsloth/models/cohere.py | 5 + unsloth/models/gemma2.py | 7 + unsloth/models/granite.py | 5 + unsloth/models/llama.py | 12 +- unsloth/models/mistral.py | 13 +- unsloth/models/qwen3.py | 5 + unsloth/models/rl_replacements.py | 263 ++++++++++++++- unsloth/utils/attention_dispatch.py | 65 ++++ unsloth/utils/prefix_grouper.py | 351 ++++++++++++++++++++ unsloth/utils/prefix_grouper_kernel.py | 436 +++++++++++++++++++++++++ 10 files changed, 1156 insertions(+), 6 deletions(-) create mode 100644 unsloth/utils/prefix_grouper.py create mode 100644 unsloth/utils/prefix_grouper_kernel.py diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index 0b7f3ab973..cb367d451e 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -22,6 +22,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, ) try: @@ -151,6 +152,9 @@ def CohereAttention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -161,6 +165,7 @@ def CohereAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 68b9ebe22f..4a0531db78 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -22,6 +22,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, SDPA, ) from .gemma import ( @@ -168,6 +169,11 @@ def Gemma2Attention_fast_forward( }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. gemma2 is + # sliding-window and softcapped: the engage gate caps spans at the window and + # excludes softcap models entirely, so PG never engages here. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -179,6 +185,7 @@ def Gemma2Attention_fast_forward( attention_mask = attention_mask, causal_mask = causal_mask, sliding_window = sliding_window, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index f5b0f57aa6..4dedf642eb 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -23,6 +23,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, SDPA, ) from .llama import ( @@ -159,6 +160,9 @@ def GraniteAttention_fast_forward( }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -169,6 +173,7 @@ def GraniteAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6bec95b577..417a88f480 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -39,6 +39,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from torch.nn.functional import scaled_dot_product_attention from transformers import __version__ as transformers_version @@ -738,6 +739,10 @@ def LlamaAttention_fast_forward( flash_dense_kwargs = {"causal": True}, flash_varlen_kwargs = {"dropout_p": 0.0, "causal": True}, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward (same route + # as packed_seq_lengths); misuse (KV cache / padding mask) raises. None => byte-identical + # default. Reuse of this forward also carries the branch to qwen2 & gemma. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -748,6 +753,7 @@ def LlamaAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = config, context = context, Q = Q, K = K, V = V) @@ -895,8 +901,10 @@ def LlamaModel_fast_forward( seq_length_with_past = seq_length # Fix out of bounds tokenization unless we were given packed metadata - allow_overlength = getattr(self, "_unsloth_allow_packed_overlength", False) or ( - "packed_seq_lengths" in kwargs + allow_overlength = ( + getattr(self, "_unsloth_allow_packed_overlength", False) + or ("packed_seq_lengths" in kwargs) + or ("prefix_seg_info" in kwargs and kwargs["prefix_seg_info"] is not None) ) if hasattr(self, "max_seq_length") and not allow_overlength: if seq_length > self.max_seq_length: diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index df2a4de5bd..4350565fe2 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -27,6 +27,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from .llama import ( LlamaRotaryEmbedding, @@ -124,6 +125,9 @@ def MistralAttention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -134,6 +138,7 @@ def MistralAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) @@ -161,7 +166,13 @@ def MistralForCausalLM_fast_forward( *args, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: - if causal_mask is None and past_key_values is None: + # PrefixGrouper brings its own mask: a synthesized causal attention_mask would trip + # resolve_prefix_seg_info on the no-xFormers path and force a fallback. + if ( + causal_mask is None + and past_key_values is None + and kwargs.get("prefix_seg_info", None) is None + ): bsz, q_len = input_ids.shape sliding_window = getattr(self.config, "sliding_window", None) diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index e28e72d3ea..0d05a2d538 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -23,6 +23,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from .llama import ( LlamaRotaryEmbedding, @@ -146,6 +147,9 @@ def Qwen3Attention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -156,6 +160,7 @@ def Qwen3Attention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 0573fd5fd8..098950de08 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -31,6 +31,7 @@ from unsloth_zoo.rl_replacements import ( left_pack_padding, create_completion_attention_mask, chunked_selective_log_softmax, + chunked_hidden_states_selective_log_softmax, _unsloth_get_mm_token_id, _unsloth_fix_mm_token_type_ids, ) @@ -65,6 +66,25 @@ try: ) except Exception: UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False +# One-time PrefixGrouper gate; any import failure degrades to "PrefixGrouper off". +_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None +UNSLOTH_GRPO_PREFIX_GROUPER_ON = os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ( + "0", + "false", + "no", + "off", +) +if UNSLOTH_GRPO_PREFIX_GROUPER_ON: + try: + from ..utils.prefix_grouper import ( + build_group_layout as _pg_build_layout, + prefix_grouper_enabled as _pg_enabled_fn, + verify_on as _pg_verify_on, + tol_ok as _pg_tol_ok, + TOL_KILL as _PG_TOL_KILL, + ) + except Exception: + UNSLOTH_GRPO_PREFIX_GROUPER_ON = False RL_EXTRA_ARGS = defaultdict(list) RL_FUNCTIONS = defaultdict(list) @@ -1380,6 +1400,166 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): # left-pad RoPE error). Self-verified against the per-row forward, re-checked as T # grows; falls back if a backend ignores packed_seq_lengths. logprobs = None + + # ---- PrefixGrouper (GRPO shared-prompt dedup; default ON, exact + self-verified) ---- + # G completions per prompt share the prefix; the packed path forwards it G times, + # PrefixGrouper stores it once (FlexAttention shared-prefix mask), cutting the trunk + # forward from G*(P+R) to P+G*R tokens. Gated by UNSLOTH_GRPO_PREFIX_GROUPER (needs + # seq-packing), tok_r auto-gate, and first-use self-verify vs the packed path + # (mismatch => fall back + mark unsafe), so a mask/isolation regression cannot ship + # silently. When off / ungrouped / unverified, the packed path below runs as before. + _pg_result = None + _pg_use = False + _pg_skip_pk = False # once a shape is PG-verified, skip the full-row forward + _pg_forward_fn = None # deferred PG forward (runs at the verify site below) + _pg_num_gen = getattr(self, "num_generations", None) + # Env gate hoisted to module level (mirrored via RL_PRE_ITEMS). Skip PG under vLLM + # (fast_inference=True): the rollout dominates the step, so PG saves little and its + # first-use self-verify is net overhead. + _pg_engage = ( + UNSLOTH_GRPO_PREFIX_GROUPER_ON + and not getattr(self, "use_vllm", False) + and not getattr(unwrapped_model, "_unsloth_prefix_grouper_nograd_disabled", False) + ) + if _pg_engage: + try: + # Skip softcap models (the flex kernel never applies attn_logit_softcapping) + # and hybrid SSM / MoE models: only the threaded attention forwards get the + # shared-prefix isolation, so a Mamba or MoE decoder that does not forward + # prefix_seg_info would leak suffixes across completions. PG also rides on + # sequence packing, so it needs the same zoo masked-column guard. + _pg_cfg = getattr(unwrapped_model, "config", None) + _pg_engage = ( + _pg_enabled_fn() + and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD + and pixel_values is None + and token_type_ids is None + and mm_token_type_ids is None + and _pg_num_gen is not None + and _pg_num_gen >= 2 + and not getattr(_pg_cfg, "attn_logit_softcapping", None) + # normal backends apply config.attention_dropout in training; the flex + # path is deterministic, so skip PG when it is set. + and not getattr(_pg_cfg, "attention_dropout", 0) + and not any( + getattr(_pg_cfg, _pg_a, None) is not None + for _pg_a in ( + "mamba_d_ssm", + "mamba_d_state", + "mamba_expand", + "num_experts", + "num_local_experts", + "n_routed_experts", + "moe_intermediate_size", + ) + ) + ) + except Exception: + _pg_engage = False + if _pg_engage: + try: + _pg_pad = self.processing_class.pad_token_id + # cap the PG span (P+max(R)) at the sliding window, like the packed _pk_sw guard. + _pg_sw = getattr( + getattr(unwrapped_model, "config", None), "sliding_window", None + ) + if not (isinstance(_pg_sw, int) and _pg_sw > 0): + _pg_sw = None + _pg_layout = _pg_build_layout( + input_ids, + logits_to_keep, + _pg_pad, + _pg_num_gen, + left_pad_tokens_per_prompt, + max_segment_cap = _pg_sw, + ) + _pg_unsafe = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None + ) + if _pg_unsafe is None: + _pg_unsafe = set() + if _pg_layout is not None and _pg_layout.signature not in _pg_unsafe: + _pg_sig = _pg_layout.signature + _pg_verified = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None + ) + if _pg_verified is None: + _pg_verified = set() + _pg_chunks = max(1, total_rows * multiplier) + + def _pg_run_forward(_pg_layout = _pg_layout, _pg_chunks = _pg_chunks): + with _get_inference_mode_context_manager(model): + with torch.amp.autocast( + device_type = "cuda", dtype = self._autocast_dtype + ): + _pg_hidden = unwrapped_model( + input_ids = _pg_layout.flat_ids, + position_ids = _pg_layout.position_ids, + prefix_seg_info = _pg_layout.prefix_seg_info, + use_cache = False, + ).logits + _pg_r = _pg_layout.extract_logps( + _pg_hidden, + lm_head, + chunked_hidden_states_selective_log_softmax, + _pg_chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + ) + _pg_hidden = None # release before any verify forward + device_synchronize() + # clip to the loss window [B, logits_to_keep+max_left_pad] + _pg_w = logits_to_keep + max_left_pad + if _pg_r.shape[1] > _pg_w: + _pg_r = _pg_r[:, -_pg_w:] + return _pg_r + + # trust only within the verified envelope: re-verify when T or the + # longest segment grows, like the packed path + _pg_T = int(_pg_layout.flat_ids.shape[1]) + _pg_maxseg = int(_pg_layout.position_ids.max()) + 1 + _pg_env = ( + _pg_verified.get(_pg_sig) if isinstance(_pg_verified, dict) else None + ) + if (not _pg_verify_on()) or ( + _pg_env is not None and _pg_T <= _pg_env[0] and _pg_maxseg <= _pg_env[1] + ): + # trusted shape: run PG now and skip the full-row forward below + _pg_result = _pg_run_forward() + _pg_use = True + _pg_skip_pk = True + else: + # unverified shape: defer the forward until the packed reference + # exists (verify site below), so a declined packed path never wastes + # a whole-batch PG forward + _pg_forward_fn = _pg_run_forward + except Exception as _pg_err: + _pg_result = None + _pg_use = False + _pg_skip_pk = False + _pg_forward_fn = None + # A FlexAttention/Triton compile failure or OOM here is GPU-wide, not + # layout-specific, so retrying the same PG forward every step just re-pays + # the failure. Persistently disable PG (mirrors the seq-packing handler + # setting _unsloth_seq_packing_nograd_ok = False); the packed/padded path + # below still produces the exact result. + unwrapped_model._unsloth_prefix_grouper_nograd_disabled = True + if isinstance(_pg_err, torch.cuda.OutOfMemoryError): + torch.cuda.empty_cache() + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) disabled (fell back to packed): {_pg_err!r}", + flush = True, + ) + + # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ---- + # One varlen [1, sum L] block-diagonal forward replaces the padded [B, Lmax] loop + # (exact per-row result; also fixes the padded path's left-pad RoPE error). + # Self-verified vs the per-row forward, re-checked as T grows; falls back if a + # backend ignores packed_seq_lengths. lm_head runs on completion positions only. _pk_result = None _pk_use = False _pk_enabled = UNSLOTH_GRPO_SEQ_PACKING_ON @@ -1388,6 +1568,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): _pk_ok = getattr(unwrapped_model, "_unsloth_seq_packing_nograd_ok", None) if ( _pk_enabled + and not _pg_skip_pk and pixel_values is None and token_type_ids is None and mm_token_type_ids is None @@ -1462,7 +1643,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): )[0] # GPT-OSS offload race guard (matches the padded loop) device_synchronize() - # scatter each completion logprob back to its (row, col) so [:, -_pk_W:] matches padded + # scatter each logprob back to its (row, col) so [:, -_pk_W:] matches padded _pk_tgt = (_pk_nz_idx[1:, 0] * _pk_L + _pk_nz_idx[1:, 1])[_pk_ctgt] _pk_result = ( torch.zeros( @@ -1574,7 +1755,73 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): f"[Unsloth] GRPO sequence-packing (no-grad) disabled (fell back to padded): {_pk_err!r}", flush = True, ) - if _pk_use and _pk_result is not None: + # ---- PrefixGrouper first-use self-verify (no-grad) ---- + # Compare the untrusted PG result to the full-row packed result (itself verified vs + # per-row) over the completion mask: < tol_ok -> trust the structure; >= TOL_KILL -> + # unsafe forever; borderline -> fall back this shape. + if _pg_forward_fn is not None and not _pg_use: + if _pk_use and _pk_result is not None: + try: + # deferred PG forward, run only now that the packed reference exists + _pg_result = _pg_forward_fn() + _pg_W2 = logits_to_keep + max_left_pad + _pg_cm = create_completion_attention_mask( + input_ids[:, -_pg_W2:], + left_pad_tokens_per_prompt, + max_left_pad, + self.processing_class.pad_token_id, + ).float() + _pg_a = _pg_result[:, -_pg_W2:].float() + _pg_b = _pk_result[:, -_pg_W2:].float() + _pg_diff = float(((_pg_a - _pg_b).abs() * _pg_cm).max()) + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) verify: sig={_pg_layout.signature} " + f"shared-prefix vs full-row-packed max|d|={_pg_diff:.4f}", + flush = True, + ) + if _pg_diff < _pg_tol_ok(): + _pg_v = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None + ) + if not isinstance(_pg_v, dict): + _pg_v = {} + _pg_vT = int(_pg_layout.flat_ids.shape[1]) + _pg_vS = int(_pg_layout.position_ids.max()) + 1 + _pg_old = _pg_v.get(_pg_layout.signature, (0, 0)) + _pg_v[_pg_layout.signature] = ( + max(_pg_vT, _pg_old[0]), + max(_pg_vS, _pg_old[1]), + ) + unwrapped_model._unsloth_prefix_grouper_nograd_verified = _pg_v + _pg_use = True + else: + _pg_u = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None + ) + if _pg_u is None: + _pg_u = set() + if _pg_diff >= _PG_TOL_KILL: + _pg_u.add(_pg_layout.signature) + unwrapped_model._unsloth_prefix_grouper_nograd_unsafe = _pg_u + _pg_use = False + except Exception as _pg_err3: + _pg_result = None + _pg_use = False + if isinstance(_pg_err3, torch.cuda.OutOfMemoryError): + torch.cuda.empty_cache() + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) verify failed (fell back to packed): {_pg_err3!r}", + flush = True, + ) + # else: no packed reference (packing off/failed) -> cannot verify; fall back. + + if _pg_use and _pg_result is not None: + logprobs = _pg_result # PrefixGrouper verified/trusted -> skip the loop + zipped_inputs = [] + elif _pk_use and _pk_result is not None: logprobs = _pk_result # verified -> skip the loop zipped_inputs = [] else: @@ -1752,7 +1999,7 @@ RL_PRE_ITEMS["grpo_trainer"].append( "import os as _unsloth_os\n" "UNSLOTH_ENABLE_LOGGING = _unsloth_os.environ.get('UNSLOTH_ENABLE_LOGGING', '0') in ('1', 'True', 'true')\n" ) -# One-time sequence-packing gates, same values as the module-top constants above. +# Sequence-packing gates, same values as the module-top constants. RL_PRE_ITEMS["grpo_trainer"].append( "UNSLOTH_GRPO_SEQ_PACKING_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_SEQ_PACKING', '1').lower() not in ('0', 'false', 'no', 'off')\n" ) @@ -1764,6 +2011,16 @@ RL_PRE_ITEMS["grpo_trainer"].append( "except Exception:\n" " UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False\n" ) +# PrefixGrouper gate, same shape as the module-top constants. +RL_PRE_ITEMS["grpo_trainer"].append( + "_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None\n" + "UNSLOTH_GRPO_PREFIX_GROUPER_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_PREFIX_GROUPER', '1').lower() not in ('0', 'false', 'no', 'off')\n" + "if UNSLOTH_GRPO_PREFIX_GROUPER_ON:\n" + " try:\n" + " from unsloth.utils.prefix_grouper import build_group_layout as _pg_build_layout, prefix_grouper_enabled as _pg_enabled_fn, verify_on as _pg_verify_on, tol_ok as _pg_tol_ok, TOL_KILL as _PG_TOL_KILL\n" + " except Exception:\n" + " UNSLOTH_GRPO_PREFIX_GROUPER_ON = False\n" +) # Edit _get_per_token_logps to handle mixed precision diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 2e984bad0a..68fb33dad9 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -17,6 +17,7 @@ from __future__ import annotations +import os from dataclasses import dataclass from typing import Any, Optional, Tuple @@ -42,6 +43,17 @@ if HAS_XFORMERS and torch.cuda.is_available(): HAS_XFORMERS = False SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "") +# PrefixGrouper kernel, resolved once when the env gate is on so PG-off users never load +# torch flex_attention. +_flex_shared_prefix_attention = None +if os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ("0", "false", "no", "off"): + try: + from .prefix_grouper_kernel import ( + flex_shared_prefix_attention as _flex_shared_prefix_attention, + ) + except Exception: + _flex_shared_prefix_attention = None + FLASH_VARLEN = "flash_varlen" FLASH_DENSE = "flash_dense" XFORMERS = "xformers" @@ -84,6 +96,9 @@ class AttentionContext: attention_mask: Optional[Tensor] causal_mask: Optional[Any] sliding_window: Optional[int] = None + # PrefixGrouper: non-None routes Q/K/V through the FlexAttention shared-prefix kernel; + # None leaves every existing construction/behavior unchanged. + prefix_seg_info: Optional[Any] = None def select_attention_backend(use_varlen: bool = False) -> str: @@ -99,6 +114,33 @@ def select_attention_backend(use_varlen: bool = False) -> str: return SDPA +def resolve_prefix_seg_info(kwargs, past_key_value, attention_mask): + """PrefixGrouper shared-prefix segment table resolver for the arch attention forwards. + + The GRPO PrefixGrouper packed path rides a ``PrefixSegInfo`` in through ``**kwargs`` + (same route as ``packed_seq_lengths``). When present, the forward must route Q/K/V + through the FlexAttention shared-prefix kernel via ``AttentionContext.prefix_seg_info``. + + Returns the seg table (or ``None`` when PrefixGrouper did not group this batch -- the + unchanged path). Hardened: the shared-prefix stream is NOT a plain causal sequence, so running + it under a KV cache or an explicit padding mask would silently produce wrong logprobs. + That combination can only arise from misuse (PrefixGrouper only rides in via the GRPO + logprob forward, which is mask-free prefill), so we RAISE loudly instead of degrading + to a wrong result. + + Factored here so every arch (llama/mistral/qwen3/gemma2/cohere/granite/falcon_h1) + shares one implementation and cannot drift. + """ + seg = kwargs.get("prefix_seg_info", None) + if seg is not None and (past_key_value is not None or attention_mask is not None): + raise RuntimeError( + "PrefixGrouper: prefix_seg_info requires prefill with no KV cache and no " + f"attention_mask (got past_key_value={past_key_value is not None}, " + f"attention_mask={attention_mask is not None})." + ) + return seg + + def run_attention( *, config: AttentionConfig, context: AttentionContext, Q: Tensor, K: Tensor, V: Tensor ) -> Tensor: @@ -111,6 +153,28 @@ def run_attention( and SDPA handle packing via a block-diagonal mask. """ + # PrefixGrouper shared-prefix attention (GRPO dedup). Q/K/V here are [bsz, H, T, D]; + # the kernel takes/returns [1, T, H, D], matching the other backends. The field is + # only set when the env gate is on and grouping succeeded; None keeps every backend + # byte-identical. + if context.prefix_seg_info is not None: + flex_shared_prefix_attention = _flex_shared_prefix_attention + if flex_shared_prefix_attention is None: + # gate flipped on after import (or one-time load failed): resolve lazily. + from ..utils.prefix_grouper_kernel import flex_shared_prefix_attention + + scale = None + if config.flash_varlen_kwargs: + scale = config.flash_varlen_kwargs.get("softmax_scale") + A = flex_shared_prefix_attention( + Q.transpose(1, 2), + K.transpose(1, 2), + V.transpose(1, 2), + context.prefix_seg_info, + scale = scale, + ) + return A # [1, T, n_heads, head_dim] + backend = config.backend if backend == FLASH_VARLEN and context.seq_info is None: backend = FLASH_DENSE if HAS_FLASH_ATTENTION else SDPA @@ -337,5 +401,6 @@ __all__ = [ "AttentionConfig", "AttentionContext", "select_attention_backend", + "resolve_prefix_seg_info", "run_attention", ] diff --git a/unsloth/utils/prefix_grouper.py b/unsloth/utils/prefix_grouper.py new file mode 100644 index 0000000000..4e6ff9672c --- /dev/null +++ b/unsloth/utils/prefix_grouper.py @@ -0,0 +1,351 @@ +# 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 . + +"""PrefixGrouper layout builder + completion-logprob extraction for the Unsloth GRPO +packed path (all archs that route through the varlen attention dispatch). + +Given the de-padded, LEFT-PACKED input_ids the packed GRPO path already works with, this +module: + + 1. Detects consecutive ``num_generations`` rows that share a prompt prefix (byte- + identical prompt precondition; falls back / returns None otherwise). + 2. Builds ONE flat shared-prefix stream across all groups + ``[ prefix_g0, suf_g0_0 .. suf_g0_{G-1}, prefix_g1, ... ]`` with position_ids that + continue each prefix positionally, plus a ``PrefixSegInfo`` segment table for the + FlexAttention shared-prefix kernel. + 3. Extracts completion logprobs via the index map (completion pos ``j==0`` predicted + from the shared prefix's last token; ``j>=1`` from the preceding suffix token) and + scatters them back into ``[total_rows, W]`` EXACTLY where the full-row packed path + puts them (dest = ``orig_row*L + orig_col``), so grpo_compute_loss / completion_mask + / TIS / metrics are byte-untouched. + +The flat stream is built by GATHERING original (row, col) coordinates out of input_ids, +so the grad path's autograd flows to the same embedding rows as today (the shared prefix +now contributes grad once = the sum of the G repeats, which is mathematically identical). + +``chunked_hidden_states_selective_log_softmax`` (from unsloth_zoo, passed in) is reused +verbatim over the gathered predicting-position hidden states, so fp32 accumulation, +logit_scale/softcapping/temperature are all preserved. + +Env: + UNSLOTH_GRPO_PREFIX_GROUPER=1 engage (default ON; set 0 to disable). Auto-off under vLLM. + UNSLOTH_GRPO_PREFIX_GROUPER_TOKR=1.3 tok_r auto-gate threshold (env-overridable) + UNSLOTH_GRPO_PREFIX_GROUPER_VERIFY=1 first-step self-verify (default ON) + UNSLOTH_GRPO_PREFIX_GROUPER_TOL=0.7 self-verify PASS band (nats) +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch + +from .prefix_grouper_kernel import build_seg_info_multigroup, PrefixSegInfo + + +# --------------------------------------------------------------------------- +# Env helpers +# --------------------------------------------------------------------------- +def env_on(name: str, default: str = "0") -> bool: + return os.environ.get(name, default).lower() not in ("0", "false", "no", "off") + + +# One-time env reads; the helpers stay callable since unsloth_zoo imports and calls them. +_ENABLED = env_on("UNSLOTH_GRPO_SEQ_PACKING", "1") and env_on("UNSLOTH_GRPO_PREFIX_GROUPER", "1") +_VERIFY_ON = env_on("UNSLOTH_GRPO_PREFIX_GROUPER_VERIFY", "1") +_TOKR_THRESHOLD = float(os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER_TOKR", "1.3")) +_TOL_OK = float(os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER_TOL", "0.7")) + + +def prefix_grouper_enabled() -> bool: + """PrefixGrouper requires seq-packing on (it reuses its de-pad + scatter machinery).""" + return _ENABLED + + +def verify_on() -> bool: + return _VERIFY_ON + + +def tokr_threshold() -> float: + return _TOKR_THRESHOLD + + +def tol_ok() -> float: + return _TOL_OK + + +# diff >= TOL_KILL = broken mask/isolation -> structure permanently unsafe; between +# tol_ok and TOL_KILL -> fall back for this shape but keep trying others. +TOL_KILL = 1.5 + + +@dataclass +class GroupLayout: + """Everything the GRPO forward needs to run + extract the shared-prefix path.""" + + flat_ids: torch.Tensor # [1, T] (T == seg.T) + position_ids: torch.Tensor # [1, T] + prefix_seg_info: PrefixSegInfo + # per completion target token, aligned 1:1: + tgt_rows: torch.Tensor # [N] original row index + tgt_cols: torch.Tensor # [N] original padded column in that row + tgt_pred: torch.Tensor # [N] flat predicting index (into the T stream) + tgt_flat: torch.Tensor # [N] flat index of the target token itself (into T) + total_rows: int + L: int # original padded seq length (input_ids.shape[1]) + W: int # logits_to_keep + max_left_pad (scatter width) + tok_r: float + signature: Tuple + + def extract_logps( + self, + hidden, + lm_head, + chunked_fn, + chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + ) -> torch.Tensor: + """hidden: [1, T, Hdim] (pre-lm_head hidden states, UNSLOTH_RETURN_HIDDEN_STATES=1). + Returns [total_rows, W] float32, byte-compatible with the packed path result.""" + # In a sharded model hidden may live on the lm-head device; move the small index + # maps to hidden.device before indexing. + device = hidden.device + pred_h = hidden[0, self.tgt_pred.to(device), :].unsqueeze(0) # [1, N, Hdim] + tgt_ids = self.flat_ids[0, self.tgt_flat].to(device).unsqueeze(0) # [1, N] + sel = chunked_fn( + pred_h, + lm_head, + tgt_ids, + chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + )[0] # [N] logprobs + dest = self.tgt_rows.to(device) * self.L + self.tgt_cols.to(device) + result = ( + torch.zeros(self.total_rows * self.L, dtype = torch.float32, device = device) + .index_put((dest,), sel.to(torch.float32)) + .view(self.total_rows, self.L)[:, -self.W :] + ) + return result + + +def _build_groups(ids_cpu, real_cols_cpu, cstart_cpu, num_generations, total_rows): + """CPU-side grouping. Returns group dicts or None. Mirrors the packed _pk_* partition. + + A row's REAL tokens are the columns where input != pad. Its completion region (what + the packed path scatters, then completion_mask masks) is the real columns with + original col >= cstart_r, where cstart_r = (L - logits_to_keep) - left_pad_r. The + prompt is the real columns < cstart_r. Within a GRPO group all G rows share the same + prompt => same left_pad => same cstart => the prompt real columns are BYTE-IDENTICAL + across the group (the shared prefix). We require that byte-identity (falls back + otherwise). No prompt-tail special-casing: every suffix token is scattered exactly + like the packed path; completion_mask masks the leading prompt-tail positions. + """ + G = num_generations + if G is None or G < 2 or total_rows % G != 0: + return None + groups = [] + for g0 in range(0, total_rows, G): + rows = list(range(g0, g0 + G)) + prompt_cols_per_row = [] # real cols < cstart + prompt_toks_per_row = [] + comp_cols_per_row = [] # real cols >= cstart (the completion region packed scatters) + for r in rows: + cs = cstart_cpu[r] + rc = real_cols_cpu[r] + p_cols = [c for c in rc if c < cs] + c_cols = [c for c in rc if c >= cs] + prompt_cols_per_row.append(p_cols) + prompt_toks_per_row.append([ids_cpu[r][c] for c in p_cols]) + comp_cols_per_row.append(c_cols) + if any(len(p) == 0 for p in prompt_toks_per_row): + return None + # require BYTE-IDENTICAL prompts across the group (shared-prefix precondition). + P = len(prompt_toks_per_row[0]) + if any(len(prompt_toks_per_row[k]) != P for k in range(1, G)): + return None + p0 = prompt_toks_per_row[0] + if any(prompt_toks_per_row[k] != p0 for k in range(1, G)): + return None + if P == 0: + return None + R_list = [len(c) for c in comp_cols_per_row] + if sum(R_list) == 0: + return None + groups.append( + dict( + rows = rows, + P = P, + prefix_cols = prompt_cols_per_row[0], # shared prompt real columns (row0) + prefix_row = rows[0], + R_list = R_list, + suf_cols = comp_cols_per_row, # per-row completion-region real columns + ) + ) + return groups + + +def _tok_r(groups) -> float: + tok_full = 0 + tok_sp = 0 + for gm in groups: + P = gm["P"] + Rs = gm["R_list"] + tok_full += sum(P + r for r in Rs) # G*P + sumR + tok_sp += P + sum(Rs) # P + sumR + return (tok_full / tok_sp) if tok_sp else 1.0 + + +def build_group_layout( + input_ids, + logits_to_keep, + pad_id, + num_generations, + left_pad_tokens_per_prompt, + *, + apply_tokr_gate = True, + max_segment_cap = None, +): + """Build the shared-prefix GroupLayout, or return None to fall back to the packed path. + + input_ids : [B, L]. GRPO's layout is left-padded in the prompt and right-padded in + the completion. Real tokens of a row are a contiguous run not necessarily + starting at column 0. + logits_to_keep : int + left_pad_tokens_per_prompt : [B] long tensor (per-row left-pad count in the prompt). + """ + device = input_ids.device + total_rows, L = input_ids.shape + keep = input_ids != pad_id + # completion start column per row (matches create_completion_attention_mask / _pk_cstart). + cstart = ((L - logits_to_keep) - left_pad_tokens_per_prompt).to(torch.long) + cstart_cpu = cstart.tolist() + ids_cpu = input_ids.tolist() + # per-row real (non-pad) columns. GRPO rows are one contiguous real run, so derive + # [first, first+n) on GPU; the O(B*L) scan is only a non-contiguous fallback. + n_real = keep.sum(dim = 1) + first = torch.argmax(keep.to(torch.int8), dim = 1) + ar = torch.arange(L, device = device) + contiguous = bool( + (keep == ((ar >= first.unsqueeze(1)) & (ar < (first + n_real).unsqueeze(1)))).all() + ) + if contiguous: + real_cols_cpu = [list(range(f, f + n)) for f, n in zip(first.tolist(), n_real.tolist())] + else: + keep_cpu = keep.tolist() + real_cols_cpu = [[c for c in range(L) if keep_cpu[r][c]] for r in range(total_rows)] + + groups = _build_groups(ids_cpu, real_cols_cpu, cstart_cpu, num_generations, total_rows) + if groups is None: + return None + + # sliding-window guard: a group's PG span is P + max(R); fall back if it exceeds the window. + if max_segment_cap is not None: + for gm in groups: + if gm["P"] + max(gm["R_list"]) > max_segment_cap: + return None + + tok_r = _tok_r(groups) + if apply_tokr_gate and tok_r < tokr_threshold(): + return None # low reuse -> not worth it; use the full-row packed path + + # Build flat stream by gathering original (row, col) coordinates. + group_specs = [(gm["P"], gm["R_list"]) for gm in groups] + seg, group_meta = build_seg_info_multigroup(group_specs, device) + + flat_src_rows: List[int] = [] + flat_src_cols: List[int] = [] + pos_list: List[int] = [] + tgt_rows: List[int] = [] + tgt_cols: List[int] = [] + tgt_pred: List[int] = [] + tgt_flat: List[int] = [] + + for gm, meta in zip(groups, group_meta): + rows = gm["rows"] + P = gm["P"] + r0 = gm["prefix_row"] + prefix_cols = gm["prefix_cols"] # ORIGINAL real prompt columns (len P) of row0 + plast = meta["prefix_last_index"] # base + P - 1 + # gather the shared prefix once, from row0. + flat_src_rows.extend([r0] * P) + flat_src_cols.extend(prefix_cols) + pos_list.extend(range(P)) + # suffixes: every suffix token is a completion-region target (scattered like the + # packed path; completion_mask hides prompt-tail positions). + for i, r in enumerate(rows): + cols = gm["suf_cols"][i] + r_i = len(cols) + s, e = meta["suffix_slices"][i] # flat offsets [s, e) + flat_src_rows.extend([r] * r_i) + flat_src_cols.extend(cols) + pos_list.extend(range(P, P + r_i)) + for j in range(r_i): + # pos 0 is predicted from the prefix's last token; j>=1 from the previous suffix token. + pred = plast if j == 0 else (s + j - 1) + tgt_rows.append(r) + tgt_cols.append(cols[j]) # ORIGINAL padded column in row r + tgt_pred.append(pred) + tgt_flat.append(s + j) # flat index of the target token itself + + T = len(flat_src_rows) + assert T == seg.T, f"flat stream len {T} != seg.T {seg.T}" + fr = torch.tensor(flat_src_rows, device = device, dtype = torch.long) + fc = torch.tensor(flat_src_cols, device = device, dtype = torch.long) + flat_ids = input_ids[fr, fc].unsqueeze(0) # [1, T] (grad-safe gather) + position_ids = torch.tensor(pos_list, device = device, dtype = torch.long).unsqueeze(0) + + max_left_pad = int(left_pad_tokens_per_prompt.max().item()) if total_rows else 0 + W = logits_to_keep + max_left_pad + + # self-verify cache key: the mask/index-map/scatter logic is structural, so key on + # (num_groups, group_sizes), not exact lengths -- GRPO lengths change every step and + # keying on T would re-verify forever ("verify once, then trust", like the packed path). + grp_sizes = tuple(sorted(len(gm["R_list"]) for gm in groups)) + sig = (len(groups), grp_sizes) + + return GroupLayout( + flat_ids = flat_ids, + position_ids = position_ids, + prefix_seg_info = seg, + tgt_rows = torch.tensor(tgt_rows, device = device, dtype = torch.long), + tgt_cols = torch.tensor(tgt_cols, device = device, dtype = torch.long), + tgt_pred = torch.tensor(tgt_pred, device = device, dtype = torch.long), + tgt_flat = torch.tensor(tgt_flat, device = device, dtype = torch.long), + total_rows = total_rows, + L = L, + W = W, + tok_r = tok_r, + signature = sig, + ) + + +__all__ = [ + "GroupLayout", + "build_group_layout", + "prefix_grouper_enabled", + "verify_on", + "tokr_threshold", + "tol_ok", + "TOL_KILL", + "env_on", +] diff --git a/unsloth/utils/prefix_grouper_kernel.py b/unsloth/utils/prefix_grouper_kernel.py new file mode 100644 index 0000000000..9a9719b015 --- /dev/null +++ b/unsloth/utils/prefix_grouper_kernel.py @@ -0,0 +1,436 @@ +# 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", +]