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

perf(cuda): eliminate NVRTC runtime compilation — all kernels now cubin

Browse Source 
Replaced 4 cudarc::nvrtc::compile_ptx() calls in ml-ppo cuda_nn with
compile_ptx_for_device() — native cubin via nvcc -O3 with disk cache.

Before: virtual PTX → driver JIT (no -O3, no arch targeting, ~100ms first launch)
After: native SASS for exact sm_XX, cached to disk, <10ms load

Files: lstm.rs (2 kernels), linear.rs (1), softmax.rs (1)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-18 07:55:11 +01:00
parent 94b968d800
commit f607335820
3 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/ml-ppo/src/cuda_nn/linear.rs
View File

@@ -73,7 +73,7 @@ extern "C" __global__ void bias_add(
fn compile_bias_add(ctx: &GpuContext) -> Result<cudarc::driver::CudaFunction, MLError> {
let context = ctx.stream.context();
let ptx_result = LINEAR_BIAS_ADD_PTX.get_or_init(|| {
cudarc::nvrtc::compile_ptx(BIAS_ADD_KERNEL)
ml_core::cuda_compile::compile_ptx_for_device(BIAS_ADD_KERNEL, &context)
.map_err(|e| format!("Failed to compile bias_add kernel: {e}"))
});
let ptx = ptx_result.as_ref().map_err(|e| {

4
crates/ml-ppo/src/cuda_nn/lstm.rs
View File

@@ -156,7 +156,7 @@ impl CudaLSTM {
let context = stream.context();
let gate_ptx_result = LSTM_GATE_PTX.get_or_init(|| {
cudarc::nvrtc::compile_ptx(LSTM_GATE_KERNEL)
ml_core::cuda_compile::compile_ptx_for_device(LSTM_GATE_KERNEL, &context)
.map_err(|e| format!("compile gate kernel: {e}"))
});
let gate_ptx = gate_ptx_result.as_ref().map_err(|e| {
@@ -170,7 +170,7 @@ impl CudaLSTM {
})?;
let bias_ptx_result = LSTM_BIAS_ADD_PTX.get_or_init(|| {
cudarc::nvrtc::compile_ptx(BIAS_ADD_KERNEL)
ml_core::cuda_compile::compile_ptx_for_device(BIAS_ADD_KERNEL, &context)
.map_err(|e| format!("compile bias_add: {e}"))
});
let bias_ptx = bias_ptx_result.as_ref().map_err(|e| {

2
crates/ml-ppo/src/cuda_nn/softmax.rs
View File

@@ -119,7 +119,7 @@ struct SoftmaxKernels {
fn compile_softmax_kernels(stream: &Arc<CudaStream>) -> Result<SoftmaxKernels, MLError> {
let context = stream.context();
let ptx_result = SOFTMAX_PTX.get_or_init(|| {
cudarc::nvrtc::compile_ptx(SOFTMAX_KERNEL)
ml_core::cuda_compile::compile_ptx_for_device(SOFTMAX_KERNEL, &context)
.map_err(|e| format!("Failed to compile softmax kernels: {e}"))
});
let ptx = ptx_result.as_ref().map_err(|e| {