Question: Why is there is a big difference in Token throughput for Deepseek R1?

[arul-lm]

Thanks for providing an open source implementation of Deepseek R1. It is simple and easy to understand.

The benchmark results you have shared for TPU-v5e-64 logs throughput at 71 tokens/sec.

https://github.com/jax-ml/jax-llm-examples/tree/main/deepseek_r1_jax#inference-performance-results

The LMSys + SGLang implementation of Deepseek R1 places the token throughput per GPU at 5600 tokens/sec/GPU

1. Why is there a significant difference in throughput? (Just first order approximation is enough)
2. Is your throughput number tokens/sec or tokens/sec/TPU

[rdyro]

The benchmarks presented in this repo are aimed at the low-latency setting with the goal of minimizing inter-token latency, not the token throughput.

It's entirely possible to optimize the shardings for maximum token throughput instead!

[carlesoctav]

What’s the difference besides more optimized kernels? When I look at TPU inference, it feels like the sharding is basically the same. We mainly use tensor parallelism, for column-parallel layers we do an all-gather on input_ids (since they’re sharded across TP), and for row-parallel layers we do an all-reduce instead of a reduce-scatter. Is that right?