diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 584773683..5c70c8c4c 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -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 */ diff --git a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu index dac92f36b..37989ec7a 100644 --- a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu @@ -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;