Enable running PyTorch models

[masahi]

This is the first cut toward running PT models under mlc-serve. Things are functional but rough. All we need to do to run a PT model is to replace --local-id llama-2-13b-chat-hf-q0f16 with --local-id models/llama-2-13b-chat-hf (the path to the HF model directory) etc.

• I haven’t verified that it supports all sampler features that we recently added (logprobs, json etc). Parallel sampling works but evicting parallel-sampling request requires model change.
• Right now it only supports vLLM models. They are directly loaded via from vllm.model_executor.models.llama import LlamaForCausalLM etc. Models such as phi-2, Qwen run out of the box, in addition to llama-based models and Mixtral that the TVM model also supports. vLLM needs to be installed from https://github.com/octoml/vllm/tree/for-mlc-serve
• Single-gpu perf is not good. For llama 13B, using benchmark_throughput.py, TVM gets 22.98 requests/s while PT gets 13.78 requests/s. Since TVM and vLLM models should be using the same kernels for matmul and attention, there shouldn't be a big difference in performance. I haven't looked into this issue deeply.
• Multi-gpu performance is terrible. Currently there is a huge overhead in serialization of inputs from main to worker processes, which can be seen below (large poll between two decode).
  

@sunggg @yelite @elvin-n @vvchernov @jroesch @binarybana

[sunggg]

Great work, Masa! A couple comments.

[sunggg]

Now that we started considering vllm's tensor layout, what do you think about unifying it? It seems like upstream mlc-llm also uses 2D inputs.

[sunggg]

And this also could help our cuda graph integration?

[masahi]

We haven't verified if 2D inputs is better for performance, and how much cuda graph actually helps.

The upstream input looks like 2D but it is always either (1, num_total_token) or (batch_size, 1). So their 2D input is essentially 1D.

[sunggg]

Yeah, I think it is worth visiting imo. But not now, in the future. Although there might not be performance boost, it would be nice to unify the layout with upstream unless there is reason.

[masahi]

Oh I just realized that, after I updated our vLLM work, which made their input representation 2D, memory profiling is broken.

Right now we only have max_num_batched_tokens as a parameter to memory profiling, and "the biggest input" to memory profiling is a batch of max_num_batched_tokens sequences each of which is length 1. But with 2D input, we also need max_seq_len param, since prefill inputs are padded to the max seq len in a batch. So memory profiling severely underestimates memory usage.

We could look into how vLLM does memory profiling currently, but I'm inclined to revert their 2D rep change instead. Either way, this is left for future work.

[masahi]

I opened a ticket to track many TODO items after this PR. #217

I think this is ready to merge for now @sunggg