diff --git a/vllm_omni/edge/weight_path.py b/vllm_omni/edge/weight_path.py index 19c83ad16..d3d8044ad 100644 --- a/vllm_omni/edge/weight_path.py +++ b/vllm_omni/edge/weight_path.py @@ -12,12 +12,23 @@ one needs a different weight layout, and the layouts are mutually exclusive: budget (see below) -- and it is the fastest thing measured here; * an x86 core without AMX falls back to PyTorch's tinygemm path, which wants its own packing and keeps activations in bf16 -- slower, but more faithful. - This is *not* the only option there, as this module previously claimed: - vLLM 0.28 also has `ops.cpu_gemm_wna16`, a vector-ISA W4A16 kernel it - selects whenever `layer.use_w4a8` is off + This module used to call that "the only option", which was false: vLLM 0.28 + also has `ops.cpu_gemm_wna16`, selected whenever `layer.use_w4a8` is off (`vllm/model_executor/kernels/linear/mixed_precision/cpu.py:202`). It has - never been measured here, so the selector does not route to it yet -- but - "no AMX means tinygemm" is an untested assumption, not a fact; + now been measured, on the same GPTQ checkpoint with only + `VLLM_CPU_INT4_W4A8` differing, 3 interleaved passes on 24 pinned cores: + + int4_scaled_mm_cpu (W4A8, AMX) 74.4 tok/s spread 1.1% + tinygemm W4A16 (same cores) 71.6 tok/s + cpu_gemm_wna16 32.3 tok/s spread 1.5% + + So the fallback is real but **2.2x slower than tinygemm**, and the routing + here was right for a reason nobody had checked. One caveat on reading that + number outside x86: `_get_isa_hint` returns `"amx"` whenever AMX is present + and activations are bf16, *regardless* of the W4A8 flag, so this measured + `cpu_gemm_wna16` on its AMX path -- its best case on this host. The `"vec"` + path a genuinely AMX-less machine takes is untested and presumably no + faster, which only strengthens the choice; * a phone NPU does not run vLLM at all. The deliverable there is an exported graph, and its quantization is decided by the export toolchain, not by us. diff --git a/vllm_omni/model_executor/layers/quantization/cpu_int4.py b/vllm_omni/model_executor/layers/quantization/cpu_int4.py index bd01ea9c7..39caa2cf9 100644 --- a/vllm_omni/model_executor/layers/quantization/cpu_int4.py +++ b/vllm_omni/model_executor/layers/quantization/cpu_int4.py @@ -168,13 +168,27 @@ one chunk at a time: verified byte-identical to packing the whole matrix at 13312x2048, 2048x6656, 3072x2048 and 131072x2048. A multiple of ``TINYGEMM_N_BLOCK``; the loop rounds down to one. +``VLLM_OMNI_CPU_INT4_REPACK_ROWS`` overrides it, and 0 means "the whole matrix +at once", which restores the old behaviour for an A/B. Unlike the prefill-path +switch this only affects load time and produces identical tensors, so it does +not perturb the compiled graph or its cache key. """ +def _repack_rows() -> int: + raw = os.environ.get("VLLM_OMNI_CPU_INT4_REPACK_ROWS") + return REPACK_ROWS if raw is None else int(raw) + + def _pack_blocked(packed: torch.Tensor) -> torch.Tensor: """Kernel layout, without ever expanding the whole matrix to int32.""" n = packed.shape[0] - step = max(REPACK_ROWS - REPACK_ROWS % TINYGEMM_N_BLOCK, TINYGEMM_N_BLOCK) + rows = _repack_rows() + if rows <= 0: # explicit opt-out: old behaviour + return torch.ops.aten._convert_weight_to_int4pack_for_cpu( + unpack_nibbles(packed), INNER_K_TILES + ) + step = max(rows - rows % TINYGEMM_N_BLOCK, TINYGEMM_N_BLOCK) if n <= step: return torch.ops.aten._convert_weight_to_int4pack_for_cpu( unpack_nibbles(packed), INNER_K_TILES @@ -206,7 +220,8 @@ def to_split_halves(packed: torch.Tensor) -> torch.Tensor: lane shuffle: measured 3.1 ms per layer instead of 4.5. """ n = packed.shape[0] - step = max(REPACK_ROWS - REPACK_ROWS % TINYGEMM_N_BLOCK, TINYGEMM_N_BLOCK) + rows = _repack_rows() + step = n if rows <= 0 else max(rows - rows % TINYGEMM_N_BLOCK, TINYGEMM_N_BLOCK) out = [] for i in range(0, n, step): # bounded transient, as above codes = unpack_nibbles(packed[i : i + step].contiguous())