feat(llama.cu): 适配 mamba 模型

[CearX]

No description provided.

[CearX]

启元人工智能大赛 推理系统报告

1. 参与赛道

九源大模型推理引擎开发

• 参与了模型适配赛题 选用 InfiniLM-Rust 框架，适配 T2-1-3：Mamba模型
• 核心目标：实现非 Transformer 架构模型的高效推理支持

2. 成果阐述

本实现基于 InfiniLM-Rust 框架 对 Mamba 模型（非 Transformer 架构） 进行了适配与优化工作。通过新增模型结构支持、算子实现及缓存优化，显著提升了推理效率。在 NVIDIA A800 硬件上, 并发 245 次请求，达到 633.66 tokens/s 的推理速度，并在 Wikitext 数据集上取得19.65 PPL 的优异表现，本实现的 F16 精度模型与 PyTorch F32 基准模型的 PPL 相对误差仅为4e-5，远低于大赛评审 1% 的误差要求，满足大赛满分标准。本工作为高效部署状态空间模型（SSM）提供了可行方案，并为后续支持 Chat 模式提供兼容性基础。

2.1 模型适配与功能扩展

• 新增模型结构：突破原框架仅支持 Transformer 模型的限制，完整支持了 Mamba 的 SSM（状态空间模型）架构。
• 关键算子实现：新增causal-conv1d、selective-scan、silu、element-mul等核心算子，确保模型计算完整性。
• 服务层优化：重构引擎接口，支持 OpenAI 标准 API 测试 PPL（困惑度），为后续 Chat 模式提供兼容性基础。

2.2 性能验证

采用 Mamba_F16_2.8B 模型，在 Wikitext 数据集上进行测试，与 HF 的 Mamba_F32_2.8B 模型 PyTorch 实现进行精度对比：

硬件平台	推理速度 (tokens/s)	PPL (Wikitext-61 samples)	相对误差
NVIDIA A800	633.66	19.65	4e-5

PPL 值 19.65 表明模型在语言建模任务中具有稳定表现，F16 模型与 F32 基准模型的相对误差仅 4e-5，完全满足大赛满分要求

3. 技术亮点

3.1 高效缓存管理

• 虚存优化缓存机制：设计基于虚存的mamba cache，通过捕获conv和scan算子动态插入缓存操作，显著减少内存冗余访问。
• 架构突破：突破原框架仅支持 KV Cache 的限制，扩展至 Mamba 特有缓存模式，实现高效序列处理。

3.2 引擎层与服务层重构

• ​​Mamba引擎重构：重构引擎逻辑，支持灵活接入不同模型类型，后续可支持 Chat 模式。
• 标准化接口设计：通过 OpenAI 兼容接口提供 PPL 测试服务，支持外部应用无缝调用。

4. 性能结果

4.1 测试命令

1. 推理测试：
   srun --nodelist=gpua20 cargo test --release --package llama-cu --lib -- exec::mamba::tests::test_mamba_infer_decode --exact --show-output
2. 服务启动：
   srun --nodelist=gpua20 cargo run -p xtask -- mamba-service "/home/cearx/Mamba-2.8B-hf-v1.0-F16.gguf" --port 8080 --gpus "0" --name "model"
3. 并发测试：
   TEST_PORT=8080 srun --nodelist=gpua20 cargo test --release test_post_send_multi -- --nocapture
4. PPL测试：
   srun --nodelist=gpua20 python3 test_ppl.py --api-base http://127.0.0.1:8080 --model model --config wikitext-2-raw-v1 --split 'test[:100]' --max-samples 100

• 4.2 性能数据与可视化

测试基于 Mamba_F16_2.8B 模型，与 HF Mamba_F32_2.8B 模型的 Pytorch 实现对比 PPL

• NVIDIA A800性能截图：
• 图1：A800 单次推理示例，max_token = 1024, 推理速度 109.14 tokens/s
  
• 图2：A800 并发 245 次请求，推理速度 633.66 tokens/s
  
• 图3：Wikitext 测试 PPL=19.65（61 samples）

5. 时间线

阶段	时间周期	核心任务
知识学习与技术调研	8月04日-8月11日	研究 Mamba 模型结构、SSM 原理及 Rust 框架集成方案
编码实现	8月11日-8月17日	完成算子添加、缓存模块开发、引擎适配层实现
调试与验证	8月17日-8月24日	完成API接口扩展、性能测试、PPL验证、多硬件兼容性检查

6. 应用场景与未来展望

6.1 可行性应用场景

• 边缘设备部署：Mamba 的 SSM 架构通过​​线性计算复杂度​​显著降低计算开销与内存占用，结合本优化方案，适配边缘设备有限算力（如英特尔 Core Ultra NPU）的实时推理需求。
• 长序列处理：SSM架构的​​选择性扫描机制​​（selective-scan）可高效捕捉长程依赖关系，在基因组分析、金融时序预测等长上下文任务中具显著优势。

6.2 未来工作

• 扩展 Chat 模式支持：基于现有 API 接口实现对话式交互功能。
• 量化优化：引入 INT8 量化进一步提升推理速度与能效比。
• 国产芯片适配： 持续推进天数沐曦等国产芯片的适配，通过软硬件协同优化提升推理性能。

[github-advanced-security]

clippy found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.