LoRA-Aware Inference

Extends BitNet inference to apply LoRA adapters during generation.

LoRA corrections are injected inside the attention layer, between BitLinear Q/V projections and RoPE. This ensures the LoRA contribution flows through the full attention computation (RoPE → KV cache → scores → softmax).

Uses Attention.forwardWithCacheLoRA and TransformerBlock.forwardWithCacheLoRA which inject LoRA at the correct point in the forward pass.

structure Hesper.LoRA.Inference.LoRAInferenceState : Type

Temporary buffers needed for LoRA inference and training backward

• hBuf : WebGPU.Buffer

  Intermediate h = A @ x buffer [rank]

• yBufQ : WebGPU.Buffer

  Temporary y buffer for Q [dim]

• yBufV : WebGPU.Buffer

  Temporary y buffer for V [kvDim]

• dAttnBuf : Option WebGPU.Buffer

  Attention backward buffers (only allocated for training)

• dScoresBuf : Option WebGPU.Buffer
• dQBuf : Option WebGPU.Buffer
• dQPreBuf : Option WebGPU.Buffer
• savedNormed : Array WebGPU.Buffer

  Per-layer saved normedBuf for multi-layer backward. savedNormed[i] = copy of normedBuf after RMSNorm, before attention layer i. This is the input to LoRA Q/V projections and is needed for gradient computation.

• savedAttn : Array WebGPU.Buffer

  Per-layer saved attention weights for softmax backward. savedAttn[i] = copy of attnBuf (softmax output) for layer i. Needed for correct softmax backward: dScores = attn * (dAttn - Σ attn*dAttn)

• savedAttnOut : Array WebGPU.Buffer

  Per-layer saved attention output (before sub-norm) for RMSNorm backward. savedAttnOut[i] = copy of qRotBuf after attention apply (= input to sub-norm). Needed for RMSNorm backward in the attention chain.

• dAttnOutBuf : Option WebGPU.Buffer

  Scratch buffer for dAttnOut (gradient after O backward, before RMSNorm backward)

• savedGate : Array WebGPU.Buffer

  Per-layer saved FFN activations for FFN backward

• savedUp : Array WebGPU.Buffer
• savedHidden : Array WebGPU.Buffer
• savedResidual1 : Array WebGPU.Buffer
• dFFNNormed : Option WebGPU.Buffer

  Scratch buffers for FFN backward

• dFFNHidden : Option WebGPU.Buffer
• dGateBuf : Option WebGPU.Buffer
• dUpBuf : Option WebGPU.Buffer
• dNormed2Buf : Option WebGPU.Buffer

def Hesper.LoRA.Inference.createLoRAInferenceState (device : WebGPU.Device) (adapter : Adapter) (dim kvDim : Nat) : IO LoRAInferenceState

Create LoRA inference state (inference only, no backward buffers)

Equations

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

def Hesper.LoRA.Inference.createLoRATrainingState (device : WebGPU.Device) (adapter : Adapter) (dim kvDim numHeads headDim maxSeqLen numLayers : Nat) : IO LoRAInferenceState

Create LoRA inference state with training backward buffers

Equations

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

def Hesper.LoRA.Inference.forwardSingleTokenWithLoRA (device : WebGPU.Device) (model : Models.BitNet.BitNetModel WebGPU.Buffer WGSL.Execute.PreparedDispatch WGSL.Execute.CompiledKernel) (tokenId pos : Nat) (cacheState : Models.BitNet.KVCacheState WebGPU.Buffer WGSL.Execute.PreparedDispatch) (adapter : Adapter) (loraState : LoRAInferenceState) : IO Unit

Single-token forward pass with LoRA. Uses TransformerBlock.forwardWithCacheLoRA which injects LoRA inside the attention layer (between BitLinear Q/V and RoPE).

Equations

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

def Hesper.LoRA.Inference.forwardAndBackwardBatched (device : WebGPU.Device) (model : Models.BitNet.BitNetModel WebGPU.Buffer WGSL.Execute.PreparedDispatch WGSL.Execute.CompiledKernel) (tokenId pos : Nat) (cacheState : Models.BitNet.KVCacheState WebGPU.Buffer WGSL.Execute.PreparedDispatch) (adapter : Adapter) (loraState : LoRAInferenceState) (isOutputToken : Bool) (targetBuf lossAccumBuf dLogitsBuf dHiddenBuf : WebGPU.Buffer) (grads : AdapterGrad) (trainState : Training.TrainLoop.TrainState) (startLayer : Nat) : IO Unit

Combined forward + backward in a SINGLE GPU batch. All dispatches (forward 30 layers + loss + backward) are recorded into one command buffer and submitted as a single GPU submit. This eliminates ~20 GPU sync points per token compared to separate forward/backward calls.

@param isOutputToken If true, compute loss + backward after forward. If false (prompt tokens), only forward is executed. @param targetBuf Pre-uploaded target token ID [1] u32 @param lossAccumBuf GPU-side loss accumulator (added to, not overwritten) @param dLogitsBuf Scratch buffer for dLogits [vocabSize] @param dHiddenBuf Scratch buffer for dHidden [dim] @param grads Gradient accumulators for LoRA weights @param startLayer First layer to compute LoRA backward for @param trainState Training state with temp buffers

Equations

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

def Hesper.LoRA.Inference.generateWithLoRA (device : WebGPU.Device) (model : Models.BitNet.BitNetModel WebGPU.Buffer WGSL.Execute.PreparedDispatch WGSL.Execute.CompiledKernel) (adapter : Adapter) (loraState : LoRAInferenceState) (promptTokens : Array Nat) (maxTokens : Nat) (strategy : Inference.Sampling.Strategy := Inference.Sampling.Strategy.Greedy) (eosToken : Option Nat := none) (repetitionPenalty : Float := 1.1) : IO (Array Nat)

Generate text with LoRA adapter applied. Same interface as BitNetModel.generate but with LoRA corrections.