From 2d863b835404ba6a9d2971dfeb50f32c310c9429 Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Mon, 7 Sep 2026 14:26:02 +0800 Subject: [PATCH 1/6] perf(gfx1250): label SMI replays per benchmark call --- aiter/test_common.py | 81 ++++++++++++++++++++++- op_tests/bench_gfx1250_combo.py | 102 +++++++---------------------- op_tests/test_pa_sparse_prefill.py | 6 +- 3 files changed, 105 insertions(+), 84 deletions(-) diff --git a/aiter/test_common.py b/aiter/test_common.py index bf3061e27b..e3afd6ca85 100644 --- a/aiter/test_common.py +++ b/aiter/test_common.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import copy +import inspect import json import multiprocessing as mp import os +from contextvars import ContextVar +from enum import Enum import numpy as np import pandas as pd @@ -14,6 +17,7 @@ pd.set_option("display.max_rows", 200) _SMI_LABEL_COUNTS = {} +_SMI_CALL_LABEL = ContextVar("aiter_smi_call_label", default=None) ## debug ## # pd.set_option("display.max_rows", None) # pd.set_option("display.max_columns", None) @@ -46,6 +50,67 @@ def print_json_table(name, rows, keep=None): print(json.dumps({"name": name, "rows": records}), flush=True) +def _smi_label_value(value): + """Return a compact, stable label value, or None for opaque arguments.""" + if isinstance(value, torch.Tensor): + shape = "x".join(map(str, value.shape)) or "scalar" + return f"{shape}:{str(value.dtype).removeprefix('torch.')}" + if isinstance(value, torch.dtype): + return str(value).removeprefix("torch.") + if isinstance(value, Enum): + return str(value.value) + if value is None or isinstance(value, (str, bool, int, float, np.generic)): + return str(value) + if isinstance(value, (tuple, list)): + items = [_smi_label_value(item) for item in value] + if all(item is not None for item in items): + return ",".join(items) + return None + + +def _smi_call_tag(func, callargs): + """Build a call-local SMI label from @benchmark's named arguments.""" + source = os.path.splitext(os.path.basename(func.__code__.co_filename))[0] + parts = [f"{source}.{func.__name__}"] + aliases = {"m": "M", "n": "N", "k": "K", "t": "T", "h": "H", "d": "D"} + for name, value in callargs.items(): + formatted = _smi_label_value(value) + if formatted is None: + continue + formatted = formatted.replace("/", "_").replace("\n", "") + parts.append(f"{aliases.get(name, name)}={formatted}") + return "/".join(parts) + + +def _smi_perftest_tag(func, args, kwargs): + """Return scalar call details that distinguish one perftest invocation.""" + try: + signature = inspect.signature(func) + callargs = signature.bind(*args, **kwargs) + callargs.apply_defaults() + except (TypeError, ValueError): + return None + + parts = [] + for name, parameter in signature.parameters.items(): + if parameter.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + continue + value = callargs.arguments.get(name) + # The outer @benchmark call already supplies tensor-derived shapes. + # Keep this suffix for lightweight selectors such as a backend name. + if value is None or isinstance(value, torch.Tensor) or callable(value): + continue + formatted = _smi_label_value(value) + if formatted is None: + continue + formatted = formatted.replace("/", "_").replace("\n", "") + parts.append(f"{name}={formatted}") + return "/".join(parts) or None + + def ensure_spawn_method(): """ Ensure multiprocessing uses 'spawn' start method. @@ -186,7 +251,12 @@ def replay(): replay_us = avg - case_label = os.environ.get("AITER_SMI_LABEL", "benchmark_case") + case_label = _SMI_CALL_LABEL.get() or os.environ.get( + "AITER_SMI_LABEL", "benchmark_case" + ) + perftest_tag = _smi_perftest_tag(func, args, kwargs) + if perftest_tag: + case_label = f"{case_label}/{perftest_tag}" label_key = (case_label, fn_name) occurrence = _SMI_LABEL_COUNTS.get(label_key, 0) + 1 _SMI_LABEL_COUNTS[label_key] = occurrence @@ -208,7 +278,14 @@ def benchmark(): def decorator(func): def wrapper(*args, **kwargs): callargs = log_args(func, *args, **kwargs) - ret = func(*args, **kwargs) + token = None + if os.environ.get("AITER_SMI_MONITOR", "0") == "1": + token = _SMI_CALL_LABEL.set(_smi_call_tag(func, callargs)) + try: + ret = func(*args, **kwargs) + finally: + if token is not None: + _SMI_CALL_LABEL.reset(token) if ret is not None: callargs.update(ret) return callargs diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index de4a0c734e..d3528ca310 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -328,7 +328,7 @@ def _silence(): @contextlib.contextmanager def _smi_case(label): - """Set the case label consumed by the common perftest hook.""" + """Set the fallback label for calls without @benchmark context.""" old = os.environ.get("AITER_SMI_LABEL") if os.environ.get("AITER_SMI_MONITOR") == "1": os.environ["AITER_SMI_LABEL"] = label @@ -885,8 +885,8 @@ def _run_child(name, cmd, cwd, env=None, extract=None, timeout=None, tail=30, when the child fails or emits nothing recognisable. """ extract = extract or _json_tables - # Give every child invocation its combo-owned SMI case label. UTs remain - # unaware of telemetry; the common perftest hook reads this environment. + # Give calls without @benchmark context a combo-owned fallback label. + # Decorated UT calls derive their per-case label from their actual args. if smi and os.environ.get("AITER_SMI_MONITOR") == "1": env = os.environ.copy() if env is None else env.copy() env["AITER_SMI_LABEL"] = name @@ -1173,25 +1173,12 @@ def run_case(tokens, shapes, init_pairs, label): init_pairs = _init_pairs( args, defaults=(("constant", "constant"), ("uniform", "auto")) ) - if args.smi_monitor: - for m, (n, k), pair in itertools.product( - _A8W8_BLOCKSCALE_TOKENS, nk_shapes, init_pairs - ): - data_init, scale_init = pair - run_case( - (m,), - ((n, k),), - (pair,), - f"a8w8_blockscale/M={m}/N={n}/K={k}/data={data_init}/" - f"scale={scale_init}/seed={args.seed}", - ) - else: - run_case( - _A8W8_BLOCKSCALE_TOKENS, - nk_shapes, - init_pairs, - "gemm_a8w8_blockscale (DSv4)", - ) + run_case( + _A8W8_BLOCKSCALE_TOKENS, + nk_shapes, + init_pairs, + "gemm_a8w8_blockscale (DSv4)", + ) def run_a16w16(args): @@ -1376,15 +1363,7 @@ def run_case(tokens, data_inits, label): ) data_inits = args.data_init or ["norm"] - if args.smi_monitor: - for m, data_init in itertools.product(_TOKENS, data_inits): - run_case( - (m,), - (data_init,), - f"mhc/M={m}/N=7168/fuse_rmsnorm=1/data={data_init}/seed={args.seed}", - ) - else: - run_case(_TOKENS, data_inits, "mhc (DSv4, fused RMSNorm)") + run_case(_TOKENS, data_inits, "mhc (DSv4, fused RMSNorm)") def run_qk_norm(args): @@ -1424,12 +1403,8 @@ def run_case(tokens, data_init): structured=True, ) - if args.smi_monitor: - for token, data_init in itertools.product(_TOKENS, data_inits): - run_case((token,), data_init) - else: - for data_init in data_inits: - run_case(_TOKENS, data_init) + for data_init in data_inits: + run_case(_TOKENS, data_init) def run_score_qk(args): @@ -1743,20 +1718,11 @@ def run_case(tokens, data_inits, label): ) data_inits = args.data_init or ["norm"] - if args.smi_monitor: - for tokens, data_init in itertools.product(_INVERSE_ROPE_TOKENS, data_inits): - run_case( - (tokens,), - (data_init,), - f"inverse_rope/s={tokens}/heads=128/groups=16/layout=n32k4/" - f"group_size=32/data={data_init}/seed={args.seed}", - ) - else: - run_case( - _INVERSE_ROPE_TOKENS, - data_inits, - "inverse_rope_group_quant (DSv4, tp1)", - ) + run_case( + _INVERSE_ROPE_TOKENS, + data_inits, + "inverse_rope_group_quant (DSv4, tp1)", + ) def run_mla_v4_prefill(args): @@ -1800,38 +1766,16 @@ def run_case(tokens, pages, precs, modes, backends, init_values, label): structured=True, ) - if args.smi_monitor: - backend_by_prec = {"fp8": ("opus", "asm"), "bf16": ("opus", "triton")} - for tokens, pages, prec, mode, data_init in itertools.product( - _MLA_PREFILL_TOKENS, + for tokens in _MLA_PREFILL_TOKENS: + run_case( + tokens, (4096, 16384), ("fp8", "bf16"), ("dense", "sparse"), + ("opus", "asm", "triton"), data_inits, - ): - for backend in backend_by_prec[prec]: - run_case( - tokens, - (pages,), - (prec,), - (mode,), - (backend,), - (data_init,), - f"mla_v4_prefill/M={tokens}/H=128/D=512/pages={pages}/" - f"total_tokens={tokens}/prec={prec}/mode={mode}/backend={backend}/" - f"data={data_init}/seed={args.seed}", - ) - else: - for tokens in _MLA_PREFILL_TOKENS: - run_case( - tokens, - (4096, 16384), - ("fp8", "bf16"), - ("dense", "sparse"), - ("opus", "asm", "triton"), - data_inits, - f"mla_v4 prefill (M={tokens}, prec=fp8/bf16, pages=4096/16384)", - ) + f"mla_v4 prefill (M={tokens}, prec=fp8/bf16, pages=4096/16384)", + ) OPS = { diff --git a/op_tests/test_pa_sparse_prefill.py b/op_tests/test_pa_sparse_prefill.py index 345c55a49c..335e3fcda5 100644 --- a/op_tests/test_pa_sparse_prefill.py +++ b/op_tests/test_pa_sparse_prefill.py @@ -491,8 +491,8 @@ def _csr(total_rows: int, seed_offset: int): @perftest() -def _profile_func(target_func, *args, **kwargs): - return target_func(*args, **kwargs) +def _profile_func(target_func, *, backend: str): + return target_func() # --------------------------------------------------------------------------- @@ -641,7 +641,7 @@ def run_pa_sparse_prefill( ) if bench: - _, lat_us = _profile_func(invoke) # (data, avg_us_per_iter) + _, lat_us = _profile_func(invoke, backend=name) flops = 4.0 * h * total_nnz * d tflops = flops / max(lat_us * 1e-6, 1e-12) / 1e12 row[f"{name} us"] = round(float(lat_us), 2) From d9460e8900410b20a41f927fa53ea4dcb307d1cd Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Mon, 7 Sep 2026 16:08:54 +0800 Subject: [PATCH 2/6] perf(gfx1250): cover saturated DSv4 MLA decode KV --- op_tests/bench_gfx1250_combo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index d3528ca310..b671834e58 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -632,13 +632,13 @@ def _int_quad(s): ] _MLA_V4_DSV4_SHAPES = [ (128, 512, kv_seq_lens, num_kv_splits) - for kv_seq_lens in (256, 512, 1024) + for kv_seq_lens in (256, 512, 1024, 1152) for num_kv_splits in (1, 2, 4) ] + [ (128, tokens, kv_seq_lens, num_kv_splits) for tokens in _MLA_DECODE_TOKENS if tokens != 512 - for kv_seq_lens in (256, 512, 1024) + for kv_seq_lens in (256, 512, 1024, 1152) for num_kv_splits in (1, 2, 4) ] _MLA_V4_COMPARE_KEEP = [ From 5a1a745bef6523b6acd6e083950e7976821ec7af Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Mon, 7 Sep 2026 16:20:03 +0800 Subject: [PATCH 3/6] perf(gfx1250): sweep DSv4 score-QK decode batches --- op_tests/bench_gfx1250_combo.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index b671834e58..d5d5585733 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -447,7 +447,9 @@ def _unused_scale_init(args, op): # against rather than in the kernel under test. AITER_BENCH_TOKENS still wins if # set -- what an explicit request sweeps is the caller's business. _INVERSE_ROPE_TOKENS = _tokens((1, 16, 32, 64, 128, 256, 512, 1024, 2048, 16384)) -_SCORE_QK_TOKENS = _tokens() +# Score-QK runs once per decode step. With MTP disabled, this axis is both the +# number of concurrent sequences and the number of query tokens in the launch. +_SCORE_QK_TOKENS = _tokens((1, 16, 32, 64, 128, 256, 512, 1024)) # score_qk is decode, so its KV length is the average context a decode step # scans: input + output/2, then CSA's 4x compression. # 1K in / 1K out -> (1024 + 512) / 4 = 384 @@ -1424,18 +1426,18 @@ def run_score_qk(args): "--blocksize", "64", ] - # None => let the UT pick the batch, so run the KV lengths once each. for tokens, (label, kv_length), data_init in itertools.product( - _SCORE_QK_TOKENS or (None,), + _SCORE_QK_TOKENS, _SCORE_QK_KV_LENGTHS, args.data_init or ["norm"], ): _run_child( - f"score_qk (decode, B={tokens or 'UT default'}, {label} " + f"score_qk (decode, B={tokens}, {label} " f"CSA KV={kv_length}, init={data_init}, seed={args.seed})", [ *base_cmd, - *(["--batch", str(tokens)] if tokens else []), + "--batch", + str(tokens), "-kv_length", kv_length, "--data-init", From bc25ec3d3b9d3582a3d458bb368118f954603df8 Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Mon, 7 Sep 2026 17:33:18 +0800 Subject: [PATCH 4/6] perf(gfx1250): cover TP4 inverse RoPE shape --- op_tests/bench_gfx1250_combo.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index d5d5585733..95f64fd805 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -56,9 +56,9 @@ mla_v4_decode 1..1024. Decode carries one token per sequence, so the axis is really the batch, and 65536 is not a shape the model runs. - inverse_rope 1..16384. The axis is -s at a fixed -b 128,16, and 65536 - faults -- in the triton reference the UT compares against, - not in the kernel under test. + inverse_rope 1..16384. The axis is -s at fixed TP1/TP4 shapes + -b 128,16 32,4, and 65536 faults -- in the triton reference + the UT compares against, not in the kernel under test. mega_moe 1..2048. 65536 cannot allocate its symmetric arena; see _MEGA_MOE_TOKENS. a8w8_blockscale 512..65536. M=512 covers a DSv4 decode batch of 512; @@ -193,11 +193,11 @@ (sparse draws a random nnz per row, dense fills every row) under --seed, so nnz is an outcome, not an input. -The ``inverse_rope`` op runs the tp1 attention-output shape (-b is -(n_local_heads, n_local_groups); 128,16 is V4-Pro at dp/tp1): +The ``inverse_rope`` op runs the TP1 and TP4 attention-output shapes (-b is +(n_local_heads, n_local_groups); 128,16 and 32,4 are V4-Pro at TP1 and TP4): python3 op_tests/test_inverse_rope_group_quant.py \ - -b 128,16 -s -l n32k4 --group-size 32 + -b 128,16 32,4 -s -l n32k4 --group-size 32 The ``a8w8_blockscale`` op runs: @@ -1690,11 +1690,11 @@ def run_mla_v4_decode(args): def run_inverse_rope(args): - """Run DSv4 inverse RoPE + group quant at the tp1 attention-output shape.""" + """Run DSv4 inverse RoPE + group quant at TP1/TP4 attention shapes.""" _unused_scale_init(args, "inverse_rope") - # -b is (n_local_heads, n_local_groups); 128,16 is V4-Pro at dp/tp1. The UT - # defaults to the two smallest configs instead, which never reach the shape - # the model runs, so name it explicitly. + # -b is (n_local_heads, n_local_groups); 128,16 and 32,4 are V4-Pro at TP1 + # and TP4. The UT defaults to the two smallest configs instead, which never + # reach these model shapes, so name them explicitly. def run_case(tokens, data_inits, label): _run_child( label, @@ -1703,6 +1703,7 @@ def run_case(tokens, data_inits, label): "op_tests/test_inverse_rope_group_quant.py", "-b", "128,16", + "32,4", "-s", *map(str, tokens), "-l", @@ -1723,7 +1724,7 @@ def run_case(tokens, data_inits, label): run_case( _INVERSE_ROPE_TOKENS, data_inits, - "inverse_rope_group_quant (DSv4, tp1)", + "inverse_rope_group_quant (DSv4, TP1/TP4)", ) From c370d257f68472d025647f4d64ec84c255d717a5 Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Mon, 7 Sep 2026 18:17:31 +0800 Subject: [PATCH 5/6] perf(gfx1250): add M256 blockscale coverage --- op_tests/bench_gfx1250_combo.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index 95f64fd805..ec5e6bffb3 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -61,7 +61,7 @@ the UT compares against, not in the kernel under test. mega_moe 1..2048. 65536 cannot allocate its symmetric arena; see _MEGA_MOE_TOKENS. - a8w8_blockscale 512..65536. M=512 covers a DSv4 decode batch of 512; + a8w8_blockscale 256..65536. M=256/512 cover DSv4 decode batches; smaller M stays out because of a UT bug; see DSV4_OPS. mla_v4_prefill 1024..16384, the DSv4 prefill chunk. 65536 faults; see _MLA_PREFILL_TOKENS. @@ -202,7 +202,7 @@ The ``a8w8_blockscale`` op runs: python3 op_tests/test_gemm_a8w8_blockscale.py \ - -m 512 \ + -m 256 512 1024 2048 4096 8192 16384 65536 \ -nk 2048,7168 7168,16384 6144,7168 \ 7168,3072 65536,1536 8192,1536 \ --ck_preshuffle True --flydsl @@ -463,18 +463,20 @@ def _unused_scale_init(args, op): # Was unset, which let the UT sweep its own 27-value default down to M=1. Two # reasons to set it. First, M here is the token count of one step, so the small # end of that default is decode batch and the large end is prefill chunk. This -# list retains the model-real decode point M=512, then covers the prefill side +# list retains the model-real decode points M=256/512, then covers the prefill side # up to the 65536 the other DSv4 ops sweep and past the UT default's own ceiling # of 10240. Second, the tiny M are what walk into # the UT bug described at "a8w8_blockscale" below: get_CKGEMM_config retries the # lookup as M -> get_padded_m(gl=0) -> nextPow2, so anything in [1, 16] or # [33, 64] can land on one of #4773's M=16/M=64 gluon rows (gemm_common.cu:13). -# Starting at 512 clears both ranges by a wide margin. +# Starting at 256 clears both ranges by a wide margin. # # Two things remain outside coverage, both worth remembering: decode-side M -# below 512, and the 11 tuned rows that are the only shapes dispatching to +# below 256, and the 11 tuned rows that are the only shapes dispatching to # gluon. This is a way around the UT bug, not a fix for it. -_A8W8_BLOCKSCALE_TOKENS = _tokens((512, 1024, 2048, 4096, 8192, 16384, 65536)) +_A8W8_BLOCKSCALE_TOKENS = _tokens( + (256, 512, 1024, 2048, 4096, 8192, 16384, 65536) +) # Decode carries one token per sequence, so this axis is the batch, not a token # count; past 1024 it stops being a shape the model runs, hence its own default # rather than _TOKENS. AITER_BENCH_TOKENS overrides it like everywhere else. @@ -1130,7 +1132,7 @@ def run_f8gemm(args): def run_a8w8_blockscale(args): - """Run DSv4 FP8 blockscale linear projections at M=512.""" + """Run DSv4 FP8 blockscale linear projections across decode/prefill M.""" # AITER_LOG_MORE=1 is set at module scope for the FlyDSL MoE ops, and a # child started with env=None inherits this process's whole environ. In this # UT that turned a clean sweep into an intermittent HSA memory fault, so @@ -1860,10 +1862,11 @@ def run_case(tokens, pages, precs, modes, backends, init_values, label): # -m 16 -nk 2048,7168 --ck_preshuffle True passes the strided check with # the line untouched, and only adding --flydsl makes it crash. # - # Back in the sweep because _A8W8_BLOCKSCALE_TOKENS now starts at 512, + # Back in the sweep because _A8W8_BLOCKSCALE_TOKENS now starts at 256, # which keeps every shape clear of the problematic tiny-M ranges while - # retaining a real DSv4 decode batch. M=512 was verified above across all - # six (n,k). The previous 1024..65536 sweep was verified on 20260828, + # retaining real DSv4 decode batches. M=512 was verified above across all + # six (n,k); M=256 is also included in the workload sweep. The previous + # 1024..65536 sweep was verified on 20260828, # rocm/fw-bringup:gfx1250-atom--20260827-ubench: 36/36 cases, err=0 on all, # 2207-7003 TFLOPS. That run also clears M=10240, the shape the earlier # sweep faulted on -- more evidence that fault was cross-case state and not From 0bbc66acad4380d0f8df8acd2cbfd7352f4b7a50 Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Mon, 7 Sep 2026 19:08:31 +0800 Subject: [PATCH 6/6] perf(gfx1250): add 64K Mega MoE coverage --- op_tests/bench_gfx1250_combo.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/op_tests/bench_gfx1250_combo.py b/op_tests/bench_gfx1250_combo.py index ec5e6bffb3..d9cf4d73bd 100644 --- a/op_tests/bench_gfx1250_combo.py +++ b/op_tests/bench_gfx1250_combo.py @@ -59,8 +59,8 @@ inverse_rope 1..16384. The axis is -s at fixed TP1/TP4 shapes -b 128,16 32,4, and 65536 faults -- in the triton reference the UT compares against, not in the kernel under test. - mega_moe 1..2048. 65536 cannot allocate its symmetric arena; see - _MEGA_MOE_TOKENS. + mega_moe 1..2048 plus 65536. The 65536 tier is expected to expose + the current cco symmetric-arena limit; see _MEGA_MOE_TOKENS. a8w8_blockscale 256..65536. M=256/512 cover DSv4 decode batches; smaller M stays out because of a UT bug; see DSV4_OPS. mla_v4_prefill 1024..16384, the DSv4 prefill chunk. 65536 faults; see @@ -116,12 +116,12 @@ ``--scale-init`` is reported as not applicable for operators without a scale operand. -mega_moe at tokens/rank=65536 fails in setup(), asking 7.5 GB for cco's VMM -arena against a 4 GiB default. MORI_SHMEM_HEAP_SIZE does not reach that arena -(see run_mega_moe), so exporting it changes nothing -- and exporting it -sweep-wide takes the machine down, because that heap is preallocated per rank -for every case. The tier is out of the sweep; fixing it means passing -per_rank_vmm at Communicator.init(). +mega_moe at tokens/rank=65536 has previously failed in setup(), asking 7.5 GB +for cco's VMM arena against a 4 GiB default. MORI_SHMEM_HEAP_SIZE does not reach +that arena (see run_mega_moe), so exporting it changes nothing -- and exporting +it sweep-wide takes the machine down, because that heap is preallocated per +rank for every case. The tier is included to expose the limit; fixing it means +passing per_rank_vmm at Communicator.init(). Failures do not stop the sweep: a case that aborts is recorded and the run moves to the next one, with a "N failed, M ops selected" list at the end and a @@ -497,12 +497,12 @@ def _unused_scale_init(args, op): # that is fixed these are timings from an unverified kernel -- the same footing # as a16w16's M=65536 rows before _A16W16_MAX_ERR caught them. _MLA_PREFILL_TOKENS = _tokens((1024, 2048, 4096, 8192, 16384)) -# Default stops at 2048: tokens/rank=65536 dies in pipe.setup() building the -# symmetric arena -- cco sizes it from Communicator.DEFAULT_PER_RANK_VMM (4 GiB) +# tokens/rank=65536 has previously died in pipe.setup() while building the +# symmetric arena: cco sizes it from Communicator.DEFAULT_PER_RANK_VMM (4 GiB) # and asks for 7.5 GB. That is a per_rank_vmm the UT never passes, not something -# MORI_SHMEM_HEAP_SIZE reaches, so the tier cannot run from here. Ask for it via -# AITER_BENCH_TOKENS anyway and you get it, along with that failure. -_MEGA_MOE_TOKENS = _tokens((1, 16, 32, 64, 128, 256, 512, 1024, 2048)) +# MORI_SHMEM_HEAP_SIZE reaches. Keep the tier in the sweep so the limitation is +# visible in the structured failure output rather than silently unmeasured. +_MEGA_MOE_TOKENS = _tokens((1, 16, 32, 64, 128, 256, 512, 1024, 2048, 65536)) # What dispatch puts on the wire; combine is always bf16, so anything but bf16 # is an asymmetric pair. fp4 is the wire DSv4 actually serves on -- the receiver # hands the payload straight to the expert GEMM as its A operand, and that GEMM