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

Adding user defined hooks support (#13217)

Browse Source
This commit is contained in:
Carlo Mussolini
2025-11-17 23:07:37 +08:00
committed by GitHub
parent 15db5497d3
commit df56139226
7 changed files with 554 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/forward_hooks.md
+297
View File
@@ -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
+5
View File
@@ -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
+82
View File
@@ -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
+4
View File
@@ -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
+13 -1
View File
@@ -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
+1
View File
@@ -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
+152
View File
@@ -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()