Function num_compute_units

src/fused_mm_sampling/qitra.py:27–29

(device_index: int | None = None)

def num_compute_units(device_index: int | None = None) -> int:
    props = torch.cuda.get_device_properties(device_index)
    return props.multi_processor_count


_TRITON_TABLE_CACHE: dict[tuple[torch.device], tuple[torch.Tensor, torch.Tensor]] = {}