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

tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py

@@ -326,60 +326,59 @@ 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.
+        # 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 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)."
+            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()
+                # 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

@@ -1911,25 +1911,17 @@ 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`` / ``DeepseekV32ForCausalLM`` / ``GlmMoeDsaForCausalLM``
+        (GLM 5.2) use MLA attention, which transfers a large latent KV that the Python
+        (v2) transceiver handles better in disaggregated serving. The Python transceiver
+        also supports GLM 5.2's per-layer masked DSA indexer k-cache pool (cross-layer
+        indexer sharing), so every checkpoint sharing this implementation prefers the
+        Python transceiver. 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]):

tests/unittest/disaggregated/test_extractor.py

@@ -208,8 +208,17 @@ def test_build_page_table():
     manager.shutdown()
 
 
-def _make_v1_dsa_manager(pp_size: int = 1, pp_rank: int = 0) -> KVCacheManager:
-    """V1 KVCacheManager with the DSA indexer K cache enabled (MLA-style)."""
+def _make_v1_dsa_manager(
+    pp_size: int = 1,
+    pp_rank: int = 0,
+    indexer_k_cache_layer_mask=None,
+) -> KVCacheManager:
+    """V1 KVCacheManager with the DSA indexer K cache enabled (MLA-style).
+
+    ``indexer_k_cache_layer_mask`` is a global per-model ``list[bool]`` marking
+    the "full" indexer-owning layers (cross-layer indexer sharing, e.g. GLM
+    5.2); ``None`` keeps the dense layout where every layer owns an indexer row.
+    """
     return KVCacheManager(
         kv_cache_config=KvCacheConfig(
             max_tokens=512,
@@ -228,6 +237,7 @@ def _make_v1_dsa_manager(pp_size: int = 1, pp_rank: int = 0) -> KVCacheManager:
         enable_indexer_k_cache=True,
         indexer_k_cache_quant_block_size=128,
         indexer_k_cache_index_head_dim=128,
+        indexer_k_cache_layer_mask=indexer_k_cache_layer_mask,
     )
 
 
@@ -263,6 +273,44 @@ def test_v1_dsa_indexer_page_table_is_replicated_with_per_layer_entries():
         manager.shutdown()
 
 
+@pytest.mark.cuda
+def test_v1_dsa_masked_indexer_page_table_covers_owning_layers():
+    """The masked indexer view covers only the owning layers.
+
+    A per-layer indexer mask (cross-layer indexer sharing, e.g. GLM 5.2) gives
+    only the owning layers a pool row, so the REPLICATED indexer view covers
+    exactly that subset -- one entry per owning layer mapped to its packed row
+    -- instead of one entry per LG layer.
+    """
+    # Of the 4 layers, only local layers 0 and 2 own an indexer K cache row.
+    manager = _make_v1_dsa_manager(indexer_k_cache_layer_mask=[True, False, True, False])
+    try:
+        page_table = build_page_table(manager)
+        lg = page_table.layer_groups[0]
+        assert len(lg.pool_views) == 2
+        _, idx_view = lg.pool_views
+        assert idx_view.mapper_kind == MapperKind.REPLICATED
+        assert idx_view.pool_role == frozenset({"indexer_k"})
+
+        # Only the two owning layers appear -- a strict subset of the LG.
+        owning = sorted(int(e["local_layer_id"]) for e in idx_view.buffer_entries)
+        assert owning == [0, 2]
+
+        # The pool holds one row per owning layer; entries pack contiguously in
+        # owning (local-layer) order, so layer 0 -> row 0, layer 2 -> row 1.
+        idx_pool = get_physical_pool(page_table, 0, idx_view.pool_idx)
+        sizes = {int(e["size"]) for e in idx_view.buffer_entries}
+        assert len(sizes) == 1
+        per_layer = sizes.pop()
+        assert per_layer * len(idx_view.buffer_entries) == idx_pool.slot_bytes
+        offset_by_layer = {
+            int(e["local_layer_id"]): int(e["offset"]) for e in idx_view.buffer_entries
+        }
+        assert offset_by_layer == {0: 0, 2: per_layer}
+    finally:
+        manager.shutdown()
+
+
 @pytest.mark.cuda
 def test_v1_dsa_indexer_replicated_transfer_across_pp():
     """PP1 ctx sends the DSA indexer K cache into two PP2 gen ranks.

tests/unittest/llmapi/test_llm_args.py

@@ -3871,13 +3871,7 @@ def test_resolve_default_backend_env_priority(self, monkeypatch):
 
 
 class TestDeepseekTransceiverPreference:
-    """Per-architecture preferred KV-cache transceiver runtime.
-
-    DeepseekV3ForCausalLM and DeepseekV32ForCausalLM prefer the Python KV-cache
-    transceiver, while GlmMoeDsaForCausalLM (GLM 5.2) requires the C++ transceiver
-    because its per-layer masked DSA indexer k-cache pool is not supported by the
-    Python (v2) transceiver.
-    """
+    """DeepseekV3ForCausalLM, DeepseekV32ForCausalLM, and GlmMoeDsaForCausalLM all prefer the Python KV-cache transceiver."""
 
     @staticmethod
     def _pretrained_config(architectures, model_type):
@@ -3886,19 +3880,18 @@ def _pretrained_config(architectures, model_type):
         cfg.model_type = model_type
         return cfg
 
-    @pytest.mark.parametrize("architectures,model_type,expected", [
-        (["GlmMoeDsaForCausalLM"], "glm_moe_dsa", "CPP"),
-        (["DeepseekV3ForCausalLM"], "deepseek_v3", "PYTHON"),
-        (["DeepseekV32ForCausalLM"], "deepseek_v32", "PYTHON"),
+    @pytest.mark.parametrize("architectures,model_type", [
+        (["GlmMoeDsaForCausalLM"], "glm_moe_dsa"),
+        (["DeepseekV3ForCausalLM"], "deepseek_v3"),
+        (["DeepseekV32ForCausalLM"], "deepseek_v32"),
     ])
     def test_preference_per_architecture(self, architectures: list[str],
-                                         model_type: str,
-                                         expected: str) -> None:
+                                         model_type: str) -> None:
         from tensorrt_llm._torch.models.modeling_deepseekv3 import \
             DeepseekV3ForCausalLM
         cfg = self._pretrained_config(architectures, model_type)
         assert DeepseekV3ForCausalLM.get_preferred_transceiver_runtime(
-            cfg) == expected
+            cfg) == "PYTHON"
 
     def test_prefers_python_without_config(self) -> None:
         """Preference is unconditional without a pretrained config."""