89 lines
4.0 KiB
C++
89 lines
4.0 KiB
C++
|
// ix_full_bridge_fused_ar.cpp — Bridge to ixformer C++ functions + fused linear+allreduce
|
||
|
|
//
|
||
|
|
// Adds ixformer_linear_allreduce (linear + NCCL allreduce fused in one kernel launch)
|
||
|
|
//
|
||
|
|
// Symbols from _ixformer_torch.cpython-310-x86_64-linux-gnu.so (nm -D verified):
|
||
|
|
// ixformer_torch_ext::silu_and_mul_forward(at::Tensor&, at::Tensor&)
|
||
|
|
// ixformer_torch_ext::rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double)
|
||
|
|
// ixformer_torch_ext::fused_add_rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double, double)
|
||
|
|
// ixformer_torch_ext::ixformer_linear(at::Tensor&, at::Tensor&, optional<Tensor>, optional<Tensor>)
|
||
|
|
// ixformer_torch_ext::ixformer_linear_ex(at::Tensor&, at::Tensor&, optional<Tensor>)
|
||
|
|
// ixformer_torch_ext::ixformer_linear_allreduce(at::Tensor&, at::Tensor&, optional<Tensor>) ← NEW
|
||
|
|
|
||
|
|
#include <torch/extension.h>
|
||
|
|
#include <optional>
|
||
|
|
|
||
|
|
namespace ixformer_torch_ext {
|
||
|
|
|
||
|
|
void silu_and_mul_forward(at::Tensor& input, at::Tensor& output);
|
||
|
|
|
||
|
|
void rms_norm_forward(at::Tensor& input, at::Tensor& weight,
|
||
|
|
at::Tensor& output, double eps);
|
||
|
|
|
||
|
|
void fused_add_rms_norm_forward(at::Tensor& input, at::Tensor& residual,
|
||
|
|
at::Tensor& weight, double eps, double alpha);
|
||
|
|
|
||
|
|
at::Tensor ixformer_linear(at::Tensor& input, at::Tensor& weight,
|
||
|
|
const c10::optional<at::Tensor>& bias,
|
||
|
|
const c10::optional<at::Tensor>& out);
|
||
|
|
|
||
|
|
at::Tensor ixformer_linear_ex(at::Tensor& input, at::Tensor& weight,
|
||
|
|
const c10::optional<at::Tensor>& bias);
|
||
|
|
|
||
|
|
// fused linear + allreduce
|
||
|
|
// Signature from nm -D on _ixformer_torch.cpython-310-x86_64-linux-gnu.so:
|
||
|
|
// 000000000003c410 T ixformer_torch_ext::ixformer_linear_allreduce(at::Tensor&, at::Tensor&, c10::optional<at::Tensor> const&)
|
||
|
|
at::Tensor ixformer_linear_allreduce(at::Tensor& input, at::Tensor& weight,
|
||
|
|
const c10::optional<at::Tensor>& bias);
|
||
|
|
|
||
|
|
} // namespace ixformer_torch_ext
|
||
|
|
|
||
|
|
|
||
|
|
// ============================================================================
|
||
|
|
// Python wrappers
|
||
|
|
// ============================================================================
|
||
|
|
|
||
|
|
torch::Tensor ix_silu_and_mul(torch::Tensor input) {
|
||
|
|
int64_t half_dim = input.size(-1) / 2;
|
||
|
|
auto output = input.new_empty({input.size(0), half_dim});
|
||
|
|
ixformer_torch_ext::silu_and_mul_forward(input, output);
|
||
|
|
return output;
|
||
|
|
}
|
||
|
|
|
||
|
|
void ix_rms_norm(torch::Tensor output, torch::Tensor input,
|
||
|
|
torch::Tensor weight, double eps) {
|
||
|
|
ixformer_torch_ext::rms_norm_forward(input, weight, output, eps);
|
||
|
|
}
|
||
|
|
|
||
|
|
void ix_fused_add_rms_norm(torch::Tensor input, torch::Tensor residual,
|
||
|
|
torch::Tensor weight, double eps) {
|
||
|
|
ixformer_torch_ext::fused_add_rms_norm_forward(input, residual, weight, eps, 1.0);
|
||
|
|
}
|
||
|
|
|
||
|
|
torch::Tensor ix_linear(torch::Tensor input, torch::Tensor weight,
|
||
|
|
const c10::optional<torch::Tensor>& bias) {
|
||
|
|
auto input_2d = input.view({-1, input.size(-1)});
|
||
|
|
int64_t m = input_2d.size(0);
|
||
|
|
if (m <= 1 && !bias.has_value()) {
|
||
|
|
return ixformer_torch_ext::ixformer_linear_ex(input, weight, bias);
|
||
|
|
}
|
||
|
|
return ixformer_torch_ext::ixformer_linear(input, weight, bias,
|
||
|
|
c10::optional<at::Tensor>());
|
||
|
|
}
|
||
|
|
|
||
|
|
// fused linear + allreduce — single kernel launch
|
||
|
|
// Replaces: output = linear(input, weight); output = allreduce(output)
|
||
|
|
torch::Tensor ix_linear_allreduce(torch::Tensor input, torch::Tensor weight,
|
||
|
|
const c10::optional<torch::Tensor>& bias) {
|
||
|
|
return ixformer_torch_ext::ixformer_linear_allreduce(input, weight, bias);
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||
|
|
m.def("silu_and_mul", &ix_silu_and_mul, "Fused SiLU+mul activation");
|
||
|
|
m.def("rms_norm", &ix_rms_norm, "RMSNorm");
|
||
|
|
m.def("fused_add_rms_norm", &ix_fused_add_rms_norm, "Fused residual + RMSNorm");
|
||
|
|
m.def("linear", &ix_linear, "ixformer GEMM (linear/linear_ex auto-dispatch)");
|
||
|
|
m.def("linear_allreduce", &ix_linear_allreduce,
|
||
|
|
"Fused linear + allreduce (saves one kernel launch per RowParallel op)");
|
||
|
|
}
|