# TMA Alignment Requirement for DeepGEMM Tensors: 16-Byte Hardware Constraint Explained

> Understand the TMA alignment requirement for DeepGEMM tensors. Learn about the 16-byte hardware constraint and how it's enforced in the C++ backend accessed via check sf layout.

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

---

**DeepGEMM requires all tensors participating in Tensor Memory Access (TMA) operations to be strictly 16-byte aligned, enforced through the `kNumTMAAlignmentBytes` constant and validated via `check_sf_layout` in the library's C++ backend.**

DeepGEMM is an FP8 GEMM (General Matrix Multiply) library optimized for NVIDIA Hopper GPUs. When working with TMA-enabled kernels in this repository, understanding the **TMA alignment requirement** is critical for avoiding runtime assertion failures and ensuring hardware-efficient memory access patterns.

## What Is the TMA Alignment Requirement in DeepGEMM?

The TMA alignment requirement in DeepGEMM mandates that any tensor accessed via Tensor Memory Access must start at a memory address that is a multiple of 16 bytes. This constraint exists because the TMA hardware engine on NVIDIA Hopper (SM90) and later architectures fetches data in 16-byte chunks.

The library defines this requirement explicitly in [`csrc/utils/math.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/utils/math.hpp):

```cpp
// In csrc/utils/math.hpp, lines 24-27
constexpr int kNumTMAAlignmentBytes = 16;

```

When calculating strides for TMA operations, DeepGEMM converts this byte alignment into element units based on the tensor's data type size.

## Where DeepGEMM Enforces TMA Alignment

### Constant Definition in math.hpp

The foundation of TMA alignment logic resides in [`csrc/utils/math.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/utils/math.hpp). Here, the constant `kNumTMAAlignmentBytes` is set to 16, establishing the hardware constraint that all subsequent calculations reference.

The file also provides the `get_tma_aligned_size()` utility function, which aligns any given size to the nearest multiple of 16 bytes divided by the element size:

```cpp
int aligned_stride = get_tma_aligned_size(mn, sf.element_size());

```

### Stride Validation in layout.hpp

