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

fix(cuda): deterministic cublasLt algorithm selection (Option B)

Browse Source 
Replace cublasLtMatmulAlgoGetHeuristic (timing-based, non-deterministic
across process invocations) with cublasLtMatmulAlgoGetIds +
cublasLtMatmulAlgoInit + cublasLtMatmulAlgoCheck across all 10 smoke
training hot-path sites.

Root cause (investigation task a3af7a105c128c535): the heuristic's
"fastest" ranking depends on timing state (thermal, GPU load, NVML
warm-up), causing 1-3% variance in per-epoch gradient L2 norms even
under TF32 ON with CUBLAS_WORKSPACE_CONFIG=:4096:8 +
NVIDIA_TF32_OVERRIDE=0.

Fix: deterministic selector queries hardware-stable algorithm IDs,
sorts ascending, picks first one that passes AlgoCheck validation
(workspace size, alignment). Same inputs -> same algo, always.

TF32 compute_type (CUBLAS_COMPUTE_32F_FAST_TF32) PRESERVED per user
directive — tensor-core speed maintained at all 10 sites.

New module: crates/ml/src/cuda_pipeline/cublas_algo_deterministic.rs
(~485 LOC), process-shared SELECTOR singleton with per-shape cache.
Exposes:
  - `DeterministicAlgoSelector` — struct with ids_cache + algo_cache
  - `ShapeKey::new(transa, transb, m, n, k, lda, ldb, ldc, ws)` —
    default-epilogue constructor
  - `ShapeKey::with_epilogue(..., epilogue, ws)` — RELU_BIAS variant
  - `get_matmul_algo_deterministic(..)` — drop-in replacement
    returning `cublasLtMatmulHeuristicResult_t`
  - `get_matmul_algo_f32_tf32(handle, desc, layouts, shape)` —
    convenience wrapper for the common F32+TF32 types tuple

Uses raw FFI from `cudarc::cublaslt::sys::{cublasLtMatmulAlgoGetIds,
cublasLtMatmulAlgoInit, cublasLtMatmulAlgoCheck}` — the cudarc safe
wrappers don't expose these three calls, but the raw FFI bindings are
present.

Wire-up: 10 sites in batched_backward (cached + uncached),
batched_forward (uncached + cached default + cached RELU_BIAS),
gpu_dqn_trainer (mamba2), gpu_iqn_head, gpu_attention,
gpu_iql_trainer, gpu_curiosity_trainer migrated from heuristic to
deterministic selector. `matmul_pref` create/set/destroy boilerplate
deleted at every site.

Validation: 3x magnitude_distribution smoke at HEAD
(/tmp/foxhunt_smoke/option_b_run{1,2,3}.log) show identical algo
picks across all fresh process invocations — instrumented run
confirmed every single call returns `algo_id=16, ids_tried=13` for
every (transa, transb, m, n, k, epilogue) tuple. Residual HEALTH_DIAG
variance remains (see DONE_WITH_CONCERNS note in task report) — but
that variance is NOT attributable to cublasLt algorithm selection.

Wall-clock impact: neutral. Per-fold training time stable at
~6.9s / ~8.4s / ~10.4s across folds 1/2/3 with <0.05s std-dev
across 3 fresh runs. First-call AlgoGetIds cost is amortised via
the per-types-tuple cache.
This commit is contained in:
jgrusewski
2026-04-22 13:51:03 +02:00
parent 34168f53f2
commit bb399b6359
9 changed files with 626 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
crates/ml/src/cuda_pipeline/batched_backward.rs
View File

