diff --git a/.github/workflows/windows-vulkan.yml b/.github/workflows/windows-vulkan.yml new file mode 100644 index 0000000..309eba9 --- /dev/null +++ b/.github/workflows/windows-vulkan.yml @@ -0,0 +1,269 @@ +name: Windows Vulkan build + +on: + push: + branches: ["main"] + paths: + - ".github/workflows/windows-vulkan.yml" + - "native-kernels/orbitquant-vulkan/**" + pull_request: + paths: + - ".github/workflows/windows-vulkan.yml" + - "native-kernels/orbitquant-vulkan/**" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: windows-vulkan-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +env: + EXECUTORCH_COMMIT: a6d812a082df57898b8608f56c867140cc9da32c + ORBITQUANT_SOURCE_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + +jobs: + build: + name: MSVC Vulkan source build + runs-on: windows-2025 + timeout-minutes: 90 + steps: + - name: Check out OrbitQuant + uses: actions/checkout@v7.0.0 + + - name: Set up Python + uses: actions/setup-python@v6.3.0 + with: + python-version: "3.12" + + - name: Build the ExecuTorch and OrbitQuant Vulkan targets + shell: pwsh + run: | + $ErrorActionPreference = "Stop" + $PSNativeCommandUseErrorActionPreference = $true + + $executorch = Join-Path $env:RUNNER_TEMP "executorch" + $executorchBuild = Join-Path $env:RUNNER_TEMP "executorch-build" + $executorchInstall = Join-Path $env:RUNNER_TEMP "executorch-install" + $orbitquantBuild = Join-Path $env:RUNNER_TEMP "orbitquant-vulkan-build" + $python = Join-Path $env:pythonLocation "python.exe" + + git config --global core.autocrlf false + git config --global core.symlinks true + git clone --filter=blob:none --no-checkout ` + https://github.com/pytorch/executorch.git $executorch + git -C $executorch fetch --depth=1 origin $env:EXECUTORCH_COMMIT + git -C $executorch config core.autocrlf false + git -C $executorch config core.symlinks true + git -C $executorch checkout --detach $env:EXECUTORCH_COMMIT + if ((git -C $executorch rev-parse HEAD) -ne $env:EXECUTORCH_COMMIT) { + throw "ExecuTorch checkout does not match the pinned commit" + } + git -C $executorch submodule update --init ` + backends/vulkan/third-party/Vulkan-Headers ` + backends/vulkan/third-party/VulkanMemoryAllocator ` + backends/vulkan/third-party/volk + git -C $executorch submodule update --init + + $env:PATH = ` + "$env:CONDA;$env:CONDA\Scripts;$env:CONDA\Library\bin;$env:PATH" + Push-Location $executorch + & .ci\scripts\setup-vulkan-windows-deps.ps1 + Pop-Location + $glslc = (Get-Command glslc.exe -ErrorAction Stop).Source + & $python -m pip install PyYAML==6.0.1 + & $python -m pip install torch==2.12.1 + $env:PYTHONPATH = ` + "$executorch\src;$executorch;$(Split-Path $executorch -Parent)" + & $python -c ` + "import executorch.codegen, torchgen; print('codegen imports ok')" + + $vswhere = ` + "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" + $vsRoot = & $vswhere ` + -latest ` + -products * ` + -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 ` + -property installationPath + if (-not $vsRoot) { + throw "MSVC x86_64 tools were not found" + } + $vsVersion = & $vswhere ` + -latest ` + -products * ` + -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 ` + -property installationVersion + $vsMajor = [int](($vsVersion -split "\.")[0]) + $cmakeGenerator = switch ($vsMajor) { + 17 { "Visual Studio 17 2022" } + 18 { "Visual Studio 18 2026" } + default { throw "Unsupported Visual Studio major version $vsMajor" } + } + $vsDevCmd = Join-Path $vsRoot "Common7\Tools\VsDevCmd.bat" + $envDump = cmd.exe /d /s /c ` + "call `"$vsDevCmd`" -arch=x64 -host_arch=x64 >nul && set" + foreach ($line in $envDump) { + $parts = $line -split "=", 2 + if ($parts.Length -eq 2) { + Set-Item -Path "Env:$($parts[0])" -Value $parts[1] + } + } + Get-Command cl.exe -ErrorAction Stop | Format-List Source + Write-Host "Using CMake generator: $cmakeGenerator" + $glslcVersion = (& $glslc --version | Select-Object -First 1) + Write-Host $glslcVersion + cmake --version + + cmake -S $executorch -B $executorchBuild ` + -G $cmakeGenerator -A x64 ` + -DCMAKE_BUILD_TYPE=Release ` + "-DCMAKE_INSTALL_PREFIX=$executorchInstall" ` + "-DPYTHON_EXECUTABLE=$python" ` + -DEXECUTORCH_BUILD_VULKAN=ON ` + -DEXECUTORCH_BUILD_PORTABLE_OPS=ON ` + -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=OFF ` + -DEXECUTORCH_BUILD_TESTS=OFF ` + -DEXECUTORCH_BUILD_PYBIND=OFF ` + -DEXECUTORCH_BUILD_XNNPACK=OFF ` + "-DGLSLC_PATH=$glslc" + cmake --build $executorchBuild --config Release ` + --target install --parallel 4 + + cmake -S native-kernels\orbitquant-vulkan -B $orbitquantBuild ` + -G $cmakeGenerator -A x64 ` + -DCMAKE_BUILD_TYPE=Release ` + "-DCMAKE_PREFIX_PATH=$executorchInstall" ` + "-DPYTHON_EXECUTABLE=$python" ` + "-DORBITQUANT_EXECUTORCH_ROOT=$executorch" ` + "-DGLSLC_PATH=$glslc" ` + -DORBITQUANT_VULKAN_BUILD_TESTS=ON ` + -DORBITQUANT_VULKAN_BUILD_PTE_RUNNER=ON + cmake --build $orbitquantBuild --config Release ` + --target orbitquant_vulkan_test orbitquant_vulkan_pte_runner ` + --parallel 4 + + Get-Item ` + "$orbitquantBuild\Release\orbitquant_vulkan_test.exe", ` + "$orbitquantBuild\Release\orbitquant_vulkan_pte_runner.exe" + + $bundle = Join-Path ` + $env:RUNNER_TEMP "orbitquant-vulkan-windows-x86_64" + New-Item -ItemType Directory -Force -Path $bundle | Out-Null + Copy-Item ` + "$orbitquantBuild\Release\orbitquant_vulkan_test.exe", ` + "$orbitquantBuild\Release\orbitquant_vulkan_pte_runner.exe" ` + -Destination $bundle + + $runtimeDirs = @( + "$orbitquantBuild\Release", + "$executorchInstall\bin" + ) + foreach ($runtimeDir in $runtimeDirs) { + if (Test-Path $runtimeDir) { + Get-ChildItem $runtimeDir -Filter "*.dll" -File | + Copy-Item -Destination $bundle -Force + } + } + + $licenseFiles = @( + @{ + Source = "$env:GITHUB_WORKSPACE\LICENSE" + Destination = "LICENSE-OrbitQuant.txt" + }, + @{ + Source = "$executorch\LICENSE" + Destination = "LICENSE-ExecuTorch.txt" + }, + @{ + Source = "$executorch\third-party\flatbuffers\LICENSE" + Destination = "LICENSE-FlatBuffers.txt" + }, + @{ + Source = "$executorch\third-party\gflags\COPYING.txt" + Destination = "LICENSE-gflags.txt" + }, + @{ + Source = "$executorch\backends\vulkan\third-party\Vulkan-Headers\LICENSE.md" + Destination = "LICENSE-Vulkan-Headers.md" + }, + @{ + Source = "$executorch\backends\vulkan\third-party\VulkanMemoryAllocator\LICENSE.txt" + Destination = "LICENSE-VulkanMemoryAllocator.txt" + }, + @{ + Source = "$executorch\backends\vulkan\third-party\volk\LICENSE.md" + Destination = "LICENSE-volk.md" + } + ) + foreach ($licenseFile in $licenseFiles) { + if (-not (Test-Path $licenseFile.Source)) { + throw "Required license is missing: $($licenseFile.Source)" + } + Copy-Item ` + $licenseFile.Source ` + (Join-Path $bundle $licenseFile.Destination) + } + + @( + "Install the current Microsoft Visual C++ v14 x64 Redistributable:", + "https://aka.ms/vc14/vc_redist.x64.exe", + "", + "Prerequisite: install the hardware vendor Vulkan driver.", + "The bundle intentionally does not ship vulkan-1.dll.", + "", + "Run native correctness and realistic-shape benchmarks:", + ".\orbitquant_vulkan_test.exe", + "", + "Run an exported PTE with one raw tensor input:", + ".\orbitquant_vulkan_pte_runner.exe --model_path model.pte --inputs input0.bin --output_file actual --print_output=none", + "The output tensor is written to actual-0.bin." + ) | Set-Content ` + -Encoding ascii (Join-Path $bundle "README.txt") + + $dependencyReport = Join-Path $bundle "DEPENDENCIES.txt" + foreach ($executable in Get-ChildItem $bundle -Filter "*.exe") { + "[$($executable.Name)]" | Add-Content $dependencyReport + dumpbin.exe /dependents $executable.FullName | + Add-Content $dependencyReport + } + @( + "orbitquant_source_commit=$env:ORBITQUANT_SOURCE_SHA", + "orbitquant_checkout_commit=$(git rev-parse HEAD)", + "executorch_commit=$env:EXECUTORCH_COMMIT", + "cmake_generator=$cmakeGenerator", + "visual_studio_version=$vsVersion", + "python_version=3.12", + "pytorch_version=2.12.1", + "glslc_version=$glslcVersion" + ) | Set-Content ` + -Encoding ascii (Join-Path $bundle "BUILD-METADATA.txt") + $checksumPath = Join-Path $bundle "SHA256SUMS" + Get-ChildItem $bundle -File -Recurse | + Where-Object { $_.FullName -ne $checksumPath } | + Sort-Object FullName | + ForEach-Object { + $hash = (Get-FileHash $_.FullName -Algorithm SHA256).Hash + $relativePath = [IO.Path]::GetRelativePath( + $bundle, $_.FullName + ).Replace("\", "/") + "$hash $relativePath" + } | Set-Content -Encoding ascii $checksumPath + + $bundleBytes = ( + Get-ChildItem $bundle -File -Recurse | Measure-Object Length -Sum + ).Sum + if ($bundleBytes -gt 64MB) { + throw "Windows Vulkan runner bundle exceeds 64 MiB: $bundleBytes" + } + Write-Host "Windows Vulkan runner bundle bytes: $bundleBytes" + Get-ChildItem $bundle | Format-Table Name, Length + + - name: Upload Windows Vulkan runner bundle + uses: actions/upload-artifact@v7.0.1 + with: + name: orbitquant-vulkan-windows-x86_64-${{ env.ORBITQUANT_SOURCE_SHA }} + path: ${{ runner.temp }}/orbitquant-vulkan-windows-x86_64 + if-no-files-found: error + retention-days: 14 + compression-level: 9 diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..1f9cbc0 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,79 @@ +# CLAUDE.md + +Guidance for coding agents working in this repository (OrbitQuant: calibration-free +post-training quantizer for transformer linear projections). + +## Environment + +```bash +uv sync --extra hf --extra dev +``` + +- Tests: `uv run pytest -ra` +- Lint: `uv run ruff check .` (line-length 100) + +## Formatting + +- NEVER mass-reformat. `ruff format .` is forbidden: the tree has historical + formatting drift and a blanket reformat produces unreviewable diffs. +- Format only the lines you touch. + +## Gates Before Claiming Done + +Run all of these and confirm they pass: + +```bash +uv run ruff check . +uv run pytest -ra +scripts/run_paper_methodology_checks.sh +scripts/run_hf_compat_checks.sh --mode all +``` + +Backend gates (run on matching hardware only): + +```bash +scripts/run_cuda_kernel_checks.sh # CUDA host +PYTHON_BIN="$(uv python find)" scripts/run_mps_kernel_checks.sh # Apple Silicon +``` + +The MPS gate requires bash >= 4; on macOS use `/opt/homebrew/bin/bash`. + +## Native Kernels + +The native package lives in `native-kernels/orbitquant-packed-matmul` and uses the +Hugging Face kernel-builder, pinned to rev +`d43de01d0b43285d8e5061ca4380c2bd1c40ae3b` (see `.github/workflows/ci.yml`). + +Local build: + +```bash +cd native-kernels/orbitquant-packed-matmul +kernel-builder create-pyproject -f . +CMAKE_ARGS=-DGPU_LANG=CPU python setup.py build_kernel +``` + +or via Nix: + +```bash +nix --option sandbox relaxed run .#build-and-copy -L +``` + +Build artifacts inside the kernel directory are GENERATED — never commit them: +`CMakeLists.txt`, `cmake/`, `compat.py`, `metadata-*.json`, `pyproject.toml`, +`setup.py`, `build/`, `torch-ext/registration.h`, +`torch-ext/orbitquant_packed_matmul/_ops.py`, and `uv.lock` inside the kernel +directory. Delete them before running the repo test suite: two ABI3 lint tests +scan the kernel directory and fail on leftover generated files. + +## Documentation Policy + +- Docs describe the current artifact and user contract only. No development + chronicles, no raw logs. +- The source of truth for backend status is + `orbitquant.kernels.dispatch.backend_capabilities()`; the audited backend + contract is `docs/kernel-audit.md`. + +## Hard Boundaries + +- Hugging Face Kernel Hub publication is OUT OF SCOPE. Never post there. +- Do not commit, push, or publish anything without explicit operator instruction. diff --git a/README.md b/README.md index cde4717..4938880 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,12 @@ The implementation is clean-room and Apache-2.0 licensed. `from_pretrained()` integration. - Packed-weight CUDA, Triton, and Metal inference paths that avoid a full dequantized weight matrix. +- `torch.compile(fullgraph=True)` support: quantized forwards run behind a + registered custom op with a fake implementation, so compiled models match + eager output exactly. The forward is also CUDA-Graph-capturable; + `torch.compile(mode="reduce-overhead")` removes the per-launch overhead that + dominates small decode batches (measured 1.6-5.5x at 1-32 rows on + A40/RTX 4090). Embeddings, timestep modules, task heads, and common final projections are kept in source precision by default. Every automatic decision is available as diff --git a/docs/kernel-audit.md b/docs/kernel-audit.md index 22c3cd7..ee975dc 100644 --- a/docs/kernel-audit.md +++ b/docs/kernel-audit.md @@ -36,9 +36,9 @@ Other explicit modes are `native_packed_matmul`, `triton_packed_matmul`, | CUDA | Optimized packed inference | Native RPBH/quantization, chunked packed-weight decode plus CUTLASS INT8 matmul, direct packed CUDA MMA fallback, and generic Triton packed fallback | | MPS/Metal | Optimized packed inference | Native Metal packed matmul and Metal activation quantization stages | | CPU | Hardware-verified native source and CI wheel builds; wheel publication pending | Runtime ISA dispatch across scalar, AVX2/FMA, AVX-512/BF16, and ARM64 NEON; exact packed activation, packed low-bit matmul, and packed INT4 group-64 AdaLN | -| Vulkan | Unverified | No runtime backend | -| ROCm | Unsupported | No release backend | -| XPU | Unsupported | No release backend | +| Vulkan | Experimental source path; hardware-verified on AMD Cezanne/RADV Linux, Windows unverified | ExecuTorch whole-graph delegate adapter with exact packed W4A4 activation and matmul shaders | +| ROCm | Experimental source candidate; supported-hardware proof pending | Exact Triton activation, low-bit pack/unpack, packed matmul, weight quantization, and packed INT4 AdaLN; explicit-only | +| XPU | Experimental source candidate; Intel hardware proof pending | Exact Triton activation, low-bit pack/unpack, packed matmul, weight quantization, and packed INT4 AdaLN on `torch.xpu`; explicit-only | ## Local Native Package @@ -107,17 +107,297 @@ export LOCAL_KERNELS="WaveCut/orbitquant-packed-matmul=$PWD/build/` specialization. Against the generic eight-row baseline, +annotated disassembly reduced the hot body from 896 to 703 instructions, +increased `vdpbf16ps` sites from 8 to 30, and reduced ZMM stack references from +37 to zero. The provider exposed `perf_event_paranoid=4`, so hardware PMU +counters were unavailable; the evidence uses wall-clock distributions, native +stack sampling, and ISA disassembly instead. Two threads pinned to the same +physical core were 6.8% and 8.3% slower at dimensions 1536 and 3072, so one +worker per physical core is selected. This does not establish multi-core or +NUMA scaling. + +A native-only 32x9216 memory probe constructed the 42,467,328-byte packed +payload directly, without an unpacked index tensor or reference weight. The +first packed call increased process peak RSS by 2,129,920 bytes and 20 calls +increased it by 3,440,640 bytes; a materialized BF16 weight for the same +projection would require 169,869,312 bytes. The optimized call therefore +retained bounded output and allocator scratch instead of a full floating-point +weight. + +W2, W3, and W6 packed matmuls dispatch to dedicated SIMD decoders on every +tier: AVX-512 covers all three widths, AVX2 covers W2 (`vpermps` lookup) and +W6 (`vgatherdps` lookup) while W3 stays scalar there, and NEON covers all +three with vectorized FMA over scalar table decode. Row starts must be +byte-aligned (`in_features % 4 == 0` for W2/W6, `% 8 == 0` for W3); other +shapes keep the scalar fallback. Multi-threaded execution uses a lazily +created persistent worker pool instead of per-call thread spawning, with +outputs bitwise identical to single-threaded runs. When at least 16 activation +rows are present, every SIMD tier decodes each packed column once into a +thread-local buffer and reuses it across the row tile instead of re-decoding +per row. Measured with six unpinned worker threads, 32 BF16 rows, 21 timed +iterations (hot-loop medians), and square projections on an AMD EPYC 9354 +(Zen 4) Secure Cloud cpuset, and with one thread on an Apple M2 Max +aarch64-linux container: + +| Bits | Dimension | Scalar (Zen 4) | AVX2/FMA | AVX-512 | Scalar (M2) | NEON | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| W2 | 3072 / 1536 | 81.76 ms | 5.04 ms | 4.96 ms | 91.66 ms | 7.27 ms | +| W3 | 3072 / 1536 | 82.10 ms | scalar | 7.04 ms | 92.32 ms | 7.49 ms | +| W6 | 3072 / 1536 | 83.98 ms | 8.69 ms | 5.49 ms | 92.31 ms | 8.15 ms | +| W4 | 3072 / 1536 | 82.35 ms | 6.87 ms | 2.53 ms | 91.63 ms | 7.30 ms | + +Zen 4 columns use dimension 3072 and the M2 columns use dimension 1536. A +higher-clock EPYC 9684X cpuset measured the same AVX-512 shapes at +1.92/2.86/2.03 ms (W2/W3/W6); absolute numbers track the host, the tier +ordering does not. + +The AVX2-only tier was separately profiled on an AMD Ryzen 5 5600G Cezanne +(family 19h/model 50h), GCC 13.3, and Torch 2.11.0. Runtime dispatch uses a +16-row primary accumulator tile with 8-row and 4-row tails only for BF16 W4, +at least 16 rows, dimensions 1536, 1920, or 3072, and that exact verified +CPUID. Every other AVX2 processor and shape retains the generic eight-row +path. The exact Cezanne guard prevents a machine-specific result from changing +the portable AVX2 default. + +The finite same-host search compared 4, 8, 10, 12, and 16-row tiles, +`vgatherdps` versus two-register `vpermps` lookup, K-loop unrolling, software +prefetch, LTO, and `-mtune=znver3`. Gather lookup, four-row tiling, and software +prefetch were materially slower; unrolling and LTO did not improve both tested +dimensions. The selected 16/8/4 topology amortizes packed decode across more +activation rows without changing the artifact or numerical path. Disassembly +confirmed AVX2 `vpermps` lookup and FMA accumulation. PMU counters were +unavailable because the host sets `perf_event_paranoid=4`. + +The final ABBA repeat pinned one logical CPU with its SMT sibling idle and used +10 warmups plus 101 timed calls per leg. Values below average the two medians; +the selected p95 is the slower of its two legs. + +| Rows | Dimension | Eight-row median | Selected median | Selected p95 | Resident BF16 median | Median improvement | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 32 | 1536 | 3.9535 ms | 3.5735 ms | 3.6396 ms | 10.1302 ms | 9.61% | +| 32 | 1920 | 6.1341 ms | 5.5369 ms | 5.6402 ms | 15.6631 ms | 9.73% | +| 32 | 3072 | 15.5206 ms | 14.8918 ms | 15.4647 ms | 39.7343 ms | 4.05% | + +Activation plus matmul measured 3.7806, 5.7961, and 15.2280 ms at those +dimensions, improvements of 9.24%, 9.74%, and 4.91%. Matmul accounted for +94.5%, 95.5%, and 97.8% of the selected full-pipeline medians, identifying +packed decode/LUT/FMA as the remaining bottleneck rather than activation norm +or RPBH/FWHT. Relative RMSE against the independently materialized BF16 +reference was 0.002315, 0.002292, and 0.002250; the optimized and baseline +AVX2 paths had the same error. + +Using the six physical cores, with the container cpuset restricted to one +hardware thread per core, improved packed-matmul medians by 4.43x, 4.87x, and +5.33x. Twelve SMT workers were slower than six physical-core workers on every +shape. The host also serves production workloads, so the measured p95 includes +their scheduling interference rather than representing an isolated-server +tail-latency claim. + +| Dimension | One thread | Six physical cores | Six-core p95 | Twelve SMT threads | +| ---: | ---: | ---: | ---: | ---: | +| 1536 | 3.5735 ms | 0.8063 ms | 1.4194 ms | 1.0395 ms | +| 1920 | 5.5369 ms | 1.1366 ms | 2.1565 ms | 1.2310 ms | +| 3072 | 14.8918 ms | 2.7934 ms | 5.5808 ms | 3.9207 ms | + +Independent hardware checks covered exact activation RPBH/FWHT and exhaustive +centroid assignment, packed W2/W3/W4/W6 with bias, realistic 16/24/32/33-row +W4 shapes, and packed INT4 group-64 AdaLN. A native-only 32x9216 probe used a +42,467,328-byte packed payload and 18,432 bytes of row norms. Its peak RSS grew +by 3,796,992 bytes across 20 calls, versus 169,869,312 bytes for a full BF16 +weight, so the Cezanne specialization retains bounded scratch and does not +materialize a floating-point weight matrix. For an exact W4A4 projection with 32 activation rows and dimension 1536, the native activation stage measured 0.1484 ms median, packed matmul measured @@ -164,17 +444,17 @@ under Torch 2.12.1. This verifies the current LibTorch Stable ABI boundary for those two runtimes; it does not make the wheel independent of the platform GLIBC, C++ runtime, or CPU ISA. -CI run `29162193254` separately verified installable Python 3.9 ABI3 wheels on -Linux and Windows. The 101,453-byte Linux wheel has both -`manylinux_2_24_x86_64` and `manylinux_2_28_x86_64` tags, passed strict +CI run `29162460269` separately verified installable Python 3.9 ABI3 wheels on +Linux and Windows at source commit `769c9b8`. The 101,455-byte Linux wheel has +both `manylinux_2_24_x86_64` and `manylinux_2_28_x86_64` tags, passed strict `abi3audit` with a computed Python 3.9 baseline, and has SHA-256 -`5cd19321e11281e03b9b1126f40571ad2c1e956060fc0bc5173c90628045d4df`. +`5458d49de7cb6b15db7f88313d96cb34827dea2b3d4cc21a39c9574f3c928b16`. Inside the manylinux container, 73 native tests passed and 33 CUDA/MPS tests were skipped; a separate clean environment with OrbitQuant 0.4.0 and Torch 2.12.1+cpu passed 33 integration/oracle tests with two CUDA/Triton skips. -The 53,178-byte Windows wheel has the `cp39-abi3-win_amd64` tag and SHA-256 -`cf200774d7388496ce3b7f554dd3e757eca470d99ef9986fe7a7bb0ea1e0e297`. +The 52,805-byte Windows wheel has the `cp39-abi3-win_amd64` tag and SHA-256 +`832aa3e89836c3ec8de44a0ec3bd160fe90e162d8a89c7cdca76c657df2e6728`. It was built with MSVC 14.51 on Windows Server 2025 and executed on an AMD EPYC 9V74 runner with two cores and four logical processors. Its clean native suite passed 73 tests with 33 CUDA/MPS skips, and the OrbitQuant integration/oracle diff --git a/docs/paper-methodology-audit.md b/docs/paper-methodology-audit.md index b2b4653..a234f59 100644 --- a/docs/paper-methodology-audit.md +++ b/docs/paper-methodology-audit.md @@ -171,11 +171,15 @@ metrics; they do not by themselves claim GenEval or VBench scores. | Backend | Status | Evidence | Claim boundary | | --- | --- | --- | --- | -| CPU | Pass as reference | `src/orbitquant/kernels/dispatch.py`, `src/orbitquant/functional.py` | Correctness baseline only; no optimized CPU kernel claim. | +| CPU | Pass for native packed inference when the optional CPU package is importable; reference fallback otherwise | `native-kernels/orbitquant-packed-matmul`, `src/orbitquant/kernels/dispatch.py`, `src/orbitquant/kernels/native_packed_matmul.py`, `src/orbitquant/functional.py`, `docs/kernel-audit.md` | Runtime ISA dispatch covers scalar, AVX2/FMA, AVX-512/BF16, and ARM64 NEON for exact packed activation, packed low-bit matmul, and packed INT4 group-64 AdaLN. Native builds are hardware-verified with CI wheels; wheel publication is pending. Without the native package, CPU remains the correctness reference baseline. | | CUDA/native/Triton | Pass for optimized W4A4 and packed fallback | `native-kernels/orbitquant-packed-matmul`, `src/orbitquant/kernels/native_packed_matmul.py`, `src/orbitquant/kernels/triton_cuda.py`, `tests/test_native_packed_matmul.py`, `tests/test_kernels.py`, `tests/test_orbit_linear.py`, `docs/kernel-audit.md` | The selected W4A4 path fuses norm/RPBH/FWHT/codebook assignment to INT8 surrogate activations, decodes bounded W4 chunks, uses CUTLASS INT8 matmul, and applies a fused epilogue. Direct packed CUDA MMA and generic Triton packed matmul remain fallbacks. No path selected by `auto_fused` materializes a full BF16/FP16 weight matrix. | | MPS/Metal | Pass for native packed inference | `src/orbitquant/kernels/mps.py`, `src/orbitquant/kernels/dispatch.py`, `tests/test_kernels.py`, `tests/test_orbit_linear.py`, `docs/kernel-audit.md` | A fused Metal shader performs activation norm, RPBH/FWHT, codebook lookup, and rescale. The native package performs packed matmul for generic leading dimensions, including short decode rows and partial matrix tiles, without full weight materialization. Offline weight and AdaLN quantization remain reference paths on MPS. | -| ROCm | Blocked for backend claim | No implementation in current tree | Do not claim ROCm optimization. | -| XPU | Blocked for backend claim | No implementation in current tree | Do not claim XPU optimization. | +| ROCm | Experimental source candidate; supported-hardware proof pending | `src/orbitquant/kernels/triton_cuda.py`, `src/orbitquant/kernels/dispatch.py`, `docs/kernel-audit.md` | The exact Triton pipeline is reused on HIP, explicit-only, never selected by `auto_fused`, and reports `experimental_unverified`. Do not claim verified ROCm optimization without evidence on supported AMD hardware. | +| XPU | Experimental source candidate; Intel hardware proof pending | `src/orbitquant/kernels/triton_cuda.py`, `src/orbitquant/kernels/dispatch.py`, `docs/kernel-audit.md` | The exact Triton pipeline is reused on `torch.xpu`, explicit-only, never selected by `auto_fused`, and reports `experimental_unverified`. Do not claim verified XPU optimization without evidence on real Intel GPU hardware. | + +Backend status is canonically reported by +`orbitquant.kernels.dispatch.backend_capabilities()` and audited in +[`docs/kernel-audit.md`](kernel-audit.md). ## Acceleration Claim Boundary @@ -229,7 +233,7 @@ not silently substitute a different quantization method. | Published checkpoints use converged Lloyd-Max codebook version 2 and `activation_eps=1e-10`. | Pass | All 14 canonical FLUX.2 4B, FLUX.1-schnell, Z-Image-Turbo, and Wan2.1 manifests expose these values. The FLUX.2 Klein 9B comparison checkpoint records them in both quantized component configs. Legacy version 1 artifacts remain loadable but are no longer the published release checkpoints. | | Full config-derived inventories are audit artifacts, not committed source files. | Accepted artifact hygiene choice | Inventory summaries are recorded above; raw JSON may remain unpublished to avoid turning the repository into an artifact store. | | Release-grade GenEval/VBench metrics are required only for metric claims. | Accepted claim boundary | Missing full metrics block paper metric/reproduction claims only. | -| ROCm and XPU kernels are not implemented. | Backend claim blocker | The release must either implement and verify them or explicitly exclude them. | +| ROCm and XPU are experimental source candidates without supported-hardware proof. | Accepted claim boundary | Both reuse the exact Triton pipeline, stay explicit-only, and report `experimental_unverified`. Verified backend claims require correctness, memory, profiler, and performance evidence on supported AMD or Intel hardware. | GenEval/VBench scores are not claimed without their corresponding external metric runs. Any changed checkpoint must pass manifest, policy-inventory, diff --git a/native-kernels/orbitquant-packed-matmul/CARD.md b/native-kernels/orbitquant-packed-matmul/CARD.md index 8196fa8..3376f7a 100644 --- a/native-kernels/orbitquant-packed-matmul/CARD.md +++ b/native-kernels/orbitquant-packed-matmul/CARD.md @@ -138,9 +138,15 @@ python benchmarks/benchmark.py \ --iters 20 ``` +`--rows` accepts a comma-separated sweep (the default covers decode-bound +small batches and GEMM-bound large batches), and `--dtype` selects the +activation dtype explicitly. Headline timings are hot-loop medians. + The script prints JSON with `packed_seconds_per_iter`, `predequantized_f_linear_seconds_per_iter`, -`dequantize_then_f_linear_seconds_per_iter`, +`dequantize_then_f_linear_seconds_per_iter` (all hot-loop medians), the +per-path `*_hot_mean_seconds`, `*_hot_median_seconds`, and +`*_hot_p95_seconds` distributions, `packed_vs_predequantized_f_linear_speedup`, `packed_vs_dequantize_then_f_linear_speedup`, compatibility aliases `reference_seconds_per_iter` and `packed_vs_reference_speedup`, and diff --git a/native-kernels/orbitquant-packed-matmul/benchmarks/benchmark.py b/native-kernels/orbitquant-packed-matmul/benchmarks/benchmark.py index 0e32038..8a0a0aa 100644 --- a/native-kernels/orbitquant-packed-matmul/benchmarks/benchmark.py +++ b/native-kernels/orbitquant-packed-matmul/benchmarks/benchmark.py @@ -51,48 +51,37 @@ def _time_distribution(device: str, iters: int, fn) -> dict[str, float]: } -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--device", choices=["cpu", "cuda", "mps"], default="cuda") - parser.add_argument("--bits", type=int, default=4) - parser.add_argument("--rows", type=int, default=4096) - parser.add_argument("--in-features", type=int, default=3072) - parser.add_argument("--out-features", type=int, default=3072) - parser.add_argument("--iters", type=int, default=20) - parser.add_argument("--warmup", type=int, default=3) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--threads", type=int, default=0) - parser.add_argument("--with-bias", action="store_true") - args = parser.parse_args() +def _parse_rows(raw: str) -> list[int]: + rows_values = [] + for chunk in raw.split(","): + chunk = chunk.strip() + if not chunk: + continue + rows = int(chunk) + if rows <= 0: + raise argparse.ArgumentTypeError("--rows values must be positive") + rows_values.append(rows) + if not rows_values: + raise argparse.ArgumentTypeError("--rows must list at least one row count") + return rows_values - if args.threads < 0: - parser.error("--threads must be non-negative") - if args.iters <= 0 or args.warmup < 0: - parser.error("--iters must be positive and --warmup must be non-negative") - if args.device == "cpu" and args.threads > 0: - os.environ["ORBITQUANT_CPU_THREADS"] = str(args.threads) - torch.set_num_threads(args.threads) - torch.manual_seed(args.seed) - dtype = torch.float16 if args.device == "mps" else torch.bfloat16 - x = torch.randn(args.rows, args.in_features, device=args.device, dtype=dtype) - indices = torch.randint( - 0, - 2**args.bits, - (args.out_features, args.in_features), - dtype=torch.uint8, - ) - packed = _pack(indices, args.bits).to(args.device) - row_norms = torch.ones(args.out_features, device=args.device, dtype=torch.bfloat16) - centroids = torch.linspace(-1.0, 1.0, 2**args.bits, device=args.device) +_DTYPES = { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, +} + + +def _benchmark_rows(args, rows: int, dtype: torch.dtype, weights) -> dict: + packed, row_norms, centroids, indices_device = weights + x = torch.randn(rows, args.in_features, device=args.device, dtype=dtype) bias = ( torch.randn(args.out_features, device=args.device, dtype=dtype) if args.with_bias else None ) - indices_device = indices.long().to(args.device) - def materialize_reference_weight() -> torch.Tensor: return (row_norms[:, None] * centroids[indices_device]).to(dtype) @@ -140,9 +129,11 @@ def dequantize_then_linear_call() -> torch.Tensor: args.iters, dequantize_then_linear_call, ) - packed_seconds = packed_distribution["mean"] - predequantized_linear_seconds = predequantized_distribution["mean"] - dequantize_then_linear_seconds = dequantize_then_distribution["mean"] + # Headline numbers are hot-loop medians; the mean is retained alongside the + # median/p95 so noisy first-iteration outliers cannot skew comparisons. + packed_seconds = packed_distribution["median"] + predequantized_linear_seconds = predequantized_distribution["median"] + dequantize_then_linear_seconds = dequantize_then_distribution["median"] packed_output = packed_call() reference_output = predequantized_linear_call() @@ -153,7 +144,7 @@ def dequantize_then_linear_call() -> torch.Tensor: reference_rms = reference_output.float().square().mean().sqrt().item() relative_rmse = rmse / max(reference_rms, 1e-12) - payload = { + return { "device": args.device, "device_name": ( torch.cuda.get_device_name(0) @@ -165,7 +156,7 @@ def dequantize_then_linear_call() -> torch.Tensor: ), "dtype": str(dtype).replace("torch.", ""), "bits": args.bits, - "rows": args.rows, + "rows": rows, "in_features": args.in_features, "out_features": args.out_features, "iters": args.iters, @@ -179,14 +170,17 @@ def dequantize_then_linear_call() -> torch.Tensor: "with_bias": args.with_bias, "packed_seconds_per_iter": packed_seconds, "packed_first_call_seconds": packed_first_call_seconds, + "packed_hot_mean_seconds": packed_distribution["mean"], "packed_hot_median_seconds": packed_distribution["median"], "packed_hot_p95_seconds": packed_distribution["p95"], "predequantized_f_linear_seconds_per_iter": predequantized_linear_seconds, "predequantized_first_call_seconds": predequantized_first_call_seconds, + "predequantized_hot_mean_seconds": predequantized_distribution["mean"], "predequantized_hot_median_seconds": predequantized_distribution["median"], "predequantized_hot_p95_seconds": predequantized_distribution["p95"], "dequantize_then_f_linear_seconds_per_iter": dequantize_then_linear_seconds, "dequantize_then_first_call_seconds": dequantize_then_first_call_seconds, + "dequantize_then_hot_mean_seconds": dequantize_then_distribution["mean"], "dequantize_then_hot_median_seconds": dequantize_then_distribution["median"], "dequantize_then_hot_p95_seconds": dequantize_then_distribution["p95"], "packed_weight_indices_bytes": packed_weight_indices_bytes, @@ -213,6 +207,7 @@ def dequantize_then_linear_call() -> torch.Tensor: "max_abs_error": max_abs_error, "rmse": rmse, "relative_rmse": relative_rmse, + "timing_headline": "hot-loop median seconds per iteration", "reference": ( "predequantized PyTorch F.linear over a materialized dequantized " "weight matrix" @@ -221,7 +216,64 @@ def dequantize_then_linear_call() -> torch.Tensor: "materialize the dequantized weight matrix, then call PyTorch F.linear" ), } - print(json.dumps(payload, indent=2, sort_keys=True)) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--device", choices=["cpu", "cuda", "mps"], default="cuda") + parser.add_argument("--bits", type=int, default=4) + parser.add_argument( + "--rows", + type=_parse_rows, + default=[1, 8, 512, 4096], + help="comma-separated row counts to sweep (default covers decode-bound " + "small batches and GEMM-bound large batches)", + ) + parser.add_argument( + "--dtype", + choices=["auto", *sorted(_DTYPES)], + default="auto", + help="activation dtype; auto picks float16 on mps and bfloat16 elsewhere", + ) + parser.add_argument("--in-features", type=int, default=3072) + parser.add_argument("--out-features", type=int, default=3072) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--threads", type=int, default=0) + parser.add_argument("--with-bias", action="store_true") + args = parser.parse_args() + + if args.threads < 0: + parser.error("--threads must be non-negative") + if args.iters <= 0 or args.warmup < 0: + parser.error("--iters must be positive and --warmup must be non-negative") + if args.device == "cpu" and args.threads > 0: + os.environ["ORBITQUANT_CPU_THREADS"] = str(args.threads) + torch.set_num_threads(args.threads) + + torch.manual_seed(args.seed) + if args.dtype == "auto": + dtype = torch.float16 if args.device == "mps" else torch.bfloat16 + else: + dtype = _DTYPES[args.dtype] + indices = torch.randint( + 0, + 2**args.bits, + (args.out_features, args.in_features), + dtype=torch.uint8, + ) + packed = _pack(indices, args.bits).to(args.device) + row_norms = torch.ones(args.out_features, device=args.device, dtype=torch.bfloat16) + centroids = torch.linspace(-1.0, 1.0, 2**args.bits, device=args.device) + indices_device = indices.long().to(args.device) + weights = (packed, row_norms, centroids, indices_device) + + payloads = [_benchmark_rows(args, rows, dtype, weights) for rows in args.rows] + if len(payloads) == 1: + print(json.dumps(payloads[0], indent=2, sort_keys=True)) + else: + print(json.dumps(payloads, indent=2, sort_keys=True)) if __name__ == "__main__": diff --git a/native-kernels/orbitquant-packed-matmul/build.toml b/native-kernels/orbitquant-packed-matmul/build.toml index fb9efe1..f1d1cbf 100644 --- a/native-kernels/orbitquant-packed-matmul/build.toml +++ b/native-kernels/orbitquant-packed-matmul/build.toml @@ -26,6 +26,8 @@ include = ["orbitquant_packed_matmul_cpu"] src = [ "orbitquant_packed_matmul_cpu/cpu_isa.cpp", "orbitquant_packed_matmul_cpu/cpu_kernel_args.h", + "orbitquant_packed_matmul_cpu/cpu_pool.cpp", + "orbitquant_packed_matmul_cpu/cpu_pool.h", "orbitquant_packed_matmul_cpu/cpu_threads.cpp", "orbitquant_packed_matmul_cpu/cpu_threads.h", "orbitquant_packed_matmul_cpu/packed_adaln_cpu.cpp", diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_kernel_args.h b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_kernel_args.h index c18e5e2..a028db6 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_kernel_args.h +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_kernel_args.h @@ -16,7 +16,10 @@ enum class ActivationIsa : std::uint8_t { struct ActivationArgs { void *out; void const *x; + // Exactly one of the permutation pointers is set (int64 checkpoints and the + // int32 hot-path constants are both accepted). std::int64_t const *permutation; + std::int32_t const *permutation_i32; std::int8_t const *signs; float const *centroids; float const *boundaries; diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_pool.cpp b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_pool.cpp new file mode 100644 index 0000000..c4236e0 --- /dev/null +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_pool.cpp @@ -0,0 +1,126 @@ +#include "cpu_pool.h" +#include "cpu_threads.h" + +#include +#include +#include +#include +#include + +namespace orbitquant::cpu { +namespace { + +struct Job { + std::vector> const *ranges; + std::function const *fn; + std::atomic next_range{0}; + std::atomic pending_workers{0}; +}; + +class WorkerPool { + public: + static WorkerPool &instance() { + // Intentionally leaked: joining workers during static destruction can + // deadlock interpreter teardown, so the pool lives for the process. + static WorkerPool *pool = new WorkerPool(requested_threads() - 1); + return *pool; + } + + int worker_count() const { + return worker_count_; + } + + void run(Job &job) { + { + std::lock_guard lock(mutex_); + job.pending_workers.store(worker_count_, std::memory_order_relaxed); + job_ = &job; + ++generation_; + } + wake_cv_.notify_all(); + drain(job); + std::unique_lock lock(mutex_); + done_cv_.wait(lock, [&job] { + return job.pending_workers.load(std::memory_order_acquire) == 0; + }); + job_ = nullptr; + } + + private: + explicit WorkerPool(int requested_workers) + : worker_count_(requested_workers > 0 ? requested_workers : 0) { + for (int worker = 0; worker < worker_count_; ++worker) { + std::thread([this] { worker_loop(); }).detach(); + } + } + + static void drain(Job &job) { + const std::size_t range_count = job.ranges->size(); + for (;;) { + const std::size_t index = + job.next_range.fetch_add(1, std::memory_order_relaxed); + if (index >= range_count) { + return; + } + auto const &range = (*job.ranges)[index]; + (*job.fn)(range.first, range.second); + } + } + + void worker_loop() { + std::uint64_t seen_generation = 0; + for (;;) { + Job *job = nullptr; + { + std::unique_lock lock(mutex_); + wake_cv_.wait(lock, [this, seen_generation] { + return generation_ != seen_generation; + }); + seen_generation = generation_; + job = job_; + } + if (job == nullptr) { + continue; + } + drain(*job); + if (job->pending_workers.fetch_sub(1, std::memory_order_acq_rel) == 1) { + std::lock_guard lock(mutex_); + done_cv_.notify_all(); + } + } + } + + const int worker_count_; + std::mutex mutex_; + std::condition_variable wake_cv_; + std::condition_variable done_cv_; + Job *job_ = nullptr; + std::uint64_t generation_ = 0; +}; + +} // namespace + +void run_ranges( + std::vector> const &ranges, + std::function const &fn) { + if (ranges.empty()) { + return; + } + if (ranges.size() == 1) { + fn(ranges.front().first, ranges.front().second); + return; + } + WorkerPool &pool = WorkerPool::instance(); + if (pool.worker_count() == 0) { + for (auto const &range : ranges) { + fn(range.first, range.second); + } + return; + } + Job job; + job.ranges = &ranges; + job.fn = &fn; + pool.run(job); +} + +} // namespace orbitquant::cpu diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_pool.h b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_pool.h new file mode 100644 index 0000000..ad5bafe --- /dev/null +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/cpu_pool.h @@ -0,0 +1,18 @@ +#pragma once + +#include +#include +#include +#include + +namespace orbitquant::cpu { + +// Executes fn over each [start, end) range using a lazily created persistent +// worker pool plus the calling thread. Ranges are disjoint, so results stay +// deterministic regardless of which thread runs which range. A single range +// (or an empty pool) runs inline on the caller. +void run_ranges( + std::vector> const &ranges, + std::function const &fn); + +} // namespace orbitquant::cpu diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_adaln_cpu.cpp b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_adaln_cpu.cpp index f5872cd..d78d6cc 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_adaln_cpu.cpp +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_adaln_cpu.cpp @@ -1,3 +1,4 @@ +#include "cpu_pool.h" #include "cpu_threads.h" #include "cpu_kernel_args.h" #include "packed_matmul_cpu.h" @@ -12,7 +13,7 @@ #include #include #include -#include +#include #include #if defined(__aarch64__) || defined(_M_ARM64) @@ -780,8 +781,8 @@ void parallel_packed_adaln( return; } - std::vector workers; - workers.reserve(threads); + std::vector> ranges; + ranges.reserve(threads); const std::int64_t columns_per_thread = (args.out_features + threads - 1) / threads; for (int thread = 0; thread < threads; ++thread) { @@ -792,13 +793,13 @@ void parallel_packed_adaln( if (start >= end) { break; } - workers.emplace_back([&args, function, start, end] { - function(args, start, end); - }); - } - for (auto &worker : workers) { - worker.join(); + ranges.emplace_back(start, end); } + orbitquant::cpu::run_ranges( + ranges, + [&args, function](std::int64_t start, std::int64_t end) { + function(args, start, end); + }); } } // namespace diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_cpu.cpp b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_cpu.cpp index 8eb7ddb..787b385 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_cpu.cpp +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_cpu.cpp @@ -1,3 +1,4 @@ +#include "cpu_pool.h" #include "cpu_threads.h" #include "packed_matmul_cpu.h" #include "../torch-ext/torch_binding.h" @@ -9,7 +10,7 @@ #include #include #include -#include +#include #include namespace { @@ -47,8 +48,8 @@ void parallel_packed_matmul( return; } - std::vector workers; - workers.reserve(threads); + std::vector> ranges; + ranges.reserve(threads); const std::int64_t columns_per_thread = (args.out_features + threads - 1) / threads; for (int thread = 0; thread < threads; ++thread) { @@ -59,13 +60,13 @@ void parallel_packed_matmul( if (start >= end) { break; } - workers.emplace_back([&args, function, start, end] { - function(args, start, end); - }); - } - for (auto &worker : workers) { - worker.join(); + ranges.emplace_back(start, end); } + orbitquant::cpu::run_ranges( + ranges, + [&args, function](std::int64_t start, std::int64_t end) { + function(args, start, end); + }); } orbitquant::cpu::PackedMatmulRangeFn select_packed_matmul() { diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_neon.cpp b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_neon.cpp index 7fff79f..876173e 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_neon.cpp +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_neon.cpp @@ -7,6 +7,7 @@ #include #include +#include namespace orbitquant::cpu { namespace { @@ -126,6 +127,255 @@ inline void packed_matmul_neon_w4_rows( } } +template +inline std::uint32_t unpack_index_generic( + std::uint8_t const *packed_row, + std::int64_t value_index) { + const std::int64_t bit_start = value_index * Bits; + const std::int64_t byte_index = bit_start >> 3; + const unsigned bit_offset = static_cast(bit_start & 7); + std::uint32_t raw = packed_row[byte_index]; + if (bit_offset + static_cast(Bits) > 8) { + raw |= static_cast(packed_row[byte_index + 1]) << 8; + } + return (raw >> bit_offset) & ((1u << Bits) - 1u); +} + +template < + typename scalar_t, + float32x4_t (*load4)(void const *, std::int64_t), + int Bits, + int row_tile> +inline void packed_matmul_neon_lowbit_rows( + PackedMatmulArgs const &args, + std::uint8_t const *packed_row, + std::int64_t out_col, + std::int64_t row_start) { + float32x4_t accumulator0[row_tile]; + float32x4_t accumulator1[row_tile]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + accumulator0[row] = vdupq_n_f32(0.0f); + accumulator1[row] = vdupq_n_f32(0.0f); + } + + std::int64_t k = 0; + for (; k + 8 <= args.in_features; k += 8) { + float weights0[4]; + float weights1[4]; +#pragma clang loop unroll(full) + for (int value = 0; value < 4; ++value) { + weights0[value] = + args.centroids[unpack_index_generic(packed_row, k + value)]; + weights1[value] = + args.centroids[unpack_index_generic(packed_row, k + 4 + value)]; + } + const float32x4_t weight0 = vld1q_f32(weights0); + const float32x4_t weight1 = vld1q_f32(weights1); +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_offset = + (row_start + row) * args.in_features + k; + accumulator0[row] = + vfmaq_f32(accumulator0[row], load4(args.x, input_offset), weight0); + accumulator1[row] = vfmaq_f32( + accumulator1[row], load4(args.x, input_offset + 4), weight1); + } + } + + const float row_norm = args.row_norms[out_col]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_row_offset = + (row_start + row) * args.in_features; + float accumulator = + horizontal_sum(vaddq_f32(accumulator0[row], accumulator1[row])); + for (std::int64_t tail = k; tail < args.in_features; ++tail) { + const std::uint32_t index = + unpack_index_generic(packed_row, tail); + if constexpr (std::is_same_v) { + accumulator += + static_cast(args.x)[input_row_offset + tail] * + args.centroids[index]; + } else { + accumulator += static_cast( + static_cast( + args.x)[input_row_offset + tail]) * + args.centroids[index]; + } + } + accumulator *= row_norm; + if (args.has_bias) { + accumulator += args.bias[out_col]; + } + store_value( + args.out, + (row_start + row) * args.out_features + out_col, + accumulator); + } +} + +template +inline void packed_matmul_neon_buffered_rows( + PackedMatmulArgs const &args, + float const *decoded_row, + std::int64_t out_col, + std::int64_t row_start) { + float32x4_t accumulator0[row_tile]; + float32x4_t accumulator1[row_tile]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + accumulator0[row] = vdupq_n_f32(0.0f); + accumulator1[row] = vdupq_n_f32(0.0f); + } + + std::int64_t k = 0; + for (; k + 8 <= args.in_features; k += 8) { + const float32x4_t weight0 = vld1q_f32(decoded_row + k); + const float32x4_t weight1 = vld1q_f32(decoded_row + k + 4); +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_offset = + (row_start + row) * args.in_features + k; + accumulator0[row] = + vfmaq_f32(accumulator0[row], load4(args.x, input_offset), weight0); + accumulator1[row] = vfmaq_f32( + accumulator1[row], load4(args.x, input_offset + 4), weight1); + } + } + + const float row_norm = args.row_norms[out_col]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_row_offset = + (row_start + row) * args.in_features; + float accumulator = + horizontal_sum(vaddq_f32(accumulator0[row], accumulator1[row])); + for (std::int64_t tail = k; tail < args.in_features; ++tail) { + if constexpr (std::is_same_v) { + accumulator += + static_cast(args.x)[input_row_offset + tail] * + decoded_row[tail]; + } else { + accumulator += static_cast( + static_cast( + args.x)[input_row_offset + tail]) * + decoded_row[tail]; + } + } + accumulator *= row_norm; + if (args.has_bias) { + accumulator += args.bias[out_col]; + } + store_value( + args.out, + (row_start + row) * args.out_features + out_col, + accumulator); + } +} + +template < + typename scalar_t, + float32x4_t (*load4)(void const *, std::int64_t), + int Bits> +void packed_matmul_neon_lowbit_typed( + PackedMatmulArgs const &args, + std::int64_t out_start, + std::int64_t out_end) { + const std::int64_t packed_row_bytes = args.in_features * Bits / 8; + // The NEON decode goes through scalar table lookups, so expanding the + // column once per >=16-row call amortizes the costliest stage. + const bool use_decoded_buffer = args.rows >= 16; + thread_local std::vector decoded_row_storage; + if (use_decoded_buffer && + decoded_row_storage.size() < static_cast(args.in_features)) { + decoded_row_storage.resize(static_cast(args.in_features)); + } + for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) { + const auto *packed_row = + args.packed_weight_indices + out_col * packed_row_bytes; + if (use_decoded_buffer) { + float *decoded_row = decoded_row_storage.data(); + for (std::int64_t k = 0; k < args.in_features; ++k) { + decoded_row[k] = + args.centroids[unpack_index_generic(packed_row, k)]; + } + std::int64_t row = 0; + for (; row + 8 <= args.rows; row += 8) { + packed_matmul_neon_buffered_rows( + args, decoded_row, out_col, row); + } + for (; row + 4 <= args.rows; row += 4) { + packed_matmul_neon_buffered_rows( + args, decoded_row, out_col, row); + } + switch (args.rows - row) { + case 3: + packed_matmul_neon_buffered_rows( + args, decoded_row, out_col, row); + break; + case 2: + packed_matmul_neon_buffered_rows( + args, decoded_row, out_col, row); + break; + case 1: + packed_matmul_neon_buffered_rows( + args, decoded_row, out_col, row); + break; + default: + break; + } + continue; + } + std::int64_t row = 0; + for (; row + 8 <= args.rows; row += 8) { + packed_matmul_neon_lowbit_rows( + args, packed_row, out_col, row); + } + for (; row + 4 <= args.rows; row += 4) { + packed_matmul_neon_lowbit_rows( + args, packed_row, out_col, row); + } + switch (args.rows - row) { + case 3: + packed_matmul_neon_lowbit_rows( + args, packed_row, out_col, row); + break; + case 2: + packed_matmul_neon_lowbit_rows( + args, packed_row, out_col, row); + break; + case 1: + packed_matmul_neon_lowbit_rows( + args, packed_row, out_col, row); + break; + default: + break; + } + } +} + +template +void packed_matmul_neon_lowbit_dispatch( + PackedMatmulArgs const &args, + std::int64_t out_start, + std::int64_t out_end) { + switch (args.scalar_kind) { + case ScalarKind::Float32: + packed_matmul_neon_lowbit_typed( + args, out_start, out_end); + return; + case ScalarKind::Float16: + packed_matmul_neon_lowbit_typed( + args, out_start, out_end); + return; + case ScalarKind::BFloat16: + packed_matmul_neon_lowbit_typed( + args, out_start, out_end); + return; + } +} + template void packed_matmul_neon_w4_typed( PackedMatmulArgs const &args, @@ -176,6 +426,24 @@ void packed_matmul_neon_range( PackedMatmulArgs const &args, std::int64_t out_start, std::int64_t out_end) { + if (args.bits == 4 && args.in_features % 2 == 0 && args.rows >= 16) { + // The buffered generic path decodes each column once, which beats the + // per-tile scalar decode as soon as several row tiles reuse it. + packed_matmul_neon_lowbit_dispatch<4>(args, out_start, out_end); + return; + } + if (args.bits == 2 && args.in_features % 4 == 0) { + packed_matmul_neon_lowbit_dispatch<2>(args, out_start, out_end); + return; + } + if (args.bits == 3 && args.in_features % 8 == 0) { + packed_matmul_neon_lowbit_dispatch<3>(args, out_start, out_end); + return; + } + if (args.bits == 6 && args.in_features % 4 == 0) { + packed_matmul_neon_lowbit_dispatch<6>(args, out_start, out_end); + return; + } if (args.bits != 4 || args.in_features % 2 != 0) { packed_matmul_scalar_range(args, out_start, out_end); return; diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_x86.cpp b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_x86.cpp index ca317ac..b6121e8 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_x86.cpp +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_x86.cpp @@ -2,6 +2,11 @@ #if defined(__x86_64__) || defined(_M_X64) #include +#if defined(_MSC_VER) +#include +#else +#include +#endif #include #include @@ -9,6 +14,7 @@ #include #include #include +#include #if defined(_MSC_VER) #define ORBITQUANT_TARGET_AVX2 @@ -145,20 +151,385 @@ ORBITQUANT_TARGET_AVX2 inline void packed_matmul_avx2_w4_rows( } } -template -ORBITQUANT_TARGET_AVX2 ORBITQUANT_NOINLINE void packed_matmul_avx2_w4_typed( +template +inline std::uint32_t unpack_index_generic( + std::uint8_t const *packed_row, + std::int64_t value_index) { + const std::int64_t bit_start = value_index * Bits; + const std::int64_t byte_index = bit_start >> 3; + const unsigned bit_offset = static_cast(bit_start & 7); + std::uint32_t raw = packed_row[byte_index]; + if (bit_offset + static_cast(Bits) > 8) { + raw |= static_cast(packed_row[byte_index + 1]) << 8; + } + return (raw >> bit_offset) & ((1u << Bits) - 1u); +} + +// Each decoder turns 8 consecutive packed indices into 8 fp32 centroid +// values. Rows are byte-aligned because dispatch requires in_features % 4 == 0. +struct W2Avx2Decoder { + static constexpr int kBits = 2; + struct Tables { + __m256 lut; + }; + + ORBITQUANT_TARGET_AVX2 static inline Tables load_tables( + float const *centroids) { + const __m128 lut4 = _mm_loadu_ps(centroids); + return Tables{_mm256_set_m128(lut4, lut4)}; + } + + ORBITQUANT_TARGET_AVX2 static inline __m256 decode( + std::uint8_t const *packed_row, + std::int64_t k, + Tables const &tables) { + std::uint16_t packed_bits; + std::memcpy(&packed_bits, packed_row + (k >> 2), sizeof(packed_bits)); + const __m128i bytes = _mm_cvtsi32_si128(packed_bits); + const __m128i replicated = _mm_shuffle_epi8( + bytes, + _mm_setr_epi8(0, 0, 0, 0, 1, 1, 1, 1, -1, -1, -1, -1, -1, -1, -1, -1)); + const __m256i widened = _mm256_cvtepu8_epi32(replicated); + const __m256i shifts = _mm256_setr_epi32(0, 2, 4, 6, 0, 2, 4, 6); + const __m256i indices = _mm256_and_si256( + _mm256_srlv_epi32(widened, shifts), + _mm256_set1_epi32(3)); + return _mm256_permutevar8x32_ps(tables.lut, indices); + } +}; + +struct W6Avx2Decoder { + static constexpr int kBits = 6; + struct Tables { + float const *centroids; + }; + + ORBITQUANT_TARGET_AVX2 static inline Tables load_tables( + float const *centroids) { + return Tables{centroids}; + } + + ORBITQUANT_TARGET_AVX2 static inline __m256 decode( + std::uint8_t const *packed_row, + std::int64_t k, + Tables const &tables) { + std::uint64_t raw_bits = 0; + std::memcpy(&raw_bits, packed_row + (k * 6 >> 3), 6); + const __m128i raw = + _mm_cvtsi64_si128(static_cast(raw_bits)); + const __m128i windows = _mm_shuffle_epi8( + raw, + _mm_setr_epi8(0, 1, 0, 1, 1, 2, 2, 3, 3, 4, 3, 4, 4, 5, 5, 6)); + const __m256i widened = _mm256_cvtepu16_epi32(windows); + const __m256i shifts = _mm256_setr_epi32(0, 6, 4, 2, 0, 6, 4, 2); + const __m256i indices = _mm256_and_si256( + _mm256_srlv_epi32(widened, shifts), + _mm256_set1_epi32(63)); + return _mm256_i32gather_ps(tables.centroids, indices, 4); + } +}; + +template < + typename scalar_t, + __m256 (*load8)(void const *, std::int64_t), + typename decoder_t, + int row_tile> +ORBITQUANT_TARGET_AVX2 inline void packed_matmul_avx2_lowbit_rows( + PackedMatmulArgs const &args, + std::uint8_t const *packed_row, + std::int64_t out_col, + std::int64_t row_start) { + __m256 accumulators[row_tile]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + accumulators[row] = _mm256_setzero_ps(); + } + const typename decoder_t::Tables tables = + decoder_t::load_tables(args.centroids); + + std::int64_t k = 0; + for (; k + 8 <= args.in_features; k += 8) { + const __m256 weight = decoder_t::decode(packed_row, k, tables); +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_offset = + (row_start + row) * args.in_features + k; + accumulators[row] = _mm256_fmadd_ps( + load8(args.x, input_offset), + weight, + accumulators[row]); + } + } + + const float row_norm = args.row_norms[out_col]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_row_offset = + (row_start + row) * args.in_features; + float accumulator = horizontal_sum(accumulators[row]); + for (std::int64_t tail = k; tail < args.in_features; ++tail) { + const std::uint32_t index = + unpack_index_generic(packed_row, tail); + if constexpr (std::is_same_v) { + accumulator += + static_cast(args.x)[input_row_offset + tail] * + args.centroids[index]; + } else { + accumulator += static_cast( + static_cast( + args.x)[input_row_offset + tail]) * + args.centroids[index]; + } + } + accumulator *= row_norm; + if (args.has_bias) { + accumulator += args.bias[out_col]; + } + store_value( + args.out, + (row_start + row) * args.out_features + out_col, + accumulator); + } +} + +template +ORBITQUANT_TARGET_AVX2 inline void packed_matmul_avx2_buffered_rows( + PackedMatmulArgs const &args, + float const *decoded_row, + std::int64_t out_col, + std::int64_t row_start) { + __m256 accumulators[row_tile]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + accumulators[row] = _mm256_setzero_ps(); + } + + std::int64_t k = 0; + for (; k + 8 <= args.in_features; k += 8) { + const __m256 weight = _mm256_loadu_ps(decoded_row + k); +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_offset = + (row_start + row) * args.in_features + k; + accumulators[row] = _mm256_fmadd_ps( + load8(args.x, input_offset), + weight, + accumulators[row]); + } + } + + const float row_norm = args.row_norms[out_col]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_row_offset = + (row_start + row) * args.in_features; + float accumulator = horizontal_sum(accumulators[row]); + for (std::int64_t tail = k; tail < args.in_features; ++tail) { + if constexpr (std::is_same_v) { + accumulator += + static_cast(args.x)[input_row_offset + tail] * + decoded_row[tail]; + } else { + accumulator += static_cast( + static_cast( + args.x)[input_row_offset + tail]) * + decoded_row[tail]; + } + } + accumulator *= row_norm; + if (args.has_bias) { + accumulator += args.bias[out_col]; + } + store_value( + args.out, + (row_start + row) * args.out_features + out_col, + accumulator); + } +} + +template < + typename scalar_t, + __m256 (*load8)(void const *, std::int64_t), + typename decoder_t> +ORBITQUANT_TARGET_AVX2 ORBITQUANT_NOINLINE void +packed_matmul_avx2_lowbit_typed( PackedMatmulArgs const &args, std::int64_t out_start, std::int64_t out_end) { constexpr int kPrimaryRowTile = 8; - const std::int64_t packed_row_bytes = args.in_features / 2; + const std::int64_t packed_row_bytes = + args.in_features * decoder_t::kBits / 8; + // Two or more row tiles amortize the packed decode: expand the column once + // into a per-thread scratch row and stream plain FMA tiles from it. + const bool use_decoded_buffer = args.rows >= 16; + thread_local std::vector decoded_row_storage; + if (use_decoded_buffer && + decoded_row_storage.size() < static_cast(args.in_features)) { + decoded_row_storage.resize(static_cast(args.in_features)); + } for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) { const auto *packed_row = args.packed_weight_indices + out_col * packed_row_bytes; + if (use_decoded_buffer) { + float *decoded_row = decoded_row_storage.data(); + const typename decoder_t::Tables tables = + decoder_t::load_tables(args.centroids); + std::int64_t k = 0; + for (; k + 8 <= args.in_features; k += 8) { + _mm256_storeu_ps(decoded_row + k, decoder_t::decode(packed_row, k, tables)); + } + for (; k < args.in_features; ++k) { + decoded_row[k] = + args.centroids[unpack_index_generic(packed_row, k)]; + } + std::int64_t row = 0; + for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) { + packed_matmul_avx2_buffered_rows( + args, decoded_row, out_col, row); + } + if (row + 4 <= args.rows) { + packed_matmul_avx2_buffered_rows( + args, decoded_row, out_col, row); + row += 4; + } + switch (args.rows - row) { + case 3: + packed_matmul_avx2_buffered_rows( + args, decoded_row, out_col, row); + break; + case 2: + packed_matmul_avx2_buffered_rows( + args, decoded_row, out_col, row); + break; + case 1: + packed_matmul_avx2_buffered_rows( + args, decoded_row, out_col, row); + break; + default: + break; + } + continue; + } std::int64_t row = 0; for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) { - packed_matmul_avx2_w4_rows( + packed_matmul_avx2_lowbit_rows( + args, packed_row, out_col, row); + } + if (row + 4 <= args.rows) { + packed_matmul_avx2_lowbit_rows( + args, packed_row, out_col, row); + row += 4; + } + switch (args.rows - row) { + case 3: + packed_matmul_avx2_lowbit_rows( + args, packed_row, out_col, row); + break; + case 2: + packed_matmul_avx2_lowbit_rows( + args, packed_row, out_col, row); + break; + case 1: + packed_matmul_avx2_lowbit_rows( + args, packed_row, out_col, row); + break; + default: + break; + } + } +} + +template +void packed_matmul_avx2_lowbit_dispatch( + PackedMatmulArgs const &args, + std::int64_t out_start, + std::int64_t out_end) { + switch (args.scalar_kind) { + case ScalarKind::Float32: + packed_matmul_avx2_lowbit_typed( + args, out_start, out_end); + return; + case ScalarKind::Float16: + packed_matmul_avx2_lowbit_typed( + args, out_start, out_end); + return; + case ScalarKind::BFloat16: + packed_matmul_avx2_lowbit_typed( + args, out_start, out_end); + return; + } +} + +template +bool use_verified_amd_cezanne_row_tile(PackedMatmulArgs const &args) { + if constexpr (!std::is_same_v) { + return false; + } + const bool tuned_dimension = args.in_features == 1536 || + args.in_features == 1920 || args.in_features == 3072; + if (args.rows < 16 || !tuned_dimension) { + return false; + } + static const bool verified_cpu = [] { + unsigned int eax = 0; + unsigned int ebx = 0; + unsigned int ecx = 0; + unsigned int edx = 0; +#if defined(_MSC_VER) + int registers[4]{}; + __cpuid(registers, 0); + eax = static_cast(registers[0]); + ebx = static_cast(registers[1]); + ecx = static_cast(registers[2]); + edx = static_cast(registers[3]); + if (ebx != 0x68747541u || edx != 0x69746e65u || + ecx != 0x444d4163u) { + return false; + } + __cpuid(registers, 1); + eax = static_cast(registers[0]); +#else + // CPUID vendor registers spell "AuthenticAMD" in EBX, EDX, ECX order. + if (!__get_cpuid(0, &eax, &ebx, &ecx, &edx) || + ebx != 0x68747541u || edx != 0x69746e65u || ecx != 0x444d4163u || + !__get_cpuid(1, &eax, &ebx, &ecx, &edx)) { + return false; + } +#endif + const unsigned int base_family = (eax >> 8) & 0xfu; + const unsigned int family = base_family == 0xfu + ? base_family + ((eax >> 20) & 0xffu) + : base_family; + // The 16-row tile was measured on Zen 3 (Ryzen 5 5600G); apply it to the + // whole AVX2-only Zen 3 family (19h without AVX-512) instead of pinning + // the one benchmarked model. + return family == 0x19u && !packed_matmul_x86_avx512_available(); + }(); + return verified_cpu; +} + +template < + typename scalar_t, + __m256 (*load8)(void const *, std::int64_t), + int primary_row_tile> +ORBITQUANT_TARGET_AVX2 ORBITQUANT_NOINLINE void packed_matmul_avx2_w4_typed( + PackedMatmulArgs const &args, + std::int64_t out_start, + std::int64_t out_end) { + static_assert(primary_row_tile == 8 || primary_row_tile == 16); + const std::int64_t packed_row_bytes = args.in_features / 2; + for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) { + const auto *packed_row = + args.packed_weight_indices + out_col * packed_row_bytes; + std::int64_t row = 0; + for (; row + primary_row_tile <= args.rows; row += primary_row_tile) { + packed_matmul_avx2_w4_rows( + args, packed_row, out_col, row); + } + if (row + 8 <= args.rows) { + packed_matmul_avx2_w4_rows( args, packed_row, out_col, row); + row += 8; } if (row + 4 <= args.rows) { packed_matmul_avx2_w4_rows( @@ -190,21 +561,39 @@ void packed_matmul_x86_avx2_range( PackedMatmulArgs const &args, std::int64_t out_start, std::int64_t out_end) { - if (!packed_matmul_x86_avx2_available() || args.bits != 4 || - args.in_features % 2 != 0) { + if (!packed_matmul_x86_avx2_available()) { + packed_matmul_scalar_range(args, out_start, out_end); + return; + } + if (args.bits == 2 && args.in_features % 4 == 0) { + packed_matmul_avx2_lowbit_dispatch(args, out_start, out_end); + return; + } + if (args.bits == 6 && args.in_features % 4 == 0) { + packed_matmul_avx2_lowbit_dispatch(args, out_start, out_end); + return; + } + if (args.bits != 4 || args.in_features % 2 != 0) { packed_matmul_scalar_range(args, out_start, out_end); return; } switch (args.scalar_kind) { case ScalarKind::Float32: - packed_matmul_avx2_w4_typed(args, out_start, out_end); + packed_matmul_avx2_w4_typed( + args, out_start, out_end); return; case ScalarKind::Float16: - packed_matmul_avx2_w4_typed(args, out_start, out_end); + packed_matmul_avx2_w4_typed( + args, out_start, out_end); return; case ScalarKind::BFloat16: - packed_matmul_avx2_w4_typed( - args, out_start, out_end); + if (use_verified_amd_cezanne_row_tile(args)) { + packed_matmul_avx2_w4_typed( + args, out_start, out_end); + } else { + packed_matmul_avx2_w4_typed( + args, out_start, out_end); + } return; } } diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_x86_avx512.cpp b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_x86_avx512.cpp index 1bbeed0..8210321 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_x86_avx512.cpp +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/packed_matmul_x86_avx512.cpp @@ -1,6 +1,7 @@ #include "packed_matmul_cpu.h" #if (defined(__x86_64__) || defined(_M_X64)) && !defined(_MSC_VER) +#include #include #include @@ -9,6 +10,7 @@ #include #include #include +#include #define ORBITQUANT_TARGET_AVX512 \ __attribute__((target("avx512f,avx512dq,avx512bw,avx512vl,fma,f16c"))) @@ -17,6 +19,7 @@ "avx512f,avx512dq,avx512bw,avx512vl,avx512bf16,fma,f16c"))) #define ORBITQUANT_HAS_AVX512_BF16_INTRINSICS 1 #define ORBITQUANT_NOINLINE __attribute__((noinline)) +#define ORBITQUANT_ALWAYS_INLINE __attribute__((always_inline)) namespace orbitquant::cpu { namespace { @@ -177,8 +180,408 @@ ORBITQUANT_TARGET_AVX512 ORBITQUANT_NOINLINE void packed_matmul_avx512_w4_typed( } } +template +inline std::uint32_t unpack_index_generic( + std::uint8_t const *packed_row, + std::int64_t value_index) { + const std::int64_t bit_start = value_index * Bits; + const std::int64_t byte_index = bit_start >> 3; + const unsigned bit_offset = static_cast(bit_start & 7); + std::uint32_t raw = packed_row[byte_index]; + if (bit_offset + static_cast(Bits) > 8) { + raw |= static_cast(packed_row[byte_index + 1]) << 8; + } + return (raw >> bit_offset) & ((1u << Bits) - 1u); +} + +// Each decoder turns 16 consecutive packed indices into 16 fp32 centroid +// values. Rows are byte-aligned because dispatch requires in_features +// divisibility (W2/W6: in % 4 == 0, W3: in % 8 == 0). +struct W2Avx512Decoder { + static constexpr int kBits = 2; + struct Tables { + __m512 lut; + }; + + ORBITQUANT_TARGET_AVX512 static inline Tables load_tables( + float const *centroids) { + return Tables{_mm512_broadcast_f32x4(_mm_loadu_ps(centroids))}; + } + + ORBITQUANT_TARGET_AVX512 static inline __m512 decode( + std::uint8_t const *packed_row, + std::int64_t k, + Tables const &tables) { + std::uint32_t packed_bits; + std::memcpy(&packed_bits, packed_row + (k >> 2), sizeof(packed_bits)); + const __m128i bytes = _mm_cvtsi32_si128(static_cast(packed_bits)); + const __m128i replicated = _mm_shuffle_epi8( + bytes, + _mm_setr_epi8(0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3)); + const __m512i widened = _mm512_cvtepu8_epi32(replicated); + const __m512i shifts = _mm512_set_epi32( + 6, 4, 2, 0, 6, 4, 2, 0, 6, 4, 2, 0, 6, 4, 2, 0); + const __m512i indices = _mm512_and_si512( + _mm512_srlv_epi32(widened, shifts), + _mm512_set1_epi32(3)); + return _mm512_permutexvar_ps(indices, tables.lut); + } +}; + +struct W3Avx512Decoder { + static constexpr int kBits = 3; + struct Tables { + __m512 lut; + }; + + ORBITQUANT_TARGET_AVX512 static inline Tables load_tables( + float const *centroids) { + return Tables{_mm512_broadcast_f32x8(_mm256_loadu_ps(centroids))}; + } + + ORBITQUANT_TARGET_AVX512 static inline __m512 decode( + std::uint8_t const *packed_row, + std::int64_t k, + Tables const &tables) { + std::uint64_t raw_bits = 0; + std::memcpy(&raw_bits, packed_row + (k * 3 >> 3), 6); + const __m128i raw = _mm_cvtsi64_si128(static_cast(raw_bits)); + const __m128i window_low = _mm_shuffle_epi8( + raw, + _mm_setr_epi8(0, 1, 0, 1, 0, 1, 1, 2, 1, 2, 1, 2, 2, 3, 2, 3)); + const __m128i window_high = _mm_shuffle_epi8( + raw, + _mm_setr_epi8(3, 4, 3, 4, 3, 4, 4, 5, 4, 5, 4, 5, 5, 6, 5, 6)); + const __m512i widened = _mm512_cvtepu16_epi32( + _mm256_set_m128i(window_high, window_low)); + const __m512i shifts = _mm512_set_epi32( + 5, 2, 7, 4, 1, 6, 3, 0, 5, 2, 7, 4, 1, 6, 3, 0); + const __m512i indices = _mm512_and_si512( + _mm512_srlv_epi32(widened, shifts), + _mm512_set1_epi32(7)); + return _mm512_permutexvar_ps(indices, tables.lut); + } +}; + +struct W6Avx512Decoder { + static constexpr int kBits = 6; + struct Tables { + __m512 lut0; + __m512 lut1; + __m512 lut2; + __m512 lut3; + }; + + ORBITQUANT_TARGET_AVX512 static inline Tables load_tables( + float const *centroids) { + return Tables{ + _mm512_loadu_ps(centroids), + _mm512_loadu_ps(centroids + 16), + _mm512_loadu_ps(centroids + 32), + _mm512_loadu_ps(centroids + 48)}; + } + + ORBITQUANT_TARGET_AVX512 static inline __m512 decode( + std::uint8_t const *packed_row, + std::int64_t k, + Tables const &tables) { + std::uint64_t low_bytes; + std::uint32_t high_bytes; + std::uint8_t const *source = packed_row + (k * 6 >> 3); + std::memcpy(&low_bytes, source, sizeof(low_bytes)); + std::memcpy(&high_bytes, source + 8, sizeof(high_bytes)); + const __m128i raw = _mm_set_epi64x( + static_cast(high_bytes), + static_cast(low_bytes)); + const __m128i window_low = _mm_shuffle_epi8( + raw, + _mm_setr_epi8(0, 1, 0, 1, 1, 2, 2, 3, 3, 4, 3, 4, 4, 5, 5, 6)); + const __m128i window_high = _mm_shuffle_epi8( + raw, + _mm_setr_epi8(6, 7, 6, 7, 7, 8, 8, 9, 9, 10, 9, 10, 10, 11, 11, 12)); + const __m512i widened = _mm512_cvtepu16_epi32( + _mm256_set_m128i(window_high, window_low)); + const __m512i shifts = _mm512_set_epi32( + 2, 4, 6, 0, 2, 4, 6, 0, 2, 4, 6, 0, 2, 4, 6, 0); + const __m512i indices = _mm512_and_si512( + _mm512_srlv_epi32(widened, shifts), + _mm512_set1_epi32(63)); + const __m512 low_pair = _mm512_permutex2var_ps(tables.lut0, indices, tables.lut1); + const __m512 high_pair = _mm512_permutex2var_ps(tables.lut2, indices, tables.lut3); + const __mmask16 use_high = + _mm512_test_epi32_mask(indices, _mm512_set1_epi32(32)); + return _mm512_mask_blend_ps(use_high, low_pair, high_pair); + } +}; + +template < + typename scalar_t, + __m512 (*load16)(void const *, std::int64_t), + typename decoder_t, + int row_tile> +ORBITQUANT_TARGET_AVX512 inline void packed_matmul_avx512_lowbit_rows( + PackedMatmulArgs const &args, + std::uint8_t const *packed_row, + std::int64_t out_col, + std::int64_t row_start) { + __m512 accumulators[row_tile]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + accumulators[row] = _mm512_setzero_ps(); + } + const typename decoder_t::Tables tables = + decoder_t::load_tables(args.centroids); + + std::int64_t k = 0; + for (; k + 16 <= args.in_features; k += 16) { + const __m512 weight = decoder_t::decode(packed_row, k, tables); +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_offset = + (row_start + row) * args.in_features + k; + accumulators[row] = _mm512_fmadd_ps( + load16(args.x, input_offset), + weight, + accumulators[row]); + } + } + + const float row_norm = args.row_norms[out_col]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_row_offset = + (row_start + row) * args.in_features; + float accumulator = horizontal_sum(accumulators[row]); + for (std::int64_t tail = k; tail < args.in_features; ++tail) { + const std::uint32_t index = + unpack_index_generic(packed_row, tail); + if constexpr (std::is_same_v) { + accumulator += + static_cast(args.x)[input_row_offset + tail] * + args.centroids[index]; + } else { + accumulator += static_cast( + static_cast( + args.x)[input_row_offset + tail]) * + args.centroids[index]; + } + } + accumulator *= row_norm; + if (args.has_bias) { + accumulator += args.bias[out_col]; + } + store_value( + args.out, + (row_start + row) * args.out_features + out_col, + accumulator); + } +} + +template +ORBITQUANT_TARGET_AVX512 inline void packed_matmul_avx512_buffered_rows( + PackedMatmulArgs const &args, + float const *decoded_row, + std::int64_t out_col, + std::int64_t row_start) { + __m512 accumulators[row_tile]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + accumulators[row] = _mm512_setzero_ps(); + } + + std::int64_t k = 0; + for (; k + 16 <= args.in_features; k += 16) { + const __m512 weight = _mm512_loadu_ps(decoded_row + k); +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_offset = + (row_start + row) * args.in_features + k; + accumulators[row] = _mm512_fmadd_ps( + load16(args.x, input_offset), + weight, + accumulators[row]); + } + } + + const float row_norm = args.row_norms[out_col]; +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_row_offset = + (row_start + row) * args.in_features; + float accumulator = horizontal_sum(accumulators[row]); + for (std::int64_t tail = k; tail < args.in_features; ++tail) { + if constexpr (std::is_same_v) { + accumulator += + static_cast(args.x)[input_row_offset + tail] * + decoded_row[tail]; + } else { + accumulator += static_cast( + static_cast( + args.x)[input_row_offset + tail]) * + decoded_row[tail]; + } + } + accumulator *= row_norm; + if (args.has_bias) { + accumulator += args.bias[out_col]; + } + store_value( + args.out, + (row_start + row) * args.out_features + out_col, + accumulator); + } +} + +template < + typename scalar_t, + __m512 (*load16)(void const *, std::int64_t), + typename decoder_t> +ORBITQUANT_TARGET_AVX512 ORBITQUANT_NOINLINE void +packed_matmul_avx512_lowbit_typed( + PackedMatmulArgs const &args, + std::int64_t out_start, + std::int64_t out_end) { + constexpr int kPrimaryRowTile = 8; + const std::int64_t packed_row_bytes = + args.in_features * decoder_t::kBits / 8; + // Two or more row tiles amortize the packed decode: expand the column once + // into a per-thread scratch row and stream plain FMA tiles from it. + const bool use_decoded_buffer = args.rows >= 16; + thread_local std::vector decoded_row_storage; + if (use_decoded_buffer && + decoded_row_storage.size() < static_cast(args.in_features)) { + decoded_row_storage.resize(static_cast(args.in_features)); + } + for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) { + const auto *packed_row = + args.packed_weight_indices + out_col * packed_row_bytes; + if (use_decoded_buffer) { + float *decoded_row = decoded_row_storage.data(); + const typename decoder_t::Tables tables = + decoder_t::load_tables(args.centroids); + std::int64_t k = 0; + for (; k + 16 <= args.in_features; k += 16) { + _mm512_storeu_ps(decoded_row + k, decoder_t::decode(packed_row, k, tables)); + } + for (; k < args.in_features; ++k) { + decoded_row[k] = + args.centroids[unpack_index_generic(packed_row, k)]; + } + std::int64_t row = 0; + for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) { + packed_matmul_avx512_buffered_rows( + args, decoded_row, out_col, row); + } + if (row + 4 <= args.rows) { + packed_matmul_avx512_buffered_rows( + args, decoded_row, out_col, row); + row += 4; + } + switch (args.rows - row) { + case 3: + packed_matmul_avx512_buffered_rows( + args, decoded_row, out_col, row); + break; + case 2: + packed_matmul_avx512_buffered_rows( + args, decoded_row, out_col, row); + break; + case 1: + packed_matmul_avx512_buffered_rows( + args, decoded_row, out_col, row); + break; + default: + break; + } + continue; + } + std::int64_t row = 0; + for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) { + packed_matmul_avx512_lowbit_rows( + args, packed_row, out_col, row); + } + if (row + 4 <= args.rows) { + packed_matmul_avx512_lowbit_rows( + args, packed_row, out_col, row); + row += 4; + } + switch (args.rows - row) { + case 3: + packed_matmul_avx512_lowbit_rows( + args, packed_row, out_col, row); + break; + case 2: + packed_matmul_avx512_lowbit_rows( + args, packed_row, out_col, row); + break; + case 1: + packed_matmul_avx512_lowbit_rows( + args, packed_row, out_col, row); + break; + default: + break; + } + } +} + +template +ORBITQUANT_TARGET_AVX512 void packed_matmul_avx512_lowbit_dispatch( + PackedMatmulArgs const &args, + std::int64_t out_start, + std::int64_t out_end) { + switch (args.scalar_kind) { + case ScalarKind::Float32: + packed_matmul_avx512_lowbit_typed( + args, out_start, out_end); + return; + case ScalarKind::Float16: + packed_matmul_avx512_lowbit_typed( + args, out_start, out_end); + return; + case ScalarKind::BFloat16: + packed_matmul_avx512_lowbit_typed( + args, out_start, out_end); + return; + } +} + #if defined(ORBITQUANT_HAS_AVX512_BF16_INTRINSICS) -template +template +ORBITQUANT_TARGET_AVX512_BF16 ORBITQUANT_ALWAYS_INLINE inline void +accumulate_bf16_w4_chunk( + PackedMatmulArgs const &args, + std::uint8_t const *packed_row, + std::int64_t row_start, + std::int64_t k, + __m512i centroid_lut_words, + __m128i nibble_mask, + __m512 (&accumulators)[row_tile]) { + const std::int64_t byte_offset = k / 2; + const __m128i bytes = _mm_loadu_si128( + reinterpret_cast<__m128i const *>(packed_row + byte_offset)); + const __m128i low = _mm_and_si128(bytes, nibble_mask); + const __m128i high = _mm_and_si128( + _mm_srli_epi16(bytes, 4), + nibble_mask); + const __m256i packed_indices = _mm256_set_m128i( + _mm_unpackhi_epi8(low, high), + _mm_unpacklo_epi8(low, high)); + const __m512i indices = _mm512_cvtepu8_epi16(packed_indices); + const __m512bh weights = (__m512bh)_mm512_permutexvar_epi16( + indices, + centroid_lut_words); +#pragma clang loop unroll(full) + for (int row = 0; row < row_tile; ++row) { + const std::int64_t input_offset = + (row_start + row) * args.in_features + k; + const __m512bh activations = (__m512bh)_mm512_loadu_si512( + static_cast(args.x) + input_offset); + accumulators[row] = + _mm512_dpbf16_ps(accumulators[row], activations, weights); + } +} + +template ORBITQUANT_TARGET_AVX512_BF16 inline void packed_matmul_avx512_bf16_w4_rows( PackedMatmulArgs const &args, std::uint8_t const *packed_row, @@ -199,29 +602,32 @@ ORBITQUANT_TARGET_AVX512_BF16 inline void packed_matmul_avx512_bf16_w4_rows( const __m128i nibble_mask = _mm_set1_epi8(15); std::int64_t k = 0; - for (; k + 32 <= args.in_features; k += 32) { - const std::int64_t byte_offset = k / 2; - const __m128i bytes = _mm_loadu_si128( - reinterpret_cast<__m128i const *>(packed_row + byte_offset)); - const __m128i low = _mm_and_si128(bytes, nibble_mask); - const __m128i high = _mm_and_si128( - _mm_srli_epi16(bytes, 4), - nibble_mask); - const __m256i packed_indices = _mm256_set_m128i( - _mm_unpackhi_epi8(low, high), - _mm_unpacklo_epi8(low, high)); - const __m512i indices = _mm512_cvtepu8_epi16(packed_indices); - const __m512bh weights = (__m512bh)_mm512_permutexvar_epi16( - indices, - centroid_lut_words); -#pragma clang loop unroll(full) - for (int row = 0; row < row_tile; ++row) { - const std::int64_t input_offset = - (row_start + row) * args.in_features + k; - const __m512bh activations = (__m512bh)_mm512_loadu_si512( - static_cast(args.x) + input_offset); - accumulators[row] = - _mm512_dpbf16_ps(accumulators[row], activations, weights); + if constexpr (unroll_k) { +#if defined(__clang__) +#pragma clang loop unroll_count(2) +#elif defined(__GNUC__) +#pragma GCC unroll 2 +#endif + for (; k + 32 <= args.in_features; k += 32) { + accumulate_bf16_w4_chunk( + args, + packed_row, + row_start, + k, + centroid_lut_words, + nibble_mask, + accumulators); + } + } else { + for (; k + 32 <= args.in_features; k += 32) { + accumulate_bf16_w4_chunk( + args, + packed_row, + row_start, + k, + centroid_lut_words, + nibble_mask, + accumulators); } } @@ -252,38 +658,38 @@ ORBITQUANT_TARGET_AVX512_BF16 inline void packed_matmul_avx512_bf16_w4_rows( } } -template +template ORBITQUANT_TARGET_AVX512_BF16 ORBITQUANT_NOINLINE void packed_matmul_avx512_bf16_w4_typed( PackedMatmulArgs const &args, std::int64_t out_start, std::int64_t out_end) { - constexpr int kPrimaryRowTile = 8; + static_assert(primary_row_tile == 4 || primary_row_tile == 8); const std::int64_t packed_row_bytes = args.in_features / 2; for (std::int64_t out_col = out_start; out_col < out_end; ++out_col) { const auto *packed_row = args.packed_weight_indices + out_col * packed_row_bytes; std::int64_t row = 0; - for (; row + kPrimaryRowTile <= args.rows; row += kPrimaryRowTile) { - packed_matmul_avx512_bf16_w4_rows( + for (; row + primary_row_tile <= args.rows; row += primary_row_tile) { + packed_matmul_avx512_bf16_w4_rows( args, packed_row, out_col, row); } if (row + 4 <= args.rows) { - packed_matmul_avx512_bf16_w4_rows<4, aligned_k>( + packed_matmul_avx512_bf16_w4_rows<4, aligned_k, unroll_k>( args, packed_row, out_col, row); row += 4; } switch (args.rows - row) { case 3: - packed_matmul_avx512_bf16_w4_rows<3, aligned_k>( + packed_matmul_avx512_bf16_w4_rows<3, aligned_k, unroll_k>( args, packed_row, out_col, row); break; case 2: - packed_matmul_avx512_bf16_w4_rows<2, aligned_k>( + packed_matmul_avx512_bf16_w4_rows<2, aligned_k, unroll_k>( args, packed_row, out_col, row); break; case 1: - packed_matmul_avx512_bf16_w4_rows<1, aligned_k>( + packed_matmul_avx512_bf16_w4_rows<1, aligned_k, unroll_k>( args, packed_row, out_col, row); break; default: @@ -327,6 +733,30 @@ bool runtime_has_avx512_bf16() { #endif } +bool runtime_has_verified_amd_bf16_tuning() { + static const bool available = [] { + unsigned int eax = 0; + unsigned int ebx = 0; + unsigned int ecx = 0; + unsigned int edx = 0; + // CPUID vendor registers spell "AuthenticAMD" in EBX, EDX, ECX order. + if (!__get_cpuid(0, &eax, &ebx, &ecx, &edx) || + ebx != 0x68747541u || edx != 0x69746e65u || ecx != 0x444d4163u || + !__get_cpuid(1, &eax, &ebx, &ecx, &edx)) { + return false; + } + const unsigned int base_family = (eax >> 8) & 0xfu; + const unsigned int family = base_family == 0xfu + ? base_family + ((eax >> 20) & 0xffu) + : base_family; + // The four-row unrolled tile was measured on Zen 4 (EPYC 4564P); apply it + // to the AMD families that ship AVX-512 BF16 (Zen 4 = 19h, Zen 5 = 1Ah) + // instead of pinning the one benchmarked model. + return (family == 0x19u || family == 0x1au) && runtime_has_avx512_bf16(); + }(); + return available; +} + } // namespace bool packed_matmul_x86_avx512_available() { @@ -338,8 +768,26 @@ void packed_matmul_x86_avx512_range( PackedMatmulArgs const &args, std::int64_t out_start, std::int64_t out_end) { - if (!packed_matmul_x86_avx512_available() || args.bits != 4 || - args.in_features % 2 != 0) { + if (!packed_matmul_x86_avx512_available()) { + packed_matmul_scalar_range(args, out_start, out_end); + return; + } + if (args.bits == 2 && args.in_features % 4 == 0) { + packed_matmul_avx512_lowbit_dispatch( + args, out_start, out_end); + return; + } + if (args.bits == 3 && args.in_features % 8 == 0) { + packed_matmul_avx512_lowbit_dispatch( + args, out_start, out_end); + return; + } + if (args.bits == 6 && args.in_features % 4 == 0) { + packed_matmul_avx512_lowbit_dispatch( + args, out_start, out_end); + return; + } + if (args.bits != 4 || args.in_features % 2 != 0) { packed_matmul_scalar_range(args, out_start, out_end); return; } @@ -355,11 +803,21 @@ void packed_matmul_x86_avx512_range( case ScalarKind::BFloat16: #if defined(ORBITQUANT_HAS_AVX512_BF16_INTRINSICS) if (runtime_has_avx512_bf16()) { + const bool tuned_dimension = args.in_features == 1536 || + args.in_features == 1920 || args.in_features == 3072; + const bool use_tuned_shape = + runtime_has_verified_amd_bf16_tuning() && args.rows >= 16 && + tuned_dimension; if (args.in_features % 32 == 0) { - packed_matmul_avx512_bf16_w4_typed( - args, out_start, out_end); + if (use_tuned_shape) { + packed_matmul_avx512_bf16_w4_typed<4, true, true>( + args, out_start, out_end); + } else { + packed_matmul_avx512_bf16_w4_typed<8, true, false>( + args, out_start, out_end); + } } else { - packed_matmul_avx512_bf16_w4_typed( + packed_matmul_avx512_bf16_w4_typed<8, false, false>( args, out_start, out_end); } return; diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/quantize_activations_cpu.cpp b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/quantize_activations_cpu.cpp index 6236663..981bfd2 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/quantize_activations_cpu.cpp +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cpu/quantize_activations_cpu.cpp @@ -1,3 +1,4 @@ +#include "cpu_pool.h" #include "cpu_threads.h" #include "cpu_kernel_args.h" #include "packed_matmul_cpu.h" @@ -10,12 +11,14 @@ #include #include +#include #include #include #include #include -#include +#include #include +#include #include #if defined(__aarch64__) || defined(_M_ARM64) @@ -128,21 +131,30 @@ inline void fwht_block_portable(float *values, std::int64_t block_size) { #if defined(__aarch64__) || defined(_M_ARM64) inline void fwht_block_neon(float *values, std::int64_t block_size) { - for (std::int64_t half = 1; half < block_size; half *= 2) { + if (block_size < 4) { + fwht_block_portable(values, block_size); + return; + } + const uint32x4_t odd_lanes = {0u, 0xFFFFFFFFu, 0u, 0xFFFFFFFFu}; + const uint32x4_t high_lanes = {0u, 0u, 0xFFFFFFFFu, 0xFFFFFFFFu}; + for (std::int64_t base = 0; base < block_size; base += 4) { + float32x4_t value = vld1q_f32(values + base); + float32x4_t swapped = vrev64q_f32(value); + value = vbslq_f32( + odd_lanes, vsubq_f32(swapped, value), vaddq_f32(value, swapped)); + swapped = vextq_f32(value, value, 2); + value = vbslq_f32( + high_lanes, vsubq_f32(swapped, value), vaddq_f32(value, swapped)); + vst1q_f32(values + base, value); + } + for (std::int64_t half = 4; half < block_size; half *= 2) { for (std::int64_t base = 0; base < block_size; base += 2 * half) { - std::int64_t offset = 0; - for (; offset + 4 <= half; offset += 4) { + for (std::int64_t offset = 0; offset + 4 <= half; offset += 4) { const float32x4_t left = vld1q_f32(values + base + offset); const float32x4_t right = vld1q_f32(values + base + half + offset); vst1q_f32(values + base + offset, vaddq_f32(left, right)); vst1q_f32(values + base + half + offset, vsubq_f32(left, right)); } - for (; offset < half; ++offset) { - const float left = values[base + offset]; - const float right = values[base + half + offset]; - values[base + offset] = left + right; - values[base + half + offset] = left - right; - } } } } @@ -152,10 +164,26 @@ inline void fwht_block_neon(float *values, std::int64_t block_size) { ORBITQUANT_TARGET_AVX2 void fwht_block_avx2( float *values, std::int64_t block_size) { - for (std::int64_t half = 1; half < block_size; half *= 2) { + if (block_size < 8) { + fwht_block_portable(values, block_size); + return; + } + for (std::int64_t base = 0; base < block_size; base += 8) { + __m256 value = _mm256_loadu_ps(values + base); + __m256 swapped = _mm256_permute_ps(value, 0xB1); + value = _mm256_blend_ps( + _mm256_add_ps(value, swapped), _mm256_sub_ps(swapped, value), 0xAA); + swapped = _mm256_permute_ps(value, 0x4E); + value = _mm256_blend_ps( + _mm256_add_ps(value, swapped), _mm256_sub_ps(swapped, value), 0xCC); + swapped = _mm256_permute2f128_ps(value, value, 0x01); + value = _mm256_blend_ps( + _mm256_add_ps(value, swapped), _mm256_sub_ps(swapped, value), 0xF0); + _mm256_storeu_ps(values + base, value); + } + for (std::int64_t half = 8; half < block_size; half *= 2) { for (std::int64_t base = 0; base < block_size; base += 2 * half) { - std::int64_t offset = 0; - for (; offset + 8 <= half; offset += 8) { + for (std::int64_t offset = 0; offset + 8 <= half; offset += 8) { const __m256 left = _mm256_loadu_ps(values + base + offset); const __m256 right = _mm256_loadu_ps(values + base + half + offset); @@ -164,12 +192,6 @@ ORBITQUANT_TARGET_AVX2 void fwht_block_avx2( values + base + half + offset, _mm256_sub_ps(left, right)); } - for (; offset < half; ++offset) { - const float left = values[base + offset]; - const float right = values[base + half + offset]; - values[base + offset] = left + right; - values[base + half + offset] = left - right; - } } } } @@ -177,10 +199,32 @@ ORBITQUANT_TARGET_AVX2 void fwht_block_avx2( ORBITQUANT_TARGET_AVX512 void fwht_block_avx512( float *values, std::int64_t block_size) { - for (std::int64_t half = 1; half < block_size; half *= 2) { + if (block_size < 16) { + fwht_block_portable(values, block_size); + return; + } + // Butterfly widths 1-8 stay inside one 16-float register: swap the paired + // lanes with shuffles and blend the +/- results, one pass over memory + // instead of four. + for (std::int64_t base = 0; base < block_size; base += 16) { + __m512 value = _mm512_loadu_ps(values + base); + __m512 swapped = _mm512_permute_ps(value, 0xB1); + value = _mm512_mask_blend_ps( + 0xAAAA, _mm512_add_ps(value, swapped), _mm512_sub_ps(swapped, value)); + swapped = _mm512_permute_ps(value, 0x4E); + value = _mm512_mask_blend_ps( + 0xCCCC, _mm512_add_ps(value, swapped), _mm512_sub_ps(swapped, value)); + swapped = _mm512_shuffle_f32x4(value, value, 0xB1); + value = _mm512_mask_blend_ps( + 0xF0F0, _mm512_add_ps(value, swapped), _mm512_sub_ps(swapped, value)); + swapped = _mm512_shuffle_f32x4(value, value, 0x4E); + value = _mm512_mask_blend_ps( + 0xFF00, _mm512_add_ps(value, swapped), _mm512_sub_ps(swapped, value)); + _mm512_storeu_ps(values + base, value); + } + for (std::int64_t half = 16; half < block_size; half *= 2) { for (std::int64_t base = 0; base < block_size; base += 2 * half) { - std::int64_t offset = 0; - for (; offset + 16 <= half; offset += 16) { + for (std::int64_t offset = 0; offset + 16 <= half; offset += 16) { const __m512 left = _mm512_loadu_ps(values + base + offset); const __m512 right = _mm512_loadu_ps(values + base + half + offset); @@ -189,12 +233,6 @@ ORBITQUANT_TARGET_AVX512 void fwht_block_avx512( values + base + half + offset, _mm512_sub_ps(left, right)); } - for (; offset < half; ++offset) { - const float left = values[base + offset]; - const float right = values[base + half + offset]; - values[base + offset] = left + right; - values[base + half + offset] = left - right; - } } } } @@ -591,12 +629,16 @@ void quantize_lookup( quantize_lookup_scalar(args, scratch, output_offset, norm); } -template +template void quantize_activation_range( ActivationArgs const &args, + index_t const *permutation, std::int64_t row_start, std::int64_t row_end) { - std::vector scratch(static_cast(args.dim)); + thread_local std::vector scratch; + if (scratch.size() < static_cast(args.dim)) { + scratch.resize(static_cast(args.dim)); + } for (std::int64_t row = row_start; row < row_end; ++row) { const std::int64_t input_offset = row * args.dim; const float norm_squared = @@ -606,7 +648,7 @@ void quantize_activation_range( for (std::int64_t index = 0; index < args.dim; ++index) { const float value = load_scalar( args.x, - input_offset + args.permutation[index]); + input_offset + static_cast(permutation[index])); scratch[index] = value * static_cast(args.signs[index]) * inverse_norm; } for (std::int64_t block = 0; block < args.dim; block += args.block_size) { @@ -621,23 +663,37 @@ void quantize_activation_range( } } -void quantize_activation_dispatch( +template +void quantize_activation_dispatch_indexed( ActivationArgs const &args, + index_t const *permutation, std::int64_t row_start, std::int64_t row_end) { switch (args.scalar_kind) { case orbitquant::cpu::ScalarKind::Float32: - quantize_activation_range(args, row_start, row_end); + quantize_activation_range(args, permutation, row_start, row_end); return; case orbitquant::cpu::ScalarKind::Float16: - quantize_activation_range(args, row_start, row_end); + quantize_activation_range(args, permutation, row_start, row_end); return; case orbitquant::cpu::ScalarKind::BFloat16: - quantize_activation_range(args, row_start, row_end); + quantize_activation_range(args, permutation, row_start, row_end); return; } } +void quantize_activation_dispatch( + ActivationArgs const &args, + std::int64_t row_start, + std::int64_t row_end) { + if (args.permutation_i32 != nullptr) { + quantize_activation_dispatch_indexed( + args, args.permutation_i32, row_start, row_end); + return; + } + quantize_activation_dispatch_indexed(args, args.permutation, row_start, row_end); +} + void parallel_quantize_activations(ActivationArgs const &args) { const std::int64_t values = args.rows * args.dim; const int threads = values < 16'384 @@ -652,8 +708,8 @@ void parallel_quantize_activations(ActivationArgs const &args) { return; } - std::vector workers; - workers.reserve(threads); + std::vector> ranges; + ranges.reserve(threads); const std::int64_t rows_per_thread = (args.rows + threads - 1) / threads; for (int thread = 0; thread < threads; ++thread) { const std::int64_t start = thread * rows_per_thread; @@ -661,13 +717,13 @@ void parallel_quantize_activations(ActivationArgs const &args) { if (start >= end) { break; } - workers.emplace_back([&args, start, end] { - quantize_activation_dispatch(args, start, end); - }); - } - for (auto &worker : workers) { - worker.join(); + ranges.emplace_back(start, end); } + orbitquant::cpu::run_ranges( + ranges, + [&args](std::int64_t start, std::int64_t end) { + quantize_activation_dispatch(args, start, end); + }); } } // namespace @@ -705,8 +761,9 @@ void quantize_activations_cpu( STD_TORCH_CHECK(boundaries.is_contiguous(), "boundaries must be contiguous"); STD_TORCH_CHECK(out.scalar_type() == x.scalar_type(), "out dtype must match x dtype"); STD_TORCH_CHECK( - permutation.scalar_type() == ScalarType::Long, - "permutation must be int64"); + permutation.scalar_type() == ScalarType::Long || + permutation.scalar_type() == ScalarType::Int, + "permutation must be int32 or int64"); STD_TORCH_CHECK(signs.scalar_type() == ScalarType::Char, "signs must be int8"); STD_TORCH_CHECK(centroids.scalar_type() == ScalarType::Float, "centroids must be float32"); STD_TORCH_CHECK(boundaries.scalar_type() == ScalarType::Float, "boundaries must be float32"); @@ -729,15 +786,61 @@ void quantize_activations_cpu( STD_TORCH_CHECK(eps >= 0.0, "eps must be non-negative"); STD_TORCH_CHECK(inv_sqrt_block > 0.0, "inv_sqrt_block must be positive"); - const auto *permutation_values = permutation.const_data_ptr(); + const bool int32_permutation = permutation.scalar_type() == ScalarType::Int; + const std::int64_t *permutation_values = + int32_permutation ? nullptr : permutation.const_data_ptr(); + const std::int32_t *permutation_values_i32 = + int32_permutation ? permutation.const_data_ptr() : nullptr; const auto *sign_values = signs.const_data_ptr(); - for (int64_t index = 0; index < dim; ++index) { - STD_TORCH_CHECK( - permutation_values[index] >= 0 && permutation_values[index] < dim, - "permutation contains an out-of-range index"); - STD_TORCH_CHECK( - sign_values[index] == -1 || sign_values[index] == 1, - "signs must contain only -1 and 1"); + const auto permutation_at = [&](std::int64_t index) -> std::int64_t { + return int32_permutation + ? static_cast(permutation_values_i32[index]) + : permutation_values[index]; + }; + // The permutation/sign buffers are immutable module constants, so cache a + // fingerprint of validated buffers instead of re-scanning them per forward. + struct ValidatedEntry { + void const *permutation; + void const *signs; + std::int64_t dim; + std::int64_t head; + std::int64_t tail; + }; + static std::mutex validated_mutex; + static std::array validated_entries{}; + static std::size_t validated_cursor = 0; + const ValidatedEntry candidate{ + int32_permutation ? static_cast(permutation_values_i32) + : static_cast(permutation_values), + sign_values, + dim, + permutation_at(0), + permutation_at(dim - 1)}; + bool already_validated = false; + { + std::lock_guard lock(validated_mutex); + for (auto const &entry : validated_entries) { + if (entry.permutation == candidate.permutation && + entry.signs == candidate.signs && entry.dim == candidate.dim && + entry.head == candidate.head && entry.tail == candidate.tail) { + already_validated = true; + break; + } + } + } + if (!already_validated) { + for (int64_t index = 0; index < dim; ++index) { + const std::int64_t source_index = permutation_at(index); + STD_TORCH_CHECK( + source_index >= 0 && source_index < dim, + "permutation contains an out-of-range index"); + STD_TORCH_CHECK( + sign_values[index] == -1 || sign_values[index] == 1, + "signs must contain only -1 and 1"); + } + std::lock_guard lock(validated_mutex); + validated_entries[validated_cursor % validated_entries.size()] = candidate; + ++validated_cursor; } if (x.numel() == 0) { return; @@ -747,6 +850,7 @@ void quantize_activations_cpu( out.mutable_data_ptr(), x.const_data_ptr(), permutation_values, + permutation_values_i32, sign_values, centroids.const_data_ptr(), boundaries.const_data_ptr(), diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu index b52e0d4..15bc082 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu @@ -11,9 +11,34 @@ #include #include +#include using namespace nvcuda; +namespace { + +// Escape hatch (and A/B benchmarking toggle) for the cp.async-pipelined mma64 +// path: set ORBITQUANT_MMA64_DISABLE_PIPELINE=1 to force the legacy kernel. +inline bool orbitquant_mma64_pipeline_disabled() { + static const bool disabled = []() { + char const *value = std::getenv("ORBITQUANT_MMA64_DISABLE_PIPELINE"); + return value != nullptr && value[0] != '\0' && value[0] != '0'; + }(); + return disabled; +} + +// ORBITQUANT_MMA64_FORCE_PIPELINE=1 extends the pipeline to every eligible +// width (W4/W6) for re-benchmarking on other GPUs. +inline bool orbitquant_mma64_pipeline_forced() { + static const bool forced = []() { + char const *value = std::getenv("ORBITQUANT_MMA64_FORCE_PIPELINE"); + return value != nullptr && value[0] != '\0' && value[0] != '0'; + }(); + return forced; +} + +} // namespace + __device__ __forceinline__ uint32_t unpack_lowbit_index( uint8_t const *__restrict__ packed_weight_indices, int64_t value_offset, @@ -99,6 +124,43 @@ __device__ __forceinline__ void decode_mma64_weight_segment( } } +template +__device__ __forceinline__ void decode_mma64_weight_segment_from_stage( + mma_t *__restrict__ destination, + uint8_t const *__restrict__ staged_bytes, + float const *__restrict__ centroids, + float norm, + bool valid_segment) { + constexpr uint32_t mask = (1u << Bits) - 1u; + constexpr int word_count = (Values * Bits + 31) / 32; + uint32_t packed_words[word_count] = {}; + if (valid_segment) { + auto const *words = reinterpret_cast(staged_bytes); +#pragma unroll + for (int word = 0; word < word_count; ++word) { + packed_words[word] = words[word]; + } + } + +#pragma unroll + for (int index_offset = 0; index_offset < Values; ++index_offset) { + const int bit_start = index_offset * Bits; + const int word_index = bit_start >> 5; + const int shift = bit_start & 31; + uint32_t raw = packed_words[word_index] >> shift; + if (shift + Bits > 32) { + raw |= packed_words[word_index + 1] << (32 - shift); + } + const uint32_t codebook_index = raw & mask; + const float value = valid_segment ? norm * centroids[codebook_index] : 0.0f; + destination[index_offset] = orbitquant_mma_from_float(value); + } +} + +// WMMA fragments (and the bf16 variants in particular) only exist on sm80+. +// Multi-architecture builds still instantiate these templates for older +// targets, so the bodies compile to empty stubs there; the host dispatch +// requires compute capability >= 8 before launching either kernel. template __global__ void orbitquant_packed_matmul_mma64_kernel( storage_t *__restrict__ out, @@ -111,6 +173,7 @@ __global__ void orbitquant_packed_matmul_mma64_kernel( int64_t rows, int64_t out_features, int64_t in_features) { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 constexpr int tile_m = 128; constexpr int tile_n = 128; constexpr int tile_k = 64; @@ -228,6 +291,7 @@ __global__ void orbitquant_packed_matmul_mma64_kernel( } __syncwarp(); } +#endif // __CUDA_ARCH__ >= 800 } __device__ __forceinline__ void copy_async_16( @@ -257,6 +321,238 @@ __device__ __forceinline__ void wait_for_async_copies() { #endif } +// cp.async-pipelined variant of the mma64 kernel (sm80+, byte-aligned k-tile +// strides, i.e. every supported width except W3). Interior blocks double-buffer +// the X tile and the packed weight bytes so global loads overlap the MMA work; +// edge blocks keep the guarded synchronous path. +template +__global__ void orbitquant_packed_matmul_mma64_pipelined_kernel( + storage_t *__restrict__ out, + storage_t const *__restrict__ x, + uint8_t const *__restrict__ packed_weight_indices, + c10::BFloat16 const *__restrict__ row_norms, + float const *__restrict__ centroids, + storage_t const *__restrict__ bias, + bool has_bias, + int64_t rows, + int64_t out_features, + int64_t in_features) { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 + constexpr int tile_m = 128; + constexpr int tile_n = 128; + constexpr int tile_k = 64; + constexpr int padded_k = 72; + constexpr int warps_per_block = 8; + constexpr int warp_tile = 16; + constexpr int col_tiles = tile_n / warp_tile; + constexpr int x_vector_values = 8; + constexpr int x_vectors_per_row = tile_k / x_vector_values; + constexpr int seg_stride = tile_k * Bits / 8; + constexpr int weight_segment_values = 16; + constexpr int weight_segments_per_row = tile_k / weight_segment_values; + constexpr int segment_bytes = weight_segment_values * Bits / 8; + static_assert(sizeof(storage_t) == sizeof(mma_t)); + static_assert(Bits == 2 || Bits == 4 || Bits == 6); + + extern __shared__ __align__(16) uint8_t dynamic_shared[]; + mma_t *x_tiles = reinterpret_cast(dynamic_shared); + mma_t *weight_tile = x_tiles + 2 * tile_m * padded_k; + float *accumulator_tile = + reinterpret_cast(weight_tile + tile_n * padded_k); + uint8_t *packed_stage = reinterpret_cast( + accumulator_tile + warps_per_block * warp_tile * warp_tile); + + const int warp_id = threadIdx.x / warpSize; + const int lane = threadIdx.x & (warpSize - 1); + const int64_t block_row = int64_t(blockIdx.y) * tile_m; + const int64_t block_col = int64_t(blockIdx.x) * tile_n; + const int64_t packed_row_bytes = in_features * Bits / 8; + + wmma::fragment + accumulators[col_tiles]; +#pragma unroll + for (int col_tile = 0; col_tile < col_tiles; ++col_tile) { + wmma::fill_fragment(accumulators[col_tile], 0.0f); + } + + const bool interior = + block_row + tile_m <= rows && block_col + tile_n <= out_features; + + if (interior) { + constexpr int x_copies = tile_m * x_vectors_per_row; + constexpr int w_copies = tile_n * seg_stride / 16; + auto stage_tile = [&](int buffer, int64_t k_start) { + mma_t *x_destination = x_tiles + buffer * tile_m * padded_k; + for (int task = threadIdx.x; task < x_copies; task += blockDim.x) { + const int local_row = task / x_vectors_per_row; + const int local_vector = task - local_row * x_vectors_per_row; + copy_async_16( + x_destination + local_row * padded_k + local_vector * x_vector_values, + x + (block_row + local_row) * in_features + k_start + + local_vector * x_vector_values); + } + uint8_t *stage_destination = packed_stage + buffer * tile_n * seg_stride; + for (int task = threadIdx.x; task < w_copies; task += blockDim.x) { + const int byte_offset = task * 16; + const int local_col = byte_offset / seg_stride; + const int seg_byte = byte_offset - local_col * seg_stride; + copy_async_16( + stage_destination + byte_offset, + packed_weight_indices + (block_col + local_col) * packed_row_bytes + + (k_start * Bits) / 8 + seg_byte); + } + commit_async_copies(); + }; + + stage_tile(0, 0); + wait_for_async_copies(); + __syncthreads(); + + for (int64_t k_start = 0; k_start < in_features; k_start += tile_k) { + const int buffer = static_cast((k_start / tile_k) & 1); + const bool has_next = k_start + tile_k < in_features; + if (has_next) { + stage_tile(buffer ^ 1, k_start + tile_k); + } + + uint8_t const *stage_source = packed_stage + buffer * tile_n * seg_stride; + constexpr int weight_tasks = tile_n * weight_segments_per_row; + for (int task = threadIdx.x; task < weight_tasks; task += blockDim.x) { + const int local_col = task / weight_segments_per_row; + const int local_segment = task - local_col * weight_segments_per_row; + decode_mma64_weight_segment_from_stage( + weight_tile + local_col * padded_k + + local_segment * weight_segment_values, + stage_source + local_col * seg_stride + local_segment * segment_bytes, + centroids, + static_cast(row_norms[block_col + local_col]), + true); + } + __syncthreads(); + + mma_t const *x_source = x_tiles + buffer * tile_m * padded_k; +#pragma unroll + for (int local_k = 0; local_k < tile_k; local_k += warp_tile) { + wmma::fragment + lhs; + wmma::load_matrix_sync( + lhs, + x_source + warp_id * warp_tile * padded_k + local_k, + padded_k); +#pragma unroll + for (int col_tile = 0; col_tile < col_tiles; ++col_tile) { + wmma::fragment + rhs; + wmma::load_matrix_sync( + rhs, + weight_tile + col_tile * warp_tile * padded_k + local_k, + padded_k); + wmma::mma_sync( + accumulators[col_tile], lhs, rhs, accumulators[col_tile]); + } + } + __syncthreads(); + if (has_next) { + wait_for_async_copies(); + __syncthreads(); + } + } + } else { + mma_t *x_tile = x_tiles; + for (int64_t k_start = 0; k_start < in_features; k_start += tile_k) { + constexpr int x_vector_tasks = tile_m * x_vectors_per_row; + for (int task = threadIdx.x; task < x_vector_tasks; task += blockDim.x) { + const int local_row = task / x_vectors_per_row; + const int local_vector = task - local_row * x_vectors_per_row; + const int local_k = local_vector * x_vector_values; + const int64_t global_row = block_row + local_row; + auto *destination = reinterpret_cast( + x_tile + local_row * padded_k + local_k); + if (global_row < rows) { + auto const *source = reinterpret_cast( + x + global_row * in_features + k_start + local_k); + *destination = *source; + } else { + *destination = make_uint4(0, 0, 0, 0); + } + } + + constexpr int weight_tasks = tile_n * weight_segments_per_row; + for (int weight_task = threadIdx.x; weight_task < weight_tasks; + weight_task += blockDim.x) { + const int local_col = weight_task / weight_segments_per_row; + const int local_segment = + weight_task - local_col * weight_segments_per_row; + const int local_k = local_segment * weight_segment_values; + const int64_t global_col = block_col + local_col; + decode_mma64_weight_segment( + weight_tile + local_col * padded_k + local_k, + packed_weight_indices, + row_norms, + centroids, + global_col, + k_start + local_k, + in_features, + global_col < out_features); + } + __syncthreads(); + +#pragma unroll + for (int local_k = 0; local_k < tile_k; local_k += warp_tile) { + wmma::fragment + lhs; + wmma::load_matrix_sync( + lhs, + x_tile + warp_id * warp_tile * padded_k + local_k, + padded_k); +#pragma unroll + for (int col_tile = 0; col_tile < col_tiles; ++col_tile) { + wmma::fragment + rhs; + wmma::load_matrix_sync( + rhs, + weight_tile + col_tile * warp_tile * padded_k + local_k, + padded_k); + wmma::mma_sync( + accumulators[col_tile], lhs, rhs, accumulators[col_tile]); + } + } + __syncthreads(); + } + } + + float *warp_accumulator = accumulator_tile + warp_id * warp_tile * warp_tile; +#pragma unroll + for (int col_tile = 0; col_tile < col_tiles; ++col_tile) { + wmma::store_matrix_sync( + warp_accumulator, + accumulators[col_tile], + warp_tile, + wmma::mem_row_major); + __syncwarp(); + for (int offset = lane; offset < warp_tile * warp_tile; offset += warpSize) { + const int local_row = offset / warp_tile; + const int local_col = offset - local_row * warp_tile; + const int64_t global_row = block_row + warp_id * warp_tile + local_row; + const int64_t global_col = block_col + col_tile * warp_tile + local_col; + if (global_row < rows && global_col < out_features) { + float value = warp_accumulator[offset]; + if (has_bias) { + value += static_cast(bias[global_col]); + } + out[global_row * out_features + global_col] = + static_cast(value); + } + } + __syncwarp(); + } +#endif // __CUDA_ARCH__ >= 800 +} + __device__ __forceinline__ uint8_t orbitquant_bucketize_w4( float value, float const *__restrict__ boundaries) { @@ -267,12 +563,12 @@ __device__ __forceinline__ uint8_t orbitquant_bucketize_w4( return static_cast(index); } -template +template __global__ void orbitquant_rpbh_quantize_pack_w4_kernel( uint8_t *__restrict__ packed_out, float *__restrict__ norms_out, storage_t const *__restrict__ x, - int64_t const *__restrict__ permutation, + index_t const *__restrict__ permutation, int8_t const *__restrict__ signs, float const *__restrict__ boundaries, float eps, @@ -345,12 +641,12 @@ __global__ void orbitquant_rpbh_quantize_pack_w4_kernel( } } -template +template __global__ void orbitquant_rpbh_quantize_int8_kernel( int8_t *__restrict__ int8_out, float *__restrict__ norms_out, storage_t const *__restrict__ x, - int64_t const *__restrict__ permutation, + index_t const *__restrict__ permutation, int8_t const *__restrict__ signs, float const *__restrict__ boundaries, int8_t const *__restrict__ codes, @@ -451,6 +747,7 @@ __global__ void orbitquant_packed_w4a4_int8_mma_kernel( int64_t rows, int64_t out_features, int64_t in_features) { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 constexpr int tile_k = 64; constexpr int packed_tile_k = tile_k / 2; constexpr int padded_k = 80; @@ -693,6 +990,7 @@ __global__ void orbitquant_packed_w4a4_int8_mma_kernel( } __syncwarp(); } +#endif // __CUDA_ARCH__ >= 800 } template @@ -707,6 +1005,7 @@ __global__ void orbitquant_packed_matmul_wmma_bf16_kernel( int64_t rows, int64_t out_features, int64_t in_features) { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 constexpr int tile = 16; constexpr int col_tiles = 4; constexpr int warps_per_block = 8; @@ -788,6 +1087,7 @@ __global__ void orbitquant_packed_matmul_wmma_bf16_kernel( } } } +#endif // __CUDA_ARCH__ >= 800 } template @@ -1081,6 +1381,7 @@ void matmul_packed_weight( sizeof(float); const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaDeviceProp *mma_properties = at::cuda::getCurrentDeviceProperties(); if (x.size(0) <= 8) { constexpr int channels_per_warp = 4; @@ -1109,7 +1410,8 @@ void matmul_packed_weight( return; } - if (x.scalar_type() == at::kBFloat16 && x.size(0) >= 9) { + if (x.scalar_type() == at::kBFloat16 && x.size(0) >= 9 && + mma_properties->major >= 8) { if (in_features % 64 == 0 && (bits == 2 || bits == 3 || bits == 4 || bits == 6)) { constexpr int mma_tile_m = 128; @@ -1118,6 +1420,53 @@ void matmul_packed_weight( const dim3 mma_grid( (out_features + mma_tile_n - 1) / mma_tile_n, (x.size(0) + mma_tile_m - 1) / mma_tile_m); + +#define ORBITQUANT_LAUNCH_MMA64_PIPELINED(STORAGE_TYPE, MMA_TYPE, BITS_VALUE) \ + do { \ + constexpr int kSegStride = 64 * (BITS_VALUE) / 8; \ + const int pipelined_shared_bytes = static_cast( \ + 2 * 128 * 72 * sizeof(MMA_TYPE) + 128 * 72 * sizeof(MMA_TYPE) + \ + 8 * 16 * 16 * sizeof(float) + 2 * 128 * kSegStride); \ + if (pipelined_shared_bytes > mma_properties->sharedMemPerBlock) { \ + C10_CUDA_CHECK(cudaFuncSetAttribute( \ + orbitquant_packed_matmul_mma64_pipelined_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + pipelined_shared_bytes)); \ + } \ + orbitquant_packed_matmul_mma64_pipelined_kernel \ + <<>>( \ + reinterpret_cast(out.data_ptr()), \ + reinterpret_cast(x.data_ptr()), \ + packed_weight_indices.data_ptr(), \ + row_norms.data_ptr(), \ + centroids.data_ptr(), \ + has_bias ? bias.data_ptr() : nullptr, \ + has_bias, \ + x.size(0), \ + out_features, \ + in_features); \ + } while (0) +// The cp.async pipeline wins when the launch is latency-bound: with cold L2 +// (each layer's weights are evicted between calls in a real model) it is +// 1.35-1.46x faster for W2/W4/W6 on RTX 4060 Ti, A40, and RTX 4090 whenever +// the grid fits in one wave (blocks <= SM count), and it regresses up to 15% +// once the grid oversubscribes the device. Gate on grid size; measured +// 2026-07 across the three architectures above. +// ORBITQUANT_MMA64_FORCE_PIPELINE=1 bypasses the grid gate, +// ORBITQUANT_MMA64_DISABLE_PIPELINE=1 forces the legacy kernel everywhere. +#define ORBITQUANT_MMA64_USE_PIPELINE(BITS_VALUE) \ + ((orbitquant_mma64_pipeline_forced() || \ + static_cast(mma_grid.x) * mma_grid.y <= \ + mma_properties->multiProcessorCount) && \ + !orbitquant_mma64_pipeline_disabled() && \ + mma_properties->major >= 8 && \ + (2 * 128 * 72 * 2 + 128 * 72 * 2 + 8 * 16 * 16 * 4 + \ + 2 * 128 * (64 * (BITS_VALUE) / 8)) <= \ + static_cast(mma_properties->sharedMemPerBlockOptin)) + #define ORBITQUANT_LAUNCH_MMA64_BF16(BITS_VALUE) \ orbitquant_packed_matmul_mma64_kernel<<= 9) { + if (x.scalar_type() == at::kHalf && x.size(0) >= 9 && + mma_properties->major >= 8) { if (in_features % 64 == 0 && (bits == 2 || bits == 3 || bits == 4 || bits == 6)) { constexpr int mma_tile_m = 128; @@ -1225,19 +1587,33 @@ void matmul_packed_weight( in_features) switch (bits) { case 2: - ORBITQUANT_LAUNCH_MMA64_HALF(2); + if (ORBITQUANT_MMA64_USE_PIPELINE(2)) { + ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::Half, half, 2); + } else { + ORBITQUANT_LAUNCH_MMA64_HALF(2); + } break; case 3: ORBITQUANT_LAUNCH_MMA64_HALF(3); break; case 4: - ORBITQUANT_LAUNCH_MMA64_HALF(4); + if (ORBITQUANT_MMA64_USE_PIPELINE(4)) { + ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::Half, half, 4); + } else { + ORBITQUANT_LAUNCH_MMA64_HALF(4); + } break; case 6: - ORBITQUANT_LAUNCH_MMA64_HALF(6); + if (ORBITQUANT_MMA64_USE_PIPELINE(6)) { + ORBITQUANT_LAUNCH_MMA64_PIPELINED(c10::Half, half, 6); + } else { + ORBITQUANT_LAUNCH_MMA64_HALF(6); + } break; } #undef ORBITQUANT_LAUNCH_MMA64_HALF +#undef ORBITQUANT_MMA64_USE_PIPELINE +#undef ORBITQUANT_LAUNCH_MMA64_PIPELINED C10_CUDA_KERNEL_LAUNCH_CHECK(); return; } @@ -1520,8 +1896,9 @@ void quantize_activations_packed_w4( x.scalar_type() == torch::kBFloat16 || x.scalar_type() == torch::kHalf, "x must be bfloat16 or float16"); TORCH_CHECK( - permutation.scalar_type() == torch::kLong, - "permutation must be int64"); + permutation.scalar_type() == torch::kLong || + permutation.scalar_type() == torch::kInt, + "permutation must be int32 or int64"); TORCH_CHECK(signs.scalar_type() == torch::kChar, "signs must be int8"); TORCH_CHECK(boundaries.scalar_type() == torch::kFloat, "boundaries must be float32"); TORCH_CHECK(x.dim() == 2, "x must be rank 2"); @@ -1561,49 +1938,59 @@ void quantize_activations_packed_w4( const dim3 grid(static_cast(rows)); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); -#define ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, DIM_VALUE) \ +#define ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, DIM_VALUE) \ if (shared_bytes > properties->sharedMemPerBlock) { \ C10_CUDA_CHECK(cudaFuncSetAttribute( \ - orbitquant_rpbh_quantize_pack_w4_kernel, \ + orbitquant_rpbh_quantize_pack_w4_kernel, \ cudaFuncAttributeMaxDynamicSharedMemorySize, \ shared_bytes)); \ } \ - orbitquant_rpbh_quantize_pack_w4_kernel \ + orbitquant_rpbh_quantize_pack_w4_kernel \ <<>>( \ packed_out.data_ptr(), \ norms_out.data_ptr(), \ reinterpret_cast(x.data_ptr()), \ - permutation.data_ptr(), \ + permutation.data_ptr(), \ signs.data_ptr(), \ boundaries.data_ptr(), \ static_cast(eps), \ static_cast(inv_sqrt_block), \ rows) -#define ORBITQUANT_DISPATCH_RPBH_PACK_W4(STORAGE_TYPE) \ +#define ORBITQUANT_DISPATCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE) \ switch (dim) { \ case 512: \ - ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, 512); \ + ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 512); \ break; \ case 1024: \ - ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, 1024); \ + ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 1024); \ break; \ case 2048: \ - ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, 2048); \ + ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 2048); \ break; \ case 4096: \ - ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, 4096); \ + ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 4096); \ break; \ case 8192: \ - ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, 8192); \ + ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 8192); \ break; \ case 16384: \ - ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, 16384); \ + ORBITQUANT_LAUNCH_RPBH_PACK_W4(STORAGE_TYPE, INDEX_TYPE, 16384); \ break; \ } + const bool int32_permutation = permutation.scalar_type() == torch::kInt; if (x.scalar_type() == torch::kBFloat16) { - ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::BFloat16); + if (int32_permutation) { + ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::BFloat16, int32_t); + } else { + ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::BFloat16, int64_t); + } } else { - ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::Half); + if (int32_permutation) { + ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::Half, int32_t); + } else { + ORBITQUANT_DISPATCH_RPBH_PACK_W4(c10::Half, int64_t); + } } #undef ORBITQUANT_DISPATCH_RPBH_PACK_W4 #undef ORBITQUANT_LAUNCH_RPBH_PACK_W4 @@ -1641,8 +2028,9 @@ void quantize_activations_int8( x.scalar_type() == torch::kBFloat16 || x.scalar_type() == torch::kHalf, "x must be bfloat16 or float16"); TORCH_CHECK( - permutation.scalar_type() == torch::kLong, - "permutation must be int64"); + permutation.scalar_type() == torch::kLong || + permutation.scalar_type() == torch::kInt, + "permutation must be int32 or int64"); TORCH_CHECK(signs.scalar_type() == torch::kChar, "signs must be int8"); TORCH_CHECK(boundaries.scalar_type() == torch::kFloat, "boundaries must be float32"); TORCH_CHECK(codes.scalar_type() == torch::kChar, "codes must be int8"); @@ -1686,55 +2074,64 @@ void quantize_activations_int8( const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); #define ORBITQUANT_LAUNCH_RPBH_INT8( \ - STORAGE_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE) \ + STORAGE_TYPE, INDEX_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE) \ if (shared_bytes > properties->sharedMemPerBlock) { \ C10_CUDA_CHECK(cudaFuncSetAttribute( \ orbitquant_rpbh_quantize_int8_kernel< \ - STORAGE_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE>, \ + STORAGE_TYPE, INDEX_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE>, \ cudaFuncAttributeMaxDynamicSharedMemorySize, \ shared_bytes)); \ } \ orbitquant_rpbh_quantize_int8_kernel< \ - STORAGE_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE> \ + STORAGE_TYPE, INDEX_TYPE, DIM_VALUE, ORBIT_BLOCK_VALUE> \ <<>>( \ int8_out.data_ptr(), \ norms_out.data_ptr(), \ reinterpret_cast(x.data_ptr()), \ - permutation.data_ptr(), \ + permutation.data_ptr(), \ signs.data_ptr(), \ boundaries.data_ptr(), \ codes.data_ptr(), \ static_cast(eps), \ static_cast(inv_sqrt_block), \ rows) -#define ORBITQUANT_DISPATCH_RPBH_INT8(STORAGE_TYPE) \ +#define ORBITQUANT_DISPATCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE) \ switch (dim) { \ case 512: \ - ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, 512, 512); \ + ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 512, 512); \ break; \ case 1024: \ - ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, 1024, 1024); \ + ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 1024, 1024); \ break; \ case 2048: \ - ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, 2048, 2048); \ + ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 2048, 2048); \ break; \ case 4096: \ - ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, 4096, 4096); \ + ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 4096, 4096); \ break; \ case 8192: \ - ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, 8192, 8192); \ + ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 8192, 8192); \ break; \ case 12288: \ - ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, 12288, 4096); \ + ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 12288, 4096); \ break; \ case 16384: \ - ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, 16384, 16384); \ + ORBITQUANT_LAUNCH_RPBH_INT8(STORAGE_TYPE, INDEX_TYPE, 16384, 16384); \ break; \ } + const bool int32_permutation = permutation.scalar_type() == torch::kInt; if (x.scalar_type() == torch::kBFloat16) { - ORBITQUANT_DISPATCH_RPBH_INT8(c10::BFloat16); + if (int32_permutation) { + ORBITQUANT_DISPATCH_RPBH_INT8(c10::BFloat16, int32_t); + } else { + ORBITQUANT_DISPATCH_RPBH_INT8(c10::BFloat16, int64_t); + } } else { - ORBITQUANT_DISPATCH_RPBH_INT8(c10::Half); + if (int32_permutation) { + ORBITQUANT_DISPATCH_RPBH_INT8(c10::Half, int32_t); + } else { + ORBITQUANT_DISPATCH_RPBH_INT8(c10::Half, int64_t); + } } #undef ORBITQUANT_DISPATCH_RPBH_INT8 #undef ORBITQUANT_LAUNCH_RPBH_INT8 diff --git a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.metal b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.metal index 81ae2f2..31b12b7 100644 --- a/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.metal +++ b/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.metal @@ -54,13 +54,13 @@ inline void packed_matmul_small_rows_value( raw |= uint(packed_weight_indices[byte_index + 1]) << 8; } const uint index = (raw >> bit_offset) & mask; - accumulators[col_offset] += x_value * norms[col_offset] * centroids[index]; + accumulators[col_offset] += x_value * centroids[index]; } } #pragma clang loop unroll(full) for (uint col_offset = 0; col_offset < packed_small_cols; ++col_offset) { - const float value = simd_sum(accumulators[col_offset]); + const float value = simd_sum(accumulators[col_offset]) * norms[col_offset]; const long col = col_start + long(col_offset); if (lane == 0 && col < params.out_features) { out[row * params.out_features + col] = diff --git a/native-kernels/orbitquant-packed-matmul/tests/test_packed_matmul.py b/native-kernels/orbitquant-packed-matmul/tests/test_packed_matmul.py index 4464456..44716b0 100644 --- a/native-kernels/orbitquant-packed-matmul/tests/test_packed_matmul.py +++ b/native-kernels/orbitquant-packed-matmul/tests/test_packed_matmul.py @@ -67,11 +67,7 @@ def _runnable_cpu_isas() -> list[str]: isas = ["scalar"] if machine in {"x86_64", "amd64"} and capability in {"AVX2", "AVX512"}: isas.append("avx2") - if ( - machine in {"x86_64", "amd64"} - and capability == "AVX512" - and platform.system() != "Windows" - ): + if machine in {"x86_64", "amd64"} and capability == "AVX512" and platform.system() != "Windows": # The MSVC wheel currently ships the separately compiled AVX2 TU; the # AVX-512 implementation uses GCC/Clang per-function target attributes. isas.append("avx512") @@ -182,12 +178,107 @@ def test_cpu_runtime_isa_dispatch_matches_scalar_reference(monkeypatch) -> None: ) for isa in _runnable_cpu_isas()[1:]: - torch.testing.assert_close( - activations[isa], activations["scalar"], atol=2e-6, rtol=2e-6 + torch.testing.assert_close(activations[isa], activations["scalar"], atol=2e-6, rtol=2e-6) + torch.testing.assert_close(outputs[isa], outputs["scalar"], atol=2e-5, rtol=2e-5) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("bits", [2, 3, 6]) +@pytest.mark.parametrize("in_features", [64, 84]) +def test_cpu_isa_matmul_matches_scalar_for_low_bit_widths( + monkeypatch, bits: int, in_features: int +) -> None: + if not supports_device("cpu"): + pytest.skip("the built variant has no native CPU backend") + torch.manual_seed(47 + bits + in_features) + rows = 9 + out_features = 11 + levels = 2**bits + x = torch.randn(rows, in_features) + indices = torch.randint(0, levels, (out_features, in_features), dtype=torch.uint8) + packed = _pack(indices, bits) + row_norms = torch.linspace(0.5, 1.5, out_features) + centroids = torch.tanh(torch.linspace(-1.7, 1.7, levels)) + bias = torch.randn(out_features) + + outputs = {} + for isa in _runnable_cpu_isas(): + monkeypatch.setenv("ORBITQUANT_CPU_ISA", isa) + outputs[isa] = matmul_packed_weight( + x, + packed, + row_norms, + centroids, + bits=bits, + out_features=out_features, + in_features=in_features, + bias=bias, ) + + for isa in _runnable_cpu_isas()[1:]: torch.testing.assert_close(outputs[isa], outputs["scalar"], atol=2e-5, rtol=2e-5) +@pytest.mark.kernels_ci +@pytest.mark.parametrize("rows", [16, 24, 32, 33]) +def test_cpu_avx2_bf16_realistic_row_tiles_match_reference(monkeypatch, rows: int) -> None: + if "avx2" not in _runnable_cpu_isas() or not supports_device("cpu"): + pytest.skip("the built variant has no runnable AVX2 CPU path") + monkeypatch.setenv("ORBITQUANT_CPU_ISA", "avx2") + torch.manual_seed(59 + rows) + in_features = 1536 + out_features = 17 + x = torch.randn(rows, in_features, dtype=torch.bfloat16) + indices = torch.randint(0, 16, (out_features, in_features), dtype=torch.uint8) + packed = _pack(indices, 4) + row_norms = torch.linspace(0.5, 1.5, out_features) + centroids = torch.tanh(torch.linspace(-1.7, 1.7, 16)) + + expected_weight = (row_norms[:, None] * centroids[indices.long()]).to(torch.bfloat16) + expected = torch.nn.functional.linear(x, expected_weight) + actual = matmul_packed_weight( + x, + packed, + row_norms, + centroids, + bits=4, + out_features=out_features, + in_features=in_features, + ) + + torch.testing.assert_close(actual, expected, atol=0.25, rtol=3e-2) + + +@pytest.mark.kernels_ci +def test_cpu_avx512_bf16_realistic_row_tile_matches_reference(monkeypatch) -> None: + if "avx512" not in _runnable_cpu_isas() or not supports_device("cpu"): + pytest.skip("the built variant has no runnable AVX-512 CPU path") + monkeypatch.setenv("ORBITQUANT_CPU_ISA", "avx512") + torch.manual_seed(59) + rows = 32 + in_features = 1536 + out_features = 17 + x = torch.randn(rows, in_features, dtype=torch.bfloat16) + indices = torch.randint(0, 16, (out_features, in_features), dtype=torch.uint8) + packed = _pack(indices, 4) + row_norms = torch.linspace(0.5, 1.5, out_features) + centroids = torch.tanh(torch.linspace(-1.7, 1.7, 16)) + + expected_weight = (row_norms[:, None] * centroids[indices.long()]).to(torch.bfloat16) + expected = torch.nn.functional.linear(x, expected_weight) + actual = matmul_packed_weight( + x, + packed, + row_norms, + centroids, + bits=4, + out_features=out_features, + in_features=in_features, + ) + + torch.testing.assert_close(actual, expected, atol=0.125, rtol=3e-2) + + @pytest.mark.kernels_ci @pytest.mark.parametrize("group_size", [8, 64]) @pytest.mark.parametrize("with_bias", [False, True]) @@ -215,9 +306,9 @@ def test_matmul_packed_adaln_cpu_matches_independent_bf16_reference( bias = torch.randn(out_features, dtype=torch.bfloat16) if with_bias else None signed = indices.to(torch.int16).sub(8).float() - weight = (signed * scales.float()[..., None]).reshape( - out_features, padded_in_features - )[:, :in_features] + weight = (signed * scales.float()[..., None]).reshape(out_features, padded_in_features)[ + :, :in_features + ] expected = torch.nn.functional.linear(x, weight.to(torch.bfloat16), bias) actual = matmul_packed_adaln_int4_cpu( x, @@ -268,9 +359,7 @@ def test_cpu_adaln_runtime_isa_dispatch_matches_scalar_reference(monkeypatch) -> ) for isa in _runnable_cpu_isas()[1:]: - torch.testing.assert_close( - outputs[isa], outputs["scalar"], atol=3e-2, rtol=3e-2 - ) + torch.testing.assert_close(outputs[isa], outputs["scalar"], atol=3e-2, rtol=3e-2) @pytest.mark.kernels_ci @@ -347,9 +436,7 @@ def test_matmul_packed_weight_short_sequence_matches_reference( bias.cpu(), ).float() else: - expected = torch.nn.functional.linear( - x.float().cpu(), expected_weight, bias.float().cpu() - ) + expected = torch.nn.functional.linear(x.float().cpu(), expected_weight, bias.float().cpu()) actual = matmul_packed_weight( x, packed, @@ -646,6 +733,43 @@ def test_quantize_activations_packed_w4_matches_torch_reference( assert torch.allclose(norms, expected_norms, atol=1e-6, rtol=1e-6) +@pytest.mark.kernels_ci +def test_quantize_activations_packed_w4_accepts_int32_permutation() -> None: + device = _cuda_device() + dim = 512 + torch.manual_seed(0) + x = torch.randn((3, dim), device=device, dtype=torch.bfloat16) + permutation = torch.randperm(dim, device=device) + signs = torch.where( + torch.arange(dim, device=device) % 2 == 0, + torch.ones(dim, device=device, dtype=torch.int8), + -torch.ones(dim, device=device, dtype=torch.int8), + ) + boundaries = torch.linspace(-0.2, 0.2, 15, device=device) + + packed_int64, norms_int64 = quantize_activations_packed_w4( + x, + permutation, + signs, + boundaries, + eps=1e-12, + inv_sqrt_block=dim**-0.5, + threads=256, + ) + packed_int32, norms_int32 = quantize_activations_packed_w4( + x, + permutation.to(torch.int32), + signs, + boundaries, + eps=1e-12, + inv_sqrt_block=dim**-0.5, + threads=256, + ) + + assert torch.equal(packed_int32, packed_int64) + assert torch.equal(norms_int32, norms_int64) + + @pytest.mark.kernels_ci @pytest.mark.parametrize(("dim", "threads"), [(512, 128), (4096, 256), (16384, 512)]) def test_quantize_activations_int8_matches_packed_codes( @@ -736,13 +860,9 @@ def test_quantize_activations_int8_matches_blocked_rpbh_reference() -> None: blocks = work.reshape(rows, -1, width * 2) left = blocks[..., :width] right = blocks[..., width:] - work = torch.cat((left + right, left - right), dim=-1).reshape( - rows, -1, block_size - ) + work = torch.cat((left + right, left - right), dim=-1).reshape(rows, -1, block_size) width *= 2 - indices = torch.bucketize( - work.reshape(rows, dim) * (block_size**-0.5), boundaries - ) + indices = torch.bucketize(work.reshape(rows, dim) * (block_size**-0.5), boundaries) expected = codes[indices] assert torch.equal(quantized, expected) diff --git a/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py b/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py index f798ec3..2534ecf 100644 --- a/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py +++ b/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py @@ -87,7 +87,8 @@ def quantize_activations_cpu( original_shape = x.shape values = x.contiguous().reshape(-1, dim) out = torch.empty_like(values) - permutation_values = permutation.to(device="cpu", dtype=torch.int64).contiguous() + permutation_dtype = torch.int32 if permutation.dtype == torch.int32 else torch.int64 + permutation_values = permutation.to(device="cpu", dtype=permutation_dtype).contiguous() sign_values = signs.to(device="cpu", dtype=torch.int8).contiguous() centroid_values = centroids.to(device="cpu", dtype=torch.float32).contiguous() boundary_values = boundaries.to(device="cpu", dtype=torch.float32).contiguous() @@ -312,7 +313,8 @@ def quantize_activations_packed_w4( values = x.contiguous().reshape(-1, dim) packed = torch.empty((values.shape[0], dim // 2), device=x.device, dtype=torch.uint8) norms = torch.empty(values.shape[0], device=x.device, dtype=torch.float32) - permutation_values = permutation.to(device=x.device, dtype=torch.int64).contiguous() + permutation_dtype = torch.int32 if permutation.dtype == torch.int32 else torch.int64 + permutation_values = permutation.to(device=x.device, dtype=permutation_dtype).contiguous() sign_values = signs.to(device=x.device, dtype=torch.int8).contiguous() boundary_values = boundaries.to(device=x.device, dtype=torch.float32).contiguous() ops.quantize_activations_packed_w4( @@ -366,7 +368,8 @@ def quantize_activations_int8( values = x.contiguous().reshape(-1, dim) quantized = torch.empty(values.shape, device=x.device, dtype=torch.int8) norms = torch.empty(values.shape[0], device=x.device, dtype=torch.float32) - permutation_values = permutation.to(device=x.device, dtype=torch.int64).contiguous() + permutation_dtype = torch.int32 if permutation.dtype == torch.int32 else torch.int64 + permutation_values = permutation.to(device=x.device, dtype=permutation_dtype).contiguous() sign_values = signs.to(device=x.device, dtype=torch.int8).contiguous() boundary_values = boundaries.to(device=x.device, dtype=torch.float32).contiguous() code_values = codes.to(device=x.device, dtype=torch.int8).contiguous() diff --git a/native-kernels/orbitquant-vulkan/CMakeLists.txt b/native-kernels/orbitquant-vulkan/CMakeLists.txt new file mode 100644 index 0000000..7fb15fd --- /dev/null +++ b/native-kernels/orbitquant-vulkan/CMakeLists.txt @@ -0,0 +1,158 @@ +cmake_minimum_required(VERSION 3.19) +project(orbitquant_vulkan LANGUAGES C CXX) + +set( + ORBITQUANT_EXECUTORCH_ROOT + "" + CACHE PATH "Path to the ExecuTorch source checkout used to build Vulkan" +) +if(NOT ORBITQUANT_EXECUTORCH_ROOT) + message( + FATAL_ERROR + "Set ORBITQUANT_EXECUTORCH_ROOT to a post-2026-07-01 ExecuTorch source checkout." + ) +endif() + +find_package(executorch CONFIG REQUIRED COMPONENTS vulkan_backend) + +# ExecuTorch's shader helper reads this source-tree variable internally. +set(EXECUTORCH_ROOT ${ORBITQUANT_EXECUTORCH_ROOT}) +if(NOT PYTHON_EXECUTABLE) + find_package(Python3 REQUIRED COMPONENTS Interpreter) + set(PYTHON_EXECUTABLE ${Python3_EXECUTABLE}) +endif() + +include(${ORBITQUANT_EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) +include( + ${ORBITQUANT_EXECUTORCH_ROOT}/backends/vulkan/cmake/ShaderLibrary.cmake +) + +set( + VULKAN_THIRD_PARTY_PATH + ${ORBITQUANT_EXECUTORCH_ROOT}/backends/vulkan/third-party +) +set(VULKAN_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/Vulkan-Headers/include) +set(VOLK_PATH ${VULKAN_THIRD_PARTY_PATH}/volk) +set(VMA_PATH ${VULKAN_THIRD_PARTY_PATH}/VulkanMemoryAllocator) + +set(VULKAN_CXX_FLAGS "$<$>:-fexceptions>") +list(APPEND VULKAN_CXX_FLAGS "$<$:/EHsc>") +set( + VULKAN_COMPILE_DEFINITIONS + USE_VULKAN_WRAPPER + USE_VULKAN_VOLK + "$<$:NOMINMAX>" + "$<$:WIN32_LEAN_AND_MEAN>" +) + +gen_vulkan_shader_lib_cpp(${CMAKE_CURRENT_SOURCE_DIR}/glsl) +vulkan_shader_lib(orbitquant_vulkan_shaderlib ${generated_spv_cpp}) +target_compile_definitions( + orbitquant_vulkan_shaderlib PRIVATE ${VULKAN_COMPILE_DEFINITIONS} +) +target_include_directories( + orbitquant_vulkan_shaderlib PRIVATE ${ORBITQUANT_EXECUTORCH_ROOT}/src +) + +add_library(orbitquant_vulkan_ops STATIC OrbitLinear.cpp) +target_compile_features(orbitquant_vulkan_ops PRIVATE cxx_std_17) +target_compile_options(orbitquant_vulkan_ops PRIVATE ${VULKAN_CXX_FLAGS}) +target_compile_definitions( + orbitquant_vulkan_ops PRIVATE ${VULKAN_COMPILE_DEFINITIONS} +) +target_include_directories( + orbitquant_vulkan_ops + PRIVATE + ${ORBITQUANT_EXECUTORCH_ROOT}/src + ${VULKAN_HEADERS_PATH} + ${VOLK_PATH} + ${VMA_PATH} +) +target_link_libraries( + orbitquant_vulkan_ops + PRIVATE executorch_core vulkan_backend orbitquant_vulkan_shaderlib +) +executorch_target_link_options_shared_lib(orbitquant_vulkan_ops) + +install( + TARGETS orbitquant_vulkan_ops orbitquant_vulkan_shaderlib + ARCHIVE DESTINATION lib +) + +option(ORBITQUANT_VULKAN_BUILD_TESTS "Build the Vulkan hardware test" OFF) +if(ORBITQUANT_VULKAN_BUILD_TESTS) + add_executable( + orbitquant_vulkan_test + tests/OrbitLinearTest.cpp + ${ORBITQUANT_EXECUTORCH_ROOT}/backends/vulkan/test/custom_ops/utils.cpp + ${ORBITQUANT_EXECUTORCH_ROOT}/backends/vulkan/test/custom_ops/conv2d_utils.cpp + ${ORBITQUANT_EXECUTORCH_ROOT}/backends/vulkan/test/custom_ops/cm_utils.cpp + ) + target_compile_features(orbitquant_vulkan_test PRIVATE cxx_std_17) + target_compile_options(orbitquant_vulkan_test PRIVATE ${VULKAN_CXX_FLAGS}) + target_compile_definitions( + orbitquant_vulkan_test PRIVATE ${VULKAN_COMPILE_DEFINITIONS} + ) + target_include_directories( + orbitquant_vulkan_test + PRIVATE + ${ORBITQUANT_EXECUTORCH_ROOT}/src + ${ORBITQUANT_EXECUTORCH_ROOT}/backends/vulkan/test/custom_ops + ${VULKAN_HEADERS_PATH} + ${VOLK_PATH} + ${VMA_PATH} + ) + target_link_libraries( + orbitquant_vulkan_test + PRIVATE + executorch_core + vulkan_backend + orbitquant_vulkan_shaderlib + orbitquant_vulkan_ops + ) +endif() + +option( + ORBITQUANT_VULKAN_BUILD_PTE_RUNNER + "Build an ExecuTorch PTE runner with the OrbitQuant Vulkan op registered" + OFF +) +if(ORBITQUANT_VULKAN_BUILD_PTE_RUNNER) + set(GFLAGS_BUILD_TESTING OFF) + set(GFLAGS_BUILD_PACKAGING OFF) + set(CMAKE_POLICY_VERSION_MINIMUM 3.5) + add_subdirectory( + ${ORBITQUANT_EXECUTORCH_ROOT}/third-party/gflags + ${CMAKE_CURRENT_BINARY_DIR}/third-party/gflags + EXCLUDE_FROM_ALL + ) + + add_executable( + orbitquant_vulkan_pte_runner + ${ORBITQUANT_EXECUTORCH_ROOT}/examples/portable/executor_runner/executor_runner.cpp + ) + target_compile_features(orbitquant_vulkan_pte_runner PRIVATE cxx_std_17) + target_compile_options( + orbitquant_vulkan_pte_runner PRIVATE ${VULKAN_CXX_FLAGS} + ) + target_compile_definitions( + orbitquant_vulkan_pte_runner PRIVATE ${VULKAN_COMPILE_DEFINITIONS} + ) + target_include_directories( + orbitquant_vulkan_pte_runner PRIVATE ${ORBITQUANT_EXECUTORCH_ROOT}/.. + ) + target_link_libraries( + orbitquant_vulkan_pte_runner + PRIVATE + executorch + extension_data_loader + extension_evalue_util + extension_flat_tensor + extension_runner_util + gflags + portable_ops_lib + vulkan_backend + orbitquant_vulkan_shaderlib + orbitquant_vulkan_ops + ) +endif() diff --git a/native-kernels/orbitquant-vulkan/OrbitLinear.cpp b/native-kernels/orbitquant-vulkan/OrbitLinear.cpp new file mode 100644 index 0000000..bf42247 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/OrbitLinear.cpp @@ -0,0 +1,243 @@ +#include + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace vkcompute { +namespace { + +constexpr uint32_t kNormWorkers = 64u; +constexpr uint32_t kRpbhWorkers = 64u; +constexpr uint32_t kMatmulPrefillLocalX = 16u; +constexpr uint32_t kMatmulPrefillLocalY = 4u; +constexpr uint32_t kMatmulPrefillOutputTileM = 8u; +constexpr uint32_t kMatmulPrefillOutputTileN = 2u; + +ValueRef prepack_buffer(ComputeGraph& graph, const ValueRef data) { + return prepack_standard( + graph, data, utils::kBuffer, utils::kWidthPacked, false); +} + +ValueRef prepack_w4_transposed( + ComputeGraph& graph, + const ValueRef data, + const int64_t N, + const int64_t K) { + const int64_t words_per_row = K / 8; + const ValueRef packed = graph.add_tensor( + {N * words_per_row}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + const utils::uvec3 global_workgroup_size = { + utils::safe_downcast(N), + utils::safe_downcast(words_per_row), + 1u}; + const utils::ivec4 problem_sizes = { + 0, + utils::safe_downcast(N), + utils::safe_downcast(K), + 0}; + graph.prepack_nodes().emplace_back(new PrepackNode( + graph, + VK_KERNEL_FROM_STR("orbitquant_pack_w4_transposed"), + global_workgroup_size, + {8u, 8u, 1u}, + data, + packed, + {}, + {}, + {PushConstantDataInfo(&problem_sizes, sizeof(problem_sizes))})); + return packed; +} + +std::string dtype_shader_name( + std::string kernel_name, + const vkapi::ScalarType dtype) { + add_dtype_suffix(kernel_name, dtype); + return kernel_name; +} + +void check_float_constant(ComputeGraph& graph, const ValueRef ref) { + VK_CHECK_COND(graph.dtype_of(ref) == vkapi::kFloat); +} + +int64_t logical_numel(ComputeGraph& graph, const ValueRef ref) { + int64_t numel = 1; + for (const int64_t size : graph.sizes_of(ref)) { + numel *= size; + } + return numel; +} + +} // namespace + +void orbitquant_linear_w4a4_exact( + ComputeGraph& graph, + const std::vector& args) { + int32_t idx = 0; + const ValueRef input = args.at(idx++); + const ValueRef packed_weight_data = args.at(idx++); + const ValueRef row_norms_data = args.at(idx++); + const ValueRef permutation_data = args.at(idx++); + const ValueRef signs_data = args.at(idx++); + const ValueRef activation_boundaries_data = args.at(idx++); + const ValueRef pair_lut_data = args.at(idx++); + const ValueRef block_size_ref = args.at(idx++); + const ValueRef eps_ref = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx); + + const vkapi::ScalarType input_dtype = graph.dtype_of(input); + VK_CHECK_COND( + input_dtype == vkapi::kFloat || input_dtype == vkapi::kHalf, + "OrbitQuant Vulkan W4A4 requires float32 or float16 activations"); + VK_CHECK_COND(graph.dtype_of(output) == input_dtype); + VK_CHECK_COND(graph.dtype_of(packed_weight_data) == vkapi::kByte); + VK_CHECK_COND(graph.dtype_of(permutation_data) == vkapi::kInt); + VK_CHECK_COND(graph.dtype_of(signs_data) == vkapi::kInt); + check_float_constant(graph, row_norms_data); + check_float_constant(graph, activation_boundaries_data); + check_float_constant(graph, pair_lut_data); + + const std::vector input_sizes = graph.sizes_of(input); + VK_CHECK_COND(!input_sizes.empty()); + const int64_t K = input_sizes.back(); + const int64_t M = graph.numel_of(input) / K; + const int64_t N = logical_numel(graph, row_norms_data); + const int64_t block_size = graph.extract_scalar(block_size_ref); + const float eps = graph.extract_scalar(eps_ref); + + VK_CHECK_COND(K > 0 && M > 0 && N > 0); + VK_CHECK_COND(K % 8 == 0, "OrbitQuant Vulkan W4A4 requires K divisible by 8"); + VK_CHECK_COND( + block_size >= 8 && block_size <= 4096 && + (block_size & (block_size - 1)) == 0 && K % block_size == 0, + "OrbitQuant Vulkan W4A4 requires a power-of-two block_size in " + "[8, 4096] dividing K"); + VK_CHECK_COND( + logical_numel(graph, packed_weight_data) == N * K / 2, + "packed W4 weight length does not match N*K/2"); + VK_CHECK_COND(logical_numel(graph, permutation_data) == K); + VK_CHECK_COND(logical_numel(graph, signs_data) == K); + VK_CHECK_COND(logical_numel(graph, activation_boundaries_data) == 15); + VK_CHECK_COND(logical_numel(graph, pair_lut_data) == 256); + VK_CHECK_COND(graph.numel_of(output) == M * N); + + const ValueRef packed_weight = + prepack_w4_transposed(graph, packed_weight_data, N, K); + const ValueRef row_norms = prepack_buffer(graph, row_norms_data); + const ValueRef permutation = prepack_buffer(graph, permutation_data); + const ValueRef signs = prepack_buffer(graph, signs_data); + const ValueRef activation_boundaries = + prepack_buffer(graph, activation_boundaries_data); + const ValueRef pair_lut = prepack_buffer(graph, pair_lut_data); + + TmpTensor token_norms( + &graph, {M}, vkapi::kFloat, utils::kBuffer, utils::kWidthPacked); + TmpTensor packed_activations( + &graph, + {M * K / 8}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + + const utils::ivec4 problem_sizes = { + utils::safe_downcast(M), + utils::safe_downcast(N), + utils::safe_downcast(K), + utils::safe_downcast(block_size)}; + + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + VK_KERNEL_FROM_STR(dtype_shader_name("orbitquant_token_norm", input_dtype)), + {kNormWorkers, utils::safe_downcast(M), 1u}, + {kNormWorkers, 1u, 1u}, + {{token_norms.vref, vkapi::kWrite}, {input, vkapi::kRead}}, + {}, + {PushConstantDataInfo(&problem_sizes, sizeof(problem_sizes))})); + + const float inv_sqrt_block = 1.0f / std::sqrt(static_cast(block_size)); + std::string rpbh_kernel = + "orbitquant_rpbh_w4_b" + std::to_string(block_size); + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + VK_KERNEL_FROM_STR(dtype_shader_name(rpbh_kernel, input_dtype)), + {utils::safe_downcast(K / block_size) * kRpbhWorkers, + utils::safe_downcast(M), + 1u}, + {kRpbhWorkers, 1u, 1u}, + {{packed_activations.vref, vkapi::kWrite}, + {{input, + token_norms.vref, + permutation, + signs, + activation_boundaries}, + vkapi::kRead}}, + {}, + {PushConstantDataInfo(&problem_sizes, sizeof(problem_sizes)), + PushConstantDataInfo(&eps, sizeof(eps)), + PushConstantDataInfo(&inv_sqrt_block, sizeof(inv_sqrt_block))})); + + TmpTensor dummy_bias( + &graph, {}, vkapi::kFloat, utils::kBuffer, utils::kWidthPacked); + ValueRef bias = dummy_bias.vref; + int32_t apply_bias = 0; + if (graph.val_is_not_none(bias_data)) { + check_float_constant(graph, bias_data); + VK_CHECK_COND(logical_numel(graph, bias_data) == N); + bias = prepack_buffer(graph, bias_data); + apply_bias = 1; + } + + std::string matmul_kernel = "orbitquant_w4a4_lut"; + utils::uvec3 matmul_global = { + utils::safe_downcast(N), + utils::safe_downcast(M), + 1u}; + utils::uvec3 matmul_local = {64u, 1u, 1u}; + if (M >= 8) { + matmul_kernel = "orbitquant_w4a4_lut_tiled_t8"; + matmul_global = { + utils::safe_downcast( + (N + kMatmulPrefillOutputTileN - 1u) / + kMatmulPrefillOutputTileN), + utils::safe_downcast( + (M + kMatmulPrefillOutputTileM - 1u) / + kMatmulPrefillOutputTileM), + 1u}; + matmul_local = {kMatmulPrefillLocalX, kMatmulPrefillLocalY, 1u}; + } + + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + VK_KERNEL_FROM_STR(dtype_shader_name(matmul_kernel, input_dtype)), + matmul_global, + matmul_local, + {{output, vkapi::kWrite}, + {{packed_activations.vref, + packed_weight, + token_norms.vref, + row_norms, + pair_lut, + bias}, + vkapi::kRead}}, + {}, + {PushConstantDataInfo(&problem_sizes, sizeof(problem_sizes)), + PushConstantDataInfo(&apply_bias, sizeof(apply_bias))})); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP( + orbitquant_vulkan.linear_w4a4_exact.default, + orbitquant_linear_w4a4_exact); +} + +} // namespace vkcompute diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_pack_w4_transposed.glsl b/native-kernels/orbitquant-vulkan/glsl/orbitquant_pack_w4_transposed.glsl new file mode 100644 index 0000000..b6a2724 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_pack_w4_transposed.glsl @@ -0,0 +1,28 @@ +#version 450 core + +${define_required_extensions("buffer", "int")} +${define_active_storage_type("buffer")} + +#define PRECISION highp + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_output", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_input", "uint", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 problem_sizes; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const uint n = gl_GlobalInvocationID.x; + const uint word = gl_GlobalInvocationID.y; + const uint N = uint(problem_sizes.y); + const uint words_per_row = uint(problem_sizes.z) >> 3u; + if (n >= N || word >= words_per_row) { + return; + } + t_output[word * N + n] = int(t_input[n * words_per_row + word]); +} diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_pack_w4_transposed.yaml b/native-kernels/orbitquant-vulkan/glsl/orbitquant_pack_w4_transposed.yaml new file mode 100644 index 0000000..4699fb1 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_pack_w4_transposed.yaml @@ -0,0 +1,4 @@ +orbitquant_pack_w4_transposed: + parameter_names_with_default_values: {} + shader_variants: + - NAME: orbitquant_pack_w4_transposed diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_w4.glsl b/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_w4.glsl new file mode 100644 index 0000000..0fbed45 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_w4.glsl @@ -0,0 +1,83 @@ +#version 450 core + +${define_required_extensions("buffer", DTYPE)} + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} +#define BLOCK_SIZE ${BLOCK_SIZE} + +${define_active_storage_type("buffer")} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_packed", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_input", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_norms", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_permutation", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_signs", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_boundaries", "float", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 problem_sizes; + float eps; + float inv_sqrt_block; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +shared float values[BLOCK_SIZE]; + +void main() { + const uint block = gl_WorkGroupID.x; + const uint row = gl_WorkGroupID.y; + const uint lane = gl_LocalInvocationID.x; + const uint rows = uint(problem_sizes.x); + const uint K = uint(problem_sizes.z); + const uint block_size = uint(problem_sizes.w); + if (row >= rows || block_size != uint(BLOCK_SIZE)) { + return; + } + + const uint block_offset = block * block_size; + const float denominator = t_norms[row] + eps; + for (uint i = lane; i < block_size; i += 64u) { + const uint rotated_index = block_offset + i; + const uint source_index = uint(t_permutation[rotated_index]); + const float sign = float(t_signs[rotated_index]); + values[i] = float(t_input[row * K + source_index]) * sign / denominator; + } + memoryBarrierShared(); + barrier(); + + for (uint span = 1u; span < block_size; span <<= 1u) { + const uint butterflies = block_size >> 1u; + for (uint i = lane; i < butterflies; i += 64u) { + const uint group = i / span; + const uint offset = i - group * span; + const uint left = group * (span << 1u) + offset; + const uint right = left + span; + const float a = values[left]; + const float b = values[right]; + values[left] = a + b; + values[right] = a - b; + } + memoryBarrierShared(); + barrier(); + } + + const uint words_per_block = block_size >> 3u; + const uint words_per_row = K >> 3u; + for (uint word = lane; word < words_per_block; word += 64u) { + uint packed = 0u; + for (uint nibble = 0u; nibble < 8u; ++nibble) { + const float value = values[word * 8u + nibble] * inv_sqrt_block; + uint code = 0u; + for (uint boundary = 0u; boundary < 15u; ++boundary) { + code += value > t_boundaries[boundary] ? 1u : 0u; + } + packed |= code << (nibble * 4u); + } + const uint output_word = block * words_per_block + word; + t_packed[output_word * rows + row] = int(packed); + } +} diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_w4.yaml b/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_w4.yaml new file mode 100644 index 0000000..db151ff --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_rpbh_w4.yaml @@ -0,0 +1,29 @@ +orbitquant_rpbh_w4: + parameter_names_with_default_values: + BLOCK_SIZE: 8 + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: orbitquant_rpbh_w4_b8 + BLOCK_SIZE: 8 + - NAME: orbitquant_rpbh_w4_b16 + BLOCK_SIZE: 16 + - NAME: orbitquant_rpbh_w4_b32 + BLOCK_SIZE: 32 + - NAME: orbitquant_rpbh_w4_b64 + BLOCK_SIZE: 64 + - NAME: orbitquant_rpbh_w4_b128 + BLOCK_SIZE: 128 + - NAME: orbitquant_rpbh_w4_b256 + BLOCK_SIZE: 256 + - NAME: orbitquant_rpbh_w4_b512 + BLOCK_SIZE: 512 + - NAME: orbitquant_rpbh_w4_b1024 + BLOCK_SIZE: 1024 + - NAME: orbitquant_rpbh_w4_b2048 + BLOCK_SIZE: 2048 + - NAME: orbitquant_rpbh_w4_b4096 + BLOCK_SIZE: 4096 diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_token_norm.glsl b/native-kernels/orbitquant-vulkan/glsl/orbitquant_token_norm.glsl new file mode 100644 index 0000000..997df26 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_token_norm.glsl @@ -0,0 +1,52 @@ +#version 450 core + +${define_required_extensions("buffer", DTYPE)} + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type("buffer")} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_norms", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_input", DTYPE, "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 problem_sizes; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +shared float sums[64]; + +void main() { + const uint row = gl_GlobalInvocationID.y; + const uint lane = gl_LocalInvocationID.x; + const uint rows = uint(problem_sizes.x); + const uint K = uint(problem_sizes.z); + + float sum = 0.0; + if (row < rows) { + const uint row_offset = row * K; + for (uint k = lane; k < K; k += 64u) { + const float value = float(t_input[row_offset + k]); + sum += value * value; + } + } + sums[lane] = sum; + memoryBarrierShared(); + barrier(); + + for (uint stride = 32u; stride > 0u; stride >>= 1u) { + if (lane < stride) { + sums[lane] += sums[lane + stride]; + } + memoryBarrierShared(); + barrier(); + } + + if (lane == 0u && row < rows) { + t_norms[row] = sqrt(sums[0]); + } +} diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_token_norm.yaml b/native-kernels/orbitquant-vulkan/glsl/orbitquant_token_norm.yaml new file mode 100644 index 0000000..9946d48 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_token_norm.yaml @@ -0,0 +1,9 @@ +orbitquant_token_norm: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: orbitquant_token_norm diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut.glsl b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut.glsl new file mode 100644 index 0000000..0a48802 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut.glsl @@ -0,0 +1,94 @@ +#version 450 core + +${define_required_extensions("buffer", DTYPE)} +#extension GL_EXT_control_flow_attributes : require + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type("buffer")} + +#define TILE_M 1 +#define TILE_N 1 + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_output", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_activations", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_weight", "uint", "buffer")} +${layout_declare_tensor(B, "r", "t_token_norms", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_row_norms", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_pair_lut", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_bias", "float", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 problem_sizes; + int apply_bias; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const uint n_base = gl_GlobalInvocationID.x * TILE_N; + const uint m_base = gl_GlobalInvocationID.y * TILE_M; + const uint M = uint(problem_sizes.x); + const uint N = uint(problem_sizes.y); + const uint K = uint(problem_sizes.z); + if (m_base >= M || n_base >= N) { + return; + } + + const uint words_per_row = K >> 3u; + float accumulators[TILE_M][TILE_N]; + [[unroll]] for (uint mi = 0u; mi < TILE_M; ++mi) { + [[unroll]] for (uint ni = 0u; ni < TILE_N; ++ni) { + accumulators[mi][ni] = 0.0; + } + } + + for (uint word = 0u; word < words_per_row; ++word) { + uint activation_codes[TILE_M]; + uint weight_codes[TILE_N]; + [[unroll]] for (uint mi = 0u; mi < TILE_M; ++mi) { + const uint m = min(m_base + mi, M - 1u); + activation_codes[mi] = + uint(t_activations[word * M + m]); + } + [[unroll]] for (uint ni = 0u; ni < TILE_N; ++ni) { + const uint n = min(n_base + ni, N - 1u); + weight_codes[ni] = t_weight[word * N + n]; + } + + for (uint nibble = 0u; nibble < 8u; ++nibble) { + const uint shift = nibble * 4u; + [[unroll]] for (uint mi = 0u; mi < TILE_M; ++mi) { + const uint activation_code = + (activation_codes[mi] >> shift) & 15u; + [[unroll]] for (uint ni = 0u; ni < TILE_N; ++ni) { + const uint weight_code = (weight_codes[ni] >> shift) & 15u; + accumulators[mi][ni] += + t_pair_lut[activation_code * 16u + weight_code]; + } + } + } + } + + [[unroll]] for (uint mi = 0u; mi < TILE_M; ++mi) { + const uint m = m_base + mi; + if (m >= M) { + continue; + } + [[unroll]] for (uint ni = 0u; ni < TILE_N; ++ni) { + const uint n = n_base + ni; + if (n >= N) { + continue; + } + float value = + accumulators[mi][ni] * t_token_norms[m] * t_row_norms[n]; + if (apply_bias != 0) { + value += t_bias[n]; + } + t_output[m * N + n] = T(value); + } + } +} diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut.yaml b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut.yaml new file mode 100644 index 0000000..e90716a --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut.yaml @@ -0,0 +1,9 @@ +orbitquant_w4a4_lut: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: orbitquant_w4a4_lut diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut_tiled.glsl b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut_tiled.glsl new file mode 100644 index 0000000..2653ffe --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut_tiled.glsl @@ -0,0 +1,140 @@ +#version 450 core + +${define_required_extensions("buffer", DTYPE)} +#extension GL_EXT_control_flow_attributes : require + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} +#define K_TILE_WORDS ${K_TILE_WORDS} +#define USE_SHARED_LUT ${USE_SHARED_LUT} +#define TILE_M 8u +#define TILE_N 2u +#define WORKGROUP_M 32u +#define WORKGROUP_N 32u +#define WORKGROUP_SIZE 64u + +${define_active_storage_type("buffer")} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_output", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_activations", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_weight", "uint", "buffer")} +${layout_declare_tensor(B, "r", "t_token_norms", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_row_norms", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_pair_lut", "float", "buffer")} +${layout_declare_tensor(B, "r", "t_bias", "float", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 problem_sizes; + int apply_bias; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +shared uint activation_tile[K_TILE_WORDS][WORKGROUP_M]; +shared uint weight_tile[K_TILE_WORDS][WORKGROUP_N]; +#if USE_SHARED_LUT +shared float pair_lut_tile[256]; +#endif + +void main() { + const uint M = uint(problem_sizes.x); + const uint N = uint(problem_sizes.y); + const uint K = uint(problem_sizes.z); + const uint words_per_row = K >> 3u; + const uint lane = gl_LocalInvocationID.y * 16u + gl_LocalInvocationID.x; + const uint group_m_base = gl_WorkGroupID.y * WORKGROUP_M; + const uint group_n_base = gl_WorkGroupID.x * WORKGROUP_N; + const uint m_base = group_m_base + gl_LocalInvocationID.y * TILE_M; + const uint n_base = group_n_base + gl_LocalInvocationID.x * TILE_N; + +#if USE_SHARED_LUT + for (uint index = lane; index < 256u; index += WORKGROUP_SIZE) { + pair_lut_tile[index] = t_pair_lut[index]; + } + memoryBarrierShared(); + barrier(); +#endif + + float accumulators[TILE_M][TILE_N]; + [[unroll]] for (uint mi = 0u; mi < TILE_M; ++mi) { + [[unroll]] for (uint ni = 0u; ni < TILE_N; ++ni) { + accumulators[mi][ni] = 0.0; + } + } + + for (uint word_base = 0u; word_base < words_per_row; + word_base += K_TILE_WORDS) { + for (uint index = lane; index < K_TILE_WORDS * WORKGROUP_M; + index += WORKGROUP_SIZE) { + const uint tile_word = index / WORKGROUP_M; + const uint row = index - tile_word * WORKGROUP_M; + const uint word = word_base + tile_word; + const uint m = group_m_base + row; + activation_tile[tile_word][row] = + word < words_per_row && m < M + ? uint(t_activations[word * M + m]) + : 0u; + const uint n = group_n_base + row; + weight_tile[tile_word][row] = + word < words_per_row && n < N ? t_weight[word * N + n] : 0u; + } + memoryBarrierShared(); + barrier(); + + [[unroll]] for (uint tile_word = 0u; tile_word < K_TILE_WORDS; + ++tile_word) { + if (word_base + tile_word >= words_per_row) { + continue; + } + uint activation_codes[TILE_M]; + uint weight_codes[TILE_N]; + [[unroll]] for (uint mi = 0u; mi < TILE_M; ++mi) { + activation_codes[mi] = activation_tile[tile_word] + [gl_LocalInvocationID.y * TILE_M + mi]; + } + [[unroll]] for (uint ni = 0u; ni < TILE_N; ++ni) { + weight_codes[ni] = + weight_tile[tile_word][gl_LocalInvocationID.x * TILE_N + ni]; + } + [[unroll]] for (uint nibble = 0u; nibble < 8u; ++nibble) { + const uint shift = nibble * 4u; + [[unroll]] for (uint mi = 0u; mi < TILE_M; ++mi) { + const uint activation_code = + (activation_codes[mi] >> shift) & 15u; + [[unroll]] for (uint ni = 0u; ni < TILE_N; ++ni) { + const uint weight_code = (weight_codes[ni] >> shift) & 15u; + const uint lut_index = activation_code * 16u + weight_code; +#if USE_SHARED_LUT + accumulators[mi][ni] += pair_lut_tile[lut_index]; +#else + accumulators[mi][ni] += t_pair_lut[lut_index]; +#endif + } + } + } + } + memoryBarrierShared(); + barrier(); + } + + [[unroll]] for (uint mi = 0u; mi < TILE_M; ++mi) { + const uint m = m_base + mi; + if (m >= M) { + continue; + } + [[unroll]] for (uint ni = 0u; ni < TILE_N; ++ni) { + const uint n = n_base + ni; + if (n >= N) { + continue; + } + float value = + accumulators[mi][ni] * t_token_norms[m] * t_row_norms[n]; + if (apply_bias != 0) { + value += t_bias[n]; + } + t_output[m * N + n] = T(value); + } + } +} diff --git a/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut_tiled.yaml b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut_tiled.yaml new file mode 100644 index 0000000..dd41443 --- /dev/null +++ b/native-kernels/orbitquant-vulkan/glsl/orbitquant_w4a4_lut_tiled.yaml @@ -0,0 +1,13 @@ +orbitquant_w4a4_lut_tiled: + parameter_names_with_default_values: + DTYPE: float + K_TILE_WORDS: 8 + USE_SHARED_LUT: 1 + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: orbitquant_w4a4_lut_tiled_t8 + K_TILE_WORDS: 8 + USE_SHARED_LUT: 1 diff --git a/native-kernels/orbitquant-vulkan/tests/OrbitLinearTest.cpp b/native-kernels/orbitquant-vulkan/tests/OrbitLinearTest.cpp new file mode 100644 index 0000000..618c79e --- /dev/null +++ b/native-kernels/orbitquant-vulkan/tests/OrbitLinearTest.cpp @@ -0,0 +1,425 @@ +#include "utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace executorch::vulkan::prototyping; + +namespace { + +struct TestShape { + int32_t rows; + int32_t outputs; + int32_t input_dim; + int32_t block_size; + const char* name; +}; + +constexpr TestShape kSmallShape = {3, 7, 24, 8, "small"}; +constexpr TestShape kDit1536Shape = {2, 9, 1536, 512, "dit1536"}; +constexpr TestShape kTiledPartialShape = {9, 17, 24, 8, "tiled_partial"}; +constexpr TestShape kDit1536TiledShape = {8, 9, 1536, 512, "dit1536_tiled"}; +constexpr TestShape kDit1536DecodeShape = { + 1, 1536, 1536, 512, "dit1536_decode"}; +constexpr TestShape kDit1536PrefillShape = { + 32, 1536, 1536, 512, "dit1536_prefill"}; +constexpr TestShape kDit3072PrefillShape = { + 32, 3072, 3072, 1024, "dit3072_prefill"}; + +bool profile_enabled() { + const char* value = std::getenv("ORBITQUANT_VULKAN_PROFILE"); + return value != nullptr && std::string(value) == "1"; +} + +std::string profile_case_filter() { + const char* value = std::getenv("ORBITQUANT_VULKAN_PROFILE_CASE"); + return value == nullptr ? std::string() : std::string(value); +} + +float percentile(std::vector values, const float quantile) { + if (values.empty()) { + return 0.0f; + } + std::sort(values.begin(), values.end()); + const size_t index = static_cast( + std::ceil(quantile * static_cast(values.size())) - 1.0f); + return values.at(std::min(index, values.size() - 1)); +} + +void print_profile_percentiles(const TestResult& result) { + for (const BenchmarkResult& benchmark : result.get_results()) { + std::cout << "ORBITQUANT_PROFILE case=" << benchmark.get_kernel_name() + << " median_us=" + << percentile(benchmark.get_iter_timings(), 0.5f) + << " p95_us=" + << percentile(benchmark.get_iter_timings(), 0.95f) + << std::endl; + for (const ShaderTiming& shader : benchmark.get_shader_timings()) { + std::cout << "ORBITQUANT_PROFILE case=" << benchmark.get_kernel_name() + << " shader=" << shader.shader_name << " median_us=" + << percentile(shader.iter_timings_us, 0.5f) << " p95_us=" + << percentile(shader.iter_timings_us, 0.95f) << std::endl; + } + } +} + +ValueSpec tensor_spec( + const std::vector& sizes, + const vkapi::ScalarType dtype, + const bool constant) { + ValueSpec spec( + sizes, + dtype, + utils::kBuffer, + utils::kWidthPacked, + DataGenType::ZEROS); + if (constant) { + spec.set_constant(true); + } else { + spec.ensure_data_generated(); + } + return spec; +} + +TestCase make_test_case( + const TestShape shape, + const vkapi::ScalarType input_dtype, + const bool has_bias) { + const std::string dtype_name = + input_dtype == vkapi::kHalf ? "f16" : "f32"; + TestCase test_case( + std::string(shape.name) + "_" + dtype_name + + (has_bias ? "_bias" : "_no_bias")); + test_case.set_operator_name("orbitquant_vulkan.linear_w4a4_exact.default"); + const float tolerance = input_dtype == vkapi::kHalf ? 8.0e-3f : 2.0e-4f; + test_case.set_abs_tolerance(tolerance); + test_case.set_rel_tolerance(tolerance); + test_case.set_op_invocations_per_execute(1); + + ValueSpec input = + tensor_spec({shape.rows, shape.input_dim}, input_dtype, false); + const int32_t first_nonzero_row = shape.rows == 1 ? 0 : 1; + for (int32_t row = first_nonzero_row; row < shape.rows; ++row) { + for (int32_t k = 0; k < shape.input_dim; ++k) { + const float value = std::sin(float(row * 19 + k * 7) * 0.13f); + const int32_t index = row * shape.input_dim + k; + if (input_dtype == vkapi::kHalf) { + input.get_half_data()[index] = float_to_half(value); + } else { + input.get_float_data()[index] = value; + } + } + } + + ValueSpec packed_weight = + tensor_spec({shape.outputs, shape.input_dim / 2}, vkapi::kByte, true); + for (int32_t n = 0; n < shape.outputs; ++n) { + for (int32_t k = 0; k < shape.input_dim; k += 2) { + const uint8_t low = static_cast((n * 3 + k * 5) & 15); + const uint8_t high = + static_cast((n * 3 + (k + 1) * 5) & 15); + packed_weight.get_uint8_data()[n * (shape.input_dim / 2) + k / 2] = + static_cast(low | (high << 4)); + } + } + + ValueSpec row_norms = tensor_spec({shape.outputs}, vkapi::kFloat, true); + for (int32_t n = 0; n < shape.outputs; ++n) { + row_norms.get_float_data()[n] = 0.55f + float(n % 17) * 0.007f; + } + + ValueSpec permutation = tensor_spec({shape.input_dim}, vkapi::kInt, true); + ValueSpec signs = tensor_spec({shape.input_dim}, vkapi::kInt, true); + for (int32_t k = 0; k < shape.input_dim; ++k) { + permutation.get_int32_data()[k] = (k * 5) % shape.input_dim; + signs.get_int32_data()[k] = (k % 3 == 0) ? -1 : 1; + } + + std::vector activation_centroids(16); + std::vector weight_centroids(16); + for (int32_t index = 0; index < 16; ++index) { + activation_centroids[index] = -0.75f + float(index) * 0.1f; + weight_centroids[index] = -0.9f + float(index) * 0.12f; + } + + ValueSpec boundaries = tensor_spec({15}, vkapi::kFloat, true); + for (int32_t index = 0; index < 15; ++index) { + boundaries.get_float_data()[index] = + 0.5f * (activation_centroids[index] + activation_centroids[index + 1]); + } + + ValueSpec pair_lut = tensor_spec({256}, vkapi::kFloat, true); + for (int32_t activation = 0; activation < 16; ++activation) { + for (int32_t weight = 0; weight < 16; ++weight) { + pair_lut.get_float_data()[activation * 16 + weight] = + activation_centroids[activation] * weight_centroids[weight]; + } + } + + ValueSpec bias = tensor_spec({shape.outputs}, vkapi::kFloat, true); + if (has_bias) { + for (int32_t n = 0; n < shape.outputs; ++n) { + bias.get_float_data()[n] = float(n - 3) * 0.025f; + } + } else { + bias.set_none(true); + } + + test_case.add_input_spec(input); + test_case.add_input_spec(packed_weight); + test_case.add_input_spec(row_norms); + test_case.add_input_spec(permutation); + test_case.add_input_spec(signs); + test_case.add_input_spec(boundaries); + test_case.add_input_spec(pair_lut); + test_case.add_input_spec(ValueSpec(shape.block_size)); + test_case.add_input_spec(ValueSpec(1.0e-10f)); + test_case.add_input_spec(bias); + test_case.add_output_spec( + tensor_spec({shape.rows, shape.outputs}, input_dtype, false)); + return test_case; +} + +TestCase make_fp16_linear_test_case(const TestShape shape) { + TestCase test_case(std::string("fp16_linear_") + shape.name); + test_case.set_operator_name("aten.linear.default"); + test_case.set_abs_tolerance(2.0e-2f); + test_case.set_rel_tolerance(2.0e-2f); + test_case.set_op_invocations_per_execute(1); + + ValueSpec input = + tensor_spec({shape.rows, shape.input_dim}, vkapi::kHalf, false); + for (int32_t row = 0; row < shape.rows; ++row) { + for (int32_t k = 0; k < shape.input_dim; ++k) { + const float value = std::sin(float(row * 19 + k * 7) * 0.013f); + input.get_half_data()[row * shape.input_dim + k] = + float_to_half(value); + } + } + + ValueSpec weight = tensor_spec( + {shape.outputs, shape.input_dim}, vkapi::kHalf, true); + for (int32_t n = 0; n < shape.outputs; ++n) { + for (int32_t k = 0; k < shape.input_dim; ++k) { + const float value = + std::sin(float(n * 13 + k * 11) * 0.017f) * 0.01f; + weight.get_half_data()[n * shape.input_dim + k] = + float_to_half(value); + } + } + + ValueSpec bias = tensor_spec({shape.outputs}, vkapi::kHalf, true); + for (int32_t n = 0; n < shape.outputs; ++n) { + bias.get_half_data()[n] = float_to_half(float(n % 17 - 8) * 0.001f); + } + + test_case.add_input_spec(input); + test_case.add_input_spec(weight); + test_case.add_input_spec(bias); + test_case.add_output_spec( + tensor_spec({shape.rows, shape.outputs}, vkapi::kHalf, false)); + return test_case; +} + +std::vector generate_test_cases() { + std::vector test_cases = { + make_test_case(kSmallShape, vkapi::kFloat, false), + make_test_case(kSmallShape, vkapi::kFloat, true), + make_test_case(kSmallShape, vkapi::kHalf, false), + make_test_case(kSmallShape, vkapi::kHalf, true), + make_test_case(kDit1536Shape, vkapi::kFloat, false), + make_test_case(kTiledPartialShape, vkapi::kHalf, true), + make_test_case(kDit1536TiledShape, vkapi::kHalf, true), + }; + if (profile_enabled()) { + const std::string filter = profile_case_filter(); + if (!filter.empty()) { + test_cases.clear(); + if (filter == "packed_decode1536") { + test_cases.emplace_back( + make_test_case(kDit1536DecodeShape, vkapi::kHalf, true)); + } else if (filter == "packed_prefill1536") { + test_cases.emplace_back( + make_test_case(kDit1536PrefillShape, vkapi::kHalf, true)); + } else if (filter == "packed_prefill3072") { + test_cases.emplace_back( + make_test_case(kDit3072PrefillShape, vkapi::kHalf, true)); + } else if (filter == "fp16_decode1536") { + test_cases.emplace_back( + make_fp16_linear_test_case(kDit1536DecodeShape)); + } else if (filter == "fp16_prefill1536") { + test_cases.emplace_back( + make_fp16_linear_test_case(kDit1536PrefillShape)); + } else if (filter == "fp16_prefill3072") { + test_cases.emplace_back( + make_fp16_linear_test_case(kDit3072PrefillShape)); + } else { + throw std::invalid_argument( + "unknown ORBITQUANT_VULKAN_PROFILE_CASE=" + filter); + } + return test_cases; + } + test_cases.emplace_back( + make_test_case(kDit1536DecodeShape, vkapi::kHalf, true)); + test_cases.emplace_back( + make_test_case(kDit1536PrefillShape, vkapi::kHalf, true)); + test_cases.emplace_back( + make_test_case(kDit3072PrefillShape, vkapi::kHalf, true)); + test_cases.emplace_back(make_fp16_linear_test_case(kDit1536DecodeShape)); + test_cases.emplace_back(make_fp16_linear_test_case(kDit1536PrefillShape)); + test_cases.emplace_back(make_fp16_linear_test_case(kDit3072PrefillShape)); + } + return test_cases; +} + +void fwht(std::vector& values, const int32_t block_size) { + for (int32_t span = 1; span < block_size; span *= 2) { + for (int32_t base = 0; base < block_size; base += span * 2) { + for (int32_t offset = 0; offset < span; ++offset) { + const float a = values[base + offset]; + const float b = values[base + span + offset]; + values[base + offset] = a + b; + values[base + span + offset] = a - b; + } + } + } +} + +void reference_compute(TestCase& test_case) { + if (test_case.operator_name() == "aten.linear.default") { + const ValueSpec& input = test_case.inputs().at(0); + const ValueSpec& weight = test_case.inputs().at(1); + const ValueSpec& bias = test_case.inputs().at(2); + const int32_t M = + static_cast(input.get_tensor_sizes().at(0)); + const int32_t K = + static_cast(input.get_tensor_sizes().at(1)); + const int32_t N = + static_cast(weight.get_tensor_sizes().at(0)); + auto& output = test_case.outputs().at(0).get_ref_float_data(); + output.assign(M * N, 0.0f); + for (int32_t m = 0; m < M; ++m) { + for (int32_t n = 0; n < N; ++n) { + float accumulator = half_to_float(bias.get_half_data().at(n)); + for (int32_t k = 0; k < K; ++k) { + accumulator += + half_to_float(input.get_half_data().at(m * K + k)) * + half_to_float(weight.get_half_data().at(n * K + k)); + } + output[m * N + n] = accumulator; + } + } + return; + } + + const ValueSpec& input_spec = test_case.inputs().at(0); + const auto& packed_weight = test_case.inputs().at(1).get_uint8_data(); + const auto& row_norms = test_case.inputs().at(2).get_float_data(); + const auto& permutation = test_case.inputs().at(3).get_int32_data(); + const auto& signs = test_case.inputs().at(4).get_int32_data(); + const auto& boundaries = test_case.inputs().at(5).get_float_data(); + const auto& pair_lut = test_case.inputs().at(6).get_float_data(); + const ValueSpec& bias_spec = test_case.inputs().at(9); + const auto& bias = bias_spec.get_float_data(); + const int32_t block_size = test_case.inputs().at(7).get_int_value(); + const int32_t M = static_cast(input_spec.get_tensor_sizes().at(0)); + const int32_t K = static_cast(input_spec.get_tensor_sizes().at(1)); + const int32_t N = static_cast(row_norms.size()); + auto& output = test_case.outputs().at(0).get_ref_float_data(); + output.assign(M * N, 0.0f); + + const auto input_value = [&input_spec](const int32_t index) { + if (input_spec.dtype == vkapi::kHalf) { + return half_to_float(input_spec.get_half_data().at(index)); + } + return input_spec.get_float_data().at(index); + }; + + std::vector activation_codes(M * K); + std::vector token_norms(M); + for (int32_t row = 0; row < M; ++row) { + float squared_norm = 0.0f; + for (int32_t k = 0; k < K; ++k) { + const float value = input_value(row * K + k); + squared_norm += value * value; + } + token_norms[row] = std::sqrt(squared_norm); + for (int32_t block = 0; block < K / block_size; ++block) { + std::vector values(block_size); + for (int32_t k = 0; k < block_size; ++k) { + const int32_t rotated_index = block * block_size + k; + values[k] = + input_value(row * K + permutation[rotated_index]) * + signs[rotated_index] / + (token_norms[row] + 1.0e-10f); + } + fwht(values, block_size); + for (int32_t k = 0; k < block_size; ++k) { + const float value = values[k] / std::sqrt(float(block_size)); + uint8_t code = 0; + for (const float boundary : boundaries) { + code += value > boundary ? 1 : 0; + } + activation_codes[row * K + block * block_size + k] = code; + } + } + } + + for (int32_t row = 0; row < M; ++row) { + for (int32_t n = 0; n < N; ++n) { + float accumulator = 0.0f; + for (int32_t k = 0; k < K; ++k) { + const uint8_t packed = packed_weight[n * (K / 2) + k / 2]; + const uint8_t weight_code = + (k & 1) == 0 ? packed & 15 : (packed >> 4) & 15; + accumulator += pair_lut[ + activation_codes[row * K + k] * 16 + weight_code]; + } + float value = accumulator * token_norms[row] * row_norms[n]; + if (!bias_spec.is_none()) { + value += bias[n]; + } + output[row * N + n] = value; + } + } +} + +int64_t flop_count(const TestCase& test_case) { + const auto& input_sizes = test_case.inputs().at(0).get_tensor_sizes(); + const int64_t outputs = test_case.inputs().at(2).numel(); + return int64_t{2} * input_sizes.at(0) * outputs * input_sizes.at(1); +} + +} // namespace + +int main() { + const bool profile = profile_enabled(); + set_print_output(false); + set_print_latencies(profile); + set_use_gpu_timestamps(true); + try { + api::context()->initialize_querypool(); + } catch (const std::exception& error) { + std::cerr << "Failed to initialize Vulkan: " << error.what() << std::endl; + return 1; + } + + const TestResult result = execute_test_cases( + generate_test_cases, + flop_count, + "OrbitQuant exact W4A4", + profile ? 10 : 1, + profile ? 31 : 3, + reference_compute); + if (profile) { + result.print_detailed_results(); + print_profile_percentiles(result); + } + return result.get_correctness_passed() ? 0 : 1; +} diff --git a/pyproject.toml b/pyproject.toml index a098001..8294018 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.4.0" +version = "0.5.0" description = "Calibration-free OrbitQuant for transformer linear projections" readme = "README.md" requires-python = ">=3.11" diff --git a/scripts/run_mps_kernel_checks.sh b/scripts/run_mps_kernel_checks.sh index 53b69d1..0cf9c84 100755 --- a/scripts/run_mps_kernel_checks.sh +++ b/scripts/run_mps_kernel_checks.sh @@ -161,7 +161,7 @@ orbitquant kernel-bench \ --activation-bits 4 \ --block-size "$BLOCK_SIZE" \ --activation-kernel-backend mps \ - "${BENCH_RUNTIME_ARGS[@]}" \ + ${BENCH_RUNTIME_ARGS[@]+"${BENCH_RUNTIME_ARGS[@]}"} \ --device mps \ --dtype bfloat16 \ --warmup "$WARMUP" \ diff --git a/src/orbitquant/__init__.py b/src/orbitquant/__init__.py index e14e412..f2057ff 100644 --- a/src/orbitquant/__init__.py +++ b/src/orbitquant/__init__.py @@ -18,7 +18,7 @@ from orbitquant.quantizer import register_hf_quantizers from orbitquant.recipes import recipe -__version__ = "0.4.0" +__version__ = "0.5.0" register_hf_quantizers() diff --git a/src/orbitquant/adaln.py b/src/orbitquant/adaln.py index 1418c85..ebb6ae8 100644 --- a/src/orbitquant/adaln.py +++ b/src/orbitquant/adaln.py @@ -40,7 +40,16 @@ def _quantize_adaln_weight( *, group_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: - if weight.is_cuda: + if weight.device.type == "xpu": + try: + from orbitquant.kernels.triton_cuda import quantize_adaln_weight_with_triton + except Exception as exc: + raise RuntimeError( + "XPU AdaLN quantization requires the Intel XPU Triton backend; " + "use CPU quantization for the reference path" + ) from exc + return quantize_adaln_weight_with_triton(weight, group_size=group_size) + if weight.device.type == "cuda": try: from orbitquant.kernels.triton_cuda import quantize_adaln_weight_with_triton except Exception: @@ -134,7 +143,23 @@ def __init__( self.bias = nn.Parameter(bias.detach().clone(), requires_grad=False) self._dequantized_weight_cache: torch.Tensor | None = None self._dequantized_weight_cache_key: tuple[str, torch.dtype] | None = None + self._bias_cache: tuple[tuple[str, torch.dtype], torch.Tensor] | None = None self.last_effective_runtime_mode: str | None = None + from orbitquant.layers import _compile_registry_register + + self._compile_handle = _compile_registry_register(self) + + def _bias_for(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None: + if self.bias is None: + return None + if self.bias.device == device and self.bias.dtype == dtype: + return self.bias + key = (str(device), dtype) + if self._bias_cache is not None and self._bias_cache[0] == key: + return self._bias_cache[1] + value = self.bias.to(device=device, dtype=dtype) + self._bias_cache = (key, value) + return value def clear_dequantized_cache(self) -> None: self._dequantized_weight_cache = None @@ -255,7 +280,7 @@ def _dequantize_weight(self, *, device: torch.device, dtype: torch.dtype) -> tor ): return self._dequantized_weight_cache - if device.type == "cuda": + if device.type in {"cuda", "xpu"}: try: from orbitquant.kernels.triton_cuda import dequantize_adaln_weight_with_triton except Exception: @@ -285,7 +310,79 @@ def _dequantize_weight(self, *, device: torch.device, dtype: torch.dtype) -> tor return self._remember_dequantized_weight(dequantized, cache_key) def forward(self, x: torch.Tensor) -> torch.Tensor: + if torch.compiler.is_compiling(): + return torch.ops.orbitquant.packed_linear_forward( + x, self._compile_handle, self.out_features, True + ) + return self._forward_impl(x) + + def _forward_impl(self, x: torch.Tensor) -> torch.Tensor: activation = x.to(torch.bfloat16) + uses_hip = bool(getattr(torch.version, "hip", None)) + if x.device.type == "xpu": + if self.runtime_mode == "auto_fused": + raise RuntimeError( + "auto_fused XPU AdaLN dispatch remains experimental until real " + "Intel GPU correctness, memory, profiler, and performance proof is " + "recorded. For explicit validation, set " + "runtime_mode='triton_packed_matmul'. Set " + "runtime_mode='dequant_bf16' for the reference/debug path." + ) + if self.runtime_mode == "native_packed_matmul": + raise RuntimeError( + "native_packed_matmul AdaLN has no XPU kernel. Use " + "runtime_mode='triton_packed_matmul' for explicit XPU validation " + "or runtime_mode='dequant_bf16' for the reference/debug path." + ) + if x.device.type == "cuda" and uses_hip: + if self.runtime_mode == "auto_fused": + raise RuntimeError( + "auto_fused ROCm AdaLN dispatch remains experimental until real " + "supported AMD hardware correctness, memory, and performance proof " + "is recorded. For explicit validation, set " + "runtime_mode='triton_packed_matmul'. Set " + "runtime_mode='dequant_bf16' for the reference/debug path." + ) + if self.runtime_mode == "native_packed_matmul": + raise RuntimeError( + "native_packed_matmul AdaLN currently provides CUDA, not HIP, " + "accelerator kernels. Use runtime_mode='triton_packed_matmul' for " + "explicit ROCm validation or runtime_mode='dequant_bf16' for the " + "reference/debug path." + ) + if x.device.type in {"cuda", "xpu"} and self.runtime_mode in { + "auto_fused", + "native_packed_matmul", + "triton_packed_matmul", + }: + try: + from orbitquant.kernels.triton_cuda import ( + matmul_packed_adaln_int4_with_triton, + ) + except Exception as exc: + backend = ( + "XPU" + if x.device.type == "xpu" + else "ROCm" + if getattr(torch.version, "hip", None) + else "CUDA" + ) + raise RuntimeError( + f"{self.runtime_mode} AdaLN on {backend} requires the Triton " + "packed INT4 group kernel and will not materialize a full BF16 " + "weight. Install the matching PyTorch Triton package or use " + "runtime_mode='dequant_bf16'." + ) from exc + self.last_effective_runtime_mode = "triton_packed_adaln_int4" + return matmul_packed_adaln_int4_with_triton( + activation, + self.packed_weight, + self.scales, + out_features=self.out_features, + in_features=self.in_features, + group_size=self.group_size, + bias=self._bias_for(x.device, torch.bfloat16), + ) if x.device.type == "cpu" and self.runtime_mode in { "auto_fused", "native_packed_matmul", @@ -325,5 +422,5 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.last_effective_runtime_mode = "dequant_bf16" weight = self._dequantize_weight(device=x.device, dtype=torch.bfloat16) - bias = None if self.bias is None else self.bias.to(device=x.device, dtype=torch.bfloat16) + bias = self._bias_for(x.device, torch.bfloat16) return F.linear(activation, weight, bias) diff --git a/src/orbitquant/benchmarks.py b/src/orbitquant/benchmarks.py index 25dc9d4..9f042db 100644 --- a/src/orbitquant/benchmarks.py +++ b/src/orbitquant/benchmarks.py @@ -37,6 +37,10 @@ def _synchronize(device: torch.device) -> None: torch.cuda.synchronize(device) elif device.type == "mps" and torch.backends.mps.is_available(): torch.mps.synchronize() + elif device.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is not None and xpu.is_available(): + xpu.synchronize(device) def _mean_time_ms( @@ -64,6 +68,17 @@ def _mean_time_ms( end.record() end.synchronize() return float(start.elapsed_time(end) / iterations) + if device.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is not None and xpu.is_available(): + start = xpu.Event(enable_timing=True) + end = xpu.Event(enable_timing=True) + start.record() + for _ in range(iterations): + fn() + end.record() + end.synchronize() + return float(start.elapsed_time(end) / iterations) started_at = time.perf_counter() for _ in range(iterations): @@ -82,6 +97,16 @@ def _time_once_ms(fn: Callable[[], Any], *, device: torch.device) -> tuple[Any, end.record() end.synchronize() return result, float(start.elapsed_time(end)) + if device.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is not None and xpu.is_available(): + start = xpu.Event(enable_timing=True) + end = xpu.Event(enable_timing=True) + start.record() + result = fn() + end.record() + end.synchronize() + return result, float(start.elapsed_time(end)) started_at = time.perf_counter() result = fn() @@ -94,13 +119,22 @@ def _device_name(device: torch.device) -> str: return torch.cuda.get_device_name(device) if device.type == "mps": return "mps" + if device.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is not None and xpu.is_available(): + return str(xpu.get_device_name(device)) + return "xpu" return "cpu" def _peak_memory_bytes(device: torch.device) -> int | None: - if device.type != "cuda" or not torch.cuda.is_available(): - return None - return int(torch.cuda.max_memory_allocated(device)) + if device.type == "cuda" and torch.cuda.is_available(): + return int(torch.cuda.max_memory_allocated(device)) + if device.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is not None and xpu.is_available(): + return int(xpu.max_memory_allocated(device)) + return None def _runtime_uses_packed_matmul(runtime_mode: str, device: torch.device) -> bool: @@ -110,10 +144,14 @@ def _runtime_uses_packed_matmul(runtime_mode: str, device: torch.device) -> bool def _weight_quantization_backend_label(device: torch.device) -> str: - if device.type == "cuda" and backend_capabilities()["triton_cuda"]["available"]: - return "triton_cuda" + if device.type == "cuda": + backend = "triton_rocm" if getattr(torch.version, "hip", None) else "triton_cuda" + if backend_capabilities()[backend]["available"]: + return backend if device.type == "mps": return "torch_reference_mps" + if device.type == "xpu" and backend_capabilities()["triton_xpu"]["available"]: + return "triton_xpu" return "torch_reference" @@ -147,11 +185,18 @@ def benchmark_orbit_linear( raise RuntimeError("CUDA benchmark requested but CUDA is not available") if target_device.type == "mps" and not torch.backends.mps.is_available(): raise RuntimeError("MPS benchmark requested but MPS is not available") + if target_device.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is None or not xpu.is_available(): + raise RuntimeError("XPU benchmark requested but XPU is not available") torch.manual_seed(seed) if target_device.type == "cuda": torch.cuda.manual_seed_all(seed) torch.cuda.reset_peak_memory_stats(target_device) + elif target_device.type == "xpu": + torch.xpu.manual_seed_all(seed) + torch.xpu.reset_peak_memory_stats(target_device) config = OrbitQuantConfig( weight_bits=weight_bits, @@ -319,6 +364,7 @@ def forward_prewarmed() -> torch.Tensor: "peak_memory_bytes": _peak_memory_bytes(target_device), "quantization_buffers": { "source_weight_device": str(source.weight.device), + "source_weight_device_type": source.weight.device.type, "source_weight_is_cuda": bool(source.weight.is_cuda), "packed_weight_indices_device": ( None @@ -336,16 +382,24 @@ def forward_prewarmed() -> torch.Tensor: if quantized.packed_weight_indices is None else bool(quantized.packed_weight_indices.is_cuda) ), + "packed_weight_indices_device_type": ( + None + if quantized.packed_weight_indices is None + else quantized.packed_weight_indices.device.type + ), "row_norms_is_cuda": ( None if quantized.row_norms is None else bool(quantized.row_norms.is_cuda) ), + "row_norms_device_type": ( + None if quantized.row_norms is None else quantized.row_norms.device.type + ), }, "backend_capabilities": backend_capabilities(), "notes": ( "weight_quantize_pack_cold_ms includes first-use backend compilation " - "where applicable. On Triton/CUDA this can be CPU-heavy and show low " - "GPU utilization; weight_quantize_pack_hot_ms measures the already " - "compiled CUDA path. In dequant_bf16 mode, forward_prewarmed_ms uses " + "where applicable. On Triton accelerators this can be CPU-heavy and show " + "low device utilization; weight_quantize_pack_hot_ms measures the already " + "compiled path. In dequant_bf16 mode, forward_prewarmed_ms uses " "OrbitQuant activation kernels plus cached dequantized weights and " "PyTorch linear. In auto_fused and packed matmul runtime modes, " "forward_prewarmed_ms uses packed-weight matmul instead of the cached " @@ -433,12 +487,20 @@ def benchmark_model_quantization( raise RuntimeError("CUDA source device requested but CUDA is not available") if target.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA quantization device requested but CUDA is not available") + for label, selected in (("source", source), ("quantization", target)): + if selected.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is None or not xpu.is_available(): + raise RuntimeError(f"XPU {label} device requested but XPU is not available") torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) if target.type == "cuda": torch.cuda.reset_peak_memory_stats(target) + elif target.type == "xpu": + torch.xpu.manual_seed_all(seed) + torch.xpu.reset_peak_memory_stats(target) model = _QuantizeBenchModel( layers=layers, @@ -453,7 +515,15 @@ def benchmark_model_quantization( activation_bits=activation_bits, block_size=block_size, target_policy="flux2", - activation_kernel_backend="triton_cuda" if target.type == "cuda" else "auto", + activation_kernel_backend=( + "triton_rocm" + if target.type == "cuda" and getattr(torch.version, "hip", None) + else "triton_cuda" + if target.type == "cuda" + else "triton_xpu" + if target.type == "xpu" + else "auto" + ), runtime_mode="dequant_bf16", ) @@ -488,7 +558,7 @@ def benchmark_model_quantization( "notes": ( "This benchmark measures the full model replacement loop, including " "host-to-device staging when source_device is cpu and quantization_device " - "is cuda. The OrbitQuant/AdaLN compute fields exclude explicit transfer " + "is an accelerator. The OrbitQuant/AdaLN compute fields exclude explicit transfer " "time; device_transfer_seconds is reported separately." ), } diff --git a/src/orbitquant/cli/main.py b/src/orbitquant/cli/main.py index a32df50..1e8f76d 100644 --- a/src/orbitquant/cli/main.py +++ b/src/orbitquant/cli/main.py @@ -162,6 +162,10 @@ def _release_generation_pipeline(pipeline: Any, *, device: str) -> None: empty_cache = getattr(torch.mps, "empty_cache", None) if callable(empty_cache): empty_cache() + elif torch_device.type == "xpu" and hasattr(torch, "xpu"): + empty_cache = getattr(torch.xpu, "empty_cache", None) + if callable(empty_cache): + empty_cache() def _with_runtime_overrides( @@ -749,7 +753,7 @@ def main(argv: list[str] | None = None) -> int: kernel_bench_parser.add_argument( "--activation-kernel-backend", default="auto", - choices=["auto", "cpu", "mps", "triton_cuda"], + choices=["auto", "cpu", "mps", "triton_cuda", "triton_rocm", "triton_xpu"], ) kernel_bench_parser.add_argument( "--runtime-mode", @@ -874,7 +878,7 @@ def main(argv: list[str] | None = None) -> int: native_script_parser.add_argument( "--activation-kernel-backend", default="triton_cuda", - choices=["auto", "cpu", "mps", "triton_cuda"], + choices=["auto", "cpu", "mps", "triton_cuda", "triton_rocm", "triton_xpu"], ) native_script_parser.add_argument( "--runtime-mode", @@ -907,7 +911,7 @@ def main(argv: list[str] | None = None) -> int: quantize_parser.add_argument( "--activation-kernel-backend", default="auto", - choices=["auto", "cpu", "mps", "triton_cuda"], + choices=["auto", "cpu", "mps", "triton_cuda", "triton_rocm", "triton_xpu"], ) quantize_parser.add_argument("--device", default="auto") quantize_parser.add_argument( @@ -1160,7 +1164,7 @@ def main(argv: list[str] | None = None) -> int: generate_parser.add_argument( "--activation-kernel-backend", default="auto", - choices=["auto", "cpu", "mps", "triton_cuda"], + choices=["auto", "cpu", "mps", "triton_cuda", "triton_rocm", "triton_xpu"], ) generate_parser.add_argument( "--no-prewarm", @@ -1206,7 +1210,7 @@ def main(argv: list[str] | None = None) -> int: compare_native_parser.add_argument( "--activation-kernel-backend", default="auto", - choices=["auto", "cpu", "mps", "triton_cuda"], + choices=["auto", "cpu", "mps", "triton_cuda", "triton_rocm", "triton_xpu"], ) compare_native_parser.add_argument( "--no-prewarm", diff --git a/src/orbitquant/config.py b/src/orbitquant/config.py index 6e390ce..12387af 100644 --- a/src/orbitquant/config.py +++ b/src/orbitquant/config.py @@ -25,7 +25,14 @@ class QuantizationConfigMixin: # type: ignore[no-redef] "triton_packed_matmul", "native_packed_matmul", } -_SUPPORTED_ACTIVATION_KERNEL_BACKENDS = {"auto", "cpu", "mps", "triton_cuda"} +_SUPPORTED_ACTIVATION_KERNEL_BACKENDS = { + "auto", + "cpu", + "mps", + "triton_cuda", + "triton_rocm", + "triton_xpu", +} _SUPPORTED_TARGET_POLICIES = { "auto", "universal", @@ -86,6 +93,9 @@ class OrbitQuantConfig(QuantizationConfigMixin): packed_matmul_block_k: int = 128 packed_matmul_num_warps: int = 4 weight_row_tile_size: int = 256 + # Opt-in: keep a persistent INT8 copy of each W4A4 weight on CUDA (twice + # the packed size) so the cuBLASLt path skips its per-forward decode. + w4a4_int8_weight_cache: bool = False def __post_init__(self) -> None: if self.weight_bits not in _SUPPORTED_BITS: @@ -131,6 +141,8 @@ def __post_init__(self) -> None: raise ValueError(f"target_policy must be one of {sorted(_SUPPORTED_TARGET_POLICIES)}") if self.adaln_group_size <= 0: raise ValueError("adaln_group_size must be positive") + if not isinstance(self.w4a4_int8_weight_cache, bool): + raise ValueError("w4a4_int8_weight_cache must be a boolean") normalized_dtype_dict: dict[str, list[str]] = {} for dtype_name, module_names in self.modules_dtype_dict.items(): normalized_dtype = dtype_name.lower() diff --git a/src/orbitquant/kernels/dispatch.py b/src/orbitquant/kernels/dispatch.py index 7ab3982..4098353 100644 --- a/src/orbitquant/kernels/dispatch.py +++ b/src/orbitquant/kernels/dispatch.py @@ -1,5 +1,7 @@ from __future__ import annotations +from dataclasses import dataclass + import torch from orbitquant.codebooks import LloydMaxCodebook @@ -13,14 +15,24 @@ _MPS_NATIVE_STAGE = "packed_weight_matmul" _MPS_IMPLEMENTED_STAGE = f"{_MPS_SHADER_STAGE},{_MPS_NATIVE_STAGE}" _CPU_IMPLEMENTED_STAGE = ( - "activation_norm_rpbh_quant_rescale,packed_weight_matmul," - "adaln_rtn_packed_matmul" + "activation_norm_rpbh_quant_rescale,packed_weight_matmul,adaln_rtn_packed_matmul" ) _TRITON_CUDA_IMPLEMENTED_STAGE = ( "activation_norm_rpbh_quant_rescale,packed_weight_dequant," "packed_weight_matmul,lowbit_pack,lowbit_unpack,weight_rotation_fwht_quant_pack," - "adaln_rtn_quant_pack,adaln_rtn_dequant" + "adaln_rtn_quant_pack,adaln_rtn_dequant,adaln_rtn_packed_matmul" ) +_TRITON_ROCM_IMPLEMENTED_STAGE = _TRITON_CUDA_IMPLEMENTED_STAGE +_TRITON_XPU_IMPLEMENTED_STAGE = _TRITON_CUDA_IMPLEMENTED_STAGE + + +def _torch_uses_hip() -> bool: + return bool(getattr(torch.version, "hip", None)) + + +def _xpu_available() -> bool: + xpu = getattr(torch, "xpu", None) + return bool(xpu is not None and xpu.is_available()) def _triton_available() -> bool: @@ -87,198 +99,312 @@ def _join_stages(stages: list[str]) -> str | None: return ",".join(stages) if stages else None +_AVAILABLE_BACKENDS_CACHE: BackendAvailability | None = None + + +def clear_backend_availability_cache() -> None: + """Reset the memoized backend availability (tests and diagnostics).""" + global _AVAILABLE_BACKENDS_CACHE + _AVAILABLE_BACKENDS_CACHE = None + + def available_backends() -> BackendAvailability: - return { - "cpu": True, - "mps": bool(torch.backends.mps.is_available()), - "triton_cuda": bool(torch.cuda.is_available() and _triton_available()), + # Probing availability imports triton and touches every accelerator + # runtime; none of that changes within a process, and this runs on the + # per-forward dispatch path, so memoize the result. + global _AVAILABLE_BACKENDS_CACHE + if _AVAILABLE_BACKENDS_CACHE is None: + triton_available = _triton_available() + triton_gpu_available = bool(torch.cuda.is_available() and triton_available) + triton_xpu_available = bool(_xpu_available() and triton_available) + uses_hip = _torch_uses_hip() + _AVAILABLE_BACKENDS_CACHE = { + "cpu": True, + "mps": bool(torch.backends.mps.is_available()), + "triton_cuda": bool(triton_gpu_available and not uses_hip), + "triton_rocm": bool(triton_gpu_available and uses_hip), + "triton_xpu": triton_xpu_available, + } + return dict(_AVAILABLE_BACKENDS_CACHE) + + +def _capability_entry( + *, + available: bool, + claim_status: str, + optimized: bool, + implemented_stage: str, + optimized_stage: str | None, + device_types: list[str], + implementation: str, + package_format: str, + hf_kernel_builder_compliant: bool, + notes: str, + weight_dequant_optimized: bool = False, + weight_pack_optimized: bool = False, + lowbit_unpack_optimized: bool = False, + weight_quant_optimized: bool = False, + adaln_quant_optimized: bool = False, + adaln_dequant_optimized: bool = False, + adaln_packed_matmul_optimized: bool = False, + upstream_native_mps_op: bool | None = None, +) -> dict[str, object]: + """Assemble one backend capability payload with the canonical key set and order.""" + entry: dict[str, object] = { + "available": available, + "claim_status": claim_status, + "optimized": optimized, + "full_fusion": False, + "implemented_stage": implemented_stage, + "optimized_stage": optimized_stage, + "weight_dequant_optimized": weight_dequant_optimized, + "weight_pack_optimized": weight_pack_optimized, + "lowbit_unpack_optimized": lowbit_unpack_optimized, + "weight_quant_optimized": weight_quant_optimized, + "adaln_quant_optimized": adaln_quant_optimized, + "adaln_dequant_optimized": adaln_dequant_optimized, + "adaln_packed_matmul_optimized": adaln_packed_matmul_optimized, + "device_types": device_types, + "implementation": implementation, + "package_format": package_format, } + if upstream_native_mps_op is not None: + entry["upstream_native_mps_op"] = upstream_native_mps_op + entry["hf_kernel_builder_compliant"] = hf_kernel_builder_compliant + entry["notes"] = notes + return entry + + +# Per-stage CPU rows: (optimized stage name, note when native, note when reference). +# Row order matches the (activation, matmul, adaln) flag tuple in _cpu_capability. +_CPU_STAGE_TABLE: tuple[tuple[str, str, str], ...] = ( + ( + "activation_norm_rpbh_quant_rescale", + "The native exact activation pipeline performs FP32 token norm, " + "RPBH/FWHT, codebook assignment, and rescale.", + "Activation quantization uses the reference PyTorch CPU path.", + ), + ( + "packed_weight_matmul", + "Native exact packed matmul consumes low-bit weights without a full " + "floating-point cache.", + "The main linear matmul uses the reference PyTorch CPU path.", + ), + ( + "adaln_rtn_packed_matmul", + "Native packed INT4 AdaLN matmul consumes group-64 weights without a " + "full floating-point cache.", + "AdaLN uses the explicit BF16 dequantization path.", + ), +) - -def backend_capabilities(backends: BackendAvailability | None = None) -> BackendCapabilities: - available = available_backends() if backends is None else backends - cpu_native_matmul_optimized = bool( - available["cpu"] and _native_cpu_packed_matmul_available() - ) - cpu_native_activation_optimized = bool( - available["cpu"] and _native_cpu_activation_available() - ) - cpu_native_adaln_optimized = bool( - available["cpu"] and _native_cpu_adaln_available() - ) - cpu_optimized_stages = _join_stages( +# (activation optimized, matmul optimized) -> implementation label. +_CPU_IMPLEMENTATIONS: dict[tuple[bool, bool], str] = { + (True, True): "native_exact_activation+native_exact_packed_matmul", + (True, False): "native_exact_activation+reference_weight_matmul", + (False, True): "native_exact_packed_matmul+torch_activation_reference", + (False, False): "torch_reference", +} + + +def _cpu_capability(available: bool) -> dict[str, object]: + matmul_optimized = bool(available and _native_cpu_packed_matmul_available()) + activation_optimized = bool(available and _native_cpu_activation_available()) + adaln_optimized = bool(available and _native_cpu_adaln_available()) + stage_flags = (activation_optimized, matmul_optimized, adaln_optimized) + optimized_stages = _join_stages( [ - *( - ["activation_norm_rpbh_quant_rescale"] - if cpu_native_activation_optimized - else [] - ), - *(["packed_weight_matmul"] if cpu_native_matmul_optimized else []), - *( - ["adaln_rtn_packed_matmul"] - if cpu_native_adaln_optimized - else [] - ), + stage + for enabled, (stage, _, _) in zip(stage_flags, _CPU_STAGE_TABLE, strict=True) + if enabled ] ) - cpu_implementation = ( - "native_exact_activation+native_exact_packed_matmul" - if cpu_native_activation_optimized and cpu_native_matmul_optimized - else "native_exact_packed_matmul+torch_activation_reference" - if cpu_native_matmul_optimized - else "native_exact_activation+reference_weight_matmul" - if cpu_native_activation_optimized - else "torch_reference" - ) - if cpu_native_adaln_optimized: - cpu_implementation += "+native_packed_adaln_int4" - cpu_notes = [] - if cpu_native_activation_optimized: - cpu_notes.append( - "The native exact activation pipeline performs FP32 token norm, " - "RPBH/FWHT, codebook assignment, and rescale." + optimized = optimized_stages is not None + implementation = _CPU_IMPLEMENTATIONS[(activation_optimized, matmul_optimized)] + if adaln_optimized: + implementation += "+native_packed_adaln_int4" + notes = " ".join( + native_note if enabled else reference_note + for enabled, (_, native_note, reference_note) in zip( + stage_flags, _CPU_STAGE_TABLE, strict=True ) - else: - cpu_notes.append("Activation quantization uses the reference PyTorch CPU path.") - if cpu_native_matmul_optimized: - cpu_notes.append( - "Native exact packed matmul consumes low-bit weights without a full " - "floating-point cache." - ) - else: - cpu_notes.append("The main linear matmul uses the reference PyTorch CPU path.") - if cpu_native_adaln_optimized: - cpu_notes.append( - "Native packed INT4 AdaLN matmul consumes group-64 weights without a " - "full floating-point cache." - ) - else: - cpu_notes.append("AdaLN uses the explicit BF16 dequantization path.") - mps_shader_optimized = bool(available["mps"] and _mps_metal_available()) - mps_native_matmul_optimized = bool( - available["mps"] and _native_packed_matmul_available() ) - mps_optimized_stages = _join_stages( + return _capability_entry( + available=available, + claim_status="partial_optimized" if optimized else "reference_only", + optimized=optimized, + implemented_stage=_CPU_IMPLEMENTED_STAGE, + optimized_stage=optimized_stages, + adaln_dequant_optimized=adaln_optimized, + adaln_packed_matmul_optimized=adaln_optimized, + device_types=["cpu"], + implementation=implementation, + package_format=( + "native_kernel_package_torch_stable_abi" if optimized else "torch_reference" + ), + hf_kernel_builder_compliant=optimized, + notes=notes, + ) + + +# (shader optimized, native packed matmul optimized) +# -> (implementation, package_format, notes). +_MPS_VARIANTS: dict[tuple[bool, bool], tuple[str, str, str]] = { + (True, True): ( + "torch_mps_compile_shader_fused_activation+native_packed_matmul", + "torch.mps.compile_shader,native_kernel_package", + "Fused Metal shaders handle activation norm, RPBH/FWHT rotation, codebook " + "lookup/rescale, and packed weight dequant. The native packed matmul " + "package handles packed low-bit matmul.", + ), + (True, False): ( + "torch_mps_compile_shader_fused_activation", + "torch.mps.compile_shader", + "Fused Metal shaders handle activation norm, RPBH/FWHT rotation, codebook " + "lookup/rescale, and packed weight dequant. The native packed matmul " + "package is not available in this environment.", + ), + (False, True): ( + "native_packed_matmul", + "native_kernel_package", + "The native packed matmul package handles packed low-bit matmul. " + "Activation quantization helpers use the reference PyTorch path in this " + "environment.", + ), + (False, False): ( + "torch_reference_mps", + "torch_reference", + "Runs the reference PyTorch path on MPS tensors; native Metal shader " + "support and the native packed matmul package are not available in this " + "environment.", + ), +} + + +def _mps_capability(available: bool) -> dict[str, object]: + shader_optimized = bool(available and _mps_metal_available()) + native_matmul_optimized = bool(available and _native_packed_matmul_available()) + optimized_stages = _join_stages( [ - *([_MPS_SHADER_STAGE] if mps_shader_optimized else []), - *([_MPS_NATIVE_STAGE] if mps_native_matmul_optimized else []), + *([_MPS_SHADER_STAGE] if shader_optimized else []), + *([_MPS_NATIVE_STAGE] if native_matmul_optimized else []), ] ) - mps_optimized = mps_optimized_stages is not None - if mps_shader_optimized and mps_native_matmul_optimized: - mps_notes = ( - "Fused Metal shaders handle activation norm, RPBH/FWHT rotation, codebook " - "lookup/rescale, and packed weight dequant. The native packed matmul " - "package handles packed low-bit matmul." - ) - elif mps_shader_optimized: - mps_notes = ( - "Fused Metal shaders handle activation norm, RPBH/FWHT rotation, codebook " - "lookup/rescale, and packed weight dequant. The native packed matmul " - "package is not available in this environment." - ) - elif mps_native_matmul_optimized: - mps_notes = ( - "The native packed matmul package handles packed low-bit matmul. " - "Activation quantization helpers use the reference PyTorch path in this " - "environment." - ) - else: - mps_notes = ( - "Runs the reference PyTorch path on MPS tensors; native Metal shader " - "support and the native packed matmul package are not available in this " - "environment." - ) - return { - "cpu": { - "available": bool(available["cpu"]), - "claim_status": "partial_optimized" - if cpu_optimized_stages is not None - else "reference_only", - "optimized": cpu_optimized_stages is not None, - "full_fusion": False, - "implemented_stage": _CPU_IMPLEMENTED_STAGE, - "optimized_stage": cpu_optimized_stages, - "weight_dequant_optimized": False, - "weight_pack_optimized": False, - "lowbit_unpack_optimized": False, - "weight_quant_optimized": False, - "adaln_quant_optimized": False, - "adaln_dequant_optimized": cpu_native_adaln_optimized, - "device_types": ["cpu"], - "implementation": cpu_implementation, - "package_format": "native_kernel_package_torch_stable_abi" - if cpu_optimized_stages is not None - else "torch_reference", - "hf_kernel_builder_compliant": cpu_optimized_stages is not None, - "notes": " ".join(cpu_notes), - }, - "mps": { - "available": bool(available["mps"]), - "claim_status": "partial_optimized" if mps_optimized else "reference_only", - "optimized": mps_optimized, - "full_fusion": False, - "implemented_stage": _MPS_IMPLEMENTED_STAGE, - "optimized_stage": mps_optimized_stages, - "weight_dequant_optimized": mps_shader_optimized, - "weight_pack_optimized": False, - "lowbit_unpack_optimized": False, - "weight_quant_optimized": False, - "adaln_quant_optimized": False, - "adaln_dequant_optimized": False, - "device_types": ["mps"], - "implementation": ( - "torch_mps_compile_shader_fused_activation+native_packed_matmul" - if mps_shader_optimized and mps_native_matmul_optimized - else "torch_mps_compile_shader_fused_activation" - if mps_shader_optimized - else "native_packed_matmul" - if mps_native_matmul_optimized - else "torch_reference_mps" - ), - "package_format": ( - "torch.mps.compile_shader,native_kernel_package" - if mps_shader_optimized and mps_native_matmul_optimized - else "torch.mps.compile_shader" - if mps_shader_optimized - else "native_kernel_package" - if mps_native_matmul_optimized - else "torch_reference" - ), - "upstream_native_mps_op": False, - "hf_kernel_builder_compliant": False, - "notes": mps_notes, - }, - "triton_cuda": { - "available": bool(available["triton_cuda"]), - "claim_status": "partial_optimized" if available["triton_cuda"] else "unavailable", - "optimized": bool(available["triton_cuda"]), - "full_fusion": False, - "implemented_stage": _TRITON_CUDA_IMPLEMENTED_STAGE, - "optimized_stage": _TRITON_CUDA_IMPLEMENTED_STAGE - if available["triton_cuda"] - else None, - "weight_dequant_optimized": bool(available["triton_cuda"]), - "weight_pack_optimized": bool(available["triton_cuda"]), - "lowbit_unpack_optimized": bool(available["triton_cuda"]), - "weight_quant_optimized": bool(available["triton_cuda"]), - "adaln_quant_optimized": bool(available["triton_cuda"]), - "adaln_dequant_optimized": bool(available["triton_cuda"]), - "device_types": ["cuda"], - "implementation": "python_triton_orbitquant_pipeline", - "package_format": "python_triton", - "hf_kernel_builder_compliant": False, - "notes": ( - "Triton handles runtime activation norm, RPBH/FWHT rotation, codebook " - "lookup/rescale, packed weight dequant, packed weight matmul, " - "offline low-bit pack/unpack, offline weight RPBH/FWHT codebook indexing " - "with direct low-bit packing, and AdaLN INT4 RTN quantize/pack/dequant. " - "The default auto_fused runtime prefers packed low-bit matmul when a " - "native or Triton packed kernel is available; full-model speedup claims " - "still require separate benchmark artifacts." - ), - }, + optimized = optimized_stages is not None + implementation, package_format, notes = _MPS_VARIANTS[ + (shader_optimized, native_matmul_optimized) + ] + return _capability_entry( + available=available, + claim_status="partial_optimized" if optimized else "reference_only", + optimized=optimized, + implemented_stage=_MPS_IMPLEMENTED_STAGE, + optimized_stage=optimized_stages, + weight_dequant_optimized=shader_optimized, + device_types=["mps"], + implementation=implementation, + package_format=package_format, + upstream_native_mps_op=False, + hf_kernel_builder_compliant=False, + notes=notes, + ) + + +@dataclass(frozen=True) +class _TritonBackendSpec: + """Static description of one Triton backend flavor.""" + + implemented_stage: str + optimized_when_available: bool + available_claim_status: str + device_type: str + implementation: str + package_format: str + notes: str + + +_TRITON_SPECS: dict[str, _TritonBackendSpec] = { + "triton_cuda": _TritonBackendSpec( + implemented_stage=_TRITON_CUDA_IMPLEMENTED_STAGE, + optimized_when_available=True, + available_claim_status="partial_optimized", + device_type="cuda", + implementation="python_triton_orbitquant_pipeline", + package_format="python_triton", + notes=( + "Triton handles runtime activation norm, RPBH/FWHT rotation, codebook " + "lookup/rescale, packed weight dequant, packed weight matmul, " + "offline low-bit pack/unpack, offline weight RPBH/FWHT codebook indexing " + "with direct low-bit packing, and AdaLN INT4 RTN quantize/pack, " + "dequant, and packed matmul. " + "The default auto_fused runtime prefers packed low-bit matmul when a " + "native or Triton packed kernel is available; full-model speedup claims " + "still require separate benchmark artifacts." + ), + ), + "triton_rocm": _TritonBackendSpec( + implemented_stage=_TRITON_ROCM_IMPLEMENTED_STAGE, + optimized_when_available=False, + available_claim_status="experimental_unverified", + device_type="cuda", + implementation="python_triton_orbitquant_pipeline_rocm_candidate", + package_format="python_triton_rocm", + notes=( + "PyTorch exposes HIP tensors through the cuda device type. The candidate " + "reuses the exact packed Triton pipeline without loading CUDA-native " + "extensions. It remains experimental until correctness, memory, profiler, " + "and performance evidence is recorded on supported AMD hardware." + ), + ), + "triton_xpu": _TritonBackendSpec( + implemented_stage=_TRITON_XPU_IMPLEMENTED_STAGE, + optimized_when_available=False, + available_claim_status="experimental_unverified", + device_type="xpu", + implementation="python_triton_orbitquant_pipeline_xpu_candidate", + package_format="python_triton_xpu", + notes=( + "The candidate reuses the exact packed Triton pipeline on torch.xpu. " + "It remains explicit-only and experimental until correctness, memory, " + "profiler, and performance evidence is recorded on supported Intel GPU " + "hardware." + ), + ), +} + + +def _triton_capability(spec: _TritonBackendSpec, available: bool) -> dict[str, object]: + optimized = bool(available and spec.optimized_when_available) + return _capability_entry( + available=available, + claim_status=spec.available_claim_status if available else "unavailable", + optimized=optimized, + implemented_stage=spec.implemented_stage, + optimized_stage=spec.implemented_stage if optimized else None, + weight_dequant_optimized=optimized, + weight_pack_optimized=optimized, + lowbit_unpack_optimized=optimized, + weight_quant_optimized=optimized, + adaln_quant_optimized=optimized, + adaln_dequant_optimized=optimized, + adaln_packed_matmul_optimized=optimized, + device_types=[spec.device_type], + implementation=spec.implementation, + package_format=spec.package_format, + hf_kernel_builder_compliant=False, + notes=spec.notes, + ) + + +def backend_capabilities(backends: BackendAvailability | None = None) -> BackendCapabilities: + available = available_backends() if backends is None else backends + capabilities: BackendCapabilities = { + "cpu": _cpu_capability(bool(available["cpu"])), + "mps": _mps_capability(bool(available["mps"])), } + for name, spec in _TRITON_SPECS.items(): + capabilities[name] = _triton_capability(spec, bool(available.get(name, False))) + return capabilities def select_backend( @@ -290,7 +416,11 @@ def select_backend( requested_device = torch.device(device) available = available_backends() if backends is None else backends if requested == "auto": - if requested_device.type == "cuda" and available["triton_cuda"]: + if ( + requested_device.type == "cuda" + and not _torch_uses_hip() + and available.get("triton_cuda", False) + ): return "triton_cuda" if requested_device.type == "mps" and available["mps"]: return "mps" @@ -298,9 +428,19 @@ def select_backend( if requested == "cpu": return "cpu" if requested == "triton_cuda": - if not available["triton_cuda"]: + if not available.get("triton_cuda", False): raise RuntimeError("CUDA/Triton activation kernel is not available in this environment") return "triton_cuda" + if requested == "triton_rocm": + if not available.get("triton_rocm", False): + raise RuntimeError("ROCm/Triton activation kernel is not available in this environment") + return "triton_rocm" + if requested == "triton_xpu": + if not available.get("triton_xpu", False): + raise RuntimeError("XPU/Triton activation kernel is not available in this environment") + if requested_device.type != "xpu": + raise RuntimeError("XPU/Triton activation kernel requires an XPU device") + return "triton_xpu" if requested == "mps": if not available["mps"]: raise RuntimeError("MPS activation kernel is not available in this environment") @@ -339,7 +479,7 @@ def _mps_quantize_activations( ) -def _triton_cuda_quantize_activations( +def _triton_quantize_activations( x: torch.Tensor, *, rotation: RPBHRotation, @@ -378,8 +518,8 @@ def quantize_activations_kernel( eps=eps, constant_tensors=constant_tensors, ) - if selected == "triton_cuda": - return _triton_cuda_quantize_activations( + if selected in {"triton_cuda", "triton_rocm", "triton_xpu"}: + return _triton_quantize_activations( x, rotation=rotation, codebook=codebook, diff --git a/src/orbitquant/kernels/executorch_vulkan.py b/src/orbitquant/kernels/executorch_vulkan.py new file mode 100644 index 0000000..5a072d3 --- /dev/null +++ b/src/orbitquant/kernels/executorch_vulkan.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from torch import nn + +from orbitquant.packing import unpack_lowbit +from orbitquant.rotations.fwht import fwht + +if TYPE_CHECKING: + from orbitquant.layers import OrbitQuantLinear + + +def _linear_w4a4_exact_reference( + x: torch.Tensor, + packed_weight: torch.Tensor, + row_norms: torch.Tensor, + permutation: torch.Tensor, + signs: torch.Tensor, + activation_boundaries: torch.Tensor, + pair_lut: torch.Tensor, + block_size: int, + eps: float, + bias: torch.Tensor | None, +) -> torch.Tensor: + in_features = x.shape[-1] + out_features = row_norms.numel() + work = x.to(dtype=torch.float32) + token_norms = torch.linalg.vector_norm(work, dim=-1, keepdim=True) + rotated = work.index_select(-1, permutation.to(dtype=torch.int64)) + rotated = rotated * signs.to(dtype=torch.float32) + rotated = rotated / (token_norms + eps) + rotated = rotated.reshape(*rotated.shape[:-1], in_features // block_size, block_size) + rotated = fwht(rotated) * (block_size**-0.5) + rotated = rotated.reshape(*work.shape) + activation_indices = torch.bucketize(rotated, activation_boundaries) + + weight_indices = unpack_lowbit( + packed_weight, + bits=4, + length=out_features * in_features, + ).to(dtype=torch.int64) + weight_indices = weight_indices.reshape(out_features, in_features) + flat_activation_indices = activation_indices.reshape(-1, in_features) + flat_token_norms = token_norms.reshape(-1) + output = torch.empty( + (flat_activation_indices.shape[0], out_features), + device=x.device, + dtype=torch.float32, + ) + for out_feature in range(out_features): + lut_indices = flat_activation_indices * 16 + weight_indices[out_feature] + values = pair_lut[lut_indices].sum(dim=-1) + values = values * flat_token_norms * row_norms[out_feature] + if bias is not None: + values = values + bias[out_feature] + output[:, out_feature] = values + output = output.reshape(*x.shape[:-1], out_features) + return output.to(dtype=x.dtype) + + +@torch.library.custom_op("orbitquant_vulkan::linear_w4a4_exact", mutates_args=()) +def linear_w4a4_exact( + x: torch.Tensor, + packed_weight: torch.Tensor, + row_norms: torch.Tensor, + permutation: torch.Tensor, + signs: torch.Tensor, + activation_boundaries: torch.Tensor, + pair_lut: torch.Tensor, + block_size: int, + eps: float, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + return _linear_w4a4_exact_reference( + x, + packed_weight, + row_norms, + permutation, + signs, + activation_boundaries, + pair_lut, + block_size, + eps, + bias, + ) + + +@linear_w4a4_exact.register_fake +def _linear_w4a4_exact_fake( + x: torch.Tensor, + packed_weight: torch.Tensor, + row_norms: torch.Tensor, + permutation: torch.Tensor, + signs: torch.Tensor, + activation_boundaries: torch.Tensor, + pair_lut: torch.Tensor, + block_size: int, + eps: float, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + del ( + packed_weight, + permutation, + signs, + activation_boundaries, + pair_lut, + block_size, + eps, + bias, + ) + return x.new_empty((*x.shape[:-1], row_norms.numel())) + + +class ExecuTorchVulkanW4A4Linear(nn.Module): + """Export-only exact W4A4 OrbitLinear for the ExecuTorch Vulkan delegate.""" + + def __init__(self, layer: OrbitQuantLinear) -> None: + super().__init__() + if layer.weight_bits != 4 or layer.activation_bits != 4: + raise ValueError("ExecuTorch Vulkan currently implements exact W4A4 only") + if layer.packed_weight_indices is None or layer.row_norms is None: + raise ValueError("OrbitQuantLinear is missing packed weights or row norms") + if layer.in_features % 8 != 0: + raise ValueError("ExecuTorch Vulkan W4A4 requires in_features divisible by 8") + block_size = int(layer.rotation.block_size) + if block_size < 8 or block_size > 4096: + raise ValueError( + "ExecuTorch Vulkan W4A4 currently requires an RPBH block size in [8, 4096]" + ) + if block_size & (block_size - 1) or layer.in_features % block_size != 0: + raise ValueError( + "ExecuTorch Vulkan W4A4 requires a power-of-two RPBH block size " + "dividing in_features" + ) + + self.in_features = layer.in_features + self.out_features = layer.out_features + self.block_size = block_size + self.eps = float(layer.activation_eps) + self.register_buffer( + "packed_weight", + layer.packed_weight_indices.detach().to(device="cpu", dtype=torch.uint8).contiguous(), + ) + self.register_buffer( + "row_norms", + layer.row_norms.detach().to(device="cpu", dtype=torch.float32).contiguous(), + ) + self.register_buffer( + "permutation", + layer.rotation.permutation.detach().to(device="cpu", dtype=torch.int32).contiguous(), + ) + self.register_buffer( + "signs", + layer.rotation.signs.detach().to(device="cpu", dtype=torch.int32).contiguous(), + ) + activation_centroids = layer.activation_codebook.centroids.detach().to( + device="cpu", dtype=torch.float32 + ) + weight_centroids = layer.weight_codebook.centroids.detach().to( + device="cpu", dtype=torch.float32 + ) + self.register_buffer( + "activation_boundaries", + layer.activation_codebook.boundaries.detach() + .to(device="cpu", dtype=torch.float32) + .contiguous(), + ) + self.register_buffer( + "pair_lut", + (activation_centroids[:, None] * weight_centroids[None, :]).flatten().contiguous(), + ) + if layer.bias is None: + self.register_buffer("bias", None) + else: + self.register_buffer( + "bias", + layer.bias.detach().to(device="cpu", dtype=torch.float32).contiguous(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.dtype not in {torch.float16, torch.float32}: + raise TypeError( + "ExecuTorch Vulkan accepts float16 or float32 activations; " + f"got {x.dtype}. Convert the exported model explicitly." + ) + if x.shape[-1] != self.in_features: + raise ValueError(f"expected last dimension {self.in_features}, got {x.shape[-1]}") + return linear_w4a4_exact( + x, + self.packed_weight, + self.row_norms, + self.permutation, + self.signs, + self.activation_boundaries, + self.pair_lut, + self.block_size, + self.eps, + self.bias, + ) + + +def prepare_executorch_vulkan_w4a4_model(model: nn.Module) -> nn.Module: + """Replace loaded OrbitQuant W4A4 linears with export-only Vulkan modules.""" + + from orbitquant.layers import OrbitQuantLinear + + if isinstance(model, OrbitQuantLinear): + return ExecuTorchVulkanW4A4Linear(model) + + def replace_children(module: nn.Module, prefix: str) -> None: + for child_name, child in list(module.named_children()): + qualified_name = f"{prefix}.{child_name}" if prefix else child_name + if isinstance(child, OrbitQuantLinear): + try: + replacement = ExecuTorchVulkanW4A4Linear(child) + except (TypeError, ValueError) as exc: + raise ValueError( + f"cannot prepare {qualified_name!r} for ExecuTorch Vulkan: {exc}" + ) from exc + module._modules[child_name] = replacement + else: + replace_children(child, qualified_name) + + replace_children(model, "") + return model + + +def register_executorch_vulkan_w4a4() -> object: + """Register the custom op with the current ExecuTorch Vulkan partitioner.""" + + try: + from executorch.backends.vulkan import utils + from executorch.backends.vulkan.op_registry import ( + OpFeatures, + update_features, + vulkan_supported_ops, + ) + from executorch.exir.dialects._ops import ops as exir_ops + except ImportError as exc: + raise RuntimeError( + "ExecuTorch Vulkan export requires an ExecuTorch source build containing " + "desktop Vulkan support (upstream after 2026-07-01)." + ) from exc + + edge_op = exir_ops.edge.orbitquant_vulkan.linear_w4a4_exact.default + if edge_op not in vulkan_supported_ops: + + @update_features(edge_op) + def register_features() -> OpFeatures: + return OpFeatures( + inputs_storage=utils.CONTIGUOUS_BUFFER, + inputs_dtypes=utils.ALL_T, + outputs_dtypes=utils.FP_T, + supports_highdim=True, + supports_prepacking=True, + ) + + return edge_op + + +__all__ = [ + "ExecuTorchVulkanW4A4Linear", + "linear_w4a4_exact", + "prepare_executorch_vulkan_w4a4_model", + "register_executorch_vulkan_w4a4", +] diff --git a/src/orbitquant/kernels/mps.py b/src/orbitquant/kernels/mps.py index 9e947f2..2539282 100644 --- a/src/orbitquant/kernels/mps.py +++ b/src/orbitquant/kernels/mps.py @@ -15,7 +15,7 @@ constant uint orbitquant_activation_threads = 256; constant uint orbitquant_wide_activation_threads = 512; -constant uint orbitquant_max_rpbh_block_size = 4096; +constant uint orbitquant_max_rpbh_block_size = __ORBITQUANT_MAX_BLOCK__; template inline void orbitquant_row_norm_impl( @@ -26,7 +26,9 @@ threadgroup float* scratch, uint row, uint thread_index, - uint thread_count) { + uint thread_count, + uint simd_lane, + uint simd_group) { if (row >= uint(rows)) { return; } @@ -37,17 +39,20 @@ const float value = float(input[row_offset + col]); sum += value * value; } - scratch[thread_index] = sum; - threadgroup_barrier(mem_flags::mem_threadgroup); - - for (uint stride = thread_count >> 1; stride > 0; stride >>= 1) { - if (thread_index < stride) { - scratch[thread_index] += scratch[thread_index + stride]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); + // One simdgroup reduction plus a single cross-warp pass replaces the + // log2(threads) barrier tree. + const float warp_sum = simd_sum(sum); + if (simd_lane == 0) { + scratch[simd_group] = warp_sum; } + threadgroup_barrier(mem_flags::mem_threadgroup); if (thread_index == 0) { - norms[row] = sqrt(scratch[0]); + float total = 0.0f; + const uint warps = (thread_count + 31) / 32; + for (uint warp = 0; warp < warps; ++warp) { + total += scratch[warp]; + } + norms[row] = sqrt(total); } } @@ -122,10 +127,13 @@ constant int& dim [[buffer(3)]], \ uint3 group_id [[threadgroup_position_in_grid]], \ uint thread_index [[thread_index_in_threadgroup]], \ + uint simd_lane [[thread_index_in_simdgroup]], \ + uint simd_group [[simdgroup_index_in_threadgroup]], \ uint3 group_size [[threads_per_threadgroup]]) { \ - threadgroup float scratch[THREADS]; \ + threadgroup float scratch[(THREADS + 31) / 32]; \ orbitquant_row_norm_impl( \ - input, norms, rows, dim, scratch, group_id.x, thread_index, group_size.x); \ + input, norms, rows, dim, scratch, group_id.x, thread_index, \ + group_size.x, simd_lane, simd_group); \ } ORBITQUANT_ROW_NORM_KERNEL( @@ -232,11 +240,21 @@ def mps_metal_available() -> bool: return bool(torch.backends.mps.is_available() and hasattr(torch.mps, "compile_shader")) -@lru_cache(maxsize=1) -def _mps_shader(): +@lru_cache(maxsize=8) +def _mps_shader(max_block_size: int = 512): + """Compile the shader library with the fused-kernel threadgroup array + sized for the requested RPBH block bucket (over-allocating to the 4096 + worst case would cap occupancy on every smaller block).""" if not mps_metal_available(): raise RuntimeError("MPS Metal shader backend is not available in this environment") - return torch.mps.compile_shader(_MPS_KERNEL_SOURCE) + source = _MPS_KERNEL_SOURCE.replace( + "__ORBITQUANT_MAX_BLOCK__", str(int(max_block_size)) + ) + return torch.mps.compile_shader(source) + + +def _fused_block_bucket(block_size: int) -> int: + return max(128, int(block_size)) def _activation_kernel_names(dtype: torch.dtype) -> tuple[str, str, str]: @@ -310,7 +328,7 @@ def quantize_activations_with_mps( norms = torch.empty(rows, device=x.device, dtype=torch.float32) output = torch.empty_like(flat) norm_name, wide_norm_name, activation_name = _activation_kernel_names(x.dtype) - shader = _mps_shader() + shader = _mps_shader(_fused_block_bucket(rotation.block_size)) activation_threads = ( 512 if rotation.block_size == 4096 and rotation.num_blocks == 1 else 256 ) diff --git a/src/orbitquant/kernels/triton_cuda.py b/src/orbitquant/kernels/triton_cuda.py index 8e9f499..ad8aa88 100644 --- a/src/orbitquant/kernels/triton_cuda.py +++ b/src/orbitquant/kernels/triton_cuda.py @@ -5,13 +5,62 @@ from orbitquant.codebooks import LloydMaxCodebook from orbitquant.rotations import RPBHRotation +_TRITON_DEVICE_TYPES = {"cuda", "xpu"} + + +def _xpu_available() -> bool: + xpu = getattr(torch, "xpu", None) + return bool(xpu is not None and xpu.is_available()) + + +def _triton_backend_label(device: torch.device | None = None) -> str: + if device is not None and device.type == "xpu": + return "triton_xpu" + if getattr(torch.version, "hip", None): + return "triton_rocm" + return "triton_cuda" + + +def _is_triton_device_tensor(tensor: torch.Tensor) -> bool: + return tensor.device.type in _TRITON_DEVICE_TYPES + + +def _triton_device_available(device: torch.device) -> bool: + if device.type == "cuda": + return torch.cuda.is_available() + if device.type == "xpu": + return _xpu_available() + return False + + +def _require_triton_tensor(tensor: torch.Tensor, operation: str) -> None: + if not _is_triton_device_tensor(tensor): + backend = _triton_backend_label(tensor.device) + raise RuntimeError( + f"{backend} {operation} requires a CUDA, HIP, or XPU tensor" + ) + + +def _require_triton_device(device: torch.device | str, operation: str) -> torch.device: + target_device = torch.device(device) + if target_device.type not in _TRITON_DEVICE_TYPES: + backend = _triton_backend_label(target_device) + raise RuntimeError( + f"{backend} {operation} requires a CUDA, HIP, or XPU device" + ) + if not _triton_device_available(target_device): + backend = _triton_backend_label(target_device) + raise RuntimeError(f"{backend} device {target_device} is not available") + return _normalize_triton_device(target_device) + def _load_triton(): try: import triton import triton.language as tl except Exception as exc: # pragma: no cover - environment dependent - raise RuntimeError("Triton is required for the triton_cuda backend") from exc + backend = "triton_xpu" if _xpu_available() else _triton_backend_label() + raise RuntimeError(f"Triton is required for the {backend} backend") from exc return triton, tl @@ -21,10 +70,12 @@ def _load_triton(): _INT8_CENTROID_SURROGATE_CACHE: dict[tuple[float, ...], tuple[torch.Tensor, float]] = {} -def _normalize_cuda_device(device: torch.device | str) -> torch.device: +def _normalize_triton_device(device: torch.device | str) -> torch.device: torch_device = torch.device(device) if torch_device.type == "cuda" and torch_device.index is None: return torch.device("cuda", torch.cuda.current_device()) + if torch_device.type == "xpu" and torch_device.index is None: + return torch.device("xpu", torch.xpu.current_device()) return torch_device @@ -32,14 +83,14 @@ def clear_triton_constant_cache() -> None: _WEIGHT_QUANT_CONSTANT_CACHE.clear() -def _cached_cuda_tensor( +def _cached_triton_tensor( key: tuple[object, ...], source: torch.Tensor, *, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: - device = _normalize_cuda_device(device) + device = _normalize_triton_device(device) cached = _WEIGHT_QUANT_CONSTANT_CACHE.get(key) if cached is not None and cached.device == device and cached.dtype == dtype: return cached @@ -54,7 +105,7 @@ def _weight_quantization_constants( codebook: LloydMaxCodebook, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - device = _normalize_cuda_device(device) + device = _normalize_triton_device(device) device_key = str(device) rotation_key = (int(rotation.dim), int(rotation.seed), int(rotation.block_size), device_key) codebook_key = ( @@ -63,19 +114,19 @@ def _weight_quantization_constants( int(codebook.algorithm_version), device_key, ) - permutation = _cached_cuda_tensor( + permutation = _cached_triton_tensor( ("rotation_permutation", *rotation_key), rotation.permutation, device=device, - dtype=torch.int64, + dtype=torch.int32, ) - signs = _cached_cuda_tensor( + signs = _cached_triton_tensor( ("rotation_signs", *rotation_key), rotation.signs, device=device, dtype=torch.int8, ) - boundaries = _cached_cuda_tensor( + boundaries = _cached_triton_tensor( ("codebook_boundaries", *codebook_key), codebook.boundaries, device=device, @@ -118,13 +169,11 @@ def _permute_sign_normalize_activation_kernel( mask = offsets < total rows = offsets // dim cols = offsets % dim - source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int64) + source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int32) signs = tl.load(signs_ptr + cols, mask=mask, other=1).to(tl.float32) norms = tl.load(norms_ptr + rows, mask=mask, other=1.0).to(tl.float32) denom = norms + eps - values = tl.load(input_ptr + rows * dim + source_cols, mask=mask, other=0.0).to( - tl.float32 - ) + values = tl.load(input_ptr + rows * dim + source_cols, mask=mask, other=0.0).to(tl.float32) tl.store(work_ptr + offsets, values * signs / denom, mask=mask) @@ -192,7 +241,7 @@ def _fused_rpbh_quantize_activation_kernel( local_cols = tl.arange(0, orbit_block_size) cols = orbit_block * orbit_block_size + local_cols mask = (row < rows) & (cols < dim) - source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int64) + source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int32) signs = tl.load(signs_ptr + cols, mask=mask, other=1).to(tl.float32) norm = tl.load(norms_ptr + row, mask=row < rows, other=0.0).to(tl.float32) denominator = norm + eps @@ -260,7 +309,7 @@ def _fused_rpbh_quantize_activation_pack_w4_kernel( local_cols = tl.arange(0, orbit_block_size) cols = orbit_block * orbit_block_size + local_cols mask = (row < rows) & (cols < dim) - source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int64) + source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int32) signs = tl.load(signs_ptr + cols, mask=mask, other=1).to(tl.float32) norm = tl.load(norms_ptr + row, mask=row < rows, other=0.0).to(tl.float32) values = tl.load( @@ -290,6 +339,10 @@ def _fused_rpbh_quantize_activation_pack_w4_kernel( values = _fwht_local_stage(values, orbit_block_size, 256) if fwht_stages > 9: values = _fwht_local_stage(values, orbit_block_size, 512) + if fwht_stages > 10: + values = _fwht_local_stage(values, orbit_block_size, 1024) + if fwht_stages > 11: + values = _fwht_local_stage(values, orbit_block_size, 2048) values *= inv_sqrt_block indices = tl.zeros((orbit_block_size,), dtype=tl.int32) @@ -300,11 +353,7 @@ def _fused_rpbh_quantize_activation_pack_w4_kernel( low, high = tl.split(pairs) packed = low | (high << 4) local_bytes = tl.arange(0, orbit_block_size // 2) - output_offsets = ( - row * (dim // 2) - + orbit_block * (orbit_block_size // 2) - + local_bytes - ) + output_offsets = row * (dim // 2) + orbit_block * (orbit_block_size // 2) + local_bytes tl.store(output_ptr + output_offsets, packed.to(tl.uint8), mask=row < rows) @@ -330,7 +379,7 @@ def _fused_rpbh_fwht_1024_chunk_kernel( local_cols = tl.arange(0, 1024) cols = orbit_block * orbit_block_size + chunk * 1024 + local_cols mask = (row < rows) & (cols < dim) - source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int64) + source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int32) signs = tl.load(signs_ptr + cols, mask=mask, other=1).to(tl.float32) norm = tl.load(norms_ptr + row, mask=row < rows, other=0.0).to(tl.float32) values = tl.load( @@ -391,9 +440,7 @@ def _quantize_activation_work_pack_w4_kernel( mask = byte_offsets < packed_total value_offsets = byte_offsets * 2 low_values = tl.load(work_ptr + value_offsets, mask=mask, other=0.0).to(tl.float32) - high_values = tl.load(work_ptr + value_offsets + 1, mask=mask, other=0.0).to( - tl.float32 - ) + high_values = tl.load(work_ptr + value_offsets + 1, mask=mask, other=0.0).to(tl.float32) low_values *= inv_sqrt_block high_values *= inv_sqrt_block low_indices = tl.zeros((block_size,), dtype=tl.int32) @@ -491,9 +538,7 @@ def _dequantize_packed_weight_kernel( bit_offsets = bit_starts % 8 low = tl.load(packed_ptr + byte_indices, mask=mask, other=0).to(tl.uint32) needs_high = bit_offsets + bits > 8 - high = tl.load(packed_ptr + byte_indices + 1, mask=mask & needs_high, other=0).to( - tl.uint32 - ) + high = tl.load(packed_ptr + byte_indices + 1, mask=mask & needs_high, other=0).to(tl.uint32) raw = low | (high << 8) indices = (raw >> bit_offsets) & ((1 << bits) - 1) rows = offsets // in_features @@ -548,9 +593,7 @@ def _matmul_packed_weight_kernel( ).to(tl.uint32) raw = low | (high << 8) indices = (raw >> bit_offsets) & ((1 << bits) - 1) - centroids = tl.load(centroids_ptr + indices, mask=weight_mask, other=0.0).to( - tl.float32 - ) + centroids = tl.load(centroids_ptr + indices, mask=weight_mask, other=0.0).to(tl.float32) norms = tl.load(row_norms_ptr + offs_n, mask=offs_n < out_features, other=0.0).to( tl.float32 ) @@ -615,12 +658,8 @@ def _matmul_packed_weight_w4_kernel( ) indices0 = packed_values & 15 indices1 = (packed_values >> 4) & 15 - centroids0 = tl.load(centroids_ptr + indices0, mask=weight_mask, other=0.0).to( - tl.float32 - ) - centroids1 = tl.load(centroids_ptr + indices1, mask=weight_mask, other=0.0).to( - tl.float32 - ) + centroids0 = tl.load(centroids_ptr + indices0, mask=weight_mask, other=0.0).to(tl.float32) + centroids1 = tl.load(centroids_ptr + indices1, mask=weight_mask, other=0.0).to(tl.float32) norms = tl.load(row_norms_ptr + offs_n, mask=offs_n < out_features, other=0.0).to( tl.float32 ) @@ -640,6 +679,145 @@ def _matmul_packed_weight_w4_kernel( ) +@triton.jit +def _matmul_packed_adaln_int4_group_kernel( + input_ptr, + packed_ptr, + scales_ptr, + bias_ptr, + output_ptr, + rows: tl.constexpr, + out_features: tl.constexpr, + in_features: tl.constexpr, + padded_in_features: tl.constexpr, + num_groups: tl.constexpr, + group_size: tl.constexpr, + has_bias: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_k: tl.constexpr, +): + """Group-aligned AdaLN matmul: every k-tile stays inside one scale group, + so scales load as one row vector instead of a (block_k, block_n) gather.""" + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * block_m + tl.arange(0, block_m) + offs_n = pid_n * block_n + tl.arange(0, block_n) + offs_k = tl.arange(0, block_k) + col_mask = offs_n < out_features + accumulator = tl.zeros((block_m, block_n), dtype=tl.float32) + + for k_start in range(0, in_features, block_k): + k = k_start + offs_k + input_mask = (offs_m[:, None] < rows) & (k[None, :] < in_features) + input_values = tl.load( + input_ptr + offs_m[:, None] * in_features + k[None, :], + mask=input_mask, + other=0.0, + ) + + value_offsets = offs_n[None, :] * padded_in_features + k[:, None] + weight_mask = col_mask[None, :] & (k[:, None] < in_features) + packed = tl.load( + packed_ptr + value_offsets // 2, + mask=weight_mask, + other=0, + ).to(tl.uint32) + shifts = (value_offsets & 1) * 4 + unsigned = (packed >> shifts) & 15 + signed = unsigned.to(tl.int32) - 8 + group = k_start // group_size + scales = tl.load( + scales_ptr + offs_n * num_groups + group, + mask=col_mask, + other=0.0, + ).to(tl.float32) + weights = (signed.to(tl.float32) * scales[None, :]).to(input_values.dtype) + accumulator += tl.dot(input_values, weights) + + if has_bias: + bias = tl.load( + bias_ptr + offs_n, + mask=col_mask, + other=0.0, + ).to(tl.float32) + accumulator += bias[None, :] + + tl.store( + output_ptr + offs_m[:, None] * out_features + offs_n[None, :], + accumulator, + mask=(offs_m[:, None] < rows) & col_mask[None, :], + ) + + +@triton.jit +def _matmul_packed_adaln_int4_kernel( + input_ptr, + packed_ptr, + scales_ptr, + bias_ptr, + output_ptr, + rows: tl.constexpr, + out_features: tl.constexpr, + in_features: tl.constexpr, + padded_in_features: tl.constexpr, + num_groups: tl.constexpr, + group_size: tl.constexpr, + has_bias: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_k: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * block_m + tl.arange(0, block_m) + offs_n = pid_n * block_n + tl.arange(0, block_n) + offs_k = tl.arange(0, block_k) + accumulator = tl.zeros((block_m, block_n), dtype=tl.float32) + + for k_start in range(0, in_features, block_k): + k = k_start + offs_k + input_mask = (offs_m[:, None] < rows) & (k[None, :] < in_features) + input_values = tl.load( + input_ptr + offs_m[:, None] * in_features + k[None, :], + mask=input_mask, + other=0.0, + ) + + value_offsets = offs_n[None, :] * padded_in_features + k[:, None] + weight_mask = (offs_n[None, :] < out_features) & (k[:, None] < in_features) + packed = tl.load( + packed_ptr + value_offsets // 2, + mask=weight_mask, + other=0, + ).to(tl.uint32) + shifts = (value_offsets & 1) * 4 + unsigned = (packed >> shifts) & 15 + signed = unsigned.to(tl.int32) - 8 + groups = k[:, None] // group_size + scales = tl.load( + scales_ptr + offs_n[None, :] * num_groups + groups, + mask=weight_mask, + other=0.0, + ).to(tl.float32) + weights = (signed.to(tl.float32) * scales).to(input_values.dtype) + accumulator += tl.dot(input_values, weights) + + if has_bias: + bias = tl.load( + bias_ptr + offs_n, + mask=offs_n < out_features, + other=0.0, + ).to(tl.float32) + accumulator += bias[None, :] + + tl.store( + output_ptr + offs_m[:, None] * out_features + offs_n[None, :], + accumulator, + mask=(offs_m[:, None] < rows) & (offs_n[None, :] < out_features), + ) + + @triton.jit def _decode_packed_w4_codes_kernel( packed_ptr, @@ -679,12 +857,8 @@ def _scale_int32_matmul_output_chunk_kernel( accumulator_offsets = offs_m[:, None] * chunk_out_features + offs_n[None, :] output_columns = output_col_start + offs_n output_offsets = offs_m[:, None] * out_features + output_columns[None, :] - accumulator = tl.load( - accumulator_ptr + accumulator_offsets, mask=mask, other=0 - ).to(tl.float32) - token_norms = tl.load( - token_norms_ptr + offs_m, mask=offs_m < rows, other=0.0 - ).to(tl.float32) + accumulator = tl.load(accumulator_ptr + accumulator_offsets, mask=mask, other=0).to(tl.float32) + token_norms = tl.load(token_norms_ptr + offs_m, mask=offs_m < rows, other=0.0).to(tl.float32) row_norms = tl.load( row_norms_ptr + output_columns, mask=output_columns < out_features, @@ -743,9 +917,7 @@ def _unpack_lowbit_kernel( bit_offsets = bit_starts % 8 low = tl.load(packed_ptr + byte_indices, mask=mask, other=0).to(tl.uint32) needs_high = bit_offsets + bits > 8 - high = tl.load(packed_ptr + byte_indices + 1, mask=mask & needs_high, other=0).to( - tl.uint32 - ) + high = tl.load(packed_ptr + byte_indices + 1, mask=mask & needs_high, other=0).to(tl.uint32) raw = low | (high << 8) values = (raw >> bit_offsets) & ((1 << bits) - 1) tl.store(values_ptr + offsets, values.to(tl.uint8), mask=mask) @@ -765,7 +937,7 @@ def _permute_sign_weight_kernel( mask = offsets < total rows = offsets // in_features cols = offsets % in_features - source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int64) + source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int32) signs = tl.load(signs_ptr + cols, mask=mask, other=1).to(tl.float32) values = tl.load( weight_ptr + rows * in_features + source_cols, @@ -893,6 +1065,15 @@ def _adaln_group_scales_kernel( tl.store(scales_ptr + rows * num_groups + groups, max_abs / 7.0) +@triton.jit +def _round_to_nearest_even(values): + lower = tl.floor(values) + lower_int = lower.to(tl.int32) + fraction = values - lower + round_up = (fraction > 0.5) | ((fraction == 0.5) & ((lower_int & 1) != 0)) + return lower_int + round_up.to(tl.int32) + + @triton.jit def _adaln_quantize_pack_int4_kernel( weight_ptr, @@ -927,14 +1108,7 @@ def _adaln_quantize_pack_int4_kernel( other=1.0, ).to(tl.float32) scaled0 = values0 / scales0 - rounded0 = tl.inline_asm_elementwise( - "cvt.rni.s32.f32 $0, $1;", - "=r,f", - [scaled0], - dtype=tl.int32, - is_pure=True, - pack=1, - ) + rounded0 = _round_to_nearest_even(scaled0) clamped0 = tl.minimum(tl.maximum(rounded0, -8), 7) + 8 rows1 = value_offsets1 // padded_in_features @@ -953,14 +1127,7 @@ def _adaln_quantize_pack_int4_kernel( other=1.0, ).to(tl.float32) scaled1 = values1 / scales1 - rounded1 = tl.inline_asm_elementwise( - "cvt.rni.s32.f32 $0, $1;", - "=r,f", - [scaled1], - dtype=tl.int32, - is_pure=True, - pack=1, - ) + rounded1 = _round_to_nearest_even(scaled1) clamped1 = tl.minimum(tl.maximum(rounded1, -8), 7) + 8 clamped1 = tl.where(valid1, clamped1, 0) @@ -991,9 +1158,7 @@ def _dequantize_adaln_weight_kernel( packed = tl.load(packed_ptr + byte_offsets, mask=mask, other=0).to(tl.uint32) unsigned = (packed >> shifts) & 15 signed = unsigned.to(tl.int32) - 8 - scales = tl.load(scales_ptr + rows * num_groups + groups, mask=mask, other=0.0).to( - tl.float32 - ) + scales = tl.load(scales_ptr + rows * num_groups + groups, mask=mask, other=0.0).to(tl.float32) tl.store(output_ptr + offsets, signed.to(tl.float32) * scales, mask=mask) @@ -1002,8 +1167,7 @@ def quantize_rotated_activations_with_triton( norms: torch.Tensor, codebook: LloydMaxCodebook, ) -> torch.Tensor: - if not rotated.is_cuda: - raise RuntimeError("triton_cuda backend requires CUDA tensors") + _require_triton_tensor(rotated, "activation quantization") triton, _ = _load_triton() rotated_contiguous = rotated.contiguous() flat = rotated_contiguous.reshape(-1) @@ -1035,8 +1199,7 @@ def quantize_activations_with_triton( eps: float, constant_tensors: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: - if not x.is_cuda: - raise RuntimeError("triton_cuda activation quantization requires CUDA tensors") + _require_triton_tensor(x, "activation quantization") if x.shape[-1] != rotation.dim: raise ValueError(f"expected last dimension {rotation.dim}, got {x.shape[-1]}") @@ -1056,7 +1219,7 @@ def quantize_activations_with_triton( signs_source = constants.get("signs", rotation.signs) centroids_source = constants.get("centroids", codebook.centroids) boundaries_source = constants.get("boundaries", codebook.boundaries) - permutation = permutation_source.to(device=x.device, dtype=torch.int64).contiguous() + permutation = permutation_source.to(device=x.device, dtype=torch.int32).contiguous() signs = signs_source.to(device=x.device, dtype=torch.int8).contiguous() centroids = centroids_source.to(device=x.device, dtype=torch.float32).contiguous() boundaries = boundaries_source.to(device=x.device, dtype=torch.float32).contiguous() @@ -1074,7 +1237,7 @@ def quantize_activations_with_triton( orbit_block_size = int(rotation.block_size) num_blocks = int(rotation.num_blocks) - if orbit_block_size <= 1024: + if orbit_block_size <= 4096: fwht_stages = orbit_block_size.bit_length() - 1 _fused_rpbh_quantize_activation_kernel[(rows * num_blocks,)]( input_contiguous, @@ -1099,9 +1262,7 @@ def quantize_activations_with_triton( element_block_size = 256 work = torch.empty(total, device=x.device, dtype=torch.float32) chunks_per_orbit_block = orbit_block_size // 1024 - _fused_rpbh_fwht_1024_chunk_kernel[ - (rows * num_blocks * chunks_per_orbit_block,) - ]( + _fused_rpbh_fwht_1024_chunk_kernel[(rows * num_blocks * chunks_per_orbit_block,)]( input_contiguous, norms, permutation, @@ -1141,9 +1302,7 @@ def quantize_activations_with_triton( block_size=element_block_size, ) - _quantize_activation_work_rescale_kernel[ - (triton.cdiv(total, element_block_size),) - ]( + _quantize_activation_work_rescale_kernel[(triton.cdiv(total, element_block_size),)]( work, norms, centroids, @@ -1167,8 +1326,7 @@ def quantize_activations_packed_w4_with_triton( constant_tensors: dict[str, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Quantize RPBH activations directly to packed W4 indices and token norms.""" - if not x.is_cuda: - raise RuntimeError("triton_cuda packed activation quantization requires CUDA tensors") + _require_triton_tensor(x, "packed activation quantization") if codebook.bits != 4: raise ValueError("packed W4 activation quantization requires a 4-bit codebook") if x.shape[-1] != rotation.dim: @@ -1185,20 +1343,22 @@ def quantize_activations_packed_w4_with_triton( packed = torch.empty(rows * (dim // 2), device=x.device, dtype=torch.uint8) norms = torch.empty(rows, device=x.device, dtype=torch.float32) if total == 0: - return packed.reshape(*original_shape[:-1], dim // 2), norms.reshape( - original_shape[:-1] - ) + return packed.reshape(*original_shape[:-1], dim // 2), norms.reshape(original_shape[:-1]) constants = {} if constant_tensors is None else constant_tensors - permutation = constants.get("permutation", rotation.permutation).to( - device=x.device, dtype=torch.int64 - ).contiguous() - signs = constants.get("signs", rotation.signs).to( - device=x.device, dtype=torch.int8 - ).contiguous() - boundaries = constants.get("boundaries", codebook.boundaries).to( - device=x.device, dtype=torch.float32 - ).contiguous() + permutation = ( + constants.get("permutation", rotation.permutation) + .to(device=x.device, dtype=torch.int32) + .contiguous() + ) + signs = ( + constants.get("signs", rotation.signs).to(device=x.device, dtype=torch.int8).contiguous() + ) + boundaries = ( + constants.get("boundaries", codebook.boundaries) + .to(device=x.device, dtype=torch.float32) + .contiguous() + ) norm_block_size = triton.next_power_of_2(dim) _row_norm_kernel[(rows,)]( @@ -1213,7 +1373,7 @@ def quantize_activations_packed_w4_with_triton( orbit_block_size = int(rotation.block_size) num_blocks = int(rotation.num_blocks) - if orbit_block_size <= 1024: + if orbit_block_size <= 4096: fwht_stages = orbit_block_size.bit_length() - 1 _fused_rpbh_quantize_activation_pack_w4_kernel[(rows * num_blocks,)]( input_contiguous, @@ -1232,16 +1392,12 @@ def quantize_activations_packed_w4_with_triton( inv_sqrt_block=float(rotation.normalization), num_warps=8 if orbit_block_size >= 512 else 4, ) - return packed.reshape(*original_shape[:-1], dim // 2), norms.reshape( - original_shape[:-1] - ) + return packed.reshape(*original_shape[:-1], dim // 2), norms.reshape(original_shape[:-1]) element_block_size = 256 work = torch.empty(total, device=x.device, dtype=torch.float32) chunks_per_orbit_block = orbit_block_size // 1024 - _fused_rpbh_fwht_1024_chunk_kernel[ - (rows * num_blocks * chunks_per_orbit_block,) - ]( + _fused_rpbh_fwht_1024_chunk_kernel[(rows * num_blocks * chunks_per_orbit_block,)]( input_contiguous, norms, permutation, @@ -1279,9 +1435,7 @@ def quantize_activations_packed_w4_with_triton( stage_width=stage_width, block_size=element_block_size, ) - _quantize_activation_work_pack_w4_kernel[ - (triton.cdiv(packed.numel(), element_block_size),) - ]( + _quantize_activation_work_pack_w4_kernel[(triton.cdiv(packed.numel(), element_block_size),)]( work, boundaries, packed, @@ -1290,9 +1444,7 @@ def quantize_activations_packed_w4_with_triton( inv_sqrt_block=float(rotation.normalization), block_size=element_block_size, ) - return packed.reshape(*original_shape[:-1], dim // 2), norms.reshape( - original_shape[:-1] - ) + return packed.reshape(*original_shape[:-1], dim // 2), norms.reshape(original_shape[:-1]) def quantize_adaln_weight_with_triton( @@ -1300,8 +1452,7 @@ def quantize_adaln_weight_with_triton( *, group_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: - if not weight.is_cuda: - raise RuntimeError("triton_cuda AdaLN RTN quantization requires CUDA tensors") + _require_triton_tensor(weight, "AdaLN RTN quantization") if weight.ndim != 2: raise ValueError("weight must be a rank-2 tensor") if group_size <= 0: @@ -1355,11 +1506,7 @@ def dequantize_adaln_weight_with_triton( group_size: int, device: torch.device | str = "cuda", ) -> torch.Tensor: - if not torch.cuda.is_available(): - raise RuntimeError("triton_cuda backend requires CUDA tensors") - target_device = torch.device(device) - if target_device.type != "cuda": - raise RuntimeError("triton_cuda AdaLN dequant requires a CUDA device") + target_device = _require_triton_device(device, "AdaLN dequantization") if group_size <= 0: raise ValueError("group_size must be positive") @@ -1398,12 +1545,8 @@ def dequantize_packed_weight_with_triton( in_features: int, device: torch.device | str = "cuda", ) -> torch.Tensor: - if not torch.cuda.is_available(): - raise RuntimeError("triton_cuda backend requires CUDA tensors") triton, _ = _load_triton() - target_device = torch.device(device) - if target_device.type != "cuda": - raise RuntimeError("triton_cuda weight dequant requires a CUDA device") + target_device = _require_triton_device(device, "weight dequantization") total = out_features * in_features packed = packed_weight_indices.to(device=target_device, dtype=torch.uint8).contiguous() norms = row_norms.to(device=target_device, dtype=torch.float32).contiguous() @@ -1427,6 +1570,39 @@ def dequantize_packed_weight_with_triton( return output.reshape(out_features, in_features) +_PACKED_MATMUL_LEGACY_TILE = (64, 64, 128, 4) + + +def _resolve_packed_matmul_tile( + rows: int, + bits: int, + block_m: int, + block_n: int, + block_k: int, + num_warps: int, +) -> tuple[int, int, int, int]: + """Map the historical default tile to a measured shape-aware choice. + + Callers that pass any non-default tile keep full manual control. The + historical default (64, 64, 128, 4) measured 1.9-2.4x slower than + (128, 64, 64, 4) for large-row W3/W4 packed matmuls on RTX 4090-class + hardware, so it is treated as "auto". + """ + if (block_m, block_n, block_k, num_warps) != _PACKED_MATMUL_LEGACY_TILE: + return block_m, block_n, block_k, num_warps + if rows >= 128 and bits != 4: + return 128, 128, 64, 8 + return 128, 64, 64, 4 + + +def _triton_out_of_resources_error() -> type[Exception] | None: + try: + from triton.runtime.errors import OutOfResources + except Exception: # pragma: no cover - triton version dependent + return None + return OutOfResources + + def matmul_packed_weight_with_triton( x: torch.Tensor, packed_weight_indices: torch.Tensor, @@ -1446,10 +1622,7 @@ def matmul_packed_weight_with_triton( raise ValueError("bits must be in [1, 8]") if block_m <= 0 or block_n <= 0 or block_k <= 0 or num_warps <= 0: raise ValueError("packed matmul tile sizes and num_warps must be positive") - if not torch.cuda.is_available(): - raise RuntimeError("triton_cuda backend requires CUDA tensors") - if not x.is_cuda: - raise RuntimeError("triton_cuda packed matmul requires CUDA input tensors") + _require_triton_tensor(x, "packed matmul") if x.shape[-1] != in_features: raise ValueError(f"expected input last dimension {in_features}, got {x.shape[-1]}") @@ -1472,45 +1645,182 @@ def matmul_packed_weight_with_triton( bias_tensor = bias.to(device=x.device).contiguous() has_bias = True + block_m, block_n, block_k, num_warps = _resolve_packed_matmul_tile( + rows, bits, block_m, block_n, block_k, num_warps + ) effective_block_m = min(block_m, 16 if rows < 16 else triton.next_power_of_2(rows)) - grid = (triton.cdiv(rows, effective_block_m), triton.cdiv(out_features, block_n)) - if bits == 4 and in_features % 2 == 0 and block_k >= 2: - block_k_bytes = max(1, block_k // 2) - _matmul_packed_weight_w4_kernel[grid]( - input_contiguous, - packed, - norms, - centroids, - bias_tensor, - output, - rows=rows, - out_features=out_features, - in_features=in_features, - has_bias=has_bias, - block_m=effective_block_m, - block_n=block_n, - block_k_bytes=block_k_bytes, - num_warps=num_warps, + out_of_resources = _triton_out_of_resources_error() + attempts: list[tuple[int, int]] = [(block_n, block_k)] + for candidate in ( + (block_n, max(32, block_k // 2)), + (max(64, block_n // 2), max(32, block_k // 2)), + ): + if candidate not in attempts: + attempts.append(candidate) + last_error: Exception | None = None + for attempt_n, attempt_k in attempts: + grid = (triton.cdiv(rows, effective_block_m), triton.cdiv(out_features, attempt_n)) + try: + if bits == 4 and in_features % 2 == 0 and attempt_k >= 2: + block_k_bytes = max(1, attempt_k // 2) + _matmul_packed_weight_w4_kernel[grid]( + input_contiguous, + packed, + norms, + centroids, + bias_tensor, + output, + rows=rows, + out_features=out_features, + in_features=in_features, + has_bias=has_bias, + block_m=effective_block_m, + block_n=attempt_n, + block_k_bytes=block_k_bytes, + num_warps=num_warps, + ) + else: + _matmul_packed_weight_kernel[grid]( + input_contiguous, + packed, + norms, + centroids, + bias_tensor, + output, + rows=rows, + out_features=out_features, + in_features=in_features, + bits=bits, + has_bias=has_bias, + block_m=effective_block_m, + block_n=attempt_n, + block_k=attempt_k, + num_warps=num_warps, + ) + return output.reshape(*original_shape[:-1], out_features) + except Exception as exc: # noqa: BLE001 - re-raised unless a tile overflow + if out_of_resources is None or not isinstance(exc, out_of_resources): + raise + last_error = exc + raise RuntimeError( + "packed matmul exceeded the device shared-memory limit for tile " + f"attempts {attempts}; pass smaller block_n/block_k values" + ) from last_error + + +def matmul_packed_adaln_int4_with_triton( + x: torch.Tensor, + packed_weight: torch.Tensor, + scales: torch.Tensor, + *, + out_features: int, + in_features: int, + group_size: int, + bias: torch.Tensor | None = None, + block_m: int = 64, + block_n: int = 128, + block_k: int = 64, + num_warps: int = 4, +) -> torch.Tensor: + _require_triton_tensor(x, "packed AdaLN matmul") + if x.shape[-1] != in_features: + raise ValueError(f"expected input last dimension {in_features}, got {x.shape[-1]}") + if group_size <= 0: + raise ValueError("group_size must be positive") + tile_sizes = (block_m, block_n, block_k) + if any(value <= 0 or value & (value - 1) for value in tile_sizes): + raise ValueError("packed AdaLN tile sizes must be positive powers of two") + if num_warps <= 0: + raise ValueError("num_warps must be positive") + + triton, _ = _load_triton() + original_shape = x.shape + input_contiguous = x.to(dtype=torch.bfloat16).contiguous().reshape(-1, in_features) + rows = input_contiguous.shape[0] + output = torch.empty( + (rows, out_features), + device=x.device, + dtype=torch.bfloat16, + ) + if rows == 0 or out_features == 0: + return output.reshape(*original_shape[:-1], out_features) + + num_groups = (in_features + group_size - 1) // group_size + padded_in_features = num_groups * group_size + expected_packed_bytes = (out_features * padded_in_features + 1) // 2 + if packed_weight.numel() != expected_packed_bytes: + raise ValueError( + f"expected {expected_packed_bytes} packed AdaLN bytes, got {packed_weight.numel()}" ) + if tuple(scales.shape) != (out_features, num_groups): + raise ValueError( + f"expected AdaLN scales shape {(out_features, num_groups)}, got {tuple(scales.shape)}" + ) + + packed = packed_weight.to(device=x.device, dtype=torch.uint8).contiguous().flatten() + scale_values = scales.to(device=x.device).contiguous() + if bias is None: + bias_values = output + has_bias = False else: - _matmul_packed_weight_kernel[grid]( - input_contiguous, - packed, - norms, - centroids, - bias_tensor, - output, - rows=rows, - out_features=out_features, - in_features=in_features, - bits=bits, - has_bias=has_bias, - block_m=effective_block_m, - block_n=block_n, - block_k=block_k, - num_warps=num_warps, + bias_values = bias.to( + device=x.device, + dtype=torch.bfloat16, + ).contiguous() + has_bias = True + + effective_block_m = min(block_m, 16 if rows < 16 else triton.next_power_of_2(rows)) + out_of_resources = _triton_out_of_resources_error() + attempts: list[tuple[int, int]] = [(block_n, block_k)] + for candidate in ( + (block_n, max(16, block_k // 2)), + (max(64, block_n // 2), max(16, block_k // 2)), + ): + if candidate not in attempts: + attempts.append(candidate) + last_error: Exception | None = None + for attempt_n, attempt_k in attempts: + grid = ( + triton.cdiv(rows, effective_block_m), + triton.cdiv(out_features, attempt_n), ) - return output.reshape(*original_shape[:-1], out_features) + # When every k-tile stays inside one scale group, the group-aligned + # kernel loads scales as a row vector; otherwise keep the generic + # per-element gather kernel. + use_group_kernel = attempt_k <= group_size and group_size % attempt_k == 0 + kernel = ( + _matmul_packed_adaln_int4_group_kernel + if use_group_kernel + else _matmul_packed_adaln_int4_kernel + ) + try: + kernel[grid]( + input_contiguous, + packed, + scale_values, + bias_values, + output, + rows=rows, + out_features=out_features, + in_features=in_features, + padded_in_features=padded_in_features, + num_groups=num_groups, + group_size=group_size, + has_bias=has_bias, + block_m=effective_block_m, + block_n=attempt_n, + block_k=attempt_k, + num_warps=num_warps, + ) + return output.reshape(*original_shape[:-1], out_features) + except Exception as exc: # noqa: BLE001 - re-raised unless a tile overflow + if out_of_resources is None or not isinstance(exc, out_of_resources): + raise + last_error = exc + raise RuntimeError( + "packed AdaLN matmul exceeded the device shared-memory limit for tile " + f"attempts {attempts}; pass smaller block_n/block_k values" + ) from last_error def fit_int8_centroid_surrogate(centroids: torch.Tensor) -> tuple[torch.Tensor, float]: @@ -1529,18 +1839,612 @@ def fit_int8_centroid_surrogate(centroids: torch.Tensor) -> tuple[torch.Tensor, return result[0].clone(), result[1] base_scale = max_abs / 127.0 candidate_scales = base_scale * torch.linspace(0.5, 2.0, 3001, dtype=torch.float64) - candidate_codes = torch.round(values[None, :] / candidate_scales[:, None]).clamp( - -127, 127 - ) - errors = ( - candidate_codes * candidate_scales[:, None] - values[None, :] - ).square().mean(dim=1) + candidate_codes = torch.round(values[None, :] / candidate_scales[:, None]).clamp(-127, 127) + errors = (candidate_codes * candidate_scales[:, None] - values[None, :]).square().mean(dim=1) best = int(errors.argmin()) result = (candidate_codes[best].to(torch.int8), float(candidate_scales[best])) _INT8_CENTROID_SURROGATE_CACHE[cache_key] = result return result[0].clone(), result[1] +@triton.jit +def _matmul_packed_w4a4_fused_int8_kernel( + packed_act_ptr, + packed_weight_ptr, + token_norms_ptr, + row_norms_ptr, + act_codes_ptr, + weight_codes_ptr, + bias_ptr, + output_ptr, + rows: tl.constexpr, + out_features: tl.constexpr, + in_features: tl.constexpr, + combined_scale: tl.constexpr, + has_bias: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_k: tl.constexpr, +): + """W4A4 GEMM that reads packed nibbles directly: decode both operands to + INT8 surrogate codes in-register and accumulate with tensor-core int8 dot, + applying the norm/scale epilogue in the same kernel.""" + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * block_m + tl.arange(0, block_m) + offs_n = pid_n * block_n + tl.arange(0, block_n) + packed_k: tl.constexpr = in_features // 2 + packed_block_k: tl.constexpr = block_k // 2 + offs_kb = tl.arange(0, packed_block_k) + m_mask = offs_m < rows + n_mask = offs_n < out_features + accumulator = tl.zeros((block_m, block_n), dtype=tl.int32) + + for kb_start in range(0, packed_k, packed_block_k): + kb = kb_start + offs_kb + kb_mask = kb < packed_k + act_bytes = tl.load( + packed_act_ptr + offs_m[:, None] * packed_k + kb[None, :], + mask=m_mask[:, None] & kb_mask[None, :], + other=0, + ).to(tl.int32) + act_low = tl.load(act_codes_ptr + (act_bytes & 15)) + act_high = tl.load(act_codes_ptr + ((act_bytes >> 4) & 15)) + act_tile = tl.reshape( + tl.join(act_low, act_high), (block_m, block_k) + ).to(tl.int8) + + weight_bytes = tl.load( + packed_weight_ptr + offs_n[:, None] * packed_k + kb[None, :], + mask=n_mask[:, None] & kb_mask[None, :], + other=0, + ).to(tl.int32) + weight_low = tl.load(weight_codes_ptr + (weight_bytes & 15)) + weight_high = tl.load(weight_codes_ptr + ((weight_bytes >> 4) & 15)) + weight_tile = tl.reshape( + tl.join(weight_low, weight_high), (block_n, block_k) + ).to(tl.int8) + + accumulator = tl.dot( + act_tile, tl.trans(weight_tile), accumulator, out_dtype=tl.int32 + ) + + result = accumulator.to(tl.float32) + token_norms = tl.load(token_norms_ptr + offs_m, mask=m_mask, other=0.0).to(tl.float32) + row_norms = tl.load(row_norms_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + result = result * token_norms[:, None] * row_norms[None, :] * combined_scale + if has_bias: + bias = tl.load(bias_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + result += bias[None, :] + tl.store( + output_ptr + offs_m[:, None] * out_features + offs_n[None, :], + result, + mask=m_mask[:, None] & n_mask[None, :], + ) + + +@triton.jit +def _matmul_int8act_packed_lowbit_fused_kernel( + act_ptr, + packed_weight_ptr, + token_norms_ptr, + row_norms_ptr, + weight_codes_ptr, + bias_ptr, + output_ptr, + rows: tl.constexpr, + out_features: tl.constexpr, + in_features: tl.constexpr, + combined_scale: tl.constexpr, + has_bias: tl.constexpr, + weight_bits: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_k: tl.constexpr, +): + """Fused low-bit GEMM over INT8-surrogate activations: the activation + tile loads directly (no nibble unpack) and the packed W2/W4 weights + decode to INT8 surrogates in-register.""" + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * block_m + tl.arange(0, block_m) + offs_n = pid_n * block_n + tl.arange(0, block_n) + offs_k = tl.arange(0, block_k) + values_per_byte: tl.constexpr = 8 // weight_bits + packed_k: tl.constexpr = in_features // values_per_byte + packed_block_k: tl.constexpr = block_k // values_per_byte + offs_kb = tl.arange(0, packed_block_k) + m_mask = offs_m < rows + n_mask = offs_n < out_features + accumulator = tl.zeros((block_m, block_n), dtype=tl.int32) + + for k_start in range(0, in_features, block_k): + k = k_start + offs_k + k_mask = k < in_features + act_tile = tl.load( + act_ptr + offs_m[:, None] * in_features + k[None, :], + mask=m_mask[:, None] & k_mask[None, :], + other=0, + ) + + kb = (k_start // values_per_byte) + offs_kb + kb_mask = kb < packed_k + weight_bytes = tl.load( + packed_weight_ptr + offs_n[:, None] * packed_k + kb[None, :], + mask=n_mask[:, None] & kb_mask[None, :], + other=0, + ).to(tl.int32) + if weight_bits == 4: + weight_low = tl.load(weight_codes_ptr + (weight_bytes & 15)) + weight_high = tl.load(weight_codes_ptr + ((weight_bytes >> 4) & 15)) + weight_tile = tl.reshape( + tl.join(weight_low, weight_high), (block_n, block_k) + ).to(tl.int8) + else: + w0 = tl.load(weight_codes_ptr + (weight_bytes & 3)) + w1 = tl.load(weight_codes_ptr + ((weight_bytes >> 2) & 3)) + w2 = tl.load(weight_codes_ptr + ((weight_bytes >> 4) & 3)) + w3 = tl.load(weight_codes_ptr + ((weight_bytes >> 6) & 3)) + weight_tile = tl.reshape( + tl.join(tl.join(w0, w2), tl.join(w1, w3)), (block_n, block_k) + ).to(tl.int8) + + accumulator = tl.dot( + act_tile, tl.trans(weight_tile), accumulator, out_dtype=tl.int32 + ) + + result = accumulator.to(tl.float32) + token_norms = tl.load(token_norms_ptr + offs_m, mask=m_mask, other=0.0).to(tl.float32) + row_norms = tl.load(row_norms_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + result = result * token_norms[:, None] * row_norms[None, :] * combined_scale + if has_bias: + bias = tl.load(bias_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + result += bias[None, :] + tl.store( + output_ptr + offs_m[:, None] * out_features + offs_n[None, :], + result, + mask=m_mask[:, None] & n_mask[None, :], + ) + + +@triton.jit +def _matmul_packed_w2a4_fused_int8_kernel( + packed_act_ptr, + packed_weight_ptr, + token_norms_ptr, + row_norms_ptr, + act_codes_ptr, + weight_codes_ptr, + bias_ptr, + output_ptr, + rows: tl.constexpr, + out_features: tl.constexpr, + in_features: tl.constexpr, + combined_scale: tl.constexpr, + has_bias: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_k: tl.constexpr, +): + """W2A4 GEMM over packed operands: 4-bit activation nibbles and 2-bit + weight fields decode to INT8 surrogates in-register.""" + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * block_m + tl.arange(0, block_m) + offs_n = pid_n * block_n + tl.arange(0, block_n) + packed_act_k: tl.constexpr = in_features // 2 + act_block_k: tl.constexpr = block_k // 2 + offs_akb = tl.arange(0, act_block_k) + packed_w_k: tl.constexpr = in_features // 4 + w_block_k: tl.constexpr = block_k // 4 + offs_wkb = tl.arange(0, w_block_k) + m_mask = offs_m < rows + n_mask = offs_n < out_features + accumulator = tl.zeros((block_m, block_n), dtype=tl.int32) + + for k_start in range(0, in_features, block_k): + akb = (k_start // 2) + offs_akb + akb_mask = akb < packed_act_k + act_bytes = tl.load( + packed_act_ptr + offs_m[:, None] * packed_act_k + akb[None, :], + mask=m_mask[:, None] & akb_mask[None, :], + other=0, + ).to(tl.int32) + act_low = tl.load(act_codes_ptr + (act_bytes & 15)) + act_high = tl.load(act_codes_ptr + ((act_bytes >> 4) & 15)) + act_tile = tl.reshape( + tl.join(act_low, act_high), (block_m, block_k) + ).to(tl.int8) + + wkb = (k_start // 4) + offs_wkb + wkb_mask = wkb < packed_w_k + weight_bytes = tl.load( + packed_weight_ptr + offs_n[:, None] * packed_w_k + wkb[None, :], + mask=n_mask[:, None] & wkb_mask[None, :], + other=0, + ).to(tl.int32) + w0 = tl.load(weight_codes_ptr + (weight_bytes & 3)) + w1 = tl.load(weight_codes_ptr + ((weight_bytes >> 2) & 3)) + w2 = tl.load(weight_codes_ptr + ((weight_bytes >> 4) & 3)) + w3 = tl.load(weight_codes_ptr + ((weight_bytes >> 6) & 3)) + weight_tile = tl.reshape( + tl.join(tl.join(w0, w2), tl.join(w1, w3)), (block_n, block_k) + ).to(tl.int8) + + accumulator = tl.dot( + act_tile, tl.trans(weight_tile), accumulator, out_dtype=tl.int32 + ) + + result = accumulator.to(tl.float32) + token_norms = tl.load(token_norms_ptr + offs_m, mask=m_mask, other=0.0).to(tl.float32) + row_norms = tl.load(row_norms_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + result = result * token_norms[:, None] * row_norms[None, :] * combined_scale + if has_bias: + bias = tl.load(bias_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + result += bias[None, :] + tl.store( + output_ptr + offs_m[:, None] * out_features + offs_n[None, :], + result, + mask=m_mask[:, None] & n_mask[None, :], + ) + + +def _fused_w4a4_tile_config(rows: int) -> tuple[int, int, int, int]: + """Measured tile table (A40/sm_86 and RTX 4090/sm_89 grids, 2026-07): + large batches want a taller M tile and a narrower K tile.""" + if rows >= 512: + return 128, 128, 64, 4 + return 64, 128, 128, 4 + + +def matmul_packed_w4a4_fused_with_triton( + packed_activations: torch.Tensor, + packed_weight_indices: torch.Tensor, + token_norms: torch.Tensor, + row_norms: torch.Tensor, + activation_codes: torch.Tensor, + weight_codes: torch.Tensor, + *, + activation_scale: float, + weight_scale: float, + out_features: int, + in_features: int, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.bfloat16, + block_m: int | None = None, + block_n: int | None = None, + block_k: int | None = None, + num_warps: int | None = None, +) -> torch.Tensor: + """Single-kernel W4A4 matmul over packed operands (no INT8 materialization).""" + if not packed_activations.is_cuda: + raise RuntimeError("fused packed W4A4 matmul requires CUDA tensors") + if output_dtype not in {torch.float16, torch.bfloat16}: + raise ValueError("output_dtype must be float16 or bfloat16") + if in_features <= 0 or in_features % 64 != 0: + raise ValueError("in_features must be positive and divisible by 64") + if out_features <= 0 or out_features % 16 != 0: + raise ValueError("out_features must be positive and divisible by 16") + + triton, _ = _load_triton() + original_shape = packed_activations.shape + packed_k = in_features // 2 + if original_shape[-1] != packed_k: + raise ValueError( + f"expected packed activation last dimension {packed_k}, got {original_shape[-1]}" + ) + activations = packed_activations.to(dtype=torch.uint8).contiguous().reshape(-1, packed_k) + rows = activations.shape[0] + device = packed_activations.device + if packed_weight_indices.numel() != out_features * packed_k: + raise ValueError("packed weights have an unexpected size") + output = torch.empty((rows, out_features), device=device, dtype=output_dtype) + if rows == 0: + return output.reshape(*original_shape[:-1], out_features) + + weights = packed_weight_indices.to(device=device, dtype=torch.uint8).contiguous().flatten() + activation_code_values = activation_codes.to(device=device, dtype=torch.int8).contiguous() + weight_code_values = weight_codes.to(device=device, dtype=torch.int8).contiguous() + token_norm_values = token_norms.to(device=device, dtype=torch.float32).contiguous().reshape(-1) + row_norm_values = row_norms.to(device=device, dtype=torch.float32).contiguous().reshape(-1) + if bias is None: + bias_values = output + has_bias = False + else: + bias_values = bias.to(device=device, dtype=output_dtype).contiguous() + has_bias = True + combined_scale = float(activation_scale * weight_scale) + + default_m, default_n, default_k, default_warps = _fused_w4a4_tile_config(rows) + block_m = default_m if block_m is None else block_m + block_n = default_n if block_n is None else block_n + block_k = default_k if block_k is None else block_k + num_warps = default_warps if num_warps is None else num_warps + + effective_block_m = min(block_m, 16 if rows < 16 else triton.next_power_of_2(rows)) + out_of_resources = _triton_out_of_resources_error() + attempts: list[tuple[int, int]] = [(block_n, block_k)] + for candidate in ( + (block_n, max(32, block_k // 2)), + (max(64, block_n // 2), max(32, block_k // 2)), + ): + if candidate not in attempts: + attempts.append(candidate) + last_error: Exception | None = None + for attempt_n, attempt_k in attempts: + grid = ( + triton.cdiv(rows, effective_block_m), + triton.cdiv(out_features, attempt_n), + ) + try: + _matmul_packed_w4a4_fused_int8_kernel[grid]( + activations, + weights, + token_norm_values, + row_norm_values, + activation_code_values, + weight_code_values, + bias_values, + output, + rows=rows, + out_features=out_features, + in_features=in_features, + combined_scale=combined_scale, + has_bias=has_bias, + block_m=effective_block_m, + block_n=attempt_n, + block_k=attempt_k, + num_warps=num_warps, + ) + return output.reshape(*original_shape[:-1], out_features) + except Exception as exc: # noqa: BLE001 - re-raised unless a tile overflow + if out_of_resources is None or not isinstance(exc, out_of_resources): + raise + last_error = exc + raise RuntimeError( + "fused packed W4A4 matmul exceeded the device shared-memory limit for " + f"tile attempts {attempts}" + ) from last_error + + +def matmul_int8_activations_packed_lowbit_fused_with_triton( + int8_activations: torch.Tensor, + packed_weight_indices: torch.Tensor, + token_norms: torch.Tensor, + row_norms: torch.Tensor, + weight_codes: torch.Tensor, + *, + weight_bits: int, + activation_scale: float, + weight_scale: float, + out_features: int, + in_features: int, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Fused low-bit GEMM over INT8-surrogate activations (W2/W4 weights). + + Measured on RTX 4090: skipping the packed-activation round trip is + 13-20% faster than the packed-activation fused kernel at 64-2048 rows. + """ + if not int8_activations.is_cuda: + raise RuntimeError("fused low-bit matmul requires CUDA tensors") + if weight_bits not in (2, 4): + raise ValueError("fused low-bit matmul supports weight_bits 2 or 4") + if output_dtype not in {torch.float16, torch.bfloat16}: + raise ValueError("output_dtype must be float16 or bfloat16") + if in_features <= 0 or in_features % 64 != 0: + raise ValueError("in_features must be positive and divisible by 64") + if out_features <= 0 or out_features % 16 != 0: + raise ValueError("out_features must be positive and divisible by 16") + + triton, _ = _load_triton() + original_shape = int8_activations.shape + if original_shape[-1] != in_features: + raise ValueError( + f"expected activation last dimension {in_features}, got {original_shape[-1]}" + ) + activations = int8_activations.to(dtype=torch.int8).contiguous().reshape(-1, in_features) + rows = activations.shape[0] + device = int8_activations.device + values_per_byte = 8 // weight_bits + packed_k = in_features // values_per_byte + if packed_weight_indices.numel() != out_features * packed_k: + raise ValueError("packed weights have an unexpected size") + output = torch.empty((rows, out_features), device=device, dtype=output_dtype) + if rows == 0: + return output.reshape(*original_shape[:-1], out_features) + + weights = packed_weight_indices.to(device=device, dtype=torch.uint8).contiguous().flatten() + weight_code_values = weight_codes.to(device=device, dtype=torch.int8).contiguous() + token_norm_values = token_norms.to(device=device, dtype=torch.float32).contiguous().reshape(-1) + row_norm_values = row_norms.to(device=device, dtype=torch.float32).contiguous().reshape(-1) + if bias is None: + bias_values = output + has_bias = False + else: + bias_values = bias.to(device=device, dtype=output_dtype).contiguous() + has_bias = True + combined_scale = float(activation_scale * weight_scale) + + block_m, block_n, block_k, num_warps = _fused_w4a4_tile_config(rows) + effective_block_m = min(block_m, 16 if rows < 16 else triton.next_power_of_2(rows)) + out_of_resources = _triton_out_of_resources_error() + attempts: list[tuple[int, int]] = [(block_n, block_k)] + for candidate in ( + (block_n, max(32, block_k // 2)), + (max(64, block_n // 2), max(32, block_k // 2)), + ): + if candidate not in attempts: + attempts.append(candidate) + last_error: Exception | None = None + for attempt_n, attempt_k in attempts: + grid = ( + triton.cdiv(rows, effective_block_m), + triton.cdiv(out_features, attempt_n), + ) + try: + _matmul_int8act_packed_lowbit_fused_kernel[grid]( + activations, + weights, + token_norm_values, + row_norm_values, + weight_code_values, + bias_values, + output, + rows=rows, + out_features=out_features, + in_features=in_features, + combined_scale=combined_scale, + has_bias=has_bias, + weight_bits=weight_bits, + block_m=effective_block_m, + block_n=attempt_n, + block_k=attempt_k, + num_warps=num_warps, + ) + return output.reshape(*original_shape[:-1], out_features) + except Exception as exc: # noqa: BLE001 - re-raised unless a tile overflow + if out_of_resources is None or not isinstance(exc, out_of_resources): + raise + last_error = exc + raise RuntimeError( + "fused low-bit matmul exceeded the device shared-memory limit for " + f"tile attempts {attempts}" + ) from last_error + + +def matmul_packed_w2a4_fused_with_triton( + packed_activations: torch.Tensor, + packed_weight_indices: torch.Tensor, + token_norms: torch.Tensor, + row_norms: torch.Tensor, + activation_codes: torch.Tensor, + weight_codes: torch.Tensor, + *, + activation_scale: float, + weight_scale: float, + out_features: int, + in_features: int, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Single-kernel W2A4 matmul over packed operands (no INT8 materialization).""" + if not packed_activations.is_cuda: + raise RuntimeError("fused packed W2A4 matmul requires CUDA tensors") + if output_dtype not in {torch.float16, torch.bfloat16}: + raise ValueError("output_dtype must be float16 or bfloat16") + if in_features <= 0 or in_features % 64 != 0: + raise ValueError("in_features must be positive and divisible by 64") + if out_features <= 0 or out_features % 16 != 0: + raise ValueError("out_features must be positive and divisible by 16") + + triton, _ = _load_triton() + original_shape = packed_activations.shape + packed_act_k = in_features // 2 + if original_shape[-1] != packed_act_k: + raise ValueError( + f"expected packed activation last dimension {packed_act_k}, got {original_shape[-1]}" + ) + activations = packed_activations.to(dtype=torch.uint8).contiguous().reshape(-1, packed_act_k) + rows = activations.shape[0] + device = packed_activations.device + packed_w_k = in_features // 4 + if packed_weight_indices.numel() != out_features * packed_w_k: + raise ValueError("packed weights have an unexpected size") + output = torch.empty((rows, out_features), device=device, dtype=output_dtype) + if rows == 0: + return output.reshape(*original_shape[:-1], out_features) + + weights = packed_weight_indices.to(device=device, dtype=torch.uint8).contiguous().flatten() + activation_code_values = activation_codes.to(device=device, dtype=torch.int8).contiguous() + weight_code_values = weight_codes.to(device=device, dtype=torch.int8).contiguous() + token_norm_values = token_norms.to(device=device, dtype=torch.float32).contiguous().reshape(-1) + row_norm_values = row_norms.to(device=device, dtype=torch.float32).contiguous().reshape(-1) + if bias is None: + bias_values = output + has_bias = False + else: + bias_values = bias.to(device=device, dtype=output_dtype).contiguous() + has_bias = True + combined_scale = float(activation_scale * weight_scale) + + block_m, block_n, block_k, num_warps = _fused_w4a4_tile_config(rows) + effective_block_m = min(block_m, 16 if rows < 16 else triton.next_power_of_2(rows)) + out_of_resources = _triton_out_of_resources_error() + attempts: list[tuple[int, int]] = [(block_n, block_k)] + for candidate in ( + (block_n, max(32, block_k // 2)), + (max(64, block_n // 2), max(32, block_k // 2)), + ): + if candidate not in attempts: + attempts.append(candidate) + last_error: Exception | None = None + for attempt_n, attempt_k in attempts: + grid = ( + triton.cdiv(rows, effective_block_m), + triton.cdiv(out_features, attempt_n), + ) + try: + _matmul_packed_w2a4_fused_int8_kernel[grid]( + activations, + weights, + token_norm_values, + row_norm_values, + activation_code_values, + weight_code_values, + bias_values, + output, + rows=rows, + out_features=out_features, + in_features=in_features, + combined_scale=combined_scale, + has_bias=has_bias, + block_m=effective_block_m, + block_n=attempt_n, + block_k=attempt_k, + num_warps=num_warps, + ) + return output.reshape(*original_shape[:-1], out_features) + except Exception as exc: # noqa: BLE001 - re-raised unless a tile overflow + if out_of_resources is None or not isinstance(exc, out_of_resources): + raise + last_error = exc + raise RuntimeError( + "fused packed W2A4 matmul exceeded the device shared-memory limit for " + f"tile attempts {attempts}" + ) from last_error + + +def decode_packed_w4_weight_to_int8( + packed_weight_indices: torch.Tensor, + weight_codes: torch.Tensor, + *, + out_features: int, + in_features: int, +) -> torch.Tensor: + """Materialize the INT8 surrogate weight once (for the opt-in cache).""" + if not packed_weight_indices.is_cuda: + raise RuntimeError("packed W4 weight decode requires CUDA tensors") + triton, _ = _load_triton() + device = packed_weight_indices.device + packed = packed_weight_indices.to(dtype=torch.uint8).contiguous().flatten() + codes = weight_codes.to(device=device, dtype=torch.int8).contiguous() + total = out_features * in_features + decoded = torch.empty((out_features, in_features), device=device, dtype=torch.int8) + if total == 0: + return decoded + decode_block_size = 256 + _decode_packed_w4_codes_kernel[(triton.cdiv(total, decode_block_size),)]( + packed, + codes, + decoded, + total=total, + block_size=decode_block_size, + ) + return decoded + + def matmul_packed_w4a4_with_int_mm( packed_activations: torch.Tensor, packed_weight_indices: torch.Tensor, @@ -1559,8 +2463,13 @@ def matmul_packed_w4a4_with_int_mm( epilogue_block_n: int = 256, chunk_out_features: int | None = None, activations_are_int8: bool = False, + decoded_weight: torch.Tensor | None = None, ) -> torch.Tensor: - """Decode active W4/A4 tiles to INT8 and use PyTorch's CUTLASS GEMM.""" + """Decode active W4/A4 tiles to INT8 and use PyTorch's CUTLASS GEMM. + + ``decoded_weight`` optionally supplies a persistent INT8 surrogate weight + (``decode_packed_w4_weight_to_int8``) so the per-forward decode is skipped. + """ if not packed_activations.is_cuda: raise RuntimeError("packed W4A4 int_mm requires CUDA tensors") if output_dtype not in {torch.float16, torch.bfloat16}: @@ -1571,9 +2480,7 @@ def matmul_packed_w4a4_with_int_mm( raise ValueError("out_features must be positive and divisible by 16") if min(epilogue_block_m, epilogue_block_n) <= 0: raise ValueError("epilogue tile sizes must be positive") - if chunk_out_features is not None and ( - chunk_out_features <= 0 or chunk_out_features % 16 != 0 - ): + if chunk_out_features is not None and (chunk_out_features <= 0 or chunk_out_features % 16 != 0): raise ValueError("chunk_out_features must be positive and divisible by 16") triton, _ = _load_triton() @@ -1584,8 +2491,7 @@ def matmul_packed_w4a4_with_int_mm( raise ValueError("INT8 activations must have dtype int8") if original_shape[-1] != in_features: raise ValueError( - f"expected INT8 activation last dimension {in_features}, " - f"got {original_shape[-1]}" + f"expected INT8 activation last dimension {in_features}, got {original_shape[-1]}" ) activation_values = packed_activations.contiguous().reshape(-1, in_features) activations = None @@ -1593,12 +2499,9 @@ def matmul_packed_w4a4_with_int_mm( else: if original_shape[-1] != packed_k: raise ValueError( - f"expected packed activation last dimension {packed_k}, " - f"got {original_shape[-1]}" + f"expected packed activation last dimension {packed_k}, got {original_shape[-1]}" ) - activations = ( - packed_activations.to(dtype=torch.uint8).contiguous().reshape(-1, packed_k) - ) + activations = packed_activations.to(dtype=torch.uint8).contiguous().reshape(-1, packed_k) rows = activations.shape[0] activation_values = torch.empty( (rows, in_features), device=activations.device, dtype=torch.int8 @@ -1617,24 +2520,14 @@ def matmul_packed_w4a4_with_int_mm( dtype=output_dtype, ) - weights = packed_weight_indices.to( - device=device, dtype=torch.uint8 - ).contiguous().flatten() - activation_code_values = activation_codes.to( - device=device, dtype=torch.int8 - ).contiguous() - weight_code_values = weight_codes.to( - device=device, dtype=torch.int8 - ).contiguous() + weights = packed_weight_indices.to(device=device, dtype=torch.uint8).contiguous().flatten() + activation_code_values = activation_codes.to(device=device, dtype=torch.int8).contiguous() + weight_code_values = weight_codes.to(device=device, dtype=torch.int8).contiguous() if activation_code_values.numel() != 16 or weight_code_values.numel() != 16: raise ValueError("packed W4A4 int_mm requires 16 activation and weight codes") - token_norm_values = token_norms.to( - device=device, dtype=torch.float32 - ).contiguous().reshape(-1) - row_norm_values = row_norms.to( - device=device, dtype=torch.float32 - ).contiguous().reshape(-1) + token_norm_values = token_norms.to(device=device, dtype=torch.float32).contiguous().reshape(-1) + row_norm_values = row_norms.to(device=device, dtype=torch.float32).contiguous().reshape(-1) if token_norm_values.numel() != rows: raise ValueError("token_norms must contain one value per activation row") if row_norm_values.numel() != out_features: @@ -1643,9 +2536,7 @@ def matmul_packed_w4a4_with_int_mm( decode_block_size = 256 if not activations_are_int8: activation_total = rows * in_features - _decode_packed_w4_codes_kernel[ - (triton.cdiv(activation_total, decode_block_size),) - ]( + _decode_packed_w4_codes_kernel[(triton.cdiv(activation_total, decode_block_size),)]( activations, activation_code_values, activation_values, @@ -1653,9 +2544,21 @@ def matmul_packed_w4a4_with_int_mm( block_size=decode_block_size, ) - output = torch.empty( - (rows, out_features), device=device, dtype=output_dtype - ) + # cuBLASLt INT8 GEMM rejects some row counts (every m <= 16 and, on some + # torch/CUDA stacks, shape-dependent m values that are not multiples of + # 32), so run the GEMM on zero-padded rows and let the epilogue write only + # the real ones. + gemm_rows = max(32, -(-rows // 32) * 32) + if gemm_rows != rows: + padded_activations = torch.zeros( + (gemm_rows, in_features), device=device, dtype=torch.int8 + ) + padded_activations[:rows].copy_(activation_values) + gemm_activations = padded_activations + else: + gemm_activations = activation_values + + output = torch.empty((rows, out_features), device=device, dtype=output_dtype) if bias is None: bias_values = output has_bias = False @@ -1664,28 +2567,41 @@ def matmul_packed_w4a4_with_int_mm( has_bias = True combined_scale = float(activation_scale * weight_scale) + if decoded_weight is not None: + if decoded_weight.dtype != torch.int8 or decoded_weight.shape != ( + out_features, + in_features, + ): + raise ValueError( + "decoded_weight must be an int8 tensor of shape " + f"({out_features}, {in_features})" + ) + if not decoded_weight.is_contiguous(): + raise ValueError("decoded_weight must be contiguous") + chunk_size = min(chunk_out_features or out_features, out_features) for output_col_start in range(0, out_features, chunk_size): current_chunk = min(chunk_size, out_features - output_col_start) - packed_offset = output_col_start * packed_k - packed_count = current_chunk * packed_k - packed_weight_chunk = weights.narrow(0, packed_offset, packed_count) - weight_storage = torch.empty( - (current_chunk, in_features), - device=device, - dtype=torch.int8, - ) - weight_total = current_chunk * in_features - _decode_packed_w4_codes_kernel[ - (triton.cdiv(weight_total, decode_block_size),) - ]( - packed_weight_chunk, - weight_code_values, - weight_storage, - total=weight_total, - block_size=decode_block_size, - ) - accumulator = torch._int_mm(activation_values, weight_storage.t()) + if decoded_weight is not None: + weight_storage = decoded_weight.narrow(0, output_col_start, current_chunk) + else: + packed_offset = output_col_start * packed_k + packed_count = current_chunk * packed_k + packed_weight_chunk = weights.narrow(0, packed_offset, packed_count) + weight_storage = torch.empty( + (current_chunk, in_features), + device=device, + dtype=torch.int8, + ) + weight_total = current_chunk * in_features + _decode_packed_w4_codes_kernel[(triton.cdiv(weight_total, decode_block_size),)]( + packed_weight_chunk, + weight_code_values, + weight_storage, + total=weight_total, + block_size=decode_block_size, + ) + accumulator = torch._int_mm(gemm_activations, weight_storage.t()) del weight_storage epilogue_grid = ( triton.cdiv(rows, epilogue_block_m), @@ -1708,6 +2624,7 @@ def matmul_packed_w4a4_with_int_mm( num_warps=4, ) del accumulator + del gemm_activations del activation_values return output.reshape(*original_shape[:-1], out_features) @@ -1728,6 +2645,7 @@ def matmul_int8_activations_packed_w4_with_int_mm( epilogue_block_m: int = 32, epilogue_block_n: int = 256, chunk_out_features: int | None = None, + decoded_weight: torch.Tensor | None = None, ) -> torch.Tensor: return matmul_packed_w4a4_with_int_mm( int8_activations, @@ -1746,6 +2664,7 @@ def matmul_int8_activations_packed_w4_with_int_mm( epilogue_block_n=epilogue_block_n, chunk_out_features=chunk_out_features, activations_are_int8=True, + decoded_weight=decoded_weight, ) @@ -1754,12 +2673,13 @@ def pack_lowbit_with_triton( ) -> torch.Tensor: if bits <= 0 or bits > 8: raise ValueError("bits must be in [1, 8]") - if not values.is_cuda: - raise RuntimeError("triton_cuda pack requires CUDA tensors") + _require_triton_tensor(values, "low-bit packing") max_value = (1 << bits) - 1 flat_raw = values.detach().flatten() - if validate and flat_raw.numel() and bool( - torch.any((flat_raw < 0) | (flat_raw > max_value)).item() + if ( + validate + and flat_raw.numel() + and bool(torch.any((flat_raw < 0) | (flat_raw > max_value)).item()) ): raise ValueError(f"all values must fit in {bits} bits") flat = flat_raw.to(dtype=torch.uint8).contiguous() @@ -1787,8 +2707,7 @@ def unpack_lowbit_with_triton(packed: torch.Tensor, *, bits: int, length: int) - raise ValueError("bits must be in [1, 8]") if length < 0: raise ValueError("length must be non-negative") - if not packed.is_cuda: - raise RuntimeError("triton_cuda unpack requires CUDA tensors") + _require_triton_tensor(packed, "low-bit unpacking") triton, _ = _load_triton() packed_contiguous = packed.detach().to(dtype=torch.uint8).contiguous().flatten() @@ -1809,8 +2728,7 @@ def unpack_lowbit_with_triton(packed: torch.Tensor, *, bits: int, length: int) - def row_norms_with_triton(weight: torch.Tensor, *, eps: float) -> torch.Tensor: - if not weight.is_cuda: - raise RuntimeError("triton_cuda row norm requires CUDA tensors") + _require_triton_tensor(weight, "row norm") if weight.ndim != 2: raise ValueError("weight must be a rank-2 tensor") if not weight.is_floating_point(): @@ -1844,8 +2762,7 @@ def quantize_weight_indices_with_triton( codebook: LloydMaxCodebook, eps: float = 1e-10, ) -> torch.Tensor: - if not weight.is_cuda: - raise RuntimeError("triton_cuda weight quantization requires CUDA tensors") + _require_triton_tensor(weight, "weight quantization") if weight.ndim != 2: raise ValueError("weight must be a rank-2 tensor") if weight.shape[-1] != rotation.dim: @@ -1930,8 +2847,7 @@ def quantize_weight_packed_with_triton( ) -> torch.Tensor: if bits <= 0 or bits > 8: raise ValueError("bits must be in [1, 8]") - if not weight.is_cuda: - raise RuntimeError("triton_cuda packed weight quantization requires CUDA tensors") + _require_triton_tensor(weight, "packed weight quantization") if weight.ndim != 2: raise ValueError("weight must be a rank-2 tensor") if weight.shape[-1] != rotation.dim: diff --git a/src/orbitquant/layers.py b/src/orbitquant/layers.py index ce524aa..6614b37 100644 --- a/src/orbitquant/layers.py +++ b/src/orbitquant/layers.py @@ -1,5 +1,7 @@ from __future__ import annotations +import weakref + import torch from torch import nn from torch.nn import functional as F @@ -13,6 +15,53 @@ from orbitquant.streaming import accelerate_hook_offloads, iter_aligned_row_tiles _PACKED_MATMUL_PROBE_MISSING = object() + +# Measured crossover on RTX 4090-class hardware: below this row count the +# per-forward INT8 weight decode dominates torch._int_mm, so the fused Triton +# kernel that reads packed nibbles directly wins. +# Measured crossover vs the int_mm path with the tuned tile table (RTX 4090, +# 2026-07): the fused kernel stays ahead through 2048 rows. +_W4A4_FUSED_MAX_ROWS = 2048 + +# torch.compile support: the runtime dispatch below is Python-heavy, so the +# whole quantized forward is exposed to Dynamo as one opaque custom op. The +# registry maps a stable integer handle (burned into the traced graph) back to +# the live module; weakref finalizers drop entries when modules die. +_COMPILE_REGISTRY: dict[int, object] = {} + + +def _compile_registry_lookup(handle: int): + module = _COMPILE_REGISTRY.get(handle) + if module is None: + raise RuntimeError( + "orbitquant compiled-forward handle is stale; rebuild the module " + "before calling the compiled graph" + ) + return module + + +def _compile_registry_register(module) -> int: + handle = id(module) + if handle not in _COMPILE_REGISTRY: + _COMPILE_REGISTRY[handle] = module + weakref.finalize(module, _COMPILE_REGISTRY.pop, handle, None) + return handle + + +@torch.library.custom_op("orbitquant::packed_linear_forward", mutates_args=()) +def _packed_linear_forward_op( + x: torch.Tensor, handle: int, out_features: int, bf16_output: bool +) -> torch.Tensor: + module = _compile_registry_lookup(handle) + return module._forward_impl(x) + + +@_packed_linear_forward_op.register_fake +def _packed_linear_forward_fake( + x: torch.Tensor, handle: int, out_features: int, bf16_output: bool +) -> torch.Tensor: + dtype = torch.bfloat16 if bf16_output else x.dtype + return x.new_empty((*x.shape[:-1], out_features), dtype=dtype) _NATIVE_PACKED_MATMUL_LOAD_ERROR: object | Exception | None = _PACKED_MATMUL_PROBE_MISSING _NATIVE_CPU_PACKED_MATMUL_LOAD_ERROR: object | Exception | None = ( _PACKED_MATMUL_PROBE_MISSING @@ -20,6 +69,10 @@ _TRITON_PACKED_MATMUL_IMPORT_ERROR: object | Exception | None = _PACKED_MATMUL_PROBE_MISSING +def _torch_uses_hip() -> bool: + return bool(getattr(torch.version, "hip", None)) + + def _packed_length(value_count: int, bits: int) -> int: return (value_count * bits + 7) // 8 @@ -46,7 +99,7 @@ def _quantize_weight_indices( codebook, eps: float, ) -> torch.Tensor: - if weight.is_cuda: + if weight.device.type in {"cuda", "xpu"}: from orbitquant.kernels.triton_cuda import quantize_weight_indices_with_triton return quantize_weight_indices_with_triton( @@ -73,7 +126,7 @@ def _quantize_weight_pack( bits: int, eps: float, ) -> torch.Tensor: - if weight.is_cuda: + if weight.device.type in {"cuda", "xpu"}: from orbitquant.kernels.triton_cuda import quantize_weight_packed_with_triton return quantize_weight_packed_with_triton( @@ -131,7 +184,7 @@ def _quantize_weight_bounded( row_tile_size, ): tile = weight[row_start:row_end].to(device=work_device) - if tile.is_cuda: + if tile.device.type in {"cuda", "xpu"}: from orbitquant.kernels.triton_cuda import row_norms_with_triton row_norms = row_norms_with_triton(tile, eps=eps) @@ -270,6 +323,7 @@ def __init__( self.weight_bits = config.weight_bits self.activation_bits = config.activation_bits self.runtime_mode = config.runtime_mode + self.w4a4_int8_weight_cache = config.w4a4_int8_weight_cache self.activation_kernel_backend = config.activation_kernel_backend self.packed_matmul_block_m = config.packed_matmul_block_m self.packed_matmul_block_n = config.packed_matmul_block_n @@ -295,7 +349,9 @@ def __init__( ) self.register_buffer( "_rotation_permutation", - _clone_constant(self.rotation.permutation, device=constant_device), + _clone_constant( + self.rotation.permutation.to(torch.int32), device=constant_device + ), persistent=False, ) self.register_buffer( @@ -343,6 +399,10 @@ def __init__( self._int8_surrogate_cache: tuple[str, torch.Tensor, float, torch.Tensor, float] | None = ( None ) + self._runtime_probe_cache: dict[str, bool] = {} + self._bias_cache: tuple[tuple[str, torch.dtype], torch.Tensor] | None = None + self._int8_weight_cache: tuple[str, torch.Tensor] | None = None + self._compile_handle = _compile_registry_register(self) self.last_effective_runtime_mode: str | None = None self.last_activation_kernel_backend: str | None = None self.last_forward_device_type: str | None = None @@ -510,14 +570,55 @@ def _apply(self, fn, recurse: bool = True): result = super()._apply(fn, recurse=recurse) self._derived_constants_valid = False self._int8_surrogate_cache = None + self._runtime_probe_cache = {} + self._bias_cache = None + self._int8_weight_cache = None self.clear_dequantized_cache() return result + def _int8_weight_for(self, device: torch.device) -> torch.Tensor | None: + """Opt-in persistent INT8 surrogate weight (skips per-forward decode).""" + if not self.w4a4_int8_weight_cache or accelerate_hook_offloads(self): + return None + key = str(device) + if self._int8_weight_cache is not None and self._int8_weight_cache[0] == key: + return self._int8_weight_cache[1] + from orbitquant.kernels.triton_cuda import decode_packed_w4_weight_to_int8 + + _, _, weight_codes, _ = self._int8_surrogate_constants(device) + decoded = decode_packed_w4_weight_to_int8( + self.packed_weight_indices, + weight_codes, + out_features=self.out_features, + in_features=self.in_features, + ) + self._int8_weight_cache = (key, decoded) + return decoded + + def _cached_runtime_probe(self, key: str, probe) -> bool: + cached = self._runtime_probe_cache.get(key) + if cached is None: + cached = bool(probe()) + self._runtime_probe_cache[key] = cached + return cached + + def _bias_for(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None: + if self.bias is None: + return None + if self.bias.device == device and self.bias.dtype == dtype: + return self.bias + key = (str(device), dtype) + if self._bias_cache is not None and self._bias_cache[0] == key: + return self._bias_cache[1] + value = self.bias.to(device=device, dtype=dtype) + self._bias_cache = (key, value) + return value + def _ensure_derived_constants(self, device: torch.device) -> None: if self._derived_constants_valid and self._rotation_permutation.device == device: return constants = ( - ("_rotation_permutation", self.rotation.permutation, torch.int64), + ("_rotation_permutation", self.rotation.permutation, torch.int32), ("_rotation_signs", self.rotation.signs, torch.int8), ( "_activation_codebook_centroids", @@ -556,7 +657,7 @@ def _constant_buffer( def _activation_kernel_constant_tensors(self, device: torch.device) -> dict[str, torch.Tensor]: return { "permutation": self._constant_buffer( - "_rotation_permutation", device=device, dtype=torch.int64 + "_rotation_permutation", device=device, dtype=torch.int32 ), "signs": self._constant_buffer("_rotation_signs", device=device, dtype=torch.int8), "centroids": self._constant_buffer( @@ -568,84 +669,136 @@ def _activation_kernel_constant_tensors(self, device: torch.device) -> dict[str, } def _validate_triton_packed_matmul_input(self, x: torch.Tensor) -> None: - if x.device.type != "cuda": + if x.device.type not in {"cuda", "xpu"}: raise RuntimeError( - f"triton_packed_matmul runtime requires CUDA input tensors; got {x.device.type}." + "triton_packed_matmul runtime requires CUDA, HIP, or XPU input tensors; " + f"got {x.device.type}." ) def _validate_native_packed_matmul_input(self, x: torch.Tensor) -> None: + if x.device.type == "cuda" and _torch_uses_hip(): + raise RuntimeError( + "native_packed_matmul currently provides CUDA, not HIP, accelerator " + "kernels. Use runtime_mode='auto_fused' or 'triton_packed_matmul' on " + "ROCm, or runtime_mode='dequant_bf16' for the explicit reference path." + ) if x.device.type not in {"cpu", "cuda", "mps"}: raise RuntimeError( "native_packed_matmul runtime requires CPU, CUDA, or MPS input tensors; " f"got {x.device.type}." ) + def _fused_w2a4_available(self, x: torch.Tensor) -> bool: + if ( + x.device.type != "cuda" + or _torch_uses_hip() + or x.dtype not in {torch.bfloat16, torch.float16} + or self.weight_bits != 2 + or self.activation_bits != 4 + or self.in_features % 64 != 0 + or self.out_features % 16 != 0 + or self.packed_weight_indices is None + or self.row_norms is None + ): + return False + rows = x.numel() // self.in_features + if not (32 <= rows < _W4A4_FUSED_MAX_ROWS): + return False + + def probe() -> bool: + try: + from orbitquant.kernels.triton_cuda import _load_triton + + _load_triton() + except Exception: + return False + return True + + return self._cached_runtime_probe("triton_fused_lowbit", probe) + def _native_w4a4_available(self, x: torch.Tensor) -> bool: if ( x.device.type != "cuda" + or _torch_uses_hip() or x.dtype not in {torch.bfloat16, torch.float16} or self.weight_bits != 4 or self.activation_bits != 4 or self.in_features % 64 != 0 ): return False - try: - from orbitquant.kernels.native_packed_matmul import ( - native_packed_w4a4_available, - ) + def probe() -> bool: + try: + from orbitquant.kernels.native_packed_matmul import ( + native_packed_w4a4_available, + ) - if not native_packed_w4a4_available(): + if not native_packed_w4a4_available(): + return False + from orbitquant.kernels.triton_cuda import ( # noqa: F401 + quantize_activations_packed_w4_with_triton, + ) + except (ImportError, RuntimeError): return False - from orbitquant.kernels.triton_cuda import ( # noqa: F401 - quantize_activations_packed_w4_with_triton, - ) - except (ImportError, RuntimeError): - return False - return True + return True + + return self._cached_runtime_probe("native_w4a4_kernels", probe) def _native_cpu_activation_available(self, x: torch.Tensor) -> bool: if x.device.type != "cpu": return False - try: - from orbitquant.kernels.native_packed_matmul import ( - native_cpu_activation_available, - ) - return native_cpu_activation_available() - except (ImportError, RuntimeError): - return False + def probe() -> bool: + try: + from orbitquant.kernels.native_packed_matmul import ( + native_cpu_activation_available, + ) + + return native_cpu_activation_available() + except (ImportError, RuntimeError): + return False + + return self._cached_runtime_probe("native_cpu_activation", probe) def _native_w4_activation_available(self, x: torch.Tensor) -> bool: if ( not x.is_cuda + or _torch_uses_hip() or self.rotation.block_size != self.in_features or self.in_features not in {512, 1024, 2048, 4096, 8192, 16384} ): return False - try: - from orbitquant.kernels.native_packed_matmul import ( - native_packed_w4_activation_available, - ) - return native_packed_w4_activation_available() - except (ImportError, RuntimeError): - return False + def probe() -> bool: + try: + from orbitquant.kernels.native_packed_matmul import ( + native_packed_w4_activation_available, + ) + + return native_packed_w4_activation_available() + except (ImportError, RuntimeError): + return False + + return self._cached_runtime_probe("native_w4_activation", probe) def _native_int8_activation_available(self, x: torch.Tensor) -> bool: supported_rotation = ( self.rotation.block_size == self.in_features and self.in_features in {512, 1024, 2048, 4096, 8192, 16384} ) or (self.in_features, self.rotation.block_size) == (12288, 4096) - if not x.is_cuda or not supported_rotation: + if not x.is_cuda or _torch_uses_hip() or not supported_rotation: return False - try: - from orbitquant.kernels.native_packed_matmul import ( - native_int8_activation_available, - ) - return native_int8_activation_available() - except (ImportError, RuntimeError): - return False + def probe() -> bool: + try: + from orbitquant.kernels.native_packed_matmul import ( + native_int8_activation_available, + ) + + return native_int8_activation_available() + except (ImportError, RuntimeError): + return False + + return self._cached_runtime_probe("native_int8_activation", probe) def _int8_surrogate_constants( self, device: torch.device @@ -701,6 +854,17 @@ def _auto_fused_unavailable_error( "or make the `orbitquant_packed_matmul` package importable." ) if device_type == "cuda": + if _torch_uses_hip(): + triton_hint = ( + "Install the ROCm-compatible PyTorch Triton package to use " + "triton_packed_matmul." + ) + return RuntimeError( + "auto_fused runtime requires packed low-bit matmul on ROCm and will " + "not silently materialize a full dequantized weight matrix. " + f"triton_packed_matmul failed: {triton_error}. {triton_hint} " + f"{reference_hint}" + ) triton_hint = "Install a CUDA-compatible Triton stack to use triton_packed_matmul." return RuntimeError( "auto_fused runtime requires packed low-bit matmul on CUDA and will not " @@ -715,6 +879,14 @@ def _auto_fused_unavailable_error( f"weight matrix. native_packed_matmul failed: {native_error}. " f"{native_hint} {reference_hint}" ) + if device_type == "xpu": + return RuntimeError( + "auto_fused XPU dispatch remains experimental until real Intel GPU " + "correctness, memory, profiler, and performance proof is recorded. " + "For explicit validation, set runtime_mode='triton_packed_matmul' and " + "activation_kernel_backend='triton_xpu'. " + f"{reference_hint}" + ) return RuntimeError( f"auto_fused runtime does not support device type {device_type!r}. {reference_hint}" ) @@ -726,6 +898,21 @@ def _resolve_auto_fused_runtime(self, x: torch.Tensor) -> str: return "native_packed_matmul" return "dequant_bf16" if device_type == "cuda": + if _torch_uses_hip(): + triton_error = _triton_packed_matmul_import_error() + if triton_error is not None: + raise self._auto_fused_unavailable_error( + device_type=device_type, + native_error=None, + triton_error=triton_error, + ) + raise RuntimeError( + "auto_fused ROCm dispatch remains experimental until real AMD " + "hardware correctness, memory, and performance proof is recorded. " + "For explicit validation, set runtime_mode='triton_packed_matmul' " + "and activation_kernel_backend='triton_rocm'. Set " + "runtime_mode='dequant_bf16' for the reference/debug path." + ) native_error = _native_packed_matmul_load_error() if native_error is None: return "native_packed_matmul" @@ -745,6 +932,11 @@ def _resolve_auto_fused_runtime(self, x: torch.Tensor) -> str: device_type=device_type, native_error=native_error, ) + if device_type == "xpu": + raise self._auto_fused_unavailable_error( + device_type=device_type, + native_error=None, + ) raise self._auto_fused_unavailable_error( device_type=device_type, native_error=None, @@ -784,7 +976,7 @@ def _dequantize_weight(self, *, device: torch.device, dtype: torch.dtype) -> tor ) dequantized = weight.to(dtype=dtype) return self._remember_dequantized_weight(dequantized, cache_key) - if device.type == "cuda": + if device.type in {"cuda", "xpu"}: try: from orbitquant.kernels.triton_cuda import dequantize_packed_weight_with_triton except Exception: @@ -815,6 +1007,13 @@ def _dequantize_weight(self, *, device: torch.device, dtype: torch.dtype) -> tor return self._remember_dequantized_weight(dequantized, cache_key) def forward(self, x: torch.Tensor) -> torch.Tensor: + if torch.compiler.is_compiling(): + return torch.ops.orbitquant.packed_linear_forward( + x, self._compile_handle, self.out_features, False + ) + return self._forward_impl(x) + + def _forward_impl(self, x: torch.Tensor) -> torch.Tensor: self._ensure_derived_constants(x.device) runtime_mode = ( self._resolve_auto_fused_runtime(x) @@ -832,6 +1031,75 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: elif runtime_mode == "native_packed_matmul": self._validate_native_packed_matmul_input(x) + if runtime_mode == "native_packed_matmul" and self._fused_w2a4_available(x): + activation_constants = self._activation_kernel_constant_tensors(x.device) + activation_codes, activation_scale, weight_codes, weight_scale = ( + self._int8_surrogate_constants(x.device) + ) + bias = self._bias_for(x.device, x.dtype) + if self._native_int8_activation_available(x): + from orbitquant.kernels.native_packed_matmul import ( + quantize_activations_int8_with_native_kernel, + ) + from orbitquant.kernels.triton_cuda import ( + matmul_int8_activations_packed_lowbit_fused_with_triton, + ) + + capability = torch.cuda.get_device_capability(x.device) + int8_x, token_norms = quantize_activations_int8_with_native_kernel( + x, + activation_constants["permutation"], + activation_constants["signs"], + activation_constants["boundaries"], + activation_codes, + eps=self.activation_eps, + inv_sqrt_block=self.rotation.normalization, + threads=512 if capability == (8, 9) else 256, + ) + self.last_activation_kernel_backend = "native_cuda_int8_surrogate" + self.last_native_int8_activation_enabled = True + return matmul_int8_activations_packed_lowbit_fused_with_triton( + int8_x, + self.packed_weight_indices, + token_norms, + self.row_norms, + weight_codes, + weight_bits=2, + activation_scale=activation_scale, + weight_scale=weight_scale, + out_features=self.out_features, + in_features=self.in_features, + bias=bias, + output_dtype=x.dtype, + ) + from orbitquant.kernels.triton_cuda import ( + matmul_packed_w2a4_fused_with_triton, + quantize_activations_packed_w4_with_triton, + ) + + packed_x, token_norms = quantize_activations_packed_w4_with_triton( + x, + rotation=self.rotation, + codebook=self.activation_codebook, + eps=self.activation_eps, + constant_tensors=activation_constants, + ) + self.last_activation_kernel_backend = "triton_cuda_packed_w4" + return matmul_packed_w2a4_fused_with_triton( + packed_x, + self.packed_weight_indices, + token_norms, + self.row_norms, + activation_codes, + weight_codes, + activation_scale=activation_scale, + weight_scale=weight_scale, + out_features=self.out_features, + in_features=self.in_features, + bias=bias, + output_dtype=x.dtype, + ) + if runtime_mode == "native_packed_matmul" and self._native_w4a4_available(x): if self.packed_weight_indices is None or self.row_norms is None: raise RuntimeError("OrbitQuantLinear is missing quantized weight buffers") @@ -840,11 +1108,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: activation_codes, activation_scale, weight_codes, weight_scale = ( self._int8_surrogate_constants(x.device) ) + rows = x.numel() // self.in_features + # cuBLASLt INT8 GEMM needs padded rows below 32; the fused native + # tile kernel reads the packed weights directly and wins there. use_cutlass_tn = ( capability[0] >= 8 + and rows >= 32 and self.out_features % 16 == 0 and callable(getattr(torch, "_int_mm", None)) ) + int8_weight = self._int8_weight_for(x.device) if use_cutlass_tn else None + # Mid-size row counts pay more for the per-forward INT8 weight + # decode than the GEMM itself; the fused Triton kernel reads the + # packed nibbles directly there. + use_fused_triton = ( + use_cutlass_tn and int8_weight is None and rows < _W4A4_FUSED_MAX_ROWS + ) int8_x: torch.Tensor | None = None if use_cutlass_tn and self._native_int8_activation_available(x): from orbitquant.kernels.native_packed_matmul import ( @@ -893,13 +1172,54 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: constant_tensors=activation_constants, ) self.last_activation_kernel_backend = "triton_cuda_packed_w4" - bias = None if self.bias is None else self.bias.to(device=x.device, dtype=x.dtype) + bias = self._bias_for(x.device, x.dtype) self.last_native_w4a4_enabled = True + if use_fused_triton: + if int8_x is not None: + from orbitquant.kernels.triton_cuda import ( + matmul_int8_activations_packed_lowbit_fused_with_triton, + ) + + return matmul_int8_activations_packed_lowbit_fused_with_triton( + int8_x, + self.packed_weight_indices, + token_norms, + self.row_norms, + weight_codes, + weight_bits=4, + activation_scale=activation_scale, + weight_scale=weight_scale, + out_features=self.out_features, + in_features=self.in_features, + bias=bias, + output_dtype=x.dtype, + ) + from orbitquant.kernels.triton_cuda import ( + matmul_packed_w4a4_fused_with_triton, + ) + + return matmul_packed_w4a4_fused_with_triton( + packed_x, + self.packed_weight_indices, + token_norms, + self.row_norms, + activation_codes, + weight_codes, + activation_scale=activation_scale, + weight_scale=weight_scale, + out_features=self.out_features, + in_features=self.in_features, + bias=bias, + output_dtype=x.dtype, + ) if use_cutlass_tn: - if self.in_features >= 8192 and self.out_features <= 4096: - cutlass_chunk_out_features = 4096 + # Chunking bounds the transient INT8 weight but re-reads the + # activations per chunk; it only measured faster for large row + # counts against wide projections. + if rows >= 2048 and self.out_features >= 8192: + cutlass_chunk_out_features: int | None = 2048 else: - cutlass_chunk_out_features = 2048 + cutlass_chunk_out_features = None if int8_x is not None: from orbitquant.kernels.triton_cuda import ( matmul_int8_activations_packed_w4_with_int_mm, @@ -918,6 +1238,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: bias=bias, output_dtype=x.dtype, chunk_out_features=cutlass_chunk_out_features, + decoded_weight=int8_weight, ) from orbitquant.kernels.triton_cuda import ( matmul_packed_w4a4_with_int_mm, @@ -937,12 +1258,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: bias=bias, output_dtype=x.dtype, chunk_out_features=cutlass_chunk_out_features, + decoded_weight=int8_weight, ) from orbitquant.kernels.native_packed_matmul import ( matmul_packed_w4a4_int8_with_native_kernel, ) - rows = x.numel() // self.in_features tile_m, tile_n = self._native_w4a4_tile( rows=rows, out_features=self.out_features, @@ -1007,15 +1328,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: constant_tensors=self._activation_kernel_constant_tensors(x.device), ) - bias = None if self.bias is None else self.bias.to(device=x.device, dtype=rotated_x.dtype) + bias = self._bias_for(x.device, rotated_x.dtype) if runtime_mode == "triton_packed_matmul": if self.packed_weight_indices is None or self.row_norms is None: raise RuntimeError("OrbitQuantLinear is missing quantized weight buffers") try: from orbitquant.kernels.triton_cuda import matmul_packed_weight_with_triton except Exception as exc: + backend = ( + "XPU" + if x.device.type == "xpu" + else "ROCm" + if _torch_uses_hip() + else "CUDA" + ) raise RuntimeError( - "triton_packed_matmul runtime requires the Triton CUDA backend" + f"triton_packed_matmul runtime requires the Triton {backend} backend" ) from exc return matmul_packed_weight_with_triton( rotated_x, diff --git a/src/orbitquant/modeling.py b/src/orbitquant/modeling.py index da8faf2..c3a02da 100644 --- a/src/orbitquant/modeling.py +++ b/src/orbitquant/modeling.py @@ -164,6 +164,10 @@ def _quantization_device(device: str | torch.device | None) -> torch.device | No raise RuntimeError("CUDA quantization device requested but CUDA is not available") if torch_device.type == "mps" and not torch.backends.mps.is_available(): raise RuntimeError("MPS quantization device requested but MPS is not available") + if torch_device.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is None or not xpu.is_available(): + raise RuntimeError("XPU quantization device requested but XPU is not available") return torch_device @@ -171,14 +175,25 @@ def _weight_quantization_backend(device: torch.device | None) -> str: if device is None: return "module_device" if device.type == "cuda": - if not available_backends()["triton_cuda"]: + backend = "triton_rocm" if getattr(torch.version, "hip", None) else "triton_cuda" + if not available_backends().get(backend, False): raise RuntimeError( - "CUDA quantization requires the Triton CUDA backend. Install triton " - "or use --device cpu for the reference quantization path." + f"GPU quantization requires the {backend} backend. Install the " + "matching PyTorch Triton package or use --device cpu for the " + "reference quantization path." ) - return "triton_cuda" + return backend if device.type == "mps": return "torch_reference_mps" + if device.type == "xpu": + backend = "triton_xpu" + if not available_backends().get(backend, False): + raise RuntimeError( + "XPU quantization requires the triton_xpu backend. Install the " + "PyTorch XPU wheel with its matching Intel Triton package or use " + "--device cpu for the reference quantization path." + ) + return backend return "torch_reference" @@ -396,6 +411,10 @@ def _synchronize_if_needed(device: torch.device) -> None: torch.cuda.synchronize(device) elif device.type == "mps" and torch.backends.mps.is_available(): torch.mps.synchronize() + elif device.type == "xpu": + xpu = getattr(torch, "xpu", None) + if xpu is not None and xpu.is_available(): + xpu.synchronize(device) def prewarm_quantized_linear_modules( @@ -426,7 +445,7 @@ def prewarm_quantized_linear_modules( if not ( isinstance(module, OrbitQuantLinear) and module.runtime_mode == "auto_fused" - and module_device.type in {"cuda", "mps"} + and module_device.type in {"cuda", "mps", "xpu"} ): module._dequantize_weight(device=module_device, dtype=module_dtype) last_device = module_device diff --git a/src/orbitquant/packing/bitpack.py b/src/orbitquant/packing/bitpack.py index 6485223..1b4ac8b 100644 --- a/src/orbitquant/packing/bitpack.py +++ b/src/orbitquant/packing/bitpack.py @@ -14,6 +14,22 @@ def _pack_lowbit_cpu(values: torch.Tensor, bits: int, *, validate: bool) -> torc raise ValueError(f"all values must fit in {bits} bits") values_cpu = values_cpu_raw.to(dtype=torch.uint8) + # Byte-aligned widths pack with a couple of vectorized shift/or passes + # instead of the generic per-bit scatter_add loop below. + if bits in {2, 4, 8}: + fields_per_byte = 8 // bits + count = values_cpu.numel() + remainder = count % fields_per_byte + if remainder: + padded = torch.zeros(count + fields_per_byte - remainder, dtype=torch.uint8) + padded[:count] = values_cpu + values_cpu = padded + fields = values_cpu.reshape(-1, fields_per_byte).to(torch.int16) + packed_bytes = fields[:, 0].clone() + for field in range(1, fields_per_byte): + packed_bytes |= fields[:, field] << (field * bits) + return packed_bytes.to(dtype=torch.uint8) + total_bits = values_cpu.numel() * bits packed = torch.zeros((total_bits + 7) // 8, dtype=torch.int16) bit_ids = torch.arange(bits, dtype=torch.int64) @@ -58,6 +74,15 @@ def unpack_lowbit(packed: torch.Tensor, bits: int, length: int) -> torch.Tensor: return unpack_lowbit_with_triton(packed, bits=bits, length=length) packed_cpu = packed.detach().to(device="cpu", dtype=torch.uint8).flatten() + + if bits in {2, 4, 8}: + fields_per_byte = 8 // bits + mask = (1 << bits) - 1 + expanded = torch.empty(packed_cpu.numel() * fields_per_byte, dtype=torch.uint8) + for field in range(fields_per_byte): + expanded[field::fields_per_byte] = (packed_cpu >> (field * bits)) & mask + return expanded[:length] + values = torch.zeros(length, dtype=torch.uint8) bit_ids = torch.arange(bits, dtype=torch.int64) diff --git a/src/orbitquant/py.typed b/src/orbitquant/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/src/orbitquant/quantizer.py b/src/orbitquant/quantizer.py index 2e180a7..aa17ec6 100644 --- a/src/orbitquant/quantizer.py +++ b/src/orbitquant/quantizer.py @@ -221,11 +221,22 @@ def adjust_target_dtype(self, torch_dtype: torch.dtype | None) -> torch.dtype | def update_device_map(self, device_map: Any | None) -> Any | None: return device_map + def _classified_decisions(self, model: Any) -> dict: + # classify_linear_modules walks every module; Hugging Face calls the + # per-parameter hooks once per checkpoint tensor, so memoize the walk + # for the model instance being loaded. + cached = getattr(self, "_classification_cache", None) + if cached is not None and cached[0] is model: + return cached[1] + decisions = classify_linear_modules(model, self.quantization_config) + self._classification_cache = (model, decisions) + return decisions + def _param_action(self, model: Any, param_name: str) -> str | None: if not param_name.endswith(".weight"): return None module_name = param_name.removesuffix(".weight") - decisions = classify_linear_modules(model, self.quantization_config) + decisions = self._classified_decisions(model) decision = decisions.get(module_name) return None if decision is None else decision.action @@ -422,6 +433,7 @@ def is_compileable(self) -> bool: def _process_model_before_weight_loading(self, model: Any, *args: Any, **kwargs: Any) -> Any: model.quantization_config = self.quantization_config + self._classification_cache = None checkpoint_files = kwargs.get("checkpoint_files") self._transformers_postload_quantization = ( not self.pre_quantized and checkpoint_files is not None @@ -438,7 +450,7 @@ def _process_model_before_weight_loading(self, model: Any, *args: Any, **kwargs: self._transformers_adaln_module_names = [] self._transformers_base_model_prefix = str(getattr(model, "base_model_prefix", "") or "") if self._transformers_streaming_quantization or self._diffusers_streaming_quantization: - decisions = classify_linear_modules(model, self.quantization_config) + decisions = self._classified_decisions(model) self._transformers_orbit_module_names = [ name for name, decision in decisions.items() if decision.action == "orbitquant" ] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..0b6f8dd --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,12 @@ +import pytest + +import orbitquant.kernels.dispatch as _dispatch_module + + +@pytest.fixture(autouse=True) +def _clear_backend_availability_cache(): + """Backend availability is memoized per process; tests monkeypatch the + underlying probes, so reset the cache around every test.""" + _dispatch_module.clear_backend_availability_cache() + yield + _dispatch_module.clear_backend_availability_cache() diff --git a/tests/hub_helpers.py b/tests/hub_helpers.py new file mode 100644 index 0000000..b29e4ed --- /dev/null +++ b/tests/hub_helpers.py @@ -0,0 +1,340 @@ +"""Shared Hugging Face fakes and artifact helpers for the test_hub_* modules.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import torch + +import orbitquant.hub as hub_module +from orbitquant.artifacts import save_orbitquant_artifact +from orbitquant.config import OrbitQuantConfig +from orbitquant.modeling import quantize_linear_modules + + +class TinyHubArtifactModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + + +class FakeHfApi: + def __init__(self): + self.create_repo_calls = [] + self.upload_folder_calls = [] + self.model_info_calls = [] + + def create_repo(self, **kwargs): + self.create_repo_calls.append(kwargs) + return f"https://huggingface.co/{kwargs['repo_id']}" + + def upload_folder(self, **kwargs): + self.upload_folder_calls.append(kwargs) + return SimpleNamespace( + oid="uploaded-sha", + commit_url=f"https://huggingface.co/{kwargs['repo_id']}/commit/uploaded-sha", + pr_url=None, + ) + + def model_info(self, repo_id, *, revision=None): + self.model_info_calls.append({"repo_id": repo_id, "revision": revision}) + return SimpleNamespace(sha=revision or "main-sha", private=True, gated=False) + + +class FakeAuditHfApi: + def __init__(self, siblings_by_repo): + self.siblings_by_repo = siblings_by_repo + self.model_info_calls = [] + + def model_info(self, repo_id, *, revision=None, files_metadata=False): + self.model_info_calls.append( + {"repo_id": repo_id, "revision": revision, "files_metadata": files_metadata} + ) + if repo_id not in self.siblings_by_repo: + raise RuntimeError("missing repo") + siblings = [] + for name, metadata in self.siblings_by_repo[repo_id].items(): + if isinstance(metadata, dict): + lfs_sha256 = metadata.get("lfs_sha256") + lfs = ( + SimpleNamespace(sha256=lfs_sha256, size=metadata.get("size")) + if lfs_sha256 is not None + else None + ) + siblings.append(SimpleNamespace(rfilename=name, size=metadata.get("size"), lfs=lfs)) + else: + siblings.append(SimpleNamespace(rfilename=name, size=metadata, lfs=None)) + return SimpleNamespace( + sha="remote-sha", + private=True, + gated=False, + siblings=siblings, + ) + + +class FakeCommitHfApi: + def __init__(self): + self.create_commit_calls = [] + + def create_commit(self, **kwargs): + self.create_commit_calls.append(kwargs) + return SimpleNamespace( + oid="repair-sha", + commit_url=f"https://huggingface.co/{kwargs['repo_id']}/commit/repair-sha", + pr_url=None, + ) + + +class FakeCleanupHfApi(FakeCommitHfApi): + def __init__(self, artifact_dir): + super().__init__() + self.artifact_dir = artifact_dir + self.model_info_calls = [] + + def model_info(self, repo_id, *, revision=None, files_metadata=False): + self.model_info_calls.append( + {"repo_id": repo_id, "revision": revision, "files_metadata": files_metadata} + ) + siblings = [ + SimpleNamespace(rfilename=path.relative_to(self.artifact_dir).as_posix(), size=1) + for path in sorted(self.artifact_dir.rglob("*")) + if path.is_file() + ] + return SimpleNamespace(sha="remote-sha", private=True, gated=False, siblings=siblings) + + +def _write_artifact(tmp_path): + model = TinyHubArtifactModel() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config, quantization_device=None) + return save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + +def _required_remote_files(): + return {name: 1 for name in hub_module._REQUIRED_ARTIFACT_FILES} + + +def _remote_file_map(repo_id, artifact_dir): + return { + (repo_id, "orbitquant_manifest.json"): artifact_dir / "orbitquant_manifest.json", + (repo_id, "model_index.json"): artifact_dir / "model_index.json", + (repo_id, "benchmark/summary.json"): artifact_dir / "benchmark" / "summary.json", + (repo_id, "quantization_config.json"): artifact_dir / "quantization_config.json", + (repo_id, "README.md"): artifact_dir / "README.md", + (repo_id, "SHA256SUMS"): artifact_dir / "SHA256SUMS", + } + + +def _remote_path(artifact_dir, filename): + return artifact_dir / filename + + +def make_fake_download(file_map=None, *, artifact_dir=None): + """Build an ``hf_hub_download`` replacement for monkeypatching. + + Resolves ``(repo_id, filename)`` keys through ``file_map`` when given, + otherwise resolves ``filename`` directly inside ``artifact_dir``. + """ + if (file_map is None) == (artifact_dir is None): + raise ValueError("provide exactly one of file_map or artifact_dir") + + def fake_download(repo, filename, **kwargs): + if file_map is not None: + return str(file_map[(repo, filename)]) + return str(_remote_path(artifact_dir, filename)) + + return fake_download + + +def _write_remote_model_index( + path, + *, + activation_eps=None, + quantization_device="unknown", + weight_quantization_backend="unknown", + quantization_staging_mode="unknown", +): + payload = { + "_class_name": "OrbitQuantArtifact", + "quantization_device": quantization_device, + "weight_quantization_backend": weight_quantization_backend, + "quantization_staging_mode": quantization_staging_mode, + } + if activation_eps is not None: + payload["activation_eps"] = activation_eps + path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + + +def _audit_file_map( + repo_id, + *, + manifest_path, + sha256sums_path, + summary_path, + model_index_path, + readme_path=None, + quantization_config_path=None, +): + if readme_path is None: + readme_path = Path(manifest_path).parent / "README.md" + manifest_payload = json.loads(Path(manifest_path).read_text(encoding="utf-8")) + manifest_payload.setdefault("source_revision", "unknown") + manifest_payload.setdefault("source_license", "unknown") + manifest = hub_module.OrbitQuantManifest.from_dict(manifest_payload) + benchmark_summary = json.loads(Path(summary_path).read_text(encoding="utf-8")) + Path(readme_path).write_text( + hub_module.render_model_card(manifest, benchmark_summary=benchmark_summary), + encoding="utf-8", + ) + file_map = { + (repo_id, "orbitquant_manifest.json"): manifest_path, + (repo_id, "model_index.json"): model_index_path, + (repo_id, "SHA256SUMS"): sha256sums_path, + (repo_id, "benchmark/summary.json"): summary_path, + (repo_id, "README.md"): readme_path, + } + if quantization_config_path is not None: + file_map[(repo_id, "quantization_config.json")] = quantization_config_path + return file_map + + +def _native_smoke_summary( + suite, + *, + release_metrics=None, + activation_eps=None, + quantization_device="unknown", + weight_quantization_backend="unknown", + quantization_staging_mode="unknown", +): + native_settings = hub_module._native_smoke_expected_settings(suite) + generated_frames = 0 if suite.frames is None else suite.frames + metrics = {"generated_samples": 1} + if generated_frames: + metrics["generated_frames"] = generated_frames + for key, value in (release_metrics or {}).items(): + metrics[key] = value + split_proof = { + "records": 1, + "generated_samples": 1, + "generated_frames": generated_frames, + "nonempty_output_count": 1, + "seeds": ["0"], + "prompt_ids": ["simple-object"], + "pair_keys": [[suite.name, "0", "simple-object"]], + "native_settings": [native_settings], + } + payload = { + "published_summary": "compact", + "raw_generation_records": "local-only", + "quantization_device": quantization_device, + "weight_quantization_backend": weight_quantization_backend, + "quantization_staging_mode": quantization_staging_mode, + "metrics": { + "original": {"records": 1, "latest_metrics": metrics}, + "orbitquant": {"records": 1, "latest_metrics": metrics}, + }, + "native_smoke": { + "proof_format": "orbitquant-native-smoke-v1", + "comparison_asset_path": "assets/image_generation_comparison_matrix.webp", + "paired_prompt_seed_count": 1, + "paired_prompt_seed_keys": [[suite.name, "0", "simple-object"]], + "splits": { + "original": split_proof, + "orbitquant": split_proof, + }, + }, + } + if activation_eps is not None: + payload["activation_eps"] = activation_eps + return json.dumps(payload, indent=2) + + +def _legacy_compact_summary_without_native_smoke( + suite, + *, + generated_samples=1, + generated_frames=0, +): + metrics = {"generated_samples": generated_samples} + if generated_frames: + metrics["generated_frames"] = generated_frames + return json.dumps( + { + "published_summary": "compact", + "raw_generation_records": "local-only", + "activation_eps": 1e-10, + "quantization_device": "cuda", + "weight_quantization_backend": "triton_cuda", + "quantization_staging_mode": "component", + "metrics": { + "original": {"records": generated_samples, "latest_metrics": metrics}, + "orbitquant": {"records": generated_samples, "latest_metrics": metrics}, + }, + }, + indent=2, + ) + + +def _write_native_smoke_backup( + backup_root, + *, + repo_id, + suite, + bit_setting, + prompt_ids=("simple-object", "counting"), +): + assets = backup_root / repo_id.replace("/", "__") / "assets" + assets.mkdir(parents=True, exist_ok=True) + extension = "mp4" if suite.frames is not None else "png" + for prompt_id in prompt_ids: + for split_token, quantization in ( + ("original", None), + ( + bit_setting, + { + "config": { + "weight_bits": int(bit_setting.split("A", maxsplit=1)[0][1:]), + "activation_bits": int(bit_setting.split("A", maxsplit=1)[1]), + } + }, + ), + ): + media_path = assets / f"{suite.name}_seed0_{split_token}_{prompt_id}.{extension}" + media_path.write_bytes(b"not empty") + metadata = { + "suite": suite.name, + "model_id": suite.model_id, + "prompt": f"prompt for {prompt_id}", + "seed": 0, + "height": suite.height, + "width": suite.width, + "frames": suite.frames, + "steps": suite.steps, + "guidance": suite.guidance, + "quantization": quantization, + } + if suite.export_fps is not None: + metadata["export_fps"] = suite.export_fps + (media_path.with_suffix(media_path.suffix + ".json")).write_text( + json.dumps(metadata, indent=2) + "\n", + encoding="utf-8", + ) + + +def _expected_missing_geneval_metrics(): + return [ + {"split": split, "metric": metric} + for metric in hub_module._GENEVAL_REQUIRED_METRICS + for split in ("original", "orbitquant") + ] diff --git a/tests/test_adaln_rtn.py b/tests/test_adaln_rtn.py index 63036a0..cb6406c 100644 --- a/tests/test_adaln_rtn.py +++ b/tests/test_adaln_rtn.py @@ -15,17 +15,19 @@ def _independent_packed_adaln_weight(layer: RTNInt4Linear) -> torch.Tensor: packed = layer.packed_weight.detach().cpu() unsigned = torch.tensor( [ - int(packed[index // 2]) & 15 - if index % 2 == 0 - else (int(packed[index // 2]) >> 4) & 15 + int(packed[index // 2]) & 15 if index % 2 == 0 else (int(packed[index // 2]) >> 4) & 15 for index in range(count) ], dtype=torch.int16, ) - signed = unsigned.sub(8).float().reshape( - layer.out_features, - layer.num_groups, - layer.group_size, + signed = ( + unsigned.sub(8) + .float() + .reshape( + layer.out_features, + layer.num_groups, + layer.group_size, + ) ) weight = signed * layer.scales.detach().cpu().float()[..., None] return weight.reshape(layer.out_features, -1)[:, : layer.in_features].to(torch.bfloat16) @@ -158,34 +160,54 @@ def test_int4_rtn_rejects_non_positive_group_size(): assert quantized.packed_weight.dtype == torch.uint8 -def test_int4_rtn_cuda_quantize_path_matches_reference(): - if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: - pytest.skip("CUDA/Triton backend is not available") +def _active_triton_device_and_backend() -> tuple[torch.device, str] | None: + if torch.cuda.is_available(): + backend = "triton_rocm" if getattr(torch.version, "hip", None) else "triton_cuda" + return torch.device("cuda"), backend + xpu = getattr(torch, "xpu", None) + if xpu is not None and xpu.is_available(): + return torch.device("xpu"), "triton_xpu" + return None + + +def test_int4_rtn_triton_quantize_path_matches_reference(): + active = _active_triton_device_and_backend() + if active is None: + pytest.skip("no Triton accelerator is available") + device, backend = active + if not available_backends().get(backend, False): + pytest.skip(f"{backend} backend is not available") from orbitquant.kernels.triton_cuda import quantize_adaln_weight_with_triton torch.manual_seed(1) - weight = torch.randn(7, 19, device="cuda", dtype=torch.float32) + weight_cpu = torch.randn(7, 19, dtype=torch.float32) + weight = weight_cpu.to(device) expected_packed, expected_scales = _quantize_adaln_weight_reference( - weight.cpu(), group_size=8 + weight_cpu, + group_size=8, ) actual_packed, actual_scales = quantize_adaln_weight_with_triton(weight, group_size=8) - assert actual_packed.is_cuda - assert actual_scales.is_cuda + assert actual_packed.device.type == device.type + assert actual_scales.device.type == device.type assert torch.equal(actual_packed.cpu(), expected_packed) assert torch.allclose(actual_scales.cpu(), expected_scales) -def test_int4_rtn_cuda_dequant_path_matches_reference(): - if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: - pytest.skip("CUDA/Triton backend is not available") +def test_int4_rtn_triton_dequant_path_matches_reference(): + active = _active_triton_device_and_backend() + if active is None: + pytest.skip("no Triton accelerator is available") + device, backend = active + if not available_backends().get(backend, False): + pytest.skip(f"{backend} backend is not available") from orbitquant.kernels.triton_cuda import dequantize_adaln_weight_with_triton torch.manual_seed(2) - source = torch.nn.Linear(19, 7).to("cuda") + source = torch.nn.Linear(19, 7).to(device) config = OrbitQuantConfig(adaln_group_size=8) quantized = RTNInt4Linear.from_linear(source, config=config, module_name="block.modulation") @@ -197,8 +219,42 @@ def test_int4_rtn_cuda_dequant_path_matches_reference(): out_features=quantized.out_features, in_features=quantized.in_features, group_size=quantized.group_size, - device="cuda", + device=device, ) - assert actual.is_cuda + assert actual.device.type == device.type assert torch.allclose(actual.cpu(), expected) + + +def test_int4_rtn_triton_packed_forward_matches_independent_reference_without_dense_cache(): + active = _active_triton_device_and_backend() + if active is None: + pytest.skip("no Triton accelerator is available") + device, backend = active + if not available_backends().get(backend, False): + pytest.skip(f"{backend} backend is not available") + + torch.manual_seed(4) + source = torch.nn.Linear(65, 33).to(device) + x = torch.randn(3, 5, 65, dtype=torch.float32).to(device) + config = OrbitQuantConfig( + adaln_group_size=8, + runtime_mode="triton_packed_matmul", + ) + quantized = RTNInt4Linear.from_linear( + source, + config=config, + module_name="block.modulation", + ) + expected = F.linear( + x.cpu().to(torch.bfloat16), + _independent_packed_adaln_weight(quantized), + source.bias.detach().cpu().to(torch.bfloat16), + ) + + actual = quantized(x) + + torch.testing.assert_close(actual.cpu(), expected, atol=2e-2, rtol=2e-2) + assert quantized.last_effective_runtime_mode == "triton_packed_adaln_int4" + assert quantized._dequantized_weight_cache is None + assert quantized._dequantized_weight_cache_key is None diff --git a/tests/test_artifact_writer.py b/tests/test_artifact_writer.py index 82bdb48..86e300d 100644 --- a/tests/test_artifact_writer.py +++ b/tests/test_artifact_writer.py @@ -24,6 +24,10 @@ write_sha256sums_from_manifest, ) from orbitquant.config import OrbitQuantConfig +from orbitquant.kernels.executorch_vulkan import ( + ExecuTorchVulkanW4A4Linear, + prepare_executorch_vulkan_w4a4_model, +) from orbitquant.layers import OrbitQuantLinear from orbitquant.modeling import inspect_linear_module_policy, quantize_linear_modules @@ -341,6 +345,36 @@ def test_load_orbitquant_artifact_restores_quantized_modules_into_matching_model assert restored_state[packed_key].dtype == torch.uint8 +def test_loaded_artifact_can_be_prepared_for_executorch_vulkan_without_requantizing( + tmp_path, +): + torch.manual_seed(11) + source = TinyArtifactModel() + config = OrbitQuantConfig(block_size=8, runtime_mode="dequant_bf16") + summary = quantize_linear_modules(source, config) + save_orbitquant_artifact( + source, + tmp_path, + config=config, + source_model_id="example/model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + restored = TinyArtifactModel() + load_orbitquant_artifact(restored, tmp_path) + loaded_layer = restored.transformer_blocks[0]["attn"]["to_q"] + x = torch.randn(2, 3, 8) + expected = loaded_layer(x) + + prepared = prepare_executorch_vulkan_w4a4_model(restored) + prepared_layer = prepared.transformer_blocks[0]["attn"]["to_q"] + + assert prepared is restored + assert isinstance(prepared_layer, ExecuTorchVulkanW4A4Linear) + torch.testing.assert_close(prepared_layer(x), expected, atol=1e-6, rtol=1e-6) + + def test_validate_orbitquant_artifact_reports_eval_ready_required_files(tmp_path): source = TinyArtifactModel() config = OrbitQuantConfig(block_size=4, adaln_group_size=32) diff --git a/tests/test_backend_capabilities_golden.py b/tests/test_backend_capabilities_golden.py new file mode 100644 index 0000000..f2c051d --- /dev/null +++ b/tests/test_backend_capabilities_golden.py @@ -0,0 +1,1286 @@ +"""Golden snapshot tests for backend_capabilities(). + +The EXPECTED literals below were captured from the hand-built implementation +of orbitquant.kernels.dispatch.backend_capabilities prior to its table-driven +refactor. They freeze the full reported capability payload (values and key +order) for a representative grid of availability combinations. Only regenerate +them when the reported capabilities are intentionally changed. +""" + +import pytest + +import orbitquant.kernels.dispatch as dispatch_module +from orbitquant.kernels.dispatch import backend_capabilities + +# Frozen copies of long space-free payload strings (deliberately NOT imported +# from orbitquant.kernels.dispatch, so the snapshot stays independent). +_CPU_STAGES = "activation_norm_rpbh_quant_rescale,packed_weight_matmul,adaln_rtn_packed_matmul" +_MPS_STAGES = "activation_norm_rpbh_quant_rescale,packed_weight_dequant,packed_weight_matmul" +_TRITON_STAGES = ( + "activation_norm_rpbh_quant_rescale,packed_weight_dequant,packed_weight_matmul," + "lowbit_pack,lowbit_unpack,weight_rotation_fwht_quant_pack,adaln_rtn_quant_pack," + "adaln_rtn_dequant,adaln_rtn_packed_matmul" +) +_CPU_IMPL_NATIVE_ALL = "native_exact_activation+native_exact_packed_matmul+native_packed_adaln_int4" +_CPU_IMPL_MATMUL_ADALN = ( + "native_exact_packed_matmul+torch_activation_reference+native_packed_adaln_int4" +) + +# Each case: (availability probes monkeypatched onto dispatch_module, +# `backends` argument passed to backend_capabilities). +CASES: dict[str, tuple[dict[str, bool], dict[str, bool]]] = { + "all_false": ( + {'_native_cpu_packed_matmul_available': False, + '_native_cpu_activation_available': False, + '_native_cpu_adaln_available': False, + '_mps_metal_available': False, + '_native_packed_matmul_available': False}, + {'cpu': False, + 'mps': False, + 'triton_cuda': False, + 'triton_rocm': False, + 'triton_xpu': False}, + ), + "cpu_native_all": ( + {'_native_cpu_packed_matmul_available': True, + '_native_cpu_activation_available': True, + '_native_cpu_adaln_available': True, + '_mps_metal_available': False, + '_native_packed_matmul_available': False}, + {'cpu': True, + 'mps': False, + 'triton_cuda': False, + 'triton_rocm': False, + 'triton_xpu': False}, + ), + "cpu_activation_only": ( + {'_native_cpu_packed_matmul_available': False, + '_native_cpu_activation_available': True, + '_native_cpu_adaln_available': False, + '_mps_metal_available': False, + '_native_packed_matmul_available': False}, + {'cpu': True, + 'mps': False, + 'triton_cuda': False, + 'triton_rocm': False, + 'triton_xpu': False}, + ), + "cpu_adaln_only": ( + {'_native_cpu_packed_matmul_available': False, + '_native_cpu_activation_available': False, + '_native_cpu_adaln_available': True, + '_mps_metal_available': False, + '_native_packed_matmul_available': False}, + {'cpu': True, + 'mps': False, + 'triton_cuda': False, + 'triton_rocm': False, + 'triton_xpu': False}, + ), + "mps_shader_and_native_matmul": ( + {'_native_cpu_packed_matmul_available': False, + '_native_cpu_activation_available': False, + '_native_cpu_adaln_available': False, + '_mps_metal_available': True, + '_native_packed_matmul_available': True}, + {'cpu': True, + 'mps': True, + 'triton_cuda': False, + 'triton_rocm': False, + 'triton_xpu': False}, + ), + "mps_native_matmul_only": ( + {'_native_cpu_packed_matmul_available': False, + '_native_cpu_activation_available': False, + '_native_cpu_adaln_available': False, + '_mps_metal_available': False, + '_native_packed_matmul_available': True}, + {'cpu': True, + 'mps': True, + 'triton_cuda': False, + 'triton_rocm': False, + 'triton_xpu': False}, + ), + "triton_cuda": ( + {'_native_cpu_packed_matmul_available': False, + '_native_cpu_activation_available': False, + '_native_cpu_adaln_available': False, + '_mps_metal_available': False, + '_native_packed_matmul_available': False}, + {'cpu': True, + 'mps': False, + 'triton_cuda': True, + 'triton_rocm': False, + 'triton_xpu': False}, + ), + "triton_rocm": ( + {'_native_cpu_packed_matmul_available': False, + '_native_cpu_activation_available': False, + '_native_cpu_adaln_available': False, + '_mps_metal_available': False, + '_native_packed_matmul_available': False}, + {'cpu': True, + 'mps': False, + 'triton_cuda': False, + 'triton_rocm': True, + 'triton_xpu': False}, + ), + "triton_xpu": ( + {'_native_cpu_packed_matmul_available': False, + '_native_cpu_activation_available': False, + '_native_cpu_adaln_available': False, + '_mps_metal_available': False, + '_native_packed_matmul_available': False}, + {'cpu': True, + 'mps': False, + 'triton_cuda': False, + 'triton_rocm': False, + 'triton_xpu': True}, + ), + "mixed": ( + {'_native_cpu_packed_matmul_available': True, + '_native_cpu_activation_available': False, + '_native_cpu_adaln_available': True, + '_mps_metal_available': True, + '_native_packed_matmul_available': False}, + {'cpu': True, 'mps': True, 'triton_cuda': True, 'triton_rocm': False, 'triton_xpu': True}, + ), +} + +EXPECTED: dict[str, dict[str, dict[str, object]]] = { + "all_false": { + "cpu": {'available': False, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cpu'], + 'implementation': 'torch_reference', + 'package_format': 'torch_reference', + 'hf_kernel_builder_compliant': False, + 'notes': 'Activation quantization uses the reference PyTorch CPU path. The main ' + 'linear matmul uses the reference PyTorch CPU path. AdaLN uses the ' + 'explicit BF16 dequantization path.'}, + "mps": {'available': False, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_reference_mps', + 'package_format': 'torch_reference', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Runs the reference PyTorch path on MPS tensors; native Metal shader ' + 'support and the native packed matmul package are not available in this ' + 'environment.'}, + "triton_cuda": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "cpu_native_all": { + "cpu": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': _CPU_STAGES, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': True, + 'adaln_packed_matmul_optimized': True, + 'device_types': ['cpu'], + 'implementation': _CPU_IMPL_NATIVE_ALL, + 'package_format': 'native_kernel_package_torch_stable_abi', + 'hf_kernel_builder_compliant': True, + 'notes': 'The native exact activation pipeline performs FP32 token norm, ' + 'RPBH/FWHT, codebook assignment, and rescale. Native exact packed ' + 'matmul consumes low-bit weights without a full floating-point cache. ' + 'Native packed INT4 AdaLN matmul consumes group-64 weights without a ' + 'full floating-point cache.'}, + "mps": {'available': False, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_reference_mps', + 'package_format': 'torch_reference', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Runs the reference PyTorch path on MPS tensors; native Metal shader ' + 'support and the native packed matmul package are not available in this ' + 'environment.'}, + "triton_cuda": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "cpu_activation_only": { + "cpu": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': 'activation_norm_rpbh_quant_rescale', + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cpu'], + 'implementation': 'native_exact_activation+reference_weight_matmul', + 'package_format': 'native_kernel_package_torch_stable_abi', + 'hf_kernel_builder_compliant': True, + 'notes': 'The native exact activation pipeline performs FP32 token norm, ' + 'RPBH/FWHT, codebook assignment, and rescale. The main linear matmul ' + 'uses the reference PyTorch CPU path. AdaLN uses the explicit BF16 ' + 'dequantization path.'}, + "mps": {'available': False, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_reference_mps', + 'package_format': 'torch_reference', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Runs the reference PyTorch path on MPS tensors; native Metal shader ' + 'support and the native packed matmul package are not available in this ' + 'environment.'}, + "triton_cuda": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "cpu_adaln_only": { + "cpu": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': 'adaln_rtn_packed_matmul', + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': True, + 'adaln_packed_matmul_optimized': True, + 'device_types': ['cpu'], + 'implementation': 'torch_reference+native_packed_adaln_int4', + 'package_format': 'native_kernel_package_torch_stable_abi', + 'hf_kernel_builder_compliant': True, + 'notes': 'Activation quantization uses the reference PyTorch CPU path. The main ' + 'linear matmul uses the reference PyTorch CPU path. Native packed INT4 ' + 'AdaLN matmul consumes group-64 weights without a full floating-point ' + 'cache.'}, + "mps": {'available': False, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_reference_mps', + 'package_format': 'torch_reference', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Runs the reference PyTorch path on MPS tensors; native Metal shader ' + 'support and the native packed matmul package are not available in this ' + 'environment.'}, + "triton_cuda": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "mps_shader_and_native_matmul": { + "cpu": {'available': True, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cpu'], + 'implementation': 'torch_reference', + 'package_format': 'torch_reference', + 'hf_kernel_builder_compliant': False, + 'notes': 'Activation quantization uses the reference PyTorch CPU path. The main ' + 'linear matmul uses the reference PyTorch CPU path. AdaLN uses the ' + 'explicit BF16 dequantization path.'}, + "mps": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': _MPS_STAGES, + 'weight_dequant_optimized': True, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_mps_compile_shader_fused_activation+native_packed_matmul', + 'package_format': 'torch.mps.compile_shader,native_kernel_package', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Fused Metal shaders handle activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, and packed weight dequant. The native packed ' + 'matmul package handles packed low-bit matmul.'}, + "triton_cuda": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "mps_native_matmul_only": { + "cpu": {'available': True, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cpu'], + 'implementation': 'torch_reference', + 'package_format': 'torch_reference', + 'hf_kernel_builder_compliant': False, + 'notes': 'Activation quantization uses the reference PyTorch CPU path. The main ' + 'linear matmul uses the reference PyTorch CPU path. AdaLN uses the ' + 'explicit BF16 dequantization path.'}, + "mps": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': 'packed_weight_matmul', + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'native_packed_matmul', + 'package_format': 'native_kernel_package', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'The native packed matmul package handles packed low-bit matmul. ' + 'Activation quantization helpers use the reference PyTorch path in this ' + 'environment.'}, + "triton_cuda": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "triton_cuda": { + "cpu": {'available': True, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cpu'], + 'implementation': 'torch_reference', + 'package_format': 'torch_reference', + 'hf_kernel_builder_compliant': False, + 'notes': 'Activation quantization uses the reference PyTorch CPU path. The main ' + 'linear matmul uses the reference PyTorch CPU path. AdaLN uses the ' + 'explicit BF16 dequantization path.'}, + "mps": {'available': False, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_reference_mps', + 'package_format': 'torch_reference', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Runs the reference PyTorch path on MPS tensors; native Metal shader ' + 'support and the native packed matmul package are not available in this ' + 'environment.'}, + "triton_cuda": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': _TRITON_STAGES, + 'weight_dequant_optimized': True, + 'weight_pack_optimized': True, + 'lowbit_unpack_optimized': True, + 'weight_quant_optimized': True, + 'adaln_quant_optimized': True, + 'adaln_dequant_optimized': True, + 'adaln_packed_matmul_optimized': True, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "triton_rocm": { + "cpu": {'available': True, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cpu'], + 'implementation': 'torch_reference', + 'package_format': 'torch_reference', + 'hf_kernel_builder_compliant': False, + 'notes': 'Activation quantization uses the reference PyTorch CPU path. The main ' + 'linear matmul uses the reference PyTorch CPU path. AdaLN uses the ' + 'explicit BF16 dequantization path.'}, + "mps": {'available': False, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_reference_mps', + 'package_format': 'torch_reference', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Runs the reference PyTorch path on MPS tensors; native Metal shader ' + 'support and the native packed matmul package are not available in this ' + 'environment.'}, + "triton_cuda": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': True, + 'claim_status': 'experimental_unverified', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "triton_xpu": { + "cpu": {'available': True, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cpu'], + 'implementation': 'torch_reference', + 'package_format': 'torch_reference', + 'hf_kernel_builder_compliant': False, + 'notes': 'Activation quantization uses the reference PyTorch CPU path. The main ' + 'linear matmul uses the reference PyTorch CPU path. AdaLN uses the ' + 'explicit BF16 dequantization path.'}, + "mps": {'available': False, + 'claim_status': 'reference_only', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_reference_mps', + 'package_format': 'torch_reference', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Runs the reference PyTorch path on MPS tensors; native Metal shader ' + 'support and the native packed matmul package are not available in this ' + 'environment.'}, + "triton_cuda": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': True, + 'claim_status': 'experimental_unverified', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, + "mixed": { + "cpu": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _CPU_STAGES, + 'optimized_stage': 'packed_weight_matmul,adaln_rtn_packed_matmul', + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': True, + 'adaln_packed_matmul_optimized': True, + 'device_types': ['cpu'], + 'implementation': _CPU_IMPL_MATMUL_ADALN, + 'package_format': 'native_kernel_package_torch_stable_abi', + 'hf_kernel_builder_compliant': True, + 'notes': 'Activation quantization uses the reference PyTorch CPU path. Native ' + 'exact packed matmul consumes low-bit weights without a full ' + 'floating-point cache. Native packed INT4 AdaLN matmul consumes ' + 'group-64 weights without a full floating-point cache.'}, + "mps": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _MPS_STAGES, + 'optimized_stage': 'activation_norm_rpbh_quant_rescale,packed_weight_dequant', + 'weight_dequant_optimized': True, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['mps'], + 'implementation': 'torch_mps_compile_shader_fused_activation', + 'package_format': 'torch.mps.compile_shader', + 'upstream_native_mps_op': False, + 'hf_kernel_builder_compliant': False, + 'notes': 'Fused Metal shaders handle activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, and packed weight dequant. The native packed ' + 'matmul package is not available in this environment.'}, + "triton_cuda": {'available': True, + 'claim_status': 'partial_optimized', + 'optimized': True, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': _TRITON_STAGES, + 'weight_dequant_optimized': True, + 'weight_pack_optimized': True, + 'lowbit_unpack_optimized': True, + 'weight_quant_optimized': True, + 'adaln_quant_optimized': True, + 'adaln_dequant_optimized': True, + 'adaln_packed_matmul_optimized': True, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline', + 'package_format': 'python_triton', + 'hf_kernel_builder_compliant': False, + 'notes': 'Triton handles runtime activation norm, RPBH/FWHT rotation, ' + 'codebook lookup/rescale, packed weight dequant, packed weight ' + 'matmul, offline low-bit pack/unpack, offline weight RPBH/FWHT ' + 'codebook indexing with direct low-bit packing, and AdaLN INT4 ' + 'RTN quantize/pack, dequant, and packed matmul. The default ' + 'auto_fused runtime prefers packed low-bit matmul when a native ' + 'or Triton packed kernel is available; full-model speedup ' + 'claims still require separate benchmark artifacts.'}, + "triton_rocm": {'available': False, + 'claim_status': 'unavailable', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['cuda'], + 'implementation': 'python_triton_orbitquant_pipeline_rocm_candidate', + 'package_format': 'python_triton_rocm', + 'hf_kernel_builder_compliant': False, + 'notes': 'PyTorch exposes HIP tensors through the cuda device type. The ' + 'candidate reuses the exact packed Triton pipeline without ' + 'loading CUDA-native extensions. It remains experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported AMD hardware.'}, + "triton_xpu": {'available': True, + 'claim_status': 'experimental_unverified', + 'optimized': False, + 'full_fusion': False, + 'implemented_stage': _TRITON_STAGES, + 'optimized_stage': None, + 'weight_dequant_optimized': False, + 'weight_pack_optimized': False, + 'lowbit_unpack_optimized': False, + 'weight_quant_optimized': False, + 'adaln_quant_optimized': False, + 'adaln_dequant_optimized': False, + 'adaln_packed_matmul_optimized': False, + 'device_types': ['xpu'], + 'implementation': 'python_triton_orbitquant_pipeline_xpu_candidate', + 'package_format': 'python_triton_xpu', + 'hf_kernel_builder_compliant': False, + 'notes': 'The candidate reuses the exact packed Triton pipeline on ' + 'torch.xpu. It remains explicit-only and experimental until ' + 'correctness, memory, profiler, and performance evidence is ' + 'recorded on supported Intel GPU hardware.'}, + }, +} + + +@pytest.mark.parametrize("case_id", list(CASES)) +def test_backend_capabilities_matches_golden_snapshot( + case_id: str, monkeypatch: pytest.MonkeyPatch +) -> None: + probe_overrides, backend_flags = CASES[case_id] + for attr, value in probe_overrides.items(): + monkeypatch.setattr(dispatch_module, attr, lambda value=value: value) + result = backend_capabilities(backends=dict(backend_flags)) + expected = EXPECTED[case_id] + assert result == expected + # Dict equality ignores insertion order; freeze the key order as well. + assert list(result) == list(expected) + for backend_name, entry in result.items(): + assert list(entry) == list(expected[backend_name]) + + +def test_golden_grid_is_consistent() -> None: + assert set(CASES) == set(EXPECTED) diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py index 567da94..062229c 100644 --- a/tests/test_benchmarks.py +++ b/tests/test_benchmarks.py @@ -40,6 +40,9 @@ def test_benchmark_orbit_linear_reports_stage_timings_on_cpu(): assert result["timings_ms"][key] >= 0.0 assert result["quantization_buffers"]["packed_weight_indices_device"] == "cpu" assert result["quantization_buffers"]["row_norms_device"] == "cpu" + assert result["quantization_buffers"]["source_weight_device_type"] == "cpu" + assert result["quantization_buffers"]["packed_weight_indices_device_type"] == "cpu" + assert result["quantization_buffers"]["row_norms_device_type"] == "cpu" assert result["quantization_buffers"]["packed_weight_indices_is_cuda"] is False assert result["quantization_buffers"]["row_norms_is_cuda"] is False @@ -119,3 +122,24 @@ def fake_forward(self, x): def test_benchmark_weight_backend_label_separates_reference_mps_from_cpu(): assert _weight_quantization_backend_label(torch.device("cpu")) == "torch_reference" assert _weight_quantization_backend_label(torch.device("mps")) == "torch_reference_mps" + + +def test_benchmark_weight_backend_label_reports_rocm_on_hip(monkeypatch): + monkeypatch.setattr(benchmarks_module.torch.version, "hip", "7.2", raising=False) + monkeypatch.setattr( + benchmarks_module, + "backend_capabilities", + lambda: {"triton_rocm": {"available": True}}, + ) + + assert _weight_quantization_backend_label(torch.device("cuda")) == "triton_rocm" + + +def test_benchmark_weight_backend_label_reports_explicit_xpu(monkeypatch): + monkeypatch.setattr( + benchmarks_module, + "backend_capabilities", + lambda: {"triton_xpu": {"available": True}}, + ) + + assert _weight_quantization_backend_label(torch.device("xpu")) == "triton_xpu" diff --git a/tests/test_cli.py b/tests/test_cli.py deleted file mode 100644 index b5ea201..0000000 --- a/tests/test_cli.py +++ /dev/null @@ -1,4050 +0,0 @@ -import json -import sys -from types import SimpleNamespace - -import numpy as np -import pytest -import torch -from PIL import Image - -import orbitquant.cli.main as cli_main -from orbitquant import __version__ -from orbitquant.artifacts import save_orbitquant_artifact, validate_orbitquant_artifact -from orbitquant.artifacts.checksums import write_sha256sums_from_manifest -from orbitquant.cli.main import main -from orbitquant.config import OrbitQuantConfig -from orbitquant.eval import get_native_suite -from orbitquant.eval.prompts import default_prompt_payload, select_prompt_record -from orbitquant.modeling import quantize_linear_modules - - -def test_cli_version_prints_version(capsys): - assert main(["--version"]) == 0 - - output = capsys.readouterr().out - assert __version__ in output - - -def test_cli_packed_matmul_runtime_modes_skip_dequant_prewarm(): - assert cli_main._should_prewarm_quantized_weights(None) is True - assert cli_main._should_prewarm_quantized_weights(OrbitQuantConfig()) is False - assert ( - cli_main._should_prewarm_quantized_weights( - OrbitQuantConfig(runtime_mode="dequant_bf16") - ) - is True - ) - assert ( - cli_main._should_prewarm_quantized_weights( - OrbitQuantConfig(runtime_mode="triton_packed_matmul") - ) - is False - ) - assert ( - cli_main._should_prewarm_quantized_weights( - OrbitQuantConfig(runtime_mode="native_packed_matmul") - ) - is False - ) - - -def test_cli_generation_placement_uses_model_cpu_offload_when_requested(): - class TinyPipeline: - def __init__(self): - self.calls = [] - - def to(self, device): - self.calls.append(("to", device)) - return self - - def enable_model_cpu_offload(self, *, device): - self.calls.append(("offload", device)) - - pipeline = TinyPipeline() - cli_main._place_pipeline_for_generation( - pipeline, - device="cuda", - enable_model_cpu_offload=True, - ) - - assert pipeline.calls == [("offload", "cuda")] - - -def test_cli_generation_placement_fails_loud_when_offload_is_missing(): - class TinyPipeline: - def to(self, device): - return self - - with pytest.raises(RuntimeError, match="enable_model_cpu_offload"): - cli_main._place_pipeline_for_generation( - TinyPipeline(), - device="cuda", - enable_model_cpu_offload=True, - ) - - -def test_cli_native_suites_lists_no_range_smoke_settings(capsys): - assert main(["native-suites"]) == 0 - - output = capsys.readouterr().out - assert "flux2-native" in output - assert "wan-native" in output - assert "range" not in output.lower() - - -def test_cli_kernel_info_reports_backend_capabilities(capsys, monkeypatch): - monkeypatch.setattr( - cli_main, - "backend_capabilities", - lambda: { - "cpu": { - "available": True, - "claim_status": "reference_only", - "optimized": False, - "implemented_stage": None, - "optimized_stage": None, - "weight_dequant_optimized": False, - "weight_pack_optimized": False, - "weight_quant_optimized": False, - "adaln_quant_optimized": False, - "adaln_dequant_optimized": False, - "hf_kernel_builder_compliant": False, - }, - "mps": { - "claim_status": "partial_optimized", - "implementation": "torch_mps_compile_shader_codebook_rescale+native_packed_matmul", - "package_format": "torch.mps.compile_shader,native_kernel_package", - "implemented_stage": ( - "codebook_lookup_rescale,packed_weight_dequant,packed_weight_matmul" - ), - "optimized_stage": ( - "codebook_lookup_rescale,packed_weight_dequant,packed_weight_matmul" - ), - "weight_dequant_optimized": True, - "weight_pack_optimized": False, - "weight_quant_optimized": False, - "adaln_quant_optimized": False, - "adaln_dequant_optimized": False, - "full_fusion": False, - "upstream_native_mps_op": False, - "hf_kernel_builder_compliant": False, - }, - "triton_cuda": { - "claim_status": "partial_optimized", - "implemented_stage": ( - "activation_norm_rpbh_quant_rescale,packed_weight_dequant," - "packed_weight_matmul,lowbit_pack,lowbit_unpack," - "weight_rotation_fwht_quant_pack," - "adaln_rtn_quant_pack,adaln_rtn_dequant" - ), - "optimized_stage": ( - "activation_norm_rpbh_quant_rescale,packed_weight_dequant," - "packed_weight_matmul,lowbit_pack,lowbit_unpack," - "weight_rotation_fwht_quant_pack," - "adaln_rtn_quant_pack,adaln_rtn_dequant" - ), - "implementation": "python_triton_orbitquant_pipeline", - "package_format": "python_triton", - "weight_dequant_optimized": True, - "weight_pack_optimized": True, - "weight_quant_optimized": True, - "adaln_quant_optimized": True, - "adaln_dequant_optimized": True, - "full_fusion": False, - "hf_kernel_builder_compliant": False, - }, - }, - ) - - assert main(["kernel-info"]) == 0 - - payload = json.loads(capsys.readouterr().out) - assert payload["cpu"]["available"] is True - assert payload["cpu"]["claim_status"] == "reference_only" - assert payload["cpu"]["optimized"] is False - assert payload["cpu"]["weight_dequant_optimized"] is False - assert payload["cpu"]["weight_pack_optimized"] is False - assert payload["cpu"]["weight_quant_optimized"] is False - assert payload["cpu"]["adaln_quant_optimized"] is False - assert payload["cpu"]["adaln_dequant_optimized"] is False - assert ( - payload["mps"]["implementation"] - == "torch_mps_compile_shader_codebook_rescale+native_packed_matmul" - ) - assert payload["mps"]["package_format"] == "torch.mps.compile_shader,native_kernel_package" - assert payload["mps"]["claim_status"] == "partial_optimized" - assert ( - payload["mps"]["implemented_stage"] - == "codebook_lookup_rescale,packed_weight_dequant,packed_weight_matmul" - ) - assert ( - payload["mps"]["optimized_stage"] - == "codebook_lookup_rescale,packed_weight_dequant,packed_weight_matmul" - ) - assert payload["mps"]["weight_dequant_optimized"] is True - assert payload["mps"]["weight_pack_optimized"] is False - assert payload["mps"]["weight_quant_optimized"] is False - assert payload["mps"]["adaln_quant_optimized"] is False - assert payload["mps"]["adaln_dequant_optimized"] is False - assert payload["mps"]["full_fusion"] is False - assert payload["mps"]["upstream_native_mps_op"] is False - assert payload["mps"]["hf_kernel_builder_compliant"] is False - expected_triton_stage = ( - "activation_norm_rpbh_quant_rescale,packed_weight_dequant," - "packed_weight_matmul,lowbit_pack,lowbit_unpack," - "weight_rotation_fwht_quant_pack," - "adaln_rtn_quant_pack,adaln_rtn_dequant" - ) - assert payload["triton_cuda"]["implemented_stage"] == expected_triton_stage - assert payload["triton_cuda"]["optimized_stage"] == expected_triton_stage - assert payload["triton_cuda"]["implementation"] == "python_triton_orbitquant_pipeline" - assert payload["triton_cuda"]["package_format"] == "python_triton" - assert payload["triton_cuda"]["claim_status"] == "partial_optimized" - assert payload["triton_cuda"]["weight_dequant_optimized"] is True - assert payload["triton_cuda"]["weight_pack_optimized"] is True - assert payload["triton_cuda"]["weight_quant_optimized"] is True - assert payload["triton_cuda"]["adaln_quant_optimized"] is True - assert payload["triton_cuda"]["adaln_dequant_optimized"] is True - assert payload["triton_cuda"]["full_fusion"] is False - assert payload["triton_cuda"]["hf_kernel_builder_compliant"] is False - - -def test_cli_kernel_bench_prints_stage_timings(capsys): - assert ( - main( - [ - "kernel-bench", - "--tokens", - "4", - "--in-features", - "16", - "--out-features", - "8", - "--block-size", - "8", - "--activation-kernel-backend", - "cpu", - "--runtime-mode", - "debug_no_activation_quant", - "--device", - "cpu", - "--dtype", - "float32", - "--warmup", - "0", - "--iterations", - "1", - ] - ) - == 0 - ) - - payload = json.loads(capsys.readouterr().out) - assert payload["device"] == "cpu" - assert payload["runtime_mode"] == "debug_no_activation_quant" - assert payload["full_fusion"] is False - assert payload["prewarm"]["total_modules"] == 1 - assert payload["timings_ms"]["weight_quantize_pack_cold_ms"] >= 0.0 - assert payload["timings_ms"]["weight_quantize_pack_hot_ms"] >= 0.0 - assert payload["timings_ms"]["forward_prewarmed_ms"] >= 0.0 - assert payload["selected_activation_kernel_backend"] == "cpu" - assert payload["weight_quantization_backend"] == "torch_reference" - assert payload["quantization_buffers"]["source_weight_device"] == "cpu" - assert payload["quantization_buffers"]["source_weight_is_cuda"] is False - assert payload["quantization_buffers"]["packed_weight_indices_device"] == "cpu" - - -@pytest.mark.parametrize("runtime_mode", ["triton_packed_matmul", "native_packed_matmul"]) -def test_cli_kernel_bench_passes_packed_matmul_tile_options( - monkeypatch, capsys, runtime_mode -): - seen_kwargs = [] - - def fake_benchmark_orbit_linear(**kwargs): - seen_kwargs.append(kwargs) - return {"ok": True} - - monkeypatch.setattr(cli_main, "benchmark_orbit_linear", fake_benchmark_orbit_linear) - - assert ( - main( - [ - "kernel-bench", - "--runtime-mode", - runtime_mode, - "--packed-matmul-block-m", - "32", - "--packed-matmul-block-n", - "64", - "--packed-matmul-block-k", - "64", - "--packed-matmul-num-warps", - "8", - ] - ) - == 0 - ) - - assert json.loads(capsys.readouterr().out) == {"ok": True} - assert seen_kwargs[0]["runtime_mode"] == runtime_mode - assert seen_kwargs[0]["packed_matmul_block_m"] == 32 - assert seen_kwargs[0]["packed_matmul_block_n"] == 64 - assert seen_kwargs[0]["packed_matmul_block_k"] == 64 - assert seen_kwargs[0]["packed_matmul_num_warps"] == 8 - - -def test_cli_quantize_bench_prints_full_model_staging_timings(capsys): - assert ( - main( - [ - "quantize-bench", - "--layers", - "1", - "--in-features", - "16", - "--hidden-features", - "32", - "--block-size", - "8", - "--source-device", - "cpu", - "--quantization-device", - "cpu", - "--staging-mode", - "component", - "--dtype", - "float32", - ] - ) - == 0 - ) - - payload = json.loads(capsys.readouterr().out) - assert payload["source_device"] == "cpu" - assert payload["quantization_device"] == "cpu" - assert payload["staging_mode"] == "component" - assert payload["synchronize_per_module"] is False - assert payload["summary"]["quantization_staging_mode"] == "component" - assert payload["summary"]["synchronize_per_module"] is False - assert payload["summary"]["source_linear_device_counts"]["cpu"] == 7 - assert payload["summary"]["device_transfer_seconds"] >= 0.0 - assert payload["summary"]["quantized_modules"] - - -@pytest.mark.parametrize("runtime_mode", ["triton_packed_matmul", "native_packed_matmul"]) -def test_cli_native_plan_lists_full_target_bit_matrix_without_range_smoke( - capsys, tmp_path, runtime_mode -): - assert ( - main( - [ - "native-plan", - "--output-root", - str(tmp_path / "artifacts"), - "--seeds", - "0", - "--runtime-mode", - runtime_mode, - ] - ) - == 0 - ) - - payload = json.loads(capsys.readouterr().out) - jobs = payload["jobs"] - assert payload["job_count"] == 14 - assert {job["suite"] for job in jobs} == { - "flux2-native", - "flux1-schnell-native", - "z-image-native", - "wan-native", - } - assert sum(1 for job in jobs if job["suite"] == "wan-native") == 2 - assert {job["runtime_mode"] for job in jobs} == {runtime_mode} - wan_w4a6 = next( - job - for job in jobs - if job["suite"] == "wan-native" and job["bit_setting"] == "W4A6" - ) - flux1_w3a3 = next( - job - for job in jobs - if job["suite"] == "flux1-schnell-native" and job["bit_setting"] == "W3A3" - ) - assert wan_w4a6["width"] == 832 - assert wan_w4a6["height"] == 480 - assert wan_w4a6["frames"] == 81 - assert wan_w4a6["metric"] == "vbench" - assert flux1_w3a3["width"] == 1024 - assert flux1_w3a3["height"] == 1024 - assert flux1_w3a3["guidance"] == 0.0 - assert flux1_w3a3["artifact_dir"].endswith("flux1-schnell-native-w3a3") - assert "range" not in json.dumps(payload).lower() - - -def test_cli_native_plan_defaults_to_auto_fused_runtime(capsys, tmp_path): - assert ( - main( - [ - "native-plan", - "--output-root", - str(tmp_path / "artifacts"), - "--seeds", - "0", - ] - ) - == 0 - ) - - payload = json.loads(capsys.readouterr().out) - - assert payload["job_count"] == 14 - assert {job["runtime_mode"] for job in payload["jobs"]} == {"auto_fused"} - - -def test_cli_external_eval_plan_lists_metric_runner_import_commands(capsys, tmp_path): - assert ( - main( - [ - "external-eval-plan", - "--suite", - "wan-native", - "--output-root", - str(tmp_path / "artifacts"), - "--metrics-root", - str(tmp_path / "metrics"), - ] - ) - == 0 - ) - - payload = json.loads(capsys.readouterr().out) - assert payload["job_count"] == 4 - job = next( - item - for item in payload["jobs"] - if item["bit_setting"] == "W4A6" and item["split"] == "orbitquant" - ) - assert job["suite"] == "wan-native" - assert job["metric"] == "vbench" - assert job["artifact_dir"].endswith("wan-native-w4a6") - assert job["metrics_json"].endswith("wan-native-w4a6_orbitquant_vbench.json") - assert "orbitquant export-vbench" in job["export_command"] - assert "vbench evaluate" in job["eval_command"] - assert "orbitquant summarize-vbench-results" in job["summarize_command"] - assert "orbitquant record-metrics" in job["import_command"] - assert "--metric-prefix vbench" in job["import_command"] - assert "--split orbitquant" in job["import_command"] - assert "range" not in json.dumps(payload).lower() - - -def test_cli_external_eval_script_prints_metric_runner_script(capsys, tmp_path): - assert ( - main( - [ - "external-eval-script", - "--suite", - "wan-native", - "--output-root", - str(tmp_path / "artifacts"), - "--metrics-root", - str(tmp_path / "metrics"), - "--report-output", - str(tmp_path / "reports"), - ] - ) - == 0 - ) - - script = capsys.readouterr().out - assert script.startswith("#!/usr/bin/env bash\nset -euo pipefail\n") - assert "if ! command -v vbench >/dev/null 2>&1; then" in script - assert f"if [ ! -d {tmp_path / 'artifacts' / 'wan-native-w4a6'} ]; then" in script - assert "orbitquant export-vbench" in script - assert "vbench evaluate" in script - assert "orbitquant summarize-vbench-results" in script - assert "orbitquant record-metrics" in script - assert "--metric-prefix vbench" in script - assert "orbitquant report" in script - assert f"--output {tmp_path / 'reports'}" in script - assert "--fail-on-missing-required" in script - assert "range" not in script.lower() - - -def test_cli_external_eval_script_resume_skips_existing_metric_summaries(capsys, tmp_path): - assert ( - main( - [ - "external-eval-script", - "--suite", - "wan-native", - "--output-root", - str(tmp_path / "artifacts"), - "--metrics-root", - str(tmp_path / "metrics"), - "--report-output", - str(tmp_path / "reports"), - "--resume", - ] - ) - == 0 - ) - - script = capsys.readouterr().out - metrics_json = tmp_path / "metrics" / "wan-native-w4a4_original_vbench.json" - assert f"if [ -s {metrics_json} ]; then" in script - assert "stage_log SKIP 'wan-native W4A4 original vbench metrics'" in script - assert "fi\nstage_log START 'wan-native W4A4 original import vbench'" in script - assert script.count("stage_log SKIP") == 4 - assert script.count("orbitquant record-metrics") == 4 - assert "--fail-on-missing-required" in script - - -def test_cli_report_can_fail_on_missing_required_metrics(capsys, monkeypatch, tmp_path): - result = SimpleNamespace( - report_path=tmp_path / "report.md", - table_paths={}, - rows=[], - missing_required_metrics=[{"metric": "geneval_overall"}], - ) - monkeypatch.setattr(cli_main, "generate_native_eval_report", lambda *args, **kwargs: result) - - assert ( - main(["report", "--artifact", str(tmp_path / "artifact"), "--output", str(tmp_path)]) - == 0 - ) - assert json.loads(capsys.readouterr().out)["missing_required_metric_count"] == 1 - - assert ( - main( - [ - "report", - "--artifact", - str(tmp_path / "artifact"), - "--output", - str(tmp_path), - "--fail-on-missing-required", - ] - ) - == 1 - ) - assert json.loads(capsys.readouterr().out)["missing_required_metric_count"] == 1 - - -def test_cli_native_script_groups_quantize_and_generate_pack_commands(capsys, tmp_path): - assert ( - main( - [ - "native-script", - "--suite", - "wan-native", - "--output-root", - str(tmp_path / "artifacts"), - "--seeds", - "0,1", - "--prompt-limit", - "1", - "--device", - "cuda", - "--dtype", - "bfloat16", - "--activation-kernel-backend", - "triton_cuda", - "--runtime-mode", - "triton_packed_matmul", - ] - ) - == 0 - ) - - script = capsys.readouterr().out - assert script.startswith("#!/usr/bin/env bash\nset -euo pipefail\n") - assert "stage_log() {" in script - assert "stage_log START preflight" in script - assert "stage_log START 'kernel preflight'" in script - assert "stage_log START 'policy inventories'" in script - assert "stage_log START 'wan-native W4A6 quantize'" in script - assert "stage_log START 'wan-native W4A6 original generate-pack'" in script - assert "stage_log END 'native eval report'" in script - assert "hf auth whoami" in script - assert "hf env" in script - assert "torch.cuda.is_available()" in script - assert "shutil.disk_usage" in script - assert "hf models info Wan-AI/Wan2.1-T2V-1.3B-Diffusers --format json >/dev/null" in script - assert "orbitquant kernel-info" in script - assert script.count("orbitquant kernel-bench") == 2 - assert script.count("orbitquant inspect-policy") == 1 - assert "--output reports/native/module-inventories/wan-native-policy.json" in script - assert "--tokens 256 --in-features 3072 --out-features 3072" in script - assert script.count("orbitquant quantize") == 2 - assert "--suite wan-native" in script - assert "--weight-bits 4 --activation-bits 6" in script - assert "--weight-bits 4 --activation-bits 4" in script - assert "--activation-kernel-backend triton_cuda" in script - assert "--runtime-mode triton_packed_matmul" in script - assert ( - "--policy-inventory reports/native/module-inventories/wan-native-policy.json " - "--runtime-mode triton_packed_matmul" - ) in script - assert "--staging-mode component" in script - assert script.count("\norbitquant generate-pack") == 4 - assert "--split original" in script - assert "--split orbitquant" in script - assert "--seeds 0,1" in script - assert "--prompt-limit 1" in script - assert script.count("orbitquant validate-artifact") == 4 - assert ( - "--policy-inventory reports/native/module-inventories/wan-native-policy.json" - in script - ) - assert script.count("orbitquant report") == 1 - assert f"--artifact {tmp_path / 'artifacts' / 'wan-native-w4a6'}" in script - assert f"--artifact {tmp_path / 'artifacts' / 'wan-native-w4a4'}" in script - assert "--output reports/native" in script - assert "range" not in script.lower() - - -def test_cli_native_script_can_generate_geneval_metadata_pack(capsys, tmp_path): - metadata_jsonl = tmp_path / "evaluation_metadata.jsonl" - assert ( - main( - [ - "native-script", - "--suite", - "flux1-schnell-native", - "--output-root", - str(tmp_path / "artifacts"), - "--prompt-metadata-jsonl", - str(metadata_jsonl), - "--seeds", - "0", - "--prompt-limit", - "2", - ] - ) - == 0 - ) - - script = capsys.readouterr().out - assert "--prompt-metadata-jsonl" in script - assert str(metadata_jsonl) in script - assert "--prompt-limit 2" in script - assert "--prompt-pack" not in script - assert "stage_log START 'flux1-schnell-native W4A4 original generate-pack'" in script - assert "stage_log START 'flux1-schnell-native W4A4 orbitquant generate-pack'" in script - assert script.count("\norbitquant generate-pack") == 8 - - -def test_cli_native_script_resume_skips_valid_existing_artifacts(capsys, tmp_path): - assert ( - main( - [ - "native-script", - "--suite", - "flux2-native", - "--output-root", - str(tmp_path / "artifacts"), - "--seeds", - "0", - "--resume", - ] - ) - == 0 - ) - - script = capsys.readouterr().out - artifact_dir = tmp_path / "artifacts" / "flux2-native-w4a4" - inventory_path = "reports/native/module-inventories/flux2-native-policy.json" - assert ( - f"if orbitquant validate-artifact --artifact {artifact_dir} " - f"--policy-inventory {inventory_path} --runtime-mode auto_fused" - ) in script - assert f"echo 'Skipping existing valid artifact: {artifact_dir}'" in script - assert "else\norbitquant quantize --suite flux2-native" in script - assert "\nfi\nstage_log END 'flux2-native W4A4 quantize'" in script - assert "stage_log START 'flux2-native W4A4 validate quantized artifact'" in script - assert "\norbitquant generate-pack --suite flux2-native" in script - assert "--resume-existing" in script - - -def test_cli_generate_requires_prompt_and_output(): - try: - main(["generate", "--suite", "flux2-native"]) - except (SystemExit, ValueError) as exc: - assert "prompt" in str(exc) or "output" in str(exc) or getattr(exc, "code", 0) != 0 - else: - raise AssertionError("generate accepted missing prompt/output arguments") - - -def test_cli_record_metrics_imports_nested_json_into_artifact(capsys, tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="flux") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="black-forest-labs/FLUX.1-schnell", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - metrics_path = tmp_path / "geneval.json" - metrics_path.write_text( - json.dumps( - { - "overall": 0.71, - "per_task": { - "single_object": 0.92, - "counting": 0.47, - }, - "notes": "ignored", - } - ) - ) - - assert ( - main( - [ - "record-metrics", - "--artifact", - str(tmp_path), - "--split", - "orbitquant", - "--metrics-json", - str(metrics_path), - "--metric-prefix", - "geneval", - "--suite", - "flux1-schnell-native", - "--seed", - "0", - "--bit-setting", - "W4A4", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - record = json.loads((tmp_path / "benchmark" / "orbitquant.metrics.jsonl").read_text()) - csv_text = (tmp_path / "benchmark" / "orbitquant.metrics.csv").read_text() - manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) - assert output["metrics"] == { - "geneval_overall": 0.71, - "geneval_per_task_counting": 0.47, - "geneval_per_task_single_object": 0.92, - } - assert record["metrics"] == output["metrics"] - assert record["metadata"] == { - "suite": "flux1-schnell-native", - "seed": 0, - "bit_setting": "W4A4", - "metrics_source": str(metrics_path), - } - assert "geneval_per_task_single_object,0.92" in csv_text - assert "benchmark/orbitquant.metrics.jsonl" in manifest["checksums"] - - report_dir = tmp_path / "reports" - assert ( - main(["report", "--artifact", str(tmp_path), "--output", str(report_dir)]) - == 0 - ) - - report_output = json.loads(capsys.readouterr().out) - missing_table = (report_dir / "tables" / "missing_required_metrics.csv").read_text() - assert report_output["missing_required_metric_count"] == 11 - assert "orbitquant,flux1-schnell-native,geneval_overall" not in missing_table - assert "orbitquant,flux1-schnell-native,geneval_per_task_single_object" not in ( - missing_table - ) - assert "orbitquant,flux1-schnell-native,geneval_per_task_counting" not in ( - missing_table - ) - assert "original,flux1-schnell-native,geneval_overall" in missing_table - assert "orbitquant,flux1-schnell-native,geneval_per_task_color_attr" in missing_table - - -def test_cli_generate_dry_run_prints_native_request(capsys, tmp_path): - assert ( - main( - [ - "generate", - "--suite", - "flux2-native", - "--prompt", - "A native prompt", - "--output", - str(tmp_path), - "--seed", - "9", - "--device", - "cpu", - "--enable-model-cpu-offload", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["model_id"] == "black-forest-labs/FLUX.2-klein-4B" - assert output["pipeline_kwargs"]["height"] == 1024 - assert output["pipeline_kwargs"]["width"] == 1024 - assert output["enable_model_cpu_offload"] is True - - -@pytest.mark.parametrize("runtime_mode", ["triton_packed_matmul", "native_packed_matmul"]) -def test_cli_generate_dry_run_prints_quantized_native_request(capsys, tmp_path, runtime_mode): - assert ( - main( - [ - "generate", - "--suite", - "wan-native", - "--prompt", - "A native prompt", - "--output", - str(tmp_path), - "--seed", - "9", - "--device", - "cpu", - "--bit-setting", - "W4A6", - "--activation-kernel-backend", - "cpu", - "--runtime-mode", - runtime_mode, - "--dry-run", - ] - ) - == 0 - ) - - output = capsys.readouterr().out - assert '"height": 480' in output - assert '"width": 832' in output - assert '"weight_bits": 4' in output - assert '"activation_bits": 6' in output - assert '"activation_kernel_backend": "cpu"' in output - assert f'"runtime_mode": "{runtime_mode}"' in output - assert '"target_policy": "wan"' in output - - -def test_cli_compare_native_dry_run_prints_artifact_and_native_settings(capsys, tmp_path): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - assert ( - main( - [ - "compare-native", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--source-model", - str(tmp_path / "source-model"), - "--prompt", - "A native prompt", - "--output", - str(tmp_path / "comparison"), - "--seed", - "11", - "--device", - "cpu", - "--dtype", - "float32", - "--runtime-mode", - "triton_packed_matmul", - "--activation-kernel-backend", - "cpu", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["model_id"] == "example/artifact-model" - assert output["source_model"] == str(tmp_path / "source-model") - assert output["bit_setting"] == "W4A4" - assert output["runtime_mode"] == "triton_packed_matmul" - assert output["activation_kernel_backend"] == "cpu" - assert output["pipeline_kwargs"]["height"] == 1024 - assert output["pipeline_kwargs"]["width"] == 1024 - - -def test_cli_compare_native_runs_original_and_quantized_side_by_side( - monkeypatch, - capsys, - tmp_path, -): - monkeypatch.setattr( - "orbitquant.layers._native_cpu_packed_matmul_load_error", - lambda: RuntimeError("reference-path test"), - ) - - class TinyPipeline: - def __init__(self, color="red"): - self.color = color - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - self.device = None - - def to(self, device): - self.device = device - return self - - def __call__(self, **kwargs): - self.kwargs = kwargs - layer = self.transformer.transformer_blocks[0]["attn"]["to_q"] - layer(torch.zeros(1, 1, 8)) - return SimpleNamespace(images=[Image.new("RGB", (16, 16), self.color)]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path / "artifact", - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - pipelines = iter([TinyPipeline("red"), TinyPipeline("blue")]) - - def fake_load_pipeline_for_suite(suite, *, model_id=None, **kwargs): - assert suite.name == "flux2-native" - assert model_id == "example/artifact-model" - assert kwargs["torch_dtype"] is torch.float32 - return next(pipelines) - - validation_calls = [] - - def fake_validate_compare_native_bundle(bundle_dir): - validation_calls.append(bundle_dir) - return { - "valid": True, - "orbitquant_runtime": { - "runtime_mode_counts": { - "dequant_bf16": 1, - }, - }, - } - - monkeypatch.setattr(cli_main, "load_pipeline_for_suite", fake_load_pipeline_for_suite) - monkeypatch.setattr( - cli_main, - "_validate_compare_native_bundle", - fake_validate_compare_native_bundle, - ) - - assert ( - main( - [ - "compare-native", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path / "artifact"), - "--prompt", - "A native prompt", - "--output", - str(tmp_path / "comparison"), - "--seed", - "4", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - summary = json.loads((tmp_path / "comparison" / "summary.json").read_text()) - comparison_path = ( - tmp_path - / "comparison" - / "flux2-native_seed4_W4A4_original_vs_orbitquant.webp" - ) - assert output["summary_path"] == str(tmp_path / "comparison" / "summary.json") - assert output["comparison_path"] == str(comparison_path) - assert output["validation"]["valid"] is True - assert output["validation"]["orbitquant_runtime"]["runtime_mode_counts"] == { - "dequant_bf16": 1 - } - assert validation_calls == [tmp_path / "comparison"] - assert comparison_path.is_file() - assert summary["original"]["output_path"].endswith("flux2-native_seed4_original.png") - assert summary["orbitquant"]["output_path"].endswith("flux2-native_seed4_W4A4.png") - assert summary["model_id"] == "example/artifact-model" - assert summary["source_model"] == "example/artifact-model" - assert summary["runtime_mode"] == "auto_fused" - assert summary["enable_model_cpu_offload"] is False - assert summary["available_backends"]["cpu"] is True - assert summary["orbitquant"]["runtime"] == { - "orbitquant_linear_count": 1, - "executed_module_count": 1, - "runtime_mode_counts": {"dequant_bf16": 1}, - "activation_kernel_backend_counts": {"cpu": 1}, - "forward_device_type_counts": {"cpu": 1}, - "unexecuted_module_count": 0, - "unexecuted_module_sample": [], - } - assert not any((tmp_path / "artifact" / "assets").iterdir()) - - -def test_cli_compare_native_can_skip_post_generation_validation( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self, color="red"): - self.color = color - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - return self - - def __call__(self, **kwargs): - layer = self.transformer.transformer_blocks[0]["attn"]["to_q"] - layer(torch.zeros(1, 1, 8)) - return SimpleNamespace(images=[Image.new("RGB", (16, 16), self.color)]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path / "artifact", - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - pipelines = iter([TinyPipeline("red"), TinyPipeline("blue")]) - - monkeypatch.setattr( - cli_main, - "load_pipeline_for_suite", - lambda *args, **kwargs: next(pipelines), - ) - monkeypatch.setattr( - cli_main, - "_validate_compare_native_bundle", - lambda *args, **kwargs: (_ for _ in ()).throw( - AssertionError("post-generation validation should be skipped") - ), - ) - - assert ( - main( - [ - "compare-native", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path / "artifact"), - "--prompt", - "A native prompt", - "--output", - str(tmp_path / "comparison"), - "--device", - "cpu", - "--dtype", - "float32", - "--skip-comparison-validation", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["validation"] is None - assert (tmp_path / "comparison" / "summary.json").is_file() - - -def test_cli_validate_comparison_accepts_copied_runpod_bundle(capsys, tmp_path): - suite = get_native_suite("flux2-native") - original_path = tmp_path / "flux2-native_seed4_original.png" - orbitquant_path = tmp_path / "flux2-native_seed4_W4A4.png" - comparison_path = tmp_path / "flux2-native_seed4_W4A4_original_vs_orbitquant.webp" - original_image = Image.effect_noise((suite.width, suite.height), 32).convert("RGB") - orbitquant_image = Image.effect_noise((suite.width, suite.height), 32).convert("RGB") - original_image.save(original_path) - orbitquant_image.save(orbitquant_path) - comparison_sheet = Image.new("RGB", (suite.width * 2, suite.height + 24), "white") - comparison_sheet.paste(original_image, (0, 24)) - comparison_sheet.paste(orbitquant_image, (suite.width, 24)) - comparison_sheet.save(comparison_path) - - for path, bit_setting in ( - (original_path, "original"), - (orbitquant_path, "W4A4"), - ): - metadata = { - "suite": suite.name, - "model_id": "example/artifact-model", - "prompt": "A native prompt", - "seed": 4, - "height": suite.height, - "width": suite.width, - "frames": suite.frames, - "export_fps": suite.export_fps, - "steps": suite.steps, - "guidance": suite.guidance, - "quantization": None - if bit_setting == "original" - else {"config": {"weight_bits": 4, "activation_bits": 4}}, - } - path.with_suffix(path.suffix + ".json").write_text( - json.dumps(metadata) + "\n", - encoding="utf-8", - ) - - (tmp_path / "summary.json").write_text( - json.dumps( - { - "suite": suite.name, - "model_id": "example/artifact-model", - "source_model": "/workspace/hf-models/example", - "artifact": "/workspace/artifacts/example", - "component": "transformer", - "prompt": "A native prompt", - "seed": 4, - "height": suite.height, - "width": suite.width, - "frames": suite.frames, - "steps": suite.steps, - "guidance": suite.guidance, - "dtype": "bfloat16", - "device": "cuda", - "bit_setting": "W4A4", - "runtime_mode": "triton_packed_matmul", - "activation_kernel_backend": "triton_cuda", - "enable_model_cpu_offload": True, - "available_backends": {"cpu": True, "mps": False, "triton_cuda": True}, - "original": { - "output_path": f"/workspace/run/{original_path.name}", - "metadata_path": f"/workspace/run/{original_path.name}.json", - "wall_time_seconds": 10.0, - "peak_vram_bytes": 1000, - }, - "orbitquant": { - "output_path": f"/workspace/run/{orbitquant_path.name}", - "metadata_path": f"/workspace/run/{orbitquant_path.name}.json", - "wall_time_seconds": 15.0, - "peak_vram_bytes": 900, - "runtime": { - "runtime_mode_counts": {"triton_packed_matmul": 1}, - "activation_kernel_backend_counts": {"triton_cuda": 1}, - }, - }, - "comparison_path": f"/workspace/run/{comparison_path.name}", - } - ) - + "\n", - encoding="utf-8", - ) - - assert main(["validate-comparison", "--input", str(tmp_path)]) == 0 - - output = json.loads(capsys.readouterr().out) - assert output["valid"] is True - assert output["speed_ratio_orbitquant_over_original"] == 1.5 - assert output["comparison"]["size"] == [suite.width * 2, suite.height + 24] - assert output["orbitquant_runtime"]["runtime_mode_counts"] == { - "triton_packed_matmul": 1 - } - - -def test_cli_validate_comparison_accepts_wan_video_contact_sheets(capsys, tmp_path): - from orbitquant.eval.assets import create_image_comparison_sheet - - suite = get_native_suite("wan-native") - original_path = tmp_path / "wan-native_seed4_original.mp4" - orbitquant_path = tmp_path / "wan-native_seed4_W4A4.mp4" - original_contact_sheet = tmp_path / "wan-native_seed4_original_contact_sheet.webp" - orbitquant_contact_sheet = tmp_path / "wan-native_seed4_W4A4_contact_sheet.webp" - comparison_path = tmp_path / "wan-native_seed4_W4A4_original_vs_orbitquant.webp" - contact_sheet_size = (suite.width * 4, suite.height * 3) - - for path in (original_path, orbitquant_path): - path.write_bytes(b"fake-mp4") - for path, colors in ( - (original_contact_sheet, ((0, 0, 0), (255, 255, 255))), - (orbitquant_contact_sheet, ((255, 0, 0), (0, 0, 255))), - ): - sheet = Image.new("RGB", contact_sheet_size, colors[0]) - sheet.paste(colors[1], (0, 0, contact_sheet_size[0] // 2, contact_sheet_size[1])) - sheet.save(path) - create_image_comparison_sheet( - original_contact_sheet, - orbitquant_contact_sheet, - comparison_path, - labels=("BF16", "OrbitQuant W4A4"), - ) - - for path, contact_sheet, bit_setting in ( - (original_path, original_contact_sheet, "original"), - (orbitquant_path, orbitquant_contact_sheet, "W4A4"), - ): - metadata = { - "suite": suite.name, - "model_id": "example/wan-artifact-model", - "prompt": "A native video prompt", - "seed": 4, - "height": suite.height, - "width": suite.width, - "frames": suite.frames, - "export_fps": suite.export_fps, - "steps": suite.steps, - "guidance": suite.guidance, - "contact_sheet_path": f"/workspace/run/{contact_sheet.name}", - "quantization": None - if bit_setting == "original" - else {"config": {"weight_bits": 4, "activation_bits": 4}}, - } - path.with_suffix(path.suffix + ".json").write_text( - json.dumps(metadata) + "\n", - encoding="utf-8", - ) - - (tmp_path / "summary.json").write_text( - json.dumps( - { - "suite": suite.name, - "model_id": "example/wan-artifact-model", - "source_model": "/workspace/hf-models/example-wan", - "artifact": "/workspace/artifacts/example-wan", - "component": "transformer", - "prompt": "A native video prompt", - "seed": 4, - "height": suite.height, - "width": suite.width, - "frames": suite.frames, - "steps": suite.steps, - "guidance": suite.guidance, - "dtype": "bfloat16", - "device": "cuda", - "bit_setting": "W4A4", - "runtime_mode": "triton_packed_matmul", - "activation_kernel_backend": "triton_cuda", - "enable_model_cpu_offload": True, - "available_backends": {"cpu": True, "mps": False, "triton_cuda": True}, - "original": { - "output_path": f"/workspace/run/{original_path.name}", - "metadata_path": f"/workspace/run/{original_path.name}.json", - "wall_time_seconds": 10.0, - "peak_vram_bytes": 1000, - }, - "orbitquant": { - "output_path": f"/workspace/run/{orbitquant_path.name}", - "metadata_path": f"/workspace/run/{orbitquant_path.name}.json", - "wall_time_seconds": 14.0, - "peak_vram_bytes": 900, - "runtime": { - "runtime_mode_counts": {"triton_packed_matmul": 1}, - "activation_kernel_backend_counts": {"triton_cuda": 1}, - }, - }, - "comparison_path": f"/workspace/run/{comparison_path.name}", - } - ) - + "\n", - encoding="utf-8", - ) - - assert main(["validate-comparison", "--input", str(tmp_path)]) == 0 - - output = json.loads(capsys.readouterr().out) - assert output["valid"] is True - assert output["native_settings"]["frames"] == 81 - assert output["original"]["contact_sheet"]["size"] == [ - contact_sheet_size[0], - contact_sheet_size[1], - ] - assert output["orbitquant"]["contact_sheet"]["size"] == [ - contact_sheet_size[0], - contact_sheet_size[1], - ] - assert output["comparison"]["size"] == [ - contact_sheet_size[0] * 2, - contact_sheet_size[1] + 24, - ] - assert output["speed_ratio_orbitquant_over_original"] == 1.4 - - -def test_cli_validate_comparison_rejects_blank_images(tmp_path): - suite = get_native_suite("flux2-native") - for name, quantization in ( - ("flux2-native_seed4_original.png", None), - ( - "flux2-native_seed4_W4A4.png", - {"config": {"weight_bits": 4, "activation_bits": 4}}, - ), - ): - path = tmp_path / name - Image.new("RGB", (suite.width, suite.height), "black").save(path) - path.with_suffix(path.suffix + ".json").write_text( - json.dumps( - { - "suite": suite.name, - "model_id": "example/artifact-model", - "prompt": "A native prompt", - "seed": 4, - "height": suite.height, - "width": suite.width, - "frames": suite.frames, - "export_fps": suite.export_fps, - "steps": suite.steps, - "guidance": suite.guidance, - "quantization": quantization, - } - ) - + "\n", - encoding="utf-8", - ) - Image.new("RGB", (suite.width * 2, suite.height + 24), "black").save( - tmp_path / "flux2-native_seed4_W4A4_original_vs_orbitquant.webp" - ) - (tmp_path / "summary.json").write_text( - json.dumps( - { - "suite": suite.name, - "model_id": "example/artifact-model", - "prompt": "A native prompt", - "seed": 4, - "bit_setting": "W4A4", - "original": { - "output_path": "/workspace/run/flux2-native_seed4_original.png", - "metadata_path": "/workspace/run/flux2-native_seed4_original.png.json", - }, - "orbitquant": { - "output_path": "/workspace/run/flux2-native_seed4_W4A4.png", - "metadata_path": "/workspace/run/flux2-native_seed4_W4A4.png.json", - }, - "comparison_path": ( - "/workspace/run/" - "flux2-native_seed4_W4A4_original_vs_orbitquant.webp" - ), - } - ) - + "\n", - encoding="utf-8", - ) - - with pytest.raises(RuntimeError, match="comparison image looks blank"): - main(["validate-comparison", "--input", str(tmp_path)]) - - -def test_cli_generate_with_packed_runtime_on_the_fly_quantization_skips_dequant_prewarm( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - - def to(self, device): - return self - - def __call__(self, **kwargs): - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "blue")]) - - def fake_apply_quantization(pipeline, suite, config): - assert config.runtime_mode == "triton_packed_matmul" - return SimpleNamespace( - quantized_modules=["transformer_blocks.0.attn.to_q"], - adaln_modules=[], - skipped_modules=[], - ) - - def fail_prewarm(*args, **kwargs): - raise AssertionError("packed runtime must not materialize dequant prewarm") - - monkeypatch.setattr(cli_main, "load_pipeline_for_suite", lambda *args, **kwargs: TinyPipeline()) - monkeypatch.setattr(cli_main, "apply_quantization_to_pipeline", fake_apply_quantization) - monkeypatch.setattr(cli_main, "_prewarm_pipeline_component", fail_prewarm) - - assert ( - main( - [ - "generate", - "--suite", - "flux2-native", - "--prompt", - "A native prompt", - "--output", - str(tmp_path), - "--bit-setting", - "W4A4", - "--runtime-mode", - "triton_packed_matmul", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["output_path"].endswith("flux2-native_seed0_W4A4.png") - - -def test_cli_validate_generation_reports_valid_native_output(capsys, tmp_path): - output_path = tmp_path / "flux2-native_seed5_W4A4_simple-object.png" - Image.new("RGB", (16, 16), "red").save(output_path) - output_path.with_suffix(".png.json").write_text( - json.dumps( - { - "suite": "flux2-native", - "model_id": "black-forest-labs/FLUX.2-klein-4B", - "prompt": "A native image", - "seed": 5, - "height": 1024, - "width": 1024, - "frames": None, - "steps": 4, - "guidance": 1.0, - "quantization": { - "config": { - "weight_bits": 4, - "activation_bits": 4, - } - }, - } - ) - + "\n" - ) - - assert ( - main( - [ - "validate-generation", - "--suite", - "flux2-native", - "--output", - str(output_path), - "--seed", - "5", - "--bit-setting", - "W4A4", - "--prompt", - "A native image", - ] - ) - == 0 - ) - - payload = json.loads(capsys.readouterr().out) - assert payload["valid"] is True - assert payload["suite"] == "flux2-native" - assert payload["bit_setting"] == "W4A4" - - -def test_cli_generate_dry_run_uses_artifact_source_and_default_assets_output(capsys, tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - assert ( - main( - [ - "generate", - "--suite", - "flux2-native", - "--prompt", - "A native prompt", - "--artifact", - str(tmp_path), - "--device", - "cpu", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["model_id"] == "example/artifact-model" - assert output["artifact"] == str(tmp_path) - assert output["output"] == str(tmp_path / "assets") - assert output["quantization_config"]["weight_bits"] == 4 - assert output["quantization_config"]["activation_bits"] == 4 - - -def test_cli_generate_dry_run_selects_prompt_from_artifact_pack(capsys, tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - assert ( - main( - [ - "generate", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-id", - "english-text-rendering", - "--device", - "cpu", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["prompt_record"]["id"] == "english-text-rendering" - expected_prompt = select_prompt_record( - default_prompt_payload("flux2"), prompt_id="english-text-rendering" - )["prompt"] - assert output["pipeline_kwargs"]["prompt"] == expected_prompt - - -def test_cli_generate_with_artifact_loads_component_and_records_metrics( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - self.device = None - - def to(self, device): - self.device = device - return self - - def __call__(self, **kwargs): - self.kwargs = kwargs - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "blue")]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - restored = TinyPipeline() - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - assert kwargs["torch_dtype"] is torch.float32 - return restored - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - - assert ( - main( - [ - "generate", - "--suite", - "flux2-native", - "--prompt", - "A native prompt", - "--artifact", - str(tmp_path), - "--seed", - "3", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - record = json.loads((tmp_path / "benchmark" / "orbitquant.metrics.jsonl").read_text()) - assert output["artifact"] == str(tmp_path) - assert output["output_path"].endswith("flux2-native_seed3_W4A4.png") - assert (tmp_path / "assets" / "flux2-native_seed3_W4A4.png").is_file() - assert record["split"] == "orbitquant" - assert record["metrics"]["generated_samples"] == 1 - assert record["metrics"]["wall_time_seconds"] >= 0.0 - assert record["metadata"]["output_path"] == output["output_path"] - assert record["metadata"]["device"] == "cpu" - assert record["metadata"]["dtype"] == "float32" - sample_metadata = json.loads( - (tmp_path / "assets" / "flux2-native_seed3_W4A4.png.json").read_text() - ) - assert sample_metadata["device"] == "cpu" - assert sample_metadata["dtype"] == "float32" - assert sample_metadata["pipeline_class"] == "TinyPipeline" - manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) - assert "assets/flux2-native_seed3_W4A4.png" in manifest["checksums"] - assert "assets/flux2-native_seed3_W4A4.png.json" in manifest["checksums"] - assert restored.device == "cpu" - - -def test_cli_generate_with_packed_runtime_artifact_skips_dequant_prewarm( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - return self - - def __call__(self, **kwargs): - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "blue")]) - - source = TinyPipeline() - config = OrbitQuantConfig( - block_size=4, - target_policy="generic_dit", - runtime_mode="triton_packed_matmul", - ) - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - return TinyPipeline() - - def fail_prewarm(*args, **kwargs): - raise AssertionError("packed runtime must not materialize dequant prewarm") - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - monkeypatch.setattr(cli_main, "_prewarm_pipeline_component", fail_prewarm) - - assert ( - main( - [ - "generate", - "--suite", - "flux2-native", - "--prompt", - "A native prompt", - "--artifact", - str(tmp_path), - "--seed", - "3", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["output_path"].endswith("flux2-native_seed3_W4A4.png") - - -def test_cli_generate_with_artifact_original_split_skips_quantized_component( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - return self - - def __call__(self, **kwargs): - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "yellow")]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - return TinyPipeline() - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - monkeypatch.setattr( - cli_main, - "load_quantized_pipeline_component", - lambda *args, **kwargs: (_ for _ in ()).throw( - AssertionError("original split must not load quantized component") - ), - ) - - assert ( - main( - [ - "generate", - "--suite", - "flux2-native", - "--prompt", - "A native prompt", - "--artifact", - str(tmp_path), - "--split", - "original", - "--seed", - "3", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - record = json.loads((tmp_path / "benchmark" / "original.metrics.jsonl").read_text()) - assert output["output_path"].endswith("flux2-native_seed3_original.png") - assert record["split"] == "original" - assert record["metadata"]["bit_setting"] == "original" - assert not (tmp_path / "benchmark" / "orbitquant.metrics.jsonl").exists() - - -def test_cli_generate_creates_comparison_when_original_pair_exists( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - color = "red" - - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - return self - - def __call__(self, **kwargs): - return SimpleNamespace(images=[Image.new("RGB", (16, 16), self.color)]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - return TinyPipeline() - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - - common_args = [ - "generate", - "--suite", - "flux2-native", - "--prompt", - "A native prompt", - "--artifact", - str(tmp_path), - "--seed", - "4", - "--device", - "cpu", - "--dtype", - "float32", - ] - assert main([*common_args, "--split", "original"]) == 0 - capsys.readouterr() - TinyPipeline.color = "blue" - assert main(common_args) == 0 - - output = json.loads(capsys.readouterr().out) - manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) - comparison = "assets/original_vs_orbitquant_flux2-native_seed4_W4A4_prompt.webp" - assert comparison in manifest["checksums"] - assert (tmp_path / comparison).is_file() - assert output["artifact_comparisons"] == [comparison] - - -def test_cli_generate_with_video_artifact_records_contact_sheet_asset( - monkeypatch, - capsys, - tmp_path, -): - class TinyWanPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - self.device = device - return self - - def __call__(self, **kwargs): - self.kwargs = kwargs - return SimpleNamespace(frames=np.zeros((1, 2, 8, 8, 3), dtype=np.uint8)) - - export_calls = [] - - def fake_export_to_video(frames, path, **kwargs): - export_calls.append(kwargs) - with open(path, "wb") as output: - output.write(b"fake mp4") - - source = TinyWanPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - restored = TinyWanPipeline() - - class FakeWanPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - assert kwargs["torch_dtype"] is torch.float32 - return restored - - monkeypatch.setitem(sys.modules, "diffusers", SimpleNamespace(WanPipeline=FakeWanPipeline)) - monkeypatch.setitem( - sys.modules, - "diffusers.utils", - SimpleNamespace(export_to_video=fake_export_to_video), - ) - - assert ( - main( - [ - "generate", - "--suite", - "wan-native", - "--prompt", - "A native video prompt", - "--artifact", - str(tmp_path), - "--seed", - "2", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) - metrics_record = json.loads((tmp_path / "benchmark" / "orbitquant.metrics.jsonl").read_text()) - metadata = json.loads((tmp_path / "assets" / "wan-native_seed2_W4A4.mp4.json").read_text()) - assert output["output_path"].endswith("wan-native_seed2_W4A4.mp4") - assert export_calls == [{"fps": 16}] - assert metadata["export_fps"] == 16 - assert metrics_record["metadata"]["export_fps"] == 16 - assert "assets/wan-native_seed2_W4A4.mp4" in manifest["checksums"] - assert "assets/wan-native_seed2_W4A4.mp4.json" in manifest["checksums"] - assert "assets/wan-native_seed2_W4A4_contact_sheet.webp" in manifest["checksums"] - - -def test_cli_generate_pack_dry_run_lists_prompt_seed_jobs(capsys, tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-limit", - "2", - "--seeds", - "0,1", - "--device", - "cpu", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["job_count"] == 4 - assert output["jobs"][0]["prompt_record"]["id"] == "simple-object" - assert output["jobs"][0]["seed"] == 0 - assert output["jobs"][2]["seed"] == 1 - assert output["output"] == str(tmp_path / "assets") - - -def test_cli_generate_pack_dry_run_accepts_geneval_smoke_prompt_pack(capsys, tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux1-schnell-native", - "--artifact", - str(tmp_path), - "--prompt-pack", - "geneval-smoke", - "--prompt-limit", - "2", - "--seeds", - "0", - "--device", - "cpu", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["prompt_pack"] == "geneval_smoke_v1" - assert output["job_count"] == 2 - assert output["jobs"][0]["prompt_record"]["id"].startswith("geneval-00000-") - assert output["jobs"][0]["prompt_record"]["geneval"]["tag"] == "single_object" - - -def test_cli_generate_pack_dry_run_accepts_geneval_metadata_jsonl(capsys, tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - metadata_jsonl = tmp_path / "evaluation_metadata.jsonl" - metadata_jsonl.write_text( - json.dumps( - { - "tag": "single_object", - "include": [{"class": "bench", "count": 1}], - "prompt": "a photo of a bench", - } - ) - + "\n", - encoding="utf-8", - ) - - assert ( - main( - [ - "generate-pack", - "--suite", - "z-image-native", - "--artifact", - str(tmp_path), - "--prompt-metadata-jsonl", - str(metadata_jsonl), - "--seeds", - "4", - "--device", - "cpu", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["prompt_pack"] == "geneval_metadata_jsonl" - assert output["job_count"] == 1 - assert output["jobs"][0]["seed"] == 4 - assert output["jobs"][0]["prompt_record"]["geneval"]["include"] == [ - {"class": "bench", "count": 1} - ] - - -def test_cli_generate_pack_runs_jobs_once_per_prompt_seed_and_records_artifacts( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - self.calls = [] - - def to(self, device): - self.device = device - return self - - def __call__(self, **kwargs): - self.calls.append(kwargs) - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - restored = TinyPipeline() - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - assert kwargs["torch_dtype"] is torch.float32 - return restored - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-id", - "simple-object", - "--prompt-id", - "counting", - "--seeds", - "3", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) - metrics_rows = (tmp_path / "benchmark" / "orbitquant.metrics.jsonl").read_text().splitlines() - assert output["job_count"] == 2 - assert len(restored.calls) == 2 - assert "assets/flux2-native_seed3_W4A4_simple-object.png" in manifest["checksums"] - assert "assets/flux2-native_seed3_W4A4_counting.png" in manifest["checksums"] - assert len(metrics_rows) == 2 - assert json.loads(metrics_rows[0])["metadata"]["prompt_record"]["id"] == "simple-object" - - -def test_cli_generate_pack_with_packed_runtime_artifact_skips_dequant_prewarm( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - return self - - def __call__(self, **kwargs): - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) - - source = TinyPipeline() - config = OrbitQuantConfig( - block_size=4, - target_policy="generic_dit", - runtime_mode="triton_packed_matmul", - ) - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - return TinyPipeline() - - def fail_prewarm(*args, **kwargs): - raise AssertionError("packed runtime must not materialize dequant prewarm") - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - monkeypatch.setattr(cli_main, "_prewarm_pipeline_component", fail_prewarm) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-id", - "simple-object", - "--seeds", - "3", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["run_count"] == 1 - assert output["outputs"][0]["output_path"].endswith( - "flux2-native_seed3_W4A4_simple-object.png" - ) - - -def test_cli_generate_pack_skip_checksums_refreshes_artifact_once_at_end( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - self.calls = [] - - def to(self, device): - return self - - def __call__(self, **kwargs): - self.calls.append(kwargs) - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - restored = TinyPipeline() - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - return restored - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-id", - "simple-object", - "--seeds", - "3", - "--resume-existing", - "--skip-artifact-checksums", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) - validation = validate_orbitquant_artifact(tmp_path) - image_path = "assets/flux2-native_seed3_W4A4_simple-object.png" - metadata_path = "assets/flux2-native_seed3_W4A4_simple-object.png.json" - assert output["job_count"] == 1 - assert output["run_count"] == 1 - assert output["checksum_refresh"]["checksum_count"] == len(manifest["checksums"]) - assert validation["valid"] is True - assert image_path in manifest["checksums"] - assert metadata_path in manifest["checksums"] - assert "benchmark/orbitquant.metrics.jsonl" in manifest["checksums"] - assert output["artifact_comparisons"] == [] - assert output["outputs"][0]["comparisons"] == [] - - -def test_cli_generate_pack_defers_comparison_creation_until_after_jobs( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - return self - - def __call__(self, **kwargs): - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - return TinyPipeline() - - comparison_calls = [] - - def fake_create_comparisons(artifact_dir, **kwargs): - comparison_calls.append({"artifact_dir": artifact_dir, "kwargs": kwargs}) - return ["assets/fake-comparison.webp"] - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - monkeypatch.setattr(cli_main, "create_artifact_image_comparisons", fake_create_comparisons) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-id", - "simple-object", - "--prompt-id", - "counting", - "--seeds", - "3", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert len(comparison_calls) == 1 - assert comparison_calls[0]["kwargs"]["comparison_keys"] == { - ("flux2-native", 3, "simple-object"), - ("flux2-native", 3, "counting"), - } - assert output["run_count"] == 2 - assert output["artifact_comparisons"] == ["assets/fake-comparison.webp"] - assert [item["comparisons"] for item in output["outputs"]] == [[], []] - - -def test_cli_generate_pack_prompt_metadata_disables_comparisons_by_default( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - return self - - def __call__(self, **kwargs): - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - metadata_jsonl = tmp_path / "evaluation_metadata.jsonl" - metadata_jsonl.write_text( - json.dumps( - { - "tag": "single_object", - "include": [{"class": "bench", "count": 1}], - "prompt": "a photo of a bench", - } - ) - + "\n", - encoding="utf-8", - ) - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - return TinyPipeline() - - def fail_create_comparisons(*args, **kwargs): - raise AssertionError("GenEval metadata packs must not create comparison sheets by default") - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - monkeypatch.setattr(cli_main, "create_artifact_image_comparisons", fail_create_comparisons) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-metadata-jsonl", - str(metadata_jsonl), - "--seeds", - "3", - "--skip-artifact-checksums", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["run_count"] == 1 - assert output["artifact_comparisons"] == [] - assert output["outputs"][0]["comparisons"] == [] - - -def test_cli_generate_pack_resume_existing_skips_completed_outputs( - monkeypatch, - capsys, - tmp_path, -): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - output_path = tmp_path / "assets" / "flux2-native_seed3_W4A4_simple-object.png" - Image.new("RGB", (16, 16), "green").save(output_path) - output_path.with_suffix(".png.json").write_text( - json.dumps( - { - "suite": "flux2-native", - "model_id": "example/artifact-model", - "prompt": select_prompt_record( - default_prompt_payload("flux2"), prompt_id="simple-object" - )["prompt"], - "seed": 3, - "height": 1024, - "width": 1024, - "frames": None, - "steps": 4, - "guidance": 1.0, - "quantization": { - "config": { - "weight_bits": 4, - "activation_bits": 4, - } - }, - } - ) - + "\n" - ) - monkeypatch.setattr( - cli_main, - "load_pipeline_for_suite", - lambda *args, **kwargs: (_ for _ in ()).throw( - AssertionError("completed resume jobs must not load a pipeline") - ), - ) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-id", - "simple-object", - "--seeds", - "3", - "--resume-existing", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["job_count"] == 1 - assert output["run_count"] == 0 - assert output["skipped_count"] == 1 - assert output["skipped_outputs"] == [str(output_path)] - - -def test_cli_generate_pack_resume_existing_reruns_invalid_metadata( - monkeypatch, - capsys, - tmp_path, -): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - self.calls = [] - - def to(self, device): - return self - - def __call__(self, **kwargs): - self.calls.append(kwargs) - return SimpleNamespace(images=[Image.new("RGB", (16, 16), "purple")]) - - source = TinyPipeline() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(source.transformer, config) - save_orbitquant_artifact( - source.transformer, - tmp_path, - config=config, - source_model_id="example/artifact-model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - output_path = tmp_path / "assets" / "flux2-native_seed3_W4A4_simple-object.png" - Image.new("RGB", (16, 16), "green").save(output_path) - output_path.with_suffix(".png.json").write_text('{"status":"complete"}\n') - restored = TinyPipeline() - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/artifact-model" - return restored - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), - ) - - assert ( - main( - [ - "generate-pack", - "--suite", - "flux2-native", - "--artifact", - str(tmp_path), - "--prompt-id", - "simple-object", - "--seeds", - "3", - "--resume-existing", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["job_count"] == 1 - assert output["run_count"] == 1 - assert output["skipped_count"] == 0 - assert len(restored.calls) == 1 - - -def test_cli_quantize_saves_transformer_component_artifact(monkeypatch, capsys, tmp_path): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - self.device = None - - def to(self, device): - self.device = device - return self - - pipeline = TinyPipeline() - - class FakeDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "example/model" - assert kwargs["revision"] == "main" - assert kwargs["torch_dtype"] is torch.float32 - return pipeline - - monkeypatch.setitem( - sys.modules, "diffusers", SimpleNamespace(DiffusionPipeline=FakeDiffusionPipeline) - ) - monkeypatch.setattr( - cli_main, - "inspect_model_metadata", - lambda model_id, revision=None: { - "sha": "abc123", - "license": "apache-2.0", - }, - ) - - assert ( - main( - [ - "quantize", - "--model-id", - "example/model", - "--revision", - "main", - "--output", - str(tmp_path), - "--component", - "transformer", - "--target-policy", - "generic_dit", - "--weight-bits", - "4", - "--activation-bits", - "4", - "--block-size", - "4", - "--activation-kernel-backend", - "cpu", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["artifact_dir"] == str(tmp_path) - assert output["quantization_staging_mode"] == "streaming" - assert output["synchronize_per_module"] is False - assert output["device_transfer_seconds"] >= 0.0 - assert output["module_device_transfer_count"] >= 0 - assert output["source_linear_device_counts"] - assert output["artifact_save_elapsed_seconds"] >= 0.0 - assert output["quantized_modules"] == ["transformer_blocks.0.attn.to_q"] - assert (tmp_path / "model.safetensors").exists() - assert (tmp_path / "orbitquant_manifest.json").exists() - assert (tmp_path / "SHA256SUMS").exists() - quantization_config = json.loads((tmp_path / "quantization_config.json").read_text()) - assert quantization_config["activation_kernel_backend"] == "cpu" - - -def test_cli_quantize_with_suite_uses_named_native_pipeline(monkeypatch, capsys, tmp_path): - class TinyPipeline: - def __init__(self): - self.transformer = torch.nn.Module() - self.transformer.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def to(self, device): - return self - - class FakeFlux2KleinPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - assert model_id == "black-forest-labs/FLUX.2-klein-4B" - assert kwargs["torch_dtype"] is torch.float32 - return TinyPipeline() - - class WrongDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - raise AssertionError("suite-specific pipeline should be used") - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace( - Flux2KleinPipeline=FakeFlux2KleinPipeline, - DiffusionPipeline=WrongDiffusionPipeline, - ), - ) - monkeypatch.setattr( - cli_main, - "inspect_model_metadata", - lambda model_id, revision=None: { - "sha": "abc123", - "license": "apache-2.0", - }, - ) - - assert ( - main( - [ - "quantize", - "--suite", - "flux2-native", - "--output", - str(tmp_path), - "--component", - "transformer", - "--weight-bits", - "4", - "--activation-bits", - "4", - "--block-size", - "4", - "--device", - "cpu", - "--dtype", - "float32", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - quantization_config = json.loads((tmp_path / "quantization_config.json").read_text()) - assert output["source_model_id"] == "black-forest-labs/FLUX.2-klein-4B" - assert quantization_config["target_policy"] == "flux2" - - -def test_cli_inspect_policy_with_suite_writes_module_inventory(monkeypatch, capsys, tmp_path): - calls = [] - - class FakeFlux2Transformer2DModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.transformer = torch.nn.Module() - self.double_stream_modulation_img = torch.nn.ModuleDict( - {"linear": torch.nn.Linear(8, 16)} - ) - self.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - @classmethod - def load_config(cls, model_id, **kwargs): - assert model_id == "black-forest-labs/FLUX.2-klein-4B" - assert kwargs["subfolder"] == "transformer" - assert kwargs["local_files_only"] is False - calls.append({"method": "load_config", "kwargs": kwargs}) - return {"hidden_size": 8} - - @classmethod - def from_config(cls, config): - assert config == {"hidden_size": 8} - calls.append({"method": "from_config"}) - return cls() - - class WrongDiffusionPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - raise AssertionError("inspect-policy default must not load full pipeline weights") - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace( - Flux2Transformer2DModel=FakeFlux2Transformer2DModel, - DiffusionPipeline=WrongDiffusionPipeline, - ), - ) - output_path = tmp_path / "flux2-policy-inventory.json" - - assert ( - main( - [ - "inspect-policy", - "--suite", - "flux2-native", - "--dtype", - "float32", - "--output", - str(output_path), - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - saved = json.loads(output_path.read_text()) - - assert output["output"] == str(output_path) - assert output["linear_module_count"] == 2 - assert output["quantized_module_count"] == 1 - assert output["adaln_module_count"] == 1 - assert output["skipped_module_count"] == 0 - assert "modules" not in output - assert saved["source_model_id"] == "black-forest-labs/FLUX.2-klein-4B" - assert saved["suite"] == "flux2-native" - assert saved["component"] == "transformer" - assert saved["load_mode"] == "config" - assert saved["pipeline_class"] is None - assert saved["component_class"] == "FakeFlux2Transformer2DModel" - assert saved["target_policy"] == "flux2" - assert saved["action_counts"] == { - "orbitquant": 1, - "adaln_int4_rtn": 1, - "bf16_skip": 0, - } - assert saved["quantized_modules"] == ["transformer_blocks.0.attn.to_q"] - assert saved["adaln_modules"] == ["double_stream_modulation_img.linear"] - assert calls == [ - { - "method": "load_config", - "kwargs": {"subfolder": "transformer", "local_files_only": False}, - }, - {"method": "from_config"}, - ] - - -@pytest.mark.parametrize( - ("suite", "transformer_class", "model_id", "target_policy", "module_family"), - [ - ( - "flux2-native", - "Flux2Transformer2DModel", - "black-forest-labs/FLUX.2-klein-4B", - "flux2", - "flux", - ), - ( - "flux1-schnell-native", - "FluxTransformer2DModel", - "black-forest-labs/FLUX.1-schnell", - "flux", - "flux", - ), - ( - "z-image-native", - "ZImageTransformer2DModel", - "Tongyi-MAI/Z-Image-Turbo", - "z_image", - "z_image", - ), - ( - "wan-native", - "WanTransformer3DModel", - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "wan", - "wan", - ), - ], -) -def test_cli_inspect_policy_config_mode_covers_all_native_target_suites( - monkeypatch, - capsys, - tmp_path, - suite, - transformer_class, - model_id, - target_policy, - module_family, -): - calls = [] - - def build_module(self): - if module_family == "wan": - self.blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn1": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - elif module_family == "z_image": - self.layers = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attention": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - else: - self.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} - ) - ] - ) - - def load_config(cls, requested_model_id, **kwargs): - assert requested_model_id == model_id - assert kwargs["subfolder"] == "transformer" - calls.append(("load_config", transformer_class)) - return {"hidden_size": 8} - - def from_config(cls, config): - assert config == {"hidden_size": 8} - calls.append(("from_config", transformer_class)) - return cls() - - def init(self): - torch.nn.Module.__init__(self) - build_module(self) - - fake_transformer_cls = type( - transformer_class, - (torch.nn.Module,), - { - "__init__": init, - "load_config": classmethod(load_config), - "from_config": classmethod(from_config), - }, - ) - - class WrongDiffusionPipeline: - @classmethod - def from_pretrained(cls, requested_model_id, **kwargs): - raise AssertionError("config inspect-policy must not load full pipeline weights") - - monkeypatch.setitem( - sys.modules, - "diffusers", - SimpleNamespace( - **{ - transformer_class: fake_transformer_cls, - "DiffusionPipeline": WrongDiffusionPipeline, - } - ), - ) - output_path = tmp_path / f"{suite}-policy-inventory.json" - - assert main(["inspect-policy", "--suite", suite, "--output", str(output_path)]) == 0 - - output = json.loads(capsys.readouterr().out) - saved = json.loads(output_path.read_text()) - - assert output["output"] == str(output_path) - assert saved["suite"] == suite - assert saved["source_model_id"] == model_id - assert saved["component"] == "transformer" - assert saved["load_mode"] == "config" - assert saved["pipeline_class"] is None - assert saved["component_class"] == transformer_class - assert saved["target_policy"] == target_policy - assert saved["action_counts"]["orbitquant"] == 1 - assert saved["quantized_modules"] - assert calls == [ - ("load_config", transformer_class), - ("from_config", transformer_class), - ] - - -def test_cli_validate_artifact_reports_valid_component_artifact(capsys, tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - assert main(["validate-artifact", "--artifact", str(tmp_path)]) == 0 - - output = json.loads(capsys.readouterr().out) - assert output["valid"] is True - assert output["source_model_id"] == "example/model" - assert output["tensor_count"] > 0 - assert output["quantized_module_count"] == 1 - - -@pytest.mark.parametrize("runtime_mode", ["triton_packed_matmul", "native_packed_matmul"]) -def test_cli_validate_artifact_rejects_unexpected_runtime_mode(capsys, tmp_path, runtime_mode): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - assert ( - main( - [ - "validate-artifact", - "--artifact", - str(tmp_path), - "--runtime-mode", - runtime_mode, - ] - ) - == 1 - ) - - assert "runtime_mode mismatch" in capsys.readouterr().err - - -def test_cli_validate_artifact_checks_policy_inventory(capsys, tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - inventory_path = tmp_path / "policy-inventory.json" - inventory_path.write_text( - json.dumps( - { - "source_model_id": "example/model", - "target_policy": "generic_dit", - "component": "transformer", - "load_mode": "config", - "linear_module_count": 1, - "action_counts": { - "orbitquant": 1, - "adaln_int4_rtn": 0, - "bf16_skip": 0, - }, - "quantized_modules": summary.quantized_modules, - "adaln_modules": summary.adaln_modules, - "skipped_modules": summary.skipped_modules, - } - ) - + "\n" - ) - - assert ( - main( - [ - "validate-artifact", - "--artifact", - str(tmp_path), - "--policy-inventory", - str(inventory_path), - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - policy_validation = output["policy_inventory_validation"] - assert policy_validation["valid"] is True - assert policy_validation["quantized_module_count"] == 1 - assert policy_validation["inventory_path"] == str(inventory_path) - - -def test_cli_validate_artifact_rejects_policy_inventory_mismatch(tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - inventory_path = tmp_path / "policy-inventory.json" - inventory_path.write_text( - json.dumps( - { - "source_model_id": "example/model", - "target_policy": "generic_dit", - "quantized_modules": ["transformer_blocks.0.attn.to_k"], - "adaln_modules": [], - "skipped_modules": [], - } - ) - + "\n" - ) - - with pytest.raises(RuntimeError, match="policy inventory mismatch"): - main( - [ - "validate-artifact", - "--artifact", - str(tmp_path), - "--policy-inventory", - str(inventory_path), - ] - ) - - -def test_cli_validate_artifact_rejects_policy_inventory_component_mismatch(tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - component="transformer", - ) - inventory_path = tmp_path / "policy-inventory.json" - inventory_path.write_text( - json.dumps( - { - "source_model_id": "example/model", - "target_policy": "generic_dit", - "component": "denoiser", - "quantized_modules": summary.quantized_modules, - "adaln_modules": summary.adaln_modules, - "skipped_modules": summary.skipped_modules, - } - ) - + "\n" - ) - - with pytest.raises(RuntimeError, match="component"): - main( - [ - "validate-artifact", - "--artifact", - str(tmp_path), - "--policy-inventory", - str(inventory_path), - ] - ) - - -def test_cli_validate_artifact_rejects_manifest_policy_config_mismatch(tmp_path): - model = torch.nn.Module() - model.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config) - save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - manifest_path = tmp_path / "orbitquant_manifest.json" - manifest = json.loads(manifest_path.read_text()) - manifest["target_policy"] = "flux" - manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") - write_sha256sums_from_manifest(tmp_path, manifest["checksums"]) - inventory_path = tmp_path / "policy-inventory.json" - inventory_path.write_text( - json.dumps( - { - "source_model_id": "example/model", - "target_policy": "flux", - "component": "transformer", - "quantized_modules": summary.quantized_modules, - "adaln_modules": summary.adaln_modules, - "skipped_modules": summary.skipped_modules, - } - ) - + "\n" - ) - - with pytest.raises(RuntimeError, match="quantization_config mismatch.*target_policy"): - main( - [ - "validate-artifact", - "--artifact", - str(tmp_path), - "--policy-inventory", - str(inventory_path), - ] - ) - - -def test_cli_repair_artifact_metadata_wires_provenance_options( - capsys, tmp_path, monkeypatch -): - seen = {} - - def fake_repair_artifact_metadata(artifact, **kwargs): - seen["artifact"] = artifact - seen["kwargs"] = kwargs - return { - "artifact_dir": artifact, - "updated": { - "quantization_device": kwargs["quantization_device"], - "weight_quantization_backend": kwargs["weight_quantization_backend"], - "quantization_staging_mode": kwargs["quantization_staging_mode"], - }, - } - - monkeypatch.setattr(cli_main, "repair_artifact_metadata", fake_repair_artifact_metadata) - - assert ( - main( - [ - "repair-artifact-metadata", - "--artifact", - str(tmp_path), - "--quantization-device", - "cuda", - "--weight-quantization-backend", - "triton_cuda", - "--quantization-staging-mode", - "component", - "--skip-tensor-validation", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["updated"]["quantization_device"] == "cuda" - assert output["updated"]["weight_quantization_backend"] == "triton_cuda" - assert output["updated"]["quantization_staging_mode"] == "component" - assert seen == { - "artifact": str(tmp_path), - "kwargs": { - "quantization_device": "cuda", - "weight_quantization_backend": "triton_cuda", - "quantization_staging_mode": "component", - "validate_tensors": False, - }, - } - - -def test_cli_upload_artifact_wires_validation_and_hf_options(capsys, tmp_path, monkeypatch): - seen = {} - - def fake_upload_artifact(artifact, **kwargs): - seen["artifact"] = artifact - seen["kwargs"] = kwargs - return { - "artifact_dir": artifact, - "repo_id": kwargs["repo_id"], - "private": kwargs["private"], - "dry_run": kwargs["dry_run"], - "validation": {"valid": True}, - } - - monkeypatch.setattr(cli_main, "upload_orbitquant_artifact", fake_upload_artifact) - - assert ( - main( - [ - "upload-artifact", - "--artifact", - str(tmp_path), - "--repo-id", - "WaveCut/example-orbitquant", - "--revision", - "main", - "--commit-message", - "upload artifact", - "--public", - "--no-create-repo", - "--replace-repo-files", - "--skip-tensor-validation", - "--upload-profile", - "compact", - "--report-dir", - "/tmp/orbitquant-report", - "--staging-dir", - "/tmp/orbitquant-stage", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["repo_id"] == "WaveCut/example-orbitquant" - assert output["private"] is False - assert output["dry_run"] is True - assert seen == { - "artifact": str(tmp_path), - "kwargs": { - "repo_id": "WaveCut/example-orbitquant", - "private": False, - "create_repo": False, - "revision": "main", - "commit_message": "upload artifact", - "replace_repo_files": True, - "validate_tensors": False, - "upload_profile": "compact", - "report_dirs": ["/tmp/orbitquant-report"], - "staging_dir": "/tmp/orbitquant-stage", - "dry_run": True, - }, - } - - -def test_cli_upload_artifact_defaults_to_compact_profile(capsys, tmp_path, monkeypatch): - seen = {} - - def fake_upload_artifact(artifact, **kwargs): - seen["artifact"] = artifact - seen["kwargs"] = kwargs - return { - "artifact_dir": artifact, - "repo_id": kwargs["repo_id"], - "private": kwargs["private"], - "dry_run": kwargs["dry_run"], - "validation": {"valid": True}, - } - - monkeypatch.setattr(cli_main, "upload_orbitquant_artifact", fake_upload_artifact) - - assert ( - main( - [ - "upload-artifact", - "--artifact", - str(tmp_path), - "--repo-id", - "WaveCut/example-orbitquant", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["repo_id"] == "WaveCut/example-orbitquant" - assert seen["artifact"] == str(tmp_path) - assert seen["kwargs"]["upload_profile"] == "compact" - assert seen["kwargs"]["replace_repo_files"] is True - assert seen["kwargs"]["report_dirs"] is None - assert seen["kwargs"]["staging_dir"] is None - - -def test_cli_upload_artifact_rejects_disabling_remote_file_replacement( - capsys, tmp_path, monkeypatch -): - seen = {} - - def fake_upload_artifact(artifact, **kwargs): - seen["artifact"] = artifact - seen["kwargs"] = kwargs - return { - "artifact_dir": artifact, - "repo_id": kwargs["repo_id"], - "private": kwargs["private"], - "dry_run": kwargs["dry_run"], - "validation": {"valid": True}, - } - - monkeypatch.setattr(cli_main, "upload_orbitquant_artifact", fake_upload_artifact) - - with pytest.raises(SystemExit): - main( - [ - "upload-artifact", - "--artifact", - str(tmp_path), - "--repo-id", - "WaveCut/example-orbitquant", - "--no-replace-repo-files", - "--dry-run", - ] - ) - - assert seen == {} - assert "unrecognized arguments: --no-replace-repo-files" in capsys.readouterr().err - - -def test_cli_repair_hf_artifact_metadata_wires_single_repo_options(capsys, monkeypatch): - seen = {} - - def fake_repair_hf_artifact_metadata(**kwargs): - seen.update(kwargs) - return { - "repo_id": kwargs["repo_id"], - "dry_run": kwargs["dry_run"], - "changed_files": ["orbitquant_manifest.json"], - } - - monkeypatch.setattr( - cli_main, "repair_hf_artifact_metadata", fake_repair_hf_artifact_metadata - ) - - assert ( - main( - [ - "repair-hf-artifact-metadata", - "--repo-id", - "WaveCut/example-orbitquant", - "--revision", - "main", - "--commit-message", - "repair metadata", - "--quantization-device", - "cuda", - "--weight-quantization-backend", - "triton_cuda", - "--quantization-staging-mode", - "component", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["repo_id"] == "WaveCut/example-orbitquant" - assert output["dry_run"] is True - assert seen == { - "repo_id": "WaveCut/example-orbitquant", - "quantization_device": "cuda", - "weight_quantization_backend": "triton_cuda", - "quantization_staging_mode": "component", - "revision": "main", - "commit_message": "repair metadata", - "dry_run": True, - } - - -def test_cli_repair_hf_native_smoke_proof_wires_single_repo_options( - capsys, monkeypatch -): - seen = {} - - def fake_repair_hf_native_smoke_proof(**kwargs): - seen.update(kwargs) - return { - "repo_id": kwargs["repo_id"], - "suite": kwargs["suite"].name, - "dry_run": kwargs["dry_run"], - "changed_files": ["benchmark/summary.json"], - } - - monkeypatch.setattr( - cli_main, "repair_hf_native_smoke_proof", fake_repair_hf_native_smoke_proof - ) - - assert ( - main( - [ - "repair-hf-native-smoke-proof", - "--repo-id", - "WaveCut/example-orbitquant", - "--suite", - "flux2-native", - "--revision", - "main", - "--commit-message", - "repair native smoke proof", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["repo_id"] == "WaveCut/example-orbitquant" - assert output["suite"] == "flux2-native" - assert output["dry_run"] is True - assert seen["repo_id"] == "WaveCut/example-orbitquant" - assert seen["suite"].name == "flux2-native" - assert seen["revision"] == "main" - assert seen["commit_message"] == "repair native smoke proof" - assert seen["dry_run"] is True - - -def test_cli_cleanup_hf_artifact_reports_wires_single_repo_options( - capsys, monkeypatch -): - seen = {} - - def fake_cleanup_hf_artifact_reports(**kwargs): - seen.update(kwargs) - return { - "repo_id": kwargs["repo_id"], - "revision": kwargs["revision"], - "dry_run": kwargs["dry_run"], - "report_file_count": 2, - "promoted_assets": ["assets/image_generation_comparison_matrix.webp"], - "delete_paths": ["reports"], - } - - monkeypatch.setattr( - cli_main, "cleanup_hf_artifact_reports", fake_cleanup_hf_artifact_reports - ) - - assert ( - main( - [ - "cleanup-hf-artifact-reports", - "--repo-id", - "WaveCut/example-orbitquant", - "--revision", - "main", - "--commit-message", - "cleanup reports", - "--dry-run", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["repo_id"] == "WaveCut/example-orbitquant" - assert output["promoted_assets"] == ["assets/image_generation_comparison_matrix.webp"] - assert output["delete_paths"] == ["reports"] - assert seen == { - "repo_id": "WaveCut/example-orbitquant", - "revision": "main", - "commit_message": "cleanup reports", - "dry_run": True, - } - - -def test_cli_audit_hf_artifacts_writes_json_report(capsys, tmp_path, monkeypatch): - seen = {} - - def fake_audit_hf_artifacts(*, namespace, suites, revision, policy_inventory_root): - seen["namespace"] = namespace - seen["suites"] = [suite.name for suite in suites] - seen["revision"] = revision - seen["policy_inventory_root"] = policy_inventory_root - return { - "namespace": namespace, - "policy_inventory_root": policy_inventory_root, - "repo_count": 1, - "existing_count": 1, - "artifact_ready_count": 1, - "native_smoke_ready_count": 1, - "release_eval_ready_count": 0, - "rows": [{"repo_id": "WaveCut/example"}], - } - - monkeypatch.setattr(cli_main, "audit_hf_artifact_repos", fake_audit_hf_artifacts) - output_path = tmp_path / "reports" / "native" / "audit.json" - markdown_output_path = tmp_path / "reports" / "native" / "audit.md" - - assert ( - main( - [ - "audit-hf-artifacts", - "--namespace", - "WaveCut", - "--suite", - "flux2-native", - "--revision", - "main", - "--policy-inventory-root", - str(tmp_path / "inventories"), - "--output", - str(output_path), - "--markdown-output", - str(markdown_output_path), - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - written = json.loads(output_path.read_text()) - markdown = markdown_output_path.read_text() - assert output_path.parent.is_dir() - assert output == written - assert output["repo_count"] == 1 - assert "# OrbitQuant HF Artifact Audit" in markdown - assert "| `WaveCut/example` |" in markdown - assert seen == { - "namespace": "WaveCut", - "suites": ["flux2-native"], - "revision": "main", - "policy_inventory_root": str(tmp_path / "inventories"), - } - - -def test_cli_audit_hf_artifacts_summary_only_omits_full_rows( - capsys, tmp_path, monkeypatch -): - def fake_audit_hf_artifacts(*, namespace, suites, revision, policy_inventory_root): - return { - "namespace": namespace, - "policy_inventory_root": policy_inventory_root, - "repo_count": 1, - "existing_count": 1, - "artifact_ready_count": 1, - "native_smoke_ready_count": 1, - "metadata_complete_ready_count": 1, - "policy_inventory_ready_count": 1, - "policy_inventory_error_count": 0, - "release_eval_applicable_count": 1, - "release_eval_ready_count": 0, - "missing_required_metric_count": 2, - "manifest_warning_count": 0, - "metadata_missing_count": 0, - "remote_checksum_mismatch_count": 0, - "readme_mismatch_count": 0, - "forbidden_file_count": 0, - "rows": [ - { - "repo_id": "WaveCut/example", - "suite": "flux1-schnell-native", - "bit_setting": "W4A4", - "artifact_ready": True, - "native_smoke_ready": True, - "release_eval_applicable": True, - "release_eval_ready": False, - "metadata_complete_ready": True, - "policy_inventory_ready": True, - "forbidden_file_count": 0, - "missing_required_metrics": [ - {"split": "original", "metric": "geneval_overall"}, - {"split": "orbitquant", "metric": "geneval_overall"}, - ], - "large_remote_row_detail": {"not": "printed"}, - } - ], - } - - monkeypatch.setattr(cli_main, "audit_hf_artifact_repos", fake_audit_hf_artifacts) - output_path = tmp_path / "audit-summary.json" - - assert ( - main( - [ - "audit-hf-artifacts", - "--summary-only", - "--output", - str(output_path), - "--fail-on-artifact-regression", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - written = json.loads(output_path.read_text()) - assert output == written - assert "rows" not in output - assert output["repo_count"] == 1 - assert output["row_count"] == 1 - assert output["repos"] == [ - { - "repo_id": "WaveCut/example", - "suite": "flux1-schnell-native", - "bit_setting": "W4A4", - "artifact_ready": True, - "native_smoke_ready": True, - "release_eval_applicable": True, - "release_eval_ready": False, - "metadata_complete_ready": True, - "policy_inventory_ready": True, - "forbidden_file_count": 0, - "missing_required_metric_count": 2, - } - ] - assert "large_remote_row_detail" not in json.dumps(output) - - -def test_cli_audit_hf_artifacts_artifact_regression_gate_ignores_release_metrics( - capsys, tmp_path, monkeypatch -): - def fake_audit_hf_artifacts(*, namespace, suites, revision, policy_inventory_root): - return { - "namespace": namespace, - "policy_inventory_root": policy_inventory_root, - "repo_count": 1, - "existing_count": 1, - "artifact_ready_count": 1, - "native_smoke_ready_count": 1, - "metadata_complete_ready_count": 1, - "policy_inventory_ready_count": 1, - "policy_inventory_error_count": 0, - "release_eval_ready_count": 0, - "missing_required_metric_count": 14, - "manifest_warning_count": 0, - "metadata_missing_count": 0, - "remote_checksum_mismatch_count": 0, - "forbidden_file_count": 0, - "rows": [{"repo_id": "WaveCut/example"}], - } - - monkeypatch.setattr(cli_main, "audit_hf_artifact_repos", fake_audit_hf_artifacts) - - assert ( - main( - [ - "audit-hf-artifacts", - "--policy-inventory-root", - str(tmp_path / "inventories"), - "--fail-on-artifact-regression", - ] - ) - == 0 - ) - - output = json.loads(capsys.readouterr().out) - assert output["missing_required_metric_count"] == 14 - - -def test_cli_audit_hf_artifacts_artifact_regression_gate_fails_on_remote_hygiene( - capsys, tmp_path, monkeypatch -): - def fake_audit_hf_artifacts(*, namespace, suites, revision, policy_inventory_root): - return { - "namespace": namespace, - "policy_inventory_root": policy_inventory_root, - "repo_count": 1, - "existing_count": 1, - "artifact_ready_count": 1, - "native_smoke_ready_count": 1, - "metadata_complete_ready_count": 1, - "policy_inventory_ready_count": 1, - "policy_inventory_error_count": 0, - "release_eval_ready_count": 0, - "missing_required_metric_count": 0, - "manifest_warning_count": 0, - "metadata_missing_count": 0, - "remote_checksum_mismatch_count": 0, - "forbidden_file_count": 1, - "rows": [{"repo_id": "WaveCut/example"}], - } - - monkeypatch.setattr(cli_main, "audit_hf_artifact_repos", fake_audit_hf_artifacts) - - assert ( - main( - [ - "audit-hf-artifacts", - "--policy-inventory-root", - str(tmp_path / "inventories"), - "--fail-on-artifact-regression", - ] - ) - == 1 - ) - - captured = capsys.readouterr() - assert json.loads(captured.out)["forbidden_file_count"] == 1 - assert "forbidden_file_count=1 expected 0" in captured.err - - -def test_cli_fetch_hf_artifacts_wires_suite_and_download_options( - capsys, tmp_path, monkeypatch -): - seen = {} - - def fake_fetch_hf_artifacts(**kwargs): - seen.update(kwargs) - if kwargs["stage_logger"] is not None: - kwargs["stage_logger"]("START", "example fetch") - return { - "namespace": kwargs["namespace"], - "output_root": str(kwargs["output_root"]), - "repo_count": 1, - "downloaded_count": 0, - "skipped_existing_count": 0, - "dry_run": kwargs["dry_run"], - "rows": [{"repo_id": "WaveCut/example"}], - } - - monkeypatch.setattr(cli_main, "fetch_hf_artifacts", fake_fetch_hf_artifacts) - - assert ( - main( - [ - "fetch-hf-artifacts", - "--namespace", - "WaveCut", - "--suite", - "flux1-schnell-native", - "--output-root", - str(tmp_path / "artifacts"), - "--revision", - "main", - "--no-resume", - "--force-download", - "--local-files-only", - "--validate-checksums", - "--validate-tensors", - ] - ) - == 0 - ) - - captured = capsys.readouterr() - output = json.loads(captured.out) - assert output["repo_count"] == 1 - assert "example fetch" in captured.err - assert seen["namespace"] == "WaveCut" - assert [suite.name for suite in seen["suites"]] == ["flux1-schnell-native"] - assert seen["output_root"] == str(tmp_path / "artifacts") - assert seen["revision"] == "main" - assert seen["resume"] is False - assert seen["force_download"] is True - assert seen["local_files_only"] is True - assert seen["validate_checksums"] is True - assert seen["validate_tensors"] is True - assert seen["dry_run"] is False - assert seen["stage_logger"] is not None - - -def test_cli_fetch_hf_artifacts_dry_run_suppresses_stage_log(capsys, monkeypatch): - seen = {} - - def fake_fetch_hf_artifacts(**kwargs): - seen.update(kwargs) - return { - "namespace": kwargs["namespace"], - "output_root": str(kwargs["output_root"]), - "repo_count": 0, - "downloaded_count": 0, - "skipped_existing_count": 0, - "dry_run": kwargs["dry_run"], - "rows": [], - } - - monkeypatch.setattr(cli_main, "fetch_hf_artifacts", fake_fetch_hf_artifacts) - - assert main(["fetch-hf-artifacts", "--dry-run"]) == 0 - - captured = capsys.readouterr() - output = json.loads(captured.out) - assert output["dry_run"] is True - assert captured.err == "" - assert seen["dry_run"] is True - assert seen["stage_logger"] is None diff --git a/tests/test_cli_artifact_validation.py b/tests/test_cli_artifact_validation.py new file mode 100644 index 0000000..1ec806b --- /dev/null +++ b/tests/test_cli_artifact_validation.py @@ -0,0 +1,313 @@ +import json + +import pytest +import torch + +import orbitquant.cli.main as cli_main +from orbitquant.artifacts import save_orbitquant_artifact +from orbitquant.artifacts.checksums import write_sha256sums_from_manifest +from orbitquant.cli.main import main +from orbitquant.config import OrbitQuantConfig +from orbitquant.modeling import quantize_linear_modules + + +def test_cli_validate_artifact_reports_valid_component_artifact(capsys, tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + assert main(["validate-artifact", "--artifact", str(tmp_path)]) == 0 + + output = json.loads(capsys.readouterr().out) + assert output["valid"] is True + assert output["source_model_id"] == "example/model" + assert output["tensor_count"] > 0 + assert output["quantized_module_count"] == 1 + + +@pytest.mark.parametrize("runtime_mode", ["triton_packed_matmul", "native_packed_matmul"]) +def test_cli_validate_artifact_rejects_unexpected_runtime_mode(capsys, tmp_path, runtime_mode): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + assert ( + main( + [ + "validate-artifact", + "--artifact", + str(tmp_path), + "--runtime-mode", + runtime_mode, + ] + ) + == 1 + ) + + assert "runtime_mode mismatch" in capsys.readouterr().err + + +def test_cli_validate_artifact_checks_policy_inventory(capsys, tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + inventory_path = tmp_path / "policy-inventory.json" + inventory_path.write_text( + json.dumps( + { + "source_model_id": "example/model", + "target_policy": "generic_dit", + "component": "transformer", + "load_mode": "config", + "linear_module_count": 1, + "action_counts": { + "orbitquant": 1, + "adaln_int4_rtn": 0, + "bf16_skip": 0, + }, + "quantized_modules": summary.quantized_modules, + "adaln_modules": summary.adaln_modules, + "skipped_modules": summary.skipped_modules, + } + ) + + "\n" + ) + + assert ( + main( + [ + "validate-artifact", + "--artifact", + str(tmp_path), + "--policy-inventory", + str(inventory_path), + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + policy_validation = output["policy_inventory_validation"] + assert policy_validation["valid"] is True + assert policy_validation["quantized_module_count"] == 1 + assert policy_validation["inventory_path"] == str(inventory_path) + + +def test_cli_validate_artifact_rejects_policy_inventory_mismatch(tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + inventory_path = tmp_path / "policy-inventory.json" + inventory_path.write_text( + json.dumps( + { + "source_model_id": "example/model", + "target_policy": "generic_dit", + "quantized_modules": ["transformer_blocks.0.attn.to_k"], + "adaln_modules": [], + "skipped_modules": [], + } + ) + + "\n" + ) + + with pytest.raises(RuntimeError, match="policy inventory mismatch"): + main( + [ + "validate-artifact", + "--artifact", + str(tmp_path), + "--policy-inventory", + str(inventory_path), + ] + ) + + +def test_cli_validate_artifact_rejects_policy_inventory_component_mismatch(tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + component="transformer", + ) + inventory_path = tmp_path / "policy-inventory.json" + inventory_path.write_text( + json.dumps( + { + "source_model_id": "example/model", + "target_policy": "generic_dit", + "component": "denoiser", + "quantized_modules": summary.quantized_modules, + "adaln_modules": summary.adaln_modules, + "skipped_modules": summary.skipped_modules, + } + ) + + "\n" + ) + + with pytest.raises(RuntimeError, match="component"): + main( + [ + "validate-artifact", + "--artifact", + str(tmp_path), + "--policy-inventory", + str(inventory_path), + ] + ) + + +def test_cli_validate_artifact_rejects_manifest_policy_config_mismatch(tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + manifest_path = tmp_path / "orbitquant_manifest.json" + manifest = json.loads(manifest_path.read_text()) + manifest["target_policy"] = "flux" + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + write_sha256sums_from_manifest(tmp_path, manifest["checksums"]) + inventory_path = tmp_path / "policy-inventory.json" + inventory_path.write_text( + json.dumps( + { + "source_model_id": "example/model", + "target_policy": "flux", + "component": "transformer", + "quantized_modules": summary.quantized_modules, + "adaln_modules": summary.adaln_modules, + "skipped_modules": summary.skipped_modules, + } + ) + + "\n" + ) + + with pytest.raises(RuntimeError, match="quantization_config mismatch.*target_policy"): + main( + [ + "validate-artifact", + "--artifact", + str(tmp_path), + "--policy-inventory", + str(inventory_path), + ] + ) + + +def test_cli_repair_artifact_metadata_wires_provenance_options(capsys, tmp_path, monkeypatch): + seen = {} + + def fake_repair_artifact_metadata(artifact, **kwargs): + seen["artifact"] = artifact + seen["kwargs"] = kwargs + return { + "artifact_dir": artifact, + "updated": { + "quantization_device": kwargs["quantization_device"], + "weight_quantization_backend": kwargs["weight_quantization_backend"], + "quantization_staging_mode": kwargs["quantization_staging_mode"], + }, + } + + monkeypatch.setattr(cli_main, "repair_artifact_metadata", fake_repair_artifact_metadata) + + assert ( + main( + [ + "repair-artifact-metadata", + "--artifact", + str(tmp_path), + "--quantization-device", + "cuda", + "--weight-quantization-backend", + "triton_cuda", + "--quantization-staging-mode", + "component", + "--skip-tensor-validation", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["updated"]["quantization_device"] == "cuda" + assert output["updated"]["weight_quantization_backend"] == "triton_cuda" + assert output["updated"]["quantization_staging_mode"] == "component" + assert seen == { + "artifact": str(tmp_path), + "kwargs": { + "quantization_device": "cuda", + "weight_quantization_backend": "triton_cuda", + "quantization_staging_mode": "component", + "validate_tensors": False, + }, + } diff --git a/tests/test_cli_compare_and_validate.py b/tests/test_cli_compare_and_validate.py new file mode 100644 index 0000000..dea9ad3 --- /dev/null +++ b/tests/test_cli_compare_and_validate.py @@ -0,0 +1,606 @@ +import json +from types import SimpleNamespace + +import pytest +import torch +from PIL import Image + +import orbitquant.cli.main as cli_main +from orbitquant.artifacts import save_orbitquant_artifact +from orbitquant.cli.main import main +from orbitquant.config import OrbitQuantConfig +from orbitquant.eval import get_native_suite +from orbitquant.modeling import quantize_linear_modules + + +def test_cli_compare_native_dry_run_prints_artifact_and_native_settings(capsys, tmp_path): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + assert ( + main( + [ + "compare-native", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--source-model", + str(tmp_path / "source-model"), + "--prompt", + "A native prompt", + "--output", + str(tmp_path / "comparison"), + "--seed", + "11", + "--device", + "cpu", + "--dtype", + "float32", + "--runtime-mode", + "triton_packed_matmul", + "--activation-kernel-backend", + "cpu", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["model_id"] == "example/artifact-model" + assert output["source_model"] == str(tmp_path / "source-model") + assert output["bit_setting"] == "W4A4" + assert output["runtime_mode"] == "triton_packed_matmul" + assert output["activation_kernel_backend"] == "cpu" + assert output["pipeline_kwargs"]["height"] == 1024 + assert output["pipeline_kwargs"]["width"] == 1024 + + +def test_cli_compare_native_runs_original_and_quantized_side_by_side( + monkeypatch, + capsys, + tmp_path, +): + monkeypatch.setattr( + "orbitquant.layers._native_cpu_packed_matmul_load_error", + lambda: RuntimeError("reference-path test"), + ) + + class TinyPipeline: + def __init__(self, color="red"): + self.color = color + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + self.device = None + + def to(self, device): + self.device = device + return self + + def __call__(self, **kwargs): + self.kwargs = kwargs + layer = self.transformer.transformer_blocks[0]["attn"]["to_q"] + layer(torch.zeros(1, 1, 8)) + return SimpleNamespace(images=[Image.new("RGB", (16, 16), self.color)]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path / "artifact", + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + pipelines = iter([TinyPipeline("red"), TinyPipeline("blue")]) + + def fake_load_pipeline_for_suite(suite, *, model_id=None, **kwargs): + assert suite.name == "flux2-native" + assert model_id == "example/artifact-model" + assert kwargs["torch_dtype"] is torch.float32 + return next(pipelines) + + validation_calls = [] + + def fake_validate_compare_native_bundle(bundle_dir): + validation_calls.append(bundle_dir) + return { + "valid": True, + "orbitquant_runtime": { + "runtime_mode_counts": { + "dequant_bf16": 1, + }, + }, + } + + monkeypatch.setattr(cli_main, "load_pipeline_for_suite", fake_load_pipeline_for_suite) + monkeypatch.setattr( + cli_main, + "_validate_compare_native_bundle", + fake_validate_compare_native_bundle, + ) + + assert ( + main( + [ + "compare-native", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path / "artifact"), + "--prompt", + "A native prompt", + "--output", + str(tmp_path / "comparison"), + "--seed", + "4", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + summary = json.loads((tmp_path / "comparison" / "summary.json").read_text()) + comparison_path = ( + tmp_path / "comparison" / "flux2-native_seed4_W4A4_original_vs_orbitquant.webp" + ) + assert output["summary_path"] == str(tmp_path / "comparison" / "summary.json") + assert output["comparison_path"] == str(comparison_path) + assert output["validation"]["valid"] is True + assert output["validation"]["orbitquant_runtime"]["runtime_mode_counts"] == {"dequant_bf16": 1} + assert validation_calls == [tmp_path / "comparison"] + assert comparison_path.is_file() + assert summary["original"]["output_path"].endswith("flux2-native_seed4_original.png") + assert summary["orbitquant"]["output_path"].endswith("flux2-native_seed4_W4A4.png") + assert summary["model_id"] == "example/artifact-model" + assert summary["source_model"] == "example/artifact-model" + assert summary["runtime_mode"] == "auto_fused" + assert summary["enable_model_cpu_offload"] is False + assert summary["available_backends"]["cpu"] is True + assert summary["orbitquant"]["runtime"] == { + "orbitquant_linear_count": 1, + "executed_module_count": 1, + "runtime_mode_counts": {"dequant_bf16": 1}, + "activation_kernel_backend_counts": {"cpu": 1}, + "forward_device_type_counts": {"cpu": 1}, + "unexecuted_module_count": 0, + "unexecuted_module_sample": [], + } + assert not any((tmp_path / "artifact" / "assets").iterdir()) + + +def test_cli_compare_native_can_skip_post_generation_validation( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self, color="red"): + self.color = color + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + return self + + def __call__(self, **kwargs): + layer = self.transformer.transformer_blocks[0]["attn"]["to_q"] + layer(torch.zeros(1, 1, 8)) + return SimpleNamespace(images=[Image.new("RGB", (16, 16), self.color)]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path / "artifact", + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + pipelines = iter([TinyPipeline("red"), TinyPipeline("blue")]) + + monkeypatch.setattr( + cli_main, + "load_pipeline_for_suite", + lambda *args, **kwargs: next(pipelines), + ) + monkeypatch.setattr( + cli_main, + "_validate_compare_native_bundle", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("post-generation validation should be skipped") + ), + ) + + assert ( + main( + [ + "compare-native", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path / "artifact"), + "--prompt", + "A native prompt", + "--output", + str(tmp_path / "comparison"), + "--device", + "cpu", + "--dtype", + "float32", + "--skip-comparison-validation", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["validation"] is None + assert (tmp_path / "comparison" / "summary.json").is_file() + + +def test_cli_validate_comparison_accepts_copied_runpod_bundle(capsys, tmp_path): + suite = get_native_suite("flux2-native") + original_path = tmp_path / "flux2-native_seed4_original.png" + orbitquant_path = tmp_path / "flux2-native_seed4_W4A4.png" + comparison_path = tmp_path / "flux2-native_seed4_W4A4_original_vs_orbitquant.webp" + original_image = Image.effect_noise((suite.width, suite.height), 32).convert("RGB") + orbitquant_image = Image.effect_noise((suite.width, suite.height), 32).convert("RGB") + original_image.save(original_path) + orbitquant_image.save(orbitquant_path) + comparison_sheet = Image.new("RGB", (suite.width * 2, suite.height + 24), "white") + comparison_sheet.paste(original_image, (0, 24)) + comparison_sheet.paste(orbitquant_image, (suite.width, 24)) + comparison_sheet.save(comparison_path) + + for path, bit_setting in ( + (original_path, "original"), + (orbitquant_path, "W4A4"), + ): + metadata = { + "suite": suite.name, + "model_id": "example/artifact-model", + "prompt": "A native prompt", + "seed": 4, + "height": suite.height, + "width": suite.width, + "frames": suite.frames, + "export_fps": suite.export_fps, + "steps": suite.steps, + "guidance": suite.guidance, + "quantization": None + if bit_setting == "original" + else {"config": {"weight_bits": 4, "activation_bits": 4}}, + } + path.with_suffix(path.suffix + ".json").write_text( + json.dumps(metadata) + "\n", + encoding="utf-8", + ) + + (tmp_path / "summary.json").write_text( + json.dumps( + { + "suite": suite.name, + "model_id": "example/artifact-model", + "source_model": "/workspace/hf-models/example", + "artifact": "/workspace/artifacts/example", + "component": "transformer", + "prompt": "A native prompt", + "seed": 4, + "height": suite.height, + "width": suite.width, + "frames": suite.frames, + "steps": suite.steps, + "guidance": suite.guidance, + "dtype": "bfloat16", + "device": "cuda", + "bit_setting": "W4A4", + "runtime_mode": "triton_packed_matmul", + "activation_kernel_backend": "triton_cuda", + "enable_model_cpu_offload": True, + "available_backends": {"cpu": True, "mps": False, "triton_cuda": True}, + "original": { + "output_path": f"/workspace/run/{original_path.name}", + "metadata_path": f"/workspace/run/{original_path.name}.json", + "wall_time_seconds": 10.0, + "peak_vram_bytes": 1000, + }, + "orbitquant": { + "output_path": f"/workspace/run/{orbitquant_path.name}", + "metadata_path": f"/workspace/run/{orbitquant_path.name}.json", + "wall_time_seconds": 15.0, + "peak_vram_bytes": 900, + "runtime": { + "runtime_mode_counts": {"triton_packed_matmul": 1}, + "activation_kernel_backend_counts": {"triton_cuda": 1}, + }, + }, + "comparison_path": f"/workspace/run/{comparison_path.name}", + } + ) + + "\n", + encoding="utf-8", + ) + + assert main(["validate-comparison", "--input", str(tmp_path)]) == 0 + + output = json.loads(capsys.readouterr().out) + assert output["valid"] is True + assert output["speed_ratio_orbitquant_over_original"] == 1.5 + assert output["comparison"]["size"] == [suite.width * 2, suite.height + 24] + assert output["orbitquant_runtime"]["runtime_mode_counts"] == {"triton_packed_matmul": 1} + + +def test_cli_validate_comparison_accepts_wan_video_contact_sheets(capsys, tmp_path): + from orbitquant.eval.assets import create_image_comparison_sheet + + suite = get_native_suite("wan-native") + original_path = tmp_path / "wan-native_seed4_original.mp4" + orbitquant_path = tmp_path / "wan-native_seed4_W4A4.mp4" + original_contact_sheet = tmp_path / "wan-native_seed4_original_contact_sheet.webp" + orbitquant_contact_sheet = tmp_path / "wan-native_seed4_W4A4_contact_sheet.webp" + comparison_path = tmp_path / "wan-native_seed4_W4A4_original_vs_orbitquant.webp" + contact_sheet_size = (suite.width * 4, suite.height * 3) + + for path in (original_path, orbitquant_path): + path.write_bytes(b"fake-mp4") + for path, colors in ( + (original_contact_sheet, ((0, 0, 0), (255, 255, 255))), + (orbitquant_contact_sheet, ((255, 0, 0), (0, 0, 255))), + ): + sheet = Image.new("RGB", contact_sheet_size, colors[0]) + sheet.paste(colors[1], (0, 0, contact_sheet_size[0] // 2, contact_sheet_size[1])) + sheet.save(path) + create_image_comparison_sheet( + original_contact_sheet, + orbitquant_contact_sheet, + comparison_path, + labels=("BF16", "OrbitQuant W4A4"), + ) + + for path, contact_sheet, bit_setting in ( + (original_path, original_contact_sheet, "original"), + (orbitquant_path, orbitquant_contact_sheet, "W4A4"), + ): + metadata = { + "suite": suite.name, + "model_id": "example/wan-artifact-model", + "prompt": "A native video prompt", + "seed": 4, + "height": suite.height, + "width": suite.width, + "frames": suite.frames, + "export_fps": suite.export_fps, + "steps": suite.steps, + "guidance": suite.guidance, + "contact_sheet_path": f"/workspace/run/{contact_sheet.name}", + "quantization": None + if bit_setting == "original" + else {"config": {"weight_bits": 4, "activation_bits": 4}}, + } + path.with_suffix(path.suffix + ".json").write_text( + json.dumps(metadata) + "\n", + encoding="utf-8", + ) + + (tmp_path / "summary.json").write_text( + json.dumps( + { + "suite": suite.name, + "model_id": "example/wan-artifact-model", + "source_model": "/workspace/hf-models/example-wan", + "artifact": "/workspace/artifacts/example-wan", + "component": "transformer", + "prompt": "A native video prompt", + "seed": 4, + "height": suite.height, + "width": suite.width, + "frames": suite.frames, + "steps": suite.steps, + "guidance": suite.guidance, + "dtype": "bfloat16", + "device": "cuda", + "bit_setting": "W4A4", + "runtime_mode": "triton_packed_matmul", + "activation_kernel_backend": "triton_cuda", + "enable_model_cpu_offload": True, + "available_backends": {"cpu": True, "mps": False, "triton_cuda": True}, + "original": { + "output_path": f"/workspace/run/{original_path.name}", + "metadata_path": f"/workspace/run/{original_path.name}.json", + "wall_time_seconds": 10.0, + "peak_vram_bytes": 1000, + }, + "orbitquant": { + "output_path": f"/workspace/run/{orbitquant_path.name}", + "metadata_path": f"/workspace/run/{orbitquant_path.name}.json", + "wall_time_seconds": 14.0, + "peak_vram_bytes": 900, + "runtime": { + "runtime_mode_counts": {"triton_packed_matmul": 1}, + "activation_kernel_backend_counts": {"triton_cuda": 1}, + }, + }, + "comparison_path": f"/workspace/run/{comparison_path.name}", + } + ) + + "\n", + encoding="utf-8", + ) + + assert main(["validate-comparison", "--input", str(tmp_path)]) == 0 + + output = json.loads(capsys.readouterr().out) + assert output["valid"] is True + assert output["native_settings"]["frames"] == 81 + assert output["original"]["contact_sheet"]["size"] == [ + contact_sheet_size[0], + contact_sheet_size[1], + ] + assert output["orbitquant"]["contact_sheet"]["size"] == [ + contact_sheet_size[0], + contact_sheet_size[1], + ] + assert output["comparison"]["size"] == [ + contact_sheet_size[0] * 2, + contact_sheet_size[1] + 24, + ] + assert output["speed_ratio_orbitquant_over_original"] == 1.4 + + +def test_cli_validate_comparison_rejects_blank_images(tmp_path): + suite = get_native_suite("flux2-native") + for name, quantization in ( + ("flux2-native_seed4_original.png", None), + ( + "flux2-native_seed4_W4A4.png", + {"config": {"weight_bits": 4, "activation_bits": 4}}, + ), + ): + path = tmp_path / name + Image.new("RGB", (suite.width, suite.height), "black").save(path) + path.with_suffix(path.suffix + ".json").write_text( + json.dumps( + { + "suite": suite.name, + "model_id": "example/artifact-model", + "prompt": "A native prompt", + "seed": 4, + "height": suite.height, + "width": suite.width, + "frames": suite.frames, + "export_fps": suite.export_fps, + "steps": suite.steps, + "guidance": suite.guidance, + "quantization": quantization, + } + ) + + "\n", + encoding="utf-8", + ) + Image.new("RGB", (suite.width * 2, suite.height + 24), "black").save( + tmp_path / "flux2-native_seed4_W4A4_original_vs_orbitquant.webp" + ) + (tmp_path / "summary.json").write_text( + json.dumps( + { + "suite": suite.name, + "model_id": "example/artifact-model", + "prompt": "A native prompt", + "seed": 4, + "bit_setting": "W4A4", + "original": { + "output_path": "/workspace/run/flux2-native_seed4_original.png", + "metadata_path": "/workspace/run/flux2-native_seed4_original.png.json", + }, + "orbitquant": { + "output_path": "/workspace/run/flux2-native_seed4_W4A4.png", + "metadata_path": "/workspace/run/flux2-native_seed4_W4A4.png.json", + }, + "comparison_path": ( + "/workspace/run/flux2-native_seed4_W4A4_original_vs_orbitquant.webp" + ), + } + ) + + "\n", + encoding="utf-8", + ) + + with pytest.raises(RuntimeError, match="comparison image looks blank"): + main(["validate-comparison", "--input", str(tmp_path)]) + + +def test_cli_validate_generation_reports_valid_native_output(capsys, tmp_path): + output_path = tmp_path / "flux2-native_seed5_W4A4_simple-object.png" + Image.new("RGB", (16, 16), "red").save(output_path) + output_path.with_suffix(".png.json").write_text( + json.dumps( + { + "suite": "flux2-native", + "model_id": "black-forest-labs/FLUX.2-klein-4B", + "prompt": "A native image", + "seed": 5, + "height": 1024, + "width": 1024, + "frames": None, + "steps": 4, + "guidance": 1.0, + "quantization": { + "config": { + "weight_bits": 4, + "activation_bits": 4, + } + }, + } + ) + + "\n" + ) + + assert ( + main( + [ + "validate-generation", + "--suite", + "flux2-native", + "--output", + str(output_path), + "--seed", + "5", + "--bit-setting", + "W4A4", + "--prompt", + "A native image", + ] + ) + == 0 + ) + + payload = json.loads(capsys.readouterr().out) + assert payload["valid"] is True + assert payload["suite"] == "flux2-native" + assert payload["bit_setting"] == "W4A4" diff --git a/tests/test_cli_generate.py b/tests/test_cli_generate.py new file mode 100644 index 0000000..4ea9c58 --- /dev/null +++ b/tests/test_cli_generate.py @@ -0,0 +1,748 @@ +import json +import sys +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image + +import orbitquant.cli.main as cli_main +from orbitquant.artifacts import save_orbitquant_artifact +from orbitquant.cli.main import main +from orbitquant.config import OrbitQuantConfig +from orbitquant.eval.prompts import default_prompt_payload, select_prompt_record +from orbitquant.modeling import quantize_linear_modules + + +def test_cli_generate_requires_prompt_and_output(): + try: + main(["generate", "--suite", "flux2-native"]) + except (SystemExit, ValueError) as exc: + assert "prompt" in str(exc) or "output" in str(exc) or getattr(exc, "code", 0) != 0 + else: + raise AssertionError("generate accepted missing prompt/output arguments") + + +def test_cli_record_metrics_imports_nested_json_into_artifact(capsys, tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="flux") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="black-forest-labs/FLUX.1-schnell", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + metrics_path = tmp_path / "geneval.json" + metrics_path.write_text( + json.dumps( + { + "overall": 0.71, + "per_task": { + "single_object": 0.92, + "counting": 0.47, + }, + "notes": "ignored", + } + ) + ) + + assert ( + main( + [ + "record-metrics", + "--artifact", + str(tmp_path), + "--split", + "orbitquant", + "--metrics-json", + str(metrics_path), + "--metric-prefix", + "geneval", + "--suite", + "flux1-schnell-native", + "--seed", + "0", + "--bit-setting", + "W4A4", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + record = json.loads((tmp_path / "benchmark" / "orbitquant.metrics.jsonl").read_text()) + csv_text = (tmp_path / "benchmark" / "orbitquant.metrics.csv").read_text() + manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) + assert output["metrics"] == { + "geneval_overall": 0.71, + "geneval_per_task_counting": 0.47, + "geneval_per_task_single_object": 0.92, + } + assert record["metrics"] == output["metrics"] + assert record["metadata"] == { + "suite": "flux1-schnell-native", + "seed": 0, + "bit_setting": "W4A4", + "metrics_source": str(metrics_path), + } + assert "geneval_per_task_single_object,0.92" in csv_text + assert "benchmark/orbitquant.metrics.jsonl" in manifest["checksums"] + + report_dir = tmp_path / "reports" + assert main(["report", "--artifact", str(tmp_path), "--output", str(report_dir)]) == 0 + + report_output = json.loads(capsys.readouterr().out) + missing_table = (report_dir / "tables" / "missing_required_metrics.csv").read_text() + assert report_output["missing_required_metric_count"] == 11 + assert "orbitquant,flux1-schnell-native,geneval_overall" not in missing_table + assert "orbitquant,flux1-schnell-native,geneval_per_task_single_object" not in (missing_table) + assert "orbitquant,flux1-schnell-native,geneval_per_task_counting" not in (missing_table) + assert "original,flux1-schnell-native,geneval_overall" in missing_table + assert "orbitquant,flux1-schnell-native,geneval_per_task_color_attr" in missing_table + + +def test_cli_generate_dry_run_prints_native_request(capsys, tmp_path): + assert ( + main( + [ + "generate", + "--suite", + "flux2-native", + "--prompt", + "A native prompt", + "--output", + str(tmp_path), + "--seed", + "9", + "--device", + "cpu", + "--enable-model-cpu-offload", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["model_id"] == "black-forest-labs/FLUX.2-klein-4B" + assert output["pipeline_kwargs"]["height"] == 1024 + assert output["pipeline_kwargs"]["width"] == 1024 + assert output["enable_model_cpu_offload"] is True + + +@pytest.mark.parametrize("runtime_mode", ["triton_packed_matmul", "native_packed_matmul"]) +def test_cli_generate_dry_run_prints_quantized_native_request(capsys, tmp_path, runtime_mode): + assert ( + main( + [ + "generate", + "--suite", + "wan-native", + "--prompt", + "A native prompt", + "--output", + str(tmp_path), + "--seed", + "9", + "--device", + "cpu", + "--bit-setting", + "W4A6", + "--activation-kernel-backend", + "cpu", + "--runtime-mode", + runtime_mode, + "--dry-run", + ] + ) + == 0 + ) + + output = capsys.readouterr().out + assert '"height": 480' in output + assert '"width": 832' in output + assert '"weight_bits": 4' in output + assert '"activation_bits": 6' in output + assert '"activation_kernel_backend": "cpu"' in output + assert f'"runtime_mode": "{runtime_mode}"' in output + assert '"target_policy": "wan"' in output + + +def test_cli_generate_with_packed_runtime_on_the_fly_quantization_skips_dequant_prewarm( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + + def to(self, device): + return self + + def __call__(self, **kwargs): + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "blue")]) + + def fake_apply_quantization(pipeline, suite, config): + assert config.runtime_mode == "triton_packed_matmul" + return SimpleNamespace( + quantized_modules=["transformer_blocks.0.attn.to_q"], + adaln_modules=[], + skipped_modules=[], + ) + + def fail_prewarm(*args, **kwargs): + raise AssertionError("packed runtime must not materialize dequant prewarm") + + monkeypatch.setattr(cli_main, "load_pipeline_for_suite", lambda *args, **kwargs: TinyPipeline()) + monkeypatch.setattr(cli_main, "apply_quantization_to_pipeline", fake_apply_quantization) + monkeypatch.setattr(cli_main, "_prewarm_pipeline_component", fail_prewarm) + + assert ( + main( + [ + "generate", + "--suite", + "flux2-native", + "--prompt", + "A native prompt", + "--output", + str(tmp_path), + "--bit-setting", + "W4A4", + "--runtime-mode", + "triton_packed_matmul", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["output_path"].endswith("flux2-native_seed0_W4A4.png") + + +def test_cli_generate_dry_run_uses_artifact_source_and_default_assets_output(capsys, tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + assert ( + main( + [ + "generate", + "--suite", + "flux2-native", + "--prompt", + "A native prompt", + "--artifact", + str(tmp_path), + "--device", + "cpu", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["model_id"] == "example/artifact-model" + assert output["artifact"] == str(tmp_path) + assert output["output"] == str(tmp_path / "assets") + assert output["quantization_config"]["weight_bits"] == 4 + assert output["quantization_config"]["activation_bits"] == 4 + + +def test_cli_generate_dry_run_selects_prompt_from_artifact_pack(capsys, tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + assert ( + main( + [ + "generate", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-id", + "english-text-rendering", + "--device", + "cpu", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["prompt_record"]["id"] == "english-text-rendering" + expected_prompt = select_prompt_record( + default_prompt_payload("flux2"), prompt_id="english-text-rendering" + )["prompt"] + assert output["pipeline_kwargs"]["prompt"] == expected_prompt + + +def test_cli_generate_with_artifact_loads_component_and_records_metrics( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + self.device = None + + def to(self, device): + self.device = device + return self + + def __call__(self, **kwargs): + self.kwargs = kwargs + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "blue")]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + restored = TinyPipeline() + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + assert kwargs["torch_dtype"] is torch.float32 + return restored + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + + assert ( + main( + [ + "generate", + "--suite", + "flux2-native", + "--prompt", + "A native prompt", + "--artifact", + str(tmp_path), + "--seed", + "3", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + record = json.loads((tmp_path / "benchmark" / "orbitquant.metrics.jsonl").read_text()) + assert output["artifact"] == str(tmp_path) + assert output["output_path"].endswith("flux2-native_seed3_W4A4.png") + assert (tmp_path / "assets" / "flux2-native_seed3_W4A4.png").is_file() + assert record["split"] == "orbitquant" + assert record["metrics"]["generated_samples"] == 1 + assert record["metrics"]["wall_time_seconds"] >= 0.0 + assert record["metadata"]["output_path"] == output["output_path"] + assert record["metadata"]["device"] == "cpu" + assert record["metadata"]["dtype"] == "float32" + sample_metadata = json.loads( + (tmp_path / "assets" / "flux2-native_seed3_W4A4.png.json").read_text() + ) + assert sample_metadata["device"] == "cpu" + assert sample_metadata["dtype"] == "float32" + assert sample_metadata["pipeline_class"] == "TinyPipeline" + manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) + assert "assets/flux2-native_seed3_W4A4.png" in manifest["checksums"] + assert "assets/flux2-native_seed3_W4A4.png.json" in manifest["checksums"] + assert restored.device == "cpu" + + +def test_cli_generate_with_packed_runtime_artifact_skips_dequant_prewarm( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + return self + + def __call__(self, **kwargs): + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "blue")]) + + source = TinyPipeline() + config = OrbitQuantConfig( + block_size=4, + target_policy="generic_dit", + runtime_mode="triton_packed_matmul", + ) + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + return TinyPipeline() + + def fail_prewarm(*args, **kwargs): + raise AssertionError("packed runtime must not materialize dequant prewarm") + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + monkeypatch.setattr(cli_main, "_prewarm_pipeline_component", fail_prewarm) + + assert ( + main( + [ + "generate", + "--suite", + "flux2-native", + "--prompt", + "A native prompt", + "--artifact", + str(tmp_path), + "--seed", + "3", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["output_path"].endswith("flux2-native_seed3_W4A4.png") + + +def test_cli_generate_with_artifact_original_split_skips_quantized_component( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + return self + + def __call__(self, **kwargs): + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "yellow")]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + return TinyPipeline() + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + monkeypatch.setattr( + cli_main, + "load_quantized_pipeline_component", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("original split must not load quantized component") + ), + ) + + assert ( + main( + [ + "generate", + "--suite", + "flux2-native", + "--prompt", + "A native prompt", + "--artifact", + str(tmp_path), + "--split", + "original", + "--seed", + "3", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + record = json.loads((tmp_path / "benchmark" / "original.metrics.jsonl").read_text()) + assert output["output_path"].endswith("flux2-native_seed3_original.png") + assert record["split"] == "original" + assert record["metadata"]["bit_setting"] == "original" + assert not (tmp_path / "benchmark" / "orbitquant.metrics.jsonl").exists() + + +def test_cli_generate_creates_comparison_when_original_pair_exists( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + color = "red" + + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + return self + + def __call__(self, **kwargs): + return SimpleNamespace(images=[Image.new("RGB", (16, 16), self.color)]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + return TinyPipeline() + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + + common_args = [ + "generate", + "--suite", + "flux2-native", + "--prompt", + "A native prompt", + "--artifact", + str(tmp_path), + "--seed", + "4", + "--device", + "cpu", + "--dtype", + "float32", + ] + assert main([*common_args, "--split", "original"]) == 0 + capsys.readouterr() + TinyPipeline.color = "blue" + assert main(common_args) == 0 + + output = json.loads(capsys.readouterr().out) + manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) + comparison = "assets/original_vs_orbitquant_flux2-native_seed4_W4A4_prompt.webp" + assert comparison in manifest["checksums"] + assert (tmp_path / comparison).is_file() + assert output["artifact_comparisons"] == [comparison] + + +def test_cli_generate_with_video_artifact_records_contact_sheet_asset( + monkeypatch, + capsys, + tmp_path, +): + class TinyWanPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + self.device = device + return self + + def __call__(self, **kwargs): + self.kwargs = kwargs + return SimpleNamespace(frames=np.zeros((1, 2, 8, 8, 3), dtype=np.uint8)) + + export_calls = [] + + def fake_export_to_video(frames, path, **kwargs): + export_calls.append(kwargs) + with open(path, "wb") as output: + output.write(b"fake mp4") + + source = TinyWanPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + restored = TinyWanPipeline() + + class FakeWanPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + assert kwargs["torch_dtype"] is torch.float32 + return restored + + monkeypatch.setitem(sys.modules, "diffusers", SimpleNamespace(WanPipeline=FakeWanPipeline)) + monkeypatch.setitem( + sys.modules, + "diffusers.utils", + SimpleNamespace(export_to_video=fake_export_to_video), + ) + + assert ( + main( + [ + "generate", + "--suite", + "wan-native", + "--prompt", + "A native video prompt", + "--artifact", + str(tmp_path), + "--seed", + "2", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) + metrics_record = json.loads((tmp_path / "benchmark" / "orbitquant.metrics.jsonl").read_text()) + metadata = json.loads((tmp_path / "assets" / "wan-native_seed2_W4A4.mp4.json").read_text()) + assert output["output_path"].endswith("wan-native_seed2_W4A4.mp4") + assert export_calls == [{"fps": 16}] + assert metadata["export_fps"] == 16 + assert metrics_record["metadata"]["export_fps"] == 16 + assert "assets/wan-native_seed2_W4A4.mp4" in manifest["checksums"] + assert "assets/wan-native_seed2_W4A4.mp4.json" in manifest["checksums"] + assert "assets/wan-native_seed2_W4A4_contact_sheet.webp" in manifest["checksums"] diff --git a/tests/test_cli_generate_pack.py b/tests/test_cli_generate_pack.py new file mode 100644 index 0000000..ab6817e --- /dev/null +++ b/tests/test_cli_generate_pack.py @@ -0,0 +1,762 @@ +import json +import sys +from types import SimpleNamespace + +import torch +from PIL import Image + +import orbitquant.cli.main as cli_main +from orbitquant.artifacts import save_orbitquant_artifact, validate_orbitquant_artifact +from orbitquant.cli.main import main +from orbitquant.config import OrbitQuantConfig +from orbitquant.eval.prompts import default_prompt_payload, select_prompt_record +from orbitquant.modeling import quantize_linear_modules + + +def test_cli_generate_pack_dry_run_lists_prompt_seed_jobs(capsys, tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-limit", + "2", + "--seeds", + "0,1", + "--device", + "cpu", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["job_count"] == 4 + assert output["jobs"][0]["prompt_record"]["id"] == "simple-object" + assert output["jobs"][0]["seed"] == 0 + assert output["jobs"][2]["seed"] == 1 + assert output["output"] == str(tmp_path / "assets") + + +def test_cli_generate_pack_dry_run_accepts_geneval_smoke_prompt_pack(capsys, tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux1-schnell-native", + "--artifact", + str(tmp_path), + "--prompt-pack", + "geneval-smoke", + "--prompt-limit", + "2", + "--seeds", + "0", + "--device", + "cpu", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["prompt_pack"] == "geneval_smoke_v1" + assert output["job_count"] == 2 + assert output["jobs"][0]["prompt_record"]["id"].startswith("geneval-00000-") + assert output["jobs"][0]["prompt_record"]["geneval"]["tag"] == "single_object" + + +def test_cli_generate_pack_dry_run_accepts_geneval_metadata_jsonl(capsys, tmp_path): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + metadata_jsonl = tmp_path / "evaluation_metadata.jsonl" + metadata_jsonl.write_text( + json.dumps( + { + "tag": "single_object", + "include": [{"class": "bench", "count": 1}], + "prompt": "a photo of a bench", + } + ) + + "\n", + encoding="utf-8", + ) + + assert ( + main( + [ + "generate-pack", + "--suite", + "z-image-native", + "--artifact", + str(tmp_path), + "--prompt-metadata-jsonl", + str(metadata_jsonl), + "--seeds", + "4", + "--device", + "cpu", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["prompt_pack"] == "geneval_metadata_jsonl" + assert output["job_count"] == 1 + assert output["jobs"][0]["seed"] == 4 + assert output["jobs"][0]["prompt_record"]["geneval"]["include"] == [ + {"class": "bench", "count": 1} + ] + + +def test_cli_generate_pack_runs_jobs_once_per_prompt_seed_and_records_artifacts( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + self.calls = [] + + def to(self, device): + self.device = device + return self + + def __call__(self, **kwargs): + self.calls.append(kwargs) + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + restored = TinyPipeline() + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + assert kwargs["torch_dtype"] is torch.float32 + return restored + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-id", + "simple-object", + "--prompt-id", + "counting", + "--seeds", + "3", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) + metrics_rows = (tmp_path / "benchmark" / "orbitquant.metrics.jsonl").read_text().splitlines() + assert output["job_count"] == 2 + assert len(restored.calls) == 2 + assert "assets/flux2-native_seed3_W4A4_simple-object.png" in manifest["checksums"] + assert "assets/flux2-native_seed3_W4A4_counting.png" in manifest["checksums"] + assert len(metrics_rows) == 2 + assert json.loads(metrics_rows[0])["metadata"]["prompt_record"]["id"] == "simple-object" + + +def test_cli_generate_pack_with_packed_runtime_artifact_skips_dequant_prewarm( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + return self + + def __call__(self, **kwargs): + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) + + source = TinyPipeline() + config = OrbitQuantConfig( + block_size=4, + target_policy="generic_dit", + runtime_mode="triton_packed_matmul", + ) + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + return TinyPipeline() + + def fail_prewarm(*args, **kwargs): + raise AssertionError("packed runtime must not materialize dequant prewarm") + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + monkeypatch.setattr(cli_main, "_prewarm_pipeline_component", fail_prewarm) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-id", + "simple-object", + "--seeds", + "3", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["run_count"] == 1 + assert output["outputs"][0]["output_path"].endswith("flux2-native_seed3_W4A4_simple-object.png") + + +def test_cli_generate_pack_skip_checksums_refreshes_artifact_once_at_end( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + self.calls = [] + + def to(self, device): + return self + + def __call__(self, **kwargs): + self.calls.append(kwargs) + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + restored = TinyPipeline() + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + return restored + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-id", + "simple-object", + "--seeds", + "3", + "--resume-existing", + "--skip-artifact-checksums", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + manifest = json.loads((tmp_path / "orbitquant_manifest.json").read_text()) + validation = validate_orbitquant_artifact(tmp_path) + image_path = "assets/flux2-native_seed3_W4A4_simple-object.png" + metadata_path = "assets/flux2-native_seed3_W4A4_simple-object.png.json" + assert output["job_count"] == 1 + assert output["run_count"] == 1 + assert output["checksum_refresh"]["checksum_count"] == len(manifest["checksums"]) + assert validation["valid"] is True + assert image_path in manifest["checksums"] + assert metadata_path in manifest["checksums"] + assert "benchmark/orbitquant.metrics.jsonl" in manifest["checksums"] + assert output["artifact_comparisons"] == [] + assert output["outputs"][0]["comparisons"] == [] + + +def test_cli_generate_pack_defers_comparison_creation_until_after_jobs( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + return self + + def __call__(self, **kwargs): + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + return TinyPipeline() + + comparison_calls = [] + + def fake_create_comparisons(artifact_dir, **kwargs): + comparison_calls.append({"artifact_dir": artifact_dir, "kwargs": kwargs}) + return ["assets/fake-comparison.webp"] + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + monkeypatch.setattr(cli_main, "create_artifact_image_comparisons", fake_create_comparisons) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-id", + "simple-object", + "--prompt-id", + "counting", + "--seeds", + "3", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert len(comparison_calls) == 1 + assert comparison_calls[0]["kwargs"]["comparison_keys"] == { + ("flux2-native", 3, "simple-object"), + ("flux2-native", 3, "counting"), + } + assert output["run_count"] == 2 + assert output["artifact_comparisons"] == ["assets/fake-comparison.webp"] + assert [item["comparisons"] for item in output["outputs"]] == [[], []] + + +def test_cli_generate_pack_prompt_metadata_disables_comparisons_by_default( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + return self + + def __call__(self, **kwargs): + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "green")]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + metadata_jsonl = tmp_path / "evaluation_metadata.jsonl" + metadata_jsonl.write_text( + json.dumps( + { + "tag": "single_object", + "include": [{"class": "bench", "count": 1}], + "prompt": "a photo of a bench", + } + ) + + "\n", + encoding="utf-8", + ) + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + return TinyPipeline() + + def fail_create_comparisons(*args, **kwargs): + raise AssertionError("GenEval metadata packs must not create comparison sheets by default") + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + monkeypatch.setattr(cli_main, "create_artifact_image_comparisons", fail_create_comparisons) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-metadata-jsonl", + str(metadata_jsonl), + "--seeds", + "3", + "--skip-artifact-checksums", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["run_count"] == 1 + assert output["artifact_comparisons"] == [] + assert output["outputs"][0]["comparisons"] == [] + + +def test_cli_generate_pack_resume_existing_skips_completed_outputs( + monkeypatch, + capsys, + tmp_path, +): + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList( + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] + ) + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(model, config) + save_orbitquant_artifact( + model, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + output_path = tmp_path / "assets" / "flux2-native_seed3_W4A4_simple-object.png" + Image.new("RGB", (16, 16), "green").save(output_path) + output_path.with_suffix(".png.json").write_text( + json.dumps( + { + "suite": "flux2-native", + "model_id": "example/artifact-model", + "prompt": select_prompt_record( + default_prompt_payload("flux2"), prompt_id="simple-object" + )["prompt"], + "seed": 3, + "height": 1024, + "width": 1024, + "frames": None, + "steps": 4, + "guidance": 1.0, + "quantization": { + "config": { + "weight_bits": 4, + "activation_bits": 4, + } + }, + } + ) + + "\n" + ) + monkeypatch.setattr( + cli_main, + "load_pipeline_for_suite", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("completed resume jobs must not load a pipeline") + ), + ) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-id", + "simple-object", + "--seeds", + "3", + "--resume-existing", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["job_count"] == 1 + assert output["run_count"] == 0 + assert output["skipped_count"] == 1 + assert output["skipped_outputs"] == [str(output_path)] + + +def test_cli_generate_pack_resume_existing_reruns_invalid_metadata( + monkeypatch, + capsys, + tmp_path, +): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + self.calls = [] + + def to(self, device): + return self + + def __call__(self, **kwargs): + self.calls.append(kwargs) + return SimpleNamespace(images=[Image.new("RGB", (16, 16), "purple")]) + + source = TinyPipeline() + config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") + summary = quantize_linear_modules(source.transformer, config) + save_orbitquant_artifact( + source.transformer, + tmp_path, + config=config, + source_model_id="example/artifact-model", + source_revision="abc123", + source_license="apache-2.0", + summary=summary, + ) + output_path = tmp_path / "assets" / "flux2-native_seed3_W4A4_simple-object.png" + Image.new("RGB", (16, 16), "green").save(output_path) + output_path.with_suffix(".png.json").write_text('{"status":"complete"}\n') + restored = TinyPipeline() + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/artifact-model" + return restored + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace(Flux2KleinPipeline=FakeDiffusionPipeline), + ) + + assert ( + main( + [ + "generate-pack", + "--suite", + "flux2-native", + "--artifact", + str(tmp_path), + "--prompt-id", + "simple-object", + "--seeds", + "3", + "--resume-existing", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["job_count"] == 1 + assert output["run_count"] == 1 + assert output["skipped_count"] == 0 + assert len(restored.calls) == 1 diff --git a/tests/test_cli_hub.py b/tests/test_cli_hub.py new file mode 100644 index 0000000..44465ab --- /dev/null +++ b/tests/test_cli_hub.py @@ -0,0 +1,591 @@ +import json + +import pytest + +import orbitquant.cli.main as cli_main +from orbitquant.cli.main import main + + +def test_cli_upload_artifact_wires_validation_and_hf_options(capsys, tmp_path, monkeypatch): + seen = {} + + def fake_upload_artifact(artifact, **kwargs): + seen["artifact"] = artifact + seen["kwargs"] = kwargs + return { + "artifact_dir": artifact, + "repo_id": kwargs["repo_id"], + "private": kwargs["private"], + "dry_run": kwargs["dry_run"], + "validation": {"valid": True}, + } + + monkeypatch.setattr(cli_main, "upload_orbitquant_artifact", fake_upload_artifact) + + assert ( + main( + [ + "upload-artifact", + "--artifact", + str(tmp_path), + "--repo-id", + "WaveCut/example-orbitquant", + "--revision", + "main", + "--commit-message", + "upload artifact", + "--public", + "--no-create-repo", + "--replace-repo-files", + "--skip-tensor-validation", + "--upload-profile", + "compact", + "--report-dir", + "/tmp/orbitquant-report", + "--staging-dir", + "/tmp/orbitquant-stage", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["repo_id"] == "WaveCut/example-orbitquant" + assert output["private"] is False + assert output["dry_run"] is True + assert seen == { + "artifact": str(tmp_path), + "kwargs": { + "repo_id": "WaveCut/example-orbitquant", + "private": False, + "create_repo": False, + "revision": "main", + "commit_message": "upload artifact", + "replace_repo_files": True, + "validate_tensors": False, + "upload_profile": "compact", + "report_dirs": ["/tmp/orbitquant-report"], + "staging_dir": "/tmp/orbitquant-stage", + "dry_run": True, + }, + } + + +def test_cli_upload_artifact_defaults_to_compact_profile(capsys, tmp_path, monkeypatch): + seen = {} + + def fake_upload_artifact(artifact, **kwargs): + seen["artifact"] = artifact + seen["kwargs"] = kwargs + return { + "artifact_dir": artifact, + "repo_id": kwargs["repo_id"], + "private": kwargs["private"], + "dry_run": kwargs["dry_run"], + "validation": {"valid": True}, + } + + monkeypatch.setattr(cli_main, "upload_orbitquant_artifact", fake_upload_artifact) + + assert ( + main( + [ + "upload-artifact", + "--artifact", + str(tmp_path), + "--repo-id", + "WaveCut/example-orbitquant", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["repo_id"] == "WaveCut/example-orbitquant" + assert seen["artifact"] == str(tmp_path) + assert seen["kwargs"]["upload_profile"] == "compact" + assert seen["kwargs"]["replace_repo_files"] is True + assert seen["kwargs"]["report_dirs"] is None + assert seen["kwargs"]["staging_dir"] is None + + +def test_cli_upload_artifact_rejects_disabling_remote_file_replacement( + capsys, tmp_path, monkeypatch +): + seen = {} + + def fake_upload_artifact(artifact, **kwargs): + seen["artifact"] = artifact + seen["kwargs"] = kwargs + return { + "artifact_dir": artifact, + "repo_id": kwargs["repo_id"], + "private": kwargs["private"], + "dry_run": kwargs["dry_run"], + "validation": {"valid": True}, + } + + monkeypatch.setattr(cli_main, "upload_orbitquant_artifact", fake_upload_artifact) + + with pytest.raises(SystemExit): + main( + [ + "upload-artifact", + "--artifact", + str(tmp_path), + "--repo-id", + "WaveCut/example-orbitquant", + "--no-replace-repo-files", + "--dry-run", + ] + ) + + assert seen == {} + assert "unrecognized arguments: --no-replace-repo-files" in capsys.readouterr().err + + +def test_cli_repair_hf_artifact_metadata_wires_single_repo_options(capsys, monkeypatch): + seen = {} + + def fake_repair_hf_artifact_metadata(**kwargs): + seen.update(kwargs) + return { + "repo_id": kwargs["repo_id"], + "dry_run": kwargs["dry_run"], + "changed_files": ["orbitquant_manifest.json"], + } + + monkeypatch.setattr(cli_main, "repair_hf_artifact_metadata", fake_repair_hf_artifact_metadata) + + assert ( + main( + [ + "repair-hf-artifact-metadata", + "--repo-id", + "WaveCut/example-orbitquant", + "--revision", + "main", + "--commit-message", + "repair metadata", + "--quantization-device", + "cuda", + "--weight-quantization-backend", + "triton_cuda", + "--quantization-staging-mode", + "component", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["repo_id"] == "WaveCut/example-orbitquant" + assert output["dry_run"] is True + assert seen == { + "repo_id": "WaveCut/example-orbitquant", + "quantization_device": "cuda", + "weight_quantization_backend": "triton_cuda", + "quantization_staging_mode": "component", + "revision": "main", + "commit_message": "repair metadata", + "dry_run": True, + } + + +def test_cli_repair_hf_native_smoke_proof_wires_single_repo_options(capsys, monkeypatch): + seen = {} + + def fake_repair_hf_native_smoke_proof(**kwargs): + seen.update(kwargs) + return { + "repo_id": kwargs["repo_id"], + "suite": kwargs["suite"].name, + "dry_run": kwargs["dry_run"], + "changed_files": ["benchmark/summary.json"], + } + + monkeypatch.setattr(cli_main, "repair_hf_native_smoke_proof", fake_repair_hf_native_smoke_proof) + + assert ( + main( + [ + "repair-hf-native-smoke-proof", + "--repo-id", + "WaveCut/example-orbitquant", + "--suite", + "flux2-native", + "--revision", + "main", + "--commit-message", + "repair native smoke proof", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["repo_id"] == "WaveCut/example-orbitquant" + assert output["suite"] == "flux2-native" + assert output["dry_run"] is True + assert seen["repo_id"] == "WaveCut/example-orbitquant" + assert seen["suite"].name == "flux2-native" + assert seen["revision"] == "main" + assert seen["commit_message"] == "repair native smoke proof" + assert seen["dry_run"] is True + + +def test_cli_cleanup_hf_artifact_reports_wires_single_repo_options(capsys, monkeypatch): + seen = {} + + def fake_cleanup_hf_artifact_reports(**kwargs): + seen.update(kwargs) + return { + "repo_id": kwargs["repo_id"], + "revision": kwargs["revision"], + "dry_run": kwargs["dry_run"], + "report_file_count": 2, + "promoted_assets": ["assets/image_generation_comparison_matrix.webp"], + "delete_paths": ["reports"], + } + + monkeypatch.setattr(cli_main, "cleanup_hf_artifact_reports", fake_cleanup_hf_artifact_reports) + + assert ( + main( + [ + "cleanup-hf-artifact-reports", + "--repo-id", + "WaveCut/example-orbitquant", + "--revision", + "main", + "--commit-message", + "cleanup reports", + "--dry-run", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["repo_id"] == "WaveCut/example-orbitquant" + assert output["promoted_assets"] == ["assets/image_generation_comparison_matrix.webp"] + assert output["delete_paths"] == ["reports"] + assert seen == { + "repo_id": "WaveCut/example-orbitquant", + "revision": "main", + "commit_message": "cleanup reports", + "dry_run": True, + } + + +def test_cli_audit_hf_artifacts_writes_json_report(capsys, tmp_path, monkeypatch): + seen = {} + + def fake_audit_hf_artifacts(*, namespace, suites, revision, policy_inventory_root): + seen["namespace"] = namespace + seen["suites"] = [suite.name for suite in suites] + seen["revision"] = revision + seen["policy_inventory_root"] = policy_inventory_root + return { + "namespace": namespace, + "policy_inventory_root": policy_inventory_root, + "repo_count": 1, + "existing_count": 1, + "artifact_ready_count": 1, + "native_smoke_ready_count": 1, + "release_eval_ready_count": 0, + "rows": [{"repo_id": "WaveCut/example"}], + } + + monkeypatch.setattr(cli_main, "audit_hf_artifact_repos", fake_audit_hf_artifacts) + output_path = tmp_path / "reports" / "native" / "audit.json" + markdown_output_path = tmp_path / "reports" / "native" / "audit.md" + + assert ( + main( + [ + "audit-hf-artifacts", + "--namespace", + "WaveCut", + "--suite", + "flux2-native", + "--revision", + "main", + "--policy-inventory-root", + str(tmp_path / "inventories"), + "--output", + str(output_path), + "--markdown-output", + str(markdown_output_path), + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + written = json.loads(output_path.read_text()) + markdown = markdown_output_path.read_text() + assert output_path.parent.is_dir() + assert output == written + assert output["repo_count"] == 1 + assert "# OrbitQuant HF Artifact Audit" in markdown + assert "| `WaveCut/example` |" in markdown + assert seen == { + "namespace": "WaveCut", + "suites": ["flux2-native"], + "revision": "main", + "policy_inventory_root": str(tmp_path / "inventories"), + } + + +def test_cli_audit_hf_artifacts_summary_only_omits_full_rows(capsys, tmp_path, monkeypatch): + def fake_audit_hf_artifacts(*, namespace, suites, revision, policy_inventory_root): + return { + "namespace": namespace, + "policy_inventory_root": policy_inventory_root, + "repo_count": 1, + "existing_count": 1, + "artifact_ready_count": 1, + "native_smoke_ready_count": 1, + "metadata_complete_ready_count": 1, + "policy_inventory_ready_count": 1, + "policy_inventory_error_count": 0, + "release_eval_applicable_count": 1, + "release_eval_ready_count": 0, + "missing_required_metric_count": 2, + "manifest_warning_count": 0, + "metadata_missing_count": 0, + "remote_checksum_mismatch_count": 0, + "readme_mismatch_count": 0, + "forbidden_file_count": 0, + "rows": [ + { + "repo_id": "WaveCut/example", + "suite": "flux1-schnell-native", + "bit_setting": "W4A4", + "artifact_ready": True, + "native_smoke_ready": True, + "release_eval_applicable": True, + "release_eval_ready": False, + "metadata_complete_ready": True, + "policy_inventory_ready": True, + "forbidden_file_count": 0, + "missing_required_metrics": [ + {"split": "original", "metric": "geneval_overall"}, + {"split": "orbitquant", "metric": "geneval_overall"}, + ], + "large_remote_row_detail": {"not": "printed"}, + } + ], + } + + monkeypatch.setattr(cli_main, "audit_hf_artifact_repos", fake_audit_hf_artifacts) + output_path = tmp_path / "audit-summary.json" + + assert ( + main( + [ + "audit-hf-artifacts", + "--summary-only", + "--output", + str(output_path), + "--fail-on-artifact-regression", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + written = json.loads(output_path.read_text()) + assert output == written + assert "rows" not in output + assert output["repo_count"] == 1 + assert output["row_count"] == 1 + assert output["repos"] == [ + { + "repo_id": "WaveCut/example", + "suite": "flux1-schnell-native", + "bit_setting": "W4A4", + "artifact_ready": True, + "native_smoke_ready": True, + "release_eval_applicable": True, + "release_eval_ready": False, + "metadata_complete_ready": True, + "policy_inventory_ready": True, + "forbidden_file_count": 0, + "missing_required_metric_count": 2, + } + ] + assert "large_remote_row_detail" not in json.dumps(output) + + +def test_cli_audit_hf_artifacts_artifact_regression_gate_ignores_release_metrics( + capsys, tmp_path, monkeypatch +): + def fake_audit_hf_artifacts(*, namespace, suites, revision, policy_inventory_root): + return { + "namespace": namespace, + "policy_inventory_root": policy_inventory_root, + "repo_count": 1, + "existing_count": 1, + "artifact_ready_count": 1, + "native_smoke_ready_count": 1, + "metadata_complete_ready_count": 1, + "policy_inventory_ready_count": 1, + "policy_inventory_error_count": 0, + "release_eval_ready_count": 0, + "missing_required_metric_count": 14, + "manifest_warning_count": 0, + "metadata_missing_count": 0, + "remote_checksum_mismatch_count": 0, + "forbidden_file_count": 0, + "rows": [{"repo_id": "WaveCut/example"}], + } + + monkeypatch.setattr(cli_main, "audit_hf_artifact_repos", fake_audit_hf_artifacts) + + assert ( + main( + [ + "audit-hf-artifacts", + "--policy-inventory-root", + str(tmp_path / "inventories"), + "--fail-on-artifact-regression", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["missing_required_metric_count"] == 14 + + +def test_cli_audit_hf_artifacts_artifact_regression_gate_fails_on_remote_hygiene( + capsys, tmp_path, monkeypatch +): + def fake_audit_hf_artifacts(*, namespace, suites, revision, policy_inventory_root): + return { + "namespace": namespace, + "policy_inventory_root": policy_inventory_root, + "repo_count": 1, + "existing_count": 1, + "artifact_ready_count": 1, + "native_smoke_ready_count": 1, + "metadata_complete_ready_count": 1, + "policy_inventory_ready_count": 1, + "policy_inventory_error_count": 0, + "release_eval_ready_count": 0, + "missing_required_metric_count": 0, + "manifest_warning_count": 0, + "metadata_missing_count": 0, + "remote_checksum_mismatch_count": 0, + "forbidden_file_count": 1, + "rows": [{"repo_id": "WaveCut/example"}], + } + + monkeypatch.setattr(cli_main, "audit_hf_artifact_repos", fake_audit_hf_artifacts) + + assert ( + main( + [ + "audit-hf-artifacts", + "--policy-inventory-root", + str(tmp_path / "inventories"), + "--fail-on-artifact-regression", + ] + ) + == 1 + ) + + captured = capsys.readouterr() + assert json.loads(captured.out)["forbidden_file_count"] == 1 + assert "forbidden_file_count=1 expected 0" in captured.err + + +def test_cli_fetch_hf_artifacts_wires_suite_and_download_options(capsys, tmp_path, monkeypatch): + seen = {} + + def fake_fetch_hf_artifacts(**kwargs): + seen.update(kwargs) + if kwargs["stage_logger"] is not None: + kwargs["stage_logger"]("START", "example fetch") + return { + "namespace": kwargs["namespace"], + "output_root": str(kwargs["output_root"]), + "repo_count": 1, + "downloaded_count": 0, + "skipped_existing_count": 0, + "dry_run": kwargs["dry_run"], + "rows": [{"repo_id": "WaveCut/example"}], + } + + monkeypatch.setattr(cli_main, "fetch_hf_artifacts", fake_fetch_hf_artifacts) + + assert ( + main( + [ + "fetch-hf-artifacts", + "--namespace", + "WaveCut", + "--suite", + "flux1-schnell-native", + "--output-root", + str(tmp_path / "artifacts"), + "--revision", + "main", + "--no-resume", + "--force-download", + "--local-files-only", + "--validate-checksums", + "--validate-tensors", + ] + ) + == 0 + ) + + captured = capsys.readouterr() + output = json.loads(captured.out) + assert output["repo_count"] == 1 + assert "example fetch" in captured.err + assert seen["namespace"] == "WaveCut" + assert [suite.name for suite in seen["suites"]] == ["flux1-schnell-native"] + assert seen["output_root"] == str(tmp_path / "artifacts") + assert seen["revision"] == "main" + assert seen["resume"] is False + assert seen["force_download"] is True + assert seen["local_files_only"] is True + assert seen["validate_checksums"] is True + assert seen["validate_tensors"] is True + assert seen["dry_run"] is False + assert seen["stage_logger"] is not None + + +def test_cli_fetch_hf_artifacts_dry_run_suppresses_stage_log(capsys, monkeypatch): + seen = {} + + def fake_fetch_hf_artifacts(**kwargs): + seen.update(kwargs) + return { + "namespace": kwargs["namespace"], + "output_root": str(kwargs["output_root"]), + "repo_count": 0, + "downloaded_count": 0, + "skipped_existing_count": 0, + "dry_run": kwargs["dry_run"], + "rows": [], + } + + monkeypatch.setattr(cli_main, "fetch_hf_artifacts", fake_fetch_hf_artifacts) + + assert main(["fetch-hf-artifacts", "--dry-run"]) == 0 + + captured = capsys.readouterr() + output = json.loads(captured.out) + assert output["dry_run"] is True + assert captured.err == "" + assert seen["dry_run"] is True + assert seen["stage_logger"] is None diff --git a/tests/test_cli_info_and_bench.py b/tests/test_cli_info_and_bench.py new file mode 100644 index 0000000..019049c --- /dev/null +++ b/tests/test_cli_info_and_bench.py @@ -0,0 +1,319 @@ +import json + +import pytest + +import orbitquant.cli.main as cli_main +from orbitquant import __version__ +from orbitquant.cli.main import main +from orbitquant.config import OrbitQuantConfig + + +def test_cli_version_prints_version(capsys): + assert main(["--version"]) == 0 + + output = capsys.readouterr().out + assert __version__ in output + + +def test_cli_packed_matmul_runtime_modes_skip_dequant_prewarm(): + assert cli_main._should_prewarm_quantized_weights(None) is True + assert cli_main._should_prewarm_quantized_weights(OrbitQuantConfig()) is False + assert ( + cli_main._should_prewarm_quantized_weights(OrbitQuantConfig(runtime_mode="dequant_bf16")) + is True + ) + assert ( + cli_main._should_prewarm_quantized_weights( + OrbitQuantConfig(runtime_mode="triton_packed_matmul") + ) + is False + ) + assert ( + cli_main._should_prewarm_quantized_weights( + OrbitQuantConfig(runtime_mode="native_packed_matmul") + ) + is False + ) + + +def test_cli_generation_placement_uses_model_cpu_offload_when_requested(): + class TinyPipeline: + def __init__(self): + self.calls = [] + + def to(self, device): + self.calls.append(("to", device)) + return self + + def enable_model_cpu_offload(self, *, device): + self.calls.append(("offload", device)) + + pipeline = TinyPipeline() + cli_main._place_pipeline_for_generation( + pipeline, + device="cuda", + enable_model_cpu_offload=True, + ) + + assert pipeline.calls == [("offload", "cuda")] + + +def test_cli_generation_placement_fails_loud_when_offload_is_missing(): + class TinyPipeline: + def to(self, device): + return self + + with pytest.raises(RuntimeError, match="enable_model_cpu_offload"): + cli_main._place_pipeline_for_generation( + TinyPipeline(), + device="cuda", + enable_model_cpu_offload=True, + ) + + +def test_cli_native_suites_lists_no_range_smoke_settings(capsys): + assert main(["native-suites"]) == 0 + + output = capsys.readouterr().out + assert "flux2-native" in output + assert "wan-native" in output + assert "range" not in output.lower() + + +def test_cli_kernel_info_reports_backend_capabilities(capsys, monkeypatch): + monkeypatch.setattr( + cli_main, + "backend_capabilities", + lambda: { + "cpu": { + "available": True, + "claim_status": "reference_only", + "optimized": False, + "implemented_stage": None, + "optimized_stage": None, + "weight_dequant_optimized": False, + "weight_pack_optimized": False, + "weight_quant_optimized": False, + "adaln_quant_optimized": False, + "adaln_dequant_optimized": False, + "hf_kernel_builder_compliant": False, + }, + "mps": { + "claim_status": "partial_optimized", + "implementation": "torch_mps_compile_shader_codebook_rescale+native_packed_matmul", + "package_format": "torch.mps.compile_shader,native_kernel_package", + "implemented_stage": ( + "codebook_lookup_rescale,packed_weight_dequant,packed_weight_matmul" + ), + "optimized_stage": ( + "codebook_lookup_rescale,packed_weight_dequant,packed_weight_matmul" + ), + "weight_dequant_optimized": True, + "weight_pack_optimized": False, + "weight_quant_optimized": False, + "adaln_quant_optimized": False, + "adaln_dequant_optimized": False, + "full_fusion": False, + "upstream_native_mps_op": False, + "hf_kernel_builder_compliant": False, + }, + "triton_cuda": { + "claim_status": "partial_optimized", + "implemented_stage": ( + "activation_norm_rpbh_quant_rescale,packed_weight_dequant," + "packed_weight_matmul,lowbit_pack,lowbit_unpack," + "weight_rotation_fwht_quant_pack," + "adaln_rtn_quant_pack,adaln_rtn_dequant," + "adaln_rtn_packed_matmul" + ), + "optimized_stage": ( + "activation_norm_rpbh_quant_rescale,packed_weight_dequant," + "packed_weight_matmul,lowbit_pack,lowbit_unpack," + "weight_rotation_fwht_quant_pack," + "adaln_rtn_quant_pack,adaln_rtn_dequant," + "adaln_rtn_packed_matmul" + ), + "implementation": "python_triton_orbitquant_pipeline", + "package_format": "python_triton", + "weight_dequant_optimized": True, + "weight_pack_optimized": True, + "weight_quant_optimized": True, + "adaln_quant_optimized": True, + "adaln_dequant_optimized": True, + "full_fusion": False, + "hf_kernel_builder_compliant": False, + }, + }, + ) + + assert main(["kernel-info"]) == 0 + + payload = json.loads(capsys.readouterr().out) + assert payload["cpu"]["available"] is True + assert payload["cpu"]["claim_status"] == "reference_only" + assert payload["cpu"]["optimized"] is False + assert payload["cpu"]["weight_dequant_optimized"] is False + assert payload["cpu"]["weight_pack_optimized"] is False + assert payload["cpu"]["weight_quant_optimized"] is False + assert payload["cpu"]["adaln_quant_optimized"] is False + assert payload["cpu"]["adaln_dequant_optimized"] is False + assert ( + payload["mps"]["implementation"] + == "torch_mps_compile_shader_codebook_rescale+native_packed_matmul" + ) + assert payload["mps"]["package_format"] == "torch.mps.compile_shader,native_kernel_package" + assert payload["mps"]["claim_status"] == "partial_optimized" + assert ( + payload["mps"]["implemented_stage"] + == "codebook_lookup_rescale,packed_weight_dequant,packed_weight_matmul" + ) + assert ( + payload["mps"]["optimized_stage"] + == "codebook_lookup_rescale,packed_weight_dequant,packed_weight_matmul" + ) + assert payload["mps"]["weight_dequant_optimized"] is True + assert payload["mps"]["weight_pack_optimized"] is False + assert payload["mps"]["weight_quant_optimized"] is False + assert payload["mps"]["adaln_quant_optimized"] is False + assert payload["mps"]["adaln_dequant_optimized"] is False + assert payload["mps"]["full_fusion"] is False + assert payload["mps"]["upstream_native_mps_op"] is False + assert payload["mps"]["hf_kernel_builder_compliant"] is False + expected_triton_stage = ( + "activation_norm_rpbh_quant_rescale,packed_weight_dequant," + "packed_weight_matmul,lowbit_pack,lowbit_unpack," + "weight_rotation_fwht_quant_pack," + "adaln_rtn_quant_pack,adaln_rtn_dequant,adaln_rtn_packed_matmul" + ) + assert payload["triton_cuda"]["implemented_stage"] == expected_triton_stage + assert payload["triton_cuda"]["optimized_stage"] == expected_triton_stage + assert payload["triton_cuda"]["implementation"] == "python_triton_orbitquant_pipeline" + assert payload["triton_cuda"]["package_format"] == "python_triton" + assert payload["triton_cuda"]["claim_status"] == "partial_optimized" + assert payload["triton_cuda"]["weight_dequant_optimized"] is True + assert payload["triton_cuda"]["weight_pack_optimized"] is True + assert payload["triton_cuda"]["weight_quant_optimized"] is True + assert payload["triton_cuda"]["adaln_quant_optimized"] is True + assert payload["triton_cuda"]["adaln_dequant_optimized"] is True + assert payload["triton_cuda"]["full_fusion"] is False + assert payload["triton_cuda"]["hf_kernel_builder_compliant"] is False + + +def test_cli_kernel_bench_prints_stage_timings(capsys): + assert ( + main( + [ + "kernel-bench", + "--tokens", + "4", + "--in-features", + "16", + "--out-features", + "8", + "--block-size", + "8", + "--activation-kernel-backend", + "cpu", + "--runtime-mode", + "debug_no_activation_quant", + "--device", + "cpu", + "--dtype", + "float32", + "--warmup", + "0", + "--iterations", + "1", + ] + ) + == 0 + ) + + payload = json.loads(capsys.readouterr().out) + assert payload["device"] == "cpu" + assert payload["runtime_mode"] == "debug_no_activation_quant" + assert payload["full_fusion"] is False + assert payload["prewarm"]["total_modules"] == 1 + assert payload["timings_ms"]["weight_quantize_pack_cold_ms"] >= 0.0 + assert payload["timings_ms"]["weight_quantize_pack_hot_ms"] >= 0.0 + assert payload["timings_ms"]["forward_prewarmed_ms"] >= 0.0 + assert payload["selected_activation_kernel_backend"] == "cpu" + assert payload["weight_quantization_backend"] == "torch_reference" + assert payload["quantization_buffers"]["source_weight_device"] == "cpu" + assert payload["quantization_buffers"]["source_weight_is_cuda"] is False + assert payload["quantization_buffers"]["packed_weight_indices_device"] == "cpu" + + +@pytest.mark.parametrize("runtime_mode", ["triton_packed_matmul", "native_packed_matmul"]) +def test_cli_kernel_bench_passes_packed_matmul_tile_options(monkeypatch, capsys, runtime_mode): + seen_kwargs = [] + + def fake_benchmark_orbit_linear(**kwargs): + seen_kwargs.append(kwargs) + return {"ok": True} + + monkeypatch.setattr(cli_main, "benchmark_orbit_linear", fake_benchmark_orbit_linear) + + assert ( + main( + [ + "kernel-bench", + "--runtime-mode", + runtime_mode, + "--packed-matmul-block-m", + "32", + "--packed-matmul-block-n", + "64", + "--packed-matmul-block-k", + "64", + "--packed-matmul-num-warps", + "8", + ] + ) + == 0 + ) + + assert json.loads(capsys.readouterr().out) == {"ok": True} + assert seen_kwargs[0]["runtime_mode"] == runtime_mode + assert seen_kwargs[0]["packed_matmul_block_m"] == 32 + assert seen_kwargs[0]["packed_matmul_block_n"] == 64 + assert seen_kwargs[0]["packed_matmul_block_k"] == 64 + assert seen_kwargs[0]["packed_matmul_num_warps"] == 8 + + +def test_cli_quantize_bench_prints_full_model_staging_timings(capsys): + assert ( + main( + [ + "quantize-bench", + "--layers", + "1", + "--in-features", + "16", + "--hidden-features", + "32", + "--block-size", + "8", + "--source-device", + "cpu", + "--quantization-device", + "cpu", + "--staging-mode", + "component", + "--dtype", + "float32", + ] + ) + == 0 + ) + + payload = json.loads(capsys.readouterr().out) + assert payload["source_device"] == "cpu" + assert payload["quantization_device"] == "cpu" + assert payload["staging_mode"] == "component" + assert payload["synchronize_per_module"] is False + assert payload["summary"]["quantization_staging_mode"] == "component" + assert payload["summary"]["synchronize_per_module"] is False + assert payload["summary"]["source_linear_device_counts"]["cpu"] == 7 + assert payload["summary"]["device_transfer_seconds"] >= 0.0 + assert payload["summary"]["quantized_modules"] diff --git a/tests/test_cli_plans_and_scripts.py b/tests/test_cli_plans_and_scripts.py new file mode 100644 index 0000000..f709c73 --- /dev/null +++ b/tests/test_cli_plans_and_scripts.py @@ -0,0 +1,337 @@ +import json +from types import SimpleNamespace + +import pytest + +import orbitquant.cli.main as cli_main +from orbitquant.cli.main import main + + +@pytest.mark.parametrize("runtime_mode", ["triton_packed_matmul", "native_packed_matmul"]) +def test_cli_native_plan_lists_full_target_bit_matrix_without_range_smoke( + capsys, tmp_path, runtime_mode +): + assert ( + main( + [ + "native-plan", + "--output-root", + str(tmp_path / "artifacts"), + "--seeds", + "0", + "--runtime-mode", + runtime_mode, + ] + ) + == 0 + ) + + payload = json.loads(capsys.readouterr().out) + jobs = payload["jobs"] + assert payload["job_count"] == 14 + assert {job["suite"] for job in jobs} == { + "flux2-native", + "flux1-schnell-native", + "z-image-native", + "wan-native", + } + assert sum(1 for job in jobs if job["suite"] == "wan-native") == 2 + assert {job["runtime_mode"] for job in jobs} == {runtime_mode} + wan_w4a6 = next( + job for job in jobs if job["suite"] == "wan-native" and job["bit_setting"] == "W4A6" + ) + flux1_w3a3 = next( + job + for job in jobs + if job["suite"] == "flux1-schnell-native" and job["bit_setting"] == "W3A3" + ) + assert wan_w4a6["width"] == 832 + assert wan_w4a6["height"] == 480 + assert wan_w4a6["frames"] == 81 + assert wan_w4a6["metric"] == "vbench" + assert flux1_w3a3["width"] == 1024 + assert flux1_w3a3["height"] == 1024 + assert flux1_w3a3["guidance"] == 0.0 + assert flux1_w3a3["artifact_dir"].endswith("flux1-schnell-native-w3a3") + assert "range" not in json.dumps(payload).lower() + + +def test_cli_native_plan_defaults_to_auto_fused_runtime(capsys, tmp_path): + assert ( + main( + [ + "native-plan", + "--output-root", + str(tmp_path / "artifacts"), + "--seeds", + "0", + ] + ) + == 0 + ) + + payload = json.loads(capsys.readouterr().out) + + assert payload["job_count"] == 14 + assert {job["runtime_mode"] for job in payload["jobs"]} == {"auto_fused"} + + +def test_cli_external_eval_plan_lists_metric_runner_import_commands(capsys, tmp_path): + assert ( + main( + [ + "external-eval-plan", + "--suite", + "wan-native", + "--output-root", + str(tmp_path / "artifacts"), + "--metrics-root", + str(tmp_path / "metrics"), + ] + ) + == 0 + ) + + payload = json.loads(capsys.readouterr().out) + assert payload["job_count"] == 4 + job = next( + item + for item in payload["jobs"] + if item["bit_setting"] == "W4A6" and item["split"] == "orbitquant" + ) + assert job["suite"] == "wan-native" + assert job["metric"] == "vbench" + assert job["artifact_dir"].endswith("wan-native-w4a6") + assert job["metrics_json"].endswith("wan-native-w4a6_orbitquant_vbench.json") + assert "orbitquant export-vbench" in job["export_command"] + assert "vbench evaluate" in job["eval_command"] + assert "orbitquant summarize-vbench-results" in job["summarize_command"] + assert "orbitquant record-metrics" in job["import_command"] + assert "--metric-prefix vbench" in job["import_command"] + assert "--split orbitquant" in job["import_command"] + assert "range" not in json.dumps(payload).lower() + + +def test_cli_external_eval_script_prints_metric_runner_script(capsys, tmp_path): + assert ( + main( + [ + "external-eval-script", + "--suite", + "wan-native", + "--output-root", + str(tmp_path / "artifacts"), + "--metrics-root", + str(tmp_path / "metrics"), + "--report-output", + str(tmp_path / "reports"), + ] + ) + == 0 + ) + + script = capsys.readouterr().out + assert script.startswith("#!/usr/bin/env bash\nset -euo pipefail\n") + assert "if ! command -v vbench >/dev/null 2>&1; then" in script + assert f"if [ ! -d {tmp_path / 'artifacts' / 'wan-native-w4a6'} ]; then" in script + assert "orbitquant export-vbench" in script + assert "vbench evaluate" in script + assert "orbitquant summarize-vbench-results" in script + assert "orbitquant record-metrics" in script + assert "--metric-prefix vbench" in script + assert "orbitquant report" in script + assert f"--output {tmp_path / 'reports'}" in script + assert "--fail-on-missing-required" in script + assert "range" not in script.lower() + + +def test_cli_external_eval_script_resume_skips_existing_metric_summaries(capsys, tmp_path): + assert ( + main( + [ + "external-eval-script", + "--suite", + "wan-native", + "--output-root", + str(tmp_path / "artifacts"), + "--metrics-root", + str(tmp_path / "metrics"), + "--report-output", + str(tmp_path / "reports"), + "--resume", + ] + ) + == 0 + ) + + script = capsys.readouterr().out + metrics_json = tmp_path / "metrics" / "wan-native-w4a4_original_vbench.json" + assert f"if [ -s {metrics_json} ]; then" in script + assert "stage_log SKIP 'wan-native W4A4 original vbench metrics'" in script + assert "fi\nstage_log START 'wan-native W4A4 original import vbench'" in script + assert script.count("stage_log SKIP") == 4 + assert script.count("orbitquant record-metrics") == 4 + assert "--fail-on-missing-required" in script + + +def test_cli_report_can_fail_on_missing_required_metrics(capsys, monkeypatch, tmp_path): + result = SimpleNamespace( + report_path=tmp_path / "report.md", + table_paths={}, + rows=[], + missing_required_metrics=[{"metric": "geneval_overall"}], + ) + monkeypatch.setattr(cli_main, "generate_native_eval_report", lambda *args, **kwargs: result) + + assert ( + main(["report", "--artifact", str(tmp_path / "artifact"), "--output", str(tmp_path)]) == 0 + ) + assert json.loads(capsys.readouterr().out)["missing_required_metric_count"] == 1 + + assert ( + main( + [ + "report", + "--artifact", + str(tmp_path / "artifact"), + "--output", + str(tmp_path), + "--fail-on-missing-required", + ] + ) + == 1 + ) + assert json.loads(capsys.readouterr().out)["missing_required_metric_count"] == 1 + + +def test_cli_native_script_groups_quantize_and_generate_pack_commands(capsys, tmp_path): + assert ( + main( + [ + "native-script", + "--suite", + "wan-native", + "--output-root", + str(tmp_path / "artifacts"), + "--seeds", + "0,1", + "--prompt-limit", + "1", + "--device", + "cuda", + "--dtype", + "bfloat16", + "--activation-kernel-backend", + "triton_cuda", + "--runtime-mode", + "triton_packed_matmul", + ] + ) + == 0 + ) + + script = capsys.readouterr().out + assert script.startswith("#!/usr/bin/env bash\nset -euo pipefail\n") + assert "stage_log() {" in script + assert "stage_log START preflight" in script + assert "stage_log START 'kernel preflight'" in script + assert "stage_log START 'policy inventories'" in script + assert "stage_log START 'wan-native W4A6 quantize'" in script + assert "stage_log START 'wan-native W4A6 original generate-pack'" in script + assert "stage_log END 'native eval report'" in script + assert "hf auth whoami" in script + assert "hf env" in script + assert "torch.cuda.is_available()" in script + assert "shutil.disk_usage" in script + assert "hf models info Wan-AI/Wan2.1-T2V-1.3B-Diffusers --format json >/dev/null" in script + assert "orbitquant kernel-info" in script + assert script.count("orbitquant kernel-bench") == 2 + assert script.count("orbitquant inspect-policy") == 1 + assert "--output reports/native/module-inventories/wan-native-policy.json" in script + assert "--tokens 256 --in-features 3072 --out-features 3072" in script + assert script.count("orbitquant quantize") == 2 + assert "--suite wan-native" in script + assert "--weight-bits 4 --activation-bits 6" in script + assert "--weight-bits 4 --activation-bits 4" in script + assert "--activation-kernel-backend triton_cuda" in script + assert "--runtime-mode triton_packed_matmul" in script + assert ( + "--policy-inventory reports/native/module-inventories/wan-native-policy.json " + "--runtime-mode triton_packed_matmul" + ) in script + assert "--staging-mode component" in script + assert script.count("\norbitquant generate-pack") == 4 + assert "--split original" in script + assert "--split orbitquant" in script + assert "--seeds 0,1" in script + assert "--prompt-limit 1" in script + assert script.count("orbitquant validate-artifact") == 4 + assert "--policy-inventory reports/native/module-inventories/wan-native-policy.json" in script + assert script.count("orbitquant report") == 1 + assert f"--artifact {tmp_path / 'artifacts' / 'wan-native-w4a6'}" in script + assert f"--artifact {tmp_path / 'artifacts' / 'wan-native-w4a4'}" in script + assert "--output reports/native" in script + assert "range" not in script.lower() + + +def test_cli_native_script_can_generate_geneval_metadata_pack(capsys, tmp_path): + metadata_jsonl = tmp_path / "evaluation_metadata.jsonl" + assert ( + main( + [ + "native-script", + "--suite", + "flux1-schnell-native", + "--output-root", + str(tmp_path / "artifacts"), + "--prompt-metadata-jsonl", + str(metadata_jsonl), + "--seeds", + "0", + "--prompt-limit", + "2", + ] + ) + == 0 + ) + + script = capsys.readouterr().out + assert "--prompt-metadata-jsonl" in script + assert str(metadata_jsonl) in script + assert "--prompt-limit 2" in script + assert "--prompt-pack" not in script + assert "stage_log START 'flux1-schnell-native W4A4 original generate-pack'" in script + assert "stage_log START 'flux1-schnell-native W4A4 orbitquant generate-pack'" in script + assert script.count("\norbitquant generate-pack") == 8 + + +def test_cli_native_script_resume_skips_valid_existing_artifacts(capsys, tmp_path): + assert ( + main( + [ + "native-script", + "--suite", + "flux2-native", + "--output-root", + str(tmp_path / "artifacts"), + "--seeds", + "0", + "--resume", + ] + ) + == 0 + ) + + script = capsys.readouterr().out + artifact_dir = tmp_path / "artifacts" / "flux2-native-w4a4" + inventory_path = "reports/native/module-inventories/flux2-native-policy.json" + assert ( + f"if orbitquant validate-artifact --artifact {artifact_dir} " + f"--policy-inventory {inventory_path} --runtime-mode auto_fused" + ) in script + assert f"echo 'Skipping existing valid artifact: {artifact_dir}'" in script + assert "else\norbitquant quantize --suite flux2-native" in script + assert "\nfi\nstage_log END 'flux2-native W4A4 quantize'" in script + assert "stage_log START 'flux2-native W4A4 validate quantized artifact'" in script + assert "\norbitquant generate-pack --suite flux2-native" in script + assert "--resume-existing" in script diff --git a/tests/test_cli_quantize_and_policy.py b/tests/test_cli_quantize_and_policy.py new file mode 100644 index 0000000..c4a075b --- /dev/null +++ b/tests/test_cli_quantize_and_policy.py @@ -0,0 +1,398 @@ +import json +import sys +from types import SimpleNamespace + +import pytest +import torch + +import orbitquant.cli.main as cli_main +from orbitquant.cli.main import main + + +def test_cli_quantize_saves_transformer_component_artifact(monkeypatch, capsys, tmp_path): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + self.device = None + + def to(self, device): + self.device = device + return self + + pipeline = TinyPipeline() + + class FakeDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "example/model" + assert kwargs["revision"] == "main" + assert kwargs["torch_dtype"] is torch.float32 + return pipeline + + monkeypatch.setitem( + sys.modules, "diffusers", SimpleNamespace(DiffusionPipeline=FakeDiffusionPipeline) + ) + monkeypatch.setattr( + cli_main, + "inspect_model_metadata", + lambda model_id, revision=None: { + "sha": "abc123", + "license": "apache-2.0", + }, + ) + + assert ( + main( + [ + "quantize", + "--model-id", + "example/model", + "--revision", + "main", + "--output", + str(tmp_path), + "--component", + "transformer", + "--target-policy", + "generic_dit", + "--weight-bits", + "4", + "--activation-bits", + "4", + "--block-size", + "4", + "--activation-kernel-backend", + "cpu", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + assert output["artifact_dir"] == str(tmp_path) + assert output["quantization_staging_mode"] == "streaming" + assert output["synchronize_per_module"] is False + assert output["device_transfer_seconds"] >= 0.0 + assert output["module_device_transfer_count"] >= 0 + assert output["source_linear_device_counts"] + assert output["artifact_save_elapsed_seconds"] >= 0.0 + assert output["quantized_modules"] == ["transformer_blocks.0.attn.to_q"] + assert (tmp_path / "model.safetensors").exists() + assert (tmp_path / "orbitquant_manifest.json").exists() + assert (tmp_path / "SHA256SUMS").exists() + quantization_config = json.loads((tmp_path / "quantization_config.json").read_text()) + assert quantization_config["activation_kernel_backend"] == "cpu" + + +def test_cli_quantize_with_suite_uses_named_native_pipeline(monkeypatch, capsys, tmp_path): + class TinyPipeline: + def __init__(self): + self.transformer = torch.nn.Module() + self.transformer.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def to(self, device): + return self + + class FakeFlux2KleinPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + assert model_id == "black-forest-labs/FLUX.2-klein-4B" + assert kwargs["torch_dtype"] is torch.float32 + return TinyPipeline() + + class WrongDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + raise AssertionError("suite-specific pipeline should be used") + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace( + Flux2KleinPipeline=FakeFlux2KleinPipeline, + DiffusionPipeline=WrongDiffusionPipeline, + ), + ) + monkeypatch.setattr( + cli_main, + "inspect_model_metadata", + lambda model_id, revision=None: { + "sha": "abc123", + "license": "apache-2.0", + }, + ) + + assert ( + main( + [ + "quantize", + "--suite", + "flux2-native", + "--output", + str(tmp_path), + "--component", + "transformer", + "--weight-bits", + "4", + "--activation-bits", + "4", + "--block-size", + "4", + "--device", + "cpu", + "--dtype", + "float32", + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + quantization_config = json.loads((tmp_path / "quantization_config.json").read_text()) + assert output["source_model_id"] == "black-forest-labs/FLUX.2-klein-4B" + assert quantization_config["target_policy"] == "flux2" + + +def test_cli_inspect_policy_with_suite_writes_module_inventory(monkeypatch, capsys, tmp_path): + calls = [] + + class FakeFlux2Transformer2DModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.transformer = torch.nn.Module() + self.double_stream_modulation_img = torch.nn.ModuleDict( + {"linear": torch.nn.Linear(8, 16)} + ) + self.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + @classmethod + def load_config(cls, model_id, **kwargs): + assert model_id == "black-forest-labs/FLUX.2-klein-4B" + assert kwargs["subfolder"] == "transformer" + assert kwargs["local_files_only"] is False + calls.append({"method": "load_config", "kwargs": kwargs}) + return {"hidden_size": 8} + + @classmethod + def from_config(cls, config): + assert config == {"hidden_size": 8} + calls.append({"method": "from_config"}) + return cls() + + class WrongDiffusionPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + raise AssertionError("inspect-policy default must not load full pipeline weights") + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace( + Flux2Transformer2DModel=FakeFlux2Transformer2DModel, + DiffusionPipeline=WrongDiffusionPipeline, + ), + ) + output_path = tmp_path / "flux2-policy-inventory.json" + + assert ( + main( + [ + "inspect-policy", + "--suite", + "flux2-native", + "--dtype", + "float32", + "--output", + str(output_path), + ] + ) + == 0 + ) + + output = json.loads(capsys.readouterr().out) + saved = json.loads(output_path.read_text()) + + assert output["output"] == str(output_path) + assert output["linear_module_count"] == 2 + assert output["quantized_module_count"] == 1 + assert output["adaln_module_count"] == 1 + assert output["skipped_module_count"] == 0 + assert "modules" not in output + assert saved["source_model_id"] == "black-forest-labs/FLUX.2-klein-4B" + assert saved["suite"] == "flux2-native" + assert saved["component"] == "transformer" + assert saved["load_mode"] == "config" + assert saved["pipeline_class"] is None + assert saved["component_class"] == "FakeFlux2Transformer2DModel" + assert saved["target_policy"] == "flux2" + assert saved["action_counts"] == { + "orbitquant": 1, + "adaln_int4_rtn": 1, + "bf16_skip": 0, + } + assert saved["quantized_modules"] == ["transformer_blocks.0.attn.to_q"] + assert saved["adaln_modules"] == ["double_stream_modulation_img.linear"] + assert calls == [ + { + "method": "load_config", + "kwargs": {"subfolder": "transformer", "local_files_only": False}, + }, + {"method": "from_config"}, + ] + + +@pytest.mark.parametrize( + ("suite", "transformer_class", "model_id", "target_policy", "module_family"), + [ + ( + "flux2-native", + "Flux2Transformer2DModel", + "black-forest-labs/FLUX.2-klein-4B", + "flux2", + "flux", + ), + ( + "flux1-schnell-native", + "FluxTransformer2DModel", + "black-forest-labs/FLUX.1-schnell", + "flux", + "flux", + ), + ( + "z-image-native", + "ZImageTransformer2DModel", + "Tongyi-MAI/Z-Image-Turbo", + "z_image", + "z_image", + ), + ( + "wan-native", + "WanTransformer3DModel", + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "wan", + "wan", + ), + ], +) +def test_cli_inspect_policy_config_mode_covers_all_native_target_suites( + monkeypatch, + capsys, + tmp_path, + suite, + transformer_class, + model_id, + target_policy, + module_family, +): + calls = [] + + def build_module(self): + if module_family == "wan": + self.blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn1": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + elif module_family == "z_image": + self.layers = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attention": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + else: + self.transformer_blocks = torch.nn.ModuleList( + [ + torch.nn.ModuleDict( + {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})} + ) + ] + ) + + def load_config(cls, requested_model_id, **kwargs): + assert requested_model_id == model_id + assert kwargs["subfolder"] == "transformer" + calls.append(("load_config", transformer_class)) + return {"hidden_size": 8} + + def from_config(cls, config): + assert config == {"hidden_size": 8} + calls.append(("from_config", transformer_class)) + return cls() + + def init(self): + torch.nn.Module.__init__(self) + build_module(self) + + fake_transformer_cls = type( + transformer_class, + (torch.nn.Module,), + { + "__init__": init, + "load_config": classmethod(load_config), + "from_config": classmethod(from_config), + }, + ) + + class WrongDiffusionPipeline: + @classmethod + def from_pretrained(cls, requested_model_id, **kwargs): + raise AssertionError("config inspect-policy must not load full pipeline weights") + + monkeypatch.setitem( + sys.modules, + "diffusers", + SimpleNamespace( + **{ + transformer_class: fake_transformer_cls, + "DiffusionPipeline": WrongDiffusionPipeline, + } + ), + ) + output_path = tmp_path / f"{suite}-policy-inventory.json" + + assert main(["inspect-policy", "--suite", suite, "--output", str(output_path)]) == 0 + + output = json.loads(capsys.readouterr().out) + saved = json.loads(output_path.read_text()) + + assert output["output"] == str(output_path) + assert saved["suite"] == suite + assert saved["source_model_id"] == model_id + assert saved["component"] == "transformer" + assert saved["load_mode"] == "config" + assert saved["pipeline_class"] is None + assert saved["component_class"] == transformer_class + assert saved["target_policy"] == target_policy + assert saved["action_counts"]["orbitquant"] == 1 + assert saved["quantized_modules"] + assert calls == [ + ("load_config", transformer_class), + ("from_config", transformer_class), + ] diff --git a/tests/test_config.py b/tests/test_config.py index 0acae0d..34d401f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -93,6 +93,18 @@ def test_orbit_quant_config_accepts_triton_packed_matmul_runtime_mode(): assert config.runtime_mode == "triton_packed_matmul" +def test_orbit_quant_config_accepts_explicit_triton_rocm_activation_backend(): + config = OrbitQuantConfig(activation_kernel_backend="triton_rocm") + + assert config.activation_kernel_backend == "triton_rocm" + + +def test_orbit_quant_config_accepts_explicit_triton_xpu_activation_backend(): + config = OrbitQuantConfig(activation_kernel_backend="triton_xpu") + + assert config.activation_kernel_backend == "triton_xpu" + + def test_orbit_quant_config_accepts_native_packed_matmul_runtime_mode(): config = OrbitQuantConfig(runtime_mode="native_packed_matmul") diff --git a/tests/test_hub.py b/tests/test_hub.py deleted file mode 100644 index c5341bd..0000000 --- a/tests/test_hub.py +++ /dev/null @@ -1,2888 +0,0 @@ -import json -import shutil -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch -from huggingface_hub import CommitOperationAdd, CommitOperationDelete - -import orbitquant.hub as hub_module -from orbitquant.artifacts import ( - record_artifact_metrics, - refresh_artifact_checksums, - save_orbitquant_artifact, -) -from orbitquant.artifacts.checksums import read_sha256sums, write_sha256sums -from orbitquant.config import OrbitQuantConfig -from orbitquant.eval.native_settings import NativeSuite -from orbitquant.hub import ( - audit_hf_artifact_repos, - cleanup_hf_artifact_reports, - cleanup_hf_artifact_reports_matrix, - fetch_hf_artifacts, - render_hf_artifact_audit_markdown, - repair_hf_artifact_metadata, - repair_hf_artifact_metadata_matrix, - repair_hf_native_smoke_proof, - repair_hf_native_smoke_proof_matrix, - stage_compact_upload_artifact, - upload_orbitquant_artifact, -) -from orbitquant.modeling import quantize_linear_modules - - -class TinyHubArtifactModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.transformer_blocks = torch.nn.ModuleList( - [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(8, 8)})})] - ) - - -class FakeHfApi: - def __init__(self): - self.create_repo_calls = [] - self.upload_folder_calls = [] - self.model_info_calls = [] - - def create_repo(self, **kwargs): - self.create_repo_calls.append(kwargs) - return f"https://huggingface.co/{kwargs['repo_id']}" - - def upload_folder(self, **kwargs): - self.upload_folder_calls.append(kwargs) - return SimpleNamespace( - oid="uploaded-sha", - commit_url=f"https://huggingface.co/{kwargs['repo_id']}/commit/uploaded-sha", - pr_url=None, - ) - - def model_info(self, repo_id, *, revision=None): - self.model_info_calls.append({"repo_id": repo_id, "revision": revision}) - return SimpleNamespace(sha=revision or "main-sha", private=True, gated=False) - - -class FakeAuditHfApi: - def __init__(self, siblings_by_repo): - self.siblings_by_repo = siblings_by_repo - self.model_info_calls = [] - - def model_info(self, repo_id, *, revision=None, files_metadata=False): - self.model_info_calls.append( - {"repo_id": repo_id, "revision": revision, "files_metadata": files_metadata} - ) - if repo_id not in self.siblings_by_repo: - raise RuntimeError("missing repo") - siblings = [] - for name, metadata in self.siblings_by_repo[repo_id].items(): - if isinstance(metadata, dict): - lfs_sha256 = metadata.get("lfs_sha256") - lfs = ( - SimpleNamespace(sha256=lfs_sha256, size=metadata.get("size")) - if lfs_sha256 is not None - else None - ) - siblings.append(SimpleNamespace(rfilename=name, size=metadata.get("size"), lfs=lfs)) - else: - siblings.append(SimpleNamespace(rfilename=name, size=metadata, lfs=None)) - return SimpleNamespace( - sha="remote-sha", - private=True, - gated=False, - siblings=siblings, - ) - - -class FakeCommitHfApi: - def __init__(self): - self.create_commit_calls = [] - - def create_commit(self, **kwargs): - self.create_commit_calls.append(kwargs) - return SimpleNamespace( - oid="repair-sha", - commit_url=f"https://huggingface.co/{kwargs['repo_id']}/commit/repair-sha", - pr_url=None, - ) - - -class FakeCleanupHfApi(FakeCommitHfApi): - def __init__(self, artifact_dir): - super().__init__() - self.artifact_dir = artifact_dir - self.model_info_calls = [] - - def model_info(self, repo_id, *, revision=None, files_metadata=False): - self.model_info_calls.append( - {"repo_id": repo_id, "revision": revision, "files_metadata": files_metadata} - ) - siblings = [ - SimpleNamespace(rfilename=path.relative_to(self.artifact_dir).as_posix(), size=1) - for path in sorted(self.artifact_dir.rglob("*")) - if path.is_file() - ] - return SimpleNamespace(sha="remote-sha", private=True, gated=False, siblings=siblings) - - -def _write_artifact(tmp_path): - model = TinyHubArtifactModel() - config = OrbitQuantConfig(block_size=4, target_policy="generic_dit") - summary = quantize_linear_modules(model, config, quantization_device=None) - return save_orbitquant_artifact( - model, - tmp_path, - config=config, - source_model_id="example/model", - source_revision="abc123", - source_license="apache-2.0", - summary=summary, - ) - - -def _required_remote_files(): - return {name: 1 for name in hub_module._REQUIRED_ARTIFACT_FILES} - - -def _remote_file_map(repo_id, artifact_dir): - return { - (repo_id, "orbitquant_manifest.json"): artifact_dir / "orbitquant_manifest.json", - (repo_id, "model_index.json"): artifact_dir / "model_index.json", - (repo_id, "benchmark/summary.json"): artifact_dir / "benchmark" / "summary.json", - (repo_id, "quantization_config.json"): artifact_dir / "quantization_config.json", - (repo_id, "README.md"): artifact_dir / "README.md", - (repo_id, "SHA256SUMS"): artifact_dir / "SHA256SUMS", - } - - -def _remote_path(artifact_dir, filename): - return artifact_dir / filename - - -def _write_remote_model_index( - path, - *, - activation_eps=None, - quantization_device="unknown", - weight_quantization_backend="unknown", - quantization_staging_mode="unknown", -): - payload = { - "_class_name": "OrbitQuantArtifact", - "quantization_device": quantization_device, - "weight_quantization_backend": weight_quantization_backend, - "quantization_staging_mode": quantization_staging_mode, - } - if activation_eps is not None: - payload["activation_eps"] = activation_eps - path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") - - -def _audit_file_map( - repo_id, - *, - manifest_path, - sha256sums_path, - summary_path, - model_index_path, - readme_path=None, - quantization_config_path=None, -): - if readme_path is None: - readme_path = Path(manifest_path).parent / "README.md" - manifest_payload = json.loads(Path(manifest_path).read_text(encoding="utf-8")) - manifest_payload.setdefault("source_revision", "unknown") - manifest_payload.setdefault("source_license", "unknown") - manifest = hub_module.OrbitQuantManifest.from_dict(manifest_payload) - benchmark_summary = json.loads(Path(summary_path).read_text(encoding="utf-8")) - Path(readme_path).write_text( - hub_module.render_model_card(manifest, benchmark_summary=benchmark_summary), - encoding="utf-8", - ) - file_map = { - (repo_id, "orbitquant_manifest.json"): manifest_path, - (repo_id, "model_index.json"): model_index_path, - (repo_id, "SHA256SUMS"): sha256sums_path, - (repo_id, "benchmark/summary.json"): summary_path, - (repo_id, "README.md"): readme_path, - } - if quantization_config_path is not None: - file_map[(repo_id, "quantization_config.json")] = quantization_config_path - return file_map - - -def _native_smoke_summary( - suite, - *, - release_metrics=None, - activation_eps=None, - quantization_device="unknown", - weight_quantization_backend="unknown", - quantization_staging_mode="unknown", -): - native_settings = hub_module._native_smoke_expected_settings(suite) - generated_frames = 0 if suite.frames is None else suite.frames - metrics = {"generated_samples": 1} - if generated_frames: - metrics["generated_frames"] = generated_frames - for key, value in (release_metrics or {}).items(): - metrics[key] = value - split_proof = { - "records": 1, - "generated_samples": 1, - "generated_frames": generated_frames, - "nonempty_output_count": 1, - "seeds": ["0"], - "prompt_ids": ["simple-object"], - "pair_keys": [[suite.name, "0", "simple-object"]], - "native_settings": [native_settings], - } - payload = { - "published_summary": "compact", - "raw_generation_records": "local-only", - "quantization_device": quantization_device, - "weight_quantization_backend": weight_quantization_backend, - "quantization_staging_mode": quantization_staging_mode, - "metrics": { - "original": {"records": 1, "latest_metrics": metrics}, - "orbitquant": {"records": 1, "latest_metrics": metrics}, - }, - "native_smoke": { - "proof_format": "orbitquant-native-smoke-v1", - "comparison_asset_path": "assets/image_generation_comparison_matrix.webp", - "paired_prompt_seed_count": 1, - "paired_prompt_seed_keys": [[suite.name, "0", "simple-object"]], - "splits": { - "original": split_proof, - "orbitquant": split_proof, - }, - }, - } - if activation_eps is not None: - payload["activation_eps"] = activation_eps - return json.dumps(payload, indent=2) - - -def _legacy_compact_summary_without_native_smoke( - suite, - *, - generated_samples=1, - generated_frames=0, -): - metrics = {"generated_samples": generated_samples} - if generated_frames: - metrics["generated_frames"] = generated_frames - return json.dumps( - { - "published_summary": "compact", - "raw_generation_records": "local-only", - "activation_eps": 1e-10, - "quantization_device": "cuda", - "weight_quantization_backend": "triton_cuda", - "quantization_staging_mode": "component", - "metrics": { - "original": {"records": generated_samples, "latest_metrics": metrics}, - "orbitquant": {"records": generated_samples, "latest_metrics": metrics}, - }, - }, - indent=2, - ) - - -def test_native_smoke_expected_settings_include_video_export_fps_only_when_defined(): - image_suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - video_suite = NativeSuite( - name="wan-native", - model_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - pipeline="WanPipeline", - width=832, - height=480, - frames=81, - export_fps=16, - steps=50, - guidance=5.0, - bit_settings=["W4A4"], - ) - - assert "export_fps" not in hub_module._native_smoke_expected_settings(image_suite) - assert hub_module._native_smoke_expected_settings(video_suite)["export_fps"] == 16 - - -def _write_native_smoke_backup( - backup_root, - *, - repo_id, - suite, - bit_setting, - prompt_ids=("simple-object", "counting"), -): - assets = backup_root / repo_id.replace("/", "__") / "assets" - assets.mkdir(parents=True, exist_ok=True) - extension = "mp4" if suite.frames is not None else "png" - for prompt_id in prompt_ids: - for split_token, quantization in ( - ("original", None), - ( - bit_setting, - { - "config": { - "weight_bits": int(bit_setting.split("A", maxsplit=1)[0][1:]), - "activation_bits": int(bit_setting.split("A", maxsplit=1)[1]), - } - }, - ), - ): - media_path = assets / f"{suite.name}_seed0_{split_token}_{prompt_id}.{extension}" - media_path.write_bytes(b"not empty") - metadata = { - "suite": suite.name, - "model_id": suite.model_id, - "prompt": f"prompt for {prompt_id}", - "seed": 0, - "height": suite.height, - "width": suite.width, - "frames": suite.frames, - "steps": suite.steps, - "guidance": suite.guidance, - "quantization": quantization, - } - if suite.export_fps is not None: - metadata["export_fps"] = suite.export_fps - (media_path.with_suffix(media_path.suffix + ".json")).write_text( - json.dumps(metadata, indent=2) + "\n", - encoding="utf-8", - ) - - -def _expected_missing_geneval_metrics(): - return [ - {"split": split, "metric": metric} - for metric in hub_module._GENEVAL_REQUIRED_METRICS - for split in ("original", "orbitquant") - ] - - -def test_recover_native_smoke_proof_from_compact_summary_requires_raw_pair_evidence(): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - summary = json.loads(_legacy_compact_summary_without_native_smoke(suite)) - file_names = {"assets/image_generation_comparison_matrix.webp"} - - proof, reason = hub_module._recover_native_smoke_proof_from_compact_summary( - summary, - suite=suite, - file_names=file_names, - ) - - assert proof is None - assert reason == "raw_paired_native_smoke_evidence_missing" - - -def test_native_smoke_proof_status_rejects_recovered_pair_claims(): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - summary = json.loads(_native_smoke_summary(suite)) - summary["native_smoke"]["proof_source"] = ( - "recovered_from_compact_summary_and_published_comparison_matrix" - ) - - status = hub_module._native_smoke_proof_status( - summary, - suite=suite, - file_names={"assets/image_generation_comparison_matrix.webp"}, - ) - - assert status["ready"] is False - assert "native_smoke.raw_paired_native_smoke_evidence" in status["missing"] - - -def test_recover_native_smoke_proof_requires_uploaded_comparison_asset(): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - summary = json.loads(_legacy_compact_summary_without_native_smoke(suite)) - summary.pop("raw_generation_records") - - proof, reason = hub_module._recover_native_smoke_proof_from_compact_summary( - summary, - suite=suite, - file_names=set(), - ) - - assert proof is None - assert reason == "comparison_asset_missing" - - -def test_recover_native_smoke_proof_rejects_missing_generated_samples(): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - summary = json.loads(_legacy_compact_summary_without_native_smoke(suite, generated_samples=0)) - summary.pop("raw_generation_records") - - proof, reason = hub_module._recover_native_smoke_proof_from_compact_summary( - summary, - suite=suite, - file_names={"assets/image_generation_comparison_matrix.webp"}, - ) - - assert proof is None - assert reason == "original.generated_samples_missing" - - -def test_recover_native_smoke_proof_rejects_video_without_generated_frames(): - suite = NativeSuite( - name="wan-native", - model_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - pipeline="WanPipeline", - width=832, - height=480, - frames=81, - steps=50, - guidance=5.0, - bit_settings=["W4A4"], - ) - summary = json.loads(_legacy_compact_summary_without_native_smoke(suite)) - summary.pop("raw_generation_records") - - proof, reason = hub_module._recover_native_smoke_proof_from_compact_summary( - summary, - suite=suite, - file_names={"assets/video_generation_comparison_matrix.webp"}, - ) - - assert proof is None - assert reason == "original.generated_frames_insufficient" - - -def test_upload_orbitquant_artifact_dry_run_validates_without_hub_calls(tmp_path): - _write_artifact(tmp_path) - fake_api = FakeHfApi() - - result = upload_orbitquant_artifact( - tmp_path, - repo_id="WaveCut/example-orbitquant", - dry_run=True, - api=fake_api, - ) - - assert result["dry_run"] is True - assert result["repo_id"] == "WaveCut/example-orbitquant" - assert result["private"] is True - assert result["validation"]["valid"] is True - assert result["upload"] is None - assert fake_api.create_repo_calls == [] - assert fake_api.upload_folder_calls == [] - assert fake_api.model_info_calls == [] - - -def test_stage_compact_upload_artifact_omits_raw_eval_reports_and_promotes_matrices( - tmp_path, -): - _write_artifact(tmp_path) - raw_eval_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_geneval-00000.png" - raw_eval_metadata = raw_eval_asset.with_suffix(".png.json") - visual_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" - comparison_asset = ( - tmp_path / "assets" / "original_vs_orbitquant_flux2-native_seed0_W4A4_simple-object.webp" - ) - raw_video_asset = tmp_path / "assets" / "wan-native_seed0_W4A4_simple-motion.mp4" - raw_video_metadata = raw_video_asset.with_suffix(".mp4.json") - raw_eval_asset.write_bytes(b"raw eval image") - raw_eval_metadata.write_text('{"prompt_id":"geneval-00000"}\n', encoding="utf-8") - visual_asset.write_bytes(b"visual image") - comparison_asset.write_bytes(b"comparison") - raw_video_asset.write_bytes(b"raw video") - raw_video_metadata.write_text('{"prompt_id":"simple-motion"}\n', encoding="utf-8") - report_dir = tmp_path.parent / f"{tmp_path.name}-report" - report_matrix = report_dir / "assets" / "image_generation_comparison_matrix.webp" - report_matrix.parent.mkdir(parents=True) - report_matrix.write_bytes(b"report matrix") - report_markdown = report_dir / "orbitquant-native-eval.md" - report_markdown.write_text("# local report log\n", encoding="utf-8") - report_table = report_dir / "tables" / "perf.csv" - report_table.parent.mkdir(parents=True) - report_table.write_text("metric,value\n", encoding="utf-8") - in_artifact_report = ( - tmp_path - / "reports" - / "native" - / "flux2-w4a4" - / "assets" - / "image_generation_comparison_matrix.webp" - ) - in_artifact_report.parent.mkdir(parents=True) - in_artifact_report.write_bytes(b"in artifact report matrix") - in_artifact_markdown = in_artifact_report.parent.parent / "orbitquant-native-eval.md" - in_artifact_markdown.write_text("# local artifact report log\n", encoding="utf-8") - refresh_artifact_checksums(tmp_path) - - stage_dir = tmp_path.parent / f"{tmp_path.name}-stage" - result = stage_compact_upload_artifact( - tmp_path, - stage_dir, - report_dirs=[report_dir], - validate_tensors=False, - ) - - staged_checksums = read_sha256sums(stage_dir / "SHA256SUMS") - assert result["validation"]["valid"] is True - assert result["omitted_raw_eval_asset_count"] == 6 - assert result["omitted_report_file_count"] == 2 - assert result["copied_report_asset_count"] == 1 - assert not (stage_dir / raw_eval_asset.relative_to(tmp_path)).exists() - assert not (stage_dir / raw_eval_metadata.relative_to(tmp_path)).exists() - assert not (stage_dir / visual_asset.relative_to(tmp_path)).exists() - assert not (stage_dir / comparison_asset.relative_to(tmp_path)).exists() - assert not (stage_dir / raw_video_asset.relative_to(tmp_path)).exists() - assert not (stage_dir / raw_video_metadata.relative_to(tmp_path)).exists() - assert not (stage_dir / in_artifact_report.relative_to(tmp_path)).exists() - assert not (stage_dir / in_artifact_markdown.relative_to(tmp_path)).exists() - staged_report = stage_dir / "assets" / "image_generation_comparison_matrix.webp" - staged_collision_report = ( - stage_dir / "assets" / f"{report_dir.name}_image_generation_comparison_matrix.webp" - ) - assert staged_report.is_file() - assert staged_report.read_bytes() == b"in artifact report matrix" - assert not staged_collision_report.exists() - assert raw_eval_asset.relative_to(tmp_path).as_posix() not in staged_checksums - assert visual_asset.relative_to(tmp_path).as_posix() not in staged_checksums - assert comparison_asset.relative_to(tmp_path).as_posix() not in staged_checksums - assert raw_video_asset.relative_to(tmp_path).as_posix() not in staged_checksums - assert not any(path.startswith("reports/") for path in staged_checksums) - assert staged_report.relative_to(stage_dir).as_posix() in staged_checksums - - -def test_stage_compact_upload_artifact_keeps_one_direct_comparison_matrix(tmp_path): - _write_artifact(tmp_path) - image_matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - video_matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" - image_matrix.write_bytes(b"image matrix") - video_matrix.write_bytes(b"video matrix") - report_matrix = ( - tmp_path - / "reports" - / "native" - / "flux2-w4a4" - / "assets" - / "image_generation_comparison_matrix.webp" - ) - report_matrix.parent.mkdir(parents=True) - report_matrix.write_bytes(b"report matrix") - refresh_artifact_checksums(tmp_path) - - stage_dir = tmp_path.parent / f"{tmp_path.name}-stage" - result = stage_compact_upload_artifact( - tmp_path, - stage_dir, - validate_tensors=False, - ) - - staged_checksums = read_sha256sums(stage_dir / "SHA256SUMS") - assert result["copied_report_asset_count"] == 0 - assert (stage_dir / "assets" / "image_generation_comparison_matrix.webp").read_bytes() == ( - b"image matrix" - ) - assert not (stage_dir / "assets" / "video_generation_comparison_matrix.webp").exists() - assert "assets/image_generation_comparison_matrix.webp" in staged_checksums - assert "assets/video_generation_comparison_matrix.webp" not in staged_checksums - - -def test_stage_compact_upload_artifact_enforces_asset_allowlist(tmp_path): - _write_artifact(tmp_path) - allowed_matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - extra_matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" - raw_png = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" - raw_png_sidecar = raw_png.with_suffix(".png.json") - raw_webp = tmp_path / "assets" / "original_vs_orbitquant_seed0.webp" - contact_sheet = tmp_path / "assets" / "wan_contact_sheet.webp" - raw_video = tmp_path / "assets" / "wan-native_seed0_W4A4_motion.mp4" - raw_video_sidecar = raw_video.with_suffix(".mp4.json") - nested_raw = tmp_path / "assets" / "nested" / "debug_frame.png" - nested_matrix = tmp_path / "assets" / "nested" / "debug_generation_comparison_matrix.webp" - unexpected_file = tmp_path / "debug.txt" - ignored_cache = tmp_path / ".cache" / "upload.tmp" - for path, payload in ( - (allowed_matrix, b"matrix"), - (extra_matrix, b"video matrix"), - (raw_png, b"png"), - (raw_png_sidecar, b"json"), - (raw_webp, b"webp"), - (contact_sheet, b"contact"), - (raw_video, b"mp4"), - (raw_video_sidecar, b"mp4 json"), - (nested_raw, b"nested"), - (nested_matrix, b"nested matrix"), - (unexpected_file, b"debug"), - (ignored_cache, b"ignored"), - ): - path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes(payload) - refresh_artifact_checksums(tmp_path) - - stage_dir = tmp_path.parent / f"{tmp_path.name}-asset-allowlist-stage" - result = stage_compact_upload_artifact(tmp_path, stage_dir, validate_tensors=False) - - staged_files = { - path.relative_to(stage_dir).as_posix() for path in stage_dir.rglob("*") if path.is_file() - } - expected_files = set(hub_module._REQUIRED_ARTIFACT_FILES) | { - "assets/image_generation_comparison_matrix.webp" - } - assert staged_files == expected_files - - staged_manifest = json.loads((stage_dir / "orbitquant_manifest.json").read_text()) - manifest_checksum_files = set(staged_manifest["checksums"]) - sha256sum_files = set(read_sha256sums(stage_dir / "SHA256SUMS")) - assert manifest_checksum_files == staged_files - { - "README.md", - "SHA256SUMS", - "orbitquant_manifest.json", - } - assert sha256sum_files == staged_files - {"SHA256SUMS"} - - staged_summary = json.loads((stage_dir / "benchmark" / "summary.json").read_text()) - assert staged_summary["published_summary"] == "compact" - assert staged_summary["raw_generation_records"] == "local-only" - assert not any(path.endswith(".png") for path in staged_files) - assert not any(path.endswith(".png.json") for path in staged_files) - assert not any(path.endswith(".mp4") for path in staged_files) - assert "assets/video_generation_comparison_matrix.webp" not in staged_files - assert "assets/original_vs_orbitquant_seed0.webp" not in staged_files - assert "assets/wan_contact_sheet.webp" not in staged_files - assert "assets/nested/debug_frame.png" not in staged_files - assert "assets/nested/debug_generation_comparison_matrix.webp" not in staged_files - assert "debug.txt" not in staged_files - assert ".cache/upload.tmp" not in staged_files - assert result["omitted_unexpected_file_count"] == 1 - assert result["omitted_unexpected_files"] == ["debug.txt"] - assert set(result["omitted_raw_eval_assets"]) >= { - "assets/video_generation_comparison_matrix.webp", - "assets/flux2-native_seed0_W4A4_simple-object.png", - "assets/flux2-native_seed0_W4A4_simple-object.png.json", - "assets/original_vs_orbitquant_seed0.webp", - "assets/wan_contact_sheet.webp", - "assets/wan-native_seed0_W4A4_motion.mp4", - "assets/wan-native_seed0_W4A4_motion.mp4.json", - "assets/nested/debug_frame.png", - "assets/nested/debug_generation_comparison_matrix.webp", - } - - -def test_stage_compact_upload_artifact_writes_native_smoke_proof(tmp_path): - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - original_output = tmp_path / "assets" / "flux2-native_seed0_original_simple-object.png" - orbitquant_output = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" - for path, payload in ( - (matrix, b"matrix"), - (original_output, b"original image"), - (orbitquant_output, b"orbitquant image"), - ): - path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes(payload) - for split, output, bit_setting in ( - ("original", original_output, "original"), - ("orbitquant", orbitquant_output, "W4A4"), - ): - record_artifact_metrics( - tmp_path, - split=split, - metrics={"generated_samples": 1}, - metadata={ - "suite": "flux2-native", - "prompt": "A native prompt", - "prompt_record": {"id": "simple-object"}, - "seed": 0, - "height": 1024, - "width": 1024, - "frames": None, - "steps": 4, - "guidance": 1.0, - "bit_setting": bit_setting, - "output_path": str(output), - }, - validate_checksums_enabled=False, - refresh_checksums_enabled=False, - ) - refresh_artifact_checksums(tmp_path) - - stage_dir = tmp_path.parent / f"{tmp_path.name}-native-smoke-stage" - stage_compact_upload_artifact(tmp_path, stage_dir, validate_tensors=False) - - staged_summary = json.loads((stage_dir / "benchmark" / "summary.json").read_text()) - staged_readme = (stage_dir / "README.md").read_text(encoding="utf-8") - assert staged_summary["raw_generation_records"] == "local-only" - assert "latest" not in staged_summary["metrics"]["original"] - assert "## Native Validation Evidence" in staged_readme - assert "| Comparison matrix | `assets/image_generation_comparison_matrix.webp` |" in ( - staged_readme - ) - assert "| Paired prompt/seed count | `1` |" in staged_readme - assert "| BF16 source nonempty outputs | `1` |" in staged_readme - assert "| OrbitQuant nonempty outputs | `1` |" in staged_readme - assert "original.metrics.jsonl" not in staged_readme - assert staged_summary["native_smoke"] == { - "proof_format": "orbitquant-native-smoke-v1", - "comparison_asset_path": "assets/image_generation_comparison_matrix.webp", - "paired_prompt_seed_count": 1, - "paired_prompt_seed_keys": [["flux2-native", "0", "simple-object"]], - "splits": { - "original": { - "records": 1, - "generated_samples": 1, - "generated_frames": 0, - "nonempty_output_count": 1, - "seeds": ["0"], - "prompt_ids": ["simple-object"], - "pair_keys": [["flux2-native", "0", "simple-object"]], - "native_settings": [ - { - "suite": "flux2-native", - "height": 1024, - "width": 1024, - "frames": None, - "steps": 4, - "guidance": 1.0, - } - ], - }, - "orbitquant": { - "records": 1, - "generated_samples": 1, - "generated_frames": 0, - "nonempty_output_count": 1, - "seeds": ["0"], - "prompt_ids": ["simple-object"], - "pair_keys": [["flux2-native", "0", "simple-object"]], - "native_settings": [ - { - "suite": "flux2-native", - "height": 1024, - "width": 1024, - "frames": None, - "steps": 4, - "guidance": 1.0, - } - ], - }, - }, - } - - -def test_stage_compact_upload_artifact_promotes_compare_native_bundle(tmp_path): - _write_artifact(tmp_path) - bundle_dir = tmp_path.parent / f"{tmp_path.name}-compare-native" - bundle_dir.mkdir() - comparison = bundle_dir / "flux2-native_seed0_W4A4_original_vs_orbitquant.webp" - original = bundle_dir / "flux2-native_seed0_original.png" - orbitquant = bundle_dir / "flux2-native_seed0_W4A4.png" - for path, payload in ( - (comparison, b"comparison"), - (original, b"original"), - (orbitquant, b"orbitquant"), - ): - path.write_bytes(payload) - metadata = { - "suite": "flux2-native", - "prompt": "A native prompt", - "prompt_record": {"id": "simple-object"}, - "seed": 0, - "height": 1024, - "width": 1024, - "frames": None, - "steps": 4, - "guidance": 1.0, - } - for output in (original, orbitquant): - output.with_suffix(output.suffix + ".json").write_text( - json.dumps(metadata) + "\n", encoding="utf-8" - ) - (bundle_dir / "summary.json").write_text( - json.dumps( - { - **metadata, - "model_id": "example/model", - "bit_setting": "W4A4", - "comparison_path": "/workspace/run/" + comparison.name, - "original": { - "output_path": "/workspace/run/" + original.name, - "metadata_path": "/workspace/run/" + original.name + ".json", - }, - "orbitquant": { - "output_path": "/workspace/run/" + orbitquant.name, - "metadata_path": "/workspace/run/" + orbitquant.name + ".json", - }, - } - ) - + "\n", - encoding="utf-8", - ) - - stage_dir = tmp_path.parent / f"{tmp_path.name}-compare-native-stage" - result = stage_compact_upload_artifact( - tmp_path, - stage_dir, - report_dirs=[bundle_dir], - validate_tensors=False, - ) - - matrix_path = stage_dir / "assets" / "image_generation_comparison_matrix.webp" - summary = json.loads((stage_dir / "benchmark" / "summary.json").read_text()) - proof = summary["native_smoke"] - assert result["copied_report_assets"] == ["assets/image_generation_comparison_matrix.webp"] - assert matrix_path.read_bytes() == b"comparison" - assert proof["proof_source"] == "local_compare_native_bundle" - assert proof["comparison_asset_path"] == ("assets/image_generation_comparison_matrix.webp") - assert proof["paired_prompt_seed_keys"] == [["flux2-native", "0", "simple-object"]] - assert proof["splits"]["original"]["nonempty_output_count"] == 1 - assert proof["splits"]["orbitquant"]["nonempty_output_count"] == 1 - assert not (stage_dir / original.name).exists() - assert not (stage_dir / orbitquant.name).exists() - - -def test_repair_hf_artifact_metadata_dry_run_preserves_large_file_checksum( - tmp_path, - monkeypatch, -): - repo_id = "WaveCut/example-orbitquant" - _write_artifact(tmp_path) - old_checksums = read_sha256sums(tmp_path / "SHA256SUMS") - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCommitHfApi() - - result = repair_hf_artifact_metadata( - repo_id=repo_id, - quantization_device="cuda", - weight_quantization_backend="triton_cuda", - quantization_staging_mode="component", - dry_run=True, - api=fake_api, - ) - - assert result["dry_run"] is True - assert result["updated"]["quantization_staging_mode"] == "component" - assert result["commit"] is None - assert "model.safetensors" in result["preserved_checksum_entries"] - assert old_checksums["model.safetensors"] - assert fake_api.create_commit_calls == [] - - -def test_repair_hf_artifact_metadata_commits_only_metadata_files_and_sha256sums( - tmp_path, - monkeypatch, -): - repo_id = "WaveCut/example-orbitquant" - _write_artifact(tmp_path) - stale_cache_metadata = tmp_path / ".cache" / "huggingface" / "download" / "README.md.metadata" - stale_cache_metadata.parent.mkdir(parents=True) - stale_cache_metadata.write_text("transient hub metadata", encoding="utf-8") - with (tmp_path / "SHA256SUMS").open("a", encoding="utf-8") as handle: - handle.write("0" * 64 + " .cache/huggingface/download/README.md.metadata\n") - old_checksums = read_sha256sums(tmp_path / "SHA256SUMS") - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCommitHfApi() - - result = repair_hf_artifact_metadata( - repo_id=repo_id, - quantization_device="cuda", - weight_quantization_backend="triton_cuda", - quantization_staging_mode="component", - revision="main", - commit_message="repair metadata", - api=fake_api, - ) - - assert result["commit"]["commit_oid"] == "repair-sha" - assert len(fake_api.create_commit_calls) == 1 - commit_call = fake_api.create_commit_calls[0] - assert commit_call["repo_id"] == repo_id - assert commit_call["repo_type"] == "model" - assert commit_call["revision"] == "main" - assert commit_call["commit_message"] == "repair metadata" - operation_paths = {operation.path_in_repo for operation in commit_call["operations"]} - assert operation_paths == { - "orbitquant_manifest.json", - "model_index.json", - "benchmark/summary.json", - "README.md", - "SHA256SUMS", - } - assert "model.safetensors" not in operation_paths - sha_operation = next( - operation - for operation in commit_call["operations"] - if operation.path_in_repo == "SHA256SUMS" - ) - sha_entries = hub_module._parse_sha256sums_bytes(sha_operation.path_or_fileobj) - assert sha_entries["model.safetensors"] == old_checksums["model.safetensors"] - assert sha_entries["orbitquant_manifest.json"] != old_checksums["orbitquant_manifest.json"] - assert ".cache/huggingface/download/README.md.metadata" not in sha_entries - manifest_operation = next( - operation - for operation in commit_call["operations"] - if operation.path_in_repo == "orbitquant_manifest.json" - ) - assert b'"quantization_staging_mode": "component"' in manifest_operation.path_or_fileobj - assert b'"activation_eps": 1e-10' in manifest_operation.path_or_fileobj - - -def test_cleanup_hf_artifact_reports_promotes_matrices_and_deletes_report_folder( - tmp_path, - monkeypatch, -): - repo_id = "WaveCut/example-orbitquant" - _write_artifact(tmp_path) - raw_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" - raw_sidecar = raw_asset.with_suffix(".png.json") - raw_asset.write_bytes(b"raw generated image") - raw_sidecar.write_text('{"prompt_id":"simple-object"}\n', encoding="utf-8") - report_matrix = ( - tmp_path - / "reports" - / "native" - / "flux2-w4a4" - / "assets" - / "image_generation_comparison_matrix.webp" - ) - report_matrix.parent.mkdir(parents=True) - report_matrix.write_bytes(b"report matrix") - report_markdown = report_matrix.parent.parent / "orbitquant-native-eval.md" - report_markdown.write_text("# local report log\n", encoding="utf-8") - refresh_artifact_checksums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_path(tmp_path, filename)) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = cleanup_hf_artifact_reports( - repo_id=repo_id, - commit_message="cleanup reports", - api=fake_api, - ) - - assert result["report_file_count"] == 2 - assert result["report_matrix_count"] == 1 - assert result["forbidden_file_count"] == 4 - assert result["promoted_assets"] == ["assets/image_generation_comparison_matrix.webp"] - assert result["delete_paths"] == [ - "assets/flux2-native_seed0_W4A4_simple-object.png", - "assets/flux2-native_seed0_W4A4_simple-object.png.json", - "reports", - ] - assert result["commit"]["commit_oid"] == "repair-sha" - assert len(fake_api.create_commit_calls) == 1 - - commit_call = fake_api.create_commit_calls[0] - assert commit_call["commit_message"] == "cleanup reports" - operations = commit_call["operations"] - added = { - operation.path_in_repo: operation.path_or_fileobj - for operation in operations - if isinstance(operation, CommitOperationAdd) - } - deleted = [ - operation for operation in operations if isinstance(operation, CommitOperationDelete) - ] - assert "model.safetensors" not in added - assert added["assets/image_generation_comparison_matrix.webp"] == b"report matrix" - assert len(deleted) == 3 - deleted_paths = {operation.path_in_repo: operation.is_folder for operation in deleted} - assert deleted_paths == { - "assets/flux2-native_seed0_W4A4_simple-object.png": False, - "assets/flux2-native_seed0_W4A4_simple-object.png.json": False, - "reports": True, - } - - next_manifest = json.loads(added["orbitquant_manifest.json"].decode("utf-8")) - next_summary = json.loads(added["benchmark/summary.json"].decode("utf-8")) - next_readme = added["README.md"].decode("utf-8") - next_sha = added["SHA256SUMS"].decode("utf-8") - assert "assets/image_generation_comparison_matrix.webp" in next_manifest["checksums"] - assert "benchmark/summary.json" in next_manifest["checksums"] - assert not any(path.endswith(".metrics.jsonl") for path in next_manifest["checksums"]) - assert not any(path.endswith(".metrics.csv") for path in next_manifest["checksums"]) - assert not any(path.startswith("reports/") for path in next_manifest["checksums"]) - assert next_summary["published_summary"] == "compact" - assert next_summary["raw_generation_records"] == "local-only" - assert "assets/image_generation_comparison_matrix.webp" in next_readme - assert "## Validation Status" in next_readme - assert "reports/native" not in next_readme - assert "assets/image_generation_comparison_matrix.webp" in next_sha - assert "benchmark/original.metrics.jsonl" not in next_sha - assert "reports/native" not in next_sha - - -def test_cleanup_hf_artifact_reports_deletes_extra_comparison_matrix_assets( - tmp_path, - monkeypatch, -): - repo_id = "WaveCut/example-orbitquant" - _write_artifact(tmp_path) - image_matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - video_matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" - image_matrix.write_bytes(b"image matrix") - video_matrix.write_bytes(b"video matrix") - refresh_artifact_checksums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_path(tmp_path, filename)) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = cleanup_hf_artifact_reports( - repo_id=repo_id, - commit_message="cleanup extra matrix", - api=fake_api, - ) - - assert result["promoted_assets"] == [] - assert result["forbidden_files"] == [ - "assets/video_generation_comparison_matrix.webp", - ] - assert "assets/video_generation_comparison_matrix.webp" in result["delete_paths"] - commit_call = fake_api.create_commit_calls[0] - operations = commit_call["operations"] - added = { - operation.path_in_repo: operation.path_or_fileobj - for operation in operations - if isinstance(operation, CommitOperationAdd) - } - deleted_paths = { - operation.path_in_repo - for operation in operations - if isinstance(operation, CommitOperationDelete) - } - assert "assets/video_generation_comparison_matrix.webp" in deleted_paths - next_manifest = json.loads(added["orbitquant_manifest.json"].decode("utf-8")) - next_readme = added.get("README.md", (tmp_path / "README.md").read_bytes()).decode("utf-8") - assert "assets/image_generation_comparison_matrix.webp" in next_manifest["checksums"] - assert "assets/video_generation_comparison_matrix.webp" not in next_manifest["checksums"] - assert "assets/image_generation_comparison_matrix.webp" in next_readme - assert "assets/video_generation_comparison_matrix.webp" not in next_readme - - -def test_cleanup_hf_artifact_reports_matrix_cleans_expected_suite_repo( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - _write_artifact(tmp_path) - report_matrix = ( - tmp_path - / "reports" - / "native" - / "flux2-w4a4" - / "assets" - / "image_generation_comparison_matrix.webp" - ) - report_matrix.parent.mkdir(parents=True) - report_matrix.write_bytes(b"report matrix") - refresh_artifact_checksums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_path(tmp_path, filename)) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = cleanup_hf_artifact_reports_matrix( - suites=[suite], - dry_run=True, - api=fake_api, - ) - - assert result["repo_count"] == 1 - assert result["changed_repo_count"] == 1 - assert result["error_count"] == 0 - assert result["report_file_count"] == 1 - assert result["promoted_asset_count"] == 1 - assert result["rows"][0]["repo_id"] == repo_id - assert result["rows"][0]["suite"] == "flux2-native" - assert result["rows"][0]["bit_setting"] == "W4A4" - - -def test_repair_hf_artifact_metadata_matrix_repairs_expected_suite_repo( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - _write_artifact(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCommitHfApi() - - result = repair_hf_artifact_metadata_matrix( - suites=[suite], - quantization_device="cuda", - weight_quantization_backend="triton_cuda", - quantization_staging_mode="component", - dry_run=True, - api=fake_api, - ) - - assert result["repo_count"] == 1 - assert result["error_count"] == 0 - assert result["rows"][0]["repo_id"] == repo_id - assert result["rows"][0]["suite"] == "flux2-native" - assert result["rows"][0]["bit_setting"] == "W4A4" - - -def test_repair_hf_native_smoke_proof_skips_compact_summary_without_raw_pair_evidence( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - matrix.parent.mkdir(parents=True, exist_ok=True) - matrix.write_bytes(b"matrix") - summary_path = tmp_path / "benchmark" / "summary.json" - summary_path.write_text( - _legacy_compact_summary_without_native_smoke(suite), - encoding="utf-8", - ) - refresh_artifact_checksums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = repair_hf_native_smoke_proof( - repo_id=repo_id, - suite=suite, - revision="main", - commit_message="repair native smoke proof", - api=fake_api, - ) - - assert result["commit"] is None - assert result["repair_skipped_reason"] == "raw_paired_native_smoke_evidence_missing" - assert result["changed_files"] == [] - assert fake_api.create_commit_calls == [] - - -def test_repair_hf_native_smoke_proof_removes_recovered_pair_claim( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - matrix.parent.mkdir(parents=True, exist_ok=True) - matrix.write_bytes(b"matrix") - summary_path = tmp_path / "benchmark" / "summary.json" - summary = json.loads(_native_smoke_summary(suite)) - summary["native_smoke"]["proof_source"] = ( - "recovered_from_compact_summary_and_published_comparison_matrix" - ) - summary_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") - refresh_artifact_checksums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = repair_hf_native_smoke_proof( - repo_id=repo_id, - suite=suite, - revision="main", - api=fake_api, - ) - - assert result["commit"]["commit_oid"] == "repair-sha" - assert result["removed_invalid_native_smoke"] is True - assert result["changed_files"] == [ - "benchmark/summary.json", - "orbitquant_manifest.json", - "README.md", - "SHA256SUMS", - ] - operation_by_path = { - operation.path_in_repo: operation.path_or_fileobj - for operation in fake_api.create_commit_calls[0]["operations"] - } - repaired_summary = json.loads(operation_by_path["benchmark/summary.json"]) - repaired_readme = operation_by_path["README.md"].decode("utf-8") - repaired_manifest = json.loads(operation_by_path["orbitquant_manifest.json"]) - sha_entries = hub_module._parse_sha256sums_bytes(operation_by_path["SHA256SUMS"]) - - assert "native_smoke" not in repaired_summary - assert "## Native Validation Proof" not in repaired_readme - assert ( - repaired_manifest["checksums"]["benchmark/summary.json"] - == sha_entries["benchmark/summary.json"] - ) - assert sha_entries["README.md"] == hub_module._sha256_bytes(operation_by_path["README.md"]) - - -def test_repair_hf_native_smoke_proof_restores_from_local_backup_without_raw_upload( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - matrix.parent.mkdir(parents=True, exist_ok=True) - matrix.write_bytes(b"matrix") - (tmp_path / "benchmark" / "summary.json").write_text( - _legacy_compact_summary_without_native_smoke(suite), - encoding="utf-8", - ) - refresh_artifact_checksums(tmp_path) - backup_root = tmp_path / "native-backup" - _write_native_smoke_backup( - backup_root, - repo_id=repo_id, - suite=suite, - bit_setting="W4A4", - ) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = repair_hf_native_smoke_proof( - repo_id=repo_id, - suite=suite, - native_smoke_backup_root=backup_root, - revision="main", - commit_message="restore native smoke proof", - api=fake_api, - ) - - assert result["commit"]["commit_oid"] == "repair-sha" - assert result["backup_skipped_reason"] is None - assert result["proof_source"] == "local_paired_native_smoke_backup" - assert result["changed_files"] == [ - "benchmark/summary.json", - "orbitquant_manifest.json", - "README.md", - "SHA256SUMS", - ] - operation_by_path = { - operation.path_in_repo: operation.path_or_fileobj - for operation in fake_api.create_commit_calls[0]["operations"] - } - assert sorted(operation_by_path) == sorted(result["changed_files"]) - - repaired_summary = json.loads(operation_by_path["benchmark/summary.json"]) - proof = repaired_summary["native_smoke"] - assert proof["proof_source"] == "local_paired_native_smoke_backup" - assert proof["comparison_asset_path"] == "assets/image_generation_comparison_matrix.webp" - assert proof["paired_prompt_seed_count"] == 2 - assert proof["paired_prompt_seed_keys"] == [ - ["flux2-native", "0", "counting"], - ["flux2-native", "0", "simple-object"], - ] - assert proof["splits"]["original"]["generated_samples"] == 2 - assert proof["splits"]["orbitquant"]["generated_samples"] == 2 - assert proof["splits"]["original"]["nonempty_output_count"] == 2 - assert proof["splits"]["orbitquant"]["nonempty_output_count"] == 2 - assert proof["splits"]["original"]["native_settings"] == [ - hub_module._native_smoke_expected_settings(suite) - ] - assert "## Native Validation Evidence" in operation_by_path["README.md"].decode("utf-8") - - -def test_repair_hf_native_smoke_proof_refreshes_stale_readme_when_proof_exists( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - matrix.parent.mkdir(parents=True, exist_ok=True) - matrix.write_bytes(b"matrix") - summary_path = tmp_path / "benchmark" / "summary.json" - summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") - refresh_artifact_checksums(tmp_path) - manifest = hub_module.OrbitQuantManifest.from_dict( - json.loads((tmp_path / "orbitquant_manifest.json").read_text()) - ) - (tmp_path / "README.md").write_text(hub_module.render_model_card(manifest), encoding="utf-8") - write_sha256sums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = repair_hf_native_smoke_proof( - repo_id=repo_id, - suite=suite, - revision="main", - api=fake_api, - ) - - assert result["existing_native_smoke_ready"] is True - assert result["changed_files"] == ["README.md", "SHA256SUMS"] - operation_by_path = { - operation.path_in_repo: operation.path_or_fileobj - for operation in fake_api.create_commit_calls[0]["operations"] - } - repaired_readme = operation_by_path["README.md"].decode("utf-8") - assert "## Native Validation Evidence" in repaired_readme - assert "| Paired prompt/seed count | `1` |" in repaired_readme - - -def test_repair_hf_native_smoke_proof_normalizes_legacy_video_export_fps( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="wan-native", - model_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - pipeline="WanPipeline", - width=832, - height=480, - frames=81, - export_fps=16, - steps=50, - guidance=5.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/Wan2.1-T2V-1.3B-Diffusers-OrbitQuant-W4A4" - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" - matrix.parent.mkdir(parents=True, exist_ok=True) - matrix.write_bytes(b"matrix") - summary_path = tmp_path / "benchmark" / "summary.json" - summary = json.loads(_native_smoke_summary(suite)) - summary["native_smoke"]["comparison_asset_path"] = ( - "assets/video_generation_comparison_matrix.webp" - ) - summary["native_smoke"]["proof_source"] = "local_paired_native_smoke_backup" - for split in ("original", "orbitquant"): - del summary["native_smoke"]["splits"][split]["native_settings"][0]["export_fps"] - summary_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") - refresh_artifact_checksums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = repair_hf_native_smoke_proof( - repo_id=repo_id, - suite=suite, - revision="main", - api=fake_api, - ) - - assert result["existing_native_smoke_ready"] is False - assert result["normalized_native_smoke_settings"] is True - assert set(result["changed_files"]) >= { - "benchmark/summary.json", - "orbitquant_manifest.json", - "SHA256SUMS", - } - assert "model.safetensors" not in result["changed_files"] - operation_by_path = { - operation.path_in_repo: operation.path_or_fileobj - for operation in fake_api.create_commit_calls[0]["operations"] - } - repaired_summary = json.loads(operation_by_path["benchmark/summary.json"]) - expected_settings = hub_module._native_smoke_expected_settings(suite) - for split in ("original", "orbitquant"): - assert repaired_summary["native_smoke"]["splits"][split]["native_settings"] == [ - expected_settings - ] - status = hub_module._native_smoke_proof_status( - repaired_summary, - suite=suite, - file_names={"assets/video_generation_comparison_matrix.webp"}, - ) - assert status["ready"] is True - - -def test_repair_hf_native_smoke_proof_skips_when_video_frames_are_insufficient( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="wan-native", - model_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - pipeline="WanPipeline", - width=832, - height=480, - frames=81, - steps=50, - guidance=5.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/Wan2.1-T2V-1.3B-Diffusers-OrbitQuant-W4A4" - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" - matrix.parent.mkdir(parents=True, exist_ok=True) - matrix.write_bytes(b"matrix") - (tmp_path / "benchmark" / "summary.json").write_text( - _legacy_compact_summary_without_native_smoke( - suite, - generated_samples=1, - generated_frames=1, - ), - encoding="utf-8", - ) - refresh_artifact_checksums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = repair_hf_native_smoke_proof( - repo_id=repo_id, - suite=suite, - dry_run=True, - api=fake_api, - ) - - assert result["commit"] is None - assert result["changed_files"] == [] - assert result["repair_skipped_reason"] == "raw_paired_native_smoke_evidence_missing" - assert fake_api.create_commit_calls == [] - - -def test_repair_hf_native_smoke_proof_matrix_skips_compact_summary_without_raw_pair_evidence( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - matrix.parent.mkdir(parents=True, exist_ok=True) - matrix.write_bytes(b"matrix") - (tmp_path / "benchmark" / "summary.json").write_text( - _legacy_compact_summary_without_native_smoke(suite), - encoding="utf-8", - ) - refresh_artifact_checksums(tmp_path) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = repair_hf_native_smoke_proof_matrix( - suites=[suite], - dry_run=True, - api=fake_api, - ) - - assert result["repo_count"] == 1 - assert result["changed_repo_count"] == 0 - assert result["skipped_repo_count"] == 1 - assert result["error_count"] == 0 - assert result["rows"][0]["repo_id"] == repo_id - assert result["rows"][0]["suite"] == "flux2-native" - assert result["rows"][0]["bit_setting"] == "W4A4" - assert result["rows"][0]["repair_skipped_reason"] == ( - "raw_paired_native_smoke_evidence_missing" - ) - - -def test_repair_hf_native_smoke_proof_matrix_uses_local_backup( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - _write_artifact(tmp_path) - matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" - matrix.parent.mkdir(parents=True, exist_ok=True) - matrix.write_bytes(b"matrix") - (tmp_path / "benchmark" / "summary.json").write_text( - _legacy_compact_summary_without_native_smoke(suite), - encoding="utf-8", - ) - refresh_artifact_checksums(tmp_path) - backup_root = tmp_path / "native-backup" - _write_native_smoke_backup( - backup_root, - repo_id=repo_id, - suite=suite, - bit_setting="W4A4", - prompt_ids=("simple-object",), - ) - - def fake_download(repo, filename, **kwargs): - return str(_remote_file_map(repo_id, tmp_path)[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - fake_api = FakeCleanupHfApi(tmp_path) - - result = repair_hf_native_smoke_proof_matrix( - suites=[suite], - native_smoke_backup_root=backup_root, - dry_run=True, - api=fake_api, - ) - - assert result["repo_count"] == 1 - assert result["changed_repo_count"] == 1 - assert result["skipped_repo_count"] == 0 - assert result["error_count"] == 0 - assert result["rows"][0]["repo_id"] == repo_id - assert result["rows"][0]["proof_source"] == "local_paired_native_smoke_backup" - assert result["rows"][0]["backup_skipped_reason"] is None - assert result["rows"][0]["changed_files"] == [ - "benchmark/summary.json", - "orbitquant_manifest.json", - "README.md", - "SHA256SUMS", - ] - - -def test_upload_orbitquant_artifact_rejects_raw_full_upload_profile(tmp_path): - _write_artifact(tmp_path) - fake_api = FakeHfApi() - - with pytest.raises(ValueError, match="unsupported upload profile: full"): - upload_orbitquant_artifact( - tmp_path, - repo_id="WaveCut/example-orbitquant", - private=False, - revision="main", - commit_message="upload test artifact", - replace_repo_files=True, - validate_tensors=False, - upload_profile="full", - api=fake_api, - ) - - assert fake_api.create_repo_calls == [] - assert fake_api.upload_folder_calls == [] - - -def test_upload_orbitquant_artifact_can_upload_compact_staged_copy(tmp_path): - _write_artifact(tmp_path) - raw_eval_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_geneval-00000.png" - visual_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" - raw_eval_asset.write_bytes(b"raw eval image") - visual_asset.write_bytes(b"visual image") - refresh_artifact_checksums(tmp_path) - stage_dir = tmp_path.parent / f"{tmp_path.name}-upload-stage" - fake_api = FakeHfApi() - - result = upload_orbitquant_artifact( - tmp_path, - repo_id="WaveCut/example-orbitquant", - upload_profile="compact", - staging_dir=stage_dir, - validate_tensors=False, - api=fake_api, - ) - - assert result["upload_profile"] == "compact" - assert result["staging"]["enabled"] is True - assert result["staging"]["omitted_raw_eval_asset_count"] == 2 - assert len(fake_api.upload_folder_calls) == 1 - assert fake_api.upload_folder_calls[0]["folder_path"] == str(stage_dir) - assert not (stage_dir / raw_eval_asset.relative_to(tmp_path)).exists() - assert not (stage_dir / visual_asset.relative_to(tmp_path)).exists() - assert not (stage_dir / "benchmark" / "original.metrics.jsonl").exists() - assert (stage_dir / "benchmark" / "summary.json").is_file() - staged_summary = json.loads((stage_dir / "benchmark" / "summary.json").read_text()) - assert staged_summary["published_summary"] == "compact" - assert staged_summary["raw_generation_records"] == "local-only" - assert result["validation"]["valid"] is True - assert result["upload"]["commit_oid"] == "uploaded-sha" - - -def test_upload_orbitquant_artifact_defaults_to_compact_staged_copy(tmp_path): - _write_artifact(tmp_path) - raw_eval_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_geneval-00000.png" - visual_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" - raw_eval_asset.write_bytes(b"raw eval image") - visual_asset.write_bytes(b"visual image") - refresh_artifact_checksums(tmp_path) - fake_api = FakeHfApi() - - result = upload_orbitquant_artifact( - tmp_path, - repo_id="WaveCut/example-orbitquant", - validate_tensors=False, - api=fake_api, - ) - - assert result["upload_profile"] == "compact" - assert result["replace_repo_files"] is True - assert result["upload_kwargs"]["delete_patterns"] == "*" - assert result["staging"]["enabled"] is True - assert result["staging"]["omitted_raw_eval_asset_count"] == 2 - assert set(result["staging"]["omitted_raw_eval_assets"]) == { - "assets/flux2-native_seed0_W4A4_geneval-00000.png", - "assets/flux2-native_seed0_W4A4_simple-object.png", - } - upload_path = Path(fake_api.upload_folder_calls[0]["folder_path"]) - assert upload_path != tmp_path - - -def test_upload_orbitquant_artifact_rejects_remote_file_replacement_opt_out(tmp_path): - _write_artifact(tmp_path) - fake_api = FakeHfApi() - - with pytest.raises(ValueError, match="must replace remote files"): - upload_orbitquant_artifact( - tmp_path, - repo_id="WaveCut/example-orbitquant", - replace_repo_files=False, - validate_tensors=False, - api=fake_api, - ) - - assert fake_api.create_repo_calls == [] - assert fake_api.upload_folder_calls == [] - assert fake_api.model_info_calls == [] - - -def test_upload_orbitquant_artifact_rejects_invalid_artifact_before_hub_calls(tmp_path): - fake_api = FakeHfApi() - - with pytest.raises(RuntimeError, match="required artifact file missing"): - upload_orbitquant_artifact( - tmp_path, - repo_id="WaveCut/bad-artifact", - validate_tensors=False, - api=fake_api, - ) - - assert fake_api.create_repo_calls == [] - assert fake_api.upload_folder_calls == [] - assert fake_api.model_info_calls == [] - - -def test_fetch_hf_artifacts_dry_run_reports_native_artifact_layout(tmp_path): - suite = NativeSuite( - name="flux1-schnell-native", - model_id="black-forest-labs/FLUX.1-schnell", - pipeline="FluxPipeline", - width=1024, - height=1024, - steps=4, - guidance=0.0, - bit_settings=["W4A4", "W3A3"], - metric="geneval", - ) - - result = fetch_hf_artifacts( - namespace="WaveCut", - suites=[suite], - output_root=tmp_path / "artifacts", - dry_run=True, - ) - - assert result["repo_count"] == 2 - assert result["downloaded_count"] == 0 - assert result["dry_run"] is True - assert result["rows"][0]["repo_id"] == "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4" - assert result["rows"][0]["artifact_dir"].endswith("flux1-schnell-native-w4a4") - assert result["rows"][1]["repo_id"] == "WaveCut/FLUX.1-schnell-OrbitQuant-W3A3" - assert result["rows"][1]["artifact_dir"].endswith("flux1-schnell-native-w3a3") - - -def test_fetch_hf_artifacts_downloads_and_validates_artifact( - tmp_path, - monkeypatch, -): - source_dir = tmp_path / "source" - _write_artifact(source_dir) - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - calls = [] - stages = [] - - def fake_snapshot_download(**kwargs): - calls.append(kwargs) - destination = Path(kwargs["local_dir"]) - shutil.copytree(source_dir, destination, dirs_exist_ok=True) - return str(destination) - - monkeypatch.setattr(hub_module, "snapshot_download", fake_snapshot_download) - - result = fetch_hf_artifacts( - namespace="WaveCut", - suites=[suite], - output_root=tmp_path / "artifacts", - revision="main", - stage_logger=lambda event, label: stages.append((event, label)), - ) - - artifact_dir = tmp_path / "artifacts" / "flux2-native-w4a4" - assert result["downloaded_count"] == 1 - assert result["skipped_existing_count"] == 0 - assert result["rows"][0]["artifact_dir"] == str(artifact_dir) - assert result["rows"][0]["validation"]["valid"] is True - assert calls == [ - { - "repo_id": "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4", - "repo_type": "model", - "revision": "main", - "local_dir": artifact_dir, - "force_download": False, - "local_files_only": False, - } - ] - assert stages == [ - ("START", "flux2-native W4A4 fetch WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4"), - ("END", "flux2-native W4A4 fetch WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4"), - ] - - -def test_fetch_hf_artifacts_resume_skips_valid_existing_artifact(tmp_path, monkeypatch): - artifact_dir = tmp_path / "artifacts" / "flux2-native-w4a4" - _write_artifact(artifact_dir) - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - - def fail_snapshot_download(**kwargs): - raise AssertionError("resume should not download a valid existing artifact") - - monkeypatch.setattr(hub_module, "snapshot_download", fail_snapshot_download) - - result = fetch_hf_artifacts( - namespace="WaveCut", - suites=[suite], - output_root=tmp_path / "artifacts", - ) - - assert result["downloaded_count"] == 0 - assert result["skipped_existing_count"] == 1 - assert result["rows"][0]["skipped_existing"] is True - assert result["rows"][0]["validation"]["valid"] is True - - -def test_fetch_hf_artifacts_logs_error_stage_on_download_failure(tmp_path, monkeypatch): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - stages = [] - - def fail_snapshot_download(**kwargs): - raise RuntimeError("download failed") - - monkeypatch.setattr(hub_module, "snapshot_download", fail_snapshot_download) - - with pytest.raises(RuntimeError, match="download failed"): - fetch_hf_artifacts( - namespace="WaveCut", - suites=[suite], - output_root=tmp_path / "artifacts", - stage_logger=lambda event, label: stages.append((event, label)), - ) - - assert stages == [ - ("START", "flux2-native W4A4 fetch WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4"), - ("ERROR", "flux2-native W4A4 fetch WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4"), - ] - - -def test_audit_hf_artifact_repos_flags_native_smoke_ready_but_missing_release_metrics( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux1-schnell-native", - model_id="black-forest-labs/FLUX.1-schnell", - pipeline="FluxPipeline", - width=1024, - height=1024, - steps=4, - guidance=0.0, - bit_settings=["W4A4"], - metric="geneval", - ) - repo_id = "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "black-forest-labs/FLUX.1-schnell", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": ["block.modulation"], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index(model_index_path) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - - assert result["repo_count"] == 1 - assert result["existing_count"] == 1 - assert result["artifact_ready_count"] == 1 - assert result["native_smoke_ready_count"] == 1 - assert result["release_eval_applicable_count"] == 1 - assert result["release_eval_not_applicable_count"] == 0 - assert result["release_eval_ready_count"] == 0 - assert result["forbidden_file_count"] == 0 - assert result["missing_required_metric_count"] == len(_expected_missing_geneval_metrics()) - row = result["rows"][0] - assert row["repo_id"] == repo_id - assert row["artifact_ready"] is True - assert row["native_smoke_ready"] is True - assert row["release_eval_applicable"] is True - assert row["release_eval_ready"] is False - assert row["manifest_warnings"] == [ - "quantization_device_missing", - "weight_quantization_backend_missing", - ] - assert row["metadata_complete_ready"] is False - assert row["metadata_missing"] == [ - "manifest.activation_eps_missing", - "manifest.quantization_device_missing", - "manifest.weight_quantization_backend_missing", - "manifest.quantization_staging_mode_missing", - "model_index.activation_eps_missing", - "model_index.quantization_device_missing", - "model_index.weight_quantization_backend_missing", - "model_index.quantization_staging_mode_missing", - "benchmark_summary.activation_eps_missing", - "benchmark_summary.quantization_device_missing", - "benchmark_summary.weight_quantization_backend_missing", - "benchmark_summary.quantization_staging_mode_missing", - ] - assert row["remote_checksum_mismatches"] == [] - assert row["missing_required_metrics"] == _expected_missing_geneval_metrics() - assert row["native_smoke_proof_ready"] is True - assert row["native_smoke_missing_evidence"] == [] - - -def test_audit_hf_artifact_repos_marks_complete_metadata_when_provenance_matches( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "black-forest-labs/FLUX.2-klein-4B", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux2", - "activation_eps": 1e-10, - "quantization_device": "cuda", - "weight_quantization_backend": "triton_cuda", - "quantization_staging_mode": "component", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": [], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index( - model_index_path, - activation_eps=1e-10, - quantization_device="cuda", - weight_quantization_backend="triton_cuda", - quantization_staging_mode="component", - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text( - _native_smoke_summary( - suite, - activation_eps=1e-10, - quantization_device="cuda", - weight_quantization_backend="triton_cuda", - quantization_staging_mode="component", - ), - encoding="utf-8", - ) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - - row = result["rows"][0] - assert result["artifact_ready_count"] == 1 - assert result["metadata_complete_ready_count"] == 1 - assert result["metadata_missing_count"] == 0 - assert row["metadata_complete_ready"] is True - assert row["metadata_missing"] == [] - - -def test_audit_hf_artifact_repos_validates_policy_inventory_without_tensor_download( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - config = OrbitQuantConfig(target_policy="flux2") - quantized_modules = ["block.attn.to_q"] - adaln_modules = ["block.modulation"] - skipped_modules = ["proj_out"] - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - json.dumps( - { - "source_model_id": suite.model_id, - "source_revision": "remote-sha", - "source_license": "apache-2.0", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux2", - "quantized_modules": quantized_modules, - "adaln_modules": adaln_modules, - "skipped_modules": skipped_modules, - "checksums": { - "model.safetensors": "a" * 64, - }, - }, - indent=2, - ) - + "\n", - encoding="utf-8", - ) - quantization_config_path = tmp_path / "quantization_config.json" - quantization_config_path.write_text( - json.dumps(config.to_dict(), indent=2) + "\n", - encoding="utf-8", - ) - model_index_path = tmp_path / "model_index.json" - model_index_path.write_text( - json.dumps({"component": "transformer"}, indent=2) + "\n", - encoding="utf-8", - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text(f"{'a' * 64} model.safetensors\n", encoding="utf-8") - summary_path = tmp_path / "summary.json" - summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") - inventory_root = tmp_path / "inventories" - inventory_root.mkdir() - (inventory_root / "flux2-native-policy.json").write_text( - json.dumps( - { - "source_model_id": suite.model_id, - "target_policy": "flux2", - "component": "transformer", - "load_mode": "config", - "linear_module_count": 3, - "action_counts": { - "orbitquant": 1, - "adaln_int4_rtn": 1, - "bf16_skip": 1, - }, - "quantized_modules": quantized_modules, - "adaln_modules": adaln_modules, - "skipped_modules": skipped_modules, - }, - indent=2, - ) - + "\n", - encoding="utf-8", - ) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - quantization_config_path=quantization_config_path, - ) - downloaded_files = [] - - def fake_download(repo, filename, **kwargs): - downloaded_files.append(filename) - if filename == "model.safetensors": - raise AssertionError("metadata-only audit must not download tensors") - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos( - suites=[suite], - api=api, - policy_inventory_root=inventory_root, - ) - - row = result["rows"][0] - assert result["policy_inventory_ready_count"] == 1 - assert result["policy_inventory_error_count"] == 0 - assert result["artifact_ready_count"] == 1 - assert row["policy_inventory_ready"] is True - assert row["policy_inventory_error"] is None - assert row["policy_inventory_validation"]["inventory_path"] == str( - inventory_root / "flux2-native-policy.json" - ) - assert "quantization_config.json" in downloaded_files - assert "model.safetensors" not in downloaded_files - - -def test_audit_hf_artifact_repos_requires_native_smoke_proof_block( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux1-schnell-native", - model_id="black-forest-labs/FLUX.1-schnell", - pipeline="FluxPipeline", - width=1024, - height=1024, - steps=4, - guidance=0.0, - bit_settings=["W4A4"], - metric="geneval", - ) - repo_id = "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "black-forest-labs/FLUX.1-schnell", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": ["block.modulation"], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text( - """{ - "metrics": { - "original": {"records": 1, "latest_metrics": {"generated_samples": 1}}, - "orbitquant": {"records": 1, "latest_metrics": {"generated_samples": 1}} - } - }""" - ) - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index(model_index_path) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - - row = result["rows"][0] - assert result["artifact_ready_count"] == 1 - assert result["native_smoke_ready_count"] == 0 - assert result["release_eval_ready_count"] == 0 - assert row["artifact_ready"] is True - assert row["native_smoke_ready"] is False - assert row["native_smoke_proof_ready"] is False - assert row["native_smoke_missing_evidence"] == ["native_smoke_missing"] - - -def test_audit_hf_artifact_repos_requires_geneval_per_task_metrics( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="z-image-native", - model_id="Tongyi-MAI/Z-Image-Turbo", - pipeline="ZImagePipeline", - width=1024, - height=1024, - steps=10, - guidance=0.0, - bit_settings=["W4A4"], - metric="geneval", - ) - repo_id = "WaveCut/Z-Image-Turbo-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "Tongyi-MAI/Z-Image-Turbo", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "z_image", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": ["block.modulation"], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text( - _native_smoke_summary(suite, release_metrics={"geneval_overall": 0.71}), - encoding="utf-8", - ) - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index(model_index_path) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - - row = result["rows"][0] - assert result["artifact_ready_count"] == 1 - assert result["native_smoke_ready_count"] == 1 - assert result["release_eval_ready_count"] == 0 - assert result["missing_required_metric_count"] == 12 - assert row["release_eval_ready"] is False - assert row["missing_required_metrics"] == [ - {"split": split, "metric": metric} - for metric in hub_module._GENEVAL_REQUIRED_METRICS - if metric != "geneval_overall" - for split in ("original", "orbitquant") - ] - - -def test_audit_hf_artifact_repos_flags_extra_comparison_matrix_assets( - tmp_path, - monkeypatch, -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - "assets/video_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "black-forest-labs/FLUX.2-klein-4B", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux2", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": ["block.modulation"], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text('{"metrics": {}}\n') - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index(model_index_path) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - - assert result["artifact_ready_count"] == 0 - assert result["forbidden_file_count"] == 1 - row = result["rows"][0] - assert row["artifact_ready"] is False - assert row["asset_count"] == 2 - assert row["forbidden_files"] == ["assets/video_generation_comparison_matrix.webp"] - - -def test_audit_hf_artifact_repos_rejects_lfs_checksum_mismatch(tmp_path, monkeypatch): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings["model.safetensors"] = {"size": 123, "lfs_sha256": "b" * 64} - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "black-forest-labs/FLUX.2-klein-4B", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux2", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": [], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index(model_index_path) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - - row = result["rows"][0] - assert result["artifact_ready_count"] == 0 - assert result["remote_checksum_mismatch_count"] == 2 - assert result["release_eval_applicable_count"] == 0 - assert row["release_eval_applicable"] is False - assert row["artifact_ready"] is False - assert row["remote_checksum_mismatches"] == [ - ( - "manifest/LFS mismatch for model.safetensors: expected " - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, got " - "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" - ), - ( - "SHA256SUMS/LFS mismatch for model.safetensors: expected " - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, got " - "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" - ), - ] - - -def test_render_hf_artifact_audit_markdown_reports_checksum_mismatch_count(): - markdown = render_hf_artifact_audit_markdown( - { - "namespace": "WaveCut", - "repo_count": 1, - "existing_count": 1, - "artifact_ready_count": 0, - "native_smoke_ready_count": 0, - "release_eval_applicable_count": 1, - "release_eval_not_applicable_count": 0, - "release_eval_ready_count": 0, - "missing_required_metric_count": 0, - "manifest_warning_count": 0, - "remote_checksum_mismatch_count": 2, - "rows": [], - } - ) - - assert "- Remote checksum mismatches: 2" in markdown - - -def test_audit_hf_artifact_repos_flags_stale_remote_readme(tmp_path, monkeypatch): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "black-forest-labs/FLUX.2-klein-4B", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux2", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": [], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index(model_index_path) - readme_path = tmp_path / "README.md" - readme_path.write_text( - "# Stale card\n\nRunPod logs and reports/native/debug output.\n", - encoding="utf-8", - ) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - readme_path=readme_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - row = result["rows"][0] - - assert result["artifact_ready_count"] == 0 - assert result["readme_mismatch_count"] == 1 - assert row["artifact_ready"] is False - assert row["readme_mismatches"] == ["README.md does not match generated OrbitQuant model card"] - - -def test_audit_hf_artifact_repos_marks_visual_only_extra_target_not_applicable( - tmp_path, monkeypatch -): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - metric="visual+optional-geneval", - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "black-forest-labs/FLUX.2-klein-4B", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux2", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": [], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index(model_index_path) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - - row = result["rows"][0] - assert result["artifact_ready_count"] == 1 - assert result["native_smoke_ready_count"] == 1 - assert result["release_eval_applicable_count"] == 0 - assert result["release_eval_not_applicable_count"] == 1 - assert result["release_eval_ready_count"] == 0 - assert result["forbidden_file_count"] == 0 - assert row["release_eval_applicable"] is False - assert row["release_eval_ready"] is False - assert row["missing_required_metrics"] == [] - - -def test_audit_hf_artifact_repos_flags_forbidden_remote_assets(tmp_path, monkeypatch): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" - siblings = _required_remote_files() - siblings.update( - { - "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, - "assets/image_generation_comparison_matrix.webp": 10, - "assets/flux2-native_seed0_W4A4_simple-object.png": 10, - "assets/nested/debug_generation_comparison_matrix.webp": 10, - } - ) - api = FakeAuditHfApi({repo_id: siblings}) - manifest_path = tmp_path / "manifest.json" - manifest_path.write_text( - """{ - "source_model_id": "black-forest-labs/FLUX.2-klein-4B", - "weight_bits": 4, - "activation_bits": 4, - "target_policy": "flux2", - "quantized_modules": ["block.attn.to_q"], - "adaln_modules": [], - "checksums": { - "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - } - }""" - ) - sha256sums_path = tmp_path / "SHA256SUMS" - sha256sums_path.write_text( - "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" - ) - summary_path = tmp_path / "summary.json" - summary_path.write_text( - """{ - "metrics": { - "original": {"records": 1, "latest_metrics": {"generated_samples": 1}}, - "orbitquant": {"records": 1, "latest_metrics": {"generated_samples": 1}} - } - }""" - ) - model_index_path = tmp_path / "model_index.json" - _write_remote_model_index(model_index_path) - file_map = _audit_file_map( - repo_id, - manifest_path=manifest_path, - sha256sums_path=sha256sums_path, - summary_path=summary_path, - model_index_path=model_index_path, - ) - - def fake_download(repo, filename, **kwargs): - return str(file_map[(repo, filename)]) - - monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - row = result["rows"][0] - - assert result["forbidden_file_count"] == 2 - assert result["artifact_ready_count"] == 0 - assert row["artifact_ready"] is False - assert row["forbidden_files"] == [ - "assets/flux2-native_seed0_W4A4_simple-object.png", - "assets/nested/debug_generation_comparison_matrix.webp", - ] - - -def test_audit_hf_artifact_repos_reports_missing_repo_without_downloading(): - suite = NativeSuite( - name="flux2-native", - model_id="black-forest-labs/FLUX.2-klein-4B", - pipeline="Flux2KleinPipeline", - width=1024, - height=1024, - steps=4, - guidance=1.0, - bit_settings=["W4A4"], - ) - api = FakeAuditHfApi({}) - - result = audit_hf_artifact_repos(suites=[suite], api=api) - - assert result["repo_count"] == 1 - assert result["existing_count"] == 0 - assert result["rows"][0]["exists"] is False - assert "missing repo" in result["rows"][0]["error"] - - -def test_render_hf_artifact_audit_markdown_summarizes_ready_and_metric_gaps(): - payload = { - "namespace": "WaveCut", - "repo_count": 2, - "existing_count": 2, - "artifact_ready_count": 2, - "metadata_complete_ready_count": 1, - "native_smoke_ready_count": 2, - "release_eval_applicable_count": 1, - "release_eval_not_applicable_count": 1, - "release_eval_ready_count": 0, - "missing_required_metric_count": 2, - "manifest_warning_count": 0, - "metadata_missing_count": 3, - "rows": [ - { - "suite": "flux2-native", - "bit_setting": "W4A4", - "repo_id": "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4", - "private": True, - "artifact_ready": True, - "metadata_complete_ready": True, - "native_smoke_ready": True, - "release_eval_applicable": False, - "release_eval_ready": False, - "sha": "abcdef1234567890", - "missing_required_metrics": [], - }, - { - "suite": "flux1-schnell-native", - "bit_setting": "W4A4", - "repo_id": "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4", - "private": True, - "artifact_ready": True, - "metadata_complete_ready": False, - "native_smoke_ready": True, - "release_eval_applicable": True, - "release_eval_ready": False, - "sha": "123456abcdef7890", - "missing_required_metrics": [ - {"split": "original", "metric": "geneval_overall"}, - {"split": "orbitquant", "metric": "geneval_overall"}, - ], - }, - ], - } - - markdown = render_hf_artifact_audit_markdown(payload) - - assert "# OrbitQuant HF Artifact Audit" in markdown - assert "- Metadata complete: 1 / 2" in markdown - assert "- Policy inventory ready: not checked" in markdown - assert "- Release eval applicable: 1 / 2" in markdown - assert "- Release eval ready: 0 / 1" in markdown - assert "- Missing release metrics: 2" in markdown - assert "- Metadata missing fields: 3" in markdown - assert "## Readiness Semantics" in markdown - assert "no forbidden raw files" in markdown - assert "it is not a GenEval or VBench result" in markdown - assert "Metadata complete means activation normalization epsilon" in markdown - assert "paper metric or reproduction claims" in markdown - assert "Missing required metrics" not in markdown - assert ( - "| flux2-native | W4A4 | `WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4` | " - "yes | yes | yes | yes | n/a | n/a | abcdef123456 | | |" - ) in markdown - assert ( - "| flux1-schnell-native | W4A4 | `WaveCut/FLUX.1-schnell-OrbitQuant-W4A4` | " - "yes | yes | no | yes | n/a | no | 123456abcdef | 2 release metrics missing | |" - ) in markdown - assert ( - "`WaveCut/FLUX.1-schnell-OrbitQuant-W4A4`: " - "orbitquant:geneval_overall, original:geneval_overall" - ) in markdown diff --git a/tests/test_hub_audit.py b/tests/test_hub_audit.py new file mode 100644 index 0000000..97fb6d4 --- /dev/null +++ b/tests/test_hub_audit.py @@ -0,0 +1,934 @@ +import json + +from hub_helpers import ( + FakeAuditHfApi, + _audit_file_map, + _expected_missing_geneval_metrics, + _native_smoke_summary, + _required_remote_files, + _write_remote_model_index, + make_fake_download, +) + +import orbitquant.hub as hub_module +from orbitquant.config import OrbitQuantConfig +from orbitquant.eval.native_settings import NativeSuite +from orbitquant.hub import audit_hf_artifact_repos, render_hf_artifact_audit_markdown + + +def test_audit_hf_artifact_repos_flags_native_smoke_ready_but_missing_release_metrics( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux1-schnell-native", + model_id="black-forest-labs/FLUX.1-schnell", + pipeline="FluxPipeline", + width=1024, + height=1024, + steps=4, + guidance=0.0, + bit_settings=["W4A4"], + metric="geneval", + ) + repo_id = "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "black-forest-labs/FLUX.1-schnell", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": ["block.modulation"], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index(model_index_path) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + + assert result["repo_count"] == 1 + assert result["existing_count"] == 1 + assert result["artifact_ready_count"] == 1 + assert result["native_smoke_ready_count"] == 1 + assert result["release_eval_applicable_count"] == 1 + assert result["release_eval_not_applicable_count"] == 0 + assert result["release_eval_ready_count"] == 0 + assert result["forbidden_file_count"] == 0 + assert result["missing_required_metric_count"] == len(_expected_missing_geneval_metrics()) + row = result["rows"][0] + assert row["repo_id"] == repo_id + assert row["artifact_ready"] is True + assert row["native_smoke_ready"] is True + assert row["release_eval_applicable"] is True + assert row["release_eval_ready"] is False + assert row["manifest_warnings"] == [ + "quantization_device_missing", + "weight_quantization_backend_missing", + ] + assert row["metadata_complete_ready"] is False + assert row["metadata_missing"] == [ + "manifest.activation_eps_missing", + "manifest.quantization_device_missing", + "manifest.weight_quantization_backend_missing", + "manifest.quantization_staging_mode_missing", + "model_index.activation_eps_missing", + "model_index.quantization_device_missing", + "model_index.weight_quantization_backend_missing", + "model_index.quantization_staging_mode_missing", + "benchmark_summary.activation_eps_missing", + "benchmark_summary.quantization_device_missing", + "benchmark_summary.weight_quantization_backend_missing", + "benchmark_summary.quantization_staging_mode_missing", + ] + assert row["remote_checksum_mismatches"] == [] + assert row["missing_required_metrics"] == _expected_missing_geneval_metrics() + assert row["native_smoke_proof_ready"] is True + assert row["native_smoke_missing_evidence"] == [] + + +def test_audit_hf_artifact_repos_marks_complete_metadata_when_provenance_matches( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "black-forest-labs/FLUX.2-klein-4B", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux2", + "activation_eps": 1e-10, + "quantization_device": "cuda", + "weight_quantization_backend": "triton_cuda", + "quantization_staging_mode": "component", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": [], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index( + model_index_path, + activation_eps=1e-10, + quantization_device="cuda", + weight_quantization_backend="triton_cuda", + quantization_staging_mode="component", + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text( + _native_smoke_summary( + suite, + activation_eps=1e-10, + quantization_device="cuda", + weight_quantization_backend="triton_cuda", + quantization_staging_mode="component", + ), + encoding="utf-8", + ) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + + row = result["rows"][0] + assert result["artifact_ready_count"] == 1 + assert result["metadata_complete_ready_count"] == 1 + assert result["metadata_missing_count"] == 0 + assert row["metadata_complete_ready"] is True + assert row["metadata_missing"] == [] + + +def test_audit_hf_artifact_repos_validates_policy_inventory_without_tensor_download( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + config = OrbitQuantConfig(target_policy="flux2") + quantized_modules = ["block.attn.to_q"] + adaln_modules = ["block.modulation"] + skipped_modules = ["proj_out"] + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + json.dumps( + { + "source_model_id": suite.model_id, + "source_revision": "remote-sha", + "source_license": "apache-2.0", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux2", + "quantized_modules": quantized_modules, + "adaln_modules": adaln_modules, + "skipped_modules": skipped_modules, + "checksums": { + "model.safetensors": "a" * 64, + }, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + quantization_config_path = tmp_path / "quantization_config.json" + quantization_config_path.write_text( + json.dumps(config.to_dict(), indent=2) + "\n", + encoding="utf-8", + ) + model_index_path = tmp_path / "model_index.json" + model_index_path.write_text( + json.dumps({"component": "transformer"}, indent=2) + "\n", + encoding="utf-8", + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text(f"{'a' * 64} model.safetensors\n", encoding="utf-8") + summary_path = tmp_path / "summary.json" + summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") + inventory_root = tmp_path / "inventories" + inventory_root.mkdir() + (inventory_root / "flux2-native-policy.json").write_text( + json.dumps( + { + "source_model_id": suite.model_id, + "target_policy": "flux2", + "component": "transformer", + "load_mode": "config", + "linear_module_count": 3, + "action_counts": { + "orbitquant": 1, + "adaln_int4_rtn": 1, + "bf16_skip": 1, + }, + "quantized_modules": quantized_modules, + "adaln_modules": adaln_modules, + "skipped_modules": skipped_modules, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + quantization_config_path=quantization_config_path, + ) + downloaded_files = [] + + def fake_download(repo, filename, **kwargs): + downloaded_files.append(filename) + if filename == "model.safetensors": + raise AssertionError("metadata-only audit must not download tensors") + return str(file_map[(repo, filename)]) + + monkeypatch.setattr(hub_module, "hf_hub_download", fake_download) + + result = audit_hf_artifact_repos( + suites=[suite], + api=api, + policy_inventory_root=inventory_root, + ) + + row = result["rows"][0] + assert result["policy_inventory_ready_count"] == 1 + assert result["policy_inventory_error_count"] == 0 + assert result["artifact_ready_count"] == 1 + assert row["policy_inventory_ready"] is True + assert row["policy_inventory_error"] is None + assert row["policy_inventory_validation"]["inventory_path"] == str( + inventory_root / "flux2-native-policy.json" + ) + assert "quantization_config.json" in downloaded_files + assert "model.safetensors" not in downloaded_files + + +def test_audit_hf_artifact_repos_requires_native_smoke_proof_block( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux1-schnell-native", + model_id="black-forest-labs/FLUX.1-schnell", + pipeline="FluxPipeline", + width=1024, + height=1024, + steps=4, + guidance=0.0, + bit_settings=["W4A4"], + metric="geneval", + ) + repo_id = "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "black-forest-labs/FLUX.1-schnell", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": ["block.modulation"], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text( + """{ + "metrics": { + "original": {"records": 1, "latest_metrics": {"generated_samples": 1}}, + "orbitquant": {"records": 1, "latest_metrics": {"generated_samples": 1}} + } + }""" + ) + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index(model_index_path) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + + row = result["rows"][0] + assert result["artifact_ready_count"] == 1 + assert result["native_smoke_ready_count"] == 0 + assert result["release_eval_ready_count"] == 0 + assert row["artifact_ready"] is True + assert row["native_smoke_ready"] is False + assert row["native_smoke_proof_ready"] is False + assert row["native_smoke_missing_evidence"] == ["native_smoke_missing"] + + +def test_audit_hf_artifact_repos_requires_geneval_per_task_metrics( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="z-image-native", + model_id="Tongyi-MAI/Z-Image-Turbo", + pipeline="ZImagePipeline", + width=1024, + height=1024, + steps=10, + guidance=0.0, + bit_settings=["W4A4"], + metric="geneval", + ) + repo_id = "WaveCut/Z-Image-Turbo-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "Tongyi-MAI/Z-Image-Turbo", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "z_image", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": ["block.modulation"], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text( + _native_smoke_summary(suite, release_metrics={"geneval_overall": 0.71}), + encoding="utf-8", + ) + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index(model_index_path) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + + row = result["rows"][0] + assert result["artifact_ready_count"] == 1 + assert result["native_smoke_ready_count"] == 1 + assert result["release_eval_ready_count"] == 0 + assert result["missing_required_metric_count"] == 12 + assert row["release_eval_ready"] is False + assert row["missing_required_metrics"] == [ + {"split": split, "metric": metric} + for metric in hub_module._GENEVAL_REQUIRED_METRICS + if metric != "geneval_overall" + for split in ("original", "orbitquant") + ] + + +def test_audit_hf_artifact_repos_flags_extra_comparison_matrix_assets( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + "assets/video_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "black-forest-labs/FLUX.2-klein-4B", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux2", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": ["block.modulation"], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text('{"metrics": {}}\n') + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index(model_index_path) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + + assert result["artifact_ready_count"] == 0 + assert result["forbidden_file_count"] == 1 + row = result["rows"][0] + assert row["artifact_ready"] is False + assert row["asset_count"] == 2 + assert row["forbidden_files"] == ["assets/video_generation_comparison_matrix.webp"] + + +def test_audit_hf_artifact_repos_rejects_lfs_checksum_mismatch(tmp_path, monkeypatch): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings["model.safetensors"] = {"size": 123, "lfs_sha256": "b" * 64} + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "black-forest-labs/FLUX.2-klein-4B", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux2", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": [], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index(model_index_path) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + + row = result["rows"][0] + assert result["artifact_ready_count"] == 0 + assert result["remote_checksum_mismatch_count"] == 2 + assert result["release_eval_applicable_count"] == 0 + assert row["release_eval_applicable"] is False + assert row["artifact_ready"] is False + assert row["remote_checksum_mismatches"] == [ + ( + "manifest/LFS mismatch for model.safetensors: expected " + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, got " + "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + ), + ( + "SHA256SUMS/LFS mismatch for model.safetensors: expected " + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, got " + "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + ), + ] + + +def test_render_hf_artifact_audit_markdown_reports_checksum_mismatch_count(): + markdown = render_hf_artifact_audit_markdown( + { + "namespace": "WaveCut", + "repo_count": 1, + "existing_count": 1, + "artifact_ready_count": 0, + "native_smoke_ready_count": 0, + "release_eval_applicable_count": 1, + "release_eval_not_applicable_count": 0, + "release_eval_ready_count": 0, + "missing_required_metric_count": 0, + "manifest_warning_count": 0, + "remote_checksum_mismatch_count": 2, + "rows": [], + } + ) + + assert "- Remote checksum mismatches: 2" in markdown + + +def test_audit_hf_artifact_repos_flags_stale_remote_readme(tmp_path, monkeypatch): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "black-forest-labs/FLUX.2-klein-4B", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux2", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": [], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index(model_index_path) + readme_path = tmp_path / "README.md" + readme_path.write_text( + "# Stale card\n\nRunPod logs and reports/native/debug output.\n", + encoding="utf-8", + ) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + readme_path=readme_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + row = result["rows"][0] + + assert result["artifact_ready_count"] == 0 + assert result["readme_mismatch_count"] == 1 + assert row["artifact_ready"] is False + assert row["readme_mismatches"] == ["README.md does not match generated OrbitQuant model card"] + + +def test_audit_hf_artifact_repos_marks_visual_only_extra_target_not_applicable( + tmp_path, monkeypatch +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + metric="visual+optional-geneval", + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "black-forest-labs/FLUX.2-klein-4B", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux2", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": [], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index(model_index_path) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + + row = result["rows"][0] + assert result["artifact_ready_count"] == 1 + assert result["native_smoke_ready_count"] == 1 + assert result["release_eval_applicable_count"] == 0 + assert result["release_eval_not_applicable_count"] == 1 + assert result["release_eval_ready_count"] == 0 + assert result["forbidden_file_count"] == 0 + assert row["release_eval_applicable"] is False + assert row["release_eval_ready"] is False + assert row["missing_required_metrics"] == [] + + +def test_audit_hf_artifact_repos_flags_forbidden_remote_assets(tmp_path, monkeypatch): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + siblings = _required_remote_files() + siblings.update( + { + "model.safetensors": {"size": 123, "lfs_sha256": "a" * 64}, + "assets/image_generation_comparison_matrix.webp": 10, + "assets/flux2-native_seed0_W4A4_simple-object.png": 10, + "assets/nested/debug_generation_comparison_matrix.webp": 10, + } + ) + api = FakeAuditHfApi({repo_id: siblings}) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + """{ + "source_model_id": "black-forest-labs/FLUX.2-klein-4B", + "weight_bits": 4, + "activation_bits": 4, + "target_policy": "flux2", + "quantized_modules": ["block.attn.to_q"], + "adaln_modules": [], + "checksums": { + "model.safetensors": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }""" + ) + sha256sums_path = tmp_path / "SHA256SUMS" + sha256sums_path.write_text( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa model.safetensors\n" + ) + summary_path = tmp_path / "summary.json" + summary_path.write_text( + """{ + "metrics": { + "original": {"records": 1, "latest_metrics": {"generated_samples": 1}}, + "orbitquant": {"records": 1, "latest_metrics": {"generated_samples": 1}} + } + }""" + ) + model_index_path = tmp_path / "model_index.json" + _write_remote_model_index(model_index_path) + file_map = _audit_file_map( + repo_id, + manifest_path=manifest_path, + sha256sums_path=sha256sums_path, + summary_path=summary_path, + model_index_path=model_index_path, + ) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(file_map)) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + row = result["rows"][0] + + assert result["forbidden_file_count"] == 2 + assert result["artifact_ready_count"] == 0 + assert row["artifact_ready"] is False + assert row["forbidden_files"] == [ + "assets/flux2-native_seed0_W4A4_simple-object.png", + "assets/nested/debug_generation_comparison_matrix.webp", + ] + + +def test_audit_hf_artifact_repos_reports_missing_repo_without_downloading(): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + api = FakeAuditHfApi({}) + + result = audit_hf_artifact_repos(suites=[suite], api=api) + + assert result["repo_count"] == 1 + assert result["existing_count"] == 0 + assert result["rows"][0]["exists"] is False + assert "missing repo" in result["rows"][0]["error"] + + +def test_render_hf_artifact_audit_markdown_summarizes_ready_and_metric_gaps(): + payload = { + "namespace": "WaveCut", + "repo_count": 2, + "existing_count": 2, + "artifact_ready_count": 2, + "metadata_complete_ready_count": 1, + "native_smoke_ready_count": 2, + "release_eval_applicable_count": 1, + "release_eval_not_applicable_count": 1, + "release_eval_ready_count": 0, + "missing_required_metric_count": 2, + "manifest_warning_count": 0, + "metadata_missing_count": 3, + "rows": [ + { + "suite": "flux2-native", + "bit_setting": "W4A4", + "repo_id": "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4", + "private": True, + "artifact_ready": True, + "metadata_complete_ready": True, + "native_smoke_ready": True, + "release_eval_applicable": False, + "release_eval_ready": False, + "sha": "abcdef1234567890", + "missing_required_metrics": [], + }, + { + "suite": "flux1-schnell-native", + "bit_setting": "W4A4", + "repo_id": "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4", + "private": True, + "artifact_ready": True, + "metadata_complete_ready": False, + "native_smoke_ready": True, + "release_eval_applicable": True, + "release_eval_ready": False, + "sha": "123456abcdef7890", + "missing_required_metrics": [ + {"split": "original", "metric": "geneval_overall"}, + {"split": "orbitquant", "metric": "geneval_overall"}, + ], + }, + ], + } + + markdown = render_hf_artifact_audit_markdown(payload) + + assert "# OrbitQuant HF Artifact Audit" in markdown + assert "- Metadata complete: 1 / 2" in markdown + assert "- Policy inventory ready: not checked" in markdown + assert "- Release eval applicable: 1 / 2" in markdown + assert "- Release eval ready: 0 / 1" in markdown + assert "- Missing release metrics: 2" in markdown + assert "- Metadata missing fields: 3" in markdown + assert "## Readiness Semantics" in markdown + assert "no forbidden raw files" in markdown + assert "it is not a GenEval or VBench result" in markdown + assert "Metadata complete means activation normalization epsilon" in markdown + assert "paper metric or reproduction claims" in markdown + assert "Missing required metrics" not in markdown + assert ( + "| flux2-native | W4A4 | `WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4` | " + "yes | yes | yes | yes | n/a | n/a | abcdef123456 | | |" + ) in markdown + assert ( + "| flux1-schnell-native | W4A4 | `WaveCut/FLUX.1-schnell-OrbitQuant-W4A4` | " + "yes | yes | no | yes | n/a | no | 123456abcdef | 2 release metrics missing | |" + ) in markdown + assert ( + "`WaveCut/FLUX.1-schnell-OrbitQuant-W4A4`: " + "orbitquant:geneval_overall, original:geneval_overall" + ) in markdown diff --git a/tests/test_hub_fetch.py b/tests/test_hub_fetch.py new file mode 100644 index 0000000..4659ac0 --- /dev/null +++ b/tests/test_hub_fetch.py @@ -0,0 +1,157 @@ +import shutil +from pathlib import Path + +import pytest +from hub_helpers import _write_artifact + +import orbitquant.hub as hub_module +from orbitquant.eval.native_settings import NativeSuite +from orbitquant.hub import fetch_hf_artifacts + + +def test_fetch_hf_artifacts_dry_run_reports_native_artifact_layout(tmp_path): + suite = NativeSuite( + name="flux1-schnell-native", + model_id="black-forest-labs/FLUX.1-schnell", + pipeline="FluxPipeline", + width=1024, + height=1024, + steps=4, + guidance=0.0, + bit_settings=["W4A4", "W3A3"], + metric="geneval", + ) + + result = fetch_hf_artifacts( + namespace="WaveCut", + suites=[suite], + output_root=tmp_path / "artifacts", + dry_run=True, + ) + + assert result["repo_count"] == 2 + assert result["downloaded_count"] == 0 + assert result["dry_run"] is True + assert result["rows"][0]["repo_id"] == "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4" + assert result["rows"][0]["artifact_dir"].endswith("flux1-schnell-native-w4a4") + assert result["rows"][1]["repo_id"] == "WaveCut/FLUX.1-schnell-OrbitQuant-W3A3" + assert result["rows"][1]["artifact_dir"].endswith("flux1-schnell-native-w3a3") + + +def test_fetch_hf_artifacts_downloads_and_validates_artifact( + tmp_path, + monkeypatch, +): + source_dir = tmp_path / "source" + _write_artifact(source_dir) + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + calls = [] + stages = [] + + def fake_snapshot_download(**kwargs): + calls.append(kwargs) + destination = Path(kwargs["local_dir"]) + shutil.copytree(source_dir, destination, dirs_exist_ok=True) + return str(destination) + + monkeypatch.setattr(hub_module, "snapshot_download", fake_snapshot_download) + + result = fetch_hf_artifacts( + namespace="WaveCut", + suites=[suite], + output_root=tmp_path / "artifacts", + revision="main", + stage_logger=lambda event, label: stages.append((event, label)), + ) + + artifact_dir = tmp_path / "artifacts" / "flux2-native-w4a4" + assert result["downloaded_count"] == 1 + assert result["skipped_existing_count"] == 0 + assert result["rows"][0]["artifact_dir"] == str(artifact_dir) + assert result["rows"][0]["validation"]["valid"] is True + assert calls == [ + { + "repo_id": "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4", + "repo_type": "model", + "revision": "main", + "local_dir": artifact_dir, + "force_download": False, + "local_files_only": False, + } + ] + assert stages == [ + ("START", "flux2-native W4A4 fetch WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4"), + ("END", "flux2-native W4A4 fetch WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4"), + ] + + +def test_fetch_hf_artifacts_resume_skips_valid_existing_artifact(tmp_path, monkeypatch): + artifact_dir = tmp_path / "artifacts" / "flux2-native-w4a4" + _write_artifact(artifact_dir) + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + + def fail_snapshot_download(**kwargs): + raise AssertionError("resume should not download a valid existing artifact") + + monkeypatch.setattr(hub_module, "snapshot_download", fail_snapshot_download) + + result = fetch_hf_artifacts( + namespace="WaveCut", + suites=[suite], + output_root=tmp_path / "artifacts", + ) + + assert result["downloaded_count"] == 0 + assert result["skipped_existing_count"] == 1 + assert result["rows"][0]["skipped_existing"] is True + assert result["rows"][0]["validation"]["valid"] is True + + +def test_fetch_hf_artifacts_logs_error_stage_on_download_failure(tmp_path, monkeypatch): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + stages = [] + + def fail_snapshot_download(**kwargs): + raise RuntimeError("download failed") + + monkeypatch.setattr(hub_module, "snapshot_download", fail_snapshot_download) + + with pytest.raises(RuntimeError, match="download failed"): + fetch_hf_artifacts( + namespace="WaveCut", + suites=[suite], + output_root=tmp_path / "artifacts", + stage_logger=lambda event, label: stages.append((event, label)), + ) + + assert stages == [ + ("START", "flux2-native W4A4 fetch WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4"), + ("ERROR", "flux2-native W4A4 fetch WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4"), + ] diff --git a/tests/test_hub_native_smoke_proof.py b/tests/test_hub_native_smoke_proof.py new file mode 100644 index 0000000..46d3b7b --- /dev/null +++ b/tests/test_hub_native_smoke_proof.py @@ -0,0 +1,157 @@ +import json + +from hub_helpers import _legacy_compact_summary_without_native_smoke, _native_smoke_summary + +import orbitquant.hub as hub_module +from orbitquant.eval.native_settings import NativeSuite + + +def test_native_smoke_expected_settings_include_video_export_fps_only_when_defined(): + image_suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + video_suite = NativeSuite( + name="wan-native", + model_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + pipeline="WanPipeline", + width=832, + height=480, + frames=81, + export_fps=16, + steps=50, + guidance=5.0, + bit_settings=["W4A4"], + ) + + assert "export_fps" not in hub_module._native_smoke_expected_settings(image_suite) + assert hub_module._native_smoke_expected_settings(video_suite)["export_fps"] == 16 + + +def test_recover_native_smoke_proof_from_compact_summary_requires_raw_pair_evidence(): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + summary = json.loads(_legacy_compact_summary_without_native_smoke(suite)) + file_names = {"assets/image_generation_comparison_matrix.webp"} + + proof, reason = hub_module._recover_native_smoke_proof_from_compact_summary( + summary, + suite=suite, + file_names=file_names, + ) + + assert proof is None + assert reason == "raw_paired_native_smoke_evidence_missing" + + +def test_native_smoke_proof_status_rejects_recovered_pair_claims(): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + summary = json.loads(_native_smoke_summary(suite)) + summary["native_smoke"]["proof_source"] = ( + "recovered_from_compact_summary_and_published_comparison_matrix" + ) + + status = hub_module._native_smoke_proof_status( + summary, + suite=suite, + file_names={"assets/image_generation_comparison_matrix.webp"}, + ) + + assert status["ready"] is False + assert "native_smoke.raw_paired_native_smoke_evidence" in status["missing"] + + +def test_recover_native_smoke_proof_requires_uploaded_comparison_asset(): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + summary = json.loads(_legacy_compact_summary_without_native_smoke(suite)) + summary.pop("raw_generation_records") + + proof, reason = hub_module._recover_native_smoke_proof_from_compact_summary( + summary, + suite=suite, + file_names=set(), + ) + + assert proof is None + assert reason == "comparison_asset_missing" + + +def test_recover_native_smoke_proof_rejects_missing_generated_samples(): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + summary = json.loads(_legacy_compact_summary_without_native_smoke(suite, generated_samples=0)) + summary.pop("raw_generation_records") + + proof, reason = hub_module._recover_native_smoke_proof_from_compact_summary( + summary, + suite=suite, + file_names={"assets/image_generation_comparison_matrix.webp"}, + ) + + assert proof is None + assert reason == "original.generated_samples_missing" + + +def test_recover_native_smoke_proof_rejects_video_without_generated_frames(): + suite = NativeSuite( + name="wan-native", + model_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + pipeline="WanPipeline", + width=832, + height=480, + frames=81, + steps=50, + guidance=5.0, + bit_settings=["W4A4"], + ) + summary = json.loads(_legacy_compact_summary_without_native_smoke(suite)) + summary.pop("raw_generation_records") + + proof, reason = hub_module._recover_native_smoke_proof_from_compact_summary( + summary, + suite=suite, + file_names={"assets/video_generation_comparison_matrix.webp"}, + ) + + assert proof is None + assert reason == "original.generated_frames_insufficient" diff --git a/tests/test_hub_repair.py b/tests/test_hub_repair.py new file mode 100644 index 0000000..e49b1ce --- /dev/null +++ b/tests/test_hub_repair.py @@ -0,0 +1,802 @@ +import json + +from hub_helpers import ( + FakeCleanupHfApi, + FakeCommitHfApi, + _legacy_compact_summary_without_native_smoke, + _native_smoke_summary, + _remote_file_map, + _write_artifact, + _write_native_smoke_backup, + make_fake_download, +) +from huggingface_hub import CommitOperationAdd, CommitOperationDelete + +import orbitquant.hub as hub_module +from orbitquant.artifacts import refresh_artifact_checksums +from orbitquant.artifacts.checksums import read_sha256sums, write_sha256sums +from orbitquant.eval.native_settings import NativeSuite +from orbitquant.hub import ( + cleanup_hf_artifact_reports, + cleanup_hf_artifact_reports_matrix, + repair_hf_artifact_metadata, + repair_hf_artifact_metadata_matrix, + repair_hf_native_smoke_proof, + repair_hf_native_smoke_proof_matrix, +) + + +def test_repair_hf_artifact_metadata_dry_run_preserves_large_file_checksum( + tmp_path, + monkeypatch, +): + repo_id = "WaveCut/example-orbitquant" + _write_artifact(tmp_path) + old_checksums = read_sha256sums(tmp_path / "SHA256SUMS") + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCommitHfApi() + + result = repair_hf_artifact_metadata( + repo_id=repo_id, + quantization_device="cuda", + weight_quantization_backend="triton_cuda", + quantization_staging_mode="component", + dry_run=True, + api=fake_api, + ) + + assert result["dry_run"] is True + assert result["updated"]["quantization_staging_mode"] == "component" + assert result["commit"] is None + assert "model.safetensors" in result["preserved_checksum_entries"] + assert old_checksums["model.safetensors"] + assert fake_api.create_commit_calls == [] + + +def test_repair_hf_artifact_metadata_commits_only_metadata_files_and_sha256sums( + tmp_path, + monkeypatch, +): + repo_id = "WaveCut/example-orbitquant" + _write_artifact(tmp_path) + stale_cache_metadata = tmp_path / ".cache" / "huggingface" / "download" / "README.md.metadata" + stale_cache_metadata.parent.mkdir(parents=True) + stale_cache_metadata.write_text("transient hub metadata", encoding="utf-8") + with (tmp_path / "SHA256SUMS").open("a", encoding="utf-8") as handle: + handle.write("0" * 64 + " .cache/huggingface/download/README.md.metadata\n") + old_checksums = read_sha256sums(tmp_path / "SHA256SUMS") + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCommitHfApi() + + result = repair_hf_artifact_metadata( + repo_id=repo_id, + quantization_device="cuda", + weight_quantization_backend="triton_cuda", + quantization_staging_mode="component", + revision="main", + commit_message="repair metadata", + api=fake_api, + ) + + assert result["commit"]["commit_oid"] == "repair-sha" + assert len(fake_api.create_commit_calls) == 1 + commit_call = fake_api.create_commit_calls[0] + assert commit_call["repo_id"] == repo_id + assert commit_call["repo_type"] == "model" + assert commit_call["revision"] == "main" + assert commit_call["commit_message"] == "repair metadata" + operation_paths = {operation.path_in_repo for operation in commit_call["operations"]} + assert operation_paths == { + "orbitquant_manifest.json", + "model_index.json", + "benchmark/summary.json", + "README.md", + "SHA256SUMS", + } + assert "model.safetensors" not in operation_paths + sha_operation = next( + operation + for operation in commit_call["operations"] + if operation.path_in_repo == "SHA256SUMS" + ) + sha_entries = hub_module._parse_sha256sums_bytes(sha_operation.path_or_fileobj) + assert sha_entries["model.safetensors"] == old_checksums["model.safetensors"] + assert sha_entries["orbitquant_manifest.json"] != old_checksums["orbitquant_manifest.json"] + assert ".cache/huggingface/download/README.md.metadata" not in sha_entries + manifest_operation = next( + operation + for operation in commit_call["operations"] + if operation.path_in_repo == "orbitquant_manifest.json" + ) + assert b'"quantization_staging_mode": "component"' in manifest_operation.path_or_fileobj + assert b'"activation_eps": 1e-10' in manifest_operation.path_or_fileobj + + +def test_cleanup_hf_artifact_reports_promotes_matrices_and_deletes_report_folder( + tmp_path, + monkeypatch, +): + repo_id = "WaveCut/example-orbitquant" + _write_artifact(tmp_path) + raw_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" + raw_sidecar = raw_asset.with_suffix(".png.json") + raw_asset.write_bytes(b"raw generated image") + raw_sidecar.write_text('{"prompt_id":"simple-object"}\n', encoding="utf-8") + report_matrix = ( + tmp_path + / "reports" + / "native" + / "flux2-w4a4" + / "assets" + / "image_generation_comparison_matrix.webp" + ) + report_matrix.parent.mkdir(parents=True) + report_matrix.write_bytes(b"report matrix") + report_markdown = report_matrix.parent.parent / "orbitquant-native-eval.md" + report_markdown.write_text("# local report log\n", encoding="utf-8") + refresh_artifact_checksums(tmp_path) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(artifact_dir=tmp_path)) + fake_api = FakeCleanupHfApi(tmp_path) + + result = cleanup_hf_artifact_reports( + repo_id=repo_id, + commit_message="cleanup reports", + api=fake_api, + ) + + assert result["report_file_count"] == 2 + assert result["report_matrix_count"] == 1 + assert result["forbidden_file_count"] == 4 + assert result["promoted_assets"] == ["assets/image_generation_comparison_matrix.webp"] + assert result["delete_paths"] == [ + "assets/flux2-native_seed0_W4A4_simple-object.png", + "assets/flux2-native_seed0_W4A4_simple-object.png.json", + "reports", + ] + assert result["commit"]["commit_oid"] == "repair-sha" + assert len(fake_api.create_commit_calls) == 1 + + commit_call = fake_api.create_commit_calls[0] + assert commit_call["commit_message"] == "cleanup reports" + operations = commit_call["operations"] + added = { + operation.path_in_repo: operation.path_or_fileobj + for operation in operations + if isinstance(operation, CommitOperationAdd) + } + deleted = [ + operation for operation in operations if isinstance(operation, CommitOperationDelete) + ] + assert "model.safetensors" not in added + assert added["assets/image_generation_comparison_matrix.webp"] == b"report matrix" + assert len(deleted) == 3 + deleted_paths = {operation.path_in_repo: operation.is_folder for operation in deleted} + assert deleted_paths == { + "assets/flux2-native_seed0_W4A4_simple-object.png": False, + "assets/flux2-native_seed0_W4A4_simple-object.png.json": False, + "reports": True, + } + + next_manifest = json.loads(added["orbitquant_manifest.json"].decode("utf-8")) + next_summary = json.loads(added["benchmark/summary.json"].decode("utf-8")) + next_readme = added["README.md"].decode("utf-8") + next_sha = added["SHA256SUMS"].decode("utf-8") + assert "assets/image_generation_comparison_matrix.webp" in next_manifest["checksums"] + assert "benchmark/summary.json" in next_manifest["checksums"] + assert not any(path.endswith(".metrics.jsonl") for path in next_manifest["checksums"]) + assert not any(path.endswith(".metrics.csv") for path in next_manifest["checksums"]) + assert not any(path.startswith("reports/") for path in next_manifest["checksums"]) + assert next_summary["published_summary"] == "compact" + assert next_summary["raw_generation_records"] == "local-only" + assert "assets/image_generation_comparison_matrix.webp" in next_readme + assert "## Validation Status" in next_readme + assert "reports/native" not in next_readme + assert "assets/image_generation_comparison_matrix.webp" in next_sha + assert "benchmark/original.metrics.jsonl" not in next_sha + assert "reports/native" not in next_sha + + +def test_cleanup_hf_artifact_reports_deletes_extra_comparison_matrix_assets( + tmp_path, + monkeypatch, +): + repo_id = "WaveCut/example-orbitquant" + _write_artifact(tmp_path) + image_matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + video_matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" + image_matrix.write_bytes(b"image matrix") + video_matrix.write_bytes(b"video matrix") + refresh_artifact_checksums(tmp_path) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(artifact_dir=tmp_path)) + fake_api = FakeCleanupHfApi(tmp_path) + + result = cleanup_hf_artifact_reports( + repo_id=repo_id, + commit_message="cleanup extra matrix", + api=fake_api, + ) + + assert result["promoted_assets"] == [] + assert result["forbidden_files"] == [ + "assets/video_generation_comparison_matrix.webp", + ] + assert "assets/video_generation_comparison_matrix.webp" in result["delete_paths"] + commit_call = fake_api.create_commit_calls[0] + operations = commit_call["operations"] + added = { + operation.path_in_repo: operation.path_or_fileobj + for operation in operations + if isinstance(operation, CommitOperationAdd) + } + deleted_paths = { + operation.path_in_repo + for operation in operations + if isinstance(operation, CommitOperationDelete) + } + assert "assets/video_generation_comparison_matrix.webp" in deleted_paths + next_manifest = json.loads(added["orbitquant_manifest.json"].decode("utf-8")) + next_readme = added.get("README.md", (tmp_path / "README.md").read_bytes()).decode("utf-8") + assert "assets/image_generation_comparison_matrix.webp" in next_manifest["checksums"] + assert "assets/video_generation_comparison_matrix.webp" not in next_manifest["checksums"] + assert "assets/image_generation_comparison_matrix.webp" in next_readme + assert "assets/video_generation_comparison_matrix.webp" not in next_readme + + +def test_cleanup_hf_artifact_reports_matrix_cleans_expected_suite_repo( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + _write_artifact(tmp_path) + report_matrix = ( + tmp_path + / "reports" + / "native" + / "flux2-w4a4" + / "assets" + / "image_generation_comparison_matrix.webp" + ) + report_matrix.parent.mkdir(parents=True) + report_matrix.write_bytes(b"report matrix") + refresh_artifact_checksums(tmp_path) + + monkeypatch.setattr(hub_module, "hf_hub_download", make_fake_download(artifact_dir=tmp_path)) + fake_api = FakeCleanupHfApi(tmp_path) + + result = cleanup_hf_artifact_reports_matrix( + suites=[suite], + dry_run=True, + api=fake_api, + ) + + assert result["repo_count"] == 1 + assert result["changed_repo_count"] == 1 + assert result["error_count"] == 0 + assert result["report_file_count"] == 1 + assert result["promoted_asset_count"] == 1 + assert result["rows"][0]["repo_id"] == repo_id + assert result["rows"][0]["suite"] == "flux2-native" + assert result["rows"][0]["bit_setting"] == "W4A4" + + +def test_repair_hf_artifact_metadata_matrix_repairs_expected_suite_repo( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + _write_artifact(tmp_path) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCommitHfApi() + + result = repair_hf_artifact_metadata_matrix( + suites=[suite], + quantization_device="cuda", + weight_quantization_backend="triton_cuda", + quantization_staging_mode="component", + dry_run=True, + api=fake_api, + ) + + assert result["repo_count"] == 1 + assert result["error_count"] == 0 + assert result["rows"][0]["repo_id"] == repo_id + assert result["rows"][0]["suite"] == "flux2-native" + assert result["rows"][0]["bit_setting"] == "W4A4" + + +def test_repair_hf_native_smoke_proof_skips_compact_summary_without_raw_pair_evidence( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + matrix.parent.mkdir(parents=True, exist_ok=True) + matrix.write_bytes(b"matrix") + summary_path = tmp_path / "benchmark" / "summary.json" + summary_path.write_text( + _legacy_compact_summary_without_native_smoke(suite), + encoding="utf-8", + ) + refresh_artifact_checksums(tmp_path) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCleanupHfApi(tmp_path) + + result = repair_hf_native_smoke_proof( + repo_id=repo_id, + suite=suite, + revision="main", + commit_message="repair native smoke proof", + api=fake_api, + ) + + assert result["commit"] is None + assert result["repair_skipped_reason"] == "raw_paired_native_smoke_evidence_missing" + assert result["changed_files"] == [] + assert fake_api.create_commit_calls == [] + + +def test_repair_hf_native_smoke_proof_removes_recovered_pair_claim( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + matrix.parent.mkdir(parents=True, exist_ok=True) + matrix.write_bytes(b"matrix") + summary_path = tmp_path / "benchmark" / "summary.json" + summary = json.loads(_native_smoke_summary(suite)) + summary["native_smoke"]["proof_source"] = ( + "recovered_from_compact_summary_and_published_comparison_matrix" + ) + summary_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + refresh_artifact_checksums(tmp_path) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCleanupHfApi(tmp_path) + + result = repair_hf_native_smoke_proof( + repo_id=repo_id, + suite=suite, + revision="main", + api=fake_api, + ) + + assert result["commit"]["commit_oid"] == "repair-sha" + assert result["removed_invalid_native_smoke"] is True + assert result["changed_files"] == [ + "benchmark/summary.json", + "orbitquant_manifest.json", + "README.md", + "SHA256SUMS", + ] + operation_by_path = { + operation.path_in_repo: operation.path_or_fileobj + for operation in fake_api.create_commit_calls[0]["operations"] + } + repaired_summary = json.loads(operation_by_path["benchmark/summary.json"]) + repaired_readme = operation_by_path["README.md"].decode("utf-8") + repaired_manifest = json.loads(operation_by_path["orbitquant_manifest.json"]) + sha_entries = hub_module._parse_sha256sums_bytes(operation_by_path["SHA256SUMS"]) + + assert "native_smoke" not in repaired_summary + assert "## Native Validation Proof" not in repaired_readme + assert ( + repaired_manifest["checksums"]["benchmark/summary.json"] + == sha_entries["benchmark/summary.json"] + ) + assert sha_entries["README.md"] == hub_module._sha256_bytes(operation_by_path["README.md"]) + + +def test_repair_hf_native_smoke_proof_restores_from_local_backup_without_raw_upload( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + matrix.parent.mkdir(parents=True, exist_ok=True) + matrix.write_bytes(b"matrix") + (tmp_path / "benchmark" / "summary.json").write_text( + _legacy_compact_summary_without_native_smoke(suite), + encoding="utf-8", + ) + refresh_artifact_checksums(tmp_path) + backup_root = tmp_path / "native-backup" + _write_native_smoke_backup( + backup_root, + repo_id=repo_id, + suite=suite, + bit_setting="W4A4", + ) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCleanupHfApi(tmp_path) + + result = repair_hf_native_smoke_proof( + repo_id=repo_id, + suite=suite, + native_smoke_backup_root=backup_root, + revision="main", + commit_message="restore native smoke proof", + api=fake_api, + ) + + assert result["commit"]["commit_oid"] == "repair-sha" + assert result["backup_skipped_reason"] is None + assert result["proof_source"] == "local_paired_native_smoke_backup" + assert result["changed_files"] == [ + "benchmark/summary.json", + "orbitquant_manifest.json", + "README.md", + "SHA256SUMS", + ] + operation_by_path = { + operation.path_in_repo: operation.path_or_fileobj + for operation in fake_api.create_commit_calls[0]["operations"] + } + assert sorted(operation_by_path) == sorted(result["changed_files"]) + + repaired_summary = json.loads(operation_by_path["benchmark/summary.json"]) + proof = repaired_summary["native_smoke"] + assert proof["proof_source"] == "local_paired_native_smoke_backup" + assert proof["comparison_asset_path"] == "assets/image_generation_comparison_matrix.webp" + assert proof["paired_prompt_seed_count"] == 2 + assert proof["paired_prompt_seed_keys"] == [ + ["flux2-native", "0", "counting"], + ["flux2-native", "0", "simple-object"], + ] + assert proof["splits"]["original"]["generated_samples"] == 2 + assert proof["splits"]["orbitquant"]["generated_samples"] == 2 + assert proof["splits"]["original"]["nonempty_output_count"] == 2 + assert proof["splits"]["orbitquant"]["nonempty_output_count"] == 2 + assert proof["splits"]["original"]["native_settings"] == [ + hub_module._native_smoke_expected_settings(suite) + ] + assert "## Native Validation Evidence" in operation_by_path["README.md"].decode("utf-8") + + +def test_repair_hf_native_smoke_proof_refreshes_stale_readme_when_proof_exists( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + matrix.parent.mkdir(parents=True, exist_ok=True) + matrix.write_bytes(b"matrix") + summary_path = tmp_path / "benchmark" / "summary.json" + summary_path.write_text(_native_smoke_summary(suite), encoding="utf-8") + refresh_artifact_checksums(tmp_path) + manifest = hub_module.OrbitQuantManifest.from_dict( + json.loads((tmp_path / "orbitquant_manifest.json").read_text()) + ) + (tmp_path / "README.md").write_text(hub_module.render_model_card(manifest), encoding="utf-8") + write_sha256sums(tmp_path) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCleanupHfApi(tmp_path) + + result = repair_hf_native_smoke_proof( + repo_id=repo_id, + suite=suite, + revision="main", + api=fake_api, + ) + + assert result["existing_native_smoke_ready"] is True + assert result["changed_files"] == ["README.md", "SHA256SUMS"] + operation_by_path = { + operation.path_in_repo: operation.path_or_fileobj + for operation in fake_api.create_commit_calls[0]["operations"] + } + repaired_readme = operation_by_path["README.md"].decode("utf-8") + assert "## Native Validation Evidence" in repaired_readme + assert "| Paired prompt/seed count | `1` |" in repaired_readme + + +def test_repair_hf_native_smoke_proof_normalizes_legacy_video_export_fps( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="wan-native", + model_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + pipeline="WanPipeline", + width=832, + height=480, + frames=81, + export_fps=16, + steps=50, + guidance=5.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/Wan2.1-T2V-1.3B-Diffusers-OrbitQuant-W4A4" + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" + matrix.parent.mkdir(parents=True, exist_ok=True) + matrix.write_bytes(b"matrix") + summary_path = tmp_path / "benchmark" / "summary.json" + summary = json.loads(_native_smoke_summary(suite)) + summary["native_smoke"]["comparison_asset_path"] = ( + "assets/video_generation_comparison_matrix.webp" + ) + summary["native_smoke"]["proof_source"] = "local_paired_native_smoke_backup" + for split in ("original", "orbitquant"): + del summary["native_smoke"]["splits"][split]["native_settings"][0]["export_fps"] + summary_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + refresh_artifact_checksums(tmp_path) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCleanupHfApi(tmp_path) + + result = repair_hf_native_smoke_proof( + repo_id=repo_id, + suite=suite, + revision="main", + api=fake_api, + ) + + assert result["existing_native_smoke_ready"] is False + assert result["normalized_native_smoke_settings"] is True + assert set(result["changed_files"]) >= { + "benchmark/summary.json", + "orbitquant_manifest.json", + "SHA256SUMS", + } + assert "model.safetensors" not in result["changed_files"] + operation_by_path = { + operation.path_in_repo: operation.path_or_fileobj + for operation in fake_api.create_commit_calls[0]["operations"] + } + repaired_summary = json.loads(operation_by_path["benchmark/summary.json"]) + expected_settings = hub_module._native_smoke_expected_settings(suite) + for split in ("original", "orbitquant"): + assert repaired_summary["native_smoke"]["splits"][split]["native_settings"] == [ + expected_settings + ] + status = hub_module._native_smoke_proof_status( + repaired_summary, + suite=suite, + file_names={"assets/video_generation_comparison_matrix.webp"}, + ) + assert status["ready"] is True + + +def test_repair_hf_native_smoke_proof_skips_when_video_frames_are_insufficient( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="wan-native", + model_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + pipeline="WanPipeline", + width=832, + height=480, + frames=81, + steps=50, + guidance=5.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/Wan2.1-T2V-1.3B-Diffusers-OrbitQuant-W4A4" + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" + matrix.parent.mkdir(parents=True, exist_ok=True) + matrix.write_bytes(b"matrix") + (tmp_path / "benchmark" / "summary.json").write_text( + _legacy_compact_summary_without_native_smoke( + suite, + generated_samples=1, + generated_frames=1, + ), + encoding="utf-8", + ) + refresh_artifact_checksums(tmp_path) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCleanupHfApi(tmp_path) + + result = repair_hf_native_smoke_proof( + repo_id=repo_id, + suite=suite, + dry_run=True, + api=fake_api, + ) + + assert result["commit"] is None + assert result["changed_files"] == [] + assert result["repair_skipped_reason"] == "raw_paired_native_smoke_evidence_missing" + assert fake_api.create_commit_calls == [] + + +def test_repair_hf_native_smoke_proof_matrix_skips_compact_summary_without_raw_pair_evidence( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + matrix.parent.mkdir(parents=True, exist_ok=True) + matrix.write_bytes(b"matrix") + (tmp_path / "benchmark" / "summary.json").write_text( + _legacy_compact_summary_without_native_smoke(suite), + encoding="utf-8", + ) + refresh_artifact_checksums(tmp_path) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCleanupHfApi(tmp_path) + + result = repair_hf_native_smoke_proof_matrix( + suites=[suite], + dry_run=True, + api=fake_api, + ) + + assert result["repo_count"] == 1 + assert result["changed_repo_count"] == 0 + assert result["skipped_repo_count"] == 1 + assert result["error_count"] == 0 + assert result["rows"][0]["repo_id"] == repo_id + assert result["rows"][0]["suite"] == "flux2-native" + assert result["rows"][0]["bit_setting"] == "W4A4" + assert result["rows"][0]["repair_skipped_reason"] == ( + "raw_paired_native_smoke_evidence_missing" + ) + + +def test_repair_hf_native_smoke_proof_matrix_uses_local_backup( + tmp_path, + monkeypatch, +): + suite = NativeSuite( + name="flux2-native", + model_id="black-forest-labs/FLUX.2-klein-4B", + pipeline="Flux2KleinPipeline", + width=1024, + height=1024, + steps=4, + guidance=1.0, + bit_settings=["W4A4"], + ) + repo_id = "WaveCut/FLUX.2-klein-4B-OrbitQuant-W4A4" + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + matrix.parent.mkdir(parents=True, exist_ok=True) + matrix.write_bytes(b"matrix") + (tmp_path / "benchmark" / "summary.json").write_text( + _legacy_compact_summary_without_native_smoke(suite), + encoding="utf-8", + ) + refresh_artifact_checksums(tmp_path) + backup_root = tmp_path / "native-backup" + _write_native_smoke_backup( + backup_root, + repo_id=repo_id, + suite=suite, + bit_setting="W4A4", + prompt_ids=("simple-object",), + ) + + monkeypatch.setattr( + hub_module, "hf_hub_download", make_fake_download(_remote_file_map(repo_id, tmp_path)) + ) + fake_api = FakeCleanupHfApi(tmp_path) + + result = repair_hf_native_smoke_proof_matrix( + suites=[suite], + native_smoke_backup_root=backup_root, + dry_run=True, + api=fake_api, + ) + + assert result["repo_count"] == 1 + assert result["changed_repo_count"] == 1 + assert result["skipped_repo_count"] == 0 + assert result["error_count"] == 0 + assert result["rows"][0]["repo_id"] == repo_id + assert result["rows"][0]["proof_source"] == "local_paired_native_smoke_backup" + assert result["rows"][0]["backup_skipped_reason"] is None + assert result["rows"][0]["changed_files"] == [ + "benchmark/summary.json", + "orbitquant_manifest.json", + "README.md", + "SHA256SUMS", + ] diff --git a/tests/test_hub_upload.py b/tests/test_hub_upload.py new file mode 100644 index 0000000..c68d258 --- /dev/null +++ b/tests/test_hub_upload.py @@ -0,0 +1,515 @@ +import json +from pathlib import Path + +import pytest +from hub_helpers import FakeHfApi, _write_artifact + +import orbitquant.hub as hub_module +from orbitquant.artifacts import record_artifact_metrics, refresh_artifact_checksums +from orbitquant.artifacts.checksums import read_sha256sums +from orbitquant.hub import stage_compact_upload_artifact, upload_orbitquant_artifact + + +def test_upload_orbitquant_artifact_dry_run_validates_without_hub_calls(tmp_path): + _write_artifact(tmp_path) + fake_api = FakeHfApi() + + result = upload_orbitquant_artifact( + tmp_path, + repo_id="WaveCut/example-orbitquant", + dry_run=True, + api=fake_api, + ) + + assert result["dry_run"] is True + assert result["repo_id"] == "WaveCut/example-orbitquant" + assert result["private"] is True + assert result["validation"]["valid"] is True + assert result["upload"] is None + assert fake_api.create_repo_calls == [] + assert fake_api.upload_folder_calls == [] + assert fake_api.model_info_calls == [] + + +def test_stage_compact_upload_artifact_omits_raw_eval_reports_and_promotes_matrices( + tmp_path, +): + _write_artifact(tmp_path) + raw_eval_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_geneval-00000.png" + raw_eval_metadata = raw_eval_asset.with_suffix(".png.json") + visual_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" + comparison_asset = ( + tmp_path / "assets" / "original_vs_orbitquant_flux2-native_seed0_W4A4_simple-object.webp" + ) + raw_video_asset = tmp_path / "assets" / "wan-native_seed0_W4A4_simple-motion.mp4" + raw_video_metadata = raw_video_asset.with_suffix(".mp4.json") + raw_eval_asset.write_bytes(b"raw eval image") + raw_eval_metadata.write_text('{"prompt_id":"geneval-00000"}\n', encoding="utf-8") + visual_asset.write_bytes(b"visual image") + comparison_asset.write_bytes(b"comparison") + raw_video_asset.write_bytes(b"raw video") + raw_video_metadata.write_text('{"prompt_id":"simple-motion"}\n', encoding="utf-8") + report_dir = tmp_path.parent / f"{tmp_path.name}-report" + report_matrix = report_dir / "assets" / "image_generation_comparison_matrix.webp" + report_matrix.parent.mkdir(parents=True) + report_matrix.write_bytes(b"report matrix") + report_markdown = report_dir / "orbitquant-native-eval.md" + report_markdown.write_text("# local report log\n", encoding="utf-8") + report_table = report_dir / "tables" / "perf.csv" + report_table.parent.mkdir(parents=True) + report_table.write_text("metric,value\n", encoding="utf-8") + in_artifact_report = ( + tmp_path + / "reports" + / "native" + / "flux2-w4a4" + / "assets" + / "image_generation_comparison_matrix.webp" + ) + in_artifact_report.parent.mkdir(parents=True) + in_artifact_report.write_bytes(b"in artifact report matrix") + in_artifact_markdown = in_artifact_report.parent.parent / "orbitquant-native-eval.md" + in_artifact_markdown.write_text("# local artifact report log\n", encoding="utf-8") + refresh_artifact_checksums(tmp_path) + + stage_dir = tmp_path.parent / f"{tmp_path.name}-stage" + result = stage_compact_upload_artifact( + tmp_path, + stage_dir, + report_dirs=[report_dir], + validate_tensors=False, + ) + + staged_checksums = read_sha256sums(stage_dir / "SHA256SUMS") + assert result["validation"]["valid"] is True + assert result["omitted_raw_eval_asset_count"] == 6 + assert result["omitted_report_file_count"] == 2 + assert result["copied_report_asset_count"] == 1 + assert not (stage_dir / raw_eval_asset.relative_to(tmp_path)).exists() + assert not (stage_dir / raw_eval_metadata.relative_to(tmp_path)).exists() + assert not (stage_dir / visual_asset.relative_to(tmp_path)).exists() + assert not (stage_dir / comparison_asset.relative_to(tmp_path)).exists() + assert not (stage_dir / raw_video_asset.relative_to(tmp_path)).exists() + assert not (stage_dir / raw_video_metadata.relative_to(tmp_path)).exists() + assert not (stage_dir / in_artifact_report.relative_to(tmp_path)).exists() + assert not (stage_dir / in_artifact_markdown.relative_to(tmp_path)).exists() + staged_report = stage_dir / "assets" / "image_generation_comparison_matrix.webp" + staged_collision_report = ( + stage_dir / "assets" / f"{report_dir.name}_image_generation_comparison_matrix.webp" + ) + assert staged_report.is_file() + assert staged_report.read_bytes() == b"in artifact report matrix" + assert not staged_collision_report.exists() + assert raw_eval_asset.relative_to(tmp_path).as_posix() not in staged_checksums + assert visual_asset.relative_to(tmp_path).as_posix() not in staged_checksums + assert comparison_asset.relative_to(tmp_path).as_posix() not in staged_checksums + assert raw_video_asset.relative_to(tmp_path).as_posix() not in staged_checksums + assert not any(path.startswith("reports/") for path in staged_checksums) + assert staged_report.relative_to(stage_dir).as_posix() in staged_checksums + + +def test_stage_compact_upload_artifact_keeps_one_direct_comparison_matrix(tmp_path): + _write_artifact(tmp_path) + image_matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + video_matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" + image_matrix.write_bytes(b"image matrix") + video_matrix.write_bytes(b"video matrix") + report_matrix = ( + tmp_path + / "reports" + / "native" + / "flux2-w4a4" + / "assets" + / "image_generation_comparison_matrix.webp" + ) + report_matrix.parent.mkdir(parents=True) + report_matrix.write_bytes(b"report matrix") + refresh_artifact_checksums(tmp_path) + + stage_dir = tmp_path.parent / f"{tmp_path.name}-stage" + result = stage_compact_upload_artifact( + tmp_path, + stage_dir, + validate_tensors=False, + ) + + staged_checksums = read_sha256sums(stage_dir / "SHA256SUMS") + assert result["copied_report_asset_count"] == 0 + assert (stage_dir / "assets" / "image_generation_comparison_matrix.webp").read_bytes() == ( + b"image matrix" + ) + assert not (stage_dir / "assets" / "video_generation_comparison_matrix.webp").exists() + assert "assets/image_generation_comparison_matrix.webp" in staged_checksums + assert "assets/video_generation_comparison_matrix.webp" not in staged_checksums + + +def test_stage_compact_upload_artifact_enforces_asset_allowlist(tmp_path): + _write_artifact(tmp_path) + allowed_matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + extra_matrix = tmp_path / "assets" / "video_generation_comparison_matrix.webp" + raw_png = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" + raw_png_sidecar = raw_png.with_suffix(".png.json") + raw_webp = tmp_path / "assets" / "original_vs_orbitquant_seed0.webp" + contact_sheet = tmp_path / "assets" / "wan_contact_sheet.webp" + raw_video = tmp_path / "assets" / "wan-native_seed0_W4A4_motion.mp4" + raw_video_sidecar = raw_video.with_suffix(".mp4.json") + nested_raw = tmp_path / "assets" / "nested" / "debug_frame.png" + nested_matrix = tmp_path / "assets" / "nested" / "debug_generation_comparison_matrix.webp" + unexpected_file = tmp_path / "debug.txt" + ignored_cache = tmp_path / ".cache" / "upload.tmp" + for path, payload in ( + (allowed_matrix, b"matrix"), + (extra_matrix, b"video matrix"), + (raw_png, b"png"), + (raw_png_sidecar, b"json"), + (raw_webp, b"webp"), + (contact_sheet, b"contact"), + (raw_video, b"mp4"), + (raw_video_sidecar, b"mp4 json"), + (nested_raw, b"nested"), + (nested_matrix, b"nested matrix"), + (unexpected_file, b"debug"), + (ignored_cache, b"ignored"), + ): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(payload) + refresh_artifact_checksums(tmp_path) + + stage_dir = tmp_path.parent / f"{tmp_path.name}-asset-allowlist-stage" + result = stage_compact_upload_artifact(tmp_path, stage_dir, validate_tensors=False) + + staged_files = { + path.relative_to(stage_dir).as_posix() for path in stage_dir.rglob("*") if path.is_file() + } + expected_files = set(hub_module._REQUIRED_ARTIFACT_FILES) | { + "assets/image_generation_comparison_matrix.webp" + } + assert staged_files == expected_files + + staged_manifest = json.loads((stage_dir / "orbitquant_manifest.json").read_text()) + manifest_checksum_files = set(staged_manifest["checksums"]) + sha256sum_files = set(read_sha256sums(stage_dir / "SHA256SUMS")) + assert manifest_checksum_files == staged_files - { + "README.md", + "SHA256SUMS", + "orbitquant_manifest.json", + } + assert sha256sum_files == staged_files - {"SHA256SUMS"} + + staged_summary = json.loads((stage_dir / "benchmark" / "summary.json").read_text()) + assert staged_summary["published_summary"] == "compact" + assert staged_summary["raw_generation_records"] == "local-only" + assert not any(path.endswith(".png") for path in staged_files) + assert not any(path.endswith(".png.json") for path in staged_files) + assert not any(path.endswith(".mp4") for path in staged_files) + assert "assets/video_generation_comparison_matrix.webp" not in staged_files + assert "assets/original_vs_orbitquant_seed0.webp" not in staged_files + assert "assets/wan_contact_sheet.webp" not in staged_files + assert "assets/nested/debug_frame.png" not in staged_files + assert "assets/nested/debug_generation_comparison_matrix.webp" not in staged_files + assert "debug.txt" not in staged_files + assert ".cache/upload.tmp" not in staged_files + assert result["omitted_unexpected_file_count"] == 1 + assert result["omitted_unexpected_files"] == ["debug.txt"] + assert set(result["omitted_raw_eval_assets"]) >= { + "assets/video_generation_comparison_matrix.webp", + "assets/flux2-native_seed0_W4A4_simple-object.png", + "assets/flux2-native_seed0_W4A4_simple-object.png.json", + "assets/original_vs_orbitquant_seed0.webp", + "assets/wan_contact_sheet.webp", + "assets/wan-native_seed0_W4A4_motion.mp4", + "assets/wan-native_seed0_W4A4_motion.mp4.json", + "assets/nested/debug_frame.png", + "assets/nested/debug_generation_comparison_matrix.webp", + } + + +def test_stage_compact_upload_artifact_writes_native_smoke_proof(tmp_path): + _write_artifact(tmp_path) + matrix = tmp_path / "assets" / "image_generation_comparison_matrix.webp" + original_output = tmp_path / "assets" / "flux2-native_seed0_original_simple-object.png" + orbitquant_output = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" + for path, payload in ( + (matrix, b"matrix"), + (original_output, b"original image"), + (orbitquant_output, b"orbitquant image"), + ): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(payload) + for split, output, bit_setting in ( + ("original", original_output, "original"), + ("orbitquant", orbitquant_output, "W4A4"), + ): + record_artifact_metrics( + tmp_path, + split=split, + metrics={"generated_samples": 1}, + metadata={ + "suite": "flux2-native", + "prompt": "A native prompt", + "prompt_record": {"id": "simple-object"}, + "seed": 0, + "height": 1024, + "width": 1024, + "frames": None, + "steps": 4, + "guidance": 1.0, + "bit_setting": bit_setting, + "output_path": str(output), + }, + validate_checksums_enabled=False, + refresh_checksums_enabled=False, + ) + refresh_artifact_checksums(tmp_path) + + stage_dir = tmp_path.parent / f"{tmp_path.name}-native-smoke-stage" + stage_compact_upload_artifact(tmp_path, stage_dir, validate_tensors=False) + + staged_summary = json.loads((stage_dir / "benchmark" / "summary.json").read_text()) + staged_readme = (stage_dir / "README.md").read_text(encoding="utf-8") + assert staged_summary["raw_generation_records"] == "local-only" + assert "latest" not in staged_summary["metrics"]["original"] + assert "## Native Validation Evidence" in staged_readme + assert "| Comparison matrix | `assets/image_generation_comparison_matrix.webp` |" in ( + staged_readme + ) + assert "| Paired prompt/seed count | `1` |" in staged_readme + assert "| BF16 source nonempty outputs | `1` |" in staged_readme + assert "| OrbitQuant nonempty outputs | `1` |" in staged_readme + assert "original.metrics.jsonl" not in staged_readme + assert staged_summary["native_smoke"] == { + "proof_format": "orbitquant-native-smoke-v1", + "comparison_asset_path": "assets/image_generation_comparison_matrix.webp", + "paired_prompt_seed_count": 1, + "paired_prompt_seed_keys": [["flux2-native", "0", "simple-object"]], + "splits": { + "original": { + "records": 1, + "generated_samples": 1, + "generated_frames": 0, + "nonempty_output_count": 1, + "seeds": ["0"], + "prompt_ids": ["simple-object"], + "pair_keys": [["flux2-native", "0", "simple-object"]], + "native_settings": [ + { + "suite": "flux2-native", + "height": 1024, + "width": 1024, + "frames": None, + "steps": 4, + "guidance": 1.0, + } + ], + }, + "orbitquant": { + "records": 1, + "generated_samples": 1, + "generated_frames": 0, + "nonempty_output_count": 1, + "seeds": ["0"], + "prompt_ids": ["simple-object"], + "pair_keys": [["flux2-native", "0", "simple-object"]], + "native_settings": [ + { + "suite": "flux2-native", + "height": 1024, + "width": 1024, + "frames": None, + "steps": 4, + "guidance": 1.0, + } + ], + }, + }, + } + + +def test_stage_compact_upload_artifact_promotes_compare_native_bundle(tmp_path): + _write_artifact(tmp_path) + bundle_dir = tmp_path.parent / f"{tmp_path.name}-compare-native" + bundle_dir.mkdir() + comparison = bundle_dir / "flux2-native_seed0_W4A4_original_vs_orbitquant.webp" + original = bundle_dir / "flux2-native_seed0_original.png" + orbitquant = bundle_dir / "flux2-native_seed0_W4A4.png" + for path, payload in ( + (comparison, b"comparison"), + (original, b"original"), + (orbitquant, b"orbitquant"), + ): + path.write_bytes(payload) + metadata = { + "suite": "flux2-native", + "prompt": "A native prompt", + "prompt_record": {"id": "simple-object"}, + "seed": 0, + "height": 1024, + "width": 1024, + "frames": None, + "steps": 4, + "guidance": 1.0, + } + for output in (original, orbitquant): + output.with_suffix(output.suffix + ".json").write_text( + json.dumps(metadata) + "\n", encoding="utf-8" + ) + (bundle_dir / "summary.json").write_text( + json.dumps( + { + **metadata, + "model_id": "example/model", + "bit_setting": "W4A4", + "comparison_path": "/workspace/run/" + comparison.name, + "original": { + "output_path": "/workspace/run/" + original.name, + "metadata_path": "/workspace/run/" + original.name + ".json", + }, + "orbitquant": { + "output_path": "/workspace/run/" + orbitquant.name, + "metadata_path": "/workspace/run/" + orbitquant.name + ".json", + }, + } + ) + + "\n", + encoding="utf-8", + ) + + stage_dir = tmp_path.parent / f"{tmp_path.name}-compare-native-stage" + result = stage_compact_upload_artifact( + tmp_path, + stage_dir, + report_dirs=[bundle_dir], + validate_tensors=False, + ) + + matrix_path = stage_dir / "assets" / "image_generation_comparison_matrix.webp" + summary = json.loads((stage_dir / "benchmark" / "summary.json").read_text()) + proof = summary["native_smoke"] + assert result["copied_report_assets"] == ["assets/image_generation_comparison_matrix.webp"] + assert matrix_path.read_bytes() == b"comparison" + assert proof["proof_source"] == "local_compare_native_bundle" + assert proof["comparison_asset_path"] == ("assets/image_generation_comparison_matrix.webp") + assert proof["paired_prompt_seed_keys"] == [["flux2-native", "0", "simple-object"]] + assert proof["splits"]["original"]["nonempty_output_count"] == 1 + assert proof["splits"]["orbitquant"]["nonempty_output_count"] == 1 + assert not (stage_dir / original.name).exists() + assert not (stage_dir / orbitquant.name).exists() + + +def test_upload_orbitquant_artifact_rejects_raw_full_upload_profile(tmp_path): + _write_artifact(tmp_path) + fake_api = FakeHfApi() + + with pytest.raises(ValueError, match="unsupported upload profile: full"): + upload_orbitquant_artifact( + tmp_path, + repo_id="WaveCut/example-orbitquant", + private=False, + revision="main", + commit_message="upload test artifact", + replace_repo_files=True, + validate_tensors=False, + upload_profile="full", + api=fake_api, + ) + + assert fake_api.create_repo_calls == [] + assert fake_api.upload_folder_calls == [] + + +def test_upload_orbitquant_artifact_can_upload_compact_staged_copy(tmp_path): + _write_artifact(tmp_path) + raw_eval_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_geneval-00000.png" + visual_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" + raw_eval_asset.write_bytes(b"raw eval image") + visual_asset.write_bytes(b"visual image") + refresh_artifact_checksums(tmp_path) + stage_dir = tmp_path.parent / f"{tmp_path.name}-upload-stage" + fake_api = FakeHfApi() + + result = upload_orbitquant_artifact( + tmp_path, + repo_id="WaveCut/example-orbitquant", + upload_profile="compact", + staging_dir=stage_dir, + validate_tensors=False, + api=fake_api, + ) + + assert result["upload_profile"] == "compact" + assert result["staging"]["enabled"] is True + assert result["staging"]["omitted_raw_eval_asset_count"] == 2 + assert len(fake_api.upload_folder_calls) == 1 + assert fake_api.upload_folder_calls[0]["folder_path"] == str(stage_dir) + assert not (stage_dir / raw_eval_asset.relative_to(tmp_path)).exists() + assert not (stage_dir / visual_asset.relative_to(tmp_path)).exists() + assert not (stage_dir / "benchmark" / "original.metrics.jsonl").exists() + assert (stage_dir / "benchmark" / "summary.json").is_file() + staged_summary = json.loads((stage_dir / "benchmark" / "summary.json").read_text()) + assert staged_summary["published_summary"] == "compact" + assert staged_summary["raw_generation_records"] == "local-only" + assert result["validation"]["valid"] is True + assert result["upload"]["commit_oid"] == "uploaded-sha" + + +def test_upload_orbitquant_artifact_defaults_to_compact_staged_copy(tmp_path): + _write_artifact(tmp_path) + raw_eval_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_geneval-00000.png" + visual_asset = tmp_path / "assets" / "flux2-native_seed0_W4A4_simple-object.png" + raw_eval_asset.write_bytes(b"raw eval image") + visual_asset.write_bytes(b"visual image") + refresh_artifact_checksums(tmp_path) + fake_api = FakeHfApi() + + result = upload_orbitquant_artifact( + tmp_path, + repo_id="WaveCut/example-orbitquant", + validate_tensors=False, + api=fake_api, + ) + + assert result["upload_profile"] == "compact" + assert result["replace_repo_files"] is True + assert result["upload_kwargs"]["delete_patterns"] == "*" + assert result["staging"]["enabled"] is True + assert result["staging"]["omitted_raw_eval_asset_count"] == 2 + assert set(result["staging"]["omitted_raw_eval_assets"]) == { + "assets/flux2-native_seed0_W4A4_geneval-00000.png", + "assets/flux2-native_seed0_W4A4_simple-object.png", + } + upload_path = Path(fake_api.upload_folder_calls[0]["folder_path"]) + assert upload_path != tmp_path + + +def test_upload_orbitquant_artifact_rejects_remote_file_replacement_opt_out(tmp_path): + _write_artifact(tmp_path) + fake_api = FakeHfApi() + + with pytest.raises(ValueError, match="must replace remote files"): + upload_orbitquant_artifact( + tmp_path, + repo_id="WaveCut/example-orbitquant", + replace_repo_files=False, + validate_tensors=False, + api=fake_api, + ) + + assert fake_api.create_repo_calls == [] + assert fake_api.upload_folder_calls == [] + assert fake_api.model_info_calls == [] + + +def test_upload_orbitquant_artifact_rejects_invalid_artifact_before_hub_calls(tmp_path): + fake_api = FakeHfApi() + + with pytest.raises(RuntimeError, match="required artifact file missing"): + upload_orbitquant_artifact( + tmp_path, + repo_id="WaveCut/bad-artifact", + validate_tensors=False, + api=fake_api, + ) + + assert fake_api.create_repo_calls == [] + assert fake_api.upload_folder_calls == [] + assert fake_api.model_info_calls == [] diff --git a/tests/test_kernels.py b/tests/test_kernels.py index d70e581..637d1d4 100644 --- a/tests/test_kernels.py +++ b/tests/test_kernels.py @@ -11,12 +11,46 @@ quantize_activations_kernel, select_backend, ) +from orbitquant.kernels.executorch_vulkan import ExecuTorchVulkanW4A4Linear from orbitquant.layers import OrbitQuantLinear from orbitquant.modeling import quantize_linear_modules from orbitquant.packing import pack_lowbit, unpack_lowbit from orbitquant.rotations import RPBHRotation +def _make_executorch_vulkan_export_layer() -> ExecuTorchVulkanW4A4Linear: + source = torch.nn.Linear(24, 7, bias=True) + layer = OrbitQuantLinear.from_linear( + source, + config=OrbitQuantConfig( + weight_bits=4, + activation_bits=4, + block_size="paper", + runtime_mode="dequant_bf16", + ), + module_name="probe", + ) + return ExecuTorchVulkanW4A4Linear(layer) + + +def test_executorch_vulkan_export_keeps_exact_w4a4_as_one_custom_op(): + exported = torch.export.export( + _make_executorch_vulkan_export_layer(), + (torch.randn(2, 3, 24, dtype=torch.float32),), + strict=True, + ) + + call_targets = {node.target for node in exported.graph.nodes if node.op == "call_function"} + assert torch.ops.orbitquant_vulkan.linear_w4a4_exact.default in call_targets + + +def test_executorch_vulkan_export_rejects_unsupported_bfloat16_activations(): + layer = _make_executorch_vulkan_export_layer() + + with pytest.raises(TypeError, match="float16 or float32"): + layer(torch.randn(2, 24, dtype=torch.bfloat16)) + + def test_cpu_activation_kernel_matches_reference_functional_path(): torch.manual_seed(0) x = torch.randn(2, 3, 16) @@ -42,9 +76,7 @@ def test_activation_quantization_is_per_token_scale_equivariant_and_batch_indepe codebook = get_codebook(dim=16, bits=4) baseline = quantize_activations(x, rotation=rotation, codebook=codebook, eps=1e-12) - scaled = quantize_activations( - x * scales, rotation=rotation, codebook=codebook, eps=1e-12 - ) + scaled = quantize_activations(x * scales, rotation=rotation, codebook=codebook, eps=1e-12) assert torch.allclose(scaled, baseline * scales, atol=1e-6, rtol=1e-6) @@ -117,16 +149,13 @@ def test_backend_capabilities_report_partial_and_fallback_kernel_status(monkeypa monkeypatch.setattr(dispatch_module, "_native_cpu_packed_matmul_available", lambda: False) monkeypatch.setattr(dispatch_module, "_native_cpu_activation_available", lambda: False) monkeypatch.setattr(dispatch_module, "_native_cpu_adaln_available", lambda: False) - capabilities = backend_capabilities( - backends={"cpu": True, "mps": True, "triton_cuda": True} - ) + capabilities = backend_capabilities(backends={"cpu": True, "mps": True, "triton_cuda": True}) assert capabilities["cpu"]["available"] is True assert capabilities["cpu"]["claim_status"] == "reference_only" assert capabilities["cpu"]["optimized"] is False assert capabilities["cpu"]["implemented_stage"] == ( - "activation_norm_rpbh_quant_rescale,packed_weight_matmul," - "adaln_rtn_packed_matmul" + "activation_norm_rpbh_quant_rescale,packed_weight_matmul,adaln_rtn_packed_matmul" ) assert capabilities["cpu"]["optimized_stage"] is None assert capabilities["cpu"]["weight_dequant_optimized"] is False @@ -160,7 +189,7 @@ def test_backend_capabilities_report_partial_and_fallback_kernel_status(monkeypa expected_triton_stage = ( "activation_norm_rpbh_quant_rescale,packed_weight_dequant," "packed_weight_matmul,lowbit_pack,lowbit_unpack,weight_rotation_fwht_quant_pack," - "adaln_rtn_quant_pack,adaln_rtn_dequant" + "adaln_rtn_quant_pack,adaln_rtn_dequant,adaln_rtn_packed_matmul" ) assert capabilities["triton_cuda"]["implemented_stage"] == expected_triton_stage assert capabilities["triton_cuda"]["optimized_stage"] == expected_triton_stage @@ -170,6 +199,7 @@ def test_backend_capabilities_report_partial_and_fallback_kernel_status(monkeypa assert capabilities["triton_cuda"]["weight_quant_optimized"] is True assert capabilities["triton_cuda"]["adaln_quant_optimized"] is True assert capabilities["triton_cuda"]["adaln_dequant_optimized"] is True + assert capabilities["triton_cuda"]["adaln_packed_matmul_optimized"] is True assert capabilities["triton_cuda"]["full_fusion"] is False assert capabilities["triton_cuda"]["implementation"] == "python_triton_orbitquant_pipeline" assert capabilities["triton_cuda"]["package_format"] == "python_triton" @@ -183,9 +213,7 @@ def test_backend_capabilities_report_mps_metal_partial_kernel(monkeypatch): monkeypatch.setattr(dispatch_module, "_native_cpu_activation_available", lambda: False) monkeypatch.setattr(dispatch_module, "_native_cpu_adaln_available", lambda: False) - capabilities = backend_capabilities( - backends={"cpu": True, "mps": True, "triton_cuda": False} - ) + capabilities = backend_capabilities(backends={"cpu": True, "mps": True, "triton_cuda": False}) assert capabilities["mps"]["available"] is True assert capabilities["mps"]["claim_status"] == "partial_optimized" @@ -194,10 +222,7 @@ def test_backend_capabilities_report_mps_metal_partial_kernel(monkeypatch): capabilities["mps"]["implementation"] == "torch_mps_compile_shader_fused_activation+native_packed_matmul" ) - assert ( - capabilities["mps"]["package_format"] - == "torch.mps.compile_shader,native_kernel_package" - ) + assert capabilities["mps"]["package_format"] == "torch.mps.compile_shader,native_kernel_package" assert ( capabilities["mps"]["optimized_stage"] == "activation_norm_rpbh_quant_rescale,packed_weight_dequant,packed_weight_matmul" @@ -221,7 +246,7 @@ def test_backend_capabilities_report_mps_metal_partial_kernel(monkeypatch): assert capabilities["triton_cuda"]["implemented_stage"] == ( "activation_norm_rpbh_quant_rescale,packed_weight_dequant," "packed_weight_matmul,lowbit_pack,lowbit_unpack,weight_rotation_fwht_quant_pack," - "adaln_rtn_quant_pack,adaln_rtn_dequant" + "adaln_rtn_quant_pack,adaln_rtn_dequant,adaln_rtn_packed_matmul" ) assert capabilities["triton_cuda"]["optimized_stage"] is None assert capabilities["triton_cuda"]["weight_dequant_optimized"] is False @@ -230,6 +255,7 @@ def test_backend_capabilities_report_mps_metal_partial_kernel(monkeypatch): assert capabilities["triton_cuda"]["weight_quant_optimized"] is False assert capabilities["triton_cuda"]["adaln_quant_optimized"] is False assert capabilities["triton_cuda"]["adaln_dequant_optimized"] is False + assert capabilities["triton_cuda"]["adaln_packed_matmul_optimized"] is False assert capabilities["triton_cuda"]["hf_kernel_builder_compliant"] is False @@ -240,16 +266,11 @@ def test_backend_capabilities_report_mps_shader_without_native_matmul(monkeypatc monkeypatch.setattr(dispatch_module, "_native_cpu_activation_available", lambda: False) monkeypatch.setattr(dispatch_module, "_native_cpu_adaln_available", lambda: False) - capabilities = backend_capabilities( - backends={"cpu": True, "mps": True, "triton_cuda": False} - ) + capabilities = backend_capabilities(backends={"cpu": True, "mps": True, "triton_cuda": False}) assert capabilities["mps"]["claim_status"] == "partial_optimized" assert capabilities["mps"]["optimized"] is True - assert ( - capabilities["mps"]["implementation"] - == "torch_mps_compile_shader_fused_activation" - ) + assert capabilities["mps"]["implementation"] == "torch_mps_compile_shader_fused_activation" assert capabilities["mps"]["package_format"] == "torch.mps.compile_shader" assert ( capabilities["mps"]["optimized_stage"] @@ -269,9 +290,7 @@ def test_backend_capabilities_report_mps_native_matmul_without_shader(monkeypatc monkeypatch.setattr(dispatch_module, "_native_cpu_activation_available", lambda: False) monkeypatch.setattr(dispatch_module, "_native_cpu_adaln_available", lambda: False) - capabilities = backend_capabilities( - backends={"cpu": True, "mps": True, "triton_cuda": False} - ) + capabilities = backend_capabilities(backends={"cpu": True, "mps": True, "triton_cuda": False}) assert capabilities["mps"]["claim_status"] == "partial_optimized" assert capabilities["mps"]["optimized"] is True @@ -288,9 +307,7 @@ def test_backend_capabilities_label_cpu_native_matmul_as_partial_not_fused(monke monkeypatch.setattr(dispatch_module, "_native_cpu_activation_available", lambda: False) monkeypatch.setattr(dispatch_module, "_native_cpu_adaln_available", lambda: False) - capabilities = backend_capabilities( - backends={"cpu": True, "mps": False, "triton_cuda": False} - ) + capabilities = backend_capabilities(backends={"cpu": True, "mps": False, "triton_cuda": False}) assert capabilities["cpu"]["claim_status"] == "partial_optimized" assert capabilities["cpu"]["optimized"] is True @@ -309,9 +326,7 @@ def test_backend_capabilities_report_native_cpu_activation_and_matmul(monkeypatc monkeypatch.setattr(dispatch_module, "_native_cpu_activation_available", lambda: True) monkeypatch.setattr(dispatch_module, "_native_cpu_adaln_available", lambda: False) - capabilities = backend_capabilities( - backends={"cpu": True, "mps": False, "triton_cuda": False} - ) + capabilities = backend_capabilities(backends={"cpu": True, "mps": False, "triton_cuda": False}) assert capabilities["cpu"]["claim_status"] == "partial_optimized" assert capabilities["cpu"]["optimized"] is True @@ -332,13 +347,10 @@ def test_backend_capabilities_report_complete_native_cpu_kernel_surface(monkeypa monkeypatch.setattr(dispatch_module, "_native_cpu_activation_available", lambda: True) monkeypatch.setattr(dispatch_module, "_native_cpu_adaln_available", lambda: True) - capabilities = backend_capabilities( - backends={"cpu": True, "mps": False, "triton_cuda": False} - ) + capabilities = backend_capabilities(backends={"cpu": True, "mps": False, "triton_cuda": False}) assert capabilities["cpu"]["optimized_stage"] == ( - "activation_norm_rpbh_quant_rescale,packed_weight_matmul," - "adaln_rtn_packed_matmul" + "activation_norm_rpbh_quant_rescale,packed_weight_matmul,adaln_rtn_packed_matmul" ) assert capabilities["cpu"]["adaln_quant_optimized"] is False assert capabilities["cpu"]["adaln_dequant_optimized"] is True @@ -348,16 +360,159 @@ def test_backend_capabilities_report_complete_native_cpu_kernel_surface(monkeypa assert "without a full floating-point cache" in capabilities["cpu"]["notes"] -def test_backend_selection_accepts_injected_availability_for_gpu_paths(): +def test_backend_selection_accepts_injected_availability_for_gpu_paths(monkeypatch): + monkeypatch.setattr(dispatch_module, "_torch_uses_hip", lambda: False) backends = {"cpu": True, "mps": False, "triton_cuda": True} assert ( - select_backend(torch.device("cuda"), requested="auto", backends=backends) - == "triton_cuda" + select_backend(torch.device("cuda"), requested="auto", backends=backends) == "triton_cuda" ) assert select_backend(torch.device("mps"), requested="auto", backends=backends) == "cpu" +def test_backend_selection_separates_hip_from_cuda(monkeypatch): + monkeypatch.setattr(dispatch_module, "_torch_uses_hip", lambda: True) + backends = { + "cpu": True, + "mps": False, + "triton_cuda": False, + "triton_rocm": True, + } + + assert select_backend(torch.device("cuda"), requested="auto", backends=backends) == "cpu" + assert ( + select_backend(torch.device("cuda"), requested="triton_rocm", backends=backends) + == "triton_rocm" + ) + + +def test_backend_selection_keeps_xpu_explicit_until_hardware_proof(): + backends = { + "cpu": True, + "mps": False, + "triton_cuda": False, + "triton_rocm": False, + "triton_xpu": True, + } + + assert select_backend(torch.device("xpu"), requested="auto", backends=backends) == "cpu" + assert ( + select_backend(torch.device("xpu"), requested="triton_xpu", backends=backends) + == "triton_xpu" + ) + + +def test_available_backends_reports_xpu_only_with_device_and_triton(monkeypatch): + monkeypatch.setattr(dispatch_module, "_xpu_available", lambda: True) + monkeypatch.setattr(dispatch_module, "_triton_available", lambda: True) + + assert available_backends()["triton_xpu"] is True + + +def test_available_backends_does_not_report_cuda_triton_on_hip(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(dispatch_module, "_triton_available", lambda: True) + monkeypatch.setattr(dispatch_module, "_torch_uses_hip", lambda: True) + + backends = available_backends() + + assert backends["triton_cuda"] is False + assert backends["triton_rocm"] is True + + +def test_rocm_capability_remains_experimental_until_hardware_proof(): + capabilities = backend_capabilities( + backends={ + "cpu": True, + "mps": False, + "triton_cuda": False, + "triton_rocm": True, + } + ) + + rocm = capabilities["triton_rocm"] + assert rocm["available"] is True + assert rocm["claim_status"] == "experimental_unverified" + assert rocm["optimized"] is False + assert rocm["optimized_stage"] is None + assert rocm["weight_dequant_optimized"] is False + assert rocm["adaln_packed_matmul_optimized"] is False + assert "without loading CUDA-native extensions" in rocm["notes"] + + +def test_xpu_capability_remains_experimental_until_hardware_proof(): + capabilities = backend_capabilities( + backends={ + "cpu": True, + "mps": False, + "triton_cuda": False, + "triton_rocm": False, + "triton_xpu": True, + } + ) + + xpu = capabilities["triton_xpu"] + assert xpu["available"] is True + assert xpu["claim_status"] == "experimental_unverified" + assert xpu["optimized"] is False + assert xpu["optimized_stage"] is None + assert xpu["weight_dequant_optimized"] is False + assert xpu["adaln_packed_matmul_optimized"] is False + assert xpu["device_types"] == ["xpu"] + assert "explicit-only" in xpu["notes"] + + +@pytest.mark.parametrize("bits", [2, 3, 4, 6]) +def test_triton_rocm_lowbit_pack_unpack_matches_reference(bits): + if not torch.cuda.is_available() or not getattr(torch.version, "hip", None): + pytest.skip("a PyTorch ROCm runtime is required") + if not available_backends().get("triton_rocm", False): + pytest.skip("the ROCm-compatible Triton package is not importable") + + from orbitquant.kernels.triton_cuda import ( + pack_lowbit_with_triton, + unpack_lowbit_with_triton, + ) + + values = torch.arange(0, 1 << bits, dtype=torch.uint8).repeat(5)[:37] + expected_packed = pack_lowbit(values, bits=bits) + actual_packed = pack_lowbit_with_triton(values.to("cuda"), bits=bits) + actual_values = unpack_lowbit_with_triton( + actual_packed, + bits=bits, + length=values.numel(), + ) + + assert torch.equal(actual_packed.cpu(), expected_packed) + assert torch.equal(actual_values.cpu(), values) + + +@pytest.mark.parametrize("bits", [2, 3, 4, 6]) +def test_triton_xpu_lowbit_pack_unpack_matches_reference(bits): + xpu = getattr(torch, "xpu", None) + if xpu is None or not xpu.is_available(): + pytest.skip("a PyTorch XPU runtime is required") + if not available_backends().get("triton_xpu", False): + pytest.skip("the Intel XPU Triton package is not importable") + + from orbitquant.kernels.triton_cuda import ( + pack_lowbit_with_triton, + unpack_lowbit_with_triton, + ) + + values = torch.arange(0, 1 << bits, dtype=torch.uint8).repeat(5)[:37] + expected_packed = pack_lowbit(values, bits=bits) + actual_packed = pack_lowbit_with_triton(values.to("xpu"), bits=bits) + actual_values = unpack_lowbit_with_triton( + actual_packed, + bits=bits, + length=values.numel(), + ) + + assert torch.equal(actual_packed.cpu(), expected_packed) + assert torch.equal(actual_values.cpu(), values) + + def test_triton_availability_requires_importable_triton_modules(monkeypatch): real_import = __import__ @@ -389,7 +544,7 @@ def test_mps_shader_source_declares_fused_activation_and_dequant_kernels(): assert "linear" not in source.lower() -def test_triton_cuda_dispatch_uses_backend_function_with_reference_equivalent_output(monkeypatch): +def test_triton_dispatch_uses_backend_function_with_reference_equivalent_output(monkeypatch): torch.manual_seed(0) x = torch.randn(2, 3, 16) rotation = RPBHRotation(dim=16, seed=3, block_size=8) @@ -401,7 +556,7 @@ def fake_triton_backend(input_tensor, *, rotation, codebook, eps, constant_tenso assert constant_tensors is None return quantize_activations(input_tensor, rotation=rotation, codebook=codebook, eps=eps) - monkeypatch.setattr(dispatch_module, "_triton_cuda_quantize_activations", fake_triton_backend) + monkeypatch.setattr(dispatch_module, "_triton_quantize_activations", fake_triton_backend) monkeypatch.setattr( dispatch_module, "available_backends", @@ -506,9 +661,7 @@ def fail_reference(*args, **kwargs): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) -@pytest.mark.parametrize( - "dim,block_size", [(16, 8), (3072, "paper"), (4096, "paper")] -) +@pytest.mark.parametrize("dim,block_size", [(16, 8), (3072, "paper"), (4096, "paper")]) def test_mps_fused_activation_matches_reference_without_torch_rotation( monkeypatch, dtype, dim, block_size ): @@ -940,11 +1093,7 @@ def test_cuda_quantize_linear_modules_keeps_packed_buffers_on_gpu_until_serializ torch.manual_seed(3) model = torch.nn.Module() model.transformer_blocks = torch.nn.ModuleList( - [ - torch.nn.ModuleDict( - {"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(32, 32)})} - ) - ] + [torch.nn.ModuleDict({"attn": torch.nn.ModuleDict({"to_q": torch.nn.Linear(32, 32)})})] ) config = OrbitQuantConfig(block_size=8, target_policy="generic_dit") @@ -962,3 +1111,588 @@ def test_cuda_quantize_linear_modules_keeps_packed_buffers_on_gpu_until_serializ assert quantized._rotation_signs.is_cuda assert quantized._activation_codebook_centroids.is_cuda assert quantized._activation_codebook_boundaries.is_cuda + + +def _build_w4a4_triton_fixture(in_features: int, out_features: int) -> dict: + from orbitquant.kernels.triton_cuda import ( + fit_int8_centroid_surrogate, + quantize_weight_packed_with_triton, + row_norms_with_triton, + ) + + torch.manual_seed(7) + weight = torch.randn(out_features, in_features, device="cuda", dtype=torch.bfloat16) + rotation = RPBHRotation(dim=in_features, seed=0, block_size="paper") + weight_codebook = get_codebook(dim=in_features, bits=4) + activation_codebook = get_codebook(dim=in_features, bits=4) + row_norms = row_norms_with_triton(weight, eps=1e-10) + packed = quantize_weight_packed_with_triton( + weight, row_norms, rotation=rotation, codebook=weight_codebook, bits=4 + ) + activation_codes, activation_scale = fit_int8_centroid_surrogate( + activation_codebook.centroids + ) + weight_codes, weight_scale = fit_int8_centroid_surrogate(weight_codebook.centroids) + return { + "rotation": rotation, + "weight_codebook": weight_codebook, + "activation_codebook": activation_codebook, + "row_norms": row_norms.to(torch.bfloat16), + "packed": packed, + "activation_codes": activation_codes.to("cuda"), + "activation_scale": activation_scale, + "weight_codes": weight_codes.to("cuda"), + "weight_scale": weight_scale, + "in_features": in_features, + "out_features": out_features, + } + + +@pytest.mark.parametrize("rows", [1, 16, 17, 31, 33, 77]) +def test_w4a4_int_mm_path_supports_arbitrary_row_counts(rows): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + dequantize_packed_weight_with_triton, + matmul_packed_w4a4_with_int_mm, + quantize_activations_packed_w4_with_triton, + quantize_activations_with_triton, + ) + + fixture = _build_w4a4_triton_fixture(in_features=512, out_features=512) + x = torch.randn(rows, fixture["in_features"], device="cuda", dtype=torch.bfloat16) + packed_x, token_norms = quantize_activations_packed_w4_with_triton( + x, + rotation=fixture["rotation"], + codebook=fixture["activation_codebook"], + eps=1e-10, + ) + + actual = matmul_packed_w4a4_with_int_mm( + packed_x, + fixture["packed"], + token_norms, + fixture["row_norms"], + fixture["activation_codes"], + fixture["weight_codes"], + activation_scale=fixture["activation_scale"], + weight_scale=fixture["weight_scale"], + out_features=fixture["out_features"], + in_features=fixture["in_features"], + output_dtype=torch.bfloat16, + ) + + dequantized = dequantize_packed_weight_with_triton( + fixture["packed"], + fixture["row_norms"], + fixture["weight_codebook"], + bits=4, + out_features=fixture["out_features"], + in_features=fixture["in_features"], + ).to(torch.bfloat16) + quantized_values = quantize_activations_with_triton( + x, + rotation=fixture["rotation"], + codebook=fixture["activation_codebook"], + eps=1e-10, + ) + expected = torch.nn.functional.linear(quantized_values.to(torch.bfloat16), dequantized) + + assert actual.shape == (rows, fixture["out_features"]) + relative_error = (actual.float() - expected.float()).norm() / expected.float().norm() + assert relative_error <= 3e-2 + + +@pytest.mark.parametrize("rows", [32, 400]) +def test_int8_activation_lowbit_fused_matches_packed_fused(rows): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + matmul_int8_activations_packed_lowbit_fused_with_triton, + matmul_packed_w4a4_fused_with_triton, + quantize_activations_packed_w4_with_triton, + ) + + fixture = _build_w4a4_triton_fixture(in_features=512, out_features=512) + x = torch.randn(rows, fixture["in_features"], device="cuda", dtype=torch.bfloat16) + packed_x, token_norms = quantize_activations_packed_w4_with_triton( + x, + rotation=fixture["rotation"], + codebook=fixture["activation_codebook"], + eps=1e-10, + ) + bias = torch.randn(fixture["out_features"], device="cuda", dtype=torch.bfloat16) + packed_result = matmul_packed_w4a4_fused_with_triton( + packed_x, + fixture["packed"], + token_norms, + fixture["row_norms"], + fixture["activation_codes"], + fixture["weight_codes"], + activation_scale=fixture["activation_scale"], + weight_scale=fixture["weight_scale"], + out_features=fixture["out_features"], + in_features=fixture["in_features"], + bias=bias, + output_dtype=torch.bfloat16, + ) + + # Decode the packed nibbles into the same INT8 surrogates the fused + # kernel would produce, then run the direct-activation kernel. + low = (packed_x & 15).to(torch.long) + high = (packed_x >> 4).to(torch.long) + int8_x = torch.empty( + (rows, fixture["in_features"]), device="cuda", dtype=torch.int8 + ) + int8_x[:, 0::2] = fixture["activation_codes"][low] + int8_x[:, 1::2] = fixture["activation_codes"][high] + + direct_result = matmul_int8_activations_packed_lowbit_fused_with_triton( + int8_x, + fixture["packed"], + token_norms, + fixture["row_norms"], + fixture["weight_codes"], + weight_bits=4, + activation_scale=fixture["activation_scale"], + weight_scale=fixture["weight_scale"], + out_features=fixture["out_features"], + in_features=fixture["in_features"], + bias=bias, + output_dtype=torch.bfloat16, + ) + torch.testing.assert_close(direct_result, packed_result) + + +@pytest.mark.parametrize("rows", [32, 400]) +def test_w2a4_fused_triton_matmul_matches_int8_emulation(rows): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + matmul_int8_activations_packed_lowbit_fused_with_triton, + matmul_packed_w2a4_fused_with_triton, + ) + + torch.manual_seed(0) + in_features, out_features = 512, 512 + act_indices = torch.randint(0, 16, (rows, in_features), dtype=torch.uint8) + weight_indices = torch.randint(0, 4, (out_features, in_features), dtype=torch.uint8) + act_flat = act_indices.flatten() + packed_x = (act_flat[0::2] | (act_flat[1::2] << 4)).reshape( + rows, in_features // 2 + ).cuda() + w_flat = weight_indices.flatten() + packed_w = ( + w_flat[0::4] | (w_flat[1::4] << 2) | (w_flat[2::4] << 4) | (w_flat[3::4] << 6) + ).cuda() + activation_codes = torch.arange(-8, 8, dtype=torch.int8, device="cuda") + weight_codes = torch.tensor([-7, -2, 2, 7], dtype=torch.int8, device="cuda") + token_norms = (torch.rand(rows, device="cuda") + 0.5).float() + row_norms = (torch.rand(out_features, device="cuda") + 0.5).float() + scale = 0.01 + act_i8 = activation_codes[act_indices.long().cuda()] + w_i8 = weight_codes[weight_indices.long().cuda()] + expected = ( + act_i8.float() @ w_i8.float().t() + * token_norms[:, None] * row_norms[None, :] * scale + ) + + packed_result = matmul_packed_w2a4_fused_with_triton( + packed_x, + packed_w, + token_norms, + row_norms, + activation_codes, + weight_codes, + activation_scale=scale, + weight_scale=1.0, + out_features=out_features, + in_features=in_features, + output_dtype=torch.bfloat16, + ) + assert (packed_result.float() - expected).abs().max().item() < 0.5 + + direct_result = matmul_int8_activations_packed_lowbit_fused_with_triton( + act_i8, + packed_w, + token_norms, + row_norms, + weight_codes, + weight_bits=2, + activation_scale=scale, + weight_scale=1.0, + out_features=out_features, + in_features=in_features, + output_dtype=torch.bfloat16, + ) + torch.testing.assert_close(direct_result, packed_result) + + +@pytest.mark.parametrize("rows", [17, 48]) +def test_w4a4_int_mm_path_supports_unaligned_rows_on_large_shape(rows): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + dequantize_packed_weight_with_triton, + matmul_packed_w4a4_with_int_mm, + quantize_activations_packed_w4_with_triton, + quantize_activations_with_triton, + ) + + fixture = _build_w4a4_triton_fixture(in_features=4096, out_features=2048) + x = torch.randn(rows, fixture["in_features"], device="cuda", dtype=torch.bfloat16) + packed_x, token_norms = quantize_activations_packed_w4_with_triton( + x, + rotation=fixture["rotation"], + codebook=fixture["activation_codebook"], + eps=1e-10, + ) + + actual = matmul_packed_w4a4_with_int_mm( + packed_x, + fixture["packed"], + token_norms, + fixture["row_norms"], + fixture["activation_codes"], + fixture["weight_codes"], + activation_scale=fixture["activation_scale"], + weight_scale=fixture["weight_scale"], + out_features=fixture["out_features"], + in_features=fixture["in_features"], + output_dtype=torch.bfloat16, + chunk_out_features=2048, + ) + + dequantized = dequantize_packed_weight_with_triton( + fixture["packed"], + fixture["row_norms"], + fixture["weight_codebook"], + bits=4, + out_features=fixture["out_features"], + in_features=fixture["in_features"], + ).to(torch.bfloat16) + quantized_values = quantize_activations_with_triton( + x, + rotation=fixture["rotation"], + codebook=fixture["activation_codebook"], + eps=1e-10, + ) + expected = torch.nn.functional.linear(quantized_values.to(torch.bfloat16), dequantized) + relative_error = (actual.float() - expected.float()).norm() / expected.float().norm() + assert relative_error <= 3e-2 + + +def test_adaln_triton_default_config_compiles_on_flux_modulation_shape(): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + dequantize_adaln_weight_with_triton, + matmul_packed_adaln_int4_with_triton, + quantize_adaln_weight_with_triton, + ) + + torch.manual_seed(3) + in_features, out_features = 3072, 18432 + weight = torch.randn(out_features, in_features, device="cuda", dtype=torch.float32) + packed, scales = quantize_adaln_weight_with_triton(weight, group_size=64) + x = torch.randn(1, in_features, device="cuda", dtype=torch.bfloat16) + + actual = matmul_packed_adaln_int4_with_triton( + x, + packed, + scales, + out_features=out_features, + in_features=in_features, + group_size=64, + ) + + reference_weight = dequantize_adaln_weight_with_triton( + packed, + scales, + out_features=out_features, + in_features=in_features, + group_size=64, + ).to(torch.bfloat16) + expected = torch.nn.functional.linear(x, reference_weight) + relative_error = (actual.float() - expected.float()).norm() / expected.float().norm() + assert relative_error <= 5e-3 + + +@pytest.mark.parametrize("rows", [1, 3, 16, 17]) +def test_adaln_triton_small_rows_match_dequantized_reference(rows): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + dequantize_adaln_weight_with_triton, + matmul_packed_adaln_int4_with_triton, + quantize_adaln_weight_with_triton, + ) + + torch.manual_seed(4) + in_features, out_features = 256, 192 + weight = torch.randn(out_features, in_features, device="cuda", dtype=torch.float32) + packed, scales = quantize_adaln_weight_with_triton(weight, group_size=64) + bias = torch.randn(out_features, device="cuda", dtype=torch.bfloat16) + x = torch.randn(rows, in_features, device="cuda", dtype=torch.bfloat16) + + actual = matmul_packed_adaln_int4_with_triton( + x, + packed, + scales, + out_features=out_features, + in_features=in_features, + group_size=64, + bias=bias, + ) + + reference_weight = dequantize_adaln_weight_with_triton( + packed, + scales, + out_features=out_features, + in_features=in_features, + group_size=64, + ).to(torch.bfloat16) + expected = torch.nn.functional.linear(x, reference_weight, bias) + relative_error = (actual.float() - expected.float()).norm() / expected.float().norm() + assert relative_error <= 5e-3 + + +@pytest.mark.parametrize("dim", [2048, 4096]) +def test_triton_cuda_activation_kernel_matches_reference_for_large_blocks(dim): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + torch.manual_seed(5) + x = torch.randn(3, dim, device="cuda", dtype=torch.float32) + rotation = RPBHRotation(dim=dim, seed=0, block_size="paper") + codebook = get_codebook(dim=dim, bits=4) + expected = quantize_activations(x, rotation=rotation, codebook=codebook, eps=1e-12) + actual = quantize_activations_kernel( + x, rotation=rotation, codebook=codebook, eps=1e-12, backend="triton_cuda" + ) + + assert rotation.block_size == dim + assert torch.allclose(actual, expected, atol=1e-4, rtol=1e-4) + + +@pytest.mark.parametrize("dim", [2048, 4096]) +def test_triton_packed_w4_activation_quant_matches_values_path_for_large_blocks(dim): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + quantize_activations_packed_w4_with_triton, + quantize_activations_with_triton, + ) + + torch.manual_seed(6) + rows = 3 + x = torch.randn(rows, dim, device="cuda", dtype=torch.bfloat16) + rotation = RPBHRotation(dim=dim, seed=0, block_size="paper") + codebook = get_codebook(dim=dim, bits=4) + + packed_x, token_norms = quantize_activations_packed_w4_with_triton( + x, rotation=rotation, codebook=codebook, eps=1e-10 + ) + values = quantize_activations_with_triton( + x, rotation=rotation, codebook=codebook, eps=1e-10 + ) + + codes = torch.empty(rows, dim, dtype=torch.uint8, device="cuda") + codes[:, 0::2] = packed_x & 15 + codes[:, 1::2] = packed_x >> 4 + centroids = codebook.centroids.to(device="cuda") + reconstructed = (centroids[codes.long()] * token_norms[:, None]).to(values.dtype) + + assert torch.equal(reconstructed, values) + + +def test_triton_activation_kernel_accepts_int32_permutation_constants(): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import quantize_activations_with_triton + + torch.manual_seed(8) + dim = 64 + x = torch.randn(4, dim, device="cuda", dtype=torch.float32) + rotation = RPBHRotation(dim=dim, seed=1, block_size=32) + codebook = get_codebook(dim=dim, bits=4) + expected = quantize_activations(x, rotation=rotation, codebook=codebook, eps=1e-12) + + constants = { + "permutation": rotation.permutation.to(device="cuda", dtype=torch.int32), + "signs": rotation.signs.to(device="cuda"), + "centroids": codebook.centroids.to(device="cuda"), + "boundaries": codebook.boundaries.to(device="cuda"), + } + actual = quantize_activations_with_triton( + x, rotation=rotation, codebook=codebook, eps=1e-12, constant_tensors=constants + ) + + assert torch.allclose(actual, expected, atol=1e-5, rtol=1e-5) + + +def test_triton_packed_matmul_recovers_from_shared_memory_overflow(): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import matmul_packed_weight_with_triton + + torch.manual_seed(9) + in_features, out_features, rows = 4096, 512, 64 + codebook = get_codebook(dim=in_features, bits=4) + indices = ( + torch.arange(out_features * in_features, dtype=torch.int64) % 16 + ).to(torch.uint8) + packed = pack_lowbit(indices, bits=4).to("cuda") + row_norms = torch.linspace(0.5, 1.5, out_features, dtype=torch.bfloat16, device="cuda") + x = torch.randn(rows, in_features, device="cuda", dtype=torch.bfloat16) + + weight = row_norms.float().cpu()[:, None] * codebook.centroids[ + indices.reshape(out_features, in_features).to(torch.long) + ] + expected = torch.nn.functional.linear(x.float().cpu(), weight) + + actual = matmul_packed_weight_with_triton( + x, + packed, + row_norms, + codebook, + bits=4, + out_features=out_features, + in_features=in_features, + block_m=64, + block_n=256, + block_k=128, + num_warps=8, + ) + + assert torch.allclose(actual.float().cpu(), expected, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize("rows", [32, 64, 400]) +def test_w4a4_fused_triton_matmul_matches_dequantized_reference(rows): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + dequantize_packed_weight_with_triton, + matmul_packed_w4a4_fused_with_triton, + quantize_activations_packed_w4_with_triton, + quantize_activations_with_triton, + ) + + fixture = _build_w4a4_triton_fixture(in_features=512, out_features=512) + x = torch.randn(rows, fixture["in_features"], device="cuda", dtype=torch.bfloat16) + packed_x, token_norms = quantize_activations_packed_w4_with_triton( + x, + rotation=fixture["rotation"], + codebook=fixture["activation_codebook"], + eps=1e-10, + ) + bias = torch.randn(fixture["out_features"], device="cuda", dtype=torch.bfloat16) + + actual = matmul_packed_w4a4_fused_with_triton( + packed_x, + fixture["packed"], + token_norms, + fixture["row_norms"], + fixture["activation_codes"], + fixture["weight_codes"], + activation_scale=fixture["activation_scale"], + weight_scale=fixture["weight_scale"], + out_features=fixture["out_features"], + in_features=fixture["in_features"], + bias=bias, + output_dtype=torch.bfloat16, + ) + + dequantized = dequantize_packed_weight_with_triton( + fixture["packed"], + fixture["row_norms"], + fixture["weight_codebook"], + bits=4, + out_features=fixture["out_features"], + in_features=fixture["in_features"], + ).to(torch.bfloat16) + quantized_values = quantize_activations_with_triton( + x, + rotation=fixture["rotation"], + codebook=fixture["activation_codebook"], + eps=1e-10, + ) + expected = torch.nn.functional.linear( + quantized_values.to(torch.bfloat16), dequantized, bias + ) + relative_error = (actual.float() - expected.float()).norm() / expected.float().norm() + assert relative_error <= 3e-2 + + +def test_w4a4_int_mm_decoded_weight_cache_matches_per_forward_decode(): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + from orbitquant.kernels.triton_cuda import ( + decode_packed_w4_weight_to_int8, + matmul_packed_w4a4_with_int_mm, + quantize_activations_packed_w4_with_triton, + ) + + fixture = _build_w4a4_triton_fixture(in_features=512, out_features=512) + x = torch.randn(64, fixture["in_features"], device="cuda", dtype=torch.bfloat16) + packed_x, token_norms = quantize_activations_packed_w4_with_triton( + x, + rotation=fixture["rotation"], + codebook=fixture["activation_codebook"], + eps=1e-10, + ) + shared_kwargs = dict( + activation_scale=fixture["activation_scale"], + weight_scale=fixture["weight_scale"], + out_features=fixture["out_features"], + in_features=fixture["in_features"], + output_dtype=torch.bfloat16, + ) + + baseline = matmul_packed_w4a4_with_int_mm( + packed_x, + fixture["packed"], + token_norms, + fixture["row_norms"], + fixture["activation_codes"], + fixture["weight_codes"], + **shared_kwargs, + ) + decoded = decode_packed_w4_weight_to_int8( + fixture["packed"], + fixture["weight_codes"], + out_features=fixture["out_features"], + in_features=fixture["in_features"], + ) + cached = matmul_packed_w4a4_with_int_mm( + packed_x, + fixture["packed"], + token_norms, + fixture["row_norms"], + fixture["activation_codes"], + fixture["weight_codes"], + decoded_weight=decoded, + **shared_kwargs, + ) + + assert torch.equal(baseline, cached) + + +def test_config_w4a4_int8_weight_cache_flag_roundtrips(): + config = OrbitQuantConfig(w4a4_int8_weight_cache=True) + assert config.w4a4_int8_weight_cache is True + restored = OrbitQuantConfig.from_dict(config.to_dict()) + assert restored.w4a4_int8_weight_cache is True + assert OrbitQuantConfig().w4a4_int8_weight_cache is False diff --git a/tests/test_model_quantization.py b/tests/test_model_quantization.py index 147cb71..91c2484 100644 --- a/tests/test_model_quantization.py +++ b/tests/test_model_quantization.py @@ -39,7 +39,12 @@ def test_quantize_linear_modules_replaces_orbit_and_adaln_targets_only(): assert summary.adaln_modules == ["transformer_blocks.0.modulation"] assert summary.skipped_modules == ["proj_out"] assert summary.quantization_device == ("cuda" if torch.cuda.is_available() else "cpu") - assert summary.weight_quantization_backend in {"torch_reference", "triton_cuda"} + assert summary.weight_quantization_backend in { + "torch_reference", + "triton_cuda", + "triton_rocm", + "triton_xpu", + } assert summary.quantization_staging_mode == "streaming" assert summary.synchronize_per_module is False assert summary.elapsed_seconds >= 0.0 @@ -143,6 +148,19 @@ def test_quantize_linear_modules_fails_loud_for_unavailable_cuda_device(monkeypa raise AssertionError("unavailable CUDA quantization device was accepted") +def test_quantize_linear_modules_fails_loud_for_unavailable_xpu_device(monkeypatch): + model = TinyPipelineTransformer() + config = OrbitQuantConfig(block_size=8) + monkeypatch.setattr(torch.xpu, "is_available", lambda: False) + + try: + quantize_linear_modules(model, config, quantization_device="xpu") + except RuntimeError as exc: + assert "XPU quantization device requested" in str(exc) + else: + raise AssertionError("unavailable XPU quantization device was accepted") + + def test_quantize_linear_modules_fails_loud_for_cuda_without_triton(monkeypatch): model = TinyPipelineTransformer() config = OrbitQuantConfig(block_size=8) @@ -150,22 +168,60 @@ def test_quantize_linear_modules_fails_loud_for_cuda_without_triton(monkeypatch) monkeypatch.setattr( modeling_module, "available_backends", - lambda: {"cpu": True, "mps": False, "triton_cuda": False}, + lambda: { + "cpu": True, + "mps": False, + "triton_cuda": False, + "triton_rocm": False, + }, ) try: quantize_linear_modules(model, config, quantization_device="cuda") except RuntimeError as exc: - assert "Triton CUDA backend" in str(exc) + assert "triton_cuda backend" in str(exc) else: raise AssertionError("CUDA quantization without Triton was accepted") -def test_auto_quantization_device_prefers_cuda_and_falls_back_to_cpu(monkeypatch): +def test_weight_quantization_backend_labels_hip_as_rocm(monkeypatch): + monkeypatch.setattr(modeling_module.torch.version, "hip", "7.2", raising=False) + monkeypatch.setattr( + modeling_module, + "available_backends", + lambda: { + "cpu": True, + "mps": False, + "triton_cuda": False, + "triton_rocm": True, + }, + ) + + assert modeling_module._weight_quantization_backend(torch.device("cuda")) == "triton_rocm" + + +def test_weight_quantization_backend_labels_explicit_xpu(monkeypatch): + monkeypatch.setattr( + modeling_module, + "available_backends", + lambda: { + "cpu": True, + "mps": False, + "triton_cuda": False, + "triton_rocm": False, + "triton_xpu": True, + }, + ) + + assert modeling_module._weight_quantization_backend(torch.device("xpu")) == "triton_xpu" + + +def test_auto_quantization_device_prefers_verified_cuda_and_keeps_xpu_explicit(monkeypatch): monkeypatch.setattr(torch.cuda, "is_available", lambda: True) assert modeling_module._quantization_device("auto") == torch.device("cuda") monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(torch.xpu, "is_available", lambda: True) assert modeling_module._quantization_device("auto") == torch.device("cpu") diff --git a/tests/test_native_kernel_package.py b/tests/test_native_kernel_package.py index 82616e9..363ccb7 100644 --- a/tests/test_native_kernel_package.py +++ b/tests/test_native_kernel_package.py @@ -38,6 +38,8 @@ def test_kernel_builder_manifest_targets_cpu_cuda_and_metal() -> None: assert kernels["packed_matmul_cpu"]["src"] == [ "orbitquant_packed_matmul_cpu/cpu_isa.cpp", "orbitquant_packed_matmul_cpu/cpu_kernel_args.h", + "orbitquant_packed_matmul_cpu/cpu_pool.cpp", + "orbitquant_packed_matmul_cpu/cpu_pool.h", "orbitquant_packed_matmul_cpu/cpu_threads.cpp", "orbitquant_packed_matmul_cpu/cpu_threads.h", "orbitquant_packed_matmul_cpu/packed_adaln_cpu.cpp", diff --git a/tests/test_orbit_linear.py b/tests/test_orbit_linear.py index 8d5a38f..1787f75 100644 --- a/tests/test_orbit_linear.py +++ b/tests/test_orbit_linear.py @@ -423,10 +423,58 @@ def fail_activation_quantization(*args, **kwargs): fail_activation_quantization, ) - with pytest.raises(RuntimeError, match="requires CUDA input tensors"): + with pytest.raises(RuntimeError, match="requires CUDA, HIP, or XPU input tensors"): quantized(x) +def test_orbit_linear_auto_fused_on_hip_stays_disabled_until_hardware_proof(monkeypatch): + source = torch.nn.Linear(16, 7) + quantized = OrbitQuantLinear.from_linear( + source, + config=OrbitQuantConfig(block_size=8), + module_name="block.ff.linear", + ) + fake_hip_input = SimpleNamespace(device=torch.device("cuda")) + + monkeypatch.setattr(layers_module, "_torch_uses_hip", lambda: True) + monkeypatch.setattr( + layers_module, + "_native_packed_matmul_load_error", + lambda: (_ for _ in ()).throw(AssertionError("CUDA native probe must be skipped on HIP")), + ) + monkeypatch.setattr(layers_module, "_triton_packed_matmul_import_error", lambda: None) + + with pytest.raises(RuntimeError, match="remains experimental"): + quantized._resolve_auto_fused_runtime(fake_hip_input) + + +def test_orbit_linear_auto_fused_on_xpu_stays_disabled_until_hardware_proof(): + source = torch.nn.Linear(16, 7) + quantized = OrbitQuantLinear.from_linear( + source, + config=OrbitQuantConfig(block_size=8), + module_name="block.ff.linear", + ) + fake_xpu_input = SimpleNamespace(device=torch.device("xpu")) + + with pytest.raises(RuntimeError, match="auto_fused XPU dispatch remains experimental"): + quantized._resolve_auto_fused_runtime(fake_xpu_input) + + +def test_orbit_linear_explicit_native_runtime_rejects_hip(monkeypatch): + source = torch.nn.Linear(16, 7) + quantized = OrbitQuantLinear.from_linear( + source, + config=OrbitQuantConfig(block_size=8, runtime_mode="native_packed_matmul"), + module_name="block.ff.linear", + ) + fake_hip_input = SimpleNamespace(device=torch.device("cuda")) + monkeypatch.setattr(layers_module, "_torch_uses_hip", lambda: True) + + with pytest.raises(RuntimeError, match="not HIP"): + quantized._validate_native_packed_matmul_input(fake_hip_input) + + def test_orbit_linear_triton_packed_matmul_runtime_avoids_weight_dequant_cache(monkeypatch): torch.manual_seed(5) source = torch.nn.Linear(16, 7) diff --git a/tests/test_paper_reference_oracle.py b/tests/test_paper_reference_oracle.py index 7d3859d..2568a73 100644 --- a/tests/test_paper_reference_oracle.py +++ b/tests/test_paper_reference_oracle.py @@ -7,6 +7,7 @@ from orbitquant.adaln import RTNInt4Linear from orbitquant.codebooks import get_codebook from orbitquant.config import OrbitQuantConfig +from orbitquant.kernels.executorch_vulkan import ExecuTorchVulkanW4A4Linear from orbitquant.layers import OrbitQuantLinear from orbitquant.rotations import RPBHRotation @@ -202,6 +203,199 @@ def test_reference_runtime_matches_independent_algorithm_1_oracle( torch.testing.assert_close(actual[0, 0], expected_zero_token, atol=0, rtol=0) +def test_executorch_vulkan_w4a4_export_op_matches_independent_algorithm_1_oracle(): + torch.manual_seed(211) + source = torch.nn.Linear(24, 7, bias=True) + activations = torch.randn(2, 3, 24, dtype=torch.float32) + activations[0, 0].zero_() + config = OrbitQuantConfig( + weight_bits=4, + activation_bits=4, + rotation_seed=31, + block_size="paper", + runtime_mode="dequant_bf16", + ) + layer = OrbitQuantLinear.from_linear( + source, + config=config, + module_name="transformer_blocks.0.attn.to_q", + ) + export_layer = ExecuTorchVulkanW4A4Linear(layer) + + row_rotation = _dense_paper_row_rotation(layer.rotation) + rotated_activations = activations @ row_rotation + token_norms = rotated_activations.norm(dim=-1, keepdim=True) + activation_directions = rotated_activations / (token_norms + config.activation_eps) + _, quantized_activation_directions = _nearest_centroid( + activation_directions, + layer.activation_codebook.centroids.to(torch.float32), + ) + quantized_activations = token_norms * quantized_activation_directions + + weight_indices = _unpack_little_endian_indices( + layer.packed_weight_indices, + bits=4, + count=source.out_features * source.in_features, + ).reshape(source.out_features, source.in_features) + weight_centroids = layer.weight_codebook.centroids.to(torch.float32) + quantized_weight = layer.row_norms.to(torch.float32)[:, None] * weight_centroids[weight_indices] + expected = F.linear( + quantized_activations, + quantized_weight, + source.bias.detach().to(torch.float32), + ) + + actual = export_layer(activations) + + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5) + torch.testing.assert_close(actual[0, 0], source.bias.detach().to(torch.float32), atol=0, rtol=0) + full_weight_numel = source.in_features * source.out_features + assert export_layer.packed_weight.numel() == full_weight_numel // 2 + assert not any( + tensor.is_floating_point() and tensor.numel() == full_weight_numel + for tensor in export_layer.state_dict().values() + ) + + +def test_triton_rocm_w4a4_matches_independent_algorithm_1_oracle(): + if not torch.cuda.is_available() or not getattr(torch.version, "hip", None): + pytest.skip("a PyTorch ROCm runtime is required") + from orbitquant.kernels import available_backends + + if not available_backends()["triton_rocm"]: + pytest.skip("the ROCm-compatible Triton package is not importable") + + torch.manual_seed(223) + source = torch.nn.Linear(24, 7, bias=True) + activations = torch.randn(2, 3, 24, dtype=torch.bfloat16) + activations[0, 0].zero_() + config = OrbitQuantConfig( + weight_bits=4, + activation_bits=4, + rotation_seed=37, + block_size="paper", + runtime_mode="triton_packed_matmul", + activation_kernel_backend="triton_rocm", + ) + layer = OrbitQuantLinear.from_linear( + source, + config=config, + module_name="transformer_blocks.0.attn.to_q", + ) + + row_rotation = _dense_paper_row_rotation(layer.rotation) + rotated_weight = source.weight.detach().to(torch.float32) @ row_rotation + raw_row_norms = rotated_weight.norm(dim=-1) + weight_directions = rotated_weight / raw_row_norms.clamp_min(config.activation_eps)[:, None] + _, quantized_weight_directions = _nearest_centroid( + weight_directions, + layer.weight_codebook.centroids.to(torch.float32), + ) + quantized_weight = ( + raw_row_norms.to(torch.bfloat16).float()[:, None] * quantized_weight_directions + ).to(torch.bfloat16) + + rotated_activations = activations.float() @ row_rotation + token_norms = rotated_activations.norm(dim=-1, keepdim=True) + activation_directions = rotated_activations / (token_norms + config.activation_eps) + _, quantized_activation_directions = _nearest_centroid( + activation_directions, + layer.activation_codebook.centroids.to(torch.float32), + ) + quantized_activations = (token_norms * quantized_activation_directions).to(torch.bfloat16) + expected = F.linear( + quantized_activations, + quantized_weight, + source.bias.detach().to(torch.bfloat16), + ) + + layer = layer.to("cuda") + actual = layer(activations.to("cuda")).cpu() + error = actual.float() - expected.float() + relative_rmse = error.square().mean().sqrt() / expected.float().square().mean().sqrt() + + torch.testing.assert_close(actual, expected, atol=0.125, rtol=3e-2) + assert float(relative_rmse) < 1e-2 + assert layer.last_effective_runtime_mode == "triton_packed_matmul" + assert layer.last_activation_kernel_backend == "triton_rocm" + assert layer._dequantized_weight_cache is None + full_weight_numel = source.in_features * source.out_features + assert not any( + tensor.is_floating_point() and tensor.numel() == full_weight_numel + for tensor in layer.state_dict().values() + ) + + +def test_triton_xpu_w4a4_matches_independent_algorithm_1_oracle(): + xpu = getattr(torch, "xpu", None) + if xpu is None or not xpu.is_available(): + pytest.skip("a PyTorch XPU runtime is required") + from orbitquant.kernels import available_backends + + if not available_backends()["triton_xpu"]: + pytest.skip("the Intel XPU Triton package is not importable") + + torch.manual_seed(227) + source = torch.nn.Linear(24, 7, bias=True) + activations = torch.randn(2, 3, 24, dtype=torch.bfloat16) + activations[0, 0].zero_() + config = OrbitQuantConfig( + weight_bits=4, + activation_bits=4, + rotation_seed=41, + block_size="paper", + runtime_mode="triton_packed_matmul", + activation_kernel_backend="triton_xpu", + ) + layer = OrbitQuantLinear.from_linear( + source, + config=config, + module_name="transformer_blocks.0.attn.to_q", + ) + + row_rotation = _dense_paper_row_rotation(layer.rotation) + rotated_weight = source.weight.detach().to(torch.float32) @ row_rotation + raw_row_norms = rotated_weight.norm(dim=-1) + weight_directions = rotated_weight / raw_row_norms.clamp_min(config.activation_eps)[:, None] + _, quantized_weight_directions = _nearest_centroid( + weight_directions, + layer.weight_codebook.centroids.to(torch.float32), + ) + quantized_weight = ( + raw_row_norms.to(torch.bfloat16).float()[:, None] * quantized_weight_directions + ).to(torch.bfloat16) + + rotated_activations = activations.float() @ row_rotation + token_norms = rotated_activations.norm(dim=-1, keepdim=True) + activation_directions = rotated_activations / (token_norms + config.activation_eps) + _, quantized_activation_directions = _nearest_centroid( + activation_directions, + layer.activation_codebook.centroids.to(torch.float32), + ) + quantized_activations = (token_norms * quantized_activation_directions).to(torch.bfloat16) + expected = F.linear( + quantized_activations, + quantized_weight, + source.bias.detach().to(torch.bfloat16), + ) + + layer = layer.to("xpu") + actual = layer(activations.to("xpu")).cpu() + error = actual.float() - expected.float() + relative_rmse = error.square().mean().sqrt() / expected.float().square().mean().sqrt() + + torch.testing.assert_close(actual, expected, atol=0.125, rtol=3e-2) + assert float(relative_rmse) < 1e-2 + assert layer.last_effective_runtime_mode == "triton_packed_matmul" + assert layer.last_activation_kernel_backend == "triton_xpu" + assert layer._dequantized_weight_cache is None + full_weight_numel = source.in_features * source.out_features + assert not any( + tensor.is_floating_point() and tensor.numel() == full_weight_numel + for tensor in layer.state_dict().values() + ) + + @pytest.mark.parametrize("bits", [2, 3, 4, 6]) def test_lloyd_max_v2_satisfies_independent_beta_centroid_condition(bits: int): dim = 64 diff --git a/tests/test_torch_compile.py b/tests/test_torch_compile.py new file mode 100644 index 0000000..7ea5d53 --- /dev/null +++ b/tests/test_torch_compile.py @@ -0,0 +1,90 @@ +import pytest +import torch + +from orbitquant.adaln import RTNInt4Linear +from orbitquant.config import OrbitQuantConfig +from orbitquant.layers import OrbitQuantLinear + + +def _build_layers() -> tuple[OrbitQuantLinear, RTNInt4Linear]: + torch.manual_seed(0) + config = OrbitQuantConfig( + weight_bits=4, + activation_bits=4, + block_size=8, + runtime_mode="dequant_bf16", + ) + orbit = OrbitQuantLinear.from_linear( + torch.nn.Linear(32, 24, bias=True), + config=config, + module_name="orbit", + ) + adaln = RTNInt4Linear.from_linear( + torch.nn.Linear(32, 16, bias=True), + config=config, + module_name="adaln", + ) + return orbit, adaln + + +def test_orbit_linear_compiles_fullgraph_and_matches_eager(): + orbit, _ = _build_layers() + x = torch.randn(3, 32) + eager = orbit(x) + + compiled = torch.compile(orbit, fullgraph=True) + compiled_out = compiled(x) + + assert compiled_out.shape == eager.shape + assert compiled_out.dtype == eager.dtype + assert torch.equal(compiled_out, eager) + + +def test_adaln_linear_compiles_fullgraph_and_matches_eager(): + _, adaln = _build_layers() + x = torch.randn(3, 32) + eager = adaln(x) + + compiled = torch.compile(adaln, fullgraph=True) + compiled_out = compiled(x) + + assert compiled_out.shape == eager.shape + assert compiled_out.dtype == torch.bfloat16 + assert torch.equal(compiled_out, eager) + + +def test_compiled_stack_of_quantized_layers_matches_eager(): + orbit, adaln = _build_layers() + + class Stack(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.orbit = orbit + self.adaln = adaln + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.adaln(x).float().sum() + self.orbit(x).float().sum() + + model = Stack() + x = torch.randn(2, 32) + eager = model(x) + compiled = torch.compile(model, fullgraph=True) + out = compiled(x) + assert torch.allclose(out, eager) + + +def test_compiled_handle_survives_repeated_calls(): + orbit, _ = _build_layers() + compiled = torch.compile(orbit, fullgraph=True) + x = torch.randn(4, 32) + first = compiled(x) + second = compiled(x) + assert torch.equal(first, second) + + +@pytest.mark.parametrize("rows", [1, 5]) +def test_compiled_orbit_linear_handles_leading_batch_dims(rows): + orbit, _ = _build_layers() + compiled = torch.compile(orbit, fullgraph=True) + x = torch.randn(rows, 2, 32) + assert torch.equal(compiled(x), orbit(x))