[XPU][CI] Make test_mamba_prefix_cache block-size agnostic (#58269)
Signed-off-by: Lin, Fanli <fanli.lin@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
Fanli Lin
@@ -55,7 +55,10 @@ async_scheduling_mode = False
|
||||
num_accepted_tokens = 1
|
||||
prompt_token_ids: list[int] = []
|
||||
MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
|
||||
BLOCK_SIZE = 560
|
||||
# Must be a value no platform rewrites, since the step actions below encode
|
||||
# where block boundaries fall. XPU rounds GDN block sizes up to a multiple of 64
|
||||
# (XPUPlatform.update_block_size_for_backend), so use one.
|
||||
BLOCK_SIZE = 576
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
NUM_HIDDEN_LAYERS = 1
|
||||
cur_step_action_idx = 0
|
||||
@@ -541,26 +544,27 @@ def get_mamba_prefix_cache_step_configs(
|
||||
async_scheduling: bool = False,
|
||||
) -> dict[str, TestConfig]:
|
||||
a = async_scheduling
|
||||
B = BLOCK_SIZE
|
||||
tests = {
|
||||
"accept_1": TestConfig(
|
||||
num_prompt_tokens=554,
|
||||
num_prompt_tokens=B - 6,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=1,
|
||||
step_actions=[
|
||||
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
555, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
|
||||
B - 5, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
|
||||
),
|
||||
StepAction(
|
||||
556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
|
||||
B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
|
||||
),
|
||||
StepAction(557, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(558, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
|
||||
StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 3, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(B - 2, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 1, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
|
||||
StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
561,
|
||||
B + 1,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
@@ -570,19 +574,19 @@ def get_mamba_prefix_cache_step_configs(
|
||||
),
|
||||
# test case 2.1: no hit, accept 2 tokens
|
||||
"accept_2_1": TestConfig(
|
||||
num_prompt_tokens=554,
|
||||
num_prompt_tokens=B - 6,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=2,
|
||||
step_actions=[
|
||||
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
|
||||
B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
|
||||
),
|
||||
StepAction(558, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)),
|
||||
StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 2, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)),
|
||||
StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
562,
|
||||
B + 2,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
@@ -592,56 +596,56 @@ def get_mamba_prefix_cache_step_configs(
|
||||
),
|
||||
# test case 2.2: no hit, accept 2 tokens
|
||||
"accept_2_2": TestConfig(
|
||||
num_prompt_tokens=555,
|
||||
num_prompt_tokens=B - 5,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=2,
|
||||
step_actions=[
|
||||
StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
|
||||
StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
|
||||
StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 3, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
|
||||
StepAction(B - 1, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
|
||||
StepAction(
|
||||
561,
|
||||
B + 1,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B + 3, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_3_1": TestConfig(
|
||||
num_prompt_tokens=553,
|
||||
num_prompt_tokens=B - 7,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=3,
|
||||
step_actions=[
|
||||
StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(0, B - 7, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 7, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
|
||||
B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
|
||||
),
|
||||
StepAction(559, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)),
|
||||
StepAction(B - 1, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)),
|
||||
StepAction(
|
||||
562,
|
||||
B + 2,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B + 5, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_3_2": TestConfig(
|
||||
num_prompt_tokens=554,
|
||||
num_prompt_tokens=B - 6,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=3,
|
||||
step_actions=[
|
||||
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
|
||||
StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 3, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
|
||||
StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
563,
|
||||
B + 3,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
@@ -650,87 +654,87 @@ def get_mamba_prefix_cache_step_configs(
|
||||
],
|
||||
),
|
||||
"accept_3_3": TestConfig(
|
||||
num_prompt_tokens=555,
|
||||
num_prompt_tokens=B - 5,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=3,
|
||||
step_actions=[
|
||||
StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
|
||||
StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 2, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
|
||||
StepAction(
|
||||
561,
|
||||
B + 1,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B + 4, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_4_1": TestConfig(
|
||||
num_prompt_tokens=553,
|
||||
num_prompt_tokens=B - 7,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
|
||||
StepAction(0, B - 7, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 7, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 3, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
|
||||
StepAction(
|
||||
561,
|
||||
B + 1,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B + 5, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_4_2": TestConfig(
|
||||
num_prompt_tokens=554,
|
||||
num_prompt_tokens=B - 6,
|
||||
num_generated_tokens=25,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
|
||||
StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 2, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
|
||||
StepAction(
|
||||
562,
|
||||
B + 2,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B + 6, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_4_3": TestConfig(
|
||||
num_prompt_tokens=555,
|
||||
num_prompt_tokens=B - 5,
|
||||
num_generated_tokens=25,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
|
||||
StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 1, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
|
||||
StepAction(
|
||||
563,
|
||||
B + 3,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B + 7, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_4_4": TestConfig(
|
||||
num_prompt_tokens=556,
|
||||
num_prompt_tokens=B - 4,
|
||||
num_generated_tokens=25,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(556, 4, [1, 1, 1, 1], (-1, -1), (3, 0)),
|
||||
StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(0, B - 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B - 4, 4, [1, 1, 1, 1], (-1, -1), (3, 0)),
|
||||
StepAction(B, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(
|
||||
564,
|
||||
B + 4,
|
||||
4,
|
||||
[1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
@@ -739,23 +743,23 @@ def get_mamba_prefix_cache_step_configs(
|
||||
],
|
||||
),
|
||||
"prompt_block_size": TestConfig(
|
||||
num_prompt_tokens=560,
|
||||
num_prompt_tokens=B,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"prompt_2_block_size": TestConfig(
|
||||
num_prompt_tokens=560 * 2,
|
||||
num_prompt_tokens=B * 2,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B, B, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(
|
||||
560 * 2,
|
||||
B * 2,
|
||||
4,
|
||||
[1, 1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1, 1],
|
||||
(1, 2),
|
||||
@@ -764,14 +768,14 @@ def get_mamba_prefix_cache_step_configs(
|
||||
],
|
||||
),
|
||||
"prompt_2_block_size_10": TestConfig(
|
||||
num_prompt_tokens=560 * 2 + 10,
|
||||
num_prompt_tokens=B * 2 + 10,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
|
||||
StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B, B + 10, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
|
||||
StepAction(
|
||||
560 * 2 + 10,
|
||||
B * 2 + 10,
|
||||
4,
|
||||
[1, 0, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
@@ -780,14 +784,14 @@ def get_mamba_prefix_cache_step_configs(
|
||||
],
|
||||
),
|
||||
"prompt_3_block_size": TestConfig(
|
||||
num_prompt_tokens=560 * 3,
|
||||
num_prompt_tokens=B * 3,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
|
||||
StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B * 2, B, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
|
||||
StepAction(
|
||||
560 * 3,
|
||||
B * 3,
|
||||
4,
|
||||
[0, 1, 1, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1, 1],
|
||||
(2, 3),
|
||||
@@ -796,14 +800,14 @@ def get_mamba_prefix_cache_step_configs(
|
||||
],
|
||||
),
|
||||
"prompt_3_block_size_10": TestConfig(
|
||||
num_prompt_tokens=560 * 3 + 10,
|
||||
num_prompt_tokens=B * 3 + 10,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
|
||||
StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(B * 2, B + 10, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
|
||||
StepAction(
|
||||
560 * 3 + 10,
|
||||
B * 3 + 10,
|
||||
4,
|
||||
[0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1],
|
||||
(-1, -1),
|
||||
@@ -812,21 +816,21 @@ def get_mamba_prefix_cache_step_configs(
|
||||
],
|
||||
),
|
||||
"prompt_10_block_size": TestConfig(
|
||||
num_prompt_tokens=560 * 10,
|
||||
num_prompt_tokens=B * 10,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
560 * 5,
|
||||
560 * 4,
|
||||
B * 5,
|
||||
B * 4,
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
|
||||
(4, 8),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(
|
||||
560 * 9,
|
||||
560,
|
||||
B * 9,
|
||||
B,
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1]
|
||||
if a
|
||||
else [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
|
||||
@@ -834,7 +838,7 @@ def get_mamba_prefix_cache_step_configs(
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(
|
||||
560 * 10,
|
||||
B * 10,
|
||||
4,
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
|
||||
if a
|
||||
@@ -845,21 +849,21 @@ def get_mamba_prefix_cache_step_configs(
|
||||
],
|
||||
),
|
||||
"prompt_10_block_size_10": TestConfig(
|
||||
num_prompt_tokens=560 * 10 + 10,
|
||||
num_prompt_tokens=B * 10 + 10,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
560 * 5,
|
||||
560 * 4,
|
||||
B * 5,
|
||||
B * 4,
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
|
||||
(4, 8),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(
|
||||
560 * 9,
|
||||
560 + 10,
|
||||
B * 9,
|
||||
B + 10,
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1]
|
||||
if a
|
||||
else [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
|
||||
|
||||
@@ -464,13 +464,16 @@ class XPUPlatform(Platform):
|
||||
)
|
||||
cache_config.block_size = new_block_size
|
||||
logger.info(
|
||||
"[XPU]Setting attention block size to %d tokens to ensure multiple of %d, "
|
||||
"set mamba_page_size_padded to %d bytes accordingly, before was %d bytes.",
|
||||
"[XPU]Setting attention block size to %d tokens to ensure multiple of %d.",
|
||||
new_block_size,
|
||||
kernel_block_size,
|
||||
cache_config.mamba_page_size_padded,
|
||||
original_mamba_page_size_padded,
|
||||
)
|
||||
if original_mamba_page_size_padded is not None:
|
||||
logger.info(
|
||||
"[XPU]Scaled mamba_page_size_padded from %d to %d bytes accordingly.",
|
||||
original_mamba_page_size_padded,
|
||||
cache_config.mamba_page_size_padded,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def support_hybrid_kv_cache(cls) -> bool:
|
||||
|
||||
Reference in New Issue
Block a user