diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 87ca109d5..afe4f4866 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -710,6 +710,7 @@ extern "C" { GGML_OP_HC_PRE, GGML_OP_HC_POST, GGML_OP_MASK_TO_IDX, + GGML_OP_LATENT_ATTN, GGML_OP_COUNT, }; @@ -2500,6 +2501,53 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * sinks); + // Latent attention over a packed K/V cache with an independently-visible K/V prefix. + // The value vector of cache row n is cache[dv_off .. dv_off+dv, n] (MLA "absorbed" + // layout: openPangu packs [ckv|k_pe] -> dv_off=0; DeepSeek/GLM_DSA packs [k_pe|ckv] + // rope-first -> dv_off=64). The prefix K/V rows are always visible (bias 0); the mask + // applies to the cache segment only. + // + // q: [Dk, T, H] F32, contiguous + // cache: [Dk, N] F32 / F16: row stride nb[1] arbitrary (windowed views ok); + // Q8_0: rows must be packed (nb[1] == row size) + // prefix_k: [Dk, P] F32 / F16, contiguous; NULL iff P == 0 + // prefix_v: [P, Dv] F32 / F16, contiguous (value-transposed: ne0 = P); NULL iff P == 0 + // mask: [N, >=T] F32, additive, cache segment only; NULL = unmasked cache + // returns: [Dv, T, H] F32, contiguous + // Precondition: every query column must have at least one visible position (a prefix row, + // or a cache row not masked to -inf). Violating it is undefined: a fully-masked column with + // P == 0 yields NaN on CUDA and on the CPU dense path, and zeros on the CPU indexed path, so + // callers must not depend on the value. Inference-only: a gradient-bearing input is rejected. + GGML_API struct ggml_tensor * ggml_latent_attn_prefix_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * cache, + struct ggml_tensor * prefix_k, + struct ggml_tensor * prefix_v, + struct ggml_tensor * mask, + int dv, + int dv_off, + float scale, + float max_bias); + + // Indexed latent attention over absolute cache row ids shared across heads. + // indices: [topk, T] I32, contiguous; mask remains [N, >=T] and is gathered by index. + // Preconditions: each index must lie in [0, N); an out-of-range id aborts on CPU and is + // undefined behavior (an out-of-bounds device read) on CUDA. The visible-position and + // inference-only preconditions above apply here as well. + GGML_API struct ggml_tensor * ggml_latent_attn_indexed_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * cache, + struct ggml_tensor * prefix_k, + struct ggml_tensor * prefix_v, + struct ggml_tensor * mask, + struct ggml_tensor * indices, + int dv, + int dv_off, + float scale, + float max_bias); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 315c7e310..49876eb95 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -912,6 +912,11 @@ GGML_CALL static bool ggml_backend_cpu_supports_op(ggml_backend_t backend, const #else return false; #endif + case GGML_OP_LATENT_ATTN: + // Scalar reference forward exists and is dispatched, so support is truthful. + // Whether to ADOPT the op on a CPU-resident layer is performance policy, and + // that lives in the openPangu builder gate, which requires a non-CPU backend. + return true; default: return true; } diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index f5d604002..015277ada 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -57,6 +57,7 @@ #include "ggml-cuda/tri.cuh" #include "ggml-cuda/delta-net.cuh" #include "ggml-cuda/sinkhorn.cuh" +#include "ggml-cuda/latent_attn.cuh" #include "ggml-cuda/blend.cuh" #include "ggml-cuda/indexer_topk.cuh" @@ -4134,6 +4135,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_SINKHORN: ggml_cuda_op_sinkhorn(ctx, dst); break; + case GGML_OP_LATENT_ATTN: + ggml_cuda_op_latent_attn(ctx, dst); + break; case GGML_OP_HC_PRE: ggml_cuda_op_hc_pre(ctx, dst); break; @@ -4481,7 +4485,7 @@ static bool ggml_graph_node_has_matching_properties(ggml_tensor * node, ggml_gra } } - if (node->op == GGML_OP_SCALE && + if ((node->op == GGML_OP_SCALE || node->op == GGML_OP_LATENT_ATTN) && memcmp(graph_node_properties->op_params, node->op_params, GGML_MAX_OP_PARAMS) != 0) { return false; } @@ -5070,6 +5074,8 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && sink_s >= 1 && sink_s <= 8 && op->src[0]->ne[0] == (int64_t) sink_s*sink_s; } + case GGML_OP_LATENT_ATTN: + return ggml_cuda_latent_attn_is_supported(op); case GGML_OP_FLASH_ATTN_EXT: #if defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__) return (op->src[0]->ne[0] == 64 && op->src[1]->type == GGML_TYPE_F16) || op->src[0]->ne[0] == 128; diff --git a/ggml/src/ggml-cuda/latent_attn.cu b/ggml/src/ggml-cuda/latent_attn.cu new file mode 100644 index 000000000..dad9fb178 --- /dev/null +++ b/ggml/src/ggml-cuda/latent_attn.cu @@ -0,0 +1,698 @@ +#include "common.cuh" +#include "convert.cuh" +#include "latent_attn.cuh" + +#include + +// Latent attention over a packed K/V cache with an independently-visible K/V prefix +// (ggml_latent_attn_prefix_ext, dense mode 0). CUDA path for F32, F16, and Q8_0 caches +// (Q8_0 is dequantized once to a contiguous F16 buffer, then takes the F16 path). +// +// K is shared across every query (MLA: the latent cache does not depend on the head), +// so the whole score matrix is two plain GEMMs, not a per-head batch: +// scores[P+N, QT] = scale * [ prefix_k^T ; cache^T ] @ Q (QT = T*H flattened) +// then a fused mask+softmax down each column, then two value GEMMs: +// out[Dv, QT] = cacheV @ W_cache + prefix_v^T @ W_prefix +// where cacheV is the value slice cache[dv_off .. dv_off+Dv, :] read in place (no +// transpose, no cont). Query columns are tiled to bound the [P+N, cw] score buffer. + +static __device__ float latent_block_reduce_max(float value, float * buf) { + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; + + value = warp_reduce_max(value); + if (blockDim.x > WARP_SIZE) { + __syncthreads(); + if (warp == 0) { + buf[lane] = -INFINITY; + } + __syncthreads(); + if (lane == 0) { + buf[warp] = value; + } + __syncthreads(); + value = warp_reduce_max(buf[lane]); + } + return value; +} + +static __device__ float latent_block_reduce_sum(float value, float * buf) { + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; + + value = warp_reduce_sum(value); + if (blockDim.x > WARP_SIZE) { + __syncthreads(); + if (warp == 0) { + buf[lane] = 0.0f; + } + __syncthreads(); + if (lane == 0) { + buf[warp] = value; + } + __syncthreads(); + value = warp_reduce_sum(buf[lane]); + } + return value; +} + +// Fused mask-add + column softmax. One block per score column. scores is column-major +// [PN, cw] (a column's PN logits are contiguous). The cache segment (rows >= P) gets the +// additive mask for that column's token t; the prefix segment (rows < P) is always visible. +// Writes normalized weights into wout (float or half) with the same layout. For the half +// path, restore[col] reverses the power-of-two query scaling before mask and softmax. +template +static __global__ void k_latent_mask_softmax( + const float * __restrict__ scores, Tout * __restrict__ wout, + const float * __restrict__ restore, + const float * __restrict__ mask, int PN, int P, int T, int c0, + int64_t mask_nb1_f /* row stride of mask in floats */) { + const int col = blockIdx.x; + const int tid = threadIdx.x; + const int nth = blockDim.x; + const float * s = scores + (size_t) col*PN; + Tout * w = wout + (size_t) col*PN; + const float r = Restore ? restore[col] : 1.0f; + const bool do_restore = Restore && r != 1.0f; + const int t = (c0 + col) % T; + const float * mrow = mask ? mask + (size_t) t*mask_nb1_f : nullptr; + + extern __shared__ float shbuf[]; // WARP_SIZE floats + + // pass 1: max of (logit + mask) + float local_max = -INFINITY; + for (int k = tid; k < PN; k += nth) { + float v = s[k]; + if (do_restore) v *= r; + if (mrow && k >= P) v += mrow[k - P]; + local_max = fmaxf(local_max, v); + } + const float mx = latent_block_reduce_max(local_max, shbuf); + + // pass 2: sum of exp + float local_sum = 0.0f; + for (int k = tid; k < PN; k += nth) { + float v = s[k]; + if (do_restore) v *= r; + if (mrow && k >= P) v += mrow[k - P]; + local_sum += expf(v - mx); + } + const float sum = latent_block_reduce_sum(local_sum, shbuf); + const float inv = sum > 0.0f ? 1.0f/sum : 0.0f; + + // pass 3: normalized weights + for (int k = tid; k < PN; k += nth) { + float v = s[k]; + if (do_restore) v *= r; + if (mrow && k >= P) v += mrow[k - P]; + w[k] = (Tout) (expf(v - mx) * inv); + } +} + +// Mode 1 packs query columns token-major ([Dk, H*rows]) so a gathered cache tile for +// one token can be paired with all H heads in one strided-batched GEMM. +template +static __global__ void k_latent_pack_q_indexed( + const float * __restrict__ q, Tout * __restrict__ qout, + int Dk, int H, int first, int rows, int64_t q_nb1_f, int64_t q_nb2_f) { + const int col = blockIdx.x; + if (col >= rows*H) { + return; + } + const int tr = col/H; + const int h = col%H; + const float * qrow = q + (size_t) (first + tr)*q_nb1_f + (size_t) h*q_nb2_f; + Tout * dst = qout + (size_t) col*Dk; + for (int d = threadIdx.x; d < Dk; d += blockDim.x) { + dst[d] = (Tout) qrow[d]; + } +} + +// Range-safe F32 -> F16 query pack. Each query column is divided by the smallest +// power of two that puts its largest finite magnitude in the F16 domain. The score +// softmax kernel multiplies the resulting logits by restore[col] before adding masks. +static __global__ void k_latent_pack_q_scaled( + const float * __restrict__ q, half * __restrict__ qout, + float * __restrict__ restore, int Dk, int H, int first, int rows, + int64_t q_nb1_f, int64_t q_nb2_f, bool indexed) { + const int col = blockIdx.x; + const int tid = threadIdx.x; + if (col >= rows*H) { + return; + } + + const int tr = indexed ? col/H : col; + const int h = indexed ? col%H : 0; + const float * qrow = q + (size_t) (first + tr)*q_nb1_f + (size_t) h*q_nb2_f; + + // Dk <= 4*nth covers the in-tree MLA shapes (OpenPangu/GLM use Dk=576). + // Keeping those values in registers avoids rereading q after the reduction; + // larger generic shapes retain the bounded two-read fallback. + constexpr int kValsPerThread = 4; + const bool cache_values = Dk <= blockDim.x*kValsPerThread; + float qvals[kValsPerThread]; + int nvals = 0; + float amax = 0.0f; + for (int d = tid; d < Dk; d += blockDim.x) { + const float v = qrow[d]; + if (cache_values) qvals[nvals++] = v; + amax = fmaxf(amax, fabsf(v)); + } + + extern __shared__ float shbuf[]; + amax = latent_block_reduce_max(amax, shbuf); + __syncthreads(); // reduce slots must be quiesced before shbuf[0] is reused below + if (tid == 0) { + float r = 1.0f; + if (isfinite(amax)) { + while (amax/r > 65504.0f) { + r *= 2.0f; + } + } + restore[col] = r; + shbuf[0] = r; + } + __syncthreads(); + + const float r = shbuf[0]; + half * dst = qout + (size_t) col*Dk; + nvals = 0; + for (int d = tid; d < Dk; d += blockDim.x) { + const float v = cache_values ? qvals[nvals++] : qrow[d]; + dst[d] = (half) (v/r); + } +} + +// F32/F16 row gather. Each selected cache row is copied whole; values are later read +// from the block-aligned [dv_off, dv_off+dv) slice of this full-row tile. +template +static __global__ void k_latent_gather_rows_indexed( + const char * __restrict__ cache, int64_t cache_nb1, + const int32_t * __restrict__ indices, int64_t indices_nb1_i, + T * __restrict__ gathered, int Dk, int topk, int first, int rows) { + const int row = blockIdx.x; + if (row >= rows*topk) { + return; + } + const int tr = row/topk; + const int k = row%topk; + const int32_t idx = indices[(size_t) (first + tr)*indices_nb1_i + k]; + const T * src = (const T *) (cache + (size_t) idx*cache_nb1); + T * dst = gathered + (size_t) row*Dk; + for (int d = threadIdx.x; d < Dk; d += blockDim.x) { + dst[d] = src[d]; + } +} + +// Q8_0 is decoded only for selected rows. The complete packed row is the source and the +// complete dequantized row is the destination; no narrowed quantized view is formed. +static __global__ void k_latent_gather_rows_q8_0_indexed( + const char * __restrict__ cache, int64_t cache_nb1, + const int32_t * __restrict__ indices, int64_t indices_nb1_i, + half * __restrict__ gathered, int Dk, int topk, int first, int rows) { + const int row = blockIdx.x; + if (row >= rows*topk) { + return; + } + const int tr = row/topk; + const int k = row%topk; + const int32_t idx = indices[(size_t) (first + tr)*indices_nb1_i + k]; + const block_q8_0 * src = (const block_q8_0 *) (cache + (size_t) idx*cache_nb1); + half * dst = gathered + (size_t) row*Dk; + for (int d = threadIdx.x; d < Dk; d += blockDim.x) { + const block_q8_0 & block = src[d/QK8_0]; + dst[d] = (half) ((float) block.d * block.qs[d%QK8_0]); + } +} + +// Indexed fused mask-add + softmax. Score columns are token-major, so col/H selects +// the token and col%H selects the head. Only gathered cache logits consult the mask. +template +static __global__ void k_latent_mask_softmax_indexed( + const float * __restrict__ scores, Tout * __restrict__ wout, + const float * __restrict__ restore, + const float * __restrict__ mask, const int32_t * __restrict__ indices, + int M, int P, int H, int first, int64_t mask_nb1_f, int64_t indices_nb1_i) { + const int col = blockIdx.x; + const int tid = threadIdx.x; + const int nth = blockDim.x; + const float * s = scores + (size_t) col*M; + Tout * w = wout + (size_t) col*M; + const float r = Restore ? restore[col] : 1.0f; + const bool do_restore = Restore && r != 1.0f; + const int t = first + col/H; + const float * mrow = mask ? mask + (size_t) t*mask_nb1_f : nullptr; + const int32_t * idx = indices + (size_t) t*indices_nb1_i; + + extern __shared__ float shbuf[]; + + float local_max = -INFINITY; + for (int k = tid; k < M; k += nth) { + float v = s[k]; + if (do_restore) v *= r; + if (mrow && k >= P) v += mrow[idx[k - P]]; + local_max = fmaxf(local_max, v); + } + const float mx = latent_block_reduce_max(local_max, shbuf); + + float local_sum = 0.0f; + for (int k = tid; k < M; k += nth) { + float v = s[k]; + if (do_restore) v *= r; + if (mrow && k >= P) v += mrow[idx[k - P]]; + local_sum += expf(v - mx); + } + const float sum = latent_block_reduce_sum(local_sum, shbuf); + const float inv = sum > 0.0f ? 1.0f/sum : 0.0f; + + for (int k = tid; k < M; k += nth) { + float v = s[k]; + if (do_restore) v *= r; + if (mrow && k >= P) v += mrow[idx[k - P]]; + w[k] = (Tout) (expf(v - mx) * inv); + } +} + +// Convert token-major tile output [Dv, H*rows] back to ggml's [Dv, T, H] strides. +static __global__ void k_latent_copy_out_indexed( + const float * __restrict__ src, float * __restrict__ dst, + int Dv, int H, int first, int rows, int64_t dst_nb1_f, int64_t dst_nb2_f) { + const int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + const int64_t ne = (int64_t) Dv*H*rows; + if (i >= ne) { + return; + } + const int d = i%Dv; + const int col = i/Dv; + const int tr = col/H; + const int h = col%H; + dst[(size_t) (first + tr)*dst_nb1_f + (size_t) h*dst_nb2_f + d] = src[i]; +} + +// ---- host-side small helpers ------------------------------------------------------- + +bool ggml_cuda_latent_attn_is_supported(const ggml_tensor * op) { + const ggml_tensor * cache = op->src[1]; + if (cache == nullptr) { + return false; + } + + const int mode = op->op_params[4]; + if (cache->type != GGML_TYPE_F32 && cache->type != GGML_TYPE_F16 && cache->type != GGML_TYPE_Q8_0) { + return false; + } + + // Both dense (cuBLAS) and indexed (row-gather) readers index each cache row as a packed + // element array. The dense path additionally requires an int-sized cuBLAS leading dimension. + if (cache->type == GGML_TYPE_F32 || cache->type == GGML_TYPE_F16) { + const size_t element_size = ggml_type_size(cache->type); + if (cache->nb[0] != element_size) { + return false; + } + if (mode == 0) { + if (cache->nb[1] % element_size != 0) { + return false; + } + const size_t lda = cache->nb[1] / element_size; + if (lda < (size_t) cache->ne[0] || lda > (size_t) INT_MAX) { + return false; + } + } + } + if (ggml_is_quantized(cache->type) && + cache->nb[1] != ggml_row_size(cache->type, cache->ne[0])) { + return false; + } + return true; +} + +// f32 GEMM path: C[m,n] = alpha * op(A) @ op(B) + beta * C, all float. +static void sgemm(ggml_backend_cuda_context & ctx, cublasOperation_t ta, cublasOperation_t tb, + int m, int n, int k, float alpha, const float * A, int lda, + const float * B, int ldb, float beta, float * C, int ldc) { + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(ctx.device), ctx.stream())); + CUBLAS_CHECK(cublasSgemm(ctx.cublas_handle(ctx.device), ta, tb, m, n, k, + &alpha, A, lda, B, ldb, &beta, C, ldc)); +} + +// f16-input GEMM with f32 accumulate into an f32 C. +static void hgemm_f32acc(ggml_backend_cuda_context & ctx, cublasOperation_t ta, cublasOperation_t tb, + int m, int n, int k, float alpha, const half * A, int lda, + const half * B, int ldb, float beta, float * C, int ldc) { + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(ctx.device), ctx.stream())); + CUBLAS_CHECK(cublasGemmEx(ctx.cublas_handle(ctx.device), ta, tb, m, n, k, + &alpha, A, CUDA_R_16F, lda, B, CUDA_R_16F, ldb, + &beta, C, CUDA_R_32F, ldc, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +} + +static void sgemm_strided_batched(ggml_backend_cuda_context & ctx, cublasOperation_t ta, cublasOperation_t tb, + int m, int n, int k, float alpha, const float * A, int lda, int64_t stride_a, + const float * B, int ldb, int64_t stride_b, float beta, float * C, int ldc, + int64_t stride_c, int batch_count) { + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(ctx.device), ctx.stream())); + CUBLAS_CHECK(cublasSgemmStridedBatched(ctx.cublas_handle(ctx.device), ta, tb, m, n, k, + &alpha, A, lda, stride_a, B, ldb, stride_b, &beta, C, ldc, stride_c, batch_count)); +} + +static void hgemm_f32acc_strided_batched( + ggml_backend_cuda_context & ctx, cublasOperation_t ta, cublasOperation_t tb, + int m, int n, int k, float alpha, const half * A, int lda, int64_t stride_a, + const half * B, int ldb, int64_t stride_b, float beta, float * C, int ldc, + int64_t stride_c, int batch_count) { + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(ctx.device), ctx.stream())); + CUBLAS_CHECK(cublasGemmStridedBatchedEx(ctx.cublas_handle(ctx.device), ta, tb, m, n, k, + &alpha, A, CUDA_R_16F, lda, stride_a, B, CUDA_R_16F, ldb, stride_b, + &beta, C, CUDA_R_32F, ldc, stride_c, batch_count, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +} + +static void ggml_cuda_op_latent_attn_indexed(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * cache = dst->src[1]; + const ggml_tensor * pk = dst->src[2]; + const ggml_tensor * pv = dst->src[3]; + const ggml_tensor * mask = dst->src[4]; + const ggml_tensor * indices = dst->src[5]; + + const int Dk = q->ne[0]; + const int T = q->ne[1]; + const int H = q->ne[2]; + const int P = pk ? pk->ne[1] : 0; + const int topk = indices->ne[0]; + const int M = P + topk; + + float scale; + memcpy(&scale, &dst->op_params[0], sizeof(float)); + const int dv = dst->op_params[2]; + const int dv_off = dst->op_params[3]; + + const bool f16 = cache->type != GGML_TYPE_F32; + const int64_t q_nb1_f = q->nb[1]/sizeof(float); + const int64_t q_nb2_f = q->nb[2]/sizeof(float); + const int64_t mask_nb1_f = mask ? mask->nb[1]/sizeof(float) : 0; + const int64_t indices_nb1_i = indices->nb[1]/sizeof(int32_t); + const int64_t dst_nb1_f = dst->nb[1]/sizeof(float); + const int64_t dst_nb2_f = dst->nb[2]/sizeof(float); + cudaStream_t stream = ctx.stream(); + + // Score scratch is [P+topk, H*rows] f32. Cap it at 16M floats and retain the + // DSA max_rows=32 knob: rows=min(T, 32, floor(16M/((P+topk)*H))), clamped to 1. + constexpr int kMaxRows = 32; + constexpr int64_t kMaxScoreElems = 16*1024*1024; + const int64_t score_elems_per_row = (int64_t) M*H; + int max_rows = std::max(1, kMaxScoreElems/score_elems_per_row); + max_rows = std::min(max_rows, kMaxRows); + max_rows = std::min(max_rows, T); + + // Prefix tensors are converted once to the cache compute type. F32 cache stays on + // SGEMM; F16 and gathered Q8_0 use half inputs with f32 accumulation. + ggml_cuda_pool_alloc pk_f32(ctx.pool()), pv_f32(ctx.pool()); + ggml_cuda_pool_alloc pk_f16(ctx.pool()), pv_f16(ctx.pool()); + const float * pk_f = nullptr; + const float * pv_f = nullptr; + const half * pk_h = nullptr; + const half * pv_h = nullptr; + if (P > 0) { + if (!f16) { + if (pk->type == GGML_TYPE_F32) { + pk_f = (const float *) pk->data; + } else { + pk_f32.alloc((int64_t) Dk*P); + ggml_get_to_fp32_cuda(pk->type)(pk->data, pk_f32.get(), (int64_t) Dk*P, 1, stream); + pk_f = pk_f32.get(); + } + if (pv->type == GGML_TYPE_F32) { + pv_f = (const float *) pv->data; + } else { + pv_f32.alloc((int64_t) P*dv); + ggml_get_to_fp32_cuda(pv->type)(pv->data, pv_f32.get(), (int64_t) P*dv, 1, stream); + pv_f = pv_f32.get(); + } + } else { + if (pk->type == GGML_TYPE_F16) { + pk_h = (const half *) pk->data; + } else { + pk_f16.alloc((int64_t) Dk*P); + ggml_get_to_fp16_cuda(pk->type)(pk->data, pk_f16.get(), (int64_t) Dk*P, 1, stream); + pk_h = pk_f16.get(); + } + if (pv->type == GGML_TYPE_F16) { + pv_h = (const half *) pv->data; + } else { + pv_f16.alloc((int64_t) P*dv); + ggml_get_to_fp16_cuda(pv->type)(pv->data, pv_f16.get(), (int64_t) P*dv, 1, stream); + pv_h = pv_f16.get(); + } + } + } + + const int64_t tile_cols = (int64_t) H*max_rows; + ggml_cuda_pool_alloc scores(ctx.pool(), (int64_t) M*tile_cols); + ggml_cuda_pool_alloc outbuf(ctx.pool(), (int64_t) dv*tile_cols); + ggml_cuda_pool_alloc q_f32(ctx.pool()), gathered_f32(ctx.pool()), w_f32(ctx.pool()); + ggml_cuda_pool_alloc q_f16(ctx.pool()), gathered_f16(ctx.pool()), w_f16(ctx.pool()); + ggml_cuda_pool_alloc q_restore(ctx.pool()); + if (!f16) { + q_f32.alloc((int64_t) Dk*tile_cols); + gathered_f32.alloc((int64_t) Dk*topk*max_rows); + w_f32.alloc((int64_t) M*tile_cols); + } else { + q_f16.alloc((int64_t) Dk*tile_cols); + gathered_f16.alloc((int64_t) Dk*topk*max_rows); + w_f16.alloc((int64_t) M*tile_cols); + q_restore.alloc(tile_cols); + } + + constexpr int nth = 256; + const size_t softmax_shmem = WARP_SIZE*sizeof(float); + for (int first = 0; first < T; first += max_rows) { + const int rows = std::min(max_rows, T - first); + const int cols = rows*H; + + if (!f16) { + k_latent_pack_q_indexed<<>>( + (const float *) q->data, q_f32.get(), Dk, H, first, rows, q_nb1_f, q_nb2_f); + k_latent_gather_rows_indexed<<>>( + (const char *) cache->data, cache->nb[1], (const int32_t *) indices->data, + indices_nb1_i, gathered_f32.get(), Dk, topk, first, rows); + CUDA_CHECK(cudaGetLastError()); + + if (P > 0) { + sgemm(ctx, CUBLAS_OP_T, CUBLAS_OP_N, P, cols, Dk, scale, + pk_f, Dk, q_f32.get(), Dk, 0.0f, scores.get(), M); + } + sgemm_strided_batched(ctx, CUBLAS_OP_T, CUBLAS_OP_N, topk, H, Dk, scale, + gathered_f32.get(), Dk, (int64_t) Dk*topk, + q_f32.get(), Dk, (int64_t) Dk*H, 0.0f, scores.get() + P, M, + (int64_t) M*H, rows); + } else { + k_latent_pack_q_scaled<<>>( + (const float *) q->data, q_f16.get(), q_restore.get(), Dk, H, first, rows, + q_nb1_f, q_nb2_f, true); + if (cache->type == GGML_TYPE_F16) { + k_latent_gather_rows_indexed<<>>( + (const char *) cache->data, cache->nb[1], (const int32_t *) indices->data, + indices_nb1_i, gathered_f16.get(), Dk, topk, first, rows); + } else { + k_latent_gather_rows_q8_0_indexed<<>>( + (const char *) cache->data, cache->nb[1], (const int32_t *) indices->data, + indices_nb1_i, gathered_f16.get(), Dk, topk, first, rows); + } + CUDA_CHECK(cudaGetLastError()); + + if (P > 0) { + hgemm_f32acc(ctx, CUBLAS_OP_T, CUBLAS_OP_N, P, cols, Dk, scale, + pk_h, Dk, q_f16.get(), Dk, 0.0f, scores.get(), M); + } + hgemm_f32acc_strided_batched(ctx, CUBLAS_OP_T, CUBLAS_OP_N, topk, H, Dk, scale, + gathered_f16.get(), Dk, (int64_t) Dk*topk, + q_f16.get(), Dk, (int64_t) Dk*H, 0.0f, scores.get() + P, M, + (int64_t) M*H, rows); + } + + if (!f16) { + k_latent_mask_softmax_indexed<<>>( + scores.get(), w_f32.get(), nullptr, mask ? (const float *) mask->data : nullptr, + (const int32_t *) indices->data, M, P, H, first, mask_nb1_f, indices_nb1_i); + } else { + k_latent_mask_softmax_indexed<<>>( + scores.get(), w_f16.get(), q_restore.get(), mask ? (const float *) mask->data : nullptr, + (const int32_t *) indices->data, M, P, H, first, mask_nb1_f, indices_nb1_i); + } + CUDA_CHECK(cudaGetLastError()); + + if (!f16) { + sgemm_strided_batched(ctx, CUBLAS_OP_N, CUBLAS_OP_N, dv, H, topk, 1.0f, + gathered_f32.get() + dv_off, Dk, (int64_t) Dk*topk, + w_f32.get() + P, M, (int64_t) M*H, 0.0f, outbuf.get(), dv, + (int64_t) dv*H, rows); + if (P > 0) { + sgemm(ctx, CUBLAS_OP_T, CUBLAS_OP_N, dv, cols, P, 1.0f, + pv_f, P, w_f32.get(), M, 1.0f, outbuf.get(), dv); + } + } else { + hgemm_f32acc_strided_batched(ctx, CUBLAS_OP_N, CUBLAS_OP_N, dv, H, topk, 1.0f, + gathered_f16.get() + dv_off, Dk, (int64_t) Dk*topk, + w_f16.get() + P, M, (int64_t) M*H, 0.0f, outbuf.get(), dv, + (int64_t) dv*H, rows); + if (P > 0) { + hgemm_f32acc(ctx, CUBLAS_OP_T, CUBLAS_OP_N, dv, cols, P, 1.0f, + pv_h, P, w_f16.get(), M, 1.0f, outbuf.get(), dv); + } + } + + const int64_t ne = (int64_t) dv*cols; + k_latent_copy_out_indexed<<<(ne + nth - 1)/nth, nth, 0, stream>>>( + outbuf.get(), (float *) dst->data, dv, H, first, rows, dst_nb1_f, dst_nb2_f); + CUDA_CHECK(cudaGetLastError()); + } +} + +void ggml_cuda_op_latent_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + GGML_ASSERT(ggml_cuda_latent_attn_is_supported(dst)); + if (dst->op_params[4] == 1) { + ggml_cuda_op_latent_attn_indexed(ctx, dst); + return; + } + + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * cache = dst->src[1]; + const ggml_tensor * pk = dst->src[2]; + const ggml_tensor * pv = dst->src[3]; + const ggml_tensor * mask = dst->src[4]; + + const int64_t Dk = q->ne[0]; + const int64_t T = q->ne[1]; + const int64_t H = q->ne[2]; + const int64_t N = cache->ne[1]; + const int64_t P = pk ? pk->ne[1] : 0; + const int64_t QT = T*H; + const int64_t PN = P + N; + + float scale; + memcpy(&scale, &dst->op_params[0], sizeof(float)); + const int dv = dst->op_params[2]; + const int dv_off = dst->op_params[3]; + + // Quantized caches dequantize to a contiguous F16 buffer once, then take the F16 path. + const bool cache_quant = ggml_is_quantized(cache->type); + const bool f16 = cache_quant || cache->type == GGML_TYPE_F16; + + const int64_t mask_nb1_f = mask ? mask->nb[1] / sizeof(float) : 0; + + cudaStream_t stream = ctx.stream(); + + // resolve the cache into a compute-type pointer + element stride + ggml_cuda_pool_alloc cache_deq(ctx.pool()); + const half * cache_h = nullptr; + const float * cache_f = nullptr; + int64_t cache_stride = 0; + if (cache->type == GGML_TYPE_F32) { + cache_f = (const float *) cache->data; + cache_stride = cache->nb[1] / sizeof(float); + } else if (cache->type == GGML_TYPE_F16) { + cache_h = (const half *) cache->data; + cache_stride = cache->nb[1] / sizeof(half); + } else { + cache_deq.alloc(Dk*N); + ggml_get_to_fp16_cuda(cache->type)(cache->data, cache_deq.get(), N, Dk, stream); + CUDA_CHECK(cudaGetLastError()); + cache_h = cache_deq.get(); + cache_stride = Dk; + } + + // column tile: bound the [PN, cw] score buffer to ~16M floats + const int64_t kMaxScoreElems = 16*1024*1024; + int64_t cw = PN > 0 ? kMaxScoreElems / PN : QT; + if (cw < 1) cw = 1; + if (cw > QT) cw = QT; + + // ---- prefix K/V converted once to the compute type (if needed) ---- + ggml_cuda_pool_alloc pk_f32(ctx.pool()), pv_f32(ctx.pool()); + ggml_cuda_pool_alloc pk_f16(ctx.pool()), pv_f16(ctx.pool()); + const float * pk_f = nullptr; const float * pv_f = nullptr; + const half * pk_h = nullptr; const half * pv_h = nullptr; + if (P > 0) { + if (!f16) { + if (pk->type == GGML_TYPE_F32) { pk_f = (const float *) pk->data; } + else { pk_f32.alloc(Dk*P); ggml_get_to_fp32_cuda(pk->type)(pk->data, pk_f32.get(), Dk*P, 1, stream); pk_f = pk_f32.get(); } + if (pv->type == GGML_TYPE_F32) { pv_f = (const float *) pv->data; } + else { pv_f32.alloc(P*dv); ggml_get_to_fp32_cuda(pv->type)(pv->data, pv_f32.get(), P*dv, 1, stream); pv_f = pv_f32.get(); } + } else { + if (pk->type == GGML_TYPE_F16) { pk_h = (const half *) pk->data; } + else { pk_f16.alloc(Dk*P); ggml_get_to_fp16_cuda(pk->type)(pk->data, pk_f16.get(), Dk*P, 1, stream); pk_h = pk_f16.get(); } + if (pv->type == GGML_TYPE_F16) { pv_h = (const half *) pv->data; } + else { pv_f16.alloc(P*dv); ggml_get_to_fp16_cuda(pv->type)(pv->data, pv_f16.get(), P*dv, 1, stream); pv_h = pv_f16.get(); } + } + } + + // ---- per-tile buffers ---- + ggml_cuda_pool_alloc scores(ctx.pool(), PN*cw); + ggml_cuda_pool_alloc outbuf(ctx.pool(), (int64_t) dv*cw); + ggml_cuda_pool_alloc q_f16(ctx.pool()); // f16 path: range-scaled q tile + ggml_cuda_pool_alloc w_f16(ctx.pool()); // f16 path: normalized weights + ggml_cuda_pool_alloc w_f32(ctx.pool()); // f32 path: weights (== softmax out) + ggml_cuda_pool_alloc q_restore(ctx.pool()); + if (f16) { + q_f16.alloc(Dk*cw); + w_f16.alloc(PN*cw); + q_restore.alloc(cw); + } else { + w_f32.alloc(PN*cw); + } + + const int sm_nth = 256; + const size_t sm_shmem = WARP_SIZE*sizeof(float); + + for (int64_t c0 = 0; c0 < QT; c0 += cw) { + const int64_t cwn = std::min(cw, QT - c0); + + // ---- scores = scale * K^T @ Q (prefix rows [0,P), cache rows [P,PN)) ---- + const float * qtile = (const float *) q->data + c0*Dk; // contiguous [Dk, cwn] + if (!f16) { + if (P > 0) sgemm(ctx, CUBLAS_OP_T, CUBLAS_OP_N, P, cwn, Dk, scale, + pk_f, Dk, qtile, Dk, 0.0f, scores.get(), PN); + sgemm(ctx, CUBLAS_OP_T, CUBLAS_OP_N, N, cwn, Dk, scale, + cache_f, cache_stride, qtile, Dk, 0.0f, scores.get() + P, PN); + } else { + k_latent_pack_q_scaled<<>>( + qtile, q_f16.get(), q_restore.get(), Dk, 1, 0, cwn, Dk, 0, false); + CUDA_CHECK(cudaGetLastError()); + if (P > 0) hgemm_f32acc(ctx, CUBLAS_OP_T, CUBLAS_OP_N, P, cwn, Dk, scale, + pk_h, Dk, q_f16.get(), Dk, 0.0f, scores.get(), PN); + hgemm_f32acc(ctx, CUBLAS_OP_T, CUBLAS_OP_N, N, cwn, Dk, scale, + cache_h, cache_stride, q_f16.get(), Dk, 0.0f, scores.get() + P, PN); + } + + // ---- fused range restoration + mask + column softmax -> weights ---- + if (f16) { + k_latent_mask_softmax<<>>( + scores.get(), w_f16.get(), q_restore.get(), + mask ? (const float *) mask->data : nullptr, PN, P, T, c0, mask_nb1_f); + } else { + k_latent_mask_softmax<<>>( + scores.get(), w_f32.get(), nullptr, + mask ? (const float *) mask->data : nullptr, PN, P, T, c0, mask_nb1_f); + } + CUDA_CHECK(cudaGetLastError()); + + // ---- values: out = cacheV @ W_cache (+ prefix_v^T @ W_prefix) ---- + if (!f16) { + const float * cacheV = cache_f + dv_off; // row offset within packed row + sgemm(ctx, CUBLAS_OP_N, CUBLAS_OP_N, dv, cwn, N, 1.0f, + cacheV, cache_stride, w_f32.get() + P, PN, 0.0f, outbuf.get(), dv); + if (P > 0) sgemm(ctx, CUBLAS_OP_T, CUBLAS_OP_N, dv, cwn, P, 1.0f, + pv_f, P, w_f32.get(), PN, 1.0f, outbuf.get(), dv); + } else { + const half * cacheV = cache_h + dv_off; + hgemm_f32acc(ctx, CUBLAS_OP_N, CUBLAS_OP_N, dv, cwn, N, 1.0f, + cacheV, cache_stride, w_f16.get() + P, PN, 0.0f, outbuf.get(), dv); + if (P > 0) hgemm_f32acc(ctx, CUBLAS_OP_T, CUBLAS_OP_N, dv, cwn, P, 1.0f, + pv_h, P, w_f16.get(), PN, 1.0f, outbuf.get(), dv); + } + + // ---- copy tile -> dst (both column-major [Dv, .] contiguous) ---- + CUDA_CHECK(cudaMemcpyAsync((float *) dst->data + c0*dv, outbuf.get(), + (size_t) dv*cwn*sizeof(float), cudaMemcpyDeviceToDevice, stream)); + } +} diff --git a/ggml/src/ggml-cuda/latent_attn.cuh b/ggml/src/ggml-cuda/latent_attn.cuh new file mode 100644 index 000000000..0a3bed50e --- /dev/null +++ b/ggml/src/ggml-cuda/latent_attn.cuh @@ -0,0 +1,4 @@ +#include "common.cuh" + +bool ggml_cuda_latent_attn_is_supported(const ggml_tensor * op); +void ggml_cuda_op_latent_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h index d1d164315..ec24fcd3c 100644 --- a/ggml/src/ggml-cuda/vendors/hip.h +++ b/ggml/src/ggml-cuda/vendors/hip.h @@ -30,6 +30,7 @@ #define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS #define cublasSetStream hipblasSetStream #define cublasSgemm hipblasSgemm +#define cublasSgemmStridedBatched hipblasSgemmStridedBatched #define cublasStatus_t hipblasStatus_t #define cudaDataType_t hipblasDatatype_t //deprecated, new hipblasDatatype not in 5.6 #define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer diff --git a/ggml/src/ggml-cuda/vendors/musa.h b/ggml/src/ggml-cuda/vendors/musa.h index e50a103ac..d23732b9e 100644 --- a/ggml/src/ggml-cuda/vendors/musa.h +++ b/ggml/src/ggml-cuda/vendors/musa.h @@ -25,6 +25,7 @@ #define cublasSetMathMode mublasSetMathMode #define cublasSetStream mublasSetStream #define cublasSgemm mublasSgemm +#define cublasSgemmStridedBatched mublasSgemmStridedBatched #define cublasStatus_t mublasStatus_t #define cublasGetStatusString mublasStatus_to_string #define cudaDataType_t musaDataType_t diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index c3bcf1815..fe82fa463 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -4340,9 +4340,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "HC_PRE", "HC_POST", "MASK_TO_IDX", + "LATENT_ATTN", }; -static_assert(GGML_OP_COUNT == 109, "GGML_OP_COUNT != 109"); +static_assert(GGML_OP_COUNT == 110, "GGML_OP_COUNT != 110"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -4467,10 +4468,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "hc_pre(x,s,b)", "hc_post(x,p,r,c)", "mask_to_idx(masl)", + "latent_attn_prefix(q,c,pk,pv,mask)", }; -static_assert(GGML_OP_COUNT == 109, "GGML_OP_COUNT != 109"); +static_assert(GGML_OP_COUNT == 110, "GGML_OP_COUNT != 110"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -10480,6 +10482,150 @@ struct ggml_tensor * ggml_top_k_thresh( return result; } +// ggml_latent_attn_ext_impl + +static struct ggml_tensor * ggml_latent_attn_ext_impl( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * cache, + struct ggml_tensor * prefix_k, + struct ggml_tensor * prefix_v, + struct ggml_tensor * mask, + struct ggml_tensor * indices, + int dv, + int dv_off, + float scale, + float max_bias, + int mode) { + // Inference-only op: reject a differentiable graph at construction rather than silently + // detaching it (there is no backward pass; see the GGML_OP_LATENT_ATTN backward dispatch). + if ((q && q->grad) || (cache && cache->grad) || (prefix_k && prefix_k->grad) || + (prefix_v && prefix_v->grad) || (mask && mask->grad) || (indices && indices->grad)) { + GGML_ABORT("ggml_latent_attn does not support automatic differentiation"); + } + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(q)); + GGML_ASSERT(q->ne[0] == cache->ne[0]); // Dk match + // Both backends compute exactly [Dv, T, H]: unused outer dimensions must be 1 on every + // operand, so a batched or higher-rank view is rejected instead of silently truncated. + GGML_ASSERT(q->ne[3] == 1); + GGML_ASSERT(cache->ne[2] == 1 && cache->ne[3] == 1); + GGML_ASSERT(dv > 0 && dv_off >= 0 && (int64_t) dv_off + dv <= cache->ne[0]); + GGML_ASSERT(max_bias == 0.0f && "ggml_latent_attn: ALiBi not implemented"); + GGML_ASSERT(cache->type == GGML_TYPE_F32 || cache->type == GGML_TYPE_F16 || + cache->type == GGML_TYPE_Q8_0); + GGML_ASSERT(mode == 0 || mode == 1); + GGML_ASSERT((indices != NULL) == (mode == 1)); + + // prefix is all-or-nothing + GGML_ASSERT((prefix_k == NULL) == (prefix_v == NULL)); + if (prefix_k) { + GGML_ASSERT(prefix_k->ne[0] == cache->ne[0]); // Dk + GGML_ASSERT(prefix_v->ne[0] == prefix_k->ne[1]); // value-transposed: ne0 = P + GGML_ASSERT(prefix_v->ne[1] == dv); + GGML_ASSERT(ggml_is_contiguous(prefix_k)); + GGML_ASSERT(ggml_is_contiguous(prefix_v)); + GGML_ASSERT(prefix_k->type == GGML_TYPE_F32 || prefix_k->type == GGML_TYPE_F16); + GGML_ASSERT(prefix_v->type == GGML_TYPE_F32 || prefix_v->type == GGML_TYPE_F16); + GGML_ASSERT(prefix_k->ne[2] == 1 && prefix_k->ne[3] == 1); + GGML_ASSERT(prefix_v->ne[2] == 1 && prefix_v->ne[3] == 1); + } + + if (mask) { + GGML_ASSERT(mask->type == GGML_TYPE_F32); + GGML_ASSERT(mask->nb[0] == sizeof(float)); + // Both backends step mask rows by float-pointer arithmetic (nb[1]/sizeof(float)), which + // cannot express a row stride that is not a whole number of floats. Reject it here so + // construction and both backends agree; canonical OpenPangu masks are aligned. + GGML_ASSERT(mask->nb[1] % sizeof(float) == 0 && + "ggml_latent_attn: mask row stride must be a whole number of floats"); + GGML_ASSERT(mask->ne[0] == cache->ne[1]); // N + GGML_ASSERT(mask->ne[1] >= q->ne[1]); // >= T + GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); + } + + if (indices) { + GGML_ASSERT(indices->type == GGML_TYPE_I32); + GGML_ASSERT(indices->ne[0] >= 1); // topk + GGML_ASSERT(indices->ne[1] >= q->ne[1]); // >= T + GGML_ASSERT(ggml_is_contiguous(indices)); + GGML_ASSERT(indices->ne[2] == 1 && indices->ne[3] == 1); + } + + // block alignment so quantized-cache dequant and the value slice stay on block rows, and + // packed rows so the constructor accepts exactly what the backends implement + if (ggml_is_quantized(cache->type)) { + const int64_t blk = ggml_blck_size(cache->type); + GGML_ASSERT(cache->ne[0] % blk == 0); + GGML_ASSERT(dv % blk == 0); + GGML_ASSERT(dv_off % blk == 0); + GGML_ASSERT(cache->nb[1] == ggml_row_size(cache->type, cache->ne[0]) && + "ggml_latent_attn: quantized cache rows must be packed"); + } + + // The dense (cuBLAS) and indexed (row-gather) readers both index each F32/F16 cache row as a + // packed element array (src[d]), so a non-packed inner stride would read the wrong elements. + // The row stride nb[1] may still be a windowed view. + if (cache->type == GGML_TYPE_F32 || cache->type == GGML_TYPE_F16) { + GGML_ASSERT(cache->nb[0] == ggml_type_size(cache->type) && + "ggml_latent_attn: cache inner dimension must be packed"); + } + + int64_t ne[4] = { dv, q->ne[1], q->ne[2], 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_f32(result, 0, scale); + ggml_set_op_params_f32(result, 1, max_bias); + ggml_set_op_params_i32(result, 2, dv); + ggml_set_op_params_i32(result, 3, dv_off); + ggml_set_op_params_i32(result, 4, mode); // mode 0 = dense; mode 1 = indexed + + result->op = GGML_OP_LATENT_ATTN; + result->src[0] = q; + result->src[1] = cache; + result->src[2] = prefix_k; + result->src[3] = prefix_v; + result->src[4] = mask; + result->src[5] = indices; + + return result; +} + +// ggml_latent_attn_prefix_ext + +struct ggml_tensor * ggml_latent_attn_prefix_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * cache, + struct ggml_tensor * prefix_k, + struct ggml_tensor * prefix_v, + struct ggml_tensor * mask, + int dv, + int dv_off, + float scale, + float max_bias) { + return ggml_latent_attn_ext_impl(ctx, q, cache, prefix_k, prefix_v, mask, NULL, + dv, dv_off, scale, max_bias, 0); +} + +// ggml_latent_attn_indexed_ext + +struct ggml_tensor * ggml_latent_attn_indexed_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * cache, + struct ggml_tensor * prefix_k, + struct ggml_tensor * prefix_v, + struct ggml_tensor * mask, + struct ggml_tensor * indices, + int dv, + int dv_off, + float scale, + float max_bias) { + return ggml_latent_attn_ext_impl(ctx, q, cache, prefix_k, prefix_v, mask, indices, + dv, dv_off, scale, max_bias, 1); +} + // ggml_flash_attn_ext struct ggml_tensor * ggml_flash_attn_ext( @@ -23911,6 +24057,239 @@ static void ggml_compute_forward_mask_to_idx(const struct ggml_compute_params * } +// ggml_compute_forward_latent_attn + +// Latent attention over a packed K/V cache with an independently-visible K/V prefix. +// Reference (correctness-first) implementation: parallelize over the T*H query rows; for +// each row compute prefix+cache logits, one joint f32 softmax over [prefix | cache], then +// the two value contractions. Quantized/F16 cache rows are dequantized whole into a per- +// thread scratch row (one row at a time -> the full F32 cache tensor is never materialized). +// The value slice of cache row n is channels [dv_off, dv_off+dv); the prefix value tensor is +// value-transposed [P, Dv] (ne0=P) exactly like the openPangu s_lat_t operand. The fused +// single-dequant-per-row optimization lives in the CUDA kernel. In dense mode this reference +// dequantizes each non-F32 cache row once for the K dot and again for a nonzero V contribution. +// A masked-out row has softmax weight exactly 0, so its value contribution is skipped. +// Mode 1 gathers indices[k,t] shared across heads and applies mask[idx,t], then jointly +// normalizes prefix and gathered rows. Non-F32 gathered rows are dequantized whole once. +static void ggml_compute_forward_latent_attn_f32( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * q = dst->src[0]; + const struct ggml_tensor * cache = dst->src[1]; + const struct ggml_tensor * prefix_k = dst->src[2]; + const struct ggml_tensor * prefix_v = dst->src[3]; + const struct ggml_tensor * mask = dst->src[4]; + const struct ggml_tensor * indices = dst->src[5]; + + const int64_t Dk = q->ne[0]; + const int64_t T = q->ne[1]; + const int64_t H = q->ne[2]; + const int64_t N = cache->ne[1]; + const int64_t P = prefix_k ? prefix_k->ne[1] : 0; + const int64_t topk = indices ? indices->ne[0] : 0; + + float scale; + memcpy(&scale, &dst->op_params[0], sizeof(float)); + const int dv = dst->op_params[2]; + const int dv_off = dst->op_params[3]; + const int mode = dst->op_params[4]; + GGML_ASSERT(mode == 0 || mode == 1); + GGML_ASSERT((indices != NULL) == (mode == 1)); + GGML_ASSERT(dst->ne[0] == dv && dst->ne[1] == T && dst->ne[2] == H); + + const int ith = params->ith; + const int nth = params->nth; + + // per-thread scratch: [P + dense-or-gathered rows] scores followed by one [Dk] dequant row + const int64_t score_count = P + (mode == 0 ? N : topk); + const int64_t stride = score_count + Dk; + float * scores = (float *) params->wdata + (size_t) ith*stride; + float * rowbuf = scores + score_count; + + const bool cache_is_f32 = cache->type == GGML_TYPE_F32; + ggml_to_float_t cache_to_float = cache_is_f32 ? NULL : type_traits[cache->type].to_float; + + const int64_t QR = T*H; // total query rows + const int64_t r0 = (QR * ith ) / nth; + const int64_t r1 = (QR * (ith+1)) / nth; + + for (int64_t r = r0; r < r1; ++r) { + const int64_t t = r % T; + const int64_t h = r / T; + + const float * qrow = (const float *)((const char *)q->data + t*q->nb[1] + h*q->nb[2]); // [Dk] + + if (mode == 1) { + // Online accumulation lets one whole-row dequant serve both the K dot and V slice. + float * out = (float *)((char *)dst->data + t*dst->nb[1] + h*dst->nb[2]); // [Dv] + for (int64_t d = 0; d < dv; ++d) out[d] = 0.0f; + + float mx = -INFINITY; + for (int64_t p = 0; p < P; ++p) { + const char * kcol = (const char *)prefix_k->data + p*prefix_k->nb[1]; + float s = 0.0f; + if (prefix_k->type == GGML_TYPE_F32) { + const float * kf = (const float *)kcol; + for (int64_t i = 0; i < Dk; ++i) s += qrow[i]*kf[i]; + } else { + const ggml_fp16_t * kf = (const ggml_fp16_t *)kcol; + for (int64_t i = 0; i < Dk; ++i) s += qrow[i]*GGML_FP16_TO_FP32(kf[i]); + } + scores[p] = scale*s; + mx = MAX(mx, scores[p]); + } + + float sum = 0.0f; + for (int64_t p = 0; p < P; ++p) { + scores[p] = expf(scores[p] - mx); + sum += scores[p]; + } + for (int64_t d = 0; d < dv; ++d) { + const char * pvd = prefix_v ? (const char *)prefix_v->data + d*prefix_v->nb[1] : NULL; + float acc = 0.0f; + if (P > 0) { + if (prefix_v->type == GGML_TYPE_F32) { + const float * pv = (const float *)pvd; + for (int64_t p = 0; p < P; ++p) acc += scores[p]*pv[p]; + } else { + const ggml_fp16_t * pv = (const ggml_fp16_t *)pvd; + for (int64_t p = 0; p < P; ++p) acc += scores[p]*GGML_FP16_TO_FP32(pv[p]); + } + } + out[d] = acc; + } + + for (int64_t k = 0; k < topk; ++k) { + const int32_t idx = *(const int32_t *)((const char *)indices->data + t*indices->nb[1] + k*indices->nb[0]); + GGML_ASSERT(idx >= 0 && idx < N); + + const char * crow = (const char *)cache->data + idx*cache->nb[1]; + const float * cf; + if (cache_is_f32) { + cf = (const float *)crow; + } else { + cache_to_float(crow, rowbuf, Dk); + cf = rowbuf; + } + + float s = 0.0f; + for (int64_t i = 0; i < Dk; ++i) s += qrow[i]*cf[i]; + s *= scale; + if (mask) { + s += *(const float *)((const char *)mask->data + t*mask->nb[1] + idx*mask->nb[0]); + } + + if (mx == -INFINITY && s == -INFINITY) { + continue; + } + if (s > mx) { + const float rescale = expf(mx - s); + sum = sum*rescale + 1.0f; + for (int64_t d = 0; d < dv; ++d) out[d] = out[d]*rescale + cf[dv_off + d]; + mx = s; + } else { + const float e = expf(s - mx); + sum += e; + for (int64_t d = 0; d < dv; ++d) out[d] += e*cf[dv_off + d]; + } + } + + const float inv = sum > 0.0f ? 1.0f/sum : 0.0f; + for (int64_t d = 0; d < dv; ++d) out[d] *= inv; + continue; + } + + // --- logits: prefix rows (always visible) then cache rows (masked) --- + for (int64_t p = 0; p < P; ++p) { + const char * kcol = (const char *)prefix_k->data + p*prefix_k->nb[1]; + float s = 0.0f; + if (prefix_k->type == GGML_TYPE_F32) { + const float * kf = (const float *)kcol; + for (int64_t i = 0; i < Dk; ++i) s += qrow[i]*kf[i]; + } else { + const ggml_fp16_t * kf = (const ggml_fp16_t *)kcol; + for (int64_t i = 0; i < Dk; ++i) s += qrow[i]*GGML_FP16_TO_FP32(kf[i]); + } + scores[p] = scale*s; + } + for (int64_t n = 0; n < N; ++n) { + const char * crow = (const char *)cache->data + n*cache->nb[1]; + const float * cf; + if (cache_is_f32) { + cf = (const float *)crow; + } else { + cache_to_float(crow, rowbuf, Dk); + cf = rowbuf; + } + float s = 0.0f; + for (int64_t i = 0; i < Dk; ++i) s += qrow[i]*cf[i]; + s *= scale; + if (mask) { + s += *(const float *)((const char *)mask->data + t*mask->nb[1] + n*mask->nb[0]); + } + scores[P + n] = s; + } + + // --- one joint softmax over [0, P+N) --- + const int64_t M = P + N; + float mx = -INFINITY; + for (int64_t k = 0; k < M; ++k) mx = MAX(mx, scores[k]); + float sum = 0.0f; + for (int64_t k = 0; k < M; ++k) { const float e = expf(scores[k] - mx); scores[k] = e; sum += e; } + const float inv = sum > 0.0f ? 1.0f/sum : 0.0f; + for (int64_t k = 0; k < M; ++k) scores[k] *= inv; + + // --- values: out[d] = sum_p w[p]*prefix_v[p,d] + sum_n w[P+n]*cache[dv_off+d, n] --- + float * out = (float *)((char *)dst->data + t*dst->nb[1] + h*dst->nb[2]); // [Dv] + for (int64_t d = 0; d < dv; ++d) out[d] = 0.0f; + + // prefix values: prefix_v is [P, Dv] (ne0=P), so row d holds the P prefix values for d + for (int64_t d = 0; d < dv; ++d) { + const char * pvd = prefix_v ? (const char *)prefix_v->data + d*prefix_v->nb[1] : NULL; + float acc = 0.0f; + if (P > 0) { + if (prefix_v->type == GGML_TYPE_F32) { + const float * pv = (const float *)pvd; + for (int64_t p = 0; p < P; ++p) acc += scores[p]*pv[p]; + } else { + const ggml_fp16_t * pv = (const ggml_fp16_t *)pvd; + for (int64_t p = 0; p < P; ++p) acc += scores[p]*GGML_FP16_TO_FP32(pv[p]); + } + } + out[d] = acc; + } + + // cache values: re-dequant each row, take channels [dv_off, dv_off+dv). A masked row + // has softmax weight exactly 0, so its contribution vanishes without a branch. + for (int64_t n = 0; n < N; ++n) { + const float w = scores[P + n]; + if (w == 0.0f) continue; + const char * crow = (const char *)cache->data + n*cache->nb[1]; + const float * cf; + if (cache_is_f32) { + cf = (const float *)crow; + } else { + cache_to_float(crow, rowbuf, Dk); + cf = rowbuf; + } + const float * cv = cf + dv_off; + for (int64_t d = 0; d < dv; ++d) out[d] += w * cv[d]; + } + } +} + +static void ggml_compute_forward_latent_attn( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: + ggml_compute_forward_latent_attn_f32(params, dst); + break; + default: + GGML_ABORT("fatal error"); + } +} + // ggml_compute_forward_win_part static void ggml_compute_forward_win_part_f32( @@ -25673,6 +26052,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml { ggml_compute_forward_mask_to_idx(params, tensor); } break; + case GGML_OP_LATENT_ATTN: + { + ggml_compute_forward_latent_attn(params, tensor); + } break; case GGML_OP_INDEXER_TOPK: { if (!iqk_indexer_topk(tensor, params->wdata, (barrier_t)ggml_barrier, (void *)params->shared, params->ith, params->nth)) { @@ -26750,6 +27133,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor case GGML_OP_HC_PRE: case GGML_OP_HC_POST: case GGML_OP_MASK_TO_IDX: + case GGML_OP_LATENT_ATTN: { GGML_ABORT("fatal error"); // TODO: not implemented } @@ -27497,6 +27881,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_HC_PRE: case GGML_OP_HC_POST: case GGML_OP_MASK_TO_IDX: + case GGML_OP_LATENT_ATTN: { n_tasks = n_threads; } break; @@ -27805,6 +28190,17 @@ struct ggml_cplan ggml_graph_plan(const struct ggml_cgraph * cgraph, int n_threa cur = MAX(cur, size); #endif } break; + case GGML_OP_LATENT_ATTN: + { + // per thread: [P+dense-or-gathered rows] f32 scores + one [Dk] dequant row + const int64_t Dk = node->src[1]->ne[0]; + const int64_t N = node->src[1]->ne[1]; + const int64_t P = node->src[2] ? node->src[2]->ne[1] : 0; + const int mode = node->op_params[4]; + GGML_ASSERT(mode == 0 || mode == 1); + const int64_t rows = mode == 0 ? N : node->src[5]->ne[0]; + cur = sizeof(float)*(P + rows + Dk)*n_tasks; + } break; case GGML_OP_FLASH_ATTN_BACK: { const int64_t D = node->src[0]->ne[0]; diff --git a/src/graphs/build_openpangu.cpp b/src/graphs/build_openpangu.cpp index 6961dbb0d..99a546170 100644 --- a/src/graphs/build_openpangu.cpp +++ b/src/graphs/build_openpangu.cpp @@ -226,6 +226,48 @@ static ggml_tensor * openpangu_causal_conv(ggml_context * ctx, ggml_cgraph * gf, // x_normed = input-layernormed hidden [n_embd, T]; returns post-o_proj output [n_embd, T]. // conv_state is the recurrent MoME state slot. seq_qnext is the [1, T] sequence-id input // used by ggml_ssm_conv and is shared by all three conv sites. + +// Raw latent-cache view (native type: f32/f16/q8_0) for the fused latent-attention op, +// which dequantizes internally -- no F32 cast, no value transpose. +static ggml_tensor * openpangu_build_k_latent_raw( + ggml_context * ctx, const llama_kv_cache & kv_self, int il, + int64_t n_kv_view, int64_t win_off) { + ggml_tensor * kl = kv_self.k_l[il]; + return ggml_view_2d(ctx, kl, kl->ne[0], n_kv_view, kl->nb[1], (size_t) win_off*kl->nb[1]); +} + +// Whether to route dense/SWA spans or indexed DSA gathers through the latent-attention op. +// Capability-gated at the call site (openpangu_backend_supports_fused_attn). The op has no +// ALiBi, so a configured max-bias keeps the explicit chain. +static bool openpangu_fused_attn_enabled(const llama_hparams & hparams) { + return hparams.f_max_alibi_bias == 0.0f; +} + +// Confirm the scheduled backend for this layer's attention output projection can execute +// the fused op before adopting it, so an unsupported backend (or a CUDA layout the op +// rejects) keeps the exact unfused chain instead of forcing a scheduler CPU island. The +// projection weight anchors the scheduled placement. The CPU backend truthfully reports +// support (it carries the scalar reference forward), so the refusal below is performance +// policy, not capability: adopting the reference on a CPU-resident layer would swap the +// vectorized unfused chain for a correctness-oriented loop. The latent cache must also be +// resident on the chosen backend: under --no-kv-offload it stays in host memory while the +// projection weights sit on the device, and adopting the op there would re-upload the +// whole cache view every graph, so that case keeps the unfused chain too. +static bool openpangu_backend_supports_fused_attn( + llama_context & lctx, const llama_hparams & hparams, + ggml_tensor * placement, ggml_tensor * kv_cache, ggml_tensor * candidate) { + if (!openpangu_fused_attn_enabled(hparams) || placement == nullptr || kv_cache == nullptr || + candidate == nullptr) { + return false; + } + ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(lctx.sched, placement); + if (backend == nullptr || ggml_backend_is_cpu(backend) || + !ggml_backend_supports_op(backend, candidate)) { + return false; + } + return ggml_backend_sched_get_tensor_backend(lctx.sched, kv_cache) == backend; +} + ggml_tensor * llm_build_context::build_openpangu_attention( ggml_cgraph * gf, const llama_layer & layer, int il, ggml_tensor * x_normed, ggml_tensor * KQ_mask, ggml_tensor * inp_pos, @@ -511,67 +553,118 @@ ggml_tensor * llm_build_context::build_openpangu_attention( return vl_all; }; - ggml_tensor * k_gath = nullptr; - ggml_tensor * kq_cache = nullptr; + // Capability probe: build a representative dense latent-attention candidate and adopt the + // fused path only when this layer's scheduled backend can execute it. When the unfused + // chain is chosen (unset probe, unsupported backend, or forced unfused mode), the + // probe node is left unreferenced, so it never enters the built graph. + ggml_tensor * fused_attn_probe = nullptr; + if (openpangu_fused_attn_enabled(hparams)) { + ggml_tensor * kl_raw_probe = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv_attn, win_off); + ggml_tensor * mask_probe = KQ_mask ? + ggml_view_2d(ctx0, KQ_mask, n_kv_attn, n_tokens, KQ_mask->nb[1], 0) : nullptr; + if (kl_raw_probe->type == GGML_TYPE_F32 || kl_raw_probe->type == GGML_TYPE_F16 || + kl_raw_probe->type == GGML_TYPE_Q8_0) { + fused_attn_probe = ggml_latent_attn_prefix_ext( + ctx0, q_all, kl_raw_probe, sink_blk, s_lat_t, mask_probe, + kv_lora_rank, 0, kq_scale, hparams.f_max_alibi_bias); + } + } + const bool use_fused_attn = + openpangu_backend_supports_fused_attn(lctx, hparams, layer.wv_b, kv_self.k_l[il], fused_attn_probe); if (use_dsa_gather) { - ggml_tensor * kq_sinks = ggml_mul_mat(ctx0, sink_blk, q_all); // [NS, T, H] - if (n_tokens == 1) { - ggml_tensor * kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], kv_lora_rank + n_embd_head_qk_rope, n_kv, - kv_self.k_l[il]->nb[1], 0); - ggml_tensor * sel_idx_flat = ggml_cont_2d(ctx0, sel_idx, dsa_topk, 1); // [topk] i32 - k_gath = ggml_get_rows(ctx0, kl_full, sel_idx_flat); // [576, topk] f32 - k_gath = openpangu_cast_gathered_latent_for_cache_type(ctx0, k_gath, kv_self.k_l[il]); - k_gath = ggml_reshape_3d(ctx0, k_gath, kv_lora_rank + n_embd_head_qk_rope, dsa_topk, 1); - kq_cache = ggml_mul_mat(ctx0, - ggml_reshape_2d(ctx0, k_gath, kv_lora_rank + n_embd_head_qk_rope, dsa_topk), - q_all); // [topk, 1, H] + ggml_tensor * kqv = nullptr; + if (use_fused_attn) { + // sel_idx contains absolute cache rows selected after the causal indexer mask, + // so gathered attention is maskless just like the unfused chain. + ggml_tensor * kl_raw = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv, 0); + ggml_tensor * sel_idx_cont = ggml_cont(ctx0, sel_idx); // [topk, T] i32 + kqv = ggml_latent_attn_indexed_ext(ctx0, q_all, kl_raw, sink_blk, s_lat_t, nullptr, + sel_idx_cont, kv_lora_rank, 0, + kq_scale, 0.0f); // [512, T, H] } else { - ggml_tensor * kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], - kv_lora_rank + n_embd_head_qk_rope, n_kv, - kv_self.k_l[il]->nb[1], 0); - ggml_tensor * sel_idx_flat = ggml_cont_2d(ctx0, sel_idx, dsa_topk*n_tokens, 1); // [topk*T] i32 - k_gath = ggml_get_rows(ctx0, kl_full, sel_idx_flat); // [576, topk*T] f32 - k_gath = openpangu_cast_gathered_latent_for_cache_type(ctx0, k_gath, kv_self.k_l[il]); - k_gath = ggml_reshape_3d(ctx0, k_gath, kv_lora_rank + n_embd_head_qk_rope, - dsa_topk, n_tokens); // [576, topk, T] - ggml_tensor * q_gath = ggml_cont(ctx0, ggml_permute(ctx0, q_all, 0, 2, 1, 3)); // [576, H, T] - kq_cache = ggml_mul_mat(ctx0, k_gath, q_gath); // [topk, H, T] - kq_cache = ggml_cont(ctx0, ggml_permute(ctx0, kq_cache, 0, 2, 1, 3)); // [topk, T, H] - } - ggml_tensor * kq = ggml_concat(ctx0, kq_sinks, kq_cache, 0); // [NS+topk, T, H] - // sel_idx came from scores after adding the causal KQ_mask. The engagement bound gives - // every token at least topk valid positions, so all gathered rows are visible here. - kq = ggml_soft_max_ext(ctx0, kq, nullptr, kq_scale, hparams.f_max_alibi_bias); + // Unfused fallback: the complete pre-gather + explicit attention chain, kept for + // backends and configs the op declines. + ggml_tensor * k_gath = nullptr; + ggml_tensor * kq_cache = nullptr; + ggml_tensor * kq_sinks = ggml_mul_mat(ctx0, sink_blk, q_all); // [NS, T, H] + if (n_tokens == 1) { + ggml_tensor * kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], + kv_lora_rank + n_embd_head_qk_rope, n_kv, + kv_self.k_l[il]->nb[1], 0); + ggml_tensor * sel_idx_flat = ggml_cont_2d(ctx0, sel_idx, dsa_topk, 1); // [topk] i32 + k_gath = ggml_get_rows(ctx0, kl_full, sel_idx_flat); // [576, topk] f32 + k_gath = openpangu_cast_gathered_latent_for_cache_type( + ctx0, k_gath, kv_self.k_l[il]); + k_gath = ggml_reshape_3d(ctx0, k_gath, + kv_lora_rank + n_embd_head_qk_rope, dsa_topk, 1); + kq_cache = ggml_mul_mat(ctx0, + ggml_reshape_2d(ctx0, k_gath, + kv_lora_rank + n_embd_head_qk_rope, dsa_topk), + q_all); // [topk, 1, H] + } else { + ggml_tensor * kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], + kv_lora_rank + n_embd_head_qk_rope, n_kv, + kv_self.k_l[il]->nb[1], 0); + ggml_tensor * sel_idx_flat = ggml_cont_2d( + ctx0, sel_idx, dsa_topk*n_tokens, 1); // [topk*T] i32 + k_gath = ggml_get_rows(ctx0, kl_full, sel_idx_flat); // [576, topk*T] f32 + k_gath = openpangu_cast_gathered_latent_for_cache_type( + ctx0, k_gath, kv_self.k_l[il]); + k_gath = ggml_reshape_3d(ctx0, k_gath, + kv_lora_rank + n_embd_head_qk_rope, + dsa_topk, n_tokens); // [576, topk, T] + ggml_tensor * q_gath = ggml_cont(ctx0, + ggml_permute(ctx0, q_all, 0, 2, 1, 3)); // [576, H, T] + kq_cache = ggml_mul_mat(ctx0, k_gath, q_gath); // [topk, H, T] + kq_cache = ggml_cont(ctx0, + ggml_permute(ctx0, kq_cache, 0, 2, 1, 3)); // [topk, T, H] + } + ggml_tensor * kq = ggml_concat(ctx0, kq_sinks, kq_cache, 0); // [NS+topk, T, H] + // sel_idx came from scores after adding the causal KQ_mask. The engagement bound + // gives every token at least topk valid positions, so all gathered rows are visible. + kq = ggml_soft_max_ext(ctx0, kq, nullptr, kq_scale, hparams.f_max_alibi_bias); - ggml_tensor * kq_s = ggml_view_3d(ctx0, kq, NS, n_tokens, n_head, kq->nb[1], kq->nb[2], 0); - ggml_tensor * kq_c = ggml_view_3d(ctx0, kq, n_kv_attn, n_tokens, n_head, kq->nb[1], kq->nb[2], - NS*ggml_element_size(kq)); - ggml_tensor * kqv_cache = nullptr; - GGML_ASSERT(k_gath != nullptr); - if (n_tokens == 1) { - ggml_tensor * v_gath = ggml_view_2d(ctx0, k_gath, kv_lora_rank, dsa_topk, - k_gath->nb[1], 0); // [512, topk] - ggml_tensor * v_gath_t = ggml_cont(ctx0, ggml_transpose(ctx0, v_gath)); // [topk, 512] - kqv_cache = ggml_mul_mat(ctx0, v_gath_t, kq_c); // [512, 1, H] - } else { - ggml_tensor * v_gath = ggml_view_3d(ctx0, k_gath, kv_lora_rank, dsa_topk, n_tokens, - k_gath->nb[1], k_gath->nb[2], 0); // [512, topk, T] - ggml_tensor * v_gath_t = ggml_cont(ctx0, ggml_permute(ctx0, v_gath, 1, 0, 2, 3)); // [topk, 512, T] - ggml_tensor * kq_c_gath = ggml_cont(ctx0, ggml_permute(ctx0, kq_c, 0, 2, 1, 3)); // [topk, H, T] - kqv_cache = ggml_mul_mat(ctx0, v_gath_t, kq_c_gath); // [512, H, T] - kqv_cache = ggml_cont(ctx0, ggml_permute(ctx0, kqv_cache, 0, 2, 1, 3)); // [512, T, H] + ggml_tensor * kq_s = ggml_view_3d(ctx0, kq, NS, n_tokens, n_head, + kq->nb[1], kq->nb[2], 0); + ggml_tensor * kq_c = ggml_view_3d(ctx0, kq, n_kv_attn, n_tokens, n_head, + kq->nb[1], kq->nb[2], + NS*ggml_element_size(kq)); + ggml_tensor * kqv_cache = nullptr; + GGML_ASSERT(k_gath != nullptr); + if (n_tokens == 1) { + ggml_tensor * v_gath = ggml_view_2d(ctx0, k_gath, kv_lora_rank, dsa_topk, + k_gath->nb[1], 0); // [512, topk] + ggml_tensor * v_gath_t = ggml_cont(ctx0, ggml_transpose(ctx0, v_gath)); // [topk, 512] + kqv_cache = ggml_mul_mat(ctx0, v_gath_t, kq_c); // [512, 1, H] + } else { + ggml_tensor * v_gath = ggml_view_3d(ctx0, k_gath, kv_lora_rank, dsa_topk, + n_tokens, k_gath->nb[1], k_gath->nb[2], 0); // [512, topk, T] + ggml_tensor * v_gath_t = ggml_cont(ctx0, + ggml_permute(ctx0, v_gath, 1, 0, 2, 3)); // [topk, 512, T] + ggml_tensor * kq_c_gath = ggml_cont(ctx0, + ggml_permute(ctx0, kq_c, 0, 2, 1, 3)); // [topk, H, T] + kqv_cache = ggml_mul_mat(ctx0, v_gath_t, kq_c_gath); // [512, H, T] + kqv_cache = ggml_cont(ctx0, + ggml_permute(ctx0, kqv_cache, 0, 2, 1, 3)); // [512, T, H] + } + kqv = ggml_add(ctx0, + ggml_mul_mat(ctx0, s_lat_t, kq_s), + kqv_cache); // [512, T, H] } - ggml_tensor * kqv = ggml_add(ctx0, - ggml_mul_mat(ctx0, s_lat_t, kq_s), - kqv_cache); // [512, T, H] ggml_tensor * wv_b3 = ggml_reshape_3d(ctx0, layer.wv_b, kv_lora_rank, n_embd_head_v, n_head); ggml_tensor * out_h = ggml_mul_mat(ctx0, wv_b3, kqv); // [128, T, H] ggml_tensor * merged = ggml_cont(ctx0, ggml_permute(ctx0, out_h, 0, 2, 1, 3)); // [128, H, T] cur = ggml_reshape_2d(ctx0, merged, n_embd_head_v * n_head, n_tokens); } else { - const bool chunk_att = openpangu_att_score_should_chunk(n_kv_attn, NS, n_head, n_tokens, - OPENPANGU_ATT_SCORE_CHUNK, OPENPANGU_ATT_FULL_KQ_MAX_MIB); const bool use_dsa_sel_idx_mask = sel_idx != nullptr && dsa_topk > 0; + const bool chunk_att = + openpangu_att_score_should_chunk(n_kv_attn, NS, n_head, n_tokens, + OPENPANGU_ATT_SCORE_CHUNK, OPENPANGU_ATT_FULL_KQ_MAX_MIB) && + // Dense/SWA and MTP fused attention tiles the full T internally. DSA prefill + // keeps this loop because it also owns the gathered path, the unfused + // fallback's subchunks, and deferred selection-mask construction for + // non-gather chunks. + (!use_fused_attn || use_dsa_sel_idx_mask); ggml_tensor * kqv = nullptr; if (chunk_att) { const bool can_prefill_gather = @@ -579,8 +672,12 @@ ggml_tensor * llm_build_context::build_openpangu_attention( hparams.f_max_alibi_bias == 0.0f && openpangu_dsa_gather_rows_fit_cuda(dsa_topk, 1); ggml_tensor * kl_full = nullptr; - if (can_prefill_gather) { - kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], kv_lora_rank + n_embd_head_qk_rope, n_kv, + ggml_tensor * kl_raw_full = nullptr; + if (can_prefill_gather && use_fused_attn) { + kl_raw_full = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv, 0); + } else if (can_prefill_gather) { + kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], + kv_lora_rank + n_embd_head_qk_rope, n_kv, kv_self.k_l[il]->nb[1], 0); } @@ -591,7 +688,8 @@ ggml_tensor * llm_build_context::build_openpangu_attention( (size_t) c0*q_all->nb[1]); q_all_c = ggml_cont(ctx0, q_all_c); - ggml_tensor * kq_sinks_c = ggml_mul_mat(ctx0, sink_blk, q_all_c); // [NS, Tc, H] + ggml_tensor * kq_sinks_c = use_fused_attn ? nullptr : + ggml_mul_mat(ctx0, sink_blk, q_all_c); // [NS, Tc, H] const bool prefill_gather_chunk = can_prefill_gather && @@ -599,46 +697,98 @@ ggml_tensor * llm_build_context::build_openpangu_attention( openpangu_kv_cache_pad(cparams)); ggml_tensor * kqv_c = nullptr; if (prefill_gather_chunk) { - const int64_t gather_token_chunk = openpangu_dsa_gather_tokens_per_get_rows(dsa_topk); - for (int64_t g0 = 0; g0 < tc; g0 += gather_token_chunk) { - const int64_t tg = std::min(gather_token_chunk, tc - g0); - ggml_tensor * q_all_g = ggml_view_3d(ctx0, q_all_c, - kv_lora_rank + n_embd_head_qk_rope, - tg, n_head, q_all_c->nb[1], q_all_c->nb[2], - (size_t) g0*q_all_c->nb[1]); - q_all_g = ggml_cont(ctx0, q_all_g); - ggml_tensor * kq_sinks_g = ggml_view_3d(ctx0, kq_sinks_c, NS, tg, n_head, - kq_sinks_c->nb[1], kq_sinks_c->nb[2], - (size_t) g0*kq_sinks_c->nb[1]); - ggml_tensor * sel_idx_g = ggml_view_2d(ctx0, sel_idx, dsa_topk, tg, - sel_idx->nb[1], (size_t) (c0 + g0)*sel_idx->nb[1]); - ggml_tensor * sel_idx_flat_g = ggml_cont_2d(ctx0, sel_idx_g, dsa_topk*tg, 1); - ggml_tensor * k_gath_g = ggml_get_rows(ctx0, kl_full, sel_idx_flat_g); // [576, topk*tg] - k_gath_g = openpangu_cast_gathered_latent_for_cache_type(ctx0, k_gath_g, kv_self.k_l[il]); - k_gath_g = ggml_reshape_3d(ctx0, k_gath_g, kv_lora_rank + n_embd_head_qk_rope, - dsa_topk, tg); // [576, topk, tg] - ggml_tensor * q_gath_g = ggml_cont(ctx0, ggml_permute(ctx0, q_all_g, 0, 2, 1, 3)); // [576, H, tg] - ggml_tensor * kq_cache_g = ggml_mul_mat(ctx0, k_gath_g, q_gath_g); // [topk, H, tg] - kq_cache_g = ggml_cont(ctx0, ggml_permute(ctx0, kq_cache_g, 0, 2, 1, 3)); // [topk, tg, H] - ggml_tensor * kq_g_all = ggml_concat(ctx0, kq_sinks_g, kq_cache_g, 0); // [NS+topk, tg, H] - kq_g_all = ggml_soft_max_ext(ctx0, kq_g_all, nullptr, kq_scale, hparams.f_max_alibi_bias); + if (use_fused_attn) { + GGML_ASSERT(kl_raw_full != nullptr); + ggml_tensor * sel_idx_c = ggml_view_2d(ctx0, sel_idx, dsa_topk, tc, + sel_idx->nb[1], + (size_t) c0*sel_idx->nb[1]); + sel_idx_c = ggml_cont(ctx0, sel_idx_c); // [topk, Tc] i32 + // Selection consumed the causal mask before top-k, matching the + // unfused maskless softmax. The indexed op therefore receives no mask. + kqv_c = ggml_latent_attn_indexed_ext(ctx0, q_all_c, kl_raw_full, + sink_blk, s_lat_t, nullptr, + sel_idx_c, kv_lora_rank, 0, + kq_scale, 0.0f); // [512, Tc, H] + } else { + // Unfused fallback: CUDA-grid-safe get_rows subchunks, kept for backends + // and configs the op declines. + GGML_ASSERT(kl_full != nullptr && kq_sinks_c != nullptr); + const int64_t gather_token_chunk = + openpangu_dsa_gather_tokens_per_get_rows(dsa_topk); + for (int64_t g0 = 0; g0 < tc; g0 += gather_token_chunk) { + const int64_t tg = std::min(gather_token_chunk, tc - g0); + ggml_tensor * q_all_g = ggml_view_3d(ctx0, q_all_c, + kv_lora_rank + n_embd_head_qk_rope, + tg, n_head, q_all_c->nb[1], q_all_c->nb[2], + (size_t) g0*q_all_c->nb[1]); + q_all_g = ggml_cont(ctx0, q_all_g); + ggml_tensor * kq_sinks_g = ggml_view_3d(ctx0, kq_sinks_c, NS, tg, n_head, + kq_sinks_c->nb[1], kq_sinks_c->nb[2], + (size_t) g0*kq_sinks_c->nb[1]); + ggml_tensor * sel_idx_g = ggml_view_2d(ctx0, sel_idx, dsa_topk, tg, + sel_idx->nb[1], + (size_t) (c0 + g0)*sel_idx->nb[1]); + ggml_tensor * sel_idx_flat_g = ggml_cont_2d(ctx0, sel_idx_g, dsa_topk*tg, 1); + ggml_tensor * k_gath_g = ggml_get_rows(ctx0, kl_full, sel_idx_flat_g); // [576, topk*tg] + k_gath_g = openpangu_cast_gathered_latent_for_cache_type( + ctx0, k_gath_g, kv_self.k_l[il]); + k_gath_g = ggml_reshape_3d(ctx0, k_gath_g, + kv_lora_rank + n_embd_head_qk_rope, + dsa_topk, tg); // [576, topk, tg] + ggml_tensor * q_gath_g = ggml_cont(ctx0, + ggml_permute(ctx0, q_all_g, 0, 2, 1, 3)); // [576, H, tg] + ggml_tensor * kq_cache_g = ggml_mul_mat(ctx0, k_gath_g, q_gath_g); // [topk, H, tg] + kq_cache_g = ggml_cont(ctx0, + ggml_permute(ctx0, kq_cache_g, 0, 2, 1, 3)); // [topk, tg, H] + ggml_tensor * kq_g_all = ggml_concat(ctx0, kq_sinks_g, + kq_cache_g, 0); // [NS+topk, tg, H] + kq_g_all = ggml_soft_max_ext(ctx0, kq_g_all, nullptr, + kq_scale, hparams.f_max_alibi_bias); - ggml_tensor * kq_s_g = ggml_view_3d(ctx0, kq_g_all, NS, tg, n_head, - kq_g_all->nb[1], kq_g_all->nb[2], 0); - ggml_tensor * kq_cache_soft_g = ggml_view_3d(ctx0, kq_g_all, dsa_topk, tg, n_head, - kq_g_all->nb[1], kq_g_all->nb[2], - NS*ggml_element_size(kq_g_all)); - ggml_tensor * v_gath_g = ggml_view_3d(ctx0, k_gath_g, kv_lora_rank, dsa_topk, tg, - k_gath_g->nb[1], k_gath_g->nb[2], 0); - ggml_tensor * v_gath_t_g = ggml_cont(ctx0, ggml_permute(ctx0, v_gath_g, 1, 0, 2, 3)); // [topk, 512, tg] - ggml_tensor * kq_cache_gath_g = ggml_cont(ctx0, ggml_permute(ctx0, kq_cache_soft_g, 0, 2, 1, 3)); // [topk, H, tg] - ggml_tensor * kqv_cache_g = ggml_mul_mat(ctx0, v_gath_t_g, kq_cache_gath_g); // [512, H, tg] - kqv_cache_g = ggml_cont(ctx0, ggml_permute(ctx0, kqv_cache_g, 0, 2, 1, 3)); // [512, tg, H] - ggml_tensor * kqv_g = ggml_add(ctx0, - ggml_mul_mat(ctx0, s_lat_t, kq_s_g), - kqv_cache_g); // [512, tg, H] - kqv_c = kqv_c == nullptr ? kqv_g : ggml_concat(ctx0, kqv_c, kqv_g, 1); + ggml_tensor * kq_s_g = ggml_view_3d(ctx0, kq_g_all, NS, tg, n_head, + kq_g_all->nb[1], kq_g_all->nb[2], 0); + ggml_tensor * kq_cache_soft_g = ggml_view_3d( + ctx0, kq_g_all, dsa_topk, tg, n_head, + kq_g_all->nb[1], kq_g_all->nb[2], + NS*ggml_element_size(kq_g_all)); + ggml_tensor * v_gath_g = ggml_view_3d(ctx0, k_gath_g, + kv_lora_rank, dsa_topk, tg, + k_gath_g->nb[1], k_gath_g->nb[2], 0); + ggml_tensor * v_gath_t_g = ggml_cont(ctx0, + ggml_permute(ctx0, v_gath_g, 1, 0, 2, 3)); // [topk, 512, tg] + ggml_tensor * kq_cache_gath_g = ggml_cont(ctx0, + ggml_permute(ctx0, kq_cache_soft_g, 0, 2, 1, 3)); // [topk, H, tg] + ggml_tensor * kqv_cache_g = ggml_mul_mat( + ctx0, v_gath_t_g, kq_cache_gath_g); // [512, H, tg] + kqv_cache_g = ggml_cont(ctx0, + ggml_permute(ctx0, kqv_cache_g, 0, 2, 1, 3)); // [512, tg, H] + ggml_tensor * kqv_g = ggml_add(ctx0, + ggml_mul_mat(ctx0, s_lat_t, kq_s_g), + kqv_cache_g); // [512, tg, H] + kqv_c = kqv_c == nullptr ? kqv_g : + ggml_concat(ctx0, kqv_c, kqv_g, 1); + } } + } else if (use_fused_attn) { + // DSA unfused fallback inside the gather-owned outer loop. Dense/SWA/MTP + // fused attention bypasses the loop and uses the full-span call below. + ggml_tensor * kl_raw = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv_attn, win_off); + ggml_tensor * mask_eff = ggml_view_2d(ctx0, KQ_mask, n_kv_attn, tc, + KQ_mask->nb[1], (size_t) c0*KQ_mask->nb[1]); + if (sel_mask) { + ggml_tensor * sel_mask_c = ggml_view_2d(ctx0, sel_mask, n_kv_attn, tc, + sel_mask->nb[1], (size_t) c0*sel_mask->nb[1]); + mask_eff = ggml_add(ctx0, mask_eff, sel_mask_c); + } else if (use_dsa_sel_idx_mask) { + ggml_tensor * sel_idx_c = ggml_view_2d(ctx0, sel_idx, dsa_topk, tc, + sel_idx->nb[1], (size_t) c0*sel_idx->nb[1]); + ggml_tensor * base_c = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv_attn, tc), -1e30f); + ggml_tensor * zeros_c = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, dsa_topk, tc), 0.0f); + ggml_tensor * sel_mask_c = ggml_reshape_2d(ctx0, ggml_set_rows(ctx0, base_c, zeros_c, sel_idx_c), n_kv_attn, tc); + mask_eff = ggml_add(ctx0, mask_eff, sel_mask_c); + } + kqv_c = ggml_latent_attn_prefix_ext(ctx0, q_all_c, kl_raw, sink_blk, s_lat_t, mask_eff, + kv_lora_rank, 0, kq_scale, hparams.f_max_alibi_bias); // [512, Tc, H] } else { ggml_tensor * kq_cache_c = ggml_mul_mat(ctx0, kl_all, q_all_c); // [n_kv_attn, Tc, H] ggml_tensor * kq_c_all = ggml_concat(ctx0, kq_sinks_c, kq_cache_c, 0); // [NS+n_kv_attn, Tc, H] @@ -681,10 +831,22 @@ ggml_tensor * llm_build_context::build_openpangu_attention( } kqv = kqv == nullptr ? kqv_c : ggml_concat(ctx0, kqv, kqv_c, 1); } + } else if (use_fused_attn) { + // Fused latent attention: one op replaces [prefix|cache] QK, joint softmax, the + // two value contractions, the zero-prefix mask, and the value transpose. The op + // reads the raw (possibly q8) cache and dequantizes internally. + GGML_ASSERT(!use_dsa_sel_idx_mask || sel_mask != nullptr); + ggml_tensor * kl_raw = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv_attn, win_off); + ggml_tensor * mask_eff = ggml_view_2d(ctx0, KQ_mask, n_kv_attn, n_tokens, KQ_mask->nb[1], 0); + if (sel_mask) { + mask_eff = ggml_add(ctx0, mask_eff, sel_mask); + } + kqv = ggml_latent_attn_prefix_ext(ctx0, q_all, kl_raw, sink_blk, s_lat_t, mask_eff, + kv_lora_rank, 0, kq_scale, hparams.f_max_alibi_bias); // [512, T, H] } else { GGML_ASSERT(!use_dsa_sel_idx_mask || sel_mask != nullptr); ggml_tensor * kq_sinks = ggml_mul_mat(ctx0, sink_blk, q_all); // [NS, T, H] - kq_cache = ggml_mul_mat(ctx0, kl_all, q_all); // [n_kv_attn, T, H] + ggml_tensor * kq_cache = ggml_mul_mat(ctx0, kl_all, q_all); // [n_kv_attn, T, H] ggml_tensor * kq = ggml_concat(ctx0, kq_sinks, kq_cache, 0); // [NS+n_kv_attn, T, H] // mask: sinks always visible (0) ++ the causal/SWA KQ_mask (+ the DSA top-k selection