fix(models): pass contiguous router-replay experts to compiled layers

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

@@ -367,7 +367,9 @@ def forward(
 
         for layer_idx, decoder_layer in enumerate(self.layers):
             mask = causal_mask_mapping[decoder_layer.attention_type] if causal_mask_mapping is not None else None
-            routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None
+            routed_experts_layer = (
+                routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None
+            )
 
             hidden_states = decoder_layer(
                 hidden_states,

src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py

@@ -267,7 +267,9 @@ def forward(
 
         mhc_states = inputs_embeds.unsqueeze(2).expand(-1, -1, self.config.hc_mult, -1).contiguous()
         for layer_idx, decoder_layer in enumerate(self.layers):
-            routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None
+            routed_experts_layer = (
+                routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None
+            )
             mhc_states = decoder_layer(
                 mhc_states,
                 input_ids=input_ids,

src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py

@@ -241,7 +241,9 @@ def forward(
         cached_indices = None
         use_index_cache = getattr(self.config, "use_index_cache", False)
         for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
-            routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None
+            routed_experts_layer = (
+                routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None
+            )
             hidden_states, next_cached_indices = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,

src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py

@@ -148,7 +148,7 @@ def forward(
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
         for layer_idx, decoder_layer in enumerate(self.layers):
-            layer_routed_experts = routed_experts[:, :, layer_idx] if routed_experts is not None else None
+            layer_routed_experts = routed_experts[:, :, layer_idx].contiguous() if routed_experts is not None else None
             hidden_states = decoder_layer(
                 hidden_states,
                 position_embeddings,

src/prime_rl/trainer/models/laguna/modeling_laguna.py

@@ -329,7 +329,9 @@ def forward(
         }
 
         for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
-            routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None
+            routed_experts_layer = (
+                routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None
+            )
             hidden_states = decoder_layer(
                 hidden_states,
                 attention_mask=causal_mask_mapping[self.config.layer_types[layer_idx]],

src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py

@@ -177,7 +177,9 @@ def forward(
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
-            routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None
+            routed_experts_layer = (
+                routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None
+            )
             hidden_states = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,

src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py

@@ -192,7 +192,7 @@ def forward(
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
 
         for layer_idx, decoder_layer in enumerate(self.layers):
-            layer_routed_experts = routed_experts[:, :, layer_idx] if routed_experts is not None else None
+            layer_routed_experts = routed_experts[:, :, layer_idx].contiguous() if routed_experts is not None else None
             hidden_states = decoder_layer(
                 hidden_states,
                 cu_seqlens,

src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py

@@ -210,7 +210,9 @@ def forward(
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
-            routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None
+            routed_experts_layer = (
+                routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None
+            )
             hidden_states = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,