From ad131bb6322ede95a12408b6ca24930b58492246 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Fri, 14 Aug 2026 15:38:21 -0700 Subject: [PATCH 1/2] Use explicit SiTU activation in MegaMoE --- README.md | 4 +- csrc/apis/mega.hpp | 13 +- .../impls/sm100_fp8_fp4_mega_moe.hpp | 5 + .../impls/sm100_fp8_fp4_mega_moe.cuh | 9 +- deep_gemm/mega/__init__.py | 22 ++- sgl_deep_gemm/run_tests.sh | 2 + sgl_deep_gemm/tests/test_mega_moe_situ.py | 145 ++++++++++++++++++ 7 files changed, 188 insertions(+), 12 deletions(-) create mode 100644 sgl_deep_gemm/tests/test_mega_moe_situ.py diff --git a/README.md b/README.md index 6ef705ffce..0da277efe2 100644 --- a/README.md +++ b/README.md @@ -113,7 +113,7 @@ For more details and the paged version `fp8_paged_mqa_logits`, please refer to ` #### Mega MoE -Mega MoE fuses and overlaps EP dispatch, linear 1 (FP8xFP4), SwiGLU, linear 2 (FP8xFP4), and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory. Usage: +Mega MoE fuses and overlaps EP dispatch, linear 1 (FP8xFP4), activation, linear 2 (FP8xFP4), and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory. Usage: ```python # Allocate symmetric memory buffer @@ -137,6 +137,8 @@ y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer) ``` +The FP8xFP4 kernel accepts `activation="swiglu"` (default) and `activation="situ"`. Pass the same activation to `get_symm_buffer_for_mega_moe`, `transform_weights_for_mega_moe`, and `fp8_fp4_mega_moe`. The SiTU path uses the Kimi-K3 constants beta=4.0 and linear beta=25.0; `activation_clamp` applies only to SwiGLU. + For the full example with multi-process setup and benchmarking, please refer to `tests/test_mega_moe.py`. #### Utilities 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/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 165ff769c3..19accfce91 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -15,6 +15,14 @@ from .. import _C _MAX_CANDIDATE_BLOCK_M = 192 +_FP8_FP4_ACTIVATIONS = ('swiglu', 'situ') + + +def _validate_fp8_fp4_activation(activation: str) -> None: + assert activation in _FP8_FP4_ACTIVATIONS, ( + f'FP8xFP4 MegaMoE activation must be one of ' + f'{_FP8_FP4_ACTIVATIONS}, got `{activation}`' + ) class SymmBuffer: @@ -25,7 +33,13 @@ def __init__(self, group: dist.ProcessGroup, num_ring_tokens: int, mma_type: str = 'fp8xfp4', activation: str = 'swiglu'): - assert activation == 'swiglu', f'Only `swiglu` activation is supported, got `{activation}`' + if mma_type == 'fp8xfp4': + _validate_fp8_fp4_activation(activation) + else: + assert activation == 'swiglu', ( + f'Only `swiglu` activation is supported for `{mma_type}`, ' + f'got `{activation}`' + ) self.group = group self.num_experts = num_experts self.num_max_tokens_per_rank = num_max_tokens_per_rank @@ -145,7 +159,7 @@ def transform_weights_for_mega_moe( activation: str = 'swiglu' ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]: - assert activation == 'swiglu', f'Only `swiglu` activation is supported, got `{activation}`' + _validate_fp8_fp4_activation(activation) if isinstance(l1_weights, tuple): # FP8: interleave gate/up for weight and SF, then transpose L1 SF for UTCCP l1_w = _interleave_weights(l1_weights[0]) @@ -170,6 +184,10 @@ def fp8_fp4_mega_moe(y: torch.Tensor, activation: str = 'swiglu', activation_clamp: Optional[float] = None, fast_math: bool = True): + _validate_fp8_fp4_activation(activation) + assert activation != 'situ' or activation_clamp is None, ( + '`activation_clamp` is not supported with `activation="situ"`' + ) (l1_weights_data, l1_weights_sf) = l1_weights (l2_weights_data, l2_weights_sf) = l2_weights _C.fp8_fp4_mega_moe( 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..7854ff9ad0 --- /dev/null +++ b/sgl_deep_gemm/tests/test_mega_moe_situ.py @@ -0,0 +1,145 @@ +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), + activation="situ", + ) + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + activation="situ", + ) + 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 AssertionError 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, + ) From 0e9020bad619840782ff874998fdad9670933147 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Fri, 14 Aug 2026 16:23:59 -0700 Subject: [PATCH 2/2] Keep SiTU selection at the kernel boundary --- README.md | 4 +--- deep_gemm/mega/__init__.py | 22 ++-------------------- sgl_deep_gemm/tests/test_mega_moe_situ.py | 4 +--- 3 files changed, 4 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index 0da277efe2..6ef705ffce 100644 --- a/README.md +++ b/README.md @@ -113,7 +113,7 @@ For more details and the paged version `fp8_paged_mqa_logits`, please refer to ` #### Mega MoE -Mega MoE fuses and overlaps EP dispatch, linear 1 (FP8xFP4), activation, linear 2 (FP8xFP4), and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory. Usage: +Mega MoE fuses and overlaps EP dispatch, linear 1 (FP8xFP4), SwiGLU, linear 2 (FP8xFP4), and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory. Usage: ```python # Allocate symmetric memory buffer @@ -137,8 +137,6 @@ y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer) ``` -The FP8xFP4 kernel accepts `activation="swiglu"` (default) and `activation="situ"`. Pass the same activation to `get_symm_buffer_for_mega_moe`, `transform_weights_for_mega_moe`, and `fp8_fp4_mega_moe`. The SiTU path uses the Kimi-K3 constants beta=4.0 and linear beta=25.0; `activation_clamp` applies only to SwiGLU. - For the full example with multi-process setup and benchmarking, please refer to `tests/test_mega_moe.py`. #### Utilities diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 19accfce91..165ff769c3 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -15,14 +15,6 @@ from .. import _C _MAX_CANDIDATE_BLOCK_M = 192 -_FP8_FP4_ACTIVATIONS = ('swiglu', 'situ') - - -def _validate_fp8_fp4_activation(activation: str) -> None: - assert activation in _FP8_FP4_ACTIVATIONS, ( - f'FP8xFP4 MegaMoE activation must be one of ' - f'{_FP8_FP4_ACTIVATIONS}, got `{activation}`' - ) class SymmBuffer: @@ -33,13 +25,7 @@ def __init__(self, group: dist.ProcessGroup, num_ring_tokens: int, mma_type: str = 'fp8xfp4', activation: str = 'swiglu'): - if mma_type == 'fp8xfp4': - _validate_fp8_fp4_activation(activation) - else: - assert activation == 'swiglu', ( - f'Only `swiglu` activation is supported for `{mma_type}`, ' - f'got `{activation}`' - ) + assert activation == 'swiglu', f'Only `swiglu` activation is supported, got `{activation}`' self.group = group self.num_experts = num_experts self.num_max_tokens_per_rank = num_max_tokens_per_rank @@ -159,7 +145,7 @@ def transform_weights_for_mega_moe( activation: str = 'swiglu' ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]: - _validate_fp8_fp4_activation(activation) + assert activation == 'swiglu', f'Only `swiglu` activation is supported, got `{activation}`' if isinstance(l1_weights, tuple): # FP8: interleave gate/up for weight and SF, then transpose L1 SF for UTCCP l1_w = _interleave_weights(l1_weights[0]) @@ -184,10 +170,6 @@ def fp8_fp4_mega_moe(y: torch.Tensor, activation: str = 'swiglu', activation_clamp: Optional[float] = None, fast_math: bool = True): - _validate_fp8_fp4_activation(activation) - assert activation != 'situ' or activation_clamp is None, ( - '`activation_clamp` is not supported with `activation="situ"`' - ) (l1_weights_data, l1_weights_sf) = l1_weights (l2_weights_data, l2_weights_sf) = l2_weights _C.fp8_fp4_mega_moe( diff --git a/sgl_deep_gemm/tests/test_mega_moe_situ.py b/sgl_deep_gemm/tests/test_mega_moe_situ.py index 7854ff9ad0..df2c891428 100644 --- a/sgl_deep_gemm/tests/test_mega_moe_situ.py +++ b/sgl_deep_gemm/tests/test_mega_moe_situ.py @@ -62,7 +62,6 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): 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), - activation="situ", ) buffer = deep_gemm.get_symm_buffer_for_mega_moe( group, @@ -71,7 +70,6 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): num_topk, hidden, intermediate_hidden, - activation="situ", ) cumulative_recv_stats = torch.zeros((num_experts,), dtype=torch.int, device="cuda") @@ -121,7 +119,7 @@ def run(activation: str, activation_clamp=None): try: run("situ", activation_clamp=0.03125) - except AssertionError as error: + except RuntimeError as error: assert "activation_clamp" in str(error) else: raise AssertionError("SiTU must reject the unrelated activation_clamp option")