[AMD] Fuse shared_expert_gate GEMV into the MoE append kernel (HIP/aiter) (#28666)
Co-authored-by: sogalin_codegen <39478626+sogalin@users.noreply.github.com>
This commit is contained in:
jacky.cheng
co-authored by
sogalin_codegen
sogalin_codegen
parent
0a6bbbe128
commit
240a12b302
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user