From 966d7123f5bcc5b4bb7e5ffee4c26b1e1dec3aa8 Mon Sep 17 00:00:00 2001 From: Codex H100 Validation Date: Sun, 30 Aug 2026 04:46:09 +0000 Subject: [PATCH] feat: publish strict deterministic runtime and collective optimization --- csrc/cuda/activation.cu | 113 ++ csrc/cuda/batch_invariant_logp_kernel_sm90.cu | 22 +- .../distributed/deterministic_collective.cu | 983 +++++++++++++++++- csrc/cuda/gemm/det_gemm_kernel.cu | 8 + csrc/ops.cpp | 33 + .../aligned_python_entrypoint.sh | 41 + rl_engine/_C.pyi | 7 + rl_engine/distributed/collectives.py | 197 +++- rl_engine/integrations/framework_operators.py | 878 +++++++++++----- rl_engine/integrations/linear_logp.py | 113 +- rl_engine/integrations/megatron_runtime.py | 109 +- rl_engine/integrations/runtime.py | 139 ++- rl_engine/integrations/vllm_runtime.py | 289 ++++- .../kernels/ops/cuda/attention/cp_comm.py | 14 +- .../kernels/ops/cuda/attention/flash_attn.py | 173 ++- .../ops/cuda/attention/strict_runtime.py | 190 ++-- .../kernels/ops/cuda/loss/linear_logp.py | 124 ++- rl_engine/kernels/ops/cuda/matmul/det_gemm.py | 299 +++++- rl_engine/kernels/ops/pytorch/ffn/ffn.py | 452 ++++++-- .../kernels/ops/pytorch/loss/linear_logp.py | 102 +- .../kernels/ops/pytorch/norm/rms_norm.py | 70 ++ rl_engine/runtime_mode.py | 98 ++ tests/test_framework_runtime_adapters.py | 15 +- tests/test_qwen_ffn.py | 5 +- 24 files changed, 3848 insertions(+), 626 deletions(-) create mode 100755 examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh create mode 100644 rl_engine/runtime_mode.py diff --git a/csrc/cuda/activation.cu b/csrc/cuda/activation.cu index 8f2c62399..505654f69 100644 --- a/csrc/cuda/activation.cu +++ b/csrc/cuda/activation.cu @@ -93,6 +93,47 @@ __global__ void swiglu_backward_kernel( d_gate[idx] = static_cast(dyv * uv * silu_grad_f32(gv)); } +template +__global__ void swiglu_packed_forward_kernel( + const scalar_t* __restrict__ gate_up, + scalar_t* __restrict__ y, + const int64_t n, + const int64_t width) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const int64_t row = idx / width; + const int64_t column = idx - row * width; + const int64_t gate_index = row * (2 * width) + column; + const float gv = static_cast(gate_up[gate_index]); + const float uv = static_cast(gate_up[gate_index + width]); + y[idx] = static_cast(silu_f32(gv) * uv); +} + +template +__global__ void swiglu_packed_backward_kernel( + const scalar_t* __restrict__ dy, + const scalar_t* __restrict__ gate_up, + scalar_t* __restrict__ d_gate, + scalar_t* __restrict__ d_up, + const int64_t n, + const int64_t width) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const int64_t row = idx / width; + const int64_t column = idx - row * width; + const int64_t gate_index = row * (2 * width) + column; + const float dyv = static_cast(dy[idx]); + const float gv = static_cast(gate_up[gate_index]); + const float uv = static_cast(gate_up[gate_index + width]); + const float s = silu_f32(gv); + d_up[idx] = static_cast(dyv * s); + d_gate[idx] = static_cast(dyv * uv * silu_grad_f32(gv)); +} + static void launch_1d(int64_t n, int& threads, int64_t& blocks) { threads = 256; blocks = (n + threads - 1) / threads; @@ -251,3 +292,75 @@ std::vector swiglu_backward_cuda( C10_CUDA_KERNEL_LAUNCH_CHECK(); return {d_gate, d_up}; } + +torch::Tensor swiglu_packed_forward_cuda(torch::Tensor gate_up) { + check_cuda_contig(gate_up, "gate_up"); + TORCH_CHECK(gate_up.dim() == 2, "gate_up must be [rows, 2 * intermediate]"); + TORCH_CHECK(gate_up.size(1) % 2 == 0, "gate_up width must be even"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(gate_up)); + const int64_t width = gate_up.size(1) / 2; + auto y = torch::empty({gate_up.size(0), width}, gate_up.options()); + const int64_t n = y.numel(); + if (n == 0) { + return y; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + gate_up.scalar_type(), + "swiglu_packed_forward_cuda", + [&] { + swiglu_packed_forward_kernel<<>>( + gate_up.data_ptr(), y.data_ptr(), n, width); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return y; +} + +std::vector swiglu_packed_backward_cuda( + torch::Tensor dy, + torch::Tensor gate_up) { + check_cuda_contig(dy, "dy"); + check_cuda_contig(gate_up, "gate_up"); + check_same_device(dy, gate_up, "dy", "gate_up"); + TORCH_CHECK(gate_up.dim() == 2, "gate_up must be [rows, 2 * intermediate]"); + TORCH_CHECK(gate_up.size(1) % 2 == 0, "gate_up width must be even"); + TORCH_CHECK(dy.dim() == 2, "dy must be [rows, intermediate]"); + TORCH_CHECK( + dy.size(0) == gate_up.size(0) && dy.size(1) * 2 == gate_up.size(1), + "dy shape must match the packed gate/up halves"); + TORCH_CHECK(dy.scalar_type() == gate_up.scalar_type(), "dy and gate_up must share dtype"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(gate_up)); + auto d_gate = torch::empty_like(dy); + auto d_up = torch::empty_like(dy); + const int64_t n = dy.numel(); + if (n == 0) { + return {d_gate, d_up}; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + gate_up.scalar_type(), + "swiglu_packed_backward_cuda", + [&] { + swiglu_packed_backward_kernel<<>>( + dy.data_ptr(), + gate_up.data_ptr(), + d_gate.data_ptr(), + d_up.data_ptr(), + n, + dy.size(1)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {d_gate, d_up}; +} diff --git a/csrc/cuda/batch_invariant_logp_kernel_sm90.cu b/csrc/cuda/batch_invariant_logp_kernel_sm90.cu index a6c67db7d..fb1f15331 100644 --- a/csrc/cuda/batch_invariant_logp_kernel_sm90.cu +++ b/csrc/cuda/batch_invariant_logp_kernel_sm90.cu @@ -18,6 +18,20 @@ namespace { #define TMA_BOX 256 static constexpr int LOADS_PER_TILE = SMEM_TILE / TMA_BOX; +// CUDA 13's CCCL no longer exposes cub::Max/Sum as public functors. Keep the +// reduction tree explicit so the SM90 build remains portable and deterministic. +struct DeterministicMax { + __device__ __forceinline__ float operator()(float lhs, float rhs) const { + return lhs > rhs ? lhs : rhs; + } +}; + +struct DeterministicSum { + __device__ __forceinline__ float operator()(float lhs, float rhs) const { + return lhs + rhs; + } +}; + template __device__ __forceinline__ float to_float(T x); template <> __device__ __forceinline__ float to_float(nv_bfloat16 x) { return __bfloat162float(x); @@ -81,7 +95,7 @@ __global__ void batch_invariant_logp_sm90_kernel(const __grid_constant__ CUtenso for (int i = tid; i < current_tile_size; i += NUM_THREADS) { tile_max = max(tile_max, to_float(smem_logits[i])); } - float block_tile_max = BlockReduce(temp_storage).Reduce(tile_max, cub::Max()); + float block_tile_max = BlockReduce(temp_storage).Reduce(tile_max, DeterministicMax()); if (tid == 0) s_tile_max = block_tile_max; __syncthreads(); @@ -91,7 +105,7 @@ __global__ void batch_invariant_logp_sm90_kernel(const __grid_constant__ CUtenso for (int i = tid; i < current_tile_size; i += NUM_THREADS) { tile_sum += expf(to_float(smem_logits[i]) - s_tile_max); } - float block_tile_sum = BlockReduce(temp_storage).Reduce(tile_sum, cub::Sum()); + float block_tile_sum = BlockReduce(temp_storage).Reduce(tile_sum, DeterministicSum()); // Online log-sum-exp merge of this tile into the running row state. if (tid == 0) { @@ -110,7 +124,7 @@ __global__ void batch_invariant_logp_sm90_kernel(const __grid_constant__ CUtenso for (int i = tid; i < tail; i += NUM_THREADS) { tail_max = max(tail_max, to_float(logits_gmem[base + i])); } - float block_tail_max = BlockReduce(temp_storage).Reduce(tail_max, cub::Max()); + float block_tail_max = BlockReduce(temp_storage).Reduce(tail_max, DeterministicMax()); if (tid == 0) s_tile_max = block_tail_max; __syncthreads(); @@ -119,7 +133,7 @@ __global__ void batch_invariant_logp_sm90_kernel(const __grid_constant__ CUtenso for (int i = tid; i < tail; i += NUM_THREADS) { tail_sum += expf(to_float(logits_gmem[base + i]) - s_tile_max); } - float block_tail_sum = BlockReduce(temp_storage).Reduce(tail_sum, cub::Sum()); + float block_tail_sum = BlockReduce(temp_storage).Reduce(tail_sum, DeterministicSum()); if (tid == 0) { const float new_max = max(row_max, s_tile_max); row_sum = row_sum * expf(row_max - new_max) + block_tail_sum * expf(s_tile_max - new_max); diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu index 63a21f94c..72d3f8745 100644 --- a/csrc/cuda/distributed/deterministic_collective.cu +++ b/csrc/cuda/distributed/deterministic_collective.cu @@ -13,10 +13,14 @@ #include #include +#include #include #include #include +#include #include +#include +#include #include namespace { @@ -24,15 +28,371 @@ namespace { constexpr int kMaxDeterministicWorldSize = 8; constexpr int kThreads = 256; constexpr int kMaxBlocks = 4096; +constexpr int kStagingFrames = 3; +constexpr int kFusedStagingSlots = 2; +constexpr int64_t kSequenceHeaderBytes = 3 * sizeof(uint64_t); +// Small tensor collectives are launch-bound. Keep the existing DMA path for +// larger transfers, where replacing cudaMemcpyAsync with a one-block copy +// would reduce bandwidth and hurt overlap. +constexpr int64_t kSingleBlockFastPathMaxBytes = 256 * 1024; +// The trace's hot grid=64 collectives are dominated by asymmetric remote +// reads. Reduce medium payloads once on logical rank 0 and distribute the +// canonical result, while leaving true large transfers on the established +// parallel path. +constexpr int64_t kOwnerReduceMaxBytes = 4 * 1024 * 1024; struct PeerPointers { const void* values[kMaxDeterministicWorldSize]; + const uint64_t* stage_sequences[kMaxDeterministicWorldSize]; + const uint64_t* done_sequences[kMaxDeterministicWorldSize]; +}; + +struct FusedPeerPointers { + PeerPointers slots[kFusedStagingSlots]; }; bool is_supported_world_size(int64_t world_size) { return world_size == 1 || world_size == 2 || world_size == 4 || world_size == 8; } +__device__ __forceinline__ uint64_t load_acquire_system( + const uint64_t* address) { + uint64_t value; + asm volatile( + "ld.acquire.sys.global.u64 %0, [%1];" + : "=l"(value) + : "l"(address) + : "memory"); + return value; +} + +__device__ __forceinline__ void store_release_system( + uint64_t* address, + uint64_t value) { + asm volatile( + "st.release.sys.global.u64 [%0], %1;" + : + : "l"(address), "l"(value) + : "memory"); +} + +__global__ void wait_for_previous_done_kernel( + PeerPointers peers, + int world_size, + const uint64_t* local_stage_sequence) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + const uint64_t sequence = load_acquire_system(local_stage_sequence); + for (int peer = 0; peer < world_size; ++peer) { + while (load_acquire_system(peers.done_sequences[peer]) < sequence) { + __nanosleep(64); + } + } +} + +__global__ void publish_next_stage_sequence_kernel(uint64_t* stage_sequence) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + const uint64_t current = load_acquire_system(stage_sequence); + store_release_system(stage_sequence, current + 1); +} + +__global__ void wait_for_staged_peers_kernel( + PeerPointers peers, + int world_size, + const uint64_t* local_stage_sequence) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + const uint64_t sequence = load_acquire_system(local_stage_sequence); + for (int peer = 0; peer < world_size; ++peer) { + while (load_acquire_system(peers.stage_sequences[peer]) < sequence) { + __nanosleep(64); + } + } +} + +__global__ void publish_done_sequence_kernel( + uint64_t* done_sequence, + const uint64_t* stage_sequence) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + store_release_system(done_sequence, load_acquire_system(stage_sequence)); +} + +__global__ void wait_for_owner_done_kernel( + PeerPointers peers, + const uint64_t* local_stage_sequence) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + const uint64_t sequence = load_acquire_system(local_stage_sequence); + while (load_acquire_system(peers.done_sequences[0]) < sequence) { + __nanosleep(64); + } +} + +__device__ __forceinline__ void wait_for_stage_sequence( + const PeerPointers& peers, + int world_size, + const uint64_t* local_stage_sequence) { + const uint64_t sequence = load_acquire_system(local_stage_sequence); + for (int peer = 0; peer < world_size; ++peer) { + while (load_acquire_system(peers.stage_sequences[peer]) < sequence) { + __nanosleep(64); + } + } +} + +// The launch-bound path keeps the protocol and the payload operation in one +// block. It is used only for small messages, where the copy cost is negligible +// compared with two extra control-kernel launches. +__global__ void stage_payload_fast_kernel( + PeerPointers peers, + int world_size, + uint64_t* local_stage_sequence, + const uint8_t* input, + uint8_t* payload, + int64_t input_bytes) { + if (threadIdx.x == 0) { + const uint64_t sequence = load_acquire_system(local_stage_sequence); + for (int peer = 0; peer < world_size; ++peer) { + while (load_acquire_system(peers.done_sequences[peer]) < sequence) { + __nanosleep(64); + } + } + } + __syncthreads(); + + if (((reinterpret_cast(input) | + reinterpret_cast(payload) | + static_cast(input_bytes)) & 15u) == 0u) { + const auto* source = reinterpret_cast(input); + auto* destination = reinterpret_cast(payload); + const int64_t vector_count = input_bytes / sizeof(uint4); + for (int64_t index = threadIdx.x; index < vector_count; index += blockDim.x) { + destination[index] = source[index]; + } + } else { + for (int64_t index = threadIdx.x; index < input_bytes; index += blockDim.x) { + payload[index] = input[index]; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + __threadfence_system(); + store_release_system( + local_stage_sequence, + load_acquire_system(local_stage_sequence) + 1); + } +} + +template +__device__ __forceinline__ T fixed_tree_reduce( + const PeerPointers& peers, + int64_t index); + +template +__global__ void deterministic_all_reduce_fast_kernel( + PeerPointers peers, + const uint64_t* local_stage_sequence, + uint64_t* local_done_sequence, + T* output, + int64_t element_count) { + if (threadIdx.x == 0) { + wait_for_stage_sequence(peers, WorldSize, local_stage_sequence); + } + __syncthreads(); + + for (int64_t index = threadIdx.x; index < element_count; index += blockDim.x) { + output[index] = fixed_tree_reduce(peers, index); + } + __syncthreads(); + + if (threadIdx.x == 0) { + __threadfence_system(); + store_release_system( + local_done_sequence, + load_acquire_system(local_stage_sequence)); + } +} + +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template +__device__ __forceinline__ nv_bfloat162 fixed_tree_reduce_bf16x2( + const PeerPointers& peers, + int64_t pair_index) { + const auto* rank0 = static_cast(peers.values[0]); + if constexpr (WorldSize == 1) { + return rank0[pair_index]; + } else { + const auto* rank1 = static_cast(peers.values[1]); + const nv_bfloat162 sum01 = __hadd2(rank0[pair_index], rank1[pair_index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const auto* rank2 = static_cast(peers.values[2]); + const auto* rank3 = static_cast(peers.values[3]); + const nv_bfloat162 sum23 = __hadd2(rank2[pair_index], rank3[pair_index]); + const nv_bfloat162 sum03 = __hadd2(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const auto* rank4 = static_cast(peers.values[4]); + const auto* rank5 = static_cast(peers.values[5]); + const auto* rank6 = static_cast(peers.values[6]); + const auto* rank7 = static_cast(peers.values[7]); + const nv_bfloat162 sum45 = __hadd2(rank4[pair_index], rank5[pair_index]); + const nv_bfloat162 sum67 = __hadd2(rank6[pair_index], rank7[pair_index]); + const nv_bfloat162 sum47 = __hadd2(sum45, sum67); + return __hadd2(sum03, sum47); + } + } + } +} +#endif + +// Fuse the small-message stage and fixed-tree reduction into one launch. Two +// IPC payload slots let a faster rank publish operation N+1 while a peer is +// still retiring N; a slot is only reused after every peer has completed its +// prior use. The canonical reduction order itself is unchanged. +template +__global__ void deterministic_all_reduce_fused_fast_kernel( + FusedPeerPointers peer_slots, + uint64_t* local_next_sequence, + int rank, + const uint8_t* input, + T* output, + int64_t element_count, + int64_t input_bytes) { + __shared__ uint64_t sequence; + __shared__ int slot; + if (threadIdx.x == 0) { + sequence = load_acquire_system(local_next_sequence) + 1; + slot = static_cast(sequence & 1u); + if (sequence > kFusedStagingSlots) { + const uint64_t reuse_after = sequence - kFusedStagingSlots; + const PeerPointers& reuse_peers = peer_slots.slots[slot]; + for (int peer = 0; peer < WorldSize; ++peer) { + while (load_acquire_system(reuse_peers.done_sequences[peer]) < + reuse_after) { + __nanosleep(64); + } + } + } + } + __syncthreads(); + + const PeerPointers peers = peer_slots.slots[slot]; + auto* payload = const_cast( + static_cast(peers.values[rank])); + + if (((reinterpret_cast(input) | + reinterpret_cast(payload) | + static_cast(input_bytes)) & 15u) == 0u) { + const auto* source = reinterpret_cast(input); + auto* destination = reinterpret_cast(payload); + const int64_t vector_count = input_bytes / sizeof(uint4); + for (int64_t index = threadIdx.x; index < vector_count; index += blockDim.x) { + destination[index] = source[index]; + } + } else { + for (int64_t index = threadIdx.x; index < input_bytes; index += blockDim.x) { + payload[index] = input[index]; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + __threadfence_system(); + store_release_system( + const_cast(peers.stage_sequences[rank]), sequence); + for (int peer = 0; peer < WorldSize; ++peer) { + while (load_acquire_system(peers.stage_sequences[peer]) < sequence) { + __nanosleep(64); + } + } + } + __syncthreads(); + + if constexpr (std::is_same_v) { +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + const int64_t pair_count = element_count / 2; + auto* pair_output = reinterpret_cast(output); + for (int64_t pair_index = threadIdx.x; pair_index < pair_count; + pair_index += blockDim.x) { + pair_output[pair_index] = + fixed_tree_reduce_bf16x2(peers, pair_index); + } + if ((element_count & 1) != 0 && threadIdx.x == 0) { + output[element_count - 1] = + fixed_tree_reduce(peers, element_count - 1); + } +#endif + } else { + for (int64_t index = threadIdx.x; index < element_count; index += blockDim.x) { + output[index] = fixed_tree_reduce(peers, index); + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + __threadfence_system(); + store_release_system( + const_cast(peers.done_sequences[rank]), sequence); + store_release_system(local_next_sequence, sequence); + } +} + +template +__global__ void deterministic_reduce_scatter_fast_kernel( + PeerPointers peers, + const uint64_t* local_stage_sequence, + uint64_t* local_done_sequence, + T* output, + int64_t output_element_count, + int rank) { + if (threadIdx.x == 0) { + wait_for_stage_sequence(peers, WorldSize, local_stage_sequence); + } + __syncthreads(); + + const int64_t input_offset = static_cast(rank) * output_element_count; + for (int64_t index = threadIdx.x; index < output_element_count; index += blockDim.x) { + output[index] = fixed_tree_reduce(peers, input_offset + index); + } + __syncthreads(); + + if (threadIdx.x == 0) { + __threadfence_system(); + store_release_system( + local_done_sequence, + load_acquire_system(local_stage_sequence)); + } +} + +__global__ void deterministic_all_gather_fast_kernel( + PeerPointers peers, + int world_size, + const uint64_t* local_stage_sequence, + uint64_t* local_done_sequence, + uint8_t* output, + int64_t input_bytes) { + if (threadIdx.x == 0) { + wait_for_stage_sequence(peers, world_size, local_stage_sequence); + } + __syncthreads(); + + const int64_t output_bytes = input_bytes * world_size; + for (int64_t index = threadIdx.x; index < output_bytes; index += blockDim.x) { + const int peer = static_cast(index / input_bytes); + const int64_t peer_offset = index - static_cast(peer) * input_bytes; + output[index] = static_cast(peers.values[peer])[peer_offset]; + } + __syncthreads(); + + if (threadIdx.x == 0) { + __threadfence_system(); + store_release_system( + local_done_sequence, + load_acquire_system(local_stage_sequence)); + } +} + template __device__ __forceinline__ T ordered_add(T lower, T upper); @@ -136,6 +496,177 @@ void launch_all_reduce( } } +template +void launch_all_reduce_fast( + const PeerPointers& peers, + const uint64_t* local_stage_sequence, + uint64_t* local_done_sequence, + T* output, + int64_t element_count, + int64_t world_size, + cudaStream_t stream) { + switch (world_size) { + case 1: + deterministic_all_reduce_fast_kernel<<<1, kThreads, 0, stream>>>( + peers, + local_stage_sequence, + local_done_sequence, + output, + element_count); + break; + case 2: + deterministic_all_reduce_fast_kernel<<<1, kThreads, 0, stream>>>( + peers, + local_stage_sequence, + local_done_sequence, + output, + element_count); + break; + case 4: + deterministic_all_reduce_fast_kernel<<<1, kThreads, 0, stream>>>( + peers, + local_stage_sequence, + local_done_sequence, + output, + element_count); + break; + case 8: + deterministic_all_reduce_fast_kernel<<<1, kThreads, 0, stream>>>( + peers, + local_stage_sequence, + local_done_sequence, + output, + element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } +} + +template +void launch_all_reduce_fused_fast( + const FusedPeerPointers& peer_slots, + uint64_t* local_next_sequence, + int rank, + const uint8_t* input, + T* output, + int64_t element_count, + int64_t input_bytes, + int64_t world_size, + cudaStream_t stream) { + switch (world_size) { + case 1: + deterministic_all_reduce_fused_fast_kernel + <<<1, kThreads, 0, stream>>>( + peer_slots, local_next_sequence, rank, input, + output, element_count, input_bytes); + break; + case 2: + deterministic_all_reduce_fused_fast_kernel + <<<1, kThreads, 0, stream>>>( + peer_slots, local_next_sequence, rank, input, + output, element_count, input_bytes); + break; + case 4: + deterministic_all_reduce_fused_fast_kernel + <<<1, kThreads, 0, stream>>>( + peer_slots, local_next_sequence, rank, input, + output, element_count, input_bytes); + break; + case 8: + deterministic_all_reduce_fused_fast_kernel + <<<1, kThreads, 0, stream>>>( + peer_slots, local_next_sequence, rank, input, + output, element_count, input_bytes); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } +} + +__global__ void deterministic_all_gather_fused_fast_kernel( + PeerPointers peers, + int world_size, + uint64_t* local_stage_sequence, + uint64_t* local_done_sequence, + const uint8_t* input, + uint8_t* payload, + uint8_t* output, + int64_t input_bytes) { + uint64_t sequence = 0; + if (threadIdx.x == 0) { + sequence = load_acquire_system(local_stage_sequence); + for (int peer = 0; peer < world_size; ++peer) { + while (load_acquire_system(peers.done_sequences[peer]) < sequence) { + __nanosleep(64); + } + } + } + __syncthreads(); + + if (((reinterpret_cast(input) | + reinterpret_cast(payload) | + static_cast(input_bytes)) & 15u) == 0u) { + const auto* source = reinterpret_cast(input); + auto* destination = reinterpret_cast(payload); + const int64_t vector_count = input_bytes / sizeof(uint4); + for (int64_t index = threadIdx.x; index < vector_count; index += blockDim.x) { + destination[index] = source[index]; + } + } else { + for (int64_t index = threadIdx.x; index < input_bytes; index += blockDim.x) { + payload[index] = input[index]; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + __threadfence_system(); + store_release_system(local_stage_sequence, sequence + 1); + sequence += 1; + for (int peer = 0; peer < world_size; ++peer) { + while (load_acquire_system(peers.stage_sequences[peer]) < sequence) { + __nanosleep(64); + } + } + } + __syncthreads(); + + const int64_t output_bytes = input_bytes * world_size; + for (int64_t index = threadIdx.x; index < output_bytes; index += blockDim.x) { + const int peer = static_cast(index / input_bytes); + const int64_t peer_offset = index - static_cast(peer) * input_bytes; + output[index] = static_cast(peers.values[peer])[peer_offset]; + } + __syncthreads(); + + if (threadIdx.x == 0) { + __threadfence_system(); + store_release_system(local_done_sequence, sequence); + } +} + +void launch_all_gather_fused_fast( + const PeerPointers& peers, + uint64_t* local_stage_sequence, + uint64_t* local_done_sequence, + const uint8_t* input, + uint8_t* payload, + uint8_t* output, + int64_t input_bytes, + int64_t world_size, + cudaStream_t stream) { + deterministic_all_gather_fused_fast_kernel<<<1, kThreads, 0, stream>>>( + peers, + static_cast(world_size), + local_stage_sequence, + local_done_sequence, + input, + payload, + output, + input_bytes); +} + template __global__ void deterministic_reduce_scatter_kernel( PeerPointers peers, @@ -182,6 +713,62 @@ void launch_reduce_scatter( } } +template +void launch_reduce_scatter_fast( + const PeerPointers& peers, + const uint64_t* local_stage_sequence, + uint64_t* local_done_sequence, + T* output, + int64_t output_element_count, + int rank, + int64_t world_size, + cudaStream_t stream) { + switch (world_size) { + case 1: + deterministic_reduce_scatter_fast_kernel + <<<1, kThreads, 0, stream>>>( + peers, + local_stage_sequence, + local_done_sequence, + output, + output_element_count, + rank); + break; + case 2: + deterministic_reduce_scatter_fast_kernel + <<<1, kThreads, 0, stream>>>( + peers, + local_stage_sequence, + local_done_sequence, + output, + output_element_count, + rank); + break; + case 4: + deterministic_reduce_scatter_fast_kernel + <<<1, kThreads, 0, stream>>>( + peers, + local_stage_sequence, + local_done_sequence, + output, + output_element_count, + rank); + break; + case 8: + deterministic_reduce_scatter_fast_kernel + <<<1, kThreads, 0, stream>>>( + peers, + local_stage_sequence, + local_done_sequence, + output, + output_element_count, + rank); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } +} + __global__ void deterministic_all_gather_kernel( PeerPointers peers, uint8_t* output, @@ -208,13 +795,23 @@ class DeterministicCollectiveState { : rank_(rank), world_size_(handles.size()), device_index_(staging.get_device()), - capacity_bytes_(staging.numel() * staging.element_size()) { + frame_bytes_( + staging.numel() * staging.element_size() / kStagingFrames), + capacity_bytes_( + frame_bytes_ - kSequenceHeaderBytes) { TORCH_CHECK(staging.is_cuda(), "collective staging buffer must be CUDA"); TORCH_CHECK(staging.is_contiguous(), "collective staging buffer must be contiguous"); TORCH_CHECK( staging.scalar_type() == torch::kUInt8, "collective staging buffer must have dtype torch.uint8"); - TORCH_CHECK(capacity_bytes_ > 0, "collective staging capacity must be positive"); + TORCH_CHECK( + staging.numel() * staging.element_size() % kStagingFrames == 0, + "collective staging allocation must contain exactly ", + kStagingFrames, + " equal frames"); + TORCH_CHECK( + capacity_bytes_ > 0, + "collective staging must reserve sequence metadata and positive payload capacity"); TORCH_CHECK( is_supported_world_size(world_size_), "deterministic collectives require world size 1, 2, 4, or 8; got ", @@ -228,8 +825,15 @@ class DeterministicCollectiveState { world_size_, ")"); - for (auto& peer : peers_.values) { - peer = nullptr; + for (int peer = 0; peer < kMaxDeterministicWorldSize; ++peer) { + peers_.values[peer] = nullptr; + peers_.stage_sequences[peer] = nullptr; + peers_.done_sequences[peer] = nullptr; + for (int slot = 0; slot < kFusedStagingSlots; ++slot) { + fused_peer_slots_.slots[slot].values[peer] = nullptr; + fused_peer_slots_.slots[slot].stage_sequences[peer] = nullptr; + fused_peer_slots_.slots[slot].done_sequences[peer] = nullptr; + } } imported_bases_.fill(nullptr); try { @@ -240,29 +844,51 @@ class DeterministicCollectiveState { peer); TORCH_CHECK(offsets[peer] >= 0, "negative CUDA IPC offset for rank ", peer); + void* allocation = nullptr; if (peer == rank_) { - peers_.values[peer] = staging.data_ptr(); - continue; - } + allocation = staging.data_ptr(); + } else { + cudaIpcMemHandle_t handle{}; + auto* raw_handle = reinterpret_cast(&handle); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + TORCH_CHECK( + handles[peer][byte] >= 0 && handles[peer][byte] <= 255, + "invalid CUDA IPC handle byte for rank ", + peer); + raw_handle[byte] = static_cast(handles[peer][byte]); + } - cudaIpcMemHandle_t handle{}; - auto* raw_handle = reinterpret_cast(&handle); - for (size_t byte = 0; byte < sizeof(handle); ++byte) { - TORCH_CHECK( - handles[peer][byte] >= 0 && handles[peer][byte] <= 255, - "invalid CUDA IPC handle byte for rank ", - peer); - raw_handle[byte] = static_cast(handles[peer][byte]); + void* base = nullptr; + AT_CUDA_CHECK(cudaIpcOpenMemHandle( + &base, + handle, + cudaIpcMemLazyEnablePeerAccess)); + imported_bases_[peer] = base; + allocation = static_cast(base) + offsets[peer]; + } + auto* bytes = static_cast(allocation); + peers_.stage_sequences[peer] = + reinterpret_cast(bytes); + peers_.done_sequences[peer] = + reinterpret_cast(bytes + sizeof(uint64_t)); + peers_.values[peer] = bytes + kSequenceHeaderBytes; + for (int slot = 0; slot < kFusedStagingSlots; ++slot) { + auto* frame = bytes + (slot + 1) * frame_bytes_; + fused_peer_slots_.slots[slot].stage_sequences[peer] = + reinterpret_cast(frame); + fused_peer_slots_.slots[slot].done_sequences[peer] = + reinterpret_cast(frame + sizeof(uint64_t)); + fused_peer_slots_.slots[slot].values[peer] = + frame + kSequenceHeaderBytes; } - - void* base = nullptr; - AT_CUDA_CHECK(cudaIpcOpenMemHandle( - &base, - handle, - cudaIpcMemLazyEnablePeerAccess)); - imported_bases_[peer] = base; - peers_.values[peer] = static_cast(base) + offsets[peer]; } + auto* local_bytes = static_cast(staging.data_ptr()); + local_stage_sequence_ = reinterpret_cast(local_bytes); + local_done_sequence_ = + reinterpret_cast(local_bytes + sizeof(uint64_t)); + local_fused_sequence_ = + reinterpret_cast(local_bytes + 2 * sizeof(uint64_t)); + } catch (...) { close_imports(); throw; @@ -291,20 +917,39 @@ class DeterministicCollectiveState { input_bytes, " bytes but staging capacity is ", capacity_bytes_); - if (input_bytes > 0) { - AT_CUDA_CHECK(cudaMemcpyAsync( - const_cast(peers_.values[rank_]), - input.data_ptr(), - input_bytes, - cudaMemcpyDeviceToDevice, - stream)); + staged_fast_path_ = input_bytes <= kSingleBlockFastPathMaxBytes; + staged_owner_path_ = input_bytes <= kOwnerReduceMaxBytes; + if (staged_fast_path_) { + stage_payload_fast_kernel<<<1, kThreads, 0, stream>>>( + peers_, + world_size_, + local_stage_sequence_, + static_cast(input.data_ptr()), + const_cast(static_cast(peers_.values[rank_])), + input_bytes); + AT_CUDA_CHECK(cudaGetLastError()); + } else { + wait_for_previous_done_kernel<<<1, 1, 0, stream>>>( + peers_, world_size_, local_stage_sequence_); + AT_CUDA_CHECK(cudaGetLastError()); + if (input_bytes > 0) { + AT_CUDA_CHECK(cudaMemcpyAsync( + const_cast(peers_.values[rank_]), + input.data_ptr(), + input_bytes, + cudaMemcpyDeviceToDevice, + stream)); + } + publish_next_stage_sequence_kernel<<<1, 1, 0, stream>>>( + local_stage_sequence_); + AT_CUDA_CHECK(cudaGetLastError()); } staged_bytes_ = input_bytes; staged_scalar_type_ = input.scalar_type(); has_staged_input_ = true; } - void all_reduce(torch::Tensor& output, cudaStream_t stream) const { + void all_reduce(torch::Tensor& output, cudaStream_t stream) { check_tensor(output, "output"); TORCH_CHECK(has_staged_input_, "stage() must be called before all_reduce()"); TORCH_CHECK( @@ -315,7 +960,71 @@ class DeterministicCollectiveState { "all-reduce output size must match the staged input size"); const int64_t element_count = output.numel(); + if (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 + // pushes the canonical bytes into each follower's local IPC frame. + auto* owner_output = const_cast( + static_cast(peers_.values[0])); + wait_for_staged_peers(stream); + if (element_count > 0) { + const int blocks = static_cast(std::min( + kMaxBlocks, (element_count + kThreads - 1) / kThreads)); + switch (output.scalar_type()) { + case at::ScalarType::Float: + launch_all_reduce( + peers_, reinterpret_cast(owner_output), + element_count, blocks, world_size_, stream); + break; + case at::ScalarType::Half: + launch_all_reduce( + peers_, reinterpret_cast(owner_output), + element_count, blocks, world_size_, stream); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + launch_all_reduce( + peers_, reinterpret_cast(owner_output), + element_count, blocks, 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()); + for (int peer = 1; peer < world_size_; ++peer) { + AT_CUDA_CHECK(cudaMemcpyAsync( + const_cast(peers_.values[peer]), owner_output, + staged_bytes_, cudaMemcpyDeviceToDevice, stream)); + } + AT_CUDA_CHECK(cudaMemcpyAsync( + output.data_ptr(), owner_output, staged_bytes_, + cudaMemcpyDeviceToDevice, stream)); + } + publish_done(stream); + } else { + // Owner completion is published only after its push into this rank's + // local IPC frame. Followers therefore perform only a local copy. + wait_for_owner_done_kernel<<<1, 1, 0, stream>>>( + peers_, local_stage_sequence_); + AT_CUDA_CHECK(cudaGetLastError()); + if (staged_bytes_ > 0) { + AT_CUDA_CHECK(cudaMemcpyAsync( + output.data_ptr(), peers_.values[rank_], staged_bytes_, + cudaMemcpyDeviceToDevice, stream)); + } + publish_done(stream); + } + return; + } + + wait_for_staged_peers(stream); if (element_count == 0) { + publish_done(stream); return; } const int blocks = static_cast(std::min( @@ -359,9 +1068,71 @@ class DeterministicCollectiveState { output.scalar_type()); } AT_CUDA_CHECK(cudaGetLastError()); + publish_done(stream); + } + + void all_reduce_fused( + torch::Tensor& input, + torch::Tensor& output, + cudaStream_t stream) { + check_tensor(input, "input"); + check_tensor(output, "output"); + TORCH_CHECK(!has_staged_input_, "cannot fuse all-reduce with a pending stage()"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK( + input_bytes <= capacity_bytes_, + "input requires ", input_bytes, + " bytes but staging capacity is ", capacity_bytes_); + TORCH_CHECK( + output.scalar_type() == input.scalar_type(), + "all-reduce output dtype must match the input dtype"); + TORCH_CHECK( + 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. + stage(input, stream); + all_reduce(output, stream); } - void reduce_scatter(torch::Tensor& output, cudaStream_t stream) const { + void all_gather_fused( + torch::Tensor& input, + torch::Tensor& output, + cudaStream_t stream) { + check_tensor(input, "input"); + check_tensor(output, "output"); + TORCH_CHECK(!has_staged_input_, "cannot fuse all-gather with a pending stage()"); + const int64_t input_bytes = input.numel() * input.element_size(); + const int64_t output_bytes = output.numel() * output.element_size(); + TORCH_CHECK( + input_bytes <= capacity_bytes_, + "input requires ", input_bytes, + " bytes but staging capacity is ", capacity_bytes_); + TORCH_CHECK( + output_bytes == input_bytes * world_size_, + "all-gather output size must contain one input per rank"); + + if (output_bytes > kSingleBlockFastPathMaxBytes) { + stage(input, stream); + all_gather(output, stream); + return; + } + + launch_all_gather_fused_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(input.data_ptr()), + const_cast(static_cast(peers_.values[rank_])), + static_cast(output.data_ptr()), + input_bytes, + world_size_, + stream); + AT_CUDA_CHECK(cudaGetLastError()); + } + + void reduce_scatter(torch::Tensor& output, cudaStream_t stream) { check_tensor(output, "output"); TORCH_CHECK(has_staged_input_, "stage() must be called before reduce_scatter()"); TORCH_CHECK( @@ -372,7 +1143,57 @@ class DeterministicCollectiveState { "reduce-scatter output must contain one world-size fraction of the staged input"); const int64_t output_element_count = output.numel(); + if (staged_fast_path_) { + switch (output.scalar_type()) { + case at::ScalarType::Float: + launch_reduce_scatter_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + output_element_count, + static_cast(rank_), + world_size_, + stream); + break; + case at::ScalarType::Half: + launch_reduce_scatter_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + output_element_count, + static_cast(rank_), + world_size_, + stream); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + launch_reduce_scatter_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + output_element_count, + static_cast(rank_), + world_size_, + stream); + break; +#endif + default: + TORCH_CHECK( + false, + "deterministic reduce-scatter supports float32, float16, and bfloat16; got ", + output.scalar_type()); + } + AT_CUDA_CHECK(cudaGetLastError()); + has_staged_input_ = false; + return; + } + + wait_for_staged_peers(stream); if (output_element_count == 0) { + publish_done(stream); return; } const int blocks = static_cast(std::min( @@ -419,9 +1240,10 @@ class DeterministicCollectiveState { output.scalar_type()); } AT_CUDA_CHECK(cudaGetLastError()); + publish_done(stream); } - void all_gather(torch::Tensor& output, cudaStream_t stream) const { + void all_gather(torch::Tensor& output, cudaStream_t stream) { check_tensor(output, "output"); TORCH_CHECK(has_staged_input_, "stage() must be called before all_gather()"); TORCH_CHECK( @@ -433,7 +1255,22 @@ class DeterministicCollectiveState { "all-gather output must contain one staged input per rank"); const int64_t output_bytes = output.numel() * output.element_size(); + if (staged_fast_path_ && output_bytes <= kSingleBlockFastPathMaxBytes) { + deterministic_all_gather_fast_kernel<<<1, kThreads, 0, stream>>>( + peers_, + world_size_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + staged_bytes_); + AT_CUDA_CHECK(cudaGetLastError()); + has_staged_input_ = false; + return; + } + + wait_for_staged_peers(stream); if (output_bytes == 0) { + publish_done(stream); return; } const int blocks = static_cast(std::min( @@ -445,9 +1282,23 @@ class DeterministicCollectiveState { staged_bytes_, world_size_); AT_CUDA_CHECK(cudaGetLastError()); + publish_done(stream); } private: + void wait_for_staged_peers(cudaStream_t stream) const { + wait_for_staged_peers_kernel<<<1, 1, 0, stream>>>( + peers_, world_size_, local_stage_sequence_); + AT_CUDA_CHECK(cudaGetLastError()); + } + + void publish_done(cudaStream_t stream) { + publish_done_sequence_kernel<<<1, 1, 0, stream>>>( + local_done_sequence_, local_stage_sequence_); + AT_CUDA_CHECK(cudaGetLastError()); + has_staged_input_ = false; + } + void check_tensor(const torch::Tensor& tensor, const char* name) const { TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); @@ -472,17 +1323,36 @@ class DeterministicCollectiveState { int64_t rank_; int64_t world_size_; int device_index_; + int64_t frame_bytes_; int64_t capacity_bytes_; int64_t staged_bytes_{0}; at::ScalarType staged_scalar_type_{at::ScalarType::Undefined}; bool has_staged_input_{false}; + bool staged_fast_path_{false}; + bool staged_owner_path_{false}; + uint64_t* local_stage_sequence_{nullptr}; + uint64_t* local_done_sequence_{nullptr}; + uint64_t* local_fused_sequence_{nullptr}; PeerPointers peers_{}; + FusedPeerPointers fused_peer_slots_{}; std::array imported_bases_{}; }; -DeterministicCollectiveState* state_from_handle(int64_t handle) { +std::atomic next_collective_handle{1}; +std::mutex collective_states_mutex; +std::unordered_map> + collective_states; + +std::shared_ptr state_from_handle(int64_t handle) { TORCH_CHECK(handle != 0, "deterministic collective handle is closed"); - return reinterpret_cast(handle); + std::lock_guard lock(collective_states_mutex); + const auto it = collective_states.find(handle); + TORCH_CHECK( + it != collective_states.end(), + "deterministic collective handle ", + handle, + " is unknown or already closed"); + return it->second; } } // namespace @@ -535,16 +1405,35 @@ int64_t deterministic_collective_create( const std::vector& offsets, int64_t rank) { const c10::cuda::CUDAGuard device_guard(staging.device()); - auto state = std::make_unique( + auto state = std::make_shared( staging, handles, offsets, rank); - return reinterpret_cast(state.release()); + const int64_t handle = next_collective_handle.fetch_add( + 1, std::memory_order_relaxed); + TORCH_CHECK(handle > 0, "deterministic collective handle space exhausted"); + { + std::lock_guard lock(collective_states_mutex); + const bool inserted = collective_states.emplace(handle, state).second; + TORCH_CHECK(inserted, "duplicate deterministic collective handle ", handle); + } + return handle; } void deterministic_collective_destroy(int64_t handle) { - delete state_from_handle(handle); + std::shared_ptr state; + { + std::lock_guard lock(collective_states_mutex); + const auto it = collective_states.find(handle); + TORCH_CHECK( + it != collective_states.end(), + "deterministic collective handle ", + handle, + " is unknown or already closed"); + state = std::move(it->second); + collective_states.erase(it); + } } void deterministic_collective_stage(int64_t handle, torch::Tensor& input) { @@ -559,6 +1448,24 @@ void deterministic_collective_all_reduce(int64_t handle, torch::Tensor& output) state_from_handle(handle)->all_reduce(output, stream); } +void deterministic_collective_all_reduce_fused( + int64_t handle, + torch::Tensor& input, + torch::Tensor& output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + state_from_handle(handle)->all_reduce_fused(input, output, stream); +} + +void deterministic_collective_all_gather_fused( + int64_t handle, + torch::Tensor& input, + torch::Tensor& output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + state_from_handle(handle)->all_gather_fused(input, output, stream); +} + void deterministic_collective_reduce_scatter(int64_t handle, torch::Tensor& output) { const c10::cuda::CUDAGuard device_guard(output.device()); auto stream = c10::cuda::getCurrentCUDAStream().stream(); diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 34b080303..1b535d2a2 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -512,6 +512,14 @@ torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& rhs, } // anonymous namespace +bool det_gemm_sm90_compiled() { +#if defined(RL_KERNEL_ENABLE_SM90) + return true; +#else + return false; +#endif +} + torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b) { check_in(a, "A"); check_in(b, "B"); a = a.contiguous(); b = b.contiguous(); diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f1e66dba8..aecc30ed8 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -104,10 +104,15 @@ int64_t deterministic_collective_create( void deterministic_collective_destroy(int64_t handle); void deterministic_collective_stage(int64_t handle, torch::Tensor& input); void deterministic_collective_all_reduce(int64_t handle, torch::Tensor& output); +void deterministic_collective_all_reduce_fused( + int64_t handle, torch::Tensor& input, torch::Tensor& output); void deterministic_collective_reduce_scatter(int64_t handle, torch::Tensor& output); void deterministic_collective_all_gather(int64_t handle, torch::Tensor& output); +void deterministic_collective_all_gather_fused( + int64_t handle, torch::Tensor& input, torch::Tensor& output); // Batch-Invariant Deterministic GEMM Declarations +bool det_gemm_sm90_compiled(); torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); torch::Tensor det_gemm_fwd_rhs_transposed(torch::Tensor a, torch::Tensor bt); torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b); @@ -121,6 +126,10 @@ std::vector swiglu_backward_cuda( torch::Tensor dy, torch::Tensor gate, torch::Tensor up); +torch::Tensor swiglu_packed_forward_cuda(torch::Tensor gate_up); +std::vector swiglu_packed_backward_cuda( + torch::Tensor dy, + torch::Tensor gate_up); // RMSNorm Declarations & Wrappers @@ -274,6 +283,16 @@ std::vector swiglu_backward( return swiglu_backward_cuda(dy, gate, up); } +torch::Tensor swiglu_packed_forward(torch::Tensor gate_up) { + return swiglu_packed_forward_cuda(gate_up); +} + +std::vector swiglu_packed_backward( + torch::Tensor dy, + torch::Tensor gate_up) { + return swiglu_packed_backward_cuda(dy, gate_up); +} + // Deterministic standard-softmax attention (issue #147) std::vector deterministic_attention_forward( torch::Tensor q, @@ -427,6 +446,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_collective_all_reduce", &deterministic_collective_all_reduce, "Run the TP=8 deterministic fixed-tree all-reduce kernel"); + m.def( + "deterministic_collective_all_reduce_fused", + &deterministic_collective_all_reduce_fused, + "Run a fused small-message deterministic fixed-tree all-reduce"); m.def( "deterministic_collective_reduce_scatter", &deterministic_collective_reduce_scatter, @@ -435,6 +458,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_collective_all_gather", &deterministic_collective_all_gather, "Run the TP=8 deterministic rank-ordered all-gather kernel"); + m.def( + "deterministic_collective_all_gather_fused", + &deterministic_collective_all_gather_fused, + "Run a fused small-message deterministic rank-ordered all-gather"); // Prefix-shared attention uses NVIDIA PTX and falls back to PyTorch SDPA on ROCm. #if !defined(USE_ROCM) @@ -442,6 +469,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #endif // registry Batch-Invariant Deterministic GEMM + m.def("det_gemm_sm90_compiled", &det_gemm_sm90_compiled, + "Whether the extension contains the SM90 deterministic GEMM implementation"); m.def("det_gemm_fwd", &det_gemm_fwd, "Batch-invariant deterministic GEMM forward (C=A@B)"); m.def( "det_gemm_fwd_rhs_transposed", @@ -469,6 +498,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("silu_backward", &silu_backward, "Batch-invariant SiLU backward CUDA"); m.def("swiglu_forward", &swiglu_forward, "Batch-invariant SwiGLU forward CUDA"); m.def("swiglu_backward", &swiglu_backward, "Batch-invariant SwiGLU backward CUDA"); + m.def("swiglu_packed_forward", &swiglu_packed_forward, + "Batch-invariant SwiGLU forward for [rows, 2 * intermediate]"); + m.def("swiglu_packed_backward", &swiglu_packed_backward, + "Batch-invariant SwiGLU backward for [rows, 2 * intermediate]"); // Deterministic standard-softmax attention (issue #147) m.def( diff --git a/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh b/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh new file mode 100755 index 000000000..4b3a601cc --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +set -euo pipefail + +REAL_PYTHON="${RL_KERNEL_REAL_PYTHON:?RL_KERNEL_REAL_PYTHON must name the real Python executable}" +RL_KERNEL_ROOT="${RL_KERNEL_ROOT:-$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)}" +export RL_KERNEL_ROOT +export PYTHONPATH="${RL_KERNEL_ROOT}:${PYTHONPATH:-}" + +if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then + # Ray job submission does not always forward the shell exports used by the + # outer launcher. A strict linear-logp job must still install the matching + # 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 + previous_arg="" + for current_arg in "$@"; do + if [[ "${previous_arg}" == "--linear-logp-provider-mode" && "${current_arg}" == "strict" ]]; then + strict_linear_logp=1 + break + fi + previous_arg="${current_arg}" + done + 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}" + fi + exec "${REAL_PYTHON}" "$@" \ + --seed 1234 \ + --rollout-seed 42 \ + --vllm-enable-deterministic-inference \ + --vllm-attention-backend flash_attn \ + --vllm-disable-custom-all-reduce \ + --deterministic-mode \ + --accumulate-allreduce-grads-in-fp32 +fi + +exec "${REAL_PYTHON}" "$@" diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8727a7c97..8e0e865af 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -14,8 +14,14 @@ def deterministic_collective_create( def deterministic_collective_destroy(handle: int) -> None: ... def deterministic_collective_stage(handle: int, input: torch.Tensor) -> None: ... def deterministic_collective_all_reduce(handle: int, output: torch.Tensor) -> None: ... +def deterministic_collective_all_reduce_fused( + handle: int, input: torch.Tensor, output: torch.Tensor +) -> None: ... def deterministic_collective_reduce_scatter(handle: int, output: torch.Tensor) -> None: ... def deterministic_collective_all_gather(handle: int, output: torch.Tensor) -> None: ... +def deterministic_collective_all_gather_fused( + handle: int, input: torch.Tensor, output: torch.Tensor +) -> None: ... def fused_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: ... def fused_logp_sm90(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: ... def batch_invariant_logp_sm90( @@ -182,6 +188,7 @@ def deterministic_attention_backward( scale: float, key_padding_mask: torch.Tensor | None, ) -> list[torch.Tensor]: ... +def det_gemm_sm90_compiled() -> bool: ... def det_gemm_fwd(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ... def det_gemm_fwd_rhs_transposed( a: torch.Tensor, diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index 54b00fa5e..0fed9ecaa 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -13,7 +13,39 @@ _SUPPORTED_WORLD_SIZES = (1, 2, 4, 8) _DEFAULT_MAX_SIZE_BYTES = 64 * 1024 * 1024 +_COLLECTIVE_STAGING_FRAMES = 3 +_COLLECTIVE_FRAME_METADATA_BYTES = 3 * 8 _REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) +_COLLECTIVES: dict[tuple[int, int, int, int], DeterministicCollective] = {} +DETERMINISTIC_ALL_REDUCE_OP = "rl_kernel::deterministic_all_reduce_" + + +@torch.library.custom_op(DETERMINISTIC_ALL_REDUCE_OP, mutates_args={"input"}) +def _deterministic_all_reduce_(input: torch.Tensor, collective_handle: int) -> None: + """Expose the stateful IPC reduction as an explicit graph mutation.""" + + from rl_engine import _C + + _C.deterministic_collective_all_reduce_fused(collective_handle, input, input) + + +@_deterministic_all_reduce_.register_fake +def _deterministic_all_reduce_fake( + input: torch.Tensor, + collective_handle: int, +) -> None: + del input, collective_handle + + +def deterministic_all_reduce_inplace( + input: torch.Tensor, + *, + collective_handle: int, +) -> torch.Tensor: + """Run the graph-visible deterministic all-reduce in place.""" + + _deterministic_all_reduce_(input, collective_handle) + return input class DeterministicCollective: @@ -26,9 +58,9 @@ class DeterministicCollective: node evaluates the lower logical subtree before the higher one. One instance owns a symmetric CUDA IPC staging buffer. All ranks must call - its methods in the same order with matching shapes and dtypes. Calls are - host-synchronizing by design; the first version prioritizes determinism and - lifetime safety over overlap or throughput. + its methods in the same order with matching shapes and dtypes. Device-side + IPC sequence fences order staging and payload access without a steady-state + host barrier and advance correctly during CUDA Graph replay. """ def __init__( @@ -39,7 +71,9 @@ def __init__( max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, ) -> None: if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("torch.distributed must be initialized before collectives") + raise RuntimeError( + "torch.distributed must be initialized before collectives" + ) if not torch.cuda.is_available(): raise RuntimeError("deterministic collectives require CUDA") if max_size_bytes <= 0: @@ -61,7 +95,9 @@ def __init__( else: normalized_device = torch.device(device) if normalized_device.type != "cuda": - raise ValueError(f"deterministic collectives require a CUDA device, got {device!r}") + raise ValueError( + f"deterministic collectives require a CUDA device, got {device!r}" + ) if normalized_device.index is None: normalized_device = torch.device("cuda", torch.cuda.current_device()) if normalized_device.index != torch.cuda.current_device(): @@ -98,13 +134,17 @@ def __init__( self._extension = _C self._lock = threading.Lock() self._handle = 0 - self._staging = torch.empty( - self.max_size_bytes, + self._validated_signatures: set[tuple[Any, ...]] = set() + self._staging = torch.zeros( + _COLLECTIVE_STAGING_FRAMES + * (self.max_size_bytes + _COLLECTIVE_FRAME_METADATA_BYTES), dtype=torch.uint8, device=self.device, ) - handle, offset = self._extension.deterministic_collective_ipc_meta(self._staging) + handle, offset = self._extension.deterministic_collective_ipc_meta( + self._staging + ) local_meta = { "handle": handle, "offset": int(offset), @@ -138,6 +178,7 @@ def all_reduce( input: torch.Tensor, *, out: torch.Tensor | None = None, + validate_signature: bool = True, ) -> torch.Tensor: """Return the TBIK-compatible fixed-tree sum on every rank. @@ -153,11 +194,11 @@ def all_reduce( self._validate_output(out, input) with self._lock: - self._validate_matching_signature("all_reduce", input) - self._extension.deterministic_collective_stage(self._handle, input) - self._synchronize_ranks() - self._extension.deterministic_collective_all_reduce(self._handle, out) - self._synchronize_ranks() + if validate_signature: + self._validate_matching_signature("all_reduce", input) + self._extension.deterministic_collective_all_reduce_fused( + self._handle, input, out + ) return out def all_gather( @@ -165,6 +206,7 @@ def all_gather( input: torch.Tensor, *, out: torch.Tensor | None = None, + validate_signature: bool = True, ) -> torch.Tensor: """Gather rank-ordered input bit patterns along dimension 0. @@ -180,18 +222,34 @@ def all_gather( self._validate_sharded_output(out, input, output_shape) with self._lock: - self._validate_matching_signature("all_gather", input) - self._extension.deterministic_collective_stage(self._handle, input) - self._synchronize_ranks() - self._extension.deterministic_collective_all_gather(self._handle, out) - self._synchronize_ranks() + if validate_signature: + self._validate_matching_signature("all_gather", input) + self._extension.deterministic_collective_all_gather_fused( + self._handle, input, out + ) return out + def all_gather_many( + self, + inputs: tuple[torch.Tensor, ...], + *, + validate_signature: bool = True, + ) -> tuple[torch.Tensor, ...]: + """Gather several tensors through the available single-tensor ABI.""" + + if not inputs: + raise ValueError("all_gather_many requires at least one input") + return tuple( + self.all_gather(input, validate_signature=validate_signature) + for input in inputs + ) + def reduce_scatter( self, input: torch.Tensor, *, out: torch.Tensor | None = None, + validate_signature: bool = True, ) -> torch.Tensor: """TBIK-compatible fixed-tree sum, then a rank-ordered dimension-0 scatter. @@ -214,13 +272,27 @@ def reduce_scatter( self._validate_sharded_output(out, input, output_shape) with self._lock: - self._validate_matching_signature("reduce_scatter", input) + if validate_signature: + self._validate_matching_signature("reduce_scatter", input) self._extension.deterministic_collective_stage(self._handle, input) - self._synchronize_ranks() self._extension.deterministic_collective_reduce_scatter(self._handle, out) - self._synchronize_ranks() return out + def reduce_scatter_many( + self, + inputs: tuple[torch.Tensor, ...], + *, + validate_signature: bool = True, + ) -> tuple[torch.Tensor, ...]: + """Reduce-scatter several tensors through the single-tensor ABI.""" + + if not inputs: + raise ValueError("reduce_scatter_many requires at least one input") + return tuple( + self.reduce_scatter(input, validate_signature=validate_signature) + for input in inputs + ) + def close(self) -> None: """Release imported CUDA IPC mappings after the last collective call.""" @@ -303,18 +375,55 @@ def _validate_sharded_output( if output.dtype != input.dtype: raise TypeError("out must have the same dtype as input") if output.shape != output_shape: - raise ValueError(f"out must have shape {output_shape}, got {tuple(output.shape)}") + raise ValueError( + f"out must have shape {output_shape}, got {tuple(output.shape)}" + ) if not output.is_contiguous(): raise ValueError("out must be contiguous") def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> None: signature = (op_name, tuple(input.shape), str(input.dtype), input.numel()) + if signature in self._validated_signatures: + return signatures: list[tuple[Any, ...] | None] = [None] * self.world_size dist.all_gather_object(signatures, signature, group=self.group) if any(peer_signature != signature for peer_signature in signatures): raise ValueError( f"all ranks must call {op_name} with matching shapes and dtypes; got {signatures}" ) + self._validated_signatures.add(signature) + + def _validate_matching_many_signature( + self, + op_name: str, + inputs: tuple[torch.Tensor, ...], + ) -> None: + signature = ( + op_name, + tuple( + (tuple(input.shape), str(input.dtype), input.numel()) + for input in inputs + ), + ) + if signature in self._validated_signatures: + return + signatures: list[tuple[Any, ...] | None] = [None] * self.world_size + dist.all_gather_object(signatures, signature, group=self.group) + if any(peer_signature != signature for peer_signature in signatures): + raise ValueError( + f"all ranks must call {op_name} with matching tensors; got {signatures}" + ) + self._validated_signatures.add(signature) + + def _validate_many_capacity(self, inputs: tuple[torch.Tensor, ...]) -> None: + total_bytes = 0 + for input in inputs: + total_bytes = (total_bytes + 15) & ~15 + total_bytes += input.numel() * input.element_size() + if total_bytes > self.max_size_bytes: + raise ValueError( + f"inputs require {total_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) def _synchronize_ranks(self) -> None: torch.cuda.synchronize(self.device) @@ -323,3 +432,47 @@ def _synchronize_ranks(self) -> None: dist.barrier(group=self.group, device_ids=[self.device.index]) else: dist.barrier(group=self.group) + + +def collective_for_group( + group: dist.ProcessGroup | None, + *, + min_size_bytes: int = 0, + minimum_capacity_bytes: int = _DEFAULT_MAX_SIZE_BYTES, + device: torch.device | str | int | None = None, +) -> DeterministicCollective | None: + """Return the process-local RL-Kernel collective shared by hot-path ops.""" + + if group is None: + return None + if min_size_bytes < 0: + raise ValueError("min_size_bytes must be non-negative") + if minimum_capacity_bytes <= 0: + raise ValueError("minimum_capacity_bytes must be positive") + + rank = dist.get_rank(group=group) + world_size = dist.get_world_size(group=group) + if device is None: + device_index = torch.cuda.current_device() + else: + normalized_device = torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + device_index = ( + torch.cuda.current_device() + if normalized_device.index is None + else normalized_device.index + ) + key = (id(group), rank, world_size, device_index) + cached = _COLLECTIVES.get(key) + if cached is not None and cached.max_size_bytes >= min_size_bytes: + return cached + # Borrowers such as Attention autograd contexts can outlive this cache + # entry. Replacing an undersized entry must not invalidate those live + # references; normal Python ownership closes it after the last borrower. + + collective = DeterministicCollective( + group=group, + device=device_index, + max_size_bytes=max(minimum_capacity_bytes, min_size_bytes), + ) + _COLLECTIVES[key] = collective + return collective diff --git a/rl_engine/integrations/framework_operators.py b/rl_engine/integrations/framework_operators.py index d5c7c7217..b97daf6a7 100644 --- a/rl_engine/integrations/framework_operators.py +++ b/rl_engine/integrations/framework_operators.py @@ -29,12 +29,10 @@ AttentionDType, AttentionMode, AttentionRole, + SplitKVSpec, ) 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, @@ -43,13 +41,16 @@ ) 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.cuda.matmul.det_gemm import det_gemm_backend_id 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 ( + 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 +from rl_engine.runtime_mode import strict_contract_enabled ATTENTION_BACKEND_ID = "rlkernel.attention.deterministic.v1" FFN_BACKEND_ID = "rlkernel.ffn.qwen3.deterministic.v1" @@ -75,38 +76,15 @@ def _optional_env_int(name: str) -> int | None: 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 _tensor_metadata(tensor: torch.Tensor) -> dict[str, Any]: + """Describe a hot-path tensor without reading CUDA values on the host.""" - -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 + return { + "shape": list(tensor.shape), + "dtype": _dtype_name(tensor), + "device": str(tensor.device), + "numel": int(tensor.numel()), + } def _attention_dtype(tensor: torch.Tensor) -> AttentionDType: @@ -136,8 +114,9 @@ def __init__(self, *, target: str, semantic_op: str, backend_id: str) -> None: 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._runtime_device: torch.device | None = None + self._runtime_dtype: torch.dtype | None = None self._lock = Lock() def get( @@ -147,12 +126,6 @@ def get( 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 @@ -165,14 +138,20 @@ def get( # 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 + if ( + tensor.device != self._runtime_device + or tensor.dtype != self._runtime_dtype ): raise RuntimeError( f"{self.semantic_op} runtime device/dtype changed after resolution" ) return self._instance + requirements = OperatorRequirements( + device=_device_name(tensor), + dtype=_dtype_name(tensor), + topology=topology, + alignment_properties={"deterministic": True}, + ) target = cast(Any, self.target) resolved = self._bridge.resolve_override( OperatorOverride.for_target( @@ -201,8 +180,9 @@ def get( f"{actual_backend!r}" ) self._instance = instance - self._requirements = requirements self._provenance = provenance.to_dict() + self._runtime_device = tensor.device + self._runtime_dtype = tensor.dtype return instance @property @@ -219,9 +199,40 @@ def _weight(module: Any, name: str) -> torch.Tensor: return value -def _split_gate_up(weight: torch.Tensor, name: str) -> tuple[torch.Tensor, torch.Tensor]: +def _fused_rms_norm_input( + projection: Any, + hidden_states: torch.Tensor, + name: str, +) -> torch.Tensor: + """Recover the RMSNorm hidden by TE's LayerNormLinear wrapper.""" + + weight = getattr(projection, "layer_norm_weight", None) + if weight is None: + return hidden_states + if not isinstance(weight, torch.Tensor): + raise RuntimeError(f"{name}.layer_norm_weight must be a tensor") + if getattr(projection, "normalization", None) != "RMSNorm": + raise RuntimeError(f"strict {name} requires fused RMSNorm") + if getattr(projection, "layer_norm_bias", None) is not None: + raise RuntimeError(f"strict {name} RMSNorm must be bias-free") + if bool(getattr(projection, "zero_centered_gamma", False)): + raise RuntimeError(f"strict {name} does not support zero-centered gamma") + eps = float(getattr(projection, "eps")) + return torch.nn.functional.rms_norm( + hidden_states, + (hidden_states.shape[-1],), + weight, + eps, + ) + + +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") + 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() @@ -264,37 +275,63 @@ def _packed_local_sequence_layout( """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'") + 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") + 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() + 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() + 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") + 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) + 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") + 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") + raise RuntimeError( + "packed Attention cu_seqlens do not cover the local Q/KV token rows" + ) return tuple(local_offsets), global_lengths +def _tensor_cache_token(tensor: torch.Tensor) -> tuple[Any, ...]: + """Identify one tensor value while it is reused across framework layers.""" + + return ( + tensor.data_ptr(), + tuple(tensor.shape), + tensor.dtype, + tensor.device.type, + tensor.device.index, + int(tensor._version), + ) + + def _compact_attention_provenance(value: Mapping[str, Any]) -> dict[str, Any]: """Keep strict backend identity without retaining one record per token.""" @@ -373,10 +410,57 @@ def __init__(self, handle: SemanticOperatorHandle | None = None) -> None: target="training", semantic_op="attention", backend_id=self.backend_id ) self._last_provenance: dict[str, Any] = {} + self._packed_layout_owner: Any | None = None + self._packed_layout_key: tuple[Any, ...] | None = None + self._packed_layout_value: tuple[tuple[int, ...], tuple[int, ...]] | None = None + + def _packed_layout( + self, + packed_seq_params: Any, + *, + cp_world_size: int, + local_query_tokens: int, + local_kv_tokens: int, + ) -> tuple[tuple[int, ...], tuple[int, ...]]: + 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" + ) + key = ( + _tensor_cache_token(query_cu), + _tensor_cache_token(kv_cu), + cp_world_size, + local_query_tokens, + local_kv_tokens, + ) + if ( + self._packed_layout_owner is packed_seq_params + and self._packed_layout_key == key + ): + if self._packed_layout_value is None: + raise RuntimeError("packed Attention layout cache is empty") + return self._packed_layout_value + value = _packed_local_sequence_layout( + packed_seq_params, + cp_world_size=cp_world_size, + local_query_tokens=local_query_tokens, + local_kv_tokens=local_kv_tokens, + ) + self._packed_layout_owner = packed_seq_params + self._packed_layout_key = key + self._packed_layout_value = value + return value @property def provenance(self) -> Mapping[str, Any]: return { + "interface": "megatron.attention.forward", + "operator": self.backend_id, + "fallback": False, "semantic_instance": self._handle.provenance, "execution": dict(self._last_provenance), } @@ -399,9 +483,15 @@ def __call__( 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") + 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: + 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") @@ -430,10 +520,12 @@ def execute_sequence( *, 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, + 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, @@ -486,14 +578,13 @@ def execute_sequence( "operator": _compact_attention_provenance(result.provenance), } else: - local_offsets, global_lengths = _packed_local_sequence_layout( + local_offsets, global_lengths = self._packed_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]] = [] + grouped_sequences: dict[tuple[int, int], list[tuple[int, int, int]]] = {} for sequence_index, (start, end, global_length) in enumerate( zip( local_offsets[:-1], @@ -502,30 +593,68 @@ def execute_sequence( 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() + grouped_sequences.setdefault((end - start, global_length), []).append( + (sequence_index, start, end) + ) + + outputs: list[torch.Tensor | None] = [None] * len(global_lengths) + sequence_provenance: list[dict[str, Any] | None] = [None] * len( + global_lengths + ) + launch_group_count = 0 + for (local_length, global_length), sequences in grouped_sequences.items(): + q_ready = ( + torch.stack( + [query[start:end] for _index, start, end in sequences], dim=0 + ) + .permute(0, 2, 1, 3) + .contiguous() + ) + k_ready = ( + torch.stack( + [key[start:end] for _index, start, end in sequences], dim=0 + ) + .permute(0, 2, 1, 3) + .contiguous() + ) + v_ready = ( + torch.stack( + [value[start:end] for _index, start, end in sequences], dim=0 + ) + .permute(0, 2, 1, 3) + .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) + group_output = ( + result.out.permute(0, 2, 1, 3).contiguous().flatten(start_dim=2) ) - sequence_provenance.append( - { + operator_provenance = _compact_attention_provenance(result.provenance) + for batch_index, (sequence_index, _start, _end) in enumerate(sequences): + outputs[sequence_index] = group_output[batch_index] + sequence_provenance[sequence_index] = { "sequence_index": sequence_index, - "local_tokens": end - start, + "local_tokens": local_length, "global_tokens": global_length, - "operator": _compact_attention_provenance(result.provenance), + "operator": operator_provenance, } + launch_group_count += 1 + if any(item is None for item in outputs) or any( + item is None for item in sequence_provenance + ): + raise RuntimeError( + "packed Attention failed to materialize every sequence" ) - output = torch.cat(outputs, dim=0) + output = torch.cat(cast(list[torch.Tensor], outputs), dim=0) execution_provenance = { "packed_sequence_count": len(global_lengths), - "sequences": sequence_provenance, + "launch_group_count": launch_group_count, + "sequence_batching": "equal_length_rows", + "sequences": cast(list[dict[str, Any]], sequence_provenance), } self._last_provenance = { "framework_layout": ( @@ -555,6 +684,9 @@ def __init__(self, handle: SemanticOperatorHandle | None = None) -> None: @property def provenance(self) -> Mapping[str, Any]: return { + "interface": "megatron.mlp.forward", + "operator": self.backend_id, + "fallback": False, "semantic_instance": self._handle.provenance, "execution": dict(self._last_provenance), } @@ -576,16 +708,26 @@ def __call__( if name != "padding_mask" or value is not None } if unexpected: - raise RuntimeError("strict dense Qwen3 FFN does not accept expert token scaling") + 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") + 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") + hidden_states = _fused_rms_norm_input( + module.linear_fc1, + hidden_states, + "linear_fc1", + ) + fused_gate_up = _weight(module.linear_fc1, "linear_fc1").contiguous() + gate, up = _split_gate_up(fused_gate_up, "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()) @@ -615,6 +757,9 @@ def __call__( "tp_world_size": tp_world, "runtime_platform": "cuda", "actual_backend": "rlkernel.cuda.det_gemm_swiglu", + "gemm_backend": det_gemm_backend_id(), + "fallback": False, + "gate_up_projection": "separate_strict_launches", "triton_used": False, } return output, None @@ -646,7 +791,8 @@ def _vllm_kv_cache_views( 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]" + "vLLM FlashAttention KV cache must use " + "[blocks, kv_heads, block, 2 * head_size]" ) @@ -660,32 +806,36 @@ def __init__(self, handle: SemanticOperatorHandle | None = None) -> None: target="rollout", semantic_op="attention", backend_id=self.backend_id ) self._last_provenance: dict[str, Any] = {} - self._prime_semantic_handle() + self._tp_coordinates: tuple[int, int, Any] | None = None + self._metadata_cache_key: tuple[Any, ...] | None = None + self._metadata_cache_owners: set[int] = set() + self._metadata_cache_value: ( + tuple[list[dict[str, Any]], dict[str, Any]] | None + ) = None - 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. + def bind_inference(self) -> None: + """Resolve the backend after vLLM has selected the worker CUDA device.""" + + if self._tp_coordinates is not None: return + tp_world, tp_rank, tp_group = _vllm_tp_coordinates() + device = torch.device("cuda", torch.cuda.current_device()) + self._handle.get( + torch.empty((1,), device=device, dtype=torch.bfloat16), + topology={ + "world_size": tp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": 1, + }, + ) + self._tp_coordinates = (tp_world, tp_rank, tp_group) @property def provenance(self) -> Mapping[str, Any]: return { + "interface": "vllm.attention.forward", + "operator": self.backend_id, + "fallback": False, "semantic_instance": self._handle.provenance, "execution": dict(self._last_provenance), } @@ -698,6 +848,114 @@ def _metadata_tensor(metadata: Any, *names: str) -> torch.Tensor: return value raise RuntimeError(f"vLLM Attention metadata is missing {'/'.join(names)}") + def _materialization_groups( + self, + attn_metadata: Any, + *, + query: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + num_actual: int, + cache_owner: Any | None = None, + ) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Build one paged launch from vLLM's graph-replayable GPU metadata.""" + + query_starts_source = self._metadata_tensor( + attn_metadata, "query_start_loc", "query_start_loc_cpu" + ) + seq_lens_source = self._metadata_tensor( + attn_metadata, "seq_lens", "seq_lens_cpu" + ) + if num_actual == 0: + return [], { + "row_count": 0, + "query_position_range": [None, None], + "kv_token_range": [None, None], + "launch_group_count": 0, + "metadata_source": "vllm_gpu", + } + + query_starts = query_starts_source.to( + device=query.device, dtype=torch.int32 + ) + seq_lens = seq_lens_source.to(device=query.device, dtype=torch.int32) + query_indices = torch.arange( + num_actual, dtype=torch.int32, device=query.device + ) + query_ends = query_starts[1:] + request_indices = torch.searchsorted( + query_ends, query_indices, right=True + ).to(dtype=torch.long) + active_queries = query_indices < query_starts[-1] + request_indices = request_indices.clamp_max(seq_lens.numel() - 1) + request_query_ends = query_ends.index_select(0, request_indices) + request_seq_lens = seq_lens.index_select(0, request_indices) + seqused_k = request_seq_lens - (request_query_ends - query_indices) + 1 + seqused_k = torch.where( + active_queries, + seqused_k, + torch.ones_like(seqused_k), + ) + + max_seq_len = int( + getattr(attn_metadata, "max_seq_len", block_table.size(1) * block_size) + ) + page_count = min( + block_table.size(1), (max_seq_len + block_size - 1) // block_size + ) + cache_key = ( + _tensor_cache_token(query_starts_source), + _tensor_cache_token(seq_lens_source), + _tensor_cache_token(block_table), + query.device.type, + query.device.index, + num_actual, + block_size, + page_count, + ) + owner_id = id(cache_owner) if cache_owner is not None else None + if ( + owner_id is not None + and self._metadata_cache_key == cache_key + and owner_id not in self._metadata_cache_owners + and self._metadata_cache_value is not None + ): + self._metadata_cache_owners.add(owner_id) + groups, cached_summary = self._metadata_cache_value + return groups, {**cached_summary, "metadata_reused_across_layers": True} + + pages = ( + block_table.index_select(0, request_indices)[:, :page_count] + .to(dtype=torch.int32) + .contiguous() + ) + groups = [ + { + "page_count": page_count, + "pages": pages, + "query_indices": query_indices, + "query_start": 0, + "query_count": num_actual, + "query_contiguous": True, + "seqused_k": seqused_k, + } + ] + summary = { + "row_count": num_actual, + "query_position_range": "device_dynamic", + "kv_token_range": "device_dynamic", + "launch_group_count": 1, + "metadata_source": "vllm_gpu", + "metadata_reused_across_layers": False, + } + if owner_id is not None: + # The first repeated layer marks the next model forward (or graph + # capture), so dynamic metadata is re-derived instead of frozen. + self._metadata_cache_key = cache_key + self._metadata_cache_owners = {owner_id} + self._metadata_cache_value = (groups, summary) + return groups, summary + def __call__( self, impl: Any, @@ -711,9 +969,11 @@ def __call__( output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: - del layer, key, value + del 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") + 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") @@ -728,26 +988,25 @@ def __call__( 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 + block_table = self._metadata_tensor( + attn_metadata, "block_table", "block_table_tensor" ) num_actual = int(getattr(attn_metadata, "num_actual_tokens", query.size(0))) + if num_actual < 0 or num_actual > query.size(0): + raise RuntimeError("vLLM num_actual_tokens is outside the query buffer") if output is None: output = torch.empty( (query.size(0), impl.num_heads * impl.head_size), dtype=query.dtype, device=query.device, ) - output.zero_() + if output.size(0) < num_actual: + raise RuntimeError("vLLM output buffer is smaller than num_actual_tokens") 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() + if self._tp_coordinates is None: + self._tp_coordinates = _vllm_tp_coordinates() + tp_world, tp_rank, tp_group = self._tp_coordinates operator = self._handle.get( query, topology={ @@ -756,97 +1015,62 @@ def __call__( "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 + runtime = operator.bind_cuda_runtime() + groups, metadata_summary = self._materialization_groups( + attn_metadata, + query=query, + block_table=block_table, + block_size=block_size, + num_actual=num_actual, + cache_owner=layer, + ) 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) + next_query_row = 0 + for group in groups: + if not group["query_contiguous"] or group["query_start"] != next_query_row: + break + next_query_row += int(group["query_count"]) + direct_output_buffer = bool(groups) and next_query_row == num_actual + if not direct_output_buffer: + output.zero_() + elif num_actual < output.size(0): + output[num_actual:].zero_() + for group in groups: + page_count = int(group["page_count"]) + pages = group["pages"] + query_indices = group["query_indices"] + if group["query_contiguous"]: + query_start = int(group["query_start"]) + query_count = int(group["query_count"]) + q_ready = query.narrow(0, query_start, query_count).unsqueeze(2) + else: + q_ready = query.index_select(0, query_indices).unsqueeze(2).contiguous() + result = runtime.forward_paged_with_lse( + q_ready, + key_cache, + value_cache, + page_table=pages, + seqused_k=group["seqused_k"], + max_seqlen_k=page_count * block_size, + scale=float(impl.scale), + ) + result_output = result.out.squeeze(2) + if group["query_contiguous"]: + output_heads.narrow( + 0, int(group["query_start"]), int(group["query_count"]) + ).copy_(result_output) + else: + output_heads.index_copy_(0, query_indices, result.out.squeeze(2)) + last_operator_provenance = _compact_attention_provenance(result.provenance) self._last_provenance = { "framework_layout": "vllm_paged_kv", - "materialization": "one_causal_prefix_per_query_row", + "materialization": "direct_paged_fa4", "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], + "direct_output_buffer": direct_output_buffer, + **metadata_summary, "operator": last_operator_provenance, } return output @@ -860,65 +1084,116 @@ def __init__(self, handle: SemanticOperatorHandle | None = None) -> None: 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 + self._packed_inference_binding: tuple[int, int, str] | None = None + self._tp_coordinates: tuple[int, int, Any] | None = None @property def provenance(self) -> Mapping[str, Any]: return { + "interface": "vllm.qwen3.mlp.forward", + "operator": self.backend_id, + "fallback": False, "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() + def bind_packed_inference(self, module: Any) -> tuple[int, int]: + """Bind vLLM's TP group before torch.compile captures the model.""" + + fused_gate_up = _weight(module.gate_up_proj, "gate_up_proj") + down = _weight(module.down_proj, "down_proj") + tp_world, tp_rank, tp_group = _vllm_tp_coordinates() + self._tp_coordinates = (tp_world, tp_rank, tp_group) operator = self._handle.get( - hidden_states, + fused_gate_up, topology={ "world_size": tp_world, "tensor_parallel_size": tp_world, "context_parallel_size": 1, }, ) - output = operator( - hidden_states, - gate, - up, + prepare = getattr(operator, "prepare_packed_inference", None) + if not callable(prepare): + raise RuntimeError("strict FFN backend lacks packed inference preparation") + collective_handle, bound_tp_world = prepare( + fused_gate_up, down, tp_group=tp_group, - sequence_parallel=False, - deterministic=True, ) + backend_id = "deterministic_all_reduce.ipc_localized_fixed_tree.v1" + resolve_backend = getattr(operator, "packed_inference_backend_id", None) + if callable(resolve_backend) and collective_handle: + backend_id = str(resolve_backend(collective_handle)) + self._packed_inference_binding = ( + collective_handle, + bound_tp_world, + backend_id, + ) + self._set_runtime_provenance(tp_world, backend_id) + return collective_handle, bound_tp_world + + def _set_runtime_provenance( + self, tp_world: int, deterministic_all_reduce_backend: str = "unbound" + ) -> None: self._last_provenance = { "framework_layout": "vllm_tensor_parallel", "tp_world_size": tp_world, "runtime_platform": "cuda", "actual_backend": "rlkernel.cuda.det_gemm_swiglu", + "gemm_backend": det_gemm_backend_id(), + "fallback": False, + "gate_up_projection": "packed_single_launch", + "deterministic_all_reduce_backend": deterministic_all_reduce_backend, "triton_used": False, } + + def __call__(self, module: Any, hidden_states: torch.Tensor) -> torch.Tensor: + _require_nvidia_cuda(hidden_states, "FFN") + fused_gate_up = _weight(module.gate_up_proj, "gate_up_proj").contiguous() + down = _weight(module.down_proj, "down_proj").contiguous() + if self._tp_coordinates is None: + self._tp_coordinates = _vllm_tp_coordinates() + tp_world, _tp_rank, tp_group = self._tp_coordinates + operator = self._handle.get( + hidden_states, + topology={ + "world_size": tp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": 1, + }, + ) + bound_backend = "unbound" + if self._packed_inference_binding is not None: + bound_backend = self._packed_inference_binding[2] + self._set_runtime_provenance(tp_world, bound_backend) + if self._packed_inference_binding is not None: + collective_handle, bound_tp_world, _ = self._packed_inference_binding + if bound_tp_world != tp_world: + raise RuntimeError( + "packed rollout FFN TP topology changed after binding" + ) + output = operator.packed_inference( + hidden_states, + fused_gate_up, + down, + collective_handle=collective_handle, + tp_world_size=bound_tp_world, + ) + else: + if torch._dynamo.is_compiling(): + raise RuntimeError( + "packed rollout FFN was not bound before torch.compile capture" + ) + gate, up = _split_gate_up(fused_gate_up, "gate_up_proj") + output = operator( + hidden_states, + gate, + up, + down, + tp_group=tp_group, + sequence_parallel=False, + deterministic=True, + ) return output @@ -948,19 +1223,20 @@ def provenance(self) -> Mapping[str, Any]: def __call__(self, request: Any) -> Any: logits = getattr(request, "logits", None) context = getattr(request, "context", None) - hidden = getattr(context, "hidden", None) - strict = os.getenv("VIME_RL_KERNEL_STRICT", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } + hidden = getattr(request, "hidden", None) + if not isinstance(hidden, torch.Tensor): + hidden = getattr(context, "hidden", None) + strict = strict_contract_enabled() 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 = { + "interface": "vime.selected_logprob_provider", + "operator": self.backend_id, + "actual_backend": self.backend_id, + "fallback": False, "runtime_platform": "cuda", "triton_used": False, "provider": dict(getattr(result, "provenance", {})), @@ -977,6 +1253,10 @@ def __call__(self, request: Any) -> Any: _require_nvidia_cuda(logits, "Logp") result = self._provider(request) self._last_provenance = { + "interface": "vime.selected_logprob_provider", + "operator": self.backend_id, + "actual_backend": self.backend_id, + "fallback": False, "runtime_platform": "cuda", "triton_used": False, "provider": dict(getattr(result, "provenance", {})), @@ -1095,24 +1375,31 @@ def _replace_sampled_value( "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() + every_sample_present = matches.any(dim=1).all() + if every_sample_present.is_cuda: + torch._assert_async( + every_sample_present, + "vLLM logprob result does not contain every sampled token", + ) + elif not bool(every_sample_present): + raise RuntimeError( + "vLLM logprob result does not contain every sampled token" + ) # vLLM prepends the sampled token and then appends top-k tokens. When # the sample is also in top-k, its API conversion keeps the last # duplicate, so every matching column must carry the strict value. values = torch.where(matches, selected.unsqueeze(1), values) self._last_provenance = { **dict(provenance), - "sampled_token_ids": _tensor_debug_stats(token_ids), + "interface": "vllm.sampler.selected_logprob", + "operator": self.backend_id, + "actual_backend": self.backend_id, + "fallback": False, + "sampled_token_ids": _tensor_metadata(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), + "strict_selected_logp": _tensor_metadata(selected), + "native_reference_compared": False, } if hasattr(logprobs_tensors, "_replace"): updated_tensors = logprobs_tensors._replace(logprobs=values) @@ -1128,8 +1415,59 @@ def __call__( predict_bonus_token: bool = False, logprobs_mode_override: Any = None, ) -> Any: - source_logits = logits.clone() + # Native sampling owns logits; strict selected-logp preserves its raw + # local shard before the sampler applies in-place transformations. + source_logits = logits _require_nvidia_cuda(source_logits, "Logp") + context = None + local_logits = None + 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 source_logits.size(1) < context.real_vocab_size: + raise RuntimeError( + "strict rollout source logits do not cover the real vocabulary: " + f"{source_logits.size(1)} < {context.real_vocab_size}" + ) + local_vocab = int(context.lm_head_weight.size(0)) + available = max( + 0, + min( + local_vocab, + source_logits.size(1) - context.vocab_start_index, + ), + ) + # Preserve raw model logits before vLLM's sampler transforms its + # input in place (temperature, penalties, and masking). + if available == local_vocab: + # Rank 0 normally has a complete local shard. Narrowing first + # avoids a fill kernel followed by a second device copy. + local_logits = source_logits.narrow( + 1, + context.vocab_start_index, + local_vocab, + ).contiguous() + else: + # The final TP shard may include padded rows absent from the + # serving logits; retain the exact -inf padding contract. + local_logits = source_logits.new_full( + (source_logits.size(0), local_vocab), + float("-inf"), + ) + if available: + local_logits[:, :available].copy_( + source_logits.narrow( + 1, + context.vocab_start_index, + available, + ) + ) if self._worker_sampler: result = self._native_forward(sampler, logits, sampling_metadata) else: @@ -1149,25 +1487,16 @@ def __call__( 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)}" - ) + assert context is not None and local_logits is not None 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, + selected = self._linear_logp.from_local_logits( + local_logits, 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, @@ -1175,16 +1504,28 @@ def __call__( temperature=float(os.getenv("RL_KERNEL_VLLM_TEMPERATURE", "1.0")), target="rollout", ) + strict_provenance = self._linear_logp.provenance + if ( + strict_provenance.get("deterministic_linear_logp") is not True + or strict_provenance.get("actual_backend") != self._linear_logp.backend_id + or strict_provenance.get("strict_entrypoint") + != "sm90_deterministic_logp_from_local_logits_tp" + ): + raise RuntimeError( + "strict vLLM rollout linear_logp did not execute the deterministic " + "linear-logp entry point" + ) provenance = { - **dict(self._linear_logp.provenance), + **dict(strict_provenance), "runtime_platform": "cuda", "triton_used": False, "execution": { "role": "vllm_rollout_linear_logprob", "strict_backend": True, - "sampling_logits_source": "native_vllm", + "sampling_logits_source": "rlkernel_det_gemm_vllm", "logits_materialized": True, "padded_lm_head_alignment": True, + "duplicate_lm_head_gemm": False, }, "source_logits_shape": list(source_logits.shape), "source_logits_dtype": _dtype_name(source_logits), @@ -1204,7 +1545,11 @@ def __call__( f"{DEFAULT_NUM_VOCAB_TILES}" ) contract = _full_vocab_contract(source_logits) - dispatch = kernel_registry.get_logprob_op(contract, requested_backend=self.backend_id) + from rl_engine.kernels.registry import kernel_registry + + dispatch = kernel_registry.get_logprob_op( + contract, requested_backend=self.backend_id + ) if ( dispatch.provenance["actual_backend"] != self.backend_id or dispatch.provenance["fallback"] @@ -1260,7 +1605,10 @@ def __call__( logprobs_tensors.logprobs.size(0), device=logprobs_tensors.logprobs.device, ), - (logprobs_tensors.logprob_token_ids == token_ids.unsqueeze(1)) + ( + logprobs_tensors.logprob_token_ids + == token_ids.unsqueeze(1) + ) .to(torch.int64) .argmax(dim=1), ], diff --git a/rl_engine/integrations/linear_logp.py b/rl_engine/integrations/linear_logp.py index 4704403ee..21c7c8fad 100644 --- a/rl_engine/integrations/linear_logp.py +++ b/rl_engine/integrations/linear_logp.py @@ -122,8 +122,22 @@ def _resolve(self, *, tensor_parallel: bool) -> Any: "sm90_deterministic_linear_logp entry points" ) from exc op = ( - sm90_deterministic_linear_logp_tp if tensor_parallel else sm90_deterministic_linear_logp + sm90_deterministic_linear_logp_tp + if tensor_parallel + else sm90_deterministic_linear_logp ) + # This wrapper is the strict boundary. Fail closed if a refactor ever + # swaps in the native/provider logp path under the same API. + expected_name = ( + "sm90_deterministic_linear_logp_tp" + if tensor_parallel + else "sm90_deterministic_linear_logp" + ) + if getattr(op, "__name__", "") != expected_name: + raise RuntimeError( + "strict linear_logp resolved a non-deterministic implementation: " + f"{getattr(op, '__name__', op)!r}" + ) if tensor_parallel: self._tp_op = op else: @@ -157,7 +171,10 @@ def _temperature_tensor( raise ValueError( f"temperature must be scalar or one value per row, got {value.numel()}" ) - if bool((value <= 0).any().item()): + positive = torch.all(value > 0) + if positive.is_cuda: + torch._assert_async(positive, "temperature must be positive") + elif not bool(positive): raise ValueError("temperature must be positive") return value value = float(temperature) @@ -175,7 +192,13 @@ def _validate_targets(target_ids: torch.Tensor, *, rows: int, real_vocab_size: i 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()): + valid = torch.all((target_long >= 0) & (target_long < real_vocab_size)) + if valid.is_cuda: + torch._assert_async( + valid, + f"linear_logp target_ids must be in [0, {real_vocab_size})", + ) + elif not bool(valid): raise ValueError( f"linear_logp target_ids must be in [0, {real_vocab_size}), got invalid ids" ) @@ -246,6 +269,84 @@ def _validate_contract( cls._validate_targets(target_ids, rows=hidden.size(0), real_vocab_size=real) return tensor_parallel, requested_global, real + def from_local_logits( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + tp_group: Any, + vocab_start_index: int, + global_vocab_size: int, + real_vocab_size: int, + target: str = "rollout", + temperature: float | torch.Tensor | None = None, + ) -> torch.Tensor: + """Score deterministic local LM-head logits without a duplicate GEMM.""" + + if local_logits.ndim != 2 or local_logits.dtype != torch.bfloat16: + raise TypeError("strict reused LM-head logits must be 2-D bfloat16") + if not local_logits.is_cuda or torch.version.hip is not None: + raise RuntimeError("strict reused LM-head logits require NVIDIA CUDA") + if target_ids.device != local_logits.device: + raise ValueError("linear_logp target_ids must share the logits device") + rank, world = self._tp_coordinates(tp_group) + if tp_group is None or world <= 1: + raise ValueError("reused rollout LM-head logits require a multi-rank TP group") + local_vocab = int(local_logits.size(1)) + requested_global = int(global_vocab_size) + expected_global = local_vocab * world + expected_start = rank * local_vocab + if requested_global != expected_global or int(vocab_start_index) != expected_start: + raise ValueError( + "reused rollout LM-head logits must use equal rank-ordered TP shards" + ) + real = int(real_vocab_size) + if not 0 < real <= requested_global: + raise ValueError("real_vocab_size must be within the padded global vocabulary") + self._validate_targets( + target_ids, + rows=local_logits.size(0), + real_vocab_size=real, + ) + temperature_tensor = self._temperature_tensor( + temperature, + rows=local_logits.size(0), + device=local_logits.device, + ) + from rl_engine.kernels.ops.cuda.loss.linear_logp import ( + sm90_deterministic_logp_from_local_logits_tp, + ) + + result, _lse = sm90_deterministic_logp_from_local_logits_tp( + local_logits.contiguous(), + target_ids, + tp_group=tp_group, + vocab_start_index=int(vocab_start_index), + global_vocab_size=requested_global, + real_vocab_size=real, + temperature=temperature_tensor, + ) + self._last_provenance = { + **self._mismatch_provenance(), + "target": target, + "runtime_platform": "cuda", + "triton_used": False, + "actual_backend": self.backend_id, + "deterministic_linear_logp": True, + "strict_entrypoint": "sm90_deterministic_logp_from_local_logits_tp", + "local_logits_shape": list(local_logits.shape), + "target_shape": list(target_ids.shape), + "tp_group_present": True, + "vocab_start_index": int(vocab_start_index), + "global_vocab_size": requested_global, + "real_vocab_size": real, + "temperature": None if temperature is None else "provided", + "contract_version": "cuda-det-gemm-linear-logp-sm90-contract-v2", + "logits_materialized": True, + "lm_head_result_reused": True, + } + return result + @staticmethod def _mismatch_provenance() -> dict[str, Any]: case_id = os.getenv("RL_KERNEL_LOGP_CASE", "P/P").strip().upper() @@ -342,6 +443,12 @@ def __call__( "runtime_platform": "cuda", "triton_used": False, "actual_backend": actual_backend, + "deterministic_linear_logp": True, + "strict_entrypoint": ( + "sm90_deterministic_linear_logp_tp" + if tensor_parallel + else "sm90_deterministic_linear_logp" + ), "hidden_shape": list(hidden.shape), "hidden_dtype": str(hidden.dtype).replace("torch.", ""), "lm_head_weight_shape": list(lm_head_weight.shape), diff --git a/rl_engine/integrations/megatron_runtime.py b/rl_engine/integrations/megatron_runtime.py index 52e490c98..2f6b5b737 100644 --- a/rl_engine/integrations/megatron_runtime.py +++ b/rl_engine/integrations/megatron_runtime.py @@ -6,7 +6,8 @@ from __future__ import annotations import importlib -from collections.abc import Iterable +from collections.abc import Callable, Iterable +from types import MethodType from typing import Any import torch @@ -16,6 +17,7 @@ MegatronAttentionOperator, MegatronFFNOperator, MegatronLogpOperator, + _fused_rms_norm_input, ) from rl_engine.integrations.linear_logp import LinearLogpWrapper from rl_engine.integrations.megatron import MegatronIntegration @@ -25,6 +27,7 @@ _PATCH_MARKER = "__rl_kernel_original_forward__" _STRICT_ATTENTION_PATCH_MARKER = "__rl_kernel_original_strict_attention_init__" _STRICT_ATTENTION_PROJECTION_MARKER = "__rl_kernel_strict_attention_projection__" +_STRICT_TE_RMS_NORM_PATCH_MARKER = "__rl_kernel_original_strict_rms_norm_forward__" def _optional_class(path: str) -> type[Any] | None: @@ -82,6 +85,8 @@ def _patch_strict_attention_projections( column_linear_cls: type[Any] | None = None, row_linear_cls: type[Any] | None = None, det_gemm: Any | None = None, + copy_to_tp: Callable[[torch.Tensor], torch.Tensor] | None = None, + reduce_from_tp: Callable[[torch.Tensor], torch.Tensor] | None = None, ) -> None: """Use the shared deterministic GEMM for Megatron Attention projections.""" @@ -106,22 +111,60 @@ def _patch_strict_attention_projections( column_forward_impl = column_linear_cls._forward_impl row_forward_impl = row_linear_cls._forward_impl + def tp_mappings() -> tuple[ + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], + ]: + if copy_to_tp is not None and reduce_from_tp is not None: + return copy_to_tp, reduce_from_tp + from megatron.core.tensor_parallel.mappings import ( + copy_to_tensor_model_parallel_region, + reduce_from_tensor_model_parallel_region, + ) + + return ( + copy_to_tensor_model_parallel_region, + reduce_from_tensor_model_parallel_region, + ) + def deterministic_projection( input_value: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, ) -> torch.Tensor: input_2d = input_value.reshape(-1, input_value.shape[-1]) - output_2d = det_gemm(input_2d, weight.t().contiguous()) + linear = getattr(det_gemm, "linear", None) + output_2d = ( + linear(input_2d, weight) + if linear is not None + else det_gemm(input_2d, weight.t().contiguous()) + ) output = output_2d.reshape(*input_value.shape[:-1], weight.shape[0]) return output if bias is None else output + bias def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: attention_init(instance, *args, **kwargs) - setattr(instance.linear_qkv, _STRICT_ATTENTION_PROJECTION_MARKER, "qkv") - setattr(instance.linear_proj, _STRICT_ATTENTION_PROJECTION_MARKER, "o_proj") - # copy_to_tensor_model_parallel_region already owns this TP dgrad reduction. - instance.linear_qkv.allreduce_dgrad = False + qkv = instance.linear_qkv + projection = instance.linear_proj + setattr(qkv, _STRICT_ATTENTION_PROJECTION_MARKER, "qkv") + setattr(projection, _STRICT_ATTENTION_PROJECTION_MARKER, "o_proj") + if hasattr(qkv, "layer_norm_weight"): + tp_copy, tp_reduce = tp_mappings() + + def te_qkv_forward(module: Any, input_value: torch.Tensor) -> Any: + normalized = _fused_rms_norm_input(module, input_value, "linear_qkv") + normalized = tp_copy(normalized) + return deterministic_projection(normalized, module.weight, None), None + + def te_projection_forward(module: Any, input_value: torch.Tensor) -> Any: + output = deterministic_projection(input_value, module.weight, None) + return tp_reduce(output), None + + qkv.forward = MethodType(te_qkv_forward, qkv) + projection.forward = MethodType(te_projection_forward, projection) + else: + # The local ColumnParallelLinear wrapper already owns TP dgrad. + qkv.allreduce_dgrad = False def column_forward_impl_wrapped( instance: Any, @@ -151,6 +194,36 @@ def row_forward_impl_wrapped( row_linear_cls._forward_impl = row_forward_impl_wrapped +def _patch_strict_te_rms_norm(rms_norm_cls: type[Any] | None = None) -> None: + """Match standalone TE RMSNorm modules to the strict rollout arithmetic.""" + + if rms_norm_cls is None: + try: + from transformer_engine.pytorch import RMSNorm + except ImportError: + return + + rms_norm_cls = RMSNorm + if hasattr(rms_norm_cls, _STRICT_TE_RMS_NORM_PATCH_MARKER): + return + original = rms_norm_cls.forward + + def wrapped(instance: Any, input_value: torch.Tensor, *args: Any, **kwargs: Any) -> Any: + if args or kwargs: + raise RuntimeError("strict TE RMSNorm does not accept extra forward arguments") + if bool(getattr(instance, "zero_centered_gamma", False)): + raise RuntimeError("strict TE RMSNorm does not support zero-centered gamma") + return torch.nn.functional.rms_norm( + input_value, + (input_value.shape[-1],), + instance.weight, + float(instance.eps), + ) + + setattr(rms_norm_cls, _STRICT_TE_RMS_NORM_PATCH_MARKER, original) + rms_norm_cls.forward = wrapped + + def install_megatron_integration( plan: IntegrationPlan, *, @@ -167,17 +240,17 @@ def install_megatron_integration( 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(), - ), - }, - ) + rl_kernel_operators: dict[str, Any] = {} + if plan.implementation_for("attention", "training") is Implementation.RL_KERNEL: + rl_kernel_operators["attention"] = MegatronAttentionOperator() + if plan.implementation_for("ffn", "training") is Implementation.RL_KERNEL: + rl_kernel_operators["ffn"] = MegatronFFNOperator() + if plan.implementation_for("logp", "training") is Implementation.RL_KERNEL: + rl_kernel_operators["logp"] = MegatronLogpOperator( + _provider_impl, + linear_logp=LinearLogpWrapper(), + ) + integration = MegatronIntegration(plan, rl_kernel_operators=rl_kernel_operators) resolved_attention = _unique_classes( _discover_attention_classes() if attention_classes is None else attention_classes ) @@ -193,6 +266,8 @@ def install_megatron_integration( set_active_integration("megatron", integration) if plan.implementation_for("attention", "training") is Implementation.RL_KERNEL: _patch_strict_attention_projections() + if plan.implementation_for("ffn", "training") is Implementation.RL_KERNEL: + _patch_strict_te_rms_norm() for cls in resolved_attention: _patch_forward(cls, integration=integration, module="attention") if resolved_attention: diff --git a/rl_engine/integrations/runtime.py b/rl_engine/integrations/runtime.py index 9dce3da75..6e0cc13d6 100644 --- a/rl_engine/integrations/runtime.py +++ b/rl_engine/integrations/runtime.py @@ -5,6 +5,7 @@ from __future__ import annotations +import atexit import json import os from collections import Counter @@ -27,6 +28,7 @@ class OperatorReadback: implementation: str backend_id: str call_count: int + execution_mode: str = "eager" provenance: Mapping[str, Any] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: @@ -51,10 +53,15 @@ def __init__( self.plan = plan self._rl_kernel_operators = dict(rl_kernel_operators) self._counts: Counter[str] = Counter() + self._profile_counts: Counter[str] = Counter() self._readbacks: dict[str, OperatorReadback] = {} self._installed_hooks: dict[str, str] = {} self._fallbacks: list[dict[str, str]] = [] + self._reported_routes: set[str] = set() + self._persisted_route_fingerprints: dict[str, tuple[Any, ...]] = {} self._lock = Lock() + if os.getenv("RL_KERNEL_READBACK_DIR", "").strip(): + atexit.register(self._persist_readback) def install_operator(self, module: str, operator: Callable[..., Any]) -> None: normalized = module.strip().lower() @@ -112,6 +119,45 @@ def execute( # inside the captured graph. if torch._dynamo.is_compiling(): return result + raw_provenance = getattr(selected, "provenance", {}) + if _contains_profile_only(raw_provenance): + self.record_profile_call(normalized) + else: + self.record_execution(normalized, selected) + return result + + def record_profile_call(self, module: str) -> None: + """Record framework shape/profile work without treating it as execution.""" + + normalized = module.strip().lower() + if normalized not in self.plan.cases: + raise ValueError(f"unknown integration module {module!r}") + with self._lock: + self._profile_counts[normalized] += 1 + first_call = self._profile_counts[normalized] == 1 + if first_call: + self._persist_readback() + + def record_execution( + self, + module: str, + selected: Callable[..., Any], + *, + execution_mode: str = "eager", + ) -> None: + """Record one eager or custom-op execution outside Dynamo tracing.""" + + normalized = module.strip().lower() + implementation = self.plan.implementation_for(normalized, self.target) + if implementation is Implementation.RL_KERNEL and ( + selected is not self._rl_kernel_operators.get(normalized) + ): + raise RuntimeError( + f"{self.framework} {normalized} execution evidence did not use " + "the installed RL-Kernel operator" + ) + if execution_mode not in {"eager", "compiled_cuda_graph"}: + raise ValueError(f"unknown execution mode {execution_mode!r}") backend_id = getattr(selected, "backend_id", None) if not isinstance(backend_id, str) or not backend_id.strip(): backend_id = ( @@ -121,6 +167,17 @@ def execute( ) raw_provenance = getattr(selected, "provenance", {}) provenance = dict(raw_provenance) if isinstance(raw_provenance, Mapping) else {} + actual_backend = _actual_backend(provenance) or backend_id + fallback = _contains_fallback(provenance) + should_report = False + route_changed = False + route_fingerprint = ( + implementation.value, + backend_id, + actual_backend, + fallback, + _runtime_platform(provenance), + ) with self._lock: self._counts[normalized] += 1 case = self.plan.cases[normalized] @@ -132,10 +189,30 @@ def execute( implementation=implementation.value, backend_id=backend_id, call_count=self._counts[normalized], + execution_mode=execution_mode, provenance=provenance, ) - self._persist_readback() - return result + if normalized not in self._reported_routes: + self._reported_routes.add(normalized) + should_report = _route_report_enabled() + if self._persisted_route_fingerprints.get(normalized) != route_fingerprint: + self._persisted_route_fingerprints[normalized] = route_fingerprint + route_changed = True + if route_changed: + self._persist_readback() + if should_report: + print( + "[RL-Kernel][route] " + f"framework={self.framework} target={self.target} module={normalized} " + f"requested={implementation.value} actual={actual_backend} " + f"fallback={str(fallback).lower()}", + flush=True, + ) + if implementation is Implementation.RL_KERNEL and fallback: + self.record_fallback(normalized, f"operator provenance selected {actual_backend}") + raise RuntimeError( + f"{self.framework} {normalized} strict RL-Kernel route reported fallback" + ) def readback(self) -> dict[str, Any]: with self._lock: @@ -145,6 +222,7 @@ def readback(self) -> dict[str, Any]: "plan": self.plan.to_dict(), "installed_hooks": dict(self._installed_hooks), "fallbacks": list(self._fallbacks), + "profile_calls": dict(self._profile_counts), "operators": { module: readback.to_dict() for module, readback in self._readbacks.items() }, @@ -222,6 +300,63 @@ def _contains_triton(value: Any) -> bool: return False +def _contains_fallback(value: Any) -> bool: + if isinstance(value, Mapping): + for key, item in value.items(): + normalized_key = str(key).strip().lower() + if normalized_key in { + "fallback", + "fallback_used", + "split_kv_fallback", + "used_fallback", + } and item not in (False, None, "", 0): + return True + if _contains_fallback(item): + return True + return False + if isinstance(value, (list, tuple)): + return any(_contains_fallback(item) for item in value) + return False + + +def _contains_profile_only(value: Any) -> bool: + if not isinstance(value, Mapping): + return False + if value.get("profile_only") is True: + return True + return any(_contains_profile_only(item) for item in value.values()) + + +def _actual_backend(value: Any) -> str | None: + if not isinstance(value, Mapping): + return None + direct = value.get("actual_backend") + if isinstance(direct, str) and direct.strip(): + return direct.strip() + for item in value.values(): + nested = _actual_backend(item) + if nested is not None: + return nested + return None + + +def _route_report_enabled() -> bool: + enabled = os.getenv("RL_KERNEL_ROUTE_REPORT", "1").strip().lower() + if enabled in {"0", "false", "no", "off"}: + return False + all_ranks = os.getenv("RL_KERNEL_ROUTE_REPORT_ALL_RANKS", "0").strip().lower() + if all_ranks in {"1", "true", "yes", "on"}: + return True + for name in ("RANK", "LOCAL_RANK"): + value = os.getenv(name, "").strip() + if value: + try: + return int(value) == 0 + except ValueError: + continue + return True + + def _runtime_platform(value: Any) -> str | None: if not isinstance(value, Mapping): return None diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index efddf3e8c..be83c743b 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -12,6 +12,7 @@ import torch +from rl_engine.distributed.collectives import DETERMINISTIC_ALL_REDUCE_OP from rl_engine.integrations.ablation import ( Implementation, IntegrationPlan, @@ -29,12 +30,22 @@ ) from rl_engine.integrations.state import get_active_integration, set_active_integration from rl_engine.integrations.vllm import VllmIntegration +from rl_engine.kernels.ops.pytorch.ffn.ffn import register_packed_inference_observer +from rl_engine.kernels.ops.pytorch.norm.rms_norm import ( + strict_add_rms_norm, + strict_rms_norm, +) _PATCH_MARKER = "__rl_kernel_original_forward__" _STRICT_MODEL_PATCH_MARKER = "__rl_kernel_original_strict_model_init__" _STRICT_PROJECTION_MARKER = "__rl_kernel_strict_attention_projection__" +_STRICT_FFN_INIT_MARKER = "__rl_kernel_original_strict_ffn_init__" +_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__" _RLK_ATTENTION_BACKEND: type[Any] | None = None _RLK_ATTENTION_IMPL: type[Any] | None = None +_RLK_ATTENTION_BUILDER: type[Any] | None = None def _is_worker_sampler_profile_batch(value: Any) -> bool: @@ -85,7 +96,9 @@ def plan_from_environment() -> IntegrationPlan: return integration_plan_from_environment() -def configure_vllm_environment(plan: IntegrationPlan, *, readback_dir: str | None = None) -> None: +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" @@ -123,7 +136,9 @@ def _patch_qwen_lm_head_padding() -> None: original = ParallelLMHead.__init__ original_weight_loader = VocabParallelEmbedding.weight_loader - def strict_weight_loader(instance: Any, param: Any, loaded_weight: torch.Tensor) -> None: + 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) @@ -168,7 +183,9 @@ def wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: setattr(ParallelLMHead, _PATCH_MARKER, original) if not hasattr(VocabParallelEmbedding, "__rl_kernel_original_weight_loader__"): - VocabParallelEmbedding.__rl_kernel_original_weight_loader__ = original_weight_loader + VocabParallelEmbedding.__rl_kernel_original_weight_loader__ = ( + original_weight_loader + ) VocabParallelEmbedding.weight_loader = strict_weight_loader ParallelLMHead.__init__ = wrapped @@ -193,14 +210,23 @@ def wrapped( *, _original: Any = original, ) -> Any: + lm_head = getattr(instance, "lm_head", None) + if lm_head is None: + raise RuntimeError("strict rollout linear_logp requires a Qwen LM-head") + setattr(lm_head, _STRICT_PROJECTION_MARKER, "lm_head") 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") + 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( @@ -211,7 +237,9 @@ def wrapped( 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) + getattr( + lm_head, "_rl_kernel_real_vocab_size", lm_head.org_vocab_size + ) ), ) return logits @@ -223,14 +251,111 @@ def wrapped( integration.record_installed_hook("logp", ",".join(installed)) +def _patch_strict_lm_head_linear( + *, + linear_method_cls: type[Any] | None = None, + det_gemm: Any | None = None, +) -> None: + """Route vLLM's existing LM-head projection through RL-Kernel det_gemm.""" + + if linear_method_cls is None: + from vllm.model_executor.layers.vocab_parallel_embedding import ( + UnquantizedEmbeddingMethod, + ) + + linear_method_cls = UnquantizedEmbeddingMethod + if hasattr(linear_method_cls, _STRICT_LM_HEAD_LINEAR_PATCH_MARKER): + return + previous_apply = linear_method_cls.apply + if det_gemm is None: + from rl_engine.kernels.ops.cuda.matmul.det_gemm import DetGemmOp + + det_gemm = DetGemmOp() + + def wrapped( + method: Any, + layer: Any, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if getattr(layer, _STRICT_PROJECTION_MARKER, None) != "lm_head": + return previous_apply(method, layer, x, bias) + x_2d = x.reshape(-1, x.shape[-1]) + output_2d = det_gemm.linear(x_2d, layer.weight) + output = output_2d.reshape(*x.shape[:-1], layer.weight.size(0)) + return output if bias is None else output + bias + + setattr( + linear_method_cls, + _STRICT_LM_HEAD_LINEAR_PATCH_MARKER, + previous_apply, + ) + linear_method_cls.apply = wrapped + + +def _configure_strict_ffn_compilation(vllm_config: Any | None = None) -> None: + """Keep the graph-safe TP reduction inside vLLM CUDA graphs.""" + + if vllm_config is None: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + raise RuntimeError( + "strict TP FFN initialization requires an active vLLM config" + ) + + compilation = vllm_config.compilation_config + splitting_ops = compilation.splitting_ops + if splitting_ops is None: + raise RuntimeError( + "vLLM splitting operators were not finalized before model init" + ) + # Older strict runtimes split this op because their host-owned sequence + # value was frozen at graph capture. Sequence allocation is now device + # owned, so remove stale entries and preserve the user's graph mode. + splitting_ops[:] = [ + op for op in splitting_ops if op != DETERMINISTIC_ALL_REDUCE_OP + ] + + def _patch_qwen_ffn(integration: VllmIntegration) -> None: from vllm.model_executor.models.qwen2 import Qwen2MLP - integration.install_operator("ffn", VllmFFNOperator()) + operator: VllmFFNOperator | None = None + if ( + integration.plan.implementation_for("ffn", "rollout") + is Implementation.RL_KERNEL + ): + operator = VllmFFNOperator() + integration.install_operator("ffn", operator) if hasattr(Qwen2MLP, _PATCH_MARKER): raise RuntimeError("vLLM Qwen2MLP is already RL-Kernel patched") original = Qwen2MLP.forward + if operator is not None: + if hasattr(Qwen2MLP, _STRICT_FFN_INIT_MARKER): + raise RuntimeError("vLLM Qwen2MLP init is already RL-Kernel patched") + original_init = Qwen2MLP.__init__ + compiled_evidence_armed = False + + def wrapped_init(instance: Any, *args: Any, **kwargs: Any) -> None: + nonlocal compiled_evidence_armed + original_init(instance, *args, **kwargs) + _handle, tp_world_size = operator.bind_packed_inference(instance) + if not compiled_evidence_armed: + register_packed_inference_observer( + lambda: integration.record_execution( + "ffn", operator, execution_mode="compiled_cuda_graph" + ) + ) + compiled_evidence_armed = True + if tp_world_size > 1: + _configure_strict_ffn_compilation() + + setattr(Qwen2MLP, _STRICT_FFN_INIT_MARKER, original_init) + Qwen2MLP.__init__ = wrapped_init + def wrapped(instance: Any, hidden_states: Any) -> Any: def native(_module: Any, value: Any) -> Any: return original(instance, value) @@ -239,7 +364,9 @@ def native(_module: Any, value: Any) -> Any: setattr(Qwen2MLP, _PATCH_MARKER, original) Qwen2MLP.forward = wrapped - integration.record_installed_hook("ffn", "vllm.model_executor.models.qwen2.Qwen2MLP.forward") + integration.record_installed_hook( + "ffn", "vllm.model_executor.models.qwen2.Qwen2MLP.forward" + ) def _install_flash_attn_ops_compatibility() -> None: @@ -263,22 +390,29 @@ def _install_flash_attn_ops_compatibility() -> None: def _patch_qwen3_strict_model( *, rms_norm_cls: type[Any] | None = None, + rotary_cls: type[Any] | None = None, linear_method_cls: type[Any] | None = None, attention_cls: type[Any] | None = None, det_gemm: Any | None = None, ) -> None: """Align vLLM's RMSNorm and Attention projections with Megatron.""" - if rms_norm_cls is None or linear_method_cls is None or attention_cls is None: + 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.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 - if hasattr(attention_cls, _STRICT_MODEL_PATCH_MARKER): - return + assert rms_norm_cls is not None + assert linear_method_cls is not None + assert attention_cls is not None if det_gemm is None: from rl_engine.kernels.ops.cuda.matmul.det_gemm import DetGemmOp @@ -286,6 +420,7 @@ def _patch_qwen3_strict_model( attention_init = attention_cls.__init__ unquantized_apply = linear_method_cls.apply + original_rms_forward_native = rms_norm_cls.forward_native def deterministic_linear_apply( method: Any, @@ -296,7 +431,12 @@ def deterministic_linear_apply( if not hasattr(layer, _STRICT_PROJECTION_MARKER): return unquantized_apply(method, layer, x, bias) x_2d = x.reshape(-1, x.shape[-1]) - output_2d = det_gemm(x_2d, layer.weight.t().contiguous()) + linear = getattr(det_gemm, "linear", None) + output_2d = ( + linear(x_2d, layer.weight) + if linear is not None + else det_gemm(x_2d, layer.weight.t().contiguous()) + ) output = output_2d.reshape(*x.shape[:-1], layer.weight.shape[0]) return output if bias is None else output + bias @@ -306,20 +446,46 @@ def strict_rms_norm_forward_cuda( residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if instance.variance_size_override is not None: - return instance.forward_native(x, residual) - if residual is not None: - residual = x + residual - x = residual + return original_rms_forward_native(instance, x, residual) weight = instance.weight.data if instance.has_weight else None - normalized = torch.nn.functional.rms_norm( + if weight is None: + return original_rms_forward_native(instance, x, residual) + if residual is None: + return strict_rms_norm( + x, + weight, + eps=instance.variance_epsilon, + ) + return strict_add_rms_norm( x, - (instance.hidden_size,), + residual, weight, - instance.variance_epsilon, + eps=instance.variance_epsilon, ) - if residual is None: - return normalized - return normalized, residual + + if not hasattr(rms_norm_cls, _STRICT_RMS_NORM_INIT_MARKER): + rms_norm_init = rms_norm_cls.__init__ + + def rms_norm_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: + rms_norm_init(instance, *args, **kwargs) + instance._forward_method = instance.forward_cuda + + setattr(rms_norm_cls, _STRICT_RMS_NORM_INIT_MARKER, rms_norm_init) + rms_norm_cls.__init__ = rms_norm_init_wrapped + rms_norm_cls.forward_cuda = strict_rms_norm_forward_cuda + + if rotary_cls is not None and not hasattr(rotary_cls, _STRICT_ROTARY_INIT_MARKER): + rotary_init = rotary_cls.__init__ + + def rotary_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: + rotary_init(instance, *args, **kwargs) + instance._forward_method = instance.forward_cuda + + setattr(rotary_cls, _STRICT_ROTARY_INIT_MARKER, rotary_init) + rotary_cls.__init__ = rotary_init_wrapped + + if hasattr(attention_cls, _STRICT_MODEL_PATCH_MARKER): + return def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: attention_init(instance, *args, **kwargs) @@ -327,7 +493,6 @@ def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: setattr(instance.o_proj, _STRICT_PROJECTION_MARKER, "o_proj") setattr(attention_cls, _STRICT_MODEL_PATCH_MARKER, attention_init) - rms_norm_cls.forward_cuda = strict_rms_norm_forward_cuda linear_method_cls.apply = deterministic_linear_apply attention_cls.__init__ = attention_init_wrapped @@ -338,8 +503,9 @@ def _patch_sampler(integration: VllmIntegration, *, strict_linear_logp: bool) -> 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) + if strict_linear_logp: + operator = VllmLogpOperator(original, strict_linear_logp=True) + integration.install_operator("logp", operator) def wrapped(instance: Any, *args: Any, **kwargs: Any) -> Any: sampling_metadata = kwargs.get("sampling_metadata") @@ -364,7 +530,9 @@ def native(_sampler: Any, *call_args: Any, **call_kwargs: Any) -> Any: integration.record_installed_hook("logp", "vllm.v1.sample.sampler.Sampler.forward") -def _patch_worker_sampler(integration: VllmIntegration, *, strict_linear_logp: bool) -> None: +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 @@ -391,53 +559,85 @@ 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): + 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__") + 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 + global _RLK_ATTENTION_BACKEND, _RLK_ATTENTION_BUILDER, _RLK_ATTENTION_IMPL from vllm.v1.attention.backends.flash_attn import ( FlashAttentionBackend, FlashAttentionImpl, + FlashAttentionMetadataBuilder, ) from vllm.v1.attention.backends.registry import ( AttentionBackendEnum, register_backend, ) - operator = VllmAttentionOperator() - integration.install_operator("attention", operator) + operator: VllmAttentionOperator | None = None + if ( + integration.plan.implementation_for("attention", "rollout") + is Implementation.RL_KERNEL + ): + operator = VllmAttentionOperator() + integration.install_operator("attention", operator) class RlKernelFlashAttentionImpl(FlashAttentionImpl): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + if operator is not None: + operator.bind_inference() + 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") + 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 RlKernelFlashAttentionMetadataBuilder(FlashAttentionMetadataBuilder): + pass + class RlKernelFlashAttentionBackend(FlashAttentionBackend): @staticmethod def get_impl_cls() -> type[Any]: return RlKernelFlashAttentionImpl + @staticmethod + def get_builder_cls() -> type[Any]: + return RlKernelFlashAttentionMetadataBuilder + RlKernelFlashAttentionImpl.__module__ = __name__ RlKernelFlashAttentionImpl.__qualname__ = "RlKernelFlashAttentionImpl" + RlKernelFlashAttentionMetadataBuilder.__module__ = __name__ + RlKernelFlashAttentionMetadataBuilder.__qualname__ = ( + "RlKernelFlashAttentionMetadataBuilder" + ) RlKernelFlashAttentionBackend.__module__ = __name__ RlKernelFlashAttentionBackend.__qualname__ = "RlKernelFlashAttentionBackend" _RLK_ATTENTION_IMPL = RlKernelFlashAttentionImpl + _RLK_ATTENTION_BUILDER = RlKernelFlashAttentionMetadataBuilder _RLK_ATTENTION_BACKEND = RlKernelFlashAttentionBackend globals()["RlKernelFlashAttentionImpl"] = RlKernelFlashAttentionImpl + globals()[ + "RlKernelFlashAttentionMetadataBuilder" + ] = RlKernelFlashAttentionMetadataBuilder 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 @@ -446,7 +646,9 @@ def get_impl_cls() -> type[Any]: AttentionBackendEnum.FLASH_ATTN, f"{__name__}.RlKernelFlashAttentionBackend", ) - integration.record_installed_hook("attention", f"{__name__}.RlKernelFlashAttentionBackend") + integration.record_installed_hook( + "attention", f"{__name__}.RlKernelFlashAttentionBackend" + ) def install_vllm_integration(plan: IntegrationPlan) -> VllmIntegration: @@ -457,18 +659,27 @@ def install_vllm_integration(plan: IntegrationPlan) -> VllmIntegration: 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") + 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 - strict_attention = plan.implementation_for("attention", "rollout") is Implementation.RL_KERNEL - if strict_attention or strict_linear_logp: - _install_flash_attn_ops_compatibility() + # vLLM imports this legacy rotary path even when every operator is routed + # to production. FA4-only installations need the bundled compatibility + # namespace before any attention backend is imported. + _install_flash_attn_ops_compatibility() + strict_linear_logp = ( + plan.implementation_for("logp", "rollout") is Implementation.RL_KERNEL + ) + strict_attention = ( + plan.implementation_for("attention", "rollout") is Implementation.RL_KERNEL + ) if strict_attention: _patch_qwen3_strict_model() if strict_linear_logp: _patch_qwen_lm_head_padding() + _patch_strict_lm_head_linear() _patch_qwen_compute_logits(integration) # Patch every boundary so P/R cases also produce production readback. The diff --git a/rl_engine/kernels/ops/cuda/attention/cp_comm.py b/rl_engine/kernels/ops/cuda/attention/cp_comm.py index d1470bfda..2519cbd73 100644 --- a/rl_engine/kernels/ops/cuda/attention/cp_comm.py +++ b/rl_engine/kernels/ops/cuda/attention/cp_comm.py @@ -364,16 +364,24 @@ def _get_collective(self, plan: AttentionCPCommunicationPlan): if self._collective is not None: return self._collective try: - from rl_engine.distributed import DeterministicCollective + from rl_engine.distributed.collectives import collective_for_group except ImportError as exc: raise AttentionCPCommunicationUnavailable( "self-owned CUDA AG/RS requires PR311/PR312 DeterministicCollective" ) from exc try: - self._collective = DeterministicCollective( - group=self._process_group, + dist = self._dist() + group = ( + self._process_group + if self._process_group is not None + else dist.group.WORLD + ) + self._collective = collective_for_group( + group=group, device=torch.device("cuda", torch.cuda.current_device()), ) + if self._collective is None: + raise RuntimeError("the CP process group is unavailable") except (RuntimeError, ValueError, TypeError) as exc: raise AttentionCPCommunicationUnavailable( f"self-owned CUDA AG/RS is unavailable: {exc}" diff --git a/rl_engine/kernels/ops/cuda/attention/flash_attn.py b/rl_engine/kernels/ops/cuda/attention/flash_attn.py index 1cf594164..3282c82fd 100644 --- a/rl_engine/kernels/ops/cuda/attention/flash_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/flash_attn.py @@ -29,6 +29,9 @@ _FA4_REQUIRED_PARAMETERS = frozenset( {"softmax_scale", "causal", "num_splits", "pack_gqa", "deterministic", "return_lse"} ) +_FA4_PAGED_REQUIRED_PARAMETERS = _FA4_REQUIRED_PARAMETERS | frozenset( + {"page_table", "seqused_k", "max_seqlen_k"} +) class StrictFlashAttentionUnavailable(RuntimeError): @@ -81,6 +84,7 @@ def __init__( *, split_kv: SplitKVSpec | None = None, _op: Callable[..., Any] | None = None, + _paged_op: Callable[..., Any] | None = None, _package_version: str | None = None, ) -> None: requested = SplitKVSpec.disabled() if split_kv is None else split_kv @@ -90,15 +94,22 @@ def __init__( raise ValueError("strict FA4 Attention requires Split-KV to be disabled") if _op is None: op, package_version, api_source = _load_fa4_cute_op() + paged_op = getattr( + importlib.import_module(api_source), "flash_attn_varlen_func", None + ) else: op = _op + paged_op = _paged_op package_version = "test-double" if _package_version is None else _package_version api_source = _FA4_API_SOURCE self._validate_api(op) + if paged_op is not None: + self._validate_paged_api(paged_op) self.split_kv = requested self.package_version = package_version self.api_source = api_source self._op = op + self._paged_op = paged_op @staticmethod def _validate_api(op: Callable[..., Any]) -> None: @@ -114,6 +125,21 @@ def _validate_api(op: Callable[..., Any]) -> None: "FlashAttention CuTe API is missing strict controls: " + ", ".join(missing) ) + @staticmethod + def _validate_paged_api(op: Callable[..., Any]) -> None: + try: + parameters = inspect.signature(op).parameters + except (TypeError, ValueError) as exc: + raise StrictFlashAttentionUnavailable( + "cannot inspect the FlashAttention CuTe paged API signature" + ) from exc + missing = sorted(_FA4_PAGED_REQUIRED_PARAMETERS.difference(parameters)) + if missing: + raise StrictFlashAttentionUnavailable( + "FlashAttention CuTe paged API is missing strict controls: " + + ", ".join(missing) + ) + def __call__( self, q: torch.Tensor, @@ -255,6 +281,141 @@ def forward_bshd_with_lse( }, ) + def forward_paged_bshd_with_lse( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + max_seqlen_k: int, + scale: float | None = None, + output_dtype: torch.dtype | None = None, + out: torch.Tensor | None = None, + ) -> DeterministicAttentionCoreResult: + """Run the strict FA4 schedule directly over a paged KV cache.""" + + if self._paged_op is None: + raise StrictFlashAttentionUnavailable( + "strict paged Attention requires flash_attn_varlen_func" + ) + self._validate_paged_bshd_inputs( + q, + k_cache, + v_cache, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_k=max_seqlen_k, + ) + resolved_dtype = q.dtype if output_dtype is None else output_dtype + if resolved_dtype != q.dtype: + raise ValueError("strict Attention output_dtype must match the Q/K/V input dtype") + if out is not None: + if out.shape != q.shape: + raise ValueError("paged Attention out must have the same shape as q") + if out.dtype != resolved_dtype or out.device != q.device: + raise ValueError("paged Attention out must match the requested dtype and device") + if not out.is_contiguous(): + raise ValueError("paged Attention out must be contiguous") + + paged_kwargs = { + "page_table": page_table, + "seqused_k": seqused_k, + "max_seqlen_k": max_seqlen_k, + "softmax_scale": scale, + "causal": False, + "num_splits": self.num_splits, + "pack_gqa": q.size(2) > k_cache.size(2), + "deterministic": self.deterministic_backward, + "return_lse": True, + } + if out is not None: + paged_kwargs["out"] = out + out_fa, lse = self._paged_op( + q, + k_cache, + v_cache, + **paged_kwargs, + ) + expected_lse_shape = (q.size(0), q.size(2), q.size(1)) + if not isinstance(out_fa, torch.Tensor) or tuple(out_fa.shape) != tuple(q.shape): + raise StrictFlashAttentionUnavailable( + f"paged FlashAttention output must have shape {tuple(q.shape)}" + ) + if not isinstance(lse, torch.Tensor) or tuple(lse.shape) != expected_lse_shape: + raise StrictFlashAttentionUnavailable( + f"paged FlashAttention LSE must have shape {expected_lse_shape}" + ) + if out_fa.dtype != resolved_dtype or lse.dtype != torch.float32: + raise StrictFlashAttentionUnavailable( + "paged FlashAttention returned an incompatible output dtype" + ) + if out is not None and out_fa.data_ptr() != out.data_ptr(): + raise StrictFlashAttentionUnavailable( + "paged FlashAttention did not write to the requested output buffer" + ) + return DeterministicAttentionCoreResult( + out=out_fa, + lse=lse.contiguous(), + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "attention_backend": f"{self.backend_id}.paged", + "fa_api_source": self.api_source, + "fa_package_version": self.package_version, + "num_splits": self.num_splits, + "deterministic_backward": self.deterministic_backward, + "dropout_p": 0.0, + "split_kv": self.split_kv.resolve( + max_seqlen_k, backend=self.backend_id + ).to_dict(), + "merge_order": self.merge_order, + "accum_dtype": self.accum_dtype, + "downcast_at": self.downcast_at, + "kv_layout": "paged_direct", + "output_buffer_reused": out is not None, + "fallback": False, + "fallback_reason": None, + "native_attention_arithmetic": True, + "production_ready": True, + "reference_only": False, + }, + ) + + @staticmethod + def _validate_paged_bshd_inputs( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + max_seqlen_k: int, + ) -> None: + if q.ndim != 4 or q.size(1) != 1: + raise ValueError("paged q must use [B, 1, H, D]") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("paged k/v must use [pages, page_size, H, D]") + if q.size(2) % k_cache.size(2) != 0 or q.size(3) != k_cache.size(3): + raise ValueError("paged q/k head counts or head dimensions are incompatible") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict paged Attention supports FP16/BF16 only") + if k_cache.dtype != q.dtype or v_cache.dtype != q.dtype: + raise ValueError("paged q/k/v must share one dtype") + if not (q.is_cuda and k_cache.is_cuda and v_cache.is_cuda): + raise ValueError("strict paged Attention requires CUDA tensors") + if not (q.device == k_cache.device == v_cache.device): + raise ValueError("paged q/k/v must be on one CUDA device") + if page_table.shape[0] != q.size(0) or page_table.dtype != torch.int32: + raise ValueError("page_table must be int32 with one row per query") + if seqused_k.shape != (q.size(0),) or seqused_k.dtype != torch.int32: + raise ValueError("seqused_k must be int32 with one length per query") + if page_table.device != q.device or seqused_k.device != q.device: + raise ValueError("paged Attention metadata must be on the Q device") + if max_seqlen_k <= 0 or max_seqlen_k > page_table.size(1) * k_cache.size(1): + raise ValueError("max_seqlen_k exceeds the page table capacity") + @staticmethod def _validate_bshd_inputs( q: torch.Tensor, @@ -268,9 +429,9 @@ def _validate_bshd_inputs( ) if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: raise ValueError("q/k/v must be 4-D [B, S, H, D]") - if q.size(0) != 1: - raise ValueError("strict FA4 core executes one logical batch row at a time") - if k.size(0) != 1 or v.size(0) != 1: + if q.size(0) < 1: + raise ValueError("strict FA4 core requires at least one logical batch row") + if k.size(0) != q.size(0) or v.size(0) != q.size(0): raise ValueError("q/k/v batch sizes must match") if k.shape != v.shape or q.size(3) != k.size(3): raise ValueError("k/v shapes and q/k/v head dimensions must match") @@ -298,9 +459,9 @@ def _validate_inputs( ) if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: raise ValueError("q/k/v must be 4-D [B, H, S, D]") - if q.size(0) != 1: - raise ValueError("strict FA4 core executes one logical batch row at a time") - if k.size(0) != 1 or v.size(0) != 1: + if q.size(0) < 1: + raise ValueError("strict FA4 core requires at least one logical batch row") + if k.size(0) != q.size(0) or v.size(0) != q.size(0): raise ValueError("q/k/v batch sizes must match") if k.shape != v.shape or q.size(3) != k.size(3): raise ValueError("k/v shapes and q/k/v head dimensions must match") diff --git a/rl_engine/kernels/ops/cuda/attention/strict_runtime.py b/rl_engine/kernels/ops/cuda/attention/strict_runtime.py index 90b20c971..009b16bf5 100644 --- a/rl_engine/kernels/ops/cuda/attention/strict_runtime.py +++ b/rl_engine/kernels/ops/cuda/attention/strict_runtime.py @@ -73,6 +73,7 @@ def forward_with_lse( cp_world_size: int, query_position_ids: torch.Tensor, key_position_ids: torch.Tensor, + positions_are_sorted: bool = False, ) -> StrictCUDAAttentionResult: self._require_nvidia_cuda(q) if cp_world_size != contract.sharding.cp_world_size: @@ -97,10 +98,21 @@ def forward_with_lse( 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) + if positions_are_sorted: + if cp_world_size != 1: + raise RuntimeError("pre-sorted Attention positions are supported only at CP=1") + q_sorted, k_sorted, v_sorted = global_q, global_k, global_v + q_positions_sorted, k_positions_sorted = global_q_positions, global_k_positions + q_sort = None + else: + q_sorted, q_positions_sorted, q_sort = self._sort_by_position( + global_q, global_q_positions + ) + k_sorted, k_positions_sorted, k_sort = self._sort_by_position( + global_k, global_k_positions + ) + v_sorted = self._gather_sequence(global_v, k_sort) + 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. @@ -108,57 +120,31 @@ def forward_with_lse( 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) + # FA4's full causal schedule produces the same output, LSE, and dQ as + # launching one single-query prefix at a time. Keeping all query rows + # in one launch removes O(sequence_length) Python and CUDA dispatches; + # deterministic=True still pins the backward implementation. + result = self._core.forward_bshd_with_lse( + q_fa, + k_fa, + v_fa, + causal=causal, + scale=scale, + query_position_ids=q_positions_sorted, + key_position_ids=k_positions_sorted, + output_dtype=q.dtype, + ) + out_sorted = result.out.transpose(1, 2).contiguous() + lse_sorted = result.lse + backend = ( + result.provenance.get("actual_backend") + or result.provenance.get("attention_backend") + or getattr(self._core, "backend_id", None) + ) if cp_world_size > 1: + if q_sort is None: + raise RuntimeError("CP Attention requires a framework position reorder") 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) @@ -187,8 +173,82 @@ def forward_with_lse( "reference_only": False, "split_kv": "disabled", "framework_position_reorder": True, - "query_schedule": "single_query_causal_prefix", - "core_rows": core_rows, + "query_schedule": "full_sequence_causal_single_launch", + "backward_schedule": "fa4_deterministic_full_sequence", + "core_row_count": q_sorted.size(0) * q_sorted.size(2), + "core_launch_count": 1, + "core_batch_size": q_sorted.size(0), + "core_query_length": q_sorted.size(2), + "core_actual_backends": [] if backend is None else [str(backend)], + }, + ) + + def forward_paged_with_lse( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + max_seqlen_k: int, + scale: float | None, + out: torch.Tensor | None = None, + ) -> StrictCUDAAttentionResult: + """Run strict inference Attention without materializing paged KV rows.""" + + self._require_nvidia_cuda(q) + if out is not None and ( + out.shape != q.shape or out.dtype != q.dtype or out.device != q.device + ): + raise ValueError("paged Attention out must match q shape, dtype, and device") + paged_out = None + if out is not None: + paged_out = torch.empty( + q.transpose(1, 2).shape, + dtype=q.dtype, + device=q.device, + ) + result = self._core.forward_paged_bshd_with_lse( + q.transpose(1, 2).contiguous(), + k_cache, + v_cache, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_k=max_seqlen_k, + scale=scale, + output_dtype=q.dtype, + out=paged_out, + ) + self.communication_executed = False + result_out = result.out.transpose(1, 2).contiguous() + if out is not None: + out.copy_(result_out) + result_out = out + return StrictCUDAAttentionResult( + out=result_out, + lse=result.lse, + provenance={ + **result.provenance, + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "communication_backend": "none", + "communication_executed": False, + "native_attention_arithmetic": True, + "production_ready": True, + "fallback": False, + "fallback_reason": None, + "reference_only": False, + "query_schedule": "paged_single_query_batch", + "core_row_count": q.size(0), + "core_launch_count": 1, + "core_batch_size": q.size(0), + "core_actual_backends": [ + result.provenance.get( + "attention_backend", "flash_attention_4.cute.paged" + ) + ], }, ) @@ -285,13 +345,19 @@ def _validate_global_positions( (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") + if positions.size(1) > 1: + torch._assert_async( + torch.all(positions[:, 1:] > positions[:, :-1]), + f"global {name} positions must be unique and increasing", + ) + if causal: + torch._assert_async( + torch.all( + query_positions + == key_positions[:, -query_positions.size(1) :] + ), + "causal Attention queries must be the trailing global KV positions", + ) __all__ = ["StrictCUDAAttentionResult", "StrictCUDAAttentionRuntime"] diff --git a/rl_engine/kernels/ops/cuda/loss/linear_logp.py b/rl_engine/kernels/ops/cuda/loss/linear_logp.py index 635829dfc..3b3f1431e 100644 --- a/rl_engine/kernels/ops/cuda/loss/linear_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/linear_logp.py @@ -9,7 +9,20 @@ import torch from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.matmul.det_gemm import ( + det_gemm_linear, + det_gemm_linear_input_gradient, + det_gemm_linear_weight_gradient, +) from rl_engine.kernels.ops.pytorch.loss.linear_logp import ( + _merge_tp_local_logp, + _require_distributed_initialized, + _validate_global_targets, + _validate_tp_targets_enabled, + _validate_tp_vocab_partition_cached, + _assert_global_targets_async, + _deterministic_tp_all_reduce_, + _validate_even_tp_vocab_partition_local, _merge_tp_local_logp, _require_distributed_initialized, _validate_global_targets, @@ -1079,7 +1092,6 @@ def _strict_tp_run( tp_group, ): required = ( - "det_gemm_fwd", "linear_logp_local_bf16_forward", "linear_logp_logits_bf16_to_dlogits", ) @@ -1091,24 +1103,29 @@ def _strict_tp_run( target = target.reshape(-1).to(device=hidden.device, dtype=torch.long).contiguous() if hidden_2d.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: raise TypeError("strict TP linear_logp requires BF16 hidden and weight") - global_vocab = _validate_tp_vocab_partition_cached( + global_vocab = _validate_even_tp_vocab_partition_local( tp_group=tp_group, - device=hidden.device, vocab_start_index=int(vocab_start), local_vocab_size=weight.size(0), global_vocab_size=int(global_vocab), ) if not 0 < int(real_vocab) <= global_vocab: raise ValueError(f"invalid real_vocab_size={real_vocab} for padded vocab={global_vocab}") - _validate_global_targets(target, int(real_vocab), tp_group) - dist = _require_distributed_initialized() - owners = ( - (target >= int(vocab_start)) & (target < int(vocab_start) + weight.size(0)) - ).to(torch.int32) - dist.all_reduce(owners, op=dist.ReduceOp.SUM, group=tp_group) - if bool((owners != 1).any().item()): - raise ValueError("each selected target must have exactly one TP LM-head owner") - logits = _C.det_gemm_fwd(hidden_2d, weight.t().contiguous()) + if _validate_tp_targets_enabled(): + # Opt-in diagnostics retain the synchronized cross-rank checks. The + # cached contiguous [0, V) partition plus a valid target already prove + # unique ownership during normal execution. + _validate_global_targets(target, int(real_vocab), tp_group) + dist = _require_distributed_initialized() + owners = ( + (target >= int(vocab_start)) & (target < int(vocab_start) + weight.size(0)) + ).to(torch.int32) + dist.all_reduce(owners, op=dist.ReduceOp.SUM, group=tp_group) + if bool((owners != 1).any().item()): + raise ValueError("each selected target must have exactly one TP LM-head owner") + else: + _assert_global_targets_async(target, int(real_vocab)) + logits = det_gemm_linear(hidden_2d, weight) temp = hidden_2d.new_empty(0, dtype=torch.float32) if bias is not None: logits = (logits.float() + bias.float().reshape(1, -1)).to(torch.bfloat16) @@ -1116,8 +1133,12 @@ def _strict_tp_run( temp = temperature.to(device=hidden.device, dtype=torch.float32).reshape(-1) if temp.numel() == 1: temp = temp.expand(hidden_2d.size(0)).contiguous() - if temp.numel() != hidden_2d.size(0) or bool((temp <= 0).any().item()): + if temp.numel() != hidden_2d.size(0): raise ValueError("temperature must be positive and scalar or per-token") + torch._assert_async( + (temp > 0).all(), + "temperature must be positive and scalar or per-token", + ) logits = (logits.float() / temp.reshape(-1, 1)).to(torch.bfloat16) local_real = max(0, min(weight.size(0), int(real_vocab) - int(vocab_start))) if local_real < weight.size(0): @@ -1131,6 +1152,74 @@ def _strict_tp_run( return logp, lse, hidden_2d, weight, target, logits, temp +def sm90_deterministic_logp_from_local_logits_tp( + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + vocab_start_index: int, + global_vocab_size: int, + real_vocab_size: int = -1, + temperature: Optional[torch.Tensor] = None, + tp_group: Any, +) -> tuple[torch.Tensor, torch.Tensor]: + """Merge a deterministic TP LM-head result without recomputing its GEMM.""" + + required = "linear_logp_local_bf16_forward" + if not (_EXT_AVAILABLE and hasattr(_C, required)): + raise RuntimeError(f"strict TP logits logp missing CUDA symbol: {required}") + if local_logits.ndim != 2 or local_logits.dtype != torch.bfloat16: + raise TypeError("strict TP logits logp requires 2-D BF16 local logits") + if not local_logits.is_cuda or not local_logits.is_contiguous(): + raise ValueError("strict TP logits logp requires contiguous CUDA logits") + + vocab_start = int(vocab_start_index) + global_vocab = _validate_even_tp_vocab_partition_local( + tp_group=tp_group, + vocab_start_index=vocab_start, + local_vocab_size=local_logits.size(1), + global_vocab_size=int(global_vocab_size), + ) + real_vocab = global_vocab if int(real_vocab_size) < 0 else int(real_vocab_size) + if not 0 < real_vocab <= global_vocab: + raise ValueError( + f"invalid real_vocab_size={real_vocab} for padded vocab={global_vocab}" + ) + target = ( + target_ids.reshape(-1) + .to(device=local_logits.device, dtype=torch.long) + .contiguous() + ) + if target.numel() != local_logits.size(0): + raise ValueError("target_ids must contain one id per local-logits row") + _assert_global_targets_async(target, real_vocab) + + logits = local_logits + if temperature is not None: + temp = temperature.to(device=logits.device, dtype=torch.float32).reshape(-1) + if temp.numel() == 1: + temp = temp.expand(logits.size(0)).contiguous() + if temp.numel() != logits.size(0): + raise ValueError("temperature must be positive and scalar or per-token") + torch._assert_async( + (temp > 0).all(), + "temperature must be positive and scalar or per-token", + ) + logits = (logits.float() / temp.reshape(-1, 1)).to(torch.bfloat16) + local_real = max(0, min(logits.size(1), real_vocab - vocab_start)) + if local_real < logits.size(1): + logits = logits.clone() + logits[:, local_real:] = float("-inf") + logits = logits.contiguous() + local_target, local_lse = _C.linear_logp_local_bf16_forward( + logits, + target, + vocab_start, + ) + logp, lse = _merge_tp_local_logp(local_lse, local_target, tp_group=tp_group) + lead_shape = target_ids.shape + return logp.reshape(lead_shape), lse.reshape(lead_shape) + + class _StrictTensorParallelLinearLogpAutograd(torch.autograd.Function): @staticmethod def forward( @@ -1188,14 +1277,15 @@ def backward(ctx, grad_logp, grad_lse): dlogits = dlogits.contiguous() grad_hidden = grad_weight = None if ctx.needs_input_grad[0]: - grad_hidden = _C.det_gemm_fwd(dlogits, weight) - dist = _require_distributed_initialized() - dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group) + grad_hidden = det_gemm_linear_input_gradient(dlogits, weight) + _deterministic_tp_all_reduce_(grad_hidden, ctx.tp_group) grad_hidden = grad_hidden.reshape(ctx.lead_shape + (weight.size(1),)).to( ctx.hidden_dtype ) if ctx.needs_input_grad[1]: - grad_weight = _C.det_gemm_fwd(dlogits.t().contiguous(), hidden).to(ctx.weight_dtype) + grad_weight = det_gemm_linear_weight_gradient(hidden, dlogits).to( + ctx.weight_dtype + ) grad_bias = None if ctx.has_bias and ctx.needs_input_grad[2]: grad_bias = dlogits.float().sum(dim=0).to(ctx.bias_dtype) diff --git a/rl_engine/kernels/ops/cuda/matmul/det_gemm.py b/rl_engine/kernels/ops/cuda/matmul/det_gemm.py index c410fbb23..df725e6ac 100644 --- a/rl_engine/kernels/ops/cuda/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/cuda/matmul/det_gemm.py @@ -1,17 +1,213 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Batch-invariant deterministic GEMM, CUDA path (WS1 #146). +"""Batch-invariant deterministic GEMM backends (WS1 #146). Hand-written kernel (csrc/cuda/gemm/det_gemm_kernel.cu): fixed K-accumulation -order, FP32 accumulation, no split-K. A row's output is invariant to batch size, -chunked-prefill, and padding. No PyTorch fallback -- a generic matmul (cuBLAS) -would silently break invariance (see NativeGemmOp). Tensor-parallel GEMM is WS2. +order, FP32 accumulation, no split-K. On Hopper, an explicitly configured +cuBLASLt no-split-K backend is also available. Both are strict implementations; +backend selection never falls back silently. Tensor-parallel GEMM is WS2. """ +import os +from collections.abc import Callable +from threading import Lock + import torch from rl_engine.kernels.ops.backward_runtime import record_backward from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE -from rl_engine.utils.logger import logger +from rl_engine.runtime_mode import rl_kernel_mode, route_report_enabled + +_BACKEND_ENV = "RL_KERNEL_DET_GEMM_BACKEND" +_AUTO_BACKEND = "auto" +_SM90_BACKEND = "sm90" +_CUBLASLT_BACKEND = "cublaslt_nosplitk" +_CUBLASLT_CONFIGURED = False +_ROUTE_REPORTED = False +_ROUTE_REPORT_LOCK = Lock() + + +def _requested_det_gemm_backend() -> str: + value = os.getenv(_BACKEND_ENV, _AUTO_BACKEND).strip().lower() + aliases = { + "cuda": _SM90_BACKEND, + "cublaslt": _CUBLASLT_BACKEND, + } + value = aliases.get(value, value) + if value not in {_AUTO_BACKEND, _SM90_BACKEND, _CUBLASLT_BACKEND}: + raise RuntimeError( + f"{_BACKEND_ENV} must be '{_AUTO_BACKEND}', '{_SM90_BACKEND}', " + f"or '{_CUBLASLT_BACKEND}', got {value!r}" + ) + return value + + +def det_gemm_backend() -> str: + """Return the selected strict GEMM implementation. + + Auto mode prefers the same Hopper cuBLASLt path as vLLM and uses the + self-contained SM90 kernel when the process was not started with the + no-split-K cuBLASLt contract. + """ + + value = _requested_det_gemm_backend() + if value == _AUTO_BACKEND: + return _CUBLASLT_BACKEND if _cublaslt_contract_ready() else _SM90_BACKEND + return value + + +def _cublaslt_contract_ready() -> bool: + return ( + os.getenv("CUBLAS_WORKSPACE_CONFIG", "").strip() == ":16:8" + and os.getenv("CUBLASLT_WORKSPACE_SIZE", "").strip() == "1" + and torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 9 + ) + + +def det_gemm_fallback_reason() -> str | None: + requested = _requested_det_gemm_backend() + if requested == _AUTO_BACKEND and not _cublaslt_contract_ready(): + return "cuBLASLt no-split-K startup contract is unavailable" + return None + + +def det_gemm_backend_id() -> str: + backend = det_gemm_backend() + if backend == _CUBLASLT_BACKEND: + return "rlkernel.det_gemm.cublaslt_nosplitk.v1" + return "rlkernel.det_gemm.sm90.v1" + + +def _report_strict_route_once() -> None: + global _ROUTE_REPORTED + if torch._dynamo.is_compiling(): + return + if not route_report_enabled(): + return + with _ROUTE_REPORT_LOCK: + if _ROUTE_REPORTED: + return + _ROUTE_REPORTED = True + fallback_reason = det_gemm_fallback_reason() + print( + f"[RL-Kernel][route] mode={rl_kernel_mode().value} module=gemm " + f"requested={_requested_det_gemm_backend()} " + f"actual={det_gemm_backend_id()} " + f"fallback={str(fallback_reason is not None).lower()}" + f"{'' if fallback_reason is None else ' reason=' + fallback_reason}", + flush=True, + ) + + +def _configure_cublaslt_nosplitk(a: torch.Tensor | None = None) -> None: + """Validate the Hopper batch-invariant cuBLASLt contract once per process.""" + + global _CUBLASLT_CONFIGURED + if _CUBLASLT_CONFIGURED: + return + if torch._dynamo.is_compiling(): + raise RuntimeError( + "cublaslt_nosplitk must be configured before torch.compile tracing" + ) + device = None if a is None else a.device + if torch.cuda.get_device_capability(device)[0] != 9: + raise RuntimeError("cublaslt_nosplitk is currently validated only on SM90") + workspace = os.getenv("CUBLAS_WORKSPACE_CONFIG", "").strip() + if workspace != ":16:8": + raise RuntimeError( + "cublaslt_nosplitk requires CUBLAS_WORKSPACE_CONFIG=:16:8 before " + "the process starts" + ) + if os.getenv("CUBLASLT_WORKSPACE_SIZE", "").strip() != "1": + raise RuntimeError( + "cublaslt_nosplitk requires CUBLASLT_WORKSPACE_SIZE=1 before the " + "process starts" + ) + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False + torch.backends.cuda.preferred_blas_library("cublaslt") + _CUBLASLT_CONFIGURED = True + + +def _require_sm90_backend() -> None: + marker = getattr(_C, "det_gemm_sm90_compiled", None) + if not _EXT_AVAILABLE or not callable(marker) or not marker(): + raise RuntimeError( + "strict RL-Kernel SM90 GEMM was requested, but the extension was " + "built without KERNEL_ALIGN_DET_GEMM_SM90=1; refusing naive fallback" + ) + + +def det_gemm_linear( + a: torch.Tensor, + weight: torch.Tensor, + *, + native_op: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, +) -> torch.Tensor: + """Apply a native [N,K] weight through the selected strict backend.""" + + if det_gemm_backend() == _CUBLASLT_BACKEND: + _configure_cublaslt_nosplitk(a) + _report_strict_route_once() + return torch.mm(a, weight.t()) + _require_sm90_backend() + _report_strict_route_once() + if native_op is not None: + return native_op(a, weight) + return _det_gemm_fwd_weight(a, weight) + + +def det_gemm_linear_input_gradient( + grad_output: torch.Tensor, + weight: torch.Tensor, + *, + native_op: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, +) -> torch.Tensor: + """Compute ``dX = dY @ weight`` through the selected strict backend.""" + + if det_gemm_backend() == _CUBLASLT_BACKEND: + _configure_cublaslt_nosplitk(grad_output) + _report_strict_route_once() + return torch.mm(grad_output, weight) + _require_sm90_backend() + _report_strict_route_once() + if native_op is not None: + return native_op(grad_output, weight) + return _C.det_gemm_fwd(grad_output, weight) + + +def det_gemm_linear_weight_gradient( + a: torch.Tensor, + grad_output: torch.Tensor, + *, + native_op: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, +) -> torch.Tensor: + """Compute ``dWeight = dY.T @ X`` through the selected strict backend.""" + + if det_gemm_backend() == _CUBLASLT_BACKEND: + _configure_cublaslt_nosplitk(a) + _report_strict_route_once() + return torch.mm(grad_output.t(), a) + _require_sm90_backend() + _report_strict_route_once() + if native_op is not None: + return native_op(a, grad_output) + return _C.det_gemm_db_transposed(a, grad_output) + + +@torch.library.custom_op("rl_kernel::det_gemm_fwd_weight", mutates_args=()) +def _det_gemm_fwd_weight(a: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Expose the native-weight SM90 GEMM through the PyTorch dispatcher.""" + + return _C.det_gemm_fwd_rhs_transposed(a, weight) + + +@_det_gemm_fwd_weight.register_fake +def _det_gemm_fwd_weight_fake( + a: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + return a.new_empty((*a.shape[:-1], weight.shape[0])) class _DetGemmFn(torch.autograd.Function): @@ -20,7 +216,9 @@ def forward(ctx, a, b, output_fp32=False): ctx.save_for_backward(a, b) if output_fp32: if not hasattr(_C, "det_gemm_fwd_fp32"): - raise RuntimeError("FP32 deterministic GEMM output requires the rebuilt extension") + raise RuntimeError( + "FP32 deterministic GEMM output requires the rebuilt extension" + ) return _C.det_gemm_fwd_fp32(a, b) return _C.det_gemm_fwd(a, b) @@ -41,6 +239,45 @@ def backward(ctx, grad_out): return da, db, None +class _DetLinearFn(torch.autograd.Function): + @staticmethod + def forward(ctx, a, weight): + ctx.save_for_backward(a, weight) + return det_gemm_linear(a, weight) + + @staticmethod + def backward(ctx, grad_out): + a, weight = ctx.saved_tensors + grad_out = grad_out.contiguous() + if grad_out.dtype != torch.bfloat16: + grad_out = grad_out.to(torch.bfloat16) + da = ( + det_gemm_linear_input_gradient( + grad_out, + weight, + native_op=_C.det_gemm_fwd, + ) + if ctx.needs_input_grad[0] + else None + ) + dweight = ( + det_gemm_linear_weight_gradient( + a, + grad_out, + native_op=_C.det_gemm_db_transposed, + ) + if ctx.needs_input_grad[1] + else None + ) + record_backward( + "det_gemm", + kernel_id=det_gemm_backend_id(), + impl="strict_det_gemm", + family="cuda", + ) + return da, dweight + + class _DetGemmAccumFn(torch.autograd.Function): @staticmethod def forward(ctx, a, b): @@ -81,15 +318,34 @@ class DetGemmOp: def __init__(self): self.has_hardware_op = False - if _EXT_AVAILABLE and hasattr(_C, "det_gemm_fwd"): - self.op = _C.det_gemm_fwd - self.has_hardware_op = True - logger.info("Successfully linked to RL-Kernel _C.det_gemm_fwd.") - else: - logger.warning( - "RL-Kernel _C.det_gemm_fwd unavailable; DetGemmOp requires the " - "compiled CUDA extension and has no batch-invariant fallback." + backend = det_gemm_backend() + if backend == _CUBLASLT_BACKEND: + _configure_cublaslt_nosplitk() + required = ( + "det_gemm_fwd", + "det_gemm_fwd_rhs_transposed", + "det_gemm_db_transposed", + ) + missing = ( + list(required) + if _C is None + else [name for name in required if not hasattr(_C, name)] + ) + if not _EXT_AVAILABLE or _C is None or missing: + detail = ( + "compiled extension unavailable" + if not _EXT_AVAILABLE or _C is None + else f"missing {', '.join(missing)}" + ) + raise RuntimeError( + "strict RL-Kernel GEMM requires the compiled CUDA extension; " + f"{detail}; refusing non-strict fallback" ) + if backend == _SM90_BACKEND: + _require_sm90_backend() + self.op = _C.det_gemm_fwd + self.has_hardware_op = True + _report_strict_route_once() def __call__(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, "BF16 only" @@ -101,6 +357,21 @@ def __call__(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ) return _DetGemmFn.apply(a.contiguous(), b.contiguous(), False) + def linear(self, a: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Apply a native [N,K] linear weight without materializing weight.T.""" + assert a.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, "BF16 only" + assert a.is_cuda and weight.is_cuda, "Inputs must be on CUDA device" + required = ( + "det_gemm_fwd_rhs_transposed", + "det_gemm_fwd", + "det_gemm_db_transposed", + ) + if not self.has_hardware_op or any(not hasattr(_C, name) for name in required): + raise RuntimeError( + "DetGemmOp.linear requires the rebuilt native-weight CUDA extension" + ) + return _DetLinearFn.apply(a.contiguous(), weight.contiguous()) + def forward_fp32(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, "BF16 only" assert a.is_cuda and b.is_cuda, "Inputs must be on CUDA device" diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py index bc0ddcefe..1b8290638 100644 --- a/rl_engine/kernels/ops/pytorch/ffn/ffn.py +++ b/rl_engine/kernels/ops/pytorch/ffn/ffn.py @@ -4,12 +4,23 @@ from __future__ import annotations +from collections.abc import Callable from typing import Any import torch from torch import Tensor +from rl_engine.distributed.collectives import ( + _COLLECTIVES as _SHARED_COLLECTIVES, + collective_for_group, + deterministic_all_reduce_inplace, +) from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.matmul.det_gemm import ( + det_gemm_linear, + det_gemm_linear_input_gradient, + det_gemm_linear_weight_gradient, +) QWEN3_8B_HIDDEN_SIZE = 4096 QWEN3_8B_INTERMEDIATE_SIZE = 12288 @@ -24,13 +35,98 @@ "swiglu_forward", "swiglu_backward", ) +_PACKED_SWIGLU_SYMBOLS = ( + "swiglu_packed_forward", + "swiglu_packed_backward", +) _REQUIRED_SYMBOLS = _DET_GEMM_SYMBOLS + _SWIGLU_SYMBOLS _COLLECTIVE_MIN_CAPACITY_BYTES = 64 * 1024 * 1024 -_COLLECTIVES: dict[tuple[int, int, int, int], Any] = {} +# Backward-compatible test hook; ownership lives in the shared communication layer. +_COLLECTIVES = _SHARED_COLLECTIVES +_PACKED_INFERENCE_OBSERVERS: list[Callable[[], None]] = [] + + +def register_packed_inference_observer(callback: Callable[[], None]) -> None: + """Arm one execution callback for the graph-captured rollout custom op.""" + if not callable(callback): + raise TypeError("packed inference observer must be callable") + _PACKED_INFERENCE_OBSERVERS.append(callback) -def _require_ffn_kernels(*, disable_split_k: bool) -> None: + +def _notify_packed_inference_observers() -> None: + callbacks = tuple(_PACKED_INFERENCE_OBSERVERS) + _PACKED_INFERENCE_OBSERVERS.clear() + for callback in callbacks: + callback() + + +@torch.library.custom_op("rl_kernel::qwen3_ffn_packed_inference", mutates_args=()) +def _qwen3_ffn_packed_inference( + rmsnorm_output: Tensor, + fused_gate_up_weight: Tensor, + down_weight: Tensor, +) -> Tensor: + """Run the graph-safe compute portion of the strict rollout FFN.""" + + _notify_packed_inference_observers() + input_shape = rmsnorm_output.shape + hidden_2d = rmsnorm_output.reshape(-1, input_shape[-1]).contiguous() + gate_up = det_gemm_linear( + hidden_2d, + fused_gate_up_weight, + native_op=_C.det_gemm_fwd_rhs_transposed, + ) + activated = _C.swiglu_packed_forward(gate_up) + output = det_gemm_linear( + activated, + down_weight, + native_op=_C.det_gemm_fwd_rhs_transposed, + ) + return output.reshape(*input_shape[:-1], output.size(-1)) + + +@_qwen3_ffn_packed_inference.register_fake +def _qwen3_ffn_packed_inference_fake( + rmsnorm_output: Tensor, + fused_gate_up_weight: Tensor, + down_weight: Tensor, +) -> Tensor: + del fused_gate_up_weight + return rmsnorm_output.new_empty((*rmsnorm_output.shape[:-1], down_weight.shape[0])) + + +def qwen3_ffn_packed_inference( + rmsnorm_output: Tensor, + fused_gate_up_weight: Tensor, + down_weight: Tensor, + *, + collective_handle: int = 0, + tp_world_size: int = 1, +) -> Tensor: + """Inference-only packed FFN entry compatible with torch.compile.""" + + output = _qwen3_ffn_packed_inference( + rmsnorm_output, + fused_gate_up_weight, + down_weight, + ) + if tp_world_size <= 1: + return output + if collective_handle <= 0: + raise RuntimeError("packed rollout FFN requires a bound TP collective") + return deterministic_all_reduce_inplace( + output, + collective_handle=collective_handle, + ) + + +def _require_ffn_kernels( + *, disable_split_k: bool, packed_gate_up: bool = False +) -> None: required = _REQUIRED_SYMBOLS if disable_split_k else _SWIGLU_SYMBOLS + if packed_gate_up: + required += _PACKED_SWIGLU_SYMBOLS missing = [name for name in required if not hasattr(_C, name)] if not _EXT_AVAILABLE or _C is None or missing: suffix = f" Missing symbols: {', '.join(missing)}." if missing else "" @@ -42,38 +138,38 @@ def _require_ffn_kernels(*, disable_split_k: bool) -> None: raise RuntimeError(f"qwen3_ffn requires the {needed}.{suffix}") -def _gemm_fwd(a: Tensor, b: Tensor, *, disable_split_k: bool) -> Tensor: +def _linear_fwd(a: Tensor, weight: Tensor, *, disable_split_k: bool) -> Tensor: if disable_split_k: - return _C.det_gemm_fwd(a, b) + return det_gemm_linear( + a, + weight, + native_op=_C.det_gemm_fwd_rhs_transposed, + ) # cuBLASLt / CUTLASS: may use split-K. Detach so Autograd.Function owns backward. with torch.no_grad(): - return torch.matmul(a, b) + return torch.nn.functional.linear(a, weight) -def _gemm_fwd_rhs_transposed( - a: Tensor, - bt: Tensor, - *, - disable_split_k: bool, -) -> Tensor: +def _linear_da(grad_output: Tensor, weight: Tensor, *, disable_split_k: bool) -> Tensor: if disable_split_k: - return _C.det_gemm_fwd_rhs_transposed(a, bt) + return det_gemm_linear_input_gradient( + grad_output, + weight, + native_op=_C.det_gemm_fwd, + ) with torch.no_grad(): - return torch.matmul(a, bt.t()) + return torch.matmul(grad_output, weight) -def _gemm_db_transposed( - a: Tensor, - grad_output: Tensor, - *, - disable_split_k: bool, -) -> Tensor: +def _linear_dw(a: Tensor, grad_output: Tensor, *, disable_split_k: bool) -> Tensor: if disable_split_k: - return _C.det_gemm_db_transposed(a, grad_output) + return det_gemm_linear_weight_gradient( + a, + grad_output, + native_op=_C.det_gemm_db_transposed, + ) with torch.no_grad(): - # dW is naturally [out,in]; let the production GEMM create that layout - # instead of materializing [in,out] and transposing the result. - return torch.matmul(grad_output.t(), a) + return torch.matmul(grad_output.t().contiguous(), a) def _require_parallel_group(group: Any, name: str): @@ -85,7 +181,9 @@ def _require_parallel_group(group: Any, name: str): if not dist.is_available(): raise RuntimeError(f"{name}-parallel FFN requires torch.distributed.") if not dist.is_initialized(): - raise RuntimeError(f"{name}-parallel FFN requires an initialized process group.") + raise RuntimeError( + f"{name}-parallel FFN requires an initialized process group." + ) if dist.get_world_size(group=group) <= 1: raise ValueError(f"{name}_group must contain at least two ranks.") return dist @@ -96,6 +194,7 @@ def _validate_ffn_inputs( gate_weight: Tensor, up_weight: Tensor, down_weight: Tensor, + fused_gate_up_weight: Tensor | None, ) -> None: tensors = { "rmsnorm_output": rmsnorm_output, @@ -103,6 +202,8 @@ def _validate_ffn_inputs( "up_weight": up_weight, "down_weight": down_weight, } + if fused_gate_up_weight is not None: + tensors["fused_gate_up_weight"] = fused_gate_up_weight for name, tensor in tensors.items(): if not isinstance(tensor, Tensor): raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor)!r}.") @@ -126,6 +227,8 @@ def _validate_ffn_inputs( "up_weight": (intermediate_size, hidden_size), "down_weight": (hidden_size, intermediate_size), } + if fused_gate_up_weight is not None: + expected_shapes["fused_gate_up_weight"] = (2 * intermediate_size, hidden_size) for name, expected in expected_shapes.items(): actual = tuple(tensors[name].shape) if actual != expected: @@ -135,7 +238,9 @@ def _validate_ffn_inputs( if tensor.dtype != torch.bfloat16: raise TypeError(f"{name} must have dtype bfloat16, got {tensor.dtype}.") if not tensor.is_cuda: - raise RuntimeError(f"{name} must be on a CUDA device, got '{tensor.device}'.") + raise RuntimeError( + f"{name} must be on a CUDA device, got '{tensor.device}'." + ) if tensor.device != rmsnorm_output.device: raise RuntimeError( f"all FFN inputs must be on {rmsnorm_output.device}, " @@ -144,35 +249,27 @@ def _validate_ffn_inputs( def _collective_for_group(group: Any, *, min_size_bytes: int): - if group is None: - return None - - import torch.distributed as dist - - from rl_engine.distributed import DeterministicCollective - - rank = dist.get_rank(group=group) - world_size = dist.get_world_size(group=group) - device_index = torch.cuda.current_device() - key = (id(group), rank, world_size, device_index) - cached = _COLLECTIVES.get(key) - if cached is not None and cached.max_size_bytes >= min_size_bytes: - return cached - if cached is not None: - cached.close() - - collective = DeterministicCollective( + return collective_for_group( group=group, - max_size_bytes=max(_COLLECTIVE_MIN_CAPACITY_BYTES, min_size_bytes), + min_size_bytes=min_size_bytes, + minimum_capacity_bytes=_COLLECTIVE_MIN_CAPACITY_BYTES, ) - _COLLECTIVES[key] = collective - return collective def _all_gather_tokens(tensor: Tensor, collective: Any) -> Tensor: return collective.all_gather(tensor.contiguous()) +def _all_gather_packed_tokens( + *tensors: Tensor, collective: Any +) -> tuple[Tensor, ...]: + """Gather same-row tensors with one handshake and no repacking copies.""" + + return collective.all_gather_many( + tuple(tensor.contiguous() for tensor in tensors) + ) + + def _reduce_scatter_tokens(tensor: Tensor, collective: Any) -> Tensor: world_size = collective.world_size if tensor.size(0) % world_size != 0: @@ -195,6 +292,7 @@ def forward( gate_weight: Tensor, up_weight: Tensor, down_weight: Tensor, + fused_gate_up_weight: Tensor | None, tp_group: Any, cp_group: Any, sequence_parallel: bool, @@ -217,6 +315,15 @@ def forward( up_weight.numel() * element_size, down_weight.numel() * element_size, ) + if cp_group is not None: + packed_width = ( + 2 * rmsnorm_output_2d.size(1) + + 3 * gate_weight.size(0) + ) + min_size_bytes = max( + min_size_bytes, + gemm_tokens * packed_width * element_size, + ) # Create TP before CP so every rank follows the same group order. tp_collective = _collective_for_group(tp_group, min_size_bytes=min_size_bytes) cp_collective = _collective_for_group(cp_group, min_size_bytes=min_size_bytes) @@ -224,58 +331,81 @@ def forward( if sequence_parallel: rmsnorm_output_2d = _all_gather_tokens(rmsnorm_output_2d, tp_collective) - # The SM90 TMA kernel consumes physical B^T=[N,K]. Canonical model - # weights already have exactly that layout, so no prepared copy is - # needed and optimizer updates cannot leave a stale cache behind. - gate = _gemm_fwd_rhs_transposed( - rmsnorm_output_2d, - gate_weight, - disable_split_k=disable_split_k, - ) - up = _gemm_fwd_rhs_transposed( - rmsnorm_output_2d, - up_weight, - disable_split_k=disable_split_k, - ) - activated = _C.swiglu_forward(gate, up) - output = _gemm_fwd_rhs_transposed( - activated, - down_weight, - disable_split_k=disable_split_k, - ) + packed_gate_up = fused_gate_up_weight is not None and disable_split_k + if packed_gate_up: + gate_up = _linear_fwd( + rmsnorm_output_2d, + fused_gate_up_weight, + disable_split_k=True, + ) + activated = _C.swiglu_packed_forward(gate_up) + else: + gate = _linear_fwd( + rmsnorm_output_2d, + gate_weight, + disable_split_k=disable_split_k, + ) + up = _linear_fwd( + rmsnorm_output_2d, + up_weight, + disable_split_k=disable_split_k, + ) + activated = _C.swiglu_forward(gate, up) + output = _linear_fwd(activated, down_weight, disable_split_k=disable_split_k) if sequence_parallel: output = _reduce_scatter_tokens(output, tp_collective) elif tp_collective is not None: output = _all_reduce_inplace(output, tp_collective) - ctx.save_for_backward( - rmsnorm_output_2d, - gate, - up, - activated, - gate_weight, - up_weight, - down_weight, - ) + if packed_gate_up: + ctx.save_for_backward( + rmsnorm_output_2d, + gate_up, + activated, + gate_weight, + up_weight, + down_weight, + ) + else: + ctx.save_for_backward( + rmsnorm_output_2d, + gate, + up, + activated, + gate_weight, + up_weight, + down_weight, + ) ctx.input_shape = input_shape ctx.tp_collective = tp_collective ctx.cp_collective = cp_collective ctx.sequence_parallel = sequence_parallel ctx.disable_split_k = disable_split_k + ctx.packed_gate_up = packed_gate_up return output.reshape(*input_shape[:-1], output.size(-1)) @staticmethod def backward(ctx, grad_output: Tensor): - ( - rmsnorm_output, - gate, - up, - activated, - gate_weight, - up_weight, - down_weight, - ) = ctx.saved_tensors + if ctx.packed_gate_up: + ( + rmsnorm_output, + gate_up, + activated, + gate_weight, + up_weight, + down_weight, + ) = ctx.saved_tensors + else: + ( + rmsnorm_output, + gate, + up, + activated, + gate_weight, + up_weight, + down_weight, + ) = ctx.saved_tensors tp_collective = ctx.tp_collective cp_collective = ctx.cp_collective disable_split_k = ctx.disable_split_k @@ -283,56 +413,74 @@ def backward(ctx, grad_output: Tensor): if ctx.sequence_parallel: grad_output = _all_gather_tokens(grad_output, tp_collective) - # Down weight gradients must see every CP token so the wgrad K-tree - # matches CP=1. Local dW + AllReduce is a different parenthesization - # whenever T is not a complete mid-split tree of 32-wide leaves. + # Down input-gradient shards concatenate across TP; no TP reduction. + grad_activated = _linear_da( + grad_output, + down_weight, + disable_split_k=disable_split_k, + ) + if ctx.packed_gate_up: + grad_gate, grad_up = _C.swiglu_packed_backward(grad_activated, gate_up) + else: + grad_gate, grad_up = _C.swiglu_backward(grad_activated, gate, up) + + # Weight gradients must see every CP token so the GEMM K-tree matches + # CP=1. These payloads become available before any weight-gradient GEMM, + # so one rank-ordered gather preserves the arithmetic contract while + # avoiding four redundant collective handshakes per layer. if cp_collective is not None: - activated_full = _all_gather_tokens(activated, cp_collective) - grad_output_full = _all_gather_tokens(grad_output, cp_collective) - grad_down_weight = _gemm_db_transposed( + ( activated_full, grad_output_full, - disable_split_k=disable_split_k, - ) - else: - grad_down_weight = _gemm_db_transposed( + rmsnorm_full, + grad_gate_full, + grad_up_full, + ) = _all_gather_packed_tokens( activated, grad_output, + rmsnorm_output, + grad_gate, + grad_up, + collective=cp_collective, + ) + grad_down_weight = _linear_dw( + activated_full, + grad_output_full, disable_split_k=disable_split_k, ) - - # Down input-gradient shards concatenate across TP; no TP reduction. - grad_activated = _gemm_fwd(grad_output, down_weight, disable_split_k=disable_split_k) - grad_gate, grad_up = _C.swiglu_backward(grad_activated, gate, up) - - if cp_collective is not None: - rmsnorm_full = _all_gather_tokens(rmsnorm_output, cp_collective) - grad_gate_full = _all_gather_tokens(grad_gate, cp_collective) - grad_up_full = _all_gather_tokens(grad_up, cp_collective) - grad_gate_weight = _gemm_db_transposed( + grad_gate_weight = _linear_dw( rmsnorm_full, grad_gate_full, disable_split_k=disable_split_k, ) - grad_up_weight = _gemm_db_transposed( + grad_up_weight = _linear_dw( rmsnorm_full, grad_up_full, disable_split_k=disable_split_k, ) else: - grad_gate_weight = _gemm_db_transposed( + grad_down_weight = _linear_dw( + activated, + grad_output, + disable_split_k=disable_split_k, + ) + grad_gate_weight = _linear_dw( rmsnorm_output, grad_gate, disable_split_k=disable_split_k, ) - grad_up_weight = _gemm_db_transposed( + grad_up_weight = _linear_dw( rmsnorm_output, grad_up, disable_split_k=disable_split_k, ) # Gate/Up input gradients reduce across TP, then add locally. - grad_rmsnorm_from_gate = _gemm_fwd(grad_gate, gate_weight, disable_split_k=disable_split_k) + grad_rmsnorm_from_gate = _linear_da( + grad_gate, + gate_weight, + disable_split_k=disable_split_k, + ) if ctx.sequence_parallel: grad_rmsnorm_from_gate = _reduce_scatter_tokens( grad_rmsnorm_from_gate, @@ -344,7 +492,11 @@ def backward(ctx, grad_output: Tensor): tp_collective, ) - grad_rmsnorm_from_up = _gemm_fwd(grad_up, up_weight, disable_split_k=disable_split_k) + grad_rmsnorm_from_up = _linear_da( + grad_up, + up_weight, + disable_split_k=disable_split_k, + ) if ctx.sequence_parallel: grad_rmsnorm_from_up = _reduce_scatter_tokens( grad_rmsnorm_from_up, @@ -366,6 +518,7 @@ def backward(ctx, grad_output: Tensor): None, None, None, + None, ) @@ -375,6 +528,7 @@ def qwen3_ffn( up_weight: Tensor, down_weight: Tensor, *, + fused_gate_up_weight: Tensor | None = None, tp_group: Any = None, cp_group: Any = None, sequence_parallel: bool = False, @@ -391,6 +545,11 @@ def qwen3_ffn( ``[I_local, H]``. down_weight: Down projection weight in ``[out, in]`` layout, shape ``[H, I_local]``. + fused_gate_up_weight: Optional existing framework weight in + ``[2 * I_local, H]`` layout. Strict CUDA execution consumes this + with one GEMM launch while returning gradients through the + separate gate/up views, so the framework parameter layout stays + unchanged. tp_group: Optional tensor-parallel process group. Gate and Up are column-parallel; Down is row-parallel. Reductions use the deterministic fixed-tree collectives rather than NCCL. @@ -414,13 +573,23 @@ def qwen3_ffn( if not isinstance(sequence_parallel, bool): raise TypeError("sequence_parallel must be a bool.") deterministic = _resolve_deterministic_mode(deterministic, disable_split_k) - _validate_ffn_inputs(rmsnorm_output, gate_weight, up_weight, down_weight) - _require_ffn_kernels(disable_split_k=deterministic) + _validate_ffn_inputs( + rmsnorm_output, + gate_weight, + up_weight, + down_weight, + fused_gate_up_weight, + ) + _require_ffn_kernels( + disable_split_k=deterministic, + packed_gate_up=fused_gate_up_weight is not None and deterministic, + ) return _DeterministicFFNFunction.apply( rmsnorm_output, gate_weight, up_weight, down_weight, + fused_gate_up_weight, tp_group, cp_group, sequence_parallel, @@ -441,7 +610,9 @@ def _resolve_deterministic_mode( and disable_split_k is not None and deterministic != disable_split_k ): - raise ValueError("deterministic and disable_split_k select conflicting FFN backends.") + raise ValueError( + "deterministic and disable_split_k select conflicting FFN backends." + ) if deterministic is not None: return deterministic if disable_split_k is not None: @@ -456,6 +627,61 @@ class Qwen3FFNOp: is_batch_invariant = True backend_id = BACKEND_ID + def __init__(self) -> None: + # Keep graph-bound IPC resources alive independently of the module-level + # lookup cache. CUDA Graphs retain only the small opaque C++ handle. + self._packed_inference_collectives: dict[int, Any] = {} + + def prepare_packed_inference( + self, + fused_gate_up_weight: Tensor, + down_weight: Tensor, + *, + tp_group: Any, + ) -> tuple[int, int]: + """Create the rollout TP resource before Dynamo captures the model.""" + + if tp_group is None: + return 0, 1 + dist = _require_parallel_group(tp_group, "tensor") + if dist is None: + return 0, 1 + tp_world_size = int(dist.get_world_size(group=tp_group)) + if fused_gate_up_weight.size(0) % 2: + raise ValueError("fused gate/up weight must contain two equal shards") + element_size = fused_gate_up_weight.element_size() + min_size_bytes = ( + max( + fused_gate_up_weight.numel() // 2, + down_weight.numel(), + ) + * element_size + ) + collective = _collective_for_group( + tp_group, + min_size_bytes=min_size_bytes, + ) + collective_handle = int(collective._handle) + self._packed_inference_collectives[collective_handle] = collective + return collective_handle, tp_world_size + + def packed_inference( + self, + rmsnorm_output: Tensor, + fused_gate_up_weight: Tensor, + down_weight: Tensor, + *, + collective_handle: int, + tp_world_size: int, + ) -> Tensor: + return qwen3_ffn_packed_inference( + rmsnorm_output, + fused_gate_up_weight, + down_weight, + collective_handle=collective_handle, + tp_world_size=tp_world_size, + ) + def __call__( self, rmsnorm_output: Tensor, @@ -463,6 +689,7 @@ def __call__( up_weight: Tensor, down_weight: Tensor, *, + fused_gate_up_weight: Tensor | None = None, tp_group: Any = None, cp_group: Any = None, sequence_parallel: bool = False, @@ -474,6 +701,7 @@ def __call__( gate_weight, up_weight, down_weight, + fused_gate_up_weight=fused_gate_up_weight, tp_group=tp_group, cp_group=cp_group, sequence_parallel=sequence_parallel, @@ -488,6 +716,7 @@ def apply( up_weight: Tensor, down_weight: Tensor, *, + fused_gate_up_weight: Tensor | None = None, tp_group: Any = None, cp_group: Any = None, sequence_parallel: bool = False, @@ -499,6 +728,7 @@ def apply( gate_weight, up_weight, down_weight, + fused_gate_up_weight=fused_gate_up_weight, tp_group=tp_group, cp_group=cp_group, sequence_parallel=sequence_parallel, diff --git a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py index 5cea932dd..fecdc8738 100644 --- a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py @@ -8,6 +8,8 @@ import torch +from rl_engine.distributed.collectives import collective_for_group + # Backward token-chunk target: process at most this many ``[chunk, V]`` logit # elements per cuBLAS step so peak backward memory stays ~``chunk*V`` instead of # ``N*V``. @@ -176,6 +178,37 @@ def _validate_tp_vocab_partition_cached( return int(covered) +def _validate_even_tp_vocab_partition_local( + *, + tp_group: Any, + vocab_start_index: int, + local_vocab_size: int, + global_vocab_size: Optional[int], +) -> int: + """Validate the strict equal-shard contract without a data-plane collective.""" + + dist = _require_distributed_initialized() + rank = dist.get_rank(group=tp_group) + world_size = dist.get_world_size(group=tp_group) + expected_global_size = local_vocab_size * world_size + expected_start = rank * local_vocab_size + if local_vocab_size <= 0: + raise ValueError("strict TP linear_logp requires a non-empty vocab shard.") + if vocab_start_index != expected_start: + raise ValueError( + "strict TP linear_logp requires equal rank-ordered vocab shards: " + f"rank={rank} expected vocab_start_index={expected_start}, " + f"got {vocab_start_index}." + ) + if global_vocab_size is not None and int(global_vocab_size) != expected_global_size: + raise ValueError( + "strict TP linear_logp global_vocab_size must equal " + f"local_vocab_size * world_size={expected_global_size}, " + f"got {global_vocab_size}." + ) + return expected_global_size + + def _validate_global_targets( target_1d: torch.Tensor, global_vocab_size: int, @@ -211,6 +244,18 @@ def _validate_global_targets( ) +def _assert_global_targets_async( + target_1d: torch.Tensor, + global_vocab_size: int, +) -> None: + """Range-check targets without synchronizing the CUDA hot path.""" + torch._assert_async( + ((target_1d >= 0) & (target_1d < global_vocab_size)).all(), + f"target_ids out of range: expected [0, {global_vocab_size - 1}]. " + "Mask or filter padding / ignore-index values (e.g. -100) before this op.", + ) + + def _chunked_local_linear_logp_stats( hidden_2d: torch.Tensor, weight: torch.Tensor, @@ -276,17 +321,30 @@ def _merge_tp_local_logp( return local_target_logit - global_lse, global_lse local_stats = torch.stack((local_lse, local_target_logit), dim=0).contiguous() - gathered = torch.empty( - (world_size, *local_stats.shape), - device=local_stats.device, - dtype=local_stats.dtype, - ) - try: - dist.all_gather_into_tensor(gathered, local_stats, group=tp_group) - except (AttributeError, RuntimeError): - chunks = [torch.empty_like(local_stats) for _ in range(world_size)] - dist.all_gather(chunks, local_stats, group=tp_group) - gathered = torch.stack(chunks, dim=0) + if local_stats.is_cuda: + collective = collective_for_group( + tp_group, + min_size_bytes=local_stats.numel() * local_stats.element_size(), + device=local_stats.device, + ) + if collective is None: + raise RuntimeError("tensor-parallel linear_logp requires an RL-Kernel collective") + gathered = collective.all_gather( + local_stats, + validate_signature=False, + ).reshape(world_size, *local_stats.shape) + else: + gathered = torch.empty( + (world_size, *local_stats.shape), + device=local_stats.device, + dtype=local_stats.dtype, + ) + try: + dist.all_gather_into_tensor(gathered, local_stats, group=tp_group) + except (AttributeError, RuntimeError): + chunks = [torch.empty_like(local_stats) for _ in range(world_size)] + dist.all_gather(chunks, local_stats, group=tp_group) + gathered = torch.stack(chunks, dim=0) # [contract] Explicit rank-ordered reduction trees. ``torch.logsumexp`` and # ``Tensor.sum`` own their internal reduction order, which is not part of @@ -310,6 +368,25 @@ def _merge_tp_local_logp( return target_logit - global_lse, global_lse +def _deterministic_tp_all_reduce_(tensor: torch.Tensor, tp_group: Any) -> torch.Tensor: + if not tensor.is_cuda: + dist = _require_distributed_initialized() + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=tp_group) + return tensor + collective = collective_for_group( + tp_group, + min_size_bytes=tensor.numel() * tensor.element_size(), + device=tensor.device, + ) + if collective is None: + raise RuntimeError("tensor-parallel linear_logp requires an RL-Kernel collective") + return collective.all_reduce( + tensor.contiguous(), + out=tensor, + validate_signature=False, + ) + + def tensor_parallel_linear_logp_backward( grad_logp: torch.Tensor, hidden_2d: torch.Tensor, @@ -330,7 +407,6 @@ def tensor_parallel_linear_logp_backward( compute_grad_bias: bool = True, chunk_elems: int = BWD_CHUNK_ELEMS, ): - dist = _require_distributed_initialized() n, d = hidden_2d.shape local_vocab = weight.shape[0] dt = weight.dtype @@ -384,7 +460,7 @@ def tensor_parallel_linear_logp_backward( grad_hidden = None if grad_h is not None: - dist.all_reduce(grad_h, op=dist.ReduceOp.SUM, group=tp_group) + _deterministic_tp_all_reduce_(grad_h, tp_group) grad_hidden = grad_h.to(hidden_dtype).reshape((*tuple(lead_shape), d)) grad_weight = grad_w.to(weight_dtype) if grad_w is not None else None grad_bias = grad_b.to(bias_dtype) if grad_b is not None else None diff --git a/rl_engine/kernels/ops/pytorch/norm/rms_norm.py b/rl_engine/kernels/ops/pytorch/norm/rms_norm.py index f2a87625a..b891a7cfb 100644 --- a/rl_engine/kernels/ops/pytorch/norm/rms_norm.py +++ b/rl_engine/kernels/ops/pytorch/norm/rms_norm.py @@ -6,6 +6,76 @@ import torch +@torch.library.custom_op("rl_kernel::strict_rms_norm", mutates_args=()) +def _strict_rms_norm( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Preserve PyTorch eager RMSNorm arithmetic across graph compilation.""" + + return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight, eps) + + +@_strict_rms_norm.register_fake +def _strict_rms_norm_fake( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + del weight, eps + return torch.empty_like(x) + + +@torch.library.custom_op("rl_kernel::strict_add_rms_norm", mutates_args=()) +def _strict_add_rms_norm( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Preserve vLLM's eager residual-add and RMSNorm contract.""" + + updated_residual = x + residual + normalized = torch.nn.functional.rms_norm( + updated_residual, + (updated_residual.shape[-1],), + weight, + eps, + ) + return normalized, updated_residual + + +@_strict_add_rms_norm.register_fake +def _strict_add_rms_norm_fake( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + del residual, weight, eps + return torch.empty_like(x), torch.empty_like(x) + + +def strict_rms_norm( + x: torch.Tensor, + weight: torch.Tensor, + *, + eps: float, +) -> torch.Tensor: + return _strict_rms_norm(x, weight, eps) + + +def strict_add_rms_norm( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + *, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + return _strict_add_rms_norm(x, residual, weight, eps) + + class NativeRMSNormOp: """ Pure Pytorch native RMSNorm reference diff --git a/rl_engine/runtime_mode.py b/rl_engine/runtime_mode.py new file mode 100644 index 000000000..7db729721 --- /dev/null +++ b/rl_engine/runtime_mode.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""User-facing RL-Kernel runtime policy.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from enum import Enum + + +class RLKernelMode(str, Enum): + STRICT = "strict" + AUDIT = "audit" + AUTO = "auto" + OFF = "off" + + +@dataclass(frozen=True) +class RLKernelModePolicy: + mode: RLKernelMode + enabled: bool + case_id: str + provider_mode: str | None + fail_on_fallback: bool + report_only: bool + + +_POLICIES = { + RLKernelMode.STRICT: RLKernelModePolicy( + RLKernelMode.STRICT, True, "R/R", "strict", True, False + ), + RLKernelMode.AUDIT: RLKernelModePolicy( + RLKernelMode.AUDIT, True, "R/R", "strict", False, True + ), + RLKernelMode.AUTO: RLKernelModePolicy( + RLKernelMode.AUTO, True, "P/P", "auto", False, False + ), + RLKernelMode.OFF: RLKernelModePolicy( + RLKernelMode.OFF, False, "P/P", None, False, False + ), +} + + +def rl_kernel_mode(value: str | RLKernelMode | None = None) -> RLKernelMode: + raw = os.getenv("RL_KERNEL_MODE", "strict") if value is None else value + if isinstance(raw, RLKernelMode): + return raw + normalized = str(raw).strip().lower() + try: + return RLKernelMode(normalized) + except ValueError as exc: + choices = ", ".join(mode.value for mode in RLKernelMode) + raise RuntimeError( + f"RL_KERNEL_MODE must be one of {choices}; got {raw!r}" + ) from exc + + +def rl_kernel_mode_policy( + value: str | RLKernelMode | None = None, +) -> RLKernelModePolicy: + return _POLICIES[rl_kernel_mode(value)] + + +def strict_contract_enabled() -> bool: + configured = os.getenv("RL_KERNEL_MODE") + if configured is not None: + return rl_kernel_mode(configured) in {RLKernelMode.STRICT, RLKernelMode.AUDIT} + legacy = os.getenv("VIME_RL_KERNEL_STRICT", "").strip().lower() + return legacy in {"1", "true", "yes", "on"} + + +def route_report_enabled() -> bool: + enabled = os.getenv("RL_KERNEL_ROUTE_REPORT", "1").strip().lower() + if enabled in {"0", "false", "no", "off"}: + return False + all_ranks = os.getenv("RL_KERNEL_ROUTE_REPORT_ALL_RANKS", "0").strip().lower() + if all_ranks in {"1", "true", "yes", "on"}: + return True + for name in ("RANK", "LOCAL_RANK"): + value = os.getenv(name, "").strip() + if value: + try: + return int(value) == 0 + except ValueError: + continue + return True + + +__all__ = [ + "RLKernelMode", + "RLKernelModePolicy", + "rl_kernel_mode", + "rl_kernel_mode_policy", + "route_report_enabled", + "strict_contract_enabled", +] diff --git a/tests/test_framework_runtime_adapters.py b/tests/test_framework_runtime_adapters.py index 4ae0dcc3e..9787689cd 100644 --- a/tests/test_framework_runtime_adapters.py +++ b/tests/test_framework_runtime_adapters.py @@ -241,15 +241,14 @@ def get(self, tensor, *, topology): ) assert output.shape == (8, 8) - assert len(calls) == 2 + assert len(calls) == 1 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]], ] + assert calls[0]["query_position_ids"].tolist() == [ + [0, 1, 6, 7], + [0, 1, 6, 7], + ] def test_vllm_current_flash_attention_kv_cache_layout_is_materialized(): @@ -465,8 +464,8 @@ def forward_bshd_with_lse(self, q, k, v, *, causal, **kwargs): 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" + assert calls == [(4, 4, True)] + assert result.provenance["query_schedule"] == "full_sequence_causal_single_launch" class _ReadbackOperator: diff --git a/tests/test_qwen_ffn.py b/tests/test_qwen_ffn.py index 515abc7ad..cc70219fd 100644 --- a/tests/test_qwen_ffn.py +++ b/tests/test_qwen_ffn.py @@ -655,7 +655,7 @@ def _cache_worker(rank, world_size, init_method, result_queue): assert len(ffn_module._COLLECTIVES) == 1 grown_collective = next(iter(ffn_module._COLLECTIVES.values())) assert grown_collective is not first_collective - assert first_collective._handle == 0 + assert first_collective._handle == first_handle assert grown_collective.max_size_bytes > first_capacity assert grown_collective._handle != 0 @@ -674,6 +674,7 @@ def _cache_worker(rank, world_size, init_method, result_queue): assert len(ffn_module._COLLECTIVES) == 2 _close_ffn_collectives() + first_collective.close() result_queue.put({"ok": True, "rank": rank}) except Exception: # pragma: no cover - forwarded to the parent process. result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) @@ -1111,7 +1112,7 @@ def test_qwen_ffn_world8_tp4_cp2_match_tp1_cp1_bitwise(): _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_TP4_CP2_CONFIGS,), timeout=120) -def test_qwen_ffn_collective_cache_reuses_grows_closes_and_rebuilds_group(): +def test_qwen_ffn_collective_cache_growth_preserves_borrowers_and_rebuilds_group(): _spawn_nccl_workers(_cache_worker, 2, timeout=120)