698f8c13d9db6a2d3c2aa6eb4d55532c2fdaadaa
4
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 tzhouamandClaude Opus 5
|
6d7b2bfc1c |
cpu_int4: repack weights in n-blocks, to bound the load-time peak
I0 of the reviewed plan: the peak, not the settled figure, is what an edge
device has to survive. Measured earlier this session, W4A16 loads at 5722-5837
MB and settles near 4480 after the allocator is asked for its arenas back -- a
1.2-1.3 GB gap, and a device that cannot hold the peak never reaches the
release.
A large part of that gap is in our own repack path. unpack_nibbles expands
4-bit codes to int32, which is 4 bytes per weight and eight times the packed
form, and the kernel's packer wants int32. Done for a whole matrix that
transient is the largest object in the process:
gate_up 13312 x 2048 104 MiB
down 2048 x 6656 52 MiB
q_k_v 3072 x 2048 24 MiB
embedding 131072 x 2048 1024 MiB <- tied, so it is also the head
against 24 MiB of packed weights for an entire layer. The embedding alone is a
~1 GB spike, which is about the size of the whole peak-vs-settled gap.
The packed layout is [N/64][K][32 bytes], so n-blocks are independent and the
expansion can cover one chunk at a time. Verified byte-identical to packing the
whole matrix at 13312x2048, 2048x6656, 3072x2048 and 131072x2048, plus a ragged
last chunk. The int32 transient drops from 1024 MiB to 4 MiB for the embedding
and from 104 to 4 MiB for gate_up. to_split_halves is chunked the same way for
the fallback path.
Not yet claimed: the end-to-end peak reduction. The arithmetic says ~1 GB, but
this study has been wrong often enough about memory that the number goes in
only after RssAnon is sampled through a real load. Tests: 21 passed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
||
|
tzhouamandClaude Opus 5
|
f3629320e7 |
cpu_int4: dequantize prefill weights from the tinygemm layout, not a copy
The W4A16 path kept two packed weight layouts because the tinygemm layout was treated as opaque: the prefill path needed something row-major to dequantize from, so `_finalize_int4` built `weight_halves` plus cloned scale/zero beside the kernel's own buffer. That duplicate is the 740 MB the report attributes to "the dual weight layout". The layout is not opaque. Measured on this host by packing an index pattern and reading back where each nibble landed: it is [N/64][K][32 bytes] -- K is the middle axis, so a K range is contiguous inside each n-block -- and within every 32-byte run, byte d carries output channel d in its low nibble and d+32 in its high one, the same permutation for every k and every block. Decoding it is a deinterleave. _dequantize_from_int4pack does that, bit-exact against the old dequantizer at 2048x2048, 3072x2048 and gate_up's real 13312x2048, and the duplicate is no longer allocated. Measured with the two paths toggled in otherwise identical code, 3 arms x 2 interleaved passes, RssAnon summed over the process tree, 1536-token context: 615 MB of resident memory and ~915 MB of peak, landing statistically indistinguishable from prefill_dequant: false (the harness refused to rank them, 0.4% apart inside the spread). Config arithmetic predicts 764 MB; ~150 MB is allocator retention malloc_trim does not return. It is not free, and the docstring says so. The duplicate was not only a copy: to_split_halves de-interleaves once at load, so each prefill call read two contiguous planes, while this redoes the de-interleave every call. pp512 drops 2272 -> 1210 tok/s. Dequantizing to [K, N] and using matmul to dodge the transpose was tried and is worse (20.3 ms vs 14.4 per gate_up at 8 threads); recovering the prefill wants a C++ de-interleave, not an ATen expression chain. So it is a dial. Against prefill_dequant: false it is strictly better -- same memory, 2.7x the prefill, and dequant_threshold stays in force so multi-row batches never reach the tinygemm kernel and the documented SIGILL hazard. That setting is now dominated. Against keeping both layouts it trades 615 MB for 1.9x prefill; default on, because this path exists for devices chosen for their memory ceiling. supports_direct_dequant gates the reshape on N tiling the 64-channel block, which every Spark linear does (2048, 3072, 13312, 131072). An N like 80 packs as 64 + a 16 tail the reshape would mis-read, so those layers keep the copy rather than being quietly wrong. The choice is a `direct_dequant` field of the checkpoint's quantization_config, not only an env var, and that matters: the two paths build different graphs (one passes weight_halves into the custom op, the other passes None), so they must not share a compiled artifact. ModelConfig.compute_hash hashes config fields, so a config field reaches vLLM's compile-cache key; an env var does not. Toggling the env var against a warm cache loaded the wrong graph and killed three benchmark passes with `KeyError: 'weight_halves'` inside the AOT-loaded forward -- found because the harness records failures instead of dropping them. VLLM_OMNI_CPU_INT4_DIRECT_DEQUANT remains as an experiment override, documented as needing its own VLLM_CACHE_ROOT. Tests: 18 passed, including bit-exactness at three shapes, a pin on the packed layout itself (so a torch release that repacks differently fails loudly rather than returning wrong weights), the N=80 fallback, and both paths producing identical output. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> |
||
|
tzhouamandClaude Opus 5
|
3dac97e0bd |
[Edge][CPU] W4A8 via vLLM's own AMX kernel, a per-device weight path, and stop recompiling on every start
Four changes, each measured on a Xeon 8480C with Spark-X2.5-1.7B at 24 threads. Numbers are best-of-interleaved-passes; this host is shared and a single sequential grid swung the llama.cpp control 30%, so anything under ~8% is not treated as a result. **vLLM's CPU backend does have a 4-bit kernel.** An earlier note in this tree said it did not, on the strength of `woq_int4_linear`, `int4_scaled_mm_with_quant` and `da8w4_linear` all being absent. Those are simply not the names it uses: `CPUWNA16LinearKernel` dispatches GPTQ/AWQ checkpoints to `int4_scaled_mm_cpu`, an AMX W4A8 GEMM, whenever activations are bf16 and the checkpoint carries no act-order. Four-bit weights times eight-bit activations is the arithmetic llama.cpp's Q4_K_M already does. `gptq_cpu_export` builds checkpoints for it; the kernel pins the zero point at 8, so the grid must be symmetric, and a per-group scale search recovers most of what that costs. **A per-device weight path** (`edge/weight_path.py`). The fast kernel on each device needs a different, mutually exclusive weight layout, so the choice has to come from a probe rather than a flag: AMX x86 -> W4A8/GPTQ; x86 without AMX, or priority="fidelity" -> W4A16/tinygemm (bf16 activations, the more faithful of the two); ARM -> W4A16 but flagged unverified, because PyTorch's 4-bit CPU GEMM dispatches a different kernel there than the AVX-512 one measured here; phone -> runs_in_vllm=False, the deliverable is an exported graph with calibrated w8a16. **The CPU compile cache never loaded.** `_try_load_aot_compiled_fn` finishes with `compiled_fn.finalize_loading(cfg)`, which the CUDA wrapper object has and the inductor CPU backend's plain function does not. The AttributeError was swallowed as a cache miss, so every engine start recompiled the model: warm start 49.8 s against 14.9 s with the missing no-op attached. Self- disabling if upstream gives the CPU path a real one. **The head gate stops being a GEMM.** A 2048x8 projection cannot amortize a thread fan-out: dispatched as `aten::mm` it costs 1.155 ms of a decode step, nearly all barrier. As an eager reduction behind an opaque op it costs 0.444 ms. It has to be opaque -- written inline the reduction is traced, and inductor's loop made the whole step 1.8x slower (78.8 -> 44.5 tok/s). Also included and *off by default*: `cpu_ir_ops` registers the fused `_C::rms_norm` / `_C::fused_add_rms_norm` kernels that vLLM's CPU build ships but gates on `GPGPU_DEVICE`. In isolation they are 3.6x and 4.5x faster than the native path; end to end they are slightly slower (72.7 vs 74.0 tok/s), because inductor was already fusing the native norm into its neighbours and an opaque call breaks that fusion. Kept with the measurement recorded so the next person does not repeat it. `gelu_and_mul` is deliberately not registered: its fused CPU kernel is 2.7x slower at Spark's GEGLU shape. Decode, W4A16 g64 -> W4A8 g128 with the gate fix: 69.1 -> 84.5 tok/s (+22%), against llama.cpp Q4_K_M's 90.0; prefill 3347 against 1121 (3.0x). The two 4-bit paths are not equally faithful: against each build's own bf16 over 12 greedy prompts, Q4_K_M reproduces 4/12, W4A16 2/12 and W4A8 1/12. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> |
||
|
tzhouamandClaude Opus 5
|
7869a8dc0d |
Add cpu_int4, an asymmetric 4-bit weight path for vLLM's CPU backend
Single-stream CPU decode reads every weight to emit one token, so bits per weight set the token rate. This registers a W4A16 group-wise method out-of-tree through register_quantization_config -- vLLM itself is untouched -- using PyTorch's _weight_int4pack_mm_for_cpu tinygemm kernel. Two paths over one set of weights, because that kernel dequantizes once per row of activations and so is linear in the batch (12.9x slower than bf16 at 512 rows). Single-row steps take tinygemm; wider batches unpack the weight to bf16 once and use the ordinary oneDNN/AMX GEMM. Shipping only the first measured pp512 = 557 tok/s against bf16's 4455; with both, 3589. The batch threshold is 2 rather than the ~8 rows where the unpack path starts winning on speed, because running tinygemm on a multi-row batch and then vLLM's cpu_attention_with_kv_cache in the same forward dies with SIGILL inside the attention kernel -- an AMX instruction with no valid tile config loaded. The same prompt is fine in bf16 and fine once every multi-row GEMM takes the unpack path. Positioning, measured on a Xeon 8480C with 32 pinned cores and a warm cache (analysis/spark_cpu_int4.md): vLLM's own GPTQ W4A8 AMX path is faster -- 88.7 tok/s decode and 4355 prefill against 85.5 and 3589 here -- but its config rejects sym=False at 4 bits, so it can only be fed a symmetric grid. This one is asymmetric and costs +12.5% perplexity against that path's +27.2%, and needs no AMX. llama.cpp Q4_K_M, for reference, is 98.0 tok/s decode and 1356 prefill at +7.6% and 5.15 bits per weight. Expanding a quantized checkpoint back to bf16 offline and scoring it reproduces the quantized run to within 0.2% perplexity, so both GEMM paths, the nibble packing and the quantized tied embedding are faithful. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GXhbAfkSNRaeqZrus5Z2fx |