Implementation:Predibase Lorax SGMV Expand Slice Kernel

Knowledge Sources	Predibase_Lorax
Domains	GPU_Kernels, LoRA
Last Updated	2026-02-08 00:00 GMT

Overview

Triton kernel implementing the Segmented Gather Matrix-Vector multiply (SGMV) expand operation with slice offset support, enabling LoRA B weight application to a specific slice of the output tensor for variable-length sequences during multi-tenant LoRA adapter inference.

Description

The SGMV expand slice kernel extends the standard SGMV expand operation by adding a slice_offset parameter that controls where in the output tensor the results are written. This is the variable-length sequence counterpart to the BGMV expand slice kernel. For each batch b and token m in that batch's segment, it computes output[start_b + m][slice_offset:slice_offset+N] += input[start_b + m] * lora_B[idx[b]].

As noted in the source code, this kernel is separated from the standard sgmv_expand rather than being parameterized because a future fused operator may combine multiple slice operations into a single kernel launch. The kernel uses GroupGEMM tiling with fixed block sizes (BLOCK_M=32, BLOCK_N=32, BLOCK_K=16) and a 2D grid. It supports mixed-precision type casting and optional additive accumulation. The @libentry() decorator is applied for reduced Triton launch overhead. Based on the Punica paper (Chen et al., 2023).

Usage

This kernel is invoked during LoRA adapter inference in the prefill phase when applying LoRA B matrices to sliced output tensors. It handles fused linear layers (e.g., combined QKV projections or gate/up projections in MLP layers) where each LoRA adapter's B matrix corresponds to a specific slice of the combined output. The kernel writes results at the designated slice_offset within the output tensor.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/utils/ops/sgmv_expand_slice.py
• Lines: 1-205

Signature

@libentry()
@triton.jit
def _sgmv_expand_slice_kernel(
    input_ptr,
    lora_ptr,
    out_ptr,
    N,
    K,
    b_seq_start_loc,
    seq_lens,
    lora_indices,
    xm_stride,
    xk_stride,
    l0_stride,
    lora_k_stride,
    lora_n_stride,
    cm_stride,
    cn_stride,
    slice_offset,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    EVEN_K: tl.constexpr,
    ADD_INPUTS: tl.constexpr,
    CAST_TYPE: tl.constexpr,
):

@torch.inference_mode()
def sgmv_expand_slice(
    inputs: torch.Tensor,
    lora_b_weights: torch.Tensor,
    output_tensor: torch.Tensor,
    b_seq_start_loc: torch.Tensor,
    seq_len_tensor: torch.Tensor,
    lora_indices_tensor: torch.Tensor,
    batches: int,
    max_seq_length: int,
    slice_offset: int,
    slice_size: int,
    add_inputs: bool = False,
):

Import

from lorax_server.utils.ops.sgmv_expand_slice import sgmv_expand_slice

I/O Contract

Inputs

Name	Type	Required	Description
inputs	torch.Tensor	Yes	Input tensor of shape (total_tokens, rank), dtype float16/bfloat16/float32. The low-rank intermediate activations from the LoRA A projection, packed across all sequences.
lora_b_weights	torch.Tensor	Yes	LoRA B weight matrices of shape (lora_num, slice_size, rank) or (lora_num, 1, slice_size, rank), dtype float16/bfloat16. The per-adapter up-projection weights for the specific slice.
output_tensor	torch.Tensor	Yes	Output tensor of shape (total_tokens, full_hidden_size), modified in-place at the designated slice region. Must be contiguous.
b_seq_start_loc	torch.Tensor	Yes	Tensor of shape (batch_size,) containing cumulative sequence start positions. E.g., for sequence lengths [4, 6], this would be [0, 4].
seq_len_tensor	torch.Tensor	Yes	Tensor of shape (batch_size,) containing the sequence length for each batch entry.
lora_indices_tensor	torch.Tensor	Yes	Tensor of shape (batch_size,) mapping each batch entry to its LoRA adapter index. A value of -1 means no LoRA is applied.
batches	int	Yes	Number of batch entries (sequences).
max_seq_length	int	Yes	Maximum sequence length in the batch, used to size the Triton grid.
slice_offset	int	Yes	Starting column offset in the output tensor where results are written.
slice_size	int	Yes	Size of the output slice. Must equal lora_b_weights.size(-2).
add_inputs	bool	No	Whether to add results to existing output values (default False) or overwrite them.

Outputs

Name	Type	Description
output_tensor	torch.Tensor	Modified in-place. For each batch b and token m, the slice output_tensor[start_b + m][slice_offset:slice_offset+slice_size] is updated with the matrix product of the input and lora_b_weights[lora_indices[b]].

Usage Examples

# Called internally by LoRA kernel orchestrator (punica.py)
from lorax_server.utils.ops.sgmv_expand_slice import sgmv_expand_slice

# Apply LoRA B weights to the Q slice of a fused QKV output during prefill
sgmv_expand_slice(
    inputs=lora_a_output,           # (total_tokens, rank)
    lora_b_weights=lora_b_q,        # (num_loras, head_dim, rank)
    output_tensor=qkv_output,       # (total_tokens, 3*head_dim), modified in-place
    b_seq_start_loc=seq_starts,     # (batch_size,)
    seq_len_tensor=seq_lens,        # (batch_size,)
    lora_indices_tensor=indices,    # (batch_size,)
    batches=batch_size,
    max_seq_length=max_seq_len,
    slice_offset=0,                 # Write to Q portion
    slice_size=head_dim,
    add_inputs=True,
)

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime