review-fixes(dist_reuse): consolidated reviewer feedback

Squashes three prior review-fix commits into one. Addresses reviewer comments on the dist_reuse feature commit (22bc183):

F1. is_nsa source: read directly from model_config.is_nsa instead of reverse-deriving from enable_nsa_prefill_context_parallel (the latter is a CP toggle, orthogonal to NSA architecture).

F2. Control-plane / rank-plane separation: SharingDomainKey.from_model_config now takes an optional rank_info=RankInfo argument and reads pp_rank / tp_node_idx from it; the control plane (KVManager) only constructs self-SD via default() and enumerates peers via enumerate_peers() — no fake rank fabrication.

F3. RankTopology factory dropped; reuse the existing RankInfo end-to-end. Integration adapters (vLLM v1 / TRT-LLM / SGLang) plumb the real rank_info through.

F4. shell / TransferManagerOnRemote decoupling: revert start_dist_reuse_serving.sh changes; TransferManagerOnRemote stays per-node and is tagged via set_target_sd_key on each handle.

F5. delete unused flexkv.integration.multinode_policy module and its is_multinode_tp / is_multinode_cp / is_multinode_pp helpers; CP never participates in sd_key (attention all-gather makes per-cp_rank pools bit-wise identical), and TP-cross-node is encoded in SharingDomainKey.tp_node_count directly. Verified no external references in the SGLang FlexKVConnector codebase.

Tests: full dist_reuse suite (363/363) passes on both GPU executors.

Author: phaedonsun

flexkv/common/config.py

@@ -104,16 +104,16 @@ 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.tp_node_count > 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"tp_node_count={self.tp_node_count} "
                 f"(tp_size={self.tp_size}, gpus_per_node={self.gpus_per_node})"
             )
-        if self.tp_size % self.nnodes_per_tp_group != 0:
+        if self.tp_size % self.tp_node_count != 0:
             raise ValueError(
                 f"[ModelConfig] tp_size={self.tp_size} not divisible by "
-                f"nnodes_per_tp_group={self.nnodes_per_tp_group}"
+                f"tp_node_count={self.tp_node_count}"
             )
         if self.instance_num < 1:
             raise ValueError(
@@ -160,13 +160,23 @@ def nnodes_per_pp_rank(self) -> int:
 
     @property
     def nnodes_per_tp_group(self) -> int:
-        """Number of nodes spanned by one TP group."""
-        return self.nnodes_per_pp_rank
+        """Number of nodes spanned by one TP group.
+
+        .. deprecated::
+            Kept as a stable alias of :pyattr:`tp_node_count` for
+            backwards compatibility with adapter code that pre-dates
+            the SD-key naming convention.  New code should read
+            ``tp_node_count`` directly — that property carries the
+            authoritative semantic ("the TP-axis node-count entering
+            ``SharingDomainKey``") and is the value tracked in the
+            redis schema (``docs/dist_reuse/redis_schema.md``).
+        """
+        return self.tp_node_count
 
     @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
+        return self.tp_size // self.tp_node_count
 
     @property
     def attn_dp_size(self) -> int:
@@ -183,7 +193,7 @@ def attn_tp_size(self) -> int:
     @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
+        return self.attn_tp_size // self.tp_node_count
 
     @property
     def attn_cp_size_per_node(self) -> int:
@@ -242,6 +252,43 @@ def is_multinode_tp(self) -> bool:
         """
         return self.tp_node_count > 1
 
+    @property
+    def is_multinode_pp(self) -> bool:
+        """PP is the dimension that makes *this instance* cross node boundaries.
+
+        Returns True iff:
+
+        * ``pp_size > 1``      — PP is actually deployed,
+        * ``nnodes > 1``       — the instance occupies more than one node,
+        * ``tp_node_count == 1`` — TP **does not** cross nodes
+          (otherwise TP-multinode is the dominant axis and would
+          already drive the SD-Remote decision; classifying the same
+          deployment as "multinode-PP" too would double-count).
+
+        This is the missing third axis next to :pyattr:`is_multinode_tp`
+        and :pyattr:`is_multinode_cp`.  It exists so the connector's
+        runtime launch logic can stop folding "PP-only crosses nodes"
+        into the off-master fall-through branch.
+
+        Worked examples:
+
+        * ``pp=4, nnodes=2, tp=8, gpus_per_node=8`` → True
+          (PP=4 stages × tp=8 = 32 GPUs across 2 nodes; each node
+          owns 2 PP stages; TP stays inside one node).
+        * ``pp=1, nnodes=2, tp=16`` → False
+          (PP single-stage; TP is the multinode axis).
+        * ``pp=2, nnodes=2, tp=16`` → False
+          (TP already crosses; PP is *not* the dominant axis here \u2014
+          we leave the multinode-TP branch to handle this).
+        * ``pp=2, nnodes=1`` → False
+          (single node; PP fits in-host).
+        """
+        return (
+            self.pp_size > 1
+            and self.nnodes > 1
+            and self.tp_node_count == 1
+        )
+
     @property
     def is_multinode_cp(self) -> bool:
         """CP > 1 *and* the CP group spans more than one physical node.
@@ -300,9 +347,17 @@ def num_kv_heads_per_node(self) -> int:
     # ------------------------------------------------------------------
     @property
     def tp_node_count(self) -> int:
-        """Number of physical nodes one TP group spans (=
-        ``nnodes_per_tp_group``).  ``1`` when TP fits on a single node."""
-        return self.nnodes_per_tp_group
+        """Number of physical nodes one TP group spans.
+
+        Authoritative source for the TP-axis node-count used in
+        :class:`SharingDomainKey` and ``docs/dist_reuse/redis_schema.md``.
+        ``1`` when TP fits on a single node.  Deprecated alias:
+        :pyattr:`nnodes_per_tp_group`.
+        """
+        # PP and TP groups share the same per-rank node assignment in the
+        # current topology (one TP group sits on the same set of nodes as
+        # one PP stage), so ``nnodes_per_tp_group == nnodes_per_pp_rank``.
+        return self.nnodes_per_pp_rank
 
     # NOTE: ``tp_node_idx`` is a per-rank concept and was moved to
     # ``RankInfo`` in PR #165 (separate-per-rank-state-into-RankInfo).

flexkv/common/dist_reuse/sharing_domain.py

@@ -360,8 +360,39 @@ def from_model_config(
                         getattr(model_config, "tp_node_idx", 0))
             )
         else:
