[Perf] Tune the W4AFP8 DeepEP low-latency requant launch geometry (#35760)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alex Nails
co-authored by
Claude Opus 5
Claude Opus 5
parent
3a0f1a1344
commit
00fbb6e8ac
@@ -0,0 +1,135 @@
|
||||
"""Benchmark the W4AFP8 DeepEP low-latency requant against its previous geometry.
|
||||
|
||||
``legacy-geometry`` launches the current kernel with its previous launch
|
||||
parameters (1024-element tile, 8 warps, 32 programs per expert); ``tuned`` goes
|
||||
through the wrapper. ``skew`` concentrates rows on one hot expert, which the
|
||||
m-grid cannot see because ``expected_m`` is a dispatch-wide average.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.kernels.jit.benchmark import marker
|
||||
from sglang.kernels.ops.moe.ep_moe_kernels import (
|
||||
_fp8_per_token_quant_to_per_tensor_quant_kernel,
|
||||
fp8_per_token_to_per_tensor_quant_triton,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=45, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
FP8 = torch.float8_e4m3fn
|
||||
K_SCALE_BLOCK_SIZE = 128
|
||||
LEGACY_G_BLOCK = 8 # 1024 hidden elements
|
||||
LEGACY_WARPS = 8
|
||||
LEGACY_M_GRID = 32
|
||||
|
||||
|
||||
def _expected_m(num_experts, dispatched_rows):
|
||||
"""``dispatch_a`` reports ``(rows + num_experts) // num_experts``, one high
|
||||
at exact averages; benchmark with what production would pass."""
|
||||
return (dispatched_rows + num_experts) // num_experts
|
||||
|
||||
|
||||
def _row_counts(num_experts, rows, skew):
|
||||
"""One hot expert at ``skew * rows``, the rest share the fixed remainder; a
|
||||
dispatch redistributes rows, so skew cannot exceed ``num_experts``."""
|
||||
if skew == 1:
|
||||
return [rows] * num_experts
|
||||
total = num_experts * rows
|
||||
counts = [(total - rows * skew) // (num_experts - 1)] * num_experts
|
||||
counts[0] = rows * skew
|
||||
return counts
|
||||
|
||||
|
||||
def _build(num_experts, m, k, rows, skew=1):
|
||||
x = (torch.randn(num_experts, m, k, device="cuda") * 4).to(FP8)
|
||||
# DeepEP returns the last two scale dims column-major (for TMA).
|
||||
x_scale = (
|
||||
torch.rand(num_experts, m, k // K_SCALE_BLOCK_SIZE, device="cuda")
|
||||
.add_(0.5)
|
||||
.permute(0, 2, 1)
|
||||
.contiguous()
|
||||
.permute(0, 2, 1)
|
||||
)
|
||||
counts = _row_counts(num_experts, rows, skew)
|
||||
masked_m = torch.tensor(counts, dtype=torch.int32, device="cuda")
|
||||
output_scale = torch.tensor([2.0], dtype=torch.float32, device="cuda")
|
||||
output = torch.empty((num_experts, m, k), dtype=FP8, device="cuda")
|
||||
return (
|
||||
x,
|
||||
x_scale,
|
||||
masked_m,
|
||||
output_scale,
|
||||
output,
|
||||
_expected_m(num_experts, sum(counts)),
|
||||
)
|
||||
|
||||
|
||||
def _tuned(x, x_scale, masked_m, output_scale, output, expected_m):
|
||||
fp8_per_token_to_per_tensor_quant_triton(
|
||||
x=x,
|
||||
x_scale=x_scale,
|
||||
masked_m=masked_m,
|
||||
output_scale=output_scale,
|
||||
output=output,
|
||||
expected_rows=expected_m,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def _legacy_geometry(x, x_scale, masked_m, output_scale, output, expected_m):
|
||||
num_groups = x.size(2) // K_SCALE_BLOCK_SIZE
|
||||
grid = (triton.cdiv(num_groups, LEGACY_G_BLOCK), LEGACY_M_GRID, x.size(0))
|
||||
_fp8_per_token_quant_to_per_tensor_quant_kernel[grid](
|
||||
x,
|
||||
x_scale,
|
||||
*x_scale.stride(),
|
||||
masked_m,
|
||||
output_scale,
|
||||
output,
|
||||
x.size(1),
|
||||
x.size(2),
|
||||
x.size(0),
|
||||
# row_cap = m keeps every row on its own expert, as the old launch did.
|
||||
x.size(1),
|
||||
K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE,
|
||||
G_BLOCK_SIZE=LEGACY_G_BLOCK,
|
||||
HAS_G_TAIL=(num_groups % LEGACY_G_BLOCK != 0),
|
||||
EXPERT_BLOCK=triton.next_power_of_2(x.size(0)),
|
||||
num_warps=LEGACY_WARPS,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
FN_MAP = {"tuned": _tuned, "legacy-geometry": _legacy_geometry}
|
||||
|
||||
# (hidden, local experts, padded rows): DeepSeek-V3 at EP8, then a 3584 hidden
|
||||
# size at a low and a high local-expert count.
|
||||
SHAPES = [(7168, 8, 1024), (3584, 8, 1024), (3584, 56, 256)]
|
||||
|
||||
|
||||
@marker.parametrize("hidden,num_experts,m", SHAPES, [(7168, 8, 1024), (3584, 56, 256)])
|
||||
@marker.parametrize("rows", [8, 32, 128, 256], [8, 32, 256])
|
||||
@marker.parametrize("skew", [1, 4, 16], [1, 16])
|
||||
@marker.benchmark("impl", ["tuned", "legacy-geometry"])
|
||||
def benchmark(hidden: int, num_experts: int, m: int, rows: int, skew: int, impl: str):
|
||||
if skew > num_experts:
|
||||
marker.skip("one expert cannot hold more than the whole dispatch")
|
||||
if rows * skew > m:
|
||||
marker.skip("more live rows than the payload holds")
|
||||
args = _build(num_experts, m, hidden, rows, skew)
|
||||
return marker.do_bench(
|
||||
FN_MAP[impl],
|
||||
input_args=args,
|
||||
graph_clone_args=(0, 1),
|
||||
memory_args=None,
|
||||
# Tensor-size bandwidth would be off by the padding factor.
|
||||
disable_log_bandwidth=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
@@ -1,22 +1,21 @@
|
||||
"""Unit test for ``fp8_per_token_to_per_tensor_quant_triton`` across hidden sizes.
|
||||
"""Unit test for ``fp8_per_token_to_per_tensor_quant_triton``.
|
||||
|
||||
W4AFP8 DeepEP low-latency requantizes the fp8 dispatch payload with this kernel
|
||||
before the first CUTLASS grouped GEMM. The payload's hidden size is only
|
||||
guaranteed to be a multiple of the fp8 scale-group size (128) -- e.g. 3584 for
|
||||
Kimi-K3 -- so the kernel must handle a ``k`` tail that does not fill a whole
|
||||
``K_BLOCK_SIZE`` (1024) block, and must still leave the rows past ``masked_m``
|
||||
untouched.
|
||||
The hidden size is only guaranteed to be a multiple of the scale group (128),
|
||||
rows past ``masked_m`` must stay untouched, and every launch geometry must
|
||||
produce the same bytes.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.kernels.ops.moe.ep_moe_kernels import (
|
||||
_fp8_per_token_quant_to_per_tensor_quant_kernel,
|
||||
fp8_per_token_to_per_tensor_quant_triton,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
dev = "cuda"
|
||||
FP8 = torch.float8_e4m3fn
|
||||
@@ -27,15 +26,17 @@ SENTINEL = 0.375
|
||||
OUTPUT_SCALE = 2.0
|
||||
|
||||
|
||||
def _build(num_experts, m, k, seed):
|
||||
def _build(num_experts, m, k, seed, column_major_scales=False):
|
||||
g = torch.Generator(device="cpu").manual_seed(seed)
|
||||
# Integers in [-8, 8] with power-of-two per-token-group scales keep every
|
||||
# intermediate exactly representable in e4m3, so the reference below matches
|
||||
# bit-for-bit regardless of the rounding mode of the final cast.
|
||||
# Integers in [-8, 8] with power-of-two group scales keep every intermediate
|
||||
# exactly representable in e4m3, so the reference matches bit-for-bit.
|
||||
x = torch.randint(-8, 9, (num_experts, m, k), generator=g).float()
|
||||
exps = torch.randint(-1, 2, (num_experts, m, k // K_SCALE_BLOCK_SIZE), generator=g)
|
||||
x_scale = torch.pow(2.0, exps.float())
|
||||
return x.to(dev).to(FP8), x_scale.to(dev)
|
||||
x_scale = torch.pow(2.0, exps.float()).to(dev)
|
||||
if column_major_scales:
|
||||
# DeepEP returns the last two scale dims column-major (for TMA).
|
||||
x_scale = x_scale.permute(0, 2, 1).contiguous().permute(0, 2, 1)
|
||||
return x.to(dev).to(FP8), x_scale
|
||||
|
||||
|
||||
def _ref(x, x_scale):
|
||||
@@ -43,14 +44,23 @@ def _ref(x, x_scale):
|
||||
return (dequant * (1.0 / OUTPUT_SCALE)).to(FP8)
|
||||
|
||||
|
||||
# 7168: exact multiple of K_BLOCK_SIZE (the DeepSeek-V3 hidden size).
|
||||
# 3584 / 1152: only 128-aligned, so the last k block is partially masked.
|
||||
@pytest.mark.parametrize("k", [7168, 3584, 1152])
|
||||
def test_masked_rows_and_k_tail(k):
|
||||
num_experts, m = 4, 48
|
||||
masked = [0, 1, 17, m]
|
||||
def _assert_output(output, x, x_scale, masked):
|
||||
ref = _ref(x, x_scale)
|
||||
for e, valid in enumerate(masked):
|
||||
torch.testing.assert_close(
|
||||
output[e, :valid].float(), ref[e, :valid].float(), rtol=0, atol=0
|
||||
)
|
||||
# Rows past masked_m must stay as the caller left them.
|
||||
padding = output[e, valid:].float()
|
||||
torch.testing.assert_close(
|
||||
padding, torch.full_like(padding, SENTINEL), rtol=0, atol=0
|
||||
)
|
||||
|
||||
x, x_scale = _build(num_experts, m, k, seed=k)
|
||||
|
||||
def _run_and_check(num_experts, m, k, masked, expected_rows, column_major_scales):
|
||||
x, x_scale = _build(
|
||||
num_experts, m, k, seed=k + num_experts, column_major_scales=column_major_scales
|
||||
)
|
||||
masked_m = torch.tensor(masked, dtype=torch.int32, device=dev)
|
||||
output_scale = torch.tensor([OUTPUT_SCALE], dtype=torch.float32, device=dev)
|
||||
output = torch.full((num_experts, m, k), SENTINEL, device=dev).to(FP8)
|
||||
@@ -61,19 +71,100 @@ def test_masked_rows_and_k_tail(k):
|
||||
masked_m=masked_m,
|
||||
output_scale=output_scale,
|
||||
output=output,
|
||||
expected_rows=expected_rows,
|
||||
)
|
||||
|
||||
ref = _ref(x, x_scale)
|
||||
for e, valid in enumerate(masked):
|
||||
torch.testing.assert_close(
|
||||
output[e, :valid].float(), ref[e, :valid].float(), rtol=0, atol=0
|
||||
)
|
||||
# Padding rows are not part of any expert's GEMM problem size and must
|
||||
# stay as the caller left them.
|
||||
padding = output[e, valid:].float()
|
||||
torch.testing.assert_close(
|
||||
padding, torch.full_like(padding, SENTINEL), rtol=0, atol=0
|
||||
)
|
||||
_assert_output(output, x, x_scale, masked)
|
||||
|
||||
|
||||
# 7168: fills every tile of scale groups (the DeepSeek-V3 hidden size).
|
||||
# 3584 / 1152: only 128-aligned, so the last tile is partially masked.
|
||||
@pytest.mark.parametrize("k", [7168, 3584, 1152])
|
||||
# None: shape-independent grid; 4 and 64: both ends of the m-grid heuristic.
|
||||
@pytest.mark.parametrize("expected_rows", [None, 4, 64])
|
||||
@pytest.mark.parametrize("column_major_scales", [False, True])
|
||||
def test_masked_rows_and_group_tail(k, expected_rows, column_major_scales):
|
||||
_run_and_check(
|
||||
num_experts=4,
|
||||
m=48,
|
||||
k=k,
|
||||
masked=[0, 1, 17, 48],
|
||||
expected_rows=expected_rows,
|
||||
column_major_scales=column_major_scales,
|
||||
)
|
||||
|
||||
|
||||
# Row estimates chosen so the expert-count cap binds: 40 experts cap at 16
|
||||
# programs, 128 at 8; uncapped these would be 32 and 16.
|
||||
@pytest.mark.parametrize("num_experts,expected_rows", [(40, 32), (128, 16)])
|
||||
def test_many_experts(num_experts, expected_rows):
|
||||
m = 48
|
||||
# Every expert gets a different row count, as a real dispatch would.
|
||||
masked = [(e * 7) % (m + 1) for e in range(num_experts)]
|
||||
_run_and_check(
|
||||
num_experts=num_experts,
|
||||
m=m,
|
||||
k=3584,
|
||||
masked=masked,
|
||||
expected_rows=expected_rows,
|
||||
column_major_scales=True,
|
||||
)
|
||||
|
||||
|
||||
# The wrapper only launches the running vendor's tile, so drive the kernel
|
||||
# directly across every width either vendor can pick, plus the degenerate
|
||||
# single-group tile.
|
||||
@pytest.mark.parametrize("g_block", [1, 8, 16, 32])
|
||||
@pytest.mark.parametrize("k", [7168, 3584])
|
||||
@pytest.mark.parametrize("m_grid", [1, 4, 32])
|
||||
# 0 sends every row to the shared overflow path, 48 keeps every row on its own
|
||||
# expert, and 4 splits the batch across both.
|
||||
@pytest.mark.parametrize("row_cap", [0, 4, 48])
|
||||
def test_every_launch_geometry_agrees(g_block, k, m_grid, row_cap):
|
||||
num_experts, m = 4, 48
|
||||
masked = [0, 1, 17, 48]
|
||||
x, x_scale = _build(num_experts, m, k, seed=k + g_block, column_major_scales=True)
|
||||
masked_m = torch.tensor(masked, dtype=torch.int32, device=dev)
|
||||
output_scale = torch.tensor([OUTPUT_SCALE], dtype=torch.float32, device=dev)
|
||||
output = torch.full((num_experts, m, k), SENTINEL, device=dev).to(FP8)
|
||||
|
||||
num_groups = k // K_SCALE_BLOCK_SIZE
|
||||
grid = (triton.cdiv(num_groups, g_block), m_grid, num_experts)
|
||||
_fp8_per_token_quant_to_per_tensor_quant_kernel[grid](
|
||||
x,
|
||||
x_scale,
|
||||
*x_scale.stride(),
|
||||
masked_m,
|
||||
output_scale,
|
||||
output,
|
||||
m,
|
||||
k,
|
||||
num_experts,
|
||||
row_cap,
|
||||
K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE,
|
||||
G_BLOCK_SIZE=g_block,
|
||||
HAS_G_TAIL=(num_groups % g_block != 0),
|
||||
EXPERT_BLOCK=triton.next_power_of_2(num_experts),
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
_assert_output(output, x, x_scale, masked)
|
||||
|
||||
|
||||
# Experts with no live rows must be stepped over by the prefix-sum mapping,
|
||||
# the case most likely to be off by one.
|
||||
@pytest.mark.parametrize(
|
||||
"masked", [[0, 0, 0, 0], [0, 5, 0, 7], [9, 0, 0, 0], [0, 0, 0, 9]]
|
||||
)
|
||||
def test_experts_with_no_rows_are_skipped(masked):
|
||||
_run_and_check(
|
||||
num_experts=4,
|
||||
m=48,
|
||||
k=3584,
|
||||
masked=masked,
|
||||
expected_rows=2,
|
||||
column_major_scales=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
"""CPU tests for the W4AFP8 low-latency requant launch geometry."""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.kernels.ops.moe.ep_moe_kernels import requant_launch_geometry
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
DSV3_GROUPS = 7168 // 128 # 56
|
||||
K3_GROUPS = 3584 // 128 # 28
|
||||
PREVIOUS_FIXED_M_GRID = 32
|
||||
ROW_CAP_SLACK = 2
|
||||
|
||||
|
||||
class TestRequantLaunchGeometry(CustomTestCase):
|
||||
def test_cap_leaves_ordinary_variation_to_the_owning_expert(self):
|
||||
"""Rows below the cap stay on their expert; the shared path costs a
|
||||
lookup per row and only pays under real imbalance."""
|
||||
for expected_rows in (1, 4, 16, 64, 256):
|
||||
_, _, row_cap = requant_launch_geometry(
|
||||
DSV3_GROUPS, 64, expected_rows=expected_rows
|
||||
)
|
||||
self.assertGreaterEqual(row_cap, expected_rows * ROW_CAP_SLACK)
|
||||
|
||||
def test_cap_never_exceeds_the_payload(self):
|
||||
"""A cap past the padded rows would leave the shared path unreachable."""
|
||||
for max_rows in (1, 8, 128):
|
||||
for expected_rows in (1, 64, 4096):
|
||||
_, _, row_cap = requant_launch_geometry(
|
||||
DSV3_GROUPS, 64, expected_rows=expected_rows, max_rows=max_rows
|
||||
)
|
||||
self.assertLessEqual(row_cap, max_rows)
|
||||
|
||||
def test_unknown_row_count_keeps_every_row_with_its_expert(self):
|
||||
"""With no estimate there is nothing to place a cap against."""
|
||||
for num_experts in (8, 56):
|
||||
_, m_grid, row_cap = requant_launch_geometry(
|
||||
K3_GROUPS, num_experts, max_rows=128
|
||||
)
|
||||
self.assertEqual(m_grid, PREVIOUS_FIXED_M_GRID)
|
||||
self.assertEqual(row_cap, 128)
|
||||
|
||||
def test_m_grid_never_exceeds_the_previous_fixed_grid(self):
|
||||
"""The estimate only ever shrinks the grid, so no batch can regress."""
|
||||
for expected_rows in (1, 8, 32, 33, 1024):
|
||||
for num_experts in (8, 56, 256):
|
||||
_, m_grid, _ = requant_launch_geometry(
|
||||
DSV3_GROUPS, num_experts, expected_rows=expected_rows
|
||||
)
|
||||
self.assertLessEqual(m_grid, PREVIOUS_FIXED_M_GRID)
|
||||
|
||||
def test_m_grid_shrinks_as_the_expert_axis_fills_the_grid(self):
|
||||
"""Hundreds of experts already saturate the grid without 32 programs each."""
|
||||
scarce = [
|
||||
requant_launch_geometry(DSV3_GROUPS, num_experts, expected_rows=32)[1]
|
||||
for num_experts in (8, 64, 128, 256)
|
||||
]
|
||||
self.assertEqual(scarce, sorted(scarce, reverse=True))
|
||||
self.assertEqual(scarce[0], PREVIOUS_FIXED_M_GRID)
|
||||
self.assertLess(scarce[-1], scarce[0])
|
||||
|
||||
def test_expert_cap_lifts_once_rows_carry_the_work(self):
|
||||
"""Past the row threshold the extra programs are not just early exits."""
|
||||
for num_experts in (8, 128, 512):
|
||||
_, m_grid, _ = requant_launch_geometry(
|
||||
DSV3_GROUPS, num_experts, expected_rows=1024
|
||||
)
|
||||
self.assertEqual(m_grid, PREVIOUS_FIXED_M_GRID)
|
||||
|
||||
def test_dispatcher_round_up_does_not_bump_the_grid(self):
|
||||
"""dispatch_a reports (rows + num_experts) // num_experts, one high;
|
||||
rounding up as well would double the launch at every power of two."""
|
||||
for rows in (4, 8, 16, 32):
|
||||
exact = requant_launch_geometry(DSV3_GROUPS, 8, expected_rows=rows)[1]
|
||||
reported = requant_launch_geometry(DSV3_GROUPS, 8, expected_rows=rows + 1)[
|
||||
1
|
||||
]
|
||||
self.assertEqual(reported, exact, f"rows={rows}")
|
||||
|
||||
def test_m_grid_is_bounded_and_monotonic(self):
|
||||
previous = 0
|
||||
for expected_rows in range(1, 512):
|
||||
_, m_grid, _ = requant_launch_geometry(
|
||||
K3_GROUPS, 8, expected_rows=expected_rows
|
||||
)
|
||||
# The floor keeps a one-row batch from serializing an expert into
|
||||
# one program (measured several times slower).
|
||||
self.assertGreaterEqual(m_grid, 4)
|
||||
self.assertLessEqual(m_grid, PREVIOUS_FIXED_M_GRID)
|
||||
self.assertGreaterEqual(m_grid, previous)
|
||||
previous = m_grid
|
||||
|
||||
def test_tile_never_exceeds_the_payload(self):
|
||||
"""A 512-wide hidden size is 4 groups; a wider tile would be mostly masked."""
|
||||
for num_groups in (1, 4, 12):
|
||||
for num_experts in (8, 56):
|
||||
g_block, _, _ = requant_launch_geometry(
|
||||
num_groups, num_experts, expected_rows=64
|
||||
)
|
||||
self.assertLessEqual(g_block, num_groups)
|
||||
|
||||
def test_tile_holds_bytes_per_lane_across_warp_widths(self):
|
||||
"""The tuned unit is bytes per lane: 2048 elements at warp 32 must become
|
||||
4096 at warp 64, or a wave64 part gets half the bytes per lane."""
|
||||
for warp_size, want_elems in ((32, 2048), (64, 4096)):
|
||||
for group_size in (64, 128, 256, 512):
|
||||
g_block, _, _ = requant_launch_geometry(
|
||||
num_groups=7168 // group_size,
|
||||
num_experts=56,
|
||||
group_size=group_size,
|
||||
expected_rows=16,
|
||||
warp_size=warp_size,
|
||||
)
|
||||
self.assertEqual(
|
||||
g_block * group_size, want_elems, f"warp_size={warp_size}"
|
||||
)
|
||||
|
||||
def test_few_experts_halve_the_tile_on_either_warp_width(self):
|
||||
"""A grid too small to fill the part buys k-blocks by halving the tile."""
|
||||
for warp_size, want_elems in ((32, 1024), (64, 2048)):
|
||||
g_block, _, _ = requant_launch_geometry(
|
||||
DSV3_GROUPS, 8, expected_rows=16, warp_size=warp_size
|
||||
)
|
||||
self.assertEqual(g_block * 128, want_elems, f"warp_size={warp_size}")
|
||||
|
||||
def test_warp_width_does_not_move_the_m_grid(self):
|
||||
"""The two knobs are independent: the m-grid answers to rows and experts."""
|
||||
for num_experts in (8, 56, 256):
|
||||
for expected_rows in (1, 8, 32, 1024):
|
||||
grids = {
|
||||
requant_launch_geometry(
|
||||
DSV3_GROUPS,
|
||||
num_experts,
|
||||
expected_rows=expected_rows,
|
||||
warp_size=warp_size,
|
||||
)[1]
|
||||
for warp_size in (32, 64)
|
||||
}
|
||||
self.assertEqual(len(grids), 1, f"E={num_experts} rows={expected_rows}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user