"""sagellm_core.pd — PD 分离引擎模块.

本模块实现 Prefill-Decode (PD) 分离推理体系：

- **KVTransferChannel**: 通过 sagellm-comm 在 PrefillEngine 和 DecodeEngine
  之间传输 KV 块
- **PrefillEngine**: 独立 Prefill 引擎，执行前向传播 + 发送 KV
- **DecodeEngine**: 独立 Decode 引擎，接收 KV + 自回归生成

Quick Start::

    from sagellm_core.pd import PrefillEngine, DecodeEngine, KVTransferChannel
    from sagellm_comm import get_comm_backend

    # CPU-mock 模式（同进程测试，无需真实 comm）
    channel = KVTransferChannel(comm=None, local_rank=0, remote_rank=1)

    prefill_eng = PrefillEngine(
        engine=your_base_engine,
        channel=channel,
        decode_engine_id="decode-0",
        decode_rank=1,
    )
    decode_eng = DecodeEngine(
        engine=your_base_engine,
        channel=channel,
        prefill_engine_id=prefill_eng.engine_id,
        prefill_rank=0,
    )

    # Prefill 和 Decode 并发运行
    import asyncio
    prefill_task = asyncio.create_task(prefill_eng.run_prefill(request))
    # transfer_req 需同步，实际由 PDCoordinator / PDRouter 传递
    decode_task  = asyncio.create_task(decode_eng.run_decode(request, transfer_req))
    prefill_result, decode_result = await asyncio.gather(prefill_task, decode_task)
"""

from __future__ import annotations

from sagellm_core.pd.decode_engine import DecodeEngine, DecodeResult
from sagellm_core.pd.kv_channel import KVTransferChannel
from sagellm_core.pd.prefill_engine import PrefillEngine, PrefillResult

__all__ = [
    "KVTransferChannel",
    "PrefillEngine",
    "PrefillResult",
    "DecodeEngine",
    "DecodeResult",
]
