[BUG]: `@compiler.register` graph compile fails on H200 (sm_90), works on RTX 5090 (sm_120)

Bug description

Summary

I have been having a lot of issues with my kernels working on an RTX 5090 but completely hanging/freezing and not executing on a H200 (no error or return). I started trying very simple kernels to see if it was a niche error but I found out that this seems generalised to any GPU kernel I could create when called through the python interop for pytorch. With a very simple kernel I at least managed to see an error, opposed to just freezing completely. I tried 26.2 before moving to the nightlies but no luck. I share here a simple example of this happening.

A trivial custom op registered via @compiler.register and dispatched from PyTorch via @graph_op fails with RuntimeError: Invalid argument (or terminate called after throwing 'std::system_error': Invalid argument, exit code 134) on NVIDIA H200 (sm_90, Hopper). The exact same code runs successfully on NVIDIA RTX 5090 (sm_120, Blackwell).

The kernel's execute() method is never called on H200, a print(...) placed at the top of execute() produces no output, indicating the failure happens inside MAX's graph compile / dispatch layer, before our kernel runs.

Standalone Mojo GPU code (no MAX, no @compiler.register) runs fine on the same H200 with the same driver.

Notes

• The bug appears to be in MAX's graph build / compile layer, not in our kernel. Since standalone Mojo GPU works on the same H200, the Mojo compiler / runtime / driver path itself is fine, the issue is upstream of execute().
• Happy to test patches / diagnostic builds.

Steps to reproduce

trivial.mojo

"""Minimal @compiler.register kernel for sm_90 hang isolation."""

import compiler
from gpu import thread_idx, block_idx
from math import ceildiv
from runtime.asyncrt import DeviceContextPtr
from tensor import OutputTensor
from memory import UnsafePointer

comptime block_size: Int = 256


fn fill_const_kernel(
    out_ptr: UnsafePointer[Scalar[DType.float32], MutAnyOrigin],
    n: Int,
):
    var i = Int(block_idx.x) * block_size + Int(thread_idx.x)
    if i < n:
        out_ptr[i] = 42.0


@compiler.register("trivial_fill")
struct TrivialFill:
    @staticmethod
    fn execute[
        target: StaticString,
    ](
        result: OutputTensor[dtype=DType.float32, rank=1, ...],
        ctx: DeviceContextPtr,
    ) raises:
        var N = result.dim_size(0)
        var out_ptr = rebind[UnsafePointer[Scalar[DType.float32], MutAnyOrigin]](
            result.to_layout_tensor().ptr
        )

        # If this print fires, execute() was reached. On H200 it does NOT fire,
        # confirming the failure happens inside MAX before dispatch.
        print("[trivial_fill] execute() called, target=", target, ", N=", N)

        var gpu_ctx = ctx.get_device_context()
        var grid = ceildiv(N, block_size)
        gpu_ctx.enqueue_function_unchecked[fill_const_kernel](
            out_ptr, N,
            grid_dim=grid, block_dim=block_size,
        )

trivial_max_test.py

"""Minimal @compiler.register + @graph_op repro."""

import time
from pathlib import Path

import torch

from max.experimental.torch import graph_op
from max.graph import ops as graph_ops, TensorType as GraphTensorType, DeviceRef
from max.dtype import DType as MaxDType


_mojo_kernels = Path(__file__).resolve().parent / "kernels"  # dir with trivial.mojo
_gpu = DeviceRef.GPU()


@graph_op(
    name="trivial_fill_op",
    kernel_library=_mojo_kernels,
    input_types=[],
    output_types=[GraphTensorType(MaxDType.float32, (16,), device=_gpu)],
)
def _trivial_fill_graph():
    return graph_ops.custom(
        "trivial_fill",
        _gpu,
        [],
        out_types=[GraphTensorType(MaxDType.float32, (16,), device=_gpu)],
    )


def main():
    print("[py] start", flush=True)
    result = torch.empty(16, dtype=torch.float32, device="cuda")
    print("[py] calling graph_op", flush=True)
    t0 = time.time()
    _trivial_fill_graph(result)
    torch.cuda.synchronize()
    print(f"[py] graph_op done in {time.time() - t0:.2f}s", flush=True)
    print(f"[py] result[:5]={result[:5].tolist()}", flush=True)


if __name__ == "__main__":
    main()

Run

python trivial_max_test.py

Expected behavior

[py] start
[py] calling graph_op
[trivial_fill] execute() called, target=gpu , N= 16
[py] graph_op done in 0.05s
[py] result[:5]=[42.0, 42.0, 42.0, 42.0, 42.0]

(This is what we get on RTX 5090.)

Actual behavior on H200

[py] start
[py] calling graph_op
terminate called after throwing an instance of 'std::system_error'
  what():  Invalid argument
exit_code=134

Or with a slightly different graph variant (in-place via InputTensor + dummy OutputTensor):

[py] start
[py] calling graph_op
Traceback (most recent call last):
  ...
  File "/.../max/engine/api.py", line 131, in _Model_execute
    return self._execute_device_tensors(input_impls)
RuntimeError: Invalid argument

In neither case does the print(...) from inside execute() fire, the failure happens in MAX's graph compile, before our kernel is dispatched.

System information

- **GPU (failing):** NVIDIA H200, sm_90, Hopper
- **GPU (working):** NVIDIA RTX 5090, sm_120, Blackwell
- **Modular version:** `26.3.0.dev2026042605`
- **NVIDIA driver:** 580.126.09 (CUDA 13.0)
- **Python:** 3.12

Start with the minimal reproducer in trivial.mojo and trivial_max_test.py, then run python trivial_max_test.py on the H200 and RTX 5090 environments described. Trace the @compiler.register and @graph_op path before execute() is reached, using the reported Invalid argument failure as the checkpoint. Done means the H200 reproducer reaches execute(), completes graph_op, and produces the expected values without aborting.