feat(rl-kernel): add ROCm linear logp provider integration

tests/test_linear_logp_provider.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import sys
+import types
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vime.backends.megatron_utils.linear_logp_provider import (
+    LinearLogpContext,
+    LinearLogpProviderUnavailable,
+    LinearLogpRequest,
+    LinearLogpResult,
+    LinearProjection,
+    TokenLayout,
+    VocabPartition,
+    compute_linear_logp,
+)
+
+
+def _context() -> LinearLogpContext:
+    return LinearLogpContext(
+        hidden=torch.randn(3, 4),
+        projection=LinearProjection(weight=torch.randn(5, 4), bias=torch.randn(5)),
+        vocab_partition=VocabPartition(local_start=0, local_size=5, real_size=4, padded_size=5),
+    )
+
+
+def _request(*, with_entropy=False, context=None) -> LinearLogpRequest:
+    logits = torch.randn(3, 5, requires_grad=True)
+    return LinearLogpRequest(
+        logits=logits,
+        target_ids=torch.tensor([1, 2, 3]),
+        tensor_parallel_group=None,
+        token_layout=TokenLayout(world_size=2, rank=1, layout="zigzag"),
+        with_entropy=with_entropy,
+        with_entropy_grad=with_entropy,
+        chunk_size=64,
+        context=context,
+    )
+
+
+def _native(logits, *_args, with_entropy, **_kwargs):
+    return logits[:, :1], logits.sum(dim=-1) if with_entropy else None
+
+
+def _args(path=None, mode="auto"):
+    return SimpleNamespace(linear_logp_provider=path, linear_logp_provider_mode=mode)
+
+
+def _install_provider(monkeypatch, provider):
+    module = types.ModuleType("linear_logp_provider_fixture")
+    module.provider = provider
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    return f"{module.__name__}.provider"
+
+
+@pytest.mark.unit
+def test_unconfigured_provider_uses_native_path():
+    request = _request()
+
+    logp, entropy = compute_linear_logp(args=_args(), request=request, native=_native)
+
+    torch.testing.assert_close(logp, request.logits[:, :1])
+    assert entropy is None
+
+
+@pytest.mark.unit
+def test_provider_receives_structural_context(monkeypatch):
+    context = _context()
+    request = _request(with_entropy=True, context=context)
+    observed = {}
+
+    def provider(actual_request):
+        observed["context"] = actual_request.context
+        return LinearLogpResult(
+            logp=actual_request.logits[:, :1],
+            entropy=actual_request.logits.sum(dim=-1),
+            backend_id="fixture",
+            contract_id="fixture.v1",
+        )
+
+    path = _install_provider(monkeypatch, provider)
+    logp, entropy = compute_linear_logp(args=_args(path, "strict"), request=request, native=_native)
+
+    assert observed["context"] is context
+    torch.testing.assert_close(logp, request.logits[:, :1])
+    torch.testing.assert_close(entropy, request.logits.sum(dim=-1))
+
+
+@pytest.mark.unit
+def test_auto_falls_back_only_for_explicit_unavailability(monkeypatch):
+    request = _request()
+
+    def unavailable(_request):
+        raise LinearLogpProviderUnavailable("unsupported")
+
+    path = _install_provider(monkeypatch, unavailable)
+    logp, _ = compute_linear_logp(args=_args(path), request=request, native=_native)
+    torch.testing.assert_close(logp, request.logits[:, :1])
+
+    def broken(_request):
+        raise RuntimeError("provider bug")
+
+    path = _install_provider(monkeypatch, broken)
+    with pytest.raises(RuntimeError, match="provider bug"):
+        compute_linear_logp(args=_args(path), request=request, native=_native)
+
+
+@pytest.mark.unit
+def test_strict_mode_fails_when_provider_is_unavailable(monkeypatch):
+    def provider(_request):
+        raise LinearLogpProviderUnavailable("unsupported")
+
+    path = _install_provider(monkeypatch, provider)
+    with pytest.raises(RuntimeError, match="is unavailable"):
+        compute_linear_logp(args=_args(path, "strict"), request=_request(), native=_native)
+
+
+@pytest.mark.unit
+def test_strict_mode_validates_identity_and_autograd(monkeypatch):
+    request = _request()
+
+    def provider(actual_request):
+        return {
+            "logp": actual_request.logits[:, :1].detach(),
+            "entropy": None,
+            "backend_id": "fixture",
+            "contract_id": "fixture.v1",
+            "provenance": {},
+        }
+
+    path = _install_provider(monkeypatch, provider)
+    with pytest.raises(ValueError, match="detached"):
+        compute_linear_logp(args=_args(path, "strict"), request=request, native=_native)
+
+
+@pytest.mark.unit
+def test_structural_context_rejects_misaligned_projection():
+    with pytest.raises(ValueError, match="projection and vocabulary"):
+        LinearLogpContext(
+            hidden=torch.randn(3, 4),
+            projection=LinearProjection(weight=torch.randn(5, 4)),
+            vocab_partition=VocabPartition(local_start=0, local_size=4, real_size=5, padded_size=5),
+        )

tests/test_logprob_response_spans.py

