Keep CP compute padding out of sparse MoE
CP shared-KV compute padding creates per-request lane slots, so valid rows are not a simple prefix/suffix mask. DeepEP MoE was still seeing dummy rows and using scalar non-padded semantics, which let padding participate in gate/topk and corrupted cache-hit tiny-extend inference.\n\nThe fix compacts CP-local valid rows before MoE dispatch and restores the compact output back to the compute-padded row layout before downstream layer communication. The local GSM8K investigation ledger is now removed from the tracked tree and ignored so future debug notes stay local.\n\nConstraint: CP shared-KV compute-padding layout must keep downstream communicator shapes stable.\nRejected: Disable bs>1/current reuse/cache-hit fast paths | hides the semantic bug and loses the intended performance path.\nRejected: Use num_token_non_padded for MoE under compute padding | valid rows are interleaved with dummy lane slots, not suffix-padded.\nConfidence: high\nScope-risk: moderate\nDirective: Do not feed compute-padded dummy rows into sparse MoE gate/topk; compact valid rows at the MoE boundary and restore shape afterward.\nTested: python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/models/deepseek_v2.py\nTested: remote focused CP utils tests passed, 4 tests.\nTested: remote GSM8K 50-question smoke accuracy 0.960; 200-question runs accuracy 0.955 and 0.965; full 1319-question run accuracy 0.952.\nNot-tested: Long-running production traffic beyond GSM8K after this commit.
This commit is contained in:
@@ -60,9 +60,11 @@ from sglang.srt.layers.attention.nsa.utils import (
|
||||
cp_collect_last_token_hidden,
|
||||
cp_split_and_rebuild_data,
|
||||
cp_split_and_rebuild_position,
|
||||
restore_cp_local_valid_rows_for_moe,
|
||||
is_nsa_enable_prefill_cp,
|
||||
nsa_use_prefill_cp,
|
||||
prepare_input_dp_with_cp_dsa,
|
||||
select_cp_local_valid_rows_for_cache_write,
|
||||
)
|
||||
from sglang.srt.layers.communicator import (
|
||||
LayerCommunicator,
|
||||
@@ -755,6 +757,22 @@ class DeepseekV2MoE(nn.Module):
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
) -> torch.Tensor:
|
||||
local_compute_hidden_states = None
|
||||
if nsa_use_prefill_cp(forward_batch):
|
||||
plan = getattr(
|
||||
getattr(forward_batch, "nsa_cp_metadata", None),
|
||||
"batch_plan",
|
||||
None,
|
||||
)
|
||||
if plan is not None and bool(
|
||||
getattr(plan, "compute_padding_enabled", False)
|
||||
):
|
||||
local_compute_hidden_states = hidden_states
|
||||
hidden_states = select_cp_local_valid_rows_for_cache_write(
|
||||
forward_batch,
|
||||
hidden_states,
|
||||
)
|
||||
|
||||
shared_output = None
|
||||
sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn
|
||||
sbo_overlap_dispatch_flag = (
|
||||
@@ -959,6 +977,13 @@ class DeepseekV2MoE(nn.Module):
|
||||
):
|
||||
final_hidden_states *= self.routed_scaling_factor
|
||||
|
||||
if local_compute_hidden_states is not None:
|
||||
final_hidden_states = restore_cp_local_valid_rows_for_moe(
|
||||
forward_batch,
|
||||
final_hidden_states,
|
||||
local_compute_hidden_states,
|
||||
)
|
||||
|
||||
return final_hidden_states
|
||||
|
||||
def _forward_shared_experts(
|
||||
|
||||
Reference in New Issue
Block a user