[SM100] Select SiTU through the activation argument

csrc/apis/mega.hpp

@@ -38,7 +38,11 @@ get_symm_buffer_size_for_mega_moe(
     const std::string& mma_type, const std::string& activation,
     const int& num_ring_tokens) {
     DG_HOST_ASSERT(num_experts % num_ranks == 0);
-    DG_HOST_ASSERT(activation == "swiglu");
+
+    // SiTU is implemented only by the SM100 FP8xFP4 MegaMoE kernel.
+    const auto mma_kind = parse_mma_kind(mma_type);
+    DG_HOST_ASSERT(activation == "swiglu" or
+                   (mma_kind == MmaKind::MXFP8FP4 and activation == "situ"));
 
     // Pool capacity must fit at least one full wave (one expert per wave) and aligned to block size
     const auto num_experts_per_rank = num_experts / num_ranks;
@@ -48,7 +52,6 @@ get_symm_buffer_size_for_mega_moe(
     DG_HOST_ASSERT(num_min_ring_tokens <= num_ring_tokens and num_ring_tokens <= num_max_ring_tokens);
 
     // Parse MMA type
-    const auto mma_kind = parse_mma_kind(mma_type);
     const auto num_mma_elem_bytes = get_num_mma_elem_bytes(mma_kind);
     const auto with_sf = is_mma_with_sf(mma_kind);
 
@@ -220,7 +223,9 @@ static void fp8_fp4_mega_moe(
     const auto num_tokens = static_cast<int>(y.size(0));
     const auto [rm, rn, rk] = recipe;
     DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32);
-    DG_HOST_ASSERT(activation == "swiglu");
+    DG_HOST_ASSERT(activation == "swiglu" or activation == "situ");
+    DG_HOST_ASSERT(activation != "situ" or not activation_clamp_opt.has_value());
+    const bool use_situ = activation == "situ";
 
     // Activation checks
     const auto activation_clamp =
@@ -302,7 +307,7 @@ static void fp8_fp4_mega_moe(
                                num_experts_per_rank,
                                num_tokens, num_topk,
                                hidden, intermediate_hidden,
-                               activation_clamp, fast_math,
+                               activation_clamp, use_situ, fast_math,
                                use_fp4_acts, use_mxf4_kind, use_fp8_combine);
     } else {
         DG_HOST_UNREACHABLE("Unsupported architecture");

csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp

@@ -24,6 +24,7 @@ class SM100FP8FP4MegaMoERuntime final : public LaunchRuntime<SM100FP8FP4MegaMoER
         int num_experts, num_topk;
         int num_ranks;
         float activation_clamp;
+        bool use_situ;
         bool fast_math;
         // Stream A0.1: enable FP4 (E2M1) activations from L1 epilogue.
         // Default false — keeps the FP8-acts baseline byte-identical.
@@ -79,6 +80,7 @@ static void __instantiate_kernel() {{
         {},
         {},
         {},
+        {},
         {}, {}, {},
         {}, {},
         {},
@@ -102,6 +104,7 @@ static void __instantiate_kernel() {{
     args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads,
     args.launch_args.grid_dim.first, args.num_ranks,
     to_string(args.activation_clamp),
+    args.use_situ ? "true" : "false",
     args.fast_math ? "true" : "false",
     args.use_fp4_acts ? "true" : "false",
     args.use_mxf4_kind ? "true" : "false",
@@ -141,6 +144,7 @@ static void sm100_fp8_fp4_mega_moe(
     const int& num_tokens, const int& num_topk,
     const int& hidden, const int& intermediate_hidden,
     const float& activation_clamp,
+    const bool& use_situ,
     const bool& fast_math,
     const bool& use_fp4_acts = false,
     const bool& use_mxf4_kind = false,
@@ -302,6 +306,7 @@ static void sm100_fp8_fp4_mega_moe(
         .num_experts = num_experts, .num_topk = num_topk,
         .num_ranks = num_ranks,
         .activation_clamp = activation_clamp,
+        .use_situ = use_situ,
         .fast_math = fast_math,
         .use_fp4_acts = use_fp4_acts,
         .use_mxf4_kind = use_mxf4_kind,

deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh

@@ -34,6 +34,7 @@ template <
     uint32_t kNumEpilogueThreads,
     uint32_t kNumSMs, uint32_t kNumRanks,
     float kActivationClamp,
+    bool kUseSitu,
     bool kFastMath,
     // ====== Stream A0.1 — DG_USE_FP4_ACTS ======
     // When true, the L1 epilogue quantizes its SwiGLU outputs to E2M1 (FP4) +
@@ -1328,14 +1329,12 @@ sm100_fp8_fp4_mega_moe_impl(void* y,
                             tmem_empty_barriers[accum_stage_idx]->arrive(0u);
                         }
 
-                        // Apply activation: SwiGLU, or Kimi-K3 SiTU via sentinel
+                        // Apply the explicitly selected SwiGLU or Kimi-K3 SiTU activation.
                         // Gate/up pairs: (0, 2), (1, 3), (4, 6), (5, 7)
-                        // K3-SITU-PATCH: kActivationClamp == 0.03125f (2^-5 magic;
-                        // host asserts clamp >= 0 so negatives can't sentinel) selects SiTU:
+                        // SiTU:
                         //   act = kSituBeta * tanh(gate/kSituBeta) * sigmoid(gate)
                         //   up' = kSituLinearBeta * tanh(up/kSituLinearBeta)
                         // K3 config constants baked in (activation_situ_{beta,linear_beta}).
-                        constexpr bool kUseSitu = (kActivationClamp == 0.03125f);
                         constexpr float kSituBeta = 4.0f;
                         constexpr float kSituLinearBeta = 25.0f;
                         auto fp32_values = reinterpret_cast<float*>(values);
@@ -1365,7 +1364,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y,
                             }
                             auto up = __bfloat1622float2(bf16_up);
                             if constexpr (kUseSitu) {
-                                // K3-SITU-PATCH: tanh-bounded gate, soft-clipped up
+                                // Tanh-bounded gate and soft-clipped up branch.
                                 gate = {kSituBeta * tanhf(gate.x / kSituBeta) * sig.x,
                                         kSituBeta * tanhf(gate.y / kSituBeta) * sig.y};
                                 up = {kSituLinearBeta * tanhf(up.x / kSituLinearBeta),

sgl_deep_gemm/run_tests.sh

@@ -153,6 +153,7 @@ fi
 # pre_dispatch tests use deep_gemm's own symmetric buffer.
 MEGA_MOE_BLACKWELL=(
   test_mega_moe.py
+  test_mega_moe_situ.py
   test_mega_moe_l1_fp4_accuracy.py
   test_mega_moe_l1_sentinel.py
   test_mega_moe_pre_dispatch.py
@@ -181,6 +182,7 @@ elif [ "${ARCH_MAJOR}" -ge 10 ]; then
       skip_test test_mega_moe.py "deep_ep with ElasticBuffer not installed"
     fi
   fi
+  [ -f "${TESTS_DIR}/test_mega_moe_situ.py" ] && run_test test_mega_moe_situ.py --num-processes 1
   L1_NPROC="${NPROC}"
   if [ "${L1_NPROC}" -gt 2 ]; then
     L1_NPROC=2

sgl_deep_gemm/tests/test_mega_moe_situ.py

@@ -0,0 +1,143 @@
+import argparse
+import os
+import re
+import tempfile
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+
+import deep_gemm
+from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8
+from deep_gemm.utils.dist import dist_print, init_dist
+
+
+def _cast_grouped_weights_to_fp4(weights: torch.Tensor):
+    num_groups, n, k = weights.shape
+    data = torch.empty((num_groups, n, k // 2), dtype=torch.int8, device="cuda")
+    scale = torch.empty((num_groups, n, k // 32), dtype=torch.float, device="cuda")
+    for group_idx in range(num_groups):
+        data[group_idx], scale[group_idx] = per_token_cast_to_fp4(
+            weights[group_idx], use_ue8m0=True, gran_k=32
+        )
+    scale = deep_gemm.transform_sf_into_required_layout(
+        scale, n, k, (1, 32), num_groups
+    )
+    return data, scale
+
+
+def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
+    jit_cache = tempfile.TemporaryDirectory(prefix="deep-gemm-situ-")
+    os.environ["DG_JIT_CACHE_DIR"] = jit_cache.name
+    _, _, group = init_dist(local_rank, num_local_ranks)
+    os.environ["DG_COMM_KERNEL_DEBUG"] = "0"
+
+    num_tokens = 32
+    num_max_tokens_per_rank = 8192
+    hidden = 1024
+    intermediate_hidden = 512
+    num_experts = 8
+    num_topk = 2
+
+    torch.manual_seed(0)
+    x = torch.full((num_tokens, hidden), 0.25, dtype=torch.bfloat16, device="cuda")
+    l1_weights = torch.full(
+        (num_experts, 2 * intermediate_hidden, hidden),
+        0.25,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    l2_weights = torch.full(
+        (num_experts, hidden, intermediate_hidden),
+        0.25,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device="cuda")
+    topk_weights, topk_idx = torch.topk(
+        scores, num_topk, dim=-1, largest=True, sorted=False
+    )
+
+    x_fp8 = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
+    transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(
+        _cast_grouped_weights_to_fp4(l1_weights),
+        _cast_grouped_weights_to_fp4(l2_weights),
+    )
+    buffer = deep_gemm.get_symm_buffer_for_mega_moe(
+        group,
+        num_experts,
+        num_max_tokens_per_rank,
+        num_topk,
+        hidden,
+        intermediate_hidden,
+    )
+    cumulative_recv_stats = torch.zeros((num_experts,), dtype=torch.int, device="cuda")
+
+    def run(activation: str, activation_clamp=None):
+        buffer.x[:num_tokens].copy_(x_fp8[0])
+        buffer.x_sf[:num_tokens].copy_(x_fp8[1])
+        buffer.topk_idx[:num_tokens].copy_(topk_idx)
+        buffer.topk_weights[:num_tokens].copy_(topk_weights)
+        cumulative_recv_stats.zero_()
+        y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
+        deep_gemm.fp8_fp4_mega_moe(
+            y,
+            transformed_l1,
+            transformed_l2,
+            buffer,
+            cumulative_local_expert_recv_stats=cumulative_recv_stats,
+            activation=activation,
+            activation_clamp=activation_clamp,
+            fast_math=True,
+        )
+        torch.cuda.synchronize()
+        return y
+
+    situ = run("situ")
+    clamped_swiglu = run("swiglu", activation_clamp=0.03125)
+
+    assert torch.isfinite(situ).all()
+    assert torch.isfinite(clamped_swiglu).all()
+
+    kernel_sources = [
+        path.read_text()
+        for path in Path(jit_cache.name).glob(
+            "cache/kernel.sm100_fp8_fp4_mega_moe.*/kernel.cu"
+        )
+    ]
+    assert len(kernel_sources) == 2
+    assert any(
+        re.search(
+            r"cute::numeric_limits<float>::infinity\(\),\s+true,\s+true,",
+            source,
+        )
+        for source in kernel_sources
+    ), "explicit SiTU did not instantiate kUseSitu=true"
+    assert any(
+        re.search(r"0x1p-5f,\s+false,\s+true,", source) for source in kernel_sources
+    ), "activation_clamp=0.03125 still instantiated kUseSitu=true"
+
+    try:
+        run("situ", activation_clamp=0.03125)
+    except RuntimeError as error:
+        assert "activation_clamp" in str(error)
+    else:
+        raise AssertionError("SiTU must reject the unrelated activation_clamp option")
+
+    dist_print(
+        "Explicit SiTU and tightly clamped SwiGLU both launched", once_in_node=True
+    )
+    buffer.destroy()
+    dist.destroy_process_group()
+    jit_cache.cleanup()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-processes", type=int, default=1)
+    parsed_args = parser.parse_args()
+    torch.multiprocessing.spawn(
+        test,
+        args=(parsed_args.num_processes, parsed_args),
+        nprocs=parsed_args.num_processes,
+    )