perf(bagel): accelerate agent SFT replay with FlexAttention

examples/sft/bagel_agent_sft.yaml

@@ -17,6 +17,7 @@
 #   BAGEL_PATH      HF id or local ckpt (default ByteDance-Seed/BAGEL-7B-MoT)
 #   SFT_DATA        train manifest (searchgen prepare_sft train.jsonl)
 #   SFT_EVAL_DATA   eval manifest (defaults to SFT_DATA — supply a real split!)
+#   BAGEL_REPLAY_ATTENTION_BACKEND  sdpa (baseline) or flex (default: flex)
 #
 # Download/cook guide: datasets/searchgen/README.md (manifests shared with qwen_vl_agent_sft).
 #
@@ -69,6 +70,7 @@ pipeline:
   shift: 3.0
   max_prompt_length: 24576
   replay_mode: train                 # full-grad und replay (interleaved packed forward_train)
+  replay_attention_backend: ${oc.env:BAGEL_REPLAY_ATTENTION_BACKEND,flex}
   strategy:
     _target_: unirl.sde.kernels.FlowSDEStrategy
 

unirl/models/bagel/ar.py

@@ -81,6 +81,7 @@ def __init__(
         autocast_precision: str = "bf16",
         logprob_precision: str = "fp32",
         replay_mode: str = "train",
+        replay_attention_backend: str = "sdpa",
     ) -> None:
         self.model = model
         self.autocast_dtype = parse_torch_dtype(autocast_precision, field_name="BagelARStage.autocast_precision")
@@ -90,6 +91,11 @@ def __init__(
             self.replay_mode in ("train", "inference"),
             f"BagelARStage: replay_mode must be 'train' or 'inference'; got {replay_mode!r}.",
         )
+        self.replay_attention_backend = str(replay_attention_backend).strip().lower()
+        require(
+            self.replay_attention_backend in ("sdpa", "flex"),
+            f"BagelARStage: replay_attention_backend must be 'sdpa' or 'flex'; got {replay_attention_backend!r}.",
+        )
         llm_cfg = getattr(getattr(getattr(model, "model", None), "config", None), "llm_config", None)
         require(
             not getattr(llm_cfg, "freeze_und", False),
@@ -267,6 +273,7 @@ def _replay_train(
                     splits=splits,
                     response_input=response_input,
                     device=device,
+                    attention_backend=self.replay_attention_backend,
                 )
                 logits = rl_ops.und_replay_logits(bagel, packed)
                 logp_full = torch.log_softmax(logits.float() / temp, dim=-1)

unirl/models/bagel/pipeline.py

@@ -61,6 +61,7 @@ def __init__(
         logprob_precision: str = "fp32",
         shift: float = 3.0,
         replay_mode: str = "train",
+        replay_attention_backend: str = "sdpa",
         max_prompt_length: int = 8192,
         cache_t2i_contexts: Optional[bool] = None,
         context_cache_size: Optional[int] = None,
@@ -84,6 +85,7 @@ def __init__(
             autocast_precision=autocast_precision,
             logprob_precision=logprob_precision,
             replay_mode=replay_mode,
+            replay_attention_backend=replay_attention_backend,
         )
         self.autocast_precision = autocast_precision
         self.shift = shift

unirl/models/bagel/rl_ops.py

@@ -443,18 +443,57 @@ def _chunk_logp(h: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
     return torch.cat(parts, dim=0)
 
 
+def _build_und_attention_mask(
+    model: Any,
+    *,
+    sample_lens: List[int],
+    split_lens: List[int],
+    attn_modes: List[str],
+    device: torch.device,
+    attention_backend: str,
+) -> Any:
+    """Build the train-replay mask while preserving identical visibility rules."""
+    backend = str(attention_backend).strip().lower()
+    if backend == "sdpa":
+        from .vendor.data.data_utils import prepare_attention_mask_per_sample
+
+        return [prepare_attention_mask_per_sample(split_lens, attn_modes, device=device)]
+    if backend == "flex":
+        try:
+            from torch.nn.attention.flex_attention import create_block_mask
+        except ImportError as exc:
+            raise RuntimeError(
+                "pack_und_forward_inputs: attention_backend='flex' requires PyTorch >= 2.5 "
+                "with torch.nn.attention.flex_attention."
+            ) from exc
+        from .vendor.data.data_utils import create_sparse_mask
+
+        seqlen = sum(sample_lens)
+        mask_mod = create_sparse_mask(sample_lens, split_lens, attn_modes, device)
+        return create_block_mask(
+            mask_mod,
+            B=1,
+            H=model.num_heads,
+            Q_LEN=seqlen,
+            KV_LEN=seqlen,
+            device=device,
+            BLOCK_SIZE=128,
+            _compile=True,
+        )
+    raise ValueError(f"pack_und_forward_inputs: attention_backend must be 'sdpa' or 'flex'; got {attention_backend!r}.")
+
+
 def pack_und_forward_inputs(
     model: Any,
     *,
     new_token_ids: Dict[str, Any],
     splits: List[Dict[str, Any]],
     response_input: torch.Tensor,
     device: torch.device,
+    attention_backend: str = "sdpa",
     vit_transform: Callable[[Any], Any] = lambda x: x,
 ) -> Dict[str, Any]:
-    """Train-mode packing: one und sample ``[*ordered splits | response_input]`` with a nested attention mask."""
-    from .vendor.data.data_utils import prepare_attention_mask_per_sample
-
+    """Pack one und sample and build either the baseline dense mask or Flex BlockMask."""
     text_ids: List[int] = []
     text_indexes: List[int] = []
     position_ids: List[int] = []
@@ -510,15 +549,24 @@ def _append_text_block(ids: List[int]) -> None:
     ce_loss_indexes = list(range(resp_start, resp_start + int(response_input.shape[0])))
 
     seqlen = pos
-    nested_mask = prepare_attention_mask_per_sample(split_lens, attn_modes, device=device)
+    attention_mask = _build_und_attention_mask(
+        model,
+        sample_lens=[seqlen],
+        split_lens=split_lens,
+        attn_modes=attn_modes,
+        device=device,
+        attention_backend=attention_backend,
+    )
 
     return {
         "seqlen": seqlen,
         "sample_lens": [seqlen],
         "packed_text_ids": torch.tensor(text_ids, dtype=torch.long, device=device),
         "packed_text_indexes": torch.tensor(text_indexes, dtype=torch.long, device=device),
         "packed_position_ids": torch.tensor(position_ids, dtype=torch.long, device=device),
-        "nested_attention_masks": [nested_mask],
+        # The vendored Navit attention dispatches List[Tensor] to SDPA and
+        # BlockMask to FlexAttention; retain the key for compatibility.
+        "nested_attention_masks": attention_mask,
         "packed_vit_tokens": (
             torch.cat(vit_tokens_parts, dim=0).to(device=device, dtype=model.dtype) if vit_tokens_parts else None
         ),