Enable exact-batch CUDA Graph for strict rollout

csrc/cuda/distributed/deterministic_collective.cu

@@ -949,7 +949,10 @@ class DeterministicCollectiveState {
     has_staged_input_ = true;
   }
 
-  void all_reduce(torch::Tensor& output, cudaStream_t stream) {
+  void all_reduce(
+      torch::Tensor& output,
+      cudaStream_t stream,
+      bool allow_owner_path = true) {
     check_tensor(output, "output");
     TORCH_CHECK(has_staged_input_, "stage() must be called before all_reduce()");
     TORCH_CHECK(
@@ -960,7 +963,10 @@ class DeterministicCollectiveState {
         "all-reduce output size must match the staged input size");
 
     const int64_t element_count = output.numel();
-    if (staged_owner_path_) {
+    // Graph-replayed calls must avoid the owner-push branch because it writes
+    // to remote IPC frames. Callers that use the fused ABI pass false here;
+    // direct staged collectives retain the eager owner optimization.
+    if (allow_owner_path && staged_owner_path_) {
       if (rank_ == 0) {
         // Logical rank 0 is the topology-favorable reader in both traced TP
         // engines. It evaluates the original fixed tree exactly once, then
@@ -1022,6 +1028,58 @@ class DeterministicCollectiveState {
       return;
     }
 
+    // Small collectives are launch-bound. The fast kernel folds the peer
+    // stage wait, fixed-tree reduction, and completion publication into one
+    // launch while preserving the exact reduction order.
+    if (staged_fast_path_ && staged_bytes_ <= kSingleBlockFastPathMaxBytes) {
+      if (element_count > 0) {
+        switch (output.scalar_type()) {
+          case at::ScalarType::Float:
+            launch_all_reduce_fast<float>(
+                peers_,
+                local_stage_sequence_,
+                local_done_sequence_,
+                static_cast<float*>(output.data_ptr()),
+                element_count,
+                world_size_,
+                stream);
+            break;
+          case at::ScalarType::Half:
+            launch_all_reduce_fast<half>(
+                peers_,
+                local_stage_sequence_,
+                local_done_sequence_,
+                static_cast<half*>(output.data_ptr()),
+                element_count,
+                world_size_,
+                stream);
+            break;
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+          case at::ScalarType::BFloat16:
+            launch_all_reduce_fast<nv_bfloat16>(
+                peers_,
+                local_stage_sequence_,
+                local_done_sequence_,
+                static_cast<nv_bfloat16*>(output.data_ptr()),
+                element_count,
+                world_size_,
+                stream);
+            break;
+#endif
+          default:
+            TORCH_CHECK(
+                false,
+                "deterministic all-reduce supports float32, float16, and bfloat16; got ",
+                output.scalar_type());
+        }
+        AT_CUDA_CHECK(cudaGetLastError());
+      } else {
+        publish_done(stream);
+      }
+      has_staged_input_ = false;
+      return;
+    }
+
     wait_for_staged_peers(stream);
     if (element_count == 0) {
       publish_done(stream);
@@ -1090,10 +1148,11 @@ class DeterministicCollectiveState {
         output.numel() == input.numel(),
         "all-reduce output size must match the input size");
 
-    // Route both ABI entry points through the same staged protocol so the
-    // topology-aware owner reduction is always applied.
+    // Use the graph-safe staged protocol for every message size. The fused
+    // two-slot protocol is intentionally kept available in the extension for
+    // experiments, but is not safe to replay across vLLM's many graph shapes.
     stage(input, stream);
-    all_reduce(output, stream);
+    all_reduce(output, stream, /*allow_owner_path=*/false);
   }
 
   void all_gather_fused(

examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh

@@ -12,21 +12,80 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then
   # vLLM hooks; otherwise it silently falls back to native attention/FFN and
   # loses the R/R performance path. Preserve explicit ablation selections.
   strict_linear_logp=0
+  rollout_batch_size=""
+  n_samples_per_prompt="1"
+  explicit_vllm_execution_config=0
   previous_arg=""
   for current_arg in "$@"; do
     if [[ "${previous_arg}" == "--linear-logp-provider-mode" && "${current_arg}" == "strict" ]]; then
       strict_linear_logp=1
-      break
+    elif [[ "${previous_arg}" == "--rollout-batch-size" ]]; then
+      rollout_batch_size="${current_arg}"
+    elif [[ "${previous_arg}" == "--n-samples-per-prompt" ]]; then
+      n_samples_per_prompt="${current_arg}"
     fi
+
+    case "${current_arg}" in
+      --linear-logp-provider-mode=strict)
+        strict_linear_logp=1
+        ;;
+      --rollout-batch-size=*)
+        rollout_batch_size="${current_arg#*=}"
+        ;;
+      --n-samples-per-prompt=*)
+        n_samples_per_prompt="${current_arg#*=}"
+        ;;
+      --vllm-enforce-eager|--vllm-optimization-level|--vllm-optimization-level=*|--vllm-compilation-config|--vllm-compilation-config=*)
+        explicit_vllm_execution_config=1
+        ;;
+    esac
     previous_arg="${current_arg}"
   done
+
+  strict_cudagraph_args=()
   if [[ "${strict_linear_logp}" == "1" ]]; then
     export RL_KERNEL_VLLM_INTEGRATION="${RL_KERNEL_VLLM_INTEGRATION:-1}"
     export RL_KERNEL_CUDA_ONLY="${RL_KERNEL_CUDA_ONLY:-1}"
     export VIME_RL_KERNEL_STRICT="${VIME_RL_KERNEL_STRICT:-1}"
     export RL_KERNEL_ATTENTION_CASE="${RL_KERNEL_ATTENTION_CASE:-R/R}"
     export RL_KERNEL_FFN_CASE="${RL_KERNEL_FFN_CASE:-R/R}"
     export RL_KERNEL_LOGP_CASE="${RL_KERNEL_LOGP_CASE:-R/R}"
+
+    # Strict rollout kernels preserve their arithmetic order under CUDA Graph.
+    # Capturing the complete decode graph removes the per-layer host-launch
+    # gaps that otherwise dominate small decode batches. Capture every exact
+    # batch size: padding a strict custom kernel to a larger sparse graph can
+    # access invalid slots and, more importantly, changes the tested contract.
+    # Explicit vLLM execution flags always win so callers can opt out.
+    if [[ "${explicit_vllm_execution_config}" == "0" ]]; then
+      if [[ "${rollout_batch_size}" =~ ^[1-9][0-9]*$ && "${n_samples_per_prompt}" =~ ^[1-9][0-9]*$ ]]; then
+        max_capture_size=$((rollout_batch_size * n_samples_per_prompt))
+        if [[ -n "${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE:-}" ]]; then
+          max_capture_size="${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE}"
+        fi
+        if ! [[ "${max_capture_size}" =~ ^[1-9][0-9]*$ ]]; then
+          echo "RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE must be a positive integer" >&2
+          exit 2
+        fi
+
+        capture_sizes="["
+        for ((batch_size = 1; batch_size <= max_capture_size; batch_size++)); do
+          if ((batch_size > 1)); then
+            capture_sizes+=","
+          fi
+          capture_sizes+="${batch_size}"
+        done
+        capture_sizes+="]"
+        compilation_config="{\"cudagraph_mode\":\"FULL_DECODE_ONLY\",\"cudagraph_capture_sizes\":${capture_sizes},\"max_cudagraph_capture_size\":${max_capture_size}}"
+        strict_cudagraph_args=(
+          --vllm-optimization-level 0
+          --vllm-compilation-config "${compilation_config}"
+        )
+        echo "[RL-Kernel] strict vLLM full-decode CUDA Graph capture sizes: ${capture_sizes}" >&2
+      else
+        echo "[RL-Kernel] strict CUDA Graph disabled: rollout batch size is unavailable" >&2
+      fi
+    fi
   fi
   exec "${REAL_PYTHON}" "$@" \
     --seed 1234 \
@@ -35,7 +94,8 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then
     --vllm-attention-backend flash_attn \
     --vllm-disable-custom-all-reduce \
     --deterministic-mode \
-    --accumulate-allreduce-grads-in-fp32
+    --accumulate-allreduce-grads-in-fp32 \
+    "${strict_cudagraph_args[@]}"
 fi
 
 exec "${REAL_PYTHON}" "$@"

rl_engine/integrations/megatron_runtime.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import importlib
+import os
 from collections.abc import Callable, Iterable
 from types import MethodType
 from typing import Any
@@ -290,7 +291,48 @@ def initialize_from_environment(_args: Any = None) -> MegatronIntegration:
 
     from rl_engine.integrations.ablation import integration_plan_from_environment
 
-    return install_megatron_integration(integration_plan_from_environment())
+    plan = integration_plan_from_environment()
+    integration = install_megatron_integration(plan)
+    if plan.implementation_for("attention", "training") is Implementation.RL_KERNEL:
+        _precompile_strict_attention_training(_args)
+    return integration
+
+
+def _precompile_strict_attention_training(args: Any) -> None:
+    """Warm FA4 CuTe fwd/bwd JIT outside Vime's actor_train timer."""
+
+    if args is None or torch.version.hip is not None or not torch.cuda.is_available():
+        return
+    if os.getenv("RL_KERNEL_PRECOMPILE_FA4", "1") == "0":
+        return
+
+    from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core
+
+    attention_heads = int(getattr(args, "num_attention_heads", 0) or 0)
+    query_groups = int(getattr(args, "num_query_groups", 0) or attention_heads)
+    tp_size = int(getattr(args, "tensor_model_parallel_size", 1) or 1)
+    head_dim = int(
+        getattr(args, "kv_channels", 0)
+        or (int(getattr(args, "hidden_size", 0) or 0) // attention_heads)
+    )
+    if (
+        attention_heads <= 0
+        or query_groups <= 0
+        or tp_size <= 0
+        or attention_heads % tp_size
+        or query_groups % tp_size
+        or head_dim <= 0
+    ):
+        return
+
+    params_dtype = getattr(args, "params_dtype", None)
+    dtype = params_dtype if params_dtype in (torch.float16, torch.bfloat16) else torch.bfloat16
+    StrictFlashAttention4Core.precompile_training(
+        q_heads=attention_heads // tp_size,
+        kv_heads=query_groups // tp_size,
+        head_dim=head_dim,
+        dtype=dtype,
+    )
 
 
 __all__ = ["initialize_from_environment", "install_megatron_integration"]

rl_engine/integrations/vllm_runtime.py

@@ -12,7 +12,11 @@
 
 import torch
 
-from rl_engine.distributed.collectives import DETERMINISTIC_ALL_REDUCE_OP
+from rl_engine.distributed.collectives import (
+    DETERMINISTIC_ALL_REDUCE_OP,
+    collective_for_group,
+    deterministic_all_reduce_inplace,
+)
 from rl_engine.integrations.ablation import (
     Implementation,
     IntegrationPlan,
@@ -40,6 +44,8 @@
 _STRICT_RMS_NORM_INIT_MARKER = "__rl_kernel_original_strict_rms_norm_init__"
 _STRICT_ROTARY_INIT_MARKER = "__rl_kernel_original_strict_rotary_init__"
 _STRICT_LM_HEAD_LINEAR_PATCH_MARKER = "__rl_kernel_original_lm_head_linear_apply__"
+_STRICT_O_PROJ_COLLECTIVE_MARKER = "__rl_kernel_o_proj_collective__"
+_STRICT_ROW_PARALLEL_PATCH_MARKER = "__rl_kernel_original_row_parallel_forward__"
 _RLK_ATTENTION_BACKEND: type[Any] | None = None
 _RLK_ATTENTION_IMPL: type[Any] | None = None
 _RLK_ATTENTION_BUILDER: type[Any] | None = None
@@ -367,21 +373,23 @@ def _patch_qwen3_strict_model(
     rotary_cls: type[Any] | None = None,
     linear_method_cls: type[Any] | None = None,
     attention_cls: type[Any] | None = None,
+    row_parallel_cls: type[Any] | None = None,
     det_gemm: Any | None = None,
 ) -> None:
     """Align vLLM's RMSNorm and Attention projections with Megatron."""
 
     production_classes = rms_norm_cls is None or linear_method_cls is None or attention_cls is None
     if production_classes:
         from vllm.model_executor.layers.layernorm import RMSNorm
-        from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+        from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod
         from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
         from vllm.model_executor.models.qwen3 import Qwen3Attention
 
         rms_norm_cls = RMSNorm
         rotary_cls = RotaryEmbedding
         linear_method_cls = UnquantizedLinearMethod
         attention_cls = Qwen3Attention
+        row_parallel_cls = RowParallelLinear
     assert rms_norm_cls is not None
     assert linear_method_cls is not None
     assert attention_cls is not None
@@ -435,6 +443,62 @@ def strict_rms_norm_forward_cuda(
             eps=instance.variance_epsilon,
         )
 
+    def bind_o_proj_collective(module: Any) -> None:
+        if int(getattr(module, "tp_size", 1)) <= 1:
+            return
+        from vllm.distributed.parallel_state import get_tp_group
+
+        coordinator = get_tp_group()
+        group = getattr(coordinator, "device_group", coordinator)
+        collective = collective_for_group(group)
+        if collective is None:
+            raise RuntimeError("strict rollout o_proj requires an initialized TP process group")
+        setattr(module, _STRICT_O_PROJ_COLLECTIVE_MARKER, collective)
+
+    if row_parallel_cls is not None and not hasattr(
+        row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER
+    ):
+        row_parallel_forward = row_parallel_cls.forward
+
+        def strict_row_parallel_forward(instance: Any, input_: torch.Tensor) -> Any:
+            collective = getattr(instance, _STRICT_O_PROJ_COLLECTIVE_MARKER, None)
+            if collective is None:
+                return row_parallel_forward(instance, input_)
+
+            if instance.input_is_parallel:
+                input_parallel = input_
+            else:
+                from vllm.distributed import split_tensor_along_last_dim
+
+                input_parallel = split_tensor_along_last_dim(
+                    input_, num_partitions=instance.tp_size
+                )[instance.tp_rank].contiguous()
+
+            assert instance.quant_method is not None
+            bias_ = (
+                None
+                if (instance.tp_rank > 0 or instance.skip_bias_add)
+                else instance.bias
+            )
+            output_parallel = instance.quant_method.apply(instance, input_parallel, bias_)
+
+            if instance.reduce_results and instance.tp_size > 1:
+                deterministic_all_reduce_inplace(
+                    output_parallel,
+                    collective_handle=int(collective._handle),
+                )
+                output = output_parallel
+            else:
+                output = output_parallel
+
+            if not instance.return_bias:
+                return output
+            output_bias = instance.bias if instance.skip_bias_add else None
+            return output, output_bias
+
+        setattr(row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER, row_parallel_forward)
+        row_parallel_cls.forward = strict_row_parallel_forward
+
     if not hasattr(rms_norm_cls, _STRICT_RMS_NORM_INIT_MARKER):
         rms_norm_init = rms_norm_cls.__init__
 
@@ -463,6 +527,7 @@ def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None:
         attention_init(instance, *args, **kwargs)
         setattr(instance.qkv_proj, _STRICT_PROJECTION_MARKER, "qkv")
         setattr(instance.o_proj, _STRICT_PROJECTION_MARKER, "o_proj")
+        bind_o_proj_collective(instance.o_proj)
 
     setattr(attention_cls, _STRICT_MODEL_PATCH_MARKER, attention_init)
     linear_method_cls.apply = deterministic_linear_apply