update release v0.2.4

Author: zhe-pang

.pre-commit-config.yaml

@@ -66,3 +66,12 @@ repos:
       # Run the formatter.
       - id: ruff-format
         types_or: [ python, pyi ]
+
+  - repo: local
+    hooks:
+      - id: check-api-docs
+        name: check-api-docs
+        entry: python scripts/check_api_docs.py
+        language: system
+        pass_filenames: false
+        files: ^(docs/api_manifest\.yaml|docs/source/api_reference\.rst|docs/source/api/.*\.rst|docs/gdn\.md|mate/__init__\.py|mate/(deep_gemm|flashmla|gdn_decode|gdn_prefill|gemm|hash_topk|hyperconnection|kda|mha_interface|moe_fused_gate|sage_attention_interface)\.py|scripts/check_api_docs\.py|tests/test_check_api_docs\.py)$

README.md

@@ -7,72 +7,90 @@ MATE (**M**USA **A**I **T**ensor **E**ngine) is a centralized library for Genera
 ## Highlights
 
 - High-performance attention and GEMM operators for MUSA
+- Mixed-dtype W4A8 MoE GEMM APIs for ragged and masked MoE paths in `0.2.4`
 - Compatibility wrappers for `flash_attn_3`, `sageattention`, `flash_mla`,
-  `flash_kda`, and `deep_gemm`
+  `flash_kda`, and `deep-gemm`
 - CLI tools for environment checks, configuration inspection, and replay
 
 ## Requirements
 
 | Component | Requirement |
 | --- | --- |
+| Python | `3.10` or later |
 | MUSA Toolkit | `4.3.6` or later |
 | TorchMUSA | `2.7` or later |
 | Architecture | `Pinghu (MP31)` |
 
+The current external delivery source mainly covers `x86_64` and Python `3.10`
+/ `3.12` wheels.
+
 ## Recommended Workflow
 
 For Moore Threads platforms, the normal integration flow is wrapper-first:
 
-1. Install MATE on top of an existing MUSA-enabled `torch` / `torch_musa` stack
-2. Install the wrapper package that matches the Python package surface your framework already expects
-3. Keep the upstream import path and high-level API shape as stable as possible
-4. Use `mate check`, `mate show-config`, `mate env`, logging, and replay if something fails
-5. Use native MATE APIs only when no wrapper matches your workload or the wrapper does not cover the feature you need
+1. Start with an existing MUSA-enabled `torch` / `torch_musa` stack
+2. If a supported wrapper matches your framework surface, install that wrapper from the MUSA wheel source
+3. The wrapper install pulls in the matching `mate` package automatically
+4. Install `mate` directly only when no wrapper matches your workload or you need direct MATE APIs
+5. Keep the upstream import path and high-level API shape as stable as possible
+6. Use `mate check`, `mate show-config`, `mate env`, logging, and replay if something fails
 
 ## Quick Start
 
-Use these commands after the MUSA-enabled `torch` / `torch_musa` stack is
-installed. Keep dependency resolution disabled for local builds so pip does not
-replace that stack with upstream PyPI packages.
+### Install from the MUSA Wheel Source
+
+For delivered builds, use the external delivery source after the MUSA-enabled
+`torch` / `torch_musa` stack is already installed.
 
-- Use `--no-build-isolation` for source installs.
-- Use `--no-isolation` for wheel builds.
-- Use `--no-deps` when installing local builds.
+```bash
+python -m pip install flash_attn_3 \
+  --index-url https://dl.mthreads.com/repo/api/pypi/pypi/simple
+```
 
-### Development Install
+This installs the matching `mate` dependency automatically. Install other
+wrapper packages such as `flash_mla`, `flash_kda`, `deep-gemm`, or
+`sageattention` the same way.
+
+MUSA compatibility wrapper releases use the PEP 440 local version suffix
+`+musa`, for example `0.2.4+musa`. Check the installed distribution with
+`python -m pip show flash_attn_3`; a version ending in `+musa` identifies the
+MATE-backed MUSA wrapper rather than the native implementation. The matching
+`mate` dependency keeps its normal version without this wrapper-only suffix.
+
+Install `mate` directly only when you need direct MATE APIs without a wrapper:
 
 ```bash
-git clone https://github.com/MooreThreads/mate.git --recursive
-cd mate
-pip install --no-build-isolation --no-deps -e . -v
+python -m pip install mate \
+  --index-url https://dl.mthreads.com/repo/api/pypi/pypi/simple
 ```
 
