Specscheduler evaluation support code

[goliaro]

Description of changes:

This PR does the following:

• LLAMA 3 speculation support:

  • Add support for LLAMA 3.1 and 3.2
  • Benchmark performance of LLAMA-3.1-70B with small models: Zhuominc/Llama-3-330M, meta-llama/Llama-3.2-1B-Instruct, meta-llama/Llama-3.2-3B-Instruct, meta-llama/Llama-3.1-8B-Instruct (tl;dr meta-llama/Llama-3.2-1B-Instruct is the best)
  • Add support for serving SSMs with TP_degre > 1

• Benchmarking

  • Gather fine-grained performance statistics allowing us to compute additional metrics later, e.g. P90 latency
  • Fix issue causing to underestimate average number of generation tokens per step
  • Added code to benchmark speculation accuracy and end-to-end performance for specinfer and incr decoding with various SSMs and arrival rates.
  • Plots available below:
    average_accepted_tokens.pdf
    throughput_vs_tpot.pdf
    ttft_vs_arrival_rate.pdf
    queueing_time_vs_arrival_rate.pdf

• Make evaluation easier/faster to run:

  • Add code to load all the weights in parallel, fixing context issue discussed with Legion team here
  • Record memory usage breakdown when passing --log-instance-creation. Add script to debug issues related to insufficient memory by device and task. See here.

• Bug fixes

  • Remove all reduce deadlock by adding Legion barriers
  • Detection of EOS tokens when produced in the middle of speculation (instead of at the end) and early stop to prevent infinite generation (until max sequence length) when the EOS token is in middle of verified sequence

Related Issues:

Linked Issues:

• Issue #

Issues closed by this PR:

• Closes #

---

This change is