nvfp4 KV: the store kernel writes K block scales linearly and upstream's own reference dequant un-swizzles them
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 423
- Forks
- 53
- Avg merge
- 20h 26m
- Merged PRs (30d)
- 310
Description
Row: KV-NVFP4-TURBO
Owed under ## Owed in .agents/specs/nvfp4-kv-cache.md.
What was measured
At the parity pin 5559679229bc961848b121ccdeaa8fa5d79bec98, upstream's NVFP4 KV
store kernel writes the K block scales linearly and the V block scales
4x4-swizzled:
csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu:157-172
if (kv == 0) {
scale_dst = scale_block + head * scale_head_stride +
block_offset * scale_block_offset_stride + scale_idx;
} else {
int swizzled_offset = swizzle_scale_offset(block_offset, scale_idx, scale_dim);
...
}
The swizzle (:26-39) is stored_t = (t/4)*4 + s/(S/4),
stored_s = (s%(S/4))*4 + t%4, where S = head_size/16.
Upstream's own reference dequant un-swizzles both sides:
tests/kernels/quantization/nvfp4_utils.py:90-146 (dequant_nvfp4_kv_cache)
sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4)
perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
That permutation is exactly the inverse of swizzle_scale_offset, and
tests/kernels/attention/test_cache.py:357-372 feeds the K side through it:
result_key_cache = dequant_nvfp4_cache_nhd(
nvfp4_key_data, key_scale_cache, k_scale.item()
)
The swizzle is the identity only where S == 4 and s == t % 4, so for any
head_size other than 64 the store and the reference disagree about where a K
block scale lives. One of the two is wrong.
Why it is not settled here
test_cache.py:205-207 skips the entire NVFP4 case below
current_platform.has_device_capability(100), and
FlashInferBackend.supports_kv_cache_dtype
(vllm/v1/attention/backends/flashinfer.py:875-882) admits nvfp4 only on
capability family 100. GATE_DEVICE on this fleet is a GB10, sm_121a, family
121. The disagreement has therefore never executed on any machine reachable
from here, and the pinned oracle refuses the dtype on our only GPU.
Why it blocks the port
KV-NVFP4-TURBO cannot mirror a layout upstream does not agree with itself
about. Choosing one reading is inventing a format, which is what "mirror vLLM"
exists to prevent. This must be settled by running the pinned vLLM on an SM100
box, or by reading the trtllm-gen kernel's own block-scale addressing, before
any nvfp4 KV store lands here.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Compare csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu with tests/kernels/quantization/nvfp4_utils.py and tests/kernels/attention/test_cache.py, starting with the scale addressing and inverse permutation. Check whether the pinned vLLM runs on an SM100 system, or inspect the trtllm-gen kernel's block-scale addressing. Done means resolving the K-scale layout disagreement before the NVFP4 KV store is implemented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100