-### Build a Wheel
+Use exactly one package index per install step. Do not mix the MUSA wheel
+source with a public PyPI mirror in the same `pip install` invocation.
+
+### Build from Source
 
 ```bash
 git clone https://github.com/MooreThreads/mate.git --recursive
 cd mate
-python -m build --wheel --no-isolation
-python -m pip install --no-deps dist/mate-*.whl
+python -m pip install --no-build-isolation --no-deps -e . -v
 ```
 
-### Optional: Pre-Build AOT Kernels
+For a local wheel instead of an editable install:
 
 ```bash
-MATE_MUSA_ARCH_LIST=3.1 python -m mate.aot
 python -m build --wheel --no-isolation
+python -m pip install --no-deps dist/mate-*.whl
 ```
 
-Customize AOT coverage when needed:
+See [`docs/source/install.rst`](docs/source/install.rst) for the full
+wheel-source workflow, wrapper package commands, `PIP_INDEX_URL`, `pip config`,
+`pip index versions`, direct wheel URLs, and local AOT build options.
 
-```bash
-python -m mate.aot --attention-aot-level 0 --add-gemm true --add-moe false
-```
+### Choose a Wrapper
 
-### Install a Wrapper
-
-After MATE is installed, install the wrapper that matches your framework's
-expected Python package surface.
+Install the wrapper that matches your framework's expected Python package
+surface. Each delivered wrapper install pulls in the matching `mate`
+dependency automatically.
 
 | Wrapper directory | Package | Import path | Typical use |
 | --- | --- | --- | --- |
@@ -82,24 +100,8 @@ expected Python package surface.
 | `wrappers/DeepGEMM` | `deep-gemm` | `deep_gemm` | DeepGEMM style integration |
 | `wrappers/SageAttention` | `sageattention` | `sageattention` | SageAttention style integration |
 
-Generic editable install pattern:
-
-```bash
-cd wrappers/flash-attention
-pip install --no-build-isolation -e .
-```
-
-Generic wheel install pattern:
-
-```bash
-cd wrappers/flash-attention
-python -m build --wheel
-pip install dist/flash_attn_3-*.whl
-```
-
-Repeat the same workflow for `wrappers/FlashMLA`, `wrappers/FlashKDA`,
-`wrappers/DeepGEMM`, or `wrappers/SageAttention` when those package surfaces
-match your framework.
+For wrapper install commands, use the same external index or follow the local
+build workflow in [`docs/source/install.rst`](docs/source/install.rst).
 
 ### Verify and Diagnose
 
@@ -153,20 +155,6 @@ complete environment variable reference.
 
 MATE includes a guarded MUSA allocator that can replace the default `torch_musa` allocator during debugging to help localize out-of-bounds reads and writes across MUSA workloads. Start with `mate guard-run --mode tail -- python your_script.py`, or enable it in tests with `pytest --guard-alloc`. The detailed workflow, pytest defaults, and limitations are documented in [docs/guard_allocator.md](docs/guard_allocator.md).
 
-## Wrappers
-
-MATE uses the packages under `wrappers/` as a compatibility layer for CUDA-oriented software stacks on MUSA. These wrappers preserve familiar package names and high-level APIs while routing execution to MATE operators and kernels on MUSA, which helps existing integrations migrate with smaller code changes.
-
-For the guided wrapper-first documentation path, start with [docs/source/overview.rst](docs/source/overview.rst) and [docs/source/wrapper_tutorials.rst](docs/source/wrapper_tutorials.rst) for the wrapper quickstart flow.
-
-| Wrapper | Package | Import Path | Purpose | Documentation |
-| --- | --- | --- | --- | --- |
-| `wrappers/flash-attention` | `flash_attn_3` | `flash_attn_interface` | FlashAttention-3-compatible APIs on top of MATE attention operators on MUSA | [wrapper README](wrappers/flash-attention/README.md), [compatibility summary](docs/source/wrappers/flash_attention_forward_compatibility.md) |
-| `wrappers/SageAttention` | `sageattention` | `sageattention` | SageAttention-compatible dense quantized attention wrapper on top of MATE on MUSA | [wrapper README](wrappers/SageAttention/README.md) |
-| `wrappers/FlashMLA` | `flash_mla` | `flash_mla` | FlashMLA-compatible MLA dense/sparse decode and sparse prefill APIs on top of MATE MLA operators on MUSA | [wrapper README](wrappers/FlashMLA/README.md) |
-| `wrappers/FlashKDA` | `flash_kda` | `flash_kda` | FlashKDA-compatible KDA forward APIs on top of MATE KDA operators on MUSA | [wrapper README](wrappers/FlashKDA/README.md) |
-| `wrappers/DeepGEMM` | `deep-gemm` | `deep_gemm` | DeepGEMM-compatible APIs on top of MATE GEMM operators on MUSA | [wrapper README](wrappers/DeepGEMM/README.md) |
-
 ## Repository Layout
 
 | Path | Purpose |
