Skip to content

Commit 0eae311

Browse files
committed
resolve rebase conflict
Signed-off-by: zhongboz <[email protected]>
1 parent cce3d2f commit 0eae311

File tree

1 file changed

+2
-0
lines changed
  • transformer_engine/pytorch/csrc/extensions

1 file changed

+2
-0
lines changed

transformer_engine/pytorch/csrc/extensions/pybind.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
200200
py::arg("weight"), py::arg("eps"), py::arg("ln_out"), py::arg("quantizer"),
201201
py::arg("otype"), py::arg("sm_margin"), py::arg("zero_centered_gamma"));
202202
m.def("rmsnorm_bwd", &transformer_engine::pytorch::rmsnorm_bwd, "Backward of RMSNorm");
203+
m.def("fused_bulk_alloc_outputs", &transformer_engine::pytorch::fused_bulk_alloc_outputs, "Fused Bulk Alloc Outputs",
204+
py::arg("input_view"), py::arg("m_splits"), py::arg("quantizer_list"));
203205
m.def("fused_multi_quantize", &transformer_engine::pytorch::fused_multi_quantize,
204206
"Fused Multi-tensor Cast + Transpose", py::arg("input_list"), py::arg("output_list"),
205207
py::arg("quantizer_list"), py::arg("otype"));

0 commit comments

Comments
 (0)