Harden real-GPU paths and add end-to-end GPU test suite

.github/workflows/ci.yml

@@ -27,7 +27,7 @@ jobs:
       - name: Install ruff
         run: |
           python -m pip install --upgrade pip
-          pip install "ruff>=0.5.0"
+          pip install "ruff==0.15.21"  # pinned: newer ruff also reformats Markdown code blocks
       - name: Run ruff check
         run: ruff check .
       - name: Run ruff format check
@@ -55,7 +55,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
@@ -67,19 +67,23 @@ jobs:
           python -m pip install --upgrade pip
           pip install --index-url https://download.pytorch.org/whl/cpu "torch>=2.0.0"
           pip install \
-            "transformers>=4.36.0" \
+            "transformers>=4.46.0" \
             "datasets>=2.14.0" \
             "accelerate>=0.24.0" \
-            "peft>=0.6.0" \
+            "peft>=0.12.0" \
             "scipy>=1.10.0" \
             "scikit-learn>=1.3.0" \
             "tqdm>=4.65.0" \
             "rich>=13.0.0" \
             "huggingface_hub>=0.20.0" \
-            "bitsandbytes>=0.41.0" \
+            "bitsandbytes>=0.43.0" \
             "psutil" \
             "gguf" \
             "py-cpuinfo" \
+            "fastapi>=0.100.0" \
+            "uvicorn>=0.23.0" \
+            "pydantic>=2.0.0" \
+            "httpx>=0.24.0" \
             "pytest>=7.4.0"
       - name: Install QuantLLM (no deps)
         run: pip install --no-deps -e .

.github/workflows/gpu.yml

@@ -0,0 +1,45 @@
+name: GPU E2E
+
+# Runs tests/gpu on a real GPU through Modal. Manual, or on PRs labelled
+# "run-gpu-tests" so normal pushes don't spend GPU credit.
+# Needs repo secrets MODAL_TOKEN_ID and MODAL_TOKEN_SECRET.
+
+on:
+  workflow_dispatch:
+    inputs:
+      gpu:
+        description: "GPU type (T4, L4, A10G, H100)"
+        default: "T4"
+      slow:
+        description: "Include slow tests (GGUF via llama.cpp, ONNX)"
+        type: boolean
+        default: false
+  pull_request:
+    types: [labeled, synchronize]
+
+permissions:
+  contents: read
+
+concurrency:
+  group: gpu-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  gpu-tests:
+    if: github.event_name == 'workflow_dispatch' || contains(github.event.pull_request.labels.*.name, 'run-gpu-tests')
+    runs-on: ubuntu-latest
+    timeout-minutes: 90
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - run: pip install "modal>=1.0"
+      - name: Run GPU suite on Modal
+        run: |
+          ARGS="--gpu ${{ inputs.gpu || 'T4' }}"
+          if [ "${{ inputs.slow }}" = "true" ]; then ARGS="$ARGS --slow"; fi
+          modal run scripts/modal_gpu_tests.py $ARGS

.gitignore

@@ -190,3 +190,7 @@ venv
 upcoming.md
 post.md
 modal_custom_workflow.py
+hub_staging/
+gpu-test-results/
+gpu-probe-*/
+gpu-test-results-*/

.pre-commit-config.yaml

@@ -29,3 +29,4 @@ repos:
     hooks:
       - id: ruff
         args: ["--fix"]
+      - id: ruff-format

CLAUDE.md

@@ -157,6 +157,7 @@ quantllm/
 
 - Tests use CPU-only PyTorch (installed via `--index-url https://download.pytorch.org/whl/cpu` in CI)
 - Test `QUANTLLM_BANNER=0` to suppress banner
