[aoti-et] Enable multimodal runner for Voxtral on CUDA

[larryliu0820]

This pull request introduces changes to the CUDA workflow, model artifact handling, and multimodal runner logic. The main changes include restructuring the GitHub Actions workflow to separate model export, benchmarking, and end-to-end testing for the Voxtral CUDA pipeline, improving artifact management and reproducibility. Additionally, the multimodal runner now supports automatic conversion of audio tensors to bfloat16, ensuring compatibility with expected input types. There are also enhancements to caching and symbol registration in the CUDA backend, and build system updates to support linking the CUDA backend.

Workflow and Artifact Management Improvements:

• Refactored .github/workflows/cuda.yml to split the Voxtral CUDA pipeline into three jobs: export-voxtral-cuda-artifact (exports and stores model artifacts), benchmark-voxtral-cuda (benchmarks using exported artifacts), and test-voxtral-cuda-e2e (runs full end-to-end tests with artifact download and audio input). Improved artifact handling, reproducibility, and added explicit checks for required files. [1] [2] [3] [4] [5]

Multimodal Runner Logic:

• Added automatic conversion of audio tensors to bfloat16 in MultimodalPrefiller::prefill and implemented a helper function convert_to_bfloat16 in util.h to support this. This ensures that audio inputs match the expected dtype for the encoder, improving robustness for multimodal inference. [1] [2]

CUDA Backend and Caching Enhancements:

• Improved caching logic in common_shims.cpp for tensor strides and sizes by validating cached values and updating them when necessary. This prevents stale cache issues and ensures correct tensor metadata. [1] [2]
• Added dynamic symbol re-registration in CudaBackend to handle multiple shared objects in the same process, ensuring correct execution when switching between models.
• Removed redundant logging statements in CUDA backend for cleaner output. [1] [2]

Build System Updates:

• Updated CMakeLists.txt and executorch-config.cmake to include and link the CUDA backend (aoti_cuda) when building Voxtral and other components, improving build flexibility and CUDA support. [1] [2]

Debugging and Tuning Options:

• Added support for enabling debug compilation in cuda_backend.py via the DEBUG environment variable, allowing easier troubleshooting and development.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/14980

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Macos CI runners unavailable

❌ 6 New Failures, 4 Pending

As of commit afc2159 with merge base 66c3dea ():

NEW FAILURES - The following jobs have failed:

• pull / test-multimodal-linux (gemma3-4b) / linux-job (gh)
  RuntimeError: Command docker exec -t d391e429644cb32020a4e8ce8f03033022daa68aead16d3c25e5694b88a55879 /exec failed with exit code 139
• pull / test-qnn-wheel-packages-linux (3.10) / linux-job (gh)
  RuntimeError: Command docker exec -t 23591f13064fef1f13167b5629a02b4dc2db5fb7cf13b49de25114a255052092 /exec failed with exit code 1
• pull / test-qnn-wheel-packages-linux (3.11) / linux-job (gh)
  RuntimeError: Command docker exec -t 39047d1ddaa896eb401358e930a31b722846a0bd7e6fb0774f99edd45f0965af /exec failed with exit code 1
• pull / test-qnn-wheel-packages-linux (3.12) / linux-job (gh)
  RuntimeError: Command docker exec -t eca9076bb71e07d14ee304e2eba56bfc5a23a4bafd68ab22c23bbfe690ae9c3b /exec failed with exit code 1
• pull / test-samsung-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 35f714861cc3aec1fbb27e02610892e6c0f90c55f7dc43d78d14521a3a1080bc /exec failed with exit code 1
• Test CUDA Builds / check-all-cuda-builds (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Gasoonjia]

Thansk for your great work!
The size/stride change for me is pretty strange: i con't image a case that the tensor ptr keeps the same while its size/stride got changed

[larryliu0820]

@swolchok take another look?

[swolchok]

objections withdrawn. I think with some work you can further simplify the sizes()/strides() update stuff, up to you how much of it you want to do right now

[mergennachin]

See inline

[mergennachin]

Can you make docblock something like this?

// CRITICAL: Multimodal models reuse tensors with different shapes across
// executions (e.g., variable-length audio). We MUST validate cached metadata
// matches current tensor state, or CUDA kernels will receive incorrect shapes
// leading to memory corruption and segfaults.

[mergennachin]

If we're loading the model once and doing execute/inference multiple times, it will register multiple times, no?

Can you do something like this?

  void* last_registered_handle = nullptr;

  if (handle->so_handle != last_registered_handle) {
      ET_CHECK_OK_OR_RETURN_ERROR(
          register_shared_library_functions(handle->so_handle));
      last_registered_handle = handle->so_handle;
  }

[larryliu0820]

So the so_handle won't change. It's just we are mapping the symbols differently, especially AOTInductorModelContainerRun. Let's say we do the following:

1. load(token_embedding)
2. load(audio_encoder)
3. load(text_decoder)
4. run(audio_encoder) <-- here AOTInductorModelContainerRun maps to the symbol in text_decoder.so, so we need to remap the symbol to audio_encoder.so

[mergennachin]

@larryliu0820

Can you store the AOTInductorModelContainerRunFunc inside AOTIDelegateHandle?

  struct AOTIDelegateHandle {
      void* so_handle;
      std::string so_path;
      AOTInductorModelContainerHandle container_handle;
      void* cuda_stream;

      AOTInductorModelContainerRunFunc run_func;
      // ... etc for all symbols
  };

 Result<DelegateHandle*> init(...) const override {
      AOTIDelegateHandle* handle = new AOTIDelegateHandle();
      handle->so_handle = so_handle;

      // Load symbols into THIS handle's struct (not global)
      handle->run_func = reinterpret_cast<AOTInductorModelContainerRunFunc>(
          dlsym(so_handle, "AOTInductorModelContainerRun"));
      // ... etc

      ET_CHECK_OR_RETURN_ERROR(
          handle->run_func != nullptr,
          AccessFailed,
          "Failed to load AOTInductorModelContainerRun");

      return (DelegateHandle*)handle;
  }

  Error execute(..., DelegateHandle* handle_, ...) const override {
      AOTIDelegateHandle* handle = (AOTIDelegateHandle*)handle_;

      // NO re-registration, use the handle's local symbols
      AOTIRuntimeError error = handle->run_func(
          ...)

      // ... rest of execution ...
  }

[mergennachin]

Also can you update the https://github.com/pytorch/executorch/blob/main/examples/models/voxtral/README.md to include additional CUDA instructions too?