Run tensor primitives on a GPU

The problem: matrix work outgrows a scalar loop

The decoder repeatedly multiplies feature vectors by learned matrices. Our CPU implementation computes each output cell by multiplying matching features and adding their products. Once the inputs are fixed, different output cells do not depend on one another. A GPU matrix library can use that available parallel work while the model keeps the same mathematical operation.

Selecting a GPU also introduces a memory boundary. The Rust tensors currently hold f64 values in host memory, used by the CPU. A GPU needs values in its device memory, and the Rust caller needs a completed result it can read. We will select matrix products and sum/mean reductions explicitly, preserving their checked shapes and the existing text model’s gradient paths.

Explain one product before changing its execution

Use two rows of three features and a matrix mapping those features to two output columns:

Each output cell pairs one row of $A$ with one column of $W$. Row zero and column one give the complete contraction

The other three cells follow the same rule:

The row-major output stores the first row before the second, so its literal values are [58,64,139,154]. The value $64$ occupies flat index $1$. Summing $A$ over feature axis $1$ gives $[6,15]$; dividing each row sum by its three features gives means $[2,5]$. These exact small integers also appear in the captured CPU demonstration. They are an arithmetic fixture, not learned weights or evidence of useful language generation.

Keep one matrix rule for both execution paths

For notation, flatten the broadcast leading batch axes into $B$ effective matrices. Let $M$ count output rows, $K$ contracted features and $N$ output columns in each matrix. Effective operands have shapes $[B,M,K]$ and $[B,K,N]$; their product has shape $[B,M,N]$:

All indices are zero-based. Index $b$ selects an effective batch matrix, $i$ an output row, $j$ an output column and $k$ a contracted feature. $q$ is that cell’s flat row-major output index. In the tiny example, $B=1$, $M=N=2$ and $K=3$, so row zero, column one has $q=(0\cdot2+0)\cdot2+1=1$.

The independence is between output cells. Within one cell, all $K$ products must still contribute to its sum. A device library chooses its tiling, thread mapping and reduction strategy; the formula does not promise a particular schedule or simultaneous execution of every cell.

Preserve views, broadcasting and axis meaning

MatmulPlan::new still checks the operands before either execution path. The final two axes are matrix axes; only the leading batch axes broadcast. The existing TensorView offset traversal determines which stored value belongs to each logical coordinate, including transposed and sliced views.

Packing copies those checked logical values into contiguous row-major matrices for the device primitive. A view need not already be contiguous. Each effective batch gets packed matrices, including repeated values from a broadcast source. Packing changes the storage supplied to the primitive; it preserves the mathematical transpose, output shape and validation rules.

The cuBLAS 12.9 manual’s layout description specifies column-major matrices. Our adapter interprets the packed row-major bytes as transposed column-major matrices and reverses the operands:

The native dimensions are $m=N$, $n=M$ and $k=K$. The native leading dimensions lda, ldb and ldc are respectively $N$, $K$ and $N$ elements. The cuBLAS uniform strided-batch API receives successive right, left and output matrices at strides $KN$, $MK$ and $MN$ elements. This is a storage interpretation of the same product, not an extra mathematical transpose of the answer.

Express sum and mean with the same primitive

For an axis reduction, the existing checked traversal supplies one group for each output value. Pack those groups as rows of a matrix $R$, with $K$ selected-axis values per row. Multiplication by a column of ones gives the row sums:

Here $g$ selects an output group, $s_g$ is its sum, and $\mathbf1$ is a column containing $K$ ones. For a nonempty axis, its mean is $s_g/K$. The CUDA mean request supplies scale $1/K$ to the matrix primitive; the CPU reference sums first and then divides. The real-number formulas agree, but their represented f64 results need not have identical bits.

An empty selected sum axis retains its zero identity. An empty mean axis is an error. A zero-length matrix contraction returns zeros, and an empty output returns no values. These cases require no native arithmetic and must not be counted as executed GPU work.

Pass execution separately from model meaning

TensorExecution::cpu() selects the reference path. TensorExecution::cuda(ordinal) initializes the requested CUDA device and can fail if the feature, libraries or device are unavailable. The executor’s matmul, matmul_with_transpose, sum_axis and mean_axis methods pass that explicit choice into the shared tensor operations.

A failed CUDA request returns an error rather than becoming a reported CUDA run performed on the CPU. Ordinary tensor functions and backward() retain their CPU behavior, including after a CUDA call. The execution choice belongs to the calls that receive it; it does not change a thread’s or process’s default device.

Default CPU builds and inference need no NVIDIA driver, toolkit or GPU. The optional Cargo feature cuda dynamically loads the fixed CUDA 12.9 API; it can compile without nvcc or an installed toolkit. Running that path still needs compatible NVIDIA hardware, a driver and the external libraries listed in the reproduction instructions below. Other GPU backends are not implemented here.

/// Owns an explicitly selected matmul and sum/mean executor.
/// Other operations keep their CPU implementation. Tensors remain host-backed.
/// Pass this choice through AutogradContext and explicitly to backward.
#[derive(Clone, Debug, Default)]
pub struct TensorExecution {
    cuda: Option<Rc<CudaExecutor>>,
    matmul_precision: MatmulPrecision,
    attention_tiles: Option<AttentionTiles>,
}

impl TensorExecution {
    pub const fn cpu() -> Self {
        Self::cpu_with_precision(MatmulPrecision::F64)
    }

    /// Selects represented matrix arithmetic without changing host storage.
    pub const fn cpu_with_precision(matmul_precision: MatmulPrecision) -> Self {
        Self {
            cuda: None,
            matmul_precision,
            attention_tiles: None,
        }
    }

    /// Initializes the requested device now; unavailable CUDA is an error.
    pub fn cuda(ordinal: usize) -> Result<Self, ExecutionError> {
        Ok(Self {
            cuda: Some(Rc::new(CudaExecutor::new(ordinal)?)),
            matmul_precision: MatmulPrecision::F64,
            attention_tiles: None,
        })
    }

    /// Refuses unsupported arithmetic before loading a driver or opening a GPU.
    pub fn cuda_with_precision(
        ordinal: usize,
        precision: MatmulPrecision,
    ) -> Result<Self, ExecutionError> {
        if precision != MatmulPrecision::F64 {
            return Err(ExecutionError::new(
                "precision",
                "CUDA mixed-precision matrix arithmetic is not implemented; select the CPU numerical reference explicitly",
            ));
        }
        Self::cuda(ordinal)
    }

    pub const fn matmul_precision(&self) -> MatmulPrecision {
        self.matmul_precision
    }

    /// Reuses the selected device with the precision recorded by a matrix node.
    pub fn with_matmul_precision(
        &self,
        precision: MatmulPrecision,
    ) -> Result<Self, ExecutionError> {
        if self.cuda.is_some() && precision != MatmulPrecision::F64 {
            return Err(ExecutionError::new(
                "precision",
                "CUDA mixed-precision matrix arithmetic is not implemented",
            ));
        }
        Ok(Self {
            cuda: self.cuda.clone(),
            matmul_precision: precision,
            attention_tiles: self.attention_tiles,
        })
    }

    pub fn device(&self) -> ExecutionDevice {
        self.cuda
            .as_ref()
            .map_or(ExecutionDevice::Cpu, |c| ExecutionDevice::Cuda {
                ordinal: c.ordinal(),
            })
    }

    // region:explicit-attention-execution
    /// Select compact, tiled full attention for output-oriented decoder paths.
    /// Inspectable attention APIs retain their dense reference reports.
    pub fn with_tiled_attention(mut self, tiles: AttentionTiles) -> Self {
        self.attention_tiles = Some(tiles);
        self
    }

    pub const fn attention_tiles(&self) -> Option<AttentionTiles> {
        self.attention_tiles
    }
    // endregion:explicit-attention-execution

    pub fn matmul(
        &self,
        left: &TensorView<'_>,
        right: &TensorView<'_>,
    ) -> Result<Tensor, MatmulError> {
        self.matmul_with_transpose(left, right, false, false)
    }

