Add batch-invariant h_aggregate kernel

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
@@ -1090,10 +1096,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(

csrc/cuda/mhc/mhc_pre_h_aggregate.cu

@@ -0,0 +1,62 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+#include <torch/extension.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <limits>
+
+#include "mhc_pre_h_aggregate_kernel.cuh"
+
+torch::Tensor mhc_pre_h_aggregate_cuda(torch::Tensor residual,
+                                       torch::Tensor pre) {
+  TORCH_CHECK(residual.is_cuda() && pre.is_cuda(),
+              "residual and pre must be CUDA tensors");
+  TORCH_CHECK(residual.device() == pre.device(),
+              "residual and pre must be on the same CUDA device");
+  TORCH_CHECK(residual.is_contiguous() && pre.is_contiguous(),
+              "residual and pre must be contiguous");
+  TORCH_CHECK(residual.scalar_type() == torch::kBFloat16,
+              "residual must be bfloat16");
+  TORCH_CHECK(pre.scalar_type() == torch::kFloat32,
+              "pre must be float32");
+  TORCH_CHECK(residual.dim() == 3 && residual.size(1) == 4,
+              "residual must have shape [num_tokens, 4, hidden_size]");
+  TORCH_CHECK(pre.dim() == 2 && pre.size(1) == 4,
+              "pre must have shape [num_tokens, 4]");
+  TORCH_CHECK(residual.size(0) == pre.size(0),
+              "residual and pre must have the same num_tokens");
+
+  int64_t const num_tokens = residual.size(0);
+  int64_t const hidden_size = residual.size(2);
+  TORCH_CHECK(num_tokens <= std::numeric_limits<unsigned int>::max(),
+              "num_tokens exceeds the CUDA grid limit");
+
+  auto output = torch::empty({num_tokens, hidden_size}, residual.options());
+  if (num_tokens == 0 || hidden_size == 0) {
+    return output;
+  }
+
+  c10::cuda::CUDAGuard const device_guard(residual.device());
+  int device = 0;
+  int major = 0;
+  C10_CUDA_CHECK(cudaGetDevice(&device));
+  C10_CUDA_CHECK(cudaDeviceGetAttribute(
+      &major, cudaDevAttrComputeCapabilityMajor, device));
+  TORCH_CHECK(major >= 8, "mhc_pre_h_aggregate requires SM80 or newer");
+
+  auto const* residual_ptr = reinterpret_cast<__nv_bfloat16 const*>(
+      residual.data_ptr<at::BFloat16>());
+  auto const* pre_ptr = pre.data_ptr<float>();
+  auto* output_ptr = reinterpret_cast<__nv_bfloat16*>(
+      output.data_ptr<at::BFloat16>());
+
+  cudaError_t const status = rl_kernel::mhc::launch_mhc_pre_h_aggregate(
+      residual_ptr, pre_ptr, output_ptr, num_tokens, hidden_size,
+      at::cuda::getCurrentCUDAStream(), major >= 9);
+  C10_CUDA_CHECK(status);
+  return output;
+}

csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh

@@ -0,0 +1,107 @@
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace rl_kernel::mhc {
+
+constexpr int kMhcPreHAggregateDecodeThreads = 1024;
+constexpr int kMhcPreHAggregateBatchThreads = 512;
+
+__global__ void mhc_pre_h_aggregate_kernel(__nv_bfloat16 const* residual,
+                                           float const* pre,
+                                           __nv_bfloat16* output,
+                                           int64_t hidden_size) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaGridDependencySynchronize();
+#endif
+
+  __shared__ float weights[4];
+  if (threadIdx.x < 4) {
+    weights[threadIdx.x] =
+        pre[static_cast<int64_t>(blockIdx.x) * 4 + threadIdx.x];
+  }
+  __syncthreads();
+
+  int64_t const token_offset =
+      static_cast<int64_t>(blockIdx.x) * 4 * hidden_size;
+  int64_t const output_offset = static_cast<int64_t>(blockIdx.x) * hidden_size;
+  if ((hidden_size & 1) == 0) {
+    auto const* residual_pairs =
+        reinterpret_cast<__nv_bfloat162 const*>(residual + token_offset);
+    auto* output_pairs =
+        reinterpret_cast<__nv_bfloat162*>(output + output_offset);
+    int64_t const pair_count = hidden_size / 2;
+    for (int64_t hidden_pair = threadIdx.x; hidden_pair < pair_count;
+         hidden_pair += blockDim.x) {
+      // store two bf16 to a 32-bit reg
+      float2 const value_0 = __bfloat1622float2(residual_pairs[hidden_pair]);
+      float2 const value_1 =
+          __bfloat1622float2(residual_pairs[pair_count + hidden_pair]);
+      float2 const value_2 =
+          __bfloat1622float2(residual_pairs[2 * pair_count + hidden_pair]);
+      float2 const value_3 =
+          __bfloat1622float2(residual_pairs[3 * pair_count + hidden_pair]);
+      float2 result;
+      float const left_x = __fadd_rn(__fmul_rn(weights[0], value_0.x),
+                                     __fmul_rn(weights[1], value_1.x));
+      float const right_x = __fadd_rn(__fmul_rn(weights[2], value_2.x),
+                                      __fmul_rn(weights[3], value_3.x));
+      result.x = __fadd_rn(left_x, right_x);
+      float const left_y = __fadd_rn(__fmul_rn(weights[0], value_0.y),
+                                     __fmul_rn(weights[1], value_1.y));
+      float const right_y = __fadd_rn(__fmul_rn(weights[2], value_2.y),
+                                      __fmul_rn(weights[3], value_3.y));
+      result.y = __fadd_rn(left_y, right_y);
+      output_pairs[hidden_pair] = __floats2bfloat162_rn(result.x, result.y);
+    }
+  } else {
+    for (int64_t hidden = threadIdx.x; hidden < hidden_size;
+         hidden += blockDim.x) {
+      float const product_0 = __fmul_rn(
+          weights[0], __bfloat162float(residual[token_offset + hidden]));
+      float const product_1 = __fmul_rn(
+          weights[1],
+          __bfloat162float(residual[token_offset + hidden_size + hidden]));
+      float const product_2 = __fmul_rn(
+          weights[2], __bfloat162float(
+                          residual[token_offset + 2 * hidden_size + hidden]));
+      float const product_3 = __fmul_rn(
+          weights[3], __bfloat162float(
+                          residual[token_offset + 3 * hidden_size + hidden]));
+      float const left = __fadd_rn(product_0, product_1);
+      float const right = __fadd_rn(product_2, product_3);
+      output[output_offset + hidden] = __float2bfloat16_rn(__fadd_rn(left, right));
+    }
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
+inline cudaError_t launch_mhc_pre_h_aggregate(
+    __nv_bfloat16 const* residual, float const* pre, __nv_bfloat16* output,
+    int64_t num_tokens, int64_t hidden_size, cudaStream_t stream,
+    bool enable_pdl) {
+  cudaLaunchConfig_t config{};
+  config.gridDim = dim3(static_cast<unsigned int>(num_tokens));
+  config.blockDim = dim3(num_tokens <= 128 ? kMhcPreHAggregateDecodeThreads
+                                           : kMhcPreHAggregateBatchThreads);
+  config.stream = stream;
+
+  cudaLaunchAttribute attribute{};
+  if (enable_pdl) {
+    attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attribute.val.programmaticStreamSerializationAllowed = 1;
+    config.attrs = &attribute;
+    config.numAttrs = 1;
+  }
+
+  return cudaLaunchKernelEx(&config, mhc_pre_h_aggregate_kernel, residual, pre,
+                            output, hidden_size);
+}
+
+}

csrc/ops.cpp

@@ -164,6 +164,10 @@ int64_t rmsnorm_backward_dw_chunks_cuda(int64_t rows);
 void reduce_rows_fp32_left_fold_cuda(
   torch::Tensor rows,
   torch::Tensor output);
+
+torch::Tensor mhc_pre_h_aggregate_cuda(
+  torch::Tensor residual,
+  torch::Tensor pre);
 #endif
 
 static void rmsnorm_check_input(const torch::Tensor& x, const char* name) {
@@ -491,6 +495,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         "reduce_rows_fp32_left_fold",
         &reduce_rows_fp32_left_fold,
         "Ascending-row FP32 left-fold reduction CUDA");
+    m.def(
+        "mhc_pre_h_aggregate",
+        &mhc_pre_h_aggregate_cuda,
+        "Batch-invariant MHC H Aggregate CUDA");
 #endif
 
     // registry SiLU / SwiGLU (elementwise activation)

rl_engine/_C.pyi

@@ -2,6 +2,10 @@
 # This file is a type stub for the compiled C++ extension module.
 import torch
 
+def mhc_pre_h_aggregate(
+    residual: torch.Tensor,
+    pre: torch.Tensor,
+) -> torch.Tensor: ...
 def deterministic_collective_ipc_meta(
     tensor: torch.Tensor,
 ) -> tuple[list[int], int]: ...

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

setup.py

@@ -145,6 +145,7 @@ def get_extensions():
             # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync).
             # The ROCm dispatcher falls back to PyTorch SDPA for this operator.
             cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu")
+            cuda_sources.append("csrc/cuda/mhc/mhc_pre_h_aggregate.cu")
 
         nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"]
         if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH):