diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 2ccc34a3cda4..c95aac6c8147 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -28,7 +28,7 @@ import copy import math import os -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Literal, Optional, Tuple import torch import triton @@ -1893,6 +1893,30 @@ def forward( class DeepseekV3ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV3Model, PretrainedConfig]): + @classmethod + def get_preferred_transceiver_runtime(cls, + pretrained_config: Any = None + ) -> Optional[Literal["PYTHON"]]: + """GLM-5 family checkpoints default to the Python (v2) KV-cache transceiver. + + This implementation class is shared by DeepSeek-V3/V3.2 and the GLM-5 family — both + GLM-5 and GLM-5.2 declare ``GlmMoeDsaForCausalLM`` / ``glm_moe_dsa`` — so the preference + is differentiated per checkpoint: only GLM checkpoints opt into the Python transceiver. + The MLA backbone transfers a large latent KV, which the Python transceiver handles better + in disaggregated serving. This is only adopted when the user leaves + ``cache_transceiver_config.transceiver_runtime`` at 'auto' and the effective backend is + NIXL; otherwise the C++ transceiver is used. + """ + if pretrained_config is None: + return 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 "PYTHON" + return None + def __init__(self, model_config: ModelConfig[PretrainedConfig]): self.mapping_with_cp = None # Note: Currently the usage of mapping is all over the place making its usage brittle diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp4ep4_gentp4ep4_glm5_nvfp4_dp_tllm.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp4ep4_gentp4ep4_glm5_nvfp4_dp_tllm.yaml index 57cf5bcbebed..fc790618269b 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp4ep4_gentp4ep4_glm5_nvfp4_dp_tllm.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp4ep4_gentp4ep4_glm5_nvfp4_dp_tllm.yaml @@ -20,7 +20,8 @@ context_servers: cuda_graph_config: null print_iter_log: true cache_transceiver_config: - backend: DEFAULT + backend: NIXL + transceiver_runtime: PYTHON max_tokens_in_buffer: 16384 generation_servers: num_instances: 1 @@ -55,5 +56,6 @@ generation_servers: - 1024 print_iter_log: true cache_transceiver_config: - backend: DEFAULT + backend: NIXL + transceiver_runtime: PYTHON max_tokens_in_buffer: 16384 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con1_ctx1_dep4_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con1_ctx1_dep4_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml index 6678fdfc4f8b..9a514f62dc65 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con1_ctx1_dep4_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con1_ctx1_dep4_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml @@ -61,6 +61,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -89,5 +90,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con4096_ctx1_dep4_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con4096_ctx1_dep4_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml index a54e7d0335cf..4b01710438c2 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con4096_ctx1_dep4_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con4096_ctx1_dep4_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml @@ -66,6 +66,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -93,5 +94,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con512_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con512_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml index e8bcff454293..1b427f4f2c05 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con512_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_1k1k_con512_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml @@ -63,6 +63,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -90,5 +91,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con1024_ctx1_dep4_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con1024_ctx1_dep4_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml index 37d162efda13..9dc48b40047f 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con1024_ctx1_dep4_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con1024_ctx1_dep4_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml @@ -66,6 +66,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -93,5 +94,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index 3761761e4e9d..d9fef4a05711 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -61,6 +61,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -89,5 +90,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con512_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con512_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml index 876121e49a57..6f95b245621a 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con512_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_glm-5-fp4_8k1k_con512_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml @@ -63,6 +63,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -90,5 +91,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml index 7e5aa487facc..e1f98578eda1 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml @@ -61,6 +61,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -89,5 +90,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml index f6334f686d1c..e2d50f587b7b 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml @@ -66,6 +66,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -93,5 +94,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml index 8ef0b3db0473..ebc92211e8ee 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml @@ -63,6 +63,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -90,5 +91,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml index 22a64298673b..0260accf37ea 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml @@ -66,6 +66,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -93,5 +94,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index 1e6c43b2ddf7..67ae7c9d5c3a 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -61,6 +61,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -89,5 +90,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml index 777987a01492..8a8ddd2f7e96 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml @@ -63,6 +63,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: &id001 decoding_type: MTP @@ -90,5 +91,6 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 16384 backend: NIXL + transceiver_runtime: PYTHON disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index b5410a8edfdb..89e2cdcfab86 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3385,3 +3385,60 @@ def test_resolve_default_backend_env_priority(self, monkeypatch): # An explicit backend bypasses the env vars entirely. assert CacheTransceiverConfig( backend="UCX")._resolve_default_backend() == ("UCX", None) + + +class TestGlm5TransceiverPreference: + """GLM-5 defaults to the Python KV-cache transceiver in disagg. + + DeepseekV3ForCausalLM is shared by DeepSeek-V3/V3.2 and GLM-5 + (GlmMoeDsaForCausalLM); the preference must apply to GLM checkpoints + only. + """ + + @staticmethod + def _pretrained_config(architectures, model_type): + from transformers import PretrainedConfig + cfg = PretrainedConfig(architectures=architectures) + cfg.model_type = model_type + return cfg + + @pytest.mark.parametrize( + "architectures,model_type,expected", + [ + (["GlmMoeDsaForCausalLM"], "glm_moe_dsa", "PYTHON"), + (["DeepseekV3ForCausalLM"], "deepseek_v3", None), + (["DeepseekV32ForCausalLM"], "deepseek_v32", None), + # Each predicate in isolation: the architecture match and the + # model_type fallback must each suffice on their own. + (["GlmMoeDsaForCausalLM"], "deepseek_v32", "PYTHON"), + (["DeepseekV32ForCausalLM"], "glm_moe_dsa", "PYTHON"), + ]) + def test_preference_per_architecture(self, architectures, model_type, + expected): + from tensorrt_llm._torch.models.modeling_deepseekv3 import \ + DeepseekV3ForCausalLM + cfg = self._pretrained_config(architectures, model_type) + assert DeepseekV3ForCausalLM.get_preferred_transceiver_runtime( + cfg) == expected + + def test_no_config_defers_to_cpp(self): + """Without a pretrained config the class defers to the C++ default.""" + from tensorrt_llm._torch.models.modeling_deepseekv3 import \ + DeepseekV3ForCausalLM + assert DeepseekV3ForCausalLM.get_preferred_transceiver_runtime() is None + + def test_glm5_resolves_auto_to_python_on_nixl(self): + """GLM-5 on NIXL adopts the Python transceiver from 'auto'. + + End-to-end through _resolve_transceiver_runtime_auto with the real + model class and a GLM pretrained config. + """ + from tensorrt_llm._torch.models.modeling_deepseekv3 import \ + DeepseekV3ForCausalLM + args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig(backend="NIXL"), + ) + cfg = self._pretrained_config(["GlmMoeDsaForCausalLM"], "glm_moe_dsa") + _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg) + assert args.cache_transceiver_config.transceiver_runtime == "PYTHON"