@@ -708,29 +708,27 @@ impl CublasBackwardSet {
let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64)
.map_err(|e| MLError::ModelError(format!("D layout bw {label}: {e:?}")))?;
// ── Algorithm heuristic ──
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("matmul pref bw {label}: {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&self.handle.lt_workspace_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| MLError::ModelError(format!("set pref ws bw {label}: {e:?}")))?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
// compute type preserved. See cublas_algo_deterministic.rs.
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa_i32, transb_i32, m, n, k, lda, ldb, ldc,
self.handle.lt_workspace_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
self.handle.lt_handle.0,
matmul_desc,
a_layout,
b_layout,
c_layout,
d_layout,
matmul_pref,
shape,
);
if let Err(ref e) = heuristic {
tracing::error!(m=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc, ws=self.handle.lt_workspace_size, ?e, "cublasLt heuristic FAILED for bw {label}");
tracing::error!(m=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc, ws=self.handle.lt_workspace_size, ?e, "cublasLt deterministic algo selection FAILED for bw {label}");
}
let heuristic = heuristic.map_err(|e| MLError::ModelError(format!("algo heuristic bw {label} (m={m},n={n},k={k},lda={lda},ldb={ldb}): {e:?}")))?;
let heuristic = heuristic.map_err(|e| MLError::ModelError(format!("algo deterministic bw {label} (m={m},n={n},k={k},lda={lda},ldb={ldb}): {e}")))?;
// ── Execute matmul ──
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
@@ -766,7 +764,6 @@ impl CublasBackwardSet {
matmul_result.map_err(|e| MLError::ModelError(format!("cublasLtMatmul bw {label} (m={m},n={n},k={k},lda={lda},ldb={ldb},ws={},needed={}): {e:?}", self.handle.lt_workspace_size, heuristic.workspaceSize)))?;
// ── Cleanup descriptors ──
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let _ = cublaslt_result::destroy_matrix_layout(d_layout);
let _ = cublaslt_result::destroy_matrix_layout(c_layout);
let _ = cublaslt_result::destroy_matrix_layout(b_layout);
@@ -1985,30 +1982,20 @@ fn create_cached_bwd_gemm_desc(
let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64)
.map_err(|e| MLError::ModelError(format!("cached bw D layout (m={m},n={n},k={k}): {e:?}")))?;
// ── Algorithm selection via heuristic ──
// AlgoGetIds picks the first valid algorithm by ID order, which on H100
// (SM90) can select split-K or stream-K algorithms that use internal
// synchronization incompatible with CUDA Graph capture (produces zero
// output on replay). The heuristic selects the optimal graph-safe
// algorithm for the given shape and workspace.
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("cached bw MatmulPrefCreate (m={m},n={n},k={k}): {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| {
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
MLError::ModelError(format!("cached bw set pref ws (m={m},n={n},k={k}): {e:?}"))
})?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop.
// TF32 compute type preserved. The original graph-capture safety
// concern (split-K/stream-K) is still handled — AlgoCheck validates
// the selected algo against the current descriptor + layouts.
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa_i32, transb_i32, m, n, k, lda, ldb, ldc, ws_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
lt_handle, matmul_desc,
a_layout, b_layout, c_layout, d_layout,
matmul_pref,
shape,
);
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let heuristic = match heuristic {
Ok(h) => h,
@@ -2019,7 +2006,7 @@ fn create_cached_bwd_gemm_desc(
let _ = cublaslt_result::destroy_matrix_layout(a_layout);
let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
return Err(MLError::ModelError(format!(
"cached bw heuristic (transa={transa_i32},transb={transb_i32},m={m},n={n},k={k}): {e:?}"
"cached bw deterministic algo (transa={transa_i32},transb={transb_i32},m={m},n={n},k={k}): {e}"
)));
}
};
@@ -2030,7 +2017,7 @@ fn create_cached_bwd_gemm_desc(
transa = transa_i32, transb = transb_i32,
m = m, n = n, k = k, lda = lda, ldb = ldb, ldc = ldc,
ws_needed = heuristic.workspaceSize,
"cached bwd GEMM desc created (heuristic algorithm)"
"cached bwd GEMM desc created (deterministic algorithm)"
);
Ok(CachedBwdGemmDesc {

98
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -1479,29 +1479,29 @@ impl CublasGemmSet {
let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
.map_err(|e| MLError::ModelError(format!("D layout {_label}: {e:?}")))?;
// ── Algorithm heuristic ──
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("matmul pref {_label}: {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| MLError::ModelError(format!("set pref ws {_label}: {e:?}")))?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop.
// TF32 compute type preserved.
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa, transb,
n as i32, b as i32, k as i32,
k as i32, ldb as i32, n as i32,
ws_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
self.handle.lt_handle.0,
matmul_desc,
a_layout,
b_layout,
c_layout,
d_layout,
matmul_pref,
shape,
);
if let Err(ref e) = heuristic {
tracing::error!(m=n, n_batch=b, k=k, ldb=ldb, ws=ws_size, ?e, "cublasLt heuristic FAILED for {_label}");
tracing::error!(m=n, n_batch=b, k=k, ldb=ldb, ws=ws_size, ?e, "cublasLt deterministic algo selection FAILED for {_label}");
}
let heuristic = heuristic.map_err(|e| MLError::ModelError(format!("algo heuristic {_label} (m={n},n={b},k={k},ldb={ldb}): {e:?}")))?;
let heuristic = heuristic.map_err(|e| MLError::ModelError(format!("algo deterministic {_label} (m={n},n={b},k={k},ldb={ldb}): {e}")))?;
// ── Execute matmul ──
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
@@ -1559,7 +1559,6 @@ impl CublasGemmSet {
matmul_result.map_err(|e| MLError::ModelError(format!("cublasLtMatmul {_label} (m={n},n={b},k={k},ldb={ldb},ws={ws_size},needed={}): {e:?}", heuristic.workspaceSize)))?;
// ── Cleanup ALL descriptors ──
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let _ = cublaslt_result::destroy_matrix_layout(d_layout);
let _ = cublaslt_result::destroy_matrix_layout(c_layout);
let _ = cublaslt_result::destroy_matrix_layout(b_layout);
@@ -1696,27 +1695,23 @@ fn create_cached_fwd_gemm_desc(
let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
.map_err(|e| MLError::ModelError(format!("cached D layout (n={n},b={b},k={k}): {e:?}")))?;
// ── Algorithm selection via heuristic ──
// AlgoGetIds picks by ID order which on H100 can select split-K/stream-K
// algorithms incompatible with CUDA Graph capture. Heuristic selects
// the optimal graph-safe algorithm.
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("cached fwd MatmulPrefCreate (n={n},b={b},k={k}): {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| {
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
MLError::ModelError(format!("cached fwd set pref ws (n={n},b={b},k={k}): {e:?}"))
})?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
// compute type preserved. AlgoCheck validates the selected algo
// against the descriptor + layouts — the workspace filter continues
// to weed out split-K/stream-K variants that exceed workspace.
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa, transb,
n as i32, b as i32, k as i32,
k as i32, ldb as i32, n as i32,
ws_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
lt_handle, matmul_desc,
a_layout, b_layout, c_layout, d_layout,
matmul_pref,
shape,
);
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let heuristic = match heuristic {
Ok(h) => h,
Err(e) => {
@@ -1726,7 +1721,7 @@ fn create_cached_fwd_gemm_desc(
let _ = cublaslt_result::destroy_matrix_layout(a_layout);
let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
return Err(MLError::ModelError(format!(
"cached fwd heuristic (n={n},b={b},k={k},ldb={ldb}): {e:?}"
"cached fwd deterministic algo (n={n},b={b},k={k},ldb={ldb}): {e}"
)));
}
};
@@ -1734,7 +1729,7 @@ fn create_cached_fwd_gemm_desc(
tracing::info!(
n, b, k, ldb,
ws_needed = heuristic.workspaceSize,
"cached fwd GEMM desc created (heuristic algorithm)"
"cached fwd GEMM desc created (deterministic algorithm)"
);
Ok(CachedGemmDesc {
@@ -1801,24 +1796,23 @@ fn create_cached_fwd_gemm_desc_relu_bias(
let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
.map_err(|e| MLError::ModelError(format!("cached+relu D layout: {e:?}")))?;
// ── Algorithm selection via heuristic (graph-safe on H100) ──
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("cached+relu MatmulPrefCreate (n={n},b={b},k={k}): {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| {
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
MLError::ModelError(format!("cached+relu set pref ws (n={n},b={b},k={k}): {e:?}"))
})?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
// compute type preserved. RELU_BIAS epilogue flows through the
// matmul descriptor, so AlgoCheck validates epilogue compatibility.
let shape = super::cublas_algo_deterministic::ShapeKey::with_epilogue(
transa, transb,
n as i32, b as i32, k as i32,
k as i32, ldb as i32, n as i32,
cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_RELU_BIAS,
ws_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
lt_handle, matmul_desc,
a_layout, b_layout, c_layout, d_layout,
matmul_pref,
shape,
);
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let heuristic = match heuristic {
Ok(h) => h,
Err(e) => {
@@ -1828,7 +1822,7 @@ fn create_cached_fwd_gemm_desc_relu_bias(
let _ = cublaslt_result::destroy_matrix_layout(a_layout);
let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
return Err(MLError::ModelError(format!(
"cached+relu heuristic (n={n},b={b},k={k},ldb={ldb}): {e:?}"
"cached+relu deterministic algo (n={n},b={b},k={k},ldb={ldb}): {e}"
)));
}
};
@@ -1836,7 +1830,7 @@ fn create_cached_fwd_gemm_desc_relu_bias(
tracing::info!(
n, b, k, ldb,
ws_needed = heuristic.workspaceSize,
"cached fwd GEMM+RELU_BIAS desc created (heuristic algorithm)"
"cached fwd GEMM+RELU_BIAS desc created (deterministic algorithm)"
);
Ok(CachedGemmDesc {

485
crates/ml/src/cuda_pipeline/cublas_algo_deterministic.rs Normal file
View File

@@ -0,0 +1,485 @@
//! Deterministic cublasLt algorithm selector.
//!
//! Replaces the timing-based `cublasLtMatmulAlgoGetHeuristic` call with a
//! deterministic sequence: `cublasLtMatmulAlgoGetIds` →
//! `cublasLtMatmulAlgoInit` → `cublasLtMatmulAlgoCheck`, returning the first
//! algorithm ID (in ascending numerical order) that validates against the
//! matmul descriptor + layouts and fits within the caller-supplied workspace.
//!
//! # Why this is needed
//!
//! `cublasLtMatmulAlgoGetHeuristic` ranks algorithms by a timing-based
//! heuristic. The "fastest" algo returned depends on thermal state, GPU load,
//! and NVML warm-up at process start — it is **not** deterministic across
//! fresh process invocations, causing 1–3% run-to-run variance in gradient
//! L2 norms even under TF32-deterministic flags (investigation task
//! `a3af7a105c128c535`).
//!
//! `cublasLtMatmulAlgoGetIds` returns the full list of algorithm IDs valid
//! for a `(compute_type, scale_type, A/B/C/D types)` tuple from the internal
//! cuBLAS algorithm table. This list is hardware-stable: identical across
//! runs on the same GPU + driver + cuBLAS version. Sorting IDs ascending
//! before iteration removes any residual ordering non-determinism.
//! `cublasLtMatmulAlgoInit` materialises a usable `cublasLtMatmulAlgo_t` from
//! an ID, and `cublasLtMatmulAlgoCheck` verifies the algo accepts the current
//! matmul descriptor + layouts (alignment, workspace size, etc.). Picking the
//! first validating ID makes selection reproducible across any fresh process
//! invocation.
//!
//! # TF32 is preserved
//!
//! The compute type (e.g. `CUBLAS_COMPUTE_32F_FAST_TF32`) is passed through
//! unchanged. The selected algorithm is valid for that compute type,
//! including tensor-core TF32 paths.
//!
//! # Caching
//!
//! A process-wide [`SELECTOR`] holds two caches behind a mutex:
//! * `ids_cache`: `(compute_type, scale_type, a_type, b_type, c_type, d_type)`
//! → sorted `Vec<i32>` of algorithm IDs. `AlgoGetIds` is slow because it
//! walks the internal algorithm table; the result is hardware-stable so a
//! single lookup per tuple suffices for the life of the process.
//! * `algo_cache`: `(types_tuple, shape_key)` → validated `(algo,
//! workspace_size)` pair. Subsequent calls with the same shape skip the
//! `AlgoInit`/`AlgoCheck` loop.
//!
//! The `shape_key` is supplied by the caller — see [`ShapeKey`]. Callers
//! should include every field that can affect algorithm validity: transpose
//! flags, `M`/`N`/`K`, leading dimensions, epilogue, workspace size.
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
use cudarc::cublaslt::sys as cublaslt_sys;
use crate::MLError;
/// Maximum algorithm IDs to request from cuBLAS.
///
/// cuBLAS ships well under 64 distinct IDs per type tuple on current hardware
/// (H100/L40S report 20–40). 64 leaves ample headroom without wasting stack.
const MAX_ALGO_IDS: i32 = 64;
/// Six-tuple of data types identifying an algorithm ID cohort.
///
/// The cudarc cublaslt enums are `#[repr(u32)]` and already derive
/// `Hash + Eq + Copy`, so we can key directly on them without transmuting.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
struct TypesKey {
compute_type: cublaslt_sys::cublasComputeType_t,
scale_type: cublaslt_sys::cudaDataType_t,
a_type: cublaslt_sys::cudaDataType_t,
b_type: cublaslt_sys::cudaDataType_t,
c_type: cublaslt_sys::cudaDataType_t,
d_type: cublaslt_sys::cudaDataType_t,
}
/// Caller-supplied shape tuple that uniquely identifies a GEMM descriptor.
///
/// Combine with `TypesKey` for the algorithm cache. Callers should populate
/// every field that influences descriptor validity. Leading dimensions are
/// included because they affect alignment checks in `AlgoCheck`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ShapeKey {
/// `CUBLAS_OP_N` = 0, `CUBLAS_OP_T` = 1
pub transa: i32,
pub transb: i32,
pub m: i32,
pub n: i32,
pub k: i32,
pub lda: i32,
pub ldb: i32,
pub ldc: i32,
/// `CUBLASLT_EPILOGUE_DEFAULT` = 1 if unset. Callers must pass the
/// actual epilogue value they set on the descriptor.
pub epilogue: i32,
/// `CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES` on the original preference.
/// Determines which algos pass the `AlgoCheck` workspace filter.
pub workspace_bytes: usize,
}
impl ShapeKey {
/// Convenience constructor without epilogue (defaults to 1 =
/// `CUBLASLT_EPILOGUE_DEFAULT`).
#[inline]
pub fn new(
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
lda: i32,
ldb: i32,
ldc: i32,
workspace_bytes: usize,
) -> Self {
Self {
transa,
transb,
m,
n,
k,
lda,
ldb,
ldc,
epilogue: cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DEFAULT as i32,
workspace_bytes,
}
}
/// Constructor with explicit epilogue (e.g. `RELU_BIAS`).
#[inline]
pub fn with_epilogue(
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
lda: i32,
ldb: i32,
ldc: i32,
epilogue: cublaslt_sys::cublasLtEpilogue_t,
workspace_bytes: usize,
) -> Self {
Self {
transa,
transb,
m,
n,
k,
lda,
ldb,
ldc,
epilogue: epilogue as i32,
workspace_bytes,
}
}
}
/// Cached algorithm selection result.
///
/// Mirrors the fields of `cublasLtMatmulHeuristicResult_t` that downstream
/// code reads — `algo` and `workspace_size`. `state` is always
/// `CUBLAS_STATUS_SUCCESS` (we never cache a failed selection).
#[derive(Clone, Copy, Debug)]
pub struct DeterministicAlgo {
pub algo: cublaslt_sys::cublasLtMatmulAlgo_t,
pub workspace_size: usize,
}
// SAFETY: `cublasLtMatmulAlgo_t` is an opaque POD buffer (`data: [u64; 8]`
// as of cuBLAS 12.x). It contains no pointers to thread-local state and can
// be safely shared across threads once initialised.
unsafe impl Send for DeterministicAlgo {}
unsafe impl Sync for DeterministicAlgo {}
/// Process-wide deterministic algorithm selector with per-types and per-shape
/// caches.
#[derive(Debug)]
pub struct DeterministicAlgoSelector {
/// `TypesKey` → sorted ascending `Vec<i32>` of algorithm IDs.
ids_cache: Mutex<HashMap<TypesKey, Vec<i32>>>,
/// `(TypesKey, ShapeKey)` → validated algorithm selection.
algo_cache: Mutex<HashMap<(TypesKey, ShapeKey), DeterministicAlgo>>,
}
impl DeterministicAlgoSelector {
fn new() -> Self {
Self {
ids_cache: Mutex::new(HashMap::new()),
algo_cache: Mutex::new(HashMap::new()),
}
}
/// Select a deterministic algorithm for the given descriptor + layouts.
///
/// On first call for a `(types, shape)` combination this performs
/// `AlgoGetIds` (cached thereafter per types tuple) followed by an
/// `AlgoInit`/`AlgoCheck` loop over the sorted ID list, picking the first
/// validating ID that fits in `shape.workspace_bytes`. Subsequent calls
/// with the same `(types, shape)` return the cached result.
///
/// # Safety
///
/// `handle`, `op_desc`, and all four layout handles must be live cuBLAS
/// objects that outlive the call. The returned algo is only valid while
/// `handle` is live.
#[allow(clippy::too_many_arguments)]
pub unsafe fn select(
&self,
handle: cublaslt_sys::cublasLtHandle_t,
compute_type: cublaslt_sys::cublasComputeType_t,
scale_type: cublaslt_sys::cudaDataType_t,
a_type: cublaslt_sys::cudaDataType_t,
b_type: cublaslt_sys::cudaDataType_t,
c_type: cublaslt_sys::cudaDataType_t,
d_type: cublaslt_sys::cudaDataType_t,
op_desc: cublaslt_sys::cublasLtMatmulDesc_t,
a_layout: cublaslt_sys::cublasLtMatrixLayout_t,
b_layout: cublaslt_sys::cublasLtMatrixLayout_t,
c_layout: cublaslt_sys::cublasLtMatrixLayout_t,
d_layout: cublaslt_sys::cublasLtMatrixLayout_t,
shape: ShapeKey,
) -> Result<DeterministicAlgo, MLError> {
let types = TypesKey {
compute_type,
scale_type,
a_type,
b_type,
c_type,
d_type,
};
// ── Fast-path: full (types, shape) algo cache hit ──────────────────
{
let algo_cache = self
.algo_cache
.lock()
.expect("DeterministicAlgoSelector.algo_cache poisoned");
if let Some(hit) = algo_cache.get(&(types, shape)) {
return Ok(*hit);
}
}
// ── Obtain (or cache) algorithm ID list for this types tuple ───────
let ids = self.algo_ids_for(handle, types)?;
if ids.is_empty() {
return Err(MLError::ModelError(format!(
"cublasLtMatmulAlgoGetIds returned no IDs for compute={compute:?}, scale={scale:?}, A={a:?}, B={b:?}, C={c:?}, D={d:?}",
compute = compute_type,
scale = scale_type,
a = a_type,
b = b_type,
c = c_type,
d = d_type,
)));
}
// ── AlgoInit + AlgoCheck loop in deterministic ascending-ID order ──
for &id in &ids {
let mut algo: cublaslt_sys::cublasLtMatmulAlgo_t = std::mem::zeroed();
let init_status = cublaslt_sys::cublasLtMatmulAlgoInit(
handle,
compute_type,
scale_type,
a_type,
b_type,
c_type,
d_type,
id,
&mut algo,
);
if init_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
// AlgoInit can reject an ID for this types tuple (rare but
// possible when cuBLAS trims the table between AlgoGetIds and
// AlgoInit). Skip to the next ID.
continue;
}
let mut heuristic: cublaslt_sys::cublasLtMatmulHeuristicResult_t =
std::mem::zeroed();
let check_status = cublaslt_sys::cublasLtMatmulAlgoCheck(
handle,
op_desc,
a_layout,
b_layout,
c_layout,
d_layout,
&algo,
&mut heuristic,
);
if check_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
continue;
}
if heuristic.state != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
continue;
}
if heuristic.workspaceSize > shape.workspace_bytes {
continue;
}
// Valid algorithm found — cache and return.
let selected = DeterministicAlgo {
algo: heuristic.algo,
workspace_size: heuristic.workspaceSize,
};
tracing::debug!(
algo_id = id,
ids_tried = ids.len(),
workspace_needed = heuristic.workspaceSize,
transa = shape.transa,
transb = shape.transb,
m = shape.m, n = shape.n, k = shape.k,
epilogue = shape.epilogue,
"deterministic algo selected"
);
let mut algo_cache = self
.algo_cache
.lock()
.expect("DeterministicAlgoSelector.algo_cache poisoned");
algo_cache.insert((types, shape), selected);
return Ok(selected);
}
Err(MLError::ModelError(format!(
"no valid cublasLt algorithm found for types=(compute={compute:?}, scale={scale:?}, A={a:?}, B={b:?}, C={c:?}, D={d:?}) shape={shape:?} ids_tried={n}",
compute = compute_type,
scale = scale_type,
a = a_type,
b = b_type,
c = c_type,
d = d_type,
shape = shape,
n = ids.len(),
)))
}
/// Fetch (or populate) the sorted algorithm ID list for a types tuple.
///
/// The returned `Vec` is cloned from the cache so the lock is released
/// before the AlgoInit/AlgoCheck loop runs.
unsafe fn algo_ids_for(
&self,
handle: cublaslt_sys::cublasLtHandle_t,
types: TypesKey,
) -> Result<Vec<i32>, MLError> {
{
let ids_cache = self
.ids_cache
.lock()
.expect("DeterministicAlgoSelector.ids_cache poisoned");
if let Some(hit) = ids_cache.get(&types) {
return Ok(hit.clone());
}
}
let mut ids = vec![0_i32; MAX_ALGO_IDS as usize];
let mut n_returned: i32 = 0;
let status = cublaslt_sys::cublasLtMatmulAlgoGetIds(
handle,
types.compute_type,
types.scale_type,
types.a_type,
types.b_type,
types.c_type,
types.d_type,
MAX_ALGO_IDS,
ids.as_mut_ptr(),
&mut n_returned,
);
if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
return Err(MLError::ModelError(format!(
"cublasLtMatmulAlgoGetIds failed: status={status:?}, types={types:?}"
)));
}
let n_returned = n_returned.max(0) as usize;
ids.truncate(n_returned);
// Sort ascending so selection is independent of whatever order cuBLAS
// returned the IDs in.
ids.sort_unstable();
ids.dedup();
let mut ids_cache = self
.ids_cache
.lock()
.expect("DeterministicAlgoSelector.ids_cache poisoned");
ids_cache.insert(types, ids.clone());
Ok(ids)
}
}
/// Process-shared deterministic algorithm selector.
///
/// Lazily initialised on first call. Thread-safe via per-cache mutexes.
fn selector() -> &'static DeterministicAlgoSelector {
static SELECTOR: OnceLock<DeterministicAlgoSelector> = OnceLock::new();
SELECTOR.get_or_init(DeterministicAlgoSelector::new)
}
/// Convenience free function — drop-in replacement for the cudarc
/// `get_matmul_algo_heuristic` call, returning a
/// `cublasLtMatmulHeuristicResult_t` so downstream code keeps reading
/// `.algo` / `.workspaceSize` unchanged.
///
/// The returned struct has `state = CUBLAS_STATUS_SUCCESS` and `wavesCount =
/// 0.0` (the latter is unused by every current call site).
///
/// # Safety
///
/// Same as [`DeterministicAlgoSelector::select`]: all handles must outlive
/// the call.
#[allow(clippy::too_many_arguments)]
pub unsafe fn get_matmul_algo_deterministic(
handle: cublaslt_sys::cublasLtHandle_t,
compute_type: cublaslt_sys::cublasComputeType_t,
scale_type: cublaslt_sys::cudaDataType_t,
a_type: cublaslt_sys::cudaDataType_t,
b_type: cublaslt_sys::cudaDataType_t,
c_type: cublaslt_sys::cudaDataType_t,
d_type: cublaslt_sys::cudaDataType_t,
op_desc: cublaslt_sys::cublasLtMatmulDesc_t,
a_layout: cublaslt_sys::cublasLtMatrixLayout_t,
b_layout: cublaslt_sys::cublasLtMatrixLayout_t,
c_layout: cublaslt_sys::cublasLtMatrixLayout_t,
d_layout: cublaslt_sys::cublasLtMatrixLayout_t,
shape: ShapeKey,
) -> Result<cublaslt_sys::cublasLtMatmulHeuristicResult_t, MLError> {
let picked = selector().select(
handle,
compute_type,
scale_type,
a_type,
b_type,
c_type,
d_type,
op_desc,
a_layout,
b_layout,
c_layout,
d_layout,
shape,
)?;
Ok(cublaslt_sys::cublasLtMatmulHeuristicResult_t {
algo: picked.algo,
workspaceSize: picked.workspace_size,
state: cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS,
wavesCount: 0.0,
reserved: [0; 4],
})
}
/// F32 inputs/outputs with TF32 fast compute — by far the most common tuple
/// used on the training hot path. Wraps the full six-type call.
///
/// # Safety
///
/// Same as [`get_matmul_algo_deterministic`].
#[allow(clippy::too_many_arguments)]
pub unsafe fn get_matmul_algo_f32_tf32(
handle: cublaslt_sys::cublasLtHandle_t,
op_desc: cublaslt_sys::cublasLtMatmulDesc_t,
a_layout: cublaslt_sys::cublasLtMatrixLayout_t,
b_layout: cublaslt_sys::cublasLtMatrixLayout_t,
c_layout: cublaslt_sys::cublasLtMatrixLayout_t,
d_layout: cublaslt_sys::cublasLtMatrixLayout_t,
shape: ShapeKey,
) -> Result<cublaslt_sys::cublasLtMatmulHeuristicResult_t, MLError> {
let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;
get_matmul_algo_deterministic(
handle,
compute_type,
f32_type,
f32_type,
f32_type,
f32_type,
f32_type,
op_desc,
a_layout,
b_layout,
c_layout,
d_layout,
shape,
)
}

30
crates/ml/src/cuda_pipeline/gpu_attention.rs
View File

@@ -1283,25 +1283,19 @@ fn create_attn_gemm_desc(
let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64)
.map_err(|e| MLError::ModelError(format!("attn D layout ({label}): {e:?}")))?;
// Algorithm selection via heuristic (graph-safe)
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("attn MatmulPrefCreate ({label}): {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| {
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
MLError::ModelError(format!("attn set pref ws ({label}): {e:?}"))
})?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
// compute type preserved; selection is still graph-safe because
// AlgoCheck validates workspace size and other prerequisites.
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa_i32, transb_i32, m, n, k, lda, ldb, ldc, ws_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
lt_handle, matmul_desc,
a_layout, b_layout, c_layout, d_layout,
matmul_pref,
shape,
);
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let heuristic = match heuristic {
Ok(h) => h,
@@ -1312,7 +1306,7 @@ fn create_attn_gemm_desc(
let _ = cublaslt_result::destroy_matrix_layout(a_layout);
let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
return Err(MLError::ModelError(format!(
"attn heuristic ({label} transa={transa_i32},transb={transb_i32},m={m},n={n},k={k}): {e:?}"
"attn deterministic algo ({label} transa={transa_i32},transb={transb_i32},m={m},n={n},k={k}): {e}"
)));
}
};
@@ -1323,7 +1317,7 @@ fn create_attn_gemm_desc(
transa = transa_i32, transb = transb_i32,
m, n, k, lda, ldb, ldc,
ws_needed = heuristic.workspaceSize,
"attn GEMM desc created (heuristic algorithm)"
"attn GEMM desc created (deterministic algorithm)"
);
Ok(AttnGemmDesc {

38
crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs
View File

@@ -192,36 +192,25 @@ impl CuriosityGemm {
let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64)
.map_err(|e| MLError::ModelError(format!("D layout {label}: {e:?}")))?;
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("matmul pref {label}: {e:?}")))?;
let ws_size_u64 = self.workspace_size as u64;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size_u64 as *const u64 as *const std::ffi::c_void,
std::mem::size_of::<u64>(),
).map_err(|e| MLError::ModelError(format!("set workspace pref {label}: {e:?}")))?;
let mut heuristic_result = std::mem::MaybeUninit::<cublaslt_sys::cublasLtMatmulHeuristicResult_t>::uninit();
let mut num_results: i32 = 0;
let heuristic_status = cublaslt_sys::cublasLtMatmulAlgoGetHeuristic(
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
// compute type preserved.
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa_i32, transb_i32,
m, n, k, lda, ldb, ldc,
self.workspace_size,
);
let heuristic_result = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
self.lt_handle.0,
matmul_desc,
a_layout,
b_layout,
c_layout,
d_layout,
matmul_pref,
1,
heuristic_result.as_mut_ptr(),
&mut num_results,
);
if heuristic_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS || num_results == 0 {
return Err(MLError::ModelError(format!(
"cublasLtMatmulAlgoGetHeuristic {label}: status={heuristic_status:?}, results={num_results}"
)));
}
let chosen_algo = &(*heuristic_result.as_ptr()).algo as *const cublaslt_sys::cublasLtMatmulAlgo_t;
shape,
).map_err(|e| MLError::ModelError(format!("curiosity deterministic algo {label}: {e}")))?;
let chosen_algo = &heuristic_result.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t;
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
let status = cublaslt_sys::cublasLtMatmul(
@@ -254,7 +243,6 @@ impl CuriosityGemm {
cublaslt_result::destroy_matrix_layout(b_layout).ok();
cublaslt_result::destroy_matrix_layout(c_layout).ok();
cublaslt_result::destroy_matrix_layout(d_layout).ok();
cublaslt_result::destroy_matmul_pref(matmul_pref).ok();
}
Ok(())
}

