How to use Tensor Core to accelerate Speculative Sampling?

[zhaoyang-star]

When using KV Cache, the input ids shape becomes [1, seq_len], where:

• seq_len could only be 1 or 2 for draft model
• seq_len could only be gamma+1 for target model

At this time, we could only use CUDA Core rather than Tensor Core to accelerate the decoding kernel for draft model inference.

As tensor core is much faster, is there any method to make seq_len larger to use tensor core?

[feifeibear]

Sorry for my late response. Decoding is a notoriously memory-bound problem, using tensor core sometimes is a waste.

See this question on zhihu.com (in chinese) https://www.zhihu.com/question/636533414/answer/3345355574