KDA MTP (Multi-Token Prediction) support

[icavan]

Description

Add Multi-Token Prediction (MTP) support to cuLA's inference kernels.

Context

MTP is an inference optimization technique that predicts multiple tokens simultaneously, improving throughput for autoregressive generation. Supporting MTP in cuLA's linear attention kernels would enable faster inference for models using this technique.

Tasks

• Design MTP integration for linear attention inference kernels
• Implement MTP support in relevant kernels (KDA, Lightning Attention)
• Add tests and benchmarks
• Document usage

[icavan]

Actually we need a KDA decode kernel first OTL

[zhouaihui]

Hi, I'd like to work on this issue. Could you assign it to me?