[MiniMax-M3.1] FP4 index cache and NVFP4 sparse KV cache

atom/config.py

@@ -1919,6 +1919,74 @@ def qrep_for_step(
     return qrep_enabled and not use_seg_mla and (not is_prefill or prefill_qrep)
 
 
+def m3_reuses_index_topk(text_config) -> bool:
+    """Whether any MiniMax-M3 sparse layer reuses the previous layer's top-k:
+    the model's own schedule (``models/minimax_m3/index_schedule``) over its
+    sparse layers, ``index_skip_topk_offset`` and the pattern's bounds included."""
+    from atom.models.minimax_m3.index_schedule import reuses_index_topk
+
+    return reuses_index_topk(text_config)
+
+
+def m3_fp4_cache_unsupported_reason(
+    gfx: str,
+    index_fp4: bool,
+    kv_nvfp4: bool,
+    tp_size: int,
+    num_heads: int,
+    num_kv_heads: int,
+    reuses_topk: bool,
+    plugin: bool = False,
+) -> str | None:
+    """Why MiniMax-M3's FP4 index cache / NVFP4 sparse KV cannot run here, or None.
+
+    Pure for the same reason as ``indexer_cp_unsupported_reason``. Both caches
+    decode e2m1 with a gfx950-only instruction. The NVFP4 attention kernel
+    serves 16 q heads per kv head on a rank (M3.1 at TP4: 16 / 1) and only layers
+    that run their own indexer: a layer reusing the previous top-k writes no
+    NVFP4 K/V. The vLLM / SGLang plugin bridges bind their own caches and know
+    neither layout, so a plugin deployment would silently keep FP8.
+    """
+    if not (index_fp4 or kv_nvfp4):
+        return None
+    if plugin:
+        return "FP4/NVFP4 MiniMax-M3 caches are not supported in plugin mode"
+    if gfx != "gfx950":
+        return f"FP4/NVFP4 MiniMax-M3 caches need gfx950, not {gfx}"
+    if kv_nvfp4:
+        q_per_kv = (num_heads // tp_size) // max(1, num_kv_heads // tp_size)
+        if q_per_kv != 16:
+            return (
+                "ATOM_M3_SPARSE_KV_NVFP4 needs 16 q heads per kv head on a rank, "
+                f"got {q_per_kv} at TP{tp_size}"
+            )
+        if reuses_topk:
+            return "ATOM_M3_SPARSE_KV_NVFP4 does not cover top-k-reusing layers"
+    return None
+
+
+def m3_nvfp4_stage_unsupported_reason(stage_pages: int, index_pages: int) -> str | None:
+    """Why NVFP4 sparse KV cannot run with this persistent (FP8 / BF16) index
+    pool, or None.
+
+    The NVFP4 K/V is written through a token-major staging buffer of
+    ``stage_pages`` 128-token pages, sized for the largest step a forward can
+    run (max_num_batched_tokens rows, which also bounds every captured decode
+    bucket). AITER's page-128 writer checks the index cache against the KV
+    pages it is given, and a captured bucket can run more rows than the pool
+    holds, so the pool must hold the largest step. An FP4 index is staged too
+    and is not coupled this way.
+    """
+    if index_pages >= stage_pages:
+        return None
+    return (
+        f"ATOM_M3_SPARSE_KV_NVFP4 with an FP8/BF16 index cache needs the KV pool "
+        f"to hold one full step: {index_pages} pages < {stage_pages} staging "
+        "pages. Lower --max-num-batched-tokens / --max-num-seqs, give the KV "
+        "cache more memory, or use --index-cache-dtype fp4."
+    )
+
+
 def indexer_cp_unsupported_reason(
     arches,
     tp_size: int,
@@ -2659,6 +2727,11 @@ def __post_init__(self):
                 self.enable_tbo or self.enable_tbo_decode,
                 is_plugin_mode(),
             )
+            if indexer_cp_off is None and self.index_cache_dtype == "fp4":
+                # The CP context scorer reads FP8/BF16 index keys only.
+                indexer_cp_off = (
+                    "--index-cache-dtype fp4 is not supported by the CP scorer"
+                )
             if indexer_cp_off is not None:
                 logger.warning(
                     "dcp_config.indexer_dcp_only disabled: %s.", indexer_cp_off
@@ -2701,6 +2774,25 @@ def __post_init__(self):
             else:
                 self.index_cache_dtype = self.kv_cache_dtype
 
+        if _is_minimax_m3_config(self.hf_config):
+            text_cfg = getattr(self.hf_config, "text_config", self.hf_config)
+            index_fp4 = self.index_cache_dtype == "fp4"
+            if index_fp4 or envs.ATOM_M3_SPARSE_KV_NVFP4:
+                from aiter.jit.utils.chip_info import get_gfx
+
+                fp4_off = m3_fp4_cache_unsupported_reason(
+                    get_gfx(),
+                    index_fp4,
+                    envs.ATOM_M3_SPARSE_KV_NVFP4,
+                    self.tensor_parallel_size,
+                    text_cfg.num_attention_heads,
+                    text_cfg.num_key_value_heads,
+                    m3_reuses_index_topk(text_cfg),
+                    plugin=self.plugin_config is not None,
+                )
+                if fp4_off is not None:
+                    raise ValueError(fp4_off)
+
         if self.hf_config.model_type == "deepseek_v41_text":
             # CSA2's index plane is paged by the same scorer GLM's is, so a
             # PAGE again has to hold a whole number of `MQA_LOGITS_PRESHUFFLE_ROWS`

atom/kv_transfer/offload/config.py

@@ -379,6 +379,13 @@ def build_page_namespace(
             getattr(config, "speculative_config", None)
         ),
     }
+    from atom.utils import envs
+
+    if envs.ATOM_M3_SPARSE_KV_NVFP4:
+        # NVFP4 sparse K/V bytes at the same block size: only the env says so,
+        # and an FP8 tier's chunk would load into them. Added only when set,
+        # so every other server's namespace is unchanged.
+        document["sparse_kv_format"] = "nvfp4"
     canonical = json.dumps(
         document,
         sort_keys=True,

atom/kv_transfer/offload/mooncake_store/keys.py

@@ -79,13 +79,17 @@ def kv_layout_selectors() -> dict[str, Any]:
     it spawns read the same environment. A switch a model does not use costs a
     miss across servers that differ in it, never another layout's KV.
     """
-    return {
+    selectors = {
         "mla_page_size": int(envs.ATOM_MLA_PAGE_SIZE),
         "triton_mla": bool(envs.ATOM_USE_TRITON_MLA),
         "triton_mla_shuffle_kv": bool(envs.ATOM_USE_TRITON_MLA_SHUFFLE_KV),
         "unified_attn": bool(envs.ATOM_USE_UNIFIED_ATTN),
         "force_attn_triton": bool(envs.ATOM_FORCE_ATTN_TRITON),
     }
+    # MiniMax-M3's NVFP4 sparse K/V; present only when set, so other keys keep.
+    if envs.ATOM_M3_SPARSE_KV_NVFP4:
+        selectors["m3_sparse_kv_nvfp4"] = True
+    return selectors
 
 
 def kv_storage_formats(config: Any) -> dict[str, str]: