edge: measure cpu_gemm_wna16 instead of asserting it away (I2)

The review caught weight_path.py claiming tinygemm is "the only option" for an
x86 core without AMX. That was false -- vLLM 0.28 selects ops.cpu_gemm_wna16
whenever layer.use_w4a8 is off -- and the honest response was to measure the
kernel, not to reword the claim.

Same GPTQ checkpoint, only VLLM_CPU_INT4_W4A8 differing, isolated compile
caches, 3 interleaved passes on 24 pinned node-local 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%

The fallback is real but 2.2x slower than tinygemm, so the selector's routing
was right for a reason nobody had checked. It stays as it is; what changes is
that the module now cites a measurement rather than an assumption.

Recorded with the caveat that matters for reading it off-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 conclusion, but no number is
claimed for it.

Also adds VLLM_OMNI_CPU_INT4_REPACK_ROWS so the blocked repack can be A/B'd
against whole-matrix packing in identical code. Unlike the prefill-path switch
this is load-time only and yields identical tensors, so it does not perturb the
compiled graph or its cache key.

tests/edge: 8 passed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

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.
 

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())