Fix examples and pytest, run ruff (#3230)
Co-authored-by: dePaul Miller <23461061+depaulmillz@users.noreply.github.com>
This commit is contained in:
dePaul Miller
co-authored by
dePaul Miller
dePaul Miller
parent
982cb9e718
commit
546c3efa89
@@ -26,6 +26,7 @@
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
config.default_SMs[__file__] = "90a"
|
||||
config.addinivalue_line(
|
||||
|
||||
@@ -55,7 +55,7 @@ Coverage
|
||||
|
||||
import pytest
|
||||
import cutlass
|
||||
from hopper.dense_gemm_fp8_2xacc import run
|
||||
from hopper.kernel.dense_gemm.dense_gemm_fp8_2xacc import run
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Type aliases
|
||||
@@ -169,8 +169,8 @@ def _run_benchmark(
|
||||
[
|
||||
pytest.param((128, 256), (2048, 2048, 2048, 1), id="tile128x256"),
|
||||
pytest.param((128, 128), (2048, 2048, 2048, 1), id="tile128x128"),
|
||||
pytest.param((128, 64), (2048, 2048, 2048, 1), id="tile128x64"),
|
||||
pytest.param((64, 64), (2048, 2048, 2048, 1), id="tile64x64"),
|
||||
pytest.param((128, 64), (2048, 2048, 2048, 1), id="tile128x64"),
|
||||
pytest.param((64, 64), (2048, 2048, 2048, 1), id="tile64x64"),
|
||||
],
|
||||
)
|
||||
def test_l0_tile_shapes(tile_shape_mn, mnkl):
|
||||
@@ -195,8 +195,11 @@ def test_l0_tile_shapes(tile_shape_mn, mnkl):
|
||||
)
|
||||
def test_l0_cluster_shapes(cluster_shape_mn):
|
||||
"""All valid cluster shapes compile (tile 128x128, 2048^3)."""
|
||||
_run_compile(mnkl=(2048, 2048, 2048, 1), tile_shape_mn=(128, 128),
|
||||
cluster_shape_mn=cluster_shape_mn)
|
||||
_run_compile(
|
||||
mnkl=(2048, 2048, 2048, 1),
|
||||
tile_shape_mn=(128, 128),
|
||||
cluster_shape_mn=cluster_shape_mn,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -208,8 +211,8 @@ def test_l0_cluster_shapes(cluster_shape_mn):
|
||||
@pytest.mark.parametrize(
|
||||
"c_dtype",
|
||||
[
|
||||
pytest.param(F16, id="Float16"),
|
||||
pytest.param(F32, id="Float32"),
|
||||
pytest.param(F16, id="Float16"),
|
||||
pytest.param(F32, id="Float32"),
|
||||
pytest.param(F8E4, id="Float8E4M3FN"),
|
||||
],
|
||||
)
|
||||
@@ -227,8 +230,8 @@ def test_l0_output_dtypes(c_dtype):
|
||||
@pytest.mark.parametrize(
|
||||
"mma_promotion_interval",
|
||||
[
|
||||
pytest.param(4, id="interval4"),
|
||||
pytest.param(8, id="interval8"),
|
||||
pytest.param(4, id="interval4"),
|
||||
pytest.param(8, id="interval8"),
|
||||
pytest.param(16, id="interval16"),
|
||||
],
|
||||
)
|
||||
@@ -249,8 +252,8 @@ def test_l0_mma_promotion_intervals(mma_promotion_interval):
|
||||
[
|
||||
pytest.param((128, 256), (2048, 2048, 2048, 1), id="tile128x256"),
|
||||
pytest.param((128, 128), (2048, 2048, 2048, 1), id="tile128x128"),
|
||||
pytest.param((128, 64), (2048, 2048, 2048, 1), id="tile128x64"),
|
||||
pytest.param((64, 64), (2048, 2048, 2048, 1), id="tile64x64"),
|
||||
pytest.param((128, 64), (2048, 2048, 2048, 1), id="tile128x64"),
|
||||
pytest.param((64, 64), (2048, 2048, 2048, 1), id="tile64x64"),
|
||||
],
|
||||
)
|
||||
def test_l1_tile_shapes(tile_shape_mn, mnkl):
|
||||
@@ -276,8 +279,11 @@ def test_l1_tile_shapes(tile_shape_mn, mnkl):
|
||||
)
|
||||
def test_l1_cluster_shapes(cluster_shape_mn):
|
||||
"""All cluster shapes (including A/B multicast paths) produce correct results."""
|
||||
_run_correctness(mnkl=(2048, 2048, 2048, 1), tile_shape_mn=(128, 128),
|
||||
cluster_shape_mn=cluster_shape_mn)
|
||||
_run_correctness(
|
||||
mnkl=(2048, 2048, 2048, 1),
|
||||
tile_shape_mn=(128, 128),
|
||||
cluster_shape_mn=cluster_shape_mn,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -290,9 +296,9 @@ def test_l1_cluster_shapes(cluster_shape_mn):
|
||||
@pytest.mark.parametrize(
|
||||
"c_dtype, tolerance",
|
||||
[
|
||||
pytest.param(F16, 0.1, id="Float16"),
|
||||
pytest.param(F32, 0.1, id="Float32"),
|
||||
pytest.param(F8E4, 0.5, id="Float8E4M3FN"),
|
||||
pytest.param(F16, 0.1, id="Float16"),
|
||||
pytest.param(F32, 0.1, id="Float32"),
|
||||
pytest.param(F8E4, 0.5, id="Float8E4M3FN"),
|
||||
],
|
||||
)
|
||||
def test_l1_output_dtypes(c_dtype, tolerance):
|
||||
@@ -310,8 +316,8 @@ def test_l1_output_dtypes(c_dtype, tolerance):
|
||||
@pytest.mark.parametrize(
|
||||
"mma_promotion_interval",
|
||||
[
|
||||
pytest.param(4, id="interval4"),
|
||||
pytest.param(8, id="interval8"),
|
||||
pytest.param(4, id="interval4"),
|
||||
pytest.param(8, id="interval8"),
|
||||
pytest.param(16, id="interval16"),
|
||||
],
|
||||
)
|
||||
@@ -330,9 +336,9 @@ def test_l1_mma_promotion_intervals(mma_promotion_interval):
|
||||
@pytest.mark.parametrize(
|
||||
"scale_a_val, scale_b_val",
|
||||
[
|
||||
pytest.param(0.5, 2.0, id="scale_a0.5_b2.0"),
|
||||
pytest.param(0.25, 4.0, id="scale_a0.25_b4.0"),
|
||||
pytest.param(2.0, 0.5, id="scale_a2.0_b0.5"),
|
||||
pytest.param(0.5, 2.0, id="scale_a0.5_b2.0"),
|
||||
pytest.param(0.25, 4.0, id="scale_a0.25_b4.0"),
|
||||
pytest.param(2.0, 0.5, id="scale_a2.0_b0.5"),
|
||||
],
|
||||
)
|
||||
def test_l1_scale_factors(scale_a_val, scale_b_val):
|
||||
@@ -351,7 +357,7 @@ def test_l1_scale_factors(scale_a_val, scale_b_val):
|
||||
"mnkl",
|
||||
[
|
||||
pytest.param((1024, 1024, 1024, 2), id="L2"),
|
||||
pytest.param((512, 512, 512, 4), id="L4"),
|
||||
pytest.param((512, 512, 512, 4), id="L4"),
|
||||
],
|
||||
)
|
||||
def test_l1_batched(mnkl):
|
||||
@@ -370,25 +376,118 @@ def test_l1_batched(mnkl):
|
||||
"mnkl, tile_shape_mn, cluster_shape_mn, mma_promotion_interval, label",
|
||||
[
|
||||
# Square 4096^3 — tile / cluster sweep
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 128), (1, 1), 4, "4096^3 tile=128x128 cluster=1x1", id="4096-128x128-1x1"),
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 128), (1, 2), 4, "4096^3 tile=128x128 cluster=1x2", id="4096-128x128-1x2"),
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 128), (2, 2), 4, "4096^3 tile=128x128 cluster=2x2", id="4096-128x128-2x2"),
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 256), (1, 2), 4, "4096^3 tile=128x256 cluster=1x2", id="4096-128x256-1x2"),
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 256), (2, 2), 4, "4096^3 tile=128x256 cluster=2x2", id="4096-128x256-2x2"),
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 64), (1, 2), 4, "4096^3 tile=128x64 cluster=1x2", id="4096-128x64-1x2"),
|
||||
pytest.param((4096, 4096, 4096, 1), (64, 64), (1, 2), 4, "4096^3 tile=64x64 cluster=1x2", id="4096-64x64-1x2"),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 128),
|
||||
(1, 1),
|
||||
4,
|
||||
"4096^3 tile=128x128 cluster=1x1",
|
||||
id="4096-128x128-1x1",
|
||||
),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 128),
|
||||
(1, 2),
|
||||
4,
|
||||
"4096^3 tile=128x128 cluster=1x2",
|
||||
id="4096-128x128-1x2",
|
||||
),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 128),
|
||||
(2, 2),
|
||||
4,
|
||||
"4096^3 tile=128x128 cluster=2x2",
|
||||
id="4096-128x128-2x2",
|
||||
),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 256),
|
||||
(1, 2),
|
||||
4,
|
||||
"4096^3 tile=128x256 cluster=1x2",
|
||||
id="4096-128x256-1x2",
|
||||
),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 256),
|
||||
(2, 2),
|
||||
4,
|
||||
"4096^3 tile=128x256 cluster=2x2",
|
||||
id="4096-128x256-2x2",
|
||||
),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 64),
|
||||
(1, 2),
|
||||
4,
|
||||
"4096^3 tile=128x64 cluster=1x2",
|
||||
id="4096-128x64-1x2",
|
||||
),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(64, 64),
|
||||
(1, 2),
|
||||
4,
|
||||
"4096^3 tile=64x64 cluster=1x2",
|
||||
id="4096-64x64-1x2",
|
||||
),
|
||||
# LLM-like: 8192x8192x4096
|
||||
pytest.param((8192, 8192, 4096, 1), (128, 128), (1, 2), 4, "8192x8192x4096 tile=128x128 cluster=1x2", id="llm-128x128-1x2"),
|
||||
pytest.param((8192, 8192, 4096, 1), (128, 256), (2, 2), 4, "8192x8192x4096 tile=128x256 cluster=2x2", id="llm-128x256-2x2"),
|
||||
pytest.param(
|
||||
(8192, 8192, 4096, 1),
|
||||
(128, 128),
|
||||
(1, 2),
|
||||
4,
|
||||
"8192x8192x4096 tile=128x128 cluster=1x2",
|
||||
id="llm-128x128-1x2",
|
||||
),
|
||||
pytest.param(
|
||||
(8192, 8192, 4096, 1),
|
||||
(128, 256),
|
||||
(2, 2),
|
||||
4,
|
||||
"8192x8192x4096 tile=128x256 cluster=2x2",
|
||||
id="llm-128x256-2x2",
|
||||
),
|
||||
# mma_promotion_interval sweep (shows precision/performance trade-off)
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 128), (1, 2), 4, "4096^3 interval=4", id="4096-interval4"),
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 128), (1, 2), 8, "4096^3 interval=8", id="4096-interval8"),
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 128), (1, 2), 16, "4096^3 interval=16", id="4096-interval16"),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 128),
|
||||
(1, 2),
|
||||
4,
|
||||
"4096^3 interval=4",
|
||||
id="4096-interval4",
|
||||
),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 128),
|
||||
(1, 2),
|
||||
8,
|
||||
"4096^3 interval=8",
|
||||
id="4096-interval8",
|
||||
),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 128),
|
||||
(1, 2),
|
||||
16,
|
||||
"4096^3 interval=16",
|
||||
id="4096-interval16",
|
||||
),
|
||||
# FP8 output
|
||||
pytest.param((4096, 4096, 4096, 1), (128, 128), (1, 2), 4, "4096^3 out=FP8E4M3", id="4096-fp8-out"),
|
||||
pytest.param(
|
||||
(4096, 4096, 4096, 1),
|
||||
(128, 128),
|
||||
(1, 2),
|
||||
4,
|
||||
"4096^3 out=FP8E4M3",
|
||||
id="4096-fp8-out",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_bench(mnkl, tile_shape_mn, cluster_shape_mn, mma_promotion_interval, label, capsys):
|
||||
def test_bench(
|
||||
mnkl, tile_shape_mn, cluster_shape_mn, mma_promotion_interval, label, capsys
|
||||
):
|
||||
"""
|
||||
Performance benchmark — run with: pytest -m bench -s
|
||||
|
||||
|
||||
@@ -59,7 +59,7 @@ os.environ.setdefault("GROUPED_GEMM_FORCE_CUTE_COPY", "0")
|
||||
import pytest
|
||||
import cutlass
|
||||
import cutlass.utils as utils
|
||||
from hopper.grouped_gemm import run
|
||||
from hopper.kernel.grouped_gemm.grouped_gemm import run
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -203,8 +203,8 @@ def _run_case(
|
||||
[
|
||||
pytest.param((128, 256), [(128, 256, 64, 1)], id="tile128x256"),
|
||||
pytest.param((128, 128), [(128, 128, 64, 1)], id="tile128x128"),
|
||||
pytest.param((128, 64), [(128, 64, 64, 1)], id="tile128x64"),
|
||||
pytest.param((64, 64), [(64, 64, 64, 1)], id="tile64x64"),
|
||||
pytest.param((128, 64), [(128, 64, 64, 1)], id="tile128x64"),
|
||||
pytest.param((64, 64), [(64, 64, 64, 1)], id="tile64x64"),
|
||||
],
|
||||
)
|
||||
def test_l0_tile_shapes(tile_shape_mn, problem_sizes_mnkl, tmap_mode):
|
||||
@@ -222,15 +222,20 @@ def test_l0_tile_shapes(tile_shape_mn, problem_sizes_mnkl, tmap_mode):
|
||||
@pytest.mark.parametrize(
|
||||
"num_groups, problem_sizes_mnkl",
|
||||
[
|
||||
pytest.param(2, [(128, 256, 64, 1)] * 2, id="2g-uniform"),
|
||||
pytest.param(4, [(128, 256, 64, 1), (64, 128, 64, 1),
|
||||
(256, 128, 64, 1), (192, 256, 64, 1)], id="4g-mixed"),
|
||||
pytest.param(8, [(128, 256, 64, 1)] * 8, id="8g-uniform"),
|
||||
pytest.param(2, [(128, 256, 64, 1)] * 2, id="2g-uniform"),
|
||||
pytest.param(
|
||||
4,
|
||||
[(128, 256, 64, 1), (64, 128, 64, 1), (256, 128, 64, 1), (192, 256, 64, 1)],
|
||||
id="4g-mixed",
|
||||
),
|
||||
pytest.param(8, [(128, 256, 64, 1)] * 8, id="8g-uniform"),
|
||||
],
|
||||
)
|
||||
def test_l0_group_counts(num_groups, problem_sizes_mnkl, tmap_mode):
|
||||
"""Various group counts compile for tile (128,256) fp16."""
|
||||
_run_compile(num_groups, problem_sizes_mnkl, (128, 256), tensormap_update_mode=tmap_mode)
|
||||
_run_compile(
|
||||
num_groups, problem_sizes_mnkl, (128, 256), tensormap_update_mode=tmap_mode
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -244,28 +249,43 @@ def test_l0_group_counts(num_groups, problem_sizes_mnkl, tmap_mode):
|
||||
"a_dtype, b_dtype, c_dtype, acc_dtype, problem_sizes_mnkl",
|
||||
[
|
||||
# fp16 → fp16 output
|
||||
pytest.param(F16, F16, F16, F32, [(128, 256, 64, 1)], id="fp16-fp16-fp16-fp32"),
|
||||
pytest.param(F16, F16, F16, F32, [(128, 256, 64, 1)], id="fp16-fp16-fp16-fp32"),
|
||||
# fp16 → fp32 output
|
||||
pytest.param(F16, F16, F32, F32, [(128, 256, 64, 1)], id="fp16-fp16-fp32-fp32"),
|
||||
pytest.param(F16, F16, F32, F32, [(128, 256, 64, 1)], id="fp16-fp16-fp32-fp32"),
|
||||
# fp16 with fp16 accumulator
|
||||
pytest.param(F16, F16, F16, F16, [(128, 256, 64, 1)], id="fp16-fp16-fp16-fp16"),
|
||||
pytest.param(F16, F16, F16, F16, [(128, 256, 64, 1)], id="fp16-fp16-fp16-fp16"),
|
||||
# fp8 E4M3 → fp16 output (K must be multiple of 16 for fp8 alignment)
|
||||
pytest.param(F8E4, F8E4, F16, F32, [(128, 256, 128, 1)], id="fp8e4-fp8e4-fp16-fp32"),
|
||||
pytest.param(
|
||||
F8E4, F8E4, F16, F32, [(128, 256, 128, 1)], id="fp8e4-fp8e4-fp16-fp32"
|
||||
),
|
||||
# fp8 E5M2 → fp16 output
|
||||
pytest.param(F8E5, F8E5, F16, F32, [(128, 256, 128, 1)], id="fp8e5-fp8e5-fp16-fp32"),
|
||||
pytest.param(
|
||||
F8E5, F8E5, F16, F32, [(128, 256, 128, 1)], id="fp8e5-fp8e5-fp16-fp32"
|
||||
),
|
||||
# mixed fp8: E4M3 × E5M2
|
||||
pytest.param(F8E4, F8E5, F16, F32, [(128, 256, 128, 1)], id="fp8e4-fp8e5-fp16-fp32"),
|
||||
pytest.param(
|
||||
F8E4, F8E5, F16, F32, [(128, 256, 128, 1)], id="fp8e4-fp8e5-fp16-fp32"
|
||||
),
|
||||
# int8 → int32 output (K must be multiple of 16)
|
||||
pytest.param(I8, I8, I32, I32, [(128, 256, 128, 1)], id="int8-int8-int32-int32"),
|
||||
pytest.param(
|
||||
I8, I8, I32, I32, [(128, 256, 128, 1)], id="int8-int8-int32-int32"
|
||||
),
|
||||
# uint8 → int32 output
|
||||
pytest.param(U8, U8, I32, I32, [(128, 256, 128, 1)], id="uint8-uint8-int32-int32"),
|
||||
pytest.param(
|
||||
U8, U8, I32, I32, [(128, 256, 128, 1)], id="uint8-uint8-int32-int32"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_l0_dtypes(a_dtype, b_dtype, c_dtype, acc_dtype, problem_sizes_mnkl, tmap_mode):
|
||||
"""Data type combinations compile for tile (128,256)."""
|
||||
_run_compile(
|
||||
1, problem_sizes_mnkl, (128, 256),
|
||||
a_dtype=a_dtype, b_dtype=b_dtype, c_dtype=c_dtype, acc_dtype=acc_dtype,
|
||||
1,
|
||||
problem_sizes_mnkl,
|
||||
(128, 256),
|
||||
a_dtype=a_dtype,
|
||||
b_dtype=b_dtype,
|
||||
c_dtype=c_dtype,
|
||||
acc_dtype=acc_dtype,
|
||||
tensormap_update_mode=tmap_mode,
|
||||
)
|
||||
|
||||
@@ -289,14 +309,22 @@ def test_l0_dtypes(a_dtype, b_dtype, c_dtype, acc_dtype, problem_sizes_mnkl, tma
|
||||
# m-major C output (M must be multiple of 8)
|
||||
pytest.param("k", "k", "m", [(128, 256, 64, 1)], (128, 256), id="akm-bkm-cmaj"),
|
||||
# m-major A + n-major B
|
||||
pytest.param("m", "n", "n", [(128, 128, 64, 1)], (128, 128), id="amaj-bnmaj-cn"),
|
||||
pytest.param(
|
||||
"m", "n", "n", [(128, 128, 64, 1)], (128, 128), id="amaj-bnmaj-cn"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_l0_major_modes(a_major, b_major, c_major, problem_sizes_mnkl, tile_shape_mn, tmap_mode):
|
||||
def test_l0_major_modes(
|
||||
a_major, b_major, c_major, problem_sizes_mnkl, tile_shape_mn, tmap_mode
|
||||
):
|
||||
"""Matrix major mode combinations compile."""
|
||||
_run_compile(
|
||||
1, problem_sizes_mnkl, tile_shape_mn,
|
||||
a_major=a_major, b_major=b_major, c_major=c_major,
|
||||
1,
|
||||
problem_sizes_mnkl,
|
||||
tile_shape_mn,
|
||||
a_major=a_major,
|
||||
b_major=b_major,
|
||||
c_major=c_major,
|
||||
tensormap_update_mode=tmap_mode,
|
||||
)
|
||||
|
||||
@@ -321,10 +349,14 @@ def test_l0_major_modes(a_major, b_major, c_major, problem_sizes_mnkl, tile_shap
|
||||
pytest.param((2, 2), [(256, 512, 64, 1)], (128, 256), id="cluster2x2"),
|
||||
],
|
||||
)
|
||||
def test_l0_cluster_shapes(cluster_shape_mn, problem_sizes_mnkl, tile_shape_mn, tmap_mode):
|
||||
def test_l0_cluster_shapes(
|
||||
cluster_shape_mn, problem_sizes_mnkl, tile_shape_mn, tmap_mode
|
||||
):
|
||||
"""Cluster shapes including multicast paths compile."""
|
||||
_run_compile(
|
||||
1, problem_sizes_mnkl, tile_shape_mn,
|
||||
1,
|
||||
problem_sizes_mnkl,
|
||||
tile_shape_mn,
|
||||
cluster_shape_mn=cluster_shape_mn,
|
||||
tensormap_update_mode=tmap_mode,
|
||||
)
|
||||
@@ -341,18 +373,20 @@ def test_l0_cluster_shapes(cluster_shape_mn, problem_sizes_mnkl, tile_shape_mn,
|
||||
"num_groups, problem_sizes_mnkl",
|
||||
[
|
||||
# groups with very different shapes
|
||||
pytest.param(4, [(64, 64, 64, 1),
|
||||
(128, 128, 64, 1),
|
||||
(256, 128, 64, 1),
|
||||
(128, 256, 64, 1)], id="4g-all-tiles"),
|
||||
pytest.param(
|
||||
4,
|
||||
[(64, 64, 64, 1), (128, 128, 64, 1), (256, 128, 64, 1), (128, 256, 64, 1)],
|
||||
id="4g-all-tiles",
|
||||
),
|
||||
# tiny vs large
|
||||
pytest.param(2, [(64, 64, 64, 1),
|
||||
(512, 512, 64, 1)], id="2g-tiny-large"),
|
||||
pytest.param(2, [(64, 64, 64, 1), (512, 512, 64, 1)], id="2g-tiny-large"),
|
||||
],
|
||||
)
|
||||
def test_l0_mixed_problem_sizes(num_groups, problem_sizes_mnkl, tmap_mode):
|
||||
"""Heterogeneous per-group problem sizes compile."""
|
||||
_run_compile(num_groups, problem_sizes_mnkl, (128, 256), tensormap_update_mode=tmap_mode)
|
||||
_run_compile(
|
||||
num_groups, problem_sizes_mnkl, (128, 256), tensormap_update_mode=tmap_mode
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -386,13 +420,15 @@ def test_l1_fp16_4g_mixed(tmap_mode):
|
||||
[
|
||||
pytest.param((128, 256), [(128, 256, 64, 1)], id="tile128x256"),
|
||||
pytest.param((128, 128), [(128, 128, 64, 1)], id="tile128x128"),
|
||||
pytest.param((128, 64), [(128, 64, 64, 1)], id="tile128x64"),
|
||||
pytest.param((64, 64), [(64, 64, 64, 1)], id="tile64x64"),
|
||||
pytest.param((128, 64), [(128, 64, 64, 1)], id="tile128x64"),
|
||||
pytest.param((64, 64), [(64, 64, 64, 1)], id="tile64x64"),
|
||||
],
|
||||
)
|
||||
def test_l1_tile_shapes_fp16(tile_shape_mn, problem_sizes_mnkl, tmap_mode):
|
||||
"""All tile shapes produce correct results."""
|
||||
_run_correctness(1, problem_sizes_mnkl, tile_shape_mn, tensormap_update_mode=tmap_mode)
|
||||
_run_correctness(
|
||||
1, problem_sizes_mnkl, tile_shape_mn, tensormap_update_mode=tmap_mode
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -429,8 +465,11 @@ def test_l1_group_count_scaling(num_groups, tmap_mode):
|
||||
def test_l1_fp16_c_fp32(tmap_mode):
|
||||
"""fp16 inputs with fp32 output are numerically correct."""
|
||||
_run_correctness(
|
||||
1, [(128, 256, 64, 1)], (128, 256),
|
||||
c_dtype=F32, acc_dtype=F32,
|
||||
1,
|
||||
[(128, 256, 64, 1)],
|
||||
(128, 256),
|
||||
c_dtype=F32,
|
||||
acc_dtype=F32,
|
||||
tensormap_update_mode=tmap_mode,
|
||||
)
|
||||
|
||||
@@ -441,8 +480,13 @@ def test_l1_fp16_c_fp32(tmap_mode):
|
||||
def test_l1_fp8_e4m3(tmap_mode):
|
||||
"""fp8 E4M3FN inputs are numerically correct (K=128 for 16B alignment)."""
|
||||
_run_correctness(
|
||||
1, [(128, 256, 128, 1)], (128, 256),
|
||||
a_dtype=F8E4, b_dtype=F8E4, c_dtype=F16, acc_dtype=F32,
|
||||
1,
|
||||
[(128, 256, 128, 1)],
|
||||
(128, 256),
|
||||
a_dtype=F8E4,
|
||||
b_dtype=F8E4,
|
||||
c_dtype=F16,
|
||||
acc_dtype=F32,
|
||||
tensormap_update_mode=tmap_mode,
|
||||
tolerance=0.5,
|
||||
)
|
||||
@@ -454,8 +498,13 @@ def test_l1_fp8_e4m3(tmap_mode):
|
||||
def test_l1_fp8_mixed(tmap_mode):
|
||||
"""Mixed fp8 inputs (E4M3 × E5M2) are numerically correct."""
|
||||
_run_correctness(
|
||||
1, [(128, 256, 128, 1)], (128, 256),
|
||||
a_dtype=F8E4, b_dtype=F8E5, c_dtype=F16, acc_dtype=F32,
|
||||
1,
|
||||
[(128, 256, 128, 1)],
|
||||
(128, 256),
|
||||
a_dtype=F8E4,
|
||||
b_dtype=F8E5,
|
||||
c_dtype=F16,
|
||||
acc_dtype=F32,
|
||||
tensormap_update_mode=tmap_mode,
|
||||
tolerance=0.5,
|
||||
)
|
||||
@@ -467,8 +516,13 @@ def test_l1_fp8_mixed(tmap_mode):
|
||||
def test_l1_int8(tmap_mode):
|
||||
"""int8 inputs with int32 accumulator are correct."""
|
||||
_run_correctness(
|
||||
1, [(128, 256, 128, 1)], (128, 256),
|
||||
a_dtype=I8, b_dtype=I8, c_dtype=I32, acc_dtype=I32,
|
||||
1,
|
||||
[(128, 256, 128, 1)],
|
||||
(128, 256),
|
||||
a_dtype=I8,
|
||||
b_dtype=I8,
|
||||
c_dtype=I32,
|
||||
acc_dtype=I32,
|
||||
tensormap_update_mode=tmap_mode,
|
||||
tolerance=0,
|
||||
)
|
||||
@@ -485,7 +539,9 @@ def test_l1_int8(tmap_mode):
|
||||
def test_l1_c_m_major(tmap_mode):
|
||||
"""m-major C output is correct."""
|
||||
_run_correctness(
|
||||
1, [(128, 256, 64, 1)], (128, 256),
|
||||
1,
|
||||
[(128, 256, 64, 1)],
|
||||
(128, 256),
|
||||
c_major="m",
|
||||
tensormap_update_mode=tmap_mode,
|
||||
)
|
||||
@@ -498,8 +554,12 @@ def test_l1_c_m_major(tmap_mode):
|
||||
def test_l1_all_non_default_majors(tmap_mode):
|
||||
"""m-major A, n-major B, m-major C together are correct."""
|
||||
_run_correctness(
|
||||
1, [(64, 64, 64, 1)], (128, 128),
|
||||
a_major="m", b_major="n", c_major="m",
|
||||
1,
|
||||
[(64, 64, 64, 1)],
|
||||
(128, 128),
|
||||
a_major="m",
|
||||
b_major="n",
|
||||
c_major="m",
|
||||
tensormap_update_mode=tmap_mode,
|
||||
)
|
||||
|
||||
@@ -521,7 +581,9 @@ def test_l1_all_non_default_majors(tmap_mode):
|
||||
def test_l1_cluster_shapes(cluster_shape_mn, problem_sizes_mnkl, tmap_mode):
|
||||
"""Multicast cluster shapes produce correct results."""
|
||||
_run_correctness(
|
||||
1, problem_sizes_mnkl, (128, 256),
|
||||
1,
|
||||
problem_sizes_mnkl,
|
||||
(128, 256),
|
||||
cluster_shape_mn=cluster_shape_mn,
|
||||
tensormap_update_mode=tmap_mode,
|
||||
)
|
||||
@@ -540,14 +602,14 @@ def test_l1_8g_mixed_sizes(tmap_mode):
|
||||
_run_correctness(
|
||||
8,
|
||||
[
|
||||
(128, 256, 64, 1),
|
||||
(64, 128, 64, 1),
|
||||
(256, 128, 64, 1),
|
||||
(128, 256, 64, 1),
|
||||
(64, 128, 64, 1),
|
||||
(256, 128, 64, 1),
|
||||
(128, 128, 128, 1),
|
||||
(192, 256, 64, 1),
|
||||
(64, 64, 64, 1),
|
||||
(192, 256, 64, 1),
|
||||
(64, 64, 64, 1),
|
||||
(128, 256, 128, 1),
|
||||
(256, 256, 64, 1),
|
||||
(256, 256, 64, 1),
|
||||
],
|
||||
(128, 256),
|
||||
tensormap_update_mode=tmap_mode,
|
||||
|
||||
Reference in New Issue
Block a user