Adding user defined hooks support (#13217)

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