afmoe ring attention support.

CHANGELOG.md

@@ -62,5 +62,5 @@ Documenting changes which affect configuration usage patterns (added/moved/remov
 - **`inference.model.tool_call_parser`**: Changed default from `"hermes"` to auto-detection from model name. Uses `MODEL_TOOL_CALL_PARSER` dict to infer the correct vLLM tool call parser (e.g. Qwen3→`hermes`, GLM-4.5→`glm45`, GLM-4.7→`glm47`, MiniMax-M2→`minimax_m2`, INTELLECT-3→`hermes`). Unknown models default to `None`. Explicit values still take priority. (#1795, 2026-02-16)
 - **`orchestrator.eval.cancel_inflight_rollouts_on_eval`**: Added flag to optionally cancel in-flight training rollouts before starting online evals. When enabled, avoids congestion by preventing training and eval rollouts from running simultaneously, but slows training as the rollout pipeline must refill after each eval (default: False) (2026-02-16)
 - **`orchestrator.use_token_client`**: Added flag to use the token-in-token-out (TITO) client for training across all environments. When enabled, uses `openai_chat_completions_token` client type instead of `openai_chat_completions`. Only use when environments have linear history and the chat template has the extension property (default: False) (2026-02-21)
-- **`inference.enable_expert_parallel`**, **`inference.all2all_backend`**, and **`inference.enable_eplb`**: Added expert-parallel inference controls passed to vLLM as `--enable-expert-parallel`, `--all2all-backend`, and `--enable-eplb` (defaults: `False`, `"allgather_reducescatter"`, `False`) (2026-02-23)
-
+- **`model.cp` + AFMoE**: Context parallelism now works with AFMoE models via unified `substitute_ring_attn` which patches `_compute_attention` on both `FlashAttention` and `AfmoeFlashAttention` to use ring attention. Sliding window layers automatically get per-layer `window_size`; full attention layers default to `(-1, -1)`. Also plumbed `window_size` through the FA3 ring attention wrapper (`ring_fa3_varlen_func`). (2026-02-21)
+- **`inference.enable_expert_parallel`**, **`inference.all2all_backend`**, and **`inference.enable_eplb`**: Added expert-parallel inference controls passed to vLLM as `--enable-expert-parallel`, `--all2all-backend`, and `--enable-eplb` (defaults: `False`, `"allgather_reducescatter"`, `False`) (2026-0

src/prime_rl/trainer/models/afmoe/modeling_afmoe.py

@@ -189,6 +189,19 @@ def __init__(self, config: AfmoeAttentionConfig, flash_attn_version: int = 4):
         if self._flash_attn_version == 4:
             self._flash_attn_call = torch._dynamo.disable(self.func)
 
+    def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen):
+        """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim]."""
+        args = [q, k, v, cu_seqlens, cu_seqlens]
+        if self._flash_attn_version != 4:
+            args.extend([max_seqlen, max_seqlen])
+        kwargs: dict = {"causal": True}
+        if self.sliding_window is not None:
+            kwargs["window_size"] = (self.sliding_window - 1, 0)
+        out = self._flash_attn_call(*args, **kwargs)
+        if isinstance(out, tuple):
+            out = out[0]
+        return out
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -209,26 +222,14 @@ def forward(
         key_states = self.k_norm(key_states)
 
         if self.is_local_attention:
-            # apply_rotary_pos_emb expects [batch, heads, seq, dim]
             query_states = query_states.transpose(1, 2)
             key_states = key_states.transpose(1, 2)
             cos, sin = position_embeddings
             query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
             query_states = query_states.transpose(1, 2)
             key_states = key_states.transpose(1, 2)
 
-        # Flash attention varlen expects [total_tokens, heads, dim]
-        args = [query_states[0], key_states[0], value_states[0], cu_seqlens, cu_seqlens]
-        if self._flash_attn_version != 4:
-            args.extend([max_seqlen, max_seqlen])
-
-        kwargs: dict = {"causal": True}
-        if self.sliding_window is not None:
-            kwargs["window_size"] = (self.sliding_window - 1, 0)
-
-        out = self._flash_attn_call(*args, **kwargs)
-        if isinstance(out, tuple):
-            out = out[0]
+        out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen)
 
         attn_output = out.contiguous().view(*input_shape, -1)
         attn_output = attn_output * torch.sigmoid(gate_states)

src/prime_rl/trainer/models/layers/attn.py

@@ -91,6 +91,20 @@ def __init__(self, config: AttentionConfig, flash_attn_version: int = 2):
         if self._flash_attn_version == 4:
             self._flash_attn_call = torch._dynamo.disable(self.func)
 
+    def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen):
+        """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim]."""
+        args = [q, k, v, cu_seqlens, cu_seqlens]
+        if self._flash_attn_version != 4:
+            args.extend([max_seqlen, max_seqlen])
+        kwargs: dict = {"causal": True}
+        sliding_window = getattr(self, "sliding_window", None)
+        if sliding_window is not None:
+            kwargs["window_size"] = (sliding_window - 1, 0)
+        out = self._flash_attn_call(*args, **kwargs)
+        if isinstance(out, tuple):
+            out = out[0]
+        return out
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -129,22 +143,7 @@ def forward(
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)
 
-        args = [
-            query_states[0],
-            key_states[0],
-            value_states[0],
-            cu_seqlens,
-            cu_seqlens,
-        ]
-        if self._flash_attn_version != 4:
-            args.extend([max_seqlen, max_seqlen])
-
-        out = self._flash_attn_call(
-            *args,
-            causal=True,
-        )
-        if isinstance(out, tuple):
-            out = out[0]
+        out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen)
 
         out = out.contiguous()
         attn_output = out.view(1, out.shape[0], -1)
