Commit aa85ddc7 authored by Lei Wang's avatar Lei Wang Committed by LeiWang1999
Browse files

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

* Add deepseek_mla to documentation index (#380)

* lint fix
parent 35e9b47c
...@@ -40,6 +40,7 @@ deeplearning_operators/matmul ...@@ -40,6 +40,7 @@ deeplearning_operators/matmul
deeplearning_operators/matmul_dequant deeplearning_operators/matmul_dequant
deeplearning_operators/flash_attention deeplearning_operators/flash_attention
deeplearning_operators/flash_linear_attention deeplearning_operators/flash_linear_attention
deeplearning_operators/deepseek_mla
deeplearning_operators/convolution deeplearning_operators/convolution
deeplearning_operators/tmac_gpu deeplearning_operators/tmac_gpu
::: :::
......
...@@ -9,6 +9,7 @@ import itertools ...@@ -9,6 +9,7 @@ import itertools
torch.random.manual_seed(0) torch.random.manual_seed(0)
def get_configs(): def get_configs():
block_N = [64, 128] block_N = [64, 128]
block_H = [64] block_H = [64]
......
...@@ -45,14 +45,18 @@ ROCM_HOME = envs.ROCM_HOME ...@@ -45,14 +45,18 @@ ROCM_HOME = envs.ROCM_HOME
# Check if both CUDA and ROCM are enabled # Check if both CUDA and ROCM are enabled
if USE_ROCM and not ROCM_HOME: 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: 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 # Ensure one of CUDA or ROCM is available
if not (CUDA_HOME or ROCM_HOME): 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 # TileLang only supports Linux platform
assert sys.platform.startswith("linux"), "TileLang only supports Linux platform (including WSL)." assert sys.platform.startswith("linux"), "TileLang only supports Linux platform (including WSL)."
...@@ -105,7 +109,8 @@ def get_rocm_version(): ...@@ -105,7 +109,8 @@ def get_rocm_version():
return LooseVersion(match.group(1)) return LooseVersion(match.group(1))
else: else:
rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") 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): if os.path.exists(rocm_version_file):
with open(rocm_version_file, "r") as f: with open(rocm_version_file, "r") as f:
content = f.read() content = f.read()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment