Understanding the head_splits Parameter in fp8_gemm_nt_skip_head_mid for DeepGEMM

The head_splits parameter is a tuple of three integers (left, mid, right) that instructs the DeepGEMM FP8 kernel to reshape the output matrix by inserting zero-padding in the middle region, enabling efficient sparse attention patterns by skipping computation on masked areas.

The fp8_gemm_nt_skip_head_mid function in the DeepGEMM repository provides optimized FP8 General Matrix Multiply (GEMM) kernels specifically designed for transformer attention mechanisms. This CUDA kernel implements a specialized "skip head mid" pattern that avoids unnecessary computation on masked regions of attention matrices. The head_splits parameter serves as the critical configuration mechanism that defines how the output matrix is logically partitioned to enable this sparse computation pattern.

What Is the head_splits Parameter?

Tuple Structure and Definition

head_splits is defined as std::tuple<int, int, int> representing (left, mid, right):

• left: Number of columns in the left block that undergo normal computation
• mid: Number of columns in the middle region that are skipped (filled with zeros)
• right: Number of columns in the right block that undergo normal computation

Role in Sparse Attention Patterns

This parameter enables the kernel to handle attention masks where a central region is masked out, such as in certain causal attention implementations. Instead of computing the full dense matrix and then applying a mask, the kernel logically expands the output dimensions to insert zero-padding, skipping computation on the mid region entirely. This approach eliminates wasted compute cycles on regions known to be zero.

Technical Implementation in DeepGEMM

Source Code Location and API Signature

The parameter is consumed in csrc/apis/attention.hpp within the fp8_gemm_nt_skip_head_mid function:

static void fp8_gemm_nt_skip_head_mid(
    const std::pair<torch::Tensor, torch::Tensor>& a,
    const std::pair<torch::Tensor, torch::Tensor>& b,
    const torch::Tensor& d,
    const std::tuple<int, int, int>& head_splits,
    … ) {
    const auto [left, mid, right] = head_splits;
    // Validation and kernel dispatch logic
}

Dimensional Constraints and Validation

The implementation enforces strict dimensional constraints through assertions in csrc/apis/attention.hpp:

DG_HOST_ASSERT(n % (left + right) == 0
               && n_ == n + n / (left + right) * mid);

These constraints ensure:

1. Divisibility: The original column count n must be divisible by (left + right), ensuring an integer number of head-tail pairs exists in the matrix.
2. Expansion calculation: The expanded width n_ must equal n plus the additional padding columns calculated as n/(left+right) * mid.

The kernel assumes the original output matrix d has shape [M, N]. With head_splits = (L, M_id, R), the kernel internally works on a virtual matrix of shape [M, N + (N/(L+R))*M_id] where a mid-wide padding is inserted after every (L+R) columns.

Practical Code Examples

Python API Usage

Invoke the kernel with specific split dimensions from Python:

import torch
import deep_gemm

# FP8 input tensors (M, K) and (N, K)

a = (torch.randn(128, 32768, dtype=torch.float8_e4m3fn, device='cuda'), None)
b = (torch.randn(8192, 32768, dtype=torch.float8_e4m3fn, device='cuda'), None)

# Output buffer [M, N]

d = torch.randn(128, 8192, dtype=torch.bfloat16, device='cuda')

# Define split pattern: 128 left, 64 mid (skip), 128 right

head_splits = (128, 64, 128)

# Execute kernel

deep_gemm.fp8_gemm_nt_skip_head_mid(
    a, b, d, head_splits,
    disable_ue8m0_cast=True
)

Reference Implementation for Testing

The test suite in tests/test_attention.py provides a Python reference implementation that mirrors the kernel's behavior:

def apply_skip_head_mid(d: torch.Tensor, head_splits: Tuple[int, int, int]) -> torch.Tensor:
    left, mid, right = head_splits
    m, n = d.shape
    assert n % (left + right) == 0, "N must be divisible by (left + right)"
    num_heads = n // (left + right)
    
    # Reshape to separate head groups

    d = d.view(m, num_heads, left + right)
    d_left = d[:, :, :left]
    d_right = d[:, :, -right:]
    
    # Create zero padding for middle region

    d_mid = torch.zeros(m, num_heads, mid, dtype=d.dtype, device=d.device)
    
    # Concatenate left | zeros (mid) | right

    return torch.cat([d_left, d_mid, d_right], dim=2).reshape(m, -1)

This reference constructs the expected output layout by explicitly inserting zero tensors between the left and right blocks, matching the virtual expansion performed by the CUDA kernel.

Performance Implications

Using head_splits eliminates unnecessary computation on masked attention regions. By skipping the mid region entirely rather than computing and then masking, the kernel reduces:

• Memory bandwidth: No writes occur to the middle region, conserving memory bandwidth
• Compute cycles: Skipped MMA (Matrix Multiply Accumulate) operations on zero-masked regions
• Memory footprint: The logical expansion pattern allows efficient handling of sparse attention without materializing full dense tensors in memory

Typical values like (128, 64, 128) reflect common transformer attention head configurations where attention is split into left and right token groups with a central masked region.

Summary

• The head_splits parameter in fp8_gemm_nt_skip_head_mid is a (left, mid, right) tuple that defines column block sizes for sparse attention patterns.
• Located in csrc/apis/attention.hpp, it enforces dimensional constraints via DG_HOST_ASSERT to ensure valid matrix reshaping.
• The kernel logically expands the output matrix to insert zero-padding in the mid region, skipping computation on that area to improve performance.
• Typical values like (128, 64, 128) reflect common transformer attention head configurations with masked middle regions.
• The reference implementation in tests/test_attention.py provides a Python equivalent for verifying kernel correctness.

Frequently Asked Questions

What happens if the head_splits tuple doesn't satisfy the divisibility constraint?

The kernel will trigger an assertion error via DG_HOST_ASSERT(n % (left + right) == 0). The original column dimension n must be evenly divisible by the sum of left and right to ensure an integer number of head-tail pairs exists in the matrix. If this condition fails, the kernel cannot properly reshape the output tensor and will halt execution.

Can I use head_splits with any matrix dimensions?

No, the dimensions must satisfy the specific constraint that n_ == n + n / (left + right) * mid, where n_ is the expanded width of the destination tensor. This ensures the logical expansion matches the actual memory layout expected by the kernel. Additionally, the batch dimensions and inner product dimensions must align with standard FP8 GEMM requirements as defined in the DeepGEMM API.

Why skip the middle region instead of computing everything and masking later?

Skipping the middle region via head_splits saves both compute and memory bandwidth. The kernel avoids performing Matrix Multiply Accumulate (MMA) operations on the masked region and eliminates unnecessary memory writes, providing significant speedups for sparse attention patterns in large transformer models. This approach is particularly effective when the middle region is known to be zero, such as in certain causal attention implementations with local attention windows.

Where can I find the reference implementation for testing head_splits?

The reference implementation is located in tests/test_attention.py in the apply_skip_head_mid function. This Python function mirrors the kernel's behavior by explicitly reshaping the tensor, splitting into left and right blocks, inserting zero padding for the middle region, and concatenating the results. This reference is used to verify that the CUDA kernel produces identical output to the Python implementation, ensuring correctness across different hardware configurations.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}