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..0e611af0 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,21 @@ 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.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)"`. ### 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`. + +The cuLA wheel does not install FLA. Install `flash-linear-attention` before importing `chunk_kda`: + + pip install flash-linear-attention ### Build from Source @@ -41,10 +45,20 @@ 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.8.0 --index-url https://download.pytorch.org/whl/cu129 +pip install 'nvidia-cutlass-dsl>=4.6.2,<4.9' +``` + +For CUDA 13: ```bash -pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu129 +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' ``` **Install cuLA & dependencies:** @@ -57,10 +71,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/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/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/_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/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/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/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/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, 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/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, 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 d93af578..4326c178 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 @@ -22,12 +22,17 @@ 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. + +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. @@ -116,5 +121,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/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 520c3e80..9ecfeb42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,21 +1,21 @@ [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" +requires-python = ">=3.12" 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 = { text = "Apache-2.0" } +license = "Apache-2.0" [project.optional-dependencies] dev = [ @@ -25,6 +25,7 @@ dev = [ "build", "twine", ] +cu13 = ["nvidia-cutlass-dsl[cu13]>=4.6.2,<4.9"] [tool.setuptools.packages.find] where = ["."] @@ -33,10 +34,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/scripts/modal_validate.py b/scripts/modal_validate.py index 29ac097d..55b3579b 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 @@ -24,9 +24,9 @@ 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.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/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_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_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_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 diff --git a/tests/test_kda_flashkda_pipeline.py b/tests/test_kda_flashkda_pipeline.py new file mode 100644 index 00000000..2ee6b5e0 --- /dev/null +++ b/tests/test_kda_flashkda_pipeline.py @@ -0,0 +1,65 @@ +# 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", "expected_tma_bytes"), + [ + ( + "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, 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) + 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 = 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 + + 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 + ) 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 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)