Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions packages/prime-rl-configs/src/prime_rl/utils/parsers.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,8 @@
(re.compile(r"^poolside/Laguna"), "poolside_v1"),
(re.compile(r"^PrimeIntellect/INTELLECT-3"), "deepseek_r1"),
(re.compile(r"^nvidia/NVIDIA-Nemotron-3-Super"), "nemotron_v3"),
(re.compile(r"^nvidia/NVIDIA-Nemotron-3-Nano"), "nano_v3"),
(re.compile(r"^nvidia/NVIDIA-Nemotron-3\.5"), "nano_v3"),
(re.compile(r"^nvidia/NVIDIA-Nemotron-3-Nano"), "nemotron_v3"),
(re.compile(r"^nvidia/NVIDIA-Nemotron-3\.5"), "nemotron_v3"),
(re.compile(r"^stepfun-ai/Step-3\.5"), "step3p5"),
# Only Qwen3 Thinking models reason — Instruct/base models do not.
(re.compile(r"^Qwen/Qwen3-.*Thinking"), "deepseek_r1"),
Expand Down
28 changes: 16 additions & 12 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ dependencies = [
"jaxtyping>=0.3.2",
"loguru>=0.7.3",
"numpy>=2.2.6",
"transformers==5.6.2",
"transformers==5.19.0",
"wandb>=0.26.1",
"wandb-workspaces>=0.4.3",
"prime>=0.8.0",
Expand Down Expand Up @@ -55,7 +55,7 @@ gpu = [
"torchvision ; sys_platform == 'linux'",
"torchaudio ; sys_platform == 'linux'",
"torchdata>=0.11.0 ; sys_platform == 'linux'",
"vllm>=0.31.0 ; sys_platform == 'linux'",
"vllm>=0.32.0 ; sys_platform == 'linux'",
"mooncake-transfer-engine>=0.3.10.post2 ; sys_platform == 'linux'",
"ring-flash-attn>=0.1.8 ; sys_platform == 'linux'",
"torchtitan ; sys_platform == 'linux'",
Expand Down Expand Up @@ -149,10 +149,9 @@ override-dependencies = [
# Shared by vLLM, TileLang, and the pinned FlashAttention 4 runtime.
"apache-tvm-ffi==0.1.11",
"nvidia-cudnn-cu13>=9.15 ; sys_platform == 'linux'",
# vLLM 0.31.0 requires this exact CUTLASS DSL version and quack-kernels 0.6.5.
# vLLM 0.32.0 requires this exact CUTLASS DSL version and quack-kernels 0.6.5.
# Keep this pin aligned with quack-kernels and the flash-attn-4 revision below.
"nvidia-cutlass-dsl==4.7.1",
"transformers==5.6.2",
"nvidia-cutlass-dsl==4.8.0",
"torch>=2.13.0",
"openenv-core",
]
Expand All @@ -176,9 +175,11 @@ members = [
]
# Tau tasksets require litellm<1.82.7, while Harbor requires litellm>=1.92.0.
# Apex pins a verifiers git revision instead of the editable workspace member.
# ALE CLI requires numpy>=2.4.4, while vLLM 0.31.0 (via mistral-common) requires numpy<2.4.
# ALE CLI requires numpy>=2.4.4, while vLLM 0.32.0 (via mistral-common) requires numpy<2.4.
# IFBench requires setuptools<78, while vLLM 0.32.0 requires setuptools>=78.1.1.
# A workspace shares one lock, so these packages need separate environments.
exclude = [
"deps/prime-envs/environments/if/ifbench",
"deps/prime-envs/environments/office/ale_cli",
"deps/prime-envs/environments/tool_use/tau2_bench",
"deps/prime-envs/environments/tool_use/tau2_synth",
Expand Down Expand Up @@ -222,19 +223,22 @@ requires-dist = ["nixl-cu13==0.10.1"]
# so resolve that variant without installing both runtime packages.
[[tool.uv.dependency-metadata]]
name = "nvidia-cutlass-dsl"
version = "4.7.1"
version = "4.8.0"
requires-dist = [
"nvidia-cutlass-dsl-libs-base==4.7.1",
"nvidia-cutlass-dsl-libs-cu13==4.7.1",
"nvidia-cutlass-dsl-libs-base==4.8.0",
"nvidia-cutlass-dsl-libs-cu13==4.8.0",
]

[tool.uv.exclude-newer-package]
# verifiers pins harbor==0.24.0, published inside the rolling window;
# bounded cutoff mirroring the one in deps/verifiers/pyproject.toml
harbor = "2026-10-08T00:00:00Z"
tokenspeed-mla = false
# vLLM 0.31.0 pins flashinfer-python 0.7.0.post1, published inside the rolling window
# vLLM 0.32.0 pins flashinfer-python 0.7.0.post1, published inside the rolling window
flashinfer-python = false
# vLLM 0.32.0 pins these, published inside the rolling window
instanttensor = false
humming-kernels = false
fastokens = false
flash_attn_3 = false
# PrimeIntellect-published packages in this project's dependency closure —
Expand Down Expand Up @@ -284,8 +288,8 @@ vllm-router = [
{ url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.2.3/vllm_router-0.2.3-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
]
vllm = [
{ url = "https://github.com/vllm-project/vllm/releases/download/v0.31.0/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
{ url = "https://github.com/vllm-project/vllm/releases/download/v0.31.0/vllm-0.31.0-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
{ url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "platform_machine == 'x86_64'" },
{ url = "https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
]
deep-ep = [
{ url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
Expand Down
196 changes: 6 additions & 190 deletions src/prime_rl/inference/patches.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,16 +16,12 @@ def apply_shared_vllm_patches():

patch_gpt_oss_weight_loading()
patch_qwen38_weight_loading()
monkey_patch_shm_buffer_ownership()
_patch_lora_key_prefix()
_patch_qwen35_moe_lora_format()
monkey_patch_nano_v3_reasoning_parser()
monkey_patch_minimax_m2_think_end_passthrough()
monkey_patch_routed_experts_with_pd_connectors()
monkey_patch_kv_xfer_finished_tolerate_freed()
monkey_patch_online_fp8_parameter_cast()
monkey_patch_deepseek_v4_allowed_layer_types()
monkey_patch_tokenize_params_validation()
monkey_patch_strip_routed_experts_from_chat()
monkey_patch_dp_coordinator_startup_timeout()
monkey_patch_minimax_m2_for_lora()
Expand All @@ -37,65 +33,6 @@ def apply_shared_vllm_patches():
monkey_patch_deepseek_v4_attn_sink_loading()


def monkey_patch_shm_buffer_ownership():
"""Own out-of-band buffers before releasing a reusable vLLM SHM slot."""
import pickle

import zmq
from vllm.distributed.device_communicators import shm_broadcast

# Backport https://github.com/vllm-project/vllm/pull/61000; remove when our
# minimum vLLM version includes it. Retained routed-expert arrays otherwise
# alias ring slots that subsequent messages can overwrite.
MessageQueue = shm_broadcast.MessageQueue
if getattr(MessageQueue.dequeue, "_prime_rl_owns_shm_buffers", False):
return

def dequeue(self, timeout: float | None = None, indefinite: bool = False):
if self._is_local_reader:
with self.acquire_read(timeout, indefinite) as buf:
overflow = buf[0] == 1
if not overflow:
offset = 3
buf_count = shm_broadcast.from_bytes_big(buf[1:offset])
all_buffers = []
for i in range(buf_count):
buf_offset = offset + 4
buf_len = shm_broadcast.from_bytes_big(buf[offset:buf_offset])
offset = buf_offset + buf_len
buffer = buf[buf_offset:offset]
all_buffers.append(buffer if i == 0 else bytearray(buffer))
obj = pickle.loads(all_buffers[0], buffers=all_buffers[1:])
if overflow:
obj = MessageQueue.recv(self.local_socket, timeout)
elif self._is_remote_reader:
obj = MessageQueue.recv(self.remote_socket, timeout)
else:
raise RuntimeError("Only readers can dequeue")
return obj

def _rebuild_tensor(buf, shape: tuple[int, ...], dtype_str: str) -> torch.Tensor:
dtype = getattr(torch, dtype_str)
assert isinstance(dtype, torch.dtype)
if isinstance(buf, zmq.Frame):
try:
return torch.frombuffer(buf, dtype=torch.uint8).view(dtype).view(shape)
except ValueError:
pass
raw = buf if isinstance(buf, bytearray) else bytearray(buf)
if not raw:
assert 0 in shape
return torch.empty(shape, dtype=dtype)
return torch.frombuffer(raw, dtype=torch.uint8).view(dtype).view(shape)

# Pickle resolves this reducer by its module and name in receiving processes.
_rebuild_tensor.__module__ = shm_broadcast.__name__
_rebuild_tensor.__qualname__ = "_rebuild_tensor"
dequeue._prime_rl_owns_shm_buffers = True
shm_broadcast._rebuild_tensor = _rebuild_tensor
MessageQueue.dequeue = dequeue


@dataclass
class PackedSamplingMask:
"""A finished request's sampling masks in CSR form: ``ids[sum(counts[:i]):][:counts[i]]``
Expand Down Expand Up @@ -147,27 +84,10 @@ def _new_completion_output(self, token_ids, finish_reason, stop_reason):
RequestState._new_completion_output = _new_completion_output


def monkey_patch_deepseek_v4_allowed_layer_types():
"""Let vLLM's `DeepseekV4Config` construct under a transformers pin below 5.15.

`vllm/transformers_utils/configs/deepseek_v4.py` leaves `layer_types` to
`PretrainedConfig`, whose validator checks it against a vocabulary that only gained V4's
compressed attention names in transformers 5.15, so `AutoConfig.from_pretrained` on any
DeepSeek V4 checkpoint raises. vLLM declares `transformers>=5.5.3`, so this is a gap in
upstream's own version floor rather than something prime-rl introduced.

`EngineArgs.__post_init__` loads general plugins before `create_model_config`, which is why
patching here is early enough. The trainer applies the same shim on its own import path.
"""
from prime_rl.utils.transformers_compat import allow_deepseek_v4_layer_types

allow_deepseek_v4_layer_types()


def monkey_patch_deepseek_v4_c128_boundary():
"""Fix DeepSeek V4 long-context KV corruption under piecewise CUDA graphs.

On vLLM 0.30+, including the pinned 0.31.0, the piecewise CUDA graphs are captured on a step
On vLLM 0.30+, including the pinned 0.32.0, the piecewise CUDA graphs are captured on a step
that skips writing C128 compressed KV, so every replay skips it too. Mixed batches of up to
512 tokens never store their C128 entries, later queries attend over stale compressed KV, and
long-context rollouts drift from the trainer, inflating mismatch KL.
Expand Down Expand Up @@ -294,6 +214,7 @@ def monkey_patch_kv_xfer_finished_tolerate_freed():
"""
from vllm.logger import init_logger
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.metrics.stats import KV_FETCH_COMPLETED_WAITING
from vllm.v1.request import RequestStatus

logger = init_logger("vllm.v1.core.sched.scheduler")
Expand All @@ -314,6 +235,7 @@ def _update_from_kv_xfer_finished(self, kv_connector_output):
req = self.requests[req_id]
if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
self.finished_recving_kv_req_ids.add(req_id)
self._set_kv_fetch_stage(req, KV_FETCH_COMPLETED_WAITING)
else:
assert RequestStatus.is_finished(req.status)
self._free_blocks(self.requests[req_id])
Expand All @@ -330,23 +252,6 @@ def _update_from_kv_xfer_finished(self, kv_connector_output):
logger.warning("Patched Scheduler._update_from_kv_xfer_finished to tolerate freed (aborted) KV-transfer reqs.")


def monkey_patch_nano_v3_reasoning_parser():
from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager
from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser

class NanoV3ReasoningParser(DeepSeekR1ReasoningParser):
def extract_reasoning(self, model_output, request):
reasoning_content, final_content = super().extract_reasoning(model_output, request)
chat_template_kwargs = getattr(request, "chat_template_kwargs", None)

if chat_template_kwargs and chat_template_kwargs.get("enable_thinking") is False and final_content is None:
reasoning_content, final_content = final_content, reasoning_content

return reasoning_content, final_content

ReasoningParserManager.register_module("nano_v3", module=NanoV3ReasoningParser)


def monkey_patch_minimax_m2_think_end_passthrough():
"""Keep the literal ``</think>`` in MiniMax-M2 content on tool-calling turns.

Expand Down Expand Up @@ -384,7 +289,7 @@ def _patched_config():
def monkey_patch_routed_experts_with_pd_connectors():
"""Let vLLM's routed-experts capture (AuxOutput connector) run on NIXL P/D instances.

vLLM 0.31 rejects routed-experts capture with P/D connectors, because a decode instance never
vLLM (0.32 included) rejects routed-experts capture with P/D connectors, because a decode instance never
computes a remotely prefilled prompt, so it has no rows for it. This carries our upstream fix
(AuxOutput patch 02): a request on a decode instance emits rows only from its first forward on
this instance, and the worker does not buffer the
Expand Down Expand Up @@ -544,7 +449,7 @@ def _patch_qwen35_moe_lora_format():
def _patch_lora_key_prefix():
"""Accept both bare-suffix and fully-qualified expert module names in LoRA adapters.

Copy of vLLM 0.24.0's ``LoRAModel.from_local_checkpoint`` with one change: the
Copy of vLLM 0.32.0's ``LoRAModel.from_local_checkpoint`` with one change: the
``.experts`` branch of ``check_unexpected_modules`` accepts either the bare suffix
(``down_proj``) or the qualified per-expert name (``experts.N.down_proj``), where
upstream only accepts the qualified form. Our trainer's 2D per-expert adapters
Expand All @@ -556,7 +461,6 @@ def _patch_lora_key_prefix():
LoRAModel,
MoEEPLoadSpec,
PEFTHelper,
TensorizerConfig,
WeightsMapper,
_is_remote_expert_key,
get_lora_id,
Expand All @@ -577,7 +481,6 @@ def _patched_from_local_checkpoint(
dtype: torch.dtype | None = None,
model_vocab_size: int | None = None,
weights_mapper: WeightsMapper | None = None,
tensorizer_config_dict: dict | None = None,
skip_prefixes: list[str] | None = None,
moe_ep_spec: MoEEPLoadSpec | None = None,
) -> "LoRAModel":
Expand Down Expand Up @@ -644,24 +547,7 @@ def check_unexpected_modules(modules: dict):
f" Please verify that the loaded LoRA module is correct"
)

if tensorizer_config_dict:
from tensorizer import TensorDeserializer

tensorizer_config = TensorizerConfig(**tensorizer_config_dict)
tensorizer_dir = tensorizer_config.tensorizer_dir
if tensorizer_dir is None:
raise ValueError("tensorizer_dir must be set in tensorizer config.")
lora_tensor_path = os.path.join(tensorizer_dir, "adapter_model.tensors")
tensorizer_args = tensorizer_config._construct_tensorizer_args()
tensors = TensorDeserializer(
lora_tensor_path,
dtype=tensorizer_config.dtype,
device=device,
**tensorizer_args.deserialization_kwargs,
)
check_unexpected_modules(tensors)

elif os.path.isfile(lora_tensor_path):
if os.path.isfile(lora_tensor_path):
# Find unexpected modules.
# Use safetensor key as a source of truth to find expected modules.
# in peft if you have target_modules A, B, C and C does not exist
Expand Down Expand Up @@ -702,76 +588,6 @@ def check_unexpected_modules(modules: dict):
LoRAModel.from_local_checkpoint = classmethod(_patched_from_local_checkpoint)


# Monkeypatch TokenizeParams to fix overly conservative validation
def monkey_patch_tokenize_params_validation():
"""
Patch TokenizeParams validation to only reject requests where the prompt
itself exceeds max_model_len, not where prompt + max_tokens > max_model_len.

Original behavior:
- Rejects if prompt_len > (max_model_len - max_tokens)

Patched behavior:
- Only rejects if prompt_len > max_model_len
- Lets the engine naturally cap generation at max_model_len
"""
from vllm.exceptions import VLLMValidationError
from vllm.renderers.params import TokenizeParams

def _patched_token_len_check(self, tokenizer, tokens):
"""Only validate that prompt fits in max_model_len, not prompt+max_tokens"""
if self.max_total_tokens is not None and len(tokens) > self.max_total_tokens:
raise VLLMValidationError(
f"The prompt is {len(tokens)} tokens, which exceeds the "
f"model's maximum context length of {self.max_total_tokens} tokens. "
f"Please reduce the length of the input prompt.",
parameter="input_tokens",
value=len(tokens),
)
return tokens

def _patched_text_len_check(self, tokenizer, text):
"""Only validate text length against max_model_len, not max_input_tokens"""
if self.max_total_tokens is None or tokenizer is None:
return text

if self.truncate_prompt_tokens is None:
max_chars = self.max_total_tokens * tokenizer.max_chars_per_token
if len(text) > max_chars:
raise VLLMValidationError(
f"You passed {len(text)} input characters. "
f"However, the model's context length is only "
f"{self.max_total_tokens} tokens "
f"(at most {max_chars} characters). "
f"Please reduce the length of the input prompt.",
parameter="input_text",
value=len(text),
)
return text

def _patched_get_encode_kwargs(self):
"""Use max_total_tokens (max_model_len) instead of max_input_tokens for HF tokenizer truncation.

The original uses max_input_tokens (= max_model_len - max_tokens) + 1, which causes HuggingFace's
tokenizer.encode() to left-truncate prompts before _token_len_check even runs.
"""
max_length = self.truncate_prompt_tokens
if max_length is not None and max_length < 0:
max_length = self.max_total_tokens
elif max_length is None and self.max_total_tokens is not None:
max_length = self.max_total_tokens + 1

return dict(
truncation=max_length is not None,
max_length=max_length,
add_special_tokens=self.add_special_tokens,
)

TokenizeParams._token_len_check = _patched_token_len_check
TokenizeParams._text_len_check = _patched_text_len_check
TokenizeParams.get_encode_kwargs = _patched_get_encode_kwargs


def monkey_patch_minimax_m2_for_lora():
"""Patch vLLM's MiniMaxM2 model for LoRA compatibility.

Expand Down
Loading
Loading