FP8 KV Cache Support on SM89 (RTX 4090) — Working Solution

I got FP8 KV cache working on 8x RTX 4090D (SM89) with the patched vLLM fork from this repo. Sharing the fix for anyone interested.

Problem

The fused CUDA kernel fused_minimax_m3_qknorm_rope_kv_insert hardcodes bf16 cache writes (assertion: "kv_cache dtype must match qkv (bf16 cache only)"), so passing --kv-cache-dtype fp8_e4m3 crashes immediately.

Solution

The AMD model path (vllm/models/minimax_m3/amd/model.py) already implements an FP8 bypass — I ported the same pattern to the NVIDIA model.

File to modify

<vllm-install-path>/vllm/models/minimax_m3/nvidia/model.py

In the docker image from this repo (based on toncao/vllm minimax-m3-compressed-tensors branch), the full path is:

/workspace/vllm-awq/vllm/models/minimax_m3/nvidia/model.py

Patch (unified diff)

Apply with patch -p1 from the vllm root, or manually edit the three sections below.

--- a/vllm/models/minimax_m3/nvidia/model.py
+++ b/vllm/models/minimax_m3/nvidia/model.py
@@ -514,6 +514,11 @@
         self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
             self.kv_cache_dtype, vllm_config.model_config
         )
+        # fp8 main-K/V cache: the fused qknorm+rope+kv-insert op is bf16-cache-only
+        # (asserts kv_cache dtype == qkv), so on the fp8 path we run it in
+        # norm+rope-only mode and write the cache via the fp8-capable
+        # reshape_and_cache_flash in _insert_kv. (index cache stays bf16.)
+        self._fp8_kv = "fp8" in self.kv_cache_dtype
         # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main
         # cache (--attention-config '{"indexer_kv_dtype": ...}').
         self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype
@@ -574,27 +579,25 @@
 
     def _insert_kv(
-        self, key: torch.Tensor, value: torch.Tensor, index_key: torch.Tensor
+        self,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        index_key: torch.Tensor,
+        main_slot_mapping: torch.Tensor,
+        index_slot_mapping: torch.Tensor,
     ) -> None:
-        """Write main K/V and index-K into their paged caches.
+        """Write main K/V (fp8-quantizing) and index-K into their paged caches.
 
-        No-op during the profiling run, where caches are not yet bound and
-        ``attn_metadata`` is None.
+        Used only on the fp8-KV path: the fused op is bf16-cache-only, so it
+        runs in norm+rope-only mode and the (already normed/roped) k/v/index_k
+        are written here via ``reshape_and_cache_flash`` (which honors
+        kv_cache_dtype). The index cache stays bf16 (no quant).
         """
-        attn_metadata = get_forward_context().attn_metadata
-        if not isinstance(attn_metadata, dict):
-            return
-        main_meta = attn_metadata[self.layer_name]
-        index_meta = attn_metadata[self.indexer.index_cache.prefix]
-        assert isinstance(main_meta, MiniMaxM3SparseMetadata)
-        assert isinstance(index_meta, MiniMaxM3IndexerMetadata)
-
-        # Identity scale: unused for the bf16 cache, required arg of the op.
         key_cache, value_cache = self.kv_cache.unbind(1)
         scale = torch.ones((), device=key.device)
         ops.reshape_and_cache_flash(
             key.view(-1, self.num_kv_heads, self.head_dim),
             value.view(-1, self.num_kv_heads, self.head_dim),
             key_cache,
             value_cache,
-            main_meta.slot_mapping,
+            main_slot_mapping,
             self.kv_cache_dtype,
             scale,
             scale,
         )
-
         # Index-key cache: single vector per token, scatter by slot.
         idx_cache = self.indexer.index_cache.kv_cache.view(-1, self.idx_head_dim)
-        idx_cache[index_meta.slot_mapping] = index_key.to(idx_cache.dtype)
+        idx_cache[index_slot_mapping] = index_key.to(idx_cache.dtype)
 
@@ -649,6 +652,12 @@
         main_slot_mapping = fwd_slot_mapping[self.layer_name]
         index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
         q = qkv.new_empty((num_tokens, self.q_size))
         index_q = qkv.new_empty((num_tokens, self.index_q_size))
+        # On the fp8-KV path the fused op cannot write the (fp8) cache, so pass
+        # kv_cache/index_cache = None -> insert_kv=False (norm+rope only): it still
+        # de-interleaves q/index_q and rewrites the normed/roped k & index_k in
+        # place in qkv, leaving v raw (correct -- v is never normed/roped). We then
+        # write the cache via _insert_kv below.
+        insert_via_fused = not self._fp8_kv
         ops.fused_minimax_m3_qknorm_rope_kv_insert(
             qkv,
             self.q_norm.weight,
@@ -665,11 +674,24 @@
             main_slot_mapping,
             index_slot_mapping,
-            self.kv_cache,
-            self.indexer.index_cache.kv_cache,
+            self.kv_cache if insert_via_fused else None,
+            self.indexer.index_cache.kv_cache if insert_via_fused else None,
             self.kv_cache.size(2),  # paged-cache block size
             q,
             index_q,
         )
+        if not insert_via_fused:
+            # Extract the normed/roped k, raw v, normed/roped index_k from qkv
+            # ([q | k | v | index_q | index_k], all head_dim=128) and fp8-insert.
+            kv = self.num_kv_heads * self.head_dim
+            # These are strided views into qkv (row stride = full qkv width), but
+            # their last dim is contiguous, so `_insert_kv`'s `.view(-1, nkv,
+            # head_dim)` works on them and `reshape_and_cache_flash` honors the
+            # input stride -- no `.contiguous()` needed.
+            k = qkv[:, self.q_size : self.q_size + kv]
+            v = qkv[:, self.q_size + kv : self.q_size + 2 * kv]
+            ik0 = self.q_size + 2 * kv + self.index_q_size
+            index_k = qkv[:, ik0 : ik0 + self.num_idx_heads * self.idx_head_dim]
+            self._insert_kv(k, v, index_k, main_slot_mapping, index_slot_mapping)
 
         output = torch.empty_like(q)

Launch command

vllm serve cyankiwi/MiniMax-M3-AWQ-INT4 \
    --block-size 128 \
    --kv-cache-dtype fp8_e4m3 \
    --attention-backend TRITON_ATTN \
    --tensor-parallel-size 8

Important: --attention-backend TRITON_ATTN is required on SM89 because FlashInfer does not support block_size=128 on non-Blackwell GPUs (it only advertises [16, 32, 64] when can_use_trtllm_attention returns False). The Triton backend supports MultipleOf(16) and has explicit fp8 support for SM89+.

How it works

The fused CUDA kernel fused_minimax_m3_qknorm_rope_kv_insert combines QK-norm + RoPE + KV cache insertion into one kernel launch. However it only supports writing bf16 to the KV cache.

The fix splits the fp8 path into two steps:

1. Fused kernel runs in norm+rope-only mode (passing kv_cache=None → template parameter INSERT=false)
2. reshape_and_cache_flash (a generic vLLM op that already supports fp8 quantization) handles the KV cache write separately

This is the exact same pattern already used in the AMD/ROCm model path (vllm/models/minimax_m3/amd/model.py).

Hardware requirements

Benefit

FP8 KV cache halves the cache memory usage compared to bf16, allowing significantly longer context lengths on memory-constrained GPUs like the RTX 4090D (48GB).