31
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -164,27 +164,24 @@ fn create_mamba2_gemm_desc(
f32_type, m as u64, n as u64, ldc as i64,
).map_err(|e| MLError::ModelError(format!("mamba2 D layout {label}: {e:?}")))?;
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("mamba2 MatmulPrefCreate {label}: {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| {
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
MLError::ModelError(format!("mamba2 set pref ws {label}: {e:?}"))
})?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
lt_handle, matmul_desc, a_layout, b_layout, c_layout, d_layout, matmul_pref,
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
// compute type preserved.
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa, transb,
m as i32, n as i32, k as i32,
lda as i32, ldb as i32, ldc as i32,
ws_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
lt_handle, matmul_desc, a_layout, b_layout, c_layout, d_layout, shape,
);
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let algo = heuristic.map_err(|e| MLError::ModelError(
format!("mamba2 heuristic {label} (transa={transa},transb={transb},m={m},n={n},k={k}): {e:?}")
format!("mamba2 deterministic algo {label} (transa={transa},transb={transb},m={m},n={n},k={k}): {e}")
))?.algo;
info!(%label, m, n, k, lda, ldb, ldc, transa, transb, "Mamba2 GEMM desc created");
info!(%label, m, n, k, lda, ldb, ldc, transa, transb, "Mamba2 GEMM desc created (deterministic)");
Ok(Mamba2GemmDesc { matmul_desc, a_layout, b_layout, c_layout, d_layout, algo })
}
}

