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

fix(cuda): use proper mbarrier sync for TMA cp.async.bulk on H100

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-12 21:33:22 +01:00
parent 28e0af7bbe
commit 8494853eb0
3 changed files with 121 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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
}
/**

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

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

2
crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs
View File

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