macac fa impl

dlblas/kernels/fa_macac/README.md

@@ -0,0 +1,132 @@
+# fa_hdim128 — Fully-Builtin Flash-Attention Forward (MetaX C500 / xcore1000)
+
+A from-scratch reimplementation of dense flash-attention **forward** (D=128, fp16,
+non-causal, even-MN/K) that depends **only on the official MACA system SDK**
+(cute + mctlass via the `MACA_PATH` install root) plus the headers in this
+directory. It does **not** reference any external source tree (no csrc copy /
+`/home/...` paths).
+
+- Algorithm: online-softmax + tiled MMA (`MACA_16x16x16_F32F16F16F32` atom) +
+  swizzled shared memory + Q-in-regs + K reg-staged prefetch + direct-V gmem→smem.
+- Config: B=1, H=32, D=128, blockM=128, blockN=64, 4 warps (256 threads, wave=64).
+- Correctness: element-wise `allclose` (atol 1e-2) vs the torch flash_attn wheel.
+- This is the kernel body from the project's Round-8 async-hybrid; only the
+  traits / utils / softmax / params headers were reauthored on top of the
+  system cute/mctlass (the fa_src copies were dropped).
+
+## Directory layout
+
+```
+builtin/
+├── build.sh                 # self-locating build (uses $MACA_PATH)
+├── README.md
+├── src/
+│   ├── fa_my.cu             # harness: alloc / init / time / dump, calls the kernel
+│   ├── my_compute.cuh       # kernel body: compute_attn_myimpl + my_flash_fwd_kernel
+│   └── builtin/
+│       ├── fa_params.h      # Flash_fwd_params struct (namespace mcFlashAttn)
+│       ├── fa_traits.cuh     # concrete Flash_fwd_kernel_traits<128,128,64,4,true,true,half_t,128>
+│       ├── fa_utils.cuh      # flash::{gemm,gemm_rs,copy_*,clear,barrier*,softmax helpers,...}
+│       └── fa_softmax.cuh    # flash::Softmax (softmax_rescale_o + normalize_softmax_lse)
+└── tests/
+    ├── cmp_my.py             # correctness vs torch _flash_attn_forward
+    ├── bench_torch.py        # torch flash_attn timing (single run)
+    └── bench_torch_multi.py  # torch flash_attn timing (min/med/max, 5 runs)
+```
+
+## Prerequisites
+
+- MetaX MACA SDK installed (provides cute, mctlass, cu-bridge `cucc`, runtime libs).
+  The install root is exposed via the **`MACA_PATH`** env var
+  (e.g. `export MACA_PATH=/opt/maca`). It must contain:
+  - `$MACA_PATH/include` — cute + mctlass headers
+  - `$MACA_PATH/tools/cu-bridge/bin/cucc` — the compiler
+  - `$MACA_PATH/tools/cu-bridge/include` — CUDA shim headers
+  - `$MACA_PATH/lib` — runtime libs (`libmcruntime`, `libmxc-runtime64`, …)
+- A MetaX C500 (xcore1000) device.
+- For testing only: torch + the `flash_attn` wheel (`_flash_attn_forward`).
+
+> In the `metax_gemm_opt` container, `MACA_PATH=/opt/maca` is already set.
+
+## How to compile
+
+```bash
+export MACA_PATH=/opt/maca      # your MACA install root
+bash build.sh
+# -> produces ./fa_my_builtin
+```
+
+`build.sh` is self-locating (finds `./src` relative to itself) and reads the MACA
+SDK only from `$MACA_PATH`. No hardcoded absolute paths.
+
+## How to test
+
+The binary takes: `./fa_my_builtin <S> <warmup> <iters> <dump>`
+(defaults: S=512, warmup=5, iters=30, dump=0). When `dump=1` it writes
+`$FA_DUMP_DIR/fa_my_<S>.bin` (defaults to `./` if `FA_DUMP_DIR` unset) containing
+B,H,S,D + Q,K,V,O in fp16.
+
+### Correctness (vs torch flash_attn)
+
+```bash
+export FA_DUMP_DIR=/tmp        # optional; default is current dir
+./fa_my_builtin 512 5 30 1     # dump
+./fa_my_builtin 1024 5 30 1
+python tests/cmp_my.py         # reads $FA_DUMP_DIR/fa_my_{512,1024}.bin, compares to torch
+# expect: allclose(1e-2)=True, max_diff < 0.01
+```
+
+### Performance vs torch
+
+```bash
+# Our kernel (median over iters):
+for S in 512 1024 10240 102400; do ./fa_my_builtin $S 5 30 0; done
+
+# torch flash_attn API (median of 5 runs):
+python tests/bench_torch.py            # S in {512,1024,10240,102400}
+python tests/bench_torch_multi.py      # min/med/max for small/medium shapes
+```
+
+## Measured results (MetaX C500, 2026-07-22, container metax_gemm_opt)
+
+| Shape (1x32xSx128) | fa_my_builtin (ms) | torch API median (ms) |
+|---|---|---|
+| 512    | 0.0826  | 0.0666 |
+| 1024   | 0.2266  | 0.1739 |
+| 10240  | 16.052  | 17.86  |
+| 102400 | 1564.0  | 1195.2 |
+
+Correctness: `allclose(1e-2)=True` at S=512 (max_diff≈0.0065) and S=1024 (≈0.0042).
+
+## Notes
+
+- All source/build/test files are path-free (portable). System SDK location comes
+  from `$MACA_PATH`; dump location from `$FA_DUMP_DIR` (default `./`).
+- The kernel reuses no external flash_attn source headers — traits / utils /
+  softmax / params are reauthored here on top of the official cute/mctlass.
+  The one fa_src-local cute extension (`get_swizzle_offset`) is avoided by using
+  the standard `Swizzle<3,4,3>` + `composition` (the kernel never needed it).
+
+
+====
+
+  编译 / 测试（README 已详述）
+
+  export MACA_PATH=/opt/maca     # MACA SDK 根 (容器已设)
+  bash build.sh                  # → ./fa_my_builtin
+
+  export FA_DUMP_DIR=/tmp        # 可选, 默认 ./
+  ./fa_my_builtin 512 5 30 1     # dump
+  ./fa_my_builtin 1024 5 30 1
+  python tests/cmp_my.py         # 正确性 vs torch, 期望 allclose(1e-2)=True
+
+  for S in 512 1024 10240 102400; do ./fa_my_builtin $S 5 30 0; done   # 计时
+  python tests/bench_torch.py                                          # torch 对照
+
+  验证结果（scrub 后临时副本实测，已清理）
+
+  - 编译：BUILD_OK, Exit 0
+  - 正确性：S=512 max_diff=0.0065、S=1024 max_diff=0.0042，allclose(1e-2)=True ✅
+  - 计时：512→0.0825 / 1024→0.2267 / 10240→16.052 / 1024
+
+====

