!git clone https://github.com/deepseek-ai/FlashMLA.git!cd FlashMLA && python setup.py develop

#如果前一个cell执行成功，这个cell无需执行，可直接跳过!wget https://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/flashmla/FlashMLA.tgz && tar -xvzf FlashMLA.tgz!cd FlashMLA && python setup.py develop

# 安装最新版本的vllm，来获取部分FlashMLA运行需要的fix (build on commit id 4a8cfc75516f3df2ead816c0270b63470eb8e4ee)!pip install https://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/flashmla/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl# 安装最新flashinfer用于MLA性能比较 (build on commit id 56e56ea4e22fc7cf5ca3b41f8143e8847eb5f3fa)!pip install https://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/flashmla/flashinfer_python-0.2.2-py3-none-any.whl

# 下载、解压并替换vLLM库中的部分文件!wget https://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/flashmla/vllm_patch.tar && tar -xvf vllm_patch.tar!cp -r vllm-patch/vllm/* /usr/local/lib/python3.10/site-packages/vllm/

import osdsw_region = os.environ.get("dsw_region")url_link = {"cn-shanghai": "https://atp-modelzoo-sh.oss-cn-shanghai-internal.aliyuncs.com/release/tutorials/flashmla/DeepSeek_v2_lite_chat.tar"}path = url_link[dsw_region] if dsw_region in url_link else "https://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/flashmla/DeepSeek_v2_lite_chat.tar"os.environ['LINK_CHAT'] = path!wget $LINK_CHAT!tar -xvf DeepSeek_v2_lite_chat.tar

import sysimport ossys.path.append(os.path.join(os.getcwd(), 'FlashMLA'))from benchmark.bench_flash_mla import *import matplotlib.pyplot as pltimport pandas as pdfrom collections import defaultdict
# 测试函数def compare_a(target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype):# print(f"{target}: {b=}, {s_q=}, mean_seqlens={cache_seqlens.float().mean()}, {h_q=}, {h_kv=}, {d=}, {dv=}, {causal=}, {dtype=}")torch.set_default_dtype(dtype)device = torch.device("cuda:0")torch.set_default_device(device)torch.cuda.set_device(device)torch.manual_seed(0)random.seed(0)assert target in FUNC_TABLEtarget_func = FUNC_TABLE[target]total_seqlens = cache_seqlens.sum().item()mean_seqlens = cache_seqlens.float().mean().int().item()max_seqlen = cache_seqlens.max().item()max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256# print(f"{total_seqlens=}, {mean_seqlens=}, {max_seqlen=}")q = torch.randn(b, s_q, h_q, d)block_size = 64block_table = torch.arange(b * max_seqlen_pad // block_size, dtype=torch.int32).view(b, max_seqlen_pad // block_size)blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d)out_b, lse_b, perf_b = target_func(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype)FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d + b * s_q * h_q * dv) * (torch.finfo(dtype).bits // 8)# print(f"perf {target}: {perf_b:.3f} ms, {FLOPS / 10 ** 9 / perf_b:.0f} TFLOPS, {bytes / 10 ** 6 / perf_b:.0f} GB/s")return bytes / 10 ** 6 / perf_bdata = defaultdict(lambda: defaultdict(list))for shape in shape_configs:for target in available_targets:perf = compare_a(target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"])# batchsize 128 & head 128data[target][shape["cache_seqlens"].float().mean().cpu().item()] = perffor name, dt in data.items():k, v = list(dt.keys()), list(dt.values())plt.plot(k, v, label=name)plt.title('bandwidth')plt.xlabel('seqlen')plt.ylabel('bw (GB/s)')plt.legend()plt.show()

[cuda.py:176] Using Triton MLA backend.[cuda.py:173] Using FlashMLA backend.

(Triton MLA) Generate 527 tokens in 17.97 secs(FlashMLA) Generate 515 tokens in 16.64 secs

