Auto max prefill #2797
Auto max prefill #2797
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| fn compute_optimal(config: Option<&Config>, compute: Option<&ComputeType>) -> Option<usize> { | ||
| if let (Some(config), Some(compute)) = (config, compute) { | ||
| if let (Some(f16_max_compute), Some(model_compute)) = (compute.f16_flop(), config.flop()) { | ||
| tracing::debug!("MAx compute {f16_max_compute} model compute {model_compute}"); |
There was a problem hiding this comment.
Nit:
tracing::debug!("Max compute {f16_max_compute} model compute {model_compute}");
| let q_flops = 2 * num_heads * head_dim * hidden_size; | ||
| let k_flops = 2 * num_kv_heads * head_dim * hidden_size; | ||
| let v_flops = 2 * num_kv_heads * head_dim * hidden_size; | ||
| let attn_flops = 2 * num_heads * head_dim * hidden_size; | ||
| let o_flops = 2 * num_heads * head_dim * hidden_size; |
There was a problem hiding this comment.
We should probably adjust this for FP8 on >=9.0. (For <9.0, this is still fine, since FP8-Marlin computes in FP16).
There was a problem hiding this comment.
No it should be OK. Since I use the flops on hardware in f16 already.
Both are divided by 2 for fp8 so the calculation should hold, it's just expressed in f16 unit.
| let head_dim = self.head_dim? as u64; | ||
| let hidden_size = self.hidden_size? as u64; | ||
| let intermediate_size = if let Some(experts) = self.experts { | ||
| (self.intermediate_size? * experts) as u64 |
There was a problem hiding this comment.
In the case of Deepseek v2, we should do: (num_experts_per_tok + n_shared_experts) * moe_intermediate_size
(Or put differently, Phi 3.5 MoE and Mixtral are n_shared_experts == 0 && moe_intermediate_size == intermediate_size.)
To make it more fun, the first layer in Deepseek v2 does not use MoE (but the layer is the same size practically, since moe_intermediate_size is intermediate_size / (num_experts_per_tok + n_shared_experts).
What does this PR do?
Fixes # (issue)
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.