From 2c94ad512b1d9117dcee86984904bd164f4acef2 Mon Sep 17 00:00:00 2001 From: Codex H100 Validation Date: Tue, 25 Aug 2026 02:47:53 +0000 Subject: [PATCH] feat(integrations): add H100 framework runtime adapters --- .../validate_artifacts.py | 211 +++ pyproject.toml | 3 + rl_engine/integrations/__init__.py | 8 + rl_engine/integrations/ablation.py | 44 +- rl_engine/integrations/framework_operators.py | 1297 +++++++++++++++++ rl_engine/integrations/linear_logp.py | 368 +++++ rl_engine/integrations/megatron_runtime.py | 141 ++ rl_engine/integrations/runtime.py | 136 +- rl_engine/integrations/state.py | 57 + rl_engine/integrations/vime/logp.py | 143 +- rl_engine/integrations/vllm_runtime.py | 398 +++++ .../ops/cuda/attention/strict_runtime.py | 297 ++++ .../kernels/ops/pytorch/attention/ablation.py | 32 +- tests/test_framework_runtime_adapters.py | 377 +++++ tests/test_vime_validation_artifacts.py | 77 + 15 files changed, 3577 insertions(+), 12 deletions(-) create mode 100644 examples/vime_qwen3_8b_tp2_cp2/validate_artifacts.py create mode 100644 rl_engine/integrations/framework_operators.py create mode 100644 rl_engine/integrations/linear_logp.py create mode 100644 rl_engine/integrations/megatron_runtime.py create mode 100644 rl_engine/integrations/state.py create mode 100644 rl_engine/integrations/vllm_runtime.py create mode 100644 rl_engine/kernels/ops/cuda/attention/strict_runtime.py create mode 100644 tests/test_framework_runtime_adapters.py create mode 100644 tests/test_vime_validation_artifacts.py diff --git a/examples/vime_qwen3_8b_tp2_cp2/validate_artifacts.py b/examples/vime_qwen3_8b_tp2_cp2/validate_artifacts.py new file mode 100644 index 00000000..a8c41446 --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/validate_artifacts.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Validate CUDA-only framework readbacks and Vime train/rollout Logp dumps.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any, Mapping + +import torch +from rl_engine.integrations.runtime import _contains_triton, _runtime_platform + +_FRAMEWORKS = (("megatron", "training"), ("vllm", "rollout")) +_MODULES = ("attention", "ffn", "logp") +_STRICT_LOGP_BACKEND = "rlkernel.linear_logp.bitwise.v1" +_BACKEND_PREFIXES = ("rlkernel.", "pytorch-vocab-parallel-logp") + + +def load_readbacks(directory: Path) -> list[dict[str, Any]]: + values: list[dict[str, Any]] = [] + for path in sorted(directory.glob("*.json")): + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"readback must contain an object: {path}") + value["_path"] = str(path) + values.append(value) + if not values: + raise ValueError(f"no framework readbacks found in {directory}") + return values + + +def validate_readbacks(readbacks: list[dict[str, Any]]) -> dict[str, Any]: + errors: list[str] = [] + frameworks: dict[str, Any] = {} + for framework, target in _FRAMEWORKS: + matching = [ + value + for value in readbacks + if value.get("framework") == framework and value.get("target") == target + ] + label = f"{framework}/{target}" + if not matching: + errors.append(f"missing {label} readback") + continue + module_summary: dict[str, Any] = {} + for value in matching: + if value.get("fallbacks"): + errors.append(f"{label} recorded fallback: {value['fallbacks']}") + for module in _MODULES: + hook_count = sum(module in value.get("installed_hooks", {}) for value in matching) + records = [ + value["operators"][module] + for value in matching + if isinstance(value.get("operators"), Mapping) and module in value["operators"] + ] + call_count = sum(int(record.get("call_count", 0)) for record in records) + if hook_count == 0: + errors.append(f"{label} {module} hook was not installed") + if call_count == 0: + errors.append(f"{label} {module} had zero calls") + backends = sorted({str(record.get("backend_id", "")) for record in records}) + for record in records: + backend = str(record.get("backend_id", "")) + if module == "logp" and backend != _STRICT_LOGP_BACKEND: + errors.append( + f"{label} logp used {backend!r}, expected {_STRICT_LOGP_BACKEND!r}" + ) + elif not backend.startswith(_BACKEND_PREFIXES): + errors.append(f"{label} {module} used unexpected backend {backend!r}") + if _contains_triton(record): + errors.append(f"{label} {module} used Triton") + if _runtime_platform(record.get("provenance")) != "cuda": + errors.append(f"{label} {module} did not report CUDA execution") + module_summary[module] = { + "installed_processes": hook_count, + "call_count": call_count, + "backend_ids": backends, + } + frameworks[label] = { + "readback_count": len(matching), + "modules": module_summary, + } + return {"passed": not errors, "errors": errors, "frameworks": frameworks} + + +def _load_train_dump(path: Path) -> Mapping[str, Any]: + value = torch.load(path, map_location="cpu", weights_only=False) + if not isinstance(value, Mapping): + raise ValueError(f"train dump must contain a mapping: {path}") + return value + + +def compare_train_rollout_logps(paths: list[Path]) -> dict[str, Any]: + sample_count = 0 + element_count = 0 + mismatch_count = 0 + max_abs_diff = 0.0 + errors: list[str] = [] + for path in paths: + payload = _load_train_dump(path) + samples = payload.get("samples") + if not isinstance(samples, list): + rollout_data = payload.get("rollout_data") + if not isinstance(rollout_data, Mapping): + errors.append(f"{path} has neither samples nor rollout_data") + continue + training_values = rollout_data.get("log_probs") + rollout_values = rollout_data.get("rollout_log_probs") + if not isinstance(training_values, (list, tuple)) or not isinstance( + rollout_values, (list, tuple) + ): + errors.append(f"{path} rollout_data lacks list log_probs/rollout_log_probs") + continue + if len(training_values) != len(rollout_values): + errors.append( + f"{path} logprob list length mismatch: " + f"{len(training_values)} != {len(rollout_values)}" + ) + samples = [ + {"log_probs": training, "rollout_log_probs": rollout} + for training, rollout in zip(training_values, rollout_values, strict=False) + ] + for sample_index, sample in enumerate(samples): + if not isinstance(sample, Mapping): + errors.append(f"{path} sample {sample_index} is not a mapping") + continue + training = sample.get("log_probs") + rollout = sample.get("rollout_log_probs") + if not isinstance(training, torch.Tensor) or not isinstance(rollout, torch.Tensor): + errors.append( + f"{path} sample {sample_index} lacks tensor log_probs/rollout_log_probs" + ) + continue + sample_count += 1 + if training.shape != rollout.shape: + errors.append( + f"{path} sample {sample_index} shape mismatch: " + f"{tuple(training.shape)} != {tuple(rollout.shape)}" + ) + continue + if training.dtype != rollout.dtype: + errors.append( + f"{path} sample {sample_index} dtype mismatch: " + f"{training.dtype} != {rollout.dtype}" + ) + element_count += training.numel() + mismatch_count += int(torch.ne(training, rollout).sum().item()) + if training.numel(): + diff = (training.float() - rollout.float()).abs() + if not bool(torch.isfinite(diff).all().item()): + errors.append(f"{path} sample {sample_index} has non-finite drift") + else: + max_abs_diff = max(max_abs_diff, float(diff.max().item())) + if not paths: + errors.append("no Vime train dump was found") + if sample_count == 0: + errors.append("no comparable train/rollout samples were found") + torch_equal = not errors and mismatch_count == 0 + return { + "passed": torch_equal and max_abs_diff == 0.0, + "torch_equal": torch_equal, + "mismatch_count": mismatch_count, + "max_abs_diff": max_abs_diff if math.isfinite(max_abs_diff) else None, + "sample_count": sample_count, + "element_count": element_count, + "errors": errors, + "artifacts": [str(path) for path in paths], + } + + +def validate_artifacts(readback_dir: Path, train_data_dir: Path) -> dict[str, Any]: + readbacks = validate_readbacks(load_readbacks(readback_dir)) + train_paths = sorted(train_data_dir.glob("*.pt")) + bitwise = compare_train_rollout_logps(train_paths) + return { + "schema_version": "rlkernel.vime_cuda_bitwise_validation.v1", + "passed": bool(readbacks["passed"] and bitwise["passed"]), + "runtime_policy": {"platform": "cuda", "triton_allowed": False}, + "readbacks": readbacks, + "train_rollout_logp": bitwise, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--readback-dir", type=Path, required=True) + parser.add_argument("--train-data-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args(argv) + + try: + report = validate_artifacts(args.readback_dir, args.train_data_dir) + except Exception as exc: + report = { + "schema_version": "rlkernel.vime_cuda_bitwise_validation.v1", + "passed": False, + "runtime_policy": {"platform": "cuda", "triton_allowed": False}, + "error": str(exc), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if report["passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 3aa5ded1..ca3b0c5d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,9 @@ dependencies = [ "transformers==5.13.1", ] +[project.entry-points."vllm.general_plugins"] +rl_kernel = "rl_engine.integrations.vllm_runtime:register_vllm_plugin" + [project.optional-dependencies] cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] rocm = ["aiter"] diff --git a/rl_engine/integrations/__init__.py b/rl_engine/integrations/__init__.py index 24eac10c..49b7193d 100644 --- a/rl_engine/integrations/__init__.py +++ b/rl_engine/integrations/__init__.py @@ -7,11 +7,15 @@ Implementation, IntegrationPlan, OperatorAblationCase, + configure_integration_environment, + integration_plan_from_environment, operator_ablation_case, operator_ablation_cases, ) from rl_engine.integrations.megatron import MegatronIntegration +from rl_engine.integrations.megatron_runtime import install_megatron_integration from rl_engine.integrations.vllm import VllmIntegration +from rl_engine.integrations.vllm_runtime import configure_vllm_environment __all__ = [ "Implementation", @@ -19,6 +23,10 @@ "MegatronIntegration", "OperatorAblationCase", "VllmIntegration", + "configure_integration_environment", + "configure_vllm_environment", + "install_megatron_integration", + "integration_plan_from_environment", "operator_ablation_case", "operator_ablation_cases", ] diff --git a/rl_engine/integrations/ablation.py b/rl_engine/integrations/ablation.py index a2ccdb82..71097675 100644 --- a/rl_engine/integrations/ablation.py +++ b/rl_engine/integrations/ablation.py @@ -5,6 +5,7 @@ from __future__ import annotations +import os from dataclasses import dataclass from enum import Enum from types import MappingProxyType @@ -49,8 +50,18 @@ def to_dict(self) -> dict[str, Any]: _CASE_DEFINITIONS = ( ("P/P", Implementation.PRODUCTION, Implementation.PRODUCTION, "native baseline"), ("R/R", Implementation.RL_KERNEL, Implementation.RL_KERNEL, "RL-Kernel control"), - ("P/R", Implementation.PRODUCTION, Implementation.RL_KERNEL, "rollout-only mismatch"), - ("R/P", Implementation.RL_KERNEL, Implementation.PRODUCTION, "training-only mismatch"), + ( + "P/R", + Implementation.PRODUCTION, + Implementation.RL_KERNEL, + "rollout-only mismatch", + ), + ( + "R/P", + Implementation.RL_KERNEL, + Implementation.PRODUCTION, + "training-only mismatch", + ), ) @@ -121,10 +132,39 @@ def to_dict(self) -> dict[str, Any]: } +def integration_plan_from_environment() -> IntegrationPlan: + """Materialize the one plan inherited by framework worker processes.""" + + return IntegrationPlan.from_case_ids( + attention=os.getenv("RL_KERNEL_ATTENTION_CASE", "P/P"), + ffn=os.getenv("RL_KERNEL_FFN_CASE", "P/P"), + logp=os.getenv("RL_KERNEL_LOGP_CASE", "P/P"), + ) + + +def configure_integration_environment( + plan: IntegrationPlan, + *, + readback_dir: str | None = None, +) -> None: + """Export one plan for Megatron actors and vLLM subprocesses.""" + + for module, variable in ( + ("attention", "RL_KERNEL_ATTENTION_CASE"), + ("ffn", "RL_KERNEL_FFN_CASE"), + ("logp", "RL_KERNEL_LOGP_CASE"), + ): + os.environ[variable] = plan.cases[module].case_id + if readback_dir: + os.environ["RL_KERNEL_READBACK_DIR"] = readback_dir + + __all__ = [ "Implementation", "IntegrationPlan", "OperatorAblationCase", + "configure_integration_environment", + "integration_plan_from_environment", "operator_ablation_case", "operator_ablation_cases", ] diff --git a/rl_engine/integrations/framework_operators.py b/rl_engine/integrations/framework_operators.py new file mode 100644 index 00000000..16e11784 --- /dev/null +++ b/rl_engine/integrations/framework_operators.py @@ -0,0 +1,1297 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tensor-layout adapters from framework boundaries to semantic operators. + +This module owns no numerical kernel selection. Attention and FFN instances +are resolved by :class:`OperatorBridge`; Logp uses the existing contract-aware +``KernelRegistry`` dispatch shared with the Vime provider. +""" + +from __future__ import annotations + +import os +from dataclasses import replace +from threading import Lock +from typing import Any, Mapping, cast + +import torch +from rl_engine.alignment.cross_config.operators import OperatorBridge, OperatorOverride +from rl_engine.integrations.linear_logp import LinearLogpWrapper, take_rollout_linear_logp_context +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_PRODUCTION_CORE_ID, + AttentionContract, + AttentionDType, + AttentionMode, + AttentionRole, +) +from rl_engine.kernels.attention_contract import ReductionSpec as AttentionReductionSpec +from rl_engine.kernels.attention_contract import ShardingSpec as AttentionShardingSpec +from rl_engine.kernels.attention_contract import SplitKVSpec +from rl_engine.kernels.logprob_contract import LogprobContract, LogprobDType, LogprobRole, MaskSpec +from rl_engine.kernels.logprob_contract import ReductionSpec as LogprobReductionSpec +from rl_engine.kernels.logprob_contract import ShardingSpec as LogprobShardingSpec +from rl_engine.kernels.ops.pytorch.attention.ablation import AttentionAblationConfig +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import BACKEND_ID as LOGP_BACKEND_ID +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import DEFAULT_NUM_VOCAB_TILES +from rl_engine.kernels.registry import kernel_registry +from rl_engine.kernels.semantic_registry import OperatorRequirements + +ATTENTION_BACKEND_ID = "rlkernel.attention.deterministic.v1" +FFN_BACKEND_ID = "rlkernel.ffn.qwen3.deterministic.v1" + + +def _device_name(tensor: torch.Tensor) -> str: + if tensor.device.type == "cuda" and torch.version.hip is not None: + return "rocm" + return tensor.device.type + + +def _dtype_name(tensor: torch.Tensor) -> str: + return str(tensor.dtype).replace("torch.", "") + + +def _optional_env_int(name: str) -> int | None: + value = os.getenv(name, "").strip() + if not value: + return None + try: + return int(value) + except ValueError: + return None + + +def _tensor_debug_stats(tensor: torch.Tensor) -> dict[str, Any]: + detached = tensor.detach() + stats: dict[str, Any] = { + "shape": list(detached.shape), + "dtype": _dtype_name(detached), + "device": str(detached.device), + "numel": int(detached.numel()), + } + if detached.numel() == 0: + return stats + values = detached.float() + stats.update( + { + "min": float(values.min().item()), + "max": float(values.max().item()), + "mean": float(values.mean().item()), + } + ) + return stats + + +def _diff_debug_stats(left: torch.Tensor, right: torch.Tensor) -> dict[str, Any]: + if left.shape != right.shape: + return { + "shape_mismatch": True, + "left_shape": list(left.shape), + "right_shape": list(right.shape), + } + diff = (left.detach().float() - right.detach().float()).abs() + stats = _tensor_debug_stats(diff) + stats["mismatch_count"] = int(torch.ne(left.detach(), right.detach()).sum().item()) + return stats + + +def _attention_dtype(tensor: torch.Tensor) -> AttentionDType: + try: + return { + torch.bfloat16: AttentionDType.BF16, + torch.float16: AttentionDType.FP16, + torch.float32: AttentionDType.FP32, + }[tensor.dtype] + except KeyError as exc: + raise RuntimeError(f"unsupported Attention dtype {tensor.dtype}") from exc + + +def _require_nvidia_cuda(tensor: torch.Tensor, module: str) -> None: + if tensor.device.type != "cuda" or torch.version.hip is not None: + raise RuntimeError(f"strict {module} R/R requires NVIDIA CUDA tensors") + + +class SemanticOperatorHandle: + """Resolve one exact semantic backend once for one framework process.""" + + def __init__(self, *, target: str, semantic_op: str, backend_id: str) -> None: + if target not in {"training", "rollout"}: + raise ValueError("target must be 'training' or 'rollout'") + self.target = target + self.semantic_op = semantic_op + self.backend_id = backend_id + self._bridge = OperatorBridge() + self._instance: Any | None = None + self._requirements: OperatorRequirements | None = None + self._provenance: dict[str, Any] | None = None + self._lock = Lock() + + def get( + self, + tensor: torch.Tensor, + *, + topology: Mapping[str, Any], + factory_kwargs: Mapping[str, Any] | None = None, + ) -> Any: + requirements = OperatorRequirements( + device=_device_name(tensor), + dtype=_dtype_name(tensor), + topology=topology, + alignment_properties={"deterministic": True}, + ) + # This method is called from vLLM model forwards that may be captured + # by torch.compile. A Lock context manager is unsupported in a + # Dynamo fullgraph, so keep the hot-path lookup lock-free. Handles + # are constructed per framework worker and resolution is idempotent; + # duplicate first-call resolution is harmless and the bridge caches + # the resulting semantic instance. + if self._instance is not None: + # vLLM constructs the plugin before its worker TP group exists, so + # the eager prime may observe TP=1 while the first model call has + # TP=2. The operator receives the live group at invocation time; + # keep device and dtype strict, but do not reject this topology + # transition after the semantic instance is resolved. + if self._requirements is not None and ( + requirements.device != self._requirements.device + or requirements.dtype != self._requirements.dtype + ): + raise RuntimeError( + f"{self.semantic_op} runtime device/dtype changed after resolution" + ) + return self._instance + target = cast(Any, self.target) + resolved = self._bridge.resolve_override( + OperatorOverride.for_target( + semantic_op=self.semantic_op, + backend_id=self.backend_id, + target=target, + ), + requirements={self.target: requirements}, + strict=True, + ) + instance = self._bridge.instantiate( + resolved, + target=target, + factory_kwargs=factory_kwargs, + cache=True, + ) + provenance = self._bridge.instance_provenance( + resolved, + target=target, + instance=instance, + ) + actual_backend = getattr(instance, "backend_id", None) + if actual_backend != self.backend_id: + raise RuntimeError( + f"semantic registry resolved {self.backend_id!r} but instantiated " + f"{actual_backend!r}" + ) + self._instance = instance + self._requirements = requirements + self._provenance = provenance.to_dict() + return instance + + @property + def provenance(self) -> Mapping[str, Any] | None: + return None if self._provenance is None else dict(self._provenance) + + +def _weight(module: Any, name: str) -> torch.Tensor: + value = getattr(module, "weight", None) + if not isinstance(value, torch.Tensor): + raise RuntimeError(f"{name} must expose an unquantized torch.Tensor weight") + if value.ndim != 2: + raise RuntimeError(f"{name}.weight must be two-dimensional") + return value + + +def _split_gate_up(weight: torch.Tensor, name: str) -> tuple[torch.Tensor, torch.Tensor]: + if weight.size(0) % 2: + raise RuntimeError(f"{name}.weight first dimension must contain equal gate/up shards") + gate, up = weight.chunk(2, dim=0) + return gate.contiguous(), up.contiguous() + + +def _megatron_parallel_state() -> Any: + try: + from megatron.core import parallel_state + except ImportError as exc: # pragma: no cover - exercised in framework environment + raise RuntimeError("Megatron parallel_state is unavailable") from exc + return parallel_state + + +def _megatron_zigzag_layout( + local_tokens: int, + *, + cp_rank: int, + cp_world_size: int, +) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + if cp_world_size == 1: + positions = tuple(range(local_tokens)) + return positions, (0,), (0,), (0, local_tokens) + if local_tokens % 2: + raise RuntimeError("Megatron zigzag CP requires an even local sequence length") + chunk_size = local_tokens // 2 + second_index = 2 * cp_world_size - cp_rank - 1 + starts = (cp_rank * chunk_size, second_index * chunk_size) + positions = tuple(range(starts[0], starts[0] + chunk_size)) + tuple( + range(starts[1], starts[1] + chunk_size) + ) + return positions, (cp_rank, second_index), starts, (0, chunk_size, local_tokens) + + +def _packed_local_sequence_layout( + packed_seq_params: Any, + *, + cp_world_size: int, + local_query_tokens: int, + local_kv_tokens: int, +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Recover local THD sequence offsets from Megatron's global cu_seqlens.""" + + if str(getattr(packed_seq_params, "qkv_format", "")).lower() != "thd": + raise RuntimeError("strict RL-Kernel packed Attention requires qkv_format='thd'") + query_cu = getattr(packed_seq_params, "cu_seqlens_q", None) + kv_cu = getattr(packed_seq_params, "cu_seqlens_kv", None) + if not isinstance(query_cu, torch.Tensor) or not isinstance(kv_cu, torch.Tensor): + raise RuntimeError("packed Attention requires tensor cu_seqlens_q/cu_seqlens_kv") + query_offsets = tuple( + int(value) for value in query_cu.detach().to(device="cpu", dtype=torch.int64).tolist() + ) + kv_offsets = tuple( + int(value) for value in kv_cu.detach().to(device="cpu", dtype=torch.int64).tolist() + ) + if query_offsets != kv_offsets: + raise RuntimeError("strict self-Attention requires identical Q and KV cu_seqlens") + if len(query_offsets) < 2 or query_offsets[0] != 0: + raise RuntimeError("packed Attention cu_seqlens must start at zero") + global_lengths = tuple( + right - left for left, right in zip(query_offsets[:-1], query_offsets[1:], strict=True) + ) + if any(length <= 0 for length in global_lengths): + raise RuntimeError("packed Attention cu_seqlens must be strictly increasing") + if any(length % cp_world_size for length in global_lengths): + raise RuntimeError("packed Attention sequence lengths must be divisible by CP size") + local_lengths = tuple(length // cp_world_size for length in global_lengths) + local_offsets = [0] + for length in local_lengths: + local_offsets.append(local_offsets[-1] + length) + if local_offsets[-1] != local_query_tokens or local_offsets[-1] != local_kv_tokens: + raise RuntimeError("packed Attention cu_seqlens do not cover the local Q/KV token rows") + return tuple(local_offsets), global_lengths + + +def _compact_attention_provenance(value: Mapping[str, Any]) -> dict[str, Any]: + """Keep strict backend identity without retaining one record per token.""" + + compact = dict(value) + rows = compact.pop("core_rows", None) + if isinstance(rows, (list, tuple)): + compact["core_row_count"] = len(rows) + backends = sorted( + { + str(row["actual_backend"]) + for row in rows + if isinstance(row, Mapping) and row.get("actual_backend") + } + ) + compact["core_actual_backends"] = backends + return compact + + +def _dense_attention_contract( + query: torch.Tensor, + key: torch.Tensor, + *, + role: AttentionRole, + causal: bool, + tp_rank: int, + tp_world_size: int, + cp_rank: int = 0, + cp_world_size: int = 1, + mode: AttentionMode | None = None, + global_sequence_length: int | None = None, + global_block_indices: tuple[int, ...] = (0,), + global_block_token_starts: tuple[int, ...] = (0,), + local_block_offsets: tuple[int, ...] | None = None, +) -> AttentionContract: + batch, q_heads, query_tokens, head_dim = query.shape + kv_heads = key.size(1) + return AttentionContract( + role=role, + mode=mode or AttentionMode.PREFILL, + dtype=_attention_dtype(query), + batch_size=batch, + query_sequence_length=query_tokens, + head_dim=head_dim, + causal=causal, + causal_offsets=(0,) * batch if causal else None, + sharding=AttentionShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=q_heads * tp_world_size, + global_kv_heads=kv_heads * tp_world_size, + local_q_head_start=tp_rank * q_heads, + local_q_heads=q_heads, + local_kv_head_start=tp_rank * kv_heads, + local_kv_heads=kv_heads, + global_sequence_length=global_sequence_length or query_tokens, + local_sequence_length=query_tokens, + global_block_indices=global_block_indices, + global_block_token_starts=global_block_token_starts, + local_block_offsets=local_block_offsets or (0, query_tokens), + ), + reduction=AttentionReductionSpec(), + split_kv=SplitKVSpec.disabled(), + export_lse=True, + ) + + +class MegatronAttentionOperator: + """Materialize Megatron layout, then call the registered Attention wrapper.""" + + backend_id = ATTENTION_BACKEND_ID + + def __init__(self, handle: SemanticOperatorHandle | None = None) -> None: + self._handle = handle or SemanticOperatorHandle( + target="training", semantic_op="attention", backend_id=self.backend_id + ) + self._last_provenance: dict[str, Any] = {} + + @property + def provenance(self) -> Mapping[str, Any]: + return { + "semantic_instance": self._handle.provenance, + "execution": dict(self._last_provenance), + } + + def __call__( + self, + module: Any, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + attn_mask_type: Any = None, + attention_bias: torch.Tensor | None = None, + packed_seq_params: Any = None, + num_splits: int | None = None, + ) -> torch.Tensor: + del attn_mask_type + if attention_bias is not None: + raise RuntimeError("strict RL-Kernel Attention does not accept bias") + if num_splits not in (None, 1): + raise RuntimeError("strict RL-Kernel Attention requires num_splits=1") + if attention_mask is not None and attention_mask.numel() > 1: + raise RuntimeError("strict RL-Kernel Attention supports only its causal contract") + expected_ndim = 3 if packed_seq_params is not None else 4 + if query.ndim != expected_ndim or key.ndim != expected_ndim or value.ndim != expected_ndim: + layout = "[T, H, D]" if packed_seq_params is not None else "[S, B, H, D]" + raise RuntimeError(f"Megatron Attention Q/K/V must use {layout}") + _require_nvidia_cuda(query, "Attention") + + parallel_state = _megatron_parallel_state() + cp_world = int(parallel_state.get_context_parallel_world_size()) + cp_rank = int(parallel_state.get_context_parallel_rank()) + tp_world = int(parallel_state.get_tensor_model_parallel_world_size()) + tp_rank = int(parallel_state.get_tensor_model_parallel_rank()) + cp_group = parallel_state.get_context_parallel_group() if cp_world > 1 else None + operator = self._handle.get( + query, + topology={ + "world_size": tp_world * cp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": cp_world, + }, + ) + operator.bind_cuda_runtime(process_group=cp_group) + scale = float(getattr(module, "softmax_scale", query.size(-1) ** -0.5)) + + def execute_sequence( + q_ready: torch.Tensor, + k_ready: torch.Tensor, + v_ready: torch.Tensor, + *, + global_sequence_length: int, + ) -> Any: + positions, block_indices, block_starts, block_offsets = _megatron_zigzag_layout( + q_ready.size(2), + cp_rank=cp_rank, + cp_world_size=cp_world, + ) + position_ids = torch.tensor( + positions, + dtype=torch.int64, + device=q_ready.device, + ).repeat(q_ready.size(0), 1) + return operator( + q_ready, + k_ready, + v_ready, + contract=_dense_attention_contract( + q_ready, + k_ready, + role=AttentionRole.TRAIN, + causal=True, + tp_rank=tp_rank, + tp_world_size=tp_world, + cp_rank=cp_rank, + cp_world_size=cp_world, + global_sequence_length=global_sequence_length, + global_block_indices=block_indices, + global_block_token_starts=block_starts, + local_block_offsets=block_offsets, + ), + config=AttentionAblationConfig( + strict_core_id=STRICT_ATTENTION_PRODUCTION_CORE_ID, + strict_schedule=STRICT_ATTENTION_FA4_SCHEDULE_ID, + ), + return_lse=True, + communication_backend="cuda_ag_rs" if cp_world > 1 else "none", + query_position_ids=position_ids, + key_position_ids=position_ids, + scale=scale, + ) + + if packed_seq_params is None: + q_ready = query.permute(1, 2, 0, 3).contiguous() + k_ready = key.permute(1, 2, 0, 3).contiguous() + v_ready = value.permute(1, 2, 0, 3).contiguous() + result = execute_sequence( + q_ready, + k_ready, + v_ready, + global_sequence_length=q_ready.size(2) * cp_world, + ) + context = result.out.permute(2, 0, 1, 3).contiguous() + output = context.flatten(start_dim=2) + execution_provenance: dict[str, Any] = { + "packed_sequence_count": 0, + "operator": _compact_attention_provenance(result.provenance), + } + else: + local_offsets, global_lengths = _packed_local_sequence_layout( + packed_seq_params, + cp_world_size=cp_world, + local_query_tokens=query.size(0), + local_kv_tokens=key.size(0), + ) + outputs: list[torch.Tensor] = [] + sequence_provenance: list[dict[str, Any]] = [] + for sequence_index, (start, end, global_length) in enumerate( + zip( + local_offsets[:-1], + local_offsets[1:], + global_lengths, + strict=True, + ) + ): + q_ready = query[start:end].permute(1, 0, 2).unsqueeze(0).contiguous() + k_ready = key[start:end].permute(1, 0, 2).unsqueeze(0).contiguous() + v_ready = value[start:end].permute(1, 0, 2).unsqueeze(0).contiguous() + result = execute_sequence( + q_ready, + k_ready, + v_ready, + global_sequence_length=global_length, + ) + outputs.append( + result.out.squeeze(0).permute(1, 0, 2).contiguous().flatten(start_dim=1) + ) + sequence_provenance.append( + { + "sequence_index": sequence_index, + "local_tokens": end - start, + "global_tokens": global_length, + "operator": _compact_attention_provenance(result.provenance), + } + ) + output = torch.cat(outputs, dim=0) + execution_provenance = { + "packed_sequence_count": len(global_lengths), + "sequences": sequence_provenance, + } + self._last_provenance = { + "framework_layout": ( + "megatron_thd_packed_zigzag_cp" + if packed_seq_params is not None + else "megatron_sbh_zigzag_cp" + ), + "materialization": "owner_local_zigzag_cuda_ag_rs", + "cp_world_size": cp_world, + "tp_world_size": tp_world, + "runtime_platform": "cuda", + "triton_used": False, + **execution_provenance, + } + return output + + +class MegatronFFNOperator: + backend_id = FFN_BACKEND_ID + + def __init__(self, handle: SemanticOperatorHandle | None = None) -> None: + self._handle = handle or SemanticOperatorHandle( + target="training", semantic_op="ffn", backend_id=self.backend_id + ) + self._last_provenance: dict[str, Any] = {} + + @property + def provenance(self) -> Mapping[str, Any]: + return { + "semantic_instance": self._handle.provenance, + "execution": dict(self._last_provenance), + } + + def __call__( + self, + module: Any, + hidden_states: torch.Tensor, + per_token_scale: torch.Tensor | None = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, None]: + # Megatron's dense MLP path forwards padding_mask=None even when + # no token padding is active. It is a routing-only argument and must + # not be treated as expert-token scaling in the strict dense wrapper. + if kwargs: + unexpected = { + name: value + for name, value in kwargs.items() + if name != "padding_mask" or value is not None + } + if unexpected: + raise RuntimeError("strict dense Qwen3 FFN does not accept expert token scaling") + if per_token_scale is not None: + raise RuntimeError("strict dense Qwen3 FFN does not accept expert token scaling") + _require_nvidia_cuda(hidden_states, "FFN") + config = module.config + if bool(getattr(config, "add_bias_linear", False)): + raise RuntimeError("strict Qwen3 FFN requires bias-free projections") + if not bool(getattr(config, "gated_linear_unit", False)): + raise RuntimeError("strict Qwen3 FFN requires a gated linear unit") + gate, up = _split_gate_up(_weight(module.linear_fc1, "linear_fc1"), "linear_fc1") + down = _weight(module.linear_fc2, "linear_fc2").contiguous() + parallel_state = _megatron_parallel_state() + cp_world = int(parallel_state.get_context_parallel_world_size()) + tp_world = int(parallel_state.get_tensor_model_parallel_world_size()) + cp_group = parallel_state.get_context_parallel_group() if cp_world > 1 else None + operator = self._handle.get( + hidden_states, + topology={ + "world_size": tp_world * cp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": cp_world, + }, + ) + output = operator( + hidden_states, + gate, + up, + down, + tp_group=getattr(module, "tp_group", None), + cp_group=cp_group, + sequence_parallel=bool(getattr(config, "sequence_parallel", False)), + deterministic=True, + ) + self._last_provenance = { + "framework_layout": "megatron_sequence_parallel", + "cp_world_size": cp_world, + "tp_world_size": tp_world, + "runtime_platform": "cuda", + "actual_backend": "rlkernel.cuda.det_gemm_swiglu", + "triton_used": False, + } + return output, None + + +def _vllm_tp_coordinates() -> tuple[int, int, Any]: + try: + from vllm.distributed.parallel_state import get_tp_group + + coordinator = get_tp_group() + group = getattr(coordinator, "device_group", coordinator) + rank = int(getattr(coordinator, "rank_in_group", 0)) + world = int(getattr(coordinator, "world_size", 1)) + return world, rank, group + except (ImportError, AssertionError): + return 1, 0, None + + +def _vllm_kv_cache_views( + kv_cache: torch.Tensor, + *, + head_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return vLLM's paged cache as [blocks, block, kv_heads, head].""" + + if kv_cache.ndim == 4 and kv_cache.size(-1) == 2 * head_size: + return kv_cache.transpose(1, 2).split(head_size, dim=-1) + if kv_cache.ndim == 5 and kv_cache.size(0) == 2: + key_cache, value_cache = kv_cache.unbind(0) + return key_cache, value_cache + raise RuntimeError( + "vLLM FlashAttention KV cache must use " "[blocks, kv_heads, block, 2 * head_size]" + ) + + +class VllmAttentionOperator: + """Materialize logical rows from vLLM paged KV, then call the wrapper.""" + + backend_id = ATTENTION_BACKEND_ID + + def __init__(self, handle: SemanticOperatorHandle | None = None) -> None: + self._handle = handle or SemanticOperatorHandle( + target="rollout", semantic_op="attention", backend_id=self.backend_id + ) + self._last_provenance: dict[str, Any] = {} + self._prime_semantic_handle() + + def _prime_semantic_handle(self) -> None: + # vLLM wraps model execution in torch.compile. Resolve the semantic + # descriptor before graph capture so JSON/inspect-based provenance + # never runs inside Dynamo's fullgraph region. + if not torch.cuda.is_available(): + return + try: + tp_world, _rank, _group = _vllm_tp_coordinates() + self._handle.get( + torch.empty((1,), device="cuda", dtype=torch.bfloat16), + topology={ + "world_size": tp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": 1, + }, + ) + except (RuntimeError, ValueError, ImportError): + # API-server/plugin construction can precede worker CUDA setup; + # the worker retries during initialization before graph capture. + return + + @property + def provenance(self) -> Mapping[str, Any]: + return { + "semantic_instance": self._handle.provenance, + "execution": dict(self._last_provenance), + } + + @staticmethod + def _metadata_tensor(metadata: Any, *names: str) -> torch.Tensor: + for name in names: + value = getattr(metadata, name, None) + if isinstance(value, torch.Tensor): + return value + raise RuntimeError(f"vLLM Attention metadata is missing {'/'.join(names)}") + + def __call__( + self, + impl: Any, + layer: Any, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: Any, + output: torch.Tensor | None = None, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + del layer, key, value + if output_scale is not None or output_block_scale is not None: + raise RuntimeError("strict vLLM Attention does not support quantized output") + if attn_metadata is None: + if output is None: + raise RuntimeError("vLLM profiling Attention requires an output buffer") + return output.zero_() + if query.ndim != 3: + raise RuntimeError("vLLM query must use [tokens, heads, head_dim]") + _require_nvidia_cuda(query, "Attention") + key_cache, value_cache = _vllm_kv_cache_views( + kv_cache, + head_size=int(impl.head_size), + ) + if key_cache.dtype != query.dtype or value_cache.dtype != query.dtype: + raise RuntimeError("strict vLLM Attention requires an unquantized KV cache") + + query_starts = self._metadata_tensor( + attn_metadata, "query_start_loc", "query_start_loc_cpu" + ).to(device="cpu", dtype=torch.long) + seq_lens = self._metadata_tensor(attn_metadata, "seq_lens").to( + device="cpu", dtype=torch.long + ) + block_table = self._metadata_tensor(attn_metadata, "block_table", "block_table_tensor").to( + device="cpu", dtype=torch.long + ) + num_actual = int(getattr(attn_metadata, "num_actual_tokens", query.size(0))) + if output is None: + output = torch.empty( + (query.size(0), impl.num_heads * impl.head_size), + dtype=query.dtype, + device=query.device, + ) + output.zero_() + output_heads = output.view(output.size(0), impl.num_heads, impl.head_size) + block_size = key_cache.size(1) + tp_world, tp_rank, tp_group = _vllm_tp_coordinates() + operator = self._handle.get( + query, + topology={ + "world_size": tp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": 1, + }, + ) + operator.bind_cuda_runtime() + row_count = 0 + first_query_position: int | None = None + last_query_position: int | None = None + min_kv_tokens: int | None = None + max_kv_tokens: int | None = None + last_operator_provenance: dict[str, Any] = {} + for request_index in range(seq_lens.numel()): + q_start = int(query_starts[request_index]) + q_end = min(int(query_starts[request_index + 1]), num_actual) + if q_start >= q_end: + continue + final_seq_len = int(seq_lens[request_index]) + first_query_position = final_seq_len - (q_end - q_start) + for token_offset, query_index in enumerate(range(q_start, q_end)): + kv_len = first_query_position + token_offset + 1 + page_count = (kv_len + block_size - 1) // block_size + page_ids = block_table[request_index, :page_count].tolist() + if any(page < 0 for page in page_ids): + raise RuntimeError("vLLM block table contains an unallocated page") + pages = torch.tensor(page_ids, device=key_cache.device, dtype=torch.long) + k_row = key_cache.index_select(0, pages).reshape( + -1, impl.num_kv_heads, impl.head_size + )[:kv_len] + v_row = value_cache.index_select(0, pages).reshape( + -1, impl.num_kv_heads, impl.head_size + )[:kv_len] + q_ready = query[query_index : query_index + 1].permute(1, 0, 2).unsqueeze(0) + k_ready = k_row.permute(1, 0, 2).unsqueeze(0).contiguous() + v_ready = v_row.permute(1, 0, 2).unsqueeze(0).contiguous() + query_positions = torch.tensor( + [[kv_len - 1]], + dtype=torch.int64, + device=query.device, + ) + key_positions = torch.arange( + kv_len, + dtype=torch.int64, + device=query.device, + ).unsqueeze(0) + result = operator( + q_ready.contiguous(), + k_ready, + v_ready, + contract=_dense_attention_contract( + q_ready, + k_ready, + role=AttentionRole.INFER, + # The causal prefix is already materialized as one dense row. + causal=False, + tp_rank=tp_rank, + tp_world_size=tp_world, + mode=AttentionMode.CHUNKED_PREFILL, + global_sequence_length=kv_len, + global_block_token_starts=(kv_len - 1,), + ), + config=AttentionAblationConfig( + strict_core_id=STRICT_ATTENTION_PRODUCTION_CORE_ID, + strict_schedule=STRICT_ATTENTION_FA4_SCHEDULE_ID, + ), + return_lse=True, + query_position_ids=query_positions, + key_position_ids=key_positions, + scale=float(impl.scale), + ) + output_heads[query_index].copy_(result.out[0, :, 0, :]) + query_position = kv_len - 1 + row_count += 1 + first_query_position = ( + query_position + if first_query_position is None + else min(first_query_position, query_position) + ) + last_query_position = ( + query_position + if last_query_position is None + else max(last_query_position, query_position) + ) + min_kv_tokens = kv_len if min_kv_tokens is None else min(min_kv_tokens, kv_len) + max_kv_tokens = kv_len if max_kv_tokens is None else max(max_kv_tokens, kv_len) + last_operator_provenance = _compact_attention_provenance(result.provenance) + self._last_provenance = { + "framework_layout": "vllm_paged_kv", + "materialization": "one_causal_prefix_per_query_row", + "tp_world_size": tp_world, + "tp_group_bound": tp_group is not None, + "runtime_platform": "cuda", + "triton_used": False, + "row_count": row_count, + "query_position_range": [first_query_position, last_query_position], + "kv_token_range": [min_kv_tokens, max_kv_tokens], + "operator": last_operator_provenance, + } + return output + + +class VllmFFNOperator: + backend_id = FFN_BACKEND_ID + + def __init__(self, handle: SemanticOperatorHandle | None = None) -> None: + self._handle = handle or SemanticOperatorHandle( + target="rollout", semantic_op="ffn", backend_id=self.backend_id + ) + self._last_provenance: dict[str, Any] = {} + self._prime_semantic_handle() + + def _prime_semantic_handle(self) -> None: + # vLLM wraps model execution in torch.compile. Resolve the semantic + # descriptor before graph capture so JSON/inspect-based provenance + # never runs inside Dynamo's fullgraph region. + if not torch.cuda.is_available(): + return + try: + tp_world, _rank, _group = _vllm_tp_coordinates() + self._handle.get( + torch.empty((1,), device="cuda", dtype=torch.bfloat16), + topology={ + "world_size": tp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": 1, + }, + ) + except (RuntimeError, ValueError, ImportError): + # API-server/plugin construction can precede worker CUDA setup; + # the worker retries during initialization before graph capture. + return + + @property + def provenance(self) -> Mapping[str, Any]: + return { + "semantic_instance": self._handle.provenance, + "execution": dict(self._last_provenance), + } + + def __call__(self, module: Any, hidden_states: torch.Tensor) -> torch.Tensor: + _require_nvidia_cuda(hidden_states, "FFN") + gate, up = _split_gate_up(_weight(module.gate_up_proj, "gate_up_proj"), "gate_up_proj") + down = _weight(module.down_proj, "down_proj").contiguous() + tp_world, _tp_rank, tp_group = _vllm_tp_coordinates() + operator = self._handle.get( + hidden_states, + topology={ + "world_size": tp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": 1, + }, + ) + output = operator( + hidden_states, + gate, + up, + down, + tp_group=tp_group, + sequence_parallel=False, + deterministic=True, + ) + self._last_provenance = { + "framework_layout": "vllm_tensor_parallel", + "tp_world_size": tp_world, + "runtime_platform": "cuda", + "actual_backend": "rlkernel.cuda.det_gemm_swiglu", + "triton_used": False, + } + return output + + +class MegatronLogpOperator: + """Require CUDA while reusing the structural Vime Logp provider.""" + + def __init__( + self, + provider: Any, + *, + linear_logp: LinearLogpWrapper | None = None, + ) -> None: + self._provider = provider + self._linear_logp = linear_logp + self._last_provenance: dict[str, Any] = {} + + @property + def backend_id(self) -> str: + if self._linear_logp is not None: + return self._linear_logp.backend_id + return LOGP_BACKEND_ID + + @property + def provenance(self) -> Mapping[str, Any]: + return dict(self._last_provenance) + + def __call__(self, request: Any) -> Any: + logits = getattr(request, "logits", None) + hidden = getattr(request, "hidden", None) + strict = os.getenv("VIME_RL_KERNEL_STRICT", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + if isinstance(hidden, torch.Tensor): + if self._linear_logp is None: + raise RuntimeError("Megatron linear_logp route is not installed") + _require_nvidia_cuda(hidden, "linear_logp") + result = self._provider(request, linear_logp=self._linear_logp) + self._last_provenance = { + "runtime_platform": "cuda", + "triton_used": False, + "provider": dict(getattr(result, "provenance", {})), + "linear_logp": dict(self._linear_logp.provenance), + "logits_materialized": False, + } + return result + if strict: + raise RuntimeError( + "strict Megatron linear_logp request is missing hidden/LM-head structural inputs" + ) + if not isinstance(logits, torch.Tensor): + raise RuntimeError("Megatron Logp request must expose logits or hidden") + _require_nvidia_cuda(logits, "Logp") + result = self._provider(request) + self._last_provenance = { + "runtime_platform": "cuda", + "triton_used": False, + "provider": dict(getattr(result, "provenance", {})), + } + return result + + +def _full_vocab_contract(logits: torch.Tensor) -> LogprobContract: + dtype = { + torch.bfloat16: LogprobDType.BF16, + torch.float16: LogprobDType.FP16, + torch.float32: LogprobDType.FP32, + }.get(logits.dtype) + if dtype is None: + raise RuntimeError(f"unsupported vLLM logit dtype {logits.dtype}") + tokens, vocab = logits.shape + return LogprobContract( + role=LogprobRole.INFER, + dtype=dtype, + mask=MaskSpec(num_tokens=tokens, active_mask=(True,) * tokens), + sharding=LogprobShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, vocab),), + real_vocab_size=vocab, + padded_vocab_size=vocab, + ), + reduction=LogprobReductionSpec(), + ) + + +def _diagnostic_vocab_contract( + logits: torch.Tensor, + *, + real_vocab_size: int | None, + padded_vocab_size: int | None, +) -> LogprobContract | None: + if real_vocab_size is None or padded_vocab_size is None: + return None + if real_vocab_size <= 0 or padded_vocab_size <= 0: + return None + if real_vocab_size > padded_vocab_size: + return None + tokens, vocab = logits.shape + if padded_vocab_size != vocab: + return None + if padded_vocab_size % DEFAULT_NUM_VOCAB_TILES: + return None + dtype = { + torch.bfloat16: LogprobDType.BF16, + torch.float16: LogprobDType.FP16, + torch.float32: LogprobDType.FP32, + }.get(logits.dtype) + if dtype is None: + return None + return LogprobContract( + role=LogprobRole.INFER, + dtype=dtype, + mask=MaskSpec(num_tokens=tokens, active_mask=(True,) * tokens), + sharding=LogprobShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, padded_vocab_size),), + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + ), + reduction=LogprobReductionSpec(), + ) + + +class VllmLogpOperator: + """Keep sampling, then replace sampled-token logp with the selected contract.""" + + def __init__( + self, + native_forward: Any, + *, + worker_sampler: bool = False, + strict_linear_logp: bool = False, + ) -> None: + self._native_forward = native_forward + self._worker_sampler = worker_sampler + self._strict_linear_logp = strict_linear_logp + self._linear_logp = LinearLogpWrapper() if strict_linear_logp else None + self._last_provenance: dict[str, Any] = {} + + @property + def backend_id(self) -> str: + if self._strict_linear_logp: + assert self._linear_logp is not None + return self._linear_logp.backend_id + return LOGP_BACKEND_ID + + @property + def provenance(self) -> Mapping[str, Any]: + return dict(self._last_provenance) + + def _replace_sampled_value( + self, + result: Any, + *, + token_ids: torch.Tensor, + selected: torch.Tensor, + provenance: Mapping[str, Any], + ) -> Any: + logprobs_tensors = getattr(result, "logprobs_tensors", None) + if logprobs_tensors is None: + raise RuntimeError( + "strict vLLM Logp requires sampled-token logprob tensors; " + "native sampling returned none" + ) + ids = logprobs_tensors.logprob_token_ids + values = logprobs_tensors.logprobs.clone() + if selected.numel() != token_ids.numel(): + raise RuntimeError( + "strict vLLM logp selected output is not aligned with sampled tokens" + ) + matches = ids == token_ids.unsqueeze(1) + if not bool(matches.any(dim=1).all().item()): + raise RuntimeError("vLLM logprob result does not contain every sampled token") + columns = matches.to(torch.int64).argmax(dim=1) + native_selected = values[ + torch.arange(values.size(0), device=values.device), columns + ].clone() + values[torch.arange(values.size(0), device=values.device), columns] = selected + self._last_provenance = { + **dict(provenance), + "sampled_token_ids": _tensor_debug_stats(token_ids), + "logprobs_shape": list(values.shape), + "logprob_token_ids_shape": list(ids.shape), + "native_selected_logp_stats": _tensor_debug_stats(native_selected), + "rlkernel_selected_logp_stats": _tensor_debug_stats(selected), + "native_vs_rlkernel_selected_diff": _diff_debug_stats(native_selected, selected), + } + if hasattr(logprobs_tensors, "_replace"): + updated_tensors = logprobs_tensors._replace(logprobs=values) + else: + updated_tensors = replace(logprobs_tensors, logprobs=values) + return replace(result, logprobs_tensors=updated_tensors) + + def __call__( + self, + sampler: Any, + logits: torch.Tensor, + sampling_metadata: Any, + predict_bonus_token: bool = False, + logprobs_mode_override: Any = None, + ) -> Any: + source_logits = logits.clone() + _require_nvidia_cuda(source_logits, "Logp") + if self._worker_sampler: + result = self._native_forward(sampler, logits, sampling_metadata) + else: + result = self._native_forward( + sampler, + logits, + sampling_metadata, + predict_bonus_token=predict_bonus_token, + logprobs_mode_override=logprobs_mode_override, + ) + logprobs_tensors = getattr(result, "logprobs_tensors", None) + if logprobs_tensors is None: + raise RuntimeError( + "strict vLLM Logp requires sampled-token logprob tensors; " + "native sampling returned none" + ) + token_ids = result.sampled_token_ids.reshape(-1).to(torch.long) + + if self._strict_linear_logp: + context = take_rollout_linear_logp_context() + if source_logits.ndim != 2: + raise RuntimeError("vLLM sampler logits must be [tokens, vocab]") + if context.hidden.size(0) != source_logits.size(0): + raise RuntimeError( + "strict rollout linear_logp hidden/logits row mismatch: " + f"{context.hidden.size(0)} != {source_logits.size(0)}" + ) + if context.hidden.size(0) != token_ids.numel(): + raise RuntimeError( + "strict rollout linear_logp hidden/sample row mismatch: " + f"{context.hidden.size(0)} != {token_ids.numel()}" + ) + assert self._linear_logp is not None + selected = self._linear_logp( + context.hidden, + context.lm_head_weight, + token_ids, + context.lm_head_bias, + 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", + ) + provenance = { + **dict(self._linear_logp.provenance), + "runtime_platform": "cuda", + "triton_used": False, + "execution": { + "role": "vllm_rollout_linear_logprob", + "strict_backend": True, + "sampling_logits_source": "native_vllm", + "logits_materialized": True, + "padded_lm_head_alignment": True, + }, + "source_logits_shape": list(source_logits.shape), + "source_logits_dtype": _dtype_name(source_logits), + } + return self._replace_sampled_value( + result, + token_ids=token_ids, + selected=selected, + provenance=provenance, + ) + + if source_logits.ndim != 2: + raise RuntimeError("vLLM sampler logits must be [tokens, vocab]") + if source_logits.size(1) % DEFAULT_NUM_VOCAB_TILES: + raise RuntimeError( + "strict vLLM logp requires the padded vocabulary to be divisible by " + f"{DEFAULT_NUM_VOCAB_TILES}" + ) + contract = _full_vocab_contract(source_logits) + dispatch = kernel_registry.get_logprob_op(contract, requested_backend=self.backend_id) + if ( + dispatch.provenance["actual_backend"] != self.backend_id + or dispatch.provenance["fallback"] + ): + raise RuntimeError("strict vLLM Logp dispatch changed backend") + selected, _lse = dispatch.op( + source_logits, + token_ids, + contract=contract, + num_vocab_tiles=DEFAULT_NUM_VOCAB_TILES, + deterministic=True, + ) + real_vocab_env = _optional_env_int("RL_KERNEL_VLLM_REAL_VOCAB_SIZE") + padded_vocab_env = _optional_env_int("RL_KERNEL_VLLM_PADDED_VOCAB_SIZE") + masked_vocab_diagnostic: dict[str, Any] = { + "env_real_vocab_size": real_vocab_env, + "env_padded_vocab_size": padded_vocab_env, + "status": ( + "not_requested" + if real_vocab_env is None and padded_vocab_env is None + else "skipped" + ), + } + diagnostic_contract = _diagnostic_vocab_contract( + source_logits, + real_vocab_size=real_vocab_env, + padded_vocab_size=padded_vocab_env, + ) + if diagnostic_contract is not None: + try: + diagnostic_dispatch = kernel_registry.get_logprob_op( + diagnostic_contract, requested_backend=self.backend_id + ) + masked_selected, _masked_lse = diagnostic_dispatch.op( + source_logits, + token_ids, + contract=diagnostic_contract, + num_vocab_tiles=DEFAULT_NUM_VOCAB_TILES, + deterministic=True, + ) + masked_vocab_diagnostic.update( + { + "status": "computed_not_applied", + "contract_real_vocab_size": diagnostic_contract.sharding.real_vocab_size, + "contract_padded_vocab_size": ( + diagnostic_contract.sharding.padded_vocab_size + ), + "selected_stats": _tensor_debug_stats(masked_selected), + "diff_vs_native_selected": _diff_debug_stats( + masked_selected, + logprobs_tensors.logprobs[ + torch.arange( + logprobs_tensors.logprobs.size(0), + device=logprobs_tensors.logprobs.device, + ), + (logprobs_tensors.logprob_token_ids == token_ids.unsqueeze(1)) + .to(torch.int64) + .argmax(dim=1), + ], + ), + "diff_vs_current_rlkernel_selected": _diff_debug_stats( + masked_selected, selected + ), + } + ) + except Exception as exc: # pragma: no cover - diagnostic only + masked_vocab_diagnostic.update( + {"status": "error", "error": f"{type(exc).__name__}: {exc}"} + ) + elif real_vocab_env is not None or padded_vocab_env is not None: + masked_vocab_diagnostic.update( + { + "source_logits_vocab_width": int(source_logits.size(1)), + "reason": ( + "env metadata must be positive, real<=padded, and " + "padded_vocab_size must match vLLM logits width" + ), + } + ) + + provenance = { + **dict(dispatch.provenance), + "runtime_platform": "cuda", + "triton_used": False, + "source_logits_shape": list(source_logits.shape), + "source_logits_dtype": _dtype_name(source_logits), + "contract_real_vocab_size": contract.sharding.real_vocab_size, + "contract_padded_vocab_size": contract.sharding.padded_vocab_size, + "masked_vocab_diagnostic": masked_vocab_diagnostic, + } + return self._replace_sampled_value( + result, + token_ids=token_ids, + selected=selected, + provenance=provenance, + ) + + +__all__ = [ + "MegatronAttentionOperator", + "MegatronFFNOperator", + "MegatronLogpOperator", + "SemanticOperatorHandle", + "VllmAttentionOperator", + "VllmFFNOperator", + "VllmLogpOperator", +] diff --git a/rl_engine/integrations/linear_logp.py b/rl_engine/integrations/linear_logp.py new file mode 100644 index 00000000..4704403e --- /dev/null +++ b/rl_engine/integrations/linear_logp.py @@ -0,0 +1,368 @@ +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch +from rl_engine.integrations.ablation import operator_ablation_case + + +@dataclass(frozen=True) +class RolloutLinearLogpContext: + """The exact hidden/LM-head contract for the immediately following sample.""" + + hidden: torch.Tensor + lm_head_weight: torch.Tensor + lm_head_bias: torch.Tensor | None + tp_group: Any + vocab_start_index: int + global_vocab_size: int + real_vocab_size: int + + +_ROLLOUT_CONTEXT: RolloutLinearLogpContext | None = None + + +def publish_rollout_linear_logp_context( + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + lm_head_bias: torch.Tensor | None, + *, + tp_group: Any, + vocab_start_index: int, + global_vocab_size: int, + real_vocab_size: int, +) -> None: + """Publish one model-output context for the next vLLM sampler invocation.""" + + global _ROLLOUT_CONTEXT + if hidden.ndim != 2: + raise ValueError( + f"rollout linear_logp hidden must be [tokens, hidden], got {tuple(hidden.shape)}" + ) + if lm_head_weight.ndim != 2: + raise ValueError("rollout linear_logp LM-head weight must be [vocab_local, hidden]") + if hidden.device != lm_head_weight.device: + raise ValueError("rollout linear_logp hidden and LM-head must share a device") + if hidden.size(1) != lm_head_weight.size(1): + raise ValueError("rollout linear_logp hidden and LM-head hidden widths must match") + if lm_head_bias is not None and ( + lm_head_bias.ndim != 1 + or lm_head_bias.size(0) != lm_head_weight.size(0) + or lm_head_bias.device != hidden.device + ): + raise ValueError("rollout linear_logp bias must match the complete local LM-head shard") + if int(global_vocab_size) <= 0 or int(real_vocab_size) <= 0: + raise ValueError("rollout linear_logp vocab sizes must be positive") + if int(real_vocab_size) > int(global_vocab_size): + raise ValueError("rollout linear_logp real vocab cannot exceed padded global vocab") + _ROLLOUT_CONTEXT = RolloutLinearLogpContext( + hidden=hidden, + lm_head_weight=lm_head_weight, + lm_head_bias=lm_head_bias, + tp_group=tp_group, + vocab_start_index=int(vocab_start_index), + global_vocab_size=int(global_vocab_size), + real_vocab_size=int(real_vocab_size), + ) + + +def clear_rollout_linear_logp_context() -> None: + """Discard a context when vLLM computes logits without sampling from them.""" + + global _ROLLOUT_CONTEXT + _ROLLOUT_CONTEXT = None + + +def take_rollout_linear_logp_context() -> RolloutLinearLogpContext: + """Consume the context published by the matching ``compute_logits`` call.""" + + global _ROLLOUT_CONTEXT + context = _ROLLOUT_CONTEXT + _ROLLOUT_CONTEXT = None + if context is None: + raise RuntimeError( + "strict rollout linear_logp sampler has no matching hidden/LM-head context" + ) + return context + + +class LinearLogpWrapper: + """PR230 integration adapter for the existing TP-aware linear_logp op.""" + + # This is deliberately a different backend from the existing high- + # performance linear_logp op. PR1 exposes these strict entry points + # without changing the old registry route. + backend_id = "rlkernel.linear_logp.bitwise.v1" + + def __init__(self) -> None: + self._op: Any | None = None + self._tp_op: Any | None = None + self._last_provenance: dict[str, Any] = {} + + @property + def provenance(self) -> Mapping[str, Any]: + return dict(self._last_provenance) + + def _resolve(self, *, tensor_parallel: bool) -> Any: + if tensor_parallel and self._tp_op is not None: + return self._tp_op + if not tensor_parallel and self._op is not None: + return self._op + try: + from rl_engine.kernels.ops.cuda.loss.linear_logp import ( + sm90_deterministic_linear_logp, + sm90_deterministic_linear_logp_tp, + ) + except (ImportError, AttributeError) as exc: + raise RuntimeError( + "strict bitwise linear_logp requires PR1's separate " + "sm90_deterministic_linear_logp entry points" + ) from exc + op = ( + sm90_deterministic_linear_logp_tp if tensor_parallel else sm90_deterministic_linear_logp + ) + if tensor_parallel: + self._tp_op = op + else: + self._op = op + return op + + @staticmethod + def _tp_coordinates(tp_group: Any) -> tuple[int, int]: + if tp_group is None: + return 0, 1 + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + raise RuntimeError("strict TP linear_logp requires initialized torch.distributed") + return ( + int(torch.distributed.get_rank(group=tp_group)), + int(torch.distributed.get_world_size(group=tp_group)), + ) + + @staticmethod + def _temperature_tensor( + temperature: float | torch.Tensor | None, *, rows: int, device: torch.device + ) -> torch.Tensor | None: + if temperature is None: + return None + if isinstance(temperature, torch.Tensor): + value = temperature.to(device=device, dtype=torch.float32).reshape(-1) + if value.numel() == 1: + value = value.expand(rows).contiguous() + elif value.numel() == rows: + value = value.contiguous() + else: + raise ValueError( + f"temperature must be scalar or one value per row, got {value.numel()}" + ) + if bool((value <= 0).any().item()): + raise ValueError("temperature must be positive") + return value + value = float(temperature) + if value <= 0.0: + raise ValueError(f"temperature must be positive, got {value}") + return torch.full((rows,), value, device=device, dtype=torch.float32) + + @staticmethod + def _validate_targets(target_ids: torch.Tensor, *, rows: int, real_vocab_size: int) -> None: + if target_ids.ndim != 1 or target_ids.numel() != rows: + raise ValueError( + f"linear_logp target_ids must be [rows]={rows}, got {tuple(target_ids.shape)}" + ) + if target_ids.is_floating_point() or target_ids.is_complex(): + raise TypeError("linear_logp target_ids must use an integer dtype") + if target_ids.numel(): + target_long = target_ids.to(dtype=torch.long) + if bool(((target_long < 0) | (target_long >= real_vocab_size)).any().item()): + raise ValueError( + f"linear_logp target_ids must be in [0, {real_vocab_size}), got invalid ids" + ) + + @classmethod + def _validate_contract( + cls, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: torch.Tensor | None, + *, + tp_group: Any, + vocab_start_index: int, + global_vocab_size: int | None, + real_vocab_size: int | None, + ) -> tuple[bool, int, int]: + if hidden.ndim != 2: + raise ValueError( + f"linear_logp wrapper expects hidden [tokens, hidden], got {tuple(hidden.shape)}" + ) + if not hidden.is_cuda or torch.version.hip is not None: + raise RuntimeError("strict linear_logp requires NVIDIA CUDA tensors") + if lm_head_weight.ndim != 2: + raise ValueError("linear_logp LM-head weight must be [vocab_local, hidden]") + if hidden.dtype != torch.bfloat16 or lm_head_weight.dtype != torch.bfloat16: + raise TypeError("strict SM90 linear_logp requires bfloat16 hidden and LM-head") + if lm_head_weight.device != hidden.device: + raise ValueError("linear_logp hidden and LM-head must share a device") + if hidden.size(1) != lm_head_weight.size(1): + raise ValueError("linear_logp hidden width must match LM-head width") + if bias is not None and ( + bias.ndim != 1 or bias.size(0) != lm_head_weight.size(0) or bias.device != hidden.device + ): + raise ValueError("linear_logp bias must match the complete local LM-head shard") + + local_vocab = int(lm_head_weight.size(0)) + rank, world = cls._tp_coordinates(tp_group) + tensor_parallel = tp_group is not None or int(vocab_start_index) != 0 + requested_global = ( + local_vocab * world + if global_vocab_size is None and tensor_parallel + else local_vocab if global_vocab_size is None else int(global_vocab_size) + ) + if requested_global <= 0: + raise ValueError("linear_logp global_vocab_size must be positive") + if tensor_parallel: + if tp_group is None: + raise ValueError("TP linear_logp requires an explicit TP process group") + if requested_global != local_vocab * world: + raise ValueError( + "TP linear_logp requires complete equal shards: " + f"local={local_vocab}, world={world}, global={requested_global}" + ) + expected_start = rank * local_vocab + if int(vocab_start_index) != expected_start: + raise ValueError( + "TP linear_logp vocab_start_index does not match rank-local shard: " + f"got {vocab_start_index}, expected {expected_start}" + ) + elif requested_global != local_vocab or int(vocab_start_index) != 0: + raise ValueError("non-TP linear_logp must use the complete local vocab at offset zero") + real = requested_global if real_vocab_size is None else int(real_vocab_size) + if real <= 0 or real > requested_global: + raise ValueError( + f"linear_logp real_vocab_size must be in [1, {requested_global}], got {real}" + ) + cls._validate_targets(target_ids, rows=hidden.size(0), real_vocab_size=real) + return tensor_parallel, requested_global, real + + @staticmethod + def _mismatch_provenance() -> dict[str, Any]: + case_id = os.getenv("RL_KERNEL_LOGP_CASE", "P/P").strip().upper() + case = operator_ablation_case("logp", case_id) + cross_engine_mismatch = case.training is not case.rollout + return { + "module": "logp", + "case_id": case.case_id, + "training_implementation": case.training.value, + "rollout_implementation": case.rollout.value, + "mismatch_axes": { + # PR230 L1-L3 axes. R/R closes all of them by sharing the + # padded TP LM-head and the rank-ordered LSE merge. + "vocab_shard_ownership": cross_engine_mismatch, + "selected_token_identity": cross_engine_mismatch, + "vocab_lse_reduction": cross_engine_mismatch, + }, + "arithmetic_contract": { + "batch_invariant": True, + "fixed_k_reduction": True, + "vocab_group_width": 64, + "summary_merge_tree": "rank_ordered_fixed", + "vocab_reduction_axis": "TP", + "cp_is_merge_axis": False, + "output_dtype": "fp32", + }, + } + + def __call__( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: torch.Tensor | None = None, + *, + tp_group: Any = None, + vocab_start_index: int = 0, + global_vocab_size: int | None = None, + real_vocab_size: int | None = None, + target: str = "training", + temperature: float | torch.Tensor | None = None, + ) -> torch.Tensor: + if target_ids.device != hidden.device: + raise ValueError("linear_logp target_ids must share hidden device") + tensor_parallel, requested_global_vocab, real = self._validate_contract( + hidden, + lm_head_weight, + target_ids, + bias, + tp_group=tp_group, + vocab_start_index=int(vocab_start_index), + global_vocab_size=global_vocab_size, + real_vocab_size=real_vocab_size, + ) + effective_weight = lm_head_weight + effective_bias = bias + temperature_tensor = self._temperature_tensor( + temperature, rows=hidden.size(0), device=hidden.device + ) + op = self._resolve(tensor_parallel=tensor_parallel) + if tensor_parallel: + result, _lse = op( + hidden, + effective_weight, + target_ids, + effective_bias, + tp_group=tp_group, + vocab_start_index=int(vocab_start_index), + global_vocab_size=requested_global_vocab, + real_vocab_size=real, + temperature=temperature_tensor, + ) + else: + result, _lse = op( + hidden, + effective_weight, + target_ids, + effective_bias, + real_vocab_size=real, + temperature=temperature_tensor, + ) + if not isinstance(result, torch.Tensor): + raise RuntimeError("linear_logp backend returned a non-tensor result") + if result.numel() != hidden.size(0): + raise RuntimeError( + "linear_logp backend returned an unexpected token count: " + f"{result.numel()} != {hidden.size(0)}" + ) + + actual_backend = self.backend_id + self._last_provenance = { + **self._mismatch_provenance(), + "target": target, + "runtime_platform": "cuda", + "triton_used": False, + "actual_backend": actual_backend, + "hidden_shape": list(hidden.shape), + "hidden_dtype": str(hidden.dtype).replace("torch.", ""), + "lm_head_weight_shape": list(lm_head_weight.shape), + "target_shape": list(target_ids.shape), + "tp_group_present": tp_group is not None, + "vocab_start_index": int(vocab_start_index), + "global_vocab_size": requested_global_vocab, + "requested_global_vocab_size": requested_global_vocab, + "real_vocab_size": real, + "requested_real_vocab_size": None if real_vocab_size is None else int(real_vocab_size), + "temperature": None if temperature is None else "provided", + "contract_version": "cuda-det-gemm-linear-logp-sm90-contract-v2", + "logits_materialized": True, + } + return result + + +__all__ = [ + "LinearLogpWrapper", + "RolloutLinearLogpContext", + "clear_rollout_linear_logp_context", + "publish_rollout_linear_logp_context", + "take_rollout_linear_logp_context", +] diff --git a/rl_engine/integrations/megatron_runtime.py b/rl_engine/integrations/megatron_runtime.py new file mode 100644 index 00000000..551c314f --- /dev/null +++ b/rl_engine/integrations/megatron_runtime.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Megatron runtime hooks installed without editing Megatron source files.""" + +from __future__ import annotations + +import importlib +from collections.abc import Iterable +from typing import Any + +from rl_engine.integrations.ablation import Implementation, IntegrationPlan +from rl_engine.integrations.framework_operators import ( + MegatronAttentionOperator, + MegatronFFNOperator, + MegatronLogpOperator, +) +from rl_engine.integrations.linear_logp import LinearLogpWrapper +from rl_engine.integrations.megatron import MegatronIntegration +from rl_engine.integrations.state import get_active_integration, set_active_integration +from rl_engine.integrations.vime.logp import _provider_impl + +_PATCH_MARKER = "__rl_kernel_original_forward__" + + +def _optional_class(path: str) -> type[Any] | None: + module_name, _, class_name = path.rpartition(".") + try: + value = getattr(importlib.import_module(module_name), class_name) + except (ImportError, AttributeError): + return None + return value if isinstance(value, type) else None + + +def _discover_attention_classes() -> tuple[type[Any], ...]: + paths = ( + "megatron.core.transformer.dot_product_attention.DotProductAttention", + "megatron.core.extensions.transformer_engine.TEDotProductAttention", + "megatron.core.extensions.transformer_engine.TEDotProductAttentionWithCP", + ) + return tuple(value for path in paths if (value := _optional_class(path)) is not None) + + +def _discover_ffn_classes() -> tuple[type[Any], ...]: + paths = ("megatron.core.transformer.mlp.MLP",) + return tuple(value for path in paths if (value := _optional_class(path)) is not None) + + +def _unique_classes(values: Iterable[type[Any]]) -> tuple[type[Any], ...]: + return tuple(dict.fromkeys(values)) + + +def _patch_forward( + cls: type[Any], + *, + integration: MegatronIntegration, + module: str, +) -> None: + if hasattr(cls, _PATCH_MARKER): + raise RuntimeError(f"{cls.__module__}.{cls.__name__} is already RL-Kernel patched") + original = cls.forward + + def wrapped(instance: Any, *args: Any, **kwargs: Any) -> Any: + def native(_module: Any, *call_args: Any, **call_kwargs: Any) -> Any: + return original(instance, *call_args, **call_kwargs) + + return integration.execute(module, native, instance, *args, **kwargs) + + wrapped.__name__ = getattr(original, "__name__", "forward") + wrapped.__doc__ = getattr(original, "__doc__", None) + setattr(cls, _PATCH_MARKER, original) + cls.forward = wrapped + + +def install_megatron_integration( + plan: IntegrationPlan, + *, + attention_classes: Iterable[type[Any]] | None = None, + ffn_classes: Iterable[type[Any]] | None = None, +) -> MegatronIntegration: + """Install Attention, dense FFN and structural Logp routes in one actor.""" + + existing = get_active_integration("megatron") + if existing is not None: + if not isinstance(existing, MegatronIntegration): + raise RuntimeError("active Megatron integration has an unexpected type") + if existing.plan != plan: + raise RuntimeError("Megatron integration is already installed with another plan") + return existing + + integration = MegatronIntegration( + plan, + rl_kernel_operators={ + "attention": MegatronAttentionOperator(), + "ffn": MegatronFFNOperator(), + "logp": MegatronLogpOperator( + _provider_impl, + linear_logp=LinearLogpWrapper(), + ), + }, + ) + resolved_attention = _unique_classes( + _discover_attention_classes() if attention_classes is None else attention_classes + ) + resolved_ffn = _unique_classes(_discover_ffn_classes() if ffn_classes is None else ffn_classes) + if ( + plan.implementation_for("attention", "training") is Implementation.RL_KERNEL + and not resolved_attention + ): + raise RuntimeError("R/R Megatron Attention selected but no supported class was found") + if plan.implementation_for("ffn", "training") is Implementation.RL_KERNEL and not resolved_ffn: + raise RuntimeError("R/R Megatron FFN selected but no supported class was found") + + set_active_integration("megatron", integration) + for cls in resolved_attention: + _patch_forward(cls, integration=integration, module="attention") + if resolved_attention: + integration.record_installed_hook( + "attention", + ",".join(f"{cls.__module__}.{cls.__name__}.forward" for cls in resolved_attention), + ) + for cls in resolved_ffn: + _patch_forward(cls, integration=integration, module="ffn") + if resolved_ffn: + integration.record_installed_hook( + "ffn", + ",".join(f"{cls.__module__}.{cls.__name__}.forward" for cls in resolved_ffn), + ) + integration.record_installed_hook("logp", "rl_engine.integrations.vime.logp.provider") + return integration + + +def initialize_from_environment(_args: Any = None) -> MegatronIntegration: + """Vime-compatible custom-init entry point backed by the shared plan env.""" + + from rl_engine.integrations.ablation import integration_plan_from_environment + + return install_megatron_integration(integration_plan_from_environment()) + + +__all__ = ["initialize_from_environment", "install_megatron_integration"] diff --git a/rl_engine/integrations/runtime.py b/rl_engine/integrations/runtime.py index 7ede1224..9dce3da7 100644 --- a/rl_engine/integrations/runtime.py +++ b/rl_engine/integrations/runtime.py @@ -5,12 +5,16 @@ from __future__ import annotations +import json +import os from collections import Counter from collections.abc import Callable, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field +from pathlib import Path from threading import Lock from typing import Any +import torch from rl_engine.integrations.ablation import Implementation, IntegrationPlan @@ -23,9 +27,12 @@ class OperatorReadback: implementation: str backend_id: str call_count: int + provenance: Mapping[str, Any] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: - return self.__dict__.copy() + payload = self.__dict__.copy() + payload["provenance"] = dict(self.provenance) + return payload class FrameworkOperatorIntegration: @@ -45,8 +52,37 @@ def __init__( self._rl_kernel_operators = dict(rl_kernel_operators) self._counts: Counter[str] = Counter() self._readbacks: dict[str, OperatorReadback] = {} + self._installed_hooks: dict[str, str] = {} + self._fallbacks: list[dict[str, str]] = [] self._lock = Lock() + def install_operator(self, module: str, operator: Callable[..., Any]) -> None: + normalized = module.strip().lower() + if normalized not in self.plan.cases: + raise ValueError(f"unknown integration module {module!r}") + if not callable(operator): + raise TypeError("operator must be callable") + with self._lock: + self._rl_kernel_operators[normalized] = operator + + def record_installed_hook(self, module: str, hook_id: str) -> None: + normalized = module.strip().lower() + if normalized not in self.plan.cases: + raise ValueError(f"unknown integration module {module!r}") + if not isinstance(hook_id, str) or not hook_id.strip(): + raise ValueError("hook_id must be a non-empty string") + with self._lock: + self._installed_hooks[normalized] = hook_id.strip() + self._persist_readback() + + def record_fallback(self, module: str, reason: str) -> None: + normalized = module.strip().lower() + if normalized not in self.plan.cases: + raise ValueError(f"unknown integration module {module!r}") + with self._lock: + self._fallbacks.append({"module": normalized, "reason": str(reason)}) + self._persist_readback() + def execute( self, module: str, @@ -70,6 +106,12 @@ def execute( ) selected = rl_kernel_operator result = selected(*args, **kwargs) + # vLLM's model forward is captured with a Dynamo fullgraph. Readback + # bookkeeping performs Python locking and JSON I/O, so defer it until + # the non-compiled execution path while keeping the selected operator + # inside the captured graph. + if torch._dynamo.is_compiling(): + return result backend_id = getattr(selected, "backend_id", None) if not isinstance(backend_id, str) or not backend_id.strip(): backend_id = ( @@ -77,6 +119,8 @@ def execute( if implementation is Implementation.PRODUCTION else f"rlkernel.{normalized}.unidentified" ) + raw_provenance = getattr(selected, "provenance", {}) + provenance = dict(raw_provenance) if isinstance(raw_provenance, Mapping) else {} with self._lock: self._counts[normalized] += 1 case = self.plan.cases[normalized] @@ -88,7 +132,9 @@ def execute( implementation=implementation.value, backend_id=backend_id, call_count=self._counts[normalized], + provenance=provenance, ) + self._persist_readback() return result def readback(self) -> dict[str, Any]: @@ -97,10 +143,96 @@ def readback(self) -> dict[str, Any]: "framework": self.framework, "target": self.target, "plan": self.plan.to_dict(), + "installed_hooks": dict(self._installed_hooks), + "fallbacks": list(self._fallbacks), "operators": { module: readback.to_dict() for module, readback in self._readbacks.items() }, } + def assert_strict_ready(self) -> None: + """Fail unless every selected RL-Kernel route was installed and executed.""" + + payload = self.readback() + missing_hooks: list[str] = [] + missing_calls: list[str] = [] + wrong_backends: list[str] = [] + for module in self.plan.cases: + selected = self.plan.implementation_for(module, self.target) + if selected is not Implementation.RL_KERNEL: + continue + if module not in payload["installed_hooks"]: + missing_hooks.append(module) + operator = payload["operators"].get(module) + if operator is None or int(operator["call_count"]) <= 0: + missing_calls.append(module) + elif not str(operator["backend_id"]).startswith( + ("rlkernel.", "pytorch-vocab-parallel-logp") + ): + wrong_backends.append(f"{module}={operator['backend_id']}") + elif _contains_triton(operator): + wrong_backends.append(f"{module}=triton") + elif _runtime_platform(operator.get("provenance")) != "cuda": + wrong_backends.append(f"{module}=non-cuda") + failures: list[str] = [] + if missing_hooks: + failures.append("missing hooks: " + ", ".join(missing_hooks)) + if missing_calls: + failures.append("zero calls: " + ", ".join(missing_calls)) + if wrong_backends: + failures.append("unexpected backends: " + ", ".join(wrong_backends)) + if payload["fallbacks"]: + failures.append(f"fallbacks: {payload['fallbacks']}") + if failures: + raise RuntimeError( + f"{self.framework} {self.target} integration is not strict-ready: " + + "; ".join(failures) + ) + + def _persist_readback(self) -> None: + directory = os.getenv("RL_KERNEL_READBACK_DIR", "").strip() + if not directory: + return + target = Path(directory) + target.mkdir(parents=True, exist_ok=True) + path = target / f"{self.framework}-{self.target}-{os.getpid()}.json" + temporary = path.with_suffix(".tmp") + temporary.write_text( + json.dumps(self.readback(), indent=2, sort_keys=True), + encoding="utf-8", + ) + temporary.replace(path) + + +def _contains_triton(value: Any) -> bool: + if isinstance(value, str): + return "triton" in value.lower() + if isinstance(value, Mapping): + for key, item in value.items(): + normalized_key = str(key).strip().lower() + if normalized_key in {"triton_used", "uses_triton"}: + if item is True: + return True + continue + if _contains_triton(item): + return True + return False + if isinstance(value, (list, tuple)): + return any(_contains_triton(item) for item in value) + return False + + +def _runtime_platform(value: Any) -> str | None: + if not isinstance(value, Mapping): + return None + direct = value.get("runtime_platform") + if isinstance(direct, str): + return direct + for item in value.values(): + nested = _runtime_platform(item) + if nested is not None: + return nested + return None + __all__ = ["FrameworkOperatorIntegration", "OperatorReadback"] diff --git a/rl_engine/integrations/state.py b/rl_engine/integrations/state.py new file mode 100644 index 00000000..1d5002e0 --- /dev/null +++ b/rl_engine/integrations/state.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Process-local framework integration state. + +Ray actors and vLLM workers are separate processes, so each process owns one +integration object and emits its own readback file. Keeping the registry here +also lets the Vime logprob provider join the same Megatron call accounting as +the Attention and FFN hooks. +""" + +from __future__ import annotations + +from threading import Lock +from typing import Literal + +from rl_engine.integrations.runtime import FrameworkOperatorIntegration + +FrameworkName = Literal["megatron", "vllm"] + +_ACTIVE: dict[FrameworkName, FrameworkOperatorIntegration] = {} +_LOCK = Lock() + + +def set_active_integration( + framework: FrameworkName, + integration: FrameworkOperatorIntegration, +) -> None: + if integration.framework != framework: + raise ValueError( + f"integration framework {integration.framework!r} does not match {framework!r}" + ) + with _LOCK: + existing = _ACTIVE.get(framework) + if existing is not None and existing is not integration: + raise RuntimeError(f"{framework} integration is already installed in this process") + _ACTIVE[framework] = integration + + +def get_active_integration( + framework: FrameworkName, +) -> FrameworkOperatorIntegration | None: + with _LOCK: + return _ACTIVE.get(framework) + + +def clear_active_integration(framework: FrameworkName) -> None: + with _LOCK: + _ACTIVE.pop(framework, None) + + +__all__ = [ + "FrameworkName", + "clear_active_integration", + "get_active_integration", + "set_active_integration", +] diff --git a/rl_engine/integrations/vime/logp.py b/rl_engine/integrations/vime/logp.py index ef1b6d19..ec44cdc8 100644 --- a/rl_engine/integrations/vime/logp.py +++ b/rl_engine/integrations/vime/logp.py @@ -18,7 +18,7 @@ from typing import Any import torch - +from rl_engine.integrations.ablation import Implementation, operator_ablation_case from rl_engine.kernels.logprob_contract import ( LogprobContract, LogprobDType, @@ -55,6 +55,18 @@ class ProviderResult: provenance: Mapping[str, Any] +_DEFAULT_STRICT_LINEAR_LOGP: Any | None = None + + +def _default_strict_linear_logp() -> Any: + global _DEFAULT_STRICT_LINEAR_LOGP + if _DEFAULT_STRICT_LINEAR_LOGP is None: + from rl_engine.integrations.linear_logp import LinearLogpWrapper + + _DEFAULT_STRICT_LINEAR_LOGP = LinearLogpWrapper() + return _DEFAULT_STRICT_LINEAR_LOGP + + def _as_positive_int(value: Any, name: str) -> int: if isinstance(value, bool) or not isinstance(value, int) or value <= 0: raise SelectedLogprobProviderUnavailable( @@ -189,7 +201,7 @@ def _contract_for_request(request: Any) -> tuple[LogprobContract, int]: return contract, _tile_count(metadata, padded_vocab_size) -def provider(request: Any) -> ProviderResult: +def _provider_impl(request: Any, *, linear_logp: Any = None) -> ProviderResult: """Compute Vime selected logprobs on the explicit WS2 TP/CP contract. Top-p replay is deliberately unavailable until it has a separately @@ -198,6 +210,104 @@ def provider(request: Any) -> ProviderResult: distribution semantics. """ + strict = os.getenv("VIME_RL_KERNEL_STRICT", "").strip().lower() in {"1", "true", "yes", "on"} + if strict and not isinstance(getattr(request, "hidden", None), torch.Tensor): + raise RuntimeError( + "strict Vime linear_logp request is missing hidden/LM-head structural inputs" + ) + 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(getattr(request, "hidden", None), torch.Tensor) + ): + linear_logp = _default_strict_linear_logp() + if linear_logp is not None and isinstance(getattr(request, "hidden", None), torch.Tensor): + hidden = request.hidden + weight = getattr(request, "lm_head_weight", None) + if not isinstance(weight, torch.Tensor): + raise RuntimeError("linear_logp request must expose lm_head_weight") + selected = linear_logp( + hidden, + weight, + request.target_ids, + getattr(request, "lm_head_bias", None), + tp_group=getattr(request, "tensor_parallel_group", None), + vocab_start_index=int(getattr(request, "vocab_start_index", 0)), + global_vocab_size=getattr(request, "global_vocab_size", None), + real_vocab_size=getattr(request, "metadata", {}).get("real_vocab_size"), + temperature=getattr(request, "temperature", None), + ) + entropy = None + entropy_provenance: dict[str, Any] = {} + if getattr(request, "with_entropy", False): + # Entropy is a separate loss metric from selected logprob. Keep + # the strict selected path on linear_logp, while using the + # explicit TP vocab reduction for the full-vocabulary entropy + # requested by Vime's policy loss. + entropy_contract, entropy_tiles = _contract_for_request(request) + entropy_dispatch = kernel_registry.get_logprob_op( + entropy_contract, requested_backend=BACKEND_ID + ) + if ( + entropy_dispatch.provenance["actual_backend"] != BACKEND_ID + or entropy_dispatch.provenance["fallback"] + ): + raise RuntimeError( + "explicit WS2 entropy dispatch changed during strict linear_logp execution" + ) + _, _, entropy = entropy_dispatch.op.apply_with_entropy( + request.logits, + request.target_ids, + contract=entropy_contract, + tp_group=getattr(request, "tensor_parallel_group", None), + num_vocab_tiles=entropy_tiles, + with_entropy_grad=bool(getattr(request, "with_entropy_grad", False)), + ) + entropy_provenance = { + "backend_id": entropy_dispatch.capability.backend_id, + "contract_id": entropy_contract.cross_rank_fingerprint(), + "num_vocab_tiles": entropy_tiles, + "logits_materialized": True, + } + provenance = dict(getattr(linear_logp, "provenance", {})) + provenance.update( + { + "execution": { + "role": "vime_training_linear_logprob", + "strict_backend": True, + "top_p_replay": False, + "cp_is_merge_axis": False, + "logits_materialized": bool(getattr(request, "with_entropy", False)), + "entropy": entropy_provenance, + }, + "request": { + "hidden_shape": list(hidden.shape), + "hidden_dtype": str(hidden.dtype).replace("torch.", ""), + "target_shape": list(request.target_ids.shape), + "tp_world_size": int(getattr(request, "metadata", {}).get("tp_world_size", 1)), + "tp_rank": int(getattr(request, "metadata", {}).get("tp_rank", 0)), + "cp_world_size": int(getattr(request, "context_parallel", None).world_size), + "cp_rank": int(getattr(request, "context_parallel", None).rank), + }, + } + ) + backend_id = str(provenance.get("actual_backend", linear_logp.backend_id)) + contract_id = ( + "linear_logp:" + f"tp={getattr(request, 'metadata', {}).get('tp_world_size', 1)}:" + f"cp={getattr(request, 'context_parallel', None).world_size}:" + f"vocab={getattr(request, 'global_vocab_size', None)}" + ) + return ProviderResult( + selected_logprobs=selected.reshape(-1, 1), + entropy=entropy, + backend_id=backend_id, + contract_id=contract_id, + provenance=provenance, + ) + if getattr(request, "log_prob_keep_mask", None) is not None: raise SelectedLogprobProviderUnavailable( "RL-Kernel WS2 logprob does not yet materialize Vime top-p replay masks" @@ -246,7 +356,7 @@ def provider(request: Any) -> ProviderResult: provenance["cp_row_ownership"] = { "cp_rank": contract.sharding.cp_rank, "cp_world_size": contract.sharding.cp_world_size, - "layout": getattr(request.context_parallel, "layout"), + "layout": request.context_parallel.layout, "local_token_rows": int(request.logits.shape[0]), "cp_is_merge_axis": False, } @@ -260,4 +370,31 @@ def provider(request: Any) -> ProviderResult: ) +_provider_impl.backend_id = BACKEND_ID # type: ignore[attr-defined] + + +def provider(request: Any) -> ProviderResult: + """Route Vime training logp through the active Megatron integration.""" + + # PR230's P/R axis is selected independently for training and rollout. + # The Megatron provider is only the training boundary, so a production + # training side must bypass the RL-Kernel integration entirely. + case = operator_ablation_case("logp", os.getenv("RL_KERNEL_LOGP_CASE", "P/P")) + if case.training is Implementation.PRODUCTION: + return _provider_impl(request) + + from rl_engine.integrations.state import get_active_integration + + integration = get_active_integration("megatron") + if integration is None: + return _provider_impl(request) + + def native_unavailable(_request: Any) -> ProviderResult: + raise RuntimeError( + "the structural provider was invoked for a production Megatron logp route" + ) + + return integration.execute("logp", native_unavailable, request) + + __all__ = ["ProviderResult", "SelectedLogprobProviderUnavailable", "provider"] diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py new file mode 100644 index 00000000..d4eda2e9 --- /dev/null +++ b/rl_engine/integrations/vllm_runtime.py @@ -0,0 +1,398 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""vLLM plugin hooks installed without editing vLLM source files.""" + +from __future__ import annotations + +import os +from typing import Any + +import torch +from rl_engine.integrations.ablation import ( + Implementation, + IntegrationPlan, + configure_integration_environment, + integration_plan_from_environment, +) +from rl_engine.integrations.framework_operators import ( + VllmAttentionOperator, + VllmFFNOperator, + VllmLogpOperator, +) +from rl_engine.integrations.linear_logp import ( + clear_rollout_linear_logp_context, + publish_rollout_linear_logp_context, +) +from rl_engine.integrations.state import get_active_integration, set_active_integration +from rl_engine.integrations.vllm import VllmIntegration + +_PATCH_MARKER = "__rl_kernel_original_forward__" +_RLK_ATTENTION_BACKEND: type[Any] | None = None +_RLK_ATTENTION_IMPL: type[Any] | None = None + + +def _is_worker_sampler_profile_batch(value: Any) -> bool: + """Return true for vLLM's KV-cache profiling dummy sampler batch.""" + + req_ids = getattr(value, "req_ids", None) + is_padding = getattr(value, "is_padding", None) + if not isinstance(req_ids, list) or not req_ids: + return False + if not all(isinstance(item, str) and item.startswith("req_") for item in req_ids): + return False + if is_padding is None or not hasattr(is_padding, "all"): + return False + try: + return bool(is_padding.all().item()) + except (AttributeError, RuntimeError, TypeError): + return False + + +def _is_v1_sampler_profile_batch(value: Any) -> bool: + """Identify the no-logprob dummy batch used to size vLLM's KV cache.""" + + if getattr(value, "max_num_logprobs", object()) is not None: + return False + output_token_ids = getattr(value, "output_token_ids", None) + spec_token_ids = getattr(value, "spec_token_ids", None) + if not ( + isinstance(output_token_ids, list) + and output_token_ids + and all(isinstance(ids, list) and not ids for ids in output_token_ids) + and isinstance(spec_token_ids, list) + and len(spec_token_ids) == len(output_token_ids) + and all(isinstance(ids, list) and not ids for ids in spec_token_ids) + ): + return False + return ( + bool(getattr(value, "no_penalties", False)) + and getattr(value, "prompt_token_ids", None) is None + and getattr(value, "allowed_token_ids_mask", None) is None + and getattr(value, "bad_words_token_ids", None) == {} + and getattr(value, "generators", None) == {} + ) + + +def plan_from_environment() -> IntegrationPlan: + """Compatibility alias for the shared process plan loader.""" + + return integration_plan_from_environment() + + +def configure_vllm_environment(plan: IntegrationPlan, *, readback_dir: str | None = None) -> None: + """Export the plan inherited by Vime's vLLM server subprocess.""" + + os.environ["RL_KERNEL_VLLM_INTEGRATION"] = "1" + configure_integration_environment(plan, readback_dir=readback_dir) + if plan.implementation_for("attention", "rollout") is Implementation.RL_KERNEL: + os.environ["VLLM_ATTENTION_BACKEND"] = "FLASH_ATTN" + if plan.implementation_for("logp", "rollout") is Implementation.RL_KERNEL: + real_vocab = os.getenv("RL_KERNEL_VLLM_REAL_VOCAB_SIZE", "").strip() + padded_vocab = os.getenv("RL_KERNEL_VLLM_PADDED_VOCAB_SIZE", "").strip() + if not real_vocab or not padded_vocab: + raise RuntimeError( + "strict rollout linear_logp requires " + "RL_KERNEL_VLLM_REAL_VOCAB_SIZE and " + "RL_KERNEL_VLLM_PADDED_VOCAB_SIZE" + ) + + +def _patch_qwen_lm_head_padding() -> None: + """Make vLLM's TP LM-head partition identical to Megatron's padded layout.""" + + from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, + ) + + if hasattr(ParallelLMHead, _PATCH_MARKER): + return + real_vocab = int(os.environ["RL_KERNEL_VLLM_REAL_VOCAB_SIZE"]) + padded_vocab = int(os.environ["RL_KERNEL_VLLM_PADDED_VOCAB_SIZE"]) + if padded_vocab <= real_vocab or padded_vocab % 2: + raise RuntimeError( + f"invalid strict rollout vocab contract: real={real_vocab}, padded={padded_vocab}" + ) + padding_size = padded_vocab - real_vocab + original = ParallelLMHead.__init__ + original_weight_loader = VocabParallelEmbedding.weight_loader + + def strict_weight_loader(instance: Any, param: Any, loaded_weight: torch.Tensor) -> None: + strict_real_vocab = getattr(instance, "_rl_kernel_real_vocab_size", None) + output_dim = getattr(param, "output_dim", None) + packed_dim = getattr(param, "packed_dim", None) + if ( + strict_real_vocab is not None + and output_dim is not None + and packed_dim is None + and loaded_weight.shape[output_dim] == int(strict_real_vocab) + ): + start_idx = int(instance.shard_indices.org_vocab_start_index) + shard_size = int(instance.shard_indices.org_vocab_end_index - start_idx) + available = max(0, min(shard_size, int(strict_real_vocab) - start_idx)) + if available: + loaded_slice = loaded_weight.narrow(output_dim, start_idx, available) + param[:available].data.copy_(loaded_slice) + if available < shard_size: + param[available:shard_size].data.fill_(0) + if shard_size < param.shape[0]: + param[shard_size:].data.fill_(0) + return + original_weight_loader(instance, param, loaded_weight) + + def wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: + if args: + args = (padded_vocab, *args[1:]) + else: + kwargs["num_embeddings"] = padded_vocab + # vLLM normally stores original-vocab and added-vocab rows as separate + # segments. Strict R/R needs Megatron's single contiguous padded vocab. + kwargs["org_num_embeddings"] = padded_vocab + kwargs["padding_size"] = padding_size + original(instance, *args, **kwargs) + instance._rl_kernel_real_vocab_size = real_vocab + if ( + int(getattr(instance, "org_vocab_size", -1)) != padded_vocab + or int(getattr(instance, "num_embeddings_padded", -1)) != padded_vocab + ): + raise RuntimeError( + "strict rollout LM-head padding did not produce the Megatron " + f"padded vocab layout {padded_vocab}" + ) + + setattr(ParallelLMHead, _PATCH_MARKER, original) + if not hasattr(VocabParallelEmbedding, "__rl_kernel_original_weight_loader__"): + VocabParallelEmbedding.__rl_kernel_original_weight_loader__ = original_weight_loader + VocabParallelEmbedding.weight_loader = strict_weight_loader + ParallelLMHead.__init__ = wrapped + + +def _patch_qwen_compute_logits(integration: VllmIntegration) -> None: + """Publish the exact hidden and padded LM-head shard to the sampler.""" + + from vllm.distributed import get_pp_group, get_tp_group + from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM + from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM + + classes = tuple(dict.fromkeys((Qwen3ForCausalLM, Qwen2ForCausalLM))) + installed: list[str] = [] + for cls in classes: + if hasattr(cls, _PATCH_MARKER): + continue + original = cls.compute_logits + + def wrapped( + instance: Any, + hidden_states: Any, + *, + _original: Any = original, + ) -> Any: + logits = _original(instance, hidden_states) + if not get_pp_group().is_last_rank: + return logits + lm_head = getattr(instance, "lm_head", None) + shard_indices = getattr(lm_head, "shard_indices", None) + weight = getattr(lm_head, "weight", None) + if lm_head is None or shard_indices is None or not isinstance(weight, torch.Tensor): + raise RuntimeError("strict rollout linear_logp requires a real Qwen LM-head shard") + tp = get_tp_group() + tp_world = int(tp.world_size) + publish_rollout_linear_logp_context( + hidden_states, + weight, + getattr(lm_head, "bias", None), + tp_group=tp.device_group if tp_world > 1 else None, + vocab_start_index=int(shard_indices.padded_org_vocab_start_index), + global_vocab_size=int(lm_head.num_embeddings_padded), + real_vocab_size=int( + getattr(lm_head, "_rl_kernel_real_vocab_size", lm_head.org_vocab_size) + ), + ) + return logits + + setattr(cls, _PATCH_MARKER, original) + cls.compute_logits = wrapped + installed.append(f"{cls.__module__}.{cls.__name__}.compute_logits") + if installed: + integration.record_installed_hook("logp", ",".join(installed)) + + +def _patch_qwen_ffn(integration: VllmIntegration) -> None: + from vllm.model_executor.models.qwen2 import Qwen2MLP + + integration.install_operator("ffn", VllmFFNOperator()) + if hasattr(Qwen2MLP, _PATCH_MARKER): + raise RuntimeError("vLLM Qwen2MLP is already RL-Kernel patched") + original = Qwen2MLP.forward + + def wrapped(instance: Any, hidden_states: Any) -> Any: + def native(_module: Any, value: Any) -> Any: + return original(instance, value) + + return integration.execute("ffn", native, instance, hidden_states) + + setattr(Qwen2MLP, _PATCH_MARKER, original) + Qwen2MLP.forward = wrapped + integration.record_installed_hook("ffn", "vllm.model_executor.models.qwen2.Qwen2MLP.forward") + + +def _patch_sampler(integration: VllmIntegration, *, strict_linear_logp: bool) -> None: + from vllm.v1.sample.sampler import Sampler + + if hasattr(Sampler, _PATCH_MARKER): + raise RuntimeError("vLLM Sampler is already RL-Kernel patched") + original = Sampler.forward + operator = VllmLogpOperator(original, strict_linear_logp=strict_linear_logp) + integration.install_operator("logp", operator) + + def wrapped(instance: Any, *args: Any, **kwargs: Any) -> Any: + sampling_metadata = kwargs.get("sampling_metadata") + if sampling_metadata is None and len(args) >= 2: + sampling_metadata = args[1] + if strict_linear_logp and _is_v1_sampler_profile_batch(sampling_metadata): + # gpu_model_runner._dummy_sampler_run computes logits with a + # synthetic no-logprob batch. It is an allocator warmup, not a + # rollout scoring event, so do not route or count it as logp. + try: + return original(instance, *args, **kwargs) + finally: + clear_rollout_linear_logp_context() + + def native(_sampler: Any, *call_args: Any, **call_kwargs: Any) -> Any: + return original(instance, *call_args, **call_kwargs) + + return integration.execute("logp", native, instance, *args, **kwargs) + + setattr(Sampler, _PATCH_MARKER, original) + Sampler.forward = wrapped + integration.record_installed_hook("logp", "vllm.v1.sample.sampler.Sampler.forward") + + +def _patch_worker_sampler(integration: VllmIntegration, *, strict_linear_logp: bool) -> None: + """Patch the CUDA worker sampler used by vLLM 0.27's V1 runner. + + vLLM has two sampler implementations: the graph-friendly + ``vllm.v1.sample.Sampler`` and the CUDA worker sampler used by the + production GPU runner. Qwen3 rollout on vLLM 0.27 uses the latter. + """ + + try: + from vllm.v1.worker.gpu.sample.sampler import Sampler + except ImportError: + return + + if hasattr(Sampler, _PATCH_MARKER): + raise RuntimeError("vLLM GPU worker Sampler is already RL-Kernel patched") + original = Sampler.__call__ + operator = VllmLogpOperator( + original, + worker_sampler=True, + strict_linear_logp=strict_linear_logp, + ) + integration.install_operator("logp", operator) + + def wrapped(instance: Any, *args: Any, **kwargs: Any) -> Any: + def native(_sampler: Any, *call_args: Any, **call_kwargs: Any) -> Any: + return original(instance, *call_args, **call_kwargs) + + if args and _is_worker_sampler_profile_batch(args[1] if len(args) > 1 else None): + return original(instance, *args, **kwargs) + return integration.execute("logp", native, instance, *args, **kwargs) + + setattr(Sampler, _PATCH_MARKER, original) + Sampler.__call__ = wrapped + integration.record_installed_hook("logp", "vllm.v1.worker.gpu.sample.sampler.Sampler.__call__") + + +def _register_attention_backend(integration: VllmIntegration) -> None: + global _RLK_ATTENTION_BACKEND, _RLK_ATTENTION_IMPL + + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl + from vllm.v1.attention.backends.registry import AttentionBackendEnum, register_backend + + operator = VllmAttentionOperator() + integration.install_operator("attention", operator) + + class RlKernelFlashAttentionImpl(FlashAttentionImpl): + def forward(self, *args: Any, **kwargs: Any) -> Any: + active = get_active_integration("vllm") + if active is not integration: + raise RuntimeError("vLLM Attention executed without its installed integration") + + def native(_impl: Any, *call_args: Any, **call_kwargs: Any) -> Any: + return FlashAttentionImpl.forward(self, *call_args, **call_kwargs) + + return integration.execute("attention", native, self, *args, **kwargs) + + class RlKernelFlashAttentionBackend(FlashAttentionBackend): + @staticmethod + def get_impl_cls() -> type[Any]: + return RlKernelFlashAttentionImpl + + RlKernelFlashAttentionImpl.__module__ = __name__ + RlKernelFlashAttentionImpl.__qualname__ = "RlKernelFlashAttentionImpl" + RlKernelFlashAttentionBackend.__module__ = __name__ + RlKernelFlashAttentionBackend.__qualname__ = "RlKernelFlashAttentionBackend" + _RLK_ATTENTION_IMPL = RlKernelFlashAttentionImpl + _RLK_ATTENTION_BACKEND = RlKernelFlashAttentionBackend + globals()["RlKernelFlashAttentionImpl"] = RlKernelFlashAttentionImpl + globals()["RlKernelFlashAttentionBackend"] = RlKernelFlashAttentionBackend + # vLLM 0.27 selects FLASH_ATTN for Qwen on CUDA before custom third-party + # names are considered. Override the selected enum in-place so the launcher + # does not need to edit vLLM source or pass version-specific CLI flags. + register_backend( + AttentionBackendEnum.FLASH_ATTN, + f"{__name__}.RlKernelFlashAttentionBackend", + ) + integration.record_installed_hook("attention", f"{__name__}.RlKernelFlashAttentionBackend") + + +def install_vllm_integration(plan: IntegrationPlan) -> VllmIntegration: + """Install vLLM paged Attention, Qwen dense FFN and sampler Logp routes.""" + + existing = get_active_integration("vllm") + if existing is not None: + if not isinstance(existing, VllmIntegration): + raise RuntimeError("active vLLM integration has an unexpected type") + if existing.plan != plan: + raise RuntimeError("vLLM integration is already installed with another plan") + return existing + integration = VllmIntegration(plan, rl_kernel_operators={}) + set_active_integration("vllm", integration) + strict_linear_logp = plan.implementation_for("logp", "rollout") is Implementation.RL_KERNEL + if strict_linear_logp: + _patch_qwen_lm_head_padding() + _patch_qwen_compute_logits(integration) + + # Patch every boundary so P/R cases also produce production readback. The + # integration object chooses native versus RL-Kernel per module. + _register_attention_backend(integration) + _patch_qwen_ffn(integration) + _patch_sampler(integration, strict_linear_logp=strict_linear_logp) + # vLLM 0.16's GPU model runner invokes vllm.v1.sample.sampler.Sampler. + # Its separate worker sampler has an incompatible seven-argument API; it + # must not replace the single logp route used by the active V1 runner. + if strict_linear_logp: + integration.record_installed_hook( + "logp", + "vllm.model_executor.models.qwen3.Qwen3ForCausalLM.compute_logits," + "vllm.v1.sample.sampler.Sampler.forward", + ) + return integration + + +def register_vllm_plugin() -> None: + """vLLM general-plugin entry point; inactive unless Vime exported a plan.""" + + if os.getenv("RL_KERNEL_VLLM_INTEGRATION", "").strip() not in {"1", "true", "True"}: + return + install_vllm_integration(plan_from_environment()) + + +__all__ = [ + "configure_vllm_environment", + "install_vllm_integration", + "plan_from_environment", + "register_vllm_plugin", +] diff --git a/rl_engine/kernels/ops/cuda/attention/strict_runtime.py b/rl_engine/kernels/ops/cuda/attention/strict_runtime.py new file mode 100644 index 00000000..90b20c97 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/attention/strict_runtime.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Framework-neutral strict CUDA Attention runtime. + +The runtime composes the production FA4 core and the self-owned CUDA AG/RS +transport. Framework integrations provide only local layout metadata and +logical position IDs. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_PRODUCTION_CORE_ID, + AttentionContract, +) +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionParallelSpec, + CUDAAGRSAttentionCPCommunication, +) +from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core + + +@dataclass(frozen=True) +class StrictCUDAAttentionResult: + out: torch.Tensor + lse: torch.Tensor + provenance: dict[str, Any] + + +class StrictCUDAAttentionRuntime: + """Run one FA4 arithmetic identity at CP=1 or through CUDA AG/RS.""" + + backend_id = "rlkernel.cuda.attention.fa4_ag_rs.v1" + core_id = STRICT_ATTENTION_PRODUCTION_CORE_ID + strict_schedule = STRICT_ATTENTION_FA4_SCHEDULE_ID + + def __init__( + self, + *, + process_group: Any = None, + core: Any | None = None, + communication: Any | None = None, + ) -> None: + self._core = StrictFlashAttention4Core() if core is None else core + self._communication = ( + CUDAAGRSAttentionCPCommunication(process_group=process_group) + if communication is None + else communication + ) + if getattr(self._core, "core_id", None) != self.core_id: + raise RuntimeError("strict CUDA Attention runtime requires the FA4 production core") + if getattr(self._core, "strict_schedule", None) != self.strict_schedule: + raise RuntimeError("strict CUDA Attention runtime requires the FA4 fixed schedule") + self.communication_executed = False + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + contract: AttentionContract, + causal: bool, + scale: float | None, + cp_world_size: int, + query_position_ids: torch.Tensor, + key_position_ids: torch.Tensor, + ) -> StrictCUDAAttentionResult: + self._require_nvidia_cuda(q) + if cp_world_size != contract.sharding.cp_world_size: + raise RuntimeError("runtime CP world size does not match AttentionContract") + self._validate_local_positions(q, k, query_position_ids, key_position_ids) + + if cp_world_size == 1: + global_q, global_k, global_v = q, k, v + global_q_positions = query_position_ids + global_k_positions = key_position_ids + communication_backend = "none" + self.communication_executed = False + else: + plan = self._communication_plan(contract, q.size(2), k.size(2)) + global_q = self._communication.all_gather_query(q, plan) + global_k, global_v = self._communication.all_gather_kv(k, v, plan) + global_q_positions, global_k_positions = self._communication.all_gather_position_ids( + query_position_ids, + key_position_ids, + plan, + ) + communication_backend = "cuda_ag_rs" + self.communication_executed = True + + q_sorted, q_positions_sorted, q_sort = self._sort_by_position(global_q, global_q_positions) + k_sorted, k_positions_sorted, _ = self._sort_by_position(global_k, global_k_positions) + v_sorted = self._gather_sequence(global_v, torch.argsort(global_k_positions, dim=1)) + self._validate_global_positions(q_positions_sorted, k_positions_sorted, causal) + + # FA4 consumes [B, S, H, D]. Materialize that layout once for every + # logical sequence instead of once per causal prefix. + q_fa = q_sorted.transpose(1, 2).contiguous() + k_fa = k_sorted.transpose(1, 2).contiguous() + v_fa = v_sorted.transpose(1, 2).contiguous() + + batch_outputs: list[torch.Tensor] = [] + batch_lses: list[torch.Tensor] = [] + core_rows: list[dict[str, Any]] = [] + for batch_index in range(q_sorted.size(0)): + query_outputs: list[torch.Tensor] = [] + query_lses: list[torch.Tensor] = [] + for query_index in range(q_sorted.size(2)): + query_position = q_positions_sorted[batch_index, query_index] + if causal: + prefix_tokens = int( + torch.searchsorted( + k_positions_sorted[batch_index], + query_position, + right=True, + ).item() + ) + else: + prefix_tokens = k_sorted.size(2) + result = self._core.forward_bshd_with_lse( + q_fa[batch_index : batch_index + 1, query_index : query_index + 1], + k_fa[batch_index : batch_index + 1, :prefix_tokens], + v_fa[batch_index : batch_index + 1, :prefix_tokens], + # Prefix materialization is the mask. Both framework sides + # therefore launch the exact same one-query FA4 shape. + causal=False, + scale=scale, + query_position_ids=q_positions_sorted[ + batch_index : batch_index + 1, + query_index : query_index + 1, + ], + key_position_ids=k_positions_sorted[ + batch_index : batch_index + 1, + :prefix_tokens, + ], + output_dtype=q.dtype, + ) + query_outputs.append(result.out) + query_lses.append(result.lse) + core_rows.append( + { + **dict(result.provenance), + "query_position": int(query_position.item()), + "kv_tokens": prefix_tokens, + } + ) + batch_outputs.append(torch.cat(query_outputs, dim=1)) + batch_lses.append(torch.cat(query_lses, dim=2)) + out_sorted = torch.cat(batch_outputs, dim=0).transpose(1, 2).contiguous() + lse_sorted = torch.cat(batch_lses, dim=0) + + if cp_world_size > 1: + inverse_q_sort = torch.argsort(q_sort, dim=1) + out_rank_packed = self._gather_sequence(out_sorted, inverse_q_sort) + lse_rank_packed = self._gather_sequence(lse_sorted, inverse_q_sort) + shard = self._communication.reduce_scatter_strict_result( + out_rank_packed, + lse_rank_packed, + plan, + ) + out, lse = shard.out, shard.lse + else: + out, lse = out_sorted, lse_sorted + + return StrictCUDAAttentionResult( + out=out, + lse=lse, + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "communication_backend": communication_backend, + "communication_executed": self.communication_executed, + "native_attention_arithmetic": True, + "production_ready": True, + "fallback": False, + "fallback_reason": None, + "reference_only": False, + "split_kv": "disabled", + "framework_position_reorder": True, + "query_schedule": "single_query_causal_prefix", + "core_rows": core_rows, + }, + ) + + @staticmethod + def _require_nvidia_cuda(tensor: torch.Tensor) -> None: + if tensor.device.type != "cuda" or torch.version.hip is not None: + raise RuntimeError("strict Attention R/R requires NVIDIA CUDA tensors") + + @staticmethod + def _communication_plan( + contract: AttentionContract, + local_q_tokens: int, + local_kv_tokens: int, + ) -> AttentionCPCommunicationPlan: + sharding = contract.sharding + parallel = AttentionParallelSpec( + tp_world_size=sharding.tp_world_size, + tp_rank=sharding.tp_rank, + cp_world_size=sharding.cp_world_size, + cp_rank=sharding.cp_rank, + ) + query_ranges = tuple( + (rank * local_q_tokens, (rank + 1) * local_q_tokens) + for rank in range(sharding.cp_world_size) + ) + blocks = tuple( + AttentionCPBlockMetadata( + global_block_index=rank, + kv_block_start=rank * local_kv_tokens, + kv_block_end=(rank + 1) * local_kv_tokens, + owner_cp_rank=rank, + owner_tp_rank=sharding.tp_rank, + ) + for rank in range(sharding.cp_world_size) + ) + return AttentionCPCommunicationPlan( + parallel=parallel, + backend="cuda_ag_rs", + status="implemented", + expected_blocks=blocks, + expected_kv_token_range=(0, local_kv_tokens * sharding.cp_world_size), + query_token_ranges=query_ranges, + ) + + @staticmethod + def _validate_local_positions( + q: torch.Tensor, + k: torch.Tensor, + query_positions: torch.Tensor, + key_positions: torch.Tensor, + ) -> None: + if query_positions.shape != (q.size(0), q.size(2)): + raise RuntimeError("query_position_ids must describe every local query token") + if key_positions.shape != (k.size(0), k.size(2)): + raise RuntimeError("key_position_ids must describe every local KV token") + if query_positions.device != q.device or key_positions.device != k.device: + raise RuntimeError("Attention position IDs must be on the Q/K CUDA device") + if query_positions.dtype not in (torch.int32, torch.int64) or ( + key_positions.dtype not in (torch.int32, torch.int64) + ): + raise RuntimeError("Attention position IDs must contain integers") + + @staticmethod + def _sort_by_position( + tensor: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + order = torch.argsort(positions, dim=1) + return ( + StrictCUDAAttentionRuntime._gather_sequence(tensor, order), + torch.gather(positions, 1, order), + order, + ) + + @staticmethod + def _gather_sequence(tensor: torch.Tensor, order: torch.Tensor) -> torch.Tensor: + if tensor.ndim == 4: + index = order[:, None, :, None].expand( + tensor.size(0), tensor.size(1), order.size(1), tensor.size(3) + ) + elif tensor.ndim == 3: + index = order[:, None, :].expand(tensor.size(0), tensor.size(1), order.size(1)) + else: + raise RuntimeError("strict Attention sequence reorder expects a 3-D or 4-D tensor") + return torch.gather(tensor, 2, index).contiguous() + + @staticmethod + def _validate_global_positions( + query_positions: torch.Tensor, + key_positions: torch.Tensor, + causal: bool, + ) -> None: + for positions, name in ( + (query_positions, "query"), + (key_positions, "key"), + ): + if positions.size(1) > 1 and bool((positions[:, 1:] <= positions[:, :-1]).any()): + raise RuntimeError(f"global {name} positions must be unique and increasing") + if causal and not torch.equal( + query_positions, + key_positions[:, -query_positions.size(1) :], + ): + raise RuntimeError("causal Attention queries must be the trailing global KV positions") + + +__all__ = ["StrictCUDAAttentionResult", "StrictCUDAAttentionRuntime"] diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index bea6fb52..7edbe847 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -20,8 +20,6 @@ from typing import Any, Callable, Mapping, cast import torch -from torch import Tensor - from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, STRICT_ATTENTION_SCHEDULE_ID, @@ -30,6 +28,7 @@ AttentionDType, SplitKVMode, ) +from torch import Tensor BACKEND_ID = "rlkernel.attention.deterministic.v1" REFERENCE_BACKEND_ID = "rlkernel.attention.reference.v1" @@ -121,7 +120,7 @@ def readback(self) -> dict[str, Any]: "out_shape": list(self.out.shape), "out_dtype": str(self.out.dtype).replace("torch.", ""), "lse_shape": None if self.lse is None else list(self.lse.shape), - "lse_dtype": None if self.lse is None else str(self.lse.dtype).replace("torch.", ""), + "lse_dtype": (None if self.lse is None else str(self.lse.dtype).replace("torch.", "")), "gradients": { "dq": self.dq is not None, "dk": self.dk is not None, @@ -163,6 +162,26 @@ def __init__( # fallback to the PyTorch reference when AG/RS is required. self.cp_backend = cp_backend self.communication_backend = communication_backend.strip() + self._cuda_runtime_group: Any = None + self._cuda_runtime_bound = False + + def bind_cuda_runtime(self, *, process_group: Any = None) -> Any: + """Bind the shared production CUDA core/transport once per process.""" + + if self._cuda_runtime_bound: + if process_group is not self._cuda_runtime_group: + raise AttentionContractError( + "Attention CUDA runtime is already bound to another process group" + ) + return self.core + from rl_engine.kernels.ops.cuda.attention.strict_runtime import StrictCUDAAttentionRuntime + + runtime = StrictCUDAAttentionRuntime(process_group=process_group) + self.core = runtime + self.cp_backend = runtime + self._cuda_runtime_group = process_group + self._cuda_runtime_bound = True + return runtime def __call__( self, @@ -277,6 +296,7 @@ def apply( ) call_kwargs = dict(kwargs) + call_kwargs.setdefault("contract", contract) call_kwargs.setdefault("causal", contract.causal) call_kwargs.setdefault("scale", 1.0 / math.sqrt(contract.head_dim)) call_kwargs.setdefault("cp_world_size", contract.sharding.cp_world_size) @@ -305,13 +325,15 @@ def apply( "strict_schedule": cfg.strict_schedule if cfg.deterministic else None, "core_id": cfg.strict_core_id if cfg.deterministic else selected_id, "backend_deterministic": cfg.deterministic, - "native_attention_arithmetic": False if cfg.deterministic else selected_id == "native", + "native_attention_arithmetic": ( + False if cfg.deterministic else selected_id == "native" + ), "communication_backend": cfg.communication_backend, "communication_executed": bool(getattr(selected, "communication_executed", False)), "split_kv": contract.split_kv.to_dict(), "actual_split_kv": _actual_split_provenance( contract, - total_kv_tokens=k.size(2), + total_kv_tokens=k.size(2) * contract.sharding.cp_world_size, backend=selected_id, ), "reduction": _reduction_provenance(contract), diff --git a/tests/test_framework_runtime_adapters.py b/tests/test_framework_runtime_adapters.py new file mode 100644 index 00000000..e3a71b62 --- /dev/null +++ b/tests/test_framework_runtime_adapters.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import ast +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest +import rl_engine.integrations.framework_operators as framework_operators +import torch +from rl_engine.integrations.ablation import ( + IntegrationPlan, + configure_integration_environment, + integration_plan_from_environment, +) +from rl_engine.integrations.framework_operators import ( + MegatronAttentionOperator, + SemanticOperatorHandle, + _megatron_zigzag_layout, + _packed_local_sequence_layout, + _vllm_kv_cache_views, +) +from rl_engine.integrations.megatron_runtime import install_megatron_integration +from rl_engine.integrations.runtime import FrameworkOperatorIntegration +from rl_engine.integrations.state import clear_active_integration +from rl_engine.integrations.vllm_runtime import configure_vllm_environment +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_PRODUCTION_CORE_ID, + AttentionContract, + AttentionDType, + AttentionMode, + AttentionRole, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.cuda.attention.strict_runtime import StrictCUDAAttentionRuntime + + +def test_framework_adapters_do_not_construct_registered_kernels_directly(): + source_path = ( + Path(__file__).parents[1] / "rl_engine" / "integrations" / "framework_operators.py" + ) + tree = ast.parse(source_path.read_text(encoding="utf-8")) + forbidden = { + "AttentionAblationOp", + "DeterministicCPAttentionReferenceOp", + "Qwen3FFNOp", + "StrictCUDAAttentionRuntime", + "VocabParallelLogprobOp", + } + constructed = { + node.func.id + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + + assert constructed.isdisjoint(forbidden) + + +def test_semantic_handle_uses_operator_bridge_and_exposes_instance_provenance(): + handle = SemanticOperatorHandle( + target="training", + semantic_op="attention", + backend_id="rlkernel.attention.deterministic.v1", + ) + tensor = torch.zeros(1, 1, 1, 8) + + first = handle.get( + tensor, + topology={ + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + ) + second = handle.get( + tensor, + topology={ + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + ) + + assert first is second + assert first.backend_id == "rlkernel.attention.deterministic.v1" + assert handle.provenance is not None + assert handle.provenance["semantic_op"] == "attention" + assert handle.provenance["backend_id"] == first.backend_id + + +def test_plan_environment_is_shared_by_both_framework_installers(monkeypatch, tmp_path): + for variable in ( + "RL_KERNEL_ATTENTION_CASE", + "RL_KERNEL_FFN_CASE", + "RL_KERNEL_LOGP_CASE", + "RL_KERNEL_READBACK_DIR", + ): + monkeypatch.setenv(variable, "previous") + plan = IntegrationPlan.from_case_ids(attention="P/R", ffn="R/P", logp="R/R") + configure_integration_environment(plan, readback_dir=str(tmp_path)) + + assert integration_plan_from_environment() == plan + assert Path(os.environ["RL_KERNEL_READBACK_DIR"]) == tmp_path + + +def test_vllm_rlkernel_attention_overrides_selected_flash_attn_backend( + monkeypatch, +): + monkeypatch.delenv("VLLM_ATTENTION_BACKEND", raising=False) + plan = IntegrationPlan.from_case_ids(attention="P/R") + + configure_vllm_environment(plan) + + assert os.environ["VLLM_ATTENTION_BACKEND"] == "FLASH_ATTN" + + +def test_megatron_install_is_idempotent_in_one_actor(): + class Attention: + def forward(self, value): + return value + + class FFN: + def forward(self, value): + return value + + plan = IntegrationPlan.from_case_ids() + clear_active_integration("megatron") + try: + first = install_megatron_integration( + plan, + attention_classes=(Attention,), + ffn_classes=(FFN,), + ) + second = install_megatron_integration( + plan, + attention_classes=(Attention,), + ffn_classes=(FFN,), + ) + assert first is second + finally: + clear_active_integration("megatron") + + +def test_megatron_zigzag_positions_preserve_global_cp_ownership(): + rank_zero = _megatron_zigzag_layout(4, cp_rank=0, cp_world_size=2) + rank_one = _megatron_zigzag_layout(4, cp_rank=1, cp_world_size=2) + + assert rank_zero == ((0, 1, 6, 7), (0, 3), (0, 6), (0, 2, 4)) + assert rank_one == ((2, 3, 4, 5), (1, 2), (2, 4), (0, 2, 4)) + assert sorted(rank_zero[0] + rank_one[0]) == list(range(8)) + + +def test_packed_layout_recovers_local_offsets_from_global_cu_seqlens(): + packed = SimpleNamespace( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 8, 16], dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, 8, 16], dtype=torch.int32), + ) + + local_offsets, global_lengths = _packed_local_sequence_layout( + packed, + cp_world_size=2, + local_query_tokens=8, + local_kv_tokens=8, + ) + + assert local_offsets == (0, 4, 8) + assert global_lengths == (8, 8) + + +def test_megatron_packed_attention_runs_each_sequence_in_thd_order(monkeypatch): + calls: list[dict[str, object]] = [] + + class Operator: + def bind_cuda_runtime(self, *, process_group=None): + assert process_group == "cp-group" + + def __call__(self, q, k, v, **kwargs): + del k, v + calls.append(kwargs) + return SimpleNamespace( + out=q.clone(), + provenance={ + "actual_backend": "rlkernel.cuda.attention.fa4_ag_rs.v1", + "core_rows": [{"actual_backend": "rlkernel.cuda.attention.fa4.v1"}], + }, + ) + + operator = Operator() + + class Handle: + provenance = {} + + def get(self, tensor, *, topology): + assert tensor.shape == (8, 2, 4) + assert topology["context_parallel_size"] == 2 + return operator + + parallel_state = SimpleNamespace( + get_context_parallel_world_size=lambda: 2, + get_context_parallel_rank=lambda: 0, + get_tensor_model_parallel_world_size=lambda: 2, + get_tensor_model_parallel_rank=lambda: 0, + get_context_parallel_group=lambda: "cp-group", + ) + monkeypatch.setattr(framework_operators, "_megatron_parallel_state", lambda: parallel_state) + monkeypatch.setattr(framework_operators, "_require_nvidia_cuda", lambda tensor, module: None) + packed = SimpleNamespace( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 8, 16], dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, 8, 16], dtype=torch.int32), + ) + query = torch.zeros(8, 2, 4, dtype=torch.bfloat16) + key = torch.zeros(8, 1, 4, dtype=torch.bfloat16) + + output = MegatronAttentionOperator(handle=Handle())( + SimpleNamespace(softmax_scale=0.5), + query, + key, + key, + None, + packed_seq_params=packed, + num_splits=1, + ) + + assert output.shape == (8, 8) + assert len(calls) == 2 + assert [call["contract"].sharding.global_sequence_length for call in calls] == [ + 8, + 8, + ] + assert [call["query_position_ids"].tolist() for call in calls] == [ + [[0, 1, 6, 7]], + [[0, 1, 6, 7]], + ] + + +def test_vllm_current_flash_attention_kv_cache_layout_is_materialized(): + cache = torch.arange(2 * 3 * 4 * 10).reshape(2, 3, 4, 10) + + key, value = _vllm_kv_cache_views(cache, head_size=5) + + assert key.shape == (2, 4, 3, 5) + assert value.shape == (2, 4, 3, 5) + assert torch.equal(key, cache.transpose(1, 2)[..., :5]) + assert torch.equal(value, cache.transpose(1, 2)[..., 5:]) + + +def _cp1_contract(tokens: int) -> AttentionContract: + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=tokens, + head_dim=4, + causal=True, + causal_offsets=(0,), + sharding=ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=1, + global_q_heads=2, + global_kv_heads=1, + local_q_head_start=0, + local_q_heads=2, + local_kv_head_start=0, + local_kv_heads=1, + global_sequence_length=tokens, + local_sequence_length=tokens, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, tokens), + ), + reduction=ReductionSpec(), + ) + + +def test_strict_cuda_runtime_pins_training_to_single_query_prefixes(monkeypatch): + calls: list[tuple[int, int, bool]] = [] + + class Core: + core_id = STRICT_ATTENTION_PRODUCTION_CORE_ID + strict_schedule = STRICT_ATTENTION_FA4_SCHEDULE_ID + + def forward_bshd_with_lse(self, q, k, v, *, causal, **kwargs): + del v, kwargs + calls.append((q.size(1), k.size(1), causal)) + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros((q.size(0), q.size(2), q.size(1)), dtype=torch.float32), + provenance={"actual_backend": "fake.cuda.fa4"}, + ) + + monkeypatch.setattr( + StrictCUDAAttentionRuntime, "_require_nvidia_cuda", lambda self, tensor: None + ) + runtime = StrictCUDAAttentionRuntime(core=Core(), communication=object()) + q = torch.zeros(1, 2, 4, 4, dtype=torch.bfloat16) + k = torch.zeros(1, 1, 4, 4, dtype=torch.bfloat16) + positions = torch.arange(4).unsqueeze(0) + + result = runtime.forward_with_lse( + q, + k, + k, + contract=_cp1_contract(4), + causal=True, + scale=0.5, + cp_world_size=1, + query_position_ids=positions, + key_position_ids=positions, + ) + + assert calls == [(1, 1, False), (1, 2, False), (1, 3, False), (1, 4, False)] + assert result.provenance["query_schedule"] == "single_query_causal_prefix" + + +class _ReadbackOperator: + backend_id = "rlkernel.attention.test" + + def __init__(self, provenance): + self.provenance = provenance + + def __call__(self, value): + return value + + +@pytest.mark.parametrize( + ("provenance", "match"), + [ + ({"runtime_platform": "cpu"}, "non-cuda"), + ({"runtime_platform": "cuda", "actual_backend": "triton.attention"}, "triton"), + ({"runtime_platform": "cuda", "triton_used": True}, "triton"), + ], +) +def test_strict_readback_rejects_non_cuda_and_triton(provenance, match): + plan = IntegrationPlan.from_case_ids(attention="R/R") + integration = FrameworkOperatorIntegration( + framework="megatron", + target="training", + plan=plan, + rl_kernel_operators={"attention": _ReadbackOperator(provenance)}, + ) + integration.record_installed_hook("attention", "test.attention") + integration.execute("attention", lambda value: value, "x") + + with pytest.raises(RuntimeError, match=match): + integration.assert_strict_ready() + + +def test_strict_readback_accepts_cuda_without_triton(): + plan = IntegrationPlan.from_case_ids(attention="R/R") + integration = FrameworkOperatorIntegration( + framework="megatron", + target="training", + plan=plan, + rl_kernel_operators={ + "attention": _ReadbackOperator( + { + "runtime_platform": "cuda", + "actual_backend": "rlkernel.cuda.fa4", + "triton_used": False, + } + ) + }, + ) + integration.record_installed_hook("attention", "test.attention") + integration.execute("attention", lambda value: value, "x") + + integration.assert_strict_ready() diff --git a/tests/test_vime_validation_artifacts.py b/tests/test_vime_validation_artifacts.py new file mode 100644 index 00000000..36dd3122 --- /dev/null +++ b/tests/test_vime_validation_artifacts.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json + +import torch +from examples.vime_qwen3_8b_tp2_cp2.validate_artifacts import validate_artifacts + + +def _write_readback(directory, framework, target, *, triton_used=False): + operators = {} + hooks = {} + for module in ("attention", "ffn", "logp"): + backend = ( + "rlkernel.linear_logp.bitwise.v1" if module == "logp" else f"rlkernel.{module}.test" + ) + operators[module] = { + "module": module, + "implementation": "rl_kernel", + "backend_id": backend, + "call_count": 2, + "provenance": { + "runtime_platform": "cuda", + "actual_backend": f"rlkernel.cuda.{module}", + "triton_used": triton_used, + }, + } + hooks[module] = f"test.{module}" + payload = { + "framework": framework, + "target": target, + "installed_hooks": hooks, + "fallbacks": [], + "operators": operators, + } + (directory / f"{framework}-{target}.json").write_text(json.dumps(payload), encoding="utf-8") + + +def test_vime_artifacts_require_exact_cuda_logp_equality(tmp_path): + readbacks = tmp_path / "readbacks" + train_data = tmp_path / "train-data" + readbacks.mkdir() + train_data.mkdir() + _write_readback(readbacks, "megatron", "training") + _write_readback(readbacks, "vllm", "rollout") + values = torch.tensor([-1.25, -2.5], dtype=torch.float32) + torch.save( + {"samples": [{"log_probs": values.clone(), "rollout_log_probs": values.clone()}]}, + train_data / "0.pt", + ) + + report = validate_artifacts(readbacks, train_data) + + assert report["passed"] is True + assert report["train_rollout_logp"]["torch_equal"] is True + assert report["train_rollout_logp"]["mismatch_count"] == 0 + assert report["train_rollout_logp"]["max_abs_diff"] == 0.0 + + +def test_vime_artifacts_reject_triton_even_when_values_match(tmp_path): + readbacks = tmp_path / "readbacks" + train_data = tmp_path / "train-data" + readbacks.mkdir() + train_data.mkdir() + _write_readback(readbacks, "megatron", "training") + _write_readback(readbacks, "vllm", "rollout", triton_used=True) + values = torch.tensor([-1.25], dtype=torch.float32) + torch.save( + {"samples": [{"log_probs": values, "rollout_log_probs": values.clone()}]}, + train_data / "0.pt", + ) + + report = validate_artifacts(readbacks, train_data) + + assert report["passed"] is False + assert "vllm/rollout attention used Triton" in report["readbacks"]["errors"]