From 39c282525609cdbf1a83ebeb1ee74f08b15b809c Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 5 Oct 2026 01:09:05 +0000
Subject: [PATCH 1/9] refactor: MFU the torchtitan way, relative packing cost

MFU: 6 * active params of the built model (routed experts weighted by
top_k / num_experts, embeddings only when tied, vision encoder and LoRA
adapters excluded) plus each layer's attention term from an
`attention_flops_per_token` hook on its attention / linear-attention
module, as torchtitan's get_nparams_and_flops.

Packing: relative per-bin cost 24*L*d^2 * sum(n) + L_attn*h*(qk+v) * sum(n^2)
instead of the Megatron-style estimator.
---
 src/prime_rl/orchestrator/batch.py            | 138 +++------------
 .../trainer/models/afmoe/modeling_afmoe.py    |  10 ++
 .../trainer/models/deepseek_v4/attention.py   |  22 +++
 .../glm_moe_dsa/sparse_mla_attention.py       |  14 ++
 .../trainer/models/gpt_oss/attention.py       |  11 +-
 src/prime_rl/trainer/models/layers/attn.py    |  32 ++++
 .../trainer/models/nemotron_h/mamba.py        |   5 +
 .../trainer/models/qwen3_5/gated_delta_net.py |   6 +
 .../models/qwen3_8_flash_next/attention.py    |  13 ++
 .../qwen3_8_flash_next/gated_delta_net.py     |   6 +
 src/prime_rl/trainer/perf.py                  | 160 ++++++------------
 tests/unit/train/test_perf.py                 |  36 ++--
 12 files changed, 195 insertions(+), 258 deletions(-)

diff --git a/src/prime_rl/orchestrator/batch.py b/src/prime_rl/orchestrator/batch.py
index 1417ed2bdea..69750926769 100644
--- a/src/prime_rl/orchestrator/batch.py
+++ b/src/prime_rl/orchestrator/batch.py
@@ -17,132 +17,32 @@
 STREAM_FILL = {"rl_weights": 1.0, "ce_weights": 0.0, "ref_kl_weights": 0.0}
 
 
