From 9049bdbf64a71b8a154933b9349cf779915aff7a Mon Sep 17 00:00:00 2001 From: ganyi Date: Sun, 5 Jul 2026 12:48:01 +0000 Subject: [PATCH 01/11] fuse score topk Signed-off-by: ganyi --- atom/model_ops/minimax_m3/index_topk.py | 200 +++++++++++++++++------- 1 file changed, 147 insertions(+), 53 deletions(-) diff --git a/atom/model_ops/minimax_m3/index_topk.py b/atom/model_ops/minimax_m3/index_topk.py index be74e05379..a345c13bfd 100644 --- a/atom/model_ops/minimax_m3/index_topk.py +++ b/atom/model_ops/minimax_m3/index_topk.py @@ -420,20 +420,20 @@ 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] 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 +582,127 @@ 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, + ).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,47 +1029,6 @@ 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, @@ -956,7 +1036,9 @@ def minimax_m3_index_topk_decode( ) # 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 +1062,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 +1091,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 From 44b930c9f8c660fd08dafce2c914df9d28859830 Mon Sep 17 00:00:00 2001 From: ganyi Date: Mon, 6 Jul 2026 05:19:10 +0000 Subject: [PATCH 02/11] fp8 index cache Signed-off-by: ganyi --- atom/model_ops/attentions/aiter_attention.py | 7 ++++++- atom/model_ops/minimax_m3/index_topk.py | 15 ++++++++++----- .../plugin/vllm/attention/minimax_m3_attnetion.py | 2 +- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index d24e71f747..66fa2eb682 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -540,12 +540,17 @@ def allocate_kv_cache_tensors( sparse_layers = sum( 1 for enabled in sparse_cfg.get("sparse_attention_freq", []) if enabled ) + index_cache_dtype = ( + dtypes.d_dtypes[config.kv_cache_dtype] + if str(config.kv_cache_dtype).startswith("fp8") + else config.torch_dtype + ) 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 diff --git a/atom/model_ops/minimax_m3/index_topk.py b/atom/model_ops/minimax_m3/index_topk.py index a345c13bfd..dd7acd013e 100644 --- a/atom/model_ops/minimax_m3/index_topk.py +++ b/atom/model_ops/minimax_m3/index_topk.py @@ -129,6 +129,7 @@ def _index_block_score_kernel( order=(1, 0), ) q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") + q_f32 = q.to(tl.float32) q_start = prefix_len + pid_q * BLOCK_SIZE_Q off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len @@ -152,7 +153,11 @@ 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(): + k = k.to(tl.float32) + qk = tl.dot(q_f32, 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")) @@ -427,9 +432,8 @@ def _decode_index_score_kernel( + off_d[:, None] * stride_ik_d, 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] if (blk + 1) * BLOCK_SIZE_K > causal_len: pos = blk * BLOCK_SIZE_K + off_k @@ -665,7 +669,8 @@ def _decode_index_score_topk_partial_kernel( + off_d[:, None] * stride_ik_d, mask=d_mask[:, None], other=0.0, - ).to(tl.float32) + ) + 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 diff --git a/atom/plugin/vllm/attention/minimax_m3_attnetion.py b/atom/plugin/vllm/attention/minimax_m3_attnetion.py index 51923c3de2..3fb9b2bbce 100644 --- a/atom/plugin/vllm/attention/minimax_m3_attnetion.py +++ b/atom/plugin/vllm/attention/minimax_m3_attnetion.py @@ -236,7 +236,7 @@ def __init__( self.index_cache_layer = MiniMaxM3SparseIndexerCache( layer_name=f"{self.layer_name}.index_cache", head_dim=index_head_dim, - kv_cache_dtype="auto", + kv_cache_dtype=cache_dtype if str(cache_dtype).startswith("fp8") else "auto", ) _register_vllm_static_forward_context(self) From 291ed54911331fa66105d325ab9518e1bc56c668 Mon Sep 17 00:00:00 2001 From: ganyi Date: Mon, 6 Jul 2026 07:24:06 +0000 Subject: [PATCH 03/11] fp8 index score calculate in fp8 Signed-off-by: ganyi --- atom/model_ops/minimax_m3/index_topk.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/atom/model_ops/minimax_m3/index_topk.py b/atom/model_ops/minimax_m3/index_topk.py index dd7acd013e..71f457b82e 100644 --- a/atom/model_ops/minimax_m3/index_topk.py +++ b/atom/model_ops/minimax_m3/index_topk.py @@ -129,7 +129,6 @@ def _index_block_score_kernel( order=(1, 0), ) q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") - q_f32 = q.to(tl.float32) q_start = prefix_len + pid_q * BLOCK_SIZE_Q off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len @@ -154,8 +153,7 @@ def _index_block_score_kernel( + off_d[:, None] * stride_ik_d, ) if k.dtype.is_fp8(): - k = k.to(tl.float32) - qk = tl.dot(q_f32, k, out_dtype=tl.float32) * sm_scale_log2e + 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 From 95512d94de4473caf9feb113a2ee4f040c521c03 Mon Sep 17 00:00:00 2001 From: ganyi Date: Tue, 7 Jul 2026 02:21:51 +0000 Subject: [PATCH 04/11] black Signed-off-by: ganyi --- atom/plugin/vllm/attention/minimax_m3_attnetion.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/atom/plugin/vllm/attention/minimax_m3_attnetion.py b/atom/plugin/vllm/attention/minimax_m3_attnetion.py index 3fb9b2bbce..1c99782579 100644 --- a/atom/plugin/vllm/attention/minimax_m3_attnetion.py +++ b/atom/plugin/vllm/attention/minimax_m3_attnetion.py @@ -236,7 +236,9 @@ def __init__( self.index_cache_layer = MiniMaxM3SparseIndexerCache( layer_name=f"{self.layer_name}.index_cache", head_dim=index_head_dim, - kv_cache_dtype=cache_dtype if str(cache_dtype).startswith("fp8") else "auto", + kv_cache_dtype=( + cache_dtype if str(cache_dtype).startswith("fp8") else "auto" + ), ) _register_vllm_static_forward_context(self) From e432f541bdeee9840576ff43c43c8fd7593666aa Mon Sep 17 00:00:00 2001 From: ganyi Date: Tue, 7 Jul 2026 04:25:37 +0000 Subject: [PATCH 05/11] add index_cache_dtype as parameter Signed-off-by: ganyi --- atom/config.py | 2 ++ atom/model_engine/arg_utils.py | 11 ++++++++++ atom/model_ops/attention_mha.py | 3 +++ atom/model_ops/attentions/aiter_attention.py | 23 +++++++++++++++----- atom/models/minimax_m3.py | 7 ++++++ 5 files changed, 40 insertions(+), 6 deletions(-) diff --git a/atom/config.py b/atom/config.py index 52cf7652e6..0a0fcce95e 100644 --- a/atom/config.py +++ b/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( ( diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 2ec68f4cb5..bba853982c 100644 --- a/atom/model_engine/arg_utils.py +++ b/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." ) diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 0ef0b38fe7..0856529523 100644 --- a/atom/model_ops/attention_mha.py +++ b/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, diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index 66fa2eb682..1518973682 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/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_minimax_m3_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_minimax_m3_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,11 +555,7 @@ def allocate_kv_cache_tensors( sparse_layers = sum( 1 for enabled in sparse_cfg.get("sparse_attention_freq", []) if enabled ) - index_cache_dtype = ( - dtypes.d_dtypes[config.kv_cache_dtype] - if str(config.kv_cache_dtype).startswith("fp8") - else config.torch_dtype - ) + index_cache_dtype = _resolve_minimax_m3_index_cache_dtype(config) tensors["sparse_attention_index_cache"] = torch.zeros( sparse_layers, runner.num_physical_kvcache_blocks, diff --git a/atom/models/minimax_m3.py b/atom/models/minimax_m3.py index 31efec5ef9..111c282b1d 100644 --- a/atom/models/minimax_m3.py +++ b/atom/models/minimax_m3.py @@ -331,6 +331,7 @@ def __init__( quant_config: Optional[QuantizationConfig] = None, prefix: str = "", cache_config: str = "bf16", + index_cache_config: str = "auto", ) -> None: super().__init__() self.layer_num = layer_id @@ -409,6 +410,7 @@ def __init__( quant_config: Optional[QuantizationConfig] = None, prefix: str = "", cache_config: str = "bf16", + index_cache_config: str = "auto", ) -> None: super().__init__() self.is_indexed_sparse_attention = True @@ -522,6 +524,7 @@ def __init__( local_blocks=self.local_blocks, skip_index_topk=self.skip_index_topk, sparse_layer_ordinal=self.sparse_layer_ordinal, + index_cache_dtype=index_cache_config, ) def forward( @@ -553,6 +556,7 @@ def __init__( config: PretrainedConfig, prefix: str, cache_config: str = "bf16", + index_cache_config: str = "auto", quant_config: Optional[QuantizationConfig] = None, params_dtype: Optional[torch.dtype] = None, layer_num: int = 0, @@ -569,6 +573,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", cache_config=cache_config, + index_cache_config=index_cache_config, ) self.is_moe_layer = _is_moe_layer(config, layer_num) @@ -667,6 +672,7 @@ def __init__( config = _get_text_config(atom_config.hf_config) self.config = config cache_config = atom_config.kv_cache_dtype + index_cache_config = atom_config.index_cache_dtype quant_config = atom_config.quant_config if get_pp_group().is_first_rank: @@ -683,6 +689,7 @@ def __init__( config, prefix, cache_config=cache_config, + index_cache_config=index_cache_config, quant_config=quant_config, layer_num=layer_num, params_dtype=atom_config.torch_dtype, From fab5a7abee716ce4b7f064b267178a2ed87c91ed Mon Sep 17 00:00:00 2001 From: ganyi Date: Tue, 7 Jul 2026 04:55:46 +0000 Subject: [PATCH 06/11] add fp8 kv cache and fp8 index cache to recipe Signed-off-by: ganyi --- recipes/MiniMax-M3.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/recipes/MiniMax-M3.md b/recipes/MiniMax-M3.md index cd1f89b593..b8872c96c2 100644 --- a/recipes/MiniMax-M3.md +++ b/recipes/MiniMax-M3.md @@ -30,6 +30,8 @@ python -m atom.entrypoints.openai_server \ --max-model-len 32768 \ --max-num-seqs 128 \ --max-num-batched-tokens 32768 \ + --kv_cache_dtype fp8 \ + --index-cache-dtype fp8 \ --online_quant_config '{"global_quant_config": "ptpc_fp8", "exclude_layer": ["lm_head", "model.embed_tokens", "vision_tower", "multi_modal_projector", "patch_merge_mlp", "*block_sparse_moe"]}' \ --no-enable_prefix_caching \ --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' 2>&1 | tee "${run_name}-server.log" @@ -57,6 +59,8 @@ python -m atom.entrypoints.openai_server \ --block-size 128 \ --max-model-len 32768 \ --max-num-seqs 128 \ + --kv_cache_dtype fp8 \ + --index-cache-dtype fp8 \ --max-num-batched-tokens 32768 \ --online_quant_config '{"global_quant_config": "ptpc_fp8", "exclude_layer": ["lm_head", "model.embed_tokens", "vision_tower", "multi_modal_projector", "patch_merge_mlp", "*block_sparse_moe"]}' \ --no-enable_prefix_caching \ From ba2030aa922949f3470f3a71a07b10d7ee6de219 Mon Sep 17 00:00:00 2001 From: ganyi Date: Wed, 15 Jul 2026 09:42:57 +0000 Subject: [PATCH 07/11] resolve comments Signed-off-by: ganyi --- atom/model_ops/attentions/aiter_attention.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index 1518973682..8afde37b29 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -42,7 +42,7 @@ def _is_indexed_sparse_attention(module) -> bool: return bool(getattr(impl, "is_indexed_sparse_attention", False)) -def _resolve_minimax_m3_index_cache_dtype(config) -> torch.dtype: +def _resolve_index_cache_dtype(config) -> torch.dtype: from aiter import dtypes index_cache_dtype = getattr(config, "index_cache_dtype", "auto") @@ -494,7 +494,7 @@ 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_minimax_m3_index_cache_dtype(config) + index_cache_dtype = _resolve_index_cache_dtype(config) block_bytes += ( sparse_layers * runner.physical_block_size @@ -555,7 +555,7 @@ def allocate_kv_cache_tensors( sparse_layers = sum( 1 for enabled in sparse_cfg.get("sparse_attention_freq", []) if enabled ) - index_cache_dtype = _resolve_minimax_m3_index_cache_dtype(config) + index_cache_dtype = _resolve_index_cache_dtype(config) tensors["sparse_attention_index_cache"] = torch.zeros( sparse_layers, runner.num_physical_kvcache_blocks, From e57b2976e6eef3646c283a52afd65e60f1ee83af Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 15 Jul 2026 09:47:40 +0000 Subject: [PATCH 08/11] Fix SGLang readiness check to validate expected model --- .github/scripts/atom_sglang_test.sh | 50 +++++++++++++++++++++++------ 1 file changed, 41 insertions(+), 9 deletions(-) diff --git a/.github/scripts/atom_sglang_test.sh b/.github/scripts/atom_sglang_test.sh index b8c227d5d6..35fcd867f8 100644 --- a/.github/scripts/atom_sglang_test.sh +++ b/.github/scripts/atom_sglang_test.sh @@ -107,17 +107,10 @@ emit_new_sglang_logs() { } wait_server_ready() { + local expected_model_path="${1:-}" echo "" echo "========== Waiting for SGLang server (${MODEL_NAME}) ==========" for ((i=1; i<=MAX_WAIT_RETRIES; i++)); do - if curl -fsS "http://127.0.0.1:${SGLANG_PORT}/v1/models" >/dev/null 2>&1; then - emit_new_sglang_logs - echo "SGLang server is ready for ${MODEL_NAME}." - return 0 - fi - - emit_new_sglang_logs - if [[ -f "${SGLANG_PID_FILE}" ]]; then local pid pid=$(cat "${SGLANG_PID_FILE}") @@ -129,6 +122,45 @@ wait_server_ready() { fi fi + local models_response + models_response=$(curl -fsS "http://127.0.0.1:${SGLANG_PORT}/v1/models" 2>/dev/null || true) + if [[ -n "${models_response}" ]]; then + if [[ -z "${expected_model_path}" ]] || \ + MODELS_RESPONSE="${models_response}" EXPECTED_MODEL_PATH="${expected_model_path}" python3 - <<'PY' +import json +import os +import sys + + +def normalize_model_id(model_id: str) -> str: + model_id = model_id.rstrip("/") + if model_id.startswith("/models/"): + return model_id[len("/models/") :] + return model_id + + +expected_model = normalize_model_id(os.environ["EXPECTED_MODEL_PATH"]) +try: + payload = json.loads(os.environ["MODELS_RESPONSE"]) +except Exception: + sys.exit(1) + +served_models = { + normalize_model_id(item.get("id", "")) + for item in payload.get("data", []) + if isinstance(item, dict) and isinstance(item.get("id"), str) +} + +sys.exit(0 if expected_model in served_models else 1) +PY + then + emit_new_sglang_logs + echo "SGLang server is ready for ${MODEL_NAME}." + return 0 + fi + fi + + emit_new_sglang_logs echo "Waiting for SGLang server... (${i}/${MAX_WAIT_RETRIES})" sleep "${WAIT_INTERVAL_SEC}" done @@ -217,7 +249,7 @@ PY echo "Server PID: $(cat "${SGLANG_PID_FILE}")" if [[ "${wait_for_ready}" == "1" ]]; then - wait_server_ready + wait_server_ready "${resolved_model_path}" fi } From d745f4530e6600abb520f518dfb0b17d2b00124e Mon Sep 17 00:00:00 2001 From: ganyi Date: Fri, 17 Jul 2026 02:00:24 +0000 Subject: [PATCH 09/11] make index-cache-dtype default as same as kv-cache-dtype Signed-off-by: ganyi --- atom/config.py | 5 ++- atom/model_engine/arg_utils.py | 16 +++++---- atom/model_ops/attention_mha.py | 6 ++-- atom/model_ops/attentions/aiter_attention.py | 12 +++---- tests/test_arg_utils_spec.py | 36 ++++++++++++++++++++ 5 files changed, 59 insertions(+), 16 deletions(-) diff --git a/atom/config.py b/atom/config.py index 0a0fcce95e..54e688e585 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1047,7 +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" + index_cache_dtype: str | None = None enable_prefix_caching: bool = True enable_chunked_prefill: bool = True port: int = 8006 @@ -1099,6 +1099,9 @@ def _set_cudagraph_sizes(self): self.graph_bs = cuda_graph_sizes def __post_init__(self): + if self.index_cache_dtype is None: + self.index_cache_dtype = self.kv_cache_dtype + if isinstance(self.compilation_config, dict): self.compilation_config = CompilationConfig(**self.compilation_config) # assert os.path.isdir(self.model) diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index bba853982c..57f9d07ce9 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -36,7 +36,7 @@ class EngineArgs: enable_prefix_caching: bool = True port: int = 8006 kv_cache_dtype: str = "bf16" - index_cache_dtype: str = "auto" + index_cache_dtype: Optional[str] = None block_size: int = 16 max_model_len: Optional[int] = None max_num_batched_tokens: int = 16384 @@ -63,6 +63,10 @@ class EngineArgs: online_quant_config: Optional[dict] = None hf_overrides: Optional[dict] = None + def __post_init__(self) -> None: + if self.index_cache_dtype is None: + self.index_cache_dtype = self.kv_cache_dtype + @staticmethod def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: """Add engine arguments to an argument parser.""" @@ -124,13 +128,11 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: ) parser.add_argument( "--index-cache-dtype", - choices=["auto", "fp8"], + "--index_cache_dtype", + choices=["bf16", "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." - ), + default=None, + help="Index cache type. Defaults to --kv_cache_dtype.", ) parser.add_argument( "--block-size", type=int, default=16, help="KV cache block size." diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 0856529523..10ad232b6e 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -958,7 +958,7 @@ def __init__( local_blocks: int = 0, skip_index_topk: bool = False, sparse_layer_ordinal: int = -1, - index_cache_dtype: str = "auto", + index_cache_dtype: str | None = None, **kwargs, ): super().__init__( @@ -996,7 +996,9 @@ 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 + self.index_cache_dtype = ( + index_cache_dtype if index_cache_dtype is not None else kv_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 diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index 8afde37b29..83b8e85de6 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -45,13 +45,13 @@ def _is_indexed_sparse_attention(module) -> bool: 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"] + index_cache_dtype = getattr(config, "index_cache_dtype", None) + if index_cache_dtype is None: + index_cache_dtype = getattr(config, "kv_cache_dtype", "bf16") + if index_cache_dtype in ("bf16", "fp8"): + return dtypes.d_dtypes[index_cache_dtype] raise ValueError( - "MiniMax-M3 index_cache_dtype must be 'auto' or 'fp8', " + "index_cache_dtype must be 'bf16' or 'fp8', " f"got {index_cache_dtype!r}" ) diff --git a/tests/test_arg_utils_spec.py b/tests/test_arg_utils_spec.py index 9f859dec37..b1f2cf0554 100644 --- a/tests/test_arg_utils_spec.py +++ b/tests/test_arg_utils_spec.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: MIT # Regression tests for speculative-config validation in EngineArgs._get_engine_kwargs. +import argparse import sys from unittest.mock import MagicMock, patch @@ -73,3 +74,38 @@ def test_method_mtp_valid_tokens_builds_speculative_config(self): num_speculative_tokens=3, ) assert kwargs["speculative_config"] is fake_spec_config + + +class TestEngineArgsIndexCacheDtype: + """Regression tests for index-cache dtype CLI parity with KV cache dtype.""" + + def test_index_cache_dtype_defaults_to_kv_cache_dtype(self): + parser = argparse.ArgumentParser() + EngineArgs.add_cli_args(parser) + + args = EngineArgs.from_cli_args( + parser.parse_args(["--kv_cache_dtype", "fp8"]) + ) + + assert args.kv_cache_dtype == "fp8" + assert args.index_cache_dtype == "fp8" + + def test_index_cache_dtype_accepts_dashed_spelling(self): + parser = argparse.ArgumentParser() + EngineArgs.add_cli_args(parser) + + args = EngineArgs.from_cli_args( + parser.parse_args(["--index-cache-dtype", "fp8"]) + ) + + assert args.index_cache_dtype == "fp8" + + def test_index_cache_dtype_accepts_underscore_spelling(self): + parser = argparse.ArgumentParser() + EngineArgs.add_cli_args(parser) + + args = EngineArgs.from_cli_args( + parser.parse_args(["--index_cache_dtype", "fp8"]) + ) + + assert args.index_cache_dtype == "fp8" From 6bc1923f7fb20f84ac2a2b34f979bbe868a575f8 Mon Sep 17 00:00:00 2001 From: ganyi Date: Fri, 17 Jul 2026 02:07:53 +0000 Subject: [PATCH 10/11] convert auto for bf16 type Signed-off-by: ganyi --- atom/model_ops/attention_mha.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 10ad232b6e..80f349cba3 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -1096,6 +1096,9 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): # per-token dequant scales into k_scale / v_scale (outputs). fused_k_scale = k_scale if is_fp8 else None fused_v_scale = v_scale if is_fp8 else None + fused_index_cache_dtype = ( + self.index_cache_dtype if self.index_cache_dtype == "fp8" else "auto" + ) if self.skip_index_topk: from atom.model_ops.triton_fused_qkv_norm_rope_cache import ( @@ -1168,7 +1171,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, + index_cache_dtype=fused_index_cache_dtype, k_scale=fused_k_scale, v_scale=fused_v_scale, asm_layout=True, From 4bccc90b953703fd31a210681bb9f82bdf76c87f Mon Sep 17 00:00:00 2001 From: ganyi Date: Fri, 17 Jul 2026 02:18:22 +0000 Subject: [PATCH 11/11] black Signed-off-by: ganyi --- atom/model_ops/attentions/aiter_attention.py | 3 +-- tests/test_arg_utils_spec.py | 4 +--- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index 83b8e85de6..03e361b7ea 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -51,8 +51,7 @@ def _resolve_index_cache_dtype(config) -> torch.dtype: if index_cache_dtype in ("bf16", "fp8"): return dtypes.d_dtypes[index_cache_dtype] raise ValueError( - "index_cache_dtype must be 'bf16' or 'fp8', " - f"got {index_cache_dtype!r}" + "index_cache_dtype must be 'bf16' or 'fp8', " f"got {index_cache_dtype!r}" ) diff --git a/tests/test_arg_utils_spec.py b/tests/test_arg_utils_spec.py index b1f2cf0554..e1ac88e3ef 100644 --- a/tests/test_arg_utils_spec.py +++ b/tests/test_arg_utils_spec.py @@ -83,9 +83,7 @@ def test_index_cache_dtype_defaults_to_kv_cache_dtype(self): parser = argparse.ArgumentParser() EngineArgs.add_cli_args(parser) - args = EngineArgs.from_cli_args( - parser.parse_args(["--kv_cache_dtype", "fp8"]) - ) + args = EngineArgs.from_cli_args(parser.parse_args(["--kv_cache_dtype", "fp8"])) assert args.kv_cache_dtype == "fp8" assert args.index_cache_dtype == "fp8"