cpu_int4: dequantize prefill weights from the tinygemm layout, not a copy

The W4A16 path kept two packed weight layouts because the tinygemm layout was
treated as opaque: the prefill path needed something row-major to dequantize
from, so `_finalize_int4` built `weight_halves` plus cloned scale/zero beside
the kernel's own buffer. That duplicate is the 740 MB the report attributes to
"the dual weight layout".

The layout is not opaque. Measured on this host by packing an index pattern and
reading back where each nibble landed: it is [N/64][K][32 bytes] -- K is the
middle axis, so a K range is contiguous inside each n-block -- and within every
32-byte run, byte d carries output channel d in its low nibble and d+32 in its
high one, the same permutation for every k and every block. Decoding it is a
deinterleave. _dequantize_from_int4pack does that, bit-exact against the old
dequantizer at 2048x2048, 3072x2048 and gate_up's real 13312x2048, and the
duplicate is no longer allocated.

Measured with the two paths toggled in otherwise identical code, 3 arms x 2
interleaved passes, RssAnon summed over the process tree, 1536-token context:
615 MB of resident memory and ~915 MB of peak, landing statistically
indistinguishable from prefill_dequant: false (the harness refused to rank
them, 0.4% apart inside the spread). Config arithmetic predicts 764 MB; ~150 MB
is allocator retention malloc_trim does not return.

It is not free, and the docstring says so. The duplicate was not only a copy:
to_split_halves de-interleaves once at load, so each prefill call read two
contiguous planes, while this redoes the de-interleave every call. pp512 drops
2272 -> 1210 tok/s. Dequantizing to [K, N] and using matmul to dodge the
transpose was tried and is worse (20.3 ms vs 14.4 per gate_up at 8 threads);
recovering the prefill wants a C++ de-interleave, not an ATen expression chain.

So it is a dial. Against prefill_dequant: false it is strictly better -- same
memory, 2.7x the prefill, and dequant_threshold stays in force so multi-row
batches never reach the tinygemm kernel and the documented SIGILL hazard. That
setting is now dominated. Against keeping both layouts it trades 615 MB for
1.9x prefill; default on, because this path exists for devices chosen for
their memory ceiling.

supports_direct_dequant gates the reshape on N tiling the 64-channel block,
which every Spark linear does (2048, 3072, 13312, 131072). An N like 80 packs
as 64 + a 16 tail the reshape would mis-read, so those layers keep the copy
rather than being quietly wrong.

The choice is a `direct_dequant` field of the checkpoint's quantization_config,
not only an env var, and that matters: the two paths build different graphs
(one passes weight_halves into the custom op, the other passes None), so they
must not share a compiled artifact. ModelConfig.compute_hash hashes config
fields, so a config field reaches vLLM's compile-cache key; an env var does
not. Toggling the env var against a warm cache loaded the wrong graph and
killed three benchmark passes with `KeyError: 'weight_halves'` inside the
AOT-loaded forward -- found because the harness records failures instead of
dropping them. VLLM_OMNI_CPU_INT4_DIRECT_DEQUANT remains as an experiment
override, documented as needing its own VLLM_CACHE_ROOT.

Tests: 18 passed, including bit-exactness at three shapes, a pin on the packed
layout itself (so a torch release that repacks differently fails loudly rather
than returning wrong weights), the N=80 fallback, and both paths producing
identical output.

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

tests/model_executor/layers/quantization/test_cpu_int4.py

@@ -132,3 +132,175 @@ def test_ignore_list_falls_back_to_full_precision():
         cfg.get_quant_method(layer, "model.layers.0.mlp.down_proj"),
         UnquantizedLinearMethod,
     )
