Module lora_delta

Module lora_delta 

Source
Expand description

2026-09-25: Runtime LoRA delta y += scale * (x @ A^T) @ B^T on BF16 activations: the single-adapter path and the per-request routed path.

Owner: model-layers ops. Invariants:

  • Both stages contract at the pool’s padded rank max_rank, never at the adapter’s rank; the pool is zeroed before packing, so the pad rows of A and pad columns of B are zero (lora/loading.rs).
  • The launchers here only launch kernels: they allocate nothing and do not synchronize.
  • apply_lora_delta returns without launching when METRALE_LORA_NO_APPLY=1.

Structs§

LoraAttnWeights
2026-09-25: Per-layer attention LoRA weights, installed by copy onto Qwen3AttentionLayer.
LoraFfnWeights
2026-09-25: Per-layer dense-FFN LoRA weights, installed by copy onto DenseFfnLayer.
LoraKernels
2026-09-25: The kernel handles the LoRA launchers use, resolved when an adapter is installed. Module names follow the [modules] table of kernels/gb10/common/KERNEL.toml (dense_gemv_bf16.cu is gemv, dense_gemm_tc.cu is gemm_tc, dense_gemm_bf16.cu is gemm); residual_add, lora_bgmv, moe_lora_grouped_down and moe_lora_gather_bgmv are file stems.
LoraPair
2026-09-25: One adapted module, with the PEFT tensors in their own layouts, which are the [N, K] weight layout the dense kernels expect: a: [max_rank, k_in] row-major BF16 (PEFT lora_A [r, in_features] in the first rank rows, zero rows after); b: [n_out, max_rank] row-major BF16 (PEFT lora_B [out_features, r] in the first rank columns of each row, zero columns after). scale is lora_alpha / r, or lora_alpha / sqrt(r) with use_rslora, from the adapter’s config, where use_rslora is required (config/src/parsers/lora.rs). It is applied in the fold, not folded into B.
LoraRoute
2026-09-25: Per-(layer, module) routing tables that apply_lora_bgmv reads: the [max_loras] device tables of A and B pool addresses (a_table / b_table, 0 = the slot does not adapt this module), the [max_loras] f32 scale_table, and the projection dims. The addresses are fixed when the pool is packed, so the adapter for each row is chosen only by the per-step seq_slot buffer. Installed by copy onto the layer next to the active LoraPair, which the unrouted path uses.

Functions§

apply_lora_bgmv
2026-09-25: Per-request routed LoRA delta over n rows, each naming its adapter slot in seq_slot[n] (i32; a negative slot gets no delta): out[i, :] += scale_s * (x[i, :] @ A_s^T) @ B_s^T with s = seq_slot[i].
apply_lora_delta
2026-09-25: base_out[m, n_out] += scale * (x[m, k_in] @ a^T) @ b^T for one LoraPair.
lora_gemv_max_m
2026-09-25: Row count at or below which apply_lora_delta runs m single-row GEMV deltas instead of one GEMM: METRALE_LORA_GEMV_MAX_M, read once per process, default 48 when unset or unparsable. dense_gemm_tc tiles 16 rows (TC_TM in dense_gemm_tc.cu) and reads all of B whatever m is, so a small m pays full B traffic for a partly empty tile.
lora_no_apply
2026-09-25: METRALE_LORA_NO_APPLY equal to "1", read once per process: keep the adapter resident but skip every apply_lora_delta. A measurement lever that separates the cost of applying an adapter from the cost of having one loaded; the output is the base model’s while it is set.
lora_no_ffn
2026-09-25: METRALE_LORA_NO_FFN equal to "1", read once per process: the dense-FFN (dense_ffn.rs) and GDN out_proj (qwen3_ssm/lora.rs) skip their deltas, and the attention deltas still apply. A measurement lever; the output is wrong while it is set.