Crate ferrotensor_backend 

wgpu compute backend for Ferrotensor: GPU tensors and hand-written WGSL kernels.

WgpuContext opens a GPU (Metal, Vulkan, DX12), UnifiedTensor holds data on it and ComputeEngine compiles the kernels and dispatches them: matmul (plain and transposed), element-wise ops, softmax, RMSNorm, RoPE, fused flash attention (forward and backward, with GQA, masks, dropout, KV-cache offsets and sliding windows), losses, optimisers and dequantisation.

use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};

let ctx = pollster::block_on(WgpuContext::new());
let engine = ComputeEngine::new(&ctx);
let a = UnifiedTensor::new(&ctx, &[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let b = UnifiedTensor::new(&ctx, &[1.0, 0.0, 0.0, 1.0], &[2, 2]);
let c = UnifiedTensor::empty(&ctx, &[2, 2]);
engine.matmul(&ctx, &a, &b, &c);
assert_eq!(pollster::block_on(c.to_vec(&ctx)), vec![1.0, 2.0, 3.0, 4.0]);

Kernels are checked against CPU references in tests/gpu_kernels.rs (run with cargo test -p ferrotensor-backend --test gpu_kernels -- --ignored --test-threads=1).

Re-exports

pub use wgpu;

Modules

pool

Size-keyed reuse of GPU buffers.

Structs

AdamParams

Hyper-parameters for ComputeEngine::adamw_update; laid out to match the WGSL uniform.

AttentionOptions

Options for ComputeEngine::flash_attention_v2_with and ComputeEngine::flash_attention_v2_backward_with.

ComputeEngine

Compiled WGSL kernels and the methods that dispatch them.

UnifiedTensor

A dense, row-major tensor stored in a GPU buffer.

WgpuContext

A GPU device, its queue and a buffer pool shared by tensors.

Enums

DType

Data type for tensors - supports mixed precision training