Repository navigation
[KV cache] Add an mxfp4 kv_cache_dtype - #64358
amd-dlimpus wants to merge 4 commits into
Conversation
Per-head MXFP4 KV: E2M1 data plus one E8M0 scale per 32 elements, stored as uint8 (head_size/2 + head_size/32 bytes per head row, 68 for 128), with its own KVQuantMode.MXFP4. Backends opt in through supported_kv_cache_dtypes. Signed-off-by: Limpus, David <dlimpus@amd.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
…de.MXFP4 = 16) Move KVQuantMode.MXFP4 from 15 to 16, since 15 is MXFP4_DS_MLA in vllm-project#64120, and make the cache.py, torch_utils.py and kv_cache_interface.py hunks identical to c7f5944 in vllm-project#64198 so whichever PR lands second drops a clean duplicate. Signed-off-by: Limpus, David <dlimpus@amd.com>
…VQuantMode.MXFP4 = 16)" This reverts commit 762bcf8. Signed-off-by: Limpus, David <dlimpus@amd.com>
15 is MXFP4_DS_MLA in vllm-project#64120 and 16 is MXFP4 in vllm-project#64198. Signed-off-by: Limpus, David <dlimpus@amd.com>
Overview
Adds
kv_cache_dtype="mxfp4"for a non-MLA KV cache: E2M1 data with one E8M0 scale per 32 elements (OCP MX), stored asuint8.CacheDTypegains"mxfp4", andSTR_DTYPE_TO_TORCH_DTYPEmaps it totorch.uint8.KVQuantMode.MXFP4 = 17(15 isMXFP4_DS_MLAin [ROCm][DSv4.1] MXFP4 compressed KV cache (mxfp4_ds_mla) on gfx950 #64120, and 16 is used by another open MXFP4 PR), returned byget_kv_quant_mode("mxfp4"), with anis_mxfp4property. It is its own mode, notCUSTOM, so backends can key on it likeNVFP4.mxfp4_kv_cache_full_dim(head_size)gives the packed bytes per head:head_size / 2 + head_size / 32, which is 68 for 128. Unlikenvfp4_kv_cache_full_dim(72 for 128), there is no global scale.Motivation
MiniMax-M3.1 on MI355X (gfx950) can decode from an MXFP4 KV and index cache using the native
V_MFMA_SCALE_F32_*_F8F6F4instructions, which take E8M0 per-32 scales. That's 204 B per token and layer, against 216 B for NVFP4. An accuracy check that rounds K, V and index-K to MXFP4 on today's FP8 path showed no significant loss against the NVFP4-KV baseline (paired exact McNemar, seed 42):Stack
This is the bottom of a stack; the first consumer is the MiniMax-M3 MXFP4 page writer, which is stacked on #63232. Readers and the backend that selects this dtype follow in later PRs. Until then nothing selects
"mxfp4".Validation
tests/config/test_kv_cache_dtype.py::test_mxfp4_kv_cache_dtypecovers the dtype name,CacheConfig, the torch dtype, the quant mode and its properties, and the 68-byte full dim. It passes on CPU.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.