From f8e5403958381b903d18dd10e7870b3f383c7207 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 14:25:52 +0800 Subject: [PATCH 01/10] refactor(kda): remove native CUDA extension Signed-off-by: shuyan.ycf --- .coderabbit.yaml | 7 - .github/workflows/build-release.yml | 122 +- .github/workflows/lint.yml | 11 - .gitmodules | 3 - .pre-commit-config.yaml | 8 - MANIFEST.in | 2 - README.md | 20 +- RECOMMENDED_CODING_STYLE.md | 260 +- REPO_LAYOUT.md | 31 +- USAGE.md | 67 +- benchmarks/bench_intracard_cp_sm90.py | 448 ---- csrc/api/kda_sm90.cu | 226 -- csrc/api/pybind.cu | 62 - csrc/api/pybind_sm90.cu | 1 - csrc/cutlass | 1 - csrc/kda/sm90/collective/common.hpp | 477 ---- csrc/kda/sm90/collective/load_predicated.hpp | 193 -- csrc/kda/sm90/collective/load_tma.hpp | 177 -- csrc/kda/sm90/collective/mainloop_kda_fwd.hpp | 2101 ----------------- csrc/kda/sm90/collective/named_barriers.hpp | 44 - csrc/kda/sm90/collective/store_tma.hpp | 313 --- csrc/kda/sm90/device/device_universal.hpp | 244 -- csrc/kda/sm90/kda_fwd_sm90.cu | 189 -- csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu | 66 - csrc/kda/sm90/kernel/builder_kda_fwd.hpp | 80 - csrc/kda/sm90/kernel/kernel_kda_fwd.hpp | 639 ----- csrc/kda/sm90/kernel/options.hpp | 87 - csrc/kda/sm90/kernel/tile_scheduler.hpp | 155 -- csrc/kda/sm90/prefill_kernel.hpp | 54 - csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh | 169 -- csrc/kda/sm90/utils/common.hpp | 59 - csrc/kda/sm90/utils/debug.hpp | 149 -- csrc/kda/sm90/utils/math.hpp | 53 - csrc/kda/sm90/utils/math_order_barrier.hpp | 116 - csrc/kda/sm90/utils/type_traits.hpp | 66 - csrc/kda/sm90/utils/unused.hpp | 54 - .../kerutils/include/kerutils/common/common.h | 35 - .../include/kerutils/common/cute_ext.hpp | 77 - .../include/kerutils/common/static_switch.h | 78 - .../kerutils/include/kerutils/device/common.h | 98 - .../include/kerutils/device/device.cuh | 25 - .../include/kerutils/device/sm100/helpers.cuh | 116 - .../kerutils/device/sm100/intrinsics.cuh | 239 -- .../device/sm80/collective_inverse.hpp | 1145 --------- .../include/kerutils/device/sm80/helpers.cuh | 123 - .../include/kerutils/device/sm90/helpers.cuh | 151 -- csrc/kerutils/include/kerutils/host/host.h | 107 - csrc/kerutils/include/kerutils/kerutils.cuh | 18 - cula/backends.py | 86 - cula/cudac.py | 103 - cula/kda/README.md | 8 +- cula/kda/__init__.py | 8 +- cula/kda/auto_route.py | 113 - cula/kda/backends/flashkda.py | 73 - cula/kda/backends/fully_fused.py | 67 - cula/kda/cp_context.py | 452 ---- cula/kda/cp_h_boundary.py | 190 -- cula/kda/gate_l2norm_fused.py | 194 -- cula/kda/hopper_fused_fwd.py | 281 --- cula/kda/hopper_fused_fwd_opt.py | 391 --- cula/kda/l2norm_qk_fused.py | 119 - cula/kda/{backends/__init__.py => prefill.py} | 25 +- cula/ops/kda/cp_mode.py | 2 +- docs/kda_sm100_cutedsl_fwd_status.md | 8 +- pyproject.toml | 11 +- scripts/build_wheel.sh | 10 +- setup.py | 157 -- tests/test_intracard_cp_sm90.py | 439 ---- tests/test_kda_prefill_api.py | 61 + tests/test_kda_sm100_output.py | 2 +- tests/test_kda_sm100_recurrence.py | 2 +- tests/test_kda_sm90_intracard_cp.py | 4 +- 72 files changed, 208 insertions(+), 11564 deletions(-) delete mode 100644 benchmarks/bench_intracard_cp_sm90.py delete mode 100644 csrc/api/kda_sm90.cu delete mode 100644 csrc/api/pybind.cu delete mode 100644 csrc/api/pybind_sm90.cu delete mode 160000 csrc/cutlass delete mode 100644 csrc/kda/sm90/collective/common.hpp delete mode 100644 csrc/kda/sm90/collective/load_predicated.hpp delete mode 100644 csrc/kda/sm90/collective/load_tma.hpp delete mode 100644 csrc/kda/sm90/collective/mainloop_kda_fwd.hpp delete mode 100644 csrc/kda/sm90/collective/named_barriers.hpp delete mode 100644 csrc/kda/sm90/collective/store_tma.hpp delete mode 100644 csrc/kda/sm90/device/device_universal.hpp delete mode 100644 csrc/kda/sm90/kda_fwd_sm90.cu delete mode 100644 csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu delete mode 100644 csrc/kda/sm90/kernel/builder_kda_fwd.hpp delete mode 100644 csrc/kda/sm90/kernel/kernel_kda_fwd.hpp delete mode 100644 csrc/kda/sm90/kernel/options.hpp delete mode 100644 csrc/kda/sm90/kernel/tile_scheduler.hpp delete mode 100644 csrc/kda/sm90/prefill_kernel.hpp delete mode 100644 csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh delete mode 100644 csrc/kda/sm90/utils/common.hpp delete mode 100644 csrc/kda/sm90/utils/debug.hpp delete mode 100644 csrc/kda/sm90/utils/math.hpp delete mode 100644 csrc/kda/sm90/utils/math_order_barrier.hpp delete mode 100644 csrc/kda/sm90/utils/type_traits.hpp delete mode 100644 csrc/kda/sm90/utils/unused.hpp delete mode 100644 csrc/kerutils/include/kerutils/common/common.h delete mode 100644 csrc/kerutils/include/kerutils/common/cute_ext.hpp delete mode 100644 csrc/kerutils/include/kerutils/common/static_switch.h delete mode 100644 csrc/kerutils/include/kerutils/device/common.h delete mode 100644 csrc/kerutils/include/kerutils/device/device.cuh delete mode 100644 csrc/kerutils/include/kerutils/device/sm100/helpers.cuh delete mode 100644 csrc/kerutils/include/kerutils/device/sm100/intrinsics.cuh delete mode 100644 csrc/kerutils/include/kerutils/device/sm80/collective_inverse.hpp delete mode 100644 csrc/kerutils/include/kerutils/device/sm80/helpers.cuh delete mode 100644 csrc/kerutils/include/kerutils/device/sm90/helpers.cuh delete mode 100644 csrc/kerutils/include/kerutils/host/host.h delete mode 100644 csrc/kerutils/include/kerutils/kerutils.cuh delete mode 100644 cula/backends.py delete mode 100644 cula/cudac.py delete mode 100644 cula/kda/auto_route.py delete mode 100644 cula/kda/backends/flashkda.py delete mode 100644 cula/kda/backends/fully_fused.py delete mode 100644 cula/kda/cp_context.py delete mode 100644 cula/kda/cp_h_boundary.py delete mode 100644 cula/kda/gate_l2norm_fused.py delete mode 100644 cula/kda/hopper_fused_fwd.py delete mode 100644 cula/kda/hopper_fused_fwd_opt.py delete mode 100644 cula/kda/l2norm_qk_fused.py rename cula/kda/{backends/__init__.py => prefill.py} (78%) delete mode 100644 tests/test_intracard_cp_sm90.py create mode 100644 tests/test_kda_prefill_api.py diff --git a/.coderabbit.yaml b/.coderabbit.yaml index b734bbfe..121a1c52 100644 --- a/.coderabbit.yaml +++ b/.coderabbit.yaml @@ -16,16 +16,9 @@ reviews: mode: "off" path_filters: - - "!csrc/cutlass/**" - "!third_party/**" path_instructions: - - path: "csrc/**" - instructions: | - Review CUDA changes for synchronization and memory-ordering correctness, - shared-memory and barrier lifetimes, out-of-bounds access, divergence, - architecture guards, numerical precision, and concrete performance regressions. - - path: "cula/**" instructions: | Review CuTe DSL kernels and Python APIs for architecture dispatch, diff --git a/.github/workflows/build-release.yml b/.github/workflows/build-release.yml index 73ea4929..673ad4df 100644 --- a/.github/workflows/build-release.yml +++ b/.github/workflows/build-release.yml @@ -1,4 +1,4 @@ -name: Build & Release Wheels +name: Build & Release Wheel on: push: @@ -12,127 +12,50 @@ concurrency: jobs: build-wheel: - name: "wheel / ${{ matrix.cuda }} / cp312 / ${{ matrix.arch }}" - runs-on: ${{ matrix.arch == 'aarch64' && 'ubuntu-24.04-arm' || 'ubuntu-latest' }} - defaults: - run: - shell: bash - strategy: - fail-fast: false - matrix: - cuda: - - cu129 - - cu130 - arch: - - x86_64 - - aarch64 - container: - # UBI 8 provides the glibc 2.28 baseline required by manylinux_2_28. - image: "nvidia/cuda:${{ matrix.cuda == 'cu129' && '12.9.0' || '13.0.0' }}-devel-ubi8" - + name: wheel / py3-none-any + runs-on: ubuntu-latest steps: - - name: Free disk space - run: | - rm -rf /opt/hostedtoolcache /usr/local/lib/android /usr/share/dotnet \ - /usr/local/share/boost /opt/ghc 2>/dev/null || true - dnf clean all 2>/dev/null || true - df -h / || true - - - name: Install system dependencies - run: | - dnf install -y \ - git \ - gcc-toolset-13-gcc \ - gcc-toolset-13-gcc-c++ \ - python3.12 \ - python3.12-devel \ - python3.12-pip - dnf clean all - - name: Checkout uses: actions/checkout@v5 with: fetch-depth: 0 - submodules: recursive - - - name: Configure git safe directory - run: git config --global --add safe.directory "$GITHUB_WORKSPACE" - - name: Install Python dependencies - run: | - python3.12 -m pip install --no-cache-dir --upgrade pip - python3.12 -m pip install --no-cache-dir torch==2.9.1 --index-url ${{ matrix.cuda == 'cu129' && 'https://download.pytorch.org/whl/cu129' || 'https://download.pytorch.org/whl/cu130' }} - python3.12 -m pip install --no-cache-dir setuptools wheel "setuptools_scm>=6.0" build ninja auditwheel patchelf + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" - - name: Expose PyTorch CUDA libraries - # The nvidia/cuda devel image ships the CUDA toolkit but not cuDNN. - # torch's pip wheels bundle cuDNN inside site-packages/nvidia/*/lib/, - # which is not on the linker search path — add it so torch._C can load. - run: | - LIB_DIRS="$(python3.12 <<'EOF' - import glob, os, site - dirs = [] - for sp in site.getsitepackages(): - for p in [os.path.join(sp, "torch", "lib")] + glob.glob(os.path.join(sp, "nvidia", "*", "lib")): - if os.path.isdir(p): - dirs.append(p) - print(":".join(dirs)) - EOF - )" - echo "LD_LIBRARY_PATH=${LIB_DIRS}${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" >> "$GITHUB_ENV" + - name: Install build tools + run: python -m pip install --upgrade build setuptools wheel setuptools_scm - name: Compute version id: version + shell: bash run: | if [[ "$GITHUB_REF" == refs/tags/v* ]]; then BASE="${GITHUB_REF#refs/tags/v}" else - # Strip any local segment (+gXXX) so we get a clean base - BASE=$(python3.12 -c "from setuptools_scm import get_version; print(get_version().split('+')[0])") + BASE=$(python -c "from setuptools_scm import get_version; print(get_version().split('+')[0])") fi - echo "version=${BASE}+${{ matrix.cuda }}" >> "$GITHUB_OUTPUT" + echo "version=${BASE}" >> "$GITHUB_OUTPUT" - - name: Build fat-binary wheel + - name: Build wheel env: - CC: /opt/rh/gcc-toolset-13/root/usr/bin/gcc - CXX: /opt/rh/gcc-toolset-13/root/usr/bin/g++ - CUDAHOSTCXX: /opt/rh/gcc-toolset-13/root/usr/bin/g++ - CULA_BUILD_ALL_ARCHS: "1" SETUPTOOLS_SCM_PRETEND_VERSION: "${{ steps.version.outputs.version }}" - NVCC_THREADS: "4" - MAX_JOBS: "4" - run: | - "$CC" --version - "$CXX" --version - python3.12 -m build --wheel --no-isolation --outdir dist-raw - - - name: Repair wheel for manylinux_2_28 - run: | - # These libraries are supplied by the NVIDIA driver, PyTorch, or - # PyTorch's CUDA runtime dependency and must remain external. - python3.12 -m auditwheel repair \ - --plat manylinux_2_28_${{ matrix.arch }} \ - --exclude libcuda.so.1 \ - --exclude 'libcudart.so.*' \ - --exclude 'libc10*.so' \ - --exclude 'libtorch*.so' \ - --wheel-dir dist \ - dist-raw/*.whl + run: python -m build --wheel --outdir dist - name: Verify wheel + shell: bash run: | - echo "Built wheel:" - ls -lh dist/*.whl - ls dist/*.whl | grep -q "+${{ matrix.cuda }}" \ - || { echo "ERROR: wheel name missing +${{ matrix.cuda }} suffix"; exit 1; } - ls dist/*.whl | grep -q "manylinux_2_28_${{ matrix.arch }}" \ - || { echo "ERROR: wheel is not tagged manylinux_2_28_${{ matrix.arch }}"; exit 1; } - python3.12 -m auditwheel show dist/*.whl + python -m pip install --no-deps dist/*.whl + python -c "import cula; print(cula.__version__)" + ls dist/*.whl | grep -q -- '-py3-none-any.whl$' \ + || { echo "ERROR: wheel is not platform-independent"; exit 1; } - name: Upload wheel artifact uses: actions/upload-artifact@v6 with: - name: wheel-${{ matrix.cuda }}-${{ matrix.arch }} + name: wheel-py3-none-any path: dist/*.whl release: @@ -143,7 +66,7 @@ jobs: permissions: contents: write steps: - - name: Download all artifacts + - name: Download wheel uses: actions/download-artifact@v6 with: path: artifacts/ @@ -151,8 +74,7 @@ jobs: - name: Create release uses: softprops/action-gh-release@v3 with: - files: | - artifacts/wheel-*/*.whl + files: artifacts/wheel-py3-none-any/*.whl generate_release_notes: true draft: true prerelease: ${{ contains(github.ref, 'rc') || contains(github.ref, 'beta') || contains(github.ref, 'alpha') }} diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 82daf202..854e9304 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -13,17 +13,6 @@ concurrency: cancel-in-progress: true jobs: - clang-format: - name: clang-format - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v6 - - uses: j178/prek-action@v2 - with: - install-only: true - - name: Check clang-format - run: prek run clang-format --all-files --show-diff-on-failure --color=always - ruff: name: Ruff runs-on: ubuntu-latest diff --git a/.gitmodules b/.gitmodules index 1e7317f1..8c340d6d 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,6 +1,3 @@ -[submodule "csrc/cutlass"] - path = csrc/cutlass - url = https://github.com/NVIDIA/cutlass.git [submodule "third_party/flash-linear-attention"] path = third_party/flash-linear-attention url = https://github.com/fla-org/flash-linear-attention.git diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f0f30e7f..0911231c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,12 +1,4 @@ repos: - - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v15.0.7 - hooks: - - id: clang-format - files: ^csrc/ - types_or: [c, c++, cuda] - exclude: ^csrc/(cutlass|kerutils)/ - - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.15.0 hooks: diff --git a/MANIFEST.in b/MANIFEST.in index 688580ce..a9f4b91f 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -3,6 +3,4 @@ include LICENSE include pyproject.toml include setup.py -recursive-include csrc *.h *.hpp *.cuh *.cu *.cpp *.cc -recursive-include csrc/cutlass * recursive-include cula *.py diff --git a/README.md b/README.md index aeec8b17..93b4df22 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ # cuLA — CUDA Linear Attention -**High-performance CUDA kernels for linear attention variants, written in [CuTe DSL](https://github.com/NVIDIA/cutlass/tree/main/python/CuTeDSL) and CUTLASS C++.** +**High-performance CUDA kernels for linear attention variants, written entirely in [CuTe DSL](https://github.com/NVIDIA/cutlass/tree/main/python/CuTeDSL).** @@ -20,17 +20,17 @@ Linear attention mechanisms reformulate standard attention to use linear-time st cuLA supports both **Hopper (SM90)** and **Blackwell (SM10X)** GPUs. -> **Requirements (Hopper & Blackwell):** Python 3.12+, CUDA Toolkit 12.9+, PyTorch 2.9.1+. Building the optional SM90 native extension also requires NVCC. +> **Requirements (Hopper & Blackwell):** Python 3.12+, CUDA Toolkit 12.9+, PyTorch 2.9.1+, and a matching `nvidia-cutlass-dsl` installation. cuLA does not build a C++/CUDA extension. > **Note:** The PyTorch CUDA version must match your system CUDA Toolkit version. Check with `nvcc --version` and `python -c "import torch; print(torch.version.cuda)"`. ### Pre-built Wheels -Pre-built wheels with the SM90 native extension and CuTeDSL JIT kernels for SM100/SM103 are available on [GitHub Releases](https://github.com/inclusionAI/cuLA/releases). Linux wheels target `manylinux_2_28` and require glibc 2.28 or newer: +Platform-independent wheels containing the CuTeDSL sources are available on [GitHub Releases](https://github.com/inclusionAI/cuLA/releases): - pip install "cuda-linear-attention==+" -f https://github.com/inclusionAI/cuLA/releases/expanded_assets/ + pip install "cuda-linear-attention==" -f https://github.com/inclusionAI/cuLA/releases/expanded_assets/ -Replace `` with the release tag (e.g., `v0.2.0`), `` with the base version (e.g., `0.2.0`), and `` with your PyTorch CUDA build tag (e.g., `cu129` or `cu130`). Or download the `.whl` file directly from the [Releases page](https://github.com/inclusionAI/cuLA/releases) and install it with `pip install .whl`. +Replace `` with the release tag (e.g., `v0.2.0`) and `` with the version (e.g., `0.2.0`). Or download the `.whl` file directly from the [Releases page](https://github.com/inclusionAI/cuLA/releases) and install it with `pip install .whl`. ### Build from Source @@ -47,6 +47,12 @@ git submodule update --init --recursive pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu129 ``` +For CUDA 13, install the CuTeDSL CUDA 13 extra before installing cuLA: + +```bash +pip install 'nvidia-cutlass-dsl[cu13]>=4.4.2,<4.7,!=4.5.0' +``` + **Install cuLA & dependencies:** ```bash @@ -57,10 +63,10 @@ pip install -e third_party/flash-linear-attention pip install -e . --no-build-isolation ``` -**Build a wheel with the SM90 extension (SM100/SM103 use CuTeDSL JIT):** +**Build the platform-independent wheel:** ```bash -CULA_BUILD_ALL_ARCHS=1 python -m build --wheel --no-isolation +python -m build --wheel ``` ## Quick Start diff --git a/RECOMMENDED_CODING_STYLE.md b/RECOMMENDED_CODING_STYLE.md index a6090e4a..ddb137ca 100644 --- a/RECOMMENDED_CODING_STYLE.md +++ b/RECOMMENDED_CODING_STYLE.md @@ -1,199 +1,65 @@ # Coding Style Guide -This document outlines the coding style and conventions used in this project, which is primarily written in CUDA C++ with heavy use of CUTLASS and CUTE libraries. - -## Table of Contents -1. [General Principles](#general-principles) -2. [File Organization](#file-organization) -3. [Naming Conventions](#naming-conventions) -4. [Formatting](#formatting) -5. [CUDA-Specific Conventions](#cuda-specific-conventions) -6. [Comments](#comments) -7. [Templates and Generics](#templates-and-generics) -8. [Error Handling](#error-handling) - -## General Principles - -- Follow modern C++ standards (C++14/17) combined with CUDA extensions -- Prioritize performance and readability -- Use CUTLASS and CUTE libraries for tensor operations -- Maintain consistency with NVIDIA's CUTLASS codebase style - -## File Organization - -- Use hierarchical directory structure: `csrc/sm{xx,90,100}/decode|prefill/dense|sparse/` -- Group related files in subdirectories (e.g., `instantiations/`, `common/`) -- Separate header files (`.h`, `.cuh`) from implementation files (`.cu`, `.cpp`) -- Use `#include` guards or `#pragma once` in headers - -## Naming Conventions - -### Variables -- Use `snake_case` for local variables and function parameters -- Use descriptive names that indicate purpose -- Examples: - ```cpp - int batch_idx = blockIdx.x; - int warp_idx = threadIdx.x / 32; - float max_lse = -INFINITY; - ``` - -### Tensor Variables (CUTE Library) -- For CUTE `Tensor` objects, use abbreviated mixed case following CUTLASS/CUTE conventions -- Format: lowercase prefix + PascalCase suffix -- Prefix indicates memory space: `g` for global, `s` for shared, `r` for register -- Suffix indicates data type or purpose (single letter or abbreviation) -- Examples: - ```cpp - Tensor gA = make_tensor(...); // Global tensor A - Tensor gLse = make_tensor(...); // Global LSE tensor - Tensor sB = make_tensor(...); // Shared tensor B - Tensor rC = make_tensor(...); // Register tensor C - ``` -- This style prioritizes brevity in mathematical code while maintaining readability - -### Functions -- Use `snake_case` for function names -- Kernel functions: `fwd_combine_kernel` -- Host functions: `run_fwd_combine_kernel` -- Template instantiations: `run_fwd_combine_kernel` - -### Types and Classes -- Use `PascalCase` for type names -- Template parameters: `ElementT`, `HEAD_DIM_V` -- Structs: `CombineParams`, `DenseAttnDecodeParams` - -### Namespaces -- Use nested namespaces: `namespace smxx::decode { ... }` -- Avoid `using namespace` in header files -- Use `using namespace cute;` in implementation files when appropriate - -### Constants and Macros -- Use `SCREAMING_SNAKE_CASE` for constants -- CUDA-specific: `FLASH_DEVICE_ASSERT`, `CUTLASS_PRAGMA_UNROLL` -- Mathematical constants: `M_LOG2E`, `CUDART_L2E_F` - -## Formatting - -### Indentation -- Use 4 spaces for indentation (no tabs) -- Align continuation lines with the opening parenthesis or bracket - -### Braces -- Use K&R style (opening brace on same line) -- Always use braces for control structures, even for single statements -- Examples: - ```cpp - if (condition) { - return; - } - - for (int i = 0; i < n; ++i) { - // statements - } - ``` - -### Line Length -- Limit lines to 120 characters -- Break long lines at logical points (operators, commas) - -### Spacing -- One space around binary operators: `a + b` -- No space after unary operators: `-x` -- Space after keywords: `if (`, `for (` -- No space before semicolons -- No space after C-style casts: `(int)value` not `(int) value` - -### Template Declarations -- Always break template declarations onto new lines -- Example: - ```cpp - template - __global__ void kernel_function(...) - ``` - -## CUDA-Specific Conventions - -### Kernel Launch -- Use `__launch_bounds__(NUM_THREADS)` for kernel declarations -- Use `__grid_constant__` for constant parameters -- Example: - ```cpp - template - __global__ void __launch_bounds__(NUM_THREADS) - fwd_combine_kernel(__grid_constant__ const CombineParams params) - ``` - -### Thread Management -- Use `threadIdx`, `blockIdx`, `blockDim` for indexing -- Calculate `warp_idx = threadIdx.x / 32`, `lane_idx = threadIdx.x % 32` -- Use `__syncwarp()` for warp synchronization - -### Memory Access -- Use `__ldg()` for read-only global memory access (when compatible) -- Use `float4` for vectorized loads/stores -- Prefer shared memory for intra-block communication - -### Pragmas -- Use `CUTLASS_PRAGMA_UNROLL` for loop unrolling hints -- Use `#pragma unroll` sparingly, prefer compiler hints - -## Comments - -### Style -- Use `//` for single-line comments -- Use `/* */` for multi-line comments only when necessary -- Place comments above the code they describe - -### Content -- Explain complex algorithms and optimizations -- Document kernel grid/block dimensions: `// grid_shape: [batch_size, s_q, h_q/BLOCK_SIZE_M]` -- Note performance-critical sections -- Warn about limitations: `// NOTE: We don't use __ldg here since it is incompatible with PDL` - -## Templates and Generics - -### Template Parameters -- Use descriptive names: `typename ElementT`, `int HEAD_DIM_V` -- Order: types first, then integers, then booleans -- Example: - ```cpp - template - ``` - -### Template Instantiations -- Explicitly instantiate in `.cu` files -- Group instantiations at end of file -- Example: - ```cpp - template void run_fwd_combine_kernel(CombineParams ¶ms); - template void run_fwd_combine_kernel(CombineParams ¶ms); - ``` - -## Error Handling - -- Use `FLASH_DEVICE_ASSERT` for device-side assertions -- Use `CUTLASS_CHECK` for CUTLASS operations -- Handle edge cases gracefully (e.g., `my_num_splits == 1` early return) -- Validate inputs at kernel launch time - -## Tools - -### Code Formatting -- Use `clang-format` with the provided `.clang-format` configuration -- Run before committing: `clang-format -i ` - -### Static Analysis -- Use `clang-tidy` with the provided `.clang-tidy` configuration -- Run on modified files: `clang-tidy -- -I` -- Configure IDE integration for real-time checking -- Checks enforce naming conventions, code quality, and modernization rules -- Some CUDA-specific patterns are intentionally excluded to avoid false positives - -### Recommended Workflow -1. Write or modify code -2. Run `clang-format -i` to fix formatting -3. Run `clang-tidy` to check for issues -4. Fix warnings before committing - -This style guide should be followed for all new code contributions. For existing code that doesn't conform, consider refactoring during related changes. \ No newline at end of file +cuLA is a Python package whose GPU kernels are written in CuTeDSL, with a +small number of Triton support kernels. The package does not build a native +C++/CUDA extension. + +## General principles + +- Prefer clear Python host code and explicit CuTeDSL data movement. +- Keep architecture-specific code under `cula/ops//sm90/` or + `cula/ops//sm100/`. +- Preserve the public FLA-compatible tensor layout, dtype, and state-layout + contracts at wrapper boundaries. +- Validate unsupported architectures, shapes, dtypes, layouts, and modes + before launching a kernel. +- Document intentional precision changes and reduction-order differences. + +## Python + +- Use four spaces, type annotations for public APIs, and `snake_case` names. +- Keep imports sorted with Ruff and avoid eager imports of large kernel modules + from package `__init__.py` files. +- Use `PascalCase` for classes and `SCREAMING_SNAKE_CASE` for module constants. +- Raise `TypeError` for dtype/type mismatches, `ValueError` for invalid values + or shapes, and `NotImplementedError` for valid but unsupported modes. +- Keep compile caches bounded and key them by every property that affects the + generated kernel. + +## CuTeDSL kernels + +- Make grid shape, warp roles, pipeline stages, and shared-memory ownership + explicit near the kernel entry point. +- Keep barrier arrival counts and phase changes next to the producer/consumer + protocol they protect. +- Do not reuse shared memory until every prior consumer has completed and the + required async-proxy fence has been issued. +- Prefer vectorized, aligned global-memory accesses and state the required + alignment in host validation. +- Use descriptive memory-space prefixes for tensors (`g`, `s`, `r`) when that + matches the surrounding CuTe/CUTLASS conventions. +- Keep numerical conversion points visible. Avoid implicit precision loss in + accumulators, recurrence state, or transition products. + +## Tests and benchmarks + +- Add a focused test for every public behavior change and each corrected race + or boundary condition. +- Compare against an independent reference where possible, and distinguish + bitwise requirements from tolerance-based numerical contracts. +- Include tail, varlen, initial-state, output-state, and non-default-stream + coverage when the changed path supports them. +- Mark architecture-specific and long-running cases with the repository's + pytest markers so the default suite stays bounded. +- Benchmark the complete public boundary, warm up JIT compilation first, and + report the device, shapes, dtype, and aggregation method. + +## Tooling + +Before committing, run: + +```bash +pre-commit run --all-files +``` + +Run CPU-only checks locally and CUDA tests in the designated remote container. diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index 1d0c9e54..ad48d1f7 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -6,26 +6,17 @@ Legend: `[exp]` experimental / unwired · `[non-KDA]` other operator. cuLA/ ├── cula/ # Python package (pip install -e .) │ ├── __init__.py -│ ├── backends.py # Generic backend registry, probing, verification, dispatch -│ ├── cudac.py # Lazy proxy for the remaining SM90 CUDA extension │ ├── utils.py # Architecture, stream-buffer, and cu_seqlens helpers │ │ -│ ├── kda/ # KDA public API, wrappers, autograd, routing, and Triton support kernels +│ ├── kda/ # KDA public API, wrappers, autograd, and support kernels │ │ ├── __init__.py # Lazy exports for chunk, prefill, and decode APIs -│ │ ├── backends/ # kda_prefill runtime dispatch -│ │ │ ├── flashkda.py # Preferred SM90 CuTeDSL backend + verifier -│ │ │ ├── fully_fused.py # SM90 CUDA C++ backend + verifier -│ │ │ └── _common.py # Shared dispatch input checks -│ │ ├── auto_route.py # Fully-fused SM90 basic/optimized router +│ │ ├── prefill.py # Public FlashKDA prefill validation and routing │ │ ├── flashkda.py # SM90 CuTeDSL K1+K2 wrapper, autograd shell, and CP selection -│ │ ├── hopper_fused_fwd.py # SM90 CUDA C++ basic wrapper -│ │ ├── hopper_fused_fwd_opt.py # SM90 CUDA C++ optimized/CP wrapper │ │ ├── chunk.py # SM100 modular chunk API + autograd │ │ ├── chunk_fwd.py # SM100 forward orchestration │ │ ├── chunk_intra.py # CuTeDSL forward intra + Triton backward intra │ │ ├── chunk_bwd.py # Triton + FLA + CuTeDSL backward orchestration -│ │ ├── wy_intra.py / wy_recompute.py # Triton WY preparation kernels -│ │ └── gate_l2norm_fused.py / l2norm_qk_fused.py # SM90 Triton preprocessing kernels +│ │ └── wy_intra.py / wy_recompute.py # Triton WY preparation kernels │ │ │ ├── lightning/ # [non-KDA] Lightning Attention public wrappers │ │ @@ -52,15 +43,6 @@ cuLA/ │ │ └── prefill_kernel.py # Fixed and packed recurrent prefill kernel │ └── experimental/ # [non-KDA] unwired prototypes │ -├── csrc/ # Remaining CUDA C++ / CUTLASS kernels for SM90 -│ ├── api/ -│ │ ├── pybind.cu # Shared named binding definitions -│ │ ├── pybind_sm90.cu # SM90-specific binding translation unit -│ │ └── kda_sm90.cu # SM90 host API -│ ├── kda/sm90/ # Hopper fully-fused prefill kernels -│ ├── kerutils/include/ # Shared device and host helpers -│ └── cutlass/ # CUTLASS submodule -│ ├── benchmarks/ # Accuracy and performance runners ├── tests/ # CPU registry tests and architecture-marked GPU tests ├── docs/ # Pipeline and backend-dispatch design notes @@ -74,10 +56,7 @@ cuLA/ | Directory | Language | Description | |-----------|----------|-------------| -| `cula/kda/` | Python / Triton | Public KDA APIs, wrappers, autograd, routing, and supporting Triton kernels. `kda_prefill` dispatches SM90 calls to FlashKDA first and then the fully-fused CUDA backend; explicit backend exports remain available. See [`cula/kda/README.md`](cula/kda/README.md). | -| `cula/backends.py` · `cula/kda/backends/` | Python | Generic backend registry plus KDA-specific availability probes, capability verifiers, priorities, and adapters. | -| `cula/ops/kda/` | Python (CuTeDSL) | CuTeDSL KDA kernels organized into SM100 modular kernels, SM90 FlashKDA K1+K2 and intracard CP, decode, and experimental code. The fully-fused SM90 implementation lives under `csrc/`, not here. | -| `csrc/kda/sm90/` | CUDA C++ | Hopper fully-fused prefill kernels. SM100/SM103 KDA forward is implemented entirely in CuTeDSL under `cula/ops/kda/sm100/`. | -| `csrc/api/` · `cula/cudac.py` | CUDA C++ / Python | The remaining `_cudac_sm90` extension and its lazy compatibility proxy. Blackwell public APIs do not load a native extension. | +| `cula/kda/` | Python / Triton | Public KDA APIs, wrappers, autograd, routing, and supporting Triton kernels. `kda_prefill` calls FlashKDA directly without a native fallback. See [`cula/kda/README.md`](cula/kda/README.md). | +| `cula/ops/kda/` | Python (CuTeDSL) | CuTeDSL KDA kernels organized into SM100 modular kernels, SM90 FlashKDA K1+K2 and intracard CP, decode, and experimental code. | | `cula/ops/lightning/` · `cula/ops/experimental/` | Python (CuTeDSL) | `[non-KDA]` Lightning/linear-attention kernels and prototypes. Lightning prefill dispatches to the SM90 or SM100 backend from Q's device capability. | | `cula/ops/{inv,ptx}.py` · `cula/ops/sm100/ptx.py` | Python | Shared low-level helpers used across operators. | diff --git a/USAGE.md b/USAGE.md index aa3b04c2..ff6e0a1a 100644 --- a/USAGE.md +++ b/USAGE.md @@ -11,12 +11,14 @@ cuLA exposes several KDA entry points targeting different GPU architectures: | Kernel | GPU | Import | |---|---|---| | Modular Forward | Blackwell (SM100) | `from cula.kda import chunk_kda` | -| Fused Forward | Hopper (SM90) | `from cula.kda import kda_prefill_hopper` | | Two-Kernel Prefill | Hopper (SM90) | `from cula.kda import flashkda_prefill` | -| Auto-dispatch Prefill | Hopper (SM90) | `from cula.kda import kda_prefill` | +| Public Prefill | Hopper (SM90) / Blackwell (SM100) | `from cula.kda import kda_prefill` | -`chunk_kda`, `kda_prefill_hopper`, and `flashkda_prefill` each call a fixed implementation. -`kda_prefill` is the only one that picks an SM90 backend automatically. +`chunk_kda` and `flashkda_prefill` each call a fixed implementation. +`kda_prefill` is the public FlashKDA wrapper and has no native-extension fallback. +The removed `kda_prefill_hopper`, `kda_prefill_hopper_opt`, and +`kda_prefill_hopper_auto` names should be migrated to one of these FlashKDA +entry points. **General Notes** @@ -26,11 +28,13 @@ cuLA exposes several KDA entry points targeting different GPU architectures: --- -### Auto-dispatch Prefill (SM90 — Hopper) +### FlashKDA Prefill (SM90 / SM100) `kda_prefill` only runs forward (use it under `torch.inference_mode()` / `no_grad`). -It tries FlashKDA first; if that path cannot handle the call (for example GVA, where `V` heads ≠ `Q/K` heads), it falls back to the fully-fused CUDA C++ kernel. -Pass both `A_log` and `dt_bias` if you want FlashKDA to be eligible — without them, FlashKDA rejects the call. +It calls the CuTeDSL FlashKDA implementation directly. Pass both `A_log` and +`dt_bias`; unsupported calls fail explicitly rather than falling back to a +native extension. Grouped-value attention (`V` heads != `Q/K` heads) is not +yet supported by FlashKDA. ```python import torch @@ -57,10 +61,6 @@ with torch.inference_mode(): ) ``` -Set `CULA_BACKEND_FLASHKDA=0` to skip FlashKDA or `CULA_BACKEND_FULLY_FUSED=0` to skip the CUDA C++ backend. If every enabled backend rejects the call, the error lists each rejection reason. - ---- - ### Modular Forward (SM100 — Blackwell) The modular forward kernel replaces sub-kernels of KDA in FLA (chunk_intra, chunk_delta_h, fwd_o, etc.) for easy integration with [Kimi CP](https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/cp/README.md). @@ -108,51 +108,6 @@ print(f'Final state shape: {final_state.shape}') # [2, 32, 128, 128] --- -### Fused Forward (SM90 — Hopper) - -The fused forward kernel fuses intra-chunk attention, inter-chunk state propagation, and output computation into a single kernel for maximum throughput. **Forward-only; backward is not yet implemented.** - -#### Example - -```python -import torch -from cula.kda import kda_prefill_hopper - -B, T, H, K, V = 2, 2048, 32, 128, 128 -device = 'cuda' - -q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) -k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) -v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) -g = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) * 0.1 -beta = torch.randn(B, T, H, device=device, dtype=torch.bfloat16).sigmoid() -A_log = torch.randn(H, device=device, dtype=torch.float32) * 0.01 -dt_bias = torch.zeros(H * K, device=device, dtype=torch.float32) -init_state = torch.zeros(B, H, K, V, device=device, dtype=torch.float32) - -o, final_state = kda_prefill_hopper( - q=q, k=k, v=v, g=g, beta=beta, - A_log=A_log, dt_bias=dt_bias, - initial_state=init_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - safe_gate=True, - lower_bound=-5.0, -) - -print(f'Output shape: {o.shape}') # [2, 2048, 32, 128] -print(f'Final state shape: {final_state.shape}') # [2, 32, 128, 128] -``` - -**Notes** - -- Mainly **suitable for large-batch inference**; performance is limited when both batch size and head count are small, because we do not parallelize over the sequence-length dimension. -- **Matrix inversion uses fp16 precision**, which is faster and occupies less shared memory but introduces minor numerical differences compared to tf32 inversion. -- **Intra-subchunk attention uses g-first as anchor**, which causes some numerical differences compared with the FLA Triton implementation (FLA uses g-half as anchor in the diagonal). - ---- - ### Two-Kernel Prefill (SM90 — Hopper) The SM90 prefill is a **two-kernel pipeline** — K1 (Prepare) → K2 (Recurrence) — *not* a single fused kernel. Intra-chunk attention, inter-chunk state propagation, and output are computed across the two kernels. **Forward-only; backward is not yet implemented.** diff --git a/benchmarks/bench_intracard_cp_sm90.py b/benchmarks/bench_intracard_cp_sm90.py deleted file mode 100644 index c8f85773..00000000 --- a/benchmarks/bench_intracard_cp_sm90.py +++ /dev/null @@ -1,448 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# 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. - -"""bench_intracard_cp_sm90.py — CP-on vs CP-off (vs FLA baseline) for SM90 KDA prefill. - -Mirrors benchmarks/bench_intracard_cp.py (SM100 version) but for the Hopper -(SM90) path: - - CP_on : cula.kda.kda_prefill_hopper_auto - CP_off : cula.kda.kda_prefill_hopper - FLA : fla.ops.kda.chunk_kda (Triton baseline) - -Reports per-config `pred` (would CP fire?) and `n_sub` (CP-chunk count). When -`pred=N` we still measure CP_on to confirm the bypass adds no regression. The -`CP_on/FLA` column shows the speedup of cuLA's optimized (CP-on) kernel over -the FLA Triton baseline. - -Usage: - python benchmarks/bench_intracard_cp_sm90.py [--ncu] [--sanitizer] -""" - -import argparse -import os -import pathlib -import sys - -import torch - -sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) -os.environ.setdefault("FLA_INTRACARD_CP", "1") - -from fla.ops.common.intracard_cp import compute_subseq_len, prepare_subseq_cu_seqlens -from fla.ops.kda import chunk_kda as fla_chunk_kda - -from benchmarks.utils import ( - SEED, - exclusive_cumsum, - prepare_safe_gate_inputs, - set_seed, - time_cuda_fn, -) -from cula.kda import kda_prefill_hopper, kda_prefill_hopper_auto -from cula.kda.auto_route import _should_use_opt -from cula.kda.cp_context import _calc_cp_seqs, is_dominant_long_seq -from cula.kda.hopper_fused_fwd_opt import FUSED_GATE_L2NORM_VARLEN_AVG_SEQ, _fused_gate_l2norm_threshold -from cula.utils import get_device_sm_count - -# ============================================================ -# Constants -# ============================================================ -BT, D = 64, 128 -H_VALUES = [4, 8] -WARMUP = 10 -N_ITERS = 10 -NCU_MODE = False -SANITIZER_MODE = False - -# (tag, seq_lens) — varlen configs, run with cu_seqlens=cumsum(seq_lens) -CONFIGS = [ - # small varlen — exercises fused gate+l2norm path (packed_T*H <= 65536) - ("4x256", [256] * 4), - ("8x256", [256] * 8), - ("16x256", [256] * 16), - ("4x1K", [1024] * 4), - ("8x1K", [1024] * 8), - ("4x2K", [2048] * 4), - ("1K+512+256+128", [1024, 512, 256, 128]), - ("2K+1K+512+256", [2048, 1024, 512, 256]), - ("1K+1+63+65+129", [1024, 1, 63, 65, 129]), - # single seq - ("T=4K", [4096]), - ("T=8K", [8192]), - ("T=32K", [32768]), - ("T=64K", [65536]), - ("T=128K", [131072]), - # equal-length batches (~32K total) - ("8x4K", [4096] * 8), - ("4x8K", [8192] * 4), - ("2x16K", [16384] * 2), - # asymmetric multi-seq - ("16K+16K", [16384, 16384]), - ("24K+8K", [24576, 8192]), - ("28K+4K", [28672, 4096]), - ("32K+256+256", [32768, 256, 256]), - ("40K+1K+8K", [40960, 1024, 8192]), - ("64K+512+256+128", [65536, 512, 256, 128]), - ("128K+1K", [131072, 1024]), - ("128K+2x1K", [131072, 1024, 1024]), - ("128K+5x1K", [131072] + [1024] * 5), - ("128K+10x1K", [131072] + [1024] * 10), -] - - -# ============================================================ -# Helpers -# ============================================================ -def _bench_warmup_iters(): - warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP - n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS - return warmup, n_iters - - -def run_call(q, k, v, g, beta, scale, A_log, dt_bias, cu_seqlens, lower_bound, *, enable_cp, return_state=False): - fn = kda_prefill_hopper_auto if enable_cp else kda_prefill_hopper - out = fn( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - A_log=A_log, - dt_bias=dt_bias, - initial_state=None, - output_final_state=return_state, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - safe_gate=True, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - ) - return out - - -def run_fla_call(q, k, v, g, beta, scale, A_log, dt_bias, cu_seqlens, lower_bound, *, return_state=False): - # FLA's chunk_kda fuses A_log + dt_bias internally when use_gate_in_kernel=True; pass them via kwargs. - # Wrap in inference_mode so FLA's IntraCardCPBackend gate (intracard.py) activates. - with torch.inference_mode(): - return fla_chunk_kda( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state=None, - output_final_state=return_state, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - safe_gate=False, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - A_log=A_log, - dt_bias=dt_bias, - ) - - -def accuracy(ref, got): - if ref is None or got is None: - return float("nan"), float("nan") - diff = (ref.float() - got.float()).abs() - return diff.max().item(), diff.mean().item() - - -def predict_cp(seq_lens, H, num_sms, device): - cu = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - raw_batch = len(seq_lens) - packed_seq = sum(seq_lens) - - if raw_batch > 1: - cp_wf = (raw_batch * H <= 16 and packed_seq >= 8192) or ( - packed_seq >= 8192 and H <= 16 and is_dominant_long_seq(seq_lens, H) - ) - else: - cp_wf = (H <= 8 and packed_seq >= 4096) or (H <= 16 and packed_seq >= 4096) or (H <= 32 and packed_seq >= 16384) - if not cp_wf: - return False, 0 - - use_cp, cp_cu, *_ = _calc_cp_seqs(cu, BT, H, num_sms, raw_cu_seqlens_cpu=cu.cpu()) - if not use_cp: - return False, 0 - n_sub = int(cp_cu.numel() - 1) - if n_sub == raw_batch: # no-op split - return False, 0 - return True, n_sub - - -def predict_fla_cp(seq_lens, H, num_sms): - """Mirror fla.ops.common.intracard_cp.intracard_fwd_h gating to predict - whether FLA's intracard CP fires and how many sub-sequences result. - HV (num_v_heads) maps to H here.""" - cu = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int64) - seq_lens_t = torch.diff(cu) - max_seq_len = int(seq_lens_t.max().item()) - subseq_len = compute_subseq_len(max_seq_len, num_sms, H, BT) - if (seq_lens_t < 2 * subseq_len).all(): - return False, 0 - _, split_info, total_subseqs = prepare_subseq_cu_seqlens(cu, subseq_len, BT) - if not split_info: - return False, 0 - return True, total_subseqs - - -def predict_fused_all_pre(q, v, cu_seqlens_for_opt, *, cu_seqlens_is_none, use_gate_in_kernel, use_qk_l2norm_in_kernel): - if not _should_use_opt(q, cu_seqlens_for_opt): - return False - num_qk_heads = q.shape[-2] - num_v_heads = v.shape[-2] - if cu_seqlens_is_none: - avg_seq_ok = True - else: - N = cu_seqlens_for_opt.numel() - 1 - packed_T = q.shape[1] - avg_seq_ok = N <= 1 or packed_T <= N * FUSED_GATE_L2NORM_VARLEN_AVG_SEQ - return ( - use_gate_in_kernel - and use_qk_l2norm_in_kernel - and (q.numel() // q.shape[-1]) <= _fused_gate_l2norm_threshold(cu_seqlens_is_none) - and num_qk_heads == num_v_heads - and avg_seq_ok - ) - - -# ============================================================ -# Benchmark -# ============================================================ -SEP = " " + "─" * 180 -ROW_HEADER = ( - f" {'config':<24s} {'T':>7s} {'pred':>4s} {'sub':>4s} {'fla_cp':>6s} {'fla_sub':>7s} {'fused_pre':>5s}" - f" │ {'o max/mean':>17s} {'ht max/mean':>17s}" - f" │ {'FLA(ms)':>9s} {'CP_off(ms)':>10s} {'CP_on(ms)':>10s} {'CP_on/off':>8s} {'CP_on/FLA':>9s}" -) - - -def _format_row(r): - pred_s = "Y" if r["pred"] else "N" - fla_pred_s = "Y" if r["fla_pred"] else "N" - fused_s = "Y" if r["fused_all_pre"] else "N" - return ( - f" {r['tag']:<24s} {r['total_T']:>7d} {pred_s} {r['n_sub']:>4d} {fla_pred_s} {r['fla_n_sub']:>4d} {fused_s}" - f" │ {r['o_max']:>7.1e}/{r['o_mean']:>7.1e} {r['ht_max']:>7.1e}/{r['ht_mean']:>7.1e}" - f" │ {r['ms_fla']:>9.4f} {r['ms_off']:>10.4f} {r['ms_on']:>10.4f}" - f" {r['speedup']:>7.2f}x {r['speedup_vs_fla']:>8.2f}x" - ) - - -def bench_cp(h_values, configs): - print("\n" + "=" * 110) - print(" BENCHMARK REPORT: Intracard CP (SM90)") - print(" CP-on (kda_prefill_hopper_auto) vs CP-off (kda_prefill_hopper) vs FLA (chunk_kda)") - print(f" D={D} dtype=bf16 safe_gate=True") - wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP - ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS - mode_tag = " [NCU mode]" if NCU_MODE else (" [Sanitizer mode]" if SANITIZER_MODE else "") - print(f" Warmup={wu} Iters={ni}{mode_tag}") - print("=" * 110) - - device = torch.device("cuda") - num_sms = get_device_sm_count(device) - results = [] - - for H in h_values: - print(f"\n [H={H}]", flush=True) - print(SEP, flush=True) - print(ROW_HEADER, flush=True) - print(SEP, flush=True) - - for tag, seq_lens in configs: - set_seed(SEED) - torch.cuda.empty_cache() - - total_T = sum(seq_lens) - cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) - inputs = prepare_safe_gate_inputs(1, total_T, H, D, device, cu_seqlens=cu_seqlens, seed=SEED) - q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] - A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] - scale, lower_bound = inputs["scale"], inputs["lower_bound"] - - pred, n_sub = predict_cp(seq_lens, H, num_sms, device) - fla_pred, fla_n_sub = predict_fla_cp(seq_lens, H, num_sms) - fused_all_pre = predict_fused_all_pre( - q, - v, - cu_seqlens, - cu_seqlens_is_none=False, - use_gate_in_kernel=True, - use_qk_l2norm_in_kernel=True, - ) - - common = dict( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - A_log=A_log, - dt_bias=dt_bias, - cu_seqlens=cu_seqlens, - lower_bound=lower_bound, - ) - - try: - o_off, ht_off = run_call(**common, enable_cp=False, return_state=True) - o_on, ht_on = run_call(**common, enable_cp=True, return_state=True) - o_max, o_mean = accuracy(o_off, o_on) - ht_max, ht_mean = accuracy(ht_off, ht_on) - del o_off, ht_off, o_on, ht_on - - ms_off = time_cuda_fn(lambda: run_call(**common, enable_cp=False), *_bench_warmup_iters()) - ms_on = time_cuda_fn(lambda: run_call(**common, enable_cp=True), *_bench_warmup_iters()) - speedup = ms_off / ms_on if ms_on > 0 else float("inf") - try: - ms_fla = time_cuda_fn(lambda: run_fla_call(**common), *_bench_warmup_iters()) - speedup_vs_fla = ms_fla / ms_on if ms_on > 0 else float("inf") - except Exception: - ms_fla = float("nan") - speedup_vs_fla = float("nan") - except torch.cuda.OutOfMemoryError: - ms_off = ms_on = speedup = float("nan") - ms_fla = speedup_vs_fla = float("nan") - o_max = o_mean = ht_max = ht_mean = float("nan") - - row = { - "tag": tag, - "H": H, - "total_T": total_T, - "pred": pred, - "n_sub": n_sub, - "fla_pred": fla_pred, - "fla_n_sub": fla_n_sub, - "fused_all_pre": fused_all_pre, - "ms_off": ms_off, - "ms_on": ms_on, - "ms_fla": ms_fla, - "speedup": speedup, - "speedup_vs_fla": speedup_vs_fla, - "o_max": o_max, - "o_mean": o_mean, - "ht_max": ht_max, - "ht_mean": ht_mean, - } - results.append(row) - print(_format_row(row), flush=True) - - del q, k, v, g, beta, A_log, dt_bias, inputs - torch.cuda.empty_cache() - - print(SEP, flush=True) - - return results - - -# ============================================================ -# Report (summary only — per-row output is streamed inside bench_cp) -# ============================================================ -def print_report(results, h_values): - sep = "=" * 110 - triggered = [r for r in results if r["pred"]] - bypassed = [r for r in results if not r["pred"]] - - print() - print(sep) - print(" Summary") - print(sep) - - if triggered: - speedups = [r["speedup"] for r in triggered if r["speedup"] == r["speedup"]] # NaN filter - if speedups: - geo = 1.0 - for s in speedups: - geo *= s - geo = geo ** (1 / len(speedups)) - print( - f" CP triggered ({len(triggered)} configs): " - f"geo-mean={geo:.2f}x best={max(speedups):.2f}x worst={min(speedups):.2f}x" - ) - - if bypassed: - ratios = [r["ms_on"] / r["ms_off"] for r in bypassed if r["ms_off"] == r["ms_off"] and r["ms_off"] > 0] - if ratios: - print( - f" CP bypassed ({len(bypassed)} configs): " - f"mean overhead={sum(ratios) / len(ratios):.3f}x max={max(ratios):.3f}x " - f"(1.00 = no regression)" - ) - - # cuLA (CP-on) vs FLA speedups - fla_speedups = [r["speedup_vs_fla"] for r in results if r["speedup_vs_fla"] == r["speedup_vs_fla"]] - if fla_speedups: - geo = 1.0 - for s in fla_speedups: - geo *= s - geo = geo ** (1 / len(fla_speedups)) - print( - f" cuLA (CP-on) vs FLA ({len(fla_speedups)} configs): " - f"geo-mean={geo:.2f}x best={max(fla_speedups):.2f}x worst={min(fla_speedups):.2f}x" - ) - tri_fla = [r["speedup_vs_fla"] for r in triggered if r["speedup_vs_fla"] == r["speedup_vs_fla"]] - if tri_fla: - geo_t = 1.0 - for s in tri_fla: - geo_t *= s - geo_t = geo_t ** (1 / len(tri_fla)) - print( - f" └─ CP-triggered subset ({len(tri_fla)} configs): " - f"geo-mean={geo_t:.2f}x best={max(tri_fla):.2f}x worst={min(tri_fla):.2f}x" - ) - - o_maxes = [r["o_max"] for r in results if r["o_max"] == r["o_max"]] - ht_maxes = [r["ht_max"] for r in results if r["ht_max"] == r["ht_max"]] - if o_maxes: - print( - f" Accuracy (CP-on vs CP-off): " - f"o max={max(o_maxes):.2e} avg={sum(o_maxes) / len(o_maxes):.2e} " - f"ht max={max(ht_maxes):.2e} avg={sum(ht_maxes) / len(ht_maxes):.2e}" - ) - - print(sep) - - -# ============================================================ -# Main -# ============================================================ -def main(): - parser = argparse.ArgumentParser(description="bench_intracard_cp_sm90: CP-on vs CP-off") - parser.add_argument("--ncu", action="store_true", help="NCU profiling mode: warmup=1, iters=1") - parser.add_argument("--sanitizer", action="store_true", help="Sanitizer mode: warmup=1, iters=1") - args = parser.parse_args() - - global NCU_MODE, SANITIZER_MODE - if args.ncu: - NCU_MODE = True - print("[NCU mode] warmup=1, iters=1") - if args.sanitizer: - SANITIZER_MODE = True - print("[Sanitizer mode] warmup=1, iters=1") - - results = bench_cp(H_VALUES, CONFIGS) - print_report(results, H_VALUES) - return results - - -if __name__ == "__main__": - main() diff --git a/csrc/api/kda_sm90.cu b/csrc/api/kda_sm90.cu deleted file mode 100644 index bd7ba5ee..00000000 --- a/csrc/api/kda_sm90.cu +++ /dev/null @@ -1,226 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#include -#include -#include -#include - -#include "kda/sm90/prefill_kernel.hpp" - -using OptionalTensor = std::optional; - -std::tuple -kda_fwd_prefill( - OptionalTensor output_, - OptionalTensor output_state_, - torch::Tensor const& q, - torch::Tensor const& k, - torch::Tensor const& v, - OptionalTensor input_state_, - OptionalTensor alpha_, - OptionalTensor beta_, - torch::Tensor const& cu_seqlens, - torch::Tensor workspace_buffer, - float scale, - bool output_final_state, - bool safe_gate, - OptionalTensor cp_seq_map_, - OptionalTensor raw_cu_seqlens_) { - // Q, K: [packed_seq, num_qk_heads, D] - // V/O/g: [packed_seq, num_v_heads, D] (GVA: num_v_heads is a positive integer multiple of num_qk_heads) - auto packed_seq = q.size(0); - auto num_qk_heads = q.size(1); - auto num_v_heads = v.size(1); - auto head_size = q.size(2); - auto num_seqs = cu_seqlens.size(0) - 1; - - // Intra-card CP plumbing. - int32_t const* cp_seq_map_ptr = nullptr; - int32_t const* raw_cu_seqlens_ptr = nullptr; - int32_t raw_num_seqs = static_cast(num_seqs); - if (cp_seq_map_.has_value()) { - TORCH_CHECK(raw_cu_seqlens_.has_value(), "raw_cu_seqlens must be provided alongside cp_seq_map"); - auto const& cp_seq_map = cp_seq_map_.value(); - auto const& raw_cu_seqlens = raw_cu_seqlens_.value(); - TORCH_CHECK(cp_seq_map.device() == q.device(), "cp_seq_map must be on the same device as q"); - TORCH_CHECK(raw_cu_seqlens.device() == q.device(), "raw_cu_seqlens must be on the same device as q"); - TORCH_CHECK(cp_seq_map.dtype() == torch::kInt32, "cp_seq_map must be int32"); - TORCH_CHECK(raw_cu_seqlens.dtype() == torch::kInt32, "raw_cu_seqlens must be int32"); - TORCH_CHECK(cp_seq_map.is_contiguous(), "cp_seq_map must be contiguous"); - TORCH_CHECK(raw_cu_seqlens.is_contiguous(), "raw_cu_seqlens must be contiguous"); - TORCH_CHECK( - cp_seq_map.size(0) == num_seqs, - "cp_seq_map.size(0) must equal cu_seqlens.size(0)-1, got ", - cp_seq_map.size(0), - " vs ", - num_seqs); - cp_seq_map_ptr = cp_seq_map.data_ptr(); - raw_cu_seqlens_ptr = raw_cu_seqlens.data_ptr(); - raw_num_seqs = static_cast(raw_cu_seqlens.size(0) - 1); - } - - // GVA contract on the C++ side. Order matters: check positivity *before* the modulo to - // avoid % 0 / division-by-zero UB in case the Python layer passed a degenerate shape. - TORCH_CHECK(num_qk_heads > 0, "KDA requires num_qk_heads > 0, got ", num_qk_heads); - TORCH_CHECK(num_v_heads > 0, "KDA requires num_v_heads > 0, got ", num_v_heads); - TORCH_CHECK( - num_qk_heads == k.size(1), "KDA requires num_q_heads == num_k_heads, got ", num_qk_heads, " vs ", k.size(1)); - TORCH_CHECK( - num_v_heads % num_qk_heads == 0, - "KDA GVA requires num_v_heads to be a positive multiple of num_qk_heads, got num_v_heads=", - num_v_heads, - ", num_qk_heads=", - num_qk_heads); - TORCH_CHECK(head_size == v.size(2), "KDA requires Q and V head dim to match, got ", head_size, " vs ", v.size(2)); - - // Allocate output if not provided - torch::Tensor output = output_.has_value() ? output_.value() - : torch::empty( - {packed_seq, num_v_heads, head_size}, - torch::TensorOptions().dtype(q.dtype()).device(q.device())); - - // Allocate output state if not provided. In CP mode the state is keyed - // by raw_num_seqs (one slot per original sequence), not the inflated - // CP-chunk count. - OptionalTensor output_state = std::nullopt; - if (output_final_state) { - output_state = output_state_.has_value() - ? output_state_.value() - : torch::zeros( - {raw_num_seqs, num_v_heads, head_size, head_size}, - torch::TensorOptions().dtype(torch::kFloat32).device(q.device())); - } - - // Validate dtypes - TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bfloat16"); - TORCH_CHECK(k.dtype() == torch::kBFloat16, "k must be bfloat16"); - TORCH_CHECK(v.dtype() == torch::kBFloat16, "v must be bfloat16"); - TORCH_CHECK(cu_seqlens.dtype() == torch::kInt32, "cu_seqlens must be int32"); - - // Validate contiguity - TORCH_CHECK(q.is_contiguous(), "q must be contiguous"); - TORCH_CHECK(k.is_contiguous(), "k must be contiguous"); - TORCH_CHECK(v.is_contiguous(), "v must be contiguous"); - TORCH_CHECK(output.is_contiguous(), "output must be contiguous"); - if (output_state.has_value()) { - TORCH_CHECK(output_state->dtype() == torch::kFloat32, "output_state must be float32"); - TORCH_CHECK(output_state->is_contiguous(), "output_state must be contiguous"); - } - TORCH_CHECK(cu_seqlens.is_contiguous(), "cu_seqlens must be contiguous"); - TORCH_CHECK(workspace_buffer.is_contiguous(), "workspace_buffer must be contiguous"); - - // Extract optional pointers - float const* alpha_ptr = nullptr; - float const* input_state_ptr = nullptr; - - if (alpha_.has_value()) { - auto& alpha = alpha_.value(); - TORCH_CHECK(alpha.dtype() == torch::kFloat32, "alpha must be float32"); - TORCH_CHECK(alpha.is_contiguous(), "alpha must be contiguous"); - TORCH_CHECK( - alpha.size(0) == packed_seq && alpha.size(1) == num_v_heads && alpha.size(2) == head_size, - "alpha shape must be [packed_seq, num_v_heads, head_size]"); - alpha_ptr = alpha.data_ptr(); - } - - float* output_state_ptr = output_state.has_value() ? output_state->data_ptr() : nullptr; - if (beta_.has_value()) { - auto& beta = beta_.value(); - TORCH_CHECK( - beta.dtype() == torch::kFloat32 || beta.dtype() == torch::kBFloat16, - "beta must be float32 or bfloat16, got ", - beta.dtype()); - TORCH_CHECK(beta.is_contiguous(), "beta must be contiguous"); - TORCH_CHECK( - beta.size(0) == packed_seq && beta.size(1) == num_v_heads, "beta shape must be [packed_seq, num_v_heads]"); - } - if (input_state_.has_value()) { - auto& input_state = input_state_.value(); - TORCH_CHECK(input_state.dtype() == torch::kFloat32, "input_state must be float32"); - TORCH_CHECK(input_state.is_contiguous(), "input_state must be contiguous"); - // In CP mode the leading dim is the CP-chunk count (== num_seqs), otherwise it is the raw num_seqs. - TORCH_CHECK( - input_state.dim() == 4 && input_state.size(0) == num_seqs && input_state.size(1) == num_v_heads && - input_state.size(2) == head_size && input_state.size(3) == head_size, - "input_state shape must be [num_seqs, num_v_heads, head_size, head_size]"); - input_state_ptr = input_state.data_ptr(); - } - - // Auto-compute scale if 0 - if (scale == 0.0f) { - scale = 1.0f / std::sqrt(static_cast(head_size)); - } - - auto stream = at::cuda::getCurrentCUDAStream(); - auto sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; - - using bf16 = cute::bfloat16_t; - using Sm90 = cutlass::arch::Sm90; - - bool beta_is_bf16 = beta_.has_value() && beta_.value().dtype() == torch::kBFloat16; - - if (beta_is_bf16) { - kda::sm90::launch_kda_fwd_prefill_kernel( - stream, - reinterpret_cast(output.data_ptr()), - output_state_ptr, - reinterpret_cast(q.data_ptr()), - reinterpret_cast(k.data_ptr()), - reinterpret_cast(v.data_ptr()), - input_state_ptr, - alpha_ptr, - reinterpret_cast(beta_.value().data_ptr()), - cu_seqlens.data_ptr(), - workspace_buffer.data_ptr(), - static_cast(num_seqs), - static_cast(num_qk_heads), - static_cast(num_v_heads), - static_cast(head_size), - static_cast(packed_seq), - scale, - safe_gate, - static_cast(sm_count), - cp_seq_map_ptr, - raw_cu_seqlens_ptr, - raw_num_seqs); - } else { - float const* beta_ptr = beta_.has_value() ? beta_.value().data_ptr() : nullptr; - kda::sm90::launch_kda_fwd_prefill_kernel( - stream, - reinterpret_cast(output.data_ptr()), - output_state_ptr, - reinterpret_cast(q.data_ptr()), - reinterpret_cast(k.data_ptr()), - reinterpret_cast(v.data_ptr()), - input_state_ptr, - alpha_ptr, - beta_ptr, - cu_seqlens.data_ptr(), - workspace_buffer.data_ptr(), - static_cast(num_seqs), - static_cast(num_qk_heads), - static_cast(num_v_heads), - static_cast(head_size), - static_cast(packed_seq), - scale, - safe_gate, - static_cast(sm_count), - cp_seq_map_ptr, - raw_cu_seqlens_ptr, - raw_num_seqs); - } - - return {output, output_state}; -} diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu deleted file mode 100644 index d280b822..00000000 --- a/csrc/api/pybind.cu +++ /dev/null @@ -1,62 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#include -#include -#include -#include - -#if defined(CULA_SM90A_ENABLED) -std::tuple> -kda_fwd_prefill( - std::optional output_, - std::optional output_state_, - torch::Tensor const& q, - torch::Tensor const& k, - torch::Tensor const& v, - std::optional input_state_, - std::optional alpha_, - std::optional beta_, - torch::Tensor const& cu_seqlens, - torch::Tensor workspace_buffer, - float scale, - bool output_final_state, - bool safe_gate, - std::optional cp_seq_map_, - std::optional raw_cu_seqlens_); -#endif - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "cuLA"; -#if defined(CULA_SM90A_ENABLED) - m.def( - "kda_fwd_prefill", - &kda_fwd_prefill, - pybind11::arg("output_"), - pybind11::arg("output_state_"), - pybind11::arg("q"), - pybind11::arg("k"), - pybind11::arg("v"), - pybind11::arg("input_state_"), - pybind11::arg("alpha_"), - pybind11::arg("beta_"), - pybind11::arg("cu_seqlens"), - pybind11::arg("workspace_buffer"), - pybind11::arg("scale"), - pybind11::arg("output_final_state"), - pybind11::arg("safe_gate"), - pybind11::arg("cp_seq_map_") = std::nullopt, - pybind11::arg("raw_cu_seqlens_") = std::nullopt); -#endif -} diff --git a/csrc/api/pybind_sm90.cu b/csrc/api/pybind_sm90.cu deleted file mode 100644 index 22164b17..00000000 --- a/csrc/api/pybind_sm90.cu +++ /dev/null @@ -1 +0,0 @@ -#include "pybind.cu" diff --git a/csrc/cutlass b/csrc/cutlass deleted file mode 160000 index 73c59c05..00000000 --- a/csrc/cutlass +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 73c59c055c0fec87792470dbf33325158113db5e diff --git a/csrc/kda/sm90/collective/common.hpp b/csrc/kda/sm90/collective/common.hpp deleted file mode 100644 index 57fcb0de..00000000 --- a/csrc/kda/sm90/collective/common.hpp +++ /dev/null @@ -1,477 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include -#include - -namespace kda::sm90::collective { - -using namespace cute; - -template -CUTE_DEVICE void -gemm_reset_zero_acc(Atom& atom, TA const& tA, TB const& tB, TC&& tC) { - constexpr int rA = decltype(rank(tA))::value; - constexpr int rB = decltype(rank(tB))::value; - constexpr int rC = decltype(rank(tC))::value; - if constexpr (rA == 2 && rB == 2 && rC == 1) { - CUTE_UNROLL - for (int k_block = 0; k_block < size<1>(tA); k_block++) { - cute::gemm(atom, tA(_, k_block), tB(_, k_block), tC); - atom.accumulate_ = GMMA::ScaleOut::One; - } - } else { - static_assert(rA == 3 && rB == 3 && rC == 3); - CUTE_UNROLL - for (int k_block = 0; k_block < size<2>(tA); k_block++) { - cute::gemm(atom, tA(_, _, k_block), tB(_, _, k_block), tC); - atom.accumulate_ = GMMA::ScaleOut::One; - } - } -} - -template -CUTE_DEVICE void -gemm_zero_acc(Atom& atom, TA const& tA, TB const& tB, TC&& tC) { - atom.accumulate_ = GMMA::ScaleOut::Zero; - gemm_reset_zero_acc(atom, tA, tB, tC); -} - -template < - template - class Primitive, - cute::GMMA::Major tA, - cute::GMMA::Major tB, - cute::GMMA::ScaleIn sA, - cute::GMMA::ScaleIn sB> -CUTE_DEVICE constexpr auto -convert_to_gmma_rs(cute::MMA_Atom> const& tiled_mma) { - using Atom = cute::MMA_Atom>; - using ElementA = typename Atom::ValTypeA; - using ElementB = typename Atom::ValTypeB; - using ElementC = typename Atom::ValTypeC; - using Shape_MNK = typename Atom::Shape_MNK; - using RS = decltype(cute::GMMA::rs_op_selector()); - return cute::MMA_Atom{}; -} - -template < - template - class Primitive, - cute::GMMA::ScaleIn sA, - cute::GMMA::ScaleIn sB> -CUTE_DEVICE constexpr auto -convert_to_gmma_rs(cute::MMA_Atom> const& tiled_mma) { - using Atom = cute::MMA_Atom>; - using ElementA = typename Atom::ValTypeA; - using ElementB = typename Atom::ValTypeB; - using ElementC = typename Atom::ValTypeC; - using Shape_MNK = typename Atom::Shape_MNK; - constexpr auto tA = cute::GMMA::Major::K; - constexpr auto tB = cute::GMMA::Major::K; - using RS = decltype(cute::GMMA::rs_op_selector()); - return cute::MMA_Atom{}; -} - -template -CUTE_DEVICE constexpr auto -convert_to_gmma_rs(cute::TiledMMA const& tiled_mma) { - return cute::TiledMMA{}; -} - -template -CUTE_DEVICE constexpr auto -convert_c_layout_to_a_layout(CLayout const& c, AValueShape const& a) { - return make_layout( - make_shape(a, shape<1>(c), make_shape(shape<2>(c), size<0>(c) / size(a))), - make_stride(stride<0>(c), stride<1>(c), make_stride(stride<2>(c), size<2>(a) * stride<0, 2>(c)))); -} - -template -CUTE_DEVICE constexpr auto -unstage_smem_layout(Layout const& layout, Stages stages = {}) { - return composition(layout, make_tuple(_, _, make_layout(stages))); -} - -template -CUTE_DEVICE auto -make_acc_into_op(Accumulator const& acc, OperandLayout_TV const& operand_layout_tv) { - Tensor operand = - make_fragment_like(convert_c_layout_to_a_layout(acc.layout(), shape<1>(operand_layout_tv))); - Tensor operand_as_acc = make_tensor(operand.data(), acc.layout()); - - cute::copy(acc, operand_as_acc); - - if constexpr (sizeof(Element) == 1) { - // 00 11 22 33 00 11 22 33 acc layout - // 00 00 11 11 22 22 33 33 operand layout - // BB AA AA BB AA BB BB AA conflict-free exchange pattern - // 16-bit exchange; so process two at a time potentially - int tid = threadIdx.x % 4; - auto values_u32 = recast(operand); - - CUTE_UNROLL - for (int n = 0; n < size<1>(values_u32); n++) { - CUTE_UNROLL - for (int k = 0; k < size<2>(values_u32); k++) { - CUTE_UNROLL - for (int ii = 0; ii < 8; ii += 4) { - uint32_t values_tmp_0 = values_u32(ii / 2 + 0, n, k); - uint32_t values_tmp_1 = values_u32(ii / 2 + 1, n, k); - - // step A: - // t 1 v 0 -> t 0 v 1 - // t 2 v 0 -> t 1 v 0 - // t 0 v 1 -> t 2 v 0 - // t 3 v 1 -> t 3 v 1 - - int v_to_send = tid == 1 || tid == 2 ? 0 : 1; - int v_to_recv = v_to_send; - int t_to_recv_from = (0x3021 >> (tid * 4)) & 0xF; - - uint32_t values_tmp_a = v_to_send == 0 ? values_tmp_0 : values_tmp_1; - - values_tmp_a = __shfl_sync(0xFFFFFFFF, values_tmp_a, t_to_recv_from, 4); - - // step B: - // t 0 v 0 -> t 0 v 0 - // t 3 v 0 -> t 1 v 1 - // t 1 v 1 -> t 2 v 1 - // t 2 v 1 -> t 3 v 0 - - v_to_send = 1 - v_to_send; - v_to_recv = 1 - v_to_recv; - t_to_recv_from = (0x2130 >> (tid * 4)) & 0xF; - - uint32_t values_tmp_b = v_to_send == 0 ? values_tmp_0 : values_tmp_1; - - values_tmp_b = __shfl_sync(0xFFFFFFFF, values_tmp_b, t_to_recv_from, 4); - - values_u32(ii / 2 + 0, n, k) = - __byte_perm(values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x1054 : 0x5410); - values_u32(ii / 2 + 1, n, k) = - __byte_perm(values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x3276 : 0x7632); - } - } - } - } - - return operand; -} - -// Convert float register values from BF16 MMA operand A layout to TF32 MMA operand A layout. -// -// Both SM80_16x8x8_F32BF16BF16F32_TN and SM80_16x8x8_F32TF32TF32F32_TN have the same -// per-thread fragment shape: ((_2,_2),_1,_4):((_1,_2),_0,_4) → 16 values per thread. -// But they map different (M, K) positions to each thread. -// -// BF16 LayoutA_TV: ((_4,_8),(_2,_2)):((_32,_1),(_16,_8)) -// t0 = tid % 4, t1 = tid / 4 (note: Shape<_4,_8> → colmajor → t0 = tid%4) -// v_flat = v0 + v1*2: v0 → K offset (stride 16), v1 → M offset (stride 8) -// Thread t0 holds K = {2*t0, 2*t0+1} (consecutive K per thread) -// -// TF32 LayoutA_TV: ((_4,_8),(_2,_2)):((_16,_1),(_8,_64)) -// v_flat = v0 + v1*2: v0 → M offset (stride 8), v1 → K offset (stride 64) -// Thread t0 holds K = {t0, t0+4} (stride-4 K per thread) -// -// Algorithm (two-phase shuffle): -// For each v0_tf32 (M selector in TF32), the source data is at v1_bf16 = v0_tf32 -// in BF16 layout. Both v1_tf32 outputs need the same BF16 value index but from -// different source threads. -// -// Phase 1: shuffle bf16_frag[idx0] (v0_bf16=0) from source thread -// Phase 2: shuffle bf16_frag[idx1] (v0_bf16=1) from source thread -// Then select phase1 or phase2 result based on t0%2. -// -// Supports BK=32 (NumKAtoms=4, 16 values) and BK=64 (NumKAtoms=8, 32 values). -// The shuffle pattern is identical per k-atom; only the loop count changes. -// -// @param frag_A Float tensor with NumKAtoms*4 values in BF16 MMA A layout, converted -// in-place to TF32 MMA A layout. Values remain as float; the TF32 MMA -// hardware will truncate mantissa bits automatically during execution. -// Caller uses recast(frag_A) to obtain a typed view if needed. -// @param local_thread_idx Thread index within the MMA tile (0..63) -template -CUTE_DEVICE void -convert_bf16_to_tf32_operandA_layout(FragA& frag_A, int local_thread_idx) { - static_assert(cute::is_tensor>::value); - static_assert(NumKAtoms == 4 || NumKAtoms == 8, "Only BK=32 (4 k-atoms) and BK=64 (8 k-atoms) supported"); - static_assert(decltype(size(frag_A))::value == NumKAtoms * 4); - // Fragment must hold float values (gated results are already in float). - // MMA hardware will truncate to tf32 precision automatically. - using ElemType = typename cute::remove_cvref_t::value_type; - static_assert(cute::is_same_v, "Fragment must be float; tf32 truncation is done by MMA hw"); - - int tid = local_thread_idx % 32; // lane within warp - int t0 = tid % 4; - bool sel_odd = (t0 & 1); // t0%2: selects v0_bf16=1 result - - // Source lane for v1_tf32=0: t0_src = t0/2, lane = t0_src + (tid & ~3) - // Source lane for v1_tf32=1: t0_src = t0/2+2, lane = (t0/2+2) + (tid & ~3) - int src_lane_lo = (t0 / 2) + (tid & ~3); - int src_lane_hi = (t0 / 2 + 2) + (tid & ~3); - - // Process NumKAtoms k-iterations, each with 4 values: [4j+0, 4j+1, 4j+2, 4j+3] - // BF16 fragment layout per k-iter: (v0_bf16=0,v1_bf16=0), (v0_bf16=1,v1_bf16=0), - // (v0_bf16=0,v1_bf16=1), (v0_bf16=1,v1_bf16=1) - // TF32 output layout per k-iter: (v0_tf32=0,v1_tf32=0), (v0_tf32=1,v1_tf32=0), - // (v0_tf32=0,v1_tf32=1), (v0_tf32=1,v1_tf32=1) - CUTE_UNROLL - for (int j = 0; j < NumKAtoms; j++) { - // Read all 4 input values for this k-iter before writing any output, - // to avoid read-after-write hazard (in-place update). - float in0 = frag_A(0 + 4 * j); // v0_bf16=0, v1_bf16=0 - float in1 = frag_A(1 + 4 * j); // v0_bf16=1, v1_bf16=0 - float in2 = frag_A(2 + 4 * j); // v0_bf16=0, v1_bf16=1 - float in3 = frag_A(3 + 4 * j); // v0_bf16=1, v1_bf16=1 - - // For v0_tf32=0: M is selected by v1_bf16=0, so source values are (in0, in1) - // For v0_tf32=1: M is selected by v1_bf16=1, so source values are (in2, in3) - float out_vals[4]; - CUTE_UNROLL - for (int v0_tf32 = 0; v0_tf32 < 2; v0_tf32++) { - float val0 = (v0_tf32 == 0) ? in0 : in2; // v0_bf16=0 at chosen v1_bf16 - float val1 = (v0_tf32 == 0) ? in1 : in3; // v0_bf16=1 at chosen v1_bf16 - - // Shuffle to get values from source threads - float recv0_lo = __shfl_sync(0xFFFFFFFF, val0, src_lane_lo); - float recv1_lo = __shfl_sync(0xFFFFFFFF, val1, src_lane_lo); - float recv0_hi = __shfl_sync(0xFFFFFFFF, val0, src_lane_hi); - float recv1_hi = __shfl_sync(0xFFFFFFFF, val1, src_lane_hi); - - // Select based on t0%2: even → v0_bf16=0, odd → v0_bf16=1 - out_vals[v0_tf32 + 0] = sel_odd ? recv1_lo : recv0_lo; // v1_tf32=0 - out_vals[v0_tf32 + 2] = sel_odd ? recv1_hi : recv0_hi; // v1_tf32=1 - } - - // Write all 4 output values - frag_A(0 + 4 * j) = out_vals[0]; - frag_A(1 + 4 * j) = out_vals[1]; - frag_A(2 + 4 * j) = out_vals[2]; - frag_A(3 + 4 * j) = out_vals[3]; - } -} - -// Convert float register values from BF16 MMA operand B layout to TF32 MMA operand B layout. -// -// BF16 LayoutB_TV: ((_4,_8),_2):((_16,_1),_8) -// Thread t0 holds K = {2*t0, 2*t0+1}. v selects K offset (consecutive). -// -// TF32 LayoutB_TV: ((_4,_8),_2):((_8,_1),_32) -// Thread t0 holds K = {t0, t0+4}. v selects K offset (stride-4). -// -// Same two-phase shuffle approach as operand A, but B has only 2 values per atom -// (no M dimension in the value index). -// -// Supports BK=32 (NumKAtoms=4, 8 values) and BK=64 (NumKAtoms=8, 16 values). -// -// @param frag_B Float tensor with NumKAtoms*2 values in BF16 MMA B layout, converted -// in-place. Values remain as float; TF32 MMA hardware truncates automatically. -// Caller uses recast(frag_B) to obtain a typed view if needed. -// @param local_thread_idx Thread index within the MMA tile (0..63) -template -CUTE_DEVICE void -convert_bf16_to_tf32_operandB_layout(FragB& frag_B, int local_thread_idx) { - static_assert(cute::is_tensor>::value); - static_assert(NumKAtoms == 4 || NumKAtoms == 8, "Only BK=32 (4 k-atoms) and BK=64 (8 k-atoms) supported"); - static_assert(decltype(size(frag_B))::value == NumKAtoms * 2); - // Fragment must hold float values; MMA hardware truncates to tf32 automatically. - using ElemType = typename cute::remove_cvref_t::value_type; - static_assert(cute::is_same_v, "Fragment must be float; tf32 truncation is done by MMA hw"); - - int tid = local_thread_idx % 32; - int t0 = tid % 4; - bool sel_odd = (t0 & 1); - - int src_lane_lo = (t0 / 2) + (tid & ~3); - int src_lane_hi = (t0 / 2 + 2) + (tid & ~3); - - // Process NumKAtoms k-iterations, each with 2 values: [2j, 2j+1] - CUTE_UNROLL - for (int j = 0; j < NumKAtoms; j++) { - int idx0 = 2 * j; // BF16 v=0, K = 2*t0 - int idx1 = 2 * j + 1; // BF16 v=1, K = 2*t0+1 - - float val0 = frag_B(idx0); - float val1 = frag_B(idx1); - - // v_tf32=0: need K=t0 from src_t0=t0/2 - float recv0_lo = __shfl_sync(0xFFFFFFFF, val0, src_lane_lo); - float recv1_lo = __shfl_sync(0xFFFFFFFF, val1, src_lane_lo); - // v_tf32=1: need K=t0+4 from src_t0=t0/2+2 - float recv0_hi = __shfl_sync(0xFFFFFFFF, val0, src_lane_hi); - float recv1_hi = __shfl_sync(0xFFFFFFFF, val1, src_lane_hi); - - frag_B(idx0) = sel_odd ? recv1_lo : recv0_lo; - frag_B(idx1) = sel_odd ? recv1_hi : recv0_hi; - } -} - -// Broadcast row 0 from a BF16 MMA operand A fragment and output directly -// into a BF16 MMA operand B fragment. -// -// Combines broadcast_row0 + extract_A_to_B into one step, avoiding the -// intermediate 16-float operand A broadcast tensor (saves 8 float registers). -// -// Since g_first is broadcast (all M rows identical), operand B only needs -// the K-dimension values. We shuffle v1=0 values from the t1=0 thread -// and output the 8-value B fragment directly. -// -// BF16 A layout: t0 = tid % 4, t1 = tid / 4. Row 0 at t1=0, v1=0. -// frag_A(4j+0) = (v0=0, v1=0), frag_A(4j+1) = (v0=1, v1=0) -// BF16 B layout: frag_B(2j+0) = v=0, frag_B(2j+1) = v=1 -// Both have K = {2*t0, 2*t0+1} at same positions. -// -// Supports BK=32 (NumKAtoms=4) and BK=64 (NumKAtoms=8). -// -// Cost: 2 shuffles per k-iter × NumKAtoms k-iters. -// Saves: NumKAtoms*4-float intermediate tensor (vs broadcast_row0 + extract). -// -// @param frag_A Input: alpha[m, k] in BF16 MMA A layout (NumKAtoms*4 values) -// @param frag_B_first Output: alpha[0, k] in BF16 MMA B layout (NumKAtoms*2 values) -// @param local_thread_idx Thread index within the MMA tile (0..63) -template -CUTE_DEVICE void -broadcast_row0_operandA_to_operandB_bf16_layout(FragA const& frag_A, FragB& frag_B_first, int local_thread_idx) { - static_assert(cute::is_tensor>::value); - static_assert(cute::is_tensor>::value); - static_assert(NumKAtoms == 4 || NumKAtoms == 8, "Only BK=32 (4 k-atoms) and BK=64 (8 k-atoms) supported"); - static_assert(decltype(size(frag_A))::value == NumKAtoms * 4); - static_assert(decltype(size(frag_B_first))::value == NumKAtoms * 2); - - int tid = local_thread_idx % 32; // lane within warp - // Row 0 is at t1=0. In BF16 A layout, t0 = tid % 4, t1 = tid / 4. - // Source lane: same t0, t1=0 → src = tid % 4. - int src_lane = tid % 4; - - CUTE_UNROLL - for (int j = 0; j < NumKAtoms; j++) { - // Shuffle v1=0 values from thread with t1=0 (row 0 holder) - // frag_A(4j+0) = (v0=0, v1=0) → K=2*t0 - // frag_A(4j+1) = (v0=1, v1=0) → K=2*t0+1 - auto val0 = __shfl_sync(0xFFFFFFFF, frag_A(4 * j + 0), src_lane); - auto val1 = __shfl_sync(0xFFFFFFFF, frag_A(4 * j + 1), src_lane); - - // Output directly into B layout: frag_B(2j) = K_lo, frag_B(2j+1) = K_hi - frag_B_first(2 * j + 0) = val0; - frag_B_first(2 * j + 1) = val1; - } -} - -// Broadcast row 0 across all M rows in a BF16 MMA operand A fragment. -// -// Given frag_A holding alpha[m, k] per thread, produces frag_A_first holding alpha[0, k] -// for all m (broadcast). This eliminates a redundant S2R load of g_first from shared memory. -// -// BF16 LayoutA_TV: ((_4,_8),(_2,_2)):((_32,_1),(_16,_8)) -// tid decomposition: t0 = tid % 4, t1 = tid / 4 -// m = t1 + v1*8 (v1 selects M row within thread) -// Row 0 is held by threads with t1=0, at v1=0 positions. -// -// Algorithm: -// For each k-iter, shuffle v1=0 values from thread (tid % 4) (same t0, t1=0), -// then replicate to v1=1 positions. -// -// Supports BK=32 (NumKAtoms=4) and BK=64 (NumKAtoms=8). -// -// Cost: 2 shuffles per k-iter × NumKAtoms k-iters. -// -// @param frag_A Input: alpha[m, k] in BF16 MMA A layout (NumKAtoms*4 values) -// @param frag_A_first Output: alpha[0, k] broadcast in BF16 MMA A layout (NumKAtoms*4 values) -// @param local_thread_idx Thread index within the MMA tile (0..63) -template -CUTE_DEVICE void -broadcast_row0_operandA_bf16_layout(FragA const& frag_A, FragAFirst& frag_A_first, int local_thread_idx) { - static_assert(cute::is_tensor>::value); - static_assert(cute::is_tensor>::value); - static_assert(NumKAtoms == 4 || NumKAtoms == 8, "Only BK=32 (4 k-atoms) and BK=64 (8 k-atoms) supported"); - static_assert(decltype(size(frag_A))::value == NumKAtoms * 4); - static_assert(decltype(size(frag_A_first))::value == NumKAtoms * 4); - - int tid = local_thread_idx % 32; // lane within warp - // Row 0 is at t1=0 → tid % 4 (keep same t0, set t1=0) - int src_lane = tid % 4; - - CUTE_UNROLL - for (int j = 0; j < NumKAtoms; j++) { - // v1=0 positions hold m=t1, v1=1 positions hold m=t1+8 - // We want m=0, which is at t1=0, v1=0 → indices 4j+0 and 4j+1 - auto val0 = __shfl_sync(0xFFFFFFFF, frag_A(4 * j + 0), src_lane); // alpha[0, 2*t0] from t1=0 - auto val1 = __shfl_sync(0xFFFFFFFF, frag_A(4 * j + 1), src_lane); // alpha[0, 2*t0+1] from t1=0 - - // Broadcast to both v1=0 and v1=1 (same value, different M rows) - frag_A_first(4 * j + 0) = val0; // v0=0, v1=0 - frag_A_first(4 * j + 1) = val1; // v0=1, v1=0 - frag_A_first(4 * j + 2) = val0; // v0=0, v1=1 (broadcast) - frag_A_first(4 * j + 3) = val1; // v0=1, v1=1 (broadcast) - } -} - -// Extract BF16 MMA operand B fragment from a BF16 MMA operand A fragment, -// for data that is **broadcast across M rows** (e.g., g_first = g[row=0, :]). -// -// When the source data is broadcast (identical for all M rows), the K-dimension -// mapping is the same in both A and B BF16 MMA layouts: -// A: thread t0 holds K = {2*t0, 2*t0+1} at v0={0,1}, with v1 selecting M row -// B: thread t0 holds K = {2*t0, 2*t0+1} at v={0,1} -// -// Since M rows are identical (broadcast), we can simply pick v1=0 from A: -// frag_B(2j + 0) = frag_A(4j + 0) // v0_bf16=0 → K=2*t0 -// frag_B(2j + 1) = frag_A(4j + 1) // v0_bf16=1 → K=2*t0+1 -// -// Supports BK=32 (NumKAtoms=4) and BK=64 (NumKAtoms=8). -// -// This avoids a redundant S2R load from shared memory. -// No warp shuffles needed — purely register-local extraction. -// -// @param frag_A Float tensor with NumKAtoms*4 values in BF16 MMA A layout (broadcast data) -// @param frag_B Float tensor with NumKAtoms*2 values in BF16 MMA B layout (output) -template -CUTE_DEVICE void -extract_broadcast_operandA_to_operandB_bf16_layout(FragA const& frag_A, FragB& frag_B) { - static_assert(cute::is_tensor>::value); - static_assert(cute::is_tensor>::value); - static_assert(NumKAtoms == 4 || NumKAtoms == 8, "Only BK=32 (4 k-atoms) and BK=64 (8 k-atoms) supported"); - static_assert(decltype(size(frag_A))::value == NumKAtoms * 4); - static_assert(decltype(size(frag_B))::value == NumKAtoms * 2); - - CUTE_UNROLL - for (int j = 0; j < NumKAtoms; j++) { - // A layout per k-iter: [4j+0]=(v0=0,v1=0), [4j+1]=(v0=1,v1=0), [4j+2]=(v0=0,v1=1), [4j+3]=(v0=1,v1=1) - // B layout per k-iter: [2j+0]=v=0, [2j+1]=v=1 - // For broadcast data, v1 doesn't matter, so pick v1=0: - frag_B(2 * j + 0) = frag_A(4 * j + 0); // K = 2*t0 - frag_B(2 * j + 1) = frag_A(4 * j + 1); // K = 2*t0+1 - } -} - -} // namespace kda::sm90::collective diff --git a/csrc/kda/sm90/collective/load_predicated.hpp b/csrc/kda/sm90/collective/load_predicated.hpp deleted file mode 100644 index 8071a440..00000000 --- a/csrc/kda/sm90/collective/load_predicated.hpp +++ /dev/null @@ -1,193 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include -#include -#include - -#include "kda/sm90/utils/debug.hpp" -#include "kda/sm90/utils/unused.hpp" - -namespace kda::sm90::collective { - -using namespace cute; - -// Wraps a callable into a predicate "tensor" usable by CuTe's copy_if. -// copy_if calls pred(i) with a linear index; the wrapped function maps that index to bool. -template -struct FunctionPredTensor { - Fn fn_; - CUTE_HOST_DEVICE - FunctionPredTensor(Fn fn) : fn_(fn) { - } - template - CUTE_HOST_DEVICE bool - operator()(Idx const& i) const { - return fn_(i); - } -}; - -enum class LoadKindVector { - kAlpha, - kBeta, -}; - -CUTE_HOST_DEVICE constexpr char const* -to_string(LoadKindVector kind) { - if (kind == LoadKindVector::kAlpha) { - return "alpha"; - } else if (kind == LoadKindVector::kBeta) { - return "beta"; - } else { - return "unknown loadkind"; - } -} - -template < - LoadKindVector kKind, - class Pipeline, - class ElementSrc, - class GmemLayout, - class ElementDst, - class SmemLayout, - class VectorProcessor_ = Unused> -struct CollectiveLoadVector { - using SharedStorage = cute::array_aligned>; - using PipelineState = typename cutlass::PipelineState; - - using VectorProcessor = VectorProcessor_; - - static_assert(rank_v == 2 || rank_v == 3); - - static constexpr LoadKindVector kind = kKind; - static constexpr int VectorSize = size<0>(SmemLayout{}); - - CUTE_DEVICE - CollectiveLoadVector( - ElementSrc const* src, GmemLayout layout, ElementSrc oob_value, Pipeline& pipeline, SharedStorage& storage) - : src_(src), src_layout_(layout), src_oob_value_(oob_value), pipeline_(pipeline), storage_(storage) { - } - - template - CUTE_DEVICE auto - partition_SD(ProblemSize const& problem_size, TileShape const& tile_shape, WorkDesc const& work_desc) { - constexpr auto BlkSeqQ = decltype(get<0>(tile_shape))::value; - - Tensor g = [&] { - auto head_idx = work_desc.o_head_idx(); - DPRINTF0_W( - "slice view GMEM %s: seq_idx:%d head_idx:%d tok_offset:%lld\n", - to_string(kind), - work_desc.seq_idx, - head_idx, - work_desc.tok_offset); - Tensor m_varlen_head = make_tensor(make_gmem_ptr(src_), src_layout_); - - Tensor m_varlen = m_varlen_head(_, head_idx); // slice into current head_idx - Tensor m_offset = - domain_offset(make_coord(work_desc.tok_offset), m_varlen); // offset to start of the current sequence - Tensor g_full = flat_divide(m_offset, BlkSeqQ); // (blk, iter_blk) - return g_full; - }(); - // (blk, pipe) or (blk, pipe, N), N for feature rich preprocess, data will be stored at 0 - Tensor s = make_tensor(make_smem_ptr(storage_.data()), SmemLayout{}); - - auto thr_layout = Layout<_32>{}; - auto val_layout = Layout<_1>{}; - auto tiled_copy = make_tiled_copy(Copy_Atom, ElementDst>{}, thr_layout, val_layout); - auto thr_copy = tiled_copy.get_thread_slice(cutlass::canonical_lane_idx()); - - auto coord = thr_copy.partition_S(make_identity_tensor(Shape, _1>{})); - int seq_len = work_desc.chunk_len(); - auto len_of_last_blk = seq_len - (ceil_div(seq_len, BlkSeqQ) - 1) * BlkSeqQ; - - auto mask = FunctionPredTensor([coord, len_of_last_blk](auto frag_coord) { - auto coord_in_blk = get<0>(coord(frag_coord)); - return coord_in_blk < len_of_last_blk; - }); - - auto src = thr_copy.partition_S(g); // (cpy, iter_cpy, iter_blk) - auto dst = thr_copy.partition_D(s); // (cpy, iter_cpy, pipe) - - return make_tuple(src, dst, mask); - } - - template - CUTE_DEVICE void - step(SrcDst const& src_dst, int src_iter, PipelineState& dst_pipe, int num_iters, VectorProcessor processor = {}) { - auto src = get<0>(src_dst); - auto dst = get<1>(src_dst); - - auto regs = make_fragment_like(take<0, 2>(shape(dst))); - if constexpr (!IsTail) { - copy(src(_, _, src_iter), regs); - } else { - auto mask = get<2>(src_dst); - fill(regs, src_oob_value_); - copy_if(mask, src(_, _, src_iter), regs); - } - - int dst_pipe_idx = dst_pipe.index(); - - DPRINTF0_WG("%s pipeline.producer_acquire smem_pipe_write:%d\n", to_string(kind), dst_pipe_idx); - pipeline_.producer_acquire(dst_pipe); - cutlass::arch::fence_view_async_shared(); - - if constexpr (rank_v == 3) { - copy(regs, dst(_, _, _0{}, dst_pipe_idx)); - } else { - copy(regs, dst(_, _, dst_pipe_idx)); - } - - Tensor s = make_tensor(make_smem_ptr(storage_.data()), SmemLayout{}); - if constexpr (!std::is_same_v) { - if constexpr (rank_v == 3) { - processor(s(_, _, dst_pipe_idx)); - } else { - processor(s(_, dst_pipe_idx)); - } - } - - cutlass::arch::fence_view_async_shared(); - pipeline_.producer_commit(dst_pipe); - ++dst_pipe; - } - - private: - ElementSrc const* src_; - GmemLayout src_layout_; // in (packed_seq, H) coordinate - ElementSrc src_oob_value_; - Pipeline& pipeline_; - SharedStorage& storage_; -}; - -} // namespace kda::sm90::collective diff --git a/csrc/kda/sm90/collective/load_tma.hpp b/csrc/kda/sm90/collective/load_tma.hpp deleted file mode 100644 index 1d427b0c..00000000 --- a/csrc/kda/sm90/collective/load_tma.hpp +++ /dev/null @@ -1,177 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include -#include -#include - -#include "kda/sm90/utils/debug.hpp" - -namespace kda::sm90::collective { - -using namespace cute; - -enum class LoadKind { - kQ, - kK, - kV, - kAlpha, -}; - -CUTE_HOST_DEVICE constexpr char const* -to_string(LoadKind kind) { - if (kind == LoadKind::kQ) { - return "Q"; - } else if (kind == LoadKind::kK) { - return "K"; - } else if (kind == LoadKind::kV) { - return "V"; - } else if (kind == LoadKind::kAlpha) { - return "Alpha"; - } else { - return "unknown loadkind"; - } -} - -template -struct CollectiveLoadTma { - using SharedStorage = cute::array_aligned>; - using PipelineState = typename cutlass::PipelineState; - - static constexpr LoadKind kind = kKind; - TMA const& tma_load; - Pipeline& pipeline; - SharedStorage& storage; - - CUTE_DEVICE - CollectiveLoadTma(TMA const& tma_load, Pipeline& pipeline, SharedStorage& storage) - : tma_load(tma_load), pipeline(pipeline), storage(storage) { - } - - template - CUTE_DEVICE auto - partition_SD(ProblemSize const& problem_size, TileShape const& tile_shape, WorkDesc const& work_desc) { - constexpr auto BlkSeqQ = decltype(get<0>(tile_shape))::value; - constexpr auto BlkSeqKV = decltype(get<1>(tile_shape))::value; - constexpr auto HeadSize = decltype(get<2>(tile_shape))::value; - - Tensor g = [&] { - if constexpr (kind == LoadKind::kQ) { - DPRINTF0_W( - "slice view GMEM %s: seq_idx:%d head_idx:%d tok_offset:%lld\n", - to_string(kind), - work_desc.seq_idx, - work_desc.q_head_idx(), - work_desc.tok_offset); - // Q lives in the QK head space. - Tensor m_varlen_head = tma_load.get_tma_tensor(make_shape( - problem_size.total_seqlen, - problem_size.head_size, - problem_size.num_qk_heads)); // global view to the packed varlen sequence - Tensor m_varlen = m_varlen_head(_, _, work_desc.q_head_idx()); // slice into current head_idx - Tensor m_offset = domain_offset( - make_coord(work_desc.tok_offset, _0{}), - m_varlen); // offset to start of the current sequence - Tensor g_full = - local_tile(m_offset, make_tile(BlkSeqQ, HeadSize), make_coord(_, _0{})); // (blk, d, iter_blk) - return g_full; - } else if constexpr (kind == LoadKind::kAlpha) { - // Alpha (gate) is per V/O head under GVA. - DPRINTF0_W( - "slice view GMEM %s: seq_idx:%d head_idx:%d tok_offset:%lld\n", - to_string(kind), - work_desc.seq_idx, - work_desc.o_head_idx(), - work_desc.tok_offset); - Tensor m_varlen_head = tma_load.get_tma_tensor(make_shape( - problem_size.total_seqlen, - problem_size.head_size, - problem_size.num_v_heads)); // global view to the packed varlen sequence - Tensor m_varlen = m_varlen_head(_, _, work_desc.o_head_idx()); // slice into current head_idx - Tensor m_offset = domain_offset( - make_coord(work_desc.tok_offset, _0{}), - m_varlen); // offset to start of the current sequence - Tensor g_full = - local_tile(m_offset, make_tile(BlkSeqQ, HeadSize), make_coord(_, _0{})); // (blk, d, iter_blk) - return g_full; - } else { - // K lives in the QK head space; V lives in the V head space. - // `kind` is a static constexpr LoadKind, so the head-count selection collapses at compile time. - constexpr bool kIsK = (kind == LoadKind::kK); - auto head_idx = kIsK ? work_desc.k_head_idx() : work_desc.v_head_idx(); - auto num_kv_heads = kIsK ? problem_size.num_qk_heads : problem_size.num_v_heads; - DPRINTF0_W( - "slice view GMEM %s: seq_idx:%d head_idx:%d tok_offset:%lld\n", - to_string(kind), - work_desc.seq_idx, - head_idx, - work_desc.tok_offset); - Tensor m_varlen_head = tma_load.get_tma_tensor(make_shape( - problem_size.head_size, - problem_size.total_seqlen, - num_kv_heads)); // global view to the packed varlen sequence - Tensor m_varlen = m_varlen_head(_, _, head_idx); // slice into current head_idx - Tensor m_offset = domain_offset( - make_coord(_0{}, work_desc.tok_offset), - m_varlen); // offset to start of the current sequence - Tensor g_full = - local_tile(m_offset, make_tile(HeadSize, BlkSeqKV), make_coord(_0{}, _)); // (d, blk, iter_blk) - return g_full; - } - }(); - Tensor s = make_tensor(make_smem_ptr(storage.data()), SmemLayout{}); - - auto block_tma = tma_load.get_slice(_0{}); // do not support cluster - return make_tuple(block_tma.partition_S(g), block_tma.partition_D(s)); - } - - template - CUTE_DEVICE void - step(SrcDst const& src_dst, int src_iter, PipelineState& dst_pipe, uint32_t lane_predicate) { - if (lane_predicate == 1) { - DPRINTF_WG("%s pipeline.producer_acquire smem_pipe_write:%d\n", to_string(kind), dst_pipe.index()); - if constexpr (kAcquireBarrier) { - pipeline.producer_acquire(dst_pipe); - } - using BarrierType = typename Pipeline::ProducerBarrierType; - BarrierType* tma_barrier = pipeline.producer_get_barrier(dst_pipe); - - auto src = get<0>(src_dst); - auto dst = get<1>(src_dst); - - copy(tma_load.with(*tma_barrier), src(_, _, _, src_iter), dst(_, _, _, dst_pipe.index())); - ++dst_pipe; - } - } -}; - -} // namespace kda::sm90::collective diff --git a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp deleted file mode 100644 index 66563e5a..00000000 --- a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp +++ /dev/null @@ -1,2101 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include -#include - -#include "kerutils/kerutils.cuh" - -#include "kda/sm90/collective/common.hpp" -#include "kda/sm90/collective/load_predicated.hpp" -#include "kda/sm90/collective/load_tma.hpp" -#include "kda/sm90/collective/named_barriers.hpp" -#include "kda/sm90/collective/store_tma.hpp" -#include "kda/sm90/kernel/options.hpp" -#include "kda/sm90/utils/debug.hpp" -#include "kda/sm90/utils/math_order_barrier.hpp" -#include "kda/sm90/utils/unused.hpp" - -// #define INLINE_LAMBDA [[gnu::always_inline]] -#define INLINE_LAMBDA __attribute__((always_inline)) -// #define INLINE_LAMBDA [[msvc::forceinline]] - -#define WORKAROUND_WGMMA_PERFORMANCE_LOSS() \ - if (thread_idx > 8192) { \ - __syncwarp(); \ - } - -namespace kda::sm90::collective { - -struct KdaNamedBarriers : FlatSharedNamedBarriers { - static constexpr int StateMath = FlatSharedNamedBarriers::NumBarriersUsed + 0; - static constexpr int AuxMath = FlatSharedNamedBarriers::NumBarriersUsed + 1; - static constexpr int StateMathWG0 = FlatSharedNamedBarriers::NumBarriersUsed + 2; - // NOTE: only for debug - // used for subchunk MMA with two groups, each group has 2 warps - // static constexpr int AuxMathWarp0 = FlatSharedNamedBarriers::NumBarriersUsed + 3; - // static constexpr int AuxMathWarp1 = FlatSharedNamedBarriers::NumBarriersUsed + 4; -}; - -using ku::alignment_for_swizzle; -using ku::select_layout; -using ku::select_tensor; -using namespace cute; -using kda::sm90::kernel::find_option_t; -using kda::sm90::kernel::Tag; - -template < - class Element_, - class ElementAccumulatorQK_, - class ElementAccumulatorKV_, - class TileShape_, // (seqlen_q, seqlen_kv, d) - class LayoutQ_, - class LayoutK_, - class LayoutV_, - class LayoutO_, // (seqlen_q/k, d, h) - class Options> -struct FlatMainloopTmaWarpSpecializedKdaFwd { - using Element = Element_; - using ElementAccumulatorQK = ElementAccumulatorQK_; - using ElementAccumulatorO = ElementAccumulatorQK; - using ElementAccumulatorKV = ElementAccumulatorKV_; - using ElementO = Element; - using ElementAlpha = float; - using ElementBeta = float; // SMEM + compute stays fp32 - using ElementBetaGmem = find_option_t; // GMEM type (float or bf16) - using ElementGatedMMA = cutlass::tfloat32_t; - - using TileShape = TileShape_; - - using LayoutQ = LayoutQ_; // (seqlen_q, d, h) - using LayoutK = LayoutK_; // (seqlen_k, d, h) - using LayoutV = LayoutV_; // (seqlen_k, d, h) - using LayoutO = LayoutO_; // (seqlen_k, d, h) - using LayoutAlpha = LayoutQ_; // (seqlen_q, d, h) - - // Options - static constexpr bool kIsPersistent = find_option_t::value; - - static constexpr bool kInitStateFromInput = find_option_t::value; - - static constexpr int NumLoadWarpGroups = 1; - static constexpr int NumStateMmaWarpGroups = 2; - static constexpr int NumAuxMmaWarpGroups = 1; - - static constexpr int StageCountQ = find_option_t, Options>::value; - static constexpr int StageCountK = find_option_t, Options>::value; - static constexpr int StageCountV = find_option_t, Options>::value; - - static constexpr int NeedsAlpha = find_option_t::value; - static constexpr int NeedsBeta = find_option_t::value; - static_assert(NeedsAlpha && NeedsBeta, "Alpha and Beta are both used in KDA."); - - static constexpr int SafeGate = true; // only support safe_gate=true - - static constexpr int NumLoadThreads = NumLoadWarpGroups * 128; - static constexpr int NumStateMmaThreads = NumStateMmaWarpGroups * 128; - static constexpr int NumAuxMmaThreads = NumAuxMmaWarpGroups * 128; - - static constexpr uint32_t OrderedBarrierId0 = uint32_t(cutlass::arch::ReservedNamedBarriers::StreamkBarrier0); - static constexpr uint32_t OrderedBarrierId1 = uint32_t(cutlass::arch::ReservedNamedBarriers::StreamkBarrier1); - - using OrderedMathBarriers = std::conditional_t< - NumStateMmaWarpGroups == 2, - OrderedNamedBarriers, - OrderedNamedBarriers>; - - using StagesQ = cutlass::gemm::collective::StageCount; - using StagesK = cutlass::gemm::collective::StageCount; - using StagesV = cutlass::gemm::collective::StageCount; - using StagesQ_K_Scaled = cutlass::gemm::collective::StageCount<2>; - using StagesO = cutlass::gemm::collective::StageCount<1>; - using ClusterShape = Shape<_1, _1, _1>; - - using StagesQK = cutlass::gemm::collective::StageCount<2>; - using StagesKK = cutlass::gemm::collective::StageCount<2>; - - using StagesAlpha = cutlass::gemm::collective::StageCount<2>; - using StagesBeta = cutlass::gemm::collective::StageCount<2>; - - static constexpr int Alignment = 16 / sizeof(Element); - - static constexpr auto BlkSeqQ = get<0>(TileShape{}); // Blk_Q - static constexpr auto BlkSeqKV = get<1>(TileShape{}); // Blk_K/V - static constexpr auto HeadSize = get<2>(TileShape{}); // D (Dq, Dk, Dv all equal) - static constexpr auto HeadSizeQK = HeadSize; - static constexpr auto HeadSizeV = HeadSize; - using HeadSizeHalf = _64; - using HeadSizeQuar = _32; - - using TileShapeQK = decltype(make_shape(BlkSeqQ, BlkSeqKV, HeadSizeQK)); - // used for element-wise in compute_aux prologue, to reduce register usage - using TileShapeQK_Half = decltype(make_shape(BlkSeqQ, BlkSeqKV, HeadSizeHalf{})); - using TileShapeQK_Quar = decltype(make_shape(BlkSeqQ, BlkSeqKV, HeadSizeQuar{})); - using TileShapeKK = decltype(make_shape(BlkSeqKV, BlkSeqKV, HeadSizeQK)); - using TileShapeKV = decltype(make_shape(HeadSizeV, HeadSizeQK, BlkSeqKV)); - static_assert(std::is_same_v); - - using TileShapeO2 = decltype(make_shape(HeadSizeV, BlkSeqQ, BlkSeqKV)); - using TileShapeO1 = decltype(make_shape(HeadSizeV, BlkSeqQ, HeadSizeQK)); - - static_assert(BlkSeqQ % 64 == 0); - static_assert(BlkSeqQ == 64 || BlkSeqQ == 128); - static_assert(BlkSeqQ == BlkSeqKV); - static constexpr bool IsQKCooperative = BlkSeqQ == 128; - static constexpr bool IsKKCooperative = IsQKCooperative; - - using DummyStages = cutlass::gemm::collective::StageCount<2>; - ; - using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm90, - cutlass::arch::OpClassTensorOp, - Element, - LayoutQ, - Alignment, - Element, - LayoutK, - Alignment, - ElementAccumulatorQK, - TileShapeQK, - ClusterShape, - DummyStages, - std::conditional_t< - IsQKCooperative, - cutlass::gemm::KernelTmaWarpSpecializedCooperative, - cutlass::gemm::KernelTmaWarpSpecialized>>::CollectiveOp; - - // dummy TiledMmaQK RS for S2R/R2S layout consistency - using AtomLayoutQK = Layout, _1, _1>>; - using TiledMmaQK_RS = decltype(make_tiled_mma( - decltype(cute::GMMA::rs_op_selector()){}, AtomLayoutQK{})); - static_assert(size(TiledMmaQK_RS{}) == NumAuxMmaThreads); - using TiledMmaQK_RS_Quar = decltype(make_tiled_mma( - decltype(cute::GMMA::rs_op_selector()){}, - AtomLayoutQK{})); - static_assert(size(TiledMmaQK_RS_Quar{}) == NumAuxMmaThreads); - - using CollectiveMmaKV_G2S = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm90, - cutlass::arch::OpClassTensorOp, - Element, - decltype(select<1, 0, 2>(LayoutV{})), - Alignment, // direct TMA copy for GMEM -> SMEM - Element, - decltype(select<1, 0, 2>(LayoutK{})), - Alignment, - ElementAccumulatorKV, - TileShapeKV, - ClusterShape, - DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; - - using SmemLayoutAlphaAtom = GMMA::Layout_K_SW128_Atom; - using SmemLayoutAlpha_SD = decltype(tile_to_shape( - SmemLayoutAlphaAtom{}, - make_shape( - shape<1>(TileShapeQK{}), - shape<2>(TileShapeQK{}), - Int{}))); // (blk_kv, head_size), (64, 128) - using GmemShapeAlpha = Shape; // (seqlen_k, d, h) - using GmemStrideAlpha = Stride; - using GmemLayoutAlpha = Layout; - using GmemTiledCopyAlpha = cute::SM90_TMA_LOAD; - using TMA_Alpha = decltype(make_tma_copy( - GmemTiledCopyAlpha{}, - make_tensor(make_gmem_ptr(static_cast(nullptr)), GmemLayoutAlpha{}), - take<0, 2>(SmemLayoutAlpha_SD{}), - select<1, 2>(TileShapeQK{}), - size<0>(ClusterShape{}))); - - // raw layout for copy - using SmemLayoutQ_SD = - decltype(unstage_smem_layout(typename CollectiveMmaQK::SmemLayoutA{}, Int{})); - using SmemLayoutQ_K_Scaled_SD = - decltype(unstage_smem_layout(typename CollectiveMmaQK::SmemLayoutA{}, Int{})); - using SmemLayoutK_DS = - decltype(unstage_smem_layout(typename CollectiveMmaKV_G2S::SmemLayoutB{}, Int{})); - using SmemLayoutQ_K_Scaled_DS = - decltype(unstage_smem_layout(typename CollectiveMmaKV_G2S::SmemLayoutB{}, Int{})); - using SmemLayoutV_DS = - decltype(unstage_smem_layout(typename CollectiveMmaKV_G2S::SmemLayoutA{}, Int{})); - - // Layout for V^T - using RefLayoutV = decltype(make_layout(select<0, 2>(TileShapeKV{}), LayoutRight{})); - using CollectiveMmaKV = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm90, - cutlass::arch::OpClassTensorOp, - Element, - RefLayoutV, - Alignment, // needs a S2R transposition for MMA - Element, - decltype(select<1, 0, 2>(LayoutK{})), - Alignment, - ElementAccumulatorKV, - TileShapeKV, - ClusterShape, - DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; - - using RefLayoutKV = decltype(make_layout(select<0, 1>(TileShapeKV{}), LayoutRight{})); // (dv, dk) - using CollectiveMmaO1 = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm90, - cutlass::arch::OpClassTensorOp, - Element, - RefLayoutKV, - Alignment, // NOTE: S (KV) as operand A - Element, - LayoutQ, - Alignment, - ElementAccumulatorO, - TileShapeO1, - ClusterShape, - DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; - - // (blk_q,blk_k) to align with O2 mma, LayoutRight to align with QK mma output - using DesiredLayoutQK = decltype(make_layout(select<0, 1>(TileShapeQK{}), LayoutRight{})); - using CollectiveMmaO2 = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm90, - cutlass::arch::OpClassTensorOp, - Element, - RefLayoutV, - Alignment, // V^T - Element, - DesiredLayoutQK, - Alignment, - ElementAccumulatorO, - TileShapeO2, - ClusterShape, - DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; - - using TiledMmaQK = typename CollectiveMmaQK::TiledMma; // Q@K^t - using TiledMmaKV = decltype(convert_to_gmma_rs(typename CollectiveMmaKV::TiledMma{})); - using TiledMmaO1 = decltype(convert_to_gmma_rs(typename CollectiveMmaO1::TiledMma{})); - using TiledMmaO2 = decltype(convert_to_gmma_rs(typename CollectiveMmaO2::TiledMma{})); - - static_assert(size(TiledMmaQK{}) == NumAuxMmaThreads); - - static_assert(size(TiledMmaKV{}) == NumStateMmaThreads); - static_assert(size(TiledMmaO1{}) == NumStateMmaThreads); - static_assert(size(TiledMmaO2{}) == NumStateMmaThreads); - - using CollectiveStoreO = CollectiveStoreTma< - TileShapeO1, - ClusterShape, - ElementO, - ElementAccumulatorO, - /*Seme*/ ElementO, - decltype(select<1, 0, 2>(LayoutO{})), - StagesO::value>; - - // layout for compute - using QKSmemLayoutQ = SmemLayoutQ_SD; - using QKSmemLayoutK = decltype(select_layout<1, 0, 2>(SmemLayoutK_DS{})); - using QKScaledSmemLayoutQ = SmemLayoutQ_K_Scaled_SD; - - using KVSmemLayoutK = SmemLayoutK_DS; - using KVSmemLayoutV = SmemLayoutV_DS; - using QKScaledSmemLayoutKt = SmemLayoutQ_K_Scaled_DS; - - using QKQSmemLayoutAlpha = SmemLayoutAlpha_SD; - using QKKSmemLayoutAlpha = decltype(select_layout<1, 0, 2>(SmemLayoutAlpha_SD{})); - using KKTKSmemLayoutAlpha = SmemLayoutAlpha_SD; - using KKTKTSmemLayoutAlpha = decltype(select_layout<1, 0, 2>(SmemLayoutAlpha_SD{})); - - // layout for compute output - using SmemLayoutQK = decltype(tile_to_shape( - GMMA::Layout_K_INTER_Atom{}, - flatten(make_shape(select<0, 1>(TileShapeQK{}), Int{})), - Step<_1, _2, _3>{})); - using SmemLayoutO = typename CollectiveStoreO::SmemLayoutO; - - using SmemLayoutKK = decltype(tile_to_shape( - GMMA::Layout_K_INTER_Atom{}, - flatten(make_shape(select<0, 1>(TileShapeQK{}), Int{})), - Step<_1, _2, _3>{})); - - using InverseType = cutlass::half_t; - using CollectiveInverse = ku::CollectiveInverse; - - using ElementAccumulatorSK = float; - using TileShapeSK = decltype(make_shape(HeadSizeV, BlkSeqKV, HeadSizeQK)); - using CollectiveMmaSK = typename cutlass::gemm::collective::CollectiveBuilder< // basically the same as O1 - cutlass::arch::Sm90, - cutlass::arch::OpClassTensorOp, - Element, - RefLayoutKV, - Alignment, - Element, - LayoutK, - Alignment, - ElementAccumulatorSK, - TileShapeSK, - ClusterShape, - DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; - - using ElementAccumulatorNewV = float; - using TileShapeNewV = decltype(make_shape(HeadSizeV, BlkSeqKV, BlkSeqKV)); - using RefLayoutSK = decltype(make_layout(select<0, 2>(TileShapeNewV{}), LayoutRight{})); // (dv, Blk) - using DesiredLayoutKK = decltype(make_layout(select<1, 2>(TileShapeNewV{}), LayoutRight{})); // - using CollectiveMmaNewV = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm90, - cutlass::arch::OpClassTensorOp, - Element, - RefLayoutSK, - Alignment, - Element, - DesiredLayoutKK, - Alignment, - ElementAccumulatorKV, - TileShapeNewV, - ClusterShape, - DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; - - // FIXME: K@K^t are not exactly the same as Q@K^t, but similar enough (what does this mean??) - using TiledMmaKK = typename CollectiveMmaQK::TiledMma; // T = inv(I + strict_lower_triangular(K@K^t)) - using TiledMmaSK = decltype(convert_to_gmma_rs(typename CollectiveMmaSK::TiledMma{})); // ?? = -S@K^t + V^t - using TiledMmaNewV = decltype(convert_to_gmma_rs(typename CollectiveMmaNewV::TiledMma{})); // NewV = ??@T^t - - static_assert(size(TiledMmaKK{}) == NumAuxMmaThreads); - - using GmemStrideBeta = Stride; - using GmemLayoutBeta = Layout, GmemStrideBeta>; // (seq, head) - - // only store the last row in Alpha - using SmemLayoutAlphaLast = decltype(make_layout(make_shape(HeadSize, Int{}))); - using SmemLayoutBeta = decltype(make_layout(make_shape(BlkSeqQ, Int{}))); - - using MainloopQPipeline = cutlass::PipelineTmaAsync; - using MainloopKPipeline = cutlass::PipelineTmaAsync; - using MainloopVPipeline = cutlass::PipelineTmaAsync; - using MainloopAlphaPipeline = std::conditional_t, Unused>; - using MainloopOPipeline = typename CollectiveStoreO::Pipeline; - - using MainloopQKPipeline = cutlass::PipelineAsync; - using MainloopKKPipeline = cutlass::PipelineAsync; - - using MainloopAlphaLastPipeline = - std::conditional_t, Unused>; - - using MainloopBetaPipeline = std::conditional_t, Unused>; - - using QPipelineState = typename cutlass::PipelineState; - using KPipelineState = typename cutlass::PipelineState; - using VPipelineState = typename cutlass::PipelineState; - using OPipelineState = typename CollectiveStoreO::PipelineState; - - using QKPipelineState = cutlass::PipelineState; - using KKPipelineState = cutlass::PipelineState; - - using AlphaLastPipelineState = - std::conditional_t, Unused>; - - using AlphaPipelineState = - std::conditional_t, Unused>; - using BetaPipelineState = - std::conditional_t, Unused>; - - using BetaProcessor = Unused; - - static constexpr int LoadQBytes = size(QKSmemLayoutQ{}(_, _, _0{})) * sizeof(Element); - static constexpr int LoadKBytes = size(KVSmemLayoutK{}(_, _, _0{})) * sizeof(Element); - static constexpr int LoadVBytes = size(KVSmemLayoutV{}(_, _, _0{})) * sizeof(Element); - static constexpr int LoadAlphaBytes = size(QKQSmemLayoutAlpha{}(_, _, _0{})) * sizeof(ElementAlpha); - static constexpr int StoreOBytes = CollectiveStoreO::TmaTransactionBytes; - - using SharedStorageO = typename CollectiveStoreO::SharedStorage; - - struct SharedStorage { - alignas( - alignment_for_swizzle(QKSmemLayoutQ{})) cute::array_aligned> smem_q; - alignas( - alignment_for_swizzle(KVSmemLayoutK{})) cute::array_aligned> smem_k; - alignas( - alignment_for_swizzle(KVSmemLayoutV{})) cute::array_aligned> smem_v; - alignas(alignment_for_swizzle( - QKQSmemLayoutAlpha{})) cute::array_aligned> smem_alpha; - alignas( - alignment_for_swizzle(SmemLayoutQK{})) cute::array_aligned> smem_qk; - alignas(alignment_for_swizzle( - SmemLayoutKK{})) cute::array_aligned> smem_kk; - // smemq_k_scaled for exp(alpha) * Q and exp(alpha) * K in QS and KS, computed in Math WG2/3 - alignas(alignment_for_swizzle( - QKScaledSmemLayoutQ{})) cute::array_aligned> smem_q_k_scaled; - - SharedStorageO smem_o; - - cute::array_aligned> smem_beta; - // store last row in Alpha separately, used for S'=K^T NewV's epilogue and S+=decay(S') (one fused epilogue) - cute::array_aligned> smem_alpha_last; - }; - - using TMA_Q = typename CollectiveMmaQK::Params::TMA_A; - using TMA_K = typename CollectiveMmaKV_G2S::Params::TMA_B; - using TMA_V = typename CollectiveMmaKV_G2S::Params::TMA_A; - using TMA_O = typename CollectiveStoreO::Params::TMA_O; - - using LoadQ = CollectiveLoadTma; - using LoadK = CollectiveLoadTma; - using LoadV = CollectiveLoadTma; - using LoadAlpha = - CollectiveLoadTma; - - using LoadBeta = CollectiveLoadVector< - LoadKindVector::kBeta, - MainloopBetaPipeline, - ElementBetaGmem, - GmemLayoutBeta, - ElementBeta, - SmemLayoutBeta, - BetaProcessor>; - - struct Arguments { // clang-format off - Element const* ptr_Q; LayoutQ dQ; - Element const* ptr_K; LayoutK dK; - Element const* ptr_V; LayoutV dV; - Element* ptr_O; LayoutO dO; - float const* ptr_Alpha; LayoutAlpha dAlpha; - float* ptr_output_state; // layout fixed (kdim, vdim, num_v_heads, num_seqs):LayoutLeft{} - float const* ptr_input_state; - float scale; - ElementBetaGmem const* beta_ptr; GmemStrideBeta beta_stride; - }; // clang-format on - - struct Params { - TMA_Q tma_load_q; - TMA_K tma_load_k; - TMA_V tma_load_v; - TMA_Alpha tma_load_alpha; - TMA_O tma_store_o; - void* tensormaps; - float scale; - - float* ptr_output_state; - float const* ptr_input_state; - - ElementBetaGmem const* beta_ptr; - GmemLayoutBeta beta_layout; - }; - - template - static bool - can_implement(ProblemShape const& problem_size, Arguments const& args) { - return true && (problem_size.head_size <= get<2>(TileShape{})) && ((problem_size.head_size % Alignment) == 0); - } - - template - static Params - to_underlying_arguments(ProblemShape const& problem_size, Arguments const& args, void* workspace) { - int64_t s = problem_size.total_seqlen; - int64_t t = problem_size.total_seqlen; - int32_t d = problem_size.head_size; - - // GVA: Q/K are sized by num_qk_heads; V, alpha (gate), O, beta and the recurrent state - // are sized by num_v_heads. - auto params_qk = CollectiveMmaQK::to_underlying_arguments( - make_shape(s, t, d, problem_size.num_qk_heads), - typename CollectiveMmaQK::Arguments{ - args.ptr_Q, args.dQ, args.ptr_K, args.dK, // never used, dummy - }, - /*workspace=*/nullptr); - - auto params_kv_k = CollectiveMmaKV_G2S::to_underlying_arguments( - make_shape(d, d, s, problem_size.num_qk_heads), - typename CollectiveMmaKV_G2S::Arguments{ - args.ptr_V, - select<1, 0, 2>(args.dV), // not used - args.ptr_K, - select<1, 0, 2>(args.dK), // used as G2S for K - }, - /*workspace=*/nullptr); - - auto alpha_shape = make_shape(s, d, problem_size.num_v_heads); - auto alpha_stride = make_stride( - get<0>(args.dAlpha), // seqlen stride - get<1>(args.dAlpha), // head_dim stride - get<2>(args.dAlpha) // head stride - ); - Tensor mAlpha = make_tensor(make_gmem_ptr(args.ptr_Alpha), make_layout(alpha_shape, alpha_stride)); - TMA_Alpha tma_load_alpha = make_tma_copy( - GmemTiledCopyAlpha{}, - mAlpha, - take<0, 2>(SmemLayoutAlpha_SD{}), - select<1, 2>(TileShapeQK{}), - size<0>(ClusterShape{})); - - auto params_kv_v = CollectiveMmaKV_G2S::to_underlying_arguments( - make_shape(d, d, s, problem_size.num_v_heads), - typename CollectiveMmaKV_G2S::Arguments{ - args.ptr_V, - select<1, 0, 2>(args.dV), // used as G2S for V - args.ptr_K, - select<1, 0, 2>(args.dK), // not used - }, - /*workspace=*/nullptr); - - auto params_o = CollectiveStoreO::to_underlying_arguments( - make_shape(d, s, d, problem_size.num_v_heads), // in O1 - // make_shape(d, s, s, problem_size.num_v_heads), // in O2 - typename CollectiveStoreO::Arguments{args.ptr_O, select<1, 0, 2>(args.dO), workspace}, - workspace); - - return Params{ - .tma_load_q = params_qk.tma_load_a, - .tma_load_k = params_kv_k.tma_load_b, - .tma_load_v = params_kv_v.tma_load_a, - .tma_load_alpha = tma_load_alpha, - .tma_store_o = params_o.tma_store_o, - .tensormaps = params_o.tensormaps, - .scale = args.scale, - - .ptr_output_state = args.ptr_output_state, - .ptr_input_state = args.ptr_input_state, - - // TODO: refactor all name to varname_vartype - .beta_ptr = args.beta_ptr, - .beta_layout = make_layout(make_shape(s, problem_size.num_v_heads), args.beta_stride), - }; - } - - static size_t - get_workspace_size(Arguments const& args, int sm_count) { - return CollectiveStoreO::get_workspace_size(sm_count); - } - - template - static cutlass::Status - initialize_workspace( - ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { - return CollectiveStoreO::initialize_workspace(problem_shape, workspace, stream); - } - - CUTE_DEVICE static void - prefetch_tma_descriptors(Params const& params) { - cute::prefetch_tma_descriptor(params.tma_load_q.get_tma_descriptor()); - cute::prefetch_tma_descriptor(params.tma_load_k.get_tma_descriptor()); - cute::prefetch_tma_descriptor(params.tma_load_v.get_tma_descriptor()); - cute::prefetch_tma_descriptor(params.tma_load_alpha.get_tma_descriptor()); - cute::prefetch_tma_descriptor(params.tma_store_o.get_tma_descriptor()); - } - - template - CUTE_DEVICE void - load_qkv( - Params const& params, - ProblemShape const& problem_size, - LoadTileShape const& load_tile_shape, - WorkDesc const& work_desc, - MainloopQPipeline& q_pipeline, - QPipelineState& q_smem_pipe_write, - MainloopKPipeline& k_pipeline, - KPipelineState& k_smem_pipe_write, - MainloopVPipeline& v_pipeline, - VPipelineState& v_smem_pipe_write, - MainloopAlphaPipeline& alpha_pipeline, - AlphaPipelineState& alpha_smem_pipe_write, - SharedStorage& storage) { - int32_t num_blocks = ceil_div(work_desc.seq_len, get<0>(TileShape{})); - uint32_t lane_predicate = cute::elect_one_sync(); - - auto q_collective_load = LoadQ(params.tma_load_q, q_pipeline, storage.smem_q); - auto k_collective_load = LoadK(params.tma_load_k, k_pipeline, storage.smem_k); - auto v_collective_load = LoadV(params.tma_load_v, v_pipeline, storage.smem_v); - auto alpha_collective_load = LoadAlpha{params.tma_load_alpha, alpha_pipeline, storage.smem_alpha}; - - auto q_src_dst = q_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); - auto k_src_dst = k_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); - auto v_src_dst = v_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); - auto alpha_src_dst = alpha_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); - - CUTE_NO_UNROLL - for (int blk = 0; blk < num_blocks; ++blk) { - alpha_collective_load.step(alpha_src_dst, blk, alpha_smem_pipe_write, lane_predicate); - q_collective_load.step(q_src_dst, blk, q_smem_pipe_write, lane_predicate); - k_collective_load.step(k_src_dst, blk, k_smem_pipe_write, lane_predicate); - v_collective_load.step(v_src_dst, blk, v_smem_pipe_write, lane_predicate); - } - } - - template - CUTE_DEVICE void - load_beta( - Params const& params, - ProblemShape const& problem_size, - TileShape const& tile_shape, - WorkDesc const& work_desc, - MainloopBetaPipeline& pipeline, - BetaPipelineState& smem_pipe_write, - SharedStorage& storage) { - int32_t num_blocks = ceil_div(work_desc.seq_len, get<0>(TileShape{})); - - // fuse post inverse diag(beta) into diagonal of IKK - // auto collective_load = LoadBeta{params.beta_ptr, params.beta_layout, /*oob_value=*/1.0f, pipeline, - // storage.smem_beta}; - auto collective_load = LoadBeta{ - params.beta_ptr, params.beta_layout, /*oob_value=*/ElementBetaGmem(0), pipeline, storage.smem_beta}; - auto src_dst = collective_load.partition_SD(problem_size, tile_shape, work_desc); - - CUTE_NO_UNROLL - for (int blk = 0; blk < num_blocks - 1; ++blk) { - collective_load.step(src_dst, blk, smem_pipe_write, num_blocks); - } - collective_load.step(src_dst, num_blocks - 1, smem_pipe_write, num_blocks); - } - - template - CUTE_DEVICE void - extract_alpha_last( - Params const& params, - ProblemShape const& problem_size, - TileShape const& tile_shape, - WorkDesc const& work_desc, - MainloopAlphaPipeline& alpha_pipeline, - AlphaPipelineState& alpha_smem_pipe_read, - MainloopAlphaLastPipeline& alpha_last_pipeline, - AlphaLastPipelineState& alpha_last_smem_pipe_write, - SharedStorage& storage) { - int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarp; - - Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); - Tensor sAlast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); - - auto extract_loop_body = [&](int blk, auto is_final_block_) INLINE_LAMBDA { - constexpr bool is_final_block = decltype(is_final_block_)::value; - - int B = is_final_block ? valid_seq_len(work_desc, blk) : BlkSeqKV; - - auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); - Tensor sAlast_out = sAlast(_, alpha_last_smem_pipe_write.index()); - - alpha_pipeline.consumer_wait(alpha_smem_pipe_read); - alpha_last_pipeline.producer_acquire(alpha_last_smem_pipe_write); - - // each thread copy 4 elements, total 128 elements with one warp - CUTE_UNROLL - for (int t = thread_idx; t < HeadSize; t += 32) { - sAlast_out(t) = sAqkq_curr(B - 1, t); - } - - cutlass::arch::fence_view_async_shared(); - alpha_last_pipeline.producer_commit(alpha_last_smem_pipe_write); - ++alpha_last_smem_pipe_write; - alpha_pipeline.consumer_release(alpha_smem_pipe_read); - ++alpha_smem_pipe_read; - }; - - int32_t num_blocks = ceil_div(work_desc.seq_len, get<0>(TileShape{})); - CUTE_NO_UNROLL - for (int blk = 0; blk < num_blocks - 1; ++blk) { - extract_loop_body(blk, /*is_final_block_=*/cute::false_type{}); - } - extract_loop_body(num_blocks - 1, /*is_final_block_=*/cute::true_type{}); - } - - template - CUTE_DEVICE void - store( - TMA_O const& tma_store, - void* tensormaps, - ProblemSize const& problem_size, - StoreTileShape const& store_tile_shape, - WorkDesc const& work_desc, - MainloopOPipeline& pipeline, - PipelineState& smem_pipe_read, - SharedStorageO& storage) { - int32_t num_blocks = ceil_div(work_desc.seq_len, get<0>(TileShape{})); - uint32_t lane_predicate = cute::elect_one_sync(); - - auto collective_store = CollectiveStoreO{tma_store, pipeline, storage, tensormaps}; - auto src_dst = collective_store.partition_SD(problem_size, store_tile_shape, work_desc); - - CUTE_NO_UNROLL - for (int blk = 0; blk < num_blocks; ++blk) { - DPRINTF0_W( - "O collective_store.step smem_pipe_read:%d -> blk_idx:%d, num_blocks:%d\n", - smem_pipe_read.index(), - blk, - num_blocks); - collective_store.step(problem_size, work_desc, src_dst, smem_pipe_read, blk, num_blocks, lane_predicate); - } - } - - template - CUTE_DEVICE void - compute( - Params const& params, - ProblemShape const& problem_size, - WorkDesc const& work_desc, - MainloopQPipeline& q_pipeline, - QPipelineState& q_smem_pipe_read, - MainloopKPipeline& k_pipeline, - KPipelineState& k_smem_pipe_read, - MainloopVPipeline& v_pipeline, - VPipelineState& v_smem_pipe_read, - MainloopOPipeline& o_pipeline, - OPipelineState& o_smem_pipe_write, - MainloopQKPipeline& qk_pipeline, - QKPipelineState& qk_smem_pipe_read, - MainloopKKPipeline& kk_pipeline, - KKPipelineState& kk_smem_pipe_read, - MainloopAlphaPipeline& alpha_pipeline, - AlphaPipelineState& alpha_smem_pipe_read, - MainloopBetaPipeline& beta_pipeline, - BetaPipelineState& beta_smem_pipe_read, - MainloopAlphaLastPipeline& alpha_last_pipeline, - AlphaLastPipelineState& alpha_last_smem_pipe_read, - OrderedMathBarriers& math_barriers, - SharedStorage& storage) { - // MAKE NVCC HAPPY! - constexpr auto zero = Element{}; - - int32_t num_blocks = ceil_div(work_desc.seq_len, get<0>(TileShape{})); - DPRINTF0_WG("num_blocks: %d\n", num_blocks); - - int thread_idx = int(threadIdx.x) - NumLoadThreads; - int warpgroup_idx = thread_idx / cutlass::NumThreadsPerWarpGroup; - int thread_idx_in_wg = thread_idx % cutlass::NumThreadsPerWarpGroup; - - float scale = params.scale; - - Tensor Beta = make_tensor(make_smem_ptr(storage.smem_beta.data()), SmemLayoutBeta{}); - Tensor AlphaLast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); - - Tensor sQqk = make_tensor(make_smem_ptr(storage.smem_q.data()), QKSmemLayoutQ{}); - Tensor sKqk = make_tensor(make_smem_ptr(storage.smem_k.data()), QKSmemLayoutK{}); - Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); - Tensor sVkv = make_tensor(make_smem_ptr(storage.smem_v.data()), KVSmemLayoutV{}); - Tensor sQK = make_tensor(make_smem_ptr(storage.smem_qk.data()), SmemLayoutQK{}); - Tensor sO = make_tensor(make_smem_ptr(storage.smem_o.data()), SmemLayoutO{}); - - static_assert(sizeof(InverseType) == sizeof(Element)); - Tensor sKK_inv = make_tensor(make_smem_ptr(storage.smem_kk.data()), SmemLayoutKK{}); - Tensor sKK_opd = make_tensor(make_smem_ptr(reinterpret_cast(storage.smem_kk.data())), SmemLayoutKK{}); - - Tensor sQ_K_scaled = make_tensor(make_smem_ptr(storage.smem_q_k_scaled.data()), QKScaledSmemLayoutQ{}); - Tensor sQ_K_scaled_Kt = make_tensor(make_smem_ptr(storage.smem_q_k_scaled.data()), QKScaledSmemLayoutKt{}); - - /////////////////////////////////////////////////////////////////////////// - // Q@S, K@S, Q/K prologue - // each WG process 32 at a time, reduce peak register usage - // each WG process half head dim (64) at all - auto qk_tiled_mma_rs_quar = TiledMmaQK_RS_Quar{}; - auto qk_thr_mma_rs_quar = qk_tiled_mma_rs_quar.get_thread_slice(thread_idx_in_wg); - constexpr auto tiler_alpha = Shape<_64, Shape<_32, _1>>{}; - constexpr auto tiler_qk = Shape<_64, Shape<_32, _1>>{}; - constexpr auto tiler_alpha_last = Shape<_32>{}; - // used for Alpha S2R (float) - using CopyAlphaAtom = Copy_Atom, ElementAlpha>; - // used for Q/K S2R and R2S (fp16/bf16) - using CopyOpS2R = SM75_U32x4_LDSM_N; - using CopyOpR2S = SM90_U32x4_STSM_N; - auto tiled_load_qk_quar = make_tiled_copy_A(Copy_Atom{}, qk_thr_mma_rs_quar); - auto thr_load_qk_quar = tiled_load_qk_quar.get_thread_slice(thread_idx_in_wg); - auto tiled_store_qk_quar = make_tiled_copy_A(Copy_Atom{}, qk_thr_mma_rs_quar); - auto thr_store_qk_quar = tiled_store_qk_quar.get_thread_slice(thread_idx_in_wg); - - auto cMq_quar = make_identity_tensor(select<0, 2>(TileShapeQK_Quar{})); // (QTok, HeadDim / 2) - auto tQcMq_quar = qk_thr_mma_rs_quar.partition_A(cMq_quar); // (idx) -> (tok_q, head_dim / 2) - - /////////////////////////////////////////////////////////////////////////// - // K@K (basically I + strict_lower_triangular(K K^T) - auto kk_tiled_mma = TiledMmaKK{}; - auto kk_thr_mma = kk_tiled_mma.get_thread_slice(thread_idx_in_wg); - Tensor tKKsK = kk_thr_mma.partition_B(sKqk); - Tensor tKKrA = kk_thr_mma.make_fragment_A(tKKsK); - auto cMqk = make_identity_tensor(select<0, 1>(TileShapeQK{})); // (QTok, KTok) - auto const& cMkk = cMqk; - auto tKKcMkk = kk_thr_mma.partition_C(cMkk); - - // S@K (-S K^T + V^T) - auto sk_tiled_mma = TiledMmaSK{}; - auto sk_thr_mma = sk_tiled_mma.get_thread_slice(thread_idx); - - // tSKrV adds to tSKrSK (acc) - using SK_V_S2R = Copy_Atom; - auto tSKrV_tiled_copy = make_tiled_copy_C(SK_V_S2R{}, sk_tiled_mma); - auto tSKrV_thr_copy = tSKrV_tiled_copy.get_thread_slice(thread_idx); - - Tensor tSKsK = sk_thr_mma.partition_B(sQ_K_scaled); - Tensor tSKrK = sk_thr_mma.make_fragment_B(tSKsK); - - /////////////////////////////////////////////////////////////////////////// - // NewV = (S@K result) @ T^t - auto newv_tiled_mma = TiledMmaNewV{}; - auto newv_thr_mma = newv_tiled_mma.get_thread_slice(thread_idx); - - Tensor tNewVsB = newv_thr_mma.partition_B(sKK_opd); - Tensor tNewVrB = newv_thr_mma.make_fragment_B(tNewVsB); - - /////////////////////////////////////////////////////////////////////////// - // K@V - auto kv_tiled_mma = TiledMmaKV{}; // (V, Blk_k) @ (Blk_k, K) = (V, K) - auto kv_thr_mma = kv_tiled_mma.get_thread_slice(thread_idx); - - Tensor tKVrKV = partition_fragment_C(kv_thr_mma, select<0, 1>(TileShapeKV{})); - - // Tensor tKVrV = kv_thr_mma.partition_fragment_A(sVkv(_, _, _0{})); // mma src - // Tensor tKVrV_cv = tKVrV_thr_copy.retile_D(tKVrV); // copy view dst - // Tensor tKVsV = tKVrV_thr_copy.partition_S(sVkv); // copy view src - - auto const cV = make_identity_tensor(Shape, Int>{}); - Tensor tKVcV = kv_thr_mma.partition_A(cV); - auto const cS = make_identity_tensor(Shape, Int>{}); - Tensor tKVcS = kv_thr_mma.partition_C(cS); - - /////////////////////////////////////////////////////////////////////////// - // Q@K@V - auto o1_tiled_mma = TiledMmaO1{}; - auto o1_thr_mma = o1_tiled_mma.get_thread_slice(thread_idx); - auto o2_tiled_mma = TiledMmaO2{}; - auto o2_thr_mma = o2_tiled_mma.get_thread_slice(thread_idx); - - // A1 for Q@(KV) - // Tensor tOrKV = make_acc_into_op(tKVrKV, typename TiledMmaO1::LayoutA_TV{}); - // B1 for Q@(KV) - Tensor tOsQ = o1_thr_mma.partition_B(sQ_K_scaled); - Tensor tOrQ = o1_thr_mma.make_fragment_B(tOsQ); - - // A2 for QK@V - // Tensor tOsV = o2_thr_mma.partition_A(sVkv); - // Tensor tOrV = o2_thr_mma.make_fragment_A(tOsV); - // B2 for QK@V - Tensor tOsQK = o2_thr_mma.partition_B(sQK); - Tensor tOrQK = o2_thr_mma.make_fragment_B(tOsQK); - - using O_R2S = typename CollectiveStoreO::CopyAtomR2S; - auto tiled_copy_o = make_tiled_copy_C(O_R2S{}, o1_tiled_mma); - auto thr_copy_o = tiled_copy_o.get_thread_slice(thread_idx); - auto tOsO = thr_copy_o.partition_D(sO); - - auto const cO = make_identity_tensor(Shape, Int>{}); - Tensor tOcO = o1_thr_mma.partition_C(cO); - - auto const seq_idx = work_desc.seq_idx; - auto const q_head_idx = work_desc.q_head_idx(); - auto const k_head_idx = work_desc.k_head_idx(); - auto const v_head_idx = work_desc.v_head_idx(); - - auto sk_load_v = [&](int pipe_idx) INLINE_LAMBDA { - Tensor tSKrV = make_fragment_like(partition_fragment_C(sk_thr_mma, sVkv(_, _, _0{}))); // mma acc - Tensor tSKrV_cv = tSKrV_thr_copy.retile_D(tSKrV); // copy view dst - Tensor tSKsV = tSKrV_thr_copy.partition_S(sVkv); // copy view src - copy(tSKrV_tiled_copy, tSKsV(_, _, _, pipe_idx), tSKrV_cv); - return tSKrV; - }; - - auto kv_load = [&](auto& tKVrKV) INLINE_LAMBDA { - DPRINTF0_WG("[%d,%d,%d,%d]>> load tKVgKV -> tKVrKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); - // GVA: state is stored per V/O head. - int num_state_heads = problem_size.num_v_heads; - int state_head_idx = work_desc.o_head_idx(); - auto gKV = make_tensor( - make_gmem_ptr(params.ptr_input_state), - make_layout(make_shape(Int{}, Int{}, num_state_heads, problem_size.num_seqs)))( - _, _, state_head_idx, seq_idx); // (KDim, VDim), K-contiguous - - auto tiled_copy_kv = make_tiled_copy_C(Copy_Atom{}, kv_tiled_mma); - auto thr_copy_kv = tiled_copy_kv.get_thread_slice(thread_idx); - // transposed load state - auto tKVgKV = thr_copy_kv.partition_S(select_tensor<1, 0>(gKV)); - copy(tiled_copy_kv, tKVgKV, tKVrKV); - }; - - auto kv_store = [&]() INLINE_LAMBDA { // tKVrKV is carried over whole mainloop - // Skip the final-state write-back entirely when the caller does not need it - // (i.e. output_final_state=False on the public API). The check is uniform - // across the launch, so the branch cost is negligible while it saves a full - // [N, num_v_heads, D, D] float32 store to GMEM. - if (params.ptr_output_state == nullptr) { - DPRINTF0_WG( - "[%d,%d,%d,%d]>> skip tKVrKV -> tKVgKV (output_final_state=false)\n", - seq_idx, - q_head_idx, - k_head_idx, - v_head_idx); - return; - } - // Intra-card CP - int32_t out_seq_idx = seq_idx; - int32_t out_num_seqs = problem_size.num_seqs; - if (problem_size.cp_seq_map != nullptr) { - out_seq_idx = problem_size.cp_seq_map[seq_idx]; - out_num_seqs = problem_size.raw_num_seqs; - int32_t this_end = problem_size.cu_seqlens[seq_idx + 1]; - int32_t raw_end = problem_size.raw_cu_seqlens[out_seq_idx + 1]; - if (this_end != raw_end) { - return; - } - } - DPRINTF0_WG("[%d,%d,%d,%d]>> save tKVrKV -> tKVgKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); - // GVA: state is stored per V/O head. - int num_state_heads = problem_size.num_v_heads; - int state_head_idx = work_desc.o_head_idx(); - auto gKV = make_tensor( - make_gmem_ptr(params.ptr_output_state), - make_layout(make_shape(Int{}, Int{}, num_state_heads, out_num_seqs)))( - _, _, state_head_idx, out_seq_idx); // (KDim, VDim), K-contiguous - - auto tiled_copy_kv = make_tiled_copy_C(Copy_Atom{}, kv_tiled_mma); - auto thr_copy_kv = tiled_copy_kv.get_thread_slice(thread_idx); - // transposed store state - auto tKVgKV = thr_copy_kv.partition_D(select_tensor<1, 0>(gKV)); - copy(tiled_copy_kv, tKVrKV, tKVgKV); - }; - - auto s_decay = [&](auto& tKVrKV, auto const& alpha_last_smem_pipe_read) INLINE_LAMBDA { - Tensor alpha_last_curr = AlphaLast(_, alpha_last_smem_pipe_read.index()); - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tKVcS(i); - auto [s, t] = coord; // (head_size_v, head_size_k) - tKVrKV(i) *= exp2f(alpha_last_curr(t)); - }); - }; - - auto o1_epi = [&](auto& tOrO1) INLINE_LAMBDA { - CUTE_UNROLL - for (int i = 0; i < size(tOrO1); ++i) { - tOrO1(i) = scale * tOrO1(i); - } - }; - - auto o_store = [&](auto tOrO) INLINE_LAMBDA { - auto tOrO_cvt = make_fragment_like(tOrO); - copy(tOrO, tOrO_cvt); - - DPRINTF0_WG("compute: o_pipeline.producer_wait: smem_pipe_write:%d\n", o_smem_pipe_write.index()); - o_pipeline.producer_acquire(o_smem_pipe_write); - Tensor tOrO_cvt_cv = thr_copy_o.retile_S(tOrO_cvt); - cutlass::arch::fence_view_async_shared(); - copy(tiled_copy_o, tOrO_cvt_cv, tOsO(_, _, _, o_smem_pipe_write.index())); - cutlass::arch::fence_view_async_shared(); - o_pipeline.producer_commit(o_smem_pipe_write); - ++o_smem_pipe_write; - }; - - auto kk_inv = [&](auto const& kk_smem_pipe_read) INLINE_LAMBDA { - auto sKK_inv_pipe_slice = sKK_inv(_, _, kk_smem_pipe_read.index()); - static_assert(sizeof(Element) == 2); - using CopyOpR2S = SM90_U32x4_STSM_N; - auto tiled_store_kk = make_tiled_copy_C(Copy_Atom{}, kk_tiled_mma); - auto thr_store_kk = tiled_store_kk.get_thread_slice(thread_idx); - auto tKKsKK = thr_store_kk.partition_D(sKK_inv_pipe_slice); - // TODO: use tKKcMkk? no more allocating fragments - auto tKKrKK = kk_thr_mma.partition_fragment_C(sKK_inv_pipe_slice); - auto tKKrKK_cv = thr_store_kk.retile_S(tKKrKK); - auto collective_inverse = CollectiveInverse(KdaNamedBarriers::StateMathWG0); - collective_inverse.compute(sKK_inv_pipe_slice); - // FIXME: we can ignore core matrices above diagonal - if constexpr (NeedsBeta || !std::is_same_v) { - cutlass::arch::NamedBarrier::arrive_and_wait( - cutlass::NumThreadsPerWarpGroup, KdaNamedBarriers::StateMathWG0); - using CopyOpS2R = SM75_U32x4_LDSM_N; - auto tiled_load_kk = make_tiled_copy_C(Copy_Atom{}, kk_tiled_mma); - auto thr_load_kk = tiled_load_kk.get_thread_slice(thread_idx); - auto tKKrKK_cpy = make_fragment_like(tKKrKK_cv); - auto tKKrKK_cvt = make_fragment_like(tKKrKK_cv); - auto tKKcMkk_cv = thr_load_kk.retile_D(tKKcMkk); - copy(tiled_load_kk, thr_load_kk.partition_S(sKK_inv_pipe_slice), tKKrKK_cpy); - cute::transform(tKKrKK_cpy, tKKcMkk_cv, tKKrKK_cvt, [&](auto val, auto coord) { - auto [_, t] = coord; - if constexpr (NeedsBeta) { - return Element(float(val) * Beta(t, beta_smem_pipe_read.index())); - } else { - return Element(val); - } - }); - copy(tiled_store_kk, tKKrKK_cvt, recast(tKKsKK)); - } - }; - - auto compute_loop_body = [&](int blk, auto is_first_block_, auto is_final_block_) INLINE_LAMBDA { - constexpr bool is_first_block = decltype(is_first_block_)::value; - constexpr bool is_final_block = decltype(is_final_block_)::value; - int B = is_final_block ? valid_seq_len(work_desc, blk) : BlkSeqKV; - - auto sQqk_curr = sQqk(_, _, q_smem_pipe_read.index()); - auto sKqk_curr = sKqk(_, _, k_smem_pipe_read.index()); - auto sQ_scaled_curr = sQ_K_scaled(_, _, _0{}); - auto sK_scaled_curr = sQ_K_scaled(_, _, _1{}); - auto sAlast_curr = AlphaLast(_, alpha_last_smem_pipe_read.index()); - auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); - auto sQqk_slice = flat_divide(sQqk_curr, tiler_qk); - auto sKqk_slice = flat_divide(sKqk_curr, tiler_qk); - auto sQ_scaled_slice = flat_divide(sQ_scaled_curr, tiler_qk); - auto sK_scaled_slice = flat_divide(sK_scaled_curr, tiler_qk); - auto sAqkq_slice = flat_divide(sAqkq_curr, tiler_alpha); - auto sAlast_slice = flat_divide(sAlast_curr, tiler_alpha_last); - - if constexpr (NeedsAlpha) { - alpha_pipeline.consumer_wait(alpha_smem_pipe_read); - } - DPRINTF0_WG("compute: q_pipeline.consumer_wait: smem_pipe_read:%d\n", q_smem_pipe_read.index()); - q_pipeline.consumer_wait(q_smem_pipe_read); - DPRINTF0_WG("compute: k_pipeline.consumer_wait: smem_pipe_read:%d\n", k_smem_pipe_read.index()); - k_pipeline.consumer_wait(k_smem_pipe_read); - - // load alpha and exp2(alpha) only once - // and reuse these registers in exp(alpha) * Q/K prologue - if constexpr (!is_first_block) { - // make sure sQ_K_scaled is already consumed for previous K^@V - cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - // Each WG iterates over 2 slices of 32 elements each. - // WG0 (thread_idx < 128): wg_idx=0, processes alpha indices {0,1}, Q/K dim1=0 - // WG1 (thread_idx >= 128): wg_idx=1, processes alpha indices {2,3}, Q/K dim1=1 - { - int wg_idx = thread_idx / 128; // 0 or 1 - int alpha_base = wg_idx * 2; // 0 or 2 - - // Allocate Q/K register fragments once (reused across slices) - // Only shape/layout matters for partition_fragment_A, use compile-time indices - auto tQKrQ_wg = - qk_thr_mma_rs_quar.partition_fragment_A(sQqk_slice(_, _, _0{}, make_coord(_0{}, _0{}))); - auto tQKrK_wg = - qk_thr_mma_rs_quar.partition_fragment_A(sKqk_slice(_, _, _0{}, make_coord(_0{}, _0{}))); - auto tArA = make_fragment_like(tQKrQ_wg); - - for (int s = 0; s < 2; ++s) { - // S2R Alpha: alpha_col = wg_idx * 2 + s - int alpha_col = alpha_base + s; - auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); - auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); - copy(CopyAlphaAtom{}, tAsA_cur, tArA); - - cute::transform(tArA, [](auto g) { return exp2f(g); }); - - // S2R Q - auto sQqk_cur = sQqk_slice(_, _, _0{}, make_coord(s, wg_idx)); - auto tQKsQ_cur = thr_load_qk_quar.partition_S(sQqk_cur); - auto tQKrQ_cv = thr_load_qk_quar.retile_D(tQKrQ_wg); - copy(tiled_load_qk_quar, tQKsQ_cur, tQKrQ_cv); - - // element-wise exp(alpha) * Q - cute::transform(tQKrQ_wg, tArA, tQKrQ_wg, [&](auto q, auto alpha) { - Element dst = Element(alpha * float(q)); - return dst; - }); - - // R2S Q -> stage 0 - auto sQ_scaled_cur = sQ_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); - auto tQKsQ_out = thr_store_qk_quar.partition_D(sQ_scaled_cur); - auto tQKrQ_out_cv = thr_store_qk_quar.retile_S(tQKrQ_wg); - copy(tiled_store_qk_quar, tQKrQ_out_cv, tQKsQ_out); - - // S2R K - auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(s, wg_idx)); - auto tQKsK_cur = thr_load_qk_quar.partition_S(sKqk_cur); - auto tQKrK_cv = thr_load_qk_quar.retile_D(tQKrK_wg); - copy(tiled_load_qk_quar, tQKsK_cur, tQKrK_cv); - - // element-wise exp(alpha) * K - cute::transform(tQKrK_wg, tArA, tQKrK_wg, [&](auto k, auto alpha) { - Element dst = Element(alpha * float(k)); - return dst; - }); - - // R2S K -> stage 1 - auto sK_scaled_cur = sK_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); - auto tQKsK_out = thr_store_qk_quar.partition_D(sK_scaled_cur); - auto tQKrK_out_cv = thr_store_qk_quar.retile_S(tQKrK_wg); - copy(tiled_store_qk_quar, tQKrK_out_cv, tQKsK_out); - } - } - cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - // fence to produce data for WGMMA async proxy - cutlass::arch::fence_view_async_shared(); - // if (blk <= 1 && thread_idx == 0) { - // printf("After Q/K prologue: exp(alpha) * Q at stage 0, exp(alpha) * K at stage 1\n"); - // cute::print_tensor(sQ_K_scaled_curr); - // } - } - - // 2.1 Q @ KV, NOTE: use old KV here - - DPRINTF0_WG("[%d,%d,%d,%d]** dispatch O WGMMA\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); - auto tOrO = partition_fragment_C(o1_thr_mma, select<0, 1>(TileShapeO1{})); - if constexpr (is_first_block) { - DPRINTF0_WG("compute: q_pipeline.consumer_release: smem_pipe_read:%d\n", q_smem_pipe_read.index()); - q_pipeline.consumer_release(q_smem_pipe_read); - ++q_smem_pipe_read; - } else { - Tensor tOrKV = make_acc_into_op(tKVrKV, typename TiledMmaO1::LayoutA_TV{}); - warpgroup_fence_operand(tOrKV); - warpgroup_fence_operand(tOrO); - // ======DEBUG======= - // if (blk <= 6 && thread_idx == 0) { - // printf("=======Before Q@S, block_idx: %d, thread_idx: %d=======\n", blk, thread_idx); - // cute::print_tensor(tOrKV); - // cute::print_tensor(sQ_K_scaled_slice); - // } - - math_barriers.ordered_or_wait(warpgroup_idx); - warpgroup_arrive(); - gemm_zero_acc(o1_thr_mma, tOrKV, tOrQ(_, _, _, 0), tOrO); - warpgroup_commit_batch(); // q@kv batch - math_barriers.notify_next_blocked(warpgroup_idx); - } - if constexpr (!is_first_block) { - warpgroup_wait<0>(); // q@kv batch - // ======DEBUG======= - // if (blk <= 1 && thread_idx == 0) { - // printf("\n"); - // printf("=======O_inter after Q@S, block_idx: %d, thread_idx: %d=======\n", blk, thread_idx); - // cute::print_tensor(tOrO); - // } - DPRINTF0_WG("compute: q_pipeline.consumer_release: smem_pipe_read:%d\n", q_smem_pipe_read.index()); - q_pipeline.consumer_release(q_smem_pipe_read); - ++q_smem_pipe_read; - o1_epi(tOrO); - } - - auto tSKrSK = partition_fragment_C(sk_thr_mma, sVkv(_, _, _0{})); - if constexpr (!is_first_block) { - auto tSKrS = make_acc_into_op(tKVrKV, typename TiledMmaSK::LayoutA_TV{}); - warpgroup_fence_operand(tSKrSK); - warpgroup_fence_operand(tSKrS); - math_barriers.ordered_or_wait(warpgroup_idx); - warpgroup_arrive(); - - // ======DEBUG======= - // if (blk <= 6 && thread_idx == 0) { - // printf("=======Before K@S, block_idx: %d, thread_idx: %d=======\n", blk, thread_idx); - // cute::print_tensor(tSKrS); - // } - - // SK: K_scaled is in stage 1 of sQ_K_scaled - gemm_zero_acc(sk_tiled_mma, tSKrS, tSKrK(_, _, _, 1), tSKrSK); - warpgroup_commit_batch(); - math_barriers.notify_next_blocked(warpgroup_idx); - warpgroup_wait<0>(); - } - // ======DEBUG======= - // if (blk <= 6 && thread_idx == 0) { - // printf("=======After K@S, block_idx: %d, thread_idx: %d=======\n", blk, thread_idx); - // cute::print_tensor(tSKrSK); - // } - - DPRINTF0_WG("compute: v_pipeline.consumer_wait: smem_pipe_read:%d\n", v_smem_pipe_read.index()); - v_pipeline.consumer_wait(v_smem_pipe_read); - auto tSKrV = sk_load_v(v_smem_pipe_read.index()); - if constexpr (!is_first_block) { - // sk_epi(tSKrSK, alpha_smem_pipe_read); - // V' = V - SK - transform(tSKrV, tSKrSK, tSKrV, [](auto v, auto sk) { return v - Element(sk); }); - } - - kk_pipeline.consumer_wait(kk_smem_pipe_read); - beta_pipeline.consumer_wait(beta_smem_pipe_read); - cutlass::arch::fence_view_async_shared(); - // KK inverse - if (warpgroup_idx == 0) { - kk_inv(kk_smem_pipe_read); - } - // wait for KK inverse ready - cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - - auto tNewVrA = make_acc_into_op(tSKrV, typename TiledMmaNewV::LayoutA_TV{}); - auto tNewVrC = partition_fragment_C(newv_thr_mma, select<0, 1>(TileShapeNewV{})); - warpgroup_fence_operand(tNewVrA); - warpgroup_fence_operand(tNewVrC); - math_barriers.ordered_or_wait(warpgroup_idx); - warpgroup_arrive(); - // if constexpr (is_final_block) { - // if (thread_idx == 0) { - // printf("\n"); - // printf("=======tNewVrA, tNewVrB before V'@T, block_idx: %d, thread_idx: %d=======\n", blk, - // thread_idx); printf("tNewVrA\n"); cute::print_tensor(tNewVrA); printf("sKK_opd\n"); - // cute::print_tensor(sKK_opd); - // printf("=======tNewVrA, tNewVrB before V'@T, block_idx: %d, thread_idx: %d=======\n", blk, - // thread_idx); printf("\n"); - // } - // } - // NewV = V'T - gemm_zero_acc(o1_thr_mma, tNewVrA, tNewVrB(_, _, _, kk_smem_pipe_read.index()), tNewVrC); - warpgroup_commit_batch(); // new_v batch - math_barriers.notify_next_blocked(warpgroup_idx); - warpgroup_wait<0>(); // new_v batch - // if constexpr (is_final_block) { - // if (thread_idx == 0) { - // printf("\n"); - // printf("=======tNewVrC after V'@T, block_idx: %d, thread_idx: %d=======\n", blk, thread_idx); - // cute::print_tensor(tNewVrC); - // } - // } - DPRINTF0_WG("compute: v_pipeline.consumer_release: smem_pipe_read:%d\n", v_smem_pipe_read.index()); - ++v_smem_pipe_read; // NOTE: if we delay this increment after consumer_release, race condition happens, - // why? - v_pipeline.consumer_release(v_smem_pipe_read); - - kk_pipeline.consumer_release(kk_smem_pipe_read); - ++kk_smem_pipe_read; - beta_pipeline.consumer_release(beta_smem_pipe_read); - ++beta_smem_pipe_read; - - ///////////////////////////////////////////////////////////////////////// - // 2. compute qkv - // 2.2 QK @ V, NOTE: use old KV here and QK is scaled - qk_pipeline.consumer_wait(qk_smem_pipe_read); - auto tOrV_or_tKVrV = make_acc_into_op(tNewVrC, typename TiledMmaKV::LayoutA_TV{}); - warpgroup_fence_operand(tOrV_or_tKVrV); - warpgroup_fence_operand(tOrO); - math_barriers.ordered_or_wait(warpgroup_idx); - warpgroup_arrive(); - // (V_new)^T @ QK - if constexpr (is_first_block) { - gemm_zero_acc(o2_tiled_mma, tOrV_or_tKVrV, tOrQK(_, _, _, qk_smem_pipe_read.index()), tOrO); - } else { - gemm(o2_tiled_mma, tOrV_or_tKVrV, tOrQK(_, _, _, qk_smem_pipe_read.index()), tOrO); - } - warpgroup_commit_batch(); // qk@v batch - math_barriers.notify_next_blocked(warpgroup_idx); - warpgroup_wait<0>(); // qk@v batch - // if (blk <= 6 && thread_idx == 0) { - // printf("\n"); - // printf("=======O_intra after NewV@QK, block_idx: %d, thread_idx: %d=======\n", blk, thread_idx); - // cute::print_tensor(tOrO); - // } - qk_pipeline.consumer_release(qk_smem_pipe_read); - ++qk_smem_pipe_read; - o_store(tOrO); - - ///////////////////////////////////////////////////////////////////////// - // 3. update KV - Tensor tKVsK = kv_thr_mma.partition_B(sQ_K_scaled_Kt); - Tensor tKVrK = kv_thr_mma.make_fragment_B(tKVsK); - - if constexpr (NeedsAlpha) { - alpha_last_pipeline.consumer_wait(alpha_last_smem_pipe_read); - cutlass::arch::fence_view_async_shared(); - } - s_decay(tKVrKV, alpha_last_smem_pipe_read); - - // synchronize 2 WGs before rewriting sQ_K_scaled - cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - // exp(alpha_last - alpha) * K - // Each WG iterates over 2 slices of 32 elements each. - // WG0 (thread_idx < 128): wg_idx=0, alpha_last indices {0,1}, K/output dim1=0 - // WG1 (thread_idx >= 128): wg_idx=1, alpha_last indices {2,3}, K/output dim1=1 - { - int wg_idx = thread_idx / 128; // 0 or 1 - int alpha_base = wg_idx * 2; // 0 or 2 - - // Allocate K/Alpha register fragments once (reused across slices) - auto tQKrK_wg = qk_thr_mma_rs_quar.partition_fragment_A(sKqk_slice(_, _, _0{}, make_coord(_0{}, _0{}))); - auto tArA_wg = make_fragment_like(tQKrK_wg); - - for (int s = 0; s < 2; ++s) { - // S2R Alpha - int alpha_col = alpha_base + s; - auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); - auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); - copy(CopyAlphaAtom{}, tAsA_cur, tArA_wg); - - // S2R K - auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(s, wg_idx)); - auto tQKsK_cur = thr_load_qk_quar.partition_S(sKqk_cur); - auto tQKrK_cv = thr_load_qk_quar.retile_D(tQKrK_wg); - copy(tiled_load_qk_quar, tQKsK_cur, tQKrK_cv); - - // element-wise: exp(alpha_last - alpha) * K - int alast_idx = alpha_base + s; - auto alpha_last_cur = sAlast_slice(_, alast_idx); - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tQcMq_quar(i); - auto [seq, t] = coord; - auto alpha = tArA_wg(i); - auto k = tQKrK_wg(i); - auto alpha_last = alpha_last_cur(t); - auto k_scaled = Element(exp2f(alpha_last - alpha) * float(k)); - tQKrK_wg(i) = k_scaled; - if constexpr (is_final_block) { - if (seq >= B) { - tQKrK_wg(i) = Element(0.0f); - } - } - }); - - // R2S K -> stage 0 (reuse for KV update) - auto sQ_scaled_cur = sQ_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); - auto tQKsK_out = thr_store_qk_quar.partition_D(sQ_scaled_cur); - auto tQKrK_out_cv = thr_store_qk_quar.retile_S(tQKrK_wg); - copy(tiled_store_qk_quar, tQKrK_out_cv, tQKsK_out); - } - } - // wait for smemq_k_scaled ready - cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - // fence to produce data for WGMMA async proxy - cutlass::arch::fence_view_async_shared(); - - if constexpr (NeedsAlpha) { - alpha_last_pipeline.consumer_release(alpha_last_smem_pipe_read); - ++alpha_last_smem_pipe_read; - } - - DPRINTF0_WG("[%d,%d,%d,%d]** dispatch KV WGMMA\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); - warpgroup_fence_operand(tOrV_or_tKVrV); - warpgroup_fence_operand(tKVrKV); - - math_barriers.ordered_or_wait(warpgroup_idx); - warpgroup_arrive(); - gemm(kv_tiled_mma, tOrV_or_tKVrV, tKVrK(_, _, _, 0), tKVrKV); - warpgroup_commit_batch(); // k@v batch - math_barriers.notify_next_blocked(warpgroup_idx); - warpgroup_wait<0>(); - // if constexpr (is_final_block) { - // if (thread_idx == 0) { - // printf("\n"); - // printf("=======After K^T@NewV, block_idx: %d, thread_idx: %d=======\n", blk, thread_idx); - // printf("tKVrKV\n"); - // cute::print_tensor(tKVrKV); - // } - // } - - if constexpr (NeedsAlpha) { - alpha_pipeline.consumer_release(alpha_smem_pipe_read); - ++alpha_smem_pipe_read; - } - - DPRINTF0_WG("compute: k_pipeline.consumer_release: smem_pipe_read:%d\n", k_smem_pipe_read.index()); - k_pipeline.consumer_release(k_smem_pipe_read); - ++k_smem_pipe_read; - - // if (blk <= 6 && thread_idx == 0) { - // printf("\n"); - // printf("=======After S epilogue, block_idx: %d, thread_idx: %d, head_idx: %d=======\n", blk, - // thread_idx, q_head_idx); printf("tKVrKV\n"); cute::print_tensor(tKVrKV); - // } - }; - - if constexpr (!kInitStateFromInput) { - clear(tKVrKV); - if (num_blocks == 1) { - compute_loop_body(0, /*is_first_block_=*/cute::true_type{}, /*is_final_block_=*/cute::true_type{}); - } else { - compute_loop_body(0, /*is_first_block_=*/cute::true_type{}, /*is_final_block_=*/cute::false_type{}); - } - } else { - kv_load(tKVrKV); // GMEM -> Register, only once at the beginning - if (num_blocks == 1) { - compute_loop_body(0, /*is_first_block_=*/cute::false_type{}, /*is_final_block_=*/cute::true_type{}); - } else { - compute_loop_body(0, /*is_first_block_=*/cute::false_type{}, /*is_final_block_=*/cute::false_type{}); - } - } - CUTE_NO_UNROLL - for (int blk = 1; blk < num_blocks - 1; ++blk) { - compute_loop_body(blk, /*is_first_block_=*/cute::false_type{}, /*is_final_block_=*/cute::false_type{}); - } - if (num_blocks != 1) { - compute_loop_body( - num_blocks - 1, - /*is_first_block_=*/cute::false_type{}, - /*is_final_block_=*/cute::true_type{}); - } - if (params.ptr_output_state != nullptr) { - kv_store(); - } - } - - template - CUTE_DEVICE void - compute_aux_safe( - Params const& params, - ProblemShape const& problem_size, - WorkDesc const& work_desc, - MainloopQPipeline& q_pipeline, - QPipelineState& q_smem_pipe_read, - MainloopKPipeline& k_pipeline, - KPipelineState& k_smem_pipe_read, - MainloopQKPipeline& qk_pipeline, - QKPipelineState& qk_smem_pipe_write, - MainloopKKPipeline& kk_pipeline, - KKPipelineState& kk_smem_pipe_write, - MainloopAlphaPipeline& alpha_pipeline, - AlphaPipelineState& alpha_smem_pipe_read, - MainloopBetaPipeline& beta_pipeline, - BetaPipelineState& beta_smem_pipe_read, - MainloopAlphaLastPipeline& alpha_last_pipeline, - AlphaLastPipelineState& alpha_last_smem_pipe_write, - SharedStorage& storage) { - using TileShape_SubChunk = Shape<_16, _16, _32>; - int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; - - float scale = params.scale; - - Tensor Beta = make_tensor(make_smem_ptr(storage.smem_beta.data()), SmemLayoutBeta{}); - - Tensor sQqk = make_tensor(make_smem_ptr(storage.smem_q.data()), QKSmemLayoutQ{}); - Tensor sKqk = make_tensor(make_smem_ptr(storage.smem_k.data()), QKSmemLayoutK{}); - - Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); - Tensor sAqkk = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKKSmemLayoutAlpha{}); - Tensor sAlast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); - - Tensor sKkv = make_tensor(make_smem_ptr(storage.smem_k.data()), KVSmemLayoutK{}); - Tensor sVkv = make_tensor(make_smem_ptr(storage.smem_v.data()), KVSmemLayoutV{}); - Tensor sQK = make_tensor(make_smem_ptr(storage.smem_qk.data()), SmemLayoutQK{}); - Tensor sO = make_tensor(make_smem_ptr(storage.smem_o.data()), SmemLayoutO{}); - - static_assert(sizeof(InverseType) == sizeof(Element)); - Tensor sKK_inv = make_tensor(make_smem_ptr(storage.smem_kk.data()), SmemLayoutKK{}); - Tensor sKK_opd = make_tensor(make_smem_ptr(reinterpret_cast(storage.smem_kk.data())), SmemLayoutKK{}); - - constexpr int BK = 32; // should be same as TileShape_SubChunk - constexpr int NK = 128 / BK; - - /////////////////////////////////////////////////////////////////////////// - // Q@K - auto qk_tiled_mma_rs = TiledMmaQK_RS{}; - auto qk_thr_mma_rs = qk_tiled_mma_rs.get_thread_slice(thread_idx); - - auto cMqk = make_identity_tensor(select<0, 1>(TileShapeQK{})); // (QTok, KTok) - auto tQKcMqk = qk_thr_mma_rs.partition_C(cMqk); // (idx) -> (tok_q, tok_k) - auto cMq = make_identity_tensor(select<0, 2>(TileShapeQK{})); // (QTok, HeadDim) - auto tQcMq = qk_thr_mma_rs.partition_A(cMq); // (idx) -> (tok_q, head_dim) - - auto const seq_idx = work_desc.seq_idx; - auto const q_head_idx = work_desc.q_head_idx(); - auto const k_head_idx = work_desc.k_head_idx(); - auto const v_head_idx = work_desc.v_head_idx(); - - auto qk_kk_subchunk_mma_and_store = [&](int blk) INLINE_LAMBDA { - using CopyOp_R2S = SM90_U32x2_STSM_N; - using CopyAlphaAtom = Copy_Atom, ElementAlpha>; - // Q/K S2R: use BF16 MMA's LDSM tiled copy for efficient shared memory loads, - // then convert register layout to TF32 MMA layout via warp shuffles. - // This replaces the previous AutoVectorizingCopy<16> which caused 50% more smem traffic. - using CopyQKAtom_LDSM = Copy_Atom; - // TF32 MMA: float(Element) → tf32 → MMA, better precision than fp16/bf16 MMA - using MMA = SM80_16x8x8_F32TF32TF32F32_TN; - using TiledMma_SubChunk = - decltype(make_tiled_mma(MMA{}, Layout>{}, TileShape_SubChunk{})); - // BF16 MMA (same shape 16x8x8) used only for creating LDSM-compatible tiled copies - using MMA_BF16 = SM80_16x8x8_F32BF16BF16F32_TN; - using TiledMma_BF16_SubChunk = - decltype(make_tiled_mma(MMA_BF16{}, Layout>{}, TileShape_SubChunk{})); - - int local_thread_idx = thread_idx % 64; - auto tiledmma_subchunk = TiledMma_SubChunk{}; - auto thr_mma_subchunk = tiledmma_subchunk.get_thread_slice(local_thread_idx); - auto tiledmma_bf16_subchunk = TiledMma_BF16_SubChunk{}; - auto thr_mma_bf16_subchunk = tiledmma_bf16_subchunk.get_thread_slice(local_thread_idx); - - // Alpha S2R: load in BF16 MMA layout so gating happens before the layout shuffle, - // reducing register pressure (alpha can be freed before the shuffle). - // BF16-layout alpha copies for operand A and B (for element-wise gating) - auto alpha_Q_bf16_tiled_copy = make_tiled_copy_A(CopyAlphaAtom{}, tiledmma_bf16_subchunk); - auto alpha_Kt_bf16_tiled_copy = make_tiled_copy_B(CopyAlphaAtom{}, tiledmma_bf16_subchunk); - // Q/K S2R: LDSM copies using BF16 MMA layout for efficient ldmatrix loads - auto Q_tiled_copy = make_tiled_copy_A(CopyQKAtom_LDSM{}, tiledmma_bf16_subchunk); - auto Kt_tiled_copy = make_tiled_copy_B(CopyQKAtom_LDSM{}, tiledmma_bf16_subchunk); - // R2S copies for accumulators (C layout, same for all MMA types with same output) - auto O_tiled_copy = make_tiled_copy_C(Copy_Atom{}, tiledmma_subchunk); - auto O_tiled_copy_kk = make_tiled_copy_C(Copy_Atom{}, tiledmma_subchunk); - - auto alpha_Q_bf16_thr_copy = alpha_Q_bf16_tiled_copy.get_thread_slice(local_thread_idx); - auto alpha_Kt_bf16_thr_copy = alpha_Kt_bf16_tiled_copy.get_thread_slice(local_thread_idx); - auto Q_thr_copy = Q_tiled_copy.get_thread_slice(local_thread_idx); - auto Kt_thr_copy = Kt_tiled_copy.get_thread_slice(local_thread_idx); - auto O_thr_copy = O_tiled_copy.get_thread_slice(local_thread_idx); - auto O_thr_copy_kk = O_tiled_copy_kk.get_thread_slice(local_thread_idx); - - // index tensor - auto cMqk_subchunk = make_identity_tensor(select<0, 1>(TileShape_SubChunk{})); - auto tQKcMqk_subchunk = thr_mma_subchunk.partition_C(cMqk_subchunk); - - // do MMA at the granularity of 16x16x64 with two warps - constexpr auto tiler_subchunk_alpha = Shape<_16, Shape<_32, _1>>{}; - constexpr auto tiler_subchunk_qk = Shape<_16, Shape<_32, _1>>{}; - constexpr auto tiler_subchunk_beta = Shape<_16>{}; - auto sQqk_curr = sQqk(_, _, q_smem_pipe_read.index()); - auto sKqk_curr = sKqk(_, _, k_smem_pipe_read.index()); - auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); - Tensor sBeta_curr = Beta(_, beta_smem_pipe_read.index()); - - // (_16,(_32,_1),_4,(_2,_2)):(_64,(_1,_0),_1024,(_32,_4096)) - auto sQqk_slice = flat_divide(sQqk_curr, tiler_subchunk_qk); - auto sKqk_slice = flat_divide(sKqk_curr, tiler_subchunk_qk); - // (_16,(_32,_1),_4,(_1,_4)):(_32,(_1,_0),_512,(_0,_2048)) - auto sAqkq_slice = flat_divide(sAqkq_curr, tiler_subchunk_alpha); - auto sBeta_slice = flat_divide(sBeta_curr, tiler_subchunk_beta); - - // Acc results - constexpr auto tiler_acc_qk_kk = Shape<_16, _16>{}; - static_assert(sizeof(Element) == 2); - auto sQK_curr = sQK(_, _, qk_smem_pipe_write.index()); - auto sQK_slice = flat_divide(sQK_curr, tiler_acc_qk_kk); - auto sKK_inv_curr = sKK_inv(_, _, kk_smem_pipe_write.index()); - auto sKK_inv_slice = flat_divide(sKK_inv_curr, tiler_acc_qk_kk); - - // used for make_fragment_like in Alpha and S2R (TF32 MMA layout) - Tensor sQqk_1_0 = sQqk_slice(_, _, _1{}, make_coord(_0{}, _0{})); - Tensor sKqk_1_0 = sKqk_slice(_, _, _1{}, make_coord(_0{}, _0{})); - Tensor tQKrQ_1_0 = thr_mma_subchunk.partition_fragment_A(sQqk_1_0); - Tensor tQKrKt_1_0 = thr_mma_subchunk.partition_fragment_B(sKqk_1_0); - auto tv_layout_mma_A = tQKrQ_1_0.layout(); - auto tv_layout_mma_B = tQKrKt_1_0.layout(); - - // BF16 MMA fragment layouts for LDSM-based S2R loads (same shape, different TV mapping) - Tensor tQKrQ_bf16_1_0 = thr_mma_bf16_subchunk.partition_fragment_A(sQqk_1_0); - Tensor tQKrKt_bf16_1_0 = thr_mma_bf16_subchunk.partition_fragment_B(sKqk_1_0); - auto tv_layout_bf16_mma_A = tQKrQ_bf16_1_0.layout(); - auto tv_layout_bf16_mma_B = tQKrKt_bf16_1_0.layout(); - - // S2R Q/K/G for operand A at row r, head dim slice j, and element-wise compute. - // Loads alpha once in BF16 MMA layout, derives g_first via warp shuffle (8 shuffles, - // replaces 1 S2R load), gates Q/K before the BF16→TF32 layout conversion. - // Also extracts g_first in operand B layout for free (broadcast → register copy). - // j0 = j % 2, j1 = j / 2: precomputed by caller to avoid redundant div/mod. - // returns (tQKrQ, tQKrK, tArAfirst_kt) = (Q * exp2(g - g_first), K * exp2(g - g_first), g_first in B - // layout) - auto s2r_compute_subchunk_operandA = [&](auto r_, int j, int j0, int j1) INLINE_LAMBDA { - // S2R g_r_j in BF16 MMA operand A layout (single load) - Tensor sAqkq_r_j = sAqkq_slice(_, _, r_, make_coord(_0{}, j)); - Tensor tAsA_r_j = alpha_Q_bf16_thr_copy.partition_S(sAqkq_r_j); - Tensor tArA_r_j = make_fragment_like(tv_layout_bf16_mma_A); - Tensor tArA_r_j_cv = alpha_Q_bf16_thr_copy.retile_D(tArA_r_j); - copy(alpha_Q_bf16_tiled_copy, tAsA_r_j, tArA_r_j_cv); - - // Derive g_first (alpha[row=0, :]) from tArA_r_j via warp shuffle, - // directly into operand B layout (8 values instead of 16). - // g_first is broadcast (all M rows identical), so operand B only needs the - // v1=0 subset of operand A. We shuffle v1=0 values from t1=0 thread and - // output directly as operand B fragment, saving 8 float registers. - Tensor tArAfirst_r_j_kt = make_fragment_like(tv_layout_bf16_mma_B); - broadcast_row0_operandA_to_operandB_bf16_layout(tArA_r_j, tArAfirst_r_j_kt, local_thread_idx); - - // gqn_r_j = exp2(g_r_j - g_r_j_first[None, :]) in BF16 MMA A layout. - // g_first per k-iter is in tArAfirst_r_j_kt: frag_B(2j)=K_lo, frag_B(2j+1)=K_hi. - // In A layout: v1=0 indices (4j+0, 4j+1) have same K as v1=1 indices (4j+2, 4j+3), - // so g_first for index 4j+{0,2} = frag_B(2j), for 4j+{1,3} = frag_B(2j+1). - CUTE_UNROLL - for (int k = 0; k < 4; k++) { - auto gf_lo = tArAfirst_r_j_kt(2 * k); // g_first at K = 2*t0 - auto gf_hi = tArAfirst_r_j_kt(2 * k + 1); // g_first at K = 2*t0+1 - tArA_r_j(4 * k + 0) = exp2f(tArA_r_j(4 * k + 0) - gf_lo); // v0=0, v1=0 - tArA_r_j(4 * k + 1) = exp2f(tArA_r_j(4 * k + 1) - gf_hi); // v0=1, v1=0 - tArA_r_j(4 * k + 2) = exp2f(tArA_r_j(4 * k + 2) - gf_lo); // v0=0, v1=1 - tArA_r_j(4 * k + 3) = exp2f(tArA_r_j(4 * k + 3) - gf_hi); // v0=1, v1=1 - } - - Tensor sQqk_r_j = sQqk_slice(_, _, r_, make_coord(j0, j1)); - Tensor sKqk_r_j = sKqk_slice(_, _, r_, make_coord(j0, j1)); - - // --- Process Q --- - // S2R Q in BF16 MMA layout - Tensor tQKrQ_r_j_bf16 = make_fragment_like(tv_layout_bf16_mma_A); - Tensor tQKsQ_r_j = Q_thr_copy.partition_S(sQqk_r_j); - Tensor tQKrQ_r_j_bf16_cv = Q_thr_copy.retile_D(tQKrQ_r_j_bf16); - copy(Q_tiled_copy, tQKsQ_r_j, tQKrQ_r_j_bf16_cv); - // gate: Q * exp2(g - g_first) in BF16 MMA layout, producing float - Tensor tQKrQ_r_j_float = make_fragment_like(tv_layout_bf16_mma_A); - cute::transform( - tQKrQ_r_j_bf16, tArA_r_j, tQKrQ_r_j_float, [&](auto q, auto g) { return float(q) * g; }); - // convert BF16 MMA layout → TF32 MMA layout in-place via warp shuffles - convert_bf16_to_tf32_operandA_layout(tQKrQ_r_j_float, local_thread_idx); - // NOTE: triton tl.dot also lets MMA hardware for truncation - // recast float storage as tf32 view (zero cost, same 32-bit registers; MMA hw truncates) - auto tQKrQ_r_j = recast(tQKrQ_r_j_float); - - // --- Process K (sequential, after Q is done to reduce peak reg usage) --- - Tensor tQKrK_r_j_bf16 = make_fragment_like(tv_layout_bf16_mma_A); - Tensor tQKsK_r_j = Q_thr_copy.partition_S(sKqk_r_j); - Tensor tQKrK_r_j_bf16_cv = Q_thr_copy.retile_D(tQKrK_r_j_bf16); - copy(Q_tiled_copy, tQKsK_r_j, tQKrK_r_j_bf16_cv); - Tensor tQKrK_r_j_float = make_fragment_like(tv_layout_bf16_mma_A); - cute::transform( - tQKrK_r_j_bf16, tArA_r_j, tQKrK_r_j_float, [&](auto k, auto g) { return float(k) * g; }); - // convert BF16 MMA layout → TF32 MMA layout in-place via warp shuffles - convert_bf16_to_tf32_operandA_layout(tQKrK_r_j_float, local_thread_idx); - auto tQKrK_r_j = recast(tQKrK_r_j_float); - - return cute::make_tuple(tQKrQ_r_j, tQKrK_r_j, tArAfirst_r_j_kt); - }; - - // S2R K/G for operand B at column c, head dim slice j, and element-wise compute - // Loads alpha in BF16 MMA B layout, gates K in BF16 MMA B layout (before shuffle), - // then converts gated result to TF32 MMA B layout. - // tArAfirst_kt: pre-loaded g_first register tensor (BF16 MMA B layout) for computing gktn = exp2(g_first - - // g_c) returns tQKrKt = K_c * exp2(g_first - g_c) - auto s2r_compute_subchunk_operandB = - [&](auto c_, int j, int j0, int j1, auto const& tArAfirst_kt) INLINE_LAMBDA { - // S2R g_c_j in BF16 MMA operand B layout - Tensor sAqkq_c_j = sAqkq_slice(_, _, c_, make_coord(_0{}, j)); - Tensor tAsA_c_j = alpha_Kt_bf16_thr_copy.partition_S(sAqkq_c_j); - Tensor tArA_c_j = make_fragment_like(tv_layout_bf16_mma_B); - Tensor tArA_c_j_cv = alpha_Kt_bf16_thr_copy.retile_D(tArA_c_j); - copy(alpha_Kt_bf16_tiled_copy, tAsA_c_j, tArA_c_j_cv); - - // compute gktn_c_j = exp2(g_first - g_c_j) in BF16 MMA B layout - cute::transform( - tArA_c_j, tArAfirst_kt, tArA_c_j, [&](auto g, auto g_first) { return exp2f(g_first - g); }); - - // S2R k_c_j using BF16 LDSM - Tensor sKqk_c_j = sKqk_slice(_, _, c_, make_coord(j0, j1)); - Tensor tQKrKt_c_j_bf16 = make_fragment_like(tv_layout_bf16_mma_B); - Tensor tQKsKt_c_j = Kt_thr_copy.partition_S(sKqk_c_j); - Tensor tQKrKt_c_j_bf16_cv = Kt_thr_copy.retile_D(tQKrKt_c_j_bf16); - copy(Kt_tiled_copy, tQKsKt_c_j, tQKrKt_c_j_bf16_cv); - - // convert bf16 → float in BF16 MMA B layout - Tensor tQKrKt_c_j_float = make_fragment_like(tv_layout_bf16_mma_B); - // gate in BF16 MMA B layout (alpha and K are in the same layout) - cute::transform( - tQKrKt_c_j_bf16, tArA_c_j, tQKrKt_c_j_float, [&](auto k, auto g) { return float(k) * g; }); - - // convert BF16 MMA layout → TF32 MMA layout in-place via warp shuffles - convert_bf16_to_tf32_operandB_layout(tQKrKt_c_j_float, local_thread_idx); - auto tQKrKt_c_j = recast(tQKrKt_c_j_float); - - return tQKrKt_c_j; - }; - - // R2S (register to shared memory) store for subchunk accumulator results - // Stores both tQKrQK (QK accumulator, fp32 -> Element) and tKKrKK (KK accumulator, fp32 -> InverseType) - // into their respective shared memory tiles at position (r_, c_) - auto r2s_subchunk_acc = [&](auto r_, auto c_, auto const& tQKrQK, auto const& tKKrKK) INLINE_LAMBDA { - // R2S KK - Tensor sKK_inv_r_c = sKK_inv_slice(_, _, r_, c_); - Tensor tKKsKK_r_c = O_thr_copy_kk.partition_D(sKK_inv_r_c); - Tensor tKKrKK_cv = O_thr_copy_kk.retile_S(tKKrKK); - auto tKKrKK_cvt_cv = make_fragment_like(tKKrKK_cv); - cute::transform(tKKrKK_cv, tKKrKK_cvt_cv, [](auto v) { return InverseType(v); }); - copy(O_tiled_copy_kk, tKKrKK_cvt_cv, tKKsKK_r_c); - - // R2S QK - Tensor sQK_r_c = sQK_slice(_, _, r_, c_); - Tensor tQKsQK_r_c = O_thr_copy.partition_D(sQK_r_c); - Tensor tQKrQK_cv = O_thr_copy.retile_S(tQKrQK); - auto tQKrQK_cvt_cv = make_fragment_like(tQKrQK_cv); - cute::transform(tQKrQK_cv, tQKrQK_cvt_cv, [](auto v) { return Element(v); }); - copy(O_tiled_copy, tQKrQK_cvt_cv, tQKsQK_r_c); - }; - - // do tensor core GEMM with single 16x16x128 - // NOTE: should use safe_gate with lower_bound >= -5, otherwise overflow issues - auto gemm_tensor_core_1x16x16x128 = - [&](auto r_, auto c_, auto is_diagonal_, auto is_first_subchunk_) INLINE_LAMBDA { - constexpr bool is_first_subchunk = decltype(is_first_subchunk_)::value; - - // allocate acc_r_c [16, 16] - Tensor tQKrQK_r_c = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_r_c = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - clear(tQKrQK_r_c); - clear(tKKrKK_r_c); - // wait for data ready - if constexpr (is_first_subchunk) { - q_pipeline.consumer_wait(q_smem_pipe_read); - k_pipeline.consumer_wait(k_smem_pipe_read); - if constexpr (NeedsAlpha) { - alpha_pipeline.consumer_wait(alpha_smem_pipe_read); - } - } - - // for loop head dim - CUTE_NO_UNROLL - for (int j = 0; j < NK; ++j) { - int j0 = j % 2, j1 = j / 2; - // S2R Q/K/G/g_first for operand A and element-wise compute - auto [tQKrQ_r_j, tQKrK_r_j, tArAfirst_r_j_kt] = s2r_compute_subchunk_operandA(r_, j, j0, j1); - // S2R K/G for operand B and element-wise compute - auto tQKrKt_c_j = s2r_compute_subchunk_operandB(c_, j, j0, j1, tArAfirst_r_j_kt); - - // q_r_j/k_r_j @ k_c_j, accumulate acc_r_c - gemm(tiledmma_subchunk, tQKrQ_r_j, tQKrKt_c_j, tQKrQK_r_c); - gemm(tiledmma_subchunk, tQKrK_r_j, tQKrKt_c_j, tKKrKK_r_c); - } - - // S2R beta_j (maybe resident in register?) - // epilogue: qk^t * scale - cute::transform(tQKrQK_r_c, [scale](auto v) { return v * scale; }); - // epilogue: kk^t * beta_r - if constexpr (is_first_subchunk) { - beta_pipeline.consumer_wait(beta_smem_pipe_read); - cutlass::arch::fence_view_async_shared(); - } - Tensor sBeta_r = sBeta_slice(_, r_); - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tQKcMqk_subchunk(i); - auto [s, t] = coord; - tKKrKK_r_c(i) *= sBeta_r(s); - }); - - // R2S qk_r_c, kk_r_c, wait for current QK/KK free - if constexpr (is_first_subchunk) { - kk_pipeline.producer_acquire(kk_smem_pipe_write); - } - if constexpr (is_first_subchunk) { - qk_pipeline.producer_acquire(qk_smem_pipe_write); - } - r2s_subchunk_acc(r_, c_, tQKrQK_r_c, tKKrKK_r_c); - }; - - // zero fill for upper triangular of QK and KK, because smem is randomly initialized - auto zero_fill = [&](int row, int col) INLINE_LAMBDA { - auto sQK_r_c = sQK_slice(_, _, row, col); - auto sKK_r_c = sKK_inv_slice(_, _, row, col); - // allocate regs - Tensor tQKrQK_r_c = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tQKsQK_r_c = O_thr_copy.partition_D(sQK_r_c); - Tensor tQKrQK_r_c_cv = O_thr_copy.retile_S(tQKrQK_r_c); - Tensor tKKsKK_r_c = O_thr_copy_kk.partition_D(sKK_r_c); - Tensor tKKrKK_r_c_cv = O_thr_copy_kk.retile_S(tQKrQK_r_c); - auto tQKrQK_r_c_cvt_cv = make_fragment_like(tQKrQK_r_c_cv); - auto tKKrKK_r_c_cvt_cv = make_fragment_like(tKKrKK_r_c_cv); - // zero fill - clear(tQKrQK_r_c_cvt_cv); - clear(tKKrKK_r_c_cvt_cv); - // R2S - copy(O_tiled_copy, tQKrQK_r_c_cvt_cv, tQKsQK_r_c); - copy(O_tiled_copy_kk, tKKrKK_r_c_cvt_cv, tKKsKK_r_c); - }; - - // g_i_j/q_i_j/k_i_j: the j-th head dim slice of the i-th subchunk - if (thread_idx < 64) { - // Q/K0@K0, Q/K3@K3, Q/K3@K0, Q/K3@K1, Q/K3@K2 - - // NOTE: tensor core MMA for safe gate with lower_bound >= -5 - gemm_tensor_core_1x16x16x128( - Int<0>{}, - Int<0>{}, - /*is_diagonal_=*/cute::true_type{}, - /*is_first_subchunk_=*/cute::true_type{}); - - // Q/K3@K0, Q/K3@K1, Q/K3@K2 - // allocate acc_3_0, acc_3_1, acc_3_2 [16, 16] - Tensor tQKrQK_3_0 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_3_0 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tQKrQK_3_1 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_3_1 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tQKrQK_3_2 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_3_2 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tQKrQK_3_3 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_3_3 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - clear(tQKrQK_3_0); - clear(tKKrKK_3_0); - clear(tQKrQK_3_1); - clear(tKKrKK_3_1); - clear(tQKrQK_3_2); - clear(tKKrKK_3_2); - clear(tQKrQK_3_3); - clear(tKKrKK_3_3); - - // for loop head dim - CUTE_NO_UNROLL - for (int j = 0; j < NK; ++j) { - int j0 = j % 2, j1 = j / 2; - // S2R Q/K/G/g_first for operand A (row 3) and element-wise compute - auto [tQKrQ_3_j, tQKrK_3_j, tArAfirst_3_j_kt] = s2r_compute_subchunk_operandA(_3{}, j, j0, j1); - - // S2R K/G for operand B (col 0) and element-wise compute - auto tQKrKt_0_j = s2r_compute_subchunk_operandB(_0{}, j, j0, j1, tArAfirst_3_j_kt); - // q_3_j/k_3_j @ k_0_j, accumulate acc_3_0 - gemm(tiledmma_subchunk, tQKrQ_3_j, tQKrKt_0_j, tQKrQK_3_0); - gemm(tiledmma_subchunk, tQKrK_3_j, tQKrKt_0_j, tKKrKK_3_0); - - // S2R K/G for operand B (col 1) and element-wise compute - auto tQKrKt_1_j = s2r_compute_subchunk_operandB(_1{}, j, j0, j1, tArAfirst_3_j_kt); - // q_3_j/k_3_j @ k_1_j, accumulate acc_3_1 - gemm(tiledmma_subchunk, tQKrQ_3_j, tQKrKt_1_j, tQKrQK_3_1); - gemm(tiledmma_subchunk, tQKrK_3_j, tQKrKt_1_j, tKKrKK_3_1); - - // S2R K/G for operand B (col 2) and element-wise compute - auto tQKrKt_2_j = s2r_compute_subchunk_operandB(_2{}, j, j0, j1, tArAfirst_3_j_kt); - // q_3_j/k_3_j @ k_2_j, accumulate acc_3_2 - gemm(tiledmma_subchunk, tQKrQ_3_j, tQKrKt_2_j, tQKrQK_3_2); - gemm(tiledmma_subchunk, tQKrK_3_j, tQKrKt_2_j, tKKrKK_3_2); - - // S2R K/G for operand B (col 3) and element-wise compute - auto tQKrKt_3_j = s2r_compute_subchunk_operandB(_3{}, j, j0, j1, tArAfirst_3_j_kt); - // q_3_j/k_3_j @ k_3_j, accumulate acc_3_3 - gemm(tiledmma_subchunk, tQKrQ_3_j, tQKrKt_3_j, tQKrQK_3_3); - gemm(tiledmma_subchunk, tQKrK_3_j, tQKrKt_3_j, tKKrKK_3_3); - } - - // S2R beta (maybe resident in register?) - // epilogue: qk^t * scale - cute::transform(tQKrQK_3_0, [scale](auto v) { return v * scale; }); - cute::transform(tQKrQK_3_1, [scale](auto v) { return v * scale; }); - cute::transform(tQKrQK_3_2, [scale](auto v) { return v * scale; }); - cute::transform(tQKrQK_3_3, [scale](auto v) { return v * scale; }); - // epilogue: kk^t * beta_3 - Tensor sBeta_3 = sBeta_slice(_, _3{}); - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tQKcMqk_subchunk(i); - auto [s, t] = coord; - auto b = sBeta_3(s); - tKKrKK_3_0(i) *= b; - tKKrKK_3_1(i) *= b; - tKKrKK_3_2(i) *= b; - tKKrKK_3_3(i) *= b; - }); - - // R2S qk_3_0, kk_3_0, wait for current QK/KK free - r2s_subchunk_acc(_3{}, _0{}, tQKrQK_3_0, tKKrKK_3_0); - // R2S qk_3_1, kk_3_1 - r2s_subchunk_acc(_3{}, _1{}, tQKrQK_3_1, tKKrKK_3_1); - // R2S qk_3_2, kk_3_2 - r2s_subchunk_acc(_3{}, _2{}, tQKrQK_3_2, tKKrKK_3_2); - // R2S qk_3_3, kk_3_3 - r2s_subchunk_acc(_3{}, _3{}, tQKrQK_3_3, tKKrKK_3_3); - } else { - // Q/K1@K0, Q/K2@K0, Q/K2@K1, Q/K2@K2, Q/K1@K1 - - // Q/K1@K0 - // allocate acc_1_0 [16, 16] - Tensor tQKrQK_1_0 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_1_0 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tQKrQK_1_1 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_1_1 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - clear(tQKrQK_1_0); - clear(tKKrKK_1_0); - clear(tQKrQK_1_1); - clear(tKKrKK_1_1); - // wait for data ready - q_pipeline.consumer_wait(q_smem_pipe_read); - k_pipeline.consumer_wait(k_smem_pipe_read); - if constexpr (NeedsAlpha) { - alpha_pipeline.consumer_wait(alpha_smem_pipe_read); - } - - // for loop head dim - CUTE_NO_UNROLL - for (int j = 0; j < NK; ++j) { - int j0 = j % 2, j1 = j / 2; - // S2R Q/K/G/g_first for operand A (row 1) and element-wise compute - auto [tQKrQ_1_j, tQKrK_1_j, tArAfirst_1_j_kt] = s2r_compute_subchunk_operandA(_1{}, j, j0, j1); - - // S2R K/G for operand B (col 0) and element-wise compute - auto tQKrKt_0_j = s2r_compute_subchunk_operandB(_0{}, j, j0, j1, tArAfirst_1_j_kt); - // q_1_j/k_1_j @ k_0_j, accumulate acc_1_0 - gemm(tiledmma_subchunk, tQKrQ_1_j, tQKrKt_0_j, tQKrQK_1_0); - gemm(tiledmma_subchunk, tQKrK_1_j, tQKrKt_0_j, tKKrKK_1_0); - - // S2R K/G for operand B (col 1) and element-wise compute - auto tQKrKt_1_j = s2r_compute_subchunk_operandB(_1{}, j, j0, j1, tArAfirst_1_j_kt); - // q_1_j/k_1_j @ k_1_j, accumulate acc_1_1 - gemm(tiledmma_subchunk, tQKrQ_1_j, tQKrKt_1_j, tQKrQK_1_1); - gemm(tiledmma_subchunk, tQKrK_1_j, tQKrKt_1_j, tKKrKK_1_1); - } - - // S2R beta_j (maybe resident in register?) - // epilogue: qk^t * scale - cute::transform(tQKrQK_1_0, [scale](auto v) { return v * scale; }); - cute::transform(tQKrQK_1_1, [scale](auto v) { return v * scale; }); - // epilogue: kk^t * beta_1 - beta_pipeline.consumer_wait(beta_smem_pipe_read); - cutlass::arch::fence_view_async_shared(); - Tensor sBeta_1 = sBeta_slice(_, _1{}); - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tQKcMqk_subchunk(i); - auto [s, t] = coord; - tKKrKK_1_0(i) *= sBeta_1(s); - tKKrKK_1_1(i) *= sBeta_1(s); - }); - - // R2S qk_1_0, kk_1_0, wait for current QK/KK free - kk_pipeline.producer_acquire(kk_smem_pipe_write); - qk_pipeline.producer_acquire(qk_smem_pipe_write); - - r2s_subchunk_acc(_1{}, _0{}, tQKrQK_1_0, tKKrKK_1_0); - // R2S qk_1_1, kk_1_1 - r2s_subchunk_acc(_1{}, _1{}, tQKrQK_1_1, tKKrKK_1_1); - - // Q/K2@K0, Q/K2@K1 - // allocate acc_2_0, acc_2_1 [16, 16] - Tensor tQKrQK_2_0 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_2_0 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tQKrQK_2_1 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_2_1 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tQKrQK_2_2 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - Tensor tKKrKK_2_2 = partition_fragment_C(tiledmma_subchunk, select<0, 1>(TileShape_SubChunk{})); - clear(tQKrQK_2_0); - clear(tKKrKK_2_0); - clear(tQKrQK_2_1); - clear(tKKrKK_2_1); - clear(tQKrQK_2_2); - clear(tKKrKK_2_2); - - // for loop head dim - CUTE_NO_UNROLL - for (int j = 0; j < NK; ++j) { - int j0 = j % 2, j1 = j / 2; - // S2R Q/K/G/g_first for operand A (row 2) and element-wise compute - auto [tQKrQ_2_j, tQKrK_2_j, tArAfirst_2_j_kt] = s2r_compute_subchunk_operandA(_2{}, j, j0, j1); - - // S2R K/G for operand B (col 0) and element-wise compute - auto tQKrKt_0_j = s2r_compute_subchunk_operandB(_0{}, j, j0, j1, tArAfirst_2_j_kt); - // q_2_j/k_2_j @ k_0_j, accumulate acc_2_0 - gemm(tiledmma_subchunk, tQKrQ_2_j, tQKrKt_0_j, tQKrQK_2_0); - gemm(tiledmma_subchunk, tQKrK_2_j, tQKrKt_0_j, tKKrKK_2_0); - - // S2R K/G for operand B (col 1) and element-wise compute - auto tQKrKt_1_j = s2r_compute_subchunk_operandB(_1{}, j, j0, j1, tArAfirst_2_j_kt); - // q_2_j/k_2_j @ k_1_j, accumulate acc_2_1 - gemm(tiledmma_subchunk, tQKrQ_2_j, tQKrKt_1_j, tQKrQK_2_1); - gemm(tiledmma_subchunk, tQKrK_2_j, tQKrKt_1_j, tKKrKK_2_1); - - // S2R K/G for operand B (col 2) and element-wise compute - auto tQKrKt_2_j = s2r_compute_subchunk_operandB(_2{}, j, j0, j1, tArAfirst_2_j_kt); - // q_2_j/k_2_j @ k_2_j, accumulate acc_2_2 - gemm(tiledmma_subchunk, tQKrQ_2_j, tQKrKt_2_j, tQKrQK_2_2); - gemm(tiledmma_subchunk, tQKrK_2_j, tQKrKt_2_j, tKKrKK_2_2); - } - - // S2R beta (maybe resident in register?) - // epilogue: qk^t * scale - cute::transform(tQKrQK_2_0, [scale](auto v) { return v * scale; }); - cute::transform(tQKrQK_2_1, [scale](auto v) { return v * scale; }); - cute::transform(tQKrQK_2_2, [scale](auto v) { return v * scale; }); - // epilogue: kk^t * beta_2 - Tensor sBeta_2 = sBeta_slice(_, 2); - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tQKcMqk_subchunk(i); - auto [s, t] = coord; - auto b = sBeta_2(s); - tKKrKK_2_0(i) *= b; - tKKrKK_2_1(i) *= b; - tKKrKK_2_2(i) *= b; - }); - - // R2S qk_2_0, kk_2_0, wait for current QK/KK free - r2s_subchunk_acc(_2{}, _0{}, tQKrQK_2_0, tKKrKK_2_0); - // R2S qk_2_1, kk_2_1 - r2s_subchunk_acc(_2{}, _1{}, tQKrQK_2_1, tKKrKK_2_1); - // R2S qk_2_2, kk_2_2 - r2s_subchunk_acc(_2{}, _2{}, tQKrQK_2_2, tKKrKK_2_2); - } - }; - - auto qk_and_kk_epi = [&](auto is_final_block_, auto B /*valid seqlen*/) INLINE_LAMBDA { - using CopyOpS2R_Chunk = SM75_U32x4_LDSM_N; - using CopyOpR2S_Chunk = SM90_U32x4_STSM_N; - // S2R QK/KK - auto sQK_curr = sQK(_, _, qk_smem_pipe_write.index()); - auto sKK_inv_curr = sKK_inv(_, _, kk_smem_pipe_write.index()); - Tensor tQKrQK_ref = partition_fragment_C(TiledMmaQK_RS{}, select<0, 1>(TileShapeQK{})); - auto tiled_load_qk = make_tiled_copy_C(Copy_Atom{}, qk_tiled_mma_rs); - auto thr_load_qk = tiled_load_qk.get_thread_slice(thread_idx); - auto tiled_load_kk = make_tiled_copy_C(Copy_Atom{}, qk_tiled_mma_rs); - auto thr_load_kk = tiled_load_kk.get_thread_slice(thread_idx); - - auto tQKrQK_cv = thr_load_qk.retile_D(tQKrQK_ref); - auto tQKrQK = make_fragment_like(tQKrQK_cv); - auto tKKrKK = make_fragment_like(tQKrQK_cv); - copy(tiled_load_qk, thr_load_qk.partition_S(sQK_curr), tQKrQK); - copy(tiled_load_kk, thr_load_kk.partition_S(sKK_inv_curr), tKKrKK); - - // triangular mask and boundary mask - constexpr bool is_final_block = decltype(is_final_block_)::value; - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tQKcMqk(i); - auto [s, t] = coord; - bool pred = s >= t; - tQKrQK(i) = pred ? tQKrQK(i) : Element(0.0f); - tKKrKK(i) = - pred ? tKKrKK(i) : InverseType(0.0f); // diagonal is garbage filled, will process during inversion - if constexpr (is_final_block) { - bool pred = s < B && t < B; - tQKrQK(i) = pred ? tQKrQK(i) : Element(0.0f); - tKKrKK(i) = pred ? tKKrKK(i) : InverseType(0.0f); - } - }); - - // R2S QK/KK - auto tiled_store_qk = make_tiled_copy_C(Copy_Atom{}, qk_tiled_mma_rs); - auto thr_store_qk = tiled_store_qk.get_thread_slice(thread_idx); - auto tiled_store_kk = make_tiled_copy_C(Copy_Atom{}, qk_tiled_mma_rs); - auto thr_store_kk = tiled_store_kk.get_thread_slice(thread_idx); - - copy(tiled_store_qk, thr_store_qk.retile_S(tQKrQK), thr_store_qk.partition_D(sQK_curr)); - copy(tiled_store_kk, thr_store_kk.retile_S(tKKrKK), thr_store_kk.partition_D(sKK_inv_curr)); - }; - - auto compute_aux_loop_body = [&](int blk, auto is_final_block_) INLINE_LAMBDA { - constexpr bool is_final_block = decltype(is_final_block_)::value; - - int B = is_final_block ? valid_seq_len(work_desc, blk) : BlkSeqKV; - - // ====DEBUG===== - // maintain pipeline correctness while removing subchunk - // wait for data ready - // q_pipeline.consumer_wait(q_smem_pipe_read); - // k_pipeline.consumer_wait(k_smem_pipe_read); - // if constexpr (NeedsAlpha) { alpha_pipeline.consumer_wait(alpha_smem_pipe_read); } - // beta_pipeline.consumer_wait(beta_smem_pipe_read); - // cutlass::arch::fence_view_async_shared(); - - // qk_pipeline.producer_acquire(qk_smem_pipe_write); - // kk_pipeline.producer_acquire(kk_smem_pipe_write); - - // SubChunk MMA for QK^T and KK^T for numerical stability - // FIXME: use g_half as anchor in the diagonal subchunk to align with FLA for smaller numerical differences - qk_kk_subchunk_mma_and_store(blk); - // wait for QK/KK ready - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, KdaNamedBarriers::AuxMath); - qk_and_kk_epi(is_final_block_, B); - // =====DEBUG====== - // cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, KdaNamedBarriers::AuxMath); - // if constexpr (is_final_block) { - // if (thread_idx == 0) { - // printf("sQK\n"); - // cute::print_tensor(sQK(_, _, qk_smem_pipe_write.index())); - // printf("sKK_inv\n"); - // cute::print_tensor(sKK_inv(_, _, kk_smem_pipe_write.index())); - // } - // } - - // QK/KK is ready to consume - cutlass::arch::fence_view_async_shared(); - qk_pipeline.producer_commit(qk_smem_pipe_write); - ++qk_smem_pipe_write; - kk_pipeline.producer_commit(kk_smem_pipe_write); - ++kk_smem_pipe_write; - - k_pipeline.consumer_release(k_smem_pipe_read); - ++k_smem_pipe_read; - q_pipeline.consumer_release(q_smem_pipe_read); - ++q_smem_pipe_read; - if constexpr (NeedsAlpha) { - alpha_pipeline.consumer_release(alpha_smem_pipe_read); - ++alpha_smem_pipe_read; - } - - if constexpr (NeedsBeta) { - beta_pipeline.consumer_release(beta_smem_pipe_read); - ++beta_smem_pipe_read; - } - }; - - int32_t num_blocks = ceil_div(work_desc.seq_len, get<0>(TileShape{})); - CUTE_NO_UNROLL - for (int blk = 0; blk < num_blocks - 1; ++blk) { - compute_aux_loop_body(blk, /*is_final_block_=*/cute::false_type{}); - } - compute_aux_loop_body(num_blocks - 1, /*is_final_block_=*/cute::true_type{}); - } - - template - CUTE_DEVICE int - valid_seq_len(WorkDesc work_desc, int blk_idx) { - int remain_len = work_desc.seq_len - BlkSeqKV * blk_idx; - return remain_len <= BlkSeqKV ? remain_len : BlkSeqKV; - } -}; - -} // namespace kda::sm90::collective diff --git a/csrc/kda/sm90/collective/named_barriers.hpp b/csrc/kda/sm90/collective/named_barriers.hpp deleted file mode 100644 index 6f34a010..00000000 --- a/csrc/kda/sm90/collective/named_barriers.hpp +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -namespace kda::sm90::collective { - -struct FlatSharedNamedBarriers { - static constexpr int AllMmaThreadsSync = 0; - static constexpr int AllLdStThreadsSync = 1; - static constexpr int MmaCooperativeStore = 2; - - protected: - static constexpr int NumBarriersUsed = 4; -}; - -} // namespace kda::sm90::collective diff --git a/csrc/kda/sm90/collective/store_tma.hpp b/csrc/kda/sm90/collective/store_tma.hpp deleted file mode 100644 index 0f7f7c1a..00000000 --- a/csrc/kda/sm90/collective/store_tma.hpp +++ /dev/null @@ -1,313 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include - -#include -#include -#include - -#include "kerutils/kerutils.cuh" - -#include "kda/sm90/utils/debug.hpp" - -namespace kda::sm90::collective { - -using ku::alignment_for_swizzle; -using namespace cute; - -/* -NOTE: what we need is as follows - - using DefaultOperation = cutlass::epilogue::fusion::LinearCombination; - using CollectiveStoreO = typename cutlass::epilogue::collective::CollectiveBuilder< - cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, - TileShapeO1, ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto, - ElementAccumulatorO, ElementAccumulatorO, - void, LayoutO, Alignment, // C, not exists - ElementO, decltype(select<1,0,2>(LayoutO{})), Alignment, // D - cutlass::epilogue::TmaWarpSpecializedCooperative, DefaultOperation>::CollectiveOp; - -but unfortunately the required type alias is only useful for our purpose is private so we roll out our own. -*/ - -CUTE_DEVICE uint32_t -smid() { -#ifdef __CUDA_ARCH__ - uint32_t virtual_smid; - asm("mov.u32 %0, %%smid;" : "=r"(virtual_smid)); - return virtual_smid; -#else - return 0; -#endif -} - -template < - typename TileShape_MNK_, - typename ClusterShape, - typename ElementO, - typename ElementAccumulator, - typename SmemElementO, - typename StrideO, - int Stages> -struct CollectiveStoreTma { - static_assert(size_v == 1); - using TileShape_MNK = TileShape_MNK_; - using TileShape_MN = - decltype(select<0, 1>(TileShape_MNK{})); // Collective work on TileShape_MN, it is also the OutputTile - using SizeM = decltype(get<0>(TileShape_MNK{})); // head_size - using SizeN = decltype(get<1>(TileShape_MNK{})); // seqlen - - constexpr static bool is_m_major_O = cutlass::epilogue::collective::detail::is_m_major(); - -#if 0 - // NOTE: the following derived layout is a bit slower than the manual one, will evaluate it later - using SmemLayoutAtom = decltype(cutlass::epilogue::collective::detail::sm90_get_epilogue_smem_swizzle_layout_atom< - StrideO, ElementO, TileShape_MN>()); -#else - static_assert(sizeof(SmemElementO) == 2); - using SmemLayoutAtom = GMMA::Layout_MN_SW32_Atom; -#endif - - using SmemLayoutO = decltype(tile_to_shape( - SmemLayoutAtom{}, - make_shape(SizeM{}, SizeN{}, Int{}), - cute::conditional_t, Step<_1, _2, _3>>{})); - - constexpr static uint32_t TmaTransactionBytes = - (size(take<0, 2>(SmemLayoutO{})) * static_cast(sizeof_bits::value)) / 8; - - using CopyOpR2S = decltype(cutlass::epilogue::collective::detail:: - sm90_get_smem_store_op_for_accumulator()); - using CopyAtomR2S = Copy_Atom; - - using CopyOpS2G = SM90_TMA_STORE; - - using SharedStorage = - cute::array_aligned, alignment_for_swizzle(SmemLayoutO{})>; - using Pipeline = cutlass::PipelineAsync; // NOT PipelineTmaStore! - using PipelineState = cutlass::PipelineState; - - struct Arguments { - ElementO* ptr_O; - StrideO dO; - void* workspace; - }; - - struct Params { - using TMA_O = decltype(make_tma_copy( - CopyOpS2G{}, - make_tensor(make_gmem_ptr(nullptr), repeat_like(StrideO{}, int32_t(0)), StrideO{}), - take<0, 2>(SmemLayoutO{}), - TileShape_MN{}, - _1{})); - - TMA_O tma_store_o; - uint32_t tma_transaction_bytes = TmaTransactionBytes; - void* tensormaps; - }; - using TMA = typename Params::TMA_O; - - CUTE_DEVICE - CollectiveStoreTma(TMA const& tma_store, Pipeline& pipeline, SharedStorage& storage, void* tensormaps) - : tma_store_(tma_store), pipeline_(pipeline), storage_(storage), tensormaps_(tensormaps) { - } - - template - static Params - to_underlying_arguments(ProblemSize const& problem_size, Arguments const& args, void* workspace) { - auto problem_size_MNKL = append<4>(problem_size, 1); - auto [M, N, K, L] = problem_size_MNKL; - - Tensor tensor_o = make_tensor(make_gmem_ptr(args.ptr_O), make_layout(make_shape(M, N, L), args.dO)); - TMA tma_store_o = make_tma_copy_C_sm90(CopyOpS2G{}, tensor_o, take<0, 2>(SmemLayoutO{}), TileShape_MN{}); - - return { - .tma_store_o = tma_store_o, - .tma_transaction_bytes = TmaTransactionBytes, - .tensormaps = workspace, - }; - } - - static size_t - get_workspace_size(/*Arguments const& args,*/ int sm_count) { - // only use additional TMA desc for output tail tiles - size_t num_bytes = sizeof(cute::TmaDescriptor) * sm_count; - DPRINTF("workspace num_bytes:%zu\n", num_bytes); - return num_bytes; - } - - template - static cutlass::Status - initialize_workspace( - ProblemShape const& problem_shape, - /*Arguments const& args,*/ void* workspace, - cudaStream_t stream) { - return cutlass::Status::kSuccess; - } - - CUTE_DEVICE static void - prefetch_tma_descriptors(Params const& params) { - cute::prefetch_tma_descriptor(params.tma_store_o.get_tma_descriptor()); - } - - template - CUTE_DEVICE auto - partition_SD(ProblemSize const& problem_size, TileShape const& tile_shape, WorkDesc const& work_desc) { - constexpr auto BlkSeqQ = decltype(get<0>(tile_shape))::value; - constexpr auto HeadSize = decltype(get<2>(tile_shape))::value; - - Tensor g = [&] { - DPRINTF0_W( - "slice view GMEM O: seq_idx:%d head_idx:%d tok_offset:%lld\n", - work_desc.seq_idx, - work_desc.o_head_idx(), - work_desc.tok_offset); - Tensor m_varlen_head = tma_store_.get_tma_tensor(make_shape( - problem_size.head_size, - problem_size.total_seqlen, - problem_size.num_v_heads)); // O lives in the V/O head space under GVA - Tensor m_varlen = m_varlen_head(_, _, work_desc.o_head_idx()); // slice into current head_idx - Tensor m_offset = domain_offset( - make_coord(_0{}, work_desc.tok_offset), - m_varlen); // offset to start of the current sequence - Tensor g_full = - local_tile(m_offset, make_tile(HeadSize, BlkSeqQ), make_coord(_0{}, _)); // (d, blk, iter_blk) - return g_full; - }(); - Tensor s = make_tensor(make_smem_ptr(storage_.data()), SmemLayoutO{}); - - auto block_tma = tma_store_.get_slice(_0{}); // do not support cluster - return make_tuple(block_tma.partition_S(s), block_tma.partition_D(g)); - } - - template - CUTE_DEVICE static bool - can_process(ProblemSize const& problem_size, WorkDesc const& work_desc, int blk, int num_blocks) { - if (blk < num_blocks - 1) { - // intermediate full tiles, always use TMA - return true; - } else if (work_desc.seq_len % SizeN{} == 0 || work_desc.seq_idx == problem_size.num_seqs - 1) { - // 1. last tile but full, also use TMA - // 2. last tile but last seq, oob can be handled by TMA - return true; - } else { - return false; - } - } - - template - CUTE_DEVICE void - step( - ProblemSize const& problem_size, - WorkDesc const& work_desc, - SrcDst const& src_dst, - PipelineState& src_pipe, - int dst_iter, - int num_iters, - uint32_t lane_predicate) { - auto src = get<0>(src_dst); - auto dst = get<1>(src_dst); - - if (dst_iter == 0) { - bool can_process_tail = can_process(problem_size, work_desc, num_iters - 1, num_iters); - if (!can_process_tail) { - create_tensormap_for_tail(work_desc, lane_predicate); - } - } - - DPRINTF0_WG("pipeline.producer_acquire smem_pipe_read:%d\n", src_pipe.index()); - if constexpr (kAcquireBarrier) { - pipeline_.consumer_wait(src_pipe); - } - - if (can_process(problem_size, work_desc, dst_iter, num_iters)) { - DPRINTF0_W("store src_pipe:%d -> blk:%d\n", src_pipe.index(), dst_iter); - if (lane_predicate == 1) { - copy(tma_store_, src(_, _, _, src_pipe.index()), dst(_, _, _, dst_iter)); - } - } else { - cute::TmaDescriptor* tensormap = acquire_tensormap_for_tail(); - DPRINTF0_W("store tail with tensormap:%p src_pipe:%d -> blk:%d\n", tensormap, src_pipe.index(), dst_iter); - if (lane_predicate == 1) { - copy(tma_store_.with(tensormap), src(_, _, _, src_pipe.index()), dst(_, _, _, dst_iter)); - } - } - - if constexpr (kAcquireBarrier) { - pipeline_.consumer_release(src_pipe); - } - ++src_pipe; - } - - template - CUTE_DEVICE void - create_tensormap_for_tail(WorkDesc const& work_desc, uint32_t lane_predicate) { - namespace ptx = cuda::ptx; - constexpr int num_of_16B = sizeof(cute::TmaDescriptor) / sizeof(uint128_t); - - cute::TmaDescriptor* tensormap = static_cast(tensormaps_) + smid(); - - auto lane_idx = cutlass::canonical_lane_idx(); - if (lane_idx < num_of_16B) { - auto src = reinterpret_cast(tma_store_.get_tma_descriptor()); - auto dst = reinterpret_cast(tensormap); - - dst[lane_idx] = src[lane_idx]; - } - __syncwarp(); - - if (lane_predicate == 1) { - uint32_t new_total_seqlen = work_desc.tok_offset + work_desc.seq_len; - ptx::tensormap_replace_global_dim(ptx::space_global, tensormap, /*ord=*/ptx::n32_t<1>{}, new_total_seqlen); - } - __syncwarp(); - - ptx::fence_proxy_tensormap_generic(ptx::sem_release, ptx::scope_cta); - } - - CUTE_DEVICE cute::TmaDescriptor* - acquire_tensormap_for_tail() { - namespace ptx = cuda::ptx; - cute::TmaDescriptor* tensormap = static_cast(tensormaps_) + smid(); - ptx::fence_proxy_tensormap_generic(ptx::sem_acquire, ptx::scope_cta, tensormap, /*size=*/ptx::n32_t<128>{}); - return tensormap; - } - - private: - TMA const& tma_store_; - Pipeline& pipeline_; - SharedStorage& storage_; - void* tensormaps_; -}; - -} // namespace kda::sm90::collective diff --git a/csrc/kda/sm90/device/device_universal.hpp b/csrc/kda/sm90/device/device_universal.hpp deleted file mode 100644 index 18005ce7..00000000 --- a/csrc/kda/sm90/device/device_universal.hpp +++ /dev/null @@ -1,244 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -// common -#include -#include - -#if !defined(__CUDACC_RTC__) -#include -#include -#endif // !defined(__CUDACC_RTC__) - -namespace cutlass::device { - -template -class Universal { - public: - using Kernel = Kernel_; - - static int const kThreadCount = Kernel::MaxThreadsPerBlock; - - /// Argument structure: User API - using Arguments = typename Kernel::Arguments; - /// Argument structure: Kernel API - using Params = typename Kernel::Params; - - private: - /// Kernel API parameters object - Params params_; - - public: - /// Access the Params structure - Params const& - params() const { - return params_; - } - - /// Determines whether the GEMM can execute the given problem. - static Status - can_implement(Arguments const& args) { - if (Kernel::can_implement(args)) { - return Status::kSuccess; - } else { - return Status::kInvalid; - } - } - - /// Gets the workspace size - static size_t - get_workspace_size(Arguments const& args) { - size_t workspace_bytes = 0; - workspace_bytes += Kernel::get_workspace_size(args); - return workspace_bytes; - } - - /// Computes the grid shape - static dim3 - get_grid_shape(Params const& params) { - return Kernel::get_grid_shape(params); - } - - /// Computes the maximum number of active blocks per multiprocessor - static int - maximum_active_blocks(int /* smem_capacity */ = -1) { - CUTLASS_TRACE_HOST("Universal::maximum_active_blocks()"); - int max_active_blocks = -1; - int smem_size = Kernel::SharedStorageSize; - - // first, account for dynamic smem capacity if needed - cudaError_t result; - if (smem_size >= (48 << 10)) { - CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); - result = - cudaFuncSetAttribute(device_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); - if (cudaSuccess != result) { - result = cudaGetLastError(); // to clear the error bit - CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); - return -1; - } - } - - // query occupancy after setting smem size - result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &max_active_blocks, device_kernel, Kernel::MaxThreadsPerBlock, smem_size); - - if (cudaSuccess != result) { - result = cudaGetLastError(); // to clear the error bit - CUTLASS_TRACE_HOST( - " cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: " << cudaGetErrorString(result)); - return -1; - } - - CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); - return max_active_blocks; - } - - /// Initializes GEMM state from arguments. - Status - initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { - CUTLASS_TRACE_HOST( - "Universal::initialize() - workspace " << workspace << ", stream: " << (stream ? "non-null" : "null")); - - // Initialize the workspace - Status status = Kernel::initialize_workspace(args, workspace, stream); - if (status != Status::kSuccess) { - return status; - } - - // Initialize the Params structure - params_ = Kernel::to_underlying_arguments(args, workspace); - - // account for dynamic smem capacity if needed - int smem_size = Kernel::SharedStorageSize; - if (smem_size >= (48 << 10)) { - CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); - cudaError_t result = - cudaFuncSetAttribute(device_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); - if (cudaSuccess != result) { - result = cudaGetLastError(); // to clear the error bit - CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); - return Status::kErrorInternal; - } - } - - return Status::kSuccess; - } - - /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params. - Status - update(Arguments const& args, void* workspace = nullptr) { - CUTLASS_TRACE_HOST("Universal()::update() - workspace: " << workspace); - - size_t workspace_bytes = get_workspace_size(args); - if (workspace_bytes > 0 && nullptr == workspace) { - return Status::kErrorWorkspaceNull; - } - - params_ = Kernel::to_underlying_arguments(args, workspace); - return Status::kSuccess; - } - - /// Primary run() entry point API that is static allowing users to create and manage their own params. - /// Supplied params struct must be construct by calling Kernel::to_underling_arguments() - static Status - run(Params& params, cudaStream_t stream = nullptr) { - CUTLASS_TRACE_HOST("Universal::run()"); - dim3 const block = Kernel::get_block_shape(); - dim3 const grid = get_grid_shape(params); - - // configure smem size and carveout - int smem_size = Kernel::SharedStorageSize; - - Status launch_result; - // Use extended launch API only for mainloops that use it - if constexpr (Kernel::ArchTag::kMinComputeCapability >= 90) { - dim3 cluster( - cute::size<0>(typename Kernel::ClusterShape{}), - cute::size<1>(typename Kernel::ClusterShape{}), - cute::size<2>(typename Kernel::ClusterShape{})); - void const* kernel = (void const*)device_kernel; - void* kernel_params[] = {¶ms}; - launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params); - } else { - launch_result = Status::kSuccess; - cutlass::arch::synclog_setup(); - device_kernel<<>>(params); - } - - cudaError_t result = cudaGetLastError(); - if (cudaSuccess == result && Status::kSuccess == launch_result) { - return Status::kSuccess; - } else { - CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); - return Status::kErrorInternal; - } - } - - // - // Non-static launch overloads that first create and set the internal params struct of this kernel handle. - // - - /// Launches the kernel after first constructing Params internal state from supplied arguments. - Status - run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { - Status status = initialize(args, workspace, stream); - if (Status::kSuccess == status) { - status = run(params_, stream); - } - return status; - } - - /// Launches the kernel after first constructing Params internal state from supplied arguments. - Status - operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { - return run(args, workspace, stream); - } - - /// Overload that allows a user to re-launch the same kernel without updating internal params struct. - Status - run(cudaStream_t stream = nullptr) { - return run(params_, stream); - } - - /// Overload that allows a user to re-launch the same kernel without updating internal params struct. - Status - operator()(cudaStream_t stream = nullptr) { - return run(params_, stream); - } -}; - -//////////////////////////////////////////////////////////////////////////////// - -} // namespace cutlass::device - -//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/kda/sm90/kda_fwd_sm90.cu b/csrc/kda/sm90/kda_fwd_sm90.cu deleted file mode 100644 index ed855db6..00000000 --- a/csrc/kda/sm90/kda_fwd_sm90.cu +++ /dev/null @@ -1,189 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -// Dispatch function only — does NOT include the .cuh to avoid -// implicit instantiation of all kernel variants in one TU. -// Each SafeGate variant is explicitly instantiated in its own .cu file. - -#include -#include - -namespace kda::sm90 { - -using namespace cute; - -// Forward declaration of the per-variant launcher (defined in .cuh, instantiated in separate TUs) -template < - bool NeedsBeta, - bool NeedsAlpha, - bool InitStateFromInput, - bool SafeGate, - typename ArchTag, - typename TO, - typename TQKV, - typename TState, - typename TBeta = float> -void -launch_kda_fwd_prefill_kernel_gbai( - cudaStream_t stream, - TO* output, - TState* output_state, - TQKV const* q, - TQKV const* k, - TQKV const* v, - TState const* input_state, - float const* alpha, - TBeta const* beta, - int32_t const* cu_seqlens, - uint8_t* workspace_buffer, - int32_t num_seqs, - int32_t num_qk_heads, - int32_t num_v_heads, - int32_t head_size, - int64_t total_seqlen, - float scale, - int32_t sm_count, - int32_t const* cp_seq_map, - int32_t const* raw_cu_seqlens, - int32_t raw_num_seqs); - -template < - typename ArchTag, // TODO: hide this - typename TO, - typename TQKV, - typename TState, - typename TBeta> -void -launch_kda_fwd_prefill_kernel( - cudaStream_t stream, - TO* output, - TState* output_state, - TQKV const* q, - TQKV const* k, - TQKV const* v, - TState const* input_state, - float const* alpha, - TBeta const* beta, - int32_t const* cu_seqlens, - uint8_t* workspace_buffer, - int32_t num_seqs, - int32_t num_qk_heads, - int32_t num_v_heads, - int32_t head_size, - int64_t total_seqlen, - float scale, - bool safe_gate, - int32_t sm_count, - int32_t const* cp_seq_map, - int32_t const* raw_cu_seqlens, - int32_t raw_num_seqs) { - bool needs_beta = beta != nullptr; - bool needs_alpha = alpha != nullptr; - bool init_state = input_state != nullptr; - -#define LAUNCH(needs_beta, needs_alpha, init_state, safe_gate) \ - launch_kda_fwd_prefill_kernel_gbai( \ - stream, \ - output, \ - output_state, \ - q, \ - k, \ - v, \ - input_state, \ - alpha, \ - beta, \ - cu_seqlens, \ - workspace_buffer, \ - num_seqs, \ - num_qk_heads, \ - num_v_heads, \ - head_size, \ - total_seqlen, \ - scale, \ - sm_count, \ - cp_seq_map, \ - raw_cu_seqlens, \ - raw_num_seqs); - if (init_state) { - if (needs_beta && needs_alpha && safe_gate) { - LAUNCH(true, true, true, true); - } else { - throw std::runtime_error("unreachable"); - } - } else { - if (needs_beta && needs_alpha && safe_gate) { - LAUNCH(true, true, false, true); - } else { - throw std::runtime_error("unreachable"); - } - } - -#undef LAUNCH -} - -using bf16 = cute::bfloat16_t; - -// TBeta=float (default) -template void -launch_kda_fwd_prefill_kernel( - cudaStream_t stream, - bf16* output, - float* state, - bf16 const* q, - bf16 const* k, - bf16 const* v, - float const* input_state, - float const* alpha, - float const* beta, - int32_t const* cu_seqlens, - uint8_t* workspace_buffer, - int32_t num_seqs, - int32_t num_qk_heads, - int32_t num_v_heads, - int32_t head_size, - int64_t total_seqlen, - float scale, - bool safe_gate, - int32_t sm_count, - int32_t const* cp_seq_map, - int32_t const* raw_cu_seqlens, - int32_t raw_num_seqs); - -// TBeta=bf16 -template void -launch_kda_fwd_prefill_kernel( - cudaStream_t stream, - bf16* output, - float* state, - bf16 const* q, - bf16 const* k, - bf16 const* v, - float const* input_state, - float const* alpha, - bf16 const* beta, - int32_t const* cu_seqlens, - uint8_t* workspace_buffer, - int32_t num_seqs, - int32_t num_qk_heads, - int32_t num_v_heads, - int32_t head_size, - int64_t total_seqlen, - float scale, - bool safe_gate, - int32_t sm_count, - int32_t const* cp_seq_map, - int32_t const* raw_cu_seqlens, - int32_t raw_num_seqs); - -} // namespace kda::sm90 diff --git a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu deleted file mode 100644 index 0da2986e..00000000 --- a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu +++ /dev/null @@ -1,66 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#include -#include - -#include "kda/sm90/prefill_kernel_kda_fwd_sm90.cuh" -#include "kda/sm90/utils/common.hpp" - -namespace kda::sm90 { - -using namespace cute; -using bf16 = cute::bfloat16_t; - -#define INSTANTIATE_GBAI(NeedsBeta, NeedsAlpha, InitState, SafeGate, TBeta) \ - template void launch_kda_fwd_prefill_kernel_gbai< \ - NeedsBeta, \ - NeedsAlpha, \ - InitState, \ - SafeGate, \ - cutlass::arch::Sm90, \ - bf16, \ - bf16, \ - float, \ - TBeta>( \ - cudaStream_t, \ - bf16*, \ - float*, \ - bf16 const*, \ - bf16 const*, \ - bf16 const*, \ - float const*, \ - float const*, \ - TBeta const*, \ - int32_t const*, \ - uint8_t*, \ - int32_t, \ - int32_t, \ - int32_t, \ - int32_t, \ - int64_t, \ - float, \ - int32_t, \ - int32_t const*, \ - int32_t const*, \ - int32_t) - -INSTANTIATE_GBAI(true, true, false, true, float); -INSTANTIATE_GBAI(true, true, true, true, float); -INSTANTIATE_GBAI(true, true, false, true, bf16); -INSTANTIATE_GBAI(true, true, true, true, bf16); - -#undef INSTANTIATE_GBAI - -} // namespace kda::sm90 diff --git a/csrc/kda/sm90/kernel/builder_kda_fwd.hpp b/csrc/kda/sm90/kernel/builder_kda_fwd.hpp deleted file mode 100644 index 74e9c43a..00000000 --- a/csrc/kda/sm90/kernel/builder_kda_fwd.hpp +++ /dev/null @@ -1,80 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include "kda/sm90/collective/mainloop_kda_fwd.hpp" -#include "kda/sm90/kernel/kernel_kda_fwd.hpp" -#include "kda/sm90/kernel/options.hpp" -#include "kda/sm90/kernel/tile_scheduler.hpp" -#include "kda/sm90/utils/type_traits.hpp" - -namespace kda::sm90::kernel { - -template < - class Element_, - class ElementAccumulatorQK_, - class ElementAccumulatorPV_, - class TileShape_, // BlkSeqQO, BlkSeqKV, HeadSize - class LayoutQ_, - class LayoutK_, - class LayoutV_, - class LayoutO_, - class DispatchPolicy, - class Options = DefaultOptions> -struct FlatBuilderKdaFwd; - -template < - class Element, - class ElementAccumulatorQK, - class ElementAccumulatorPV, - class TileShape, // BlkSeqQO, BlkSeqKV, HeadSize - class LayoutQ, - class LayoutK, - class LayoutV, - class LayoutO, - class Options> -struct FlatBuilderKdaFwd< - Element, - ElementAccumulatorQK, - ElementAccumulatorPV, - TileShape, - LayoutQ, - LayoutK, - LayoutV, - LayoutO, - cutlass::gemm::KernelTmaWarpSpecializedCooperative, - Options> { - using CollectiveMainloop = kda::sm90::collective::FlatMainloopTmaWarpSpecializedKdaFwd< - Element, - ElementAccumulatorQK, - ElementAccumulatorPV, - TileShape, - LayoutQ, - LayoutK, - LayoutV, - LayoutO, - Options>; - - static constexpr bool kIsPersistent = find_option_t::value; - static_assert(!kIsPersistent, "not implemented"); - - using TileScheduler = kda::sm90::kernel::IndividualTileScheduler; - // using TileScheduler = std::conditional_t; - - using Kernel = kda::sm90::kernel::FlatKernelTmaWarpSpecializedKdaFwd; -}; - -} // namespace kda::sm90::kernel diff --git a/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp b/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp deleted file mode 100644 index ac597f7b..00000000 --- a/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp +++ /dev/null @@ -1,639 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include -#include -#include -#include - -#include "kda/sm90/kernel/options.hpp" -#include "kda/sm90/utils/common.hpp" -#include "kda/sm90/utils/unused.hpp" - -namespace kda::sm90::kernel { - -using namespace cute; - -template -constexpr T1 -round_down(T1 a, T2 b) { - return (a / b) * b; -} - -constexpr std::tuple -get_register_requirements( - uint32_t max_threads_per_block, - uint32_t min_blocks_per_multiprocessor, - uint32_t num_state_mma_warp_groups // state related mma -) { - uint32_t reg_alloc_granularity = 8; - -#if !defined(FLAT_DEBUG_PRINT) || !FLAT_DEBUG_PRINT - uint32_t load_registers = 40 - 2 * reg_alloc_granularity; -#else - uint32_t load_registers = 40; -#endif - // TODO: better tuning - uint32_t total_aux_load_budget = 176; - uint32_t aux_registers = total_aux_load_budget - load_registers; // (24 + X) or (40 + X) - - uint32_t total_registers = - round_down(64 * 1024 / min_blocks_per_multiprocessor, max_threads_per_block * reg_alloc_granularity) / - cutlass::NumThreadsPerWarpGroup; - uint32_t mma_registers = round_down( - (total_registers - load_registers - aux_registers) / num_state_mma_warp_groups, reg_alloc_granularity); - - // max reg is 255, 248 round to multiple of reg_alloc_granularity; - return {cute::min(248, load_registers), cute::min(248, mma_registers), cute::min(248, aux_registers)}; -} - -template -struct FlatKernelTmaWarpSpecializedKdaFwd { - using ArchTag = cutlass::arch::Sm90; - - static const int NumLoadWarpGroups = 1; - static constexpr int NumStateMmaWarpGroups = CollectiveMainloop::NumStateMmaWarpGroups; - static constexpr int NumAuxMmaWarpGroups = CollectiveMainloop::NumAuxMmaWarpGroups; - - static constexpr int NeedsAlpha = CollectiveMainloop::NeedsAlpha; - static constexpr int NeedsBeta = CollectiveMainloop::NeedsBeta; - static constexpr int SafeGate = CollectiveMainloop::SafeGate; - - using TileShape = typename CollectiveMainloop::TileShape; - using ClusterShape = typename CollectiveMainloop::ClusterShape; - - using MainloopQPipeline = typename CollectiveMainloop::MainloopQPipeline; - using MainloopKPipeline = typename CollectiveMainloop::MainloopKPipeline; - using MainloopVPipeline = typename CollectiveMainloop::MainloopVPipeline; - using MainloopOPipeline = typename CollectiveMainloop::MainloopOPipeline; - - using MainloopQKPipeline = typename CollectiveMainloop::MainloopQKPipeline; - using MainloopKKPipeline = typename CollectiveMainloop::MainloopKKPipeline; - - using MainloopAlphaLastPipeline = typename CollectiveMainloop::MainloopAlphaLastPipeline; - - using MainloopAlphaPipeline = typename CollectiveMainloop::MainloopAlphaPipeline; - using MainloopBetaPipeline = typename CollectiveMainloop::MainloopBetaPipeline; - - using OrderedMathBarriers = typename CollectiveMainloop::OrderedMathBarriers; - - static constexpr uint32_t StagesPerMathWarpGroup = 2; - - using MathWarpGroupOrderBarrier = cutlass::OrderedSequenceBarrier; - - struct TensorStorage { - typename CollectiveMainloop::SharedStorage mainloop; - }; - - struct SharedStorage { - TensorStorage tensors; - - using QPipelineStorage = typename MainloopQPipeline::SharedStorage; - using KPipelineStorage = typename MainloopKPipeline::SharedStorage; - using VPipelineStorage = typename MainloopVPipeline::SharedStorage; - using OPipelineStorage = typename MainloopOPipeline::SharedStorage; - - alignas(16) QPipelineStorage q_pipeline_storage; - alignas(16) KPipelineStorage k_pipeline_storage; - alignas(16) VPipelineStorage v_pipeline_storage; - alignas(16) OPipelineStorage o_pipeline_storage; - - using QKPipelineStorage = typename MainloopQKPipeline::SharedStorage; - using KKPipelineStorage = typename MainloopKKPipeline::SharedStorage; - - alignas(16) QKPipelineStorage qk_pipeline_storage; - alignas(16) KKPipelineStorage kk_pipeline_storage; - - using AlphaLastPipelineStorage = typename MainloopAlphaLastPipeline::SharedStorage; - - alignas(16) AlphaLastPipelineStorage alpha_last_pipeline_storage; - - using AlphaPipelineStorage = typename MainloopAlphaPipeline::SharedStorage; - using BetaPipelineStorage = typename MainloopBetaPipeline::SharedStorage; - alignas(16) AlphaPipelineStorage alpha_pipeline_storage; - alignas(16) BetaPipelineStorage beta_pipeline_storage; - - alignas(16) cutlass::arch::ClusterBarrier load_warp_barrier; - }; - - static constexpr int SharedStorageSize = sizeof(SharedStorage); - - struct VarlenProblemShape { - int32_t const* cu_seqlens; - int64_t total_seqlen; - int32_t num_seqs; - // GVA: Q/K share num_qk_heads; V, g, beta, O, state are sized by num_v_heads. - // Constraint: num_v_heads % num_qk_heads == 0; heads_per_group = num_v_heads / num_qk_heads. - int32_t num_qk_heads; - int32_t num_v_heads; - int32_t head_size; // d - - // For intra-card CP - int32_t const* cp_seq_map = nullptr; - int32_t const* raw_cu_seqlens = nullptr; - int32_t raw_num_seqs = 0; - }; - using ProblemShape = VarlenProblemShape; - - struct Arguments { - ProblemShape problem_size; - typename CollectiveMainloop::Arguments mainloop; - cutlass::KernelHardwareInfo hw_info; - }; - - struct Params { - ProblemShape problem_size; - typename CollectiveMainloop::Params mainloop; - typename TileScheduler::Params scheduler; - }; - - using QPipelineParams = typename MainloopQPipeline::Params; - using QPipelineState = typename cutlass::PipelineState; - - using KPipelineParams = typename MainloopKPipeline::Params; - using KPipelineState = typename cutlass::PipelineState; - - using VPipelineParams = typename MainloopVPipeline::Params; - using VPipelineState = typename cutlass::PipelineState; - - using OPipelineParams = typename MainloopOPipeline::Params; - using OPipelineState = typename cutlass::PipelineState; - - using QKPipelineParams = typename MainloopQKPipeline::Params; - using QKPipelineState = typename cutlass::PipelineState; - - using KKPipelineParams = typename MainloopKKPipeline::Params; - using KKPipelineState = typename cutlass::PipelineState; - - using AlphaLastPipelineParams = std::conditional_t; - using AlphaLastPipelineState = - std::conditional_t, Unused>; - - using AlphaPipelineParams = std::conditional_t; - using AlphaPipelineState = - std::conditional_t, Unused>; - - using BetaPipelineParams = std::conditional_t; - using BetaPipelineState = - std::conditional_t, Unused>; - - static constexpr int MinBlocksPerMultiprocessor = 1; - static constexpr int MaxThreadsPerBlock = - (NumLoadWarpGroups + NumStateMmaWarpGroups + NumAuxMmaWarpGroups) * cutlass::NumThreadsPerWarpGroup; - - static constexpr auto RegisterRequirements = - get_register_requirements(MaxThreadsPerBlock, MinBlocksPerMultiprocessor, NumStateMmaWarpGroups); - static constexpr uint32_t LdStRegisterRequirement = get<0>(RegisterRequirements); - static constexpr uint32_t StateMmaRegisterRequirement = get<1>(RegisterRequirements); - static constexpr uint32_t AuxMmaRegisterRequirement = get<2>(RegisterRequirements); - - static size_t - get_workspace_size(Arguments const& args) { - return CollectiveMainloop::get_workspace_size(args.mainloop, args.hw_info.sm_count); - } - - static cutlass::Status - initialize_workspace(Arguments const& args, void* workspace, cudaStream_t stream) { - return CollectiveMainloop::initialize_workspace(args.problem_size, args.mainloop, workspace, stream); - } - - static bool - can_implement(Arguments const& args) { - return CollectiveMainloop::can_implement(args.problem_size, args.mainloop); - } - - static dim3 - get_grid_shape(Params const& params) { - return TileScheduler::get_grid_shape(params.scheduler); - } - - static dim3 - get_block_shape() { - dim3 block(MaxThreadsPerBlock, 1, 1); - return block; - } - - static Params - to_underlying_arguments(Arguments const& args, void* workspace) { - return Params{ - args.problem_size, - CollectiveMainloop::to_underlying_arguments(args.problem_size, args.mainloop, workspace), - TileScheduler::to_underlying_arguments(args.problem_size, args.hw_info, ClusterShape{}, TileShape{})}; - } - - CUTE_DEVICE void - operator()(const Params& params, char* smem) { - enum class WarpGroupRole { - LdSt = 0, - Math0 = 1, - Math1 = 2, - MathA = 3, // auxiliary math WG - }; - - // NOTE: CollectiveInverse will have more utilization on warp 0&1 - // so we put beta and alpha preprocessing on warp 2&3 - enum class LdStWarpRole { - LoadQKV = 0, - StoreO = 1, - LoadBeta = 2, - LoadAlpha = 3, - }; - - TileScheduler scheduler{params.scheduler}; - - // Shared memory. - auto& storage = *reinterpret_cast(smem); - - int lane_idx = cutlass::canonical_lane_idx(); - int warp_idx = cutlass::canonical_warp_idx_sync(); - int warp_idx_in_wg = warp_idx % cutlass::NumWarpsPerWarpGroup; - int warp_group_idx = cutlass::canonical_warp_group_idx(); - auto warp_group_role = WarpGroupRole(warp_group_idx); - auto ldst_warp_role = LdStWarpRole(warp_idx_in_wg); - - int lane_predicate = cute::elect_one_sync(); - uint32_t block_rank_in_cluster = cute::block_rank_in_cluster(); - - // Issue Tma Descriptor Prefetch from a single thread - if ((warp_idx == 0) && lane_predicate) { - CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); - } - - constexpr int NumStateMathThreads = NumStateMmaWarpGroups * cutlass::NumThreadsPerWarpGroup; - constexpr int NumAuxMathThreads = NumAuxMmaWarpGroups * cutlass::NumThreadsPerWarpGroup; - - QPipelineParams q_pipeline_params; - q_pipeline_params.transaction_bytes = CollectiveMainloop::LoadQBytes; - q_pipeline_params.is_leader = lane_predicate && (ldst_warp_role == LdStWarpRole::LoadQKV); - q_pipeline_params.num_consumers = NumStateMathThreads + NumAuxMathThreads; - - KPipelineParams k_pipeline_params; - k_pipeline_params.transaction_bytes = CollectiveMainloop::LoadKBytes; - k_pipeline_params.is_leader = lane_predicate && (ldst_warp_role == LdStWarpRole::LoadQKV); - k_pipeline_params.num_consumers = NumStateMathThreads + NumAuxMathThreads; - - VPipelineParams v_pipeline_params; - v_pipeline_params.transaction_bytes = CollectiveMainloop::LoadVBytes; - v_pipeline_params.is_leader = lane_predicate && (ldst_warp_role == LdStWarpRole::LoadQKV); - v_pipeline_params.num_consumers = NumStateMathThreads; - - AlphaPipelineParams alpha_pipeline_params; - if constexpr (NeedsAlpha) { - alpha_pipeline_params.transaction_bytes = CollectiveMainloop::LoadAlphaBytes; - alpha_pipeline_params.is_leader = lane_predicate && (ldst_warp_role == LdStWarpRole::LoadQKV); - alpha_pipeline_params.num_consumers = NumStateMathThreads + NumAuxMathThreads + cutlass::NumThreadsPerWarp; - } - - OPipelineParams o_pipeline_params; - o_pipeline_params.producer_arv_count = NumStateMathThreads; - o_pipeline_params.consumer_arv_count = cutlass::NumThreadsPerWarp; - - QKPipelineParams qk_pipeline_params; - qk_pipeline_params.producer_arv_count = NumAuxMathThreads; - qk_pipeline_params.consumer_arv_count = NumStateMathThreads; - - KKPipelineParams kk_pipeline_params; - kk_pipeline_params.producer_arv_count = NumAuxMathThreads; - kk_pipeline_params.consumer_arv_count = NumStateMathThreads; - - AlphaLastPipelineParams alpha_last_pipeline_params; - if constexpr (NeedsAlpha) { - alpha_last_pipeline_params.producer_arv_count = cutlass::NumThreadsPerWarp; - alpha_last_pipeline_params.consumer_arv_count = NumStateMathThreads; - } - - BetaPipelineParams beta_pipeline_params; - if constexpr (NeedsBeta) { - beta_pipeline_params.producer_arv_count = cutlass::NumThreadsPerWarp; - beta_pipeline_params.consumer_arv_count = NumAuxMathThreads + NumStateMathThreads; - } - - OrderedMathBarriers math_barriers; - - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadQKV) { - DPRINTF0_W("ldst_warp_role: LoadQKV Alpha\n"); - q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Producer; - k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Producer; - v_pipeline_params.role = MainloopVPipeline::ThreadCategory::Producer; - if constexpr (NeedsAlpha) { - alpha_pipeline_params.role = MainloopAlphaPipeline::ThreadCategory::Producer; - } - } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::StoreO) { - DPRINTF0_W("ldst_warp_role: StoreO\n"); - o_pipeline_params.role = MainloopOPipeline::ThreadCategory::Consumer; - } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadBeta) { - if constexpr (NeedsBeta) { - beta_pipeline_params.role = MainloopBetaPipeline::ThreadCategory::Producer; - } - } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadAlpha) { - // LoadAlpha warp consumes alpha_pipeline (reads last row) and produces alpha_last_pipeline - if constexpr (NeedsAlpha) { - alpha_pipeline_params.role = MainloopAlphaPipeline::ThreadCategory::Consumer; - } - alpha_last_pipeline_params.role = MainloopAlphaLastPipeline::ThreadCategory::Producer; - } - if (warp_group_role == WarpGroupRole::Math0 || warp_group_role == WarpGroupRole::Math1) { - DPRINTF0_WG("warp_group_role: MathX\n"); - q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Consumer; - k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Consumer; - v_pipeline_params.role = MainloopVPipeline::ThreadCategory::Consumer; - o_pipeline_params.role = MainloopOPipeline::ThreadCategory::Producer; - - qk_pipeline_params.role = MainloopQKPipeline::ThreadCategory::Consumer; - kk_pipeline_params.role = MainloopKKPipeline::ThreadCategory::Consumer; - - if constexpr (NeedsAlpha) { - alpha_pipeline_params.role = MainloopAlphaPipeline::ThreadCategory::Consumer; - alpha_last_pipeline_params.role = MainloopAlphaLastPipeline::ThreadCategory::Consumer; - } - if constexpr (NeedsBeta) { - beta_pipeline_params.role = MainloopBetaPipeline::ThreadCategory::Consumer; - } - - math_barriers.init(warp_group_idx - 1); - } - if (warp_group_role == WarpGroupRole::MathA) { - DPRINTF0_WG("warp_group_role: MathA\n"); - q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Consumer; - k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Consumer; - - qk_pipeline_params.role = MainloopQKPipeline::ThreadCategory::Producer; - kk_pipeline_params.role = MainloopKKPipeline::ThreadCategory::Producer; - - if constexpr (NeedsAlpha) { - alpha_pipeline_params.role = MainloopAlphaPipeline::ThreadCategory::Consumer; - } - if constexpr (NeedsBeta) { - beta_pipeline_params.role = MainloopBetaPipeline::ThreadCategory::Consumer; - } - } - - MainloopQPipeline q_pipeline(storage.q_pipeline_storage, q_pipeline_params, ClusterShape{}); - MainloopKPipeline k_pipeline(storage.k_pipeline_storage, k_pipeline_params, ClusterShape{}); - MainloopVPipeline v_pipeline(storage.v_pipeline_storage, v_pipeline_params, ClusterShape{}); - MainloopAlphaPipeline alpha_pipeline(storage.alpha_pipeline_storage, alpha_pipeline_params, ClusterShape{}); - MainloopOPipeline o_pipeline(storage.o_pipeline_storage, o_pipeline_params, /*InitBarriers=*/cute::true_type{}); - - MainloopAlphaLastPipeline alpha_last_pipeline( - storage.alpha_last_pipeline_storage, - alpha_last_pipeline_params, - /*InitBarriers=*/cute::true_type{}); - - MainloopQKPipeline qk_pipeline( - storage.qk_pipeline_storage, - qk_pipeline_params, - /*InitBarriers=*/cute::true_type{}); - MainloopKKPipeline kk_pipeline( - storage.kk_pipeline_storage, - kk_pipeline_params, - /*InitBarriers=*/cute::true_type{}); - - // MainloopAlphaPipeline alpha_pipeline(storage.alpha_pipeline_storage, alpha_pipeline_params, - // /*InitBarriers=*/cute::true_type{}); - MainloopBetaPipeline beta_pipeline( - storage.beta_pipeline_storage, - beta_pipeline_params, - /*InitBarriers=*/cute::true_type{}); - - QPipelineState q_smem_pipe_read; - QPipelineState q_smem_pipe_write = cutlass::make_producer_start_state(); - KPipelineState k_smem_pipe_read; - KPipelineState k_smem_pipe_write = cutlass::make_producer_start_state(); - VPipelineState v_smem_pipe_read; - VPipelineState v_smem_pipe_write = cutlass::make_producer_start_state(); - OPipelineState o_smem_pipe_read; - OPipelineState o_smem_pipe_write = cutlass::make_producer_start_state(); - - AlphaLastPipelineState alpha_last_smem_pipe_read; - AlphaLastPipelineState alpha_last_smem_pipe_write; - if constexpr (NeedsAlpha) { - alpha_last_smem_pipe_write = cutlass::make_producer_start_state(); - } - - QKPipelineState qk_smem_pipe_read; - QKPipelineState qk_smem_pipe_write = cutlass::make_producer_start_state(); - KKPipelineState kk_smem_pipe_read; - KKPipelineState kk_smem_pipe_write = cutlass::make_producer_start_state(); - - AlphaPipelineState alpha_smem_pipe_read; - AlphaPipelineState alpha_smem_pipe_write; - if constexpr (NeedsAlpha) { - alpha_smem_pipe_write = cutlass::make_producer_start_state(); - } - BetaPipelineState beta_smem_pipe_read; - BetaPipelineState beta_smem_pipe_write; - if constexpr (NeedsBeta) { - beta_smem_pipe_write = cutlass::make_producer_start_state(); - } - - // barrier sm or cluster level for initialization - if constexpr (size(ClusterShape{}) > 1) { - cute::cluster_arrive_relaxed(); - cute::cluster_wait(); - } else { - __syncthreads(); - } - DPRINTF0_WG("warpspecialized grid initialized\n"); - - CollectiveMainloop collective_mainloop; - - if (warp_group_role == WarpGroupRole::LdSt) { - DPRINTF0_WG("LsSt warp_group_idx:%d, RegisterRequirement:%d\n", warp_group_idx, LdStRegisterRequirement); - cutlass::arch::warpgroup_reg_dealloc(); - if (ldst_warp_role == LdStWarpRole::LoadQKV) { - auto work_desc = scheduler.get_next_work(params.scheduler, params.problem_size); - CUTE_NO_UNROLL - for (; work_desc.is_valid(params.scheduler); - work_desc = scheduler.get_next_work(params.scheduler, params.problem_size)) { - DPRINTF0_WG( - "LsSt working on LoadQ/K/V, seq_idx:%d, q/k/v_head_idx:(%d,%d,%d), seq_len:%lld)\n", - work_desc.seq_idx, - work_desc.q_head_idx(), - work_desc.k_head_idx(), - work_desc.v_head_idx(), - work_desc.seq_len); - auto tile_shape = typename CollectiveMainloop::TileShape{}; - collective_mainloop.load_qkv( - params.mainloop, - params.problem_size, - tile_shape, - work_desc, - q_pipeline, - q_smem_pipe_write, - k_pipeline, - k_smem_pipe_write, - v_pipeline, - v_smem_pipe_write, - alpha_pipeline, - alpha_smem_pipe_write, - storage.tensors.mainloop); - } - } else if (ldst_warp_role == LdStWarpRole::LoadBeta) { - if constexpr (NeedsBeta) { - auto work_desc = scheduler.get_next_work(params.scheduler, params.problem_size); - CUTE_NO_UNROLL - for (; work_desc.is_valid(params.scheduler); - work_desc = scheduler.get_next_work(params.scheduler, params.problem_size)) { - DPRINTF0_WG( - "LsSt working on LoadBeta, seq_idx:%d, sab_head_idx:%d, seq_len:%lld)\n", - work_desc.seq_idx, - work_desc.o_head_idx(), - work_desc.seq_len); - auto tile_shape = typename CollectiveMainloop::TileShape{}; - collective_mainloop.load_beta( - params.mainloop, - params.problem_size, - tile_shape, - work_desc, - beta_pipeline, - beta_smem_pipe_write, - storage.tensors.mainloop); - } - } - } else if (ldst_warp_role == LdStWarpRole::LoadAlpha) { - // produce the last row of Alpha - if constexpr (NeedsAlpha) { - auto work_desc = scheduler.get_next_work(params.scheduler, params.problem_size); - CUTE_NO_UNROLL - for (; work_desc.is_valid(params.scheduler); - work_desc = scheduler.get_next_work(params.scheduler, params.problem_size)) { - DPRINTF0_WG( - "LsSt working on LoadAlpha+ExtractLast, seq_idx:%d, sab_head_idx:%d, seq_len:%lld)\n", - work_desc.seq_idx, - work_desc.o_head_idx(), - work_desc.seq_len); - auto tile_shape = typename CollectiveMainloop::TileShape{}; - collective_mainloop.extract_alpha_last( - params.mainloop, - params.problem_size, - tile_shape, - work_desc, - alpha_pipeline, - alpha_smem_pipe_read, - alpha_last_pipeline, - alpha_last_smem_pipe_write, - storage.tensors.mainloop); - } - } - } else if (ldst_warp_role == LdStWarpRole::StoreO) { - auto work_desc = scheduler.get_next_work(params.scheduler, params.problem_size); - DPRINTF0_WG( - "LsSt working on StoreO, seq_idx:%d, o_head_idx:%d, seq_len:%lld)\n", - work_desc.seq_idx, - work_desc.o_head_idx(), - work_desc.seq_len); - auto tile_shape = typename CollectiveMainloop::TileShape{}; - collective_mainloop.store( - params.mainloop.tma_store_o, - params.mainloop.tensormaps, - params.problem_size, - tile_shape, - work_desc, - o_pipeline, - o_smem_pipe_read, - storage.tensors.mainloop.smem_o); - } - } else if (warp_group_role == WarpGroupRole::Math0 || warp_group_role == WarpGroupRole::Math1) { - DPRINTF0_WG( - "Compute[state]: warp_group_idx:%d, RegisterRequirement:%d\n", - warp_group_idx, - StateMmaRegisterRequirement); - cutlass::arch::warpgroup_reg_alloc(); - auto work_desc = scheduler.get_next_work(params.scheduler, params.problem_size); - CUTE_NO_UNROLL - for (; work_desc.is_valid(params.scheduler); - work_desc = scheduler.get_next_work(params.scheduler, params.problem_size)) { - DPRINTF0_WG( - "Compute[state]: seq_idx:%d, qk/v/o_head_idx:(%d,%d,%d,%d), seq_len:%lld)\n", - work_desc.seq_idx, - work_desc.q_head_idx(), - work_desc.k_head_idx(), - work_desc.v_head_idx(), - work_desc.o_head_idx(), - work_desc.seq_len); - collective_mainloop.compute( - params.mainloop, - params.problem_size, - work_desc, - q_pipeline, - q_smem_pipe_read, - k_pipeline, - k_smem_pipe_read, - v_pipeline, - v_smem_pipe_read, - o_pipeline, - o_smem_pipe_write, - qk_pipeline, - qk_smem_pipe_read, - kk_pipeline, - kk_smem_pipe_read, - alpha_pipeline, - alpha_smem_pipe_read, - beta_pipeline, - beta_smem_pipe_read, - alpha_last_pipeline, - alpha_last_smem_pipe_read, - math_barriers, - storage.tensors.mainloop); - } - } else if (warp_group_role == WarpGroupRole::MathA) { - DPRINTF0_WG( - "Compute[aux]: warp_group_idx:%d, RegisterRequirement:%d\n", warp_group_idx, AuxMmaRegisterRequirement); - cutlass::arch::warpgroup_reg_alloc(); - auto work_desc = scheduler.get_next_work(params.scheduler, params.problem_size); - CUTE_NO_UNROLL - for (; work_desc.is_valid(params.scheduler); - work_desc = scheduler.get_next_work(params.scheduler, params.problem_size)) { - DPRINTF0_WG( - "Compute[aux]: seq_idx:%d, qk/v/o_head_idx:(%d,%d,%d,%d), seq_len:%lld)\n", - work_desc.seq_idx, - work_desc.q_head_idx(), - work_desc.k_head_idx(), - work_desc.v_head_idx(), - work_desc.o_head_idx(), - work_desc.seq_len); - collective_mainloop.compute_aux_safe( - params.mainloop, - params.problem_size, - work_desc, - q_pipeline, - q_smem_pipe_read, - k_pipeline, - k_smem_pipe_read, - qk_pipeline, - qk_smem_pipe_write, - kk_pipeline, - kk_smem_pipe_write, - alpha_pipeline, - alpha_smem_pipe_read, - beta_pipeline, - beta_smem_pipe_read, - alpha_last_pipeline, - alpha_last_smem_pipe_write, - storage.tensors.mainloop); - } - } else { - DPRINTF0_WG("Unknown warp role, warp_group_idx:%d\n", warp_group_idx); - } - - __syncthreads(); - } -}; - -} // namespace kda::sm90::kernel diff --git a/csrc/kda/sm90/kernel/options.hpp b/csrc/kda/sm90/kernel/options.hpp deleted file mode 100644 index e25fe9d4..00000000 --- a/csrc/kda/sm90/kernel/options.hpp +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include - -#include - -namespace kda::sm90::kernel { - -template -struct Option { - static constexpr auto tag = kTag; - using option_value = Value; -}; - -using DefaultOptions = std::tuple<>; - -namespace detail { - -template -struct find_option_impl; - -template -struct find_option_impl { - using option_value = Default; -}; - -template -struct find_option_impl : find_option_impl {}; - -template -struct find_option_impl - : std::conditional_t> {}; - -template -struct find_option_impl> : find_option_impl {}; - -template -struct add_option_impl; - -template -struct add_option_impl> { - using options = std::tuple; -}; - -} // namespace detail - -template -using find_option_t = typename detail::find_option_impl>::option_value; - -template -using add_option_t = typename detail::add_option_impl, std::tuple>::options; - -template -constexpr auto -add_option(Option new_option, std::tuple options_tuple) { - return add_option_t(); -} - -enum class Tag { - kIsDeltaRule, - kIsPersistent, - kNumMmaWarpGroups, - kStagesQ, - kStagesK, - kStagesV, - kNeedsAlpha, // gated delta rule - kNeedsBeta, // delta rule - kInitStateFromInput, // if true, initialize state by reading global memory instead of zero initialization. - kSafeGate, // KDA - kElementBetaGmem, // GMEM element type for beta (default float, can be bf16) -}; - -} // namespace kda::sm90::kernel diff --git a/csrc/kda/sm90/kernel/tile_scheduler.hpp b/csrc/kda/sm90/kernel/tile_scheduler.hpp deleted file mode 100644 index c70c7a63..00000000 --- a/csrc/kda/sm90/kernel/tile_scheduler.hpp +++ /dev/null @@ -1,155 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include -#include -#include - -namespace kda::sm90::kernel { - -using namespace cute; - -struct WorkDesc { - // coord - int32_t seq_idx; // which sequence to process - int32_t qk_head_idx; // head idx for Q/K (the representative of the GVA group) - int32_t head_idx; // head idx for V/O/g/beta - int64_t tok_offset; // start offset of this sequence in the packed tensor - - // shape - int64_t seq_len; // length of this sequence - - // update by mainloop - int32_t tile_idx = 0; // current tile index (mutated by the mainloop) - - template - CUTE_DEVICE bool - is_valid(Params const& params) { - return seq_idx >= 0 && seq_idx < params.num_seqs; - } - - CUTE_DEVICE int32_t - q_head_idx() const { - return qk_head_idx; - } - CUTE_DEVICE int32_t - k_head_idx() const { - return qk_head_idx; - } - CUTE_DEVICE int32_t - v_head_idx() const { - return head_idx; - } - CUTE_DEVICE int32_t - o_head_idx() const { - return head_idx; - } - - // compatible interface, for work without ChunkWiseParallel, chunk_len equals to seq_len - CUTE_DEVICE int32_t - chunk_len() const { - return seq_len; - } -}; - -// Each block handles a single (seq, v_head) work item; CTAs do not cooperate. -// GVA optimization: heads_per_group is precomputed on the host and stored in -// Params, so the device side does not redo the integer division per CTA. -struct IndividualTileScheduler { - struct Params { - dim3 grid; - int32_t num_seqs; - int32_t num_v_heads; - int32_t heads_per_group; // = num_v_heads / num_qk_heads, precomputed on host - }; - - bool scheduled = false; // a once flag - - CUTE_DEVICE - IndividualTileScheduler(Params const& params) { - } - - template - static Params - to_underlying_arguments( - ProblemSize const& problem_size, - cutlass::KernelHardwareInfo const& hw_info, - ClusterShape const& cluster_shape, - TileShape const& tile_shape) { - // Compute heads_per_group once on the host so every CTA does not have to redo - // the integer division. - int32_t const heads_per_group = problem_size.num_v_heads / problem_size.num_qk_heads; - dim3 grid(0, 1, 1); - grid.x = problem_size.num_seqs * problem_size.num_v_heads; - DPRINTF( - "to_underlying_arguments: grid:{.x:%d, .y:%d, .z:%d}, num_seqs:%d, num_qk_heads:%d, num_v_heads:%d, " - "heads_per_group:%d\n", - grid.x, - grid.y, - grid.z, - problem_size.num_seqs, - problem_size.num_qk_heads, - problem_size.num_v_heads, - heads_per_group); - return { - .grid = grid, - .num_seqs = problem_size.num_seqs, - .num_v_heads = problem_size.num_v_heads, - .heads_per_group = heads_per_group, - }; - } - - static dim3 - get_grid_shape(Params const& params) { - return params.grid; - } - - template - CUTE_DEVICE WorkDesc - get_next_work(Params params, ProblemSize const& problem_size) { - int32_t seq_idx = blockIdx.x / params.num_v_heads; - int32_t head_idx = blockIdx.x % params.num_v_heads; - // GVA: use the host-precomputed heads_per_group to avoid device-side division. - int32_t qk_head_idx = head_idx / params.heads_per_group; - - int32_t s = problem_size.cu_seqlens[seq_idx]; - int32_t e = problem_size.cu_seqlens[seq_idx + 1]; - int32_t seq_len = e - s; - - if (scheduled) { - seq_idx = -1; - } else { - scheduled = true; - DPRINTF0_W( - "get_next_work: this_work={seq_idx:%d qk_head_idx:%d head_idx:%d tok_offset:%lld seq_len:%lld}\n", - seq_idx, - qk_head_idx, - head_idx, - s, - seq_len); - } - - return { - .seq_idx = seq_idx, - .qk_head_idx = qk_head_idx, - .head_idx = head_idx, - .tok_offset = s, - .seq_len = seq_len, - }; - } -}; - -} // namespace kda::sm90::kernel diff --git a/csrc/kda/sm90/prefill_kernel.hpp b/csrc/kda/sm90/prefill_kernel.hpp deleted file mode 100644 index 6e54c3a3..00000000 --- a/csrc/kda/sm90/prefill_kernel.hpp +++ /dev/null @@ -1,54 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include - -#include - -namespace kda::sm90 { - -template < - typename ArchTag, // TODO: hide this - typename TO, - typename TQKV, - typename TState, - typename TBeta = float> -void -launch_kda_fwd_prefill_kernel( - cudaStream_t stream, - TO* output, - TState* output_state, - TQKV const* q, - TQKV const* k, - TQKV const* v, - TState const* input_state, - float const* alpha, - TBeta const* beta, - int32_t const* cu_seqlens, - uint8_t* workspace_buffer, - int32_t num_seqs, - int32_t num_qk_heads, - int32_t num_v_heads, - int32_t head_size, - int64_t total_seqlen, - float scale, - bool safe_gate, - int32_t sm_count = 0, - int32_t const* cp_seq_map = nullptr, - int32_t const* raw_cu_seqlens = nullptr, - int32_t raw_num_seqs = 0); - -} // namespace kda::sm90 diff --git a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh deleted file mode 100644 index c53f2ae3..00000000 --- a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh +++ /dev/null @@ -1,169 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include - -#include -#include -#include -#include - -#include "kda/sm90/device/device_universal.hpp" -#include "kda/sm90/kernel/builder_kda_fwd.hpp" -#include "kda/sm90/utils/common.hpp" - -namespace kda::sm90 { - -using namespace cute; - -template < - bool NeedsBeta, - bool NeedsAlpha, - bool InitStateFromInput, - bool SafeGate, - typename ArchTag, - typename TO, - typename TQKV, - typename TState, - typename TBeta = float> -void -launch_kda_fwd_prefill_kernel_gbai( - cudaStream_t stream, - TO* output, - TState* output_state, - TQKV const* q, - TQKV const* k, - TQKV const* v, - TState const* input_state, - float const* alpha, - TBeta const* beta, - int32_t const* cu_seqlens, - uint8_t* workspace_buffer, - int32_t num_seqs, - int32_t num_qk_heads, - int32_t num_v_heads, - int32_t head_size, - int64_t total_seqlen, - float scale, - int32_t sm_count, - int32_t const* cp_seq_map = nullptr, - int32_t const* raw_cu_seqlens = nullptr, - int32_t raw_num_seqs = 0) { -#if defined(CULA_SM90A_ENABLED) - constexpr bool HopperSupported = true; -#else - constexpr bool HopperSupported = false; -#endif - - if constexpr (HopperSupported) { - static_assert(std::is_same_v); - - using namespace kda::sm90::kernel; - using T = map_to_cutlass_t; - - cutlass::KernelHardwareInfo hw_info; - hw_info.sm_count = sm_count; - - using SafeGateType = std::conditional_t; - using NeedsBetaType = std::conditional_t; - using NeedsAlphaType = std::conditional_t; - using InitStateType = std::conditional_t; - using Options = decltype(add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option(Option{}, DefaultOptions{}))))))); - - using TileShape = Shape<_64, _64, _128>; - using Scheduler = cutlass::gemm::KernelTmaWarpSpecializedCooperative; - using Operation = cutlass::device::Universal, - /*LayoutK=*/cute::tuple, - /*LayoutV=*/cute::tuple, - /*LayoutO=*/cute::tuple, - Scheduler, - Options>::Kernel>; - using Arguments = typename Operation::Arguments; - - // NOTE: LayoutQ/K/V in (seq, head_size, (b,h)) coordinate semantics - // GVA: Q/K rows are packed as [packed_seq, num_qk_heads, head_size]; - // V/O/g/beta rows are packed as [packed_seq, num_v_heads, head_size]. - - int32_t qk_tok_stride = num_qk_heads * head_size; - int32_t v_tok_stride = num_v_heads * head_size; - int32_t head_stride = head_size; - - Operation op; - Arguments arguments{ - .problem_size = - { - .cu_seqlens = cu_seqlens, - .total_seqlen = total_seqlen, - .num_seqs = num_seqs, - .num_qk_heads = num_qk_heads, - .num_v_heads = num_v_heads, - .head_size = head_size, - .cp_seq_map = cp_seq_map, - .raw_cu_seqlens = raw_cu_seqlens, - .raw_num_seqs = raw_num_seqs, - }, - .mainloop = - { - // clang-format off - .ptr_Q = (T*)q, .dQ = {qk_tok_stride, _1{}, head_stride}, - .ptr_K = (T*)k, .dK = {qk_tok_stride, _1{}, head_stride}, - .ptr_V = (T*)v, .dV = {v_tok_stride, _1{}, head_stride}, - .ptr_O = (T*)output, .dO = {v_tok_stride, _1{}, head_stride}, - .ptr_Alpha = alpha, .dAlpha = {v_tok_stride, _1{}, head_stride}, - .ptr_output_state = (float*)output_state, - .ptr_input_state = (float*)input_state, - .scale = scale, - .beta_ptr = beta, .beta_stride = {num_v_heads, 1}, - }, // clang-format on - .hw_info = hw_info}; - - cutlass::Status status; - status = op.can_implement(arguments); - if (status != cutlass::Status::kSuccess) { - throw std::runtime_error("can_implement failed"); - } - - status = op.initialize(arguments, workspace_buffer, stream); - if (status != cutlass::Status::kSuccess) { - throw std::runtime_error("initialize failed"); - } - - status = op.run(stream); - if (status != cutlass::Status::kSuccess) { - throw std::runtime_error("run failed"); - } - - } else { - throw std::runtime_error("hopper not supported"); - } -} - -}; // namespace kda::sm90 diff --git a/csrc/kda/sm90/utils/common.hpp b/csrc/kda/sm90/utils/common.hpp deleted file mode 100644 index a307de68..00000000 --- a/csrc/kda/sm90/utils/common.hpp +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include -#include -#include - -#include "kda/sm90/utils/debug.hpp" - -#define FLAT_UNUSED_PARAMETER(x) (void)x - -#define CHECK(expr, msg) \ - do { \ - if (!(expr)) { \ - std::string buffer(1024, '\0'); \ - sprintf(buffer.data(), "Failed to check %s, %s at %s:%d\n", ##expr, msg __FILE__, __LINE__); \ - throw std::runtime_error(buffer.c_str()); \ - } \ - } while (0) - -#define CUDA_CHECK(expr) \ - do { \ - cudaError_t err = (expr); \ - if (err != cudaSuccess) { \ - std::string buffer(1024, '\0'); \ - sprintf( \ - buffer.data(), "CUDA Error: %s, Code: %d at %s:%d\n", cudaGetErrorName(err), err, __FILE__, __LINE__); \ - throw std::runtime_error(buffer.c_str()); \ - } \ - } while (0) diff --git a/csrc/kda/sm90/utils/debug.hpp b/csrc/kda/sm90/utils/debug.hpp deleted file mode 100644 index 244f46d8..00000000 --- a/csrc/kda/sm90/utils/debug.hpp +++ /dev/null @@ -1,149 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include - -#if DEBUG_PIPE -#define PIPE_DEBUG_PRINTF(fmt, ...) \ - if (threadIdx.x == 0) \ - printf("%s:%d " fmt, __FILE__, __LINE__, ##__VA_ARGS__) -#else -#define PIPE_DEBUG_PRINTF(...) -#endif - -#ifndef FLAT_DEBUG_PRINT -#define FLAT_DEBUG_PRINT 0 -#endif - -#if FLAT_DEBUG_PRINT -#define IS_PRINT_BLOCK cute::block(1) -#define DPRINTF(fmt, ...) \ - if (IS_PRINT_BLOCK) \ - printf("%s:%d " fmt, __FILE__, __LINE__, ##__VA_ARGS__) -#define DPRINTF0(fmt, ...) \ - if (IS_PRINT_BLOCK && threadIdx.x == 0) \ - printf("%s:%d " fmt, __FILE__, __LINE__, ##__VA_ARGS__) -#define DPRINTF_W(fmt, ...) \ - if (IS_PRINT_BLOCK) \ - printf( \ - "%s:%d [WG%d][W%d][T%-3d] " fmt, \ - __FILE__, \ - __LINE__, \ - threadIdx.x / 128, \ - threadIdx.x / 32, \ - threadIdx.x, \ - ##__VA_ARGS__) -#define DPRINTF0_W(fmt, ...) \ - if (IS_PRINT_BLOCK && threadIdx.x % 32 == 0) \ - printf( \ - "%s:%d [WG%d][W%d][T%-3d] " fmt, \ - __FILE__, \ - __LINE__, \ - threadIdx.x / 128, \ - threadIdx.x / 32, \ - threadIdx.x, \ - ##__VA_ARGS__) -#define DPRINTF_WG(fmt, ...) \ - if (IS_PRINT_BLOCK) \ - printf( \ - "%s:%d [WG%d][W%d][T%-3d] " fmt, \ - __FILE__, \ - __LINE__, \ - threadIdx.x / 128, \ - threadIdx.x / 32, \ - threadIdx.x, \ - ##__VA_ARGS__) -#define DPRINTF0_WG(fmt, ...) \ - if (IS_PRINT_BLOCK && threadIdx.x % 128 == 0) \ - printf( \ - "%s:%d [WG%d][W%d][T%-3d] " fmt, \ - __FILE__, \ - __LINE__, \ - threadIdx.x / 128, \ - threadIdx.x / 32, \ - threadIdx.x, \ - ##__VA_ARGS__) -#else -#define DPRINTF(...) -#define DPRINTF0(...) -#define DPRINTF_W(...) -#define DPRINTF0_W(...) -#define DPRINTF_WG(...) -#define DPRINTF0_WG(...) -#endif - -#if FLAT_DEBUG_PRINT -#define DPRINT_TMA_DESC(tma_dess_addr) \ - do { \ - auto p = reinterpret_cast(tma_dess_addr); \ - DPRINTF( \ - "\n" \ - "%08X%08X %08X%08X %08X%08X %08X%08X\n" \ - "%08X%08X %08X%08X %08X%08X %08X%08X\n" \ - "%08X%08X %08X%08X %08X%08X %08X%08X\n" \ - "%08X%08X %08X%08X %08X%08X %08X%08X\n", \ - p[0], \ - p[1], \ - p[2], \ - p[3], \ - p[4], \ - p[5], \ - p[6], \ - p[7], \ - p[8], \ - p[9], \ - p[10], \ - p[11], \ - p[12], \ - p[13], \ - p[14], \ - p[15], \ - p[16], \ - p[17], \ - p[18], \ - p[19], \ - p[20], \ - p[21], \ - p[22], \ - p[23], \ - p[24], \ - p[25], \ - p[26], \ - p[27], \ - p[28], \ - p[29], \ - p[30], \ - p[31]); \ - } while (0) -#else -#define DPRINT_TMA_DESC(tma_dess_addr) -#endif diff --git a/csrc/kda/sm90/utils/math.hpp b/csrc/kda/sm90/utils/math.hpp deleted file mode 100644 index 88cf206a..00000000 --- a/csrc/kda/sm90/utils/math.hpp +++ /dev/null @@ -1,53 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include - -namespace kda::sm90 { - -namespace detail { - -template -CUTE_HOST_DEVICE constexpr T -ceil_log2(T n) { - return n <= 1 ? 0 : 1 + ceil_log2((n + 1) / 2); -} - -} // namespace detail - -template -CUTE_HOST_DEVICE constexpr T -next_power_of_two(T n) { - return static_cast(1) << detail::ceil_log2(n); -} - -} // namespace kda::sm90 diff --git a/csrc/kda/sm90/utils/math_order_barrier.hpp b/csrc/kda/sm90/utils/math_order_barrier.hpp deleted file mode 100644 index 66c5a0d3..00000000 --- a/csrc/kda/sm90/utils/math_order_barrier.hpp +++ /dev/null @@ -1,116 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include -#include - -namespace kda::sm90 { - -// cutlass' OrderedSequenceBarrier uses mbarrier -template < - bool UseReservedNB_, // treat nb_id as cutlass::ReservedNamedBarriers - uint32_t... WGIdToNBIdMapping // say 6,4 is passed, means wg0 use nb6 and wg1 use nb4 - > -struct OrderedNamedBarriers { - static constexpr bool UseReservedNB = UseReservedNB_; - static constexpr int NumWG = sizeof...(WGIdToNBIdMapping); - using NBId_t = std::conditional_t; - - CUTE_DEVICE - OrderedNamedBarriers() : mapping_{NBId_t(WGIdToNBIdMapping)...} { - } - - CUTE_DEVICE - void - init(int wg_idx) { // wg_idx in among all WG participants - for (int i = wg_idx; i > 0; --i) { - cutlass::arch::NamedBarrier::arrive(cutlass::NumThreadsPerWarpGroup * NumWG, mapping_[i - 1]); - } - // with 3 WGs, init to namedbarrier_id:(arrived_wg,expected_wg) - // 0:(2,3) - // 1:(1,3) - // 2:(0,3) - } - - CUTE_DEVICE - ~OrderedNamedBarriers() { - // FIXME: this will be a problem for persistent scheduler - } - - CUTE_DEVICE - void - ordered_or_wait(int wg_idx) { // wg_idx in participants - // during first call, before - // 0:(2,3) - // 1:(1,3) - // 2:(0,3) - cutlass::arch::NamedBarrier::sync(cutlass::NumThreadsPerWarpGroup * NumWG, mapping_[wg_idx]); - // after - // 0:(3,3) immediately unblock wg0, and named barrier automatically reset to (0,3) - // 1:(2,3) - // 2:(1,3) - } - - CUTE_DEVICE - void - notify_next_blocked(int wg_idx) { // wg_idx in participants - // allways call this after ordered_or_wait - // during first call, before - // 0:(0,3) - // 1:(2,3) - // 2:(1,3) - CUTE_UNROLL - for (int i = 1; i < NumWG; ++i) { - cutlass::arch::NamedBarrier::arrive( - cutlass::NumThreadsPerWarpGroup * NumWG, mapping_[(wg_idx + i) % NumWG]); - } - // after wg0 called this function - // 0:(0,3), wg0 has not reached on second ordered_or_wait() or (1,3) wg0 wait on second ordered_or_wait() call - // 1:(0,3), unblocked wg1's first ordered_or_wait() and reset nb1 - // 2:(2,3), still wait on first ordered_or_wait() call - // - // after wg1 called this function - // 0:(1,3), wg0 has not reached on second ordered_or_wait() or (2,3) wg0 wait on second ordered_or_wait() call - // 1:(0,3), wg1 has not reached on second ordered_or_wait() or (1,3) wg1 wait on second ordered_or_wait() call - // 2:(0,3), unblocked wg2's first ordered_or_wait() and reset nb2 - // - // after wg2 called this function - // 0:(2,3), wg0 has not reached on second ordered_or_wait() or (0,3) wg0 wait on second ordered_or_wait() call, - // unblocked 1:(1,3), wg1 has not reached on second ordered_or_wait() or (2,3) wg1 wait on second - // ordered_or_wait() call, still block 2:(0,3), unblock wg0 ordered_or_wait() and reset - // - } - - private: - cute::array mapping_; -}; -} // namespace kda::sm90 diff --git a/csrc/kda/sm90/utils/type_traits.hpp b/csrc/kda/sm90/utils/type_traits.hpp deleted file mode 100644 index 14474ce4..00000000 --- a/csrc/kda/sm90/utils/type_traits.hpp +++ /dev/null @@ -1,66 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include - -#include - -namespace kda::sm90 { - -// clang-format off -template struct map_to_cutlass; -template<> struct map_to_cutlass { using type = cutlass::half_t; }; -template<> struct map_to_cutlass { using type = cutlass::bfloat16_t; }; -template<> struct map_to_cutlass { using type = cutlass::half_t; }; -template<> struct map_to_cutlass { using type = cutlass::bfloat16_t; }; - -template using map_to_cutlass_t = typename map_to_cutlass::type; -// clang-format on - -template -struct first_non_void { - static_assert(sizeof...(Ts) > 0, "all voids is not allowed"); - using type = void; -}; - -template -struct first_non_void { - using type = T; -}; - -template -struct first_non_void : first_non_void {}; - -template -using first_non_void_t = typename first_non_void::type; - -} // namespace kda::sm90 diff --git a/csrc/kda/sm90/utils/unused.hpp b/csrc/kda/sm90/utils/unused.hpp deleted file mode 100644 index af5b0a1c..00000000 --- a/csrc/kda/sm90/utils/unused.hpp +++ /dev/null @@ -1,54 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include - -namespace kda::sm90 { - -struct Unused { - using Params = Unused; - using SharedStorage = char; - static constexpr uint32_t Stages = 0; - - template - CUTE_HOST_DEVICE - Unused(Ts... vs) { - } - - template - CUTE_HOST_DEVICE Unused - operator=(T&& v) { - return Unused{}; - } -}; - -} // namespace kda::sm90 diff --git a/csrc/kerutils/include/kerutils/common/common.h b/csrc/kerutils/include/kerutils/common/common.h deleted file mode 100644 index e1bb9957..00000000 --- a/csrc/kerutils/include/kerutils/common/common.h +++ /dev/null @@ -1,35 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include -#include - -namespace kerutils {} - -#define KU_PRINTLN(fmt, ...) { cute::print(fmt, ##__VA_ARGS__); print("\n"); } - -#define CHECK_CUDA(call) \ - do { \ - cudaError_t status_ = call; \ - if (status_ != cudaSuccess) { \ - fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, cudaGetErrorString(status_)); \ - exit(1); \ - } \ - } while(0) - -#define CHECK_CUDA_KERNEL_LAUNCH() CHECK_CUDA(cudaGetLastError()) - -namespace ku = kerutils; diff --git a/csrc/kerutils/include/kerutils/common/cute_ext.hpp b/csrc/kerutils/include/kerutils/common/cute_ext.hpp deleted file mode 100644 index 72f46322..00000000 --- a/csrc/kerutils/include/kerutils/common/cute_ext.hpp +++ /dev/null @@ -1,77 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include -#include - -namespace kerutils { - -using namespace cute; - -template -__forceinline__ __host__ __device__ constexpr auto -select_layout(Layout&& l) { - if constexpr (is_composed_layout::value) { - return make_composed_layout( - l.layout_a(), - l.offset(), - select(l.layout_b()) - ); - } else { - return select(l); - } -} - -template -__forceinline__ __host__ __device__ constexpr auto -select_tensor(Tensor&& t) { - if constexpr (is_composed_layout::value) { - return make_tensor( - std::forward(t).data(), - make_composed_layout( - std::forward(t).layout().layout_a(), - std::forward(t).layout().offset(), - select(std::forward(t).layout().layout_b()) - ) - ); - } else { - return make_tensor(std::forward(t).data(), select(t.layout())); - } -} - -template -CUTE_DEVICE constexpr size_t -alignment_for_swizzle(Layout&& layout) { - return cutlass::detail::alignment_for_swizzle(std::forward(layout)); -} - -} // namespace kerutils diff --git a/csrc/kerutils/include/kerutils/common/static_switch.h b/csrc/kerutils/include/kerutils/common/static_switch.h deleted file mode 100644 index f796ac64..00000000 --- a/csrc/kerutils/include/kerutils/common/static_switch.h +++ /dev/null @@ -1,78 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -// Inspired by -// https://github.com/Dao-AILab/flash-attention/blob/main/hopper/static_switch.h -// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h -// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h - -#pragma once - -/// @param COND - a boolean expression to switch by -/// @param CONST_NAME - a name given for the constexpr bool variable -/// @param ... - code to execute for true and false -/// -/// Usage: -/// ``` -/// BOOL_SWITCH(flag, BoolConst, [&] { -/// some_function(...); -/// }); -/// ``` -#define BOOL_SWITCH(COND, CONST_NAME, ...) \ - [&] { \ - if (COND) { \ - constexpr static bool CONST_NAME = true; \ - return __VA_ARGS__(); \ - } else { \ - constexpr static bool CONST_NAME = false; \ - return __VA_ARGS__(); \ - } \ - }() - -/// @param COND - a boolean expression (true → BF16Type, false → FP32Type) -/// @param BF16Type - the type to use when COND is true -/// @param FP32Type - the type to use when COND is false -/// @param TYPE_ALIAS - a name given for the using type alias -/// @param ... - code to execute -/// -/// Usage: -/// ``` -/// FP_TYPE_SWITCH(is_bf16, __nv_bfloat16, float, ElemType, [&] { -/// some_function(...); -/// }); -/// ``` -#define FP_TYPE_SWITCH(COND, BF16Type, FP32Type, TYPE_ALIAS, ...) \ - [&] { \ - if (COND) { \ - using TYPE_ALIAS = BF16Type; \ - return __VA_ARGS__(); \ - } else { \ - using TYPE_ALIAS = FP32Type; \ - return __VA_ARGS__(); \ - } \ - }() - -/// Convenience shorthand: COND=true → __nv_bfloat16, COND=false → float -/// @param COND - a boolean expression -/// @param TYPE_ALIAS - a name given for the using type alias -/// @param ... - code to execute -/// -/// Usage: -/// ``` -/// BETA_TYPE_SWITCH(params.is_beta_bf16, BetaType, [&] { -/// some_function(...); -/// }); -/// ``` -#define BETA_TYPE_SWITCH(COND, TYPE_ALIAS, ...) \ - FP_TYPE_SWITCH(COND, __nv_bfloat16, float, TYPE_ALIAS, __VA_ARGS__) diff --git a/csrc/kerutils/include/kerutils/device/common.h b/csrc/kerutils/include/kerutils/device/common.h deleted file mode 100644 index 4e338145..00000000 --- a/csrc/kerutils/include/kerutils/device/common.h +++ /dev/null @@ -1,98 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include - -#include -#include -#include -#include - -namespace kerutils { - -using nvbf16 = __nv_bfloat16; -using nvbf16x2 = __nv_bfloat162; - -using tf32 = cutlass::tfloat32_t; -using bf16 = cutlass::bfloat16_t; -using fp16 = cutlass::half_t; -using transac_bar_t = cutlass::arch::ClusterTransactionBarrier; - -struct bf16x4 { - bf16 a, b, c, d; -}; - -struct nvbf16x4 { - __nv_bfloat162 a, b; -}; - -struct tf32x4 { - tf32 a, b, c, d; -}; - -struct bf16x8 { - __nv_bfloat162 a01; - __nv_bfloat162 a23; - __nv_bfloat162 a45; - __nv_bfloat162 a67; -}; - -struct bf16x16 { - __nv_bfloat162 a0; - __nv_bfloat162 a1; - __nv_bfloat162 a2; - __nv_bfloat162 a3; - __nv_bfloat162 a4; - __nv_bfloat162 a5; - __nv_bfloat162 a6; - __nv_bfloat162 a7; -}; - -} // namespace kerutils - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) -#define KERUTILS_ENABLE_SM80 -#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)) -static_assert(false, "kerutils doesn't support SM architectures below SM80"); -#endif - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) -#define KERUTILS_ENABLE_SM90 -#endif - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000)) -#define KERUTILS_ENABLE_SM90A -#endif - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) -#define KERUTILS_ENABLE_SM100 -#endif - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)) -#define KERUTILS_ENABLE_SM100A -#endif - -#if (defined(__CLION_IDE__) || defined(__VSCODE_IDE__)) -#define KERUTILS_ENABLE_SM80 -#define KERUTILS_ENABLE_SM90 -#define KERUTILS_ENABLE_SM90A -#define KERUTILS_ENABLE_SM100 -#define KERUTILS_ENABLE_SM100A -#endif - -#ifndef CUDA_VERSION_CHECK -#define CUDA_VERSION_CHECK (__CUDACC_VER_MAJOR__ * 1000 + __CUDACC_VER_MINOR__ * 10) -#endif diff --git a/csrc/kerutils/include/kerutils/device/device.cuh b/csrc/kerutils/include/kerutils/device/device.cuh deleted file mode 100644 index ea6e1374..00000000 --- a/csrc/kerutils/include/kerutils/device/device.cuh +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include "kerutils/common/common.h" -#include "kerutils/common/cute_ext.hpp" - -#include "common.h" -#include "sm80/helpers.cuh" -#include "sm80/collective_inverse.hpp" -#include "sm90/helpers.cuh" -#include "sm100/intrinsics.cuh" -#include "sm100/helpers.cuh" diff --git a/csrc/kerutils/include/kerutils/device/sm100/helpers.cuh b/csrc/kerutils/include/kerutils/device/sm100/helpers.cuh deleted file mode 100644 index 041c857c..00000000 --- a/csrc/kerutils/include/kerutils/device/sm100/helpers.cuh +++ /dev/null @@ -1,116 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 DeepSeek. All Rights Reserved. - * - * Licensed under the MIT License. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ - -#pragma once - -#include - -#include "kerutils/device/common.h" - -namespace kerutils { - -// Perform SS UTCMMA -// sA and sB should be shared memory tensors (i.e. make_tensor(make_shared_ptr(XXX), XXX)) while tC_frag should be tmem fragment -template< - typename TiledMMA, - typename TensorA, - typename TensorB, - typename TensorFragC -> -CUTE_DEVICE -void utcmma_ss( - TiledMMA &tiled_mma, - TensorA sA, - TensorB sB, - TensorFragC tC_frag, - bool clear_accum -) { - using namespace cute; - tiled_mma.accumulate_ = clear_accum ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; - ThrMMA thr_mma = tiled_mma.get_slice(_0{}); // Since A/B/C are already CTA-local tiles, this number does not matter - auto sA_frag = thr_mma.partition_fragment_A(sA); - auto sB_frag = thr_mma.partition_fragment_B(sB); - static_assert(size<2>(sA_frag) == size<2>(sB_frag)); - static_assert(size<1>(sA_frag) == size<1>(tC_frag)); - static_assert(size<1>(sB_frag) == size<2>(tC_frag)); - CUTE_UNROLL - for (int k = 0; k < size<2>(sA_frag); ++k) { - cute::gemm( - tiled_mma, - sA_frag(_, _, k), - sB_frag(_, _, k), - tC_frag - ); - tiled_mma.accumulate_ = UMMA::ScaleOut::One; - } -} - -// Perform TS UTCMMA -// sB should be shared memory tensors (i.e. make_tensor(make_shared_ptr(XXX), XXX)) while tA_frag and tC_frag should be tmem fragment -template< - typename TiledMMA, - typename TensorA, - typename TensorB, - typename TensorFragC -> -CUTE_DEVICE -void utcmma_ts( - TiledMMA &tiled_mma, - TensorA tA_frag, - TensorB sB, - TensorFragC tC_frag, - bool clear_accum -) { - using namespace cute; - tiled_mma.accumulate_ = clear_accum ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; - ThrMMA thr_mma = tiled_mma.get_slice(_0{}); // Since A/B/C are already CTA-local tiles, this number does not matter - auto sB_frag = thr_mma.partition_fragment_B(sB); - static_assert(size<2>(tA_frag) == size<2>(sB_frag)); - CUTE_UNROLL - for (int k = 0; k < size<2>(tA_frag); ++k) { - cute::gemm( - tiled_mma, - tA_frag(_, _, k), - sB_frag(_, _, k), - tC_frag - ); - tiled_mma.accumulate_ = UMMA::ScaleOut::One; - } -} - -} // namespace kerutils diff --git a/csrc/kerutils/include/kerutils/device/sm100/intrinsics.cuh b/csrc/kerutils/include/kerutils/device/sm100/intrinsics.cuh deleted file mode 100644 index d1605191..00000000 --- a/csrc/kerutils/include/kerutils/device/sm100/intrinsics.cuh +++ /dev/null @@ -1,239 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 DeepSeek. All Rights Reserved. - * - * Licensed under the MIT License. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ - -#pragma once - -#include - -#include "kerutils/device/common.h" - -// Adapted from https://github.com/deepseek-ai/FlashMLA/blob/main/csrc/kerutils/include/kerutils/device/sm100/intrinsics.cuh -namespace kerutils { - -// ============================================================ -// TMA gather4 intrinsics (SM100) -// ============================================================ - -// tma gather4 (https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor) -// Please pay attention that the coordinates of TMA gather4 are int32, which may lead to overflow under some scenarios -CUTE_DEVICE -void tma_gather4(const void* desc_ptr, transac_bar_t &mbar_ptr, void* smem_ptr, int col_idx, int4 row_idxs, int64_t cache_hint) { - uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); - uint32_t mbar_addr = cute::cast_smem_ptr_to_uint(&mbar_ptr); - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;\n" - : - : "r"(smem_addr), "l"(desc_ptr), "r"(col_idx), - "r"(row_idxs.x), "r"(row_idxs.y), "r"(row_idxs.z), "r"(row_idxs.w), - "r"(mbar_addr), "l"(cache_hint) - : "memory" - ); -} - -// tma gather4 prefetch (https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-prefetch-tensor) -// Please pay attention that the coordinates of TMA gather4 are int32, which may lead to overflow under some scenarios -CUTE_DEVICE -void tma_gather4_prefetch(const void* desc_ptr, int col_idx, int4 row_idxs, int64_t cache_hint) { - asm volatile( - "cp.async.bulk.prefetch.tensor.2d.L2.global.tile::gather4.L2::cache_hint [%0, {%1, %2, %3, %4, %5}], %6;\n" - : - : "l"(desc_ptr), "r"(col_idx), - "r"(row_idxs.x), "r"(row_idxs.y), "r"(row_idxs.z), "r"(row_idxs.w), - "l"(cache_hint) - ); -} - -// tma gather4 with cta_group::2, allowing for synchronization across CTAs within a pair of CTAs (https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor) -template -CUTE_DEVICE void tma_gather4_cta_group_2(const void* desc_ptr, transac_bar_t &mbar_ptr, void* smem_ptr, int col_idx, int4 row_idxs, int64_t cache_hint) { - uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); - uint32_t mbar_addr = cute::cast_smem_ptr_to_uint(&mbar_ptr); - if constexpr (USE_CTA0_MBAR) { - mbar_addr &= cute::Sm100MmaPeerBitMask; - } - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::2.L2::cache_hint [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;\n" - : - : "r"(smem_addr), "l"(desc_ptr), "r"(col_idx), - "r"(row_idxs.x), "r"(row_idxs.y), "r"(row_idxs.z), "r"(row_idxs.w), - "r"(mbar_addr), "l"(cache_hint) - : "memory" - ); -} - -// ============================================================ -// Vectorized float2 arithmetic -// ============================================================ - -// Vectorized addition for float32 (https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-add) -CUTE_DEVICE -float2 float2_add(const float2 &a, const float2 &b) { - float2 c; - asm volatile( - "add.f32x2 %0, %1, %2;\n" - : "=l"(reinterpret_cast(c)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b)) - ); - return c; -} - -// Vectorized subtraction for float32 -CUTE_DEVICE -float2 float2_sub(const float2 &a, const float2 &b) { - float2 c; - asm volatile( - "sub.f32x2 %0, %1, %2;\n" - : "=l"(reinterpret_cast(c)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b)) - ); - return c; -} - -// Vectorized multiplication for float32 (https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-mul) -CUTE_DEVICE -float2 float2_mul(const float2 &a, const float2 &b) { - float2 c; - asm volatile( - "mul.f32x2 %0, %1, %2;\n" - : "=l"(reinterpret_cast(c)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b))); - return c; -} - -// Vectorized fused multiply-add for float32 (https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma) -CUTE_DEVICE -float2 float2_fma(const float2 &a, const float2 &b, const float2 &c) { - // return a*b+c - float2 d; - asm volatile( - "fma.rn.f32x2 %0, %1, %2, %3;\n" - : "=l"(reinterpret_cast(d)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b)), - "l"(reinterpret_cast(c))); - return d; -} - -// Vectorized negation for float32 -CUTE_DEVICE -float2 float2_neg(const float2 &a) { - float2 t = {-1.0f, -1.0f}; - return float2_mul(a, t); -} - -// ============================================================ -// tcgen05 fence intrinsics (SM100) -// ============================================================ - -// tcgen05.fence::before_thread_sync (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-special-sync-operations-fence) -__device__ __forceinline__ void tcgen05_before_thread_sync() { - asm volatile("tcgen05.fence::before_thread_sync;"); -} - -// tcgen05.fence::after_thread_sync (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-special-sync-operations-fence) -__device__ __forceinline__ void tcgen05_after_thread_sync() { - asm volatile("tcgen05.fence::after_thread_sync;"); -} - -// ============================================================ -// Tensor memory (TMEM) load/store intrinsics (SM100) -// ============================================================ - -// Load from tensor memory, 32 data path lanes, 32-bit pattern, repeated N times. (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld) -template -__device__ __forceinline__ -void tmem_ld_32dp32bNx(uint32_t tmem_start, void* data_) { - uint32_t* data = (uint32_t*)data_; - static_assert(kNumElements == 1 || kNumElements == 2 || kNumElements == 4 || kNumElements == 8 || kNumElements == 16 || kNumElements == 32 || kNumElements == 64 || kNumElements == 128, "Invalid kNumElements"); - // NOTE The following code crashes VSCode intellisense engine, so we disable it -#ifndef __VSCODE_IDE__ - [&](cute::index_sequence) { - if constexpr (kNumElements == 1) { - cute::SM100_TMEM_LOAD_32dp32b1x::copy(tmem_start, data[Is]...); - } else if constexpr (kNumElements == 2) { - cute::SM100_TMEM_LOAD_32dp32b2x::copy(tmem_start, data[Is]...); - } else if constexpr (kNumElements == 4) { - cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_start, data[Is]...); - } else if constexpr (kNumElements == 8) { - cute::SM100_TMEM_LOAD_32dp32b8x::copy(tmem_start, data[Is]...); - } else if constexpr (kNumElements == 16) { - cute::SM100_TMEM_LOAD_32dp32b16x::copy(tmem_start, data[Is]...); - } else if constexpr (kNumElements == 32) { - cute::SM100_TMEM_LOAD_32dp32b32x::copy(tmem_start, data[Is]...); - } else if constexpr (kNumElements == 64) { - cute::SM100_TMEM_LOAD_32dp32b64x::copy(tmem_start, data[Is]...); - } else if constexpr (kNumElements == 128) { - cute::SM100_TMEM_LOAD_32dp32b128x::copy(tmem_start, data[Is]...); - } - }(cute::make_index_sequence{}); -#endif -} - -// Store into tensor memory, 32 data path lanes, 32-bit pattern, repeated N times. (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st) -template -__device__ __forceinline__ -void tmem_st_32dp32bNx(uint32_t tmem_start, void const* data_) { - uint32_t const* data = (uint32_t const*)data_; - static_assert(kNumElements == 1 || kNumElements == 2 || kNumElements == 4 || kNumElements == 8 || kNumElements == 16 || kNumElements == 32 || kNumElements == 64 || kNumElements == 128, "Invalid kNumElements"); -#ifndef __VSCODE_IDE__ - [&](cute::index_sequence) { - if constexpr (kNumElements == 1) { - cute::SM100_TMEM_STORE_32dp32b1x::copy(data[Is]..., tmem_start); - } else if constexpr (kNumElements == 2) { - cute::SM100_TMEM_STORE_32dp32b2x::copy(data[Is]..., tmem_start); - } else if constexpr (kNumElements == 4) { - cute::SM100_TMEM_STORE_32dp32b4x::copy(data[Is]..., tmem_start); - } else if constexpr (kNumElements == 8) { - cute::SM100_TMEM_STORE_32dp32b8x::copy(data[Is]..., tmem_start); - } else if constexpr (kNumElements == 16) { - cute::SM100_TMEM_STORE_32dp32b16x::copy(data[Is]..., tmem_start); - } else if constexpr (kNumElements == 32) { - cute::SM100_TMEM_STORE_32dp32b32x::copy(data[Is]..., tmem_start); - } else if constexpr (kNumElements == 64) { - cute::SM100_TMEM_STORE_32dp32b64x::copy(data[Is]..., tmem_start); - } else if constexpr (kNumElements == 128) { - cute::SM100_TMEM_STORE_32dp32b128x::copy(data[Is]..., tmem_start); - } - }(cute::make_index_sequence{}); -#endif -} - -} // namespace kerutils diff --git a/csrc/kerutils/include/kerutils/device/sm80/collective_inverse.hpp b/csrc/kerutils/include/kerutils/device/sm80/collective_inverse.hpp deleted file mode 100644 index 655140eb..00000000 --- a/csrc/kerutils/include/kerutils/device/sm80/collective_inverse.hpp +++ /dev/null @@ -1,1145 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -#pragma once - -#include -#include -#include - -#include "kerutils/common/cute_ext.hpp" - -namespace kerutils { - -using namespace cute; - -template -constexpr bool -is_contiguous(Layout&& layout) { - auto dim_layout = get(layout); - if constexpr(rank(dim_layout) == 0) { - return stride(dim_layout) == 1; - } else { - return stride<0>(dim_layout) == 1; - } -} - -namespace detail::SM80 { - -// ============================================================================ -// convert_c_layout_to_a_layout — Compile-time layout reshape: C-fragment → A-fragment -// ============================================================================ -// The C-fragment (accumulator) layout is (FragAtomC, MMA_M, MMA_N). -// The A-fragment (operand A) layout is (FragAtomA, MMA_M, MMA_K). -// -// When FragAtomA > FragAtomC (e.g. SM80_16x8x16 where A has 8 vals/atom but -// C has 4 vals/atom, ratio=2), the N-dimension is split: MMA_N → (ratio, MMA_N/ratio), -// and the "ratio" factor is folded into the value (FragAtom) dimension: -// C: ((2,2), MMA_M, MMA_N) → ((2,2,ratio), MMA_M, MMA_N/ratio) -// -// For SM80_16x8x8 TF32 (used in the inverse), FragAtomA == FragAtomC == 4, -// so ratio=1 and this function returns the C-layout unchanged. The function is -// still called to maintain a uniform code path and to reinterpret the fragment's -// rank-3 indexing in A-operand terms for the per-k-atom shuffle loop. -// -template -CUTE_DEVICE constexpr auto -convert_c_layout_to_a_layout(CLayout const& c, TiledMMA const& tiled_mma) { - constexpr auto c_frag_atom_size = size<0>(CLayout{}); // values per C atom (4 for SM80_16x8) - constexpr auto a_frag_atom_size = size<1>(typename TiledMMA::AtomLayoutA_TV{}); // values per A atom - static_assert(a_frag_atom_size % c_frag_atom_size == 0); - constexpr auto ratio = a_frag_atom_size / c_frag_atom_size; // 1 for 16x8x8, 2 for 16x8x16 - if constexpr (ratio == 1) { - // SM80_16x8x8: C and A have the same value count per atom → no reshape needed - return CLayout{}; - } else { - // SM80_16x8x16: A atom has 2× more values than C atom. - // Split N-dimension by ratio, fold the extra factor into the value dimension: - // C: (FragAtom, MMA_M, MMA_N) - // → logical_divide by (_, _, ratio): (FragAtom, MMA_M, (ratio, MMA_N/ratio)) - // → flatten first + third[0]: ((FragAtom, ratio), MMA_M, MMA_N/ratio) - constexpr auto tiler = make_shape(_, _, Int{}); // keep FragAtom and MMA_M, split MMA_N - constexpr auto divided = logical_divide(CLayout{}, tiler); // (FragAtom, MMA_M, (ratio, MMA_N/ratio)) - - return make_layout(flatten(make_layout(get<0>(divided), get<2, 0>(divided))), get<1>(divided), get<2, 1>(divided)); - } -} - -// ============================================================================ -// make_acc_into_op — Generic acc→operand conversion (FP16 path only) -// ============================================================================ -// Works for FP16 MMA because C-layout and A-layout have the SAME thread-to-K -// mapping ({2*t0, 2*t0+1} for both). A simple in-thread element copy suffices. -// For TF32 MMA, thread-to-K mappings DIFFER between C and A layouts, so this -// function is INCORRECT for TF32. Use convert_fp32_acc_to_tf32_operandA_layout -// instead, which performs cross-thread __shfl_sync shuffles. -template -CUTE_DEVICE auto -make_acc_into_op(Accumulator const& acc, TiledMMA const& tiled_mma) { - Tensor operand = make_fragment_like(convert_c_layout_to_a_layout(acc.layout(), tiled_mma)); - Tensor operand_as_acc = make_tensor(operand.data(), acc.layout()); - cute::copy(acc, operand_as_acc); - return operand; -} - -// ============================================================================ -// convert_fp32_acc_to_tf32_operandA_layout — Cross-thread shuffle: C-layout → A-layout (TF32) -// ============================================================================ -// Converts an MMA accumulator fragment (C-layout, float) into an A-operand fragment -// (A-layout, float or tf32) via warp-level __shfl_sync shuffles. -// -// Why this is needed: -// The SM80 TF32 MMA (16x8x8) has different thread-value (TV) mappings for -// C-layout vs A-layout. Specifically, the K-dimension ownership differs: -// C-layout: t0 selects M-row group (stride 32), v0/v1 select M sub-rows -// A-layout: t0 selects K-column pair {t0, t0+4}, v0 selects M-half, v1 selects K-half -// This means values must be redistributed ACROSS threads (not just within a thread). -// -// Algorithm overview (per k-atom, 4 values each): -// 1. Read all 4 src values into locals (avoid read-after-write hazard) -// 2. For each M-half (v0_tf32 ∈ {0,1}): -// a. Select the 2 source values for this M-half (maps to v1_bf16 in C-layout) -// b. Shuffle from source threads: src_lane = t0/2 (K-lo) or t0/2+2 (K-hi) -// c. Select v0_bf16=0 or v0_bf16=1 based on t0 parity (even/odd) -// 3. Write 4 output values (with optional TF32 rounding) -// -// Cost: 8 __shfl_sync per k-atom (2 M-halves × 4 shuffles each). -// -// Template parameters: -// NumKAtoms — number of K-atoms to process (1, 2, or 4) -// RoundingTF32 — if true, output is tfloat32_t with explicit rounding; -// if false, output is float and MMA hardware truncates -template -CUTE_DEVICE void -convert_fp32_acc_to_tf32_operandA_layout( - const FragSrc& frag_src, FragDst& frag_dst, - TiledMMA const& tiled_mma, int lane_id) -{ - using ElemSrc = typename cute::remove_cvref_t::value_type; - using ElemDst = typename cute::remove_cvref_t::value_type; - static_assert(cute::is_same_v, "Fragment input must be float with acc data dtype"); - static_assert(RoundingTF32 || cute::is_same_v, "Fragment must be float with no rounding; tf32 truncation is done by MMA hardware"); - static_assert(!RoundingTF32 || cute::is_same_v, "Fragment must be tfloat32 with rounding; tf32 truncation is done manually"); - - // Reinterpret the source fragment with A-layout indexing. - // For SM80_16x8x8 TF32, ratio=1 so this is effectively a no-op (same shape). - // For SM80_16x8x16, this would reshape (4, M, N) → (8, M, N/2) to align k-atoms. - auto frag_src_cvt = make_tensor(frag_src.data(), convert_c_layout_to_a_layout(frag_src.layout(), tiled_mma)); - - // Thread decomposition within a warp of 32 threads. - // tid = t0 + 4*t1, where t0 ∈ [0,4) indexes K-groups, t1 ∈ [0,8) indexes M-rows. - int tid = lane_id; - int t0 = tid % 4; - // Parity of t0 determines which v0_bf16 value to select after shuffle: - // even t0 → needs v0_bf16=0 (K = 2*(t0/2) = t0) - // odd t0 → needs v0_bf16=1 (K = 2*(t0/2)+1 = t0) - bool sel_odd = (t0 & 1); - - // Compute source lane IDs for the K-dimension shuffle. - // In C-layout, thread with t0_src holds K = {2*t0_src, 2*t0_src+1}. - // We need K = t0 (from t0_src = t0/2) and K = t0+4 (from t0_src = t0/2+2). - // The upper bits of lane_id (t1 portion) stay the same: (tid & ~3). - int src_lane_lo = (t0 / 2) + (tid & ~3); // source for K = t0 (v1_tf32=0) - int src_lane_hi = (t0 / 2 + 2) + (tid & ~3); // source for K = t0 + 4 (v1_tf32=1) - - // Process each k-atom independently. Each atom has 4 values indexed as: - // C-layout: [4j+0]=(v0=0,v1=0), [4j+1]=(v0=1,v1=0), [4j+2]=(v0=0,v1=1), [4j+3]=(v0=1,v1=1) - // where v0 selects K sub-position, v1 selects M sub-row - // A-layout: [4j+0]=(v0=0,v1=0), [4j+1]=(v0=1,v1=0), [4j+2]=(v0=0,v1=1), [4j+3]=(v0=1,v1=1) - // where v0 selects M half (0 or +8), v1 selects K half (lo or hi) - CUTE_UNROLL - for (int j = 0; j < NumKAtoms; j++) { - // Step 1: Read all 4 input values BEFORE writing any output. - // This prevents read-after-write hazard when src and dst share storage - // (output positions overlap with input positions within the same k-atom). - float in0 = frag_src_cvt(0 + 4*j); // C-layout: (v0=0, v1=0) → K-lo, M-row-0 - float in1 = frag_src_cvt(1 + 4*j); // C-layout: (v0=1, v1=0) → K-hi, M-row-0 - float in2 = frag_src_cvt(2 + 4*j); // C-layout: (v0=0, v1=1) → K-lo, M-row-1 - float in3 = frag_src_cvt(3 + 4*j); // C-layout: (v0=1, v1=1) → K-hi, M-row-1 - - // Step 2: Process each M-half (v0_tf32) independently. - // v0_tf32=0 (M-half 0) needs data from v1_bf16=0 → (in0, in1) - // v0_tf32=1 (M-half 1) needs data from v1_bf16=1 → (in2, in3) - float out_vals[4]; - CUTE_UNROLL - for (int v0_tf32 = 0; v0_tf32 < 2; v0_tf32++) { - // Select the two source values for this M-half. - // In C-layout, v1 selects M-row; in A-layout, v0 selects M-half. - // So v0_tf32 maps to v1_bf16 for M-row selection. - float val0 = (v0_tf32 == 0) ? in0 : in2; // v0_bf16=0 at chosen M-half - float val1 = (v0_tf32 == 0) ? in1 : in3; // v0_bf16=1 at chosen M-half - - // Step 3: Cross-thread shuffle for K-dimension remapping. - // Fetch both v0_bf16 values (=0 and =1) from the source threads, - // then select the correct one based on t0 parity. - float recv0_lo = __shfl_sync(0xFFFFFFFF, val0, src_lane_lo); // v0_bf16=0 from K-lo source - float recv1_lo = __shfl_sync(0xFFFFFFFF, val1, src_lane_lo); // v0_bf16=1 from K-lo source - float recv0_hi = __shfl_sync(0xFFFFFFFF, val0, src_lane_hi); // v0_bf16=0 from K-hi source - float recv1_hi = __shfl_sync(0xFFFFFFFF, val1, src_lane_hi); // v0_bf16=1 from K-hi source - - // Step 4: Select correct value based on t0 parity. - // Even t0 needs v0_bf16=0 (recv0), odd t0 needs v0_bf16=1 (recv1). - out_vals[v0_tf32 + 0] = sel_odd ? recv1_lo : recv0_lo; // v1_tf32=0 (K-lo half) - out_vals[v0_tf32 + 2] = sel_odd ? recv1_hi : recv0_hi; // v1_tf32=1 (K-hi half) - } - - // Step 5: Write output values in A-layout order. - // If RoundingTF32: explicit cast to tfloat32_t applies rounding (x += 0x1000u). - // If !RoundingTF32: store as float; MMA hardware truncates to TF32 precision. - if constexpr (RoundingTF32) { - frag_dst(0 + 4*j) = (ElemDst)out_vals[0]; // (v0_tf32=0, v1_tf32=0) → M-half-0, K-lo - frag_dst(1 + 4*j) = (ElemDst)out_vals[1]; // (v0_tf32=1, v1_tf32=0) → M-half-1, K-lo - frag_dst(2 + 4*j) = (ElemDst)out_vals[2]; // (v0_tf32=0, v1_tf32=1) → M-half-0, K-hi - frag_dst(3 + 4*j) = (ElemDst)out_vals[3]; // (v0_tf32=1, v1_tf32=1) → M-half-1, K-hi - } else { - frag_dst(0 + 4*j) = out_vals[0]; - frag_dst(1 + 4*j) = out_vals[1]; - frag_dst(2 + 4*j) = out_vals[2]; - frag_dst(3 + 4*j) = out_vals[3]; - } - } -} - -} - -template -struct CollectiveInverseTF32 { - static_assert(std::is_same_v); - using ElementView = std::conditional_t; - - CUTE_DEVICE - CollectiveInverseTF32(int wg_sync_named_barrier_id): wg_sync_named_barrier_id_(wg_sync_named_barrier_id) {} - - template - CUTE_DEVICE void - compute(TensorT&& sT) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 64); - static_assert(size<1>(L) == 64); - - int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; - - auto sT_view = recast(sT); - auto t16X16sT = flat_divide(sT_view, Shape<_16, _16>{}); - - if (thread_idx < 64) { // compute 16x16 inverse on diagnal directly - compute_diagonal_inverse_NxN<16>(t16X16sT(_, _, thread_idx / 16, thread_idx / 16), thread_idx % 16); - } - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - - auto t32X32sT = flat_divide(sT_view, Shape<_32, _32>{}); - if (thread_idx < 64) { // two warps for 16x16 -> 32x32 - blockwise_diagonal_inversed_16x16_to_32x32(t32X32sT(_, _, thread_idx / 32, thread_idx /32)); - } - - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - // one warpgroup for 32x32 -> 64x64 - blockwise_diagonal_inversed_32x32_to_64x64(sT_view); - } - -private: - - template - CUTE_DEVICE void - compute_diagonal_inverse_NxN(TensorT&& mat, int tid_in_group) { // group_size = N - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == N); - static_assert(size<1>(L) == N); - - using ElementCompute = float; - - using CopyOp = Copy_Atom, ElementView>; - - auto load_row = [&](int y) { - auto row = make_tensor(Shape>{}); - copy(CopyOp{}, std::forward(mat)(y, _), row); - - auto row_cvt = [&]() { - if constexpr (RoundingTF32) { - return make_tensor_like(row); - } else { - return row; - } - }(); - if constexpr (RoundingTF32) { - copy(row, row_cvt); - } - - if constexpr (GarbageFilledDiagonal || GarbageFilledUpperTriangular) { - CUTE_UNROLL - for (int i = 0; i < N; ++i) { - row_cvt(i) = i == y ? 1.0f - : (i > y ? 0.0f - : row_cvt(i)); - } - } - return row_cvt; - }; - - auto store_row = [&](int y, auto row) { - auto row_cvt = [&]() { - if constexpr (RoundingTF32) { - return make_tensor_like(row); - } else { - return row; - } - }(); - if constexpr (RoundingTF32) { - copy(row, row_cvt); - } - copy(CopyOp{}, row_cvt, std::forward(mat)(y, _)); - }; - - auto row = load_row(tid_in_group); -#define LOAD(y, x) __shfl_sync(0xffffffff, row(x), y, N) - - CUTE_UNROLL - for (int src_row = 0; src_row < N - 1; ++src_row) { // idx of src row to eliminate - auto row_scale = -row(src_row); // scale the src row - CUTE_UNROLL - for (int i = 0; i < src_row; ++i) { - auto src_row_value = LOAD(src_row, i); - row(i) = tid_in_group > src_row ? row_scale * src_row_value + row(i) : row(i); - } - row(src_row) = tid_in_group > src_row ? row_scale : row(src_row); - } - -#undef LOAD - - store_row(tid_in_group, row); - } - - /* - blockwise inverse has relation as follows - inv(| A 0 |) = | inv(A) 0 | - | C D | | -inv(D)C inv(A) inv(D) | - */ - - // not used - template - CUTE_DEVICE void - blockwise_diagonal_inversed_8x8_to_16x16(TensorT&& mat) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 16); - static_assert(size<1>(L) == 16); - - static_assert(is_contiguous<0>(L) == 1 || is_contiguous<1>(L) == 1); - constexpr bool is_col_major = is_contiguous<0>(L); - - auto mat_8x8_2x2 = flat_divide(std::forward(mat), Shape<_8, _8>{}); - using MMA = SM80_16x8x8_F32TF32TF32F32_TN; - using TiledMMA = decltype(make_tiled_mma(MMA{}, Layout>{}, Shape<_16, _8, _8>{})); - - using CopyOpTF32 = AutoVectorizingCopyWithAssumedAlignment<128>; - using CopyOpD_S2R = CopyOpTF32; - using CopyOpC_S2R = CopyOpTF32; - using CopyOpA_S2R = CopyOpTF32; -#ifdef CUTE_ARCH_STSM_SM90_ENABLED - using CopyOpO_R2S = CopyOpTF32; -#else - using CopyOpO_R2S = UniversalCopy; -#endif - - int lane_id = cutlass::canonical_lane_idx(); - auto tiled_mma = TiledMMA{}; - auto thr_mma = tiled_mma.get_thread_slice(lane_id); - - auto D_tiled_copy = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto C_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto A_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto O_tiled_copy = make_tiled_copy_C(Copy_Atom{}, tiled_mma); - - auto D_thr_copy = D_tiled_copy.get_thread_slice(lane_id); - auto C_thr_copy = C_tiled_copy.get_thread_slice(lane_id); - auto A_thr_copy = A_tiled_copy.get_thread_slice(lane_id); - auto O_thr_copy = O_tiled_copy.get_thread_slice(lane_id); - - Tensor sDinv = mat_8x8_2x2(_, _, _1{}, _1{}); - Tensor sC = select_tensor<1, 0>(mat_8x8_2x2(_, _, _1{}, _0{})); - Tensor sAinv = select_tensor<1, 0>(mat_8x8_2x2(_, _, _0{}, _0{})); - Tensor sO = mat_8x8_2x2(_, _, _1{}, _0{}); - - Tensor sDinv_m_bcast = make_tensor(sDinv.data(), logical_product(sDinv.layout(), Tile>{})); - Tensor sO_m_bcast = make_tensor(sO.data(), logical_product(sO.layout(), Tile>{})); - - Tensor tOrDinv = make_fragment_like(partition_shape_A(tiled_mma, Shape<_16, _8>{})); - Tensor tOrC = make_fragment_like(partition_shape_B(tiled_mma, Shape<_8, _8>{})); - // Tensor tOrC = thr_mma.partition_fragment_B(sC); - Tensor tOrAinv = make_fragment_like(partition_shape_B(tiled_mma, Shape<_8, _8>{})); - // Tensor tOrAinv = thr_mma.partition_fragment_B(sAinv); - - Tensor tDCrDC = partition_fragment_C(tiled_mma, Shape<_16, _8>{}); // output of -inv(D)C - Tensor tOrO = partition_fragment_C(tiled_mma, Shape<_16, _8>{}); // output of -inv(D)C inv(A) - - Tensor tOsDinv = D_thr_copy.partition_S(sDinv_m_bcast); - Tensor tOrDinv_cv = D_thr_copy.retile_D(tOrDinv); - Tensor tOsC = C_thr_copy.partition_S(sC); - Tensor tOrC_cv = C_thr_copy.retile_D(tOrC); - Tensor tOsAinv = A_thr_copy.partition_S(sAinv); - Tensor tOrAinv_cv = A_thr_copy.retile_D(tOrAinv); - Tensor tOsO = O_thr_copy.partition_D(sO_m_bcast); - Tensor tOrO_cv = O_thr_copy.retile_S(tOrO); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C - copy(D_tiled_copy, tOsDinv, tOrDinv_cv); - copy(C_tiled_copy, tOsC, tOrC_cv); - - clear(tDCrDC); - auto tOrDinv_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrDinv); - } else { - return tOrDinv; - } - }(); - auto tOrC_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrC); - } else { - return tOrC; - } - }(); - gemm(tiled_mma, tOrDinv_mma, tOrC_mma, tDCrDC); - transform(tDCrDC, [](auto v) { return -v; }); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C inv(A) - Tensor tOrDC = make_fragment_like(partition_shape_A(tiled_mma, Shape<_16, _8>{})); - // Tensor tOrDC = detail::SM80::make_acc_into_op(tDCrDC, tiled_mma); - detail::SM80::convert_fp32_acc_to_tf32_operandA_layout<1, RoundingTF32>(tDCrDC, tOrDC, tiled_mma, lane_id); - auto tOrDC_mma = [&] () { - if constexpr (!RoundingTF32) { - return recast(tOrDC); - } else { - return tOrDC; - } - }(); - - copy(A_tiled_copy, tOsAinv, tOrAinv_cv); - - clear(tOrO); - auto tOrAinv_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrAinv); - } else { - return tOrAinv; - } - }(); - gemm(tiled_mma, tOrDC_mma, tOrAinv_mma, tOrO); - - if constexpr (!RoundingTF32) { - // no need for conversion - copy(O_tiled_copy, tOrO_cv, tOsO); - } else { - auto tOrO_cv_cvt = make_tensor_like(tOrO_cv); - transform(tOrO_cv, tOrO_cv_cvt, [](auto v) { return Element(v); }); - copy(O_tiled_copy, tOrO_cv_cvt, tOsO); - } - } - - template - CUTE_DEVICE void - blockwise_diagonal_inversed_16x16_to_32x32(TensorT&& mat) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 32); - static_assert(size<1>(L) == 32); - - static_assert(is_contiguous<0>(L) == 1 || is_contiguous<1>(L) == 1); - // row-major - constexpr bool is_col_major = is_contiguous<0>(L); - - using TileShape = Shape<_16, _16, _16>; - auto mat_16x16_2x2 = flat_divide(std::forward(mat), select<0,1>(TileShape{})); - - using MMA = SM80_16x8x8_F32TF32TF32F32_TN; - using TiledMMA = decltype(make_tiled_mma(MMA{}, Layout>{}, TileShape{})); - - using CopyOpTF32 = AutoVectorizingCopyWithAssumedAlignment<128>; - using CopyOpD_S2R = CopyOpTF32; - using CopyOpC_S2R = CopyOpTF32; - using CopyOpA_S2R = CopyOpTF32; -#ifdef CUTE_ARCH_STSM_SM90_ENABLED - using CopyOpO_R2S = CopyOpTF32; -#else - using CopyOpO_R2S = UniversalCopy; -#endif - - int lane_id = cutlass::canonical_lane_idx(); - auto tiled_mma = TiledMMA{}; - auto thr_mma = tiled_mma.get_thread_slice(lane_id); - - auto D_tiled_copy = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto C_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto A_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto O_tiled_copy = make_tiled_copy_C(Copy_Atom{}, tiled_mma); - - auto D_thr_copy = D_tiled_copy.get_thread_slice(lane_id); - auto C_thr_copy = C_tiled_copy.get_thread_slice(lane_id); - auto A_thr_copy = A_tiled_copy.get_thread_slice(lane_id); - auto O_thr_copy = O_tiled_copy.get_thread_slice(lane_id); - - Tensor sDinv = mat_16x16_2x2(_, _, _1{}, _1{}); - Tensor sC = select_tensor<1, 0>(mat_16x16_2x2(_, _, _1{}, _0{})); - Tensor sAinv = select_tensor<1, 0>(mat_16x16_2x2(_, _, _0{}, _0{})); - Tensor sO = mat_16x16_2x2(_, _, _1{}, _0{}); - - Tensor tOrDinv = make_fragment_like(partition_shape_A(tiled_mma, select<0, 2>(TileShape{}))); - Tensor tOrC = make_fragment_like(partition_shape_B(tiled_mma, select<1, 2>(TileShape{}))); - Tensor tOrAinv = make_fragment_like(partition_shape_B(tiled_mma, select<1, 2>(TileShape{}))); - - Tensor tDCrDC = partition_fragment_C(tiled_mma, select<0,1>(TileShape{})); // output of -inv(D)C - Tensor tOrO = partition_fragment_C(tiled_mma, select<0,1>(TileShape{})); // output of -inv(D)C inv(A) - - Tensor tOsDinv = D_thr_copy.partition_S(sDinv); - Tensor tOrDinv_cv = D_thr_copy.retile_D(tOrDinv); - Tensor tOsC = C_thr_copy.partition_S(sC); - Tensor tOrC_cv = C_thr_copy.retile_D(tOrC); - Tensor tOsAinv = A_thr_copy.partition_S(sAinv); - Tensor tOrAinv_cv = A_thr_copy.retile_D(tOrAinv); - Tensor tOsO = O_thr_copy.partition_D(sO); - Tensor tOrO_cv = O_thr_copy.retile_S(tOrO); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C - copy(D_tiled_copy, tOsDinv, tOrDinv_cv); - copy(C_tiled_copy, tOsC, tOrC_cv); - - clear(tDCrDC); - auto tOrDinv_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrDinv); - } else { - return tOrDinv; - } - }(); - auto tOrC_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrC); - } else { - return tOrC; - } - }(); - gemm(tiled_mma, tOrDinv_mma, tOrC_mma, tDCrDC); - transform(tDCrDC, [](auto v) { return -v; }); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C inv(A) - Tensor tOrDC = make_fragment_like(partition_shape_A(tiled_mma, Shape<_16, _16>{})); - // Tensor tOrDC = detail::SM80::make_acc_into_op(tDCrDC, tiled_mma); - detail::SM80::convert_fp32_acc_to_tf32_operandA_layout<2, RoundingTF32>(tDCrDC, tOrDC, tiled_mma, lane_id); - auto tOrDC_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrDC); - } else { - return tOrDC; - } - }(); - - copy(A_tiled_copy, tOsAinv, tOrAinv_cv); - clear(tOrO); - - auto tOrAinv_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrAinv); - } else { - return tOrAinv; - } - }(); - gemm(tiled_mma, tOrDC_mma, tOrAinv_mma, tOrO); - - if constexpr (!RoundingTF32) { - // no need for conversion: MMA hardware truncates float→tf32 - copy(O_tiled_copy, tOrO_cv, tOsO); - } else { - // Explicit rounding: convert float→tfloat32_t before store to SMEM - auto tOrO_cv_cvt = make_tensor_like(tOrO_cv); - transform(tOrO_cv, tOrO_cv_cvt, [](auto v) { return Element(v); }); - copy(O_tiled_copy, tOrO_cv_cvt, tOsO); - } - } - - template - CUTE_DEVICE void - blockwise_diagonal_inversed_32x32_to_64x64(TensorT&& mat) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 64); - static_assert(size<1>(L) == 64); - - static_assert(is_contiguous<0>(L) == 1 || is_contiguous<1>(L) == 1); - constexpr bool is_col_major = is_contiguous<0>(L); - - auto mat_32x32_2x2 = flat_divide(std::forward(mat), select<0,1>(Shape<_32, _32>{})); - auto mat_16x2X16x2_2x2 = logical_divide(mat_32x32_2x2, Shape<_16, _16>{}); - - using MMA = SM80_16x8x8_F32TF32TF32F32_TN; - using TiledMMA1 = decltype(make_tiled_mma(MMA{}, Layout>{}, Shape<_16, _16, _32>{})); - using TiledMMA2 = decltype(make_tiled_mma(MMA{}, Layout>{}, Shape<_16, _32, _16>{})); - - using CopyOpTF32 = AutoVectorizingCopyWithAssumedAlignment<128>; - using CopyOpD_S2R = CopyOpTF32; - using CopyOpC_S2R = CopyOpTF32; - using CopyOpA_S2R = CopyOpTF32; - using CopyOpO_S2R = CopyOpTF32; -#ifdef CUTE_ARCH_STSM_SM90_ENABLED - using CopyOpO_R2S = CopyOpTF32; -#else - using CopyOpO_R2S = UniversalCopy; -#endif - - int warp_id_in_wg = cutlass::canonical_warp_idx() - cutlass::NumWarpsPerWarpGroup * cutlass::canonical_warp_group_idx(); - int x = warp_id_in_wg / 2; - int y = warp_id_in_wg % 2; - - int lane_id = cutlass::canonical_lane_idx(); - auto tiled_mma1 = TiledMMA1{}; - auto thr_mma1 = tiled_mma1.get_thread_slice(lane_id); - - auto tiled_mma2 = TiledMMA2{}; - auto thr_mma2 = tiled_mma2.get_thread_slice(lane_id); - - auto D_tiled_copy = make_tiled_copy_A(Copy_Atom{}, tiled_mma1); - auto C_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma1); - auto A_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma2); - auto O_tiled_s2r = make_tiled_copy_C(Copy_Atom{}, tiled_mma2); - auto O_tiled_r2s = make_tiled_copy_C(Copy_Atom{}, tiled_mma2); - - auto D_thr_copy = D_tiled_copy.get_thread_slice(lane_id); - auto C_thr_copy = C_tiled_copy.get_thread_slice(lane_id); - auto A_thr_copy = A_tiled_copy.get_thread_slice(lane_id); - auto O_thr_s2r = O_tiled_s2r.get_thread_slice(lane_id); - auto O_thr_r2s = O_tiled_r2s.get_thread_slice(lane_id); - - Tensor sDinv = mat_16x2X16x2_2x2(make_coord(_, y), _, _1{}, _1{}); - Tensor sC = select_tensor<1, 0>(mat_16x2X16x2_2x2(_, make_coord(_, x), _1{}, _0{})); - Tensor sAinv = select_tensor<1, 0>(mat_16x2X16x2_2x2(make_coord(_, x), _, _0{}, _0{})); // NOTE: not y! - Tensor sO = mat_16x2X16x2_2x2(make_coord(_, y), _, _1{}, _0{}); // needs cross-warp reduction - - Tensor tOrDinv = make_fragment_like(partition_shape_A(tiled_mma1, Shape<_16, _32>{})); - Tensor tOrC = make_fragment_like(partition_shape_B(tiled_mma1, Shape<_16, _32>{})); - Tensor tOrAinv = make_fragment_like(partition_shape_B(tiled_mma2, Shape<_32, _16>{})); - - Tensor tDCrDC = partition_fragment_C(tiled_mma1, Shape<_16, _16>{}); // output of -inv(D)C - Tensor tOrO = partition_fragment_C(tiled_mma2, Shape<_16, _32>{}); // output of -inv(D)C inv(A) - - Tensor tOsDinv = D_thr_copy.partition_S(sDinv); - Tensor tOrDinv_cv = D_thr_copy.retile_D(tOrDinv); - Tensor tOsC = C_thr_copy.partition_S(sC); - Tensor tOrC_cv = C_thr_copy.retile_D(tOrC); - Tensor tOsAinv = A_thr_copy.partition_S(sAinv); - Tensor tOrAinv_cv = A_thr_copy.retile_D(tOrAinv); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C - copy(D_tiled_copy, tOsDinv, tOrDinv_cv); - copy(C_tiled_copy, tOsC, tOrC_cv); - - clear(tDCrDC); - auto tOrDinv_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrDinv); - } else { - return tOrDinv; - } - }(); - auto tOrC_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrC); - } else { - return tOrC; - } - }(); - gemm(tiled_mma1, tOrDinv_mma, tOrC_mma, tDCrDC); - transform(tDCrDC, [](auto v) { return -v; }); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C inv(A) - Tensor tOrDC = make_fragment_like(partition_shape_A(tiled_mma2, Shape<_16, _16>{})); - detail::SM80::convert_fp32_acc_to_tf32_operandA_layout<2, RoundingTF32>(tDCrDC, tOrDC, tiled_mma2, lane_id); - auto tOrDC_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrDC); - } else { - return tOrDC; - } - }(); - - copy(A_tiled_copy, tOsAinv, tOrAinv_cv); - clear(tOrO); - - auto tOrAinv_mma = [&]() { - if constexpr (!RoundingTF32) { - return recast(tOrAinv); - } else { - return tOrAinv; - } - }(); - gemm(tiled_mma2, tOrDC_mma, tOrAinv_mma, tOrO); - - auto tOrO_cvt = [&]() { - if constexpr (!RoundingTF32) { - return tOrO; - } else { - return make_tensor_like(tOrO); - } - }(); - if constexpr (RoundingTF32) { - transform(tOrO, tOrO_cvt, [](auto v) { return Element(v); }); - } - - // ensure tOsC consumed, tOsC and tOsO are the same buffer - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - - Tensor tOsO = O_thr_r2s.partition_D(sO); - Tensor tOrO_cvt_cv = O_thr_r2s.retile_S(tOrO_cvt); - if (x == 0) { - copy(O_tiled_r2s, tOrO_cvt_cv, tOsO); - } - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - if (x == 1) { - Tensor tOrO_red = make_tensor_like(tOrO_cvt); - Tensor tOsO_s = O_thr_s2r.partition_S(sO); - Tensor tOrO_red_cv = O_thr_s2r.retile_D(tOrO_red); - copy(O_tiled_s2r, tOsO_s, tOrO_red_cv); - transform(tOrO_cvt, tOrO_red, tOrO_cvt, [](auto a, auto b) { return a + b; }); - copy(O_tiled_r2s, tOrO_cvt_cv, tOsO); - } - } - -private: - int wg_sync_named_barrier_id_; -}; - -// Adapted from https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/flat/ampere/collective/flat_collective_inverse.hpp -template -struct CollectiveInverse { - // FIXME: precision is not good due to half - static_assert(std::is_same_v || std::is_same_v, "only half is implemented"); - - CUTE_DEVICE - CollectiveInverse(int wg_sync_named_barrier_id): wg_sync_named_barrier_id_(wg_sync_named_barrier_id) {} - - template - CUTE_DEVICE void - compute(TensorT&& sT) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 64); - static_assert(size<1>(L) == 64); - - int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; - - auto t8X8sT = flat_divide(sT, Shape<_8, _8>{}); - if (thread_idx < 64) { // compute 8x8 inverse on diagnal directly - compute_diagonal_inverse_NxN<8>(t8X8sT(_, _, thread_idx / 8, thread_idx / 8), thread_idx % 8); - } - - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - - auto t16X16sT = flat_divide(sT, Shape<_16, _16>{}); - // four warps for 8x8 -> 16x16 - blockwise_diagonal_inversed_8x8_to_16x16(t16X16sT(_, _, thread_idx / 32, thread_idx / 32)); - - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - - auto t32X32sT = flat_divide(sT, Shape<_32, _32>{}); - if (thread_idx < 64) { // two warps for 16x16 -> 32x32 - blockwise_diagonal_inversed_16x16_to_32x32(t32X32sT(_, _, thread_idx / 32, thread_idx /32)); - } - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - - // one warpgroup for 32x32 -> 64x64 - blockwise_diagonal_inversed_32x32_to_64x64(sT); - } - -private: - - template - CUTE_DEVICE void - compute_diagonal_inverse_NxN(TensorT&& mat, int tid_in_group) { // group_size = N - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == N); - static_assert(size<1>(L) == N); - - using ElementCompute = float; - - using CopyOp = Copy_Atom, Element>; - - auto load_row = [&](int y) { - auto row = make_tensor(Shape>{}); - copy(CopyOp{}, std::forward(mat)(y, _), row); - - auto row_cvt = make_tensor_like(row); - copy(row, row_cvt); - - if constexpr (GarbageFilledDiagonal || GarbageFilledUpperTriangular) { - CUTE_UNROLL - for (int i = 0; i < N; ++i) { - row_cvt(i) = i == y ? 1.0f - : (i > y ? 0.0f - : row_cvt(i)); - } - } - return row_cvt; - }; - - auto store_row = [&](int y, auto row) { - auto row_cvt = make_tensor_like(row); - copy(row, row_cvt); - copy(CopyOp{}, row_cvt, std::forward(mat)(y, _)); - }; - - auto row = load_row(tid_in_group); -#define LOAD(y, x) __shfl_sync(0xffffffff, row(x), y, N) - - CUTE_UNROLL - for (int src_row = 0; src_row < N - 1; ++src_row) { // idx of src row to eliminate - auto row_scale = -row(src_row); // scale the src row - CUTE_UNROLL - for (int i = 0; i < src_row; ++i) { - auto src_row_value = LOAD(src_row, i); - row(i) = tid_in_group > src_row ? row_scale * src_row_value + row(i) : row(i); - } - row(src_row) = tid_in_group > src_row ? row_scale : row(src_row); - } - -#undef LOAD - - store_row(tid_in_group, row); - } - - /* - blockwise inverse has relation as follows - inv(| A 0 |) = | inv(A) 0 | - | C D | | -inv(D)C inv(A) inv(D) | - */ - - template - CUTE_DEVICE void - blockwise_diagonal_inversed_4x4_to_8x8(TensorT&& mat) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 8); - static_assert(size<1>(L) == 8); - auto mat_NxN_2x2 = flat_divide(std::forward(mat), Shape<_4, _4>{}); - - // FIXME: implement - } - - template - CUTE_DEVICE void - blockwise_diagonal_inversed_8x8_to_16x16(TensorT&& mat) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 16); - static_assert(size<1>(L) == 16); - - static_assert(is_contiguous<0>(L) == 1 || is_contiguous<1>(L) == 1); - constexpr bool is_col_major = is_contiguous<0>(L); - - auto mat_8x8_2x2 = flat_divide(std::forward(mat), Shape<_8, _8>{}); - using MMA = SM80_16x8x8_F32F16F16F32_TN; - using TiledMMA = decltype(make_tiled_mma(MMA{}, Layout>{}, Shape<_16, _8, _8>{})); - - using CopyOpD_S2R = std::conditional_t; - using CopyOpC_S2R = std::conditional_t; - using CopyOpA_S2R = std::conditional_t; -#ifdef CUTE_ARCH_STSM_SM90_ENABLED - using CopyOpO_R2S = std::conditional_t; -#else - using CopyOpO_R2S = UniversalCopy; -#endif - - int lane_id = cutlass::canonical_lane_idx(); - auto tiled_mma = TiledMMA{}; - auto thr_mma = tiled_mma.get_thread_slice(lane_id); - - auto D_tiled_copy = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto C_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto A_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto O_tiled_copy = make_tiled_copy_C(Copy_Atom{}, tiled_mma); - - auto D_thr_copy = D_tiled_copy.get_thread_slice(lane_id); - auto C_thr_copy = C_tiled_copy.get_thread_slice(lane_id); - auto A_thr_copy = A_tiled_copy.get_thread_slice(lane_id); - auto O_thr_copy = O_tiled_copy.get_thread_slice(lane_id); - - Tensor sDinv = mat_8x8_2x2(_, _, _1{}, _1{}); - Tensor sC = select_tensor<1, 0>(mat_8x8_2x2(_, _, _1{}, _0{})); - Tensor sAinv = select_tensor<1, 0>(mat_8x8_2x2(_, _, _0{}, _0{})); - Tensor sO = mat_8x8_2x2(_, _, _1{}, _0{}); - - Tensor sDinv_m_bcast = make_tensor(sDinv.data(), logical_product(sDinv.layout(), Tile>{})); - Tensor sO_m_bcast = make_tensor(sO.data(), logical_product(sO.layout(), Tile>{})); - - Tensor tOrDinv = make_fragment_like(partition_shape_A(tiled_mma, Shape<_16, _8>{})); - Tensor tOrC = thr_mma.partition_fragment_B(sC); - Tensor tOrAinv = thr_mma.partition_fragment_B(sAinv); - - Tensor tDCrDC = partition_fragment_C(tiled_mma, Shape<_16, _8>{}); // output of -inv(D)C - Tensor tOrO = partition_fragment_C(tiled_mma, Shape<_16, _8>{}); // output of -inv(D)C inv(A) - - Tensor tOsDinv = D_thr_copy.partition_S(sDinv_m_bcast); - Tensor tOrDinv_cv = D_thr_copy.retile_D(tOrDinv); - Tensor tOsC = C_thr_copy.partition_S(sC); - Tensor tOrC_cv = C_thr_copy.retile_D(tOrC); - Tensor tOsAinv = A_thr_copy.partition_S(sAinv); - Tensor tOrAinv_cv = A_thr_copy.retile_D(tOrAinv); - Tensor tOsO = O_thr_copy.partition_D(sO_m_bcast); - Tensor tOrO_cv = O_thr_copy.retile_S(tOrO); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C - copy(D_tiled_copy, tOsDinv(make_coord(_, _0{}), _, _), tOrDinv_cv(make_coord(_, _0{}), _, _)); - copy(C_tiled_copy, tOsC, tOrC_cv); - - clear(tDCrDC); - gemm(tiled_mma, tOrDinv, tOrC, tDCrDC); - transform(tDCrDC(make_coord(_, _0{}), _, _), [](auto v) { return -v; }); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C inv(A) - Tensor tOrDC = detail::SM80::make_acc_into_op(tDCrDC, tiled_mma); - - copy(A_tiled_copy, tOsAinv, tOrAinv_cv); - - clear(tOrO); - gemm(tiled_mma, tOrDC, tOrAinv, tOrO); - - auto tOrO_cv_cvt = make_tensor_like(tOrO_cv(make_coord(_, _0{}), _, _)); - transform(tOrO_cv(make_coord(_, _0{}), _, _), tOrO_cv_cvt, [](auto v) { return Element(v); }); - copy(O_tiled_copy, tOrO_cv_cvt, tOsO(make_coord(_, _0{}), _, _)); - } - - template - CUTE_DEVICE void - blockwise_diagonal_inversed_16x16_to_32x32(TensorT&& mat) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 32); - static_assert(size<1>(L) == 32); - - static_assert(is_contiguous<0>(L) == 1 || is_contiguous<1>(L) == 1); - constexpr bool is_col_major = is_contiguous<0>(L); - - using TileShape = Shape<_16, _16, _16>; - auto mat_16x16_2x2 = flat_divide(std::forward(mat), select<0,1>(TileShape{})); - - using MMA = SM80_16x8x16_F32F16F16F32_TN; - using TiledMMA = decltype(make_tiled_mma(MMA{}, Layout>{}, TileShape{})); - - using CopyOpD_S2R = std::conditional_t; - using CopyOpC_S2R = std::conditional_t; - using CopyOpA_S2R = std::conditional_t; -#ifdef CUTE_ARCH_STSM_SM90_ENABLED - using CopyOpO_R2S = std::conditional_t; -#else - using CopyOpO_R2S = UniversalCopy; -#endif - - int lane_id = cutlass::canonical_lane_idx(); - auto tiled_mma = TiledMMA{}; - auto thr_mma = tiled_mma.get_thread_slice(lane_id); - - auto D_tiled_copy = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto C_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto A_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto O_tiled_copy = make_tiled_copy_C(Copy_Atom{}, tiled_mma); - - auto D_thr_copy = D_tiled_copy.get_thread_slice(lane_id); - auto C_thr_copy = C_tiled_copy.get_thread_slice(lane_id); - auto A_thr_copy = A_tiled_copy.get_thread_slice(lane_id); - auto O_thr_copy = O_tiled_copy.get_thread_slice(lane_id); - - Tensor sDinv = mat_16x16_2x2(_, _, _1{}, _1{}); - Tensor sC = select_tensor<1, 0>(mat_16x16_2x2(_, _, _1{}, _0{})); - Tensor sAinv = select_tensor<1, 0>(mat_16x16_2x2(_, _, _0{}, _0{})); - Tensor sO = mat_16x16_2x2(_, _, _1{}, _0{}); - - Tensor tOrDinv = thr_mma.partition_fragment_A(sDinv); - Tensor tOrC = thr_mma.partition_fragment_B(sC); - Tensor tOrAinv = thr_mma.partition_fragment_B(sAinv); - - Tensor tDCrDC = partition_fragment_C(tiled_mma, select<0,1>(TileShape{})); // output of -inv(D)C - Tensor tOrO = partition_fragment_C(tiled_mma, select<0,1>(TileShape{})); // output of -inv(D)C inv(A) - - Tensor tOsDinv = D_thr_copy.partition_S(sDinv); - Tensor tOrDinv_cv = D_thr_copy.retile_D(tOrDinv); - Tensor tOsC = C_thr_copy.partition_S(sC); - Tensor tOrC_cv = C_thr_copy.retile_D(tOrC); - Tensor tOsAinv = A_thr_copy.partition_S(sAinv); - Tensor tOrAinv_cv = A_thr_copy.retile_D(tOrAinv); - Tensor tOsO = O_thr_copy.partition_D(sO); - Tensor tOrO_cv = O_thr_copy.retile_S(tOrO); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C - copy(D_tiled_copy, tOsDinv, tOrDinv_cv); - copy(C_tiled_copy, tOsC, tOrC_cv); - - clear(tDCrDC); - gemm(tiled_mma, tOrDinv, tOrC, tDCrDC); - transform(tDCrDC, [](auto v) { return -v; }); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C inv(A) - Tensor tOrDC = detail::SM80::make_acc_into_op(tDCrDC, tiled_mma); - - copy(A_tiled_copy, tOsAinv, tOrAinv_cv); - clear(tOrO); - gemm(tiled_mma, tOrDC, tOrAinv, tOrO); - - auto tOrO_cv_cvt = make_tensor_like(tOrO_cv); - transform(tOrO_cv, tOrO_cv_cvt, [](auto v) { return Element(v); }); - copy(O_tiled_copy, tOrO_cv_cvt, tOsO); - } - - template - CUTE_DEVICE void - blockwise_diagonal_inversed_32x32_to_64x64(TensorT&& mat) { - constexpr auto L = typename std::remove_const_t>::layout_type{}; - static_assert(rank(L) == 2); - static_assert(size<0>(L) == 64); - static_assert(size<1>(L) == 64); - - static_assert(is_contiguous<0>(L) == 1 || is_contiguous<1>(L) == 1); - constexpr bool is_col_major = is_contiguous<0>(L); - - auto mat_32x32_2x2 = flat_divide(std::forward(mat), select<0,1>(Shape<_32, _32>{})); - auto mat_16x2X16x2_2x2 = logical_divide(mat_32x32_2x2, Shape<_16, _16>{}); - - using MMA = SM80_16x8x16_F32F16F16F32_TN; - using TiledMMA1 = decltype(make_tiled_mma(MMA{}, Layout>{}, Shape<_16, _16, _32>{})); - using TiledMMA2 = decltype(make_tiled_mma(MMA{}, Layout>{}, Shape<_16, _32, _16>{})); - - using CopyOpD_S2R = std::conditional_t; - using CopyOpC_S2R = std::conditional_t; - using CopyOpA_S2R = std::conditional_t; - using CopyOpO_S2R = std::conditional_t; - using CopyOpO_S2R = std::conditional_t; -#ifdef CUTE_ARCH_STSM_SM90_ENABLED - using CopyOpO_R2S = std::conditional_t; -#else - using CopyOpO_R2S = UniversalCopy; -#endif - - int warp_id_in_wg = cutlass::canonical_warp_idx() - cutlass::NumWarpsPerWarpGroup * cutlass::canonical_warp_group_idx(); - int x = warp_id_in_wg / 2; - int y = warp_id_in_wg % 2; - - int lane_id = cutlass::canonical_lane_idx(); - auto tiled_mma1 = TiledMMA1{}; - auto thr_mma1 = tiled_mma1.get_thread_slice(lane_id); - - auto tiled_mma2 = TiledMMA2{}; - auto thr_mma2 = tiled_mma2.get_thread_slice(lane_id); - - auto D_tiled_copy = make_tiled_copy_A(Copy_Atom{}, tiled_mma1); - auto C_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma1); - auto A_tiled_copy = make_tiled_copy_B(Copy_Atom{}, tiled_mma2); - auto O_tiled_s2r = make_tiled_copy_C(Copy_Atom{}, tiled_mma2); - auto O_tiled_r2s = make_tiled_copy_C(Copy_Atom{}, tiled_mma2); - - auto D_thr_copy = D_tiled_copy.get_thread_slice(lane_id); - auto C_thr_copy = C_tiled_copy.get_thread_slice(lane_id); - auto A_thr_copy = A_tiled_copy.get_thread_slice(lane_id); - auto O_thr_s2r = O_tiled_s2r.get_thread_slice(lane_id); - auto O_thr_r2s = O_tiled_r2s.get_thread_slice(lane_id); - - Tensor sDinv = mat_16x2X16x2_2x2(make_coord(_, y), _, _1{}, _1{}); - Tensor sC = select_tensor<1, 0>(mat_16x2X16x2_2x2(_, make_coord(_, x), _1{}, _0{})); - Tensor sAinv = select_tensor<1, 0>(mat_16x2X16x2_2x2(make_coord(_, x), _, _0{}, _0{})); // NOTE: not y! - Tensor sO = mat_16x2X16x2_2x2(make_coord(_, y), _, _1{}, _0{}); // needs cross-warp reduction - - Tensor tOrDinv = thr_mma1.partition_fragment_A(sDinv); - Tensor tOrC = thr_mma1.partition_fragment_B(sC); - Tensor tOrAinv = thr_mma2.partition_fragment_B(sAinv); - - Tensor tDCrDC = partition_fragment_C(tiled_mma1, Shape<_16, _16>{}); // output of -inv(D)C - Tensor tOrO = partition_fragment_C(tiled_mma2, Shape<_16, _32>{}); // output of -inv(D)C inv(A) - - Tensor tOsDinv = D_thr_copy.partition_S(sDinv); - Tensor tOrDinv_cv = D_thr_copy.retile_D(tOrDinv); - Tensor tOsC = C_thr_copy.partition_S(sC); - Tensor tOrC_cv = C_thr_copy.retile_D(tOrC); - Tensor tOsAinv = A_thr_copy.partition_S(sAinv); - Tensor tOrAinv_cv = A_thr_copy.retile_D(tOrAinv); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C - copy(D_tiled_copy, tOsDinv, tOrDinv_cv); - copy(C_tiled_copy, tOsC, tOrC_cv); - - clear(tDCrDC); - gemm(tiled_mma1, tOrDinv, tOrC, tDCrDC); - transform(tDCrDC, [](auto v) { return -v; }); - - ///////////////////////////////////////////////////////////////////////////// - // -inv(D)C inv(A) - Tensor tOrDC = detail::SM80::make_acc_into_op(tDCrDC, tiled_mma2); - - copy(A_tiled_copy, tOsAinv, tOrAinv_cv); - clear(tOrO); - gemm(tiled_mma2, tOrDC, tOrAinv, tOrO); - - auto tOrO_cvt = make_tensor_like(tOrO); - transform(tOrO, tOrO_cvt, [](auto v) { return Element(v); }); - - // ensure tOsC consumed, tOsC and tOsO are the same buffer - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - - Tensor tOsO = O_thr_r2s.partition_D(sO); - Tensor tOrO_cvt_cv = O_thr_r2s.retile_S(tOrO_cvt); - if (x == 0) { - copy(O_tiled_r2s, tOrO_cvt_cv, tOsO); - } - cutlass::arch::NamedBarrier::arrive_and_wait(cutlass::NumThreadsPerWarpGroup, wg_sync_named_barrier_id_); - if (x == 1) { - Tensor tOrO_red = make_tensor_like(tOrO_cvt); - Tensor tOsO_s = O_thr_s2r.partition_S(sO); - Tensor tOrO_red_cv = O_thr_s2r.retile_D(tOrO_red); - copy(O_tiled_s2r, tOsO_s, tOrO_red_cv); - transform(tOrO_cvt, tOrO_red, tOrO_cvt, [](auto a, auto b) { return a + b; }); - copy(O_tiled_r2s, tOrO_cvt_cv, tOsO); - } - } - -private: - int wg_sync_named_barrier_id_; -}; - -} // namespace kerutils diff --git a/csrc/kerutils/include/kerutils/device/sm80/helpers.cuh b/csrc/kerutils/include/kerutils/device/sm80/helpers.cuh deleted file mode 100644 index a8ce452d..00000000 --- a/csrc/kerutils/include/kerutils/device/sm80/helpers.cuh +++ /dev/null @@ -1,123 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/****************************************************************************** - * Copyright (c) 2024, Tri Dao. - ******************************************************************************/ - -#pragma once - -#include - -#include "kerutils/device/common.h" - -namespace kerutils { - -// ============================================================ -// Global / shared memory store intrinsics -// ============================================================ - -CUTE_DEVICE -void _store_256b( - uint32_t const& src0, uint32_t const& src1, uint32_t const& src2, uint32_t const& src3, - uint32_t const& src4, uint32_t const& src5, uint32_t const& src6, uint32_t const& src7, - void *gmem_addr -) { - asm volatile("st.global.L1::no_allocate.v8.f32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n" - :: "l"(gmem_addr), "r"(src0), "r"(src1), "r"(src2), "r"(src3), "r"(src4), "r"(src5), "r"(src6), "r"(src7)); -} - -// reference: https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/copy_sm100.hpp#L70 -CUTE_DEVICE -void store_256b(void *src, void *dst) { - uint32_t *src_ptr = reinterpret_cast(src); - _store_256b( - src_ptr[0], src_ptr[1], src_ptr[2], src_ptr[3], src_ptr[4], src_ptr[5], src_ptr[6], src_ptr[7], dst); -} - -template -CUTE_DEVICE -void store_128b(void* smem_ptr, const T &data) { - static_assert(sizeof(T) == 16); - *(__int128*)smem_ptr = *(__int128*)&data; -} - -// Predicated copy helper (generic, works on SM80+) -// Supports predication on both MN and K dimensions with optional OOB clearing. -template -CUTLASS_DEVICE void copy_pred(cute::TiledCopy const &tiled_copy, cute::Tensor const &S, - cute::Tensor &D, cute::Tensor const &identity_MN, - cute::Tensor const &predicate_K, const int max_MN=0) { - using namespace cute; - // Decay TiledCopy to CopyAtom - auto copy_atom = static_cast(tiled_copy); - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K - // There's no case where !Clear_OOB_K && Clear_OOB_MN - static_assert(!(Clear_OOB_MN && !Clear_OOB_K)); - auto has_with_bool = cute::is_valid([](auto t)->void_t().with(true))>{}, copy_atom); - #pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - bool predicate_mn = Is_even_MN || get<0>(identity_MN(_0{}, m, _0{})) < max_MN; - // NOTE: currently only this predicate is true because we set Clear_OOB_MN=false - if constexpr (Is_even_MN || !Clear_OOB_MN) { - if (Is_even_MN || predicate_mn) { - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - if constexpr (Is_even_K || !Clear_OOB_K) { - if (Is_even_K || predicate_K(k)) { cute::copy(copy_atom, S(_, m, k), D(_, m, k)); } - } else { // Clear_OOB_K == true && Is_even_K == false - // If copy traits can be transformed with a predicate value, do it, otherwise branch here - if constexpr (has_with_bool) { - cute::copy(copy_atom.with(predicate_K(k)), S(_, m, k), D(_, m, k)); - } else { - if (predicate_K(k)) { - cute::copy(copy_atom, S(_, m, k), D(_, m, k)); - } else { - cute::clear(D(_, m, k)); - } - } - } - } - } - } else { // Clear_OOB_MN == true && Is_even_MN == false, also implies Clear_OOB_K == true - if constexpr (!has_with_bool) { - if (predicate_mn) { - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - if (Is_even_K || predicate_K(k)) { - cute::copy(copy_atom, S(_, m, k), D(_, m, k)); - } else if (Clear_OOB_K) { - cute::clear(D(_, m, k)); - } - } - } else { - cute::clear(D(_, m, _)); - } - } else { // combine the mn predicate with the k predicate - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - cute::copy(copy_atom.with(predicate_mn && (Is_even_K || predicate_K(k))), S(_, m, k), D(_, m, k)); - } - } - } - } -} - -} // namespace kerutils diff --git a/csrc/kerutils/include/kerutils/device/sm90/helpers.cuh b/csrc/kerutils/include/kerutils/device/sm90/helpers.cuh deleted file mode 100644 index 233b69c5..00000000 --- a/csrc/kerutils/include/kerutils/device/sm90/helpers.cuh +++ /dev/null @@ -1,151 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 DeepSeek. All Rights Reserved. - * - * Licensed under the MIT License. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ - -#pragma once - -#include - -#include "kerutils/device/common.h" - -// Adapted from https://github.com/deepseek-ai/FlashMLA/blob/main/csrc/kerutils/include/kerutils/device/sm90/helpers.cuh -namespace kerutils { - -// TMA copy helper with optional multicast support -template< - typename TMA, - typename Tensor0, - typename Tensor1 -> -CUTE_DEVICE -void launch_tma_copy( - const TMA &tma_copy, - Tensor0 src, - Tensor1 dst, - uint64_t &bar, - const cute::TMA::CacheHintSm90 &cache_hint = cute::TMA::CacheHintSm90::EVICT_NORMAL, - const uint16_t &multicast_mask = 0 -) { - auto thr_tma = tma_copy.get_slice(cute::_0{}); - cute::copy( - tma_copy.with(bar, multicast_mask, cache_hint), - thr_tma.partition_S(src), - thr_tma.partition_D(dst) - ); -} - -// In the layout of fragment A and fragment C during WGMMA, data each thread holds resides in two particular rows. This function converts the local_row_idx (0~2) to the actual row_idx -// You may refer to this link for the detailed layout: https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-a -CUTE_DEVICE -int get_AorC_row_idx(int local_row_idx, int idx_in_warpgroup) { - int row_idx = (idx_in_warpgroup/32)*16 + local_row_idx*8 + (idx_in_warpgroup%32/4); - return row_idx; -} - -// In the layout of fragment A and fragment C during WGMMA, data each thread holds resides in some rows. This function converts the local_elem_idx to the actual col_idx -// You may refer to this link for the detailed layout: https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-a -CUTE_DEVICE -int get_AorC_col_idx(int local_elem_idx, int idx_in_warpgroup) { - int col_idx = 8*(local_elem_idx/4) + (idx_in_warpgroup%4)*2 + (local_elem_idx&1); - return col_idx; -} - -template -CUTE_DEVICE -void wgmma(TiledMma &tiled_mma, Tensor0 const &tCrA, Tensor1 const &tCrB, Tensor2 &tCrC, bool zero_init) { - using namespace cute; - constexpr bool Is_RS = !cute::is_base_of::value; - // Need to cast away const on tCrA since warpgroup_fence_operand doesn't take const - if constexpr (Is_RS) { cute::warpgroup_fence_operand(const_cast(tCrA)); } - warpgroup_fence_operand(tCrC); - warpgroup_arrive(); - tiled_mma.accumulate_ = zero_init ? GMMA::ScaleOut::Zero : GMMA::ScaleOut::One; - // Unroll the K mode manually to set scale D to 1 - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { - cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - if constexpr (commit) { - warpgroup_commit_batch(); - } - warpgroup_fence_operand(tCrC); - if constexpr (Is_RS) { warpgroup_fence_operand(const_cast(tCrA)); } -} - -template -CUTE_DEVICE -void wgmma_ss(bool clear_accum, TiledMma tiled_mma, Tensor0 const &sA, Tensor1 const &sB, Tensor2 &rC_frag, int idx_in_warpgroup) { - using namespace cute; - ThrMMA thr_mma = tiled_mma.get_slice(idx_in_warpgroup); - Tensor sA_frag = thr_mma.partition_fragment_A(sA); - Tensor sB_frag = thr_mma.partition_fragment_B(sB); - static_assert(size<2>(sA_frag) == size<2>(sB_frag)); - - warpgroup_fence_operand(rC_frag); - warpgroup_arrive(); - tiled_mma.accumulate_ = clear_accum ? GMMA::ScaleOut::Zero : GMMA::ScaleOut::One; - CUTLASS_PRAGMA_UNROLL - for (int k = 0; k < size<2>(sA_frag); ++k) { - cute::gemm(tiled_mma, sA_frag(_, _, k), sB_frag(_, _, k), rC_frag); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - warpgroup_fence_operand(rC_frag); -} - -template -CUTE_DEVICE -void wgmma_rs(bool clear_accum, TiledMma tiled_mma, Tensor0 rA_frag, Tensor1 const &sB, Tensor2 &rC_frag, int idx_in_warpgroup) { - using namespace cute; - ThrMMA thr_mma = tiled_mma.get_slice(idx_in_warpgroup); - Tensor sB_frag = thr_mma.partition_fragment_B(sB); - static_assert(size<2>(rA_frag) == size<2>(sB_frag)); - - warpgroup_fence_operand(const_cast(rA_frag)); - warpgroup_fence_operand(rC_frag); - warpgroup_arrive(); - tiled_mma.accumulate_ = clear_accum ? GMMA::ScaleOut::Zero : GMMA::ScaleOut::One; - CUTLASS_PRAGMA_UNROLL - for (int k = 0; k < size<2>(rA_frag); ++k) { - cute::gemm(tiled_mma, rA_frag(_, _, k), sB_frag(_, _, k), rC_frag); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - } - warpgroup_fence_operand(rC_frag); - warpgroup_fence_operand(const_cast(rA_frag)); -} - -} // namespace kerutils diff --git a/csrc/kerutils/include/kerutils/host/host.h b/csrc/kerutils/include/kerutils/host/host.h deleted file mode 100644 index fc8b61c9..00000000 --- a/csrc/kerutils/include/kerutils/host/host.h +++ /dev/null @@ -1,107 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -/* - * Copyright (c) 2025 DeepSeek. All Rights Reserved. - * - * Licensed under the MIT License. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ - -#pragma once - -#include -#include -#include - -#include - -#include "kerutils/common/common.h" -#include "kerutils/common/static_switch.h" - -namespace kerutils { - -class KUException final : public std::exception { - std::string message = {}; - -public: - template - explicit KUException(const char *name, const char* file, const int line, Args&&... args) { - std::ostringstream oss; - - oss << name << " error (" << file << ":" << line << "): "; - (oss << ... << args); - message = oss.str(); - } - - const char *what() const noexcept override { - return message.c_str(); - } -}; - -#define THROW_KU_EXCEPTION(name, ...) \ - throw kerutils::KUException(name, __FILE__, __LINE__, __VA_ARGS__) - -#define KU_CUDA_CHECK(call) \ -do { \ - cudaError_t status_ = call; \ - if (status_ != cudaSuccess) { \ - fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, cudaGetErrorString(status_)); \ - THROW_KU_EXCEPTION("CUDA", "CUDA error: ", cudaGetErrorString(status_)); \ - } \ -} while(0) - -// This `KU_ASSERT` is triggered no matter if the code is compiled with `-DNDEBUG` or not. -#define KU_ASSERT(cond, ...) \ - do { \ - if (not (cond)) { \ - fprintf(stderr, "Assertion `%s` failed (%s:%d): ", #cond, __FILE__, __LINE__); \ - if constexpr (sizeof(#__VA_ARGS__) > 1) { \ - fprintf(stderr, ", " __VA_ARGS__); \ - } \ - fprintf(stderr, "\n"); \ - THROW_KU_EXCEPTION("Assertion", "Assertion `", #cond, "` failed."); \ - } \ - } while(0) - -#define KU_CHECK_KERNEL_LAUNCH() KU_CUDA_CHECK(cudaGetLastError()) - -template -inline __host__ __device__ constexpr T ceil_div(const T &a, const T &b) { - return (a + b - 1) / b; -} - -template -inline __host__ __device__ constexpr T ceil(const T &a, const T &b) { - return (a + b - 1) / b * b; -} - -} // namespace kerutils diff --git a/csrc/kerutils/include/kerutils/kerutils.cuh b/csrc/kerutils/include/kerutils/kerutils.cuh deleted file mode 100644 index a0daab01..00000000 --- a/csrc/kerutils/include/kerutils/kerutils.cuh +++ /dev/null @@ -1,18 +0,0 @@ -// Copyright 2025-2026 Ant Group Co., Ltd. -// -// 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. - -#pragma once - -#include "host/host.h" -#include "device/device.cuh" diff --git a/cula/backends.py b/cula/backends.py deleted file mode 100644 index 18fa37ed..00000000 --- a/cula/backends.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import logging -import os -from typing import ClassVar - -logger = logging.getLogger("cula.backends") - - -class BaseBackend: - backend_type: ClassVar[str] = "base" - env_var: ClassVar[str | None] = None - default_enable: ClassVar[bool] = True - priority: ClassVar[int] = 5 - - _available: tuple[bool, str | None] | None = None - - def probe(self) -> tuple[bool, str | None]: - return True, None - - def is_available(self) -> tuple[bool, str | None]: - if self._available is None: - self._available = self.probe() - return self._available - - def is_enabled(self) -> bool: - if self.env_var is None: - return True - default = "1" if self.default_enable else "0" - return os.environ.get(self.env_var, default) != "0" - - def verify(self, func_name: str, *args, **kwargs) -> tuple[bool, str | None]: - verifier = getattr(self, f"{func_name}_verifier", None) - if verifier is None: - return True, None - return verifier(*args, **kwargs) - - -class BackendRegistry: - def __init__(self, operation: str): - self.operation = operation - self._backends: list[BaseBackend] = [] - self._logged: set[str] = set() - - def register(self, backend: BaseBackend) -> None: - self._backends.append(backend) - self._backends.sort(key=lambda b: b.priority) - - def backends(self) -> list[BaseBackend]: - return list(self._backends) - - def _log_once(self, key: str, msg: str) -> None: - if key not in self._logged: - self._logged.add(key) - logger.info(msg) - - def dispatch(self, func_name: str, *args, **kwargs): - rejections: list[tuple[str, str]] = [] - for be in self._backends: - name = be.backend_type - if not be.is_enabled(): - rejections.append((name, f"disabled via {be.env_var}=0")) - continue - ok, why = be.is_available() - if not ok: - rejections.append((name, why or "unavailable")) - continue - ok, why = be.verify(func_name, *args, **kwargs) - if not ok: - rejections.append((name, why or "rejected")) - self._log_once( - f"{self.operation}:{func_name}:{name}:reject", - f"[cula backend] {self.operation}.{func_name}: {name} rejected: {why}", - ) - continue - self._log_once( - f"{self.operation}:{func_name}:{name}", - f"[cula backend] {self.operation}.{func_name} -> {name}", - ) - return getattr(be, func_name)(*args, **kwargs) - - table = "\n".join(f" {n:<12s}: {r}" for n, r in rejections) or " (no backends registered)" - raise RuntimeError(f"no {self.operation} backend accepted this call:\n{table}") diff --git a/cula/cudac.py b/cula/cudac.py deleted file mode 100644 index 68caed9b..00000000 --- a/cula/cudac.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# 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. - -"""Unified interface to per-architecture CUDA extensions. - -This module exposes the remaining SM90 native extension. SM100/SM103 KDA -forward uses the public CuTeDSL path and no longer ships a CUDA extension. - -Loading is **once per process**: the first attribute access checks the -currently active CUDA device, imports ``cula._cudac_sm90`` when applicable, -and caches the discovered callables on the module instance. -Changing the active CUDA device to a different architecture after a -process has already loaded ``cula.cudac`` will therefore not be picked -up -- callers that need a different extension must restart Python. -""" - -import importlib -import sys -import threading -from types import ModuleType - - -def _current_device_extension() -> tuple[str, str]: - try: - import torch - except ImportError as exc: - raise ImportError("cuLA CUDA extensions require PyTorch to detect the current GPU.") from exc - - if not torch.cuda.is_available(): - raise RuntimeError("cuLA CUDA extensions require a visible CUDA GPU, but torch.cuda.is_available() is False.") - - device = torch.cuda.current_device() - prop = torch.cuda.get_device_properties(device) - sm_label = f"sm_{prop.major}{prop.minor}" - if prop.major == 9 and prop.minor == 0: - return "cula._cudac_sm90", sm_label - if prop.major == 10 and prop.minor in (0, 3): - raise RuntimeError( - f"{sm_label} KDA kernels use CuTeDSL and do not provide a cula.cudac extension; " - "call the public cula.kda APIs instead." - ) - raise RuntimeError(f"Unsupported CUDA compute capability {sm_label}. The native extension supports sm_90 only.") - - -class _CudacProxy(ModuleType): - """Lazy proxy that exposes functions from the current GPU arch extension.""" - - def __init__(self): - super().__init__(__name__) - self.__path__ = [] - self._modules_loaded = False - self._funcs: dict[str, object] = {} - self._lock = threading.Lock() - - def _load(self): - if self._modules_loaded: - return - with self._lock: - if self._modules_loaded: - return - ext_name, sm_label = _current_device_extension() - try: - mod = importlib.import_module(ext_name) - for attr in dir(mod): - if not attr.startswith("_"): - self._funcs[attr] = getattr(mod, attr) - except (ImportError, AttributeError, OSError) as exc: - raise ImportError( - f"The cuLA CUDA extension for the current GPU ({sm_label}) could not be imported. " - f"Extension {ext_name} failed with: {exc}. " - "Please make sure cuLA is compiled correctly." - ) from exc - self.__dict__.update(self._funcs) - self._modules_loaded = True - - def __getattr__(self, name: str): - if name.startswith("_"): - raise AttributeError(name) - self._load() - try: - return self._funcs[name] - except KeyError: - raise AttributeError(f"module 'cula.cudac' has no attribute '{name}'") from None - - def __dir__(self): - self._load() - return list(self._funcs.keys()) - - -_proxy = _CudacProxy() -_proxy.__dict__.update({k: globals().get(k) for k in ("__spec__", "__file__", "__package__", "__loader__")}) -sys.modules[__name__] = _proxy diff --git a/cula/kda/README.md b/cula/kda/README.md index 928e372b..f326e96b 100644 --- a/cula/kda/README.md +++ b/cula/kda/README.md @@ -7,14 +7,14 @@ Public entry points for KDA. For full examples (varlen, intracard CP, per-backen | Symbol | Description | Arch | Direction | |--------|-------------|------|-----------| | `chunk_kda` | Chunked KDA prefill | SM100 (Blackwell) | Forward + backward | -| `kda_prefill_hopper` / `_opt` / `_auto` | KDA prefill (CUDA C++ fused) | SM90 (Hopper) | Forward only | | `flashkda_prefill` | KDA prefill (CuTeDSL K1+K2, intracard-CP capable) | SM90 (Hopper) | Forward only | -| `kda_prefill` | Auto backend dispatch (FlashKDA → fully_fused) | SM90 (Hopper) | Forward only | +| `kda_prefill` | Public FlashKDA prefill entry point | SM90/SM100 | Forward only | | `kda_decode` | Single-token decode | SM100 | Forward | | `fused_sigmoid_gating_delta_rule_update` | Decode state update | SM100 | Forward | -`chunk_kda`, `kda_prefill_hopper*`, and `flashkda_prefill` each call a fixed implementation. -`kda_prefill` picks an SM90 backend at runtime (disable with `CULA_BACKEND_FLASHKDA=0` / `CULA_BACKEND_FULLY_FUSED=0`). +`chunk_kda` and `flashkda_prefill` each call a fixed implementation. +`kda_prefill` calls FlashKDA directly; there is no C++/CUDA extension fallback. +FlashKDA currently requires matching Q/K/V head counts (MHA rather than GVA). ## Minimal example diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index f8159806..0cc27eec 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -26,20 +26,14 @@ "kda_flush_kvbuffer_all_layers", "kda_packed_decode", "fused_sigmoid_gating_delta_rule_update", - "kda_prefill_hopper", - "kda_prefill_hopper_opt", - "kda_prefill_hopper_auto", "flashkda_prefill", "kda_prefill", ] _LAZY = { "chunk_kda": ("cula.kda.chunk", "chunk_kda"), - "kda_prefill_hopper": ("cula.kda.hopper_fused_fwd", "cula_kda_prefill"), - "kda_prefill_hopper_opt": ("cula.kda.hopper_fused_fwd_opt", "cula_kda_prefill_opt"), - "kda_prefill_hopper_auto": ("cula.kda.auto_route", "cula_kda_prefill_auto"), "flashkda_prefill": ("cula.kda.flashkda", "cula_kda_prefill"), - "kda_prefill": ("cula.kda.backends", "kda_prefill"), + "kda_prefill": ("cula.kda.prefill", "kda_prefill"), "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), "kda_decode_mtp": ("cula.ops.kda.decode.mtp", "kda_decode_mtp"), "kda_decode_mtp_recurrent": ("cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent"), diff --git a/cula/kda/auto_route.py b/cula/kda/auto_route.py deleted file mode 100644 index 1ce8405a..00000000 --- a/cula/kda/auto_route.py +++ /dev/null @@ -1,113 +0,0 @@ -from __future__ import annotations - -import torch - -from cula.kda.cp_context import is_dominant_long_seq -from cula.kda.hopper_fused_fwd import cula_kda_prefill as _basic -from cula.kda.hopper_fused_fwd_opt import FUSED_GATE_L2NORM_TH_VARLEN -from cula.kda.hopper_fused_fwd_opt import cula_kda_prefill_opt as _opt - - -def _should_use_opt(q: torch.Tensor, cu_seqlens: torch.Tensor | None) -> bool: - """Pick opt vs basic based on H100 measurements.""" - B = q.shape[0] - T = q.shape[1] - H = q.shape[2] - - if cu_seqlens is not None: - N = cu_seqlens.numel() - 1 - if N > 1: - packed_T = q.shape[1] - if packed_T * H <= FUSED_GATE_L2NORM_TH_VARLEN: - return True - if N * H <= 16 and T >= 8192: - return True - if H <= 16 and T >= 32768 + N - 1: - cu_list = cu_seqlens.tolist() - seqlens = [cu_list[i + 1] - cu_list[i] for i in range(N)] - if is_dominant_long_seq(seqlens, H): - return True - return False - # N == 1 falls through to the single-sequence logic below. - - # Fused gate+l2norm reliably wins at very small T*H even with B>1. - if T * H <= 6000: - return True - - if B == 1: - # T=1024 H=8/16 gets a small win from fused l2norm_qk (T*H<10000). - if H <= 16 and T <= 1024: - return True - if H <= 8: - return T >= 4096 # CP kicks in - elif H <= 16: - return T >= 4096 - elif H <= 32: - return T >= 16384 - else: # H >= 64 - return False # base ties or wins - - if B == 2: - if H == 8: - return T >= 4096 - return False # B=2 H>=16 mostly ties - - # B >= 4 : B*H >= 32 already saturates a sizable fraction of SMs, CP - # buys little; basic and opt tie. Default to basic (cheaper wrapper). - return False - - -def cula_kda_prefill_auto( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = True, - use_qk_l2norm_in_kernel: bool = True, - use_gate_in_kernel: bool = True, - safe_gate: bool = True, - lower_bound: float | None = -5.0, - cu_seqlens: torch.IntTensor | None = None, - chunk_indices: torch.IntTensor | None = None, - auto_cp: bool = True, - **kwargs, -): - if _should_use_opt(q, cu_seqlens): - return _opt( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - use_gate_in_kernel=use_gate_in_kernel, - safe_gate=safe_gate, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - auto_cp=auto_cp, - **kwargs, - ) - return _basic( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - use_gate_in_kernel=use_gate_in_kernel, - safe_gate=safe_gate, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - **kwargs, - ) diff --git a/cula/kda/backends/flashkda.py b/cula/kda/backends/flashkda.py deleted file mode 100644 index fe903bbf..00000000 --- a/cula/kda/backends/flashkda.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import importlib.util - -import torch - -from cula.backends import BaseBackend -from cula.kda._flashkda_arch import is_flashkda_supported - - -class FlashKDABackend(BaseBackend): - backend_type = "flashkda" - env_var = "CULA_BACKEND_FLASHKDA" - priority = 1 - - def probe(self): - if importlib.util.find_spec("cutlass") is None: - return False, "nvidia-cutlass-dsl not installed" - return True, None - - def kda_prefill_verifier( - self, - q, - k, - v, - g, - beta, - scale=None, - initial_state=None, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - safe_gate=True, - lower_bound=-5.0, - cu_seqlens=None, - chunk_indices=None, - **kwargs, - ): - if not is_flashkda_supported(q.device): - return False, "requires an SM90 (Hopper) or SM100 (Blackwell) device" - if v.shape[2] != q.shape[2]: - return False, f"no GVA support (HV={v.shape[2]} != H={q.shape[2]})" - if any(not t.is_contiguous() for t in (q, k, v, g)): - return False, "requires contiguous q/k/v/g (no implicit copies)" - if g.dtype != torch.bfloat16: - return False, f"requires bfloat16 g (no implicit cast), got {g.dtype}" - if not use_gate_in_kernel: - return False, "requires use_gate_in_kernel=True" - if not safe_gate: - return False, "requires safe_gate=True" - if not use_qk_l2norm_in_kernel: - return False, "always applies qk l2norm in kernel" - A_log, dt_bias = kwargs.get("A_log"), kwargs.get("dt_bias") - if A_log is None or dt_bias is None: - return False, "requires A_log and dt_bias" - if isinstance(A_log, torch.Tensor) and not A_log.is_contiguous(): - return False, "requires contiguous A_log" - if isinstance(dt_bias, torch.Tensor) and not dt_bias.is_contiguous(): - return False, "requires contiguous dt_bias" - if isinstance(cu_seqlens, torch.Tensor) and not cu_seqlens.is_contiguous(): - return False, "requires contiguous cu_seqlens" - if isinstance(initial_state, torch.Tensor) and not initial_state.is_contiguous(): - return False, "requires contiguous initial_state" - return True, None - - def kda_prefill(self, q, k, v, g, beta, **kwargs): - from cula.kda.flashkda import cula_kda_prefill - - kwargs.setdefault("use_intracard_cp", "auto") - return cula_kda_prefill(q, k, v, g, beta, **kwargs) diff --git a/cula/kda/backends/fully_fused.py b/cula/kda/backends/fully_fused.py deleted file mode 100644 index 805e0cc7..00000000 --- a/cula/kda/backends/fully_fused.py +++ /dev/null @@ -1,67 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import importlib.util - -import torch - -from cula.backends import BaseBackend -from cula.ops.kda.cp_mode import CPMode - - -def _is_sm90(device: torch.device) -> bool: - return device.type == "cuda" and torch.cuda.get_device_capability(device) == (9, 0) - - -class FullyFusedBackend(BaseBackend): - backend_type = "fully_fused" - env_var = "CULA_BACKEND_FULLY_FUSED" - priority = 2 - - def probe(self): - if importlib.util.find_spec("cula._cudac_sm90") is None: - return False, "cula._cudac_sm90 extension not built (rebuild with CULA_DISABLE_SM90=0)" - return True, None - - def kda_prefill_verifier( - self, - q, - k, - v, - g, - beta, - scale=None, - initial_state=None, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - safe_gate=True, - lower_bound=-5.0, - cu_seqlens=None, - chunk_indices=None, - **kwargs, - ): - if not _is_sm90(q.device): - return False, "requires an SM90 (Hopper) device" - H, HV = q.shape[2], v.shape[2] - if HV % H != 0: - return False, f"requires HV to be a multiple of H, got HV={HV} H={H}" - if not safe_gate: - return False, "requires safe_gate=True" - if kwargs.get("use_beta_sigmoid_in_kernel", False) is not False: - return False, "does not support use_beta_sigmoid_in_kernel=True" - if kwargs.get("use_intracard_cp") is CPMode.FORCE: - return False, "does not support forced intracard CP" - if kwargs.get("out") is not None or kwargs.get("final_state") is not None: - return False, "does not support preallocated out or final_state buffers" - return True, None - - def kda_prefill(self, q, k, v, g, beta, **kwargs): - from cula.kda.auto_route import cula_kda_prefill_auto - - cp_mode = kwargs.pop("use_intracard_cp", None) - for name in ("use_beta_sigmoid_in_kernel", "out", "final_state"): - kwargs.pop(name, None) - return cula_kda_prefill_auto(q, k, v, g, beta, auto_cp=cp_mode is not CPMode.OFF, **kwargs) diff --git a/cula/kda/cp_context.py b/cula/kda/cp_context.py deleted file mode 100644 index 976870ed..00000000 --- a/cula/kda/cp_context.py +++ /dev/null @@ -1,452 +0,0 @@ -from __future__ import annotations - -import functools -import math -import weakref - -import torch -from fla.utils import tensor_cache - -from cula.utils import get_device_sm_count - - -@tensor_cache -def _create_cu_seqlens(batch_size: int, num_tokens: int, device_idx: int, dtype: torch.dtype) -> torch.Tensor: - return torch.arange(batch_size + 1, dtype=dtype, device=f"cuda:{device_idx}") * num_tokens - - -@functools.lru_cache(maxsize=32) -def _create_full_cu_seqlens_2(T: int, device_idx: int, dtype: torch.dtype) -> torch.Tensor: - return torch.tensor([0, T], dtype=dtype, device=f"cuda:{device_idx}") - - -_CP_SEQS_CACHE: dict = {} -_SLOT_MAP_CACHE: dict = {} - - -DOMINANT_LONG_SEQ_MIN_LEN = 32768 -DOMINANT_LONG_SEQ_MAX_H = 16 - - -def is_dominant_long_seq( - seqlens: list[int], - H: int, - min_long_len: int = DOMINANT_LONG_SEQ_MIN_LEN, - max_H: int = DOMINANT_LONG_SEQ_MAX_H, -) -> bool: - if not seqlens or max_H < H: - return False - longest = max(seqlens) - if longest < min_long_len: - return False - return 2 * longest >= sum(seqlens) - - -def _seqlens_from_cu(cu_seqlens: torch.Tensor, cu_seqlens_cpu: torch.Tensor | None = None) -> list[int]: - """Return per-seq lengths from a cu_seqlens tensor (CPU sync if no CPU copy).""" - src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens - cu_list = src.tolist() - return [cu_list[i + 1] - cu_list[i] for i in range(len(cu_list) - 1)] - - -def _get_slot_map(cp_cu_seqlens: torch.Tensor, T: int, chunk_size: int) -> torch.Tensor: - key = (id(cp_cu_seqlens), T, chunk_size) - hit = _SLOT_MAP_CACHE.get(key) - if hit is not None: - weak_ref, cached = hit - if weak_ref() is cp_cu_seqlens: - return cached - # id was reused after the original tensor was collected — recompute. - num_chunks = (T + chunk_size - 1) // chunk_size - cp_starts = (cp_cu_seqlens[:-1] // chunk_size).to(torch.int32) - slot_map = torch.full((num_chunks,), -1, dtype=torch.int32, device=cp_cu_seqlens.device) - slots = torch.arange(cp_starts.numel(), dtype=torch.int32, device=cp_cu_seqlens.device) - slot_map[cp_starts.long()] = slots - _SLOT_MAP_CACHE[key] = (weakref.ref(cp_cu_seqlens), slot_map) - return slot_map - - -def _calc_cp_seqs_cached(raw_cu_seqlens, chunk_size, num_v_heads, sm_count, raw_cu_seqlens_cpu=None): - key = (id(raw_cu_seqlens), chunk_size, num_v_heads, sm_count) - hit = _CP_SEQS_CACHE.get(key) - if hit is not None: - weak_ref, cached = hit - if weak_ref() is raw_cu_seqlens: - return cached - - val = _calc_cp_seqs(raw_cu_seqlens, chunk_size, num_v_heads, sm_count, raw_cu_seqlens_cpu=raw_cu_seqlens_cpu) - _CP_SEQS_CACHE[key] = (weakref.ref(raw_cu_seqlens), val) - return val - - -def _calc_cp_seqs( - raw_cu_seqlens: torch.Tensor, - chunk_size: int, - num_v_heads: int, - sm_count: int, - raw_cu_seqlens_cpu: torch.Tensor | None = None, -) -> tuple[bool, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: - """Decide whether intra-card CP pays off and, if so, build the split tables.""" - device = raw_cu_seqlens.device - seqlen_dtype = raw_cu_seqlens.dtype - - raw_cu_seqlens_list = (raw_cu_seqlens_cpu if raw_cu_seqlens_cpu is not None else raw_cu_seqlens).tolist() - raw_batch_size = len(raw_cu_seqlens_list) - 1 - seqlens = [raw_cu_seqlens_list[i + 1] - raw_cu_seqlens_list[i] for i in range(raw_batch_size)] - num_chunks = [(s + chunk_size - 1) // chunk_size for s in seqlens] - if max(num_chunks) <= 0: - return False, None, None, None, None - - H = num_v_heads - V_BLOCKS = 1 # bump to 2 once main kernel supports V-blocking - target_cp_batch = max(1, sm_count // (H * V_BLOCKS)) - total_chunks = sum(num_chunks) - # mlc * cp_batch * chunk_size ≈ T per raw seq → mlc = total_chunks / cp_batch. - target_mlc = max(1, total_chunks // (max(1, target_cp_batch))) - # Snap to nearest power of 2; clamp to ≥ 4 to keep multi-stage pipelining alive. - max_local_chunks = 2 ** round(math.log2(max(target_mlc, 1.0))) - max_local_chunks = max(max_local_chunks, 4) - max_local_tokens = max_local_chunks * chunk_size - - cp_cu_seqlens: list[int] = [] - ht_mask: list[bool] = [] - seq_map_c2r: list[int] = [] - seq_map_r2c: list[int] = [0] - - for i, c in enumerate(num_chunks): - s = raw_cu_seqlens_list[i] - e = raw_cu_seqlens_list[i + 1] - if c > max_local_chunks: - cut = s - while True: - cp_cu_seqlens.append(cut) - ht_mask.append(False) - seq_map_c2r.append(i) - remaining = e - cut - if remaining <= max_local_tokens + chunk_size: - break - cut += max_local_tokens - ht_mask[-1] = True - else: - cp_cu_seqlens.append(s) - ht_mask.append(True) - seq_map_c2r.append(i) - seq_map_r2c.append(len(cp_cu_seqlens)) - cp_cu_seqlens.append(raw_cu_seqlens_list[-1]) - - Be = total_chunks / max(num_chunks) - use_cp = (Be * H <= 40) or (Be * H <= 56 and max(num_chunks) >= 128) - # Additional cuLA-specific guard: never bother if there is only one CP-chunk - # (no split happened). - if len(cp_cu_seqlens) - 1 == raw_batch_size: - use_cp = False - - if use_cp and raw_batch_size == 1: - T_max = max(seqlens) - if H <= 8 or H <= 16: - if T_max < 4096: - use_cp = False - elif H <= 32: - if T_max < 16384: - use_cp = False - else: # H >= 64 - use_cp = False - elif use_cp and raw_batch_size > 1: - T_packed = sum(seqlens) - native_grid = raw_batch_size * H - dominant_exception = is_dominant_long_seq(seqlens, H) - - unaligned = any(s % chunk_size != 0 for s in seqlens[:-1]) or (raw_cu_seqlens_list[-1] % chunk_size != 0) - if unaligned: - use_cp = False - elif native_grid > 16 and not dominant_exception: - # Native grid already big enough that CP's lift is marginal. - use_cp = False - elif T_packed * H <= 32768: - use_cp = False - elif H <= 8: - if T_packed < 8192: - use_cp = False - elif H <= 16: - if T_packed < 4096: - use_cp = False - else: # H >= 32 with native_grid <= 16 → only B=1 fits this, handled above - use_cp = False - - if not use_cp: - return False, None, None, None, None - - cp_cu_seqlens_t = torch.tensor(cp_cu_seqlens, dtype=seqlen_dtype, device=device) - seq_map_c2r_t = torch.tensor(seq_map_c2r, dtype=seqlen_dtype, device=device) - seq_map_r2c_t = torch.tensor(seq_map_r2c, dtype=seqlen_dtype, device=device) - ht_mask_t = torch.tensor(ht_mask, dtype=torch.bool, device=device) - return True, cp_cu_seqlens_t, seq_map_r2c_t, seq_map_c2r_t, ht_mask_t - - -def _build_raw_seq_idx( - cp_cu_seqlens: torch.Tensor, seq_map_c2r: torch.Tensor, T: int, chunk_size: int -) -> tuple[torch.Tensor, list[int]]: - """Per-chunk raw seq id: raw_seq_idx[i_t] = raw_seq containing chunk i_t.""" - NT = (T + chunk_size - 1) // chunk_size - out = torch.empty(NT, dtype=torch.int32, device=cp_cu_seqlens.device) - # cp_cu_seqlens[:-1] // chunk_size gives the chunk-index start of each CP-chunk. - cp_starts = (cp_cu_seqlens[:-1] // chunk_size).to(torch.int64) - cp_ends = ((cp_cu_seqlens[1:] + chunk_size - 1) // chunk_size).to(torch.int64) - - c2r_cpu = seq_map_c2r.tolist() - starts = cp_starts.tolist() - ends = cp_ends.tolist() - out_cpu = [0] * NT - for i in range(len(starts)): - out[starts[i] : ends[i]] = c2r_cpu[i] - for j in range(starts[i], ends[i]): - out_cpu[j] = c2r_cpu[i] - return out, out_cpu - - -_RAW_SEQ_IDX_CACHE: dict = {} - - -def _get_raw_seq_idx( - cp_cu_seqlens: torch.Tensor, seq_map_c2r: torch.Tensor, T: int, chunk_size: int -) -> tuple[torch.Tensor, list[int]]: - """weakref-guarded cache for the per-chunk raw_seq_idx tensor.""" - key = (id(cp_cu_seqlens), T, chunk_size) - hit = _RAW_SEQ_IDX_CACHE.get(key) - if hit is not None: - weak_ref, cached = hit - if weak_ref() is cp_cu_seqlens: - return cached - val = _build_raw_seq_idx(cp_cu_seqlens, seq_map_c2r, T, chunk_size) - _RAW_SEQ_IDX_CACHE[key] = (weakref.ref(cp_cu_seqlens), val) - return val - - -def _compute_cp_h0_via_fla_h( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor, - beta: torch.Tensor, - scale: float, - raw_h0: torch.Tensor | None, - cp_cu_seqlens: torch.Tensor, - seq_map_c2r: torch.Tensor, - chunk_size: int, -) -> torch.Tensor: - """Compute cp_h0 directly via FLA's per-chunk h tensor — exact, no mt needed.""" - from cula.kda.cp_h_boundary import kda_cp_h0_boundary - from cula.kda.wy_intra import kda_intra_native - - cp_batch = cp_cu_seqlens.size(0) - 1 - raw_batch = int(seq_map_c2r.max().item()) + 1 if seq_map_c2r.numel() > 0 else 1 - T = k.size(1) - - if T > CP_PREPROCESS_TILE_TOKENS: - return _compute_cp_h0_via_fla_h_tiled( - k=k, - v=v, - g_cumsum=g_cumsum, - beta=beta, - raw_h0=raw_h0, - cp_cu_seqlens=cp_cu_seqlens, - seq_map_c2r=seq_map_c2r, - chunk_size=chunk_size, - ) - - w, u, _, kg = kda_intra_native( - k=k, - v=v, - gk=g_cumsum, - beta=beta, - chunk_size=chunk_size, - ) - - slot_map = _get_slot_map(cp_cu_seqlens, T, chunk_size) - - if raw_batch > 1: - if raw_h0 is None: - H_v = v.size(2) - K = k.size(3) - V = v.size(3) - raw_h0 = torch.zeros(raw_batch, H_v, V, K, dtype=torch.float32, device=k.device) - h0_chunk0 = raw_h0[0:1] # state at chunk 0 (always raw seq 0) - raw_seq_idx, _ = _get_raw_seq_idx(cp_cu_seqlens, seq_map_c2r, T, chunk_size) - else: - h0_chunk0 = raw_h0 - raw_seq_idx = None - - cp_h0 = kda_cp_h0_boundary( - kg=kg, - w=w, - u=u, - g_cumsum=g_cumsum, - h0=h0_chunk0, - slot_map=slot_map, - num_cp=cp_batch, - chunk_size=chunk_size, - raw_h0_dense=raw_h0 if raw_batch > 1 else None, - raw_seq_idx=raw_seq_idx, - ) - del w, u, kg - return cp_h0 - - -CP_PREPROCESS_TILE_TOKENS = 16384 - - -def _compute_cp_h0_via_fla_h_tiled( - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor, - beta: torch.Tensor, - raw_h0: torch.Tensor | None, - cp_cu_seqlens: torch.Tensor, - seq_map_c2r: torch.Tensor, - chunk_size: int, -) -> torch.Tensor: - from cula.kda.cp_h_boundary import kda_cp_h0_boundary - from cula.kda.wy_intra import kda_intra_native - - assert k.size(0) == 1, "tiled CP preprocess assumes packed [1, T, H, K] layout" - cp_batch = cp_cu_seqlens.size(0) - 1 - raw_batch = int(seq_map_c2r.max().item()) + 1 if seq_map_c2r.numel() > 0 else 1 - T = k.size(1) - H = k.size(2) - K = k.size(3) - V = v.size(3) - - if raw_h0 is None: - raw_h0 = torch.zeros(raw_batch, H, V, K, dtype=torch.float32, device=k.device) - device = k.device - - assert CP_PREPROCESS_TILE_TOKENS % chunk_size == 0 - tile_tokens = CP_PREPROCESS_TILE_TOKENS - - global_slot_map = _get_slot_map(cp_cu_seqlens, T, chunk_size) - # Per-chunk raw-seq map (length NT_total). For raw_batch==1 this is all - # zeros and we don't even pass it to the kernel. - if raw_batch > 1: - global_raw_seq_idx, global_raw_seq_idx_cpu = _get_raw_seq_idx(cp_cu_seqlens, seq_map_c2r, T, chunk_size) - else: - global_raw_seq_idx = None - global_raw_seq_idx_cpu = None - - cp_h0 = torch.empty(cp_batch, H, V, K, dtype=torch.float32, device=device) - - exit_state = torch.empty(H, V, K, dtype=torch.float32, device=device) - - BT = chunk_size - BC = 16 - max_tile_T = min(tile_tokens, T) - buf_w = torch.empty(1, max_tile_T, H, K, dtype=k.dtype, device=device) - buf_u = torch.empty(1, max_tile_T, H, V, dtype=v.dtype, device=device) - buf_kg = torch.empty(1, max_tile_T, H, K, dtype=k.dtype, device=device) - buf_Akkd = torch.empty(1, max_tile_T, H, BC, dtype=torch.float32, device=device) - buf_Akk = torch.empty(1, max_tile_T, H, BT, dtype=k.dtype, device=device) - - n_tiles = (T + tile_tokens - 1) // tile_tokens - for tile_idx in range(n_tiles): - s = tile_idx * tile_tokens - e = min(s + tile_tokens, T) - is_last_tile = tile_idx == n_tiles - 1 - s_chunk = s // chunk_size - e_chunk = (e + chunk_size - 1) // chunk_size - - k_t = k[:, s:e] - v_t = v[:, s:e] - g_t = g_cumsum[:, s:e] - beta_t = beta[:, s:e] - - w_t, u_t, _, kg_t = kda_intra_native( - k=k_t, - v=v_t, - gk=g_t, - beta=beta_t, - chunk_size=chunk_size, - out_w=buf_w, - out_u=buf_u, - out_kg=buf_kg, - out_Akkd=buf_Akkd, - out_Akk=buf_Akk, - ) - - slot_map_t = global_slot_map[s_chunk:e_chunk] - - if raw_batch > 1: - first_raw_in_tile = global_raw_seq_idx_cpu[s_chunk] - if tile_idx == 0: - h0_in = raw_h0[first_raw_in_tile : first_raw_in_tile + 1] - else: - prev_tile_last_raw = global_raw_seq_idx_cpu[s_chunk - 1] - if first_raw_in_tile == prev_tile_last_raw: - h0_in = exit_state - else: - h0_in = raw_h0[first_raw_in_tile : first_raw_in_tile + 1] - raw_seq_idx_t = global_raw_seq_idx[s_chunk:e_chunk] - else: - h0_in = raw_h0 if tile_idx == 0 else exit_state - raw_seq_idx_t = None - - kda_cp_h0_boundary( - kg=kg_t, - w=w_t, - u=u_t, - g_cumsum=g_t, - h0=h0_in, - slot_map=slot_map_t, - num_cp=cp_batch, - chunk_size=chunk_size, - cp_h0_out=cp_h0, - # Skip writing exit_state on the last tile — no consumer. - exit_state=None if is_last_tile else exit_state, - raw_h0_dense=raw_h0 if raw_batch > 1 else None, - raw_seq_idx=raw_seq_idx_t, - ) - - return cp_h0 - - -def intra_card_cp_preprocess( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - raw_h0: torch.Tensor | None, - raw_cu_seqlens: torch.Tensor | None, - chunk_size: int = 64, - raw_cu_seqlens_cpu: torch.Tensor | None = None, -) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: - assert k.dim() == 4 and k.size(0) == 1, "expected packed [1, T, H, K]" - num_v_heads = v.size(2) - sm_count = get_device_sm_count(k.device) - - if raw_cu_seqlens is None: - raw_cu_seqlens = _create_cu_seqlens(1, k.size(1), k.device.index, torch.int32) - - use_cp, cp_cu_seqlens, seq_map_r2c, seq_map_c2r, ht_mask = _calc_cp_seqs_cached( - raw_cu_seqlens, - chunk_size, - num_v_heads, - sm_count, - raw_cu_seqlens_cpu=raw_cu_seqlens_cpu, - ) - if not use_cp: - return None, None, None, None - - cp_h0 = _compute_cp_h0_via_fla_h( - q=q, - k=k, - v=v, - g_cumsum=g, - beta=beta, - scale=scale, - raw_h0=raw_h0, - cp_cu_seqlens=cp_cu_seqlens, - seq_map_c2r=seq_map_c2r, - chunk_size=chunk_size, - ) - - return cp_h0, cp_cu_seqlens, seq_map_c2r, raw_cu_seqlens diff --git a/cula/kda/cp_h_boundary.py b/cula/kda/cp_h_boundary.py deleted file mode 100644 index b584bee9..00000000 --- a/cula/kda/cp_h_boundary.py +++ /dev/null @@ -1,190 +0,0 @@ -from __future__ import annotations - -import torch -import triton -import triton.language as tl - - -@triton.autotune( - configs=[ - triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) - for BV in [16, 32, 64] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V", "BT"], -) -@triton.heuristics( - { - "USE_INITIAL_STATE": lambda args: args["h0"] is not None, - "STORE_EXIT_STATE": lambda args: args["exit_state"] is not None, - "MULTI_RAW": lambda args: args["raw_seq_idx"] is not None, - } -) -@triton.jit(do_not_specialize=["T"]) -def _kda_h_boundary_kernel( - kg, - w, - u, - gk, - h0, # fp32 [H, V, K] OR [1, H, V, K] — state at chunk 0 - cp_h0_out, # fp32 [num_cp, H, V, K], cuLA layout - slot_map, # int32 [NT], slot_map[i_t] = output slot or -1 - exit_state, # fp32 [H, V, K] or None — end-of-tile state for next tile - raw_h0_dense, # fp32 [raw_batch, H, V, K] — per-raw-seq h0 for cross-seq resets - raw_seq_idx, # int32 [NT] — raw_seq_idx[i_t] = which raw seq chunk i_t belongs to - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_EXIT_STATE: tl.constexpr, - MULTI_RAW: tl.constexpr, -): - # Grid: (V_tiles, H). One CTA per (V-tile, head) walks all T chunks serially. - i_v, i_h = tl.program_id(0), tl.program_id(1) - NT = tl.cdiv(T, BT) - - # State tiles: 2× [BV, 64] holds the (V_tile=BV, K=128) state split into K=64 halves. - b_h1 = tl.zeros([BV, 64], dtype=tl.float32) - b_h2 = tl.zeros([BV, 64], dtype=tl.float32) - - # Per-batch (B=1) offset into per-head buffers. - kg_ptr = kg + i_h * K - w_ptr = w + i_h * K - u_ptr = u + i_h * V - gk_ptr = gk + i_h * K - h0_ptr = h0 + i_h * V * K if USE_INITIAL_STATE else h0 # h0 [H, V, K] or [1, H, V, K] - - if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0_ptr, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - p_h0_2 = tl.make_block_ptr(h0_ptr, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) - - if MULTI_RAW: - prev_raw = tl.load(raw_seq_idx + 0) - else: - prev_raw = 0 # unused - - for i_t in range(NT): - if MULTI_RAW: - if i_t > 0: - cur_raw = tl.load(raw_seq_idx + i_t) - if cur_raw != prev_raw: - rh_base = raw_h0_dense + cur_raw.to(tl.int64) * (H * V * K) + i_h * (V * K) - p_rh_1 = tl.make_block_ptr(rh_base, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - p_rh_2 = tl.make_block_ptr(rh_base, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - b_h1 = tl.load(p_rh_1, boundary_check=(0, 1)).to(tl.float32) - b_h2 = tl.load(p_rh_2, boundary_check=(0, 1)).to(tl.float32) - prev_raw = cur_raw - - # ----- Conditional boundary store ----- - # slot = slot_map[i_t]; if slot >= 0, write h to cp_h0_out[slot, i_h, ...] - slot = tl.load(slot_map + i_t) - is_boundary = slot >= 0 - # Use slot=0 as safe target when not boundary; mask blocks the store. - safe_slot = tl.maximum(slot, 0).to(tl.int64) - out_base = cp_h0_out + safe_slot * (H * V * K) + i_h * (V * K) - p_out_1 = tl.make_block_ptr(out_base, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - p_out_2 = tl.make_block_ptr(out_base, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - if is_boundary: - tl.store(p_out_1, b_h1, boundary_check=(0, 1)) - tl.store(p_out_2, b_h2, boundary_check=(0, 1)) - - # ----- Compute v_new = u - w @ h ----- - # w [BT, K] split into 2× [BT, 64] for the matmuls - p_w1 = tl.make_block_ptr(w_ptr, (T, K), (H * K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_v_acc = tl.dot(b_w1, tl.trans(b_h1).to(b_w1.dtype)) - p_w2 = tl.make_block_ptr(w_ptr, (T, K), (H * K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - b_v_acc += tl.dot(b_w2, tl.trans(b_h2).to(b_w2.dtype)) - - # u [BT, V_tile] - p_u = tl.make_block_ptr(u_ptr, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_u, boundary_check=(0, 1)) - b_v_acc - - # ----- Apply gate decay: h *= exp2(gk_last) ----- - last_idx = tl.minimum((i_t + 1) * BT, T) - 1 - o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk_ptr + last_idx * (H * K) + o_k1, mask=(o_k1 < K), other=0.0).to(tl.float32) - b_h1 *= tl.math.exp2(b_gk_last1)[None, :] - o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk_ptr + last_idx * (H * K) + o_k2, mask=(o_k2 < K), other=0.0).to(tl.float32) - b_h2 *= tl.math.exp2(b_gk_last2)[None, :] - - # ----- State update: h += kg^T @ v_new (transpose_state_layout=True form) ----- - b_v_bf = b_v.to(kg.dtype.element_ty) - p_kg1 = tl.make_block_ptr(kg_ptr, (K, T), (1, H * K), (0, i_t * BT), (64, BT), (0, 1)) - b_kg1 = tl.load(p_kg1, boundary_check=(0, 1)) - b_h1 += tl.trans(tl.dot(b_kg1, b_v_bf)) - p_kg2 = tl.make_block_ptr(kg_ptr, (K, T), (1, H * K), (64, i_t * BT), (64, BT), (0, 1)) - b_kg2 = tl.load(p_kg2, boundary_check=(0, 1)) - b_h2 += tl.trans(tl.dot(b_kg2, b_v_bf)) - - # ----- After the loop: optionally store the end-of-tile state for the next tile ----- - if STORE_EXIT_STATE: - es_ptr = exit_state + i_h * V * K - p_es_1 = tl.make_block_ptr(es_ptr, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - p_es_2 = tl.make_block_ptr(es_ptr, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - tl.store(p_es_1, b_h1, boundary_check=(0, 1)) - tl.store(p_es_2, b_h2, boundary_check=(0, 1)) - - -def kda_cp_h0_boundary( - kg: torch.Tensor, # bf16 [1, T, H, K] - w: torch.Tensor, # bf16 [1, T, H, K] - u: torch.Tensor, # bf16 [1, T, H, V] - g_cumsum: torch.Tensor, # fp32 [1, T, H, K] - h0: torch.Tensor | None, # fp32 [H, V, K] (or [1, H, V, K]) — state at this tile's chunk 0 - slot_map: torch.Tensor, # int32 [NT], slot_map[i_t] = output slot or -1 - num_cp: int, - chunk_size: int = 64, - cp_h0_out: torch.Tensor | None = None, # pre-allocated [num_cp, H, V, K] fp32; allocated if None - exit_state: torch.Tensor | None = None, # if not None, kernel writes the end-of-T state to [H, V, K] fp32 - raw_h0_dense: torch.Tensor | None = None, # fp32 [raw_batch, H, V, K] for cross-raw-seq resets - raw_seq_idx: torch.Tensor | None = None, # int32 [NT] mapping chunk → raw seq idx (None for single-raw) -) -> torch.Tensor: - assert kg.is_contiguous() and w.is_contiguous() and u.is_contiguous() - assert kg.size(0) == 1 and w.size(0) == 1 and u.size(0) == 1 and g_cumsum.size(0) == 1 - assert slot_map.dtype == torch.int32 and slot_map.is_contiguous() - if raw_seq_idx is not None: - assert raw_h0_dense is not None, "raw_h0_dense required when raw_seq_idx is set" - assert raw_seq_idx.dtype == torch.int32 and raw_seq_idx.is_contiguous() - - T = kg.size(1) - H = kg.size(2) - K = kg.size(3) - V = u.size(3) - assert K == 128 and V == 128, f"Phase 1 kernel hard-codes K=V=128, got K={K} V={V}" - - if cp_h0_out is None: - cp_h0_out = torch.empty(num_cp, H, V, K, dtype=torch.float32, device=kg.device) - - BT = chunk_size - - # Grid depends on BV (autotune-selected); use meta-grid. - def grid(meta): - return (V // meta["BV"], H) - - _kda_h_boundary_kernel[grid]( - kg, - w, - u, - g_cumsum, - h0, - cp_h0_out, - slot_map, - exit_state, - raw_h0_dense, - raw_seq_idx, - T, - H=H, - K=K, - V=V, - BT=BT, - ) - return cp_h0_out diff --git a/cula/kda/gate_l2norm_fused.py b/cula/kda/gate_l2norm_fused.py deleted file mode 100644 index aa93f69f..00000000 --- a/cula/kda/gate_l2norm_fused.py +++ /dev/null @@ -1,194 +0,0 @@ -import torch -import triton -import triton.language as tl -from fla.ops.utils.constant import RCP_LN2 as _RCP_LN2 -from fla.ops.utils.index import prepare_chunk_indices -from fla.ops.utils.softplus import softplus - -# Triton requires module-level constants used inside @jit kernels to be -# wrapped in tl.constexpr. -RCP_LN2 = tl.constexpr(_RCP_LN2) - - -@triton.jit -def _gate_l2norm_fused_kernel( - # Pointers - g_ptr, - A_log_ptr, - dt_bias_ptr, # gate inputs - q_ptr, - k_ptr, # qk inputs - g_out_ptr, # gate output (fp32 cumsum) - yq_ptr, - yk_ptr, # qk outputs (bf16) - rstd_q_ptr, - rstd_k_ptr, # qk rstd outputs (fp32) - cu_seqlens_ptr, - chunk_indices_ptr, - # Scalars - lower_bound, - eps_l2, - T, - H: tl.constexpr, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - HAS_BIAS: tl.constexpr, - USE_LOWER_BOUND: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t = tl.program_id(0) - i_bh = tl.program_id(1) - i_h = i_bh % H - - if IS_VARLEN: - i_n = tl.load(chunk_indices_ptr + i_t * 2).to(tl.int32) - i_t_local = tl.load(chunk_indices_ptr + i_t * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens_ptr + i_n).to(tl.int32) - eos = tl.load(cu_seqlens_ptr + i_n + 1).to(tl.int32) - T_seq = eos - bos - bt_base = bos + i_t_local * BT - valid_t = (i_t_local * BT + tl.arange(0, BT)) < T_seq - else: - bt_base = (i_bh // H) * T + i_t * BT - valid_t = (i_t * BT + tl.arange(0, BT)) < T - - rows = bt_base + tl.arange(0, BT) - - cols = tl.arange(0, BD) # (BD,) - valid_d = cols < D - - offs = rows[:, None] * (H * D) + i_h * D + cols[None, :] - mask = valid_t[:, None] & valid_d[None, :] - - # ==================================================================== - # GATE: cumsum( transform(g + bias) ) * RCP_LN2 - # ==================================================================== - b_g = tl.load(g_ptr + offs, mask=mask, other=0.0).to(tl.float32) - - if HAS_BIAS: - b_bias = tl.load(dt_bias_ptr + i_h * D + cols, mask=valid_d, other=0.0).to(tl.float32) - b_g = b_g + b_bias[None, :] - - b_A = tl.load(A_log_ptr + i_h).to(tl.float32) - if USE_LOWER_BOUND: - b_gate = lower_bound * tl.sigmoid(tl.exp(b_A) * b_g) - else: - b_gate = -tl.exp(b_A) * softplus(b_g) - - b_gate_cs = tl.cumsum(b_gate, axis=0) * RCP_LN2 - # zero out the tokens beyond T so we don't pollute g_out tail rows - b_gate_cs = tl.where(mask, b_gate_cs, 0.0) - tl.store(g_out_ptr + offs, b_gate_cs, mask=mask) - - # ==================================================================== - # L2-NORM Q — per-row normalisation along D, rstd written to (B*T*H,) - # ==================================================================== - b_q = tl.load(q_ptr + offs, mask=mask, other=0.0).to(tl.float32) - b_q_sq = tl.sum(b_q * b_q, axis=1) # (BT,) - b_rstd_q = 1.0 / tl.sqrt(b_q_sq + eps_l2) # (BT,) - b_yq = b_q * b_rstd_q[:, None] - tl.store(yq_ptr + offs, b_yq.to(yq_ptr.dtype.element_ty), mask=mask) - - # rstd is shape (B, T, H,) contiguous in (b*T*H + t*H + h) order - rstd_offs = rows * H + i_h # (BT,) - tl.store(rstd_q_ptr + rstd_offs, b_rstd_q, mask=valid_t) - - # ==================================================================== - # L2-NORM K - # ==================================================================== - b_k = tl.load(k_ptr + offs, mask=mask, other=0.0).to(tl.float32) - b_k_sq = tl.sum(b_k * b_k, axis=1) - b_rstd_k = 1.0 / tl.sqrt(b_k_sq + eps_l2) - b_yk = b_k * b_rstd_k[:, None] - tl.store(yk_ptr + offs, b_yk.to(yk_ptr.dtype.element_ty), mask=mask) - tl.store(rstd_k_ptr + rstd_offs, b_rstd_k, mask=valid_t) - - -def gate_l2norm_fused_fwd( - g: torch.Tensor, # (B, T, H, D) bf16 - q: torch.Tensor, # (B, T, H, D) bf16 - k: torch.Tensor, # (B, T, H, D) bf16 - A_log: torch.Tensor, # (H,) fp32 - dt_bias: torch.Tensor | None, # (H*D,) fp32 or None - lower_bound: float | None, # if not None and safe_gate -> use lb*sigmoid path - chunk_size: int = 64, - eps_l2: float = 1e-6, - cu_seqlens: torch.IntTensor | None = None, - chunk_indices: torch.IntTensor | None = None, -): - """One-launch fused preprocessing. - - Returns: - g_out: (B, T, H, D) fp32 -- gate cumsum * RCP_LN2 - y_q: (B, T, H, D) bf16 -- l2-normalised q - y_k: (B, T, H, D) bf16 -- l2-normalised k - rstd_q, rstd_k: (B, T, H) fp32 -- 1/sqrt(sum^2 + eps) - """ - assert g.shape == q.shape == k.shape, f"shapes must match: g{g.shape} q{q.shape} k{k.shape}" - assert g.is_contiguous() and q.is_contiguous() and k.is_contiguous(), "all inputs must be contiguous" - B, T, H, D = g.shape - assert chunk_size == 64, "only chunk_size=64 supported (matches SM90 main kernel)" - - is_varlen = cu_seqlens is not None - if is_varlen: - assert B == 1, "varlen path expects packed B=1 layout" - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - - if dt_bias is None: - dt_bias_arg = A_log # any valid fp32 pointer; HAS_BIAS=False suppresses use - has_bias = False - else: - dt_bias_arg = dt_bias - has_bias = True - - use_lower_bound = lower_bound is not None - if not use_lower_bound: - lower_bound = 0.0 # unused, but Triton needs a value - - g_out = torch.empty_like(g, dtype=torch.float32) - y_q = torch.empty_like(q) - y_k = torch.empty_like(k) - rstd_q = torch.empty((B, T, H), dtype=torch.float32, device=q.device) - rstd_k = torch.empty((B, T, H), dtype=torch.float32, device=q.device) - - BD = triton.next_power_of_2(D) - if is_varlen: - NT = chunk_indices.shape[0] - grid = (NT, H) - cu_seqlens_arg = cu_seqlens - chunk_indices_arg = chunk_indices - else: - NT = triton.cdiv(T, chunk_size) # ceil — partial last chunk handled by mask - grid = (NT, B * H) - cu_seqlens_arg = A_log - chunk_indices_arg = A_log - num_warps = 1 if BD <= 128 else (2 if BD <= 256 else 4) - - _gate_l2norm_fused_kernel[grid]( - g_ptr=g, - A_log_ptr=A_log, - dt_bias_ptr=dt_bias_arg, - q_ptr=q, - k_ptr=k, - g_out_ptr=g_out, - yq_ptr=y_q, - yk_ptr=y_k, - rstd_q_ptr=rstd_q, - rstd_k_ptr=rstd_k, - cu_seqlens_ptr=cu_seqlens_arg, - chunk_indices_ptr=chunk_indices_arg, - lower_bound=lower_bound, - eps_l2=eps_l2, - T=T, - H=H, - D=D, - BT=chunk_size, - BD=BD, - HAS_BIAS=has_bias, - USE_LOWER_BOUND=use_lower_bound, - IS_VARLEN=is_varlen, - num_warps=num_warps, - ) - return g_out, y_q, y_k, rstd_q, rstd_k diff --git a/cula/kda/hopper_fused_fwd.py b/cula/kda/hopper_fused_fwd.py deleted file mode 100644 index d84d0aba..00000000 --- a/cula/kda/hopper_fused_fwd.py +++ /dev/null @@ -1,281 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# 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. - - -import torch -from einops import rearrange -from fla.modules.l2norm import l2norm_fwd -from fla.ops.kda.gate import kda_gate_chunk_cumsum -from fla.ops.utils import chunk_local_cumsum -from fla.ops.utils.constant import RCP_LN2 -from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - -import cula.cudac as cula_cuda -from cula.utils import _get_cache_buf, assert_hopper, get_device_sm_count, prepare_uniform_cu_seqlens - - -class HopperChunkKDAFunction(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - safe_gate: bool = False, - lower_bound: float | None = None, - cu_seqlens: torch.IntTensor | None = None, - chunk_indices: torch.IntTensor | None = None, - ): - chunk_size = 64 - # GVA: q/k share num_qk_heads; v/g/beta share num_v_heads. - # num_v_heads must be a positive multiple of num_qk_heads (heads_per_group = HV / H). - assert q.shape == k.shape, "q and k must have the same shape." - assert q.shape[:2] == v.shape[:2] == g.shape[:2], "q, k, v, g must share batch and sequence dimensions." - - batch_size, seq_len, num_qk_heads, head_dim = q.shape - num_v_heads = v.shape[-2] - assert num_qk_heads > 0, f"num_qk_heads must be positive, got {num_qk_heads}." - assert num_v_heads > 0, f"num_v_heads must be positive, got {num_v_heads}." - assert num_v_heads % num_qk_heads == 0, ( - f"num_v_heads ({num_v_heads}) must be a positive multiple of num_qk_heads ({num_qk_heads})." - ) - - if cu_seqlens is None: - cu_seqlens = prepare_uniform_cu_seqlens(batch_size, seq_len, q.device, torch.int32) - - # set batch size to 1 after handling cu_seqlens - if batch_size != 1: - q, k, v, g, beta = map(lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), (q, k, v, g, beta)) - - # gate preprocessing - if use_gate_in_kernel: - if safe_gate: - assert lower_bound is not None, "lower_bound must be set when use safe_gate" - g = kda_gate_chunk_cumsum( - g=g, - A_log=A_log, - dt_bias=dt_bias, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - lower_bound=lower_bound, - ) - else: - g = chunk_local_cumsum( - g=g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - q_rstd, k_rstd = None, None - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - - # reshape q/k to packed [T, H, K] and v/g to [T, HV, K], beta to [T, HV] for the C++ kernel - packed_seq = batch_size * seq_len - q = q.reshape(packed_seq, num_qk_heads, head_dim).contiguous() - k = k.reshape(packed_seq, num_qk_heads, head_dim).contiguous() - v = v.reshape(packed_seq, num_v_heads, head_dim).contiguous() - g = g.reshape(packed_seq, num_v_heads, head_dim).contiguous() - beta = beta.reshape(packed_seq, num_v_heads).contiguous() - - # workspace buffer for TMA Store O tensormap - sm_count = get_device_sm_count(q.device) - workspace_size = sm_count * 128 - workspace_buffer = _get_cache_buf("hopper_kda_fwd_workspace", workspace_size, q.device) - - # call the C++ kernel - # Signature: kda_fwd_prefill(output_, output_state_, q, k, v, input_state_, alpha_, beta_, cu_seqlens, - # workspace, scale, output_final_state, safe_gate) - o, final_state = cula_cuda.kda_fwd_prefill( - None, # output_ (auto-allocate) - None, # output_state_ (auto-allocate) - q, - k, - v, - initial_state, # input_state_ - g, # alpha_ - beta, # beta_ - cu_seqlens, - workspace_buffer, - scale, - output_final_state, - safe_gate, - ) - - # reshape back - o = rearrange(o, "(b t) h d -> b t h d", b=batch_size) - - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht): - raise NotImplementedError("Backward pass is not implemented yet.") - - -@torch.compiler.disable -def cula_kda_prefill( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - safe_gate: bool = False, - lower_bound: float | None = None, - cu_seqlens: torch.IntTensor | None = None, - chunk_indices: torch.IntTensor | None = None, - **kwargs, -): - r""" - Hopper (SM90) fully-fused KDA forward prefill using CUTLASS TMA warp-specialized kernel. - - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, HV, K]`. - g (torch.Tensor): - (forget) gating tensor (in log space!) of shape `[B, T, HV, K]`. - beta (torch.Tensor): - betas of shape `[B, T, HV]`. - scale (Optional[float]): - Scale factor for the KDA attention scores. - If not provided, it will default to `1 / sqrt(D)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, HV, K, K]` for `N` input sequences. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, HV, K, K]`. Default: `False`. - use_qk_l2norm_in_kernel (bool): - Whether to apply L2norm to the q,k tensor internally. Default: `False`. - use_gate_in_kernel (bool): - Whether to compute the log-space KDA decay internally. Default: `False`. - safe_gate (bool): - Whether the kernel can assume the input gate values `g` are in a safe range. - When `True`, the kernel can use M=16 TensorCore acceleration. - The safe range is approximately [-5, 0). Default: `False`. - lower_bound (Optional[float]): - Lower bound for the forget gate activation function. Default: `None`. - cu_seqlens (torch.IntTensor): - Cumulative sequence lengths of shape `[N+1]`, int32. - chunk_indices (torch.IntTensor): - Chunk indices for variable-length training. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, HV, K]`. - final_state (torch.Tensor): - Final state of shape `[N, HV, K, K]` if `output_final_state=True` else `None`. - """ - assert_hopper(q.device) - assert safe_gate, "Only support safe_gate=True." - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if initial_state is not None: - assert initial_state.dtype == torch.float32, "initial_state must be in float32." - - A_log, dt_bias = None, None - if use_gate_in_kernel: - assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." - A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") - if safe_gate: - if lower_bound is None: - raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") - if not (-5 <= lower_bound < 0): - raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") - - assert q.shape == k.shape, "q and k must have the same shape." - assert q.shape[:2] == v.shape[:2] == g.shape[:2], "q, k, v, g must share batch and sequence dimensions." - - batch_size, seq_len, num_qk_heads, head_dim = q.shape - num_v_heads = v.shape[-2] - # Order matters here: positivity *first*, modulo second, to avoid ZeroDivisionError on bad inputs. - assert num_qk_heads > 0, f"num_qk_heads must be positive, got {num_qk_heads}." - assert num_v_heads > 0, f"num_v_heads must be positive, got {num_v_heads}." - assert num_v_heads % num_qk_heads == 0, ( - f"num_v_heads ({num_v_heads}) must be a positive multiple of num_qk_heads ({num_qk_heads})." - ) - assert g.shape == (batch_size, seq_len, num_v_heads, head_dim), ( - f"g must have shape (B, T, HV, D)=({batch_size}, {seq_len}, {num_v_heads}, {head_dim}), got {tuple(g.shape)}." - ) - assert v.shape == (batch_size, seq_len, num_v_heads, head_dim), ( - f"v must have shape (B, T, HV, D)=({batch_size}, {seq_len}, {num_v_heads}, {head_dim}), got {tuple(v.shape)}." - ) - assert beta.shape == (batch_size, seq_len, num_v_heads), ( - f"beta must have shape (B, T, HV)=({batch_size}, {seq_len}, {num_v_heads}), got {tuple(beta.shape)}." - ) - if initial_state is not None: - expected_num_states = (len(cu_seqlens) - 1) if cu_seqlens is not None else batch_size - assert initial_state.shape == (expected_num_states, num_v_heads, head_dim, head_dim), ( - f"initial_state must have shape (N, HV, D, D)=" - f"({expected_num_states}, {num_v_heads}, {head_dim}, {head_dim}), got {tuple(initial_state.shape)}." - ) - assert q.dtype == k.dtype == v.dtype == torch.bfloat16, "q, k, v must be in bfloat16." - assert beta.dtype == torch.bfloat16 or beta.dtype == torch.float32, "beta must be in bfloat16 or float32." - assert q.shape[-1] == k.shape[-1] == v.shape[-1] == 128, "Currently we only support head dim of 128 for KDA" - if scale is None: - scale = k.shape[-1] ** -0.5 - with torch.cuda.device_of(q): - o, final_state = HopperChunkKDAFunction.apply( - q, - k, - v, - g, - beta, - A_log, - dt_bias, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - use_gate_in_kernel, - safe_gate, - lower_bound, - cu_seqlens, - chunk_indices, - ) - return o, final_state diff --git a/cula/kda/hopper_fused_fwd_opt.py b/cula/kda/hopper_fused_fwd_opt.py deleted file mode 100644 index a9ddd545..00000000 --- a/cula/kda/hopper_fused_fwd_opt.py +++ /dev/null @@ -1,391 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# 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. - -"""Optimized Hopper KDA prefill: fused gate+l2norm preprocessing + intra-card CP.""" - -import torch -from einops import rearrange -from fla.modules.l2norm import l2norm_fwd -from fla.ops.kda.gate import kda_gate_chunk_cumsum -from fla.ops.utils import chunk_local_cumsum -from fla.ops.utils.constant import RCP_LN2 -from fla.ops.utils.index import prepare_chunk_indices -from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - -import cula.cudac as cula_cuda -from cula.kda.cp_context import intra_card_cp_preprocess, is_dominant_long_seq -from cula.kda.gate_l2norm_fused import gate_l2norm_fused_fwd -from cula.kda.l2norm_qk_fused import l2norm_fwd_qk -from cula.utils import _get_cache_buf, assert_hopper, get_device_sm_count, prepare_uniform_cu_seqlens - -FUSED_L2NORM_QK_TH_MAX = 10000 - -FUSED_GATE_L2NORM_TH_FIXED = 16384 - -FUSED_GATE_L2NORM_TH_VARLEN = 65536 - -FUSED_GATE_L2NORM_VARLEN_AVG_SEQ = 256 - - -def _fused_gate_l2norm_threshold(cu_seqlens_is_none): - return FUSED_GATE_L2NORM_TH_FIXED if cu_seqlens_is_none else FUSED_GATE_L2NORM_TH_VARLEN - - -FUSED_GATE_L2NORM_TH_MAX = FUSED_GATE_L2NORM_TH_VARLEN - - -def _inference_forward( - q, - k, - v, - g, - beta, - A_log, - dt_bias, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - use_gate_in_kernel, - safe_gate, - lower_bound, - cu_seqlens, - chunk_indices, - auto_cp, - cu_seqlens_cpu=None, -): - chunk_size = 64 - batch_size, seq_len, num_qk_heads, head_dim = q.shape - num_v_heads = v.shape[-2] - - cu_seqlens_is_none = cu_seqlens is None - if cu_seqlens_is_none: - cu_seqlens = prepare_uniform_cu_seqlens(batch_size, seq_len, q.device, torch.int32) - if batch_size != 1: - q, k, v, g, beta = map(lambda x: rearrange(x, "b t ... -> 1 (b t) ..."), (q, k, v, g, beta)) - - if cu_seqlens_is_none: - avg_seq_ok = True - else: - N = cu_seqlens.numel() - 1 - packed_T = q.shape[1] - avg_seq_ok = N <= 1 or packed_T <= N * FUSED_GATE_L2NORM_VARLEN_AVG_SEQ - - fused_all_pre = ( - use_gate_in_kernel - and use_qk_l2norm_in_kernel - and (q.numel() // q.shape[-1]) <= _fused_gate_l2norm_threshold(cu_seqlens_is_none) - and num_qk_heads == num_v_heads - and avg_seq_ok - ) - - if fused_all_pre: - if chunk_indices is None and not cu_seqlens_is_none: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) - g_out, yq, yk, _, _ = gate_l2norm_fused_fwd( - g=g, - q=q, - k=k, - A_log=A_log, - dt_bias=dt_bias, - lower_bound=lower_bound if safe_gate else None, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens if not cu_seqlens_is_none else None, - chunk_indices=chunk_indices if not cu_seqlens_is_none else None, - ) - g, q, k = g_out, yq, yk - else: - if chunk_indices is None and not cu_seqlens_is_none: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) - if use_gate_in_kernel: - g = kda_gate_chunk_cumsum( - g=g, - A_log=A_log, - dt_bias=dt_bias, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - lower_bound=lower_bound, - ) - else: - g = chunk_local_cumsum( - g=g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - if use_qk_l2norm_in_kernel: - D = q.shape[-1] - n_rows = q.numel() // D - if n_rows <= FUSED_L2NORM_QK_TH_MAX: - q_flat = q.view(-1, D) - k_flat = k.view(-1, D) - yq_flat, yk_flat, _, _ = l2norm_fwd_qk(q_flat, k_flat) - q = yq_flat.view_as(q) - k = yk_flat.view_as(k) - else: - q, _ = l2norm_fwd(q) - k, _ = l2norm_fwd(k) - - packed_seq = batch_size * seq_len - q = q.reshape(packed_seq, num_qk_heads, head_dim).contiguous() - k = k.reshape(packed_seq, num_qk_heads, head_dim).contiguous() - v = v.reshape(packed_seq, num_v_heads, head_dim).contiguous() - g = g.reshape(packed_seq, num_v_heads, head_dim).contiguous() - beta = beta.reshape(packed_seq, num_v_heads).contiguous() - - cp_seq_map = None - raw_cu_seqlens_for_cp = None - - def _dominant_long_seq_gate() -> bool: - if cu_seqlens is None: - return False - N_ = cu_seqlens.numel() - 1 - if N_ <= 1 or N_ * num_v_heads <= 16: - return False # falls under the main branch - if num_v_heads > 16 or packed_seq < 8192: - return False - cu_src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens - cu_list = cu_src.tolist() - seqlens = [cu_list[i + 1] - cu_list[i] for i in range(N_)] - return is_dominant_long_seq(seqlens, num_v_heads) - - cp_would_fire = auto_cp and ( - # Multi-seq varlen: CP only when grid is starved AND total T amortizes. - ( - cu_seqlens is not None - and (cu_seqlens.numel() - 1) > 1 - and (cu_seqlens.numel() - 1) * num_v_heads <= 16 - and packed_seq >= 8192 - ) - # Multi-seq dominant-seq exception. - or _dominant_long_seq_gate() - or - # Single sequence: per-H T thresholds matching _calc_cp_seqs. - ( - (cu_seqlens is None or cu_seqlens.numel() - 1 == 1) - and ( - (num_v_heads <= 8 and packed_seq >= 4096) - or (num_v_heads <= 16 and packed_seq >= 4096) - or (num_v_heads <= 32 and packed_seq >= 16384) - ) - ) - ) - if cp_would_fire: - q4 = q.view(1, packed_seq, num_qk_heads, head_dim) - k4 = k.view(1, packed_seq, num_qk_heads, head_dim) - v4 = v.view(1, packed_seq, num_v_heads, head_dim) - g4 = g.view(1, packed_seq, num_v_heads, head_dim) - beta4 = beta.view(1, packed_seq, num_v_heads) - cp_h0, cp_cu_seqlens, cp_seq_map, raw_cu_seqlens_for_cp = intra_card_cp_preprocess( - q=q4, - k=k4, - v=v4, - g=g4, - beta=beta4, - scale=scale, - raw_h0=initial_state, - raw_cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - raw_cu_seqlens_cpu=cu_seqlens_cpu, - ) - del q4, k4, v4, g4, beta4 - if cp_seq_map is not None: - cu_seqlens = cp_cu_seqlens - initial_state = cp_h0 - - sm_count = get_device_sm_count(q.device) - workspace_buffer = _get_cache_buf("hopper_kda_fwd_workspace", sm_count * 128, q.device) - - o, final_state = cula_cuda.kda_fwd_prefill( - None, - None, - q, - k, - v, - initial_state, - g, - beta, - cu_seqlens, - workspace_buffer, - scale, - output_final_state, - safe_gate, - cp_seq_map_=cp_seq_map, - raw_cu_seqlens_=raw_cu_seqlens_for_cp, - ) - o = rearrange(o, "(b t) h d -> b t h d", b=batch_size) - return o.to(q.dtype), final_state - - -class HopperChunkKDAFunctionOpt(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - safe_gate: bool = False, - lower_bound: float | None = None, - cu_seqlens: torch.IntTensor | None = None, - chunk_indices: torch.IntTensor | None = None, - auto_cp: bool = True, - cu_seqlens_cpu: torch.IntTensor | None = None, - ): - return _inference_forward( - q, - k, - v, - g, - beta, - A_log, - dt_bias, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - use_gate_in_kernel, - safe_gate, - lower_bound, - cu_seqlens, - chunk_indices, - auto_cp, - cu_seqlens_cpu=cu_seqlens_cpu, - ) - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht): - raise NotImplementedError("Backward pass is not implemented yet.") - - -@torch.compiler.disable -def cula_kda_prefill_opt( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - safe_gate: bool = False, - lower_bound: float | None = None, - cu_seqlens: torch.IntTensor | None = None, - chunk_indices: torch.IntTensor | None = None, - auto_cp: bool = True, - cu_seqlens_cpu: torch.IntTensor | None = None, - **kwargs, -): - assert_hopper(q.device) - assert safe_gate, "Only support safe_gate=True." - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError(f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`.") - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if initial_state is not None: - assert initial_state.dtype == torch.float32, "initial_state must be in float32." - - A_log, dt_bias = None, None - if use_gate_in_kernel: - assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." - A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") - if safe_gate: - if lower_bound is None: - raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") - if not (-5 <= lower_bound < 0): - raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") - - assert q.shape == k.shape, "q and k must have the same shape." - assert q.shape[:2] == v.shape[:2] == g.shape[:2], "q, k, v, g must share batch and sequence dimensions." - batch_size, seq_len, num_qk_heads, head_dim = q.shape - num_v_heads = v.shape[-2] - assert num_qk_heads > 0 and num_v_heads > 0 - assert num_v_heads % num_qk_heads == 0 - assert g.shape == (batch_size, seq_len, num_v_heads, head_dim) - assert v.shape == (batch_size, seq_len, num_v_heads, head_dim) - assert beta.shape == (batch_size, seq_len, num_v_heads) - assert q.dtype == k.dtype == v.dtype == torch.bfloat16, "q, k, v must be in bfloat16." - assert beta.dtype == torch.bfloat16 or beta.dtype == torch.float32, "beta must be in bfloat16 or float32." - assert q.shape[-1] == k.shape[-1] == v.shape[-1] == 128, "Currently we only support head dim of 128 for KDA" - if scale is None: - scale = k.shape[-1] ** -0.5 - - needs_grad = torch.is_grad_enabled() and any(t.requires_grad for t in (q, k, v, g, beta) if t is not None) - with torch.cuda.device_of(q): - if not needs_grad: - o, final_state = _inference_forward( - q, - k, - v, - g, - beta, - A_log, - dt_bias, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - use_gate_in_kernel, - safe_gate, - lower_bound, - cu_seqlens, - chunk_indices, - auto_cp, - cu_seqlens_cpu=cu_seqlens_cpu, - ) - return o, (final_state if output_final_state else None) - - o, final_state = HopperChunkKDAFunctionOpt.apply( - q, - k, - v, - g, - beta, - A_log, - dt_bias, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - use_gate_in_kernel, - safe_gate, - lower_bound, - cu_seqlens, - chunk_indices, - auto_cp, - cu_seqlens_cpu, - ) - - return o, (final_state if output_final_state else None) diff --git a/cula/kda/l2norm_qk_fused.py b/cula/kda/l2norm_qk_fused.py deleted file mode 100644 index 234d57bf..00000000 --- a/cula/kda/l2norm_qk_fused.py +++ /dev/null @@ -1,119 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# 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 -"""Fused l2-norm for paired (q, k) tensors — one Triton kernel handles both. - -cuLA's baseline `cula_kda_prefill` calls `l2norm_fwd(q)` and `l2norm_fwd(k)` -as two separate Triton kernel launches. Each launch costs ~50-80 μs of CPU -overhead (Python wrapper + torch.empty + CUDA driver dispatch), even though -the actual GPU work is tiny for D=128. - -The two operations are mathematically identical and operate on disjoint -inputs/outputs. By writing one kernel whose grid covers both q and k -(distinguishing them via `tl.program_id(1)`), we cut the Python-driver -overhead in half — saving one launch (~50 μs) per fwd at any T. - -Combined with skipping the gate-stream optimization (which didn't pay off -due to torch.cuda.stream context overhead — see hopper_fused_fwd_opt.py), -this is the cleanest small-T speedup we found. -""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _l2norm_fwd_qk_kernel( - q_ptr, - k_ptr, # input pointers (T*H, D) each - yq_ptr, - yk_ptr, # output pointers (T*H, D) each - rstd_q_ptr, - rstd_k_ptr, # output rstd (T*H,) each - eps, - D, - BD: tl.constexpr, -): - i_row = tl.program_id(0) # 0..T*H-1 - i_qk = tl.program_id(1) # 0=q, 1=k - - cols = tl.arange(0, BD) - mask = cols < D - - # is_q is uniform across all threads in the block (driven by - # tl.program_id(1)), so the if/else compiles to a single conditional - # branch — no warp divergence, and only one tensor is actually loaded. - is_q = i_qk == 0 - base_off = i_row * D - if is_q: - b_x = tl.load(q_ptr + base_off + cols, mask=mask, other=0.0).to(tl.float32) - else: - b_x = tl.load(k_ptr + base_off + cols, mask=mask, other=0.0).to(tl.float32) - - b_rstd = 1.0 / tl.sqrt(tl.sum(b_x * b_x) + eps) - b_y = b_x * b_rstd - - # Symmetric stores — same uniform-branch logic as the load above. - if is_q: - tl.store(yq_ptr + base_off + cols, b_y, mask=mask) - tl.store(rstd_q_ptr + i_row, b_rstd) - else: - tl.store(yk_ptr + base_off + cols, b_y, mask=mask) - tl.store(rstd_k_ptr + i_row, b_rstd) - - -def l2norm_fwd_qk( - q: torch.Tensor, - k: torch.Tensor, - eps: float = 1e-6, -): - """L2-normalize q and k along the last dim, in a single fused kernel. - - Args: - q, k: shape (..., D). Last dim is normalised; preceding dims are - treated as a flat "row" index. q and k must have identical shape. - eps: numerical safety eps. - - Returns: - (y_q, y_k, rstd_q, rstd_k) - y_q, y_k: normalized outputs, same shape as q, k. - rstd_q, rstd_k: 1/sqrt(sum(x^2)+eps), shape q.shape[:-1]. - """ - assert q.shape == k.shape, f"q.shape {q.shape} != k.shape {k.shape}" - assert q.dtype == k.dtype - assert q.device == k.device - assert q.is_contiguous() and k.is_contiguous(), "q, k must be contiguous" - - D = q.shape[-1] - T = q.numel() // D - - y_q = torch.empty_like(q) - y_k = torch.empty_like(k) - rstd_q = torch.empty(q.shape[:-1], dtype=torch.float32, device=q.device) - rstd_k = torch.empty(k.shape[:-1], dtype=torch.float32, device=k.device) - - BD = triton.next_power_of_2(D) - if D > BD or 65536 // q.element_size() < BD: - raise RuntimeError(f"D={D} too large for fused l2norm_fwd_qk") - - # Grid: (T*H, 2) — program_id(1) picks q (0) or k (1). - # Heuristic on num_warps: small D (e.g. 128) needs only 1 warp; up to 4. - num_warps = 1 if BD <= 256 else (2 if BD <= 1024 else 4) - _l2norm_fwd_qk_kernel[(T, 2)]( - q_ptr=q, - k_ptr=k, - yq_ptr=y_q, - yk_ptr=y_k, - rstd_q_ptr=rstd_q, - rstd_k_ptr=rstd_k, - eps=eps, - D=D, - BD=BD, - num_warps=num_warps, - ) - return y_q, y_k, rstd_q, rstd_k diff --git a/cula/kda/backends/__init__.py b/cula/kda/prefill.py similarity index 78% rename from cula/kda/backends/__init__.py rename to cula/kda/prefill.py index ce79381e..fc713e12 100644 --- a/cula/kda/backends/__init__.py +++ b/cula/kda/prefill.py @@ -1,15 +1,14 @@ # Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 +"""Public FlashKDA prefill entry point.""" + from __future__ import annotations from numbers import Real import torch -from cula.backends import BackendRegistry -from cula.kda.backends.flashkda import FlashKDABackend -from cula.kda.backends.fully_fused import FullyFusedBackend from cula.ops.kda.cp_mode import CPMode _EXTRA_KWARGS = { @@ -22,10 +21,6 @@ "final_state", } -kda_prefill_registry = BackendRegistry("kda_prefill") -kda_prefill_registry.register(FlashKDABackend()) -kda_prefill_registry.register(FullyFusedBackend()) - @torch.compiler.disable def kda_prefill( @@ -45,11 +40,12 @@ def kda_prefill( chunk_indices: torch.IntTensor | None = None, **kwargs, ): - """Select the first SM90 prefill backend that accepts the call. + """Run the CuTeDSL FlashKDA forward prefill implementation. - This forward-only API does not implement backward. FlashKDA has priority - over the fully-fused CUDA path; disabling either backend with its - ``CULA_BACKEND_*`` variable makes dispatch try the remaining implementation. + This forward-only API supports SM90 and SM100. It has no native-extension + fallback: unsupported inputs are rejected by FlashKDA with a concrete + validation error. In particular, grouped-value attention is not currently + supported and the Q/K/V head counts must match. """ tensors = (q, k, v, g, beta) if any(not isinstance(t, torch.Tensor) for t in tensors): @@ -75,8 +71,9 @@ def kda_prefill( if "use_intracard_cp" in kwargs: kwargs["use_intracard_cp"] = CPMode.parse(kwargs["use_intracard_cp"]) - return kda_prefill_registry.dispatch( - "kda_prefill", + from cula.kda.flashkda import cula_kda_prefill + + return cula_kda_prefill( q, k, v, @@ -95,4 +92,4 @@ def kda_prefill( ) -__all__ = ["kda_prefill", "kda_prefill_registry"] +__all__ = ["kda_prefill"] diff --git a/cula/ops/kda/cp_mode.py b/cula/ops/kda/cp_mode.py index 6135de18..9a8a589a 100644 --- a/cula/ops/kda/cp_mode.py +++ b/cula/ops/kda/cp_mode.py @@ -1,7 +1,7 @@ # Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 -"""Intracard context-parallel dispatch vocabulary shared by the KDA backends.""" +"""Intracard context-parallel dispatch vocabulary shared by KDA paths.""" from __future__ import annotations diff --git a/docs/kda_sm100_cutedsl_fwd_status.md b/docs/kda_sm100_cutedsl_fwd_status.md index d93af578..d5a69be6 100644 --- a/docs/kda_sm100_cutedsl_fwd_status.md +++ b/docs/kda_sm100_cutedsl_fwd_status.md @@ -2,8 +2,8 @@ The SM100/SM103 modular KDA forward no longer builds or loads a C++/CUDA extension. Intra-chunk Aqk/Akk, the csrc-ordered TF32 inverse, recompute W/U, -KG, and optional QG are implemented in CuTeDSL. The remaining native cuLA -extension is SM90-only. +KG, and optional QG are implemented in CuTeDSL. cuLA no longer builds any +native extension; SM90 prefill uses FlashKDA's CuTeDSL K1+K2 pipeline. ## Compatibility and dispatch @@ -116,5 +116,5 @@ Final pure-path results: - default-stream and concurrent-stream isolation; - output lifetime, pool reuse, and bounded retained memory; - fixed, GVA, QG, packed-varlen, and tail routing; -- installation without an SM100 native extension; -- SM90 extension build retained independently. +- installation from the platform-independent wheel; +- SM90 FlashKDA prefill without a native fallback. diff --git a/pyproject.toml b/pyproject.toml index 520c3e80..60cbd2b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,12 +1,12 @@ [build-system] -requires = ["setuptools>=61", "wheel", "setuptools_scm>=6.0", "torch"] +requires = ["setuptools>=61", "wheel", "setuptools_scm>=6.0"] build-backend = "setuptools.build_meta" [project] name = "cuda-linear-attention" dynamic = ["version"] -description = "cuLA CUDA extension" +description = "High-performance CuTeDSL linear-attention kernels" readme = "README.md" authors = [ { name = "cula contributors" } ] requires-python = ">=3.10" @@ -15,7 +15,7 @@ dependencies = [ "apache-tvm-ffi>=0.1.9", "packaging>=23.2", ] -license = { text = "Apache-2.0" } +license = "Apache-2.0" [project.optional-dependencies] dev = [ @@ -33,10 +33,7 @@ include = ["cula", "cula.*"] [tool.ruff] target-version = "py312" line-length = 127 -exclude = [ - "csrc/", - "third_party/", -] +exclude = ["third_party/"] [tool.ruff.lint] select = [ diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index 79ac3305..daffe095 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -24,9 +24,9 @@ for arg in "$@"; do ISOLATION_FLAG="" echo "[build_wheel] Using isolated build environment" ;; - --fat) - export CULA_BUILD_ALL_ARCHS=1 - echo "[build_wheel] Fat binary: building for all SM architectures" + *) + echo "[build_wheel] Unknown argument: $arg" >&2 + exit 2 ;; esac done @@ -40,9 +40,7 @@ rm -rf dist build *.egg-info # Show environment info echo "[build_wheel] Python: $(python -V 2>&1)" -echo "[build_wheel] torch: $(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo 'not installed')" -echo "[build_wheel] CUDA: $(nvcc --version 2>/dev/null | grep 'release' | sed 's/.*release //' | sed 's/,.*//' || echo 'not found')" -echo "[build_wheel] Fat binary: ${CULA_BUILD_ALL_ARCHS:-0}" +echo "[build_wheel] Packaging: pure Python (CuTeDSL kernels compile at runtime)" # Build wheel echo "[build_wheel] Building wheel..." diff --git a/setup.py b/setup.py index 329d062f..c6cd7ac5 100644 --- a/setup.py +++ b/setup.py @@ -1,163 +1,6 @@ -import os -import subprocess -from pathlib import Path - from setuptools import find_packages, setup -from torch.utils.cpp_extension import ( - CUDA_HOME, - IS_WINDOWS, - BuildExtension, - CUDAExtension, -) - - -def detect_gpu_archs() -> tuple[bool, bool, bool]: - """ - Query all visible CUDA devices via torch and return three booleans: - - has_sm100: major version numbers == 10 with minor == 0 (sm100) - - has_sm103: major version numbers == 10 with minor == 3 (sm103) - - has_sm90: major version numbers == 9 with minor == 0 (sm90a) - Returns (has_sm100, has_sm103, has_sm90) - """ - try: - import torch - - if not torch.cuda.is_available(): - return False, False, False - has_sm100 = False - has_sm103 = False - has_sm90 = False - for i in range(torch.cuda.device_count()): - prop = torch.cuda.get_device_properties(i) - major, minor = prop.major, prop.minor - print(f" GPU {i}: {prop.name}, compute capability sm_{major}{minor}") - if major == 10 and minor == 0: - has_sm100 = True - if major == 10 and minor == 3: - has_sm103 = True - if major == 9 and minor == 0: - has_sm90 = True - return has_sm100, has_sm103, has_sm90 - except Exception as e: - print(f"Warning: failed to detect GPU architectures via torch: {e}") - return False, False, False - - -def resolve_disable_flag(env_name: str, detected: bool) -> bool: - """ - Resolve whether to disable a given SM target. - - If CULA_BUILD_ALL_ARCHS is set, all targets are enabled unconditionally. - - If the environment variable is explicitly set, honour it. - - Otherwise, disable the target when no matching GPU is detected. - """ - if os.getenv("CULA_BUILD_ALL_ARCHS", "0") == "1": - return False - env_val = os.getenv(env_name) - if env_val is not None: - return env_val.lower() in ["true", "1", "y", "yes"] - disable = not detected - if disable: - print(f" No matching GPU detected; auto-setting {env_name}=1 (disable). Set {env_name}=0 to override.") - return disable - - -def get_features_args(): - features_args = [] - return features_args - - -USE_FAST_MATH = os.getenv("CULA_USE_FAST_MATH", "1") == "1" - -if os.getenv("CULA_BUILD_ALL_ARCHS", "0") == "1": - print("CULA_BUILD_ALL_ARCHS=1: enabling the SM90 native extension; SM100/SM103 use CuTeDSL JIT") -else: - print("Detecting GPU architectures...") - -_has_sm100, _has_sm103, _has_sm90 = detect_gpu_archs() -DISABLE_SM90 = resolve_disable_flag("CULA_DISABLE_SM90", _has_sm90) - - -def get_nvcc_version() -> tuple[int, int]: - """Return the NVCC (major, minor) version tuple, e.g. ``(12, 9)``.""" - # NOTE The "CUDA_HOME" here is not necessarily from the `CUDA_HOME` environment variable. - # For more details, see `torch/utils/cpp_extension.py` - assert CUDA_HOME is not None, "PyTorch must be compiled with CUDA support" - nvcc_version = subprocess.check_output( - [os.path.join(CUDA_HOME, "bin", "nvcc"), "--version"], stderr=subprocess.STDOUT - ).decode("utf-8") - nvcc_version_number = nvcc_version.split("release ")[1].split(",")[0].strip() - major, minor = map(int, nvcc_version_number.split(".")) - return major, minor - - -def get_nvcc_thread_args(): - nvcc_threads = os.getenv("NVCC_THREADS") or "32" - return ["--threads", nvcc_threads] - - -this_dir = os.path.dirname(os.path.abspath(__file__)) - -if IS_WINDOWS: - cxx_args = ["/O2", "/std:c++20", "/DNDEBUG", "/W0"] -else: - cxx_args = ["-O3", "-std=c++20", "-DNDEBUG", "-Wno-deprecated-declarations"] - -nvcc_common_args = [ - "-O3", - "-std=c++20", - "-DNDEBUG", - # "-D_USE_MATH_DEFINES", - "-Wno-deprecated-declarations", - "-U__CUDA_NO_HALF_OPERATORS__", - "-U__CUDA_NO_HALF_CONVERSIONS__", - "-U__CUDA_NO_HALF2_OPERATORS__", - "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", - "--expt-relaxed-constexpr", - "--expt-extended-lambda", - "-lineinfo", - "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage", - "-diag-suppress=3189", -] - -include_dirs = [ - Path(this_dir) / "csrc", - Path(this_dir) / "csrc" / "kerutils" / "include", - Path(this_dir) / "csrc" / "cutlass" / "include", - Path(this_dir) / "csrc" / "cutlass" / "tools" / "util" / "include", -] - -ext_modules = [] - -if not DISABLE_SM90: - subprocess.run(["git", "submodule", "update", "--init", "csrc/cutlass"], check=True) - major, minor = get_nvcc_version() - print(f"Compiling the SM90 extension using NVCC {major}.{minor}") - sm90_arch_flags = ["-gencode", "arch=compute_90a,code=sm_90a", "-DCULA_SM90A_ENABLED"] - - ext_modules.append( - CUDAExtension( - name="cula._cudac_sm90", - sources=[ - "csrc/api/kda_sm90.cu", - "csrc/api/pybind_sm90.cu", - "csrc/kda/sm90/kda_fwd_sm90.cu", - "csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu", - ], - extra_compile_args={ - "cxx": cxx_args + get_features_args(), - "nvcc": nvcc_common_args - + get_features_args() - + sm90_arch_flags - + get_nvcc_thread_args() - + (["--use_fast_math"] if USE_FAST_MATH else []), - }, - include_dirs=include_dirs, - ) - ) setup( name="cuda-linear-attention", packages=find_packages(include=["cula", "cula.*"]), - ext_modules=ext_modules, - cmdclass={"build_ext": BuildExtension}, ) diff --git a/tests/test_intracard_cp_sm90.py b/tests/test_intracard_cp_sm90.py deleted file mode 100644 index 2353a66e..00000000 --- a/tests/test_intracard_cp_sm90.py +++ /dev/null @@ -1,439 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2025-2026 Ant Group Co., Ltd. -# Licensed under the Apache License, Version 2.0. -"""Tests for SM90 intra-card CP: dispatch routing + numerical accuracy. - -Mirrors tests/test_intracard_cp.py (SM100 version) but targets the Hopper -(SM90) `kda_prefill_hopper_opt` / `kda_prefill_hopper_auto` path. - -Three reference levels: - - cuLA basic (kda_prefill_hopper) — same C++ kernel, no CP scheduling - → verifies CP scheduling is value-preserving - - cuLA opt with auto_cp=False — opt Python wrapper but CP disabled - → isolates the CP code paths - - FLA chunk_kda (cross-impl reference) → source of truth for end-to-end output - -The CP path is exercised through: - - kda_prefill_hopper_auto (router picks opt when shape benefits from CP) - - kda_prefill_hopper_opt(auto_cp=True) (force CP entry; bypasses router) -""" - -from __future__ import annotations - -import math -import pathlib -import sys - -import pytest -import torch - -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from fla.ops.kda import chunk_kda as fla_chunk_kda # noqa: E402 -from fla.utils import assert_close # noqa: E402 - -from cula.kda import ( # noqa: E402 - kda_prefill_hopper, - kda_prefill_hopper_auto, - kda_prefill_hopper_opt, -) -from cula.kda.cp_context import _calc_cp_seqs # noqa: E402 -from cula.utils import get_device_sm_count # noqa: E402 - -BT, D = 64, 128 -DEVICE = "cuda" -DTYPE = torch.bfloat16 -LOWER_BOUND = -5.0 - -# Tolerances — same convention as tests/test_intracard_cp.py: -# * Same-kernel (CP scheduling only): torch.testing.assert_close -# (CP-on vs CP-off both go through cuLA kernels) -# * Cross-impl (vs FLA): fla.utils.assert_close(ratio=...) -ATOL_SAME_KERNEL = 1e-2 -RTOL_SAME_KERNEL = 1e-2 -RATIO_VS_FLA = 0.015 # bf16 cross-impl noise band (matches SM100 test) -RATIO_STRESS = 1e-6 # deterministic re-run: drift implies race - - -pytestmark = [ - pytest.mark.sm90_only, - pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), -] - - -# ============================== Helpers ============================== - - -def _cu_from_seq_lens(seq_lens, device=DEVICE): - cu = [0] - for s in seq_lens: - cu.append(cu[-1] + s) - return torch.tensor(cu, dtype=torch.int32, device=device) - - -def make_varlen_inputs(seq_lens, H, *, use_h0=False, seed=42): - """Build varlen-packed B=1 inputs for kda_prefill_hopper_*.""" - total = sum(seq_lens) - N = len(seq_lens) - cu = _cu_from_seq_lens(seq_lens) - torch.manual_seed(seed) - q = torch.randn(1, total, H, D, dtype=DTYPE, device=DEVICE) - k = torch.randn(1, total, H, D, dtype=DTYPE, device=DEVICE) - v = torch.randn(1, total, H, D, dtype=DTYPE, device=DEVICE) - g = -torch.rand(1, total, H, D, dtype=torch.float32, device=DEVICE).abs() * 0.5 - beta = torch.randn(1, total, H, dtype=torch.float32, device=DEVICE).sigmoid().to(DTYPE) - A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) - dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) - h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) * 0.1 if use_h0 else None - return q, k, v, g, beta, h0, A_log, dt_bias, cu - - -# ---- entry points under test ---- - - -def _common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu): - return dict( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=1.0 / math.sqrt(D), - A_log=A_log, - dt_bias=dt_bias, - initial_state=h0, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - safe_gate=True, - lower_bound=LOWER_BOUND, - cu_seqlens=cu, - ) - - -def run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu): - return kda_prefill_hopper(**_common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu)) - - -def run_cula_opt_no_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu): - return kda_prefill_hopper_opt( - **_common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu), - auto_cp=False, - ) - - -def run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu): - """Force CP entry through opt wrapper.""" - return kda_prefill_hopper_opt( - **_common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu), - auto_cp=True, - ) - - -def run_cula_auto(q, k, v, g, beta, h0, A_log, dt_bias, cu): - """Adaptive router — exercises the production entry point.""" - return kda_prefill_hopper_auto( - **_common_cula_kw(q, k, v, g, beta, h0, A_log, dt_bias, cu), - ) - - -def run_fla(q, k, v, g, beta, h0, A_log, dt_bias, cu): - """FLA reference. cuLA returns ht as [N, HV, V, K]; FLA's default layout - is [N, HV, K, V]. We pass ``transpose_state_layout=True`` so its output - matches cuLA's layout — no manual transpose needed before assert_close. - """ - return fla_chunk_kda( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=1.0 / math.sqrt(D), - A_log=A_log, - dt_bias=dt_bias, - initial_state=h0, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - safe_gate=True, - lower_bound=LOWER_BOUND, - cu_seqlens=cu.long(), - transpose_state_layout=True, - ) - - -# ---- assertions ---- - - -def _assert_same_kernel(name, actual, ref): - """torch.testing.assert_close with tight atol/rtol (CP-on vs CP-off use the - same C++ kernel; the only delta is per-chunk recurrence reordering).""" - if actual is None or ref is None: - assert actual is ref, f"{name}: one is None and the other isn't" - return - torch.testing.assert_close( - actual.float(), - ref.float(), - atol=ATOL_SAME_KERNEL, - rtol=RTOL_SAME_KERNEL, - msg=lambda m: f"{name}: {m}", - ) - - -def assert_cp_engages(cu, H): - """Fail fast if _calc_cp_seqs won't engage CP for this shape — without - that, the test silently checks CP-off vs CP-off. - """ - num_sms = get_device_sm_count(torch.device(DEVICE)) - use_cp, cp_cu, *_ = _calc_cp_seqs( - cu, - BT, - H, - num_sms, - raw_cu_seqlens_cpu=cu.cpu(), - ) - assert use_cp and cp_cu is not None, f"_calc_cp_seqs returned use_cp=False for cu={cu.tolist()} H={H}" - n_sub = int(cp_cu.numel() - 1) - raw_batch = int(cu.numel() - 1) - assert n_sub > raw_batch, f"CP didn't split: n_sub={n_sub} == raw_batch={raw_batch}" - - -# ====================== Dispatch path: CP vs no-CP ====================== -# Verifies kda_prefill_hopper_opt(auto_cp=True) routes through CP and matches -# the same-kernel no-CP baseline (kda_prefill_hopper). - -DISPATCH_CONFIGS = [ - # (seq_lens, H, use_h0) - ([32768], 4, False), - ([32768], 4, True), - ([65536], 4, True), - ([32768], 8, False), - ([65536], 8, True), - ([16384, 16384], 4, True), - ([28672, 4096], 4, True), - ([131072, 1024], 4, False), -] - - -@pytest.mark.parametrize("seq_lens,H,use_h0", DISPATCH_CONFIGS) -def test_cp_matches_basic_baseline(seq_lens, H, use_h0): - """CP-on (opt+auto_cp) output equals basic baseline (no-CP).""" - q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs( - seq_lens, - H, - use_h0=use_h0, - ) - assert_cp_engages(cu, H) - with torch.inference_mode(): - o_base, ht_base = run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu) - o_cp, ht_cp = run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) - _assert_same_kernel("o", o_cp, o_base) - _assert_same_kernel("ht", ht_cp, ht_base) - - -@pytest.mark.parametrize("seq_lens,H,use_h0", DISPATCH_CONFIGS) -def test_auto_router_matches_basic_baseline(seq_lens, H, use_h0): - """kda_prefill_hopper_auto output (whatever path it picks) equals basic baseline.""" - q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs( - seq_lens, - H, - use_h0=use_h0, - ) - with torch.inference_mode(): - o_base, ht_base = run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu) - o_auto, ht_auto = run_cula_auto(q, k, v, g, beta, h0, A_log, dt_bias, cu) - _assert_same_kernel("o", o_auto, o_base) - _assert_same_kernel("ht", ht_auto, ht_base) - - -def test_cp_off_matches_basic_baseline(): - """opt with auto_cp=False must match basic (no CP, no fused-pre divergence).""" - seq_lens, H, use_h0 = [32768], 4, True - q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=use_h0) - with torch.inference_mode(): - o_base, ht_base = run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu) - o_off, ht_off = run_cula_opt_no_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) - _assert_same_kernel("o", o_off, o_base) - _assert_same_kernel("ht", ht_off, ht_base) - - -# ====================== Cross-impl: CP vs FLA ====================== - -VS_FLA_CONFIGS = [ - ([32768], 4), - ([65536], 4), - ([32768], 8), - ([16384, 16384], 4), - ([28672, 4096], 4), - ([131072, 1024], 4), -] - - -# Irregular varlen lengths -IRREGULAR_VARLEN_CONFIGS = [ - ([1], 4), - ([63], 4), - ([64], 4), - ([65], 4), - ([129], 4), - ([1, 63, 64, 65, 129], 4), - ([1, 63, 64, 65, 129], 8), - ([129, 65, 64, 63, 1], 4), - ([1024, 1, 63, 65, 129], 4), - ([4096, 1, 63, 64, 65, 129], 4), - ([4096, 1, 63, 64, 65, 129], 8), - ([8192, 1, 31, 63, 65, 127, 129, 255], 4), - ([1] * 8 + [63] * 4 + [129] * 2, 4), - ([255, 257, 511, 513], 4), -] - - -@pytest.mark.parametrize("seq_lens,H", IRREGULAR_VARLEN_CONFIGS) -def test_irregular_varlen_vs_fla(seq_lens, H): - """Irregular varlen lengths""" - q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=False) - with torch.inference_mode(): - o_fla, ht_fla = run_fla(q, k, v, g, beta, None, A_log, dt_bias, cu) - o_opt, ht_opt = run_cula_opt_no_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) - assert_close(f"o (cu={cu.tolist()},H={H})", o_fla, o_opt, ratio=RATIO_VS_FLA) - assert_close(f"ht (cu={cu.tolist()},H={H})", ht_fla, ht_opt, ratio=RATIO_VS_FLA) - - -@pytest.mark.parametrize("seq_lens,H", IRREGULAR_VARLEN_CONFIGS) -def test_irregular_varlen_opt_matches_basic(seq_lens, H): - """Irregular varlen: opt path (may take fused gate+l2norm) equals basic baseline.""" - q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=False) - with torch.inference_mode(): - o_base, ht_base = run_cula_basic(q, k, v, g, beta, None, A_log, dt_bias, cu) - o_opt, ht_opt = run_cula_opt_no_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) - _assert_same_kernel("o", o_opt, o_base) - _assert_same_kernel("ht", ht_opt, ht_base) - - -@pytest.mark.parametrize("seq_lens,H", VS_FLA_CONFIGS) -def test_cp_vs_fla(seq_lens, H): - """CP output matches FLA chunk_kda reference (cross-impl).""" - q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=False) - assert_cp_engages(cu, H) - with torch.inference_mode(): - o_fla, ht_fla = run_fla(q, k, v, g, beta, None, A_log, dt_bias, cu) - o_cp, ht_cp = run_cula_opt_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) - assert_close(f"o (cu={cu.tolist()},H={H})", o_fla, o_cp, ratio=RATIO_VS_FLA) - assert_close(f"ht (cu={cu.tolist()},H={H})", ht_fla, ht_cp, ratio=RATIO_VS_FLA) - - -# ====================== Final state ht correctness ====================== -# Per-sequence ht must be independently correct for prefill→decode handoff. - -FINAL_STATE_CONFIGS = [ - ([65536], 4, False), - ([65536], 4, True), - ([65536, 16384], 4, True), - ([28672, 4096], 4, False), - ([131072, 1024], 4, True), -] - - -@pytest.mark.parametrize("seq_lens,H,use_h0", FINAL_STATE_CONFIGS) -def test_cp_final_state_per_seq(seq_lens, H, use_h0): - """Each sequence's ht matches basic baseline independently (no cross-leakage).""" - q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=use_h0) - assert_cp_engages(cu, H) - with torch.inference_mode(): - _, ht_base = run_cula_basic(q, k, v, g, beta, h0, A_log, dt_bias, cu) - _, ht_cp = run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) - assert ht_cp is not None and ht_cp.shape == ht_base.shape, ( - f"shape mismatch: cp={tuple(ht_cp.shape)} base={tuple(ht_base.shape)}" - ) - for i in range(len(seq_lens)): - _assert_same_kernel(f"ht[{i}] (len={seq_lens[i]})", ht_cp[i], ht_base[i]) - - -# ====================== Stress: race / non-determinism ====================== -# CP's per-chunk preprocess + main kernel — re-running same inputs must -# produce bit-identical outputs (no race, no order-dependence). - -STRESS_ITERS = 50 - - -@pytest.mark.parametrize( - "seq_lens,H,use_h0", - [ - pytest.param([65536], 4, True, id="single-64K-H4-h0"), - pytest.param([65536, 4096], 4, True, id="multi-64K+4K-H4-h0"), - ], -) -def test_cp_stress_repeat(seq_lens, H, use_h0): - """Run CP N times; every iter must match the first (deterministic).""" - q, k, v, g, beta, h0, A_log, dt_bias, cu = make_varlen_inputs( - seq_lens, - H, - use_h0=use_h0, - seed=20260516, - ) - assert_cp_engages(cu, H) - with torch.inference_mode(): - o_ref, ht_ref = run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) - torch.cuda.synchronize() - for i in range(STRESS_ITERS): - o_i, ht_i = run_cula_opt_cp(q, k, v, g, beta, h0, A_log, dt_bias, cu) - torch.cuda.synchronize() - assert_close(f"iter {i} o", o_ref, o_i, ratio=RATIO_STRESS) - assert_close(f"iter {i} ht", ht_ref, ht_i, ratio=RATIO_STRESS) - - -# ====================== h0=None equivalence ====================== -# We patched cp_context.py so raw_h0=None synthesizes a zero pool. Verify the -# kernel result is numerically equivalent to passing an explicit zero h0. - - -def test_cp_h0_none_equiv_h0_zeros(): - """h0=None must produce identical ht to h0=zeros (no implicit init drift).""" - seq_lens, H = [65536, 4096], 4 - assert_cp_engages(_cu_from_seq_lens(seq_lens), H) - q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs( - seq_lens, - H, - use_h0=False, - seed=20260501, - ) - h0_zeros = torch.zeros(len(seq_lens), H, D, D, dtype=torch.float32, device=DEVICE) - with torch.inference_mode(): - o_none, ht_none = run_cula_opt_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) - o_zeros, ht_zeros = run_cula_opt_cp(q, k, v, g, beta, h0_zeros, A_log, dt_bias, cu) - torch.cuda.synchronize() - o_diff = (o_none.float() - o_zeros.float()).abs().max().item() - ht_diff = (ht_none.float() - ht_zeros.float()).abs().max().item() - assert o_diff < 1e-3, f"o: h0=None vs h0=zeros max abs diff {o_diff:.4e}" - assert ht_diff < 1e-4, f"ht: h0=None vs h0=zeros max abs diff {ht_diff:.4e}" - - -# ====================== CP bypass: shapes where _calc_cp_seqs returns False ====================== -# When CP heuristic says "don't split", auto_cp=True must produce bit-identical -# output to basic (because the kernel takes the same path). - -BYPASS_CONFIGS = [ - ([2048], 8), # H=8 single seq T<=2048 → no CP - ([16384], 64), # H=64 → CP never fires (per _calc_cp_seqs H>=64 branch) - ([4096, 4096, 4096, 4096, 4096, 4096, 4096, 4096], 8), # native_grid 64 >> 16 - ([131072] + [1024] * 5, 8), # raw_batch big enough that native_grid > 16 -] - - -@pytest.mark.parametrize("seq_lens,H", BYPASS_CONFIGS) -def test_cp_bypass_matches_basic(seq_lens, H): - """When CP heuristic skips, auto_cp=True must be a no-op (same output as basic).""" - q, k, v, g, beta, _, A_log, dt_bias, cu = make_varlen_inputs(seq_lens, H, use_h0=False) - num_sms = get_device_sm_count(torch.device(DEVICE)) - use_cp, cp_cu, *_ = _calc_cp_seqs(cu, BT, H, num_sms, raw_cu_seqlens_cpu=cu.cpu()) - n_sub = int(cp_cu.numel() - 1) if cp_cu is not None else 0 - assert not use_cp or n_sub == len(seq_lens), ( - f"expected bypass for cu={cu.tolist()} H={H}, got use_cp={use_cp} n_sub={n_sub}" - ) - with torch.inference_mode(): - o_base, ht_base = run_cula_basic(q, k, v, g, beta, None, A_log, dt_bias, cu) - o_cp, ht_cp = run_cula_opt_cp(q, k, v, g, beta, None, A_log, dt_bias, cu) - _assert_same_kernel("o", o_cp, o_base) - _assert_same_kernel("ht", ht_cp, ht_base) diff --git a/tests/test_kda_prefill_api.py b/tests/test_kda_prefill_api.py new file mode 100644 index 00000000..6e575c57 --- /dev/null +++ b/tests/test_kda_prefill_api.py @@ -0,0 +1,61 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest +import torch + +import cula.kda as kda +from cula.kda.prefill import kda_prefill +from cula.ops.kda.cp_mode import CPMode + + +def _inputs(): + q = torch.empty(1, 2, 1, 128) + return q, q.clone(), q.clone(), q.clone(), torch.empty(1, 2, 1) + + +def test_public_prefill_routes_directly_to_flashkda(monkeypatch): + expected = (object(), object()) + seen = {} + + def fake_prefill(*args, **kwargs): + seen["args"] = args + seen["kwargs"] = kwargs + return expected + + monkeypatch.setitem(sys.modules, "cula.kda.flashkda", SimpleNamespace(cula_kda_prefill=fake_prefill)) + q, k, v, g, beta = _inputs() + + actual = kda_prefill( + q, + k, + v, + g, + beta, + A_log=torch.empty(1), + dt_bias=torch.empty(1, 128), + use_intracard_cp="auto", + ) + + assert actual is expected + assert all(actual is expected for actual, expected in zip(seen["args"], (q, k, v, g, beta), strict=True)) + assert seen["kwargs"]["use_intracard_cp"] is CPMode.AUTO + assert seen["kwargs"]["safe_gate"] is True + assert seen["kwargs"]["use_gate_in_kernel"] is True + assert seen["kwargs"]["use_qk_l2norm_in_kernel"] is True + + +def test_removed_native_entry_points_are_not_exported(): + assert "kda_prefill_hopper" not in kda.__all__ + assert "kda_prefill_hopper_opt" not in kda.__all__ + assert "kda_prefill_hopper_auto" not in kda.__all__ + + +def test_public_prefill_rejects_unknown_keyword_before_launch(): + with pytest.raises(TypeError, match="unexpected keyword"): + kda_prefill(*_inputs(), unknown_native_option=True) diff --git a/tests/test_kda_sm100_output.py b/tests/test_kda_sm100_output.py index e7ec698b..782f3caf 100644 --- a/tests/test_kda_sm100_output.py +++ b/tests/test_kda_sm100_output.py @@ -32,7 +32,7 @@ import pytest import torch -# Import directly from the module file to avoid cula package __init__ (requires cudac) +# Import directly from the module file so this standalone test controls loading. sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import importlib.util diff --git a/tests/test_kda_sm100_recurrence.py b/tests/test_kda_sm100_recurrence.py index f521497e..3486b2d3 100644 --- a/tests/test_kda_sm100_recurrence.py +++ b/tests/test_kda_sm100_recurrence.py @@ -17,7 +17,7 @@ # ─── FLA reference ─── from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h as fla_fwd_h -# ─── CuTe DSL kernel (via importlib to avoid cula __init__ requiring cudac) ─── +# ─── CuTe DSL kernel (loaded directly for this standalone test) ─── sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import importlib.util diff --git a/tests/test_kda_sm90_intracard_cp.py b/tests/test_kda_sm90_intracard_cp.py index 438018d7..aafdb7f7 100644 --- a/tests/test_kda_sm90_intracard_cp.py +++ b/tests/test_kda_sm90_intracard_cp.py @@ -3,8 +3,8 @@ """FlashKDA CuTeDSL intracard CP (SM90): CP vs serial, CP vs FLA, determinism. -Distinct from tests/test_intracard_cp_sm90.py, which covers the CUDA C++ -fully-fused path (kda_prefill_hopper_opt / auto_cp via cp_context). +This is the production SM90 prefill path; the former CUDA C++ implementation +and its separate CP scheduler have been removed. """ import pytest From 0c2e4af68910a864e94a89c61b4c1fea2ff22db6 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 15:13:26 +0800 Subject: [PATCH 02/10] fix(cutedsl): support modern compiler releases Raise the supported runtime to CuTeDSL 4.6.2 through 4.8 and normalize the NVVM API changes introduced in 4.8. Use the native CuTe atomic operation for persistent work scheduling. Signed-off-by: shuyan.ycf --- README.md | 2 +- cula/ops/_cutedsl_compat.py | 65 ------------ cula/ops/_mlir_compat.py | 22 ++-- cula/ops/_nvvm_compat.py | 48 +++++++++ cula/ops/kda/decode/mtp_conv.py | 21 +--- cula/ops/kda/sm100/delta_h.py | 23 +--- cula/ops/lightning/prefill_sm100.py | 24 +---- cula/ops/lightning/sm90/prefill_kernel.py | 2 +- cula/ops/sm100/ptx.py | 95 +++++------------ docs/kda_sm100_cutedsl_fwd_status.md | 11 +- docs/lightning_attn_sm90.md | 2 +- pyproject.toml | 5 +- scripts/modal_validate.py | 4 +- tests/test_cutedsl_compat.py | 123 ++++++++++++---------- tests/test_lightning_attn_prefill_sm90.py | 6 +- 15 files changed, 172 insertions(+), 281 deletions(-) delete mode 100644 cula/ops/_cutedsl_compat.py create mode 100644 cula/ops/_nvvm_compat.py diff --git a/README.md b/README.md index 93b4df22..d1f130f7 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,7 @@ pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu129 For CUDA 13, install the CuTeDSL CUDA 13 extra before installing cuLA: ```bash -pip install 'nvidia-cutlass-dsl[cu13]>=4.4.2,<4.7,!=4.5.0' +pip install 'nvidia-cutlass-dsl[cu13]>=4.6.2,<4.9' ``` **Install cuLA & dependencies:** diff --git a/cula/ops/_cutedsl_compat.py b/cula/ops/_cutedsl_compat.py deleted file mode 100644 index 66bbf0df..00000000 --- a/cula/ops/_cutedsl_compat.py +++ /dev/null @@ -1,65 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# 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. - -"""Compatibility helpers for low-level CuTeDSL bindings.""" - -from collections.abc import Callable -from dataclasses import dataclass -from inspect import Parameter, signature -from typing import Literal - - -@dataclass(frozen=True) -class Tcgen05LdStApi: - """Detected keyword interface of the generated tcgen05 load/store ops.""" - - ld_has_num: bool - st_has_num: bool - st_value_keyword: Literal["r", "val"] - - -def _parameter_names(op: Callable, op_name: str) -> set[str]: - try: - parameters = signature(op).parameters - except (TypeError, ValueError) as exc: - raise RuntimeError(f"Unable to inspect the CuTeDSL {op_name} binding") from exc - - if any(parameter.kind is Parameter.VAR_KEYWORD for parameter in parameters.values()): - raise RuntimeError(f"Unsupported CuTeDSL {op_name} signature: variadic keyword arguments are ambiguous") - return set(parameters) - - -def detect_tcgen05_ldst_api(ld_op: Callable, st_op: Callable) -> Tcgen05LdStApi: - """Detect supported tcgen05 load/store keyword variants from their signatures.""" - - ld_parameters = _parameter_names(ld_op, "tcgen05_ld") - st_parameters = _parameter_names(st_op, "tcgen05_st") - - missing_ld = {"res", "shape", "tmem_addr"} - ld_parameters - if missing_ld: - raise RuntimeError(f"Unsupported CuTeDSL tcgen05_ld signature: missing {sorted(missing_ld)}") - - missing_st = {"shape", "tmem_addr"} - st_parameters - if missing_st: - raise RuntimeError(f"Unsupported CuTeDSL tcgen05_st signature: missing {sorted(missing_st)}") - - value_keywords = {"r", "val"} & st_parameters - if len(value_keywords) != 1: - raise RuntimeError("Unsupported CuTeDSL tcgen05_st signature: expected exactly one value keyword from ['r', 'val']") - - return Tcgen05LdStApi( - ld_has_num="num" in ld_parameters, - st_has_num="num" in st_parameters, - st_value_keyword=value_keywords.pop(), - ) diff --git a/cula/ops/_mlir_compat.py b/cula/ops/_mlir_compat.py index 35c5abc8..46dbdff6 100644 --- a/cula/ops/_mlir_compat.py +++ b/cula/ops/_mlir_compat.py @@ -40,9 +40,9 @@ - Zero dependencies: version parsing is done with a small regex so this module stays usable in any environment cuLA can be installed into. -The version contract mirrors ``pyproject.toml`` (including its ``!=4.5.0`` -exclusion) and the canary probes make a broken binding fail fast with an -actionable message instead of surfacing mid-JIT. +The version contract mirrors ``pyproject.toml`` and the canary probes make a +broken binding fail fast with an actionable message instead of surfacing +mid-JIT. """ from __future__ import annotations @@ -55,9 +55,8 @@ # ``pyproject.toml``; when CuTeDSL is bumped, extend ``_SUPPORTED_MIN`` / # ``_SUPPORTED_MAX`` only after the new release has been validated against the # canaries below (and ideally against the SM90/SM100 kernel test suites). -_SUPPORTED_MIN: Final[tuple[int, ...]] = (4, 4, 2) -_SUPPORTED_MAX: Final[tuple[int, ...]] = (4, 7, 0) -_EXCLUDED_VERSIONS: Final[frozenset[tuple[int, ...]]] = frozenset({(4, 5, 0)}) +_SUPPORTED_MIN: Final[tuple[int, ...]] = (4, 6, 2) +_SUPPORTED_MAX: Final[tuple[int, ...]] = (4, 9, 0) # dialect name -> (package, attribute) inside ``cutlass``. _PRIVATE_TABLE: Final[dict[str, tuple[str, str]]] = { @@ -123,7 +122,7 @@ def _installed_version() -> tuple[int, int, int]: import cutlass # noqa: PLC0415 except ImportError: raise RuntimeError( - "cuLA requires the nvidia-cutlass-dsl package; install it with `pip install 'nvidia-cutlass-dsl>=4.4.2,<4.7'`." + "cuLA requires the nvidia-cutlass-dsl package; install it with `pip install 'nvidia-cutlass-dsl>=4.6.2,<4.9'`." ) from None version = getattr(cutlass, "__version__", None) if not isinstance(version, str): @@ -135,11 +134,6 @@ def _installed_version() -> tuple[int, int, int]: def _check_contract(version: tuple[int, ...]) -> None: - if version in _EXCLUDED_VERSIONS: - raise RuntimeError( - f"Installed CuTeDSL version {'.'.join(map(str, version))} is explicitly " - f"excluded by cuLA (see pyproject.toml). {_INCIDENT_NOTE}" - ) if not (_SUPPORTED_MIN <= version < _SUPPORTED_MAX): raise RuntimeError( f"Installed CuTeDSL version {'.'.join(map(str, version))} is outside the " @@ -231,8 +225,8 @@ def vector_extract_element(vec, position, *, loc=None, ip=None): plus a static-position array, so a constant position is ``extract(vec, [], [position], ...)``. - Preferring ``extractelement`` when both exist keeps the pre-4.6 code path - byte-identical to what cuLA shipped before the gateway. + The fallback remains for early 4.6 development wheels; supported releases + use ``extract``. """ vector_dialect = _load("vector") if _has_attribute_path(vector_dialect, ("extractelement",)): diff --git a/cula/ops/_nvvm_compat.py b/cula/ops/_nvvm_compat.py new file mode 100644 index 00000000..9d0c12ec --- /dev/null +++ b/cula/ops/_nvvm_compat.py @@ -0,0 +1,48 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize NVVM binding names across supported CuTeDSL releases.""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable, Mapping +from functools import wraps +from typing import Any + + +def adapt_keywords(function: Callable[..., Any], renames: Mapping[str, str]) -> Callable[..., Any]: + """Translate legacy keyword names only when the binding uses new names.""" + parameters = inspect.signature(function).parameters + active = {old: new for old, new in renames.items() if old not in parameters and new in parameters} + if not active: + return function + + @wraps(function) + def adapted(*args, **kwargs): + for old, new in active.items(): + if old not in kwargs: + continue + if new in kwargs: + raise TypeError(f"Both {old!r} and {new!r} were provided") + kwargs[new] = kwargs.pop(old) + return function(*args, **kwargs) + + return adapted + + +_MATRIX_KEYWORDS = {"mma_kind": "kind", "d": "matrix_d", "a": "matrix_a", "b": "matrix_b"} + + +def bind_mma_ops(nvvm): + """Return CuTeDSL 4.6-style MMA callables for a 4.6 or 4.8 NVVM module.""" + return ( + adapt_keywords( + nvvm.tcgen05_mma, + {**_MATRIX_KEYWORDS, "write_disable_mask": "disable_output_lane"}, + ), + adapt_keywords( + nvvm.tcgen05_mma_ws, + {**_MATRIX_KEYWORDS, "col_b_zero_mask": "zero_col_mask"}, + ), + ) diff --git a/cula/ops/kda/decode/mtp_conv.py b/cula/ops/kda/decode/mtp_conv.py index 96d3d25a..b153e4d9 100644 --- a/cula/ops/kda/decode/mtp_conv.py +++ b/cula/ops/kda/decode/mtp_conv.py @@ -35,9 +35,7 @@ import torch from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import Int32 -from cutlass.cutlass_dsl import T as _T -from cula.ops._mlir_compat import llvm as _llvm from cula.ops.kda.decode.cute import ( TILE_K, _get_cached_stream, @@ -56,30 +54,13 @@ NWARP2 = 2 -@cutlass.dsl_user_op -def _atomic_add_global_i32(ptr_i64, addend_i32, *, loc=None, ip=None): - """Atomically add to a global int32 counter and return its old value.""" - result = _llvm.inline_asm( - _T.i32(), - [ptr_i64, addend_i32], - "atom.global.add.s32 $0, [$1], $2;", - "=r,l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - return Int32(result) - - @cute.jit def _announce_qk_read_complete(qk_arrival_counters, counter_idx, expected_arrivals, lane): """Return a warp-uniform flag identifying the last CTA in a q/k group.""" arrival = Int32(0) if lane == 0: counter_ptr = qk_arrival_counters.iterator + counter_idx - arrival = _atomic_add_global_i32(counter_ptr.toint().ir_value(), Int32(1).ir_value()) + arrival = cute.arch.atomic_add(counter_ptr, Int32(1), sem="relaxed", scope="gpu") arrival = cute.arch.shuffle_sync(arrival, 0, mask=-1, mask_and_clamp=31) return arrival == expected_arrivals - 1 diff --git a/cula/ops/kda/sm100/delta_h.py b/cula/ops/kda/sm100/delta_h.py index 1d31f1ff..be573f9e 100644 --- a/cula/ops/kda/sm100/delta_h.py +++ b/cula/ops/kda/sm100/delta_h.py @@ -30,11 +30,9 @@ from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream from cutlass.cute.typing import Float32, Int32, Int64 -from cutlass.cutlass_dsl import T as _T from fla.ops.utils import prepare_chunk_indices, prepare_lens from fla.utils import tensor_cache -from cula.ops._mlir_compat import llvm as _llvm from cula.ops.kda.sm100.policy import sm100_intracard_cp_decision from cula.utils import USE_FAST_MATH, assert_blackwell, get_device_sm_count @@ -50,23 +48,6 @@ def prepare_chunk_offsets_i32( return F.pad(triton.cdiv(prepare_lens(cu_seqlens), chunk_size), (1, 0), value=0).cumsum(-1).int() -@cutlass.dsl_user_op -def _atomic_add_global_i32(ptr_i64, addend_i32, *, loc=None, ip=None): - """Atomic add on global memory int32. Returns old value.""" - result = _llvm.inline_asm( - _T.i32(), - [ptr_i64, addend_i32], - "atom.global.add.s32 $0, [$1], $2;", - "=r,l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - return Int32(result) - - PRINT_DEBUG = False LN2 = 0.6931471805599453 @@ -873,7 +854,7 @@ def kernel( sWorkIdx = storage.sWorkIdx.get_tensor(cute.make_layout((2,))) if warp_idx == self.load_warp_id: with cute.arch.elect_one(): - first_work_idx = _atomic_add_global_i32(workspace_iter.toint().ir_value(), Int32(1).ir_value()) + first_work_idx = cute.arch.atomic_add(workspace_iter, Int32(1), sem="relaxed", scope="gpu") sWorkIdx[(0,)] = first_work_idx cute.arch.fence_acq_rel_cta() cute.arch.mbarrier_arrive(sched_mbar_base) @@ -1059,7 +1040,7 @@ def kernel( consumed_phase1 = Int32(1) - consumed_phase1 first_wu = Int32(0) with cute.arch.elect_one(): - next_idx = _atomic_add_global_i32(workspace_iter.toint().ir_value(), Int32(1).ir_value()) + next_idx = cute.arch.atomic_add(workspace_iter, Int32(1), sem="relaxed", scope="gpu") sWorkIdx[(sched_buf,)] = next_idx cute.arch.fence_acq_rel_cta() cute.arch.mbarrier_arrive(sched_mbar_base + sched_buf) diff --git a/cula/ops/lightning/prefill_sm100.py b/cula/ops/lightning/prefill_sm100.py index 2796effb..0d71176d 100644 --- a/cula/ops/lightning/prefill_sm100.py +++ b/cula/ops/lightning/prefill_sm100.py @@ -61,29 +61,9 @@ from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream from cutlass.cute.typing import Float32, Int32, Int64 -from cutlass.cutlass_dsl import T as _T -from cula.ops._mlir_compat import llvm as _llvm from cula.utils import USE_FAST_MATH, assert_blackwell - -@cutlass.dsl_user_op -def _atomic_add_global_i32(ptr_i64, addend_i32, *, loc=None, ip=None): - """Atomic add on global memory int32. Returns old value.""" - result = _llvm.inline_asm( - _T.i32(), - [ptr_i64, addend_i32], - "atom.global.add.s32 $0, [$1], $2;", - "=r,l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - return Int32(result) - - PRINT_DEBUG = False @@ -1242,7 +1222,7 @@ def kernel( sWorkIdx = storage.sWorkIdx.get_tensor(cute.make_layout((2,))) if warp_idx == self.load_warp_id: with cute.arch.elect_one(): - first_work_idx = _atomic_add_global_i32(workspace_iter.toint().ir_value(), Int32(1).ir_value()) + first_work_idx = cute.arch.atomic_add(workspace_iter, Int32(1), sem="relaxed", scope="gpu") sWorkIdx[(0,)] = first_work_idx cute.arch.fence_acq_rel_cta() cute.arch.mbarrier_arrive(sched_mbar_base) @@ -1376,7 +1356,7 @@ def kernel( cute.arch.mbarrier_wait(sched_done_mbar, sched_done_phase) sched_done_phase = Int32(1) - sched_done_phase with cute.arch.elect_one(): - next_idx = _atomic_add_global_i32(workspace_iter.toint().ir_value(), Int32(1).ir_value()) + next_idx = cute.arch.atomic_add(workspace_iter, Int32(1), sem="relaxed", scope="gpu") sWorkIdx[(sched_buf,)] = next_idx cute.arch.fence_acq_rel_cta() cute.arch.mbarrier_arrive(sched_mbar_base + sched_buf) diff --git a/cula/ops/lightning/sm90/prefill_kernel.py b/cula/ops/lightning/sm90/prefill_kernel.py index 23672975..4441f494 100644 --- a/cula/ops/lightning/sm90/prefill_kernel.py +++ b/cula/ops/lightning/sm90/prefill_kernel.py @@ -72,7 +72,7 @@ LightningSm90PrefillSchedule, ) -SUPPORTED_CUTLASS_DSL_SPECIFIER = ">=4.4.2,<4.7,!=4.5.0" +SUPPORTED_CUTLASS_DSL_SPECIFIER = ">=4.6.2,<4.9" DECAY_LUT_ENTRIES = 65 TENSORMAP_BYTES = 128 diff --git a/cula/ops/sm100/ptx.py b/cula/ops/sm100/ptx.py index 64360938..fb86d6e7 100644 --- a/cula/ops/sm100/ptx.py +++ b/cula/ops/sm100/ptx.py @@ -60,18 +60,15 @@ from cutlass.cute.typing import Int32 from cutlass.cutlass_dsl import dsl_user_op -from cula.ops._cutedsl_compat import detect_tcgen05_ldst_api from cula.ops._mlir_compat import arith as _arith from cula.ops._mlir_compat import ir, llvm from cula.ops._mlir_compat import nvvm as _nvvm +from cula.ops._nvvm_compat import bind_mma_ops CollectorBBuffer = _nvvm.Tcgen05MMACollectorBBuffer CollectorOp = _nvvm.Tcgen05MMACollectorOp -_tcgen05_ldst_api = detect_tcgen05_ldst_api(_nvvm.tcgen05_ld, _nvvm.tcgen05_st) -_TCGEN05_LD_HAS_NUM = _tcgen05_ldst_api.ld_has_num -_TCGEN05_ST_HAS_NUM = _tcgen05_ldst_api.st_has_num -_TCGEN05_ST_USES_R = _tcgen05_ldst_api.st_value_keyword == "r" +_tcgen05_mma, _tcgen05_mma_ws = bind_mma_ops(_nvvm) def _to_ir(val, loc=None, ip=None): @@ -174,24 +171,13 @@ def _do(addr_val, *, loc=None, ip=None): ptr6_ty = llvm.PointerType.get(address_space=6) tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) vec_i32_ty = ir.VectorType.get([num], i32_ty) - if cutlass.const_expr(_TCGEN05_LD_HAS_NUM): - result = _nvvm.tcgen05_ld( - res=vec_i32_ty, - shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, - num=num, - tmem_addr=tmem_ptr, - loc=loc, - ip=ip, - ) - else: - result = _nvvm.tcgen05_ld( - res=vec_i32_ty, - shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, - tmem_addr=tmem_ptr, - loc=loc, - ip=ip, - ) - return result + return _nvvm.tcgen05_ld( + res=vec_i32_ty, + shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, + tmem_addr=tmem_ptr, + loc=loc, + ip=ip, + ) return _do(Int32(taddr)) @@ -204,42 +190,13 @@ def tcgen05_st_32x32b(num: int, taddr: int, vec): def _do(addr_val, vec_val, *, loc=None, ip=None): ptr6_ty = llvm.PointerType.get(address_space=6) tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip) - if cutlass.const_expr(_TCGEN05_ST_HAS_NUM): - if cutlass.const_expr(_TCGEN05_ST_USES_R): - _nvvm.tcgen05_st( - shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, - num=num, - tmem_addr=tmem_ptr, - r=_to_ir(vec_val, loc, ip), - loc=loc, - ip=ip, - ) - else: - _nvvm.tcgen05_st( - shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, - num=num, - tmem_addr=tmem_ptr, - val=_to_ir(vec_val, loc, ip), - loc=loc, - ip=ip, - ) - else: - if cutlass.const_expr(_TCGEN05_ST_USES_R): - _nvvm.tcgen05_st( - shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, - tmem_addr=tmem_ptr, - r=_to_ir(vec_val, loc, ip), - loc=loc, - ip=ip, - ) - else: - _nvvm.tcgen05_st( - shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, - tmem_addr=tmem_ptr, - val=_to_ir(vec_val, loc, ip), - loc=loc, - ip=ip, - ) + _nvvm.tcgen05_st( + shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B, + tmem_addr=tmem_ptr, + val=_to_ir(vec_val, loc, ip), + loc=loc, + ip=ip, + ) _do(Int32(taddr), vec) @@ -256,7 +213,7 @@ def _do(addr_val, desc_val, *, loc=None, ip=None): shape=_nvvm.Tcgen05CpShape.SHAPE_128x256b, taddr=tmem_ptr, smem_desc=_to_ir(desc_val, loc, ip), - cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + group=_nvvm.CTAGroupKind.CTA_1, loc=loc, ip=ip, ) @@ -276,7 +233,7 @@ def _do(addr_val, desc_val, *, loc=None, ip=None): shape=_nvvm.Tcgen05CpShape.SHAPE_128x128b, taddr=tmem_ptr, smem_desc=_to_ir(desc_val, loc, ip), - cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + group=_nvvm.CTAGroupKind.CTA_1, loc=loc, ip=ip, ) @@ -382,9 +339,9 @@ def _do(c_val, da_val, db_val, dv_val, sc_val, m0_val, m1_val, m2_val, m3_val, * v = llvm.InsertElementOp(v, m2_ir, idx2, loc=loc, ip=ip) mask = llvm.InsertElementOp(v, m3_ir, idx3, loc=loc, ip=ip) - _nvvm.tcgen05_mma( + _tcgen05_mma( mma_kind=_nvvm.Tcgen05MMAKind.TF32, - cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + cta_group=_nvvm.CTAGroupKind.CTA_1, d=d_ptr, a=da_ir, b=db_ir, @@ -453,9 +410,9 @@ def _do(c_val, a_val, db_val, dv_val, sc_val, m0_val, m1_val, m2_val, m3_val, *, v = llvm.InsertElementOp(v, m2_ir, idx2, loc=loc, ip=ip) mask = llvm.InsertElementOp(v, m3_ir, idx3, loc=loc, ip=ip) - _nvvm.tcgen05_mma( + _tcgen05_mma( mma_kind=_nvvm.Tcgen05MMAKind.TF32, - cta_group=_nvvm.Tcgen05GroupKind.CTA_1, + cta_group=_nvvm.CTAGroupKind.CTA_1, d=d_ptr, a=a_ptr, b=b_ir, @@ -509,7 +466,7 @@ def _do(c_val, da_val, db_val, dv_val, sc_val, *, loc=None, ip=None): sc_ir = _to_ir(sc_val, loc, ip) enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - _nvvm.tcgen05_mma_ws( + _tcgen05_mma_ws( mma_kind=_nvvm.Tcgen05MMAKind.TF32, d=d_ptr, a=da_ir, @@ -556,7 +513,7 @@ def _do(c_val, da_val, db_val, dv_val, sc_val, *, loc=None, ip=None): sc_ir = _to_ir(sc_val, loc, ip) enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - _nvvm.tcgen05_mma_ws( + _tcgen05_mma_ws( mma_kind=_nvvm.Tcgen05MMAKind.F16, d=d_ptr, a=da_ir, @@ -604,7 +561,7 @@ def _do(c_val, a_val, db_val, dv_val, sc_val, *, loc=None, ip=None): sc_ir = _to_ir(sc_val, loc, ip) enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - _nvvm.tcgen05_mma_ws( + _tcgen05_mma_ws( mma_kind=_nvvm.Tcgen05MMAKind.TF32, d=d_ptr, a=a_ptr, @@ -652,7 +609,7 @@ def _do(c_val, a_val, db_val, dv_val, sc_val, *, loc=None, ip=None): sc_ir = _to_ir(sc_val, loc, ip) enable_d = _arith.trunci(i1_ty, sc_ir, loc=loc, ip=ip) - _nvvm.tcgen05_mma_ws( + _tcgen05_mma_ws( mma_kind=_nvvm.Tcgen05MMAKind.F16, d=d_ptr, a=a_ptr, diff --git a/docs/kda_sm100_cutedsl_fwd_status.md b/docs/kda_sm100_cutedsl_fwd_status.md index d5a69be6..8e2fdc22 100644 --- a/docs/kda_sm100_cutedsl_fwd_status.md +++ b/docs/kda_sm100_cutedsl_fwd_status.md @@ -22,12 +22,11 @@ partial, and packed-varlen inputs use the same exact CuTeDSL kernels through the modular host path. Batch one deliberately avoids a composed-host compiler failure observed in CUTLASS DSL 4.6.2 and 4.6.3; it does not fall back to C++. -The package currently declares `nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0` -and this migration was validated with 4.6.2. An isolated 4.6.3 check reproduced -the batch-one composed-host compiler failure, so a patch-version upgrade does -not remove the modular fallback. Version 4.8.1 is outside the supported range -and did not provide a clean drop-in validation in the current environment; -the production path therefore does not depend on an unverified upgrade. +The package declares `nvidia-cutlass-dsl>=4.6.2,<4.9` and a CUDA 13 extra. +The low-level TMEM and NVVM wrappers use the 4.6 API directly and normalize the +MMA keyword renames introduced by 4.8. The batch-one modular fallback remains: +it is a stable dispatch choice and does not depend on the composed-host +compiler behavior of a particular patch release. `unified_gref=True` and `use_tf32_inverse=False` are rejected explicitly because neither option belongs to the bitwise-compatible production path. diff --git a/docs/lightning_attn_sm90.md b/docs/lightning_attn_sm90.md index 30227101..00b476dc 100644 --- a/docs/lightning_attn_sm90.md +++ b/docs/lightning_attn_sm90.md @@ -171,7 +171,7 @@ units before pipeline storage and TensorMap views are reused. ## Supported specialization - Target: `sm_90a` on compute capability 9.0. -- CuTe DSL: `nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0`. +- CuTe DSL: `nvidia-cutlass-dsl>=4.6.2,<4.9`. - Q, K, V, and O: BF16. - Decay and recurrent state: FP32. - Key and value dimensions: 128. diff --git a/pyproject.toml b/pyproject.toml index 60cbd2b1..0bdf8bd2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,8 +11,8 @@ readme = "README.md" authors = [ { name = "cula contributors" } ] requires-python = ">=3.10" dependencies = [ - "nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0", - "apache-tvm-ffi>=0.1.9", + "nvidia-cutlass-dsl>=4.6.2,<4.9", + "apache-tvm-ffi>=0.1.11", "packaging>=23.2", ] license = "Apache-2.0" @@ -25,6 +25,7 @@ dev = [ "build", "twine", ] +cu13 = ["nvidia-cutlass-dsl[cu13]>=4.6.2,<4.9"] [tool.setuptools.packages.find] where = ["."] diff --git a/scripts/modal_validate.py b/scripts/modal_validate.py index 29ac097d..a2006546 100644 --- a/scripts/modal_validate.py +++ b/scripts/modal_validate.py @@ -5,7 +5,7 @@ summary. The GPU and the CuTeDSL version are chosen via environment variables so a reviewer can reproduce a specific environment, e.g.:: - CULA_VALIDATE_GPU=B200 CULA_VALIDATE_CUTLASS='nvidia-cutlass-dsl==4.5.2' \\ + CULA_VALIDATE_GPU=B200 CULA_VALIDATE_CUTLASS='nvidia-cutlass-dsl==4.8.1' \\ modal run scripts/modal_validate.py SM100-only tests (``test_ptx_umma_ws.py``) self-deselect on non-Blackwell @@ -26,7 +26,7 @@ CUDA_TAG = "cu129" TORCH_VERSION = "2.9.1" GPU = os.environ.get("CULA_VALIDATE_GPU", "H100") -CUTLASS_SPEC = os.environ.get("CULA_VALIDATE_CUTLASS", "nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0") +CUTLASS_SPEC = os.environ.get("CULA_VALIDATE_CUTLASS", "nvidia-cutlass-dsl>=4.6.2,<4.9") TESTS = [ "tests/test_cutedsl_compat.py", "tests/test_lightning_attn_prefill_sm90.py", diff --git a/tests/test_cutedsl_compat.py b/tests/test_cutedsl_compat.py index 5944e502..1c475d29 100644 --- a/tests/test_cutedsl_compat.py +++ b/tests/test_cutedsl_compat.py @@ -18,56 +18,72 @@ import pytest from cula.ops import _mlir_compat -from cula.ops._cutedsl_compat import Tcgen05LdStApi, detect_tcgen05_ldst_api from cula.ops._mlir_compat import _parse_version +from cula.ops._nvvm_compat import adapt_keywords, bind_mma_ops -def _legacy_ld(res, shape, num, tmem_addr, *, pack=None, half_split_offset=None): - pass +def test_adapt_keywords_preserves_46_binding(): + def mma(mma_kind, d, a, b, *, write_disable_mask=None): + return mma_kind, d, a, b, write_disable_mask + adapted = adapt_keywords(mma, {"mma_kind": "kind", "d": "matrix_d", "write_disable_mask": "disable_output_lane"}) + assert adapted is mma + assert adapted(mma_kind="tf32", d=1, a=2, b=3, write_disable_mask=4) == ("tf32", 1, 2, 3, 4) -def _legacy_st(shape, num, tmem_addr, r, *, unpack=None, half_split_offset=None): - pass +def test_adapt_keywords_translates_48_binding(): + def mma(kind, matrix_d, matrix_a, matrix_b, *, disable_output_lane=None): + return kind, matrix_d, matrix_a, matrix_b, disable_output_lane -def _inferred_ld(res, shape, tmem_addr, *, pack=None, offset=None): - pass - - -def _inferred_st(shape, tmem_addr, val, *, unpack=None, offset=None): - pass - - -def test_detect_tcgen05_ldst_legacy_api(): - assert detect_tcgen05_ldst_api(_legacy_ld, _legacy_st) == Tcgen05LdStApi( - ld_has_num=True, - st_has_num=True, - st_value_keyword="r", - ) - - -def test_detect_tcgen05_ldst_inferred_api(): - assert detect_tcgen05_ldst_api(_inferred_ld, _inferred_st) == Tcgen05LdStApi( - ld_has_num=False, - st_has_num=False, - st_value_keyword="val", + adapted = adapt_keywords( + mma, + { + "mma_kind": "kind", + "d": "matrix_d", + "a": "matrix_a", + "b": "matrix_b", + "write_disable_mask": "disable_output_lane", + }, ) - - -def test_detect_tcgen05_ldst_mixed_api(): - assert detect_tcgen05_ldst_api(_legacy_ld, _inferred_st) == Tcgen05LdStApi( - ld_has_num=True, - st_has_num=False, - st_value_keyword="val", - ) - - -def test_detect_tcgen05_ldst_rejects_unknown_store_value_keyword(): - def unsupported_st(shape, tmem_addr, value): - pass - - with pytest.raises(RuntimeError, match="expected exactly one value keyword"): - detect_tcgen05_ldst_api(_inferred_ld, unsupported_st) + assert adapted(mma_kind="tf32", d=1, a=2, b=3, write_disable_mask=4) == ("tf32", 1, 2, 3, 4) + + +def test_adapt_keywords_rejects_duplicate_old_and_new_names(): + def mma(kind): + return kind + + adapted = adapt_keywords(mma, {"mma_kind": "kind"}) + with pytest.raises(TypeError, match="Both 'mma_kind' and 'kind'"): + adapted(mma_kind="tf32", kind="f16") + + +def test_bind_mma_ops_adapts_regular_and_weight_stationary_calls(): + def mma(kind, cta_group, matrix_d, matrix_a, matrix_b, idesc, enable_input_d, *, disable_output_lane=None): + return kind, matrix_d, matrix_a, matrix_b, disable_output_lane + + def mma_ws(kind, matrix_d, matrix_a, matrix_b, idesc, enable_input_d, *, zero_col_mask=None): + return kind, matrix_d, matrix_a, matrix_b, zero_col_mask + + regular, weight_stationary = bind_mma_ops(types.SimpleNamespace(tcgen05_mma=mma, tcgen05_mma_ws=mma_ws)) + assert regular( + mma_kind="tf32", + cta_group="one", + d=1, + a=2, + b=3, + idesc=4, + enable_input_d=True, + write_disable_mask=5, + ) == ("tf32", 1, 2, 3, 5) + assert weight_stationary( + mma_kind="f16", + d=1, + a=2, + b=3, + idesc=4, + enable_input_d=True, + col_b_zero_mask=5, + ) == ("f16", 1, 2, 3, 5) # --------------------------------------------------------------------------- @@ -186,8 +202,8 @@ def test_version_contract_rejects_out_of_range(): def test_version_contract_rejects_two_component_upper_bound(): - # "4.7" must normalize to (4, 7, 0) and not slip under the < 4.7.0 cap. - restore = _install_fake_cutlass("4.7") + # "4.9" must normalize to (4, 9, 0) and not slip under the < 4.9 cap. + restore = _install_fake_cutlass("4.9") try: with pytest.raises(RuntimeError, match="outside the range validated by cuLA"): _touch_gateway("llvm") @@ -195,18 +211,17 @@ def test_version_contract_rejects_two_component_upper_bound(): restore() -def test_version_contract_rejects_two_component_excluded(): - # "4.5" must normalize to (4, 5, 0) and hit the explicit 4.5.0 exclusion. - restore = _install_fake_cutlass("4.5") +def test_version_contract_rejects_below_minimum(): + restore = _install_fake_cutlass("4.6.1") try: - with pytest.raises(RuntimeError, match="explicitly excluded by cuLA"): + with pytest.raises(RuntimeError, match="outside the range validated by cuLA"): _touch_gateway("llvm") finally: restore() def test_missing_cutlass_reports_install_hint(): - restore = _install_fake_cutlass("4.5.3", with_mlir=False) + restore = _install_fake_cutlass("4.6.2", with_mlir=False) sys.modules["cutlass"] = None # simulate the package being absent try: with pytest.raises(RuntimeError, match="nvidia-cutlass-dsl"): @@ -216,7 +231,7 @@ def test_missing_cutlass_reports_install_hint(): def test_missing_dialect_raises_with_dialect_name(): - restore = _install_fake_cutlass("4.5.3", missing_dialect="llvm") + restore = _install_fake_cutlass("4.8.1", missing_dialect="llvm") try: with pytest.raises(RuntimeError, match="no longer exposes 'llvm'"): _touch_gateway("llvm") @@ -225,7 +240,7 @@ def test_missing_dialect_raises_with_dialect_name(): def test_missing_canary_raises_naming_the_entry_point(): - restore = _install_fake_cutlass("4.5.3", broken_canary="llvm") + restore = _install_fake_cutlass("4.8.1", broken_canary="llvm") try: with pytest.raises(RuntimeError, match="missing the canary entry point inline_asm"): _touch_gateway("llvm") @@ -234,7 +249,7 @@ def test_missing_canary_raises_naming_the_entry_point(): def test_vector_dispatch_uses_extractelement_when_extract_absent(): - restore = _install_fake_cutlass("4.5.3") + restore = _install_fake_cutlass("4.6.2") try: seen = {} vector_dialect = _mlir_compat._load("vector") @@ -254,7 +269,7 @@ def fake_extractelement(vec, *, position, loc=None, ip=None): def test_vector_dispatch_uses_extract_when_available(): - restore = _install_fake_cutlass("4.6.0") + restore = _install_fake_cutlass("4.8.1") try: vector_dialect = _mlir_compat._load("vector") # simulate the 4.6 rename: drop extractelement, add extract @@ -281,7 +296,7 @@ def test_unknown_attribute_raises_attribute_error(): def test_dialect_bindings_are_cached(): - restore = _install_fake_cutlass("4.5.3") + restore = _install_fake_cutlass("4.8.1") try: assert _mlir_compat.llvm is _mlir_compat.llvm finally: diff --git a/tests/test_lightning_attn_prefill_sm90.py b/tests/test_lightning_attn_prefill_sm90.py index f585fb29..d49b1512 100644 --- a/tests/test_lightning_attn_prefill_sm90.py +++ b/tests/test_lightning_attn_prefill_sm90.py @@ -82,10 +82,10 @@ def test_kernel_exposes_only_production_entrypoints() -> None: assert [_qualified_name(base) for base in backend.bases] == ["LightningSm90PrefillSchedule"] specifier = values["SUPPORTED_CUTLASS_DSL_SPECIFIER"] - assert specifier == ">=4.4.2,<4.7,!=4.5.0" + assert specifier == ">=4.6.2,<4.9" accepted = SpecifierSet(specifier) - assert all(version in accepted for version in ("4.4.2", "4.5.1", "4.5.2", "4.5.3", "4.6.0", "4.6.1")) - assert all(version not in accepted for version in ("4.4.1", "4.5.0", "4.7.0")) + assert all(version in accepted for version in ("4.6.2", "4.6.3", "4.7.0", "4.8.0", "4.8.1")) + assert all(version not in accepted for version in ("4.6.1", "4.9.0")) assert values["DECAY_LUT_ENTRIES"] == 65 assert SCHEDULE_PATH.exists() assert not any(isinstance(node, ast.FunctionDef) and node.name.startswith("launch_") for node in tree.body) From 50c727dfd882fce23953695ba67e4bfabe8a5d5a Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 15:25:46 +0800 Subject: [PATCH 03/10] fix(kda): load FlashKDA workspaces through TMA Preserve the raw K_INTER byte image while attaching all K2 workspace loads to the input transaction barrier. This avoids the untracked async bulk-copy path that can leave the recurrence kernel waiting indefinitely. Signed-off-by: shuyan.ycf --- cula/ops/kda/sm90/k2.py | 70 ++++++++++++++++++++++++++++------------- 1 file changed, 48 insertions(+), 22 deletions(-) diff --git a/cula/ops/kda/sm90/k2.py b/cula/ops/kda/sm90/k2.py index 1cb60cc5..bebe3bb6 100644 --- a/cula/ops/kda/sm90/k2.py +++ b/cula/ops/kda/sm90/k2.py @@ -44,7 +44,7 @@ def _make_state_smem_layout(): return cute.tile_to_shape(atom, (D, D), (0, 1)) -from cula.ops.kda.sm90._common import _stream_key, copy_async_bulk, movm_t_b16 # noqa: E402 +from cula.ops.kda.sm90._common import _stream_key, movm_t_b16 # noqa: E402 def _make_out_kinter_one_stage(): @@ -63,6 +63,12 @@ def _make_out_kinter_one_stage(): def k2_kernel( tma_atom_v: cute.CopyAtom, tma_tensor_v: cute.Tensor, + tma_atom_kd: cute.CopyAtom, + tma_tensor_kd: cute.Tensor, + tma_atom_qd: cute.CopyAtom, + tma_tensor_qd: cute.Tensor, + tma_atom_kr: cute.CopyAtom, + tma_tensor_kr: cute.Tensor, tma_atom_inv: cute.CopyAtom, tma_tensor_inv: cute.Tensor, tma_atom_mqk: cute.CopyAtom, @@ -153,6 +159,27 @@ def k2_kernel( cute.group_modes(sV, 0, 2), cute.group_modes(gSrc_v, 0, 2), ) + # K1 stores the physical K_INTER byte image directly. Load through a plain + # TMA view so the transaction participates in the input mbarrier without + # reinterpreting or permuting those bytes. + raw_qk_stage_layout = cute.make_layout((CHUNK, D, STAGES), stride=(D, 1, CHUNK * D)) + sKD_raw = cute.make_tensor(sKd.iterator, raw_qk_stage_layout) + sQD_raw = cute.make_tensor(sQd.iterator, raw_qk_stage_layout) + sKR_raw = cute.make_tensor(sKr.iterator, raw_qk_stage_layout) + + def partition_raw_workspace(atom, tensor, smem_tensor): + g_tile = cute.local_tile(tensor, (CHUNK, D), (None, None, None)) + return cpasync.tma_partition( + atom, + 0, + cute.make_layout(1), + cute.group_modes(smem_tensor, 0, 2), + cute.group_modes(g_tile, 0, 2), + ) + + tKDs, tKDg = partition_raw_workspace(tma_atom_kd, tma_tensor_kd, sKD_raw) + tQDs, tQDg = partition_raw_workspace(tma_atom_qd, tma_tensor_qd, sQD_raw) + tKRs, tKRg = partition_raw_workspace(tma_atom_kr, tma_tensor_kr, sKR_raw) gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None)) tIs, tIg = cpasync.tma_partition( tma_atom_inv, @@ -193,22 +220,6 @@ def k2_kernel( cute.group_modes(sBeta, 0, 2), cute.group_modes(gSrc_beta, 0, 2), ) - raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), cutlass.BFloat16) - raw_stage_layout = cute.make_layout( - (CHUNK * D, STAGES), - stride=(1, CHUNK * D), - ) - raw_gmem_layout = cute.make_layout( - (CHUNK * D, total_tiles * H), - stride=(1, CHUNK * D), - ) - sKD_raw = cute.make_tensor(sKd.iterator, raw_stage_layout) - sQD_raw = cute.make_tensor(sQd.iterator, raw_stage_layout) - sKR_raw = cute.make_tensor(sKr.iterator, raw_stage_layout) - gKD_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout) - gQD_raw = cute.make_tensor(ws_qd.iterator, raw_gmem_layout) - gKR_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout) - # Load initial_state -> sState[K_in, D_out] if has_initial_state: state_base = cutlass.Int32(seq_idx) * cutlass.Int32(H * D * D) + cutlass.Int32(head_idx) * cutlass.Int32(D * D) @@ -348,11 +359,9 @@ def k2_kernel( cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) else: cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - # Restore the byte-identical K_INTER images produced by K1. The - # raw-workspace transport idea is credited there to Flash-Flash-KDA. - copy_async_bulk(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) - copy_async_bulk(raw_copy_atom, gQD_raw[(None, wt_l)], sQD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) - copy_async_bulk(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_qd, tQDg[(None, 0, 0, wt_l)], tQDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_mqk, tMg[(None, 0, 0, wt_l)], tMs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) @@ -615,6 +624,7 @@ def run_k2( ): cc_smem = cute.make_layout((CHUNK, CHUNK), stride=(CHUNK, 1)) kinter_smem = _make_out_kinter_one_stage() + raw_qk_smem = cute.make_layout((CHUNK, D), stride=(D, 1)) def make_thd_atom(t, op, t_total: cutlass.Int32): view = cute.make_tensor( @@ -630,8 +640,18 @@ def make_ws_cc_atom(t): ) return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, cc_smem, (CHUNK, CHUNK)) + def make_raw_ws_qk_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), + ) + return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, raw_qk_smem, (CHUNK, D)) + tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp(), V_T_total) tma_atom_out, tma_tensor_out = make_thd_atom(out, cpasync.CopyBulkTensorTileS2GOp(), O_T_total) + tma_atom_kd, tma_tensor_kd = make_raw_ws_qk_atom(ws_kd) + tma_atom_qd, tma_tensor_qd = make_raw_ws_qk_atom(ws_qd) + tma_atom_kr, tma_tensor_kr = make_raw_ws_qk_atom(ws_kr) tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv) tma_atom_mqk, tma_tensor_mqk = make_ws_cc_atom(ws_mqk) @@ -681,6 +701,12 @@ def make_beta_atom(t): k2_kernel( tma_atom_v, tma_tensor_v, + tma_atom_kd, + tma_tensor_kd, + tma_atom_qd, + tma_tensor_qd, + tma_atom_kr, + tma_tensor_kr, tma_atom_inv, tma_tensor_inv, tma_atom_mqk, From 7e7c2306d14f0f81d07c61cc1a1065c26d8f13b7 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 15:31:06 +0800 Subject: [PATCH 04/10] fix(kda): synchronize CP workspace loads Route intracard-CP pre-scan workspace loads through plain-layout TMA so they preserve the K_INTER bytes and complete the pipeline transaction barrier. Signed-off-by: shuyan.ycf --- cula/ops/kda/sm90/cp/pre_scan.py | 57 ++++++++++++++++++++++---------- 1 file changed, 40 insertions(+), 17 deletions(-) diff --git a/cula/ops/kda/sm90/cp/pre_scan.py b/cula/ops/kda/sm90/cp/pre_scan.py index 4cc31b9a..a831b25a 100644 --- a/cula/ops/kda/sm90/cp/pre_scan.py +++ b/cula/ops/kda/sm90/cp/pre_scan.py @@ -30,7 +30,7 @@ from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream -from cula.ops.kda.sm90._common import copy_async_bulk, movm_t_b16 +from cula.ops.kda.sm90._common import movm_t_b16 from cula.ops.kda.sm90.k2 import ( CHUNK, D, @@ -48,6 +48,10 @@ def pre_scan_kernel( tma_atom_v: cute.CopyAtom, tma_tensor_v: cute.Tensor, + tma_atom_kd: cute.CopyAtom, + tma_tensor_kd: cute.Tensor, + tma_atom_kr: cute.CopyAtom, + tma_tensor_kr: cute.Tensor, tma_atom_inv: cute.CopyAtom, tma_tensor_inv: cute.Tensor, tma_atom_gt: cute.CopyAtom, @@ -119,6 +123,25 @@ def pre_scan_kernel( cute.group_modes(sV, 0, 2), cute.group_modes(gSrc_v, 0, 2), ) + # K1 stores the physical K_INTER byte image directly. Plain-layout TMA + # restores those bytes while participating in the input transaction + # barrier; raw cp.async.bulk copies can otherwise leave the wait pending. + raw_qk_stage_layout = cute.make_layout((CHUNK, D, STAGES), stride=(D, 1, CHUNK * D)) + sKD_raw = cute.make_tensor(sKd.iterator, raw_qk_stage_layout) + sKR_raw = cute.make_tensor(sKr.iterator, raw_qk_stage_layout) + + def partition_raw_workspace(atom, tensor, smem_tensor): + g_tile = cute.local_tile(tensor, (CHUNK, D), (None, None, None)) + return cpasync.tma_partition( + atom, + 0, + cute.make_layout(1), + cute.group_modes(smem_tensor, 0, 2), + cute.group_modes(g_tile, 0, 2), + ) + + tKDs, tKDg = partition_raw_workspace(tma_atom_kd, tma_tensor_kd, sKD_raw) + tKRs, tKRg = partition_raw_workspace(tma_atom_kr, tma_tensor_kr, sKR_raw) gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None)) tIs, tIg = cpasync.tma_partition( tma_atom_inv, @@ -143,20 +166,6 @@ def pre_scan_kernel( cute.group_modes(sBeta, 0, 2), cute.group_modes(gSrc_beta, 0, 2), ) - raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), cutlass.BFloat16) - raw_stage_layout = cute.make_layout( - (CHUNK * D, STAGES), - stride=(1, CHUNK * D), - ) - raw_gmem_layout = cute.make_layout( - (CHUNK * D, total_tiles * H), - stride=(1, CHUNK * D), - ) - sKD_raw = cute.make_tensor(sKd.iterator, raw_stage_layout) - sKR_raw = cute.make_tensor(sKr.iterator, raw_stage_layout) - gKD_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout) - gKR_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout) - # sState=0. The distributed FP32 M accumulator initializes sM below. if tidx < D: for e in cutlass.range_constexpr(D): @@ -285,8 +294,8 @@ def pre_scan_kernel( cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) else: cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - copy_async_bulk(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) - copy_async_bulk(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_beta, tBg[(None, 0, 0, wt_l)], tBs[(None, s_dyn_l)], tma_bar_ptr=bar_l) @@ -486,6 +495,7 @@ def run_pre_scan( ): cc_smem = cute.make_layout((CHUNK, CHUNK), stride=(CHUNK, 1)) kinter_smem = _make_out_kinter_one_stage() + raw_qk_smem = cute.make_layout((CHUNK, D), stride=(D, 1)) def make_thd_atom(t, op): view = cute.make_tensor( @@ -501,7 +511,16 @@ def make_ws_cc_atom(t): ) return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, cc_smem, (CHUNK, CHUNK)) + def make_raw_ws_qk_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), + ) + return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, raw_qk_smem, (CHUNK, D)) + tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp()) + tma_atom_kd, tma_tensor_kd = make_raw_ws_qk_atom(ws_kd) + tma_atom_kr, tma_tensor_kr = make_raw_ws_qk_atom(ws_kr) tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv) gt_smem = cute.make_layout((D, 1), stride=(1, D)) @@ -549,6 +568,10 @@ def make_beta_atom(t): pre_scan_kernel( tma_atom_v, tma_tensor_v, + tma_atom_kd, + tma_tensor_kd, + tma_atom_kr, + tma_tensor_kr, tma_atom_inv, tma_tensor_inv, tma_atom_gt, From b26a5154c411bcd1cb6c5bb1c024e6ed64a3a548 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 15:33:43 +0800 Subject: [PATCH 05/10] test(kda): lock FlashKDA workspace transport Remove the obsolete pre-4.6 bulk-copy branch, document the K2 hang signature, and add static contracts requiring barrier-accounted plain-layout TMA workspace loads. Signed-off-by: shuyan.ycf --- cula/ops/kda/sm90/_common.py | 42 ---------------------------- cula/ops/kda/sm90/k1.py | 12 ++++---- docs/kda_sm100_cutedsl_fwd_status.md | 6 ++++ tests/test_kda_flashkda_pipeline.py | 38 +++++++++++++++++++++++++ 4 files changed, 50 insertions(+), 48 deletions(-) create mode 100644 tests/test_kda_flashkda_pipeline.py diff --git a/cula/ops/kda/sm90/_common.py b/cula/ops/kda/sm90/_common.py index 7c6114a5..15db3b7b 100644 --- a/cula/ops/kda/sm90/_common.py +++ b/cula/ops/kda/sm90/_common.py @@ -3,12 +3,7 @@ """Shared low-level helpers for the SM90 FlashKDA kernels.""" -import re -from importlib.metadata import PackageNotFoundError -from importlib.metadata import version as package_version - import cutlass -import cutlass.cute as cute import torch from cutlass import Int32 from cutlass.cutlass_dsl import T as _T @@ -16,43 +11,6 @@ from cula.ops._mlir_compat import llvm as _llvm -def _parse_cutedsl_version(raw_version: str) -> tuple[int, int, int]: - """Return the numeric CuTeDSL version from a release or dev version string.""" - match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", raw_version) - if match is None: - raise RuntimeError(f"Unable to parse the installed CuTeDSL version: {raw_version!r}") - return tuple(int(component or 0) for component in match.groups()) - - -def _installed_cutedsl_version() -> tuple[int, int, int]: - """Read the CuTeDSL version at runtime without relying on a private symbol.""" - raw_version = getattr(cutlass, "__version__", None) - if raw_version is None: - try: - raw_version = package_version("nvidia-cutlass-dsl") - except PackageNotFoundError as exc: - raise RuntimeError("nvidia-cutlass-dsl is required by the SM90 FlashKDA backend") from exc - return _parse_cutedsl_version(raw_version) - - -# CuTeDSL 4.6.0 added the missing elect_one inside cute.copy for async bulk -# atoms. Older releases need an explicit elect_one, while nesting one around -# cute.copy is incorrect in 4.6+. Keep this decision as a compile-time -# constant after detecting the installed runtime version, so the same source -# supports both API behaviours. -_CUTEDSL_VERSION = _installed_cutedsl_version() -_CUTE_COPY_AUTO_ELECTS_BULK = _CUTEDSL_VERSION >= (4, 6, 0) - - -def copy_async_bulk(atom, src, dst, **kwargs) -> None: - """Issue a CuTeDSL async bulk copy across supported elect_one APIs.""" - if cutlass.const_expr(_CUTE_COPY_AUTO_ELECTS_BULK): - cute.copy(atom, src, dst, **kwargs) - else: - with cute.arch.elect_one(): - cute.copy(atom, src, dst, **kwargs) - - def _stream_key(device: torch.device) -> tuple[str, int]: return str(device), int(torch.cuda.current_stream(device).cuda_stream) diff --git a/cula/ops/kda/sm90/k1.py b/cula/ops/kda/sm90/k1.py index 34aa1e43..7edd60a9 100644 --- a/cula/ops/kda/sm90/k1.py +++ b/cula/ops/kda/sm90/k1.py @@ -30,7 +30,7 @@ from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream -from cula.ops.kda.sm90._common import _stream_key, add_f16x2_u32, copy_async_bulk, movm_t_b16 +from cula.ops.kda.sm90._common import _stream_key, add_f16x2_u32, movm_t_b16 CHUNK: int = 16 D: int = 128 @@ -522,11 +522,11 @@ def k1_kernel( # layout-aware tensor TMA stores. This raw-workspace transport idea comes # from Flash-Flash-KDA: https://github.com/Itssshikhar/Flash-Flash-KDA if warp_idx == 0: - # copy_async_bulk supplies elect_one only for pre-4.6 CuTeDSL. CuTeDSL - # 4.6+ emits it inside cute.copy, where an outer elect_one is invalid. - copy_async_bulk(raw_copy_atom, sQDws_raw, gQDws_raw[(None, ws_slot)]) - copy_async_bulk(raw_copy_atom, sKDws_raw, gKDws_raw[(None, ws_slot)]) - copy_async_bulk(raw_copy_atom, sKRws_raw, gKRws_raw[(None, ws_slot)]) + # CuTeDSL 4.6.2+ elects the issuing lane inside cute.copy for bulk + # atoms; wrapping these calls in another elect_one is invalid. + cute.copy(raw_copy_atom, sQDws_raw, gQDws_raw[(None, ws_slot)]) + cute.copy(raw_copy_atom, sKDws_raw, gKDws_raw[(None, ws_slot)]) + cute.copy(raw_copy_atom, sKRws_raw, gKRws_raw[(None, ws_slot)]) with cute.arch.elect_one(): cute.copy(tma_atom_ws_inv, tINVws_s[(None,)], tINVws_g[(None, 0, 0, ws_slot)]) cute.copy(tma_atom_ws_mqk, tMQKws_s[(None,)], tMQKws_g[(None, 0, 0, ws_slot)]) diff --git a/docs/kda_sm100_cutedsl_fwd_status.md b/docs/kda_sm100_cutedsl_fwd_status.md index 8e2fdc22..4326c178 100644 --- a/docs/kda_sm100_cutedsl_fwd_status.md +++ b/docs/kda_sm100_cutedsl_fwd_status.md @@ -28,6 +28,12 @@ MMA keyword renames introduced by 4.8. The batch-one modular fallback remains: it is a stable dispatch choice and does not depend on the composed-host compiler behavior of a particular patch release. +FlashKDA K2 and intracard-CP pre-scan restore K1's physical K_INTER workspace +images through plain-layout TMA loads. This keeps the byte layout unchanged and +accounts every load in the transaction barrier. Raw G2S bulk copies are not +mixed into that barrier: on SM100 they could leave K2 waiting indefinitely, +which appears on the host as a busy `cudaDeviceSynchronize()` call. + `unified_gref=True` and `use_tf32_inverse=False` are rejected explicitly because neither option belongs to the bitwise-compatible production path. diff --git a/tests/test_kda_flashkda_pipeline.py b/tests/test_kda_flashkda_pipeline.py new file mode 100644 index 00000000..5f2db23a --- /dev/null +++ b/tests/test_kda_flashkda_pipeline.py @@ -0,0 +1,38 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Static contracts for FlashKDA's barrier-accounted workspace transport.""" + +import ast +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +@pytest.mark.parametrize( + ("relative_path", "run_function", "workspace_names"), + [ + ("cula/ops/kda/sm90/k2.py", "run_k2", ("ws_kd", "ws_qd", "ws_kr")), + ("cula/ops/kda/sm90/cp/pre_scan.py", "run_pre_scan", ("ws_kd", "ws_kr")), + ], +) +def test_workspace_loads_use_barrier_accounted_plain_tma(relative_path, run_function, workspace_names): + source = (REPO_ROOT / relative_path).read_text(encoding="utf-8") + tree = ast.parse(source) + run = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == run_function) + run_text = ast.unparse(run) + + assert "CopyBulkG2SOp" not in source + assert "raw_qk_smem = cute.make_layout((CHUNK, D), stride=(D, 1))" in run_text + for name in workspace_names: + assert f"make_raw_ws_qk_atom({name})" in run_text + + kernel_text = ast.unparse(next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name.endswith("kernel"))) + assert "raw_qk_stage_layout = cute.make_layout((CHUNK, D, STAGES), stride=(D, 1, CHUNK * D))" in kernel_text + for atom in ("tma_atom_kd", "tma_atom_kr"): + assert f"cute.copy({atom}" in kernel_text + if "ws_qd" in workspace_names: + assert "cute.copy(tma_atom_qd" in kernel_text + assert kernel_text.count("tma_bar_ptr=bar_l") >= len(workspace_names) From 8d7bfa50cb7a48078d1d4e9730fc285553ef0017 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 15:41:01 +0800 Subject: [PATCH 06/10] style(kda): satisfy current Ruff formatting Signed-off-by: shuyan.ycf --- tests/test_kda_flashkda_pipeline.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_kda_flashkda_pipeline.py b/tests/test_kda_flashkda_pipeline.py index 5f2db23a..5f706e1f 100644 --- a/tests/test_kda_flashkda_pipeline.py +++ b/tests/test_kda_flashkda_pipeline.py @@ -29,7 +29,9 @@ def test_workspace_loads_use_barrier_accounted_plain_tma(relative_path, run_func for name in workspace_names: assert f"make_raw_ws_qk_atom({name})" in run_text - kernel_text = ast.unparse(next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name.endswith("kernel"))) + kernel_text = ast.unparse( + next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name.endswith("kernel")) + ) assert "raw_qk_stage_layout = cute.make_layout((CHUNK, D, STAGES), stride=(D, 1, CHUNK * D))" in kernel_text for atom in ("tma_atom_kd", "tma_atom_kr"): assert f"cute.copy({atom}" in kernel_text From 19ddea21da2e37275f9b0fb7d0043a50a366242f Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 18:54:57 +0800 Subject: [PATCH 07/10] fix(benchmarks): allow FlashKDA runs on SM100 Signed-off-by: shuyan.ycf --- benchmarks/bench_kda_sm90_cp.py | 11 ++++++----- benchmarks/bench_kda_sm90_prefill.py | 28 +++++++++++++++++----------- tests/test_kda_flashkda_arch.py | 10 ++++++++++ 3 files changed, 33 insertions(+), 16 deletions(-) diff --git a/benchmarks/bench_kda_sm90_cp.py b/benchmarks/bench_kda_sm90_cp.py index 33465831..83afa6a8 100644 --- a/benchmarks/bench_kda_sm90_cp.py +++ b/benchmarks/bench_kda_sm90_cp.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 """ -bench_kda_sm90_cp.py — Benchmark: SM90 FlashKDA CP-auto vs CP-off +bench_kda_sm90_cp.py — Benchmark: SM90/SM100 FlashKDA CP-auto vs CP-off Compares flashkda_prefill(use_intracard_cp="auto") against flashkda_prefill(use_intracard_cp=False) on varlen configs. Reports whether @@ -33,8 +33,9 @@ set_seed, ) from cula.kda import flashkda_prefill as cula_kda_prefill +from cula.kda._flashkda_arch import assert_flashkda_supported from cula.ops.kda.sm90.cp.plan import CHUNK, plan_auto -from cula.utils import assert_hopper, get_device_sm_count +from cula.utils import get_device_sm_count D = 128 H_VALUES = [4, 8] @@ -99,11 +100,11 @@ def _auto_engages(seq_lens, H, sm_count) -> bool: def bench_cp(h_values, configs): print("\n" + "=" * 100) - print(" SM90 FlashKDA Benchmark: CP-auto vs CP-off") + print(" SM90/SM100 FlashKDA Benchmark: CP-auto vs CP-off") print("=" * 100) device = torch.device("cuda") - assert_hopper(device) + assert_flashkda_supported(device) num_sms = get_device_sm_count(device) print(f" [Device] {torch.cuda.get_device_name(0)} SMs={num_sms}") results = [] @@ -160,7 +161,7 @@ def bench_cp(h_values, configs): def print_report(results, h_values): sep = "=" * 105 print(f"\n\n{sep}") - print(" BENCHMARK REPORT: SM90 FlashKDA") + print(" BENCHMARK REPORT: SM90/SM100 FlashKDA") print(' CP-auto (use_intracard_cp="auto") vs CP-off (False)') print(f" D={D} dtype=bf16 safe_gate=True") wu, ni = resolve_benchmark_repeats(WARMUP, N_ITERS, ncu_mode=NCU_MODE, sanitizer_mode=SANITIZER_MODE) diff --git a/benchmarks/bench_kda_sm90_prefill.py b/benchmarks/bench_kda_sm90_prefill.py index f43013e6..0591531d 100644 --- a/benchmarks/bench_kda_sm90_prefill.py +++ b/benchmarks/bench_kda_sm90_prefill.py @@ -14,10 +14,11 @@ # limitations under the License. """ -bench_kda_sm90_prefill.py — Benchmark: SM90 CuTeDSL KDA prefill vs FLA Triton baseline +bench_kda_sm90_prefill.py — Benchmark: SM90/SM100 CuTeDSL KDA prefill vs FLA Triton baseline -Hopper-only. Calls cula.kda.flashkda (cula_kda_prefill, the K1+K2 -two-kernel prefill) directly. +Calls cula.kda.flashkda (cula_kda_prefill, the K1+K2 two-kernel prefill) +directly on supported SM90 and SM100 devices. The filename is retained because +the implementation originated as the SM90 FlashKDA pipeline. Compares: - Accuracy: relative_rms_error, rel_max, mean_diff of the output o vs FLA Triton @@ -58,9 +59,9 @@ set_seed, ) from cula.kda import flashkda_prefill as cula_kda_prefill -from cula.utils import assert_hopper +from cula.kda._flashkda_arch import assert_flashkda_supported -_SM_TAG = "sm90" +_SM_TAG = "sm90/sm100" H, D = 64, 128 WARMUP = 25 @@ -159,7 +160,7 @@ def _make_common(B, T, cu_seqlens, device): def bench_fixed(configs): print("\n" + "=" * 100) - print(f" Fixed-Length Benchmark: SM90 cula_kda_prefill ({_SM_TAG}) vs FLA Triton") + print(f" Fixed-Length Benchmark: cula_kda_prefill ({_SM_TAG}) vs FLA Triton") print("=" * 100) results = [] for B, T in configs: @@ -187,7 +188,7 @@ def bench_fixed(configs): def bench_varlen(configs): print("\n" + "=" * 100) - print(f" Varlen Benchmark: SM90 cula_kda_prefill ({_SM_TAG}) vs FLA Triton") + print(f" Varlen Benchmark: cula_kda_prefill ({_SM_TAG}) vs FLA Triton") print("=" * 100) results = [] for seq_lens, total_len, dist in configs: @@ -219,7 +220,7 @@ def bench_varlen(configs): def print_report(fixed_results, varlen_results): sep = "=" * 110 print(f"\n\n{sep}") - print(" BENCHMARK REPORT: SM90 cula_kda_prefill (K1+K2 two-kernel)") + print(" BENCHMARK REPORT: cula_kda_prefill (K1+K2 two-kernel)") print(f" cuLA {_SM_TAG} vs FLA Triton") print(f" H={H} D={D} dtype=bf16 safe_gate=True has_init_state={HAS_INIT_STATE}") wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP @@ -264,7 +265,7 @@ def print_report(fixed_results, varlen_results): def main(): - parser = argparse.ArgumentParser(description="bench_kda_sm90_prefill: SM90 cula_kda_prefill vs FLA Triton") + parser = argparse.ArgumentParser(description="bench_kda_sm90_prefill: SM90/SM100 cula_kda_prefill vs FLA Triton") parser.add_argument( "--mode", type=str, @@ -279,7 +280,8 @@ def main(): parser.add_argument("--heads", type=int, default=H, help=f"Number of heads (H == HV, MHA). Default: {H}") args = parser.parse_args() - assert_hopper(torch.device("cuda")) + device = torch.device("cuda") + assert_flashkda_supported(device) global NCU_MODE, SANITIZER_MODE, HAS_INIT_STATE H = args.heads @@ -293,7 +295,11 @@ def main(): HAS_INIT_STATE = True print("[init_state] using non-zero initial state") - print(f"[Device] {torch.cuda.get_device_name(0)} compute capability {_SM_TAG} → cula_kda_prefill (SM90 K1+K2)") + major, minor = torch.cuda.get_device_capability(device) + print( + f"[Device] {torch.cuda.get_device_name(device)} compute capability sm{major}{minor} " + "→ cula_kda_prefill (SM90-derived K1+K2)" + ) fixed_configs = [ (1, 512), diff --git a/tests/test_kda_flashkda_arch.py b/tests/test_kda_flashkda_arch.py index b37b8df1..e7121cbd 100644 --- a/tests/test_kda_flashkda_arch.py +++ b/tests/test_kda_flashkda_arch.py @@ -1,6 +1,8 @@ # Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 +from pathlib import Path + import pytest import torch @@ -33,3 +35,11 @@ def test_flashkda_rejects_cpu(): assert not is_flashkda_supported(device) with pytest.raises(RuntimeError, match="requires a CUDA device"): assert_flashkda_supported(device) + + +@pytest.mark.parametrize("benchmark", ["bench_kda_sm90_prefill.py", "bench_kda_sm90_cp.py"]) +def test_flashkda_benchmark_uses_shared_architecture_guard(benchmark): + source = (Path(__file__).parents[1] / "benchmarks" / benchmark).read_text() + + assert "assert_flashkda_supported" in source + assert "assert_hopper" not in source From 167f4f5603b633a3bd7f685cc0851f5024768c8c Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 19:20:40 +0800 Subject: [PATCH 08/10] fix(install): align dependencies with CUDA versions Signed-off-by: shuyan.ycf --- README.md | 16 ++++++++++++---- scripts/modal_validate.py | 2 +- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index d1f130f7..0e611af0 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ Linear attention mechanisms reformulate standard attention to use linear-time st cuLA supports both **Hopper (SM90)** and **Blackwell (SM10X)** GPUs. -> **Requirements (Hopper & Blackwell):** Python 3.12+, CUDA Toolkit 12.9+, PyTorch 2.9.1+, and a matching `nvidia-cutlass-dsl` installation. cuLA does not build a C++/CUDA extension. +> **Requirements (Hopper & Blackwell):** Python 3.12+, CUDA Toolkit 12.9+, PyTorch 2.8.0+, and a matching `nvidia-cutlass-dsl` installation. cuLA does not build a C++/CUDA extension. > **Note:** The PyTorch CUDA version must match your system CUDA Toolkit version. Check with `nvcc --version` and `python -c "import torch; print(torch.version.cuda)"`. @@ -32,6 +32,10 @@ Platform-independent wheels containing the CuTeDSL sources are available on [Git Replace `` with the release tag (e.g., `v0.2.0`) and `` with the version (e.g., `0.2.0`). Or download the `.whl` file directly from the [Releases page](https://github.com/inclusionAI/cuLA/releases) and install it with `pip install .whl`. +The cuLA wheel does not install FLA. Install `flash-linear-attention` before importing `chunk_kda`: + + pip install flash-linear-attention + ### Build from Source **Clone cuLA & dependencies:** @@ -41,15 +45,19 @@ git clone https://github.com/inclusionAI/cuLA.git git submodule update --init --recursive ``` -**Install PyTorch:** +**Install PyTorch and CuTeDSL:** + +For CUDA 12.9: ```bash -pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu129 +pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu129 +pip install 'nvidia-cutlass-dsl>=4.6.2,<4.9' ``` -For CUDA 13, install the CuTeDSL CUDA 13 extra before installing cuLA: +For CUDA 13: ```bash +pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu130 pip install 'nvidia-cutlass-dsl[cu13]>=4.6.2,<4.9' ``` diff --git a/scripts/modal_validate.py b/scripts/modal_validate.py index a2006546..55b3579b 100644 --- a/scripts/modal_validate.py +++ b/scripts/modal_validate.py @@ -24,7 +24,7 @@ FORK = "https://github.com/bikrammajhi/cuLA.git" BRANCH = "mlir-compat-gateway" CUDA_TAG = "cu129" -TORCH_VERSION = "2.9.1" +TORCH_VERSION = "2.8.0" GPU = os.environ.get("CULA_VALIDATE_GPU", "H100") CUTLASS_SPEC = os.environ.get("CULA_VALIDATE_CUTLASS", "nvidia-cutlass-dsl>=4.6.2,<4.9") TESTS = [ From fb9544d8372bbc69ccda37c64e9e73bcdd86aa49 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 19:20:40 +0800 Subject: [PATCH 09/10] test(kda): strengthen FlashKDA barrier contracts Signed-off-by: shuyan.ycf --- tests/test_kda_flashkda_pipeline.py | 49 ++++++++++++++++++++++------- 1 file changed, 37 insertions(+), 12 deletions(-) diff --git a/tests/test_kda_flashkda_pipeline.py b/tests/test_kda_flashkda_pipeline.py index 5f706e1f..2ee6b5e0 100644 --- a/tests/test_kda_flashkda_pipeline.py +++ b/tests/test_kda_flashkda_pipeline.py @@ -12,13 +12,23 @@ @pytest.mark.parametrize( - ("relative_path", "run_function", "workspace_names"), + ("relative_path", "run_function", "workspace_names", "expected_tma_bytes"), [ - ("cula/ops/kda/sm90/k2.py", "run_k2", ("ws_kd", "ws_qd", "ws_kr")), - ("cula/ops/kda/sm90/cp/pre_scan.py", "run_pre_scan", ("ws_kd", "ws_kr")), + ( + "cula/ops/kda/sm90/k2.py", + "run_k2", + ("ws_kd", "ws_qd", "ws_kr"), + "4 * CHUNK * D * 2 + 2 * CHUNK * CHUNK * 2 + D * 4 + CHUNK * 2", + ), + ( + "cula/ops/kda/sm90/cp/pre_scan.py", + "run_pre_scan", + ("ws_kd", "ws_kr"), + "3 * CHUNK * D * 2 + CHUNK * CHUNK * 2 + D * 4 + CHUNK * 2", + ), ], ) -def test_workspace_loads_use_barrier_accounted_plain_tma(relative_path, run_function, workspace_names): +def test_workspace_loads_use_barrier_accounted_plain_tma(relative_path, run_function, workspace_names, expected_tma_bytes): source = (REPO_ROOT / relative_path).read_text(encoding="utf-8") tree = ast.parse(source) run = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == run_function) @@ -29,12 +39,27 @@ def test_workspace_loads_use_barrier_accounted_plain_tma(relative_path, run_func for name in workspace_names: assert f"make_raw_ws_qk_atom({name})" in run_text - kernel_text = ast.unparse( - next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name.endswith("kernel")) - ) + kernel = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name.endswith("kernel")) + kernel_text = ast.unparse(kernel) assert "raw_qk_stage_layout = cute.make_layout((CHUNK, D, STAGES), stride=(D, 1, CHUNK * D))" in kernel_text - for atom in ("tma_atom_kd", "tma_atom_kr"): - assert f"cute.copy({atom}" in kernel_text - if "ws_qd" in workspace_names: - assert "cute.copy(tma_atom_qd" in kernel_text - assert kernel_text.count("tma_bar_ptr=bar_l") >= len(workspace_names) + + tma_bytes = next( + node + for node in ast.walk(kernel) + if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "TMA_BYTES" + ) + expected_bytes = ast.parse(expected_tma_bytes, mode="eval").body + assert ast.dump(tma_bytes.value, include_attributes=False) == ast.dump(expected_bytes, include_attributes=False) + + copy_calls = [ + node + for node in ast.walk(kernel) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and ast.unparse(node.func) == "cute.copy" + ] + for workspace_name in workspace_names: + atom = f"tma_atom_{workspace_name.removeprefix('ws_')}" + workspace_copies = [call for call in copy_calls if call.args and ast.unparse(call.args[0]) == atom] + assert len(workspace_copies) == 1 + assert any( + keyword.arg == "tma_bar_ptr" and ast.unparse(keyword.value) == "bar_l" for keyword in workspace_copies[0].keywords + ) From 20f776687227451856cd83d0f2ea4bc89e921142 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 24 Oct 2026 19:27:35 +0800 Subject: [PATCH 10/10] fix(packaging): require supported Python version Signed-off-by: shuyan.ycf --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 0bdf8bd2..9ecfeb42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ dynamic = ["version"] description = "High-performance CuTeDSL linear-attention kernels" readme = "README.md" authors = [ { name = "cula contributors" } ] -requires-python = ">=3.10" +requires-python = ">=3.12" dependencies = [ "nvidia-cutlass-dsl>=4.6.2,<4.9", "apache-tvm-ffi>=0.1.11",