From 7178515b37ccfe38f9d7e96428b616a9ca1c2ecd Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Fri, 13 Nov 2026 23:37:47 +0800 Subject: [PATCH] fix: reject client pd kv transport params --- lightllm/server/core/objs/sampling_params.py | 2 +- .../server/core/objs/test_sampling_params.py | 15 ++++++++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py index a1fbba9cde..d2b0a859fc 100644 --- a/lightllm/server/core/objs/sampling_params.py +++ b/lightllm/server/core/objs/sampling_params.py @@ -399,8 +399,8 @@ def init(self, tokenizer, **kwargs): self.exponential_decay_length_penalty = ExponentialDecayLengthPenalty() self.exponential_decay_length_penalty.initialize(kwargs.get("exponential_decay_length_penalty", (1, 1.0))) + # The trusted PD manager populates this field after request parsing. self.pd_kv_trans_params = PDKVTransParamObj() - self.pd_kv_trans_params.set(kwargs.get("pd_kv_trans_params", None)) self.pd_master_node_id = NodeUUId() self.pd_master_node_id.initialize(kwargs.get("pd_master_node_id", 0)) diff --git a/unit_tests/server/core/objs/test_sampling_params.py b/unit_tests/server/core/objs/test_sampling_params.py index b78608bf55..e216133da4 100644 --- a/unit_tests/server/core/objs/test_sampling_params.py +++ b/unit_tests/server/core/objs/test_sampling_params.py @@ -151,7 +151,6 @@ def test_node_uuid_initialization(): def test_sampling_params_initialization(): params = SamplingParams() pd_master_node_id = 12345678901234567890 - pd_kv_trans_params = b"pd-kv-transport-params" data = { "best_of": 2, "n": 2, @@ -176,7 +175,6 @@ def test_sampling_params_initialization(): "stop_sequences": [[2, 1], [3, 4]], "exponential_decay_length_penalty": (1, 1.0), "pd_master_node_id": pd_master_node_id, - "pd_kv_trans_params": pd_kv_trans_params, } params.init(None, **data) @@ -187,7 +185,18 @@ def test_sampling_params_initialization(): assert params.temperature == 1.0 assert params.stop_sequences.size == 2 assert params.pd_master_node_id.get() == pd_master_node_id - assert params.pd_kv_trans_params.get() == pd_kv_trans_params + + +def test_client_cannot_set_pd_kv_trans_params(): + params = SamplingParams() + params.init(None, pd_kv_trans_params="AAAAdeadbeefBBBB") + + assert params.pd_kv_trans_params.data_len == 0 + assert params.pd_kv_trans_params.get() == b"" + + trusted_pd_data = b"pd-kv-transport-params" + params.pd_kv_trans_params.set(trusted_pd_data) + assert params.pd_kv_trans_params.get() == trusted_pd_data @pytest.mark.parametrize(