feat: support sglang dp attention and drop sglang-specific attn_* fields

Author: zittozhang

flexkv/cache/redis_meta.py

@@ -283,8 +283,6 @@ def register_node(self) -> Optional[int]:
                 "uuid": self.uuid,
                 "status": "active",
                 "timestamp": str(int(time.time())),
-                "pp_rank": str(getattr(self, 'pp_rank', 0)),
-                "pp_size": str(getattr(self, 'pp_size', 1)),
             })
 
             # Set TTL so the key auto-expires if the process crashes

flexkv/common/config.py

@@ -40,14 +40,8 @@ class ModelConfig:
     # ------------------------------------------------------------------
     # Attention-level parallel configs
     # ------------------------------------------------------------------
-    # enable_dp_attention: whether DP-attention is enabled (sglang
-    # ``--enable-dp-attention`` or TRT-LLM ``enable_attention_dp``).
-    # When True, the physical TP group is split into
-    # attn_tp × attn_cp × attn_dp.
-    enable_dp_attention: bool = False
-
-    # attn_cp_size: context-parallel size (global).
-    attn_cp_size: int = 1
+    # cp_size: context-parallel size (global), default 1.
+    cp_size: int = 1
 
     # ------------------------------------------------------------------
     # Topology configs (global)
@@ -94,17 +88,12 @@ def freeze(self) -> None:
                 f"[ModelConfig] cannot derive gpus_per_node: "
                 f"total_gpus={self.total_gpus} not divisible by nnodes={self.nnodes}"
             )
-        if self.nnodes_per_tp_group > 2:
+        if self.nnodes_per_pp_rank > 2:
             raise ValueError(
                 f"[ModelConfig] only support 2-nodes TP for now, but got "
-                f"nnodes_per_tp_group={self.nnodes_per_tp_group} "
+                f"nnodes_per_pp_rank={self.nnodes_per_pp_rank} "
                 f"(tp_size={self.tp_size}, gpus_per_node={self.gpus_per_node})"
             )
