submission 528485
josusanmartin · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 306 lines, June 9 Researcher Reciprocity License v1.0.
auto_v9_force_nt_large_e33.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-amd-moe-mxfp4-528485?include=source"interfacepython
Compatibility
measured onAMD Instinct MI355X
declared hardwareAMD Instinct MI355X
architecturesgfx950
dtypesbf16, fp32, fp8_e8m0, int32, mxfp4
Benchmark evidence
1 measurement across 1 GPU, fastest first.
Operation / workload
Hardware
Latency
Rank
Observed
Reported · How evidence levels are derived →
Source and license
sourceavailable
revision digestsha256:27e1a69b2b8a534d6dfa91b12416df4fc5fb4e039f2c14106209c354cb475b75
license declaredunknown
license concludedunknown
authorsjosusanmartin
imported2026-08-15
Kernel source
auto_v9_force_nt_large_e33.py306 lines
import os
import functools
import sys
from pathlib import Path
import pandas as pd
import torch
from aiter import ActivationType, QuantType
import aiter.fused_moe as fused_moe_mod
from aiter.fused_moe import fused_moe
from task import input_t, output_t
DEBUG_FMOE = os.getenv("JOSU_DEBUG_FMOE") == "1"
FORCE_REBUILD_FMOE = DEBUG_FMOE or os.getenv("JOSU_REBUILD_FMOE") == "1"
CFG_VERSION = 'auto_v9_force_nt_large_e33'
FORCE_NT_MODE = 'large_e33'
TARGET_SHAPES = (
(16, 7168, 256, 257, 9),
(128, 7168, 256, 257, 9),
(512, 7168, 256, 257, 9),
(16, 7168, 512, 33, 9),
(128, 7168, 512, 33, 9),
(512, 7168, 512, 33, 9),
(512, 7168, 2048, 33, 9),
)
def _shape_mask(frame: pd.DataFrame, shape: tuple[int, int, int, int, int]) -> pd.Series:
token, model_dim, inter_dim, expert, topk = shape
return (
(frame["token"] == token)
& (frame["model_dim"] == model_dim)
& (frame["inter_dim"] == inter_dim)
& (frame["expert"] == expert)
& (frame["topk"] == topk)
)
def _debug_rows(label: str, frame: pd.DataFrame) -> None:
if not DEBUG_FMOE:
return
cols = [
"token",
"model_dim",
"inter_dim",
"expert",
"topk",
"block_m",
"ksplit",
"run_1stage",
"us",
"kernelName1",
"kernelName2",
"_tag",
]
keep_cols = [col for col in cols if col in frame.columns]
print(f"[best.py] {label}", file=sys.stderr)
for shape in TARGET_SHAPES:
rows = frame.loc[_shape_mask(frame, shape), keep_cols].sort_values("us").head(6)
print(f"[best.py] shape={shape} rows={len(rows)}", file=sys.stderr)
if not rows.empty:
print(rows.to_string(index=False), file=sys.stderr)
def _bootstrap_cfg() -> str:
out_path = Path(f"/tmp/josu_cfg_hybrid_under150_fmoe_{CFG_VERSION}.csv")
if out_path.exists() and out_path.stat().st_size > 0 and not FORCE_REBUILD_FMOE:
return str(out_path)
source_paths = [
Path("/home/runner/aiter/aiter/configs/tuned_fmoe.csv"),
Path(
"/home/runner/aiter/aiter/configs/model_configs/"
"a8w8_blockscale_tuned_fmoe_qwen3_235b.csv"
),
Path("/home/runner/aiter/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv"),
]
frames = [pd.read_csv(path) for path in source_paths if path.exists()]
merged = pd.concat(frames, ignore_index=True)
untuned_path = Path("/home/runner/aiter/aiter/configs/untuned_fmoe.csv")
keys = pd.read_csv(untuned_path, nrows=0).columns.tolist()
if "cu_num" not in keys:
keys.append("cu_num")
defaults = {column: "" for column in merged.columns}
defaults.update(
{
"cu_num": 256,
"act_type": "ActivationType.Silu",
"dtype": "torch.bfloat16",
"q_dtype_a": "torch.float4_e2m1fn_x2",
"q_dtype_w": "torch.float4_e2m1fn_x2",
"q_type": "QuantType.per_1x32",
"use_g1u1": 1,
"doweight_stage1": 0,
"block_m": 32,
"ksplit": 0,
"us1": 0.0,
"kernelName1": "",
"err1": "0.0%",
"us2": 0.0,
"kernelName2": "",
"err2": "0.0%",
"us": 0.0,
"run_1stage": 0,
"tflops": 0.0,
"bw": 0.0,
"_tag": "",
}
)
online_tuned_rows = [{'token': 16, 'model_dim': 7168, 'inter_dim': 256, 'expert': 257, 'topk': 9, 'block_m': 32, 'ksplit': 2, 'us': 0.0, 'us1': 0.0, 'kernelName1': '', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': '', 'err2': '0.0%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 128, 'model_dim': 7168, 'inter_dim': 256, 'expert': 257, 'topk': 9, 'block_m': 32, 'ksplit': 2, 'us': 0.0, 'us1': 0.0, 'kernelName1': '', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': '', 'err2': '0.0%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 512, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 0, 'us': 0.0, 'us1': 0.0, 'kernelName1': 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': 'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'err2': '2.9%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 512, 'model_dim': 7168, 'inter_dim': 2048, 'expert': 33, 'topk': 9, 'block_m': 128, 'ksplit': 0, 'us': 0.0, 'us1': 0.0, 'kernelName1': 'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': 'moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'err2': '7.1%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}]
k2_overrides = []
default_2stage_rows = [{'token': 16, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 2}, {'token': 128, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 2}]
rows = []
for row in online_tuned_rows:
seeded = defaults.copy()
seeded.update(row)
seeded["us"] = -1.0
rows.append(seeded)
for token, model_dim, inter_dim, expert, topk in k2_overrides:
mask = (
(merged["cu_num"] == 256)
& (merged["token"] == token)
& (merged["model_dim"] == model_dim)
& (merged["inter_dim"] == inter_dim)
& (merged["expert"] == expert)
& (merged["topk"] == topk)
& (merged["act_type"] == "ActivationType.Silu")
& (merged["dtype"] == "torch.bfloat16")
& (merged["q_dtype_a"] == "torch.float4_e2m1fn_x2")
& (merged["q_dtype_w"] == "torch.float4_e2m1fn_x2")
& (merged["q_type"] == "QuantType.per_1x32")
& (merged["use_g1u1"] == 1)
& (merged["doweight_stage1"] == 0)
)
if not mask.any():
continue
row = merged.loc[mask].sort_values("us").iloc[0].to_dict()
kernel_name_1 = row.get("kernelName1")
kernel_name_2 = row.get("kernelName2")
if pd.isna(kernel_name_1) or pd.isna(kernel_name_2) or not kernel_name_1 or not kernel_name_2:
continue
row["ksplit"] = 2
row["us"] = -1.0
rows.append(row)
for row in default_2stage_rows:
seeded = defaults.copy()
seeded.update(row)
seeded["us"] = -1.0
rows.append(seeded)
merged = pd.concat([merged, pd.DataFrame(rows)], ignore_index=True)
_debug_rows("pre-dedup", merged)
merged = (
merged.sort_values("us")
.drop_duplicates(subset=keys, keep="first")
.reset_index(drop=True)
)
_debug_rows("post-dedup", merged)
out_path.parent.mkdir(parents=True, exist_ok=True)
merged.to_csv(out_path, index=False)
return str(out_path)
os.environ["AITER_CONFIG_FMOE"] = _bootstrap_cfg()
def _should_force_nt(token: int, inter_dim: int, expert: int) -> bool:
if FORCE_NT_MODE == "off":
return False
if FORCE_NT_MODE == "all":
return True
if FORCE_NT_MODE == "e33":
return expert == 33
if FORCE_NT_MODE == "large":
return token >= 512
if FORCE_NT_MODE == "large_e33":
return expert == 33 and token >= 512
if FORCE_NT_MODE == "257_only":
return expert == 257
raise ValueError("Unknown FORCE_NT_MODE=" + str(FORCE_NT_MODE))
def _install_force_nt_patch() -> None:
if FORCE_NT_MODE == "off":
return
if getattr(fused_moe_mod, "_josu_force_nt_patch", False):
return
os.environ["AITER_USE_NT"] = "1"
original_get_2stage_cfgs = fused_moe_mod.get_2stage_cfgs
def wrapped_get_2stage_cfgs(
token,
model_dim,
inter_dim,
expert,
topk,
dtype,
q_dtype_a,
q_dtype_w,
q_type,
use_g1u1,
activation,
doweight_stage1,
hidden_pad,
intermediate_pad,
is_shuffled=True,
):
meta = original_get_2stage_cfgs(
token,
model_dim,
inter_dim,
expert,
topk,
dtype,
q_dtype_a,
q_dtype_w,
q_type,
use_g1u1,
activation,
doweight_stage1,
hidden_pad,
intermediate_pad,
is_shuffled,
)
if not _should_force_nt(token, inter_dim, expert):
return meta
stage1 = meta.stage1
stage2 = meta.stage2
if isinstance(stage1, functools.partial) and stage1.func is fused_moe_mod.ck_moe_stage1:
stage1_kwargs = dict(stage1.keywords or dict())
stage1_kwargs["use_non_temporal_load"] = True
stage1 = functools.partial(stage1.func, *(stage1.args or ()), **stage1_kwargs)
if (
isinstance(stage2, functools.partial)
and stage2.func is fused_moe_mod.aiter.ck_moe_stage2_fwd
):
stage2_kwargs = dict(stage2.keywords or dict())
stage2_kwargs["use_non_temporal_load"] = True
stage2 = functools.partial(stage2.func, *(stage2.args or ()), **stage2_kwargs)
return fused_moe_mod.MOEMetadata(
stage1,
stage2,
int(meta.block_m),
int(meta.ksplit),
meta.run_1stage,
meta.has_bias,
True,
)
fused_moe_mod.get_2stage_cfgs = wrapped_get_2stage_cfgs
fused_moe_mod._josu_force_nt_patch = True
_install_force_nt_patch()
def custom_kernel(data: input_t) -> output_t:
(
hidden_states,
gate_up_weight,
down_weight,
gate_up_weight_scale,
down_weight_scale,
gate_up_weight_shuffled,
down_weight_shuffled,
gate_up_weight_scale_shuffled,
down_weight_scale_shuffled,
topk_weights,
topk_ids,
config,
) = data
hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
intermediate_pad = config["d_expert_pad"] - config["d_expert"]
return fused_moe(
hidden_states,
gate_up_weight_shuffled,
down_weight_shuffled,
topk_weights,
topk_ids,
expert_mask=None,
activation=ActivationType.Silu,
quant_type=QuantType.per_1x32,
doweight_stage1=False,
w1_scale=gate_up_weight_scale_shuffled,
w2_scale=down_weight_scale_shuffled,
a1_scale=None,
a2_scale=None,
hidden_pad=hidden_pad,
intermediate_pad=intermediate_pad,
)
scrolls · 306 lines total
Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0
Changes from previous submission
Against this author's previous submission submission 528351.
import os+ import functoolsimport sysfrom pathlib import Pathimport pandas as pdimport torchfrom aiter import ActivationType, QuantType+ import aiter.fused_moe as fused_moe_modfrom aiter.fused_moe import fused_moefrom task import input_t, output_tDEBUG_FMOE = os.getenv("JOSU_DEBUG_FMOE") == "1"FORCE_REBUILD_FMOE = DEBUG_FMOE or os.getenv("JOSU_REBUILD_FMOE") == "1"- CFG_VERSION = 'auto_baseline_v5'+ CFG_VERSION = 'auto_v9_force_nt_large_e33'+ FORCE_NT_MODE = 'large_e33'TARGET_SHAPES = ((16, 7168, 256, 257, 9),(128, 7168, 256, 257, 9),⋯ 91 unchanged lines})- online_tuned_rows = [{'token': 512, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 0, 'us': 0.0, 'us1': 0.0, 'kernelName1': 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': 'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'err2': '2.9%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 512, 'model_dim': 7168, 'inter_dim': 2048, 'expert': 33, 'topk': 9, 'block_m': 128, 'ksplit': 0, 'us': 0.0, 'us1': 0.0, 'kernelName1': 'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': 'moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'err2': '7.1%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}]- k2_overrides = [(16, 7168, 256, 257, 9), (128, 7168, 256, 257, 9)]+ online_tuned_rows = [{'token': 16, 'model_dim': 7168, 'inter_dim': 256, 'expert': 257, 'topk': 9, 'block_m': 32, 'ksplit': 2, 'us': 0.0, 'us1': 0.0, 'kernelName1': '', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': '', 'err2': '0.0%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 128, 'model_dim': 7168, 'inter_dim': 256, 'expert': 257, 'topk': 9, 'block_m': 32, 'ksplit': 2, 'us': 0.0, 'us1': 0.0, 'kernelName1': '', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': '', 'err2': '0.0%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 512, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 0, 'us': 0.0, 'us1': 0.0, 'kernelName1': 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': 'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'err2': '2.9%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 512, 'model_dim': 7168, 'inter_dim': 2048, 'expert': 33, 'topk': 9, 'block_m': 128, 'ksplit': 0, 'us': 0.0, 'us1': 0.0, 'kernelName1': 'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': 'moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'err2': '7.1%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}]+ k2_overrides = []default_2stage_rows = [{'token': 16, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 2}, {'token': 128, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 2}]rows = []⋯ 55 unchanged linesos.environ["AITER_CONFIG_FMOE"] = _bootstrap_cfg()+ def _should_force_nt(token: int, inter_dim: int, expert: int) -> bool:+ if FORCE_NT_MODE == "off":+ return False+ if FORCE_NT_MODE == "all":+ return True+ if FORCE_NT_MODE == "e33":+ return expert == 33+ if FORCE_NT_MODE == "large":+ return token >= 512+ if FORCE_NT_MODE == "large_e33":+ return expert == 33 and token >= 512+ if FORCE_NT_MODE == "257_only":+ return expert == 257+ raise ValueError("Unknown FORCE_NT_MODE=" + str(FORCE_NT_MODE))+++ def _install_force_nt_patch() -> None:+ if FORCE_NT_MODE == "off":+ return+ if getattr(fused_moe_mod, "_josu_force_nt_patch", False):+ return++ os.environ["AITER_USE_NT"] = "1"+ original_get_2stage_cfgs = fused_moe_mod.get_2stage_cfgs++ def wrapped_get_2stage_cfgs(+ token,+ model_dim,+ inter_dim,+ expert,+ topk,+ dtype,+ q_dtype_a,+ q_dtype_w,+ q_type,+ use_g1u1,+ activation,+ doweight_stage1,+ hidden_pad,+ intermediate_pad,+ is_shuffled=True,+ ):+ meta = original_get_2stage_cfgs(+ token,+ model_dim,+ inter_dim,+ expert,+ topk,+ dtype,+ q_dtype_a,+ q_dtype_w,+ q_type,+ use_g1u1,+ activation,+ doweight_stage1,+ hidden_pad,+ intermediate_pad,+ is_shuffled,+ )+ if not _should_force_nt(token, inter_dim, expert):+ return meta++ stage1 = meta.stage1+ stage2 = meta.stage2+ if isinstance(stage1, functools.partial) and stage1.func is fused_moe_mod.ck_moe_stage1:+ stage1_kwargs = dict(stage1.keywords or dict())+ stage1_kwargs["use_non_temporal_load"] = True+ stage1 = functools.partial(stage1.func, *(stage1.args or ()), **stage1_kwargs)+ if (+ isinstance(stage2, functools.partial)+ and stage2.func is fused_moe_mod.aiter.ck_moe_stage2_fwd+ ):+ stage2_kwargs = dict(stage2.keywords or dict())+ stage2_kwargs["use_non_temporal_load"] = True+ stage2 = functools.partial(stage2.func, *(stage2.args or ()), **stage2_kwargs)+ return fused_moe_mod.MOEMetadata(+ stage1,+ stage2,+ int(meta.block_m),+ int(meta.ksplit),+ meta.run_1stage,+ meta.has_bias,+ True,+ )++ fused_moe_mod.get_2stage_cfgs = wrapped_get_2stage_cfgs+ fused_moe_mod._josu_force_nt_patch = True+++ _install_force_nt_patch()++def custom_kernel(data: input_t) -> output_t:(hidden_states,
scrolls · 131 diff lines total
Best evidence level for this revision: reported
JSON