Skip to content

Commit 85f357c

Browse files
wlxjhyfwlxjhyf
committed
feat:add triton op and cpp wrapper for copy_ and to_copy
1 parent 00072f1 commit 85f357c

8 files changed

Lines changed: 380 additions & 1 deletion

File tree

‎ctests/CMakeLists.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,3 +98,8 @@ add_executable(test_triton_rwkv_mm_sparsity test_triton_rwkv_mm_sparsity.cpp)
9898
target_link_libraries(test_triton_rwkv_mm_sparsity
9999
PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
100100
add_test(NAME test_triton_rwkv_mm_sparsity COMMAND test_triton_rwkv_mm_sparsity)
101+
102+
add_executable(test_triton_copy test_triton_copy.cpp)
103+
target_link_libraries(test_triton_copy
104+
PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
105+
add_test(NAME test_triton_copy COMMAND test_triton_copy)

‎ctests/test_triton_copy.cpp‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
#include "flag_gems/operators.h"
2+
#include "gtest/gtest.h"
3+
#include "torch/torch.h"
4+
5+
TEST(CopyTest, ContiguousTensorCopy) {
6+
const torch::Device device(torch::kCUDA, 0);
7+
torch::Tensor t = torch::randn({4, 5}, torch::TensorOptions().device(device).dtype(torch::kFloat32));
8+
9+
torch::Tensor out_gems = flag_gems::to_copy(t);
10+
torch::Tensor out_ref = t.clone();
11+
12+
EXPECT_TRUE(torch::allclose(out_gems, out_ref));
13+
EXPECT_EQ(out_gems.dtype(), t.dtype());
14+
}
15+
16+
TEST(CopyTest, ContiguousTensorCopyWithDtype) {
17+
const torch::Device device(torch::kCUDA, 0);
18+
torch::Tensor t = torch::randn({3, 3}, torch::TensorOptions().device(device).dtype(torch::kFloat16));
19+
20+
torch::Tensor out_gems = flag_gems::to_copy(t, torch::kFloat32);
21+
torch::Tensor out_ref = t.to(torch::kFloat32);
22+
23+
EXPECT_TRUE(torch::allclose(out_gems, out_ref));
24+
EXPECT_EQ(out_gems.dtype(), torch::kFloat32);
25+
}
26+
27+
TEST(CopyTest, NonContiguousTensorCopy) {
28+
const torch::Device device(torch::kCUDA, 0);
29+
torch::Tensor t = torch::randn({2, 3, 4}, torch::TensorOptions().device(device));
30+
torch::Tensor t_transposed = t.transpose(0, 1);
31+
32+
torch::Tensor out_gems = flag_gems::to_copy(t_transposed);
33+
torch::Tensor out_ref = t_transposed.clone();
34+
35+
EXPECT_TRUE(torch::allclose(out_gems, out_ref));
36+
}
37+
38+
TEST(CopyTest, CopyInplaceContiguous) {
39+
const torch::Device device(torch::kCUDA, 0);
40+
torch::Tensor src = torch::randn({5, 5}, torch::TensorOptions().device(device));
41+
torch::Tensor dst = torch::empty_like(src);
42+
43+
flag_gems::copy_(dst, src);
44+
45+
EXPECT_TRUE(torch::allclose(dst, src));
46+
}
47+
48+
TEST(CopyTest, CopyInplaceNonContiguous) {
49+
const torch::Device device(torch::kCUDA, 0);
50+
torch::Tensor src = torch::randn({3, 4, 5}, torch::TensorOptions().device(device));
51+
torch::Tensor dst = torch::empty({5, 4, 3}, torch::TensorOptions().device(device));
52+
torch::Tensor src_transposed = src.transpose(0, 2);
53+
54+
flag_gems::copy_(dst, src_transposed);
55+
56+
EXPECT_TRUE(torch::allclose(dst, src_transposed));
57+
}
58+
59+
TEST(CopyTest, CopyBroadcasting) {
60+
const torch::Device device(torch::kCUDA, 0);
61+
torch::Tensor src = torch::randn({1, 5}, torch::TensorOptions().device(device));
62+
torch::Tensor dst = torch::empty({3, 5}, torch::TensorOptions().device(device));
63+
64+
flag_gems::copy_(dst, src);
65+
66+
torch::Tensor expected = src.expand_as(dst);
67+
EXPECT_TRUE(torch::allclose(dst, expected));
68+
}
69+
70+
TEST(CopyTest, EmptyTensor) {
71+
const torch::Device device(torch::kCUDA, 0);
72+
torch::Tensor src = torch::empty({0}, torch::TensorOptions().device(device));
73+
torch::Tensor dst = torch::empty_like(src);
74+
75+
flag_gems::copy_(dst, src);
76+
77+
EXPECT_EQ(dst.numel(), 0);
78+
}

‎include/flag_gems/operators.h‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,4 +226,14 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> rwkv_ka_fusion(const at::Tensor &
226226
int64_t H,
227227
int64_t N);
228228

229+
at::Tensor to_copy(const at::Tensor &self,
230+
c10::optional<at::ScalarType> dtype = c10::nullopt,
231+
c10::optional<at::Layout> layout = c10::nullopt,
232+
c10::optional<at::Device> device = c10::nullopt,
233+
c10::optional<bool> pin_memory = c10::nullopt,
234+
bool non_blocking = false,
235+
c10::optional<at::MemoryFormat> memory_format = c10::nullopt);
236+
237+
at::Tensor &copy_(at::Tensor &dst, const at::Tensor &src, bool non_blocking = false);
238+
229239
} // namespace flag_gems

‎lib/CMakeLists.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@ add_library(operators
2727
reshape_and_cache_flash.cpp
2828
flash_attn_varlen_func.cpp
2929
rwkv_mm_sparsity.cpp
30-
rwkv_ka_fusion.cpp)
30+
rwkv_ka_fusion.cpp
31+
copy.cpp)
3132

