[profiler] add record_shapes env flag SGLANG_PROFILE_RECORD_SHAPES

for profiler

Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>

Author: zejunchen-zejun

docs/references/environment_variables.md

@@ -94,6 +94,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 (capture shapes) | `true` |
 
 ## Storage & Caching
 

evaluation/README.md

@@ -108,6 +108,7 @@
     ```
     export SGLANG_TORCH_PROFILER_DIR=./
     export SGLANG_PROFILE_WITH_STACK=1
+    export SGLANG_PROFILE_RECORD_SHAPES=1
     <launch the server>
     <launch the client with the additional --profile argument>
     ```

python/sglang/srt/environ.py

@@ -117,6 +117,7 @@ class Envs:
     SGLANG_TEST_RETRACT = 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_GC_LOG = EnvBool(False)
     SGLANG_FORCE_SHUTDOWN = EnvBool(False)

python/sglang/srt/managers/tokenizer_communicator_mixin.py

@@ -310,6 +310,8 @@ async def start_profile(
         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 = False if record_shapes is False or env_record_shapes is False else True
         req = ProfileReq(
             type=ProfileReqType.START_PROFILE,
             output_dir=output_dir,