Add docstring and typehint

Signed-off-by: Chang Liu (Enterprise Products) <[email protected]>

Author: chang-l

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -574,7 +574,8 @@ def update_for_spec_dec(self):
 
 
 @maybe_compile(dynamic=True)
-def _scale(weights, q_scale, s):
+def _scale(weights: torch.Tensor, q_scale: torch.Tensor,
+           s: float) -> torch.Tensor:
     return weights * q_scale.squeeze(-1) * s
 
 

tensorrt_llm/_torch/utils.py

@@ -328,6 +328,15 @@ def get_device_uuid(device_idx: int) -> str:
 
 
 def maybe_compile(func=None, **compile_kwargs):
+    """
+    Conditionally compile a function with torch.compile.
+    If is_piecewise_running() is True, the function will be compiled with torch.compile.
+    Args:
+        func: The function to decorate (optional, for direct decoration).
+        **compile_kwargs: Keyword arguments for torch.compile.
+    Returns:
+        The conditionally compiled function..
+    """
 
     def decorator(f):