diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d567596628e1..8cfeafecf943 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1724,6 +1724,35 @@ def to_sparse_metadata_params(self, **kwargs): layout=CSA2Layout.from_hf_config(kwargs.get("pretrained_config"))) +class GlmKpoolSparseAttentionConfig(BaseSparseAttentionConfig): + """Use checkpoint-owned GLM-5.3-Flash k-pool indexer geometry with the PyTorch sparse-MLA backend. + + The indexer top-k, pool size and key width come from the checkpoint. + """ + algorithm: Literal["glm_kpool"] = "glm_kpool" + + def supports_backend(self, backend: str) -> bool: + return backend == "pytorch" + + def to_sparse_params(self, **kwargs): + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.params import \ + GlmKpoolSparseParams + from tensorrt_llm._torch.pyexecutor.config_utils import \ + unwrap_glm5_next_text_config + + pretrained_config = kwargs.get("pretrained_config") + if pretrained_config is None: + return GlmKpoolSparseParams() + config = unwrap_glm5_next_text_config(pretrained_config) + return GlmKpoolSparseParams( + index_topk=int(config.index_topk), + index_kpool=int(config.index_kpool), + index_always_select_tail=bool( + config.index_kpool_always_select_tail), + index_head_dim=int(config.index_head_dim), + ) + + class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig): """Configuration for skip softmax attention.""" algorithm: Literal["skip_softmax"] = Field(default="skip_softmax") @@ -4421,6 +4450,7 @@ def _to_pybind(self): DeepSeekSparseAttentionConfig, DeepSeekV4SparseAttentionConfig, CSA2SparseAttentionConfig, + GlmKpoolSparseAttentionConfig, SkipSoftmaxAttentionConfig, MiniMaxM3SparseAttentionConfig, ], diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 8b268a9df3f2..8caac649167f 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -2619,6 +2619,20 @@ "title": "EncodeExtraInputSpec", "type": "object" }, + "GlmKpoolSparseAttentionConfig": { + "additionalProperties": false, + "description": "Use checkpoint-owned GLM-5.3-Flash k-pool indexer geometry with the PyTorch sparse-MLA backend.\n\nThe indexer top-k, pool size and key width come from the checkpoint.", + "properties": { + "algorithm": { + "const": "glm_kpool", + "default": "glm_kpool", + "title": "Algorithm", + "type": "string" + } + }, + "title": "GlmKpoolSparseAttentionConfig", + "type": "object" + }, "GmsConfig": { "additionalProperties": false, "description": "Prototype configuration for GPU Memory Service (GMS) weight sharing.\n\nComposes orthogonally with ``checkpoint_format`` on ``TorchLlmArgs``:\nGMS controls *where* weights live (a node-shared GPU memory pool\nso multiple TRT-LLM instances zero-copy share the same bytes),\nwhile ``checkpoint_format`` controls *how* they get there (HF disk,\nMX P2P, etc.). Effective only when ``TorchLlmArgs.load_format ==\nLoadFormat.GMS``; setting any non-default field here without that\nload_format triggers a warning via :meth:`TorchLlmArgs.validate_gms_config`.\n\nTwo roles are decided at connect time by the GMS daemon, not by\nconfig:\n\n- **RW (writer)**: the first instance loads weights into the pool\n via the normal checkpoint-loader pipeline, then commits them.\n- **RO (reader)**: subsequent instances zero-copy materialize the\n already-committed layout \u2014 no disk I/O, no per-instance copies.\n\nSee :class:`tensorrt_llm._torch.memory.gpu_memory_backend.GMSBackend`\nfor the integration adapter and the ``LoadFormat.GMS`` branch in\n``ModelLoader.load`` for orchestration.", @@ -7526,6 +7540,7 @@ "csa2": "#/$defs/CSA2SparseAttentionConfig", "deepseek_v4": "#/$defs/DeepSeekV4SparseAttentionConfig", "dsa": "#/$defs/DeepSeekSparseAttentionConfig", + "glm_kpool": "#/$defs/GlmKpoolSparseAttentionConfig", "minimax_m3": "#/$defs/MiniMaxM3SparseAttentionConfig", "qsa": "#/$defs/QSASparseAttentionConfig", "rocket": "#/$defs/RocketSparseAttentionConfig", @@ -7549,6 +7564,9 @@ { "$ref": "#/$defs/CSA2SparseAttentionConfig" }, + { + "$ref": "#/$defs/GlmKpoolSparseAttentionConfig" + }, { "$ref": "#/$defs/SkipSoftmaxAttentionConfig" }, diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index c9a19899dff5..e280de115928 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -2908,6 +2908,7 @@ "csa2": "#/$defs/CSA2SparseAttentionConfig", "deepseek_v4": "#/$defs/DeepSeekV4SparseAttentionConfig", "dsa": "#/$defs/DeepSeekSparseAttentionConfig", + "glm_kpool": "#/$defs/GlmKpoolSparseAttentionConfig", "minimax_m3": "#/$defs/MiniMaxM3SparseAttentionConfig", "qsa": "#/$defs/QSASparseAttentionConfig", "rocket": "#/$defs/RocketSparseAttentionConfig", @@ -2931,6 +2932,9 @@ { "$ref": "#/$defs/CSA2SparseAttentionConfig" }, + { + "$ref": "#/$defs/GlmKpoolSparseAttentionConfig" + }, { "$ref": "#/$defs/SkipSoftmaxAttentionConfig" }, @@ -4169,6 +4173,20 @@ "title": "EncodeExtraInputSpec", "type": "object" }, + "GlmKpoolSparseAttentionConfig": { + "additionalProperties": false, + "description": "Use checkpoint-owned GLM-5.3-Flash k-pool indexer geometry with the PyTorch sparse-MLA backend.\n\nThe indexer top-k, pool size and key width come from the checkpoint.", + "properties": { + "algorithm": { + "const": "glm_kpool", + "default": "glm_kpool", + "title": "Algorithm", + "type": "string" + } + }, + "title": "GlmKpoolSparseAttentionConfig", + "type": "object" + }, "GmsConfig": { "additionalProperties": false, "description": "Prototype configuration for GPU Memory Service (GMS) weight sharing.\n\nComposes orthogonally with ``checkpoint_format`` on ``TorchLlmArgs``:\nGMS controls *where* weights live (a node-shared GPU memory pool\nso multiple TRT-LLM instances zero-copy share the same bytes),\nwhile ``checkpoint_format`` controls *how* they get there (HF disk,\nMX P2P, etc.). Effective only when ``TorchLlmArgs.load_format ==\nLoadFormat.GMS``; setting any non-default field here without that\nload_format triggers a warning via :meth:`TorchLlmArgs.validate_gms_config`.\n\nTwo roles are decided at connect time by the GMS daemon, not by\nconfig:\n\n- **RW (writer)**: the first instance loads weights into the pool\n via the normal checkpoint-loader pipeline, then commits them.\n- **RO (reader)**: subsequent instances zero-copy materialize the\n already-committed layout \u2014 no disk I/O, no per-instance copies.\n\nSee :class:`tensorrt_llm._torch.memory.gpu_memory_backend.GMSBackend`\nfor the integration adapter and the ``LoadFormat.GMS`` branch in\n``ModelLoader.load`` for orchestration.", @@ -9246,6 +9264,7 @@ "csa2": "#/$defs/CSA2SparseAttentionConfig", "deepseek_v4": "#/$defs/DeepSeekV4SparseAttentionConfig", "dsa": "#/$defs/DeepSeekSparseAttentionConfig", + "glm_kpool": "#/$defs/GlmKpoolSparseAttentionConfig", "minimax_m3": "#/$defs/MiniMaxM3SparseAttentionConfig", "qsa": "#/$defs/QSASparseAttentionConfig", "rocket": "#/$defs/RocketSparseAttentionConfig", @@ -9269,6 +9288,9 @@ { "$ref": "#/$defs/CSA2SparseAttentionConfig" }, + { + "$ref": "#/$defs/GlmKpoolSparseAttentionConfig" + }, { "$ref": "#/$defs/SkipSoftmaxAttentionConfig" }, diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e58f42598f2d..d0018f105ec4 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1380,6 +1380,7 @@ "csa2", "dsa", "deepseek_v4", + "glm_kpool", "minimax_m3", "qsa", "rocket", diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index 052f86db7c1b..31f22e079df5 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -752,6 +752,31 @@ def test_cache_manager_routing_guards(route, monkeypatch): get_kv_cache_manager_cls(config, kv, True, transceiver) +@pytest.mark.cpu_only +def test_glm_kpool_sparse_config_lowers_the_checkpoint_geometry(): + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import ( + GlmKpoolSparseAttention, + GlmKpoolSparseParams, + ) + from tensorrt_llm._torch.attention.backends.utils import get_attention_backend + from tensorrt_llm._torch.pyexecutor.config_utils import unwrap_glm5_next_text_config + from tensorrt_llm.llmapi.llm_args import GlmKpoolSparseAttentionConfig + + pretrained = _config() + text = unwrap_glm5_next_text_config(pretrained) + sparse = GlmKpoolSparseAttentionConfig() + params = sparse.to_sparse_params(pretrained_config=pretrained) + assert params == GlmKpoolSparseParams( + index_topk=text.index_topk, + index_kpool=text.index_kpool, + index_always_select_tail=text.index_kpool_always_select_tail, + index_head_dim=text.index_head_dim, + ) + assert sparse.to_sparse_params(pretrained_config=text) == params + assert sparse.supports_backend("pytorch") and not sparse.supports_backend("tensorrt") + assert get_attention_backend("TRTLLM", sparse_params=params) is GlmKpoolSparseAttention + + @pytest.mark.cpu_only @pytest.mark.parametrize("fp8_kv_cache", [False, True], ids=["bf16-kv", "fp8-kv"]) @pytest.mark.parametrize( diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index d2d1739473ff..5b1c06a57a18 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -323,7 +323,7 @@ methods: default: null status: prototype sparse_attention_config: - annotation: Union[tensorrt_llm.llmapi.llm_args.QSASparseAttentionConfig, tensorrt_llm.llmapi.llm_args.RocketSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekV4SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.CSA2SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.SkipSoftmaxAttentionConfig, tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig, NoneType] + annotation: Union[tensorrt_llm.llmapi.llm_args.QSASparseAttentionConfig, tensorrt_llm.llmapi.llm_args.RocketSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekV4SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.CSA2SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.GlmKpoolSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.SkipSoftmaxAttentionConfig, tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig, NoneType] default: null status: prototype kv_cache_compression_config: diff --git a/tests/unittest/usage/test_llmapi_config_capture.py b/tests/unittest/usage/test_llmapi_config_capture.py index 13af3416689e..b60e131aca89 100644 --- a/tests/unittest/usage/test_llmapi_config_capture.py +++ b/tests/unittest/usage/test_llmapi_config_capture.py @@ -741,6 +741,7 @@ def test_collect_llm_api_config_captures_sparse_algorithm_for_every_arm(): CSA2SparseAttentionConfig, DeepSeekSparseAttentionConfig, DeepSeekV4SparseAttentionConfig, + GlmKpoolSparseAttentionConfig, MiniMaxM3SparseAttentionConfig, QSASparseAttentionConfig, RocketSparseAttentionConfig, @@ -753,6 +754,7 @@ def test_collect_llm_api_config_captures_sparse_algorithm_for_every_arm(): DeepSeekSparseAttentionConfig(), DeepSeekV4SparseAttentionConfig(), CSA2SparseAttentionConfig(), + GlmKpoolSparseAttentionConfig(), RocketSparseAttentionConfig(), SkipSoftmaxAttentionConfig(), MiniMaxM3SparseAttentionConfig(),