From fae84ac0f9a4668a4087a68f245b5987e0e93c92 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sat, 25 Jul 2026 16:36:15 -0700 Subject: [PATCH] Fix token count localization for replicated attention-TP forwards (#32411) Co-authored-by: Xingyu Liu <38244988+charlotte12l@users.noreply.github.com> --- .../srt/model_executor/forward_batch_info.py | 25 +++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index d6d9bd08d..277d95038 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -32,7 +32,7 @@ import warnings from dataclasses import dataclass from enum import IntEnum, auto from functools import total_ordering -from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set, Tuple, Union import torch @@ -218,8 +218,29 @@ class CaptureHiddenMode(IntEnum): return self.value < other.value +# Predicate for whether a forward's sequence is sharded across the attn-TP group +# (vs. replicated on every rank). Injected at init; unset defaults to sharded. +_attn_tp_sequence_sharded_predicate: Optional[Callable[[int], bool]] = None + + +def register_attn_tp_sequence_sharded_predicate( + predicate: Callable[[int], bool], +) -> None: + """Register the predicate for whether a forward is sharded across attn-TP.""" + global _attn_tp_sequence_sharded_predicate + _attn_tp_sequence_sharded_predicate = predicate + + def _attn_tp_local_shard_bounds(num_tokens_per_dp: int) -> Tuple[int, int]: - """(tokens_per_rank, rank_offset) of this attn-TP rank's contiguous shard.""" + """(tokens_per_rank, rank_offset) of this attn-TP rank's slice of the sequence. + + A replicated (non-sharded) forward puts the whole sequence on every rank, so + the slice is the full range with no offset; localizing it as a shard would + drop real tokens on non-zero ranks. + """ + predicate = _attn_tp_sequence_sharded_predicate + if predicate is not None and not predicate(num_tokens_per_dp): + return num_tokens_per_dp, 0 parallel = get_parallel() tokens_per_rank = num_tokens_per_dp // parallel.attn_tp_size return tokens_per_rank, tokens_per_rank * parallel.attn_tp_rank