def get_tensor_accessor_alignments(func_attrs):
    """Infer the A, B, and epilogue alignments from the respective TAs."""
    input_accessors = func_attrs["input_accessors"]
    a_alignment = tensor_accessor_codegen.find_max_alignment_for_accessor(
        input_accessors[0]
    )
    b_alignment = tensor_accessor_codegen.find_max_alignment_for_accessor(
        input_accessors[1]
    )
    output_accessor = func_attrs["output_accessors"][0]
    epilogue_alignment = tensor_accessor_codegen.find_max_alignment_for_accessor(
        output_accessor
    )

    # if the last dim is dynamic, force align=1
    if not isinstance(output_accessor.original_shapes[-1], IntImm):
        epilogue_alignment = 1

    return a_alignment, b_alignment, epilogue_alignment
Infer the A, B, and epilogue alignments from the respective TAs.
python/aitemplate/backend/cuda/gemm_universal/common.py