Adding user defined hooks support (#13217)

docs/advanced_features/forward_hooks.md

@@ -0,0 +1,297 @@
+## Model Hooks
+
+SGLang supports attaching PyTorch forward hooks to specific submodules in the loaded model, configured entirely via `server_args` JSON.
+
+This is useful for:
+
+* Logging intermediate activations
+* Debugging model internals
+* Exporting hidden states to external tooling
+
+Hooks are attached once during `ModelRunner.initialize` and run on every forward pass.
+
+---
+
+### Configuration overview
+
+Hooks are configured via a `ServerArgs` field:
+
+```python
+class ServerArgs:
+    ...
+    # For forward hooks
+    hooks: Optional[List[dict[str, Any]]] = None
+````
+
+In JSON form, a minimal configuration looks like:
+
+```jsonc
+{
+  "hooks": [
+    {
+      "name": "outer_linear_hooks",
+      "target_modules": ["outer.0", "outer.1"],
+      "hook_factory": "my_project.hooks:dummy_hook_factory",
+      "config": {
+        "tag": "outer-layer"
+      }
+    }
+  ]
+}
+```
+
+#### Top-level fields
+
+* `hooks` (optional list of objects)
+  Each element is a hook spec describing:
+
+  * Which modules to target
+  * Which Python factory to call
+  * What configuration to pass into that factory
+
+---
+
+### Hook spec schema
+
+Each entry in `hooks` is a JSON object with the following shape:
+
+```jsonc
+{
+  "name": "optional-descriptive-name",
+  "target_modules": ["pattern1", "pattern2", "..."],
+  "hook_factory": "module.submodule:factory_name",
+  "config": {
+    "...": "arbitrary JSON"
+  }
+}
+```
+
+#### `name` (optional)
+
+* Human-readable name for logging.
+* Used only in log messages such as:
+
+  ```text
+  Registered forward hook 'outer_linear_hooks' on outer.0
+  ```
+
+#### `target_modules` (required)
+
+* List of **module name patterns** used to match entries in `model.named_modules()`.
+* Patterns are matched using `fnmatch.fnmatch`, so:
+
+  * `"outer.0"` matches exactly `"outer.0"`.
+  * `"outer.*"` matches `"outer.0"`, `"outer.1"`, `"outer.inner"`, etc.
+  * `"outer.inner.*"` matches children under `outer.inner`.
+
+> If no modules match the given patterns, hook registration does **not** fail.
+> Instead, SGLang logs a warning and continues:
+>
+> ```text
+> No modules matched hook spec 'name' patterns=['...']
+> ```
+
+#### `hook_factory` (required)
+
+* String path to the Python factory function that creates the hook.
+* Supported formats:
+
+  * `"package.module:factory_name"`
+  * `"package.module.submodule.factory_name"`
+
+The path is resolved via:
+
+```python
+def resolve_callable(path: Optional[str]) -> Optional[Callable]:
+    if path is None:
+        return None
+
+    if ":" in path:
+        module_name, fn_name = path.split(":", 1)
+    else:
+        parts = path.split(".")
+        if len(parts) < 2:
+            raise ValueError(
+                f"Invalid hook callable path '{path}'. "
+                "Expected 'module.submodule:factory' or 'module.submodule.factory'."
+            )
+        *mod_parts, fn_name = parts
+        module_name = ".".join(mod_parts)
+
+    module = importlib.import_module(module_name)
+    try:
+        return getattr(module, fn_name)
+    except AttributeError as e:
+        raise AttributeError(
+            f"Module '{module_name}' has no attribute '{fn_name}' "
+            f"(from hook path '{path}')"
+        ) from e
+```
+
+**Failure modes**:
+
+* If the path is malformed (not enough dots and no `:`), a `ValueError` is raised at startup.
+* If the module imports but the attribute is missing, an `AttributeError` is raised with a clear error message.
+* If the hook factory returns `None`, a warning is logged and no hook is registered for that spec (initialization continues).
+
+The first two cause initialization to fail fast with a descriptive error; the last one is non-fatal.
+
+#### `config` (optional)
+
+* Arbitrary JSON object.
+* Passed directly to the hook factory as a Python `dict`.
+* This lets you parameterize hook behavior from config (e.g. tags, log levels, sampling rates, etc.).
+
+---
+
+### Hook lifecycle and behavior
+
+Hooks are registered in `ModelRunner.initialize()`:
+
+```python
+if server_args.hooks:
+    register_hooks(self.model, server_args.hooks)
+```
+
+The actual registration logic is implemented by `register_hooks`:
+
+```python
+def register_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None:
+    """
+    hook_specs is a list of dicts from server_args.hooks.
+    Attaches forward hooks to the matching modules.
+    """
+    name_to_module = dict(model.named_modules())
+
+    for spec in hook_specs:
+        spec_name = spec.get("name", "")
+        target_patterns = spec.get("target_modules", [])
+        if not target_patterns:
+            logger.warning(
+                f"Hook spec '{spec_name}' has no 'target_modules', skipping"
+            )
+            continue
+
+        hook_factory_path = spec.get("hook_factory")
+        if not hook_factory_path:
+            logger.warning(
+                f"Hook spec '{spec_name}' has no 'hook_factory', skipping"
+            )
+            continue
+
+        config = spec.get("config") or {}
+        hook_factory = resolve_callable(hook_factory_path)
+
+        hook = hook_factory(config) if hook_factory else None
+        if hook is None:
+            logger.warning(
+                f"Hook factory '{hook_factory_path}' for spec '{spec_name}' "
+                "returned None, not registering any hook"
+            )
+            continue
+
+        # Resolve patterns like "model.layers.*.mlp"
+        matched = []
+        for name, module in name_to_module.items():
+            if any(fnmatch.fnmatch(name, pattern) for pattern in target_patterns):
+                matched.append((name, module))
+
+        if not matched:
+            logger.warning(
+                f"No modules matched hook spec '{spec_name}' "
+                f"patterns={target_patterns}"
+            )
+            continue
+
+        for module_name, module in matched:
+            if hook:
+                _ = module.register_forward_hook(hook)
+                logger.info(
+                    f"Registered forward hook '{spec_name}' "
+                    f"on {module_name}"
+                )
+```
+
+Key points:
+
+* Hooks are **forward hooks only** (via `module.register_forward_hook`).
+* They are attached once at initialization.
+* Hook handles are currently not stored on `ModelRunner` (they cannot be removed later via this API).
+* Failure to match any modules is non-fatal; a warning is logged instead.
+* If a hook factory returns `None`, a warning is logged and that spec is skipped.
+
+---
+
+### Writing a hook factory
+
+A hook factory is a regular Python function:
+
+* Takes a `config: dict` (from JSON)
+* Returns a forward hook function with signature `(module, inputs, output)`
+
+Example:
+
+```python
+HOOK_CALLS = []
+
+def dummy_hook_factory(config):
+    """Factory that returns a forward hook capturing a tag from config."""
+    tag = config.get("tag", "default")
+
+    def hook(module, inputs, output):
+        HOOK_CALLS.append(
+            {
+                "module_type": type(module).__name__,
+                "tag": tag,
+                "shape": tuple(output.shape),
+            }
+        )
+        return output  # must return output if you don’t want to modify the tensor
+
+    return hook
+```
+
+In JSON:
+
+```jsonc
+{
+  "hooks": [
+    {
+      "name": "capture_outer",
+      "target_modules": ["outer.0", "outer.1"],
+      "hook_factory": "my_project.hooks:dummy_hook_factory",
+      "config": {
+        "tag": "outer"
+      }
+    }
+  ]
+}
+```
+
+This will:
+
+* Resolve `my_project.hooks:dummy_hook_factory` to a Python callable.
+* Call it with `config = {"tag": "outer"}`.
+* Use the returned hook for all modules matching `outer.0` and `outer.1`.
+* Append metadata about each call to `HOOK_CALLS`.
+
+---
+
+### Summary
+
+* Define `hooks` as a list of specs in `ServerArgs` to turn on the feature.
+
+* Each spec:
+
+  * selects modules via `target_modules` (glob patterns over `model.named_modules()`),
+  * points to a hook factory via `hook_factory`,
+  * passes arbitrary `config` into that factory.
+
+* Hook factories are resolved via `resolve_callable`, which supports `module:factory` and `module.submodule.factory`.
+
+* Hooks are standard PyTorch forward hooks, attached once at startup and invoked on every forward pass.
+
+* Misconfiguration is either:
+
+  * **fatal and explicit** (bad path / missing attribute), or
+  * **non-fatal with clear warnings** (no targets matched, or factory returned `None`).

docs/advanced_features/server_arguments.md

@@ -398,6 +398,11 @@ Please consult the documentation below and [server_args.py](https://github.com/s
 | `--enable-attn-tp-input-scattered` | Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent.                                                                                                      | `False`  | bool flag (set to enable) |
 | `--enable-nsa-prefill-context-parallel` | Context parallelism used in the long sequence prefill phase of DeepSeek v3.2 | `False` | bool flag (set to enable) |
 
+## Forward hooks
+| Argument | Description | Defaults | Options |
+| --- | --- | --- | --- |
+| `--hooks` | JSON-formatted list of hook specifications. Each element must include `target_modules` (list of glob patterns matched against `model.named_modules()` names) and `hook_factory` (Python import path to a factory, e.g. `my_package.hooks:make_hook`). An optional `name` field is used for logging, and an optional `config` object is passed as a `dict` to the factory. | `None` | Type: JSON list |
+
 ## Debug tensor dumps
 | Argument | Description | Defaults | Options |
 | --- | --- | --- | --- |

python/sglang/srt/model_executor/hook_manager.py

@@ -0,0 +1,82 @@
+import fnmatch
+import importlib
+import logging
+from typing import Any, Callable, List, Optional
+
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+def register_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None:
+    """
+    hook_specs is a list of dicts from server_args.hooks.
+    Attaches forward hooks to the matching modules.
+    """
+    name_to_module = dict(model.named_modules())
+
+    for spec in hook_specs:
+        spec_name = spec.get("name", "")
+        target_patterns = spec.get("target_modules", [])
+        if not target_patterns:
+            logger.warning(f"Hook spec '{spec_name}' has no 'target_modules', skipping")
+            continue
+
+        hook_factory_path = spec.get("hook_factory")
+        if not hook_factory_path:
+            logger.warning(f"Hook spec '{spec_name}' has no 'hook_factory', skipping")
+            continue
+
+        config = spec.get("config") or {}
+        hook_factory = resolve_callable(hook_factory_path)
+
+        hook = hook_factory(config) if hook_factory else None
+        if hook is None:
+            logger.warning(
+                f"Hook factory '{hook_factory_path}' for spec '{spec_name}' "
+                "returned None, not registering any hook"
+            )
+            continue
+
+        # Resolve patterns like "model.layers.*.mlp"
+        matched = []
+        for name, module in name_to_module.items():
+            if any(fnmatch.fnmatch(name, pattern) for pattern in target_patterns):
+                matched.append((name, module))
+
+        if not matched:
+            logger.warning(
+                f"No modules matched hook spec '{spec_name}' "
+                f"patterns={target_patterns}"
+            )
+            continue
+
+        for module_name, module in matched:
+            _ = module.register_forward_hook(hook)
+            logger.info(f"Registered forward hook '{spec_name}' " f"on {module_name}")
+
+
+def resolve_callable(path: Optional[str]) -> Optional[Callable]:
+    if path is None:
+        return None
+
+    if ":" in path:
+        module_name, fn_name = path.split(":", 1)
+    else:
+        parts = path.split(".")
+        if len(parts) < 2:
+            raise ValueError(
+                f"Invalid hook callable path '{path}'. "
+                "Expected 'module.submodule:factory' or 'module.submodule.factory'."
+            )
+        *mod_parts, fn_name = parts
+        module_name = ".".join(mod_parts)
+
+    module = importlib.import_module(module_name)
+    try:
+        return getattr(module, fn_name)
+    except AttributeError as e:
+        raise AttributeError(
+            f"Module '{module_name}' has no attribute '{fn_name}' "
+            f"(from hook path '{path}')"
+        ) from e

python/sglang/srt/model_executor/model_runner.py

@@ -112,6 +112,7 @@ from sglang.srt.mem_cache.memory_pool import (
 from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
 from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.hook_manager import register_hooks
 from sglang.srt.model_executor.npu_graph_runner import NPUGraphRunner
 from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
     PiecewiseCudaGraphRunner,
@@ -497,6 +498,9 @@ class ModelRunner:
             self.graph_mem_usage = 0
             self.init_attention_backend()
 
+        if server_args.hooks:
+            register_hooks(self.model, server_args.hooks)
+
         # auxiliary hidden capture mode. TODO: expose this to server args?
         if self.spec_algorithm.is_eagle3() and not self.is_draft_worker:
             # load draft config

python/sglang/srt/server_args.py

@@ -22,7 +22,7 @@ import logging
 import os
 import random
 import tempfile
-from typing import Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 import orjson
 
@@ -391,6 +391,7 @@ class ServerArgs:
     speculative_token_map: Optional[str] = None
     speculative_attention_mode: str = "prefill"
     speculative_moe_runner_backend: Optional[str] = None
+
     # For ngram only
     speculative_ngram_min_match_window_size: int = 1
     speculative_ngram_max_match_window_size: int = 12
@@ -577,6 +578,9 @@ class ServerArgs:
     decrypted_config_file: Optional[str] = None
     decrypted_draft_config_file: Optional[str] = None
 
+    # For forward hooks
+    hooks: Optional[List[dict[str, Any]]] = None
+
     def __post_init__(self):
         """
         Orchestrates the handling of various server arguments, ensuring proper configuration and validation.
@@ -3725,6 +3729,14 @@ class ServerArgs:
             help="The path of the decrypted draft config file.",
         )
 
+        # For registering hooks
+        parser.add_argument(
+            "--hooks",
+            type=json_list_type,
+            default=None,
+            help="The hooks to be attached.",
+        )
+
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace):
         args.tp_size = args.tensor_parallel_size

test/srt/run_suite.py

@@ -87,6 +87,7 @@ suites = {
         TestFile("test_mla_flashinfer.py", 302),
         TestFile("test_mla_fp8.py", 93),
         TestFile("test_mla_int8_deepseek_v3.py", 300),
+        TestFile("test_model_hooks.py", 1),
         TestFile("test_modelopt_loader.py", 30),
         TestFile("test_multi_tokenizer.py", 230),
         TestFile("test_ngram_speculative_decoding.py", 290),

test/srt/test_model_hooks.py

@@ -0,0 +1,152 @@
+import argparse
+import json
+import unittest
+
+import torch
+import torch.nn as nn
+
+from sglang.srt.model_executor.hook_manager import register_hooks
+from sglang.srt.server_args import ServerArgs
+from sglang.test.test_utils import CustomTestCase
+
+HOOK_CALLS = []
+
+
+def dummy_hook_factory(config):
+    """Factory that returns a forward hook capturing a tag from config."""
+    tag = config.get("tag", "default")
+
+    def hook(module, inputs, output):
+        HOOK_CALLS.append(
+            {
+                "module_type": type(module).__name__,
+                "tag": tag,
+                "shape": tuple(output.shape),
+            }
+        )
+        return output
+
+    return hook
+
+
+class TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.inner = nn.Sequential(
+            nn.Linear(4, 2),
+            nn.ReLU(),
+        )
+        self.outer = nn.Sequential(
+            nn.Linear(4, 4),
+            nn.ReLU(),
+            self.inner,
+        )
+
+    def forward(self, x):
+        return self.outer(x)
+
+
+class TestAttachHooks(CustomTestCase):
+    """Tests for ModelRunner.register_hooks / resolve_callable integration."""
+
+    def setUp(self):
+        HOOK_CALLS.clear()
+
+    def test_hook_is_attached(self):
+        """Hook from a factory string is registered and fired."""
+        hook_specs = [
+            {
+                "target_modules": ["outer.0", "outer.1"],
+                "hook_factory": "test_model_hooks:dummy_hook_factory",
+                "config": {"tag": "forward-ok"},
+            },
+            {
+                "target_modules": ["inner.*"],
+                "hook_factory": "test_model_hooks:dummy_hook_factory",
+                "config": {"tag": "forward-ok"},
+            },
+        ]
+
+        model = TinyModel()
+        register_hooks(model, hook_specs)
+
+        x = torch.randn(3, 4)
+        _ = model(x)
+
+        self.assertEqual(
+            len(HOOK_CALLS),
+            4,
+            "Forward hook was not called correct number of times",
+        )
+        tags = {call["tag"] for call in HOOK_CALLS}
+        self.assertIn("forward-ok", tags)
+
+    def test_no_matching_modules_does_not_crash(self):
+        """Hook spec with no matching modules should not crash."""
+        model = TinyModel()
+        hook_specs = [
+            {
+                "name": "no_match",
+                "target_modules": ["does_not_exist.*"],
+                "hook_factory": "test_model_hooks:dummy_hook_factory",
+                "config": {"tag": "unused"},
+            }
+        ]
+
+        register_hooks(model, hook_specs)
+
+        x = torch.randn(3, 4)
+        _ = model(x)
+
+        # No hooks should have fired
+        self.assertEqual(len(HOOK_CALLS), 0)
+
+    def test_cli_hooks_reach_model(self):
+        """
+        Ensure that when hooks are provided via CLI, they are parsed into
+        ServerArgs, passed to ModelRunner.register_hooks, and actually
+        run during a forward pass.
+        """
+        parser = argparse.ArgumentParser()
+        ServerArgs.add_cli_args(parser)
+
+        hooks_spec = [
+            {
+                "name": "outer_and_inner_from_cli",
+                "target_modules": ["outer.0", "outer.1", "inner.*"],
+                "hook_factory": "test_model_hooks:dummy_hook_factory",
+                "config": {"tag": "cli-hook"},
+            }
+        ]
+
+        cli_args = [
+            "--model-path",
+            "Qwen/Qwen2-7B-Instruct",  # Dummy value; not used in this test
+            "--hooks",
+            json.dumps(hooks_spec),
+        ]
+
+        args = parser.parse_args(cli_args)
+        server_args = ServerArgs.from_cli_args(args)
+
+        self.assertEqual(server_args.hooks, hooks_spec)
+
+        model = TinyModel()
+        register_hooks(model, server_args.hooks)
+
+        x = torch.randn(3, 4)
+        _ = model(x)
+
+        # We expect hooks on outer.0, outer.1, inner.0, inner.1  => 4 calls
+        self.assertEqual(
+            len(HOOK_CALLS),
+            4,
+            "CLI-configured hooks did not fire expected number of times",
+        )
+
+        tags = {call["tag"] for call in HOOK_CALLS}
+        self.assertEqual(tags, {"cli-hook"})
+
+
+if __name__ == "__main__":
+    unittest.main()