@@ -29,9 +29,10 @@ class SchedulerMetricsLogger:
2929
3030 DEFAULT_DECODE_LOG_INTERVAL = 5
3131
32- def __init__ (self , enabled : bool = True , dp_rank : int = 0 ) -> None :
32+ def __init__ (self , enabled : bool = True , dp_rank : int = 0 , splitwise_role : str = "mixed" ) -> None :
3333 self .enabled = enabled
3434 self .dp_rank = dp_rank
35+ self .splitwise_role = splitwise_role
3536 decode_log_interval = envs .FD_CONSOLE_DECODE_LOG_INTERVAL
3637 if decode_log_interval <= 0 :
3738 decode_log_interval = self .DEFAULT_DECODE_LOG_INTERVAL
@@ -65,8 +66,9 @@ def on_decode_tokens(self, num_tokens: int) -> None:
6566 with self ._lock :
6667 self ._decode_tokens_since_last += num_tokens
6768
68- def log_prefill_batch (
69+ def _log_prefill_like_batch (
6970 self ,
71+ batch_name : str ,
7072 prefill_reqs : Iterable ,
7173 running_cnt : int ,
7274 queue_cnt : int ,
@@ -91,8 +93,9 @@ def log_prefill_batch(
9193 cached_tokens += getattr (req , "num_cached_tokens" , 0 ) or 0
9294
9395 msg = (
94- "Prefill batch , "
96+ f" { batch_name } , "
9597 f"dp_rank: { self .dp_rank } , "
98+ f"splitwise_role: { self .splitwise_role } , "
9699 f"#new-seq: { len (prefill_reqs )} , "
97100 f"#new-token: { new_tokens } , "
98101 f"#cached-token: { cached_tokens } , "
@@ -104,6 +107,48 @@ def log_prefill_batch(
104107 )
105108 self ._logger .info (msg )
106109
110+ def log_prefill_batch (
111+ self ,
112+ prefill_reqs : Iterable ,
113+ running_cnt : int ,
114+ queue_cnt : int ,
115+ tokens_used : int ,
116+ token_usage : float ,
117+ free_blocks : int = 0 ,
118+ evictable_blocks : int = 0 ,
119+ ) -> None :
120+ self ._log_prefill_like_batch (
121+ batch_name = "Prefill batch" ,
122+ prefill_reqs = prefill_reqs ,
123+ running_cnt = running_cnt ,
124+ queue_cnt = queue_cnt ,
125+ tokens_used = tokens_used ,
126+ token_usage = token_usage ,
127+ free_blocks = free_blocks ,
128+ evictable_blocks = evictable_blocks ,
129+ )
130+
131+ def log_decode_bootstrap_batch (
132+ self ,
133+ prefill_reqs : Iterable ,
134+ running_cnt : int ,
135+ queue_cnt : int ,
136+ tokens_used : int ,
137+ token_usage : float ,
138+ free_blocks : int = 0 ,
139+ evictable_blocks : int = 0 ,
140+ ) -> None :
141+ self ._log_prefill_like_batch (
142+ batch_name = "Decode bootstrap batch from prefill" ,
143+ prefill_reqs = prefill_reqs ,
144+ running_cnt = running_cnt ,
145+ queue_cnt = queue_cnt ,
146+ tokens_used = tokens_used ,
147+ token_usage = token_usage ,
148+ free_blocks = free_blocks ,
149+ evictable_blocks = evictable_blocks ,
150+ )
151+
107152 def log_decode_batch (
108153 self ,
109154 running_cnt : int ,
@@ -132,6 +177,7 @@ def log_decode_batch(
132177 msg = (
133178 "Decode batch, "
134179 f"dp_rank: { self .dp_rank } , "
180+ f"splitwise_role: { self .splitwise_role } , "
135181 f"#running-req: { running_cnt } , "
136182 f"#token: { tokens_used } , "
137183 f"token usage: { token_usage :.2f} , "
0 commit comments