From 9d5732be1759b08a9ecbab84b5d98092d0b078ef Mon Sep 17 00:00:00 2001 From: Takumi <3051000145@qq.com> Date: Sun, 30 Aug 2026 20:37:25 +0800 Subject: [PATCH 1/2] Optimize deterministic rollout tensor-parallel all-reduce (#365) --- .../distributed/deterministic_collective.cu | 17 +++-- rl_engine/integrations/vllm_runtime.py | 69 ++++++++++++++++++- 2 files changed, 79 insertions(+), 7 deletions(-) diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu index 72d3f874..41cea442 100644 --- a/csrc/cuda/distributed/deterministic_collective.cu +++ b/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( diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index c4da28be..ade13cab 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/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,6 +373,7 @@ 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.""" @@ -374,7 +381,7 @@ def _patch_qwen3_strict_model( 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 @@ -382,6 +389,7 @@ def _patch_qwen3_strict_model( 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 From e45f4f3c7544e9f1e01898e3057eed30dba21efd Mon Sep 17 00:00:00 2001 From: nodeeeeee Date: Sun, 30 Aug 2026 13:26:25 +0000 Subject: [PATCH 2/2] feat: add batch-invariant MHC H Aggregate kernel --- csrc/cuda/mhc/mhc_pre_h_aggregate.cu | 62 +++++++++++ csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh | 107 +++++++++++++++++++ csrc/ops.cpp | 8 ++ rl_engine/_C.pyi | 4 + setup.py | 1 + tests/test_mhc_pre_h_aggregate.py | 72 +++++++++++++ 6 files changed, 254 insertions(+) create mode 100644 csrc/cuda/mhc/mhc_pre_h_aggregate.cu create mode 100644 csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh create mode 100644 tests/test_mhc_pre_h_aggregate.py diff --git a/csrc/cuda/mhc/mhc_pre_h_aggregate.cu b/csrc/cuda/mhc/mhc_pre_h_aggregate.cu new file mode 100644 index 00000000..d109952e --- /dev/null +++ b/csrc/cuda/mhc/mhc_pre_h_aggregate.cu @@ -0,0 +1,62 @@ +#include +#include +#include +#include + +#include +#include + +#include +#include + +#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::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()); + auto const* pre_ptr = pre.data_ptr(); + auto* output_ptr = reinterpret_cast<__nv_bfloat16*>( + output.data_ptr()); + + 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; +} diff --git a/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh b/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh new file mode 100644 index 00000000..b9c13bbe --- /dev/null +++ b/csrc/cuda/mhc/mhc_pre_h_aggregate_kernel.cuh @@ -0,0 +1,107 @@ +#pragma once + +#include +#include + +#include + +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(blockIdx.x) * 4 + threadIdx.x]; + } + __syncthreads(); + + int64_t const token_offset = + static_cast(blockIdx.x) * 4 * hidden_size; + int64_t const output_offset = static_cast(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(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); +} + +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index aecc30ed..2887dea5 100644 --- a/csrc/ops.cpp +++ b/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) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8e0e865a..f2f64c8c 100644 --- a/rl_engine/_C.pyi +++ b/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]: ... diff --git a/setup.py b/setup.py index 79f882d9..b79aa3ea 100644 --- a/setup.py +++ b/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): diff --git a/tests/test_mhc_pre_h_aggregate.py b/tests/test_mhc_pre_h_aggregate.py new file mode 100644 index 00000000..6563e6a5 --- /dev/null +++ b/tests/test_mhc_pre_h_aggregate.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import pytest +import torch + + +def _kernel_available() -> bool: + if not torch.cuda.is_available(): + return False + if torch.cuda.get_device_capability()[0] < 8: + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + except Exception: + return False + return _EXT_AVAILABLE and hasattr(_C, "mhc_pre_h_aggregate") + + +requires_mhc_kernel = pytest.mark.skipif( + not _kernel_available(), + reason="mhc_pre_h_aggregate requires the CUDA extension on SM80 or newer", +) + + +def _same_bytes(left: torch.Tensor, right: torch.Tensor) -> bool: + return torch.equal(left.view(torch.uint8), right.view(torch.uint8)) + + +@requires_mhc_kernel +def test_mhc_pre_h_aggregate_is_batch_invariant_and_matches_pytorch(): + from rl_engine.kernels.ops.base import _C + + num_tokens = 129 + hidden_size = 4096 + num_runs = 100 + generator = torch.Generator(device="cuda").manual_seed(0) + residual = torch.randn( + num_tokens, + 4, + hidden_size, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + pre = torch.rand( + num_tokens, + 4, + dtype=torch.float32, + device="cuda", + generator=generator, + ) + + first = _C.mhc_pre_h_aggregate(residual, pre) + for _ in range(num_runs - 1): + repeated = _C.mhc_pre_h_aggregate(residual, pre) + assert _same_bytes(first, repeated) + + token_by_token = torch.cat( + [ + _C.mhc_pre_h_aggregate( + residual[token : token + 1], pre[token : token + 1] + ) + for token in range(num_tokens) + ] + ) + assert _same_bytes(first, token_by_token) + + reference = torch.sum( + pre.unsqueeze(-1) * residual.to(torch.float32), dim=1 + ).to(torch.bfloat16) + torch.testing.assert_close(first, reference, atol=5e-2, rtol=1e-2)