CPUAttention.cpp 和 KVCacheManager.cpp 问题

[DakeQQ]

您好，MNN团队，

我们正在尝试复现MNN-LLM项目，特别是关注llm.cpp中的源代码。我们已使用llmexport.py导出了MNN模型，并模仿了llm.cpp中的加载和推理过程。

// 加载部分
mMeta = std::make_shared<KVMeta>();
runtimeManager->setHint(MNN::Interpreter::HintMode::QKV_QUANT_OPTIONS, 0);   // Turn it off to ensure no accuracy loss during the test.
runtimeManager->setHintPtr(MNN::Interpreter::HintMode::KVCACHE_INFO, mMeta.get());

// 推理部分
mMeta->add = ids_len;
std::vector<MNN::Express::VARP> outputs = module_A->onForward({input_ids, attention_mask, position_ids, logits_index});
mMeta->sync();

我们导出的MNN模型是Qwen2.5-0.5B，包含MNN/transformers/llm/export/utils/custom.py 定义的 Attention算子，使用了8位量化。然而，在推理过程中，只有第一个输出单词是正确的，后续的单词都不正确。这似乎是因为KV缓存没有正常工作。我们已深入研究了CPUAttention.cpp和KVCacheManager.cpp，但仍然没有头绪。

llm.cpp中是否有特定细节或关键之处使得KV缓存能正常工作？虽然我们已经仔细模仿了llm.cpp主要的推理代码，但仍未能成功复现。

任何您提供的指导都会非常感谢: ) 谢谢~

[jxt1234]

看上去只有 prefill 的过程，decode 的推理代码是?

[DakeQQ]

感恩指导～

decode重复使用了这个代码，

mMeta->add = ids_len;
std::vector<MNN::Express::VARP> outputs = module_A->onForward({input_ids, attention_mask, position_ids, logits_index});
mMeta->sync();

decode時候的输入是： input_ids = embedding(最大概率ids)， attention_mask , position_ids 则是根据seq_len=1时候的get_attention_mask() , get_postion_ids()代码

decode需要特别的处理吗？我简单的理解decode跟prefill主要差异在是seq_len=1，position_ids要跟着吐字数累加， attention_mask没有作用=0

[jxt1234]

decode 的时候一般 mMeta->add = 1 （也就是 ids_len = 1）

[DakeQQ]

嗯嗯是的，decode阶段这个add也确实保持在+1，但第二个以後输出单词就完全不对。

我们尝试过一直使用prefill的方式喂给onForward()，这样也能吐出正确的单词，只是速度慢很多。我想表达的是整个算子计算应该是正确的，所以猜测kv cache没正常工作。请问kv cache在管理过程中，还有哪里需要关注的地方？或是您认为问题不出在这儿呢？

[jxt1234]

感觉没有问题，可能是其他部分的代码有误？ 可以把 mnn-llm 输入的 varp 打印出来看一下

[github-actions[bot]]

Marking as stale. No activity in 60 days.