import osimport timeimport loggingfrom transformers import AutoTokenizerfrom vllm import LLM, SamplingParamsmodel_name = "DeepSeek-V2-Lite-Chat"max_model_len, tp_size = 8192, 1def build_triton_mla_dskv2():os.environ['VLLM_ATTENTION_BACKEND'] = 'TRITON_MLA'return LLM(model=model_name, tensor_parallel_size=tp_size, max_model_len=max_model_len, trust_remote_code=True, enforce_eager=True, block_size=64)def build_flash_mla_dskv2():os.environ['VLLM_ATTENTION_BACKEND'] = 'FLASHMLA'return LLM(model=model_name, tensor_parallel_size=tp_size, max_model_len=max_model_len, trust_remote_code=True, enforce_eager=True, block_size=64)def warmup_and_infer(llm, messages_list):tokenizer = AutoTokenizer.from_pretrained(model_name)sampling_params = SamplingParams(temperature=0.3, max_tokens=2048, stop_token_ids=[tokenizer.eos_token_id])warmup_token_ids = [tokenizer.apply_chat_template([{"role": "user", "content": "What's your name?"}], add_generation_prompt=True)]llm.generate(prompt_token_ids=warmup_token_ids, sampling_params=sampling_params)prompt_token_ids = [tokenizer.apply_chat_template(messages, add_generation_prompt=True) for messages in messages_list]start_time = time.time()outputs = llm.generate(prompt_token_ids=prompt_token_ids, sampling_params=sampling_params)end_time = time.time()return [output.outputs[0].text for output in outputs], sum(len(output.outputs[0].token_ids) for output in outputs), (end_time - start_time)

messages_list = [    [{"role": "user", "content": "Write a piece of quicksort code in C++."}],]# 使用FlashMLA进行推理generated_text, n_tokens, generation_time = warmup_and_infer(build_flash_mla_dskv2(), messages_list)# 使用默认MLA后端进行推理，如果FlashMLA已经运行完成，切换至默认MLA之后，请注意刷新kernel后再重新运行# generated_text, n_tokens, generation_time = warmup_and_infer(build_triton_mla_dskv2(), messages_list)print(f'Generate {n_tokens} tokens in {generation_time:.2f} secs')print(generated_text[0])

