From 1c3ac799a995c7f8cca068465f9346ea1bbf0f7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Fri, 6 Nov 2026 16:16:16 +0300 Subject: [PATCH 1/7] Enhance leaderboard support by adding Mojo to templates and file extensions. Update assertions to include Mojo in task definitions and modify header generation accordingly. --- src/kernelbot/cogs/leaderboard_cog.py | 18 ++++++++++++++++-- src/libkernelbot/task.py | 2 +- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/kernelbot/cogs/leaderboard_cog.py b/src/kernelbot/cogs/leaderboard_cog.py index 01d5c5838..8432f79b5 100644 --- a/src/kernelbot/cogs/leaderboard_cog.py +++ b/src/kernelbot/cogs/leaderboard_cog.py @@ -239,7 +239,14 @@ async def lang_autocomplete( def add_header_to_template(lang: str, code: str, lb: LeaderboardItem): - comment_char = {"CUDA": "//", "Python": "#", "Triton": "#", "HIP": "#", "CuteDSL": "#"}[lang] + comment_char = { + "CUDA": "//", + "Python": "#", + "Triton": "#", + "HIP": "#", + "CuteDSL": "#", + "Mojo": "#", + }[lang] description_comment = [f"{comment_char} > {line}" for line in lb["description"].splitlines()] header = f""" @@ -548,7 +555,14 @@ async def get_task_template( return template = add_header_to_template(lang, templates[lang], leaderboard_item) - ext = {"CUDA": "cu", "Python": "py", "Triton": "py", "HIP": "py", "CuteDSL": "py"} + ext = { + "CUDA": "cu", + "Python": "py", + "Triton": "py", + "HIP": "py", + "CuteDSL": "py", + "Mojo": "py", + } file_name = f"{leaderboard_name}.{ext[lang]}" file = discord.File(fp=StringIO(template), filename=file_name) message = f"**Starter code for leaderboard `{leaderboard_name}`**\n" diff --git a/src/libkernelbot/task.py b/src/libkernelbot/task.py index ddd819451..7d5cb63a4 100644 --- a/src/libkernelbot/task.py +++ b/src/libkernelbot/task.py @@ -163,7 +163,7 @@ def make_task_definition(yaml_file: str | Path) -> LeaderboardDefinition: # noq # load template files templates = {} for lang, source in raw.get("templates", {}).items(): - assert lang in ["CUDA", "Python", "Triton", "HIP", "CuteDSL"] + assert lang in ["CUDA", "Python", "Triton", "HIP", "CuteDSL", "Mojo"] templates[lang] = (root / source).read_text() if templates: From 238eae723239e97c442ee9f81ddf4516cc374243 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Fri, 6 Nov 2026 16:17:26 +0300 Subject: [PATCH 2/7] Add Mojo-based kernel submissions via `max` package --- src/runners/modal_runner.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/runners/modal_runner.py b/src/runners/modal_runner.py index a457a44ca..e0af5014d 100644 --- a/src/runners/modal_runner.py +++ b/src/runners/modal_runner.py @@ -68,6 +68,7 @@ .uv_pip_install( "tinygrad~=0.10", "helion", + "max", ) # nvidia cuda packages .uv_pip_install( @@ -108,8 +109,9 @@ }) .run_commands( "python -m pip check", - 'python -c "import cuda.tile, nvmath"', + 'python -c "import cuda.tile, nvmath, mojo.importer, max"', "tileiras --version", + "mojo --version", ( "printf '#include \\n' | " "nvcc -std=c++17 -x cu -c - -o /tmp/cublasdx-smoke.o && " From 7261371f4d351a2db2de6a308f5cd382ea2740ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Fri, 6 Nov 2026 16:17:44 +0300 Subject: [PATCH 3/7] Add Mojo-based submission for vector addition kernel in `submission_mojo.py` and update test cases to include this new script. --- examples/vectoradd_py/submission_mojo.py | 110 +++++++++++++++++++++++ tests/test_modal.py | 1 + 2 files changed, 111 insertions(+) create mode 100644 examples/vectoradd_py/submission_mojo.py diff --git a/examples/vectoradd_py/submission_mojo.py b/examples/vectoradd_py/submission_mojo.py new file mode 100644 index 000000000..b25355e30 --- /dev/null +++ b/examples/vectoradd_py/submission_mojo.py @@ -0,0 +1,110 @@ +#!POPCORN leaderboard vectoradd_py-dev +import sys +from pathlib import Path + +import mojo.importer # noqa: F401 +import torch +from task import input_t, output_t + +MOJO_SOURCE = """ +from std.math import ceildiv +from std.memory import Pointer +from std.os import abort +from std.python import PythonObject +from std.python.bindings import PythonModuleBuilder +from max.gpu import global_idx +from max.gpu.host import DeviceContext +from layout import RowMajorLayout, TileTensor, row_major + +comptime BLOCK = 256 +comptime Layout1D = RowMajorLayout[Int] + + +def add_kernel( + a: TileTensor[.float16, Layout1D, ImmUntrackedOrigin], + b: TileTensor[.float16, Layout1D, ImmUntrackedOrigin], + c: TileTensor[.float16, Layout1D, MutUntrackedOrigin], + n: Int32, +): + var idx = Int32(global_idx.x) + if idx < n: + c[idx] = a[idx] + b[idx] + + +@fieldwise_init +struct Runner(Movable, Writable): + var ctx: DeviceContext + var func: type_of(DeviceContext().compile_function[add_kernel]()) + + def write_repr_to(self, mut writer: Some[Writer]): + writer.write("Runner()") + + @staticmethod + def py_init(out self: Self, args: PythonObject, kwargs: PythonObject) raises: + var ctx = DeviceContext() + var func = ctx.compile_function[add_kernel]() + self = Self(ctx^, func^) + + @staticmethod + def launch( + self_ptr: Pointer[Self, MutAnyOrigin], + a_py: PythonObject, + b_py: PythonObject, + c_py: PythonObject, + n_py: PythonObject, + ) raises: + var n = Int(py=n_py) + var a = TileTensor( + Pointer[Float16, ImmUntrackedOrigin](unsafe_from_address=Int(py=a_py)), + row_major(n), + ) + var b = TileTensor( + Pointer[Float16, ImmUntrackedOrigin](unsafe_from_address=Int(py=b_py)), + row_major(n), + ) + var c = TileTensor( + Pointer[Float16, MutUntrackedOrigin](unsafe_from_address=Int(py=c_py)), + row_major(n), + ) + self_ptr[].ctx.enqueue_function( + self_ptr[].func, + a, + b, + c, + Int32(n), + grid_dim=ceildiv(n, BLOCK), + block_dim=BLOCK, + ) + + +@export +def PyInit_mojo_vecadd_example() abi("C") -> PythonObject: + try: + var mb = PythonModuleBuilder("mojo_vecadd_example") + _ = ( + mb.add_type[Runner]("Runner") + .def_py_init[Runner.py_init]() + .def_method[Runner.launch]("launch") + ) + return mb.finalize() + except e: + abort(t"Failed to initialize mojo_vecadd_example: {e}") +""" + +_pkg = Path("/tmp/mojo_vecadd_example") +_pkg.mkdir(exist_ok=True) +(_pkg / "mojo_vecadd_example.mojo").write_text(MOJO_SOURCE) +sys.path.insert(0, str(_pkg)) +import mojo_vecadd_example # noqa: E402 + +_LAUNCH = mojo_vecadd_example.Runner().launch + + +def custom_kernel(data: input_t) -> output_t: + A, B = data + assert A.is_cuda and B.is_cuda, "Both tensors must be on GPU" + assert A.shape == B.shape, "Input tensors must have the same shape" + assert A.dtype == torch.float16 and B.dtype == torch.float16, "Input tensors must be float16" + C = torch.empty_like(A) + _LAUNCH(A.data_ptr(), B.data_ptr(), C.data_ptr(), A.numel()) + return C diff --git a/tests/test_modal.py b/tests/test_modal.py index 241edfb58..2c31fe557 100644 --- a/tests/test_modal.py +++ b/tests/test_modal.py @@ -197,6 +197,7 @@ def modal_deployment(project_root: Path): [ ("vectoradd_py", "submission_cuda_inline.py"), ("vectoradd_py", "submission_triton.py"), + ("vectoradd_py", "submission_mojo.py"), ], ) async def test_modal_launcher_python_script( From 94aedc481374cf9e1f3ae2f7370e1baa6776cd64 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Fri, 6 Nov 2026 17:08:36 +0300 Subject: [PATCH 4/7] Remove assertions for tensor properties in `custom_kernel` function of `submission_mojo.py` to allow more flexible input handling. --- examples/vectoradd_py/submission_mojo.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/examples/vectoradd_py/submission_mojo.py b/examples/vectoradd_py/submission_mojo.py index b25355e30..837d076d7 100644 --- a/examples/vectoradd_py/submission_mojo.py +++ b/examples/vectoradd_py/submission_mojo.py @@ -102,9 +102,6 @@ def PyInit_mojo_vecadd_example() abi("C") -> PythonObject: def custom_kernel(data: input_t) -> output_t: A, B = data - assert A.is_cuda and B.is_cuda, "Both tensors must be on GPU" - assert A.shape == B.shape, "Input tensors must have the same shape" - assert A.dtype == torch.float16 and B.dtype == torch.float16, "Input tensors must be float16" C = torch.empty_like(A) _LAUNCH(A.data_ptr(), B.data_ptr(), C.data_ptr(), A.numel()) return C From d6c14778ec2f1b7d17bcf47cd42810d72df73995 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Sat, 7 Nov 2026 16:20:01 +0300 Subject: [PATCH 5/7] Refactor import of `mojo_vecadd_example` in `submission_mojo.py` to use `importlib` --- examples/vectoradd_py/submission_mojo.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/vectoradd_py/submission_mojo.py b/examples/vectoradd_py/submission_mojo.py index 837d076d7..4c4447383 100644 --- a/examples/vectoradd_py/submission_mojo.py +++ b/examples/vectoradd_py/submission_mojo.py @@ -1,4 +1,5 @@ #!POPCORN leaderboard vectoradd_py-dev +import importlib import sys from pathlib import Path @@ -95,7 +96,7 @@ def PyInit_mojo_vecadd_example() abi("C") -> PythonObject: _pkg.mkdir(exist_ok=True) (_pkg / "mojo_vecadd_example.mojo").write_text(MOJO_SOURCE) sys.path.insert(0, str(_pkg)) -import mojo_vecadd_example # noqa: E402 +mojo_vecadd_example = importlib.import_module("mojo_vecadd_example") _LAUNCH = mojo_vecadd_example.Runner().launch From 92a6df0c2a9d4767df78d2452984c5efc1263216 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Sat, 7 Nov 2026 16:27:47 +0300 Subject: [PATCH 6/7] Update temporary directory creation in `submission_mojo.py` to use `tempfile.mkdtemp` for improved path management. --- examples/vectoradd_py/submission_mojo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/vectoradd_py/submission_mojo.py b/examples/vectoradd_py/submission_mojo.py index 4c4447383..eed9e6979 100644 --- a/examples/vectoradd_py/submission_mojo.py +++ b/examples/vectoradd_py/submission_mojo.py @@ -1,6 +1,7 @@ #!POPCORN leaderboard vectoradd_py-dev import importlib import sys +import tempfile from pathlib import Path import mojo.importer # noqa: F401 @@ -92,8 +93,7 @@ def PyInit_mojo_vecadd_example() abi("C") -> PythonObject: abort(t"Failed to initialize mojo_vecadd_example: {e}") """ -_pkg = Path("/tmp/mojo_vecadd_example") -_pkg.mkdir(exist_ok=True) +_pkg = Path(tempfile.mkdtemp(prefix="mojo_vecadd_example_")) (_pkg / "mojo_vecadd_example.mojo").write_text(MOJO_SOURCE) sys.path.insert(0, str(_pkg)) mojo_vecadd_example = importlib.import_module("mojo_vecadd_example") From b29cea42075c03421899258ecf56ec1fead868e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Sat, 7 Nov 2026 18:07:03 +0300 Subject: [PATCH 7/7] Refactor temporary file handling in `submission_mojo.py` to directly write the Mojo source file to the script's directory, removing the use of `tempfile` for improved clarity and path management. --- examples/vectoradd_py/submission_mojo.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/examples/vectoradd_py/submission_mojo.py b/examples/vectoradd_py/submission_mojo.py index eed9e6979..7bf2714f4 100644 --- a/examples/vectoradd_py/submission_mojo.py +++ b/examples/vectoradd_py/submission_mojo.py @@ -1,7 +1,5 @@ #!POPCORN leaderboard vectoradd_py-dev import importlib -import sys -import tempfile from pathlib import Path import mojo.importer # noqa: F401 @@ -93,9 +91,7 @@ def PyInit_mojo_vecadd_example() abi("C") -> PythonObject: abort(t"Failed to initialize mojo_vecadd_example: {e}") """ -_pkg = Path(tempfile.mkdtemp(prefix="mojo_vecadd_example_")) -(_pkg / "mojo_vecadd_example.mojo").write_text(MOJO_SOURCE) -sys.path.insert(0, str(_pkg)) +(Path(__file__).resolve().parent / "mojo_vecadd_example.mojo").write_text(MOJO_SOURCE) mojo_vecadd_example = importlib.import_module("mojo_vecadd_example") _LAUNCH = mojo_vecadd_example.Runner().launch