[BUG]: [sm_121] MOGGKernelAPI ValueError on DGX Spark with Gemma 4 MoE

Bug description

Hardware/Environment:

System: NVIDIA DGX Spark (Grace-Blackwell GB10)

Compute Capability: 12.1 (sm_121)

Memory: 128GB Unified Memory

OS: Ubuntu 22.04 / 24.04 (DGX OS)

MAX Version: 26.3.0.dev (Nightly) via Pixi

Issue 1: Kernel Resolution Failure (MoE)
When attempting to serve google/gemma-4-26B-A4B-it, the engine fails during graph compilation. It appears the custom kernels for Mixture of Experts (MoE) routing are not resolving for the sm_121 target.

Error Traceback:

Python
File "/home/prabhatkumar/quickstart/.pixi/envs/default/lib/python3.14/site-packages/max/pipelines/architectures/gemma4/gemma4.py", line 225, in call
merge_multimodal_embeddings(
inputs_embeds=h_device,
multimodal_embeddings=img_embed,
image_token_indices=img_tok_indices,
)
File "/home/prabhatkumar/quickstart/.pixi/envs/default/lib/python3.14/site-packages/max/pipelines/lib/vlm_utils.py", line 64, in merge_multimodal_embeddings
return scatter_nd_skip_oob_indices(
input=inputs_embeds,
updates=multimodal_embeddings,
indices=indices_2d,
)
File "/home/prabhatkumar/quickstart/.pixi/envs/default/lib/python3.14/site-packages/max/nn/kernels.py", line 5243, in scatter_nd_skip_oob_indices
return ops.custom(
"mo.scatter_nd.skip_neg_indices",
...
out_types=[TensorType(input.dtype, input.shape, device=input.device)],
)[0].tensor

ValueError: Failed to resolve module path for MOGGKernelAPI
Issue 2: Memory Estimation Logic on Unified Memory
The MemoryEstimator fails to recognize the 128GB Unified Memory pool on the DGX Spark, seemingly defaulting to a ~28GB "available" limit, which prevents dense models (Gemma 4 31B) from passing the Pydantic validation phase even with high utilization flags.

Error Message:

Plaintext
RuntimeError: Model size exceeds available memory (73.25 GiB > 28.68 GiB).
Model weights: 58.25 GiB, Activation memory: 15.00 GiB.
Try running a smaller model, using a smaller precision, or using a device with more memory.
Steps Taken to Resolve:

Set MODULAR_NVPTX_COMPILER_PATH to system ptxas (v8.1+).

Attempted --device-memory-utilization values from 0.7 to 0.99.

Attempted --max-length 512 to minimize KV Cache footprint.

Attempted quantization flags (q4_0, float8_e4m3fn), which currently return ValidationError for the Gemma 4 architecture.

Expected Behavior:
The MAX engine should recognize the sm_121 compute capability of the GB10 chip and utilize the full Unified Memory pool available on DGX Spark systems.

Steps to reproduce

1. Install the MAX engine in a clean Pixi environment:

Bash
pixi init reproduction
cd reproduction
pixi project channel add https://conda.modular.com/max
pixi add max
2. Attempt to serve the Gemma 4 MoE model (Triggering the MOGGKernelAPI error):

Bash
pixi run max serve --model-path google/gemma-4-26B-A4B-it --max-length 1024
3. Attempt to serve the Gemma 4 Dense model (Triggering the MemoryEstimator error):

Bash
pixi run max serve --model-path google/gemma-4-31b-it --max-length 1024 --device-memory-utilization 0.95
4. Observe Results:

Actual Result (Step 2): Process crashes with ValueError: Failed to resolve module path for MOGGKernelAPI.

Actual Result (Step 3): Process fails validation with RuntimeError: Model size exceeds available memory (73.25 GiB > 28.68 GiB).

Expected Result: The engine should correctly resolve sm_121 kernels and recognize the full 128GB Unified Memory pool on the DGX Spark.

System information

Model: NVIDIA DGX Spark (Founders Edition)

Architecture: aarch64 (ARM64)

CPU: Grace 20-core (10x Cortex-X925, 10x Cortex-A725)

GPU: NVIDIA GB10 (Blackwell Generation)

Compute Capability: sm_121

Driver Version: [Run nvidia-smi and paste the version here, e.g., 580.95.05]

CUDA Version: [Run nvcc --version or see nvidia-smi, e.g., 13.0]

Total System Memory: 128 GB LPDDR5x (Unified)

OS: DGX OS (Ubuntu 22.04 or 24.04)

Python Version: 3.14 (via Pixi)

MAX SDK Version: latest

Start by reproducing both failures with the two pixi run max serve commands on the listed DGX Spark environment. Read max/nn/kernels.py, max/pipelines/lib/vlm_utils.py, and max/pipelines/architectures/gemma4/gemma4.py, then locate the MemoryEstimator entry point. Done means the MoE kernel resolves for sm_121 and unified memory is estimated correctly for the dense model.