+
+
+# ------------------------------------------------- direct dequant (one copy)
+
+def test_direct_dequant_reproduces_the_row_major_path_exactly():
+    """The prefill path used to need a second packed copy of every weight --
+    740 MB on Spark -- purely because the tinygemm layout was treated as
+    opaque. Reading it directly has to give bit-identical weights, or the
+    saving is bought with an accuracy change nobody asked for."""
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
+        _dequantize_from_int4pack,
+        _dequantize_halves,
+        _to_tinygemm,
+        quantize_groupwise,
+        to_split_halves,
+    )
+
+    torch.manual_seed(0)
+    for n, k, group in ((2048, 2048, 128), (3072, 2048, 128), (256, 512, 64)):
+        w = torch.randn(n, k, dtype=torch.bfloat16) * 0.05
+        packed, scale, zero = quantize_groupwise(w, group)
+        int4pack, sz = _to_tinygemm(packed, scale, zero)
+        ref = _dequantize_halves(
+            to_split_halves(packed), scale.to(torch.bfloat16),
+            zero.to(torch.bfloat16), group,
+        )
+        got = _dequantize_from_int4pack(int4pack, sz, group)
+        assert torch.equal(got, ref), f"N={n} K={k} g={group}"
+
+
+def test_the_packed_layout_this_depends_on_is_what_torch_actually_emits():
+    """Pins the measured layout: [N/64][K][32 bytes], byte d holding output
+    channel d in its low nibble and d+32 in its high one. A torch release that
+    repacks differently must fail here rather than silently return wrong
+    weights."""
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
+        INNER_K_TILES,
+        TINYGEMM_N_BLOCK,
+    )
+
+    n, k = 128, 64
+    codes = torch.randint(0, 16, (n, k), dtype=torch.int32)
+    packed = torch.ops.aten._convert_weight_to_int4pack_for_cpu(codes, INNER_K_TILES)
+    blocks = packed.reshape(n // TINYGEMM_N_BLOCK, k, TINYGEMM_N_BLOCK // 2)
+    lo = (blocks & 0x0F).to(torch.int32)
+    hi = (blocks >> 4).to(torch.int32)
+    rebuilt = torch.cat([lo, hi], dim=2).permute(0, 2, 1).reshape(n, k)
+    assert torch.equal(rebuilt, codes)
+
+
+def test_shapes_that_do_not_tile_the_block_keep_the_second_copy():
+    """N=80 packs as a 64 block plus a 16 tail, which the reshape would
+    mis-read. Those layers must fall back, not be quietly wrong."""
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
+        supports_direct_dequant,
+    )
+
+    assert supports_direct_dequant(2048)
+    assert supports_direct_dequant(3072)
+    assert supports_direct_dequant(13312)
+    assert supports_direct_dequant(131072)
+    assert supports_direct_dequant(32)        # below the block, tiles exactly
+    assert not supports_direct_dequant(80)    # 64 + 16 tail
+    assert not supports_direct_dequant(144)
+
+
+def test_a_layer_with_prefill_dequant_no_longer_allocates_a_second_copy():
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
+        _finalize_int4,
+        quantize_groupwise,
+    )
+
+    n, k, group = 512, 512, 128
+    w = torch.randn(n, k, dtype=torch.bfloat16) * 0.05
+    packed, scale, zero = quantize_groupwise(w, group)
+    layer = torch.nn.Module()
+    layer.weight_packed = torch.nn.Parameter(packed, requires_grad=False)
+    layer.weight_scale = torch.nn.Parameter(scale.to(torch.bfloat16), requires_grad=False)
+    layer.weight_zero = torch.nn.Parameter(zero.to(torch.bfloat16), requires_grad=False)
+    _finalize_int4(layer, keep_row_major=False, prefill_dequant=True)
+
+    assert layer.cpu_int4_direct_dequant is True
+    assert layer.weight_halves is None
+    assert layer.weight_group_scale is None
+    assert layer.weight_group_zero is None
+    # and the prefill path still works off the remaining layout
+    x = torch.randn(8, k, dtype=torch.bfloat16)
+    out = torch.ops.vllm.cpu_int4_linear(
+        x, layer.weight_int4pack, layer.weight_scales_and_zeros,
+        None, None, None, group, 2, True,
+    )
+    assert out.shape == (8, n) and torch.isfinite(out).all()
+
+
+def test_the_escape_hatch_restores_the_second_copy(monkeypatch):
+    """The two paths trade 615 MB against 1.9x prefill, so both have to remain
+    reachable -- and the A/B that measured them needs identical code."""
+    from vllm_omni.model_executor.layers.quantization import cpu_int4
+
+    n, k, group = 512, 512, 128
+    w = torch.randn(n, k, dtype=torch.bfloat16) * 0.05
+    packed, scale, zero = cpu_int4.quantize_groupwise(w, group)
+
+    def build():
+        layer = torch.nn.Module()
+        layer.weight_packed = torch.nn.Parameter(packed, requires_grad=False)
+        layer.weight_scale = torch.nn.Parameter(
+            scale.to(torch.bfloat16), requires_grad=False)
+        layer.weight_zero = torch.nn.Parameter(
+            zero.to(torch.bfloat16), requires_grad=False)
+        cpu_int4._finalize_int4(layer, keep_row_major=False, prefill_dequant=True)
+        return layer
+
+    monkeypatch.setenv("VLLM_OMNI_CPU_INT4_DIRECT_DEQUANT", "0")
+    old = build()
+    assert old.cpu_int4_direct_dequant is False
+    assert old.weight_halves is not None
+
+    monkeypatch.setenv("VLLM_OMNI_CPU_INT4_DIRECT_DEQUANT", "1")
+    new = build()
+    assert new.cpu_int4_direct_dequant is True
+    assert new.weight_halves is None
+
+    # and both produce the same answer, which is what makes the A/B meaningful
+    x = torch.randn(8, k, dtype=torch.bfloat16)
+    a = torch.ops.vllm.cpu_int4_linear(
+        x, old.weight_int4pack, old.weight_scales_and_zeros,
+        old.weight_halves, old.weight_group_scale, old.weight_group_zero,
+        group, 2, False)
+    b = torch.ops.vllm.cpu_int4_linear(
+        x, new.weight_int4pack, new.weight_scales_and_zeros,
+        None, None, None, group, 2, True)
+    assert torch.equal(a, b)
+
+
+def test_the_path_choice_lives_in_the_config_so_it_reaches_the_cache_key():
+    """The two prefill paths build different graphs -- one passes
+    weight_halves into the custom op, the other passes None -- so they must not
+    share a compiled artifact.
+
+    An env-var-only switch does not reach vLLM's compile-cache key, and a warm
+    cache then loads the wrong graph: three passes of a benchmark failed with
+    `KeyError: 'weight_halves'` inside the AOT-loaded forward, after the first
+    pass had saved an artifact for the other path. `direct_dequant` is a config
+    field so it is hashed with the rest of the model config.
+    """
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import CPUInt4Config
+
+    on = CPUInt4Config.from_config({"group_size": 128, "direct_dequant": True})
+    off = CPUInt4Config.from_config({"group_size": 128, "direct_dequant": False})
+    assert on.direct_dequant is True
+    assert off.direct_dequant is False
+    # and it shows up in the repr, which is what lands in logs and configs
+    assert "direct_dequant=True" in repr(on)
+    assert "direct_dequant=False" in repr(off)
+    # default stays on: this path exists for memory-constrained devices
+    assert CPUInt4Config.from_config({"group_size": 128}).direct_dequant is True
+
+
+def test_env_override_still_wins_for_experiments(monkeypatch):
+    from vllm_omni.model_executor.layers.quantization.cpu_int4 import (
+        direct_dequant_enabled,
+    )
+
+    monkeypatch.delenv("VLLM_OMNI_CPU_INT4_DIRECT_DEQUANT", raising=False)
+    assert direct_dequant_enabled(True) is True
+    assert direct_dequant_enabled(False) is False
+
+    monkeypatch.setenv("VLLM_OMNI_CPU_INT4_DIRECT_DEQUANT", "0")
+    assert direct_dequant_enabled(True) is False
+    monkeypatch.setenv("VLLM_OMNI_CPU_INT4_DIRECT_DEQUANT", "1")
+    assert direct_dequant_enabled(False) is True