Fix gathered_buffer issues in tbo (#7531)

b8df43ab · Qiaolin Yu · GitHub · a1c1ebe9 · b8df43ab
Unverified Commit b8df43ab authored Jun 25, 2025 by Qiaolin Yu Committed by GitHub Jun 25, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

python/sglang/srt/two_batch_overlap.py python/sglang/srt/two_batch_overlap.py +4 -1

No files found.
--- a/python/sglang/srt/two_batch_overlap.py
+++ b/python/sglang/srt/two_batch_overlap.py
@@ -346,7 +346,10 @@ class TboForwardBatchPreparer:
        )

        # TODO improve, e.g. unify w/ `init_raw`
-        if global_server_args_dict["moe_dense_tp_size"] == 1:
+        if (
+            global_server_args_dict["moe_dense_tp_size"] == 1
+            and batch.gathered_buffer is not None
+        ):
            sum_len = end_token_index - start_token_index
            gathered_buffer = torch.zeros(
                (sum_len, batch.gathered_buffer.shape[1]),