LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

openpangu: fused latent attention op (GGML_OP_LATENT_ATTN) (#2168)

Browse Source 
Adds ggml_latent_attn_prefix_ext / ggml_latent_attn_indexed_ext: MLA
latent-cache attention with an always-visible learned K/V prefix
(openPangu's 128 param_sink rows), joint softmax over [prefix | cache],
reading the raw F32/F16/Q8_0 latent cache directly. CUDA implementation
plus a scalar CPU reference that pins the op's semantics; the CPU
backend reports support truthfully, and openPangu adopts the op only on
a non-CPU backend as builder policy.

openPangu routes its dense/SWA/MTP full-span attention and the gathered
DSA path through the op, capability-gated per layer on the attention
output projection's scheduled backend, with the latent cache required
resident on that same backend (--no-kv-offload keeps the unfused
chain); any layer whose backend cannot run the candidate keeps the
exact unfused chain.

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
This commit is contained in:
Joel Farthing
2026-07-23 14:50:45 +03:00
committed by GitHub
co-authored by Joel Farthing
parent e5357286c0
commit 31018dc511
9 changed files with 1417 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/include/ggml.h
+48
View File
@@ -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,
ggml/src/ggml-backend.cpp
+5
View File
@@ -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;
}
ggml/src/ggml-cuda.cu
+7 -1
View File
@@ -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;
ggml/src/ggml-cuda/latent_attn.cu
+698
View File
@@ -0,0 +1,698 @@
#include "common.cuh"
#include "convert.cuh"
#include "latent_attn.cuh"
#include <climits>
// 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 <typename Tout, bool Restore>
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 <typename Tout>
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 <typename T>
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 <typename Tout, bool Restore>
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<int64_t>(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<float> pk_f32(ctx.pool()), pv_f32(ctx.pool());
ggml_cuda_pool_alloc<half> 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<float> scores(ctx.pool(), (int64_t) M*tile_cols);
ggml_cuda_pool_alloc<float> outbuf(ctx.pool(), (int64_t) dv*tile_cols);
ggml_cuda_pool_alloc<float> q_f32(ctx.pool()), gathered_f32(ctx.pool()), w_f32(ctx.pool());
ggml_cuda_pool_alloc<half> q_f16(ctx.pool()), gathered_f16(ctx.pool()), w_f16(ctx.pool());
ggml_cuda_pool_alloc<float> 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<float><<<cols, nth, 0, stream>>>(
(const float *) q->data, q_f32.get(), Dk, H, first, rows, q_nb1_f, q_nb2_f);
k_latent_gather_rows_indexed<float><<<rows*topk, nth, 0, stream>>>(
(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<<<cols, nth, softmax_shmem, stream>>>(
(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<half><<<rows*topk, nth, 0, stream>>>(
(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<<<rows*topk, nth, 0, stream>>>(
(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<float, false><<<cols, nth, softmax_shmem, stream>>>(
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<half, true><<<cols, nth, softmax_shmem, stream>>>(
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<half> 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<float> pk_f32(ctx.pool()), pv_f32(ctx.pool());
ggml_cuda_pool_alloc<half> 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<float> scores(ctx.pool(), PN*cw);
ggml_cuda_pool_alloc<float> outbuf(ctx.pool(), (int64_t) dv*cw);
ggml_cuda_pool_alloc<half> q_f16(ctx.pool()); // f16 path: range-scaled q tile
ggml_cuda_pool_alloc<half> w_f16(ctx.pool()); // f16 path: normalized weights
ggml_cuda_pool_alloc<float> w_f32(ctx.pool()); // f32 path: weights (== softmax out)
ggml_cuda_pool_alloc<float> 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<int64_t>(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<<<cwn, sm_nth, sm_shmem, stream>>>(
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<half, true><<<cwn, sm_nth, sm_shmem, stream>>>(
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<float, false><<<cwn, sm_nth, sm_shmem, stream>>>(
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));
}
}
ggml/src/ggml-cuda/latent_attn.cuh
+4
View File
@@ -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);
ggml/src/ggml-cuda/vendors/hip.h
Vendored
+1
View File
@@ -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
ggml/src/ggml-cuda/vendors/musa.h
Vendored
+1
View File
@@ -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
ggml/src/ggml.c
+398 -2
View File
@@ -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];
src/graphs/build_openpangu.cpp
+255 -93
View File
@@ -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<int64_t>(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<int64_t>(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