diff --git a/rl_engine/kernels/ops/cuda/loss/linear_logp.py b/rl_engine/kernels/ops/cuda/loss/linear_logp.py index 313c2467..fae5dd97 100644 --- a/rl_engine/kernels/ops/cuda/loss/linear_logp.py +++ b/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 " diff --git a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py index 620b3c24..db06b831 100644 --- a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py +++ b/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 " diff --git a/rl_engine/kernels/ops/triton/loss/linear_logp.py b/rl_engine/kernels/ops/triton/loss/linear_logp.py index 561d81e4..35bbca7e 100644 --- a/rl_engine/kernels/ops/triton/loss/linear_logp.py +++ b/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 " diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py index 659127ce..2ae18c21 100644 --- a/tests/test_linear_logp.py +++ b/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")