| 5431 | // ggml_ssm_scan |
| 5432 | |
| 5433 | struct ggml_tensor * ggml_ssm_scan( |
| 5434 | struct ggml_context * ctx, |
| 5435 | struct ggml_tensor * s, |
| 5436 | struct ggml_tensor * x, |
| 5437 | struct ggml_tensor * dt, |
| 5438 | struct ggml_tensor * A, |
| 5439 | struct ggml_tensor * B, |
| 5440 | struct ggml_tensor * C, |
| 5441 | struct ggml_tensor * ids) { |
| 5442 | GGML_ASSERT(ggml_is_contiguous(s)); |
| 5443 | GGML_ASSERT(ggml_is_contiguous(dt)); |
| 5444 | GGML_ASSERT(ggml_is_contiguous(A)); |
| 5445 | GGML_ASSERT(x->nb[0] == ggml_type_size(x->type)); |
| 5446 | GGML_ASSERT(B->nb[0] == ggml_type_size(B->type)); |
| 5447 | GGML_ASSERT(C->nb[0] == ggml_type_size(C->type)); |
| 5448 | GGML_ASSERT(x->nb[1] == x->ne[0]*x->nb[0]); |
| 5449 | GGML_ASSERT(B->nb[1] == B->ne[0]*B->nb[0]); |
| 5450 | GGML_ASSERT(C->nb[1] == C->ne[0]*C->nb[0]); |
| 5451 | GGML_ASSERT(ggml_are_same_shape(B, C)); |
| 5452 | GGML_ASSERT(ids->type == GGML_TYPE_I32); |
| 5453 | |
| 5454 | { |
| 5455 | const int64_t d_state = s->ne[0]; |
| 5456 | const int64_t head_dim = x->ne[0]; |
| 5457 | const int64_t n_head = x->ne[1]; |
| 5458 | const int64_t n_seq_tokens = x->ne[2]; |
| 5459 | const int64_t n_seqs = x->ne[3]; |
| 5460 | |
| 5461 | GGML_ASSERT(dt->ne[0] == n_head); |
| 5462 | GGML_ASSERT(dt->ne[1] == n_seq_tokens); |
| 5463 | GGML_ASSERT(dt->ne[2] == n_seqs); |
| 5464 | GGML_ASSERT(ggml_is_3d(dt)); |
| 5465 | GGML_ASSERT(s->ne[1] == head_dim); |
| 5466 | GGML_ASSERT(s->ne[2] == n_head); |
| 5467 | GGML_ASSERT(B->ne[0] == d_state); |
| 5468 | GGML_ASSERT(B->ne[2] == n_seq_tokens); |
| 5469 | GGML_ASSERT(B->ne[3] == n_seqs); |
| 5470 | GGML_ASSERT(ids->ne[0] == n_seqs); |
| 5471 | GGML_ASSERT(ggml_is_vector(ids)); |
| 5472 | GGML_ASSERT(A->ne[1] == n_head); |
| 5473 | GGML_ASSERT(ggml_is_matrix(A)); |
| 5474 | |
| 5475 | if (A->ne[0] != 1) { |
| 5476 | // Mamba-1 has more granular decay factors |
| 5477 | GGML_ASSERT(A->ne[0] == d_state); |
| 5478 | } |
| 5479 | } |
| 5480 | |
| 5481 | // concatenated y + ssm_states |
| 5482 | struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); |
| 5483 | |
| 5484 | result->op = GGML_OP_SSM_SCAN; |
| 5485 | result->src[0] = s; |
| 5486 | result->src[1] = x; |
| 5487 | result->src[2] = dt; |
| 5488 | result->src[3] = A; |
| 5489 | result->src[4] = B; |
| 5490 | result->src[5] = C; |