INFO 02-25 15:57:54 [config.py:208] Replacing legacy 'type' key with 'rope_type'INFO 02-25 15:58:00 [config.py:569] This model supports multiple tasks: {'classify', 'reward', 'score', 'embed', 'generate'}. Defaulting to 'generate'.WARNING 02-25 15:58:00 [cuda.py:95] To see benefits of async output processing, enable CUDA graph. Since, enforce-eager is enabled, async output processor cannot be usedINFO 02-25 15:58:00 [llm_engine.py:234] Initializing a V0 LLM engine (v0.7.4.dev75+g4a8cfc75) with config: model='DeepSeek-V2-Lite-Chat', speculative_config=None, tokenizer='DeepSeek-V2-Lite-Chat', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=True, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar'), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=DeepSeek-V2-Lite-Chat, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=False, chunked_prefill_enabled=False, use_async_output_proc=False, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={"splitting_ops":[],"compile_sizes":[],"cudagraph_capture_sizes":[],"max_capture_size":0}, use_cached_outputs=False, INFO 02-25 15:58:01 [cuda.py:173] Using FlashMLA backend.INFO 02-25 15:58:01 [parallel_state.py:948] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0INFO 02-25 15:58:01 [model_runner.py:1110] Starting to load model DeepSeek-V2-Lite-Chat...INFO 02-25 15:58:01 [cuda.py:173] Using FlashMLA backend.Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]Loading safetensors checkpoint shards:  25% Completed | 1/4 [00:01<00:04,  1.37s/it]Loading safetensors checkpoint shards:  50% Completed | 2/4 [00:02<00:02,  1.45s/it]Loading safetensors checkpoint shards:  75% Completed | 3/4 [00:03<00:01,  1.26s/it]Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:05<00:00,  1.33s/it]Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:05<00:00,  1.34s/it]
INFO 02-25 15:58:07 [model_runner.py:1117] Loading model weights took 31.1253 GB and 5.736558 secondsWARNING 02-25 15:58:09 [fused_moe.py:849] Using default MoE config. Performance might be sub-optimal! Config file not found at /usr/local/lib/python3.10/site-packages/vllm/model_executor/layers/fused_moe/configs/E=64,N=1408,device_name=****.jsonINFO 02-25 15:58:09 [worker.py:267] Memory profiling takes 1.76 secondsINFO 02-25 15:58:09 [worker.py:267] the current vLLM instance can use total_gpu_memory (95.00GiB) x gpu_memory_utilization (0.90) = 85.50GiBINFO 02-25 15:58:09 [worker.py:267] model weights take 31.13GiB; non_torch_memory takes 0.12GiB; PyTorch activation peak memory takes 0.96GiB; the rest of the memory reserved for KV Cache is 53.30GiB.INFO 02-25 15:58:09 [executor_base.py:111] # cuda blocks: 25874, # CPU blocks: 1941INFO 02-25 15:58:09 [executor_base.py:116] Maximum concurrency for 8192 tokens per request: 202.14xINFO 02-25 15:58:11 [llm_engine.py:436] init engine (profile, create kv cache, warmup model) took 4.16 seconds/tmp/ipykernel_3769/871633.py:7: DeprecationWarning: The keyword arguments {'prompt_token_ids'} are deprecated and will be removed in a future update. Please use the 'prompts' parameter instead.  generated_text, n_tokens, generation_time = warmup_and_infer(build_flash_mla_dskv2(), messages_list)Processed prompts: 100%|██████████| 1/1 [00:01<00:00,  1.44s/it, est. speed input: 9.02 toks/s, output: 23.59 toks/s]/tmp/ipykernel_3769/871633.py:7: DeprecationWarning: The keyword arguments {'prompt_token_ids'} are deprecated and will be removed in a future update. Please use the 'prompts' parameter instead.  generated_text, n_tokens, generation_time = warmup_and_infer(build_flash_mla_dskv2(), messages_list)Processed prompts: 100%|██████████| 1/1 [00:17<00:00, 17.13s/it, est. speed input: 0.99 toks/s, output: 30.06 toks/s]Generate 515 tokens in 17.13 secs Here is a simple implementation of the QuickSort algorithm in C++:
```cpp#include <iostream>#include <vector>
void swap(int* a, int* b) {    int t = *a;    *a = *b;    *b = t;}
int partition (std::vector<int>& arr, int low, int high) {    int pivot = arr[high];     int i = (low - 1); 
    for (int j = low; j <= high - 1; j++) {        if (arr[j] < pivot) {            i++;             swap(&arr[i], &arr[j]);        }    }    swap(&arr[i + 1], &arr[high]);    return (i + 1);}
void quickSort(std::vector<int>& arr, int low, int high) {    if (low < high) {        int pi = partition(arr, low, high);        quickSort(arr, low, pi - 1);        quickSort(arr, pi + 1, high);    }}
void printArray(std::vector<int>& arr) {    for (int i = 0; i < arr.size(); ++i)        std::cout << arr[i] << " ";    std::cout << "\n";}
int main() {    std::vector<int> arr = {10, 7, 8, 9, 1, 5};    int n = arr.size();    quickSort(arr, 0, n - 1);    std::cout << "Sorted array: \n";    printArray(arr);    return 0;}This code sorts an array in ascending order using the QuickSort algorithm. The quickSort function is a recursive function that sorts the sub-array to the left of pi and the sub-array to the right of pi. The partition function rearranges the elements in the array so that all elements less than the pivot are to its left and all elements greater are to its right. The pivot is always the last element of the sub-array.