enable minimax m3 fp8 index cache and fuse index_score_kernel with partial_topk kernel

atom/config.py

@@ -1047,6 +1047,7 @@ class Config:
     kv_cache_block_size: int = 16
     num_kvcache_blocks: int = -1
     kv_cache_dtype: str = "bf16"
+    index_cache_dtype: str = "auto"
     enable_prefix_caching: bool = True
     enable_chunked_prefill: bool = True
     port: int = 8006
@@ -1267,6 +1268,7 @@ def compute_hash(self) -> str:
         factors.append(vllm_factors)
         factors.append(self.tensor_parallel_size)
         factors.append(self.enable_dp_attention)
+        factors.append(self.index_cache_dtype)
         text_config = getattr(self.hf_config, "text_config", self.hf_config)
         factors.append(
             (

atom/model_engine/arg_utils.py

@@ -36,6 +36,7 @@ class EngineArgs:
     enable_prefix_caching: bool = True
     port: int = 8006
     kv_cache_dtype: str = "bf16"
+    index_cache_dtype: str = "auto"
     block_size: int = 16
     max_model_len: Optional[int] = None
     max_num_batched_tokens: int = 16384
@@ -121,6 +122,16 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
             default="bf16",
             help="KV cache type. Default is 'bf16'.",
         )
+        parser.add_argument(
+            "--index-cache-dtype",
+            choices=["auto", "fp8"],
+            type=str,
+            default="auto",
+            help=(
+                "MiniMax-M3 index cache type. 'auto' uses the model dtype; "
+                "'fp8' stores the index cache in FP8 independently of KV cache dtype."
+            ),
+        )
         parser.add_argument(
             "--block-size", type=int, default=16, help="KV cache block size."
         )

atom/model_ops/attention_mha.py

@@ -958,6 +958,7 @@ def __init__(
         local_blocks: int = 0,
         skip_index_topk: bool = False,
         sparse_layer_ordinal: int = -1,
+        index_cache_dtype: str = "auto",
         **kwargs,
     ):
         super().__init__(
@@ -995,6 +996,7 @@ def __init__(
         self.local_blocks = local_blocks
         self.skip_index_topk = skip_index_topk
         self.sparse_layer_ordinal = sparse_layer_ordinal
+        self.index_cache_dtype = index_cache_dtype
         # Bound by AiterAttentionMetadataBuilder.build_kv_cache_tensor (Task 6):
         # the page-128 indexer-key cache. None until the runner binds it.
         self.index_cache: Optional[torch.Tensor] = None
@@ -1164,6 +1166,7 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext):
             index_q,
             slot_mapping,
             kv_cache_dtype=kv_cache_dtype,
+            index_cache_dtype=self.index_cache_dtype,
             k_scale=fused_k_scale,
             v_scale=fused_v_scale,
             asm_layout=True,

atom/model_ops/attentions/aiter_attention.py

@@ -42,6 +42,20 @@ def _is_indexed_sparse_attention(module) -> bool:
     return bool(getattr(impl, "is_indexed_sparse_attention", False))
 
 
+def _resolve_index_cache_dtype(config) -> torch.dtype:
+    from aiter import dtypes
+
+    index_cache_dtype = getattr(config, "index_cache_dtype", "auto")
+    if index_cache_dtype == "auto":
+        return config.torch_dtype
+    if index_cache_dtype == "fp8":
+        return dtypes.d_dtypes["fp8"]
+    raise ValueError(
+        "MiniMax-M3 index_cache_dtype must be 'auto' or 'fp8', "
+        f"got {index_cache_dtype!r}"
+    )
+
+
 @triton.jit
 def _mtp_prepare_decode_metadata_kernel(
     context_lens_ptr,
@@ -480,11 +494,12 @@ def compute_block_bytes(self) -> int:
                 1 for enabled in sparse_cfg.get("sparse_attention_freq", []) if enabled
             )
             index_dim = sparse_cfg["sparse_index_dim"]
+            index_cache_dtype = _resolve_index_cache_dtype(config)
             block_bytes += (
                 sparse_layers
                 * runner.physical_block_size
                 * index_dim
-                * torch.empty((), dtype=config.torch_dtype).element_size()
+                * torch.empty((), dtype=index_cache_dtype).element_size()
             )
         return block_bytes
 
@@ -540,12 +555,13 @@ def allocate_kv_cache_tensors(
             sparse_layers = sum(
                 1 for enabled in sparse_cfg.get("sparse_attention_freq", []) if enabled
             )
+            index_cache_dtype = _resolve_index_cache_dtype(config)
             tensors["sparse_attention_index_cache"] = torch.zeros(
                 sparse_layers,
                 runner.num_physical_kvcache_blocks,
                 runner.physical_block_size,
                 sparse_cfg["sparse_index_dim"],
-                dtype=config.torch_dtype,
+                dtype=index_cache_dtype,
                 device="cuda",
             )
             tensors["_sparse_attention_cache_next"] = 0

atom/model_ops/minimax_m3/index_topk.py

@@ -152,7 +152,10 @@ def _index_block_score_kernel(
             + off_k[None, :] * stride_ik_pos
             + off_d[:, None] * stride_ik_d,
         )
-        qk = tl.dot(q, k) * sm_scale_log2e
+        if k.dtype.is_fp8():
+            qk = tl.dot(q.to(k.dtype), k, out_dtype=tl.float32) * sm_scale_log2e
+        else:
+            qk = tl.dot(q, k, out_dtype=tl.float32) * sm_scale_log2e
         # apply causal mask as needed
         if q_start < i + BLOCK_SIZE_K:
             qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf"))
@@ -420,20 +423,19 @@ def _decode_index_score_kernel(
     )  # [D]
     for blk in tl.range(chunk_start_block, chunk_end_block):
         page = tl.load(bt_row + blk).to(tl.int64)
-        pos = blk * BLOCK_SIZE_K + off_k
-        pos_mask = pos < causal_len
         k = tl.load(
             ik_cache_ptr
             + page * stride_ik_blk
             + off_k[None, :] * stride_ik_pos
             + off_d[:, None] * stride_ik_d,
-            mask=d_mask[:, None] & pos_mask[None, :],
+            mask=d_mask[:, None],
             other=0.0,
-        ).to(
-            tl.float32
-        )  # [D, N]
+        )
+        k = k.to(tl.float32)  # [D, N]
         qk = tl.sum(q[:, None] * k, axis=0) * sm_scale_log2e  # [N]
-        qk = tl.where(pos_mask, qk, float("-inf"))
+        if (blk + 1) * BLOCK_SIZE_K > causal_len:
+            pos = blk * BLOCK_SIZE_K + off_k
+            qk = tl.where(pos < causal_len, qk, float("-inf"))
         score = tl.max(qk, axis=0)  # one score for this 128-block
         is_init = blk < init_blocks
         is_local = (blk >= local_start) & (blk < num_blocks)
@@ -582,6 +584,128 @@ def _topk_index_partial_kernel(
     tl.store(ti_ptrs, final_idx)
 
 
+@triton.heuristics(
+    {
+        "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]),
+        "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
+    }
+)
+@triton.jit
+def _decode_index_score_topk_partial_kernel(
+    q_ptr,  # idx_q: [total_q, num_idx_heads, head_dim]
+    ik_cache_ptr,  # index-K cache: [num_blocks, 128, head_dim]
+    ts_partial_ptr,  # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
+    ti_partial_ptr,  # partial idx out (1-indexed global, 0=invalid): same shape
+    block_table_ptr,  # [num_reqs, max_blocks]
+    seq_lens,  # [batch]
+    num_idx_heads,
+    total_q,  # batch * max_q (one row per query token)
+    head_dim,
+    init_blocks,
+    local_blocks,
+    sm_scale,
+    topk: tl.constexpr,
+    chunk_blocks: tl.constexpr,
+    stride_q_n,
+    stride_q_h,
+    stride_q_d,
+    stride_ik_blk,
+    stride_ik_pos,
+    stride_ik_d,
+    stride_ts_c,
+    stride_ts_h,
+    stride_ts_b,
+    stride_ts_t,
+    stride_ti_c,
+    stride_ti_h,
+    stride_ti_b,
+    stride_ti_t,
+    stride_bt_b,
+    MAX_Q: tl.constexpr,  # query tokens per request (num_spec + 1; 1 == plain decode)
+    BLOCK_SIZE_K: tl.constexpr,  # == SPARSE_BLOCK_SIZE (128)
+    BLOCK_SIZE_T: tl.constexpr,
+    BLOCK_SIZE_D: tl.constexpr,
+):
+    sm_scale_log2e = sm_scale * 1.4426950409
+    pid_t = tl.program_id(0)  # global query-token row
+    pid_h = tl.program_id(1)
+    pid_chunk = tl.program_id(2)
+
+    pid_b = pid_t // MAX_Q
+    tok = pid_t % MAX_Q
+    seq_len = tl.load(seq_lens + pid_b)
+    causal_len = seq_len - MAX_Q + tok + 1
+    num_blocks = (causal_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
+
+    chunk_start = pid_chunk * chunk_blocks
+    chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks)
+    chunk_actual = tl.maximum(chunk_end - chunk_start, 0)
+
+    off_t = tl.arange(0, BLOCK_SIZE_T)
+    off_k = tl.arange(0, BLOCK_SIZE_K)
+    off_d = tl.arange(0, BLOCK_SIZE_D)
+    d_mask = off_d < head_dim
+
+    topk_score = tl.full((BLOCK_SIZE_T,), -1e30, dtype=tl.float32)
+    topk_idx = tl.full((BLOCK_SIZE_T,), 0, dtype=tl.int32)
+
+    bt_row = block_table_ptr + pid_b * stride_bt_b
+    local_start = tl.maximum(0, num_blocks - local_blocks)
+    q = tl.load(
+        q_ptr + pid_t * stride_q_n + pid_h * stride_q_h + off_d * stride_q_d,
+        mask=d_mask,
+        other=0.0,
+    ).to(tl.float32)
+
+    for i in tl.range(0, chunk_actual):
+        blk = chunk_start + i
+        page = tl.load(bt_row + blk).to(tl.int64)
+        k = tl.load(
+            ik_cache_ptr
+            + page * stride_ik_blk
+            + off_k[None, :] * stride_ik_pos
+            + off_d[:, None] * stride_ik_d,
+            mask=d_mask[:, None],
+            other=0.0,
+        )
+        k = k.to(tl.float32)
+        qk = tl.sum(q[:, None] * k, axis=0) * sm_scale_log2e
+        if (blk + 1) * BLOCK_SIZE_K > causal_len:
+            pos = blk * BLOCK_SIZE_K + off_k
+            qk = tl.where(pos < causal_len, qk, float("-inf"))
+        score = tl.max(qk, axis=0)
+        is_init = blk < init_blocks
+        is_local = (blk >= local_start) & (blk < num_blocks)
+        score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score))
+        score = tl.where(score == score, score, -1e30)
+
+        # Maintain an unordered top-k vector. The merge kernel sorts all partial
+        # candidates later, so replacing any current minimum is sufficient.
+        min_score = tl.min(topk_score, axis=0)
+        min_mask = topk_score == min_score
+        first_min = tl.cumsum(min_mask.to(tl.int32), axis=0) == 1
+        replace = (score > min_score) & min_mask & first_min
+        topk_score = tl.where(replace, score, topk_score)
+        topk_idx = tl.where(replace, (blk + 1).to(tl.int32), topk_idx)
+
+    ts_ptrs = (
+        ts_partial_ptr
+        + pid_chunk * stride_ts_c
+        + pid_h * stride_ts_h
+        + pid_t * stride_ts_b
+        + off_t * stride_ts_t
+    )
+    ti_ptrs = (
+        ti_partial_ptr
+        + pid_chunk * stride_ti_c
+        + pid_h * stride_ti_h
+        + pid_t * stride_ti_b
+        + off_t * stride_ti_t
+    )
+    tl.store(ts_ptrs, topk_score)
+    tl.store(ti_ptrs, topk_idx)
+
+
 @triton.heuristics(
     {
         "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]),
@@ -908,55 +1032,16 @@ def minimax_m3_index_topk_decode(
         total_q % max_query_len == 0
     ), f"total_q {total_q} not divisible by max_query_len {max_query_len}"
     max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
-    score = torch.empty(
-        (num_idx_heads, total_q, max_block),
-        dtype=torch.float32,
-        device=idx_q.device,
-    )
-    # split-K over seq blocks; chunk count depends only on shape constants so
-    # the grid is fixed within a cuda graph.
-    TARGET_GRID = 4096
-    MAX_NUM_KV_CHUNKS = 256
-    target = max(
-        1, min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, total_q * num_idx_heads))
-    )
-    num_kv_chunks = 1 << (target.bit_length() - 1)
-    grid_score = (total_q * num_kv_chunks, num_idx_heads)
-    _decode_index_score_kernel[grid_score](
-        idx_q,
-        index_kv_cache,
-        score,
-        block_table,
-        seq_lens,
-        num_idx_heads,
-        total_q,
-        head_dim,
-        init_blocks,
-        local_blocks,
-        sm_scale,
-        idx_q.stride(0),
-        idx_q.stride(1),
-        idx_q.stride(2),
-        index_kv_cache.stride(0),
-        index_kv_cache.stride(1),
-        index_kv_cache.stride(2),
-        score.stride(0),
-        score.stride(1),
-        score.stride(2),
-        block_table.stride(0),
-        MAX_Q=max_query_len,
-        BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
-        NUM_KV_CHUNKS=num_kv_chunks,
-    )
-
     topk_idx = torch.empty(
         (num_idx_heads, total_q, topk),
         dtype=torch.int32,
         device=idx_q.device,
     )
     # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
     # pow2(num_topk_chunks * pow2(topk)) candidates.
-    TOPK_TARGET_GRID = 64
+    # Fused score+partial-topk uses topk chunks as the score split-K dimension.
+    # Keep enough chunks to avoid serializing long contexts inside too few CTAs.
+    TOPK_TARGET_GRID = 2048
     MAX_NUM_TOPK_CHUNKS = 16
     topk_target = max(
         1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, total_q * num_idx_heads))
