[Bugfix] Fix Kimi-Linear state transfer across heterogeneous TP (#32262)

This commit is contained in:
YAMY
2026-07-24 10:31:17 -07:00
committed by GitHub
parent 5da0b6ec39
commit 2428f56145
10 changed files with 306 additions and 130 deletions
@@ -50,6 +50,9 @@ class PDDisaggregationServerBase(CustomTestCase):
capture_per_side_logs: ClassVar[bool] = False
extra_prefill_env: ClassVar[dict[str, str]] = {}
extra_decode_env: ClassVar[dict[str, str]] = {}
prefill_tp_size: ClassVar[int] = 1
decode_tp_size: ClassVar[int] = 1
decode_base_gpu_id: ClassVar[int] = 1
_prefill_stdout_buf: ClassVar[Optional[io.StringIO]] = None
_prefill_stderr_buf: ClassVar[Optional[io.StringIO]] = None
_decode_stdout_buf: ClassVar[Optional[io.StringIO]] = None
@@ -114,7 +117,7 @@ class PDDisaggregationServerBase(CustomTestCase):
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
str(cls.prefill_tp_size),
] + list(cls.extra_prefill_args)
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
@@ -139,9 +142,9 @@ class PDDisaggregationServerBase(CustomTestCase):
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
str(cls.decode_tp_size),
"--base-gpu-id",
"1",
str(cls.decode_base_gpu_id),
] + list(cls.extra_decode_args)
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(