jundot--omlx
wehub-resource-sync
52 行
1.7 KiB
Python
52 行
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Target-phase planning for SpecPrefill sparse prefill."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SpecPrefillTargetPlan:
|
|
"""Target-prefill inputs derived from selected conversation indices."""
|
|
|
|
system_token_count: int
|
|
conversation_tokens: Sequence[int]
|
|
conversation_token_count: int
|
|
generation_kickoff_index: int
|
|
remove_kickoff_index: bool
|
|
sparse_selected_token_count: int
|
|
total_tracker_prefill_count: int
|
|
position_offset: int
|
|
|
|
|
|
def plan_specprefill_target(
|
|
*,
|
|
all_tokens: Sequence[int],
|
|
system_token_count: int,
|
|
selected_indices: Sequence[int],
|
|
position_offset: int,
|
|
) -> SpecPrefillTargetPlan:
|
|
"""Derive target-prefill inputs without touching MLX state."""
|
|
conversation_tokens = all_tokens[system_token_count:]
|
|
conversation_token_count = len(conversation_tokens)
|
|
generation_kickoff_index = conversation_token_count - 1
|
|
|
|
# BatchGenerator processes the generation-kickoff token separately.
|
|
remove_kickoff_index = generation_kickoff_index in selected_indices
|
|
sparse_selected_token_count = len(selected_indices) - int(remove_kickoff_index)
|
|
|
|
return SpecPrefillTargetPlan(
|
|
system_token_count=system_token_count,
|
|
conversation_tokens=conversation_tokens,
|
|
conversation_token_count=conversation_token_count,
|
|
generation_kickoff_index=generation_kickoff_index,
|
|
remove_kickoff_index=remove_kickoff_index,
|
|
sparse_selected_token_count=sparse_selected_token_count,
|
|
total_tracker_prefill_count=(
|
|
system_token_count + sparse_selected_token_count
|
|
),
|
|
position_offset=position_offset,
|
|
)
|