breakdown kernel update (#8334)

This commit is contained in:
Zhiqiang Xie
2025-07-24 17:33:17 -07:00
committed by GitHub
parent 145482f422
commit d40846d456
2 changed files with 44 additions and 80 deletions

View File

@@ -101,9 +101,7 @@ def test_transfer_kv(
dst_pool_kernel[layer_idx_to_test],
src_indices_device,
dst_indices_device,
io_backend="kernel",
page_size=page_size,
item_size=item_size,
item_size=item_size * dtype.itemsize,
)
transfer_kv_direct(
[src_pool_host[layer_idx_to_test]],
@@ -138,7 +136,6 @@ def test_transfer_kv(
dst_layers_device,
src_indices_device,
dst_indices_device,
io_backend="kernel",
item_size=item_size * dtype.itemsize,
num_layers=num_layers,
)
@@ -173,9 +170,7 @@ def test_transfer_kv(
dst_v_pool_kernel[layer_idx_to_test],
src_indices_device,
dst_indices_device,
io_backend="kernel",
page_size=page_size,
item_size=item_size,
item_size=item_size * dtype.itemsize,
)
transfer_kv_direct(
[src_k_pool[layer_idx_to_test], src_v_pool[layer_idx_to_test]],
@@ -235,7 +230,6 @@ def test_transfer_kv(
dst_v_layers_device,
src_indices_device,
dst_indices_device,
io_backend="kernel",
item_size=item_size * dtype.itemsize,
num_layers=num_layers,
)