Skip to content

[performance] InferEngine: drastically improve initialization speed#481

Merged
wooway777 merged 2 commits into
InfiniTensor:mainfrom
rubik-hua:InferEngine_performance
Jul 8, 2026
Merged

[performance] InferEngine: drastically improve initialization speed#481
wooway777 merged 2 commits into
InfiniTensor:mainfrom
rubik-hua:InferEngine_performance

Conversation

@rubik-hua

Copy link
Copy Markdown
Contributor
  • Replace sync set_zeros with async set_zeros_device_async to eliminate CPU buffer allocation and H2D overhead.
  • Decouple RankWorker creation from initialization wait to allow parallel startup.
  • Add explicit synchronization after KV cache allocation to ensure readiness.

提升InferEngine初始化时间,将rank_worker并行化

InferEngine初始化非常慢,尤其是在多tp场景下,更加恶化。
主要两个原因,1、显存分配多卡是串行的 2、kv-cache清零操作的时候额外写了内存,再从内存搬迁到GPU显存

以下几个图说明,显存分配的串行过程
image
image
image
image

然后在csrc/engine/infer_engine.cpp的InferEngine构造函数中增加如下时间记录打印
image

优化对比测试:
沐曦环境 InferEngine初始化时间能从157.34秒缩短到28.54秒

沐曦环境优化前
image
沐曦环境优化后
image

摩尔环境初始化时间能从11.93秒缩短到4.47秒
摩尔环境优化前
python examples/test_infer.py --device moore --model /home/rubik/models/Mistral-Large-Instruct-2411 --tp 8 --enable-paged-attn --prompt 'introduce yourself'
image
摩尔环境优化后
image

优化后显存分配是多卡并行的,参考以下几个图:
image
image
image
image

不影响模型的正常推理。
摩尔环境验证了Mistral-Large-Instruct-2411
image

沐曦环境也验证了Mistral-Large-Instruct-2411
image

N卡验证了所支持的大部分模型。
image
image
image
image
image
image

- Replace sync `set_zeros` with async `set_zeros_device_async` to eliminate CPU buffer allocation and H2D overhead.
- Decouple RankWorker creation from initialization wait to allow parallel startup.
- Add explicit synchronization after KV cache allocation to ensure readiness.
@rubik-hua rubik-hua requested a review from a team July 7, 2026 02:27
@pengcheng888 pengcheng888 requested a review from wooway777 July 7, 2026 02:58
@wooway777

Copy link
Copy Markdown
Collaborator

/test

@wooway777

Copy link
Copy Markdown
Collaborator

需要解conflict以及用scripts/format.py格式化

Comment thread csrc/cache/kv_cache.cpp Outdated
@rubik-hua rubik-hua force-pushed the InferEngine_performance branch from 97ae62a to b480569 Compare July 7, 2026 09:31
@rubik-hua rubik-hua force-pushed the InferEngine_performance branch from b480569 to 2bc048c Compare July 8, 2026 06:34
@wooway777 wooway777 merged commit 4308d54 into InfiniTensor:main Jul 8, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants