From 7f74336872a0538548f0b5afcf931b9ac23fcd14 Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Sat, 18 Jul 2026 03:18:27 -0500
Subject: [PATCH 1/9] chore: gitignore .rxignore

---
 .gitignore | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 239d70c88f1..c3d9047f893 100644
--- a/.gitignore
+++ b/.gitignore
@@ -20,4 +20,5 @@ runs/
 
 # Sphinx documentation
 docs/_build
-docs/apidocs
\ No newline at end of file
+docs/apidocs
+.rxignore

From 93e2aea6d737622f77d3a2a253d6432b95a48586 Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Sat, 18 Jul 2026 03:18:27 -0500
Subject: [PATCH 2/9] feat(top): batch-invariant kernel seam with analytic
 backward

---
 .../true_on_policy/kernels.py                 |  85 ++++++++++
 .../true_on_policy/matmul.py                  | 149 +++---------------
 miles_megatron_plugins/true_on_policy/norm.py |  10 +-
 3 files changed, 113 insertions(+), 131 deletions(-)
 create mode 100644 miles_megatron_plugins/true_on_policy/kernels.py

diff --git a/miles_megatron_plugins/true_on_policy/kernels.py b/miles_megatron_plugins/true_on_policy/kernels.py
new file mode 100644
index 00000000000..d04fbf06d0a
--- /dev/null
+++ b/miles_megatron_plugins/true_on_policy/kernels.py
@@ -0,0 +1,85 @@
+"""Adapter seam for true-on-policy parity-critical kernels.
+
+This is the ONE place the Megatron true-on-policy backend imports the invariant
+kernels; they delegate to SGLang's implementations (SGLang ships these as its
+deterministic / true-on-policy feature). The rest of the plugin calls these functions
+and never names SGLang directly, so parity comes from running the *same* kernel as
+inference -- not a reimplemented copy that can silently drift (see matmul.py history).
+
+Pattern: wrap SGLang's forward-only kernel in a ``torch.autograd.Function`` whose
+backward is the standard analytic linear vjp (plain GEMMs). The backward has no
+inference counterpart, so it needs no special kernel and is precision-agnostic.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+class _TpInvRowLinear(torch.autograd.Function):
+    """Row-linear local GEMM delegated to SGLang's exact TP-invariant kernel."""
+
+    @staticmethod
+    def forward(ctx, input_2d: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+        import sglang.srt.tp_invariant_ops  # noqa: F401  (registers torch.ops.tp_inv_ops)
+
+        ctx.save_for_backward(input_2d, weight)
+        # weight is [out, K_local] (Megatron RowParallelLinear); the kernel wants [K, N].
+        return torch.ops.tp_inv_ops.matmul_tp_inv(input_2d.contiguous(), weight.t(), None)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        input_2d, weight = ctx.saved_tensors
+        grad_input = torch.matmul(grad_output, weight)
+        grad_weight = torch.matmul(grad_output.transpose(-2, -1), input_2d)
+        return grad_input, grad_weight
+
+
+def tp_invariant_row_linear(input_2d: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    """Per-rank row-linear local GEMM, bitwise-identical to SGLang's TP-invariant kernel.
+
+    Args:
+        input_2d: ``[tokens, K_local]`` activation shard.
+        weight:   ``[out, K_local]`` (Megatron ``RowParallelLinear`` layout).
+    Returns ``[tokens, out]``; differentiable (backward = analytic linear vjp).
+    """
+    return _TpInvRowLinear.apply(input_2d, weight)
+
+
+class _RmsNormBatchInvariant(torch.autograd.Function):
+    """RMS-normalize (weight=1) via SGLang's exact batch-invariant kernel.
+
+    Forward runs the same forward-only Triton kernel SGLang uses (bitwise inference
+    parity). That kernel is non-differentiable (writes an ``empty_like`` output, no
+    grad_fn), so training needs this wrapper: backward is the analytic RMSNorm vjp
+    (plain torch), which has no inference counterpart and need not match a kernel.
+    """
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, eps: float) -> torch.Tensor:
+        from sglang.srt.batch_invariant_ops import rms_norm_batch_invariant
+
+        ones = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
+        ctx.save_for_backward(x)
+        ctx.eps = eps
+        return rms_norm_batch_invariant(x, ones, eps)
+
+    @staticmethod
+    def backward(ctx, grad_y: torch.Tensor):
+        (x,) = ctx.saved_tensors
+        # y = x / sqrt(mean(x^2) + eps); vjp over the last dim (H):
+        #   grad_x = (grad_y - x * mean(grad_y * x) / ms) / rms
+        ms = x.pow(2).mean(-1, keepdim=True) + ctx.eps
+        rms = ms.sqrt()
+        dot = (grad_y * x).mean(-1, keepdim=True)
+        grad_x = (grad_y - x * (dot / ms)) / rms
+        return grad_x, None
+
+
+def rms_norm_batch_invariant(x: torch.Tensor, eps: float) -> torch.Tensor:
+    """Batch-invariant RMS-normalize (implicit unit weight), differentiable.
+
+    Forward is SGLang's exact kernel (bitwise inference parity); backward is the analytic
+    RMSNorm vjp. The affine ``weight`` multiply stays in the caller (already autograd-safe).
+    """
+    return _RmsNormBatchInvariant.apply(x, eps)
diff --git a/miles_megatron_plugins/true_on_policy/matmul.py b/miles_megatron_plugins/true_on_policy/matmul.py
index c53a1fabfa4..94f9c03a865 100644
--- a/miles_megatron_plugins/true_on_policy/matmul.py
+++ b/miles_megatron_plugins/true_on_policy/matmul.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Iterable, List, Optional
+from typing import List, Optional
 
 import torch
 
@@ -9,135 +9,27 @@
     linear_with_grad_accumulation_and_async_allreduce,
 )
 
-_ROW_LINEAR_INV_BLOCK_K = 128
-
-
-def _fixed_tree_sum_tensors(tensors: Iterable[torch.Tensor]) -> torch.Tensor:
-    """Sum tensors in the same fixed pairwise order as SGLang."""
-    partials = list(tensors)
-    if not partials:
-        raise ValueError("at least one tensor is required")
-
-    while len(partials) > 1:
-        next_partials = []
-        for index in range(0, len(partials), 2):
-            if index + 1 < len(partials):
-                next_partials.append(partials[index] + partials[index + 1])
-            else:
-                next_partials.append(partials[index])
-        partials = next_partials
-
-    return partials[0]
-
-
-def _safe_group_size(group: Optional[torch.distributed.ProcessGroup]) -> int:
-    if group is not None:
-        return group.size()
-    try:
-        from megatron.core.parallel_state import get_tensor_model_parallel_world_size
-
-        return get_tensor_model_parallel_world_size()
-    except Exception:
-        return 1
-
-
-def _safe_tensor_context_parallel_size() -> int:
-    try:
-        from megatron.core.parallel_state import get_tensor_and_context_parallel_world_size
-
-        return get_tensor_and_context_parallel_world_size()
-    except Exception:
-        return _safe_group_size(None)
-
-
-def _rollout_row_parallel_partition_k(
-    input_: torch.Tensor, tp_group: Optional[torch.distributed.ProcessGroup]
-) -> int:
-    train_tp_size = _safe_group_size(tp_group)
-    rollout_tp_size = _safe_tensor_context_parallel_size()
-    global_k_size = input_.shape[-1] * train_tp_size
-    if rollout_tp_size <= 0 or global_k_size % rollout_tp_size != 0:
-        return input_.shape[-1]
-    return global_k_size // rollout_tp_size
-
-
-def _should_use_sglang_tp_invariant_row_linear(
-    input_: torch.Tensor, row_parallel: bool, tp_group: Optional[torch.distributed.ProcessGroup]
-) -> bool:
-    rollout_partition_k = _rollout_row_parallel_partition_k(input_, tp_group)
-    return (
-        row_parallel
-        and rollout_partition_k >= _ROW_LINEAR_INV_BLOCK_K
-        and rollout_partition_k % _ROW_LINEAR_INV_BLOCK_K == 0
-    )
+from . import kernels
 
 
 def _sglang_row_parallel_matmul(
     input_: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
 ) -> torch.Tensor:
-    """SGLang's row-linear TP-invariant matmul contract.
+    """SGLang row-linear local GEMM, delegated through the shared kernel seam.
 
