c4e61133fbb5587e0cf8972ce32ffac4e37fa739
3
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
|
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> |
||
|
|
658341531a |
[Edge][Mobile] Spark-X2.5 decode-step export, and a compile-safe head gate
Exports one Spark decode step with the KV cache as explicit tensors, sized per layer type: the 21 sliding layers hold a fixed 512 entries at any context, the 7 full-attention layers grow. That asymmetry is the whole reason the architecture is interesting on a phone, so the two types get separately sized caches, masks and rotary angles. Verified against the HF reference at 1024 context: 138.5 dB logits SNR, identical top-1 and top-5. Profiled on a real Galaxy S25 (Snapdragon 8 Elite) through AI Hub, QNN context binary on the NPU, ms per layer: sliding (512 KV, any ctx) fp16 1.689 w8a16 0.912 w4a16 0.917 full @ 1k fp16 1.559 w8a16 1.129 w4a16 1.127 full @ 4k fp16 2.964 w8a16 1.681 w4a16 1.662 Composed into a step (21 sliding + 7 full), w4a16 gives 27.2 ms at 1k and 30.9 ms at 4k, against 31.6 and 46.5 ms for the same model with every layer full-attention -- 1.16x and 1.51x. Worth noting the hybrid is 0.94x at 1k in fp16: sliding layers rotate all 256 head dims where full layers rotate 64, and below roughly 2k that RoPE cost outweighs the shorter cache. Also moves the head-wise attention gate off vLLM's CPU linear dispatch. The gate is [2048, 8], and any linear that narrow is routed to the sgl-kernel packed GEMM, whose fake-tensor rule assumes the packed weight keeps N in dim 0; VNNI packing does not at that size, so it reports K and torch.compile fails tracing the gate multiply. A plain parameter with F.linear gains nothing less -- a 2048x8 matmul has no use for packed AMX tiles -- and it keeps the model compilable on CPU and GPU alike. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ |
||
|
|
e4a565b3df |
[Edge][Model] Spark-X2.5 (Spark2_5ForCausalLM) support
Spark-X2.5-1.7B is an Apache-2.0 on-device model built for agentic tool use. vLLM 0.28 has no implementation for it, so it could not be served at all; llama.cpp already ships one, which made it the only edge engine able to run the model. Three traits separate it from a Llama-style decoder: * Hybrid attention -- layer_types interleaves three 512-token sliding layers with one full-attention layer, which is what makes the 1M-token context affordable on a phone. * Per-layer-type RoPE -- sliding layers rotate all 256 head dims with theta 1e4, full layers only the first quarter with theta 5e6. * A head-wise attention output gate -- g_proj reads the same post-norm input as QKV and scales each head's attention output by a sigmoid. The checkpoint also stores QKV pre-fused as a single q_k_v_proj, so load_weights splits it before vLLM's sharded QKVParallelLinear sees it. Verified against the official remote-code implementation on transformers 4.57.1 (the version config.json targets; it does not import on the 5.x in our venvs). In fp32 all three scenarios match exactly -- 32/32 greedy tokens with a max top-20 logprob difference of 2.2e-05. In bf16 the short and the 2465-token long-context prompts still match 32/32, which is what exercises the sliding layers; the tool-calling prompt diverges at token 23 on a near-tie, and fp32 confirms that is rounding rather than a defect. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ |