How to Configure Block Size Multiples for Kernel Alignment in DeepGEMM

Use deep_gemm.set_block_size_multiple_of() to force JIT heuristics to only select block sizes that are integer multiples of your specified values, ensuring hardware alignment constraints are met.

DeepGEMM automatically tunes GPU kernel shapes during Just‑In‑Time (JIT) compilation, but hardware units like Tensor Memory Accelerators (TMA) often require specific alignment granularities. By configuring block size multiples, you constrain the heuristic search space to only those block dimensions that satisfy your alignment requirements.

Understanding Block Size Multiples in DeepGEMM

DeepGEMM’s JIT heuristics select candidate kernel shapes based on hardware capabilities and problem dimensions. Two per‑kernel configuration fields—block_m_multiple_of and block_n_multiple_of—act as filters during this selection process.

These integers tell the runtime to only consider block sizes that are exact multiples of the supplied values. The default value for both fields is 1, meaning no additional alignment constraints are applied beyond the hardware’s intrinsic granularity.

Setting Block Size Multiples via Python

The Python API exposes a single function to configure both dimensions. In deep_gemm/__init__.py, the binding is re‑exported:

from ._C import set_block_size_multiple_of

You can call this function with either a single integer (applied to both M and N) or a tuple of two integers (specifying M and N separately):

import deep_gemm

# Force both block M and block N to be multiples of 64

deep_gemm.set_block_size_multiple_of(64)

# Different alignment for M and N: M must be 128-aligned, N must be 32-aligned

deep_gemm.set_block_size_multiple_of((128, 32))

The underlying C++ binding in csrc/apis/runtime.hpp uses std::variant to accept either format and forwards the values to the HeuristicsRuntime instance.

How the Runtime Applies Alignment Constraints

C++ Runtime Storage

The configuration is stored in the HeuristicsRuntime struct defined in csrc/jit_kernels/heuristics/runtime.hpp:

struct HeuristicsRuntime {
    int block_m_multiple_of = 1;
    int block_n_multiple_of = 1;
    
    void set_block_size_multiple_of(int m, int n) {
        block_m_multiple_of = m;
        block_n_multiple_of = n;
    }
    
    int get_block_m_multiple_of() const { return block_m_multiple_of; }
    int get_block_n_multiple_of() const { return block_n_multiple_of; }
};

SM100 Heuristic Integration

During kernel selection for SM100 (Hopper) GPUs, the heuristic computes a step size that combines the hardware’s intrinsic granularity with your configured multiples. In csrc/jit_kernels/heuristics/sm100.hpp, the code uses std::lcm to calculate this step:

For block M (which has a hardware granularity of 16):

int step = std::lcm(16, heuristics_runtime->get_block_m_multiple_of());

For block N (which has a hardware granularity of 32):

int step = std::lcm(32, heuristics_runtime->get_block_n_multiple_of());

The heuristic then generates candidate block sizes using this step value. If the problem dimensions cannot accommodate a block size that satisfies the multiple constraint, that layout candidate is discarded, and the JIT either selects a different configuration or raises an error.

Distinguishing Block Multiples from MK Alignment

Block size multiples are independent of the group‑level M/K alignment used for contiguous layouts. The set_mk_alignment_for_contiguous_layout function (exposed in deep_gemm/utils/layout.py) controls the minimum block M and block K for grouped kernels, typically set to values like 128 to ensure TMA swizzling works correctly.

While MK alignment sets a minimum granularity for the M dimension, block_m_multiple_of adds an additional constraint that further restricts the search space. For example, you might set MK alignment to 128 (ensuring all M blocks are at least 128) while also setting block_m_multiple_of to 64, which effectively allows block sizes of 128, 192, 256, etc., but ensures they are all multiples of 64.

Practical Configuration Examples

TMA Swizzling Alignment

When using Tensor Memory Accelerator (TMA) swizzling, hardware often requires 128‑byte alignment on the M dimension:

import deep_gemm

# Force 128-byte alignment for block M

deep_gemm.set_block_size_multiple_of(128)

Custom Tile Shapes

To restrict the heuristic to only consider 64×64 tiles for a specific workload:

# Both M and N must be multiples of 64

deep_gemm.set_block_size_multiple_of((64, 64))

Combining with MK Alignment

For grouped GEMM layouts that require both contiguous alignment and specific block multiples:

# Set group-level alignment first

deep_gemm.set_mk_alignment_for_contiguous_layout(128)

# Then restrict block M to multiples of 64 (effective minimum becomes 128, step is 64)

deep_gemm.set_block_size_multiple_of(64)

Summary

• Block size multiples filter JIT kernel selection to only consider block dimensions that are exact multiples of user-supplied values.
• Configuration is performed via deep_gemm.set_block_size_multiple_of(), accepting either a single integer or a tuple (m, n).
• Storage occurs in HeuristicsRuntime (csrc/jit_kernels/heuristics/runtime.hpp), with defaults of 1 (no constraint).
• Application happens in SM100ArchSpec::get_layout_candidates (csrc/jit_kernels/heuristics/sm100.hpp), where the step size is computed as lcm(hardware_granularity, multiple).
• Independence from MK alignment means you can combine set_block_size_multiple_of() with set_mk_alignment_for_contiguous_layout() to achieve both minimum group alignment and specific block granularity.

Frequently Asked Questions

What happens if my matrix dimensions are not divisible by the configured multiple?

The JIT heuristic will attempt to find a valid block size that satisfies the multiple constraint. If no such block size exists for the given problem dimensions, that layout candidate is discarded. If all candidates are discarded, the compilation will fail. Always ensure your matrix dimensions are compatible with the alignment constraints you impose, or use the default multiple of 1 to allow any block size.

Can I set different multiples for the K dimension?

No, the current API only exposes block_m_multiple_of and block_n_multiple_of. The K dimension alignment is controlled separately by the MK alignment setting (set_mk_alignment_for_contiguous_layout) and hardware-specific granularity constraints (typically 16 or 32 bytes). There is no block_k_multiple_of parameter in the current HeuristicsRuntime implementation.

How do block size multiples interact with TMA swizzling modes?

TMA swizzling requires specific byte alignments (commonly 128 bytes) on the M dimension to function correctly. By setting block_m_multiple_of to the swizzle granularity (e.g., 128), you ensure that the JIT only generates kernels with block sizes that satisfy TMA alignment requirements. This is independent of the set_mk_alignment_for_contiguous_layout setting, which controls group-level contiguous layout constraints.

Where can I verify that my multiples are being applied?

Since HeuristicsRuntime does not expose a public getter to Python, you can verify the configuration by running a dummy kernel compilation with DG_JIT_DEBUG=1 set in your environment. The debug output will display the selected block sizes for the compiled kernel. If you set block_m_multiple_of to 64, you should observe that the chosen block M values are exact multiples of 64 (e.g., 64, 128, 192) rather than arbitrary values like 80 or 96.

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