eace8bf0b9
This PR is the first step towards fixing https://github.com/vllm-project/vllm/pull/3208 It implements dynamic per-tensor scaling (see https://github.com/vllm-project/vllm/pull/4118), so users do not need to compute activation scales on a calibration dataset and they also don't need to convert their model checkpoints. It is enough to specify the `quantization="fp8"` argument. You can try out the PR like this: ```python from vllm import LLM, SamplingParams prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] sampling_params = SamplingParams(temperature=0.8, top_p=0.95) llm = LLM(model="mistralai/Mixtral-8x7B-Instruct-v0.1", tensor_parallel_size=2, quantization="fp8") outputs = llm.generate(prompts, sampling_params) # Print the outputs. for output in outputs: prompt = output.prompt generated_text = output.outputs[0].text print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") ``` **Performance**: For this PR, the focus is on making the code clean (while still trying to get reasonable performance), there is a bunch of optimizations that we will submit as a follow up PR that significantly improve the performance (similar to the numbers in https://github.com/vllm-project/vllm/pull/3954). With this PR, the results are as follows: <img width="725" alt="Screenshot 2024-04-21 at 1 31 50 PM" src="https://github.com/vllm-project/vllm/assets/113316/d8fe1118-07a0-4d4e-8530-37a77d465a03"> **Accuracy**: The accuracy with this PR on MMLU on `mistralai/Mixtral-8x7B-v0.1` is as follows: ``` | Groups |Version|Filter|n-shot|Metric|Value | |Stderr| |------------------|-------|------|-----:|------|-----:|---|-----:| |mmlu |N/A |none | 0|acc |0.7018|± |0.0036| | - humanities |N/A |none | 5|acc |0.6472|± |0.0065| | - other |N/A |none | 5|acc |0.7673|± |0.0072| | - social_sciences|N/A |none | 5|acc |0.8099|± |0.0070| | - stem |N/A |none | 5|acc |0.6131|± |0.0083| ``` this compares favorably with the fp16 results which are ``` | Groups |Version|Filter|n-shot|Metric|Value | |Stderr| |------------------|-------|------|-----:|------|-----:|---|-----:| |mmlu |N/A |none | 0|acc |0.7020|± |0.1313| | - humanities |N/A |none | 5|acc |0.6425|± |0.1349| | - other |N/A |none | 5|acc |0.7744|± |0.1038| | - social_sciences|N/A |none | 5|acc |0.8131|± |0.0695| | - stem |N/A |none | 5|acc |0.6108|± |0.1383| ``` Happy hacking!
182 lines
4.4 KiB
C++
182 lines
4.4 KiB
C++
#pragma once
|
|
|
|
#include <torch/extension.h>
|
|
|
|
void paged_attention_v1(
|
|
torch::Tensor& out,
|
|
torch::Tensor& query,
|
|
torch::Tensor& key_cache,
|
|
torch::Tensor& value_cache,
|
|
int num_kv_heads,
|
|
float scale,
|
|
torch::Tensor& block_tables,
|
|
torch::Tensor& context_lens,
|
|
int block_size,
|
|
int max_context_len,
|
|
const c10::optional<torch::Tensor>& alibi_slopes,
|
|
const std::string& kv_cache_dtype,
|
|
float kv_scale);
|
|
|
|
void paged_attention_v2(
|
|
torch::Tensor& out,
|
|
torch::Tensor& exp_sums,
|
|
torch::Tensor& max_logits,
|
|
torch::Tensor& tmp_out,
|
|
torch::Tensor& query,
|
|
torch::Tensor& key_cache,
|
|
torch::Tensor& value_cache,
|
|
int num_kv_heads,
|
|
float scale,
|
|
torch::Tensor& block_tables,
|
|
torch::Tensor& context_lens,
|
|
int block_size,
|
|
int max_context_len,
|
|
const c10::optional<torch::Tensor>& alibi_slopes,
|
|
const std::string& kv_cache_dtype,
|
|
float kv_scale);
|
|
|
|
void rms_norm(
|
|
torch::Tensor& out,
|
|
torch::Tensor& input,
|
|
torch::Tensor& weight,
|
|
float epsilon);
|
|
|
|
void fused_add_rms_norm(
|
|
torch::Tensor& input,
|
|
torch::Tensor& residual,
|
|
torch::Tensor& weight,
|
|
float epsilon);
|
|
|
|
void rotary_embedding(
|
|
torch::Tensor& positions,
|
|
torch::Tensor& query,
|
|
torch::Tensor& key,
|
|
int head_size,
|
|
torch::Tensor& cos_sin_cache,
|
|
bool is_neox);
|
|
|
|
void batched_rotary_embedding(
|
|
torch::Tensor& positions,
|
|
torch::Tensor& query,
|
|
torch::Tensor& key,
|
|
int head_size,
|
|
torch::Tensor& cos_sin_cache,
|
|
bool is_neox,
|
|
int rot_dim,
|
|
torch::Tensor& cos_sin_cache_offsets);
|
|
|
|
void silu_and_mul(
|
|
torch::Tensor& out,
|
|
torch::Tensor& input);
|
|
|
|
void gelu_and_mul(
|
|
torch::Tensor& out,
|
|
torch::Tensor& input);
|
|
|
|
void gelu_tanh_and_mul(
|
|
torch::Tensor& out,
|
|
torch::Tensor& input);
|
|
|
|
void gelu_new(
|
|
torch::Tensor& out,
|
|
torch::Tensor& input);
|
|
|
|
void gelu_fast(
|
|
torch::Tensor& out,
|
|
torch::Tensor& input);
|
|
|
|
#ifndef USE_ROCM
|
|
torch::Tensor aqlm_gemm(
|
|
const torch::Tensor& input,
|
|
const torch::Tensor& codes,
|
|
const torch::Tensor& codebooks,
|
|
const torch::Tensor& scales,
|
|
const torch::Tensor& codebook_partition_sizes,
|
|
const std::optional<torch::Tensor>& bias
|
|
);
|
|
|
|
torch::Tensor aqlm_dequant(
|
|
const torch::Tensor& codes,
|
|
const torch::Tensor& codebooks,
|
|
const torch::Tensor& codebook_partition_sizes
|
|
);
|
|
|
|
torch::Tensor awq_gemm(
|
|
torch::Tensor _in_feats,
|
|
torch::Tensor _kernel,
|
|
torch::Tensor _scaling_factors,
|
|
torch::Tensor _zeros,
|
|
int split_k_iters);
|
|
|
|
torch::Tensor awq_dequantize(
|
|
torch::Tensor _kernel,
|
|
torch::Tensor _scaling_factors,
|
|
torch::Tensor _zeros,
|
|
int split_k_iters,
|
|
int thx,
|
|
int thy);
|
|
|
|
torch::Tensor marlin_gemm(
|
|
torch::Tensor& a,
|
|
torch::Tensor& b_q_weight,
|
|
torch::Tensor& b_scales,
|
|
torch::Tensor& workspace,
|
|
int64_t size_m,
|
|
int64_t size_n,
|
|
int64_t size_k);
|
|
#endif
|
|
|
|
void squeezellm_gemm(
|
|
torch::Tensor vec,
|
|
torch::Tensor mat,
|
|
torch::Tensor mul,
|
|
torch::Tensor lookup_table);
|
|
|
|
torch::Tensor gptq_gemm(
|
|
torch::Tensor a,
|
|
torch::Tensor b_q_weight,
|
|
torch::Tensor b_gptq_qzeros,
|
|
torch::Tensor b_gptq_scales,
|
|
torch::Tensor b_g_idx,
|
|
bool use_exllama,
|
|
int bit);
|
|
|
|
void gptq_shuffle(
|
|
torch::Tensor q_weight,
|
|
torch::Tensor q_perm,
|
|
int bit);
|
|
|
|
void scaled_fp8_quant(
|
|
torch::Tensor& out,
|
|
torch::Tensor& input,
|
|
torch::Tensor& scale);
|
|
|
|
void moe_align_block_size(
|
|
torch::Tensor topk_ids,
|
|
int num_experts,
|
|
int block_size,
|
|
torch::Tensor sorted_token_ids,
|
|
torch::Tensor experts_ids,
|
|
torch::Tensor num_tokens_post_pad);
|
|
|
|
#ifndef USE_ROCM
|
|
using fptr_t = uint64_t;
|
|
fptr_t init_custom_ar(torch::Tensor &meta, torch::Tensor &rank_data,
|
|
const std::vector<std::string> &handles,
|
|
const std::vector<int64_t> &offsets, int rank,
|
|
bool full_nvlink);
|
|
bool should_custom_ar(torch::Tensor &inp, int max_size, int world_size,
|
|
bool full_nvlink);
|
|
void all_reduce_reg(fptr_t _fa, torch::Tensor &inp, torch::Tensor &out);
|
|
void all_reduce_unreg(fptr_t _fa, torch::Tensor &inp, torch::Tensor ®_buffer,
|
|
torch::Tensor &out);
|
|
void dispose(fptr_t _fa);
|
|
int meta_size();
|
|
void register_buffer(fptr_t _fa, torch::Tensor &t,
|
|
const std::vector<std::string> &handles,
|
|
const std::vector<int64_t> &offsets);
|
|
std::pair<std::vector<uint8_t>, std::vector<int64_t>> get_graph_buffer_ipc_meta(fptr_t _fa);
|
|
void register_graph_buffers(fptr_t _fa, const std::vector<std::string> &handles,
|
|
const std::vector<std::vector<int64_t>> &offsets);
|
|
#endif
|