GitLab

[Core] Cross-attention KV caching and memory-management (towards eventual encoder/decoder model support) (#4837)

Signed-off-by: pandyamarut <pandyamarut@gmail.com>

Co-authored-by: Roger Wang <ywang@roblox.com>

This PR adds Triton kernel configs for the MoE kernel for MI300X

Co-authored-by: Ruth Evans <ruthevans@Ruths-MacBook-Pro.local>

Co-authored-by: rsnm2 <rshaw@neuralmagic.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-neuralmagic@users.noreply.github.com>

Co-authored-by: beagleski <yunanzhang@microsoft.com>
Co-authored-by: bapatra <bapatra@microsoft.com>
Co-authored-by: Barun Patra <codedecde@users.noreply.github.com>
Co-authored-by: Michael Goin <michael@neuralmagic.com>

Co-authored-by: Lei Wen <wenlei03@qiyi.com>

Co-authored-by: Cody Yu <hao.yu.cody@gmail.com>

Co-authored-by: Michael Goin <michael@neuralmagic.com>

Co-authored-by: Elisei Smirnov <el.smirnov@innopolis.university>

Co-authored-by: Varun Sundar Rabindranath <varunsundar08@gmail.com>
Co-authored-by: Varun Sundar Rabindranath <varun@neuralmagic.com>

Signed-off-by: Muralidhar Andoorveedu <muralidhar.andoorveedu@centml.ai>

The 2nd PR for #4532.

This PR supports loading FP8 kv-cache scaling factors from a FP8 checkpoint (with .kv_scale parameter).

Pass the CUDA stream into the CUTLASS GEMMs, to avoid future issues with CUDA graphs