From 76b411a4e6e7ba3d9ee755f430405915a6e59ea7 Mon Sep 17 00:00:00 2001 From: samyyc Date: Sun, 1 Nov 2026 19:51:08 +0800 Subject: [PATCH 1/3] feat(cuda): add optimized SM70 Volta compatibility kernels --- .github/workflows/build.yml | 19 +- .github/workflows/release.yml | 3 + CMakeLists.txt | 9 +- CMakePresets.json | 13 ++ README.md | 5 +- cmake/modules/backend.cmake | 3 + cmake/modules/packaging.cmake | 3 + cmake/modules/sm70.cmake | 43 ++++ cmake/modules/sm75.cmake | 58 ++++-- docs/integration.md | 7 +- docs/sm70.md | 105 ++++++++++ src/backend/probe.cpp | 12 +- src/backend/probe.h | 9 +- src/kernels/attention/flash_attention.cu | 8 + src/kernels/basic/gemm.cpp | 2 +- src/kernels/gdn/chunk_rule.cu | 7 +- src/kernels/quant/int8_gemm.cu | 10 +- src/kernels/quant/int8_gemm_residual.cu | 8 + src/kernels/quant/int8_gemm_situ.cu | 8 + src/kernels/sm70/attention_sm70.cu | 241 +++++++++++++++++++++++ src/kernels/sm70/int8_gemm_sm70.cu | 237 ++++++++++++++++++++++ src/kernels/sm70/sm70.cpp | 38 ++++ src/kernels/sm70/sm70.h | 26 +++ src/pipeline/pipeline.cpp | 14 +- src/pipeline/pipeline_impl.h | 6 +- src/pipeline/session.cpp | 4 +- tests/unit/test_sm70.cpp | 168 ++++++++++++++++ tools/sm70_bench.cu | 100 ++++++++++ tools/sm75_check.cu | 26 ++- 29 files changed, 1138 insertions(+), 54 deletions(-) create mode 100644 cmake/modules/sm70.cmake create mode 100644 docs/sm70.md create mode 100644 src/kernels/sm70/attention_sm70.cu create mode 100644 src/kernels/sm70/int8_gemm_sm70.cu create mode 100644 src/kernels/sm70/sm70.cpp create mode 100644 src/kernels/sm70/sm70.h create mode 100644 tests/unit/test_sm70.cpp create mode 100644 tools/sm70_bench.cu diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index a6cf4d0..c30ae1a 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -1,5 +1,5 @@ -# Reusable build of the four release flavours (Linux/Windows x CUDA 12/13) plus Windows CUDA 13 with SM75, called by -# ci.yml and release.yml. Per flavour: configure + build with the CMakePresets.json release preset, CPU-only tests +# Reusable build of the four release flavours (Linux/Windows x CUDA 12/13), Windows CUDA 13 with SM75 and Linux +# CUDA 12 with SM70, called by ci.yml and release.yml. Per flavour: configure + build with the CMakePresets.json release preset, CPU-only tests # (hosted runners have no GPU), install, AC1 (no torch/python) + exports + AC17 (size) + cuBLASLt-through-the-shim # checks, the find_package consumer, CPack archives, and the system-CUDA variant (the install without the `cuda` # component: no cuBLASLt). @@ -33,9 +33,9 @@ jobs: # bitexact_cu13 tests are GPU tests and do not run on hosted runners. cuda: ["12.8.1", "13.3.1"] # An include entry whose os/cuda match a combination is merged into it (that is how the entries below set - # platform/flavour/nvcc); a new combination needs a value no combination has. The Turing flavour is that: - # sm75=true does not match any of the four standard jobs, so it adds a fifth one. + # platform/flavour/nvcc); separate SM75/SM70 values add compatibility builds to the four standard jobs. sm75: [false] + sm70: [false] include: - os: ubuntu-22.04 platform: linux @@ -85,6 +85,17 @@ jobs: flavour: cuda13-sm75 nvcc: "13.3.73" sub-packages: '["nvcc", "crt", "nvvm", "cudart", "cublas", "cublas_dev", "curand_dev", "thrust", "documentation"]' + - os: ubuntu-22.04 + platform: linux + cuda: "12.8.1" + sm70: true + flavour: cuda12-sm70 + nvcc: "12.8.93" + apt: >- + cuda-nvcc-12-8=12.8.93-1 cuda-nvvm-12-8=12.8.93-1 cuda-crt-12-8=12.8.93-1 + cuda-cudart-12-8=12.8.90-1 cuda-cudart-dev-12-8=12.8.90-1 cuda-cccl-12-8=12.8.90-1 + cuda-documentation-12-8=12.8.90-1 libcurand-12-8=10.3.9.90-1 libcurand-dev-12-8=10.3.9.90-1 + libcublas-12-8=12.8.4.1-1 libcublas-dev-12-8=12.8.4.1-1 env: preset: ${{ matrix.platform }}-${{ matrix.flavour }} SCCACHE_GHA_ENABLED: "true" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2236e2b..33b66d6 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -121,6 +121,9 @@ jobs: sm75="mossinfer-$ARCHIVE_VERSION-windows-cu13-x86_64-sm75.zip" test -f "$sm75" || { echo "missing SM75 package $sm75"; exit 1; } test -f "${sm75%.zip}-system-cuda.zip" || { echo "missing SM75 system-CUDA package"; exit 1; } + sm70="mossinfer-$ARCHIVE_VERSION-linux-cu12-x86_64-sm70.zip" + test -f "$sm70" || { echo "missing SM70 package $sm70"; exit 1; } + test -f "${sm70%.zip}-system-cuda.zip" || { echo "missing SM70 system-CUDA package"; exit 1; } - name: SHA256SUMS shell: bash working-directory: dist diff --git a/CMakeLists.txt b/CMakeLists.txt index 9cb46b8..569778c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,6 +31,10 @@ option(MOSSINFER_TESTS "Build unit tests" ON) # Turing (compute capability 7.5): builds src/kernels/sm75 for sm_75 and accepts those devices (see # cmake/modules/sm75.cmake). Off: the engine is Ampere-or-newer only, exactly as before. option(MOSSINFER_SM75 "Build the sm_75 (Turing) kernels and accept compute capability 7.5" OFF) +option(MOSSINFER_SM70 "Build the sm_70 (Volta) kernels and accept compute capability 7.0 (CUDA 12 only)" OFF) +if(MOSSINFER_SM70 AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) + message(FATAL_ERROR "MOSSINFER_SM70 needs CUDA 12.x: CUDA 13 removed Volta compilation and library support") +endif() find_package(CUDAToolkit 12.4 REQUIRED) # sm_120 (Blackwell GeForce) needs CUDA 12.8. @@ -181,7 +185,10 @@ if(MOSSINFER_TESTS) list(FILTER MOSSINFER_TEST_SOURCES EXCLUDE REGEX ".*/tests/unit/test_(c_api|c_api_check|c_api_inspect|component_api|component_gpu|progress_model)\\.cpp$") endif() add_executable(mossinfer_tests ${MOSSINFER_TEST_SOURCES}) - if(MOSSINFER_SM75) + if(MOSSINFER_SM70) + target_compile_definitions(mossinfer_tests PRIVATE MOSSINFER_HAVE_SM70=1) + endif() + if(MOSSINFER_SM75 OR MOSSINFER_SM70) # probe.h's minimum compute capability (test_backend_probe) and the sm_75 kernel-family smoke test. target_compile_definitions(mossinfer_tests PRIVATE MOSSINFER_HAVE_SM75=1) endif() diff --git a/CMakePresets.json b/CMakePresets.json index d899faa..38898f7 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -81,6 +81,12 @@ "inherits": [ "release-base", "linux" ], "cacheVariables": { "MOSSINFER_EXPECT_CUDA_MAJOR": "12" } }, + { + "name": "linux-cuda12-sm70", + "displayName": "Linux release, CUDA 12.8+ with Volta (sm_70)", + "inherits": "linux-cuda12", + "cacheVariables": { "MOSSINFER_SM70": "ON" } + }, { "name": "windows-msvc-cuda13", "displayName": "Windows release, MSVC + CUDA 13.x (driver >= 580)", @@ -106,6 +112,7 @@ { "name": "dev-sccache", "configurePreset": "dev-sccache" }, { "name": "linux-cuda13", "configurePreset": "linux-cuda13" }, { "name": "linux-cuda12", "configurePreset": "linux-cuda12" }, + { "name": "linux-cuda12-sm70", "configurePreset": "linux-cuda12-sm70" }, { "name": "windows-msvc-cuda13", "configurePreset": "windows-msvc-cuda13" }, { "name": "windows-msvc-cuda13-sm75", "configurePreset": "windows-msvc-cuda13-sm75" }, { "name": "windows-msvc-cuda12", "configurePreset": "windows-msvc-cuda12" } @@ -127,6 +134,11 @@ "description": "bitexact_cu13 tests compare against CUDA 13.4 libdevice and are excluded.", "filter": { "exclude": { "label": "bitexact_cu13" } } }, + { + "name": "linux-cuda12-sm70", + "inherits": "linux-cuda12", + "configurePreset": "linux-cuda12-sm70" + }, { "name": "windows-msvc-cuda13", "inherits": "test-base", "configurePreset": "windows-msvc-cuda13" }, { "name": "windows-msvc-cuda13-sm75", @@ -144,6 +156,7 @@ "packagePresets": [ { "name": "linux-cuda13", "configurePreset": "linux-cuda13" }, { "name": "linux-cuda12", "configurePreset": "linux-cuda12" }, + { "name": "linux-cuda12-sm70", "configurePreset": "linux-cuda12-sm70" }, { "name": "windows-msvc-cuda13", "configurePreset": "windows-msvc-cuda13" }, { "name": "windows-msvc-cuda13-sm75", "configurePreset": "windows-msvc-cuda13-sm75" }, { "name": "windows-msvc-cuda12", "configurePreset": "windows-msvc-cuda12" } diff --git a/README.md b/README.md index 9d0389a..8674a37 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ A native C++/CUDA inference engine for the MOSS-Image text-to-image model (5B Di Models: MOSS-Image Base and SFT only (`OpenMOSS-Team/MOSS-Image-Base`, `OpenMOSS-Team/MOSS-Image-SFT`); other models are refused ([supported models](docs/integration.md#supported-models)). -Version 2.1.1. Supported: Linux x86-64 and Windows x64 (MSVC), each with a CUDA 12 and a CUDA 13 build, as a full package and a `-system-cuda` package. Turing (compute capability 7.5, e.g. RTX 2080) is supported by builds configured with `-DMOSSINFER_SM75=ON` (the `windows-msvc-cuda13-sm75` preset packages it): fp16 tensor-core GEMMs replace the bf16 ones, INT8 and INT4 ConvRot files run the INT8 GEMMs, the exact attention and the VAE run sm_75 kernels, all three quality tiers work, and an unset attention resolves to `exact`. See [docs/sm75.md](docs/sm75.md) for what runs, the measured speed on an RTX 2080 Ti and the remaining gaps. The full support matrix is in [docs/integration.md](docs/integration.md#support-status). +Version 2.1.1. Supported: Linux x86-64 and Windows x64 (MSVC), each with a CUDA 12 and a CUDA 13 build, as a full package and a `-system-cuda` package. Turing (compute capability 7.5, e.g. RTX 2080) is supported by builds configured with `-DMOSSINFER_SM75=ON` (the `windows-msvc-cuda13-sm75` preset packages it): fp16 tensor-core GEMMs replace the bf16 ones, INT8 and INT4 ConvRot files run the INT8 GEMMs, the exact attention and the VAE run sm_75 kernels, all three quality tiers work, and an unset attention resolves to `exact`. See [docs/sm75.md](docs/sm75.md) for what runs, the measured speed on an RTX 2080 Ti and the remaining gaps. Volta (Tesla V100, sm_70) is supported by the CUDA 12 `linux-cuda12-sm70` preset; see [docs/sm70.md](docs/sm70.md) for its kernels, validation and performance. The full support matrix is in [docs/integration.md](docs/integration.md#support-status). The caller picks the quality tier: `reference` (bf16, exact), `high` (the loaded DiT's GEMMs with 8-bit attention) or `turbo` (a 4-bit DiT). An INT8 DiT with `high` is the everyday choice: 2.5 times the bf16 speed on an RTX 5090, in 40 % less VRAM ([performance](docs/integration.md#performance)). Upgrading from 1.x: [migration](docs/integration.md#migrating-from-1x). @@ -28,6 +28,7 @@ Lowest VRAM against normal use at 1024×1024. `min_memory` streams the DiT and c ![Step time against peak GPU memory at 1024×1024 in low-VRAM and normal modes for owoinfer and official PyTorch](assets/comparison_1024x1024_low_vram.png) - [docs/integration.md](docs/integration.md): building the engine into an application. Start here. +- [docs/sm70.md](docs/sm70.md): Volta (Tesla V100) support: CUDA 12 build, exact INT8 accumulation, validation and performance. - [docs/sm75.md](docs/sm75.md): Turing (compute capability 7.5) support: build and package, what runs, measured speed on an RTX 2080 Ti, limitations. - [docs/c-api.md](docs/c-api.md): the C library, `include/mossinfer/mossinfer.h`. - [docs/server-protocol.md](docs/server-protocol.md): `moss-server`, the JSON-lines sidecar process. @@ -39,6 +40,6 @@ cmake --preset linux-cuda13 && cmake --build --preset linux-cuda13 CUDA_VISIBLE_DEVICES="" ctest --preset linux-cuda13 -L '^cpu$' ``` -The other presets are `linux-cuda12`, `windows-msvc-cuda13`, `windows-msvc-cuda13-sm75` and `windows-msvc-cuda12` ([docs/integration.md](docs/integration.md#getting-the-binaries)). The `windows-msvc-cuda13-sm75` preset builds a full CUDA 13 package with Turing (sm_75) support and names it with the `-sm75` suffix; without `MOSSINFER_SM75=ON`, the engine is Ampere-or-newer as before. +The other presets are `linux-cuda12`, `linux-cuda12-sm70`, `windows-msvc-cuda13`, `windows-msvc-cuda13-sm75` and `windows-msvc-cuda12` ([docs/integration.md](docs/integration.md#getting-the-binaries)). The `windows-msvc-cuda13-sm75` preset builds a full CUDA 13 package with Turing (sm_75) support and names it with the `-sm75` suffix; the `linux-cuda12-sm70` preset adds Volta with a `-sm70` suffix. The standard presets remain Ampere-or-newer. Licensed under Apache-2.0 (`LICENSE`, `NOTICE`, `THIRD_PARTY_NOTICES.md`). diff --git a/cmake/modules/backend.cmake b/cmake/modules/backend.cmake index b41c85e..f40421f 100644 --- a/cmake/modules/backend.cmake +++ b/cmake/modules/backend.cmake @@ -13,6 +13,9 @@ if(_mi_backend_probe IN_LIST MOSSINFER_SOURCES) if(MOSSINFER_SM75) string(APPEND _mi_arch_string ",75-real") endif() + if(MOSSINFER_SM70) + string(APPEND _mi_arch_string ",70-real") + endif() set_property(SOURCE ${_mi_backend_probe} APPEND PROPERTY COMPILE_DEFINITIONS "MOSSINFER_CUDA_ARCHITECTURES=\"${_mi_arch_string}\"") unset(_mi_arch_string) diff --git a/cmake/modules/packaging.cmake b/cmake/modules/packaging.cmake index 9b491f5..986d3fd 100644 --- a/cmake/modules/packaging.cmake +++ b/cmake/modules/packaging.cmake @@ -292,6 +292,9 @@ set(_mi_sm75_suffix "") if(MOSSINFER_SM75) set(_mi_sm75_suffix -sm75) endif() +if(MOSSINFER_SM70) + string(APPEND _mi_sm75_suffix -sm70) +endif() set(CPACK_PACKAGE_FILE_NAME mossinfer-${PROJECT_VERSION}-${_mi_os}-cu${CUDAToolkit_VERSION_MAJOR}-${_mi_arch}${_mi_sm75_suffix}) unset(_mi_sm75_suffix) set(CPACK_ARCHIVE_THREADS 0) diff --git a/cmake/modules/sm70.cmake b/cmake/modules/sm70.cmake new file mode 100644 index 0000000..d2245b3 --- /dev/null +++ b/cmake/modules/sm70.cmake @@ -0,0 +1,43 @@ +# Volta tensor-core kernels live outside the sm_80+ core architecture list, like sm75 and nvfp4. CUDA 12 only. +set(_mi_sm70_sources ${MOSSINFER_SOURCES}) +list(FILTER _mi_sm70_sources INCLUDE REGEX ".*/src/kernels/sm70/[^/]*\\.cu$") +if(_mi_sm70_sources) + get_target_property(_mi_core_sources mossinfer_core SOURCES) + list(REMOVE_ITEM _mi_core_sources ${_mi_sm70_sources}) + set_property(TARGET mossinfer_core PROPERTY SOURCES ${_mi_core_sources}) + unset(_mi_core_sources) + if(MOSSINFER_SM70) + add_library(mossinfer_sm70 OBJECT ${_mi_sm70_sources}) + set_target_properties(mossinfer_sm70 PROPERTIES CUDA_ARCHITECTURES 70-real POSITION_INDEPENDENT_CODE ON + CUDA_SEPARABLE_COMPILATION OFF CUDA_VISIBILITY_PRESET hidden VISIBILITY_INLINES_HIDDEN ON) + target_include_directories(mossinfer_sm70 PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/src) + target_include_directories(mossinfer_sm70 SYSTEM PRIVATE ${CMAKE_SOURCE_DIR}/third_party/cutlass/include + ${CMAKE_SOURCE_DIR}/third_party/cutlass/tools/util/include) + target_compile_options(mossinfer_sm70 PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) + target_sources(mossinfer_core PRIVATE $) + set_property(SOURCE ${CMAKE_SOURCE_DIR}/src/backend/probe.cpp + ${CMAKE_SOURCE_DIR}/src/kernels/sm70/sm70.cpp + ${CMAKE_SOURCE_DIR}/src/kernels/attention/flash_attention.cu + ${CMAKE_SOURCE_DIR}/src/kernels/basic/gemm.cpp + ${CMAKE_SOURCE_DIR}/src/kernels/gdn/chunk_rule.cu + ${CMAKE_SOURCE_DIR}/src/kernels/quant/int8_gemm.cu + ${CMAKE_SOURCE_DIR}/src/kernels/quant/int8_gemm_residual.cu + ${CMAKE_SOURCE_DIR}/src/kernels/quant/int8_gemm_situ.cu + APPEND PROPERTY COMPILE_DEFINITIONS MOSSINFER_HAVE_SM70=1) + if(EXISTS ${CMAKE_SOURCE_DIR}/tools/sm75_check.cu) + add_executable(sm70_check ${CMAKE_SOURCE_DIR}/tools/sm75_check.cu) + set_target_properties(sm70_check PROPERTIES CUDA_ARCHITECTURES 70-real) + target_include_directories(sm70_check PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/src) + target_compile_options(sm70_check PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) + target_link_libraries(sm70_check PRIVATE mossinfer_core) + endif() + if(EXISTS ${CMAKE_SOURCE_DIR}/tools/sm70_bench.cu) + add_executable(sm70_bench ${CMAKE_SOURCE_DIR}/tools/sm70_bench.cu) + set_target_properties(sm70_bench PROPERTIES CUDA_ARCHITECTURES 70-real) + target_link_libraries(sm70_bench PRIVATE mossinfer_core) + target_compile_options(sm70_bench PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) + endif() + message(STATUS "SM70: Volta kernel family built; compute capability 7.0 accepted") + endif() +endif() +unset(_mi_sm70_sources) diff --git a/cmake/modules/sm75.cmake b/cmake/modules/sm75.cmake index 472b57c..e3e58cb 100644 --- a/cmake/modules/sm75.cmake +++ b/cmake/modules/sm75.cmake @@ -19,22 +19,40 @@ if(_mi_sm75_sources) set_property(TARGET mossinfer_core PROPERTY SOURCES ${_mi_core_sources}) unset(_mi_core_sources) - add_library(mossinfer_sm75 OBJECT ${_mi_sm75_sources}) - set_target_properties(mossinfer_sm75 PROPERTIES - CUDA_ARCHITECTURES 75-real - POSITION_INDEPENDENT_CODE ON - CUDA_SEPARABLE_COMPILATION OFF) - target_include_directories(mossinfer_sm75 PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/src) - target_include_directories(mossinfer_sm75 SYSTEM PRIVATE - ${CMAKE_SOURCE_DIR}/third_party/flash-attention/csrc/flash_attn/src - ${CMAKE_SOURCE_DIR}/third_party/cutlass/include - ${CMAKE_SOURCE_DIR}/third_party/cutlass/tools/util/include) - target_compile_options(mossinfer_sm75 PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) - target_sources(mossinfer_core PRIVATE $) - message(STATUS "SM75: Turing kernel family built (${_mi_sm75_sources})") + set(_mi_sm75_present TRUE) + if(MOSSINFER_SM70) + # Conversion and VAE decode are SIMT; share their source with Volta without compiling Turing MMA for sm_70. + set(_mi_sm7x_sources ${_mi_sm75_sources}) + list(FILTER _mi_sm7x_sources INCLUDE REGEX "/(convert|vae_sm75)\\.cu$") + if(_mi_sm7x_sources) + list(REMOVE_ITEM _mi_sm75_sources ${_mi_sm7x_sources}) + add_library(mossinfer_sm7x OBJECT ${_mi_sm7x_sources}) + set_target_properties(mossinfer_sm7x PROPERTIES CUDA_ARCHITECTURES "70-real;75-real" + POSITION_INDEPENDENT_CODE ON CUDA_SEPARABLE_COMPILATION OFF) + target_include_directories(mossinfer_sm7x PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/src) + target_compile_options(mossinfer_sm7x PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) + target_sources(mossinfer_core PRIVATE $) + endif() + unset(_mi_sm7x_sources) + endif() + if(_mi_sm75_sources) + add_library(mossinfer_sm75 OBJECT ${_mi_sm75_sources}) + set_target_properties(mossinfer_sm75 PROPERTIES + CUDA_ARCHITECTURES 75-real + POSITION_INDEPENDENT_CODE ON + CUDA_SEPARABLE_COMPILATION OFF) + target_include_directories(mossinfer_sm75 PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/src) + target_include_directories(mossinfer_sm75 SYSTEM PRIVATE + ${CMAKE_SOURCE_DIR}/third_party/flash-attention/csrc/flash_attn/src + ${CMAKE_SOURCE_DIR}/third_party/cutlass/include + ${CMAKE_SOURCE_DIR}/third_party/cutlass/tools/util/include) + target_compile_options(mossinfer_sm75 PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) + target_sources(mossinfer_core PRIVATE $) + message(STATUS "SM75: Turing kernel family built (${_mi_sm75_sources})") + endif() endif() -if(MOSSINFER_SM75 AND _mi_sm75_sources) +if((MOSSINFER_SM75 OR MOSSINFER_SM70) AND _mi_sm75_present) # Portable kernels: pure SIMT translation units (no cp.async, no tensor-core asm, no CUTLASS), which run unchanged on # Turing but were only compiled for the core's sm_80+ list. They are built for sm_75 too, in their own OBJECT library # (the core keeps its list, so the mixed-precision kernels around them never see --arch=sm_75). @@ -71,6 +89,9 @@ if(MOSSINFER_SM75 AND _mi_sm75_sources) set_property(TARGET mossinfer_core PROPERTY SOURCES ${_mi_core_sources}) unset(_mi_core_sources) set(_mi_portable_archs 75 ${CMAKE_CUDA_ARCHITECTURES}) + if(MOSSINFER_SM70) + list(PREPEND _mi_portable_archs 70-real) + endif() list(REMOVE_DUPLICATES _mi_portable_archs) add_library(mossinfer_sm75_portable OBJECT ${_mi_portable_present}) set_target_properties(mossinfer_sm75_portable PROPERTIES @@ -90,16 +111,18 @@ if(MOSSINFER_SM75 AND _mi_sm75_sources) # src/backend/probe.h: the minimum compute capability and the architecture list this build accepts. Defined on the # sources that read it rather than on the whole target, so a build without MOSSINFER_SM75 keeps its old objects and # adding the option only recompiles the dispatchers. - set_property(SOURCE ${CMAKE_SOURCE_DIR}/src/backend/probe.cpp + if(MOSSINFER_SM75 OR MOSSINFER_SM70) + set_property(SOURCE ${CMAKE_SOURCE_DIR}/src/backend/probe.cpp ${CMAKE_SOURCE_DIR}/src/kernels/attention/flash_attention.cu ${CMAKE_SOURCE_DIR}/src/kernels/basic/gemm.cpp ${CMAKE_SOURCE_DIR}/src/kernels/gdn/chunk_rule.cu APPEND PROPERTY COMPILE_DEFINITIONS MOSSINFER_HAVE_SM75=1) - message(STATUS "SM75: compute capability 7.5 accepted") + message(STATUS "SM75: compute capability 7.5 accepted") + endif() # Bring-up harness (not installed): the engine's GPU tests compile their reference kernels for the build's CUDA # architectures, so a Turing device cannot run them; this tool builds its own host reference. - if(EXISTS ${CMAKE_SOURCE_DIR}/tools/sm75_check.cu) + if(MOSSINFER_SM75 AND EXISTS ${CMAKE_SOURCE_DIR}/tools/sm75_check.cu) add_executable(sm75_check ${CMAKE_SOURCE_DIR}/tools/sm75_check.cu) set_target_properties(sm75_check PROPERTIES CUDA_ARCHITECTURES 75) target_include_directories(sm75_check PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/src) @@ -108,3 +131,4 @@ if(MOSSINFER_SM75 AND _mi_sm75_sources) endif() endif() unset(_mi_sm75_sources) +unset(_mi_sm75_present) diff --git a/docs/integration.md b/docs/integration.md index 1e8b6a2..6bc521d 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -18,8 +18,8 @@ MossInfer is a low-level engine: it runs what the caller asks for, refuses what | Area | What works | |---|---| -| Platform | Linux x86-64 and Windows x64 (MSVC), each in a `cu12` and a `cu13` flavour, as a full package and a `-system-cuda` package. CI builds, packages and runs the CPU tests of all four standard flavours, plus the Windows CUDA 13 Turing build (`-sm75`); GPU runs are verified on Linux `cu13`. | -| GPUs | NVIDIA compute capability 8.x and 12.x: RTX 30xx, RTX 40xx, RTX 50xx, A100/A10/A40, L4/L40, RTX PRO Blackwell. Compute capability 7.5 (Turing: RTX 20xx, T4, GTX 16xx) with a build configured `-DMOSSINFER_SM75=ON`: the bf16 GEMMs run as fp16 tensor-core GEMMs, INT8/INT4 ConvRot files run the INT8 GEMMs (INT4 through the INT8 unpack fallback), the exact attention and the VAE use the sm_75 kernels, and an unset attention resolves to `exact` (there is no 8-bit attention kernel on Turing). See [sm75.md](sm75.md). | +| Platform | Linux x86-64 and Windows x64 (MSVC), each in a `cu12` and a `cu13` flavour, as a full package and a `-system-cuda` package. CI builds, packages and runs the CPU tests of all four standard flavours, plus the Windows CUDA 13 Turing build (`-sm75`) and Linux CUDA 12 Volta build (`-sm70`); GPU runs are verified on Linux `cu13`. | +| GPUs | NVIDIA compute capability 8.x and 12.x: RTX 30xx, RTX 40xx, RTX 50xx, A100/A10/A40, L4/L40, RTX PRO Blackwell. Compute capability 7.5 (Turing: RTX 20xx, T4, GTX 16xx) with a build configured `-DMOSSINFER_SM75=ON`: the bf16 GEMMs run as fp16 tensor-core GEMMs, INT8/INT4 ConvRot files run the INT8 GEMMs (INT4 through the INT8 unpack fallback), the exact attention and the VAE use the sm_75 kernels, and an unset attention resolves to `exact` (there is no 8-bit attention kernel on Turing). See [sm75.md](sm75.md). Compute capability 7.0 (Volta, Tesla V100) with CUDA 12 and `-DMOSSINFER_SM70=ON`: dedicated fp16 tensor-core attention and exact INT8 GEMMs, sharing the Turing layer’s conversion and VAE paths. See [sm70.md](sm70.md). | | Models | MOSS-Image Base and SFT only: the official Hugging Face folders (bf16), read in place, and fine-tunes with the same architecture; other models are refused when they load. See [Supported models](#supported-models). | | Component model files | a single-file DiT, text encoder or VAE next to the bundle folder: ComfyUI and community files, INT8 ConvRot, INT4 ConvRot and NVFP4 DiT (NVFP4 on RTX 50 only), INT8 ConvRot and W4A8 text encoder (`dit_file`, `text_encoder_file`, `vae_file`); community GGUF DiTs (Q8_0, Q6_K, Q5_K, Q4_K, Q4_0, BF16, F16), converted once into INT8 ConvRot or bf16 ([GGUF DiT files](#gguf-dit-files)) | | Generation | text-to-image with the official app's parameters: prompt, negative prompt, size, steps, CFG, shift, seed, skip layers, skip-layer guidance | @@ -209,6 +209,7 @@ cmake --install build-linux-cuda13 --component Unspecified --prefix build-linux- ``` - **glibc:** the archive needs the glibc of the machine that built it, or a newer one. Build on the oldest distribution you want to support. CI uses Ubuntu 22.04, which gives glibc 2.35. A build from a rolling distribution such as Arch will not start on Ubuntu LTS. +- **Linux CUDA 12 with Volta:** use the `linux-cuda12-sm70` configure/build/test/package presets. This produces a `-sm70` package; CUDA 13 cannot build or run the Volta backend. See [sm70.md](sm70.md#build). - **Windows CUDA 13 with Turing:** from a Visual Studio x64 developer prompt, run `cmake --preset windows-msvc-cuda13-sm75`, `cmake --build --preset windows-msvc-cuda13-sm75`, `ctest --preset windows-msvc-cuda13-sm75`, then `cpack --config build-windows-msvc-cuda13-sm75/CPackConfig.cmake -B build-windows-msvc-cuda13-sm75`. This creates the full `-sm75.zip` package; the ordinary presets remain Ampere-and-newer. Any preset takes `-DMOSSINFER_SM75=ON`; what the build adds is in [sm75.md](sm75.md#build). - **Other presets:** `linux-cuda12`, `windows-msvc-cuda13` and `windows-msvc-cuda12`. The Windows presets run from a Visual Studio 2022 x64 developer prompt. @@ -367,6 +368,8 @@ Decision procedure for the launcher: - `""`: no usable GPU. Show `probe.message` (e.g. "no flavour recommended: no device has compute capability >= 8.0") and the per-device `unsupported_reason`. 4. Use `devices[i].free_memory_bytes` for sizing hints in the UI. The engine checks memory again itself at generation time. +**Volta machines.** Install a CUDA 12 `-sm70` build. Even with a CUDA 13-capable driver, a machine containing only Volta devices receives `recommended_flavour: "cu12"`. The standard packages do not support these devices. The Turing compatibility restrictions also apply to Volta; see [sm70.md](sm70.md). + **Turing machines.** The standard packages carry no `sm_75` kernel image, so their `hello` lists a compute capability 7.5 device as unsupported (`"this build has no kernel image for compute capability 7.5"`). Install the `-sm75` package instead (the Windows CUDA 13 one today; any build configured `-DMOSSINFER_SM75=ON` carries it): its `hello` reports the diff --git a/docs/sm70.md b/docs/sm70.md new file mode 100644 index 0000000..88374d1 --- /dev/null +++ b/docs/sm70.md @@ -0,0 +1,105 @@ +# Volta (sm_70) support + +`-DMOSSINFER_SM70=ON` adds Volta kernels, including Tesla V100, through the same object-library and runtime-dispatch +boundary as the Turing compatibility layer. The default build remains Ampere or newer. The option also carries the +existing Turing kernels: sm_70 SASS is forward-compatible within compute capability 7.x, and Turing still dispatches +to its own attention and native INT8 kernels. + +## Build + +Use CUDA **12.8 or 12.9** for the release preset, CMake 3.28 or newer, Ninja and GoogleTest. CUDA 13 removed Volta compilation and library support; +configuration rejects `MOSSINFER_SM70=ON` with that toolkit. A newer NVIDIA driver can run the CUDA 12 binary. On a +machine whose supported devices are all Volta, the probe recommends `cu12` even when the driver supports CUDA 13. + +```sh +cmake --preset linux-cuda12-sm70 -DCMAKE_CUDA_COMPILER=/usr/local/cuda-12.9/bin/nvcc +cmake --build --preset linux-cuda12-sm70 +CUDA_VISIBLE_DEVICES="" ctest --preset linux-cuda12-sm70 -L '^cpu$' +build-linux-cuda12-sm70/sm70_check all +build-linux-cuda12-sm70/sm70_check big +build-linux-cuda12-sm70/mossinfer_tests --gtest_filter='Sm70Build.*:GpuSm70.*:GpuSm75Convert.*' +cpack --config build-linux-cuda12-sm70/CPackConfig.cmake -B build-linux-cuda12-sm70 +``` + +The package name ends in `-sm70`. The build/release matrix includes the Linux CUDA 12 SM70 preset. The CUDA architecture list of the ordinary core must stay sm_80 or newer; do not +set `CMAKE_CUDA_ARCHITECTURES=70` for the whole project. The option compiles the Volta kernel family separately and +adds sm_70 images to the portable SIMT kernels. `MOSSINFER_SM70=OFF` leaves the default kernel dispatch unchanged. + +## Kernel paths + +| Component | Volta path | +|---|---| +| bf16 linears | The existing compatibility layer's fp16 weights, row-scaled fp16 activations, cuBLASLt fp32 accumulation, and bias before the final bf16/fp32 rounding | +| INT8 ConvRot | Volta fp16 tensor cores with exact integer partial sums, then the existing quantization epilogues | +| INT4 ConvRot | The existing INT4-to-INT8 unpack route, followed by the Volta INT8 kernels | +| Exact attention | Dedicated fp16 WMMA kernel with fp32 online softmax and accumulation, bf16 input/output; non-causal head dimension 128 and causal 256, varlen and GQA | +| GatedDeltaNet | The chunked fp32 rule: Volta's 96 KiB shared memory fits it, so the smaller Turing recurrent fallback is unnecessary | +| VAE decode | The compatibility layer's fp32 convolution and attention; the planner retains reference decode precision | +| Elementwise kernels | The same portable SIMT sources, compiled for sm_70 as well | + +Volta has no native INT8 tensor-core instruction. The quantized kernel keeps weights in their original one-byte +layout and converts tiles to fp16 in shared memory. Every signed byte is exact in fp16. Each fp32 partial sums at +most 1024 products, bounded by `1024 * 128 * 128 = 2^24`; it is converted exactly to int32 before the next partial. +The final integer sum rounds to fp32 once before dequantization. This avoids both a persistent expanded weight copy +and the precision loss of accumulating an arbitrarily long integer dot product in fp32. The existing plain, +gated-residual and SiTU-GLU epilogues keep their rounding order and column permutation. + +Attention converts bf16 Q/K/V as it loads shared-memory tiles, leaves the inputs unchanged, and uses a bottom-right +causal mask. Fully masked query rows return zero and a positive-infinite log-sum-exp, matching the existing contract. +Its fp16 operand range and rounding limitations are the same as [the Turing layer](sm75.md#numerics). + +The quality policy also follows Turing: an unset attention mode resolves to `exact` with an explicit reason. Explicit +8-bit attention and NVFP4 are refused. The fused bf16 block kernels, runtime LoRAs, control adapters, VAE encoding +(img2img, inpainting and references), and hires fix retain the existing sm_7x refusal/fallback policy. No sm_80 kernel +is launched as a substitute for a missing Volta implementation. + +## Verification and performance + +`sm70_check` reuses the compatibility layer's host-reference harness, compiled for sm_70. It exercises conversion, +fp16 GEMM, large activation rows and workspace chunks, INT8 GEMM, VAE convolution/attention, and both exact attention +head dimensions. `GpuSm70` also checks all three INT8 epilogues, row-stride padding, tail shapes, CUDA Graph replay, +and cancellation at K=131056 against a host int64 sum. Attention also checks log-sum-exp (including fully masked +rows) against the host reference. + +`sm70_bench [rows] [sequence_length]` measures the four DiT GEMM shapes and both attention shapes with CUDA events +after warm-up. It is a development tool, not installed. Run the correctness tests before comparing its timings. + +Measured on a Tesla V100-SXM2-32GB (80 SMs), CUDA 12.9.86, driver 580.196.07, GCC 14.2, Release, on 2026-11-01. +One GPU per measurement, default clocks, no concurrent workload on that GPU. CUDA-event means over 20 launches +following five warm-up launches, `sm70_bench 8222 4111`: + +| Kernel | Shape | Time | Effective INT8 dot-product rate | +|---|---|---:|---:| +| QKV | M=8222, N=9216, K=3072 | 8.03 ms | 58.0 TOP/s | +| Projection + gated residual | M=8222, N=3072, K=3072 | 3.15 ms | 49.2 TOP/s | +| Up/gate + SiTU-GLU | M=8222, N=16384, K=3072 | 12.93 ms | 64.0 TOP/s | +| Down + gated residual | M=8222, N=3072, K=8192 | 6.15 ms | 67.3 TOP/s | +| DiT exact attention | B=2, L=4111, H=24, D=128 | 17.41 ms | — | +| Causal text attention | B=2, L=282, H=16, KV=4, D=256 | 0.259 ms | — | + +The INT8 tile is 64x128x32, with 32x64 warps and the CUTLASS Volta shared-memory swizzle/register pipeline. +The main loop's shared storage is reused by the fused output stage. The SiTU factors use the same bf16 lookup-table +values as the Turing path. No global dequantization buffer or expanded resident weights are allocated. Attention +uses 128 query rows for D=128 and 64 for D=256, 32 keys per tile; a larger key tile was slower on this device. + +End to end: MOSS-Image-SFT, Comfy-Org INT8 ConvRot DiT and Qwen3.5 9B W4A8 text encoder, `high`, exact attention, +Euler, CFG 3, seed 42, prompt `a red apple on a wooden table, natural light, detailed photograph`: + +| Run | Measured result | +|---|---| +| 1024x1024, 20 steps, `max_speed`, embedding cache hit | 45.35 s total; 41.35 s denoise (~2.07 s/step), 3.29 s decode; 6.74 GiB device peak | +| 1024x1024, 3 steps, `max_speed`, embedding cache hit | 10.24 s total; 6.16 s denoise (~2.05 s/step), 3.28 s decode; 6.74 GiB device peak | +| 256x256, 3 steps, `min_memory`, uncached blue-teapot prompt | 2.91 s total; 1.17 s encode, 1.13 s denoise, 0.13 s decode; 1.04 GiB device peak; all 42 DiT blocks streamed | + +The initial correctness implementation used a straightforward WMMA INT8 loop and took ~13.9 s/step at 1024x1024. +The optimized kernels reduce that by ~6.8x; the three-step PNGs before and after optimization are byte-identical. +This is a comparison within this Volta implementation, not a benchmark against another engine or an upstream V100 +backend. The upstream build did not run on V100. The 20-step timing includes weight loading and decode but uses a +cached text embedding; the streamed run separately exercises cold text encoding. + +Validation on V100: `sm70_check all` and `big`, all six `Sm70Build`/`GpuSm70`/`GpuSm75Convert` tests, Compute Sanitizer +memcheck (zero errors), and racecheck on attention and all INT8 epilogues including multiple accumulation segments +(zero errors/warnings). The default-off probe build passes its nine CPU tests. CTest reports no failures across 517 selected CPU cases +(with the normal missing-fixture skips), excluding the pre-existing `cli_check_verdict` CMake-script issue on CMake +3.31: that script drops empty list elements under unset CMP0007. Running it through a wrapper that sets CMP0007 to +NEW also passes. The shared-library `moss-server` and CLI also produce identical three-step 1024x1024 PNGs. The implementation was hardware-tested on V100; other devices were not available for regression. diff --git a/src/backend/probe.cpp b/src/backend/probe.cpp index 181f478..3715461 100644 --- a/src/backend/probe.cpp +++ b/src/backend/probe.cpp @@ -134,10 +134,15 @@ void evaluate_device(DeviceInfo& dev, int driver_version, int build_cuda_version // The sm_80 path needs Ampere; a build with the sm_75 kernel family reaches down to Turing. dev.reason = "compute capability " + cc_string(dev.cc_major, dev.cc_minor) + " is below the required " + cc_string(kMinComputeMajor, kMinComputeMinor) + - (MOSSINFER_HAVE_SM75 ? " (Turing or newer)" : " (Ampere or newer)"); + (MOSSINFER_HAVE_SM70 ? " (Volta or newer)" : + MOSSINFER_HAVE_SM75 ? " (Turing or newer)" : " (Ampere or newer)"); return; } const Flavour build = flavour_of_cuda_version(build_cuda_version); + if (dev.cc_major == 7 && dev.cc_minor < 5 && build_cuda_version >= 13000) { + dev.reason = "Volta needs a CUDA 12 build; CUDA 13 removed its compilation and library support"; + return; + } const int min_driver = build == Flavour::kNone ? build_cuda_version / 1000 * 1000 : min_driver_version(build); if (driver_version < min_driver) { dev.reason = driver_version == 0 ? std::string("no NVIDIA driver is loaded") @@ -165,6 +170,7 @@ void recommend(ProbeResult& r, const std::vector& archs) { auto append = [&r](const std::string& s) { r.reason += (r.reason.empty() ? "" : "; ") + s; }; bool any_capable = false; bool any_image = false; + bool any_cuda13_image = false; std::string no_image; // "sm_90, sm_100" of capable devices without a kernel image for (const DeviceInfo& d : r.devices) { if (std::tie(d.cc_major, d.cc_minor) < std::tie(kMinComputeMajor, kMinComputeMinor)) continue; @@ -173,6 +179,7 @@ void recommend(ProbeResult& r, const std::vector& archs) { // driver is new enough for its own toolkit. if (compatible_image(archs, d.cc_major, d.cc_minor, true)) { any_image = true; + if (d.cc_major > 7 || (d.cc_major == 7 && d.cc_minor >= 5)) any_cuda13_image = true; } else { no_image += (no_image.empty() ? "sm_" : ", sm_") + std::to_string(d.cc_major) + std::to_string(d.cc_minor); } @@ -191,7 +198,8 @@ void recommend(ProbeResult& r, const std::vector& archs) { build_architectures_string() + ")"); return; } - const Flavour f = flavour_for_driver(r.driver_version); + Flavour f = flavour_for_driver(r.driver_version); + if (f == Flavour::kCuda13 && !any_cuda13_image) f = Flavour::kCuda12; if (f == Flavour::kNone) { append(r.driver_version == 0 ? std::string("no flavour recommended: no NVIDIA driver is loaded") : "no flavour recommended: the driver supports CUDA " + diff --git a/src/backend/probe.h b/src/backend/probe.h index 1cb1440..9a8b1c7 100644 --- a/src/backend/probe.h +++ b/src/backend/probe.h @@ -18,12 +18,15 @@ namespace mi::backend { // Minimum compute capability this build can run. Every sm_80 kernel needs bf16 tensor cores and cp.async; a build // with the sm_75 (Turing) kernel family (MOSSINFER_SM75, cmake/modules/sm75.cmake) also runs compute capability 7.5, -// through the fp16/INT8 fallbacks in src/kernels/sm75. +// through the fp16/INT8 fallbacks in src/kernels/sm75. MOSSINFER_SM70 adds the Volta kernels for 7.0 with CUDA 12. #ifndef MOSSINFER_HAVE_SM75 #define MOSSINFER_HAVE_SM75 0 #endif -inline constexpr int kMinComputeMajor = MOSSINFER_HAVE_SM75 ? 7 : 8; -inline constexpr int kMinComputeMinor = MOSSINFER_HAVE_SM75 ? 5 : 0; +#ifndef MOSSINFER_HAVE_SM70 +#define MOSSINFER_HAVE_SM70 0 +#endif +inline constexpr int kMinComputeMajor = (MOSSINFER_HAVE_SM70 || MOSSINFER_HAVE_SM75) ? 7 : 8; +inline constexpr int kMinComputeMinor = MOSSINFER_HAVE_SM70 ? 0 : (MOSSINFER_HAVE_SM75 ? 5 : 0); // Release flavours, named by the CUDA major version of the toolkit they are built with. enum class Flavour { kNone, kCuda12, kCuda13 }; diff --git a/src/kernels/attention/flash_attention.cu b/src/kernels/attention/flash_attention.cu index 5ef530e..5b8ba99 100644 --- a/src/kernels/attention/flash_attention.cu +++ b/src/kernels/attention/flash_attention.cu @@ -10,6 +10,7 @@ #include "core/varlen.h" #include "kernels/attention/flash_fwd_kernels.h" #include "kernels/sm75/sm75.h" +#include "kernels/sm70/sm70.h" namespace mi::attention { namespace { @@ -152,6 +153,13 @@ void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& // difference. int device = 0, smem_per_sm = 0, smem_per_block = 0; MI_CUDA_CHECK(cudaGetDevice(&device)); +#if MOSSINFER_HAVE_SM70 + if (kernels::sm70::volta()) { + kernels::sm70::flash_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen, causal, scale, workspace, + workspace_bytes, stream); + return; + } +#endif #if MOSSINFER_HAVE_SM75 if (kernels::sm75::turing()) { // Turing: the sm_75 family's own exact attention (src/kernels/sm75/attention_sm75.cu, fp16 MMA; the vendored diff --git a/src/kernels/basic/gemm.cpp b/src/kernels/basic/gemm.cpp index 5ea5912..f788a01 100644 --- a/src/kernels/basic/gemm.cpp +++ b/src/kernels/basic/gemm.cpp @@ -213,7 +213,7 @@ struct Gemm::Impl { Gemm::Gemm() : impl_(std::make_unique()) { MI_LT_CHECK(cublasLtCreate(&impl_->handle)); -#if MOSSINFER_HAVE_SM75 +#if MOSSINFER_HAVE_SM75 || MOSSINFER_HAVE_SM70 // Turing: no bf16 tensor cores, so bf16 linears run as fp16 with fp32 accumulation (see Impl::linear_f16). int device = 0, major = 0; if (cudaGetDevice(&device) == cudaSuccess && diff --git a/src/kernels/gdn/chunk_rule.cu b/src/kernels/gdn/chunk_rule.cu index b16eb45..6977f95 100644 --- a/src/kernels/gdn/chunk_rule.cu +++ b/src/kernels/gdn/chunk_rule.cu @@ -21,6 +21,7 @@ #include "kernels/gdn/checks.h" #include "kernels/gdn/gdn.h" #include "kernels/sm75/sm75.h" +#include "kernels/sm70/sm70.h" namespace mi::gdn { namespace { @@ -565,7 +566,11 @@ void chunk_gated_delta_rule(const Tensor& qkv, const Tensor& g, const Tensor& be // computes the same rule with 33 KiB; its numerics are fp32 like the chunked path's. Only an MOSSINFER_SM75 build // checks (once per device, kernels::sm75::turing()); without it the probe never accepts a 7.x device. #if MOSSINFER_HAVE_SM75 - if (kernels::sm75::turing()) { + if (kernels::sm75::turing() +#if MOSSINFER_HAVE_SM70 + && !kernels::sm70::volta() // Volta's 96 KiB fits the faster chunked rule. +#endif + ) { constexpr int kSm75Smem = (kHeadDim * kSm75Cols + 2 * kHeadDim + 4) * static_cast(sizeof(float)); const dim3 grid(kNumVHeads, num_seqs, kHeadDim / kSm75Cols); recurrent_sm75_kernel<__nv_bfloat16><<>>( diff --git a/src/kernels/quant/int8_gemm.cu b/src/kernels/quant/int8_gemm.cu index 68a7078..c9ff07b 100644 --- a/src/kernels/quant/int8_gemm.cu +++ b/src/kernels/quant/int8_gemm.cu @@ -7,6 +7,7 @@ #include "kernels/quant/int8_gemm_impl.cuh" #include "kernels/sm75/int8_sm75.h" #include "kernels/sm75/sm75.h" +#include "kernels/sm70/sm70.h" namespace mi::kernels::quant { namespace detail { @@ -94,7 +95,7 @@ Int8Tile resolve_tile(Int8Tile tile, const Problem& p) { void check_sm75_tile(const char* op, Int8Tile tile) { if (tile != Int8Tile::kAuto) { - fail(op, std::string("a forced tile (") + int8_tile_name(tile) + ") names an sm_80 kernel; sm_75 runs one tile"); + fail(op, std::string("a forced tile (") + int8_tile_name(tile) + ") names an sm_80 kernel; sm_7x selects its own tiles"); } } @@ -156,6 +157,13 @@ void int8_gemm(const Int8Operands& op, const __nv_bfloat16* bias, __nv_bfloat16* detail::check_output(name, "out", out, ldo, p.n); detail::check_vector(name, "bias", bias, /*optional=*/true); if (p.m == 0) return; +#if MOSSINFER_HAVE_SM70 + if (kernels::sm70::volta()) { + detail::check_sm75_tile(name, tile); + kernels::sm70::int8_gemm(op, bias, out, ldo, stream); + return; + } +#endif if (kernels::sm75::turing()) { // Turing's 8x8x16 INT8 MMA (src/kernels/sm75/int8_gemm_sm75.cu) detail::check_sm75_tile(name, tile); kernels::sm75::int8_gemm(op, bias, out, ldo, stream); diff --git a/src/kernels/quant/int8_gemm_residual.cu b/src/kernels/quant/int8_gemm_residual.cu index 6eb2b01..57167b7 100644 --- a/src/kernels/quant/int8_gemm_residual.cu +++ b/src/kernels/quant/int8_gemm_residual.cu @@ -2,6 +2,7 @@ #include "kernels/quant/int8_gemm_impl.cuh" #include "kernels/sm75/int8_sm75.h" #include "kernels/sm75/sm75.h" +#include "kernels/sm70/sm70.h" namespace mi::kernels::quant { namespace detail { @@ -38,6 +39,13 @@ void int8_gemm_gated_residual(const Int8Operands& op, const __nv_bfloat16* bias, detail::check_vector(name, "bias", bias, /*optional=*/true); detail::check_vector(name, "gate", gate, /*optional=*/false); if (p.m == 0) return; +#if MOSSINFER_HAVE_SM70 + if (kernels::sm70::volta()) { + detail::check_sm75_tile(name, tile); + kernels::sm70::int8_gemm_gated_residual(op, bias, gate, residual, ldr, stream); + return; + } +#endif if (kernels::sm75::turing()) { detail::check_sm75_tile(name, tile); kernels::sm75::int8_gemm_gated_residual(op, bias, gate, residual, ldr, stream); diff --git a/src/kernels/quant/int8_gemm_situ.cu b/src/kernels/quant/int8_gemm_situ.cu index 6562b6c..418e9a4 100644 --- a/src/kernels/quant/int8_gemm_situ.cu +++ b/src/kernels/quant/int8_gemm_situ.cu @@ -7,6 +7,7 @@ #include "kernels/quant/int8_gemm_impl.cuh" #include "kernels/sm75/int8_sm75.h" #include "kernels/sm75/sm75.h" +#include "kernels/sm70/sm70.h" namespace mi::kernels::quant { namespace detail { @@ -104,6 +105,13 @@ void int8_gemm_situ_glu(const Int8Operands& op, __nv_bfloat16* out, std::int64_t std::to_string(ldo)); } if (p.m == 0) return; +#if MOSSINFER_HAVE_SM70 + if (kernels::sm70::volta()) { + detail::check_sm75_tile(name, tile); + kernels::sm70::int8_gemm_situ_glu(op, out, ldo, stream); + return; + } +#endif if (kernels::sm75::turing()) { detail::check_sm75_tile(name, tile); kernels::sm75::int8_gemm_situ_glu(op, out, ldo, stream); diff --git a/src/kernels/sm70/attention_sm70.cu b/src/kernels/sm70/attention_sm70.cu new file mode 100644 index 0000000..9980d8d --- /dev/null +++ b/src/kernels/sm70/attention_sm70.cu @@ -0,0 +1,241 @@ +// Exact varlen/GQA attention on Volta: fp16 WMMA with fp32 scores, online softmax and accumulation. Q/K/V stay +// bf16 in global memory; only shared-memory tiles are converted. No quadratic score buffer or hidden allocation. +// A warp owns 16 query rows. The sm_70 accumulator row groups use lane XOR 2/8 (unlike Turing's XOR 1/2); +// coordinates are loaded through WMMA itself, so output stores follow its fragment layout. +#include +#include +#include + +#include +#include + +#include "core/cuda_check.h" +#include "kernels/sm70/sm70.h" + +namespace mi::kernels::sm70 { +namespace { + +namespace wmma = nvcuda::wmma; +constexpr float kLog2e = 1.4426950408889634f; + +struct Params { + const __nv_bfloat16* q = nullptr; + const __nv_bfloat16* k = nullptr; + const __nv_bfloat16* v = nullptr; + __nv_bfloat16* o = nullptr; + float* lse = nullptr; + std::int64_t q_row = 0, q_head = 0; + std::int64_t k_row = 0, k_head = 0; + std::int64_t v_row = 0, v_head = 0; + std::int64_t o_row = 0, o_head = 0; + const int* cu_q = nullptr; + const int* cu_k = nullptr; + int heads = 0, kv_heads = 0, total_q = 0; + float scale = 1.f; + bool causal = false; +}; + +template +__global__ void attention_kernel(const Params p) { + constexpr int kBM = D == 128 ? 128 : 64, kThreads = kBM * 2; + constexpr int kQStride = D + 8, kVStride = BN + 8; + using Acc = wmma::fragment; + using A = wmma::fragment; + using B = wmma::fragment; + extern __shared__ __align__(32) char storage[]; + __half* qs = reinterpret_cast<__half*>(storage); + __half* ks = qs + kBM * kQStride; + __half* vt = ks + BN * kQStride; + __half* ps = vt + D * kVStride; + const int tid = threadIdx.x, warp = tid / 32; + const int batch = blockIdx.y, head = blockIdx.z; + const int q_start = p.cu_q[batch], k_start = p.cu_k[batch]; + const int len_q = p.cu_q[batch + 1] - q_start, len_k = p.cu_k[batch + 1] - k_start; + const int row_base = blockIdx.x * kBM; + if (row_base >= len_q) return; + const int rows = min(kBM, len_q - row_base); + const int kv_head = head / (p.heads / p.kv_heads), shift = len_k - len_q; + + // ps is scratch until the first P tile. Encode coordinates into an accumulator fragment once. + float* coords = reinterpret_cast(ps); + for (int i = tid; i < 256; i += kThreads) coords[i] = static_cast(i); + __syncthreads(); + Acc positions; + wmma::load_matrix_sync(positions, coords, 16, wmma::mem_row_major); + int rr[Acc::num_elements], cc[Acc::num_elements]; +#pragma unroll + for (int e = 0; e < Acc::num_elements; ++e) { + rr[e] = static_cast(positions.x[e]) / 16; + cc[e] = static_cast(positions.x[e]) % 16; + } + for (int i = tid; i < kBM * D; i += kThreads) { + const int r = i / D, d = i % D; + qs[r * kQStride + d] = __float2half_rn(r < rows + ? __bfloat162float(p.q[(q_start + row_base + r) * p.q_row + head * p.q_head + d]) : 0.f); + } + __syncthreads(); + Acc o[D / 16]; +#pragma unroll + for (int d = 0; d < D / 16; ++d) wmma::fill_fragment(o[d], 0.f); + float m[2] = {-1e30f, -1e30f}, l[2] = {0.f, 0.f}; + const int end = p.causal ? min(len_k, max(0, row_base + rows + shift)) : len_k; + for (int j0 = 0; j0 < end; j0 += BN) { + for (int i = tid; i < BN * D; i += kThreads) { + const int j = i / D, d = i % D, token = k_start + j0 + j; + const bool valid = j0 + j < len_k; + ks[j * kQStride + d] = __float2half_rn(valid ? __bfloat162float(p.k[token * p.k_row + kv_head * p.k_head + d]) : 0.f); + vt[d * kVStride + j] = __float2half_rn(valid ? __bfloat162float(p.v[token * p.v_row + kv_head * p.v_head + d]) : 0.f); + } + __syncthreads(); + Acc s[BN / 16]; +#pragma unroll + for (int j = 0; j < BN / 16; ++j) wmma::fill_fragment(s[j], 0.f); +#pragma unroll + for (int d = 0; d < D; d += 16) { + A a; + wmma::load_matrix_sync(a, qs + warp * 16 * kQStride + d, kQStride); +#pragma unroll + for (int j = 0; j < BN / 16; ++j) { + B b; + wmma::load_matrix_sync(b, ks + j * 16 * kQStride + d, kQStride); + wmma::mma_sync(s[j], a, b, s[j]); + } + } + float mx[2] = {-1e30f, -1e30f}; +#pragma unroll + for (int j = 0; j < BN / 16; ++j) { +#pragma unroll + for (int e = 0; e < Acc::num_elements; ++e) { + const int key = j0 + j * 16 + cc[e], row = row_base + warp * 16 + rr[e]; + float value = s[j].x[e] * p.scale; + if (key >= len_k || (p.causal && key > row + shift)) value = -1e30f; + s[j].x[e] = value; + mx[(e & 2) / 2] = fmaxf(mx[(e & 2) / 2], value); + } + } + float alpha[2]; +#pragma unroll + for (int r = 0; r < 2; ++r) { + mx[r] = fmaxf(mx[r], __shfl_xor_sync(0xffffffffu, mx[r], 2)); + mx[r] = fmaxf(mx[r], __shfl_xor_sync(0xffffffffu, mx[r], 8)); + mx[r] = fmaxf(mx[r], m[r]); + alpha[r] = m[r] <= -1e29f ? 0.f : exp2f((m[r] - mx[r]) * kLog2e); + m[r] = mx[r]; + } + float sum[2] = {0.f, 0.f}; +#pragma unroll + for (int j = 0; j < BN / 16; ++j) { +#pragma unroll + for (int e = 0; e < Acc::num_elements; ++e) { + const int r = (e & 2) / 2; + const float value = s[j].x[e] <= -1e29f ? 0.f : exp2f((s[j].x[e] - m[r]) * kLog2e); + sum[r] += value; + ps[(warp * 16 + rr[e]) * kVStride + j * 16 + cc[e]] = __float2half_rn(value); + } + } +#pragma unroll + for (int r = 0; r < 2; ++r) { + sum[r] += __shfl_xor_sync(0xffffffffu, sum[r], 2); + sum[r] += __shfl_xor_sync(0xffffffffu, sum[r], 8); + l[r] = l[r] * alpha[r] + sum[r]; + } +#pragma unroll + for (int d = 0; d < D / 16; ++d) { +#pragma unroll + for (int e = 0; e < Acc::num_elements; ++e) o[d].x[e] *= alpha[(e & 2) / 2]; + } + __syncwarp(); // this warp alone writes and reads its P rows +#pragma unroll + for (int j = 0; j < BN; j += 16) { + A a; + wmma::load_matrix_sync(a, ps + warp * 16 * kVStride + j, kVStride); +#pragma unroll + for (int d = 0; d < D / 16; ++d) { + B b; + wmma::load_matrix_sync(b, vt + d * 16 * kVStride + j, kVStride); + wmma::mma_sync(o[d], a, b, o[d]); + } + } + __syncthreads(); + } +#pragma unroll + for (int e = 0; e < Acc::num_elements; ++e) { + const int r = (e & 2) / 2, row = warp * 16 + rr[e]; + if (row < rows) { + const float inv = l[r] > 0.f ? 1.f / l[r] : 0.f; +#pragma unroll + for (int d = 0; d < D / 16; ++d) { + p.o[(q_start + row_base + row) * p.o_row + head * p.o_head + d * 16 + cc[e]] = __float2bfloat16_rn(o[d].x[e] * inv); + } + if (cc[e] == 0 && p.lse) p.lse[static_cast(head) * p.total_q + q_start + row_base + row] = + l[r] > 0.f ? m[r] + logf(l[r]) : INFINITY; + } + } +} + +template +void launch(const Params& p, int max_seqlen, int batch, cudaStream_t stream) { + constexpr int kBM = D == 128 ? 128 : 64, kThreads = kBM * 2; + constexpr int bytes = ((kBM + BN) * (D + 8) + (D + kBM) * (BN + 8)) * sizeof(__half); + static_assert(bytes <= 96 * 1024); + auto kernel = &attention_kernel; + if constexpr (bytes > 48 * 1024) MI_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes)); + const dim3 grid(static_cast((max_seqlen + kBM - 1) / kBM), static_cast(batch), + static_cast(p.heads)); + kernel<<>>(p); + MI_CUDA_CHECK_LAUNCH(); +} + +bool head_dim_supported(int d) { return d == 128 || d == 256; } + +} // namespace + +void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& out, const Tensor& cu_seqlens_q, + const Tensor& cu_seqlens_k, int max_seqlen, bool causal, float scale, void* workspace, + std::size_t workspace_bytes, cudaStream_t stream) { + const std::int64_t total_q = q.dim(0); + if (total_q == 0) return; + const std::int64_t heads = q.dim(1); + const std::int64_t head_dim = q.dim(2); + const std::int64_t kv_heads = k.dim(1); + if (!head_dim_supported(static_cast(head_dim))) { + throw std::invalid_argument("sm70::flash_fwd: head_dim must be 128 or 256"); + } + if ((head_dim == 128 && causal) || (head_dim == 256 && !causal)) { + throw std::invalid_argument("sm70::flash_fwd: head_dim 128 is non-causal, 256 is causal"); + } + if (workspace_bytes < static_cast(total_q) * heads * sizeof(float)) { + throw std::invalid_argument("sm70::flash_fwd: workspace too small for the log-sum-exp"); + } + + Params p; + p.q = static_cast(q.data); + p.k = static_cast(k.data); + p.v = static_cast(v.data); + p.o = static_cast<__nv_bfloat16*>(out.data); + p.lse = static_cast(workspace); + p.q_row = q.stride[0]; + p.q_head = q.stride[1]; + p.k_row = k.stride[0]; + p.k_head = k.stride[1]; + p.v_row = v.stride[0]; + p.v_head = v.stride[1]; + p.o_row = out.stride[0]; + p.o_head = out.stride[1]; + p.cu_q = cu_seqlens_q.ptr(); + p.cu_k = cu_seqlens_k.ptr(); + p.heads = static_cast(heads); + p.kv_heads = static_cast(kv_heads); + p.total_q = static_cast(total_q); + p.scale = scale; + p.causal = causal; + const int batch = static_cast(cu_seqlens_q.dim(0) - 1); + + if (head_dim == 128) { + launch<128, 32>(p, max_seqlen, batch, stream); + } else { + launch<256, 32>(p, max_seqlen, batch, stream); + } +} + +} // namespace mi::kernels::sm70 diff --git a/src/kernels/sm70/int8_gemm_sm70.cu b/src/kernels/sm70/int8_gemm_sm70.cu new file mode 100644 index 0000000..2530666 --- /dev/null +++ b/src/kernels/sm70/int8_gemm_sm70.cu @@ -0,0 +1,237 @@ +// Volta INT8 ConvRot GEMMs: fp16 tensor-core partials, exact int32 sums and the existing quant epilogues. +#include "kernels/sm70/sm70.h" + +#include +#include +#include +#include + +#include "kernels/quant/int8_gemm_impl.cuh" + +namespace mi::kernels::sm70 { +namespace qd = quant::detail; + +namespace { + +// The SiTU factor tables of the quantized epilogue: the sm_80 tables (int8_gemm_situ.cu) are filled by a kernel built for +// sm_80+ only, so the sm_70 path keeps its own 256 KiB copy, filled once per device by a kernel built for 70. The +// factors are the same functions, so a lookup is bit-identical to the sm_80 table's value. +constexpr int kBf16Values = 1 << 16; +__device__ std::uint16_t g_situ_gate[kBf16Values]; +__device__ std::uint16_t g_situ_hidden[kBf16Values]; + +__global__ void situ_tables_sm70_kernel() { + const int i = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + const float x = __uint_as_float(static_cast(i) << 16); + g_situ_gate[i] = __bfloat16_as_ushort(__float2bfloat16_rn(qd::situ_gate_factor(x))); + g_situ_hidden[i] = __bfloat16_as_ushort(__float2bfloat16_rn(qd::situ_hidden_factor(x))); +} + +// Fill once per device (see quant::detail::situ_factor_tables; the same lock and capture handling). +qd::SituFactorTables situ_tables(cudaStream_t stream) { + constexpr int kMaxDevices = 64; + struct Device { + std::atomic filled{false}; + std::atomic gate{nullptr}, hidden{nullptr}; + }; + static Device devices[kMaxDevices]; + static std::mutex fill_mutex; + int dev = 0; + MI_CUDA_CHECK(cudaGetDevice(&dev)); + if (dev < 0 || dev >= kMaxDevices) throw std::runtime_error("mi::kernels::sm70: device ordinal out of range"); + Device& d = devices[dev]; + const auto fill = [&] { + constexpr int kThreads = 256; + situ_tables_sm70_kernel<<>>(); + MI_CUDA_CHECK_LAUNCH(); + }; + if (!d.filled.load(std::memory_order_acquire)) { + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + MI_CUDA_CHECK(cudaStreamIsCapturing(stream, &capture)); + if (capture != cudaStreamCaptureStatusNone) { + fill(); + } else { + const std::lock_guard lock(fill_mutex); + if (!d.filled.load(std::memory_order_relaxed)) { + fill(); + MI_CUDA_CHECK(cudaStreamSynchronize(stream)); + d.filled.store(true, std::memory_order_release); + } + } + } + if (d.gate.load(std::memory_order_acquire) == nullptr) { + void *gate = nullptr, *hidden = nullptr; + MI_CUDA_CHECK(cudaGetSymbolAddress(&gate, g_situ_gate)); + MI_CUDA_CHECK(cudaGetSymbolAddress(&hidden, g_situ_hidden)); + d.hidden.store(static_cast(hidden), std::memory_order_release); + d.gate.store(static_cast(gate), std::memory_order_release); + } + return {d.gate.load(std::memory_order_acquire), d.hidden.load(std::memory_order_acquire)}; +} + +// The main loop uses CUTLASS's Volta shared-memory swizzle and pipelined register loads. Global iterators still +// read INT8; MmaPipelined's NumericArrayConverter stages exact fp16 integers. Fold every 1024 products into int32: +// each fp32 partial is bounded by 2^24, preserving integer sums even with cancellation at the maximum supported K. +using Sm70Tile = qd::TileCfg<64, 128, 32, 32, 64, 2>; +using Sm70Instr = cutlass::gemm::GemmShape<8, 8, 4>; + +template +struct ExactMma : Base { + using typename Base::FragmentC; + using typename Base::IteratorA; + using typename Base::IteratorB; + using typename Base::SharedStorage; + SharedStorage& storage; + int thread, warp, lane; + CUTLASS_DEVICE ExactMma(SharedStorage& s, int t, int w, int l) + : Base(s, t, w, l), storage(s), thread(t), warp(w), lane(l) {} + + CUTLASS_DEVICE void operator()(int iterations, FragmentC& accum, IteratorA a, IteratorB b, const FragmentC& src) { + cutlass::Array sums; + sums.clear(); + constexpr int kIterations = 1024 / Base::Shape::kK; + for (int start = 0; start < iterations; start += kIterations) { + FragmentC partial, zero; + zero.clear(); + Base mma(storage, thread, warp, lane); + mma(min(kIterations, iterations - start), partial, a, b, zero); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < FragmentC::kElements; ++i) sums[i] += __float2int_rn(partial[i]); + __syncthreads(); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < kIterations; ++j) { + ++a; + ++b; + } + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < FragmentC::kElements; ++i) accum[i] = __int2float_rn(sums[i]) + src[i]; + } +}; + +enum class Epilogue { kPlain, kResidual, kSitu }; +using qd::add; +using qd::mul; +using qd::rbf; +using qd::dequant; + +template +struct Kernel { + using Core = cutlass::gemm::threadblock::DefaultMmaCore< + Sm70Tile::Threadblock, Sm70Tile::Warp, Sm70Instr, cutlass::half_t, qd::RowMajor, + cutlass::half_t, qd::ColumnMajor, float, qd::RowMajor, cutlass::arch::OpClassTensorOp, 2>; + using IteratorA = cutlass::transform::threadblock::PredicatedTileIterator< + cutlass::MatrixShape, std::int8_t, qd::RowMajor, 1, + typename Core::IteratorThreadMapA, 8>; + using PermuteB = std::conditional_t; + using IteratorB = cutlass::transform::threadblock::PredicatedTileIterator< + cutlass::MatrixShape, std::int8_t, qd::ColumnMajor, 0, + typename Core::IteratorThreadMapB, 8, false, PermuteB>; + using Mma = ExactMma>; +}; + +template +__global__ void int8_kernel(quant::Int8Operands p, typename Kernel::IteratorA::Params pa, + typename Kernel::IteratorB::Params pb, const __nv_bfloat16* bias, + const __nv_bfloat16* gate, __nv_bfloat16* out, std::int64_t ldo, + qd::SituFactorTables tables) { + using K = Kernel; + constexpr int BM = Sm70Tile::Threadblock::kM, BN = Sm70Tile::Threadblock::kN; + constexpr int kThreads = (BM / Sm70Tile::Warp::kM) * (BN / Sm70Tile::Warp::kN) * 32; + __shared__ union Storage { + typename K::Mma::SharedStorage mainloop; + float result[BM * BN]; + } storage; + const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32; + const int m0 = blockIdx.x * BM, n0 = blockIdx.y * BN; + typename K::IteratorA a(pa, const_cast(p.x_q), {int(p.m), int(p.k)}, tid, {m0, 0}); + typename K::IteratorB b(pb, const_cast(p.w), {int(p.k), int(p.n)}, tid, {0, n0}); + typename K::Mma mma(storage.mainloop, tid, warp, lane); + typename K::Mma::FragmentC accum, zero; + zero.clear(); + mma((p.k + Sm70Tile::Threadblock::kK - 1) / Sm70Tile::Threadblock::kK, accum, a, b, zero); + __syncthreads(); + using IteratorC = typename K::Core::MmaTensorOp::IteratorC; + IteratorC c({storage.result, qd::RowMajor(BN)}, lane); + c.add_tile_offset({warp % (BM / Sm70Tile::Warp::kM), warp / (BM / Sm70Tile::Warp::kM)}); + c.store(accum); + __syncthreads(); + const float* result = storage.result; + constexpr int kOutN = E == Epilogue::kSitu ? BN / 2 : BN; + for (int i = tid; i < BM * kOutN; i += kThreads) { + const int row = i / kOutN, c = i % kOutN; + const int r = m0 + row, col = (E == Epilogue::kSitu ? n0 / 2 : n0) + c; + if (r >= p.m || col >= (E == Epilogue::kSitu ? p.n / 2 : p.n)) continue; + const float xs = p.x_scale[r]; + float value; + if constexpr (E == Epilogue::kSitu) { + constexpr int g = fused::detail::kUpGateGroup; + const int hc = (c / g) * (2 * g) + c % g; + const float h = rbf(dequant(result[row * BN + hc], xs, p.w_scale[col], 0.f)); + const float gt = rbf(dequant(result[row * BN + hc + g], xs, p.w_scale[col + p.n / 2], 0.f)); + value = qd::mul(qd::bf16_entry(tables.hidden, qd::bf16_index(h)), + qd::bf16_entry(tables.gate, qd::bf16_index(gt))); + } else { + value = dequant(result[row * BN + c], xs, p.w_scale[col], bias ? __bfloat162float(bias[col]) : 0.f); + if constexpr (E == Epilogue::kResidual) { + value = add(__bfloat162float(out[r * ldo + col]), rbf(mul(rbf(value), __bfloat162float(gate[col])))); + } + } + out[r * ldo + col] = __float2bfloat16_rn(value); + } +} + +template +void launch(const quant::Int8Operands& p, const __nv_bfloat16* bias, const __nv_bfloat16* gate, + __nv_bfloat16* out, std::int64_t ldo, cudaStream_t stream) { + constexpr int kBM = Sm70Tile::Threadblock::kM, kBN = Sm70Tile::Threadblock::kN; + constexpr int kThreads = (kBM / Sm70Tile::Warp::kM) * (kBN / Sm70Tile::Warp::kN) * 32; + const dim3 grid(static_cast((p.m + kBM - 1) / kBM), static_cast((p.n + kBN - 1) / kBN)); + const typename Kernel::IteratorA::Params pa{qd::RowMajor(p.ldx)}; + const typename Kernel::IteratorB::Params pb{qd::ColumnMajor(p.k)}; + qd::SituFactorTables tables{}; + if constexpr (E == Epilogue::kSitu) tables = situ_tables(stream); + int8_kernel<<>>(p, pa, pb, bias, gate, out, ldo, tables); + MI_CUDA_CHECK_LAUNCH(); +} + +} // namespace + +void int8_gemm(const quant::Int8Operands& op, const __nv_bfloat16* bias, __nv_bfloat16* out, std::int64_t ldo, + cudaStream_t stream) { + constexpr const char* name = "sm70::int8_gemm"; + const qd::Problem p = qd::check_operands(name, op); + qd::check_output(name, "out", out, ldo, p.n); + qd::check_vector(name, "bias", bias, /*optional=*/true); + if (p.m == 0) return; + launch(op, bias, nullptr, out, ldo, stream); +} + +void int8_gemm_gated_residual(const quant::Int8Operands& op, const __nv_bfloat16* bias, const __nv_bfloat16* gate, + __nv_bfloat16* residual, std::int64_t ldr, cudaStream_t stream) { + constexpr const char* name = "sm70::int8_gemm_gated_residual"; + const qd::Problem p = qd::check_operands(name, op); + qd::check_output(name, "residual", residual, ldr, p.n); + qd::check_vector(name, "bias", bias, /*optional=*/true); + qd::check_vector(name, "gate", gate, /*optional=*/false); + if (p.m == 0) return; + launch(op, bias, gate, residual, ldr, stream); +} + +void int8_gemm_situ_glu(const quant::Int8Operands& op, __nv_bfloat16* out, std::int64_t ldo, cudaStream_t stream) { + constexpr const char* name = "sm70::int8_gemm_situ_glu"; + if (op.n % (2 * qd::kUpGateGroup) != 0) { + kernels::detail::fail(name, "n = 2F needs F % 4 == 0, got n " + std::to_string(op.n)); + } + const qd::Problem p = qd::check_operands(name, op); + if (out == nullptr || reinterpret_cast(out) % 8 != 0 || ldo < p.n / 2 || ldo % 4 != 0) { + kernels::detail::fail(name, "out must be 8-byte aligned with a row stride >= F and a multiple of 4, got " + + std::to_string(ldo)); + } + if (p.m == 0) return; + launch(op, nullptr, nullptr, out, ldo, stream); +} + +} // namespace mi::kernels::sm70 diff --git a/src/kernels/sm70/sm70.cpp b/src/kernels/sm70/sm70.cpp new file mode 100644 index 0000000..e9a2837 --- /dev/null +++ b/src/kernels/sm70/sm70.cpp @@ -0,0 +1,38 @@ +#include "kernels/sm70/sm70.h" + +#include + +namespace mi::kernels::sm70 { + +bool available() { +#if MOSSINFER_HAVE_SM70 + return true; +#else + return false; +#endif +} + +bool volta() { + if (!available()) return false; + static std::atomic cached[64]; + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { + cudaGetLastError(); + return false; + } + if (device < 0 || device >= 64) return false; + int cc = cached[device].load(std::memory_order_relaxed); + if (cc == 0) { + int major = 0, minor = 0; + if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess || + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device) != cudaSuccess) { + cudaGetLastError(); + return false; + } + cc = major * 10 + minor; + cached[device].store(cc, std::memory_order_relaxed); + } + return cc >= 70 && cc < 75; +} + +} // namespace mi::kernels::sm70 diff --git a/src/kernels/sm70/sm70.h b/src/kernels/sm70/sm70.h new file mode 100644 index 0000000..bb7a883 --- /dev/null +++ b/src/kernels/sm70/sm70.h @@ -0,0 +1,26 @@ +// Volta (sm_70) tensor-core kernels. The SIMT conversions and fp16 cuBLASLt path are shared with sm75; attention +// and quantized linears use Volta's fp16 MMA, with fp32 accumulation and the engine's bf16 epilogue rounding. +#pragma once + +#include "core/tensor.h" +#include "kernels/quant/quant.h" + +namespace mi::kernels::sm70 { + +bool available(); +// Cached per device; true only on compute capability 7.0/7.2 in a build carrying the Volta kernel family. +bool volta(); + +void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& out, const Tensor& cu_seqlens_q, + const Tensor& cu_seqlens_k, int max_seqlen, bool causal, float scale, void* workspace, + std::size_t workspace_bytes, cudaStream_t stream); + +// INT8 inputs are converted exactly to fp16 in shared memory. Partial integer dot products fit fp32 exactly; +// periodically folding them into int32 preserves the native INT8 path's bits for every supported K. +void int8_gemm(const quant::Int8Operands& op, const __nv_bfloat16* bias, __nv_bfloat16* out, std::int64_t ldo, + cudaStream_t stream); +void int8_gemm_gated_residual(const quant::Int8Operands& op, const __nv_bfloat16* bias, const __nv_bfloat16* gate, + __nv_bfloat16* residual, std::int64_t ldr, cudaStream_t stream); +void int8_gemm_situ_glu(const quant::Int8Operands& op, __nv_bfloat16* out, std::int64_t ldo, cudaStream_t stream); + +} // namespace mi::kernels::sm70 diff --git a/src/pipeline/pipeline.cpp b/src/pipeline/pipeline.cpp index 814f901..ff21b15 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -1505,25 +1505,25 @@ void Pipeline::Impl::explain_control(const Request::Control& control, Result& r) // inpainting, reference images, and the hires fix, whose second pass encodes the upscaled image) and the LoRA GEMMs // (runtime LoRAs, and control adapters, which carry a LoRA). The message names the reason and what does work, instead // of failing at a missing kernel image mid-request. -void detail::refuse_unsupported_on_sm75(const Request& req) { +void detail::refuse_unsupported_on_sm7x(const Request& req) { if (!kernels::sm75::turing()) return; if (req.init_image != nullptr) { - throw NotSupportedError("img2img and inpainting encode the input with the VAE encoder, which has no sm_75 kernel " + throw NotSupportedError("img2img and inpainting encode the input with the VAE encoder, which has no sm_7x kernel " "in this build; text-to-image works"); } if (req.hires.enabled()) { - throw NotSupportedError("the hires fix encodes the upscaled image with the VAE encoder, which has no sm_75 kernel " + throw NotSupportedError("the hires fix encodes the upscaled image with the VAE encoder, which has no sm_7x kernel " "in this build; text-to-image without it works"); } if (req.references.enabled()) { - throw NotSupportedError("reference images are encoded with the VAE encoder, which has no sm_75 kernel in this " + throw NotSupportedError("reference images are encoded with the VAE encoder, which has no sm_7x kernel in this " "build; text-to-image without them works"); } if (!req.loras.empty()) { - throw NotSupportedError("runtime LoRAs need sm_80+ GEMM kernels; this sm_75 build cannot apply them"); + throw NotSupportedError("runtime LoRAs need sm_80+ GEMM kernels; this sm_7x build cannot apply them"); } if (req.control.enabled()) { - throw NotSupportedError("control adapters are LoRAs and need sm_80+ GEMM kernels; this sm_75 build cannot run them"); + throw NotSupportedError("control adapters are LoRAs and need sm_80+ GEMM kernels; this sm_7x build cannot run them"); } } @@ -1532,7 +1532,7 @@ Result Pipeline::run(const Request& req, const GenerateControl& control, bool pr Impl& m = *impl_; check_user_paths(m.options, req); validate(req); - refuse_unsupported_on_sm75(req); + refuse_unsupported_on_sm7x(req); Result r; r.tier.tier = m.options.tier; // what a run cancelled before the tier is resolved reports Timings& tm = r.timings; diff --git a/src/pipeline/pipeline_impl.h b/src/pipeline/pipeline_impl.h index 64a1266..bd61867 100644 --- a/src/pipeline/pipeline_impl.h +++ b/src/pipeline/pipeline_impl.h @@ -218,8 +218,8 @@ inline void check_user_paths(const PipelineOptions& o, const Request& r) { for (const lora::LoraSpec& l : r.loras) util::check_user_path(l.path, "loras"); } -// On a Turing device: throws NotSupportedError for what still needs sm_80+ kernels (pipeline.cpp). No-op elsewhere. -void refuse_unsupported_on_sm75(const Request& req); +// On a Volta or Turing device: throws NotSupportedError for what still needs sm_80+ kernels (pipeline.cpp). No-op elsewhere. +void refuse_unsupported_on_sm7x(const Request& req); inline moss_dit::SkipSet skip_set(const std::vector& layers) { moss_dit::SkipSet s; @@ -1018,7 +1018,7 @@ struct Pipeline::Impl { moss_dit::BlockPath block_path(std::string& note) { if (options.block_path == moss_dit::BlockPath::kOpByOp) return options.block_path; if (kernels::sm75::turing()) { - note = "sm_75: DiT blocks run op by op (the fused bf16 CUTLASS GEMMs need sm_80+ tensor cores)"; + note = "sm_7x: DiT blocks run op by op (the fused bf16 CUTLASS GEMMs need sm_80+ tensor cores)"; return moss_dit::BlockPath::kOpByOp; } const moss_dit::FusedCheck c = check_fused(); diff --git a/src/pipeline/session.cpp b/src/pipeline/session.cpp index 05f7eef..7b50593 100644 --- a/src/pipeline/session.cpp +++ b/src/pipeline/session.cpp @@ -381,7 +381,7 @@ Session::Session(Pipeline& pipeline, JobGate& gate, SessionOptions options) { if (tok::python_strip(q.prompt).empty()) q.prompt = "preload"; check_user_paths(pipeline.impl_->options, q); pipeline::validate(q); - refuse_unsupported_on_sm75(q); + refuse_unsupported_on_sm7x(q); if (q.accel.step_cache_on() && !q.accel.step_cache_threshold) { throw std::invalid_argument("The step cache needs an explicit threshold in (0, 1] (step_cache_threshold)."); } @@ -466,7 +466,7 @@ void Session::validate(const RunSpec& spec, const SessionOptions& options) { const Guidance g = run_guidance(spec); const Request run_q = run_request(options, spec, g); pipeline::validate(run_q); - refuse_unsupported_on_sm75(run_q); + refuse_unsupported_on_sm7x(run_q); if (spec.sigmas.empty() && !(std::isfinite(spec.denoise_strength) && spec.denoise_strength >= 0 && spec.denoise_strength <= 1)) { throw std::invalid_argument("Denoise strength must be between 0 and 1."); } diff --git a/tests/unit/test_sm70.cpp b/tests/unit/test_sm70.cpp new file mode 100644 index 0000000..de5874a --- /dev/null +++ b/tests/unit/test_sm70.cpp @@ -0,0 +1,168 @@ +// Host references for Volta's integer Tensor Core path, including sums above fp32's exact-integer range. +#include + +#include +#include + +#include +#include +#include + +#include "backend/probe.h" +#include "core/cuda_check.h" +#include "kernels/basic/elementwise.h" +#include "kernels/sm70/sm70.h" + +namespace { + +using mi::DType; +using mi::Device; +using mi::Tensor; + +std::uint16_t bf16(float f) { + const auto u = std::bit_cast(f); + return static_cast((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} +float f32(std::uint16_t b) { return std::bit_cast(static_cast(b) << 16); } + +template +class Buffer { + public: + explicit Buffer(const std::vector& values) : n_(values.size()) { + MI_CUDA_CHECK(cudaMalloc(&p_, n_ * sizeof(T))); + MI_CUDA_CHECK(cudaMemcpy(p_, values.data(), n_ * sizeof(T), cudaMemcpyHostToDevice)); + } + ~Buffer() { cudaFree(p_); } + Buffer(const Buffer&) = delete; + Buffer& operator=(const Buffer&) = delete; + T* get() const { return p_; } + std::vector host() const { + std::vector values(n_); + MI_CUDA_CHECK(cudaMemcpy(values.data(), p_, n_ * sizeof(T), cudaMemcpyDeviceToHost)); + return values; + } + private: + T* p_ = nullptr; + std::size_t n_; +}; + +void check_gemm(int m, int n, int k, bool adversarial) { + const int ldx = k + 16, ldo = n + 8; + std::vector x(static_cast(m) * ldx, 0), w(static_cast(n) * k); + std::uint32_t rng = 42; + const auto next = [&] { + rng ^= rng << 13; + rng ^= rng >> 17; + rng ^= rng << 5; + return static_cast(static_cast(rng % 256) - 128); + }; + for (int r = 0; r < m; ++r) { + for (int d = 0; d < k; ++d) x[static_cast(r) * ldx + d] = adversarial ? -128 : next(); + } + for (int c = 0; c < n; ++c) { + for (int d = 0; d < k; ++d) { + // Near-cancelling large positive/negative partial sums expose an unbounded fp32 accumulator's lost units. + const int v = d < k / 2 ? 127 : -127; + w[static_cast(c) * k + d] = adversarial ? static_cast(d == k - 1 ? c - 8 : v) : next(); + } + } + std::vector xs(static_cast(m)), ws(static_cast(n)); + std::vector bias(static_cast(n)), gate(static_cast(n)); + for (int r = 0; r < m; ++r) xs[static_cast(r)] = 0.001f * static_cast(r % 7 + 1); + for (int c = 0; c < n; ++c) { + ws[static_cast(c)] = 0.002f * static_cast(c % 5 + 1); + bias[static_cast(c)] = bf16(static_cast(c % 13 - 6) * 0.125f); + gate[static_cast(c)] = bf16(static_cast(c % 9 - 4) * 0.25f); + } + const std::uint16_t sentinel = bf16(-123.f); + std::vector expected(static_cast(m) * ldo, sentinel); + std::vector unbiased(static_cast(m) * n); + for (int r = 0; r < m; ++r) { + for (int c = 0; c < n; ++c) { + std::int64_t acc = 0; + for (int d = 0; d < k; ++d) acc += static_cast(x[static_cast(r) * ldx + d]) * + static_cast(w[static_cast(c) * k + d]); + const float value = (static_cast(acc) * xs[static_cast(r)]) * ws[static_cast(c)]; + expected[static_cast(r) * ldo + c] = bf16(value + f32(bias[static_cast(c)])); + unbiased[static_cast(r) * n + c] = bf16(value); + } + } + Buffer dx(x), dw(w); + Buffer dxs(xs), dws(ws); + Buffer db(bias), dg(gate), out(std::vector(expected.size(), sentinel)); + const mi::kernels::quant::Int8Operands p{dx.get(), ldx, dxs.get(), dw.get(), dws.get(), m, n, k}; + const auto ptr = [](std::uint16_t* b) { return reinterpret_cast<__nv_bfloat16*>(b); }; + mi::kernels::quant::int8_gemm(p, ptr(db.get()), ptr(out.get()), ldo, nullptr); + ASSERT_EQ(out.host(), expected); + mi::kernels::quant::int8_gemm_gated_residual(p, ptr(db.get()), ptr(dg.get()), ptr(out.get()), ldo, nullptr); + for (int r = 0; r < m; ++r) { + for (int c = 0; c < n; ++c) { + auto& e = expected[static_cast(r) * ldo + c]; + e = bf16(f32(e) + f32(bf16(f32(e) * f32(gate[static_cast(c)])))); + } + } + ASSERT_EQ(out.host(), expected); + + const int f = n / 2, lds = f + 4; + Buffer linear(unbiased), situ(std::vector(static_cast(m) * lds, sentinel)); + Buffer ref(std::vector(static_cast(m) * lds, sentinel)); + Tensor st = Tensor::view(ref.get(), DType::kBF16, Device::kCUDA, {m, f}); + st.stride[0] = lds; + mi::kernels::situ_glu(Tensor::view(linear.get(), DType::kBF16, Device::kCUDA, {m, n}), st, nullptr); + mi::kernels::quant::int8_gemm_situ_glu(p, ptr(situ.get()), lds, nullptr); + ASSERT_EQ(situ.host(), ref.host()); + + // No allocation or synchronization in the GEMM: graph replay must obey its stream and preserve stride padding. + cudaStream_t stream = nullptr; + cudaGraph_t graph = nullptr; + cudaGraphExec_t exec = nullptr; + MI_CUDA_CHECK(cudaStreamCreate(&stream)); + MI_CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + mi::kernels::quant::int8_gemm_situ_glu(p, ptr(situ.get()), lds, stream); + MI_CUDA_CHECK(cudaStreamEndCapture(stream, &graph)); + MI_CUDA_CHECK(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); + for (int i = 0; i < 3; ++i) MI_CUDA_CHECK(cudaGraphLaunch(exec, stream)); + MI_CUDA_CHECK(cudaStreamSynchronize(stream)); + EXPECT_EQ(situ.host(), ref.host()); + cudaGraphExecDestroy(exec); + cudaGraphDestroy(graph); + cudaStreamDestroy(stream); +} + +} // namespace + +TEST(Sm70Build, KernelFamilyPresence) { + EXPECT_EQ(mi::kernels::sm70::available(), MOSSINFER_HAVE_SM70 != 0); +} + +TEST(Sm70Build, ProbeAndCuda12Recommendation) { + const auto archs = mi::backend::parse_cuda_architectures("70-real,80-real"); + mi::backend::DeviceInfo d; + d.cc_major = 7; + d.cc_minor = 0; + mi::backend::evaluate_device(d, 13000, 12090, archs); + EXPECT_EQ(d.supported, MOSSINFER_HAVE_SM70 != 0); + mi::backend::evaluate_device(d, 13000, 13000, archs); + EXPECT_FALSE(d.supported); + mi::backend::ProbeResult r; + r.driver_version = 13000; + r.devices = {d}; + mi::backend::recommend(r, archs); + EXPECT_EQ(r.recommended, MOSSINFER_HAVE_SM70 ? mi::backend::Flavour::kCuda12 : mi::backend::Flavour::kNone); + d.cc_major = 8; + r.devices.push_back(d); + mi::backend::recommend(r, archs); + EXPECT_EQ(r.recommended, mi::backend::Flavour::kCuda13); +} + +TEST(GpuSm70, Int8EpiloguesAndStrides) { + if (!mi::kernels::sm70::volta()) GTEST_SKIP() << "needs MOSSINFER_SM70 and a Volta device"; + check_gemm(3, 8, 16, false); + check_gemm(67, 136, 1040, false); + check_gemm(129, 256, 4096, false); +} + +TEST(GpuSm70, Int8ExactLargeCancellation) { + if (!mi::kernels::sm70::volta()) GTEST_SKIP() << "needs MOSSINFER_SM70 and a Volta device"; + check_gemm(3, 16, 131056, true); +} diff --git a/tools/sm70_bench.cu b/tools/sm70_bench.cu new file mode 100644 index 0000000..ca5b8a8 --- /dev/null +++ b/tools/sm70_bench.cu @@ -0,0 +1,100 @@ +// V100 rate probe, not installed. Run sm70_check first; timings are CUDA events after warm-up, not host launch time. +// Usage: sm70_bench [GEMM rows (default 8222)] [attention sequence length (default 4111)]. +#include +#include + +#include +#include +#include +#include + +#include "core/cuda_check.h" +#include "kernels/attention/flash_attention.h" +#include "kernels/sm70/sm70.h" + +namespace { + +class Buffer { + public: + explicit Buffer(std::size_t bytes) { + MI_CUDA_CHECK(cudaMalloc(&p_, bytes)); + MI_CUDA_CHECK(cudaMemset(p_, 0, bytes)); + } + ~Buffer() { cudaFree(p_); } + Buffer(const Buffer&) = delete; + Buffer& operator=(const Buffer&) = delete; + template T* get() const { return static_cast(p_); } + private: + void* p_ = nullptr; +}; + +template +float time_ms(F&& f) { + for (int i = 0; i < 5; ++i) f(); + cudaEvent_t start = nullptr, end = nullptr; + MI_CUDA_CHECK(cudaEventCreate(&start)); + MI_CUDA_CHECK(cudaEventCreate(&end)); + MI_CUDA_CHECK(cudaEventRecord(start)); + constexpr int kIters = 20; + for (int i = 0; i < kIters; ++i) f(); + MI_CUDA_CHECK(cudaEventRecord(end)); + MI_CUDA_CHECK(cudaEventSynchronize(end)); + float ms = 0.f; + MI_CUDA_CHECK(cudaEventElapsedTime(&ms, start, end)); + cudaEventDestroy(start); + cudaEventDestroy(end); + return ms / kIters; +} + +void gemm(int m, int n, int k, int epilogue) { + Buffer x(static_cast(m) * k), w(static_cast(n) * k); + Buffer xs(static_cast(m) * 4), ws(static_cast(n) * 4); + Buffer bias(static_cast(n) * 2), gate(static_cast(n) * 2); + Buffer out(static_cast(m) * n * 2); + mi::kernels::quant::Int8Operands p{x.get(), k, xs.get(), w.get(), ws.get(), m, n, k}; + const float ms = time_ms([&] { + if (epilogue == 2) mi::kernels::quant::int8_gemm_situ_glu(p, out.get<__nv_bfloat16>(), n / 2, nullptr); + else if (epilogue == 1) mi::kernels::quant::int8_gemm_gated_residual(p, bias.get<__nv_bfloat16>(), gate.get<__nv_bfloat16>(), + out.get<__nv_bfloat16>(), n, nullptr); + else mi::kernels::quant::int8_gemm(p, bias.get<__nv_bfloat16>(), out.get<__nv_bfloat16>(), n, nullptr); + }); + std::printf("gemm M=%d N=%d K=%d epilogue=%d: %.3f ms, %.2f TOP/s\n", m, n, k, epilogue, ms, + 2.0 * m * n * k / (ms * 1e9)); +} + +void attention(int length, int heads, int kv_heads, int dim, bool causal) { + constexpr int batch = 2; + const int tokens = batch * length; + Buffer q(static_cast(tokens) * heads * dim * 2), out(static_cast(tokens) * heads * dim * 2); + Buffer k(static_cast(tokens) * kv_heads * dim * 2), v(static_cast(tokens) * kv_heads * dim * 2); + Buffer cu((batch + 1) * sizeof(int)), lse(static_cast(tokens) * heads * sizeof(float)); + const int bounds[] = {0, length, tokens}; + MI_CUDA_CHECK(cudaMemcpy(cu.get(), bounds, sizeof(bounds), cudaMemcpyHostToDevice)); + const auto tensor = [](const Buffer& b, int t, int h, int d) { + return mi::Tensor::view(b.get(), mi::DType::kBF16, mi::Device::kCUDA, {t, h, d}); + }; + const auto tq = tensor(q, tokens, heads, dim), tk = tensor(k, tokens, kv_heads, dim), tv = tensor(v, tokens, kv_heads, dim); + const auto to = tensor(out, tokens, heads, dim); + const auto tc = mi::Tensor::view(cu.get(), mi::DType::kI32, mi::Device::kCUDA, {batch + 1}); + const float ms = time_ms([&] { + mi::attention::flash_fwd(tq, tk, tv, to, tc, tc, length, causal, 1.f / std::sqrt(static_cast(dim)), + lse.get(), static_cast(tokens) * heads * sizeof(float), nullptr); + }); + std::printf("attention L=%d H=%d KV=%d D=%d causal=%d: %.3f ms\n", length, heads, kv_heads, dim, causal, ms); +} + +} // namespace + +int main(int argc, char** argv) { + std::setvbuf(stdout, nullptr, _IONBF, 0); + if (!mi::kernels::sm70::volta()) return 1; + const int m = argc > 1 ? std::atoi(argv[1]) : 8222; + const int length = argc > 2 ? std::atoi(argv[2]) : 4111; + if (m <= 0 || length <= 0) return 2; + gemm(m, 9216, 3072, 0); + gemm(m, 3072, 3072, 1); + gemm(m, 16384, 3072, 2); + gemm(m, 3072, 8192, 1); + attention(length, 24, 24, 128, false); + attention(282, 16, 4, 256, true); +} diff --git a/tools/sm75_check.cu b/tools/sm75_check.cu index 6fc6f02..41a4c48 100644 --- a/tools/sm75_check.cu +++ b/tools/sm75_check.cu @@ -1,7 +1,7 @@ -// sm_75 bring-up harness: runs the Turing kernel family on the current device and checks it against a host -// reference. Built only with MOSSINFER_SM75 (cmake/modules/sm75.cmake), as a development tool, not installed shipped. +// sm_7x bring-up harness: runs the compatibility kernels on the current device against a host reference. +// Built as sm75_check with MOSSINFER_SM75, or sm70_check with MOSSINFER_SM70; a development tool, not installed. // -// Usage: sm75_check [attn128|attn256|convert|all] +// Usage: sm75_check / sm70_check [attn128|attn256|convert|gemm|int8|vae|big|all] // attn128 varlen non-causal attention, head_dim 128 (the DiT's joint attention shape), GQA 2:1 // attn256 varlen causal attention, head_dim 256 (the text encoder's full-attention shape), GQA 4:1, and causal // with key lengths unlike the query lengths (bottom-right alignment) @@ -97,10 +97,11 @@ class DevBuf { std::vector reference_attention(const std::vector& q, const std::vector& k, const std::vector& v, const std::vector& cu_q, const std::vector& cu_k, int heads, int kv_heads, int dim, bool causal, - float scale) { + float scale, std::vector* lse = nullptr) { const int batch = static_cast(cu_q.size()) - 1; const int total_q = cu_q.back(); std::vector out(static_cast(total_q) * heads * dim, 0.f); + if (lse) lse->assign(static_cast(total_q) * heads, INFINITY); const int group = heads / kv_heads; std::vector scores; for (int b = 0; b < batch; ++b) { @@ -129,6 +130,7 @@ std::vector reference_attention(const std::vector& q, cons scores[static_cast(j)] = std::exp(scores[static_cast(j)] - max_s); sum += scores[static_cast(j)]; } + if (lse) (*lse)[static_cast(h) * total_q + i] = max_s + std::log(sum); std::size_t base = (static_cast(i) * heads + h) * dim; for (int j = 0; j < n; ++j) { const float p = scores[static_cast(j)] / sum; @@ -201,7 +203,15 @@ int run_attention(const char* name, int heads, int kv_heads, int dim, bool causa std::fflush(stdout); const auto got = dout.download(qb.size()); - const auto want = reference_attention(qb, kb, vb, cu, cu_k, heads, kv_heads, dim, causal, scale); + std::vector expected_lse; + const auto want = reference_attention(qb, kb, vb, cu, cu_k, heads, kv_heads, dim, causal, scale, &expected_lse); + const auto actual_lse = dls.download(lse.size()); + bool lse_ok = true; + for (std::size_t i = 0; i < lse.size(); ++i) { + if (std::isinf(expected_lse[i])) lse_ok &= actual_lse[i] == INFINITY; + else lse_ok &= std::fabs(actual_lse[i] - expected_lse[i]) <= 2e-4f; + } + std::printf(" log-sum-exp: %s\n", lse_ok ? "OK" : "MISMATCH"); double dot = 0, ng = 0, nw = 0; float max_abs = 0.f; for (std::size_t i = 0; i < want.size(); ++i) { @@ -234,7 +244,7 @@ int run_attention(const char* name, int heads, int kv_heads, int dim, bool causa } } MI_CUDA_CHECK(cudaGetLastError()); - return cosine >= 0.9999 && inputs_kept ? 0 : 1; + return cosine >= 0.9999 && inputs_kept && lse_ok ? 0 : 1; } int run_convert() { @@ -534,14 +544,14 @@ int main(int argc, char** argv) { failures += run_gemm(1000, 512, 256, /*out_f32=*/true, 999, std::size_t{1} << 20); failures += run_gemm(1000, 512, 256, false, 7, std::size_t{1} << 20); } - if (what == "vae") { + if (what == "vae" || what == "all") { failures += run_vae_conv_case(1, 6, 6, 8, 8, 3, 61); failures += run_vae_conv_case(1, 8, 8, 96, 96, 3, 65); // decoder widths: several channel chunks and K blocks failures += run_vae_conv_case(1, 5, 7, 96, 3, 1, 69); // conv_out: 1x1, K = 3 failures += run_vae_conv_case(1, 4, 4, 32, 96, 1, 73); // 1x1 shortcut failures += run_vae_attn(); } - if (what == "int8") { + if (what == "int8" || what == "all") { failures += run_int8_gemm(64, 128, 256); failures += run_int8_gemm(3, 8, 64); // tiny M: one tile row failures += run_int8_gemm(257, 384, 1024); From a1126dc5e78b8b90bbbf7984c27688a45acc25fc Mon Sep 17 00:00:00 2001 From: samyyc Date: Sun, 1 Nov 2026 20:27:42 +0800 Subject: [PATCH 2/3] perf(sm70): accelerate FP32 VAE and tune Volta DiT tiles --- cmake/modules/sm70.cmake | 2 + docs/sm70.md | 88 +++++++++----- src/kernels/sm70/attention_sm70.cu | 17 +-- src/kernels/sm70/int8_gemm_sm70.cu | 58 ++++++---- src/kernels/sm70/vae_attention_sm70.cu | 154 +++++++++++++++++++++++++ src/kernels/sm70/vae_conv2d_sm70.cu | 70 +++++++++++ src/kernels/sm70/vae_sm70.h | 18 +++ src/kernels/vae/attention.cu | 17 ++- src/kernels/vae/attention.h | 5 +- src/kernels/vae/conv2d.cu | 48 ++------ src/kernels/vae/conv2d_epilogue.cuh | 46 ++++++++ src/kernels/vae/vae_kernels.h | 1 + tests/unit/test_sm70.cpp | 135 ++++++++++++++++++++++ tools/sm70_bench.cu | 43 ++++++- tools/sm75_check.cu | 4 +- 15 files changed, 605 insertions(+), 101 deletions(-) create mode 100644 src/kernels/sm70/vae_attention_sm70.cu create mode 100644 src/kernels/sm70/vae_conv2d_sm70.cu create mode 100644 src/kernels/sm70/vae_sm70.h create mode 100644 src/kernels/vae/conv2d_epilogue.cuh diff --git a/cmake/modules/sm70.cmake b/cmake/modules/sm70.cmake index d2245b3..610c21c 100644 --- a/cmake/modules/sm70.cmake +++ b/cmake/modules/sm70.cmake @@ -18,6 +18,8 @@ if(_mi_sm70_sources) set_property(SOURCE ${CMAKE_SOURCE_DIR}/src/backend/probe.cpp ${CMAKE_SOURCE_DIR}/src/kernels/sm70/sm70.cpp ${CMAKE_SOURCE_DIR}/src/kernels/attention/flash_attention.cu + ${CMAKE_SOURCE_DIR}/src/kernels/vae/attention.cu + ${CMAKE_SOURCE_DIR}/src/kernels/vae/conv2d.cu ${CMAKE_SOURCE_DIR}/src/kernels/basic/gemm.cpp ${CMAKE_SOURCE_DIR}/src/kernels/gdn/chunk_rule.cu ${CMAKE_SOURCE_DIR}/src/kernels/quant/int8_gemm.cu diff --git a/docs/sm70.md b/docs/sm70.md index 88374d1..03029cb 100644 --- a/docs/sm70.md +++ b/docs/sm70.md @@ -34,7 +34,7 @@ adds sm_70 images to the portable SIMT kernels. `MOSSINFER_SM70=OFF` leaves the | INT4 ConvRot | The existing INT4-to-INT8 unpack route, followed by the Volta INT8 kernels | | Exact attention | Dedicated fp16 WMMA kernel with fp32 online softmax and accumulation, bf16 input/output; non-causal head dimension 128 and causal 256, varlen and GQA | | GatedDeltaNet | The chunked fp32 rule: Volta's 96 KiB shared memory fits it, so the smaller Turing recurrent fallback is unnecessary | -| VAE decode | The compatibility layer's fp32 convolution and attention; the planner retains reference decode precision | +| VAE decode | FP32 SIMT implicit-GEMM convolution and cuBLASLt FP32 attention slabs; the planner retains reference decode precision | | Elementwise kernels | The same portable SIMT sources, compiled for sm_70 as well | Volta has no native INT8 tensor-core instruction. The quantized kernel keeps weights in their original one-byte @@ -59,10 +59,13 @@ is launched as a substitute for a missing Volta implementation. fp16 GEMM, large activation rows and workspace chunks, INT8 GEMM, VAE convolution/attention, and both exact attention head dimensions. `GpuSm70` also checks all three INT8 epilogues, row-stride padding, tail shapes, CUDA Graph replay, and cancellation at K=131056 against a host int64 sum. Attention also checks log-sum-exp (including fully masked -rows) against the host reference. +rows) against the host reference, including the larger query tile. VAE tests compare batched convolutions against +FP64 host sums, including spatial/channel tails and an aliased residual; attention covers both head dimensions, +strided Q/K/V, slab tails, a 16-byte-aligned workspace and Graph replay. `sm70_bench [rows] [sequence_length]` measures the four DiT GEMM shapes and both attention shapes with CUDA events -after warm-up. It is a development tool, not installed. Run the correctness tests before comparing its timings. +after warm-up. `sm70_bench vae` measures FP32 VAE convolutions and the 16384-token attention. These are development +tools, not installed. Run the correctness tests before comparing their timings. Measured on a Tesla V100-SXM2-32GB (80 SMs), CUDA 12.9.86, driver 580.196.07, GCC 14.2, Release, on 2026-11-01. One GPU per measurement, default clocks, no concurrent workload on that GPU. CUDA-event means over 20 launches @@ -70,36 +73,65 @@ following five warm-up launches, `sm70_bench 8222 4111`: | Kernel | Shape | Time | Effective INT8 dot-product rate | |---|---|---:|---:| -| QKV | M=8222, N=9216, K=3072 | 8.03 ms | 58.0 TOP/s | -| Projection + gated residual | M=8222, N=3072, K=3072 | 3.15 ms | 49.2 TOP/s | -| Up/gate + SiTU-GLU | M=8222, N=16384, K=3072 | 12.93 ms | 64.0 TOP/s | -| Down + gated residual | M=8222, N=3072, K=8192 | 6.15 ms | 67.3 TOP/s | -| DiT exact attention | B=2, L=4111, H=24, D=128 | 17.41 ms | — | +| QKV | M=8222, N=9216, K=3072 | 7.52 ms | 61.9 TOP/s | +| Projection + gated residual | M=8222, N=3072, K=3072 | 2.98 ms | 52.1 TOP/s | +| Up/gate + SiTU-GLU | M=8222, N=16384, K=3072 | 12.11 ms | 68.4 TOP/s | +| Down + gated residual | M=8222, N=3072, K=8192 | 5.92 ms | 69.9 TOP/s | +| DiT exact attention | B=2, L=4111, H=24, D=128 | 14.57 ms | — | | Causal text attention | B=2, L=282, H=16, KV=4, D=256 | 0.259 ms | — | - -The INT8 tile is 64x128x32, with 32x64 warps and the CUTLASS Volta shared-memory swizzle/register pipeline. -The main loop's shared storage is reused by the fused output stage. The SiTU factors use the same bf16 lookup-table -values as the Turing path. No global dequantization buffer or expanded resident weights are allocated. Attention -uses 128 query rows for D=128 and 64 for D=256, 32 keys per tile; a larger key tile was slower on this device. +| VAE FP32 attention | B=1, L=16384, D=384 | 34.75 ms | — | +| VAE FP32 convolution | H=W=512, C=K=384, R=3 | 46.33 ms | — | + +The INT8 tile is 128x128x32 for M >= 1024 and 64x128x32 for smaller inputs, with 32x64 warps and the CUTLASS +Volta shared-memory swizzle/register pipeline. Main-loop storage is reused by the fused output stage. The SiTU +factors use the same bf16 lookup-table values as the Turing path. No global dequantization buffer or expanded +resident weights are allocated. D=128 attention uses 192 query rows for sequences >= 1024 and 128 for shorter +sequences; D=256 keeps 64 rows. All use 32 keys per tile. The larger DiT tiles reduce redundant loads; forcing +lower register limits, smaller warp tiles, or larger attention key tiles was slower or exceeded launch resources. + +VAE convolution uses a two-stage FP32 SIMT implicit GEMM, with 128x128x8 tiles (128x32x8 for narrow outputs), and +shares the existing bias-then-residual epilogue with the sm_80 implementation. It needs no scratch or im2col tensor. +VAE mid-attention reuses dead GroupNorm scratch for score slabs: up to 2048 query rows, limited to 64 MiB of scores +at supported image sizes, plus 8 MiB for cuBLASLt and 256 bytes for alignment. Two strict FP32 SGEMMs and a row +softmax replace the old per-query loop. It uses the caller's handle, stream and arena; no hidden device allocation, +weight conversion, synchronization or persistent plan cache is introduced. The standalone no-workspace attention +API keeps its original compatibility kernel. Default-off and Turing dispatch retain their existing implementations. End to end: MOSS-Image-SFT, Comfy-Org INT8 ConvRot DiT and Qwen3.5 9B W4A8 text encoder, `high`, exact attention, Euler, CFG 3, seed 42, prompt `a red apple on a wooden table, natural light, detailed photograph`: -| Run | Measured result | -|---|---| -| 1024x1024, 20 steps, `max_speed`, embedding cache hit | 45.35 s total; 41.35 s denoise (~2.07 s/step), 3.29 s decode; 6.74 GiB device peak | -| 1024x1024, 3 steps, `max_speed`, embedding cache hit | 10.24 s total; 6.16 s denoise (~2.05 s/step), 3.28 s decode; 6.74 GiB device peak | -| 256x256, 3 steps, `min_memory`, uncached blue-teapot prompt | 2.91 s total; 1.17 s encode, 1.13 s denoise, 0.13 s decode; 1.04 GiB device peak; all 42 DiT blocks streamed | - -The initial correctness implementation used a straightforward WMMA INT8 loop and took ~13.9 s/step at 1024x1024. -The optimized kernels reduce that by ~6.8x; the three-step PNGs before and after optimization are byte-identical. -This is a comparison within this Volta implementation, not a benchmark against another engine or an upstream V100 -backend. The upstream build did not run on V100. The 20-step timing includes weight loading and decode but uses a -cached text embedding; the streamed run separately exercises cold text encoding. - -Validation on V100: `sm70_check all` and `big`, all six `Sm70Build`/`GpuSm70`/`GpuSm75Convert` tests, Compute Sanitizer -memcheck (zero errors), and racecheck on attention and all INT8 epilogues including multiple accumulation segments -(zero errors/warnings). The default-off probe build passes its nine CPU tests. CTest reports no failures across 517 selected CPU cases +The baseline is the first SM70 implementation, commit `76b411a`; both runs use the same local weights and cached +text embedding, with default clocks on GPU 0. Totals include model loading and decode. Each generation uses one GPU. + +| Run / stage | Baseline | Optimized | Time reduction | +|---|---:|---:|---:| +| 1024x1024, 20 steps, `max_speed`, total | 45.35 s | 38.93 s | 14.2% | +| 20-step denoise | 41.35 s | 37.68 s | 8.9% | +| 1024x1024, 3 steps, `max_speed`, total | 10.24 s | 6.90 s | 32.6% | +| 3-step denoise | 6.16 s | 5.63 s | 8.6% | +| 1024x1024 decode | 3.28 s | 0.53 s | 83.8% | +| Device peak | 6.74 GiB | 6.74 GiB | unchanged | + +A second three-step generation in the same process takes 6.18 s (5.63 s denoise, 0.53 s decode). A separate +256x256 three-step `min_memory` run with an uncached green-cup prompt takes 2.14 s: 0.70 s text encoding, +1.14 s denoise and 0.05 s decode, at a 1.04 GiB device peak; all 42 DiT blocks stream from disk. + +Nsight Systems identified the main avoidable costs in the baseline: 1.47 s in the VAE attention kernel and +1.77 s across its convolutions. The new attention takes ~35 ms; individual representative convolutions are +4–4.6x faster (8.2x for the measured 1x1 shortcut). DiT attention improves from 17.41 to 14.57 ms, and the four +INT8 shapes improve by about 4–7%. CPU launch overhead was small; a graph/scheduler rewrite was not warranted +by this profile. These are measurements against the previous Volta implementation, not another inference engine. + +The INT8 sums and quantization epilogues remain exact under their existing contract. The VAE keeps FP32 operands +and accumulation; its changed reduction order is not bit-identical. Against the baseline PNGs, the three-step image +has 270 of 3,145,728 color components differing by one 8-bit level (PSNR 88.79 dB); the 20-step image has 240 such +components (89.31 dB). Neither comparison has a difference larger than one level. This checks the measured prompt +and seeds, not all possible images. + +Validation on V100: `sm70_check all` and `big`, all eight `Sm70Build`/`GpuSm70`/`GpuSm75Convert` tests, Compute Sanitizer +memcheck (zero errors), and racecheck on VAE, DiT attention and all INT8 epilogues including multiple accumulation +segments (zero errors/warnings). The default-off build also compiles the changed VAE translation units and passes +its nine CPU probe tests. CTest reports no failures across 517 selected CPU cases (with the normal missing-fixture skips), excluding the pre-existing `cli_check_verdict` CMake-script issue on CMake 3.31: that script drops empty list elements under unset CMP0007. Running it through a wrapper that sets CMP0007 to NEW also passes. The shared-library `moss-server` and CLI also produce identical three-step 1024x1024 PNGs. The implementation was hardware-tested on V100; other devices were not available for regression. diff --git a/src/kernels/sm70/attention_sm70.cu b/src/kernels/sm70/attention_sm70.cu index 9980d8d..b61cb65 100644 --- a/src/kernels/sm70/attention_sm70.cu +++ b/src/kernels/sm70/attention_sm70.cu @@ -35,9 +35,9 @@ struct Params { bool causal = false; }; -template -__global__ void attention_kernel(const Params p) { - constexpr int kBM = D == 128 ? 128 : 64, kThreads = kBM * 2; +template +__global__ __launch_bounds__(kBM * 2) void attention_kernel(const Params p) { + constexpr int kThreads = kBM * 2; constexpr int kQStride = D + 8, kVStride = BN + 8; using Acc = wmma::fragment; using A = wmma::fragment; @@ -173,12 +173,12 @@ __global__ void attention_kernel(const Params p) { } } -template +template void launch(const Params& p, int max_seqlen, int batch, cudaStream_t stream) { - constexpr int kBM = D == 128 ? 128 : 64, kThreads = kBM * 2; + constexpr int kThreads = kBM * 2; constexpr int bytes = ((kBM + BN) * (D + 8) + (D + kBM) * (BN + 8)) * sizeof(__half); static_assert(bytes <= 96 * 1024); - auto kernel = &attention_kernel; + auto kernel = &attention_kernel; if constexpr (bytes > 48 * 1024) MI_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes)); const dim3 grid(static_cast((max_seqlen + kBM - 1) / kBM), static_cast(batch), static_cast(p.heads)); @@ -232,9 +232,10 @@ void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& const int batch = static_cast(cu_seqlens_q.dim(0) - 1); if (head_dim == 128) { - launch<128, 32>(p, max_seqlen, batch, stream); + if (max_seqlen >= 1024) launch<128, 32, 192>(p, max_seqlen, batch, stream); + else launch<128, 32, 128>(p, max_seqlen, batch, stream); } else { - launch<256, 32>(p, max_seqlen, batch, stream); + launch<256, 32, 64>(p, max_seqlen, batch, stream); } } diff --git a/src/kernels/sm70/int8_gemm_sm70.cu b/src/kernels/sm70/int8_gemm_sm70.cu index 2530666..e370d6b 100644 --- a/src/kernels/sm70/int8_gemm_sm70.cu +++ b/src/kernels/sm70/int8_gemm_sm70.cu @@ -72,7 +72,8 @@ qd::SituFactorTables situ_tables(cudaStream_t stream) { // The main loop uses CUTLASS's Volta shared-memory swizzle and pipelined register loads. Global iterators still // read INT8; MmaPipelined's NumericArrayConverter stages exact fp16 integers. Fold every 1024 products into int32: // each fp32 partial is bounded by 2^24, preserving integer sums even with cancellation at the maximum supported K. -using Sm70Tile = qd::TileCfg<64, 128, 32, 32, 64, 2>; +using Sm70SmallTile = qd::TileCfg<64, 128, 32, 32, 64, 2>; +using Sm70LargeTile = qd::TileCfg<128, 128, 32, 32, 64, 2>; using Sm70Instr = cutlass::gemm::GemmShape<8, 8, 4>; template @@ -115,10 +116,10 @@ using qd::mul; using qd::rbf; using qd::dequant; -template +template struct Kernel { using Core = cutlass::gemm::threadblock::DefaultMmaCore< - Sm70Tile::Threadblock, Sm70Tile::Warp, Sm70Instr, cutlass::half_t, qd::RowMajor, + typename Tile::Threadblock, typename Tile::Warp, Sm70Instr, cutlass::half_t, qd::RowMajor, cutlass::half_t, qd::ColumnMajor, float, qd::RowMajor, cutlass::arch::OpClassTensorOp, 2>; using IteratorA = cutlass::transform::threadblock::PredicatedTileIterator< cutlass::MatrixShape, std::int8_t, qd::RowMajor, 1, @@ -130,20 +131,22 @@ struct Kernel { using Mma = ExactMma>; + union SharedStorage { + typename Mma::SharedStorage mainloop; + float result[Tile::Threadblock::kM * Tile::Threadblock::kN]; + }; }; -template -__global__ void int8_kernel(quant::Int8Operands p, typename Kernel::IteratorA::Params pa, - typename Kernel::IteratorB::Params pb, const __nv_bfloat16* bias, +template +__global__ void int8_kernel(quant::Int8Operands p, typename Kernel::IteratorA::Params pa, + typename Kernel::IteratorB::Params pb, const __nv_bfloat16* bias, const __nv_bfloat16* gate, __nv_bfloat16* out, std::int64_t ldo, qd::SituFactorTables tables) { - using K = Kernel; - constexpr int BM = Sm70Tile::Threadblock::kM, BN = Sm70Tile::Threadblock::kN; - constexpr int kThreads = (BM / Sm70Tile::Warp::kM) * (BN / Sm70Tile::Warp::kN) * 32; - __shared__ union Storage { - typename K::Mma::SharedStorage mainloop; - float result[BM * BN]; - } storage; + using K = Kernel; + constexpr int BM = Tile::Threadblock::kM, BN = Tile::Threadblock::kN; + constexpr int kThreads = (BM / Tile::Warp::kM) * (BN / Tile::Warp::kN) * 32; + extern __shared__ __align__(16) char buffer[]; + auto& storage = *reinterpret_cast(buffer); const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32; const int m0 = blockIdx.x * BM, n0 = blockIdx.y * BN; typename K::IteratorA a(pa, const_cast(p.x_q), {int(p.m), int(p.k)}, tid, {m0, 0}); @@ -151,11 +154,11 @@ __global__ void int8_kernel(quant::Int8Operands p, typename Kernel::IteratorA typename K::Mma mma(storage.mainloop, tid, warp, lane); typename K::Mma::FragmentC accum, zero; zero.clear(); - mma((p.k + Sm70Tile::Threadblock::kK - 1) / Sm70Tile::Threadblock::kK, accum, a, b, zero); + mma((p.k + Tile::Threadblock::kK - 1) / Tile::Threadblock::kK, accum, a, b, zero); __syncthreads(); using IteratorC = typename K::Core::MmaTensorOp::IteratorC; IteratorC c({storage.result, qd::RowMajor(BN)}, lane); - c.add_tile_offset({warp % (BM / Sm70Tile::Warp::kM), warp / (BM / Sm70Tile::Warp::kM)}); + c.add_tile_offset({warp % (BM / Tile::Warp::kM), warp / (BM / Tile::Warp::kM)}); c.store(accum); __syncthreads(); const float* result = storage.result; @@ -183,20 +186,31 @@ __global__ void int8_kernel(quant::Int8Operands p, typename Kernel::IteratorA } } -template -void launch(const quant::Int8Operands& p, const __nv_bfloat16* bias, const __nv_bfloat16* gate, +template +void launch_tile(const quant::Int8Operands& p, const __nv_bfloat16* bias, const __nv_bfloat16* gate, __nv_bfloat16* out, std::int64_t ldo, cudaStream_t stream) { - constexpr int kBM = Sm70Tile::Threadblock::kM, kBN = Sm70Tile::Threadblock::kN; - constexpr int kThreads = (kBM / Sm70Tile::Warp::kM) * (kBN / Sm70Tile::Warp::kN) * 32; + constexpr int kBM = Tile::Threadblock::kM, kBN = Tile::Threadblock::kN; + constexpr int kThreads = (kBM / Tile::Warp::kM) * (kBN / Tile::Warp::kN) * 32; const dim3 grid(static_cast((p.m + kBM - 1) / kBM), static_cast((p.n + kBN - 1) / kBN)); - const typename Kernel::IteratorA::Params pa{qd::RowMajor(p.ldx)}; - const typename Kernel::IteratorB::Params pb{qd::ColumnMajor(p.k)}; + const typename Kernel::IteratorA::Params pa{qd::RowMajor(p.ldx)}; + const typename Kernel::IteratorB::Params pb{qd::ColumnMajor(p.k)}; qd::SituFactorTables tables{}; if constexpr (E == Epilogue::kSitu) tables = situ_tables(stream); - int8_kernel<<>>(p, pa, pb, bias, gate, out, ldo, tables); + constexpr int bytes = sizeof(typename Kernel::SharedStorage); + auto kernel = &int8_kernel; + if constexpr (bytes > 48 * 1024) MI_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes)); + kernel<<>>(p, pa, pb, bias, gate, out, ldo, tables); MI_CUDA_CHECK_LAUNCH(); } +template +void launch(const quant::Int8Operands& p, const __nv_bfloat16* bias, const __nv_bfloat16* gate, + __nv_bfloat16* out, std::int64_t ldo, cudaStream_t stream) { + // More rows reuse each weight tile during DiT prefill; short text/boundary shapes retain the smaller tile. + if (p.m >= 1024) launch_tile(p, bias, gate, out, ldo, stream); + else launch_tile(p, bias, gate, out, ldo, stream); +} + } // namespace void int8_gemm(const quant::Int8Operands& op, const __nv_bfloat16* bias, __nv_bfloat16* out, std::int64_t ldo, diff --git a/src/kernels/sm70/vae_attention_sm70.cu b/src/kernels/sm70/vae_attention_sm70.cu new file mode 100644 index 0000000..ce5da10 --- /dev/null +++ b/src/kernels/sm70/vae_attention_sm70.cu @@ -0,0 +1,154 @@ +// The VAE's single-head attention uses fp32 throughout. A bounded score slab lets cuBLASLt reuse Q/K/V tiles, +// replacing the compatibility kernel's one-dot-product-per-warp loop without narrowing its fp32 operands. +#include "kernels/sm70/vae_sm70.h" + +#include +#include +#include +#include + +#include "core/cuda_check.h" + +namespace mi::kernels::sm70 { +namespace { + +constexpr std::size_t kScoresBytes = std::size_t{64} << 20; +constexpr std::size_t kLtBytes = std::size_t{8} << 20; +constexpr int kThreads = 256; + +int slab_rows(std::int64_t length) { + if (length <= 0) return 0; + return static_cast(std::min(length, std::max(1, + std::min(2048, kScoresBytes / (static_cast(length) * sizeof(float)))))); +} + +std::size_t score_bytes(std::int64_t length) { + const auto bytes = static_cast(slab_rows(length)) * static_cast(length) * sizeof(float); + return (bytes + 255) / 256 * 256; +} + +__global__ void softmax_kernel(float* scores, int columns) { + __shared__ float partial[8]; + float* row = scores + static_cast(blockIdx.x) * columns; + const int tid = threadIdx.x, lane = tid % 32, warp = tid / 32; + float value = -INFINITY; + for (int c = tid; c < columns; c += kThreads) value = fmaxf(value, row[c]); +#pragma unroll + for (int d = 16; d > 0; d >>= 1) value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, d)); + if (lane == 0) partial[warp] = value; + __syncthreads(); + value = lane < 8 ? partial[lane] : -INFINITY; +#pragma unroll + for (int d = 16; d > 0; d >>= 1) value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, d)); + const float maximum = value; + __syncthreads(); + float sum = 0.f; + for (int c = tid; c < columns; c += kThreads) { + const float p = expf(row[c] - maximum); + row[c] = p; + sum += p; + } +#pragma unroll + for (int d = 16; d > 0; d >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, d); + if (lane == 0) partial[warp] = sum; + __syncthreads(); + sum = lane < 8 ? partial[lane] : 0.f; +#pragma unroll + for (int d = 16; d > 0; d >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, d); + const float inv = 1.f / sum; + for (int c = tid; c < columns; c += kThreads) row[c] *= inv; +} + +void check_lt(cublasStatus_t status, const char* what) { + if (status != CUBLAS_STATUS_SUCCESS) { + throw std::runtime_error(std::string("sm70::vae_attention_fp32: cuBLASLt ") + what + " failed: " + + cublasLtGetStatusString(status)); + } +} + +// Column-major C[n,m] = op(A) * B. Plans live only for this call, which also avoids a device/stream cache or lock. +struct Matmul { + cublasLtMatmulDesc_t desc = nullptr; + cublasLtMatrixLayout_t a = nullptr, b = nullptr, c = nullptr; + cublasLtMatmulPreference_t pref = nullptr; + cublasLtMatmulAlgo_t algo{}; + + Matmul() = default; + Matmul(const Matmul&) = delete; + Matmul& operator=(const Matmul&) = delete; + ~Matmul() { + if (pref) cublasLtMatmulPreferenceDestroy(pref); + if (c) cublasLtMatrixLayoutDestroy(c); + if (b) cublasLtMatrixLayoutDestroy(b); + if (a) cublasLtMatrixLayoutDestroy(a); + if (desc) cublasLtMatmulDescDestroy(desc); + } + + void init(cublasLtHandle_t lt, int m, int n, int k, std::int64_t lda, std::int64_t ldb, std::int64_t ldc, + bool transpose_a) { + check_lt(cublasLtMatmulDescCreate(&desc, CUBLAS_COMPUTE_32F, CUDA_R_32F), "descriptor"); + const cublasOperation_t op = transpose_a ? CUBLAS_OP_T : CUBLAS_OP_N; + check_lt(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &op, sizeof(op)), "transpose"); + check_lt(cublasLtMatrixLayoutCreate(&a, CUDA_R_32F, transpose_a ? k : n, transpose_a ? n : k, lda), "A layout"); + check_lt(cublasLtMatrixLayoutCreate(&b, CUDA_R_32F, k, m, ldb), "B layout"); + check_lt(cublasLtMatrixLayoutCreate(&c, CUDA_R_32F, n, m, ldc), "C layout"); + check_lt(cublasLtMatmulPreferenceCreate(&pref), "preference"); + check_lt(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &kLtBytes, + sizeof(kLtBytes)), "workspace"); + cublasLtMatmulHeuristicResult_t result{}; + int found = 0; + check_lt(cublasLtMatmulAlgoGetHeuristic(lt, desc, a, b, c, c, pref, 1, &result, &found), "heuristic"); + if (found == 0) throw std::runtime_error("sm70::vae_attention_fp32: no fp32 GEMM algorithm"); + algo = result.algo; + } + + void run(cublasLtHandle_t lt, const float* ap, const float* bp, float* cp, float alpha, void* workspace, + cudaStream_t stream) const { + const float beta = 0.f; + check_lt(cublasLtMatmul(lt, desc, &alpha, ap, a, bp, b, &beta, cp, c, cp, c, &algo, workspace, kLtBytes, stream), + "matmul"); + } +}; + +} // namespace + +std::size_t vae_attention_workspace_bytes(std::int64_t length) { + // mid_attention accepts a 16-byte-aligned base; reserve room to align the cuBLASLt scratch to 256 bytes. + return length > 0 ? score_bytes(length) + kLtBytes + 256 : 0; +} + +void vae_attention_fp32(cublasLtHandle_t lt, const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& out, + float scale, void* workspace, std::size_t workspace_bytes, cudaStream_t stream) { + const int length = static_cast(q.dim(1)), dim = static_cast(q.dim(2)); + if (length == 0 || q.dim(0) == 0) return; + if (workspace == nullptr || workspace_bytes < vae_attention_workspace_bytes(length) || + reinterpret_cast(workspace) % 16 != 0) { + throw std::invalid_argument("sm70::vae_attention_fp32: workspace too small or not 16-byte aligned"); + } + const auto padding = (256 - reinterpret_cast(workspace) % 256) % 256; + auto* aligned = static_cast(workspace) + padding; + auto* scores = reinterpret_cast(aligned); + void* scratch = aligned + score_bytes(length); + const int chunk = slab_rows(length), tail = length % chunk; + Matmul qk, pv, tail_qk, tail_pv; + qk.init(lt, chunk, length, dim, k.stride[1], q.stride[1], length, true); + pv.init(lt, chunk, dim, length, v.stride[1], length, out.stride[1], false); + if (tail != 0) { + tail_qk.init(lt, tail, length, dim, k.stride[1], q.stride[1], length, true); + tail_pv.init(lt, tail, dim, length, v.stride[1], length, out.stride[1], false); + } + for (std::int64_t batch = 0; batch < q.dim(0); ++batch) { + for (int start = 0; start < length; start += chunk) { + const int rows = std::min(chunk, length - start); + const bool full = rows == chunk; + (full ? qk : tail_qk).run(lt, k.ptr() + batch * k.stride[0], + q.ptr() + batch * q.stride[0] + start * q.stride[1], scores, scale, scratch, stream); + softmax_kernel<<>>(scores, length); + MI_CUDA_CHECK_LAUNCH(); + (full ? pv : tail_pv).run(lt, v.ptr() + batch * v.stride[0], scores, + out.ptr() + batch * out.stride[0] + start * out.stride[1], 1.f, scratch, stream); + } + } +} + +} // namespace mi::kernels::sm70 diff --git a/src/kernels/sm70/vae_conv2d_sm70.cu b/src/kernels/sm70/vae_conv2d_sm70.cu new file mode 100644 index 0000000..722ef4b --- /dev/null +++ b/src/kernels/sm70/vae_conv2d_sm70.cu @@ -0,0 +1,70 @@ +// Volta's reference VAE stays fp32. CUTLASS SIMT implicit GEMM reuses larger spatial/channel tiles than the +// compatibility stencil, retaining fp32 operands, accumulation and the shared bias-then-residual epilogue. +#include "kernels/sm70/vae_sm70.h" + +#include +#include + +#include "core/cuda_check.h" +#include "kernels/vae/conv2d_epilogue.cuh" + +#include "cutlass/conv/conv2d_problem_size.h" +#include "cutlass/conv/device/implicit_gemm_convolution.h" +#include "cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h" +#include "cutlass/gemm/threadblock/threadblock_swizzle.h" + +namespace mi::kernels::sm70 { +namespace { + +using cutlass::gemm::GemmShape; +using Layout = cutlass::layout::TensorNHWC; + +template +struct ConvConfig { + using Kernel = typename cutlass::conv::kernel::DefaultConv2dFpropWithBroadcast< + float, Layout, float, Layout, float, Layout, float, cutlass::arch::OpClassSimt, cutlass::arch::Sm70, + GemmShape<128, N, 8>, GemmShape<32, WarpN, 8>, GemmShape<1, 1, 1>, vae::detail::BiasThenResidual, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<4>, 2, cutlass::arch::OpMultiplyAdd, + cutlass::conv::IteratorAlgorithm::kOptimized, cutlass::conv::StrideSupport::kUnity, 4, 4>::Kernel; + using Op = cutlass::conv::device::ImplicitGemmConvolution; +}; + +void check_cutlass(cutlass::Status status, const char* what) { + if (status != cutlass::Status::kSuccess) + throw std::runtime_error(std::string("sm70::vae_conv2d: CUTLASS ") + what + " failed: " + cutlassGetStatusString(status)); +} + +template +void run(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, + cudaStream_t stream) { + const int n = static_cast(x.dim(0)), h = static_cast(x.dim(1)), width = static_cast(x.dim(2)); + const int c = static_cast(x.dim(3)), k = static_cast(w.dim(0)), r = static_cast(w.dim(1)); + const int pad = (r - 1) / 2; + const cutlass::conv::Conv2dProblemSize ps(n, h, width, c, k, r, r, h, width, pad, pad, 1, 1, 1, 1, + cutlass::conv::Mode::kCrossCorrelation, 1, 1); + const cutlass::TensorRef rx(x.ptr(), Layout::packed({n, h, width, c})); + const cutlass::TensorRef rw(w.ptr(), Layout::packed({k, r, r, c})); + const cutlass::TensorRef ry(y.ptr(), Layout::packed({n, h, width, k})); + const cutlass::TensorRef rr(residual ? residual->ptr() : y.ptr(), ry.layout()); + const typename Cfg::Op::EpilogueOutputOp::Params epi(1.f, residual ? 1.f : 0.f); + const typename Cfg::Op::Arguments args(ps, rx, rw, rr, ry, epi, cutlass::conv::SplitKMode::kSerial, + bias.data, nullptr, 0, 0); + check_cutlass(Cfg::Op::can_implement(args), "can_implement"); + // Serial split-K=1 uses no semaphores. Keep the compatibility API's zero-workspace contract explicit. + if (Cfg::Op::get_workspace_size(args) != 0) throw std::runtime_error("sm70::vae_conv2d: unexpected workspace"); + typename Cfg::Op op; + check_cutlass(op.initialize(args, nullptr, stream), "initialize"); + check_cutlass(op.run(stream), "run"); + MI_CUDA_CHECK_LAUNCH(); +} + +} // namespace + +void vae_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, + cudaStream_t stream) { + if (x.numel() == 0 || y.numel() == 0) return; + if (w.dim(0) >= 64) run>(x, w, bias, residual, y, stream); + else run>(x, w, bias, residual, y, stream); +} + +} // namespace mi::kernels::sm70 diff --git a/src/kernels/sm70/vae_sm70.h b/src/kernels/sm70/vae_sm70.h new file mode 100644 index 0000000..e3c5d31 --- /dev/null +++ b/src/kernels/sm70/vae_sm70.h @@ -0,0 +1,18 @@ +// Volta VAE: fp32 SIMT convolutions and bounded fp32 attention slabs; no hidden device allocation. +#pragma once + +#include + +#include "core/tensor.h" + +namespace mi::kernels::sm70 { + +// Arguments are validated by vae::conv2d. The split-K=1 implicit GEMM needs no workspace. +void vae_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, + cudaStream_t stream); + +std::size_t vae_attention_workspace_bytes(std::int64_t length); +void vae_attention_fp32(cublasLtHandle_t lt, const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& out, + float scale, void* workspace, std::size_t workspace_bytes, cudaStream_t stream); + +} // namespace mi::kernels::sm70 diff --git a/src/kernels/vae/attention.cu b/src/kernels/vae/attention.cu index 86491f0..4c6ffac 100644 --- a/src/kernels/vae/attention.cu +++ b/src/kernels/vae/attention.cu @@ -23,6 +23,8 @@ #include "core/cuda_check.h" #include "kernels/sm75/sm75.h" +#include "kernels/sm70/sm70.h" +#include "kernels/sm70/vae_sm70.h" #include "kernels/vae/attention.h" namespace mi::vae { @@ -450,6 +452,11 @@ MidLayout mid_layout(const Tensor& x) { l.gn = l.qkv + align_up(tokens * 3 * d * sizeof(float)); l.lt = l.gn + align_up(group_norm_workspace_bytes(x, kGroupNormGroups)); l.total = l.lt + kLtWorkspace; +#if MOSSINFER_HAVE_SM70 + if (kernels::sm70::volta()) { + l.total = std::max(l.total, l.gn + kernels::sm70::vae_attention_workspace_bytes(x.dim(1) * x.dim(2))); + } +#endif return l; } @@ -541,7 +548,15 @@ void mid_attention(cublasLtHandle_t lt, const Tensor& x, const MidAttentionWeigh // SDPA's default scale is 1/sqrt(E) formed in double and rounded to fp32. This fp32 expression gives the same float // for E = 512 (0x3d3504f3) but 1 ulp less for E = 384 (0x3d5105eb vs 0x3d5105ec); kept, the small decoder's outputs // are pinned to it (the attention is fp32-accurate, not bit-identical to SDPA, either way). - attention_fp32(qkv_view(0), qkv_view(1), qkv_view(2), o, 1.0f / std::sqrt(static_cast(dim)), precision, stream); +#if MOSSINFER_HAVE_SM70 + if (kernels::sm70::volta()) { + kernels::sm70::vae_attention_fp32(lt, qkv_view(0), qkv_view(1), qkv_view(2), o, + 1.0f / std::sqrt(static_cast(dim)), ws + l.gn, l.total - l.gn, stream); + } else +#endif + { + attention_fp32(qkv_view(0), qkv_view(1), qkv_view(2), o, 1.0f / std::sqrt(static_cast(dim)), precision, stream); + } if (!residual) { linear(lt, xn, dim, w.wo.ptr(), w.bo.ptr(), nullptr, 0, y.ptr(), dim, tokens, dim, dim, fast_linears, ws + l.lt, kLtWorkspace, stream); diff --git a/src/kernels/vae/attention.h b/src/kernels/vae/attention.h index 48feca1..6272335 100644 --- a/src/kernels/vae/attention.h +++ b/src/kernels/vae/attention.h @@ -1,6 +1,7 @@ // VAE mid-block attention (fp32, single head) for both FLUX.2 widths: head dim D = 384 (the small decoder) and -// D = 512 (the full decoder and the encoder). One kernel, templated on D (attention.cu); the D = 384 instance is the -// original kernel bit for bit. +// D = 512 (the full decoder and the encoder). The sm_80+ kernel is templated on D (attention.cu); its D = 384 +// instance is the original kernel bit for bit. Volta mid_attention uses bounded FP32 score slabs in caller scratch; +// the standalone no-workspace API keeps the compatibility kernel. // // The types and functions are declared in kernels/vae/vae_kernels.h together with the other VAE ops (their // comments there predate D = 512); this header is the attention API's home and states the head-dim contract: diff --git a/src/kernels/vae/conv2d.cu b/src/kernels/vae/conv2d.cu index e45ba25..c57b2f3 100644 --- a/src/kernels/vae/conv2d.cu +++ b/src/kernels/vae/conv2d.cu @@ -9,13 +9,14 @@ #include "core/cuda_check.h" #include "kernels/sm75/sm75.h" +#include "kernels/sm70/sm70.h" +#include "kernels/sm70/vae_sm70.h" +#include "kernels/vae/conv2d_epilogue.cuh" #include "kernels/vae/vae_kernels.h" #include "cutlass/conv/conv2d_problem_size.h" #include "cutlass/conv/device/implicit_gemm_convolution.h" #include "cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h" -#include "cutlass/epilogue/thread/activation.h" -#include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h" #include "cutlass/gemm/threadblock/threadblock_swizzle.h" namespace mi::vae { @@ -23,42 +24,7 @@ namespace { using cutlass::gemm::GemmShape; -// Z = (acc + bias) + residual, the reference's rounding order (cuDNN adds the bias, then ResnetBlock2D adds -// input_tensor). CUTLASS's LinearCombinationBiasElementwise evaluates (acc + residual) + bias; this subclass -// only replaces the with-source operator. beta = 1 with a residual, else 0 (the epilogue then skips reading the -// source tensor and the base's no-source path computes acc + bias). Bias stays fp32 in every mode. The alpha/beta -// copies here ignore the base's set_k_partition, so this epilogue requires split-K = 1 (problem() fixes that). -template -class BiasThenResidual - : public cutlass::epilogue::thread::LinearCombinationBiasElementwise< - T, float, float, T, T, kAlignOut, cutlass::epilogue::thread::Identity, cutlass::plus, - /*StoreT=*/false, /*ElementVector=*/float> { - public: - using Base = cutlass::epilogue::thread::LinearCombinationBiasElementwise< - T, float, float, T, T, kAlignOut, cutlass::epilogue::thread::Identity, cutlass::plus, false, float>; - using typename Base::FragmentAccumulator; - using typename Base::FragmentC; - using typename Base::FragmentCompute; - using typename Base::FragmentT; - using typename Base::FragmentZ; - using typename Base::Params; - using Base::operator(); - - CUTLASS_HOST_DEVICE explicit BiasThenResidual(Params const& p) : Base(p), alpha_(p.alpha), beta_(p.beta) {} - - CUTLASS_HOST_DEVICE void operator()(FragmentZ& frag_z, FragmentT&, FragmentAccumulator const& ab, FragmentC const& frag_c, - FragmentCompute const& v) const { - const FragmentCompute acc = cutlass::NumericArrayConverter()(ab); - const FragmentCompute res = cutlass::NumericArrayConverter()(frag_c); - FragmentCompute z; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kAlignOut; ++i) z[i] = (alpha_ * acc[i] + v[i]) + beta_ * res[i]; - frag_z = cutlass::NumericArrayConverter()(z); - } - - private: - float alpha_, beta_; -}; +using detail::BiasThenResidual; template struct ConvConfig { @@ -194,6 +160,12 @@ void conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* if (y.data == x.data) throw std::invalid_argument("conv2d: y must not alias x"); if (kernels::sm75::fp16_weights()) { if (x.dtype != DType::kF32) throw std::runtime_error("conv2d: the sm_75 path runs the VAE in fp32 (reference mode)"); +#if MOSSINFER_HAVE_SM70 + if (kernels::sm70::volta()) { + kernels::sm70::vae_conv2d(x, w, bias, residual, y, stream); + return; + } +#endif kernels::sm75::vae_conv2d(x, w, bias, residual, y, stream); return; } diff --git a/src/kernels/vae/conv2d_epilogue.cuh b/src/kernels/vae/conv2d_epilogue.cuh new file mode 100644 index 0000000..b09d6c8 --- /dev/null +++ b/src/kernels/vae/conv2d_epilogue.cuh @@ -0,0 +1,46 @@ +// Shared VAE convolution epilogue; used by Tensor Core and fp32 SIMT implementations. +#pragma once + +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h" + +namespace mi::vae::detail { + +// Z = (acc + bias) + residual, the reference's rounding order (cuDNN adds the bias, then ResnetBlock2D adds +// input_tensor). CUTLASS's LinearCombinationBiasElementwise evaluates (acc + residual) + bias; this subclass +// only replaces the with-source operator. beta = 1 with a residual, else 0 (the epilogue then skips reading the +// source tensor and the base's no-source path computes acc + bias). Bias stays fp32 in every mode. The alpha/beta +// copies here ignore the base's set_k_partition, so this epilogue requires split-K = 1 (both convolution callers fix that). +template +class BiasThenResidual + : public cutlass::epilogue::thread::LinearCombinationBiasElementwise< + T, float, float, T, T, kAlignOut, cutlass::epilogue::thread::Identity, cutlass::plus, + /*StoreT=*/false, /*ElementVector=*/float> { + public: + using Base = cutlass::epilogue::thread::LinearCombinationBiasElementwise< + T, float, float, T, T, kAlignOut, cutlass::epilogue::thread::Identity, cutlass::plus, false, float>; + using typename Base::FragmentAccumulator; + using typename Base::FragmentC; + using typename Base::FragmentCompute; + using typename Base::FragmentT; + using typename Base::FragmentZ; + using typename Base::Params; + using Base::operator(); + + CUTLASS_HOST_DEVICE explicit BiasThenResidual(Params const& p) : Base(p), alpha_(p.alpha), beta_(p.beta) {} + + CUTLASS_HOST_DEVICE void operator()(FragmentZ& frag_z, FragmentT&, FragmentAccumulator const& ab, FragmentC const& frag_c, + FragmentCompute const& v) const { + const FragmentCompute acc = cutlass::NumericArrayConverter()(ab); + const FragmentCompute res = cutlass::NumericArrayConverter()(frag_c); + FragmentCompute z; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kAlignOut; ++i) z[i] = (alpha_ * acc[i] + v[i]) + beta_ * res[i]; + frag_z = cutlass::NumericArrayConverter()(z); + } + + private: + float alpha_, beta_; +}; + +} // namespace mi::vae::detail diff --git a/src/kernels/vae/vae_kernels.h b/src/kernels/vae/vae_kernels.h index 79501d5..c19c5aa 100644 --- a/src/kernels/vae/vae_kernels.h +++ b/src/kernels/vae/vae_kernels.h @@ -59,6 +59,7 @@ inline constexpr float kBatchNormEps = 1e-4f; // vae.config.batch_norm_eps // may alias residual. Alignment: x, w 16 bytes; bias 32 bytes; y, residual 16 bytes (4 bytes when K < 64). // The fp32 path rounds operands to nearest TF32 and accumulates in the same order as cuDNN's TF32 kernels // (bit-identical to torch 2.14 / cuDNN 9.24 on the decoder's K >= 64 convs; conv_out differs by ~2e-6 rel). +// The sm_7x compatibility paths keep fp32 operands/accumulation; Volta uses a SIMT implicit GEMM with no workspace. std::size_t conv2d_workspace_bytes(const Tensor& x, const Tensor& w); void conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, void* workspace, std::size_t workspace_bytes, cudaStream_t stream); diff --git a/tests/unit/test_sm70.cpp b/tests/unit/test_sm70.cpp index de5874a..03176e0 100644 --- a/tests/unit/test_sm70.cpp +++ b/tests/unit/test_sm70.cpp @@ -4,7 +4,10 @@ #include #include +#include +#include #include +#include #include #include @@ -12,6 +15,9 @@ #include "core/cuda_check.h" #include "kernels/basic/elementwise.h" #include "kernels/sm70/sm70.h" +#include "kernels/sm70/vae_sm70.h" +#include "kernels/sm75/sm75.h" +#include "kernels/vae/vae_kernels.h" namespace { @@ -160,9 +166,138 @@ TEST(GpuSm70, Int8EpiloguesAndStrides) { check_gemm(3, 8, 16, false); check_gemm(67, 136, 1040, false); check_gemm(129, 256, 4096, false); + check_gemm(1025, 136, 1040, false); } TEST(GpuSm70, Int8ExactLargeCancellation) { if (!mi::kernels::sm70::volta()) GTEST_SKIP() << "needs MOSSINFER_SM70 and a Volta device"; check_gemm(3, 16, 131056, true); } + +#if MOSSINFER_HAVE_SM70 +TEST(GpuSm70, VaeConvolutionFp32AndResidualAlias) { + if (!mi::kernels::sm70::volta()) GTEST_SKIP() << "needs MOSSINFER_SM70 and a Volta device"; + // Spatial/channel tails, both decoder widths, 1x1/3x3, padding, batch and in-place residual epilogues. + for (const auto& shape : {std::array{2, 7, 9, 4, 3, 3}, std::array{1, 13, 11, 12, 65, 1}, + std::array{1, 9, 17, 128, 128, 3}, std::array{1, 5, 7, 384, 96, 3}, + std::array{1, 5, 7, 512, 256, 3}}) { + const auto [n, h, width, c, k, r] = shape; + SCOPED_TRACE(::testing::Message() << "C=" << c << " K=" << k << " R=" << r); + std::uint32_t rng = 42; + const auto values = [&](std::size_t count) { + std::vector v(count); + for (float& f : v) { + rng ^= rng << 13; + rng ^= rng >> 17; + rng ^= rng << 5; + f = static_cast(static_cast(rng % 2001) - 1000) / 4096.f; + } + return v; + }; + const auto hx = values(static_cast(n) * h * width * c), hw = values(static_cast(k) * r * r * c); + const auto hb = values(static_cast(k)), hr = values(static_cast(n) * h * width * k); + Buffer dx(hx), dw(hw), db(hb), out(hr); + const Tensor x = Tensor::view(dx.get(), DType::kF32, Device::kCUDA, {n, h, width, c}); + const Tensor w = Tensor::view(dw.get(), DType::kF32, Device::kCUDA, {k, r, r, c}); + const Tensor bias = Tensor::view(db.get(), DType::kF32, Device::kCUDA, {k}); + const Tensor y = Tensor::view(out.get(), DType::kF32, Device::kCUDA, {n, h, width, k}); + EXPECT_EQ(mi::vae::conv2d_workspace_bytes(x, w), 0u); + for (bool residual : {false, true}) { + MI_CUDA_CHECK(cudaMemcpy(out.get(), hr.data(), hr.size() * sizeof(float), cudaMemcpyHostToDevice)); + cudaStream_t stream = nullptr; + cudaGraph_t graph = nullptr; + cudaGraphExec_t exec = nullptr; + MI_CUDA_CHECK(cudaStreamCreate(&stream)); + MI_CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + mi::vae::conv2d(x, w, bias, residual ? &y : nullptr, y, nullptr, 0, stream); + MI_CUDA_CHECK(cudaStreamEndCapture(stream, &graph)); + MI_CUDA_CHECK(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); + MI_CUDA_CHECK(cudaGraphLaunch(exec, stream)); + MI_CUDA_CHECK(cudaStreamSynchronize(stream)); + cudaGraphExecDestroy(exec); + cudaGraphDestroy(graph); + cudaStreamDestroy(stream); + const auto actual = out.host(); + double max_error = 0.; + for (int batch = 0; batch < n; ++batch) for (int oh = 0; oh < h; ++oh) for (int ow = 0; ow < width; ++ow) { + for (int oc = 0; oc < k; ++oc) { + double sum = 0.; + for (int rh = 0; rh < r; ++rh) for (int rw = 0; rw < r; ++rw) { + const int ih = oh + rh - r / 2, iw = ow + rw - r / 2; + if (ih < 0 || ih >= h || iw < 0 || iw >= width) continue; + for (int ic = 0; ic < c; ++ic) + sum += static_cast(hx[static_cast(((batch * h + ih) * width + iw) * c + ic)]) * + hw[static_cast(((oc * r + rh) * r + rw) * c + ic)]; + } + const auto i = static_cast(((batch * h + oh) * width + ow) * k + oc); + const float expected = (static_cast(sum) + hb[static_cast(oc)]) + (residual ? hr[i] : 0.f); + ASSERT_TRUE(std::isfinite(actual[i])); + max_error = std::max(max_error, std::abs(static_cast(actual[i]) - expected)); + } + } + EXPECT_LE(max_error, 3e-5); + } + EXPECT_EQ(dx.host(), hx); + EXPECT_EQ(dw.host(), hw); + } +} + +TEST(GpuSm70, VaeAttentionFp32SlabsAndStrides) { + if (!mi::kernels::sm70::volta()) GTEST_SKIP() << "needs MOSSINFER_SM70 and a Volta device"; + struct Handle { + cublasLtHandle_t lt = nullptr; + ~Handle() { if (lt) cublasLtDestroy(lt); } + } handle; + ASSERT_EQ(cublasLtCreate(&handle.lt), CUBLAS_STATUS_SUCCESS); + for (const auto& [length, dim] : {std::pair{1, 384}, std::pair{67, 512}, std::pair{2053, 384}}) { + constexpr int batch = 2; + const int stride = 3 * dim + 4; + std::vector values(static_cast(batch) * length * stride); + for (std::size_t i = 0; i < values.size(); ++i) { + values[i] = std::sin(static_cast(i % 997) * 0.137f) * 0.7f; + } + Buffer qkv(values), out(std::vector(static_cast(batch) * length * dim, -123.f)); + Buffer reference(std::vector(static_cast(batch) * length * dim)); + const auto view = [&](int offset) { + Tensor t = Tensor::view(qkv.get() + offset, DType::kF32, Device::kCUDA, {batch, length, dim}); + t.stride[1] = stride; + t.stride[0] = static_cast(length) * stride; + return t; + }; + const Tensor q = view(0), k = view(dim), v = view(2 * dim); + const Tensor o = Tensor::view(out.get(), DType::kF32, Device::kCUDA, {batch, length, dim}); + const Tensor ref = Tensor::view(reference.get(), DType::kF32, Device::kCUDA, {batch, length, dim}); + const std::size_t bytes = mi::kernels::sm70::vae_attention_workspace_bytes(length); + Buffer allocation(std::vector(bytes / sizeof(float) + 4)); + float* scratch = allocation.get() + 4; // preserve mid_attention's 16-byte alignment contract + const float scale = 1.f / std::sqrt(static_cast(dim)); + mi::kernels::sm75::vae_attention_fp32(q, k, v, ref, scale, nullptr); + mi::kernels::sm70::vae_attention_fp32(handle.lt, q, k, v, o, scale, scratch, bytes, nullptr); + // The warmed cuBLASLt plans and caller-owned slab also work on a captured, non-default stream. + cudaStream_t stream = nullptr; + cudaGraph_t graph = nullptr; + cudaGraphExec_t exec = nullptr; + MI_CUDA_CHECK(cudaStreamCreate(&stream)); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + MI_CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + mi::kernels::sm70::vae_attention_fp32(handle.lt, q, k, v, o, scale, scratch, bytes, stream); + MI_CUDA_CHECK(cudaStreamEndCapture(stream, &graph)); + MI_CUDA_CHECK(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0)); + for (int i = 0; i < 2; ++i) MI_CUDA_CHECK(cudaGraphLaunch(exec, stream)); + MI_CUDA_CHECK(cudaStreamSynchronize(stream)); + cudaGraphExecDestroy(exec); + cudaGraphDestroy(graph); + cudaStreamDestroy(stream); + const auto actual = out.host(), expected = reference.host(); + float max_error = 0.f; + for (std::size_t i = 0; i < actual.size(); ++i) { + ASSERT_TRUE(std::isfinite(actual[i])); + max_error = std::max(max_error, std::fabs(actual[i] - expected[i])); + } + EXPECT_LE(max_error, 2e-5f) << "L=" << length << " D=" << dim; + EXPECT_EQ(qkv.host(), values); + EXPECT_THROW(mi::kernels::sm70::vae_attention_fp32(handle.lt, q, k, v, o, scale, scratch, bytes - 1, nullptr), + std::invalid_argument); + } +} +#endif diff --git a/tools/sm70_bench.cu b/tools/sm70_bench.cu index ca5b8a8..8955b4a 100644 --- a/tools/sm70_bench.cu +++ b/tools/sm70_bench.cu @@ -1,16 +1,20 @@ // V100 rate probe, not installed. Run sm70_check first; timings are CUDA events after warm-up, not host launch time. -// Usage: sm70_bench [GEMM rows (default 8222)] [attention sequence length (default 4111)]. +// Usage: sm70_bench [GEMM rows (default 8222)] [attention sequence length (default 4111)], or sm70_bench vae. #include #include #include #include #include +#include +#include #include #include "core/cuda_check.h" #include "kernels/attention/flash_attention.h" #include "kernels/sm70/sm70.h" +#include "kernels/sm70/vae_sm70.h" +#include "kernels/vae/vae_kernels.h" namespace { @@ -83,11 +87,48 @@ void attention(int length, int heads, int kv_heads, int dim, bool causal) { std::printf("attention L=%d H=%d KV=%d D=%d causal=%d: %.3f ms\n", length, heads, kv_heads, dim, causal, ms); } +void vae() { + using mi::DType; + using mi::Device; + using mi::Tensor; + for (int h : {128, 256, 512}) { + constexpr int channels = 384; + Buffer x(static_cast(h) * h * channels * sizeof(float)); + Buffer w(static_cast(channels) * 3 * 3 * channels * sizeof(float)), bias(channels * sizeof(float)); + Buffer out(static_cast(h) * h * channels * sizeof(float)); + const auto tx = Tensor::view(x.get(), DType::kF32, Device::kCUDA, {1, h, h, channels}); + const auto tw = Tensor::view(w.get(), DType::kF32, Device::kCUDA, {channels, 3, 3, channels}); + const auto tb = Tensor::view(bias.get(), DType::kF32, Device::kCUDA, {channels}); + const auto ty = Tensor::view(out.get(), DType::kF32, Device::kCUDA, {1, h, h, channels}); + const float ms = time_ms([&] { mi::vae::conv2d(tx, tw, tb, nullptr, ty, nullptr, 0, nullptr); }); + std::printf("vae fp32 conv H=W=%d C=K=%d R=3: %.3f ms\n", h, channels, ms); + } + constexpr int length = 16384, dim = 384; + Buffer qkv(static_cast(length) * dim * 3 * sizeof(float)); + Buffer out(static_cast(length) * dim * sizeof(float)); + const auto q = Tensor::view(qkv.get(), DType::kF32, Device::kCUDA, {1, length, dim}); + const auto k = Tensor::view(qkv.get() + length * dim, DType::kF32, Device::kCUDA, {1, length, dim}); + const auto v = Tensor::view(qkv.get() + 2 * length * dim, DType::kF32, Device::kCUDA, {1, length, dim}); + const auto o = Tensor::view(out.get(), DType::kF32, Device::kCUDA, {1, length, dim}); + const auto bytes = mi::kernels::sm70::vae_attention_workspace_bytes(length); + Buffer scratch(bytes); + cublasLtHandle_t lt = nullptr; + if (cublasLtCreate(<) != CUBLAS_STATUS_SUCCESS) throw std::runtime_error("cuBLASLt handle creation failed"); + const float ms = time_ms([&] { + mi::kernels::sm70::vae_attention_fp32(lt, q, k, v, o, 1.f / std::sqrt(static_cast(dim)), scratch.get(), + bytes, nullptr); + }); + cublasLtDestroy(lt); + std::printf("vae fp32 attention L=%d D=%d: %.3f ms, %.1f MiB scratch\n", length, dim, ms, + static_cast(bytes) / (1024 * 1024)); +} + } // namespace int main(int argc, char** argv) { std::setvbuf(stdout, nullptr, _IONBF, 0); if (!mi::kernels::sm70::volta()) return 1; + if (argc > 1 && std::string_view(argv[1]) == "vae") { vae(); return 0; } const int m = argc > 1 ? std::atoi(argv[1]) : 8222; const int length = argc > 2 ? std::atoi(argv[2]) : 4111; if (m <= 0 || length <= 0) return 2; diff --git a/tools/sm75_check.cu b/tools/sm75_check.cu index 41a4c48..262e1c9 100644 --- a/tools/sm75_check.cu +++ b/tools/sm75_check.cu @@ -29,6 +29,7 @@ #include "kernels/basic/gemm.h" #include "kernels/quant/quant.h" #include "kernels/sm75/sm75.h" +#include "kernels/vae/vae_kernels.h" namespace { @@ -362,7 +363,7 @@ int run_vae_conv_case(int N, int H, int W, int C, int K, int R, std::uint64_t se const Tensor tb = Tensor::view(db.get(), DType::kF32, Device::kCUDA, {K}); const Tensor tr = Tensor::view(dr.get(), DType::kF32, Device::kCUDA, {N, H, W, K}); const Tensor ty = Tensor::view(dy.get(), DType::kF32, Device::kCUDA, {N, H, W, K}); - mi::kernels::sm75::vae_conv2d(tx, tw, tb, &tr, ty, nullptr); + mi::vae::conv2d(tx, tw, tb, &tr, ty, nullptr, 0, nullptr); MI_CUDA_CHECK(cudaDeviceSynchronize()); const auto got = dy.download(rv.size()); const int pad = (R - 1) / 2; @@ -578,6 +579,7 @@ int main(int argc, char** argv) { if (what == "big") { // Engine-shaped cases: 24 heads at head_dim 128 (the DiT's joint attention) and the text encoder's 16/4 GQA at // head_dim 256. The host reference is O(L^2 D), so these stay at a few hundred tokens. + failures += run_attention("dit1025", 4, 2, 128, false, {1025, 193, 1}, 47); failures += run_attention("dit512", 24, 24, 128, false, {512, 512}, 41); failures += run_attention("te282", 16, 4, 256, true, {282}, 42); failures += run_attention("dit76", 24, 24, 128, false, {76, 200, 1}, 43); From 10d852c1b7afe0195a4510bc59db5c9767f7d985 Mon Sep 17 00:00:00 2001 From: xycld Date: Sun, 1 Nov 2026 20:44:29 -0500 Subject: [PATCH 3/3] sm70: docs keep a few results, the bench tool stays local docs/sm70.md drops the per-kernel table and the comparison with the PR's first commit, and keeps the tile and memory layout, end-to-end numbers and what the checks cover. tools/sm70_bench.cu and its CMake block are removed: bench tools are not tracked. sm75.cmake loses an inner option check that repeated its block's condition. integration.md states where Volta GPU runs happened and that the Linux -sm70 package carries the Turing kernels. --- cmake/modules/backend.cmake | 3 +- cmake/modules/sm70.cmake | 6 -- cmake/modules/sm75.cmake | 7 +- docs/integration.md | 5 +- docs/sm70.md | 113 ++++++++--------------------- tools/sm70_bench.cu | 141 ------------------------------------ 6 files changed, 39 insertions(+), 236 deletions(-) delete mode 100644 tools/sm70_bench.cu diff --git a/cmake/modules/backend.cmake b/cmake/modules/backend.cmake index f40421f..3ae5182 100644 --- a/cmake/modules/backend.cmake +++ b/cmake/modules/backend.cmake @@ -1,7 +1,8 @@ # src/backend: the probe reports which compute capabilities this build has kernel images for, so it is told the # CUDA architecture list of mossinfer_core (commas: a semicolon would split the definition). This file only. # The sm_75 kernels are their own OBJECT library (cmake/modules/sm75.cmake) outside the core's list, so a build with -# MOSSINFER_SM75 appends 75-real here: probe.h's minimum capability follows the same option. +# MOSSINFER_SM75 appends 75-real here: probe.h's minimum capability follows the same option. MOSSINFER_SM70 appends +# 70-real, whose SASS also counts for 7.5 (compatible_image in probe.cpp). # No-op when MOSSINFER_ONLY filtered the probe out. set(_mi_backend_probe ${CMAKE_SOURCE_DIR}/src/backend/probe.cpp) if(_mi_backend_probe IN_LIST MOSSINFER_SOURCES) diff --git a/cmake/modules/sm70.cmake b/cmake/modules/sm70.cmake index 610c21c..a7b5acd 100644 --- a/cmake/modules/sm70.cmake +++ b/cmake/modules/sm70.cmake @@ -33,12 +33,6 @@ if(_mi_sm70_sources) target_compile_options(sm70_check PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) target_link_libraries(sm70_check PRIVATE mossinfer_core) endif() - if(EXISTS ${CMAKE_SOURCE_DIR}/tools/sm70_bench.cu) - add_executable(sm70_bench ${CMAKE_SOURCE_DIR}/tools/sm70_bench.cu) - set_target_properties(sm70_bench PROPERTIES CUDA_ARCHITECTURES 70-real) - target_link_libraries(sm70_bench PRIVATE mossinfer_core) - target_compile_options(sm70_bench PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) - endif() message(STATUS "SM70: Volta kernel family built; compute capability 7.0 accepted") endif() endif() diff --git a/cmake/modules/sm75.cmake b/cmake/modules/sm75.cmake index e3e58cb..fbda063 100644 --- a/cmake/modules/sm75.cmake +++ b/cmake/modules/sm75.cmake @@ -10,6 +10,7 @@ # - the portable SIMT translation units are also compiled for sm_75, # - flash_attention.cu / gemm.cpp / chunk_rule.cu dispatch into the sm_75 kernels, # - the sm75_check bring-up tool is built. +# MOSSINFER_SM70 (cmake/modules/sm70.cmake) turns on the same paths, and also builds convert and vae_sm75 for sm_70. # Tolerates the sources being filtered out by MOSSINFER_ONLY. set(_mi_sm75_sources ${MOSSINFER_SOURCES}) list(FILTER _mi_sm75_sources INCLUDE REGEX ".*/src/kernels/sm75/[^/]*\\.cu$") @@ -111,14 +112,12 @@ if((MOSSINFER_SM75 OR MOSSINFER_SM70) AND _mi_sm75_present) # src/backend/probe.h: the minimum compute capability and the architecture list this build accepts. Defined on the # sources that read it rather than on the whole target, so a build without MOSSINFER_SM75 keeps its old objects and # adding the option only recompiles the dispatchers. - if(MOSSINFER_SM75 OR MOSSINFER_SM70) - set_property(SOURCE ${CMAKE_SOURCE_DIR}/src/backend/probe.cpp + set_property(SOURCE ${CMAKE_SOURCE_DIR}/src/backend/probe.cpp ${CMAKE_SOURCE_DIR}/src/kernels/attention/flash_attention.cu ${CMAKE_SOURCE_DIR}/src/kernels/basic/gemm.cpp ${CMAKE_SOURCE_DIR}/src/kernels/gdn/chunk_rule.cu APPEND PROPERTY COMPILE_DEFINITIONS MOSSINFER_HAVE_SM75=1) - message(STATUS "SM75: compute capability 7.5 accepted") - endif() + message(STATUS "SM75: compute capability 7.5 accepted") # Bring-up harness (not installed): the engine's GPU tests compile their reference kernels for the build's CUDA # architectures, so a Turing device cannot run them; this tool builds its own host reference. diff --git a/docs/integration.md b/docs/integration.md index 6bc521d..0c401b7 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -18,7 +18,7 @@ MossInfer is a low-level engine: it runs what the caller asks for, refuses what | Area | What works | |---|---| -| Platform | Linux x86-64 and Windows x64 (MSVC), each in a `cu12` and a `cu13` flavour, as a full package and a `-system-cuda` package. CI builds, packages and runs the CPU tests of all four standard flavours, plus the Windows CUDA 13 Turing build (`-sm75`) and Linux CUDA 12 Volta build (`-sm70`); GPU runs are verified on Linux `cu13`. | +| Platform | Linux x86-64 and Windows x64 (MSVC), each in a `cu12` and a `cu13` flavour, as a full package and a `-system-cuda` package. CI builds, packages and runs the CPU tests of all four standard flavours, plus the Windows CUDA 13 Turing build (`-sm75`) and Linux CUDA 12 Volta build (`-sm70`); GPU runs are verified on Linux `cu13`, and on Linux `cu12` for Volta. | | GPUs | NVIDIA compute capability 8.x and 12.x: RTX 30xx, RTX 40xx, RTX 50xx, A100/A10/A40, L4/L40, RTX PRO Blackwell. Compute capability 7.5 (Turing: RTX 20xx, T4, GTX 16xx) with a build configured `-DMOSSINFER_SM75=ON`: the bf16 GEMMs run as fp16 tensor-core GEMMs, INT8/INT4 ConvRot files run the INT8 GEMMs (INT4 through the INT8 unpack fallback), the exact attention and the VAE use the sm_75 kernels, and an unset attention resolves to `exact` (there is no 8-bit attention kernel on Turing). See [sm75.md](sm75.md). Compute capability 7.0 (Volta, Tesla V100) with CUDA 12 and `-DMOSSINFER_SM70=ON`: dedicated fp16 tensor-core attention and exact INT8 GEMMs, sharing the Turing layer’s conversion and VAE paths. See [sm70.md](sm70.md). | | Models | MOSS-Image Base and SFT only: the official Hugging Face folders (bf16), read in place, and fine-tunes with the same architecture; other models are refused when they load. See [Supported models](#supported-models). | | Component model files | a single-file DiT, text encoder or VAE next to the bundle folder: ComfyUI and community files, INT8 ConvRot, INT4 ConvRot and NVFP4 DiT (NVFP4 on RTX 50 only), INT8 ConvRot and W4A8 text encoder (`dit_file`, `text_encoder_file`, `vae_file`); community GGUF DiTs (Q8_0, Q6_K, Q5_K, Q4_K, Q4_0, BF16, F16), converted once into INT8 ConvRot or bf16 ([GGUF DiT files](#gguf-dit-files)) | @@ -372,7 +372,8 @@ Decision procedure for the launcher: **Turing machines.** The standard packages carry no `sm_75` kernel image, so their `hello` lists a compute capability 7.5 device as unsupported (`"this build has no kernel image for compute capability 7.5"`). Install the `-sm75` package -instead (the Windows CUDA 13 one today; any build configured `-DMOSSINFER_SM75=ON` carries it): its `hello` reports the +instead (the Windows CUDA 13 one today; any build configured `-DMOSSINFER_SM75=ON` carries it, and so does the Linux +`-sm70` package): its `hello` reports the device supported, while `recommended_flavour` stays `cu13` because that follows the driver, not the package. Branch on `devices[i].supported` and `unsupported_reason`, not on `recommended_flavour` alone, and prefer the `-sm75` asset when the only GPUs are 7.5. What such a build runs is in [sm75.md](sm75.md). diff --git a/docs/sm70.md b/docs/sm70.md index 03029cb..84acfb6 100644 --- a/docs/sm70.md +++ b/docs/sm70.md @@ -53,85 +53,34 @@ The quality policy also follows Turing: an unset attention mode resolves to `exa (img2img, inpainting and references), and hires fix retain the existing sm_7x refusal/fallback policy. No sm_80 kernel is launched as a substitute for a missing Volta implementation. -## Verification and performance - -`sm70_check` reuses the compatibility layer's host-reference harness, compiled for sm_70. It exercises conversion, -fp16 GEMM, large activation rows and workspace chunks, INT8 GEMM, VAE convolution/attention, and both exact attention -head dimensions. `GpuSm70` also checks all three INT8 epilogues, row-stride padding, tail shapes, CUDA Graph replay, -and cancellation at K=131056 against a host int64 sum. Attention also checks log-sum-exp (including fully masked -rows) against the host reference, including the larger query tile. VAE tests compare batched convolutions against -FP64 host sums, including spatial/channel tails and an aliased residual; attention covers both head dimensions, -strided Q/K/V, slab tails, a 16-byte-aligned workspace and Graph replay. - -`sm70_bench [rows] [sequence_length]` measures the four DiT GEMM shapes and both attention shapes with CUDA events -after warm-up. `sm70_bench vae` measures FP32 VAE convolutions and the 16384-token attention. These are development -tools, not installed. Run the correctness tests before comparing their timings. - -Measured on a Tesla V100-SXM2-32GB (80 SMs), CUDA 12.9.86, driver 580.196.07, GCC 14.2, Release, on 2026-11-01. -One GPU per measurement, default clocks, no concurrent workload on that GPU. CUDA-event means over 20 launches -following five warm-up launches, `sm70_bench 8222 4111`: - -| Kernel | Shape | Time | Effective INT8 dot-product rate | -|---|---|---:|---:| -| QKV | M=8222, N=9216, K=3072 | 7.52 ms | 61.9 TOP/s | -| Projection + gated residual | M=8222, N=3072, K=3072 | 2.98 ms | 52.1 TOP/s | -| Up/gate + SiTU-GLU | M=8222, N=16384, K=3072 | 12.11 ms | 68.4 TOP/s | -| Down + gated residual | M=8222, N=3072, K=8192 | 5.92 ms | 69.9 TOP/s | -| DiT exact attention | B=2, L=4111, H=24, D=128 | 14.57 ms | — | -| Causal text attention | B=2, L=282, H=16, KV=4, D=256 | 0.259 ms | — | -| VAE FP32 attention | B=1, L=16384, D=384 | 34.75 ms | — | -| VAE FP32 convolution | H=W=512, C=K=384, R=3 | 46.33 ms | — | - -The INT8 tile is 128x128x32 for M >= 1024 and 64x128x32 for smaller inputs, with 32x64 warps and the CUTLASS -Volta shared-memory swizzle/register pipeline. Main-loop storage is reused by the fused output stage. The SiTU -factors use the same bf16 lookup-table values as the Turing path. No global dequantization buffer or expanded -resident weights are allocated. D=128 attention uses 192 query rows for sequences >= 1024 and 128 for shorter -sequences; D=256 keeps 64 rows. All use 32 keys per tile. The larger DiT tiles reduce redundant loads; forcing -lower register limits, smaller warp tiles, or larger attention key tiles was slower or exceeded launch resources. - -VAE convolution uses a two-stage FP32 SIMT implicit GEMM, with 128x128x8 tiles (128x32x8 for narrow outputs), and -shares the existing bias-then-residual epilogue with the sm_80 implementation. It needs no scratch or im2col tensor. -VAE mid-attention reuses dead GroupNorm scratch for score slabs: up to 2048 query rows, limited to 64 MiB of scores -at supported image sizes, plus 8 MiB for cuBLASLt and 256 bytes for alignment. Two strict FP32 SGEMMs and a row -softmax replace the old per-query loop. It uses the caller's handle, stream and arena; no hidden device allocation, -weight conversion, synchronization or persistent plan cache is introduced. The standalone no-workspace attention -API keeps its original compatibility kernel. Default-off and Turing dispatch retain their existing implementations. - -End to end: MOSS-Image-SFT, Comfy-Org INT8 ConvRot DiT and Qwen3.5 9B W4A8 text encoder, `high`, exact attention, -Euler, CFG 3, seed 42, prompt `a red apple on a wooden table, natural light, detailed photograph`: - -The baseline is the first SM70 implementation, commit `76b411a`; both runs use the same local weights and cached -text embedding, with default clocks on GPU 0. Totals include model loading and decode. Each generation uses one GPU. - -| Run / stage | Baseline | Optimized | Time reduction | -|---|---:|---:|---:| -| 1024x1024, 20 steps, `max_speed`, total | 45.35 s | 38.93 s | 14.2% | -| 20-step denoise | 41.35 s | 37.68 s | 8.9% | -| 1024x1024, 3 steps, `max_speed`, total | 10.24 s | 6.90 s | 32.6% | -| 3-step denoise | 6.16 s | 5.63 s | 8.6% | -| 1024x1024 decode | 3.28 s | 0.53 s | 83.8% | -| Device peak | 6.74 GiB | 6.74 GiB | unchanged | - -A second three-step generation in the same process takes 6.18 s (5.63 s denoise, 0.53 s decode). A separate -256x256 three-step `min_memory` run with an uncached green-cup prompt takes 2.14 s: 0.70 s text encoding, -1.14 s denoise and 0.05 s decode, at a 1.04 GiB device peak; all 42 DiT blocks stream from disk. - -Nsight Systems identified the main avoidable costs in the baseline: 1.47 s in the VAE attention kernel and -1.77 s across its convolutions. The new attention takes ~35 ms; individual representative convolutions are -4–4.6x faster (8.2x for the measured 1x1 shortcut). DiT attention improves from 17.41 to 14.57 ms, and the four -INT8 shapes improve by about 4–7%. CPU launch overhead was small; a graph/scheduler rewrite was not warranted -by this profile. These are measurements against the previous Volta implementation, not another inference engine. - -The INT8 sums and quantization epilogues remain exact under their existing contract. The VAE keeps FP32 operands -and accumulation; its changed reduction order is not bit-identical. Against the baseline PNGs, the three-step image -has 270 of 3,145,728 color components differing by one 8-bit level (PSNR 88.79 dB); the 20-step image has 240 such -components (89.31 dB). Neither comparison has a difference larger than one level. This checks the measured prompt -and seeds, not all possible images. - -Validation on V100: `sm70_check all` and `big`, all eight `Sm70Build`/`GpuSm70`/`GpuSm75Convert` tests, Compute Sanitizer -memcheck (zero errors), and racecheck on VAE, DiT attention and all INT8 epilogues including multiple accumulation -segments (zero errors/warnings). The default-off build also compiles the changed VAE translation units and passes -its nine CPU probe tests. CTest reports no failures across 517 selected CPU cases -(with the normal missing-fixture skips), excluding the pre-existing `cli_check_verdict` CMake-script issue on CMake -3.31: that script drops empty list elements under unset CMP0007. Running it through a wrapper that sets CMP0007 to -NEW also passes. The shared-library `moss-server` and CLI also produce identical three-step 1024x1024 PNGs. The implementation was hardware-tested on V100; other devices were not available for regression. +## Tiles and memory + +The INT8 tile is 128x128x32 for M >= 1024 and 64x128x32 below, with 32x64 warps and the CUTLASS Volta shared-memory +swizzle and register pipeline. D=128 attention uses 192 query rows for sequences >= 1024 and 128 below; D=256 uses 64. +All attention tiles walk 32 keys. No expanded weights or global dequantization buffers are allocated. + +VAE convolution is a two-stage fp32 SIMT implicit GEMM (128x128x8 tiles, 128x32x8 for narrow outputs) and needs no +scratch. VAE mid-attention runs two fp32 SGEMMs and a row softmax over slabs of up to 2048 query rows, in the caller's +workspace: up to 64 MiB of scores, plus 8 MiB for cuBLASLt. The standalone no-workspace attention API keeps the +compatibility kernel. The VAE's reduction order differs from that kernel's, so it is fp32-accurate but not +bit-identical to it. + +## Performance + +Measured on a Tesla V100-SXM2-32GB, MOSS-Image SFT, the Comfy-Org INT8 ConvRot DiT and the Qwen3.5 9B W4A8 text +encoder, `high`, exact attention, Euler, CFG 3: + +| Mode | Size | Denoise step | VAE decode | Device peak | +|---|---|---|---|---| +| `max_speed` | 1024² | **1.88 s** | 0.53 s | 6.74 GiB | +| `min_memory` | 256² | 0.38 s | 0.05 s | **1.04 GiB** | + +A 20-step 1024² generation takes 39 s, including model loading and decode, from a cached text embedding. + +## Verification + +`sm70_check` reuses the compatibility layer's host-reference harness, compiled for sm_70: conversion, fp16 GEMM, large +activation rows and workspace chunks, INT8 GEMM, VAE convolution and attention, and both exact attention head +dimensions with their log-sum-exp. `GpuSm70` checks all three INT8 epilogues, row-stride padding, tail shapes, CUDA +Graph replay and cancellation at K=131056 against a host int64 sum, and the VAE kernels against host and compatibility +references, including an aliased residual, slab tails and a 16-byte-aligned workspace. diff --git a/tools/sm70_bench.cu b/tools/sm70_bench.cu deleted file mode 100644 index 8955b4a..0000000 --- a/tools/sm70_bench.cu +++ /dev/null @@ -1,141 +0,0 @@ -// V100 rate probe, not installed. Run sm70_check first; timings are CUDA events after warm-up, not host launch time. -// Usage: sm70_bench [GEMM rows (default 8222)] [attention sequence length (default 4111)], or sm70_bench vae. -#include -#include - -#include -#include -#include -#include -#include -#include - -#include "core/cuda_check.h" -#include "kernels/attention/flash_attention.h" -#include "kernels/sm70/sm70.h" -#include "kernels/sm70/vae_sm70.h" -#include "kernels/vae/vae_kernels.h" - -namespace { - -class Buffer { - public: - explicit Buffer(std::size_t bytes) { - MI_CUDA_CHECK(cudaMalloc(&p_, bytes)); - MI_CUDA_CHECK(cudaMemset(p_, 0, bytes)); - } - ~Buffer() { cudaFree(p_); } - Buffer(const Buffer&) = delete; - Buffer& operator=(const Buffer&) = delete; - template T* get() const { return static_cast(p_); } - private: - void* p_ = nullptr; -}; - -template -float time_ms(F&& f) { - for (int i = 0; i < 5; ++i) f(); - cudaEvent_t start = nullptr, end = nullptr; - MI_CUDA_CHECK(cudaEventCreate(&start)); - MI_CUDA_CHECK(cudaEventCreate(&end)); - MI_CUDA_CHECK(cudaEventRecord(start)); - constexpr int kIters = 20; - for (int i = 0; i < kIters; ++i) f(); - MI_CUDA_CHECK(cudaEventRecord(end)); - MI_CUDA_CHECK(cudaEventSynchronize(end)); - float ms = 0.f; - MI_CUDA_CHECK(cudaEventElapsedTime(&ms, start, end)); - cudaEventDestroy(start); - cudaEventDestroy(end); - return ms / kIters; -} - -void gemm(int m, int n, int k, int epilogue) { - Buffer x(static_cast(m) * k), w(static_cast(n) * k); - Buffer xs(static_cast(m) * 4), ws(static_cast(n) * 4); - Buffer bias(static_cast(n) * 2), gate(static_cast(n) * 2); - Buffer out(static_cast(m) * n * 2); - mi::kernels::quant::Int8Operands p{x.get(), k, xs.get(), w.get(), ws.get(), m, n, k}; - const float ms = time_ms([&] { - if (epilogue == 2) mi::kernels::quant::int8_gemm_situ_glu(p, out.get<__nv_bfloat16>(), n / 2, nullptr); - else if (epilogue == 1) mi::kernels::quant::int8_gemm_gated_residual(p, bias.get<__nv_bfloat16>(), gate.get<__nv_bfloat16>(), - out.get<__nv_bfloat16>(), n, nullptr); - else mi::kernels::quant::int8_gemm(p, bias.get<__nv_bfloat16>(), out.get<__nv_bfloat16>(), n, nullptr); - }); - std::printf("gemm M=%d N=%d K=%d epilogue=%d: %.3f ms, %.2f TOP/s\n", m, n, k, epilogue, ms, - 2.0 * m * n * k / (ms * 1e9)); -} - -void attention(int length, int heads, int kv_heads, int dim, bool causal) { - constexpr int batch = 2; - const int tokens = batch * length; - Buffer q(static_cast(tokens) * heads * dim * 2), out(static_cast(tokens) * heads * dim * 2); - Buffer k(static_cast(tokens) * kv_heads * dim * 2), v(static_cast(tokens) * kv_heads * dim * 2); - Buffer cu((batch + 1) * sizeof(int)), lse(static_cast(tokens) * heads * sizeof(float)); - const int bounds[] = {0, length, tokens}; - MI_CUDA_CHECK(cudaMemcpy(cu.get(), bounds, sizeof(bounds), cudaMemcpyHostToDevice)); - const auto tensor = [](const Buffer& b, int t, int h, int d) { - return mi::Tensor::view(b.get(), mi::DType::kBF16, mi::Device::kCUDA, {t, h, d}); - }; - const auto tq = tensor(q, tokens, heads, dim), tk = tensor(k, tokens, kv_heads, dim), tv = tensor(v, tokens, kv_heads, dim); - const auto to = tensor(out, tokens, heads, dim); - const auto tc = mi::Tensor::view(cu.get(), mi::DType::kI32, mi::Device::kCUDA, {batch + 1}); - const float ms = time_ms([&] { - mi::attention::flash_fwd(tq, tk, tv, to, tc, tc, length, causal, 1.f / std::sqrt(static_cast(dim)), - lse.get(), static_cast(tokens) * heads * sizeof(float), nullptr); - }); - std::printf("attention L=%d H=%d KV=%d D=%d causal=%d: %.3f ms\n", length, heads, kv_heads, dim, causal, ms); -} - -void vae() { - using mi::DType; - using mi::Device; - using mi::Tensor; - for (int h : {128, 256, 512}) { - constexpr int channels = 384; - Buffer x(static_cast(h) * h * channels * sizeof(float)); - Buffer w(static_cast(channels) * 3 * 3 * channels * sizeof(float)), bias(channels * sizeof(float)); - Buffer out(static_cast(h) * h * channels * sizeof(float)); - const auto tx = Tensor::view(x.get(), DType::kF32, Device::kCUDA, {1, h, h, channels}); - const auto tw = Tensor::view(w.get(), DType::kF32, Device::kCUDA, {channels, 3, 3, channels}); - const auto tb = Tensor::view(bias.get(), DType::kF32, Device::kCUDA, {channels}); - const auto ty = Tensor::view(out.get(), DType::kF32, Device::kCUDA, {1, h, h, channels}); - const float ms = time_ms([&] { mi::vae::conv2d(tx, tw, tb, nullptr, ty, nullptr, 0, nullptr); }); - std::printf("vae fp32 conv H=W=%d C=K=%d R=3: %.3f ms\n", h, channels, ms); - } - constexpr int length = 16384, dim = 384; - Buffer qkv(static_cast(length) * dim * 3 * sizeof(float)); - Buffer out(static_cast(length) * dim * sizeof(float)); - const auto q = Tensor::view(qkv.get(), DType::kF32, Device::kCUDA, {1, length, dim}); - const auto k = Tensor::view(qkv.get() + length * dim, DType::kF32, Device::kCUDA, {1, length, dim}); - const auto v = Tensor::view(qkv.get() + 2 * length * dim, DType::kF32, Device::kCUDA, {1, length, dim}); - const auto o = Tensor::view(out.get(), DType::kF32, Device::kCUDA, {1, length, dim}); - const auto bytes = mi::kernels::sm70::vae_attention_workspace_bytes(length); - Buffer scratch(bytes); - cublasLtHandle_t lt = nullptr; - if (cublasLtCreate(<) != CUBLAS_STATUS_SUCCESS) throw std::runtime_error("cuBLASLt handle creation failed"); - const float ms = time_ms([&] { - mi::kernels::sm70::vae_attention_fp32(lt, q, k, v, o, 1.f / std::sqrt(static_cast(dim)), scratch.get(), - bytes, nullptr); - }); - cublasLtDestroy(lt); - std::printf("vae fp32 attention L=%d D=%d: %.3f ms, %.1f MiB scratch\n", length, dim, ms, - static_cast(bytes) / (1024 * 1024)); -} - -} // namespace - -int main(int argc, char** argv) { - std::setvbuf(stdout, nullptr, _IONBF, 0); - if (!mi::kernels::sm70::volta()) return 1; - if (argc > 1 && std::string_view(argv[1]) == "vae") { vae(); return 0; } - const int m = argc > 1 ? std::atoi(argv[1]) : 8222; - const int length = argc > 2 ? std::atoi(argv[2]) : 4111; - if (m <= 0 || length <= 0) return 2; - gemm(m, 9216, 3072, 0); - gemm(m, 3072, 3072, 1); - gemm(m, 16384, 3072, 2); - gemm(m, 3072, 8192, 1); - attention(length, 24, 24, 128, false); - attention(282, 16, 4, 256, true); -}