The actual enforcement occurs in [`csrc/utils/layout.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/utils/layout.hpp) within the `check_sf_layout` function. When `tma_stride_check` is enabled (set to `true`), the function asserts that the tensor's last dimension stride matches the TMA-aligned size:

```cpp
// In csrc/utils/layout.hpp, lines 100-107
assert(sf.stride(-1) == get_tma_aligned_size(mn, sf.element_size()));

```

This assertion prevents misaligned tensors from reaching the TMA hardware, which would otherwise cause undefined behavior or performance degradation.

## Practical Implications for Tensor Layouts

DeepGEMM distinguishes between tensor orientations when applying TMA alignment rules.

### MN-Major Output Tensors

For output tensors (C/D matrices) that are MN-major (the standard output layout), the **leading dimension stride must be TMA-aligned**. This means:

- The stride between consecutive elements in the M-N direction must be a multiple of `16 / element_size`
- For FP8 (1 byte), the stride must be a multiple of 16 elements
- For FP16 (2 bytes), the stride must be a multiple of 8 elements
- For FP32 (4 bytes), the stride must be a multiple of 4 elements

### K-Major Input Tensors

Input tensors (A/B matrices) that are K-major do **not** require TMA alignment on their K dimension. The alignment constraint applies specifically to the MN-major output side where the TMA engine writes results.

## Code Examples: Checking TMA Alignment

### Python Validation Using DeepGEMM Utilities

```python
import torch
import deep_gemm.utils.layout as dg_layout
import deep_gemm.utils.math as dg_math

def verify_tma_alignment(tensor, mn_elements):
    """
    Verify that tensor meets DeepGEMM TMA alignment requirements.
    
    Args:
        tensor: torch.Tensor to validate
        mn_elements: Size of the MN dimension (M*N)
    """
    # Calculate required stride in elements

    required_stride = dg_math.get_tma_aligned_size(
        mn_elements, 
        tensor.element_size()
    )
    
    actual_stride = tensor.stride(-1)
    
    print(f"MN size: {mn_elements}")
    print(f"Element size: {tensor.element_size()} bytes")
    print(f"Required stride: {required_stride} elements")
    print(f"Actual stride: {actual_stride} elements")
    print(f"Aligned: {actual_stride == required_stride}")
    
    return actual_stride == required_stride

# Example: FP16 tensor with M=64, N=32 (MN=2048)

t = torch.randn(64, 32, dtype=torch.float16, device='cuda')
verify_tma_alignment(t, mn_elements=64*32)

```

### C++ Low-Level Alignment Check

```cpp
#include <torch/extension.h>
#include "utils/math.hpp"
#include "utils/layout.hpp"

bool validate_tma_tensor(
    const torch::Tensor& sf,
    int mn,
    int k,
    bool enable_stride_check = true
) {
    try {
        // This will throw if TMA alignment is violated
        deep_gemm::check_sf_layout(
            sf, 
            mn, 
            k,
            128,                 // gran_mn
            128,                 // gran_k
            std::nullopt,        // num_groups
            enable_stride_check, // tma_stride_check
            false,               // sm90_sfb_check
            std::nullopt
        );
        return true;
    } catch (const std::exception& e) {
        std::cerr << "TMA alignment check failed: " << e.what() << std::endl;
        return false;
    }
}

// Usage example
torch::Tensor output = torch::empty({64, 128}, torch::dtype(torch::kFloat16).device(torch::kCUDA));
bool is_valid = validate_tma_tensor(output, 64, 128);

```

## Summary

- **DeepGEMM enforces a strict 16-byte TMA alignment requirement** for all tensors accessed via Tensor Memory Access, defined in [`csrc/utils/math.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/utils/math.hpp) as `kNumTMAAlignmentBytes`.
- The `get_tma_aligned_size()` function converts byte alignment to element strides based on data type size.
- **MN-major output tensors** must have leading dimension strides that are multiples of `16 / element_size`.
- **K-major input tensors** are exempt from TMA alignment constraints.
- The `check_sf_layout` function in [`csrc/utils/layout.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/utils/layout.hpp) validates alignment at runtime when `tma_stride_check` is enabled.

## Frequently Asked Questions

### What happens if my tensor is not 16-byte aligned?

DeepGEMM will raise an assertion error in `check_sf_layout` (located in [`csrc/utils/layout.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/utils/layout.hpp)) when `tma_stride_check` is enabled. If alignment checks are disabled, the TMA hardware on NVIDIA Hopper GPUs may produce undefined behavior or silent data corruption when accessing misaligned memory addresses.

### Does TMA alignment affect performance or correctness?

TMA alignment is a **correctness requirement** for the NVIDIA Hopper TMA engine. While the hardware may tolerate some misalignment in specific scenarios, DeepGEMM treats 16-byte alignment as mandatory to ensure deterministic behavior. Proper alignment also enables optimal memory coalescing, maximizing throughput for FP8 GEMM operations.

### Which DeepGEMM kernels require TMA alignment?

All kernels utilizing **Tensor Memory Access** for writing output tensors require TMA alignment. This includes the main FP8 GEMM kernels in `csrc/jit_kernels/impls/` such as [`sm90_fp8_gemm.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/sm90_fp8_gemm.hpp) and [`smxx_fp8_mqa_logits.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/smxx_fp8_mqa_logits.hpp). The alignment applies specifically to the output (C/D) tensor strides when using TMA store operations.

### How do I calculate the correct stride for TMA alignment?

Use the `get_tma_aligned_size()` function from [`csrc/utils/math.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/utils/math.hpp), which calculates:

```

aligned_elements = ceil_div(total_mn_elements, 16 / element_size) * (16 / element_size)

```

For example, with FP16 (2 bytes), divide 16 by 2 to get 8 elements per alignment unit. Round your MN dimension up to the nearest multiple of 8 to determine the required stride.