@@ -3,10 +3,14 @@
 import _cp_dist_helpers  # noqa: F401
 import pytest
 import torch
-
 from megatron.core import mpu
-from vime.backends.megatron_utils.loss import _build_topp_keep_mask, get_rollout_top_p_logprob_kwargs
 
+from vime.backends.megatron_utils import loss as loss_module
+from vime.backends.megatron_utils.loss import (
+    _build_topp_keep_mask,
+    get_log_probs_and_entropy,
+    get_rollout_top_p_logprob_kwargs,
+)
 
 NUM_GPUS = 0
 
@@ -95,5 +99,52 @@ def test_top_p_mask_aligns_with_cp1_response_rows(monkeypatch):
     assert masked_rows == {2: [13], 3: [14], 5: [21], 6: [22], 7: [23]}
 
 
+@pytest.mark.unit
+def test_provider_receives_unscaled_logits_and_native_fallback_scales_once(monkeypatch):
+    _set_cp(monkeypatch, size=1, rank=0)
+    monkeypatch.setattr(mpu, "get_tensor_model_parallel_group", lambda: None, raising=False)
+    monkeypatch.setattr(mpu, "get_tensor_model_parallel_world_size", lambda: 1, raising=False)
+    observed = {}
+
+    def calculate(logits, *_args, with_entropy, **_kwargs):
+        observed["native_logits"] = logits.clone()
+        return logits[:, :1], logits.sum(dim=-1) if with_entropy else None
+
+    def dispatch(*, request, native, **_kwargs):
+        observed["request_logits"] = request.logits.clone()
+        observed["temperature"] = request.temperature
+        return native(
+            request.logits,
+            request.target_ids,
+            request.tensor_parallel_group,
+            with_entropy=request.with_entropy,
+            with_entropy_grad=request.with_entropy_grad,
+            chunk_size=request.chunk_size,
+            log_prob_keep_mask=request.log_prob_keep_mask,
+        )
+
+    monkeypatch.setattr(loss_module, "calculate_log_probs_and_entropy", calculate)
+    monkeypatch.setattr(loss_module, "compute_linear_logp", dispatch)
+    logits = torch.arange(12, dtype=torch.float32).reshape(1, 3, 4)
+    args = Namespace(
+        allgather_cp=False,
+        entropy_coef=0.0,
+        log_probs_chunk_size=-1,
+        rollout_temperature=0.5,
+    )
+
+    get_log_probs_and_entropy(
+        logits,
+        args=args,
+        unconcat_tokens=[torch.tensor([0, 1, 2])],
+        total_lengths=[3],
+        response_lengths=[2],
+    )
+
+    torch.testing.assert_close(observed["request_logits"], logits.squeeze(0))
+    torch.testing.assert_close(observed["native_logits"], logits.squeeze(0) / 0.5)
+    assert observed["temperature"] == 0.5
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__]))

tests/test_megatron_argument_validation.py

@@ -1,3 +1,4 @@
+import argparse
 import importlib.util
 import sys
 import types
@@ -167,6 +168,37 @@ def test_allgather_cp_ignores_cp_size_one(monkeypatch):
     module._validate_allgather_cp_supported(args)
 
 
+@pytest.mark.unit
+def test_strict_linear_logp_provider_requires_provider_path(monkeypatch):
+    module = load_arguments_module(monkeypatch)
+    args = types.SimpleNamespace(linear_logp_provider=None, linear_logp_provider_mode="strict")
+
+    with pytest.raises(ValueError, match="requires --linear-logp-provider"):
+        module._validate_linear_logp_provider_args(args)
+
+
+@pytest.mark.unit
+def test_linear_logp_provider_arguments_are_registered(monkeypatch):
+    module = load_vime_arguments_module(monkeypatch)
+    module.RouterArgs = types.SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser)
+    parser = argparse.ArgumentParser()
+    module.get_vime_extra_args_provider()(parser)
+
+    args = parser.parse_args(
+        [
+            "--rollout-batch-size",
+            "1",
+            "--linear-logp-provider",
+            "package.module.provider",
+            "--linear-logp-provider-mode",
+            "strict",
+        ]
+    )
+
+    assert args.linear_logp_provider == "package.module.provider"
+    assert args.linear_logp_provider_mode == "strict"
+
+
 @pytest.mark.unit
 def test_update_weight_disk_dir_required_for_disk_transport(monkeypatch):
     module = load_vime_arguments_module(monkeypatch)

vime/backends/megatron_utils/arguments.py

@@ -73,10 +73,18 @@ def _is_moe_config(hf_config):
     )
 
 
+def _validate_linear_logp_provider_args(args):
+    provider = str(getattr(args, "linear_logp_provider", "") or "").strip()
+    mode = str(getattr(args, "linear_logp_provider_mode", "auto")).strip().lower()
+    if mode == "strict" and not provider:
+        raise ValueError("--linear-logp-provider-mode strict requires --linear-logp-provider")
+
+
 def validate_args(args):
     """Run megatron's own validate_args plus vime-specific megatron validations."""
 
     _megatron_validate_args(args)
+    _validate_linear_logp_provider_args(args)
 
     # always use varlen
     args.variable_seq_lengths = True