update - move cutlass_groupgemm position

Author: alibaba-miji

rtp_llm/models_py/BUILD

@@ -25,7 +25,7 @@ requirement(flashinfer)
 
 filegroup(
     name = "cutlass_moe_config",
-    srcs = glob(["configs/cutlass_groupgemm/*"]),
+    srcs = glob(["kernels/cuda/fp8_kernel/cutlass_groupgemm/*"]),
     visibility = ["//visibility:public"],
 )
 

rtp_llm/models_py/kernels/cuda/fp8_kernel/__init__.py

@@ -0,0 +1,19 @@
+from .get_best_config import load_all_configs
+
+# load all configs once at import time
+load_all_configs()
+from .fp8_kernel import (
+    cutlass_moe_mm_fp8_scaled,
+    get_best_config_swap_ab,
+    scaled_fp8_per_tensor_quant,
+    scaled_fp8_per_token_quant,
+    sgl_per_token_group_quant_fp8,
+)
+
+__all__ = [
+    "sgl_per_token_group_quant_fp8",
+    "scaled_fp8_per_tensor_quant",
+    "scaled_fp8_per_token_quant",
+    "cutlass_moe_mm_fp8_scaled",
+    "get_best_config_swap_ab",
+]

…=5120-K=6144-device_name=NVIDIA_H20.json …=5120-K=6144-device_name=NVIDIA_H20.jsonrtp_llm/models_py/configs/cutlass_groupgemm/E=20-N=5120-K=6144-device_name=NVIDIA_H20.json renamed to rtp_llm/models_py/kernels/cuda/fp8_kernel/cutlass_groupgemm/E=20-N=5120-K=6144-device_name=NVIDIA_H20.json rtp_llm/models_py/configs/cutlass_groupgemm/E=20-N=5120-K=6144-device_name=NVIDIA_H20.json renamed to rtp_llm/models_py/kernels/cuda/fp8_kernel/cutlass_groupgemm/E=20-N=5120-K=6144-device_name=NVIDIA_H20.json

@@ -7,7 +7,9 @@
 
 import torch
 
-from rtp_llm.models_py.configs.get_best_config import get_cutlass_groupgemm_best_config
+from rtp_llm.models_py.kernels.cuda.fp8_kernel.get_best_config import (
+    get_cutlass_groupgemm_best_config,
+)
 from rtp_llm.models_py.utils.arch import is_cuda
 from rtp_llm.models_py.utils.math import align
 
@@ -20,7 +22,7 @@
         per_token_quant_fp8,
     )
 else:
-    logging.warning("can't import from rtp_llm_ops, only support cuda!")
+    logging.info("skip import fp8 quant from rtp_llm_ops for non cuda platform")
 
 logger = logging.getLogger(__name__)
 

…=6144-K=2560-device_name=NVIDIA_H20.json …=6144-K=2560-device_name=NVIDIA_H20.jsonrtp_llm/models_py/configs/cutlass_groupgemm/E=20-N=6144-K=2560-device_name=NVIDIA_H20.json renamed to rtp_llm/models_py/kernels/cuda/fp8_kernel/cutlass_groupgemm/E=20-N=6144-K=2560-device_name=NVIDIA_H20.json rtp_llm/models_py/configs/cutlass_groupgemm/E=20-N=6144-K=2560-device_name=NVIDIA_H20.json renamed to rtp_llm/models_py/kernels/cuda/fp8_kernel/cutlass_groupgemm/E=20-N=6144-K=2560-device_name=NVIDIA_H20.json

@@ -13,7 +13,7 @@
 _CUTLASS_GROUPGEMM_CONFIG_MAP: Dict[Tuple[int, int, int, str], Dict] = {}
 
 
-def _load_all_configs():
+def load_all_configs():
     """Load all cutlass groupgemm config files into the global map."""
     if _CUTLASS_GROUPGEMM_CONFIG_MAP:
         # Already loaded
@@ -23,38 +23,58 @@ def _load_all_configs():
 
     # Load open source config directory
     opensource_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), op_name)
