Fix stale forward_metadata leak in DP attn unpadded idle batch (#26313)

This commit is contained in:
Ke Bao
2026-05-25 16:04:00 -07:00
committed by GitHub
parent 2b9dd9c8b3
commit e7b12fe6fa
@@ -3101,11 +3101,17 @@ 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. DSA Indexer).
# In DP Attention, IDLE batches may be padded (batch_size > 0) for MLP
# sync. Reinit metadata for the padded case so attention kernels see
# the right batch_size (e.g. DSA Indexer). For the unpadded case
# (batch_size == 0) explicitly drop any stale forward_metadata left
# over from the previous forward — without this, attention layers
# called from the idle path can re-read a prior batch's req_pool
# indices and trigger SWA mapping use-after-free.
if forward_batch.batch_size > 0:
self.attn_backend.init_forward_metadata(forward_batch)
else:
self.attn_backend.forward_metadata = None
kwargs = {}
if self.support_pp: