module

GS::Metal::PersistentTransformer

Persistent Transformer module for efficient multi-block processing

Constants

PERSISTENT_TRANSFORMER_SOURCE = {{ read_file("/tmp/tmp.faCcNO/src/src/metal/kernels/persistent_transformer.metal") }}

Load kernel source at compile time

Instance methods

available?

Check availability

Source

ensure_initialized

Lazy initialization

Source

estimate_speedup(batch : Int32, seq_len : Int32, embed_dim : Int32, num_blocks : Int32) : Tuple(Float64, Float64)

Benchmark helper: estimate throughput improvement

Source

forward_multi_block!(input : MetalBuffer, output : MetalBuffer, all_weights : MetalBuffer, temp_buffer : MetalBuffer, batch : Int32, seq_len : Int32, embed_dim : Int32, num_heads : Int32, hidden_dim : Int32, num_blocks : Int32, eps : Float32 = 1e-5_f32) : Nil

Process multiple transformer blocks in a single kernel This is the main optimization - keeps threadgroups resident input: [batch, seq_len, embed_dim] all_weights: concatenated weights for all blocks output: [batch, seq_len, embed_dim] temp_buffer: [batch, seq_len, embed_dim] for ping-pong

Source

forward_single_block!(input : MetalBuffer, output : MetalBuffer, weights : MetalBuffer, batch : Int32, seq_len : Int32, embed_dim : Int32, num_heads : Int32, hidden_dim : Int32, eps : Float32 = 1e-5_f32) : Nil

Process a single transformer block input: [batch, seq_len, embed_dim] weights: concatenated weights for this block output: [batch, seq_len, embed_dim]

Source

get_pipeline(name : String) : ComputePipeline

Get or create pipeline

Source

weights_per_block(embed_dim : Int32, hidden_dim : Int32) : Int32

Weight buffer layout for a single transformer block Returns the size in floats for one block's weights

Source