deepgemm_m_grouped_fp8_fp4_contiguous_g4_n2048_k2048

import torch
import deep_gemm


def run(a_fp8, a_scale, b_fp4, b_scale, m_indices):
    G = b_fp4.shape[0]
    n = b_fp4.shape[1]
    m = a_fp8.shape[0]
    deep_gemm.set_mk_alignment_for_contiguous_layout(
        deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(256)
    )
    d = torch.empty(m, n, device=a_fp8.device, dtype=torch.bfloat16)
    deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(
        (a_fp8, a_scale), (b_fp4, b_scale), d, m_indices,
        disable_ue8m0_cast=False, recipe_a=(1, 128), recipe_b=(1, 32),
    )
    return d