[None][feat] Support the masked DSA indexer k-cache pool in the Python cache transceiver

tensorrt_llm/_torch/disaggregation/native/bounce/impl.py

@@ -16,9 +16,11 @@
 the contract in core.py. Holds the buffers, the gather and scatter kernels, and the scatter worker,
 and runs the side effects that drive each region's state machine. Never imports transfer.py."""
 
+from __future__ import annotations
+
 import queue
 import threading
-from typing import Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional
 
 import numpy as np
 
@@ -41,6 +43,9 @@
 from .core import BounceTransport, Disposition, Settlement, TransferContext
 from .gather_scatter import Plan, gather_contiguous, scatter_contiguous
 
+if TYPE_CHECKING:
+    from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable
+
 RidSlice = tuple  # the request id and slice id a region serves
 _MIB = 1024 * 1024
 _SCATTER_POLL_S = 0.5  # how often the scatter worker wakes to re-check the stop flag and reclaim
@@ -57,8 +62,8 @@ class VmmBounceTransport(BounceTransport):
 
     @classmethod
     def from_config(
-        cls, agent, cfg, *, device_id: int, block_bytes_per_group: List[int]
-    ) -> Optional["VmmBounceTransport"]:
+        cls, agent, cfg, *, device_id: int, block_bytes_per_group: list[int | None]
+    ) -> VmmBounceTransport | None:
         """Build a transport sized from the config and clamped to free memory, or None if not even one
         chunk fits."""
         chunk = cfg.chunk_mb * _MIB
@@ -94,7 +99,7 @@ def __init__(
         device_id: int,
         capacity_bytes: int,
         phys_chunk_size: int,
-        block_bytes_per_group: List[int],
+        block_bytes_per_group: list[int | None],
         min_bytes: int = DEFAULT_MIN_BYTES,
         min_blocks: int = 96,
         quarantine_grace_s: float = _QUARANTINE_GRACE_S,
@@ -587,21 +592,28 @@ def decode_result_tail(message):
     return None, None, None
 
 
-def block_bytes_per_group(page_table) -> list:
-    """Byte size of one cache block for each layer group, aligned with the layer-group indices a
-    recv request uses. Non-attention groups (mamba/KDA recurrent state) hold ``None``: they carry
-    no paged blocks (their KVSlice entry is always empty — see ``_create_kv_slice``) and their
-    payload is sized separately via ``MambaPolicy.payload_bytes``. Keeping them as placeholders
-    instead of truncating means a trailing (or hypothetically interleaved) mamba group can never
-    shift an attention group off the end of this list and poison the bounce gate."""
+def block_bytes_per_group(page_table: KVCachePageTable) -> list[int | None]:
+    """Return transferred bytes per cache block for each layer group.
+
+    All distinct physical pools exposed by an attention group contribute to its
+    transfer size. Multiple logical views of the same physical pool contribute
+    only once. Non-attention groups retain a ``None`` placeholder so the result
+    remains aligned with receive-request layer-group indices.
+    """
     from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup
     from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool
 
     assert page_table is not None
-    out: list = []
+    out: list[int | None] = []
     for lg_idx, lg in enumerate(page_table.layer_groups):
         if not isinstance(lg, AttentionLayerGroup):
             out.append(None)
             continue
-        out.append(int(get_physical_pool(page_table, lg_idx, 0).slot_bytes))
+        pool_indices = {pool_view.pool_idx for pool_view in lg.pool_views}
+        out.append(
+            sum(
+                int(get_physical_pool(page_table, lg_idx, pool_idx).slot_bytes)
+                for pool_idx in pool_indices
+            )
+        )
     return out

tensorrt_llm/_torch/disaggregation/native/transfer.py

@@ -61,13 +61,13 @@
 )
 from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message
 from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy
-from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar
+from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar
 from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager
 from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo
 from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip
 from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent
 from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1
-from tensorrt_llm._torch.disaggregation.resource.page import MapperKind
+from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable, MapperKind
 from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
@@ -1587,7 +1587,11 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo:
         )
 
     @staticmethod