-def _text_config(model_config: Any) -> Any:
-    """Unwrap a multimodal model's text config, else return the config unchanged."""
-    return getattr(model_config, "text_config", model_config)
-
-
-def _is_mla(config: Any) -> bool:
-    return bool(getattr(config, "multi_latent_attention", False) or hasattr(config, "q_lora_rank"))
-
-
-def _kv_channels(config: Any) -> int:
-    return getattr(config, "kv_channels", getattr(config, "head_dim", config.hidden_size // config.num_attention_heads))
-
-
-def _qkv_projection_flops_per_token(config: Any) -> int:
-    """Linear-in-seqlen FLOPs of the Q/K/V projections (per token)."""
-    hidden_size = config.hidden_size
-    num_attention_heads = config.num_attention_heads
-    kv_channels = _kv_channels(config)
-    is_mla = _is_mla(config)
-    qk_head_dim = getattr(config, "qk_head_dim", 0)
-    qk_pos_emb_head_dim = getattr(config, "qk_pos_emb_head_dim", 0)
-
-    if is_mla and getattr(config, "q_lora_rank", None) is not None:
-        q_flops = 2 * config.q_lora_rank * (hidden_size + num_attention_heads * (qk_head_dim + qk_pos_emb_head_dim))
-    else:
-        q_head_dim = (qk_head_dim + qk_pos_emb_head_dim) if is_mla else kv_channels
-        q_flops = 2 * hidden_size * num_attention_heads * q_head_dim
-
-    if is_mla and getattr(config, "kv_lora_rank", None) is not None:
-        v_head_dim = getattr(config, "v_head_dim", 0)
-        kv_flops = 2 * (
-            config.kv_lora_rank * (hidden_size + num_attention_heads * (qk_head_dim + v_head_dim))
-            + hidden_size * qk_pos_emb_head_dim
-        )
-    else:
-        num_query_groups = getattr(
-            config, "num_query_groups", getattr(config, "num_key_value_heads", num_attention_heads)
-        )
-        kv_flops = 4 * hidden_size * num_query_groups * kv_channels
-    return q_flops + kv_flops
-
-
-def _attention_flops_per_token_squared(config: Any) -> int:
-    """Quadratic-in-seqlen FLOPs of the attention scores and values (per token^2)."""
-    num_attention_heads = config.num_attention_heads
-    kv_channels = _kv_channels(config)
-    if _is_mla(config):
-        qk_head_dim = getattr(config, "qk_head_dim", 0)
-        qk_pos_emb_head_dim = getattr(config, "qk_pos_emb_head_dim", 0)
-        v_head_dim = getattr(config, "v_head_dim", kv_channels)
-        return num_attention_heads * (qk_head_dim + qk_pos_emb_head_dim) + num_attention_heads * v_head_dim
-    return 2 * num_attention_heads * kv_channels
-
-
-def _ffn_flops_per_token(hidden_size: int, ffn_hidden_size: int) -> int:
-    return 6 * hidden_size * ffn_hidden_size
-
-
-def _dense_ffn_hidden_size(config: Any) -> int:
-    return getattr(
-        config, "ffn_hidden_size", getattr(config, "intermediate_size", getattr(config, "moe_intermediate_size", 0))
-    )
-
-
-def _moe_ffn_hidden_size(config: Any) -> int:
-    """Effective FFN width of an MoE layer: routed experts (topk) plus shared experts."""
-    dense_ffn = _dense_ffn_hidden_size(config)
-    routed_topk = getattr(config, "moe_router_topk", getattr(config, "num_experts_per_tok", 1))
-    moe_ffn = getattr(config, "moe_ffn_hidden_size", getattr(config, "moe_intermediate_size", dense_ffn))
-    return moe_ffn * routed_topk + (getattr(config, "moe_shared_expert_intermediate_size", None) or 0)
-
-
-def _count_dense_and_moe_layers(config: Any) -> tuple[int, int]:
-    """Split the model's layers into (dense, moe) counts."""
-    num_experts = getattr(config, "num_experts", getattr(config, "n_routed_experts", None))
-    if num_experts is None:
-        return config.num_hidden_layers, 0
-
-    moe_layer_freq = getattr(config, "moe_layer_freq", None)
-    if isinstance(moe_layer_freq, list):
-        num_dense = sum(1 for freq in moe_layer_freq if freq == 0)
-        num_moe = sum(1 for freq in moe_layer_freq if freq > 0)
-    elif isinstance(moe_layer_freq, int):
-        num_dense = sum(1 for i in range(config.num_hidden_layers) if i % moe_layer_freq != 0)
-        num_moe = config.num_hidden_layers - num_dense
-    elif getattr(config, "first_k_dense_replace", None) is not None:
-        num_dense = config.first_k_dense_replace
-        num_moe = config.num_hidden_layers - num_dense
-    else:
-        num_dense = 0
-        num_moe = config.num_hidden_layers
-    return num_dense, num_moe
-
-
-def _packing_cost_coeffs(config: Any) -> tuple[int, int]:
-    """Return the (linear, quadratic) coefficients of the per-sequence forward FLOPs."""
-    hidden_size = config.hidden_size
-    num_dense_layers, num_moe_layers = _count_dense_and_moe_layers(config)
-    qkv_per_token = _qkv_projection_flops_per_token(config)
-    attn_per_token_squared = _attention_flops_per_token_squared(config)
-
-    def layer_linear(ffn_hidden_size: int) -> int:
-        return qkv_per_token + 2 * hidden_size * hidden_size + _ffn_flops_per_token(hidden_size, ffn_hidden_size)
-
-    linear = (
-        num_dense_layers * layer_linear(_dense_ffn_hidden_size(config))
-        + num_moe_layers * layer_linear(_moe_ffn_hidden_size(config))
-        + 2 * hidden_size * config.vocab_size
-    )
-    quadratic = (num_dense_layers + num_moe_layers) * attn_per_token_squared
-    return linear, quadratic
-
-
 def build_bin_cost(model_config: Any | None) -> Callable[[Sequence[int]], int]:
-    """Build a closure scoring a packed bin by estimated forward compute.
+    """Build a closure scoring a packed bin by its relative forward compute.
 
-    With ``model_config=None`` the linear/quadratic coefficients are ``(1, 0)``, so
-    the cost reduces to the token count and balancing falls back to sequence length.
+    Per token, ~24 * layers * hidden_size^2 FLOPs (2 FLOPs per matmul parameter, ~12 * hidden_size^2
+    parameters per layer). Per causal query-key pair, heads * (qk_head_dim + v_head_dim) FLOPs per
+    full-attention layer, i.e. that times n^2 for a document of n tokens. Only the ratio of the two
+    matters for balancing. With ``model_config=None`` the cost is the token count.
     """
     if model_config is None:
-        linear, quadratic = 1, 0
+        token_cost, pair_cost = 1, 0
     else:
-        linear, quadratic = _packing_cost_coeffs(_text_config(model_config))
+        config = getattr(model_config, "text_config", model_config)
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        qk_head_dim = getattr(config, "qk_nope_head_dim", 0) + getattr(config, "qk_rope_head_dim", 0) or head_dim
+        v_head_dim = getattr(config, "v_head_dim", None) or head_dim
+        layer_types = getattr(config, "layer_types", None)
+        num_attention_layers = (
+            sum(layer_type in ("full_attention", "attention") for layer_type in layer_types)
+            if layer_types
+            else config.num_hidden_layers
+        )
+        token_cost = 24 * config.num_hidden_layers * config.hidden_size**2
+        pair_cost = num_attention_layers * config.num_attention_heads * (qk_head_dim + v_head_dim)
 
     def bin_cost(seqlens: Sequence[int]) -> int:
-        return linear * sum(seqlens) + quadratic * sum(n * n for n in seqlens)
+        return token_cost * sum(seqlens) + pair_cost * sum(n * n for n in seqlens)
 
     return bin_cost
 
diff --git a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
index eadf1dbd80d..8556dd825ef 100644
--- a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
+++ b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
@@ -18,6 +18,7 @@
     flash_attn_2_varlen_op,
     flash_attn_3_varlen_op,
     flash_attn_4_varlen_op,
+    quadratic_attention_flops_per_token,
 )
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
@@ -82,6 +83,15 @@ def __init__(self, config: AfmoeAttentionConfig):
         self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps))
         self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps))
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        return quadratic_attention_flops_per_token(
+            num_heads=self.num_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=seq_len,
+            sliding_window=self.sliding_window,
+        )
+
 
 class AfmoeFlashAttention(AfmoeAttentionBase):
     """AFMoE attention using Flash Attention varlen functions."""
diff --git a/src/prime_rl/trainer/models/deepseek_v4/attention.py b/src/prime_rl/trainer/models/deepseek_v4/attention.py
index 9dfd91deeb3..e4d335aecc6 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/attention.py
+++ b/src/prime_rl/trainer/models/deepseek_v4/attention.py
@@ -134,6 +134,7 @@
 from prime_rl.trainer.models.kernels.deepseek_v4 import IGNORE_SLOT
 from prime_rl.trainer.models.kernels.deepseek_v4.dsv4_rope import dsv4_q_norm_rope, dsv4_rope
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
+from prime_rl.trainer.models.layers.attn import quadratic_attention_flops_per_token
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
 from prime_rl.utils.cp import CPContext, gather_for_cp
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
@@ -705,6 +706,27 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int, rotary_emb: Deepsee
 
         self.cp_context = CPContext()
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        # As torchtitan's DeepSeek V4: the local window, plus the compressed entries (CSA: the
+        # indexer scores all of them, attention reads `index_topk`).
+        flops = quadratic_attention_flops_per_token(
+            num_heads=self.num_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=seq_len,
+            sliding_window=self.config.sliding_window,
+        )
+        if self.compressor is not None:
+            compressed_len = seq_len // self.compressor.compress_rate
+            indexer = getattr(self.compressor, "indexer", None)
+            if indexer is not None:
+                flops += 6 * indexer.num_heads * indexer.head_dim * compressed_len
+                compressed_len = min(compressed_len, indexer.index_topk)
+            flops += quadratic_attention_flops_per_token(
+                num_heads=self.num_heads, qk_head_dim=self.head_dim, v_head_dim=self.head_dim, seq_len=compressed_len
+            )
+        return flops
+
     def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[torch.Tensor, None]:
         """`packed` carries the document boundaries every pathway below is clipped at."""
         # Shape keys in the comments below:
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 30ceea86b84..18ac67353d3 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -6,6 +6,7 @@
 
 from prime_rl.configs.trainer import DSABackend
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
+from prime_rl.trainer.models.layers.attn import quadratic_attention_flops_per_token
 from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import rotate_half
 from prime_rl.utils.cp import CPContext, gather_for_cp
@@ -149,6 +150,19 @@ def __init__(self, args: SparseMlaAttentionArgs):
 
         self.cp_context = CPContext()
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        # As torchtitan's DeepSeek V3/V4: MLA counted un-absorbed (qk = nope + rope), sparse attention
+        # attends `index_topk` tokens, and the indexer scores every token with its own heads.
+        flops = quadratic_attention_flops_per_token(
+            num_heads=self.num_heads,
+            qk_head_dim=self.qk_head_dim,
+            v_head_dim=self.v_head_dim,
+            seq_len=min(seq_len, self.args.index_topk),
+        )
+        if self.indexer is not None:
+            flops += 6 * self.args.index_n_heads * self.args.index_head_dim * seq_len
+        return flops
+
     def mla_latents(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         q_latent = self.q_a_layernorm(self.q_a_proj(hidden_states))
         compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
diff --git a/src/prime_rl/trainer/models/gpt_oss/attention.py b/src/prime_rl/trainer/models/gpt_oss/attention.py
index 73cf2a6db2a..2f5d344a42a 100644
--- a/src/prime_rl/trainer/models/gpt_oss/attention.py
+++ b/src/prime_rl/trainer/models/gpt_oss/attention.py
@@ -10,7 +10,7 @@
 
 from prime_rl.trainer.models.fusions import fuse_qkv_projections
 from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig
-from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func
+from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func, quadratic_attention_flops_per_token
 from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb
 
 
@@ -51,6 +51,15 @@ def __init__(self, config: GptOssConfig, layer_idx: int) -> None:
         nn.init.normal_(self.sinks, mean=0.0, std=config.initializer_range)
         self.flash_attn = flash_attn_4_varlen_func
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        return quadratic_attention_flops_per_token(
+            num_heads=self.num_attention_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=seq_len,
+            sliding_window=self.sliding_window,
+        )
+
     def compute_attention(
         self,
         query: torch.Tensor,
diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py
index 6025bed458a..1859dccbf43 100644
--- a/src/prime_rl/trainer/models/layers/attn.py
+++ b/src/prime_rl/trainer/models/layers/attn.py
@@ -404,6 +404,28 @@ def flash_attn_3_varlen_op(
     return out
 
 
+def quadratic_attention_flops_per_token(
+    *, num_heads: int, qk_head_dim: int, v_head_dim: int, seq_len: int, sliding_window: int | None = None
+) -> int:
+    """Training FLOPs per token of softmax attention, as torchtitan's `quadratic_attention_flops_per_token`.
+
+    Two matmuls per head and attended token (q @ K^T over `qk_head_dim`, scores @ V over `v_head_dim`),
+    2 FLOPs per multiply-add, 3x for forward plus backward. Causal sparsity and the flash-attention
+    recompute are not counted.
+    """
+    attended_tokens = seq_len if sliding_window is None else min(seq_len, sliding_window)
+    return 6 * num_heads * (qk_head_dim + v_head_dim) * attended_tokens
+
+
+def delta_rule_flops_per_token(*, num_heads: int, key_head_dim: int, v_head_dim: int) -> int:
+    """Training FLOPs per token of a gated delta-rule recurrence, as torchtitan's `delta_rule_flops_per_token`.
+
+    Reading, updating and querying the `[key_head_dim, v_head_dim]` state are three
+    `key_head_dim * v_head_dim` products per head.
+    """
+    return 6 * 3 * num_heads * key_head_dim * v_head_dim
+
+
 @dataclass
 class AttentionConfig:
     hidden_size: int
@@ -437,6 +459,7 @@ class FlashAttention(nn.Module):
     def __init__(self, config: AttentionConfig, flash_attn_version: int = 2):
         super().__init__()
         self.head_dim = config.head_dim
+        self.num_heads = config.num_attention_heads
         self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
         self.scaling = self.head_dim**-0.5
         self.is_causal = config.is_causal
@@ -468,6 +491,15 @@ def __init__(self, config: AttentionConfig, flash_attn_version: int = 2):
         self._flash_attn_version = flash_attn_version
         self.func = self._funcs[flash_attn_version]
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        return quadratic_attention_flops_per_token(
+            num_heads=self.num_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=seq_len,
+            sliding_window=getattr(self, "sliding_window", None),
+        )
+
     def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Query, key and value projections, from one packed GEMM when qkv is fused."""
         if self.qkv_proj is None:
diff --git a/src/prime_rl/trainer/models/nemotron_h/mamba.py b/src/prime_rl/trainer/models/nemotron_h/mamba.py
index 572ac9e257b..dc3b775e77a 100644
--- a/src/prime_rl/trainer/models/nemotron_h/mamba.py
+++ b/src/prime_rl/trainer/models/nemotron_h/mamba.py
@@ -127,6 +127,11 @@ def __init__(self, config: NemotronHConfig) -> None:
 
         self.cp_context = CPContext()
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        # Mamba2's analogue of torchtitan's `delta_rule_flops_per_token`: per head, the state update
+        # (x outer B) and the readout (state @ C) are two `head_dim * state_size` products.
+        return 6 * 2 * self.num_heads * self.head_dim * self.state_size
+
     def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
         batch_size, sequence_length, _ = hidden_states.shape
         projected_states = self.in_proj(hidden_states)
diff --git a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
index 60d690bdb0c..cb089e8c2ee 100644
--- a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
@@ -6,6 +6,7 @@
 from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 from torch import nn
 
+from prime_rl.trainer.models.layers.attn import delta_rule_flops_per_token
 from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
 from prime_rl.utils.cp import CPContext
 
@@ -52,6 +53,11 @@ def __init__(self, config: Qwen3_5TextConfig) -> None:
         self.in_proj_a = nn.Linear(config.hidden_size, self.num_value_heads, bias=False)
         self.cp_context = CPContext()
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        return delta_rule_flops_per_token(
+            num_heads=self.num_value_heads, key_head_dim=self.key_head_dim, v_head_dim=self.value_head_dim
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
index 1666d75dfd1..00a3b1ac2a3 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
@@ -1,6 +1,7 @@
 import torch
 from torch import nn
 
+from prime_rl.trainer.models.layers.attn import quadratic_attention_flops_per_token
 from prime_rl.trainer.models.qwen3_8_flash_next.indexer import SparseAttentionIndexer
 from prime_rl.trainer.models.qwen3_8_flash_next.norm import RMSNorm
 from prime_rl.trainer.models.qwen3_8_flash_next.rotary_embedding import apply_rotary_embedding
@@ -48,6 +49,18 @@ def __init__(
         # Resolve once before torch.compile traces the layer; the loader caches the module.
         self.indexed_attention = prime_kernels.load("indexed_attention").indexed_attention
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        # As torchtitan's DeepSeek V4 CSA: the indexer scores every compressed block,
+        # attention reads `token_budget` tokens.
+        indexer = self.indexer
+        indexer_flops = 6 * indexer.num_query_heads * indexer.head_dim * (seq_len // indexer.compression_ratio)
+        return indexer_flops + quadratic_attention_flops_per_token(
+            num_heads=self.num_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=min(seq_len, indexer.token_budget),
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
index 214919e6a1f..0164e369f6b 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
@@ -6,6 +6,7 @@
 from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 from torch import nn
 
+from prime_rl.trainer.models.layers.attn import delta_rule_flops_per_token
 from prime_rl.utils.cp import CPContext
 
 # FLA's context carries a process group that Dynamo cannot trace through the convolution.
@@ -56,6 +57,11 @@ def __init__(
 
         self.cp_context = CPContext()
 
+    def attention_flops_per_token(self, seq_len: int) -> int:
+        return delta_rule_flops_per_token(
+            num_heads=self.num_value_heads, key_head_dim=self.key_head_dim, v_head_dim=self.value_head_dim
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/src/prime_rl/trainer/perf.py b/src/prime_rl/trainer/perf.py
index 13ba65d6b99..e08842a6db7 100644
--- a/src/prime_rl/trainer/perf.py
+++ b/src/prime_rl/trainer/perf.py
@@ -1,15 +1,43 @@
 import time
+from fractions import Fraction
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
 
 from prime_rl.trainer.lora import has_lora_layers
 from prime_rl.trainer.models.layers.lora import MultiLoRAModule
+from prime_rl.trainer.models.layers.moe import MoE
 from prime_rl.trainer.world import get_world
 from prime_rl.utils.logger import get_logger
 
 
+def get_active_params(model: nn.Module) -> int:
+    """Matmul parameters used per token, as torchtitan's `get_nparams_and_active_nparams`.
+
+    Routed experts count by `top_k / num_experts`; embeddings count only when tied to the lm_head;
+    the vision encoder and LoRA adapters do not count.
+    """
+    weights = {id(param): Fraction(1) for param in model.parameters()}
+    for module in model.modules():
+        if isinstance(module, MoE):
+            for param in module.experts.parameters():
+                weights[id(param)] = Fraction(module.router.top_k, module.experts.num_experts)
+    lm_head_params = {id(param) for param in model.lm_head.parameters()}
+    for module in model.modules():
+        if isinstance(module, nn.Embedding):
+            for param in module.parameters(recurse=False):
+                if id(param) not in lm_head_params:
+                    weights[id(param)] = Fraction(0)
+    visual = getattr(model.model, "visual", None)
+    if visual is not None:
+        for param in visual.parameters():
+            weights[id(param)] = Fraction(0)
+    for name, param in model.named_parameters():
+        if "lora_A" in name or "lora_B" in name:
+            weights[id(param)] = Fraction(0)
+    return int(sum(param.numel() * weights[id(param)] for param in model.parameters()))
+
+
 class PerfCounter:
     """
     Computes throughput (tokens/s) and MFU.
@@ -37,7 +65,7 @@ def __init__(self, model: nn.Module, seq_len: int, window_size: int = 10):
             self.gpu_peak_flops = self._get_peak_flops(torch.cuda.get_device_name(torch.device("cuda")))
         else:
             self.gpu_peak_flops = 0
-        self.num_flop_per_token = self._get_num_flop_per_token(model.config, seq_len=seq_len)
+        self.num_flop_per_token = self._get_num_flop_per_token(seq_len)
 
     def count_tokens(self, tokens: int) -> None:
         """Push a step into the sliding window. Time is recorded internally."""
@@ -108,117 +136,25 @@ def _get_peak_flops(self, device_name: str) -> float:
             self._logger.warning(f"Peak FLOPS undefined for `{device_name}`. Falling back to A100 (312 TFLOPS)")
             return 312e12
 
-    @staticmethod
-    def get_active_mm_params(config: PretrainedConfig) -> float:
-        """Get number of active parameters per token involved in matmuls"""
-        # Handle VLM models with nested text_config (e.g., Qwen3-VL)
-        if hasattr(config, "text_config"):
-            config = config.text_config
-
-        vocab_size = config.vocab_size
-        hidden_size = config.hidden_size
-        intermediate_size = getattr(config, "intermediate_size", getattr(config, "moe_intermediate_size", 0))
-        head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-        num_attention_heads = config.num_attention_heads
-        num_hidden_layers = config.num_hidden_layers
-
-        ## Attention
-        if hasattr(config, "q_lora_rank") and hasattr(config, "kv_lora_rank"):
-            # MLA
-            q_params = num_hidden_layers * (
-                hidden_size * config.q_lora_rank + config.q_lora_rank * num_attention_heads * config.qk_head_dim
-            )
-            kv_params = num_hidden_layers * (
-                hidden_size * (config.kv_lora_rank + config.qk_rope_head_dim)
-                + config.kv_lora_rank * num_attention_heads * (config.qk_nope_head_dim + config.v_head_dim)
-            )
-            o_params = num_hidden_layers * (num_attention_heads * config.v_head_dim * hidden_size)
-        else:
-            # GQA
-            num_key_value_heads = config.num_key_value_heads
-            q_params = num_hidden_layers * hidden_size * num_attention_heads * head_dim
-            kv_params = 2 * num_hidden_layers * hidden_size * num_key_value_heads * head_dim
-            o_params = num_hidden_layers * hidden_size * num_attention_heads * head_dim
-
-        ## MLP
-        if hasattr(config, "first_k_dense_replace"):
-            num_dense_layers = config.first_k_dense_replace
-            num_sparse_layers = config.num_hidden_layers - num_dense_layers
-        elif hasattr(config, "num_experts_per_tok"):
-            num_dense_layers = 0
-            num_sparse_layers = config.num_hidden_layers
-        else:
-            num_dense_layers = config.num_hidden_layers
-            num_sparse_layers = 0
-
-        dense_mlp_params = num_dense_layers * 3 * intermediate_size * hidden_size
-        sparse_mlp_params = 0
-
-        # Some MoE models (e.g. DeepSeek) use moe_intermediate_size, others (e.g. Granite) just use intermediate_size
-        moe_intermediate_size = getattr(config, "moe_intermediate_size", None) or intermediate_size
-        if hasattr(config, "num_shared_experts") and config.num_shared_experts:  # Shared experts
-            sparse_mlp_params += num_sparse_layers * config.num_shared_experts * 3 * moe_intermediate_size * hidden_size
-        if hasattr(config, "num_experts_per_tok") and config.num_experts_per_tok:  # Routed experts
-            sparse_mlp_params += (
-                num_sparse_layers * config.num_experts_per_tok * 3 * moe_intermediate_size * hidden_size
-            )
-        if hasattr(config, "n_routed_experts"):  # DeepSeek Router
-            sparse_mlp_params += num_sparse_layers * config.n_routed_experts * hidden_size
-        elif hasattr(config, "num_experts") and config.num_experts is not None:  # Qwen Router
-            sparse_mlp_params += num_sparse_layers * config.num_experts * hidden_size
-        else:
-            sparse_mlp_params = 0
-
-        ## LM Head
-        lm_head_params = vocab_size * hidden_size
-        ## Total
-        return q_params + kv_params + o_params + dense_mlp_params + sparse_mlp_params + lm_head_params
-
-    def _get_num_flop_per_token(self, model_config: PretrainedConfig, seq_len: int) -> int:
-        # Handle VLM models with nested text_config (e.g., Qwen3-VL)
-        if hasattr(model_config, "text_config"):
-            model_config = model_config.text_config
-
-        l, h, t = (  # noqa: E741
-            model_config.num_hidden_layers,
-            model_config.num_attention_heads,
-            seq_len,
+    def _get_num_flop_per_token(self, seq_len: int) -> int:
+        """Training FLOPs per token, as torchtitan's `get_nparams_and_flops`: 6 FLOPs per active matmul
+        parameter plus each layer's own attention term (`attention_flops_per_token` hooks)."""
+        attention_flops = sum(
+            module.attention_flops_per_token(seq_len)
+            for module in self.model.modules()
+            if hasattr(module, "attention_flops_per_token")
+        )
+        active_params = get_active_params(self.model)
+        if not has_lora_layers(self.model):
+            return 6 * active_params + attention_flops
+        # LoRA: frozen base matmuls still cost forward (2x) and dX (2x); fully trainable non-LoRA
+        # params (modules_to_save) add their weight gradient (2x); adapters cost 6x.
+        return (
+            4 * active_params
+            + 2 * self._count_fully_trainable_params_excluding_lora()
+            + 6 * self._count_lora_adapter_params()
+            + attention_flops
         )
-        # Head dims as torchtitan's quadratic_attention_flops_per_token: the real head_dim (e.g. 128 for
-        # Qwen3-235B, whose hidden_size / num_attention_heads is 64), or the MLA qk / v head dims.
-        if hasattr(model_config, "qk_head_dim") and hasattr(model_config, "v_head_dim"):
-            qk_head_dim, v_head_dim = model_config.qk_head_dim, model_config.v_head_dim
-        else:
-            qk_head_dim = v_head_dim = (
-                getattr(model_config, "head_dim", None) or model_config.hidden_size // model_config.num_attention_heads
-            )
-        # Reasoning behind the factor of 6 for the self-attention part of the formula:
-        # 1. each self-attention has 2 matmul in the forward and 4 in the backward (6)
-        #    (q @ K^T over qk_head_dim, then scores @ V over v_head_dim, per head and attended token)
-        # 2. the flash attention does 1 more matmul recomputation in the backward
-        #    but recomputation should not be counted in calculating MFU           (+0)
-        # 3. each matmul performs 1 multiplication and 1 addition                 (*2)
-        # 4. we follow the convention and do not account for sparsity in causal attention
-        attention_flops = 6 * l * h * (qk_head_dim + v_head_dim) * t
-
-        if has_lora_layers(self.model):
-            # LoRA case:
-            # - Frozen base matmuls still incur dX in backward: 2×, plus forward: 2× => 4× active_mm
-            # - Fully trainable non-LoRA params (modules_to_save) cost 6×
-            # - LoRA adapter params cost 6×
-            # Combined (to avoid double counting): 4*active_mm + 2*fully_trainable + 6*lora_adapters + attention
-            active_mm_params = self.get_active_mm_params(model_config)
-            lora_adapter_params = self._count_lora_adapter_params()
-            fully_trainable_params = self._count_fully_trainable_params_excluding_lora()
-
-            flop_per_token = (
-                4 * active_mm_params + 2 * fully_trainable_params + 6 * lora_adapter_params + attention_flops
-            )
-        else:
-            # standard case: full fine-tuning, all params participate in forward (2×) and backward (4×)
-            flop_per_token = 6 * self.get_active_mm_params(model_config) + attention_flops
-
-        return flop_per_token
 
     def _count_lora_adapter_params(self) -> int:
         """Count LoRA adapter parameters (sum of lora_A and lora_B across all MultiLoRAModules)."""
diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py
index ea81d58968a..b3feb1002b0 100644
--- a/tests/unit/train/test_perf.py
+++ b/tests/unit/train/test_perf.py
@@ -1,35 +1,19 @@
 import pytest
 import torch
-from transformers import AutoConfig, AutoModelForCausalLM, GenerationConfig
+from transformers import AutoConfig
 
+from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL
 from prime_rl.trainer.perf import PerfCounter
 
 
+# Expected values are torchtitan's `get_nparams_and_flops(model, seq_len=1024)` for its matching flavors.
 @pytest.mark.parametrize(
-    "model_name, active_params, flops_per_token",
-    [("Qwen/Qwen3-0.6B", 595_984_384, 4_280_549_376), ("Jackmin108/debug-moe-0.5B", 256_442_368, 1_840_644_096)],
+    "model_name, flops_per_token",
+    [("Qwen/Qwen3-0.6B", 4_280_942_592), ("Qwen/Qwen3-30B-A3B", 20_667_125_760)],
 )
-def test_perf_counter(model_name: str, active_params: int, flops_per_token: int):
-    # This speeds up the model loading as its a fake device
+def test_flops_per_token_matches_torchtitan(model_name: str, flops_per_token: int):
+    config = AutoConfig.from_pretrained(model_name, attn_implementation="flash_attention_2")
+    config.pad_token_id = config.eos_token_id
     with torch.device("meta"):
-        config = AutoConfig.from_pretrained(model_name)
-        if getattr(config, "pad_token_id", None) is None:
-            gen_config = GenerationConfig.from_model_config(config)
-            # Use `is not None` instead of truthiness: token ID 0 is valid.
-            config.pad_token_id = next(
-                (
-                    v
-                    for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(config, "eos_token_id", None)]
-                    if v is not None
-                ),
-                None,
-            )
-        model = AutoModelForCausalLM.from_config(config)
-    perf_counter = PerfCounter(model, seq_len=1024)
-
-    assert perf_counter.get_active_mm_params(config) == active_params, (
-        f"Expected {active_params:,} active parameters, got {perf_counter.get_active_mm_params(config):,} active parameters"
-    )
-    assert perf_counter.num_flop_per_token == flops_per_token, (
-        f"Expected {flops_per_token:,} FLOPS per token, got {perf_counter.num_flop_per_token:,} FLOPS per token"
-    )
+        model = AutoModelForCausalLMPrimeRL.from_config(config)
+    assert PerfCounter(model, seq_len=1024).num_flop_per_token == flops_per_token

From 9f28e751f981ac41b90dda658617207415b89974 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 6 Oct 2026 17:16:19 +0000
Subject: [PATCH 2/9] refactor(perf): get_active_params in one module walk,
 integer expert share

---
 src/prime_rl/trainer/perf.py | 29 ++++++++++++-----------------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/src/prime_rl/trainer/perf.py b/src/prime_rl/trainer/perf.py
index e08842a6db7..1eb3b39885b 100644
--- a/src/prime_rl/trainer/perf.py
+++ b/src/prime_rl/trainer/perf.py
@@ -1,5 +1,4 @@
 import time
-from fractions import Fraction
 
 import torch
 from torch import nn
@@ -17,25 +16,21 @@ def get_active_params(model: nn.Module) -> int:
     Routed experts count by `top_k / num_experts`; embeddings count only when tied to the lm_head;
     the vision encoder and LoRA adapters do not count.
     """
-    weights = {id(param): Fraction(1) for param in model.parameters()}
-    for module in model.modules():
-        if isinstance(module, MoE):
-            for param in module.experts.parameters():
-                weights[id(param)] = Fraction(module.router.top_k, module.experts.num_experts)
+    excluded = {id(param) for name, param in model.named_parameters() if "lora_A" in name or "lora_B" in name}
+    visual = getattr(model.model, "visual", None)
+    if visual is not None:
+        excluded |= {id(param) for param in visual.parameters()}
     lm_head_params = {id(param) for param in model.lm_head.parameters()}
+    inactive_expert_params = 0
     for module in model.modules():
         if isinstance(module, nn.Embedding):
-            for param in module.parameters(recurse=False):
-                if id(param) not in lm_head_params:
-                    weights[id(param)] = Fraction(0)
-    visual = getattr(model.model, "visual", None)
-    if visual is not None:
-        for param in visual.parameters():
-            weights[id(param)] = Fraction(0)
-    for name, param in model.named_parameters():
-        if "lora_A" in name or "lora_B" in name:
-            weights[id(param)] = Fraction(0)
-    return int(sum(param.numel() * weights[id(param)] for param in model.parameters()))
+            excluded |= {id(param) for param in module.parameters(recurse=False)} - lm_head_params
+        elif isinstance(module, MoE):
+            # Expert weights lead with `num_experts`, so the division is exact.
+            num_experts, top_k = module.experts.num_experts, module.router.top_k
+            expert_params = sum(p.numel() for p in module.experts.parameters() if id(p) not in excluded)
+            inactive_expert_params += expert_params * (num_experts - top_k) // num_experts
+    return sum(param.numel() for param in model.parameters() if id(param) not in excluded) - inactive_expert_params
 
 
 class PerfCounter:

From bc7f98e25c0632dfb8edba93d439d027cc37b86f Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 13 Oct 2026 10:25:37 +0000
Subject: [PATCH 3/9] fix(perf): attention FLOPs per packed document, not per
 micro-batch length

---
 src/prime_rl/trainer/perf.py      | 45 +++++++++++++------------------
 src/prime_rl/trainer/rl/train.py  | 13 ++++-----
 src/prime_rl/trainer/sft/train.py | 12 +++++----
 tests/unit/train/test_perf.py     | 11 +++++++-
 4 files changed, 43 insertions(+), 38 deletions(-)

diff --git a/src/prime_rl/trainer/perf.py b/src/prime_rl/trainer/perf.py
index 1eb3b39885b..0a966dc2d7f 100644
--- a/src/prime_rl/trainer/perf.py
+++ b/src/prime_rl/trainer/perf.py
@@ -39,15 +39,15 @@ class PerfCounter:
 
     Two modes:
     - Sliding window over full-step wall time: `count_tokens(tokens)` each step,
-      read smoothed values via `get_tokens_per_second()` / `get_mfu()`.
+      read smoothed values via `get_tokens_per_second()`.
       Time is measured between successive `count_tokens` calls (full step).
     - Single point on a caller-provided duration: `get_step_tokens_per_second(tokens, fwd_bwd_time)`
-      / `get_step_mfu(tokens, fwd_bwd_time)`. No smoothing.
+      / `get_step_mfu(flops, fwd_bwd_time)`, with `flops` summed from `get_flops(seq_lens)`. No smoothing.
 
     Sliding window inspired by https://github.com/pytorch/torchtitan/blob/4b3f2e41a084bf79a8540068ed525539d1244edd/torchtitan/utils.py#L119
     """
 
-    def __init__(self, model: nn.Module, seq_len: int, window_size: int = 10):
+    def __init__(self, model: nn.Module, window_size: int = 10):
         self.window_size = window_size
         self.tokens: list[int] = []
         self.times: list[float] = []
@@ -60,7 +60,8 @@ def __init__(self, model: nn.Module, seq_len: int, window_size: int = 10):
             self.gpu_peak_flops = self._get_peak_flops(torch.cuda.get_device_name(torch.device("cuda")))
         else:
             self.gpu_peak_flops = 0
-        self.num_flop_per_token = self._get_num_flop_per_token(seq_len)
+        self._matmul_flops_per_token = self._get_matmul_flops_per_token()
+        self._attention_modules = [m for m in model.modules() if hasattr(m, "attention_flops_per_token")]
 
     def count_tokens(self, tokens: int) -> None:
         """Push a step into the sliding window. Time is recorded internally."""
@@ -75,21 +76,20 @@ def get_tokens_per_second(self) -> float | None:
             return None
         return sum(self.tokens[1:]) / (self.times[-1] - self.times[0])
 
-    def get_mfu(self) -> float | None:
-        tokens_per_second = self.get_tokens_per_second()
-        if tokens_per_second is None:
-            return None
-        return self._mfu_from_tps(tokens_per_second)
-
     def get_step_tokens_per_second(self, tokens: int, fwd_bwd_time: float) -> float:
         """Single-step throughput from a caller-provided duration."""
         return tokens / fwd_bwd_time
 
-    def get_step_mfu(self, tokens: int, fwd_bwd_time: float) -> float:
-        return self._mfu_from_tps(self.get_step_tokens_per_second(tokens, fwd_bwd_time))
+    def get_flops(self, seq_lens: list[int]) -> int:
+        """Training FLOPs of one packed micro-batch, as torchtitan's `get_nparams_and_flops` per document:
+        each token attends only within its own document, so the attention term uses the document length."""
+        return sum(
+            n * (self._matmul_flops_per_token + sum(m.attention_flops_per_token(n) for m in self._attention_modules))
+            for n in seq_lens
+        )
 
-    def _mfu_from_tps(self, tokens_per_second: float) -> float:
-        return 100 * self.num_flop_per_token * tokens_per_second / self.gpu_peak_flops / self._world.world_size
+    def get_step_mfu(self, flops: float, fwd_bwd_time: float) -> float:
+        return 100 * flops / fwd_bwd_time / self.gpu_peak_flops / self._world.world_size
 
     def _get_peak_flops(self, device_name: str) -> float:
         """
@@ -131,24 +131,17 @@ def _get_peak_flops(self, device_name: str) -> float:
             self._logger.warning(f"Peak FLOPS undefined for `{device_name}`. Falling back to A100 (312 TFLOPS)")
             return 312e12
 
-    def _get_num_flop_per_token(self, seq_len: int) -> int:
-        """Training FLOPs per token, as torchtitan's `get_nparams_and_flops`: 6 FLOPs per active matmul
-        parameter plus each layer's own attention term (`attention_flops_per_token` hooks)."""
-        attention_flops = sum(
-            module.attention_flops_per_token(seq_len)
-            for module in self.model.modules()
-            if hasattr(module, "attention_flops_per_token")
-        )
+    def _get_matmul_flops_per_token(self) -> int:
+        """6 FLOPs per active matmul parameter, as torchtitan's `get_nparams_and_flops`."""
         active_params = get_active_params(self.model)
         if not has_lora_layers(self.model):
-            return 6 * active_params + attention_flops
+            return 6 * active_params
         # LoRA: frozen base matmuls still cost forward (2x) and dX (2x); fully trainable non-LoRA
         # params (modules_to_save) add their weight gradient (2x); adapters cost 6x.
         return (
             4 * active_params
             + 2 * self._count_fully_trainable_params_excluding_lora()
             + 6 * self._count_lora_adapter_params()
-            + attention_flops
         )
 
     def _count_lora_adapter_params(self) -> int:
@@ -176,9 +169,9 @@ def _count_fully_trainable_params_excluding_lora(self) -> int:
 _PERF_COUNTER: PerfCounter | None = None
 
 
-def get_perf_counter(model: nn.Module, seq_len: int, window_size: int = 10) -> PerfCounter:
+def get_perf_counter(model: nn.Module, window_size: int = 10) -> PerfCounter:
     global _PERF_COUNTER
     if _PERF_COUNTER is None:
-        _PERF_COUNTER = PerfCounter(model, seq_len, window_size)
+        _PERF_COUNTER = PerfCounter(model, window_size)
 
     return _PERF_COUNTER
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index ed483fc672a..4ae90b3758f 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -296,7 +296,7 @@ def train(config: TrainerConfig):
             memory_profiler = MemoryProfiler(progress.step, config.memory_profiler_path)
 
         forward_backward_start_time = time.perf_counter()
-        seq_len = micro_batches[0]["input_ids"].shape[1]
+        perf_counter = get_perf_counter(model)
 
         # Normalize each loss component by its own global (dp_cp) denominator, so every rank
         # divides by the same value. With a per-rank denominator, ranks with fewer loss
@@ -339,8 +339,10 @@ def train(config: TrainerConfig):
 
         step_tokens_per_expert = 0
         step_local_num_tokens = 0
+        step_local_flops = 0
         for micro_step, micro_batch in enumerate(micro_batch_workers(prepare, micro_batches)):
             step_local_num_tokens += micro_batch["input_ids"].shape[1]
+            step_local_flops += perf_counter.get_flops(micro_batch["seq_lens"].tolist())
             input_ids = micro_batch["input_ids"].to("cuda")
             position_ids = micro_batch["position_ids"].to("cuda")
             advantages = micro_batch["advantages"].to("cuda")
@@ -631,14 +633,13 @@ def train(config: TrainerConfig):
             tensor_stats.update(get_expert_load_stats(step_tokens_per_expert, dp_cp_group))
 
         # The dp mesh excludes cp, whose ranks hold the same rows.
-        global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda")
-        dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group())
-        num_tokens = global_num_tokens.item()
+        global_counts = torch.tensor([step_local_num_tokens, step_local_flops], dtype=torch.float64, device="cuda")
+        dist.all_reduce(global_counts, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group())
+        num_tokens, num_flops = int(global_counts[0].item()), global_counts[1].item()
         progress.total_tokens += num_tokens
         progress.total_samples += batch_size
-        perf_counter = get_perf_counter(model, seq_len)
         throughput = perf_counter.get_step_tokens_per_second(num_tokens, forward_backward_time)
-        mfu = perf_counter.get_step_mfu(num_tokens, forward_backward_time)
+        mfu = perf_counter.get_step_mfu(num_flops, forward_backward_time)
         peak_memory = torch.cuda.max_memory_reserved() / 1024**3  # GiB
         max_peak_memory = max(max_peak_memory, peak_memory)
 
diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py
index 98727b861b9..98a3efe9fd9 100644
--- a/src/prime_rl/trainer/sft/train.py
+++ b/src/prime_rl/trainer/sft/train.py
@@ -450,9 +450,12 @@ def is_online_eval_step(step: int) -> bool:
             )
 
         step_tokens_per_expert = 0
+        perf_counter = get_perf_counter(model)
         step_local_num_tokens = 0
+        step_local_flops = 0
         for micro_step, micro_batch in enumerate(micro_batches):
             step_local_num_tokens += micro_batch["input_ids"].shape[1]
+            step_local_flops += perf_counter.get_flops(micro_batch["seq_lens"].tolist())
             if config.log.log_data:
                 print_sample(
                     micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer
@@ -551,19 +554,18 @@ def is_online_eval_step(step: int) -> bool:
             memory_profiler.step()
 
         # The dp mesh excludes cp, whose ranks hold the same rows.
-        global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda")
-        dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group())
-        num_tokens = global_num_tokens.item()
+        global_counts = torch.tensor([step_local_num_tokens, step_local_flops], dtype=torch.float64, device="cuda")
+        dist.all_reduce(global_counts, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group())
+        num_tokens, num_flops = int(global_counts[0].item()), global_counts[1].item()
         progress.total_tokens += num_tokens
         dataset_progress = get_dataset_progress(dataloader)
         progress.total_samples = dataset_progress["step"]
         # Throughput / MFU per step over the full step wall time, as torchtitan reports them with
         # log_freq=1 (tokens since last log / elapsed time), instead of a smoothed sliding window.
         step_time = time.perf_counter() - step_start_time
-        perf_counter = get_perf_counter(model, config.data.seq_len)
         perf_counter.count_tokens(num_tokens)
         throughput = perf_counter.get_step_tokens_per_second(num_tokens, step_time)
-        mfu = perf_counter.get_step_mfu(num_tokens, step_time)
+        mfu = perf_counter.get_step_mfu(num_flops, step_time)
         peak_memory = torch.cuda.max_memory_reserved() / 1024**3  # GiB
         max_peak_memory = max(max_peak_memory, peak_memory)
 
diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py
index b3feb1002b0..33aa612d351 100644
--- a/tests/unit/train/test_perf.py
+++ b/tests/unit/train/test_perf.py
@@ -16,4 +16,13 @@ def test_flops_per_token_matches_torchtitan(model_name: str, flops_per_token: in
     config.pad_token_id = config.eos_token_id
     with torch.device("meta"):
         model = AutoModelForCausalLMPrimeRL.from_config(config)
-    assert PerfCounter(model, seq_len=1024).num_flop_per_token == flops_per_token
+    assert PerfCounter(model).get_flops([1024]) == 1024 * flops_per_token
+
+
+def test_packed_documents_attend_within_document():
+    config = AutoConfig.from_pretrained("Qwen/Qwen3-0.6B", attn_implementation="flash_attention_2")
+    config.pad_token_id = config.eos_token_id
+    with torch.device("meta"):
+        model = AutoModelForCausalLMPrimeRL.from_config(config)
+    # torchtitan's FLOPs/token at seq_len 512: the 1024 value minus half its attention term, 6 * 28 * 16 * 256 * 512.
+    assert PerfCounter(model).get_flops([512, 512]) == 1024 * (4_280_942_592 - 352_321_536)

From 3fc3d36ec6c01f9bd9b0f4f8ba88326d07208bfa Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 15:48:14 +0000
Subject: [PATCH 4/9] refactor(perf): one FlopCounter for MFU and packing

- PreTrainedModelPrimeRL.num_active_params() and a runtime-checkable AttentionFlops
  protocol (attention_flops(seq_len) per document) on every sequence mixer.
- FlopCounter in trainer/perf.py, used by MFU and as the orchestrator's bin cost on a
  meta-built model; the config-parsing packing cost is removed.
- LoRA expert adapters count by top_k / num_experts.
- IndexShare derived in GlmMoeDsaConfig; shared set_pad_token_id; prime_rl.trainer.models
  imports without fla / ring_flash_attn / prime_kernels.
---
 src/prime_rl/orchestrator/batch.py            |  56 +----
 .../trainer/distributed/collectives.py        |   4 +-
 src/prime_rl/trainer/model.py                 |  37 +--
 src/prime_rl/trainer/models/__init__.py       |  31 ++-
 .../trainer/models/afmoe/modeling_afmoe.py    |   8 +-
 src/prime_rl/trainer/models/base.py           |  42 +++-
 .../trainer/models/deepseek_v4/attention.py   |  31 ++-
 .../glm_moe_dsa/configuration_glm_moe_dsa.py  |   4 +-
 .../glm_moe_dsa/sparse_mla_attention.py       |  11 +-
 .../trainer/models/gpt_oss/attention.py       |   8 +-
 src/prime_rl/trainer/models/layers/attn.py    |  33 +--
 .../trainer/models/nemotron_h/mamba.py        |   8 +-
 .../trainer/models/qwen3_5/gated_delta_net.py |  24 +-
 .../models/qwen3_5/modeling_qwen3_5.py        |   8 +-
 .../models/qwen3_8_flash_next/attention.py    |  22 +-
 .../qwen3_8_flash_next/gated_delta_net.py     |  24 +-
 .../models/qwen3_8_flash_next/indexer.py      |   9 +-
 src/prime_rl/trainer/perf.py                  | 100 +++-----
 src/prime_rl/trainer/rl/train.py              |   2 +-
 src/prime_rl/trainer/sft/train.py             |   2 +-
 src/prime_rl/utils/cp.py                      |   8 +-
 tests/unit/orchestrator/test_batch.py         |  40 ++-
 tests/unit/train/test_perf.py                 | 234 ++++++++++++++++--
 23 files changed, 467 insertions(+), 279 deletions(-)

diff --git a/src/prime_rl/orchestrator/batch.py b/src/prime_rl/orchestrator/batch.py
index 69750926769..0456c49be70 100644
--- a/src/prime_rl/orchestrator/batch.py
+++ b/src/prime_rl/orchestrator/batch.py
@@ -2,14 +2,15 @@
 import heapq
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
-from typing import Any
 
 import numpy as np
+import torch
 from transformers import AutoConfig
 
 from prime_rl.configs.orchestrator import OrchestratorConfig
+from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL, set_pad_token_id
+from prime_rl.trainer.perf import FlopCounter
 from prime_rl.transports.batch.types import MicroBatch, MMImageRef, MMRefs, RoutedExperts, SamplingMask, TrainingSample
-from prime_rl.utils.logger import get_logger
 
 # Backfill value per component weight stream when a packed sample doesn't
 # carry it: absent rl means weight 1.0 on the loss mask, absent ce/ref_kl
@@ -17,36 +18,6 @@
 STREAM_FILL = {"rl_weights": 1.0, "ce_weights": 0.0, "ref_kl_weights": 0.0}
 
 
-def build_bin_cost(model_config: Any | None) -> Callable[[Sequence[int]], int]:
-    """Build a closure scoring a packed bin by its relative forward compute.
-
-    Per token, ~24 * layers * hidden_size^2 FLOPs (2 FLOPs per matmul parameter, ~12 * hidden_size^2
-    parameters per layer). Per causal query-key pair, heads * (qk_head_dim + v_head_dim) FLOPs per
-    full-attention layer, i.e. that times n^2 for a document of n tokens. Only the ratio of the two
-    matters for balancing. With ``model_config=None`` the cost is the token count.
-    """
-    if model_config is None:
-        token_cost, pair_cost = 1, 0
-    else:
-        config = getattr(model_config, "text_config", model_config)
-        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
-        qk_head_dim = getattr(config, "qk_nope_head_dim", 0) + getattr(config, "qk_rope_head_dim", 0) or head_dim
-        v_head_dim = getattr(config, "v_head_dim", None) or head_dim
-        layer_types = getattr(config, "layer_types", None)
-        num_attention_layers = (
-            sum(layer_type in ("full_attention", "attention") for layer_type in layer_types)
-            if layer_types
-            else config.num_hidden_layers
-        )
-        token_cost = 24 * config.num_hidden_layers * config.hidden_size**2
-        pair_cost = num_attention_layers * config.num_attention_heads * (qk_head_dim + v_head_dim)
-
-    def bin_cost(seqlens: Sequence[int]) -> int:
-        return token_cost * sum(seqlens) + pair_cost * sum(n * n for n in seqlens)
-
-    return bin_cost
-
-
 @dataclass
 class _WeightedSet:
     total: int = 0
@@ -769,17 +740,16 @@ def __init__(self, config: OrchestratorConfig):
         self.seq_len = config.seq_len
         self.num_train_workers = config.num_train_workers
         self.pad_to_multiple_of = config.pad_to_multiple_of
-        try:
-            model_config = AutoConfig.from_pretrained(
-                config.model.name, trust_remote_code=config.tokenizer.trust_remote_code
-            )
-        except Exception as e:
-            get_logger().warning(
-                f"Could not load model config for {config.model.name} ({e}) - "
-                "packing balances by token count instead of estimated FLOPs"
-            )
-            model_config = None
-        self.bin_cost = build_bin_cost(model_config)
+        # The flash-attention version selects kernels, not modules, so it does not change the FLOPs.
+        model_config = AutoConfig.from_pretrained(
+            config.model.name,
+            attn_implementation="flash_attention_3",
+            trust_remote_code=config.tokenizer.trust_remote_code,
+        )
+        set_pad_token_id(model_config)
+        with torch.device("meta"):
+            model = AutoModelForCausalLMPrimeRL.from_config(model_config)
+        self.bin_cost = FlopCounter(model)
 
     def pack(self, samples: list[TrainingSample]) -> list[list[MicroBatch]]:
         return prepare_batch(
diff --git a/src/prime_rl/trainer/distributed/collectives.py b/src/prime_rl/trainer/distributed/collectives.py
index 5a576ca7f81..14ea622ae6f 100644
--- a/src/prime_rl/trainer/distributed/collectives.py
+++ b/src/prime_rl/trainer/distributed/collectives.py
@@ -1,6 +1,5 @@
 """Autograd and compile-friendly distributed collectives."""
 
-import prime_kernels
 import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
@@ -100,6 +99,9 @@ def _mxfp8_all_to_all(
     group_name: str,
     quantized: bool,
 ) -> torch.Tensor:
+    # Imported here because prime-kernels is an optional linux-only extra.
+    import prime_kernels
+
     kernel = prime_kernels.load("mxfp8_moe")
     operation = kernel.all_to_all_dispatch if quantized else kernel.all_to_all_combine
     return operation(
diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py
index 5cc9a3ebd52..b35dea6525c 100644
--- a/src/prime_rl/trainer/model.py
+++ b/src/prime_rl/trainer/model.py
@@ -23,7 +23,7 @@
 from torch.distributed.fsdp._fully_shard._fsdp_init import _get_mesh_info
 from torch.distributed.tensor import Shard
 from torch.distributed.tensor.parallel import parallelize_module
-from transformers import AutoConfig, AutoTokenizer, GenerationConfig, PretrainedConfig
+from transformers import AutoConfig, AutoTokenizer, PretrainedConfig
 from transformers.tokenization_utils import PreTrainedTokenizer
 from transformers.utils.import_utils import is_flash_attn_3_available
 
@@ -45,6 +45,7 @@
     cast_float_and_contiguous,
     get_custom_causal_lm_cls,
     get_custom_vlm_cls,
+    set_pad_token_id,
     supports_custom_impl,
 )
 from prime_rl.trainer.models.deepseek_v4.attention import DeepseekV4Indexer
@@ -356,40 +357,8 @@ def get_model(
         if subconfig is not None and hasattr(subconfig, "use_cache"):
             subconfig.use_cache = False
     model_config.dsa_backend = config.dsa_backend
-    # Auto-enable IndexShare from the model's own indexer schedule (e.g. GLM-5.2). The model
-    # reads `indexer_types` directly: shared layers reuse cached indices and carry no indexer weights.
-    indexer_types = getattr(model_config, "indexer_types", None)
-    if indexer_types and any(t == "shared" for t in indexer_types):
-        model_config.use_index_cache = True
-        logger.info(
-            f"Auto-enabled IndexShare from indexer_types schedule "
-            f"({sum(t == 'full' for t in indexer_types)}/{len(indexer_types)} full layers)"
-        )
 
-    # Ensure pad_token_id is set (some models like Qwen3MoE don't have it).
-    # In transformers v5, token IDs moved from PretrainedConfig to GenerationConfig.
-    if not hasattr(model_config, "pad_token_id") or model_config.pad_token_id is None:
-        gen_config = GenerationConfig.from_model_config(model_config)
-        # Use `is not None` instead of truthiness: token ID 0 is valid.
-        pad_token_id = next(
-            (
-                v
-                for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(model_config, "eos_token_id", None)]
-                if v is not None
-            ),
-            None,
-        )
-        # Some HF configs (e.g. Llama 3.2) set pad_token_id to a list, which
-        # crashes both huggingface_hub's strict setter and transformers'
-        # GenerationConfig.validate(). Unwrap before assigning.
-        if isinstance(pad_token_id, list):
-            pad_token_id = pad_token_id[0]
-        model_config.pad_token_id = pad_token_id
-
-    # Handle list pad_token_id that was already set on the config (not from our
-    # fallback above, but directly in the model's config.json).
-    if isinstance(getattr(model_config, "pad_token_id", None), list):
-        model_config.pad_token_id = model_config.pad_token_id[0]
+    set_pad_token_id(model_config)
 
     # NOTE: For VLM models, we do NOT propagate dtype to sub_configs.
     # The model should load in its default dtype (bf16) to match vLLM inference.
diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py
index af9a0684f14..c00e7042739 100644
--- a/src/prime_rl/trainer/models/__init__.py
+++ b/src/prime_rl/trainer/models/__init__.py
@@ -2,7 +2,7 @@
 
 from collections import OrderedDict
 
-from transformers import AutoConfig
+from transformers import AutoConfig, GenerationConfig
 from transformers.configuration_utils import PretrainedConfig
 from transformers.models.auto.auto_factory import _BaseAutoModelClass, _LazyAutoMapping, auto_class_update
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES
@@ -97,6 +97,35 @@ def get_custom_causal_lm_cls(model_config: PretrainedConfig) -> type[PreTrainedM
     return _CUSTOM_CAUSAL_LM_BY_MODEL_TYPE[model_config.model_type]
 
 
+def set_pad_token_id(model_config: PretrainedConfig) -> None:
+    """Ensure pad_token_id is set (some models like Qwen3MoE don't have it).
+
+    In transformers v5, token IDs moved from PretrainedConfig to GenerationConfig.
+    """
+    if not hasattr(model_config, "pad_token_id") or model_config.pad_token_id is None:
+        gen_config = GenerationConfig.from_model_config(model_config)
+        # Use `is not None` instead of truthiness: token ID 0 is valid.
+        pad_token_id = next(
+            (
+                v
+                for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(model_config, "eos_token_id", None)]
+                if v is not None
+            ),
+            None,
+        )
+        # Some HF configs (e.g. Llama 3.2) set pad_token_id to a list, which
+        # crashes both huggingface_hub's strict setter and transformers'
+        # GenerationConfig.validate(). Unwrap before assigning.
+        if isinstance(pad_token_id, list):
+            pad_token_id = pad_token_id[0]
+        model_config.pad_token_id = pad_token_id
+
+    # Handle list pad_token_id that was already set on the config (not from our
+    # fallback above, but directly in the model's config.json).
+    if isinstance(getattr(model_config, "pad_token_id", None), list):
+        model_config.pad_token_id = model_config.pad_token_id[0]
+
+
 def supports_custom_impl(model_config: PretrainedConfig) -> bool:
     """Check if the model configuration supports the custom PrimeRL implementation.
 
diff --git a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
index 8556dd825ef..4b7a5f25c45 100644
--- a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
+++ b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
@@ -18,7 +18,7 @@
     flash_attn_2_varlen_op,
     flash_attn_3_varlen_op,
     flash_attn_4_varlen_op,
-    quadratic_attention_flops_per_token,
+    quadratic_attention_flops,
 )
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
@@ -83,13 +83,13 @@ def __init__(self, config: AfmoeAttentionConfig):
         self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps))
         self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps))
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
-        return quadratic_attention_flops_per_token(
+    def attention_flops(self, seq_len: int) -> int:
+        return quadratic_attention_flops(
             num_heads=self.num_heads,
             qk_head_dim=self.head_dim,
             v_head_dim=self.head_dim,
             seq_len=seq_len,
-            sliding_window=self.sliding_window,
+            attended_tokens=min(seq_len, self.sliding_window or seq_len),
         )
 
 
diff --git a/src/prime_rl/trainer/models/base.py b/src/prime_rl/trainer/models/base.py
index dc5070580c3..2ecd2ba5f1b 100644
--- a/src/prime_rl/trainer/models/base.py
+++ b/src/prime_rl/trainer/models/base.py
@@ -1,15 +1,49 @@
+from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, Protocol, runtime_checkable
 
-from torch import Tensor
+from torch import Tensor, nn
 from transformers.modeling_utils import PreTrainedModel
 
+from prime_rl.trainer.models.layers.moe import MoE
 from prime_rl.utils.cp import CPContext
 
 CPStyle = Literal["ring", "ulysses"]
 ALL_CP_STYLES: frozenset[CPStyle] = frozenset({"ring", "ulysses"})
 
 
+@runtime_checkable
+class AttentionFlops(Protocol):
+    """A sequence-mixing layer (softmax, sliding-window, sparse or linear attention, Mamba)."""
+
+    def attention_flops(self, seq_len: int) -> int:
+        """Training FLOPs of this layer's sequence mixing for one document of `seq_len` tokens."""
+        ...
+
+
+def count_active_params(model: nn.Module, params: Iterable[nn.Parameter]) -> int:
+    """Elements of `params` that a token multiplies through, as torchtitan's `get_nparams_and_active_nparams`.
+
+    Routed experts count by `top_k / num_experts`; embedding tables count only when tied to the lm_head.
+    """
+    params = {id(param): param for param in params}
+    lm_head_params = {id(param) for param in model.lm_head.parameters()}
+    inactive = 0
+    for module in model.modules():
+        if isinstance(module, nn.Embedding):
+            inactive += sum(
+                param.numel()
+                for param in module.parameters(recurse=False)
+                if id(param) in params and id(param) not in lm_head_params
+            )
+        elif isinstance(module, MoE):
+            # Expert weights lead with `num_experts`, so the division is exact.
+            num_experts, top_k = module.experts.num_experts, module.router.top_k
+            expert_params = sum(param.numel() for param in module.experts.parameters() if id(param) in params)
+            inactive += expert_params * (num_experts - top_k) // num_experts
+    return sum(param.numel() for param in params.values()) - inactive
+
+
 @dataclass(frozen=True)
 class CPSupport:
     """Which context-parallel styles an architecture can train under, and why not the rest."""
@@ -77,6 +111,10 @@ def get_correct_experts_implementation(self, requested_experts: str | None) -> s
         """PrimeRL models always use eager experts implementation."""
         return "eager"
 
+    def num_active_params(self) -> int:
+        """Parameters a text token multiplies through, LoRA adapters included."""
+        return count_active_params(self, self.parameters())
+
     @classmethod
     def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool:
         """
diff --git a/src/prime_rl/trainer/models/deepseek_v4/attention.py b/src/prime_rl/trainer/models/deepseek_v4/attention.py
index e4d335aecc6..dedc797d7ab 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/attention.py
+++ b/src/prime_rl/trainer/models/deepseek_v4/attention.py
@@ -134,7 +134,7 @@
 from prime_rl.trainer.models.kernels.deepseek_v4 import IGNORE_SLOT
 from prime_rl.trainer.models.kernels.deepseek_v4.dsv4_rope import dsv4_q_norm_rope, dsv4_rope
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
-from prime_rl.trainer.models.layers.attn import quadratic_attention_flops_per_token
+from prime_rl.trainer.models.layers.attn import quadratic_attention_flops
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
 from prime_rl.utils.cp import CPContext, gather_for_cp
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
@@ -706,26 +706,25 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int, rotary_emb: Deepsee
 
         self.cp_context = CPContext()
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
-        # As torchtitan's DeepSeek V4: the local window, plus the compressed entries (CSA: the
-        # indexer scores all of them, attention reads `index_topk`).
-        flops = quadratic_attention_flops_per_token(
-            num_heads=self.num_heads,
-            qk_head_dim=self.head_dim,
-            v_head_dim=self.head_dim,
-            seq_len=seq_len,
-            sliding_window=self.config.sliding_window,
-        )
+    def attention_flops(self, seq_len: int) -> int:
+        # As torchtitan's DeepSeek V4: each query attends the local window plus the compressed entries
+        # (CSA: the indexer scores all of them, attention reads `index_topk`).
+        attended_tokens = min(seq_len, self.config.sliding_window)
+        indexer_flops = 0
         if self.compressor is not None:
             compressed_len = seq_len // self.compressor.compress_rate
             indexer = getattr(self.compressor, "indexer", None)
             if indexer is not None:
-                flops += 6 * indexer.num_heads * indexer.head_dim * compressed_len
+                indexer_flops = 6 * indexer.num_heads * indexer.head_dim * seq_len * compressed_len
                 compressed_len = min(compressed_len, indexer.index_topk)
-            flops += quadratic_attention_flops_per_token(
-                num_heads=self.num_heads, qk_head_dim=self.head_dim, v_head_dim=self.head_dim, seq_len=compressed_len
-            )
-        return flops
+            attended_tokens += compressed_len
+        return indexer_flops + quadratic_attention_flops(
+            num_heads=self.num_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=seq_len,
+            attended_tokens=attended_tokens,
+        )
 
     def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[torch.Tensor, None]:
         """`packed` carries the document boundaries every pathway below is clipped at."""
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py
index 4588b992765..a2855ecad87 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py
@@ -207,7 +207,9 @@ def __init__(
         self.index_head_dim = index_head_dim
         self.indexer_rope_interleave = indexer_rope_interleave
         self.index_topk = index_topk
-        self.use_index_cache = use_index_cache
+        # An `indexer_types` schedule with shared layers (e.g. GLM-5.2) turns on IndexShare: shared layers
+        # reuse the previous full layer's indices and carry no indexer weights.
+        self.use_index_cache = use_index_cache or "shared" in (indexer_types or ())
         self.index_topk_freq = index_topk_freq
         self.index_topk_pattern = index_topk_pattern
         self.indexer_types = indexer_types
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 18ac67353d3..8350f58d47c 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -6,7 +6,7 @@
 
 from prime_rl.configs.trainer import DSABackend
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
-from prime_rl.trainer.models.layers.attn import quadratic_attention_flops_per_token
+from prime_rl.trainer.models.layers.attn import quadratic_attention_flops
 from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import rotate_half
 from prime_rl.utils.cp import CPContext, gather_for_cp
@@ -150,17 +150,18 @@ def __init__(self, args: SparseMlaAttentionArgs):
 
         self.cp_context = CPContext()
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
+    def attention_flops(self, seq_len: int) -> int:
         # As torchtitan's DeepSeek V3/V4: MLA counted un-absorbed (qk = nope + rope), sparse attention
         # attends `index_topk` tokens, and the indexer scores every token with its own heads.
-        flops = quadratic_attention_flops_per_token(
+        flops = quadratic_attention_flops(
             num_heads=self.num_heads,
             qk_head_dim=self.qk_head_dim,
             v_head_dim=self.v_head_dim,
-            seq_len=min(seq_len, self.args.index_topk),
+            seq_len=seq_len,
+            attended_tokens=min(seq_len, self.args.index_topk),
         )
         if self.indexer is not None:
-            flops += 6 * self.args.index_n_heads * self.args.index_head_dim * seq_len
+            flops += 6 * self.args.index_n_heads * self.args.index_head_dim * seq_len * seq_len
         return flops
 
     def mla_latents(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
diff --git a/src/prime_rl/trainer/models/gpt_oss/attention.py b/src/prime_rl/trainer/models/gpt_oss/attention.py
index 2f5d344a42a..ba2e7f7dd2c 100644
--- a/src/prime_rl/trainer/models/gpt_oss/attention.py
+++ b/src/prime_rl/trainer/models/gpt_oss/attention.py
@@ -10,7 +10,7 @@
 
 from prime_rl.trainer.models.fusions import fuse_qkv_projections
 from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig
-from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func, quadratic_attention_flops_per_token
+from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func, quadratic_attention_flops
 from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb
 
 
@@ -51,13 +51,13 @@ def __init__(self, config: GptOssConfig, layer_idx: int) -> None:
         nn.init.normal_(self.sinks, mean=0.0, std=config.initializer_range)
         self.flash_attn = flash_attn_4_varlen_func
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
-        return quadratic_attention_flops_per_token(
+    def attention_flops(self, seq_len: int) -> int:
+        return quadratic_attention_flops(
             num_heads=self.num_attention_heads,
             qk_head_dim=self.head_dim,
             v_head_dim=self.head_dim,
             seq_len=seq_len,
-            sliding_window=self.sliding_window,
+            attended_tokens=min(seq_len, self.sliding_window or seq_len),
         )
 
     def compute_attention(
diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py
index 1859dccbf43..256f2694f48 100644
--- a/src/prime_rl/trainer/models/layers/attn.py
+++ b/src/prime_rl/trainer/models/layers/attn.py
@@ -404,26 +404,27 @@ def flash_attn_3_varlen_op(
     return out
 
 
-def quadratic_attention_flops_per_token(
-    *, num_heads: int, qk_head_dim: int, v_head_dim: int, seq_len: int, sliding_window: int | None = None
+def quadratic_attention_flops(
+    *, num_heads: int, qk_head_dim: int, v_head_dim: int, seq_len: int, attended_tokens: int | None = None
 ) -> int:
-    """Training FLOPs per token of softmax attention, as torchtitan's `quadratic_attention_flops_per_token`.
+    """Training FLOPs of softmax attention over one document, as torchtitan's `quadratic_attention_flops_per_token`.
 
-    Two matmuls per head and attended token (q @ K^T over `qk_head_dim`, scores @ V over `v_head_dim`),
-    2 FLOPs per multiply-add, 3x for forward plus backward. Causal sparsity and the flash-attention
-    recompute are not counted.
+    Each of the `seq_len` queries attends `attended_tokens` keys (default all `seq_len`): two matmuls per
+    head (q @ K^T over `qk_head_dim`, scores @ V over `v_head_dim`), 2 FLOPs per multiply-add, 3x for
+    forward plus backward. Causal sparsity and the flash-attention recompute are not counted.
     """
-    attended_tokens = seq_len if sliding_window is None else min(seq_len, sliding_window)
-    return 6 * num_heads * (qk_head_dim + v_head_dim) * attended_tokens
+    if attended_tokens is None:
+        attended_tokens = seq_len
+    return 6 * num_heads * (qk_head_dim + v_head_dim) * seq_len * attended_tokens
 
 
-def delta_rule_flops_per_token(*, num_heads: int, key_head_dim: int, v_head_dim: int) -> int:
-    """Training FLOPs per token of a gated delta-rule recurrence, as torchtitan's `delta_rule_flops_per_token`.
+def delta_rule_flops(*, num_heads: int, key_head_dim: int, v_head_dim: int, seq_len: int) -> int:
+    """Training FLOPs of a gated delta-rule recurrence over one document, as torchtitan's `delta_rule_flops_per_token`.
 
-    Reading, updating and querying the `[key_head_dim, v_head_dim]` state are three
-    `key_head_dim * v_head_dim` products per head.
+    Per token and head, reading, updating and querying the `[key_head_dim, v_head_dim]` state are three
+    `key_head_dim * v_head_dim` products.
     """
-    return 6 * 3 * num_heads * key_head_dim * v_head_dim
+    return 6 * 3 * num_heads * key_head_dim * v_head_dim * seq_len
 
 
 @dataclass
@@ -491,13 +492,13 @@ def __init__(self, config: AttentionConfig, flash_attn_version: int = 2):
         self._flash_attn_version = flash_attn_version
         self.func = self._funcs[flash_attn_version]
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
-        return quadratic_attention_flops_per_token(
+    def attention_flops(self, seq_len: int) -> int:
+        return quadratic_attention_flops(
             num_heads=self.num_heads,
             qk_head_dim=self.head_dim,
             v_head_dim=self.head_dim,
             seq_len=seq_len,
-            sliding_window=getattr(self, "sliding_window", None),
+            attended_tokens=min(seq_len, getattr(self, "sliding_window", None) or seq_len),
         )
 
     def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
diff --git a/src/prime_rl/trainer/models/nemotron_h/mamba.py b/src/prime_rl/trainer/models/nemotron_h/mamba.py
index dc3b775e77a..e824fd63861 100644
--- a/src/prime_rl/trainer/models/nemotron_h/mamba.py
+++ b/src/prime_rl/trainer/models/nemotron_h/mamba.py
@@ -127,10 +127,10 @@ def __init__(self, config: NemotronHConfig) -> None:
 
         self.cp_context = CPContext()
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
-        # Mamba2's analogue of torchtitan's `delta_rule_flops_per_token`: per head, the state update
-        # (x outer B) and the readout (state @ C) are two `head_dim * state_size` products.
-        return 6 * 2 * self.num_heads * self.head_dim * self.state_size
+    def attention_flops(self, seq_len: int) -> int:
+        # Mamba2's analogue of torchtitan's `delta_rule_flops_per_token`: per token and head, the state
+        # update (x outer B) and the readout (state @ C) are two `head_dim * state_size` products.
+        return 6 * 2 * self.num_heads * self.head_dim * self.state_size * seq_len
 
     def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
         batch_size, sequence_length, _ = hidden_states.shape
diff --git a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
index cb089e8c2ee..f2f8b5ce983 100644
--- a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
@@ -1,15 +1,20 @@
 import torch
 import torch.nn.functional as F
-from fla.modules import FusedRMSNormGated
-from fla.modules.conv import causal_conv1d
-from fla.ops.cp import build_cp_context
-from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 from torch import nn
 
-from prime_rl.trainer.models.layers.attn import delta_rule_flops_per_token
+from prime_rl.trainer.models.layers.attn import delta_rule_flops
 from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
 from prime_rl.utils.cp import CPContext
 
+# Guarded because flash-linear-attention ships in the linux-gated `gpu` extra; the layer needs it to build.
+try:
+    from fla.modules import FusedRMSNormGated
+    from fla.modules.conv import causal_conv1d
+    from fla.ops.cp import build_cp_context
+    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+except ImportError:
+    FusedRMSNormGated = causal_conv1d = build_cp_context = chunk_gated_delta_rule = None  # type: ignore
+
 # Dynamo lowers all-gather to concatenation, then fails to copy the result into
 # FLA's stacked output buffer. Keep CP convolution eager until this is fixed:
 # https://github.com/pytorch/pytorch/issues/155632
@@ -53,9 +58,12 @@ def __init__(self, config: Qwen3_5TextConfig) -> None:
         self.in_proj_a = nn.Linear(config.hidden_size, self.num_value_heads, bias=False)
         self.cp_context = CPContext()
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
-        return delta_rule_flops_per_token(
-            num_heads=self.num_value_heads, key_head_dim=self.key_head_dim, v_head_dim=self.value_head_dim
+    def attention_flops(self, seq_len: int) -> int:
+        return delta_rule_flops(
+            num_heads=self.num_value_heads,
+            key_head_dim=self.key_head_dim,
+            v_head_dim=self.value_head_dim,
+            seq_len=seq_len,
         )
 
     def forward(
diff --git a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py
index ab269415c45..433cf9afa77 100644
--- a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py
@@ -2,7 +2,7 @@
 from torch import Tensor, nn
 from transformers.modeling_outputs import BaseModelOutput
 
-from prime_rl.trainer.models.base import ALL_CP_STYLES, CPSupport, PreTrainedModelPrimeRL
+from prime_rl.trainer.models.base import ALL_CP_STYLES, CPSupport, PreTrainedModelPrimeRL, count_active_params
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear
 from prime_rl.trainer.models.layers.mlp import FeedForward
 from prime_rl.trainer.models.layers.moe import GroupedExperts, MoE, TokenChoiceTopKRouter
@@ -328,6 +328,12 @@ def get_input_embeddings(self) -> nn.Embedding:
     def set_input_embeddings(self, embeddings: nn.Embedding) -> None:
         self.model.set_input_embeddings(embeddings)
 
+    def num_active_params(self) -> int:
+        # The vision encoder runs per image, not per text token.
+        if not self.is_vlm:
+            return super().num_active_params()
+        return super().num_active_params() - count_active_params(self, self.model.visual.parameters())
+
     def forward(
         self,
         input_ids: torch.LongTensor | None = None,
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
index 00a3b1ac2a3..6b902f0804d 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
@@ -1,7 +1,7 @@
 import torch
 from torch import nn
 
-from prime_rl.trainer.models.layers.attn import quadratic_attention_flops_per_token
+from prime_rl.trainer.models.layers.attn import quadratic_attention_flops
 from prime_rl.trainer.models.qwen3_8_flash_next.indexer import SparseAttentionIndexer
 from prime_rl.trainer.models.qwen3_8_flash_next.norm import RMSNorm
 from prime_rl.trainer.models.qwen3_8_flash_next.rotary_embedding import apply_rotary_embedding
@@ -22,8 +22,6 @@ def __init__(
         indexer_token_budget: int,
         indexer_compression_ratio: int,
     ) -> None:
-        import prime_kernels
-
         super().__init__()
         self.num_heads = num_attention_heads
         self.num_key_value_heads = num_key_value_heads
@@ -46,19 +44,25 @@ def __init__(
         )
 
         self.cp_context = CPContext()
-        # Resolve once before torch.compile traces the layer; the loader caches the module.
-        self.indexed_attention = prime_kernels.load("indexed_attention").indexed_attention
+        # Resolve once before torch.compile traces the layer; the loader caches the module. A host
+        # without a GPU (the orchestrator builds the model on meta to count FLOPs) skips it.
+        if torch.cuda.is_available():
+            import prime_kernels
+
+            self.indexed_attention = prime_kernels.load("indexed_attention").indexed_attention
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
+    def attention_flops(self, seq_len: int) -> int:
         # As torchtitan's DeepSeek V4 CSA: the indexer scores every compressed block,
         # attention reads `token_budget` tokens.
         indexer = self.indexer
-        indexer_flops = 6 * indexer.num_query_heads * indexer.head_dim * (seq_len // indexer.compression_ratio)
-        return indexer_flops + quadratic_attention_flops_per_token(
+        num_blocks = seq_len // indexer.compression_ratio
+        indexer_flops = 6 * indexer.num_query_heads * indexer.head_dim * seq_len * num_blocks
+        return indexer_flops + quadratic_attention_flops(
             num_heads=self.num_heads,
             qk_head_dim=self.head_dim,
             v_head_dim=self.head_dim,
-            seq_len=min(seq_len, indexer.token_budget),
+            seq_len=seq_len,
+            attended_tokens=min(seq_len, indexer.token_budget),
         )
 
     def forward(
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
index 0164e369f6b..89ce2416a9a 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
@@ -1,14 +1,19 @@
 import torch
 import torch.nn.functional as F
-from fla.modules import FusedRMSNormGated
-from fla.modules.conv import causal_conv1d
-from fla.ops.cp import build_cp_context
-from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 from torch import nn
 
-from prime_rl.trainer.models.layers.attn import delta_rule_flops_per_token
+from prime_rl.trainer.models.layers.attn import delta_rule_flops
 from prime_rl.utils.cp import CPContext
 
+# Guarded because flash-linear-attention ships in the linux-gated `gpu` extra; the layer needs it to build.
+try:
+    from fla.modules import FusedRMSNormGated
+    from fla.modules.conv import causal_conv1d
+    from fla.ops.cp import build_cp_context
+    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+except ImportError:
+    FusedRMSNormGated = causal_conv1d = build_cp_context = chunk_gated_delta_rule = None  # type: ignore
+
 # FLA's context carries a process group that Dynamo cannot trace through the convolution.
 causal_conv1d_with_context_parallelism = torch.compiler.disable(causal_conv1d)
 
@@ -57,9 +62,12 @@ def __init__(
 
         self.cp_context = CPContext()
 
-    def attention_flops_per_token(self, seq_len: int) -> int:
-        return delta_rule_flops_per_token(
-            num_heads=self.num_value_heads, key_head_dim=self.key_head_dim, v_head_dim=self.value_head_dim
+    def attention_flops(self, seq_len: int) -> int:
+        return delta_rule_flops(
+            num_heads=self.num_value_heads,
+            key_head_dim=self.key_head_dim,
+            v_head_dim=self.value_head_dim,
+            seq_len=seq_len,
         )
 
     def forward(
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/indexer.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/indexer.py
index ace92d4dfb4..81123439f39 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/indexer.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/indexer.py
@@ -18,10 +18,13 @@ def __init__(
         compression_ratio: int,
         norm_eps: float,
     ) -> None:
-        import prime_kernels
-
         super().__init__()
-        self.select_indexed_blocks = prime_kernels.load("indexed_attention").select_indexed_blocks
+        # Resolved before torch.compile traces the layer. A host without a GPU (the orchestrator builds
+        # the model on meta to count FLOPs) never runs the layer and skips it.
+        if torch.cuda.is_available():
+            import prime_kernels
+
+            self.select_indexed_blocks = prime_kernels.load("indexed_attention").select_indexed_blocks
         self.num_query_heads = num_query_heads
         self.head_dim = head_dim
         self.token_budget = token_budget
diff --git a/src/prime_rl/trainer/perf.py b/src/prime_rl/trainer/perf.py
index 0a966dc2d7f..fe4c42886d8 100644
--- a/src/prime_rl/trainer/perf.py
+++ b/src/prime_rl/trainer/perf.py
@@ -1,36 +1,44 @@
 import time
+from collections.abc import Sequence
 
 import torch
 from torch import nn
 
-from prime_rl.trainer.lora import has_lora_layers
-from prime_rl.trainer.models.layers.lora import MultiLoRAModule
-from prime_rl.trainer.models.layers.moe import MoE
+from prime_rl.trainer.models.base import AttentionFlops, PreTrainedModelPrimeRL, count_active_params
 from prime_rl.trainer.world import get_world
 from prime_rl.utils.logger import get_logger
 
 
-def get_active_params(model: nn.Module) -> int:
-    """Matmul parameters used per token, as torchtitan's `get_nparams_and_active_nparams`.
+class FlopCounter:
+    """Training FLOPs of packed documents, as torchtitan's `get_nparams_and_flops` per document.
 
-    Routed experts count by `top_k / num_experts`; embeddings count only when tied to the lm_head;
-    the vision encoder and LoRA adapters do not count.
+    6 FLOPs per active parameter per token, plus each sequence-mixing layer's FLOPs per document (a
+    token attends only within its own document). The model is walked once, at construction.
     """
-    excluded = {id(param) for name, param in model.named_parameters() if "lora_A" in name or "lora_B" in name}
-    visual = getattr(model.model, "visual", None)
-    if visual is not None:
-        excluded |= {id(param) for param in visual.parameters()}
-    lm_head_params = {id(param) for param in model.lm_head.parameters()}
-    inactive_expert_params = 0
-    for module in model.modules():
-        if isinstance(module, nn.Embedding):
-            excluded |= {id(param) for param in module.parameters(recurse=False)} - lm_head_params
-        elif isinstance(module, MoE):
-            # Expert weights lead with `num_experts`, so the division is exact.
-            num_experts, top_k = module.experts.num_experts, module.router.top_k
-            expert_params = sum(p.numel() for p in module.experts.parameters() if id(p) not in excluded)
-            inactive_expert_params += expert_params * (num_experts - top_k) // num_experts
-    return sum(param.numel() for param in model.parameters() if id(param) not in excluded) - inactive_expert_params
+
+    def __init__(self, model: PreTrainedModelPrimeRL):
+        active_params = model.num_active_params()
+        lora_params = [param for name, param in model.named_parameters() if "lora_A" in name or "lora_B" in name]
+        if lora_params:
+            # Frozen base matmuls cost forward (2x) and dX (2x); fully trainable non-LoRA params
+            # (modules_to_save) add their weight gradient (2x); adapters cost 6x.
+            lora_active_params = count_active_params(model, lora_params)
+            lora_param_ids = {id(param) for param in lora_params}
+            trainable_params = sum(
+                param.numel() for param in model.parameters() if param.requires_grad and id(param) not in lora_param_ids
+            )
+            self.matmul_flops_per_token = (
+                4 * (active_params - lora_active_params) + 2 * trainable_params + 6 * lora_active_params
+            )
+        else:
+            self.matmul_flops_per_token = 6 * active_params
+        self.attention_modules = [module for module in model.modules() if isinstance(module, AttentionFlops)]
+
+    def __call__(self, seq_lens: Sequence[int]) -> int:
+        return sum(
+            n * self.matmul_flops_per_token + sum(module.attention_flops(n) for module in self.attention_modules)
+            for n in seq_lens
+        )
 
 
 class PerfCounter:
@@ -42,7 +50,7 @@ class PerfCounter:
       read smoothed values via `get_tokens_per_second()`.
       Time is measured between successive `count_tokens` calls (full step).
     - Single point on a caller-provided duration: `get_step_tokens_per_second(tokens, fwd_bwd_time)`
-      / `get_step_mfu(flops, fwd_bwd_time)`, with `flops` summed from `get_flops(seq_lens)`. No smoothing.
+      / `get_step_mfu(flops, fwd_bwd_time)`, with `flops` summed from `flop_counter(seq_lens)`. No smoothing.
 
     Sliding window inspired by https://github.com/pytorch/torchtitan/blob/4b3f2e41a084bf79a8540068ed525539d1244edd/torchtitan/utils.py#L119
     """
@@ -51,7 +59,6 @@ def __init__(self, model: nn.Module, window_size: int = 10):
         self.window_size = window_size
         self.tokens: list[int] = []
         self.times: list[float] = []
-        self.model = model
 
         self._world = get_world()
         self._logger = get_logger()
@@ -60,8 +67,7 @@ def __init__(self, model: nn.Module, window_size: int = 10):
             self.gpu_peak_flops = self._get_peak_flops(torch.cuda.get_device_name(torch.device("cuda")))
         else:
             self.gpu_peak_flops = 0
-        self._matmul_flops_per_token = self._get_matmul_flops_per_token()
-        self._attention_modules = [m for m in model.modules() if hasattr(m, "attention_flops_per_token")]
+        self.flop_counter = FlopCounter(model)
 
     def count_tokens(self, tokens: int) -> None:
         """Push a step into the sliding window. Time is recorded internally."""
@@ -80,14 +86,6 @@ def get_step_tokens_per_second(self, tokens: int, fwd_bwd_time: float) -> float:
         """Single-step throughput from a caller-provided duration."""
         return tokens / fwd_bwd_time
 
-    def get_flops(self, seq_lens: list[int]) -> int:
-        """Training FLOPs of one packed micro-batch, as torchtitan's `get_nparams_and_flops` per document:
-        each token attends only within its own document, so the attention term uses the document length."""
-        return sum(
-            n * (self._matmul_flops_per_token + sum(m.attention_flops_per_token(n) for m in self._attention_modules))
-            for n in seq_lens
-        )
-
     def get_step_mfu(self, flops: float, fwd_bwd_time: float) -> float:
         return 100 * flops / fwd_bwd_time / self.gpu_peak_flops / self._world.world_size
 
@@ -131,40 +129,6 @@ def _get_peak_flops(self, device_name: str) -> float:
             self._logger.warning(f"Peak FLOPS undefined for `{device_name}`. Falling back to A100 (312 TFLOPS)")
             return 312e12
 
-    def _get_matmul_flops_per_token(self) -> int:
-        """6 FLOPs per active matmul parameter, as torchtitan's `get_nparams_and_flops`."""
-        active_params = get_active_params(self.model)
-        if not has_lora_layers(self.model):
-            return 6 * active_params
-        # LoRA: frozen base matmuls still cost forward (2x) and dX (2x); fully trainable non-LoRA
-        # params (modules_to_save) add their weight gradient (2x); adapters cost 6x.
-        return (
-            4 * active_params
-            + 2 * self._count_fully_trainable_params_excluding_lora()
-            + 6 * self._count_lora_adapter_params()
-        )
-
-    def _count_lora_adapter_params(self) -> int:
-        """Count LoRA adapter parameters (sum of lora_A and lora_B across all MultiLoRAModules)."""
-        params = 0
-        for module in self.model.modules():
-            if isinstance(module, MultiLoRAModule):
-                adapter_params, _ = module.get_lora_param_counts()
-                params += adapter_params
-        return params
-
-    def _count_fully_trainable_params_excluding_lora(self) -> int:
-        """Count trainable parameters excluding LoRA adapter tensors.
-
-        Approximates trainable matmul params in modules_to_save by subtracting LoRA adapter params
-        from all trainable params.
-        """
-        total_trainable = 0
-        for name, param in self.model.named_parameters():
-            if param.requires_grad and ("lora_A" not in name and "lora_B" not in name):
-                total_trainable += param.numel()
-        return total_trainable
-
 
 _PERF_COUNTER: PerfCounter | None = None
 
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 4ae90b3758f..eadc9b35077 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -342,7 +342,7 @@ def train(config: TrainerConfig):
         step_local_flops = 0
         for micro_step, micro_batch in enumerate(micro_batch_workers(prepare, micro_batches)):
             step_local_num_tokens += micro_batch["input_ids"].shape[1]
-            step_local_flops += perf_counter.get_flops(micro_batch["seq_lens"].tolist())
+            step_local_flops += perf_counter.flop_counter(micro_batch["seq_lens"].tolist())
             input_ids = micro_batch["input_ids"].to("cuda")
             position_ids = micro_batch["position_ids"].to("cuda")
             advantages = micro_batch["advantages"].to("cuda")
diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py
index 98a3efe9fd9..eca29c9773e 100644
--- a/src/prime_rl/trainer/sft/train.py
+++ b/src/prime_rl/trainer/sft/train.py
@@ -455,7 +455,7 @@ def is_online_eval_step(step: int) -> bool:
         step_local_flops = 0
         for micro_step, micro_batch in enumerate(micro_batches):
             step_local_num_tokens += micro_batch["input_ids"].shape[1]
-            step_local_flops += perf_counter.get_flops(micro_batch["seq_lens"].tolist())
+            step_local_flops += perf_counter.flop_counter(micro_batch["seq_lens"].tolist())
             if config.log.log_data:
                 print_sample(
                     micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer
diff --git a/src/prime_rl/utils/cp.py b/src/prime_rl/utils/cp.py
index 3be423774ce..5f586dedd6c 100644
--- a/src/prime_rl/utils/cp.py
+++ b/src/prime_rl/utils/cp.py
@@ -1,16 +1,13 @@
 from __future__ import annotations
 
-# ruff: noqa: I001 — `prime_rl._compat` must run before `ring_flash_attn` imports below.
-import prime_rl._compat  # noqa: F401
-
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from ring_flash_attn import update_ring_flash_attn_params
 
+import prime_rl._compat  # noqa: F401  (patches transformers before `ring_flash_attn` is imported)
 from prime_rl.trainer.distributed.collectives import all_gather
 from prime_rl.utils.logger import get_logger
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
@@ -151,6 +148,9 @@ def setup_cp_attention_params(
     )
 
     if cp_style == "ring":
+        # Imported here because ring_flash_attn ships in the linux-gated `gpu` extra.
+        from ring_flash_attn import update_ring_flash_attn_params
+
         update_ring_flash_attn_params(cu_seqlens, cp_group)
     elif cp_style == "ulysses":
         # Delayed import: ulysses_attn lives under trainer.models, which imports
diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py
index d498d9bb1ba..99a1701e4ae 100644
--- a/tests/unit/orchestrator/test_batch.py
+++ b/tests/unit/orchestrator/test_batch.py
@@ -1,11 +1,8 @@
-from types import SimpleNamespace
-
 import numpy as np
 import pytest
 
 from prime_rl.orchestrator.batch import (
     _is_multimodal_sample,
-    build_bin_cost,
     pad_micro_batch,
     prepare_batch,
     prepare_sample,
@@ -69,16 +66,9 @@ def _has_loss_tokens(batch: MicroBatch) -> bool:
     return any(batch.loss_mask)
 
 
-def make_flops_config():
-    return SimpleNamespace(
-        hidden_size=16,
-        num_attention_heads=2,
-        num_key_value_heads=2,
-        vocab_size=32,
-        intermediate_size=64,
-        num_hidden_layers=2,
-        head_dim=8,
-    )
+def flop_cost(seq_lens: list[int]) -> int:
+    """A FlopCounter-shaped cost: a per-token term plus a per-document attention term."""
+    return sum(64 * n + n * n for n in seq_lens)
 
 
 def test_training_sample_requires_env_name():
@@ -104,7 +94,7 @@ def test_prepare_batch_balances_micro_batches_across_workers(
         rollouts=examples,
         seq_len=4,
         num_train_workers=num_train_workers,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
     )
 
     assert all(len(worker_batches) == expected_batches_per_worker for worker_batches in batches_per_gpu)
@@ -143,7 +133,7 @@ def test_prepare_batch_preserves_branch_identity(make_training_example):
         rollouts=examples,
         seq_len=4,
         num_train_workers=2,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
         pad_to_multiple_of=8,
     )
 
@@ -177,7 +167,7 @@ def test_randomized_packing_invariants():
         num_samples = int(rng.integers(1, 65))
         lengths = [int(x) for x in rng.integers(1, seq_len + 1, size=num_samples)]
         examples = [make_sized_training_example(length, env_name=f"env-{case_idx}") for length in lengths]
-        bin_cost = build_bin_cost(make_flops_config() if case_idx % 2 == 0 else None)
+        bin_cost = flop_cost if case_idx % 2 == 0 else sum
 
         batches_per_gpu = prepare_batch(
             rollouts=examples,
@@ -222,7 +212,7 @@ def test_split_to_align_avoids_dummy_micro_batches():
         rollouts=examples,
         seq_len=12,
         num_train_workers=4,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
     )
 
     assert all(_has_loss_tokens(batch) for batch in _flatten_batches(batches_per_gpu))
@@ -236,7 +226,7 @@ def test_pack_first_then_balance_distributes_micro_batches_by_tokens_without_mod
         rollouts=examples,
         seq_len=100,
         num_train_workers=2,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
     )
 
     assert _worker_token_sums(balanced) == [170, 170]
@@ -244,7 +234,7 @@ def test_pack_first_then_balance_distributes_micro_batches_by_tokens_without_mod
 
 def test_flop_aware_balancing_pairs_long_and_short_sequence_workloads():
     examples = [make_sized_training_example(length) for length in [32, 32, 16, 16, 16, 16]]
-    bin_cost = build_bin_cost(make_flops_config())
+    bin_cost = flop_cost
 
     balanced = prepare_batch(
         rollouts=examples,
@@ -265,7 +255,7 @@ def test_flop_aware_split_to_align_splits_heaviest_flop_bin():
         rollouts=examples,
         seq_len=64,
         num_train_workers=4,
-        bin_cost=build_bin_cost(make_flops_config()),
+        bin_cost=flop_cost,
     )
 
     real_batches = [batch for batch in _flatten_batches(batches_per_gpu) if _has_loss_tokens(batch)]
@@ -283,7 +273,7 @@ def test_prepare_batch_packs_different_temperatures(make_training_example):
         rollouts=[example1, example2],
         seq_len=16,
         num_train_workers=1,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
     )
 
     flat_batches = _flatten_batches(batches_per_gpu)
@@ -342,7 +332,7 @@ def test_prepare_batch_packs_mixed_components(make_training_example, streams_on_
         rollouts=[longer, shorter],
         seq_len=16,
         num_train_workers=1,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
     )
 
     flat_batches = _flatten_batches(batches_per_gpu)
@@ -381,7 +371,7 @@ def test_prepare_batch_aligns_ref_logprobs_in_mixed_bins(make_training_example,
         seq_len=16,
         pad_to_multiple_of=1,
         num_train_workers=1,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
     )
     flat_batches = _flatten_batches(batches_per_gpu)
     assert len(flat_batches) == 1  # both samples share one bin
@@ -488,7 +478,7 @@ def test_prepare_batch_packs_multimodal_with_text():
         rollouts=[mm_sample, text_sample],
         seq_len=8,
         num_train_workers=1,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
     )
 
     real_batches = [batch for batch in _flatten_batches(batches_per_gpu) if _has_loss_tokens(batch)]
@@ -520,7 +510,7 @@ def make_mm_sample(token: int) -> TrainingSample:
         rollouts=[make_mm_sample(10), make_mm_sample(20)],
         seq_len=8,
         num_train_workers=2,
-        bin_cost=build_bin_cost(None),
+        bin_cost=sum,
     )
 
     batches = _flatten_batches(batches_per_gpu)
diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py
index 33aa612d351..f2f4b6b54f5 100644
--- a/tests/unit/train/test_perf.py
+++ b/tests/unit/train/test_perf.py
@@ -1,28 +1,222 @@
+import subprocess
+import sys
+
 import pytest
 import torch
-from transformers import AutoConfig
-
-from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL
-from prime_rl.trainer.perf import PerfCounter
 
+from prime_rl.trainer.models import (
+    AfmoeConfig,
+    AutoModelForCausalLMPrimeRL,
+    DeepseekV4Config,
+    Glm4MoeConfig,
+    GlmMoeDsaConfig,
+    GptOssConfig,
+    LagunaConfig,
+    LlamaConfig,
+    MiniMaxM2Config,
+    NemotronHConfig,
+    Qwen3_5Config,
+    Qwen3_5MoeTextConfig,
+    Qwen3_5TextConfig,
+    Qwen3_8FlashNextTextConfig,
+    Qwen3Config,
+    Qwen3MoeConfig,
+)
+from prime_rl.trainer.models.base import AttentionFlops
+from prime_rl.trainer.perf import FlopCounter
 
-# Expected values are torchtitan's `get_nparams_and_flops(model, seq_len=1024)` for its matching flavors.
-@pytest.mark.parametrize(
-    "model_name, flops_per_token",
-    [("Qwen/Qwen3-0.6B", 4_280_942_592), ("Qwen/Qwen3-30B-A3B", 20_667_125_760)],
+N = 256  # document length; sliding windows are 32 tokens and sparse attention reads 32 tokens
+DIMS = dict(
+    hidden_size=64, intermediate_size=128, vocab_size=128, num_attention_heads=4, num_key_value_heads=2, pad_token_id=0
 )
-def test_flops_per_token_matches_torchtitan(model_name: str, flops_per_token: int):
-    config = AutoConfig.from_pretrained(model_name, attn_implementation="flash_attention_2")
-    config.pad_token_id = config.eos_token_id
-    with torch.device("meta"):
-        model = AutoModelForCausalLMPrimeRL.from_config(config)
-    assert PerfCounter(model).get_flops([1024]) == 1024 * flops_per_token
+MOE = dict(moe_intermediate_size=32, num_experts_per_tok=2)
+# Softmax attention of 4 heads of qk/v dim 16 over `attended` keys per query, 6 FLOPs per multiply-add.
+FULL = 6 * 4 * (16 + 16) * N * N
+WINDOW = 6 * 4 * (16 + 16) * N * 32
+
+MODELS = {
+    "llama": (LlamaConfig(**DIMS, head_dim=16, num_hidden_layers=2), 2 * FULL),
+    "qwen3": (Qwen3Config(**DIMS, head_dim=16, num_hidden_layers=2), 2 * FULL),
+    "qwen3_moe": (Qwen3MoeConfig(**DIMS, **MOE, num_experts=4, num_hidden_layers=2), 2 * FULL),
+    "glm4_moe": (Glm4MoeConfig(**DIMS, **MOE, n_routed_experts=4, num_hidden_layers=2), 2 * FULL),
+    "minimax_m2": (MiniMaxM2Config(**DIMS, head_dim=16, num_local_experts=4, num_hidden_layers=2), 2 * FULL),
+    "afmoe": (
+        AfmoeConfig(
+            **DIMS,
+            **MOE,
+            head_dim=16,
+            num_experts=4,
+            num_hidden_layers=2,
+            layer_types=["sliding_attention", "full_attention"],
+            sliding_window=32,
+        ),
+        WINDOW + FULL,
+    ),
+    "laguna": (
+        LagunaConfig(
+            **DIMS,
+            **MOE,
+            head_dim=16,
+            num_experts=4,
+            num_hidden_layers=2,
+            layer_types=["sliding_attention", "full_attention"],
+            sliding_window=32,
+        ),
+        WINDOW + FULL,
+    ),
+    "gpt_oss": (
+        GptOssConfig(
+            **DIMS,
+            head_dim=16,
+            num_local_experts=4,
+            num_hidden_layers=2,
+            layer_types=["sliding_attention", "full_attention"],
+            sliding_window=32,
+        ),
+        WINDOW + FULL,
+    ),
+    # Mamba2: 4 heads, head_dim 16, state 8; two head_dim * state products per token and head.
+    "nemotron_h": (
+        NemotronHConfig(
+            **DIMS,
+            head_dim=16,
+            hybrid_override_pattern="M*",
+            mamba_num_heads=4,
+            mamba_head_dim=16,
+            ssm_state_size=8,
+            n_groups=1,
+        ),
+        6 * 2 * 4 * 16 * 8 * N + FULL,
+    ),
+    # Gated delta rule: 4 value heads, key/value head dim 8; three key * value products per token and head.
+    "qwen3_5": (
+        Qwen3_5TextConfig(
+            **DIMS,
+            head_dim=16,
+            layer_types=["linear_attention", "full_attention"],
+            num_hidden_layers=2,
+            linear_num_key_heads=2,
+            linear_num_value_heads=4,
+            linear_key_head_dim=8,
+            linear_value_head_dim=8,
+        ),
+        6 * 3 * 4 * 8 * 8 * N + FULL,
+    ),
+    "qwen3_5_vlm": (
+        Qwen3_5Config(
+            text_config=dict(
+                **DIMS,
+                head_dim=16,
+                layer_types=["linear_attention", "full_attention"],
+                num_hidden_layers=2,
+                linear_num_key_heads=2,
+                linear_num_value_heads=4,
+                linear_key_head_dim=8,
+                linear_value_head_dim=8,
+            )
+        ),
+        6 * 3 * 4 * 8 * 8 * N + FULL,
+    ),
+    "qwen3_5_moe": (
+        Qwen3_5MoeTextConfig(
+            **DIMS,
+            **MOE,
+            num_experts=4,
+            shared_expert_intermediate_size=32,
+            head_dim=16,
+            layer_types=["linear_attention", "full_attention"],
+            num_hidden_layers=2,
+            linear_num_key_heads=2,
+            linear_num_value_heads=4,
+            linear_key_head_dim=8,
+            linear_value_head_dim=8,
+        ),
+        6 * 3 * 4 * 8 * 8 * N + FULL,
+    ),
+    # Indexed attention reads 32 tokens; its indexer (2 heads of dim 8) scores the N / 4 compressed blocks.
+    "qwen3_8_flash_next": (
+        Qwen3_8FlashNextTextConfig(
+            **DIMS,
+            **MOE,
+            num_experts=4,
+            shared_expert_intermediate_size=32,
+            head_dim=16,
+            layer_types=["linear_attention", "full_attention"],
+            num_hidden_layers=2,
+            linear_num_key_heads=2,
+            linear_num_value_heads=4,
+            linear_key_head_dim=8,
+            linear_value_head_dim=8,
+            indexer_n_heads=2,
+            indexer_head_dim=8,
+            indexer_budget=32,
+            indexer_compress_ratio=4,
+        ),
+        6 * 3 * 4 * 8 * 8 * N + 6 * 2 * 8 * N * (N // 4) + WINDOW,
+    ),
+    # MLA with qk = 16 nope + 8 rope, v = 16, reading index_topk = 32 tokens; the indexer (2 heads of dim 8)
+    # scores all N tokens on the "full" layer, and the "shared" layer reuses its indices.
+    "glm_moe_dsa": (
+        GlmMoeDsaConfig(
+            **DIMS,
+            **MOE,
+            n_routed_experts=4,
+            num_hidden_layers=2,
+            q_lora_rank=32,
+            kv_lora_rank=16,
+            qk_nope_head_dim=16,
+            qk_rope_head_dim=8,
+            v_head_dim=16,
+            index_n_heads=2,
+            index_head_dim=8,
+            index_topk=32,
+            indexer_types=["full", "shared"],
+        ),
+        2 * 6 * 4 * (24 + 16) * N * 32 + 6 * 2 * 8 * N * N,
+    ),
+    # Every layer attends a 32-token window. CSA (rate 4) adds index_topk = 16 of its N / 4 compressed
+    # entries, which its indexer (2 heads of dim 8) scores; HCA (rate 16) adds all N / 16 = 16 entries.
+    "deepseek_v4": (
+        DeepseekV4Config(
+            **DIMS | dict(num_attention_heads=32, num_key_value_heads=1),
+            **MOE,
+            n_routed_experts=4,
+            head_dim=32,
+            num_hidden_layers=3,
+            layer_types=["sliding_attention", "compressed_sparse_attention", "heavily_compressed_attention"],
+            compress_rates={"compressed_sparse_attention": 4, "heavily_compressed_attention": 16},
+            sliding_window=32,
+            index_n_heads=2,
+            index_head_dim=8,
+            index_topk=16,
+        ),
+        6 * 32 * (32 + 32) * N * (32 + 2 * (32 + 16)) + 6 * 2 * 8 * N * (N // 4),
+    ),
+}
 
 
-def test_packed_documents_attend_within_document():
-    config = AutoConfig.from_pretrained("Qwen/Qwen3-0.6B", attn_implementation="flash_attention_2")
-    config.pad_token_id = config.eos_token_id
+def build(config):
     with torch.device("meta"):
-        model = AutoModelForCausalLMPrimeRL.from_config(config)
-    # torchtitan's FLOPs/token at seq_len 512: the 1024 value minus half its attention term, 6 * 28 * 16 * 256 * 512.
-    assert PerfCounter(model).get_flops([512, 512]) == 1024 * (4_280_942_592 - 352_321_536)
+        return AutoModelForCausalLMPrimeRL.from_config(config, attn_implementation="flash_attention_3")
+
+
+@pytest.mark.parametrize("name", MODELS)
+def test_attention_flops_match_hand_count(name: str):
+    config, attention_flops = MODELS[name]
+    counter = FlopCounter(build(config))
+    assert counter([N]) - N * counter.matmul_flops_per_token == attention_flops > 0
+
+
+@pytest.mark.parametrize("name", MODELS)
+def test_attention_flops_modules_are_not_nested(name: str):
+    model = build(MODELS[name][0])
+    for module in model.modules():
+        if isinstance(module, AttentionFlops):
+            assert not any(isinstance(child, AttentionFlops) for child in list(module.modules())[1:])
+
+
+def test_model_package_imports_without_kernel_libraries():
+    blocked = ["fla", "flash_attn", "flash_attn_interface", "flash_mla", "mamba_ssm", "prime_kernels"]
+    blocked += ["ring_flash_attn", "tilelang", "deep_gemm", "deep_ep", "sonicmoe", "quack", "cudnn", "torchao"]
+    code = f"import sys; sys.modules.update(dict.fromkeys({blocked})); import prime_rl.trainer.models, prime_rl.trainer.perf"
+    subprocess.run([sys.executable, "-c", code], check=True)

From b30c5ddb3fb462bd0f22d927ed7ed60e577dd540 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 16:18:47 +0000
Subject: [PATCH 5/9] fix(models): build every layer without the gpu-extra
 kernels

Layers bind fla / mamba_ssm / tilelang through stand-ins that keep the parameters and
raise when run, so the orchestrator's meta build works on hosts without the  extra.
DeepSeek-V4 raises the missing-tilelang error in forward; the shape check stays in __init__.
---
 .../trainer/models/deepseek_v4/attention.py   | 27 +++++++-------
 .../trainer/models/layers/optional_kernels.py | 36 +++++++++++++++++++
 .../trainer/models/nemotron_h/mamba.py        |  8 +++--
 .../trainer/models/qwen3_5/gated_delta_net.py | 15 ++++----
 .../qwen3_8_flash_next/gated_delta_net.py     | 15 ++++----
 tests/unit/train/test_perf.py                 | 15 ++++++--
 6 files changed, 78 insertions(+), 38 deletions(-)
 create mode 100644 src/prime_rl/trainer/models/layers/optional_kernels.py

diff --git a/src/prime_rl/trainer/models/deepseek_v4/attention.py b/src/prime_rl/trainer/models/deepseek_v4/attention.py
index dedc797d7ab..c01abfc9e9d 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/attention.py
+++ b/src/prime_rl/trainer/models/deepseek_v4/attention.py
@@ -139,7 +139,8 @@
 from prime_rl.utils.cp import CPContext, gather_for_cp
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
 
-# Guarded because tilelang ships in the linux-gated `gpu` extra, so some installs lack it.
+# Guarded because tilelang ships in the linux-gated `gpu` extra, so some installs lack it. The layer still
+# builds without it (the orchestrator builds the model on meta to count FLOPs) and raises in forward.
 try:
     from prime_rl.trainer.models.kernels.deepseek_v4.dsv4_sparse_attn import dsv4_sparse_attn, sparse_attn_shape_error
 except ImportError:
@@ -147,15 +148,6 @@
     sparse_attn_shape_error = None  # type: ignore
 
 
-def _kernel_blocker(num_heads: int, head_dim: int) -> str | None:
-    """Why the fused kernel cannot run at this shape, or ``None`` if it can."""
-    if dsv4_sparse_attn is None:
-        return "the tilelang sparse-attention kernel failed to import; install the `gpu` extra"
-    # CSA gives every query head the same single KV head, so the kernel's `kv_group` is 1. The
-    # shape constraints themselves are stated once, next to the kernels they come from.
-    return sparse_attn_shape_error(num_heads, 1, head_dim)
-
-
 class DeepseekV4GroupedLinear(nn.Linear):
     """Block-diagonal grouped linear, the first half of the output projection.
 
@@ -695,11 +687,11 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int, rotary_emb: Deepsee
         self.o_b_proj = nn.Linear(config.o_groups * config.o_lora_rank, config.hidden_size, bias=False)
         self.sinks = nn.Parameter(torch.zeros(self.num_heads))
         self.rotary_emb = rotary_emb
-        # Raised here rather than from the first forward, where it would surface as an ImportError
-        # or a tilelang compile failure a long way from the config that caused it.
-        blocker = _kernel_blocker(self.num_heads, self.head_dim)
-        if blocker is not None:
-            raise ValueError(f"DeepSeek V4 cannot run the fused sparse-attention kernel: {blocker}")
+        # Raised here rather than from the first forward, where it would surface as a tilelang compile
+        # failure a long way from the config that caused it. CSA gives every query head the same single
+        # KV head, so the kernel's `kv_group` is 1; the shape constraints are stated next to the kernels.
+        if dsv4_sparse_attn is not None and (shape_error := sparse_attn_shape_error(self.num_heads, 1, self.head_dim)):
+            raise ValueError(f"DeepSeek V4 cannot run the fused sparse-attention kernel: {shape_error}")
         assert config.attention_dropout == 0.0, "the fused sparse attention kernel implements no dropout"
         compressor_class = COMPRESSOR_CLASSES[self.layer_type]
         self.compressor = compressor_class(config, rotary_emb) if compressor_class is not None else None
@@ -784,6 +776,11 @@ def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[t
             top_k_indices=top_k_indices,
             window_indices=packed.window_indices,
         )
+        if dsv4_sparse_attn is None:
+            raise ValueError(
+                "DeepSeek V4 cannot run the fused sparse-attention kernel: "
+                "the tilelang sparse-attention kernel failed to import; install the `gpu` extra"
+            )
         attn_output, _ = dsv4_sparse_attn(
             q,
             inputs.kv_buf,
diff --git a/src/prime_rl/trainer/models/layers/optional_kernels.py b/src/prime_rl/trainer/models/layers/optional_kernels.py
new file mode 100644
index 00000000000..072b2856757
--- /dev/null
+++ b/src/prime_rl/trainer/models/layers/optional_kernels.py
@@ -0,0 +1,36 @@
+"""Kernels from the linux-gated `gpu` extra that layers bind at construction.
+
+Without the extra the layers still build (the orchestrator builds the model on the meta device to
+count FLOPs) and raise when run.
+"""
+
+import torch
+from torch import nn
+
+
+def unavailable(package: str):
+    def raise_import_error(*args, **kwargs):
+        raise ImportError(f"{package} is not installed; install the `gpu` extra")
+
+    return raise_import_error
+
+
+try:
+    from fla.modules import FusedRMSNormGated
+    from fla.modules.conv import causal_conv1d
+    from fla.ops.cp import build_cp_context
+    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+    from fla.ops.utils import prepare_sequence_ids
+except ImportError:
+    causal_conv1d = build_cp_context = chunk_gated_delta_rule = prepare_sequence_ids = unavailable(
+        "flash-linear-attention"
+    )
+
+    class FusedRMSNormGated(nn.Module):
+        """Holds the weight of fla's `FusedRMSNormGated`."""
+
+        def __init__(self, hidden_size: int, eps: float = 1e-5, activation: str = "swish"):
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(hidden_size))
+
+        forward = unavailable("flash-linear-attention")
diff --git a/src/prime_rl/trainer/models/nemotron_h/mamba.py b/src/prime_rl/trainer/models/nemotron_h/mamba.py
index e824fd63861..4b19ce908e9 100644
--- a/src/prime_rl/trainer/models/nemotron_h/mamba.py
+++ b/src/prime_rl/trainer/models/nemotron_h/mamba.py
@@ -13,6 +13,7 @@
 from torch import nn
 
 from prime_rl.trainer.distributed.collectives import all_to_all_single_equal
+from prime_rl.trainer.models.layers.optional_kernels import causal_conv1d, prepare_sequence_ids, unavailable
 from prime_rl.trainer.models.nemotron_h.configuration_nemotron_h import NemotronHConfig
 from prime_rl.utils.cp import CPContext
 
@@ -78,9 +79,10 @@ def forward(
 class NemotronHMamba2(nn.Module):
     def __init__(self, config: NemotronHConfig) -> None:
         super().__init__()
-        from fla.modules.conv import causal_conv1d
-        from fla.ops.utils import prepare_sequence_ids
-        from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
+        try:
+            from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
+        except ImportError:
+            mamba_chunk_scan_combined = unavailable("mamba-ssm")
 
         self.hidden_size = config.hidden_size
         self.state_size = config.ssm_state_size
diff --git a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
index f2f8b5ce983..ae0a5d41f48 100644
--- a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
@@ -3,18 +3,15 @@
 from torch import nn
 
 from prime_rl.trainer.models.layers.attn import delta_rule_flops
+from prime_rl.trainer.models.layers.optional_kernels import (
+    FusedRMSNormGated,
+    build_cp_context,
+    causal_conv1d,
+    chunk_gated_delta_rule,
+)
 from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
 from prime_rl.utils.cp import CPContext
 
-# Guarded because flash-linear-attention ships in the linux-gated `gpu` extra; the layer needs it to build.
-try:
-    from fla.modules import FusedRMSNormGated
-    from fla.modules.conv import causal_conv1d
-    from fla.ops.cp import build_cp_context
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
-except ImportError:
-    FusedRMSNormGated = causal_conv1d = build_cp_context = chunk_gated_delta_rule = None  # type: ignore
-
 # Dynamo lowers all-gather to concatenation, then fails to copy the result into
 # FLA's stacked output buffer. Keep CP convolution eager until this is fixed:
 # https://github.com/pytorch/pytorch/issues/155632
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
index 89ce2416a9a..de2368f9f03 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
@@ -3,17 +3,14 @@
 from torch import nn
 
 from prime_rl.trainer.models.layers.attn import delta_rule_flops
+from prime_rl.trainer.models.layers.optional_kernels import (
+    FusedRMSNormGated,
+    build_cp_context,
+    causal_conv1d,
+    chunk_gated_delta_rule,
+)
 from prime_rl.utils.cp import CPContext
 
-# Guarded because flash-linear-attention ships in the linux-gated `gpu` extra; the layer needs it to build.
-try:
-    from fla.modules import FusedRMSNormGated
-    from fla.modules.conv import causal_conv1d
-    from fla.ops.cp import build_cp_context
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
-except ImportError:
-    FusedRMSNormGated = causal_conv1d = build_cp_context = chunk_gated_delta_rule = None  # type: ignore
-
 # FLA's context carries a process group that Dynamo cannot trace through the convolution.
 causal_conv1d_with_context_parallelism = torch.compiler.disable(causal_conv1d)
 
diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py
index f2f4b6b54f5..00c1844823a 100644
--- a/tests/unit/train/test_perf.py
+++ b/tests/unit/train/test_perf.py
@@ -1,5 +1,6 @@
 import subprocess
 import sys
+from pathlib import Path
 
 import pytest
 import torch
@@ -215,8 +216,18 @@ def test_attention_flops_modules_are_not_nested(name: str):
             assert not any(isinstance(child, AttentionFlops) for child in list(module.modules())[1:])
 
 
-def test_model_package_imports_without_kernel_libraries():
+def test_models_build_without_kernel_libraries():
+    """The orchestrator builds the model on meta to count FLOPs, on hosts without the `gpu` extra."""
     blocked = ["fla", "flash_attn", "flash_attn_interface", "flash_mla", "mamba_ssm", "prime_kernels"]
     blocked += ["ring_flash_attn", "tilelang", "deep_gemm", "deep_ep", "sonicmoe", "quack", "cudnn", "torchao"]
-    code = f"import sys; sys.modules.update(dict.fromkeys({blocked})); import prime_rl.trainer.models, prime_rl.trainer.perf"
+    code = f"""
+import sys
+sys.modules.update(dict.fromkeys({blocked}))
+sys.path.insert(0, {str(Path(__file__).parent)!r})
+from test_perf import MODELS, N, build
+from prime_rl.trainer.perf import FlopCounter
+for config, attention_flops in MODELS.values():
+    counter = FlopCounter(build(config))
+    assert counter([N]) - N * counter.matmul_flops_per_token == attention_flops
+"""
     subprocess.run([sys.executable, "-c", code], check=True)

From 4596d3d9d3d26789c1c9741b101ec654067292fe Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 16:41:21 +0000
Subject: [PATCH 6/9] test(perf): build without kernel libraries on a CPU-only
 host, also on GPU runners

---
 tests/unit/train/test_perf.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py
index 00c1844823a..705587d14af 100644
--- a/tests/unit/train/test_perf.py
+++ b/tests/unit/train/test_perf.py
@@ -1,3 +1,4 @@
+import os
 import subprocess
 import sys
 from pathlib import Path
@@ -230,4 +231,5 @@ def test_models_build_without_kernel_libraries():
     counter = FlopCounter(build(config))
     assert counter([N]) - N * counter.matmul_flops_per_token == attention_flops
 """
-    subprocess.run([sys.executable, "-c", code], check=True)
+    # Hidden GPUs make this a CPU-only host on GPU runners too.
+    subprocess.run([sys.executable, "-c", code], check=True, env={**os.environ, "CUDA_VISIBLE_DEVICES": ""})

From 4fa181e1724971b435f67d65cbffbe5ea2dc615e Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Wed, 14 Oct 2026 18:13:19 +0000
Subject: [PATCH 7/9] refactor(perf): address review - count_active_params
 defaults to model.parameters(), orchestrator builds via get_model, drop
 kernel-free build machinery and torchtitan refs, cut test_perf to four archs

---
 src/prime_rl/orchestrator/batch.py            |  17 +-
 .../trainer/distributed/collectives.py        |   4 +-
 src/prime_rl/trainer/model.py                 |  37 +++-
 src/prime_rl/trainer/models/__init__.py       |  31 +--
 src/prime_rl/trainer/models/base.py           |  12 +-
 .../trainer/models/deepseek_v4/attention.py   |  31 +--
 .../glm_moe_dsa/configuration_glm_moe_dsa.py  |   4 +-
 .../glm_moe_dsa/sparse_mla_attention.py       |   4 +-
 src/prime_rl/trainer/models/layers/attn.py    |   4 +-
 .../trainer/models/layers/optional_kernels.py |  36 ----
 .../trainer/models/nemotron_h/mamba.py        |  12 +-
 .../trainer/models/qwen3_5/gated_delta_net.py |  10 +-
 .../models/qwen3_8_flash_next/attention.py    |  13 +-
 .../qwen3_8_flash_next/gated_delta_net.py     |  10 +-
 .../models/qwen3_8_flash_next/indexer.py      |   9 +-
 src/prime_rl/trainer/perf.py                  |   2 +-
 src/prime_rl/utils/cp.py                      |   8 +-
 tests/unit/train/test_perf.py                 | 179 +-----------------
 18 files changed, 103 insertions(+), 320 deletions(-)
 delete mode 100644 src/prime_rl/trainer/models/layers/optional_kernels.py

diff --git a/src/prime_rl/orchestrator/batch.py b/src/prime_rl/orchestrator/batch.py
index 0456c49be70..5768e2b01e5 100644
--- a/src/prime_rl/orchestrator/batch.py
+++ b/src/prime_rl/orchestrator/batch.py
@@ -5,10 +5,10 @@
 
 import numpy as np
 import torch
-from transformers import AutoConfig
 
 from prime_rl.configs.orchestrator import OrchestratorConfig
-from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL, set_pad_token_id
+from prime_rl.configs.trainer import ModelConfig
+from prime_rl.trainer.model import get_model
 from prime_rl.trainer.perf import FlopCounter
 from prime_rl.transports.batch.types import MicroBatch, MMImageRef, MMRefs, RoutedExperts, SamplingMask, TrainingSample
 
@@ -740,16 +740,11 @@ def __init__(self, config: OrchestratorConfig):
         self.seq_len = config.seq_len
         self.num_train_workers = config.num_train_workers
         self.pad_to_multiple_of = config.pad_to_multiple_of
-        # The flash-attention version selects kernels, not modules, so it does not change the FLOPs.
-        model_config = AutoConfig.from_pretrained(
-            config.model.name,
-            attn_implementation="flash_attention_3",
-            trust_remote_code=config.tokenizer.trust_remote_code,
+        # Built as the trainer builds it. The flash-attention version selects kernels, not modules.
+        model_config = ModelConfig(
+            name=config.model.name, attn="flash_attention_3", trust_remote_code=config.tokenizer.trust_remote_code
         )
-        set_pad_token_id(model_config)
-        with torch.device("meta"):
-            model = AutoModelForCausalLMPrimeRL.from_config(model_config)
-        self.bin_cost = FlopCounter(model)
+        self.bin_cost = FlopCounter(get_model(model_config, device=torch.device("meta")))
 
     def pack(self, samples: list[TrainingSample]) -> list[list[MicroBatch]]:
         return prepare_batch(
diff --git a/src/prime_rl/trainer/distributed/collectives.py b/src/prime_rl/trainer/distributed/collectives.py
index 14ea622ae6f..5a576ca7f81 100644
--- a/src/prime_rl/trainer/distributed/collectives.py
+++ b/src/prime_rl/trainer/distributed/collectives.py
@@ -1,5 +1,6 @@
 """Autograd and compile-friendly distributed collectives."""
 
+import prime_kernels
 import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
@@ -99,9 +100,6 @@ def _mxfp8_all_to_all(
     group_name: str,
     quantized: bool,
 ) -> torch.Tensor:
-    # Imported here because prime-kernels is an optional linux-only extra.
-    import prime_kernels
-
     kernel = prime_kernels.load("mxfp8_moe")
     operation = kernel.all_to_all_dispatch if quantized else kernel.all_to_all_combine
     return operation(
diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py
index b35dea6525c..5cc9a3ebd52 100644
--- a/src/prime_rl/trainer/model.py
+++ b/src/prime_rl/trainer/model.py
@@ -23,7 +23,7 @@
 from torch.distributed.fsdp._fully_shard._fsdp_init import _get_mesh_info
 from torch.distributed.tensor import Shard
 from torch.distributed.tensor.parallel import parallelize_module
-from transformers import AutoConfig, AutoTokenizer, PretrainedConfig
+from transformers import AutoConfig, AutoTokenizer, GenerationConfig, PretrainedConfig
 from transformers.tokenization_utils import PreTrainedTokenizer
 from transformers.utils.import_utils import is_flash_attn_3_available
 
@@ -45,7 +45,6 @@
     cast_float_and_contiguous,
     get_custom_causal_lm_cls,
     get_custom_vlm_cls,
-    set_pad_token_id,
     supports_custom_impl,
 )
 from prime_rl.trainer.models.deepseek_v4.attention import DeepseekV4Indexer
@@ -357,8 +356,40 @@ def get_model(
         if subconfig is not None and hasattr(subconfig, "use_cache"):
             subconfig.use_cache = False
     model_config.dsa_backend = config.dsa_backend
+    # Auto-enable IndexShare from the model's own indexer schedule (e.g. GLM-5.2). The model
+    # reads `indexer_types` directly: shared layers reuse cached indices and carry no indexer weights.
+    indexer_types = getattr(model_config, "indexer_types", None)
+    if indexer_types and any(t == "shared" for t in indexer_types):
+        model_config.use_index_cache = True
+        logger.info(
+            f"Auto-enabled IndexShare from indexer_types schedule "
+            f"({sum(t == 'full' for t in indexer_types)}/{len(indexer_types)} full layers)"
+        )
 
-    set_pad_token_id(model_config)
+    # Ensure pad_token_id is set (some models like Qwen3MoE don't have it).
+    # In transformers v5, token IDs moved from PretrainedConfig to GenerationConfig.
+    if not hasattr(model_config, "pad_token_id") or model_config.pad_token_id is None:
+        gen_config = GenerationConfig.from_model_config(model_config)
+        # Use `is not None` instead of truthiness: token ID 0 is valid.
+        pad_token_id = next(
+            (
+                v
+                for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(model_config, "eos_token_id", None)]
+                if v is not None
+            ),
+            None,
+        )
+        # Some HF configs (e.g. Llama 3.2) set pad_token_id to a list, which
+        # crashes both huggingface_hub's strict setter and transformers'
+        # GenerationConfig.validate(). Unwrap before assigning.
+        if isinstance(pad_token_id, list):
+            pad_token_id = pad_token_id[0]
+        model_config.pad_token_id = pad_token_id
+
+    # Handle list pad_token_id that was already set on the config (not from our
+    # fallback above, but directly in the model's config.json).
+    if isinstance(getattr(model_config, "pad_token_id", None), list):
+        model_config.pad_token_id = model_config.pad_token_id[0]
 
     # NOTE: For VLM models, we do NOT propagate dtype to sub_configs.
     # The model should load in its default dtype (bf16) to match vLLM inference.
diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py
index c00e7042739..af9a0684f14 100644
--- a/src/prime_rl/trainer/models/__init__.py
+++ b/src/prime_rl/trainer/models/__init__.py
@@ -2,7 +2,7 @@
 
 from collections import OrderedDict
 
-from transformers import AutoConfig, GenerationConfig
+from transformers import AutoConfig
 from transformers.configuration_utils import PretrainedConfig
 from transformers.models.auto.auto_factory import _BaseAutoModelClass, _LazyAutoMapping, auto_class_update
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES
@@ -97,35 +97,6 @@ def get_custom_causal_lm_cls(model_config: PretrainedConfig) -> type[PreTrainedM
     return _CUSTOM_CAUSAL_LM_BY_MODEL_TYPE[model_config.model_type]
 
 
-def set_pad_token_id(model_config: PretrainedConfig) -> None:
-    """Ensure pad_token_id is set (some models like Qwen3MoE don't have it).
-
-    In transformers v5, token IDs moved from PretrainedConfig to GenerationConfig.
-    """
-    if not hasattr(model_config, "pad_token_id") or model_config.pad_token_id is None:
-        gen_config = GenerationConfig.from_model_config(model_config)
-        # Use `is not None` instead of truthiness: token ID 0 is valid.
-        pad_token_id = next(
-            (
-                v
-                for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(model_config, "eos_token_id", None)]
-                if v is not None
-            ),
-            None,
-        )
-        # Some HF configs (e.g. Llama 3.2) set pad_token_id to a list, which
-        # crashes both huggingface_hub's strict setter and transformers'
-        # GenerationConfig.validate(). Unwrap before assigning.
-        if isinstance(pad_token_id, list):
-            pad_token_id = pad_token_id[0]
-        model_config.pad_token_id = pad_token_id
-
-    # Handle list pad_token_id that was already set on the config (not from our
-    # fallback above, but directly in the model's config.json).
-    if isinstance(getattr(model_config, "pad_token_id", None), list):
-        model_config.pad_token_id = model_config.pad_token_id[0]
-
-
 def supports_custom_impl(model_config: PretrainedConfig) -> bool:
     """Check if the model configuration supports the custom PrimeRL implementation.
 
diff --git a/src/prime_rl/trainer/models/base.py b/src/prime_rl/trainer/models/base.py
index 2ecd2ba5f1b..78f1a514e94 100644
--- a/src/prime_rl/trainer/models/base.py
+++ b/src/prime_rl/trainer/models/base.py
@@ -21,12 +21,14 @@ def attention_flops(self, seq_len: int) -> int:
         ...
 
 
-def count_active_params(model: nn.Module, params: Iterable[nn.Parameter]) -> int:
-    """Elements of `params` that a token multiplies through, as torchtitan's `get_nparams_and_active_nparams`.
+def count_active_params(model: nn.Module, params: Iterable[nn.Parameter] | None = None) -> int:
+    """Elements of `params` (default: all of the model's) that a token multiplies through.
 
-    Routed experts count by `top_k / num_experts`; embedding tables count only when tied to the lm_head.
+    Every parameter counts (attention and MLP projections, norms, routers, the lm_head) except the
+    routed experts a token skips (`1 - top_k / num_experts` of them) and embedding tables, which
+    count only when tied to the lm_head.
     """
-    params = {id(param): param for param in params}
+    params = {id(param): param for param in (model.parameters() if params is None else params)}
     lm_head_params = {id(param) for param in model.lm_head.parameters()}
     inactive = 0
     for module in model.modules():
@@ -113,7 +115,7 @@ def get_correct_experts_implementation(self, requested_experts: str | None) -> s
 
     def num_active_params(self) -> int:
         """Parameters a text token multiplies through, LoRA adapters included."""
-        return count_active_params(self, self.parameters())
+        return count_active_params(self)
 
     @classmethod
     def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool:
diff --git a/src/prime_rl/trainer/models/deepseek_v4/attention.py b/src/prime_rl/trainer/models/deepseek_v4/attention.py
index c01abfc9e9d..95f17e3b454 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/attention.py
+++ b/src/prime_rl/trainer/models/deepseek_v4/attention.py
@@ -139,8 +139,7 @@
 from prime_rl.utils.cp import CPContext, gather_for_cp
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
 
-# Guarded because tilelang ships in the linux-gated `gpu` extra, so some installs lack it. The layer still
-# builds without it (the orchestrator builds the model on meta to count FLOPs) and raises in forward.
+# Guarded because tilelang ships in the linux-gated `gpu` extra, so some installs lack it.
 try:
     from prime_rl.trainer.models.kernels.deepseek_v4.dsv4_sparse_attn import dsv4_sparse_attn, sparse_attn_shape_error
 except ImportError:
@@ -148,6 +147,15 @@
     sparse_attn_shape_error = None  # type: ignore
 
 
+def _kernel_blocker(num_heads: int, head_dim: int) -> str | None:
+    """Why the fused kernel cannot run at this shape, or ``None`` if it can."""
+    if dsv4_sparse_attn is None:
+        return "the tilelang sparse-attention kernel failed to import; install the `gpu` extra"
+    # CSA gives every query head the same single KV head, so the kernel's `kv_group` is 1. The
+    # shape constraints themselves are stated once, next to the kernels they come from.
+    return sparse_attn_shape_error(num_heads, 1, head_dim)
+
+
 class DeepseekV4GroupedLinear(nn.Linear):
     """Block-diagonal grouped linear, the first half of the output projection.
 
@@ -687,11 +695,11 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int, rotary_emb: Deepsee
         self.o_b_proj = nn.Linear(config.o_groups * config.o_lora_rank, config.hidden_size, bias=False)
         self.sinks = nn.Parameter(torch.zeros(self.num_heads))
         self.rotary_emb = rotary_emb
-        # Raised here rather than from the first forward, where it would surface as a tilelang compile
-        # failure a long way from the config that caused it. CSA gives every query head the same single
-        # KV head, so the kernel's `kv_group` is 1; the shape constraints are stated next to the kernels.
-        if dsv4_sparse_attn is not None and (shape_error := sparse_attn_shape_error(self.num_heads, 1, self.head_dim)):
-            raise ValueError(f"DeepSeek V4 cannot run the fused sparse-attention kernel: {shape_error}")
+        # Raised here rather than from the first forward, where it would surface as an ImportError
+        # or a tilelang compile failure a long way from the config that caused it.
+        blocker = _kernel_blocker(self.num_heads, self.head_dim)
+        if blocker is not None:
+            raise ValueError(f"DeepSeek V4 cannot run the fused sparse-attention kernel: {blocker}")
         assert config.attention_dropout == 0.0, "the fused sparse attention kernel implements no dropout"
         compressor_class = COMPRESSOR_CLASSES[self.layer_type]
         self.compressor = compressor_class(config, rotary_emb) if compressor_class is not None else None
@@ -699,8 +707,8 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int, rotary_emb: Deepsee
         self.cp_context = CPContext()
 
     def attention_flops(self, seq_len: int) -> int:
-        # As torchtitan's DeepSeek V4: each query attends the local window plus the compressed entries
-        # (CSA: the indexer scores all of them, attention reads `index_topk`).
+        # Each query attends the local window plus the compressed entries (CSA: the indexer scores
+        # all of them, attention reads `index_topk`).
         attended_tokens = min(seq_len, self.config.sliding_window)
         indexer_flops = 0
         if self.compressor is not None:
@@ -776,11 +784,6 @@ def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[t
             top_k_indices=top_k_indices,
             window_indices=packed.window_indices,
         )
-        if dsv4_sparse_attn is None:
-            raise ValueError(
-                "DeepSeek V4 cannot run the fused sparse-attention kernel: "
-                "the tilelang sparse-attention kernel failed to import; install the `gpu` extra"
-            )
         attn_output, _ = dsv4_sparse_attn(
             q,
             inputs.kv_buf,
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py
index a2855ecad87..4588b992765 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/configuration_glm_moe_dsa.py
@@ -207,9 +207,7 @@ def __init__(
         self.index_head_dim = index_head_dim
         self.indexer_rope_interleave = indexer_rope_interleave
         self.index_topk = index_topk
-        # An `indexer_types` schedule with shared layers (e.g. GLM-5.2) turns on IndexShare: shared layers
-        # reuse the previous full layer's indices and carry no indexer weights.
-        self.use_index_cache = use_index_cache or "shared" in (indexer_types or ())
+        self.use_index_cache = use_index_cache
         self.index_topk_freq = index_topk_freq
         self.index_topk_pattern = index_topk_pattern
         self.indexer_types = indexer_types
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 8350f58d47c..1c273bbac2d 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -151,8 +151,8 @@ def __init__(self, args: SparseMlaAttentionArgs):
         self.cp_context = CPContext()
 
     def attention_flops(self, seq_len: int) -> int:
-        # As torchtitan's DeepSeek V3/V4: MLA counted un-absorbed (qk = nope + rope), sparse attention
-        # attends `index_topk` tokens, and the indexer scores every token with its own heads.
+        # MLA counted un-absorbed (qk = nope + rope); sparse attention attends `index_topk` tokens, and
+        # the indexer scores every token with its own heads.
         flops = quadratic_attention_flops(
             num_heads=self.num_heads,
             qk_head_dim=self.qk_head_dim,
diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py
index 256f2694f48..c407358f138 100644
--- a/src/prime_rl/trainer/models/layers/attn.py
+++ b/src/prime_rl/trainer/models/layers/attn.py
@@ -407,7 +407,7 @@ def flash_attn_3_varlen_op(
 def quadratic_attention_flops(
     *, num_heads: int, qk_head_dim: int, v_head_dim: int, seq_len: int, attended_tokens: int | None = None
 ) -> int:
-    """Training FLOPs of softmax attention over one document, as torchtitan's `quadratic_attention_flops_per_token`.
+    """Training FLOPs of softmax attention over one document.
 
     Each of the `seq_len` queries attends `attended_tokens` keys (default all `seq_len`): two matmuls per
     head (q @ K^T over `qk_head_dim`, scores @ V over `v_head_dim`), 2 FLOPs per multiply-add, 3x for
@@ -419,7 +419,7 @@ def quadratic_attention_flops(
 
 
 def delta_rule_flops(*, num_heads: int, key_head_dim: int, v_head_dim: int, seq_len: int) -> int:
-    """Training FLOPs of a gated delta-rule recurrence over one document, as torchtitan's `delta_rule_flops_per_token`.
+    """Training FLOPs of a gated delta-rule recurrence over one document.
 
     Per token and head, reading, updating and querying the `[key_head_dim, v_head_dim]` state are three
     `key_head_dim * v_head_dim` products.
diff --git a/src/prime_rl/trainer/models/layers/optional_kernels.py b/src/prime_rl/trainer/models/layers/optional_kernels.py
deleted file mode 100644
index 072b2856757..00000000000
--- a/src/prime_rl/trainer/models/layers/optional_kernels.py
+++ /dev/null
@@ -1,36 +0,0 @@
-"""Kernels from the linux-gated `gpu` extra that layers bind at construction.
-
-Without the extra the layers still build (the orchestrator builds the model on the meta device to
-count FLOPs) and raise when run.
-"""
-
-import torch
-from torch import nn
-
-
-def unavailable(package: str):
-    def raise_import_error(*args, **kwargs):
-        raise ImportError(f"{package} is not installed; install the `gpu` extra")
-
-    return raise_import_error
-
-
-try:
-    from fla.modules import FusedRMSNormGated
-    from fla.modules.conv import causal_conv1d
-    from fla.ops.cp import build_cp_context
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
-    from fla.ops.utils import prepare_sequence_ids
-except ImportError:
-    causal_conv1d = build_cp_context = chunk_gated_delta_rule = prepare_sequence_ids = unavailable(
-        "flash-linear-attention"
-    )
-
-    class FusedRMSNormGated(nn.Module):
-        """Holds the weight of fla's `FusedRMSNormGated`."""
-
-        def __init__(self, hidden_size: int, eps: float = 1e-5, activation: str = "swish"):
-            super().__init__()
-            self.weight = nn.Parameter(torch.ones(hidden_size))
-
-        forward = unavailable("flash-linear-attention")
diff --git a/src/prime_rl/trainer/models/nemotron_h/mamba.py b/src/prime_rl/trainer/models/nemotron_h/mamba.py
index 4b19ce908e9..54b854f6b13 100644
--- a/src/prime_rl/trainer/models/nemotron_h/mamba.py
+++ b/src/prime_rl/trainer/models/nemotron_h/mamba.py
@@ -13,7 +13,6 @@
 from torch import nn
 
 from prime_rl.trainer.distributed.collectives import all_to_all_single_equal
-from prime_rl.trainer.models.layers.optional_kernels import causal_conv1d, prepare_sequence_ids, unavailable
 from prime_rl.trainer.models.nemotron_h.configuration_nemotron_h import NemotronHConfig
 from prime_rl.utils.cp import CPContext
 
@@ -79,10 +78,9 @@ def forward(
 class NemotronHMamba2(nn.Module):
     def __init__(self, config: NemotronHConfig) -> None:
         super().__init__()
-        try:
-            from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
-        except ImportError:
-            mamba_chunk_scan_combined = unavailable("mamba-ssm")
+        from fla.modules.conv import causal_conv1d
+        from fla.ops.utils import prepare_sequence_ids
+        from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
 
         self.hidden_size = config.hidden_size
         self.state_size = config.ssm_state_size
@@ -130,8 +128,8 @@ def __init__(self, config: NemotronHConfig) -> None:
         self.cp_context = CPContext()
 
     def attention_flops(self, seq_len: int) -> int:
-        # Mamba2's analogue of torchtitan's `delta_rule_flops_per_token`: per token and head, the state
-        # update (x outer B) and the readout (state @ C) are two `head_dim * state_size` products.
+        # Per token and head, the state update (x outer B) and the readout (state @ C) are two
+        # `head_dim * state_size` products.
         return 6 * 2 * self.num_heads * self.head_dim * self.state_size * seq_len
 
     def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
diff --git a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
index ae0a5d41f48..5bc68b06df9 100644
--- a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
@@ -1,14 +1,12 @@
 import torch
 import torch.nn.functional as F
+from fla.modules import FusedRMSNormGated
+from fla.modules.conv import causal_conv1d
+from fla.ops.cp import build_cp_context
+from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 from torch import nn
 
 from prime_rl.trainer.models.layers.attn import delta_rule_flops
-from prime_rl.trainer.models.layers.optional_kernels import (
-    FusedRMSNormGated,
-    build_cp_context,
-    causal_conv1d,
-    chunk_gated_delta_rule,
-)
 from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
 from prime_rl.utils.cp import CPContext
 
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
index 6b902f0804d..61966bebf70 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
@@ -22,6 +22,8 @@ def __init__(
         indexer_token_budget: int,
         indexer_compression_ratio: int,
     ) -> None:
+        import prime_kernels
+
         super().__init__()
         self.num_heads = num_attention_heads
         self.num_key_value_heads = num_key_value_heads
@@ -44,16 +46,11 @@ def __init__(
         )
 
         self.cp_context = CPContext()
-        # Resolve once before torch.compile traces the layer; the loader caches the module. A host
-        # without a GPU (the orchestrator builds the model on meta to count FLOPs) skips it.
-        if torch.cuda.is_available():
-            import prime_kernels
-
-            self.indexed_attention = prime_kernels.load("indexed_attention").indexed_attention
+        # Resolve once before torch.compile traces the layer; the loader caches the module.
+        self.indexed_attention = prime_kernels.load("indexed_attention").indexed_attention
 
     def attention_flops(self, seq_len: int) -> int:
-        # As torchtitan's DeepSeek V4 CSA: the indexer scores every compressed block,
-        # attention reads `token_budget` tokens.
+        # The indexer scores every compressed block; attention reads `token_budget` tokens.
         indexer = self.indexer
         num_blocks = seq_len // indexer.compression_ratio
         indexer_flops = 6 * indexer.num_query_heads * indexer.head_dim * seq_len * num_blocks
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
index de2368f9f03..30b05c99135 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/gated_delta_net.py
@@ -1,14 +1,12 @@
 import torch
 import torch.nn.functional as F
+from fla.modules import FusedRMSNormGated
+from fla.modules.conv import causal_conv1d
+from fla.ops.cp import build_cp_context
+from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 from torch import nn
 
 from prime_rl.trainer.models.layers.attn import delta_rule_flops
-from prime_rl.trainer.models.layers.optional_kernels import (
-    FusedRMSNormGated,
-    build_cp_context,
-    causal_conv1d,
-    chunk_gated_delta_rule,
-)
 from prime_rl.utils.cp import CPContext
 
 # FLA's context carries a process group that Dynamo cannot trace through the convolution.
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/indexer.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/indexer.py
index 81123439f39..ace92d4dfb4 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/indexer.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/indexer.py
@@ -18,13 +18,10 @@ def __init__(
         compression_ratio: int,
         norm_eps: float,
     ) -> None:
-        super().__init__()
-        # Resolved before torch.compile traces the layer. A host without a GPU (the orchestrator builds
-        # the model on meta to count FLOPs) never runs the layer and skips it.
-        if torch.cuda.is_available():
-            import prime_kernels
+        import prime_kernels
 
-            self.select_indexed_blocks = prime_kernels.load("indexed_attention").select_indexed_blocks
+        super().__init__()
+        self.select_indexed_blocks = prime_kernels.load("indexed_attention").select_indexed_blocks
         self.num_query_heads = num_query_heads
         self.head_dim = head_dim
         self.token_budget = token_budget
diff --git a/src/prime_rl/trainer/perf.py b/src/prime_rl/trainer/perf.py
index fe4c42886d8..64df1965be7 100644
--- a/src/prime_rl/trainer/perf.py
+++ b/src/prime_rl/trainer/perf.py
@@ -10,7 +10,7 @@
 
 
 class FlopCounter:
-    """Training FLOPs of packed documents, as torchtitan's `get_nparams_and_flops` per document.
+    """Training FLOPs of packed documents.
 
     6 FLOPs per active parameter per token, plus each sequence-mixing layer's FLOPs per document (a
     token attends only within its own document). The model is walked once, at construction.
diff --git a/src/prime_rl/utils/cp.py b/src/prime_rl/utils/cp.py
index 5f586dedd6c..3be423774ce 100644
--- a/src/prime_rl/utils/cp.py
+++ b/src/prime_rl/utils/cp.py
@@ -1,13 +1,16 @@
 from __future__ import annotations
 
+# ruff: noqa: I001 — `prime_rl._compat` must run before `ring_flash_attn` imports below.
+import prime_rl._compat  # noqa: F401
+
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from ring_flash_attn import update_ring_flash_attn_params
 
-import prime_rl._compat  # noqa: F401  (patches transformers before `ring_flash_attn` is imported)
 from prime_rl.trainer.distributed.collectives import all_gather
 from prime_rl.utils.logger import get_logger
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
@@ -148,9 +151,6 @@ def setup_cp_attention_params(
     )
 
     if cp_style == "ring":
-        # Imported here because ring_flash_attn ships in the linux-gated `gpu` extra.
-        from ring_flash_attn import update_ring_flash_attn_params
-
         update_ring_flash_attn_params(cu_seqlens, cp_group)
     elif cp_style == "ulysses":
         # Delayed import: ulysses_attn lives under trainer.models, which imports
diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py
index 705587d14af..41058901dee 100644
--- a/tests/unit/train/test_perf.py
+++ b/tests/unit/train/test_perf.py
@@ -1,95 +1,26 @@
-import os
-import subprocess
-import sys
-from pathlib import Path
-
 import pytest
 import torch
 
 from prime_rl.trainer.models import (
-    AfmoeConfig,
     AutoModelForCausalLMPrimeRL,
-    DeepseekV4Config,
-    Glm4MoeConfig,
     GlmMoeDsaConfig,
-    GptOssConfig,
-    LagunaConfig,
-    LlamaConfig,
-    MiniMaxM2Config,
-    NemotronHConfig,
-    Qwen3_5Config,
-    Qwen3_5MoeTextConfig,
     Qwen3_5TextConfig,
-    Qwen3_8FlashNextTextConfig,
     Qwen3Config,
     Qwen3MoeConfig,
 )
-from prime_rl.trainer.models.base import AttentionFlops
 from prime_rl.trainer.perf import FlopCounter
 
-N = 256  # document length; sliding windows are 32 tokens and sparse attention reads 32 tokens
+N = 256  # document length; sparse attention reads 32 tokens
 DIMS = dict(
     hidden_size=64, intermediate_size=128, vocab_size=128, num_attention_heads=4, num_key_value_heads=2, pad_token_id=0
 )
 MOE = dict(moe_intermediate_size=32, num_experts_per_tok=2)
-# Softmax attention of 4 heads of qk/v dim 16 over `attended` keys per query, 6 FLOPs per multiply-add.
+# Softmax attention of 4 heads of qk/v dim 16 over all N keys per query, 6 FLOPs per multiply-add.
 FULL = 6 * 4 * (16 + 16) * N * N
-WINDOW = 6 * 4 * (16 + 16) * N * 32
 
 MODELS = {
-    "llama": (LlamaConfig(**DIMS, head_dim=16, num_hidden_layers=2), 2 * FULL),
     "qwen3": (Qwen3Config(**DIMS, head_dim=16, num_hidden_layers=2), 2 * FULL),
     "qwen3_moe": (Qwen3MoeConfig(**DIMS, **MOE, num_experts=4, num_hidden_layers=2), 2 * FULL),
-    "glm4_moe": (Glm4MoeConfig(**DIMS, **MOE, n_routed_experts=4, num_hidden_layers=2), 2 * FULL),
-    "minimax_m2": (MiniMaxM2Config(**DIMS, head_dim=16, num_local_experts=4, num_hidden_layers=2), 2 * FULL),
-    "afmoe": (
-        AfmoeConfig(
-            **DIMS,
-            **MOE,
-            head_dim=16,
-            num_experts=4,
-            num_hidden_layers=2,
-            layer_types=["sliding_attention", "full_attention"],
-            sliding_window=32,
-        ),
-        WINDOW + FULL,
-    ),
-    "laguna": (
-        LagunaConfig(
-            **DIMS,
-            **MOE,
-            head_dim=16,
-            num_experts=4,
-            num_hidden_layers=2,
-            layer_types=["sliding_attention", "full_attention"],
-            sliding_window=32,
-        ),
-        WINDOW + FULL,
-    ),
-    "gpt_oss": (
-        GptOssConfig(
-            **DIMS,
-            head_dim=16,
-            num_local_experts=4,
-            num_hidden_layers=2,
-            layer_types=["sliding_attention", "full_attention"],
-            sliding_window=32,
-        ),
-        WINDOW + FULL,
-    ),
-    # Mamba2: 4 heads, head_dim 16, state 8; two head_dim * state products per token and head.
-    "nemotron_h": (
-        NemotronHConfig(
-            **DIMS,
-            head_dim=16,
-            hybrid_override_pattern="M*",
-            mamba_num_heads=4,
-            mamba_head_dim=16,
-            ssm_state_size=8,
-            n_groups=1,
-        ),
-        6 * 2 * 4 * 16 * 8 * N + FULL,
-    ),
     # Gated delta rule: 4 value heads, key/value head dim 8; three key * value products per token and head.
     "qwen3_5": (
         Qwen3_5TextConfig(
@@ -104,58 +35,6 @@
         ),
         6 * 3 * 4 * 8 * 8 * N + FULL,
     ),
-    "qwen3_5_vlm": (
-        Qwen3_5Config(
-            text_config=dict(
-                **DIMS,
-                head_dim=16,
-                layer_types=["linear_attention", "full_attention"],
-                num_hidden_layers=2,
-                linear_num_key_heads=2,
-                linear_num_value_heads=4,
-                linear_key_head_dim=8,
-                linear_value_head_dim=8,
-            )
-        ),
-        6 * 3 * 4 * 8 * 8 * N + FULL,
-    ),
-    "qwen3_5_moe": (
-        Qwen3_5MoeTextConfig(
-            **DIMS,
-            **MOE,
-            num_experts=4,
-            shared_expert_intermediate_size=32,
-            head_dim=16,
-            layer_types=["linear_attention", "full_attention"],
-            num_hidden_layers=2,
-            linear_num_key_heads=2,
-            linear_num_value_heads=4,
-            linear_key_head_dim=8,
-            linear_value_head_dim=8,
-        ),
-        6 * 3 * 4 * 8 * 8 * N + FULL,
-    ),
-    # Indexed attention reads 32 tokens; its indexer (2 heads of dim 8) scores the N / 4 compressed blocks.
-    "qwen3_8_flash_next": (
-        Qwen3_8FlashNextTextConfig(
-            **DIMS,
-            **MOE,
-            num_experts=4,
-            shared_expert_intermediate_size=32,
-            head_dim=16,
-            layer_types=["linear_attention", "full_attention"],
-            num_hidden_layers=2,
-            linear_num_key_heads=2,
-            linear_num_value_heads=4,
-            linear_key_head_dim=8,
-            linear_value_head_dim=8,
-            indexer_n_heads=2,
-            indexer_head_dim=8,
-            indexer_budget=32,
-            indexer_compress_ratio=4,
-        ),
-        6 * 3 * 4 * 8 * 8 * N + 6 * 2 * 8 * N * (N // 4) + WINDOW,
-    ),
     # MLA with qk = 16 nope + 8 rope, v = 16, reading index_topk = 32 tokens; the indexer (2 heads of dim 8)
     # scores all N tokens on the "full" layer, and the "shared" layer reuses its indices.
     "glm_moe_dsa": (
@@ -173,63 +52,17 @@
             index_head_dim=8,
             index_topk=32,
             indexer_types=["full", "shared"],
+            use_index_cache=True,
         ),
         2 * 6 * 4 * (24 + 16) * N * 32 + 6 * 2 * 8 * N * N,
     ),
-    # Every layer attends a 32-token window. CSA (rate 4) adds index_topk = 16 of its N / 4 compressed
-    # entries, which its indexer (2 heads of dim 8) scores; HCA (rate 16) adds all N / 16 = 16 entries.
-    "deepseek_v4": (
-        DeepseekV4Config(
-            **DIMS | dict(num_attention_heads=32, num_key_value_heads=1),
-            **MOE,
-            n_routed_experts=4,
-            head_dim=32,
-            num_hidden_layers=3,
-            layer_types=["sliding_attention", "compressed_sparse_attention", "heavily_compressed_attention"],
-            compress_rates={"compressed_sparse_attention": 4, "heavily_compressed_attention": 16},
-            sliding_window=32,
-            index_n_heads=2,
-            index_head_dim=8,
-            index_topk=16,
-        ),
-        6 * 32 * (32 + 32) * N * (32 + 2 * (32 + 16)) + 6 * 2 * 8 * N * (N // 4),
-    ),
 }
 
 
-def build(config):
-    with torch.device("meta"):
-        return AutoModelForCausalLMPrimeRL.from_config(config, attn_implementation="flash_attention_3")
-
-
 @pytest.mark.parametrize("name", MODELS)
 def test_attention_flops_match_hand_count(name: str):
     config, attention_flops = MODELS[name]
-    counter = FlopCounter(build(config))
+    with torch.device("meta"):
+        model = AutoModelForCausalLMPrimeRL.from_config(config, attn_implementation="flash_attention_3")
+    counter = FlopCounter(model)
     assert counter([N]) - N * counter.matmul_flops_per_token == attention_flops > 0
-
-
-@pytest.mark.parametrize("name", MODELS)
-def test_attention_flops_modules_are_not_nested(name: str):
-    model = build(MODELS[name][0])
-    for module in model.modules():
-        if isinstance(module, AttentionFlops):
-            assert not any(isinstance(child, AttentionFlops) for child in list(module.modules())[1:])
-
-
-def test_models_build_without_kernel_libraries():
-    """The orchestrator builds the model on meta to count FLOPs, on hosts without the `gpu` extra."""
-    blocked = ["fla", "flash_attn", "flash_attn_interface", "flash_mla", "mamba_ssm", "prime_kernels"]
-    blocked += ["ring_flash_attn", "tilelang", "deep_gemm", "deep_ep", "sonicmoe", "quack", "cudnn", "torchao"]
-    code = f"""
-import sys
-sys.modules.update(dict.fromkeys({blocked}))
-sys.path.insert(0, {str(Path(__file__).parent)!r})
-from test_perf import MODELS, N, build
-from prime_rl.trainer.perf import FlopCounter
-for config, attention_flops in MODELS.values():
-    counter = FlopCounter(build(config))
-    assert counter([N]) - N * counter.matmul_flops_per_token == attention_flops
-"""
-    # Hidden GPUs make this a CPU-only host on GPU runners too.
-    subprocess.run([sys.executable, "-c", code], check=True, env={**os.environ, "CUDA_VISIBLE_DEVICES": ""})

From 7dbdf0abb6b959cd6aa04a37f28e6e8ddd6e46c2 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Fri, 16 Oct 2026 17:35:17 +0000
Subject: [PATCH 8/9] test(perf): hand-count active params of a dense model;
 drop the qwen3_moe case, whose attention repeats qwen3's

---
 tests/unit/train/test_perf.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py
index 41058901dee..776792828d9 100644
--- a/tests/unit/train/test_perf.py
+++ b/tests/unit/train/test_perf.py
@@ -6,7 +6,6 @@
     GlmMoeDsaConfig,
     Qwen3_5TextConfig,
     Qwen3Config,
-    Qwen3MoeConfig,
 )
 from prime_rl.trainer.perf import FlopCounter
 
@@ -20,7 +19,6 @@
 
 MODELS = {
     "qwen3": (Qwen3Config(**DIMS, head_dim=16, num_hidden_layers=2), 2 * FULL),
-    "qwen3_moe": (Qwen3MoeConfig(**DIMS, **MOE, num_experts=4, num_hidden_layers=2), 2 * FULL),
     # Gated delta rule: 4 value heads, key/value head dim 8; three key * value products per token and head.
     "qwen3_5": (
         Qwen3_5TextConfig(
@@ -66,3 +64,11 @@ def test_attention_flops_match_hand_count(name: str):
         model = AutoModelForCausalLMPrimeRL.from_config(config, attn_implementation="flash_attention_3")
     counter = FlopCounter(model)
     assert counter([N]) - N * counter.matmul_flops_per_token == attention_flops > 0
+
+
+def test_active_params_count_attention_and_mlp_weights():
+    with torch.device("meta"):
+        model = AutoModelForCausalLMPrimeRL.from_config(MODELS["qwen3"][0], attn_implementation="flash_attention_3")
+    # Per layer: q/o 64x64, k/v 64x32, q/k norms 16, MLP 3x64x128, two norms 64; then final norm and lm_head.
+    layer = 2 * 64 * 64 + 2 * 64 * 32 + 2 * 16 + 3 * 64 * 128 + 2 * 64
+    assert model.num_active_params() == 2 * layer + 64 + 128 * 64

From bb7ca5f6a874392a0c38de3f89db099849c1f54a Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 17 Oct 2026 17:17:21 +0000
Subject: [PATCH 9/9] perf: count indexer scoring forward-only (2x); indexers
 run under no_grad

---
 src/prime_rl/trainer/models/deepseek_v4/attention.py          | 4 ++--
 .../trainer/models/glm_moe_dsa/sparse_mla_attention.py        | 4 ++--
 src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py   | 4 ++--
 tests/unit/train/test_perf.py                                 | 4 ++--
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/prime_rl/trainer/models/deepseek_v4/attention.py b/src/prime_rl/trainer/models/deepseek_v4/attention.py
index de78f1cfa1f..ade056ebfca 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/attention.py
+++ b/src/prime_rl/trainer/models/deepseek_v4/attention.py
@@ -720,14 +720,14 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int, rotary_emb: Deepsee
 
     def attention_flops(self, seq_len: int) -> int:
         # Each query attends the local window plus the compressed entries (CSA: the indexer scores
-        # all of them, attention reads `index_topk`).
+        # all of them forward only, under no_grad; attention reads `index_topk`).
         attended_tokens = min(seq_len, self.config.sliding_window)
         indexer_flops = 0
         if self.compressor is not None:
             compressed_len = seq_len // self.compressor.compress_rate
             indexer = getattr(self.compressor, "indexer", None)
             if indexer is not None:
-                indexer_flops = 6 * indexer.num_heads * indexer.head_dim * seq_len * compressed_len
+                indexer_flops = 2 * indexer.num_heads * indexer.head_dim * seq_len * compressed_len
                 compressed_len = min(compressed_len, indexer.index_topk)
             attended_tokens += compressed_len
         return indexer_flops + quadratic_attention_flops(
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 1c273bbac2d..dfbe159a9b0 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -152,7 +152,7 @@ def __init__(self, args: SparseMlaAttentionArgs):
 
     def attention_flops(self, seq_len: int) -> int:
         # MLA counted un-absorbed (qk = nope + rope); sparse attention attends `index_topk` tokens, and
-        # the indexer scores every token with its own heads.
+        # the indexer scores every token with its own heads, forward only (it runs under no_grad).
         flops = quadratic_attention_flops(
             num_heads=self.num_heads,
             qk_head_dim=self.qk_head_dim,
@@ -161,7 +161,7 @@ def attention_flops(self, seq_len: int) -> int:
             attended_tokens=min(seq_len, self.args.index_topk),
         )
         if self.indexer is not None:
-            flops += 6 * self.args.index_n_heads * self.args.index_head_dim * seq_len * seq_len
+            flops += 2 * self.args.index_n_heads * self.args.index_head_dim * seq_len * seq_len
         return flops
 
     def mla_latents(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
index 61966bebf70..a13c02749a5 100644
--- a/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
+++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/attention.py
@@ -50,10 +50,10 @@ def __init__(
         self.indexed_attention = prime_kernels.load("indexed_attention").indexed_attention
 
     def attention_flops(self, seq_len: int) -> int:
-        # The indexer scores every compressed block; attention reads `token_budget` tokens.
+        # The indexer scores every compressed block, forward only (no_grad); attention reads `token_budget` tokens.
         indexer = self.indexer
         num_blocks = seq_len // indexer.compression_ratio
-        indexer_flops = 6 * indexer.num_query_heads * indexer.head_dim * seq_len * num_blocks
+        indexer_flops = 2 * indexer.num_query_heads * indexer.head_dim * seq_len * num_blocks
         return indexer_flops + quadratic_attention_flops(
             num_heads=self.num_heads,
             qk_head_dim=self.head_dim,
diff --git a/tests/unit/train/test_perf.py b/tests/unit/train/test_perf.py
index 776792828d9..9b600ded7df 100644
--- a/tests/unit/train/test_perf.py
+++ b/tests/unit/train/test_perf.py
@@ -34,7 +34,7 @@
         6 * 3 * 4 * 8 * 8 * N + FULL,
     ),
     # MLA with qk = 16 nope + 8 rope, v = 16, reading index_topk = 32 tokens; the indexer (2 heads of dim 8)
-    # scores all N tokens on the "full" layer, and the "shared" layer reuses its indices.
+    # scores all N tokens forward only (no_grad) on the "full" layer, and the "shared" layer reuses its indices.
     "glm_moe_dsa": (
         GlmMoeDsaConfig(
             **DIMS,
@@ -52,7 +52,7 @@
             indexer_types=["full", "shared"],
             use_index_cache=True,
         ),
-        2 * 6 * 4 * (24 + 16) * N * 32 + 6 * 2 * 8 * N * N,
+        2 * 6 * 4 * (24 + 16) * N * 32 + 2 * 2 * 8 * N * N,
     ),
 }