CUDA_ERROR_ILLEGAL_ADDRESS on NVIDIA GB10 Blackwell (sm_121a) - all GPU ops fail

MAX SDK Bug: CUDA_ERROR_ILLEGAL_ADDRESS on NVIDIA GB10 (Blackwell sm_121a)

Environment

Component	Version
Hardware	NVIDIA GB10 (DGX Spark), Compute Capability 12.1
CUDA Driver	590.48.01 (CUDA 13.1)
CUDA Toolkit	13.0 + 13.1 installed
MAX SDK	26.1.0 (pip) AND 26.2.0.dev2026030405 (pixi nightly) - both affected
OS	Linux aarch64 (Ubuntu-based)
Python	3.14.2

Minimal Reproducer

import numpy as np
from max import driver, engine
from max.dtype import DType
from max.graph import DeviceKind, DeviceRef, Graph, TensorType

dev = DeviceRef(DeviceKind.GPU, 0)
gpu = driver.Accelerator()
# gpu.architecture_name == "sm_121a"

# Simplest possible GPU graph: x + x
with Graph(name="test", input_types=[
    TensorType(DType.float32, [4], device=dev),
]) as g:
    x = g.inputs[0]
    g.output(x + x)

session = engine.InferenceSession(devices=[gpu])
model = session.load(g)

buf = driver.Buffer(np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), gpu)
result = model.execute(buf)

# This line crashes:
out = result[0].to(driver.CPU()).to_numpy()
# RuntimeError: CUDA call failed: CUDA_ERROR_ILLEGAL_ADDRESS

Scope: Affects ALL GPU operations - built-in ops, custom Mojo ops, even buffer round-trip. Not limited to custom kernels.

Diagnosis

1. compute-sanitizer memcheck

Invalid __global__ read of size 16 bytes
    at std_algorithm_functional__ele6A6A6A6A6A6A_4517b9067854c62b+0x190
    by thread (0,0,0) in block (0,0,0)
    Access to 0x348f63f0 is out of bounds
    and is 252,739,582,991,376 bytes before the nearest allocation
    at 0xe5ddb9a00000 of size 1,715,994,624 bytes

The MAX-generated kernel accesses address 0x348f63f0 while the actual GPU allocation is at 0xe5ddb9a00000. This is a ~253 TB address mismatch.

2. cuobjdump analysis of MAX-generated cubin

64-bit ELF: type=ET_EXEC, ABI=8, sm=121a, toolkit=12.9

code for sm_121a
LDCU.64 UR4, c[0x0][0x358]                    # Load descriptor from const mem
LDG.E.CONSTANT R15, desc[UR4][R4.64]          # Blackwell TMA descriptor load
LDG.E.128.CONSTANT R4, desc[UR4][R4.64]       # 128-bit descriptor load

The cubin is correct: 64-bit ELF, targets sm_121a, uses 64-bit addressing (R4.64). The SASS uses Blackwell's descriptor-based memory access mode.

3. Raw CUDA driver API works perfectly

# Direct cuMemAlloc + cuMemcpy round-trip succeeds:
# cuMemAlloc -> ptr=0xf11c7aa00000 (48-bit address)
# cuMemcpyHtoD -> success
# cuMemcpyDtoH -> success, returns [1.0, 2.0, 3.0, 4.0]

PyTorch CUDA and Triton also work correctly on this GPU.

4. All GPU allocations use 48-bit addresses

GB10 unified memory allocates at addresses like 0xf364..., 0xe5dd..., 0xf11c... - all in the 48-bit range. There is no way to force 32-bit-range allocations.

Conclusion

The kernel code and cubin are correct (64-bit). The hardware is fine (raw CUDA works). The bug is in MAX's runtime - specifically in how it populates the Blackwell TMA descriptor data in kernel constant memory. The descriptors end up pointing to wrong addresses, causing all global loads to access invalid memory.

The relevant code path appears to be in libAsyncRTMojoBindings.so which contains AsyncRT_cuda_tensorMapEncodeTiled and uses cuModuleLoadDataEx via dlopen/dlsym on libcuda.so.1.

GPU Properties

Name: NVIDIA GB10
Compute: 12.1 (sm_121a)
Unified addressing: 1
Managed memory: 1
Pageable memory access: 0
Concurrent managed access: 1
Total memory: ~79 GB (unified with CPU)

Impact

This completely blocks GPU usage on DGX Spark / GB10 Blackwell hardware with MAX SDK. We are using Triton (via PyTorch CUDA, which works fine) as a workaround for our custom FP4 dequant-matmul kernel.

Reproduce the failure with the Python graph and buffer round-trip shown in the issue, then inspect libAsyncRTMojoBindings.so around AsyncRT_cuda_tensorMapEncodeTiled and its cuModuleLoadDataEx path. Compare the MAX-generated behavior with the raw CUDA driver round-trip and compute-sanitizer output. Done means GPU operations on NVIDIA GB10 sm_121a complete without CUDA_ERROR_ILLEGAL_ADDRESS.