From 561c77adec52f3ca43387cb8da1ef9521924a146 Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sat, 31 Oct 2026 10:50:45 +0800 Subject: [PATCH 1/6] sm_75 (Turing) support: fp16 GEMMs, INT8 GEMMs, exact attention, VAE Build an sm_75 kernel family next to the sm_80+ code and dispatch to it by compute capability, behind a MOSSINFER_SM75 option (off: Ampere-or-newer, as before). - cmake: the sm_75 kernel family is an OBJECT library with 75-real SASS (src/kernels/sm75); the portable SIMT translation units are also compiled for 75; the probe accepts compute capability 7.5 and reports 75-real in its architecture list when MOSSINFER_SM75 is on. - kernels/basic: a cuBLASLt fp16 GEMM path for the bf16 linears (weights converted once at load or per streamed block, x converted in row chunks, output converted in place with the bias added in fp32). Text encoder, DiT globals and DiT blocks hook the weight conversion. - kernels/sm75: bf16<->fp16 conversions; exact attention on m16n8k8 (DiT joint attention head_dim 128 varlen GQA, Qwen causal head_dim 256); CUTLASS Sm75 m8n8k16 INT8 GEMMs that reuse the sm_80 epilogues; fp32 VAE conv2d, mid attention and conv weight layout. - gdn: a recurrent sm_75 kernel (the chunked prep kernel needs 85 KiB of shared memory, over Turing's 64 KiB). - tier/planner: an unset attention resolves to exact when the device has no 8-bit attention kernel (said in the reason); sm_75 DiT blocks always run op by op (the fused bf16 CUTLASS GEMMs need sm_80+ tensor cores). - vendored FlashAttention: if constexpr on the dropout/return-softmax guards, which otherwise instantiate dropout code that does not compile for the Turing MMA layout. - docs and tests: support matrix, sm_75 smoke/conversion tests, planner fallback test, and tools/sm75_check (bring-up harness with host references). Verified on an RTX 2080 8 GiB with the SFT bundle: the reference, high and turbo tiers generate through moss-infer and moss-server (protocol 2); the INT8 GEMM is bit-exact against a host int32 reference; MOSSINFER_SM75=OFF still builds and keeps the same CPU test baseline. --- CMakeLists.txt | 10 +- README.md | 4 +- cmake/modules/backend.cmake | 14 +- cmake/modules/sm75.cmake | 109 ++++ docs/integration.md | 4 +- src/backend/probe.cpp | 4 + src/backend/probe.h | 9 +- src/kernels/attention/flash_attention.cu | 18 +- src/kernels/basic/gemm.cpp | 111 +++- src/kernels/gdn/chunk_rule.cu | 89 +++ src/kernels/quant/int8_gemm.cu | 6 + src/kernels/quant/int8_gemm_impl.cuh | 7 +- src/kernels/quant/int8_gemm_residual.cu | 6 + src/kernels/quant/int8_gemm_situ.cu | 6 + src/kernels/sm75/attention_sm75.cu | 334 +++++++++++ src/kernels/sm75/convert.cu | 206 +++++++ src/kernels/sm75/int8_gemm_sm75.cu | 103 ++++ src/kernels/sm75/int8_sm75.h | 25 + src/kernels/sm75/sm75.h | 89 +++ src/kernels/sm75/vae_sm75.cu | 180 ++++++ src/kernels/vae/attention.cu | 7 + src/kernels/vae/conv2d.cu | 11 + src/models/moss_dit/dit_weights.cpp | 29 + src/models/moss_dit/dit_weights.h | 5 + src/models/moss_dit/moss_dit.cpp | 15 + src/models/qwen35_text/text_encoder.cpp | 21 +- src/pipeline/pipeline.cpp | 8 +- src/pipeline/pipeline_impl.h | 9 +- src/planner/planner.cpp | 21 +- src/planner/planner.h | 5 +- tests/unit/test_attention.cu | 11 + tests/unit/test_backend_probe.cpp | 8 +- tests/unit/test_planner.cpp | 8 + tests/unit/test_sm75.cpp | 126 +++++ third_party/flash-attention/VENDORED.md | 7 + .../csrc/flash_attn/src/flash_fwd_kernel.h | 8 +- tools/sm75_check.cu | 532 ++++++++++++++++++ 37 files changed, 2124 insertions(+), 41 deletions(-) create mode 100644 cmake/modules/sm75.cmake create mode 100644 src/kernels/sm75/attention_sm75.cu create mode 100644 src/kernels/sm75/convert.cu create mode 100644 src/kernels/sm75/int8_gemm_sm75.cu create mode 100644 src/kernels/sm75/int8_sm75.h create mode 100644 src/kernels/sm75/sm75.h create mode 100644 src/kernels/sm75/vae_sm75.cu create mode 100644 tests/unit/test_sm75.cpp create mode 100644 tools/sm75_check.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index c259252..9cb46b8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -28,6 +28,9 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON) option(MOSSINFER_WERROR "Treat warnings as errors" ON) option(MOSSINFER_ASAN "Build host code with AddressSanitizer" OFF) 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) find_package(CUDAToolkit 12.4 REQUIRED) # sm_120 (Blackwell GeForce) needs CUDA 12.8. @@ -178,8 +181,11 @@ 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}) - # Their reference arithmetic is compared bit for bit with the NVFP4 recipes (MOSSINFER_STRICT_FP_SOURCES) and with - # the samplers' kernels (the component API's host loops, design §5.3). + if(MOSSINFER_SM75) + # 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() + # Its reference arithmetic is compared bit for bit with the NVFP4 recipes (MOSSINFER_STRICT_FP_SOURCES). set(_mi_strict_tests ${MOSSINFER_TEST_SOURCES}) list(FILTER _mi_strict_tests INCLUDE REGEX ".*/tests/unit/test_(nvfp4_convert|component_gpu)\\.cpp$") if(_mi_strict_tests) diff --git a/README.md b/README.md index 1e16952..c6af1a7 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.0. 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. The full support matrix is in [docs/integration.md](docs/integration.md#support-status). +Version 2.1.0. 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`: fp16 tensor-core GEMMs, INT8/INT4 ConvRot files (INT4 through the INT8 unpack fallback), exact attention and the VAE then run on sm_75 kernels, and an unset attention resolves to `exact`. 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). @@ -24,6 +24,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` and `windows-msvc-cuda12` ([docs/integration.md](docs/integration.md#getting-the-binaries)). +The other presets are `linux-cuda12`, `windows-msvc-cuda13` and `windows-msvc-cuda12` ([docs/integration.md](docs/integration.md#getting-the-binaries)). Add `-DMOSSINFER_SM75=ON` to also build the Turing (sm_75) kernels and accept compute capability 7.5 devices; without it the engine is Ampere-or-newer as before. Licensed under Apache-2.0 (`LICENSE`, `NOTICE`, `THIRD_PARTY_NOTICES.md`). diff --git a/cmake/modules/backend.cmake b/cmake/modules/backend.cmake index ae0409a..b41c85e 100644 --- a/cmake/modules/backend.cmake +++ b/cmake/modules/backend.cmake @@ -1,10 +1,22 @@ # 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. # 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) + get_target_property(_mi_core_archs mossinfer_core CUDA_ARCHITECTURES) + if(NOT _mi_core_archs) + set(_mi_core_archs "") + endif() + string(JOIN "," _mi_arch_string ${_mi_core_archs}) + if(MOSSINFER_SM75) + string(APPEND _mi_arch_string ",75-real") + endif() set_property(SOURCE ${_mi_backend_probe} APPEND PROPERTY COMPILE_DEFINITIONS - "MOSSINFER_CUDA_ARCHITECTURES=\"$,$>\"") + "MOSSINFER_CUDA_ARCHITECTURES=\"${_mi_arch_string}\"") + unset(_mi_arch_string) + unset(_mi_core_archs) endif() unset(_mi_backend_probe) # src/backend/cuda_driver.cpp loads libcuda and NVML at run time (dlopen; libdl before glibc 2.34). diff --git a/cmake/modules/sm75.cmake b/cmake/modules/sm75.cmake new file mode 100644 index 0000000..7985703 --- /dev/null +++ b/cmake/modules/sm75.cmake @@ -0,0 +1,109 @@ +# Turing (sm_75) kernels: the engine's tensor-core work needs bf16 MMA, cp.async, TF32 and up to 99 KiB of shared +# memory, none of which exist on compute capability 7.x. src/kernels/sm75/ holds the fp16 fallbacks, built as SASS for +# sm_75 alone in their own OBJECT library (like cmake/modules/nvfp4.cmake does for sm_120a), so the rest of +# mossinfer_core keeps its sm_80/89/120 architecture list and none of the sm_80-only translation units see --arch=sm_75. +# +# The kernel family is built in every build: the host code of the sm_80 translation units references it unconditionally +# (and only executes it on a 7.x device). MOSSINFER_SM75 (declared in CMakeLists.txt) controls what makes 7.x usable: +# - the probe accepts compute capability 7.5 (MOSSINFER_HAVE_SM75 on probe.cpp, and 75-real in its architecture list, +# cmake/modules/backend.cmake), +# - the portable SIMT translation units are also compiled for sm_75, +# - flash_attention.cu / gemm.cpp dispatch into the sm_75 kernels, +# - the sm75_check bring-up tool is built. +# 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$") +if(_mi_sm75_sources) + get_target_property(_mi_core_sources mossinfer_core SOURCES) + list(REMOVE_ITEM _mi_core_sources ${_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})") +endif() + +if(MOSSINFER_SM75 AND _mi_sm75_sources) + # 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). + set(_mi_portable_sources + src/kernels/basic/elementwise.cu + src/kernels/basic/gram.cu + src/kernels/basic/norm.cu + src/kernels/basic/rope.cu + src/kernels/fused/qk_norm_rope.cu + src/kernels/gdn/chunk_rule.cu + src/kernels/gdn/conv_gates.cu + src/kernels/lora/lora_kernels.cu + src/kernels/quant/rotate_quantize.cu + src/kernels/quant/w4a8_decode.cu + src/kernels/quant4/convrot_quantize.cu + src/kernels/vae/elementwise.cu + src/kernels/vae/encoder_elementwise.cu + src/kernels/vae/group_norm.cu + src/models/flux2_vae/torch_group_norm.cu + src/models/latent_preview/latent_preview.cu + src/models/moss_dit/dit_kernels.cu + src/pipeline/accel/accel_kernels.cu + src/pipeline/samplers/sampler_kernels.cu + src/pipeline/samplers/torch_cuda_randn.cu) + set(_mi_portable_present "") + foreach(rel IN LISTS _mi_portable_sources) + if(${CMAKE_SOURCE_DIR}/${rel} IN_LIST MOSSINFER_SOURCES) + list(APPEND _mi_portable_present ${CMAKE_SOURCE_DIR}/${rel}) + endif() + endforeach() + if(_mi_portable_present) + get_target_property(_mi_core_sources mossinfer_core SOURCES) + list(REMOVE_ITEM _mi_core_sources ${_mi_portable_present}) + set_property(TARGET mossinfer_core PROPERTY SOURCES ${_mi_core_sources}) + unset(_mi_core_sources) + set(_mi_portable_archs 75 ${CMAKE_CUDA_ARCHITECTURES}) + list(REMOVE_DUPLICATES _mi_portable_archs) + add_library(mossinfer_sm75_portable OBJECT ${_mi_portable_present}) + set_target_properties(mossinfer_sm75_portable PROPERTIES + CUDA_ARCHITECTURES "${_mi_portable_archs}" + POSITION_INDEPENDENT_CODE ON + CUDA_SEPARABLE_COMPILATION OFF) + target_include_directories(mossinfer_sm75_portable PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/src) + target_compile_options(mossinfer_sm75_portable PRIVATE --expt-relaxed-constexpr -lineinfo + ${MOSSINFER_CUDA_HOST_WARN}) + target_sources(mossinfer_core PRIVATE $) + message(STATUS "SM75: portable kernels for ${_mi_portable_archs}") + unset(_mi_portable_archs) + endif() + unset(_mi_portable_present) + unset(_mi_portable_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 + ${CMAKE_SOURCE_DIR}/src/kernels/attention/flash_attention.cu + ${CMAKE_SOURCE_DIR}/src/kernels/basic/gemm.cpp + APPEND PROPERTY COMPILE_DEFINITIONS MOSSINFER_HAVE_SM75=1) + 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. + if(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) + target_link_libraries(sm75_check PRIVATE mossinfer_core) + target_compile_options(sm75_check PRIVATE --expt-relaxed-constexpr -lineinfo ${MOSSINFER_CUDA_HOST_WARN}) + endif() +endif() +unset(_mi_sm75_sources) diff --git a/docs/integration.md b/docs/integration.md index c545f54..2cb8dfd 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -19,7 +19,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 flavours; 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 | +| 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). | | 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 | @@ -65,7 +65,7 @@ These are designed and partly built. The launcher can prepare for them now (see | Area | Status | |---|---| -| GPUs with compute capability below 8.0 (GTX 16xx, RTX 20xx and older) | not supported: the engine needs bf16 tensor cores | +| GPUs with compute capability below 8.0 unless the build enables `MOSSINFER_SM75` (GTX 16xx, RTX 20xx and older) | not supported without that option: the engine needs bf16 tensor cores. With `-DMOSSINFER_SM75=ON` compute capability 7.5 runs (fp16 GEMMs + exact attention); older than 7.5 stays unsupported | | GPUs with compute capability 9.x, 10.x or 11.x (H100/H200, B200/GB200 and similar) | not supported: no kernels are built for them | | AMD, Intel and Apple GPUs; CPU-only machines; macOS | not supported | | 32-bit systems | not supported | diff --git a/src/backend/probe.cpp b/src/backend/probe.cpp index a8b73a1..f582712 100644 --- a/src/backend/probe.cpp +++ b/src/backend/probe.cpp @@ -133,6 +133,10 @@ void evaluate_device(DeviceInfo& dev, int driver_version, int build_cuda_version if (std::tie(dev.cc_major, dev.cc_minor) < std::tie(kMinComputeMajor, kMinComputeMinor)) { dev.reason = "compute capability " + cc_string(dev.cc_major, dev.cc_minor) + " is below the required " + cc_string(kMinComputeMajor, kMinComputeMinor) + " (Ampere or newer)"; +#if MOSSINFER_HAVE_SM75 + // The sm_80 path needs Ampere; this build's sm_75 kernel family reaches down to Turing. + dev.reason = "compute capability " + cc_string(dev.cc_major, dev.cc_minor) + " is below the required 7.5 (Turing)"; +#endif return; } const Flavour build = flavour_of_cuda_version(build_cuda_version); diff --git a/src/backend/probe.h b/src/backend/probe.h index b2200db..a5494d2 100644 --- a/src/backend/probe.h +++ b/src/backend/probe.h @@ -16,8 +16,13 @@ namespace mi::backend { -// Minimum compute capability of every kernel in the engine (bf16 tensor cores, cp.async). -inline constexpr int kMinComputeMajor = 8; +// 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. +#ifndef MOSSINFER_HAVE_SM75 +#define MOSSINFER_HAVE_SM75 0 +#endif +inline constexpr int kMinComputeMajor = MOSSINFER_HAVE_SM75 ? 7 : 8; inline constexpr int kMinComputeMinor = 0; // Release flavours, named by the CUDA major version of the toolkit they are built with. diff --git a/src/kernels/attention/flash_attention.cu b/src/kernels/attention/flash_attention.cu index a4a72c1..5af00ab 100644 --- a/src/kernels/attention/flash_attention.cu +++ b/src/kernels/attention/flash_attention.cu @@ -9,6 +9,7 @@ #include "core/cuda_check.h" #include "core/varlen.h" #include "kernels/attention/flash_fwd_kernels.h" +#include "kernels/sm75/sm75.h" namespace mi::attention { namespace { @@ -48,7 +49,9 @@ int tile_smem_bytes(FlashTile t) { std::atomic g_tile_override{FlashTile::kAuto}; // The kernels use cp.async and bf16 mma, which need sm_80+; on older parts they would compile to empty -// bodies (see MI_FLASH_PARAM_MODIFIER) and silently leave `out` unwritten. Checked once per device. +// bodies (see MI_FLASH_PARAM_MODIFIER) and silently leave `out` unwritten. A build with the sm_75 kernel family +// (MOSSINFER_SM75) dispatches compute capability 7.x to kernels::sm75::flash_fwd instead, so this check only runs on +// the sm_80 path. Checked once per device. void require_sm80(int device) { static std::array, 64> ok{}; if (device >= 0 && device < static_cast(ok.size()) && ok[device].load()) return; @@ -149,6 +152,19 @@ 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_SM75 + { + int major = 0; + MI_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + if (major < 8) { + // Turing: the vendored kernel's own sm_75 traits (fp16 MMA, no cp.async) run in the sm_75 kernel family; + // q/k/v/out are converted around the launch there. + kernels::sm75::flash_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen, causal, scale, workspace, + workspace_bytes, stream); + return; + } + } +#endif require_sm80(device); MI_CUDA_CHECK(cudaDeviceGetAttribute(&smem_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, device)); MI_CUDA_CHECK(cudaDeviceGetAttribute(&smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); diff --git a/src/kernels/basic/gemm.cpp b/src/kernels/basic/gemm.cpp index fc3369c..29d117e 100644 --- a/src/kernels/basic/gemm.cpp +++ b/src/kernels/basic/gemm.cpp @@ -9,6 +9,7 @@ #include "core/cuda_check.h" #include "kernels/basic/detail.h" +#include "kernels/sm75/sm75.h" namespace mi::kernels { namespace { @@ -35,11 +36,18 @@ std::int64_t row_stride(const Tensor& t) { return m.nd == 1 ? m.stride[0] : 0; } -cudaDataType_t cuda_type(DType t) { return t == DType::kF32 ? CUDA_R_32F : CUDA_R_16BF; } +cudaDataType_t cuda_type(DType t) { + switch (t) { + case DType::kF32: return CUDA_R_32F; + case DType::kF16: return CUDA_R_16F; + default: return CUDA_R_16BF; + } +} struct Key { std::int64_t m, n, k, ldx, ldw, ldo; int out_type, bias_type; // bias_type -1: no bias + int ab_type; // 0: bf16 operands (sm_80+), 1: fp16 operands (sm_75 path) std::size_t workspace; std::uint32_t align_x, align_w, align_out, align_bias; bool operator==(const Key&) const = default; @@ -50,9 +58,9 @@ struct KeyHash { std::size_t h = 1469598103934665603ull; const auto mix = [&h](std::uint64_t v) { h = (h ^ v) * 1099511628211ull; }; for (std::int64_t v : {key.m, key.n, key.k, key.ldx, key.ldw, key.ldo}) mix(static_cast(v)); - for (std::uint64_t v : {std::uint64_t(key.out_type + 8), std::uint64_t(key.bias_type + 8), std::uint64_t(key.workspace), - std::uint64_t(key.align_x), std::uint64_t(key.align_w), std::uint64_t(key.align_out), - std::uint64_t(key.align_bias)}) { + for (std::uint64_t v : {std::uint64_t(key.out_type + 8), std::uint64_t(key.bias_type + 8), + std::uint64_t(key.ab_type), std::uint64_t(key.workspace), std::uint64_t(key.align_x), + std::uint64_t(key.align_w), std::uint64_t(key.align_out), std::uint64_t(key.align_bias)}) { mix(v); } return h; @@ -84,6 +92,22 @@ struct Gemm::Impl { cublasLtHandle_t handle = nullptr; mutable std::mutex mutex; std::unordered_map, KeyHash> plans; + // sm_75: fp16 tensor cores with bf16 activations. x is converted to fp16 in a chunked scratch buffer, the weights + // are fp16 (converted once by the loader) and the GEMM writes fp16 (into the bf16 output buffer, or into `acc16` + // for an f32 output) which a second pass converts to the output dtype and adds the bias. + bool fp16 = false; + void* x16 = nullptr; + std::size_t x16_bytes = 0; + void* acc16 = nullptr; + std::size_t acc16_bytes = 0; + + static void ensure_scratch(void*& p, std::size_t& cap, std::size_t need) { + if (cap >= need) return; + MI_CUDA_CHECK(cudaDeviceSynchronize()); // an earlier launch on another stream may still read it + if (p != nullptr) MI_CUDA_CHECK(cudaFree(p)); + MI_CUDA_CHECK(cudaMalloc(&p, need)); + cap = need; + } Plan& plan_for(const Key& key, const void* bias_ptr) { auto it = plans.find(key); @@ -93,7 +117,7 @@ struct Gemm::Impl { const cublasOperation_t trans_a = CUBLAS_OP_T, trans_b = CUBLAS_OP_N; MI_LT_CHECK(cublasLtMatmulDescSetAttribute(plan->desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a))); MI_LT_CHECK(cublasLtMatmulDescSetAttribute(plan->desc, CUBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b))); - if (key.bias_type >= 0) { + if (key.bias_type >= 0 && key.ab_type == 0) { const cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS; const auto bias_type = static_cast(cuda_type(static_cast(key.bias_type))); MI_LT_CHECK(cublasLtMatmulDescSetAttribute(plan->desc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))); @@ -104,8 +128,9 @@ struct Gemm::Impl { sizeof(bias_ptr))); } const auto u64 = [](std::int64_t v) { return static_cast(v); }; - MI_LT_CHECK(cublasLtMatrixLayoutCreate(&plan->a, CUDA_R_16BF, u64(key.k), u64(key.n), key.ldw)); - MI_LT_CHECK(cublasLtMatrixLayoutCreate(&plan->b, CUDA_R_16BF, u64(key.k), u64(key.m), key.ldx)); + const cudaDataType_t ab_type = key.ab_type ? CUDA_R_16F : CUDA_R_16BF; + MI_LT_CHECK(cublasLtMatrixLayoutCreate(&plan->a, ab_type, u64(key.k), u64(key.n), key.ldw)); + MI_LT_CHECK(cublasLtMatrixLayoutCreate(&plan->b, ab_type, u64(key.k), u64(key.m), key.ldx)); MI_LT_CHECK(cublasLtMatrixLayoutCreate(&plan->d, cuda_type(static_cast(key.out_type)), u64(key.n), u64(key.m), key.ldo)); @@ -141,12 +166,64 @@ struct Gemm::Impl { plan->workspace = result.workspaceSize; return *plans.emplace(key, std::move(plan)).first->second; } + + // sm_75 path: x is bf16, the GEMM runs fp16 in / fp32 out on Turing's tensor cores, and the result is converted to + // the output dtype with the bias added in fp32. Row-chunked so the fp16 copies stay in a small scratch buffer. + void linear_f16(const Tensor& out, const Tensor& x, const Tensor& w, const Tensor* bias, void* workspace, + std::size_t workspace_bytes, cudaStream_t stream) { + const std::int64_t m = detail::rows_of(x); + const std::int64_t k = x.dim(-1); + const std::int64_t n = w.dim(0); + const std::int64_t ldx = row_stride(x); + const std::int64_t ldo = row_stride(out); + const std::int64_t ldw = row_stride(w); + if (m == 0 || n == 0) return; + const bool out_f32 = out.dtype == DType::kF32; + constexpr std::size_t kScratchBytes = std::size_t{16} << 20; + const std::size_t per_row = static_cast(std::max(n, k)) * 2; + std::int64_t chunk = static_cast(kScratchBytes / std::max(1, per_row)); + chunk = std::max(1, std::min(chunk, m)); + ensure_scratch(x16, x16_bytes, static_cast(chunk) * static_cast(k) * 2); + if (out_f32) { + ensure_scratch(acc16, acc16_bytes, static_cast(chunk) * static_cast(n) * 2); + } + const void* bias_ptr = bias != nullptr ? bias->data : nullptr; + const bool bias_f32 = bias != nullptr && bias->dtype == DType::kF32; + for (std::int64_t r0 = 0; r0 < m; r0 += chunk) { + const std::int64_t rows = std::min(chunk, m - r0); + kernels::sm75::convert_rows_bf16_to_f16(x16, static_cast(x.data) + r0 * ldx * 2, rows, k, ldx, k, + stream); + void* dptr = out_f32 ? acc16 : (static_cast(out.data) + r0 * ldo * 2); + const std::int64_t acc_ld = out_f32 ? n : ldo; + const Key key{rows, n, k, k, ldw, acc_ld, static_cast(DType::kF16), -1, 1, workspace_bytes, + alignment(x16), alignment(w.data), alignment(dptr), 0u}; + Plan& plan = plan_for(key, nullptr); + const float alpha = 1.0f, beta = 0.0f; + MI_LT_CHECK(cublasLtMatmul(handle, plan.desc, &alpha, w.data, plan.a, x16, plan.b, &beta, dptr, plan.d, dptr, + plan.d, &plan.algo, workspace, plan.workspace, stream)); + kernels::sm75::convert_gemm_out_f16(static_cast(out.data) + r0 * ldo * 2, dptr, bias_ptr, bias_f32, + out_f32, rows, n, ldo, acc_ld, stream); + } + } }; -Gemm::Gemm() : impl_(std::make_unique()) { MI_LT_CHECK(cublasLtCreate(&impl_->handle)); } +Gemm::Gemm() : impl_(std::make_unique()) { + MI_LT_CHECK(cublasLtCreate(&impl_->handle)); +#if MOSSINFER_HAVE_SM75 + // 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 && + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device) == cudaSuccess && major < 8) { + impl_->fp16 = true; + } + cudaGetLastError(); // the two queries above leave no error the caller should see +#endif +} Gemm::~Gemm() { impl_->plans.clear(); + if (impl_->x16 != nullptr) cudaFree(impl_->x16); + if (impl_->acc16 != nullptr) cudaFree(impl_->acc16); if (impl_->handle != nullptr) cublasLtDestroy(impl_->handle); } @@ -158,11 +235,13 @@ std::size_t Gemm::cached_plans() const { namespace { // Validates linear()'s operands and returns their plan key; m or n of 0 means there is nothing to compute. -Key linear_key(const Tensor& out, const Tensor& x, const Tensor& w, const Tensor* bias, std::size_t workspace_bytes) { +// fp16: the sm_75 path, where the weights are stored fp16 (converted once by the loader). +Key linear_key(const Tensor& out, const Tensor& x, const Tensor& w, const Tensor* bias, std::size_t workspace_bytes, + bool fp16) { constexpr const char* op = "linear"; detail::check_rows(op, "x", x, DType::kBF16); const std::int64_t k = x.dim(-1); - detail::check_rows(op, "w", w, DType::kBF16, k); + detail::check_rows(op, "w", w, fp16 ? DType::kF16 : DType::kBF16, k); if (w.rank != 2) detail::fail(op, "w must be [N, K], got " + w.describe()); const std::int64_t n = w.dim(0); if (out.dtype != DType::kBF16 && out.dtype != DType::kF32) detail::fail(op, "out must be bf16 or f32, got " + out.describe()); @@ -182,7 +261,7 @@ Key linear_key(const Tensor& out, const Tensor& x, const Tensor& w, const Tensor bias_type = static_cast(bias->dtype); } const void* bias_ptr = bias != nullptr ? bias->data : nullptr; - return Key{m, n, k, ldx, ldw, ldo, static_cast(out.dtype), bias_type, workspace_bytes, + return Key{m, n, k, ldx, ldw, ldo, static_cast(out.dtype), bias_type, fp16 ? 1 : 0, workspace_bytes, alignment(x.data), alignment(w.data), alignment(out.data), bias_ptr ? alignment(bias_ptr) : 0u}; } @@ -205,7 +284,7 @@ std::string Gemm::AlgoConfig::describe() const { Gemm::AlgoConfig Gemm::algo_config(const Tensor& out, const Tensor& x, const Tensor& w, const Tensor* bias, std::size_t workspace_bytes) { - const Key key = linear_key(out, x, w, bias, workspace_bytes); + const Key key = linear_key(out, x, w, bias, workspace_bytes, impl_->fp16); if (key.m == 0 || key.n == 0) return {}; std::lock_guard lock(impl_->mutex); const Plan& plan = impl_->plan_for(key, bias != nullptr ? bias->data : nullptr); @@ -225,12 +304,18 @@ Gemm::AlgoConfig Gemm::algo_config(const Tensor& out, const Tensor& x, const Ten void Gemm::linear(const Tensor& out, const Tensor& x, const Tensor& w, const Tensor* bias, void* workspace, std::size_t workspace_bytes, cudaStream_t stream) { - const Key key = linear_key(out, x, w, bias, workspace_bytes); + const Key key = linear_key(out, x, w, bias, workspace_bytes, impl_->fp16); if (workspace_bytes > 0 && (workspace == nullptr || alignment(workspace) < 256)) { detail::fail("linear", "workspace must be 256-byte aligned"); } if (key.m == 0 || key.n == 0) return; + if (impl_->fp16) { + std::lock_guard lock(impl_->mutex); + impl_->linear_f16(out, x, w, bias, workspace, workspace_bytes, stream); + return; + } + const void* bias_ptr = bias != nullptr ? bias->data : nullptr; std::lock_guard lock(impl_->mutex); Plan& plan = impl_->plan_for(key, bias_ptr); diff --git a/src/kernels/gdn/chunk_rule.cu b/src/kernels/gdn/chunk_rule.cu index 0ca900e..3d78238 100644 --- a/src/kernels/gdn/chunk_rule.cu +++ b/src/kernels/gdn/chunk_rule.cu @@ -13,6 +13,7 @@ #include #include +#include #include #include "core/cuda_check.h" @@ -458,6 +459,79 @@ void allow_dynamic_smem(Kernel kernel, int bytes) { done.insert(key); } +// sm_75: the chunked kernel needs ~85 KiB of shared memory (PrepSmem), over Turing's 64 KiB per block. This per-token +// recurrent form computes the same rule (the chunked algorithm is an algebraic rearrangement of it) with the 128 state +// columns split across two launches of 64: every state column is independent, so the split is exact. 33 KiB of dynamic +// shared memory, 64 threads per block. +constexpr int kSm75Cols = 64; + +template +__global__ void recurrent_sm75_kernel(const __nv_bfloat16* __restrict__ qkv, std::int64_t qkv_stride, + const float* __restrict__ g, const __nv_bfloat16* __restrict__ beta, + const int* __restrict__ cu, OutT* __restrict__ out, + float* __restrict__ final_state) { + extern __shared__ float smem_sm75[]; + float* s = smem_sm75; // [kHeadDim][kSm75Cols], state column c + float* sq = s + kHeadDim * kSm75Cols; // [kHeadDim] + float* sk = sq + kHeadDim; // [kHeadDim] + float* redq = sk + kHeadDim; // [2] + float* redk = redq + 2; // [2] + const int h = blockIdx.x, kh = h / (kNumVHeads / kNumKHeads), seq = blockIdx.y; + const int c = threadIdx.x; // 0..63 + const int col = static_cast(blockIdx.z) * kSm75Cols + c; + const int lane = c & 31, warp = c >> 5; + for (int i = 0; i < kHeadDim; ++i) s[i * kSm75Cols + c] = 0.f; + for (int t = cu[seq]; t < cu[seq + 1]; ++t) { + const __nv_bfloat16* row = qkv + t * qkv_stride; + const float q0 = bf(row[kh * kHeadDim + c]); + const float q1 = bf(row[kh * kHeadDim + kSm75Cols + c]); + const float k0 = bf(row[kKeyDim + kh * kHeadDim + c]); + const float k1 = bf(row[kKeyDim + kh * kHeadDim + kSm75Cols + c]); + float qq = q0 * q0 + q1 * q1, kk = k0 * k0 + k1 * k1; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + qq += __shfl_xor_sync(0xffffffffu, qq, off); + kk += __shfl_xor_sync(0xffffffffu, kk, off); + } + if (lane == 0) { + redq[warp] = qq; + redk[warp] = kk; + } + __syncthreads(); + const float qs = redq[0] + redq[1], ks = redk[0] + redk[1]; + sq[c] = q0 * rsqrtf(qs + kNormEps) / kSqrtHeadDim; + sq[kSm75Cols + c] = q1 * rsqrtf(qs + kNormEps) / kSqrtHeadDim; + sk[c] = k0 * rsqrtf(ks + kNormEps); + sk[kSm75Cols + c] = k1 * rsqrtf(ks + kNormEps); + const float v = bf(row[2 * kKeyDim + h * kHeadDim + col]); + __syncthreads(); + const float decay = expf(g[t * kNumVHeads + h]), b = bf(beta[t * kNumVHeads + h]); + float kv = 0.f; + for (int i = 0; i < kHeadDim; ++i) { + const float sv = s[i * kSm75Cols + c] * decay; + s[i * kSm75Cols + c] = sv; + kv = fmaf(sv, sk[i], kv); + } + const float delta = (v - kv) * b; + float o = 0.f; + for (int i = 0; i < kHeadDim; ++i) { + const float sv = fmaf(sk[i], delta, s[i * kSm75Cols + c]); + s[i * kSm75Cols + c] = sv; + o = fmaf(sv, sq[i], o); + } + if constexpr (std::is_same_v) { + out[(static_cast(t) * kNumVHeads + h) * kHeadDim + col] = o; + } else { + out[(static_cast(t) * kNumVHeads + h) * kHeadDim + col] = __float2bfloat16_rn(o); + } + __syncthreads(); // sq, sk and the reductions are rewritten by the next token + } + if (final_state != nullptr) { + float* fs = final_state + (static_cast(seq) * kNumVHeads + h) * kHeadDim * kHeadDim; + for (int i = 0; i < kHeadDim; ++i) fs[i * kHeadDim + col] = s[i * kSm75Cols + c]; + } +} + } // namespace std::size_t chunk_gated_delta_rule_workspace_bytes(std::int64_t total_tokens) { @@ -484,6 +558,21 @@ void chunk_gated_delta_rule(const Tensor& qkv, const Tensor& g, const Tensor& be Workspace ws{}; workspace_layout(tokens, workspace, &ws); + // sm_75: the chunked kernels need ~85 KiB of shared memory, over Turing's 64 KiB per block. The recurrent form above + // computes the same rule with 33 KiB; its numerics are fp32 like the chunked path's. + int sm75_device = 0, sm75_major = 0; + MI_CUDA_CHECK(cudaGetDevice(&sm75_device)); + MI_CUDA_CHECK(cudaDeviceGetAttribute(&sm75_major, cudaDevAttrComputeCapabilityMajor, sm75_device)); + if (sm75_major < 8) { + 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><<>>( + qkv.ptr<__nv_bfloat16>(), qkv.stride[0], g.ptr(), beta.ptr<__nv_bfloat16>(), cu_seqlens.ptr(), + out.ptr<__nv_bfloat16>(), final_state ? final_state->ptr() : nullptr); + MI_CUDA_CHECK_LAUNCH(); + return; + } + constexpr int kPrepSmem = sizeof(PrepSmem), kStateSmem = sizeof(StateSmem); allow_dynamic_smem(prep_kernel, kPrepSmem); allow_dynamic_smem(state_kernel, kStateSmem); diff --git a/src/kernels/quant/int8_gemm.cu b/src/kernels/quant/int8_gemm.cu index 9bb2b1d..30a8394 100644 --- a/src/kernels/quant/int8_gemm.cu +++ b/src/kernels/quant/int8_gemm.cu @@ -5,6 +5,8 @@ #include #include "kernels/quant/int8_gemm_impl.cuh" +#include "kernels/sm75/int8_sm75.h" +#include "kernels/sm75/sm75.h" namespace mi::kernels::quant { namespace detail { @@ -148,6 +150,10 @@ 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 (kernels::sm75::turing()) { // Turing's 8x8x16 INT8 MMA (src/kernels/sm75/int8_gemm_sm75.cu) + kernels::sm75::int8_gemm(op, bias, out, ldo, stream); + return; + } detail::with_tile(detail::resolve_tile(tile, p), [&](auto cfg) { detail::plain(p, op, bias, out, ldo, stream); }); } diff --git a/src/kernels/quant/int8_gemm_impl.cuh b/src/kernels/quant/int8_gemm_impl.cuh index 5b01bbd..e82a257 100644 --- a/src/kernels/quant/int8_gemm_impl.cuh +++ b/src/kernels/quant/int8_gemm_impl.cuh @@ -289,12 +289,15 @@ struct Int8Gemm : cutlass::gemm::device::GemmUniversalBase { // The kernel: DefaultGemmWithVisitor (default_gemm_universal_with_visitor.h) with a B-operand permutation, without // the stream-K selection, in BandedSwizzle's tile order (GemmBase keeps CUTLASS's swizzle: only its Mma is used). -template +// ArchTag/InstrShape default to the sm_80 16x8x32 INT8 MMA; the sm_75 kernel family (src/kernels/sm75/int8_gemm_sm75.cu) +// instantiates the same template with arch::Sm75 and the 8x8x16 MMA. +template > struct Int8Kernel { using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal< std::int8_t, RowMajor, cutlass::ComplexTransform::kNone, kAlignAB, std::int8_t, ColumnMajor, cutlass::ComplexTransform::kNone, kAlignAB, bf16, RowMajor, std::int32_t, cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm80, typename Tile::Threadblock, typename Tile::Warp, InstructionShape, + ArchTag, typename Tile::Threadblock, typename Tile::Warp, InstrShape, cutlass::epilogue::thread::LinearCombination, Swizzle, Tile::stages, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, PermuteB>::GemmKernel; diff --git a/src/kernels/quant/int8_gemm_residual.cu b/src/kernels/quant/int8_gemm_residual.cu index 797cb39..7892872 100644 --- a/src/kernels/quant/int8_gemm_residual.cu +++ b/src/kernels/quant/int8_gemm_residual.cu @@ -1,5 +1,7 @@ // int8_gemm_gated_residual (see quant.h). #include "kernels/quant/int8_gemm_impl.cuh" +#include "kernels/sm75/int8_sm75.h" +#include "kernels/sm75/sm75.h" namespace mi::kernels::quant { namespace detail { @@ -36,6 +38,10 @@ 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 (kernels::sm75::turing()) { + kernels::sm75::int8_gemm_gated_residual(op, bias, gate, residual, ldr, stream); + return; + } detail::with_tile(detail::resolve_tile(tile, p), [&](auto cfg) { detail::gated_residual(p, 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 92d155e..424e443 100644 --- a/src/kernels/quant/int8_gemm_situ.cu +++ b/src/kernels/quant/int8_gemm_situ.cu @@ -5,6 +5,8 @@ #include #include "kernels/quant/int8_gemm_impl.cuh" +#include "kernels/sm75/int8_sm75.h" +#include "kernels/sm75/sm75.h" namespace mi::kernels::quant { namespace detail { @@ -102,6 +104,10 @@ void int8_gemm_situ_glu(const Int8Operands& op, __nv_bfloat16* out, std::int64_t std::to_string(ldo)); } if (p.m == 0) return; + if (kernels::sm75::turing()) { + kernels::sm75::int8_gemm_situ_glu(op, out, ldo, stream); + return; + } detail::with_tile(detail::resolve_tile(tile, p), [&](auto cfg) { detail::situ_glu(p, op, out, ldo, stream); }); } diff --git a/src/kernels/sm75/attention_sm75.cu b/src/kernels/sm75/attention_sm75.cu new file mode 100644 index 0000000..0ee6269 --- /dev/null +++ b/src/kernels/sm75/attention_sm75.cu @@ -0,0 +1,334 @@ +// Exact attention on sm_75: a tiled FlashAttention-2-style kernel written against Turing's fp16 tensor-core MMA +// (mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32). The vendored FlashAttention-2 has an sm_75 trait branch, but its +// kernel only produces correct results on the sm_80 MMA layouts (upstream FA2 targets Ampere+), so the sm_75 image is +// this kernel instead. +// +// Algorithm: one CTA per (sequence, head, 64-query-row block), 4 warps, each warp owning 16 query rows and the full head +// dim. Keys are walked in tiles (32 for head_dim 128, 16 for head_dim 256) with an online softmax: per-thread fp32 +// scores in the m16n8k8 accumulator layout, row max/sum reduced across the four lanes of a row group, and the P·V +// product reusing the normalized scores directly as the A operand (the m16n8k8 accumulator layout of Q·Kᵀ is exactly +// the A-operand layout of P·V for the next MMA). GQA maps query head h to KV head h / (heads / kv_heads); sequences are +// varlen through cu_seqlens. +// +// Precision: the engine's tensors are bf16 and sm_75's MMA is fp16. bf16 -> fp16 is exact for values in fp16's normal +// range and fp16 -> bf16 recovers the original value, so q/k/v are converted in place before the launch and restored +// after it; `out` is written as fp16 by the kernel and converted to bf16 in place. +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" + +namespace mi::kernels::sm75 { +namespace { + +constexpr int kThreads = 128; // 4 warps +constexpr int kWarps = 4; +constexpr int kBM = 64; // query rows per CTA (16 per warp) +constexpr float kLog2e = 1.4426950408889634f; + +// m16n8k8 fp16 -> fp32 MMA. a = {row group, row group + 8} 2-half registers, b = 2-half register, d = 4 fp32. +__device__ __forceinline__ void mma_m16n8k8(float (&d)[4], std::uint32_t a0, std::uint32_t a1, std::uint32_t b) { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(b)); +} + +__device__ __forceinline__ std::uint32_t pack_half2(float lo, float hi) { + union { + __half2 h; + std::uint32_t u; + } c; + c.h = __halves2half2(__float2half_rn(lo), __float2half_rn(hi)); + return c.u; +} + +__device__ __forceinline__ std::uint32_t load_u32(const void* p) { + return *static_cast(p); +} + +__device__ __forceinline__ void store_u32(void* p, std::uint32_t v) { + *static_cast(p) = v; +} + +struct Params { + const __half* q = nullptr; + const __half* k = nullptr; + const __half* v = nullptr; + __half* 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_sm75(const Params p) { + constexpr int kQStride = D + 8; + constexpr int kVStride = BN + 8; + constexpr int kSubtiles = BN / 8; + constexpr int kDimChunks = D / 8; + + extern __shared__ __align__(16) char smem_raw[]; + __half* qs = reinterpret_cast<__half*>(smem_raw); // [kBM][D] + __half* ks = qs + kBM * kQStride; // [BN][D] + __half* vt = ks + BN * kQStride; // [D][BN] + + const int batch = blockIdx.y; + const int head = blockIdx.z; + const int q_start = p.cu_q[batch]; // token offset of this sequence in q + const int q_end = p.cu_q[batch + 1]; + const int k_start = p.cu_k[batch]; + const int k_end = p.cu_k[batch + 1]; + const int len_q = q_end - q_start; + const int len_k = k_end - k_start; + const int row_base = static_cast(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); + const int tid = threadIdx.x; + const int lane = tid % 32; + const int warp = tid / 32; + const int group = lane >> 2; // 0..7: the row of the 8x8 fragment + const int tig = lane & 3; // 0..3: two consecutive columns per lane + + // Q tile into shared memory (rows past this sequence's end are zero). + for (int idx = tid; idx < kBM * D; idx += kThreads) { + const int r = idx / D, d = idx % D; + qs[r * kQStride + d] = (r < rows) ? p.q[(q_start + row_base + r) * p.q_row + head * p.q_head + d] + : __float2half_rn(0.f); + } + __syncthreads(); + + // This warp's Q fragment: rows {warp*16 + group, +8}, dims {i*8 + 2*tig, +1}. + std::uint32_t qf[kDimChunks][2]; + const int r0 = warp * 16 + group; + const int r1 = r0 + 8; +#pragma unroll + for (int i = 0; i < kDimChunks; ++i) { + qf[i][0] = load_u32(&qs[r0 * kQStride + i * 8 + tig * 2]); + qf[i][1] = load_u32(&qs[r1 * kQStride + i * 8 + tig * 2]); + } + + float o[kDimChunks][4]; // P·V accumulators, 4 floats per 8-dim chunk +#pragma unroll + for (int i = 0; i < kDimChunks; ++i) { + o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.f; + } + float m0 = -1e30f, m1 = -1e30f; // running row max for r0 / r1 + float l0 = 0.f, l1 = 0.f; // running row sum + + const int key_tiles = p.causal ? min((len_k + BN - 1) / BN, (row_base + rows + BN - 1) / BN) + : (len_k + BN - 1) / BN; + for (int tile = 0; tile < key_tiles; ++tile) { + const int j0 = tile * BN; + // K tile and transposed V tile into shared memory. + for (int idx = tid; idx < BN * D; idx += kThreads) { + const int j = idx / D, d = idx % D; + const int token = j0 + j; + const bool valid = token < len_k; + ks[j * kQStride + d] = valid ? p.k[(k_start + token) * p.k_row + kv_head * p.k_head + d] : __float2half_rn(0.f); + vt[d * kVStride + j] = valid ? p.v[(k_start + token) * p.v_row + kv_head * p.v_head + d] : __float2half_rn(0.f); + } + __syncthreads(); + + // S = Q·Kᵀ for every 8-key subtile of this tile (fp32, scaled and masked below). + float s[kSubtiles][4]; +#pragma unroll + for (int t = 0; t < kSubtiles; ++t) { + s[t][0] = s[t][1] = s[t][2] = s[t][3] = 0.f; +#pragma unroll + for (int i = 0; i < kDimChunks; ++i) { + const std::uint32_t b = load_u32(&ks[(t * 8 + group) * kQStride + i * 8 + tig * 2]); + mma_m16n8k8(s[t], qf[i][0], qf[i][1], b); + } + } + + // Scale, mask (keys past the sequence, causal), row max across the tile and the four lanes of the row group. + const int row_q0 = row_base + r0; + const int row_q1 = row_base + r1; + float tile_max0 = -1e30f, tile_max1 = -1e30f; +#pragma unroll + for (int t = 0; t < kSubtiles; ++t) { + const int key0 = j0 + t * 8 + tig * 2; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int key = key0 + (e & 1); + const int row_q = (e < 2) ? row_q0 : row_q1; + const bool valid = key < len_k && (!p.causal || key <= row_q); + s[t][e] = valid ? s[t][e] * p.scale : -1e30f; + } + tile_max0 = fmaxf(tile_max0, fmaxf(s[t][0], s[t][1])); + tile_max1 = fmaxf(tile_max1, fmaxf(s[t][2], s[t][3])); + } + tile_max0 = fmaxf(tile_max0, __shfl_xor_sync(0xffffffffu, tile_max0, 1)); + tile_max0 = fmaxf(tile_max0, __shfl_xor_sync(0xffffffffu, tile_max0, 2)); + tile_max1 = fmaxf(tile_max1, __shfl_xor_sync(0xffffffffu, tile_max1, 1)); + tile_max1 = fmaxf(tile_max1, __shfl_xor_sync(0xffffffffu, tile_max1, 2)); + + const float m_new0 = fmaxf(m0, tile_max0); + const float m_new1 = fmaxf(m1, tile_max1); + const float alpha0 = (m0 <= -1e29f) ? 0.f : exp2f((m0 - m_new0) * kLog2e); + const float alpha1 = (m1 <= -1e29f) ? 0.f : exp2f((m1 - m_new1) * kLog2e); + + // P = exp(S - m_new) as fp16, and the running sum. + float tile_sum0 = 0.f, tile_sum1 = 0.f; +#pragma unroll + for (int t = 0; t < kSubtiles; ++t) { + const float p0 = exp2f((s[t][0] - m_new0) * kLog2e); + const float p1 = exp2f((s[t][1] - m_new0) * kLog2e); + const float p2 = exp2f((s[t][2] - m_new1) * kLog2e); + const float p3 = exp2f((s[t][3] - m_new1) * kLog2e); + tile_sum0 += p0 + p1; + tile_sum1 += p2 + p3; + s[t][0] = p0; + s[t][1] = p1; + s[t][2] = p2; + s[t][3] = p3; + } + tile_sum0 += __shfl_xor_sync(0xffffffffu, tile_sum0, 1); + tile_sum0 += __shfl_xor_sync(0xffffffffu, tile_sum0, 2); + tile_sum1 += __shfl_xor_sync(0xffffffffu, tile_sum1, 1); + tile_sum1 += __shfl_xor_sync(0xffffffffu, tile_sum1, 2); + l0 = l0 * alpha0 + tile_sum0; + l1 = l1 * alpha1 + tile_sum1; + m0 = m_new0; + m1 = m_new1; + + // O = O * alpha + P·V. +#pragma unroll + for (int i = 0; i < kDimChunks; ++i) { + o[i][0] *= alpha0; + o[i][1] *= alpha0; + o[i][2] *= alpha1; + o[i][3] *= alpha1; + } +#pragma unroll + for (int t = 0; t < kSubtiles; ++t) { + const std::uint32_t pa0 = pack_half2(s[t][0], s[t][1]); + const std::uint32_t pa1 = pack_half2(s[t][2], s[t][3]); +#pragma unroll + for (int i = 0; i < kDimChunks; ++i) { + const std::uint32_t b = load_u32(&vt[(i * 8 + group) * kVStride + t * 8 + tig * 2]); + mma_m16n8k8(o[i], pa0, pa1, b); + } + } + __syncthreads(); + } + + // Normalize, then out (fp16) and the log-sum-exp workspace (fp32 [heads, total_q], unpadded). + const float inv0 = l0 > 0.f ? 1.f / l0 : 0.f; + const float inv1 = l1 > 0.f ? 1.f / l1 : 0.f; + if (r0 < rows) { +#pragma unroll + for (int i = 0; i < kDimChunks; ++i) { + store_u32(&p.o[(q_start + row_base + r0) * p.o_row + head * p.o_head + i * 8 + tig * 2], + pack_half2(o[i][0] * inv0, o[i][1] * inv0)); + } + if (tig == 0 && p.lse != nullptr) p.lse[static_cast(head) * p.total_q + q_start + row_base + r0] = + m0 + logf(l0); + } + if (r1 < rows) { +#pragma unroll + for (int i = 0; i < kDimChunks; ++i) { + store_u32(&p.o[(q_start + row_base + r1) * p.o_row + head * p.o_head + i * 8 + tig * 2], + pack_half2(o[i][2] * inv1, o[i][3] * inv1)); + } + if (tig == 0 && p.lse != nullptr) p.lse[static_cast(head) * p.total_q + q_start + row_base + r1] = + m1 + logf(l1); + } +} + +template +void launch(const Params& p, int max_seqlen, int batch, cudaStream_t stream) { + constexpr int kQStride = D + 8; + constexpr int kVStride = BN + 8; + constexpr int kSmem = (kBM * kQStride + BN * kQStride + D * kVStride) * static_cast(sizeof(__half)); + static_assert(kSmem <= 64 * 1024, "sm_75 allows 64 KiB of opt-in shared memory per block"); + auto kernel = &attention_kernel_sm75; + if constexpr (kSmem > 48 * 1024) { + MI_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmem)); + } + 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("sm75::flash_fwd: head_dim must be 128 or 256"); + } + if ((head_dim == 128 && causal) || (head_dim == 256 && !causal)) { + throw std::invalid_argument("sm75::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("sm75::flash_fwd: workspace too small for the log-sum-exp"); + } + + convert_tensor_bf16_to_f16_inplace(q, total_q, heads, head_dim, stream); + convert_tensor_bf16_to_f16_inplace(k, k.dim(0), kv_heads, head_dim, stream); + convert_tensor_bf16_to_f16_inplace(v, v.dim(0), kv_heads, head_dim, stream); + + Params p; + p.q = static_cast(q.data); + p.k = static_cast(k.data); + p.v = static_cast(v.data); + p.o = static_cast<__half*>(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, 16>(p, max_seqlen, batch, stream); + } + + convert_tensor_f16_to_bf16_inplace(out, out.dim(0), out.dim(1), out.dim(2), stream); + convert_tensor_f16_to_bf16_inplace(q, total_q, heads, head_dim, stream); + convert_tensor_f16_to_bf16_inplace(k, k.dim(0), kv_heads, head_dim, stream); + convert_tensor_f16_to_bf16_inplace(v, v.dim(0), kv_heads, head_dim, stream); +} + +} // namespace mi::kernels::sm75 diff --git a/src/kernels/sm75/convert.cu b/src/kernels/sm75/convert.cu new file mode 100644 index 0000000..5a332a1 --- /dev/null +++ b/src/kernels/sm75/convert.cu @@ -0,0 +1,206 @@ +// sm_75 dtype conversion kernels: every tensor-core operand in the engine is bf16, while Turing's tensor cores take +// fp16. These kernels convert in place (the same 2-byte slot) or into a separate buffer; the attention wrapper relies +// on the round trip being exact for in-range values (fp16 has more mantissa bits than bf16, so bf16 -> fp16 is lossless +// and fp16 -> bf16 recovers the original value as long as it stayed in fp16's normal range). +#include +#include + +#include +#include +#include + +#include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" + +namespace mi::kernels::sm75 { +namespace { + +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 65535; + +__global__ void bf16_to_f16_flat(const __nv_bfloat16* __restrict__ src, __half* __restrict__ dst, std::int64_t n) { + const std::int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (std::int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += stride) { + dst[i] = __float2half_rn(__bfloat162float(src[i])); + } +} + +__global__ void f16_to_bf16_flat(const __half* __restrict__ src, __nv_bfloat16* __restrict__ dst, std::int64_t n) { + const std::int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (std::int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += stride) { + dst[i] = __float2bfloat16_rn(__half2float(src[i])); + } +} + +// One block per token row, `dim` contiguous elements per (token, head) run. +__global__ void bf16_to_f16_row(const __nv_bfloat16* __restrict__ src, __half* __restrict__ dst, std::int64_t stride, + std::int64_t dim) { + const __nv_bfloat16* s = src + static_cast(blockIdx.x) * stride; + __half* d = dst + static_cast(blockIdx.x) * stride; + for (std::int64_t i = threadIdx.x; i < dim; i += blockDim.x) { + d[i] = __float2half_rn(__bfloat162float(s[i])); + } +} + +__global__ void f16_to_bf16_row(const __half* __restrict__ src, __nv_bfloat16* __restrict__ dst, std::int64_t stride, + std::int64_t dim) { + const __half* s = src + static_cast(blockIdx.x) * stride; + __nv_bfloat16* d = dst + static_cast(blockIdx.x) * stride; + for (std::int64_t i = threadIdx.x; i < dim; i += blockDim.x) { + d[i] = __float2bfloat16_rn(__half2float(s[i])); + } +} + +int blocks_for(std::int64_t n) { + const std::int64_t blocks = (n + kThreads - 1) / kThreads; + return static_cast(blocks < kMaxBlocks ? (blocks > 0 ? blocks : 1) : kMaxBlocks); +} + +// The attention tensors are [tokens, heads, dim] with a contiguous last dim, so every (token, head) pair is one run of +// `dim` elements at token * stride[0] + head * stride[1]. q/k/v may be views of a fused QKV output (non-contiguous +// heads), which the per-head loop handles. +void convert_3d_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, bool to_f16, + cudaStream_t stream) { + if (rows == 0 || heads == 0 || dim == 0) return; + if (rows > kMaxBlocks) return; // the engine's varlen token counts stay far below 2^31 blocks; never reached + auto* base = static_cast(t.data); + for (std::int64_t h = 0; h < heads; ++h) { + char* p = base + h * static_cast(t.stride[1]) * 2; + if (to_f16) { + bf16_to_f16_row<<(rows), kThreads, 0, stream>>>( + reinterpret_cast(p), reinterpret_cast<__half*>(p), t.stride[0], dim); + } else { + f16_to_bf16_row<<(rows), kThreads, 0, stream>>>( + reinterpret_cast(p), reinterpret_cast<__nv_bfloat16*>(p), t.stride[0], dim); + } + } + MI_CUDA_CHECK_LAUNCH(); +} + +} // namespace + +bool available() { return true; } + +bool needed(int major) { return major < 8; } + +bool fp16_weights() { + static std::atomic cached[64]; + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { + cudaGetLastError(); + return false; + } + if (device < 0 || device >= 64) return false; + int major = cached[device].load(std::memory_order_relaxed); + if (major == 0) { + if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) { + cudaGetLastError(); + return false; + } + cached[device].store(major, std::memory_order_relaxed); + } + return major < 8; +} + +void convert_weight_inplace(Tensor& w, cudaStream_t stream) { + if (w.data == nullptr) return; + if (!fp16_weights()) return; + if (w.dtype == DType::kF16) return; + if (w.dtype != DType::kBF16) throw std::logic_error("sm75::convert_weight_inplace: weight is not bf16"); + convert_bf16_to_f16(w.data, w.data, w.numel(), stream); + w.dtype = DType::kF16; +} + +void convert_weight_storage_inplace(Tensor& w, cudaStream_t stream) { + if (w.data == nullptr) return; + if (!fp16_weights()) return; + if (w.dtype != DType::kBF16 && w.dtype != DType::kF16) { + throw std::logic_error("sm75::convert_weight_storage_inplace: weight is not bf16 storage"); + } + convert_bf16_to_f16(w.data, w.data, w.numel(), stream); + w.dtype = DType::kF16; +} + +void convert_bf16_to_f16(void* dst, const void* src, std::int64_t n, cudaStream_t stream) { + if (n <= 0) return; + bf16_to_f16_flat<<>>(static_cast(src), + static_cast<__half*>(dst), n); + MI_CUDA_CHECK_LAUNCH(); +} + +void convert_f16_to_bf16(void* dst, const void* src, std::int64_t n, cudaStream_t stream) { + if (n <= 0) return; + f16_to_bf16_flat<<>>(static_cast(src), + static_cast<__nv_bfloat16*>(dst), n); + MI_CUDA_CHECK_LAUNCH(); +} + +void convert_tensor_bf16_to_f16_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, + cudaStream_t stream) { + convert_3d_inplace(t, rows, heads, dim, /*to_f16=*/true, stream); +} + +void convert_tensor_f16_to_bf16_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, + cudaStream_t stream) { + convert_3d_inplace(t, rows, heads, dim, /*to_f16=*/false, stream); +} + +namespace { + +__global__ void bf16_to_f16_rows(const __nv_bfloat16* __restrict__ src, __half* __restrict__ dst, std::int64_t cols, + std::int64_t src_stride, std::int64_t dst_stride) { + const __nv_bfloat16* s = src + static_cast(blockIdx.x) * src_stride; + __half* d = dst + static_cast(blockIdx.x) * dst_stride; + for (std::int64_t i = threadIdx.x; i < cols; i += blockDim.x) { + d[i] = __float2half_rn(__bfloat162float(s[i])); + } +} + +// out is bf16: `acc` may alias it (the GEMM writes fp16 into the output buffer, this pass converts in place). +__global__ void gemm_out_f16_bf16(__half* __restrict__ acc, __nv_bfloat16* __restrict__ out, + const __nv_bfloat16* __restrict__ bias, std::int64_t cols, std::int64_t out_stride, + std::int64_t acc_stride) { + const __half* a = acc + static_cast(blockIdx.x) * acc_stride; + __nv_bfloat16* o = out + static_cast(blockIdx.x) * out_stride; + for (std::int64_t c = threadIdx.x; c < cols; c += blockDim.x) { + const float v = __half2float(a[c]) + (bias != nullptr ? __bfloat162float(bias[c]) : 0.f); + o[c] = __float2bfloat16_rn(v); + } +} + +__global__ void gemm_out_f16_f32(const __half* __restrict__ acc, float* __restrict__ out, const float* __restrict__ bias, + std::int64_t cols, std::int64_t out_stride, std::int64_t acc_stride) { + const __half* a = acc + static_cast(blockIdx.x) * acc_stride; + float* o = out + static_cast(blockIdx.x) * out_stride; + for (std::int64_t c = threadIdx.x; c < cols; c += blockDim.x) { + o[c] = __half2float(a[c]) + (bias != nullptr ? bias[c] : 0.f); + } +} + +} // namespace + +void convert_rows_bf16_to_f16(void* dst, const void* src, std::int64_t rows, std::int64_t cols, + std::int64_t src_stride, std::int64_t dst_stride, cudaStream_t stream) { + if (rows <= 0 || cols <= 0) return; + bf16_to_f16_rows<<(rows), kThreads, 0, stream>>>( + static_cast(src), static_cast<__half*>(dst), cols, src_stride, dst_stride); + MI_CUDA_CHECK_LAUNCH(); +} + +void convert_gemm_out_f16(void* out, const void* acc, const void* bias, bool bias_f32, bool out_f32, + std::int64_t rows, std::int64_t cols, std::int64_t out_stride, std::int64_t acc_stride, + cudaStream_t stream) { + if (rows <= 0 || cols <= 0) return; + if (out_f32) { + gemm_out_f16_f32<<(rows), kThreads, 0, stream>>>( + static_cast(acc), static_cast(out), static_cast(bias_f32 ? bias : nullptr), + cols, out_stride, acc_stride); + } else { + gemm_out_f16_bf16<<(rows), kThreads, 0, stream>>>( + static_cast<__half*>(const_cast(acc)), static_cast<__nv_bfloat16*>(out), + static_cast(bias), cols, out_stride, acc_stride); + } + MI_CUDA_CHECK_LAUNCH(); +} + +} // namespace mi::kernels::sm75 diff --git a/src/kernels/sm75/int8_gemm_sm75.cu b/src/kernels/sm75/int8_gemm_sm75.cu new file mode 100644 index 0000000..1a8dc71 --- /dev/null +++ b/src/kernels/sm75/int8_gemm_sm75.cu @@ -0,0 +1,103 @@ +// sm_75 INT8 ConvRot GEMMs (see kernels/sm75/int8_sm75.h): the sm_80 CUTLASS EVT kernel with Turing's m8n8k16 INT8 +// MMA and tiles that fit 64 KiB of shared memory. +#include "kernels/sm75/int8_sm75.h" + +#include + +#include "kernels/quant/int8_gemm_impl.cuh" + +namespace mi::kernels::sm75 { +namespace qd = quant::detail; + +namespace { + +using Sm75Tile = qd::TileCfg<64, 128, 64, 32, 64, 2>; // threadblock 64x128, warp 32x64, K 64, 2 stages: 24 KiB +using Sm75Instr = cutlass::gemm::GemmShape<8, 8, 16>; + +template +using Sm75Kernel = typename qd::Int8Kernel::Gemm; + +void plain(const qd::Problem& p, const quant::Int8Operands& a, const __nv_bfloat16* bias, __nv_bfloat16* out, + std::int64_t ldo, cudaStream_t stream) { + using N = qd::Nodes; + using Gemm = Sm75Kernel; + using namespace cute; + const typename N::Plain::Arguments cb{ + {{}, + {const_cast(a.x_scale), 0.f, {_1{}, _0{}, p.m}}, + {a.w_scale, 0.f, {_0{}, _1{}, p.n}}, + {reinterpret_cast(bias), qd::bf16(0), {_0{}, _1{}, p.n}}, + {}}, + {reinterpret_cast(out), {ldo, _1{}, ldo * p.m}}}; + qd::launch("sm75::int8_gemm", p, a, cb, stream); +} + +void gated_residual(const qd::Problem& p, const quant::Int8Operands& a, const __nv_bfloat16* bias, + const __nv_bfloat16* gate, __nv_bfloat16* residual, std::int64_t ldr, cudaStream_t stream) { + using N = qd::Nodes; + using Gemm = Sm75Kernel; + using namespace cute; + auto* r = reinterpret_cast(residual); + const typename N::GatedResidual::Arguments cb{ + {{}, + {const_cast(a.x_scale), 0.f, {_1{}, _0{}, p.m}}, + {a.w_scale, 0.f, {_0{}, _1{}, p.n}}, + {reinterpret_cast(bias), qd::bf16(0), {_0{}, _1{}, p.n}}, + {reinterpret_cast(gate), qd::bf16(0), {_0{}, _1{}, p.n}}, + {r, qd::bf16(0), {ldr, _1{}, ldr * p.m}}, + {}}, + {r, {ldr, _1{}, ldr * p.m}}}; + qd::launch("sm75::int8_gemm_gated_residual", p, a, cb, stream); +} + +void situ_glu(const qd::Problem& p, const quant::Int8Operands& a, __nv_bfloat16* out, std::int64_t ldo, + cudaStream_t stream) { + using N = qd::Nodes; + using Gemm = Sm75Kernel; + using namespace cute; + auto* d = reinterpret_cast(out); + const std::int64_t ldd = 2 * ldo; + const typename N::SituGlu::Arguments cb{ + {{}, {const_cast(a.x_scale), 0.f, {_1{}, _0{}, p.m}}, {a.w_scale}, {}}, + {d, {ldd, _1{}, ldd * p.m}}}; + qd::launch("sm75::int8_gemm_situ_glu", p, a, cb, stream); +} + +} // 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 = "sm75::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; + plain(p, op, bias, 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 = "sm75::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; + gated_residual(p, 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 = "sm75::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; + situ_glu(p, op, out, ldo, stream); +} + +} // namespace mi::kernels::sm75 diff --git a/src/kernels/sm75/int8_sm75.h b/src/kernels/sm75/int8_sm75.h new file mode 100644 index 0000000..ecb7292 --- /dev/null +++ b/src/kernels/sm75/int8_sm75.h @@ -0,0 +1,25 @@ +// sm_75 INT8 ConvRot GEMMs: the sm_80 kernels use m16n8k32 INT8 MMA (mma.sync.s8, sm_80 only). Turing's native INT8 +// shape is m8n8k16, so the same CUTLASS 2.x EVT kernel template is instantiated with arch::Sm75 and GemmShape<8,8,16>, +// with threadblock/warp tiles that fit Turing's 64 KiB of shared memory. The epilogues (dequantize + bias, gated +// residual, SiTU-GLU) are the sm_80 ones unchanged (they are arch-neutral visitor trees). +#pragma once + +#include +#include + +#include + +#include "core/tensor.h" +#include "kernels/quant/quant.h" + +namespace mi::kernels::sm75 { + +// Same contracts as kernels/quant/quant.h's int8_gemm / int8_gemm_gated_residual / int8_gemm_situ_glu (operand checks, +// output layout and numerics); the host-side dispatcher in those translation units routes compute capability 7.x here. +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::sm75 diff --git a/src/kernels/sm75/sm75.h b/src/kernels/sm75/sm75.h new file mode 100644 index 0000000..f32c031 --- /dev/null +++ b/src/kernels/sm75/sm75.h @@ -0,0 +1,89 @@ +// Turing (sm_75) kernel family: the pieces of the engine that need bf16 tensor cores, cp.async, TF32 or > 64 KiB of +// shared memory on sm_80+ have an sm_75 fallback here, and the host-side entry points of the sm_80 kernels dispatch to +// them when the current device is compute capability 7.x. +// +// This translation unit family is compiled for sm_75 only (cmake/modules/sm75.cmake, CUDA_ARCHITECTURES 75-real, in +// its own OBJECT library), so it may use half-precision MMA and must not assume any sm_80 instruction. The host code +// that dispatches into it lives in the sm_80 translation units and is arch-independent. +// +// Numerics: sm_75 has fp16 tensor cores but no bf16 ones, so every tensor-core operand is converted bf16 -> fp16 at +// the boundary (lossless for |x| in fp16's normal range; values outside it saturate/round, which the engine reports +// as the sm_75 numerics). Accumulation stays fp32 where the sm_80 kernel had fp32 accumulation. +#pragma once + +#include + +#include +#include + +#include "core/tensor.h" + +namespace mi::kernels::sm75 { + +// True when this build carries the sm_75 kernels. +bool available(); + +// True when `major` needs them (compute capability < 8). +bool needed(int major); + +// True when bf16 GEMM weights must be stored fp16: this build has the sm_75 kernels and the current device is compute +// capability 7.x (the fp16 GEMM path in kernels/basic/gemm.cpp). Cached per device. +bool fp16_weights(); + +// True when the current device runs the sm_75 kernel family (compute capability 7.x): the host dispatchers in the sm_80 +// translation units branch on this. +inline bool turing() { return fp16_weights(); } + +// Converts a bf16 GEMM weight to fp16 in place and marks the view kF16. No-op when fp16_weights() is false or the +// tensor is already fp16; throws std::logic_error for any other dtype. For resident weights (converted once). +void convert_weight_inplace(Tensor& w, cudaStream_t stream); + +// Same conversion, but always runs: streamed weight buffers are refilled with bf16 bytes on every use while their view +// keeps the kF16 label from the first conversion. +void convert_weight_storage_inplace(Tensor& w, cudaStream_t stream); + +// Elements of one contiguous bf16 -> fp16 conversion. `dst` may equal `src` (an in-place element-wise conversion). +void convert_bf16_to_f16(void* dst, const void* src, std::int64_t n, cudaStream_t stream); +void convert_f16_to_bf16(void* dst, const void* src, std::int64_t n, cudaStream_t stream); + +// bf16 -> fp16 over a 3-D tensor [rows, heads, dim] with free row/head strides, in place: every element is replaced by +// its fp16 bit pattern in the same 2-byte slot. Used by the sm_75 attention, which needs fp16 operands and restores +// bf16 after the launch so callers that reuse q/k/v are unaffected. +void convert_tensor_bf16_to_f16_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, + cudaStream_t stream); +void convert_tensor_f16_to_bf16_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, + cudaStream_t stream); + +// bf16 [rows, cols] (free row strides) -> fp16, into `dst` (dst_stride elements) or in place when dst == src. +void convert_rows_bf16_to_f16(void* dst, const void* src, std::int64_t rows, std::int64_t cols, + std::int64_t src_stride, std::int64_t dst_stride, cudaStream_t stream); + +// GEMM output conversion: `acc` is fp16 [rows, cols] (acc_stride elements; may alias `out` when `out` is bf16), +// `bias` is nullptr or fp32/bf16 [cols] (bias_f32 selects), `out` is bf16 or f32 [rows, cols] (out_stride). +// out[r][c] = out_type(float(acc[r][c]) + (bias ? bias[c] : 0)), the fp32 add before the final rounding. +void convert_gemm_out_f16(void* out, const void* acc, const void* bias, bool bias_f32, bool out_f32, + std::int64_t rows, std::int64_t cols, std::int64_t out_stride, std::int64_t acc_stride, + cudaStream_t stream); + +// VAE fp32 fallbacks (the reference decoder's kernels need TF32 tensor cores and cp.async). +// conv2d: y = conv(x) + bias + optional residual, x/y NHWC fp32 [N,H,W,C], w KRSC fp32 [K,R,R,C], stride 1, pad R/2, +// R in {1,3}. No workspace. +void vae_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, + cudaStream_t stream); +// Mid-block attention: o = softmax(q k^T * scale) v, fp32, one head, [B, L, D] with D a multiple of 32 and any row +// stride (>= D) for q/k/v; o is contiguous. +void vae_attention_fp32(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& o, float scale, + cudaStream_t stream); +// torch conv weight [K,C,R,S] fp32 -> [K,R,S,C] in dst's dtype (kF32 or kBF16). +void vae_weight_to_krsc(const Tensor& w_kcrs, const Tensor& w_krsc, cudaStream_t stream); + +// Exact (non-approximate) attention, the sm_75 counterpart of attention::flash_fwd for the same arguments and +// supported configurations (head_dim 128 non-causal for the DiT, head_dim 256 causal for the text encoder). q, k, v +// and out are bf16; they are converted to fp16 in place around the launch and restored to bf16 (the values are +// recoverable exactly for in-range inputs: fp16 has more mantissa bits than bf16). The kernel itself is the vendored +// FlashAttention-2 with its own sm_75 traits (fp16 m16n8k8 MMA, no cp.async). +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); + +} // namespace mi::kernels::sm75 diff --git a/src/kernels/sm75/vae_sm75.cu b/src/kernels/sm75/vae_sm75.cu new file mode 100644 index 0000000..02f1ff6 --- /dev/null +++ b/src/kernels/sm75/vae_sm75.cu @@ -0,0 +1,180 @@ +// sm_75 VAE kernels: the reference decoder's convolutions and mid-block attention are TF32 tensor-core kernels (CUTLASS +// implicit GEMM with cp.async, and a TF32 m16n8k8 FlashAttention), neither of which exists on Turing. These fp32 SIMT +// replacements keep the reference precision (fp32 accumulate, fp32 softmax) at Turing speeds: the VAE is a small part +// of a generation, so plain fp32 is the right trade for a first working path. +#include +#include + +#include +#include + +#include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" + +namespace mi::kernels::sm75 { +namespace { + +// ---------------------------------------------------------------- conv2d (NHWC, KRSC filters, stride 1, pad R/2) +// One block: an 8x8 output tile of 32 output channels. Input channels are walked in groups of 8 through shared memory +// (the input patch and the weights they multiply), each thread accumulating 2 pixels x 16 channels. +constexpr int kTW = 8, kTH = 8, kTK = 32, kTC = 8; +constexpr int kPX = kTW + 2, kPY = kTH + 2; +constexpr int kConvThreads = 128; + +__global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __restrict__ w, + const float* __restrict__ bias, const float* __restrict__ residual, + float* __restrict__ y, int H, int W, int C, int K, int R, int kb_count) { + __shared__ float xs[kPY * kPX * kTC]; + __shared__ float ws[kTK * 9 * kTC]; + const int n = blockIdx.z / kb_count; + const int kb = blockIdx.z % kb_count; + const int ox0 = blockIdx.x * kTW, oy0 = blockIdx.y * kTH; + const int pad = (R - 1) / 2; + const int tid = threadIdx.x; + const int pixel = tid >> 1; // 0..63 + const int cg = (tid & 1) * 16; // first / second 16 output channels + const int py = pixel / kTW, px = pixel % kTW; + float acc[16] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + + for (int c0 = 0; c0 < C; c0 += kTC) { + for (int i = tid; i < kPY * kPX * kTC; i += kConvThreads) { + const int cc = i % kTC, p = i / kTC; + const int yy = p / kPX, xx = p % kPX; + const int gy = oy0 - pad + yy, gx = ox0 - pad + xx, gc = c0 + cc; + xs[i] = (gy >= 0 && gy < H && gx >= 0 && gx < W && gc < C) + ? x[((static_cast(n) * H + gy) * W + gx) * C + gc] + : 0.f; + } + for (int i = tid; i < kTK * 9 * kTC; i += kConvThreads) { + const int cc = i % kTC, t = i / kTC, rs = t % 9, kk = t / 9; + const int kch = kb * kTK + kk, gc = c0 + cc; + ws[i] = (kch < K && gc < C) + ? w[(static_cast(kch) * R + rs / R) * R * C + (rs % R) * C + gc] + : 0.f; + } + __syncthreads(); + for (int cc = 0; cc < kTC; ++cc) { + if (c0 + cc >= C) break; + for (int r = 0; r < R; ++r) { + for (int s = 0; s < R; ++s) { + const float xv = xs[((py + r) * kPX + (px + s)) * kTC + cc]; +#pragma unroll + for (int ch = 0; ch < 16; ++ch) acc[ch] += xv * ws[((cg + ch) * 9 + r * R + s) * kTC + cc]; + } + } + } + __syncthreads(); + } + + const int ox = ox0 + px, oy = oy0 + py; + if (ox >= W || oy >= H) return; +#pragma unroll + for (int ch = 0; ch < 16; ++ch) { + const int kch = kb * kTK + cg + ch; + if (kch >= K) break; + const std::int64_t idx = ((static_cast(n) * H + oy) * W + ox) * K + kch; + float v = acc[ch] + bias[kch]; + if (residual != nullptr) v += residual[idx]; + y[idx] = v; + } +} + +// ---------------------------------------------------------------- mid-block attention (single head, fp32) +// One warp per query row: the lane holds D/32 elements of q, v and the output; the score is a warp reduction and the +// softmax runs online, so the L x L matrix never exists. D is a multiple of 32 (384 or 512). +constexpr int kAttnWarps = 4; + +__global__ void attention_f32_sm75(const float* __restrict__ q, const float* __restrict__ k, + const float* __restrict__ v, float* __restrict__ o, int L, int D, int q_rows, + int k_rows, int v_rows, std::int64_t q_batch, std::int64_t k_batch, + std::int64_t v_batch, std::int64_t o_batch, float scale) { + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int row = blockIdx.x * kAttnWarps + warp; + if (row >= L) return; + const int nlanes = D / 32; + const int b = blockIdx.y; + const float* qrow = q + static_cast(b) * q_batch + static_cast(row) * q_rows; + float qreg[16]; + float oacc[16]; + for (int i = 0; i < nlanes; ++i) { + qreg[i] = qrow[lane + 32 * i]; + oacc[i] = 0.f; + } + float m = -1e30f, s = 0.f; + for (int j = 0; j < L; ++j) { + const float* krow = k + static_cast(b) * k_batch + static_cast(j) * k_rows; + float dot = 0.f; + for (int i = 0; i < nlanes; ++i) dot += qreg[i] * krow[lane + 32 * i]; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) dot += __shfl_xor_sync(0xffffffffu, dot, off); + dot *= scale; + const float mn = fmaxf(m, dot); + const float alpha = (m > -1e29f) ? expf(m - mn) : 0.f; + const float p = expf(dot - mn); + s = s * alpha + p; + const float* vrow = v + static_cast(b) * v_batch + static_cast(j) * v_rows; + for (int i = 0; i < nlanes; ++i) oacc[i] = oacc[i] * alpha + p * vrow[lane + 32 * i]; + m = mn; + } + const float inv = 1.f / s; + float* orow = o + static_cast(b) * o_batch + static_cast(row) * D; + for (int i = 0; i < nlanes; ++i) orow[lane + 32 * i] = oacc[i] * inv; +} + +} // namespace + +void vae_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, + cudaStream_t stream) { + const int H = static_cast(x.dim(1)), W = static_cast(x.dim(2)), C = static_cast(x.dim(3)); + const int K = static_cast(w.dim(0)), R = static_cast(w.dim(1)); + const int kb_count = (K + kTK - 1) / kTK; + const dim3 grid(static_cast((W + kTW - 1) / kTW), static_cast((H + kTH - 1) / kTH), + static_cast(x.dim(0) * kb_count)); + conv2d_f32_sm75<<>>( + x.ptr(), w.ptr(), bias.ptr(), residual != nullptr ? residual->ptr() : nullptr, + y.ptr(), H, W, C, K, R, kb_count); + MI_CUDA_CHECK_LAUNCH(); +} + +void vae_attention_fp32(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& o, float scale, + cudaStream_t stream) { + const int L = static_cast(q.dim(1)), D = static_cast(q.dim(2)); + const int batch = static_cast(q.dim(0)); + const dim3 grid(static_cast((L + kAttnWarps - 1) / kAttnWarps), static_cast(batch)); + attention_f32_sm75<<>>(q.ptr(), k.ptr(), v.ptr(), + o.ptr(), L, D, static_cast(q.stride[1]), + static_cast(k.stride[1]), + static_cast(v.stride[1]), q.stride[0], k.stride[0], + v.stride[0], o.stride[0], scale); + MI_CUDA_CHECK_LAUNCH(); +} + +template +__global__ void kcrs_to_krsc_sm75(const float* __restrict__ src, Out* __restrict__ dst, int k, int c, int rs) { + const std::int64_t total = static_cast(k) * c * rs; + for (std::int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; i < total; + i += static_cast(gridDim.x) * blockDim.x) { + const int ci = static_cast(i % c); + const std::int64_t t = i / c; + const int p = static_cast(t % rs); + const std::int64_t ki = t / rs; + dst[i] = static_cast(src[(ki * c + ci) * rs + p]); + } +} + +void vae_weight_to_krsc(const Tensor& w_kcrs, const Tensor& w_krsc, cudaStream_t stream) { + const int k = static_cast(w_kcrs.dim(0)), c = static_cast(w_kcrs.dim(1)); + const int rs = static_cast(w_kcrs.dim(2) * w_kcrs.dim(3)); + const std::int64_t total = w_kcrs.numel(); + if (total == 0) return; + const int blocks = static_cast(std::min((total + 255) / 256, 4096)); + if (w_krsc.dtype == DType::kF32) { + kcrs_to_krsc_sm75<<>>(w_kcrs.ptr(), w_krsc.ptr(), k, c, rs); + } else { + kcrs_to_krsc_sm75<__nv_bfloat16><<>>(w_kcrs.ptr(), w_krsc.ptr<__nv_bfloat16>(), k, + c, rs); + } + MI_CUDA_CHECK_LAUNCH(); +} + +} // namespace mi::kernels::sm75 diff --git a/src/kernels/vae/attention.cu b/src/kernels/vae/attention.cu index 3f05908..86491f0 100644 --- a/src/kernels/vae/attention.cu +++ b/src/kernels/vae/attention.cu @@ -22,6 +22,7 @@ #include #include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" #include "kernels/vae/attention.h" namespace mi::vae { @@ -477,6 +478,12 @@ void attention_fp32(const Tensor& q, const Tensor& k, const Tensor& v, const Ten require(q.dim(1) < (1ll << 31) && q.dim(0) <= 65535, "attention_fp32: L or B too large"); const int len = static_cast(q.dim(1)); if (len == 0 || q.dim(0) == 0) return; + // sm_75: no TF32 tensor cores; the fp32 SIMT kernel is exact in fp32 (the reference precision) and the VAE's + // attention is a small part of a generation. + if (kernels::sm75::fp16_weights()) { + kernels::sm75::vae_attention_fp32(q, k, v, o, scale, stream); + return; + } AttnArgs a{q.ptr(), k.ptr(), v.ptr(), o.ptr(), q.stride[0], k.stride[0], v.stride[0], o.stride[0], q.stride[1], k.stride[1], v.stride[1], len, scale * 1.4426950408889634f}; const dim3 grid(static_cast((len + kBr - 1) / kBr), static_cast(q.dim(0))); diff --git a/src/kernels/vae/conv2d.cu b/src/kernels/vae/conv2d.cu index 12f478f..e45ba25 100644 --- a/src/kernels/vae/conv2d.cu +++ b/src/kernels/vae/conv2d.cu @@ -8,6 +8,7 @@ #include #include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" #include "kernels/vae/vae_kernels.h" #include "cutlass/conv/conv2d_problem_size.h" @@ -173,6 +174,7 @@ auto dispatch(const ConvShape& s, DType t, Fn&& fn) { std::size_t conv2d_workspace_bytes(const Tensor& x, const Tensor& w) { const ConvShape s = check_shapes(x, w); + if (x.dtype == DType::kF32 && kernels::sm75::fp16_weights()) return 0; // the sm_75 fp32 kernel needs none return dispatch(s, x.dtype, [&]() { return workspace_for(s, x, w); }); } @@ -190,6 +192,11 @@ void conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* check_out(y, "y"); if (residual) check_out(*residual, "residual"); 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)"); + kernels::sm75::vae_conv2d(x, w, bias, residual, y, stream); + return; + } dispatch(s, x.dtype, [&]() { run(s, x, w, bias, residual, y, workspace, workspace_bytes, stream); }); } @@ -227,6 +234,10 @@ void conv_weight_to_krsc(const Tensor& w_kcrs, const Tensor& w_krsc, cudaStream_ if (w_krsc.dim(0) != k || w_krsc.dim(1) != r || w_krsc.dim(2) != s || w_krsc.dim(3) != c) throw std::invalid_argument("conv_weight_to_krsc: shape mismatch " + w_kcrs.describe() + " -> " + w_krsc.describe()); const std::int64_t total = w_kcrs.numel(); + if (kernels::sm75::fp16_weights()) { + kernels::sm75::vae_weight_to_krsc(w_kcrs, w_krsc, stream); + return; + } const int blocks = static_cast(std::min((total + 255) / 256, 4096)); if (total == 0) return; if (w_krsc.dtype == DType::kF32) { diff --git a/src/models/moss_dit/dit_weights.cpp b/src/models/moss_dit/dit_weights.cpp index 8805005..9659e03 100644 --- a/src/models/moss_dit/dit_weights.cpp +++ b/src/models/moss_dit/dit_weights.cpp @@ -8,6 +8,7 @@ #include #include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" #include "models/moss_dit/moss_dit.h" #include "weights/quant/nvfp4_codec.h" @@ -456,6 +457,31 @@ BlockWeights::BlockWeights(const weights::ModelIndex& index, runtime::DeviceAren BlockWeights::~BlockWeights() { release(); } +namespace { +// sm_75: the bf16 linear weights of one block, converted to fp16 in place (the views are marked kF16). The quantized +// linears (qkv8/... , qkv_nvfp4/...) are not bf16 and are left alone; the fused bf16 qkv view covers its parts. +void convert_linear_to_f16(Tensor& t, cudaStream_t stream) { + if (t.data == nullptr) return; + kernels::sm75::convert_weight_storage_inplace(t, stream); +} + +void convert_block_to_f16(BlockTensors& b, cudaStream_t stream) { + convert_linear_to_f16(b.qkv, stream); + convert_linear_to_f16(b.o, stream); + convert_linear_to_f16(b.up_gate, stream); + convert_linear_to_f16(b.down, stream); + for (Tensor& t : b.qkv_part) convert_linear_to_f16(t, stream); +} +} // namespace + +void BlockWeights::convert_weights_to_f16(cudaStream_t stream) { + if (!kernels::sm75::fp16_weights()) return; + for (int i = 0; i < kNumBlocks; ++i) { + if (!resident(i)) continue; + convert_block_to_f16(resident_[static_cast(i)], stream); + } +} + void BlockWeights::create_streaming_resources() { const auto n = static_cast(reader_count(residency_)); const std::size_t s = streamed_.size(); @@ -669,6 +695,9 @@ const BlockTensors& BlockWeights::Pass::acquire(int block) { w_.sync_->wait_enqueued(j + 1); for (const auto& r : w_.rd_) MI_CUDA_CHECK(cudaStreamWaitEvent(stream_, r->copy_done[b], 0)); MI_CUDA_CHECK(cudaEventRecord(w_.wait_end_[static_cast(j)], stream_)); + // sm_75: the block's bf16 linear weights just landed; convert them to fp16 on the compute stream, before the block's + // kernels read them (the copy itself moves the stored bf16 bytes, so this runs for every refill). + if (kernels::sm75::fp16_weights()) convert_block_to_f16(views_[static_cast(j)], stream_); return views_[static_cast(j)]; } diff --git a/src/models/moss_dit/dit_weights.h b/src/models/moss_dit/dit_weights.h index f27f3b1..d34dcf1 100644 --- a/src/models/moss_dit/dit_weights.h +++ b/src/models/moss_dit/dit_weights.h @@ -209,6 +209,11 @@ class BlockWeights { bool resident(int block) const { return mask_.test(static_cast(block)); } int streamed_blocks() const { return static_cast(streamed_.size()); } std::size_t host_resident_bytes() const; // pinned host copy held (kHostPinned) + + // sm_75 (kernels::sm75::fp16_weights()): converts the resident blocks' bf16 linear weights to fp16 in place, so the + // fp16 GEMM path can consume them (the views are marked kF16). Streamed blocks are converted by Pass::acquire(), + // after each block's copy. No-op on sm_80+ devices and when the build has no sm_75 kernels. + void convert_weights_to_f16(cudaStream_t stream); void set_copy_hook(CopyHook hook) { hook_ = std::move(hook); } // Stats of the last finished forward() (synchronizes on its events the first time it is read) and the sum over diff --git a/src/models/moss_dit/moss_dit.cpp b/src/models/moss_dit/moss_dit.cpp index c602a02..2695b84 100644 --- a/src/models/moss_dit/moss_dit.cpp +++ b/src/models/moss_dit/moss_dit.cpp @@ -17,6 +17,7 @@ #include "kernels/basic/elementwise.h" #include "kernels/basic/norm.h" #include "kernels/basic/rope.h" +#include "kernels/sm75/sm75.h" #include "models/lora/lora_merge_device.h" #include "models/moss_dit/dit_kernels.h" #include "models/moss_dit/nvfp4_support.h" @@ -812,6 +813,20 @@ MossDit::MossDit(const weights::ModelIndex& index, const quant::QuantIndex* quan if (!merge_on_device(*merges, module, t)) merge_on_host(module, t); merge_targets_.emplace_back(module, t); } + // sm_75: every bf16 linear weight becomes fp16 in place (the fp16 GEMM path in kernels/basic/gemm.cpp). Streamed + // block buffers are converted by BlockWeights::Pass::acquire() after each copy. One-time cost at load, after the + // pre-merged adapters above so they still run their bf16 kernels. + if (kernels::sm75::fp16_weights()) { + for (Linear* l : {&video_embedder_, &time_in_, &time_out_, &block_modulation_, &head_modulation_, &text_proj_}) { + kernels::sm75::convert_weight_inplace(l->w, nullptr); + } + kernels::sm75::convert_weight_inplace(final_linear_, nullptr); + for (TextBlock& b : text_blocks_) { + for (Linear* l : {&b.qkv, &b.proj, &b.up_gate, &b.down}) kernels::sm75::convert_weight_inplace(l->w, nullptr); + } + blocks_->convert_weights_to_f16(nullptr); + MI_CUDA_CHECK(cudaStreamSynchronize(nullptr)); + } load_stats_.direct = readers.all_direct(); for (cudaEvent_t& e : block_done_) MI_CUDA_CHECK(cudaEventCreateWithFlags(&e, cudaEventDisableTiming)); load_stats_.seconds = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); diff --git a/src/models/qwen35_text/text_encoder.cpp b/src/models/qwen35_text/text_encoder.cpp index fa21411..fd45ad8 100644 --- a/src/models/qwen35_text/text_encoder.cpp +++ b/src/models/qwen35_text/text_encoder.cpp @@ -21,6 +21,7 @@ #include "kernels/basic/elementwise.h" #include "kernels/basic/gemm.h" #include "kernels/basic/norm.h" +#include "kernels/sm75/sm75.h" #include "kernels/basic/rope.h" #include "kernels/gdn/gdn.h" #include "kernels/quant/quant.h" @@ -738,6 +739,21 @@ void TextEncoder::Impl::reader_loop(std::size_t r, const std::vector& chunks = me.ring ? me.ring->stats().chunks - start.chunks : 0; } +namespace { +// sm_75: the layer's bf16 GEMM weights (the rank-2 tensors: q/k/v, o, mlp and GDN projections) are converted to fp16 +// in place, for the fp16 GEMM path in kernels/basic. Norms (rank 1), A_log/dt_bias (rank 1) and the conv1d kernel +// (rank 3) stay bf16. Always converts: the streamed buffers are refilled with the stored bf16 bytes on every layer. +void convert_linears_f16(weights::TensorMap& w, cudaStream_t stream) { + if (!kernels::sm75::fp16_weights()) return; + for (auto& entry : w) { + Tensor& t = entry.second; + if (t.rank == 2 && (t.dtype == DType::kBF16 || t.dtype == DType::kF16)) { + kernels::sm75::convert_weight_storage_inplace(t, stream); + } + } +} +} // namespace + // P2 path: load layer k, run it, synchronize, move on. void TextEncoder::Impl::run_layers_sync(const Batch& b, void* weight_buf, const Tensor& x, const Tensor& h, const Tensor& y, const EncodeOptions& options) { @@ -750,6 +766,7 @@ void TextEncoder::Impl::run_layers_sync(const Batch& b, void* weight_buf, const if (!aux.empty()) w.insert(aux[static_cast(i)].begin(), aux[static_cast(i)].end()); stats.weight_load_s += seconds_since(t0); stats.layer_bytes[static_cast(i)] = l.group->total_bytes; + convert_linears_f16(w, stream); t0 = Clock::now(); run_layer(l, w, b, x, h, y); if (options.layer_hook) options.layer_hook(i, x, stream); @@ -764,8 +781,7 @@ void TextEncoder::Impl::run_layers_sync(const Batch& b, void* weight_buf, const void TextEncoder::Impl::run_layers_streaming(const std::vector>& sequences, const Batch& b, const std::vector& bufs, const Tensor& x, const Tensor& h, const Tensor& y, const EncodeOptions& options) { - const std::vector plans = plan_layers(bufs); - const std::size_t nr = rd.size(); + std::vector plans = plan_layers(bufs); const std::size_t nr = rd.size(); std::vector> layer_bytes(nr, std::vector(kNumLayers, 0)); std::vector chunks(nr, 0); StreamSync sync(static_cast(nr)); @@ -806,6 +822,7 @@ void TextEncoder::Impl::run_layers_streaming(const std::vector(k)].w, stream); MI_CUDA_CHECK(cudaEventRecord(timing[static_cast(kEvLayerStart + k)], stream)); run_layer(layers[static_cast(k)], plans[static_cast(k)].w, b, x, h, y); if (options.layer_hook) options.layer_hook(k, x, stream); diff --git a/src/pipeline/pipeline.cpp b/src/pipeline/pipeline.cpp index 3e25722..84ee467 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -939,9 +939,15 @@ Pipeline::Impl::CallModels Pipeline::Impl::setup_models(const Request& req, Resu // the LoRA files too). The lossy accelerators are an explicit speed opt-in: the tier is resolved without them. const std::filesystem::path te_file = tes.model ? tes.path : std::filesystem::path(); const auto supported = [](planner::AttentionMode mode) { return moss_dit::attention_supported(dit_attention(mode)); }; + // sm_75 (and any device without an 8-bit attention kernel): the tier's default 8-bit attention does not exist, so + // the caller's unset attention resolves to exact instead of failing. On sm_80+ an 8-bit kernel always exists, so + // the fallback only ever runs on such devices. + const bool exact_when_no_q8 = !supported(planner::AttentionMode::kQ8Fp8) && + !supported(planner::AttentionMode::kQ8Fp16); r.tier = planner::resolve_tier(m.options.tier, m.options.attention, src.precision, supported, tes.precision != qwen35_text::Precision::kBF16, - src.precision == planner::Precision::kNvfp4 && kernels::nvfp4::available()); + src.precision == planner::Precision::kNvfp4 && kernels::nvfp4::available(), + exact_when_no_q8); if (lora_requested) r.tier.reason += "; runtime LoRAs active"; if (req.accel.step_cache_on()) { // An unset threshold takes the model's default (0.02 on the base model, told apart by its DiT's input embedder: diff --git a/src/pipeline/pipeline_impl.h b/src/pipeline/pipeline_impl.h index 2e00dec..8e7a418 100644 --- a/src/pipeline/pipeline_impl.h +++ b/src/pipeline/pipeline_impl.h @@ -25,6 +25,7 @@ #include #include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" #include "models/control/control.h" #include "models/flux2_vae/flux2_vae.h" #include "models/flux2_vae/vae_file.h" @@ -1008,9 +1009,15 @@ struct Pipeline::Impl { return *c; } - // The block path for this call: options.block_path unless the fused GEMMs failed their self-check. + // The block path for this call: options.block_path unless the fused GEMMs failed their self-check. sm_75 never runs + // the fused route: the fused bf16 CUTLASS GEMMs need sm_80 tensor cores (and the fp16 GEMM path takes fp16 weights), + // so the blocks run op by op there, said so instead of letting the self-check fail with a kernel-image error. 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)"; + return moss_dit::BlockPath::kOpByOp; + } const moss_dit::FusedCheck c = check_fused(); if (c.ok) return options.block_path; note = std::string("fused DiT blocks disabled, running op by op: ") + c.detail; diff --git a/src/planner/planner.cpp b/src/planner/planner.cpp index ff050e1..b86ae83 100644 --- a/src/planner/planner.cpp +++ b/src/planner/planner.cpp @@ -883,7 +883,8 @@ AttentionMode parse_attention(std::string_view name) { } TierDecision resolve_tier(Tier requested, std::optional attention, Precision precision, - const AttentionSupported& supported, bool te_quantized, bool nvfp4_native) { + const AttentionSupported& supported, bool te_quantized, bool nvfp4_native, + bool exact_when_no_q8) { TierDecision d; d.tier = requested; d.precision = precision; @@ -926,7 +927,8 @@ TierDecision resolve_tier(Tier requested, std::optional attention d.reason = "Reference tier: bf16 GEMMs, exact attention"; return d; } - // High and Turbo: the tier's default attention (kTierAttention) unless the caller gave one; never a silent fallback. + // High and Turbo: the tier's default attention (kTierAttention) unless the caller gave one; never a silent fallback + // (exact_when_no_q8 makes the absence explicit in the reason instead: sm_75 has no 8-bit attention kernel). const char* const tier = requested == Tier::kTurbo ? "turbo" : "high"; if (attention) { d.attention = *attention; @@ -934,17 +936,24 @@ TierDecision resolve_tier(Tier requested, std::optional attention const auto it = std::find_if(std::begin(kTierAttention), std::end(kTierAttention), [&](AttentionMode m) { return !supported || supported(m); }); if (it == std::end(kTierAttention)) { - throw std::invalid_argument(std::string("tier ") + tier + " runs 8-bit attention, which this device/build lacks; " - "pass attention exact explicitly"); + if (!exact_when_no_q8) { + throw std::invalid_argument(std::string("tier ") + tier + " runs 8-bit attention, which this device/build lacks; " + "pass attention exact explicitly"); + } + d.attention = AttentionMode::kExact; + } else { + d.attention = *it; } - d.attention = *it; } if (!int8 && d.attention == AttentionMode::kExact) { throw std::invalid_argument("tier high with attention exact on a bf16 DiT and text encoder is the Reference " "configuration; use tier reference, or 8-bit attention"); } d.reason = std::string(requested == Tier::kTurbo ? "Turbo tier" : "High tier") + ": " + weights + ", " + - attention_name(d.attention) + " attention"; + attention_name(d.attention) + " attention" + + (exact_when_no_q8 && !attention && d.attention == AttentionMode::kExact + ? " (no 8-bit attention kernel on this device)" + : ""); return d; } diff --git a/src/planner/planner.h b/src/planner/planner.h index ac13683..71f40a4 100644 --- a/src/planner/planner.h +++ b/src/planner/planner.h @@ -131,8 +131,11 @@ inline constexpr AttentionMode kTierAttention[] = {AttentionMode::kQ8Fp8, Attent // Throws std::invalid_argument (naming the conflict) for a refused combination. An empty `supported` runs every kernel. // te_quantized: the text encoder is an INT8 / W4A8 file, a lossy dimension as an INT8 DiT is (reference refused; the // exact-attention configuration is a High-tier one). nvfp4_native: the NVFP4 GEMMs run here (Capabilities::nvfp4_native). +// exact_when_no_q8: the device has no 8-bit attention kernel (sm_75), so an unset `attention` resolves to exact with a +// reason instead of refusing the tier (the default would otherwise be a silent, absent fallback). TierDecision resolve_tier(Tier tier, std::optional attention, Precision dit_precision, - const AttentionSupported& supported, bool te_quantized = false, bool nvfp4_native = false); + const AttentionSupported& supported, bool te_quantized = false, bool nvfp4_native = false, + bool exact_when_no_q8 = false); // Machine throughput the cost model uses. Defaults are conservative for a desktop PC; the pipeline replaces the // transfer rates with measured ones (pipeline.h: a probe at first use, cached on disk with the embedding cache, then diff --git a/tests/unit/test_attention.cu b/tests/unit/test_attention.cu index 1c023ac..4eda288 100644 --- a/tests/unit/test_attention.cu +++ b/tests/unit/test_attention.cu @@ -382,11 +382,21 @@ int max_smem_per_block() { return v; } +// Compute capability major of the current device: the four forced-tile tests below name sm_80 instantiations, which an +// sm_75 build does not carry (its dispatcher runs the vendored kernel's own Turing traits instead). +int compute_major() { + int device = 0, v = 0; + MI_CUDA_CHECK(cudaGetDevice(&device)); + MI_CUDA_CHECK(cudaDeviceGetAttribute(&v, cudaDevAttrComputeCapabilityMajor, device)); + return v; +} + // All four FA2 instantiations against the fp32 reference. The auto dispatcher picks only one of the // two tiles per head_dim on a given GPU; the override runs the other too. hdim256 128x64 needs 128 KiB // of shared memory per block: sm_89 allows 99 KiB, so there it must be rejected (and is only runnable on // sm_80/90-class parts), which this test asserts instead of skipping. void check_forced_tile(mi::attention::FlashTile tile, int need_smem, const Problem& p, std::uint64_t seed) { + if (compute_major() < 8) GTEST_SKIP() << "forced sm_80 tile instantiations are not built for sm_75"; ResetOverrides reset; Buffers buf(p); buf.fill_random(p, seed); @@ -417,6 +427,7 @@ TEST(GpuAttention, Hdim256Tile128x64) { } TEST(GpuAttention, ForcedTileMustMatchHeadDim) { + if (compute_major() < 8) GTEST_SKIP() << "forced sm_80 tile instantiations are not built for sm_75"; ResetOverrides reset; const Problem p = make_problem(4, 4, 128, false, {16}); Buffers buf(p); diff --git a/tests/unit/test_backend_probe.cpp b/tests/unit/test_backend_probe.cpp index dfd8997..534452a 100644 --- a/tests/unit/test_backend_probe.cpp +++ b/tests/unit/test_backend_probe.cpp @@ -122,7 +122,13 @@ TEST(BackendProbe, DeviceVerdicts) { DeviceInfo turing = device(7, 5); mi::backend::evaluate_device(turing, 13040, 13040, kRelease); - EXPECT_NE(turing.reason.find("8.0"), std::string::npos) << turing.reason; + EXPECT_FALSE(turing.supported); + if (mi::backend::kMinComputeMajor >= 8) { + EXPECT_NE(turing.reason.find("8.0"), std::string::npos) << turing.reason; + } else { + // An sm_75 build accepts compute capability 7.5 in general; this list (the release archs) has no image for it. + EXPECT_NE(turing.reason.find("no kernel image for compute capability 7.5"), std::string::npos) << turing.reason; + } DeviceInfo old_driver = device(8, 9); mi::backend::evaluate_device(old_driver, 12090, 13040, kRelease); EXPECT_NE(old_driver.reason.find("580"), std::string::npos) << old_driver.reason; diff --git a/tests/unit/test_planner.cpp b/tests/unit/test_planner.cpp index 45e8cae..acf8d10 100644 --- a/tests/unit/test_planner.cpp +++ b/tests/unit/test_planner.cpp @@ -1179,6 +1179,14 @@ TEST(Planner, TurboTier) { EXPECT_EQ(d.attention, AttentionMode::kQ8Fp8); EXPECT_EQ(d.reason, "Turbo tier: INT4 ConvRot GEMMs, q8_fp8 attention"); EXPECT_THROW(resolve_tier(Tier::kTurbo, O{}, Precision::kInt4, no_q8), std::invalid_argument); // no fallback + // sm_75: no 8-bit attention kernel exists, so an unset attention may resolve to exact (explicitly said so). + { + const TierDecision t = resolve_tier(Tier::kTurbo, O{}, Precision::kInt4, no_q8, false, false, /*exact_when_no_q8=*/true); + EXPECT_EQ(t.attention, AttentionMode::kExact); + EXPECT_NE(t.reason.find("no 8-bit attention kernel"), std::string::npos) << t.reason; + const TierDecision h = resolve_tier(Tier::kHigh, O{}, Precision::kInt8, no_q8, false, false, /*exact_when_no_q8=*/true); + EXPECT_EQ(h.attention, AttentionMode::kExact); + } EXPECT_EQ(resolve_tier(Tier::kTurbo, O{AttentionMode::kExact}, Precision::kInt4, no_q8).attention, AttentionMode::kExact); EXPECT_EQ(resolve_tier(Tier::kTurbo, O{AttentionMode::kExact}, Precision::kInt4, all).attention, AttentionMode::kExact); EXPECT_EQ(resolve_tier(Tier::kTurbo, O{}, Precision::kInt4, all, true).tier, Tier::kTurbo); // + quantized TE diff --git a/tests/unit/test_sm75.cpp b/tests/unit/test_sm75.cpp new file mode 100644 index 0000000..fcf5ade --- /dev/null +++ b/tests/unit/test_sm75.cpp @@ -0,0 +1,126 @@ +// src/kernels/sm75: the Turing kernel family's build presence and its dtype conversion (the only piece that runs on +// any device, so this file also works as a smoke test that the sm_75 object library is linked in). +#include + +#include +#include +#include + +#include +#include +#include + +#include "core/cuda_check.h" +#include "core/tensor.h" +#include "kernels/sm75/sm75.h" + +namespace { + +using mi::DType; +using mi::Device; +using mi::Tensor; + +float bf16_to_float(std::uint16_t bits) { + const std::uint32_t u = static_cast(bits) << 16; + float f = 0; + std::memcpy(&f, &u, sizeof f); + return f; +} + +std::uint16_t float_to_bf16_bits(float f) { + std::uint32_t u = 0; + std::memcpy(&u, &f, sizeof u); + // Round-to-nearest-even, as __float2bfloat16_rn: the test only uses values that are exact in bf16. + const std::uint32_t rounded = u + 0x7fffu + ((u >> 16) & 1u); + return static_cast(rounded >> 16); +} + +// A bf16 device buffer of `n` elements filled from host floats, with the values read back as bf16 bits. +struct DeviceBf16 { + void* data = nullptr; + std::int64_t n = 0; + explicit DeviceBf16(const std::vector& v) : n(static_cast(v.size())) { + std::vector bits(v.size()); + for (std::size_t i = 0; i < v.size(); ++i) bits[i] = float_to_bf16_bits(v[i]); + MI_CUDA_CHECK(cudaMalloc(&data, v.size() * sizeof(std::uint16_t))); + MI_CUDA_CHECK(cudaMemcpy(data, bits.data(), v.size() * sizeof(std::uint16_t), cudaMemcpyHostToDevice)); + } + ~DeviceBf16() { + if (data != nullptr) cudaFree(data); + } + std::vector host() const { + std::vector bits(static_cast(n)); + MI_CUDA_CHECK(cudaMemcpy(bits.data(), data, bits.size() * sizeof(std::uint16_t), cudaMemcpyDeviceToHost)); + return bits; + } + Tensor tensor() const { + return Tensor::view(data, DType::kBF16, Device::kCUDA, {n}); + } +}; + +} // namespace + +TEST(Sm75Build, KernelFamilyPresence) { +#ifndef MOSSINFER_HAVE_SM75 + GTEST_SKIP() << "built without MOSSINFER_SM75"; +#else + EXPECT_TRUE(mi::kernels::sm75::available()); + EXPECT_TRUE(mi::kernels::sm75::needed(7)); + EXPECT_FALSE(mi::kernels::sm75::needed(8)); +#endif +} + +// bf16 -> fp16 is lossless for in-range values and fp16 -> bf16 recovers the exact bf16 bits, which is what the sm_75 +// attention relies on when it converts q/k/v in place and restores them. +TEST(GpuSm75Convert, RoundTripIsExactInFp16Range) { +#ifndef MOSSINFER_HAVE_SM75 + GTEST_SKIP() << "built without MOSSINFER_SM75"; +#else + // 65504 is the largest fp16 value but not representable in bf16 (it rounds to 65536, which overflows fp16 to inf): + // values above fp16's range are the documented sm_75 numerics difference, not tested here. + const std::vector values = {0.f, 1.f, -1.f, 2.5f, -3.14159f, 123.456f, -1024.f, 32768.f, 1e-3f, -7e-2f}; + DeviceBf16 buf(values); + const Tensor t = buf.tensor(); + mi::kernels::sm75::convert_bf16_to_f16(t.data, t.data, t.numel(), nullptr); + mi::kernels::sm75::convert_f16_to_bf16(t.data, t.data, t.numel(), nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + const std::vector got = buf.host(); + for (std::size_t i = 0; i < values.size(); ++i) { + EXPECT_EQ(got[i], float_to_bf16_bits(values[i])) << "value " << values[i]; + EXPECT_EQ(bf16_to_float(got[i]), bf16_to_float(float_to_bf16_bits(values[i]))) << "value " << values[i]; + } +#endif +} + +// The attention tensors are [tokens, heads, dim] views with free row/head strides: the in-place conversion must touch +// exactly the view's elements and no neighbour. +TEST(GpuSm75Convert, StridedInPlaceConversionTouchesOnlyTheView) { +#ifndef MOSSINFER_HAVE_SM75 + GTEST_SKIP() << "built without MOSSINFER_SM75"; +#else + constexpr int kRows = 3, kHeads = 2, kDim = 8, kHeadStride = 16; // heads interleaved: [token][head][16] + std::vector storage(static_cast(kRows) * kHeads * kHeadStride, -12345.f); + DeviceBf16 buf(storage); + Tensor view = Tensor::view(buf.data, DType::kBF16, Device::kCUDA, {kRows, kHeads, kDim}); + view.stride[0] = kHeads * kHeadStride; + view.stride[1] = kHeadStride; + view.stride[2] = 1; + const std::vector before = buf.host(); + mi::kernels::sm75::convert_tensor_bf16_to_f16_inplace(view, kRows, kHeads, kDim, nullptr); + mi::kernels::sm75::convert_tensor_f16_to_bf16_inplace(view, kRows, kHeads, kDim, nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + const std::vector after = buf.host(); + for (int r = 0; r < kRows; ++r) { + for (int h = 0; h < kHeads; ++h) { + for (int d = 0; d < kDim; ++d) { + const std::size_t i = static_cast(r) * kHeads * kHeadStride + h * kHeadStride + d; + EXPECT_EQ(after[i], before[i]) << "r " << r << " h " << h << " d " << d; + } + for (int d = kDim; d < kHeadStride; ++d) { + const std::size_t i = static_cast(r) * kHeads * kHeadStride + h * kHeadStride + d; + EXPECT_EQ(after[i], before[i]) << "padding r " << r << " h " << h << " d " << d; + } + } + } +#endif +} diff --git a/third_party/flash-attention/VENDORED.md b/third_party/flash-attention/VENDORED.md index 91ba635..423593b 100644 --- a/third_party/flash-attention/VENDORED.md +++ b/third_party/flash-attention/VENDORED.md @@ -23,6 +23,13 @@ trees, and the Python package. which MossInfer never enables. 2. `flash.h`: `#include ` became `#include "mi_torch_shim.h"`. 3. `philox_unpack.cuh`: `#include ` became `#include "mi_torch_shim.h"`. +4. `flash_fwd_kernel.h`: the two `if (Is_dropout)` guards around `dropout.apply_dropout(...)` and the two + `if (Return_softmax)` guards around its `encode_dropout_in_sign_bit` call became `if constexpr`. With a runtime `if`, + the dropout code was instantiated even for `Is_dropout=false` / `Return_softmax=false`, and its + `static_assert(size<2>(tensor) % 2 == 0)` (dropout.h) holds for the sm_80 `m16n8k16` accumulator layout but not for + Turing's `m16n8k8` layout, so an sm_75 instantiation failed to compile. Both are template parameters, so `if constexpr` + skips the instantiation for every architecture and changes no generated code where they are false (the only mode + MossInfer builds). Nothing else changed: no kernel code was touched. diff --git a/third_party/flash-attention/csrc/flash_attn/src/flash_fwd_kernel.h b/third_party/flash-attention/csrc/flash_attn/src/flash_fwd_kernel.h index d492c87..b5bf0d8 100644 --- a/third_party/flash-attention/csrc/flash_attn/src/flash_fwd_kernel.h +++ b/third_party/flash-attention/csrc/flash_attn/src/flash_fwd_kernel.h @@ -347,7 +347,7 @@ inline __device__ void compute_attn_1rowblock(const Params ¶ms, const int bi Tensor rP = FLASH_NAMESPACE::convert_type(acc_s); int block_row_idx = m_block * (kBlockM / 16) + tidx / 32; int block_col_idx = n_block * (kBlockN / 32); - if (Return_softmax) { + if constexpr (Return_softmax) { Tensor rP_drop = make_fragment_like(rP); cute::copy(rP, rP_drop); dropout.template apply_dropout( @@ -356,7 +356,7 @@ inline __device__ void compute_attn_1rowblock(const Params ¶ms, const int bi cute::copy(rP_drop, tSgS); tSgS.data() = tSgS.data() + (-kBlockN); } - if (Is_dropout) { + if constexpr (Is_dropout) { dropout.apply_dropout(rP, block_row_idx, block_col_idx, kNWarps); } @@ -409,7 +409,7 @@ inline __device__ void compute_attn_1rowblock(const Params ¶ms, const int bi Tensor rP = FLASH_NAMESPACE::convert_type(acc_s); int block_row_idx = m_block * (kBlockM / 16) + tidx / 32; int block_col_idx = n_block * (kBlockN / 32); - if (Return_softmax) { + if constexpr (Return_softmax) { Tensor rP_drop = make_fragment_like(rP); cute::copy(rP, rP_drop); dropout.template apply_dropout( @@ -418,7 +418,7 @@ inline __device__ void compute_attn_1rowblock(const Params ¶ms, const int bi cute::copy(rP_drop, tSgS); tSgS.data() = tSgS.data() + (-kBlockN); } - if (Is_dropout) { + if constexpr (Is_dropout) { dropout.apply_dropout(rP, block_row_idx, block_col_idx, kNWarps); } diff --git a/tools/sm75_check.cu b/tools/sm75_check.cu new file mode 100644 index 0000000..a0daa54 --- /dev/null +++ b/tools/sm75_check.cu @@ -0,0 +1,532 @@ +// 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. +// +// Usage: sm75_check [attn128|attn256|convert|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 +// convert bf16 <-> fp16 round trip on device +// The engine's own GPU tests compile their reference kernels for the build's CUDA architectures (sm_80+), so the +// unit-test binary cannot exercise sm_75 kernels on a Turing device; this tool builds its own reference in host code. +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "backend/cuda_libs.h" +#include "backend/probe.h" +#include "core/cuda_check.h" +#include "core/tensor.h" +#include "kernels/attention/flash_attention.h" +#include "kernels/basic/gemm.h" +#include "kernels/quant/quant.h" +#include "kernels/sm75/sm75.h" + +namespace { + +using mi::DType; +using mi::Device; +using mi::Tensor; + +float bf16_to_float(std::uint16_t bits) { + const std::uint32_t u = static_cast(bits) << 16; + float f = 0; + std::memcpy(&f, &u, sizeof f); + return f; +} + +std::uint16_t float_to_bf16_bits(float f) { + std::uint32_t u = 0; + std::memcpy(&u, &f, sizeof u); + const std::uint32_t rounded = u + 0x7fffu + ((u >> 16) & 1u); + return static_cast(rounded >> 16); +} + +std::uint32_t hash_u32(std::uint64_t x) { + x ^= x >> 33; + x *= 0xff51afd7ed558ccdull; + x ^= x >> 33; + x *= 0xc4ceb9fe1a85ec53ull; + x ^= x >> 33; + return static_cast(x); +} + +std::vector normal(std::size_t n, std::uint64_t seed, float scale) { + std::vector v(n); + for (std::size_t i = 0; i < n; ++i) { + const std::uint64_t key = seed * 0x9E3779B97F4A7C15ull + i * 2; + const float u1 = (static_cast(hash_u32(key) >> 8) + 1.f) * (1.f / 16777217.f); + const float u2 = static_cast(hash_u32(key + 1) >> 8) * (1.f / 16777216.f); + v[i] = scale * std::sqrt(-2.f * std::log(u1)) * std::cos(6.283185307179586f * u2); + } + return v; +} + +class DevBuf { + public: + explicit DevBuf(std::size_t bytes) { MI_CUDA_CHECK(cudaMalloc(&p_, bytes == 0 ? 1 : bytes)); } + ~DevBuf() { cudaFree(p_); } + DevBuf(const DevBuf&) = delete; + DevBuf& operator=(const DevBuf&) = delete; + void* get() const { return p_; } + template + void upload(const std::vector& v) { + MI_CUDA_CHECK(cudaMemcpy(p_, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice)); + } + template + std::vector download(std::size_t n) const { + std::vector v(n); + MI_CUDA_CHECK(cudaMemcpy(v.data(), p_, n * sizeof(T), cudaMemcpyDeviceToHost)); + return v; + } + + private: + void* p_ = nullptr; +}; + +// Host reference: one (row, head) at a time, fp32 scores/softmax/PV, bf16 inputs as the device sees them. +std::vector reference_attention(const std::vector& q, const std::vector& k, + const std::vector& v, const std::vector& cu_q, + int heads, int kv_heads, int dim, bool causal, float scale) { + 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); + const int group = heads / kv_heads; + std::vector scores; + for (int b = 0; b < batch; ++b) { + const int q0 = cu_q[b], q1 = cu_q[b + 1]; + for (int h = 0; h < heads; ++h) { + const int kh = h / group; + for (int i = q0; i < q1; ++i) { + const int n = causal ? (i - q0 + 1) : (q1 - q0); + scores.assign(static_cast(n), 0.f); + float max_s = -INFINITY; + for (int j = 0; j < n; ++j) { + const int row = q0 + j; + float s = 0.f; + for (int d = 0; d < dim; ++d) { + s += bf16_to_float(q[(static_cast(i) * heads + h) * dim + d]) * + bf16_to_float(k[(static_cast(row) * kv_heads + kh) * dim + d]); + } + s *= scale; + scores[static_cast(j)] = s; + max_s = std::max(max_s, s); + } + float sum = 0.f; + for (int j = 0; j < n; ++j) { + scores[static_cast(j)] = std::exp(scores[static_cast(j)] - max_s); + sum += scores[static_cast(j)]; + } + 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; + const int row = q0 + j; + for (int d = 0; d < dim; ++d) { + out[base + d] += p * bf16_to_float(v[(static_cast(row) * kv_heads + kh) * dim + d]); + } + } + } + } + } + return out; +} + +int run_attention(const char* name, int heads, int kv_heads, int dim, bool causal, const std::vector& lengths, + std::uint64_t seed, bool dump = false) { + const int batch = static_cast(lengths.size()); + std::vector cu(static_cast(batch) + 1, 0); + for (int b = 0; b < batch; ++b) cu[static_cast(b) + 1] = cu[static_cast(b)] + lengths[static_cast(b)]; + const int total_q = cu.back(); + const int max_len = *std::max_element(lengths.begin(), lengths.end()); + const float scale = 1.f / std::sqrt(static_cast(dim)); + + const auto q = normal(static_cast(total_q) * heads * dim, seed, 1.f); + const auto k = normal(static_cast(total_q) * kv_heads * dim, seed + 1, 1.f); + const auto v = normal(static_cast(total_q) * kv_heads * dim, seed + 2, 1.f); + std::vector qb(q.size()), kb(k.size()), vb(v.size()); + for (std::size_t i = 0; i < q.size(); ++i) qb[i] = float_to_bf16_bits(q[i]); + for (std::size_t i = 0; i < k.size(); ++i) kb[i] = float_to_bf16_bits(k[i]); + for (std::size_t i = 0; i < v.size(); ++i) vb[i] = float_to_bf16_bits(v[i]); + + DevBuf dq(qb.size() * 2), dk(kb.size() * 2), dv(vb.size() * 2), dout(qb.size() * 2), dcu(cu.size() * sizeof(int)); + dq.upload(qb); + dk.upload(kb); + dv.upload(vb); + dcu.upload(cu); + MI_CUDA_CHECK(cudaMemset(dout.get(), 0, qb.size() * 2)); + std::vector lse(static_cast(total_q) * heads, 0.f); + DevBuf dls(lse.size() * sizeof(float)); + MI_CUDA_CHECK(cudaMemset(dls.get(), 0, lse.size() * sizeof(float))); + + auto t = [](void* p, DType dt, std::initializer_list shape) { + return Tensor::view(p, dt, Device::kCUDA, shape); + }; + const Tensor tq = t(dq.get(), DType::kBF16, {total_q, heads, dim}); + const Tensor tk = t(dk.get(), DType::kBF16, {total_q, kv_heads, dim}); + const Tensor tv = t(dv.get(), DType::kBF16, {total_q, kv_heads, dim}); + const Tensor to = t(dout.get(), DType::kBF16, {total_q, heads, dim}); + const Tensor tcu = t(dcu.get(), DType::kI32, {batch + 1}); + const std::size_t ws_bytes = mi::attention::flash_fwd_workspace_bytes(total_q, heads); + MI_CUDA_CHECK(cudaGetLastError()); + std::printf(" launch: tokens %d heads %d kv %d dim %d causal %d maxlen %d ws %zu\n", total_q, heads, kv_heads, dim, + causal ? 1 : 0, max_len, ws_bytes); + std::fflush(stdout); + mi::attention::flash_fwd(tq, tk, tv, to, tcu, tcu, max_len, causal, scale, dls.get(), ws_bytes, nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + std::printf(" kernel done\n"); + std::fflush(stdout); + + const auto got = dout.download(qb.size()); + const auto want = reference_attention(qb, kb, vb, cu, heads, kv_heads, dim, causal, scale); + double dot = 0, ng = 0, nw = 0; + float max_abs = 0.f; + for (std::size_t i = 0; i < want.size(); ++i) { + const float g = bf16_to_float(got[i]); + dot += static_cast(g) * want[i]; + ng += static_cast(g) * g; + nw += static_cast(want[i]) * want[i]; + max_abs = std::max(max_abs, std::fabs(g - want[i])); + } + const double cosine = dot / std::sqrt(ng * nw); + std::printf("%-8s tokens %5d heads %2d kv %2d dim %3d causal %d: cosine %.9f max|diff| %.3g\n", name, total_q, heads, + kv_heads, dim, causal ? 1 : 0, cosine, max_abs); + if (dump) { + for (std::size_t i = 0; i < std::min(want.size(), 12); ++i) { + std::printf(" [%zu] got %.6f want %.6f (q %.6f)\n", i, bf16_to_float(got[i]), want[i], q[i]); + } + // Diagnostic: is the mismatch a different softmax temperature? + const float log2e = 1.4426950408889634f; + for (float f : {1.f, log2e, 1.f / log2e, 2.f, 0.5f, 1.f / std::sqrt(static_cast(dim))}) { + const auto alt = reference_attention(qb, kb, vb, cu, heads, kv_heads, dim, causal, scale * f); + double d2 = 0, g2 = 0, w2 = 0; + for (std::size_t i = 0; i < alt.size(); ++i) { + const float g = bf16_to_float(got[i]); + d2 += static_cast(g) * alt[i]; + g2 += static_cast(g) * g; + w2 += static_cast(alt[i]) * alt[i]; + } + std::printf(" scale x %.6f: cosine %.9f\n", f, d2 / std::sqrt(g2 * w2)); + } + } + MI_CUDA_CHECK(cudaGetLastError()); + return cosine >= 0.9999 ? 0 : 1; +} + +int run_convert() { + const std::vector values = {0.f, 1.f, -1.f, 2.5f, -3.14159f, 123.456f, -1024.f, 32768.f, 1e-3f}; + DevBuf d(values.size() * 2); + std::vector bits(values.size()); + for (std::size_t i = 0; i < values.size(); ++i) bits[i] = float_to_bf16_bits(values[i]); + d.upload(bits); + mi::kernels::sm75::convert_bf16_to_f16(d.get(), d.get(), static_cast(values.size()), nullptr); + mi::kernels::sm75::convert_f16_to_bf16(d.get(), d.get(), static_cast(values.size()), nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + const auto got = d.download(values.size()); + int bad = 0; + for (std::size_t i = 0; i < values.size(); ++i) bad += got[i] != bits[i]; + std::printf("convert: %zu values, %d mismatches\n", values.size(), bad); + return bad == 0 ? 0 : 1; +} + +// The sm_75 GEMM (cuBLASLt fp16 in / fp32 accumulate, bf16 activation input): bf16 x, fp16 weights, bf16 out with a +// bias added in fp32. Cross-checked against a host accumulation of the stored operand values. +int run_gemm(int M, int N, int K) { + mi::backend::load_cublaslt(); + const auto xv = normal(static_cast(M) * K, 51, 1.f); + const auto wv = normal(static_cast(N) * K, 52, 0.05f); + const auto bv = normal(static_cast(N), 53, 0.01f); + std::vector xb(xv.size()), wb(wv.size()), bb(bv.size()); + for (std::size_t i = 0; i < xv.size(); ++i) xb[i] = float_to_bf16_bits(xv[i]); + for (std::size_t i = 0; i < wv.size(); ++i) { + const __half h = __float2half_rn(wv[i]); + std::memcpy(&wb[i], &h, sizeof h); + } + for (std::size_t i = 0; i < bv.size(); ++i) bb[i] = float_to_bf16_bits(bv[i]); + + DevBuf dx(xb.size() * 2), dw(wb.size() * 2), db(bb.size() * 2), dout(static_cast(M) * N * 2); + DevBuf dws(std::size_t{32} << 20); + dx.upload(xb); + dw.upload(wb); + db.upload(bb); + MI_CUDA_CHECK(cudaMemset(dout.get(), 0, static_cast(M) * N * 2)); + + const Tensor tx = Tensor::view(dx.get(), DType::kBF16, Device::kCUDA, {M, K}); + const Tensor tw = Tensor::view(dw.get(), DType::kF16, Device::kCUDA, {N, K}); + const Tensor tb = Tensor::view(db.get(), DType::kBF16, Device::kCUDA, {N}); + const Tensor to = Tensor::view(dout.get(), DType::kBF16, Device::kCUDA, {M, N}); + mi::kernels::Gemm gemm; + gemm.linear(to, tx, tw, &tb, dws.get(), std::size_t{32} << 20, nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + const auto got = dout.download(static_cast(M) * N); + + double dot = 0, ng = 0, nw = 0; + float max_abs = 0.f, max_rel = 0.f; + for (int r = 0; r < M; ++r) { + for (int c = 0; c < N; ++c) { + double acc = 0; + for (int d = 0; d < K; ++d) { + const float w = [&] { + __half h; + std::memcpy(&h, &wb[static_cast(c) * K + d], sizeof h); + return __half2float(h); + }(); + acc += static_cast(bf16_to_float(xb[static_cast(r) * K + d])) * w; + } + const float want = static_cast(acc) + bf16_to_float(bb[static_cast(c)]); + const float g = bf16_to_float(got[static_cast(r) * N + c]); + dot += static_cast(g) * want; + ng += static_cast(g) * g; + nw += static_cast(want) * want; + max_abs = std::max(max_abs, std::fabs(g - want)); + max_rel = std::max(max_rel, std::fabs(g - want) / std::max(1e-3f, std::fabs(want))); + } + } + const double cosine = dot / std::sqrt(ng * nw); + std::printf("gemm M %d N %d K %d: cosine %.9f max|diff| %.3g max|rel| %.3g\n", M, N, K, cosine, max_abs, max_rel); + return cosine >= 0.9999 && max_rel < 0.01 ? 0 : 1; +} + +// VAE sm_75 conv2d: NHWC fp32 with KRSC filters, bias + residual, against a host reference. +int run_vae_conv_case(int N, int H, int W, int C, int K, int R, std::uint64_t seed) { + const auto xv = normal(static_cast(N) * H * W * C, seed, 1.f); + const auto wv = normal(static_cast(K) * R * R * C, seed + 1, 0.2f); + const auto bv = normal(static_cast(K), seed + 2, 0.1f); + const auto rv = normal(static_cast(N) * H * W * K, seed + 3, 0.5f); + DevBuf dx(xv.size() * 4), dw(wv.size() * 4), db(bv.size() * 4), dr(rv.size() * 4), dy(rv.size() * 4); + dx.upload(xv); + dw.upload(wv); + db.upload(bv); + dr.upload(rv); + MI_CUDA_CHECK(cudaMemset(dy.get(), 0, rv.size() * 4)); + const Tensor tx = Tensor::view(dx.get(), DType::kF32, Device::kCUDA, {N, H, W, C}); + const Tensor tw = Tensor::view(dw.get(), DType::kF32, Device::kCUDA, {K, R, R, C}); + 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_CUDA_CHECK(cudaDeviceSynchronize()); + const auto got = dy.download(rv.size()); + const int pad = (R - 1) / 2; + double dot = 0, ng = 0, nw = 0; + float max_abs = 0.f; + for (int n = 0; n < N; ++n) { + for (int oy = 0; oy < H; ++oy) { + for (int ox = 0; ox < W; ++ox) { + for (int k = 0; k < K; ++k) { + double acc = 0; + for (int r = 0; r < R; ++r) { + for (int s = 0; s < R; ++s) { + const int gy = oy - pad + r, gx = ox - pad + s; + if (gy < 0 || gy >= H || gx < 0 || gx >= W) continue; + for (int c = 0; c < C; ++c) { + acc += static_cast(xv[(static_cast(n) * H + gy) * W * C + gx * C + c]) * + wv[((static_cast(k) * R + r) * R + s) * C + c]; + } + } + } + const float want = static_cast(acc) + bv[static_cast(k)] + + rv[(static_cast(n) * H + oy) * W * K + ox * K + k]; + const float g = got[(static_cast(n) * H + oy) * W * K + ox * K + k]; + dot += static_cast(g) * want; + ng += static_cast(g) * g; + nw += static_cast(want) * want; + max_abs = std::max(max_abs, std::fabs(g - want)); + } + } + } + } + const double cosine = dot / std::sqrt(ng * nw); + std::printf("vae_conv N%d H%d W%d C%d K%d R%d: cosine %.9f max|diff| %.3g\n", N, H, W, C, K, R, cosine, max_abs); + return cosine >= 0.99999 ? 0 : 1; +} + +// VAE sm_75 mid attention (fp32, single head) against a host reference. +int run_vae_attn() { + const int B = 1, L = 12, D = 64; + const auto qv = normal(static_cast(B) * L * D, 71, 1.f); + const auto kv = normal(static_cast(B) * L * D, 72, 1.f); + const auto vv = normal(static_cast(B) * L * D, 73, 1.f); + DevBuf dq(qv.size() * 4), dk(kv.size() * 4), dv(vv.size() * 4), dout(qv.size() * 4); + dq.upload(qv); + dk.upload(kv); + dv.upload(vv); + MI_CUDA_CHECK(cudaMemset(dout.get(), 0, qv.size() * 4)); + const Tensor tq = Tensor::view(dq.get(), DType::kF32, Device::kCUDA, {B, L, D}); + const Tensor tk = Tensor::view(dk.get(), DType::kF32, Device::kCUDA, {B, L, D}); + const Tensor tv = Tensor::view(dv.get(), DType::kF32, Device::kCUDA, {B, L, D}); + const Tensor to = Tensor::view(dout.get(), DType::kF32, Device::kCUDA, {B, L, D}); + const float scale = 1.f / std::sqrt(static_cast(D)); + mi::kernels::sm75::vae_attention_fp32(tq, tk, tv, to, scale, nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + const auto got = dout.download(qv.size()); + std::vector want(qv.size()); + for (int b = 0; b < B; ++b) { + for (int i = 0; i < L; ++i) { + std::vector s(L); + double m = -1e30; + for (int j = 0; j < L; ++j) { + double acc = 0; + for (int d = 0; d < D; ++d) acc += static_cast(qv[(static_cast(b) * L + i) * D + d]) * + kv[(static_cast(b) * L + j) * D + d]; + s[static_cast(j)] = acc * scale; + m = std::max(m, s[static_cast(j)]); + } + double sum = 0; + for (int j = 0; j < L; ++j) { + s[static_cast(j)] = std::exp(s[static_cast(j)] - m); + sum += s[static_cast(j)]; + } + for (int d = 0; d < D; ++d) { + double acc = 0; + for (int j = 0; j < L; ++j) { + acc += s[static_cast(j)] * vv[(static_cast(b) * L + j) * D + d]; + } + want[(static_cast(b) * L + i) * D + d] = static_cast(acc / sum); + } + } + } + double dot = 0, ng = 0, nw = 0; + float max_abs = 0.f; + for (std::size_t i = 0; i < want.size(); ++i) { + dot += static_cast(got[i]) * want[i]; + ng += static_cast(got[i]) * got[i]; + nw += static_cast(want[i]) * want[i]; + max_abs = std::max(max_abs, std::fabs(got[i] - want[i])); + } + std::printf("vae_attn B%d L%d D%d: cosine %.9f max|diff| %.3g\n", B, L, D, dot / std::sqrt(ng * nw), max_abs); + return dot / std::sqrt(ng * nw) >= 0.99999 ? 0 : 1; +} + +// sm_75 INT8 ConvRot GEMM: int8 x/w with per-row/col fp32 scales, bf16 output. The int32 dot product is exact, so the +// output must be bit-identical to the host reference of the same epilogue order ((acc * xs) * ws + bias, bf16 rounding). +int run_int8_gemm(int M, int N, int K) { + std::vector xv(static_cast(M) * K), wv(static_cast(N) * K); + std::vector xs(static_cast(M)), ws(static_cast(N)); + std::vector bv(static_cast(N)); + std::uint64_t seed = 81; + auto next = [&seed]() { + seed = seed * 6364136223846793005ull + 1442695040888963407ull; + return static_cast(seed >> 33); + }; + for (auto& v : xv) v = static_cast(static_cast(next() % 255) - 127); + for (auto& v : wv) v = static_cast(static_cast(next() % 255) - 127); + for (auto& s : xs) s = 0.001f + static_cast(next() % 1000) * 1e-5f; + for (auto& s : ws) s = 0.0005f + static_cast(next() % 1000) * 1e-5f; + const auto bias_f = normal(static_cast(N), 82, 0.05f); + std::vector bb(bias_f.size()); + for (std::size_t i = 0; i < bias_f.size(); ++i) bb[i] = float_to_bf16_bits(bias_f[i]); + + DevBuf dx(xv.size()), dw(wv.size()), dxs(xs.size() * 4), dws(ws.size() * 4), db(bb.size() * 2), + dout(static_cast(M) * N * 2); + dx.upload(xv); + dw.upload(wv); + dxs.upload(xs); + dws.upload(ws); + db.upload(bb); + MI_CUDA_CHECK(cudaMemset(dout.get(), 0, static_cast(M) * N * 2)); + + mi::kernels::quant::Int8Operands op; + op.x_q = static_cast(dx.get()); + op.ldx = K; + op.x_scale = static_cast(dxs.get()); + op.w = static_cast(dw.get()); + op.w_scale = static_cast(dws.get()); + op.m = M; + op.n = N; + op.k = K; + auto* bias = static_cast(db.get()); + auto* out = static_cast<__nv_bfloat16*>(dout.get()); + mi::kernels::quant::int8_gemm(op, bias, out, N, nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + const auto got = dout.download(static_cast(M) * N); + + int mismatches = 0; + float max_abs = 0.f; + for (int r = 0; r < M; ++r) { + for (int c = 0; c < N; ++c) { + std::int32_t acc = 0; + for (int d = 0; d < K; ++d) { + acc += static_cast(xv[static_cast(r) * K + d]) * + static_cast(wv[static_cast(c) * K + d]); + } + float v = static_cast(acc) * xs[static_cast(r)]; + v = v * ws[static_cast(c)]; + v += bf16_to_float(bb[static_cast(c)]); + const std::uint16_t want = float_to_bf16_bits(v); + const std::uint16_t g = got[static_cast(r) * N + c]; + if (g != want) { + ++mismatches; + max_abs = std::max(max_abs, std::fabs(bf16_to_float(g) - bf16_to_float(want))); + } + } + } + std::printf("int8 M %d N %d K %d: %d/%d mismatches, max|diff| %.3g\n", M, N, K, mismatches, M * N, max_abs); + return mismatches == 0 ? 0 : 1; +} + +} // namespace + +int main(int argc, char** argv) { + std::setvbuf(stdout, nullptr, _IONBF, 0); + const std::string what = argc > 1 ? argv[1] : "all"; + mi::backend::ProbeOptions options; + options.query_free_memory = false; + const mi::backend::ProbeResult probe = mi::backend::probe(options); + std::printf("%s", mi::backend::format_report(probe).c_str()); + int failures = 0; + try { + if (what == "convert" || what == "all") failures += run_convert(); + if (what == "gemm" || what == "all") { + failures += run_gemm(64, 128, 256); + failures += run_gemm(3, 8, 32); // a tail case: m not a multiple of anything + } + if (what == "vae") { + 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") { + 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); + } + if (what == "small" || what == "all") { + if (what == "small") { + failures += run_attention("one", 1, 1, 128, false, {1}, 31, /*dump=*/true); + failures += run_attention("two", 1, 1, 128, false, {2}, 32, /*dump=*/true); + failures += run_attention("small", 1, 1, 128, false, {8}, 21, /*dump=*/true); + } + } + if (what == "attn128" || what == "all") { + failures += run_attention("attn128", 4, 2, 128, false, {64, 128, 1}, 11); + failures += run_attention("attn128", 4, 2, 128, false, {257}, 12); + } + if (what == "attn256" || what == "all") { + failures += run_attention("attn256", 4, 1, 256, true, {32, 100}, 13); + failures += run_attention("attn256", 4, 1, 256, true, {256}, 14); + } + 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("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); + } + } catch (const std::exception& e) { + std::printf("EXCEPTION: %s\n", e.what()); + return 2; + } catch (...) { + std::printf("UNKNOWN EXCEPTION\n"); + return 3; + } + std::printf("%s\n", failures == 0 ? "OK" : "FAILED"); + return failures == 0 ? 0 : 1; +} From 00e247d1edda9e95497e2f973b281b0d3f6e5691 Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sat, 31 Oct 2026 15:42:42 +0800 Subject: [PATCH 2/6] docs: Turing (sm_75) support page, Windows -sm75 preset and package - docs/sm75.md: what MOSSINFER_SM75 changes (fp16 tensor-core GEMMs in place of bf16, INT8/INT4 ConvRot through the INT8 GEMMs, exact attention, fp32 VAE, the GDN recurrent kernel), how to build and package it, measured RTX 2080 performance, numerics, and what is not supported yet (8-bit attention, native INT4/NVFP4, img2img and the VAE encoder, runtime LoRAs). - README and docs: link the page; state the sm_75 tier, attention and block-path behavior; add the measured RTX 2080 numbers; document the -sm75 package and the launcher's device check for Turing machines. - CMakePresets: windows-msvc-cuda13-sm75 (configure/build/test/package presets). - packaging: append -sm75 to the CPack archive name when MOSSINFER_SM75 is on. - CI: build the Windows CUDA 13 Turing flavour as a fifth matrix job (a distinct sm75 axis value, so the four standard jobs keep their flavour and preset) and require its full and system-CUDA archives in release.yml. Verified: the preset configures with probe architectures 80-real,89-real,120,75-real and CPack name mossinfer-2.0.0-windows-cu13-x86_64-sm75; the matrix expands to five jobs; the workflow YAML parses. --- .github/workflows/build.yml | 20 +++++-- .github/workflows/release.yml | 8 ++- CMakePresets.json | 14 +++++ README.md | 5 +- cmake/modules/packaging.cmake | 7 ++- docs/c-api.md | 6 +-- docs/integration.md | 24 +++++++-- docs/sm75.md | 99 +++++++++++++++++++++++++++++++++++ 8 files changed, 166 insertions(+), 17 deletions(-) create mode 100644 docs/sm75.md diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c00c372..a6cf4d0 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -1,7 +1,8 @@ -# Reusable build of the four release flavours (Linux/Windows x CUDA 12/13), 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). +# 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 +# (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). name: build on: @@ -31,6 +32,10 @@ jobs: # Bump the 13.x entry to 13.4 once Jimver/cuda-toolkit lists it: the gate toolchain is 13.4. # 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. + sm75: [false] include: - os: ubuntu-22.04 platform: linux @@ -73,6 +78,13 @@ jobs: cuda: "13.3.1" nvcc: "13.3.73" sub-packages: '["nvcc", "crt", "nvvm", "cudart", "cublas", "cublas_dev", "curand_dev", "thrust", "documentation"]' + - os: windows-2022 + platform: windows-msvc + cuda: "13.3.1" + sm75: true + flavour: cuda13-sm75 + nvcc: "13.3.73" + sub-packages: '["nvcc", "crt", "nvvm", "cudart", "cublas", "cublas_dev", "curand_dev", "thrust", "documentation"]' env: preset: ${{ matrix.platform }}-${{ matrix.flavour }} SCCACHE_GHA_ENABLED: "true" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ee9aad2..2236e2b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1,6 +1,7 @@ # Tag vX.Y.Z (a release) or vX.Y.Z-rcN, N >= 1 (a release candidate, published as a pre-release that does not become -# "latest"): build and check all four flavours (build.yml), then publish their archives (the full package and the -# -system-cuda variant of each) and a SHA256SUMS file as a GitHub release. X.Y.Z must equal the MI_VERSION_* of +# "latest"): build and check all four standard flavours and the Windows CUDA 13 SM75 flavour (build.yml), then publish +# their archives (the full package and the -system-cuda variant of each) and a SHA256SUMS file as a GitHub release. +# X.Y.Z must equal the MI_VERSION_* of # include/mossinfer/mossinfer.h and an rc must come before vX.Y.Z (.github/scripts/release_version.py), the tagged # commit must be on main with a successful ci run (.github/scripts/release_gate.sh), and the release must not exist # yet (a re-run of the same workflow run may finish a partial upload). @@ -117,6 +118,9 @@ jobs: done echo "$n full archives, each with its -system-cuda variant, all version $ARCHIVE_VERSION" test "$n" -gt 0 + 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; } - name: SHA256SUMS shell: bash working-directory: dist diff --git a/CMakePresets.json b/CMakePresets.json index ed460da..d899faa 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -87,6 +87,13 @@ "inherits": [ "release-base", "windows" ], "cacheVariables": { "MOSSINFER_EXPECT_CUDA_MAJOR": "13" } }, + { + "name": "windows-msvc-cuda13-sm75", + "displayName": "Windows release, CUDA 13.x with Turing (sm_75)", + "inherits": "windows-msvc-cuda13", + "binaryDir": "${sourceDir}/build-${presetName}", + "cacheVariables": { "MOSSINFER_SM75": "ON" } + }, { "name": "windows-msvc-cuda12", "displayName": "Windows release, MSVC + CUDA 12.8+ (driver >= 527)", @@ -100,6 +107,7 @@ { "name": "linux-cuda13", "configurePreset": "linux-cuda13" }, { "name": "linux-cuda12", "configurePreset": "linux-cuda12" }, { "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" } ], "testPresets": [ @@ -120,6 +128,11 @@ "filter": { "exclude": { "label": "bitexact_cu13" } } }, { "name": "windows-msvc-cuda13", "inherits": "test-base", "configurePreset": "windows-msvc-cuda13" }, + { + "name": "windows-msvc-cuda13-sm75", + "inherits": "test-base", + "configurePreset": "windows-msvc-cuda13-sm75" + }, { "name": "windows-msvc-cuda12", "inherits": "test-base", @@ -132,6 +145,7 @@ { "name": "linux-cuda13", "configurePreset": "linux-cuda13" }, { "name": "linux-cuda12", "configurePreset": "linux-cuda12" }, { "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 c6af1a7..674a370 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.0. 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`: fp16 tensor-core GEMMs, INT8/INT4 ConvRot files (INT4 through the INT8 unpack fallback), exact attention and the VAE then run on sm_75 kernels, and an unset attention resolves to `exact`. The full support matrix is in [docs/integration.md](docs/integration.md#support-status). +Version 2.1.0. 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 and the remaining gaps. 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). @@ -14,6 +14,7 @@ moss-infer generate --model MOSS-Image/sft --dit moss_image_sft_int8_convrot.saf ``` - [docs/integration.md](docs/integration.md): building the engine into an application. Start here. +- [docs/sm75.md](docs/sm75.md): Turing (compute capability 7.5) support: build and package, what runs, measured speed on an RTX 2080, 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. @@ -24,6 +25,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` and `windows-msvc-cuda12` ([docs/integration.md](docs/integration.md#getting-the-binaries)). Add `-DMOSSINFER_SM75=ON` to also build the Turing (sm_75) kernels and accept compute capability 7.5 devices; without it the engine is Ampere-or-newer as before. +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. Licensed under Apache-2.0 (`LICENSE`, `NOTICE`, `THIRD_PARTY_NOTICES.md`). diff --git a/cmake/modules/packaging.cmake b/cmake/modules/packaging.cmake index c51d030..9b491f5 100644 --- a/cmake/modules/packaging.cmake +++ b/cmake/modules/packaging.cmake @@ -288,7 +288,12 @@ set(CPACK_PACKAGE_NAME mossinfer) set(CPACK_PACKAGE_VENDOR MossInfer) set(CPACK_PACKAGE_VERSION ${PROJECT_VERSION}) set(CPACK_PACKAGE_DESCRIPTION_SUMMARY "Native MOSS-Image inference engine (C API, moss-infer, moss-server)") -set(CPACK_PACKAGE_FILE_NAME mossinfer-${PROJECT_VERSION}-${_mi_os}-cu${CUDAToolkit_VERSION_MAJOR}-${_mi_arch}) +set(_mi_sm75_suffix "") +if(MOSSINFER_SM75) + set(_mi_sm75_suffix -sm75) +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) set(CPACK_PACKAGE_CHECKSUM SHA256) include(CPack) diff --git a/docs/c-api.md b/docs/c-api.md index 775de8b..2e8f667 100644 --- a/docs/c-api.md +++ b/docs/c-api.md @@ -147,10 +147,10 @@ The caller chooses the tier; the engine runs it or refuses it. The loaded weight |---|---| | `AUTO` (0) | reserved since 2.0: refused by `mi_engine_create()` (`MI_ERR_INVALID_ARGUMENT`, the message names the choices) | | `REFERENCE` (1) | the DiT's bf16 GEMMs with exact attention. Refused for a quantized DiT or text encoder file and for an 8-bit `attention`. | -| `HIGH` (2) | the loaded DiT's GEMMs (INT8 or bf16) with 8-bit attention. Refused for a 4-bit DiT, and for a bf16 DiT with exact attention (that is `REFERENCE`). | -| `TURBO` (3) | a 4-bit DiT file (`dit_file`: INT4 ConvRot, or NVFP4 on an RTX 50) with 8-bit attention. Refused for any other DiT. | +| `HIGH` (2) | the loaded DiT's GEMMs (INT8 or bf16) with 8-bit attention, or exact attention where the build has no 8-bit kernel (sm_75). Refused for a 4-bit DiT, and for a bf16 DiT with exact attention (that is `REFERENCE`). | +| `TURBO` (3) | a 4-bit DiT file (`dit_file`: INT4 ConvRot, or NVFP4 on an RTX 50) with 8-bit attention (exact where the build has no 8-bit kernel, e.g. sm_75). Refused for any other DiT. | -`HIGH` and `TURBO` default to `Q8_FP8` where the GPU has FP8 tensor cores (compute capability 8.9 and later), else `Q8_FP16`; `attention` overrides it. Capability refusals stay: an NVFP4 DiT needs FP4 tensor cores (`MI_ERR_MODEL` on other GPUs), an 8-bit `attention` the GPU or build cannot run is refused, and so are `HIGH` and `TURBO` on a GPU or build without any 8-bit kernel unless `attention` is `EXACT` (never a silent fallback). `mi_result_tier()` and the report say what ran. +`HIGH` and `TURBO` default to `Q8_FP8` where the GPU has FP8 tensor cores (compute capability 8.9 and later), else `Q8_FP16`; `attention` overrides it. Capability refusals stay: an NVFP4 DiT needs FP4 tensor cores (`MI_ERR_MODEL` on other GPUs), an 8-bit `attention` the GPU or build cannot run is refused, and so are `HIGH` and `TURBO` on a GPU or build without any 8-bit kernel unless `attention` is `EXACT` — except that an **unset** `attention` there resolves to `EXACT` (sm_75), and the tier result and the report say so. `mi_result_tier()` and the report say what ran. | `mi_attention` | Pipeline name | |---|---| diff --git a/docs/integration.md b/docs/integration.md index 2cb8dfd..cc7329d 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -18,18 +18,18 @@ 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 flavours; 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). | +| 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). | | 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 | -| Quality tier | required, the caller's choice: `reference` (the DiT's bf16 GEMMs, exact attention), `high` (the loaded DiT's INT8 or bf16 GEMMs, 8-bit attention), `turbo` (a 4-bit DiT: INT4, or NVFP4 on an RTX 50, 8-bit attention); a tier the loaded files cannot run is refused with the reason, and the tier that ran is reported | +| Quality tier | required, the caller's choice: `reference` (the DiT's bf16 GEMMs, exact attention), `high` (the loaded DiT's INT8 or bf16 GEMMs, 8-bit attention), `turbo` (a 4-bit DiT: INT4, or NVFP4 on an RTX 50, 8-bit attention); a tier the loaded files cannot run is refused with the reason, and the tier that ran is reported. On an `MOSSINFER_SM75` build all three run; 8-bit attention becomes `exact` ([sm75.md](sm75.md#tiers-on-sm_75)) | | Image input | img2img (an init image and a denoise strength) and inpainting (a mask), with crop / fit / stretch resizing; pixels (library) or PNG / JPEG files (both surfaces) | | Hires fix | a second pass at a larger size from the first pass upscaled, with its own strength and steps | | Distill bundles | few-step distilled checkpoints shipped as a bundle folder (time embedding, LoRA, data-space sampler parameters), applied per engine and per generation | | Control adapters | a control adapter file (a LoRA with a widened DiT input embedder, e.g. the MOSS-Image color / depth control) with one conditioning image per kind, each with a strength and a step window; pixels (library) or PNG / JPEG files (both surfaces). Opt-in. | | Lossy accelerators | CFG truncation and a step cache (First-Block Cache), off by default. Opt-in: they change the image. | -| Attention mode | exact (FlashAttention-2) or one of four 8-bit modes, as an override of the tier's default | +| Attention mode | exact (FlashAttention-2) or one of four 8-bit modes, as an override of the tier's default. An `-DMOSSINFER_SM75=ON` build has only `exact` (its own m16n8k8 kernel); the tier's 8-bit default resolves to `exact` with the reason saying so ([sm75.md](sm75.md#not-supported-on-sm75-yet)) | | LoRAs | a list of files with strengths per generation, applied at run time: kohya, PEFT / diffusers, LyCORIS LoHa and LoKr, and full deltas; per-file reports and warnings | | Samplers and schedulers | 9 samplers (Euler, AB2, Euler ancestral, Heun, DPM++ 2M, DPM++ 2M SDE, ER-SDE, RES multistep, data-space), 7 schedulers, noise layout and source, eta, s_noise, flow multiplier | | Preview | a latent preview image during sampling, every N steps (library: a callback with RGB8; server: a PNG rewritten in place and a `preview` event) | @@ -194,7 +194,7 @@ On top of that, the weights are read through the operating system's page cache. ## Getting the binaries -**From a release.** CI builds two archives per OS and flavour: `mossinfer---cu<12|13>-.zip` (the full package) and `mossinfer---cu<12|13>--system-cuda.zip` (the same without cuBLASLt); Linux also gets `.tar.gz` of both. A GitHub release attaches the archives with one `SHA256SUMS` file. Release candidates (tags `vX.Y.Z-rcN`) are GitHub pre-releases; their archives and binaries carry the plain `X.Y.Z` version. See [Package contents](#package-contents). +**From a release.** CI builds two archives per OS and flavour: `mossinfer---cu<12|13>-.zip` (the full package) and `mossinfer---cu<12|13>--system-cuda.zip` (the same without cuBLASLt); Linux also gets `.tar.gz` of both. The Windows CUDA 13 SM75 preset additionally produces `mossinfer--windows-cu13-x86_64-sm75.zip`, a full package with Turing support. A GitHub release attaches the archives with one `SHA256SUMS` file. Release candidates (tags `vX.Y.Z-rcN`) are GitHub pre-releases; their archives and binaries carry the plain `X.Y.Z` version. See [Package contents](#package-contents). **From source** (Linux shown): @@ -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. +- **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. ### Package contents @@ -226,6 +227,7 @@ cmake --install build-linux-cuda13 --component Unspecified --prefix build-linux- **Two variants per flavour.** - **Full** (`mossinfer-…-cu13-x86_64.zip`): bundles cuBLASLt and the CUDA EULA. It runs on any machine with a suitable driver, and uses its own cuBLASLt by default (`auto`, see [CUDA libraries](#cuda-libraries)): results are reproducible across machines. +- **SM75 full package** (`mossinfer-…-cu13-x86_64-sm75.zip`): the Windows CUDA 13 package built with `MOSSINFER_SM75=ON`; it includes bundled CUDA libraries and adds Turing support (compute capability 7.5). Keep this suffix when publishing assets so launchers can distinguish it from the standard package. - **System CUDA** (`mossinfer-…-cu13-x86_64-system-cuda.zip`): the same files without cuBLASLt and the EULA, several hundred MB smaller. It runs only where a compatible cuBLASLt is installed. It is made from the release presets, which link the CUDA runtime statically; a tree installed the same way from a build with a shared cudart (the dev preset) also needs a system cudart. - **Users with a CUDA toolkit** can use its cuBLASLt instead in two ways: download the smaller `-system-cuda` archive, or keep the full one and set `MOSSINFER_CUDA_LIBS=system`. A system copy of another cuBLASLt version can change result bits (see Reproducibility below). - **Which to offer:** the launcher decides from the probe. Run the probe of any installed copy with `MOSSINFER_CUDA_LIBS=system` (`hello`'s `probe.cublaslt`, or `mi_probe()`): when `source` is `"system"`, the machine has a compatible copy and the `-system-cuda` archive works there; otherwise offer the full one. A launcher that installs once for many machines ships the full archive. @@ -330,6 +332,11 @@ RTX 5090, SFT, 1024², 50 steps, every model resident: The gain holds at 2048²: INT8 `high` runs 1.12 s per step against 2.77 s for bf16. 8-bit attention runs at every size; `attention` picks another kernel (`exact`, or the fp32-accumulating variants) when the caller wants one. +An `-DMOSSINFER_SM75=ON` build on an RTX 2080 8 GiB (SFT, every model resident or streamed as planned) is slower but +the same configuration ladder works: 0.80 s per 256² step for the bf16 bundle at `reference`, 0.19 s for an INT8 DiT at +`high` (6.5 GiB peak), 0.22 s for an INT4 DiT at `turbo` (3.9 GiB peak); at 1024² the INT8 `high` runs 3.8 s per step and +the VAE fp32 decode costs 10–12 s. Details in [sm75.md](sm75.md#performance). + **Memory modes** trade host RAM for speed. `balanced` (the default) keeps the process near 8 GiB of private RAM, `min_memory` near 4 GiB (more streaming and parking, the lowest peaks), and `max_speed` uses what the machine has free (the text encoder held in RAM, pinned where there is room, and a larger embedding cache, so a new prompt starts sooner). **VRAM budgets.** A DiT that does not fit stays exact: its blocks stream to the GPU every step, from the page cache or a pinned copy in RAM, and each step pays the part of the copy not hidden behind compute. Streaming from disk instead of RAM costs several seconds per step, so leave RAM free for the page cache. An INT8 DiT moves half the bytes of the bf16 one: on 8 to 12 GiB cards it is also the fast choice. Leave `vram_budget_bytes` at 0 unless the GPU is shared; the engine plans against the memory that is actually free. @@ -360,6 +367,13 @@ 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. +**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 +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). + ## Driving the sidecar The full protocol, with an example of every message, is in [server-protocol.md](server-protocol.md). A launcher needs this much of it. diff --git a/docs/sm75.md b/docs/sm75.md new file mode 100644 index 0000000..9c69801 --- /dev/null +++ b/docs/sm75.md @@ -0,0 +1,99 @@ +# Turing (sm_75) support + +MossInfer builds for NVIDIA compute capability 8.x (Ampere, Ada) and 12.x (Blackwell) by default. A build configured +with `-DMOSSINFER_SM75=ON` also carries a full set of Turing kernels and accepts compute capability 7.5 devices (RTX +20xx, T4, GTX 16xx), so a card such as an RTX 2080 8 GiB can run the whole pipeline: text encoder, DiT and VAE. + +Turing has no bf16 tensor cores, no `cp.async`, no TF32 and only 64 KiB of shared memory per block, so the sm_75 +kernel family is not a flag on the sm_80 one: it is a second implementation of every tensor-core kernel, compiled for +`sm_75` alone and dispatched at run time by the device's compute capability. + +## Build + +The Windows CUDA 13 preset builds a complete package and names it with a `-sm75` suffix: + +```sh +cmake --preset windows-msvc-cuda13-sm75 +cmake --build --preset windows-msvc-cuda13-sm75 +ctest --preset windows-msvc-cuda13-sm75 -L '^cpu$' +cpack --config build-windows-msvc-cuda13-sm75/CPackConfig.cmake -B build-windows-msvc-cuda13-sm75 +# -> build-windows-msvc-cuda13-sm75/mossinfer-2.0.0-windows-cu13-x86_64-sm75.zip +``` + +Any other preset takes the option too (`-DMOSSINFER_SM75=ON`), e.g. a Linux build from a custom preset; the ordinary +presets stay Ampere-and-newer. Without the option the engine refuses compute capability 7.5 exactly as before. + +## What runs + +| Component | sm_80+ build (unchanged) | sm_75 build | +|---|---|---| +| bf16 linears (DiT blocks and globals, text refiner, text encoder) | cuBLASLt bf16 tensor cores | fp16 tensor cores through cuBLASLt: weights are converted to fp16 once at load (or per streamed block), activations are converted in row chunks, accumulation stays fp32, the output is converted back to bf16 with the bias added in fp32 | +| Exact attention | FlashAttention-2 (bf16, `cp.async`) | a dedicated `m16n8k8` fp16 kernel with an online softmax (head_dim 128 non-causal varlen GQA for the DiT, head_dim 256 causal for the text encoder) | +| INT8 ConvRot GEMMs (`high` tier) | CUTLASS `m16n8k32` | CUTLASS `m8n8k16` (Turing's INT8 MMA), same epilogues; bit-exact against a host int32 reference | +| INT4 ConvRot GEMMs (`turbo` tier) | CUTLASS `m16n8k64` | the engine's existing INT4 → INT8 unpack fallback (bit-exact, the same path Hopper and Blackwell use) | +| 8-bit attention (`q8_fp16`, `q8_fp8`) | SageAttention kernels | not built: an unset `attention` resolves to `exact` with the reason saying so (the tier is otherwise unchanged) | +| VAE decode | TF32/bf16 CUTLASS implicit-GEMM convolutions, TF32 attention | fp32 SIMT convolutions and a fp32 attention kernel (reference precision; slower, see below) | +| GatedDeltaNet (text encoder) | chunked kernel (~85 KiB of shared memory) | a per-token recurrent kernel that fits 33 KiB; the same rule, fp32 | +| Other kernels (norm, RoPE, elementwise, samplers, group norm, latent preview, …) | compiled for sm_80+ | compiled for sm_75 and sm_80+ from the same source (pure SIMT) | + +The DiT blocks run op by op on sm_75: the fused bf16 CUTLASS GEMMs need sm_80+ tensor cores, so the fused route and +its CUDA graphs are not used. Every block path produces the same image, so this only affects speed. + +### Tiers on sm_75 + +`reference` (bf16 bundle), `high` (INT8 ConvRot DiT) and `turbo` (INT4 ConvRot DiT) all run. As on any device without +an 8-bit attention kernel, an unset `attention` resolves to `exact` and the report's `quality.reason` says +"(no 8-bit attention kernel on this device)"; a caller may still pass `attention=exact` explicitly. NVFP4 DiTs are +refused as on every non-Blackwell GPU. + +## Performance + +Measured on an RTX 2080 8 GiB (Turing, 46 SMs), MOSS-Image SFT, `balanced` mode, the bf16 bundle with the text encoder +streamed: + +| Configuration | Size | Denoise step | Notes | +|---|---|---|---| +| bf16 bundle, `reference` | 256² | 0.80 s | 6.9 GiB device peak | +| INT8 DiT (`moss-infer quantize … --format int8`), `high` | 256² | **0.19 s** | 6.5 GiB peak, DiT resident | +| INT4 ConvRot DiT, `turbo` | 256² | **0.22 s** | 3.9 GiB peak | +| INT8 DiT, `high` | 512² | 0.63 s | 20 steps in 22 s total | +| INT8 DiT, `high` | 1024² | 3.8 s | 8 steps in 51 s total, 6.9 GiB peak | + +The VAE decodes in fp32 on sm_75: about 0.4 s at 256², 1.6 s at 512² and 10–12 s at 1024². The bf16 text encoder is +roughly 3–5 s per prompt when streamed (its ~16 GB do not fit 8 GiB of VRAM); it is planned like on any other device. + +## Numerics + +- bf16 → fp16 is exact for values in fp16's normal range, and fp16 → bf16 recovers the original bf16 value, so the + converted operands are the same numbers the sm_80 kernels see. Accumulation stays fp32 where the sm_80 kernel + accumulates in fp32. +- Attention, GEMMs and the VAE therefore differ from an sm_80 run only in accumulation order and rounding, not in the + values fed in. INT8 GEMM results are bit-identical (exact int32 accumulation). +- The reference vectors under `tests/data` are recorded on sm_89; as on sm_120, tests that compare against them run + with noise-class gates off that architecture. +- Values outside fp16's range (|x| > 65504 or |x| < ~6e-5) do not survive the bf16 → fp16 conversion; the model's + activations and weights are inside that range in practice. + +## Verification + +- `tools/sm75_check.cu` (not installed, built with `MOSSINFER_SM75=ON`) runs the Turing kernels against host + references: dtype conversion, exact attention at both head dimensions and engine-shaped varlen/GQA cases, the fp16 + GEMM, the INT8 GEMM (bit-exact), and the VAE convolution and attention. +- `tests/unit/test_sm75.cpp` covers the build's presence and the in-place conversions on a device. +- End to end on an RTX 2080: `moss-infer generate` and `moss-server` (protocol 2) produce the same PNG bytes for the + `reference`, `high` and `turbo` tiers; the report's `quality.gpu_arch` is `sm_75`. +- A build with `MOSSINFER_SM75=OFF` compiles and links unchanged and keeps the same CPU test baseline. + +## Not supported on sm_75 yet + +- The 8-bit attention kernels (`q8_fp16`, `q8_fp8`): the tier default resolves to `exact`; passing an 8-bit mode + explicitly is refused. +- Native INT4 MMA and NVFP4: INT4 files run through the INT8 unpack fallback (correct and bit-exact; the planner + prices it at about 1.5× a native INT4 step); NVFP4 files are refused on every non-Blackwell GPU. +- The fused DiT blocks and their CUDA graphs (the blocks run op by op). +- img2img and inpainting: the VAE **encoder** is still an sm_80 kernel, so only text-to-image and hires fix run + (`hires fix` uses the decoder). The image input parameters are accepted by the API but a request that needs the + encoder fails on sm_75. +- Runtime LoRAs and control adapters: their GEMM kernels are built for sm_80+ only. + +The VAE decode is fp32 SIMT on sm_75, so a 1024² decode costs about 10–12 s; a future fp16 path would cut that. From b5b0d2a7d246241c5ebd9050e65126556cd450ff Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sat, 31 Oct 2026 16:15:51 +0800 Subject: [PATCH 3/6] sm_75: refuse img2img/inpainting and LoRAs with not_supported The VAE encoder's convolutions and the LoRA GEMMs are still sm_80 kernels, so an sm_75 build used to fail in the middle of a request with a missing kernel image. Refuse both in Pipeline::run, before anything is planned, read or uploaded: - pipeline::NotSupportedError, mapped to MI_ERR_NOT_SUPPORTED by the C API, so moss-server reports `not_supported` in the done event. - img2img / inpainting: a request with an init image (or a mask) says the VAE encoder has no sm_75 kernel and that txt2img and the hires fix still work. - runtime LoRAs and control adapters (a LoRA): refused with the same code. - docs: sm75.md, integration.md and c-api.md say the three features are refused up front instead of failing mid-request. Verified on an RTX 2080: `moss-infer generate --init-image` and `--lora` exit 1 with the message and write nothing, moss-server answers `not_supported` in done, and txt2img still runs at the reference and high tiers. The CPU suite keeps its baseline (the same nine pre-existing Windows path failures). --- docs/c-api.md | 1 + docs/integration.md | 6 +++--- docs/sm75.md | 9 ++++++--- src/api/c_api.cpp | 2 ++ src/pipeline/pipeline.cpp | 15 +++++++++++++++ src/pipeline/pipeline.h | 8 ++++++++ 6 files changed, 35 insertions(+), 6 deletions(-) diff --git a/docs/c-api.md b/docs/c-api.md index 2e8f667..94ba51c 100644 --- a/docs/c-api.md +++ b/docs/c-api.md @@ -289,6 +289,7 @@ Limits, checked before any pixel is read: each side 1 to 16384, at most 64,000,0 - `denoise_strength` unset with an init image is 0.75. An explicit 1 with an init image and no mask is txt2img: the image is not used, and `MI_TEXT_IMAGE_NOTE` says so. 0 runs no step: the result is the image's VAE round trip. - The steps that run: of a pass's N scheduled steps (after snapping), the first `int(N - N * strength)` are skipped. `mi_result_steps_scheduled()` and the progress callback's `total_steps` count the steps that run, both passes of a hires fix included. - The `DATASPACE` sampler cannot start from an image (img2img, inpainting, hires fix). +- An `MOSSINFER_SM75` build has no VAE encoder kernel: img2img and inpainting are refused up front with `MI_ERR_NOT_SUPPORTED`, before anything is loaded or planned; txt2img and the hires fix still run. Runtime LoRAs and control adapters are refused the same way (their GEMM kernels are sm_80+) ([sm75.md](sm75.md#not-supported-on-sm75-yet)). | `mi_mask_encoding` | Values | |---|---| diff --git a/docs/integration.md b/docs/integration.md index cc7329d..9b4200d 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -24,13 +24,13 @@ MossInfer is a low-level engine: it runs what the caller asks for, refuses what | 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 | | Quality tier | required, the caller's choice: `reference` (the DiT's bf16 GEMMs, exact attention), `high` (the loaded DiT's INT8 or bf16 GEMMs, 8-bit attention), `turbo` (a 4-bit DiT: INT4, or NVFP4 on an RTX 50, 8-bit attention); a tier the loaded files cannot run is refused with the reason, and the tier that ran is reported. On an `MOSSINFER_SM75` build all three run; 8-bit attention becomes `exact` ([sm75.md](sm75.md#tiers-on-sm_75)) | -| Image input | img2img (an init image and a denoise strength) and inpainting (a mask), with crop / fit / stretch resizing; pixels (library) or PNG / JPEG files (both surfaces) | +| Image input | img2img (an init image and a denoise strength) and inpainting (a mask), with crop / fit / stretch resizing; pixels (library) or PNG / JPEG files (both surfaces). Needs the VAE encoder, so an `MOSSINFER_SM75` build refuses it with `not_supported` ([sm75.md](sm75.md#not-supported-on-sm75-yet)) | | Hires fix | a second pass at a larger size from the first pass upscaled, with its own strength and steps | | Distill bundles | few-step distilled checkpoints shipped as a bundle folder (time embedding, LoRA, data-space sampler parameters), applied per engine and per generation | -| Control adapters | a control adapter file (a LoRA with a widened DiT input embedder, e.g. the MOSS-Image color / depth control) with one conditioning image per kind, each with a strength and a step window; pixels (library) or PNG / JPEG files (both surfaces). Opt-in. | +| Control adapters | a control adapter file (a LoRA with a widened DiT input embedder, e.g. the MOSS-Image color / depth control) with one conditioning image per kind, each with a strength and a step window; pixels (library) or PNG / JPEG files (both surfaces). Opt-in. It is a LoRA, so an `MOSSINFER_SM75` build refuses it with `not_supported` | | Lossy accelerators | CFG truncation and a step cache (First-Block Cache), off by default. Opt-in: they change the image. | | Attention mode | exact (FlashAttention-2) or one of four 8-bit modes, as an override of the tier's default. An `-DMOSSINFER_SM75=ON` build has only `exact` (its own m16n8k8 kernel); the tier's 8-bit default resolves to `exact` with the reason saying so ([sm75.md](sm75.md#not-supported-on-sm75-yet)) | -| LoRAs | a list of files with strengths per generation, applied at run time: kohya, PEFT / diffusers, LyCORIS LoHa and LoKr, and full deltas; per-file reports and warnings | +| LoRAs | a list of files with strengths per generation, applied at run time: kohya, PEFT / diffusers, LyCORIS LoHa and LoKr, and full deltas; per-file reports and warnings. Their GEMM kernels are sm_80+, so an `MOSSINFER_SM75` build refuses a request that carries LoRAs with `not_supported` | | Samplers and schedulers | 9 samplers (Euler, AB2, Euler ancestral, Heun, DPM++ 2M, DPM++ 2M SDE, ER-SDE, RES multistep, data-space), 7 schedulers, noise layout and source, eta, s_noise, flow multiplier | | Preview | a latent preview image during sampling, every N steps (library: a callback with RGB8; server: a PNG rewritten in place and a `preview` event) | | PNG metadata | the generation settings embedded in the PNG (`generation` record, schema `mossinfer.generation/1`, plus `parameters`), the prompts optional | diff --git a/docs/sm75.md b/docs/sm75.md index 9c69801..03ce1ca 100644 --- a/docs/sm75.md +++ b/docs/sm75.md @@ -80,6 +80,8 @@ roughly 3–5 s per prompt when streamed (its ~16 GB do not fit 8 GiB of VRAM); references: dtype conversion, exact attention at both head dimensions and engine-shaped varlen/GQA cases, the fp16 GEMM, the INT8 GEMM (bit-exact), and the VAE convolution and attention. - `tests/unit/test_sm75.cpp` covers the build's presence and the in-place conversions on a device. +- The two refusals are exercised on the RTX 2080: `moss-infer generate --init-image …` and `--lora …` end with + `not_supported` (CLI exit status 1, nothing written), and `moss-server` reports the same code in the `done` event. - End to end on an RTX 2080: `moss-infer generate` and `moss-server` (protocol 2) produce the same PNG bytes for the `reference`, `high` and `turbo` tiers; the report's `quality.gpu_arch` is `sm_75`. - A build with `MOSSINFER_SM75=OFF` compiles and links unchanged and keeps the same CPU test baseline. @@ -92,8 +94,9 @@ roughly 3–5 s per prompt when streamed (its ~16 GB do not fit 8 GiB of VRAM); prices it at about 1.5× a native INT4 step); NVFP4 files are refused on every non-Blackwell GPU. - The fused DiT blocks and their CUDA graphs (the blocks run op by op). - img2img and inpainting: the VAE **encoder** is still an sm_80 kernel, so only text-to-image and hires fix run - (`hires fix` uses the decoder). The image input parameters are accepted by the API but a request that needs the - encoder fails on sm_75. -- Runtime LoRAs and control adapters: their GEMM kernels are built for sm_80+ only. + (`hires fix` uses the decoder). A request with `init_image` (or a `mask`) is refused up front with the code + `not_supported` and a message saying why, so nothing is loaded or planned for it. +- Runtime LoRAs and control adapters: their GEMM kernels are built for sm_80+ only. A request that carries `loras` or a + control adapter is refused up front with `not_supported`, again before any weight is read or uploaded. The VAE decode is fp32 SIMT on sm_75, so a 1024² decode costs about 10–12 s; a future fp16 path would cut that. diff --git a/src/api/c_api.cpp b/src/api/c_api.cpp index 7031a3d..05bcae4 100644 --- a/src/api/c_api.cpp +++ b/src/api/c_api.cpp @@ -565,6 +565,8 @@ mi_status guarded(F&& body) noexcept { return fail(MI_ERR_INVALID_ARGUMENT, e.what()); } catch (const mi::api::NotSupportedError& e) { return fail(MI_ERR_NOT_SUPPORTED, e.what()); + } catch (const mi::pipeline::NotSupportedError& e) { + return fail(MI_ERR_NOT_SUPPORTED, e.what()); } catch (const mi::planner::BudgetError& e) { return fail(MI_ERR_OUT_OF_MEMORY, e.what()); } catch (const mi::runtime::BudgetError& e) { diff --git a/src/pipeline/pipeline.cpp b/src/pipeline/pipeline.cpp index 84ee467..3aebb51 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -1506,6 +1506,21 @@ Result Pipeline::run(const Request& req, const GenerateControl& control, bool pr Impl& m = *impl_; check_user_paths(m.options, req); validate(req); + // Two features still need sm_80+ kernels and are refused here, before anything is planned, read or uploaded: the VAE + // encoder's convolutions (img2img and inpainting) 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. + if (kernels::sm75::turing()) { + if (req.init_image != nullptr) { + throw NotSupportedError("img2img and inpainting encode the input with the VAE encoder, which has no sm_75 kernel " + "in this build; text-to-image works, and so does the hires fix (it only decodes)"); + } + if (!req.loras.empty()) { + throw NotSupportedError("runtime LoRAs need sm_80+ GEMM kernels; this sm_75 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"); + } + } 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.h b/src/pipeline/pipeline.h index eb6fcfd..f3415d7 100644 --- a/src/pipeline/pipeline.h +++ b/src/pipeline/pipeline.h @@ -42,6 +42,7 @@ #include #include #include +#include #include #include #include @@ -71,6 +72,13 @@ inline constexpr std::int64_t kMaxPixels = 4'200'000; // Upper bound for PipelineOptions sizes taken from user input (vram_option_bytes). inline constexpr std::size_t kMaxVramBytes = std::size_t{1024} << 30; +// A feature this build or device does not implement (MI_ERR_NOT_SUPPORTED through the C API; the server reports it as +// `not_supported`). Raised before any work is planned, e.g. for img2img and runtime LoRAs on an sm_75 build, whose VAE +// encoder and LoRA GEMMs are still sm_80 kernels. +struct NotSupportedError : std::runtime_error { + using std::runtime_error::runtime_error; +}; + // One generation, mirroring the assets' parameters.json (gradio_app.py GenerationConfig) plus engine choices. struct Request { std::filesystem::path model; // bundle directory (base or sft): {text_encoder, tokenizer, transformer, vae} From 88b9fc9dd242a1783d168b1c909490d5eda8b7eb Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sat, 31 Oct 2026 17:14:39 +0800 Subject: [PATCH 4/6] sm_75: the INT8 SiTU-GLU epilogue's factor tables The sm_80 tables are filled by a kernel built for sm_80+ only, so the Turing epilogue hit cudaErrorNoKernelImageForDevice when the high tier ran a SiTU-GLU site. Fill a second, per-device copy with a kernel built for 75; the factors are the same functions, so a lookup is bit-identical to the sm_80 table. --- src/kernels/sm75/int8_gemm_sm75.cu | 61 +++++++++++++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/src/kernels/sm75/int8_gemm_sm75.cu b/src/kernels/sm75/int8_gemm_sm75.cu index 1a8dc71..74026f1 100644 --- a/src/kernels/sm75/int8_gemm_sm75.cu +++ b/src/kernels/sm75/int8_gemm_sm75.cu @@ -2,6 +2,9 @@ // MMA and tiles that fit 64 KiB of shared memory. #include "kernels/sm75/int8_sm75.h" +#include +#include +#include #include #include "kernels/quant/int8_gemm_impl.cuh" @@ -11,6 +14,62 @@ namespace qd = quant::detail; namespace { +// The SiTU factor tables of the Turing epilogue: the sm_80 tables (int8_gemm_situ.cu) are filled by a kernel built for +// sm_80+ only, so the sm_75 path keeps its own 256 KiB copy, filled once per device by a kernel built for 75. 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_sm75_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::sm75: device ordinal out of range"); + Device& d = devices[dev]; + const auto fill = [&] { + constexpr int kThreads = 256; + situ_tables_sm75_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)}; +} + using Sm75Tile = qd::TileCfg<64, 128, 64, 32, 64, 2>; // threadblock 64x128, warp 32x64, K 64, 2 stages: 24 KiB using Sm75Instr = cutlass::gemm::GemmShape<8, 8, 16>; @@ -58,7 +117,7 @@ void situ_glu(const qd::Problem& p, const quant::Int8Operands& a, __nv_bfloat16* auto* d = reinterpret_cast(out); const std::int64_t ldd = 2 * ldo; const typename N::SituGlu::Arguments cb{ - {{}, {const_cast(a.x_scale), 0.f, {_1{}, _0{}, p.m}}, {a.w_scale}, {}}, + {{}, {const_cast(a.x_scale), 0.f, {_1{}, _0{}, p.m}}, {a.w_scale}, situ_tables(stream)}, {d, {ldd, _1{}, ldd * p.m}}}; qd::launch("sm75::int8_gemm_situ_glu", p, a, cb, stream); } From 258093694e8b85f018114d9942957effffd4a3e8 Mon Sep 17 00:00:00 2001 From: xycld Date: Sat, 31 Oct 2026 06:11:37 -0400 Subject: [PATCH 5/6] sm_75: refuse every VAE-encoder request; probe minimum 7.5; device checks once per device - The hires fix and reference images also encode with the VAE encoder (sm_80+), so on Turing they are refused with not_supported up front, like img2img (the hires message used to say it worked; it failed mid-request). - conv2d_f32_sm75: a 1x1 filter staged 9 taps and read 8*C floats past the end of w. - convert_3d_inplace: rows above 65535 returned without converting (attention then read bf16 as fp16); grid.x holds 2^31 - 1 rows, and more throws. - probe: an sm_75 build's minimum is 7.5, not 7.0 (Volta passed the capability check). - chunk_rule.cu and flash_attention.cu use the cached kernels::sm75::turing(), and the GDN branch is compiled in MOSSINFER_SM75 builds only: no per-call device query on sm_80+. - The unset-attention exact fallback is keyed on the Turing device, so an sm_80+ build without the 8-bit attention kernels still refuses, as before. - test_sm75: the conversion tests skip off Turing (the kernels are 75-real), and the strided test checks the converted state, so a skipped conversion fails it. --- cmake/modules/sm75.cmake | 3 ++- docs/c-api.md | 2 +- docs/integration.md | 2 +- docs/sm75.md | 7 +++--- src/backend/probe.cpp | 8 +++---- src/backend/probe.h | 2 +- src/kernels/attention/flash_attention.cu | 16 +++++--------- src/kernels/gdn/chunk_rule.cu | 15 ++++++++----- src/kernels/sm75/convert.cu | 3 ++- src/kernels/sm75/vae_sm75.cu | 3 ++- src/pipeline/pipeline.cpp | 24 +++++++++++++------- tests/unit/test_sm75.cpp | 28 +++++++++++++++++++++--- 12 files changed, 72 insertions(+), 41 deletions(-) diff --git a/cmake/modules/sm75.cmake b/cmake/modules/sm75.cmake index 7985703..472b57c 100644 --- a/cmake/modules/sm75.cmake +++ b/cmake/modules/sm75.cmake @@ -8,7 +8,7 @@ # - the probe accepts compute capability 7.5 (MOSSINFER_HAVE_SM75 on probe.cpp, and 75-real in its architecture list, # cmake/modules/backend.cmake), # - the portable SIMT translation units are also compiled for sm_75, -# - flash_attention.cu / gemm.cpp dispatch into the sm_75 kernels, +# - flash_attention.cu / gemm.cpp / chunk_rule.cu dispatch into the sm_75 kernels, # - the sm75_check bring-up tool is built. # Tolerates the sources being filtered out by MOSSINFER_ONLY. set(_mi_sm75_sources ${MOSSINFER_SOURCES}) @@ -93,6 +93,7 @@ if(MOSSINFER_SM75 AND _mi_sm75_sources) 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") diff --git a/docs/c-api.md b/docs/c-api.md index 94ba51c..748c956 100644 --- a/docs/c-api.md +++ b/docs/c-api.md @@ -289,7 +289,7 @@ Limits, checked before any pixel is read: each side 1 to 16384, at most 64,000,0 - `denoise_strength` unset with an init image is 0.75. An explicit 1 with an init image and no mask is txt2img: the image is not used, and `MI_TEXT_IMAGE_NOTE` says so. 0 runs no step: the result is the image's VAE round trip. - The steps that run: of a pass's N scheduled steps (after snapping), the first `int(N - N * strength)` are skipped. `mi_result_steps_scheduled()` and the progress callback's `total_steps` count the steps that run, both passes of a hires fix included. - The `DATASPACE` sampler cannot start from an image (img2img, inpainting, hires fix). -- An `MOSSINFER_SM75` build has no VAE encoder kernel: img2img and inpainting are refused up front with `MI_ERR_NOT_SUPPORTED`, before anything is loaded or planned; txt2img and the hires fix still run. Runtime LoRAs and control adapters are refused the same way (their GEMM kernels are sm_80+) ([sm75.md](sm75.md#not-supported-on-sm75-yet)). +- An `MOSSINFER_SM75` build has no VAE encoder kernel: img2img, inpainting, reference images and the hires fix are refused up front with `MI_ERR_NOT_SUPPORTED`, before anything is loaded or planned; txt2img still runs. Runtime LoRAs and control adapters are refused the same way (their GEMM kernels are sm_80+) ([sm75.md](sm75.md#not-supported-on-sm75-yet)). | `mi_mask_encoding` | Values | |---|---| diff --git a/docs/integration.md b/docs/integration.md index 9b4200d..9c2f5e8 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -25,7 +25,7 @@ MossInfer is a low-level engine: it runs what the caller asks for, refuses what | Generation | text-to-image with the official app's parameters: prompt, negative prompt, size, steps, CFG, shift, seed, skip layers, skip-layer guidance | | Quality tier | required, the caller's choice: `reference` (the DiT's bf16 GEMMs, exact attention), `high` (the loaded DiT's INT8 or bf16 GEMMs, 8-bit attention), `turbo` (a 4-bit DiT: INT4, or NVFP4 on an RTX 50, 8-bit attention); a tier the loaded files cannot run is refused with the reason, and the tier that ran is reported. On an `MOSSINFER_SM75` build all three run; 8-bit attention becomes `exact` ([sm75.md](sm75.md#tiers-on-sm_75)) | | Image input | img2img (an init image and a denoise strength) and inpainting (a mask), with crop / fit / stretch resizing; pixels (library) or PNG / JPEG files (both surfaces). Needs the VAE encoder, so an `MOSSINFER_SM75` build refuses it with `not_supported` ([sm75.md](sm75.md#not-supported-on-sm75-yet)) | -| Hires fix | a second pass at a larger size from the first pass upscaled, with its own strength and steps | +| Hires fix | a second pass at a larger size from the first pass upscaled, with its own strength and steps. Encodes the upscaled image, so an `MOSSINFER_SM75` build refuses it on Turing with `not_supported` ([sm75.md](sm75.md#not-supported-on-sm75-yet)) | | Distill bundles | few-step distilled checkpoints shipped as a bundle folder (time embedding, LoRA, data-space sampler parameters), applied per engine and per generation | | Control adapters | a control adapter file (a LoRA with a widened DiT input embedder, e.g. the MOSS-Image color / depth control) with one conditioning image per kind, each with a strength and a step window; pixels (library) or PNG / JPEG files (both surfaces). Opt-in. It is a LoRA, so an `MOSSINFER_SM75` build refuses it with `not_supported` | | Lossy accelerators | CFG truncation and a step cache (First-Block Cache), off by default. Opt-in: they change the image. | diff --git a/docs/sm75.md b/docs/sm75.md index 03ce1ca..0a174f0 100644 --- a/docs/sm75.md +++ b/docs/sm75.md @@ -93,9 +93,10 @@ roughly 3–5 s per prompt when streamed (its ~16 GB do not fit 8 GiB of VRAM); - Native INT4 MMA and NVFP4: INT4 files run through the INT8 unpack fallback (correct and bit-exact; the planner prices it at about 1.5× a native INT4 step); NVFP4 files are refused on every non-Blackwell GPU. - The fused DiT blocks and their CUDA graphs (the blocks run op by op). -- img2img and inpainting: the VAE **encoder** is still an sm_80 kernel, so only text-to-image and hires fix run - (`hires fix` uses the decoder). A request with `init_image` (or a `mask`) is refused up front with the code - `not_supported` and a message saying why, so nothing is loaded or planned for it. +- img2img, inpainting, reference images and the hires fix: the VAE **encoder** is still an sm_80 kernel, so only + text-to-image runs (the hires fix's second pass encodes the upscaled image). A request with `init_image` (or a + `mask`), `references` or `hires` is refused up front with the code `not_supported` and a message saying why, so + nothing is loaded or planned for it. - Runtime LoRAs and control adapters: their GEMM kernels are built for sm_80+ only. A request that carries `loras` or a control adapter is refused up front with `not_supported`, again before any weight is read or uploaded. diff --git a/src/backend/probe.cpp b/src/backend/probe.cpp index f582712..181f478 100644 --- a/src/backend/probe.cpp +++ b/src/backend/probe.cpp @@ -131,12 +131,10 @@ void evaluate_device(DeviceInfo& dev, int driver_version, int build_cuda_version dev.supported = false; dev.reason.clear(); if (std::tie(dev.cc_major, dev.cc_minor) < std::tie(kMinComputeMajor, kMinComputeMinor)) { + // 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) + " (Ampere or newer)"; -#if MOSSINFER_HAVE_SM75 - // The sm_80 path needs Ampere; this build's sm_75 kernel family reaches down to Turing. - dev.reason = "compute capability " + cc_string(dev.cc_major, dev.cc_minor) + " is below the required 7.5 (Turing)"; -#endif + cc_string(kMinComputeMajor, kMinComputeMinor) + + (MOSSINFER_HAVE_SM75 ? " (Turing or newer)" : " (Ampere or newer)"); return; } const Flavour build = flavour_of_cuda_version(build_cuda_version); diff --git a/src/backend/probe.h b/src/backend/probe.h index a5494d2..1cb1440 100644 --- a/src/backend/probe.h +++ b/src/backend/probe.h @@ -23,7 +23,7 @@ namespace mi::backend { #define MOSSINFER_HAVE_SM75 0 #endif inline constexpr int kMinComputeMajor = MOSSINFER_HAVE_SM75 ? 7 : 8; -inline constexpr int kMinComputeMinor = 0; +inline constexpr int kMinComputeMinor = 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 5af00ab..832afd4 100644 --- a/src/kernels/attention/flash_attention.cu +++ b/src/kernels/attention/flash_attention.cu @@ -153,16 +153,12 @@ void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& int device = 0, smem_per_sm = 0, smem_per_block = 0; MI_CUDA_CHECK(cudaGetDevice(&device)); #if MOSSINFER_HAVE_SM75 - { - int major = 0; - MI_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); - if (major < 8) { - // Turing: the vendored kernel's own sm_75 traits (fp16 MMA, no cp.async) run in the sm_75 kernel family; - // q/k/v/out are converted around the launch there. - kernels::sm75::flash_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen, causal, scale, workspace, - workspace_bytes, stream); - return; - } + if (kernels::sm75::turing()) { + // Turing: the vendored kernel's own sm_75 traits (fp16 MMA, no cp.async) run in the sm_75 kernel family; + // q/k/v/out are converted around the launch there. + kernels::sm75::flash_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen, causal, scale, workspace, + workspace_bytes, stream); + return; } #endif require_sm80(device); diff --git a/src/kernels/gdn/chunk_rule.cu b/src/kernels/gdn/chunk_rule.cu index 3d78238..b16eb45 100644 --- a/src/kernels/gdn/chunk_rule.cu +++ b/src/kernels/gdn/chunk_rule.cu @@ -20,6 +20,7 @@ #include "core/varlen.h" #include "kernels/gdn/checks.h" #include "kernels/gdn/gdn.h" +#include "kernels/sm75/sm75.h" namespace mi::gdn { namespace { @@ -462,7 +463,8 @@ void allow_dynamic_smem(Kernel kernel, int bytes) { // sm_75: the chunked kernel needs ~85 KiB of shared memory (PrepSmem), over Turing's 64 KiB per block. This per-token // recurrent form computes the same rule (the chunked algorithm is an algebraic rearrangement of it) with the 128 state // columns split across two launches of 64: every state column is independent, so the split is exact. 33 KiB of dynamic -// shared memory, 64 threads per block. +// shared memory, 64 threads per block. Built in an MOSSINFER_SM75 build only (the dispatch below). +#if MOSSINFER_HAVE_SM75 constexpr int kSm75Cols = 64; template @@ -531,6 +533,7 @@ __global__ void recurrent_sm75_kernel(const __nv_bfloat16* __restrict__ qkv, std for (int i = 0; i < kHeadDim; ++i) fs[i * kHeadDim + col] = s[i * kSm75Cols + c]; } } +#endif } // namespace @@ -559,11 +562,10 @@ void chunk_gated_delta_rule(const Tensor& qkv, const Tensor& g, const Tensor& be workspace_layout(tokens, workspace, &ws); // sm_75: the chunked kernels need ~85 KiB of shared memory, over Turing's 64 KiB per block. The recurrent form above - // computes the same rule with 33 KiB; its numerics are fp32 like the chunked path's. - int sm75_device = 0, sm75_major = 0; - MI_CUDA_CHECK(cudaGetDevice(&sm75_device)); - MI_CUDA_CHECK(cudaDeviceGetAttribute(&sm75_major, cudaDevAttrComputeCapabilityMajor, sm75_device)); - if (sm75_major < 8) { + // 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()) { 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><<>>( @@ -572,6 +574,7 @@ void chunk_gated_delta_rule(const Tensor& qkv, const Tensor& g, const Tensor& be MI_CUDA_CHECK_LAUNCH(); return; } +#endif constexpr int kPrepSmem = sizeof(PrepSmem), kStateSmem = sizeof(StateSmem); allow_dynamic_smem(prep_kernel, kPrepSmem); diff --git a/src/kernels/sm75/convert.cu b/src/kernels/sm75/convert.cu index 5a332a1..af7df05 100644 --- a/src/kernels/sm75/convert.cu +++ b/src/kernels/sm75/convert.cu @@ -62,7 +62,8 @@ int blocks_for(std::int64_t n) { void convert_3d_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, bool to_f16, cudaStream_t stream) { if (rows == 0 || heads == 0 || dim == 0) return; - if (rows > kMaxBlocks) return; // the engine's varlen token counts stay far below 2^31 blocks; never reached + // One block per row on grid.x, whose limit is 2^31 - 1 (kMaxBlocks is the grid-stride cap of the flat kernels). + if (rows > 0x7fffffff) throw std::length_error("sm75::convert_3d_inplace: more rows than one grid dimension holds"); auto* base = static_cast(t.data); for (std::int64_t h = 0; h < heads; ++h) { char* p = base + h * static_cast(t.stride[1]) * 2; diff --git a/src/kernels/sm75/vae_sm75.cu b/src/kernels/sm75/vae_sm75.cu index 02f1ff6..24356da 100644 --- a/src/kernels/sm75/vae_sm75.cu +++ b/src/kernels/sm75/vae_sm75.cu @@ -48,7 +48,8 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest for (int i = tid; i < kTK * 9 * kTC; i += kConvThreads) { const int cc = i % kTC, t = i / kTC, rs = t % 9, kk = t / 9; const int kch = kb * kTK + kk, gc = c0 + cc; - ws[i] = (kch < K && gc < C) + // The tile holds 9 taps; a 1x1 filter (R = 1) has one, and reading the other 8 would run past the end of w. + ws[i] = (kch < K && gc < C && rs < R * R) ? w[(static_cast(kch) * R + rs / R) * R * C + (rs % R) * C + gc] : 0.f; } diff --git a/src/pipeline/pipeline.cpp b/src/pipeline/pipeline.cpp index 3aebb51..266130d 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -939,11 +939,10 @@ Pipeline::Impl::CallModels Pipeline::Impl::setup_models(const Request& req, Resu // the LoRA files too). The lossy accelerators are an explicit speed opt-in: the tier is resolved without them. const std::filesystem::path te_file = tes.model ? tes.path : std::filesystem::path(); const auto supported = [](planner::AttentionMode mode) { return moss_dit::attention_supported(dit_attention(mode)); }; - // sm_75 (and any device without an 8-bit attention kernel): the tier's default 8-bit attention does not exist, so - // the caller's unset attention resolves to exact instead of failing. On sm_80+ an 8-bit kernel always exists, so - // the fallback only ever runs on such devices. - const bool exact_when_no_q8 = !supported(planner::AttentionMode::kQ8Fp8) && - !supported(planner::AttentionMode::kQ8Fp16); + // sm_75: no 8-bit attention kernel exists for Turing, so the caller's unset attention resolves to exact instead of + // failing. Keyed on the device, not on the kernels' availability: an sm_80+ build without them (MOSSINFER_ONLY) + // still refuses, as before. + const bool exact_when_no_q8 = kernels::sm75::turing(); r.tier = planner::resolve_tier(m.options.tier, m.options.attention, src.precision, supported, tes.precision != qwen35_text::Precision::kBF16, src.precision == planner::Precision::kNvfp4 && kernels::nvfp4::available(), @@ -1507,12 +1506,21 @@ Result Pipeline::run(const Request& req, const GenerateControl& control, bool pr check_user_paths(m.options, req); validate(req); // Two features still need sm_80+ kernels and are refused here, before anything is planned, read or uploaded: the VAE - // encoder's convolutions (img2img and inpainting) 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. + // encoder's convolutions (img2img, 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. if (kernels::sm75::turing()) { if (req.init_image != nullptr) { throw NotSupportedError("img2img and inpainting encode the input with the VAE encoder, which has no sm_75 kernel " - "in this build; text-to-image works, and so does the hires fix (it only decodes)"); + "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 " + "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 " + "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"); diff --git a/tests/unit/test_sm75.cpp b/tests/unit/test_sm75.cpp index fcf5ade..b7a0846 100644 --- a/tests/unit/test_sm75.cpp +++ b/tests/unit/test_sm75.cpp @@ -1,5 +1,5 @@ -// src/kernels/sm75: the Turing kernel family's build presence and its dtype conversion (the only piece that runs on -// any device, so this file also works as a smoke test that the sm_75 object library is linked in). +// src/kernels/sm75: the Turing kernel family's build presence (any device, a smoke test that the sm_75 object library +// is linked in) and its dtype conversion (built 75-real only, so it runs on a Turing device alone). #include #include @@ -20,7 +20,7 @@ using mi::DType; using mi::Device; using mi::Tensor; -float bf16_to_float(std::uint16_t bits) { +[[maybe_unused]] float bf16_to_float(std::uint16_t bits) { const std::uint32_t u = static_cast(bits) << 16; float f = 0; std::memcpy(&f, &u, sizeof f); @@ -58,6 +58,13 @@ struct DeviceBf16 { } }; +[[maybe_unused]] std::uint16_t float_to_f16_bits(float f) { + const __half h = __float2half_rn(f); + std::uint16_t bits = 0; + std::memcpy(&bits, &h, sizeof bits); + return bits; +} + } // namespace TEST(Sm75Build, KernelFamilyPresence) { @@ -76,6 +83,7 @@ TEST(GpuSm75Convert, RoundTripIsExactInFp16Range) { #ifndef MOSSINFER_HAVE_SM75 GTEST_SKIP() << "built without MOSSINFER_SM75"; #else + if (!mi::kernels::sm75::turing()) GTEST_SKIP() << "the sm_75 kernels are built 75-real: needs a Turing device"; // 65504 is the largest fp16 value but not representable in bf16 (it rounds to 65536, which overflows fp16 to inf): // values above fp16's range are the documented sm_75 numerics difference, not tested here. const std::vector values = {0.f, 1.f, -1.f, 2.5f, -3.14159f, 123.456f, -1024.f, 32768.f, 1e-3f, -7e-2f}; @@ -98,6 +106,7 @@ TEST(GpuSm75Convert, StridedInPlaceConversionTouchesOnlyTheView) { #ifndef MOSSINFER_HAVE_SM75 GTEST_SKIP() << "built without MOSSINFER_SM75"; #else + if (!mi::kernels::sm75::turing()) GTEST_SKIP() << "the sm_75 kernels are built 75-real: needs a Turing device"; constexpr int kRows = 3, kHeads = 2, kDim = 8, kHeadStride = 16; // heads interleaved: [token][head][16] std::vector storage(static_cast(kRows) * kHeads * kHeadStride, -12345.f); DeviceBf16 buf(storage); @@ -107,6 +116,19 @@ TEST(GpuSm75Convert, StridedInPlaceConversionTouchesOnlyTheView) { view.stride[2] = 1; const std::vector before = buf.host(); mi::kernels::sm75::convert_tensor_bf16_to_f16_inplace(view, kRows, kHeads, kDim, nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + // Converted: the view holds fp16 bits (a different pattern from bf16 for -12345), the padding is untouched. + const std::vector mid = buf.host(); + const std::uint16_t f16 = float_to_f16_bits(bf16_to_float(before[0])); + ASSERT_NE(f16, before[0]); + for (int r = 0; r < kRows; ++r) { + for (int h = 0; h < kHeads; ++h) { + for (int d = 0; d < kHeadStride; ++d) { + const std::size_t i = static_cast(r) * kHeads * kHeadStride + h * kHeadStride + d; + EXPECT_EQ(mid[i], d < kDim ? f16 : before[i]) << "converted r " << r << " h " << h << " d " << d; + } + } + } mi::kernels::sm75::convert_tensor_f16_to_bf16_inplace(view, kRows, kHeads, kDim, nullptr); MI_CUDA_CHECK(cudaDeviceSynchronize()); const std::vector after = buf.host(); From e24030859327f66412e81dbb7e4931f75ee11fab Mon Sep 17 00:00:00 2001 From: xycld Date: Sat, 31 Oct 2026 06:27:49 -0400 Subject: [PATCH 6/6] sm_75: time-embedding modulation on fp16 weights, f32 linear row offsets, refusals for sessions too - linear_fp32_row reads fp16 weights as well: block_modulation_/head_modulation_ are converted to fp16 at load on sm_75, so the time-embedding path threw on every step. bf16 -> fp16 is exact for in-range weights, so the result is the sm_80 one. - linear_f16: an f32 output's rows past the first chunk were written at r0 * ldo * 2 bytes; f32 rows are 4 bytes per element. - The Turing refusals are one function (detail::refuse_unsupported_on_sm75), called by Pipeline::run and by a Session for its options and every run's spec: a session's LoRAs or control adapter reached the bf16 merge and LoRA GEMMs. - flash_attention.cu: the comment names the kernel that runs on sm_75 (not the vendored FA2 traits). --- src/kernels/attention/flash_attention.cu | 4 +- src/kernels/basic/gemm.cpp | 5 ++- src/models/moss_dit/dit_kernels.cu | 30 ++++++++++---- src/models/moss_dit/dit_kernels.h | 4 +- src/pipeline/pipeline.cpp | 52 +++++++++++++----------- src/pipeline/pipeline_impl.h | 3 ++ src/pipeline/session.cpp | 5 ++- 7 files changed, 63 insertions(+), 40 deletions(-) diff --git a/src/kernels/attention/flash_attention.cu b/src/kernels/attention/flash_attention.cu index 832afd4..5ef530e 100644 --- a/src/kernels/attention/flash_attention.cu +++ b/src/kernels/attention/flash_attention.cu @@ -154,8 +154,8 @@ void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& MI_CUDA_CHECK(cudaGetDevice(&device)); #if MOSSINFER_HAVE_SM75 if (kernels::sm75::turing()) { - // Turing: the vendored kernel's own sm_75 traits (fp16 MMA, no cp.async) run in the sm_75 kernel family; - // q/k/v/out are converted around the launch there. + // Turing: the sm_75 family's own exact attention (src/kernels/sm75/attention_sm75.cu, fp16 MMA; the vendored + // kernel's sm_75 traits are not used); q/k/v/out are converted around the launch there. kernels::sm75::flash_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen, causal, scale, workspace, workspace_bytes, stream); return; diff --git a/src/kernels/basic/gemm.cpp b/src/kernels/basic/gemm.cpp index 29d117e..5abff99 100644 --- a/src/kernels/basic/gemm.cpp +++ b/src/kernels/basic/gemm.cpp @@ -201,8 +201,9 @@ struct Gemm::Impl { const float alpha = 1.0f, beta = 0.0f; MI_LT_CHECK(cublasLtMatmul(handle, plan.desc, &alpha, w.data, plan.a, x16, plan.b, &beta, dptr, plan.d, dptr, plan.d, &plan.algo, workspace, plan.workspace, stream)); - kernels::sm75::convert_gemm_out_f16(static_cast(out.data) + r0 * ldo * 2, dptr, bias_ptr, bias_f32, - out_f32, rows, n, ldo, acc_ld, stream); + // The output row offset in bytes: f32 rows are 4 bytes per element (dptr above is acc16 for them). + kernels::sm75::convert_gemm_out_f16(static_cast(out.data) + r0 * ldo * (out_f32 ? 4 : 2), dptr, bias_ptr, + bias_f32, out_f32, rows, n, ldo, acc_ld, stream); } } }; diff --git a/src/models/moss_dit/dit_kernels.cu b/src/models/moss_dit/dit_kernels.cu index 7107fce..e9dc088 100644 --- a/src/models/moss_dit/dit_kernels.cu +++ b/src/models/moss_dit/dit_kernels.cu @@ -1,6 +1,7 @@ #include "models/moss_dit/dit_kernels.h" #include +#include #include #include @@ -69,9 +70,13 @@ __global__ void unpack_int4_kernel(const std::uint8_t* in, std::int64_t ld_in, s out[r * ld_out + 2 * c + 1] = hi; } -// One warp per output row; 8 bf16 weights per 16-byte load. -__global__ void linear_fp32_row_kernel(const __nv_bfloat16* w, const float* x, const float* bias, __nv_bfloat16* out, - int n, int k) { +__device__ __forceinline__ float widen(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float widen(__half v) { return __half2float(v); } + +// One warp per output row; 8 weights per 16-byte load. W is bf16, or fp16 on sm_75 (weights converted at load, exact +// for in-range values). +template +__global__ void linear_fp32_row_kernel(const W* w, const float* x, const float* bias, __nv_bfloat16* out, int n, int k) { const int row = static_cast(blockIdx.x) * (blockDim.x / 32) + static_cast(threadIdx.x) / 32; const int lane = static_cast(threadIdx.x) % 32; if (row >= n) return; @@ -79,8 +84,8 @@ __global__ void linear_fp32_row_kernel(const __nv_bfloat16* w, const float* x, c float acc = 0.0f; for (int i = lane; i < k / 8; i += 32) { const uint4 v = wr[i]; - const __nv_bfloat16* h = reinterpret_cast(&v); - for (int j = 0; j < 8; ++j) acc = fmaf(__bfloat162float(h[j]), x[8 * i + j], acc); + const W* h = reinterpret_cast(&v); + for (int j = 0; j < 8; ++j) acc = fmaf(widen(h[j]), x[8 * i + j], acc); } for (int o = 16; o > 0; o /= 2) acc += __shfl_xor_sync(0xffffffffu, acc, o); if (lane == 0) out[row] = __float2bfloat16_rn(acc + bias[row]); @@ -194,15 +199,22 @@ void unpack_int4(const void* in, std::int64_t ld_in, void* out, std::int64_t ld_ } void linear_fp32_row(const Tensor& w, const float* x, const float* bias, const Tensor& out, cudaStream_t stream) { - check(w.device == Device::kCUDA && w.dtype == DType::kBF16 && w.rank == 2 && w.is_contiguous() && w.dim(1) % 8 == 0 && + check(w.device == Device::kCUDA && (w.dtype == DType::kBF16 || w.dtype == DType::kF16) && w.rank == 2 && + w.is_contiguous() && w.dim(1) % 8 == 0 && reinterpret_cast(w.data) % 16 == 0 && x != nullptr && bias != nullptr && out.device == Device::kCUDA && out.dtype == DType::kBF16 && out.numel() == w.dim(0) && out.is_contiguous(), - "linear_fp32_row", "needs a contiguous 16-byte aligned CUDA bf16 w [n, k % 8 == 0] and a bf16 out [n]"); + "linear_fp32_row", "needs a contiguous 16-byte aligned CUDA bf16/fp16 w [n, k % 8 == 0] and a bf16 out [n]"); const auto n = static_cast(w.dim(0)); if (n == 0) return; constexpr int kWarps = 8; - linear_fp32_row_kernel<<((n + kWarps - 1) / kWarps), 32 * kWarps, 0, stream>>>( - w.ptr(), x, bias, out.ptr<__nv_bfloat16>(), n, static_cast(w.dim(1))); + const auto blocks = static_cast((n + kWarps - 1) / kWarps); + if (w.dtype == DType::kF16) { + linear_fp32_row_kernel<<>>(w.ptr(), x, bias, out.ptr<__nv_bfloat16>(), + n, static_cast(w.dim(1))); + } else { + linear_fp32_row_kernel<<>>(w.ptr(), x, bias, + out.ptr<__nv_bfloat16>(), n, static_cast(w.dim(1))); + } MI_CUDA_CHECK_LAUNCH(); } diff --git a/src/models/moss_dit/dit_kernels.h b/src/models/moss_dit/dit_kernels.h index 12ab192..fa96857 100644 --- a/src/models/moss_dit/dit_kernels.h +++ b/src/models/moss_dit/dit_kernels.h @@ -32,8 +32,8 @@ void extrapolate_latents(const Tensor& x, const Tensor& x_prev, float c, const T void unpack_int4(const void* in, std::int64_t ld_in, void* out, std::int64_t ld_out, std::int64_t rows, std::int64_t cols, cudaStream_t stream); -// out[r] = bf16(sum_c w[r, c] * x[c] + bias[r]) in fp32 (w bf16 [n, k] contiguous, widened; x, bias f32 device -// arrays; out bf16 [n]): an fp32 nn.Linear on one row, rounded once. k % 8 == 0, w 16-byte aligned. +// out[r] = bf16(sum_c w[r, c] * x[c] + bias[r]) in fp32 (w bf16, or fp16 on sm_75, [n, k] contiguous, widened; x, bias +// f32 device arrays; out bf16 [n]): an fp32 nn.Linear on one row, rounded once. k % 8 == 0, w 16-byte aligned. void linear_fp32_row(const Tensor& w, const float* x, const float* bias, const Tensor& out, cudaStream_t stream); // Copies `bytes` (<= kMaxParamBytes) host bytes to dst through kernel arguments: no pinned staging, diff --git a/src/pipeline/pipeline.cpp b/src/pipeline/pipeline.cpp index 266130d..814f901 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -1500,35 +1500,39 @@ void Pipeline::Impl::explain_control(const Request::Control& control, Result& r) "): " + (inputs.empty() ? std::string("no images (zero control latents)") : inputs) + "."; } +// Two features still need sm_80+ kernels and are refused on Turing before anything is planned, read or uploaded (by +// Pipeline::run, and by a Session for its options and every run's spec): the VAE encoder's convolutions (img2img, +// 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) { + 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 " + "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 " + "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 " + "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"); + } + if (req.control.enabled()) { + throw NotSupportedError("control adapters are LoRAs and need sm_80+ GEMM kernels; this sm_75 build cannot run them"); + } +} + // generate(), or with `preload` the same up to the plan and the arena, then the loads of the models the plan keeps. Result Pipeline::run(const Request& req, const GenerateControl& control, bool preload) { Impl& m = *impl_; check_user_paths(m.options, req); validate(req); - // Two features still need sm_80+ kernels and are refused here, before anything is planned, read or uploaded: the VAE - // encoder's convolutions (img2img, 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. - if (kernels::sm75::turing()) { - if (req.init_image != nullptr) { - throw NotSupportedError("img2img and inpainting encode the input with the VAE encoder, which has no sm_75 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 " - "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 " - "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"); - } - if (req.control.enabled()) { - throw NotSupportedError("control adapters are LoRAs and need sm_80+ GEMM kernels; this sm_75 build cannot run them"); - } - } + refuse_unsupported_on_sm75(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 8e7a418..64a1266 100644 --- a/src/pipeline/pipeline_impl.h +++ b/src/pipeline/pipeline_impl.h @@ -218,6 +218,9 @@ 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); + inline moss_dit::SkipSet skip_set(const std::vector& layers) { moss_dit::SkipSet s; for (const int l : layers) s.set(static_cast(l)); diff --git a/src/pipeline/session.cpp b/src/pipeline/session.cpp index a8f3e32..05f7eef 100644 --- a/src/pipeline/session.cpp +++ b/src/pipeline/session.cpp @@ -381,6 +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); 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)."); } @@ -463,7 +464,9 @@ void Session::validate(const RunSpec& spec, const SessionOptions& options) { } // The rules mi_generate's request shares (sizes, steps, schedule fields, CFG, SLG, control images, accelerators). const Guidance g = run_guidance(spec); - pipeline::validate(run_request(options, spec, g)); + const Request run_q = run_request(options, spec, g); + pipeline::validate(run_q); + refuse_unsupported_on_sm75(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."); }