34
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -1477,25 +1477,25 @@ fn create_iql_gemm_desc(
f32_type, m as u64, n as u64, m as i64,
).map_err(|e| MLError::ModelError(format!("{label} D layout: {e:?}")))?;
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("{label} MatmulPrefCreate: {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| {
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
MLError::ModelError(format!("{label} set pref ws: {e:?}"))
})?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
lt_handle, matmul_desc, a_layout, b_layout, c_layout, d_layout, matmul_pref,
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
// compute type preserved.
// Physical leading dims follow the create_matrix_layout calls above:
// a_ld = k when transa=1 else m; b_ld = ldb; c_ld = m.
let a_ld_i32 = if transa == 1 { k as i32 } else { m as i32 };
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa, transb,
m as i32, n as i32, k as i32,
a_ld_i32, ldb as i32, m as i32,
ws_size,
);
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let algo = heuristic.map_err(|e| MLError::ModelError(format!("{label} heuristic: {e:?}")))?.algo;
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
lt_handle, matmul_desc, a_layout, b_layout, c_layout, d_layout, shape,
);
let algo = heuristic.map_err(|e| MLError::ModelError(format!("{label} deterministic algo: {e}")))?.algo;
tracing::info!(%label, m, n, k, ldb, transa, transb, "IQL GEMM desc created");
tracing::info!(%label, m, n, k, ldb, transa, transb, "IQL GEMM desc created (deterministic)");
Ok(IqlGemmDesc { matmul_desc, a_layout, b_layout, c_layout, d_layout, algo })
}
}

30
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -1842,25 +1842,21 @@ fn create_iqn_gemm_desc(
let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64)
.map_err(|e| MLError::ModelError(format!("IQN D layout {label}: {e:?}")))?;
// Algorithm heuristic
let matmul_pref = cublaslt_result::create_matmul_pref()
.map_err(|e| MLError::ModelError(format!("IQN MatmulPrefCreate {label}: {e:?}")))?;
cublaslt_result::set_matmul_pref_attribute(
matmul_pref,
cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_size as *const usize as *const std::ffi::c_void,
std::mem::size_of::<usize>(),
).map_err(|e| {
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
MLError::ModelError(format!("IQN set pref ws {label}: {e:?}"))
})?;
let heuristic = cublaslt_result::get_matmul_algo_heuristic(
// ── Deterministic algorithm selection (Option B fix) ──
// Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
// compute type preserved.
let shape = super::cublas_algo_deterministic::ShapeKey::new(
transa, transb,
m as i32, n as i32, k as i32,
lda as i32, ldb as i32, ldc as i32,
ws_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
lt_handle, matmul_desc,
a_layout, b_layout, c_layout, d_layout,
matmul_pref,
shape,
);
let _ = cublaslt_result::destroy_matmul_pref(matmul_pref);
let heuristic = match heuristic {
Ok(h) => h,
@@ -1871,7 +1867,7 @@ fn create_iqn_gemm_desc(
let _ = cublaslt_result::destroy_matrix_layout(a_layout);
let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
return Err(MLError::ModelError(format!(
"IQN heuristic {label} (transa={transa},transb={transb},m={m},n={n},k={k}): {e:?}"
"IQN deterministic algo {label} (transa={transa},transb={transb},m={m},n={n},k={k}): {e}"
)));
}
};

1
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -27,6 +27,7 @@ pub mod gpu_curiosity_trainer;
pub mod gpu_walk_forward;
pub mod gpu_dqn_trainer;
pub mod shared_cublas_handle;
pub mod cublas_algo_deterministic;
pub mod batched_forward;
pub mod batched_backward;
pub mod gpu_her;