Element-wise Operations

Implements element-wise tensor operations for neural network layers.

Operations

1. Add (Residual Connections)

c[i] = a[i] + b[i]

Used in residual connections throughout transformer blocks.

2. Multiply

c[i] = a[i] × b[i]

Used for gating mechanisms in FFN.

3. ReLU² (Squared ReLU) Activation

ReLU²(x) = max(0, x)²

Used in BitNet b1.58 FFN layers (LLM_FFN_RELU_SQR).

4. GELU Activation

GELU(x) ≈ 0.5 × x × (1 + tanh(√(2/π) × (x + 0.044715 × x³)))

Gaussian Error Linear Unit - used in some transformer variants.

5. Scale

c[i] = a[i] × scalar

Used for normalization and attention scaling.

Performance

Element-wise ops are typically memory-bound:

Compute: 1 FLOP per element
Memory: 2 reads + 1 write = 3 × 4 bytes = 12 bytes
Arithmetic intensity: 1 / 12 ≈ 0.083 FLOP/byte

On A100:

• Peak FLOPS: 19.5 TFLOPS
• Memory bandwidth: 2039 GB/s
• Bandwidth limit: 2039 / 12 ≈ 170 GFLOPS
• Utilization: 170 / 19500 ≈ 0.9% (memory-bound!)

Optimization strategies:

1. Kernel fusion: Combine multiple element-wise ops
2. Vectorization: Process 4 elements per thread (vec4)
3. In-place: Reuse input buffer as output when possible

References

• ReLU²: Used in BitNet b1.58 (LLM_FFN_RELU_SQR in llama.cpp)
• GELU: "Gaussian Error Linear Units" (Hendrycks & Gimpel, 2016)

Configuration

structure Hesper.WGSL.Elementwise.Config : Type

Element-wise operation configuration

• numElements : Nat

instance Hesper.WGSL.Elementwise.instReprConfig : Repr Config

Equations

• Hesper.WGSL.Elementwise.instReprConfig = { reprPrec := Hesper.WGSL.Elementwise.instReprConfig.repr }

def Hesper.WGSL.Elementwise.instReprConfig.repr : Config → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

Addition (Residual Connections)

def Hesper.WGSL.Elementwise.addKernel (config : Config) : Monad.ShaderM Unit

Element-wise addition: c = a + b

Used for residual connections in transformer blocks.

@param config Operation configuration

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Elementwise.executeAdd {β : Type} [GPUBackend β] (ctx : β) (aBuf bBuf cBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute element-wise addition

@param device WebGPU device @param aBuf Input buffer A @param bBuf Input buffer B @param cBuf Output buffer C @param config Configuration

Equations

• One or more equations did not get rendered due to their size.

Multiplication (Gating)

def Hesper.WGSL.Elementwise.mulKernel (config : Config) : Monad.ShaderM Unit

Element-wise multiplication: c = a × b

Used for gating in FFN layers.

@param config Operation configuration

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Elementwise.executeMul {β : Type} [GPUBackend β] (ctx : β) (aBuf bBuf cBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute element-wise multiplication

@param device WebGPU device @param aBuf Input buffer A @param bBuf Input buffer B @param cBuf Output buffer C @param config Configuration

Equations

• One or more equations did not get rendered due to their size.

ReLU² Activation

def Hesper.WGSL.Elementwise.reluSqrKernel (config : Config) : Monad.ShaderM Unit

ReLU² (Squared ReLU) activation: y = max(0, x)²

Used in BitNet b1.58 FFN layers (LLM_FFN_RELU_SQR in llama.cpp).

@param config Operation configuration

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Elementwise.executeReluSqr {β : Type} [GPUBackend β] (ctx : β) (inputBuf outputBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute ReLU² activation

@param device WebGPU device @param inputBuf Input buffer @param outputBuf Output buffer @param config Configuration

Equations

• One or more equations did not get rendered due to their size.

GELU Activation

def Hesper.WGSL.Elementwise.geluKernel (config : Config) : Monad.ShaderM Unit

GELU activation (tanh approximation)

GELU(x) ≈ 0.5 × x × (1 + tanh(√(2/π) × (x + 0.044715 × x³)))

Used in some transformer variants (BERT, GPT-2).

@param config Operation configuration

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Elementwise.executeGELU {β : Type} [GPUBackend β] (ctx : β) (inputBuf outputBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute GELU activation

@param device WebGPU device @param inputBuf Input buffer @param outputBuf Output buffer @param config Configuration

Equations

• One or more equations did not get rendered due to their size.

Scalar Operations

def Hesper.WGSL.Elementwise.scaleKernel (config : Config) (scalar : Float) : Monad.ShaderM Unit

Scalar multiplication: c = a × scalar

Used for scaling operations.

@param config Operation configuration @param scalar Scalar value to multiply by

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Elementwise.executeScale {β : Type} [GPUBackend β] (ctx : β) (inputBuf outputBuf : GPUBackend.Buf β) (config : Config) (scalar : Float) : IO Unit

Execute scalar multiplication

@param device WebGPU device @param inputBuf Input buffer @param outputBuf Output buffer @param config Configuration @param scalar Scalar multiplier

Equations

• One or more equations did not get rendered due to their size.

Fused Operations (Optimization)

def Hesper.WGSL.Elementwise.reluSqrMulKernel (config : Config) : Monad.ShaderM Unit

Fused: ReLU² + Multiply (for BitNet FFN gating)

Computes: c = ReLU²(a) × b = max(0, a)² × b

This is the gating mechanism in BitNet b1.58 FFN:

gate = ReLU²(W_gate @ x)
up = W_up @ x
result = gate × up

Fusing ReLU² + multiply saves one global memory roundtrip.

@param config Operation configuration

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Elementwise.executeReluSqrMul {β : Type} [GPUBackend β] (ctx : β) (aBuf bBuf cBuf : GPUBackend.Buf β) (config : Config) (preparedRef : Option (IO.Ref (Option (GPUBackend.CachedDispatch β))) := none) : IO Unit

Execute fused ReLU² + multiply

@param device WebGPU device @param aBuf Input buffer A (will be ReLU²'d) @param bBuf Input buffer B (multiplier) @param cBuf Output buffer C @param config Configuration

Equations

• One or more equations did not get rendered due to their size.

Clamp (Gradient Clipping)

def Hesper.WGSL.Elementwise.clampInPlaceKernel (numElements : Nat) (minVal maxVal : Float) : Monad.ShaderM Unit

In-place clamp kernel: data[i] = clamp(data[i], minVal, maxVal) Uses single read-write buffer to avoid aliasing issues.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.Elementwise.executeClamp {β : Type} [GPUBackend β] (ctx : β) (inputBuf _outputBuf : GPUBackend.Buf β) (numElements : Nat) (minVal maxVal : Float) : IO Unit

Execute in-place clamp: buf[i] = clamp(buf[i], minVal, maxVal)

Equations

• One or more equations did not get rendered due to their size.