ShadowKVCache and ShadowKVCache_CPU produce different logits

Hello! Thanks for your work. I have multiple questions regarding efficiency of your implementation. Firstly you report only [generation throughput](https://github.com/ByteDance-Seed/ShadowKV/blob/e51904cdeab7d4d34013370f09f2cf5fcd655e15/test/e2e.py#L166) and other metrics [here](https://github.com/ByteDance-Seed/ShadowKV/blob/e51904cdeab7d4d34013370f09f2cf5fcd655e15/models/base.py#L335) such as latency. But you compute latency as different between [end](https://github.com/ByteDance-Seed/ShadowKV/blob/e51904cdeab7d4d34013370f09f2cf5fcd655e15/models/base.py#L334) and [start](https://github.com/ByteDance-Seed/ShadowKV/blob/e51904cdeab7d4d34013370f09f2cf5fcd655e15/models/base.py#L324) markers that cover only generation part. Thus naming e2e (end to end) that is often used to represent prefill and decode latency is not correct and in your implementation refers to only generation part. I decided to measure true e2e. I found out that it was very slow. I decided to take profiles using nvidia nsight systems and found a lot of bubbles in the prefill stage

It is not important to look too deep at this profile. The most important part is those CPU bound bubbles where GPU is idle. After debugging I found out [this](https://github.com/ByteDance-Seed/ShadowKV/blob/e51904cdeab7d4d34013370f09f2cf5fcd655e15/models/kv_cache.py#L494-L497) code where you basically call garbage collector inside every [get_svd](https://github.com/ByteDance-Seed/ShadowKV/blob/e51904cdeab7d4d34013370f09f2cf5fcd655e15/models/kv_cache.py#L480) call which triggers a lot of CPU work and makes GPU idle. I removed this part and bubbles disappeared. I did sanity check by computing benchmark via your [scripts](https://github.com/ByteDance-Seed/ShadowKV/blob/e51904cdeab7d4d34013370f09f2cf5fcd655e15/README.md?plain=1#L92) on some of the datasets and metrics did not changed. So I make conclusion that this code is not used for any algorithm logic and thus bug or something that you have used while debugging and forgot to remove?

After I removed that profile started to look like this. I also added some nvtx ranges

But prefill latency was still large. And SVD was relatively large compared to FA (flash attention). I did some research and investigated that in your code you cast input data to the [torch.svd](https://github.com/ByteDance-Seed/ShadowKV/blob/e51904cdeab7d4d34013370f09f2cf5fcd655e15/models/kv_cache.py#L499) to fp32 which does make sense since [torch.svd only works with floats and doubles](https://docs.pytorch.org/docs/stable/generated/torch.svd.html). The problem with fp32 datatype is that modern GPUs such as A100 and H100 has relatively small number of CUDA cores compared to tensor cores. In fact I used H100 in my experiments which has [FP32 67 TFLOPS and BF16 1979 TFLOPS](https://www.megware.com/fileadmin/user_upload/LandingPage%20NVIDIA/nvidia-h100-datasheet.pdf) which ~ 2 orders of magnitude. This proves that the prefill stage is bounded by SVD computation because of using cuda cores and not tensor cores and it bottleneck. I am interested why does your paper in appendix 6 table 12 reports that for sequence length 64k which I have used SVD takes 17ms and FA 186ms. You have used A100 while I am using H100 this is why my FA in profiles looks 2x faster and I am okay with that but my SVD looks 10x slower and I am interested how did you measures numbers in your table. Can you provide code for it please? Did you use custom SVD implementation and can you share this then?