diff --git a/benchmarking/nested_4bit_dequant.py b/benchmarking/nested_4bit_dequant.py new file mode 100644 index 000000000..1197768fd --- /dev/null +++ b/benchmarking/nested_4bit_dequant.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +"""Benchmark fused nested 4-bit dequantization against the legacy CUDA launch chain.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +import socket +import statistics + +import torch + +import bitsandbytes +from bitsandbytes import functional as F +from bitsandbytes.backends.cuda import ops as cuda_ops +from bitsandbytes.cextension import lib + +DTYPES = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} +DEFAULT_CASES = "square4096:4096:4096,wide11008:11008:4096,tall11008:4096:11008,square8192:8192:8192" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--cases", + default=DEFAULT_CASES, + help="Comma-separated name:rows:cols cases.", + ) + parser.add_argument("--dtypes", default="fp16,bf16,fp32") + parser.add_argument("--formats", default="nf4,fp4") + parser.add_argument("--blocksize", type=int, default=64) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rounds", type=int, default=7) + parser.add_argument("--repetitions", type=int, default=100) + parser.add_argument("--seed", type=int, default=17) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def parse_cases(value: str) -> list[tuple[str, int, int]]: + cases = [] + for item in value.split(","): + name, rows, cols = item.split(":") + cases.append((name, int(rows), int(cols))) + return cases + + +def require_sm103() -> torch.cuda.DeviceProperties: + if not torch.cuda.is_available() or torch.cuda.device_count() != 1: + raise RuntimeError(f"expected exactly one CUDA device, got {torch.cuda.device_count()}") + props = torch.cuda.get_device_properties(0) + capability = torch.cuda.get_device_capability(0) + if capability != (10, 3) or props.multi_processor_count != 148 or "B300" not in props.name.upper(): + raise RuntimeError( + f"expected B300 SM103 with 148 SMs, got {props.name}, cc={capability}, sms={props.multi_processor_count}" + ) + if not cuda_ops._dequantize_4bit_nested_supported(0): + raise RuntimeError("the loaded package did not select nested dequantization for SM103") + return props + + +def percentile(values: list[float], fraction: float) -> float: + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + return ordered[lower] * (upper - position) + ordered[upper] * (position - lower) + + +def timed_batch(function, repetitions: int) -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repetitions): + function() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1000.0 / repetitions + + +def benchmark_pair(baseline, candidate, warmup: int, rounds: int, repetitions: int) -> dict: + for index in range(warmup): + (baseline if index % 2 == 0 else candidate)() + torch.cuda.synchronize() + + samples = {"legacy": [], "nested": []} + for round_index in range(rounds): + order = ("legacy", "nested") if round_index % 2 == 0 else ("nested", "legacy") + for name in order: + function = baseline if name == "legacy" else candidate + samples[name].append(timed_batch(function, repetitions)) + + result = {} + for name, values in samples.items(): + result[name] = { + "samples_us": values, + "median_us": statistics.median(values), + "p10_us": percentile(values, 0.1), + "p90_us": percentile(values, 0.9), + } + result["ratio"] = result["legacy"]["median_us"] / result["nested"]["median_us"] + return result + + +def repetitions_for(numel: int, requested: int) -> int: + if numel >= 60_000_000: + return min(requested, 25) + if numel >= 40_000_000: + return min(requested, 40) + return requested + + +def raw_equal(left: torch.Tensor, right: torch.Tensor) -> bool: + return torch.equal( + left.contiguous().reshape(-1).view(torch.uint8), + right.contiguous().reshape(-1).view(torch.uint8), + ) + + +def emit(handle, record: dict) -> None: + line = json.dumps(record, sort_keys=True) + print(line, flush=True) + handle.write(line + "\n") + handle.flush() + + +def main() -> None: + args = parse_args() + if args.warmup < 1 or args.rounds < 1 or args.repetitions < 1: + raise ValueError("warmup, rounds, and repetitions must be positive") + + props = require_sm103() + cases = parse_cases(args.cases) + dtype_names = args.dtypes.split(",") + formats = args.formats.split(",") + if any(name not in DTYPES for name in dtype_names): + raise ValueError(f"unsupported dtype list: {args.dtypes}") + if any(name not in ("nf4", "fp4") for name in formats): + raise ValueError(f"unsupported format list: {args.formats}") + + package_library = Path(lib._lib._name).resolve() + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as handle: + emit( + handle, + { + "kind": "metadata", + "hostname": socket.gethostname(), + "gpu": props.name, + "capability": list(torch.cuda.get_device_capability(0)), + "sms": props.multi_processor_count, + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "bitsandbytes": bitsandbytes.__version__, + "library": str(package_library), + "cases": cases, + "dtypes": dtype_names, + "formats": formats, + "blocksize": args.blocksize, + "warmup": args.warmup, + "rounds": args.rounds, + "repetitions": args.repetitions, + "seed": args.seed, + }, + ) + + ratios = [] + for case_index, (name, rows, cols) in enumerate(cases): + numel = rows * cols + for dtype_name in dtype_names: + dtype = DTYPES[dtype_name] + for quant_type in formats: + generator = torch.Generator(device="cuda").manual_seed(args.seed + case_index) + source = torch.randn((rows, cols), generator=generator, device="cuda", dtype=dtype) + packed, state = F.quantize_4bit( + source, + blocksize=args.blocksize, + compress_statistics=True, + quant_type=quant_type, + ) + del source + + legacy_scale = torch.empty_like(state.absmax, dtype=torch.float32) + legacy_out = torch.empty((rows, cols), device="cuda", dtype=dtype) + nested_out = torch.empty_like(legacy_out) + + def legacy(): + cuda_ops._dequantize_blockwise_impl( + state.absmax, + state.state2.absmax, + state.state2.code, + state.state2.blocksize, + torch.float32, + legacy_scale, + ) + legacy_scale.add_(state.offset) + cuda_ops._dequantize_4bit_impl( + packed, + legacy_scale, + state.blocksize, + state.quant_type, + state.dtype, + legacy_out, + ) + + def nested(): + cuda_ops._dequantize_4bit_nested_impl( + packed, + state.absmax, + state.state2.absmax, + state.state2.code, + state.offset, + state.blocksize, + state.quant_type, + state.dtype, + nested_out, + ) + + repetitions = repetitions_for(numel, args.repetitions) + result = benchmark_pair(legacy, nested, args.warmup, args.rounds, repetitions) + legacy() + nested() + torch.cuda.synchronize() + if not raw_equal(legacy_out, nested_out): + raise AssertionError(f"raw output mismatch for {name}/{dtype_name}/{quant_type}") + + ratios.append(result["ratio"]) + output_bytes = numel * dtype.itemsize + emit( + handle, + { + "kind": "direct", + "case": name, + "rows": rows, + "cols": cols, + "numel": numel, + "dtype": dtype_name, + "quant_type": quant_type, + "blocksize": args.blocksize, + "repetitions": repetitions, + "output_bytes": output_bytes, + "legacy_effective_gbps": output_bytes / result["legacy"]["median_us"] / 1e3, + "nested_effective_gbps": output_bytes / result["nested"]["median_us"] / 1e3, + "raw_equal": True, + **result, + }, + ) + emit( + handle, + { + "kind": "summary", + "cells": len(ratios), + "geomean_ratio": math.exp(statistics.mean(math.log(value) for value in ratios)), + "minimum_ratio": min(ratios), + }, + ) + + +if __name__ == "__main__": + main() diff --git a/bitsandbytes/_ops.py b/bitsandbytes/_ops.py index 43efd8609..c82abf098 100644 --- a/bitsandbytes/_ops.py +++ b/bitsandbytes/_ops.py @@ -212,6 +212,114 @@ def _( torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") +def _check_dequantize_4bit_nested( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, + out: Optional[torch.Tensor] = None, +) -> None: + torch._check(blocksize in (32, 64, 128, 256, 512, 1024, 2048, 4096), lambda: f"invalid blocksize {blocksize}") + torch._check(nested_blocksize > 0, lambda: f"nested_blocksize must be positive, got {nested_blocksize}") + torch._check(quant_type in ("nf4", "fp4"), lambda: f"quant_type must be 'nf4' or 'fp4', got {quant_type!r}") + torch._check( + dtype in (torch.float16, torch.bfloat16, torch.float32), + lambda: f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}", + ) + torch._check( + absmax_8bit.device == A.device, + lambda: f"Expected absmax_8bit.device == {A.device}, got {absmax_8bit.device}", + ) + torch._check( + nested_absmax.device == A.device, + lambda: f"Expected nested_absmax.device == {A.device}, got {nested_absmax.device}", + ) + torch._check( + nested_code.device == A.device, + lambda: f"Expected nested_code.device == {A.device}, got {nested_code.device}", + ) + torch._check(offset.device == A.device, lambda: f"Expected offset.device == {A.device}, got {offset.device}") + if out is not None: + torch._check(out.shape == tuple(shape), lambda: f"Expected out.shape == {shape}, got {out.shape}") + torch._check(out.device == A.device, lambda: f"Expected out.device == {A.device}, got {out.device}") + torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") + + +torch.library.define( + "bitsandbytes::dequantize_4bit_nested", + "(Tensor A, Tensor absmax_8bit, Tensor nested_absmax, Tensor nested_code, Tensor offset, int blocksize, int nested_blocksize, str quant_type, int[] shape, ScalarType dtype) -> Tensor", +) + + +@register_fake("bitsandbytes::dequantize_4bit_nested") +def _( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, +) -> torch.Tensor: + _check_dequantize_4bit_nested( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + nested_blocksize, + quant_type, + shape, + dtype, + ) + return torch.empty(shape, dtype=dtype, device=A.device) + + +torch.library.define( + "bitsandbytes::dequantize_4bit_nested.out", + "(Tensor A, Tensor absmax_8bit, Tensor nested_absmax, Tensor nested_code, Tensor offset, int blocksize, int nested_blocksize, str quant_type, int[] shape, ScalarType dtype, Tensor! out) -> ()", +) + + +@register_fake("bitsandbytes::dequantize_4bit_nested.out") +def _( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, + out: torch.Tensor, +) -> None: + _check_dequantize_4bit_nested( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + nested_blocksize, + quant_type, + shape, + dtype, + out, + ) + + torch.library.define( "bitsandbytes::quantize_4bit", "(Tensor A, int blocksize, str quant_type, ScalarType quant_storage) -> (Tensor, Tensor)", diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index a0d9ffe83..ba5bc8e5b 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -9,7 +9,7 @@ from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, get_ptr -from ..._ops import register_kernel +from ..._ops import _check_dequantize_4bit_nested, register_kernel from ...cextension import lib @@ -27,6 +27,13 @@ def _setup_ctypes(names, argtypes, restype=None): [ct.c_void_p] * 4 + [ct.c_int32, ct.c_int32, ct.c_void_p], ) +if torch.version.hip is None: + # Nested 4-bit dequantize: (A, scale codes, nested absmax, nested code, offset, out, blocksize, numel, stream) + _setup_ctypes( + [f"cdequantize_blockwise_nested_{d}_{q}" for d in ("fp32", "bf16", "fp16") for q in ("nf4", "fp4")], + [ct.c_void_p] * 6 + [ct.c_int32, ct.c_int32, ct.c_void_p], + ) + # 4-bit GEMM: (A, B, absmax, absmax_8bit, absmax_code, absmax_offset, out, bias, M, N, K, blocksize, quant_type, stream) _setup_ctypes( [f"cgemm_4bit_{d}" for d in ("bf16", "fp16", "fp32")], @@ -491,6 +498,179 @@ def _dequantize_4bit_impl( ) +def _dequantize_4bit_nested_supported(device_index: int) -> bool: + if torch.version.hip is not None: + return False + _, major, minor = _gpu_dispatch_props(device_index) + return (major, minor) == (10, 3) + + +def _dequantize_4bit_nested_impl( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + quant_type: str, + dtype: torch.dtype, + out: torch.Tensor, +) -> None: + _check_dequantize_4bit_nested( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + 256, + quant_type, + out.shape, + dtype, + out, + ) + A = A.contiguous() + offset_f32 = offset.to(dtype=torch.float32) + + if dtype == torch.bfloat16: + dtype_name = "bf16" + elif dtype == torch.float16: + dtype_name = "fp16" + elif dtype == torch.float32: + dtype_name = "fp32" + else: + raise ValueError(f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}") + + fn = getattr(lib, f"cdequantize_blockwise_nested_{dtype_name}_{quant_type}") + with _cuda_device_of(A): + fn( + A.data_ptr(), + absmax_8bit.data_ptr(), + nested_absmax.data_ptr(), + nested_code.data_ptr(), + offset_f32.data_ptr(), + out.data_ptr(), + blocksize, + out.numel(), + _get_raw_stream(A.device.index), + ) + + +def _dequantize_4bit_nested_dispatch( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, + out: torch.Tensor, +) -> None: + _check_dequantize_4bit_nested( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + nested_blocksize, + quant_type, + shape, + dtype, + out, + ) + if nested_blocksize == 256 and _dequantize_4bit_nested_supported(A.device.index): + _dequantize_4bit_nested_impl( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + quant_type, + dtype, + out, + ) + return + + absmax = torch.empty_like(absmax_8bit, dtype=torch.float32) + _dequantize_blockwise_impl( + absmax_8bit, + nested_absmax, + nested_code, + nested_blocksize, + torch.float32, + out=absmax, + ) + _dequantize_4bit_impl(A, absmax + offset, blocksize, quant_type, dtype, out=out) + + +@register_kernel("bitsandbytes::dequantize_4bit_nested", "cuda") +def _( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, +) -> torch.Tensor: + out = torch.empty(shape, dtype=dtype, device=A.device) + _dequantize_4bit_nested_dispatch( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + nested_blocksize, + quant_type, + shape, + dtype, + out, + ) + return out + + +@register_kernel("bitsandbytes::dequantize_4bit_nested.out", "cuda") +def _( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, + out: torch.Tensor, +) -> None: + if out.shape != tuple(shape): + raise ValueError(f"Expected out.shape == {shape}, got {out.shape}") + if out.dtype != dtype: + raise ValueError(f"Expected out.dtype == {dtype}, got {out.dtype}") + _dequantize_4bit_nested_dispatch( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + nested_blocksize, + quant_type, + shape, + dtype, + out, + ) + + @register_kernel("bitsandbytes::gemv_4bit", "cuda") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int @@ -907,12 +1087,25 @@ def _dequant_linear_fallback( """Unfused fallback shared by CUDA and ROCm: reconstruct the (optionally nested) absmax, dequantize the 4-bit weight via the backend dequant impls (reusing preallocated buffers), then F.linear.""" - if absmax_8bit is not None: - absmax_dq = torch.empty_like(absmax_8bit, dtype=torch.float32) - _dequantize_blockwise_impl(absmax_8bit, absmax, absmax_code, 256, torch.float32, out=absmax_dq) - absmax = absmax_dq + absmax_offset B_dq = torch.empty(shapeB, dtype=A.dtype, device=A.device) - _dequantize_4bit_impl(B, absmax, blocksize, quant_type, A.dtype, out=B_dq) + if absmax_8bit is not None and _dequantize_4bit_nested_supported(A.device.index): + _dequantize_4bit_nested_impl( + B, + absmax_8bit, + absmax, + absmax_code, + absmax_offset, + blocksize, + quant_type, + A.dtype, + B_dq, + ) + else: + if absmax_8bit is not None: + absmax_dq = torch.empty_like(absmax_8bit, dtype=torch.float32) + _dequantize_blockwise_impl(absmax_8bit, absmax, absmax_code, 256, torch.float32, out=absmax_dq) + absmax = absmax_dq + absmax_offset + _dequantize_4bit_impl(B, absmax, blocksize, quant_type, A.dtype, out=B_dq) return torch.nn.functional.linear(A, B_dq, bias) diff --git a/bitsandbytes/backends/default/ops.py b/bitsandbytes/backends/default/ops.py index 521802922..e466a76da 100644 --- a/bitsandbytes/backends/default/ops.py +++ b/bitsandbytes/backends/default/ops.py @@ -5,7 +5,7 @@ import torch -from ..._ops import register_kernel +from ..._ops import _check_dequantize_4bit_nested, register_kernel from ..utils import _get_4bit_code @@ -300,6 +300,106 @@ def _( return _dequantize_4bit_compute(A.reshape(-1), absmax, code, blocksize, shape, dtype) +def _dequantize_4bit_nested_default_impl( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + _check_dequantize_4bit_nested( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + nested_blocksize, + quant_type, + shape, + dtype, + out, + ) + absmax = torch.ops.bitsandbytes.dequantize_blockwise.default( + absmax_8bit, + nested_absmax, + nested_code, + nested_blocksize, + torch.float32, + ) + result = torch.ops.bitsandbytes.dequantize_4bit.default( + A, + absmax + offset, + blocksize, + quant_type, + shape, + dtype, + ) + return out.copy_(result) if out is not None else result + + +@register_kernel("bitsandbytes::dequantize_4bit_nested", "default") +def _( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, +) -> torch.Tensor: + return _dequantize_4bit_nested_default_impl( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + nested_blocksize, + quant_type, + shape, + dtype, + ) + + +@register_kernel("bitsandbytes::dequantize_4bit_nested.out", "default") +def _( + A: torch.Tensor, + absmax_8bit: torch.Tensor, + nested_absmax: torch.Tensor, + nested_code: torch.Tensor, + offset: torch.Tensor, + blocksize: int, + nested_blocksize: int, + quant_type: str, + shape: Sequence[int], + dtype: torch.dtype, + out: torch.Tensor, +) -> None: + _dequantize_4bit_nested_default_impl( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset, + blocksize, + nested_blocksize, + quant_type, + shape, + dtype, + out, + ) + + @register_kernel("bitsandbytes::gemv_4bit", "default") def _( A: torch.Tensor, diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py index 33b2cd9ac..611bb7a20 100644 --- a/bitsandbytes/functional.py +++ b/bitsandbytes/functional.py @@ -1049,13 +1049,47 @@ def dequantize_4bit( if quant_state.dtype not in (torch.bfloat16, torch.float16, torch.float32): raise ValueError(f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {quant_state.dtype}") + nested_out = None if quant_state.nested: - absmax = dequantize_blockwise(quant_state.absmax, quant_state.state2) - absmax += quant_state.offset - if absmax.dtype != torch.float32: - absmax = absmax.float() + if A.is_cuda and not torch.compiler.is_compiling(): + if out is not None: + torch.ops.bitsandbytes.dequantize_4bit_nested.out( + A, + quant_state.absmax, + quant_state.state2.absmax, + quant_state.state2.code, + quant_state.offset, + quant_state.blocksize, + quant_state.state2.blocksize, + quant_state.quant_type, + quant_state.shape, + quant_state.dtype, + out=out, + ) + nested_out = out + else: + nested_out = torch.ops.bitsandbytes.dequantize_4bit_nested.default( + A, + quant_state.absmax, + quant_state.state2.absmax, + quant_state.state2.code, + quant_state.offset, + quant_state.blocksize, + quant_state.state2.blocksize, + quant_state.quant_type, + quant_state.shape, + quant_state.dtype, + ) - if out is not None: + if nested_out is None: + absmax = dequantize_blockwise(quant_state.absmax, quant_state.state2) + absmax += quant_state.offset + if absmax.dtype != torch.float32: + absmax = absmax.float() + + if nested_out is not None: + out = nested_out + elif out is not None: torch.ops.bitsandbytes.dequantize_4bit.out( A, absmax, quant_state.blocksize, quant_state.quant_type, quant_state.shape, quant_state.dtype, out=out ) diff --git a/csrc/kernels.cu b/csrc/kernels.cu index df0f093df..729dd5efd 100644 --- a/csrc/kernels.cu +++ b/csrc/kernels.cu @@ -528,6 +528,61 @@ __global__ void } } +#if BUILD_CUDA +template +__global__ void kDequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, T* out, + const int blocksize, const int n +) { + const int n_load = gridDim.x * TILE_SIZE; + const int base_idx = blockIdx.x * TILE_SIZE; + + T vals[NUM_PER_TH * 2]; + unsigned char qvals[NUM_PER_TH]; + + typedef bnb_cub::BlockLoad LoadChar; + typedef bnb_cub::BlockStore StoreT; + + __shared__ typename LoadChar::TempStorage loadchar; + __shared__ typename StoreT::TempStorage storet; + + const int packed_blocksize = blocksize / 2; + const int block_shift = 31 - __clz(packed_blocksize); + const float nested_offset = __ldg(offset); + + for (int i = base_idx; i < n_load; i += gridDim.x * TILE_SIZE) { + const int valid_items_load = min(TILE_SIZE, static_cast((static_cast(n) + 1) / 2) - i); + const int valid_items_store = min(TILE_SIZE * 2, n - i * 2); + const int packed_n = static_cast((static_cast(n) + 1) / 2); + const int packed_index = min(i + threadIdx.x * NUM_PER_TH, packed_n - 1); + const int outer_block = packed_index >> block_shift; + const unsigned char scale_code = __ldg(&absmax_8bit[outer_block]); + const float scale_product = __fmul_rn(__ldg(&nested_code[scale_code]), __ldg(&nested_absmax[outer_block >> 8])); + const float local_abs_max = __fadd_rn(scale_product, nested_offset); + + __syncthreads(); + LoadChar(loadchar).Load(&(A[i]), qvals, valid_items_load, 128); + + if (DATA_TYPE == FP4) { +#pragma unroll NUM_PER_TH + for (int j = 0; j < NUM_PER_TH; j++) { + vals[j * 2] = dDequantizeFP4Tree(qvals[j] >> 4) * local_abs_max; + vals[j * 2 + 1] = dDequantizeFP4Tree(qvals[j] & 0x0F) * local_abs_max; + } + } else { +#pragma unroll NUM_PER_TH + for (int j = 0; j < NUM_PER_TH; j++) { + vals[j * 2] = dDequantizeNF4(qvals[j] >> 4) * local_abs_max; + vals[j * 2 + 1] = dDequantizeNF4(qvals[j] & 0x0F) * local_abs_max; + } + } + + __syncthreads(); + StoreT(storet).Store(&(out[i * 2]), vals, valid_items_store); + } +} +#endif + template __launch_bounds__(BLOCK_SIZE / NUM_VALS, 1) __global__ void kPreconditionOptimizer32bit2State( T* g, T* p, float* state1, float* state2, float* unorm, const float beta1, const float beta2, const float eps, @@ -1827,6 +1882,33 @@ template __global__ void kDequantizeBlockwise( float* code, unsigned char* A, float* absmax, bnb_bfloat16* out, const int blocksize, const int n ); +#if BUILD_CUDA +template __global__ void kDequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, half* out, + const int blocksize, const int n +); +template __global__ void kDequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, half* out, + const int blocksize, const int n +); +template __global__ void kDequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, float* out, + const int blocksize, const int n +); +template __global__ void kDequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, float* out, + const int blocksize, const int n +); +template __global__ void kDequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, + bnb_bfloat16* out, const int blocksize, const int n +); +template __global__ void kDequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, + bnb_bfloat16* out, const int blocksize, const int n +); +#endif + #define MAKE_OptimizerStatic8bit2StateBlockwise(oname, gtype, block_size, num_per_thread) \ template __global__ void kOptimizerStatic8bit2StateBlockwise( \ gtype * p, gtype* __restrict__ const g, unsigned char* state1, unsigned char* state2, const float beta1, \ diff --git a/csrc/kernels.cuh b/csrc/kernels.cuh index dc511661b..ff956d8d5 100644 --- a/csrc/kernels.cuh +++ b/csrc/kernels.cuh @@ -23,6 +23,14 @@ template +__global__ void kDequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, T* out, + const int blocksize, const int n +); +#endif + template __global__ void kPreconditionOptimizer32bit2State( T* g, T* p, float* state1, float* state2, float* unorm, const float beta1, const float beta2, const float eps, diff --git a/csrc/ops.cu b/csrc/ops.cu index 16eed4e81..eae7b4d3b 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -93,6 +93,20 @@ void dequantizeBlockwise( BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } +#if BUILD_CUDA +template +void dequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, T* out, + int blocksize, const int n, bnb_stream_t stream +) { + constexpr int tile_size = 1024; + int grid_blocks = (static_cast(n) + tile_size - 1) / tile_size; + kDequantizeBlockwiseNested + <<>>(A, absmax_8bit, nested_absmax, nested_code, offset, out, blocksize, n); + BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); +} +#endif + template void optimizer32bit( T* g, T* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, const float beta1, @@ -563,6 +577,33 @@ template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, bnb_bfloat16* out, int blocksize, const int n, bnb_stream_t stream ); +#if BUILD_CUDA +template void dequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, half* out, + int blocksize, const int n, bnb_stream_t stream +); +template void dequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, half* out, + int blocksize, const int n, bnb_stream_t stream +); +template void dequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, float* out, + int blocksize, const int n, bnb_stream_t stream +); +template void dequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, float* out, + int blocksize, const int n, bnb_stream_t stream +); +template void dequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, + bnb_bfloat16* out, int blocksize, const int n, bnb_stream_t stream +); +template void dequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, + bnb_bfloat16* out, int blocksize, const int n, bnb_stream_t stream +); +#endif + #define MAKE_optimizer32bit(name, gtype) \ template void optimizer32bit( \ gtype * g, gtype * p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ diff --git a/csrc/ops.cuh b/csrc/ops.cuh index c7114bcaa..9483c984c 100644 --- a/csrc/ops.cuh +++ b/csrc/ops.cuh @@ -101,6 +101,14 @@ void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, T* out, int block_size, const int n, bnb_stream_t stream ); +#if BUILD_CUDA +template +void dequantizeBlockwiseNested( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, T* out, + int block_size, const int n, bnb_stream_t stream +); +#endif + template void optimizer32bit( T* g, T* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, float beta1, float beta2, diff --git a/csrc/pythonInterface.cpp b/csrc/pythonInterface.cpp index 24431603b..8a2a7f5f2 100644 --- a/csrc/pythonInterface.cpp +++ b/csrc/pythonInterface.cpp @@ -220,6 +220,18 @@ void dequantizeBlockwise_bf16_nf4( dequantizeBlockwise(nullptr, A, absmax, out, blocksize, n, stream); } +#if BUILD_CUDA +template +void dequantizeBlockwiseNestedTyped( + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, T* out, + int blocksize, const int n, cudaStream_t stream +) { + dequantizeBlockwiseNested( + A, absmax_8bit, nested_absmax, nested_code, offset, out, blocksize, n, stream + ); +} +#endif + int igemmlt_32( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream @@ -443,6 +455,25 @@ void cdequantize_blockwise_bf16_nf4( dequantizeBlockwise_bf16_nf4(code, A, absmax, out, blocksize, n, stream); } +#if BUILD_CUDA +#define MAKE_NESTED_DEQUANT_FUNC(dtype, dtype_name, data_type, data_type_name) \ + void cdequantize_blockwise_nested_##dtype_name##_##data_type_name( \ + unsigned char* A, unsigned char* absmax_8bit, float* nested_absmax, float* nested_code, float* offset, \ + dtype* out, int blocksize, const int n, cudaStream_t stream \ + ) { \ + dequantizeBlockwiseNestedTyped( \ + A, absmax_8bit, nested_absmax, nested_code, offset, out, blocksize, n, stream \ + ); \ + } + +MAKE_NESTED_DEQUANT_FUNC(half, fp16, FP4, fp4) +MAKE_NESTED_DEQUANT_FUNC(half, fp16, NF4, nf4) +MAKE_NESTED_DEQUANT_FUNC(float, fp32, FP4, fp4) +MAKE_NESTED_DEQUANT_FUNC(float, fp32, NF4, nf4) +MAKE_NESTED_DEQUANT_FUNC(bnb_bfloat16, bf16, FP4, fp4) +MAKE_NESTED_DEQUANT_FUNC(bnb_bfloat16, bf16, NF4, nf4) +#endif + #define MAKE_CFUNC32(name, gtype, gbits) \ void c##name##32bit_grad_##gbits( \ gtype* g, gtype* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ diff --git a/tests/test_functional.py b/tests/test_functional.py index e4cd6a128..9c60885e4 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -573,6 +573,145 @@ def test_coo_int8_vectorwise_quant(self, device, dim1, dim2): class TestQuantize4BitFunctional: + @staticmethod + def _legacy_nested_dequantize(packed, state, out=None): + absmax = F.dequantize_blockwise(state.absmax, state.state2) + absmax += state.offset + if absmax.dtype != torch.float32: + absmax = absmax.float() + if out is not None: + torch.ops.bitsandbytes.dequantize_4bit.out( + packed, + absmax, + state.blocksize, + state.quant_type, + state.shape, + state.dtype, + out=out, + ) + return out + return torch.ops.bitsandbytes.dequantize_4bit.default( + packed, + absmax, + state.blocksize, + state.quant_type, + state.shape, + state.dtype, + ) + + @pytest.mark.parametrize("device", get_available_devices()) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype) + @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) + @pytest.mark.parametrize("blocksize", [32, 64, 4096], ids=id_formatter("blocksize")) + def test_nested_dequantize_sm103_matches_legacy(self, device, dtype, quant_type, blocksize, monkeypatch): + if device != "cuda": + pytest.skip("The nested CUDA specialization is only available on SM103") + + from bitsandbytes.backends.cuda import ops as cuda_ops + + if not cuda_ops._dequantize_4bit_nested_supported(torch.cuda.current_device()): + pytest.skip("The nested CUDA specialization is only selected on SM103") + + calls = 0 + nested_impl = cuda_ops._dequantize_4bit_nested_impl + + def counted_nested_impl(*args, **kwargs): + nonlocal calls + calls += 1 + return nested_impl(*args, **kwargs) + + monkeypatch.setattr(cuda_ops, "_dequantize_4bit_nested_impl", counted_nested_impl) + + for shape in ((64, 64), (65, 63)): + source = torch.randn(shape, device=device, dtype=dtype) + packed, state = F.quantize_4bit( + source, + blocksize=blocksize, + compress_statistics=True, + quant_type=quant_type, + ) + snapshots = tuple( + tensor.clone() + for tensor in (packed, state.absmax, state.state2.absmax, state.state2.code, state.offset) + ) + + reference = self._legacy_nested_dequantize(packed, state) + actual = F.dequantize_4bit(packed, state) + assert torch.equal(actual, reference) + + out = torch.empty_like(reference) + returned = F.dequantize_4bit(packed, state, out=out) + assert returned.data_ptr() == out.data_ptr() + assert torch.equal(out, reference) + + for current, snapshot in zip( + (packed, state.absmax, state.state2.absmax, state.state2.code, state.offset), + snapshots, + strict=True, + ): + assert torch.equal(current, snapshot) + + assert calls == 4 + + @pytest.mark.parametrize("device", get_available_devices()) + def test_nested_dequantize_sm103_rejects_invalid_out(self, device, monkeypatch): + if device != "cuda": + pytest.skip("The nested CUDA specialization is only available on SM103") + + from bitsandbytes.backends.cuda import ops as cuda_ops + + if not cuda_ops._dequantize_4bit_nested_supported(torch.cuda.current_device()): + pytest.skip("The nested CUDA specialization is only selected on SM103") + + source = torch.randn((65, 63), device=device, dtype=torch.float16) + packed, state = F.quantize_4bit(source, compress_statistics=True, quant_type="nf4") + + def unexpected_nested_impl(*_args, **_kwargs): + raise AssertionError("invalid output reached the nested CUDA kernel") + + monkeypatch.setattr(cuda_ops, "_dequantize_4bit_nested_impl", unexpected_nested_impl) + + with pytest.raises(ValueError, match=r"Expected out\.shape"): + F.dequantize_4bit( + packed, + state, + out=torch.empty((65, 64), device=device, dtype=state.dtype), + ) + + with pytest.raises(ValueError, match=r"Expected out\.dtype"): + F.dequantize_4bit( + packed, + state, + out=torch.empty(state.shape, device=device, dtype=torch.float32), + ) + + with pytest.raises(RuntimeError, match=r"device"): + F.dequantize_4bit( + packed, + state, + out=torch.empty(state.shape, device="cpu", dtype=state.dtype), + ) + + @pytest.mark.parametrize("device", get_available_devices()) + def test_nested_dequantize_non_sm103_uses_legacy(self, device, monkeypatch): + if device != "cuda": + pytest.skip("CUDA dispatch guard test") + + from bitsandbytes.backends.cuda import ops as cuda_ops + + source = torch.randn((65, 63), device=device, dtype=torch.float16) + packed, state = F.quantize_4bit(source, compress_statistics=True, quant_type="nf4") + reference = self._legacy_nested_dequantize(packed, state) + + monkeypatch.setattr(cuda_ops, "_gpu_dispatch_props", lambda _device_index: (148, 10, 0)) + + def unexpected_nested_impl(*_args, **_kwargs): + raise AssertionError("non-SM103 dispatch reached the nested CUDA kernel") + + monkeypatch.setattr(cuda_ops, "_dequantize_4bit_nested_impl", unexpected_nested_impl) + actual = F.dequantize_4bit(packed, state) + assert torch.equal(actual, reference) + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) diff --git a/tests/test_ops.py b/tests/test_ops.py index 4ca60f845..a7f490d20 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -219,6 +219,76 @@ def test_dequantize_4bit(self, device, dtype, storage_dtype, quant_type, blocksi (A, absmax, blocksize, quant_type, shape, dtype), ) + @pytest.mark.parametrize("device", get_available_devices()) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) + @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) + def test_dequantize_4bit_nested(self, device, dtype, quant_type): + source = torch.randn((17, 19), dtype=dtype, device=device) + packed, state = bitsandbytes.functional.quantize_4bit( + source, + blocksize=64, + compress_statistics=True, + quant_type=quant_type, + ) + args = ( + packed, + state.absmax, + state.state2.absmax, + state.state2.code, + state.offset, + state.blocksize, + state.state2.blocksize, + state.quant_type, + state.shape, + state.dtype, + ) + absmax = bitsandbytes.functional.dequantize_blockwise(state.absmax, state.state2) + reference = torch.ops.bitsandbytes.dequantize_4bit.default( + packed, + absmax + state.offset, + state.blocksize, + state.quant_type, + state.shape, + state.dtype, + ) + + out = torch.ops.bitsandbytes.dequantize_4bit_nested.default(*args) + torch.testing.assert_close(out, reference, rtol=0, atol=0) + opcheck(torch.ops.bitsandbytes.dequantize_4bit_nested.default, args) + + out_buffer = torch.empty_like(reference) + torch.ops.bitsandbytes.dequantize_4bit_nested.out(*args, out=out_buffer) + torch.testing.assert_close(out_buffer, reference, rtol=0, atol=0) + opcheck(torch.ops.bitsandbytes.dequantize_4bit_nested.out, (*args, out_buffer)) + + def test_dequantize_4bit_nested_rejects_mismatched_devices(self): + A = torch.zeros((8, 1), dtype=torch.uint8) + absmax_8bit = torch.zeros(1, dtype=torch.uint8) + nested_absmax = torch.ones(1, dtype=torch.float32) + nested_code = torch.ones(256, dtype=torch.float32) + offset = torch.zeros((), dtype=torch.float32) + args = (A, absmax_8bit, nested_absmax, nested_code, offset, 64, 256, "nf4", (4, 4), torch.float16) + + with pytest.raises(RuntimeError, match=r"Expected out\.device"): + torch.ops.bitsandbytes.dequantize_4bit_nested.out( + *args, + out=torch.empty((4, 4), dtype=torch.float16, device="meta"), + ) + + with pytest.raises(RuntimeError, match=r"Expected offset\.device"): + torch.ops.bitsandbytes.dequantize_4bit_nested.default( + A, + absmax_8bit, + nested_absmax, + nested_code, + offset.to("meta"), + 64, + 256, + "nf4", + (4, 4), + torch.float16, + ) + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("storage_dtype", [torch.uint8, torch.bfloat16], ids=id_formatter("storage_dtype")) @@ -374,6 +444,66 @@ def test_gemm_4bit_non_float32_offset(self, device, dtype, offset_dtype): ) torch.testing.assert_close(out, ref) + @pytest.mark.parametrize("device", get_available_devices()) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) + @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) + def test_gemm_4bit_nested_fallback_sm103(self, device, dtype, quant_type, monkeypatch): + if device != "cuda": + pytest.skip("The nested CUDA specialization is only available on SM103") + + from bitsandbytes.backends.cuda import ops as cuda_ops + + if not cuda_ops._dequantize_4bit_nested_supported(torch.cuda.current_device()): + pytest.skip("The nested CUDA specialization is only selected on SM103") + + n, k, blocksize = 64, 64, 64 + activation = torch.randn((5, k), dtype=dtype, device=device) + weight = torch.randn((n, k), dtype=dtype, device=device) + packed, state = bitsandbytes.functional.quantize_4bit( + weight, + blocksize=blocksize, + quant_type=quant_type, + compress_statistics=True, + ) + bias = torch.randn((n,), dtype=dtype, device=device) + + legacy_absmax = bitsandbytes.functional.dequantize_blockwise(state.absmax, state.state2) + legacy_absmax += state.offset + legacy_weight = torch.ops.bitsandbytes.dequantize_4bit.default( + packed, + legacy_absmax.float(), + blocksize, + quant_type, + list(weight.shape), + dtype, + ) + reference = torch.nn.functional.linear(activation, legacy_weight, bias) + + calls = 0 + nested_impl = cuda_ops._dequantize_4bit_nested_impl + + def counted_nested_impl(*args, **kwargs): + nonlocal calls + calls += 1 + return nested_impl(*args, **kwargs) + + monkeypatch.setattr(cuda_ops, "_dequantize_4bit_nested_impl", counted_nested_impl) + actual = cuda_ops._dequant_linear_fallback( + activation, + packed, + list(weight.shape), + state.state2.absmax, + blocksize, + quant_type, + bias, + absmax_8bit=state.absmax, + absmax_code=state.state2.code, + absmax_offset=state.offset, + ) + + assert calls == 1 + assert torch.equal(actual, reference) + class TestNonContiguousInputs: """Regression tests for #1342 and #1690: quantization must handle non-contiguous tensors correctly."""