[Muse Glimmer][CUDA] SM80+PTX vision kernel exceeds RTX 5090 shared-memory limit; native SM120 export OOMs

🐛 Describe the bug

Summary

I followed the official Muse Glimmer ExecuTorch README on an NVIDIA GeForce RTX 5090 (compute capability 12.0 / SM120).

The published full 17G + image + DFlash sm80+ptx artifact starts correctly and text generation works, but every image request fails in the first vision_encoder invocation. CUDA Compute Sanitizer identifies the first failing API call as cuFuncSetAttribute while loading triton_tem_fused__unsafe_view_gelu_mm_t_view_28.

Inspection of the vision delegate binary shows that this kernel requests 0x20000 (131,072) bytes of dynamic shared memory for both loadKernel and the subsequent launch. The RTX 5090 reports an opt-in per-block limit of 101,376 bytes, so the attribute request returns CUDA_ERROR_INVALID_VALUE before the kernel launches.

I then tried the README's source-export path to produce an SM120-native full artifact from the official target, DFlash and projector GGUFs. The exporter is repeatedly killed by the Linux OOM killer while converting/packing the target CUDA tensors, before it loads DFlash or vision. This reproduced with WSL memory limits of 20, 24, 27 and finally 29 GB; on the 29 GB attempt all other local AI services were stopped and 29,454,942,208 bytes were available before launch.

Artifacts and revisions

• ExecuTorch commit: 43f89fb21aa8f109d66832227f2a797725fb4b9b
• Prebuilt repository: meta-models/Muse-Glimmer-30B-ExecuTorch-PTE
• Prebuilt revision: dd7f73bb7ff0d614dfe1c04d1ceb6a70a5f01498
• Variant: muse-glimmer-k-quant-17G-128K-text-image-dflash-sm80+ptx
• PTD: 28,607,362,688 bytes, SHA-256 5c61f6b21329128f114b4f40bc191f3d79b8bbaf26383a09db526eee6ac83f93
• PTE: 33,328,672 bytes, SHA-256 9fdff583d84e9ea65af43fed66cea5d6baa47ef95ae8fcc4b8d62aec14b91664
• pos_embed.bin: 6,291,456 bytes, SHA-256 b69dce2605c6ff068766a80686c603deb715129d574375a92672b5a4c70410d2
• Source GGUF repository: meta-models/Muse-Glimmer-30B-GGUF
• Source GGUF revision: 43c7eadd41352a299ea8e0a36b3157978dd63596

Reproduction: published full artifact

Build the Muse worker/runtime at the ExecuTorch commit above and start the official OpenAI-compatible server (paths abbreviated):

CUDA_FORCE_PTX_JIT=1 \
CUDA_HOME=/usr/local/cuda-13.0 \
PATH="/usr/local/cuda-13.0/bin:$VENV/bin:$PATH" \
LD_LIBRARY_PATH="/usr/local/cuda-13.0/lib64:$REPO/cmake-out/lib:/usr/lib/wsl/lib" \
python -m executorch.examples.models.muse_glimmer.serving.serve \
  --model-path "$MODEL_DIR/$MODEL_BASE.pte" \
  --data-path "$MODEL_DIR/$MODEL_BASE.ptd" \
  --pos-embed-path "$MODEL_DIR/pos_embed.bin" \
  --tokenizer-path "$TOKENIZER/tokenizer.json" \
  --hf-tokenizer "$TOKENIZER" \
  --worker-bin "$REPO/cmake-out/examples/models/muse-glimmer/muse_glimmer_worker" \
  --model-id muse-glimmer-30B \
  --tool-parser atem \
  --host 0.0.0.0 \
  --port 8000 \
  --max-context 131072 \
  --num-runners 1 \
  --max-sessions 1

Send a non-streaming chat-completions request containing a valid PNG:

{
  "model": "muse-glimmer-30B",
  "messages": [{
    "role": "user",
    "content": [
      {"type": "text", "text": "What is in this image? Answer briefly."},
      {"type": "image_url", "image_url": {"url": "data:image/png;base64,<VALID_PNG>"}}
    ]
  }],
  "max_tokens": 8,
  "temperature": 0,
  "stream": false
}

This reproduced with two different valid PNGs, including a standard 512×512 image. Text-only requests return HTTP 200 and complete normally.

Observed response/log:

HTTP/1.1 500 Internal Server Error
content-type: application/json

{"error":{"message":"Generation failed: image decode or vision encoding failed","type":"server_error","code":null}}

Error: CUDA driver error: invalid argument
POST /v1/chat/completions HTTP/1.1 500 Internal Server Error

Compute Sanitizer backtrace

Program hit CUDA_ERROR_INVALID_VALUE (error 1) due to "invalid argument"
on CUDA API call to cuFuncSetAttribute.

Saved host backtrace up to driver entry point at error:
  loadKernel(...) in ...wrapper.cpp:723
  call_triton_tem_fused__unsafe_view_gelu_mm_t_view_28(...) in ...wrapper.cpp:26060
  torch::aot_inductor::AOTInductorModel::run_impl(...) in ...wrapper.cpp:33985
  torch::aot_inductor::AOTInductorModelBase<...>::run(...) in model_base.h:581
  torch::aot_inductor::AOTInductorModelContainer::run(...) in model_container.h:203
  AOTInductorModelContainerRun(...) in ...wrapper.cpp:173
  executorch::backends::cuda::CudaBackend::execute(...)
  executorch::runtime::Method::execute_instruction()
  executorch::runtime::Method::execute()

