diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index a3f3e41af..631199317 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -99,6 +99,7 @@ SGLang supports various environment variables that can be used to configure its | --- | --- | --- | | `SGLANG_TORCH_PROFILER_DIR` | Directory for PyTorch profiler output | `/tmp` | | `SGLANG_PROFILE_WITH_STACK` | Set `with_stack` option (bool) for PyTorch profiler (capture stack trace) | `true` | +| `SGLANG_PROFILE_RECORD_SHAPES` | Set `record_shapes` option (bool) for PyTorch profiler (record shapes) | `true` | | `SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS` | Config BatchSpanProcessor.schedule_delay_millis if tracing is enabled | `500` | | `SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE` | Config BatchSpanProcessor.max_export_batch_size if tracing is enabled | `64` | diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 98b0c773a..fc0e3ef71 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -136,6 +136,7 @@ class Envs: SGLANG_IS_IN_CI_AMD = EnvBool(False) SGLANG_SET_CPU_AFFINITY = EnvBool(False) SGLANG_PROFILE_WITH_STACK = EnvBool(True) + SGLANG_PROFILE_RECORD_SHAPES = EnvBool(True) SGLANG_RECORD_STEP_TIME = EnvBool(False) SGLANG_FORCE_SHUTDOWN = EnvBool(False) SGLANG_DEBUG_MEMORY_POOL = EnvBool(False) diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index 81f8cb97c..7f49e48c4 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -319,6 +319,10 @@ class TokenizerCommunicatorMixin: self.auto_create_handle_loop() env_with_stack: bool = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true") with_stack = False if with_stack is False or env_with_stack is False else True + env_record_shapes: bool = get_bool_env_var( + "SGLANG_PROFILE_RECORD_SHAPES", "true" + ) + record_shapes = (record_shapes is not False) and env_record_shapes req = ProfileReq( type=ProfileReqType.START_PROFILE, output_dir=output_dir,