From 83eb699e352e1a321e110bb8434c40cae814707a Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:19:07 -0700 Subject: [PATCH] [None][feat] Support the masked DSA indexer k-cache pool in the Python cache transceiver The Python (v2) KV-cache transceiver previously raised NotImplementedError for the per-layer masked DSA indexer k-cache pool (cross-layer indexer sharing, e.g. GLM 5.2), forcing those checkpoints onto the C++ transceiver. Teach build_page_table's indexer REPLICATED view to cover only the indexer-owning layers: one buffer entry per owning layer, mapped to its packed row via get_indexer_k_cache_pool_layer_idx, and skip the pool for a layer group with no owning layers. The Python transfer machinery already matches peers per-pool by role + global_layer_id overlap, so a masked subset transfers correctly (including PP reshard) with no further changes; the dense/unmasked layout is byte-for-byte unchanged. With Python support in place, GlmMoeDsaForCausalLM prefers the Python transceiver again like the other DeepSeek-family checkpoints, reverting the CPP override that was added only because Python lacked masked-pool support. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../disaggregation/resource/kv_extractor.py | 99 +++++++++---------- .../_torch/models/modeling_deepseekv3.py | 26 ++--- .../unittest/disaggregated/test_extractor.py | 52 +++++++++- tests/unittest/llmapi/test_llm_args.py | 21 ++-- 4 files changed, 115 insertions(+), 83 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index bf8e91006b7d..a6eddbbf25b8 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/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 = [ diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 0ede1dc4b332..6e026a31143f 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/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]): diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 69fb84b26887..5b2f521da3ef 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/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. diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 453c91da3613..1fb188c84563 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/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."""