diff --git a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py index 631e4d98c..3989faa83 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py @@ -1342,6 +1342,8 @@ def _fused_append_shared_experts_with_weights_kernel( shared_weights_ptr, out_ids_ptr, out_weights_ptr, + hidden_ptr, + wgate_ptr, N_BASE, scale, K: tl.constexpr, @@ -1349,6 +1351,9 @@ def _fused_append_shared_experts_with_weights_kernel( BLOCK_K: tl.constexpr, BLOCK_S: tl.constexpr, APPLY_SIGMOID: tl.constexpr, + FUSE_GATE: tl.constexpr, + HIDDEN: tl.constexpr, + BLOCK_H: tl.constexpr, ): pid = tl.program_id(0) @@ -1366,12 +1371,20 @@ def _fused_append_shared_experts_with_weights_kernel( offs_s = tl.arange(0, BLOCK_S) mask_s = offs_s < S shared_ids = tl.cast(N_BASE + offs_s, ids.dtype) - shared_ws = tl.load(shared_weights_ptr + pid * S + offs_s, mask=mask_s) - if APPLY_SIGMOID: - # Fuse sigmoid(shared_gate) + dtype upcast (+ optional 1/ep_size scale) - # in-register so the raw bf16 logits stream straight into the fp32 - # output, eliminating the standalone sigmoid and bf16->fp32 copy kernels. - shared_ws = tl.sigmoid(shared_ws.to(tl.float32)) * scale + if FUSE_GATE: + offs_h = tl.arange(0, BLOCK_H) + mask_h = offs_h < HIDDEN + h = tl.load(hidden_ptr + pid * HIDDEN + offs_h, mask=mask_h, other=0.0).to( + tl.float32 + ) + w = tl.load(wgate_ptr + offs_h, mask=mask_h, other=0.0).to(tl.float32) + logit = tl.sum(h * w) + shared_val = tl.sigmoid(logit) * scale + shared_ws = tl.zeros((BLOCK_S,), dtype=tl.float32) + shared_val + else: + shared_ws = tl.load(shared_weights_ptr + pid * S + offs_s, mask=mask_s) + if APPLY_SIGMOID: + shared_ws = tl.sigmoid(shared_ws.to(tl.float32)) * scale tl.store(out_ids_ptr + out_row_ptr + K + offs_s, shared_ids, mask=mask_s) tl.store(out_weights_ptr + out_row_ptr + K + offs_s, shared_ws, mask=mask_s) @@ -1384,31 +1397,65 @@ def fused_append_shared_experts_with_weights( num_fused_shared_experts, N=None, apply_sigmoid=False, + fuse_gate=False, + hidden_states=None, + gate_weight=None, scale=1.0, ): """Like fused_append_shared_experts but accepts per-token shared weights tensor. - When ``apply_sigmoid`` is True, ``shared_weights`` are treated as raw gate - logits: the kernel applies ``sigmoid`` (in fp32) and the optional ``scale`` - in-register, so the caller can skip the separate ``sigmoid`` activation and - the bf16->fp32 cast. When False the legacy behavior is preserved exactly. + Two optional in-kernel fusions are supported (both default off → legacy + behavior is preserved byte-for-byte): + + - ``apply_sigmoid=True``: ``shared_weights`` are treated as raw gate logits; + the kernel applies ``sigmoid`` (in fp32) and the optional ``scale`` + in-register, so the caller can skip the separate ``sigmoid`` activation + and the bf16->fp32 cast. + - ``fuse_gate=True``: the shared_expert_gate GEMV + (``hidden_states @ gate_weight.T``) + sigmoid + ``scale`` are computed + *inside* the kernel, eliminating the standalone gate GEMM launch. + ``shared_weights`` is ignored; ``hidden_states`` ([M, HIDDEN]) and + ``gate_weight`` ([1, HIDDEN] or [HIDDEN]) must be provided. This subsumes + ``apply_sigmoid`` (the sigmoid is intrinsic), so the two are mutually + exclusive. """ + assert not ( + fuse_gate and apply_sigmoid + ), "fuse_gate already applies sigmoid in-kernel; do not also set apply_sigmoid" assert N is not None, "N (shared expert base id) must be provided" m, k = topk_ids.shape s = int(num_fused_shared_experts) if s <= 0: return topk_ids, topk_weights - # When fusing sigmoid in-kernel, keep the raw logits dtype (the kernel emits - # fp32 directly); otherwise match the output weight dtype as before. - shared_weights_2d = ( - shared_weights if apply_sigmoid else shared_weights.to(topk_weights.dtype) - ) - if shared_weights_2d.ndim == 1: - shared_weights_2d = shared_weights_2d.unsqueeze(-1) - if shared_weights_2d.shape[1] < s: - shared_weights_2d = shared_weights_2d.expand(m, s) - shared_weights_2d = shared_weights_2d.contiguous() + if fuse_gate: + assert ( + hidden_states is not None and gate_weight is not None + ), "fuse_gate=True requires hidden_states and gate_weight" + hidden_arg = hidden_states.contiguous() + wgate_arg = gate_weight.reshape(-1).contiguous() + hidden_dim = hidden_arg.shape[1] + block_h = triton.next_power_of_2(hidden_dim) + shared_arg = topk_weights + num_warps = 8 + else: + # When fusing sigmoid in-kernel (apply_sigmoid), keep the raw logits + # dtype (the kernel emits fp32 directly); otherwise match the output + # weight dtype as before. + shared_weights_2d = ( + shared_weights if apply_sigmoid else shared_weights.to(topk_weights.dtype) + ) + if shared_weights_2d.ndim == 1: + shared_weights_2d = shared_weights_2d.unsqueeze(-1) + if shared_weights_2d.shape[1] < s: + shared_weights_2d = shared_weights_2d.expand(m, s) + shared_arg = shared_weights_2d.contiguous() + # hidden_ptr / wgate_ptr are unused; pass placeholders. + hidden_arg = topk_weights + wgate_arg = topk_weights + hidden_dim = 1 + block_h = 1 + num_warps = 1 out_ids = torch.empty((m, k + s), dtype=topk_ids.dtype, device=topk_ids.device) out_weights = torch.empty( @@ -1421,9 +1468,11 @@ def fused_append_shared_experts_with_weights( _fused_append_shared_experts_with_weights_kernel[(m,)]( topk_ids, topk_weights, - shared_weights_2d, + shared_arg, out_ids, out_weights, + hidden_arg, + wgate_arg, N_BASE=N, scale=scale, K=k, @@ -1431,6 +1480,9 @@ def fused_append_shared_experts_with_weights( BLOCK_K=block_k, BLOCK_S=block_s, APPLY_SIGMOID=apply_sigmoid, - num_warps=1, + FUSE_GATE=fuse_gate, + HIDDEN=hidden_dim, + BLOCK_H=block_h, + num_warps=num_warps, ) return out_ids, out_weights diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 87594757a..dabfedbf1 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -394,34 +394,64 @@ class Qwen2MoeSparseMoeBlock(nn.Module): return F.sigmoid(shared_logits) * scale, 1.0 return shared_logits, scale + def _shared_expert_scale(self) -> float: + """1/ep_size pre-scale for the fused shared-expert routing weight. + + Mirrors the scaling applied in _get_shared_expert_weights; see that + method for the allreduce-EP rationale. + """ + moe_ep_size = get_parallel().moe_ep_size + if moe_ep_size > 1 and not is_deepep_class_backend(): + return 1.0 / float(moe_ep_size) + return 1.0 + def _append_shared_to_topk_output( self, topk_output: StandardTopKOutput, hidden_states: torch.Tensor, ) -> StandardTopKOutput: """Append shared expert ids and weights to topk output before fused MoE.""" - if not self.enable_shared_expert_fusion: + if not self.enable_shared_expert_fusion or self.shared_expert_gate is None: return topk_output - shared = self._get_shared_expert_weights(hidden_states) - if shared is None: - return topk_output - shared_weights, shared_scale = shared from sglang.kernels.ops.moe.fused_moe_triton_kernels import ( fused_append_shared_experts_with_weights, ) - # AITER returns raw logits + scale for in-kernel sigmoid fusion; CUDA - # returns pre-activated weights (scale already folded in) → no fusion. - fused_topk_ids, fused_topk_weights = fused_append_shared_experts_with_weights( - topk_output.topk_ids, - topk_output.topk_weights, - shared_weights, - self.num_fused_shared_experts, - N=self.num_experts, - apply_sigmoid=_use_aiter, - scale=shared_scale, - ) + if _use_aiter: + # HIP/aiter: fuse the shared_expert_gate GEMV + sigmoid + scale into + # the append kernel, eliminating the standalone gate GEMM launch. + # This subsumes the sigmoid-only fusion: there is no separate gate + # GEMM and no _get_shared_expert_weights call on this path. + fused_topk_ids, fused_topk_weights = ( + fused_append_shared_experts_with_weights( + topk_output.topk_ids, + topk_output.topk_weights, + None, + self.num_fused_shared_experts, + N=self.num_experts, + fuse_gate=True, + hidden_states=hidden_states, + gate_weight=self.shared_expert_gate.weight, + scale=self._shared_expert_scale(), + ) + ) + else: + # CUDA: _get_shared_expert_weights returns pre-activated weights + # (sigmoid + scale already folded in) → legacy append, no fusion. + shared = self._get_shared_expert_weights(hidden_states) + if shared is None: + return topk_output + shared_weights, _ = shared + fused_topk_ids, fused_topk_weights = ( + fused_append_shared_experts_with_weights( + topk_output.topk_ids, + topk_output.topk_weights, + shared_weights, + self.num_fused_shared_experts, + N=self.num_experts, + ) + ) return StandardTopKOutput( topk_weights=fused_topk_weights, topk_ids=fused_topk_ids, diff --git a/test/registered/amd/accuracy/mi35x/test_qwen35_mxfp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_qwen35_mxfp4_eval_mi35x.py new file mode 100644 index 000000000..458dab6a2 --- /dev/null +++ b/test/registered/amd/accuracy/mi35x/test_qwen35_mxfp4_eval_mi35x.py @@ -0,0 +1,249 @@ +"""MI35x Qwen3.5-397B-A17B MXFP4 GSM8K accuracy + serving-perf test (nightly). + +Covers the two AMD MXFP4 checkpoints, which differ ONLY in whether the shared +expert is quantized -- and therefore in whether the shared-expert fusion +``fuse_gate`` serving path (qwen2_moe.py ``_use_aiter`` branch + the fuse_gate +append kernel) actually runs: + + * ``Qwen3.5-397B-A17B-MXFP4`` -- shared expert EXCLUDED from quant (BF16), + so ``can_fuse_shared_expert`` returns False -> fusion OFF -> eager path. + This is the reference-accuracy checkpoint (gate: gsm8k > 0.91). + * ``Qwen3.5-397B-A17B-MoE-MXFP4`` -- shared expert IS quantized, so fusion is + ON and the aiter ``fuse_gate`` GEMV-in-kernel path runs end-to-end. No + per-PR CI exercises this path (the kernel unit test only checks the kernel + in isolation); here we require it to *run through* cleanly e2e. + +Launch config mirrors ~/run_qwen3.5_mxfp4_perf.sh (tp=2, aiter + unified attn + +flydsl, allreduce fusion). MXFP4 is MI35x-only, so this is not registered on +mi30x. + +Registry: nightly-amd-accuracy-8-gpu-mi35x-qwen35-mxfp4 suite +""" + +import os +import unittest +from types import SimpleNamespace +from typing import List + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.nightly_bench_utils import BenchmarkResult +from sglang.test.nightly_utils import NightlyBenchmarkRunner +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + _parse_int_list_env, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) + +register_amd_ci( + est_time=14400, suite="nightly-amd-accuracy-8-gpu-mi35x-qwen35-mxfp4", nightly=True +) + +# Local checkpoints on the MI35x runner; overridable for other hosts. +MXFP4_MODEL_PATH = os.environ.get( + "QWEN35_MXFP4_MODEL_PATH", "/data/amd/Qwen3.5-397B-A17B-MXFP4" +) +MOE_MXFP4_MODEL_PATH = os.environ.get( + "QWEN35_MOE_MXFP4_MODEL_PATH", "/data/amd/Qwen3.5-397B-A17B-MoE-MXFP4" +) + +SERVER_LAUNCH_TIMEOUT = 3600 +BENCH_TIMEOUT = 5400 + +# Accuracy gate for the reference (non-fused, BF16-shared-expert) checkpoint. +MXFP4_ACC_THRESHOLD = 0.91 + +# Server flags from ~/run_qwen3.5_mxfp4_perf.sh (the reference dev config). +COMMON_ARGS = [ + "--tp", + "2", + "--attention-backend", + "aiter", + "--trust-remote-code", + "--chunked-prefill-size", + "32768", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + "--watchdog-timeout", + "1200", + "--mem-fraction-static", + "0.9", + "--disable-radix-cache", + "--enable-aiter-allreduce-fusion", + "--max-running-requests", + "512", + "--page-size", + "16", +] + +COMMON_ENV = { + "SGLANG_USE_AITER": "1", + "SGLANG_USE_AITER_UNIFIED_ATTN": "1", + "AITER_FLYDSL_FORCE": "1", +} + + +def _run_gsm8k( + base_url: str, model: str, num_examples: int, max_tokens: int = 2048 +) -> dict: + """Few-shot GSM8K against a running server; returns run_eval metrics. + + ``metrics["score"]`` is the accuracy in [0, 1]. + """ + requests.get(base_url + "/flush_cache") + args = SimpleNamespace( + base_url=base_url, + model=model, + eval_name="gsm8k", + api="completion", + max_tokens=max_tokens, + num_examples=num_examples, + num_threads=256, + ) + return run_eval(args) + + +def _generate_perf_report(results: List[BenchmarkResult]) -> str: + """Compact markdown perf table (skips a leading warmup duplicate).""" + header = results[0].model_path + if results[0].run_name and results[0].run_name != "default": + header += f" ({results[0].run_name})" + header += f" [{os.getenv('GPU_CONFIG', 'MI35x')}]" + + summary = f"### {header}\n" + summary += "| batch size | input len | latency (s) | input tput (tok/s) | output tput (tok/s) | ITL (ms) |\n" + summary += "| ---------- | --------- | ----------- | ------------------ | ------------------- | -------- |\n" + + report = ( + results[1:] + if len(results) > 1 and results[0].batch_size == results[1].batch_size + else results + ) + for r in report: + itl = 1 / (r.output_throughput / r.batch_size) * 1000 + summary += ( + f"| {r.batch_size} | {r.input_len} | {r.latency:.2f} | " + f"{r.input_throughput:.2f} | {r.output_throughput:.2f} | {itl:.2f} |\n" + ) + return summary + + +def _run_perf(model_path: str, variant: str, base_url: str) -> None: + """Run the nightly serving benchmark for ``model_path`` with COMMON_ENV set.""" + batch_sizes = _parse_int_list_env("NIGHTLY_BATCH_SIZES", "1,8,16,64") + input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096")) + output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512")) + + runner = NightlyBenchmarkRunner( + f"performance_profiles_{variant}", variant, base_url + ) + runner.setup_profile_directory() + runner.full_report = f"## {variant}\n" + + old_env = {} + for key, value in COMMON_ENV.items(): + old_env[key] = os.environ.get(key) + os.environ[key] = value + try: + results, success = runner.run_benchmark_for_model( + model_path=model_path, + batch_sizes=batch_sizes, + input_lens=input_lens, + output_lens=output_lens, + other_args=COMMON_ARGS, + variant=variant, + extra_bench_args=["--trust-remote-code"], + enable_profile=False, + timeout=BENCH_TIMEOUT, + )[:2] + if results: + runner.full_report += _generate_perf_report(results) + "\n" + assert success, f"Perf benchmark failed for {model_path} on MI35x" + finally: + for key, value in old_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + runner.write_final_report() + + +class TestQwen35Mxfp4MI35x(CustomTestCase): + """Non-fused reference checkpoint (shared expert BF16): accuracy-gated.""" + + base_url = DEFAULT_URL_FOR_TEST + + def test_a_gsm8k(self): + """GSM8K accuracy must clear the reference gate (> 0.91).""" + env = os.environ.copy() + env.update(COMMON_ENV) + process = popen_launch_server( + MXFP4_MODEL_PATH, + self.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=COMMON_ARGS, + env=env, + ) + try: + metrics = _run_gsm8k(self.base_url, MXFP4_MODEL_PATH, num_examples=1319) + print(f"[{MXFP4_MODEL_PATH}] {metrics=}") + if is_in_ci(): + write_github_step_summary( + f"### gsm8k accuracy ({MXFP4_MODEL_PATH})\n" + f'score={metrics["score"]:.3f} ' + f"(threshold {MXFP4_ACC_THRESHOLD})\n" + ) + self.assertGreater(metrics["score"], MXFP4_ACC_THRESHOLD) + finally: + kill_process_tree(process.pid) + + def test_b_perf(self): + """Serving performance benchmark.""" + _run_perf(MXFP4_MODEL_PATH, "qwen35-mxfp4-mi35x", self.base_url) + + +class TestQwen35MoeMxfp4MI35x(CustomTestCase): + """Fused checkpoint (shared expert MXFP4 -> fuse_gate path): run-through smoke.""" + + base_url = DEFAULT_URL_FOR_TEST + + def test_a_gsm8k_runthrough(self): + """Exercise the aiter fuse_gate serving path e2e; require it to run and + produce valid (non-degenerate) output. No strict accuracy gate -- the + point is that the fuse_gate path serves cleanly end-to-end (the coverage + gap no per-PR CI test fills).""" + env = os.environ.copy() + env.update(COMMON_ENV) + process = popen_launch_server( + MOE_MXFP4_MODEL_PATH, + self.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=COMMON_ARGS, + env=env, + ) + try: + metrics = _run_gsm8k(self.base_url, MOE_MXFP4_MODEL_PATH, num_examples=200) + print(f"[{MOE_MXFP4_MODEL_PATH}] {metrics=}") + if is_in_ci(): + write_github_step_summary( + f"### gsm8k run-through ({MOE_MXFP4_MODEL_PATH}, fuse_gate)\n" + f'score={metrics["score"]:.3f} (run-through, no gate)\n' + ) + # Ran e2e and returned parseable answers -> fuse_gate path is healthy. + self.assertGreater(metrics["score"], 0.0) + finally: + kill_process_tree(process.pid) + + def test_b_perf(self): + """Serving performance benchmark for the fused checkpoint.""" + _run_perf(MOE_MXFP4_MODEL_PATH, "qwen35-moe-mxfp4-mi35x", self.base_url) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/moe/test_fused_append_shared_experts.py b/test/registered/moe/test_fused_append_shared_experts.py new file mode 100644 index 000000000..2edc166a6 --- /dev/null +++ b/test/registered/moe/test_fused_append_shared_experts.py @@ -0,0 +1,301 @@ +"""Unit tests for ``fused_append_shared_experts_with_weights``. + +Covers PR #28658, which folds the shared-expert ``sigmoid`` activation and the +bf16->fp32 cast into the append kernel's prologue (``apply_sigmoid=True``), and +PR #28666, which additionally folds the ``shared_expert_gate`` GEMV into the +kernel (``fuse_gate=True``). The old (``apply_sigmoid=False, fuse_gate=False``) +path must stay byte-for-byte identical; the ``apply_sigmoid`` path must equal +the eager ``sigmoid(logits.float()) * scale`` it replaces; and the ``fuse_gate`` +path must equal the eager ``sigmoid((hidden @ W_gate).float()) * scale`` GEMV it +replaces. These paths only run on the AITER shared-expert-fusion route at +serving time, so they are otherwise uncovered by CI. +""" + +import unittest + +import torch + +from sglang.kernels.ops.moe.fused_moe_triton_kernels import ( + fused_append_shared_experts_with_weights, +) +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase, empty_gpu_cache + +register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-large") +register_amd_ci(est_time=20, suite="stage-b-test-1-gpu-small-amd") + + +def _eager_append( + topk_ids, topk_weights, shared_weights, s, n_base, apply_sigmoid, scale +): + """Reference: eager equivalent of the kernel (matches the pre-PR two-step path). + + The last ``s`` columns hold the shared expert(s); ids are ``n_base + i`` and + weights are either the raw shared value cast to the output dtype (legacy) or + ``sigmoid(value.float()) * scale`` cast to the output dtype (fused). + """ + m, k = topk_ids.shape + sw = shared_weights + if sw.ndim == 1: + sw = sw.unsqueeze(-1) + if sw.shape[1] < s: + sw = sw.expand(m, s) + if apply_sigmoid: + col = (torch.sigmoid(sw.float()) * scale).to(topk_weights.dtype) + else: + col = sw.to(topk_weights.dtype) + + shared_ids = ( + (n_base + torch.arange(s, device=topk_ids.device)) + .view(1, s) + .expand(m, s) + .to(topk_ids.dtype) + ) + + out_ids = torch.cat([topk_ids, shared_ids], dim=1) + out_weights = torch.cat([topk_weights, col.contiguous()], dim=1) + return out_ids, out_weights + + +def _eager_gate(topk_ids, topk_weights, hidden, gate_weight, s, n_base, scale): + """Reference for fuse_gate=True: in-kernel GEMV + sigmoid + scale. + + ``logit[m] = hidden[m, :] . W_gate[:]`` (fp32), broadcast + ``sigmoid(logit) * scale`` to all ``s`` shared slots. + """ + m, k = topk_ids.shape + logit = (hidden.float() @ gate_weight.float().reshape(-1)).view(m, 1) + col = (torch.sigmoid(logit) * scale).expand(m, s).to(topk_weights.dtype) + + shared_ids = ( + (n_base + torch.arange(s, device=topk_ids.device)) + .view(1, s) + .expand(m, s) + .to(topk_ids.dtype) + ) + + out_ids = torch.cat([topk_ids, shared_ids], dim=1) + out_weights = torch.cat([topk_weights, col.contiguous()], dim=1) + return out_ids, out_weights + + +class TestFusedAppendSharedExperts(CustomTestCase): + MS = [1, 3, 128] + KS = [2, 6, 8] + SS = [1, 2] + SCALES = [1.0, 0.5] + + @staticmethod + def _rand_inputs(m, k, s, dtype, n_base): + device = get_device() + topk_ids = torch.randint(0, n_base, (m, k), dtype=torch.int32, device=device) + topk_weights = torch.rand(m, k, dtype=torch.float32, device=device).to(dtype) + # shared gate logits: wide range so sigmoid is meaningfully exercised + shared_logits = ( + torch.randn(m, s, dtype=torch.float32, device=device) * 4.0 + ).to(dtype) + return topk_ids, topk_weights, shared_logits + + def _check(self, out_ids, out_w, ref_ids, ref_w, dtype, msg): + # ids are integers -> must match exactly + self.assertTrue(torch.equal(out_ids, ref_ids), f"ids mismatch: {msg}") + # triton tl.sigmoid vs torch.sigmoid differ by a few ULPs, so the fused + # activation path is only bitexact-free: tolerance, tight in fp32. + if dtype == torch.float32: + torch.testing.assert_close( + out_w, ref_w, rtol=1e-5, atol=1e-6, msg=f"weights mismatch: {msg}" + ) + else: + torch.testing.assert_close( + out_w, ref_w, rtol=2e-3, atol=2e-3, msg=f"weights mismatch: {msg}" + ) + + def test_fused_sigmoid_matches_eager(self): + """apply_sigmoid=True must equal eager sigmoid(logits.float())*scale.""" + n_base = 64 + for dtype in [torch.float32, torch.bfloat16]: + for m in self.MS: + for k in self.KS: + for s in self.SS: + for scale in self.SCALES: + topk_ids, topk_w, logits = self._rand_inputs( + m, k, s, dtype, n_base + ) + out_ids, out_w = fused_append_shared_experts_with_weights( + topk_ids, + topk_w, + logits, + s, + N=n_base, + apply_sigmoid=True, + scale=scale, + ) + ref_ids, ref_w = _eager_append( + topk_ids, topk_w, logits, s, n_base, True, scale + ) + self.assertEqual(out_ids.shape, (m, k + s)) + self._check( + out_ids, + out_w, + ref_ids, + ref_w, + dtype, + f"m={m} k={k} s={s} scale={scale} dtype={dtype}", + ) + empty_gpu_cache() + + def test_legacy_path_unchanged(self): + """apply_sigmoid=False must reproduce the legacy cast-and-append exactly.""" + n_base = 64 + for dtype in [torch.float32, torch.bfloat16]: + for m in self.MS: + for k in self.KS: + for s in self.SS: + topk_ids, topk_w, shared = self._rand_inputs( + m, k, s, dtype, n_base + ) + out_ids, out_w = fused_append_shared_experts_with_weights( + topk_ids, topk_w, shared, s, N=n_base, apply_sigmoid=False + ) + ref_ids, ref_w = _eager_append( + topk_ids, topk_w, shared, s, n_base, False, 1.0 + ) + # legacy path is a pure copy/cast -> exact for all dtypes + self.assertTrue(torch.equal(out_ids, ref_ids)) + self.assertTrue(torch.equal(out_w, ref_w)) + empty_gpu_cache() + + def test_routed_columns_preserved(self): + """First k columns must be an untouched copy of the routed topk output.""" + n_base = 32 + m, k, s = 5, 6, 1 + topk_ids, topk_w, logits = self._rand_inputs(m, k, s, torch.bfloat16, n_base) + out_ids, out_w = fused_append_shared_experts_with_weights( + topk_ids, topk_w, logits, s, N=n_base, apply_sigmoid=True, scale=1.0 + ) + self.assertTrue(torch.equal(out_ids[:, :k], topk_ids)) + self.assertTrue(torch.equal(out_w[:, :k], topk_w)) + # shared ids are exactly n_base .. n_base+s-1 + self.assertTrue( + torch.equal( + out_ids[:, k], + torch.full((m,), n_base, dtype=topk_ids.dtype, device=topk_ids.device), + ) + ) + + def test_fuse_gate_matches_eager(self): + """fuse_gate=True must equal eager sigmoid((hidden @ W_gate).float())*scale.""" + n_base = 64 + device = get_device() + for dtype in [torch.float32, torch.bfloat16]: + for hidden_dim in [512, 2048]: + for m in self.MS: + for k in self.KS: + for s in self.SS: + for scale in self.SCALES: + topk_ids = torch.randint( + 0, n_base, (m, k), dtype=torch.int32, device=device + ) + topk_w = torch.rand( + m, k, dtype=torch.float32, device=device + ).to(dtype) + hidden = ( + torch.randn( + m, + hidden_dim, + dtype=torch.float32, + device=device, + ) + * 0.1 + ).to(dtype) + # shared_expert_gate weight: Linear(hidden, 1) + gate_w = ( + torch.randn( + 1, + hidden_dim, + dtype=torch.float32, + device=device, + ) + * 0.1 + ).to(dtype) + + out_ids, out_w = ( + fused_append_shared_experts_with_weights( + topk_ids, + topk_w, + None, + s, + N=n_base, + fuse_gate=True, + hidden_states=hidden, + gate_weight=gate_w, + scale=scale, + ) + ) + ref_ids, ref_w = _eager_gate( + topk_ids, topk_w, hidden, gate_w, s, n_base, scale + ) + self.assertEqual(out_ids.shape, (m, k + s)) + msg = ( + f"m={m} k={k} s={s} h={hidden_dim} " + f"scale={scale} dtype={dtype}" + ) + self.assertTrue( + torch.equal(out_ids, ref_ids), f"ids: {msg}" + ) + # fp32 GEMV reduction order differs from torch + # matmul -> tolerance (tight in fp32). + if dtype == torch.float32: + torch.testing.assert_close( + out_w, ref_w, rtol=1e-4, atol=1e-5, msg=msg + ) + else: + torch.testing.assert_close( + out_w, ref_w, rtol=2e-3, atol=2e-3, msg=msg + ) + empty_gpu_cache() + + def test_fuse_gate_and_apply_sigmoid_mutually_exclusive(self): + """fuse_gate already applies sigmoid; combining with apply_sigmoid asserts.""" + device = get_device() + topk_ids = torch.randint(0, 8, (2, 2), dtype=torch.int32, device=device) + topk_w = torch.rand(2, 2, dtype=torch.float32, device=device) + hidden = torch.randn(2, 16, dtype=torch.float32, device=device) + gate_w = torch.randn(1, 16, dtype=torch.float32, device=device) + with self.assertRaises(AssertionError): + fused_append_shared_experts_with_weights( + topk_ids, + topk_w, + None, + 1, + N=8, + apply_sigmoid=True, + fuse_gate=True, + hidden_states=hidden, + gate_weight=gate_w, + ) + + def test_fuse_gate_requires_hidden_and_gate(self): + """fuse_gate=True without hidden_states/gate_weight asserts.""" + device = get_device() + topk_ids = torch.randint(0, 8, (2, 2), dtype=torch.int32, device=device) + topk_w = torch.rand(2, 2, dtype=torch.float32, device=device) + with self.assertRaises(AssertionError): + fused_append_shared_experts_with_weights( + topk_ids, topk_w, None, 1, N=8, fuse_gate=True + ) + + def test_zero_shared_is_noop(self): + """num_fused_shared_experts <= 0 returns the inputs unchanged.""" + n_base = 8 + topk_ids, topk_w, logits = self._rand_inputs(4, 2, 1, torch.float32, n_base) + out_ids, out_w = fused_append_shared_experts_with_weights( + topk_ids, topk_w, logits, 0, N=n_base, apply_sigmoid=True + ) + self.assertIs(out_ids, topk_ids) + self.assertIs(out_w, topk_w) + + +if __name__ == "__main__": + unittest.main()