mudler / mudler/vllm.cpp

nvfp4 KV: the store kernel writes K block scales linearly and upstream's own reference dequant un-swizzles them

Open
#2,925 0 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.