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"

vllm_omni/model_executor/layers/quantization/cpu_int4.py

@@ -155,10 +155,43 @@ def quantize_groupwise(w: torch.Tensor, group: int):
     return pack_nibbles(q.reshape(n, k)), scale, zero
 
 
+REPACK_ROWS = 512
+"""Output channels repacked at a time, to bound the load-time peak.
+
+``unpack_nibbles`` expands 4-bit codes to int32 -- 4 bytes per weight, eight
+times the packed form -- and the kernel's packer wants int32. Doing that for a
+whole matrix makes the transient the largest thing in the process: 104 MiB for
+``gate_up``, and **1024 MiB for the 131072-row tied embedding**, against 24 MiB
+of packed weights for an entire layer. Since the packed layout is
+``[N/64][K][32 bytes]``, n-blocks are independent, so the expansion can cover
+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.
+"""
+
+
+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)
+    if n <= step:
+        return torch.ops.aten._convert_weight_to_int4pack_for_cpu(
+            unpack_nibbles(packed), INNER_K_TILES
+        )
+    out = []
+    for i in range(0, n, step):
+        codes = unpack_nibbles(packed[i : i + step].contiguous())
+        out.append(
+            torch.ops.aten._convert_weight_to_int4pack_for_cpu(codes, INNER_K_TILES)
+        )
+        del codes
+    return torch.cat(out, dim=0)
+
+
 def _to_tinygemm(packed: torch.Tensor, scale: torch.Tensor, zero: torch.Tensor):
     """Row-major nibbles + [N, G] scales -> kernel layout + [G, N, 2] scales."""
-    codes = unpack_nibbles(packed)
-    inner = torch.ops.aten._convert_weight_to_int4pack_for_cpu(codes, INNER_K_TILES)
+    inner = _pack_blocked(packed)
     sz = torch.stack([scale.float(), zero.float()], dim=-1)  # [N, G, 2]
     sz = sz.transpose(0, 1).contiguous().to(torch.bfloat16)  # [G, N, 2]
     return inner, sz
@@ -172,10 +205,16 @@ def to_split_halves(packed: torch.Tensor) -> torch.Tensor:
     both nibble planes contiguous runs, and the dequantize loop drops its
     lane shuffle: measured 3.1 ms per layer instead of 4.5.
     """
-    codes = unpack_nibbles(packed)
-    n, k = codes.shape
-    lo, hi = codes[:, : k // 2], codes[:, k // 2 :]
-    return (lo.to(torch.uint8) | (hi.to(torch.uint8) << 4)).contiguous()
+    n = packed.shape[0]
+    step = max(REPACK_ROWS - REPACK_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())
+        k = codes.shape[1]
+        lo, hi = codes[:, : k // 2], codes[:, k // 2 :]
+        out.append((lo.to(torch.uint8) | (hi.to(torch.uint8) << 4)).contiguous())
+        del codes
+    return torch.cat(out, dim=0) if len(out) > 1 else out[0]
 
 
 TINYGEMM_N_BLOCK = 64