Struct Gpu 

pub struct Gpu {
    pub ctx: Arc<CudaContext>,
    /* private fields */
}

GPU runtime handle: a context + stream + cuBLASLt.

Fields

ctx: Arc<CudaContext>

Implementations

impl Gpu

pub fn stream(&self) -> Arc<CudaStream>

The stream every engine op launches on: the thread’s override if one is pushed (pp2 stage scopes), else the main compute stream. By-value Arc so callers hold a stable handle across the call regardless of later pushes/pops.

pub fn blas(&self) -> Arc<CudaBlasLT>

The cuBLASLt handle bound to the same stream returned by stream().

pub fn main_stream(&self) -> &Arc<CudaStream>

The main compute stream, override-blind (graph capture pins itself here; the pp2 runtime uses it to fence stage streams against load-time state).

pub fn enter_main(&self) -> Result<GpuMainOverride, Box<dyn Error>>

Enter this GPU’s own context and matched main stream/cuBLASLt binding, even when the calling thread currently carries a pipeline-stage stream override for another engine.

Multi-context TP/EP helpers invoke rank-local engines from inside a PP stage scope. Without this nested binding, stream() would inherit the PP owner’s stream and launch rank-local pointers through the wrong CUDA context.

pub fn phase_pair( &self, p: usize, ) -> Result<(Arc<CudaStream>, Arc<CudaBlasLT>), Box<dyn Error>>

The phase-p stream/cuBLASLt pair, lazily created. Caller must have this gpu’s context current (enter_main does; external callers use enter_main first).

impl Gpu

pub fn new(ordinal: usize) -> Result<Self, Box<dyn Error>>

pub fn linear_f32( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, m_tokens: usize, in_f: usize, out_f: usize, ) -> Result<Vec<f32>, Box<dyn Error>>

GPU linear y = x @ W^T using cuBLASLt (f32), matching cpu_linear exactly.

Layout reasoning (cuBLASLt is column-major): We want y[m,out] row-major = y^T[out,m] column-major. Treat:

• x[m,in] row-major == x^T[in,m] col-major (an in×m col-major matrix)
• w[out,in] row-major == w^T[in,out] col-major (an in×out col-major matrix) Compute C[out,m] col-major = W_colmajor(out×in) * X_colmajor(in×m) => set A = w (interpreted col-major as in×out, so transa to get out×in), B = x (col-major in×m), C = y (col-major out×m == y[m,out] row-major). cfg: m_=out, n_=m_tokens, k=in. A is in×out (lda=in, transa=true -> out×in), B is in×m (ldb=in), C is out×m (ldc=out).