mxflashattn 1.4.0__tar.gz → 1.5.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/MANIFEST.in +1 -0
- mxflashattn-1.5.0/PKG-INFO +188 -0
- mxflashattn-1.5.0/README.md +321 -0
- mxflashattn-1.5.0/README_PYPI.md +166 -0
- mxflashattn-1.5.0/benchmarks/prefill.py +246 -0
- mxflashattn-1.5.0/docs/API.md +233 -0
- mxflashattn-1.5.0/docs/C500_20261022_STAGE_B_REPORT.md +331 -0
- mxflashattn-1.5.0/docs/C500_20261023_QHEAD_CTA_EXPERIMENT.md +165 -0
- mxflashattn-1.5.0/docs/C500_20261023_SPLIT_COMBINE_EXPERIMENT.md +155 -0
- mxflashattn-1.5.0/docs/C500_20261023_V15_UNROLLED_VLLM_VALIDATION.md +149 -0
- mxflashattn-1.5.0/docs/C500_RUNTIME_DIAGNOSTIC_20261020.md +64 -0
- mxflashattn-1.5.0/docs/C500_STAGE_B_RUNBOOK.md +78 -0
- mxflashattn-1.5.0/docs/C500_VLLM017_VALIDATION_20261023.md +70 -0
- mxflashattn-1.5.0/docs/EVIDENCE.md +58 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/GIT_WORKFLOW.md +4 -3
- mxflashattn-1.5.0/docs/PAPER_TO_ENGINEERING_NOTES.md +161 -0
- mxflashattn-1.5.0/docs/STAGE_B_RESEARCH_AND_NEXT_EXPERIMENTS.md +125 -0
- mxflashattn-1.5.0/docs/TECHNICAL_HIGHLIGHTS.md +141 -0
- mxflashattn-1.5.0/docs/V1.5_BACKWARD_STAGE_A.md +92 -0
- mxflashattn-1.5.0/docs/V1.5_BACKWARD_STAGE_B_C500.md +121 -0
- mxflashattn-1.5.0/docs/V1.5_LOCAL_STAGE_A.md +101 -0
- mxflashattn-1.5.0/docs/V1_5_017_RECHECK.md +95 -0
- mxflashattn-1.5.0/docs/V1_5_028_MATRIX.md +220 -0
- mxflashattn-1.5.0/docs/V1_5_EXP_BREAKTHROUGH.md +166 -0
- mxflashattn-1.5.0/docs/V1_5_GQA_SWEEP_VERDICT.md +89 -0
- mxflashattn-1.5.0/docs/V1_5_KERNEL_GAP_ROOTCAUSE.md +154 -0
- mxflashattn-1.5.0/docs/V1_5_M1_DISPATCH_EVIDENCE_FIX.md +104 -0
- mxflashattn-1.5.0/docs/V1_5_MMA_DECODE_PROTOTYPE.md +308 -0
- mxflashattn-1.5.0/docs/V1_5_MMA_KV_SPLIT.md +325 -0
- mxflashattn-1.5.0/docs/V1_5_RELEASE_VALIDATION.md +153 -0
- mxflashattn-1.5.0/docs/V1_5_RESEARCH_20261024.md +170 -0
- mxflashattn-1.5.0/docs/V1_5_STAGE_B_C500_20261019.md +249 -0
- mxflashattn-1.5.0/docs/V1_5_WARP_PROFILE_20261024.md +212 -0
- mxflashattn-1.5.0/docs/VLLM_028_HANDOVER_2026-10-22.md +108 -0
- mxflashattn-1.5.0/docs/results/c500-mma-profile-qwen3-4b-hkv8-2026-10-24.json +88 -0
- mxflashattn-1.5.0/docs/results/c500-mma-qwen3-4b-hkv8-2026-10-24.json +342 -0
- mxflashattn-1.5.0/docs/results/readme-data-2026-10-24.json +143 -0
- mxflashattn-1.5.0/docs/results/readme-mma-operator-2026-10-24.png +0 -0
- mxflashattn-1.5.0/docs/results/readme-runtime-routing.png +0 -0
- mxflashattn-1.5.0/docs/results/readme-vllm017-e2e-2026-10-23.png +0 -0
- mxflashattn-1.5.0/docs/results/readme-vllm028-delta-2026-10-23.png +0 -0
- mxflashattn-1.5.0/docs/support-matrix.md +85 -0
- mxflashattn-1.5.0/integrations/vllm/README.md +19 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/backend.py +101 -17
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/compat.py +51 -46
- mxflashattn-1.5.0/integrations/vllm/config.py +15 -0
- mxflashattn-1.5.0/integrations/vllm/decoder_adapter.py +751 -0
- mxflashattn-1.5.0/integrations/vllm/m1_config.py +64 -0
- mxflashattn-1.5.0/integrations/vllm/plugin.py +73 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/smoke_generate.py +150 -15
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/vllm_plugin.py +45 -13
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn/__init__.py +1 -1
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn/api.py +13 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn/dispatch.py +60 -0
- mxflashattn-1.5.0/mxflashattn/metrics.py +126 -0
- mxflashattn-1.5.0/mxflashattn/mma_metadata.py +66 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn/policy.py +41 -2
- mxflashattn-1.5.0/mxflashattn/profile.py +233 -0
- mxflashattn-1.5.0/mxflashattn.egg-info/PKG-INFO +188 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn.egg-info/SOURCES.txt +58 -0
- mxflashattn-1.5.0/mxflashattn.egg-info/entry_points.txt +2 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/pyproject.toml +43 -36
- mxflashattn-1.5.0/tests/test_api.py +458 -0
- mxflashattn-1.5.0/tests/test_bench_mma_components.py +89 -0
- mxflashattn-1.5.0/tests/test_build_mma_profile.py +261 -0
- mxflashattn-1.5.0/tests/test_c500_preflight.py +50 -0
- mxflashattn-1.5.0/tests/test_c500_profile_summary.py +76 -0
- mxflashattn-1.5.0/tests/test_c500_streaming_ttft.py +21 -0
- mxflashattn-1.5.0/tests/test_c500_sweep_manifest.py +106 -0
- mxflashattn-1.5.0/tests/test_decoder_adapter.py +614 -0
- mxflashattn-1.5.0/tests/test_fixed_path_e2e.py +134 -0
- mxflashattn-1.5.0/tests/test_local_prefill_benchmark.py +30 -0
- mxflashattn-1.5.0/tests/test_m1_dispatch_evidence.py +90 -0
- mxflashattn-1.5.0/tests/test_merge_mma_profiles.py +127 -0
- mxflashattn-1.5.0/tests/test_metrics.py +45 -0
- mxflashattn-1.5.0/tests/test_mma_metadata.py +34 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_prefill.py +17 -0
- mxflashattn-1.5.0/tests/test_profile.py +175 -0
- mxflashattn-1.5.0/tests/test_ttft_comparison.py +228 -0
- mxflashattn-1.5.0/tests/test_version_metadata.py +16 -0
- mxflashattn-1.5.0/tests/test_vllm_backend_selection.py +607 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_vllm_bridge.py +27 -1
- mxflashattn-1.5.0/tests/test_vllm_compat.py +89 -0
- mxflashattn-1.4.0/PKG-INFO +0 -242
- mxflashattn-1.4.0/README.md +0 -220
- mxflashattn-1.4.0/docs/API.md +0 -71
- mxflashattn-1.4.0/docs/support-matrix.md +0 -53
- mxflashattn-1.4.0/integrations/vllm/README.md +0 -7
- mxflashattn-1.4.0/integrations/vllm/config.py +0 -11
- mxflashattn-1.4.0/integrations/vllm/decoder_adapter.py +0 -286
- mxflashattn-1.4.0/mxflashattn.egg-info/PKG-INFO +0 -242
- mxflashattn-1.4.0/tests/test_api.py +0 -219
- mxflashattn-1.4.0/tests/test_decoder_adapter.py +0 -76
- mxflashattn-1.4.0/tests/test_vllm_compat.py +0 -19
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/CONTRIBUTING.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/LICENSE +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/__init__.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/compare.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/configs/c500.yaml +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/fair_matrix.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/fair_report.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/model_fair_report.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/plot_summary.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/report.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/benchmarks/run.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/csrc/mxmac/CMakeLists.txt +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/csrc/mxmac/attention.h +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/csrc/mxmac/attention_forward.cpp +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/csrc/mxmac/paged_kv.cpp +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/csrc/mxmac/workspace.cpp +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/C500_MINIMAL_RUNBOOK.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/C500_V14_EXPERIMENT_PLAN.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/C500_V14_VALIDATION_20261017.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/C500_VALIDATION_RUNBOOK.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/RELEASE_LAYOUT.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/TECHNICAL_NOTES.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/V1.4.0_DEV_PLAN.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/V1.5_OPTIMIZATION_PLAN.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/gate4/E2E_/345/244/247/346/250/241/345/236/213/344/270/216/345/256/210/345/215/253/351/252/214/350/257/201.md" +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/gate4/Gate4_/345/217/202/350/200/203/346/226/207/347/214/256/344/270/216/350/260/203/347/240/224/350/256/260/345/275/225.md" +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/gate4/Gate4_/346/224/266/345/256/230/346/212/245/345/221/212.md" +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/gate4/M3_VLLM_WIRING.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/results/c500-2026-10-05.json +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/results/c500-2026-10-05.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/results/c500-2026-10-13.md +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/results/c500-model-level-summary.png +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/docs/results/c500-speedup-summary.png +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/__init__.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/__init__.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/probe_compat.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/probe_decoder_layout.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/probe_registration.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/integrations/vllm/probe_v017.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn/accuracy.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn/prefill.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn/reference.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn/validation.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn.egg-info/dependency_links.txt +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn.egg-info/requires.txt +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/mxflashattn.egg-info/top_level.txt +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/requirements-c500.txt +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/requirements-vllm.txt +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/setup.cfg +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/setup.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_accuracy.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_benchmark.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_c500.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_c500_plan.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_dispatch.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_long_context_accuracy.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_plot_summary.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_policy.py +0 -0
- {mxflashattn-1.4.0 → mxflashattn-1.5.0}/tests/test_vllm_official_split_fallback.py +0 -0
|
@@ -0,0 +1,188 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: mxflashattn
|
|
3
|
+
Version: 1.5.0
|
|
4
|
+
Summary: Observable FlashAttention-compatible decode integration for MetaX C500 and MXMACA
|
|
5
|
+
License-Expression: Apache-2.0
|
|
6
|
+
Classifier: Development Status :: 3 - Alpha
|
|
7
|
+
Classifier: Programming Language :: Python :: 3
|
|
8
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
9
|
+
Requires-Python: >=3.10
|
|
10
|
+
Description-Content-Type: text/markdown
|
|
11
|
+
License-File: LICENSE
|
|
12
|
+
Requires-Dist: torch>=2.4
|
|
13
|
+
Requires-Dist: PyYAML<7,>=6.0
|
|
14
|
+
Provides-Extra: dev
|
|
15
|
+
Requires-Dist: pytest<10,>=8; extra == "dev"
|
|
16
|
+
Requires-Dist: ruff<1,>=0.6; extra == "dev"
|
|
17
|
+
Provides-Extra: vllm
|
|
18
|
+
Requires-Dist: vllm==0.17.0; extra == "vllm"
|
|
19
|
+
Provides-Extra: reports
|
|
20
|
+
Requires-Dist: matplotlib<4,>=3.8; extra == "reports"
|
|
21
|
+
Dynamic: license-file
|
|
22
|
+
|
|
23
|
+
# MXFlashAttn
|
|
24
|
+
|
|
25
|
+
**MXFlashAttn 1.5.0** provides a FlashAttention-compatible Python API and an
|
|
26
|
+
observable paged-decode integration for **MetaX C500 / MXMACA**.
|
|
27
|
+
|
|
28
|
+
The package combines:
|
|
29
|
+
|
|
30
|
+
- validated dense, varlen, and KV-cache forward entry points;
|
|
31
|
+
- explicit validation, dispatch information, and opt-in reference fallback;
|
|
32
|
+
- an M1 Split-KV decode path for C500 paged KV cache;
|
|
33
|
+
- an exact-profile MMA Split-KV path for calibrated FP16 decode shapes;
|
|
34
|
+
- vLLM EngineCore registration and per-dispatch evidence.
|
|
35
|
+
|
|
36
|
+
Native C500 extensions are compiled inside the target MetaX/MACA environment.
|
|
37
|
+
The PyPI artifact supplies the Python API, integration layer, profile handling,
|
|
38
|
+
and registration entry point.
|
|
39
|
+
|
|
40
|
+
## Install
|
|
41
|
+
|
|
42
|
+
```bash
|
|
43
|
+
python -m pip install mxflashattn==1.5.0
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
Base requirements:
|
|
47
|
+
|
|
48
|
+
- Python 3.10+
|
|
49
|
+
- PyTorch 2.4+
|
|
50
|
+
- PyYAML 6.x
|
|
51
|
+
|
|
52
|
+
For the vLLM 0.17 dependency contract:
|
|
53
|
+
|
|
54
|
+
```bash
|
|
55
|
+
python -m pip install "mxflashattn[vllm]==1.5.0"
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
MetaX deployments normally provide their own compatible PyTorch, vLLM, and
|
|
59
|
+
MXMACA wheels. vLLM 0.17.0 and 0.28.0 were validated on separate C500 software
|
|
60
|
+
stacks; they are not interchangeable environments.
|
|
61
|
+
|
|
62
|
+
## Python API
|
|
63
|
+
|
|
64
|
+
```python
|
|
65
|
+
import os
|
|
66
|
+
import torch
|
|
67
|
+
|
|
68
|
+
from mxflashattn import flash_attn_func, get_last_dispatch_info
|
|
69
|
+
|
|
70
|
+
# Explicitly enable the correctness/reference path when no native provider exists.
|
|
71
|
+
os.environ["MXFLASHATTN_ALLOW_FALLBACK"] = "1"
|
|
72
|
+
|
|
73
|
+
q = torch.randn(1, 32, 8, 64, dtype=torch.float16)
|
|
74
|
+
k = torch.randn(1, 128, 2, 64, dtype=q.dtype)
|
|
75
|
+
v = torch.randn_like(k)
|
|
76
|
+
|
|
77
|
+
out = flash_attn_func(q, k, v, causal=True)
|
|
78
|
+
print(out.shape)
|
|
79
|
+
print(get_last_dispatch_info())
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
The reference path is intended for API checks and correctness comparison. It
|
|
83
|
+
is not a C500 performance backend. Without
|
|
84
|
+
`MXFLASHATTN_ALLOW_FALLBACK=1`, a missing native provider raises an explicit
|
|
85
|
+
error instead of silently selecting a slow path.
|
|
86
|
+
|
|
87
|
+
Public entry points:
|
|
88
|
+
|
|
89
|
+
```python
|
|
90
|
+
from mxflashattn import (
|
|
91
|
+
flash_attn_func,
|
|
92
|
+
flash_attn_varlen_func,
|
|
93
|
+
flash_attn_with_kvcache,
|
|
94
|
+
get_last_dispatch_info,
|
|
95
|
+
)
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
## C500 Native Setup
|
|
99
|
+
|
|
100
|
+
Build native extensions in an environment matching the target MetaX PyTorch
|
|
101
|
+
and MACA installation:
|
|
102
|
+
|
|
103
|
+
```bash
|
|
104
|
+
git clone https://github.com/leo1852098393-blip/MXFlashAttn.git
|
|
105
|
+
cd MXFlashAttn
|
|
106
|
+
python -m pip install --no-build-isolation -e .
|
|
107
|
+
|
|
108
|
+
export MACA_PATH=/opt/maca
|
|
109
|
+
bash csrc_native/build_m1.sh
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
Then enable the guarded M1 route and collect dispatch evidence:
|
|
113
|
+
|
|
114
|
+
```bash
|
|
115
|
+
export MXFA_M1_DECODE=1
|
|
116
|
+
export MXFLASHATTN_DISPATCH_SAMPLE=1
|
|
117
|
+
|
|
118
|
+
python integrations/vllm/smoke_generate.py \
|
|
119
|
+
--model /path/to/Qwen3-4B \
|
|
120
|
+
--suite --prompt-count 4 --max-tokens 16 \
|
|
121
|
+
--candidate --output artifacts/m1-smoke.json
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
A performance result is attributed to MXFlashAttn only when the associated
|
|
125
|
+
`dispatch.jsonl` contains a successful `m1_split_decode` or
|
|
126
|
+
`mma_split_decode` event.
|
|
127
|
+
|
|
128
|
+
## Verified Scope
|
|
129
|
+
|
|
130
|
+
| Area | v1.5.0 status |
|
|
131
|
+
| --- | --- |
|
|
132
|
+
| M1 decode | FP16 / BF16, head dimensions 64 / 128, GQA ratios 2 / 4 / 8 |
|
|
133
|
+
| vLLM 0.17.0 | Real Qwen3-4B/8B generation, M1 dispatch, fallback, and long-decode measurements |
|
|
134
|
+
| vLLM 0.28.0 | EngineCore plugin registration, real M1 dispatch, and a Qwen3-8B shape matrix |
|
|
135
|
+
| MMA decode | FP16 single-token decode, exact calibrated profiles only |
|
|
136
|
+
| Prefill | MetaX vendor production path; chunked prefill remains a correctness reference |
|
|
137
|
+
| Backward | Explicit MetaX provider integration; no in-house backward kernel |
|
|
138
|
+
|
|
139
|
+
In four vLLM 0.17.0 long-decode workloads, M1 delivered 2.67-3.54x the
|
|
140
|
+
throughput of the project's historical v7 path and reached 79.2%-86.5% of
|
|
141
|
+
MetaX vendor throughput. In one exact Qwen3-4B operator shape, MMA median
|
|
142
|
+
latency was `43.520 us` versus `46.592 us` for vendor over 100 rounds.
|
|
143
|
+
|
|
144
|
+
In the reported
|
|
145
|
+
[vLLM 0.28.0 Qwen3-8B matrix](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/V1_5_028_MATRIX.md),
|
|
146
|
+
long-KV samples at batch 16 or larger remained 28%-46% below MetaX vendor. No
|
|
147
|
+
general M1-over-vendor crossover is claimed.
|
|
148
|
+
|
|
149
|
+
These results are bound to their recorded device, software stack, model,
|
|
150
|
+
shape, and baseline. The MMA comparison is operator-level, not an end-to-end
|
|
151
|
+
request-speedup claim.
|
|
152
|
+
|
|
153
|
+
## Runtime Boundaries
|
|
154
|
+
|
|
155
|
+
- Prefill and unprofiled decode shapes retain the MetaX vendor path.
|
|
156
|
+
- M1 requires its documented dtype, layout, head-dimension, and GQA contract.
|
|
157
|
+
- MMA requires an exact profile match, including KV heads, page size, binary
|
|
158
|
+
variant, and shared-object identity.
|
|
159
|
+
- Short or low-concurrency requests may not benefit from M1.
|
|
160
|
+
- FP8 KV cache, multi-card execution, SGLang, GQ=7, and general TTFT gains are
|
|
161
|
+
not v1.5.0 support claims.
|
|
162
|
+
- Unsupported API options fail explicitly or follow a documented fallback.
|
|
163
|
+
|
|
164
|
+
## Development and Audit
|
|
165
|
+
|
|
166
|
+
The project's own kernel implementations under `csrc_native/` were written by
|
|
167
|
+
AI coding agents. The project owner set the project goals, verification rules,
|
|
168
|
+
and release boundary, and made the final accept-or-reject decisions for code
|
|
169
|
+
changes and public claims.
|
|
170
|
+
|
|
171
|
+
Favorable measurements do not automatically become performance claims. The
|
|
172
|
+
short-KV `+15.2%` difference was the same order as repeated-run variance; the
|
|
173
|
+
long-KV `+22.7%` result was a single observation and does not establish a
|
|
174
|
+
general crossover. Neither is used for a general performance claim.
|
|
175
|
+
|
|
176
|
+
## Documentation
|
|
177
|
+
|
|
178
|
+
- [GitHub repository](https://github.com/leo1852098393-blip/MXFlashAttn)
|
|
179
|
+
- [GitLink mirror](https://gitlink.org.cn/Leo77/MXFlashAttn)
|
|
180
|
+
- [API guide](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/API.md)
|
|
181
|
+
- [Support matrix](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/support-matrix.md)
|
|
182
|
+
- [Evidence index](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/EVIDENCE.md)
|
|
183
|
+
- [Release validation](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/V1_5_RELEASE_VALIDATION.md)
|
|
184
|
+
|
|
185
|
+
## License
|
|
186
|
+
|
|
187
|
+
Apache License 2.0. MXMACA runtime, MetaX vendor packages, vLLM, model
|
|
188
|
+
weights, and datasets remain subject to their respective licenses.
|
|
@@ -0,0 +1,321 @@
|
|
|
1
|
+
# MXFlashAttn
|
|
2
|
+
|
|
3
|
+
> 面向 **MetaX C500 / MXMACA** 的 FlashAttention 兼容层与可观测
|
|
4
|
+
> paged-decode 优化工程。
|
|
5
|
+
|
|
6
|
+
中文说明 · [English](README.en.md)
|
|
7
|
+
|
|
8
|
+
[](https://github.com/leo1852098393-blip/MXFlashAttn/actions/workflows/ci.yml)
|
|
9
|
+
[](LICENSE)
|
|
10
|
+
[](pyproject.toml)
|
|
11
|
+
[](https://pypi.org/project/mxflashattn/)
|
|
12
|
+
|
|
13
|
+
**版本**:v1.5.0(2026-10-24) · **设备**:MetaX C500 64 GiB ·
|
|
14
|
+
**已验证引擎**:vLLM 0.17.0 / 0.28.0 · **许可证**:Apache-2.0
|
|
15
|
+
|
|
16
|
+
MXFlashAttn 把自研 Attention 优化接入真实 C500 推理链路:读取 vLLM paged
|
|
17
|
+
KV cache,在 EngineCore worker 中完成后端注册,按输入 shape 选择 M1 或 MMA
|
|
18
|
+
Split-KV decode,并为每次候选派发记录可审计证据。不满足条件的请求保留
|
|
19
|
+
MetaX vendor 路径。
|
|
20
|
+
|
|
21
|
+
## 已验证结果
|
|
22
|
+
|
|
23
|
+
| 结论 | 实测范围 | 证据 |
|
|
24
|
+
| --- | --- | --- |
|
|
25
|
+
| M1 达到项目 v7 路径的 **2.67-3.54x**,达到 MetaX vendor 吞吐的 **79.2%-86.5%** | vLLM 0.17.0,Qwen3-4B/8B,4 组长 decode workload | [复测报告](docs/V1_5_017_RECHECK.md) |
|
|
26
|
+
| MMA 中位延迟 `43.520 us`,vendor 为 `46.592 us`,该精确 shape 低约 **6.6%** | FP16,batch=1,KV=128,32 Q heads / 8 KV heads,GQA=4,head_dim=128,split=4 | [发布复核](docs/V1_5_RELEASE_VALIDATION.md) |
|
|
27
|
+
| 自研路径已进入真实 EngineCore worker,并可区分候选命中、回退与未标定输入 | vLLM 0.17.0 / 0.28.0 的 C500 实例 | [证据索引](docs/EVIDENCE.md) |
|
|
28
|
+
|
|
29
|
+
这里的 v7 是项目自有历史内核,不是 MetaX vendor。MMA 的 6.6% 是一个精确
|
|
30
|
+
算子 shape 的结果,不是端到端请求加速结论。
|
|
31
|
+
|
|
32
|
+
| 你想了解什么 | 从这里开始 |
|
|
33
|
+
| --- | --- |
|
|
34
|
+
| 在 C500 上安装和运行 | [快速开始](#快速开始) |
|
|
35
|
+
| 判断项目的技术工作量 | [技术难点与工程演进](#技术难点与工程演进) |
|
|
36
|
+
| 核对性能数字 | [实测结果](#实测结果) · [证据索引](docs/EVIDENCE.md) |
|
|
37
|
+
| 查看支持和回退边界 | [支持范围与性能边界](#支持范围与性能边界) |
|
|
38
|
+
|
|
39
|
+
## 为什么这项工程重要
|
|
40
|
+
|
|
41
|
+
在国产 AI GPU 上,Attention 优化不是把一个独立 kernel 编译成功就结束了。
|
|
42
|
+
要让它进入真实推理系统,还需要同时解决:
|
|
43
|
+
|
|
44
|
+
1. vLLM paged KV cache 的布局与生命周期;
|
|
45
|
+
2. 主进程和 EngineCore worker 之间的后端注册;
|
|
46
|
+
3. dtype、GQA、head dimension、KV heads、page size 等输入守卫;
|
|
47
|
+
4. Prefill、未标定 shape 和原生调用失败时的可靠回退;
|
|
48
|
+
5. “注册成功”和“自研 kernel 确实执行”之间的证据差异;
|
|
49
|
+
6. 算子延迟、端到端吞吐和 TTFT 等不同指标的归因边界。
|
|
50
|
+
|
|
51
|
+
MXFlashAttn 围绕这条完整链路建设。项目的目标不是只展示一次可运行的算子,
|
|
52
|
+
而是形成一套**可运行、可测量、可审计、可继续优化**的 C500 Attention
|
|
53
|
+
工程基础。
|
|
54
|
+
|
|
55
|
+
## 核心贡献
|
|
56
|
+
|
|
57
|
+
| 能力 | 项目实现 |
|
|
58
|
+
| --- | --- |
|
|
59
|
+
| 自研 paged decode | M1 Split-KV kernel 直接读取 vLLM paged KV cache,覆盖 FP16 / BF16、head dimension 64 / 128、GQA 2 / 4 / 8 |
|
|
60
|
+
| 张量核路径 | 自研 MMA Split-KV decode,通过设备探针恢复 fragment 映射,只对精确标定的 FP16 shape 启用 |
|
|
61
|
+
| 真实引擎接线 | 通过 vLLM general plugin 进入 EngineCore worker;已在 vLLM 0.17.0 与 0.28.0 记录真实 dispatch |
|
|
62
|
+
| 可观测运行时 | 每次候选派发记录 backend、operation、selection、fallback、split 等字段 |
|
|
63
|
+
| 守卫与回退 | dtype、shape、layout、profile 或原生调用不满足条件时回到 MetaX vendor,不让实验路径处理未知输入 |
|
|
64
|
+
| 可复现实验 | 基准脚本、脱敏数据、图表生成器、测量方法和负向结果随仓库公开 |
|
|
65
|
+
|
|
66
|
+
## 技术难点与工程演进
|
|
67
|
+
|
|
68
|
+
### 1. 从设备探针恢复 MMA fragment 契约
|
|
69
|
+
|
|
70
|
+
项目不能依赖完整公开的 `__builtin_mxc_mma_16x16x16f16` fragment 布局,
|
|
71
|
+
因此编写逐线程写入、读回和 NumPy 对照探针,确认 C500 上观察到的执行组为
|
|
72
|
+
64 线程,并恢复 accumulator 的线程到矩阵元素映射。映射写错时可能得到有限
|
|
73
|
+
但错误的数值,而不是直接崩溃,因此这一步同时是正确性门槛。
|
|
74
|
+
|
|
75
|
+
### 2. 定位 Split-KV 的非确定性越界写
|
|
76
|
+
|
|
77
|
+
早期 MMA Split-KV 结果会在相同输入下破坏不同 scratch slot。逐 slot 扫描
|
|
78
|
+
发现损坏状态具有 `m=0, l=256` 指纹,最终定位到发布条件把 `lane < GQ`
|
|
79
|
+
误写为 `lane < 16`:无效 query lane 会写穿到下一个 slot,并与合法 CTA
|
|
80
|
+
发生竞态。修复后:
|
|
81
|
+
|
|
82
|
+
| 校验 | 修复前 | 修复后 |
|
|
83
|
+
| --- | ---: | ---: |
|
|
84
|
+
| 正确的 per-split `(m, l)` slot | 4 / 8 | 8 / 8 |
|
|
85
|
+
| 合并输出相对 FP32 参考的最大误差 | `2.99e-2` | `6.81e-5` |
|
|
86
|
+
|
|
87
|
+
### 3. 把后端注册送进真正执行它的进程
|
|
88
|
+
|
|
89
|
+
vLLM 会在独立解释器中创建 EngineCore。只在启动脚本注册后端,会出现父进程
|
|
90
|
+
显示 `registered=true`、worker 仍走 vendor 的假阳性。v1.5.0 通过
|
|
91
|
+
`vllm.general_plugins` entry point 在各进程重复注册,并在报告成功前检查
|
|
92
|
+
最终解析到的实现类。
|
|
93
|
+
|
|
94
|
+
### 4. 从历史 v7 到精确 profile 路由
|
|
95
|
+
|
|
96
|
+
项目先以 v7 建立历史基线,再完成 M1 Split-KV paged decode、真实 worker
|
|
97
|
+
接线和 dispatch 记录,最后加入 MMA 张量核与精确 shape profile。当前 MMA
|
|
98
|
+
不会向邻近 shape 外推;batch、KV length、GQA、head dimension、KV heads、
|
|
99
|
+
page size、二进制变体或共享库身份任一不匹配,都会回到 vendor。
|
|
100
|
+
|
|
101
|
+
完整技术说明见 [技术亮点](docs/TECHNICAL_HIGHLIGHTS.md)。
|
|
102
|
+
|
|
103
|
+
## 实测结果
|
|
104
|
+
|
|
105
|
+
以下数据均来自 MetaX C500 64 GiB。不同 MACA、PyTorch 和 vLLM 软件栈分别
|
|
106
|
+
报告,不做跨版本混算。
|
|
107
|
+
|
|
108
|
+
### 真实 vLLM 0.17 长 decode
|
|
109
|
+
|
|
110
|
+

|
|
111
|
+
|
|
112
|
+
测试环境为 MACA 3.5.3.20、PyTorch `2.8.0+metax3.5.3.9`、vLLM 0.17.0;
|
|
113
|
+
Qwen3-4B / 8B 每组生成 1024 token,关闭 prefix caching。四组 workload
|
|
114
|
+
均记录到 `m1_split_decode`:
|
|
115
|
+
|
|
116
|
+
- M1 吞吐为项目 v7 路径的 **2.67-3.54 倍**;
|
|
117
|
+
- M1 达到 MetaX vendor 吞吐的 **79.2%-86.5%**;
|
|
118
|
+
- M1 仍比 vendor 低 **13.5%-20.8%**。
|
|
119
|
+
|
|
120
|
+
这组结果说明 M1 明显改善了项目自研 decode 路径,同时也量化了与 vendor
|
|
121
|
+
之间仍需优化的空间。
|
|
122
|
+
|
|
123
|
+
数据与方法:[vLLM 0.17 复测报告](docs/V1_5_017_RECHECK.md) ·
|
|
124
|
+
[图表原始数据](docs/results/readme-data-2026-10-24.json)
|
|
125
|
+
|
|
126
|
+
### C500 MMA 精确 shape
|
|
127
|
+
|
|
128
|
+

|
|
129
|
+
|
|
130
|
+
在 Qwen3-4B 的实际配置上,FP16、batch=1、KV=128、32 Q heads、
|
|
131
|
+
8 KV heads、GQA=4、head_dim=128、page_size=16、split=4,100 轮算子
|
|
132
|
+
计时中:
|
|
133
|
+
|
|
134
|
+
- MetaX vendor 中位延迟:`46.592 us`;
|
|
135
|
+
- MXFlashAttn MMA Split-KV 中位延迟:`43.520 us`;
|
|
136
|
+
- 该精确 shape 下 MMA 中位延迟低约 **6.6%**。
|
|
137
|
+
|
|
138
|
+
同一 profile 已在真实 vLLM 请求中记录 36 次 `mma_split_decode`,整次请求
|
|
139
|
+
共有 396 个 attention 事件。这个结果证明了“算子标定 -> 精确 profile ->
|
|
140
|
+
真实 worker 派发”的链路,不代表整个请求获得 6.6% 的端到端提升。
|
|
141
|
+
|
|
142
|
+
数据与方法:[MMA Split-KV 技术记录](docs/V1_5_MMA_KV_SPLIT.md) ·
|
|
143
|
+
[发布复核](docs/V1_5_RELEASE_VALIDATION.md) ·
|
|
144
|
+
[公开 profile](docs/results/c500-mma-profile-qwen3-4b-hkv8-2026-10-24.json)
|
|
145
|
+
|
|
146
|
+
## 运行路径
|
|
147
|
+
|
|
148
|
+

|
|
149
|
+
|
|
150
|
+
1. vLLM worker 将 Attention 请求交给 MXFlashAttn adapter。
|
|
151
|
+
2. Adapter 检查 phase、dtype、layout、shape 和 profile。
|
|
152
|
+
3. 精确 profile 命中时可选择 MMA;满足 M1 合约且显式启用时可选择 M1。
|
|
153
|
+
4. Prefill、未标定 shape、不支持输入或原生执行失败时回到 MetaX vendor。
|
|
154
|
+
5. 候选路径写入 dispatch 记录,性能数据只有在目标 kernel 确实派发后才归因。
|
|
155
|
+
|
|
156
|
+
## 快速开始
|
|
157
|
+
|
|
158
|
+
### Python API
|
|
159
|
+
|
|
160
|
+
```bash
|
|
161
|
+
python -m pip install mxflashattn==1.5.0
|
|
162
|
+
```
|
|
163
|
+
|
|
164
|
+
```python
|
|
165
|
+
import os
|
|
166
|
+
import torch
|
|
167
|
+
|
|
168
|
+
from mxflashattn import flash_attn_func, get_last_dispatch_info
|
|
169
|
+
|
|
170
|
+
# 无原生 provider 时,reference 路径必须显式开启。
|
|
171
|
+
os.environ["MXFLASHATTN_ALLOW_FALLBACK"] = "1"
|
|
172
|
+
|
|
173
|
+
q = torch.randn(1, 32, 8, 64, dtype=torch.float16)
|
|
174
|
+
k = torch.randn(1, 128, 2, 64, dtype=q.dtype)
|
|
175
|
+
v = torch.randn_like(k)
|
|
176
|
+
|
|
177
|
+
out = flash_attn_func(q, k, v, causal=True)
|
|
178
|
+
print(out.shape)
|
|
179
|
+
print(get_last_dispatch_info())
|
|
180
|
+
```
|
|
181
|
+
|
|
182
|
+
Reference fallback 用于接口调试和正确性对照,不是 C500 性能路径。
|
|
183
|
+
|
|
184
|
+
### C500 / vLLM
|
|
185
|
+
|
|
186
|
+
原生扩展需要在与目标 MetaX PyTorch / MACA 匹配的 C500 环境中编译:
|
|
187
|
+
|
|
188
|
+
```bash
|
|
189
|
+
git clone https://github.com/leo1852098393-blip/MXFlashAttn.git
|
|
190
|
+
cd MXFlashAttn
|
|
191
|
+
python -m pip install --no-build-isolation -e .
|
|
192
|
+
|
|
193
|
+
export MACA_PATH=/opt/maca
|
|
194
|
+
bash csrc_native/build_m1.sh
|
|
195
|
+
|
|
196
|
+
export MXFA_M1_DECODE=1
|
|
197
|
+
export MXFLASHATTN_DISPATCH_SAMPLE=1
|
|
198
|
+
python integrations/vllm/smoke_generate.py \
|
|
199
|
+
--model /path/to/Qwen3-4B \
|
|
200
|
+
--suite --prompt-count 4 --max-tokens 16 \
|
|
201
|
+
--candidate --output artifacts/m1-smoke.json
|
|
202
|
+
```
|
|
203
|
+
|
|
204
|
+
验收时同时检查结果 JSON 和同目录的 `dispatch.jsonl`。只有出现成功的
|
|
205
|
+
`m1_split_decode` 或 `mma_split_decode` 事件,相关吞吐或延迟才能归因于
|
|
206
|
+
MXFlashAttn 自研路径。MMA 的构建、profile 和精确匹配方式见
|
|
207
|
+
[API 文档](docs/API.md)。
|
|
208
|
+
|
|
209
|
+
## 支持范围与性能边界
|
|
210
|
+
|
|
211
|
+
| 项目 | 当前能力 |
|
|
212
|
+
| --- | --- |
|
|
213
|
+
| Python API | `flash_attn_func`、`flash_attn_varlen_func`、`flash_attn_with_kvcache` |
|
|
214
|
+
| M1 decode dtype | FP16 / BF16 |
|
|
215
|
+
| M1 shape | head dimension 64 / 128;GQA 2 / 4 / 8 |
|
|
216
|
+
| vLLM | 0.17.0 与 0.28.0 均完成 C500 worker 接线和 dispatch 验证 |
|
|
217
|
+
| MMA decode | FP16 单 token decode;只对精确标定 shape 启用 |
|
|
218
|
+
| Prefill | 生产路径使用 MetaX vendor;chunked prefill 是正确性研究原型 |
|
|
219
|
+
| Backward | 可显式调用 MetaX backward provider;不是项目自研反向 kernel |
|
|
220
|
+
| 不支持 shape | 不进入自研 kernel,回到 MetaX vendor 或明确报错 |
|
|
221
|
+
|
|
222
|
+
更完整的 dtype、shape、软件栈和 fallback 说明见
|
|
223
|
+
[支持矩阵](docs/support-matrix.md)与 [API 文档](docs/API.md)。
|
|
224
|
+
|
|
225
|
+
### vLLM 0.28 形状矩阵
|
|
226
|
+
|
|
227
|
+

|
|
228
|
+
|
|
229
|
+
这组结果来自独立软件栈:MACA 3.8.2.6、PyTorch 2.10.0、vLLM 0.28.0、
|
|
230
|
+
Qwen3-8B。它揭示了当前 M1 的适用边界:
|
|
231
|
+
|
|
232
|
+
- 短 KV 序列中有 3 个 batch 点高于 vendor(最高 `+15.2%`),但这些差异
|
|
233
|
+
仍处于重复测量方差带内,不作为性能结论;
|
|
234
|
+
- 中 KV、大 batch 样本低于 vendor 约 15%-19%;
|
|
235
|
+
- 长 KV、batch >= 16 的样本低于 vendor 约 28%-46%;
|
|
236
|
+
- batch=4 长 KV 的 `+22.7%` 是单个观测点,不作为通用 crossover。
|
|
237
|
+
|
|
238
|
+
进一步的 NP 扫描把 batch=16、序列容量 1024 的一个 workload 收敛到
|
|
239
|
+
vendor 的 `-2.2%`,但 batch=128 下所有已测 NP 仍落后 17.6%-24.2%。
|
|
240
|
+
这说明 shape 标定可以减少调度损失,但大 batch 扩展性仍需要内核级访存、
|
|
241
|
+
softmax / reduction 和 split / combine 优化。
|
|
242
|
+
|
|
243
|
+
完整矩阵与测量伪影修复记录见
|
|
244
|
+
[vLLM 0.28 测量报告](docs/V1_5_028_MATRIX.md)。
|
|
245
|
+
|
|
246
|
+
## 项目结构
|
|
247
|
+
|
|
248
|
+
```text
|
|
249
|
+
mxflashattn/ Python API、参数校验、reference、dispatch 与 profile
|
|
250
|
+
csrc_native/ M1、v7、MMA device kernel 及构建、正确性和基准工具
|
|
251
|
+
csrc/mxmac/ ATen bring-up / 正确性对照路径,不是融合优化 kernel
|
|
252
|
+
integrations/vllm/ vLLM backend、worker 注册、layout adapter 与 dispatch 记录
|
|
253
|
+
benchmarks/ 固定 workload 的基准与报告工具
|
|
254
|
+
scripts/ profile、结果汇总、发布审计和 README 图表生成
|
|
255
|
+
tests/ CPU/reference 回归与按环境启用的 C500 测试
|
|
256
|
+
docs/results/ 脱敏结果、图表及图表数据
|
|
257
|
+
docs/ 支持边界、实验方法和技术记录
|
|
258
|
+
```
|
|
259
|
+
|
|
260
|
+
## 可复现与审计
|
|
261
|
+
|
|
262
|
+
`csrc_native/` 中的项目自有内核实现由 AI 编程代理完成。项目负责人确定项目
|
|
263
|
+
目标、验证规则与发布边界,并对代码变更和公开结论作出最终的采纳或否决。
|
|
264
|
+
|
|
265
|
+
- **归因规则**:只有 dispatch 日志出现成功的 `m1_split_decode` 或
|
|
266
|
+
`mma_split_decode` 事件,性能才计入本项目;注册成功本身不算证据。
|
|
267
|
+
- **结论门槛**:对项目有利的测量也不会自动成为性能结论。短 KV 的 `+15.2%`
|
|
268
|
+
与重复测量方差同量级;长 KV 的 `+22.7%` 只是单个观测点,不能据此建立
|
|
269
|
+
通用 crossover。两者均未用于建立通用性能主张。
|
|
270
|
+
- **复现优先**:缺陷须先产出稳定复现用例,修复后该用例转绿且优化路径仍在
|
|
271
|
+
派发,才算修复完成。
|
|
272
|
+
- README 图表由
|
|
273
|
+
[`scripts/plot_readme_results.py`](scripts/plot_readme_results.py)
|
|
274
|
+
从
|
|
275
|
+
[`docs/results/readme-data-2026-10-24.json`](docs/results/readme-data-2026-10-24.json)
|
|
276
|
+
生成。
|
|
277
|
+
- [证据索引](docs/EVIDENCE.md)把每项首页结论映射到设备、软件栈、模型、
|
|
278
|
+
workload、基线、公开数据和限制。
|
|
279
|
+
- vLLM 0.17 与 0.28 使用不同软件栈,结果单独解释。
|
|
280
|
+
- 原始模型权重、租用实例日志和凭据不进入公开仓库;公开摘要保留复现配置与
|
|
281
|
+
可审计结论。
|
|
282
|
+
|
|
283
|
+
## 发展路线
|
|
284
|
+
|
|
285
|
+
### 1. M1 内核与调度联合优化
|
|
286
|
+
|
|
287
|
+
- 建立 `(batch, KV length, GQA, head_dim)` 的可复现 profile;
|
|
288
|
+
- 优化 KV staging、softmax / reduction 和 split / combine 开销;
|
|
289
|
+
- 以相同输入、独立进程、dispatch 命中、重复测量和 vendor 对照作为合入门槛。
|
|
290
|
+
|
|
291
|
+
### 2. 扩展 MMA 的有效覆盖
|
|
292
|
+
|
|
293
|
+
- 补充 BF16、长 KV、更多真实模型 shape 和误差压力测试;
|
|
294
|
+
- 解决 CUDA Graph 下 scratch 生命周期与 capture 约束;
|
|
295
|
+
- 只有在真实请求中取得稳定端到端收益后,才扩大自动 dispatch 范围。
|
|
296
|
+
|
|
297
|
+
### 3. Prefill 与服务指标
|
|
298
|
+
|
|
299
|
+
- 研究 C500 专用融合 Prefill kernel;
|
|
300
|
+
- 建立可归因的 TTFT P50 / P95 / P99、吞吐与显存联合测量;
|
|
301
|
+
- 将 Prefill、Decode 和调度开销分开报告。
|
|
302
|
+
|
|
303
|
+
### 4. 生态扩展
|
|
304
|
+
|
|
305
|
+
- 持续验证新的 vLLM / MACA 组合;
|
|
306
|
+
- 在同一证据门槛下评估 FP8 KV cache、多卡和 SGLang;
|
|
307
|
+
- 将 profile、测试协议和 dispatch 证据格式沉淀为 C500 Attention
|
|
308
|
+
优化基础设施。
|
|
309
|
+
|
|
310
|
+
研究依据与实验映射见
|
|
311
|
+
[研究笔记](docs/V1_5_RESEARCH_20261024.md)。
|
|
312
|
+
|
|
313
|
+
## 参与贡献
|
|
314
|
+
|
|
315
|
+
Issue 与 Pull Request 请附带最小复现、设备与软件栈版本、输入 shape,以及
|
|
316
|
+
dispatch 记录。开发与测试约定见 [CONTRIBUTING.md](CONTRIBUTING.md)。
|
|
317
|
+
|
|
318
|
+
## 许可证
|
|
319
|
+
|
|
320
|
+
本项目采用 [Apache License 2.0](LICENSE)。MXMACA runtime、MetaX vendor
|
|
321
|
+
wheel、vLLM、模型权重和数据集分别遵循各自许可证。
|
|
@@ -0,0 +1,166 @@
|
|
|
1
|
+
# MXFlashAttn
|
|
2
|
+
|
|
3
|
+
**MXFlashAttn 1.5.0** provides a FlashAttention-compatible Python API and an
|
|
4
|
+
observable paged-decode integration for **MetaX C500 / MXMACA**.
|
|
5
|
+
|
|
6
|
+
The package combines:
|
|
7
|
+
|
|
8
|
+
- validated dense, varlen, and KV-cache forward entry points;
|
|
9
|
+
- explicit validation, dispatch information, and opt-in reference fallback;
|
|
10
|
+
- an M1 Split-KV decode path for C500 paged KV cache;
|
|
11
|
+
- an exact-profile MMA Split-KV path for calibrated FP16 decode shapes;
|
|
12
|
+
- vLLM EngineCore registration and per-dispatch evidence.
|
|
13
|
+
|
|
14
|
+
Native C500 extensions are compiled inside the target MetaX/MACA environment.
|
|
15
|
+
The PyPI artifact supplies the Python API, integration layer, profile handling,
|
|
16
|
+
and registration entry point.
|
|
17
|
+
|
|
18
|
+
## Install
|
|
19
|
+
|
|
20
|
+
```bash
|
|
21
|
+
python -m pip install mxflashattn==1.5.0
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
Base requirements:
|
|
25
|
+
|
|
26
|
+
- Python 3.10+
|
|
27
|
+
- PyTorch 2.4+
|
|
28
|
+
- PyYAML 6.x
|
|
29
|
+
|
|
30
|
+
For the vLLM 0.17 dependency contract:
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
python -m pip install "mxflashattn[vllm]==1.5.0"
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
MetaX deployments normally provide their own compatible PyTorch, vLLM, and
|
|
37
|
+
MXMACA wheels. vLLM 0.17.0 and 0.28.0 were validated on separate C500 software
|
|
38
|
+
stacks; they are not interchangeable environments.
|
|
39
|
+
|
|
40
|
+
## Python API
|
|
41
|
+
|
|
42
|
+
```python
|
|
43
|
+
import os
|
|
44
|
+
import torch
|
|
45
|
+
|
|
46
|
+
from mxflashattn import flash_attn_func, get_last_dispatch_info
|
|
47
|
+
|
|
48
|
+
# Explicitly enable the correctness/reference path when no native provider exists.
|
|
49
|
+
os.environ["MXFLASHATTN_ALLOW_FALLBACK"] = "1"
|
|
50
|
+
|
|
51
|
+
q = torch.randn(1, 32, 8, 64, dtype=torch.float16)
|
|
52
|
+
k = torch.randn(1, 128, 2, 64, dtype=q.dtype)
|
|
53
|
+
v = torch.randn_like(k)
|
|
54
|
+
|
|
55
|
+
out = flash_attn_func(q, k, v, causal=True)
|
|
56
|
+
print(out.shape)
|
|
57
|
+
print(get_last_dispatch_info())
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
The reference path is intended for API checks and correctness comparison. It
|
|
61
|
+
is not a C500 performance backend. Without
|
|
62
|
+
`MXFLASHATTN_ALLOW_FALLBACK=1`, a missing native provider raises an explicit
|
|
63
|
+
error instead of silently selecting a slow path.
|
|
64
|
+
|
|
65
|
+
Public entry points:
|
|
66
|
+
|
|
67
|
+
```python
|
|
68
|
+
from mxflashattn import (
|
|
69
|
+
flash_attn_func,
|
|
70
|
+
flash_attn_varlen_func,
|
|
71
|
+
flash_attn_with_kvcache,
|
|
72
|
+
get_last_dispatch_info,
|
|
73
|
+
)
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
## C500 Native Setup
|
|
77
|
+
|
|
78
|
+
Build native extensions in an environment matching the target MetaX PyTorch
|
|
79
|
+
and MACA installation:
|
|
80
|
+
|
|
81
|
+
```bash
|
|
82
|
+
git clone https://github.com/leo1852098393-blip/MXFlashAttn.git
|
|
83
|
+
cd MXFlashAttn
|
|
84
|
+
python -m pip install --no-build-isolation -e .
|
|
85
|
+
|
|
86
|
+
export MACA_PATH=/opt/maca
|
|
87
|
+
bash csrc_native/build_m1.sh
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
Then enable the guarded M1 route and collect dispatch evidence:
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
export MXFA_M1_DECODE=1
|
|
94
|
+
export MXFLASHATTN_DISPATCH_SAMPLE=1
|
|
95
|
+
|
|
96
|
+
python integrations/vllm/smoke_generate.py \
|
|
97
|
+
--model /path/to/Qwen3-4B \
|
|
98
|
+
--suite --prompt-count 4 --max-tokens 16 \
|
|
99
|
+
--candidate --output artifacts/m1-smoke.json
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
A performance result is attributed to MXFlashAttn only when the associated
|
|
103
|
+
`dispatch.jsonl` contains a successful `m1_split_decode` or
|
|
104
|
+
`mma_split_decode` event.
|
|
105
|
+
|
|
106
|
+
## Verified Scope
|
|
107
|
+
|
|
108
|
+
| Area | v1.5.0 status |
|
|
109
|
+
| --- | --- |
|
|
110
|
+
| M1 decode | FP16 / BF16, head dimensions 64 / 128, GQA ratios 2 / 4 / 8 |
|
|
111
|
+
| vLLM 0.17.0 | Real Qwen3-4B/8B generation, M1 dispatch, fallback, and long-decode measurements |
|
|
112
|
+
| vLLM 0.28.0 | EngineCore plugin registration, real M1 dispatch, and a Qwen3-8B shape matrix |
|
|
113
|
+
| MMA decode | FP16 single-token decode, exact calibrated profiles only |
|
|
114
|
+
| Prefill | MetaX vendor production path; chunked prefill remains a correctness reference |
|
|
115
|
+
| Backward | Explicit MetaX provider integration; no in-house backward kernel |
|
|
116
|
+
|
|
117
|
+
In four vLLM 0.17.0 long-decode workloads, M1 delivered 2.67-3.54x the
|
|
118
|
+
throughput of the project's historical v7 path and reached 79.2%-86.5% of
|
|
119
|
+
MetaX vendor throughput. In one exact Qwen3-4B operator shape, MMA median
|
|
120
|
+
latency was `43.520 us` versus `46.592 us` for vendor over 100 rounds.
|
|
121
|
+
|
|
122
|
+
In the reported
|
|
123
|
+
[vLLM 0.28.0 Qwen3-8B matrix](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/V1_5_028_MATRIX.md),
|
|
124
|
+
long-KV samples at batch 16 or larger remained 28%-46% below MetaX vendor. No
|
|
125
|
+
general M1-over-vendor crossover is claimed.
|
|
126
|
+
|
|
127
|
+
These results are bound to their recorded device, software stack, model,
|
|
128
|
+
shape, and baseline. The MMA comparison is operator-level, not an end-to-end
|
|
129
|
+
request-speedup claim.
|
|
130
|
+
|
|
131
|
+
## Runtime Boundaries
|
|
132
|
+
|
|
133
|
+
- Prefill and unprofiled decode shapes retain the MetaX vendor path.
|
|
134
|
+
- M1 requires its documented dtype, layout, head-dimension, and GQA contract.
|
|
135
|
+
- MMA requires an exact profile match, including KV heads, page size, binary
|
|
136
|
+
variant, and shared-object identity.
|
|
137
|
+
- Short or low-concurrency requests may not benefit from M1.
|
|
138
|
+
- FP8 KV cache, multi-card execution, SGLang, GQ=7, and general TTFT gains are
|
|
139
|
+
not v1.5.0 support claims.
|
|
140
|
+
- Unsupported API options fail explicitly or follow a documented fallback.
|
|
141
|
+
|
|
142
|
+
## Development and Audit
|
|
143
|
+
|
|
144
|
+
The project's own kernel implementations under `csrc_native/` were written by
|
|
145
|
+
AI coding agents. The project owner set the project goals, verification rules,
|
|
146
|
+
and release boundary, and made the final accept-or-reject decisions for code
|
|
147
|
+
changes and public claims.
|
|
148
|
+
|
|
149
|
+
Favorable measurements do not automatically become performance claims. The
|
|
150
|
+
short-KV `+15.2%` difference was the same order as repeated-run variance; the
|
|
151
|
+
long-KV `+22.7%` result was a single observation and does not establish a
|
|
152
|
+
general crossover. Neither is used for a general performance claim.
|
|
153
|
+
|
|
154
|
+
## Documentation
|
|
155
|
+
|
|
156
|
+
- [GitHub repository](https://github.com/leo1852098393-blip/MXFlashAttn)
|
|
157
|
+
- [GitLink mirror](https://gitlink.org.cn/Leo77/MXFlashAttn)
|
|
158
|
+
- [API guide](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/API.md)
|
|
159
|
+
- [Support matrix](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/support-matrix.md)
|
|
160
|
+
- [Evidence index](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/EVIDENCE.md)
|
|
161
|
+
- [Release validation](https://github.com/leo1852098393-blip/MXFlashAttn/blob/main/docs/V1_5_RELEASE_VALIDATION.md)
|
|
162
|
+
|
|
163
|
+
## License
|
|
164
|
+
|
|
165
|
+
Apache License 2.0. MXMACA runtime, MetaX vendor packages, vLLM, model
|
|
166
|
+
weights, and datasets remain subject to their respective licenses.
|