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

.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
 }
 

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 | None = None
     enable_prefix_caching: bool = True
     enable_chunked_prefill: bool = True
     port: int = 8006
@@ -1098,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)
@@ -1267,6 +1271,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: Optional[str] = None
     block_size: int = 16
     max_model_len: Optional[int] = None
     max_num_batched_tokens: int = 16384
@@ -62,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."""
@@ -121,6 +126,14 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
             default="bf16",
             help="KV cache type. Default is 'bf16'.",
         )
+        parser.add_argument(
+            "--index-cache-dtype",
+            "--index_cache_dtype",
+            choices=["bf16", "fp8"],
+            type=str,
+            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."
         )

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 | None = None,
         **kwargs,
     ):
         super().__init__(
@@ -995,6 +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 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
@@ -1092,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 (
@@ -1164,6 +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=fused_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,19 @@ 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", 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(
+        "index_cache_dtype must be 'bf16' or 'fp8', " f"got {index_cache_dtype!r}"
+    )
+
+
 @triton.jit
 def _mtp_prepare_decode_metadata_kernel(
     context_lens_ptr,
@@ -480,11 +493,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 +554,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