    pub fn matmul_with_transpose(
        &self,
        left: &TensorView<'_>,
        right: &TensorView<'_>,
        transpose_left: bool,
        transpose_right: bool,
    ) -> Result<Tensor, MatmulError> {
        matmul_with_transpose_and_execution(left, right, transpose_left, transpose_right, self)
    }

    pub fn sum_axis(
        &self,
        input: &TensorView<'_>,
        axis: usize,
        keep_dim: bool,
    ) -> Result<Tensor, TensorOpError> {
        sum_axis_with_execution(input, axis, keep_dim, self)
    }

    pub fn mean_axis(
        &self,
        input: &TensorView<'_>,
        axis: usize,
        keep_dim: bool,
    ) -> Result<Tensor, TensorOpError> {
        mean_axis_with_execution(input, axis, keep_dim, self)
    }

    /// Counts completed native GEMM requests and logical copy bytes in this session.
    pub fn stats(&self) -> ExecutionStats {
        self.cuda
            .as_ref()
            .map_or_else(ExecutionStats::default, |c| c.stats())
    }
}

For each nonempty CUDA request, the adapter checks lengths and dimensions, allocates device buffers through cudarc, copies packed inputs, calls the general-purpose f64 matrix primitive, copies the result to host memory and completes the stream before returning. The returned Tensor owns ordinary host values. CUDA allocations, handles and addresses stay behind this boundary.

This chapter therefore has a hybrid execution path: matrix products and sum/mean primitives can use CUDA, including calls inside composed operations. RMSNorm’s existing mean-square reduction uses the selected mean primitive; its elementwise square, stabilization, log/exp and scaling keep their CPU implementations. Max, other elementary primitives and optimizer code also remain on the CPU. Transfers and completion happen at each selected primitive call. The model is not device-resident, and this is not a complete GPU trainer.

Reuse the text model and its reverse mode

The fixture reuses Chapter 10’s empty-merge BPE on ab: document IDs [0,99,100,1], inputs [0,99,100] and next-token targets [99,100,1]. It uses seed 39, vocabulary size $258$, residual width $4$, one attention head, feed-forward width $8$, two blocks and position capacity $8$. The existing RoPE base is $10000$ and RMSNorm epsilon is $10^{-6}$.

The actual path is text encoder, prepared vectors, the source-independent backbone, tied text head, indexed mean negative log-likelihood and backward. Its logits have shape $[1,3,258]$, containing $774$ values. The table used by lookup and the head remains one live parameter node. The model has $1372$ learned scalars in $20$ named tensors, counting that table once. No tokenizer or CUDA identity becomes a requirement of the generic backbone.

The demonstration borrows one model for both device snapshots and clears accumulated parameter gradients before each pass. It does not replace the learned nodes or change their values between comparisons.

AutogradContext::recording().with_execution(execution) borrows the chosen executor while retaining graph recording. That context travels through text encoding, the prepared input, backbone, head and indexed_mean_nll_with_context. The prepared input cannot outlive its borrowed executor. Device execution remains separate from whether a call records a graph. With a CUDA executor, AutogradContext::no_grad().with_execution(execution) selects CUDA for these matrix and sum/mean primitives while creating no reverse graph.

The execution choice must reach forward, loss and backward. For a product $Y=A W$ with upstream derivative $G=\partial\mathcal L/\partial Y$, the unbatched adjoints are

$\mathcal L$ is the scalar loss. Existing reverse mode already expresses these products and adds broadcasted batch contributions where needed. Those shared reverse-mode primitive calls also need the selected execution. A GPU choice is not part of the model’s learned weights or mathematical meaning.

Call backward_with_execution(execution) on the scalar loss to supply the choice to the existing reverse engine. The graph’s parameter nodes retain their existing semantics; they store no CUDA handle or device address.

