Cache device-side constants in dequant kernels

[lanslote]

Problem

Every dequantize_blocks_* function that needs a bit-shift vector builds it inline:

qs = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1))

torch.tensor(..., device=gpu) is a pageable host-to-device copy plus an implicit sync. Since dequantize runs once per quantized tensor per forward pass, this is paid thousands of times per sampling step. KVALUES.to(qs.device) in the IQ4 paths has the same issue.

I found this while profiling a stalled LTX-Video run — py-spy showed 100% of samples inside dequantize_blocks_Q4_K/Q5_K, and it turned out to be the constant allocation rather than the arithmetic.

Measured cost of a single torch.tensor([0, 4], device=gpu):

GPU	uncached	cached
RTX 4050 Laptop (internal PCIe)	307.6 us	26.0 us
RX 9060 XT (USB4 eGPU, 2.75 GB/s H2D)	178.6 us	1.8 us

It is not eGPU-specific — the internal-PCIe laptop GPU was actually worse in absolute terms. eGPU just made it impossible to ignore.

Change

Cache the constants keyed on (values, device, dtype). The literals used in this file are a fixed small set, so the cache holds a handful of tiny tensors per device — it does not grow with model size or step count. KVALUES gets a cached device copy instead of being re-uploaded on every IQ4 call.

No behavioural change otherwise; the tensors were already immutable and only ever used as shift operands.

Verification

Compared every dequant function against the current main implementation on identical random input, then benchmarked both. 4096 blocks, bf16 target, 60 iterations after 15 warmup:

qtype        upstream     patched   speedup  identical
BF16           77.4us      28.6us     2.71x       True
IQ4_NL        956.6us     305.1us     3.14x       True
IQ4_XS       1998.7us     693.1us     2.88x       True
Q2_K          918.7us     259.4us     3.54x       True
Q3_K         2414.9us     593.0us     4.07x       True
Q4_0          488.1us     155.4us     3.14x       True
Q4_1          521.0us     138.7us     3.76x       True
Q4_K          885.2us     559.9us     1.58x       True
Q5_0          942.5us     519.3us     1.81x       True
Q5_1         1023.0us     651.5us     1.57x       True
Q5_K         1729.2us     516.6us     3.35x       True
Q6_K         1377.4us     387.1us     3.56x       True
Q8_0           94.4us      52.8us     1.79x       True
TOTAL       13427.0us    4860.4us     2.76x

qtypes tested: 13
ALL BIT-IDENTICAL

Q8_0 is a useful control — it is the only qtype with no shift constant, and it moves the least.

End-to-end on a real workflow (LTX-Video 22B Q4_K_S, 8-step sampling + 3-step upscale + VAE decode, same seed and settings, --reserve-vram tuned so there is zero weight offload in both runs):

	wall time
upstream	976.4 s
patched	690.2 s

1.41x, output visually identical.

Repro script

bench_dequant_constants.py

import time
import torch

def bench(name, fn, iters=200, warmup=25):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize()
    dt = (time.perf_counter() - t0) / iters
    print("{:52s} {:10.1f} us".format(name, dt * 1e6))
    return dt

dev = torch.device("cuda:0")
print("device:", torch.cuda.get_device_name(0))

n_blocks = 4096
qs = torch.randint(0, 255, (n_blocks, 128), dtype=torch.uint8, device=dev)

bench("torch.tensor([0,4], device=gpu, uint8)",
      lambda: torch.tensor([0, 4], device=dev, dtype=torch.uint8))

cached = torch.tensor([0, 4], device=dev, dtype=torch.uint8).reshape((1, 1, 2, 1))
bench("cached constant (no alloc)", lambda: cached.reshape((1, 1, 2, 1)))

def shift_fresh():
    return (qs.reshape((n_blocks, -1, 1, 32))
            >> torch.tensor([0, 4], device=dev, dtype=torch.uint8).reshape((1, 1, 2, 1)))

def shift_cached():
    return qs.reshape((n_blocks, -1, 1, 32)) >> cached

a = bench("shift WITH fresh torch.tensor", shift_fresh)
b = bench("shift WITH cached constant", shift_cached)
print("{:52s} {:10.2f}x".format("speedup", a / b))

bench("qs & 0x0F  (reference elementwise)", lambda: qs & 0x0F)

Notes

This is orthogonal to #336 — that replaces the dequant kernels with Triton implementations, this just stops re-allocating constants inside the existing ones. If #336 lands, these call sites disappear anyway, but this is a small independent win in the meantime.

Tested on ROCm (gfx1200, torch 2.10) and CUDA (RTX 4050, torch 2.x). Happy to adjust naming or drop the inline comment block if you prefer it leaner.

[m8rr]

Okay, it looks like you did the optimization that TorchDynamo/Torch.compile was doing.

PR + Compile: 75s
Compile: 75s
PR: 78s
Base: 94s

m8rr@1c113ea
For my part, I had applied torch.compile to the dequantization function in a dynamic VRAM environment. However, the initial execution gets a bit slower, and if the code changes due to updates, it becomes much slower. Thanks to this, I think I can skip compilation altogether.

For reference, I am using Oculink PCIe 4x4, so that might be why it's so effective.

However, in text generation, there is a huge performance difference with torch.compile:
Compile: 4 it/s
PR: 1.3 it/s
Base: 1 it/s
Though ComfyUI's default text generation is already too slow compared to llama.cpp and others anyway, so it doesn't mean much.

How about expanding to low-beat quants as well?
#433