Skip to content

ops: add selectable SiTU-GLU activation #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mantissagithub wants to merge 4 commits into
base: main
Choose a base branch
from
Open

ops: add selectable SiTU-GLU activation #8

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9025bd4
Add SiTU-GLU activation
Mantissagithub Aug 6, 2026
2f4d24e
Merge upstream main into agent/add-situglu
Mantissagithub Aug 11, 2026
0c15b17
fix custom op schemas for swiglu_limit
Mantissagithub Aug 11, 2026
45addfd
Merge upstream main into add-situglu
Mantissagithub Aug 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,15 @@ MoK exposes 5 hyperparameters that can affect the performance of MoE execution.

You can set these values when creating the `MoKConfig` dataclass, which you pass to all functional-layer functions.

The `activation` field selects the GLU operation. It defaults to `"swiglu"`, preserving the existing behavior. Set it to `"situglu"` to use the bounded SiTU-GLU operation from the [Kimi K3 technical report](https://arxiv.org/abs/2607.24653):

$$
h = \left[4\tanh\left(\frac{g}{4}\right)\sigma(g)\right]
\left[25\tanh\left(\frac{u}{25}\right)\right].
$$

SiTU-GLU is compiled as a separate operation for both BF16 and MXFP8 forward and backward passes; it does not replace the existing SwiGLU operations.

### Workspace

MoK relies on PyTorch symmetric memory to allocate and manage inter-GPU symmetric buffers, or identically sized memory allocations across many GPUs (i.e., all GPUs in an EP group). These buffers serve as the source/destination of token dispatch/combine, along with a few other purposes. We call the entire collection of symmetric memory, along with the other metadata and scratchpads MoK needs, the **workspace**. We provide a data structure and functions for creating and destroying workspaces.
Expand Down
35 changes: 31 additions & 4 deletions csrc/bindings.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("mxfp8_quantize", &mxfp8_quantize::mxfp8_quantize_entrypoint, "",
pybind11::arg("x_bf16"),
pybind11::arg("return_normal"), pybind11::arg("return_transposed"));
m.def("dispatch_mlp_swiglu_combine_fwd_mxfp8", &dispatch_mlp_swiglu_combine_fwd_mxfp8, "",
auto bind_glu_fwd_mxfp8 = [&](const char *name, auto function) {
m.def(name, function, "",
pybind11::arg("x"), pybind11::arg("x_ptrs"),
pybind11::arg("combine_buffer"), pybind11::arg("combine_buffer_ptrs"),
pybind11::arg("w_shared_gate"), pybind11::arg("w_routed_gate"), pybind11::arg("w_routed_gate_sc"),
Expand All @@ -25,7 +26,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
pybind11::arg("num_tokens"), pybind11::arg("tokens_per_expert"),
pybind11::arg("topk"), pybind11::arg("swiglu_limit"),
pybind11::arg("num_comm_sms"), pybind11::arg("macrobatch_size"), pybind11::arg("minibatch_size"));
m.def("dispatch_mlp_swiglu_combine_bwd_mxfp8", &dispatch_mlp_swiglu_combine_bwd_mxfp8, "",
};
bind_glu_fwd_mxfp8("dispatch_mlp_swiglu_combine_fwd_mxfp8",
&dispatch_mlp_glu_combine_fwd_mxfp8<GluActivation::SWIGLU>);
bind_glu_fwd_mxfp8("dispatch_mlp_situglu_combine_fwd_mxfp8",
&dispatch_mlp_glu_combine_fwd_mxfp8<GluActivation::SITUGLU>);

auto bind_glu_bwd_mxfp8 = [&](const char *name, auto function) {
m.def(name, function, "",
pybind11::arg("d_y_buffer"), pybind11::arg("d_y_buffer_ptrs"),
pybind11::arg("d_x_routed_buffer"), pybind11::arg("d_x_routed_buffer_ptrs"),
pybind11::arg("router_weight_buffer"), pybind11::arg("router_weight_buffer_ptrs"),
Expand All @@ -44,7 +52,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
pybind11::arg("num_tokens"), pybind11::arg("tokens_per_expert"),
pybind11::arg("topk"), pybind11::arg("swiglu_limit"),
pybind11::arg("num_comm_sms"), pybind11::arg("macrobatch_size"), pybind11::arg("minibatch_size"));
m.def("dispatch_mlp_swiglu_combine_fwd_bf16", &dispatch_mlp_swiglu_combine_fwd_bf16, "",
};
bind_glu_bwd_mxfp8("dispatch_mlp_swiglu_combine_bwd_mxfp8",
&dispatch_mlp_glu_combine_bwd_mxfp8<GluActivation::SWIGLU>);
bind_glu_bwd_mxfp8("dispatch_mlp_situglu_combine_bwd_mxfp8",
&dispatch_mlp_glu_combine_bwd_mxfp8<GluActivation::SITUGLU>);

auto bind_glu_fwd_bf16 = [&](const char *name, auto function) {
m.def(name, function, "",
pybind11::arg("x"), pybind11::arg("x_ptrs"),
pybind11::arg("combine_buffer"), pybind11::arg("combine_buffer_ptrs"),
pybind11::arg("w_shared_gate"), pybind11::arg("w_routed_gate"),
Expand All @@ -54,7 +69,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
pybind11::arg("num_tokens"), pybind11::arg("tokens_per_expert"),
pybind11::arg("topk"), pybind11::arg("swiglu_limit"),
pybind11::arg("num_comm_sms"), pybind11::arg("macrobatch_size"), pybind11::arg("minibatch_size"));
m.def("dispatch_mlp_swiglu_combine_bwd_bf16", &dispatch_mlp_swiglu_combine_bwd_bf16, "",
};
bind_glu_fwd_bf16("dispatch_mlp_swiglu_combine_fwd_bf16",
&dispatch_mlp_glu_combine_fwd_bf16<GluActivation::SWIGLU>);
bind_glu_fwd_bf16("dispatch_mlp_situglu_combine_fwd_bf16",
&dispatch_mlp_glu_combine_fwd_bf16<GluActivation::SITUGLU>);

auto bind_glu_bwd_bf16 = [&](const char *name, auto function) {
m.def(name, function, "",
pybind11::arg("d_y_buffer"), pybind11::arg("d_y_buffer_ptrs"),
pybind11::arg("d_x_routed_buffer"), pybind11::arg("d_x_routed_buffer_ptrs"),
pybind11::arg("router_weight_buffer"), pybind11::arg("router_weight_buffer_ptrs"),
Expand All @@ -71,6 +93,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
pybind11::arg("num_tokens"), pybind11::arg("tokens_per_expert"),
pybind11::arg("topk"), pybind11::arg("swiglu_limit"),
pybind11::arg("num_comm_sms"), pybind11::arg("macrobatch_size"), pybind11::arg("minibatch_size"));
};
bind_glu_bwd_bf16("dispatch_mlp_swiglu_combine_bwd_bf16",
&dispatch_mlp_glu_combine_bwd_bf16<GluActivation::SWIGLU>);
bind_glu_bwd_bf16("dispatch_mlp_situglu_combine_bwd_bf16",
&dispatch_mlp_glu_combine_bwd_bf16<GluActivation::SITUGLU>);
m.def("fwd_epilogue", &utils::fwd_epilogue, "",
pybind11::arg("y_shared"), pybind11::arg("combine_buffer"), pybind11::arg("topk_weights"));
m.def("bwd_epilogue", &utils::bwd_epilogue, "",
Expand Down
Loading