@@ -242,84 +241,47 @@ def forward(
 }
 
 
-def substitute_prime_rl_flash_attn(
+def substitute_ring_attn(
     process_group: torch.distributed.ProcessGroup,
     heads_k_stride: int,
     attn_impl: str = "flash_attention_2",
 ) -> None:
+    """Patch _compute_attention on FlashAttention (and AfmoeFlashAttention) to use ring attention."""
     from ring_flash_attn import llama3_flash_attn_varlen_func
 
     from .ring_attn import ring_fa3_varlen_func
 
     use_fa3 = attn_impl == "flash_attention_3"
     ring_func = ring_fa3_varlen_func if use_fa3 else llama3_flash_attn_varlen_func
 
-    class RingFlashAttention(FlashAttention):
-        def forward(
-            self,
-            hidden_states: torch.Tensor,
-            position_embeddings: tuple[torch.Tensor, torch.Tensor],
-            cu_seqlens: torch.LongTensor | None = None,
-            max_seqlen: int | None = None,
-        ) -> tuple[torch.Tensor, torch.Tensor | None]:
-            input_shape = hidden_states.shape[:-1]
-            hidden_shape = (*input_shape, -1, self.head_dim)
-
-            query_states = self.q_proj(hidden_states)
-            key_states = self.k_proj(hidden_states)
-            value_states = self.v_proj(hidden_states)
-
-            if self.use_qk_norm and self.qk_norm_type == "per_layer":
-                query_states = self.q_norm(query_states)
-                key_states = self.k_norm(key_states)
-
-            query_states = query_states.view(hidden_shape)
-            key_states = key_states.view(hidden_shape)
-            value_states = value_states.view(hidden_shape)
-
-            if self.use_qk_norm and self.qk_norm_type == "per_head":
-                query_states = self.q_norm(query_states)
-                key_states = self.k_norm(key_states)
-
-            query_states = query_states.transpose(1, 2)
-            key_states = key_states.transpose(1, 2)
-            value_states = value_states.transpose(1, 2)
-
-            cos, sin = position_embeddings
-            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-            from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS
-
-            cu_seqlens_q = DATA_PARAMS["cu_seqlens_q"]
-            cu_seqlens_k = DATA_PARAMS["cu_seqlens_k"]
-            max_seqlen_q = DATA_PARAMS["max_seqlen_q"]
-            max_seqlen_k = DATA_PARAMS["max_seqlen_k"]
-            local_k_slice = DATA_PARAMS["local_k_slice"]
-
-            # TODO: Can we optimize the rotary application instead of double transpose?
-            query_states = query_states.transpose(1, 2)
-            key_states = key_states.transpose(1, 2)
-            value_states = value_states.transpose(1, 2)
-            out = ring_func(
-                query_states[0],
-                key_states[0],
-                value_states[0],
-                cu_seqlens_q=cu_seqlens_q,
-                cu_seqlens_k=cu_seqlens_k,
-                max_seqlen_q=max_seqlen_q,
-                max_seqlen_k=max_seqlen_k,
-                local_k_slice=local_k_slice,
-                causal=True,
-                group=process_group,
-                heads_k_stride=heads_k_stride,
-            )
-            if isinstance(out, tuple):
-                out = out[0]
-            out = out.contiguous()
-            attn_output = out.view(1, out.shape[0], -1)
-            attn_weights = None
-
-            attn_output = self.o_proj(attn_output)
-            return attn_output, attn_weights
-
-    FlashAttention.forward = RingFlashAttention.forward
+    def _ring_compute_attention(self, q, k, v, cu_seqlens, max_seqlen):
+        from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS
+
+        window_size = (-1, -1)
+        sliding_window = getattr(self, "sliding_window", None)
+        if sliding_window is not None:
+            window_size = (sliding_window - 1, 0)
+
+        out = ring_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q=DATA_PARAMS["cu_seqlens_q"],
+            cu_seqlens_k=DATA_PARAMS["cu_seqlens_k"],
+            max_seqlen_q=DATA_PARAMS["max_seqlen_q"],
+            max_seqlen_k=DATA_PARAMS["max_seqlen_k"],
+            local_k_slice=DATA_PARAMS["local_k_slice"],
+            causal=True,
+            window_size=window_size,
+            group=process_group,
+            heads_k_stride=heads_k_stride,
+        )
+        if isinstance(out, tuple):
+            out = out[0]
+        return out
+
+    FlashAttention._compute_attention = _ring_compute_attention
+
+    from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention
+
+    AfmoeFlashAttention._compute_attention = _ring_compute_attention

