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

vime/backends/megatron_utils/linear_logp_provider.py

@@ -73,6 +73,8 @@ class LinearLogpContext:
     hidden: torch.Tensor
     projection: LinearProjection
     vocab_partition: VocabPartition
+    reuse_local_logits: bool = False
+    local_logits: torch.Tensor | None = None
 
     def __post_init__(self) -> None:
         if not isinstance(self.hidden, torch.Tensor) or self.hidden.ndim != 2:
@@ -83,6 +85,17 @@ def __post_init__(self) -> None:
             raise ValueError("linear_logp projection and vocabulary shard widths do not match")
         if self.hidden.device != self.projection.weight.device:
             raise ValueError("linear_logp context tensors must share a device")
+        if not isinstance(self.reuse_local_logits, bool):
+            raise TypeError("linear_logp reuse_local_logits must be a bool")
+        if self.local_logits is not None:
+            if self.local_logits.ndim != 2:
+                raise ValueError("linear_logp context.local_logits must be [T, V_local]")
+            if self.local_logits.size(0) != self.hidden.size(0):
+                raise ValueError("linear_logp context local-logits rows must match hidden rows")
+            if self.local_logits.size(1) != self.vocab_partition.local_size:
+                raise ValueError("linear_logp context local-logits width must match vocabulary shard")
+            if self.local_logits.device != self.hidden.device:
+                raise ValueError("linear_logp context local logits must share the hidden device")
 
 
 @dataclass(frozen=True)

vime/backends/megatron_utils/model.py

@@ -90,6 +90,7 @@ def capture(module, inputs, kwargs, *, context_owner=owner):
                     real_size=real_vocab_size,
                     padded_size=padded_vocab_size,
                 ),
+                reuse_local_logits=bool(getattr(module, "__rl_kernel_reusable_local_logits__", False)),
             )
 
         output_layer._vime_linear_logp_capture_handle = output_layer.register_forward_pre_hook(
@@ -101,6 +102,18 @@ def _take_linear_logp_context(model_chunk):
     owner = _unwrap_model_chunk(model_chunk)
     context = getattr(owner, "_vime_linear_logp_context", None)
     owner._vime_linear_logp_context = None
+    output_layer = getattr(owner, "output_layer", None)
+    local_logits = getattr(output_layer, "_rl_kernel_local_logits", None)
+    if output_layer is not None:
+        output_layer._rl_kernel_local_logits = None
+    if context is not None and context.reuse_local_logits:
+        if not isinstance(local_logits, torch.Tensor):
+            raise RuntimeError("strict reusable LM head did not publish local logits")
+        if local_logits.ndim == 3:
+            local_logits = local_logits.transpose(0, 1).contiguous().reshape(-1, local_logits.size(-1))
+        elif local_logits.ndim != 2:
+            raise RuntimeError(f"unsupported strict reusable LM-head logits shape: {tuple(local_logits.shape)}")
+        context = dataclasses.replace(context, local_logits=local_logits)
     return context