Repository navigation
Add cross-platform packed inference paths - #5
Merged
Merged
Conversation
- Pad torch._int_mm rows to a multiple of 32: cuBLASLt INT8 rejects every m <= 16 and shape-dependent unaligned m (verified m=17..48 failures on torch 2.13.0+cu130/RTX 4090); route rows < 32 to the fused native W4A4 kernel, which also skips the per-forward INT8 weight decode there. - Rework the packed AdaLN INT4 kernel around group-aligned scale loads: the (block_k, block_n) scale gather made the default tile need 172 KB of shared memory and fail to compile on Triton 3.7 (101 KB limit). Default block_k drops to 64; M=1 on 3072->18432 goes 0.609 ms -> 0.257 ms and the default configuration compiles again. Shared-memory overflows now retry with smaller tiles instead of raising. - Extend the fused RPBH activation kernel to 2048/4096 blocks (bit-exact, 0.380 ms -> 0.199 ms for dim 4096 x 4096 rows) and add the two missing FWHT stages to the packed-W4 variant. - Store the rotation permutation as int32 everywhere on the hot path (Triton, native CUDA kernels, layer buffers): index traffic halves and dim-4096 activation quantization measures another -26%. - Auto-select measured packed-matmul tiles when callers keep the legacy (64, 64, 128, nw4) defaults: 4.503 ms -> 1.875 ms on 4096x4096 W4 prefill. - New regression tests: int_mm odd-row sweeps (1..77), AdaLN default-config compile on the FLUX modulation shape, large-block activation equivalence, int32 permutation acceptance, shared-memory overflow recovery.
Previously every non-W4 width fell back to the scalar per-element loop on
all ISAs. New decoders (16 values per iteration on AVX-512, 8 on AVX2,
generic-Bits template on NEON) keep the exact scalar-reference semantics:
- AVX-512: W2 via byte-replicate + srlv + 4-entry permutexvar LUT; W3 via
shuffled 16-bit windows + broadcast 8-entry LUT; W6 via 12-byte windows +
permutex2var over four 16-float tables with a bit-5 blend.
- AVX2: W2 via permutevar8x32; W6 via i32gather; W3 keeps scalar.
- NEON: shared generic-Bits row template (scalar table decode, vector FMA).
Dispatch admits the SIMD paths only for byte-aligned rows (W2/W6:
in %% 4 == 0, W3: in %% 8 == 0) and keeps scalar otherwise.
Measured 32 BF16 rows, six pinned Zen 4 (EPYC 9684X) cores, dim 3072:
W2 59.4 -> 2.14 ms (AVX-512), W3 61.0 -> 6.88 ms, W6 73.1 -> 2.60 ms
against the previous scalar execution; W4 is unchanged or faster. On an
M2 Max aarch64 container (dim 1536, one thread): W2 95.2 -> 9.1 ms,
W3 93.3 -> 9.8 ms, W6 92.6 -> 9.8 ms.
Tests: the bits-parametrized dequantized-reference suites now exercise the
new paths, plus a new scalar-vs-ISA equality sweep over bits {2,3,6} and
aligned/unaligned dims, and an int32-permutation acceptance case for the
CUDA activation kernels.
…dation Every native CPU matmul, activation-quantization, and AdaLN call previously spawned and joined fresh std::thread workers. The three parallel helpers now hand their unchanged contiguous ranges to a lazily created persistent pool (plain C++17, Stable-ABI-clean, intentionally leaked at exit so interpreter teardown cannot deadlock). Range ownership is unchanged, so multi-threaded outputs stay bitwise identical to single-threaded runs (asserted for bits 2/3/4/6 at 32x3072). Measured on six pinned Zen 4 cores, 32 BF16 rows: W4 AVX-512 packed matmul 0.498 -> 0.304 ms at dim 1536 and 1.257 -> 1.045 ms at dim 3072 versus the spawn-per-call baseline; ORBITQUANT_CPU_THREADS=1 bypasses the pool entirely. quantize_activations_cpu also re-validated the immutable permutation/sign buffers element-by-element on every forward; a small fingerprint cache now runs that O(dim) scan once per distinct buffer while keeping the out-of-range/invalid-sign errors on first use.
… early FWHT stages Family-level CPUID gates replace the exact-model checks (Zen3 0x19 AVX2 tuning; Zen4/5 0x19/0x1a + avx512bf16), so EPYC 9654/9684X now take the tuned path. Matmul decodes each packed column once into a thread-local buffer when rows >= 16 (AVX-512/AVX2/NEON) and NEON gets a W4 low-bit path. Activation quantization runs FWHT widths below the vector register in-register (AVX-512 widths 1-8, AVX2 1-4, NEON 1-2), reuses thread-local scratch, and accepts int32 permutations end to end. Verified: EPYC 9654 84 passed; NEON container 79 passed; evidence in plans/evidence-phase4-20260712.json.
Double-buffered packed-weight stages decode from shared memory while the next k-tile streams in. Measured on sm_89 (512x3072->3072, hot median): W2 2.726 -> 0.400 ms (6.8x). W4/W6 range from parity to a 10x regression depending on shape, so the default gate is W2-only; ORBITQUANT_MMA64_FORCE_PIPELINE=1 extends it to W4/W6 for experiments and ORBITQUANT_MMA64_DISABLE_PIPELINE=1 forces the legacy kernel.
The fused activation shader is compiled per block-size bucket instead of
always reserving the 4096-float worst case (1.15x at block 512,
4096x3072 fp16); the row-norm reduction uses simd_sum plus one cross-warp
pass; the small-rows Metal matmul multiplies by the row norm after the
simd reduction instead of inside the k-loop. The native benchmark now
sweeps rows {1,8,512,4096} by default, accepts --dtype, and headlines the
hot-loop median. run_mps_kernel_checks.sh no longer trips 'set -u' on
bash 3.2 when the runtime-args array is empty.
…orch.compile
A fused int8 Triton kernel decodes packed nibbles of both operands
in-register and applies the norm/scale epilogue in one launch (needs an
explicit out_dtype=tl.int32 for the int8 tl.dot on Triton 3.7). Layer
forwards route to it for 32 <= rows < 512 where it wins 1.07-1.29x over
the int_mm path on sm_89; the int_mm chunking heuristic is now M-aware
and an opt-in w4a4_int8_weight_cache config flag keeps the decoded INT8
weight resident (skipped under accelerate offload). Per-forward overhead
drops via cached runtime probes, per-device bias/weight caches, memoized
available_backends() (with a reset hook and autouse test fixture), cached
quantizer classification, and bits {2,4,8} fast paths in bitpack.
Forwards are wrapped in an orbitquant::packed_linear_forward custom op
with a fake impl, so torch.compile(fullgraph=True) matches eager
bit-exactly; backend_capabilities() is table-driven (byte-identical over
a 1024-combination differential check).
…yped tests/test_cli.py (130 KB) and tests/test_hub.py (103 KB) become 13 focused files plus shared hub_helpers; the suite is xdist-safe (536 passed with -n 4). A golden capability test locks the backend_capabilities() contract, docs/paper-methodology-audit.md matches docs/kernel-audit.md backend rows, the package ships py.typed, and the root CLAUDE.md records the repo conventions for coding agents.
The W2/W3/W6 table adds W4, records the buffered per-column decode, and carries a single consistent protocol (six unpinned worker threads, 21 hot-loop-median iterations, EPYC 9354 cpuset; one thread on the M2 Max container), with the faster 9684X AVX-512 run noted alongside. CARD.md documents the benchmark's rows sweep, --dtype flag, and median-first output fields; README lists torch.compile(fullgraph=True) support.
iamwavecut
marked this pull request as ready for review
July 12, 2026 18:14
…rch builds Cold-L2 A/B on RTX 4060 Ti, A40, and RTX 4090 (a real model's layers evict each other, so per-layer weights are never L2-resident) shows the pipelined kernel wins 1.35-1.46x for W2/W4/W6 whenever the grid fits one wave and regresses up to 15% once the device is oversubscribed; the earlier '10x W4 regression' was a hot-L2 microbenchmark artifact. The dispatch gate is now grid_blocks <= multiProcessorCount for every supported width, with ORBITQUANT_MMA64_FORCE_PIPELINE / ORBITQUANT_MMA64_DISABLE_PIPELINE overrides. All four WMMA kernels also compile to empty stubs below sm_80 and the host requires compute capability >= 8 before dispatching mma64: unpinned multi-arch builds on nvcc 12.x previously failed on bf16/int8 fragment instantiations for sm_70 (nvcc 13 masked this by dropping sm_70).
A 144-configuration grid on A40 (sm_86) and RTX 4090 (sm_89) shows the fixed (64,128,128,8) tile is never optimal: large batches want a taller M tile and narrower K tile. The kernel now defaults to a measured rows-keyed table ((64,128,128,4) below 512 rows; (128,128,64,4) at 512+), +10-28% at 256 rows and +30-63% at 1024. With the tuned table the fused kernel beats the int_mm path through 2048 rows on the 4090 (1.05-1.45x), so _W4A4_FUSED_MAX_ROWS rises 512 -> 2048. README documents torch.compile(mode="reduce-overhead") for CUDA-graph capture: measured 1.6-5.5x at 1-32 decode rows where launch overhead dominates.
Wave-3 experiments on RTX 4090: feeding the fused GEMM INT8-surrogate activations directly (skipping the pack->unpack round trip) is 13-20% faster than the packed-activation kernel across 64-2048 rows, so the W4A4 fused branch now prefers the native INT8 activation quantizer and keeps packed activations as the fallback. The int8-activation kernel takes constexpr weight_bits and also decodes 2-bit weights, and a packed-activation W2A4 kernel covers the no-native fallback; W2A4 layers dispatch through a new fused window (32 <= rows < 2048) instead of the bf16 mma64 path, at W4A4-fused speed (0.089-0.303 ms on 4096^2 versus 0.21+ ms for the pipelined mma64 on the same shapes). Numerical checks pin the packed and direct variants together (torch.equal-level) and against int8 emulation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Cross-platform packed inference paths plus a hardware-verified performance pass over them; version bumped to 0.5.0.
Packed inference paths (original scope)
auto_fusedruns exact packed activation, low-bit matmul, and INT4 group-64 AdaLN through scalar, AVX2/FMA, AVX-512/BF16, and ARM64 NEON runtime dispatchPerformance pass (this revision)
torch._int_mmshort-sequence crashes (row padding + rerouting), fixed AdaLN Triton compilation on Triton 3.7 with a faster small-M path (M=1 0.609 → 0.257 ms), auto-selected packed-matmul tiles (2.4x on 4096² prefill), fused packed-W4A4 int8 Triton GEMM with in-kernel nibble decode (1.07–1.29x over the int_mm path for 32 ≤ rows < 512 on sm_89), M-aware output chunking, opt-in resident INT8 weight cache (w4a4_int8_weight_cache), cp.async-pipelined mma64 for W2 (2.73 → 0.40 ms at 512×3072², 6.8x)simd_sumreductions, small-rows norm hoist; benchmark rows sweep +--dtype+ median-first outputtorch.compile(fullgraph=True)support through a registered custom op with a fake impl; compiled output matches eager bit-exactlybackend_capabilities()(byte-identical over a 1024-combination differential test)test_cli/test_hubsplit into 13 focused files, xdist-safe suite (536 passed in ~9 s at-n 4), golden capability contract test,py.typed, rootCLAUDE.mdClaim boundary and known limitations
ORBITQUANT_MMA64_FORCE_PIPELINE=1re-enables them for experiments,ORBITQUANT_MMA64_DISABLE_PIPELINE=1forces the legacy kernelkernel-builder create-pyproject+setup.py build_kernellocal path is the working alternativeValidation
ruff check .clean;pytest -n 4: 536 passed, 76 skipped;run_paper_methodology_checks.shandrun_hf_compat_checks.sh --mode allgreendocs/kernel-audit.md)run_mps_kernel_checks.shpasses end to end on stock bash 3.2b8cb1a5: CPU tests, Windows native CPU, and manylinux x86_64 native CPU pass; MSVC Vulkan source build compiles (see checks)