@@ -180,15 +168,11 @@ For the guided wrapper-first documentation path, start with [docs/source/overvie
 ## Quick Links
 
 - Official documentation: [https://mate-docs.mthreads.com/latest/](https://mate-docs.mthreads.com/latest/)
+- Wrapper quickstarts: [docs/source/wrapper_tutorials.rst](docs/source/wrapper_tutorials.rst)
 - CLI documentation: [docs/mate_cli.md](docs/mate_cli.md)
 - Guard allocator debugging: [docs/guard_allocator.md](docs/guard_allocator.md)
 - Environment variables: [docs/environment_variables.md](docs/environment_variables.md)
 - FlashAttention-3 compatibility summary: [docs/source/wrappers/flash_attention_forward_compatibility.md](docs/source/wrappers/flash_attention_forward_compatibility.md)
-- FlashAttention-3 wrapper: [wrappers/flash-attention/README.md](wrappers/flash-attention/README.md)
-- SageAttention wrapper: [wrappers/SageAttention/README.md](wrappers/SageAttention/README.md)
-- FlashMLA wrapper: [wrappers/FlashMLA/README.md](wrappers/FlashMLA/README.md)
-- FlashKDA wrapper: [wrappers/FlashKDA/README.md](wrappers/FlashKDA/README.md)
-- DeepGEMM wrapper: [wrappers/DeepGEMM/README.md](wrappers/DeepGEMM/README.md)
 
 ## Acknowledgement
 

benchmarks/bench_flash_attn.py

@@ -7,8 +7,8 @@
 from mate import flash_attn_varlen_func, flash_attn_with_kvcache
 from mate.mha_interface import get_scheduler_metadata
 
-# MATE_BENCH_FMHA_DTYPE = torch.bfloat16
-MATE_BENCH_FMHA_DTYPE = torch.float8_e4m3fn
+MATE_BENCH_FMHA_DTYPE = torch.bfloat16
+# MATE_BENCH_FMHA_DTYPE = torch.float8_e4m3fn
 MATE_BENCH_FMHA_DEVICE = "musa"
 
 MATE_BENCH_FMHA_ENABLE_TRACE = False
@@ -18,15 +18,15 @@
 def gen_bench_config() -> List[dict]:
     cases = [
         {
-            "name": "prefill",
+            "name": "Prefill",
             "batch_size": 56,
-            "seqlen_q": [4096] * 56,
-            "seqlen_kv": [4096] * 56,
+            "seqlen_q": [1024 * 4] * 56,
+            "seqlen_kv": [1024 * 4] * 56,
             "head_q": 64,
             "head_kv": 4,
             "headdim_qk": 128,
             "headdim_vo": 128,
-            "is_packgqa": True,
+            "is_packgqa": False,
             "is_causal": False,
             "window_size": (None, None),
             "page_size": None,
@@ -38,29 +38,6 @@ def gen_bench_config() -> List[dict]:
             "seqused_q": None,
             "seqused_kv": None,
         },
-        {
-            "name": "decode",
-            "batch_size": 56,
-            "seqlen_q": [1] * 56,
-            "seqlen_kv": [4096] * 56,
-            "head_q": 64,
-            "head_kv": 4,
-            "headdim_qk": 128,
-            "headdim_vo": 128,
-            "is_packgqa": True,
-            "is_causal": True,
-            "window_size": (None, None),
-            "page_size": 128
-            if MATE_BENCH_FMHA_DTYPE in [torch.float8_e4m3fn, torch.float8_e5m2]
-            else 64,
-            "num_splits": 0,
-            "backend": "mutlass",
-            "dtype": MATE_BENCH_FMHA_DTYPE,
-            "device": MATE_BENCH_FMHA_DEVICE,
-            "force_flash_attn_with_kvcache": False,
-            "seqused_q": None,
-            "seqused_kv": None,
-        },
     ]
 
     return cases
@@ -150,6 +127,7 @@ def run_bench(
             pack_gqa=cfg.is_packgqa,
             mp_margin=0,
         )
+        print(metadata.view(4, -1))
 
     def fn():
         if cfg.is_paged_kv or force_flash_attn_with_kvcache: