Skip to content

int4 per-token-head KV cache: triton_reshape_and_cache_flash_per_token_head_quant returns wrong values #53086

Description

@khluu

Observed on main @ eac636a (NVIDIA H200, driver 580.159.03, CUDA 13.0).

tests/quantization/test_per_token_kv_cache.py::test_reshape_and_cache_per_token_head fails in the int4 branch (kv_quant_mode=INT4_PER_TOKEN_HEAD, quant_min=-8.0, quant_max=7.0): the dequantized round-trip (unpack nibbles → zero-point → scale → inverse RHT) lands far outside the assertion tolerance (atol=rtol=0.5) vs the bf16 reference.

Deterministic failure — head_size=64, num_kv_heads=1, num_tokens=42 (seed 0, block_size 16), i.e. [int4-0-16-64-1-42]. Fails on every run with identical mismatch statistics (inputs are seed-deterministic):

AssertionError: Tensor-likes are not close!
Mismatched elements: 35 / 64 (54.7%)
Greatest absolute difference: 3.0546875 at index (0, 38) (up to 0.5 allowed)
Greatest relative difference: 3651.2666015625 at index (0, 52) (up to 0.5 allowed)

Flaky failures — pass/fail flips across runs on identical code:

  • head_size=128, num_kv_heads=1, num_tokens=42 ([int4-0-16-128-1-42]): 64/128 (50.0%) mismatched, abs 4.002 @ (0, 70), rel 25.5 @ (0, 101)
  • head_size=64, num_kv_heads=4, num_tokens=7 ([int4-0-16-64-4-7]): 149/256 (58.2%) mismatched, abs 3.93 @ (2, 26), rel 256 @ (0, 37)

All int8 and fp8 parametrizations, and the remaining int4 combos, pass consistently.

Minimal repro:

pytest tests/quantization/test_per_token_kv_cache.py \
  -k "test_reshape_and_cache_per_token_head and int4-0-16-64-1-42"

Analysis: the kernel under test is vllm.v1.attention.ops.triton_reshape_and_cache_flash.triton_reshape_and_cache_flash_per_token_head_quant (assertion at tests/quantization/test_per_token_kv_cache.py:256). A deterministic 54.7% element mismatch at ~6x the allowed absolute tolerance rules out tolerance noise and points to a correctness defect in the int4 pack/RHT path (candidate causes: zero-point packing or nibble ordering). The two shape-dependent flaky combos additionally suggest a race or uninitialized/leftover cache state whose manifestation depends on shape/head layout.

The qcfg fixture carries IDs int4/int8/fp8; the stacked parametrize resolves bottom-up, so [int4-0-16-64-1-42] = seed 0, block_size 16, head_size 64, num_kv_heads 1, num_tokens 42. All failing combos share seed 0 / block_size 16, but those are the only values in SEEDS/BLOCK_SIZES, so there is no signal there.

Happy to collect a failing-tensor dump if that helps isolate it.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions