From 93e53f6e0b2d214ab2c3b1bd3334e4ee3cf25b9b Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Mon, 26 May 2025 13:36:40 +0800
Subject: [PATCH] Logging and minor fixes to two batch overlap and EPLB (#6595)

---
 python/sglang/srt/managers/expert_location.py |  6 ++++--
 python/sglang/srt/two_batch_overlap.py        | 15 ++++++++++++++-
 python/sglang/srt/utils.py                    |  1 +
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/python/sglang/srt/managers/expert_location.py b/python/sglang/srt/managers/expert_location.py
index 274077beb..7391dc8c2 100644
--- a/python/sglang/srt/managers/expert_location.py
+++ b/python/sglang/srt/managers/expert_location.py
@@ -152,8 +152,10 @@ class ExpertLocationMetadata:
 
         return ExpertLocationMetadata._init_raw(
             ep_size=common["ep_size"],
-            physical_to_logical_map=physical_to_logical_map,
-            logical_to_all_physical_map=logical_to_all_physical_map,
+            physical_to_logical_map=physical_to_logical_map.to(server_args.device),
+            logical_to_all_physical_map=logical_to_all_physical_map.to(
+                server_args.device
+            ),
         )
 
     @staticmethod
diff --git a/python/sglang/srt/two_batch_overlap.py b/python/sglang/srt/two_batch_overlap.py
index 79ba76d49..c507c2637 100644
--- a/python/sglang/srt/two_batch_overlap.py
+++ b/python/sglang/srt/two_batch_overlap.py
@@ -1,4 +1,5 @@
 import dataclasses
+import logging
 from typing import TYPE_CHECKING, Dict, List, Optional, Sequence
 
 import torch
@@ -11,11 +12,15 @@ from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.operations import execute_operations, execute_overlapped_operations
 from sglang.srt.operations_strategy import OperationsStrategy
-from sglang.srt.utils import BumpAllocator, DeepEPMode
+from sglang.srt.utils import BumpAllocator, DeepEPMode, get_bool_env_var
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
 
+_tbo_debug = get_bool_env_var("SGLANG_TBO_DEBUG")
+
+logger = logging.getLogger(__name__)
+
 
 # -------------------------------- Compute Basic Info ---------------------------------------
 
@@ -178,6 +183,14 @@ class TboForwardBatchPreparer:
             extend_seq_lens=batch.extend_seq_lens_cpu,
         )
 
+        if _tbo_debug:
+            logger.info(
+                f"TboForwardBatchPreparer.prepare "
+                f"tbo_split_seq_index={batch.tbo_split_seq_index} "
+                f"tbo_split_token_index={tbo_split_token_index} "
+                f"extend_seq_lens={batch.extend_seq_lens_cpu}"
+            )
+
         assert isinstance(batch.attn_backend, TboAttnBackend)
         attn_backend_child_a, attn_backend_child_b = batch.attn_backend.children
 
diff --git a/python/sglang/srt/utils.py b/python/sglang/srt/utils.py
index 884e715fa..0b19272e1 100644
--- a/python/sglang/srt/utils.py
+++ b/python/sglang/srt/utils.py
@@ -44,6 +44,7 @@ from functools import lru_cache
 from importlib.metadata import PackageNotFoundError, version
 from importlib.util import find_spec
 from io import BytesIO
+from json import JSONDecodeError
 from multiprocessing.reduction import ForkingPickler
 from pathlib import Path
 from typing import (