chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

feat: Add Non-intrusive Tensor Dumping for Model Inference (#10566)

Browse Source
This commit is contained in:
Yuhong Guo
2025-10-31 12:04:48 +08:00
committed by GitHub
parent 34c286b811
commit ab95d35fcb
6 changed files with 272 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/sglang/srt/layers/logits_processor.py
View File

@@ -38,7 +38,6 @@ from sglang.srt.layers.dp_attention import (
get_dp_device,
get_dp_dtype,
get_dp_hidden_size,
get_local_attention_dp_size,
)
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.model_executor.forward_batch_info import (
@@ -47,7 +46,7 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardMode,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import dump_to_file, is_npu, use_intel_amx_backend
from sglang.srt.utils import is_npu, use_intel_amx_backend
logger = logging.getLogger(__name__)
@@ -252,10 +251,6 @@ class LogitsProcessor(nn.Module):
):
self.final_logit_softcapping = None
self.debug_tensor_dump_output_folder = (
get_global_server_args().debug_tensor_dump_output_folder
)
def compute_logprobs_for_multi_item_scoring(
self,
input_ids,
@@ -463,14 +458,6 @@ class LogitsProcessor(nn.Module):
logits[sample_indices] if sample_indices is not None else logits
)
if self.debug_tensor_dump_output_folder:
assert (
not self.do_tensor_parallel_all_gather
or get_local_attention_dp_size() == 1
), "dp attention + sharded lm_head doesn't support full logits"
full_logits = self._get_logits(hidden_states, lm_head, logits_metadata)
dump_to_file(self.debug_tensor_dump_output_folder, "logits", full_logits)
hidden_states_to_store: Optional[torch.Tensor] = None
if logits_metadata.capture_hidden_mode.need_capture():
if logits_metadata.capture_hidden_mode.is_full():