+            # Legacy path: PR #165 moved ``pp_rank`` / ``tp_node_idx``
+            # off ``ModelConfig`` onto :class:`RankInfo`.  These
+            # ``getattr(..., 0)`` reads therefore now return ``0``
+            # for any post-#165 ``ModelConfig`` instance — i.e. the
+            # caller silently gets the master-position SD even if it
+            # is actually on (pp_rank>0, tp_node_idx>0).  This is
+            # only safe for unit-test fakes that explicitly set
+            # ``pp_rank`` / ``tp_node_idx`` on the stub ModelConfig
+            # (see ``tests/test_sharing_domain_key.py``); production
+            # callers should pass ``rank_info=`` explicitly.  We log
+            # a one-time warning when the heuristic is exercised on
+            # a multi-rank topology so the error surfaces during
+            # bring-up instead of silently corrupting Redis keys.
             _pp_rank = int(getattr(model_config, "pp_rank", 0))
             _tp_node_idx = int(getattr(model_config, "tp_node_idx", 0))
+            _has_pp_rank = hasattr(model_config, "pp_rank")
+            _has_tp_node_idx = hasattr(model_config, "tp_node_idx")
+            if (
+                pp_size > 1 or int(getattr(model_config, "tp_node_count", 1)) > 1
+            ) and not (_has_pp_rank and _has_tp_node_idx):
+                # Local import to avoid pulling logger into module
+                # import-time graph (sharing_domain.py is imported
+                # very early by config.py via ``derive_model_id``).
+                from flexkv.common.debug import flexkv_logger
+                flexkv_logger.warning(
+                    "SharingDomainKey.from_model_config: called without "
+                    "rank_info on a multi-rank topology (pp_size=%d, "
+                    "tp_node_count=%s); per-rank fields default to 0, "
+                    "which only matches the master SD.  Pass rank_info="
+                    "<RankInfo> so the per-rank position is honoured.",
+                    pp_size,
+                    getattr(model_config, "tp_node_count", 1),
+                )
 
         _pp_node_idx = _pp_rank // pp_per_node
 

flexkv/integration/config.py

@@ -218,8 +218,11 @@ def post_init_from_sglang_config(
             sglang_config: sglang.srt.configs.model_config.ModelConfig-like object
             server_args: sglang ServerArgs — source of tp_size, dp_size,
                 nnodes, node_rank, enable_dp_attention, attn_cp_size,
-                is_nsa (read from server_args.enable_nsa_prefill_context_parallel),
-                kv_cache_dtype, dist_init_addr
+                kv_cache_dtype, dist_init_addr.  ``is_nsa`` is **not**
+                read from server_args: see the body below — it is
+                derived from ``sglang_config.index_head_dim`` instead,
+                because NSA is a model-layout property orthogonal to
+                whether CP-prefill is enabled.
             page_size: KV block size (tokens per block) used by sglang
             tp_rank: physical tensor parallel rank (runtime, from process group)
             pp_rank: pipeline parallel rank (runtime, from process group)
@@ -234,11 +237,24 @@ def post_init_from_sglang_config(
         node_rank = server_args.node_rank
         enable_dp_attention = server_args.enable_dp_attention
         attn_cp_size = server_args.attn_cp_size
-        # ``is_nsa`` (NSA model layout flag): True when the model has an
-        # extra indexer K cache buffer.  Sourced from sglang's
-        # ``enable_nsa_prefill_context_parallel`` server arg, but in dist_reuse
-        # context the flag represents the *layout*, not whether CP is on.
-        is_nsa = getattr(server_args, 'enable_nsa_prefill_context_parallel', False)
+        # ``is_nsa`` (NSA model layout flag): True when the model itself has
+        # an extra indexer K cache buffer.  This is a *layout* property of
+        # the model architecture, **independent** of whether CP is enabled
+        # at runtime — an NSA model with cp_size=1 still has the indexer K
+        # cache and must therefore be isolated from non-NSA models in the
+        # cross-instance reuse namespace (it lives in
+        # ``SharingDomainKey.serialize`` as the ``nsa<0|1>`` segment).
+        #
+        # Detection rule: an NSA/DSA model exposes a positive
+        # ``index_head_dim`` attribute on its sglang ModelConfig (the same
+        # signal already consulted ~25 lines below to size the indexer
+        # head buffer).  Falling back to
+        # ``server_args.enable_nsa_prefill_context_parallel`` was incorrect
+        # because it conflates the *runtime CP toggle* with the *static
+        # model layout* — a deployment can run an NSA model with CP=1
+        # (no prefill-CP) and still need NSA-isolated namespaces.
+        index_head_dim = getattr(sglang_config, "index_head_dim", None)
+        is_nsa = bool(index_head_dim) and int(index_head_dim) > 0
         kv_cache_dtype = getattr(server_args, 'kv_cache_dtype', None)
         dp_rank = 0 if dp_rank is None else int(dp_rank)