DeepGEMM Compatibility Wrapper (deep-gemm)¶

deep-gemm is a compatibility wrapper package that preserves the deep_gemm import path while running on MUSA through MATE GEMM operators.

Overview¶

This wrapper is designed for projects that already target DeepGEMM-style Python APIs. It helps run existing integrations on MUSA through MATE with minimal code changes.

The official DeepGEMM v2.1.1 API is the naming and call-shape baseline. This is not a claim of complete upstream behavioral parity: MATE implements the available operations with MUSA backends and documents deliberate differences and unsupported options below.

Package and import¶

Package name: deep-gemm
Import path: deep_gemm
Runtime backend: MATE GEMM and logits operators on MUSA

MUSA wrapper releases use the PEP 440 local version suffix +musa, for example 0.2.7+musa. Use python -m pip show deep-gemm to distinguish this wrapper from the native package. The Python API compatibility baseline is official DeepGEMM v2.1.1; use the MATE-specific limits in this page when porting a workload.

Requirements¶

Before using this wrapper, make sure the following are available:

TorchMUSA is installed and the MUSA runtime environment is configured.
The target workload is configured to run on MUSA devices.

Build¶

Build a wheel from the wrappers/DeepGEMM directory:

python -m build --wheel

The generated wheel will be placed under:

dist/

Installation¶

For delivered packages, install from the external MUSA wheel source:

python -m pip install deep-gemm \
  --index-url https://dl.mthreads.com/repo/api/pypi/pypi/simple

This installs the matching mate dependency automatically.

For local wrapper development, install from source:

python -m pip install --no-build-isolation --no-deps -e ../.. -v
python -m pip install --no-build-isolation --no-deps -e .

Install a built local wheel:

python -m pip install --no-deps dist/deep_gemm-*.whl

If you previously installed the legacy mate-deep-gemm package, uninstall it before installing deep-gemm so the environment does not keep stale wrapper metadata.

python -m pip uninstall -y mate-deep-gemm
python -m pip install --no-deps dist/deep_gemm-*.whl

Import¶

Import the package directly:

import deep_gemm

Import individual APIs:

from deep_gemm import (
    bf16_gemm_nn,
    bf16_gemm_nt,
    bf16_gemm_tn,
    bf16_gemm_tt,
    cublaslt_gemm_nn,
    cublaslt_gemm_nt,
    cublaslt_gemm_tn,
    cublaslt_gemm_tt,
    m_grouped_bf16_gemm_nt_contiguous,
    m_grouped_bf16_gemm_nt_masked,
    m_grouped_fp8_fp4_gemm_nt_contiguous,
    m_grouped_fp8_fp4_gemm_nt_masked,
    m_grouped_fp8_gemm_nn_contiguous,
    k_grouped_fp8_gemm_nt_contiguous,
    k_grouped_fp8_gemm_tn_contiguous,
    k_grouped_bf16_gemm_tn_contiguous,
    fp8_fp4_gemm_nt,
    fp8_fp4_gemm_nn,
    fp8_fp4_gemm_tn,
    fp8_fp4_gemm_tt,
    m_grouped_fp8_gemm_nt_contiguous,
    m_grouped_fp8_gemm_nt_masked,
    fp8_gemm_nn,
    fp8_gemm_nt,
    fp8_gemm_tn,
    fp8_gemm_tt,
    fp8_einsum,
    fp8_gemm_nt_skip_head_mid,
    tf32_hc_prenorm_gemm,
    get_paged_mqa_logits_metadata,
    fp8_paged_mqa_logits,
    fp8_mqa_logits,
)

Public APIs¶

Dense BF16 GEMM:

bf16_gemm_nt
bf16_gemm_nn
bf16_gemm_tn
bf16_gemm_tt

cuBLASLt-compatible GEMM:

cublaslt_gemm_nt
cublaslt_gemm_nn
cublaslt_gemm_tn
cublaslt_gemm_tt

Grouped GEMM:

m_grouped_bf16_gemm_nt_contiguous
m_grouped_bf16_gemm_nt_masked
m_grouped_fp8_fp4_gemm_nt_contiguous
m_grouped_fp8_fp4_gemm_nt_masked
m_grouped_fp8_gemm_nn_contiguous
m_grouped_fp8_gemm_nt_contiguous
m_grouped_fp8_gemm_nt_masked
k_grouped_fp8_gemm_nt_contiguous
k_grouped_fp8_gemm_tn_contiguous
k_grouped_bf16_gemm_tn_contiguous
Legacy aliases: fp8_m_grouped_gemm_nt_masked, bf16_m_grouped_gemm_nt_masked

Dense FP8 GEMM:

fp8_fp4_gemm_nt
fp8_fp4_gemm_nn
fp8_fp4_gemm_tn
fp8_fp4_gemm_tt
fp8_gemm_nt
fp8_gemm_nn
fp8_gemm_tn
fp8_gemm_tt
fp8_einsum
fp8_gemm_nt_skip_head_mid

Einsum:

einsum

HyperConnection prenorm GEMM:

tf32_hc_prenorm_gemm

MQA logits APIs:

get_paged_mqa_logits_metadata
fp8_paged_mqa_logits
fp8_mqa_logits

Utility helpers re-exported from deep_gemm.utils:

bench, bench_kineto, calc_diff
get_num_sms, set_num_sms
get_tc_util, set_tc_util
get_mk_alignment_for_contiguous_layout
get_col_major_tma_aligned_tensor
get_mn_major_tma_aligned_tensor
get_tma_aligned_size
get_mn_major_tma_aligned_packed_ue8m0_tensor
get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
transform_sf_into_required_layout