dlblas/kernels/fa_macac/build.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+# build.sh — compile fa_my_builtin (fully-builtin flash-attn fwd, hdim128).
+#
+# Self-locating & portable: no hardcoded absolute paths.
+#   - Our own sources are found relative to this script (./src).
+#   - The MACA system SDK is located via the MACA_PATH env var (must point to
+#     the official MACA install root, i.e. the dir containing include/,
+#     tools/, and lib/).
+set -euo pipefail
+
+: "${MACA_PATH:?Error: set MACA_PATH to your MACA SDK install root (the dir containing include/, tools/, lib/)}"
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+SRC_DIR="$SCRIPT_DIR/src"
+OUT="$SCRIPT_DIR/fa_my_builtin"
+
+CXX="$MACA_PATH/tools/cu-bridge/bin/cucc"
+
+D="-DFA_DTYPE_FP16 -DHDIM_128 -DHDIM_CONFIG=128 -DUSE_MACA -DXCORE1000 -D__FAST_HALF_CVT__ -D__MERGE_LDS_B64"
+I="-I$SRC_DIR -I$MACA_PATH/include -I$MACA_PATH/tools/cu-bridge/include"
+F="-w -Wno-format -x maca --compiler-options -fPIC -O3 -std=c++17 --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -fno-strict-aliasing -mllvm -metaxgpu-inlinescope=50 -mllvm -metaxgpu-disable-bsm-offset=0 -Xclang -menable-no-nans -gencode arch=compute_80,code=sm_80 -gencode arch=compute_90,code=sm_90 -mllvm -metaxgpu-enable-promote-kernel-arguments -mllvm -metaxgpu-const-KArgs-LDU=1 -mllvm -metaxgpu-enable-gvn-loadclobber=false -mllvm -metaxgpu-enable-unorder-dispatch -mllvm -metaxgpu-disable-early-vector-combine=true -mllvm -metaxgpu-llvm19-inline-max-bb=2000 -mllvm -metaxgpu-preisel-simplifycfg=true --offload-arch=xcore1000"
+L="-L$MACA_PATH/lib -lmcruntime -lmxc-runtime64 -lcurand -lmcToolsExt -lcudart"
+
+echo "Compiling fa_my_builtin (MACA_PATH=$MACA_PATH, SRC=$SRC_DIR)..."
+"$CXX" $D $I $F $L -o "$OUT" "$SRC_DIR/fa_my.cu" 2>&1 | tail -50
+echo "Exit: ${PIPESTATUS[0]}"
+
+if [ -f "$OUT" ]; then
+  echo "BUILD_OK -> $OUT"
+else
+  echo "BUILD_FAIL"; exit 1
+fi

