wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
from tokenspeed.runtime.pd.mooncake.receiver import (
|
|
_build_buffer_layout_pair,
|
|
)
|
|
from tokenspeed.runtime.pd.transfer_plan import (
|
|
BufferKind,
|
|
ParallelLayout,
|
|
PDTransferPlanner,
|
|
)
|
|
|
|
|
|
def test_replicated_prefill_kv_heads_transfer_to_decode_full_kv_heads():
|
|
prefill_buffer, decode_buffer = _build_buffer_layout_pair(
|
|
buffer_index=0,
|
|
buffer_kind=BufferKind.TARGET_K,
|
|
sharded_axis="kv_head",
|
|
prefill_item_len=16_384,
|
|
decode_item_len=32_768,
|
|
prefill_unit_len=256,
|
|
decode_unit_len=256,
|
|
prefill_tp_size=4,
|
|
decode_tp_size=1,
|
|
)
|
|
|
|
assert prefill_buffer.logical_size == 128
|
|
assert prefill_buffer.tp_replica_group_size == 2
|
|
assert decode_buffer.logical_size == 128
|
|
assert decode_buffer.tp_replica_group_size == 1
|
|
|
|
planner = PDTransferPlanner(
|
|
prefill_layout=ParallelLayout(role="prefill", world_size=4),
|
|
decode_layout=ParallelLayout(role="decode", world_size=1),
|
|
prefill_buffers=(prefill_buffer,),
|
|
decode_buffers=(decode_buffer,),
|
|
)
|
|
plan = planner.plan_for_decode_rank(0)
|
|
|
|
assert plan.plan_kind == "fragmented"
|
|
assert plan.target_prefill_ranks == (0, 2)
|
|
assert plan.required_prefill_response_num == 2
|
|
assert plan.required_dst_info_num_by_prefill_rank == {0: 1, 2: 1}
|
|
|
|
first_head = plan.fragments_by_prefill_rank[0][0]
|
|
second_head = plan.fragments_by_prefill_rank[2][0]
|
|
assert first_head.src_byte_offset == 0
|
|
assert first_head.dst_byte_offset == 0
|
|
assert first_head.bytes_per_page == 16_384
|
|
assert second_head.src_byte_offset == 0
|
|
assert second_head.dst_byte_offset == 16_384
|
|
assert second_head.bytes_per_page == 16_384
|