Testing helpers are implemented in mate.testing.deep_gemm and are also available from the upstream-style path:

deep_gemm.testing.bench
deep_gemm.testing.bench_kineto
deep_gemm.testing.calc_diff

0.2.7 Compatibility Additions¶

MATE 0.2.7 adds the following DeepGEMM-style exports:

Interface	Added capability	MATE-specific contract
`cublaslt_gemm_nn`, `cublaslt_gemm_nt`, `cublaslt_gemm_tn`, `cublaslt_gemm_tt`	2D dense GEMM entry points with cuBLASLt-style transpose names	Inputs and optional `c` are 2D tensors; the wrapper writes into `d` and dispatches to MATE `bmm` with the `mudnn` backend. It does not expose cuBLASLt descriptors or CUDA handles.
`m_grouped_fp8_gemm_nn_contiguous`	Contiguous M-grouped FP8 NN GEMM	Uses the MATE grouped FP8 path, defaults to recipe `(1, 128, 128)`, and requires `disable_ue8m0_cast=True`. The B operand follows NN layout, `[num_groups, K, N]`.
`k_grouped_fp8_gemm_nt_contiguous`	Contiguous K-grouped FP8 NT GEMM	Requires non-empty host `ks_cpu`, each K group divisible by `128`, recipe `(1, 1, 128)`, contiguous inputs, and a contiguous `c` tensor. PSUM layout is unsupported.
`get_tma_aligned_size`	TMA alignment helper	Aligns a logical size to the 16-byte TMA boundary for the supplied element size.
`get_mn_major_tma_aligned_packed_ue8m0_tensor`	Packed UE8M0 scale conversion	Repackages and pads scale factors for the M/N-major layout consumed by the MUSA compatibility path.
`get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor`	K-grouped packed scale conversion	Applies the packed scale conversion independently to the K-group segments.
`transform_sf_into_required_layout`	Upstream-style scale-layout helper	Preserves the upstream-style call shape but currently returns a contiguous M/N-major-compatible scale tensor; `mn`, `k`, recipe, and grouping arguments do not select another MUSA transformation.

The wrapper also accepts disable_ue8m0_cast=True on the grouped FP8 compatibility entry points. UE8M0 casting is not implemented on MUSA, so do not pass False to an entry point that validates this option.

Contiguous Grouped GEMM Alignment¶

get_mk_alignment_for_contiguous_layout() returns the M-axis padding alignment used by DeepGEMM-compatible contiguous grouped GEMM wrappers. It defaults to 128 and can be overridden with:

export MATE_DEEPGEMM_MK_ALIGNMENT=256

Only 128 and 256 are supported. Use the returned value when padding each expert segment and building grouped_layout for m_grouped_{fp8,bf16}_gemm_nt_contiguous. Set the environment variable before starting Python; the helper reads and caches the value on first use.

The grouped wrapper uses the DeepGEMM 559d79f parameter names and order, while retaining MATE-only alignment_m, backend, and overlap parameters at the end of the signatures. PSUM and zero-padding layouts remain unsupported. The wrapper keeps runnable MATE defaults for unsupported capabilities.

The grouped m_grouped_fp8_fp4_* entry points dispatch FP8-by-FP8 operands to the existing FP8 path. FP8 E4M3 A plus packed E2M1 FP4 B dispatches to MATE’s native W4A8 backend. Native MATE scale tuples work with the wrapper defaults; DeepGEMM FP32 logical scales are adapted for recipe_a=(1, 128) and recipe_b=(1, 32) when K=128. Larger K is rejected because the current MATE W4A8 kernel has one A scale per row. The dense fp8_fp4_gemm_nt entry point still rejects packed FP4 operands. For K-grouped GEMM, c=None clears D, an independent C is copied to D, and c is d keeps in-place accumulation.

Compatibility Boundaries¶

All wrapper entry points execute through MATE on MUSA; a function name such as cublaslt_gemm_* does not select a CUDA cuBLASLt implementation.
m_grouped_fp8_gemm_nn_contiguous supports only recipe (1, 128, 128).
k_grouped_fp8_gemm_nt_contiguous is a compatibility adapter over per-group MATE GEMM calls. It requires c, uses recipe (1, 1, 128), and rejects unsupported PSUM layouts.
The M-grouped contiguous APIs require each expert segment to be padded to the configured M alignment. The supported values are 128 and 256.
The grouped wrapper accepts DeepGEMM-style scale tuples, but packed FP4 adaptation is limited to FP8 E4M3 A, packed E2M1 B, FP32 logical scales, and K=128 for the DeepGEMM scale format.
get_paged_mqa_logits_metadata requires block_kv == 64, and the indices argument on the paged MQA APIs is unsupported.
fp8_gemm_nt_skip_head_mid writes into the supplied full-width d tensor and leaves the middle split untouched; it does not allocate a compact result.

Quick Start¶

Minimal import example:

import deep_gemm

An example script is provided at:

examples/run_deep_gemm.py

Examples¶

Run the bundled example:

python examples/run_deep_gemm.py

Notes¶

This wrapper preserves the DeepGEMM-style Python surface, but execution is provided by MATE on MUSA
get_paged_mqa_logits_metadata(..., block_kv, ...) currently requires block_kv == 64
fp8_gemm_nt_skip_head_mid writes the GEMM result into a full-width D tensor according to head_splits=(left, mid, right) and leaves the middle portion of each head untouched
The example script currently demonstrates the FP8 grouped GEMM path