[Feature]: Thresholds for enabling Sequence Parallelism + Async TP

[jasonlizhengjian]

🚀 The feature, motivation and pitch

SP + Async TP is beneficial on long prefills / big batches but can be detrimental below a certain threshold. With the enable of dynamic sizes in #26681 , this introduces regressions for decode.

Therefore, if SP + Async TP is to be enabled by default, we should have some sort of thresholding for when it is active / inactive.

Alternatives

No response

Additional context

Data:

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[ProExpertProg]

cc @angelayi @zou3519 @ilmarkov

[jasonlizhengjian]

I collected some more benchmarks and did some data analysis. I think that the threshold can be based off tensor size as calculated by dtype_bytes * hidden_size * num_tokens / tp_size

Here is the analysis for H100:

40 MB is a good general threshold and we could have 32 MB as a separate threshold for fp8.

[ProExpertProg]

Thanks for this! We should probably add a configurable threshold for this as well so users can tune it themselves. And it should be dependent on TP size.

Questions/notes on the plot:

• I assume the cutoff colors are inverted?
• What's the difference between different points that stack on top of the same tensor size (x-axis)? Could you post them as different series for different settings?
• How did you collect these numbers? Did you measure a single forward pass latency for a given number of tokens? Or something else?

[ProExpertProg]

It seems to me that even below the cutoff, we sometimes get up to a 40% speedup, which is nothing to scoff at, so it would be better if we could have a multi-dimensional heuristic (or come up with a better single-d metric than just per-gpu tensor size)

[jasonlizhengjian]

So with regard to the plot:

• Cutoff colors are inverted, these are just some colors automatically selected by google sheets
• The data collected is using vllm bench latency with varying input sizes (1k , 4k, 8k, 16k) and various batch sizes for 1k and 4k. Output len is 1. This is part of the reason why points stack up (i.e 4k * 2 = 8k) but these are mostly pretty close together
• The other part is that models of varying sizes were tested. For example 4k hidden size vs 8k hidden size. These are often pretty far apart (2x num tokens but with 4k hidden is much worse speedup than 8k hidden)

For the threshold metric, I think that a multi-dimensional threshold may work better but not sure how we could combine that with user-specified overrides. Maybe the user can specify a num tokens threshold which overrides the heuristic.

Here are some other graphs:

Colored by models:

Colored by hidden size:

Also attached is the raw data

H100_async_tp_benchmarks.csv

[ProExpertProg]

For the threshold metric, I think that a multi-dimensional threshold may work better but not sure how we could combine that with user-specified overrides. Maybe the user can specify a num tokens threshold which overrides the heuristic.

Yeah I think the user should specify the num_tokens threshold because that's also what will determine runtime dispatching and conditional compilation. By default, we can compute the num_tokens cutoff heuristically based on other parameters which will make the cutoff multi-dimensional.

[jasonlizhengjian]

After some further analysis. My proposed multi-d threshold would be:

if hidden_size >= 8192:
   per_gpu_tensor_size_threshold = 8
else:
   per_gpu_tensor_size_threshold = inf

The data shows that Async TP only has major perf benefits for models Llama 70B and larger so we can capture most of the upside without risking perf regressions.

Could also replace 8192 with 5120 but that seems like a pretty niche use case as I'm not aware of many relevant models which have that hidden size.