-        if self.tp_size % self.nnodes_per_tp_group != 0:
-            raise ValueError(
-                f"[ModelConfig] tp_size={self.tp_size} not divisible by "
-                f"nnodes_per_tp_group={self.nnodes_per_tp_group}"
-            )
         if self.instance_num < 1:
             raise ValueError(
                 f"[ModelConfig] instance_num must be >= 1, got {self.instance_num}"
@@ -119,8 +108,8 @@ def __setattr__(self, name: str, value) -> None:
             raise AttributeError(
                 f"ModelConfig is frozen — cannot set '{name}'. "
                 f"All primitive fields must be set during post_init_from_*(), "
-                f"after which freeze() is called.  Derived fields (attn_tp_size, "
-                f"tp_size_per_node) are @property "
+                f"after which freeze() is called.  Derived fields (effective_tp_size, "
+                f"tp_size_per_node, cp_size_per_node, nnodes_per_pp_rank) are @property "
                 f"and cannot be set at all."
             )
         object.__setattr__(self, name, value)
@@ -130,8 +119,10 @@ def __setattr__(self, name: str, value) -> None:
     # ------------------------------------------------------------------
     @property
     def total_gpus(self) -> int:
-        """Total GPUs across all nodes for one FlexKV instance."""
-        return self.dp_size * self.tp_size * self.pp_size
+        """Total GPU worker registration slots across all nodes for one FlexKV instance.
+
+        Unified formula: dp_size × tp_size × cp_size × pp_size."""
+        return self.dp_size * self.tp_size * self.cp_size * self.pp_size
 
     @property
     def total_clients(self) -> int:
@@ -140,62 +131,43 @@ def total_clients(self) -> int:
 
     @property
     def gpus_per_node(self) -> int:
-        """Total GPUs on this node (across all DP, PP stages and TP groups)."""
+        """GPU worker registration slots on this node (across all DP shards, PP stages and TP groups)."""
         return self.total_gpus // self.nnodes
 
     @property
     def nnodes_per_pp_rank(self) -> int:
         """Number of nodes spanned by one PP stage."""
         return max(self.nnodes // self.pp_size, 1)
 
-    @property
-    def nnodes_per_tp_group(self) -> int:
-        """Number of nodes spanned by one TP group."""
-        return self.nnodes_per_pp_rank
-
     @property
     def tp_size_per_node(self) -> int:
         """Number of TP ranks on this node within one TP group."""
-        return self.tp_size // self.nnodes_per_tp_group
-
-    @property
-    def attn_dp_size(self) -> int:
-        """Attention-level DP size (= dp_size when enable_dp_attention else 1)."""
-        return max(1, self.dp_size) if self.enable_dp_attention else 1
+        return max(1, self.tp_size // self.nnodes_per_pp_rank)
 
     @property
-    def attn_tp_size(self) -> int:
-        """Attention-level TP size derived from tp / attn_dp / attn_cp."""
-        attn_dp = self.attn_dp_size
-        cp = max(1, self.attn_cp_size)
-        return max(1, max(1, self.tp_size) // (attn_dp * cp))
+    def cp_size_per_node(self) -> int:
+        """CP size on this node for a single PP stage.
 
-    @property
-    def attn_tp_size_per_node(self) -> int:
-        """Attention-level TP size per node."""
-        return self.attn_tp_size // self.nnodes_per_tp_group
-
-    @property
-    def attn_cp_size_per_node(self) -> int:
-        """Attention-level CP size on this node for a single pp stage. """
-        return max(1, self.attn_cp_size // self.nnodes_per_pp_rank)
+        Used for multi-node scenarios where the CP group spans multiple nodes.
+        """
+        return max(1, self.cp_size // self.nnodes_per_pp_rank)
 
     @property
     def effective_tp_size(self) -> int:
-        """Effective tp-group size used for *data-plane* CPU slicing."""
-        return max(1, self.attn_tp_size) * max(1, self.attn_cp_size)
+        """Number of CPU block slices = tp_size × cp_size."""
+        return max(1, self.tp_size) * max(1, self.cp_size)
 
     @property
     def effective_tp_size_per_node(self) -> int:
         """Per-node counterpart of :pyattr:`effective_tp_size`."""
-        return self.attn_tp_size_per_node * self.attn_cp_size_per_node
+        return self.tp_size_per_node * self.cp_size_per_node
 
     @property
     def num_kv_heads_per_node(self) -> int:
         """Number of KV heads visible to a single node."""
         if self.use_mla:
             return self.num_kv_heads
-        return self.num_kv_heads * self.tp_size_per_node // max(1, self.attn_tp_size)
+        return self.num_kv_heads * self.tp_size_per_node // max(1, self.tp_size)
 
     @property
     def kv_dim(self) -> int:
@@ -218,7 +190,8 @@ def __str__(self) -> str:
             f", head_size={self.head_size}, use_mla={self.use_mla}"
             f", dtype={self.dtype}"
             f", tp_size={self.tp_size}, pp_size={self.pp_size}, dp_size={self.dp_size}"
-            f", attn_cp_size={self.attn_cp_size}"
+            f", cp_size={self.cp_size}"
+            f", total_gpus={self.total_gpus}"
             f", nnodes={self.nnodes}, master_host={self.master_host!r}"
             f", instance_num={self.instance_num}"
         )
@@ -230,11 +203,12 @@ class RankInfo:
     tp_rank: int = 0
     pp_rank: int = 0
     dp_rank: int = 0
-    attn_cp_rank: int = 0
+    cp_rank: int = 0
     node_rank: int = 0
     instance_id: int = 0
     pp_start_layer: int = 0
     pp_end_layer: int = -1
+    local_rank: int = -1
     @property
     def tp_rank_per_node(self) -> int:
         """TP rank index within the local node (within one TP group)."""
@@ -252,18 +226,10 @@ def dp_client_id(self) -> int:
         """
         return self.instance_id * self.model_config.dp_size + self.dp_rank
 
-    @property
-    def attn_tp_rank(self) -> int:
-        """Attention-level TP rank derived from tp_rank / attn_tp_size."""
-        return self.tp_rank % max(1, self.model_config.attn_tp_size)
-
     @property
     def effective_tp_rank(self) -> int:
         """Effective tp-rank in the *data-plane* segmentation space."""
-        if self.model_config.use_mla:
-            return self.attn_tp_rank
-        attn_tp_size = max(1, self.model_config.attn_tp_size)
-        return self.attn_cp_rank * attn_tp_size + self.attn_tp_rank
+        return self.cp_rank * max(1, self.model_config.tp_size) + self.tp_rank
 
     @property
     def pp_size_per_node(self) -> int:
@@ -276,25 +242,6 @@ def pp_rank_per_node(self) -> int:
         """This rank's PP index *within* its node."""
         return self.pp_rank % self.pp_size_per_node
 
-    @property
-    def dp_size_per_node(self) -> int:
-        """Number of DP replicas co-located on a single node."""
-        model_config = self.model_config
-        return model_config.gpus_per_node // (self.pp_size_per_node * model_config.tp_size_per_node)
-
-    @property
-    def dp_rank_per_node(self) -> int:
-        """This rank's DP index *within* its node (non-DP-attention layout)."""
-        return self.dp_rank % self.dp_size_per_node
-
-    @property
-    def local_rank(self) -> int:
-        model_config = self.model_config
-        if model_config.enable_dp_attention:
-            return self.pp_rank_per_node * model_config.tp_size_per_node + self.tp_rank_per_node
-        return (self.dp_rank_per_node * self.pp_size_per_node + self.pp_rank_per_node) \
-               * model_config.tp_size_per_node + self.tp_rank_per_node
-
     @property
     def num_layers_per_pp_stage(self) -> int:
         """Number of layers managed by this PP stage."""
@@ -315,8 +262,9 @@ def __str__(self) -> str:
         """
         return (
             f"RankInfo(tp_rank={self.tp_rank}, pp_rank={self.pp_rank}"
-            f", dp_rank={self.dp_rank}, attn_cp_rank={self.attn_cp_rank}"
+            f", dp_rank={self.dp_rank}, cp_rank={self.cp_rank}"
             f", node_rank={self.node_rank}, instance_id={self.instance_id}"
+f", local_rank={self.local_rank}, effective_tp_rank={self.effective_tp_rank}"
         )
 
 
@@ -540,14 +488,49 @@ def convert_to_block_num(size_in_GB: float, block_size_in_bytes: int) -> int:
 def update_default_config_from_user_config(rank_info: RankInfo,
                                            cache_config: CacheConfig,
                                            user_config: UserConfig) -> None:
-    block_size_in_bytes = rank_info.token_size_in_bytes_per_pp_stage * cache_config.tokens_per_block
+    main_block_size_in_bytes = (
+        rank_info.token_size_in_bytes_per_pp_stage * cache_config.tokens_per_block
+    )
+    indexer_block_size_in_bytes = 0
+    if cache_config.indexer is not None:
+        indexer_cfg = cache_config.indexer
+        # Indexer is MLA-style (single shared head set, no TP head split),
+        # so per-token bytes = num_kv_heads * head_size * dtype.itemsize.
+        indexer_bytes_per_token_per_layer = (
+            indexer_cfg.num_kv_heads
+            * indexer_cfg.head_size
+            * indexer_cfg.dtype.itemsize
+        )
+        indexer_block_size_in_bytes = (
+            rank_info.num_layers_per_pp_stage
+            * indexer_bytes_per_token_per_layer
+            * 1
+        )
+    block_size_in_bytes = main_block_size_in_bytes + indexer_block_size_in_bytes
 
     assert user_config.cpu_cache_gb > 0
     assert user_config.ssd_cache_gb >= 0
 
     cache_config.num_cpu_blocks = convert_to_block_num(user_config.cpu_cache_gb, block_size_in_bytes)
     cache_config.num_ssd_blocks = convert_to_block_num(user_config.ssd_cache_gb, block_size_in_bytes)
 
+    if cache_config.indexer is not None:
+        flexkv_logger.info(
+            f"[CacheConfig] GB->blocks conversion (with indexer): "
+            f"main_block_size={main_block_size_in_bytes} B, "
+            f"indexer_block_size={indexer_block_size_in_bytes} B, "
+            f"total_block_size={block_size_in_bytes} B; "
+            f"cpu_cache_gb={user_config.cpu_cache_gb} -> num_cpu_blocks={cache_config.num_cpu_blocks}, "
+            f"ssd_cache_gb={user_config.ssd_cache_gb} -> num_ssd_blocks={cache_config.num_ssd_blocks}"
+        )
+    else:
+        flexkv_logger.info(
+            f"[CacheConfig] GB->blocks conversion: "
+            f"block_size={block_size_in_bytes} B; "
+            f"cpu_cache_gb={user_config.cpu_cache_gb} -> num_cpu_blocks={cache_config.num_cpu_blocks}, "
+            f"ssd_cache_gb={user_config.ssd_cache_gb} -> num_ssd_blocks={cache_config.num_ssd_blocks}"
+        )
+
     cache_config.ssd_cache_dir = user_config.ssd_cache_dir
     cache_config.enable_ssd = user_config.ssd_cache_gb > 0
     cache_config.enable_gds = user_config.enable_gds