[WS2][Logp] Expose the lse_vocab_tile_reduction ablation switch on linear-logp ops

rl_engine/kernels/ops/cuda/loss/linear_logp.py

@@ -10,6 +10,8 @@
 
 from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE
 from rl_engine.kernels.ops.pytorch.loss.linear_logp import (
+    LseVocabTileReduction,
+    _check_lse_vocab_tile_reduction,
     _merge_tp_local_logp,
     _require_distributed_initialized,
     _validate_global_targets,
@@ -807,6 +809,7 @@ def __call__(
         tp_group: Any = None,
         vocab_start_index: int = 0,
         global_vocab_size: Optional[int] = None,
+        lse_vocab_tile_reduction: LseVocabTileReduction = "original",
     ) -> torch.Tensor:
         return self.apply(
             hidden,
@@ -816,6 +819,7 @@ def __call__(
             tp_group=tp_group,
             vocab_start_index=vocab_start_index,
             global_vocab_size=global_vocab_size,
+            lse_vocab_tile_reduction=lse_vocab_tile_reduction,
         )
 
     def apply(
@@ -828,8 +832,10 @@ def apply(
         tp_group: Any = None,
         vocab_start_index: int = 0,
         global_vocab_size: Optional[int] = None,
+        lse_vocab_tile_reduction: LseVocabTileReduction = "original",
     ) -> torch.Tensor:
         global _SM90_SAVE_PROBS_BF16_PATH_LOGGED, _SM90_FUSED_TILE_BF16_PATH_LOGGED
+        _check_lse_vocab_tile_reduction(lse_vocab_tile_reduction)
         if lm_head_weight.size(-1) != hidden.size(-1):
             raise ValueError(
                 f"hidden dim {hidden.size(-1)} must match lm_head_weight dim "

rl_engine/kernels/ops/pytorch/loss/linear_logp.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import os
-from typing import Any, Optional
+from typing import Any, Literal, Optional
 
 import torch
 
@@ -15,6 +15,18 @@
 _LOW_PRECISION_DTYPES = (torch.float16, torch.bfloat16)
 _TP_VOCAB_PARTITION_CACHE: dict[tuple[int, str, int, int, Optional[int], int], int] = {}
 
+# Ablation switch ``logp.lse_vocab_tile_reduction``: ``original`` is the
+# per-shard merge ``_merge_tp_local_logp`` already does.
+LseVocabTileReduction = Literal["fixed", "original"]
+
+
+def _check_lse_vocab_tile_reduction(reduction: LseVocabTileReduction) -> None:
+    if reduction == "fixed":
+        raise NotImplementedError(
+            "lse_vocab_tile_reduction='fixed': the deterministic implementation "
+            "is not introduced yet"
+        )
+
 
 def _env_flag(name: str) -> bool:
     return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"}
@@ -606,6 +618,7 @@ def __call__(
         tp_group: Any = None,
         vocab_start_index: int = 0,
         global_vocab_size: Optional[int] = None,
+        lse_vocab_tile_reduction: LseVocabTileReduction = "original",
     ) -> torch.Tensor:
         return self.apply(
             hidden,
@@ -615,6 +628,7 @@ def __call__(
             tp_group=tp_group,
             vocab_start_index=vocab_start_index,
             global_vocab_size=global_vocab_size,
+            lse_vocab_tile_reduction=lse_vocab_tile_reduction,
         )
 
     def apply(
@@ -627,8 +641,10 @@ def apply(
         tp_group: Any = None,
         vocab_start_index: int = 0,
         global_vocab_size: Optional[int] = None,
+        lse_vocab_tile_reduction: LseVocabTileReduction = "original",
     ) -> torch.Tensor:
         """Selected-token log-prob ``z[t] - logsumexp(z)``, returned in float32."""
+        _check_lse_vocab_tile_reduction(lse_vocab_tile_reduction)
         if hidden.shape[:-1] != target_ids.shape:
             raise ValueError(
                 f"hidden leading shape {tuple(hidden.shape[:-1])} must match "

rl_engine/kernels/ops/triton/loss/linear_logp.py

@@ -9,6 +9,8 @@
 import triton.language as tl
 
 from rl_engine.kernels.ops.pytorch.loss.linear_logp import (
+    LseVocabTileReduction,
+    _check_lse_vocab_tile_reduction,
     chunked_linear_logp_backward,
     should_use_tensor_parallel_linear_logp,
     tensor_parallel_linear_logp,
@@ -176,6 +178,7 @@ def __call__(
         tp_group: Any = None,
         vocab_start_index: int = 0,
         global_vocab_size: Optional[int] = None,
+        lse_vocab_tile_reduction: LseVocabTileReduction = "original",
     ) -> torch.Tensor:
         return self.apply(
             hidden,
@@ -185,6 +188,7 @@ def __call__(
             tp_group=tp_group,
             vocab_start_index=vocab_start_index,
             global_vocab_size=global_vocab_size,
+            lse_vocab_tile_reduction=lse_vocab_tile_reduction,
         )
 
     def apply(
@@ -197,7 +201,9 @@ def apply(
         tp_group: Any = None,
         vocab_start_index: int = 0,
         global_vocab_size: Optional[int] = None,
+        lse_vocab_tile_reduction: LseVocabTileReduction = "original",
     ) -> torch.Tensor:
+        _check_lse_vocab_tile_reduction(lse_vocab_tile_reduction)
         if hidden.device.type not in ("cuda", "xpu", "hip"):
             raise RuntimeError(
                 "TritonLinearLogpOp requires a GPU tensor (CUDA / ROCm / XPU), got "

tests/test_linear_logp.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2026 RL-Kernel Contributors
 
+import inspect
 import queue
 import tempfile
 import traceback
@@ -425,6 +426,32 @@ def test_native_rejects_shape_mismatch():
         native(hidden, weight, torch.zeros(_N + 1, dtype=torch.long), bias)
 
 
+def test_lse_vocab_tile_reduction_defaults_to_the_shard_order_merge():
+    """Every backend, so the ablation switch stays reachable whichever one dispatch picks."""
+    from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op
+
+    op_classes = [NativeLinearLogpOp, FusedLinearLogpSM90Op]
+    if _HAS_TRITON:
+        from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp
+
+        op_classes.append(TritonLinearLogpOp)
+
+    for op_class in op_classes:
+        for method in (op_class.__call__, op_class.apply):
+            parameter = inspect.signature(method).parameters["lse_vocab_tile_reduction"]
+            assert parameter.kind is inspect.Parameter.KEYWORD_ONLY
+            assert parameter.default == "original"
+
+
+def test_lse_vocab_tile_reduction_rejects_the_unimplemented_mode():
+    native = NativeLinearLogpOp()
+    hidden, weight, target, bias = _inputs(0, device="cpu")
+    assert native(hidden, weight, target, bias).shape == (_N,)
+
+    with pytest.raises(NotImplementedError, match="not introduced yet"):
+        native(hidden, weight, target, bias, lse_vocab_tile_reduction="fixed")
+
+
 def test_tensor_parallel_metadata_requires_multi_rank_group():
     native = NativeLinearLogpOp()
     hidden, weight, target, bias = _inputs(0, device="cpu")