TOP dense parity for Qwen3 on Blackwell: flashinfer attention, batch-invariant kernel seam, non-reentrant recompute

.gitignore

@@ -20,4 +20,5 @@ runs/
 
 # Sphinx documentation
 docs/_build
-docs/apidocs
+docs/apidocs
+.rxignore

megatron/core/models/gpt/gpt_layer_specs.py

@@ -56,7 +56,6 @@
     HAVE_KITCHEN = False
 
 from miles_megatron_plugins.true_on_policy.contracts import resolve_true_on_policy_runtime_policy
-from miles_megatron_plugins.true_on_policy.rope import enable_sglang_rope
 from miles_megatron_plugins.true_on_policy.runtime import enable_sglang_batch_invariant_mode
 from miles_megatron_plugins.true_on_policy.sglang_backend import (
     SGLangFinalRMSNorm,
@@ -330,7 +329,6 @@ def _select_local_backend(
     if use_true_on_policy_backend:
         assert not use_kitchen, "true_on_policy_contract is not compatible with use_kitchen."
         enable_sglang_batch_invariant_mode()
-        enable_sglang_rope()
         return SGLangSpecProvider(), True
     if use_kitchen:
         assert HAVE_KITCHEN

megatron/core/transformer/attention.py

@@ -19,17 +19,6 @@
     apply_rotary_pos_emb_with_cos_sin,
 )
 
-try:
-    from miles_megatron_plugins.true_on_policy.sglang_backend import (
-        is_sglang_rope_enabled,
-        sglang_apply_rotary_pos_emb_with_freqs,
-    )
-
-    HAVE_SGLANG_ROPE = True
-except ImportError:
-    HAVE_SGLANG_ROPE = False
-    is_sglang_rope_enabled = lambda: False
-    sglang_apply_rotary_pos_emb_with_freqs = None
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.parallel_state import (
     get_data_parallel_group,
@@ -1179,56 +1168,30 @@ def forward(
             if split_qkv:
                 ulysses_cp = _is_ulysses_cp(self.config)
                 if q_pos_emb is not None:
-                    use_sglang_rope = (
-                        HAVE_SGLANG_ROPE and is_sglang_rope_enabled() and packed_seq_params is None
-                    )
-                    sglang_rope_applied = False
-                    if use_sglang_rope and sglang_apply_rotary_pos_emb_with_freqs is not None:
-                        q_freqs, _ = (
-                            q_pos_emb if isinstance(q_pos_emb, tuple) else (q_pos_emb, q_pos_emb)
-                        )
-                        query = sglang_apply_rotary_pos_emb_with_freqs(
-                            query, q_freqs, self.config, layer_number=self.layer_number
-                        )
-                        sglang_rope_applied = True
-                    if not sglang_rope_applied:
-                        if inference_context is None or inference_context.is_static_batching():
-                            query = apply_rotary_pos_emb(
-                                query,
-                                q_pos_emb,
-                                config=self.config,
-                                cu_seqlens=cu_seqlens_q,
-                                mscale=_yarn_get_concentration_factor_from_config(self.config),
-                                cp_group=self.pg_collection.cp,
-                                ulysses_cp=ulysses_cp,
-                            )
-                        else:
-                            query = inference_context.apply_rotary_emb_query(
-                                query, q_pos_emb, self.config, cu_seqlens_q, self.pg_collection.cp
-                            )
-                if k_pos_emb is not None:
-                    use_sglang_rope = (
-                        HAVE_SGLANG_ROPE and is_sglang_rope_enabled() and packed_seq_params is None
-                    )
-                    sglang_rope_applied = False
-                    if use_sglang_rope and sglang_apply_rotary_pos_emb_with_freqs is not None:
-                        _, k_freqs = (
-                            k_pos_emb if isinstance(k_pos_emb, tuple) else (k_pos_emb, k_pos_emb)
-                        )
-                        key = sglang_apply_rotary_pos_emb_with_freqs(
-                            key, k_freqs, self.config, layer_number=self.layer_number
-                        )
-                        sglang_rope_applied = True
-                    if not sglang_rope_applied:
-                        key = apply_rotary_pos_emb(
-                            key,
-                            k_pos_emb,
+                    if inference_context is None or inference_context.is_static_batching():
+                        query = apply_rotary_pos_emb(
+                            query,
+                            q_pos_emb,
                             config=self.config,
-                            cu_seqlens=cu_seqlens_kv,
+                            cu_seqlens=cu_seqlens_q,
                             mscale=_yarn_get_concentration_factor_from_config(self.config),
                             cp_group=self.pg_collection.cp,
                             ulysses_cp=ulysses_cp,
                         )
+                    else:
+                        query = inference_context.apply_rotary_emb_query(
+                            query, q_pos_emb, self.config, cu_seqlens_q, self.pg_collection.cp
+                        )
+                if k_pos_emb is not None:
+                    key = apply_rotary_pos_emb(
+                        key,
+                        k_pos_emb,
+                        config=self.config,
+                        cu_seqlens=cu_seqlens_kv,
+                        mscale=_yarn_get_concentration_factor_from_config(self.config),
+                        cp_group=self.pg_collection.cp,
+                        ulysses_cp=ulysses_cp,
+                    )
             else:
                 query, key, value = apply_fused_qkv_rotary_pos_emb(
                     mixed_qkv, q_pos_emb, k_pos_emb, qkv_split_arg_list

megatron/core/transformer/transformer_block.py

@@ -124,7 +124,8 @@ def hook(grad: Tensor) -> Tensor:
 
 
 def _get_sglang_cp_recompute_mode() -> str:
-    mode = os.environ.get("MEGATRON_TRUE_ON_POLICY_SGLANG_CP_RECOMPUTE", "disabled")
+    # Default non_reentrant: reentrant checkpoint NaNs under TOP.
+    mode = os.environ.get("MEGATRON_TRUE_ON_POLICY_SGLANG_CP_RECOMPUTE", "non_reentrant")
     return mode.lower().replace("-", "_")
 
 
@@ -579,15 +580,15 @@ def custom_forward(
         def checkpoint_handler(forward_func):
             """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
             true_on_policy_policy = resolve_true_on_policy_runtime_policy(self.config)
-            if true_on_policy_policy.use_ulysses_cp_recompute_fallback:
+            if true_on_policy_policy.use_non_reentrant_recompute:
                 mode = _get_sglang_cp_recompute_mode()
                 global _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK
                 if mode in ("disabled", "disable", "off", "none", "false", "0"):
                     if not _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK:
                         logger.info(
                             "Bypassing full activation recompute for SGLang Ulysses CP. "
                             "Set MEGATRON_TRUE_ON_POLICY_SGLANG_CP_RECOMPUTE to "
-                            "'non_reentrant' or 'reentrant' to override."
+                            "'non_reentrant' (default) or 'reentrant' to override."
                         )
                         _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK = True
                     return forward_func(
@@ -601,7 +602,7 @@ def checkpoint_handler(forward_func):
                 if mode == "non_reentrant":
                     if not _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK:
                         logger.info(
-                            "Using non-reentrant torch checkpoint for SGLang Ulysses CP "
+                            "Using non-reentrant torch checkpoint for true-on-policy "
                             "full recompute."
                         )
                         _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK = True