gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(cuda): fix TMA inline asm syntax for CUDA 12.9 PTX ISA 8.7

Browse Source 
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>
This commit is contained in:
jgrusewski
2026-03-15 15:10:10 +01:00
parent 8c4861fb70
commit 8ea400e599
3 changed files with 5 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/ml/src/cuda_pipeline/common_device_functions.cuh
View File

@@ -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"

3
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -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}"

10
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -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);