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 betweenxrows (normed =h; attn_out =q_dim).out_row_stride: distance betweenbase_outrows. Contiguous O usesn_out; the STRIDED K/Vqkv_bufusesper_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).