diff --git a/docs/04-technical-guides/performance-tuning.md b/docs/04-technical-guides/performance-tuning.md index 7f62be639..22c3ad43a 100644 --- a/docs/04-technical-guides/performance-tuning.md +++ b/docs/04-technical-guides/performance-tuning.md @@ -93,6 +93,7 @@ Defaults in `primus/configs/modules/megatron/primus_turbo.yaml` are mostly `fals | Flag | Purpose | |------|---------| | `use_turbo_attention` | Optimized attention kernels. | +| `use_turbo_flex_attention` | Route `use_turbo_attention` through the Primus-Turbo `flex_attention` compat layer (see below). | | `use_turbo_parallel_linear` | Optimized tensor-parallel linear layers. | | `use_turbo_grouped_gemm` | Optimized grouped GEMM for MoE. | | `use_turbo_grouped_mlp` | Removed—use `use_turbo_grouped_gemm` (passing this key now raises an error). | @@ -103,6 +104,59 @@ Defaults in `primus/configs/modules/megatron/primus_turbo.yaml` are mostly `fals | `turbo_sync_free_moe_stage` | Sync-free MoE stages (`0`–`3`; `0` disables, stage `2` recommended for performance per patch notes). See [MoE training deep-dive](./moe-training.md). | | `use_turbo_fused_act_with_probs` | Fused activation with probabilities to reduce redundant work. | +#### Flex attention routing + +`use_turbo_flex_attention: true` (default `false`, requires `use_turbo_attention: true`) +makes the Turbo attention module call the `flex_attention` compat layer instead of +`flash_attn_func` directly. The compat layer takes a torch-FlexAttention +`block_mask` / `score_mod`, classifies it, and dispatches the recognized patterns onto +the same Turbo kernels -- so with the default causal configuration the result is +numerically identical to the direct call and the kernel is the same. What it adds is +the ability to express masks and score modifications that `flash_attn_func` has no +argument for: + +```yaml +primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + use_turbo_flex_attention: true + # optional, "package.module:attribute" paths to your own callables + turbo_flex_attention_mask_mod: null # (b, h, q_idx, kv_idx) -> bool + turbo_flex_attention_score_mod: null # (score, b, h, q_idx, kv_idx) -> score +``` + +Unsupported combinations raise at model-build time rather than silently reverting to +the direct call: `context_parallel_size > 1`, `enable_turbo_attention_float8: true`, +`deterministic_mode: true`, and `reset_attention_mask: true`. A sliding window without a +causal mask is rejected on the call itself. Nothing degrades quietly -- a rejected run +fails at startup instead of training on the old code path. + +#### Packed sequences (`qkv_format: thd`) + +This switch is currently the **only** way to run packed sequences on the Turbo attention +path. `PrimusTurboAttention` forwards `packed_seq_params.cu_seqlens_q/cu_seqlens_kv` to +`flex_attention_varlen`, which applies document-internal causal masking (block-diagonal +plus within-segment causal). The direct `flash_attn_func` binding has no `cu_seqlens` +argument, so with `use_turbo_flex_attention: false` a `thd` batch now raises instead of +attending across document boundaries with nothing reported. + +The packed path takes no `BlockMask`: Megatron already knows where the boundaries are, so +there is nothing to probe or classify, and the cache described below is not touched at +all. `turbo_flex_attention_mask_mod`, `turbo_flex_attention_score_mod` and an attention +bias are all refused on this path -- `flex_attention_varlen` takes `cu_seqlens` and +explicit `alibi_slopes`, not a programmable mask, and dropping a supplied callable +silently is exactly the failure mode this layer exists to prevent. + +`reset_attention_mask: true` is a different mechanism and is still refused: it encodes +per-document boundaries inside a *dense* sample's attention mask, and the Turbo call +signature has no argument to carry them. Use `qkv_format: thd` packing instead. + +Sequence-length variety costs more here than on the direct path: each distinct +`(mask pattern, q_len, kv_len, device)` builds a `BlockMask` and re-runs the compat +layer's classification once, then is cached (bounded LRU, 64 entries). Fixed-length +pretraining pays this once per run; a workload that hands attention a new length every +step will warn once and then rebuild on every miss. Bucket the lengths. + ### Feature flags (TorchTitan) TorchTitan presets include `primus_turbo` in `primus/configs/modules/torchtitan/pre_trainer.yaml` diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index d7d18b484..bf65e5160 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -99,6 +99,11 @@ from primus.backends.megatron.core.extensions._triton.inplace_add import ( inplace_add_triton_, ) +from primus.backends.megatron.core.extensions.turbo_flex_attention import ( + build_turbo_flex_attention, + reject_reset_attention_mask, + sink_kwargs_for, +) from primus.core.utils.module_utils import warning_rank_0 _dummy_wgrads = {} @@ -818,7 +823,24 @@ def __init__( self._num_heads_for_sinks = self.config.num_attention_heads self.offload = args.offload and "attn" in args.offload_ops - if args.enable_turbo_attention_float8: + self._attn_accepts_packed_seq = False + # The packed (THD) entry bound alongside the dense one. ``None`` means this + # configuration has no varlen kernel and qkv_format="thd" must be rejected. + self.attn_varlen = None + self._attn_is_fp8 = bool(getattr(args, "enable_turbo_attention_float8", False)) + self._attn_is_flex = bool(getattr(args, "use_turbo_flex_attention", False)) + if self._attn_is_flex: + # Route through the flex_attention compat layer. It presents the same call + # signature as flash_attn_func and dispatches onto the same Turbo kernels, + # so forward() below needs no changes. Unsupported combinations raise here, + # at model-build time, rather than silently reverting to the direct call. + self.attn = build_turbo_flex_attention(args=args, config=self.config) + # Only the flex path knows what to do with packed (THD) sequences: it can + # forward cu_seqlens to flex_attention_varlen. The direct flash_attn_func + # binding below takes no cu_seqlens argument, so the extra kwargs are gated + # on this flag and the non-flex path is left byte-for-byte unchanged. + self._attn_accepts_packed_seq = True + elif args.enable_turbo_attention_float8: self.attn = ( primus_turbo_torch.ops.flash_attn_fp8_usp_func if self.config.context_parallel_size > 1 @@ -830,6 +852,25 @@ def __init__( if self.config.context_parallel_size > 1 else primus_turbo_torch.ops.flash_attn_func ) + # Packed sequences on the direct path. flash_attn_varlen_func takes the + # document boundaries as cu_seqlens and honours them in the kernel, so THD + # no longer has to route through the flex compat layer to be correct. + # + # Deliberately bound for CP == 1 only. flash_attn_varlen_usp_func exists and + # its side of the contract is not in doubt: measured on 2x MI355 it is + # bit-identical to the single-card varlen kernel when given a local token + # shard plus the GLOBAL cu_seqlens/max_seqlen, including when a document + # straddles the shard boundary (TASK_PROGRESS EXP21 case C). + # + # What is unverifiable here is the other side: whether what Megatron puts in + # PackedSeqParams under context parallelism is those global boundaries or + # already-sharded local ones. The adapter cannot tell the two apart from the + # tensor it receives -- both are a plausible cu_seqlens -- and guessing wrong + # attends across documents silently, the exact failure this binding exists to + # remove. So CP > 1 keeps raising below until that contract is pinned down. + if self.config.context_parallel_size == 1: + self.attn_varlen = primus_turbo_torch.ops.flash_attn_varlen_func + self._attn_accepts_packed_seq = True if pg_collection is None: # For backward compatibility, remove in v0.14 and raise error # raise ValueError("TEDotProductAttention was called without ProcessGroupCollection") @@ -853,7 +894,53 @@ def __init__( # enable ring attention self.attn_kwargs["ring_group"] = dist.new_group(ranks=[dist.get_rank()]) - assert config.window_size is None, "primus_turbo does not support sliding window attention" + # reset_attention_mask has nowhere to go on this path -- forward() below reads + # no mask tensor and none of the bound kernels take one. Reject it rather than + # let tokens attend across documents silently. See the helper for the full + # reasoning, including why the *default* causal mask is safe to drop. + reject_reset_attention_mask(args, is_flex=self._attn_is_flex, where="primus_turbo attention") + + # Sliding-window attention. + # + # forward() already threads a `window_size` tuple into every attention call -- + # that is how the sink pattern applies its per-layer window -- and the bound + # entries honour it: flash_attn_func / flash_attn_varlen_func pass it to the + # aiter kernel, flash_attn_usp_func saves it on ctx and applies it in the + # backward too. So config-level SWA needs no new plumbing on the direct path; + # it only ever needed the gate opening. The flex path expresses the same window + # as a mask_mod and drives the same kernel. + # + # fp8 is the one exception: flash_attn_fp8_func routes to the Triton kernel, + # whose forward asserts window_size == (-1, -1). Rejecting it here turns that + # assert into an explanation at build time. + self._config_window_size = None + if config.window_size is not None: + if self._attn_is_fp8: + raise NotImplementedError( + "primus_turbo does not support sliding window attention on the fp8 " + f"attention path (config.window_size={tuple(config.window_size)}); the " + "Triton fp8 kernel takes full attention only. Set " + "enable_turbo_attention_float8=false, or window_size=None." + ) + left, right = config.window_size + if int(right) > 0: + # A positive right bound means attending forward. The flex path builds + # its window out of causal mask_mods and cannot express it at all; the + # direct path could pass it to the kernel, but honouring it on one path + # and not the other would make the same config mean two different + # models. Reject it on both. + raise NotImplementedError( + "primus_turbo supports causal windows only, so window_size[1] must be 0 " + f"or -1; got {tuple(config.window_size)}." + ) + if self._init_sink_attention and self.sink_sliding_window > 0: + raise NotImplementedError( + "primus_turbo: config.window_size and sink_sliding_window both request a " + f"sliding window ({tuple(config.window_size)} vs {self.sink_sliding_window}). " + "The sink window is applied per layer, so combining the two would make the " + "effective window depend on the layer index. Set only one of them." + ) + self._config_window_size = (int(left), 0) # Check version kv_channels = ( @@ -904,14 +991,48 @@ def forward( attention_bias: Tensor = None, packed_seq_params: PackedSeqParams = None, ): - """Forward.""" - packed_seq_kwargs = ( - {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params} - if packed_seq_params is not None - else {} - ) - - qkv_format = packed_seq_kwargs.get("qkv_format", self.qkv_format) + """Forward. + + ``attention_mask`` is deliberately unread. Every attention entry point bound in + ``__init__`` (flash_attn_func / its usp + fp8 variants / the flex compat layer) + takes the mask as a ``causal`` flag rather than a tensor, so there is nothing to + forward it to. That is only sound because the sole configuration in which the + tensor carries information the ``causal`` flag does not -- reset_attention_mask, + which zeroes the cross-document blocks -- is rejected at build time in + ``__init__``. With it off, Megatron's mask is exactly ``torch.tril(...)``, which + ``causal=True`` already expresses. Per-document masking is available through + packed sequences (``qkv_format="thd"``), where the boundaries arrive as + cu_seqlens and *are* forwarded. + """ + if attention_bias is not None: + # Neither flash_attn_func nor the flex compat layer is wired to an attention + # bias here (the kernel call below passes alibi_slopes=None unconditionally), + # so accepting the tensor and ignoring it would train a different model than + # the config asks for, with no diagnostic. Fail instead. + raise NotImplementedError( + "primus_turbo attention does not support attention_bias; the tensor would be " + "silently dropped. This also covers ALiBi, which Megatron delivers as an " + "attention bias -- use a position embedding the Turbo path implements." + ) + # Read qkv_format straight off the dataclass. + # + # This used to be a dict comprehension over ``self.kept_packed_seq_params`` -- + # copied from Megatron's TEDotProductAttention, which defines that attribute in + # its own __init__. This class does not inherit from it: it extends + # te.pytorch.DotProductAttention (checked: TE 2.15 does not define the attribute + # either). So every call carrying packed_seq_params raised + # + # AttributeError: 'PrimusTurboAttention' object has no attribute + # 'kept_packed_seq_params' + # + # before reaching any of the packing logic below, and the error said nothing + # about packing. Verified on MI355 (TASK_PROGRESS EXP23 case C). + # + # The dict was only ever read for one key, so there is nothing to reconstruct. + # A qkv_format of None falls back to the module default rather than propagating: + # None matches neither "sbhd" nor "thd" below, so it used to skip the layout + # permute entirely and hand the kernel a transposed tensor. + qkv_format = getattr(packed_seq_params, "qkv_format", None) or self.qkv_format mask_type = attn_mask_type.name if mask_type == AttnMaskType.causal.name: causal = True @@ -931,7 +1052,10 @@ def forward( # # Reference: gpt-oss/gpt_oss/triton/attention.py sink_tensor = None - window_size = (-1, -1) + # config.window_size, when set, applies to every layer. The sink pattern below + # is the other source of a window and __init__ rejects having both, so whichever + # one is configured is the only writer of this variable. + window_size = self._config_window_size or (-1, -1) use_sink_attn = self.use_sink_attention and self.sinks is not None @@ -972,28 +1096,169 @@ def forward( key = key.permute(0, 2, 1, 3) value = value.permute(0, 2, 1, 3) - o = self.attn( - query, - key, - value, - dropout_p=0.0, - softmax_scale=self.softmax_scale, - causal=causal, - window_size=window_size, - bias=None, - alibi_slopes=None, - deterministic=self.deterministic_mode, - return_lse=False, - return_attn_probs=False, - sink=sink_tensor, # PR 208: pass sink tensor to Primus-Turbo - **self.attn_kwargs, - ) + # Packed sequences: q/k/v are already THD [total_tokens, H, D] (no permute + # applies) and the document boundaries live in packed_seq_params. They reach the + # kernel two different ways: + # * flex -> forwarded as keyword args to the compat layer, which calls + # flex_attention_varlen itself; + # * direct-> dispatched to the flash_attn_varlen_func bound in __init__, whose + # signature takes cu_seqlens positionally. + # Either way the boundaries reach the kernel. Without this, "thd" used to reach + # a dense attention call with the boundaries dropped, so tokens attended across + # documents with no error raised anywhere. + packed_attn_kwargs = {} + use_direct_varlen = False + cu_seqlens_q = cu_seqlens_kv = None + max_seqlen_q = max_seqlen_kv = None + if qkv_format == "thd": + if not self._attn_accepts_packed_seq: + raise NotImplementedError( + "PrimusTurboAttention: qkv_format='thd' (packed sequences) is not " + "supported on this attention path -- there is no varlen kernel bound " + "for it, so the document boundaries in packed_seq_params could not be " + "forwarded and tokens would attend across documents. Packing is " + "available on the plain bf16/fp16 path (context_parallel_size=1) and " + "through use_turbo_flex_attention=true; it is not available with " + "enable_turbo_attention_float8=true (no fp8 varlen entry) or with " + "context parallelism. Disable sequence packing, or change the path." + ) + if packed_seq_params is None: + raise ValueError( + "PrimusTurboAttention: qkv_format='thd' requires packed_seq_params " + "carrying cu_seqlens_q / cu_seqlens_kv." + ) + cu_seqlens_q = getattr(packed_seq_params, "cu_seqlens_q", None) + cu_seqlens_kv = getattr(packed_seq_params, "cu_seqlens_kv", None) + max_seqlen_q = getattr(packed_seq_params, "max_seqlen_q", None) + max_seqlen_kv = getattr(packed_seq_params, "max_seqlen_kv", None) + if self._attn_is_flex: + packed_attn_kwargs = { + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_kv": cu_seqlens_kv, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_kv": max_seqlen_kv, + } + else: + use_direct_varlen = True + + if use_direct_varlen: + if self.attn_kwargs: + # self.attn_kwargs carries the context-parallel process groups + # (ulysses_group / ring_group) and is only populated when cp_size > 1, + # while self.attn_varlen is only bound when cp_size == 1 -- so today the + # two are mutually exclusive and this branch never drops anything. That + # is an invariant of two separate conditions in __init__, not something + # the call site can see. Assert it here so that relaxing either gate + # later fails loudly instead of silently running context parallelism + # without its groups, which would compute attention over the local shard + # only and quietly train a different model. + raise NotImplementedError( + "PrimusTurboAttention: the varlen (THD) entry does not take " + f"{sorted(self.attn_kwargs)}; packed sequences under context " + "parallelism are not wired up yet." + ) + if cu_seqlens_q is None: + raise ValueError( + "PrimusTurboAttention: qkv_format='thd' requires cu_seqlens_q; the " + "document boundaries are not recoverable from the packed tensor shape, " + "and guessing them would let tokens attend across documents." + ) + # Padding-aware boundaries, if Megatron supplied them, describe a different + # packing than cu_seqlens_q does: cu_seqlens_* are the real token counts + # while cu_seqlens_*_padded include the per-sequence padding. TE takes both + # and honours the distinction; the varlen entry bound here takes one pair + # and would apply the unpadded boundaries to a padded tensor, attending over + # padding and misplacing every document after the first. Refuse instead. + for _name, _padded, _plain in ( + ( + "cu_seqlens_q_padded", + getattr(packed_seq_params, "cu_seqlens_q_padded", None), + cu_seqlens_q, + ), + ( + "cu_seqlens_kv_padded", + getattr(packed_seq_params, "cu_seqlens_kv_padded", None), + cu_seqlens_kv if cu_seqlens_kv is not None else cu_seqlens_q, + ), + ): + if _padded is not None and not torch.equal(_padded, _plain): + raise NotImplementedError( + f"PrimusTurboAttention: packed_seq_params carries {_name}, which differs " + "from the unpadded boundaries, but the varlen entry bound here takes a " + "single cu_seqlens pair. Using the unpadded boundaries on a padded tensor " + "would attend over the padding and misplace every document after the " + "first. Disable THD padding (cu_seqlens_*_padded == cu_seqlens_*), or use " + "a backend that takes both." + ) + if cu_seqlens_kv is None: + cu_seqlens_kv = cu_seqlens_q + if max_seqlen_q is None or max_seqlen_kv is None: + # Derivable from the prefix sums; one small D2H copy, and in practice only + # on the first call of a shape because Megatron caches PackedSeqParams. + max_seqlen_q = int(max_seqlen_q or (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item()) + max_seqlen_kv = int(max_seqlen_kv or (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]).max().item()) + # Not every bound entry takes a sink; see sink_kwargs_for for the three cases. + varlen_sink_kwargs = sink_kwargs_for( + self.attn_varlen, sink_tensor, where="packed sequences, qkv_format='thd'" + ) + + o = self.attn_varlen( + query, + key, + value, + cu_seqlens_q, + cu_seqlens_kv, + int(max_seqlen_q), + int(max_seqlen_kv), + dropout_p=0.0, + softmax_scale=self.softmax_scale, + causal=causal, + window_size=window_size, + bias=None, + alibi_slopes=None, + deterministic=self.deterministic_mode, + return_lse=False, + return_attn_probs=False, + **varlen_sink_kwargs, + ) + else: + # PR 208 passes the sink tensor to Primus-Turbo -- but only the entries that + # have somewhere to put it. flash_attn_usp_func / flash_attn_fp8_*_func do + # not, and passing it to them raised a bare TypeError on every CP run. + sink_kwargs = sink_kwargs_for( + self.attn, sink_tensor, where=f"context_parallel_size={self.config.context_parallel_size}" + ) + + o = self.attn( + query, + key, + value, + dropout_p=0.0, + softmax_scale=self.softmax_scale, + causal=causal, + window_size=window_size, + bias=None, + alibi_slopes=None, + deterministic=self.deterministic_mode, + return_lse=False, + return_attn_probs=False, + **sink_kwargs, + **packed_attn_kwargs, + **self.attn_kwargs, + ) if qkv_format == "sbhd": o = o.permute(1, 0, 2, 3) elif qkv_format == "bhsd": o = o.permute(0, 2, 1, 3) + if qkv_format == "thd": + # THD output is [total_tokens, H, D]; the caller expects the heads folded + # into one feature dim, i.e. [total_tokens, H * D]. The generic 3-arg view + # below is a no-op on a 3D tensor, which would silently hand back the wrong + # rank, so the packed case gets its own reshape. + return o.reshape(o.shape[0], -1) + o = o.view(o.shape[0], o.shape[1], -1) return o diff --git a/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py b/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py index 0b903ef14..6e159f6ab 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py +++ b/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py @@ -91,6 +91,11 @@ def _advance_with(g): from megatron.training.global_vars import get_args from torch import Tensor +from primus.backends.megatron.core.extensions.turbo_flex_attention import ( + build_turbo_flex_attention, + reject_reset_attention_mask, +) + class PrimusTurboLocalAttention(MegatronModule): """ @@ -145,7 +150,12 @@ def __init__( # Select Primus Turbo flash attention variant args = get_args() - if args.enable_turbo_attention_float8: + if getattr(args, "use_turbo_flex_attention", False): + # Same swap as PrimusTurboAttention: the compat layer mirrors + # flash_attn_func's call signature, so forward() is unchanged. Wired here + # too so the switch is never silently ignored under the local specs. + self.attn_func = build_turbo_flex_attention(args=args, config=config) + elif args.enable_turbo_attention_float8: self.attn_func = ( pt.ops.flash_attn_fp8_usp_func if config.context_parallel_size > 1 @@ -174,6 +184,15 @@ def __init__( if config.window_size is not None: raise ValueError("PrimusTurboLocalAttention does not support sliding window attention") + # reset_attention_mask has nowhere to go here either: forward() reads no mask + # tensor and none of the bound kernels take one. See the helper for why the + # *default* causal mask is safe to drop but this one is not. + reject_reset_attention_mask( + args, + is_flex=bool(getattr(args, "use_turbo_flex_attention", False)), + where="PrimusTurboLocalAttention", + ) + def forward( self, query: Tensor, @@ -191,14 +210,45 @@ def forward( query: Query tensor [seq_len, batch, num_heads, head_dim] (sbhd) key: Key tensor [seq_len, batch, num_heads, head_dim] (sbhd) value: Value tensor [seq_len, batch, num_heads, head_dim] (sbhd) - attention_mask: Attention mask (not used by flash attention) + attention_mask: unread. The bound kernels take the mask as a ``causal`` + flag, not a tensor. Sound only because reset_attention_mask -- the + one setting under which Megatron's mask says more than ``causal`` + does -- is rejected in ``__init__``; otherwise the tensor is exactly + ``torch.tril(...)``. attn_mask_type: Type of attention mask (causal, no_mask, etc.) - attention_bias: Attention bias (not used in this implementation) - packed_seq_params: Packed sequence parameters (optional) + attention_bias: rejected. Nothing below forwards it -- see the guard. + packed_seq_params: rejected. Nothing below forwards it -- see the guard. Returns: Attention output [seq_len, batch, num_heads * head_dim] (merged heads) """ + if attention_bias is not None: + # The kernel call below passes bias=None and alibi_slopes=None + # unconditionally, and nothing between here and there reads this tensor. + # Accepting it would train a model with no bias while the config says + # there is one, and say nothing about it. Same refusal, same reason, as + # PrimusTurboAttention.forward in primus_turbo.py. + raise NotImplementedError( + "PrimusTurboLocalAttention does not support attention_bias; the tensor " + "would be silently dropped. This also covers ALiBi, which Megatron " + "delivers as an attention bias -- use a position embedding this path " + "implements." + ) + + if packed_seq_params is not None: + # Unlike PrimusTurboAttention, this class has no thd branch at all: it + # always calls the dense kernel with no cu_seqlens. A packed batch would + # therefore be attended to as ONE long sequence, letting every document + # see every earlier document -- wrong numbers, wrong gradients, clean log. + # Refuse rather than approximate. + raise NotImplementedError( + "PrimusTurboLocalAttention does not support packed sequences; the " + "document boundaries in packed_seq_params would be silently dropped and " + "the batch attended to as one contiguous sequence. Use the " + "TransformerEngine-based PrimusTurboAttention spec, which implements " + "qkv_format='thd'." + ) + query, key, value = [x.transpose(0, 1) for x in (query, key, value)] # gfx942: avoid aiter's broken strided-sbhd backward (see __init__). diff --git a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py index 60ab9b9de..bb907e129 100644 --- a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py +++ b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py @@ -112,6 +112,7 @@ def _build_default_primus_args() -> SimpleNamespace: use_turbo_gemm=False, use_turbo_norm_te_linear=False, use_turbo_attention=False, + use_turbo_flex_attention=False, use_turbo_grouped_gemm=False, moe_use_legacy_grouped_gemm=False, ) @@ -222,6 +223,17 @@ def layer_norm(self, rms_norm: bool = False, for_qk: bool = False) -> type: def core_attention(self) -> type: """Which module to use for attention""" + use_flex = getattr(self.cfg, "use_turbo_flex_attention", False) + if use_flex and (self.fallback_to_eager_attn or not self.cfg.use_turbo_attention): + # The flex switch only redirects the callable *inside* PrimusTurboAttention. + # Without that module it is a no-op, and a silently ignored performance/ + # feature switch is worse than a failed startup. + raise ValueError( + "use_turbo_flex_attention=true requires use_turbo_attention=true " + f"(currently use_turbo_attention={self.cfg.use_turbo_attention}, " + f"fallback_to_eager_attn={self.fallback_to_eager_attn}); the flex compat layer " + "is only reachable through the Primus-Turbo attention module." + ) if self.fallback_to_eager_attn: return DotProductAttention return ( diff --git a/primus/backends/megatron/core/extensions/turbo_flex_attention.py b/primus/backends/megatron/core/extensions/turbo_flex_attention.py new file mode 100644 index 000000000..001e4057b --- /dev/null +++ b/primus/backends/megatron/core/extensions/turbo_flex_attention.py @@ -0,0 +1,732 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""Route Megatron core attention through the Primus-Turbo ``flex_attention`` compat layer. + +Enabled by ``use_turbo_flex_attention: true`` in the ``primus_turbo`` module config +(default ``false``). It only takes effect on top of ``use_turbo_attention: true`` -- +it swaps the callable that :class:`PrimusTurboAttention` invokes, so the surrounding +Megatron plumbing (layout permutes, sink parameters, offload, spec selection) is +untouched. + +Why route through the compat layer at all +----------------------------------------- +``flash_attn_func`` takes a fixed set of knobs: ``causal``, ``window_size``, +``alibi_slopes``, ``bias``, ``dropout_p``, ``sink``. The compat layer takes a +``block_mask`` / ``score_mod`` -- torch-FlexAttention's *programmable* interface -- +classifies it, and dispatches the recognised patterns onto exactly those same Turbo +kernels. So with the default causal / no-mask configuration this switch is +**numerically identical** to the direct call (same kernel, same arguments); what it +buys is the ability to express masks and score modifications that have no +``flash_attn_func`` argument, without leaving Turbo's fast path. + +``turbo_flex_attention_mask_mod`` / ``turbo_flex_attention_score_mod`` (both ``null`` +by default) accept a ``"package.module:attribute"`` path to a user callable, which is +how a model config reaches that programmability. + +Deliberately strict +------------------- +Every unsupported combination raises with an explicit message instead of quietly +falling back to the direct call. A silent fallback would make the switch look like it +worked while the run behaved exactly as before -- the failure mode that is hardest to +notice in a training log. +""" + +import importlib +import inspect +import warnings +from collections import OrderedDict +from typing import Any, Callable, Optional, Tuple + +import torch + +# The compat layer ships with Primus-Turbo. Import defensively so that a container +# with an older Turbo build still imports this module (and fails with a clear message +# only if the switch is actually turned on). +_FLEX_IMPORT_ERROR: Optional[BaseException] = None +try: + from primus_turbo.pytorch.ops.attention import ( + create_block_mask as turbo_create_block_mask, + ) + from primus_turbo.pytorch.ops.attention import ( + flex_attention as turbo_flex_attention, + ) +except (ImportError, ModuleNotFoundError, AttributeError) as exc: # pragma: no cover + turbo_create_block_mask = None + turbo_flex_attention = None + _FLEX_IMPORT_ERROR = exc + +# ``flex_attention_bshd`` is the layout-native entry (bshd in, bshd out). Megatron has +# already permuted q/k/v to bshd by the time we are called, so this entry needs no +# transposes at all. Older Turbo builds do not have it; we transparently fall back to +# the torch-layout ``flex_attention`` entry there, handing it a ``transpose(1, 2)`` +# view. What that fallback costs depends on the Turbo build: +# * Turbo with the bhsd passthrough: nothing. The view is non-contiguous, so the +# compat layer restores the caller's own buffer and forwards it unchanged. +# * Older Turbo: 3 input + 1 output copies per forward, plus the backward mirrors. +# Either way the result is correct, so no version probing is needed here. +try: + from primus_turbo.pytorch.ops.attention import ( + flex_attention_bshd as turbo_flex_attention_bshd, + ) +except (ImportError, ModuleNotFoundError, AttributeError): # pragma: no cover + turbo_flex_attention_bshd = None + +# ``flex_attention_varlen`` is the packed-sequence (THD) entry: q/k/v are +# ``[total_tokens, H, D]`` with the document boundaries carried out-of-band in +# ``cu_seqlens``. This is what makes ``qkv_format="thd"`` reachable -- the dense +# entries cannot express per-document masking, and the compat layer's probe-based +# classifier is not asked to: the boundaries are given, not inferred. +try: + from primus_turbo.pytorch.ops.attention import ( + flex_attention_varlen as turbo_flex_attention_varlen, + ) +except (ImportError, ModuleNotFoundError, AttributeError): # pragma: no cover + turbo_flex_attention_varlen = None + + +def _entry_accepts(fn: Optional[Callable], name: str) -> bool: + """Whether ``fn`` has a parameter called ``name``. + + The compat-layer entries gained arguments over time (``deterministic`` is the + latest), and this module has to keep importing against an older Primus-Turbo. The + alternative -- passing the argument unconditionally -- would raise ``TypeError`` + deep inside the call on an old build; the alternative to *that* -- dropping it + silently -- is the exact failure this layer exists to prevent. So we probe, and the + caller decides: forward it when it is accepted, raise when the user asked for a + non-default value the entry cannot take. + """ + if fn is None: + return False + try: + params = inspect.signature(fn).parameters + except (TypeError, ValueError): # pragma: no cover - builtins / C callables + return False + if name in params: + return True + # A ``**kwargs`` catch-all accepts any keyword by definition (this is also what + # the test doubles look like). + return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) + + +def sink_kwargs_for(entry: Optional[Callable], sink: Optional[torch.Tensor], *, where: str) -> dict: + """Decide how to hand ``sink`` to ``entry``: forward it, omit it, or refuse. + + The attention entries Primus binds do not agree about sinks. On current + Primus-Turbo ``flash_attn_func`` takes one; ``flash_attn_usp_func``, + ``flash_attn_fp8_usp_func``, ``flash_attn_fp8_func`` and + ``flash_attn_varlen_func`` do not. The call site used to pass ``sink=`` to all + of them unconditionally, so every context-parallel run died on a bare + ``TypeError: ... unexpected keyword argument 'sink'`` -- whether or not a sink + was configured, and with nothing to say why. + + Three cases, and the middle one is the point: + + * the entry takes it -> forward it (``None`` included; that is its default). + * the entry does not take it, and there is no sink -> omit the argument. Nothing + is lost: there was nothing to pass. + * the entry does not take it, and a sink IS configured -> raise. Omitting it here + would change the softmax denominator of every query, and the run would train a + different model without a word of complaint. + + ``where`` names the path for the error message, since the caller knows which + binding it is holding and this helper does not. + """ + if _entry_accepts(entry, "sink"): + return {"sink": sink} + if sink is None: + return {} + name = getattr(entry, "__name__", repr(entry)) + raise NotImplementedError( + f"PrimusTurboAttention: an attention sink is configured, but the attention entry " + f"bound for this configuration ({name}, {where}) has no 'sink' parameter, so the " + f"sink logits cannot reach the kernel. Dropping them would change the softmax " + f"denominator for every query and silently train a different model. The " + f"context-parallel, fp8 and packed-varlen entries are the ones that lack it: run " + f"sink attention on the unpacked bf16/fp16 path with context_parallel_size=1, or " + f"upgrade Primus-Turbo to a build whose entry forwards a sink." + ) + + +# ============================================================================= +# mask_mod builders + BlockMask cache +# ============================================================================= + + +def _causal_mask_mod(b, h, q_idx, kv_idx): + return q_idx >= kv_idx + + +def _make_window_causal_mask_mod(window: int) -> Callable: + def mask_mod(b, h, q_idx, kv_idx): + return (q_idx >= kv_idx) & ((q_idx - kv_idx) <= window) + + return mask_mod + + +# BlockMask construction walks the whole mask once, and the compat layer memoises its +# classification by *object identity* -- so building one mask per (pattern, shape) and +# reusing it across layers and steps is what keeps the per-call cost at zero. Keyed by +# device string rather than the device object so replicas share entries. +_BLOCK_MASK_CACHE: "OrderedDict[Tuple, Any]" = OrderedDict() + +# Pretraining sees one or two shapes for the whole run, so the cache holds a couple of +# entries forever. Variable-length workloads (SFT with per-batch bucketing, evaluation +# sweeps) hand us a new ``(q_len, kv_len)`` almost every step, and an unbounded dict +# would then grow one BlockMask per distinct length for the lifetime of the process. +# Bounded + LRU: the hot shapes stay resident, the long tail is evicted. +_BLOCK_MASK_CACHE_MAX = 64 +_EVICTION_WARNED = False + + +def _get_block_mask( + *, + mask_mod: Callable, + cache_key: Tuple, + q_len: int, + kv_len: int, + device: torch.device, +): + global _EVICTION_WARNED + + key = (cache_key, q_len, kv_len, str(device)) + cached = _BLOCK_MASK_CACHE.get(key) + if cached is not None: + _BLOCK_MASK_CACHE.move_to_end(key) + return cached + + cached = turbo_create_block_mask( + mask_mod, + B=None, + H=None, + Q_LEN=q_len, + KV_LEN=kv_len, + device=device, + ) + _BLOCK_MASK_CACHE[key] = cached + if len(_BLOCK_MASK_CACHE) > _BLOCK_MASK_CACHE_MAX: + _BLOCK_MASK_CACHE.popitem(last=False) + if not _EVICTION_WARNED: + _EVICTION_WARNED = True + warnings.warn( + "Primus-Turbo flex attention: the BlockMask cache passed " + f"{_BLOCK_MASK_CACHE_MAX} distinct (pattern, q_len, kv_len, device) keys and " + "started evicting. Every miss rebuilds a BlockMask and re-runs the compat " + "layer's pattern classification, which is orders of magnitude more expensive " + "than the attention call itself. Bucket sequence lengths, or raise " + "turbo_flex_attention._BLOCK_MASK_CACHE_MAX if the working set is genuinely " + "this wide.", + RuntimeWarning, + stacklevel=2, + ) + return cached + + +def clear_turbo_flex_block_mask_cache() -> None: + """Drop every cached ``BlockMask`` (tests / benchmarks wanting a cold measurement).""" + global _EVICTION_WARNED + + _BLOCK_MASK_CACHE.clear() + _EVICTION_WARNED = False + + +# ============================================================================= +# Shared build-time rejections +# ============================================================================= + + +def reject_reset_attention_mask(args, *, is_flex: bool, where: str) -> None: + """Reject ``reset_attention_mask`` on any path that cannot honour it. + + ``reset_attention_mask`` asks Megatron for per-document causal masking *inside a + dense sample*: the document boundaries are baked into the ``attention_mask`` tensor + as cross-document zeros. No Primus-Turbo attention entry point takes a mask tensor + -- they take a ``causal`` flag -- so the boundaries have nowhere to go and would be + dropped, letting tokens attend across documents with nothing raised anywhere. + + The default mask is *not* the problem. ``create_attention_mask_in_dataloader`` + defaults to true, so an ordinary causal run also hands the module a tensor; that one + is exactly ``torch.tril(...)``, which the ``causal`` flag already expresses, so + dropping it changes nothing. Only ``reset_attention_mask`` adds information the flag + cannot carry (see ``megatron/core/datasets/gpt_dataset.py``, + ``_get_ltor_masks_and_position_ids``). + + ``qkv_format="thd"`` is a different, supported mechanism: there the boundaries + arrive out-of-band as ``cu_seqlens`` in ``packed_seq_params`` and are forwarded + explicitly to ``flex_attention_varlen``. + + Args: + args: the Megatron global args namespace. + is_flex: whether the caller routes through the flex compat layer, which raises + for itself in :func:`build_turbo_flex_attention` with a message naming the + THD alternative. Passing True here makes this a no-op so that the more + specific message wins. + where: class name to name in the error message. + """ + if not getattr(args, "reset_attention_mask", False) or is_flex: + return + raise NotImplementedError( + f"{where} does not support reset_attention_mask=true. The per-document " + "boundaries are carried in the attention_mask tensor, which flash_attn_func has " + "no argument for, so they would be silently dropped and tokens would attend " + "across documents. Use packed sequences (qkv_format='thd'), or set " + "reset_attention_mask=false." + ) + + +# ============================================================================= +# "package.module:attribute" resolution for the optional user hooks +# ============================================================================= + + +def resolve_dotted_callable(path: str, *, what: str) -> Callable: + """Import ``"package.module:attribute"`` and return the callable it names.""" + if not isinstance(path, str) or ":" not in path: + raise ValueError( + f"Primus-Turbo flex attention: {what} must be a 'package.module:attribute' string, " + f"got {path!r}." + ) + module_name, _, attr = path.partition(":") + try: + module = importlib.import_module(module_name) + except ImportError as exc: + raise ImportError( + f"Primus-Turbo flex attention: could not import module '{module_name}' for {what}." + ) from exc + try: + obj = getattr(module, attr) + except AttributeError as exc: + raise AttributeError( + f"Primus-Turbo flex attention: module '{module_name}' has no attribute '{attr}' " f"for {what}." + ) from exc + if not callable(obj): + raise TypeError(f"Primus-Turbo flex attention: {what} '{path}' is not callable.") + return obj + + +# ============================================================================= +# The drop-in replacement for flash_attn_func +# ============================================================================= + + +def _coerce_sink(sink: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Present Megatron's sink parameter in the dtype the compat layer requires. + + ``PrimusTurboAttention`` allocates ``self.sinks`` as a **bfloat16** Parameter (it + mirrors gpt-oss, and ``flash_attn_func`` takes it as-is), while the compat layer + validates ``sink.dtype == torch.float32`` and rejects anything else. Handing the + Parameter straight through therefore turns on ``use_sink_attention`` + + ``use_turbo_flex_attention`` into a hard ValueError on the first forward. + + ``.float()`` is a differentiable cast: autograd casts the fp32 grad back to bf16 on + the way out, so the Parameter still trains. Doing it here rather than loosening the + compat layer keeps the dtype contract explicit and confines the Megatron-specific + convention to the Megatron-side adapter. + """ + if sink is None or sink.dtype == torch.float32: + return sink + return sink.float() + + +class TurboFlexAttention: + """Callable with ``flash_attn_func``'s signature that dispatches via the compat layer. + + :class:`PrimusTurboAttention` keeps a bound ``self.attn`` and calls it with bshd + ``[B, S, H, D]`` tensors plus keyword arguments. Presenting the same interface here + means the switch is a one-line substitution in ``__init__`` and ``forward`` needs no + changes at all. + """ + + def __init__( + self, + *, + mask_mod: Optional[Callable] = None, + score_mod: Optional[Callable] = None, + mask_mod_key: str = "user", + ): + self.mask_mod = mask_mod + self.score_mod = score_mod + self.mask_mod_key = mask_mod_key + self._entry = turbo_flex_attention_bshd or turbo_flex_attention + self._entry_is_bshd = turbo_flex_attention_bshd is not None + # ``deterministic`` reached the compat layer later than the rest of the + # signature; probe once here rather than per call. + self._entry_accepts_deterministic = _entry_accepts(self._entry, "deterministic") + self._varlen_accepts_deterministic = _entry_accepts(turbo_flex_attention_varlen, "deterministic") + + # -- mask selection ----------------------------------------------------- + def _block_mask_for(self, *, causal: bool, window: int, q_len: int, kv_len: int, device): + """Pick the BlockMask for this call, or ``None`` for unmasked (full) attention. + + An explicit user ``mask_mod`` replaces the built-in choice; combining it with a + window is rejected rather than silently resolved. Otherwise the Megatron mask + type (``causal`` / ``no_mask``) and ``window_size`` decide. + """ + if self.mask_mod is not None: + if window > 0: + raise NotImplementedError( + "Primus-Turbo flex attention: turbo_flex_attention_mask_mod and a sliding " + f"window (window={window}) were both requested. A user mask_mod replaces " + "the window rather than composing with it, so honouring either one would " + "silently drop the other. Fold the window into the mask_mod itself." + ) + return _get_block_mask( + mask_mod=self.mask_mod, + cache_key=("user", self.mask_mod_key), + q_len=q_len, + kv_len=kv_len, + device=device, + ) + if causal and window > 0: + return _get_block_mask( + mask_mod=_make_window_causal_mask_mod(window), + cache_key=("window_causal", window), + q_len=q_len, + kv_len=kv_len, + device=device, + ) + if causal: + return _get_block_mask( + mask_mod=_causal_mask_mod, + cache_key=("causal",), + q_len=q_len, + kv_len=kv_len, + device=device, + ) + if window > 0: + raise NotImplementedError( + "Primus-Turbo flex attention: a sliding window without a causal mask is not a " + "pattern the compat layer can express; set attn_mask_type=causal or disable the " + "window." + ) + return None # no_mask -> full attention + + # -- packed sequences (THD) --------------------------------------------- + def _call_varlen( + self, + query, + key, + value, + *, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + dropout_p, + softmax_scale, + causal, + window_size, + bias, + alibi_slopes, + return_lse, + sink, + deterministic=False, + ): + """Dispatch packed THD input through ``flex_attention_varlen``. + + The dense entries take a ``BlockMask`` that the compat layer has to *classify* + by probing ``mask_mod`` on a grid. Packing is different in kind: Megatron + already knows exactly where the document boundaries are and hands them over as + ``cu_seqlens``, so there is nothing to infer. That is why this path bypasses + the BlockMask cache entirely -- no probing, no classification, no per-shape + cache entry, which also means the variable-length workloads that would thrash + that cache do not touch it at all. + + ``causal=True`` here means *document-internal* causal (block-diagonal plus + within-segment causal), which is exactly Megatron's intent for a packed batch + and the thing the dense path could not express. + """ + if turbo_flex_attention_varlen is None: + raise RuntimeError( + "Primus-Turbo flex attention: packed (THD) sequences need " + "primus_turbo.pytorch.ops.attention.flex_attention_varlen, which this " + "Primus-Turbo build does not provide. Upgrade Primus-Turbo or set " + "use_turbo_flex_attention=false for packed-sequence runs." + ) + if query.dim() != 3: + raise NotImplementedError( + "Primus-Turbo flex attention: cu_seqlens was supplied, so the packed (THD) " + f"path was selected, but query is {query.dim()}D; THD expects " + "[total_tokens, H, D]." + ) + if cu_seqlens_q is None: + raise NotImplementedError( + "Primus-Turbo flex attention: packed (THD) input without cu_seqlens_q. The " + "document boundaries are not recoverable from the tensor shape, and guessing " + "them would let tokens attend across documents." + ) + if self.mask_mod is not None: + # A user mask_mod is a *dense* [q_idx, kv_idx] predicate; composing it with + # packing would need the product mask, which the varlen backend cannot take. + raise NotImplementedError( + "Primus-Turbo flex attention: turbo_flex_attention_mask_mod cannot be combined " + "with packed (THD) sequences; the varlen backend takes cu_seqlens, not a " + "BlockMask. Use dense batches, or fold the pattern into the packing itself." + ) + if self.score_mod is not None: + # flex_attention_varlen takes explicit alibi_slopes/softcap rather than probing + # a score_mod, so silently dropping the callable is the one thing we must not do. + raise NotImplementedError( + "Primus-Turbo flex attention: turbo_flex_attention_score_mod is not supported on " + "the packed (THD) path; flex_attention_varlen takes explicit alibi_slopes " + "instead of probing a score_mod." + ) + if bias is not None: + raise NotImplementedError( + "Primus-Turbo flex attention: an attention bias is not supported on the packed " + "(THD) path (flex_attention_varlen has no bias parameter)." + ) + if cu_seqlens_kv is None: + cu_seqlens_kv = cu_seqlens_q + if max_seqlen_q is None or max_seqlen_kv is None: + # Derivable from the prefix sums; one small D2H copy, and only on the first + # call of a shape in practice because Megatron caches PackedSeqParams. + diffs_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item() + diffs_kv = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]).max().item() + max_seqlen_q = int(max_seqlen_q or diffs_q) + max_seqlen_kv = int(max_seqlen_kv or diffs_kv) + + varlen_kwargs = {} + if self._varlen_accepts_deterministic: + varlen_kwargs["deterministic"] = deterministic + elif deterministic: + raise NotImplementedError( + "Primus-Turbo flex attention: deterministic=True was requested on the packed " + "(THD) path, but this Primus-Turbo build's flex_attention_varlen has no " + "'deterministic' parameter, so the flag cannot reach the kernel. Dropping it " + "would run a non-deterministic backward while the config asks for a " + "deterministic one, so this raises instead. Upgrade Primus-Turbo, or set " + "deterministic_mode=false." + ) + + return turbo_flex_attention_varlen( + query, + key, + value, + cu_seqlens_q, + cu_seqlens_kv, + int(max_seqlen_q), + int(max_seqlen_kv), + causal=causal, + window_size=window_size if window_size is not None else (-1, -1), + scale=softmax_scale, + alibi_slopes=alibi_slopes, + dropout_p=dropout_p, + sink=_coerce_sink(sink), + return_lse=return_lse, + **varlen_kwargs, + ) + + # -- the call itself ---------------------------------------------------- + def __call__( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + dropout_p: float = 0.0, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size: Tuple[int, int] = (-1, -1), + bias: Optional[torch.Tensor] = None, + alibi_slopes: Optional[torch.Tensor] = None, + deterministic: bool = False, + return_lse: bool = False, + return_attn_probs: bool = False, + sink: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + **extra: Any, + ): + if return_attn_probs: + raise NotImplementedError( + "Primus-Turbo flex attention: return_attn_probs is not supported (the compat " + "layer never materialises the attention probabilities)." + ) + if deterministic and not (self._entry_accepts_deterministic and self._varlen_accepts_deterministic): + # The compat layer used to hard-code deterministic=False, so honouring the + # request was impossible and dropping it silently was unacceptable. It now + # threads the flag through; only an older Primus-Turbo still cannot, and + # that is what this guard is left for. + raise NotImplementedError( + "Primus-Turbo flex attention: deterministic=True was requested, but this " + "Primus-Turbo build's flex entry points have no 'deterministic' parameter, so " + "the flag cannot reach the kernel. Upgrade Primus-Turbo, or set " + "deterministic_mode=false / use_turbo_flex_attention=false." + ) + if extra: + # ulysses_group / ring_group arrive here when context parallelism is on. + raise NotImplementedError( + "Primus-Turbo flex attention: unsupported backend arguments " + f"{sorted(extra)} (context parallelism has no compat-layer path yet)." + ) + + if query.dim() == 3 or cu_seqlens_q is not None: + return self._call_varlen( + query, + key, + value, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + bias=bias, + alibi_slopes=alibi_slopes, + return_lse=return_lse, + sink=sink, + deterministic=deterministic, + ) + if query.dim() != 4: + raise NotImplementedError( + "Primus-Turbo flex attention: expected bshd [B,S,H,D] or packed THD " + f"[total_tokens,H,D] input, got {query.dim()}D." + ) + + left_window = int(window_size[0]) if window_size is not None else -1 + right_window = int(window_size[1]) if window_size is not None else -1 + if right_window > 0: + # The dense path expresses a window as a mask_mod, and every mask_mod the + # compat layer can build is causal (q_idx >= kv_idx). A positive right bound + # means attending forward, which none of them can represent -- and honouring + # the left bound while quietly discarding the right one is exactly the silent + # wrongness this layer exists to prevent. (The packed path forwards + # window_size to the varlen backend verbatim, so it is not restricted here.) + raise NotImplementedError( + "Primus-Turbo flex attention: only causal windows are supported on the dense " + "path, so window_size[1] must be 0 or -1; got " + f"window_size={tuple(window_size)}." + ) + block_mask = self._block_mask_for( + causal=causal, + window=max(left_window, 0), + q_len=query.shape[1], + kv_len=key.shape[1], + device=query.device, + ) + + kwargs = dict( + score_mod=self.score_mod, + block_mask=block_mask, + scale=softmax_scale, + enable_gqa=query.shape[2] != key.shape[2], + return_lse=return_lse, + alibi_slopes=alibi_slopes, + dropout_p=dropout_p, + sink=_coerce_sink(sink), + bias=bias, + ) + if self._entry_accepts_deterministic: + kwargs["deterministic"] = deterministic + if self._entry_is_bshd: + return self._entry(query, key, value, **kwargs) + # Older Turbo build: go through the torch-layout entry (see the import block + # above for what the round-trip costs on each Turbo version). + out = self._entry(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), **kwargs) + if return_lse: + out, lse = out + return out.transpose(1, 2), lse + return out.transpose(1, 2) + + +# ============================================================================= +# Construction + config validation +# ============================================================================= + + +def build_turbo_flex_attention(*, args, config) -> TurboFlexAttention: + """Validate the configuration and build the flex-routed attention callable. + + Raises (rather than degrading to the direct ``flash_attn_func`` call) whenever the + requested combination has no compat-layer path, so an unsupported run fails at + model-build time instead of silently training on the old code path. + """ + if turbo_flex_attention is None: + raise RuntimeError( + "use_turbo_flex_attention=true, but this Primus-Turbo build does not provide " + "primus_turbo.pytorch.ops.attention.flex_attention " + f"({type(_FLEX_IMPORT_ERROR).__name__}: {_FLEX_IMPORT_ERROR}). Upgrade Primus-Turbo " + "or set use_turbo_flex_attention=false." + ) + if getattr(config, "context_parallel_size", 1) > 1: + raise NotImplementedError( + "use_turbo_flex_attention=true is not supported with context parallelism " + f"(context_parallel_size={config.context_parallel_size}); the compat layer has no " + "ulysses/ring path yet. Set use_turbo_flex_attention=false for CP runs." + ) + if getattr(args, "enable_turbo_attention_float8", False): + raise NotImplementedError( + "use_turbo_flex_attention=true cannot be combined with " + "enable_turbo_attention_float8=true; the compat layer supports fp16/bf16 only." + ) + if getattr(args, "deterministic_mode", False) and not ( + _entry_accepts(turbo_flex_attention_bshd or turbo_flex_attention, "deterministic") + and _entry_accepts(turbo_flex_attention_varlen, "deterministic") + ): + # The flag is threaded through now, so this is only reachable on an older + # Primus-Turbo whose flex entries predate the parameter. PrimusTurboAttention + # .forward passes deterministic=self.deterministic_mode on every call, so + # without this it would surface as a NotImplementedError from the first + # forward -- after model init, data loading and (on a large job) minutes of + # startup. Reject it here instead, next to the other build-time rejections. + raise NotImplementedError( + "use_turbo_flex_attention=true is not supported with deterministic_mode=true on " + "this Primus-Turbo build: its flex entry points have no 'deterministic' " + "parameter, so the flag cannot reach the kernel. Upgrade Primus-Turbo, or set " + "deterministic_mode=false / use_turbo_flex_attention=false. Note that the aiter " + "backward accumulates dQ with fp32 atomics regardless, so the flag buys " + "whatever the backend gives it -- it is not an independent guarantee of " + "bit-reproducible attention gradients." + ) + if getattr(args, "reset_attention_mask", False): + # reset_attention_mask asks Megatron for per-document causal masking inside a + # *dense* sample: the boundaries are baked into an attention_mask tensor that the + # Turbo call signature has no argument for, so they would be dropped and tokens + # would attend across documents with nothing raised anywhere. + # + # Note this is a different mechanism from qkv_format="thd", which IS supported: + # there Megatron hands over cu_seqlens explicitly and the adapter forwards them + # to flex_attention_varlen. Packing via THD is the supported way to get + # per-document masking on this path. + raise NotImplementedError( + "use_turbo_flex_attention=true is not supported with reset_attention_mask=true. " + "The dense flex dispatch receives only the causal flag, so the per-document " + "boundaries encoded in the attention mask would be dropped and tokens would " + "attend across documents. Use packed sequences (qkv_format='thd', which routes " + "through flex_attention_varlen and IS supported) instead, or set " + "reset_attention_mask=false / use_turbo_flex_attention=false." + ) + + mask_mod_path = getattr(args, "turbo_flex_attention_mask_mod", None) + if mask_mod_path and getattr(args, "sink_sliding_window", 0) > 0: + # Both describe the mask, and a user mask_mod replaces the window instead of + # composing with it (see _block_mask_for). That conflict is reachable from a + # plain config, so catch it here rather than at the first forward. + raise NotImplementedError( + f"turbo_flex_attention_mask_mod={mask_mod_path!r} and sink_sliding_window=" + f"{args.sink_sliding_window} both describe the attention mask, and the mask_mod " + "replaces the window rather than composing with it. Fold the window into the " + "mask_mod, or set sink_sliding_window=0." + ) + score_mod_path = getattr(args, "turbo_flex_attention_score_mod", None) + mask_mod = ( + resolve_dotted_callable(mask_mod_path, what="turbo_flex_attention_mask_mod") + if mask_mod_path + else None + ) + score_mod = ( + resolve_dotted_callable(score_mod_path, what="turbo_flex_attention_score_mod") + if score_mod_path + else None + ) + return TurboFlexAttention( + mask_mod=mask_mod, + score_mod=score_mod, + mask_mod_key=str(mask_mod_path), + ) diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index 672b4106c..eb428673e 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/primus/configs/modules/megatron/primus_turbo.yaml @@ -35,6 +35,27 @@ sink_sliding_window: 0 sink_window_even_layers_only: true # attention float8 path enable_turbo_attention_float8: false +# Route Turbo attention through the flex_attention compat layer instead of calling +# flash_attn_func directly. Requires use_turbo_attention: true; on its own it does +# nothing (and is rejected at model-build time rather than silently ignored). +# The compat layer takes a torch-FlexAttention block_mask/score_mod, classifies it and +# dispatches the recognised patterns onto the same Turbo kernels -- so with the default +# causal configuration this is numerically identical to the direct call. What it adds is +# the ability to express masks/score modifications that flash_attn_func has no argument +# for, plus packed (THD) sequences, which this switch is currently the ONLY way to run +# on the Turbo attention path (they route to flex_attention_varlen with cu_seqlens). +# Not supported with context_parallel_size > 1, enable_turbo_attention_float8, +# deterministic_mode, or reset_attention_mask -- each of those is rejected at +# model-build time rather than silently ignored. +use_turbo_flex_attention: false +# Optional 'package.module:attribute' path to a mask_mod callable +# (b, h, q_idx, kv_idx) -> bool, used to build the BlockMask instead of the mask implied +# by attn_mask_type/window_size. Only read when use_turbo_flex_attention is on. +turbo_flex_attention_mask_mod: null +# Optional 'package.module:attribute' path to a score_mod callable +# (score, b, h, q_idx, kv_idx) -> score, applied to the pre-softmax logits. +# Only read when use_turbo_flex_attention is on. +turbo_flex_attention_score_mod: null # ===== Parallel Linear ===== # operator switch diff --git a/tests/unit_tests/backends/megatron/test_primus_turbo_local_attention.py b/tests/unit_tests/backends/megatron/test_primus_turbo_local_attention.py new file mode 100644 index 000000000..e48c8943d --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_primus_turbo_local_attention.py @@ -0,0 +1,122 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""CPU unit tests for ``PrimusTurboLocalAttention.forward``'s parameter guards. + +``forward`` accepts ``attention_bias`` and ``packed_seq_params`` because Megatron's +core-attention protocol says it must, but this class forwards neither: the kernel +call passes ``bias=None``/``alibi_slopes=None`` unconditionally and never builds +cu_seqlens. Both used to be accepted and ignored -- a silent drop, which trains a +different model than the config asks for while the log stays clean. These tests pin +the refusals, and pin that they do NOT fire on the ordinary dense call. + +No GPU and no real Turbo build are needed: ``forward`` is invoked unbound against a +lightweight stub, so nothing in ``__init__`` (device queries, process groups, +``get_args``) runs. +""" + +from types import SimpleNamespace + +import pytest +import torch +from megatron.core.transformer.enums import AttnMaskType + +from primus.backends.megatron.core.extensions.primus_turbo_local_spec import ( + PrimusTurboLocalAttention, +) + + +class _Recorder: + """Stands in for the bound ``attn_func``; records what it was handed.""" + + def __init__(self): + self.calls = [] + + def __call__(self, q, k, v, **kwargs): + self.calls.append({"q": q, "k": k, "v": v, **kwargs}) + return torch.zeros_like(q) + + +def _stub(attn_func=None): + """The exact attribute surface ``forward`` reads -- nothing more.""" + return SimpleNamespace( + attn_func=attn_func if attn_func is not None else _Recorder(), + softmax_scale=0.125, + force_contiguous_qkv=False, + attn_kwargs={}, + ) + + +def _qkv(s=4, b=2, h=2, d=8): + """sbhd, the layout ``forward`` documents as its input.""" + return [torch.randn(s, b, h, d) for _ in range(3)] + + +def _call(stub, *, attention_bias=None, packed_seq_params=None, mask_type=AttnMaskType.causal): + q, k, v = _qkv() + return PrimusTurboLocalAttention.forward( + stub, + q, + k, + v, + attention_mask=None, + attn_mask_type=mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + ) + + +def test_attention_bias_is_refused_not_dropped(): + with pytest.raises(NotImplementedError, match="attention_bias"): + _call(_stub(), attention_bias=torch.zeros(1, 2, 4, 4)) + + +def test_the_attention_bias_refusal_names_alibi(): + # ALiBi is the way this actually reaches people: Megatron delivers it as an + # attention bias, so a user who sets position_embedding_type=alibi lands here + # without ever typing the word "bias". The message has to say so. + with pytest.raises(NotImplementedError, match="(?i)alibi"): + _call(_stub(), attention_bias=torch.zeros(1, 2, 4, 4)) + + +def test_packed_seq_params_is_refused_not_dropped(): + with pytest.raises(NotImplementedError, match="packed sequences"): + _call(_stub(), packed_seq_params=SimpleNamespace(qkv_format="thd")) + + +def test_the_packed_refusal_points_at_the_spec_that_implements_it(): + # A refusal that does not say where to go is a dead end; PrimusTurboAttention + # does implement qkv_format='thd'. + with pytest.raises(NotImplementedError, match="PrimusTurboAttention"): + _call(_stub(), packed_seq_params=SimpleNamespace(qkv_format="thd")) + + +def test_a_packed_seq_params_carrying_nothing_is_still_refused(): + # Emptiness is not permission. An object with no cu_seqlens still means the + # caller believes it is packing; guessing that it is not is exactly the + # approximation this guard exists to prevent. + with pytest.raises(NotImplementedError, match="packed sequences"): + _call(_stub(), packed_seq_params=SimpleNamespace()) + + +def test_the_ordinary_dense_call_is_untouched(): + rec = _Recorder() + out = _call(_stub(rec)) + assert len(rec.calls) == 1 + call = rec.calls[0] + # The guards must not have perturbed the call the kernel actually receives. + assert call["causal"] is True + assert call["bias"] is None + assert call["alibi_slopes"] is None + assert call["window_size"] == (-1, -1) + # sbhd in -> sbhd out with heads merged. + assert out.shape == (4, 2, 2 * 8) + + +def test_no_mask_still_reaches_the_kernel_as_non_causal(): + rec = _Recorder() + _call(_stub(rec), mask_type=AttnMaskType.no_mask) + assert rec.calls[0]["causal"] is False diff --git a/tests/unit_tests/backends/megatron/test_turbo_flex_attention.py b/tests/unit_tests/backends/megatron/test_turbo_flex_attention.py new file mode 100644 index 000000000..71f72f447 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_turbo_flex_attention.py @@ -0,0 +1,850 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""CPU unit tests for the ``use_turbo_flex_attention`` routing layer. + +Everything here runs without a GPU and without a real Primus-Turbo build: the +compat-layer entry and ``create_block_mask`` are stubbed, so what is under test is +the Primus-side wiring -- config validation, mask selection, the BlockMask cache, +argument passthrough, and the explicit rejections. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from primus.backends.megatron.core.extensions import turbo_flex_attention as tfa + + +def _dummy_mask_mod(b, h, q_idx, kv_idx): + """Module-level callable used to exercise dotted-path resolution.""" + return q_idx >= kv_idx + + +NOT_CALLABLE = 42 + + +class _Recorder: + """Stands in for ``flex_attention_bshd``; records what it was handed.""" + + def __init__(self): + self.calls = [] + + def __call__(self, q, k, v, **kwargs): + self.calls.append({"q": q, "k": k, "v": v, **kwargs}) + return torch.zeros_like(q) + + @property + def last(self): + return self.calls[-1] + + +class _VarlenRecorder: + """Stands in for ``flex_attention_varlen``; records positional + keyword args.""" + + def __init__(self): + self.calls = [] + + def __call__(self, q, k, v, cu_q, cu_k, max_q, max_k, **kwargs): + self.calls.append( + { + "q": q, + "k": k, + "v": v, + "cu_seqlens_q": cu_q, + "cu_seqlens_k": cu_k, + "max_seqlen_q": max_q, + "max_seqlen_k": max_k, + **kwargs, + } + ) + return torch.zeros_like(q) + + @property + def last(self): + return self.calls[-1] + + +@pytest.fixture +def stub_backend(monkeypatch): + """Install a recording compat-layer entry + a cheap BlockMask factory.""" + recorder = _Recorder() + made = [] + + def fake_create_block_mask(mask_mod, B=None, H=None, Q_LEN=None, KV_LEN=None, device=None): + obj = SimpleNamespace(mask_mod=mask_mod, shape=(Q_LEN, KV_LEN), device=device) + made.append(obj) + return obj + + varlen = _VarlenRecorder() + monkeypatch.setattr(tfa, "turbo_flex_attention", recorder) + monkeypatch.setattr(tfa, "turbo_flex_attention_bshd", recorder) + monkeypatch.setattr(tfa, "turbo_flex_attention_varlen", varlen) + monkeypatch.setattr(tfa, "turbo_create_block_mask", fake_create_block_mask) + monkeypatch.setattr(tfa, "_FLEX_IMPORT_ERROR", None) + tfa.clear_turbo_flex_block_mask_cache() + yield SimpleNamespace(recorder=recorder, made=made, varlen=varlen) + tfa.clear_turbo_flex_block_mask_cache() + + +def _qkv(b=2, s=16, h=4, d=8, hkv=None): + hkv = h if hkv is None else hkv + return ( + torch.randn(b, s, h, d), + torch.randn(b, s, hkv, d), + torch.randn(b, s, hkv, d), + ) + + +def _args(**overrides): + base = dict( + use_turbo_flex_attention=True, + enable_turbo_attention_float8=False, + turbo_flex_attention_mask_mod=None, + turbo_flex_attention_score_mod=None, + deterministic_mode=False, + reset_attention_mask=False, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +def _config(**overrides): + base = dict(context_parallel_size=1) + base.update(overrides) + return SimpleNamespace(**base) + + +# ============================================================================= +# dotted-path resolution +# ============================================================================= + + +class TestResolveDottedCallable: + def test_resolves_module_attribute(self): + fn = tfa.resolve_dotted_callable( + "tests.unit_tests.backends.megatron.test_turbo_flex_attention:_dummy_mask_mod", + what="mask_mod", + ) + assert fn is _dummy_mask_mod + + def test_missing_colon_rejected(self): + with pytest.raises(ValueError, match="package.module:attribute"): + tfa.resolve_dotted_callable("some.module.attr", what="mask_mod") + + def test_non_string_rejected(self): + with pytest.raises(ValueError): + tfa.resolve_dotted_callable(_dummy_mask_mod, what="mask_mod") + + def test_unknown_module_reports_clearly(self): + with pytest.raises(ImportError, match="could not import module"): + tfa.resolve_dotted_callable("primus_no_such_module_xyz:fn", what="mask_mod") + + def test_unknown_attribute_reports_clearly(self): + with pytest.raises(AttributeError, match="has no attribute"): + tfa.resolve_dotted_callable( + "tests.unit_tests.backends.megatron.test_turbo_flex_attention:nope", + what="mask_mod", + ) + + def test_non_callable_rejected(self): + with pytest.raises(TypeError, match="not callable"): + tfa.resolve_dotted_callable( + "tests.unit_tests.backends.megatron.test_turbo_flex_attention:NOT_CALLABLE", + what="mask_mod", + ) + + +# ============================================================================= +# mask_mod semantics +# ============================================================================= + + +class TestMaskMods: + def test_causal_mask_mod(self): + q = torch.arange(4).view(-1, 1) + kv = torch.arange(4).view(1, -1) + assert torch.equal(tfa._causal_mask_mod(0, 0, q, kv), q >= kv) + + def test_window_causal_mask_mod(self): + q = torch.arange(8).view(-1, 1) + kv = torch.arange(8).view(1, -1) + got = tfa._make_window_causal_mask_mod(3)(0, 0, q, kv) + expected = (q >= kv) & ((q - kv) <= 3) + assert torch.equal(got, expected) + # row 7 keeps exactly the 4 positions 4..7 + assert int(got[7].sum()) == 4 + + +# ============================================================================= +# BlockMask cache +# ============================================================================= + + +class TestBlockMaskCache: + def test_same_shape_reuses_one_mask(self, stub_backend): + attn = tfa.TurboFlexAttention() + q, k, v = _qkv() + for _ in range(5): + attn(q, k, v, causal=True) + assert len(stub_backend.made) == 1 + masks = {id(c["block_mask"]) for c in stub_backend.recorder.calls} + assert len(masks) == 1 + + def test_different_seqlen_builds_a_second_mask(self, stub_backend): + attn = tfa.TurboFlexAttention() + attn(*_qkv(s=16), causal=True) + attn(*_qkv(s=32), causal=True) + assert len(stub_backend.made) == 2 + + def test_window_and_causal_are_distinct_entries(self, stub_backend): + attn = tfa.TurboFlexAttention() + q, k, v = _qkv() + attn(q, k, v, causal=True) + attn(q, k, v, causal=True, window_size=(4, 0)) + attn(q, k, v, causal=True, window_size=(8, 0)) + assert len(stub_backend.made) == 3 + + def test_clear_cache_forces_a_rebuild(self, stub_backend): + attn = tfa.TurboFlexAttention() + q, k, v = _qkv() + attn(q, k, v, causal=True) + tfa.clear_turbo_flex_block_mask_cache() + attn(q, k, v, causal=True) + assert len(stub_backend.made) == 2 + + +# ============================================================================= +# dispatch behaviour +# ============================================================================= + + +class TestDispatch: + def test_no_mask_sends_no_block_mask(self, stub_backend): + attn = tfa.TurboFlexAttention() + attn(*_qkv(), causal=False) + assert stub_backend.recorder.last["block_mask"] is None + assert stub_backend.made == [] + + def test_causal_sends_a_causal_block_mask(self, stub_backend): + attn = tfa.TurboFlexAttention() + attn(*_qkv(), causal=True) + bm = stub_backend.recorder.last["block_mask"] + assert bm is not None + q = torch.arange(4).view(-1, 1) + kv = torch.arange(4).view(1, -1) + assert torch.equal(bm.mask_mod(0, 0, q, kv), q >= kv) + + def test_window_size_becomes_a_window_mask(self, stub_backend): + attn = tfa.TurboFlexAttention() + attn(*_qkv(s=16), causal=True, window_size=(4, 0)) + bm = stub_backend.recorder.last["block_mask"] + q = torch.arange(16).view(-1, 1) + kv = torch.arange(16).view(1, -1) + expected = (q >= kv) & ((q - kv) <= 4) + assert torch.equal(bm.mask_mod(0, 0, q, kv), expected) + + def test_user_mask_mod_overrides_causal(self, stub_backend): + # No window here on purpose: this asserts mask_mod beats the built-in *causal* + # mask. mask_mod combined with a window is now a rejected conflict, covered by + # TestRejections.test_user_mask_mod_plus_window_rejected. + attn = tfa.TurboFlexAttention(mask_mod=_dummy_mask_mod, mask_mod_key="dummy") + attn(*_qkv(), causal=True) + assert stub_backend.recorder.last["block_mask"].mask_mod is _dummy_mask_mod + + def test_score_mod_is_forwarded(self, stub_backend): + def score_mod(score, b, h, q_idx, kv_idx): + return score + + attn = tfa.TurboFlexAttention(score_mod=score_mod) + attn(*_qkv(), causal=True) + assert stub_backend.recorder.last["score_mod"] is score_mod + + def test_enable_gqa_follows_head_counts(self, stub_backend): + attn = tfa.TurboFlexAttention() + attn(*_qkv(h=8, hkv=8), causal=True) + assert stub_backend.recorder.last["enable_gqa"] is False + attn(*_qkv(h=8, hkv=2), causal=True) + assert stub_backend.recorder.last["enable_gqa"] is True + + def test_optional_arguments_are_forwarded(self, stub_backend): + attn = tfa.TurboFlexAttention() + sink = torch.zeros(4) + alibi = torch.ones(4) + bias = torch.zeros(2, 4, 16, 16) + attn( + *_qkv(), + dropout_p=0.1, + softmax_scale=0.25, + causal=True, + bias=bias, + alibi_slopes=alibi, + sink=sink, + ) + call = stub_backend.recorder.last + assert call["scale"] == 0.25 + assert call["dropout_p"] == 0.1 + assert call["sink"] is sink + assert call["alibi_slopes"] is alibi + assert call["bias"] is bias + assert call["deterministic"] is False + + @pytest.mark.parametrize("flag", [True, False]) + def test_deterministic_is_forwarded(self, stub_backend, flag): + attn = tfa.TurboFlexAttention() + attn(*_qkv(), causal=True, deterministic=flag) + assert stub_backend.recorder.last["deterministic"] is flag + + @pytest.mark.parametrize("flag", [True, False]) + def test_deterministic_is_forwarded_on_the_packed_path(self, stub_backend, flag): + attn = tfa.TurboFlexAttention() + q = torch.randn(24, 4, 8) + cu = torch.tensor([0, 8, 16, 24], dtype=torch.int32) + attn(q, q.clone(), q.clone(), causal=True, cu_seqlens_q=cu, deterministic=flag) + assert stub_backend.varlen.last["deterministic"] is flag + + def test_bshd_entry_receives_tensors_uncopied(self, stub_backend): + attn = tfa.TurboFlexAttention() + q, k, v = _qkv() + attn(q, k, v, causal=True) + call = stub_backend.recorder.last + assert call["q"].data_ptr() == q.data_ptr() + assert call["k"].data_ptr() == k.data_ptr() + assert call["v"].data_ptr() == v.data_ptr() + + def test_legacy_bhsd_entry_round_trips_the_layout(self, monkeypatch, stub_backend): + """Older Turbo builds without flex_attention_bshd still get bshd out.""" + monkeypatch.setattr(tfa, "turbo_flex_attention_bshd", None) + attn = tfa.TurboFlexAttention() + assert attn._entry_is_bshd is False + q, k, v = _qkv(b=2, s=16, h=4, d=8) + out = attn(q, k, v, causal=True) + seen_q = stub_backend.recorder.last["q"] + assert seen_q.shape == (2, 4, 16, 8) # bhsd on the way in + assert out.shape == (2, 16, 4, 8) # bshd on the way out + + def test_return_lse_tuple_is_preserved(self, monkeypatch, stub_backend): + monkeypatch.setattr(tfa, "turbo_flex_attention_bshd", None) + + def entry(q, k, v, **kwargs): + return torch.zeros_like(q), torch.zeros(q.shape[0], q.shape[1], q.shape[2]) + + monkeypatch.setattr(tfa, "turbo_flex_attention", entry) + attn = tfa.TurboFlexAttention() + out, lse = attn(*_qkv(b=2, s=16, h=4, d=8), causal=True, return_lse=True) + assert out.shape == (2, 16, 4, 8) + assert lse.shape == (2, 4, 16) + + +# ============================================================================= +# explicit rejections (never a silent fallback) +# ============================================================================= + + +class TestRejections: + def test_thd_without_boundaries_rejected(self, stub_backend): + """THD is supported now (see TestPackedTHD), but only with cu_seqlens. Bare 3D + input has no recoverable document structure, so it must still be refused rather + than treated as one long sequence.""" + attn = tfa.TurboFlexAttention() + t = torch.randn(32, 4, 8) + with pytest.raises(NotImplementedError, match="cu_seqlens"): + attn(t, t, t, causal=True) + + def test_2d_input_rejected(self, stub_backend): + attn = tfa.TurboFlexAttention() + t = torch.randn(32, 8) + with pytest.raises(NotImplementedError, match="THD"): + attn(t, t, t, causal=True) + + def test_return_attn_probs_rejected(self, stub_backend): + attn = tfa.TurboFlexAttention() + with pytest.raises(NotImplementedError, match="return_attn_probs"): + attn(*_qkv(), causal=True, return_attn_probs=True) + + def test_deterministic_rejected_on_an_older_turbo(self, stub_backend, monkeypatch): + # The flag is threaded through now; it is only refused when the installed + # compat layer has no parameter to put it in. A stub with an explicit + # signature (no ``**kwargs``) stands in for that older build. + def old_entry( + q, + k, + v, + score_mod=None, + block_mask=None, + scale=None, + enable_gqa=False, + return_lse=False, + alibi_slopes=None, + dropout_p=0.0, + sink=None, + bias=None, + ): + return torch.zeros_like(q) + + monkeypatch.setattr(tfa, "turbo_flex_attention_bshd", old_entry) + monkeypatch.setattr(tfa, "turbo_flex_attention", old_entry) + attn = tfa.TurboFlexAttention() + with pytest.raises(NotImplementedError, match="deterministic"): + attn(*_qkv(), causal=True, deterministic=True) + # ...and it is a *conditional* refusal: the default still runs. + assert attn(*_qkv(), causal=True, deterministic=False) is not None + + def test_context_parallel_groups_rejected(self, stub_backend): + attn = tfa.TurboFlexAttention() + with pytest.raises(NotImplementedError, match="ulysses_group"): + attn(*_qkv(), causal=True, ulysses_group=object()) + + def test_window_without_causal_rejected(self, stub_backend): + attn = tfa.TurboFlexAttention() + with pytest.raises(NotImplementedError, match="sliding window"): + attn(*_qkv(), causal=False, window_size=(4, 0)) + + def test_forward_looking_window_rejected(self, stub_backend): + """window_size[1] > 0 asks each query to see tokens ahead of it. Every mask_mod + the dense path can build is causal, so the right bound cannot be honoured -- and + it used to be dropped on the floor while the left bound was applied, which + silently trains a different model than the config asked for.""" + attn = tfa.TurboFlexAttention() + with pytest.raises(NotImplementedError, match=r"window_size\[1\]"): + attn(*_qkv(s=16), causal=True, window_size=(4, 4)) + + def test_symmetric_window_is_not_silently_truncated(self, stub_backend): + """The specific regression: (4, 4) must not be treated as (4, 0).""" + attn = tfa.TurboFlexAttention() + with pytest.raises(NotImplementedError): + attn(*_qkv(s=16), causal=True, window_size=(4, 4)) + assert stub_backend.recorder.calls == [], "the backend was reached despite the rejection" + + def test_unbounded_right_window_still_accepted(self, stub_backend): + """(-1) and 0 both mean 'no forward bound' and must keep working.""" + attn = tfa.TurboFlexAttention() + attn(*_qkv(s=16), causal=True, window_size=(4, 0)) + attn(*_qkv(s=16), causal=True, window_size=(4, -1)) + assert len(stub_backend.recorder.calls) == 2 + + def test_mask_mod_plus_sink_window_rejected_at_build(self, stub_backend): + """The same conflict, but reachable from a plain config file -- so it has to be + caught when the model is built, not on the first forward pass.""" + with pytest.raises(NotImplementedError, match="sink_sliding_window"): + tfa.build_turbo_flex_attention( + args=_args( + turbo_flex_attention_mask_mod="mod:fn", + sink_sliding_window=512, + ), + config=_config(), + ) + + def test_user_mask_mod_plus_window_rejected(self, stub_backend): + """A user mask_mod replaces the window rather than composing with it. Resolving + that silently would drop whichever one lost, so the conflict is an error.""" + attn = tfa.TurboFlexAttention(mask_mod=lambda b, h, q, kv: q >= kv) + with pytest.raises(NotImplementedError, match="mask_mod"): + attn(*_qkv(s=16), causal=True, window_size=(4, 0)) + + +# ============================================================================= +# build-time validation +# ============================================================================= + + +class TestBuild: + def test_builds_with_defaults(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + assert isinstance(attn, tfa.TurboFlexAttention) + assert attn.mask_mod is None and attn.score_mod is None + + def test_resolves_configured_hooks(self, stub_backend): + path = "tests.unit_tests.backends.megatron.test_turbo_flex_attention:_dummy_mask_mod" + attn = tfa.build_turbo_flex_attention( + args=_args(turbo_flex_attention_mask_mod=path, turbo_flex_attention_score_mod=path), + config=_config(), + ) + assert attn.mask_mod is _dummy_mask_mod + assert attn.score_mod is _dummy_mask_mod + + def test_context_parallel_rejected(self, stub_backend): + with pytest.raises(NotImplementedError, match="context parallel"): + tfa.build_turbo_flex_attention(args=_args(), config=_config(context_parallel_size=2)) + + def test_deterministic_mode_accepted_on_a_current_turbo(self, stub_backend): + # The compat layer threads ``deterministic`` now, so deterministic_mode=true is + # no longer a build-time rejection -- it is simply forwarded. + attn = tfa.build_turbo_flex_attention(args=_args(deterministic_mode=True), config=_config()) + attn(*_qkv(), causal=True, deterministic=True) + assert stub_backend.recorder.last["deterministic"] is True + + def test_deterministic_mode_rejected_at_build_on_an_older_turbo(self, stub_backend, monkeypatch): + def old_entry( + q, + k, + v, + score_mod=None, + block_mask=None, + scale=None, + enable_gqa=False, + return_lse=False, + alibi_slopes=None, + dropout_p=0.0, + sink=None, + bias=None, + ): + return torch.zeros_like(q) + + monkeypatch.setattr(tfa, "turbo_flex_attention_bshd", old_entry) + monkeypatch.setattr(tfa, "turbo_flex_attention", old_entry) + with pytest.raises(NotImplementedError, match="deterministic"): + tfa.build_turbo_flex_attention(args=_args(deterministic_mode=True), config=_config()) + + def test_float8_rejected(self, stub_backend): + with pytest.raises(NotImplementedError, match="float8"): + tfa.build_turbo_flex_attention(args=_args(enable_turbo_attention_float8=True), config=_config()) + + def test_missing_turbo_build_reports_clearly(self, monkeypatch, stub_backend): + monkeypatch.setattr(tfa, "turbo_flex_attention", None) + monkeypatch.setattr(tfa, "_FLEX_IMPORT_ERROR", ModuleNotFoundError("No module named 'primus_turbo'")) + with pytest.raises(RuntimeError, match="does not provide"): + tfa.build_turbo_flex_attention(args=_args(), config=_config()) + + +# ============================================================================= +# the config key itself +# ============================================================================= +# packed sequences (THD) -> flex_attention_varlen +# ============================================================================= + + +def _thd(total=24, h=4, d=8, lens=(8, 10, 6)): + q = torch.randn(total, h, d) + k = torch.randn(total, h, d) + v = torch.randn(total, h, d) + cu = torch.tensor([0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32) + return q, k, v, cu, lens + + +class TestPackedTHD: + """THD input must reach flex_attention_varlen with the boundaries intact. + + Before this path existed, packed sequences either hit a NotImplementedError or -- + on the direct flash_attn_func binding -- silently lost cu_seqlens and let tokens + attend across documents. Both are covered here. + """ + + def test_thd_routes_to_varlen(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + out = attn(q, k, v, causal=True, cu_seqlens_q=cu, max_seqlen_q=max(lens), max_seqlen_kv=max(lens)) + assert out.shape == q.shape + assert len(stub_backend.varlen.calls) == 1 + assert len(stub_backend.recorder.calls) == 0, "must not go through the dense entry" + + def test_boundaries_are_forwarded_verbatim(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + attn(q, k, v, causal=True, cu_seqlens_q=cu, max_seqlen_q=max(lens), max_seqlen_kv=max(lens)) + call = stub_backend.varlen.last + assert torch.equal(call["cu_seqlens_q"], cu) + assert torch.equal(call["cu_seqlens_k"], cu), "kv defaults to q when omitted" + assert call["max_seqlen_q"] == max(lens) + assert call["causal"] is True + + def test_max_seqlen_derived_when_missing(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + attn(q, k, v, causal=True, cu_seqlens_q=cu) + call = stub_backend.varlen.last + assert call["max_seqlen_q"] == max(lens) + assert call["max_seqlen_k"] == max(lens) + + def test_separate_kv_boundaries(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + cu_kv = cu.clone() + attn( + q, + k, + v, + causal=False, + cu_seqlens_q=cu, + cu_seqlens_kv=cu_kv, + max_seqlen_q=max(lens), + max_seqlen_kv=max(lens), + ) + assert stub_backend.varlen.last["cu_seqlens_k"] is cu_kv + + def test_no_block_mask_is_built(self, stub_backend): + """Packing carries its boundaries explicitly -- nothing to probe or classify, + so the BlockMask cache must stay empty (this is why THD does not thrash it).""" + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + attn(q, k, v, causal=True, cu_seqlens_q=cu) + assert stub_backend.made == [] + assert len(tfa._BLOCK_MASK_CACHE) == 0 + + def test_sink_is_cast_on_the_packed_path_too(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + sink = torch.zeros(4, dtype=torch.bfloat16) + attn(q, k, v, causal=True, cu_seqlens_q=cu, sink=sink) + assert stub_backend.varlen.last["sink"].dtype == torch.float32 + + def test_3d_without_cu_seqlens_is_refused(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + with pytest.raises(NotImplementedError, match="cu_seqlens"): + attn(q, k, v, causal=True) + + def test_cu_seqlens_with_4d_is_refused(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v = _qkv() + cu = torch.tensor([0, 8, 16], dtype=torch.int32) + with pytest.raises(NotImplementedError, match="THD expects"): + attn(q, k, v, causal=True, cu_seqlens_q=cu) + + def test_user_mask_mod_is_refused_on_packed(self, stub_backend): + attn = tfa.build_turbo_flex_attention( + args=_args( + turbo_flex_attention_mask_mod=( + "tests.unit_tests.backends.megatron." "test_turbo_flex_attention:_dummy_mask_mod" + ) + ), + config=_config(), + ) + q, k, v, cu, lens = _thd() + with pytest.raises(NotImplementedError, match="mask_mod"): + attn(q, k, v, causal=True, cu_seqlens_q=cu) + + def test_bias_is_refused_on_packed(self, stub_backend): + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + with pytest.raises(NotImplementedError, match="bias"): + attn(q, k, v, causal=True, cu_seqlens_q=cu, bias=torch.zeros(1)) + + def test_missing_varlen_entry_reports_clearly(self, stub_backend, monkeypatch): + monkeypatch.setattr(tfa, "turbo_flex_attention_varlen", None) + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v, cu, lens = _thd() + with pytest.raises(RuntimeError, match="flex_attention_varlen"): + attn(q, k, v, causal=True, cu_seqlens_q=cu) + + def test_dense_path_is_unaffected(self, stub_backend): + """The packed branch must not perturb the ordinary bshd call.""" + attn = tfa.build_turbo_flex_attention(args=_args(), config=_config()) + q, k, v = _qkv() + attn(q, k, v, causal=True) + assert len(stub_backend.recorder.calls) == 1 + assert len(stub_backend.varlen.calls) == 0 + + +# ============================================================================= + + +class TestSinkDtype: + """Megatron allocates ``self.sinks`` in bfloat16; the compat layer demands fp32.""" + + def test_bf16_sink_is_cast_to_fp32(self, stub_backend): + attn = tfa.TurboFlexAttention() + q, k, v = _qkv(h=4) + sink = torch.zeros(4, dtype=torch.bfloat16) + attn(q, k, v, causal=True, sink=sink) + assert stub_backend.recorder.last["sink"].dtype == torch.float32 + + def test_fp32_sink_is_passed_through_untouched(self, stub_backend): + attn = tfa.TurboFlexAttention() + q, k, v = _qkv(h=4) + sink = torch.zeros(4, dtype=torch.float32) + attn(q, k, v, causal=True, sink=sink) + assert stub_backend.recorder.last["sink"] is sink + + def test_none_sink_stays_none(self, stub_backend): + attn = tfa.TurboFlexAttention() + q, k, v = _qkv() + attn(q, k, v, causal=True, sink=None) + assert stub_backend.recorder.last["sink"] is None + + def test_cast_keeps_the_parameter_trainable(self, stub_backend): + """The cast must be differentiable, or the sink Parameter stops learning.""" + sink = torch.nn.Parameter(torch.zeros(4, dtype=torch.bfloat16)) + cast = tfa._coerce_sink(sink) + assert cast.dtype == torch.float32 + cast.sum().backward() + assert sink.grad is not None + assert sink.grad.dtype == torch.bfloat16 + + +class TestBoundedCache: + def test_cache_is_bounded_and_evicts(self, stub_backend, monkeypatch): + monkeypatch.setattr(tfa, "_BLOCK_MASK_CACHE_MAX", 4) + attn = tfa.TurboFlexAttention() + for s_len in range(8, 8 + 10): + q, k, v = _qkv(s=s_len) + attn(q, k, v, causal=True) + assert len(tfa._BLOCK_MASK_CACHE) <= 4 + + def test_eviction_warns_once(self, stub_backend, monkeypatch): + monkeypatch.setattr(tfa, "_BLOCK_MASK_CACHE_MAX", 2) + attn = tfa.TurboFlexAttention() + with pytest.warns(RuntimeWarning, match="BlockMask cache"): + for s_len in range(8, 14): + q, k, v = _qkv(s=s_len) + attn(q, k, v, causal=True) + # A second burst must stay quiet -- one warning per process, not per miss. + with warnings.catch_warnings(): + warnings.simplefilter("error") + for s_len in range(20, 26): + q, k, v = _qkv(s=s_len) + attn(q, k, v, causal=True) + + def test_hot_shape_survives_a_cold_tail(self, stub_backend, monkeypatch): + """LRU, not FIFO: the shape used every step must not be evicted by the tail.""" + monkeypatch.setattr(tfa, "_BLOCK_MASK_CACHE_MAX", 4) + attn = tfa.TurboFlexAttention() + hot = _qkv(s=64) + attn(*hot, causal=True) + first = stub_backend.recorder.last["block_mask"] + for s_len in range(8, 20): + attn(*_qkv(s=s_len), causal=True) + attn(*hot, causal=True) # touched every step, so it stays hot + attn(*hot, causal=True) + assert stub_backend.recorder.last["block_mask"] is first + + +class TestBuildTimeRejections: + """These combinations must fail at model build, not on the first forward. + + ``deterministic_mode`` used to be one of them and no longer is: the compat layer + threads ``deterministic`` through, so the build only rejects it on a Turbo whose + entries predate the parameter. Both halves of that narrower contract live in + ``TestBuild`` -- ``test_deterministic_mode_accepted_on_a_current_turbo`` and + ``test_deterministic_mode_rejected_at_build_on_an_older_turbo``. + """ + + def test_reset_attention_mask_rejected_at_build(self, stub_backend): + with pytest.raises(NotImplementedError, match="reset_attention_mask"): + tfa.build_turbo_flex_attention(args=_args(reset_attention_mask=True), config=_config()) + + def test_defaults_still_build(self, stub_backend): + assert tfa.build_turbo_flex_attention(args=_args(), config=_config()) is not None + + +class TestConfigSchema: + def test_registered_and_defaults_to_false(self): + import pathlib + + import yaml + + root = pathlib.Path(__file__).resolve().parents[4] + cfg = yaml.safe_load((root / "primus/configs/modules/megatron/primus_turbo.yaml").read_text()) + assert cfg["use_turbo_flex_attention"] is False + assert cfg["turbo_flex_attention_mask_mod"] is None + assert cfg["turbo_flex_attention_score_mod"] is None + + +class TestRejectResetAttentionMask: + """The direct-path guard for the mask tensor ``forward()`` never reads. + + Before this guard, ``PrimusTurboAttention.forward`` accepted ``attention_mask`` and + never looked at it. With ``reset_attention_mask=true`` that silently trained a + different model than the config asked for: the per-document boundaries were dropped + and tokens attended across documents, with no error and no warning anywhere. + """ + + def test_off_is_a_noop(self): + tfa.reject_reset_attention_mask(_args(reset_attention_mask=False), is_flex=False, where="X") + + def test_missing_attribute_is_a_noop(self): + # Older arg namespaces predate the flag; absence must not raise. + tfa.reject_reset_attention_mask(SimpleNamespace(), is_flex=False, where="X") + + def test_on_and_direct_path_raises(self): + with pytest.raises(NotImplementedError, match="reset_attention_mask"): + tfa.reject_reset_attention_mask(_args(reset_attention_mask=True), is_flex=False, where="X") + + def test_error_names_the_caller_and_the_supported_alternative(self): + with pytest.raises(NotImplementedError) as excinfo: + tfa.reject_reset_attention_mask( + _args(reset_attention_mask=True), is_flex=False, where="MyAttention" + ) + msg = str(excinfo.value) + assert "MyAttention" in msg + assert "thd" in msg + + def test_flex_path_defers_to_its_own_message(self): + # build_turbo_flex_attention raises for itself with a message that names the THD + # alternative in flex terms; this helper must not pre-empt it. + tfa.reject_reset_attention_mask(_args(reset_attention_mask=True), is_flex=True, where="X") + + def test_flex_build_still_rejects_it(self, stub_backend): + # The deferral above is only sound while the flex side actually raises. + with pytest.raises(NotImplementedError, match="reset_attention_mask"): + tfa.build_turbo_flex_attention(args=_args(reset_attention_mask=True), config=_config()) + + +class TestSinkKwargsFor: + """The three-way answer to "can this entry take a sink?". + + Primus bound one of several attention entries and then passed ``sink=`` to all of + them. Only ``flash_attn_func`` has the parameter: the context-parallel, fp8 and + packed-varlen entries do not. So every ``context_parallel_size > 1`` run died on a + bare ``TypeError: flash_attn_usp_func() got an unexpected keyword argument 'sink'`` + -- even with no sink configured, and with nothing in the message to say why + (measured on 2x MI355, TASK_PROGRESS EXP21 case B; the entry inventory is EXP22). + """ + + @staticmethod + def _takes_sink(q, k, v, sink=None): + return None + + @staticmethod + def _no_sink(q, k, v): + return None + + @staticmethod + def _kwargs_catch_all(q, k, v, **kw): + return None + + def test_forwards_when_the_entry_takes_one(self): + sink = torch.zeros(4) + got = tfa.sink_kwargs_for(self._takes_sink, sink, where="X") + assert list(got) == ["sink"] + assert got["sink"] is sink + + def test_forwards_none_when_the_entry_takes_one(self): + # None is the parameter's own default; passing it explicitly is harmless and + # keeps the forwarding branch free of a second special case. + assert tfa.sink_kwargs_for(self._takes_sink, None, where="X") == {"sink": None} + + def test_omits_when_the_entry_cannot_take_one_and_there_is_nothing_to_lose(self): + # The bug this fixes: no sink configured, yet the call still died on the kwarg. + assert tfa.sink_kwargs_for(self._no_sink, None, where="X") == {} + + def test_raises_when_a_sink_would_be_dropped(self): + # Omitting it here would change the softmax denominator of every query and + # train a different model in silence -- so this must be loud, not lenient. + with pytest.raises(NotImplementedError, match="attention sink"): + tfa.sink_kwargs_for(self._no_sink, torch.zeros(4), where="X") + + def test_error_names_the_entry_and_the_configuration(self): + with pytest.raises(NotImplementedError) as excinfo: + tfa.sink_kwargs_for(self._no_sink, torch.zeros(4), where="context_parallel_size=2") + msg = str(excinfo.value) + assert "_no_sink" in msg + assert "context_parallel_size=2" in msg + + def test_unbound_entry_with_a_sink_raises_rather_than_calling_none(self): + # attn_varlen is None when no varlen entry was bound; a sink plus that + # combination is a configuration error, not an AttributeError later on. + with pytest.raises(NotImplementedError): + tfa.sink_kwargs_for(None, torch.zeros(4), where="X") + + def test_kwargs_catch_all_is_treated_as_accepting(self): + # A **kwargs entry can genuinely take a sink, and refusing it would ground + # wrappers that forward everything. The cost is that a permissive test double + # also passes -- acceptable, since a real entry that ignores **kwargs is a bug + # on its own side. + assert tfa.sink_kwargs_for(self._kwargs_catch_all, None, where="X") == {"sink": None}