From 70c12826460af3dd3fa5560922c3357eea68e7ad Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Mon, 7 Sep 2026 16:32:07 +0800 Subject: [PATCH 1/2] [Kernel][SM70] Integrate FP32 v37 prefill and exact E4M3 bridge Retain the qualified prefix/tail arithmetic under a unique FA2 symbol, admit tile-aligned chunk shapes, and preserve exact fallbacks. Model promotion checks continue in the owned PR. Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- cmake/external_projects/vllm_flash_attn.cmake | 19 + csrc/attention/sm70_v37/README.md | 74 ++ csrc/attention/sm70_v37/bridge.cu | 335 +++++++ .../sm70_v37/epilogue_visitor_with_softmax.h | 537 +++++++++++ csrc/attention/sm70_v37/gemm_with_softmax.h | 531 +++++++++++ csrc/attention/sm70_v37/prefill.cu | 666 ++++++++++++++ .../attention/sm70_v37/reduce_softmax_final.h | 256 ++++++ csrc/attention/sm70_v37/register.cpp | 20 + csrc/attention/sm70_v37/tail.cu | 854 ++++++++++++++++++ docs/design/sm70_flash_v37_prefill.md | 120 +++ .../attention/test_sm70_v37_prefill.py | 155 ++++ .../attention/test_sm70_flash_v100_policy.py | 18 +- tests/v1/attention/test_sm70_v37_prefill.py | 115 +++ vllm/envs.py | 4 + vllm/v1/attention/backends/flash_attn_v100.py | 110 ++- 15 files changed, 3796 insertions(+), 18 deletions(-) create mode 100644 csrc/attention/sm70_v37/README.md create mode 100644 csrc/attention/sm70_v37/bridge.cu create mode 100644 csrc/attention/sm70_v37/epilogue_visitor_with_softmax.h create mode 100644 csrc/attention/sm70_v37/gemm_with_softmax.h create mode 100644 csrc/attention/sm70_v37/prefill.cu create mode 100644 csrc/attention/sm70_v37/reduce_softmax_final.h create mode 100644 csrc/attention/sm70_v37/register.cpp create mode 100644 csrc/attention/sm70_v37/tail.cu create mode 100644 docs/design/sm70_flash_v37_prefill.md create mode 100644 tests/kernels/attention/test_sm70_v37_prefill.py create mode 100644 tests/v1/attention/test_sm70_v37_prefill.py diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index b7658da5d2..8e18f01626 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -94,6 +94,25 @@ install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_ FetchContent_MakeAvailable(vllm-flash-attn) message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}") +# Keep the precision-qualified SM70 prefill route in the parent repository. +# Its private CUTLASS visitors have distinct types and do not modify the +# legacy FA2 headers or operators, which remain available for rollback. +if(VLLM_FLASH_ATTN_SM70 AND TARGET _vllm_fa2_C) + set(SM70_V37_DIR "${CMAKE_CURRENT_LIST_DIR}/../../csrc/attention/sm70_v37") + set(SM70_V37_CUDA_SRCS + "${SM70_V37_DIR}/prefill.cu" + "${SM70_V37_DIR}/tail.cu" + "${SM70_V37_DIR}/bridge.cu") + # FA2 is created in a child directory. Source properties must be visible in + # that target's scope; setting them only in the parent silently loses SM70. + set_source_files_properties(${SM70_V37_CUDA_SRCS} + TARGET_DIRECTORY _vllm_fa2_C + PROPERTIES COMPILE_OPTIONS "-gencode=arch=compute_70,code=sm_70") + target_sources(_vllm_fa2_C PRIVATE + ${SM70_V37_CUDA_SRCS} + "${SM70_V37_DIR}/register.cpp") +endif() + # Restore the install prefix after FA's install rules install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS) install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) diff --git a/csrc/attention/sm70_v37/README.md b/csrc/attention/sm70_v37/README.md new file mode 100644 index 0000000000..e2cde9c752 --- /dev/null +++ b/csrc/attention/sm70_v37/README.md @@ -0,0 +1,74 @@ +# SM70 v37 long-prefill attention + +This is a prefill-only, FP32-accumulating implementation for one request with +FP16 activations, six query heads, one KV head and head dimension 256. It does +not change decode, speculative decoding, sampling, or the global KV format. + +## Dataflow and numerical contract + +We pack the six query heads into the GEMM row dimension. For a causal chunk, +the preceding keys form a non-causal prefix and the current chunk forms an +exactly masked tail. The prefix uses M128/N256/K32 SM70 Tensor Core GEMMs. +Its QK epilogue forms tile-local probabilities from FP32 logits and FP32 +max/sum statistics, before rounding probabilities to FP16 tensor operands. +PV rescales those probabilities and accumulates the numerator and online +partial state in FP32. The tail also retains its unnormalized output in FP32. +We combine both states before the single final FP16 output conversion. + +“FP32 accumulation” is not a claim of FP32 inputs or exact real arithmetic: +probability operands and the final output are still FP16. The explicit E4M3 +bridge expands stored bytes with their per-layer scales; it does not undo +the original KV quantization loss. + +## Admission and fallback + +- SM70 only; contiguous, 16-byte-aligned FP16 Q/K/V/output on one device. +- Batch 1, Hq 6, Hkv 1, D256, causal, scale 1/16. +- Q is 64–8192 in multiples of 64; Q < KV <= 262144; KV is a multiple of 32. +- The bridge can prepend zero query rows to reach a 64-row boundary. It + keeps the original KV length and slices away the leading outputs, so the + causal offset of every real query is unchanged. +- A partial K32 tile is deliberately not admitted. A 16-aligned but + non-32-aligned KV probe failed the FP64 gate during integration. +- CUDA Graph capture, unsupported shapes and insufficient workspace retain + the existing exact fallback. Decode never enters this prefill operator. + +The score cache reserves 768 MiB per used device (8192 queries times six +heads times 8192 prefix columns times two bytes). FP32 partials, statistics +and tail output use a transient slab. The first allocation requires this +workspace plus 128 MiB of downstream headroom. A per-device lock and CUDA +completion event serialize shared score/metadata use across caller streams. +An input-ready event starts the tail on its private stream; a completion +event joins it before the final merge. + +## Build and identify the route + +The normal SM70 FA2 CMake target includes these translation units. The +dedicated operator name, `sm70_d256_gqa_v37_fwd`, prevents an old FA2 library +from being mistaken for this implementation. If the operator is absent, the +backend warns and falls back instead of selecting the legacy architecture. + +`VLLM_FLASH_V100_PREFILL_D256_GQA_V37=0` restores the old architecture loader +and disables the new E4M3 bridge. The existing long-prefill architecture +switch remains an additional gate. Explicit `fp8_e4m3` and `fp8_e5m2` retain +their respective byte encodings; this change does not reinterpret `fp8`. + +The `prefill_dense_d256_gqa_v37` and `prefill_prefix_fp8_e4m3_bridge` route +counters identify actual execution. A requested environment variable alone +is not route-hit evidence. + +## Focused validation + +Run from an SM70-built environment: + +```bash +.venv/bin/python -m pytest tests/v1/attention/test_sm70_v37_prefill.py tests/v1/attention/test_sm70_flash_v100_policy.py -q +.venv/bin/python -m pytest tests/kernels/attention/test_sm70_v37_prefill.py -q +``` + +The GPU tests compare causal attention against a PyTorch FP64 oracle, check +leading-query padding and unsupported alignment, and cover every E4M3 byte +at unit/non-unit scales, including graph replay with changing live lengths. +Long-context model acceptance additionally requires a matched no-MTP TP4 +run with real route counters, finite logits, output checks, and separate +prefill/TTFT/decode timing. Operator tests alone do not establish that gate. diff --git a/csrc/attention/sm70_v37/bridge.cu b/csrc/attention/sm70_v37/bridge.cu new file mode 100644 index 0000000000..06bcc385e9 --- /dev/null +++ b/csrc/attention/sm70_v37/bridge.cu @@ -0,0 +1,335 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +// Exact E4M3 storage expansion with the validated eight-byte feed. +namespace { +template +__device__ __forceinline__ __half decode_byte(uint8_t raw) { + if constexpr (FORMAT == 2) { + return __ushort_as_half(static_cast(raw) << 8); + } else { + const uint16_t sign = static_cast(raw & 0x80u) << 8; + const uint16_t magnitude = raw & 0x7fu; + uint16_t bits; + if (magnitude == 0x7fu) { + bits = 0x7e00u; + } else if (magnitude < 8) { + bits = __half_as_ushort(__float2half_rn(float(magnitude) * 0x1p-9f)); + } else { + bits = (magnitude << 7) + 0x2000u; + } + return __ushort_as_half(sign | bits); + } +} + +template +__device__ __forceinline__ __half2 decode_pair(const uint16_t raw) { + if constexpr (FORMAT == 2) { + union { + uint32_t u; + __half2 h; + } pair; + pair.u = (uint32_t(raw & 0xffu) << 8) | (uint32_t(raw & 0xff00u) << 16); + return pair.h; + } else if constexpr (FORMAT >= 3) { + // Exact E4M3 conversion: the intermediate half2 equals value / 256. + // This is a power-of-two format conversion, never an FP16 accumulator. + if ((raw & 0x7fu) == 0x7fu || ((raw >> 8) & 0x7fu) == 0x7fu) { + return __halves2half2(decode_byte<1>(raw & 0xffu), + decode_byte<1>(raw >> 8)); + } + union { + uint32_t u; + __half2 h; + } pair; + pair.u = (uint32_t(raw & 0x0080u) << 8) | (uint32_t(raw & 0x007fu) << 7) | + (uint32_t(raw & 0x8000u) << 16) | (uint32_t(raw & 0x7f00u) << 15); + return __hmul2(pair.h, __float2half2_rn(256.f)); + } else { + return __halves2half2(decode_byte(raw & 0xffu), + decode_byte(raw >> 8)); + } +} +template +__device__ __forceinline__ __half2 load_pair(const void* cache, + int64_t offset) { + return decode_pair( + reinterpret_cast(cache)[offset >> 1]); +} +} // namespace + +namespace { + +constexpr int kThreads = 256; + +template +__global__ void fp8_paged_kv_to_fp16_kernel( + const void* __restrict__ key_cache, const void* __restrict__ value_cache, + const int* __restrict__ block_table, const int* __restrict__ seq_lens, + __half* __restrict__ key_out, __half* __restrict__ value_out, + int batch_size, int max_num_blocks, int input_block_size, + int output_blocks_per_seq, int output_block_size, int num_heads, + int head_dim, int64_t key_block_stride, int64_t key_token_stride, + int64_t key_head_stride, int64_t value_block_stride, + int64_t value_token_stride, int64_t value_head_stride, + int64_t key_out_block_stride, int64_t key_out_token_stride, + int64_t key_out_head_stride, int64_t value_out_block_stride, + int64_t value_out_token_stride, int64_t value_out_head_stride, + float key_scale, float value_scale) { + const int batch_idx = blockIdx.y; + if (batch_idx >= batch_size) { + return; + } + + constexpr int kPairs = FORMAT == 4 ? 4 : 1; + const int pairs_per_head = head_dim / (2 * kPairs); + const int64_t pairs_per_token = + static_cast(num_heads) * pairs_per_head; + const int max_tokens = output_blocks_per_seq * output_block_size; + const int64_t total_pairs = + static_cast(max_tokens) * pairs_per_token; + const int64_t pair_idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (pair_idx >= total_pairs) { + return; + } + + const int token_idx = static_cast(pair_idx / pairs_per_token); + const int seq_len = seq_lens[batch_idx]; + // Paged prefill loads complete 16-row WMMA tiles. Zero the partial tail so + // masked probabilities cannot multiply uninitialized workspace values. + const int padded_seq_len = min(max_tokens, (seq_len + 15) & ~15); + if (token_idx >= padded_seq_len) { + return; + } + const int pair_in_token = static_cast(pair_idx % pairs_per_token); + const int head_idx = pair_in_token / pairs_per_head; + const int element_idx = (pair_in_token % pairs_per_head) * 2 * kPairs; + + const __half2 zero = __float2half2_rn(0.0f); + __half2 key_pair[kPairs], value_pair[kPairs]; +#pragma unroll + for (int p = 0; p < kPairs; ++p) { + key_pair[p] = zero; + value_pair[p] = zero; + } + if (token_idx < seq_len) { + const int logical_block = token_idx / input_block_size; + const int input_block_offset = token_idx - logical_block * input_block_size; + const int physical_block = + __ldg(&block_table[batch_idx * max_num_blocks + logical_block]); + const int64_t key_input_offset = + static_cast(physical_block) * key_block_stride + + static_cast(input_block_offset) * key_token_stride + + static_cast(head_idx) * key_head_stride + element_idx; + const int64_t value_input_offset = + static_cast(physical_block) * value_block_stride + + static_cast(input_block_offset) * value_token_stride + + static_cast(head_idx) * value_head_stride + element_idx; + + if constexpr (kPairs == 4) { + const auto kr = *reinterpret_cast( + static_cast(key_cache) + key_input_offset); + const auto vr = *reinterpret_cast( + static_cast(value_cache) + value_input_offset); +#pragma unroll + for (int p = 0; p < kPairs; ++p) { + key_pair[p] = + decode_pair(static_cast(kr >> (16 * p))); + value_pair[p] = + decode_pair(static_cast(vr >> (16 * p))); + } + } else { + key_pair[0] = load_pair(key_cache, key_input_offset); + value_pair[0] = load_pair(value_cache, value_input_offset); + } + if constexpr (!UNIT_SCALE) { +#pragma unroll + for (int p = 0; p < kPairs; ++p) { + const float2 key_values = __half22float2(key_pair[p]); + const float2 value_values = __half22float2(value_pair[p]); + key_pair[p] = __float22half2_rn( + make_float2(key_values.x * key_scale, key_values.y * key_scale)); + value_pair[p] = __float22half2_rn(make_float2( + value_values.x * value_scale, value_values.y * value_scale)); + } + } + } + + const int output_block_offset = token_idx / output_block_size; + const int output_token_offset = + token_idx - output_block_offset * output_block_size; + const int output_block = + batch_idx * output_blocks_per_seq + output_block_offset; + const int64_t key_output_offset = + static_cast(output_block) * key_out_block_stride + + static_cast(output_token_offset) * key_out_token_stride + + static_cast(head_idx) * key_out_head_stride + element_idx; + const int64_t value_output_offset = + static_cast(output_block) * value_out_block_stride + + static_cast(output_token_offset) * value_out_token_stride + + static_cast(head_idx) * value_out_head_stride + element_idx; + if constexpr (kPairs == 4) { + union { + uint4 u; + __half2 h[4]; + } kr, vr; +#pragma unroll + for (int p = 0; p < kPairs; ++p) { + kr.h[p] = key_pair[p]; + vr.h[p] = value_pair[p]; + } + *reinterpret_cast(key_out + key_output_offset) = kr.u; + *reinterpret_cast(value_out + value_output_offset) = vr.u; + } else { + *reinterpret_cast<__half2*>(key_out + key_output_offset) = key_pair[0]; + *reinterpret_cast<__half2*>(value_out + value_output_offset) = + value_pair[0]; + } +} + +} // namespace + +void sm70_v37_e4m3_paged_kv_to_fp16(const at::Tensor& key_cache, + const at::Tensor& value_cache, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + at::Tensor& key_out, at::Tensor& value_out, + const double key_scale_arg, + const double value_scale_arg) { + constexpr int format = 4; + const float key_scale = static_cast(key_scale_arg); + const float value_scale = static_cast(value_scale_arg); + TORCH_CHECK(key_cache.is_cuda() && value_cache.is_cuda() && + key_out.is_cuda() && value_out.is_cuda(), + "FP8 KV bridge tensors must be CUDA tensors"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "FP8 KV bridge metadata must be CUDA tensors"); + TORCH_CHECK(key_cache.scalar_type() == at::kByte && + value_cache.scalar_type() == at::kByte, + "FP8 input caches must be stored as uint8"); + TORCH_CHECK(key_out.scalar_type() == at::kHalf && + value_out.scalar_type() == at::kHalf, + "FP8 KV bridge output caches must be fp16"); + TORCH_CHECK(block_table.scalar_type() == at::kInt && + seq_lens.scalar_type() == at::kInt, + "FP8 KV bridge block_table and seq_lens must be int32"); + TORCH_CHECK(key_cache.dim() == 4 && value_cache.dim() == 4 && + key_out.dim() == 4 && value_out.dim() == 4, + "FP8 KV bridge expects paged [blocks,tokens,heads,dim] caches"); + TORCH_CHECK(key_cache.sizes() == value_cache.sizes(), + "FP8 K/V input cache shapes must match"); + TORCH_CHECK(key_out.sizes() == value_out.sizes(), + "FP16 K/V output cache shapes must match"); + TORCH_CHECK(key_cache.size(2) == key_out.size(2) && + key_cache.size(3) == key_out.size(3), + "FP8 KV bridge head shape must not change"); + TORCH_CHECK(key_cache.stride(3) == 1 && value_cache.stride(3) == 1 && + key_out.stride(3) == 1 && value_out.stride(3) == 1, + "FP8 KV bridge requires contiguous head dimensions"); + TORCH_CHECK(key_cache.size(3) % 2 == 0, + "FP8 KV bridge requires an even head dimension"); + TORCH_CHECK(block_table.dim() == 2 && seq_lens.dim() == 1 && + block_table.size(0) == seq_lens.size(0), + "FP8 KV bridge metadata shape mismatch"); + TORCH_CHECK(key_scale > 0.f && value_scale > 0.f, + "FP8 KV bridge scales must be positive"); + + const int batch_size = block_table.size(0); + TORCH_CHECK(batch_size > 0, "FP8 KV bridge batch must be non-empty"); + TORCH_CHECK(key_out.size(0) % batch_size == 0, + "FP16 output blocks must divide evenly across the batch"); + const int output_blocks_per_seq = key_out.size(0) / batch_size; + const int input_capacity = block_table.size(1) * key_cache.size(1); + const int output_capacity = output_blocks_per_seq * key_out.size(1); + TORCH_CHECK(output_capacity >= input_capacity, + "FP16 output cache capacity must cover the input block table"); + + TORCH_CHECK(format >= 1 && format <= 4, + "format: 1 E4M3 scalar, 2 E5M2, 3 E4M3 packed, 4 E4M3 wide"); + TORCH_CHECK(std::isfinite(key_scale) && std::isfinite(value_scale), + "finite scales required"); + TORCH_CHECK(block_table.is_contiguous() && seq_lens.is_contiguous(), + "contiguous metadata required"); + for (const at::Tensor* tensor : std::initializer_list{ + &value_cache, &block_table, &seq_lens, &key_out, &value_out}) { + TORCH_CHECK(tensor->device() == key_cache.device(), + "all inputs must share a CUDA device"); + } + for (const at::Tensor* tensor : std::initializer_list{ + &key_cache, &value_cache, &key_out, &value_out}) { + TORCH_CHECK( + tensor->size(1) > 0 && tensor->size(2) > 0 && tensor->size(3) > 0, + "empty head/page unsupported"); + TORCH_CHECK(reinterpret_cast(tensor->data_ptr()) % 2 == 0, + "paired loads require aligned pointers"); + for (int axis = 0; axis < 3; ++axis) { + TORCH_CHECK(tensor->stride(axis) % 2 == 0, + "paired loads require even strides"); + } + if (format == 4) { + TORCH_CHECK(tensor->size(3) % 8 == 0 && + reinterpret_cast(tensor->data_ptr()) % 16 == 0, + "wide conversion requires D divisible by 8 and 16-byte " + "aligned pointers"); + for (int axis = 0; axis < 3; ++axis) + TORCH_CHECK(tensor->stride(axis) % 8 == 0, + "wide conversion requires 8-element aligned strides"); + } + } + c10::cuda::CUDAGuard device_guard(key_cache.device()); + const int64_t total_pairs = static_cast(output_capacity) * + key_cache.size(2) * + (key_cache.size(3) / (format == 4 ? 8 : 2)); + const dim3 grid( + static_cast((total_pairs + kThreads - 1) / kThreads), + batch_size); + const auto stream = at::cuda::getCurrentCUDAStream().stream(); + +#define LAUNCH_FP8_BRIDGE(FORMAT, UNIT_SCALE) \ + fp8_paged_kv_to_fp16_kernel \ + <<>>( \ + key_cache.data_ptr(), value_cache.data_ptr(), \ + block_table.data_ptr(), seq_lens.data_ptr(), \ + reinterpret_cast<__half*>(key_out.data_ptr()), \ + reinterpret_cast<__half*>(value_out.data_ptr()), batch_size, \ + block_table.size(1), key_cache.size(1), output_blocks_per_seq, \ + key_out.size(1), key_cache.size(2), key_cache.size(3), \ + key_cache.stride(0), key_cache.stride(1), key_cache.stride(2), \ + value_cache.stride(0), value_cache.stride(1), value_cache.stride(2), \ + key_out.stride(0), key_out.stride(1), key_out.stride(2), \ + value_out.stride(0), value_out.stride(1), value_out.stride(2), \ + key_scale, value_scale) + + if (format == 4) { + if (key_scale == 1.f && value_scale == 1.f) { + LAUNCH_FP8_BRIDGE(4, true); + } else { + LAUNCH_FP8_BRIDGE(4, false); + } + } +#undef LAUNCH_FP8_BRIDGE + + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +TORCH_LIBRARY_FRAGMENT(_vllm_fa2_C, m) { + m.def( + "sm70_v37_e4m3_bridge(Tensor key, Tensor value, Tensor table, Tensor " + "lengths, Tensor(a!) key_out, Tensor(b!) value_out, float key_scale, " + "float value_scale) -> ()"); +} +TORCH_LIBRARY_IMPL(_vllm_fa2_C, CUDA, m) { + m.impl("sm70_v37_e4m3_bridge", &sm70_v37_e4m3_paged_kv_to_fp16); +} diff --git a/csrc/attention/sm70_v37/epilogue_visitor_with_softmax.h b/csrc/attention/sm70_v37/epilogue_visitor_with_softmax.h new file mode 100644 index 0000000000..d911d55a3e --- /dev/null +++ b/csrc/attention/sm70_v37/epilogue_visitor_with_softmax.h @@ -0,0 +1,537 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + * this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue visitor for threadblock scoped GEMMs that process softmax + computations in epilogue. + + The epilogue finds max values in each row of the row-major output matrix and + stores them. The max values are also used for a further round of threadblock + scoped reduction operation, where the partial reduction results are stored in + a pre-allocated array and used for further full reduction. + +*/ + +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "cutlass/cutlass.h" +#include "cutlass/arch/memory.h" +#include "cutlass/arch/memory_sm75.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/fast_math.h" + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +template +class EpilogueVisitorSoftmaxV37 { + public: + using ThreadblockShape = ThreadblockShape_; + static int const kThreadCount = ThreadCount; + + using OutputTileIterator = OutputTileIterator_; + using ElementwiseFunctor = ElementwiseFunctor_; + + static int const kIterations = OutputTileIterator::kIterations; + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + using ElementOutput = typename OutputTileIterator::Element; + using LayoutOutput = cutlass::layout::RowMajor; + using ElementAccumulator = ElementAccumulator_; + + using ElementNorm = ElementNorm_; + using ElementSum = ElementSum_; + using ElementSoftmaxCompute = ElementSoftmaxCompute_; + + using AccumulatorFragment = Array; + using SoftmaxFragment = Array; + using OutputVector = Array; + using TensorRefD = TensorRef; + + static int const kThreadsPerRow = + OutputTileIterator::ThreadMap::Detail::kAccessWidth; + static bool const kHasMultiStepsInRow = + (OutputTileIterator::ThreadMap::Iterations::kColumn > 1); + static bool const kUseMasking = UseMasking_; + + /// Argument structure + struct Arguments { + typename ElementwiseFunctor::Params elementwise; + int64_t batch_stride_C; + int64_t batch_stride_D; + int64_t batch_stride_Max; + int64_t batch_stride_Sum; + + // + // Methods + // + Arguments() + : batch_stride_C(0), + batch_stride_D(0), + batch_stride_Max(0), + batch_stride_Sum(0) {} + + Arguments(typename ElementwiseFunctor::Params elementwise_) + : elementwise(elementwise_), + batch_stride_C(0), + batch_stride_D(0), + batch_stride_Max(0), + batch_stride_Sum(0) {} + + Arguments(typename ElementwiseFunctor::Params elementwise_, + int64_t batch_stride_C_, int64_t batch_stride_D_, + int64_t batch_stride_Max_, int64_t batch_stride_Sum_) + : elementwise(elementwise_), + batch_stride_C(batch_stride_C_), + batch_stride_D(batch_stride_D_), + batch_stride_Max(batch_stride_Max_), + batch_stride_Sum(batch_stride_Sum_) {} + }; + + struct Params { + typename ElementwiseFunctor::Params elementwise; + int64_t batch_stride_C; + int64_t batch_stride_D; + int64_t batch_stride_Max; + int64_t batch_stride_Sum; + // + // Methods + // + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Arguments const& args) + : elementwise(args.elementwise), + batch_stride_C(args.batch_stride_C), + batch_stride_D(args.batch_stride_D), + batch_stride_Max(args.batch_stride_Max), + batch_stride_Sum(args.batch_stride_Sum) {} + }; + + /// Shared storage + struct SharedStorage {}; + + private: + Params const& params_; + SharedStorage& shared_storage_; + MatrixCoord extent_; + MatrixCoord extent_real_; + ElementwiseFunctor elementwise_; + + OutputTileIterator iterator_C_; + OutputTileIterator iterator_D_; + typename OutputTileIterator::Fragment fragment_C_; + typename OutputTileIterator::Fragment fragment_D_; + + ElementAccumulator alpha_; + ElementAccumulator beta_; + + ElementNorm* ptr_Max_; + ElementSum* ptr_Sum_; + + int column_offset_; + + ElementSoftmaxCompute accum_max_; + ElementSoftmaxCompute accum_sum_; + // Keep unrounded logits until the complete tile-row maximum is known. + SoftmaxFragment + row_logits_[OutputTileIterator::ThreadMap::Iterations::kColumn]; + int first_row_fragment_; + + MatrixCoord thread_offset_; + + float infinity_; + + public: + CUTLASS_DEVICE + EpilogueVisitorSoftmaxV37( + Params const& params, SharedStorage& shared_storage, + cutlass::MatrixCoord const& problem_size, int thread_idx, int warp_idx, + int lane_idx, typename OutputTileIterator::Params params_C, + typename OutputTileIterator::Params params_D, + typename OutputTileIterator::Element* ptr_C, + typename OutputTileIterator::Element* ptr_D, + ElementNorm* ptr_Max = nullptr, ElementSum* ptr_Sum = nullptr, + cutlass::MatrixCoord const& threadblock_offset = cutlass::MatrixCoord(0, + 0), + int column_offset = 0, + cutlass::MatrixCoord const& problem_size_real = cutlass::MatrixCoord(0, + 0), + float infinity = 10000.0f) + : params_(params), + shared_storage_(shared_storage), + extent_(problem_size), + elementwise_(params.elementwise), + iterator_C_(params_C, ptr_C, problem_size, thread_idx, + threadblock_offset), + iterator_D_(params_D, ptr_D, problem_size, thread_idx, + threadblock_offset), + ptr_Max_(ptr_Max), + ptr_Sum_(ptr_Sum), + column_offset_(column_offset), + extent_real_(problem_size_real), + infinity_(infinity) { + alpha_ = (params.elementwise.alpha_ptr ? *params.elementwise.alpha_ptr + : params.elementwise.alpha); + beta_ = (params.elementwise.beta_ptr ? *params.elementwise.beta_ptr + : params.elementwise.beta); + + if (beta_ == ElementAccumulator()) { + iterator_C_.clear_mask(); + } + } + + /// Helper to indicate split-K behavior + CUTLASS_DEVICE + void set_k_partition( + int split_k_index, ///< Index of this threadblock within split-K + ///< partitioned scheme + int split_k_slices) { ///< Total number of split-K slices + } + + /// Called to set the batch index + CUTLASS_DEVICE + void set_batch_index(int batch_idx) { + iterator_C_.add_pointer_offset(batch_idx * params_.batch_stride_C); + iterator_D_.add_pointer_offset(batch_idx * params_.batch_stride_D); + } + + /// Called at the start of the epilogue just before iterating over accumulator + /// slices + CUTLASS_DEVICE + void begin_epilogue() {} + + /// Called at the start of one step before starting accumulator exchange + CUTLASS_DEVICE + void begin_step(int step_idx) { + fragment_D_.clear(); + fragment_C_.clear(); + + if (elementwise_.kScale != + cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling) { + iterator_C_.load(fragment_C_); + ++iterator_C_; + } + } + + /// Called at the start of a row + CUTLASS_DEVICE + void begin_row(int row_idx) { + // Clear accumulators for max and sum when starting a whole row + clear_accum_(); + } + + /// Called after accumulators have been exchanged for each accumulator vector + CUTLASS_DEVICE + void visit(int iter_idx, int row_idx, int column_idx, int frag_idx, + AccumulatorFragment const& accum) { + using Mul = cutlass::multiplies; + using Minus = cutlass::minus; + using Exp = cutlass::fast_exp_op; + + Minus minus; + Exp exponential; + + SoftmaxFragment result; + + NumericArrayConverter + source_converter; + OutputVector& source_vector = + reinterpret_cast(&fragment_C_)[frag_idx]; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElementsPerAccess; ++i) { + result[i] = + ElementSoftmaxCompute(accum[i]) * ElementSoftmaxCompute(alpha_); + if (beta_ != ElementAccumulator(0)) + result[i] += ElementSoftmaxCompute(source_vector[i]) * + ElementSoftmaxCompute(beta_); + } + if (!column_idx) first_row_fragment_ = frag_idx; + if (kHasMultiStepsInRow) row_logits_[column_idx] = result; + + thread_offset_ = iterator_D_.thread_start() + + OutputTileIterator::ThreadMap::iteration_offset(frag_idx); + + bool column_guard = (thread_offset_.column() < extent_.column()); + + if (kUseMasking) { + int elements_in_boundary = + extent_real_.column() - thread_offset_.column(); + elements_in_boundary = (elements_in_boundary > kElementsPerAccess) + ? kElementsPerAccess + : elements_in_boundary; + elementwise_padding_(result, elements_in_boundary); + } + + ElementSoftmaxCompute accum_max_prev = accum_max_; + + // Compute the maximum within one row + if (!column_idx) { + // This is the first fragment in a new row + if (column_guard) { + accum_max_ = maximum_accumulator_(result); + } + } else { + // This is an additional fragment in the same row + if (column_guard) { + accum_max_ = maximum_accumulator_(result, accum_max_); + } + } + + // proactively compute max in warps + accum_max_ = warp_reduce_max_(accum_max_); + + ElementSoftmaxCompute updater = fast_exp(accum_max_prev - accum_max_); + + SoftmaxFragment intermediate = exponential(minus(result, accum_max_)); + + if (kHasMultiStepsInRow) { + if (!column_idx) { + accum_sum_ = (column_guard) ? sum_accumulator_(intermediate) + : ElementSoftmaxCompute(0); + } else { + // Algorithm in $3.1, https://arxiv.org/pdf/2205.14135v1.pdf + // S* = S* x updater + sum_row(P'), where updater = exp(M* - M_row) + accum_sum_ = (column_guard) + ? sum_accumulator_(intermediate, accum_sum_ * updater) + : accum_sum_ * updater; + } + } else { + accum_sum_ = (column_guard) ? sum_accumulator_(intermediate, accum_sum_) + : ElementSoftmaxCompute(0); + } + + // Convert to the output + NumericArrayConverter + output_converter; + OutputVector& output = + reinterpret_cast(&fragment_D_)[frag_idx]; + if (!kHasMultiStepsInRow) { + output = output_converter(intermediate); + } + } + + /// Called at the end of a row + CUTLASS_DEVICE + void end_row(int row_idx) { + using ConvertSumOutput = + cutlass::NumericConverter; + using ConvertNormOutput = + cutlass::NumericConverter; + + ConvertSumOutput convert_sum_output; + ConvertNormOutput convert_norm_output; + + // Compute accumulate sum only in the last step + + bool is_first_thread_in_tile = ((threadIdx.x % kThreadsPerRow) == 0); + bool row_guard = thread_offset_.row() < extent_.row(); + bool is_write_thread = row_guard && is_first_thread_in_tile; + if (kHasMultiStepsInRow) { + NumericArrayConverter + convert; + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < OutputTileIterator::ThreadMap::Iterations::kColumn; + ++c) { + SoftmaxFragment values; + CUTLASS_PRAGMA_UNROLL + for (int e = 0; e < kElementsPerAccess; ++e) { + ElementSoftmaxCompute centered = row_logits_[c][e] - accum_max_; + values[e] = fast_exp(centered); + } + reinterpret_cast(&fragment_D_)[first_row_fragment_ + c] = + convert(values); + } + } + // Match the mass of the probabilities actually consumed by PV. This + // counterfactual is admitted only if it improves the PyTorch oracle error. + ElementSoftmaxCompute rounded_sum = ElementSoftmaxCompute(0); + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < OutputTileIterator::ThreadMap::Iterations::kColumn; + ++c) { + int fragment = first_row_fragment_ + c; + int column = + iterator_D_.thread_start().column() + + OutputTileIterator::ThreadMap::iteration_offset(fragment).column(); + OutputVector const& probabilities = + reinterpret_cast(&fragment_D_)[fragment]; + CUTLASS_PRAGMA_UNROLL + for (int e = 0; e < kElementsPerAccess; ++e) + if (column + e < extent_.column()) + rounded_sum += ElementSoftmaxCompute(probabilities[e]); + } + accum_sum_ = warp_reduce_sum_(rounded_sum); + // The final reduction overwrites tile zero with the block maximum. + // Preserve that one tile's FP32 maximum in an extra row slab. + static_assert( + ThreadblockShape::kN <= 8192 && 8192 % ThreadblockShape::kN == 0, + "private tiled score workspace contract"); + arch::global_store( + convert_norm_output(accum_max_), + (void*)(ptr_Max_ + (8192 / ThreadblockShape::kN) * extent_.row() + + thread_offset_.row()), + is_write_thread && column_offset_ == 0); + + int block_batch = blockIdx.z; + + ElementNorm* curr_ptr_max = ptr_Max_ + thread_offset_.row() + + column_offset_ + + block_batch * params_.batch_stride_Max; + ElementSum* curr_ptr_sum = ptr_Sum_ + thread_offset_.row() + + column_offset_ + + block_batch * params_.batch_stride_Sum; + + arch::global_store( + convert_norm_output(accum_max_), (void*)curr_ptr_max, is_write_thread); + + arch::global_store( + convert_sum_output(accum_sum_), (void*)curr_ptr_sum, is_write_thread); + + // Clear accumulators for max and sum when finishing a whole row + clear_accum_(); + } + + /// Called after all accumulator elements have been visited + CUTLASS_DEVICE + void end_step(int step_idx) { + iterator_D_.store(fragment_D_); + ++iterator_D_; + } + + /// Called after all steps have been completed + CUTLASS_DEVICE + void end_epilogue() {} + + private: + CUTLASS_DEVICE + void elementwise_padding_(SoftmaxFragment& result, int elements_in_boundary) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < SoftmaxFragment::kElements; ++i) { + result[i] = (i < elements_in_boundary) + ? result[i] + : ElementSoftmaxCompute(-infinity_); + } + } + + CUTLASS_DEVICE + ElementSoftmaxCompute warp_reduce_sum_(ElementSoftmaxCompute sum_) { + int half_thread_in_row = (kThreadsPerRow >> 1); + CUTLASS_PRAGMA_UNROLL + for (int i = half_thread_in_row; i > 0; i >>= 1) { + ElementSoftmaxCompute tmp = __shfl_xor_sync(0xFFFFFFFF, sum_, i); + sum_ += tmp; + } + return sum_; + } + + CUTLASS_DEVICE + ElementSoftmaxCompute warp_reduce_max_(ElementSoftmaxCompute max_) { + int half_thread_in_row = (kThreadsPerRow >> 1); + CUTLASS_PRAGMA_UNROLL + for (int i = half_thread_in_row; i > 0; i >>= 1) { + ElementSoftmaxCompute tmp = __shfl_xor_sync(0xFFFFFFFF, max_, i); + max_ = fast_max(max_, tmp); + } + return max_; + } + + CUTLASS_DEVICE + void clear_accum_() { + uint32_t float_max_bits = 0xff7fffff; // -FLT_MAX + float min_float = reinterpret_cast(float_max_bits); + accum_max_ = ElementSoftmaxCompute(min_float); + accum_sum_ = ElementSoftmaxCompute(0); + } + + CUTLASS_DEVICE + ElementSoftmaxCompute sum_accumulator_(SoftmaxFragment const& accum) { + ElementSoftmaxCompute sum_ = ElementSoftmaxCompute(0); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < SoftmaxFragment::kElements; ++i) { + sum_ += ElementSoftmaxCompute(accum[i]); + } + + return sum_; + } + + CUTLASS_DEVICE + ElementSoftmaxCompute sum_accumulator_(SoftmaxFragment const& accum, + ElementSoftmaxCompute sum_) { + // ElementSoftmaxCompute sum_ = ElementSoftmaxCompute(0); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < SoftmaxFragment::kElements; ++i) { + sum_ += ElementSoftmaxCompute(accum[i]); + } + + return sum_; + } + + CUTLASS_DEVICE + ElementSoftmaxCompute maximum_accumulator_(SoftmaxFragment const& accum) { + ElementSoftmaxCompute max_ = accum[0]; + + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < SoftmaxFragment::kElements; ++i) { + max_ = fast_max(max_, ElementSoftmaxCompute(accum[i])); + } + + return max_; + } + + CUTLASS_DEVICE + ElementSoftmaxCompute maximum_accumulator_(SoftmaxFragment const& accum, + ElementSoftmaxCompute max_) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < SoftmaxFragment::kElements; ++i) { + max_ = fast_max(max_, ElementSoftmaxCompute(accum[i])); + } + + return max_; + } +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass diff --git a/csrc/attention/sm70_v37/gemm_with_softmax.h b/csrc/attention/sm70_v37/gemm_with_softmax.h new file mode 100644 index 0000000000..380f63d48e --- /dev/null +++ b/csrc/attention/sm70_v37/gemm_with_softmax.h @@ -0,0 +1,531 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + * this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/** + +*/ + +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include +#include +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/arch/memory.h" +#include "cutlass/arch/memory_sm75.h" + +#include "cutlass/gemm/kernel/default_gemm.h" +#include "cutlass/gemm/kernel/default_gemm_complex.h" +#include "cutlass/gemm/device/default_gemm_configuration.h" +#include "epilogue_visitor_with_softmax.h" +#include "cutlass/epilogue/threadblock/epilogue_with_visitor.h" +#include "reduce_softmax_final.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "gemm_with_epilogue_visitor.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Kernel computes partial reduction +// +// +// 2. Sum[m, n'] = sum_n(exp(D[m, n] - N[m, 0])) +// +template > +class ApplySoftmaxV37 { + public: + using ElementD = ElementD_; + using ElementNorm = ElementNorm_; + using ElementSum = ElementSum_; + using ElementSoft = ElementSoft_; + using ElementSoftmaxCompute = ElementSoftmaxCompute_; + + static int const kAlignment = Alignment; + using ApplyShape = ApplyShape_; + + using Layout = cutlass::layout::RowMajor; + + using TensorRefD = TensorRef; + using TensorRefN = TensorRef; + using TensorRefSum = TensorRef; + using TensorRefSoft = TensorRef; + + using FragmentSoftmax = Array; + + // + // Arguments + // + + struct Arguments { + MatrixCoord extent; ///< Extent of D and Softmax matrices + int batch_count; ///< Batch count + TensorRefD ref_D; ///< D matrix computed by GEMM+Max (input) + TensorRefN ref_N; ///< Norm tensor (input) + TensorRefSum ref_S; ///< Sum tensor (input) + TensorRefSoft ref_Soft; ///< Softmax tensor (output) + int64_t batch_stride_D; ///< Batch stride for D tensor + int64_t batch_stride_N; ///< Batch stride for N tensor + int64_t batch_stride_S; ///< Batch stride for S tensor + int64_t batch_stride_Soft; ///< Batch stride for softmax tensor + + // + // Methods + // + Arguments() + : batch_count(1), + batch_stride_D(0), + batch_stride_N(0), + batch_stride_S(0), + batch_stride_Soft(0) {} + + Arguments(MatrixCoord extent_, ///< Extent of D and Softmax matrices + int batch_count_, ///< Batch count + TensorRefD ref_D_, ///< D matrix computed by GEMM+PartialReduce + TensorRefN ref_N_, ///< Output parameter for N + TensorRefSum ref_S_, ///< Output parameter for N + TensorRefSoft ref_Soft_, ///< Softmax + int64_t batch_stride_D_ = 0, int64_t batch_stride_N_ = 0, + int64_t batch_stride_S_ = 0, int64_t batch_stride_Soft_ = 0) + : extent(extent_), + batch_count(batch_count_), + ref_D(ref_D_), + ref_N(ref_N_), + ref_S(ref_S_), + ref_Soft(ref_Soft_), + batch_stride_D(batch_stride_D_), + batch_stride_N(batch_stride_N_), + batch_stride_S(batch_stride_S_), + batch_stride_Soft(batch_stride_Soft_) {} + }; + + // + // Params struct + // + + struct Params { + Arguments args; + + // + // Methods + // + Params() {} + + Params(Arguments const& args_) : args(args_) {} + }; + + // + // SharedStorage + // + + struct SharedStorage {}; + + private: + public: + CUTLASS_DEVICE + ApplySoftmaxV37() {} + + CUTLASS_DEVICE + void operator()(Params const& params, SharedStorage& shared_storage) { + apply(params, shared_storage); + } + + private: + /// Compute Softmax + CUTLASS_DEVICE + void apply(Params const& params, SharedStorage& shared_storage) { + using AccessTypeD = AlignedArray; + + int block_batch = blockIdx.z; + int block_m = blockIdx.x * ApplyShape::kRow; + int block_n = 0; + + int thread_m = threadIdx.y; + int thread_n = threadIdx.x * kAlignment; + + int idx_m = block_m + thread_m; + int idx_n = block_n + thread_n; + + int batch_offset_norm = block_batch * params.args.batch_stride_N; + int batch_offset_sum = block_batch * params.args.batch_stride_S; + + // Kill off thread if it is outside the row boundary + if (params.args.extent.row() <= idx_m) { + return; + } + + // + // Setup pointers to load D again + // + + using AccessTypeD = AlignedArray; + using AccessTypeSoft = AlignedArray; + using FragmentSoft = Array; + using ConvertSoftCompute = + cutlass::NumericArrayConverter; + using ConvertSoftOutput = + cutlass::NumericArrayConverter; + + using Mul = cutlass::multiplies; + using Minus = cutlass::minus; + using Exp = cutlass::fast_exp_op; + + ConvertSoftCompute convert_soft_compute; + ConvertSoftOutput convert_soft_output; + + Minus minus; + Mul mul; + Exp exponential; + + using ConvertSum = + cutlass::NumericConverter; + using ConvertNorm = + cutlass::NumericConverter; + + ConvertSum convert_sum; + ConvertNorm convert_norm; + + AccessTypeD* access_d = reinterpret_cast( + params.args.ref_D.data() + params.args.batch_stride_D * block_batch + + params.args.ref_D.layout()({idx_m, idx_n})); + + AccessTypeSoft* access_soft = reinterpret_cast( + params.args.ref_Soft.data() + + params.args.batch_stride_Soft * block_batch + + params.args.ref_Soft.layout()({idx_m, idx_n})); + + ElementSum inv_sum = (params.args.ref_S.data())[idx_m + batch_offset_sum]; + ElementNorm norm = (params.args.ref_N.data())[idx_m + batch_offset_norm]; + + // + // Loop + // + CUTLASS_PRAGMA_UNROLL + for (int idx = 0; idx < params.args.extent.column(); + idx += ApplyShape::kColumn * kAlignment) { + if (idx_n < params.args.extent.column()) { + AccessTypeD fetch; + arch::global_load(fetch, access_d, + true); + + FragmentSoftmax result = mul( + exponential(minus(convert_soft_compute(fetch), convert_norm(norm))), + convert_sum(inv_sum)); + FragmentSoft soft = convert_soft_output(result); + + arch::global_store( + soft, access_soft, true); + } + + access_d += ApplyShape::kColumn; + access_soft += ApplyShape::kColumn; + idx_n += ApplyShape::kColumn * kAlignment; + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// +template , + int AlignmentA_ = 128 / cutlass::sizeof_bits::value, + int AlignmentB_ = 128 / cutlass::sizeof_bits::value, + int AlignmentSoftmax_ = 128 / cutlass::sizeof_bits::value, + typename ElementNorm_ = float, typename ElementSum_ = float, + typename ElementSoftmax_ = ElementC_> +class GemmSoftmaxV37 { + public: + /////////////////////////////////////////////////////////////////////////////////////////////// + + // + // Type definitions + // + + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementC = ElementC_; + using ElementCompute = ElementCompute_; + using ElementSum = ElementSum_; + using ElementSoft = ElementSoftmax_; + using ElementSoftmaxCompute = float; + + using LayoutA = LayoutA_; + using LayoutB = LayoutB_; + + using EpilogueFunctorOp = EpilogueFunctorOp_; + using ElementNorm = ElementNorm_; + + using ApplyShape = ApplyShape_; + + // These are mandatory layouts. + using LayoutC = cutlass::layout::RowMajor; + using LayoutN = cutlass::layout::RowMajor; + using LayoutS = cutlass::layout::RowMajor; + using LayoutSoft = cutlass::layout::RowMajor; + + using TensorRefA = TensorRef; + using TensorRefB = TensorRef; + using TensorRefC = TensorRef; + using TensorRefN = TensorRef; + using TensorRefSum = TensorRef; + using TensorRefSoft = TensorRef; + + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + + static int const kStages = kStages_; + static int const AlignmentA = AlignmentA_; + static int const AlignmentB = AlignmentB_; + static int const AlignmentSoftmax = AlignmentSoftmax_; + + using ThreadblockSwizzle = + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle; + + /////////////////////////////////////////////////////////////////////////////////////////////// + + // basic GEMM kernel + using DefaultGemmKernel = typename cutlass::gemm::kernel::DefaultGemm< + ElementA, LayoutA, AlignmentA, ElementB, LayoutB, AlignmentB, ElementC, + LayoutC, ElementCompute, OperatorClass, ArchTag, ThreadblockShape, + WarpShape, InstructionShape, EpilogueFunctorOp, ThreadblockSwizzle, + kStages, true, + typename cutlass::gemm::device::DefaultGemmConfiguration< + OperatorClass, ArchTag, ElementA, ElementB, ElementC, + ElementCompute>::Operator, + cutlass::gemm::SharedMemoryClearOption::kNone>::GemmKernel; + + /////////////////////////////////////////////////////////////////////////////////////////////// + + // Epilogue visitor + using EpilogueVisitor = + typename cutlass::epilogue::threadblock::EpilogueVisitorSoftmaxV37< + ThreadblockShape, DefaultGemmKernel::kThreadCount, + typename DefaultGemmKernel::Epilogue::OutputTileIterator, + ElementCompute, ElementNorm, ElementSum, ElementSoftmaxCompute, + EpilogueFunctorOp>; + + /// Epilogue + using Epilogue = typename cutlass::epilogue::threadblock:: + EpilogueWithVisitorFromExistingEpilogue< + EpilogueVisitor, typename DefaultGemmKernel::Epilogue>::Epilogue; + + // GEMM + using GemmKernel = + gemm::kernel::GemmWithEpilogueVisitor; + + // Softmax kernel + using SoftmaxApplyKernel = + kernel::ApplySoftmaxV37; + + using ApplyFinalReductionKernel = + cutlass::reduction::kernel::ApplySoftmaxFinalReductionV37< + ElementNorm, ElementSum, ElementSoftmaxCompute, ThreadblockShape>; + + public: + /// Arguments class + struct Arguments { + typename GemmKernel::Arguments gemm; + typename SoftmaxApplyKernel::Arguments softmax; + typename ApplyFinalReductionKernel::Arguments reduction; + cutlass::gemm::GemmCoord extend; + + // + // Methods + // + Arguments() {} + + Arguments(cutlass::gemm::GemmCoord problem_size, int32_t batch_count_, + TensorRefA ref_A_, TensorRefB ref_B_, TensorRefC ref_C_, + TensorRefC ref_D_, + typename EpilogueFunctorOp::Params linear_scaling, + TensorRefN ref_N_, TensorRefSum ref_S_, + TensorRefSoft ref_Softmax_, int64_t batch_stride_A_ = 0, + int64_t batch_stride_B_ = 0, int64_t batch_stride_C_ = 0, + int64_t batch_stride_D_ = 0, int64_t batch_stride_Max_ = 0, + int64_t batch_stride_Sum_ = 0, int64_t batch_stride_Softmax_ = 0) + : gemm(cutlass::gemm::GemmUniversalMode::kBatched, problem_size, + batch_count_, ref_A_, ref_B_, ref_C_, ref_D_, ref_N_.data(), + ref_S_.data(), batch_stride_A_, batch_stride_B_, + typename EpilogueVisitor::Arguments( + linear_scaling, batch_stride_C_, batch_stride_D_, + batch_stride_Max_, batch_stride_Sum_)), + reduction(problem_size, ref_N_.data(), ref_S_.data(), + batch_stride_Max_, batch_stride_Sum_), + softmax(MatrixCoord(problem_size.m(), problem_size.n()), batch_count_, + ref_D_, ref_N_, ref_S_, ref_Softmax_, batch_stride_D_, + batch_stride_Max_, batch_stride_Sum_, batch_stride_Softmax_), + extend(problem_size) {} + }; + + struct Params { + typename GemmKernel::Params gemm; + typename SoftmaxApplyKernel::Params softmax; + typename ApplyFinalReductionKernel::Params reduction; + MatrixCoord extend; + // + // Methods + // + Params() {} + + Params(Arguments const& args) + : gemm(args.gemm), + reduction(args.reduction), + softmax(args.softmax), + extend(MatrixCoord(args.extend.m(), args.extend.n())) {} + }; + + public: + // Gemm + + // + // Methods + // + + private: + Params params_; + + public: + /// Ctor + GemmSoftmaxV37() {} + + /// Initialize + Status initialize(Arguments const& args) { + params_ = Params(args); + + return cutlass::Status::kSuccess; + } + + /// Run + Status run(cudaStream_t stream) { + // + // Launch the GEMM + max kernel + // + + dim3 gemm_grid = + ThreadblockSwizzle().get_grid_shape(params_.gemm.grid_tiled_shape); + dim3 gemm_block(GemmKernel::kThreadCount, 1, 1); + + int gemm_smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + cudaError_t result; + + if (gemm_smem_size >= (48 << 10)) { + result = cudaFuncSetAttribute(cutlass::Kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + gemm_smem_size); + + if (result != cudaSuccess) { + return Status::kErrorInternal; + } + } + + cutlass::Kernel + <<>>(params_.gemm); + + result = cudaGetLastError(); + + if (result != cudaSuccess) { + return cutlass::Status::kErrorInternal; + } + + // + // Launch the ApplyFinalReductionKernel + // + + int thread_per_block = 128; + int block_per_row = + (params_.extend.row() + thread_per_block - 1) / thread_per_block; + if (block_per_row < 4) { + thread_per_block = 32; + block_per_row = + (params_.extend.row() + thread_per_block - 1) / thread_per_block; + } + + dim3 final_reduction_grid(block_per_row, 1, + params_.softmax.args.batch_count); + dim3 final_reduction_block(thread_per_block); + + Kernel + <<>>( + params_.reduction); + + result = cudaGetLastError(); + + if (result != cudaSuccess) { + return cutlass::Status::kErrorInternal; + } + + return cutlass::Status::kSuccess; + } + + /// Function call operator + Status operator()(cudaStream_t stream = nullptr) { return run(stream); } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/attention/sm70_v37/prefill.cu b/csrc/attention/sm70_v37/prefill.cu new file mode 100644 index 0000000000..469a29b00c --- /dev/null +++ b/csrc/attention/sm70_v37/prefill.cu @@ -0,0 +1,666 @@ +// SPDX-License-Identifier: BSD-3-Clause +#undef FLASH_NAMESPACE +#define FLASH_NAMESPACE onecat_v37 + +/*************************************************************************************************** + * Precision-qualified v37 long-prefix attention for V100/SM70. + * + * Tile-local probabilities are formed from FP32 logits and statistics before + * conversion to FP16 Tensor Core operands. QK/PV accumulation, online partials + * and the exact causal-tail output stay FP32. Prefix and tail overlap without + * changing their arithmetic. The legacy kernel is a separate rollback route. + **************************************************************************************************/ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include "namespace_config.h" + +#include "cutlass/cutlass.h" +#include "cutlass/device_kernel.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/kernel/default_gemm.h" +#include "cutlass/gemm/kernel/gemm.h" +#include "cutlass/gemm/threadblock/mma_pipelined.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/numeric_conversion.h" +#include "gemm_with_softmax.h" + +namespace FLASH_NAMESPACE { +namespace { + +using Element = cutlass::half_t; +using TailOutput = float; + +using QKAccumulator = float; + +using QKLayoutA = cutlass::layout::RowMajor; +using QKLayoutB = cutlass::layout::ColumnMajor; +using QKThreadblockShape = cutlass::gemm::GemmShape<128, 256, 32>; +using QKWarpShape = cutlass::gemm::GemmShape<64, 64, 32>; +using QKInstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; +using QKLinearOutputOp = + cutlass::epilogue::thread::LinearCombination; +// A distinct functor type keeps the modified visitor's kernel symbols local +// to this variant, even when baseline and candidate share one extension. +struct QKOutputOp : QKLinearOutputOp { + CUTLASS_HOST_DEVICE + explicit QKOutputOp(Params const& params) : QKLinearOutputOp(params) {} +}; +using QKGemm = + cutlass::GemmSoftmaxV37>; + +using PVLayout = cutlass::layout::RowMajor; +using PVAccumulator = float; +using PVOutput = float; +constexpr int kPVOutputAccess = 4; +using PVThreadblockShape = cutlass::gemm::GemmShape<128, 256, 32>; +using PVWarpShape = cutlass::gemm::GemmShape<64, 64, 32>; +using PVInstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; +using PVSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; +using PVLinearOutputOp = + cutlass::epilogue::thread::LinearCombination; +static_assert(std::is_same::value && + std::is_same::value, + "Direct online state requires FP32 output and accumulation"); +CUTLASS_DEVICE int pv_output_row(int sequence); + +__device__ int g_rows = 0; + +struct PVOutputOp : PVLinearOutputOp { + using Base = PVLinearOutputOp; + struct Params : Base::Params { + float const* old_scales; + float const* block_scales; + bool initialize; + CUTLASS_HOST_DEVICE + Params(float alpha = 1.0f, float beta = 0.0f) + : Base::Params(alpha, beta), + old_scales(nullptr), + block_scales(nullptr), + initialize(true) {} + CUTLASS_HOST_DEVICE + Params(float const* old_, float const* next_, bool init) + : Base::Params(1.0f, init ? 0.0f : 1.0f), + old_scales(old_), + block_scales(next_), + initialize(init) {} + }; + float const* old_scales; + float const* block_scales; + bool initialize; + mutable int sequence = 0; + CUTLASS_HOST_DEVICE + explicit PVOutputOp(Params const& params) + : Base(params), + old_scales(params.old_scales), + block_scales(params.block_scales), + initialize(params.initialize) {} + CUTLASS_HOST_DEVICE + bool is_source_needed() const { return !initialize; } + CUTLASS_DEVICE + FragmentOutput operator()(FragmentAccumulator const& accum, + FragmentSource const& source) const { + int row = pv_output_row(sequence++); + FragmentOutput result; + float a = row < g_rows ? old_scales[row] : 0.0f; + float b = row < g_rows ? block_scales[row] : 0.0f; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) + result[i] = fmaf(source[i], a, accum[i] * b); + return result; + } + CUTLASS_DEVICE + FragmentOutput operator()(FragmentAccumulator const& accum) const { + int row = pv_output_row(sequence++); + FragmentOutput result; + float b = row < g_rows ? block_scales[row] : 0.0f; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) result[i] = accum[i] * b; + return result; + } +}; + +constexpr int kPVAlignment = 8; + +using PVDefaultKernel = typename cutlass::gemm::kernel::DefaultGemm< + Element, PVLayout, kPVAlignment, Element, PVLayout, kPVAlignment, PVOutput, + PVLayout, PVAccumulator, cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm70, PVThreadblockShape, PVWarpShape, PVInstructionShape, + PVOutputOp, PVSwizzle, 2, false, cutlass::arch::OpMultiplyAdd>::GemmKernel; + +using PVDefaultMma = typename PVDefaultKernel::Mma; +CUTLASS_DEVICE int pv_output_row(int sequence) { + using Iterator = typename PVDefaultKernel::Epilogue::OutputTileIterator; + using Map = typename Iterator::ThreadMap; + constexpr int kAccesses = + Iterator::Fragment::kElements / Iterator::kElementsPerAccess; + int step = sequence / kAccesses; + int fragment = sequence % kAccesses; + Iterator iterator(typename Iterator::Params(PVLayout(256)), nullptr, + {g_rows, 256}, threadIdx.x, + {int(blockIdx.x) * PVThreadblockShape::kM, + int(blockIdx.y) * PVThreadblockShape::kN}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Iterator::kIterations; ++i) { + if (i < step) ++iterator; + } + return iterator.thread_start().row() + Map::iteration_offset(fragment).row(); +} +using PVIteratorA = typename PVDefaultMma::IteratorA; +using PVIteratorB = typename PVDefaultMma::IteratorB; +using PVSmemIteratorA = typename PVDefaultMma::SmemIteratorA; +using PVSmemIteratorB = typename PVDefaultMma::SmemIteratorB; + +__device__ float const* g_row_max = nullptr; +__device__ float const* g_row_inv_sum = nullptr; + +struct ExpRowSumTransformA { + using InputFragment = typename PVIteratorA::Fragment; + using OutputFragment = cutlass::Array; + using ThreadMap = typename PVIteratorA::ThreadMap; + static constexpr int kAccessesPerVector = + PVIteratorA::UnderlyingIterator::kAccessesPerVector; + static constexpr int kContiguousIterations = + ThreadMap::Iterations::kContiguous; + static constexpr int kStridedIterations = ThreadMap::Iterations::kStrided; + + float row_max[kStridedIterations]; + float row_inv_sum[kStridedIterations]; + int row_index[kStridedIterations]; + int k_offset = 0; + float* shared_tile_scale; + + CUTLASS_DEVICE + ExpRowSumTransformA() { + // Query lengths are multiples of 64, so six packed heads fill M128 tiles. + constexpr int kScaleCount = (8192 / 256) * 128; + constexpr int kThreads = 32 * (128 / 64) * (256 / 64); + __shared__ float tile_scale_smem[kScaleCount]; + shared_tile_scale = tile_scale_smem; +#pragma unroll + for (int i = threadIdx.x; i < kScaleCount; i += kThreads) { + int tile = i / 128; + int row = blockIdx.x * 128 + i % 128; + int stored_tile = tile == 0 ? 8192 / 256 : tile; + float delta = g_row_max[stored_tile * g_rows + row] - g_row_max[row]; + shared_tile_scale[i] = exp2f(delta * 1.4426950408889634f); + } + __syncthreads(); + auto thread_offset = ThreadMap::initial_offset(threadIdx.x); +#pragma unroll + for (int s = 0; s < kStridedIterations; ++s) { + int row = blockIdx.x * PVThreadblockShape::kM + thread_offset.strided() + + s * ThreadMap::Delta::kStrided; + row_max[s] = row < g_rows ? g_row_max[row] : 0.0f; + row_inv_sum[s] = row < g_rows ? g_row_inv_sum[row] : 0.0f; + row_index[s] = row - blockIdx.x * 128; + } + } + + CUTLASS_DEVICE + OutputFragment operator()(InputFragment const& input) { + OutputFragment output; + constexpr float kLog2E = 1.4426950408889634f; + constexpr int kElementsPerAccess = PVIteratorA::AccessType::kElements; + auto const* input_access = + reinterpret_cast(&input); + auto* output_access = + reinterpret_cast(&output); +#pragma unroll + for (int s = 0; s < kStridedIterations; ++s) { + int tile = k_offset / 256; + int stored_tile = tile == 0 ? 8192 / 256 : tile; + // The double-buffered pipeline transforms one final masked look-ahead + // fragment. Its zero values are unused, but its metadata load must stay + // in bounds as well. Wrap that sentinel tile to the valid first tile. + int shared_tile = tile % (8192 / 256); + float probability_scale = + shared_tile_scale[shared_tile * 128 + row_index[s]]; +#pragma unroll + for (int c = 0; c < kContiguousIterations; ++c) { +#pragma unroll + for (int v = 0; v < kAccessesPerVector; ++v) { + int index = v + kAccessesPerVector * (c + s * kContiguousIterations); + typename PVIteratorA::AccessType transformed; +#pragma unroll + for (int e = 0; e < kElementsPerAccess; ++e) { + float value = static_cast(input_access[index][e]); + float weight = value * probability_scale; + transformed[e] = Element(weight); + } + output_access[index] = transformed; + } + } + } + k_offset += PVThreadblockShape::kK; + return output; + } +}; + +using PVTransformB = + cutlass::NumericArrayConverter; +using PVMma = cutlass::gemm::threadblock::MmaPipelined< + typename PVDefaultMma::Shape, PVIteratorA, PVSmemIteratorA, PVIteratorB, + PVSmemIteratorB, PVAccumulator, PVLayout, typename PVDefaultMma::Policy, + ExpRowSumTransformA, PVTransformB>; +using PVKernel = + cutlass::gemm::kernel::Gemm; + +void check(cudaError_t result, char const* operation) { + TORCH_CHECK(result == cudaSuccess, operation, ": ", + cudaGetErrorString(result)); +} + +struct PVLauncher { + typename PVKernel::Params params; + dim3 grid; + dim3 block; + int smem_bytes; + + PVLauncher(Element* scores, Element* value, PVOutput* output, int rows, int k, + float const* old_scales = nullptr, + float const* block_scales = nullptr, bool initialize = true) { + cutlass::gemm::GemmCoord problem(rows, 256, k); + PVSwizzle swizzle; + auto tiled_shape = + swizzle.get_tiled_shape(problem, + {PVThreadblockShape::kM, PVThreadblockShape::kN, + PVThreadblockShape::kK}, + 1); + params = typename PVKernel::Params( + problem, tiled_shape, {scores, PVLayout(k)}, {value, PVLayout(256)}, + {output, PVLayout(256)}, {output, PVLayout(256)}, + typename PVOutputOp::Params(old_scales, block_scales, initialize), + nullptr); + grid = swizzle.get_grid_shape(tiled_shape); + block = dim3(PVKernel::kThreadCount, 1, 1); + smem_bytes = int(sizeof(typename PVKernel::SharedStorage)); + if (smem_bytes >= 48 * 1024) { + check(cudaFuncSetAttribute(cutlass::Kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_bytes), + "set PV dynamic shared memory"); + } + } + + void launch(cudaStream_t stream) const { + cutlass::Kernel<<>>(params); + } +}; + +__global__ void prepare_prefix_update(float const* block_max, + float const* block_inv_sum, + float* prefix_max, float* prefix_sum, + float* old_scales, float* block_scales, + int rows, bool initialize) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= rows) { + return; + } + constexpr float kLog2E = 1.4426950408889634f; + float next_max = block_max[row]; + float next_mass = 1.0f / block_inv_sum[row]; + if (initialize) { + old_scales[row] = 0.0f; + block_scales[row] = 1.0f; + prefix_max[row] = next_max; + prefix_sum[row] = next_mass; + return; + } + float old_max = prefix_max[row]; + float global_max = fmaxf(old_max, next_max); + float old_scale = exp2f((old_max - global_max) * kLog2E); + float block_scale = next_mass * exp2f((next_max - global_max) * kLog2E); + old_scales[row] = old_scale; + block_scales[row] = exp2f((next_max - global_max) * kLog2E); + prefix_max[row] = global_max; + prefix_sum[row] = prefix_sum[row] * old_scale + block_scale; +} + +__global__ void merge_prefix_accumulator_tail( + float const* prefix_accumulator, float const* prefix_max, + float const* prefix_sum, TailOutput const* tail_output, + float const* tail_max, float const* tail_sum, __half* output, int rows) { + int row = blockIdx.x; + int d = threadIdx.x; + __shared__ float masses[3]; + if (d == 0) { + constexpr float kLog2E = 1.4426950408889634f; + float global_max = fmaxf(prefix_max[row], tail_max[row]); + masses[0] = exp2f((prefix_max[row] - global_max) * kLog2E); + masses[1] = exp2f((tail_max[row] - global_max) * kLog2E); + masses[2] = + 1.0f / (prefix_sum[row] * masses[0] + tail_sum[row] * masses[1]); + } + __syncthreads(); + int64_t element = int64_t(row) * 256 + d; + float numerator = prefix_accumulator[element] * masses[0] + + float(tail_output[element]) * masses[1]; + output[element] = __float2half_rn(numerator * masses[2]); +} + +struct BlockOperators { + QKGemm qk; + std::unique_ptr pv; +}; + +} // namespace +} // namespace FLASH_NAMESPACE + +extern "C" cudaError_t onecat_v37_dense_state_float_raw( + const void*, const void*, const void*, float*, float*, void*, int, int, int, + int, float, int, cudaStream_t); + +namespace FLASH_NAMESPACE { + +struct Sm70GqaScoreWorkspace { + at::Tensor scores; + cudaEvent_t completion = nullptr; + cudaStream_t tail_stream = nullptr; + cudaEvent_t tail_input_ready = nullptr; + cudaEvent_t tail_complete = nullptr; + bool completion_recorded = false; + std::mutex launch_mutex; + + Sm70GqaScoreWorkspace(const at::Tensor& q, int rows, int block_n) + : scores(at::empty({rows, block_n}, q.options())) { + auto cleanup = c10::make_scope_exit([&]() noexcept { + if (tail_stream) cudaStreamDestroy(tail_stream); + if (tail_input_ready) cudaEventDestroy(tail_input_ready); + if (tail_complete) cudaEventDestroy(tail_complete); + if (completion) cudaEventDestroy(completion); + }); + C10_CUDA_CHECK( + cudaEventCreateWithFlags(&completion, cudaEventDisableTiming)); + C10_CUDA_CHECK( + cudaStreamCreateWithFlags(&tail_stream, cudaStreamNonBlocking)); + C10_CUDA_CHECK( + cudaEventCreateWithFlags(&tail_input_ready, cudaEventDisableTiming)); + C10_CUDA_CHECK( + cudaEventCreateWithFlags(&tail_complete, cudaEventDisableTiming)); + cleanup.release(); + } + + ~Sm70GqaScoreWorkspace() { + if (tail_stream) cudaStreamDestroy(tail_stream); + if (tail_input_ready) cudaEventDestroy(tail_input_ready); + if (tail_complete) cudaEventDestroy(tail_complete); + if (completion != nullptr) { + cudaEventDestroy(completion); + } + } +}; + +using Sm70GqaScoreWorkspacePtr = std::shared_ptr; + +std::mutex& sm70_gqa_score_cache_mutex() { + static std::mutex cache_mutex; + return cache_mutex; +} + +std::map& sm70_gqa_score_cache() { + static std::map cache; + return cache; +} + +Sm70GqaScoreWorkspacePtr get_sm70_gqa_score_workspace(const at::Tensor& q, + int rows, int block_n) { + std::lock_guard lock(sm70_gqa_score_cache_mutex()); + int device = q.get_device(); + auto& cache = sm70_gqa_score_cache(); + auto& workspace = cache[device]; + if (!workspace) { + workspace = std::make_shared(q, rows, block_n); + } + return workspace; +} + +Sm70GqaScoreWorkspacePtr find_sm70_gqa_score_workspace(int device) { + std::lock_guard lock(sm70_gqa_score_cache_mutex()); + auto& cache = sm70_gqa_score_cache(); + auto found = cache.find(device); + return found == cache.end() ? nullptr : found->second; +} + +at::Tensor sm70_d256_gqa_v37_fwd(const at::Tensor& q, const at::Tensor& k, + const at::Tensor& v, at::Tensor& out, + double softmax_scale, bool causal) { + constexpr int kMaxQuery = 8192; + constexpr int kMaxRows = kMaxQuery * 6; + constexpr int kHeadDim = 256; + constexpr int kHeadsQ = 6; + constexpr int kHeadsKV = 1; + constexpr int kMaxTotalKV = 262144; + constexpr int kTotalKVStep = 32; + constexpr int kBlockN = 8192; + + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && out.is_cuda(), + "SM70 GQA architecture requires CUDA tensors"); + TORCH_CHECK(q.scalar_type() == at::ScalarType::Half && + k.scalar_type() == q.scalar_type() && + v.scalar_type() == q.scalar_type() && + out.scalar_type() == q.scalar_type(), + "SM70 GQA architecture requires FP16 q, k, v, and out"); + TORCH_CHECK(q.dim() == 4 && q.size(0) == 1 && q.size(1) >= 64 && + q.size(1) <= kMaxQuery && q.size(1) % 64 == 0 && + q.size(2) == kHeadsQ && q.size(3) == kHeadDim && + k.dim() == 4 && k.size(0) == 1 && k.size(2) == kHeadsKV && + k.size(3) == kHeadDim && v.sizes() == k.sizes() && + out.sizes() == q.sizes(), + "SM70 GQA architecture only accepts the validated " + "Q64..8192 (multiple of 64)/Hq6/Hkv1/D256 dense shape family"); + const int kQuery = static_cast(q.size(1)); + const int kRows = kQuery * kHeadsQ; + const int kTail = kQuery; + TORCH_CHECK(k.size(1) > kQuery && k.size(1) <= kMaxTotalKV, + "v37 requires a non-empty prefix and KV <= 262144"); + const int total_kv = static_cast(k.size(1)); + TORCH_CHECK(total_kv % kTotalKVStep == 0, + "SM70 GQA architecture requires a 32-token KV step, got ", + total_kv); + const int prefix = total_kv - kTail; + const int blocks = (prefix + kBlockN - 1) / kBlockN; + TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous() && + out.is_contiguous(), + "SM70 GQA architecture requires contiguous tensors"); + TORCH_CHECK(q.get_device() == k.get_device() && + q.get_device() == v.get_device() && + q.get_device() == out.get_device(), + "SM70 GQA architecture tensors must share one device"); + TORCH_CHECK(causal, "SM70 GQA architecture requires causal attention"); + TORCH_CHECK(std::abs(softmax_scale - 0.0625) < 1.0e-8, + "SM70 GQA architecture requires D256 softmax scale 1/16"); + + const at::cuda::OptionalCUDAGuard device_guard(q.device()); + const auto* properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(properties->major == 7 && properties->minor == 0, + "v37 prefill supports SM70 only"); + for (const at::Tensor* tensor : + {&q, &k, &v, static_cast(&out)}) { + TORCH_CHECK(reinterpret_cast(tensor->data_ptr()) % 16 == 0, + "v37 prefill requires 16-byte aligned tensors"); + } + const int qk_tiles_n = + (kBlockN + QKThreadblockShape::kN - 1) / QKThreadblockShape::kN; + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + cudaStreamCaptureStatus capture_status; + C10_CUDA_CHECK(cudaStreamIsCapturing(stream, &capture_status)); + TORCH_CHECK( + capture_status == cudaStreamCaptureStatusNone, + "v37 prefill does not support CUDA Graph capture; use exact fallback"); + + constexpr size_t kScratchAlignment = 256; + size_t scratch_bytes = 0; + auto reserve_scratch = [&](size_t bytes) { + scratch_bytes = + (scratch_bytes + kScratchAlignment - 1) & ~(kScratchAlignment - 1); + size_t offset = scratch_bytes; + scratch_bytes += bytes; + return offset; + }; + size_t prefix_accumulator_offset = + reserve_scratch(size_t(kRows) * kHeadDim * sizeof(float)); + size_t tail_output_offset = + reserve_scratch(size_t(kRows) * kHeadDim * sizeof(TailOutput)); + size_t qk_norm_offset = + reserve_scratch(size_t(qk_tiles_n + 1) * kRows * sizeof(float)); + size_t qk_sum_offset = + reserve_scratch(size_t(qk_tiles_n) * kRows * sizeof(float)); + size_t prefix_max_offset = reserve_scratch(size_t(kRows) * sizeof(float)); + size_t prefix_sum_offset = reserve_scratch(size_t(kRows) * sizeof(float)); + size_t old_scale_offset = reserve_scratch(size_t(kRows) * sizeof(float)); + size_t block_scale_offset = reserve_scratch(size_t(kRows) * sizeof(float)); + size_t tail_max_offset = reserve_scratch(size_t(kRows) * sizeof(float)); + size_t tail_sum_offset = reserve_scratch(size_t(kRows) * sizeof(float)); + int device = q.get_device(); + auto score_workspace = find_sm70_gqa_score_workspace(device); + if (!score_workspace) { + constexpr size_t kRequiredPostWorkspaceHeadroom = 128 * 1024 * 1024; + constexpr size_t kScoreBytes = size_t(kMaxRows) * kBlockN * sizeof(Element); + size_t free_bytes = 0; + size_t total_bytes = 0; + C10_CUDA_CHECK(cudaMemGetInfo(&free_bytes, &total_bytes)); + size_t required_bytes = + kScoreBytes + scratch_bytes + kRequiredPostWorkspaceHeadroom; + TORCH_CHECK_WITH( + OutOfMemoryError, free_bytes >= required_bytes, + "SM70 GQA architecture requires ", required_bytes / (1024 * 1024), + " MiB free before its first workspace allocation, including " + "128 MiB of downstream headroom, but only ", + free_bytes / (1024 * 1024), " MiB remains out of ", + total_bytes / (1024 * 1024), " MiB"); + score_workspace = get_sm70_gqa_score_workspace(q, kMaxRows, kBlockN); + } + std::unique_lock launch_lock(score_workspace->launch_mutex); + if (score_workspace->completion_recorded) { + C10_CUDA_CHECK(cudaStreamWaitEvent(stream, score_workspace->completion, 0)); + } + at::Tensor scratch = at::empty({static_cast(scratch_bytes)}, + q.options().dtype(at::ScalarType::Byte)); + auto* scratch_base = scratch.data_ptr(); + // On an exception, the private tail must stop using caller-owned tensors + // before scratch is released. Also order any queued prefix writes before + // a subsequent call reuses the shared score/metadata cache. + auto failed_launch = c10::make_scope_exit([&]() noexcept { + cudaStreamSynchronize(score_workspace->tail_stream); + score_workspace->completion_recorded = + cudaEventRecord(score_workspace->completion, stream) == cudaSuccess; + }); + + auto* query = reinterpret_cast(q.data_ptr()); + auto* key = reinterpret_cast(k.data_ptr()); + auto* value = reinterpret_cast(v.data_ptr()); + auto* score_ptr = + reinterpret_cast(score_workspace->scores.data_ptr()); + float* prefix_accumulator_ptr = + reinterpret_cast(scratch_base + prefix_accumulator_offset); + auto* tail_output_ptr = + reinterpret_cast(scratch_base + tail_output_offset); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + float* qk_norm_ptr = reinterpret_cast(scratch_base + qk_norm_offset); + float* qk_sum_ptr = reinterpret_cast(scratch_base + qk_sum_offset); + float* prefix_max_ptr = + reinterpret_cast(scratch_base + prefix_max_offset); + float* prefix_sum_ptr = + reinterpret_cast(scratch_base + prefix_sum_offset); + float* old_scale_ptr = + reinterpret_cast(scratch_base + old_scale_offset); + float* block_scale_ptr = + reinterpret_cast(scratch_base + block_scale_offset); + float* tail_max_ptr = + reinterpret_cast(scratch_base + tail_max_offset); + float* tail_sum_ptr = + reinterpret_cast(scratch_base + tail_sum_offset); + + auto set_pv_metadata = [&]() { + C10_CUDA_CHECK(cudaMemcpyToSymbolAsync(g_rows, &kRows, sizeof(kRows), 0, + cudaMemcpyHostToDevice, stream)); + float const* persistent_max_ptr = qk_norm_ptr; + C10_CUDA_CHECK(cudaMemcpyToSymbolAsync(g_row_max, &persistent_max_ptr, + sizeof(persistent_max_ptr), 0, + cudaMemcpyHostToDevice, stream)); + float const* persistent_inv_sum_ptr = qk_sum_ptr; + C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( + g_row_inv_sum, &persistent_inv_sum_ptr, sizeof(persistent_inv_sum_ptr), + 0, cudaMemcpyHostToDevice, stream)); + }; + // Set metadata before forking the independent causal tail. + set_pv_metadata(); + // Fork only after gathered inputs are ready; join before the sole consumer + // of tail state. Prefix/tail arithmetic and their scratch regions are + // disjoint. + C10_CUDA_CHECK(cudaEventRecord(score_workspace->tail_input_ready, stream)); + C10_CUDA_CHECK(cudaStreamWaitEvent(score_workspace->tail_stream, + score_workspace->tail_input_ready, 0)); + cudaStream_t tail_stream = score_workspace->tail_stream; + C10_CUDA_CHECK(onecat_v37_dense_state_float_raw( + query, key + size_t(prefix) * kHeadDim, value + size_t(prefix) * kHeadDim, + tail_max_ptr, tail_sum_ptr, tail_output_ptr, kTail, kTail, kHeadsQ, + kHeadsKV, static_cast(softmax_scale), 1, tail_stream)); + C10_CUDA_CHECK(cudaEventRecord(score_workspace->tail_complete, tail_stream)); + + for (int block = 0; block < blocks; ++block) { + int begin = block * kBlockN; + int width = std::min(kBlockN, prefix - begin); + BlockOperators operation; + typename QKGemm::Arguments arguments( + {kRows, width, kHeadDim}, 1, {query, QKLayoutA(kHeadDim)}, + {key + size_t(begin) * kHeadDim, QKLayoutB(kHeadDim)}, + {score_ptr, typename QKGemm::LayoutC(width)}, + {score_ptr, typename QKGemm::LayoutC(width)}, + {Element(static_cast(softmax_scale)), Element(0.0f)}, + {qk_norm_ptr, typename QKGemm::LayoutN(kRows)}, + {qk_sum_ptr, typename QKGemm::LayoutS(kRows)}, + {score_ptr, typename QKGemm::LayoutSoft(width)}); + TORCH_CHECK(operation.qk.initialize(arguments) == cutlass::Status::kSuccess, + "initialize SM70 GQA QK block ", block, " failed"); + TORCH_CHECK(operation.qk(stream) == cutlass::Status::kSuccess, + "launch SM70 GQA QK block ", block, " failed"); + prepare_prefix_update<<<(kRows + 255) / 256, 256, 0, stream>>>( + qk_norm_ptr, qk_sum_ptr, prefix_max_ptr, prefix_sum_ptr, old_scale_ptr, + block_scale_ptr, kRows, block == 0); + operation.pv = std::make_unique( + score_ptr, value + size_t(begin) * kHeadDim, prefix_accumulator_ptr, + kRows, width, old_scale_ptr, block_scale_ptr, block == 0); + operation.pv->launch(stream); + } + C10_CUDA_CHECK( + cudaStreamWaitEvent(stream, score_workspace->tail_complete, 0)); + merge_prefix_accumulator_tail<<>>( + prefix_accumulator_ptr, prefix_max_ptr, prefix_sum_ptr, + reinterpret_cast(tail_output_ptr), tail_max_ptr, + tail_sum_ptr, reinterpret_cast<__half*>(output_ptr), kRows); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + C10_CUDA_CHECK(cudaEventRecord(score_workspace->completion, stream)); + score_workspace->completion_recorded = true; + failed_launch.release(); + return out; +} + +} // namespace FLASH_NAMESPACE diff --git a/csrc/attention/sm70_v37/reduce_softmax_final.h b/csrc/attention/sm70_v37/reduce_softmax_final.h new file mode 100644 index 0000000000..4c2a3df8d7 --- /dev/null +++ b/csrc/attention/sm70_v37/reduce_softmax_final.h @@ -0,0 +1,256 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + * this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a final reduction for softmax +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/array.h" +#include "cutlass/functional.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/arch/memory.h" +#include "cutlass/arch/memory_sm75.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace reduction { +namespace kernel { + +template +class ApplySoftmaxFinalReductionV37 { + public: + using ElementNorm = ElementNorm_; + using ElementSum = ElementSum_; + using ElementSoftmaxCompute = ElementSoftmaxCompute_; + using ThreadblockShape = ThreadblockShape_; + static const bool isGroupedProblem = GroupedProblem; + + // + // Arguments + // + + struct Arguments { + cutlass::gemm::GemmCoord* problem_sizes{nullptr}; + cutlass::gemm::GemmCoord problem_size{}; + ElementNorm* block_Norm{nullptr}; + ElementSum* block_Sum{nullptr}; + int64_t* offset_Norm_Device{nullptr}; + int64_t* offset_Sum_Device{nullptr}; + int64_t batch_stride_Max{0}; + int64_t batch_stride_Sum{0}; + + // + // Methods + // + Arguments() {} + + // Non-grouped constructor without batching + Arguments(cutlass::gemm::GemmCoord problem_size, ElementNorm* block_Norm, + ElementSum* block_Sum) + : problem_size(problem_size), + block_Norm(block_Norm), + block_Sum(block_Sum), + problem_sizes(nullptr), + offset_Norm_Device(nullptr), + offset_Sum_Device(nullptr), + batch_stride_Max(0), + batch_stride_Sum(0) {} + + // Non-grouped constructor with batching + Arguments(cutlass::gemm::GemmCoord problem_size, ElementNorm* block_Norm, + ElementSum* block_Sum, int64_t batch_stride_Max, + int64_t batch_stride_Sum) + : problem_size(problem_size), + block_Norm(block_Norm), + block_Sum(block_Sum), + batch_stride_Max(batch_stride_Max), + batch_stride_Sum(batch_stride_Sum), + problem_sizes(nullptr), + offset_Norm_Device(nullptr), + offset_Sum_Device(nullptr) {} + + // Grouped constructor + Arguments(cutlass::gemm::GemmCoord* problem_sizes, ElementNorm* block_Norm, + ElementSum* block_Sum, int64_t* offset_Norm_Device, + int64_t* offset_Sum_Device) + : problem_sizes(problem_sizes), + problem_size(cutlass::gemm::GemmCoord(0, 0, 0)), + block_Norm(block_Norm), + block_Sum(block_Sum), + offset_Norm_Device(offset_Norm_Device), + offset_Sum_Device(offset_Sum_Device) {} + }; + + struct SharedStorage {}; + + // + // Params struct + // + + struct Params { + Arguments args; + + // + // Methods + // + Params() {} + + Params(Arguments const& args_) : args(args_) {} + }; + + private: + public: + CUTLASS_DEVICE + ApplySoftmaxFinalReductionV37() {} + + CUTLASS_DEVICE + void operator()(Params const& params, SharedStorage& shared_storage) { + apply(params, shared_storage); + } + + private: + /// Full reduction + CUTLASS_DEVICE + void apply(Params const& params, SharedStorage& shared_storage) { + int tid = threadIdx.x; + int bid = blockIdx.x; + int bdim = blockDim.x; + + int block_batch = blockIdx.z; + + // defining three vars for a general reduction module + cutlass::gemm::GemmCoord problem_size = isGroupedProblem + ? params.args.problem_sizes[bid] + : params.args.problem_size; + int m_dim_in_loop = isGroupedProblem ? problem_size.m() : tid + bdim; + int access_offset = isGroupedProblem ? 0 : bid * bdim; + + if (!isGroupedProblem && access_offset + tid >= problem_size.m()) return; + + ElementNorm* curr_ptr_Max = + isGroupedProblem + ? params.args.block_Norm + params.args.offset_Norm_Device[bid] + : params.args.block_Norm + + block_batch * params.args.batch_stride_Max; + ElementSum* curr_ptr_Sum = + isGroupedProblem + ? params.args.block_Sum + params.args.offset_Sum_Device[bid] + : params.args.block_Sum + + block_batch * params.args.batch_stride_Sum; + + int threadblock_num = + (problem_size.n() + ThreadblockShape::kN - 1) / ThreadblockShape::kN; + + using ConvertSumOutput = + cutlass::NumericConverter; + using ConvertNormOutput = + cutlass::NumericConverter; + + using ConvertSum = + cutlass::NumericConverter; + using ConvertNorm = + cutlass::NumericConverter; + + ConvertSum convert_sum; + ConvertNorm convert_norm; + + ConvertSumOutput convert_sum_output; + ConvertNormOutput convert_norm_output; + + uint32_t float_max_bits = 0xff7fffff; + float min_float = reinterpret_cast(float_max_bits); + + CUTLASS_PRAGMA_UNROLL + for (int idx_m = tid; idx_m < m_dim_in_loop; idx_m += bdim) { + ElementNorm* access_n = curr_ptr_Max + idx_m + access_offset; + ElementSum* access_s = curr_ptr_Sum + idx_m + access_offset; + ElementNorm* access_n_bak = access_n; + ElementSum* access_s_bak = access_s; + ElementSoftmaxCompute max_val = ElementSoftmaxCompute(min_float); + ElementSoftmaxCompute sum_val = ElementSoftmaxCompute(0); + ElementNorm fetch_n; + ElementSum fetch_s; + + CUTLASS_PRAGMA_UNROLL + for (int idx_n = 0; idx_n < threadblock_num; idx_n++) { + cutlass::arch::global_load( + fetch_n, access_n, true); + max_val = cutlass::fast_max(max_val, convert_norm(fetch_n)); + access_n += problem_size.m(); + } + + access_n = access_n_bak; + + CUTLASS_PRAGMA_UNROLL + for (int idx_n = 0; idx_n < threadblock_num; idx_n++) { + cutlass::arch::global_load( + fetch_n, access_n, true); + cutlass::arch::global_load( + fetch_s, access_s, true); + ElementSoftmaxCompute tile_scale = + cutlass::fast_exp(convert_norm(fetch_n) - max_val); + sum_val += convert_sum(fetch_s) * tile_scale; + if constexpr (PaperTileScales) { + ElementNorm* destination = + idx_n == 0 ? access_n_bak + + (8192 / ThreadblockShape::kN) * problem_size.m() + : access_n; + *destination = convert_norm_output(tile_scale); + } + access_n += problem_size.m(); + access_s += problem_size.m(); + } + + ElementSoftmaxCompute inv_sum = + cutlass::constants::one() / sum_val; + + access_n = access_n_bak; + access_s = access_s_bak; + + access_n[0] = convert_norm_output(max_val); + access_s[0] = convert_sum_output(inv_sum); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace reduction +} // namespace cutlass diff --git a/csrc/attention/sm70_v37/register.cpp b/csrc/attention/sm70_v37/register.cpp new file mode 100644 index 0000000000..8c6068a528 --- /dev/null +++ b/csrc/attention/sm70_v37/register.cpp @@ -0,0 +1,20 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include +#include + +namespace onecat_v37 { +at::Tensor sm70_d256_gqa_v37_fwd(const at::Tensor& q, const at::Tensor& k, + const at::Tensor& v, at::Tensor& out, + double scale, bool causal); +} + +// Separate symbol: a stale legacy FA2 library must never masquerade as v37. +TORCH_LIBRARY_FRAGMENT(_vllm_fa2_C, ops) { + ops.def( + "sm70_d256_gqa_v37_fwd(Tensor q, Tensor k, Tensor v, Tensor(a!) out, " + "float softmax_scale, bool causal) -> Tensor(a!)"); +} +TORCH_LIBRARY_IMPL(_vllm_fa2_C, CUDA, ops) { + ops.impl("sm70_d256_gqa_v37_fwd", &onecat_v37::sm70_d256_gqa_v37_fwd); +} diff --git a/csrc/attention/sm70_v37/tail.cu b/csrc/attention/sm70_v37/tail.cu new file mode 100644 index 0000000000..f9dadc1a6e --- /dev/null +++ b/csrc/attention/sm70_v37/tail.cu @@ -0,0 +1,854 @@ +#undef FLASH_NAMESPACE +#define FLASH_NAMESPACE onecat_v37_tail +/****************************************************************************** + * Copyright (c) 2026, 1CatAI. + ******************************************************************************/ + +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include "namespace_config.h" +#include "kernel_traits.h" +#include "utils.h" +#include "softmax.h" +#include "mask.h" + +namespace FLASH_NAMESPACE { + +using namespace cute; + +struct Sm70D256SplitDTraits { + using Element = cutlass::half_t; + using MmaAtom = MMA_Atom; + using PvMmaAtom = MMA_Atom; + + static constexpr int kHeadDim = 256; + static constexpr int kBlockM = 64; + static constexpr int kBlockN = 32; + static constexpr int kDChunk = 64; + static constexpr int kDChunks = kHeadDim / kDChunk; + static constexpr int kOwnedDChunks = kDChunks / 2; + static constexpr int kNThreads = 256; + static constexpr int kMmaThreads = 32; + static constexpr int kWarpsPerGroup = 2; + static constexpr int kMmaGroups = kNThreads / (kWarpsPerGroup * kMmaThreads); + static constexpr int kGroupRows = kBlockM / kMmaGroups; + static constexpr int kQkWarpRows = kGroupRows / kWarpsPerGroup; + static constexpr int kQkRowsPerThread = kQkWarpRows / 4; + static constexpr int kOutputRowsPerThread = kGroupRows / 4; + + // Each warp in a pair owns eight distinct Q rows for QK, then the pair + // shares the resulting P tile and each warp owns D/2 for PV. This keeps + // the standard FA2 N32 online-softmax order without duplicating QK work. + using QkTiledMma = TiledMMA>, + Tile, Int, _4>>; + using PvTiledMma = TiledMMA>, + Tile, Int, _4>>; + static_assert(decltype(size(QkTiledMma{}))::value == kMmaThreads); + static_assert(decltype(size(PvTiledMma{}))::value == kMmaThreads); + + using SmemLayoutAtom = decltype(composition( + Swizzle<3, 3, 3>{}, Layout, Stride<_64, _1>>{})); + using SmemLayoutQ = decltype(tile_to_shape( + SmemLayoutAtom{}, Shape, Int>{})); + using SmemLayoutKV = decltype(tile_to_shape( + SmemLayoutAtom{}, Shape, Int>{})); + // Volta HMMA.884 assigns one warp to four 8-thread quadpairs and services + // the 64-bit operand loads as two half-warps. Pitch 68 advances each row + // by one bank pair; the extra 16-half phase every 16 rows folds row bit 4 + // into bank-pair bit 2 so every half-warp covers all 16 bank pairs once. + using SmemLayoutK = Layout< + Shape, Int<2>>, Int>, + Stride, Int<16 * (kDChunk + 4) + 16>>, _1>>; + // TT PV consumes V as KxD. Row bit 1 is folded into D-address bit 2; + // the producer applies the inverse 64-bit-half swap before STS.128. + using SmemLayoutV = + Layout>, Shape<_32, Int>>, + Stride>, Stride<_1, _64>>>; + using SmemLayoutP = + Layout, Int>, Stride, _1>>; + + using SmemCopyAtom = + Copy_Atom, Element>; + using SmemCopyAtomTransposed = SmemCopyAtom; + + static constexpr int kGmemElemsPerLoad = 8; + static constexpr int kGmemThreadsPerRow = kDChunk / kGmemElemsPerLoad; + using GmemLayoutAtom = Layout< + Shape, Int>, + Stride, _1>>; + using GmemTiledCopy = + decltype(make_tiled_copy(Copy_Atom{}, + GmemLayoutAtom{}, Layout>{})); + static constexpr int kGmemRowsPerThread = + kBlockN / (kNThreads / kGmemThreadsPerRow); + using GmemTiledCopyPaged = decltype(make_tiled_copy( + Copy_Atom{}, GmemLayoutAtom{}, + Layout, _8>, Stride<_8, _1>>{})); + + static constexpr int kGmemKElemsPerLoad = 4; + static constexpr int kGmemKThreadsPerRow = kDChunk / kGmemKElemsPerLoad; + using GmemKLayoutAtom = Layout< + Shape, Int>, + Stride, _1>>; + using GmemKTiledCopy = + decltype(make_tiled_copy(Copy_Atom, Element>{}, + GmemKLayoutAtom{}, Layout>{})); + static constexpr int kGmemKRowsPerThread = + kBlockN / (kNThreads / kGmemKThreadsPerRow); + using GmemKTiledCopyPaged = decltype(make_tiled_copy( + Copy_Atom, Element>{}, GmemKLayoutAtom{}, + Layout, _4>, Stride<_4, _1>>{})); + + static constexpr int kQElements = size(SmemLayoutQ{}); + static constexpr int kKVElements = size(SmemLayoutKV{}); + static_assert(size(SmemLayoutV{}) == kKVElements); + // The second K stage is live only before P is materialized, so it may + // alias the beginning of the later P region. Keep the K stages disjoint + // and on the same 128-byte bank phase without moving the V/P regions. + static constexpr int kKStageElements = 2240; + static_assert(kKStageElements >= cosize(SmemLayoutK{})); + static_assert(kKStageElements + cosize(SmemLayoutK{}) <= + 2 * kKVElements + size(SmemLayoutP{})); + static constexpr int kPElements = size(SmemLayoutP{}); + static constexpr int kExchangeRows = kMmaGroups * kGroupRows; + static constexpr int kTensorSmemBytes = + (kQElements + 2 * kKVElements + kPElements) * sizeof(Element); + static constexpr int kExchangeBytes = 2 * kExchangeRows * sizeof(float); + static constexpr int kSmemBytes = kTensorSmemBytes + kExchangeBytes; + static_assert(kSmemBytes == 45568); +}; + +template +__device__ __forceinline__ void copy_even_tile(TiledCopy tiled_copy, + const SrcTensor& src, + DstTensor& dst) { + static_assert(decltype(rank(src))::value == 3); + static_assert(decltype(rank(dst))::value == 3); +#pragma unroll + for (int m = 0; m < size<1>(src); ++m) { +#pragma unroll + for (int k = 0; k < size<2>(src); ++k) { + cute::copy(tiled_copy, src(_, m, k), dst(_, m, k)); + } + } +} + +template +__device__ __forceinline__ void store_v_fragment_128_swizzled( + const RegTensor& source, SmemTensor& destination, + const CoordTensor& coordinates) { + static_assert(decltype(size<0>(source))::value == 8); + static_assert(decltype(size<0>(destination))::value == 8); + static_assert(decltype(size<1>(source))::value == + decltype(size<1>(destination))::value); + static_assert(decltype(size<2>(source))::value == + decltype(size<2>(destination))::value); +#pragma unroll + for (int k = 0; k < size<2>(source); ++k) { +#pragma unroll + for (int m = 0; m < size<1>(source); ++m) { + auto words = recast(source(_, m, k)); + const uint32_t address = static_cast( + __cvta_generic_to_shared(&destination(0, m, k))); + const int row = get<0>(coordinates(0, m, k)); + if (row & 2) { + asm volatile( + "st.shared.v4.u32 [%0], {%1, %2, %3, %4};\n" ::"r"(address), + "r"(words(2)), "r"(words(3)), "r"(words(0)), "r"(words(1))); + } else { + asm volatile( + "st.shared.v4.u32 [%0], {%1, %2, %3, %4};\n" ::"r"(address), + "r"(words(0)), "r"(words(1)), "r"(words(2)), "r"(words(3))); + } + } + } +} + +template +__device__ __forceinline__ void load_v_fragment_tt(const SmemTensor& sV, + TensorB& b_words, int phase, + int lane) { + auto* v = sV.data().get(); + const int d_lane = ((lane & 0x0c) << 1) | ((lane & 0x10) >> 2); + const int k = phase * 4 + (lane & 0x03); + const int offset = + (d_lane ^ ((k & 0x02) << 1)) | ((k & 0x01) << 5) | ((k & 0x1e) << 6); + const uint32_t address = + static_cast(__cvta_generic_to_shared(v + offset)); + uint32_t word0; + uint32_t word1; + uint32_t word2; + uint32_t word3; + asm volatile( + "ld.shared.v2.u32 {%0, %1}, [%4];\n" + "ld.shared.v2.u32 {%2, %3}, [%4+128];\n" + : "=r"(word0), "=r"(word1), "=r"(word2), "=r"(word3) + : "r"(address)); + b_words(0, 0) = word0; + b_words(1, 0) = word1; + b_words(0, 1) = word2; + b_words(1, 1) = word3; +} + +template +__device__ __forceinline__ void splitd_pv_gemm_tt_phase( + TensorO& acc_o, const TensorP& tPrP, const SmemTensor& sV, + TensorB& current_b, TensorBWords& current_b_words, TensorBNext& next_b, + TensorBNextWords& next_b_words, TiledMma tiled_mma, int lane) { + constexpr int kPhases = decltype(size<2>(tPrP))::value; + static_assert(kPhase < kPhases); + if constexpr (kPhase + 1 < kPhases) { + load_v_fragment_tt(sV, next_b_words, kPhase + 1, lane); + } + cute::gemm(tiled_mma, tPrP(_, _, kPhase), current_b, acc_o); + if constexpr (kPhase + 1 < kPhases) { + splitd_pv_gemm_tt_phase(acc_o, tPrP, sV, next_b, next_b_words, + current_b, current_b_words, tiled_mma, + lane); + } +} + +template +__device__ __forceinline__ void splitd_pv_gemm_tt(TensorO& acc_o, + const TensorP& tPrP, + const SmemTensor& sV, + TiledMma tiled_mma, + int lane) { + using Element = typename Sm70D256SplitDTraits::Element; + using BLayout = Layout, Stride<_1, _4>>; + auto b0 = make_tensor(BLayout{}); + auto b1 = make_tensor(BLayout{}); + auto b0_words = recast(b0); + auto b1_words = recast(b1); + static_assert(decltype(size<2>(tPrP))::value == + Sm70D256SplitDTraits::kBlockN / 4); + static_assert(decltype(size(b0))::value == 8); + static_assert(decltype(size<0>(b0_words))::value == 2); + static_assert(decltype(size<1>(b0_words))::value == 2); + load_v_fragment_tt(sV, b0_words, 0, lane); + splitd_pv_gemm_tt_phase<0>(acc_o, tPrP, sV, b0, b0_words, b1, b1_words, + tiled_mma, lane); +} + +template +__device__ __forceinline__ auto reshape_kv_thread_tensor(Tensor tensor) { + if constexpr (PagedKV) { + return make_tensor(tensor.data(), reshape_thread_tile(tensor.layout())); + } else { + return tensor; + } +} + +template +__device__ __forceinline__ void splitd_n32_online_softmax( + TensorScores& acc_s, + float (&o_storage)[Sm70D256SplitDTraits::kOwnedDChunks][kOElements], + OLayout o_layout, float (&row_max)[Sm70D256SplitDTraits::kQkRowsPerThread], + float (&row_sum)[Sm70D256SplitDTraits::kQkRowsPerThread], + float* row_scale_exchange, int mma_group, int n_warp, int lane, + float softmax_scale_log2, bool first_tile) { + using Traits = Sm70D256SplitDTraits; + auto scores = make_tensor( + acc_s.data(), FLASH_NAMESPACE::convert_layout_acc_rowcol(acc_s.layout())); + static_assert(decltype(size<0>(scores))::value == Traits::kQkRowsPerThread); + + float work[Traits::kQkRowsPerThread]; + auto work_tensor = make_tensor(make_rmem_ptr(&work[0]), + Shape>{}); + if (first_tile) { + FLASH_NAMESPACE::sm70_reduce_max(scores, work_tensor); + } else { +#pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { + work[slot] = row_max[slot]; + } + FLASH_NAMESPACE::sm70_reduce_max(scores, work_tensor); + } + + float scores_max[Traits::kQkRowsPerThread]; +#pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { + const float next_max = work[slot]; + const float safe_max = next_max == -INFINITY ? 0.0f : next_max; + work[slot] = first_tile + ? 1.0f + : exp2f((row_max[slot] - safe_max) * softmax_scale_log2); + row_max[slot] = next_max; + scores_max[slot] = safe_max; + if (!first_tile) { + row_sum[slot] *= work[slot]; + } + } + + if ((lane & 0x0e) == 0) { +#pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { + const int row = + FLASH_NAMESPACE::sm70_row_slot(slot, lane); + row_scale_exchange[mma_group * Traits::kGroupRows + + n_warp * Traits::kQkWarpRows + row] = work[slot]; + } + } + __syncthreads(); + + if (!first_tile) { +#pragma unroll + for (int d = 0; d < Traits::kOwnedDChunks; ++d) { + auto acc_o = make_tensor(make_rmem_ptr(&o_storage[d][0]), o_layout); + auto acc_o_rc = make_tensor( + acc_o.data(), + FLASH_NAMESPACE::convert_layout_acc_rowcol(acc_o.layout())); +#pragma unroll + for (int row = 0; row < Traits::kOutputRowsPerThread; ++row) { + const int logical_row = + FLASH_NAMESPACE::sm70_row_slot(row, lane); + const float row_scale = + row_scale_exchange[mma_group * Traits::kGroupRows + logical_row]; +#pragma unroll + for (int col = 0; col < size<1>(acc_o_rc); ++col) { + acc_o_rc(row, col) *= row_scale; + } + } + } + } + + auto scores_max_tensor = make_tensor(make_rmem_ptr(&scores_max[0]), + Shape>{}); + FLASH_NAMESPACE::sm70_scale_apply_exp2(scores, scores_max_tensor, + softmax_scale_log2); + + FLASH_NAMESPACE::sm70_reduce_sum(scores, work_tensor); +#pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { + row_sum[slot] += work[slot]; + } +} + +template +__device__ __forceinline__ int64_t +paged_kv_thread_offset(int tid, int n_block, int d_chunk, int page_size, + const int* __restrict__ block_table, int64_t page_stride, + int64_t row_stride) { + const int row_in_tile = (tid / kThreadsPerRow) * kRowsPerThread; + const int logical_row = n_block * Sm70D256SplitDTraits::kBlockN + row_in_tile; + const int physical_page = block_table[logical_row / page_size]; + return static_cast(physical_page) * page_stride + + static_cast(logical_row % page_size) * row_stride + + d_chunk * Sm70D256SplitDTraits::kDChunk + + (tid % kThreadsPerRow) * kElemsPerLoad; +} + +template +__global__ __launch_bounds__( + Sm70D256SplitDTraits::kNThreads, + 1) void sm70_d256_splitd_dense_kernel(const Element* __restrict__ q, + const Element* __restrict__ k, + const Element* __restrict__ v, + std::conditional_t< + FloatOutput, float, + Element>* __restrict__ out, + int q_batch_stride, int q_row_stride, + int q_head_stride, int k_outer_stride, + int k_row_stride, int k_head_stride, + int v_outer_stride, int v_row_stride, + int v_head_stride, int query_len, + int kv_len, int heads_q, int heads_kv, + float softmax_scale_log2, + const int* __restrict__ block_table, + int page_size, + int block_table_batch_stride, + float* __restrict__ partial_out, + float* __restrict__ partial_max, + float* __restrict__ partial_sum) { + using Traits = Sm70D256SplitDTraits; + constexpr int kBlockM = Traits::kBlockM; + constexpr int kBlockN = Traits::kBlockN; + constexpr int kDChunk = Traits::kDChunk; + + const int tid = threadIdx.x; + const int warp = tid / Traits::kMmaThreads; + const int mma_group = warp / Traits::kWarpsPerGroup; + const int lane = tid % Traits::kMmaThreads; + const int m_block = blockIdx.x; + const int split = SplitKV3 ? blockIdx.y % 3 : 0; + const int batch = SplitKV3 ? blockIdx.y / 3 : blockIdx.y; + const int head_q = blockIdx.z; + const int head_kv = head_q / (heads_q / heads_kv); + const int query_row_base = m_block * kBlockM; + const int kv_offset = kv_len - query_len; + const int* sequence_block_table = + PagedKV ? block_table + batch * block_table_batch_stride : nullptr; + + extern __shared__ __align__(128) Element smem[]; + Element* q_smem_ptr = smem; + Element* kv_smem_ptr = q_smem_ptr + Traits::kQElements; + auto sQ = + make_tensor(make_smem_ptr(q_smem_ptr), typename Traits::SmemLayoutQ{}); + + typename Traits::GmemTiledCopy gmem_copy; + auto gmem_thread = gmem_copy.get_thread_slice(tid); + using GmemVCopy = + std::conditional_t; + GmemVCopy gmem_v_copy; + auto gmem_v_thread = gmem_v_copy.get_thread_slice(tid); + using GmemKCopy = + std::conditional_t; + GmemKCopy gmem_k_copy; + auto gmem_k_thread = gmem_k_copy.get_thread_slice(tid); + { + const int64_t q_batch_offset = static_cast(batch) * q_batch_stride; + auto mQ = + make_tensor(make_gmem_ptr(q + q_batch_offset + head_q * q_head_stride), + make_shape(query_len, Int{}), + make_stride(q_row_stride, _1{})); + auto gQ = local_tile(mQ, Shape, Int>{}, + make_coord(m_block, 0)); + auto tQgQ = gmem_thread.partition_S(gQ); + auto tQsQ = gmem_thread.partition_D(sQ); + copy_even_tile(gmem_copy, tQgQ, tQsQ); + } + __syncthreads(); + + typename Traits::QkTiledMma qk_tiled_mma; + auto qk_mma_thread = qk_tiled_mma.get_thread_slice(lane); + typename Traits::PvTiledMma pv_tiled_mma; + auto pv_mma_thread = pv_tiled_mma.get_thread_slice(lane); + + using OFragment = decltype(partition_fragment_C( + pv_tiled_mma, Shape, Int>{})); + constexpr int kOElements = decltype(size(OFragment{}))::value; + using OLayout = typename OFragment::layout_type; + float o_storage[Traits::kOwnedDChunks][kOElements]; +#pragma unroll + for (int d = 0; d < Traits::kOwnedDChunks; ++d) { +#pragma unroll + for (int i = 0; i < kOElements; ++i) { + o_storage[d][i] = 0.0f; + } + } + + float row_max[Traits::kQkRowsPerThread]; + float row_sum[Traits::kQkRowsPerThread]; +#pragma unroll + for (int row = 0; row < Traits::kQkRowsPerThread; ++row) { + row_max[row] = -INFINITY; + row_sum[row] = 0.0f; + } + + const int max_kv_for_tile = query_row_base + kBlockM + kv_offset; + const int n_block_limit = max_kv_for_tile < kv_len ? max_kv_for_tile : kv_len; + const int visible_n_blocks = cute::ceil_div(n_block_limit, kBlockN); + int n_block_min = 0; + int n_block_max = visible_n_blocks - 1; + if constexpr (SplitKV3) { + n_block_min = visible_n_blocks * split / 3; + n_block_max = visible_n_blocks * (split + 1) / 3 - 1; + } + + const int64_t k_batch_offset = + PagedKV ? 0 : static_cast(batch) * k_outer_stride; + auto mK = + make_tensor(make_gmem_ptr(k + k_batch_offset + head_kv * k_head_stride), + make_shape(kv_len, Int{}), + make_stride(k_row_stride, _1{})); + auto sKFirst = + make_tensor(make_smem_ptr(kv_smem_ptr), typename Traits::SmemLayoutK{}); + auto tKsKFirstRaw = gmem_k_thread.partition_D(sKFirst); + auto tKsKFirst = reshape_kv_thread_tensor(tKsKFirstRaw); + auto tKrKNext = make_fragment_like(tKsKFirst); + auto gKFirst = local_tile(mK, Shape, Int>{}, + make_coord(n_block_max, 0)); + auto tKgKFirstRaw = gmem_k_thread.partition_S(gKFirst); + auto tKgKFirst = reshape_kv_thread_tensor(tKgKFirstRaw); + int64_t k_thread_tile_base = 0; + if constexpr (PagedKV) { + k_thread_tile_base = paged_kv_thread_offset( + tid, n_block_max, 0, page_size, sequence_block_table, k_outer_stride, + k_row_stride); + tKgKFirst.data() = mK.data() + k_thread_tile_base; + } + copy_even_tile(gmem_k_copy, tKgKFirst, tKsKFirst); + __syncthreads(); + + for (int n_block = n_block_max; n_block >= n_block_min; --n_block) { + const int n_warp = warp & 1; + const int group_row_base = mma_group * Traits::kGroupRows; + const int qk_row_base = group_row_base + n_warp * Traits::kQkWarpRows; + auto acc_s = partition_fragment_C( + qk_tiled_mma, Shape, Int>{}); + clear(acc_s); + +#pragma unroll + for (int d_chunk = 0; d_chunk < Traits::kDChunks; ++d_chunk) { + auto sK = make_tensor( + make_smem_ptr(kv_smem_ptr + (d_chunk & 1) * Traits::kKStageElements), + typename Traits::SmemLayoutK{}); + if (d_chunk + 1 < Traits::kDChunks) { + auto gKNext = local_tile(mK, Shape, Int>{}, + make_coord(n_block, d_chunk + 1)); + auto tKgKNextRaw = gmem_k_thread.partition_S(gKNext); + auto tKgKNext = reshape_kv_thread_tensor(tKgKNextRaw); + if constexpr (PagedKV) { + tKgKNext.data() = + mK.data() + k_thread_tile_base + (d_chunk + 1) * kDChunk; + } + copy_even_tile(gmem_k_copy, tKgKNext, tKrKNext); + } + + auto sQChunk = local_tile( + sQ, Shape, Int>{}, + make_coord(mma_group * Traits::kWarpsPerGroup + n_warp, d_chunk)); + auto tSrQ = qk_mma_thread.partition_fragment_A(sQChunk); + auto tSrK = qk_mma_thread.partition_fragment_B(sK); + auto tOsQ = qk_mma_thread.partition_A(sQChunk); + auto tOsK = qk_mma_thread.partition_B(sK); + auto smem_copy_q = + make_tiled_copy_A(typename Traits::SmemCopyAtom{}, qk_tiled_mma); + auto smem_copy_k = + make_tiled_copy_B(typename Traits::SmemCopyAtom{}, qk_tiled_mma); + auto smem_thread_q = smem_copy_q.get_thread_slice(lane); + auto smem_thread_k = smem_copy_k.get_thread_slice(lane); + auto tSsQ = smem_thread_q.retile_S(tOsQ); + auto tSsK = smem_thread_k.retile_S(tOsK); + FLASH_NAMESPACE::gemm( + acc_s, tSrQ, tSrK, tSsQ, tSsK, qk_tiled_mma, smem_copy_q, smem_copy_k, + smem_thread_q, smem_thread_k); + if (d_chunk + 1 < Traits::kDChunks) { + auto sKNext = make_tensor( + make_smem_ptr(kv_smem_ptr + + ((d_chunk + 1) & 1) * Traits::kKStageElements), + typename Traits::SmemLayoutK{}); + auto tKsKNextRaw = gmem_k_thread.partition_D(sKNext); + auto tKsKNext = reshape_kv_thread_tensor(tKsKNextRaw); + cute::copy(tKrKNext, tKsKNext); + __syncthreads(); + } + } + + const int64_t v_batch_offset = + PagedKV ? 0 : static_cast(batch) * v_outer_stride; + auto mV = + make_tensor(make_gmem_ptr(v + v_batch_offset + head_kv * v_head_stride), + make_shape(kv_len, Int{}), + make_stride(v_row_stride, _1{})); + auto sV0 = + make_tensor(make_smem_ptr(kv_smem_ptr), typename Traits::SmemLayoutV{}); + auto sV1 = make_tensor(make_smem_ptr(kv_smem_ptr + Traits::kKVElements), + typename Traits::SmemLayoutV{}); + auto tVsV0Raw = gmem_v_thread.partition_D(sV0); + auto tVsV1Raw = gmem_v_thread.partition_D(sV1); + auto tVsV0 = reshape_kv_thread_tensor(tVsV0Raw); + auto tVsV1 = reshape_kv_thread_tensor(tVsV1Raw); + auto tVrV0 = make_fragment_like(tVsV0); + auto tVrV1 = make_fragment_like(tVsV1); + auto cV = make_identity_tensor(Shape, Int>{}); + auto tVcVRaw = gmem_v_thread.partition_S(cV); + auto tVcV = reshape_kv_thread_tensor(tVcVRaw); + auto gV0 = local_tile(mV, Shape, Int>{}, + make_coord(n_block, 0)); + auto gV2 = local_tile(mV, Shape, Int>{}, + make_coord(n_block, Int{})); + auto tVgV0Raw = gmem_v_thread.partition_S(gV0); + auto tVgV2Raw = gmem_v_thread.partition_S(gV2); + auto tVgV0 = reshape_kv_thread_tensor(tVgV0Raw); + auto tVgV2 = reshape_kv_thread_tensor(tVgV2Raw); + int64_t v_thread_tile_base = 0; + if constexpr (PagedKV) { + v_thread_tile_base = paged_kv_thread_offset( + tid, n_block, 0, page_size, sequence_block_table, v_outer_stride, + v_row_stride); + tVgV0.data() = mV.data() + v_thread_tile_base; + tVgV2.data() = + mV.data() + v_thread_tile_base + Traits::kOwnedDChunks * kDChunk; + } + copy_even_tile(gmem_v_copy, tVgV0, tVrV0); + copy_even_tile(gmem_v_copy, tVgV2, tVrV1); + + FLASH_NAMESPACE::Mask mask(kv_len, query_len, -1, 0, + 0.0f); + mask.template apply_mask(acc_s, n_block * kBlockN, + query_row_base + qk_row_base, 0); + Element* p_smem_ptr = kv_smem_ptr + 2 * Traits::kKVElements; + float* row_scale_exchange = + reinterpret_cast(p_smem_ptr + Traits::kPElements); + splitd_n32_online_softmax(acc_s, o_storage, OLayout{}, row_max, row_sum, + row_scale_exchange, mma_group, n_warp, lane, + softmax_scale_log2, n_block == n_block_max); + + store_v_fragment_128_swizzled(tVrV0, tVsV0, tVcV); + store_v_fragment_128_swizzled(tVrV1, tVsV1, tVcV); + + auto sP = + make_tensor(make_smem_ptr(p_smem_ptr), typename Traits::SmemLayoutP{}); + auto cS = + make_identity_tensor(Shape, Int>{}); + auto tScS = qk_mma_thread.partition_C(cS); +#pragma unroll + for (int i = 0; i < size(acc_s); ++i) { + const int row = get<0>(tScS(i)); + const int col = get<1>(tScS(i)); + sP(qk_row_base + row, col) = Element(acc_s(i)); + } + __syncthreads(); + + auto gV1 = local_tile(mV, Shape, Int>{}, + make_coord(n_block, 1)); + auto gV3 = + local_tile(mV, Shape, Int>{}, + make_coord(n_block, Int{})); + auto tVgV1Raw = gmem_v_thread.partition_S(gV1); + auto tVgV3Raw = gmem_v_thread.partition_S(gV3); + auto tVgV1 = reshape_kv_thread_tensor(tVgV1Raw); + auto tVgV3 = reshape_kv_thread_tensor(tVgV3Raw); + if constexpr (PagedKV) { + tVgV1.data() = mV.data() + v_thread_tile_base + kDChunk; + tVgV3.data() = mV.data() + v_thread_tile_base + + (Traits::kOwnedDChunks + 1) * kDChunk; + } + copy_even_tile(gmem_v_copy, tVgV1, tVrV0); + copy_even_tile(gmem_v_copy, tVgV3, tVrV1); + + auto sPGroup = + local_tile(sP, Shape, Int>{}, + make_coord(mma_group, 0)); + auto tPrP = pv_mma_thread.partition_fragment_A(sPGroup); + auto tOsP = pv_mma_thread.partition_A(sPGroup); + auto smem_copy_p = + make_tiled_copy_A(typename Traits::SmemCopyAtom{}, pv_tiled_mma); + auto smem_thread_p = smem_copy_p.get_thread_slice(lane); + auto tPsP = smem_thread_p.retile_S(tOsP); + auto tPrPView = smem_thread_p.retile_D(tPrP); +#pragma unroll + for (int k_tile = 0; k_tile < size<2>(tPrP); ++k_tile) { + cute::copy(smem_copy_p, tPsP(_, _, k_tile), tPrPView(_, _, k_tile)); + } + +#pragma unroll + for (int d_local = 0; d_local < Traits::kOwnedDChunks; ++d_local) { + auto acc_o = + make_tensor(make_rmem_ptr(&o_storage[d_local][0]), OLayout{}); + auto sV = + make_tensor(make_smem_ptr(kv_smem_ptr + n_warp * Traits::kKVElements), + typename Traits::SmemLayoutV{}); + splitd_pv_gemm_tt(acc_o, tPrP, sV, pv_tiled_mma, lane); + __syncthreads(); + if (d_local + 1 < Traits::kOwnedDChunks) { + store_v_fragment_128_swizzled(tVrV0, tVsV0, tVcV); + store_v_fragment_128_swizzled(tVrV1, tVsV1, tVcV); + __syncthreads(); + if (n_block > n_block_min) { + auto gKNextBlock = local_tile(mK, Shape, Int>{}, + make_coord(n_block - 1, 0)); + auto tKgKNextBlockRaw = gmem_k_thread.partition_S(gKNextBlock); + auto tKgKNextBlock = + reshape_kv_thread_tensor(tKgKNextBlockRaw); + if constexpr (PagedKV) { + k_thread_tile_base = + paged_kv_thread_offset( + tid, n_block - 1, 0, page_size, sequence_block_table, + k_outer_stride, k_row_stride); + tKgKNextBlock.data() = mK.data() + k_thread_tile_base; + } + copy_even_tile(gmem_k_copy, tKgKNextBlock, tKrKNext); + } + } + } + if (n_block > n_block_min) { + auto sKNextBlock = make_tensor(make_smem_ptr(kv_smem_ptr), + typename Traits::SmemLayoutK{}); + auto tKsKNextBlockRaw = gmem_k_thread.partition_D(sKNextBlock); + auto tKsKNextBlock = reshape_kv_thread_tensor(tKsKNextBlockRaw); + cute::copy(tKrKNext, tKsKNextBlock); + __syncthreads(); + } + } + + const int n_warp = warp & 1; + const int group_row_base = mma_group * Traits::kGroupRows; + Element* p_smem_ptr = kv_smem_ptr + 2 * Traits::kKVElements; + float* row_sum_exchange = + reinterpret_cast(p_smem_ptr + Traits::kPElements) + + Traits::kExchangeRows; + SumOp sum_op; +#pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { + row_sum[slot] = + FLASH_NAMESPACE::sm70_row_allreduce_8(row_sum[slot], sum_op); + } + if constexpr (SplitKV3) { + const int64_t split_row_stride = + static_cast(gridDim.y / 3) * query_len * heads_q; + if ((lane & 0x0e) == 0) { +#pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { + const int row = + FLASH_NAMESPACE::sm70_row_slot(slot, lane); + const int query_row = query_row_base + group_row_base + + n_warp * Traits::kQkWarpRows + row; + const int64_t row_offset = + (static_cast(batch) * query_len + query_row) * heads_q + + head_q; + const int64_t partial_row = split * split_row_stride + row_offset; + partial_max[partial_row] = row_max[slot]; + partial_sum[partial_row] = row_sum[slot]; + } + } + +#pragma unroll + for (int d_local = 0; d_local < Traits::kOwnedDChunks; ++d_local) { + auto acc_o = + make_tensor(make_rmem_ptr(&o_storage[d_local][0]), OLayout{}); + auto cO = + make_identity_tensor(Shape, Int>{}); + auto tOcO = pv_mma_thread.partition_C(cO); +#pragma unroll + for (int i = 0; i < size(acc_o); ++i) { + const int row = get<0>(tOcO(i)); + const int col = get<1>(tOcO(i)); + const int query_row = query_row_base + group_row_base + row; + const int64_t row_offset = + (static_cast(batch) * query_len + query_row) * heads_q + + head_q; + const int64_t partial_row = split * split_row_stride + row_offset; + const int d = + (n_warp * Traits::kOwnedDChunks + d_local) * kDChunk + col; + partial_out[partial_row * Traits::kHeadDim + d] = acc_o(i); + } + } + } else { + if constexpr (StoreState) { + if ((lane & 0x0e) == 0) { +#pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { + const int row = + FLASH_NAMESPACE::sm70_row_slot(slot, lane); + const int query_row = query_row_base + group_row_base + + n_warp * Traits::kQkWarpRows + row; + const int64_t state_row = + (static_cast(batch) * query_len + query_row) * heads_q + + head_q; + // The standalone prefix path stores maxima after applying + // softmax_scale. Export the dense-tail state in the same + // natural-log coordinate so the two states can be merged + // directly with exp(prefix_max - global_max). + partial_max[state_row] = + row_max[slot] * (softmax_scale_log2 * float(M_LN2)); + partial_sum[state_row] = row_sum[slot]; + } + } + } + if ((lane & 0x0e) == 0) { +#pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { + const int row = + FLASH_NAMESPACE::sm70_row_slot(slot, lane); + row_sum_exchange[mma_group * Traits::kGroupRows + + n_warp * Traits::kQkWarpRows + row] = row_sum[slot]; + } + } + __syncthreads(); + + const int64_t out_batch_offset = + static_cast(batch) * query_len * heads_q * Traits::kHeadDim; +#pragma unroll + for (int d_local = 0; d_local < Traits::kOwnedDChunks; ++d_local) { + auto acc_o = + make_tensor(make_rmem_ptr(&o_storage[d_local][0]), OLayout{}); + auto acc_o_rc = make_tensor( + acc_o.data(), + FLASH_NAMESPACE::convert_layout_acc_rowcol(acc_o.layout())); +#pragma unroll + for (int row = 0; row < Traits::kOutputRowsPerThread; ++row) { + const int logical_row = + FLASH_NAMESPACE::sm70_row_slot(row, lane); + const float inv_sum = + 1.0f / + row_sum_exchange[mma_group * Traits::kGroupRows + logical_row]; +#pragma unroll + for (int col = 0; col < size<1>(acc_o_rc); ++col) { + if constexpr (!Unnormalized) acc_o_rc(row, col) *= inv_sum; + } + } + + auto cO = + make_identity_tensor(Shape, Int>{}); + auto tOcO = pv_mma_thread.partition_C(cO); +#pragma unroll + for (int i = 0; i < size(acc_o); ++i) { + const int row = get<0>(tOcO(i)); + const int col = get<1>(tOcO(i)); + const int query_row = query_row_base + group_row_base + row; + const int64_t offset = + out_batch_offset + + static_cast(query_row) * heads_q * Traits::kHeadDim + + head_q * Traits::kHeadDim + + (n_warp * Traits::kOwnedDChunks + d_local) * kDChunk + col; + out[offset] = std::conditional_t(acc_o(i)); + } + } + } +} + +extern "C" cudaError_t onecat_v37_dense_state_float_raw( + const void* q, const void* k, const void* v, float* state_max, + float* state_sum, void* out, int query_len, int kv_len, int heads_q, + int heads_kv, float softmax_scale, int unnormalized, cudaStream_t stream) { + if (q == nullptr || k == nullptr || v == nullptr || state_max == nullptr || + state_sum == nullptr || out == nullptr || query_len <= 0 || + kv_len < query_len || heads_q <= 0 || heads_kv <= 0 || + heads_q % heads_kv != 0 || + query_len % Sm70D256SplitDTraits::kBlockM != 0 || + kv_len % Sm70D256SplitDTraits::kBlockN != 0) { + return cudaErrorInvalidValue; + } + const dim3 block(Sm70D256SplitDTraits::kNThreads); + const dim3 grid(query_len / Sm70D256SplitDTraits::kBlockM, 1, heads_q); + auto kernel = unnormalized + ? sm70_d256_splitd_dense_kernel + : sm70_d256_splitd_dense_kernel; + cudaError_t result = + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + Sm70D256SplitDTraits::kSmemBytes); + if (result != cudaSuccess) { + return result; + } + kernel<<>>( + static_cast(q), + static_cast(k), + static_cast(v), static_cast(out), + query_len * heads_q * Sm70D256SplitDTraits::kHeadDim, + heads_q * Sm70D256SplitDTraits::kHeadDim, Sm70D256SplitDTraits::kHeadDim, + kv_len * heads_kv * Sm70D256SplitDTraits::kHeadDim, + heads_kv * Sm70D256SplitDTraits::kHeadDim, Sm70D256SplitDTraits::kHeadDim, + kv_len * heads_kv * Sm70D256SplitDTraits::kHeadDim, + heads_kv * Sm70D256SplitDTraits::kHeadDim, Sm70D256SplitDTraits::kHeadDim, + query_len, kv_len, heads_q, heads_kv, softmax_scale * float(M_LOG2E), + nullptr, 0, 0, nullptr, state_max, state_sum); + return cudaPeekAtLastError(); +} + +} // namespace FLASH_NAMESPACE diff --git a/docs/design/sm70_flash_v37_prefill.md b/docs/design/sm70_flash_v37_prefill.md new file mode 100644 index 0000000000..d924a7d51b --- /dev/null +++ b/docs/design/sm70_flash_v37_prefill.md @@ -0,0 +1,120 @@ +# FlashAttention-V100 v37 prefill integration + +## Scope + +We integrate the v37 FP32-accumulating long-prefill route and an explicit +E4M3-to-FP16 bridge into the SM70 FA2 extension. We retain ordinary decode, +model weights, sampling, SSM arithmetic, and speculative-decoding behavior. +No speculative configuration is used in the model validation below. + +The implementation and admission contract are described in +[the kernel README](../../csrc/attention/sm70_v37/README.md). This change +does not rename the global `fp8` encoding or turn an E5M2 byte cache into E4M3. +The new E4M3 bridge is selected only for explicit `fp8_e4m3` storage. + +## Why this is more than copying the private kernel + +The private endpoint accepted Q8000. The current chunked-prefill scheduler +uses a different family of query lengths, so installing that endpoint alone +does not prove that a model benefits. We replace fixed packed-row strides +with the actual query/head extent and test the admitted alignment family. +The persistent score allocation is bounded by the maximum admitted query +extent. We retain the causal offset when prepending query padding. + +The v37 symbol is separate from the legacy GQA symbol. Missing new native +code triggers a warning and an exact fallback, never a false v37 route hit. +We also join the private tail on exception before releasing its inputs and +temporary state. The legacy architecture remains available for rollback. + +Set `VLLM_FLASH_V100_PREFILL_D256_GQA_V37=0` before starting workers to use +the old architecture loader and disable the new E4M3 bridge. Runtime +environment mutation in an already initialized engine is not a rollback +mechanism. + +## Operator evidence + +The fixed-shape port matches the private v37 output bitwise on 12 complete +real-operand replays: three attention layers and KV lengths 16K, 64K, 128K, +and 256K. After replacing fixed row strides, the layer-63 Q8000 replay still +matches bitwise at all four lengths. Cropped operands are operator tests, +not fresh model executions. + +On one physical V100, 20 warmups and 100 ABBA pairs compare the dynamic port +with the retained private v37 at Q8000/Hq6/Hkv1/D256/FP16/causal: + +| KV tokens | Private v37 median ms | Port median ms | Median paired speedup | +| ---: | ---: | ---: | ---: | +| 128000 | 99.460 | 99.601 | 0.99877 | +| 256000 | 201.507 | 201.743 | 0.99948 | + +These are complete dense attention endpoints, including prefix, causal +tail, state merging and synchronization. They exclude paged gathering and +FP8 conversion. The speedup is the median of paired ratios, not the ratio +of the two reported marginal medians. This is a port-regression comparison, +not a new speedup over an old 18-TFLOP/s or 60-TFLOP/s implementation. + +A separate retained all-row FP64 audit on identical E4M3-representable KV +values found the following relative L2 errors for layer-63 Q8000 operands: + +| KV tokens | Native paged error | Private v37 error | +| ---: | ---: | ---: | +| 64000 | 0.15823% | 0.02788% | +| 128000 | 0.26685% | 0.02679% | +| 256000 | 0.44761% | 0.02581% | + +This isolates attention arithmetic from KV quantization. It does not show +that E4M3 has no quantization loss, nor that all model tokens must match an +unquantized reference. A constant-V result of exactly one on one layer is +also not a universal guarantee: other retained layers differ by one FP16 +rounding step. + +## Model comparison and interpretation + +The initial source-overlay comparison fixes Qwen3.8-27B-FP8, TP4 on four +V100s, FP8 weights, explicit E4M3 KV, FP16 activations/conv state, FP32 SSM +state, no MTP, chunk budget 8192, max length 262144, prefix caching off and +CUDA graphs on. Every worker reports the actual KV page size as 1568. +The control uses the old main E4M3 direct-paged route; the candidate adds +the E4M3 bridge and v37. All other native libraries and settings are shared. + +The initial post-warmup results are single observations per context, not a +publication-grade repeated timing study: + +| Prompt tokens | Control prefill s | Candidate prefill s | Control decode tok/s | Candidate decode tok/s | +| ---: | ---: | ---: | ---: | ---: | +| 8192 | 1.709 | 1.700 | 59.433 | 59.268 | +| 65536 | 76.328 | 16.825 | 37.650 | 37.448 | +| 128000 | 272.040 | 39.775 | 26.688 | 26.508 | +| 256000 | 1049.836 | 106.891 | 16.790 | 16.763 | + +Prefill uses the engine's scheduled-to-first-token interval. Decode excludes +the first token and uses the interval between the first and last generated +tokens. The timing request forces 64 tokens and requests five logprobs; +retrieval quality is scored only before the first EOS. Natural-EOS text +checks are recorded separately. The initial retrieval answers match at all +four lengths, and every rank records 848 v37/E4M3-bridge prefill executions. + +The large model prefill ratios combine two changes: enabling a previously +missing E4M3 bridge and selecting the v37 kernel. They must not be described +as v37-only kernel gains, reused as PR122/E5M2 baselines, or substituted into +the paper's earlier 60-TFLOP/s comparison without a matching contract. +Both initial runs omitted the existing E4M3 B1 long-context auto/wave +partition flags. Their similar decode rates therefore establish neither +historical performance parity nor decode-speed admission. The earlier +PR285 result of 50.376 tok/s at final context 262144 used this long-wave +route, with NVFP4 rather than the FP8 model weights used here. Promotion is +paused while a same-contract FP8 comparison isolates the missing dispatch. + +## Rejected paths and remaining admission work + +An initial variable-shape build retained a fixed 48000-row PV-statistic +stride and failed with an illegal access. The stride is now dynamic. +Q1664/KV4848, which has only 16-token KV alignment, failed the FP64 test with +relative L2 0.151814. We reject that shape before launch and retain the +32-token KV requirement; it is not a passed numerical result. + +The source overlay is not a complete rebuilt vLLM wheel. The final compiled +FA2 target, clean native Flash dependency, stream/memory-safety checks, +short-query latency screening, and natural long-output comparison must be +qualified before promotion. No result above by itself establishes universal +model-quality equivalence or authorizes an unrelated MTP change. diff --git a/tests/kernels/attention/test_sm70_v37_prefill.py b/tests/kernels/attention/test_sm70_v37_prefill.py new file mode 100644 index 0000000000..b903c478c0 --- /dev/null +++ b/tests/kernels/attention/test_sm70_v37_prefill.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FP64 attention and exact E4M3 storage tests for the v37 prefill route.""" + +import pytest +import torch + + +@pytest.fixture +def ops(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("requires SM70") + from vllm.v1.attention.backends.flash_attn_v100 import ( + _get_sm70_d256_gqa_architecture_op, + ) + + _get_sm70_d256_gqa_architecture_op() + ns = torch.ops._vllm_fa2_C + if not hasattr(ns, "sm70_d256_gqa_v37_fwd"): + pytest.skip("rebuild FA2 with v37") + return ns + + +@pytest.mark.parametrize( + "q_len,kv_len", + [ + (64, 96), + (64, 128), + (8192, 8224), + (64, 3136), + (384, 8192), + (1600, 3136), + (1600, 65536), + (1664, 4864), + (1600, 131072), + (1600, 262144), + (8000, 128000), + (8192, 262144), + ], +) +def test_v37_against_fp64(ops, q_len, kv_len): + torch.manual_seed(20260907) + q = torch.randn((1, q_len, 6, 256), dtype=torch.float16, device="cuda") + k = torch.randn((1, kv_len, 1, 256), dtype=torch.float16, device="cuda") + v = torch.randn_like(k) + out = torch.full_like(q, float("nan")) + ops.sm70_d256_gqa_v37_fwd(q, k, v, out, 0.0625, True) + rows = torch.tensor(sorted({0, 1, q_len // 2, q_len - 1}), device="cuda") + scores = q[0, rows].permute(1, 0, 2).double() @ k[0, :, 0].double().T * 0.0625 + scores.masked_fill_( + torch.arange(kv_len, device="cuda")[None, None] + > (kv_len - q_len + rows)[None, :, None], + -torch.inf, + ) + ref = (scores.softmax(-1) @ v[0, :, 0].double()).permute(1, 0, 2) + assert bool(torch.isfinite(out).all()) + relative_l2 = (out[0, rows].double() - ref).norm() / ref.norm() + assert float(relative_l2) < 0.001 + + +def test_v37_leading_query_padding_preserves_causal_offset(ops): + torch.manual_seed(20260907) + length, count, padded = 3136, 1568, 1600 + q = torch.randn((1, padded, 6, 256), dtype=torch.float16, device="cuda") + q[:, : padded - count].zero_() + k = torch.randn((1, length, 1, 256), dtype=torch.float16, device="cuda") + v = torch.randn_like(k) + out = torch.empty_like(q) + ops.sm70_d256_gqa_v37_fwd(q, k, v, out, 0.0625, True) + query = q[0, padded - count].double() + # First real query must see length-count+1 keys, not all of the tail. + ref = (query @ k[0, : length - count + 1, 0].double().T * 0.0625).softmax(-1) + ref = ref @ v[0, : length - count + 1, 0].double() + assert float((out[0, padded - count].double() - ref).norm() / ref.norm()) < 0.001 + + +def test_v37_rejects_misaligned_contiguous_query(ops): + raw = torch.empty(8000 * 6 * 256 + 1, dtype=torch.float16, device="cuda") + q = raw[1:].view(1, 8000, 6, 256) + k = torch.empty((1, 16000, 1, 256), dtype=torch.float16, device="cuda") + with pytest.raises(RuntimeError, match="16-byte aligned"): + ops.sm70_d256_gqa_v37_fwd(q, k, k, torch.empty_like(q), 0.0625, True) + + +def test_v37_rejects_partial_k32_tile(ops): + q = torch.empty((1, 1664, 6, 256), dtype=torch.float16, device="cuda") + k = torch.empty((1, 4848, 1, 256), dtype=torch.float16, device="cuda") + with pytest.raises(RuntimeError, match="32-token KV step"): + ops.sm70_d256_gqa_v37_fwd(q, k, k, torch.empty_like(q), 0.0625, True) + + +def test_v37_shared_workspace_orders_different_streams(ops): + torch.manual_seed(20260907) + inputs = [] + for q_len, kv_len in ((384, 8192), (1600, 3136)): + q = torch.randn((1, q_len, 6, 256), dtype=torch.float16, device="cuda") + k = torch.randn((1, kv_len, 1, 256), dtype=torch.float16, device="cuda") + v = torch.randn_like(k) + ref, out = torch.empty_like(q), torch.empty_like(q) + ops.sm70_d256_gqa_v37_fwd(q, k, v, ref, 0.0625, True) + inputs.append((q, k, v, ref, out)) + torch.accelerator.synchronize() + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + for _ in range(3): + for stream, (q, k, v, ref, out) in zip(streams, inputs): + with torch.cuda.stream(stream): + ops.sm70_d256_gqa_v37_fwd(q, k, v, out, 0.0625, True) + torch.accelerator.synchronize() + for q, k, v, ref, out in inputs: + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +@pytest.mark.parametrize("page", [800, 1568, 1616]) +@pytest.mark.parametrize("scale", [1.0, 0.003, 3.14159]) +def test_e4m3_bridge_all_codes_and_graph(ops, page, scale): + cache = torch.empty((3, 2, page, 1, 256), dtype=torch.uint8, device="cuda") + k, v = cache.unbind(1) + codes = torch.arange(256, device="cuda", dtype=torch.int32).byte() + k.copy_(codes) + v.copy_(codes.flip(0)) + table = torch.tensor([[2, 0, 1]], device="cuda", dtype=torch.int32) + seq = torch.tensor([3 * page - 5], device="cuda", dtype=torch.int32) + output = torch.empty( + ((3 * page + 31) // 32, 2, 32, 1, 256), dtype=torch.float16, device="cuda" + ) + ko, vo = output.unbind(1) + + def call(): + ops.sm70_v37_e4m3_bridge(k, v, table, seq, ko, vo, scale, scale) + + call() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + for length in (3 * page - 5, 0, 1, page + 17): + seq.fill_(length) + output.fill_(float("nan")) + graph.replay() + for src, dst in ((k, ko), (v, vo)): + expected = ( + src[table[0].long()].flatten(0, 1).view(torch.float8_e4m3fn).float() + * scale + ).half() + actual = dst.flatten(0, 1) + torch.testing.assert_close( + actual[:length], expected[:length], rtol=0, atol=0, equal_nan=True + ) + zero = expected[:length] == 0 + assert torch.equal( + torch.signbit(actual[:length])[zero], + torch.signbit(expected[:length])[zero], + ) + padded = (length + 15) // 16 * 16 + assert bool((actual[length:padded] == 0).all()) + assert bool(torch.isnan(actual[padded:]).all()) diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index a3bfe228ba..aec3c3f72d 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -681,6 +681,14 @@ def unexpected_paged(*args, **kwargs): ) monkeypatch.setattr(flash_v100, "_try_sm70_fa2_d256_prefill", exact_dense) monkeypatch.setattr(flash_v100, "_record_route", routes.append) + monkeypatch.setattr( + flash_v100, + "_uniform_cu_seqlens", + lambda *args, **kwargs: ( + torch.tensor([0, query_len], dtype=torch.int32), + torch.tensor([0, seq_len], dtype=torch.int32), + ), + ) impl.flash_attn_prefill_paged = unexpected_paged result = impl._flash_v100_prefill_with_prefix( @@ -776,9 +784,12 @@ def load_library(path: str) -> None: assert loaded == ["/tmp/stable-fa2.so"] -def test_sm70_d256_gqa_architecture_loader_is_optional(monkeypatch): +@pytest.mark.parametrize("v37", [False, True]) +def test_sm70_d256_gqa_architecture_loader_is_optional(monkeypatch, v37): import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_D256_GQA_V37", str(int(v37))) + fake_interface = types.ModuleType("vllm.vllm_flash_attn.flash_attn_interface") fake_package = types.ModuleType("vllm.vllm_flash_attn") fake_package.__dict__["flash_attn_interface"] = fake_interface @@ -793,6 +804,7 @@ def test_sm70_d256_gqa_architecture_loader_is_optional(monkeypatch): fake_ops = SimpleNamespace( _vllm_fa2_C=SimpleNamespace( sm70_d256_gqa_architecture_fwd=architecture, + sm70_d256_gqa_v37_fwd=architecture, ) ) monkeypatch.setattr(flash_v100, "torch", SimpleNamespace(ops=fake_ops)) @@ -910,6 +922,7 @@ def test_prefill_d256_gqa_architecture_policy_is_shape_family_bounded(monkeypatc import vllm.envs as envs import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_D256_GQA_V37", "0") name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" query = torch.empty((1, 8000, 6, 256), dtype=torch.float16, device="meta") key = torch.empty((1, 128000, 1, 256), dtype=torch.float16, device="meta") @@ -2232,6 +2245,7 @@ def test_flash_v100_fp8_prefill_bridge_prefers_logical_dense_exact(monkeypatch): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl impl = object.__new__(FlashAttnV100Impl) + impl.kv_cache_dtype = "fp8_e5m2" impl.scale = 256**-0.5 bridge_calls = [] exact_calls = [] @@ -2309,6 +2323,8 @@ def exact_op(query, key, value, **kwargs): def test_flash_v100_fp8_prefill_bridge_workspace_oom_falls_back(monkeypatch): from vllm.v1.attention.backends import flash_attn_v100 as mod + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + key_cache = torch.zeros((1, 1568, 1, 33), dtype=torch.uint8) mod._fp8_prefill_bridge_workspaces.clear() monkeypatch.setattr( diff --git a/tests/v1/attention/test_sm70_v37_prefill.py b/tests/v1/attention/test_sm70_v37_prefill.py new file mode 100644 index 0000000000..21187d793e --- /dev/null +++ b/tests/v1/attention/test_sm70_v37_prefill.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Version identity and explicit E4M3 prefill admission; no speculative path.""" + +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.mark.parametrize( + "q_len,kv_len,expected", + [ + (64, 3136, True), + (384, 8192, True), + (1600, 3136, True), + (1664, 4848, False), + (8000, 128000, True), + (8192, 262144, True), + (32, 3136, False), + (1568, 3136, False), + (1600, 1600, False), + (8192, 262145, False), + (8192, 131071, False), + (8256, 131072, False), + ], +) +def test_v37_tile_aligned_shape_family(monkeypatch, q_len, kv_len, expected): + from vllm.v1.attention.backends import flash_attn_v100 as mod + + monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_D256_GQA_V37", "1") + monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL", "1") + mod.envs.disable_envs_cache() + q = torch.empty((1, q_len, 6, 256), dtype=torch.float16, device="meta") + k = torch.empty((1, kv_len, 1, 256), dtype=torch.float16, device="meta") + assert ( + mod._should_use_prefill_d256_gqa_architecture( + q, + k, + k, + max_seqlen_q=q_len, + max_seqlen_k=kv_len, + softmax_scale=0.0625, + architecture_op=object(), + ) + is expected + ) + + +def test_missing_v37_does_not_select_legacy(monkeypatch): + from vllm.v1.attention.backends import flash_attn_v100 as mod + + monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_D256_GQA_V37", "1") + mod.envs.disable_envs_cache() + monkeypatch.setattr(mod, "_sm70_d256_gqa_architecture_op_checked", False) + monkeypatch.setattr(mod, "_sm70_d256_gqa_architecture_op", None) + legacy = object() + monkeypatch.setattr(mod, "_get_sm70_splitd_d256_ops", lambda: None) + monkeypatch.setattr( + mod, + "torch", + SimpleNamespace( + ops=SimpleNamespace( + _vllm_fa2_C=SimpleNamespace(sm70_d256_gqa_architecture_fwd=legacy) + ) + ), + ) + assert mod._get_sm70_d256_gqa_architecture_op() is None + + +@pytest.mark.parametrize("dtype", ["fp8_e4m3", "fp8_e5m2"]) +def test_explicit_fp8_bridge_routes(dtype): + from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl + + impl = object.__new__(FlashAttnV100Impl) + impl.use_fp8_prefill_bridge = True + impl.use_flash_v100_prefill_paged = True + impl.kv_cache_dtype = dtype + backing = torch.empty((2, 2, 1616, 1, 256), dtype=torch.uint8) + k, v = backing.unbind(1) + assert impl._should_use_fp8_prefill_bridge( + q_len=8000, + head_dim=256, + key_cache=k, + value_cache=v, + causal=True, + window_size=(-1, -1), + ) + assert not impl._should_use_fp8_prefill_bridge( + q_len=1, + head_dim=256, + key_cache=k, + value_cache=v, + causal=True, + window_size=(-1, -1), + ) + + +def test_e4m3_wide_bridge_rejects_unaligned_cache(): + from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl + + impl = object.__new__(FlashAttnV100Impl) + impl.use_fp8_prefill_bridge = True + impl.use_flash_v100_prefill_paged = True + impl.kv_cache_dtype = "fp8_e4m3" + storage = torch.empty(1616 * 256 + 1, dtype=torch.uint8) + k = storage[1:].view(1, 1616, 1, 256) + assert not impl._should_use_fp8_prefill_bridge( + q_len=8000, + head_dim=256, + key_cache=k, + value_cache=k, + causal=True, + window_size=(-1, -1), + ) diff --git a/vllm/envs.py b/vllm/envs.py index 213c7133a0..e6437c3414 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -413,6 +413,7 @@ VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV: int = 32768 VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL: bool = False VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL: bool = True + VLLM_FLASH_V100_PREFILL_D256_GQA_V37: bool = True VLLM_FLASH_V100_PREFILL_SPLIT_KV: bool = False VLLM_FLASH_V100_PREFILL_SPLIT_KV_TOKENS: int = 32768 VLLM_FLASH_V100_PREFILL_SPLIT_KV_MIN_Q: int = 1 @@ -2878,6 +2879,9 @@ def _resolve_rust_frontend_path() -> str | None: ) ) ), + "VLLM_FLASH_V100_PREFILL_D256_GQA_V37": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_PREFILL_D256_GQA_V37", "1")) + ), "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL": lambda: bool( int( os.getenv( diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index 9736dd0eca..ede90955f2 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -1354,19 +1354,30 @@ def _get_sm70_d256_gqa_architecture_op(): _sm70_d256_gqa_architecture_op_checked = True try: + op_name = ( + "sm70_d256_gqa_v37_fwd" + if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_V37 + else "sm70_d256_gqa_architecture_fwd" + ) # The Split-D loader also resolves an explicit source-overlay # sidecar. Calling it here keeps both operator families on one binary. if not hasattr( torch.ops._vllm_fa2_C, - "sm70_d256_gqa_architecture_fwd", + op_name, ): _get_sm70_splitd_d256_ops() _sm70_d256_gqa_architecture_op = getattr( torch.ops._vllm_fa2_C, - "sm70_d256_gqa_architecture_fwd", + op_name, None, ) + if _sm70_d256_gqa_architecture_op is None: + logger.warning_once( + "Requested SM70 GQA operator %s is absent; rebuild FA2. " + "Using exact dense prefill, not relabelling the old kernel.", + op_name, + ) except (AttributeError, ImportError, RuntimeError) as exc: _sm70_d256_gqa_architecture_op = None if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL: @@ -1379,6 +1390,14 @@ def _get_sm70_d256_gqa_architecture_op(): return _sm70_d256_gqa_architecture_op +def _get_sm70_v37_e4m3_bridge_op(): + """Resolve the format-specific bridge from the same FA2 runtime.""" + if not envs.VLLM_FLASH_V100_PREFILL_D256_GQA_V37: + return None + _get_sm70_splitd_d256_ops() + return getattr(torch.ops._vllm_fa2_C, "sm70_v37_e4m3_bridge", None) + + def _uniform_cu_seqlens( tensor: torch.Tensor, *, @@ -1502,19 +1521,30 @@ def _should_use_prefill_d256_gqa_architecture( softmax_scale: float, architecture_op: Callable[..., torch.Tensor] | None, ) -> bool: - """Gate the stable Q8000/KV16K..256K/Hq6/Hkv1/D256 family.""" + """Use the v37 tile-aligned family or the original rollback shape gate.""" + if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_V37: + shape_allowed = ( + 64 <= max_seqlen_q <= 8192 + and max_seqlen_q % 64 == 0 + and max_seqlen_q < max_seqlen_k <= 262144 + and max_seqlen_k % 32 == 0 + ) + else: + shape_allowed = ( + max_seqlen_q == 8000 + and 16000 <= max_seqlen_k <= 256000 + and max_seqlen_k % 8000 == 0 + ) return ( envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL and architecture_op is not None - and query.shape == (1, 8000, 6, 256) + and shape_allowed + and query.shape == (1, max_seqlen_q, 6, 256) and key.ndim == 4 and key.shape[0] == 1 and key.shape[2:] == (1, 256) and value.shape == key.shape - and max_seqlen_q == 8000 and max_seqlen_k == key.shape[1] - and 16000 <= max_seqlen_k <= 256000 - and max_seqlen_k % 8000 == 0 and query.dtype == torch.float16 and key.dtype == query.dtype and value.dtype == query.dtype @@ -1565,7 +1595,7 @@ def _try_sm70_fa2_d256_prefill( or query.shape[-1] != 256 or key.shape[-1] != 256 or value.shape[-1] != 256 - or max_seqlen_q < 1024 + or max_seqlen_q < 64 or not causal or window_size != (-1, -1) or cu_seqlens_q.device != query.device @@ -1574,6 +1604,19 @@ def _try_sm70_fa2_d256_prefill( ): return None paged_kv = block_table is not None + if max_seqlen_q < 1024: + if paged_kv or not envs.VLLM_FLASH_V100_PREFILL_D256_GQA_V37: + return None + if not _should_use_prefill_d256_gqa_architecture( + query, + key, + value, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + architecture_op=_get_sm70_d256_gqa_architecture_op(), + ): + return None if block_table is not None: if ( seqused_k is None @@ -1683,10 +1726,15 @@ def _try_sm70_fa2_d256_prefill( if not _logged_prefill_d256_gqa_architecture: logger.info( "FLASH_ATTN_V100 SM70 D256 GQA " - "8K-by-16K..256K architecture route active." + "long-prefill architecture route active (%s).", + "v37 FP32" + if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_V37 + else "legacy", ) _logged_prefill_d256_gqa_architecture = True _record_route("prefill_dense_d256_gqa_arch_long") + if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_V37: + _record_route("prefill_dense_d256_gqa_v37") if splitd_result is None and _should_use_prefill_dense_splitkv3( query, key, @@ -1718,7 +1766,7 @@ def _try_sm70_fa2_d256_prefill( ) _logged_prefill_dense_splitkv3 = True _record_route("prefill_dense_splitd_d256_splitkv3_kernel") - if splitd_result is None: + if splitd_result is None and max_seqlen_k % 32 == 0: splitd_result = dense_op( query, key, value, splitd_out, softmax_scale, True ) @@ -4385,6 +4433,11 @@ def __init__(self, *args, **kwargs): _flash_attn_grouped_verify_max_query_tokens ) self.fp8_e5m2_paged_kv_to_fp16 = _get_fp8_e5m2_paged_kv_bridge_op() + self.fp8_e4m3_paged_kv_to_fp16 = ( + _get_sm70_v37_e4m3_bridge_op() + if self.kv_cache_dtype == "fp8_e4m3" + else None + ) # V100 FA2 kernels consume fp16 Q. FP8 KV cache support is implemented # as storage compression only, with K/V dequantized inside FA2 kernels. self.supports_quant_query_input = False @@ -4418,9 +4471,10 @@ def __init__(self, *args, **kwargs): and not paged_prefill_disable ) self.use_fp8_prefill_bridge = ( - self.fp8_e5m2_paged_kv_to_fp16 is not None - and os.getenv("VLLM_FLASH_V100_FP8_PREFILL_BRIDGE", "1") != "0" - ) + self.fp8_e4m3_paged_kv_to_fp16 is not None + if self.kv_cache_dtype == "fp8_e4m3" + else self.fp8_e5m2_paged_kv_to_fp16 is not None + ) and os.getenv("VLLM_FLASH_V100_FP8_PREFILL_BRIDGE", "1") != "0" self.use_flash_v100_prefill_splitkv = ( self.flash_attn_prefill_paged_splitkv is not None and envs.VLLM_FLASH_V100_PREFILL_SPLIT_KV @@ -7348,10 +7402,20 @@ def _should_use_fp8_prefill_bridge( causal: bool, window_size: tuple[int, int], ) -> bool: + # Eight-byte input loads and 16-byte output stores. Keep layouts + # outside the native bridge contract on their existing fallback. + if self.kv_cache_dtype == "fp8_e4m3" and not all( + tensor.ndim == 4 + and tensor.stride(-1) == 1 + and tensor.data_ptr() % 16 == 0 + and all(stride % 8 == 0 for stride in tensor.stride()[:3]) + for tensor in (key_cache, value_cache) + ): + return False return ( self.use_fp8_prefill_bridge and self.use_flash_v100_prefill_paged - and self.kv_cache_dtype == "fp8_e5m2" + and self.kv_cache_dtype in ("fp8_e4m3", "fp8_e5m2") and key_cache.dtype == torch.uint8 and value_cache.dtype == torch.uint8 and key_cache.shape == value_cache.shape @@ -7398,7 +7462,14 @@ def _run_fp8_prefill_bridge( if workspace is None: return None key_out, value_out, output_block_table = workspace - self.fp8_e5m2_paged_kv_to_fp16( + bridge = ( + self.fp8_e4m3_paged_kv_to_fp16 + if self.kv_cache_dtype == "fp8_e4m3" + else self.fp8_e5m2_paged_kv_to_fp16 + ) + if bridge is None: + return None + bridge( key_cache, value_cache, active_block_table, @@ -8184,13 +8255,18 @@ def _flash_v100_prefill_with_prefix( out_seq, out_is_destination = bridge_result if not _logged_fp8_prefill_bridge: logger.info( - "FLASH_ATTN_V100 FP8 E5M2 prefill bridge " + "FLASH_ATTN_V100 %s prefill bridge " "active (one-pass dequant, shared FP16 page-%d " "workspace).", + self.kv_cache_dtype, _FP8_PREFILL_BRIDGE_PAGE_SIZE, ) _logged_fp8_prefill_bridge = True - _record_route("prefill_prefix_fp8_e5m2_bridge") + _record_route( + "prefill_prefix_fp8_e4m3_bridge" + if self.kv_cache_dtype == "fp8_e4m3" + else "prefill_prefix_fp8_e5m2_bridge" + ) else: out_seq = self.flash_attn_prefill_paged( q_seq, From 05f28efcf364e441146e58c6618a1cf294de1ed4 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:17:21 +0800 Subject: [PATCH 2/2] [Doc][SM70] Record v37 runtime qualification and model parity hold Preserve complete timing, numerical and natural-output evidence. Document the explicit E4M3 wave launch contract and retain the unresolved 256K token-parity gate instead of claiming a passed promotion. Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_flash_v37_prefill.md | 97 ++++++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 23 +++++ 2 files changed, 120 insertions(+) diff --git a/docs/design/sm70_flash_v37_prefill.md b/docs/design/sm70_flash_v37_prefill.md index d924a7d51b..9b22ae566b 100644 --- a/docs/design/sm70_flash_v37_prefill.md +++ b/docs/design/sm70_flash_v37_prefill.md @@ -31,6 +31,34 @@ the old architecture loader and disable the new E4M3 bridge. Runtime environment mutation in an already initialized engine is not a rollback mechanism. +### Preserve the E4M3 decode launch contract + +For the measured no-MTP TP4 E4M3/page1568 route, set these existing switches +before starting the engine or any worker: + +```bash +export VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH=1 +export VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO=1 +export VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS=1 +export VLLM_FLASH_V100_XQA_E4M3_G6_MERGED_WAVE_LAUNCH=1 +``` + +Use explicit `kv_cache_dtype="fp8_e4m3"`, `speculative_config=None`, TP4, +`attention_backend="FLASH_ATTN_V100"`, `max_model_len=262144`, +`max_num_batched_tokens=8192`, `max_num_seqs=1` and +`gpu_memory_utilization=0.8`. The measured activation +and convolution state are FP16; the resolved SSM state is FP32. Keep prefix +caching off when reproducing the reported prefill timing. + +The wave switches remain opt-in and obey their existing native shape/layout +gates. This PR does not globally enable them or reinterpret the `fp8` alias. +The `decode_xqa_p64_page1568` counter records a planning hint, not the final +device-selected partition. To verify the native route, additionally set +`VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO_TRACE=1` and check both the native +`merged_long=1, converter=shared-lut` message and the worker's long-context +CUDA Graph dispatch message. Setting a switch without seeing the relevant +execution is not a performance qualification. + ## Operator evidence The fixed-shape port matches the private v37 output bitwise on 12 complete @@ -105,6 +133,75 @@ PR285 result of 50.376 tok/s at final context 262144 used this long-wave route, with NVFP4 rather than the FP8 model weights used here. Promotion is paused while a same-contract FP8 comparison isolates the missing dispatch. +## Final runtime audit: model-parity hold + +PR [548](https://github.com/1CatAI/1Cat-vLLM/pull/548) remains in Draft. +Operator precision, port latency and memory-safety checks pass, but the +strict natural-output token-parity gate does not. Do not describe this as +a completed model-quality promotion. + +The final attention components are the CMake-built FA2 target and a clean +native Flash rebuild, both CUDA12.8/GCC12. Shared core/stable dependencies +remain pinned; this is not a complete newly built wheel. The matched model +uses the explicit launch contract above, Torch2.10.0+cu128, FP8 weights, +deterministic temperature0 and top-5 logprob recording. Timing requests +generate64 tokens; natural-quality requests respect EOS with a512-token cap. + +The first full comparison also changed native dependencies and found a +near-tie wording divergence at reasoning token279. The prompt has72 input +tokens, with no prefix, so that request cannot enter the v37 operator. +Both reasoning answers give the correct9 red/13 blue counts and explain +why6.5 whole balls cannot be moved. This observation is retained, not +relabelled as a passing strict-parity result or proof of degraded semantics. + +A third run holds the native Flash and paged-helper binaries identical and +compares retained FA2/JIT-v37 against the parent-owned FA2/v37 port. This is +the relevant single-library integration comparison: + +| Input tokens | Retained / release prefill s | Retained / release decode tok/s | +| ---: | ---: | ---: | +| 8192 | 1.705 / 1.699 | 58.918 / 59.001 | +| 65536 | 16.811 / 16.808 | 55.623 / 55.493 | +| 128000 | 39.747 / 39.779 | 50.447 / 50.505 | +| 256000 | 106.985 / 107.000 | 43.136 / 43.109 | +| 262080 +64 output | 110.551 / 110.560 | 42.904 / 42.941 | + +These remain single post-warmup model observations, not statistical paired +model measurements. The exact262144 final-context request completes without +non-finite recorded logprobs or a corruption flag. All timing-request +natural prefixes match; forced post-EOS tokens are excluded from quality. + +| Natural-EOS task | Retained / release output tokens | Exact token parity | +| --- | ---: | --- | +| Arithmetic | 4 / 4 | yes | +| Chinese explanation | 72 / 72 | yes | +| Reasoning | 394 / 394 | yes | +| Python function | 129 / 129 | yes | +| 128000-input summary | 180 / 180 | yes | +| 256000-input summary | 157 / 158 | **no** | + +The matched-native256K summary first differs at token109, changing a phrase +equivalent to "this unique phrase" versus "this verification phrase". +Both summaries are coherent and retrieve the required phrase, but they do +not satisfy the strict identity gate. Common top-5 logprobs are not bitwise +equal either; their maximum difference is0.203125 on the matched reasoning +stream. This is not full-vocabulary KL or a perplexity result. + +Further isolation finds bitwise-identical v37 outputs on17 real-derived +dynamic shapes, including Q64/384/1600/8192 with KV16K–256K. Queries beyond +the8000-row capture repeat recorded rows; these are derived operator tests, +not live model activation captures. The eight native D256 dense-prefill +kernel instruction dumps match, and six real no-MTP decode replays have +identical numerical metrics across native builds. These negative findings +do not establish the cause of the model-level divergence. Resolving it, +including possible independent-process variability, remains a merge gate. + +Post-rebase checks:141 CPU tests passed/1 GPU-only skip;33 GPU tests passed +(25 v37 and8 wave-route cases). The earlier CMake/OOM suite passed26 tests; +Compute Sanitizer reported zero errors. The reviewed Python output passes +its three emitted assertions and six additional cases. All model workers +assert no speculative configuration and no drafter. Only GPUs0–3 were used. + ## Rejected paths and remaining admission work An initial variable-shape build retained a fixed 48000-row PV-statistic diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index b4c181af9e..904055d3d6 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -2,6 +2,29 @@ Date: 2026-05-30 +## v37 prefill integration: model-parity hold, 2026-09-07 + +Draft [PR548](https://github.com/1CatAI/1Cat-vLLM/pull/548) integrates the +FP32 v37 prefill route and exact E4M3 bridge, independently of grouped/MTP +PR524. See [the integration audit](sm70_flash_v37_prefill.md) for admission, +rollback and the explicit E4M3 wave-decode launch recipe. Global KV encoding +and decode kernels are unchanged; every model run has MTP disabled. + +The port preserves bitwise operator output and latency, including17 +real-derived dynamic-shape comparisons. CPU141-pass/1-skip, GPU33-pass, +CMake build and memory-safety checks are recorded. In the controlled TP4 +comparison,256000-input decode is43.136 versus43.109tok/s and prefill is +106.985 versus107.000s. This is not a NVFP4-weight50.38tok/s baseline. + +Do not merge yet: with non-FA2 native binaries held identical, five of six +natural-EOS outputs match exactly, but the256K summary differs from token109 +(157 versus158 tokens). Both answers remain coherent and retrieval-correct; +the cause of the identity failure is not established. An earlier +multi-library comparison also differed on a72-input-token reasoning prompt, +which cannot enter v37. Preserve both failed comparisons. Do not repeat +kernel-only checks as a substitute for localizing model-level variation, +or use forced post-EOS continuations as quality evidence. + ## QUASAR E4M3 KV and FP32 logits, 2026-09-06 The [precision follow-up](sm70_quasar_e4m3_fp32_logits.md) adds explicit