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
deeplearning_operators/matmul_dequant
deeplearning_operators/flash_attention
deeplearning_operators/flash_linear_attention
deeplearning_operators/deepseek_mla
deeplearning_operators/convolution
deeplearning_operators/tmac_gpu
:::
......
......@@ -9,6 +9,7 @@ import itertools
torch.random.manual_seed(0)
def get_configs():
block_N = [64, 128]
block_H = [64]
......
......@@ -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()
......
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