Fix Kimi K2.5 dp attention+ spec decoding launch crash (#21391)
This commit is contained in:
@@ -150,6 +150,17 @@ class LlamaModel(nn.Module):
|
|||||||
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
if input_embeds is None:
|
if input_embeds is None:
|
||||||
|
embeds = forward_batch.mm_input_embeds
|
||||||
|
if (
|
||||||
|
forward_batch.forward_mode.is_extend()
|
||||||
|
and forward_batch.contains_mm_inputs()
|
||||||
|
and not forward_batch.forward_mode.is_draft_extend(include_v2=True)
|
||||||
|
):
|
||||||
|
assert embeds is not None
|
||||||
|
embeds = torch.cat(
|
||||||
|
[embeds[:-1], self.embed_tokens(input_ids[-1].unsqueeze(0))]
|
||||||
|
)
|
||||||
|
if embeds is None:
|
||||||
embeds = self.embed_tokens(input_ids)
|
embeds = self.embed_tokens(input_ids)
|
||||||
else:
|
else:
|
||||||
embeds = input_embeds
|
embeds = input_embeds
|
||||||
|
|||||||
@@ -38,11 +38,15 @@ class TestKimiK25(unittest.TestCase):
|
|||||||
"--speculative-num-steps=3",
|
"--speculative-num-steps=3",
|
||||||
"--speculative-eagle-topk=1",
|
"--speculative-eagle-topk=1",
|
||||||
"--speculative-num-draft-tokens=4",
|
"--speculative-num-draft-tokens=4",
|
||||||
"--mem-frac=0.85",
|
|
||||||
"--model-loader-extra-config",
|
"--model-loader-extra-config",
|
||||||
'{"enable_multithread_load": true, "num_threads": 64}',
|
'{"enable_multithread_load": true, "num_threads": 64}',
|
||||||
]
|
]
|
||||||
|
|
||||||
|
dp_attn_args = [
|
||||||
|
"--dp=8",
|
||||||
|
"--enable-dp-attention",
|
||||||
|
]
|
||||||
|
|
||||||
variants = [
|
variants = [
|
||||||
ModelLaunchSettings(
|
ModelLaunchSettings(
|
||||||
KIMI_K25_MODEL_PATH,
|
KIMI_K25_MODEL_PATH,
|
||||||
@@ -56,6 +60,12 @@ class TestKimiK25(unittest.TestCase):
|
|||||||
extra_args=base_args + eagle3_args,
|
extra_args=base_args + eagle3_args,
|
||||||
variant="TP8+MTP",
|
variant="TP8+MTP",
|
||||||
),
|
),
|
||||||
|
ModelLaunchSettings(
|
||||||
|
KIMI_K25_MODEL_PATH,
|
||||||
|
tp_size=8,
|
||||||
|
extra_args=base_args + dp_attn_args + eagle3_args,
|
||||||
|
variant="TP8+DP8+MTP",
|
||||||
|
),
|
||||||
]
|
]
|
||||||
|
|
||||||
run_combined_tests(
|
run_combined_tests(
|
||||||
|
|||||||
Reference in New Issue
Block a user