Fix NSA CP positions mismatch in eagle NextN model (#19367)

This commit is contained in:
Alison Shao
2026-02-25 20:14:33 -08:00
committed by GitHub
parent 119c91cb8b
commit 0fd44ff342

View File

@@ -29,6 +29,7 @@ from sglang.srt.layers.attention.nsa.utils import (
can_cp_split,
cp_all_gather_rerange_output,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
prepare_input_dp_with_cp_dsa,
@@ -160,6 +161,7 @@ class DeepseekModelNextN(nn.Module):
if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp):
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
residual = None
with get_global_expert_distribution_recorder().disable_this_region():
hidden_states, residual = self.decoder(