From 8ea400e59970ba88b6aa90f9f0352a190133ce2d Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sun, 15 Mar 2026 15:10:10 +0100
Subject: [PATCH] fix(cuda): fix TMA inline asm syntax for CUDA 12.9 PTX ISA
 8.7
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The CI toolkit (CUDA 12.9, PTX ISA 8.7) requires:
1. cp.async.bulk needs .mbarrier::complete_tx::bytes completion mechanism
   (mandatory since PTX ISA 8.3 / CUDA 12.3)
2. mbarrier.try_wait.parity.acquire needs .cta scope qualifier between
   .acquire and .shared::cta

Reverts the DISABLE_TMA workaround — TMA now compiles natively to cubin.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../ml/src/cuda_pipeline/common_device_functions.cuh   |  4 ++--
 .../ml/src/cuda_pipeline/gpu_experience_collector.rs   |  3 +--
 crates/ml/src/cuda_pipeline/mod.rs                     | 10 ++--------
 3 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh
index 72eeaff5f..75d52ccb1 100644
--- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh
+++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh
@@ -304,7 +304,7 @@ __device__ __forceinline__ void cooperative_load_tile_tma(
             int chunk = remaining < CHUNK_FLOATS ? remaining : CHUNK_FLOATS;
             int chunk_bytes = chunk * (int)sizeof(float);
             asm volatile(
-                "cp.async.bulk.shared::cta.global [%0], [%1], %2, [%3];\n"
+                "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n"
                 :
                 : "r"((unsigned int)__cvta_generic_to_shared(shmem_dst + offset)),
                   "l"((unsigned long long)(global_src + offset)),
@@ -337,7 +337,7 @@ __device__ __forceinline__ void cooperative_load_tile_tma(
             asm volatile(
                 "{\n"
                 "  .reg .pred p;\n"
-                "  mbarrier.try_wait.parity.acquire.shared::cta.b64 p, [%1], 0;\n"
+                "  mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%1], 0;\n"
                 "  selp.b32 %0, 1, 0, p;\n"
                 "}\n"
                 : "=r"(_tma_done) : "r"(mbar_addr) : "memory"
diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
index 1b6d42c29..905832437 100644
--- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
@@ -441,8 +441,7 @@ impl GpuExperienceCollector {
         // kernel (7.5 KB stack/thread) — only the warp-cooperative kernel is compiled.
         //
         // Compilation produces native cubin (SASS), not PTX.
-        // TMA disabled globally in compile_ptx_for_device — CI toolkit can't assemble
-        // cp.async.bulk instructions. float4 cooperative loads are used instead.
+        // TMA (cp.async.bulk) on sm_90+ compiles natively via nvcc — no driver JIT.
         let cubin = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
             .map_err(|e| MLError::ModelError(format!(
                 "CUDA experience kernel compilation failed: {e}"
diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs
index 6f97413f5..54ebd73b7 100644
--- a/crates/ml/src/cuda_pipeline/mod.rs
+++ b/crates/ml/src/cuda_pipeline/mod.rs
@@ -91,8 +91,8 @@ pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u3
 ///
 /// Compiles to **cubin** (native SASS binary), NOT PTX. This eliminates driver
 /// JIT entirely — `cuModuleLoad` loads the cubin directly. TMA instructions
-/// (cp.async.bulk) are disabled via `DISABLE_TMA` — the CI CUDA toolkit cannot
-/// assemble them. float4 cooperative loads are used as the working fallback.
+/// (cp.async.bulk on sm_90+) compile natively because nvcc handles them during
+/// offline compilation; the driver never sees PTX.
 ///
 /// On first compilation, nvcc compiles the source to cubin and caches the result
 /// to `$CARGO_TARGET_DIR/.cubin_cache/` (CI PVC) or `/tmp/.cubin_cache/` (fallback).
@@ -139,12 +139,6 @@ pub fn compile_ptx_for_device(
         _ => "sm_80",        // safe default
     };
 
-    // Disable TMA (cp.async.bulk) unconditionally — the CI CUDA toolkit's ptxas
-    // cannot assemble TMA instructions on sm_90. float4 cooperative loads are
-    // the working fallback path (same performance on non-TMA workloads).
-    let src_with_tma_disabled = format!("#define DISABLE_TMA 1\n{src}");
-    let src = src_with_tma_disabled.as_str();
-
     // Try loading cached cubin first
     if let Some(cached) = load_cached_cubin(arch_str, src) {
         return Ok(cached);