From 1e53253d54165e59b5f7d9885cc450f4171ace01 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 16 Jun 2026 18:49:59 +0000 Subject: [PATCH 1/6] Added nki swiglu kernel --- benchmarks/operators/swiglu/impl_nki.py | 42 +++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 benchmarks/operators/swiglu/impl_nki.py diff --git a/benchmarks/operators/swiglu/impl_nki.py b/benchmarks/operators/swiglu/impl_nki.py new file mode 100644 index 00000000..307fd40a --- /dev/null +++ b/benchmarks/operators/swiglu/impl_nki.py @@ -0,0 +1,42 @@ +import torch +from torch_xla.core import xla_model as xm + +try: + import neuronxcc.nki as nki + import neuronxcc.nki.language as nl + import neuronxcc.nki.isa as nisa + PMAX = nl.tile_size.pmax +except ImportError: + nki = None + +if nki is not None: + @nki.jit + def swiglu_kernel(x_input, y_input): + num_blocks = (x_input.shape[0] + PMAX - 1) // PMAX + + hbm_result_tile = nl.ndarray(x_input.shape, dtype=x_input.dtype, buffer=nl.hbm) + + for i in range(num_blocks): + offset = i*PMAX + + partition_index = nl.arange(PMAX)[:, None] + free_dim_index = nl.arange(x_input.shape[1])[None, :] + + mask = partition_index < (x_input.shape[0] - offset) + + x_tile = nl.load(x_input[offset + partition_index, free_dim_index], mask=mask) + y_tile = nl.load(y_input[offset + partition_index, free_dim_index], mask=mask) + + silu_x = nl.silu(x_tile, mask=mask) + result_tile = nl.multiply(silu_x, y_tile, mask=mask) + + nl.store(hbm_result_tile[offset + partition_index, free_dim_index], value=result_tile, mask=mask) + + return hbm_result_tile + +def run(x: torch.Tensor, y: torch.Tensor, block_size: int = 1024, autotune: bool = False, **kwargs) -> torch.Tensor: + return swiglu_kernel(x, y) + + +def get_last_config() -> dict | None: + return None From 6c2ca6ed7c82b0981bde03435cb8f9e6bbbbba16 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 18 Jun 2026 16:33:05 +0000 Subject: [PATCH 2/6] removed unnecessary import --- benchmarks/operators/swiglu/impl_nki.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/benchmarks/operators/swiglu/impl_nki.py b/benchmarks/operators/swiglu/impl_nki.py index 307fd40a..103a7ce6 100644 --- a/benchmarks/operators/swiglu/impl_nki.py +++ b/benchmarks/operators/swiglu/impl_nki.py @@ -1,5 +1,4 @@ import torch -from torch_xla.core import xla_model as xm try: import neuronxcc.nki as nki @@ -37,6 +36,5 @@ def swiglu_kernel(x_input, y_input): def run(x: torch.Tensor, y: torch.Tensor, block_size: int = 1024, autotune: bool = False, **kwargs) -> torch.Tensor: return swiglu_kernel(x, y) - def get_last_config() -> dict | None: return None From 13421f72c4297f64dff1cb408fcf7b8c49fe436c Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 1 Jul 2026 20:35:40 +0000 Subject: [PATCH 3/6] modified json parsing in _total_time_ms to fit trainium conventions --- core/nki_timer.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/core/nki_timer.py b/core/nki_timer.py index ce3711a4..463ac737 100644 --- a/core/nki_timer.py +++ b/core/nki_timer.py @@ -130,9 +130,17 @@ def _total_time_ms(summary_json_text: str) -> float: total_time across rows and convert seconds -> ms. """ data = json.loads(summary_json_text) - rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) - if isinstance(rows, dict): - rows = [rows] + + #rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) + #if isinstance(rows, dict): + #rows = [rows] + if isinstance(data, dict) and not any(k in data for k in ["total_time", "summary", "rows"]): + rows = [v for v in data.values() if isinstance(v, dict)] + else: + rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) + if isinstance(rows, dict): + rows = [rows] + times = [float(r["total_time"]) for r in rows if isinstance(r, dict) and "total_time" in r] if not times: raise RuntimeError(f"no 'total_time' in neuron-profile summary-json: {summary_json_text[:300]}") From 44b94e4b5ac6e58f1b7e83840cc45e5760e983a9 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 1 Jul 2026 20:39:38 +0000 Subject: [PATCH 4/6] set NEURON_RT_NUM_CORES=1 to resolve neuron-profile core allocation error on trn2.3xlarge --- scripts/run_bench.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/scripts/run_bench.py b/scripts/run_bench.py index 7e06f091..d6e98da4 100644 --- a/scripts/run_bench.py +++ b/scripts/run_bench.py @@ -1,6 +1,7 @@ import argparse import csv import json +import os from pathlib import Path from core.engine import run_benchmark_suite @@ -23,6 +24,8 @@ def _split(results: list[dict], active: list[str]) -> tuple[list[dict], list[dic def main(): + os.environ["NEURON_RT_NUM_CORES"] = "1" + parser = argparse.ArgumentParser(description="Run TileBench benchmarks") parser.add_argument("--operator", type=str, default="vector_add", help="Operator to benchmark") From ac8d899e4eccfb0ee5798de60a9bd337d5b02102 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 7 Jul 2026 18:49:38 +0000 Subject: [PATCH 5/6] Added free dim tiling --- benchmarks/operators/swiglu/impl_nki.py | 27 ++++++++++++++++--------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/benchmarks/operators/swiglu/impl_nki.py b/benchmarks/operators/swiglu/impl_nki.py index 103a7ce6..c777413d 100644 --- a/benchmarks/operators/swiglu/impl_nki.py +++ b/benchmarks/operators/swiglu/impl_nki.py @@ -11,25 +11,32 @@ if nki is not None: @nki.jit def swiglu_kernel(x_input, y_input): - num_blocks = (x_input.shape[0] + PMAX - 1) // PMAX + free_tile_size = 16384 + + num_blocks = (x_input.shape[0] +(PMAX - 1)) // PMAX + num_free_blocks = (x_input.shape[1] + free_tile_size - 1) // free_tile_size hbm_result_tile = nl.ndarray(x_input.shape, dtype=x_input.dtype, buffer=nl.hbm) for i in range(num_blocks): - offset = i*PMAX - + offset = i * PMAX partition_index = nl.arange(PMAX)[:, None] - free_dim_index = nl.arange(x_input.shape[1])[None, :] + mask_p = partition_index < (x_input.shape[0] - offset) + + for j in range(num_free_blocks): + free_offset = j * free_tile_size + free_dim_index = nl.arange(free_tile_size)[None, :] - mask = partition_index < (x_input.shape[0] - offset) + mask_f = free_dim_index < (x_input.shape[1] - free_offset) + mask = mask_p & mask_f - x_tile = nl.load(x_input[offset + partition_index, free_dim_index], mask=mask) - y_tile = nl.load(y_input[offset + partition_index, free_dim_index], mask=mask) + x_tile = nl.load(x_input[offset + partition_index, free_offset + free_dim_index], mask=mask) + y_tile = nl.load(y_input[offset + partition_index, free_offset + free_dim_index], mask=mask) - silu_x = nl.silu(x_tile, mask=mask) - result_tile = nl.multiply(silu_x, y_tile, mask=mask) + silu_x = nl.silu(x_tile, mask=mask) + result_tile = nl.multiply(silu_x, y_tile, mask=mask) - nl.store(hbm_result_tile[offset + partition_index, free_dim_index], value=result_tile, mask=mask) + nl.store(hbm_result_tile[offset + partition_index, free_offset + free_dim_index], value=result_tile, mask=mask) return hbm_result_tile From 6ddad58a04c2cb0cfcb08e56bc4a6287dd0ee51e Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Sun, 9 Aug 2026 21:47:09 +0000 Subject: [PATCH 6/6] removed code in core and scripts --- core/nki_timer.py | 14 +++----------- scripts/run_bench.py | 3 --- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/core/nki_timer.py b/core/nki_timer.py index 463ac737..ce3711a4 100644 --- a/core/nki_timer.py +++ b/core/nki_timer.py @@ -130,17 +130,9 @@ def _total_time_ms(summary_json_text: str) -> float: total_time across rows and convert seconds -> ms. """ data = json.loads(summary_json_text) - - #rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) - #if isinstance(rows, dict): - #rows = [rows] - if isinstance(data, dict) and not any(k in data for k in ["total_time", "summary", "rows"]): - rows = [v for v in data.values() if isinstance(v, dict)] - else: - rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) - if isinstance(rows, dict): - rows = [rows] - + rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) + if isinstance(rows, dict): + rows = [rows] times = [float(r["total_time"]) for r in rows if isinstance(r, dict) and "total_time" in r] if not times: raise RuntimeError(f"no 'total_time' in neuron-profile summary-json: {summary_json_text[:300]}") diff --git a/scripts/run_bench.py b/scripts/run_bench.py index d6e98da4..7e06f091 100644 --- a/scripts/run_bench.py +++ b/scripts/run_bench.py @@ -1,7 +1,6 @@ import argparse import csv import json -import os from pathlib import Path from core.engine import run_benchmark_suite @@ -24,8 +23,6 @@ def _split(results: list[dict], active: list[str]) -> tuple[list[dict], list[dic def main(): - os.environ["NEURON_RT_NUM_CORES"] = "1" - parser = argparse.ArgumentParser(description="Run TileBench benchmarks") parser.add_argument("--operator", type=str, default="vector_add", help="Operator to benchmark")