3233
if (TRITON_GE_3P5)
3334
target_compile_definitions(operators PRIVATE TRITON_GE_3P5)

‎lib/copy.cpp‎

Lines changed: 232 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,232 @@
1+
#include <c10/core/DispatchKeySet.h>
2+
#include <vector>
3+
#include "c10/cuda/CUDAStream.h"
4+
#include "flag_gems/utils.h"
5+
#include "torch/torch.h"
6+
#include "triton_jit/triton_jit_function.h"
7+
8+
namespace flag_gems {
9+
10+
using namespace triton_jit;
11+
12+
std::vector<int64_t> broadcasted_stride(const std::vector<int64_t>& shape,
13+
const std::vector<int64_t>& stride,
14+
const std::vector<int64_t>& target_shape) {
15+
int ndim_diff = target_shape.size() - shape.size();
16+
TORCH_CHECK(ndim_diff >= 0, "cannot broadcast to fewer dimensions");
17+
18+
std::vector<int64_t> full_shape(ndim_diff, 1);
19+
full_shape.insert(full_shape.end(), shape.begin(), shape.end());
20+
21+
std::vector<int64_t> full_stride(ndim_diff, 0);
22+
full_stride.insert(full_stride.end(), stride.begin(), stride.end());
23+
24+
std::vector<int64_t> out_stride(target_shape.size());
25+
26+
for (size_t i = 0; i < target_shape.size(); ++i) {
27+
if (full_shape[i] == target_shape[i]) {
28+
out_stride[i] = full_stride[i];
29+
} else if (full_shape[i] == 1) {
30+
out_stride[i] = 0;
31+
} else {
32+
TORCH_CHECK(false, "illegal broadcast at dim ", i);
33+
}
34+
}
35+
36+
return out_stride;
37+
}
38+
39+
static bool _can_use_triton_copy(const at::Tensor& dst, const at::Tensor& src, bool non_blocking) {
40+
if (!dst.is_cuda() || !src.is_cuda()) return false;
41+
if (dst.device() != src.device()) return false;
42+
if (non_blocking) return false;
43+
return true;
44+
}
45+
46+
static at::Tensor& redispatch_copy_fallback(at::Tensor& dst, const at::Tensor& src, bool non_blocking) {
47+
static auto op = c10::Dispatcher::singleton()
48+
.findSchemaOrThrow("aten::copy_", "")
49+
.typed<at::Tensor&(at::Tensor&, const at::Tensor&, bool)>();
50+
51+
constexpr c10::DispatchKeySet fallback_keyset =
52+
c10::DispatchKeySet(c10::DispatchKey::CompositeExplicitAutograd);
53+
54+
return op.redispatch(fallback_keyset, dst, src, non_blocking);
55+
}
56+
57+
static at::Tensor redispatch_to_copy_fallback(const at::Tensor& src,
58+
c10::optional<at::ScalarType> dtype,
59+
c10::optional<at::Layout> layout,
60+
c10::optional<at::Device> device,
61+
c10::optional<bool> pin_memory,
62+
bool non_blocking,
63+
c10::optional<at::MemoryFormat> memory_format) {
64+
static auto op = c10::Dispatcher::singleton()
65+
.findSchemaOrThrow("aten::_to_copy", "")
66+
.typed<at::Tensor(const at::Tensor&,
67+
c10::optional<at::ScalarType>,
68+
c10::optional<at::Layout>,
69+
c10::optional<at::Device>,
70+
c10::optional<bool>,
71+
bool,
72+
c10::optional<at::MemoryFormat>)>();
73+
74+
constexpr c10::DispatchKeySet fallback_keyset =
75+
c10::DispatchKeySet(c10::DispatchKey::CompositeExplicitAutograd);
76+
77+
return op.redispatch(fallback_keyset, src, dtype, layout, device, pin_memory, non_blocking, memory_format);
78+
}
79+
80+
at::Tensor to_copy(const at::Tensor& x,
81+
c10::optional<at::ScalarType> dtype = c10::nullopt,
82+
c10::optional<at::Layout> layout = c10::nullopt,
83+
c10::optional<at::Device> device = c10::nullopt,
84+
c10::optional<bool> pin_memory = c10::nullopt,
85+
bool non_blocking = false,
86+
c10::optional<at::MemoryFormat> memory_format = c10::nullopt) {
87+
TORCH_WARN("[flag_gems][to_copy] gems::to_copy");
88+
TORCH_CHECK(x.layout() == at::Layout::Strided, "Only strided tensors are supported");
89+
TORCH_CHECK(!x.is_quantized(), "Quantized tensors are not supported");
90+
if (layout.has_value()) {
91+
TORCH_CHECK(layout.value() == x.layout(), "to_copy: layout conversion is not supported");
92+
}
93+
TORCH_CHECK(!pin_memory.has_value(), "to_copy: pin_memory is not supported");
94+
TORCH_CHECK(!non_blocking, "to_copy: non_blocking copy is not supported");
95+
96+
auto target_dtype = dtype.has_value() ? dtype.value() : x.scalar_type();
97+
auto target_device = device.has_value() ? device.value() : x.device();
98+
auto target_memory_format = memory_format.has_value() ? memory_format.value() : at::MemoryFormat::Preserve;
99+
100+
at::Tensor out =
101+
at::empty_like(x, x.options().dtype(target_dtype).device(target_device), target_memory_format);
102+
103+
// if (!_can_use_triton_copy(out, x, non_blocking)) {
104+
// return redispatch_to_copy_fallback(x, dtype, layout, device, pin_memory, non_blocking, memory_format);
105+
// }
106+
107+
const int64_t numel = x.numel();
108+
if (numel == 0) return out;
109+
110+
constexpr int BLOCK_SIZE = 1024;
111+
const unsigned int grid_x = (numel + BLOCK_SIZE - 1) / BLOCK_SIZE;
112+
113+
c10::DeviceGuard guard(target_device);
114+
c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();
115+
CUstream raw_stream = static_cast<CUstream>(stream.stream());
116+
117+
// at::Tensor x_linear = (x.scalar_type() != target_dtype) ? x.to(target_dtype) : x;
118+
at::Tensor x_linear = x;
119+
if (x.scalar_type() != target_dtype) {
120+
return redispatch_to_copy_fallback(x, dtype, layout, device, pin_memory, non_blocking, memory_format);
121+
}
122+
if (x_linear.is_contiguous() && out.is_contiguous() && numel <= std::numeric_limits<int32_t>::max()) {
123+
const TritonJITFunction& kernel_linear =
124+
TritonJITFunction::get_instance((utils::get_triton_src_path() / "copy.py").string(),
125+
"copy_kernel_linear");
126+
kernel_linear(raw_stream, grid_x, 1, 1, 4, 0, x_linear, out, numel, BLOCK_SIZE);
127+
return out;
128+
}
129+
130+
std::vector<int64_t> task_shape(out.sizes().begin(), out.sizes().end());
131+
int NDIMS = task_shape.size();
132+
133+
std::vector<int64_t> src_stride =
134+
broadcasted_stride(std::vector<int64_t>(x_linear.sizes().begin(), x_linear.sizes().end()),
135+
std::vector<int64_t>(x_linear.strides().begin(), x_linear.strides().end()),
136+
task_shape);
137+
std::vector<int64_t> dst_stride =
138+
broadcasted_stride(std::vector<int64_t>(out.sizes().begin(), out.sizes().end()),
139+
std::vector<int64_t>(out.strides().begin(), out.strides().end()),
140+
task_shape);
141+
142+
const TritonJITFunction& kernel_nd =
143+
TritonJITFunction::get_instance((utils::get_triton_src_path() / "copy.py").string(), "copy_kernel_nd");
144+
kernel_nd(raw_stream,
145+
grid_x,
146+
1,
147+
1,
148+
4,
149+
0,
150+
x_linear,
151+
out,
152+
torch::tensor(task_shape, torch::TensorOptions().dtype(torch::kInt64).device(out.device())),
153+
torch::tensor(src_stride, torch::TensorOptions().dtype(torch::kInt64).device(out.device())),
154+
torch::tensor(dst_stride, torch::TensorOptions().dtype(torch::kInt64).device(out.device())),
155+
numel,
156+
NDIMS,
157+
BLOCK_SIZE);
158+
159+
return out;
160+
}
161+
162+
at::Tensor& copy_(at::Tensor& dst, const at::Tensor& src, bool non_blocking = false) {
163+
TORCH_WARN("[flag_gems][copy_] gems::copy_");
164+
if (!_can_use_triton_copy(dst, src, non_blocking)) {
165+
return redispatch_copy_fallback(dst, src, non_blocking);
166+
}
167+
TORCH_CHECK(!dst._is_zerotensor(), "ZeroTensors are immutable");
168+
if (src._is_zerotensor()) {
169+
dst.zero_();
170+
return dst;
171+
}
172+
173+
if (dst.data_ptr() == src.data_ptr()) return dst;
174+
175+
TORCH_CHECK(src.sizes().size() <= dst.sizes().size(), "src cannot be broadcasted to dst");
176+
for (size_t i = 0; i < src.dim(); ++i) {
177+
TORCH_CHECK(src.size(i) == dst.size(dst.dim() - src.dim() + i) || src.size(i) == 1,
178+
"src cannot be broadcasted to dst");
179+
}
180+
181+
const int64_t numel = dst.numel();
182+
183+
constexpr int BLOCK_SIZE = 1024;
184+
const unsigned int grid_x = (numel + BLOCK_SIZE - 1) / BLOCK_SIZE;
185+
186+
c10::DeviceGuard guard(dst.device());
187+
c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();
188+
CUstream raw_stream = static_cast<CUstream>(stream.stream());
189+
190+
bool no_broadcast = src.sizes().equals(dst.sizes());
191+
192+
if (dst.is_contiguous() && src.is_contiguous() && no_broadcast &&
193+
numel <= std::numeric_limits<int32_t>::max()) {
194+
const TritonJITFunction& kernel_linear =
195+
TritonJITFunction::get_instance((utils::get_triton_src_path() / "copy.py").string(),
196+
"copy_kernel_linear");
197+
kernel_linear(raw_stream, grid_x, 1, 1, 4, 0, src, dst, numel, BLOCK_SIZE);
198+
return dst;
199+
}
200+
201+
std::vector<int64_t> task_shape(dst.sizes().begin(), dst.sizes().end());
202+
int NDIMS = task_shape.size();
203+
204+
std::vector<int64_t> src_stride =
205+
broadcasted_stride(std::vector<int64_t>(src.sizes().begin(), src.sizes().end()),
206+
std::vector<int64_t>(src.strides().begin(), src.strides().end()),
207+
task_shape);
208+
std::vector<int64_t> dst_stride =
209+
broadcasted_stride(std::vector<int64_t>(dst.sizes().begin(), dst.sizes().end()),
210+
std::vector<int64_t>(dst.strides().begin(), dst.strides().end()),
211+
task_shape);
212+
213+
const TritonJITFunction& kernel_nd =
214+
TritonJITFunction::get_instance((utils::get_triton_src_path() / "copy.py").string(), "copy_kernel_nd");
215+
kernel_nd(raw_stream,
216+
grid_x,
217+
1,
218+
1,
219+
4,
220+
0,
221+
src,
222+
dst,
223+
torch::tensor(task_shape, torch::TensorOptions().dtype(torch::kInt64).device(dst.device())),
224+
torch::tensor(src_stride, torch::TensorOptions().dtype(torch::kInt64).device(dst.device())),
225+
torch::tensor(dst_stride, torch::TensorOptions().dtype(torch::kInt64).device(dst.device())),
226+
numel,
227+
NDIMS,
228+
BLOCK_SIZE);
229+
return dst;
230+
}
231+
232+
} // namespace flag_gems