/// Reuses one borrowed model on either device. Reset gradients for each comparison.
/// Forward and backward receive the numerical executor explicitly.
pub fn snapshot_with_model(
    execution: &TensorExecution,
    model: &DecoderModel,
) -> DemoResult<Snapshot> {
    for parameter in model.parameters() {
        parameter.tensor().zero_grad()?;
    }
    let a = Tensor::from_vec(vec![2, 3], vec![1., 2., 3., 4., 5., 6.])?;
    let b = Tensor::from_vec(vec![3, 2], vec![7., 8., 9., 10., 11., 12.])?;
    let matrix = execution.matmul(&a.view(), &b.view())?;
    let sums = execution.sum_axis(&a.view(), 1, false)?;
    let means = execution.mean_axis(&a.view(), 1, false)?;

    let tokenizer = BpeTokenizer::from_merge_pairs(&[])?;
    let document = tokenizer.encode_utf8_document("ab");
    let context = AutogradContext::recording().with_execution(execution);
    let backbone = model.backbone();
    let encoded =
        TextInputEncoder::tokens(model, &document[..3], &[1, 3]).encode(context, &backbone)?;
    let hidden = backbone.forward(encoded)?;
    let head = model.text_output_head();
    let logits = head.forward(context, hidden.hidden())?;
    let targets = document[1..]
        .iter()
        .map(|&id| id as usize)
        .collect::<Vec<_>>();
    let loss = logits.indexed_mean_nll_with_context(context, 2, &targets)?;
    loss.backward_with_execution(execution)?;
    let gradients = model
        .parameters()
        .iter()
        .map(|parameter| {
            let gradient = parameter
                .tensor()
                .gradient_snapshot()
                .ok_or("missing named parameter gradient")?;
            Ok((parameter.name().to_owned(), gradient))
        })
        .collect::<DemoResult<BTreeMap<_, _>>>()?;
    Ok(Snapshot {
        matrix,
        sums,
        means,
        logits: logits.value().clone(),
        loss: loss.value().as_slice()[0],
        gradients,
        scalar_parameters: model.parameter_count(),
        tied_alias: head
            .table()
            .is_same_node(model.embedding().table().tensor()),
    })
}

Compare complete represented results

The CPU reference loss is $5.668398842431571$ nats, with f64 bits 4016ac70bf03b620. Changing execution can change floating-point rounding without changing the real-number formula. The comparison therefore checks every logit, the scalar loss, every named gradient’s shape and all $1372$ gradient values, rather than comparing rounded display samples.

For finite CPU value $c$ and CUDA value $u$, the criterion fixed before measurement is

The absolute term covers values near zero; the relative term scales with the reference magnitude. Shape mismatches and nonfinite entries fail. This conservative allowance belongs to the bounded f64 fixture. It is neither a universal error bound nor a measure of language-model quality. Neither execution path is an exact-real oracle for arbitrary inputs.

Completed native-work counters report GEMM API requests and logical bytes copied from host to device and back. They stop growing at usize::MAX. They do not report vendor kernel counts, exact bus traffic, elapsed time, occupancy, request overlap or peak memory. The three tiny arithmetic operations alone request $240$ input bytes and $64$ output bytes; the text model adds further requests. Packing, copies and completion can dominate small examples, so these observations establish no speedup.

From recurrent sequence steps to matrix work

Sutskever, Vinyals and Le’s 2014 sequence-to-sequence model uses a multilayer LSTM to encode an input sequence and another LSTM to decode its output. The recurrent state carries sequence information through successive positions.

Recurrence makes each hidden state depend on its predecessor, restricting within-example parallelism. The 2017 Transformer paper proposes an attention model without recurrence. The Transformer paper’s scaled dot-product attention section groups queries, keys and values into matrices.

The existing decoder supplies projections, attention products and adjoints to the same numerical primitives, preserving causal dependencies when execution changes. Teacher forcing supplies the training sequence, so multiple supplied positions can be processed together under their causal mask. Autoregressive generation still needs previously generated tokens before it can supply the next input. A GPU does not remove that dependency. The course’s decoder-only RMSNorm, RoPE and SwiGLU architecture is not a reproduction of the paper’s encoder-decoder model, and we add no recurrent model here.

Retain the CPU contraction as a numerical contrast

