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>

tests/model_executor/layers/quantization/test_cpu_int4.py

@@ -304,3 +304,61 @@ def test_env_override_still_wins_for_experiments(monkeypatch):
     assert direct_dequant_enabled(True) is False
     monkeypatch.setenv("VLLM_OMNI_CPU_INT4_DIRECT_DEQUANT", "1")
     assert direct_dequant_enabled(False) is True
+
+
+def test_blocked_repacking_matches_packing_the_whole_matrix():
+    """unpack_nibbles expands 4-bit codes to int32 -- 8x the packed form -- and
+    doing that for a whole matrix makes the transient the largest thing in the
+    process: 1024 MiB for the 131072-row tied embedding. n-blocks of the packed
+    layout are independent, so it can be chunked; the result must be identical
+    or the weights are silently wrong."""
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
+        INNER_K_TILES,
+        _pack_blocked,
+        quantize_groupwise,
+        unpack_nibbles,
+    )
+
+    torch.manual_seed(0)
+    for n, k in ((1024, 512), (2048, 1024), (640, 256)):
+        w = torch.randn(n, k, dtype=torch.bfloat16) * 0.05
+        packed, _, _ = quantize_groupwise(w, 128)
+        whole = torch.ops.aten._convert_weight_to_int4pack_for_cpu(
+            unpack_nibbles(packed), INNER_K_TILES
+        )
+        assert torch.equal(_pack_blocked(packed), whole), f"N={n} K={k}"
+
+
+def test_blocked_repacking_handles_a_ragged_last_chunk():
+    """N not a multiple of the chunk size must still round-trip."""
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
+        INNER_K_TILES,
+        _pack_blocked,
+        quantize_groupwise,
+        unpack_nibbles,
+    )
+
+    n, k = 576, 256           # 576 = 512 + 64, one full chunk plus a short one
+    w = torch.randn(n, k, dtype=torch.bfloat16) * 0.05
+    packed, _, _ = quantize_groupwise(w, 128)
+    whole = torch.ops.aten._convert_weight_to_int4pack_for_cpu(
+        unpack_nibbles(packed), INNER_K_TILES
+    )
+    assert torch.equal(_pack_blocked(packed), whole)
+
+
+def test_split_halves_still_round_trips_after_chunking():
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
+        _dequantize_halves,
+        quantize_groupwise,
+        to_split_halves,
+    )
+
+    n, k, group = 1088, 512, 128          # not a multiple of the chunk size
+    w = torch.randn(n, k, dtype=torch.bfloat16) * 0.05
+    packed, scale, zero = quantize_groupwise(w, group)
+    deq = _dequantize_halves(to_split_halves(packed),
+                             scale.to(torch.bfloat16), zero.to(torch.bfloat16), group)
+    err = (deq.float() - w.float())
+    snr = 10 * torch.log10((w.float() ** 2).mean() / err.pow(2).mean())
+    assert snr > 20, f"SNR {snr:.1f} dB"