Compute Sanitizer also reported that full device instrumentation is unsupported under the WSL WDDM debugger interface, but CUDA API error reporting remained active and captured the failing driver call and host backtrace above.

Shared-memory proof from the published vision binary

The extracted vision_encoder AOTInductor shared object contains SM80 cubins plus PTX and the named kernel. Relevant disassembly:

# Kernel-specific loadKernel call:
24025a: ba 00 00 02 00        mov $0x20000,%edx
240262: ...                    <__triton_tem_fused__unsafe_view_gelu_mm_t_view_28_start>
24026c: ...                    call loadKernel(...)

# loadKernel applies CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES (enum 8):
d98fc: ...                     mov <requested bytes>,%edx
d9909: be 08 00 00 00         mov $0x8,%esi
d9911: ...                     call cuFuncSetAttribute@plt

# The same kernel launch is configured for 0x20000 shared bytes:
240342: 41 b9 00 00 02 00     mov $0x20000,%r9d
240351: ...                    call launchKernel(...)

Hardware query on the RTX 5090:

compute capability: 12.0
shared_memory_per_block: 49152
shared_memory_per_block_optin: 101376

Therefore the artifact requests 131,072 bytes, while this GPU permits at most 101,376 bytes for the attribute.

Reproduction: native SM120 export fallback

Official input files were hash-verified, then I ran the documented full export sequentially with CUDA 13 and native SM120 selected:

CUDA_HOME=/usr/local/cuda-13.0 \
CUDACXX=/usr/local/cuda-13.0/bin/nvcc \
PATH="/usr/local/cuda-13.0/bin:$VENV/bin:$PATH" \
LD_LIBRARY_PATH="/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-}" \
TORCH_CUDA_ARCH_LIST=12.0 \
MAX_JOBS=1 \
TORCHINDUCTOR_COMPILE_THREADS=1 \
python -m executorch.examples.models.muse_glimmer.export.export_dflash \
  --target-gguf muse-glimmer-30B-kquant-17gb.gguf \
  --draft-gguf dflash-kquant.gguf \
  --mmproj mmproj-kquant.gguf \
  --backend cuda \
  --output-dir ./muse-glimmer-sm120-native

All attempts end during the same stage:

============================================================
Loading target model...
============================================================
Loading GGUF from .../muse-glimmer-30B-kquant-17gb.gguf...
Building fused model on meta device...
Converting quantized tensors for CUDA...
Killed
EXPORT_FAILED rc=137

OOM evidence across the bounded attempts:

20 GB WSL limit:
  MemTotal: 20484428 kB
  Out of memory: Killed process ... (python) ... anon-rss:19365380kB ...

24 GB WSL limit:
  MemTotal: 24613196 kB
  Out of memory: Killed process ... (python) ... anon-rss:23174952kB ...

27 GB WSL limit:
  MemTotal: 27709772 kB
  Out of memory: Killed process ... (python) ... anon-rss:26301172kB ...

27 GB with unrelated AI services stopped:
  Out of memory: Killed process ... (python) ... anon-rss:26821068kB ...

Final 29 GB WSL limit with unrelated AI services stopped:
  MemTotal: 30488731648 bytes
  available before export: 29454942208 bytes
  swap configured: 12884901888 bytes
  result: killed during "Converting quantized tensors for CUDA", rc=137
  DFlash and vision had not yet loaded; output directory remained empty.

As a diagnostic only, I also released the loader's atomic_sd and fused_sd references immediately after assign_state_dict and forced a collection before CUDA packing. This slightly changed the peak but the isolated export still OOMed, including at the final 29 GB limit. That patch was reverted; the checkout is clean.

Expected behavior

One of the following would make the documented full-feature path usable on an RTX 5090:

1. The published sm80+ptx vision kernel runs within the SM120 shared-memory limit; or
2. An official SM120-native full PTE/PTD is published; or
3. The native export completes within a documented host-memory requirement or has a lower-peak-memory path.

Could you confirm whether an SM120-native full artifact already exists, and whether more than 32 GB of host RAM is expected for export_dflash with the 17G target?

Versions

ExecuTorch commit/package: 43f89fb21aa8f109d66832227f2a797725fb4b9b / executorch 1.5.0+43f89fb
OS: Ubuntu 24.04.3 LTS under WSL2
Kernel: 5.15.167.4-microsoft-standard-WSL2
GPU: NVIDIA GeForce RTX 5090, 32607 MiB, compute capability 12.0
NVIDIA driver: 610.47
Host physical RAM: 32 GB
PyTorch: 2.13.0+cu130
CUDA used to build PyTorch: 13.0
Triton: 3.7.1+gitf797708c
Python: 3.12.3
GCC: 13.3.0
CUDA compiler used by the Muse build/export: /usr/local/cuda-13.0/bin/nvcc
nvcc: CUDA 13.0, build cuda_13.0.r13.0/compiler.36424714_0

torch.utils.collect_env sees a legacy system /usr/bin/nvcc 12.0 first when it probes the default PATH; the Muse build, runtime and export commands explicitly use /usr/local/cuda-13.0, CUDACXX=/usr/local/cuda-13.0/bin/nvcc, CUDA 13 libraries, and PyTorch +cu130 as shown above.

cc @larryliu0820 @mergennachin @cccclai @helunwencser @jackzhxng @digantdesai @Gasoonjia