The existing Rust loop visits contracted features in ascending order:

    let precision = execution.matmul_precision();
    for (left_cell_offset, right_cell_offset) in left_cell_offsets.zip(right_cell_offsets) {
        let mut sum = 0.0;
        let mut sum_f32 = 0.0_f32;
        let mut left_offset = left_cell_offset;
        let mut right_offset = right_cell_offset;
        for inner_index in 0..plan.inner {
            let left_value = left.value_at_storage_offset(left_offset);
            let right_value = right.value_at_storage_offset(right_offset);
            match precision {
                MatmulPrecision::F64 => sum += left_value * right_value,
                MatmulPrecision::F16F32 | MatmulPrecision::Bf16F32 => {
                    sum_f32 +=
                        precision.operand_f32(left_value) * precision.operand_f32(right_value);
                }
            }

            if inner_index + 1 < plan.inner {
                left_offset = left_offset
                    .checked_add(plan.left_inner_stride)
                    .expect("a checked matmul plan cannot overflow along the left inner axis");
                right_offset = right_offset
                    .checked_add(plan.right_inner_stride)
                    .expect("a checked matmul plan cannot overflow along the right inner axis");
            }
        }
        values.push(match precision {
            MatmulPrecision::F64 => sum,
            MatmulPrecision::F16F32 | MatmulPrecision::Bf16F32 => {
                precision.round(f64::from(sum_f32))
            }
        });
    }

NVIDIA’s floating-point guide, sections 2.2–3.2 explains why regrouping a dot product or using fused multiply-add can change rounding. This supports using a declared numerical comparison; it does not identify the GPU library’s internal reduction tree or make one backend universally more accurate. The cuBLAS reproducibility conditions also do not promise CPU/GPU bitwise equivalence.

Follow a matrix product across the memory boundary

The formula stays fixed while its values move. The selected cell makes the feature reduction visible inside the larger independent output matrix.

Optional practice: reproduce and explain the execution

COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch11-accelerated-tensor-execution

COURSE_CUDA_LIBRARY_DIR=/absolute/path/to/cuda/lib64

COURSE_CORPUS=false ./course run --cuda --cuda-library-dir "$COURSE_CUDA_LIBRARY_DIR" -- cargo run --quiet --offline --locked -p practical-ch11-accelerated-tensor-execution --features cuda -- --device cuda --ordinal 0

COURSE_CORPUS=false ./course run --cuda --cuda-library-dir "$COURSE_CUDA_LIBRARY_DIR" -- cargo test --offline --locked -p practical-ch11-accelerated-tensor-execution --features cuda --test cuda -- --ignored --test-threads=1

use clap::{Parser, ValueEnum};
use llm_from_scratch_practical::tensor::execution::TensorExecution;

#[derive(Clone, Copy, ValueEnum)]
enum Device {
    Cpu,
    Cuda,
}

#[derive(Parser)]
#[command(about = "Compare shared f64 tensor and text-model operations on CPU or CUDA")]
struct Args {
    #[arg(long, value_enum, default_value = "cpu")]
    device: Device,
    #[arg(long, default_value_t = 0)]
    ordinal: usize,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();
    let execution = match args.device {
        Device::Cpu => TensorExecution::cpu(),
        Device::Cuda => TensorExecution::cuda(args.ordinal)?,
    };
    println!(
        "{}",
        serde_json::to_string(&practical_ch11_accelerated_tensor_execution::demonstrate(
            &execution
        )?)?
    );
    Ok(())
}

1. Use the matrix result above to trace output row zero, column one back to its three input products and its flat output index. Explain which work is independent and which work is a reduction.
2. Explain why a transposed or sliced tensor view must be packed according to its existing checked plan, and why a successful CUDA matrix request still returns a host-backed tensor in this chapter.
3. Using the complete-output comparison above, explain why the execution choice must reach backward and why matching the tiny integer product would be insufficient evidence for the text model.

Carry a numerical reference into mixed precision

The text encoder, source-independent backbone, tied head and reverse mode now reuse selected tensor primitives. Device selection stays outside model semantics, and CPU inference remains available without CUDA. Portable weight saving and loading belong to Chapter 15; this chapter does not implement that later GPU-to-CPU persistence demonstration.

Chapter 12 will introduce mixed precision: FP16/BF16 represented values, master weights, protected reductions, loss scaling and overflow or update-skip semantics. Our fixed f64 reference and explicit comparison criterion give those deliberate dtype choices a numerical starting point.