flashinfer.group_gemm#

This module provides a set of functions to group GEMM operations.

class flashinfer.group_gemm.SegmentGEMMWrapper(workspace_buffer: torch.Tensor)#

Wrapper for segment GEMM kernels.

__init__(workspace_buffer: torch.Tensor)#

Initialize the wrapper.

Parameters:

workspace_buffer (torch.Tensor) – The workspace buffer for the kernels, we use it to store the metadata for the segment GEMM whose size is proportional to the number of segments (batch size), 1MB workspace is enough for most cases.

forward(x: torch.Tensor, weights: torch.Tensor, batch_size: int, weight_column_major: bool, seg_lens: torch.Tensor | None = None, seg_indptr: torch.Tensor | None = None, weight_indices: torch.Tensor | None = None) → torch.Tensor#

Forward pass of segment GEMM.

Compute the matrix multiplication between a batch of input tensor (with variable number of rows, but fixed number of columns) and a batch of weight tensor with fixed number of rows and columns:

\[y[i] = x[i] \times W[i]\]

if weight_indices is provided, we will select the weight tensor based on the indices in the weight_indices tensor:

\[y[i] = x[i] \times W[\text{weight_indices}[i]]\]

We use Ragged Tensor to represent the input tensor x and the output tensor y, and each x[i] is a segment of the concatenated tensor. Please see Ragged Tensor tutorial for more details. We use a seg_len or seg_indptr tensor (either would work) to indicate the start and end of each segment, where the seg_indptr is the cumulative sum of the seg_lens tensor (with an additional 0 at the beginning):

\[\text{seg_indptr}[i] = \sum_{j=0}^{i-1} \text{seg_lens}[j], \quad \text{seg_indptr}[0] = 0\]
  • If seg_lens is provided, then x has shape (sum(seg_lens), d_in) and y has shape

    (sum(seg_lens), d_out), where d_in is the number of columns of the input tensor and d_out is the number of columns of the output tensor.

  • If seg_indptr is provided, then x has shape (seg_indptr[-1], d_in) and y has shape

    (seg_indptr[-1], d_out).

Parameters:

  • x (torch.Tensor) – The input tensor with shape (sum(seg_lens), d_in).

  • weights (torch.Tensor) – The 3D weight tensor with shape (num_weights, d_in, d_out) if weight_column_major is False, or (num_weights, d_out, d_in) if weight_column_major is True.

  • batch_size (int) – The number of segments.

  • weight_column_major (bool) – Whether the weight tensor is column major.

  • seg_lens (Optional[torch.Tensor]) – The length of each segment, with shape (batch_size,), expects a 1D tensor of dtype torch.int64.

  • seg_indptr (Optional[torch.Tensor]) – The indptr of the segments, with shape (batch_size + 1,), expects a 1D tensor of dtype torch.int64. If this is provided, then seg_lens will be ignored, otherwise seg_indptr will be computed internally from seg_lens.

  • weight_indices (Optional[torch.Tensor]) – The indices of the weight tensor to be selected for each segment, with shape (batch_size,). Expects a 1D tensor of dtype torch.int64. If this is provided, then the weight tensor will be selected based on the indices in this tensor.

Returns:

The output tensor with shape (sum(seg_lens), d_out).

Return type:

torch.Tensor

reset_workspace_buffer(new_workspace_buffer: torch.Tensor)#

Reset the workspace buffer.

Parameters:

new_workspace_buffer (torch.Tensor) – The new workspace buffer, the device of the new workspace buffer should be the same as the device of the input tensors.