diff --git a/README.md b/README.md index b7aa6ea..cfc5409 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/csrc/bindings.cu b/csrc/bindings.cu index c4a7d85..9f389ba 100644 --- a/csrc/bindings.cu +++ b/csrc/bindings.cu @@ -15,7 +15,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("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_entrypoint, "", + 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"), @@ -25,7 +26,13 @@ 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_entrypoint, "", + }; + bind_glu_fwd_mxfp8("dispatch_mlp_swiglu_combine_fwd_mxfp8", + &dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint); + bind_glu_fwd_mxfp8("dispatch_mlp_situglu_combine_fwd_mxfp8", + &dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint); + 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"), @@ -44,7 +51,13 @@ 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_entrypoint, "", + }; + bind_glu_bwd_mxfp8("dispatch_mlp_swiglu_combine_bwd_mxfp8", + &dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint); + bind_glu_bwd_mxfp8("dispatch_mlp_situglu_combine_bwd_mxfp8", + &dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint); + 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"), @@ -54,7 +67,13 @@ 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_entrypoint, "", + }; + bind_glu_fwd_bf16("dispatch_mlp_swiglu_combine_fwd_bf16", + &dispatch_mlp_swiglu_combine_fwd_bf16_entrypoint); + bind_glu_fwd_bf16("dispatch_mlp_situglu_combine_fwd_bf16", + &dispatch_mlp_swiglu_combine_fwd_bf16_entrypoint); + 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"), @@ -71,6 +90,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_swiglu_combine_bwd_bf16_entrypoint); + bind_glu_bwd_bf16("dispatch_mlp_situglu_combine_bwd_bf16", + &dispatch_mlp_swiglu_combine_bwd_bf16_entrypoint); m.def("fwd_epilogue", &utils::fwd_epilogue::fwd_epilogue_entrypoint, "", pybind11::arg("y_shared"), pybind11::arg("combine_buffer"), pybind11::arg("topk_weights")); m.def("bwd_epilogue", &utils::bwd_epilogue::bwd_epilogue_entrypoint, "", diff --git a/csrc/megakernel/entrypoints.cuh b/csrc/megakernel/entrypoints.cuh index 4b02df8..7bc3edc 100644 --- a/csrc/megakernel/entrypoints.cuh +++ b/csrc/megakernel/entrypoints.cuh @@ -5,6 +5,7 @@ #include "megakernel.cuh" +template static __host__ std::tuple @@ -43,7 +44,7 @@ dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint( switch (num_devices) { case 1: - return dispatch_mlp_swiglu_combiner<1>::dispatch_mlp_swiglu_combine_fwd_mxfp8( + return dispatch_mlp_swiglu_combiner<1, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_mxfp8( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_routed_gate_sc, w_shared_up, w_routed_up, w_routed_up_sc, @@ -51,7 +52,7 @@ dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 4: - return dispatch_mlp_swiglu_combiner<4>::dispatch_mlp_swiglu_combine_fwd_mxfp8( + return dispatch_mlp_swiglu_combiner<4, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_mxfp8( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_routed_gate_sc, w_shared_up, w_routed_up, w_routed_up_sc, @@ -59,7 +60,7 @@ dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 8: - return dispatch_mlp_swiglu_combiner<8>::dispatch_mlp_swiglu_combine_fwd_mxfp8( + return dispatch_mlp_swiglu_combiner<8, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_mxfp8( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_routed_gate_sc, w_shared_up, w_routed_up, w_routed_up_sc, @@ -67,7 +68,7 @@ dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 16: - return dispatch_mlp_swiglu_combiner<16>::dispatch_mlp_swiglu_combine_fwd_mxfp8( + return dispatch_mlp_swiglu_combiner<16, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_mxfp8( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_routed_gate_sc, w_shared_up, w_routed_up, w_routed_up_sc, @@ -75,7 +76,7 @@ dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 32: - return dispatch_mlp_swiglu_combiner<32>::dispatch_mlp_swiglu_combine_fwd_mxfp8( + return dispatch_mlp_swiglu_combiner<32, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_mxfp8( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_routed_gate_sc, w_shared_up, w_routed_up, w_routed_up_sc, @@ -83,7 +84,7 @@ dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 64: - return dispatch_mlp_swiglu_combiner<64>::dispatch_mlp_swiglu_combine_fwd_mxfp8( + return dispatch_mlp_swiglu_combiner<64, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_mxfp8( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_routed_gate_sc, w_shared_up, w_routed_up, w_routed_up_sc, @@ -96,6 +97,7 @@ dispatch_mlp_swiglu_combine_fwd_mxfp8_entrypoint( } } +template static __host__ std::tuple dispatch_mlp_swiglu_combine_fwd_bf16_entrypoint( @@ -122,37 +124,37 @@ dispatch_mlp_swiglu_combine_fwd_bf16_entrypoint( const int num_devices = static_cast(x_ptrs.size()); switch (num_devices) { case 1: - return dispatch_mlp_swiglu_combiner<1, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_fwd_bf16( + return dispatch_mlp_swiglu_combiner<1, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_bf16( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 4: - return dispatch_mlp_swiglu_combiner<4, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_fwd_bf16( + return dispatch_mlp_swiglu_combiner<4, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_bf16( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 8: - return dispatch_mlp_swiglu_combiner<8, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_fwd_bf16( + return dispatch_mlp_swiglu_combiner<8, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_bf16( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 16: - return dispatch_mlp_swiglu_combiner<16, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_fwd_bf16( + return dispatch_mlp_swiglu_combiner<16, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_bf16( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 32: - return dispatch_mlp_swiglu_combiner<32, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_fwd_bf16( + return dispatch_mlp_swiglu_combiner<32, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_bf16( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 64: - return dispatch_mlp_swiglu_combiner<64, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_fwd_bf16( + return dispatch_mlp_swiglu_combiner<64, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_fwd_bf16( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, @@ -163,6 +165,7 @@ dispatch_mlp_swiglu_combine_fwd_bf16_entrypoint( } } +template static __host__ std::tuple @@ -226,7 +229,7 @@ dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint( switch (num_devices) { case 1: - return dispatch_mlp_swiglu_combiner<1>::dispatch_mlp_swiglu_combine_bwd_mxfp8( + return dispatch_mlp_swiglu_combiner<1, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_mxfp8( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate_T, w_routed_gate_T_sc, @@ -242,7 +245,7 @@ dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 4: - return dispatch_mlp_swiglu_combiner<4>::dispatch_mlp_swiglu_combine_bwd_mxfp8( + return dispatch_mlp_swiglu_combiner<4, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_mxfp8( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate_T, w_routed_gate_T_sc, @@ -258,7 +261,7 @@ dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 8: - return dispatch_mlp_swiglu_combiner<8>::dispatch_mlp_swiglu_combine_bwd_mxfp8( + return dispatch_mlp_swiglu_combiner<8, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_mxfp8( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate_T, w_routed_gate_T_sc, @@ -274,7 +277,7 @@ dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 16: - return dispatch_mlp_swiglu_combiner<16>::dispatch_mlp_swiglu_combine_bwd_mxfp8( + return dispatch_mlp_swiglu_combiner<16, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_mxfp8( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate_T, w_routed_gate_T_sc, @@ -290,7 +293,7 @@ dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 32: - return dispatch_mlp_swiglu_combiner<32>::dispatch_mlp_swiglu_combine_bwd_mxfp8( + return dispatch_mlp_swiglu_combiner<32, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_mxfp8( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate_T, w_routed_gate_T_sc, @@ -306,7 +309,7 @@ dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 64: - return dispatch_mlp_swiglu_combiner<64>::dispatch_mlp_swiglu_combine_bwd_mxfp8( + return dispatch_mlp_swiglu_combiner<64, RoutedPrecision::MXFP8, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_mxfp8( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate_T, w_routed_gate_T_sc, @@ -327,6 +330,7 @@ dispatch_mlp_swiglu_combine_bwd_mxfp8_entrypoint( } } +template static __host__ std::tuple @@ -367,7 +371,7 @@ dispatch_mlp_swiglu_combine_bwd_bf16_entrypoint( const int num_devices = static_cast(x_ptrs.size()); switch (num_devices) { case 1: - return dispatch_mlp_swiglu_combiner<1, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_bwd_bf16( + return dispatch_mlp_swiglu_combiner<1, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_bf16( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, @@ -375,7 +379,7 @@ dispatch_mlp_swiglu_combine_bwd_bf16_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 4: - return dispatch_mlp_swiglu_combiner<4, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_bwd_bf16( + return dispatch_mlp_swiglu_combiner<4, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_bf16( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, @@ -383,7 +387,7 @@ dispatch_mlp_swiglu_combine_bwd_bf16_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 8: - return dispatch_mlp_swiglu_combiner<8, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_bwd_bf16( + return dispatch_mlp_swiglu_combiner<8, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_bf16( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, @@ -391,7 +395,7 @@ dispatch_mlp_swiglu_combine_bwd_bf16_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 16: - return dispatch_mlp_swiglu_combiner<16, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_bwd_bf16( + return dispatch_mlp_swiglu_combiner<16, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_bf16( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, @@ -399,7 +403,7 @@ dispatch_mlp_swiglu_combine_bwd_bf16_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 32: - return dispatch_mlp_swiglu_combiner<32, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_bwd_bf16( + return dispatch_mlp_swiglu_combiner<32, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_bf16( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, @@ -407,7 +411,7 @@ dispatch_mlp_swiglu_combine_bwd_bf16_entrypoint( schedule_peer_rank, schedule_peer_token_idx, num_tokens, tokens_per_expert, topk, swiglu_limit, num_comm_sms, macrobatch_size, minibatch_size); case 64: - return dispatch_mlp_swiglu_combiner<64, RoutedPrecision::BF16>::dispatch_mlp_swiglu_combine_bwd_bf16( + return dispatch_mlp_swiglu_combiner<64, RoutedPrecision::BF16, ACTIVATION>::dispatch_mlp_swiglu_combine_bwd_bf16( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, diff --git a/csrc/megakernel/megakernel.cuh b/csrc/megakernel/megakernel.cuh index 94eefc6..eb65be0 100644 --- a/csrc/megakernel/megakernel.cuh +++ b/csrc/megakernel/megakernel.cuh @@ -5,6 +5,7 @@ #include "../mxfp8.cuh" #include "../utils.cuh" +#include "situglu.cuh" #include #include @@ -17,10 +18,13 @@ enum class RoutedPrecision { MXFP8, }; -template +template struct dispatch_mlp_swiglu_combiner { static constexpr bool USE_MXFP8 = ROUTED_PRECISION == RoutedPrecision::MXFP8; +static constexpr bool USE_SITUGLU = ACTIVATION == GluActivation::SITUGLU; #include "types.cuh" diff --git a/csrc/megakernel/situglu.cuh b/csrc/megakernel/situglu.cuh new file mode 100644 index 0000000..2873675 --- /dev/null +++ b/csrc/megakernel/situglu.cuh @@ -0,0 +1,43 @@ +#pragma once + +enum class GluActivation { + SWIGLU, + SITUGLU, +}; + +// beta1 and beta2 are the constants used by situ-glu in the kimi k3 technical report. +static constexpr float SITUGLU_GATE_BETA = 4.0f; +static constexpr float SITUGLU_UP_BETA = 25.0f; + +struct situglu_backward_terms { + float hidden; + float d_gate; + float d_up; +}; + +static __device__ __forceinline__ float situglu_sigmoid(float x) { + return 1.0f / (1.0f + __expf(-x)); +} + +static __device__ __forceinline__ float situglu_forward(float gate, float up) { + // h = [beta1 * tanh(g / beta1) * sigmoid(g)] * [beta2 * tanh(u / beta2)]. + const float gate_factor = SITUGLU_GATE_BETA * tanhf(gate / SITUGLU_GATE_BETA) * situglu_sigmoid(gate); + const float up_factor = SITUGLU_UP_BETA * tanhf(up / SITUGLU_UP_BETA); + return gate_factor * up_factor; +} + +static __device__ __forceinline__ situglu_backward_terms situglu_backward(float gate, float up) { + const float gate_tanh = tanhf(gate / SITUGLU_GATE_BETA); + const float up_tanh = tanhf(up / SITUGLU_UP_BETA); + const float sigmoid = situglu_sigmoid(gate); + const float gate_factor = SITUGLU_GATE_BETA * gate_tanh * sigmoid; + const float up_factor = SITUGLU_UP_BETA * up_tanh; + + // d gate_factor / dg = (1 - tanh(g / beta1)^2) * sigmoid(g) + // + beta1 * tanh(g / beta1) * sigmoid(g) * (1 - sigmoid(g)). + const float d_gate_factor = (1.0f - gate_tanh * gate_tanh) * sigmoid + + gate_factor * (1.0f - sigmoid); + // d up_factor / du = 1 - tanh(u / beta2)^2. + const float d_up_factor = 1.0f - up_tanh * up_tanh; + return {gate_factor * up_factor, d_gate_factor * up_factor, gate_factor * d_up_factor}; +} diff --git a/csrc/megakernel/swiglu.cuh b/csrc/megakernel/swiglu.cuh index d022a4c..d68ed35 100644 --- a/csrc/megakernel/swiglu.cuh +++ b/csrc/megakernel/swiglu.cuh @@ -129,22 +129,38 @@ static __device__ __forceinline__ void swiglu_fwd_kernel( } if constexpr (!USE_ROUTED_MXFP8) { - rt_fl gate, up, denominator; - compute_group::load(gate, gate_smem[stage]); - compute_group::load(up, up_smem[stage]); - if constexpr (IS_CLAMPED) { - compute_group::min(gate, gate, swiglu_limit); - compute_group::max(up, up, -swiglu_limit); - compute_group::min(up, up, swiglu_limit); + if constexpr (USE_SITUGLU) { + if (threadIdx.x == 0) tma::store_async_read_wait(); + __syncthreads(); + const auto *gate_pairs = reinterpret_cast(gate_smem[stage].data); + const auto *up_pairs = reinterpret_cast(up_smem[stage].data); + auto *hidden_pairs = reinterpret_cast(hidden_smem.data); + #pragma unroll + for (int i = threadIdx.x; i < config::SWIGLU_Mb * config::SWIGLU_Nb / 2; i += config::NUM_THREADS) { + const float2 gate = __bfloat1622float2(gate_pairs[i]); + const float2 up = __bfloat1622float2(up_pairs[i]); + hidden_pairs[i] = __floats2bfloat162_rn( + situglu_forward(gate.x, up.x), + situglu_forward(gate.y, up.y)); + } + } else { + rt_fl gate, up, denominator; + compute_group::load(gate, gate_smem[stage]); + compute_group::load(up, up_smem[stage]); + if constexpr (IS_CLAMPED) { + compute_group::min(gate, gate, swiglu_limit); + compute_group::max(up, up, -swiglu_limit); + compute_group::min(up, up, swiglu_limit); + } + compute_group::mul(denominator, gate, -1.0f); + compute_group::exp(denominator, denominator); + compute_group::add(denominator, denominator, 1.0f); + compute_group::div(gate, gate, denominator); + compute_group::mul(gate, gate, up); + if (threadIdx.x == 0) tma::store_async_read_wait(); + __syncthreads(); + compute_group::store(hidden_smem, gate); } - compute_group::mul(denominator, gate, -1.0f); - compute_group::exp(denominator, denominator); - compute_group::add(denominator, denominator, 1.0f); - compute_group::div(gate, gate, denominator); - compute_group::mul(gate, gate, up); - if (threadIdx.x == 0) tma::store_async_read_wait(); - __syncthreads(); - compute_group::store(hidden_smem, gate); } else { const auto *gate_pairs = reinterpret_cast(gate_smem[stage].data); const auto *up_pairs = reinterpret_cast(up_smem[stage].data); @@ -153,15 +169,19 @@ static __device__ __forceinline__ void swiglu_fwd_kernel( for (int i = threadIdx.x; i < config::SWIGLU_Mb * config::SWIGLU_Nb / 2; i += config::NUM_THREADS) { float2 gate = __bfloat1622float2(gate_pairs[i]); float2 up = __bfloat1622float2(up_pairs[i]); - if constexpr (IS_CLAMPED) { - gate = {fminf(gate.x, swiglu_limit), fminf(gate.y, swiglu_limit)}; - up = {fminf(fmaxf(up.x, -swiglu_limit), swiglu_limit), fminf(fmaxf(up.y, -swiglu_limit), swiglu_limit)}; + if constexpr (USE_SITUGLU) { + gate = float2{situglu_forward(gate.x, up.x), situglu_forward(gate.y, up.y)}; + } else { + if constexpr (IS_CLAMPED) { + gate = {fminf(gate.x, swiglu_limit), fminf(gate.y, swiglu_limit)}; + up = {fminf(fmaxf(up.x, -swiglu_limit), swiglu_limit), fminf(fmaxf(up.y, -swiglu_limit), swiglu_limit)}; + } + float2 denominator = base_ops::mul::op(gate, float2{-1.0f, -1.0f}); + denominator = base_ops::exp::op(denominator); + denominator = base_ops::sum::op(denominator, float2{1.0f, 1.0f}); + gate = base_ops::div::op(gate, denominator); + gate = base_ops::mul::op(gate, up); } - float2 denominator = base_ops::mul::op(gate, float2{-1.0f, -1.0f}); - denominator = base_ops::exp::op(denominator); - denominator = base_ops::sum::op(denominator, float2{1.0f, 1.0f}); - gate = base_ops::div::op(gate, denominator); - gate = base_ops::mul::op(gate, up); hidden_pairs[i] = __floats2bfloat162_rn(gate.x, gate.y); } } @@ -383,12 +403,28 @@ static __device__ __forceinline__ void swiglu_bwd_kernel( mxfp8::dequantize_single_block(gate_fp8, (static_cast(gate_scale_word) >> (k_block_idx * 8)) & 0xFF, gate_fp32); mxfp8::dequantize_single_block(up_fp8, (static_cast(up_scale_word) >> (k_block_idx * 8)) & 0xFF, up_fp32); - // Apply SwiGLU backward + // Apply the selected GLU derivative in fp32 before requantizing the gradients bf16_2 d_gate_bf16[16], d_up_bf16[16]; #pragma unroll for (int k = 0; k < 16; ++k) { - const float2 hidden = swiglu_bwd_pair(gate_fp32[k], up_fp32[k], d_hidden_fp32[k], swiglu_limit, d_gate_bf16[k], d_up_bf16[k]); - router_grad_partial += d_hidden_fp32[k].x * inv_router_weight * hidden.x + d_hidden_fp32[k].y * inv_router_weight * hidden.y; + const float d_hidden_x = d_hidden_fp32[k].x; + const float d_hidden_y = d_hidden_fp32[k].y; + if constexpr (USE_SITUGLU) { + const situglu_backward_terms terms_x = situglu_backward(gate_fp32[k].x, up_fp32[k].x); + const situglu_backward_terms terms_y = situglu_backward(gate_fp32[k].y, up_fp32[k].y); + router_grad_partial += d_hidden_x * inv_router_weight * terms_x.hidden + + d_hidden_y * inv_router_weight * terms_y.hidden; + d_gate_bf16[k] = __floats2bfloat162_rn( + terms_x.d_gate * d_hidden_x, + terms_y.d_gate * d_hidden_y); + d_up_bf16[k] = __floats2bfloat162_rn( + terms_x.d_up * d_hidden_x, + terms_y.d_up * d_hidden_y); + } else { + const float2 hidden = swiglu_bwd_pair(gate_fp32[k], up_fp32[k], d_hidden_fp32[k], swiglu_limit, d_gate_bf16[k], d_up_bf16[k]); + router_grad_partial += d_hidden_x * inv_router_weight * hidden.x + + d_hidden_y * inv_router_weight * hidden.y; + } } // Quantize both gradients; also stage them in BF16 for the transpose-quantize below @@ -447,7 +483,25 @@ static __device__ __forceinline__ void swiglu_bwd_kernel( } __syncthreads(); } else if constexpr (IS_SHARED) { - if constexpr (IS_CLAMPED) { + if constexpr (USE_SITUGLU) { + const auto *d_hidden_pairs = reinterpret_cast(d_hidden_smem[stage].data); + auto *gate_pairs = reinterpret_cast(gate_smem[stage].data); + auto *up_pairs = reinterpret_cast(up_smem[stage].data); + #pragma unroll + for (int i = threadIdx.x; i < config::SWIGLU_Mb * config::SWIGLU_Nb / 2; i += config::NUM_THREADS) { + const float2 gate = __bfloat1622float2(gate_pairs[i]); + const float2 up = __bfloat1622float2(up_pairs[i]); + const float2 d_hidden = __bfloat1622float2(d_hidden_pairs[i]); + const situglu_backward_terms terms_x = situglu_backward(gate.x, up.x); + const situglu_backward_terms terms_y = situglu_backward(gate.y, up.y); + gate_pairs[i] = __floats2bfloat162_rn( + terms_x.d_gate * d_hidden.x, + terms_y.d_gate * d_hidden.y); + up_pairs[i] = __floats2bfloat162_rn( + terms_x.d_up * d_hidden.x, + terms_y.d_up * d_hidden.y); + } + } else if constexpr (IS_CLAMPED) { rt_bf gate, up, d_hidden; compute_group::load(gate, gate_smem[stage]); compute_group::load(up, up_smem[stage]); @@ -520,8 +574,22 @@ static __device__ __forceinline__ void swiglu_bwd_kernel( const float2 gate = __bfloat1622float2(gate_pairs[j]); const float2 up = __bfloat1622float2(up_pairs[j]); const float2 d_hidden = __bfloat1622float2(d_hidden_pairs[j]); - const float2 hidden = swiglu_bwd_pair(gate, up, d_hidden, swiglu_limit, d_gate_pairs[j], d_up_pairs[j]); - router_grad_partial += d_hidden.x * inv_router_weight * hidden.x + d_hidden.y * inv_router_weight * hidden.y; + if constexpr (USE_SITUGLU) { + const situglu_backward_terms terms_x = situglu_backward(gate.x, up.x); + const situglu_backward_terms terms_y = situglu_backward(gate.y, up.y); + router_grad_partial += d_hidden.x * inv_router_weight * terms_x.hidden + + d_hidden.y * inv_router_weight * terms_y.hidden; + d_gate_pairs[j] = __floats2bfloat162_rn( + terms_x.d_gate * d_hidden.x, + terms_y.d_gate * d_hidden.y); + d_up_pairs[j] = __floats2bfloat162_rn( + terms_x.d_up * d_hidden.x, + terms_y.d_up * d_hidden.y); + } else { + const float2 hidden = swiglu_bwd_pair(gate, up, d_hidden, swiglu_limit, d_gate_pairs[j], d_up_pairs[j]); + router_grad_partial += d_hidden.x * inv_router_weight * hidden.x + + d_hidden.y * inv_router_weight * hidden.y; + } } const auto *d_gate_words = reinterpret_cast(d_gate_pairs); const auto *d_up_words = reinterpret_cast(d_up_pairs); diff --git a/mok/_fake_impls.py b/mok/_fake_impls.py index a3450a5..1e9482b 100644 --- a/mok/_fake_impls.py +++ b/mok/_fake_impls.py @@ -381,6 +381,21 @@ def _dispatch_mlp_swiglu_combine_bwd_bf16_fake( ) +# situ-glu has the same saved-tensor and gradient layouts as swiglu. +torch.library.register_fake("mok::dispatch_mlp_situglu_combine_fwd_mxfp8")( + _dispatch_mlp_swiglu_combine_fwd_mxfp8_fake +) +torch.library.register_fake("mok::dispatch_mlp_situglu_combine_fwd_bf16")( + _dispatch_mlp_swiglu_combine_fwd_bf16_fake +) +torch.library.register_fake("mok::dispatch_mlp_situglu_combine_bwd_mxfp8")( + _dispatch_mlp_swiglu_combine_bwd_mxfp8_fake +) +torch.library.register_fake("mok::dispatch_mlp_situglu_combine_bwd_bf16")( + _dispatch_mlp_swiglu_combine_bwd_bf16_fake +) + + @torch.library.register_fake("mok::fwd_epilogue") def _fwd_epilogue_fake( y_shared: torch.Tensor, diff --git a/mok/functional.py b/mok/functional.py index 7160a9a..f4d3509 100644 --- a/mok/functional.py +++ b/mok/functional.py @@ -1,6 +1,6 @@ import math from dataclasses import dataclass -from typing import Any +from typing import Any, Literal import torch import torch.distributed as dist @@ -16,6 +16,10 @@ recompute_forward_context_bf16, dispatch_mlp_swiglu_combine_fwd_mxfp8, dispatch_mlp_swiglu_combine_fwd_bf16, + dispatch_mlp_situglu_combine_bwd_mxfp8, + dispatch_mlp_situglu_combine_bwd_bf16, + dispatch_mlp_situglu_combine_fwd_mxfp8, + dispatch_mlp_situglu_combine_fwd_bf16, fwd_epilogue, schedule, ) @@ -29,6 +33,7 @@ class MoKConfig: macrobatch_size: int = 131072 schedule_capacity_multiplier: float = 0.5 all_gather_top_experts_chunk_bytes: int = 2048 + activation: Literal["swiglu", "situglu"] = "swiglu" @dataclass(frozen=True, slots=True) @@ -48,6 +53,7 @@ class MoKForwardContext: up_routed: torch.Tensor | tuple[torch.Tensor, torch.Tensor] hidden_shared: torch.Tensor hidden_routed: torch.Tensor | tuple[torch.Tensor, torch.Tensor] + activation: Literal["swiglu", "situglu"] = "swiglu" @dataclass(slots=True) @@ -368,6 +374,8 @@ def build_schedule( raise TypeError("workspace must be a MoKWorkspace") if not isinstance(config, MoKConfig): raise TypeError("config must be a MoKConfig") + if config.activation not in ("swiglu", "situglu"): + raise ValueError("activation must be 'swiglu' or 'situglu'") device_properties = torch.cuda.get_device_properties(workspace.device) if type(config.fwd_num_comm_sms) is not int or config.fwd_num_comm_sms <= 0: raise ValueError("fwd_num_comm_sms must be a positive integer") @@ -448,6 +456,8 @@ def validate_inputs( """ if not isinstance(config, MoKConfig): raise TypeError("config must be a MoKConfig") + if config.activation not in ("swiglu", "situglu"): + raise ValueError("activation must be 'swiglu' or 'situglu'") if not isinstance(workspace, MoKWorkspace): raise TypeError("workspace must be a MoKWorkspace") if not isinstance(schedule, MoKSchedule): @@ -505,6 +515,16 @@ def forward( forward_context: MoKForwardContext """ validate_inputs(config, workspace, schedule, x, router_weights) + fwd_mxfp8_op = ( + dispatch_mlp_situglu_combine_fwd_mxfp8 + if config.activation == "situglu" + else dispatch_mlp_swiglu_combine_fwd_mxfp8 + ) + fwd_bf16_op = ( + dispatch_mlp_situglu_combine_fwd_bf16 + if config.activation == "situglu" + else dispatch_mlp_swiglu_combine_fwd_bf16 + ) workspace.x_buffer.copy_(x) # TODO: we can remove this workspace.router_weight_buffer.copy_(router_weights) @@ -519,7 +539,7 @@ def forward( gate_shared, gate_fp8_routed, gate_sc_routed, up_shared, up_fp8_routed, up_sc_routed, hidden_shared, hidden_fp8_t_routed, hidden_sc_t_routed, - y_shared, y_routed) = dispatch_mlp_swiglu_combine_fwd_mxfp8( + y_shared, y_routed) = fwd_mxfp8_op( workspace.x_buffer, workspace.x_buffer_ptrs, workspace.combine_buffer, workspace.combine_buffer_ptrs, shared_gate_weights, routed_gate_weights_fp8, routed_gate_weights_sc, @@ -538,10 +558,11 @@ def forward( up_routed=(up_fp8_routed, up_sc_routed), hidden_shared=hidden_shared, hidden_routed=(hidden_fp8_t_routed, hidden_sc_t_routed), + activation=config.activation, ) else: (x_routed, gate_shared, gate_routed, up_shared, up_routed, - hidden_shared, hidden_routed, y_shared, y_routed) = dispatch_mlp_swiglu_combine_fwd_bf16( + hidden_shared, hidden_routed, y_shared, y_routed) = fwd_bf16_op( workspace.x_buffer, workspace.x_buffer_ptrs, workspace.combine_buffer, workspace.combine_buffer_ptrs, shared_gate_weights, routed_gate_weights, @@ -560,6 +581,7 @@ def forward( up_routed=up_routed, hidden_shared=hidden_shared, hidden_routed=hidden_routed, + activation=config.activation, ) barrier_all(workspace.barrier_buffer, workspace.barrier_buffer_ptrs, @@ -596,6 +618,8 @@ def recompute_forward_context( forward_context: MoKForwardContext """ validate_inputs(config, workspace, schedule, x) + if config.activation != "swiglu": + raise ValueError("recompute_forward_context only supports the 'swiglu' activation") if isinstance(routed_gate_weights, tuple) != isinstance(routed_up_weights, tuple): raise TypeError("routed gate and up weights must use the same precision representation") if isinstance(routed_gate_weights, tuple) and (len(routed_gate_weights) != 2 or len(routed_up_weights) != 2): @@ -646,6 +670,7 @@ def recompute_forward_context( up_routed=up_routed, hidden_shared=hidden_shared, hidden_routed=hidden_routed, + activation=config.activation, ) @@ -705,6 +730,18 @@ def backward( validate_inputs(config, workspace, schedule, x, router_weights, grad_output) if not isinstance(forward_context, MoKForwardContext): raise TypeError("forward_context must be a MoKForwardContext") + if forward_context.activation != config.activation: + raise ValueError("forward_context activation does not match config activation") + bwd_mxfp8_op = ( + dispatch_mlp_situglu_combine_bwd_mxfp8 + if config.activation == "situglu" + else dispatch_mlp_swiglu_combine_bwd_mxfp8 + ) + bwd_bf16_op = ( + dispatch_mlp_situglu_combine_bwd_bf16 + if config.activation == "situglu" + else dispatch_mlp_swiglu_combine_bwd_bf16 + ) workspace.d_y_buffer.copy_(grad_output) # TODO: we can remove this workspace.x_buffer.copy_(x) # TODO: we can remove this @@ -726,7 +763,7 @@ def backward( d_up_shared, d_up_fp8_routed, d_up_sc_routed, d_hidden_shared, d_hidden_routed, d_y_fp8_routed, d_y_sc_routed, d_w_shared_gate, d_w_routed_gate, d_w_shared_up, d_w_routed_up, - d_w_shared_down, d_w_routed_down) = dispatch_mlp_swiglu_combine_bwd_mxfp8( + d_w_shared_down, d_w_routed_down) = bwd_mxfp8_op( workspace.d_y_buffer, workspace.d_y_buffer_ptrs, workspace.d_x_routed_buffer, workspace.d_x_routed_buffer_ptrs, workspace.router_weight_buffer, workspace.router_weight_buffer_ptrs, @@ -754,7 +791,7 @@ def backward( (d_x_shared, d_x_routed, d_gate_shared, d_gate_routed, d_up_shared, d_up_routed, d_hidden_shared, d_hidden_routed, d_y_routed, d_w_shared_gate, d_w_routed_gate, d_w_shared_up, d_w_routed_up, - d_w_shared_down, d_w_routed_down) = dispatch_mlp_swiglu_combine_bwd_bf16( + d_w_shared_down, d_w_routed_down) = bwd_bf16_op( workspace.d_y_buffer, workspace.d_y_buffer_ptrs, workspace.d_x_routed_buffer, workspace.d_x_routed_buffer_ptrs, workspace.router_weight_buffer, workspace.router_weight_buffer_ptrs, diff --git a/mok/ops.py b/mok/ops.py index 79fa836..f9c7a4b 100644 --- a/mok/ops.py +++ b/mok/ops.py @@ -197,8 +197,8 @@ def mxfp8_quantize( return _C.mxfp8_quantize(x_bf16, return_normal, return_transposed) -@torch.library.custom_op("mok::dispatch_mlp_swiglu_combine_fwd_mxfp8", mutates_args=("combine_buffer",)) -def dispatch_mlp_swiglu_combine_fwd_mxfp8( +def _dispatch_mlp_glu_combine_fwd_mxfp8_impl( + native_op, x: torch.Tensor, x_ptrs: list[int], combine_buffer: torch.Tensor, @@ -364,7 +364,7 @@ def dispatch_mlp_swiglu_combine_fwd_mxfp8( if tuple(tokens_per_expert.shape) != (num_local_experts,): raise ValueError("tokens_per_expert must have shape (num_local_experts,)") - return _C.dispatch_mlp_swiglu_combine_fwd_mxfp8( + return native_op( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_routed_gate_sc, w_shared_up, w_routed_up, w_routed_up_sc, @@ -374,8 +374,8 @@ def dispatch_mlp_swiglu_combine_fwd_mxfp8( ) -@torch.library.custom_op("mok::dispatch_mlp_swiglu_combine_fwd_bf16", mutates_args=("combine_buffer",)) -def dispatch_mlp_swiglu_combine_fwd_bf16( +def _dispatch_mlp_glu_combine_fwd_bf16_impl( + native_op, x: torch.Tensor, x_ptrs: list[int], combine_buffer: torch.Tensor, @@ -511,7 +511,7 @@ def dispatch_mlp_swiglu_combine_fwd_bf16( if tensor.device != x.device: raise ValueError(f"{tensor_name} must be on {x.device}") - return _C.dispatch_mlp_swiglu_combine_fwd_bf16( + return native_op( x, x_ptrs, combine_buffer, combine_buffer_ptrs, w_shared_gate, w_routed_gate, w_shared_up, w_routed_up, w_shared_down, w_routed_down, @@ -782,14 +782,8 @@ def recompute_forward_context_bf16( ) -@torch.library.custom_op( - "mok::dispatch_mlp_swiglu_combine_bwd_mxfp8", - mutates_args=("d_x_routed_buffer", "d_router_weight_buffer", - "x_fp8_t_routed", "x_sc_t_routed", - "gate_fp8_routed", "gate_sc_routed", "up_fp8_routed", "up_sc_routed", - "hidden_fp8_t_routed", "hidden_sc_t_routed"), -) -def dispatch_mlp_swiglu_combine_bwd_mxfp8( +def _dispatch_mlp_glu_combine_bwd_mxfp8_impl( + native_op, d_y_buffer: torch.Tensor, d_y_buffer_ptrs: list[int], d_x_routed_buffer: torch.Tensor, @@ -1059,7 +1053,7 @@ def dispatch_mlp_swiglu_combine_bwd_mxfp8( if tuple(tokens_per_expert.shape) != (num_local_experts,): raise ValueError("tokens_per_expert must have shape (num_local_experts,)") - return _C.dispatch_mlp_swiglu_combine_bwd_mxfp8( + return native_op( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, @@ -1076,11 +1070,8 @@ def dispatch_mlp_swiglu_combine_bwd_mxfp8( ) -@torch.library.custom_op( - "mok::dispatch_mlp_swiglu_combine_bwd_bf16", - mutates_args=("d_x_routed_buffer", "d_router_weight_buffer", "x_routed", "gate_routed", "up_routed", "hidden_routed"), -) -def dispatch_mlp_swiglu_combine_bwd_bf16( +def _dispatch_mlp_glu_combine_bwd_bf16_impl( + native_op, d_y_buffer: torch.Tensor, d_y_buffer_ptrs: list[int], d_x_routed_buffer: torch.Tensor, @@ -1262,7 +1253,7 @@ def dispatch_mlp_swiglu_combine_bwd_bf16( if schedule_peer_rank.device != x.device: raise ValueError(f"schedule_peer_rank must be on {x.device}") - return _C.dispatch_mlp_swiglu_combine_bwd_bf16( + return native_op( d_y_buffer, d_y_buffer_ptrs, d_x_routed_buffer, d_x_routed_buffer_ptrs, router_weight_buffer, router_weight_buffer_ptrs, d_router_weight_buffer, d_router_weight_buffer_ptrs, @@ -1276,6 +1267,150 @@ def dispatch_mlp_swiglu_combine_bwd_bf16( ) +def _register_glu_op(name, native_op, implementation, *, schema, mutates_args): + def dispatch(*args): + return implementation(native_op, *args) + + dispatch.__name__ = name.rsplit("::", 1)[-1] + dispatch.__doc__ = implementation.__doc__ + return torch.library.custom_op( + name, + mutates_args=mutates_args, + schema=schema, + )(dispatch) + + +_FWD_MXFP8_SCHEMA = """( + Tensor x, SymInt[] x_ptrs, Tensor(a0!) combine_buffer, SymInt[] combine_buffer_ptrs, + Tensor w_shared_gate, Tensor w_routed_gate, Tensor w_routed_gate_sc, + Tensor w_shared_up, Tensor w_routed_up, Tensor w_routed_up_sc, + Tensor w_shared_down, Tensor w_routed_down, Tensor w_routed_down_sc, + Tensor schedule_peer_rank, Tensor schedule_peer_token_idx, + Tensor num_tokens, Tensor tokens_per_expert, + SymInt topk, float? swiglu_limit, SymInt num_comm_sms, SymInt macrobatch_size, SymInt minibatch_size +) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, + Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)""" + +_FWD_BF16_SCHEMA = """( + Tensor x, SymInt[] x_ptrs, Tensor(a0!) combine_buffer, SymInt[] combine_buffer_ptrs, + Tensor w_shared_gate, Tensor w_routed_gate, + Tensor w_shared_up, Tensor w_routed_up, + Tensor w_shared_down, Tensor w_routed_down, + Tensor schedule_peer_rank, Tensor schedule_peer_token_idx, + Tensor num_tokens, Tensor tokens_per_expert, + SymInt topk, float? swiglu_limit, SymInt num_comm_sms, SymInt macrobatch_size, SymInt minibatch_size +) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)""" + +_BWD_MXFP8_SCHEMA = """( + Tensor d_y_buffer, SymInt[] d_y_buffer_ptrs, + Tensor(a0!) d_x_routed_buffer, SymInt[] d_x_routed_buffer_ptrs, + Tensor router_weight_buffer, SymInt[] router_weight_buffer_ptrs, + Tensor(a1!) d_router_weight_buffer, SymInt[] d_router_weight_buffer_ptrs, + Tensor w_shared_gate, Tensor w_routed_gate_T, Tensor w_routed_gate_T_sc, + Tensor w_shared_up, Tensor w_routed_up_T, Tensor w_routed_up_T_sc, + Tensor w_shared_down, Tensor w_routed_down_T, Tensor w_routed_down_T_sc, + Tensor(a2!) x_fp8_t_routed, Tensor(a3!) x_sc_t_routed, + Tensor gate_shared, Tensor(a4!) gate_fp8_routed, Tensor(a5!) gate_sc_routed, + Tensor up_shared, Tensor(a6!) up_fp8_routed, Tensor(a7!) up_sc_routed, + Tensor hidden_shared, Tensor(a8!) hidden_fp8_t_routed, Tensor(a9!) hidden_sc_t_routed, + Tensor x, SymInt[] x_ptrs, + Tensor w_routed_gate, Tensor w_routed_gate_sc, + Tensor w_routed_up, Tensor w_routed_up_sc, + Tensor schedule_peer_rank, Tensor schedule_peer_token_idx, + Tensor num_tokens, Tensor tokens_per_expert, + SymInt topk, float? swiglu_limit, SymInt num_comm_sms, SymInt macrobatch_size, SymInt minibatch_size +) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, + Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)""" + +_BWD_BF16_SCHEMA = """( + Tensor d_y_buffer, SymInt[] d_y_buffer_ptrs, + Tensor(a0!) d_x_routed_buffer, SymInt[] d_x_routed_buffer_ptrs, + Tensor router_weight_buffer, SymInt[] router_weight_buffer_ptrs, + Tensor(a1!) d_router_weight_buffer, SymInt[] d_router_weight_buffer_ptrs, + Tensor w_shared_gate, Tensor w_routed_gate, + Tensor w_shared_up, Tensor w_routed_up, + Tensor w_shared_down, Tensor w_routed_down, + Tensor(a2!) x_routed, + Tensor gate_shared, Tensor(a3!) gate_routed, + Tensor up_shared, Tensor(a4!) up_routed, + Tensor hidden_shared, Tensor(a5!) hidden_routed, + Tensor x, SymInt[] x_ptrs, + Tensor schedule_peer_rank, Tensor schedule_peer_token_idx, + Tensor num_tokens, Tensor tokens_per_expert, + SymInt topk, float? swiglu_limit, SymInt num_comm_sms, SymInt macrobatch_size, SymInt minibatch_size +) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, + Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)""" + +_FWD_MUTATES_ARGS = ("combine_buffer",) +_BWD_MXFP8_MUTATES_ARGS = ( + "d_x_routed_buffer", "d_router_weight_buffer", + "x_fp8_t_routed", "x_sc_t_routed", + "gate_fp8_routed", "gate_sc_routed", "up_fp8_routed", "up_sc_routed", + "hidden_fp8_t_routed", "hidden_sc_t_routed", +) +_BWD_BF16_MUTATES_ARGS = ( + "d_x_routed_buffer", "d_router_weight_buffer", + "x_routed", "gate_routed", "up_routed", "hidden_routed", +) + +dispatch_mlp_swiglu_combine_fwd_mxfp8 = _register_glu_op( + "mok::dispatch_mlp_swiglu_combine_fwd_mxfp8", + _C.dispatch_mlp_swiglu_combine_fwd_mxfp8, + _dispatch_mlp_glu_combine_fwd_mxfp8_impl, + schema=_FWD_MXFP8_SCHEMA, + mutates_args=_FWD_MUTATES_ARGS, +) +dispatch_mlp_situglu_combine_fwd_mxfp8 = _register_glu_op( + "mok::dispatch_mlp_situglu_combine_fwd_mxfp8", + _C.dispatch_mlp_situglu_combine_fwd_mxfp8, + _dispatch_mlp_glu_combine_fwd_mxfp8_impl, + schema=_FWD_MXFP8_SCHEMA, + mutates_args=_FWD_MUTATES_ARGS, +) +dispatch_mlp_swiglu_combine_fwd_bf16 = _register_glu_op( + "mok::dispatch_mlp_swiglu_combine_fwd_bf16", + _C.dispatch_mlp_swiglu_combine_fwd_bf16, + _dispatch_mlp_glu_combine_fwd_bf16_impl, + schema=_FWD_BF16_SCHEMA, + mutates_args=_FWD_MUTATES_ARGS, +) +dispatch_mlp_situglu_combine_fwd_bf16 = _register_glu_op( + "mok::dispatch_mlp_situglu_combine_fwd_bf16", + _C.dispatch_mlp_situglu_combine_fwd_bf16, + _dispatch_mlp_glu_combine_fwd_bf16_impl, + schema=_FWD_BF16_SCHEMA, + mutates_args=_FWD_MUTATES_ARGS, +) +dispatch_mlp_swiglu_combine_bwd_mxfp8 = _register_glu_op( + "mok::dispatch_mlp_swiglu_combine_bwd_mxfp8", + _C.dispatch_mlp_swiglu_combine_bwd_mxfp8, + _dispatch_mlp_glu_combine_bwd_mxfp8_impl, + schema=_BWD_MXFP8_SCHEMA, + mutates_args=_BWD_MXFP8_MUTATES_ARGS, +) +dispatch_mlp_situglu_combine_bwd_mxfp8 = _register_glu_op( + "mok::dispatch_mlp_situglu_combine_bwd_mxfp8", + _C.dispatch_mlp_situglu_combine_bwd_mxfp8, + _dispatch_mlp_glu_combine_bwd_mxfp8_impl, + schema=_BWD_MXFP8_SCHEMA, + mutates_args=_BWD_MXFP8_MUTATES_ARGS, +) +dispatch_mlp_swiglu_combine_bwd_bf16 = _register_glu_op( + "mok::dispatch_mlp_swiglu_combine_bwd_bf16", + _C.dispatch_mlp_swiglu_combine_bwd_bf16, + _dispatch_mlp_glu_combine_bwd_bf16_impl, + schema=_BWD_BF16_SCHEMA, + mutates_args=_BWD_BF16_MUTATES_ARGS, +) +dispatch_mlp_situglu_combine_bwd_bf16 = _register_glu_op( + "mok::dispatch_mlp_situglu_combine_bwd_bf16", + _C.dispatch_mlp_situglu_combine_bwd_bf16, + _dispatch_mlp_glu_combine_bwd_bf16_impl, + schema=_BWD_BF16_SCHEMA, + mutates_args=_BWD_BF16_MUTATES_ARGS, +) + + @torch.library.custom_op("mok::fwd_epilogue", mutates_args=()) def fwd_epilogue( y_shared: torch.Tensor, diff --git a/tests/test_functional_e2e.py b/tests/test_functional_e2e.py index 7a54110..594f8e2 100644 --- a/tests/test_functional_e2e.py +++ b/tests/test_functional_e2e.py @@ -500,3 +500,123 @@ def test_e2e_mxfp8_recomputed_forward_context( MXFP8_TOLERANCE, print_stats=rank == 0, ) + + +def test_e2e_situglu(context: tuple[int, int, torch.device]) -> None: + rank, world_size, device = context + num_experts = world_size + num_local_experts = 1 + hidden_dim = 256 + intermediate_dim = 256 + topk = 1 + num_local_tokens = 512 + inputs = generate_inputs( + rank, + device, + num_experts, + num_local_experts, + topk, + num_local_tokens, + hidden_dim, + intermediate_dim, + ) + ( + x, + topk_experts, + router_weights, + w_shared_gate, + w_shared_up, + w_shared_down, + w_routed_gate, + w_routed_up, + w_routed_down, + d_output, + ) = inputs + reference_results = run_reference_bf16(*inputs, activation="situglu") + config = functional.MoKConfig( + fwd_num_comm_sms=2, + bwd_num_comm_sms=2, + minibatch_size=256, + macrobatch_size=256, + schedule_capacity_multiplier=1.5, + all_gather_top_experts_chunk_bytes=16, + activation="situglu", + ) + workspace = get_workspace( + config, + dist.group.WORLD, + device=device, + num_local_tokens=num_local_tokens, + hidden_size=hidden_dim, + topk=topk, + ) + schedule = functional.build_schedule( + workspace, + config, + topk_experts, + num_local_experts=num_local_experts, + ) + + for precision in ("bf16", "mxfp8"): + if precision == "bf16": + forward_weights = (w_routed_gate, w_routed_up, w_routed_down) + backward_weights = forward_weights + tolerance = BF16_TOLERANCE + else: + gate_fp8, gate_sc, gate_t_fp8, gate_t_sc = mxfp8_quantize( + w_routed_gate, True, True + ) + up_fp8, up_sc, up_t_fp8, up_t_sc = mxfp8_quantize( + w_routed_up, True, True + ) + down_fp8, down_sc, down_t_fp8, down_t_sc = mxfp8_quantize( + w_routed_down, True, True + ) + forward_weights = ( + (gate_fp8, gate_sc), + (up_fp8, up_sc), + (down_fp8, down_sc), + ) + backward_weights = ( + (gate_fp8, gate_sc, gate_t_fp8, gate_t_sc), + (up_fp8, up_sc, up_t_fp8, up_t_sc), + (down_t_fp8, down_t_sc), + ) + tolerance = MXFP8_TOLERANCE + + output, forward_context = functional.forward( + config, + workspace, + schedule, + x, + router_weights, + w_shared_gate, + w_shared_up, + w_shared_down, + *forward_weights, + ) + gradients = functional.backward( + config, + workspace, + schedule, + forward_context, + d_output, + x, + router_weights, + w_shared_gate, + w_shared_up, + w_shared_down, + *backward_weights, + ) + results = (output, *gradients) + + for name, reference, actual in zip( + RESULT_NAMES, reference_results, results, strict=True + ): + check_correctness( + f"situ-glu/{precision}/{name}", + reference, + actual, + tolerance, + print_stats=rank == 0, + ) diff --git a/tests/test_misc.py b/tests/test_misc.py index 7c729dc..ac34f0d 100644 --- a/tests/test_misc.py +++ b/tests/test_misc.py @@ -1223,6 +1223,8 @@ def test_custom_op_mutation_schemas() -> None: "mxfp8_quantize": set(), "dispatch_mlp_swiglu_combine_fwd_mxfp8": {"combine_buffer"}, "dispatch_mlp_swiglu_combine_fwd_bf16": {"combine_buffer"}, + "dispatch_mlp_situglu_combine_fwd_mxfp8": {"combine_buffer"}, + "dispatch_mlp_situglu_combine_fwd_bf16": {"combine_buffer"}, "recompute_forward_context_mxfp8": set(), "recompute_forward_context_bf16": set(), "dispatch_mlp_swiglu_combine_bwd_mxfp8": { @@ -1245,6 +1247,26 @@ def test_custom_op_mutation_schemas() -> None: "up_routed", "hidden_routed", }, + "dispatch_mlp_situglu_combine_bwd_mxfp8": { + "d_x_routed_buffer", + "d_router_weight_buffer", + "x_fp8_t_routed", + "x_sc_t_routed", + "gate_fp8_routed", + "gate_sc_routed", + "up_fp8_routed", + "up_sc_routed", + "hidden_fp8_t_routed", + "hidden_sc_t_routed", + }, + "dispatch_mlp_situglu_combine_bwd_bf16": { + "d_x_routed_buffer", + "d_router_weight_buffer", + "x_routed", + "gate_routed", + "up_routed", + "hidden_routed", + }, "fwd_epilogue": set(), "bwd_epilogue": set(), } diff --git a/tests/test_situglu.py b/tests/test_situglu.py new file mode 100644 index 0000000..005390a --- /dev/null +++ b/tests/test_situglu.py @@ -0,0 +1,59 @@ +import torch + +from .utils import apply_glu, situglu + + +def test_situglu_is_bounded() -> None: + values = torch.linspace(-1000.0, 1000.0, 257, dtype=torch.float64) + hidden = situglu(values[:, None], values[None, :]) + + assert torch.isfinite(hidden).all() + assert hidden.abs().max() <= 100.0 + + +def test_situglu_matches_swiglu_near_origin() -> None: + gate = torch.tensor([-1e-4, 0.0, 1e-4], dtype=torch.float64) + up = torch.tensor([1e-4, -1e-4, 1e-4], dtype=torch.float64) + + torch.testing.assert_close( + situglu(gate, up), + torch.nn.functional.silu(gate) * up, + rtol=1e-8, + atol=1e-16, + ) + + +def test_situglu_closed_form_derivatives() -> None: + gate = torch.tensor([-7.0, -0.5, 3.0], dtype=torch.float64, requires_grad=True) + up = torch.tensor([-30.0, 0.25, 40.0], dtype=torch.float64, requires_grad=True) + situglu(gate, up).sum().backward() + + gate_tanh = torch.tanh(gate.detach() / 4.0) + up_tanh = torch.tanh(up.detach() / 25.0) + sigmoid = torch.sigmoid(gate.detach()) + gate_factor = 4.0 * gate_tanh * sigmoid + up_factor = 25.0 * up_tanh + d_gate_factor = (1.0 - gate_tanh.square()) * sigmoid + gate_factor * (1.0 - sigmoid) + d_up_factor = 1.0 - up_tanh.square() + + torch.testing.assert_close(gate.grad, d_gate_factor * up_factor) + torch.testing.assert_close(up.grad, gate_factor * d_up_factor) + + +def test_situglu_gradcheck() -> None: + gate = torch.randn(2, 3, dtype=torch.float64, requires_grad=True) + up = torch.randn(2, 3, dtype=torch.float64, requires_grad=True) + + assert torch.autograd.gradcheck(situglu, (gate, up)) + + +def test_apply_glu_rejects_unknown_activation() -> None: + gate = torch.zeros(1) + up = torch.zeros(1) + + try: + apply_glu(gate, up, "unknown") + except ValueError as error: + assert str(error) == "activation must be 'swiglu' or 'situglu'" + else: + raise AssertionError("unknown activation was accepted") diff --git a/tests/utils.py b/tests/utils.py index 3771f78..b29f1fe 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -8,6 +8,26 @@ MXFP8_TOLERANCE = (1.0, 0.1) +def situglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + # h = [beta1 * tanh(g / beta1) * sigmoid(g)] * [beta2 * tanh(u / beta2)]. + gate_factor = 4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate) + up_factor = 25.0 * torch.tanh(up / 25.0) + return gate_factor * up_factor + + +def apply_glu( + gate: torch.Tensor, + up: torch.Tensor, + activation: str, + swiglu_limit: float | None = None, +) -> torch.Tensor: + if activation == "swiglu": + return run_swiglu_reference(gate, up, swiglu_limit) + if activation == "situglu": + return situglu(gate, up) + raise ValueError("activation must be 'swiglu' or 'situglu'") + + def shapes(world_size: int) -> tuple[tuple[str, int, int, int, int, int], ...]: return ( # (name, routed experts, hidden dim, intermediate dim, top-k, num local tokens) ("Kimi K2.7 Code", 384, 7168, 2048, 8, 7168), @@ -245,6 +265,8 @@ def run_forward_reference_bf16( w_routed_up: torch.Tensor, # [E, I, H] w_routed_down: torch.Tensor, # [E, H, I] swiglu_limit: float | None = None, + *, + activation: str = "swiglu", ) -> tuple[ torch.Tensor, # combine_buffer torch.Tensor, # gate_shared @@ -280,7 +302,7 @@ def run_forward_reference_bf16( expert_x = recv_x[rows] gate = expert_x @ w_routed_gate[expert_idx].T up = expert_x @ w_routed_up[expert_idx].T - hidden_activations = run_swiglu_reference(gate, up, swiglu_limit) + hidden_activations = apply_glu(gate, up, activation, swiglu_limit) recv_output = recv_output.index_copy( 0, rows, hidden_activations @ w_routed_down[expert_idx].T) @@ -290,7 +312,7 @@ def run_forward_reference_bf16( gate_shared = x @ w_shared_gate.T up_shared = x @ w_shared_up.T - hidden_shared = run_swiglu_reference(gate_shared, up_shared, swiglu_limit) + hidden_shared = apply_glu(gate_shared, up_shared, activation, swiglu_limit) y_shared = hidden_shared @ w_shared_down.T return combine_buffer, gate_shared, up_shared, hidden_shared, y_shared @@ -330,6 +352,7 @@ def run_reference_bf16( d_output: torch.Tensor, # [T, H] swiglu_limit: float | None = None, *, + activation: str = "swiglu", group: dist.ProcessGroup | None = None, ) -> tuple[ torch.Tensor, # output @@ -375,7 +398,7 @@ def run_reference_bf16( expert_x = recv_x[rows] gate = expert_x @ w_routed_gate[expert_idx].T up = expert_x @ w_routed_up[expert_idx].T - hidden_activations = run_swiglu_reference(gate, up, swiglu_limit) + hidden_activations = apply_glu(gate, up, activation, swiglu_limit) recv_output = recv_output.index_copy(0, rows, hidden_activations @ w_routed_down[expert_idx].T) # Routed token weighted sum @@ -391,7 +414,7 @@ def run_reference_bf16( w_shared_down = w_shared_down.detach().requires_grad_() gate_shared = x_shared @ w_shared_gate.T up_shared = x_shared @ w_shared_up.T - shared_output = run_swiglu_reference(gate_shared, up_shared, swiglu_limit) @ w_shared_down.T + shared_output = apply_glu(gate_shared, up_shared, activation, swiglu_limit) @ w_shared_down.T # Final sum output = (routed_output + shared_output.float()).to(torch.bfloat16)