‎src/flag_gems/csrc/aten_patch.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ TORCH_LIBRARY_IMPL(aten, CUDA, m) {
3838
REGISTER_AND_LOG("zeros", zeros);
3939
REGISTER_AND_LOG("fill.Scalar", fill_scalar);
4040
REGISTER_AND_LOG("fill_.Scalar", fill_scalar_);
41+
// REGISTER_AND_LOG("_to_copy", to_copy);
42+
// REGISTER_AND_LOG("copy_", copy_);
4143
}
4244

4345
} // namespace flag_gems

‎src/flag_gems/csrc/cstub.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ PYBIND11_MODULE(c_operators, m) {
3636
m.def("remainder_.Tensor", &flag_gems::remainder_);
3737
m.def("rwkv_mm_sparsity", &flag_gems::rwkv_mm_sparsity);
3838
m.def("rwkv_ka_fusion", &flag_gems::rwkv_ka_fusion);
39+
m.def("copy_", &flag_gems::copy_);
40+
m.def("to_copy", &flag_gems::to_copy);
3941
}
4042
namespace flag_gems {
4143
TORCH_LIBRARY(flag_gems, m) {
@@ -123,6 +125,10 @@ TORCH_LIBRARY(flag_gems, m) {
123125

124126
m.def("rwkv_mm_sparsity(Tensor k, Tensor v) -> Tensor");
125127
m.def("rwkv_ka_fusion(Tensor k, Tensor kk, Tensor a, Tensor ka, int H, int N) -> (Tensor, Tensor, Tensor)");
128+
m.def("copy_(Tensor(a!) dst, Tensor src, bool non_blocking=False) -> Tensor(a!)");
129+
m.def(
130+
"to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? "
131+
"pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor");
126132
}
127133

128134
TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
@@ -195,5 +201,7 @@ TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
195201
m.impl("flash_attn_varlen_func", TORCH_FN(flash_attn_varlen_func));
196202
m.impl("rwkv_mm_sparsity", TORCH_FN(rwkv_mm_sparsity));
197203
m.impl("rwkv_ka_fusion", TORCH_FN(rwkv_ka_fusion));
204+
m.impl("to_copy", TORCH_FN(to_copy));
205+
m.impl("copy_", TORCH_FN(copy_));
198206
}
199207
} // namespace flag_gems

0 commit comments

Comments
 (0)