[llama] Produce real values to run the compiled vmfb on (for testing/benchmarking)

[ScottTodd]

Some parts of the llama model divide by input values, so providing empty inputs can crash the program or produce NaNs. We should generate real input values from tokenized prompts for testing and benchmarking.

Signature for quantized model we're focusing on:

• @prefill_bs2(%arg0: !torch.vtensor<[2,?],si64>, %arg1: !torch.vtensor<[2],si64>, %arg2: !torch.vtensor<[2,?],si64>, %arg3: !torch.tensor<[?,65536],f16>) -> !torch.vtensor<[2,?,32000],f16>
• @decode_bs2(%arg0: !torch.vtensor<[2,1],si64>, %arg1: !torch.vtensor<[2],si64>, %arg2: !torch.vtensor<[2],si64>, %arg3: !torch.vtensor<[2,?],si64>, %arg4: !torch.tensor<[?,65536],f16>)

Python code to start from is in https://github.com/nod-ai/sharktank/blob/main/shortfin/shortfin/llm/impl/service_v1_cli.py . I've added logging there and in https://github.com/nod-ai/sharktank/blob/main/shortfin/shortfin/llm/impl/service_v1.py for debugging, but I'd like a way to export to .bin (or .npy) files for use with iree-run-module and iree-benchmark-module. Maybe a python script for prompt -> input files? Could also have a script for validating a prompt -> reference model and prompt -> compiled model.

[ScottTodd]

This could be useful: https://github.com/nod-ai/sharktank/blob/main/sharktank/sharktank/models/punet/tools/sample_data.py

[ScottTodd]

Trying to write a standalone script but it might also make sense to teach service_v1_cli to just dump its inputs to disk. Might do that too / instead, once I go through enough of this code to understand what's happening.

BTW, I found https://github.com/nod-ai/sharktank/blob/main/shortfin/shortfin/llm/config.py has a nice comment at the top - I'd find that useful in docs outside of the source or closer to the model definition in sharktank/

[ScottTodd]

Here's a WIP patch adding debug file dumping to service_v1.py: main...ScottTodd:sharktank:llama-generate-data-2 . That gets the real values used in the service. Note that the service code itself is still under construction, but it at least models the stateful nature of this program.

Still need to read the cache from device memory to host memory and dump it, then try the .bin files with iree-run-module to see if it all works as expected.