diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge
index 8c46dc42590..70c5dab8db8 160000
--- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge
+++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge
@@ -1 +1 @@
-Subproject commit 8c46dc4259080c510b7455f43e836fdff222c5d3
+Subproject commit 70c5dab8db83e3f00c2bb766f738bcc892de1799
diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml
index 1ac2301025b..f895b3f9aa1 100644
--- a/examples/configs/grpo_math_1B.yaml
+++ b/examples/configs/grpo_math_1B.yaml
@@ -188,6 +188,14 @@ policy:
     # parallelism, sequence packing, or top-k/top-p training-time filtering.
     use_fused_linear_logprobs: false
     fused_linear_logprobs_chunk_size: 256
+    # Opt in to bitwise-identical Megatron generation and policy logprobs.
+    # Requires Megatron generation, BF16, matching training/generation TP, CP=1,
+    # attention_backend=flash, and flash_attention_version=3/4. Sampling supports
+    # the usual temperature, top-k, and top-p settings.
+    batch_invariant_mode: false
+    batch_invariant_backend: "te_native"  # Options: "te_native", "triton", "deepgemm".
+    batch_invariant_collective: "ordered"  # Options: "ordered", "multimem" (NVLS).
+    flash_attention_version: null  # Set to 3 or 4 when enabling the mode.
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false
diff --git a/nemo_rl/algorithms/loss/utils.py b/nemo_rl/algorithms/loss/utils.py
index ca543a9485c..e1bfc58a7e4 100644
--- a/nemo_rl/algorithms/loss/utils.py
+++ b/nemo_rl/algorithms/loss/utils.py
@@ -51,8 +51,10 @@ def prepare_loss_input(
         logits: Logits from the model.
         data: Microbatch data. Will be updated if sampling_params is not None.
         loss_fn: Loss function.
-        vocab_parallel_rank: Vocab parallel rank.
-        vocab_parallel_group: Vocab parallel group.
+        vocab_parallel_rank: Vocab parallel rank. Leave unset with
+            vocab_parallel_group when logits contain the full vocabulary.
+        vocab_parallel_group: Vocab parallel group. Leave unset with
+            vocab_parallel_rank when logits contain the full vocabulary.
         context_parallel_group: Context parallel group.
         sampling_params: Sampling parameters.
         d2t: Draft to target token mapping.
@@ -304,11 +306,12 @@ def prepare_packed_loss_input(
             f"got {loss_fn.input_type}. Use SequencePackingLossWrapper with "
             f"prepare_loss_input for other types."
         )
-    assert vocab_parallel_group is not None, (
-        "prepare_packed_loss_input requires vocab_parallel_group (Megatron TP)."
+    assert (vocab_parallel_group is None) == (vocab_parallel_rank is None), (
+        "vocab_parallel_rank and vocab_parallel_group must either both be provided "
+        "or both be None for full-vocabulary logits."
     )
-    assert vocab_parallel_rank is not None, (
-        "vocab_parallel_rank must be provided with vocab_parallel_group."
+    resolved_vocab_parallel_rank = (
+        0 if vocab_parallel_rank is None else vocab_parallel_rank
     )
 
     input_ids = data["input_ids"]
@@ -345,8 +348,8 @@ def prepare_packed_loss_input(
         packed_rolled_targets,
         cu_seqlens_q_padded,
         unpacked_seqlen,
-        vocab_start_index=vocab_parallel_rank * logits.shape[-1],
-        vocab_end_index=(vocab_parallel_rank + 1) * logits.shape[-1],
+        vocab_start_index=resolved_vocab_parallel_rank * logits.shape[-1],
+        vocab_end_index=(resolved_vocab_parallel_rank + 1) * logits.shape[-1],
         group=vocab_parallel_group,
         inference_only=False,
         cp_group=context_parallel_group,
@@ -371,8 +374,9 @@ def prepare_packed_loss_input(
                     packed_rolled_targets,
                     cu_seqlens_q_padded,
                     unpacked_seqlen,
-                    vocab_start_index=vocab_parallel_rank * logits.shape[-1],
-                    vocab_end_index=(vocab_parallel_rank + 1) * logits.shape[-1],
+                    vocab_start_index=resolved_vocab_parallel_rank * logits.shape[-1],
+                    vocab_end_index=(resolved_vocab_parallel_rank + 1)
+                    * logits.shape[-1],
                     group=vocab_parallel_group,
                     inference_only=False,
                     cp_group=context_parallel_group,
diff --git a/nemo_rl/distributed/model_utils.py b/nemo_rl/distributed/model_utils.py
index 7195592e25f..e98c4a3f1ed 100644
--- a/nemo_rl/distributed/model_utils.py
+++ b/nemo_rl/distributed/model_utils.py
@@ -1049,7 +1049,7 @@ def from_parallel_logits_to_logprobs_packed_sequences(
     unpacked_seqlen: int,
     vocab_start_index: int,
     vocab_end_index: int,
-    group: torch.distributed.ProcessGroup,
+    group: Optional[torch.distributed.ProcessGroup],
     inference_only: bool = False,
     cp_group: Optional[torch.distributed.ProcessGroup] = None,
     chunk_size: Optional[int] = None,
@@ -1070,7 +1070,8 @@ def from_parallel_logits_to_logprobs_packed_sequences(
         unpacked_seqlen (int): The length of the unpacked sequence tensor.
         vocab_start_index (int): Starting vocabulary index for this worker's partition.
         vocab_end_index (int): Ending vocabulary index for this worker's partition.
-        group (torch.distributed.ProcessGroup): Process group for distributed communication.
+        group (torch.distributed.ProcessGroup, optional): Tensor-parallel process
+            group. If None, logits already contain the full vocabulary.
         inference_only (bool, optional): If True, tensors won't be saved for backward pass. Defaults to False.
         cp_group (torch.distributed.ProcessGroup, optional): Context parallelism process group. Defaults to None.
         chunk_size (int, optional): Sequence dimension chunk size for computing the log probabilities.
@@ -1110,8 +1111,23 @@ def from_parallel_logits_to_logprobs_packed_sequences(
         target = rolled_targets.unsqueeze(0)
         vocab_parallel_logits = vocab_parallel_logits.unsqueeze(0)
 
+    # Batch-invariant Megatron scoring gathers the full vocabulary before using
+    # the same canonical FP32 log_softmax as generation. Keeping this branch in
+    # the packed helper lets packed and unpacked callers share that contract.
+    if group is None:
+        logits = vocab_parallel_logits.to(torch.float32)
+        logits, _ = apply_top_k_top_p(
+            logits,
+            top_k=sampling_params.top_k if sampling_params is not None else None,
+            top_p=sampling_params.top_p if sampling_params is not None else 1.0,
+        )
+        probs = (
+            torch.nn.functional.log_softmax(logits, dim=-1)
+            .gather(dim=-1, index=target.unsqueeze(-1))
+            .squeeze(-1)
+        )
     # Apply distributed log probability computation
-    if need_top_k_or_top_p_filtering(sampling_params):
+    elif need_top_k_or_top_p_filtering(sampling_params):
         if chunk_size is not None:
             probs: torch.Tensor = ChunkedDistributedLogprobWithSampling.apply(  # type: ignore
                 vocab_parallel_logits,
diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index fb01e70f190..a3e1e8d91d3 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -291,7 +291,14 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None:
             ),
             logging_step_interval=logging_step_interval,
             num_speculative_tokens=num_speculative_tokens,
-            logprobs_mode="processed_logprobs",
+            # Sampling parameters control token selection, but batch-invariant
+            # generation reports raw model logprobs. Policy scoring mirrors this
+            # contract so parity is independent of temperature/top-k/top-p.
+            logprobs_mode=(
+                "raw_logprobs"
+                if self.cfg["megatron_cfg"].get("batch_invariant_mode")
+                else "processed_logprobs"
+            ),
             max_requests=max_requests,
         )
 
diff --git a/nemo_rl/models/megatron/batch_invariant.py b/nemo_rl/models/megatron/batch_invariant.py
new file mode 100644
index 00000000000..f16bfa36132
--- /dev/null
+++ b/nemo_rl/models/megatron/batch_invariant.py
@@ -0,0 +1,31 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+
+def batch_invariant_token_multiple(configured_multiple: int, tp_size: int) -> int:
+    """Return a token multiple compatible with batch-invariant MCore inference."""
+    if configured_multiple < 1:
+        raise ValueError("configured_multiple must be positive.")
+    if tp_size < 1:
+        raise ValueError("tp_size must be positive.")
+
+    # Import lazily so non-Megatron policy drivers do not import Megatron-Core.
+    # TOKEN_ROUNDER is MCore's single source of truth for eager and graphed
+    # batch-invariant inference token alignment.
+    from megatron.core.inference.batch_dimensions_utils import TOKEN_ROUNDER
+
+    inference_multiple = ((TOKEN_ROUNDER + tp_size - 1) // tp_size) * tp_size
+    return math.lcm(configured_multiple, inference_multiple)
diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py
index c55740b1b7e..b811f0f2f63 100644
--- a/nemo_rl/models/megatron/data.py
+++ b/nemo_rl/models/megatron/data.py
@@ -30,6 +30,7 @@
 from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank
+from nemo_rl.models.megatron.batch_invariant import batch_invariant_token_multiple
 from nemo_rl.models.megatron.common import _round_up_to_multiple
 from nemo_rl.utils.r3_trace import (
     r3_trace_verify_forward_enabled,
@@ -213,6 +214,17 @@ def get_microbatch_iterator(
             cfg["make_sequence_length_divisible_by"],
             pack_seq_dim_size,
         )
+        if cfg["megatron_cfg"].get("batch_invariant_mode"):
+            # Packed scoring only needs its total token dimension aligned; padding
+            # every constituent sequence to the inference bucket size wastes work.
+            pad_packed_seq_to_multiple_of = batch_invariant_token_multiple(
+                pad_packed_seq_to_multiple_of,
+                cfg["megatron_cfg"]["tensor_model_parallel_size"],
+            )
+            if pad_full_seq_to is not None:
+                pad_full_seq_to = _round_up_to_multiple(
+                    pad_full_seq_to, pad_packed_seq_to_multiple_of
+                )
         micro_batch_size = 1
     else:
         raw_iterator = data.make_microbatch_iterator(mbs)
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index bada4a7f258..9d9b66fef34 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -228,6 +228,7 @@ def _sync_distrib_opt(distrib_opt):
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.models.generation.megatron.config import (
     dedicated_inference_megatron_cfg,
+    merged_inference_megatron_cfg,
 )
 from nemo_rl.models.megatron.community_import import (
     import_model_from_hf_name,
@@ -259,6 +260,109 @@ def _sync_distrib_opt(distrib_opt):
 TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase)
 
 
+def enable_batch_invariant_mode(config: PolicyConfig) -> None:
+    """Enable Megatron-Core batch-invariant kernels before CUDA initialization.
+
+    The mode is deliberately limited to the topology for which Megatron
+    generation and policy scoring can execute the same arithmetic.
+    Megatron-Core performs the remaining model-specific validation when the
+    provider is finalized. Sampling parameters do not affect this validation:
+    generation samples from the processed distribution while generation and
+    policy scoring both report raw model logprobs.
+
+    Args:
+        config: Policy configuration for this Megatron worker.
+
+    Raises:
+        ValueError: If batch-invariant mode is requested with an unsupported
+            NeMo-RL topology, generation backend, or precision.
+        AssertionError: If the installed Transformer Engine cannot pin the
+            requested FlashAttention version.
+    """
+    megatron_cfg = config["megatron_cfg"]
+    if not megatron_cfg.get("batch_invariant_mode"):
+        return
+
+    required_fields = (
+        "batch_invariant_backend",
+        "batch_invariant_collective",
+        "flash_attention_version",
+    )
+    missing_fields = [field for field in required_fields if field not in megatron_cfg]
+    if missing_fields:
+        raise ValueError(
+            "batch_invariant_mode=True requires policy.megatron_cfg fields: "
+            f"{', '.join(missing_fields)}."
+        )
+
+    if config["precision"] != "bfloat16":
+        raise ValueError(
+            "batch_invariant_mode=True requires policy.precision='bfloat16'."
+        )
+    if megatron_cfg["context_parallel_size"] != 1:
+        raise ValueError(
+            "batch_invariant_mode=True currently requires training context "
+            "parallel size 1."
+        )
+    if megatron_cfg.get("use_fused_linear_logprobs"):
+        raise ValueError(
+            "batch_invariant_mode=True is incompatible with "
+            "use_fused_linear_logprobs=True because generation parity requires "
+            "the shared float-log_softmax-gather path."
+        )
+    if megatron_cfg.get("attention_backend") != "flash":
+        raise ValueError(
+            "batch_invariant_mode=True requires "
+            "policy.megatron_cfg.attention_backend='flash'."
+        )
+    if megatron_cfg["flash_attention_version"] not in (3, 4):
+        raise ValueError(
+            "batch_invariant_mode=True requires "
+            "policy.megatron_cfg.flash_attention_version to be 3 or 4."
+        )
+
+    generation_cfg = config.get("generation")
+    if generation_cfg is None or generation_cfg["backend"] != "megatron":
+        raise ValueError(
+            "batch_invariant_mode=True requires policy.generation.backend='megatron'."
+        )
+    inference_cfg = merged_inference_megatron_cfg(config)
+    matching_fields = (
+        "tensor_model_parallel_size",
+        "context_parallel_size",
+        "batch_invariant_mode",
+        "batch_invariant_backend",
+        "batch_invariant_collective",
+        "attention_backend",
+        "flash_attention_version",
+    )
+    mismatched_fields = [
+        field
+        for field in matching_fields
+        if inference_cfg[field] != megatron_cfg[field]
+    ]
+    if mismatched_fields:
+        raise ValueError(
+            "Training and generation must use the same Megatron settings for "
+            f"batch invariance: {', '.join(mismatched_fields)}."
+        )
+
+    collective = megatron_cfg["batch_invariant_collective"]
+
+    # Keep optional Megatron GPU kernels out of imports when the mode is disabled.
+    from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+        assert_te_supports_batch_invariant_attention,
+    )
+    from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+        enable_batch_invariant_mode as enable_mcore_batch_invariant_mode,
+    )
+
+    assert_te_supports_batch_invariant_attention()
+    enable_mcore_batch_invariant_mode(
+        backend=megatron_cfg["batch_invariant_backend"], collective=collective
+    )
+
+
 def destroy_parallel_state():
     """Safely destroy parallel state and reset async call tracking.
 
@@ -350,12 +454,15 @@ def validate_and_set_config(
         generation_cfg = config["generation"]
         # set generation colocated
         is_generation_colocated = generation_cfg["colocated"]["enabled"]
-        # set sampling params
-        sampling_params = TrainingSamplingParams(
-            top_k=generation_cfg["top_k"],
-            top_p=generation_cfg["top_p"],
-            temperature=generation_cfg["temperature"],
-        )
+        # Batch-invariant Megatron inference returns raw model logprobs even
+        # when token sampling uses temperature, top-k, or top-p. Match
+        # Megatron-RL by recomputing raw training logprobs as well.
+        if not config["megatron_cfg"].get("batch_invariant_mode"):
+            sampling_params = TrainingSamplingParams(
+                top_k=generation_cfg["top_k"],
+                top_p=generation_cfg["top_p"],
+                temperature=generation_cfg["temperature"],
+            )
 
     # Explicitly set NCCL_CUMEM_ENABLE to 1 to avoid the P2P initialization error for PyNCCLCommunicator.
     # See https://github.com/NVIDIA-NeMo/RL/issues/564 for more details.
@@ -1140,6 +1247,21 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None:
                 f"Available backends are: {list(AttnBackend.__members__.keys())}"
             )
 
+    if "batch_invariant_mode" in config["megatron_cfg"]:
+        model_cfg.batch_invariant_mode = config["megatron_cfg"]["batch_invariant_mode"]
+    if "batch_invariant_backend" in config["megatron_cfg"]:
+        model_cfg.batch_invariant_backend = config["megatron_cfg"][
+            "batch_invariant_backend"
+        ]
+    if "batch_invariant_collective" in config["megatron_cfg"]:
+        model_cfg.batch_invariant_collective = config["megatron_cfg"][
+            "batch_invariant_collective"
+        ]
+    if "flash_attention_version" in config["megatron_cfg"]:
+        model_cfg.flash_attention_version = config["megatron_cfg"][
+            "flash_attention_version"
+        ]
+
     # These overrides need to be applied before the workers spawn.
     if "transformer_impl" in config["megatron_cfg"]:
         model_cfg.transformer_impl = config["megatron_cfg"]["transformer_impl"]
diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py
index 6e2cfd6e6a1..c41922edbbf 100644
--- a/nemo_rl/models/megatron/train.py
+++ b/nemo_rl/models/megatron/train.py
@@ -44,7 +44,7 @@
     prepare_packed_loss_input,
     wrap_loss_fn_with_input_preparation,
 )
-from nemo_rl.algorithms.loss.interfaces import LossFunction
+from nemo_rl.algorithms.loss.interfaces import LossFunction, LossInputType
 from nemo_rl.algorithms.utils import mask_out_neg_inf_logprobs
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.model_utils import (
@@ -52,6 +52,7 @@
     distributed_vocab_topk,
     from_parallel_logits_to_logprobs,
     from_parallel_logits_to_logprobs_packed_sequences,
+    get_next_token_logprobs_from_logits,
 )
 from nemo_rl.models.megatron.config import MegatronModule
 from nemo_rl.models.megatron.data import ProcessedMicrobatch
@@ -128,6 +129,7 @@ def model_forward(
     straggler_timer: Optional[StragglerDetector] = None,
     use_fused_linear_logprobs: bool = False,
     media_token_validity_mask: Optional[torch.Tensor] = None,
+    gather_output: bool = False,
 ) -> torch.Tensor:
     """Perform a single forward pass through the model.
 
@@ -147,6 +149,7 @@ def model_forward(
         media_token_validity_mask: Which media-token positions actually anchor a
             projected feature, already in this model's token layout. Only passed
             when the model accepts it; otherwise the model derives its own.
+        gather_output: Gather vocabulary-parallel logits on every TP rank.
 
     Returns:
         torch.Tensor: Output tensor from the model (logits)
@@ -178,6 +181,8 @@ def model_forward(
         # Only pass this kwarg when linear CE fusion is enabled. Older Megatron-LM
         # GPTModel.forward signatures do not accept it.
         additional_kwargs["return_logprobs_for_linear_ce_fusion"] = True
+    if gather_output:
+        additional_kwargs["runtime_gather_output"] = True
 
     with straggler_timer() if straggler_timer is not None else nullcontext():
         output_tensor = model(
@@ -270,6 +275,22 @@ def forward_with_post_processing_fn(
     routed_experts_cp_sharded = processed_mb.routed_experts_cp_sharded
     media_token_validity_mask = processed_mb.media_token_validity_mask
 
+    # Megatron generation materializes the full vocabulary before its canonical
+    # FP32 log_softmax. Mirror that arithmetic for batch-invariant policy
+    # logprobs at any TP degree. Other consumers retain sharded logits.
+    megatron_cfg = post_processing_fn.cfg.get("megatron_cfg")
+    gather_output = bool(
+        megatron_cfg
+        and megatron_cfg.get("batch_invariant_mode")
+        and (
+            isinstance(post_processing_fn, LogprobsPostProcessor)
+            or (
+                isinstance(post_processing_fn, LossPostProcessor)
+                and post_processing_fn.loss_fn.input_type == LossInputType.LOGPROB
+            )
+        )
+    )
+
     if use_router_replay:
         if routed_experts_cp_sharded is None:
             raise RuntimeError(
@@ -293,6 +314,7 @@ def forward_with_post_processing_fn(
                 straggler_timer=straggler_timer,
                 use_fused_linear_logprobs=use_fused_linear_logprobs,
                 media_token_validity_mask=media_token_validity_mask,
+                gather_output=gather_output,
             )
     except Exception:
         # The forward above armed the router-replay action (set_router_replay_forward);
@@ -502,6 +524,18 @@ def __call__(
         """
         # A custom prepare_fn (e.g. value models) overrides the default logit prep.
         logprob_chunk_size = self.cfg.get("logprob_chunk_size", None)
+        megatron_cfg = self.cfg.get("megatron_cfg")
+        full_vocab_logits = bool(
+            megatron_cfg
+            and megatron_cfg.get("batch_invariant_mode")
+            and self.loss_fn.input_type == LossInputType.LOGPROB
+        )
+        vocab_parallel_rank = (
+            None if full_vocab_logits else get_tensor_model_parallel_rank()
+        )
+        vocab_parallel_group = (
+            None if full_vocab_logits else get_tensor_model_parallel_group()
+        )
         if self.prepare_fn is not None:
             prepare_loss_input_wrapped = self.prepare_fn
         else:
@@ -540,8 +574,8 @@ def __call__(
                 prepare_fn=prepare_fn,
                 cu_seqlens_q=packed_seq_params.cu_seqlens_q,
                 cu_seqlens_q_padded=packed_seq_params.cu_seqlens_q_padded,
-                vocab_parallel_rank=get_tensor_model_parallel_rank(),
-                vocab_parallel_group=get_tensor_model_parallel_group(),
+                vocab_parallel_rank=vocab_parallel_rank,
+                vocab_parallel_group=vocab_parallel_group,
                 context_parallel_group=get_context_parallel_group(),
             )
         else:
@@ -549,8 +583,8 @@ def __call__(
                 wrap_loss_fn_with_input_preparation,
                 loss_fn=self.loss_fn,
                 prepare_fn=prepare_loss_input_wrapped,
-                vocab_parallel_rank=get_tensor_model_parallel_rank(),
-                vocab_parallel_group=get_tensor_model_parallel_group(),
+                vocab_parallel_rank=vocab_parallel_rank,
+                vocab_parallel_group=vocab_parallel_group,
                 context_parallel_group=get_context_parallel_group(),
             )
             if "student_logits" in data_dict:
@@ -559,8 +593,8 @@ def __call__(
                     prepare_fn=prepare_loss_input_wrapped,
                     data_dict=data_dict,
                     loss_weight=float(self.cfg["draft"]["loss_weight"]),
-                    vocab_parallel_rank=get_tensor_model_parallel_rank(),
-                    vocab_parallel_group=get_tensor_model_parallel_group(),
+                    vocab_parallel_rank=vocab_parallel_rank,
+                    vocab_parallel_group=vocab_parallel_group,
                     context_parallel_group=get_context_parallel_group(),
                 )
 
@@ -628,14 +662,20 @@ def __call__(
         """
         unpacked_input_ids = data_dict["input_ids"]
         original_seq_length = unpacked_input_ids.shape[1]
+        megatron_cfg = self.cfg.get("megatron_cfg")
+        full_vocab_logits = bool(
+            megatron_cfg and megatron_cfg.get("batch_invariant_mode")
+        )
 
         def processor_fn_inner(output_tensor):
             if self.use_fused_linear_logprobs:
                 token_logprobs = output_tensor.to(torch.float32)
                 token_logprobs = token_logprobs[:, : original_seq_length - 1]
             elif self.cfg["sequence_packing"]["enabled"]:
-                tp_grp = get_tensor_model_parallel_group()
-                tp_rank = get_tensor_model_parallel_rank()
+                tp_grp = (
+                    None if full_vocab_logits else get_tensor_model_parallel_group()
+                )
+                tp_rank = 0 if full_vocab_logits else get_tensor_model_parallel_rank()
                 logprob_chunk_size = self.cfg.get("logprob_chunk_size", None)
                 token_logprobs = from_parallel_logits_to_logprobs_packed_sequences(
                     output_tensor,
@@ -651,19 +691,26 @@ def processor_fn_inner(output_tensor):
                     sampling_params=self.sampling_params,
                 )
             else:
-                tp_grp = get_tensor_model_parallel_group()
-                tp_rank = get_tensor_model_parallel_rank()
                 logprob_chunk_size = self.cfg.get("logprob_chunk_size", None)
-                token_logprobs = from_parallel_logits_to_logprobs(
-                    output_tensor,
-                    target=unpacked_input_ids,
-                    vocab_start_index=tp_rank * output_tensor.shape[-1],
-                    vocab_end_index=(tp_rank + 1) * output_tensor.shape[-1],
-                    tp_group=tp_grp,
-                    inference_only=True,
-                    chunk_size=logprob_chunk_size,
-                    sampling_params=self.sampling_params,
-                )
+                if full_vocab_logits:
+                    token_logprobs = get_next_token_logprobs_from_logits(
+                        input_ids=unpacked_input_ids,
+                        next_token_logits=output_tensor,
+                        sampling_params=self.sampling_params,
+                    )
+                else:
+                    tp_grp = get_tensor_model_parallel_group()
+                    tp_rank = get_tensor_model_parallel_rank()
+                    token_logprobs = from_parallel_logits_to_logprobs(
+                        output_tensor,
+                        target=unpacked_input_ids,
+                        vocab_start_index=tp_rank * output_tensor.shape[-1],
+                        vocab_end_index=(tp_rank + 1) * output_tensor.shape[-1],
+                        tp_group=tp_grp,
+                        inference_only=True,
+                        chunk_size=logprob_chunk_size,
+                        sampling_params=self.sampling_params,
+                    )
 
             # Prepend 0 logprob for first token to maintain same sequence length as input
             token_logprobs = torch.cat(
diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py
index c40d7d51d07..bba8b5d9d1a 100644
--- a/nemo_rl/models/policy/__init__.py
+++ b/nemo_rl/models/policy/__init__.py
@@ -378,6 +378,22 @@ class MegatronConfig(TypedDict):
     # Attention backend available values:
     # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/enums.py#L60
     attention_backend: NotRequired[str]
+    # Enable Megatron-Core's batch-invariant kernels for bitwise-identical
+    # Megatron generation and policy logprobs. This mode currently supports
+    # BF16 with matching training/generation TP, CP=1, and FlashAttention 3/4.
+    # NeMo-RL automatically aligns the scoring token dimension with MCore's
+    # batch-invariant inference buckets; no recipe-level padding override is needed.
+    # Sampling may use any supported temperature, top-k, or top-p settings;
+    # generation and training compare the corresponding raw model logprobs.
+    batch_invariant_mode: NotRequired[bool]
+    # Megatron-Core kernel backend used by batch_invariant_mode. "te_native"
+    # is the performant default; "triton" and "deepgemm" are legacy options.
+    batch_invariant_backend: NotRequired[Literal["deepgemm", "te_native", "triton"]]
+    # Cross-rank EP combine. "ordered" is portable; "multimem" uses NVLS.
+    batch_invariant_collective: NotRequired[Literal["multimem", "ordered"]]
+    # Pin the FlashAttention generation used by both training and inference.
+    # batch_invariant_mode requires version 3 or 4.
+    flash_attention_version: NotRequired[Literal[2, 3, 4] | None]
     moe_per_layer_logging: bool
     # Set to true to enable DeepEP for expert parallel communication
     # Must set moe_token_dispatcher_type to 'flex'
diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py
index 1d7d62fd354..696e1917b15 100644
--- a/nemo_rl/models/policy/lm_policy.py
+++ b/nemo_rl/models/policy/lm_policy.py
@@ -38,6 +38,7 @@
     GenerationInterface,
     GenerationOutputSpec,
 )
+from nemo_rl.models.megatron.batch_invariant import batch_invariant_token_multiple
 from nemo_rl.models.policy import PolicyConfig
 from nemo_rl.models.policy.interfaces import (
     ColocatablePolicyInterface,
@@ -142,6 +143,27 @@ def __init__(
                 "Disable policy.sequence_packing.enabled or policy.draft."
             )
         if megatron_enable:
+            if (
+                config["megatron_cfg"].get("batch_invariant_mode")
+                and not config["sequence_packing"]["enabled"]
+            ):
+                # Native TE kernels are invariant when eager policy scoring uses
+                # the same aligned token dimension as MCore generation buckets.
+                # Sequence packing aligns its total token count in megatron.data.
+                tp_size = config["megatron_cfg"]["tensor_model_parallel_size"]
+                config["make_sequence_length_divisible_by"] = (
+                    batch_invariant_token_multiple(
+                        config["make_sequence_length_divisible_by"], tp_size
+                    )
+                )
+                if config["dynamic_batching"]["enabled"]:
+                    config["dynamic_batching"]["sequence_length_round"] = (
+                        batch_invariant_token_multiple(
+                            config["dynamic_batching"]["sequence_length_round"],
+                            tp_size,
+                        )
+                    )
+
             worker_builder_cls_fqn = resolve_policy_worker_cls(
                 "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker",
                 config,
@@ -151,6 +173,15 @@ def __init__(
             cp_size = config["megatron_cfg"]["context_parallel_size"]
 
             env_vars = dict(config["megatron_cfg"].get("env_vars") or {})
+            if (
+                config["megatron_cfg"].get("batch_invariant_mode")
+                and config["megatron_cfg"].get("batch_invariant_backend") == "te_native"
+            ):
+                # te_native obtains invariance by disabling cuBLASLt split-K.
+                # This must be set in the actor runtime environment: importing
+                # the Megatron worker can initialize CUDA before __init__ calls
+                # enable_batch_invariant_mode().
+                env_vars["CUBLASLT_WORKSPACE_SIZE"] = "0"
 
             if "TORCH_CUDA_ARCH_LIST" not in os.environ:
                 raise RuntimeError(
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index 2df57f7cf8f..979b1bfa255 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -82,6 +82,7 @@
 from nemo_rl.models.megatron.router_replay import router_replay_enabled
 from nemo_rl.models.megatron.setup import (
     build_inference_model,
+    enable_batch_invariant_mode,
     finalize_megatron_setup,
     handle_model_import,
     setup_distributed,
@@ -441,6 +442,7 @@ def __init__(
         gpu_ids = ray.get_gpu_ids()
         local_rank = int(gpu_ids[0])
         os.environ["LOCAL_RANK"] = str(local_rank)
+        enable_batch_invariant_mode(config)
         torch.cuda.set_device(local_rank)
 
         # Apply patch from https://github.com/NVIDIA/TransformerEngine/pull/2286/files
diff --git a/pyproject.toml b/pyproject.toml
index 00a6822e923..a51a2f7ba2b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -187,6 +187,11 @@ mcore = [
   "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'",
   "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'",
   "flash-attn==2.8.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')",
+  # FlashAttention 4 is the batch-invariant attention implementation for CUDA 13.
+  "flash-attn-4[cu13]==4.0.0b20 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'x86_64')",
+  # FA4 b20 requires this pre-release; pin it here instead of allowing
+  # pre-releases globally across NeMo-RL's dependency graph.
+  "nvidia-cutlass-dsl[cu13]==4.6.0.dev0 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'x86_64')",
   # mamba-ssm/causal-conv1d LOOK redundant (megatron-bridge[ssm] pulls them), but do not
   # remove them: [tool.uv.sources] git pins only attach to direct requirements, and the
   # conflicting extras resolve in separate forks — without these direct deps the mcore fork
@@ -342,7 +347,9 @@ default-groups = ["dev", "build"]
 #  --link-mode=copy (slower but more reliable; supresses warning)
 #  --link-mode=symlink (fastest option when uv cache and venv on different file-system; caveat: venv is brittle since it depends on the environment/container)
 link-mode = "copy"
-# The TE override is needed to pin the TE version across all extras
+# The TE override is needed to pin the TE version across all extras. Keep it
+# aligned with Megatron-Core: TE 2.18 adds the explicit FlashAttention version
+# selection required by batch-invariant mode.
 # The opencv-python-headless override excludes it from the shipped container (sys_platform == 'never')
 # because its bundled FFmpeg codec libs (H.264, H.265, AAC) may incur royalties. The CVE floor
 # (>=5.0.0) is kept separately in constraint-dependencies so it applies to anyone who does install it.
@@ -350,10 +357,11 @@ link-mode = "copy"
 # vllm defaults to opencv for video IO but falls back to torchcodec.
 # The timm override is needed because sglang requires timm==1.0.16.
 override-dependencies = [
-  "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.15",
+  "transformer-engine[pytorch,core_cu13]==2.18.0",
   "nvidia-cublas==13.5.1.27; sys_platform != 'darwin'",
   "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'",
-  "nvidia-cudnn-frontend==1.23.0",
+  # TE 2.18 requires >=1.25; 1.26 matches Megatron-Core and Megatron-Bridge.
+  "nvidia-cudnn-frontend==1.26.0",
   # NCCL 2.30.4 fixes a context-parallel hang on GB200 when GPUs are non-contiguous;
   # 2.30.7 additionally fixes hangs in the NCCL m2n reshard (<link to NCCL release note / internal bug>).
   # An override is required because torch 2.11.0+cu130 pins nvidia-nccl-cu13==2.28.9.
@@ -376,11 +384,11 @@ override-dependencies = [
   # xgrammar.normalize_tool_choice (vLLM requires xgrammar>=0.2.1,<1.0.0; the
   # override would otherwise force an incompatible downgrade).
   "xgrammar>=0.2.1,<1.0.0",
-  # tensorrt-llm 1.3.0rc21 pins apache-tvm-ffi==0.1.6, but the xgrammar the
-  # line above selects needs >=0.1.9, and the xgrammar releases that would
-  # accept 0.1.6 ship no cp313 wheels. Take the newer one, as with llguidance
-  # above; without this the trtllm extra is unsatisfiable on vLLM 0.25.1.
-  "apache-tvm-ffi>=0.1.9",
+  # tensorrt-llm 1.3.0rc21 pins apache-tvm-ffi==0.1.6, while xgrammar needs
+  # >=0.1.9 and FlashAttention 4.0.0b20 needs >=0.1.12. Take the newer API;
+  # FA4's CUTLASS wrapper passes map_dataclass_to_tuple, which older releases
+  # do not accept.
+  "apache-tvm-ffi>=0.1.12,<0.2",
   # Override dependencies to address CVEs
   "mlflow>=3.13.0",
   # Override langchain-core to address CVE GHSA-qh6h-p6c9-ff54 (path traversal in load_prompt)
@@ -410,7 +418,7 @@ constraint-dependencies = [
   # whose [tool.uv.sources] pins them to git. uv requires such URL deps to also appear as a
   # direct requirement or constraint of the root project. Keep these URLs/revs in sync with
   # Megatron-LM's [tool.uv.sources] (uv errors loudly on mismatch after a submodule bump).
-  "emerging-optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.2.0",
+  "emerging-optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.3.0",
   "fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git@f134af63deb2df17e1171a9ec1ea4a7d8604d5ca",
   # megatron-energon pulls s3fs; without a floor uv picks the 2020-era s3fs 0.4.2
   # (which needs only plain botocore) instead of a modern aiobotocore-based release.
@@ -561,16 +569,6 @@ name = "deep_gemm"
 version = "2.5.0"
 requires-dist = ["torch", "packaging", "ninja"]
 
-[[tool.uv.dependency-metadata]]
-name = "transformer-engine"
-version = "2.15.0+42b8400"
-requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"]
-
-[[tool.uv.dependency-metadata]]
-name = "transformer-engine-torch"
-version = "2.15.0+42b8400"
-requires-dist = ["torch", "transformer-engine"]
-
 [[tool.uv.dependency-metadata]]
 name = "nv-grouped-gemm"
 # This version has to match the version in the commit/rev/tag used
diff --git a/pyrefly.toml b/pyrefly.toml
index 7a1691b0877..6eef162a25a 100644
--- a/pyrefly.toml
+++ b/pyrefly.toml
@@ -219,6 +219,7 @@ project-includes = [
   "nemo_rl/models/generation/vllm/worker_utils.py",
   "nemo_rl/models/huggingface/__init__.py",
   "nemo_rl/models/megatron/__init__.py",
+  "nemo_rl/models/megatron/batch_invariant.py",
   "nemo_rl/models/megatron/draft/__init__.py",
   "nemo_rl/models/megatron/memory_saver.py",
   "nemo_rl/models/policy/__init__.py",
diff --git a/tests/unit/distributed/test_model_utils.py b/tests/unit/distributed/test_model_utils.py
index 72f626c001f..fe3d1116d3d 100644
--- a/tests/unit/distributed/test_model_utils.py
+++ b/tests/unit/distributed/test_model_utils.py
@@ -42,6 +42,34 @@
 from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup
 
 
+def test_packed_full_vocab_logprobs_use_canonical_log_softmax():
+    """Gathered packed logits use differentiable FP32 log_softmax without TP ops."""
+    logits = torch.randn(1, 6, 8, dtype=torch.bfloat16, requires_grad=True)
+    targets = torch.tensor([[1, 2, 3, 4, 5, 6]])
+    cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32)
+
+    actual = from_parallel_logits_to_logprobs_packed_sequences(
+        logits,
+        targets,
+        cu_seqlens,
+        unpacked_seqlen=3,
+        vocab_start_index=0,
+        vocab_end_index=8,
+        group=None,
+    )
+    full_logprobs = torch.nn.functional.log_softmax(logits.float(), dim=-1)
+    expected = torch.stack(
+        (
+            torch.stack((full_logprobs[0, 0, 2], full_logprobs[0, 1, 3])),
+            torch.stack((full_logprobs[0, 3, 5], full_logprobs[0, 4, 6])),
+        )
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    actual.sum().backward()
+    assert logits.grad is not None
+
+
 @ray.remote(num_gpus=1)
 class ModelUtilsTestActor:
     def __init__(self, tp_size, cp_size, sharding):
diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py
index e8fbbf26e9b..f0beaade390 100644
--- a/tests/unit/models/generation/test_megatron_generation.py
+++ b/tests/unit/models/generation/test_megatron_generation.py
@@ -251,6 +251,32 @@ def _assert_valid_generation_output(outputs, input_data, require_generation=True
     ), "logprobs[:, 0] should be the 0.0 placeholder"
 
 
+def _assert_exact_generation_logprobs(policy, outputs, prompt_lengths):
+    """Assert exact generation-vs-policy logprobs on generated tokens."""
+    fprop_data = BatchedDataDict(
+        {
+            "input_ids": outputs["output_ids"],
+            "input_lengths": outputs["unpadded_sequence_lengths"],
+        }
+    )
+    policy.prepare_for_lp_inference()
+    policy_logprobs = policy.get_logprobs(fprop_data)["logprobs"]
+
+    generated_mask = torch.zeros_like(outputs["logprobs"], dtype=torch.bool)
+    for row, (start, end) in enumerate(
+        zip(prompt_lengths, outputs["unpadded_sequence_lengths"])
+    ):
+        generated_mask[row, start:end] = True
+
+    generation = outputs["logprobs"].masked_select(generated_mask)
+    scoring = policy_logprobs.masked_select(generated_mask)
+    assert torch.equal(generation, scoring), (
+        f"generation and policy logprobs differ at "
+        f"{torch.count_nonzero(generation != scoring).item()}/{generation.numel()} "
+        f"generated tokens (max abs diff {(generation - scoring).abs().max().item()})"
+    )
+
+
 async def _generate_async(mg, tokenizer, test_input_data, greedy=False):
     """Drive ``generate_async`` over single-sample microbatches and reassemble in order."""
     collected = []
@@ -557,6 +583,80 @@ def test_megatron_generation_colocated(
         torch.cuda.empty_cache()
 
 
+@pytest.mark.mcore
+@pytest.mark.timeout(900)
+@pytest.mark.parametrize(
+    ("temperature", "top_k", "top_p"),
+    [
+        pytest.param(1.0, None, 1.0, id="unfiltered"),
+        pytest.param(0.7, 32, 1.0, id="temperature-top-k"),
+        pytest.param(1.3, None, 0.9, id="temperature-top-p"),
+    ],
+)
+def test_batch_invariant_generation_logprobs_are_exact(
+    test_input_data, tokenizer, temperature, top_k, top_p
+):
+    """Raw generation and policy logprobs are bitwise equal for TP=2 sampling."""
+    from megatron.core.transformer import attention as mcore_attention
+
+    if mcore_attention.HAVE_FA4:
+        flash_attention_version = 4
+    elif mcore_attention.HAVE_FA3:
+        flash_attention_version = 3
+    else:
+        pytest.skip("batch-invariant parity requires FlashAttention 3 or 4")
+
+    cluster = RayVirtualCluster(
+        bundle_ct_per_node_list=[2],
+        use_gpus=True,
+        max_colocated_worker_groups=2,
+        num_gpus_per_node=2,
+        name="megatron-batch-invariant-test-cluster",
+    )
+    config = deepcopy(basic_megatron_test_config)
+    config["generation"]["colocated"]["enabled"] = True
+    config["generation"].update(
+        {
+            "temperature": temperature,
+            "top_k": top_k,
+            "top_p": top_p,
+        }
+    )
+    config["megatron_cfg"].update(
+        {
+            "attention_backend": "flash",
+            "batch_invariant_mode": True,
+            "batch_invariant_backend": "te_native",
+            "batch_invariant_collective": "ordered",
+            "flash_attention_version": flash_attention_version,
+            "tensor_model_parallel_size": 2,
+            "sequence_parallel": True,
+            "use_fused_linear_logprobs": False,
+        }
+    )
+
+    policy = None
+    try:
+        policy = Policy(cluster=cluster, config=config, tokenizer=tokenizer)
+        generation = MegatronGeneration(
+            policy=policy, config=config, tokenizer=tokenizer
+        )
+        generation.prepare_for_generation()
+        outputs = generation.generate(test_input_data, greedy=False)
+        _assert_valid_generation_output(outputs, test_input_data)
+
+        generation.finish_generation(release_gpu=True)
+        _assert_exact_generation_logprobs(
+            policy, outputs, test_input_data["input_lengths"]
+        )
+    finally:
+        if policy is not None:
+            policy.shutdown()
+        cluster.shutdown()
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
 @pytest.mark.mcore
 @pytest.mark.timeout(900)
 @pytest.mark.parametrize("skip_weight_load", [False, True])
diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py
index 8e321d29003..a246f39c0ed 100644
--- a/tests/unit/models/megatron/test_megatron_data.py
+++ b/tests/unit/models/megatron/test_megatron_data.py
@@ -1254,6 +1254,52 @@ def test_get_microbatch_iterator_sequence_packing(
         assert micro_batch_size == 1
         assert data_iterator_len == 10
 
+    @patch("nemo_rl.models.megatron.data.get_and_validate_seqlen")
+    @patch("nemo_rl.models.megatron.data.make_processed_microbatch_iterator")
+    @patch("nemo_rl.models.megatron.data._get_pack_sequence_parameters_for_megatron")
+    def test_batch_invariant_sequence_packing_aligns_only_total_tokens(
+        self, mock_get_params, mock_make_iterator, mock_get_and_validate_seqlen
+    ):
+        """Batch invariance aligns the packed M dimension, not every sequence."""
+        from nemo_rl.models.megatron.data import get_microbatch_iterator
+
+        mock_get_and_validate_seqlen.return_value = (1, 257)
+        mock_get_params.return_value = (2, 16, 513)
+        mock_data = MagicMock()
+        mock_data.make_microbatch_iterator_for_packable_sequences.return_value = iter(
+            []
+        )
+        mock_data.get_microbatch_iterator_for_packable_sequences_len.return_value = (
+            1,
+            513,
+        )
+
+        cfg = {
+            "dynamic_batching": {"enabled": False},
+            "sequence_packing": {"enabled": True},
+            "megatron_cfg": {
+                "batch_invariant_mode": True,
+                "tensor_model_parallel_size": 2,
+                "sequence_parallel": True,
+                "pipeline_model_parallel_size": 2,
+                "context_parallel_size": 1,
+            },
+            "make_sequence_length_divisible_by": 2,
+        }
+
+        *_, padded_seq_length = get_microbatch_iterator(
+            data=mock_data,
+            cfg=cfg,
+            mbs=1,
+            straggler_timer=MagicMock(),
+        )
+
+        call_kwargs = mock_make_iterator.call_args.kwargs
+        assert call_kwargs["pad_individual_seqs_to_multiple_of"] == 2
+        assert call_kwargs["pad_packed_seq_to_multiple_of"] == 64
+        assert call_kwargs["pad_full_seq_to"] == 576
+        assert padded_seq_length == 576
+
     @patch("nemo_rl.models.megatron.data.get_and_validate_seqlen")
     @patch("nemo_rl.models.megatron.data.make_processed_microbatch_iterator")
     def test_get_microbatch_iterator_regular(
diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py
index 0fcdc8bc485..3356e545ae1 100644
--- a/tests/unit/models/megatron/test_megatron_setup.py
+++ b/tests/unit/models/megatron/test_megatron_setup.py
@@ -1021,6 +1021,28 @@ def _config(*, attention_backend=None):
             megatron_cfg["attention_backend"] = attention_backend
         return {"megatron_cfg": megatron_cfg}
 
+    def test_batch_invariant_fields_are_forwarded(self):
+        """Batch-invariant settings reach the Megatron model provider."""
+        from nemo_rl.models.megatron.setup import _apply_performance_config
+
+        model_cfg = SimpleNamespace(gated_linear_unit=True)
+        config = self._config(attention_backend="flash")
+        config["megatron_cfg"].update(
+            {
+                "batch_invariant_mode": True,
+                "batch_invariant_backend": "triton",
+                "batch_invariant_collective": "multimem",
+                "flash_attention_version": 3,
+            }
+        )
+
+        _apply_performance_config(model_cfg, config)
+
+        assert model_cfg.batch_invariant_mode is True
+        assert model_cfg.batch_invariant_backend == "triton"
+        assert model_cfg.batch_invariant_collective == "multimem"
+        assert model_cfg.flash_attention_version == 3
+
     def test_cuda_graph_training_values_are_forwarded(self):
         """Explicit training CUDA Graph settings are normalized on assignment."""
         from megatron.core.transformer.enums import CudaGraphModule
@@ -2239,6 +2261,172 @@ def test_fraction_and_transfer_overlap_are_forwarded(self, transfer_overlap):
             assert optimizer_kwargs["overlap_cpu_optimizer_d2h_h2d"] is transfer_overlap
 
 
+@pytest.mark.mcore
+class TestBatchInvariantMode:
+    """Tests for early batch-invariant activation and validation."""
+
+    @staticmethod
+    def _config() -> dict[str, Any]:
+        return {
+            "precision": "bfloat16",
+            "megatron_cfg": {
+                "batch_invariant_mode": True,
+                "batch_invariant_backend": "te_native",
+                "batch_invariant_collective": "ordered",
+                "flash_attention_version": 3,
+                "attention_backend": "flash",
+                "tensor_model_parallel_size": 1,
+                "context_parallel_size": 1,
+                "use_fused_linear_logprobs": False,
+            },
+            "generation": {
+                "backend": "megatron",
+                "temperature": 1.0,
+                "top_k": None,
+                "top_p": 1.0,
+                "mcore_generation_config": {},
+            },
+        }
+
+    def test_enables_selected_mcore_backend(self):
+        """A valid config activates the requested MCore backend."""
+        from nemo_rl.models.megatron.setup import enable_batch_invariant_mode
+
+        support_check = (
+            "megatron.core.transformer.custom_layers.batch_invariant_kernels."
+            "assert_te_supports_batch_invariant_attention"
+        )
+        with (
+            patch(
+                "megatron.core.transformer.custom_layers.batch_invariant_kernels.enable_batch_invariant_mode"
+            ) as enable_mcore,
+            patch(support_check) as assert_te_support,
+        ):
+            enable_batch_invariant_mode(self._config())
+
+        assert_te_support.assert_called_once_with()
+        enable_mcore.assert_called_once_with(backend="te_native", collective="ordered")
+
+    @pytest.mark.parametrize(
+        "missing_field",
+        [
+            "batch_invariant_backend",
+            "batch_invariant_collective",
+            "flash_attention_version",
+        ],
+    )
+    def test_rejects_missing_required_field(self, missing_field: str):
+        """Mode-specific settings must come from the resolved recipe config."""
+        from nemo_rl.models.megatron.setup import enable_batch_invariant_mode
+
+        config = self._config()
+        config["megatron_cfg"].pop(missing_field)
+
+        with pytest.raises(ValueError, match=missing_field):
+            enable_batch_invariant_mode(config)
+
+    def test_disabled_mode_is_a_noop(self):
+        """Absent opt-in does not activate batch-invariant kernels."""
+        from nemo_rl.models.megatron.setup import enable_batch_invariant_mode
+
+        with patch(
+            "megatron.core.transformer.custom_layers.batch_invariant_kernels.enable_batch_invariant_mode"
+        ) as enable_mcore:
+            enable_batch_invariant_mode({"megatron_cfg": {}})
+
+        enable_mcore.assert_not_called()
+
+    def test_supports_matching_tensor_parallelism(self):
+        """Batch-invariant logprobs support TP>1 when generation uses the same TP."""
+        from nemo_rl.models.megatron.setup import enable_batch_invariant_mode
+
+        config = self._config()
+        config["megatron_cfg"]["tensor_model_parallel_size"] = 2
+        config["generation"]["mcore_generation_config"][
+            "tensor_model_parallel_size"
+        ] = 2
+
+        support_check = (
+            "megatron.core.transformer.custom_layers.batch_invariant_kernels."
+            "assert_te_supports_batch_invariant_attention"
+        )
+        with (
+            patch(
+                "megatron.core.transformer.custom_layers.batch_invariant_kernels.enable_batch_invariant_mode"
+            ) as enable_mcore,
+            patch(support_check),
+        ):
+            enable_batch_invariant_mode(config)
+
+        enable_mcore.assert_called_once_with(backend="te_native", collective="ordered")
+
+    @pytest.mark.parametrize(
+        ("field", "value"),
+        [
+            ("temperature", 0.7),
+            ("top_k", 32),
+            ("top_p", 0.9),
+        ],
+    )
+    def test_supports_sampling_parameters(self, field: str, value: Any):
+        """Sampling and raw-logprob parity are independent contracts."""
+        from nemo_rl.models.megatron.setup import enable_batch_invariant_mode
+
+        config = self._config()
+        config["generation"][field] = value
+
+        support_check = (
+            "megatron.core.transformer.custom_layers.batch_invariant_kernels."
+            "assert_te_supports_batch_invariant_attention"
+        )
+        with (
+            patch(
+                "megatron.core.transformer.custom_layers.batch_invariant_kernels.enable_batch_invariant_mode"
+            ) as enable_mcore,
+            patch(support_check),
+        ):
+            enable_batch_invariant_mode(config)
+
+        enable_mcore.assert_called_once_with(backend="te_native", collective="ordered")
+
+    def test_rejects_mismatched_tensor_parallelism(self):
+        """TP degree may be greater than one but must match across both models."""
+        from nemo_rl.models.megatron.setup import enable_batch_invariant_mode
+
+        config = self._config()
+        config["megatron_cfg"]["tensor_model_parallel_size"] = 2
+        config["generation"]["mcore_generation_config"][
+            "tensor_model_parallel_size"
+        ] = 4
+
+        with pytest.raises(ValueError, match="tensor_model_parallel_size"):
+            enable_batch_invariant_mode(config)
+
+    @pytest.mark.parametrize(
+        ("section", "field", "value", "error"),
+        [
+            ("generation", "backend", "vllm", "generation.backend='megatron'"),
+            (
+                "megatron_cfg",
+                "use_fused_linear_logprobs",
+                True,
+                "incompatible with use_fused_linear_logprobs",
+            ),
+        ],
+    )
+    def test_rejects_non_exact_configurations(
+        self, section: str, field: str, value: Any, error: str
+    ) -> None:
+        """Unsupported topology choices fail before CUDA setup."""
+        from nemo_rl.models.megatron.setup import enable_batch_invariant_mode
+
+        config = self._config()
+        config[section][field] = value
+
+        with pytest.raises(ValueError, match=error):
+            enable_batch_invariant_mode(config)
+
+
 @pytest.mark.mcore
 class TestValidateAndSetConfig:
     """Tests for validate_and_set_config function."""
@@ -2316,6 +2504,49 @@ def test_generation_colocation_detection(self):
 
                 assert runtime_config.is_generation_colocated is True
                 assert runtime_config.offload_optimizer_for_refit is True
+                assert runtime_config.sampling_params is not None
+
+    def test_batch_invariant_mode_recomputes_raw_logprobs(self):
+        """Sampling parameters select tokens but do not process policy logprobs."""
+        from nemo_rl.models.megatron.setup import validate_and_set_config
+
+        config = {
+            "generation": {
+                "temperature": 0.7,
+                "top_p": 0.9,
+                "top_k": None,
+                "colocated": {"enabled": True},
+            },
+            "precision": "bfloat16",
+            "megatron_cfg": {
+                "batch_invariant_mode": True,
+                "optimizer": {"optimizer_cpu_offload": False},
+                "tensor_model_parallel_size": 2,
+            },
+            "offload_optimizer_for_logprob": False,
+        }
+
+        with (
+            patch("nemo_rl.models.megatron.setup.setup_model_config") as setup_config,
+            patch(
+                "nemo_rl.models.megatron.setup.calculate_padded_vocab_size",
+                return_value=32000,
+            ),
+        ):
+            megatron_cfg = MagicMock()
+            megatron_cfg.model.vocab_size = 32000
+            setup_config.return_value = (megatron_cfg, MagicMock())
+
+            runtime_config = validate_and_set_config(
+                config=config,
+                rank=0,
+                hf_model_name="test-model",
+                pretrained_path="/path/to/model",
+                weights_path=None,
+                optimizer_path=None,
+            )
+
+        assert runtime_config.sampling_params is None
 
 
 @pytest.mark.mcore
diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py
index 130cbf9825b..ce16163715d 100644
--- a/tests/unit/models/megatron/test_train.py
+++ b/tests/unit/models/megatron/test_train.py
@@ -143,6 +143,25 @@ def test_model_forward_with_defer_fp32_logits(self):
         call_kwargs = mock_model.call_args[1]
         assert call_kwargs["fp32_output"] is False
 
+    def test_model_forward_can_gather_tensor_parallel_logits(self):
+        """The logprob path can request full-vocabulary logits at runtime."""
+        from nemo_rl.models.megatron.train import model_forward
+
+        mock_model = MagicMock(return_value=torch.randn(1, 3, 100))
+        mock_data_dict = MagicMock()
+        mock_data_dict.get_multimodal_dict.return_value = {}
+
+        model_forward(
+            model=mock_model,
+            data_dict=mock_data_dict,
+            input_ids_cp_sharded=torch.tensor([[1, 2, 3]]),
+            position_ids=torch.tensor([[0, 1, 2]]),
+            attention_mask=torch.ones(1, 3),
+            gather_output=True,
+        )
+
+        assert mock_model.call_args.kwargs["runtime_gather_output"] is True
+
     def test_model_forward_clears_position_ids_for_multimodal(self):
         """Test model_forward sets position_ids to None for multimodal data."""
         from nemo_rl.models.megatron.train import model_forward
@@ -288,7 +307,10 @@ def test_forward_with_logprobs_post_processor(self, mock_model_forward):
         )
 
         data_iterator = iter([processed_mb])
-        cfg = {"sequence_packing": {"enabled": False}}
+        cfg = {
+            "sequence_packing": {"enabled": False},
+            "megatron_cfg": {"batch_invariant_mode": True},
+        }
         post_processor = LogprobsPostProcessor(cfg=cfg)
 
         with patch.object(post_processor, "__call__", return_value=MagicMock()):
@@ -299,6 +321,7 @@ def test_forward_with_logprobs_post_processor(self, mock_model_forward):
             )
 
         mock_model_forward.assert_called_once()
+        assert mock_model_forward.call_args.kwargs["gather_output"] is True
 
     @patch("nemo_rl.models.megatron.train.model_forward")
     def test_forward_with_topk_post_processor(self, mock_model_forward):
@@ -1113,6 +1136,42 @@ def test_loss_post_processor_no_packing(
         assert isinstance(metrics, dict)
         assert len(metrics) == 1 and metrics["loss"] == 0.5
 
+    @patch("nemo_rl.models.megatron.train.wrap_loss_fn_with_input_preparation")
+    @patch("nemo_rl.models.megatron.train.get_tensor_model_parallel_group")
+    @patch("nemo_rl.models.megatron.train.get_tensor_model_parallel_rank")
+    @patch("nemo_rl.models.megatron.train.get_context_parallel_group")
+    @patch(
+        "nemo_rl.models.megatron.train.get_context_parallel_world_size", return_value=1
+    )
+    def test_batch_invariant_logprob_loss_uses_full_vocab_logits(
+        self,
+        mock_cp_size,
+        mock_cp_group,
+        mock_tp_rank,
+        mock_tp_group,
+        mock_wrap_loss,
+    ):
+        """Gathered logits must not be interpreted as a TP vocabulary shard."""
+        from nemo_rl.models.megatron.train import LossPostProcessor
+
+        mock_loss_fn = MagicMock()
+        mock_loss_fn.input_type = LossInputType.LOGPROB
+        mock_wrap_loss.return_value = (torch.tensor(0.5), {})
+        cfg = {
+            "sequence_packing": {"enabled": False},
+            "megatron_cfg": {"batch_invariant_mode": True},
+        }
+
+        wrapped_fn = LossPostProcessor(
+            loss_fn=mock_loss_fn, cfg=cfg, cp_normalize=False
+        )(data_dict=MagicMock())
+        wrapped_fn(torch.randn(1, 3, 8))
+
+        assert mock_wrap_loss.call_args.kwargs["vocab_parallel_rank"] is None
+        assert mock_wrap_loss.call_args.kwargs["vocab_parallel_group"] is None
+        mock_tp_rank.assert_not_called()
+        mock_tp_group.assert_not_called()
+
     @patch(
         "nemo_rl.models.megatron.train.get_tensor_model_parallel_rank", return_value=0
     )
@@ -1225,6 +1284,33 @@ def test_logprobs_post_processor_no_packing(
         # Logprobs should be prepended with a 0
         assert result["logprobs"].shape[1] == 5
 
+    @patch("nemo_rl.models.megatron.train.get_tensor_model_parallel_group")
+    @patch("nemo_rl.models.megatron.train.get_tensor_model_parallel_rank")
+    @patch("nemo_rl.models.megatron.train.from_parallel_logits_to_logprobs")
+    @patch("nemo_rl.models.megatron.train.get_next_token_logprobs_from_logits")
+    def test_batch_invariant_logprobs_use_full_vocab_path(
+        self, mock_full_vocab, mock_distributed, mock_tp_rank, mock_tp_group
+    ):
+        """Batch-invariant rescoring mirrors generation's full-vocab log_softmax."""
+        from nemo_rl.models.megatron.train import LogprobsPostProcessor
+
+        cfg = {
+            "sequence_packing": {"enabled": False},
+            "megatron_cfg": {"batch_invariant_mode": True},
+        }
+        processor = LogprobsPostProcessor(cfg=cfg)
+        input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        data = MagicMock()
+        data.__getitem__ = MagicMock(return_value=input_ids)
+        mock_full_vocab.return_value = torch.randn(1, 4)
+
+        processor(data, input_ids, None)(torch.randn(1, 5, 16))
+
+        mock_full_vocab.assert_called_once()
+        mock_distributed.assert_not_called()
+        mock_tp_rank.assert_not_called()
+        mock_tp_group.assert_not_called()
+
     @patch("nemo_rl.models.megatron.train.get_tensor_model_parallel_group")
     @patch(
         "nemo_rl.models.megatron.train.get_tensor_model_parallel_rank", return_value=0
diff --git a/tests/unit/models/policy/test_policy_validation.py b/tests/unit/models/policy/test_policy_validation.py
index 6a4c6a8d3ab..2ad6c75be24 100644
--- a/tests/unit/models/policy/test_policy_validation.py
+++ b/tests/unit/models/policy/test_policy_validation.py
@@ -131,6 +131,7 @@ def create_megatron_config(
         "logprob_batch_size": 1,
         "precision": "float32",
         "offload_optimizer_for_logprob": False,
+        "make_sequence_length_divisible_by": tp,
         "generation": {
             "backend": "hf",
             "temperature": 1.0,
@@ -446,3 +447,35 @@ def test_world_size_validation_megatron(
             )
         # For failing cases, worker group should not be created
         mock_ray_worker_group.assert_not_called()
+
+
+@pytest.mark.mcore
+@patch("nemo_rl.models.policy.lm_policy.RayQueue")
+@patch("nemo_rl.models.policy.lm_policy.RayWorkerGroup")
+def test_batch_invariant_te_native_pins_worker_cublaslt_workspace(
+    mock_ray_worker_group,
+    _mock_ray_queue,
+):
+    """The te_native backend pins cuBLASLt before actor CUDA imports."""
+    cluster = create_mock_cluster(world_size=2)
+    tokenizer = create_mock_tokenizer()
+    config = create_megatron_config("unused-model", tp=2)
+    config["megatron_cfg"].update(
+        {
+            "batch_invariant_mode": True,
+            "batch_invariant_backend": "te_native",
+            "env_vars": {"CUBLASLT_WORKSPACE_SIZE": "1048576"},
+        }
+    )
+
+    with patch(
+        "nemo_rl.models.policy.lm_policy.get_default_hf_config",
+        side_effect=ValueError("unused in this test"),
+    ):
+        Policy(cluster=cluster, config=config, tokenizer=tokenizer)
+
+    assert mock_ray_worker_group.call_args.kwargs["env_vars"] == {
+        "CUBLASLT_WORKSPACE_SIZE": "0"
+    }
+    assert config["make_sequence_length_divisible_by"] == 64
+    assert config["dynamic_batching"]["sequence_length_round"] == 64
diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml
index 344c5820806..089cde2b6be 100644
--- a/tests/unit/reference_configs/grpo_math_1B.yaml
+++ b/tests/unit/reference_configs/grpo_math_1B.yaml
@@ -187,6 +187,10 @@ policy:
       ckpt_assume_constant_structure: true
     use_fused_linear_logprobs: false
     fused_linear_logprobs_chunk_size: 256
+    batch_invariant_mode: false
+    batch_invariant_backend: "te_native"
+    batch_invariant_collective: "ordered"
+    flash_attention_version: null
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false