Introduces a cuda version of the GPTQ kernel. This use triton to write a custom kernel for this code. On my machine it is about 2x faster than CPU :( was expecting it to be faster, but I guess CPU is pretty good. Main features

• Triton kernel is implemented in python and exported as a PTX file
• Called from Rust using fp16 and the CUST frameworks. These are gated in the installation.
• Had to move things around to make them both work. Main call is with a new feature gpu.

Right now this copies a bunch of code from the GPTQ datastructure. There is likely a cleaner way to do this.

This almost killed me. Triton is so buggy and hard to use.

There is also a test that checks that the CPU and GPU give the same results. However I need to read up on Rust testing to understand how to make this a real test.

This one is maybe just an idea - Cargo doesn't seem to support key/value features, but that's the best way to let crates that depend on this choose the model size. So here I set Cargo features (e.g. "7B") and then the build.rs converts that to a k/v pair (model_size = "7B"). I'm not sure the conversion is necessary but I didn't want to mess too much with your existing setup.

Separately this sets AVX512 as a enable-able feature that errors if the host machine doesn't support AVX512. My home computer doesn't have it, so I haven't checked it.

I am trying to run export.py, but I'm running into the following error:

│Traceback (most recent call last):                                                                                                                           │
│  File "/home/veda/Downloads/llama2.rs/export.py", line 134, in <module>                                                                                     │
│    load_and_export("", output_path)                                                                                                                         │
│  File "/home/veda/Downloads/llama2.rs/export.py", line 126, in load_and_export                                                                              │
│    export(model, output_path)                                                                                                                               │
│  File "/home/veda/Downloads/llama2.rs/export.py", line 57, in export                                                                                        │
│    hidden_dim = model.layers[0].mlp.up_proj.build()[0].shape[1]                                                                                             │
│  File "/home/veda/miniconda3/envs/llama2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1695, in __getattr__                                │
│    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")                                                                        │
│AttributeError: 'GeneralQuantLinear' object has no attribute 'build'                                                                                         │

I believe this is related to having the incorrect version of some library, but I can't determine the correct version (I'm not really much of a python person), my versions are:

Python 3.10.9
torch==2.1.0.dev20230812+rocm5.5 
transformers==4.31.0
auto-gptq==0.3.0

I would suggest either updating the readme to include the exact versions to use, or adding a requirements.txt.

@rachtsingh

For the Cuda code, I convert the model to a different form let's call it QCudaWeights. However, since it is not mmap'ed the it now has an owner. However I'm not sure how to create a LLamaModel from this.

pub struct LlamaModel {
    mmap: Mmap,
    pub config: Config,
    pub weights: &'static TWeights,
}

Any ideas for how to have a LlamaModel that can handle either owned or non-owned weights? Mmap is pretty interesting as a concept that breaks borrowing in Rust.

Playing around with a CUDA kernel for GPT-Q in triton. Was able to get a version working but it was too slow. Will probably get around to finishing this next week.

Hey Sasha,

I was trying out llama2.rs and wanted to swap between the 7B/13B versions on the fly, and I think using conditional compilation here makes it a bit easier.

I also spent some time trying to optimize the SIMD but it seems really fast! I couldn't find any easy optimizations.

Two main things I did:

• As all the dimensions you're vectorizing are perfectly matching the lane sizes you can move from chunks() to chunks_exact() this avoids a branching operation allowing the compiler to evict cmp, jmp and generating epilog "in case"

• Emit a prefetching (x86_64 only - _mm_prefetech ) operation in the most inner loop doing the dequantization from zero point + scaling. It seems the compiler (or the CPU prefetech) is not catching the linear access over qzeros and generates cache misses. With the prefetch instruction we ask to load the next chunk into L1/L2[/L3] registers

This wraps a bunch of the Python scaffolding behind a feature python so you can compile the binary without pyo3 at all. The default pip install . and maturin build commands use the python feature so there's no change necessary to README, I think.

Shrinks the binary from 38MB to 31MB on my machine.

Does the following:

1. Creates a Python package llama2_rs in addition to the existing binary (also called llama2_rs) using PyO3, which exposes an interface for doing inference from Python. You can see an example in tests/test_model_can_run.py
2. Moves some of the code needed by both paths (e.g. prefill and generate) into a new inference.rs so it can be used from both places.
3. Wraps up the &TWeights into a new struct LlamaModel that holds the mmap as well.

Notes:

• Includes the change in the export.py PR since it was useful for testing.
• I think some of the names could be changed, in particular maybe just llama2 instead of llama2_rs?
• Note the pinned auto-gptq version... I thought it safest until we figure out what's going on with https://github.com/srush/llama2.rs/issues/23
• cargo build --release does still generate a binary here, but if you want the Python package installed you want to run: pip install maturin; maturin develop --release, so if you decide to merge, might need a docs change.

Next up write a check that the matvec results are the same in AutoGPTQ and here?

• There is no command for install llama2-70b-q.bin and does'nt work

pip install torch transformers [auto-gptq](https://github.com/PanQiWei/AutoGPTQ) python export.py llama2-70b-q.bin

• this works

pip install torch transformers auto-gptq python export.py llama2-70b-q.bin

Cuts off some of the additional tokens that are added for llama fine tunes (will support in a future PR).

https://github.com/srush/llama2.rs/issues/22