From 8494853eb06af89dcacdffee56e4e0a30d34ed5d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 12 Mar 2026 21:33:22 +0100 Subject: [PATCH] fix(cuda): use proper mbarrier sync for TMA cp.async.bulk on H100 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The TMA inline PTX had three invalid instructions causing CUDA_ERROR_INVALID_PTX at driver JIT time on sm_90 (H100): 1. cp.async.bulk missing required 4th mbarrier operand 2. cp.async.bulk.commit_group — does not exist in any PTX ISA 3. cp.async.bulk.wait_group — does not exist in any PTX ISA These confused two different CUDA instruction families: - cp.async (Ampere, uses commit_group/wait_group, 16B per op) - cp.async.bulk (Hopper TMA, uses mbarrier, up to 256KB per op) Fix: rewrite cooperative_load_tile_tma() with correct Hopper protocol: mbarrier.init → mbarrier.arrive.expect_tx → cp.async.bulk [mbar] → mbarrier.try_wait.parity Also adds 16-byte alignment guard (cp.async.bulk requires size % 16 == 0) with float4 fallback for unaligned bias tiles (e.g. NUM_ACTIONS=5 → 20B). Retains DISABLE_TMA retry in gpu_experience_collector as safety net. Co-Authored-By: Claude Opus 4.6 --- .../cuda_pipeline/common_device_functions.cuh | 102 +++++++++++------- .../cuda_pipeline/gpu_experience_collector.rs | 72 +++++++++---- .../dqn/smoke_tests/training_stability.rs | 2 + 3 files changed, 121 insertions(+), 55 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 374627221..584773683 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -200,73 +200,103 @@ __device__ __forceinline__ void cooperative_load_tile_float4( /** * TMA-accelerated tile load from global to shared memory (Hopper sm_90+). * - * Uses cp.async.bulk.shared::cta.global to perform a bulk asynchronous copy - * from global (HBM) to shared memory via the Tensor Memory Accelerator. - * Only thread 0 issues the copy instruction; all other threads are freed - * for compute overlap. The caller's __syncthreads() ensures completion - * visibility (TMA commits to the async group, wait_group 0 drains it). + * Uses cp.async.bulk.shared::cta.global with mbarrier synchronization to + * perform bulk asynchronous copies from global (HBM) to shared memory via + * the Tensor Memory Accelerator. Only thread 0 issues the copy + mbarrier + * instructions; all other threads are freed for compute overlap. * - * cp.async.bulk transfers up to 128 bytes per instruction in the base form, - * so we issue ceil(byte_count / 128) sequential bulk copies from thread 0. - * This still wins over the float4 scatter because: - * (a) TMA bypasses L1/L2 pollution — direct HBM→SMEM path - * (b) No warp divergence — 31 threads skip straight to compute - * (c) Overlaps with any independent arithmetic on the other warps + * Hopper TMA uses mbarrier (NOT commit_group/wait_group — those are for the + * older cp.async non-bulk instruction family from Ampere). The protocol: + * 1. mbarrier.init — set expected arrival count (1 = thread 0) + * 2. mbarrier.arrive.expect_tx — arrive AND declare expected TX bytes + * 3. cp.async.bulk ... [mbar] — each copy auto-adds bytes to TX counter + * 4. mbarrier.try_wait.parity — spin until arrivals + TX bytes complete * - * Fallback: tiles larger than practical TMA dispatch (>64KB, rare) fall - * through to the float4 cooperative path to avoid excessive PTX overhead. + * Falls back to float4 for: + * - Non-16-byte-aligned sizes (cp.async.bulk requires size % 16 == 0) + * - Very large tiles (>256KB, should not happen in practice) * * Constraints: * - src must be 16-byte aligned (guaranteed by cudarc device alloc) * - shmem_dst must be 16-byte aligned (guaranteed by __shared__ layout) * - count must be > 0 - * - Caller MUST issue __syncthreads() (or mbarrier) after return + * - Caller MUST issue __syncthreads() after return */ __device__ __forceinline__ void cooperative_load_tile_tma( float* __restrict__ shmem_dst, const float* __restrict__ global_src, int count /* total floats to load */ ) { - /* cp.async.bulk max per-instruction is limited; we chunk at 16KB - * (4096 floats) which is well within the hardware capability and - * keeps PTX emission tractable. Each chunk is a single bulk copy. */ - const int CHUNK_FLOATS = 4096; /* 16 KB per chunk */ const int byte_count = count * (int)sizeof(float); - /* Fallback for very large tiles (>256KB): should not happen with - * SHMEM_TILE_ROWS=64 × SHMEM_MAX_IN_DIM=256 = 64KB, but guard. */ - if (byte_count > 256 * 1024) { + /* cp.async.bulk requires byte count to be a multiple of 16. + * Bias tiles (e.g., NUM_ACTIONS=5 → 20 bytes) violate this. + * Also guard against very large tiles (>256KB). */ + if ((byte_count & 15) != 0 || byte_count > 256 * 1024) { cooperative_load_tile_float4(shmem_dst, global_src, count); return; } + /* Shared mbarrier for TMA completion tracking. + * Re-initialized per call; callers separate calls with __syncthreads(). */ + __shared__ __align__(8) uint64_t __tma_mbar; + if (threadIdx.x == 0) { - /* Issue bulk async copies in chunks. - * cp.async.bulk.shared::cta.global [dst], [src], size; - * size operand is in bytes, immediate or register. */ + const unsigned int mbar_addr = + (unsigned int)__cvta_generic_to_shared(&__tma_mbar); + + /* Initialize mbarrier: 1 expected arrival (thread 0). */ + asm volatile( + "mbarrier.init.shared::cta.b64 [%0], 1;\n" + : : "r"(mbar_addr) : "memory" + ); + + /* Arrive with expected TX byte count. + * This counts as the 1 arrival AND sets the expected TX bytes. + * The mbarrier will complete once all cp.async.bulk copies deliver + * exactly this many bytes. */ + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + : : "r"(mbar_addr), "r"((unsigned int)byte_count) : "memory" + ); + + /* Issue bulk copies in 16KB chunks. + * cp.async.bulk.shared::cta.global [dst], [src], size, [mbar]; + * Each copy auto-increments the mbarrier TX counter on completion. */ + const int CHUNK_FLOATS = 4096; /* 16 KB per chunk */ int remaining = count; int offset = 0; while (remaining > 0) { int chunk = remaining < CHUNK_FLOATS ? remaining : CHUNK_FLOATS; int chunk_bytes = chunk * (int)sizeof(float); - /* Inline PTX: cp.async.bulk.shared::cta.global - * Operands: .dst (shared), .src (global), .size (bytes) */ asm volatile( - "cp.async.bulk.shared::cta.global [%0], [%1], %2;\n" + "cp.async.bulk.shared::cta.global [%0], [%1], %2, [%3];\n" : : "r"((unsigned int)__cvta_generic_to_shared(shmem_dst + offset)), "l"((unsigned long long)(global_src + offset)), - "r"((unsigned int)chunk_bytes) + "r"((unsigned int)chunk_bytes), + "r"(mbar_addr) : "memory" ); offset += chunk; remaining -= chunk; } - /* Commit all issued bulk copies to an async group and wait. */ - asm volatile("cp.async.bulk.commit_group;\n" ::: "memory"); - asm volatile("cp.async.bulk.wait_group 0;\n" ::: "memory"); + + /* Wait for all bulk copies to complete. + * mbarrier completes when: arrivals == init_count (1) + * AND tx_bytes_delivered == expected_tx (byte_count). + * Phase 0 is correct for a freshly-initialized mbarrier. */ + asm volatile( + "{\n" + ".reg .pred p;\n" + "TMA_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n" + "@!p bra TMA_WAIT;\n" + "}\n" + : : "r"(mbar_addr) : "memory" + ); } - /* Caller issues __syncthreads() to broadcast completion to all threads. + /* Caller issues __syncthreads() to broadcast shared memory to all threads. * This is already the case in every TILE_LAYER_* macro and inline site. */ } #endif /* __CUDA_ARCH__ >= 900 */ @@ -283,11 +313,11 @@ __device__ __forceinline__ void cooperative_load_tile( const float* __restrict__ global_src, int count /* total floats to load */ ) { - /* Always use float4 cooperative loads. TMA (cp.async.bulk) inline PTX - * causes CUDA_ERROR_INVALID_PTX on H100 sm_90 driver JIT due to PTX ISA - * version mismatch between NVRTC-generated PTX and the host CUDA driver. - * float4 loads are still fast: all 32 warp threads load in parallel. */ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && !defined(DISABLE_TMA) + cooperative_load_tile_tma(shmem_dst, global_src, count); +#else cooperative_load_tile_float4(shmem_dst, global_src, count); +#endif } /** diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 48aba4c5c..d5e5139ed 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -432,25 +432,59 @@ impl GpuExperienceCollector { // On Hopper (sm_90+) __CUDA_ARCH__=900 excludes the standard per-thread // kernel (7.5 KB stack/thread causes CUDA_ERROR_INVALID_PTX on sm_90 JIT). // Only the warp-cooperative kernel (200 bytes/lane) is compiled. - let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) - .map_err(|e| { - MLError::ModelError(format!( - "CUDA dqn_experience_kernel compilation failed (is NVRTC available?): {e}" - )) - })?; - info!( - source_len = full_source.len(), - sm_major, - use_warp_kernel, - "NVRTC compilation succeeded, loading PTX module..." - ); - let module = context.load_module(ptx).map_err(|e| { - MLError::ModelError(format!( - "Failed to load dqn_experience_kernel module (sm_{sm_major}0, \ - source={} chars, warp={use_warp_kernel}): {e}", - full_source.len() - )) - })?; + // + // TMA fallback: On sm_90+ we first try with cp.async.bulk TMA tile loads. + // If the driver JIT rejects the PTX (ISA version mismatch or driver bug), + // we retry with DISABLE_TMA which falls back to float4 cooperative loads. + let module = { + let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) + .map_err(|e| { + MLError::ModelError(format!( + "CUDA dqn_experience_kernel compilation failed (is NVRTC available?): {e}" + )) + })?; + info!( + source_len = full_source.len(), + sm_major, + use_warp_kernel, + "NVRTC compilation succeeded, loading PTX module (TMA enabled)..." + ); + match context.load_module(ptx) { + Ok(m) => { + info!("PTX module loaded successfully (TMA enabled)"); + m + } + Err(tma_err) => { + // TMA PTX failed — retry without TMA inline asm. + // cp.async.bulk can fail on older drivers or certain H100 firmware. + tracing::warn!( + "PTX module load failed with TMA (sm_{sm_major}0): {tma_err}. \ + Retrying with DISABLE_TMA (float4 cooperative loads)..." + ); + let disable_tma = "#define DISABLE_TMA 1\n"; + let full_source_no_tma = format!( + "{common_src}\n{disable_tma}{dim_overrides}\n{kernel_src}" + ); + let ptx_no_tma = crate::cuda_pipeline::compile_ptx_for_device( + &full_source_no_tma, &context, + ) + .map_err(|e| { + MLError::ModelError(format!( + "CUDA dqn_experience_kernel re-compilation (no TMA) failed: {e}" + )) + })?; + let m = context.load_module(ptx_no_tma).map_err(|e| { + MLError::ModelError(format!( + "Failed to load dqn_experience_kernel module even without TMA \ + (sm_{sm_major}0, source={} chars): {e}", + full_source_no_tma.len() + )) + })?; + info!("PTX module loaded successfully (TMA disabled, float4 fallback)"); + m + } + } + }; // Load kernel functions based on compute capability. // sm_90+ (H100/Hopper): only warp kernel exists in PTX (standard excluded by #if guard). diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index 47fb8377f..3292f884c 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -120,6 +120,7 @@ async fn test_no_nan_loss_short_training() -> anyhow::Result<()> { } /// PER importance-sampling weights must be finite and positive. +#[cfg(feature = "cuda")] #[tokio::test] async fn test_per_weights_in_valid_range() -> anyhow::Result<()> { use candle_core::{DType, Device, Tensor}; @@ -161,6 +162,7 @@ async fn test_per_weights_in_valid_range() -> anyhow::Result<()> { } /// PER sampled indices must be within buffer bounds. +#[cfg(feature = "cuda")] #[tokio::test] async fn test_per_indices_in_valid_range() -> anyhow::Result<()> { use candle_core::{DType, Device, Tensor};