@@ -980,18 +1065,27 @@ def minimax_m3_index_topk_decode(
         dtype=torch.int32,
         device=idx_q.device,
     )
-    _topk_index_partial_kernel[(total_q, num_idx_heads, num_topk_chunks)](
-        score,
+    _decode_index_score_topk_partial_kernel[(total_q, num_idx_heads, num_topk_chunks)](
+        idx_q,
+        index_kv_cache,
         topk_score_partial,
         topk_idx_partial,
+        block_table,
         seq_lens,
-        SPARSE_BLOCK_SIZE,
+        num_idx_heads,
+        total_q,
+        head_dim,
+        init_blocks,
+        local_blocks,
+        sm_scale,
         topk,
         chunk_blocks,
-        max_query_len,  # MAX_Q
-        score.stride(0),
-        score.stride(1),
-        score.stride(2),
+        idx_q.stride(0),
+        idx_q.stride(1),
+        idx_q.stride(2),
+        index_kv_cache.stride(0),
+        index_kv_cache.stride(1),
+        index_kv_cache.stride(2),
         topk_score_partial.stride(0),
         topk_score_partial.stride(1),
         topk_score_partial.stride(2),
@@ -1000,6 +1094,9 @@ def minimax_m3_index_topk_decode(
         topk_idx_partial.stride(1),
         topk_idx_partial.stride(2),
         topk_idx_partial.stride(3),
+        block_table.stride(0),
+        MAX_Q=max_query_len,
+        BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
     )
     # The fused emit now produces one row per (token, kv-head): the ASM/gluon path
     # collapses kv-head into the row dim. sparse_bt/ctx are sized total_q*num_idx_heads