src/prime_rl/trainer/models/layers/ring_attn.py

@@ -15,6 +15,7 @@ def _fa3_varlen_forward(
     max_seqlen_k: int,
     softmax_scale: float,
     causal: bool,
+    window_size: tuple[int, int] = (-1, -1),
 ) -> tuple[torch.Tensor, torch.Tensor]:
     from flash_attn_interface import _flash_attn_forward
 
@@ -32,6 +33,8 @@ def _fa3_varlen_forward(
             "causal": causal,
         }
     )
+    if "window_size" in params:
+        params["window_size"] = window_size
     out, lse, _, _ = _flash_attn_forward(**params)
     return out, lse
 
@@ -52,6 +55,7 @@ def _fa3_varlen_backward(
     dv: torch.Tensor,
     softmax_scale: float,
     causal: bool,
+    window_size: tuple[int, int] = (-1, -1),
 ) -> None:
     from flash_attn_interface import _flash_attn_backward
 
@@ -75,6 +79,8 @@ def _fa3_varlen_backward(
             "causal": causal,
         }
     )
+    if "window_size" in params:
+        params["window_size"] = window_size
     _flash_attn_backward(**params)
 
 
@@ -100,6 +106,8 @@ def forward(
         heads_k_stride: int,
         causal: bool,
         group_name: str,
+        window_size_left: int = -1,
+        window_size_right: int = -1,
     ) -> torch.Tensor:
         group = dist.group.WORLD
         for pg in dist.distributed_c10d._world.pg_map:
