Repository: unslothai/unsloth
File: unsloth/kernels/utils.py (Python)

Code before:
    if BLOCK_SIZE > MAX_FUSED_SIZE:
        raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds "\
                           f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.")
    num_warps : int = 4
    if   BLOCK_SIZE >= 32768: num_warps = 32
    elif BLOCK_SIZE >=  8192: num_warps = 16
    elif BLOCK_SIZE >=  2048: num_warps = 8
    return BLOCK_SIZE, num_warps
pass

if DEVICE_TYPE == "cuda":
    import bitsandbytes as bnb
    # https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files
    HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3")
    get_ptr = bnb.functional.get_ptr


if DEVICE_TYPE == "cuda":
    if torch.cuda.device_count() > 1:
        torch_gpu_device = torch.cuda.device
    else:
        from contextlib import nullcontext
        def torch_gpu_device(device): return nullcontext()
    pass
    _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream
    c_void_p = ctypes.c_void_p
    def _get_tensor_stream(tensor: torch_Tensor) -> c_void_p:
        return c_void_p(_gpu_getCurrentRawStream(tensor.device.index))
    pass
elif DEVICE_TYPE == "xpu":
    if torch.xpu.device_count() > 1:
        torch_gpu_device = torch.xpu.device
    else:
        from contextlib import nullcontext
        def torch_gpu_device(device): return nullcontext()
    pass
    _gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream
    c_void_p = ctypes.c_void_p
    def _get_tensor_stream(tensor: torch_Tensor) -> c_void_p:
        return c_void_p(_gpu_getCurrentRawStream(tensor.device.index))


# Get array of CUDA streams and other buffers
global GPU_STREAMS
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS

if DEVICE_TYPE == "cuda":
    _CUDA_STREAMS = {
        (index := torch.cuda.device(i).idx) : ctypes.c_void_p(torch._C._cuda_getCurrentRawStream(index))
        for i in range(torch.cuda.device_count())

Review comment (suggestion):
`get_ptr` needs to be defined. Seems like 'xpu' would fall into the else category for the quantization function below in this file. As I understand it, bitsandbytes supports intel backends now. Is there a plan to integrate it? Also might be best to have "cuda" as the default option in theses cases.

Code after:
def calculate_settings(n : int) -> (int, int,):
    BLOCK_SIZE : int = next_power_of_2(n)
    if BLOCK_SIZE > MAX_FUSED_SIZE:
        raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds "\
                           f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.")
    num_warps : int = 4
    if   BLOCK_SIZE >= 32768: num_warps = 32
    elif BLOCK_SIZE >=  8192: num_warps = 16
    elif BLOCK_SIZE >=  2048: num_warps = 8
    return BLOCK_SIZE, num_warps
pass

HAS_CUDA_STREAM = False
# INTEL GPU specific logic
if DEVICE_TYPE == "xpu":
    # TODO: Changed here after adding XPU BNB support
    HAS_XPU_STREAM = False
    def get_ptr(x: Optional[torch.Tensor]):
        raise RuntimeError("XPU BNB support is not implemented yet. This function should not be called.")
else:
    # NVIDIA-GPU logic here as default
    import bitsandbytes as bnb
    # https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files
    HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3")
    get_ptr = bnb.functional.get_ptr


if DEVICE_TYPE == "cuda" and torch.cuda.device_count() > 1:
    torch_gpu_device = torch.cuda.device
elif DEVICE_TYPE == "xpu" and torch.xpu.device_count() > 1:
    torch_gpu_device = torch.xpu.device
else:
    from contextlib import nullcontext
    def torch_gpu_device(device): return nullcontext()
    pass

# INTEL GPU Specific Logic
if DEVICE_TYPE == "xpu":
    _gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream 
# NVIDIA GPU Default Logic
else:
    _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream

c_void_p = ctypes.c_void_p
def _get_tensor_stream(tensor: torch_Tensor) -> c_void_p:
    return c_void_p(_gpu_getCurrentRawStream(tensor.device.index))
pass


# Get array of CUDA streams and other buffers
global CUDA_STREAMS