Function bsz_h

src/fused_mm_sampling/tl_fused_mm_topk.py:110–115

(H: int)

def bsz_h(H: int) -> int:  # noqa: N803
    if H <= 16:
        return 16
    elif H <= 32:
        return 32
    return 64


@torch.compiler.disable