[Bugfix] Move extract_layer_index back inside is_v32 guard (#46199)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
891cc4b9c5
commit
77148992cf
@@ -1001,24 +1001,30 @@ class DeepseekV2MLAAttention(nn.Module):
|
||||
# IndexCache config
|
||||
# Refer: https://arxiv.org/abs/2603.12201 for more details.
|
||||
_skip_topk = False
|
||||
_index_topk_freq = getattr(config, "index_topk_freq", 1)
|
||||
_index_topk_pattern = getattr(config, "index_topk_pattern", None)
|
||||
_index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
|
||||
layer_id = extract_layer_index(prefix)
|
||||
is_mtp_layer = False
|
||||
if self.is_v32:
|
||||
_index_topk_freq = getattr(config, "index_topk_freq", 1)
|
||||
_index_topk_pattern = getattr(config, "index_topk_pattern", None)
|
||||
_index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
|
||||
layer_id = extract_layer_index(prefix)
|
||||
|
||||
if _index_topk_pattern is None:
|
||||
_skip_topk = (
|
||||
max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq != 0
|
||||
if _index_topk_pattern is None:
|
||||
_skip_topk = (
|
||||
max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq
|
||||
!= 0
|
||||
)
|
||||
elif 0 <= layer_id < len(_index_topk_pattern):
|
||||
_skip_topk = _index_topk_pattern[layer_id] == "S"
|
||||
|
||||
# The skip pattern only governs backbone layers. MTP/nextn
|
||||
# layers (layer_id >= num_hidden_layers) always build a full
|
||||
# indexer: they compute indices at draft step 0 and toggle
|
||||
# at runtime via set_skip_topk
|
||||
# (index_share_for_mtp_iteration).
|
||||
_num_hidden_layers = getattr(config, "num_hidden_layers", None)
|
||||
is_mtp_layer = (
|
||||
_num_hidden_layers is not None and layer_id >= _num_hidden_layers
|
||||
)
|
||||
elif 0 <= layer_id < len(_index_topk_pattern):
|
||||
_skip_topk = _index_topk_pattern[layer_id] == "S"
|
||||
|
||||
# The skip pattern only governs backbone layers. MTP/nextn layers
|
||||
# (layer_id >= num_hidden_layers) always build a full indexer: they
|
||||
# compute indices at draft step 0 and toggle at runtime via
|
||||
# set_skip_topk (index_share_for_mtp_iteration).
|
||||
_num_hidden_layers = getattr(config, "num_hidden_layers", None)
|
||||
is_mtp_layer = _num_hidden_layers is not None and layer_id >= _num_hidden_layers
|
||||
|
||||
if self.is_v32 and (not _skip_topk or is_mtp_layer):
|
||||
self.indexer_rope_emb = get_rope(
|
||||
|
||||
Reference in New Issue
Block a user