-    def _fanin_bounce_safe(overlap, peer_ri) -> bool:
+    def _fanin_bounce_safe(
+        overlap: PeerOverlap,
+        peer_ri: RankInfo,
+        receiver_page_table: Optional[KVCachePageTable],
+    ) -> bool:
         """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap.
         The split assumes every writer contributes the same size, which holds when:
           * duplicate_head_factor == 1 -- else some ranks don't send KV (should_send_kv) yet still
@@ -1607,15 +1611,20 @@ def _fanin_bounce_safe(overlap, peer_ri) -> bool:
                 return False
         # Replicated pools (e.g. MiniMax M3 index-key) are sent by one elected
         # fan-in owner only, so with multiple writers their contributions
-        # differ in size and the equal split is invalid.
-        if len(overlap.ranks) > 1 and peer_ri.page_table is not None:
-            for layer_group in peer_ri.page_table.layer_groups:
-                for pool_view in getattr(layer_group, "pool_views", ()):
-                    if pool_view.mapper_kind == MapperKind.REPLICATED:
-                        return False
+        # differ in size and the equal split is invalid. Inspect both endpoints:
+        # a masked PP stage may advertise no replicated view even though another
+        # stage owns one that is visible in the receiver's page table.
+        if len(overlap.ranks) > 1:
+            for page_table in (peer_ri.page_table, receiver_page_table):
+                if page_table is None:
+                    continue
+                for layer_group in page_table.layer_groups:
+                    for pool_view in getattr(layer_group, "pool_views", ()):
+                        if pool_view.mapper_kind == MapperKind.REPLICATED:
+                            return False
         return True
 
-    def dispatch_task(self, task: KVRecvTask):
+    def dispatch_task(self, task: KVRecvTask) -> None:
         params = task._params
         logger.debug(
             f"Receiver.dispatch_task: unique_rid={task._unique_rid}, ctx_dp_rank={params.ctx_dp_rank}"
@@ -1657,7 +1666,12 @@ def dispatch_task(self, task: KVRecvTask):
         # None), where the real writer count exceeds expected_transfers and would overflow the slot.
         topo_overlap = peer_overlap if sender_dp_rank is not None else dp0_overlap
         allow_bounce = task.expected_transfers == 1 or (
-            sender_dp_rank is not None and self._fanin_bounce_safe(topo_overlap, peer_infos)
+            sender_dp_rank is not None
+            and self._fanin_bounce_safe(
+                topo_overlap,
+                peer_infos,
+                self._registrar.self_extractor.page_table,
+            )
         )
         # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks (the sender
         # appends its MambaPolicy fragments in _build_kv_write_meta), so the bounce region must be

tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py

@@ -326,60 +326,65 @@ def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable:
         pool_views = [kv_view]
 
         # Indexer K cache support. The DSA indexer K cache is identical on
-        # every TP rank (single index head), so its view is REPLICATED with
-        # one synthesized buffer entry per local layer: the slot packs the
-        # layers equal-sized in local-layer order.
-        if getattr(kv_cache_manager, "enable_indexer_k_cache", False):
-            local_indexer_mask = getattr(kv_cache_manager, "indexer_k_cache_local_layer_mask", None)
-            if local_indexer_mask is not None and not all(
-                local_indexer_mask[lid] for lid in local_layer_ids
-            ):
-                raise NotImplementedError(
-                    "The Python KV transceiver runtime does not support a "
-                    "per-layer masked indexer k-cache pool yet: "
-                    f"{sum(local_indexer_mask[lid] for lid in local_layer_ids)}"
-                    f" of {len(local_layer_ids)} layers in this layer group "
-                    "own an indexer k-cache. Use the C++ cache transceiver "
-                    "for models with cross-layer indexer sharing (e.g. "
-                    "GLM 5.2)."
+        # every TP rank (single index head), so its view is REPLICATED. With a
+        # per-layer indexer mask (cross-layer indexer sharing, e.g. GLM 5.2)
+        # only the "full" indexer-owning layers get a pool row, so the view
+        # covers that subset: one buffer entry per owning layer, each mapped to
+        # its packed row in the (possibly masked) pool. When the mask is absent
+        # every layer owns a row (dense/legacy layout) and this reduces to the
+        # equal-sized packing in local-layer order.
+        if kv_cache_manager.enable_indexer_k_cache:
+            local_indexer_mask = kv_cache_manager.indexer_k_cache_local_layer_mask
+            owning_layer_ids = [
+                lid
+                for lid in local_layer_ids
+                if local_indexer_mask is None or local_indexer_mask[lid]
+            ]
+            # A layer group whose layers are all masked out owns no indexer pool
+            # row on this rank (the pool getter would raise); skip it so the peer
+            # simply transfers nothing for this rank's indexer.
+            if owning_layer_ids:
+                indexer_pool = kv_cache_manager.impl.get_indexer_k_cache_pool()
+                if indexer_pool.shape[1] != len(owning_layer_ids):
+                    raise RuntimeError(
+                        "The DSA indexer K-cache pool row count does not match "
+                        "the number of indexer-owning layers in its layer group: "
+                        f"{indexer_pool.shape[1]} rows for {len(owning_layer_ids)} layers"
+                    )
+                # indexer_pool shape: (numBlocks, numIndexerLayers, kvFactor,
+                # blockSize), dtype=UINT8. numIndexerLayers is the number of
+                # owning layers on this rank (== the attention layer count when
+                # unmasked). slot_bytes packs every owning-layer row.
+                per_block_elems = 1
+                for d in indexer_pool.shape[1:]:  # skip numBlocks dim
+                    per_block_elems *= d
+                indexer_slot_bytes = per_block_elems * indexer_pool.element_size()
+                indexer_bytes_per_layer = indexer_slot_bytes // indexer_pool.shape[1]
+                indexer_physical = PhysicalPool(
+                    base_address=int(indexer_pool.data_ptr()),
+                    slot_bytes=indexer_slot_bytes,
+                    num_slots=num_blocks,
                 )
-            indexer_pool = kv_cache_manager.impl.get_indexer_k_cache_pool()
-            # indexer_pool shape: (numBlocks, numLayers, kvFactor, blockSize), dtype=UINT8
-            # slot_bytes = numLayers * kvFactor * blockSize * element_size
-            if indexer_pool.shape[1] != len(local_layer_ids):
-                raise NotImplementedError(
-                    "Disaggregated KV transfer does not support a per-layer "
-                    "masked indexer k-cache pool yet: the indexer "
-                    f"pool holds {indexer_pool.shape[1]} layer rows but the "
-                    f"layer group has {len(local_layer_ids)} layers. Disable "
-                    "disaggregated serving for models with cross-layer "
-                    "indexer sharing."
+                indexer_view = PoolView(
+                    pool_idx=len(physical_pools),
+                    buffer_entries=np.array(
+                        [
+                            (
+                                lid,
+                                kv_cache_manager.impl.get_indexer_k_cache_pool_layer_idx(lid)
+                                * indexer_bytes_per_layer,
+                                indexer_bytes_per_layer,
+                            )
+                            for lid in owning_layer_ids
+                        ],
+                        dtype=BUFFER_ENTRY_DTYPE,
+                    ),
+                    pool_role=frozenset({"indexer_k"}),
+                    mapper_kind=MapperKind.REPLICATED,
+                    bytes_per_layer=indexer_bytes_per_layer,
                 )
-            per_block_elems = 1
-            for d in indexer_pool.shape[1:]:  # skip numBlocks dim
-                per_block_elems *= d
-            indexer_slot_bytes = per_block_elems * indexer_pool.element_size()
-            indexer_physical = PhysicalPool(
-                base_address=int(indexer_pool.data_ptr()),
-                slot_bytes=indexer_slot_bytes,
-                num_slots=num_blocks,
-            )
-            indexer_bytes_per_layer = indexer_slot_bytes // len(local_layer_ids)
-            indexer_view = PoolView(
-                pool_idx=1,
-                buffer_entries=np.array(
-                    [
-                        (lid, i * indexer_bytes_per_layer, indexer_bytes_per_layer)
-                        for i, lid in enumerate(local_layer_ids)
-                    ],
-                    dtype=BUFFER_ENTRY_DTYPE,
-                ),
-                pool_role=frozenset({"indexer_k"}),
-                mapper_kind=MapperKind.REPLICATED,
-                bytes_per_layer=indexer_bytes_per_layer,
-            )
-            physical_pools.append(indexer_physical)
-            pool_views.append(indexer_view)
+                physical_pools.append(indexer_physical)
+                pool_views.append(indexer_view)
 
         pool_groups.append(PhysicalPoolGroup(pools=physical_pools))
         local_layers = [

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -1901,7 +1901,6 @@ def forward(
         return hidden_states
 
 
-@register_auto_model("GlmMoeDsaForCausalLM")
 @register_auto_model("DeepseekV32ForCausalLM")
 @register_auto_model("DeepseekV3ForCausalLM")
 class DeepseekV3ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV3Model,
@@ -1919,25 +1918,14 @@ def get_preferred_transceiver_runtime(
             cls,
             pretrained_config: Any = None
     ) -> Optional[Literal["CPP", "PYTHON"]]:
-        """Preferred KV-cache transceiver runtime, differentiated per checkpoint.
-
-        ``DeepseekV3ForCausalLM`` / ``DeepseekV32ForCausalLM`` use MLA attention, which transfers
-        a large latent KV that the Python (v2) transceiver handles better in disaggregated
-        serving, so they prefer the Python transceiver. GLM 5.2 (``GlmMoeDsaForCausalLM`` /
-        ``glm_moe_dsa``) uses a per-layer masked DSA indexer k-cache pool (cross-layer indexer
-        sharing) that the Python transceiver does not support, so GLM checkpoints must use the
-        C++ transceiver, which handles both the masked pool and dense indexer layouts. Applied
-        only when ``cache_transceiver_config.transceiver_runtime`` is 'auto'; an explicit runtime
+        """Preferred KV-cache transceiver runtime.
+
+        ``DeepseekV3ForCausalLM`` and ``DeepseekV32ForCausalLM`` use MLA
+        attention, which transfers a large latent KV that the Python (v2)
+        transceiver handles better in disaggregated serving. Applied only when
+        ``cache_transceiver_config.transceiver_runtime`` is 'auto'; an explicit runtime
         is always respected.
         """
-        if pretrained_config is not None:
-            architectures = getattr(pretrained_config, 'architectures',
-                                    None) or []
-            # model_type is checked as a fallback: it is 'glm_moe_dsa' on GLM
-            # checkpoints until __init__ rewrites it to 'deepseek_v32'.
-            if ("GlmMoeDsaForCausalLM" in architectures or getattr(
-                    pretrained_config, 'model_type', None) == 'glm_moe_dsa'):
-                return "CPP"
         return "PYTHON"
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig]):
@@ -2103,3 +2091,16 @@ def setup_aliases(self) -> None:
                         layer.mlp.experts.fuse_shared_expert(
                             layer.mlp.shared_experts)
                         layer.mlp.shared_experts = None
+
+
+@register_auto_model("GlmMoeDsaForCausalLM")
+class GlmMoeDsaForCausalLM(DeepseekV3ForCausalLM):
+    """GLM 5.2 model flavor with an independent transceiver preference."""
+
+    @classmethod
+    def get_preferred_transceiver_runtime(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Optional[Literal["CPP", "PYTHON"]]:
+        """Prefer Python for GLM 5.2's masked DSA indexer K-cache transfer."""
+        return "PYTHON"

tests/integration/test_lists/qa/llm_function_core.txt

@@ -830,6 +830,8 @@ accuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_nvfp4[mtp_nextn=0
 accuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_nvfp4[mtp_nextn=3] TIMEOUT (120)
 accuracy/test_llm_api_pytorch_multimodal.py::TestVILA1_5_3B::test_auto_dtype
 accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray
+unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_v1_masked_dsa_indexer_across_asymmetric_pp
+unittest/disaggregated/test_openai_disagg_server.py
 disaggregated/test_ad_disagg.py::test_async_eagle3_full_model_handoff
 disaggregated/test_ad_disagg.py::test_async_generation_matches_aggregate
 disaggregated/test_ad_disagg.py::test_async_generation_no_overlap_matches_aggregate

tests/integration/test_lists/test-db/l0_h100.yml

@@ -97,6 +97,7 @@ l0_h100:
   - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_boundary_lengths -k "v2"
   # DSA indexer-K side cache (V1, REPLICATED).
   - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_v1_dsa_indexer
+  - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_v1_masked_dsa_indexer_across_asymmetric_pp
   - unittest/disaggregated/test_cache_transceiver_harness_report.py
   - unittest/disaggregated/test_cache_transceiver_harness.py
   - unittest/disaggregated/test_cache_transceiver_precheck_e2e.py