diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index abc308302b7a..1022d471e675 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1357,9 +1357,13 @@ def load_model(self): f"mem usage={self.weight_load_mem_usage:.2f} GB." ) if self.server_args.debug_tensor_dump_output_folder is not None: + dump_folder = self.server_args.debug_tensor_dump_output_folder + if self.spec_algorithm.is_eagle(): + role = "draft" if self.is_draft_worker else "target" + dump_folder = os.path.join(dump_folder, role) register_forward_hook_for_model( self.model, - self.server_args.debug_tensor_dump_output_folder, + dump_folder, self.server_args.debug_tensor_dump_layers, self.tp_size, self.tp_rank, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 22c676ea5192..dffe0532b26f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -6122,7 +6122,11 @@ def add_cli_args(parser: argparse.ArgumentParser): "--debug-tensor-dump-output-folder", type=str, default=ServerArgs.debug_tensor_dump_output_folder, - help="The output folder for dumping tensors.", + help=( + "The output folder for dumping tensors. " + "In Eagle mode, tensor outputs from draft and target models " + "are stored in separate subdirectories ('draft' and 'target')." + ), ) parser.add_argument( "--debug-tensor-dump-layers",