mla : make the V tensor a view of K (#18986)
* mla : pass V as a view of K to the FA op * cuda : adjust mla logic to new layout * kv-cache : fix rope shift * tests : remove comment * cuda : fix reusable_cutoff Co-authored-by: Johannes Gäßler <johannesg@5d6.de> --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
This commit is contained in:
co-authored by
Johannes Gäßler
parent
e2baf02162
commit
a5eaa1d6a3
@@ -778,12 +778,15 @@ void launch_fattn(
|
||||
) {
|
||||
constexpr int ncols = ncols1 * ncols2;
|
||||
|
||||
const bool is_mla = DV == 512; // TODO better parameterization
|
||||
|
||||
const ggml_tensor * Q = dst->src[0];
|
||||
const ggml_tensor * K = dst->src[1];
|
||||
const ggml_tensor * V = dst->src[2];
|
||||
|
||||
// TODO: make this more generic by removing the notion of "MLA".
|
||||
// for example "is V a view of K?" so we can skip loading it.
|
||||
// V strides should be driven by V itself and avoid assumption of the data layout
|
||||
const bool is_mla = V->op == GGML_OP_VIEW && V->src[0] == K;
|
||||
|
||||
GGML_ASSERT(V || is_mla);
|
||||
|
||||
const ggml_tensor * mask = dst->src[3];
|
||||
|
||||
Reference in New Issue
Block a user