dlblas/kernels/fa_macac/src/builtin/fa_params.h

@@ -0,0 +1,126 @@
+#pragma once
+// fa_params.h — Flash attention forward params (reauthored, builtin).
+// Plain data struct; field layout mirrors the reference (Tri Dao flash_attn)
+// so our kernel body is unchanged. Only includes <cstdint> — no external deps.
+// Namespace mcFlashAttn kept so the harness (fa_my.cu) is unchanged.
+
+#include <cstdint>
+
+namespace mcFlashAttn {
+
+struct Qkv_params {
+    using index_t = int64_t;
+
+    void *__restrict__ q_ptr;
+    void *__restrict__ k_ptr;
+    void *__restrict__ v_ptr;
+
+    index_t q_batch_stride;
+    index_t k_batch_stride;
+    index_t v_batch_stride;
+    index_t q_row_stride;
+    index_t k_row_stride;
+    index_t v_row_stride;
+    index_t q_head_stride;
+    index_t k_head_stride;
+    index_t v_head_stride;
+
+    int h, h_k;
+    int h_h_k_ratio;
+};
+
+struct Flash_fwd_params : public Qkv_params {
+    void * __restrict__ o_ptr;
+    void * __restrict__ oaccum_ptr;
+
+    index_t o_batch_stride;
+    index_t o_row_stride;
+    index_t o_head_stride;
+
+    void * __restrict__ p_ptr;
+    void * __restrict__ softmax_lse_ptr;
+    void * __restrict__ softmax_lseaccum_ptr;
+    void * __restrict__ max_logit_ptr;
+
+    int b, seqlen_q, seqlen_k, seqlen_knew, d, seqlen_q_rounded, seqlen_k_rounded, d_rounded, rotary_dim, total_q;
+    uint32_t ngroups;
+
+    float scale_softmax;
+    float scale_softmax_log2;
+
+    int * __restrict__ cu_seqlens_q;
+    int * __restrict__ cu_seqlens_k;
+    int * __restrict__ leftpad_k;
+    int * __restrict__ seqused_k;
+    int *__restrict__ blockmask;
+
+    void * __restrict__ knew_ptr;
+    void * __restrict__ vnew_ptr;
+    index_t knew_batch_stride;
+    index_t vnew_batch_stride;
+    index_t knew_row_stride;
+    index_t vnew_row_stride;
+    index_t knew_head_stride;
+    index_t vnew_head_stride;
+
+    index_t kscale_batch_stride;
+    index_t vscale_batch_stride;
+    index_t kscale_row_stride;
+    index_t vscale_row_stride;
+    index_t kscale_head_stride;
+    index_t vscale_head_stride;
+
+    void * __restrict__ rotary_cos_ptr;
+    void * __restrict__ rotary_sin_ptr;
+    int * __restrict__ cache_batch_idx;
+    int * __restrict__ block_table;
+    index_t block_table_batch_stride;
+    int page_block_size;
+    int dequant_group;
+    void *__restrict__ k_scale_ptr;
+    void *__restrict__ v_scale_ptr;
+
+    float p_dropout;
+    uint8_t p_dropout_in_uint8_t;
+    float rp_dropout;
+    float scale_softmax_rp_dropout;
+
+    int window_size_left, window_size_right;
+    float softcap;
+
+    uint64_t rng_state_seed = 0;
+    uint64_t rng_state_offset = 0;
+
+    bool is_bf16;
+    bool is_causal;
+    bool is_seqlens_k_cumulative;
+    bool is_rotary_interleaved;
+
+    int num_splits;
+    void * __restrict__ alibi_slopes_ptr;
+    index_t alibi_slopes_batch_stride;
+    bool custom_alibi = false;
+
+    bool has_attn_mask;
+    void * __restrict__ attn_mask_ptr = nullptr;
+    index_t attn_mask_batch_stride = 0;
+    index_t attn_mask_nheads_stride = 0;
+    index_t attn_mask_row_stride = 0;
+    index_t attn_mask_col_stride = 1;
+    index_t attn_mask_batch_shape = 1;
+    index_t attn_mask_nheads_shape = 1;
+    index_t attn_mask_row_shape = 1;
+    index_t attn_mask_col_shape = 1;
+
+    bool unpadded_lse;
+    bool seqlenq_ngroups_swapped;
+
+    int d_value;
+    int d_value_rounded;
+    bool is_support_splitkv = false;
+    int arch;
+
+    void *__restrict__ s_aux_ptr;
+};
+
+} // namespace mcFlashAttn