+- **GPU end-to-end suite** lives in `tests/gpu/` (real model, default `Qwen/Qwen2.5-0.5B-Instruct`, override with `QUANTLLM_GPU_TEST_MODEL`). It skips itself without CUDA. Free option: `python scripts/kaggle_gpu_tests.py [--slow] [-k expr]` packs the working tree, runs it on Kaggle "GPU T4 x2", and downloads results to `gpu-test-results/` (needs `~/.kaggle/access_token` and a phone-verified account). T4 has no native bf16/FlashAttention-2, so those paths need an Ampere+ GPU. ONNX tests skip on transformers 5.x (optimum-onnx caps transformers<4.58). Paid/credit option: `modal run scripts/modal_gpu_tests.py [--gpu T4|L4|A10G|H100] [--slow] [--k Finetune]`, on any NVIDIA box/Colab/Kaggle with `bash scripts/run_gpu_tests.sh`, or in GitHub Actions via the `GPU E2E` workflow (manual or `run-gpu-tests` PR label; needs `MODAL_TOKEN_ID`/`MODAL_TOKEN_SECRET` secrets)
 - Test files mirror core modules: `test_hqq.py`, `test_cli.py`, `test_benchmark.py`, `test_server.py`, `test_registry.py`, `test_cpu_kernels.py`, `test_generation.py`, `test_quantization_state.py`, `test_resolve_model_type.py`, `test_export_push_config.py`, `test_architecture_fallback.py`
 
 ## Common Tasks

README.md

@@ -189,6 +189,7 @@ Complete end-to-end pipeline demonstrating QuantLLM on a new/public model, run o
 ### 1. Load & Auto-Quantize
 ```python
 from quantllm import turbo
+
 model = turbo("Qwen/Qwen2.5-7B-Instruct")  # 7.62B params, BnB 8-bit on H100
 ```
 

mypy.ini

@@ -1,5 +1,5 @@
 [mypy]
-python_version = 3.10
+python_version = 3.12
 ignore_missing_imports = true
 disallow_untyped_defs = false
 warn_unused_ignores = false

pyproject.toml

