[WIP]feat(turbo): route attention through Primus-Turbo's flex_attention when enabled

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`

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__).

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 (