[Feature]: NVFP4 KV Cache Support

[yewentao256]

🚀 The feature, motivation and pitch

We haven't supported nvfp4 kv cache yet, let's discuss the plans and possible actions in this issue.

class TritonAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    supported_dtypes: ClassVar[list[torch.dtype]] = [
        torch.float16,
        torch.bfloat16,
        torch.float32,
    ]
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "fp8",
        "fp8_e4m3",
        "fp8_e5m2",
    ]

class FlashInferBackend(AttentionBackend):
    accept_output_buffer: bool = True
    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "fp8",
        "fp8_e4m3",
        "fp8_e5m2",
    ]

ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod
ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE
ModelOptNvFp4Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[mgoin]

What attention kernels are available that support this?

[mgoin]

Reference in TRTLLM NVIDIA/TensorRT-LLM#6244

[baonudesifeizhai]

 python -m vllm.entrypoints.openai.api_server \
  --model nvidia/DeepSeek-R1-NVFP4 \
  --quantization modelopt_fp4 \
  --dtype auto \
  --kv-cache-dtype nvfp4 \
  --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.90 \
  --max-model-len 1024 \
  --enforce-eager \
  --port 8000

will test the kv cahce ...

could response curl chat with normally output

[baonudesifeizhai]

flashinfer-ai/flashinfer#2293 or should just change the code here?

[pavanimajety]

FI Attention Kernels: flashinfer-ai/flashinfer#2363

[pavanimajety]

We should consider several more aspects for adding NVFP4 KV Cache support --

How do we represent NVFP4 KV Cache?
NVFP4 datatype has the actual FP4 values and blockscales - how do we represent a block with this datatype and its corresponding metadata

What does it mean for the storage datatype to be NVFP4 -

1. Support at kernel level in vLLM to store KV values in NVFP4

   • Store BF16 K,V vectors to the NVFP4 KV Cache data-structure after quant to nvfp4
   • Store NVFP4 KV and FP8 Blockscale to NVFP4 KV Cache data-structure
   • Possible optimizations with Gemm Epilogues after projections to directly emit NVFP4 KV for storage

2. Adding support for attention kernels that take NVFP4 KV Cache and do FP8/FP16/BF16 attention.

   • GQA
   • MLA

3. Accuracy / Perf evals for different kernel backends.

What other features need to be supported / tested? (off the top of my mind)

1. Chunked Prefill
2. Prefix Cache
3. Any changes required for scheduling
4. Prefill / Disagg with NVFP4 Cache Transfer
5. Speculative Decoding / Draft model cache
6. General co-design with MXFP4 datatype