[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/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

@@ -1912,25 +1912,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_bounce.py

@@ -369,6 +369,33 @@ def _recv_req(block_counts, rid=1, slice_id=0):
 
 @pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings")
 class TestFanInReserve:
+    def test_block_bytes_include_distinct_physical_pools_once(self) -> None:
+        entries = np.array([], dtype=BUFFER_ENTRY_DTYPE)
+        page_table = KVCachePageTable(
+            tokens_per_block=32,
+            layer_groups=[
+                AttentionLayerGroup(
+                    pool_group_idx=0,
+                    local_layers=[LocalLayer(local_layer_id=0, global_layer_id=0)],
+                    pool_views=[
+                        PoolView(pool_idx=0, buffer_entries=entries),
+                        PoolView(pool_idx=1, buffer_entries=entries),
+                        PoolView(pool_idx=1, buffer_entries=entries),
+                    ],
+                )
+            ],
+            pool_groups=[
+                PhysicalPoolGroup(
+                    pools=[
+                        PhysicalPool(base_address=0x200000, slot_bytes=100, num_slots=8),
+                        PhysicalPool(base_address=0x300000, slot_bytes=25, num_slots=8),
+                    ]
+                )
+            ],
+        )
+
+        assert btr.block_bytes_per_group(page_table) == [125]
+
     def test_reserve_stamps_base_and_per_writer(self, monkeypatch):
         t = _make_transport(monkeypatch, block_bytes_per_group=[100])
         req = _recv_req([2])  # total = 2 * 100 = 200