apply_lora_bgmv

Function apply_lora_bgmv 

Source
pub fn apply_lora_bgmv(
    gpu: &dyn GpuBackend,
    kernels: &LoraKernels,
    route: &LoraRoute,
    x: DevicePtr,
    base_out: DevicePtr,
    seq_slot: DevicePtr,
    n: u32,
    x_row_stride: u32,
    out_row_stride: u32,
    lora_xa: DevicePtr,
    stream: u64,
) -> Result<()>
Expand description

M2 per-request routed LoRA delta over a batch of n decode rows, each naming its own adapter slot via seq_slot[n] (i32, <0 = base/no delta). Two launches — shrink then expand+fold — reading the module’s frozen route.a_table/route.b_table/route.scale_table ([max_loras] device arrays, NULL/0 = base-only slot) at the load-time-fixed pool addresses.

out[i, :] += scale_s * (x[i, :] @ A_s^T) @ B_s^T where s = seq_slot[i].

BYTE-IDENTICAL to n sequential apply_lora_delta(m=1) calls for the same (x_i, s_i) (the on-hardware oracle): kernel 1 is dense_gemv_bf16 with a per-row A-base gather (emits BF16 xa = the oracle’s lora_xa boundary); kernel 2 is the same body reading BF16 xa back, then the oracle’s fold (round delta→BF16, then base += scale*bf16(delta)), so per-slot scale is applied in fp32 AFTER the BF16 delta rounding. Contraction runs at route.max_rank (never true rank) — pad rows/cols are zero, bit-identical.

STRIDES (elements, not bytes):

  • x_row_stride : distance between x rows (normed = h; attn_out = q_dim).
  • out_row_stride : distance between base_out rows. Contiguous O uses n_out; the STRIDED K/V qkv_buf uses per_seq_qkv/2 (BF16 elements) so the fold lands inside the interleaved [Q|K|V] layout without corrupting it.

GRAPH-SAFE: only pointer/value-stable args — x/base_out are the fixed forward buffers, the tables are load-time-fixed, xa is a fixed arena scratch (>= n*max_rank BF16), and seq_slot is a fixed-address buffer whose CONTENTS are re-uploaded each decode step (like positions/block_table). No alloc/sync — captures inside the decode graph.

ARG ORDER is in lockstep with lora_bgmv.cu (cuLaunchKernel is type-blind; the byte-identity oracle is the only guard — keep both in sync).