@@ -33,11 +33,11 @@ keywords = [
 ]
 dependencies = [
     "torch>=2.0.0",
-    "transformers>=4.36.0",
+    "transformers>=4.46.0",  # Trainer(processing_class=...)
     "datasets>=2.14.0",
     "accelerate>=0.24.0",
-    "peft>=0.6.0",
-    "bitsandbytes>=0.41.0",
+    "peft>=0.12.0",
+    "bitsandbytes>=0.43.0; platform_system != 'Darwin'",
     "scipy>=1.10.0",
     "scikit-learn>=1.3.0",
     "tqdm>=4.65.0",
@@ -47,7 +47,6 @@ dependencies = [
     "einops>=0.7.0",
     "psutil>=5.9.0",
     "huggingface-hub>=0.19.0",
-    "rich>=13.0.0",
     "gguf>=0.10.0",
 ]
 
@@ -56,30 +55,35 @@ gguf = ["gguf>=0.10.0"]
 onnx = [
     "onnx>=1.14.0",
     "onnxruntime>=1.16.0",
-    "optimum[onnxruntime]>=1.14.0",
+    # ONNX Runtime export lives in optimum-onnx; it caps transformers<4.58 for now
+    "optimum-onnx[onnxruntime]>=0.1.0",
     "onnxscript>=0.1.0",
 ]
 mlx = ["mlx>=0.1.0", "mlx-lm>=0.1.0"]
 triton = ["triton>=2.1.0"]
 flash = ["flash-attn>=2.3.0"]
-hub = ["hf-lifecycle @ git+https://github.com/codewithdark-git/huggingface-lifecycle.git"]
+# NOTE: the former `hub` extra pointed at a git URL (hf-lifecycle), which PyPI
+# rejects on upload. Install it manually if needed:
+#   pip install "hf-lifecycle @ git+https://github.com/codewithdark-git/huggingface-lifecycle.git"
 full = [
     "gguf>=0.10.0",
     "onnx>=1.14.0",
     "onnxruntime>=1.16.0",
-    "optimum[onnxruntime]>=1.14.0",
+    "optimum-onnx[onnxruntime]>=0.1.0",
     "onnxscript>=0.1.0",
     "triton>=2.1.0",
     "wandb>=0.15.0",
     "tensorboard>=2.14.0",
     "evaluate>=0.4.0",
     "flash-attn>=2.3.0",
+    "rich>=13.0.0",
 ]
 server = ["fastapi>=0.100.0", "uvicorn>=0.23.0", "pydantic>=2.0.0"]
 bench = ["pandas>=1.5.0"]
 dev = [
     "pytest>=7.4.0",
     "pytest-cov>=4.1.0",
+    "httpx>=0.24.0",
     "black>=23.0.0",
     "isort>=5.12.0",
     "flake8>=6.1.0",
@@ -138,33 +142,20 @@ ignore = [
     "E741",   # ambiguous variable names
 ]
 
-[tool.myft.lint.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "tests/*" = ["F401", "F811", "F841"]
 "quantllm/__init__.py" = ["F401", "E402"]
 "quantllm/core/__init__.py" = ["F401"]
 
-[tool.mypy]
-python_version = "3.12"
-ignore_missing_imports = true
-disallow_untyped_defs = false
-warn_unused_ignores = false
-
-[[tool.mypy.overrides]]
-module = "numpy"
-ignore_errors = true
-follow_imports = "skip"
-
-[[tool.mypy.overrides]]
-module = "numpy.*"
-ignore_errors = true
-follow_imports = "skip"
+# mypy is configured in mypy.ini (single source of truth).
 
 [tool.pytest.ini_options]
 minversion = "7.0"
 testpaths = ["tests"]
 addopts = "-ra --strict-markers"
 markers = [
     "integration: marks tests as integration tests requiring GPU",
+    "gpu: end-to-end tests under tests/gpu (need CUDA; run via scripts/modal_gpu_tests.py)",
     "slow: marks tests as slow",
     "large_model: marks tests requiring large models",
     "stress: marks tests as stress/performance tests",

quantllm/benchmark/auto_bench.py

@@ -86,19 +86,25 @@ def _micro_bench(self) -> Dict[str, float]:
             times = []
             for n in [1024, 2048]:
                 a, b = torch.randn(n, n, device="cuda"), torch.randn(n, n, device="cuda")
-                start = time.time()
+                for _ in range(3):  # warmup (kernel selection / lazy init)
+                    a @ b
+                torch.cuda.synchronize()
+                start = time.perf_counter()
                 for _ in range(10):
                     a @ b
                 torch.cuda.synchronize()
-                times.append((time.time() - start) / 10)
+                times.append((time.perf_counter() - start) / 10)
             results["matmul_ms"] = sum(times) / len(times) * 1000
 
             x = torch.randn(1024, 1024, device="cuda")
-            start = time.time()
+            for _ in range(3):  # warmup
+                x.clone()
+            torch.cuda.synchronize()
+            start = time.perf_counter()
             for _ in range(100):
                 x.clone()
             torch.cuda.synchronize()
-            results["mem_bw_gb_s"] = (1024 * 1024 * 4 * 2 * 100) / (time.time() - start) / (1024**3)
+            results["mem_bw_gb_s"] = (1024 * 1024 * 4 * 2 * 100) / (time.perf_counter() - start) / (1024**3)
         return results
 
     def _score_quant(self, micro: Dict[str, float]) -> List[QuantScore]:

quantllm/benchmark/compare.py

@@ -6,7 +6,7 @@
 
 from typing import Any, Dict
 
-from .runner import BenchmarkRunner, BenchmarkResult
+from .runner import BenchmarkRunner, BenchmarkResult, release_memory
 
 
 class BenchmarkComparer:
@@ -27,6 +27,9 @@ def run_comparison(self, prompt: str = "The meaning of life is", max_new_tokens:
         runner = BenchmarkRunner(verbose=self.verbose)
 
         for name, bits, _ in configs:
+            # Drop the previous config's model before loading the next one (otherwise VRAM/peak stats accumulate).
+            m = None
+            release_memory()
             try:
                 m = turbo(self.model_name, bits=bits, verbose=False)
                 r = runner.benchmark_model(m.model, m.tokenizer, m.config, prompt=prompt, max_new_tokens=max_new_tokens)
@@ -37,6 +40,8 @@ def run_comparison(self, prompt: str = "The meaning of life is", max_new_tokens:
                     from ..utils import print_warning
 
                     print_warning(f"{name} failed: {e}")
+        m = None
+        release_memory()
         return self.results
 
     def generate_comparison_report(self) -> str:

quantllm/benchmark/runner.py

@@ -49,6 +49,36 @@ def get(self, key: str, default=None):
         return getattr(self, key, default)
 
 
+def _sync() -> None:
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def release_memory() -> None:
+    """Free cached CUDA memory and reset the peak-memory counter (call after dropping the previous model)."""
+    import gc
+
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+
+
+def _timed_forward(model, input_ids, num_runs: int) -> float:
+    """Average wall time (s) of a synchronized forward pass, after one warmup pass."""
+    runs = max(num_runs, 1)
+    times = []
+    with torch.inference_mode():
+        model(input_ids=input_ids)
+        for _ in range(runs):
+            _sync()
+            start = time.perf_counter()
+            model(input_ids=input_ids)
+            _sync()
+            times.append(time.perf_counter() - start)
+    return sum(times) / len(times)
+
+
 class BenchmarkRunner:
     def __init__(self, verbose: bool = True):
         self.verbose = verbose
@@ -76,9 +106,11 @@ def benchmark_model(
             pass
         r.model_size_gb = r.num_params_billion * (r.bits / 8) * 1.05
 
+        release_memory()  # reset the peak counter so VRAM reflects this model only
+
         inputs = tokenizer(prompt, return_tensors="pt")
         input_ids = inputs["input_ids"].to(model.device)
-        prefill_len = input_ids.shape[1]
+        prompt_len = input_ids.shape[1]
 
         for _ in range(num_warmup):
             with torch.inference_mode():
@@ -90,31 +122,36 @@ def benchmark_model(
                     eos_token_id=tokenizer.eos_token_id,
                 )
 
-        start = time.time()
-        with torch.inference_mode():
-            model.generate(
-                input_ids,
-                max_new_tokens=1,
-                do_sample=False,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-            )
-        r.prefill_tokens_per_sec = prefill_len / (time.time() - start) if (time.time() - start) > 0 else 0
-
-        decode_times = []
+        # Prefill: a plain forward pass over a ~512-token prompt (the short prompt repeated).
+        max_pos = getattr(getattr(model, "config", None), "max_position_embeddings", None) or 512
+        prefill_target = min(512, max_pos)
+        reps = prefill_target // max(prompt_len, 1) + 1
+        long_ids = tokenizer(" ".join([prompt] * reps), return_tensors="pt")["input_ids"][:, :prefill_target]
+        long_ids = long_ids.to(model.device)
+        prefill_time = _timed_forward(model, long_ids, num_runs)
+        r.prefill_tokens_per_sec = long_ids.shape[1] / prefill_time if prefill_time > 0 else 0
+
+        # Decode: count tokens actually generated and exclude the prefill of the benchmark prompt.
+        short_prefill_time = _timed_forward(model, input_ids, num_runs)
+        decode_rates = []
         for _ in range(num_runs):
-            start = time.time()
+            _sync()
+            start = time.perf_counter()
             with torch.inference_mode():
-                model.generate(
+                out = model.generate(
                     input_ids,
                     max_new_tokens=max_new_tokens,
                     do_sample=False,
                     pad_token_id=tokenizer.pad_token_id,
                     eos_token_id=tokenizer.eos_token_id,
                 )
-            elapsed = time.time() - start
-            decode_times.append(max_new_tokens / elapsed if elapsed > 0 else 0)
-        r.decode_tokens_per_sec = sum(decode_times) / len(decode_times)
+            _sync()
+            elapsed = time.perf_counter() - start
+            new_tokens = out.shape[1] - prompt_len
+            decode_time = elapsed - short_prefill_time
+            # The first new token comes from the prefill pass; the remaining ones are decode steps.
+            decode_rates.append((new_tokens - 1) / decode_time if new_tokens > 1 and decode_time > 0 else 0)
+        r.decode_tokens_per_sec = sum(decode_rates) / len(decode_rates)
         r.total_latency_ms = (1000.0 / r.decode_tokens_per_sec) if r.decode_tokens_per_sec > 0 else 0
         r.throughput_tokens_per_sec = r.decode_tokens_per_sec