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

fix(cuda): fix TMA duplicate-label PTX error, guard dead per-thread functions on sm_90

Browse Source 
Three root causes of CUDA_ERROR_INVALID_PTX on H100 addressed:

1. TMA busy-wait label: cooperative_load_tile_tma used a named PTX label
   (TMA_WAIT:) in a __forceinline__ function inlined at ~28 call sites,
   producing duplicate labels in the same PTX function. Replaced with a
   C while loop + selp.b32 to extract the try_wait predicate — no PTX
   labels emitted.

2. mbarrier wait scope: only thread 0 waited for TMA completion while
   lanes 1-31 skipped the entire function body. Now all 32 lanes
   participate in mbarrier.try_wait.parity.acquire, with __syncwarp()
   fences before and after to handle Hopper independent thread scheduling.

3. Dead per-thread functions: q_forward_dueling, q_forward_branching,
   q_forward_distributional, q_forward_dueling_noisy, and their _shmem
   variants were compiled into sm_90 PTX despite never being called by
   the warp kernel. q_forward_distributional alone allocates 600 floats
   (2.4KB) per thread. Guarded all three groups with
   #if __CUDA_ARCH__ < 900 to eliminate ~7KB dead stack from PTX.

Also: removed NUM_ATOMS_MAX=51 cap (no longer needed since distributional
per-thread function is excluded), moved dim_overrides before common_src
to avoid NVRTC macro redefinition warnings.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-12 21:39:06 +01:00
parent 2cfbc2633c
commit a2f61d370e
2 changed files with 66 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -238,13 +238,15 @@ __device__ __forceinline__ void cooperative_load_tile_tma(
}
/* Shared mbarrier for TMA completion tracking.
* Re-initialized per call; callers separate calls with __syncthreads(). */
* Re-initialized per call; callers separate calls with __syncwarp().
* Static __shared__ — one per block, all inline sites share it. */
__shared__ __align__(8) uint64_t __tma_mbar;
if (threadIdx.x == 0) {
const unsigned int mbar_addr =
(unsigned int)__cvta_generic_to_shared(&__tma_mbar);
/* Compute mbar shared address for all threads (same address). */
const unsigned int mbar_addr =
(unsigned int)__cvta_generic_to_shared(&__tma_mbar);
if (threadIdx.x == 0) {
/* Initialize mbarrier: 1 expected arrival (thread 0). */
asm volatile(
"mbarrier.init.shared::cta.b64 [%0], 1;\n"
@@ -254,9 +256,13 @@ __device__ __forceinline__ void cooperative_load_tile_tma(
/* 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. */
* exactly this many bytes.
* Use a temp .b64 register for the required output operand. */
asm volatile(
"mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n"
"{\n"
" .reg .b64 _st;\n"
" mbarrier.arrive.expect_tx.shared::cta.b64 _st, [%0], %1;\n"
"}\n"
: : "r"(mbar_addr), "r"((unsigned int)byte_count) : "memory"
);
@@ -281,23 +287,39 @@ __device__ __forceinline__ void cooperative_load_tile_tma(
offset += chunk;
remaining -= chunk;
}
/* 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 shared memory to all threads.
* This is already the case in every TILE_LAYER_* macro and inline site. */
/* Ensure all lanes see the mbarrier init before waiting.
* Critical on Hopper with independent thread scheduling — without
* this fence, lanes 1-31 could reach try_wait before thread 0 writes
* the mbarrier init. */
__syncwarp(0xFFFFFFFF);
/* ALL lanes wait for the bulk copies to complete.
* mbarrier completes when: arrivals == init_count (1)
* AND tx_bytes_delivered == expected_tx (byte_count).
* Phase parity 0 matches the freshly-initialized mbarrier.
*
* Uses a C loop instead of a PTX label-based loop to avoid
* duplicate-label errors when this __forceinline__ function is
* inlined at multiple call sites in the same kernel. */
{
int _tma_done = 0;
while (!_tma_done) {
asm volatile(
"{\n"
" .reg .pred p;\n"
" mbarrier.try_wait.parity.acquire.shared::cta.b64 p, [%1], 0;\n"
" selp.b32 %0, 1, 0, p;\n"
"}\n"
: "=r"(_tma_done) : "r"(mbar_addr) : "memory"
);
}
}
/* Ensure all lanes completed the wait before caller can reinitialize
* the mbarrier for the next tile in the same forward pass. */
__syncwarp(0xFFFFFFFF);
}
#endif /* __CUDA_ARCH__ >= 900 */

22
crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu
View File

@@ -39,6 +39,13 @@
/* DQN-Specific Device Functions */
/* ------------------------------------------------------------------ */
/* Per-thread forward pass functions — excluded on Hopper (sm_90+).
* The warp kernel uses the _warp_shmem variants instead.
* Excluding these avoids compiling massive stack frames (e.g.,
* q_forward_distributional: 600 floats = 2.4KB/thread) into sm_90
* PTX, which can cause CUDA_ERROR_INVALID_PTX during driver JIT. */
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 900
/**
* Dueling Q-network forward pass (standard mode — no noise, no distributional).
*
@@ -197,6 +204,8 @@ __device__ void q_forward_branching(
q_urgency[i] = value + adv_u[i] - mean_u;
}
#endif /* !sm_90: per-thread q_forward_dueling, q_forward_branching */
/** Argmax over an array of given length. */
__device__ __forceinline__ int argmax_arr(const float* vals, int len) {
int best = 0;
@@ -319,6 +328,10 @@ __device__ __forceinline__ float max_arr(const float* vals, int len) {
} \
} while (0)
/* Per-thread shmem-tiled functions — excluded on Hopper (warp kernel uses
* _warp_shmem variants with distributed vectors instead). */
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 900
/**
* Shared-memory dueling Q-network forward pass.
*
@@ -552,6 +565,8 @@ __device__ void q_forward_distributional_shmem(
}
}
#endif /* !sm_90: per-thread shmem-tiled functions */
/* ------------------------------------------------------------------ */
/* Warp-Cooperative Dueling Forward Passes */
/* ------------------------------------------------------------------ */
@@ -1032,6 +1047,11 @@ __device__ void q_forward_distributional_warp_shmem(
}
}
/* Per-thread noisy/distributional functions — excluded on Hopper.
* Warp kernel uses q_forward_dueling_noisy_warp_shmem and
* q_forward_distributional_warp_shmem instead. */
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 900
/* ------------------------------------------------------------------ */
/* D5: NoisyNet Forward Pass */
/* ------------------------------------------------------------------ */
@@ -1228,6 +1248,8 @@ __device__ void q_forward_distributional(
}
}
#endif /* !sm_90: per-thread noisy/distributional functions */
/** Argmax over NUM_ACTIONS Q-values. */
__device__ __forceinline__ int argmax_q(const float* q_values) {
int best = 0;