-    if os.path.exists(opensource_dir):
-        pattern = os.path.join(opensource_dir, "E=*-N=*-K=*-device_name=*.json")
-        for config_file in glob.glob(pattern):
-            filename = os.path.basename(config_file)
-            try:
-                # Parse filename: E={E}-N={N}-K={K}-device_name={device_name}.json
-                parts = filename.replace(".json", "").split("-")
-                E = int(parts[0].split("=")[1])
-                N = int(parts[1].split("=")[1])
-                K = int(parts[2].split("=")[1])
-                device_name = parts[3].split("=")[1]
-
-                # Load config
-                with open(config_file) as f:
-                    config_data = json.load(f)
-                    # Convert string keys to int
-                    config_data = {int(key): val for key, val in config_data.items()}
-
-                # Store in global map
-                key = (E, N, K, device_name)
-                _CUTLASS_GROUPGEMM_CONFIG_MAP[key] = config_data
-                logging.debug(f"Loaded config from {config_file}")
-            except Exception as e:
-                logging.warning(f"Failed to load config from {config_file}: {e}")
+
+    # Try to get internal source config directory
+    # Collect all config directories to load
+    config_dirs = [opensource_dir]
+    try:
+        import internal_source.rtp_llm.models_py.kernels.cuda.fp8_kernel
+
+        internalsource_dir = os.path.join(
+            os.path.dirname(
+                os.path.realpath(
+                    internal_source.rtp_llm.models_py.kernels.cuda.fp8_kernel.__file__
+                )
+            ),
+            op_name,
+        )
+        config_dirs.append(internalsource_dir)
+    except ImportError:
+        logging.info("internal_source not found")
+
+    # Load configs from all directories
+    for config_dir in config_dirs:
+        if os.path.exists(config_dir):
+            logging.info(f"Loading configs from {config_dir}")
+            pattern = os.path.join(config_dir, "E=*-N=*-K=*-device_name=*.json")
+            for config_file in glob.glob(pattern):
+                filename = os.path.basename(config_file)
+                try:
+                    # Parse filename: E={E}-N={N}-K={K}-device_name={device_name}.json
+                    parts = filename.replace(".json", "").split("-")
+                    E = int(parts[0].split("=")[1])
+                    N = int(parts[1].split("=")[1])
+                    K = int(parts[2].split("=")[1])
+                    device_name = parts[3].split("=")[1]
+
+                    # Load config
+                    with open(config_file) as f:
+                        config_data = json.load(f)
+                        # Convert string keys to int
+                        config_data = {
+                            int(key): val for key, val in config_data.items()
+                        }
+
+                    # Store in global map
+                    key = (E, N, K, device_name)
+                    _CUTLASS_GROUPGEMM_CONFIG_MAP[key] = config_data
+                    logging.debug(f"Loaded config from {config_file}")
+                except Exception as e:
+                    logging.warning(f"Failed to load config from {config_file}: {e}")
 
     logging.info(
         f"Loaded {len(_CUTLASS_GROUPGEMM_CONFIG_MAP)} cutlass groupgemm configurations"
     )
-    try:
-        import internal_source.rtp_llm.utils.register_cutlass_configs
-    except:
-        logging.info("internal_source not found")
 
 
 def register_cutlass_groupgemm_config(
@@ -90,9 +110,6 @@ def get_cutlass_groupgemm_best_config(E: int, N: int, K: int) -> Optional[Dict]:
         Configuration dictionary mapping batch sizes to tile configurations,
         or None if no configuration is found.
     """
-    # Load all configs if not already loaded
-    _load_all_configs()
-
     device_name = torch.cuda.get_device_name().replace("-", "_").replace(" ", "_")
     key = (E, N, K, device_name)
 

…=3072-K=4096-device_name=NVIDIA_H20.json …=3072-K=4096-device_name=NVIDIA_H20.jsonrtp_llm/models_py/configs/cutlass_groupgemm/E=32-N=3072-K=4096-device_name=NVIDIA_H20.json renamed to rtp_llm/models_py/kernels/cuda/fp8_kernel/cutlass_groupgemm/E=32-N=3072-K=4096-device_name=NVIDIA_H20.json rtp_llm/models_py/configs/cutlass_groupgemm/E=32-N=3072-K=4096-device_name=NVIDIA_H20.json renamed to rtp_llm/models_py/kernels/cuda/fp8_kernel/cutlass_groupgemm/E=32-N=3072-K=4096-device_name=NVIDIA_H20.json

@@ -8,11 +8,11 @@
 from torch import dtype as _dtype
 from torch.profiler import ProfilerActivity, profile, record_function
 
-from rtp_llm.models_py.kernels.cuda.fp8_kernel import (
+from rtp_llm.models_py.utils.arch import is_hip
+from rtp_llm.ops.compute_ops import (
     per_token_group_quant_fp8,
     per_token_group_quant_int8,
 )
-from rtp_llm.models_py.utils.arch import is_hip
 
 _is_hip = is_hip()