diff --git a/.github/workflows/gptq-pro-cuda-compile.yml b/.github/workflows/gptq-pro-cuda-compile.yml new file mode 100644 index 000000000..ac7c0c582 --- /dev/null +++ b/.github/workflows/gptq-pro-cuda-compile.yml @@ -0,0 +1,139 @@ +name: GPTQ-Pro Kernel CI + +on: + pull_request: + paths: + - "gptqmodel_ext/gptq_pro/**" + - "gptqmodel/utils/gptq_pro.py" + - "gptqmodel/nn_modules/qlinear/gptq_pro.py" + - "scripts/benchmark_gptq_pro_kernel.py" + - "tests/kernels/test_gptq_pro_ampere_pipeline.py" + - ".github/workflows/gptq-pro-cuda-compile.yml" + push: + branches: + - perf/ampere-kernel-pipeline + paths: + - "gptqmodel_ext/gptq_pro/**" + - "gptqmodel/utils/gptq_pro.py" + - "gptqmodel/nn_modules/qlinear/gptq_pro.py" + - "scripts/benchmark_gptq_pro_kernel.py" + - "tests/kernels/test_gptq_pro_ampere_pipeline.py" + - ".github/workflows/gptq-pro-cuda-compile.yml" + workflow_dispatch: + +permissions: + contents: read + +jobs: + python-style: + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 + with: + python-version: "3.13" + - name: Install Ruff + run: python -m pip install ruff==0.13.0 + - name: Check modified Python sources + run: | + ruff check \ + gptqmodel/utils/gptq_pro.py \ + gptqmodel/nn_modules/qlinear/gptq_pro.py \ + scripts/benchmark_gptq_pro_kernel.py \ + tests/kernels/test_gptq_pro_ampere_pipeline.py + ruff format --check \ + gptqmodel/utils/gptq_pro.py \ + gptqmodel/nn_modules/qlinear/gptq_pro.py \ + scripts/benchmark_gptq_pro_kernel.py \ + tests/kernels/test_gptq_pro_ampere_pipeline.py + python -m py_compile \ + gptqmodel/utils/gptq_pro.py \ + gptqmodel/nn_modules/qlinear/gptq_pro.py \ + scripts/benchmark_gptq_pro_kernel.py \ + tests/kernels/test_gptq_pro_ampere_pipeline.py + + cuda-compile: + runs-on: ubuntu-24.04 + container: + image: nvidia/cuda:12.6.3-devel-ubuntu22.04 + steps: + - uses: actions/checkout@v6 + - name: Install host compiler + run: | + apt-get update + apt-get install -y --no-install-recommends build-essential + rm -rf /var/lib/apt/lists/* + - name: Compile standalone validator for Ampere targets + working-directory: gptqmodel_ext/gptq_pro + run: | + nvcc -O3 -std=c++17 -lineinfo \ + -U__CUDA_NO_HALF_OPERATORS__ \ + -U__CUDA_NO_HALF_CONVERSIONS__ \ + -gencode arch=compute_80,code=sm_80 \ + -gencode arch=compute_86,code=sm_86 \ + -gencode arch=compute_87,code=sm_87 \ + -gencode arch=compute_87,code=compute_87 \ + gptq_pro_validate.cu gptq_pro_kernel.cu \ + -o /tmp/gptq_pro_validate + - name: Record binary metadata + run: | + test -x /tmp/gptq_pro_validate + ls -lh /tmp/gptq_pro_validate + + pytorch-extension-compile: + runs-on: ubuntu-24.04 + container: + image: nvidia/cuda:12.6.3-devel-ubuntu22.04 + env: + TORCH_CUDA_ARCH_LIST: "8.0;8.6;8.7+PTX" + steps: + - uses: actions/checkout@v6 + - name: Install isolated Python build environment + run: | + apt-get update + apt-get install -y --no-install-recommends \ + build-essential python3 python3-pip python3-dev python3-venv ninja-build + rm -rf /var/lib/apt/lists/* + python3 -m venv /tmp/gptq-pro-ci + /tmp/gptq-pro-ci/bin/python -m pip install --upgrade pip setuptools wheel ninja + /tmp/gptq-pro-ci/bin/python -m pip install \ + --index-url https://download.pytorch.org/whl/cu126 \ + torch==2.8.0 + /tmp/gptq-pro-ci/bin/python -m pip install pytest + - name: Compile and import PyTorch extension + run: | + mkdir -p /tmp/gptq_pro_extension_ci + /tmp/gptq-pro-ci/bin/python - <<'PY' + from pathlib import Path + from torch.utils.cpp_extension import load + + root = Path.cwd() + extension = load( + name="gptqmodel_gptq_pro_ci", + sources=[ + str(root / "gptqmodel_ext/gptq_pro/gptq_pro_torch.cpp"), + str(root / "gptqmodel_ext/gptq_pro/gptq_pro_kernel.cu"), + ], + extra_cflags=["-O3", "-std=c++17"], + extra_cuda_cflags=[ + "-O3", + "-std=c++17", + "-lineinfo", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-gencode=arch=compute_80,code=sm_80", + "-gencode=arch=compute_86,code=sm_86", + "-gencode=arch=compute_87,code=sm_87", + "-gencode=arch=compute_87,code=compute_87", + ], + build_directory="/tmp/gptq_pro_extension_ci", + verbose=True, + ) + assert hasattr(extension, "gptq_pro_gemm") + print(extension.gptq_pro_gemm.__doc__) + PY + - name: Run CPU packing and dispatch contracts + run: | + /tmp/gptq-pro-ci/bin/python -m pytest -q \ + --confcutdir=tests/kernels \ + tests/kernels/test_gptq_pro_ampere_pipeline.py diff --git a/README.md b/README.md index 7a814cbaf..619a81ffb 100644 --- a/README.md +++ b/README.md @@ -2,8 +2,8 @@ **GPTQ-Pro** is an experimental fork of [ModelCloud/GPTQModel](https://github.com/ModelCloud/GPTQModel) focused on one -runtime path: a custom symmetric INT4 GPTQ CUDA kernel for modern NVIDIA GPUs, -with native build targets for Ampere (`sm_80`, `sm_86`, and `sm_87`). +runtime path: symmetric INT4 GPTQ inference on modern NVIDIA GPUs, with native +build targets for Ampere (`sm_80`, `sm_86`, and `sm_87`). > This is not the official GPTQModel release. Use upstream GPTQModel when you > need its full multi-backend feature set. This fork deliberately removes AWQ, @@ -21,38 +21,72 @@ The Python distribution and import names remain **`GPTQModel`** and | Checkpoint formats | `FORMAT.GPTQ`, `FORMAT.GPTQ_V2` | | Runtime selectors | `BACKEND.AUTO`, `BACKEND.GPTQ_PRO` | | Runtime platform | Linux + NVIDIA CUDA, compute capability 8.0 or newer | -| Runtime weights | Symmetric 4-bit GPTQ, `desc_act=False`, `int32` packing | +| Runtime weights | Symmetric 4-bit GPTQ, `desc_act=False`, native `int32 qweight` packing | | Runtime activations | FP16; non-FP16 inputs are converted to FP16 | | Group sizes | `-1`, `16`, `32`, `64`, `128`, `256`, `512`, `1024` | | Shape constraints | input features divisible by 16; output features divisible by 8 | | Adapters | LoRA on the GPTQ-Pro linear path | `BACKEND.AUTO_TRAINABLE` remains in the compatibility enum, but the only runtime -kernel in this fork declares `SUPPORTS_TRAINING=False`; there is currently no -trainable quantized backend. - -The CUDA extension embeds native cubins for `sm_80`/`sm_86`/`sm_87` and an -`compute_87` PTX fallback. Ada and Hopper devices may therefore JIT-compile the -PTX, but Ampere is the primary development and validation target. ROCm, MPS, -CPU inference, asymmetric weights, act-order checkpoints, native BF16, and -non-4-bit inference are not supported by the GPTQ-Pro kernel. - -## Kernel status - -`gptqmodel_ext/gptq_pro/` contains the custom INT4 dequantization GEMM. It uses -Tensor Core `mma.sync` with FP32 accumulation, but it is still a research -scaffold: - -- one warp per CTA; -- no `cp.async` multi-stage pipeline; -- no `ldmatrix` data path; -- scalar INT4 decode; -- no dedicated small-`M`/GEMV decode kernel; -- no native BF16 or training path. - -No repository benchmark currently proves that it is faster than mature external -kernels. `BACKEND.AUTO` selects GPTQ-Pro because it is the only inference kernel -shipped here, not because a performance lead has been established. See +kernel declares `SUPPORTS_TRAINING=False`; there is currently no trainable +quantized backend. + +The extension embeds native cubins for `sm_80`/`sm_86`/`sm_87` and a +`compute_87` PTX fallback. Ada and Hopper devices may JIT-compile that PTX, but +Ampere is the primary development target. ROCm, MPS, CPU inference, asymmetric +weights, act-order checkpoints, native BF16, and non-4-bit inference are not +supported by the GPTQ-Pro runtime. + +## Ampere kernel architecture + +`gptqmodel_ext/gptq_pro/` contains three dispatch paths: + +1. **Decode / very small batches (`M <= 4`)** + - one CUDA thread computes one output column; + - native `qweight` reads are coalesced across each warp; + - four adjacent activation pairs are decoded from each GPTQ `int32` word; + - FP16 dequantization is preserved before FP32 accumulation; + - designed for the memory-bandwidth-bound token-generation regime. +2. **Aligned medium-batch and prefill GEMM** + - four warps per CTA, producing a `16 × 256` output region; + - double-buffered `cp.async` global-to-shared pipeline; + - coalesced 16-byte activation, native-qweight, and scale copies; + - Tensor Core `mma.sync.m16n8k16` with FP32 accumulation; + - LOP3-assisted INT4-to-FP16 fragment conversion; + - vectorized FP16 output stores. +3. **General-shape fallback** + - retains the validator-backed one-warp implementation; + - handles compatible edge shapes that do not satisfy the optimized path's + `N % 16 == 0` requirement. + +All three paths consume the checkpoint's original `int32 qweight` directly. The +runtime no longer materializes a second pair-packed weight tensor, avoiding a +persistent duplicate of the model's quantized weights and eliminating a packing +transform during module initialization. + +`AUTO` selects the small-`M` kernel first, then the pipelined Tensor Core path, +then the general fallback. Experts can force a path for parity testing: + +```bash +export GPTQMODEL_GPTQ_PRO_KERNEL=gemv # auto, gemv, ampere, or legacy +``` + +The override is intended for diagnostics and benchmarking. Forced modes reject +incompatible shapes rather than silently changing behavior. + +### Remaining kernel work + +The runtime is materially more Ampere-aware than the original scaffold, but the +following are still open engineering targets: + +- GPU-measured crossover tuning instead of the initial `M <= 4` decode rule; +- `ldmatrix`-based shared-to-register loading and shared-memory swizzling; +- additional CTA tile families and architecture-specific autotuning; +- native BF16 Tensor Core execution; +- profiler-guided scale caching and wider prefetch strategies; +- production GPU validation across the full advertised shape/group matrix. + +No performance lead is claimed without checked-in measurements. See [`docs/ASSESSMENT_AND_ROADMAP.md`](docs/ASSESSMENT_AND_ROADMAP.md). ## Installation @@ -62,9 +96,9 @@ shipped here, not because a performance lead has been established. See - Linux; - Python 3.10 or newer; - an NVIDIA GPU with compute capability 8.0 or newer; -- a CUDA toolkit containing `nvcc` that is compatible with the installed PyTorch; +- a CUDA toolkit containing `nvcc` compatible with the installed PyTorch; - a C++17 compiler and Ninja; -- sufficient free disk space for extension builds and optional quantization offload. +- enough free disk space for extension builds and optional quantization offload. ### From source @@ -79,9 +113,9 @@ python -m pip install --upgrade pip wheel setuptools python -m pip install -e . ``` -The GPTQ-Pro CUDA extension is loaded from a compatible prebuilt module when one -is available; otherwise it is JIT-compiled on first use. Set -`GPTQMODEL_EXT_BUILD=/path/to/cache` to choose its build directory and +The CUDA extension is loaded from a compatible prebuilt module when available; +otherwise it is JIT-compiled on first use. Set +`GPTQMODEL_EXT_BUILD=/path/to/cache` to choose the build directory and `GPTQMODEL_EXT_VERBOSE=1` for verbose compilation logs. The public extension registry manages only the generic CPU helper extensions: @@ -128,15 +162,16 @@ model.quantize(calibration, batch_size=1, backend=BACKEND.AUTO) model.save("model-gptq-pro-4bit") ``` -Use representative, sufficiently varied calibration samples. Large models can -use disk offload and explicit dense/MoE device pools through `QuantizeConfig`. -Avoid a Transformers `device_map="auto"` during quantization; placement is -managed by GPTQ-Pro's quantization loop. +Use representative, varied calibration samples. Large models can use disk +offload and explicit dense/MoE device pools through `QuantizeConfig`. Avoid a +Transformers `device_map="auto"` during quantization; placement is managed by +GPTQ-Pro's quantization loop. ## Quantization recipes -The recipe controls **offline quantization**. The runtime kernel is selected -separately when loading the quantized checkpoint. +The recipe controls **offline quantization quality**. CUDA kernel optimization +changes inference latency and throughput; it does not improve an already-created +checkpoint's numerical quality. | Recipe | Enabled behavior | |---|---| @@ -147,12 +182,38 @@ separately when loading the quantized checkpoint. Important details: -- `max_quality_4bit()` does **not** enable FOEM or rotation automatically. - Supply those explicitly when experimenting with them. +- `max_quality_4bit()` is the highest-quality named 4-bit preset currently + provided, but it does **not** automatically enable FOEM or rotation. +- `quality_4bit()` is the recommended balance when quantization time matters. - The generic `fast_4bit()` preset inherits the base act-order default unless - `desc_act=False` is passed. Act-order checkpoints cannot run on GPTQ-Pro. -- The 3-bit preset emits a standard GPTQ checkpoint for external experimentation, - but this single-backend fork has no local 3-bit inference kernel. + `desc_act=False` is supplied. Act-order checkpoints cannot run on GPTQ-Pro. +- The 3-bit preset is export-only in this fork because no local 3-bit runtime is + shipped. + +## Raw-kernel benchmark + +Measure the CUDA extension without model-loading or tokenization overhead: + +```bash +python scripts/benchmark_gptq_pro_kernel.py \ + --m-values 1,4,8,16,64,256 \ + --n 4096 --k 4096 --group-size 128 \ + --warmup 20 --iterations 100 \ + --output kernel-results.json +``` + +The benchmark: + +- reports median, mean, and p95 kernel time; +- reports effective dense TFLOP/s and a conservative bandwidth lower bound; +- compares AUTO against forced specialized and legacy modes; +- checks a column subset against a PyTorch FP32 reference using FP16-dequantized + native `qweight` values; +- records GPU, compute capability, PyTorch, and CUDA runtime information. + +Commit raw JSON results when making performance claims. Use Nsight Compute to +confirm whether changes improve memory throughput, Tensor Core utilization, +occupancy, and stall reasons rather than relying only on wall-clock timing. ## Qwen 3.5 and Qwen 3.6 @@ -180,7 +241,7 @@ CUDA_VISIBLE_DEVICES=0 python scripts/quant_qwen36_obliterated_gptqpro.py \ --calibration-jsonl /path/to/calibration.jsonl \ --preset quality --nsample 64 --seqlen 512 --offload-disk --dry-run -# Qwen3.5/Qwen3.6 MoE; auto chooses image calibration for multimodal layouts +# Qwen3.5/Qwen3.6 MoE CUDA_VISIBLE_DEVICES=0 python scripts/quant_qwen3_5_moe.py \ --model /path/or/hf-id \ --out /path/to/output \ @@ -194,8 +255,7 @@ repository-provided model code. ## Selective / mixed-precision quantization `QuantizeConfig.dynamic` accepts PCRE module-name overrides. A `-:` prefix skips -matching modules even when the model definition normally marks them as -quantizable. +matching modules even when the model definition marks them as quantizable. ```python qcfg = QuantizeConfig.quality_4bit(group_size=128) @@ -212,27 +272,46 @@ anchored skip pattern. ## Validation -CPU-only regression checks for the documented contracts: +CPU-side checks: ```bash pytest -q \ tests/qcfg/test_gptq_pro.py \ tests/kernels/test_selection.py \ + tests/kernels/test_gptq_pro_ampere_pipeline.py \ tests/models/test_qwen3_5_invariants.py \ tests/models/test_qwen3_5_vision.py \ tests/test_qwen3_6_support.py \ tests/test_extension_registry.py ``` -A real CUDA build, numerical comparison against the dense source checkpoint, -and generation/perplexity checks are still required before treating a new -quantized model as production-ready. +The kernel CI performs three independent gates: + +1. targeted Ruff/style and Python syntax checks; +2. standalone CUDA compilation for `sm_80`, `sm_86`, `sm_87`, and PTX; +3. compilation and import of the actual PyTorch C++/CUDA binding, followed by + CPU packing/dispatch tests. + +On a real GPU, build and run the standalone numerical validator: + +```bash +cd gptqmodel_ext/gptq_pro +nvcc -O3 -std=c++17 -arch=sm_80 \ + -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ \ + gptq_pro_validate.cu gptq_pro_kernel.cu -o gptq_pro_validate +./gptq_pro_validate +``` + +A real CUDA run, numerical comparison against the dense source checkpoint, and +generation/perplexity checks remain required before treating a newly quantized +model as production-ready. ## Repository map - `gptqmodel/` — model loading, quantization, packing, and runtime integration; - `gptqmodel_ext/gptq_pro/` — GPTQ-Pro CUDA/C++ extension sources; -- `scripts/` — quantization and validation helpers; +- `scripts/benchmark_gptq_pro_kernel.py` — raw CUDA benchmark and parity checks; +- `scripts/` — model quantization and validation helpers; - `tests/` — CPU and CUDA regression coverage; - `docs/ASSESSMENT_AND_ROADMAP.md` — current engineering status and roadmap; - `docs/QWEN35_QWEN36.md` — Qwen architecture and driver guide. diff --git a/docs/ASSESSMENT_AND_ROADMAP.md b/docs/ASSESSMENT_AND_ROADMAP.md index c7f3c0083..5b35a3d96 100644 --- a/docs/ASSESSMENT_AND_ROADMAP.md +++ b/docs/ASSESSMENT_AND_ROADMAP.md @@ -1,41 +1,42 @@ # GPTQ-Pro — current assessment and roadmap -> Updated for the single-backend repository state on 11 July 2026. +> Updated for the Ampere kernel pipeline on 11 July 2026. > -> This document distinguishes code that exists in this fork from external -> comparison baselines. Marlin, ExLlama, Machete, vLLM, SGLang, AWQ, BitBLAS, -> MLX, and the other removed paths are **not shipped by GPTQ-Pro**. +> This document distinguishes code shipped by this fork from external comparison +> baselines. Marlin, ExLlama, Machete, vLLM, SGLang, AWQ, BitBLAS, MLX, and the +> other removed paths are **not bundled by GPTQ-Pro**. ## Executive summary -GPTQ-Pro is now a deliberately narrow research fork: +GPTQ-Pro is a deliberately narrow research fork: - quantization method: GPTQ only; - checkpoint formats: GPTQ and GPTQ_V2; -- inference kernel: GPTQ-Pro only; +- inference backend: GPTQ-Pro only; - primary hardware: Linux NVIDIA GPUs, especially Ampere; - local runtime format: symmetric 4-bit, `desc_act=False`, FP16 activations, - `int32` packing. + native `int32 qweight` packing. -The quantization code has a strong set of optional quality mechanisms, but the -runtime kernel is still a performance scaffold. The project is useful for -controlled experimentation, Qwen 3.5/3.6 quantization work, and kernel -engineering. It should not yet be presented as a production-proven or -benchmark-leading runtime. +The runtime is no longer a single one-warp kernel. AUTO now dispatches among: -The highest-priority work is no longer “select the best bundled backend.” There -is no bundled fallback backend. The priority is to prove correctness and -performance of GPTQ-Pro itself, then close the measured gaps. +- a coalesced small-`M` decode kernel; +- a four-warp, double-buffered `cp.async` Tensor Core kernel; +- a general-shape correctness fallback. + +Every path consumes the checkpoint's native `qweight` tensor directly, avoiding +a second persistent copy of quantized model weights. This closes several obvious +Ampere design and memory-footprint gaps, but implementation is not the same as +proof. CUDA compilation is continuously checked; real GPU execution, profiler +evidence, numerical matrices, and checked-in benchmark results remain release +gates. Until those gates pass, GPTQ-Pro should not be described as benchmark- +leading or production-proven. ## Current implementation ### 1. Public package surface The distribution and import names remain `GPTQModel` and `gptqmodel` for -upstream compatibility. The project metadata points to this fork and no longer -advertises removed backend extras. - -Current enums: +upstream API and checkpoint compatibility. | Dimension | Values | |---|---| @@ -43,30 +44,25 @@ Current enums: | `FORMAT` | `GPTQ`, `GPTQ_V2` | | `BACKEND` | `AUTO`, `AUTO_TRAINABLE`, `GPTQ_PRO` | -`AUTO_TRAINABLE` is retained as an API compatibility selector, but the only -runtime kernel declares `SUPPORTS_TRAINING=False`. It cannot currently resolve -to a trainable quantized implementation. +`AUTO_TRAINABLE` remains for compatibility, but the only runtime implementation +sets `SUPPORTS_TRAINING=False`. ### 2. Quantization quality mechanisms -The repository retains the following GPTQ-side capabilities: +The repository retains: - Hessian/OBS-style GPTQ error compensation; - group-wise quantization; - group-aware reordering (`act_group_aware`); - optional act-order (`desc_act`); -- MSE scale search; -- activation-weighted MSE; +- MSE and activation-weighted MSE scale search; - adaptive Cholesky damping; - fallback smoothing for poorly sampled blocks; - GPTAQ activation-aware error feedback; - optional FOEM; - optional Hadamard/random rotation; -- dynamic per-module overrides and skip rules; -- dense and MoE placement controls; -- disk offload during quantization. - -These features are not all enabled by one preset. +- dynamic per-module overrides; +- dense/MoE placement controls and disk offload. | Preset | Exact intent | |---|---| @@ -75,304 +71,307 @@ These features are not all enabled by one preset. | `max_quality_4bit()` | `quality_4bit()` plus GPTAQ | | `experimental_3bit_rotation()` | 3-bit GPTQ + GPTAQ + Hadamard rotation for supported architectures | -Important constraints: +Constraints: -- `max_quality_4bit()` does not automatically enable FOEM or rotation. -- The generic `fast_4bit()` preset inherits the base act-order default unless - `desc_act=False` is supplied. -- The 3-bit preset can be used for standard-GPTQ export experiments, but the - local GPTQ-Pro kernel is 4-bit only and cannot execute the result. +- `max_quality_4bit()` does not automatically enable FOEM or rotation; +- `fast_4bit()` requires explicit `desc_act=False` for this runtime; +- the local kernel is 4-bit only, so the 3-bit preset is export-only here; +- faster CUDA execution does not improve the numerical quality of an existing + quantized checkpoint. Runtime and quantization quality must be evaluated + separately. -### 3. GPTQ-Pro runtime kernel +### 3. Runtime contract -The runtime class is `GptqProQuantLinear`. Its validated contract is: +`GptqProQuantLinear` validates: -- Linux; -- CUDA, non-ROCm; +- Linux CUDA, non-ROCm; - compute capability 8.0 or newer; - FP16 activations; - symmetric 4-bit weights; - `desc_act=False` and sequential `g_idx`; -- `int32` packing; +- native GPTQ `int32 qweight`; - input features divisible by 16; - output features divisible by 8; -- group size `-1` or one of `16`, `32`, `64`, `128`, `256`, `512`, `1024`. +- group size `-1` or `16`, `32`, `64`, `128`, `256`, `512`, `1024`. -The extension build contains native SASS for `sm_80`, `sm_86`, and `sm_87`, -plus `compute_87` PTX for driver-side JIT on newer architectures. +The extension contains native SASS for `sm_80`, `sm_86`, and `sm_87`, plus +`compute_87` PTX for driver-side JIT on newer architectures. -Implemented kernel characteristics: +### 4. Kernel dispatch -- Tensor Core `mma.sync`; -- FP32 accumulation; -- group-based symmetric INT4 dequantization; -- one warp per CTA; -- one general GEMM path for both prompt processing and decode. +#### Small-`M` decode path -Missing production-kernel characteristics: +For `M <= 4` and qword-compatible `K`, AUTO selects a dedicated GEMV-style +kernel: -- no dedicated GEMV/small-`M` decode path; -- no `cp.async` multi-stage shared-memory pipeline; -- no `ldmatrix` path; -- no fused LOP3-style INT4 decode; -- no broad vectorized load strategy; -- no multi-warp tile family or autotuning; -- no native BF16 path; -- no training support. +- one thread owns one output column; +- native qweight and scale loads are contiguous across each warp; +- each `int32` word supplies four adjacent `half2` activation/weight pairs; +- weight dequantization rounds to FP16 before FP32 accumulation, matching the + runtime's Tensor Core numerical contract. -`BACKEND.AUTO` chooses GPTQ-Pro because it is the only local inference kernel, -not because a benchmark proves it is faster than external alternatives. +The initial crossover is conservative and must be tuned from real GPU data. -### 4. Extension loading +#### Ampere Tensor Core path -There are two distinct extension mechanisms: +For aligned `K` and `N % 16 == 0`, AUTO selects a CTA tile that uses: -1. `gptqmodel.extension` manages the generic CPU helper extensions - `pack_block_cpu` and `floatx_cpu`. -2. `gptqmodel.utils.gptq_pro.ensure_gptq_pro_loaded()` loads or JIT-compiles the - CUDA inference extension on first use. +- four warps per CTA; +- a `16 × 256` output region; +- two shared-memory stages; +- `cp.async` 16-byte global-to-shared copies; +- overlap of the next K tile's movement with current-tile MMA work; +- a shared A tile and warp-private native-qweight/scale tiles; +- `mma.sync.m16n8k16` with FP32 accumulation; +- LOP3-assisted INT4-to-FP16 conversion; +- vectorized `half2` stores. -The removed `marlin` extension alias is not part of the public registry. +A 16-wide K tile reads two native qweight rows. Each lane moves four adjacent +`int32` words, preserving coalesced 16-byte transactions without repacking the +model. -### 5. Qwen 3.5 / Qwen 3.6 +#### General-shape fallback -Qwen 3.6 reuses Qwen 3.5 Transformers model types. The repository has explicit -routing for dense and MoE, multimodal and text-only, plus nested -`language_model_only=true` MoE conversions. +The original one-warp implementation remains available for compatible edge +shapes and is selectable explicitly as `legacy`. It is not the preferred +production path, but retaining it avoids converting an optimization change into +a silent shape-support regression. -Current guarantees: +#### Expert dispatch override -- both linear-attention and full-attention projections are represented in the - quantization module trees; -- Q/K norms, recurrent helpers, convolution, routers, and layer norms stay in - source precision; -- shared-expert traversal precedes routed experts; -- multimodal vision towers are materialized for calibration but not quantized; -- `mtp.*` tensors are preserved by the writer even when Transformers does not - instantiate them; -- the LM-only wrapper removes the unused vision tower before source-model - calibration. +`GPTQMODEL_GPTQ_PRO_KERNEL` accepts `auto`, `gemv`, `ampere`, or `legacy`. +Forced modes fail on incompatible shapes. The override is intended for testing, +not as a normal user requirement. -See [`QWEN35_QWEN36.md`](QWEN35_QWEN36.md). +### 5. Native qweight memory path -### 6. Validation status +GPTQ stores eight K-axis nibbles in every `int32 qweight` word. Earlier versions +of this optimization branch materialized a second pair-packed byte tensor. The +final design removes that buffer: decode, Tensor Core, and legacy kernels all +extract the required bytes from native qweight words directly. -CPU-only tests cover configuration behavior, kernel selection/rejection, -Qwen definition routing, vision lifecycle invariants, MTP declarations, and the -extension registry. +Benefits: -The following remain environment-dependent and are not proven merely by unit -tests: +- no persistent duplicate of quantized weights; +- no module-initialization repacking transform; +- no eight-times-larger temporary nibble broadcast; +- unchanged global bytes per 16-wide K tile; +- simpler checkpoint/runtime memory accounting. -- successful JIT compilation with each supported CUDA/PyTorch/toolkit pair; -- numerical parity across all supported shapes and group sizes; -- long-context generation stability; -- real multimodal generation after quantization; -- perplexity/task-quality regression bounds; -- throughput and latency versus the dense source model or external kernels; -- multi-GPU memory behavior for very large MoE layers. +The historical byte-view helper remains available for inspection tooling but is +not used by inference. Channelwise `group_size=-1` is normalized to +`in_features` before validating `g_idx` and invoking CUDA. -## Discrepancies removed in this documentation pass +### 6. Extension and build validation -The previous roadmap was written before the repository was reduced to one -backend and still described deleted code as locally available. The following -claims are no longer made: +There are two extension mechanisms: -- selecting `BACKEND.MARLIN` or falling through to Marlin/Machete; -- building a bundled Marlin extension through `gptqmodel.extension`; -- using vendored AutoRound or ParoQuant trees; -- serving through bundled vLLM or SGLang paths; -- exporting through bundled MLX support; -- treating `max_quality_4bit()` as enabling FOEM and rotation; -- treating the experimental 3-bit recipe as runnable by GPTQ-Pro; -- claiming measured superiority without a checked-in benchmark result. +1. `gptqmodel.extension` manages generic CPU helpers; +2. `ensure_gptq_pro_loaded()` loads or JIT-compiles the CUDA runtime. -External projects can still be useful comparison baselines, but references to -them below do not imply that their code ships in this fork. +Kernel CI contains independent gates for: -## Prioritized roadmap +- targeted Python lint, formatting, and syntax; +- standalone CUDA compilation for native `sm_80`/`sm_86`/`sm_87` plus PTX; +- compilation and import of the actual PyTorch C++/CUDA extension binding; +- CPU tests for qweight layout and dispatch contracts. -### P0 — correctness and reproducibility +Compilation proves that source and binding signatures are accepted by the +configured toolchains; it does not execute CUDA kernels. -#### C1. CUDA build matrix +The standalone validator covers: -Add reproducible build smoke tests for at least: +- LOP3 fragment decode; +- forced GEMV; +- forced pipelined Ampere mode; +- AUTO decode and AUTO GEMM selection; +- M and N tails; +- odd edge shapes through the legacy fallback; +- native qweight extraction in every path. -- RTX 3090 / `sm_86`; -- A100 / `sm_80`; -- one PTX-JIT device such as Ada or Hopper; -- supported Python and PyTorch ranges; -- clean build and cached rebuild paths. +Those numerical checks must still be run on real hardware. -Acceptance criteria: +### 7. Raw benchmark -- extension compiles from a clean checkout; -- a second process loads the cached artifact; -- failure messages identify compiler, CUDA, or architecture incompatibilities; -- artifacts and environment versions are recorded. +`scripts/benchmark_gptq_pro_kernel.py` records: -#### C2. Numerical reference harness +- device and compute capability; +- PyTorch and CUDA runtime versions; +- warm-up and iteration counts; +- median, mean, and p95 latency; +- effective dense TFLOP/s; +- a conservative memory-bandwidth lower bound; +- numerical error and cosine similarity against a PyTorch reference; +- forced AUTO/GEMV/Ampere/legacy comparisons where applicable. -Create a shape and group-size matrix comparing the CUDA result with a clear -PyTorch reference implementation. +The benchmark generates and dequantizes native `int32 qweight`, disables TF32 +for the FP32 reference, and includes only one quantized-weight tensor in its +bandwidth accounting. Performance claims should include raw JSON and the exact +commit. -Cover: +### 8. Qwen 3.5 / Qwen 3.6 -- `M` in `{1, 8, 64, 512}`; -- odd and boundary-compatible `N`/`K` sizes; -- every advertised group size; -- bias and LoRA paths; -- multiple random seeds and scale distributions; -- finite-value, absolute-error, relative-error, and cosine-similarity checks. +Qwen 3.6 reuses Qwen 3.5 Transformers model types. The repository explicitly +routes dense and MoE, multimodal and text-only, plus nested +`language_model_only=true` MoE conversions. -Acceptance criteria must be encoded in tests rather than documented informally. +Current guarantees include hybrid linear/full-attention module trees, source- +precision Q/K norms and recurrent helpers, shared-before-routed expert order, +source-precision vision towers, and writer preservation of `mtp.*` tensors. +See [`QWEN35_QWEN36.md`](QWEN35_QWEN36.md). -#### C3. Checkpoint round-trip tests +## Validation status -For dense, MoE, text-only, and multimodal fixtures: +### Proven by checked-in automation + +- targeted Python lint/format/syntax checks; +- qweight layout and kernel-mode normalization tests; +- source-contract tests for specialized dispatch and no duplicate weight buffer; +- standalone CUDA compilation for Ampere SASS targets and PTX; +- actual PyTorch extension compilation/import gate; +- Qwen routing and preservation invariants; +- package/manifest consistency checks. + +### Not yet proven in this environment -- quantize; -- save; -- reload through `BACKEND.AUTO`; -- compare deterministic generation; -- verify all expected `mtp.*` tensors; -- verify the vision tower remains in source precision; -- verify the quantization metadata matches the runtime contract. +- standalone validator execution on `sm_80`, `sm_86`, and `sm_87`; +- numerical parity across every group size and production shape; +- Nsight Compute counters for memory throughput, occupancy, Tensor Core use, + register pressure, and stall reasons; +- measured decode/prefill crossover points; +- model-level generation, perplexity, and task-quality regression; +- long-context and multimodal stability; +- multi-GPU behavior for very large MoE layers. -#### C4. Package validation +## Prioritized roadmap -Build both sdist and wheel artifacts and verify they contain: +### P0 — correctness and reproducibility -- `gptqmodel_ext/gptq_pro` CUDA/C++ sources; -- `pack_block_cpu.cpp` and `floatx_cpu.cpp`; -- no deleted backend source trees; -- correct project URLs and optional-dependency metadata. +#### C1. Execute the CUDA validation matrix -### P1 — performance +Run the standalone validator on at least: -#### K1. Dedicated decode kernel +- RTX 3090 / `sm_86`; +- A100 / `sm_80`; +- Jetson Orin / `sm_87` when available; +- an Ada or Hopper PTX-JIT target. -Implement a fused dequantization GEMV or small-`M` kernel and dispatch below a -measured crossover point. Single-token decode is primarily memory-bandwidth -bound and should not use the same tile strategy as large prompt batches. +Record driver, toolkit, compiler, PyTorch, clock/power state, and exact commit. -#### K2. Ampere asynchronous pipeline +#### C2. Expand the numerical matrix -Add staged global-to-shared copies with `cp.async`, double or multi-buffering, -and overlap between memory movement and MMA work. +Test: -#### K3. Vectorized data path +- all advertised group sizes; +- `M` around dispatch boundaries; +- aligned and tail `N` values; +- multiple K sizes and random seeds; +- bias and LoRA paths; +- finite-value, absolute/relative-error, and cosine bounds. -Evaluate: +Acceptance criteria must be executable tests, not prose. -- 128-bit global loads; -- `ldmatrix` shared-to-register movement; -- fused bit manipulation/dequantization; -- scale prefetching and reduced repeated staging. +#### C3. Model round trips -Each change must be justified by profiler counters, not only kernel timing. +For dense, MoE, text-only, and multimodal fixtures: -#### K4. Multi-warp tiling and dispatch +- quantize with `quality_4bit()` and `max_quality_4bit()`; +- save and reload through AUTO; +- compare deterministic generation and perplexity; +- verify `mtp.*`, vision precision, metadata, and kernel compatibility. -Introduce multiple tile families for decode, medium batches, and prefill. Select -using dimensions and measured architecture-specific crossovers. Avoid claiming -one universal kernel configuration. +#### C4. Wheel and sdist validation -#### K5. Native BF16 +Build artifacts and verify that CUDA sources and CPU helpers are present, removed +backend trees are absent, and a clean environment can JIT-compile the runtime. -Add a BF16 MMA path instead of unconditionally converting activations to FP16. -Validate overflow-sensitive layers and compare accuracy and performance with the -FP16 path. +### P1 — Ampere performance -#### K6. Benchmark protocol +#### K1. Measure and tune decode crossover -Check in a benchmark that records: +Benchmark `M` from 1 through at least 64 on RTX 3090 and A100. Replace the +initial `M <= 4` rule with architecture/shape-aware thresholds only after stable +measurements. -- GPU, driver, CUDA toolkit, PyTorch, and compiler versions; -- model, sequence length, batch size, prompt/decode split; -- warm-up and synchronization procedure; -- median and percentile latency; -- tokens/s; -- peak allocated/reserved memory; -- numerical error against the reference. +#### K2. Profile the async pipeline -External Marlin or ExLlama results may be included only when those dependencies -are installed separately and clearly labeled as external baselines. +Use Nsight Compute to verify: -### P2 — quantization quality +- global-load coalescing; +- asynchronous-copy overlap; +- shared-memory bank behavior; +- Tensor Core utilization; +- eligible warps and occupancy; +- register spills and long-scoreboard stalls. -#### Q1. Calibration guidance and fixtures +#### K3. Add `ldmatrix` and shared-memory swizzling -Provide representative public calibration fixtures for coding, general text, -multilingual text, and vision-language workloads. Measure sensitivity to sample -count and sequence length. +The current path still manually packs A fragments from shared memory. Evaluate a +bank-conflict-safe `ldmatrix` layout and retain it only if profiler and numerical +results justify the complexity. -#### Q2. Preset evaluation +#### K4. Cache scales across K tiles -Evaluate `fast_4bit`, `quality_4bit`, and `max_quality_4bit` on the same models -and tasks. Record quantization time, disk usage, perplexity/task quality, and -runtime compatibility. +Group sizes usually span several 16-wide K tiles. Reduce repeated scale staging +by caching one group's scales in registers or a persistent shared region. -#### Q3. FOEM policy +#### K5. Add tile families and autotuning -FOEM is available but not part of `max_quality_4bit()`. Determine whether it -provides repeatable gains across model families before adding a named preset or -enabling it by default. +Evaluate CTA shapes for: -#### Q4. Low-bit runtime decision +- decode and micro-batch; +- medium batch; +- long-prompt prefill; +- narrow and wide output dimensions. -The repository exposes a 3-bit rotation recipe but has no 3-bit runtime. Choose -one of two explicit directions: +Autotuning must be cached and bounded so startup cost does not dominate normal +use. -1. add and validate a compatible 3-bit kernel; or -2. keep the recipe as export-only and mark it experimental throughout the API. +#### K6. Native BF16 -Do not imply local inference support until one of those paths is complete. +Add BF16 MMA and output handling rather than converting all activations to FP16. +Validate overflow-sensitive layers and compare quality and throughput. -#### Q5. Model-specific mixed precision +#### K7. Kernel launch and graph integration + +Measure launch overhead during token generation. Evaluate CUDA Graph capture and +persistent metadata without compromising dynamic-shape correctness. + +### P2 — quantization quality + +#### Q1. Calibration fixtures + +Provide public coding, general-text, multilingual, and vision-language fixtures. +Measure sample-count and sequence-length sensitivity. + +#### Q2. Preset evaluation -Build reproducible sensitivity tooling for preserving selected embeddings, -output heads, attention outputs, or expert projections in source precision. -Store generated skip lists with model/config fingerprints so they cannot be -silently reused on incompatible layouts. +For representative dense and MoE models, compare FP16/BF16 source, +`fast_4bit(desc_act=False)`, `quality_4bit()`, and `max_quality_4bit()` using +perplexity, tasks, deterministic prompts, and calibration time/memory. -### P3 — maintainability and releases +#### Q3. FOEM and rotation policy -- Add a repository consistency test that rejects removed backend names in - package metadata and user-facing documentation. -- Add link and command checks for Markdown. -- Add a release checklist containing sdist/wheel inspection, clean-container - install, CUDA build, quantize/save/load, and documentation verification. -- Record the upstream GPTQModel commit used for every sync. -- Keep historical credits while clearly distinguishing inherited history from - code currently shipped. -- Avoid publishing the compatibility distribution name to a public index unless - the release process explicitly addresses collision with upstream GPTQModel. +FOEM and rotation remain opt-in. Promote either into a named 4-bit preset only +after repeatable improvements across multiple models and workloads. -## Recommended release gate +#### Q4. Selective precision recipes -A release should not be described as production-ready until all of the following -are true: +Add reproducible module-sensitivity tooling and documented BF16 preservation +recipes for embeddings, output heads, recurrent helpers, expert routers, and +model-specific outliers. -1. clean installation succeeds in the documented container; -2. CUDA build and cache reuse pass on `sm_80` and `sm_86`; -3. numerical reference tests pass across the advertised shape/group matrix; -4. at least one dense and one MoE model complete quantize/save/reload/generate; -5. Qwen multimodal validation includes a real image prompt; -6. MTP preservation is checked from the saved safetensor index; -7. benchmark methodology and raw results are committed; -8. known limitations are repeated in release notes and model cards. +## Release gates -## External references +Before describing a release as production-ready: -These are research or comparison references, not bundled components: +1. CUDA compilation succeeds from clean wheel/sdist installs; +2. hardware validators pass on the advertised Ampere targets; +3. numerical bounds pass across the support matrix; +4. model round trips and generation checks pass; +5. benchmark JSON and profiler summaries are checked in; +6. quality regressions remain within documented thresholds; +7. README claims match those artifacts exactly. -- [GPTQ](https://arxiv.org/abs/2210.17323) -- [Marlin](https://github.com/IST-DASLab/marlin) -- [ExLlamaV2](https://github.com/turboderp-org/exllamav2) -- [NVIDIA Ampere architecture](https://developer.nvidia.com/blog/nvidia-ampere-architecture-in-depth/) -- [CUDA asynchronous copy](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copies) -- [QuIP](https://arxiv.org/abs/2307.13304) -- [QuaRot](https://arxiv.org/abs/2404.00456) -- [SpinQuant](https://arxiv.org/abs/2405.16406) -- [SmoothQuant](https://arxiv.org/abs/2211.10438) +Until then, GPTQ-Pro remains an experimental Ampere-focused quantization and +kernel-engineering fork. diff --git a/gptqmodel/nn_modules/qlinear/gptq_pro.py b/gptqmodel/nn_modules/qlinear/gptq_pro.py index ef07f3ef8..774e5a3c8 100644 --- a/gptqmodel/nn_modules/qlinear/gptq_pro.py +++ b/gptqmodel/nn_modules/qlinear/gptq_pro.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium -from typing import List, Optional, Tuple +from typing import Optional, Tuple import torch @@ -16,25 +16,21 @@ _validate_gptq_pro_device_support, apply_gptq_pro_linear, ensure_gptq_pro_loaded, - gptq_pro_qweight_to_b_packed, ) from ...utils.rocm import IS_ROCM -# GPTQ-Pro is the only inference kernel shipped by this fork. Priority 120 keeps -# it first in the local AUTO registry wherever its validation contract is met. -# This is a repository policy, not a benchmarked-performance claim: the kernel -# remains a Tensor-Core scaffold (one warp/CTA, no cp.async/ldmatrix/multi-stage -# pipeline, scalar INT4 decode, and no dedicated GEMV path). Unsupported -# configurations fail validation cleanly; because there is no bundled fallback -# kernel, callers must use a compatible checkpoint/device or another project. +# GPTQ-Pro is the only inference kernel shipped by this fork. AUTO dispatches +# small-M workloads to a coalesced decode kernel, aligned GEMMs to a four-warp +# double-buffered cp.async Tensor Core path, and unusual compatible shapes to a +# validator-backed general-shape fallback. All paths read the native int32 +# qweight tensor, avoiding a second persistent copy of quantized model weights. _GPTQ_PRO_AUTO_PRIORITY = 120 class GptqProQuantLinear(PackableQuantLinear): SUPPORTS_BACKENDS = [BACKEND.GPTQ_PRO] SUPPORTS_METHODS = [METHOD.GPTQ] - # Priority 120 keeps GPTQ-Pro first in the single-backend AUTO registry. SUPPORTS_FORMATS = { FORMAT.GPTQ: _GPTQ_PRO_AUTO_PRIORITY, FORMAT.GPTQ_V2: _GPTQ_PRO_AUTO_PRIORITY, @@ -152,10 +148,15 @@ def _validate( and (effective_group_size % 16) != 0 ): return False, NotImplementedError( - f"{cls} requires group_size to be a positive multiple of 16: actual group_size = `{effective_group_size}`" + f"{cls} requires group_size to be a positive multiple of 16: " + f"actual group_size = `{effective_group_size}`" ) return True, None + @property + def effective_group_size(self) -> int: + return self.in_features if self.group_size == -1 else self.group_size + def post_init(self): ensure_gptq_pro_loaded() @@ -164,33 +165,26 @@ def post_init(self): "GPTQ-Pro backend requires CUDA-resident packed weights before post_init()." ) + if not self.qweight.is_contiguous(): + self.qweight = self.qweight.contiguous() + if not self.scales.is_contiguous(): + self.scales = self.scales.contiguous() + expected_g_idx = ( torch.arange( self.in_features, device=self.g_idx.device, dtype=self.g_idx.dtype, ) - // self.group_size + // self.effective_group_size ) if not torch.equal(self.g_idx, expected_g_idx): raise ValueError( "GPTQ-Pro backend only supports sequential g_idx / desc_act=False checkpoints." ) - b_packed = gptq_pro_qweight_to_b_packed(self.qweight) - if "b_packed" not in self._buffers: - self.register_buffer("b_packed", b_packed, persistent=False) - else: - self.b_packed = b_packed - super().post_init() - def list_buffers(self) -> List: - buf = super().list_buffers() - if hasattr(self, "b_packed") and self.b_packed is not None: - buf.append(self.b_packed) - return buf - def forward(self, x: torch.Tensor): if x.shape[0] == 0: return torch.empty( @@ -201,7 +195,8 @@ def forward(self, x: torch.Tensor): x = x.reshape(-1, x.shape[-1]) if x.shape[-1] != self.in_features: raise ValueError( - f"GPTQ-Pro backend expected input dim {self.in_features}, got {x.shape[-1]}." + f"GPTQ-Pro backend expected input dim {self.in_features}, " + f"got {x.shape[-1]}." ) if x.dtype != torch.float16: @@ -209,9 +204,9 @@ def forward(self, x: torch.Tensor): out = apply_gptq_pro_linear( input=x.contiguous(), - b_packed=self.b_packed, + qweight=self.qweight, scales=self.scales, - group_size=self.group_size, + group_size=self.effective_group_size, ) if self.bias is not None: diff --git a/gptqmodel/utils/gptq_pro.py b/gptqmodel/utils/gptq_pro.py index c97293601..b9aebc9fa 100644 --- a/gptqmodel/utils/gptq_pro.py +++ b/gptqmodel/utils/gptq_pro.py @@ -22,15 +22,22 @@ log = setup_logger() +_GPTQ_PRO_EXTENSION_NAME = "gptqmodel_gptq_pro_kernels_v2" _GPTQ_PRO_LOCK = threading.Lock() _GPTQ_PRO_MODULE = None _GPTQ_PRO_INITIALISED = False _GPTQ_PRO_BUILD_PREPARED = False gptq_pro_import_exception: Optional[str] = None +_KERNEL_MODES = {"auto", "gemv", "ampere", "legacy"} + def _validate_gptq_pro_device_support() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 and not IS_ROCM + return ( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] >= 8 + and not IS_ROCM + ) def _gptq_pro_sources() -> tuple[Path, Path]: @@ -44,9 +51,11 @@ def _prepare_build_directory(verbose: bool) -> str: build_dir_env = os.getenv("GPTQMODEL_EXT_BUILD") if build_dir_env: - build_directory = Path(build_dir_env) / "gptqmodel_gptq_pro_kernels" + build_directory = Path(build_dir_env) / _GPTQ_PRO_EXTENSION_NAME else: - build_directory = Path(_get_build_directory("gptqmodel_gptq_pro_kernels", verbose=verbose)) + build_directory = Path( + _get_build_directory(_GPTQ_PRO_EXTENSION_NAME, verbose=verbose) + ) if not _GPTQ_PRO_BUILD_PREPARED and build_directory.exists(): shutil.rmtree(build_directory, ignore_errors=True) @@ -63,7 +72,7 @@ def _build_gptq_pro_extension(verbose: bool): build_directory = _prepare_build_directory(verbose=verbose) return load( - name="gptqmodel_gptq_pro_kernels", + name=_GPTQ_PRO_EXTENSION_NAME, sources=[str(source_cpp), str(source_cu)], extra_cflags=["-O3", "-std=c++17"], extra_cuda_cflags=[ @@ -74,15 +83,12 @@ def _build_gptq_pro_extension(verbose: bool): "-U__CUDA_NO_HALF_CONVERSIONS__", # Target all Ampere SM variants with native SASS cubins: # sm_80 — A100, A30, GA100 (data-centre Ampere) - # sm_86 — RTX 3090/3080/A6000, GA102/GA104/GA106 (consumer + pro Ampere) + # sm_86 — RTX 3090/3080/A6000, GA102/GA104/GA106 # sm_87 — Jetson Orin / embedded Ampere "-gencode arch=compute_80,code=sm_80", "-gencode arch=compute_86,code=sm_86", "-gencode arch=compute_87,code=sm_87", - # Embed sm_87 PTX as a forward-compatible fallback so the kernel can - # also be loaded on post-Ampere devices (Ada sm_89, Hopper sm_90, …) - # that pass the major >= 8 capability check. The CUDA driver will - # JIT-compile the PTX to native code on first use for those GPUs. + # Embed PTX for driver-side JIT on newer architectures. "-gencode arch=compute_87,code=compute_87", ], build_directory=build_directory, @@ -107,7 +113,7 @@ def ensure_gptq_pro_loaded(*, verbose: Optional[bool] = None): errors = [] try: - _GPTQ_PRO_MODULE = load_extension_module("gptqmodel_gptq_pro_kernels") + _GPTQ_PRO_MODULE = load_extension_module(_GPTQ_PRO_EXTENSION_NAME) gptq_pro_import_exception = None _GPTQ_PRO_INITIALISED = True return _GPTQ_PRO_MODULE @@ -116,7 +122,8 @@ def ensure_gptq_pro_loaded(*, verbose: Optional[bool] = None): if not _validate_gptq_pro_device_support(): gptq_pro_import_exception = ( - "GPTQ-Pro kernel requires Linux CUDA with compute capability >= 8.0 and does not support ROCm." + "GPTQ-Pro kernel requires Linux CUDA with compute capability >= 8.0 " + "and does not support ROCm." ) _GPTQ_PRO_INITIALISED = True raise ImportError(gptq_pro_import_exception) @@ -133,27 +140,54 @@ def ensure_gptq_pro_loaded(*, verbose: Optional[bool] = None): raise ImportError(gptq_pro_import_exception) from exc +def normalize_gptq_pro_kernel_mode(kernel_mode: Optional[str] = None) -> str: + if kernel_mode is None: + kernel_mode = os.getenv("GPTQMODEL_GPTQ_PRO_KERNEL", "auto") + normalized = str(kernel_mode).strip().lower() + if normalized not in _KERNEL_MODES: + raise ValueError( + "GPTQ-Pro kernel mode must be one of " + f"{sorted(_KERNEL_MODES)}, got `{kernel_mode}`." + ) + return normalized + + def gptq_pro_qweight_to_b_packed(qweight: torch.Tensor) -> torch.Tensor: + """Return the historical pair-packed byte view of GPTQ qweight. + + The optimized runtime no longer uses this conversion; kernels consume the + native int32 qweight tensor directly. The helper is retained for checkpoint + inspection and backwards-compatible tooling. + """ if qweight.dtype != torch.int32: raise ValueError(f"Expected int32 qweight tensor, got `{qweight.dtype}`.") if qweight.dim() != 2: - raise ValueError(f"Expected 2D qweight tensor, got shape `{tuple(qweight.shape)}`.") + raise ValueError( + f"Expected 2D qweight tensor, got shape `{tuple(qweight.shape)}`." + ) qweight = qweight.contiguous() - shifts = torch.arange(0, 32, 4, device=qweight.device, dtype=qweight.dtype).view(1, 8, 1) - unpacked = torch.bitwise_and(torch.bitwise_right_shift(qweight.unsqueeze(1), shifts), 0xF).to(torch.uint8) - unpacked = unpacked.reshape(-1, qweight.shape[1]) - return (unpacked[0::2] | (unpacked[1::2] << 4)).contiguous() + rows, columns = qweight.shape + bytes_by_word = qweight.view(torch.uint8).view(rows, columns, 4) + return bytes_by_word.permute(0, 2, 1).reshape(rows * 4, columns).contiguous() def apply_gptq_pro_linear( input: torch.Tensor, - b_packed: torch.Tensor, + qweight: torch.Tensor, scales: torch.Tensor, group_size: int, + kernel_mode: Optional[str] = None, ) -> torch.Tensor: module = ensure_gptq_pro_loaded() - return module.gptq_pro_gemm(input, b_packed, scales, int(group_size)) + mode = normalize_gptq_pro_kernel_mode(kernel_mode) + return module.gptq_pro_gemm( + input, + qweight, + scales, + int(group_size), + mode, + ) __all__ = [ @@ -162,4 +196,5 @@ def apply_gptq_pro_linear( "ensure_gptq_pro_loaded", "gptq_pro_import_exception", "gptq_pro_qweight_to_b_packed", + "normalize_gptq_pro_kernel_mode", ] diff --git a/gptqmodel_ext/gptq_pro/gptq_pro_kernel.cu b/gptqmodel_ext/gptq_pro/gptq_pro_kernel.cu index 103242e8f..2489e30b4 100644 --- a/gptqmodel_ext/gptq_pro/gptq_pro_kernel.cu +++ b/gptqmodel_ext/gptq_pro/gptq_pro_kernel.cu @@ -1,117 +1,145 @@ /* - * Standalone gptq_pro INT4 dequantized GEMM kernel + * GPTQ-Pro INT4 dequantized matrix multiplication for Ampere-class GPUs. * - * This file implements the current end-to-end functional scaffold: - * - one warp per CTA - * - symmetric INT4 weights with implicit zero-point 8 - * - explicit shared-memory staging for A / scales / B fragments - * - Tensor Core math via mma.sync with FP32 accumulation + * AUTO dispatch uses: + * 1. a coalesced FP32-accumulating GEMV kernel for M <= 4; + * 2. a four-warp, double-buffered cp.async Tensor Core kernel for aligned + * production shapes; + * 3. the original one-warp general-shape implementation as a fallback. * - * It is intentionally smaller and more conservative than the future Marlin-like - * multi-warp cp.async/ldmatrix pipeline discussed in Project.md/progress.md. - * The validator now covers both fragment helpers and the full kernel result. + * All paths consume GPTQ's native int32 qweight layout directly. No persistent + * pair-packed duplicate of the quantized weights is required. */ #include "gptq_pro_kernel.cuh" -struct __align__(16) GptqProSmem { - half A[GPTQ_PRO_M_PER_WARP * GPTQ_PRO_K_PER_WARP]; - half S[GPTQ_PRO_N_PER_WARP]; +struct __align__(16) GptqProLegacySmem { + half A[GPTQ_PRO_M_PER_WARP * GPTQ_PRO_K_PER_WARP]; + half S[GPTQ_PRO_N_PER_WARP]; uint32_t Bfrag[GPTQ_PRO_BFRAG_WORDS_PER_BUF]; }; -__device__ __forceinline__ -half zero_half() { +struct __align__(16) GptqProAmpereSmem { + // A is shared by the four N-warps in a CTA. Q and scales are warp-private. + half A[GPTQ_PRO_PIPE][GPTQ_PRO_M_PER_WARP * GPTQ_PRO_K_PER_WARP]; + half S[GPTQ_PRO_WARPS_PER_CTA][GPTQ_PRO_PIPE][GPTQ_PRO_N_PER_WARP]; + uint32_t + Q[GPTQ_PRO_WARPS_PER_CTA][GPTQ_PRO_PIPE] + [GPTQ_PRO_QWORD_ROWS_PER_K_TILE * GPTQ_PRO_N_PER_WARP]; +}; + +__device__ __forceinline__ half zero_half() { return __float2half(0.0f); } -__device__ __forceinline__ -uint8_t load_b_pair_byte(const uint8_t* __restrict__ B_packed, - int K, int N, - int k_even, int n_col) { +// --------------------------------------------------------------------------- +// General-shape fallback helpers +// --------------------------------------------------------------------------- +__device__ __forceinline__ uint8_t load_qweight_pair_byte( + const int32_t* __restrict__ Q, + int K, + int N, + int k_even, + int n_col) { if (n_col >= N || k_even >= K) { return 0x88u; } - const int packed_row = k_even >> 1; - uint8_t byte = B_packed[packed_row * N + n_col]; + const int qword_row = k_even >> 3; + const int byte_index = (k_even & 7) >> 1; + const uint32_t word = static_cast(Q[qword_row * N + n_col]); + uint8_t byte = static_cast((word >> (byte_index * 8)) & 0xFFu); if (k_even + 1 >= K) { byte = static_cast((byte & 0x0Fu) | 0x80u); } return byte; } -__device__ __forceinline__ -uint16_t pack_lane_bfrag(const uint8_t* __restrict__ B_packed, - int N, int K, - int k_base, int n_base, - int j, int lane) { +__device__ __forceinline__ uint16_t pack_lane_bfrag( + const int32_t* __restrict__ Q, + int N, + int K, + int k_base, + int n_base, + int j, + int lane) { const int group_id = lane >> 2; - const int tid4 = lane & 3; + const int thread_id = lane & 3; const int n_col = n_base + j * 8 + group_id; if (n_col >= N) { return 0x8888u; } - const int k01 = k_base + 2 * tid4; + const int k01 = k_base + 2 * thread_id; const int k89 = k01 + 8; - const uint8_t byte01 = load_b_pair_byte(B_packed, K, N, k01, n_col); - const uint8_t byte89 = load_b_pair_byte(B_packed, K, N, k89, n_col); - return static_cast(byte01) - | (static_cast(byte89) << 8); + const uint8_t byte01 = load_qweight_pair_byte(Q, K, N, k01, n_col); + const uint8_t byte89 = load_qweight_pair_byte(Q, K, N, k89, n_col); + return static_cast(byte01) | + (static_cast(byte89) << 8); } -__device__ __forceinline__ -void stage_a_tile(GptqProSmem* __restrict__ smem, - const half* __restrict__ A, - int M, int K, - int m_base, int k_base) { - for (int idx = threadIdx.x; idx < GPTQ_PRO_M_PER_WARP * GPTQ_PRO_K_PER_WARP; idx += blockDim.x) { - const int row = idx / GPTQ_PRO_K_PER_WARP; - const int col = idx % GPTQ_PRO_K_PER_WARP; +__device__ __forceinline__ void stage_a_tile_legacy( + GptqProLegacySmem* __restrict__ smem, + const half* __restrict__ A, + int M, + int K, + int m_base, + int k_base) { + for (int index = threadIdx.x; + index < GPTQ_PRO_M_PER_WARP * GPTQ_PRO_K_PER_WARP; + index += blockDim.x) { + const int row = index / GPTQ_PRO_K_PER_WARP; + const int col = index % GPTQ_PRO_K_PER_WARP; const int global_m = m_base + row; const int global_k = k_base + col; - smem->A[idx] = (global_m < M && global_k < K) - ? A[global_m * K + global_k] - : zero_half(); + smem->A[index] = + (global_m < M && global_k < K) + ? A[global_m * K + global_k] + : zero_half(); } } -__device__ __forceinline__ -void stage_scale_row(GptqProSmem* __restrict__ smem, - const half* __restrict__ S, - int N, - int k_base, - int group_size, - int n_base) { +__device__ __forceinline__ void stage_scale_row_legacy( + GptqProLegacySmem* __restrict__ smem, + const half* __restrict__ S, + int N, + int k_base, + int group_size, + int n_base) { const int group_idx = k_base / group_size; - for (int idx = threadIdx.x; idx < GPTQ_PRO_N_PER_WARP; idx += blockDim.x) { - const int global_n = n_base + idx; - smem->S[idx] = (global_n < N) - ? S[group_idx * N + global_n] - : zero_half(); + for (int index = threadIdx.x; index < GPTQ_PRO_N_PER_WARP; + index += blockDim.x) { + const int global_n = n_base + index; + smem->S[index] = + (global_n < N) ? S[group_idx * N + global_n] : zero_half(); } } -__device__ __forceinline__ -void stage_bfrag_tiles(GptqProSmem* __restrict__ smem, - const uint8_t* __restrict__ B_packed, - int N, int K, - int k_base, int n_base) { - for (int idx = threadIdx.x; idx < GPTQ_PRO_BFRAG_WORDS_PER_BUF; idx += blockDim.x) { - const int j = idx / GPTQ_PRO_BFRAG_WORDS_PER_TILE; - const int lane_pair = idx % GPTQ_PRO_BFRAG_WORDS_PER_TILE; +__device__ __forceinline__ void stage_bfrag_tiles_legacy( + GptqProLegacySmem* __restrict__ smem, + const int32_t* __restrict__ Q, + int N, + int K, + int k_base, + int n_base) { + for (int index = threadIdx.x; + index < GPTQ_PRO_BFRAG_WORDS_PER_BUF; + index += blockDim.x) { + const int j = index / GPTQ_PRO_BFRAG_WORDS_PER_TILE; + const int lane_pair = index % GPTQ_PRO_BFRAG_WORDS_PER_TILE; const int even_lane = lane_pair * 2; - const uint16_t even_p16 = pack_lane_bfrag(B_packed, N, K, k_base, n_base, j, even_lane); - const uint16_t odd_p16 = pack_lane_bfrag(B_packed, N, K, k_base, n_base, j, even_lane + 1); - smem->Bfrag[idx] = static_cast(even_p16) - | (static_cast(odd_p16) << 16); + const uint16_t even_p16 = + pack_lane_bfrag(Q, N, K, k_base, n_base, j, even_lane); + const uint16_t odd_p16 = + pack_lane_bfrag(Q, N, K, k_base, n_base, j, even_lane + 1); + smem->Bfrag[index] = static_cast(even_p16) | + (static_cast(odd_p16) << 16); } } -__device__ __forceinline__ -void do_mma_inner_loop(const GptqProSmem* __restrict__ smem, - float RC[GPTQ_PRO_J_TILES][4]) { +__device__ __forceinline__ void do_mma_legacy( + const GptqProLegacySmem* __restrict__ smem, + float RC[GPTQ_PRO_J_TILES][4]) { const int lane = threadIdx.x & (GPTQ_PRO_WARP_SIZE - 1); const int group_id = lane >> 2; const half zero_point = __float2half(8.0f); @@ -119,37 +147,228 @@ void do_mma_inner_loop(const GptqProSmem* __restrict__ smem, uint32_t RA[4]; load_a_fragment_rowmajor(smem->A, lane, RA); - #pragma unroll +#pragma unroll for (int j = 0; j < GPTQ_PRO_J_TILES; ++j) { const half scale = smem->S[j * 8 + group_id]; - const uint16_t packed_16 = fetch_bfrag_packed16(smem->Bfrag, 0, 0, j, lane); + const uint16_t packed_16 = + fetch_bfrag_packed16(smem->Bfrag, 0, 0, j, lane); + + uint32_t RB[2]; + decode_bfrag_to_rb(packed_16, scale, zero_point, RB); + mma_f32_m16n8k16(RA, RB, RC[j]); + } +} + +// --------------------------------------------------------------------------- +// Ampere cp.async pipeline helpers +// --------------------------------------------------------------------------- +__device__ __forceinline__ void prefetch_ampere_tile( + GptqProAmpereSmem* __restrict__ smem, + const half* __restrict__ A, + const int32_t* __restrict__ Q, + const half* __restrict__ S, + int M, + int N, + int K, + int group_size, + int m_base, + int n_base, + int k_base, + int buffer, + int warp_id, + int lane) { + // One warp cooperatively stages the common 16x16 A tile. Each lane moves a + // naturally aligned 16-byte segment; invalid M rows are zero-filled by + // cp.async's source-size operand. + if (warp_id == 0) { + const int row = lane >> 1; + const int segment = lane & 1; + const int global_m = m_base + row; + half* destination = + &smem->A[buffer][row * GPTQ_PRO_K_PER_WARP + segment * 8]; + const bool valid = global_m < M; + const half* source = + valid ? &A[global_m * K + k_base + segment * 8] : A; + cp_async_ca_16(destination, source, valid ? 16 : 0); + } + // Each warp stages two 64-column qweight rows. Every lane copies four + // adjacent int32 words (16 bytes), so all 128 threads issue coalesced loads. + const int qword_row_local = lane >> 4; + const int n_segment = lane & 15; + const int global_n = n_base + n_segment * 4; + uint32_t* q_destination = + &smem->Q[warp_id][buffer] + [qword_row_local * GPTQ_PRO_N_PER_WARP + n_segment * 4]; + const bool q_valid = global_n + 4 <= N; + const int32_t* q_source = + q_valid + ? &Q[((k_base >> 3) + qword_row_local) * N + global_n] + : Q; + cp_async_cg_16(q_destination, q_source, q_valid ? 16 : 0); + + // Eight lanes per warp stage the 64 FP16 scales in 16-byte vectors. + if (lane < 8) { + const int scale_n = n_base + lane * 8; + half* scale_destination = &smem->S[warp_id][buffer][lane * 8]; + const bool valid = scale_n + 8 <= N; + const int group_idx = k_base / group_size; + const half* scale_source = + valid ? &S[group_idx * N + scale_n] : S; + cp_async_ca_16(scale_destination, scale_source, valid ? 16 : 0); + } +} + +__device__ __forceinline__ void do_mma_ampere( + const half* __restrict__ smem_a, + const half* __restrict__ smem_s, + const uint32_t* __restrict__ smem_qweight, + int lane, + float RC[GPTQ_PRO_J_TILES][4]) { + const int group_id = lane >> 2; + const half zero_point = __float2half(8.0f); + + uint32_t RA[4]; + load_a_fragment_rowmajor(smem_a, lane, RA); + +#pragma unroll + for (int j = 0; j < GPTQ_PRO_J_TILES; ++j) { + const half scale = smem_s[j * 8 + group_id]; + const uint16_t packed_16 = + load_qweight_bfrag_packed16(smem_qweight, j, lane); uint32_t RB[2]; decode_bfrag_to_rb(packed_16, scale, zero_point, RB); mma_f32_m16n8k16(RA, RB, RC[j]); } } +__device__ __forceinline__ void store_mma_output( + half* __restrict__ C, + int M, + int N, + int m_base, + int n_base, + int lane, + float RC[GPTQ_PRO_J_TILES][4]) { + const int row_base = lane >> 2; + const int col_pair = 2 * (lane & 3); + +#pragma unroll + for (int j = 0; j < GPTQ_PRO_J_TILES; ++j) { + const int m0 = m_base + row_base; + const int m1 = m_base + row_base + 8; + const int n0 = n_base + j * 8 + col_pair; + const int n1 = n0 + 1; + + if (n1 < N) { + if (m0 < M) { + *reinterpret_cast(&C[m0 * N + n0]) = + __floats2half2_rn(RC[j][0], RC[j][1]); + } + if (m1 < M) { + *reinterpret_cast(&C[m1 * N + n0]) = + __floats2half2_rn(RC[j][2], RC[j][3]); + } + } else if (n0 < N) { + // The optimized path always has even N. This scalar tail keeps the + // general-shape fallback correct for standalone odd-N validation. + if (m0 < M) { + C[m0 * N + n0] = __float2half_rn(RC[j][0]); + } + if (m1 < M) { + C[m1 * N + n0] = __float2half_rn(RC[j][2]); + } + } + } +} + #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 -__global__ void gptq_pro_gemm_kernel( - const half* __restrict__ A, - const uint8_t* __restrict__ B_packed, - const half* __restrict__ S, - half* __restrict__ C, - int M, int N, int K, int group_size) -{ - extern __shared__ uint8_t raw_smem[]; - GptqProSmem* smem = reinterpret_cast(raw_smem); - const int warp_m = blockIdx.x; - const int warp_n = blockIdx.y; +// --------------------------------------------------------------------------- +// Dedicated small-M decode path +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(GPTQ_PRO_GEMV_THREADS, 4) +void gptq_pro_gemv_kernel( + const half* __restrict__ A, + const int32_t* __restrict__ Q, + const half* __restrict__ S, + half* __restrict__ C, + int M, + int N, + int K, + int group_size) { + const int n = blockIdx.x * blockDim.x + threadIdx.x; + const int m = blockIdx.y; + if (n >= N || m >= M) { + return; + } + + const half2* a_pairs = reinterpret_cast(&A[m * K]); + const int groups = (K + group_size - 1) / group_size; + float accumulator = 0.0f; + + for (int group = 0; group < groups; ++group) { + const int k_begin = group * group_size; + const int k_end_candidate = k_begin + group_size; + const int k_end = k_end_candidate < K ? k_end_candidate : K; + const half2 scale2 = __half2half2(S[group * N + n]); + const int qword_begin = k_begin >> 3; + const int qword_end = (k_end + 7) >> 3; + + for (int qword_row = qword_begin; qword_row < qword_end; + ++qword_row) { + const uint32_t word = + static_cast(Q[qword_row * N + n]); +#pragma unroll + for (int byte_index = 0; byte_index < 4; ++byte_index) { + const int pair_index = qword_row * 4 + byte_index; + const float2 activation = + __half22float2(a_pairs[pair_index]); + const uint8_t packed = static_cast( + (word >> (byte_index * 8)) & 0xFFu); + const int weight0 = static_cast(packed & 0x0Fu) - 8; + const int weight1 = + static_cast((packed >> 4) & 0x0Fu) - 8; + const half2 signed_weights = __floats2half2_rn( + static_cast(weight0), + static_cast(weight1)); + const float2 dequantized = + __half22float2(__hmul2(scale2, signed_weights)); + accumulator = + fmaf(activation.x, dequantized.x, accumulator); + accumulator = + fmaf(activation.y, dequantized.y, accumulator); + } + } + } + + C[m * N + n] = __float2half_rn(accumulator); +} + +// --------------------------------------------------------------------------- +// Four-warp, double-buffered Ampere Tensor Core path +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(GPTQ_PRO_THREADS_PER_CTA, 4) +void gptq_pro_gemm_kernel_ampere( + const half* __restrict__ A, + const int32_t* __restrict__ Q, + const half* __restrict__ S, + half* __restrict__ C, + int M, + int N, + int K, + int group_size) { + extern __shared__ uint8_t raw_smem[]; + auto* smem = reinterpret_cast(raw_smem); - const int m_base = warp_m * GPTQ_PRO_M_PER_WARP; - const int n_base = warp_n * GPTQ_PRO_N_PER_WARP; + const int warp_id = threadIdx.x / GPTQ_PRO_WARP_SIZE; const int lane = threadIdx.x & (GPTQ_PRO_WARP_SIZE - 1); + const int m_base = blockIdx.x * GPTQ_PRO_M_PER_WARP; + const int n_base = + blockIdx.y * GPTQ_PRO_N_PER_CTA + warp_id * GPTQ_PRO_N_PER_WARP; float RC[GPTQ_PRO_J_TILES][4]; - #pragma unroll +#pragma unroll for (int j = 0; j < GPTQ_PRO_J_TILES; ++j) { RC[j][0] = 0.0f; RC[j][1] = 0.0f; @@ -157,53 +376,136 @@ __global__ void gptq_pro_gemm_kernel( RC[j][3] = 0.0f; } - const int num_k_tiles = (K + GPTQ_PRO_K_PER_WARP - 1) / GPTQ_PRO_K_PER_WARP; - for (int t = 0; t < num_k_tiles; ++t) { - const int k_base = t * GPTQ_PRO_K_PER_WARP; + const int num_k_tiles = K / GPTQ_PRO_K_PER_WARP; + int read_buffer = 0; + prefetch_ampere_tile( + smem, + A, + Q, + S, + M, + N, + K, + group_size, + m_base, + n_base, + 0, + read_buffer, + warp_id, + lane); + cp_async_commit_group(); - stage_a_tile(smem, A, M, K, m_base, k_base); - stage_scale_row(smem, S, N, k_base, group_size, n_base); - stage_bfrag_tiles(smem, B_packed, N, K, k_base, n_base); - __syncthreads(); + for (int tile = 0; tile < num_k_tiles; ++tile) { + const int next_tile = tile + 1; + const int write_buffer = read_buffer ^ 1; + + if (next_tile < num_k_tiles) { + prefetch_ampere_tile( + smem, + A, + Q, + S, + M, + N, + K, + group_size, + m_base, + n_base, + next_tile * GPTQ_PRO_K_PER_WARP, + write_buffer, + warp_id, + lane); + cp_async_commit_group(); + cp_async_wait_group<1>(); + } else { + cp_async_wait_group<0>(); + } - do_mma_inner_loop(smem, RC); __syncthreads(); + do_mma_ampere( + smem->A[read_buffer], + smem->S[warp_id][read_buffer], + smem->Q[warp_id][read_buffer], + lane, + RC); + // No thread may overwrite this buffer until every warp has finished + // reading the common A tile. + __syncthreads(); + read_buffer = write_buffer; } - const int row_base = lane >> 2; - const int col_pair = 2 * (lane & 3); - #pragma unroll + store_mma_output(C, M, N, m_base, n_base, lane, RC); +} + +// --------------------------------------------------------------------------- +// Original general-shape fallback +// --------------------------------------------------------------------------- +__global__ void gptq_pro_gemm_kernel_legacy( + const half* __restrict__ A, + const int32_t* __restrict__ Q, + const half* __restrict__ S, + half* __restrict__ C, + int M, + int N, + int K, + int group_size) { + extern __shared__ uint8_t raw_smem[]; + auto* smem = reinterpret_cast(raw_smem); + + const int m_base = blockIdx.x * GPTQ_PRO_M_PER_WARP; + const int n_base = blockIdx.y * GPTQ_PRO_N_PER_WARP; + const int lane = threadIdx.x & (GPTQ_PRO_WARP_SIZE - 1); + + float RC[GPTQ_PRO_J_TILES][4]; +#pragma unroll for (int j = 0; j < GPTQ_PRO_J_TILES; ++j) { - const int m0 = m_base + row_base; - const int m1 = m_base + row_base + 8; - const int n0 = n_base + j * 8 + col_pair + 0; - const int n1 = n_base + j * 8 + col_pair + 1; + RC[j][0] = 0.0f; + RC[j][1] = 0.0f; + RC[j][2] = 0.0f; + RC[j][3] = 0.0f; + } - if (m0 < M) { - if (n0 < N) C[m0 * N + n0] = __float2half_rn(RC[j][0]); - if (n1 < N) C[m0 * N + n1] = __float2half_rn(RC[j][1]); - } - if (m1 < M) { - if (n0 < N) C[m1 * N + n0] = __float2half_rn(RC[j][2]); - if (n1 < N) C[m1 * N + n1] = __float2half_rn(RC[j][3]); - } + const int num_k_tiles = + (K + GPTQ_PRO_K_PER_WARP - 1) / GPTQ_PRO_K_PER_WARP; + for (int tile = 0; tile < num_k_tiles; ++tile) { + const int k_base = tile * GPTQ_PRO_K_PER_WARP; + stage_a_tile_legacy(smem, A, M, K, m_base, k_base); + stage_scale_row_legacy( + smem, S, N, k_base, group_size, n_base); + stage_bfrag_tiles_legacy(smem, Q, N, K, k_base, n_base); + __syncthreads(); + do_mma_legacy(smem, RC); + __syncthreads(); } + + store_mma_output(C, M, N, m_base, n_base, lane, RC); } -#else // sm80 stub -__global__ void gptq_pro_gemm_kernel( - const half*, const uint8_t*, const half*, half*, int, int, int, int) {} -#define GPTQ_PRO_SM80_STUB 1 +#else + +__global__ void gptq_pro_gemv_kernel( + const half*, const int32_t*, const half*, half*, int, int, int, int) {} +__global__ void gptq_pro_gemm_kernel_ampere( + const half*, const int32_t*, const half*, half*, int, int, int, int) {} +__global__ void gptq_pro_gemm_kernel_legacy( + const half*, const int32_t*, const half*, half*, int, int, int, int) {} + #endif cudaError_t gptq_pro_gemm( - const half* A, - const uint8_t* B_packed, - const half* S, - half* C, - int M, int N, int K, int group_size, - cudaStream_t stream) -{ + const half* A, + const int32_t* Q, + const half* S, + half* C, + int M, + int N, + int K, + int group_size, + cudaStream_t stream, + int kernel_mode) { + if (M <= 0 || N <= 0 || K <= 0) { + return cudaSuccess; + } if (group_size <= 0) { group_size = K; } @@ -211,14 +513,61 @@ cudaError_t gptq_pro_gemm( return cudaErrorInvalidValue; } - dim3 grid( + const bool gemv_compatible = (K % GPTQ_PRO_QWORD_VALUES_PER_WORD) == 0; + const bool ampere_compatible = + (K % GPTQ_PRO_K_PER_WARP) == 0 && (N % 16) == 0; + + int selected_mode = kernel_mode; + if (selected_mode == GPTQ_PRO_KERNEL_AUTO) { + if (M <= GPTQ_PRO_GEMV_MAX_M && gemv_compatible) { + selected_mode = GPTQ_PRO_KERNEL_GEMV; + } else if (ampere_compatible) { + selected_mode = GPTQ_PRO_KERNEL_AMPERE; + } else { + selected_mode = GPTQ_PRO_KERNEL_LEGACY; + } + } + + if (selected_mode == GPTQ_PRO_KERNEL_GEMV) { + if (!gemv_compatible) { + return cudaErrorInvalidValue; + } + const dim3 grid( + (N + GPTQ_PRO_GEMV_THREADS - 1) / GPTQ_PRO_GEMV_THREADS, + M, + 1); + const dim3 block(GPTQ_PRO_GEMV_THREADS, 1, 1); + gptq_pro_gemv_kernel<<>>( + A, Q, S, C, M, N, K, group_size); + return cudaGetLastError(); + } + + if (selected_mode == GPTQ_PRO_KERNEL_AMPERE) { + if (!ampere_compatible) { + return cudaErrorInvalidValue; + } + const dim3 grid( + (M + GPTQ_PRO_M_PER_WARP - 1) / GPTQ_PRO_M_PER_WARP, + (N + GPTQ_PRO_N_PER_CTA - 1) / GPTQ_PRO_N_PER_CTA, + 1); + const dim3 block(GPTQ_PRO_THREADS_PER_CTA, 1, 1); + const size_t smem_bytes = sizeof(GptqProAmpereSmem); + gptq_pro_gemm_kernel_ampere<<>>( + A, Q, S, C, M, N, K, group_size); + return cudaGetLastError(); + } + + if (selected_mode != GPTQ_PRO_KERNEL_LEGACY) { + return cudaErrorInvalidValue; + } + + const dim3 grid( (M + GPTQ_PRO_M_PER_WARP - 1) / GPTQ_PRO_M_PER_WARP, (N + GPTQ_PRO_N_PER_WARP - 1) / GPTQ_PRO_N_PER_WARP, 1); - dim3 block(GPTQ_PRO_WARP_SIZE, 1, 1); - const size_t smem_bytes = sizeof(GptqProSmem); - - gptq_pro_gemm_kernel<<>>( - A, B_packed, S, C, M, N, K, group_size); + const dim3 block(GPTQ_PRO_WARP_SIZE, 1, 1); + const size_t smem_bytes = sizeof(GptqProLegacySmem); + gptq_pro_gemm_kernel_legacy<<>>( + A, Q, S, C, M, N, K, group_size); return cudaGetLastError(); } diff --git a/gptqmodel_ext/gptq_pro/gptq_pro_kernel.cuh b/gptqmodel_ext/gptq_pro/gptq_pro_kernel.cuh index 5039401f6..197f2b2cc 100644 --- a/gptqmodel_ext/gptq_pro/gptq_pro_kernel.cuh +++ b/gptqmodel_ext/gptq_pro/gptq_pro_kernel.cuh @@ -1,15 +1,10 @@ /* - * Standalone gptq_pro Tensor Core scaffold for Ampere. + * GPTQ-Pro Ampere INT4 kernel primitives. * - * Current scope: - * - one warp per CTA - * - symmetric INT4 weights packed as unsigned nibbles with implicit zero-point 8 - * - explicit shared-memory staging for the A tile, per-column scales, and B fragments - * - FP32 accumulation via mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 - * - * This kernel is validator-backed and end-to-end functional, but it is still the - * compact standalone scaffold referenced in README/progress.md rather than the - * future multi-warp cp.async/ldmatrix pipeline discussed in Project.md. + * The runtime consumes GPTQ's native int32 qweight layout directly. Every word + * stores eight K-axis nibbles for one output column. A 16-wide K tile therefore + * needs two coalesced qweight rows—the same 512 global bytes as the pair-packed + * representation, without a second persistent weight buffer. */ #pragma once @@ -20,123 +15,231 @@ #include // --------------------------------------------------------------------------- -// Tile dimensions +// Tile and dispatch constants // --------------------------------------------------------------------------- -static constexpr int GPTQ_PRO_PIPE = 1; // current scaffold uses one staged tile -static constexpr int GPTQ_PRO_KS_TILES = 1; // one mma.sync k16 step per outer K tile -static constexpr int GPTQ_PRO_J_TILES = 8; // 8 x n8 slices -> 64 output cols / warp +static constexpr int GPTQ_PRO_PIPE = 2; +static constexpr int GPTQ_PRO_KS_TILES = 1; +static constexpr int GPTQ_PRO_J_TILES = 8; static constexpr int GPTQ_PRO_WARP_SIZE = 32; static constexpr int GPTQ_PRO_M_PER_WARP = 16; static constexpr int GPTQ_PRO_N_PER_WARP = GPTQ_PRO_J_TILES * 8; // 64 static constexpr int GPTQ_PRO_K_PER_WARP = GPTQ_PRO_KS_TILES * 16; // 16 -// Number of uint32_t words per (ks,j) tile in Bfrag smem (lane-pair packing). -static constexpr int GPTQ_PRO_BFRAG_WORDS_PER_TILE = - GPTQ_PRO_WARP_SIZE / 2; // 16 words +static constexpr int GPTQ_PRO_WARPS_PER_CTA = 4; +static constexpr int GPTQ_PRO_THREADS_PER_CTA = + GPTQ_PRO_WARPS_PER_CTA * GPTQ_PRO_WARP_SIZE; +static constexpr int GPTQ_PRO_N_PER_CTA = + GPTQ_PRO_WARPS_PER_CTA * GPTQ_PRO_N_PER_WARP; // 256 + +static constexpr int GPTQ_PRO_QWORD_ROWS_PER_K_TILE = + GPTQ_PRO_K_PER_WARP / 8; // 2 int32 rows +static constexpr int GPTQ_PRO_QWORD_VALUES_PER_WORD = 8; +static constexpr int GPTQ_PRO_QWORD_BYTES_PER_WARP_TILE = + GPTQ_PRO_QWORD_ROWS_PER_K_TILE * GPTQ_PRO_N_PER_WARP * + sizeof(uint32_t); // 512 -// Total uint32_t words for all (ks,j) tiles in one smem buffer. +static constexpr int GPTQ_PRO_GEMV_THREADS = 128; +static constexpr int GPTQ_PRO_GEMV_MAX_M = 4; + +// Number of uint32_t words per legacy (ks,j) B-fragment shared tile. +static constexpr int GPTQ_PRO_BFRAG_WORDS_PER_TILE = + GPTQ_PRO_WARP_SIZE / 2; static constexpr int GPTQ_PRO_BFRAG_WORDS_PER_BUF = - GPTQ_PRO_KS_TILES * GPTQ_PRO_J_TILES * GPTQ_PRO_BFRAG_WORDS_PER_TILE; + GPTQ_PRO_KS_TILES * GPTQ_PRO_J_TILES * + GPTQ_PRO_BFRAG_WORDS_PER_TILE; + +enum GptqProKernelMode : int { + GPTQ_PRO_KERNEL_AUTO = 0, + GPTQ_PRO_KERNEL_GEMV = 1, + GPTQ_PRO_KERNEL_AMPERE = 2, + GPTQ_PRO_KERNEL_LEGACY = 3, +}; // --------------------------------------------------------------------------- // Helper types // --------------------------------------------------------------------------- union Half2Reg { - half2 h2; + half2 h2; uint32_t u32; uint16_t u16[2]; }; -__device__ __forceinline__ -uint32_t pack_half2_reg(half lo, half hi) { +__device__ __forceinline__ uint32_t pack_half2_reg(half lo, half hi) { Half2Reg reg; reg.h2 = __halves2half2(lo, hi); return reg.u32; } // --------------------------------------------------------------------------- -// Shared-memory layout helpers for the B fragment. -// -// The current scaffold stages one k16 slice at a time, so only ks=0 is used in -// practice, but the helper keeps the (buf, ks, j, lane) contract so the decode -// validator continues to exercise the exact same lane-pair packing logic. +// Ampere asynchronous copy helpers // --------------------------------------------------------------------------- -__device__ __forceinline__ -uint32_t bfrag_smem_addr(const uint32_t* __restrict__ smem_bfrag_base, - int buf, int ks, int j, int lane) { - const int tile_idx = ks * GPTQ_PRO_J_TILES + j; - const int buf_words = GPTQ_PRO_BFRAG_WORDS_PER_BUF; - const int word_idx = buf * buf_words - + tile_idx * GPTQ_PRO_BFRAG_WORDS_PER_TILE - + (lane >> 1); - return static_cast(__cvta_generic_to_shared(smem_bfrag_base)) - + static_cast(word_idx * sizeof(uint32_t)); +__device__ __forceinline__ void cp_async_ca_16( + void* smem_ptr, const void* global_ptr, int source_bytes = 16) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + const uint32_t smem = + static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "cp.async.ca.shared.global [%0], [%1], 16, %2;\n" + : + : "r"(smem), "l"(global_ptr), "r"(source_bytes)); +#else + (void)smem_ptr; + (void)global_ptr; + (void)source_bytes; +#endif } -__device__ __forceinline__ -uint32_t ld_shared_u32(uint32_t smem_addr) { - uint32_t val; - asm volatile("ld.shared.u32 %0, [%1];" : "=r"(val) : "r"(smem_addr)); - return val; +__device__ __forceinline__ void cp_async_cg_16( + void* smem_ptr, const void* global_ptr, int source_bytes = 16) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + const uint32_t smem = + static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "cp.async.cg.shared.global [%0], [%1], 16, %2;\n" + : + : "r"(smem), "l"(global_ptr), "r"(source_bytes)); +#else + (void)smem_ptr; + (void)global_ptr; + (void)source_bytes; +#endif } -__device__ __forceinline__ -uint16_t fetch_bfrag_packed16(const uint32_t* __restrict__ smem_bfrag, - int buf, int ks, int j, int lane) { - const uint32_t addr = bfrag_smem_addr(smem_bfrag, buf, ks, j, lane); - const uint32_t word = ld_shared_u32(addr); - return static_cast((lane & 1) ? (word >> 16) : (word & 0xFFFFu)); +__device__ __forceinline__ void cp_async_commit_group() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.commit_group;\n" ::); +#endif } +template +__device__ __forceinline__ void cp_async_wait_group() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group %0;\n" : : "n"(PendingGroups)); +#endif +} + +// --------------------------------------------------------------------------- +// LOP3-assisted INT4 decode // --------------------------------------------------------------------------- -// INT4 nibble decode -> FP16 (scale * (w - 8)). -// -// The standalone scaffold currently models the symmetric GPTQ-style runtime -// where 4-bit weights are stored as unsigned nibbles with an implicit zero-point -// of 8 and a per-group/per-column FP16 scale. -// --------------------------------------------------------------------------- -__device__ __forceinline__ -void decode_bfrag_to_rb(uint16_t packed_16, - half scale, half zero_point, - uint32_t (&RB)[2]) { - const uint32_t p = static_cast(packed_16); - const uint32_t w0 = (p >> 0) & 0xFu; - const uint32_t w1 = (p >> 4) & 0xFu; - const uint32_t w2 = (p >> 8) & 0xFu; - const uint32_t w3 = (p >> 12) & 0xFu; - - const half2 vals01 = __halves2half2(__int2half_rn(static_cast(w0)), - __int2half_rn(static_cast(w1))); - const half2 vals23 = __halves2half2(__int2half_rn(static_cast(w2)), - __int2half_rn(static_cast(w3))); - const half2 zp_h2 = __halves2half2(zero_point, zero_point); - const half2 sc_h2 = __halves2half2(scale, scale); - - Half2Reg rb0, rb1; - rb0.h2 = __hmul2(sc_h2, __hsub2(vals01, zp_h2)); - rb1.h2 = __hmul2(sc_h2, __hsub2(vals23, zp_h2)); - RB[0] = rb0.u32; - RB[1] = rb1.u32; +template +__device__ __forceinline__ uint32_t lop3_u32( + uint32_t a, uint32_t b, uint32_t c) { + uint32_t result; + asm volatile( + "lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(result) + : "r"(a), "r"(b), "r"(c), "n"(Lut)); + return result; +} + +// Convert four packed nibbles to the two FP16 B-fragment registers expected by +// mma.sync. The bit trick follows the Apache-2.0 Marlin/FasterTransformer-style +// conversion and fuses the symmetric -8 offset before scale multiplication. +__device__ __forceinline__ void decode_bfrag_to_rb( + uint16_t packed_16, half scale, half zero_point, uint32_t (&RB)[2]) { + (void)zero_point; // GPTQ-Pro's runtime contract always uses zero-point 8. + + const uint32_t p = static_cast(packed_16); + // Reorder [w0,w1,w2,w3] so LOP3 creates half2(w0,w1) and + // half2(w2,w3), matching the PTX B-fragment register order. + const uint32_t q = + (p & 0x0000000Fu) | ((p >> 4) & 0x000000F0u) | + ((p & 0x000000F0u) << 12) | ((p & 0x0000F000u) << 8); + + constexpr uint32_t LO = 0x000f000f; + constexpr uint32_t HI = 0x00f000f0; + constexpr uint32_t EX = 0x64006400; + constexpr uint32_t SUB = 0x64086408; + constexpr uint32_t MUL = 0x2c002c00; + constexpr uint32_t ADD = 0xd480d480; + + constexpr int AND_OR_LUT = (0xf0 & 0xcc) | 0xaa; + Half2Reg lo_reg; + Half2Reg hi_reg; + Half2Reg sub_reg; + Half2Reg mul_reg; + Half2Reg add_reg; + Half2Reg out0; + Half2Reg out1; + + lo_reg.u32 = lop3_u32(q, LO, EX); + hi_reg.u32 = lop3_u32(q, HI, EX); + sub_reg.u32 = SUB; + mul_reg.u32 = MUL; + add_reg.u32 = ADD; + + const half2 scale2 = __half2half2(scale); + out0.h2 = __hmul2(scale2, __hsub2(lo_reg.h2, sub_reg.h2)); + out1.h2 = __hmul2( + scale2, __hfma2(hi_reg.h2, mul_reg.h2, add_reg.h2)); + RB[0] = out0.u32; + RB[1] = out1.u32; +} + +// --------------------------------------------------------------------------- +// Native qweight shared-memory layout used by the optimized pipeline +// --------------------------------------------------------------------------- +__device__ __forceinline__ uint16_t load_qweight_bfrag_packed16( + const uint32_t* __restrict__ smem_qweight, int j, int lane) { + const int group_id = lane >> 2; + const int thread_id = lane & 3; + const int n_local = j * 8 + group_id; + const uint32_t word0 = smem_qweight[n_local]; + const uint32_t word1 = + smem_qweight[GPTQ_PRO_N_PER_WARP + n_local]; + const uint8_t byte01 = static_cast( + (word0 >> (thread_id * 8)) & 0xFFu); + const uint8_t byte89 = static_cast( + (word1 >> (thread_id * 8)) & 0xFFu); + return static_cast(byte01) | + (static_cast(byte89) << 8); +} + +// --------------------------------------------------------------------------- +// Legacy B-fragment helpers retained for the general-shape fallback and +// standalone fragment validator. +// --------------------------------------------------------------------------- +__device__ __forceinline__ uint32_t bfrag_smem_addr( + const uint32_t* __restrict__ smem_bfrag_base, + int buf, + int ks, + int j, + int lane) { + const int tile_idx = ks * GPTQ_PRO_J_TILES + j; + const int buf_words = GPTQ_PRO_BFRAG_WORDS_PER_BUF; + const int word_idx = + buf * buf_words + tile_idx * GPTQ_PRO_BFRAG_WORDS_PER_TILE + + (lane >> 1); + return static_cast( + __cvta_generic_to_shared(smem_bfrag_base)) + + static_cast(word_idx * sizeof(uint32_t)); +} + +__device__ __forceinline__ uint32_t ld_shared_u32(uint32_t smem_addr) { + uint32_t value; + asm volatile("ld.shared.u32 %0, [%1];" : "=r"(value) : "r"(smem_addr)); + return value; +} + +__device__ __forceinline__ uint16_t fetch_bfrag_packed16( + const uint32_t* __restrict__ smem_bfrag, + int buf, + int ks, + int j, + int lane) { + const uint32_t address = + bfrag_smem_addr(smem_bfrag, buf, ks, j, lane); + const uint32_t word = ld_shared_u32(address); + return static_cast( + (lane & 1) ? (word >> 16) : (word & 0xFFFFu)); } // --------------------------------------------------------------------------- // A-fragment packing for mma.sync.aligned.m16n8k16.row.col -// -// This loader follows the same lane ownership used in the validator's scalar -// reference: -// groupID = lane >> 2 -// tid4 = lane & 3 -// rows = {groupID, groupID + 8} -// cols = {2*tid4, 2*tid4 + 1, 2*tid4 + 8, 2*tid4 + 9} -// -// Using explicit register packing avoids the invalid/misaligned ldmatrix path -// that the earlier scaffold emitted for this compact one-warp layout. -// --------------------------------------------------------------------------- -__device__ __forceinline__ -void load_a_fragment_rowmajor(const half* __restrict__ smem_a, - int lane, - uint32_t (&RA)[4]) { +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_a_fragment_rowmajor( + const half* __restrict__ smem_a, int lane, uint32_t (&RA)[4]) { const int group_id = lane >> 2; const int thread_id = lane & 3; const int a_col_lo = 2 * thread_id; @@ -157,12 +260,10 @@ void load_a_fragment_rowmajor(const half* __restrict__ smem_a, } // --------------------------------------------------------------------------- -// FP32 accumulating MMA: RC += RA x RB +// FP32 accumulating Tensor Core MMA // --------------------------------------------------------------------------- -__device__ __forceinline__ -void mma_f32_m16n8k16(const uint32_t RA[4], - const uint32_t RB[2], - float RC[4]) { +__device__ __forceinline__ void mma_f32_m16n8k16( + const uint32_t RA[4], const uint32_t RB[2], float RC[4]) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "{%0, %1, %2, %3}, " @@ -170,14 +271,22 @@ void mma_f32_m16n8k16(const uint32_t RA[4], "{%8, %9}, " "{%0, %1, %2, %3};\n" : "+f"(RC[0]), "+f"(RC[1]), "+f"(RC[2]), "+f"(RC[3]) - : "r"(RA[0]), "r"(RA[1]), "r"(RA[2]), "r"(RA[3]), - "r"(RB[0]), "r"(RB[1])); + : "r"(RA[0]), + "r"(RA[1]), + "r"(RA[2]), + "r"(RA[3]), + "r"(RB[0]), + "r"(RB[1])); } cudaError_t gptq_pro_gemm( - const half* A, - const uint8_t* B_packed, - const half* S, - half* C, - int M, int N, int K, int group_size, - cudaStream_t stream); + const half* A, + const int32_t* Q, + const half* S, + half* C, + int M, + int N, + int K, + int group_size, + cudaStream_t stream, + int kernel_mode = GPTQ_PRO_KERNEL_AUTO); diff --git a/gptqmodel_ext/gptq_pro/gptq_pro_torch.cpp b/gptqmodel_ext/gptq_pro/gptq_pro_torch.cpp index e692195f1..2d129caa2 100644 --- a/gptqmodel_ext/gptq_pro/gptq_pro_torch.cpp +++ b/gptqmodel_ext/gptq_pro/gptq_pro_torch.cpp @@ -1,5 +1,5 @@ /* - * PyTorch extension entrypoint for the standalone gptq_pro CUDA scaffold. + * PyTorch extension entrypoint for GPTQ-Pro's Ampere INT4 kernels. */ #include @@ -9,77 +9,150 @@ #include #include +#include +#include +#include + cudaError_t gptq_pro_gemm( - const half* A, - const uint8_t* B_packed, - const half* S, - half* C, - int M, int N, int K, int group_size, - cudaStream_t stream); + const half* A, + const int32_t* Q, + const half* S, + half* C, + int M, + int N, + int K, + int group_size, + cudaStream_t stream, + int kernel_mode); namespace { -void check_inputs(const torch::Tensor& a, - const torch::Tensor& b_packed, - const torch::Tensor& scales, - int64_t group_size) { +constexpr int kKernelAuto = 0; +constexpr int kKernelGemv = 1; +constexpr int kKernelAmpere = 2; +constexpr int kKernelLegacy = 3; + +int parse_kernel_mode(const std::string& kernel_mode) { + if (kernel_mode == "auto") { + return kKernelAuto; + } + if (kernel_mode == "gemv") { + return kKernelGemv; + } + if (kernel_mode == "ampere") { + return kKernelAmpere; + } + if (kernel_mode == "legacy") { + return kKernelLegacy; + } + TORCH_CHECK( + false, + "gptq_pro_gemm: kernel_mode must be one of auto, gemv, ampere, or legacy; got `", + kernel_mode, + "`."); + return kKernelAuto; +} + +void check_inputs( + const torch::Tensor& a, + const torch::Tensor& qweight, + const torch::Tensor& scales, + int64_t group_size) { TORCH_CHECK(a.is_cuda(), "gptq_pro_gemm: activations must be CUDA tensors."); - TORCH_CHECK(b_packed.is_cuda(), "gptq_pro_gemm: packed weights must be CUDA tensors."); + TORCH_CHECK( + qweight.is_cuda(), "gptq_pro_gemm: qweight must be a CUDA tensor."); TORCH_CHECK(scales.is_cuda(), "gptq_pro_gemm: scales must be CUDA tensors."); - TORCH_CHECK(a.scalar_type() == torch::kFloat16, "gptq_pro_gemm: activations must be float16."); - TORCH_CHECK(b_packed.scalar_type() == torch::kUInt8, "gptq_pro_gemm: packed weights must be uint8."); - TORCH_CHECK(scales.scalar_type() == torch::kFloat16, "gptq_pro_gemm: scales must be float16."); + TORCH_CHECK( + a.scalar_type() == torch::kFloat16, + "gptq_pro_gemm: activations must be float16."); + TORCH_CHECK( + qweight.scalar_type() == torch::kInt32, + "gptq_pro_gemm: qweight must use GPTQ int32 packing."); + TORCH_CHECK( + scales.scalar_type() == torch::kFloat16, + "gptq_pro_gemm: scales must be float16."); TORCH_CHECK(a.dim() == 2, "gptq_pro_gemm: activations must be 2D [M, K]."); - TORCH_CHECK(b_packed.dim() == 2, "gptq_pro_gemm: packed weights must be 2D [(K+1)/2, N]."); - TORCH_CHECK(scales.dim() == 2, "gptq_pro_gemm: scales must be 2D [groups, N]."); + TORCH_CHECK( + qweight.dim() == 2, + "gptq_pro_gemm: qweight must be 2D [ceil(K / 8), N]."); + TORCH_CHECK( + scales.dim() == 2, + "gptq_pro_gemm: scales must be 2D [groups, N]."); TORCH_CHECK(a.is_contiguous(), "gptq_pro_gemm: activations must be contiguous."); - TORCH_CHECK(b_packed.is_contiguous(), "gptq_pro_gemm: packed weights must be contiguous."); + TORCH_CHECK(qweight.is_contiguous(), "gptq_pro_gemm: qweight must be contiguous."); TORCH_CHECK(scales.is_contiguous(), "gptq_pro_gemm: scales must be contiguous."); - TORCH_CHECK(a.device() == b_packed.device() && a.device() == scales.device(), - "gptq_pro_gemm: all tensors must live on the same CUDA device."); - TORCH_CHECK(group_size > 0 && (group_size % 16) == 0, - "gptq_pro_gemm: group_size must be a positive multiple of 16."); + TORCH_CHECK( + a.device() == qweight.device() && a.device() == scales.device(), + "gptq_pro_gemm: all tensors must live on the same CUDA device."); + TORCH_CHECK( + group_size > 0 && (group_size % 16) == 0, + "gptq_pro_gemm: group_size must be a positive multiple of 16."); + const auto m = a.size(0); const auto k = a.size(1); - const auto packed_rows = b_packed.size(0); - TORCH_CHECK(packed_rows == (k + 1) / 2, - "gptq_pro_gemm: packed weights shape does not match activation K dimension."); - TORCH_CHECK(scales.size(1) == b_packed.size(1), - "gptq_pro_gemm: scales second dimension must equal packed weight N dimension."); - TORCH_CHECK(scales.size(0) == (k + group_size - 1) / group_size, - "gptq_pro_gemm: scales first dimension must equal ceil(K / group_size)."); + const auto n = qweight.size(1); + const auto int_max = static_cast(std::numeric_limits::max()); + TORCH_CHECK( + m <= int_max && n <= int_max && k <= int_max && group_size <= int_max, + "gptq_pro_gemm: dimensions exceed the CUDA launcher's 32-bit range."); + + TORCH_CHECK( + qweight.size(0) == (k + 7) / 8, + "gptq_pro_gemm: qweight first dimension must equal ceil(K / 8)."); + TORCH_CHECK( + scales.size(1) == n, + "gptq_pro_gemm: scales second dimension must equal qweight N dimension."); + TORCH_CHECK( + scales.size(0) == (k + group_size - 1) / group_size, + "gptq_pro_gemm: scales first dimension must equal ceil(K / group_size)."); } } // namespace -torch::Tensor gptq_pro_gemm_torch(torch::Tensor a, - torch::Tensor b_packed, - torch::Tensor scales, - int64_t group_size) { - check_inputs(a, b_packed, scales, group_size); +torch::Tensor gptq_pro_gemm_torch( + torch::Tensor a, + torch::Tensor qweight, + torch::Tensor scales, + int64_t group_size, + const std::string& kernel_mode) { + check_inputs(a, qweight, scales, group_size); + const int mode = parse_kernel_mode(kernel_mode); - auto out = torch::empty({a.size(0), b_packed.size(1)}, a.options()); + auto out = torch::empty({a.size(0), qweight.size(1)}, a.options()); const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.device().index()); + const cudaStream_t stream = + at::cuda::getCurrentCUDAStream(a.device().index()); const auto status = gptq_pro_gemm( reinterpret_cast(a.data_ptr()), - b_packed.data_ptr(), + qweight.data_ptr(), reinterpret_cast(scales.data_ptr()), reinterpret_cast(out.data_ptr()), static_cast(a.size(0)), - static_cast(b_packed.size(1)), + static_cast(qweight.size(1)), static_cast(a.size(1)), static_cast(group_size), - stream); + stream, + mode); - TORCH_CHECK(status == cudaSuccess, - "gptq_pro_gemm launch failed: ", - cudaGetErrorString(status)); + TORCH_CHECK( + status == cudaSuccess, + "gptq_pro_gemm launch failed in `", + kernel_mode, + "` mode: ", + cudaGetErrorString(status)); return out; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("gptq_pro_gemm", &gptq_pro_gemm_torch, "GPTQ-Pro FP16xINT4 matmul."); + m.def( + "gptq_pro_gemm", + &gptq_pro_gemm_torch, + "GPTQ-Pro FP16 x native GPTQ INT4 matmul.", + pybind11::arg("activations"), + pybind11::arg("qweight"), + pybind11::arg("scales"), + pybind11::arg("group_size"), + pybind11::arg("kernel_mode") = "auto"); } diff --git a/gptqmodel_ext/gptq_pro/gptq_pro_validate.cu b/gptqmodel_ext/gptq_pro/gptq_pro_validate.cu index 25b26e8a3..855578363 100644 --- a/gptqmodel_ext/gptq_pro/gptq_pro_validate.cu +++ b/gptqmodel_ext/gptq_pro/gptq_pro_validate.cu @@ -1,595 +1,378 @@ /* - * GPTQ-Pro kernel validation harness + * Standalone GPTQ-Pro CUDA validation harness. * - * Implements the two validation milestones from the design TODO list: + * Build without PyTorch: + * nvcc -O3 -std=c++17 -arch=sm_80 \ + * gptq_pro_validate.cu gptq_pro_kernel.cu -o gptq_pro_validate * - * TODO 1 — Validate decode-only against a scalar host/device reference for - * one warp fragment. - * - * TODO 2 — Validate one full ks/j MMA step against a reference with FP32 - * accumulation semantics and FP16 inputs. - * - * Each milestone produces a per-thread pass/fail flag stored in a result - * buffer that can be inspected from the host. Both kernels target sm80+. - * - * Build (standalone, no PyTorch): - * nvcc -arch=sm_80 -std=c++17 gptq_pro_validate.cu gptq_pro_kernel.cu -o gptq_pro_validate + * Run on an Ampere-or-newer CUDA GPU. The harness checks the LOP3 fragment + * decoder and compares every dispatch path with a scalar FP16-dequant/FP32- + * accumulation reference using GPTQ's native int32 qweight layout. */ #include "gptq_pro_kernel.cuh" -#include #include +#include #include -#include #include -cudaError_t gptq_pro_gemm( - const half* A, - const uint8_t* B_packed, - const half* S, - half* C, - int M, int N, int K, int group_size, - cudaStream_t stream); - -#define CHECK_CUDA(expr) \ - do { \ - cudaError_t _err = (expr); \ - if (_err != cudaSuccess) { \ - fprintf(stderr, "CUDA error at %s:%d: %s\n", \ - __FILE__, __LINE__, cudaGetErrorString(_err)); \ - return 1; \ - } \ +#define CHECK_CUDA(expr) \ + do { \ + const cudaError_t error = (expr); \ + if (error != cudaSuccess) { \ + std::fprintf( \ + stderr, \ + "CUDA error at %s:%d: %s\n", \ + __FILE__, \ + __LINE__, \ + cudaGetErrorString(error)); \ + return false; \ + } \ } while (0) -// ============================================================ -// TODO 1 — Decode-only scalar reference -// ============================================================ - -/// Scalar (lane-independent) decode of one lane-local 4-nibble B fragment. -/// Mirrors exactly what decode_bfrag_to_rb() does on-device, using only -/// host-visible arithmetic so the result can be used as ground truth. -/// -/// @param w0 unsigned 4-bit integer [0, 15] for RB[0] low half -/// @param w1 unsigned 4-bit integer [0, 15] for RB[0] high half -/// @param w2 unsigned 4-bit integer [0, 15] for RB[1] low half -/// @param w3 unsigned 4-bit integer [0, 15] for RB[1] high half -/// @param scale_f FP32 version of the per-group scale -/// @param zp_f FP32 version of the per-group zero-point -/// @param out_rb_f [out] dequantized values in order {w0, w1, w2, w3} -inline void scalar_decode_bfrag(uint32_t w0, uint32_t w1, - uint32_t w2, uint32_t w3, - float scale_f, float zp_f, - float (&out_rb_f)[4]) { - out_rb_f[0] = scale_f * (static_cast(w0) - zp_f); - out_rb_f[1] = scale_f * (static_cast(w1) - zp_f); - out_rb_f[2] = scale_f * (static_cast(w2) - zp_f); - out_rb_f[3] = scale_f * (static_cast(w3) - zp_f); +namespace { + +const char* mode_name(int mode) { + switch (mode) { + case GPTQ_PRO_KERNEL_AUTO: + return "auto"; + case GPTQ_PRO_KERNEL_GEMV: + return "gemv"; + case GPTQ_PRO_KERNEL_AMPERE: + return "ampere"; + case GPTQ_PRO_KERNEL_LEGACY: + return "legacy"; + default: + return "unknown"; + } } -// --------------------------------------------------------------------------- -// Device-side TODO 1 validation kernel -// -// Each thread (= one warp lane): -// 1. Reads its packed INT4 data from a pre-filled Bfrag smem region. -// 2. Calls fetch_bfrag_packed16() (the real ld.shared.u32 path). -// 3. Calls decode_bfrag_to_rb() to obtain RB[0], RB[1]. -// 4. Compares against a pre-computed float reference stored in ref_rb -// (4 floats per lane, computed by scalar_decode_bfrag on host). -// 5. Sets result[lane] = 1 if all four decoded values match. -// --------------------------------------------------------------------------- __global__ void validate_decode_kernel( - const uint32_t* __restrict__ bfrag_smem_src, // GPTQ_PRO_BFRAG_WORDS_PER_BUF words (global, will be copied to smem) - float scale_f, - float zp_f, - const float* ref_rb, // [WARP_SIZE * 4] ground-truth per lane - int* result) // [WARP_SIZE] 1=pass, 0=fail -{ - // One warp handles one tile at ks=0, j=0. - const int lane = threadIdx.x & (GPTQ_PRO_WARP_SIZE - 1); - - // ---- Stage Bfrag into shared memory ---- - extern __shared__ uint32_t smem_bfrag[]; - // Each thread copies one word. - if (lane < GPTQ_PRO_BFRAG_WORDS_PER_BUF) { - smem_bfrag[lane] = bfrag_smem_src[lane]; + const uint16_t* __restrict__ packed, + const half* __restrict__ scales, + float* __restrict__ decoded) { + const int lane = threadIdx.x; + if (lane >= GPTQ_PRO_WARP_SIZE) { + return; } - __syncthreads(); - - // ---- TODO 1: fetch via ld.shared.u32 path ---- - const int ks = 0, j = 0, buf = 0; - uint16_t packed_16 = fetch_bfrag_packed16(smem_bfrag, buf, ks, j, lane); - - // ---- Decode nibbles ---- - const half scale = __float2half(scale_f); - const half zp = __float2half(zp_f); - uint32_t RB[2]; - decode_bfrag_to_rb(packed_16, scale, zp, RB); - - // ---- Extract decoded FP16 values and compare with reference ---- - Half2Reg rb0_r, rb1_r; - rb0_r.u32 = RB[0]; - rb1_r.u32 = RB[1]; - - float got_rb0_lo = __half2float(__low2half(rb0_r.h2)); - float got_rb0_hi = __half2float(__high2half(rb0_r.h2)); - float got_rb1_lo = __half2float(__low2half(rb1_r.h2)); - float got_rb1_hi = __half2float(__high2half(rb1_r.h2)); - - // FP16 has ~1e-3 relative error; use 2 ULP tolerance in FP16 space. - const float tol = 2.0f * __half2float(__float2half(1.0f)) * 1e-3f; - - bool ok = (fabsf(got_rb0_lo - ref_rb[lane * 4 + 0]) <= tol + 1e-5f) && - (fabsf(got_rb0_hi - ref_rb[lane * 4 + 1]) <= tol + 1e-5f) && - (fabsf(got_rb1_lo - ref_rb[lane * 4 + 2]) <= tol + 1e-5f) && - (fabsf(got_rb1_hi - ref_rb[lane * 4 + 3]) <= tol + 1e-5f); - result[lane] = ok ? 1 : 0; -} -// ============================================================ -// TODO 2 — Scalar FP32-accumulating MMA reference -// ============================================================ - -/// Scalar reference for mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32. -/// -/// Interprets A and B as FP16 inputs and accumulates into FP32, matching the -/// hardware accumulation semantics of the tensor-core instruction. -/// -/// PTX fragment ownership for mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32: -/// A: row = groupID for ai in {0,1,4,5}, else groupID + 8 -/// col = 2 * threadID + (i & 1) [+8 for i >= 4] -/// B: row = 2 * threadID + (i & 1) [+8 for i >= 2] -/// col = groupID -/// C/D: row = groupID for ci in {0,1}, else groupID + 8 -/// col = 2 * threadID + (i & 1) -/// where groupID = lane >> 2 and threadID = lane & 3. -/// -/// This scalar function operates on the unpacked FP32 equivalents; the caller -/// is responsible for unpacking and repacking. -inline float fma_f32_from_f16_inputs_proxy(float a, float b, float c) { - const half ha = __float2half(a); - const half hb = __float2half(b); - return c + __half2float(ha) * __half2float(hb); + uint32_t registers[2]; + decode_bfrag_to_rb( + packed[lane], scales[lane], __float2half(8.0f), registers); + Half2Reg first; + Half2Reg second; + first.u32 = registers[0]; + second.u32 = registers[1]; + decoded[lane * 4 + 0] = __half2float(__low2half(first.h2)); + decoded[lane * 4 + 1] = __half2float(__high2half(first.h2)); + decoded[lane * 4 + 2] = __half2float(__low2half(second.h2)); + decoded[lane * 4 + 3] = __half2float(__high2half(second.h2)); } -inline float mma_ref_a_value(int m_row, int k_col) { - return 0.03125f * static_cast(m_row + 1) - + 0.001953125f * static_cast(k_col + 1); -} +bool validate_fragment_decode() { + std::vector packed(GPTQ_PRO_WARP_SIZE); + std::vector scales(GPTQ_PRO_WARP_SIZE); + std::vector decoded(GPTQ_PRO_WARP_SIZE * 4); -inline float mma_ref_b_value(int k_row, int group_id) { - return 0.0625f * static_cast(k_row + 1) - + 0.0078125f * static_cast(group_id); -} + for (int lane = 0; lane < GPTQ_PRO_WARP_SIZE; ++lane) { + const uint16_t w0 = static_cast((lane + 0) & 0xF); + const uint16_t w1 = static_cast((lane + 3) & 0xF); + const uint16_t w2 = static_cast((lane + 7) & 0xF); + const uint16_t w3 = static_cast((lane + 11) & 0xF); + packed[lane] = static_cast( + w0 | (w1 << 4) | (w2 << 8) | (w3 << 12)); + scales[lane] = + __float2half(0.015625f * static_cast(1 + lane % 4)); + } -inline uint32_t pack_half2(float lo, float hi) { - Half2Reg reg; - const half hlo = __float2half(lo); - const half hhi = __float2half(hi); - memcpy(®.u16[0], &hlo, sizeof(uint16_t)); - memcpy(®.u16[1], &hhi, sizeof(uint16_t)); - return reg.u32; -} + uint16_t* device_packed = nullptr; + half* device_scales = nullptr; + float* device_decoded = nullptr; + CHECK_CUDA(cudaMalloc(&device_packed, packed.size() * sizeof(uint16_t))); + CHECK_CUDA(cudaMalloc(&device_scales, scales.size() * sizeof(half))); + CHECK_CUDA(cudaMalloc(&device_decoded, decoded.size() * sizeof(float))); + CHECK_CUDA(cudaMemcpy( + device_packed, + packed.data(), + packed.size() * sizeof(uint16_t), + cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy( + device_scales, + scales.data(), + scales.size() * sizeof(half), + cudaMemcpyHostToDevice)); -// --------------------------------------------------------------------------- -// Device-side TODO 2 validation kernel -// -// One thread block = one warp. Validates a single (ks=0, j=0) MMA step. -// Steps: -// 1. Loads pre-dequantized RA[4] and RB[2] from global memory. -// 2. Zeroes RC[4]. -// 3. Calls mma_f32_m16n8k16() (the real tensor-core MMA). -// 4. Computes a scalar reference using fma_f32_from_f16_inputs_proxy(). -// 5. Compares RC output against reference; writes 1=pass / 0=fail. -// --------------------------------------------------------------------------- -__global__ void validate_mma_step_kernel( - const uint32_t* __restrict__ ra_global, // [4] per lane - const uint32_t* __restrict__ rb_global, // [2] per lane - const float* __restrict__ ref_rc, // [4] per lane (float proxy) - int* result) // [WARP_SIZE] 1=pass, 0=fail -{ - const int lane = threadIdx.x & (GPTQ_PRO_WARP_SIZE - 1); - - // Load fragment registers for this lane. - uint32_t RA[4], RB[2]; - float RC[4]; - RA[0] = ra_global[lane * 4 + 0]; - RA[1] = ra_global[lane * 4 + 1]; - RA[2] = ra_global[lane * 4 + 2]; - RA[3] = ra_global[lane * 4 + 3]; - RB[0] = rb_global[lane * 2 + 0]; - RB[1] = rb_global[lane * 2 + 1]; - RC[0] = 0.0f; - RC[1] = 0.0f; - RC[2] = 0.0f; - RC[3] = 0.0f; - - // ---- TODO 2: real tensor-core MMA ---- - mma_f32_m16n8k16(RA, RB, RC); - - // FP32 accumulation with exact FP16 inputs should agree closely with the - // scalar proxy. Allow a tiny epsilon for instruction-order differences. - const float tol = 1e-6f; - bool ok = true; - for (int i = 0; i < 4; ++i) { - if (fabsf(RC[i] - ref_rc[lane * 4 + i]) > tol) { - ok = false; - } - } - result[lane] = ok ? 1 : 0; -} + validate_decode_kernel<<<1, GPTQ_PRO_WARP_SIZE>>>( + device_packed, device_scales, device_decoded); + CHECK_CUDA(cudaGetLastError()); + CHECK_CUDA(cudaDeviceSynchronize()); + CHECK_CUDA(cudaMemcpy( + decoded.data(), + device_decoded, + decoded.size() * sizeof(float), + cudaMemcpyDeviceToHost)); -// ============================================================ -// Host-side driver — fills test data, launches kernels, checks results -// ============================================================ - -/// Fill Bfrag shared-memory image with deterministic lane-local INT4 payloads. -/// packed_16 for lane `l` in tile (0,0) packs four distinct nibbles: -/// bits [3:0] = (4*l + 0) & 0xF -/// bits [7:4] = (4*l + 1) & 0xF -/// bits [11:8] = (4*l + 2) & 0xF -/// bits [15:12] = (4*l + 3) & 0xF -static void fill_bfrag_test_image(uint32_t* words) { - for (int w = 0; w < GPTQ_PRO_BFRAG_WORDS_PER_BUF; ++w) { - words[w] = 0u; - } - // Tile ks=0, j=0, buf=0 starts at word offset 0. + cudaFree(device_packed); + cudaFree(device_scales); + cudaFree(device_decoded); + + int mismatches = 0; for (int lane = 0; lane < GPTQ_PRO_WARP_SIZE; ++lane) { - uint32_t w0 = (4u * lane + 0u) & 0xFu; - uint32_t w1 = (4u * lane + 1u) & 0xFu; - uint32_t w2 = (4u * lane + 2u) & 0xFu; - uint32_t w3 = (4u * lane + 3u) & 0xFu; - uint16_t p16 = static_cast( - w0 | (w1 << 4) | (w2 << 8) | (w3 << 12)); - int word_idx = lane >> 1; // lane pair - if (lane & 1) { - words[word_idx] = (words[word_idx] & 0x0000FFFFu) - | (static_cast(p16) << 16); - } else { - words[word_idx] = (words[word_idx] & 0xFFFF0000u) - | static_cast(p16); + const float scale = __half2float(scales[lane]); + for (int index = 0; index < 4; ++index) { + const int nibble = (packed[lane] >> (index * 4)) & 0xF; + const float expected = __half2float(__float2half( + scale * static_cast(nibble - 8))); + const float actual = decoded[lane * 4 + index]; + if (std::fabs(actual - expected) > 1e-4f) { + ++mismatches; + } } } + + if (mismatches != 0) { + std::fprintf(stderr, " decode mismatches: %d\n", mismatches); + return false; + } + std::printf(" PASS LOP3 fragment decode\n"); + return true; } -static void fill_end_to_end_a(std::vector& a, int M, int K) { +void fill_activations(std::vector& activations, int M, int K) { for (int m = 0; m < M; ++m) { for (int k = 0; k < K; ++k) { - const float value = 0.125f * static_cast(m + 1) - + 0.03125f * static_cast((k % 7) + 1); - a[m * K + k] = __float2half(value); + const float value = + 0.03125f * static_cast((m % 11) - 5) + + 0.0078125f * static_cast((k % 13) - 6); + activations[m * K + k] = __float2half(value); } } } -static void fill_end_to_end_b(std::vector& b_packed, int K, int N) { - const int packed_rows = (K + 1) / 2; - for (int kp = 0; kp < packed_rows; ++kp) { - const int k0 = kp * 2; +void fill_qweight(std::vector& qweight, int K, int N) { + const int qword_rows = (K + 7) / 8; + for (int row = 0; row < qword_rows; ++row) { for (int n = 0; n < N; ++n) { - const uint8_t lo = static_cast(8 + ((k0 + 2 * n) % 3)); - uint8_t hi = static_cast(8 + (((k0 + 1) + 2 * n) % 3)); - if (k0 + 1 >= K) { - hi = 8; + uint32_t word = 0; + for (int nibble_index = 0; nibble_index < 8; ++nibble_index) { + const int k = row * 8 + nibble_index; + const uint32_t nibble = + k < K + ? static_cast((k * 3 + n * 5 + 1) & 0xF) + : 8u; + word |= nibble << (nibble_index * 4); } - b_packed[kp * N + n] = static_cast(lo | (hi << 4)); + qweight[row * N + n] = static_cast(word); } } } -static void fill_end_to_end_s(std::vector& scales, int groups, int N) { - for (int g = 0; g < groups; ++g) { +void fill_scales(std::vector& scales, int groups, int N) { + for (int group = 0; group < groups; ++group) { for (int n = 0; n < N; ++n) { - const float scale = 0.125f * static_cast(1 + ((g + n) % 4)); - scales[g * N + n] = __float2half(scale); + const float value = + 0.0078125f * static_cast(1 + ((group * 3 + n) % 8)); + scales[group * N + n] = __float2half(value); } } } -static float dequant_weight_ref(const std::vector& b_packed, - const std::vector& scales, - int N, int group_size, - int k, int n) { - const uint8_t byte = b_packed[(k >> 1) * N + n]; - const uint32_t nibble = (k & 1) ? ((byte >> 4) & 0xFu) : (byte & 0xFu); +float reference_weight( + const std::vector& qweight, + const std::vector& scales, + int N, + int group_size, + int k, + int n) { + const uint32_t word = + static_cast(qweight[(k >> 3) * N + n]); + const int nibble = static_cast((word >> ((k & 7) * 4)) & 0xFu); const float scale = __half2float(scales[(k / group_size) * N + n]); - return scale * (static_cast(nibble) - 8.0f); + return __half2float( + __float2half(scale * static_cast(nibble - 8))); } -static bool run_end_to_end_case(int M, int N, int K, int group_size, const char* label) { - const int packed_rows = (K + 1) / 2; +bool run_case( + int M, + int N, + int K, + int group_size, + int mode, + const char* label) { + const int qword_rows = (K + 7) / 8; const int groups = (K + group_size - 1) / group_size; - std::vector h_a(M * K); - std::vector h_b(packed_rows * N); - std::vector h_s(groups * N); - std::vector h_c(M * N, __float2half(0.0f)); - - fill_end_to_end_a(h_a, M, K); - fill_end_to_end_b(h_b, K, N); - fill_end_to_end_s(h_s, groups, N); - - half* d_a = nullptr; - half* d_s = nullptr; - half* d_c = nullptr; - uint8_t* d_b = nullptr; - - auto fail = [&](const char* what, cudaError_t err) { - std::fprintf(stderr, " FAIL %s: %s\n", what, cudaGetErrorString(err)); - if (d_a) cudaFree(d_a); - if (d_s) cudaFree(d_s); - if (d_c) cudaFree(d_c); - if (d_b) cudaFree(d_b); + std::vector activations(M * K); + std::vector qweight(qword_rows * N); + std::vector scales(groups * N); + std::vector output(M * N, __float2half(0.0f)); + fill_activations(activations, M, K); + fill_qweight(qweight, K, N); + fill_scales(scales, groups, N); + + half* device_activations = nullptr; + int32_t* device_qweight = nullptr; + half* device_scales = nullptr; + half* device_output = nullptr; + CHECK_CUDA(cudaMalloc( + &device_activations, activations.size() * sizeof(half))); + CHECK_CUDA(cudaMalloc( + &device_qweight, qweight.size() * sizeof(int32_t))); + CHECK_CUDA(cudaMalloc(&device_scales, scales.size() * sizeof(half))); + CHECK_CUDA(cudaMalloc(&device_output, output.size() * sizeof(half))); + CHECK_CUDA(cudaMemcpy( + device_activations, + activations.data(), + activations.size() * sizeof(half), + cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy( + device_qweight, + qweight.data(), + qweight.size() * sizeof(int32_t), + cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy( + device_scales, + scales.data(), + scales.size() * sizeof(half), + cudaMemcpyHostToDevice)); + + const cudaError_t launch_status = gptq_pro_gemm( + device_activations, + device_qweight, + device_scales, + device_output, + M, + N, + K, + group_size, + 0, + mode); + if (launch_status != cudaSuccess) { + std::fprintf( + stderr, + " FAIL %s (%s) launch: %s\n", + label, + mode_name(mode), + cudaGetErrorString(launch_status)); + cudaFree(device_activations); + cudaFree(device_qweight); + cudaFree(device_scales); + cudaFree(device_output); return false; - }; + } + CHECK_CUDA(cudaDeviceSynchronize()); + CHECK_CUDA(cudaMemcpy( + output.data(), + device_output, + output.size() * sizeof(half), + cudaMemcpyDeviceToHost)); - cudaError_t err = cudaMalloc(&d_a, h_a.size() * sizeof(half)); - if (err != cudaSuccess) return fail("cudaMalloc(d_a)", err); - err = cudaMalloc(&d_s, h_s.size() * sizeof(half)); - if (err != cudaSuccess) return fail("cudaMalloc(d_s)", err); - err = cudaMalloc(&d_c, h_c.size() * sizeof(half)); - if (err != cudaSuccess) return fail("cudaMalloc(d_c)", err); - err = cudaMalloc(&d_b, h_b.size() * sizeof(uint8_t)); - if (err != cudaSuccess) return fail("cudaMalloc(d_b)", err); - - err = cudaMemcpy(d_a, h_a.data(), h_a.size() * sizeof(half), cudaMemcpyHostToDevice); - if (err != cudaSuccess) return fail("cudaMemcpy(d_a)", err); - err = cudaMemcpy(d_s, h_s.data(), h_s.size() * sizeof(half), cudaMemcpyHostToDevice); - if (err != cudaSuccess) return fail("cudaMemcpy(d_s)", err); - err = cudaMemcpy(d_b, h_b.data(), h_b.size() * sizeof(uint8_t), cudaMemcpyHostToDevice); - if (err != cudaSuccess) return fail("cudaMemcpy(d_b)", err); - err = cudaMemset(d_c, 0, h_c.size() * sizeof(half)); - if (err != cudaSuccess) return fail("cudaMemset(d_c)", err); - - err = gptq_pro_gemm(d_a, d_b, d_s, d_c, M, N, K, group_size, 0); - if (err != cudaSuccess) return fail("gptq_pro_gemm launch", err); - err = cudaDeviceSynchronize(); - if (err != cudaSuccess) return fail("cudaDeviceSynchronize()", err); - err = cudaMemcpy(h_c.data(), d_c, h_c.size() * sizeof(half), cudaMemcpyDeviceToHost); - if (err != cudaSuccess) return fail("cudaMemcpy(h_c)", err); + cudaFree(device_activations); + cudaFree(device_qweight); + cudaFree(device_scales); + cudaFree(device_output); int mismatches = 0; + float maximum_error = 0.0f; for (int m = 0; m < M; ++m) { for (int n = 0; n < N; ++n) { - float acc = 0.0f; + float reference = 0.0f; for (int k = 0; k < K; ++k) { - const float a = __half2float(h_a[m * K + k]); - const float w = dequant_weight_ref(h_b, h_s, N, group_size, k, n); - acc = fma_f32_from_f16_inputs_proxy(a, w, acc); + const float activation = + __half2float(activations[m * K + k]); + const float weight = reference_weight( + qweight, scales, N, group_size, k, n); + reference = std::fma(activation, weight, reference); } - const float expect = __half2float(__float2half(acc)); - const float got = __half2float(h_c[m * N + n]); - if (fabsf(got - expect) > 1e-3f) { + const float expected = __half2float(__float2half(reference)); + const float actual = __half2float(output[m * N + n]); + const float error = std::fabs(actual - expected); + maximum_error = error > maximum_error ? error : maximum_error; + const float tolerance = 0.02f + 0.002f * std::fabs(expected); + if (!std::isfinite(actual) || error > tolerance) { if (mismatches < 8) { - std::fprintf(stderr, - " FAIL %s at (%d, %d): got=%f expect=%f\n", - label, m, n, got, expect); + std::fprintf( + stderr, + " FAIL %s (%s) at (%d,%d): got=%f expected=%f error=%f\n", + label, + mode_name(mode), + m, + n, + actual, + expected, + error); } ++mismatches; } } } - cudaFree(d_a); - cudaFree(d_s); - cudaFree(d_c); - cudaFree(d_b); - if (mismatches != 0) { - std::fprintf(stderr, " %s mismatches: %d\n", label, mismatches); + std::fprintf( + stderr, + " %s (%s) mismatches=%d max_error=%f\n", + label, + mode_name(mode), + mismatches, + maximum_error); return false; } - std::printf(" PASS %s\n", label); + std::printf( + " PASS %-24s mode=%-7s shape=%dx%dx%d g=%d max_error=%f\n", + label, + mode_name(mode), + M, + N, + K, + group_size, + maximum_error); return true; } -#ifndef GPTQ_PRO_VALIDATE_SKIP_MAIN +} // namespace int main() { - const float scale_f = 0.015625f; // 2^-6, exact in FP16 - const float zp_f = 8.0f; // unsigned-to-signed shift - - // ----------------------------------------------------------------------- - // TODO 1 — Decode validation - // ----------------------------------------------------------------------- - printf("=== TODO 1: decode-only validation ===\n"); - - // Build host Bfrag image. - const int bfrag_words = GPTQ_PRO_BFRAG_WORDS_PER_BUF; - uint32_t h_bfrag[bfrag_words]; - fill_bfrag_test_image(h_bfrag); - - // Compute scalar reference for every lane. - float h_ref_rb[GPTQ_PRO_WARP_SIZE * 4]; - for (int lane = 0; lane < GPTQ_PRO_WARP_SIZE; ++lane) { - uint32_t w0 = (4u * lane + 0u) & 0xFu; - uint32_t w1 = (4u * lane + 1u) & 0xFu; - uint32_t w2 = (4u * lane + 2u) & 0xFu; - uint32_t w3 = (4u * lane + 3u) & 0xFu; - float decoded[4]; - scalar_decode_bfrag(w0, w1, w2, w3, scale_f, zp_f, decoded); - for (int i = 0; i < 4; ++i) { - h_ref_rb[lane * 4 + i] = decoded[i]; - } - } - - // Allocate device memory. - uint32_t* d_bfrag; - float* d_ref_rb; - int* d_result1; - CHECK_CUDA(cudaMalloc(&d_bfrag, bfrag_words * sizeof(uint32_t))); - CHECK_CUDA(cudaMalloc(&d_ref_rb, GPTQ_PRO_WARP_SIZE * 4 * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_result1, GPTQ_PRO_WARP_SIZE * sizeof(int))); - - CHECK_CUDA(cudaMemcpy( - d_bfrag, h_bfrag, bfrag_words * sizeof(uint32_t), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy( - d_ref_rb, h_ref_rb, GPTQ_PRO_WARP_SIZE * 4 * sizeof(float), - cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemset(d_result1, 0, GPTQ_PRO_WARP_SIZE * sizeof(int))); - - size_t smem_bytes = GPTQ_PRO_BFRAG_WORDS_PER_BUF * sizeof(uint32_t); - validate_decode_kernel<<<1, GPTQ_PRO_WARP_SIZE, smem_bytes>>>( - d_bfrag, scale_f, zp_f, d_ref_rb, d_result1); - CHECK_CUDA(cudaGetLastError()); - CHECK_CUDA(cudaDeviceSynchronize()); - - int h_result1[GPTQ_PRO_WARP_SIZE]; - CHECK_CUDA(cudaMemcpy( - h_result1, d_result1, GPTQ_PRO_WARP_SIZE * sizeof(int), - cudaMemcpyDeviceToHost)); - - int pass1 = 0; - for (int i = 0; i < GPTQ_PRO_WARP_SIZE; ++i) { - const bool ok = (h_result1[i] == 1); - pass1 += ok ? 1 : 0; - if (!ok) printf(" FAIL lane %d\n", i); - } - printf(" %d / %d lanes passed\n", pass1, GPTQ_PRO_WARP_SIZE); - - CHECK_CUDA(cudaFree(d_bfrag)); - CHECK_CUDA(cudaFree(d_ref_rb)); - CHECK_CUDA(cudaFree(d_result1)); - - // ----------------------------------------------------------------------- - // TODO 2 — MMA step validation - // ----------------------------------------------------------------------- - printf("=== TODO 2: ks/j MMA step validation ===\n"); - - // Build synthetic RA, RB fragments from the PTX-defined fragment ownership. - // Both A and B vary with their logical coordinates, so row/column mix-ups in - // either fragment contract now perturb the final D fragment immediately. - float h_a_tile[GPTQ_PRO_M_PER_WARP][GPTQ_PRO_K_PER_WARP]; - float h_b_tile[GPTQ_PRO_K_PER_WARP][8]; - for (int m = 0; m < GPTQ_PRO_M_PER_WARP; ++m) { - for (int k = 0; k < GPTQ_PRO_K_PER_WARP; ++k) { - h_a_tile[m][k] = mma_ref_a_value(m, k); - } - } - for (int k = 0; k < GPTQ_PRO_K_PER_WARP; ++k) { - for (int n = 0; n < 8; ++n) { - h_b_tile[k][n] = mma_ref_b_value(k, n); - } - } + int passed = 0; + int total = 0; + + ++total; + passed += validate_fragment_decode() ? 1 : 0; + + struct TestCase { + int M; + int N; + int K; + int group_size; + int mode; + const char* label; + }; - uint32_t h_ra[GPTQ_PRO_WARP_SIZE * 4]; - uint32_t h_rb[GPTQ_PRO_WARP_SIZE * 2]; - for (int lane = 0; lane < GPTQ_PRO_WARP_SIZE; ++lane) { - const int group_id = lane >> 2; - const int thread_id = lane & 3; - const int a_col_lo = 2 * thread_id; - const int a_col_hi = a_col_lo + 8; - const int row0 = 2 * thread_id + 0; - const int row1 = 2 * thread_id + 1; - - h_ra[lane * 4 + 0] = pack_half2( - h_a_tile[group_id + 0][a_col_lo + 0], - h_a_tile[group_id + 0][a_col_lo + 1]); - h_ra[lane * 4 + 1] = pack_half2( - h_a_tile[group_id + 8][a_col_lo + 0], - h_a_tile[group_id + 8][a_col_lo + 1]); - h_ra[lane * 4 + 2] = pack_half2( - h_a_tile[group_id + 0][a_col_hi + 0], - h_a_tile[group_id + 0][a_col_hi + 1]); - h_ra[lane * 4 + 3] = pack_half2( - h_a_tile[group_id + 8][a_col_hi + 0], - h_a_tile[group_id + 8][a_col_hi + 1]); - - h_rb[lane * 2 + 0] = pack_half2( - h_b_tile[row0 + 0][group_id], - h_b_tile[row1 + 0][group_id]); - h_rb[lane * 2 + 1] = pack_half2( - h_b_tile[row0 + 8][group_id], - h_b_tile[row1 + 8][group_id]); - } + const TestCase cases[] = { + {1, 128, 128, 32, GPTQ_PRO_KERNEL_GEMV, "decode-m1"}, + {4, 80, 128, 32, GPTQ_PRO_KERNEL_GEMV, "decode-m4"}, + {1, 128, 128, 32, GPTQ_PRO_KERNEL_AUTO, "auto-decode"}, + {16, 64, 128, 32, GPTQ_PRO_KERNEL_AMPERE, "ampere-single-warp-n"}, + {31, 320, 256, 64, GPTQ_PRO_KERNEL_AMPERE, "ampere-tail-mn"}, + {32, 320, 256, 64, GPTQ_PRO_KERNEL_AUTO, "auto-ampere"}, + {13, 41, 29, 16, GPTQ_PRO_KERNEL_LEGACY, "legacy-edge"}, + {13, 41, 29, 16, GPTQ_PRO_KERNEL_AUTO, "auto-legacy"}, + }; - // Scalar reference: each D[m][n] = sum_{k=0}^{15} A[m][k] * B[k][n]. - // For m16n8k16.row.col.f32 the lane-local D fragment is a 2-row x 2-column - // tile, so compare against the full 16x8 GEMM reference rather than a - // collapsed per-column reduction. - // rows = {lane >> 2, lane >> 2 + 8} - // cols = {2 * (lane & 3), 2 * (lane & 3) + 1} - float h_ref_rc[GPTQ_PRO_WARP_SIZE * 4]; - { - float h_d_tile[GPTQ_PRO_M_PER_WARP][8]; - for (int m = 0; m < GPTQ_PRO_M_PER_WARP; ++m) { - for (int n = 0; n < 8; ++n) { - float acc = 0.0f; - for (int k = 0; k < GPTQ_PRO_K_PER_WARP; ++k) { - acc = fma_f32_from_f16_inputs_proxy( - h_a_tile[m][k], h_b_tile[k][n], acc); - } - h_d_tile[m][n] = acc; - } - } - for (int lane = 0; lane < GPTQ_PRO_WARP_SIZE; ++lane) { - const int row_base = lane >> 2; - const int col_base = 2 * (lane & 3); - h_ref_rc[lane * 4 + 0] = h_d_tile[row_base + 0][col_base + 0]; - h_ref_rc[lane * 4 + 1] = h_d_tile[row_base + 0][col_base + 1]; - h_ref_rc[lane * 4 + 2] = h_d_tile[row_base + 8][col_base + 0]; - h_ref_rc[lane * 4 + 3] = h_d_tile[row_base + 8][col_base + 1]; - } + for (const auto& test_case : cases) { + ++total; + passed += run_case( + test_case.M, + test_case.N, + test_case.K, + test_case.group_size, + test_case.mode, + test_case.label) + ? 1 + : 0; } - uint32_t* d_ra; - uint32_t* d_rb; - float* d_ref_rc; - int* d_result2; - CHECK_CUDA(cudaMalloc(&d_ra, GPTQ_PRO_WARP_SIZE * 4 * sizeof(uint32_t))); - CHECK_CUDA(cudaMalloc(&d_rb, GPTQ_PRO_WARP_SIZE * 2 * sizeof(uint32_t))); - CHECK_CUDA(cudaMalloc(&d_ref_rc, GPTQ_PRO_WARP_SIZE * 4 * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_result2, GPTQ_PRO_WARP_SIZE * sizeof(int))); - - CHECK_CUDA(cudaMemcpy( - d_ra, h_ra, GPTQ_PRO_WARP_SIZE * 4 * sizeof(uint32_t), - cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy( - d_rb, h_rb, GPTQ_PRO_WARP_SIZE * 2 * sizeof(uint32_t), - cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy( - d_ref_rc, h_ref_rc, GPTQ_PRO_WARP_SIZE * 4 * sizeof(float), - cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemset(d_result2, 0, GPTQ_PRO_WARP_SIZE * sizeof(int))); - - validate_mma_step_kernel<<<1, GPTQ_PRO_WARP_SIZE>>>( - d_ra, d_rb, d_ref_rc, d_result2); - CHECK_CUDA(cudaGetLastError()); - CHECK_CUDA(cudaDeviceSynchronize()); - - int h_result2[GPTQ_PRO_WARP_SIZE]; - CHECK_CUDA(cudaMemcpy( - h_result2, d_result2, GPTQ_PRO_WARP_SIZE * sizeof(int), - cudaMemcpyDeviceToHost)); - - int pass2 = 0; - for (int i = 0; i < GPTQ_PRO_WARP_SIZE; ++i) { - const bool ok = (h_result2[i] == 1); - pass2 += ok ? 1 : 0; - if (!ok) printf(" FAIL lane %d\n", i); - } - printf(" %d / %d lanes passed\n", pass2, GPTQ_PRO_WARP_SIZE); - - CHECK_CUDA(cudaFree(d_ra)); - CHECK_CUDA(cudaFree(d_rb)); - CHECK_CUDA(cudaFree(d_ref_rc)); - CHECK_CUDA(cudaFree(d_result2)); - - // ----------------------------------------------------------------------- - // TODO 3 — End-to-end kernel validation - // ----------------------------------------------------------------------- - printf("=== TODO 3: end-to-end kernel validation ===\n"); - int pass3 = 0; - pass3 += run_end_to_end_case(16, 64, 16, 16, "aligned-16x64x16") ? 1 : 0; - pass3 += run_end_to_end_case(13, 41, 29, 16, "edge-13x41x29") ? 1 : 0; - printf(" %d / %d cases passed\n", pass3, 2); - - // ----------------------------------------------------------------------- - int total = pass1 + pass2 + pass3; - int total_max = 2 * GPTQ_PRO_WARP_SIZE + 2; - printf("\n=== Overall: %d / %d checks passed ===\n", total, total_max); - return (total == total_max) ? 0 : 1; + std::printf("\n=== GPTQ-Pro CUDA validation: %d / %d passed ===\n", passed, total); + return passed == total ? 0 : 1; } - -#endif // GPTQ_PRO_VALIDATE_SKIP_MAIN diff --git a/scripts/benchmark_gptq_pro_kernel.py b/scripts/benchmark_gptq_pro_kernel.py new file mode 100644 index 000000000..b383d86cf --- /dev/null +++ b/scripts/benchmark_gptq_pro_kernel.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 +"""Benchmark GPTQ-Pro's raw CUDA kernels without model-loading overhead. + +The script measures AUTO, legacy, and applicable specialized paths for a set of +M values. It checks a column subset against a PyTorch FP32 reference using +FP16-dequantized weights, matching the kernel's numerical contract. + +Example: + + python scripts/benchmark_gptq_pro_kernel.py \ + --m-values 1,4,8,16,64,256 --n 4096 --k 4096 --group-size 128 \ + --warmup 20 --iterations 100 --output kernel-results.json +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from pathlib import Path + +import torch + +from gptqmodel.utils.gptq_pro import ensure_gptq_pro_loaded + + +def parse_int_list(value: str) -> list[int]: + values = [int(item.strip()) for item in value.split(",") if item.strip()] + if not values or any(item <= 0 for item in values): + raise argparse.ArgumentTypeError( + "expected a comma-separated list of positive integers" + ) + return values + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--m-values", type=parse_int_list, default=parse_int_list("1,4,8,16,64,256") + ) + parser.add_argument("--n", type=int, default=4096) + parser.add_argument("--k", type=int, default=4096) + parser.add_argument("--group-size", type=int, default=128) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--check-columns", type=int, default=256) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--output", type=Path, default=None) + args = parser.parse_args() + + for name in ("n", "k", "group_size", "warmup", "iterations"): + if getattr(args, name) <= 0: + parser.error(f"--{name.replace('_', '-')} must be greater than zero") + if args.k % 8 != 0: + parser.error("--k must be divisible by 8 for native GPTQ int32 packing") + if args.group_size % 16 != 0: + parser.error("--group-size must be a multiple of 16") + if args.k % args.group_size != 0: + parser.error("--k must be divisible by --group-size for this benchmark") + if args.check_columns <= 0: + parser.error("--check-columns must be greater than zero") + return args + + +def make_problem(m: int, n: int, k: int, group_size: int, device: torch.device): + activations = torch.randn((m, k), device=device, dtype=torch.float16) + qweight_bytes = torch.randint( + 0, + 256, + (k // 8, n, 4), + device=device, + dtype=torch.uint8, + ) + qweight = qweight_bytes.view(torch.int32).squeeze(-1).contiguous() + scales = ( + torch.rand((k // group_size, n), device=device, dtype=torch.float16) * 0.08 + + 0.001 + ) + return activations, qweight, scales + + +def dequantize_reference( + qweight: torch.Tensor, + scales: torch.Tensor, + k: int, + group_size: int, +) -> torch.Tensor: + shifts = torch.arange( + 0, + 32, + 4, + device=qweight.device, + dtype=torch.int32, + ).view(1, 8, 1) + values = ( + torch.bitwise_and(torch.bitwise_right_shift(qweight.unsqueeze(1), shifts), 0xF) + .reshape(k, qweight.shape[1]) + .to(torch.float16) + - 8 + ) + group_index = torch.arange(k, device=qweight.device) // group_size + return values * scales.index_select(0, group_index) + + +def check_numerics( + module, + activations: torch.Tensor, + qweight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + mode: str, + columns: int, +) -> dict[str, float]: + columns = min(columns, qweight.shape[1]) + if columns >= 16: + columns = max(16, columns - columns % 16) + qweight_subset = qweight[:, :columns].contiguous() + scales_subset = scales[:, :columns].contiguous() + + actual = module.gptq_pro_gemm( + activations, + qweight_subset, + scales_subset, + group_size, + mode, + ) + weights = dequantize_reference( + qweight_subset, + scales_subset, + activations.shape[1], + group_size, + ) + expected = torch.matmul(activations.float(), weights.float()).to(torch.float16) + difference = actual.float() - expected.float() + maximum_absolute_error = difference.abs().max().item() + mean_absolute_error = difference.abs().mean().item() + cosine = torch.nn.functional.cosine_similarity( + actual.float().reshape(1, -1), + expected.float().reshape(1, -1), + ).item() + if not torch.isfinite(actual).all(): + raise RuntimeError(f"non-finite output from kernel mode {mode}") + return { + "columns": columns, + "max_abs_error": maximum_absolute_error, + "mean_abs_error": mean_absolute_error, + "cosine_similarity": cosine, + } + + +def benchmark_mode( + module, + activations: torch.Tensor, + qweight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + mode: str, + warmup: int, + iterations: int, +) -> dict[str, float | str]: + for _ in range(warmup): + module.gptq_pro_gemm(activations, qweight, scales, group_size, mode) + torch.cuda.synchronize(activations.device) + + times_ms: list[float] = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + module.gptq_pro_gemm(activations, qweight, scales, group_size, mode) + end.record() + end.synchronize() + times_ms.append(float(start.elapsed_time(end))) + + median_ms = statistics.median(times_ms) + mean_ms = statistics.fmean(times_ms) + sorted_times = sorted(times_ms) + p95_ms = sorted_times[min(len(sorted_times) - 1, int(len(sorted_times) * 0.95))] + m, k = activations.shape + n = qweight.shape[1] + dense_operations = 2.0 * m * n * k + tflops = dense_operations / (median_ms / 1000.0) / 1.0e12 + minimum_bytes = ( + activations.numel() * activations.element_size() + + qweight.numel() * qweight.element_size() + + scales.numel() * scales.element_size() + + m * n * activations.element_size() + ) + minimum_bandwidth_gbs = minimum_bytes / (median_ms / 1000.0) / 1.0e9 + return { + "mode": mode, + "median_ms": median_ms, + "mean_ms": mean_ms, + "p95_ms": p95_ms, + "effective_dense_tflops": tflops, + "minimum_bandwidth_gbs": minimum_bandwidth_gbs, + } + + +def applicable_modes(m: int, n: int, k: int) -> list[str]: + modes = ["auto", "legacy"] + if m <= 4 and k % 8 == 0: + modes.append("gemv") + if n % 16 == 0 and k % 16 == 0: + modes.append("ampere") + return modes + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required") + major, minor = torch.cuda.get_device_capability() + if major < 8: + raise SystemExit(f"compute capability 8.0+ is required, got {major}.{minor}") + + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + torch.backends.cuda.matmul.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + device = torch.device("cuda:0") + module = ensure_gptq_pro_loaded(verbose=True) + properties = torch.cuda.get_device_properties(device) + report = { + "created_unix": time.time(), + "torch_version": torch.__version__, + "cuda_runtime": torch.version.cuda, + "device": properties.name, + "compute_capability": f"{major}.{minor}", + "total_memory_bytes": properties.total_memory, + "shape": {"n": args.n, "k": args.k, "group_size": args.group_size}, + "warmup": args.warmup, + "iterations": args.iterations, + "results": [], + } + + print( + f"GPTQ-Pro raw kernel benchmark on {properties.name} (sm_{major}{minor}), " + f"N={args.n}, K={args.k}, group={args.group_size}" + ) + for m in args.m_values: + activations, qweight, scales = make_problem( + m, args.n, args.k, args.group_size, device + ) + for mode in applicable_modes(m, args.n, args.k): + numerical = check_numerics( + module, + activations, + qweight, + scales, + args.group_size, + mode, + args.check_columns, + ) + timing = benchmark_mode( + module, + activations, + qweight, + scales, + args.group_size, + mode, + args.warmup, + args.iterations, + ) + result = {"m": m, **timing, "numerical": numerical} + report["results"].append(result) + print( + f"M={m:4d} mode={mode:7s} median={timing['median_ms']:.4f} ms " + f"p95={timing['p95_ms']:.4f} ms " + f"TFLOP/s={timing['effective_dense_tflops']:.2f} " + f"cos={numerical['cosine_similarity']:.8f}" + ) + del activations, qweight, scales + torch.cuda.empty_cache() + + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2), encoding="utf-8") + print(f"Wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/test_gptq_pro_ampere_pipeline.py b/tests/kernels/test_gptq_pro_ampere_pipeline.py new file mode 100644 index 000000000..887ef28c3 --- /dev/null +++ b/tests/kernels/test_gptq_pro_ampere_pipeline.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import importlib.util +import logging +import sys +import types +from pathlib import Path + +import pytest +import torch + + +ROOT = Path(__file__).resolve().parents[2] + + +def _module(name: str, **attributes): + module = types.ModuleType(name) + for attribute, value in attributes.items(): + setattr(module, attribute, value) + sys.modules[name] = module + return module + + +def _load_gptq_pro_runtime_module(): + """Load the helper without importing the repository's package root.""" + namespace = "_gptq_pro_kernel_test" + root_package = _module(namespace) + root_package.__path__ = [str(ROOT / "gptqmodel")] + utils_package = _module(f"{namespace}.utils") + utils_package.__path__ = [str(ROOT / "gptqmodel/utils")] + + _module( + f"{namespace}.utils._extension_loader", + load_extension_module=lambda _name: (_ for _ in ()).throw( + ImportError("test stub") + ), + ) + _module(f"{namespace}.utils.env", env_flag=lambda *_args, **_kwargs: False) + _module( + f"{namespace}.utils.logger", + setup_logger=lambda: logging.getLogger("gptq-pro-kernel-test"), + ) + _module(f"{namespace}.utils.rocm", IS_ROCM=False) + + module_name = f"{namespace}.utils.gptq_pro" + spec = importlib.util.spec_from_file_location( + module_name, + ROOT / "gptqmodel/utils/gptq_pro.py", + ) + assert spec is not None and spec.loader is not None + runtime = importlib.util.module_from_spec(spec) + sys.modules[module_name] = runtime + spec.loader.exec_module(runtime) + return runtime + + +RUNTIME = _load_gptq_pro_runtime_module() + + +def test_compatibility_byte_view_matches_gptq_nibble_pair_layout(): + qweight = torch.tensor( + [ + [0x76543210, 0x01234567], + [0x11111111, 0x22222222], + ], + dtype=torch.int32, + ) + + packed = RUNTIME.gptq_pro_qweight_to_b_packed(qweight) + + expected = torch.tensor( + [ + [0x10, 0x67], + [0x32, 0x45], + [0x54, 0x23], + [0x76, 0x01], + [0x11, 0x22], + [0x11, 0x22], + [0x11, 0x22], + [0x11, 0x22], + ], + dtype=torch.uint8, + ) + assert torch.equal(packed, expected) + assert packed.is_contiguous() + + +def test_kernel_mode_normalization_and_environment(monkeypatch): + assert RUNTIME.normalize_gptq_pro_kernel_mode(" AMPERE ") == "ampere" + monkeypatch.setenv("GPTQMODEL_GPTQ_PRO_KERNEL", "gemv") + assert RUNTIME.normalize_gptq_pro_kernel_mode() == "gemv" + + with pytest.raises(ValueError, match="kernel mode"): + RUNTIME.normalize_gptq_pro_kernel_mode("unknown") + + +def test_ampere_source_contains_specialized_dispatch_and_pipeline(): + header = (ROOT / "gptqmodel_ext/gptq_pro/gptq_pro_kernel.cuh").read_text( + encoding="utf-8" + ) + source = (ROOT / "gptqmodel_ext/gptq_pro/gptq_pro_kernel.cu").read_text( + encoding="utf-8" + ) + + assert "GPTQ_PRO_WARPS_PER_CTA = 4" in header + assert "GPTQ_PRO_PIPE = 2" in header + assert "GPTQ_PRO_QWORD_ROWS_PER_K_TILE" in header + assert "load_qweight_bfrag_packed16" in header + assert "cp.async.ca.shared.global" in header + assert "cp.async.cg.shared.global" in header + assert "lop3.b32" in header + assert "gptq_pro_gemv_kernel" in source + assert "gptq_pro_gemm_kernel_ampere" in source + assert "gptq_pro_gemm_kernel_legacy" in source + assert "cp_async_wait_group<1>()" in source + + +def test_runtime_uses_native_qweight_without_duplicate_buffer(): + qlinear = (ROOT / "gptqmodel/nn_modules/qlinear/gptq_pro.py").read_text( + encoding="utf-8" + ) + binding = (ROOT / "gptqmodel_ext/gptq_pro/gptq_pro_torch.cpp").read_text( + encoding="utf-8" + ) + + assert "qweight=self.qweight" in qlinear + assert "b_packed" not in qlinear + assert "qweight.data_ptr()" in binding + assert "torch::kInt32" in binding + assert "ceil(K / 8)" in binding + + +def test_cuda_compile_workflow_covers_ampere_targets(): + workflow = (ROOT / ".github/workflows/gptq-pro-cuda-compile.yml").read_text( + encoding="utf-8" + ) + + for architecture in ("sm_80", "sm_86", "sm_87"): + assert architecture in workflow + assert "gptq_pro_validate.cu gptq_pro_kernel.cu" in workflow + assert "gptq_pro_torch.cpp" in workflow