diff --git a/vllm_omni/diffusion/distributed/a2a_permute.py b/vllm_omni/diffusion/distributed/a2a_permute.py index 4584fbfd2..d6ea05218 100644 --- a/vllm_omni/diffusion/distributed/a2a_permute.py +++ b/vllm_omni/diffusion/distributed/a2a_permute.py @@ -172,7 +172,9 @@ def ulysses_qkv_fwd(x: torch.Tensor, group_name: str, world_size: int) -> torch. rows = B * s_local symm_in = _get_symm_buffer((rows, p, lc), x.dtype, x.device, group_name) # (B, S_local, H, D) -> (rows, p, lc); H is row-major so column block r = heads [r*Hl:(r+1)*Hl] - symm_in.copy_(x.reshape(rows, p, lc)) + # Fused QKV projections produce row-strided views. Stage complete rows with + # the copy engine instead of materializing the view with TensorIterator. + torch.ops.a2ap.copy_rows(x.view(rows, p, lc), symm_in) out = torch.empty(p, rows, lc, device=x.device, dtype=x.dtype) torch.ops.a2ap.all_to_all_permute(symm_in, out, 1, 0, group_name) # (p, rows, lc) -> (B, S_global, H_local, D), sequence ordered rank-major diff --git a/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu b/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu index ce5dba770..1c218a1cf 100644 --- a/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu +++ b/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu @@ -7,6 +7,7 @@ // cached per group — the PR's keyed-devcomm manager API is not in 2.11). #include +#include #include #include #include @@ -29,6 +30,37 @@ using namespace c10d::symmetric_memory; #error "Build requires NCCL >= 2.28 with symmetric-memory device support" #endif +void copy_rows(const at::Tensor& input, at::Tensor& out) { + TORCH_CHECK(input.is_cuda() && out.is_cuda(), "a2a_permute: copy_rows requires CUDA tensors"); + TORCH_CHECK(input.device() == out.device(), "a2a_permute: copy_rows device mismatch"); + TORCH_CHECK(input.scalar_type() == out.scalar_type(), "a2a_permute: copy_rows dtype mismatch"); + TORCH_CHECK(input.dim() == 3, "a2a_permute: copy_rows input must be 3-D"); + TORCH_CHECK(out.sizes() == input.sizes(), "a2a_permute: copy_rows shape mismatch"); + TORCH_CHECK(out.is_contiguous(), "a2a_permute: copy_rows output must be contiguous"); + TORCH_CHECK( + input.stride(2) == 1 && input.stride(1) == input.size(2), + "a2a_permute: copy_rows requires contiguous rows"); + + const int64_t row_elements = input.size(1) * input.size(2); + TORCH_CHECK(input.stride(0) >= row_elements, "a2a_permute: copy_rows input rows overlap"); + const size_t rows = static_cast(input.size(0)); + const size_t row_bytes = static_cast(row_elements) * input.element_size(); + const size_t input_pitch = + static_cast(input.stride(0) * input.element_size()); + + c10::cuda::CUDAGuard guard(input.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + C10_CUDA_CHECK(cudaMemcpy2DAsync( + out.data_ptr(), + row_bytes, + input.data_ptr(), + input_pitch, + row_bytes, + rows, + cudaMemcpyDeviceToDevice, + stream)); +} + // ---- kernel (verbatim from PR) ---------------------------------------------- __device__ inline void copy_bytes_vec16_aligned( const char* src, char* dst, size_t nbytes, size_t tid, size_t stride) { @@ -227,8 +259,10 @@ void all_to_all_permute( } // namespace TORCH_LIBRARY_FRAGMENT(a2ap, m) { + m.def("copy_rows(Tensor input, Tensor(a!) out) -> ()"); m.def("all_to_all_permute(Tensor input, Tensor(a!) out, int scatter_dim, int gather_dim, str group_name) -> ()"); } TORCH_LIBRARY_IMPL(a2ap, CUDA, m) { + m.impl("copy_rows", TORCH_FN(copy_rows)); m.impl("all_to_all_permute", TORCH_FN(all_to_all_permute)); }