@@ -108,6 +116,7 @@ def forward(
                 break
 
         local_k_slice = slice(local_k_slice_start, local_k_slice_stop)
+        window_size = (window_size_left, window_size_right)
         softmax_scale = q.shape[-1] ** (-0.5)
         out_list = []
         lse_list = []
@@ -146,6 +155,7 @@ def forward(
                 max_seqlen_k=max_seqlen_k,
                 softmax_scale=softmax_scale,
                 causal=causal,
+                window_size=window_size,
             )
             out_list.append(out_i)
             lse_list.append(lse_i)
@@ -161,6 +171,7 @@ def forward(
         ctx.heads_k_stride = heads_k_stride
         ctx.causal = causal
         ctx.group_name = group_name
+        ctx.window_size = window_size
         return out
 
     @staticmethod
@@ -234,6 +245,7 @@ def backward(ctx, dout: torch.Tensor):
                 dv=dv_i,
                 softmax_scale=ctx.softmax_scale,
                 causal=causal,
+                window_size=ctx.window_size,
             )
 
             if heads_k_stride != nheads_k:
@@ -250,8 +262,9 @@ def backward(ctx, dout: torch.Tensor):
                 dv[:, i : i + heads_k_stride] = dv_i
 
         # Grads for: q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
-        #            local_k_slice_start, local_k_slice_stop, heads_k_stride, causal, group_name
-        return dq, dk, dv, None, None, None, None, None, None, None, None, None
+        #            local_k_slice_start, local_k_slice_stop, heads_k_stride, causal, group_name,
+        #            window_size_left, window_size_right
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None
 
 
 def ring_fa3_varlen_func(
@@ -266,6 +279,7 @@ def ring_fa3_varlen_func(
     causal: bool,
     heads_k_stride: int,
     group: dist.ProcessGroup,
+    window_size: tuple[int, int] = (-1, -1),
 ) -> torch.Tensor:
     return _RingFA3Varlen.apply(
         q,
@@ -280,4 +294,6 @@ def ring_fa3_varlen_func(
         heads_k_stride,
         causal,
         group.group_name,
+        window_size[0],
+        window_size[1],
     )

src/prime_rl/trainer/rl/train.py

@@ -5,7 +5,7 @@
 # Import environment before any other imports
 # ruff: noqa: I001
 
-from prime_rl.trainer.models.layers.attn import substitute_prime_rl_flash_attn
+from prime_rl.trainer.models.layers.attn import substitute_ring_attn
 from prime_rl.trainer.rl.broadcast import setup_weight_broadcast
 from prime_rl.utils.act_offloading import maybe_activation_offloading
 import torch
@@ -172,7 +172,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
 
     if parallel_dims.cp_enabled:
         substitute_hf_flash_attn(parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1)
-        substitute_prime_rl_flash_attn(
+        substitute_ring_attn(
             parallel_dims.world_mesh["cp"].get_group(),
             heads_k_stride=1,
             attn_impl=config.model.attn,

src/prime_rl/trainer/sft/train.py

@@ -8,7 +8,7 @@
 # Import environment before any other imports
 # ruff: noqa: I001
 
-from prime_rl.trainer.models.layers.attn import substitute_prime_rl_flash_attn
+from prime_rl.trainer.models.layers.attn import substitute_ring_attn
 from prime_rl.utils.act_offloading import maybe_activation_offloading
 import torch
 from torch.profiler import profile, ProfilerActivity, record_function
@@ -94,7 +94,7 @@ def train(config: SFTTrainerConfig):
     if parallel_dims.cp_enabled:
         assert config.data.seq_len % parallel_dims.cp == 0, "Sequence length must be divisible by CP degree"
         substitute_hf_flash_attn(parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1)
-        substitute_prime_rl_flash_attn(
+        substitute_ring_attn(
             parallel_dims.world_mesh["cp"].get_group(),
             heads_k_stride=1,
             attn_impl=config.model.attn,