diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 18fc98a39d..1d1e208c3a 100644 --- a/csrc/apis/mega.hpp +++ b/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(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"); diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index d0e973d7d0..423c459c2b 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -24,6 +24,7 @@ class SM100FP8FP4MegaMoERuntime final : public LaunchRuntimearrive(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(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), diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh index 86dba4b4c9..ab983d0dd1 100755 --- a/sgl_deep_gemm/run_tests.sh +++ b/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 diff --git a/sgl_deep_gemm/tests/test_mega_moe_situ.py b/sgl_deep_gemm/tests/test_mega_moe_situ.py new file mode 100644 index 0000000000..df2c891428 --- /dev/null +++ b/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::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, + )