Minor update before v1.0.4 (#152)

* update ci

* update ci

* update ci

* minor polished interface

* updated gitignore

* updated pre-commit-config

---------

Co-authored-by: Strivin0311 <hyp@smail.nju.edu.cn>

Author: littsk

.github/workflows/sphinx.yaml

@@ -7,7 +7,8 @@ permissions:
   contents: write
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on:
+      group: default
     container:
       image: nvcr.io/nvidia/pytorch:25.05-py3
     steps:

.gitignore

@@ -6,7 +6,7 @@ magi_attention/flex_flash_attn*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
-*.py[cod]
+*.py[codz]
 *$py.class
 
 # C extensions
@@ -52,7 +52,7 @@ htmlcov/
 nosetests.xml
 coverage.xml
 *.cover
-*.py,cover
+*.py.cover
 .hypothesis/
 .pytest_cache/
 cover/
@@ -103,20 +103,35 @@ ipython_config.py
 #   install all needed dependencies.
 #Pipfile.lock
 
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
 # poetry
 #   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
 #   This is especially recommended for binary packages to ensure reproducibility, and is more
 #   commonly ignored for libraries.
 #   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
 #poetry.lock
+#poetry.toml
 
 # pdm
 #   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
 #pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/#use-with-ide
-.pdm.toml
+pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
 
 # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
 __pypackages__/
@@ -130,6 +145,7 @@ celerybeat.pid
 
 # Environments
 .env
+.envrc
 .venv
 env/
 venv/
@@ -168,5 +184,33 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-
-third_party/TransformerEngine/
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/

.pre-commit-config.yaml

@@ -21,10 +21,11 @@ repos:
         files: \.(txt|md|yaml|c|cc|cxx|cpp|cu|cuh|h|hpp|hxx|proto|py|pyi|sh)$
     -   id: csrc_code_formatter
         name: check for csrc code format
-        entry: clang-format-21
+        entry: bash scripts/run_csrc_code_formatter.sh
         language: system
         files: ^magi_attention/csrc/.*\.(h|cuh|hpp|cpp|cu)$
         args: ['-i']
+        require_serial: true
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.4.0
     hooks:

magi_attention/api/magi_attn_interface.py

@@ -79,6 +79,7 @@ def magi_attn_varlen_key(
             which represents ``[window_size_left, window_size_right]``. The parameter is effective only
             when ``causal`` is ``False``; when ``causal`` is ``True``, it is required to be ``(-1, -1)``.
             Defaults to be ``(-1, -1)``.
+
         dist_attn_config (DistAttnConfig): dist attn config.
 
     Returns:
@@ -208,6 +209,7 @@ def magi_attn_varlen_dispatch(
             which represents ``[window_size_left, window_size_right]``. The parameter is effective only
             when ``causal`` is ``False``; when ``causal`` is ``True``, it is required to be ``(-1, -1)``.
             Defaults to be ``(-1, -1)``.
+
         dist_attn_config (DistAttnConfig): dist attn config.
 
     Returns:
@@ -306,7 +308,6 @@ def magi_attn_flex_key(
     calculate DistAttnRuntimeKey and generate the corr. inner DistAttnRuntimeMgr.
 
     Args:
-        x (torch.Tensor): input tensor
         q_ranges (AttnRanges): the global query ranges
         k_ranges (AttnRanges): the global key ranges
         attn_mask_type (str | AttnMaskType | list[str | AttnMaskType]):
@@ -518,6 +519,7 @@ def magi_attn_flex_dispatch(
 
     Args:
         x (torch.Tensor): input tensor
+
         q_ranges (AttnRanges): the global query ranges
         k_ranges (AttnRanges): the global key ranges
         attn_mask_type (str | AttnMaskType | list[str | AttnMaskType]):
@@ -817,12 +819,14 @@ def make_varlen_key_for_new_mask_after_dispatch(
     Args:
         cu_seqlens_q (torch.Tensor): Cumulative sequence lengths for queries.
         cu_seqlens_k (torch.Tensor): Cumulative sequence lengths for keys.
+
         key_for_dispatch (DistAttnRuntimeKey): the key used for dispatch
         causal (bool, optional): whether the varlen attention mask is causal. Defaults to ``False``.
         window_size (tuple[int, int], optional): window_size of sliding window mask
             which represents ``[window_size_left, window_size_right]``. The parameter is effective only
             when ``causal`` is ``False``; when ``causal`` is ``True``, it is required to be ``(-1, -1)``.
             Defaults to be ``(-1, -1)``.
+
         dist_attn_config (DistAttnConfig, optional): the optional new dist attn config,
 
             NOTE: if not provided, we will use the same config as the ``key_for_dispatch``,
@@ -959,7 +963,9 @@ def make_flex_key_for_new_mask_after_dispatch(
         attn_mask_type (str | AttnMaskType | list[str | AttnMaskType]):
             the global attn mask type (list)
             represented by str or enum ``AttnMaskType`` or their mixed combination
+
         key_for_dispatch (DistAttnRuntimeKey): the key used for dispatch
+
         dist_attn_config (DistAttnConfig, optional): the optional new dist attn config,
 
             NOTE: if not provided, we will use the same config as the ``key_for_dispatch``,

scripts/run_csrc_code_formatter.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+
+# Copyright (c) 2025 SandAI. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Define the path to the clang-format executable.
+# You might need to adjust the version number (e.g., clang-format-15, clang-format-17)
+# or just "clang-format" if it's a generic link.
+CLANG_FORMAT_BIN="clang-format-21"
+
+INSTALL_CMD="bash scripts/install_clang_format.sh"
+
+# Check if the clang-format executable is available in the system's PATH.
+# 'command -v' attempts to locate the command.
+# '&> /dev/null' redirects both stdout and stderr to null, suppressing output during the check.
+if ! command -v "$CLANG_FORMAT_BIN" &> /dev/null
+then
+    echo "Error: '$CLANG_FORMAT_BIN' command not found."
+    echo "This pre-commit hook requires '$CLANG_FORMAT_BIN' to be installed."
+    echo "We recommend running '$INSTALL_CMD' to install it."
+    exit 1 # Exit with a non-zero status code, signaling failure to pre-commit.
+fi
+
+# If clang-format is found, execute it with all arguments passed to this script.
+# 'exec' replaces the current shell process with the clang-format process,
+# meaning clang-format's exit status will become this script's exit status.
+# '"$@"' expands to all positional parameters passed to this script, preserving whitespace.
+exec "$CLANG_FORMAT_BIN" "$@"