diff --git a/csrc/cuda/fused_linear_logp_sm90.cu b/csrc/cuda/fused_linear_logp_sm90.cu index 40f34bb2..ffd3048c 100644 --- a/csrc/cuda/fused_linear_logp_sm90.cu +++ b/csrc/cuda/fused_linear_logp_sm90.cu @@ -9,6 +9,8 @@ #include #include #include +#include +#include #include #include #include @@ -684,11 +686,17 @@ __global__ void linear_logp_probs_bf16_forward_kernel( __syncthreads(); const float row_max = row_max_shared; + // A top-p replay mask can leave an entire TP vocab shard empty for a + // particular row. Treat that shard as the additive logsumexp identity + // (LSE=-inf, selected-logit=0) instead of evaluating -inf - -inf below. + // At least one TP rank still owns a finite nucleus entry for every valid + // sampled row, so the global rank-ordered merge remains well-defined. + const bool empty_row = row_max == -CUDART_INF_F; float local_sum = 0.0f; for (int col = tid; col < V; col += blockDim.x) { const float val = __bfloat162float(logits[static_cast(row) * logits_stride0 + col]); - local_sum += __expf(val - row_max); + local_sum += empty_row ? 0.0f : __expf(val - row_max); } reduce[tid] = local_sum; __syncthreads(); @@ -702,25 +710,179 @@ __global__ void linear_logp_probs_bf16_forward_kernel( __syncthreads(); if (probs != nullptr) { - const float inv_sum = 1.0f / row_sum_shared; + const float inv_sum = empty_row ? 0.0f : 1.0f / row_sum_shared; for (int col = tid; col < V; col += blockDim.x) { const float val = __bfloat162float(logits[static_cast(row) * logits_stride0 + col]); probs[static_cast(row) * probs_stride0 + col] = - __float2bfloat16(__expf(val - row_max) * inv_sum); + __float2bfloat16(empty_row ? 0.0f : __expf(val - row_max) * inv_sum); } } if (tid == 0) { - const float lse = row_max + logf(row_sum_shared); + const float lse = empty_row ? -CUDART_INF_F : row_max + logf(row_sum_shared); + const float target_logit = empty_row ? 0.0f : target_logit_shared; if (out_logp != nullptr) - out_logp[row] = target_logit_shared - lse; + out_logp[row] = empty_row ? -CUDART_INF_F : target_logit - lse; if (out_target_logit != nullptr) - out_target_logit[row] = target_logit_shared; + out_target_logit[row] = target_logit; if (out_lse != nullptr) out_lse[row] = lse; } } +// Top-p replay normally contains only a handful of tokens, but materializing a +// [N, V_local] boolean mask and scanning the complete vocabulary on every +// decode step makes that sparse case unnecessarily expensive. This kernel +// consumes the compact replay set directly. The ids are sorted in shared +// memory and assigned to the same col % 256 lanes as the dense kernel above, +// so each lane visits its finite logits in the same ascending-column order and +// the final 256-lane reduction remains bitwise identical to the dense masked +// path. +__global__ void linear_logp_top_p_local_bf16_forward_kernel( + const nv_bfloat16 *__restrict__ logits, + const int *__restrict__ target, + const int *__restrict__ replay_ids, + const float *__restrict__ replay_logprobs, + const float *__restrict__ temperature, + float *__restrict__ out_target_logit, + float *__restrict__ out_lse, + int N, + int V, + int K, + int64_t logits_stride0, + int64_t replay_ids_stride0, + int64_t replay_logprobs_stride0, + int vocab_start_index) { + constexpr int THREADS = 256; + constexpr int SORT_SIZE = 64; + __shared__ int sorted_cols[SORT_SIZE]; + __shared__ float reduce[THREADS]; + __shared__ float row_max_shared; + __shared__ float row_sum_shared; + __shared__ float target_logit_shared; + + const int row = blockIdx.x; + const int tid = threadIdx.x; + if (row >= N) + return; + + if (tid < SORT_SIZE) { + int col = INT_MAX; + if (tid < K) { + const int64_t replay_offset = + static_cast(row) * replay_ids_stride0 + tid; + const int64_t value_offset = + static_cast(row) * replay_logprobs_stride0 + tid; + const int global_id = replay_ids[replay_offset]; + if (isfinite(replay_logprobs[value_offset]) && + global_id >= vocab_start_index && + global_id < vocab_start_index + V) { + col = global_id - vocab_start_index; + } + } + sorted_cols[tid] = col; + } + __syncthreads(); + + // Fixed-size bitonic sort. Invalid entries are INT_MAX and naturally move + // to the end. Duplicate ids (the sampled token is also present in top-k) + // become adjacent and are consumed once below, matching boolean-mask + // semantics. + for (int size = 2; size <= SORT_SIZE; size <<= 1) { + for (int stride = size >> 1; stride > 0; stride >>= 1) { + if (tid < SORT_SIZE) { + const int peer = tid ^ stride; + if (peer > tid) { + const int a = sorted_cols[tid]; + const int b = sorted_cols[peer]; + const bool ascending = (tid & size) == 0; + if ((ascending && a > b) || (!ascending && a < b)) { + sorted_cols[tid] = b; + sorted_cols[peer] = a; + } + } + } + __syncthreads(); + } + } + + const int tgt = target[row] - vocab_start_index; + const float temp = temperature[row]; + float local_max = -CUDART_INF_F; + float local_target = 0.0f; + for (int index = 0; index < SORT_SIZE; ++index) { + const int col = sorted_cols[index]; + if (col == INT_MAX) + break; + if (index > 0 && col == sorted_cols[index - 1]) + continue; + if ((col & (THREADS - 1)) != tid) + continue; + const float raw = __bfloat162float( + logits[static_cast(row) * logits_stride0 + col]); + const float val = __bfloat162float(__float2bfloat16(raw / temp)); + local_max = fmaxf(local_max, val); + if (col == tgt) + local_target = val; + } + reduce[tid] = local_max; + __syncthreads(); + for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) { + if (tid < offset) + reduce[tid] = fmaxf(reduce[tid], reduce[tid + offset]); + __syncthreads(); + } + if (tid == 0) { + row_max_shared = reduce[0]; + target_logit_shared = 0.0f; + } + __syncthreads(); + if (local_target != 0.0f || + (tgt >= 0 && tgt < V && (tgt & (THREADS - 1)) == tid)) { + // Only the target-owning lane writes. The second condition preserves + // a legitimate target logit of exactly zero. + for (int index = 0; index < SORT_SIZE; ++index) { + const int col = sorted_cols[index]; + if (col == INT_MAX) + break; + if (col == tgt) { + target_logit_shared = local_target; + break; + } + } + } + __syncthreads(); + + const float row_max = row_max_shared; + const bool empty_row = row_max == -CUDART_INF_F; + float local_sum = 0.0f; + for (int index = 0; index < SORT_SIZE; ++index) { + const int col = sorted_cols[index]; + if (col == INT_MAX) + break; + if (index > 0 && col == sorted_cols[index - 1]) + continue; + if ((col & (THREADS - 1)) != tid) + continue; + const float raw = __bfloat162float( + logits[static_cast(row) * logits_stride0 + col]); + const float val = __bfloat162float(__float2bfloat16(raw / temp)); + local_sum += empty_row ? 0.0f : __expf(val - row_max); + } + reduce[tid] = local_sum; + __syncthreads(); + for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) { + if (tid < offset) + reduce[tid] += reduce[tid + offset]; + __syncthreads(); + } + if (tid == 0) { + row_sum_shared = reduce[0]; + out_lse[row] = empty_row ? -CUDART_INF_F : row_max + logf(row_sum_shared); + out_target_logit[row] = empty_row ? 0.0f : target_logit_shared; + } +} + __global__ void linear_logp_probs_bf16_to_dlogits_kernel(nv_bfloat16 *__restrict__ probs, const int *__restrict__ target, const float *__restrict__ grad_logp, @@ -1512,6 +1674,54 @@ std::vector linear_logp_local_bf16_forward_impl( return {local_target_logit, local_lse}; } +std::vector linear_logp_top_p_local_bf16_forward_impl( + torch::Tensor logits, + torch::Tensor target, + torch::Tensor replay_ids, + torch::Tensor replay_logprobs, + torch::Tensor temperature, + int64_t vocab_start_index) { + TORCH_CHECK(logits.is_cuda() && target.is_cuda() && replay_ids.is_cuda() && + replay_logprobs.is_cuda() && temperature.is_cuda(), + "top-p logits, target, replay tensors, and temperature must be CUDA tensors"); + TORCH_CHECK(logits.scalar_type() == at::kBFloat16, + "top-p local logp requires bf16 logits"); + TORCH_CHECK(replay_ids.scalar_type() == at::kInt, + "top-p replay ids must be int32"); + TORCH_CHECK(replay_logprobs.scalar_type() == at::kFloat && + temperature.scalar_type() == at::kFloat, + "top-p replay logprobs and temperature must be fp32"); + TORCH_CHECK(logits.dim() == 2 && replay_ids.dim() == 2 && replay_logprobs.dim() == 2, + "top-p local logp expects 2-D logits and replay tensors"); + const int N = logits.size(0); + const int V = logits.size(1); + const int K = replay_ids.size(1); + TORCH_CHECK(target.numel() == N && replay_ids.size(0) == N && + replay_logprobs.sizes() == replay_ids.sizes(), + "top-p replay rows must align with logits and target"); + TORCH_CHECK(K > 0 && K <= 64, "top-p replay width must be in [1, 64]"); + TORCH_CHECK(temperature.numel() == N, "temperature must have one value per row"); + TORCH_CHECK(logits.is_contiguous() && target.is_contiguous() && + replay_ids.is_contiguous() && replay_logprobs.is_contiguous() && + temperature.is_contiguous(), + "top-p local logp inputs must be contiguous"); + + c10::cuda::CUDAGuard device_guard(logits.device()); + auto target_i = target.reshape({N}).to(torch::kInt32).contiguous(); + auto opts_f = logits.options().dtype(torch::kFloat); + auto local_target_logit = torch::empty({N}, opts_f); + auto local_lse = torch::empty({N}, opts_f); + linear_logp_top_p_local_bf16_forward_kernel<<< + N, 256, 0, at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(logits.data_ptr()), + target_i.data_ptr(), replay_ids.data_ptr(), + replay_logprobs.data_ptr(), temperature.data_ptr(), + local_target_logit.data_ptr(), local_lse.data_ptr(), N, V, K, + logits.stride(0), replay_ids.stride(0), replay_logprobs.stride(0), + static_cast(vocab_start_index)); + return {local_target_logit, local_lse}; +} + torch::Tensor linear_logp_probs_bf16_to_dlogits_impl(torch::Tensor probs, torch::Tensor target, torch::Tensor grad_logp, @@ -1789,6 +1999,17 @@ std::vector linear_logp_local_bf16_forward(torch::Tensor logits, return linear_logp_local_bf16_forward_impl(logits, target, vocab_start_index); } +std::vector linear_logp_top_p_local_bf16_forward( + torch::Tensor logits, + torch::Tensor target, + torch::Tensor replay_ids, + torch::Tensor replay_logprobs, + torch::Tensor temperature, + int64_t vocab_start_index) { + return linear_logp_top_p_local_bf16_forward_impl( + logits, target, replay_ids, replay_logprobs, temperature, vocab_start_index); +} + torch::Tensor linear_logp_probs_bf16_to_dlogits_(torch::Tensor probs, torch::Tensor target, torch::Tensor grad_logp, diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 026ea23d..5ec80e38 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -50,6 +50,13 @@ std::vector linear_logp_local_probs_bf16_forward(torch::Tensor lo std::vector linear_logp_local_bf16_forward(torch::Tensor logits, torch::Tensor target, int64_t vocab_start_index); +std::vector linear_logp_top_p_local_bf16_forward( + torch::Tensor logits, + torch::Tensor target, + torch::Tensor replay_ids, + torch::Tensor replay_logprobs, + torch::Tensor temperature, + int64_t vocab_start_index); torch::Tensor linear_logp_probs_bf16_to_dlogits_(torch::Tensor probs, torch::Tensor target, torch::Tensor grad_logp, @@ -503,6 +510,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Build local bf16 softmax probabilities, target logits, and lse from bf16 logits"); m.def("linear_logp_local_bf16_forward", &linear_logp_local_bf16_forward, "Build local target logits and lse from bf16 logits without saving probabilities"); + m.def("linear_logp_top_p_local_bf16_forward", + &linear_logp_top_p_local_bf16_forward, + "Build local target logits and lse from a compact top-p replay set"); m.def("linear_logp_probs_bf16_to_dlogits_", &linear_logp_probs_bf16_to_dlogits_, "In-place bf16 probs -> dlogits for selected log-prob backward"); m.def("linear_logp_local_probs_bf16_to_dlogits_", diff --git a/rl_engine/integrations/framework_operators.py b/rl_engine/integrations/framework_operators.py index 08b361c6..a90356d9 100644 --- a/rl_engine/integrations/framework_operators.py +++ b/rl_engine/integrations/framework_operators.py @@ -2360,28 +2360,65 @@ def __call__( f"{context.hidden.size(0)} != {token_ids.numel()}" ) assert self._linear_logp is not None - selected = self._linear_logp.from_local_logits( - local_logits, - token_ids, - tp_group=context.tp_group, - vocab_start_index=context.vocab_start_index, - global_vocab_size=context.global_vocab_size, - real_vocab_size=context.real_vocab_size, - temperature=float(os.getenv("RL_KERNEL_VLLM_TEMPERATURE", "1.0")), - target="rollout", - diagnostics_hidden=context.hidden, - diagnostics_lm_head_weight=context.lm_head_weight, - ) + top_p_replay = False + if getattr(sampling_metadata, "top_p", None) is not None: + # Native vLLM has already applied temperature/top-p before it + # builds LogprobsTensors. Reconstruct the exact finite nucleus + # on each TP shard so the strict replacement keeps processed + # logprob semantics instead of accidentally substituting a + # full-vocabulary selected logprob. + replay_ids = logprobs_tensors.logprob_token_ids + replay_values = logprobs_tensors.logprobs + if replay_ids.shape != replay_values.shape: + raise RuntimeError( + "vLLM top-p replay ids and logprobs must have matching shapes" + ) + if replay_ids.size(0) != local_logits.size(0): + raise RuntimeError( + "vLLM top-p replay rows are not aligned with strict local logits" + ) + top_p_replay = True + if top_p_replay: + selected = self._linear_logp.from_local_logits_top_p( + local_logits, + token_ids, + replay_ids, + replay_values, + tp_group=context.tp_group, + vocab_start_index=context.vocab_start_index, + global_vocab_size=context.global_vocab_size, + real_vocab_size=context.real_vocab_size, + temperature=float(os.getenv("RL_KERNEL_VLLM_TEMPERATURE", "1.0")), + target="rollout", + ) + else: + selected = self._linear_logp.from_local_logits( + local_logits, + token_ids, + tp_group=context.tp_group, + vocab_start_index=context.vocab_start_index, + global_vocab_size=context.global_vocab_size, + real_vocab_size=context.real_vocab_size, + temperature=float(os.getenv("RL_KERNEL_VLLM_TEMPERATURE", "1.0")), + target="rollout", + diagnostics_hidden=context.hidden, + diagnostics_lm_head_weight=context.lm_head_weight, + ) strict_provenance = self._linear_logp.provenance - expected_entrypoint = ( + expected_entrypoints = { "rocm_vocab_parallel_logp_from_local_logits_tp" if torch.version.hip is not None else "sm90_deterministic_logp_from_local_logits_tp" - ) + } + if top_p_replay and torch.version.hip is None: + expected_entrypoints.add( + "sm90_deterministic_top_p_logp_from_local_logits_tp" + ) if ( strict_provenance.get("deterministic_linear_logp") is not True or strict_provenance.get("actual_backend") != self._linear_logp.backend_id - or strict_provenance.get("strict_entrypoint") != expected_entrypoint + or strict_provenance.get("strict_entrypoint") + not in expected_entrypoints ): raise RuntimeError( "strict vLLM rollout linear_logp did not execute the deterministic " @@ -2398,6 +2435,7 @@ def __call__( "logits_materialized": True, "padded_lm_head_alignment": True, "duplicate_lm_head_gemm": False, + "top_p_replay": top_p_replay, }, "source_logits_shape": list(source_logits.shape), "source_logits_dtype": _dtype_name(source_logits), diff --git a/rl_engine/integrations/linear_logp.py b/rl_engine/integrations/linear_logp.py index 60887442..2ace263e 100644 --- a/rl_engine/integrations/linear_logp.py +++ b/rl_engine/integrations/linear_logp.py @@ -630,6 +630,79 @@ def from_local_logits( } return (result, lse) if return_lse else result + def from_local_logits_top_p( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + replay_ids: torch.Tensor, + replay_logprobs: torch.Tensor, + *, + tp_group: Any, + vocab_start_index: int, + global_vocab_size: int, + real_vocab_size: int, + target: str = "rollout", + temperature: float | torch.Tensor | None = None, + ) -> torch.Tensor: + """Score compact rollout top-p support without a full-vocab replay mask.""" + + if local_logits.ndim != 2 or local_logits.dtype != torch.bfloat16: + raise TypeError("strict reused LM-head logits must be 2-D bfloat16") + if not local_logits.is_cuda or torch.version.hip is not None: + raise RuntimeError("strict reused LM-head logits require NVIDIA CUDA") + rank, world = self._tp_coordinates(tp_group) + if tp_group is None or world <= 1: + raise ValueError("reused rollout LM-head logits require a multi-rank TP group") + local_vocab = int(local_logits.size(1)) + requested_global = int(global_vocab_size) + if requested_global != local_vocab * world: + raise ValueError("reused rollout LM-head logits must use equal TP shards") + if int(vocab_start_index) != rank * local_vocab: + raise ValueError("reused rollout LM-head logits use a wrong TP shard offset") + real = int(real_vocab_size) + if not 0 < real <= requested_global: + raise ValueError("real_vocab_size must be within the padded global vocabulary") + self._validate_targets(target_ids, rows=local_logits.size(0), real_vocab_size=real) + temperature_tensor = self._temperature_tensor( + temperature, rows=local_logits.size(0), device=local_logits.device + ) + from rl_engine.kernels.ops.cuda.loss.linear_logp import ( + sm90_deterministic_top_p_logp_from_local_logits_tp, + ) + + result, _lse = sm90_deterministic_top_p_logp_from_local_logits_tp( + local_logits.contiguous(), + target_ids, + replay_ids, + replay_logprobs, + tp_group=tp_group, + vocab_start_index=int(vocab_start_index), + global_vocab_size=requested_global, + real_vocab_size=real, + temperature=temperature_tensor, + ) + self._last_provenance = { + **self._mismatch_provenance(), + "target": target, + "runtime_platform": "cuda", + "triton_used": False, + "actual_backend": self.backend_id, + "deterministic_linear_logp": True, + "strict_entrypoint": "sm90_deterministic_top_p_logp_from_local_logits_tp", + "local_logits_shape": list(local_logits.shape), + "replay_shape": list(replay_ids.shape), + "tp_group_present": True, + "vocab_start_index": int(vocab_start_index), + "global_vocab_size": requested_global, + "real_vocab_size": real, + "temperature": None if temperature is None else "provided", + "contract_version": "cuda-det-gemm-linear-logp-sm90-top-p-sparse-v1", + "logits_materialized": True, + "lm_head_result_reused": True, + "top_p_replay": True, + } + return result + @staticmethod def _mismatch_provenance() -> dict[str, Any]: case_id = os.getenv("RL_KERNEL_LOGP_CASE", "P/P").strip().upper() diff --git a/rl_engine/integrations/vime/linear_logp_provider.py b/rl_engine/integrations/vime/linear_logp_provider.py index 9a181dff..3abbc29b 100644 --- a/rl_engine/integrations/vime/linear_logp_provider.py +++ b/rl_engine/integrations/vime/linear_logp_provider.py @@ -283,8 +283,6 @@ def _provider_impl(request: Any, *, linear_logp: Any = None) -> LinearLogpResult partition = getattr(context, "vocab_partition", None) if strict and not isinstance(hidden, torch.Tensor): raise RuntimeError("strict Vime linear_logp request is missing structural context") - if strict and getattr(request, "log_prob_keep_mask", None) is not None: - raise RuntimeError("strict Vime linear_logp does not support top-p replay in this contract") if linear_logp is None and strict and isinstance(hidden, torch.Tensor): linear_logp = _default_strict_linear_logp() if linear_logp is not None and isinstance(hidden, torch.Tensor): @@ -313,6 +311,11 @@ def _provider_impl(request: Any, *, linear_logp: Any = None) -> LinearLogpResult ) if materialized_local_logits: reuse_local_logits = True + keep_mask = getattr(request, "log_prob_keep_mask", None) + if keep_mask is not None and not reuse_local_logits: + raise RuntimeError( + "strict top-p replay requires reusable materialized local logits" + ) with_entropy = bool(getattr(request, "with_entropy", False)) with_entropy_grad = bool(getattr(request, "with_entropy_grad", False)) local_logits_temperature = _local_logits_temperature(request) @@ -321,6 +324,7 @@ def _provider_impl(request: Any, *, linear_logp: Any = None) -> LinearLogpResult and with_entropy and not with_entropy_grad and _is_identity_temperature(local_logits_temperature) + and keep_mask is None ) strict_lse = None if reuse_local_logits: @@ -329,6 +333,10 @@ def _provider_impl(request: Any, *, linear_logp: Any = None) -> LinearLogpResult local_logits = request_logits if not isinstance(local_logits, torch.Tensor): raise RuntimeError("strict reusable LM-head context is missing local logits") + if keep_mask is not None: + if keep_mask.shape != local_logits.shape: + raise RuntimeError("strict top-p replay mask must match local logits") + local_logits = local_logits.masked_fill(~keep_mask, float("-inf")) from_local_logits = getattr(linear_logp, "from_local_logits", None) if not callable(from_local_logits): raise RuntimeError( diff --git a/rl_engine/kernels/ops/cuda/loss/linear_logp.py b/rl_engine/kernels/ops/cuda/loss/linear_logp.py index e7e4a5b0..275eba68 100644 --- a/rl_engine/kernels/ops/cuda/loss/linear_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/linear_logp.py @@ -1287,6 +1287,81 @@ def sm90_deterministic_logp_from_local_logits_tp( return logp.reshape(lead_shape), lse.reshape(lead_shape) +def sm90_deterministic_top_p_logp_from_local_logits_tp( + local_logits: torch.Tensor, + target_ids: torch.Tensor, + replay_ids: torch.Tensor, + replay_logprobs: torch.Tensor, + *, + vocab_start_index: int, + global_vocab_size: int, + real_vocab_size: int = -1, + temperature: Optional[torch.Tensor] = None, + tp_group: Any, +) -> tuple[torch.Tensor, torch.Tensor]: + """Score a compact top-p replay set with the dense kernel's reduction order. + + This is deliberately inference-only: training retains the dense masked + logits so its existing autograd contract is unchanged. + """ + + required = "linear_logp_top_p_local_bf16_forward" + if not (_EXT_AVAILABLE and hasattr(_C, required)): + raise RuntimeError(f"strict TP top-p logits logp missing CUDA symbol: {required}") + if torch.is_grad_enabled() and local_logits.requires_grad: + raise RuntimeError("compact top-p local-logits logp is inference-only") + if local_logits.ndim != 2 or local_logits.dtype != torch.bfloat16: + raise TypeError("strict TP top-p logits logp requires 2-D BF16 local logits") + if not local_logits.is_cuda or not local_logits.is_contiguous(): + raise ValueError("strict TP top-p logits logp requires contiguous CUDA logits") + if replay_ids.ndim != 2 or replay_logprobs.shape != replay_ids.shape: + raise ValueError("top-p replay ids/logprobs must be aligned 2-D tensors") + if replay_ids.size(0) != local_logits.size(0): + raise ValueError("top-p replay rows must align with local logits") + if replay_ids.size(1) > 64: + raise ValueError("top-p replay width exceeds the compact kernel limit of 64") + + vocab_start = int(vocab_start_index) + global_vocab = _validate_even_tp_vocab_partition_local( + tp_group=tp_group, + vocab_start_index=vocab_start, + local_vocab_size=local_logits.size(1), + global_vocab_size=int(global_vocab_size), + ) + real_vocab = global_vocab if int(real_vocab_size) < 0 else int(real_vocab_size) + if not 0 < real_vocab <= global_vocab: + raise ValueError(f"invalid real_vocab_size={real_vocab} for padded vocab={global_vocab}") + target = target_ids.reshape(-1).to(device=local_logits.device, dtype=torch.long).contiguous() + if target.numel() != local_logits.size(0): + raise ValueError("target_ids must contain one id per local-logits row") + _assert_global_targets_async(target, real_vocab) + + if temperature is None: + temp = torch.ones( + local_logits.size(0), device=local_logits.device, dtype=torch.float32 + ) + else: + temp = temperature.to(device=local_logits.device, dtype=torch.float32).reshape(-1) + if temp.numel() == 1: + temp = temp.expand(local_logits.size(0)).contiguous() + if temp.numel() != local_logits.size(0): + raise ValueError("temperature must be positive and scalar or per-token") + torch._assert_async((temp > 0).all(), "temperature must be positive") + + ids = replay_ids.to(device=local_logits.device, dtype=torch.int32).contiguous() + values = replay_logprobs.to(device=local_logits.device, dtype=torch.float32).contiguous() + local_target, local_lse = _C.linear_logp_top_p_local_bf16_forward( + local_logits, + target, + ids, + values, + temp, + vocab_start, + ) + logp, lse = _merge_tp_local_logp(local_lse, local_target, tp_group=tp_group) + return logp.reshape(target_ids.shape), lse.reshape(target_ids.shape) + + class _StrictTensorParallelLinearLogpAutograd(torch.autograd.Function): @staticmethod def forward( diff --git a/tests/test_vime_linear_logp_provider.py b/tests/test_vime_linear_logp_provider.py index c97dadd1..656b6c4a 100644 --- a/tests/test_vime_linear_logp_provider.py +++ b/tests/test_vime_linear_logp_provider.py @@ -254,11 +254,39 @@ def provider(actual_request, *, linear_logp): assert result.logp.shape == (3, 1) -def test_provider_rejects_top_p_replay_without_changing_its_semantics(): - request = _request(keep_mask=torch.ones((3, 8), dtype=torch.bool)) +def test_provider_replays_top_p_mask_on_reused_local_logits(monkeypatch): + monkeypatch.setenv("VIME_RL_KERNEL_STRICT", "1") - with pytest.raises(LinearLogpProviderUnavailable, match="top-p replay"): - provider(request) + class FakeLinearLogp: + backend_id = "fake-linear-logp" + provenance = {"actual_backend": "fake-linear-logp"} + + def from_local_logits(self, local_logits, target_ids, **_kwargs): + return torch.log_softmax(local_logits[:, :7], dim=-1)[ + torch.arange(target_ids.size(0)), target_ids + ] + + import rl_engine.integrations.vime.linear_logp_provider as provider_module + + monkeypatch.setattr(provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp()) + request = _structural_request() + request.context.reuse_local_logits = True + request.context.local_logits = request.logits + request.log_prob_keep_mask = torch.tensor( + [ + [True, False, True, False, False, False, False, False], + [False, True, False, False, False, True, False, False], + [True, True, False, False, False, False, False, False], + ] + ) + + result = provider(request) + masked = request.logits.masked_fill(~request.log_prob_keep_mask, float("-inf")) + expected = torch.log_softmax(masked[:, :7], dim=-1)[ + torch.arange(request.target_ids.size(0)), request.target_ids + ] + + torch.testing.assert_close(result.logp.squeeze(-1), expected) def test_provider_rejects_local_vocab_metadata_that_cannot_describe_tp_ownership():