Skip to content
KernelIndex
Search⌘K

submission 528303

josusanmartin · python · License unknown

Use it

Vendorable · source mirrored · license unknownView source →

No package. Vendor the mirrored source: 303 lines, June 9 Researcher Reciprocity License v1.0.

best.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-amd-moe-mxfp4-528303?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
AMD MXFP4 MoEsuite of 7 cases
AMD Instinct MI355X
152.4µs
#208 of 782
2026-03-11

Reported · How evidence levels are derived →

Source and license

sourceavailable
revision digestsha256:f5988cf90fbe897b1fcc2e74b51e4ba183004e5a6736c1802932150d690c980f
license declaredunknown
license concludedunknown
authorsjosusanmartin
imported2026-08-15

Kernel source

best.py303 lines
import os
import sys
from pathlib import Path

import pandas as pd
import torch
from aiter import ActivationType, QuantType
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 = "v9"
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": 129.6056,
            "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": 782.9,
            "bw": 2888.31,
        },
        {
            "token": 512,
            "model_dim": 7168,
            "inter_dim": 2048,
            "expert": 33,
            "topk": 9,
            "block_m": 128,
            "ksplit": 0,
            "us": 275.9788,
            "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": 1470.67,
            "bw": 5305.97,
        },
    ]

    k2_overrides = []
    default_2stage_rows = [
        (16, 7168, 512, 33, 9),
        (128, 7168, 512, 33, 9),
    ]

    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 token, model_dim, inter_dim, expert, topk in default_2stage_rows:
        row = defaults.copy()
        row.update(
            {
                "token": token,
                "model_dim": model_dim,
                "inter_dim": inter_dim,
                "expert": expert,
                "topk": topk,
                "block_m": 32,
                "ksplit": 2,
                "us": -1.0,
                "run_1stage": 0,
            }
        )
        rows.append(row)

    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 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 · 303 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 527260.

import os
+ import sys
from pathlib import Path
import pandas as pd
⋯ 3 unchanged lines
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 = "v9"
+ 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("/tmp/josu_cfg_hybrid_under150_fmoe.csv")
- if out_path.exists() and out_path.stat().st_size > 0:
+ 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 = [
⋯ 43 unchanged lines
{
"token": 16,
"model_dim": 7168,
- "inter_dim": 512,
- "expert": 33,
+ "inter_dim": 256,
+ "expert": 257,
"topk": 9,
"block_m": 32,
- "ksplit": 0,
- "us": 47.3336,
+ "ksplit": 2,
+ "us": 0.0,
"us1": 0.0,
- "kernelName1": "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
+ "kernelName1": "",
"err1": "0.0%",
"us2": 0.0,
- "kernelName2": "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16",
- "err2": "2.9%",
+ "kernelName2": "",
+ "err2": "0.0%",
"run_1stage": 0,
- "tflops": 66.99,
- "bw": 7683.25,
+ "tflops": 0.0,
+ "bw": 0.0,
},
{
"token": 128,
"model_dim": 7168,
- "inter_dim": 512,
- "expert": 33,
+ "inter_dim": 256,
+ "expert": 257,
"topk": 9,
"block_m": 32,
- "ksplit": 0,
- "us": 58.5681,
+ "ksplit": 2,
+ "us": 0.0,
"us1": 0.0,
- "kernelName1": "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
+ "kernelName1": "",
"err1": "0.0%",
"us2": 0.0,
- "kernelName2": "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16",
- "err2": "2.9%",
+ "kernelName2": "",
+ "err2": "0.0%",
"run_1stage": 0,
- "tflops": 433.12,
- "bw": 6250.57,
+ "tflops": 0.0,
+ "bw": 0.0,
},
{
"token": 512,
⋯ 35 unchanged lines
},
]
- k2_overrides = [
- (16, 7168, 256, 257, 9),
- (128, 7168, 256, 257, 9),
+ k2_overrides = []
+ default_2stage_rows = [
(16, 7168, 512, 33, 9),
(128, 7168, 512, 33, 9),
]
⋯ 2 unchanged lines
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:
⋯ 12 unchanged lines
& (merged["use_g1u1"] == 1)
& (merged["doweight_stage1"] == 0)
)
- if mask.any():
- row = merged.loc[mask].sort_values("us").iloc[0].to_dict()
- else:
- row = defaults.copy()
- row.update(
- {
- "token": token,
- "model_dim": model_dim,
- "inter_dim": inter_dim,
- "expert": expert,
- "topk": topk,
- }
- )
+ 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 token, model_dim, inter_dim, expert, topk in default_2stage_rows:
+ row = defaults.copy()
+ row.update(
+ {
+ "token": token,
+ "model_dim": model_dim,
+ "inter_dim": inter_dim,
+ "expert": expert,
+ "topk": topk,
+ "block_m": 32,
+ "ksplit": 2,
+ "us": -1.0,
+ "run_1stage": 0,
+ }
+ )
+ rows.append(row)
+
merged = pd.concat([merged, pd.DataFrame(rows)], ignore_index=True)
- dedup_keys = keys + (["_tag"] if "_tag" in merged.columns else [])
+ _debug_rows("pre-dedup", merged)
merged = (
merged.sort_values("us")
- .drop_duplicates(subset=dedup_keys, keep="first")
+ .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)
scrolls · 208 diff lines total

Best evidence level for this revision: reported

JSON