# Understanding the head_splits Parameter in fp8_gemm_nt_skip_head_mid for DeepGEMM

> Learn how the head_splits parameter in DeepGEMM's FP8 kernel efficiently handles sparse attention by inserting zero-padding and skipping masked computations.

- Repository: [DeepSeek/DeepGEMM](https://github.com/deepseek-ai/DeepGEMM)
- Tags: deep-dive
- Published: 2026-04-19

---

**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`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/apis/attention.hpp) within the `fp8_gemm_nt_skip_head_mid` function:

```cpp
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`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/apis/attention.hpp):

```cpp
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:

```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`](https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py) provides a Python reference implementation that mirrors the kernel's behavior:

```python
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`](https://github.com/deepseek-ai/DeepGEMM/blob/main/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`](https://github.com/deepseek-ai/DeepGEMM/blob/main/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`](https://github.com/deepseek-ai/DeepGEMM/blob/main/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.