-    SGLang chunks the K dimension into 128-wide products, casts each product to
-    the input dtype, then combines those partials with a fixed binary tree.
-    Mirroring that order is required before the TP tree all-reduce can be
-    bitwise identical.
+    Routes to ``kernels.tp_invariant_row_linear`` -> SGLang's ``matmul_tp_inv`` (a
+    two-level-tree TP-invariant matmul), so this rank's partial is bitwise-identical to
+    inference; the subsequent TP all-reduce (a no-op at TP=1) then combines identical
+    partials.
     """
     input_shape = input_.shape
     input_2d = input_.reshape(-1, input_shape[-1])
-    weight_t = weight.t()
-    partials = []
-
-    for start in range(0, input_2d.shape[1], _ROW_LINEAR_INV_BLOCK_K):
-        end = min(start + _ROW_LINEAR_INV_BLOCK_K, input_2d.shape[1])
-        partials.append(input_2d[:, start:end] @ weight_t[start:end, :])
-
-    output = _fixed_tree_sum_tensors(partials).to(input_.dtype)
+    output = kernels.tp_invariant_row_linear(input_2d, weight)
     if bias is not None:
         output = output + bias
     return output.reshape(*input_shape[:-1], weight.shape[0])
 
 
-def _sglang_rollout_partition_row_parallel_matmul(
-    input_: torch.Tensor,
-    weight: torch.Tensor,
-    bias: Optional[torch.Tensor],
-    *,
-    tp_group: Optional[torch.distributed.ProcessGroup],
-) -> torch.Tensor:
-    """Mirror SGLang rollout row-linear shards when train TP is smaller than rollout TP."""
-    rollout_partition_k = _rollout_row_parallel_partition_k(input_, tp_group)
-    if (
-        rollout_partition_k <= 0
-        or rollout_partition_k >= input_.shape[-1]
-        or input_.shape[-1] % rollout_partition_k != 0
-    ):
-        return _linear_reference_matmul(input_, weight, bias)
-
-    input_shape = input_.shape
-    input_2d = input_.reshape(-1, input_shape[-1])
-    weight_t = weight.t()
-    partials = []
-
-    for start in range(0, input_2d.shape[1], rollout_partition_k):
-        end = start + rollout_partition_k
-        partials.append(input_2d[:, start:end] @ weight_t[start:end, :])
-
-    output = _fixed_tree_sum_tensors(partials).to(input_.dtype)
-    if bias is not None:
-        output = output + bias
-    return output.reshape(*input_shape[:-1], weight.shape[0])
-
-
-def _linear_reference_matmul(
-    input_: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
-) -> torch.Tensor:
-    output = input_.reshape(-1, input_.shape[-1]) @ weight.t()
-    output = output.reshape(*input_.shape[:-1], weight.shape[0])
-    if bias is not None:
-        output = output + bias
-    return output
-
-
 def sglang_reference_matmul(
     input_: torch.Tensor,
     weight: torch.Tensor,
@@ -151,25 +43,26 @@ def sglang_reference_matmul(
     tp_group: Optional[torch.distributed.ProcessGroup] = None,
     row_parallel: bool = False,
 ) -> torch.Tensor:
-    """Reference TP matmul entrypoint for the SGLang-compatible backend.
-
-    PR 6 keeps Megatron on the same local numerical path by default and introduces a
-    single surface that later PRs can specialize for TP-invariant ordering. The
-    implementation intentionally delegates to the existing Megatron kernels so enabling
-    the backend flag does not yet change the training contract.
+    """Matmul entry point for the SGLang-compatible (true-on-policy) backend -- a router.
+
+    Row-parallel linears (o_proj, down_proj) reduce partials across TP ranks, so their
+    reduction order must match SGLang's inference GEMM bitwise. They ALWAYS route through
+    ``matmul_tp_inv`` (``_sglang_row_parallel_matmul``); SGLang does the same
+    unconditionally. Because ``matmul_tp_inv`` is TP-degree-invariant, both engines then
+    agree for every (train_tp, rollout_tp) -- including TP=1/TP=1, where the tree is
+    marginally slower than a single GEMM but never mismatches. Using it always (rather
+    than gating on tp) makes divergence impossible and needs no cross-engine tp signal.
+
+    Column-parallel linears and frozen weights have no cross-rank K-reduction and fall
+    through to the stock Megatron kernels, numerically unchanged.
     """
-
     if input_.dtype != weight.dtype:
         input_ = input_.to(weight.dtype)
     if bias is not None and bias.dtype != weight.dtype:
         bias = bias.to(weight.dtype)
 
-    if _should_use_sglang_tp_invariant_row_linear(input_, row_parallel, tp_group):
-        return _sglang_row_parallel_matmul(input_, weight, bias)
     if row_parallel:
-        return _sglang_rollout_partition_row_parallel_matmul(
-            input_, weight, bias, tp_group=tp_group
-        )
+        return _sglang_row_parallel_matmul(input_, weight, bias)
 
     if weight.requires_grad:
         return linear_with_grad_accumulation_and_async_allreduce(
diff --git a/miles_megatron_plugins/true_on_policy/norm.py b/miles_megatron_plugins/true_on_policy/norm.py
index 0ce2da3514e..3dac0dfd5e3 100644
--- a/miles_megatron_plugins/true_on_policy/norm.py
+++ b/miles_megatron_plugins/true_on_policy/norm.py
@@ -6,6 +6,10 @@
 import torch.nn.functional as F
 
 from megatron.core.transformer.transformer_config import TransformerConfig
+# kernels.rms_norm_batch_invariant runs SGLang's exact batch-invariant RMSNorm kernel
+# (bitwise inference parity) wrapped in an autograd.Function with the analytic RMSNorm vjp;
+# the bare kernel is forward-only (no grad_fn), so it must be routed through the seam.
+from . import kernels
 from .contracts import resolve_true_on_policy_runtime_policy
 
 
@@ -92,7 +96,7 @@ def forward(
                 x_float = x_float + post_residual_addition.float()
             residual = x_float.to(orig_dtype)
 
-        output = x_float * torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps)
+        output = kernels.rms_norm_batch_invariant(x_float, self.eps)
         if self.cast_x_before_out_mul:
             output = self.weight.float() * output.to(orig_dtype)
         else:
@@ -135,7 +139,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         orig_dtype = x.dtype
         x_float = x.to(torch.float32)
-        x_float = x_float * torch.rsqrt(x_float.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        x_float = kernels.rms_norm_batch_invariant(x_float, self.eps)
 
         if self.cast_x_before_out_mul:
             return self.weight.float() * x_float.to(orig_dtype)
@@ -184,7 +188,7 @@ def forward(
             residual = x.clone()
 
         x_float = x.to(torch.float32)
-        x_float = x_float * torch.rsqrt(x_float.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        x_float = kernels.rms_norm_batch_invariant(x_float, self.eps)
         output = self.weight * x_float.to(orig_dtype)
 
         if residual is not None:

From 4e8f82cf15a0b817282006db0e31b1c1c7311f59 Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Sat, 18 Jul 2026 03:18:27 -0500
Subject: [PATCH 3/9] =?UTF-8?q?feat(top):=20dense=20parity=20=E2=80=94=20O?=
 =?UTF-8?q?(L)=20attention=20backward=20+=20non-reentrant=20recompute?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 megatron/core/models/gpt/gpt_layer_specs.py   |   2 -
 megatron/core/transformer/attention.py        |  75 ++-----
 .../core/transformer/transformer_block.py     |   9 +-
 .../true_on_policy/contracts.py               |   8 +-
 .../true_on_policy/provider.py                |   2 +-
 miles_megatron_plugins/true_on_policy/rope.py |  87 --------
 .../{attention_fa3.py => sglang_attention.py} | 185 ++++++++++++++----
 .../true_on_policy/sglang_backend.py          |  14 +-
 .../extension/test_sglang_extension.py        |  21 +-
 9 files changed, 190 insertions(+), 213 deletions(-)
 delete mode 100644 miles_megatron_plugins/true_on_policy/rope.py
 rename miles_megatron_plugins/true_on_policy/{attention_fa3.py => sglang_attention.py} (51%)

diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
index 581e3fb9207..f41829e780c 100755
--- a/megatron/core/models/gpt/gpt_layer_specs.py
+++ b/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
diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index cc5ab22c648..3f78856d6bb 100644
--- a/megatron/core/transformer/attention.py
+++ b/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
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index f42e7321872..52611b504e3 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/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,7 +580,7 @@ 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"):
@@ -587,7 +588,7 @@ def checkpoint_handler(forward_func):
                         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
diff --git a/miles_megatron_plugins/true_on_policy/contracts.py b/miles_megatron_plugins/true_on_policy/contracts.py
index 8b2e82d87f7..83cd68db276 100644
--- a/miles_megatron_plugins/true_on_policy/contracts.py
+++ b/miles_megatron_plugins/true_on_policy/contracts.py
@@ -41,7 +41,7 @@ class MegatronTrueOnPolicyRuntimePolicy:
     apply_logits_contract: bool
     use_sglang_final_norm: bool
     use_sglang_residual_pair: bool
-    use_ulysses_cp_recompute_fallback: bool
+    use_non_reentrant_recompute: bool
 
 
 DEFAULT_RUNTIME_POLICY = MegatronTrueOnPolicyRuntimePolicy(
@@ -62,7 +62,7 @@ class MegatronTrueOnPolicyRuntimePolicy:
     apply_logits_contract=False,
     use_sglang_final_norm=False,
     use_sglang_residual_pair=False,
-    use_ulysses_cp_recompute_fallback=False,
+    use_non_reentrant_recompute=False,
 )
 
 
@@ -98,7 +98,9 @@ def policy_for(self, config) -> MegatronTrueOnPolicyRuntimePolicy:
             apply_logits_contract=True,
             use_sglang_final_norm=True,
             use_sglang_residual_pair=True,
-            use_ulysses_cp_recompute_fallback=uses_ulysses_cp,
+            # Non-reentrant activation recompute for all TOP runs (reentrant NaNs
+            # under TOP; validated tp=4 abs_diff==0). Generalizes the CP-scoped fallback.
+            use_non_reentrant_recompute=True,
         )
 
 
diff --git a/miles_megatron_plugins/true_on_policy/provider.py b/miles_megatron_plugins/true_on_policy/provider.py
index 039ce3604e1..9e038e21fe9 100644
--- a/miles_megatron_plugins/true_on_policy/provider.py
+++ b/miles_megatron_plugins/true_on_policy/provider.py
@@ -9,7 +9,7 @@
 from megatron.core.transformer.mlp import MLPSubmodules
 from megatron.core.transformer.moe.experts import GroupedMLP, SequentialMLP
 from megatron.core.transformer.spec_utils import ModuleSpec
-from .attention_fa3 import SGLangCoreAttention
+from .sglang_attention import SGLangCoreAttention
 from .linear import SGLangColumnParallelLinear, SGLangRowParallelLinear
 from .norm import SGLangNorm, SGLangQKRMSNorm
 
diff --git a/miles_megatron_plugins/true_on_policy/rope.py b/miles_megatron_plugins/true_on_policy/rope.py
deleted file mode 100644
index 1e40468f37c..00000000000
--- a/miles_megatron_plugins/true_on_policy/rope.py
+++ /dev/null
@@ -1,87 +0,0 @@
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-from torch import Tensor
-
-from megatron.core.transformer.transformer_config import TransformerConfig
-
-_USE_SGLANG_ROPE = False
-
-
-def enable_sglang_rope() -> None:
-    """Enable the SGLang-compatible RoPE path used by dense true-on-policy."""
-
-    global _USE_SGLANG_ROPE
-    _USE_SGLANG_ROPE = True
-
-
-def disable_sglang_rope() -> None:
-    global _USE_SGLANG_ROPE
-    _USE_SGLANG_ROPE = False
-
-
-def is_sglang_rope_enabled() -> bool:
-    return _USE_SGLANG_ROPE
-
-
-def sglang_apply_rotary_pos_emb(
-    x: Tensor, cos: Tensor, sin: Tensor, is_neox_style: bool = True
-) -> Tensor:
-    if cos.dim() == 2:
-        cos = cos.unsqueeze(-2)
-        sin = sin.unsqueeze(-2)
-
-    orig_dtype = x.dtype
-    x = x.float()
-    cos = cos.float()
-    sin = sin.float()
-
-    rotary_dim = cos.shape[-1] * 2
-    if rotary_dim < x.shape[-1]:
-        x_rot = x[..., :rotary_dim]
-        x_pass = x[..., rotary_dim:]
-        x_rot = sglang_apply_rotary_pos_emb(x_rot, cos, sin, is_neox_style)
-        return torch.cat((x_rot, x_pass), dim=-1).to(orig_dtype)
-
-    if is_neox_style:
-        x1, x2 = torch.chunk(x, 2, dim=-1)
-    else:
-        x1 = x[..., ::2]
-        x2 = x[..., 1::2]
-
-    o1 = x1 * cos - x2 * sin
-    o2 = x2 * cos + x1 * sin
-
-    if is_neox_style:
-        return torch.cat((o1, o2), dim=-1).to(orig_dtype)
-
-    return torch.stack((o1, o2), dim=-1).flatten(-2).to(orig_dtype)
-
-
-def sglang_apply_rotary_pos_emb_with_freqs(
-    x: Tensor, freqs: Tensor, config: TransformerConfig, layer_number: Optional[int] = None
-) -> Tensor:
-    del layer_number
-
-    x_seq_len = x.shape[0]
-    freqs_seq_len = freqs.shape[0]
-
-    freqs_flat = freqs.squeeze(1).squeeze(1)
-    head_dim = x.shape[-1]
-    raw_angles = freqs_flat[..., : head_dim // 2]
-    cos = torch.cos(raw_angles)
-    sin = torch.sin(raw_angles)
-    is_neox_style = not getattr(config, "rotary_interleaved", False)
-
-    if x_seq_len == freqs_seq_len:
-        return sglang_apply_rotary_pos_emb(x, cos, sin, is_neox_style)
-    if freqs_seq_len < x_seq_len:
-        x_valid = x[:freqs_seq_len]
-        x_valid = sglang_apply_rotary_pos_emb(x_valid, cos, sin, is_neox_style)
-        return torch.cat([x_valid, x[freqs_seq_len:]], dim=0)
-
-    cos = cos[:x_seq_len]
-    sin = sin[:x_seq_len]
-    return sglang_apply_rotary_pos_emb(x, cos, sin, is_neox_style)
diff --git a/miles_megatron_plugins/true_on_policy/attention_fa3.py b/miles_megatron_plugins/true_on_policy/sglang_attention.py
similarity index 51%
rename from miles_megatron_plugins/true_on_policy/attention_fa3.py
rename to miles_megatron_plugins/true_on_policy/sglang_attention.py
index 7e5c7fe62f0..f888d283520 100644
--- a/miles_megatron_plugins/true_on_policy/attention_fa3.py
+++ b/miles_megatron_plugins/true_on_policy/sglang_attention.py
@@ -30,6 +30,97 @@
         fa3_varlen_func = None
 
 
+_FI_RAGGED_WRAPPER = None
+
+
+def _get_flashinfer_ragged_wrapper(device):
+    """One shared deterministic ragged-prefill wrapper (Blackwell / flashinfer)."""
+    global _FI_RAGGED_WRAPPER
+    if _FI_RAGGED_WRAPPER is None:
+        from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
+
+        # 1 GiB workspace: the deterministic ragged prefill needs batch_prefill_tmp_v
+        # scratch that scales with heads x head_dim x fixed_split tiles; 256 MiB overflows
+        # at the canonical config (8192 response window, batch 256) -> needs ~406 MiB.
+        ws = torch.empty(1024 * 1024 * 1024, dtype=torch.uint8, device=device)
+        _FI_RAGGED_WRAPPER = BatchPrefillWithRaggedKVCacheWrapper(ws, kv_layout="NHD")
+    return _FI_RAGGED_WRAPPER
+
+
+class _FlashinferRaggedAttn(torch.autograd.Function):
+    """Deterministic flashinfer ragged prefill. Forward = flashinfer (matches sglang
+    inference bitwise). flashinfer's run is forward-only, so backward is the fused
+    mem-efficient attention backward dispatched per packed segment (needn't bit-match
+    inference)."""
+
+    @staticmethod
+    def forward(ctx, q, k, v, cu_q, cu_k, num_q_heads, num_kv_heads, head_dim, scale):
+        w = _get_flashinfer_ragged_wrapper(q.device)
+        # Match sglang's call exactly: plan WITHOUT sm_scale, apply sm_scale at forward-time
+        # via .forward (sglang uses fast_prefill_plan (no sm_scale) + .forward(..., sm_scale=...)).
+        w.plan(
+            cu_q, cu_k, num_q_heads, num_kv_heads, head_dim,
+            causal=True,
+            q_data_type=q.dtype, kv_data_type=k.dtype, fixed_split_size=4096,
+        )
+        out = w.forward(q, k, v, causal=True, sm_scale=scale, logits_soft_cap=0.0)
+        ctx.save_for_backward(q, k, v, cu_q)
+        ctx.num_q_heads, ctx.num_kv_heads, ctx.scale = num_q_heads, num_kv_heads, scale
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        # Fused mem-efficient attention backward, dispatched DIRECTLY per packed segment.
+        # Two deliberate properties:
+        #   * O(L) memory: each segment uses is_causal=True with NO dense [L, L] mask
+        #     tensor, so the fused backend tiles the attention (never materializes the
+        #     score matrix). Peak memory is O(max segment length), not O(seq_len^2) --
+        #     essential for long single segments (a full-length sample) and long context.
+        #   * NO torch.autograd.grad: dispatching aten's fused backward kernel directly is
+        #     reentrant-safe, unlike a recompute-then-autograd.grad, which nests a second
+        #     autograd traversal inside the Function.backward. Same rule as the norm and
+        #     row-linear backwards.
+        # Flashinfer is forward-only and FA3/flash-attn varlen is Hopper-only, so on
+        # Blackwell the portable fused varlen backward is aten's mem-efficient kernel run
+        # per segment with is_causal (the packed sequence's segments are single samples).
+        q, k, v, cu_q = ctx.saved_tensors
+        H, HKV, scale = ctx.num_q_heads, ctx.num_kv_heads, ctx.scale
+        rep = H // HKV
+        head_dim = q.shape[-1]
+        dq = torch.zeros_like(q)
+        dk = torch.zeros_like(k)
+        dv = torch.zeros_like(v)
+        cu = cu_q.long()
+        eff_attn = torch.ops.aten._scaled_dot_product_efficient_attention
+        eff_attn_bwd = torch.ops.aten._scaled_dot_product_efficient_attention_backward
+        for i in range(cu.numel() - 1):
+            s, e = int(cu[i]), int(cu[i + 1])
+            seg_len = e - s
+            if seg_len == 0:
+                continue
+            # [1, H, L, D]; expand GQA kv heads to the query-head count for the dense kernel.
+            qs = q[s:e].transpose(0, 1).unsqueeze(0)
+            ks = (k[s:e].repeat_interleave(rep, dim=1) if rep > 1 else k[s:e]).transpose(0, 1).unsqueeze(0)
+            vs = (v[s:e].repeat_interleave(rep, dim=1) if rep > 1 else v[s:e]).transpose(0, 1).unsqueeze(0)
+            gos = grad_out[s:e].transpose(0, 1).unsqueeze(0)
+            # is_causal=True, attn_bias=None -> no O(L^2) mask; compute_log_sumexp=True for the bwd.
+            out, lse, philox_seed, philox_offset = eff_attn(qs, ks, vs, None, True, 0.0, True, scale=scale)
+            dqi, dki, dvi, _ = eff_attn_bwd(
+                gos, qs, ks, vs, None, out, lse, philox_seed, philox_offset,
+                0.0, [True, True, True, False], True, scale=scale,
+            )
+            dq[s:e] = dqi.squeeze(0).transpose(0, 1)
+            dki = dki.squeeze(0).transpose(0, 1)  # [L, H, D] (expanded query heads)
+            dvi = dvi.squeeze(0).transpose(0, 1)
+            if rep > 1:  # reduce expanded query heads back to the kv-head count (GQA)
+                dk[s:e] = dki.reshape(seg_len, HKV, rep, head_dim).sum(dim=2)
+                dv[s:e] = dvi.reshape(seg_len, HKV, rep, head_dim).sum(dim=2)
+            else:
+                dk[s:e] = dki
+                dv[s:e] = dvi
+        return dq, dk, dv, None, None, None, None, None, None
+
+
 class SGLangFlashAttention(MegatronModule):
     """SGLang-compatible FA3 attention path with packed-sequence support."""
 
@@ -87,6 +178,13 @@ def __init__(
         self.attention_dropout = (
             config.attention_dropout if attention_dropout is None else attention_dropout
         )
+        # Attention backend is an explicit flag (TOP_ATTN_BACKEND), NOT a GPU arch-gate:
+        # miles' TrueOnPolicyKernelPolicy sets it to match sglang's --sglang-attention-backend
+        # so both engines agree. "fa3" (Hopper, faster) or "flashinfer" (Blackwell / no-FA3,
+        # and usable on Hopper for e2e parity tests). flashinfer does GQA natively (no KV repeat)
+        # + takes fixed_split_size=4096 to match sglang's deterministic ragged prefill bitwise.
+        # Fallback "fa3" applies only when run standalone without the flag set.
+        self.attention_backend = os.environ.get("TOP_ATTN_BACKEND", "fa3")
 
     def forward(
         self,
@@ -104,13 +202,15 @@ def forward(
         assert (
             attn_mask_type is None or attn_mask_type == AttnMaskType.causal
         ), "Only causal attention is supported for SGLangFlashAttention"
-        if not HAVE_FA3_VARLEN or fa3_varlen_func is None:
+        if self.attention_backend == "fa3" and (not HAVE_FA3_VARLEN or fa3_varlen_func is None):
             raise ImportError("Flash Attention 3 varlen is required for SGLangFlashAttention")
 
         is_packed = packed_seq_params is not None
         input_ndim = query.dim()
         head_dim_idx = -2 if is_packed and input_ndim >= 3 else 2
-        if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:
+        if self.attention_backend != "flashinfer" and (
+            self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1
+        ):
             repeat_factor = (
                 self.num_attention_heads_per_partition // self.num_query_groups_per_partition
             )
@@ -166,39 +266,54 @@ def forward(
             key = self.cp_layout.sequence_to_head_parallel(key, cu_seqlens_k)
             value = self.cp_layout.sequence_to_head_parallel(value, cu_seqlens_k)
 
-        sig = inspect.signature(fa3_varlen_func)
-        fa3_kwargs = {
-            "q": query,
-            "k": key,
-            "v": value,
-            "cu_seqlens_q": cu_seqlens_q,
-            "cu_seqlens_k": cu_seqlens_k,
-            "max_seqlen_q": max_seqlen_q,
-            "max_seqlen_k": max_seqlen_k,
-            "softmax_scale": self.softmax_scale,
-            "causal": True,
-        }
-        if "dropout_p" in sig.parameters:
-            fa3_kwargs["dropout_p"] = self.attention_dropout if self.training else 0.0
-        if "window_size" in sig.parameters:
-            fa3_kwargs["window_size"] = (-1, -1)
-        if "softcap" in sig.parameters:
-            fa3_kwargs["softcap"] = 0.0
-        if "return_attn_probs" in sig.parameters:
-            fa3_kwargs["return_attn_probs"] = False
-        if "return_softmax_lse" in sig.parameters:
-            fa3_kwargs["return_softmax_lse"] = False
-        if "num_splits" in sig.parameters:
-            fa3_kwargs["num_splits"] = 1
-        if (
-            "deterministic" in sig.parameters
-            and os.environ.get("MEGATRON_TRUE_ON_POLICY_FA3_DETERMINISTIC_BWD") == "1"
-        ):
-            fa3_kwargs["deterministic"] = True
-
-        output = fa3_varlen_func(**fa3_kwargs)
-        if isinstance(output, tuple):
-            output = output[0]
+        if self.attention_backend == "flashinfer":
+            # Deterministic ragged prefill matching sglang's flashinfer TOP path (GQA-native,
+            # fixed_split_size=4096); autograd.Function supplies the backward (flashinfer is fwd-only).
+            output = _FlashinferRaggedAttn.apply(
+                query,
+                key,
+                value,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                self.num_attention_heads_per_partition,
+                self.num_query_groups_per_partition,
+                self.hidden_size_per_attention_head,
+                self.softmax_scale,
+            )
+        else:
+            sig = inspect.signature(fa3_varlen_func)
+            fa3_kwargs = {
+                "q": query,
+                "k": key,
+                "v": value,
+                "cu_seqlens_q": cu_seqlens_q,
+                "cu_seqlens_k": cu_seqlens_k,
+                "max_seqlen_q": max_seqlen_q,
+                "max_seqlen_k": max_seqlen_k,
+                "softmax_scale": self.softmax_scale,
+                "causal": True,
+            }
+            if "dropout_p" in sig.parameters:
+                fa3_kwargs["dropout_p"] = self.attention_dropout if self.training else 0.0
+            if "window_size" in sig.parameters:
+                fa3_kwargs["window_size"] = (-1, -1)
+            if "softcap" in sig.parameters:
+                fa3_kwargs["softcap"] = 0.0
+            if "return_attn_probs" in sig.parameters:
+                fa3_kwargs["return_attn_probs"] = False
+            if "return_softmax_lse" in sig.parameters:
+                fa3_kwargs["return_softmax_lse"] = False
+            if "num_splits" in sig.parameters:
+                fa3_kwargs["num_splits"] = 1
+            if (
+                "deterministic" in sig.parameters
+                and os.environ.get("MEGATRON_TRUE_ON_POLICY_FA3_DETERMINISTIC_BWD") == "1"
+            ):
+                fa3_kwargs["deterministic"] = True
+
+            output = fa3_varlen_func(**fa3_kwargs)
+            if isinstance(output, tuple):
+                output = output[0]
 
         if self.cp_size > 1:
             assert self.cp_layout is not None
diff --git a/miles_megatron_plugins/true_on_policy/sglang_backend.py b/miles_megatron_plugins/true_on_policy/sglang_backend.py
index 1c621f3ac3d..1e296122991 100644
--- a/miles_megatron_plugins/true_on_policy/sglang_backend.py
+++ b/miles_megatron_plugins/true_on_policy/sglang_backend.py
@@ -1,6 +1,6 @@
 """Compatibility facade for the Megatron true-on-policy SGLang backend."""
 
-from .attention_fa3 import (
+from .sglang_attention import (
     HAVE_FA3_VARLEN,
     SGLangCoreAttention,
     SGLangFlashAttention,
@@ -19,13 +19,6 @@
 from .linear import SGLangColumnParallelLinear, SGLangRowParallelLinear
 from .norm import SGLangFinalRMSNorm, SGLangNorm, SGLangQKRMSNorm
 from .provider import SGLangSpecProvider
-from .rope import (
-    disable_sglang_rope,
-    enable_sglang_rope,
-    is_sglang_rope_enabled,
-    sglang_apply_rotary_pos_emb,
-    sglang_apply_rotary_pos_emb_with_freqs,
-)
 from .runtime import (
     enable_sglang_batch_invariant_mode,
     ensure_batch_invariant_mode_from_config,
@@ -47,13 +40,8 @@
     "SGLangSpecProvider",
     "SGLangUlyssesCPLayout",
     "_sglang_bias_dropout_add",
-    "disable_sglang_rope",
     "enable_sglang_batch_invariant_mode",
-    "enable_sglang_rope",
     "fa3_varlen_func",
     "get_sglang_bias_dropout_add",
-    "is_sglang_rope_enabled",
     "resolve_true_on_policy_runtime_policy",
-    "sglang_apply_rotary_pos_emb",
-    "sglang_apply_rotary_pos_emb_with_freqs",
 ]
diff --git a/tests/unit_tests/extension/test_sglang_extension.py b/tests/unit_tests/extension/test_sglang_extension.py
index 6fe06c787d1..6486223a9c0 100644
--- a/tests/unit_tests/extension/test_sglang_extension.py
+++ b/tests/unit_tests/extension/test_sglang_extension.py
@@ -21,9 +21,7 @@
     SGLangRowParallelLinear,
     SGLangSpecProvider,
     _ensure_batch_invariant_mode_from_config,
-    disable_sglang_rope,
     get_sglang_bias_dropout_add,
-    is_sglang_rope_enabled,
     resolve_true_on_policy_runtime_policy,
 )
 from miles_megatron_plugins.true_on_policy.contracts import get_true_on_policy_contract
@@ -107,21 +105,20 @@ def test_legacy_sglang_backend_imports_match_true_on_policy_namespace():
     from megatron.core.extensions import sglang as legacy_backend
     from megatron.core.tensor_parallel import matmul_tp_inv as legacy_matmul
     from miles_megatron_plugins.true_on_policy import (
-        attention_fa3,
         bias_dropout,
         cp_layout,
         linear,
         norm,
         provider,
-        rope,
         runtime,
+        sglang_attention,
     )
     from miles_megatron_plugins.true_on_policy import matmul, sglang_backend
 
     assert legacy_backend.SGLangNorm is sglang_backend.SGLangNorm
     assert legacy_backend.SGLangRowParallelLinear is sglang_backend.SGLangRowParallelLinear
     assert sglang_backend.SGLangColumnParallelLinear is linear.SGLangColumnParallelLinear
-    assert sglang_backend.SGLangCoreAttention is attention_fa3.SGLangCoreAttention
+    assert sglang_backend.SGLangCoreAttention is sglang_attention.SGLangCoreAttention
     assert sglang_backend.SGLangUlyssesCPLayout is cp_layout.SGLangUlyssesCPLayout
     assert sglang_backend.SGLangNorm is norm.SGLangNorm
     assert sglang_backend.SGLangSpecProvider is provider.SGLangSpecProvider
@@ -130,7 +127,6 @@ def test_legacy_sglang_backend_imports_match_true_on_policy_namespace():
         sglang_backend.enable_sglang_batch_invariant_mode
         is runtime.enable_sglang_batch_invariant_mode
     )
-    assert sglang_backend.sglang_apply_rotary_pos_emb is rope.sglang_apply_rotary_pos_emb
     assert (
         sglang_backend.resolve_true_on_policy_runtime_policy
         is resolve_true_on_policy_runtime_policy
@@ -166,7 +162,7 @@ def test_true_on_policy_contract_resolves_megatron_runtime_policy():
     assert policy.batch_invariant_mode
     assert policy.attention_backend == "fa3_varlen"
     assert policy.cp_layout == "ulysses_a2a"
-    assert policy.use_ulysses_cp_recompute_fallback
+    assert policy.use_non_reentrant_recompute
 
 
 def test_contract_object_owns_megatron_runtime_policy_values():
@@ -196,7 +192,9 @@ def test_contract_object_owns_megatron_runtime_policy_values():
     assert policy.apply_logits_contract
     assert policy.use_sglang_final_norm
     assert policy.use_sglang_residual_pair
-    assert not policy.use_ulysses_cp_recompute_fallback
+    # All TOP runs use non-reentrant recompute (reentrant CheckpointFunction NaNs
+    # the TOP kernels), independent of context parallelism.
+    assert policy.use_non_reentrant_recompute
 
 
 def test_qwen3_dense_contract_only_marks_ulysses_a2a_as_cp_layout():
@@ -211,7 +209,8 @@ def test_qwen3_dense_contract_only_marks_ulysses_a2a_as_cp_layout():
     )
 
     assert policy.cp_layout is None
-    assert not policy.use_ulysses_cp_recompute_fallback
+    # Non-reentrant recompute applies to every TOP run, not only Ulysses-a2a CP.
+    assert policy.use_non_reentrant_recompute
 
 
 def test_missing_true_on_policy_contract_uses_default_policy():
@@ -241,14 +240,12 @@ def test_default_backend_selection_is_unchanged():
 
 
 def test_true_on_policy_contract_selects_sglang_backend():
-    disable_sglang_rope()
     config = _make_config(true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1, qk_layernorm=True)
 
     layer_spec = get_gpt_decoder_layer_specs(
         config, use_transformer_engine=False, normalization=config.normalization
     )[0]
 
-    assert is_sglang_rope_enabled()
     assert isinstance(layer_spec.submodules.input_layernorm, ModuleSpec)
     assert layer_spec.submodules.input_layernorm.module is SGLangNorm
     assert layer_spec.submodules.input_layernorm.params["override_orig_dtype"] is torch.float32
@@ -394,7 +391,7 @@ def test_sglang_reference_matmul_matches_sglang_mixed_dtype_contract():
     torch.testing.assert_close(actual, expected)
 
 
-def test_sglang_row_parallel_matmul_uses_fixed_k_block_order():
+def test_sglang_row_parallel_matmul_routes_to_tp_invariant_kernel():
     input_ = torch.randn(2, 3, 256)
     weight = torch.randn(5, 256)
     bias = torch.randn(5)

From d462714388f85dab1829d970021de370ab538440 Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Sat, 18 Jul 2026 03:18:27 -0500
Subject: [PATCH 4/9] refactor(top): rename flashinfer ragged wrapper +
 configurable workspace

---
 .../true_on_policy/sglang_attention.py        | 29 ++++++++++++-------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/miles_megatron_plugins/true_on_policy/sglang_attention.py b/miles_megatron_plugins/true_on_policy/sglang_attention.py
index f888d283520..1f70869bca1 100644
--- a/miles_megatron_plugins/true_on_policy/sglang_attention.py
+++ b/miles_megatron_plugins/true_on_policy/sglang_attention.py
@@ -30,21 +30,28 @@
         fa3_varlen_func = None
 
 
-_FI_RAGGED_WRAPPER = None
+_ragged_prefill_wrapper = None
 
+# Scratch for flashinfer's deterministic ragged prefill (batch_prefill_tmp_v etc.). Sized to
+# the TRAINING side's parallelism/window, independent of sglang's rollout-side
+# SGLANG_FLASHINFER_WORKSPACE_SIZE: the requirement scales with per-rank heads
+# (total_heads / TP) and batching, which differ between train and rollout, and the size is
+# scratch that does not affect parity. 256 MiB overflowed the canonical config (8192 window,
+# batch 256, ~406 MiB), so default 1 GiB.
+_WORKSPACE_SIZE = int(
+    os.environ.get("MEGATRON_TRUE_ON_POLICY_FLASHINFER_WORKSPACE_SIZE", 1024 * 1024 * 1024)
+)
 
-def _get_flashinfer_ragged_wrapper(device):
+
+def _get_ragged_prefill_wrapper(device):
     """One shared deterministic ragged-prefill wrapper (Blackwell / flashinfer)."""
-    global _FI_RAGGED_WRAPPER
-    if _FI_RAGGED_WRAPPER is None:
+    global _ragged_prefill_wrapper
+    if _ragged_prefill_wrapper is None:
         from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
 
-        # 1 GiB workspace: the deterministic ragged prefill needs batch_prefill_tmp_v
-        # scratch that scales with heads x head_dim x fixed_split tiles; 256 MiB overflows
-        # at the canonical config (8192 response window, batch 256) -> needs ~406 MiB.
-        ws = torch.empty(1024 * 1024 * 1024, dtype=torch.uint8, device=device)
-        _FI_RAGGED_WRAPPER = BatchPrefillWithRaggedKVCacheWrapper(ws, kv_layout="NHD")
-    return _FI_RAGGED_WRAPPER
+        ws = torch.empty(_WORKSPACE_SIZE, dtype=torch.uint8, device=device)
+        _ragged_prefill_wrapper = BatchPrefillWithRaggedKVCacheWrapper(ws, kv_layout="NHD")
+    return _ragged_prefill_wrapper
 
 
 class _FlashinferRaggedAttn(torch.autograd.Function):
@@ -55,7 +62,7 @@ class _FlashinferRaggedAttn(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, q, k, v, cu_q, cu_k, num_q_heads, num_kv_heads, head_dim, scale):
-        w = _get_flashinfer_ragged_wrapper(q.device)
+        w = _get_ragged_prefill_wrapper(q.device)
         # Match sglang's call exactly: plan WITHOUT sm_scale, apply sm_scale at forward-time
         # via .forward (sglang uses fast_prefill_plan (no sm_scale) + .forward(..., sm_scale=...)).
         w.plan(

From 52193c0b4a8906eaabded49389d7d8ccc8ce2348 Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Sat, 18 Jul 2026 03:18:27 -0500
Subject: [PATCH 5/9] test(top): guard Megatron RoPE == SGLang RoPE bitwise

---
 .../extension/test_sglang_extension.py        | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/tests/unit_tests/extension/test_sglang_extension.py b/tests/unit_tests/extension/test_sglang_extension.py
index 6486223a9c0..c238d7c4cff 100644
--- a/tests/unit_tests/extension/test_sglang_extension.py
+++ b/tests/unit_tests/extension/test_sglang_extension.py
@@ -349,6 +349,50 @@ def test_ulysses_rope_keeps_unsplit_positions_for_full_sequence_layout():
     torch.testing.assert_close(actual, expected)
 
 
+def _sglang_reference_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    """SGLang's forward_native neox rotary — the formula the now-removed custom sglang-rope
+    path forced Megatron to use. ``freqs`` is Megatron's [seq, 1, 1, head_dim] cache of
+    duplicated half-angles; SGLang takes the first half as the angles. This is the reference the
+    standard Megatron ``apply_rotary_pos_emb`` must match bitwise for the custom-path removal to
+    be policy-preserving."""
+    head_dim = x.shape[-1]
+    angles = freqs[..., : head_dim // 2]
+    cos = torch.cos(angles).float()
+    sin = torch.sin(angles).float()
+    orig_dtype = x.dtype
+    x1, x2 = torch.chunk(x.float(), 2, dim=-1)
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    return torch.cat((o1, o2), dim=-1).to(orig_dtype)
+
+
+def test_standard_megatron_rope_matches_sglang_rope_bitwise():
+    # The dense contract drops the custom sglang-rope path and relies on Megatron's own
+    # apply_rotary_pos_emb (unfused, non-interleaved, mscale=1) already being BITWISE-identical
+    # to SGLang's rotary. Guard that equivalence so the removal can't silently change the policy.
+    config = _make_config(
+        true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1,
+        apply_rope_fusion=False,
+        rotary_interleaved=False,
+    )
+    seq, batch, heads, head_dim = 6, 2, 4, 8
+    cp_group = _FakeCPGroup(size=1, rank=0)  # bshd path ignores it; avoids parallel_state init
+    for dtype in (torch.float32, torch.bfloat16):
+        torch.manual_seed(0)
+        half = torch.randn(seq, 1, 1, head_dim // 2)
+        freqs = torch.cat((half, half), dim=-1)  # Megatron duplicates half-angles across head_dim
+        x = torch.randn(seq, batch, heads, head_dim, dtype=dtype)
+
+        megatron_out = apply_rotary_pos_emb(x, freqs, config=config, cp_group=cp_group)
+        sglang_out = _sglang_reference_rope(x, freqs)
+
+        assert torch.equal(megatron_out, sglang_out), (
+            f"Megatron RoPE != SGLang RoPE at dtype={dtype}: dropping the custom sglang-rope "
+            f"path would change the policy (max abs diff "
+            f"{(megatron_out.float() - sglang_out.float()).abs().max().item()})"
+        )
+
+
 def test_local_attention_still_rejects_ulysses_context_parallel_without_true_on_policy():
     with pytest.raises(ValueError, match="only supports all_gather"):
         _make_config(tensor_model_parallel_size=1, context_parallel_size=2, cp_comm_type="a2a")

From d332f60011528f4f082c326213ca41648e8f9887 Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Sat, 18 Jul 2026 14:43:17 -0500
Subject: [PATCH 6/9] test(top): grad tests for hand-written RMSNorm +
 attention backwards

---
 .../unit_tests/extension/test_top_backward.py | 104 ++++++++++++++++++
 1 file changed, 104 insertions(+)
 create mode 100644 tests/unit_tests/extension/test_top_backward.py

diff --git a/tests/unit_tests/extension/test_top_backward.py b/tests/unit_tests/extension/test_top_backward.py
new file mode 100644
index 00000000000..46fe1616912
--- /dev/null
+++ b/tests/unit_tests/extension/test_top_backward.py
@@ -0,0 +1,104 @@
+"""Gradient tests for the hand-written true-on-policy backward Functions.
+
+The analytic backwards — the RMSNorm vjp (``kernels._RmsNormBatchInvariant``) and the O(L)
+per-segment attention vjp (``sglang_attention._FlashinferRaggedAttn``) — have no inference
+counterpart, so they are hand-written and need direct validation. Their forwards are
+CUDA-kernel-bound (SGLang batch-invariant Triton; flashinfer, Blackwell), and those kernels
+run in fp32/bf16, not fp64 — so a finite-difference ``torch.autograd.gradcheck`` on the real
+forward isn't possible. Instead each test invokes the real Function and compares its gradient
+to a trusted, differentiable torch reference (whose own autograd is gradcheck-validated by
+PyTorch). CUDA + the kernels are required, so these skip on CPU CI and run on GPU.
+"""
+
+import pytest
+import torch
+
+CUDA = torch.cuda.is_available()
+
+
+def _rms_norm_ref(x: torch.Tensor, eps: float) -> torch.Tensor:
+    # Pure-torch differentiable RMSNorm with implicit unit weight (same math as the kernel).
+    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+
+
+@pytest.mark.skipif(not CUDA, reason="TOP batch-invariant RMSNorm kernel is CUDA-only")
+def test_rms_norm_batch_invariant_backward_matches_reference():
+    pytest.importorskip("sglang.srt.batch_invariant_ops")
+    from sglang.srt.batch_invariant_ops import rms_norm_batch_invariant as bare_kernel
+
+    from miles_megatron_plugins.true_on_policy import kernels
+
+    H, N, eps = 128, 8, 1e-6
+    x = torch.randn(N, H, device="cuda", dtype=torch.float32)
+
+    # Forward must stay bitwise-identical to the bare kernel (the wrapper preserves TOP parity).
+    ones = torch.ones(H, device="cuda", dtype=torch.float32)
+    with torch.no_grad():
+        assert torch.equal(kernels.rms_norm_batch_invariant(x, eps), bare_kernel(x, ones, eps))
+
+    # Backward: our analytic vjp reproduces the fp32 torch RMSNorm reference's autograd.
+    xa = x.clone().requires_grad_(True)
+    xb = x.clone().requires_grad_(True)
+    go = torch.randn_like(x)
+    kernels.rms_norm_batch_invariant(xa, eps).backward(go)
+    _rms_norm_ref(xb, eps).backward(go)
+    rel = ((xa.grad - xb.grad).norm() / xb.grad.norm().clamp_min(1e-12)).item()
+    assert torch.isfinite(xa.grad).all()
+    assert rel < 1e-3, f"RMSNorm analytic backward vs reference rel err {rel:.2e}"
+
+
+@pytest.mark.skipif(not CUDA, reason="flashinfer ragged prefill + CUDA required")
+def test_flashinfer_ragged_attn_backward_matches_sdpa():
+    pytest.importorskip("flashinfer")
+    import torch.nn.functional as F
+
+    from miles_megatron_plugins.true_on_policy.sglang_attention import _FlashinferRaggedAttn
+
+    H, HKV, D = 32, 8, 128
+    scale = 1.0 / (D ** 0.5)
+    rep = H // HKV
+    # Packed multi-segment input including length-1 segments (the real rollout edge case).
+    cu = torch.tensor([0, 137, 138, 400, 801, 1300, 1301, 2048], dtype=torch.int32, device="cuda")
+    T = int(cu[-1])
+
+    torch.manual_seed(0)
+    q = torch.randn(T, H, D, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(T, HKV, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(T, HKV, D, device="cuda", dtype=torch.bfloat16)
+    go = torch.randn(T, H, D, device="cuda", dtype=torch.bfloat16)
+
+    # Block-diagonal causal reference via SDPA (GQA heads expanded), differentiable.
+    seg = torch.zeros(T, dtype=torch.long, device="cuda")
+    cl = cu.long()
+    for i in range(cl.numel() - 1):
+        seg[cl[i]:cl[i + 1]] = i
+    idx = torch.arange(T, device="cuda")
+    mask = (seg[:, None] == seg[None, :]) & (idx[:, None] >= idx[None, :])
+
+    def _bhtd(t):  # [T, heads, D] -> [1, heads, T, D]
+        return t.transpose(0, 1).unsqueeze(0)
+
+    def _expand(t):  # GQA: expand kv heads to the query-head count
+        return t.repeat_interleave(rep, 1) if rep > 1 else t
+
+    qr, kr, vr = (t.clone().requires_grad_(True) for t in (q, k, v))
+    o_ref = (
+        F.scaled_dot_product_attention(
+            _bhtd(qr), _bhtd(_expand(kr)), _bhtd(_expand(vr)), attn_mask=mask[None, None], scale=scale
+        )
+        .squeeze(0)
+        .transpose(0, 1)
+    )
+    gq_ref, gk_ref, gv_ref = torch.autograd.grad(o_ref, [qr, kr, vr], go)
+
+    qp, kp, vp = (t.clone().requires_grad_(True) for t in (q, k, v))
+    out = _FlashinferRaggedAttn.apply(qp, kp, vp, cu, cu, H, HKV, D, scale)
+    out.backward(go)
+
+    def _rel(a, b):
+        return ((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12)).item()
+
+    assert _rel(out, o_ref) < 5e-3, "flashinfer forward diverged from the SDPA reference"
+    for name, gp, gr in (("dq", qp.grad, gq_ref), ("dk", kp.grad, gk_ref), ("dv", vp.grad, gv_ref)):
+        assert torch.isfinite(gp).all(), f"attention backward {name} has non-finite grads"
+        assert _rel(gp, gr) < 5e-3, f"attention backward {name} vs SDPA reference rel err too large"

From de26aa17e513287358ec86b70d0d92661064373e Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Sat, 18 Jul 2026 20:48:05 -0500
Subject: [PATCH 7/9] fix(top): pin flashinfer prefill backend to sglang's
 (contractual parity)

---
 .../true_on_policy/sglang_attention.py             | 13 ++++++++++++-
 tests/unit_tests/extension/test_top_backward.py    | 14 ++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/miles_megatron_plugins/true_on_policy/sglang_attention.py b/miles_megatron_plugins/true_on_policy/sglang_attention.py
index 1f70869bca1..bf849c031cc 100644
--- a/miles_megatron_plugins/true_on_policy/sglang_attention.py
+++ b/miles_megatron_plugins/true_on_policy/sglang_attention.py
@@ -43,6 +43,15 @@
 )
 
 
+def _fmha_backend(device) -> str:
+    """The flashinfer prefill backend, matching sglang's rollout-side choice EXACTLY so training
+    and rollout run the IDENTICAL kernel (contractual parity, not coincidental). sglang's
+    flashinfer_backend.py uses "cutlass" on SM100 (Blackwell), "auto" otherwise. Omitting the
+    backend only accidentally matched sglang on Blackwell (both resolved to cutlass) and diverged
+    on Hopper (train "auto" default != rollout "auto" kernel -> abs_diff 0.017)."""
+    return "cutlass" if torch.cuda.get_device_capability(device)[0] >= 10 else "auto"
+
+
 def _get_ragged_prefill_wrapper(device):
     """One shared deterministic ragged-prefill wrapper (Blackwell / flashinfer)."""
     global _ragged_prefill_wrapper
@@ -50,7 +59,9 @@ def _get_ragged_prefill_wrapper(device):
         from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
 
         ws = torch.empty(_WORKSPACE_SIZE, dtype=torch.uint8, device=device)
-        _ragged_prefill_wrapper = BatchPrefillWithRaggedKVCacheWrapper(ws, kv_layout="NHD")
+        _ragged_prefill_wrapper = BatchPrefillWithRaggedKVCacheWrapper(
+            ws, kv_layout="NHD", backend=_fmha_backend(device)
+        )
     return _ragged_prefill_wrapper
 
 
diff --git a/tests/unit_tests/extension/test_top_backward.py b/tests/unit_tests/extension/test_top_backward.py
index 46fe1616912..b483cc31de2 100644
--- a/tests/unit_tests/extension/test_top_backward.py
+++ b/tests/unit_tests/extension/test_top_backward.py
@@ -102,3 +102,17 @@ def _rel(a, b):
     for name, gp, gr in (("dq", qp.grad, gq_ref), ("dk", kp.grad, gk_ref), ("dv", vp.grad, gv_ref)):
         assert torch.isfinite(gp).all(), f"attention backward {name} has non-finite grads"
         assert _rel(gp, gr) < 5e-3, f"attention backward {name} vs SDPA reference rel err too large"
+
+
+def test_flashinfer_prefill_backend_matches_sglang(monkeypatch):
+    # CPU guard (no GPU): the training-side flashinfer prefill backend MUST match sglang's
+    # rollout-side choice (flashinfer_backend.py: cutlass on SM100/Blackwell, auto otherwise).
+    # If it drifts, train and rollout run different flashinfer kernels and parity breaks SILENTLY
+    # -- exactly what happened when the backend was left unset: Blackwell coincidentally matched,
+    # Hopper diverged (abs_diff 0.017). This locks the two together by rule.
+    from miles_megatron_plugins.true_on_policy import sglang_attention
+
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 0))
+    assert sglang_attention._fmha_backend("cuda") == "cutlass"  # SM100 / Blackwell
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (9, 0))
+    assert sglang_attention._fmha_backend("cuda") == "auto"  # Hopper and earlier

From a9b2342c70faf5e92e95df27fe75e110e26d1ec4 Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Sat, 18 Jul 2026 21:00:43 -0500
Subject: [PATCH 8/9] test(top): grad test for tp-invariant row-linear backward

---
 .../unit_tests/extension/test_top_backward.py | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tests/unit_tests/extension/test_top_backward.py b/tests/unit_tests/extension/test_top_backward.py
index b483cc31de2..baf4c2e97ac 100644
--- a/tests/unit_tests/extension/test_top_backward.py
+++ b/tests/unit_tests/extension/test_top_backward.py
@@ -47,6 +47,37 @@ def test_rms_norm_batch_invariant_backward_matches_reference():
     assert rel < 1e-3, f"RMSNorm analytic backward vs reference rel err {rel:.2e}"
 
 
+@pytest.mark.skipif(not CUDA, reason="TOP tp-invariant row-linear kernel is CUDA-only")
+def test_tp_invariant_row_linear_backward_matches_reference():
+    pytest.importorskip("sglang.srt.tp_invariant_ops")
+    from miles_megatron_plugins.true_on_policy import kernels
+
+    tokens, K, out = 64, 256, 128
+    torch.manual_seed(0)
+    x = torch.randn(tokens, K, device="cuda", dtype=torch.bfloat16)
+    w = torch.randn(out, K, device="cuda", dtype=torch.bfloat16)  # [out, K_local], RowParallelLinear layout
+    go = torch.randn(tokens, out, device="cuda", dtype=torch.bfloat16)
+
+    # Forward delegates to SGLang's matmul_tp_inv (deterministic TP-invariant reduction order),
+    # so it matches a plain x @ Wᵀ reference within kernel tolerance, not bitwise. Backward is the
+    # hand-written linear vjp (grad_x = go @ W, grad_w = goᵀ @ x) and must reproduce the reference's
+    # autograd — same pattern as the RMSNorm test: real Function vs differentiable torch reference.
+    xa, wa = (t.clone().requires_grad_(True) for t in (x, w))
+    xb, wb = (t.clone().requires_grad_(True) for t in (x, w))
+    out_k = kernels.tp_invariant_row_linear(xa, wa)
+    out_ref = xb @ wb.t()
+    out_k.backward(go)
+    out_ref.backward(go)
+
+    def _rel(a, b):
+        return ((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12)).item()
+
+    assert _rel(out_k, out_ref) < 5e-3, "tp-invariant row-linear forward diverged from reference"
+    for name, gk, gr in (("dx", xa.grad, xb.grad), ("dw", wa.grad, wb.grad)):
+        assert torch.isfinite(gk).all(), f"row-linear backward {name} has non-finite grads"
+        assert _rel(gk, gr) < 5e-3, f"row-linear backward {name} vs reference rel err too large"
+
+
 @pytest.mark.skipif(not CUDA, reason="flashinfer ragged prefill + CUDA required")
 def test_flashinfer_ragged_attn_backward_matches_sdpa():
     pytest.importorskip("flashinfer")

From aeb932ffc2de61ed797ce66fab667326f8a72757 Mon Sep 17 00:00:00 2001
From: Yiming Li <li-ym98@qq.com>
Date: Mon, 20 Jul 2026 06:53:12 -0500
Subject: [PATCH 9/9] fix(top): flashinfer prefill backend fa2; CP-correct head
 counts + output-parity test

---
 .../true_on_policy/sglang_attention.py        | 19 ++++---
 .../unit_tests/extension/test_top_backward.py | 57 ++++++++++++++-----
 2 files changed, 54 insertions(+), 22 deletions(-)

diff --git a/miles_megatron_plugins/true_on_policy/sglang_attention.py b/miles_megatron_plugins/true_on_policy/sglang_attention.py
index bf849c031cc..09f6042dcc0 100644
--- a/miles_megatron_plugins/true_on_policy/sglang_attention.py
+++ b/miles_megatron_plugins/true_on_policy/sglang_attention.py
@@ -44,12 +44,10 @@
 
 
 def _fmha_backend(device) -> str:
-    """The flashinfer prefill backend, matching sglang's rollout-side choice EXACTLY so training
-    and rollout run the IDENTICAL kernel (contractual parity, not coincidental). sglang's
-    flashinfer_backend.py uses "cutlass" on SM100 (Blackwell), "auto" otherwise. Omitting the
-    backend only accidentally matched sglang on Blackwell (both resolved to cutlass) and diverged
-    on Hopper (train "auto" default != rollout "auto" kernel -> abs_diff 0.017)."""
-    return "cutlass" if torch.cuda.get_device_capability(device)[0] >= 10 else "auto"
+    """Flashinfer prefill backend for TOP: fa2, matching sglang's deterministic prefill
+    (its paged wrapper uses fa2). Wrapper choice (ragged vs paged) is irrelevant — bitwise-equal
+    under the same backend."""
+    return "fa2"
 
 
 def _get_ragged_prefill_wrapper(device):
@@ -293,8 +291,13 @@ def forward(
                 value,
                 cu_seqlens_q,
                 cu_seqlens_k,
-                self.num_attention_heads_per_partition,
-                self.num_query_groups_per_partition,
+                # Head counts must match the ACTUAL q/k tensors fed to flashinfer's plan().
+                # Read them from the tensors (as fa3 does) rather than the tp-local partition
+                # counts: under Ulysses CP the a2a above leaves heads/cp per rank, so the
+                # partition counts are cp x too large (crashes flashinfer's reshape). At cp=1
+                # query.shape[-2] == num_attention_heads_per_partition, so this is a no-op there.
+                query.shape[-2],
+                key.shape[-2],
                 self.hidden_size_per_attention_head,
                 self.softmax_scale,
             )
diff --git a/tests/unit_tests/extension/test_top_backward.py b/tests/unit_tests/extension/test_top_backward.py
index baf4c2e97ac..8c5f28b7ca5 100644
--- a/tests/unit_tests/extension/test_top_backward.py
+++ b/tests/unit_tests/extension/test_top_backward.py
@@ -58,10 +58,8 @@ def test_tp_invariant_row_linear_backward_matches_reference():
     w = torch.randn(out, K, device="cuda", dtype=torch.bfloat16)  # [out, K_local], RowParallelLinear layout
     go = torch.randn(tokens, out, device="cuda", dtype=torch.bfloat16)
 
-    # Forward delegates to SGLang's matmul_tp_inv (deterministic TP-invariant reduction order),
-    # so it matches a plain x @ Wᵀ reference within kernel tolerance, not bitwise. Backward is the
-    # hand-written linear vjp (grad_x = go @ W, grad_w = goᵀ @ x) and must reproduce the reference's
-    # autograd — same pattern as the RMSNorm test: real Function vs differentiable torch reference.
+    # Forward = SGLang's matmul_tp_inv (matches x @ Wᵀ within tolerance); backward = the hand-written
+    # linear vjp (grad_x = go @ W, grad_w = goᵀ @ x), checked against the reference's autograd.
     xa, wa = (t.clone().requires_grad_(True) for t in (x, w))
     xb, wb = (t.clone().requires_grad_(True) for t in (x, w))
     out_k = kernels.tp_invariant_row_linear(xa, wa)
@@ -135,15 +133,46 @@ def _rel(a, b):
         assert _rel(gp, gr) < 5e-3, f"attention backward {name} vs SDPA reference rel err too large"
 
 
-def test_flashinfer_prefill_backend_matches_sglang(monkeypatch):
-    # CPU guard (no GPU): the training-side flashinfer prefill backend MUST match sglang's
-    # rollout-side choice (flashinfer_backend.py: cutlass on SM100/Blackwell, auto otherwise).
-    # If it drifts, train and rollout run different flashinfer kernels and parity breaks SILENTLY
-    # -- exactly what happened when the backend was left unset: Blackwell coincidentally matched,
-    # Hopper diverged (abs_diff 0.017). This locks the two together by rule.
+def test_flashinfer_prefill_backend_is_fa2():
+    # Cheap rule guard; the real check is the output-parity test below.
     from miles_megatron_plugins.true_on_policy import sglang_attention
 
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 0))
-    assert sglang_attention._fmha_backend("cuda") == "cutlass"  # SM100 / Blackwell
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (9, 0))
-    assert sglang_attention._fmha_backend("cuda") == "auto"  # Hopper and earlier
+    assert sglang_attention._fmha_backend("cuda") == "fa2"
+
+
+@pytest.mark.skipif(not CUDA, reason="flashinfer required")
+def test_flashinfer_output_matches_deterministic_paged_reference():
+    # Our training flashinfer path (ragged + _fmha_backend) must be bitwise-equal to sglang's
+    # deterministic prefill kernel (paged + fa2) on identical q/k/v. Catches a backend drift.
+    pytest.importorskip("flashinfer")
+    from flashinfer import (
+        BatchPrefillWithRaggedKVCacheWrapper as Ragged,
+        BatchPrefillWithPagedKVCacheWrapper as Paged,
+    )
+    from miles_megatron_plugins.true_on_policy import sglang_attention
+
+    dev, HD, NQ, NKV, L = "cuda", 128, 16, 4, 1024
+    torch.manual_seed(0)
+    q = torch.randn(L, NQ, HD, device=dev, dtype=torch.bfloat16)
+    k = torch.randn(L, NKV, HD, device=dev, dtype=torch.bfloat16)
+    v = torch.randn(L, NKV, HD, device=dev, dtype=torch.bfloat16)
+    cu = torch.tensor([0, L], dtype=torch.int32, device=dev)
+
+    # our training path: ragged wrapper with the backend _fmha_backend selects
+    be = sglang_attention._fmha_backend(dev)
+    wr = Ragged(torch.empty(256 * 1024 * 1024, dtype=torch.uint8, device=dev), kv_layout="NHD", backend=be)
+    wr.plan(cu, cu, NQ, NKV, HD, causal=True, q_data_type=torch.bfloat16,
+            kv_data_type=torch.bfloat16, fixed_split_size=4096)
+    o_ours = wr.run(q, k, v)
+
+    # sglang deterministic reference: paged wrapper + "fa2", page_size=1 over the same K/V
+    wp = Paged(torch.empty(256 * 1024 * 1024, dtype=torch.uint8, device=dev), kv_layout="NHD", backend="fa2")
+    wp.plan(cu, cu, torch.arange(L, dtype=torch.int32, device=dev),
+            torch.tensor([1], dtype=torch.int32, device=dev), NQ, NKV, HD, 1, causal=True,
+            q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16, fixed_split_size=4096)
+    o_ref = wp.run(q, (k.view(L, 1, NKV, HD).contiguous(), v.view(L, 1, NKV, HD).contiguous()))
+
+    assert torch.equal(o_ours, o_ref), (
+        f"training flashinfer backend {be!r} does not bitwise-match sglang deterministic paged+fa2 "
+        f"(max abs diff {(o_ours.float() - o_ref.float()).abs().max().item():.3e})"
+    )