[IDLE FORWARD][Indexer] Fix forward_idle bs mismatch issue in DeepseekV3.2's NSAIndexer (#15227)

This commit is contained in:
YAMY
2026-01-10 02:14:30 -08:00
committed by GitHub
parent 3a8b44fe89
commit 2d088b85d9

View File

@@ -2164,6 +2164,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
def forward_idle(
self, forward_batch: ForwardBatch, pp_proxy_tensors=None
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
# In DP Attention, IDLE batches are padded (batch_size > 0) for MLP sync.
# in this case, we need to reinit the forward metadata, otherwise the stale
# metadata causes batch_size mismatch in attention kernel(e.g. NSA Indexer).
if forward_batch.batch_size > 0:
self.attn_backend.init_forward_metadata(forward_batch)
kwargs = {}
if self.support_pp:
kwargs["pp_proxy_tensors"] = pp_proxy_tensors