deepgemm_m_grouped_fp8_fp4_masked_g4_n4096_k2048

import torch
import deep_gemm


def run(a_fp8, a_scale, b_fp4, b_scale, masked_m, expected_m):
    G, max_m, k = a_fp8.shape
    n = b_fp4.shape[1]
    em = int(expected_m)
    deep_gemm.set_mk_alignment_for_contiguous_layout(
        deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(em)
    )
    d = torch.empty(G, max_m, n, device=a_fp8.device, dtype=torch.bfloat16)
    deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked(
        (a_fp8, a_scale), (b_fp4, b_scale), d, masked_m, em,
        disable_ue8m0_cast=False, recipe_a=(1, 128), recipe_b=(1, 32),
    )
    out = torch.zeros_like(d)
    for g in range(G):
        mm = int(masked_m[g])
        if mm > 0:
            out[g, :mm] = d[g, :mm]
    return out