[Doc] Add deepseek_mla to documentation index (#380)

* Add deepseek_mla to documentation index (#380)

* lint fix

docs/index.md

@@ -40,6 +40,7 @@ deeplearning_operators/matmul
 deeplearning_operators/matmul_dequant
 deeplearning_operators/flash_attention
 deeplearning_operators/flash_linear_attention
+deeplearning_operators/deepseek_mla
 deeplearning_operators/convolution
 deeplearning_operators/tmac_gpu
 :::

examples/flash_decoding/example_gqa_decode.py

@@ -9,6 +9,7 @@ import itertools
 
 torch.random.manual_seed(0)
 
+
 def get_configs():
     block_N = [64, 128]
     block_H = [64]

setup.py

@@ -45,14 +45,18 @@ ROCM_HOME = envs.ROCM_HOME
 
 # Check if both CUDA and ROCM are enabled
 if USE_ROCM and not ROCM_HOME:
-    raise ValueError("ROCM support is enabled (USE_ROCM=True) but ROCM_HOME is not set or detected.")
+    raise ValueError(
+        "ROCM support is enabled (USE_ROCM=True) but ROCM_HOME is not set or detected.")
 
 if not USE_ROCM and not CUDA_HOME:
-    raise ValueError("CUDA support is enabled by default (USE_ROCM=False) but CUDA_HOME is not set or detected.")
+    raise ValueError(
+        "CUDA support is enabled by default (USE_ROCM=False) but CUDA_HOME is not set or detected.")
 
 # Ensure one of CUDA or ROCM is available
 if not (CUDA_HOME or ROCM_HOME):
-    raise ValueError("Failed to automatically detect CUDA or ROCM installation. Please set the CUDA_HOME or ROCM_HOME environment variable manually (e.g., export CUDA_HOME=/usr/local/cuda or export ROCM_HOME=/opt/rocm).")
+    raise ValueError(
+        "Failed to automatically detect CUDA or ROCM installation. Please set the CUDA_HOME or ROCM_HOME environment variable manually (e.g., export CUDA_HOME=/usr/local/cuda or export ROCM_HOME=/opt/rocm)."
+    )
 
 # TileLang only supports Linux platform
 assert sys.platform.startswith("linux"), "TileLang only supports Linux platform (including WSL)."
@@ -105,7 +109,8 @@ def get_rocm_version():
         return LooseVersion(match.group(1))
     else:
         rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
-        rocm_version_file = os.path.join(rocm_path, "lib", "cmake", "rocm", "rocm-config-version.cmake")
+        rocm_version_file = os.path.join(rocm_path, "lib", "cmake", "rocm",
+                                         "rocm-config-version.cmake")
         if os.path.exists(rocm_version_file):
             with open(rocm_version_file, "r") as f:
                 content = f.read()

tilelang/jit/adapter/cython/adapter.py

@@ -110,7 +110,7 @@ with open(cython_wrapper_path, "r") as f:
         try:
             with open(md5_path, "w") as f:
                 f.write(code_hash)
-            
+
             # compile the cython_wrapper.pyx file into .cpp
             cython = get_cython_compiler()
             if cython is None:
@@ -120,7 +120,7 @@ with open(cython_wrapper_path, "r") as f:
             cc = get_cplus_compiler()
             command = f"{cc} -shared -pthread -fPIC -fwrapv -O2 -Wall -fno-strict-aliasing -I{python_include_path} {source_path} -o {temp_path}"
             os.system(command)
-            
+
             # rename the temp file to the library file
             temp_path.rename(library_path)
         except Exception as e: