Add comments in response to feedback

Author: bkryu

flashinfer/gemm/gemm_base.py

@@ -1993,12 +1993,15 @@ def _heuristic_func_mm_fp4(
     use_nvfp4: bool = True,
 ):
     r"""
-    Heuristic function for mm_fp4 backend selection. Routes to either cudnn or cutlass, but not trtllm.
+    Heuristic function for mm_fp4 backend selection. Routes to either cudnn or cutlass.
+    Note: trtllm is not considered in the backend selection because it requires a specific
+    input quantization (swizzling/shuffling) that differs from the preparation used
+    for cudnn and cutlass backends.
 
     Logic for which comes first:
     - If cuda version is 12 - use cutlass.
-    - If cuda version is 13 and cudnn version is less than 9.14 - use cutlass.
-    - If cuda version is 13 and cudnn version is 9.14 or greater - use cudnn.
+    - If cuda version is 13 and cudnn version is less than 9.15 - use cutlass.
+    - If cuda version is 13 and cudnn version is 9.15 or greater - use cudnn.
 
     """
     cuda_major, _ = get_cuda_version()

flashinfer/utils.py

@@ -921,6 +921,12 @@ def backend_requirement(
         backends. Should accept the same arguments as the decorated function and return
         True if requirements are met, False otherwise.
         In the case where the kernel function does not have any specific backends, this can be decorated with @supported_compute_capability to specify the function's supported compute capabilities.
+    heuristic_func : callable, optional
+        An optional function that performs heuristic backend selection when backend is "auto". Does not do anything if backend is not "auto".
+        Should accept the same arguments as the decorated function.
+        Should return an ordered list of runnable backends with the most preferred backend first.
+        When decorated function is not autotuned, the first backend in the heuristic list will be run.
+        When decorated function is autotuned, the backends in the heuristic list will be autotuned over to find the best backend.
 
     Returns
     -------