From ac56e4e3329d637b40476fe478578e80d8653830 Mon Sep 17 00:00:00 2001 From: yingcanw Date: Tue, 24 Dec 2024 03:48:57 -0800 Subject: [PATCH 01/12] Add README --- examples/deepseek_v3/README.md | 269 +++++++++++++++++++++++++++++++++ 1 file changed, 269 insertions(+) create mode 100644 examples/deepseek_v3/README.md diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md new file mode 100644 index 000000000000..34e13da5a13f --- /dev/null +++ b/examples/deepseek_v3/README.md @@ -0,0 +1,269 @@ +# Deepseek-v3 + +This document shows how to build and run `DeepSeek-v3` model in TensorRT-LLM. + + +***Current branch is a preview version only for supporting DeepSeek-V3 in bf16 precision, which will be part of the TRT-LLM official release in 0.18+ version.*** + + +- [Deepseek-V3](#deepseek-v3) + - [Prerequisite](#prerequisite) + - [Hardware](#hardware) + - [Overview](#overview) + - [Support Matrix](#support-matrix) + - [Usage](#usage) + - [Build TensorRT engine(s)](#build-tensorrt-engines) + +## Prerequisite + +First, please download DeepSeek-V3 weights from HF https://huggingface.co/deepseek-ai/DeepSeek-V3. + +```bash +git lfs install +git clone https://huggingface.co/deepseek-ai/DeepSeek-V3 +``` + +## Hardware + +The DeepSeek-V3 model requires least 32x80G GPU memory, model contains 660B parameters roughly 1.3TB memory (with BF16 precision). + +***Caution: Current TRT-LLM MLA kernel only support Hopper architecture (SM90). Ampere architecture (SM80 & SM86) will be supported in next release.*** + +## Overview + +The TensorRT-LLM DeepSeek-V3 implementation can be found in [tensorrt_llm/models/deepseek_v2/model.py](../../tensorrt_llm/models/deepseek_v2/model.py). The TensorRT-LLM Deepseek-V3 example code is located in [`example/deepseek_v3`](./). There is one main file: + +* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the DeepSeek-V3 model into tensorrt-llm checkpoint format. + +In addition, there are three shared files in the parent folder [`examples`](../) can be used for inference and evaluation: + +* [`../run.py`](../run.py) to run the model inference output by given an input text. +* [`../summarize.py`](../summarize.py) to summarize the article from [cnn_dailmail](https://huggingface.co/datasets/cnn_dailymail) dataset, it can running the summarize from HF model and TensorRT-LLM model. +* [`../mmlu.py`](../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT-LLM model on the MMLU dataset. + +## Support Matrix + +- [x] BF16 +- [ ] FP8 +- [x] Tensor Parallelism +- [ ] Expert Parallelism + +***Caution: prefer using BF16 over FP16 for DeepSeek-V3 since model original training precision is FP8 and we found direct convert FP8 -> FP16 may cause unknown accuracy issues.*** + +## Usage + +The TensorRT-LLM DeepSeek-V3 example code locates at [examples/deepseek_v3](./). It takes PyTorch weights as input, and builds corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. + +### Build TensorRT engine(s) + +Below is the step-by-step to run DeepSeek-V3 with TensorRT-LLM. + +First the checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be build with TensorRT-LLM checkpoint. + +```bash +# Convert Deepseek-v HF weights to TensorRT-LLM checkpoint format. +python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --dtype bfloat16 \ + --tp_size 32 + +# With '--load_model_on_cpu' option if total GPU memory is insufficient +python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_cpu_bf16 \ + --dtype bfloat16 \ + --tp_size 32 \ + --load_model_on_cpu +``` + + +We observe the checkpoint conversion time took exceed 10 hours, while use CPU memory required >= 1.3TB. + +After the checkpoint conversion, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. + +```bash +# Build engine +trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --output_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ + --gpt_attention_plugin bfloat16 \ + --gemm_plugin bfloat16 \ + --max_batch_size 4 \ + --max_seq_len 4096 \ + --max_input_len 2048 \ + --use_paged_context_fmha enable +``` + +***Caution: `--max_batch_size` and `--max_seq_len` are the main factors to determine how many GPU memory will be used during runtime, so later when try to run e.g., `summarize.py` or `mmlu.py` or `gptManagerBenchmark.cpp`may need adjust `--max_batch_size` and `--max_seq_len` accordingly to avoid OOM.(meaning rebuild TensorRT engine with smaller `--max_batch_size` and `--max_seq_len` if needed based on GPU memory size), there is beautiful technical log perf-best-practices.md (https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/performance/perf-best-practices.md) explained the mechanism.*** + +Test the engine with [run.py](../run.py) script: + +```bash +mpirun --allow-run-as-root -np 32 python ../run.py --engine_dir ./trtllm_engines/deepseek_3/bf16/tp32-sel4096-isl2048-bs4 \ + --tokenizer_dir ./DeepSeek-V3 \ + --max_output_len 40 \ + --input_text "The president of the United States is person who" +``` + +and the output will be like: + +``` +[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31490468978882 sec +[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31163835525513 sec +[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31164216995239 sec +[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31491041183472 sec +[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3116364479065 sec +[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3118085861206 sec +[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3118691444397 sec +[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31516337394714 sec +Input [Text 0]: "<｜begin▁of▁sentence｜>The president of the United States is person who" +Output [Text 0 Beam 0]: " is elected by the people of the United States to lead the country. The president is the head of the executive branch of the government. The president is also the commander in chief of the armed forces." +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +``` + +If we want to evaluate the model summarization ability, we can use [summarize.py](../summarize.py) script: + +```bash +mpirun --allow-run-as-root -np 32 python ../summarize.py --engine_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ + --hf_model_dir ./DeepSeek-V3 \ + --data_type bfloat16 \ + --batch_size 1 \ + --test_trt_llm \ + --test_hf +``` + +and the output will be like: + + +``` +[10/28/2024-16:46:22] [TRT-LLM] [I] HF Generated : +[10/28/2024-16:46:22] [TRT-LLM] [I] Input : ['(CNN)James Best, best known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," died Monday after a brief illness. He was 88. Best died in hospice in Hickory, North Carolina, of complications from pneumonia, said Steve Latshaw, a longtime friend and Hollywood colleague. Although he\'d been a busy actor for decades in theater and in Hollywood, Best didn\'t become famous until 1979, when "The Dukes of Hazzard\'s" cornpone charms began beaming into millions of American homes almost every Friday night. For seven seasons, Best\'s Rosco P. Coltrane chased the moonshine-running Duke boys back and forth across the back roads of fictitious Hazzard County, Georgia, although his "hot pursuit" usually ended with him crashing his patrol car. Although Rosco was slow-witted and corrupt, Best gave him a childlike enthusiasm that got laughs and made him endearing. His character became known for his distinctive "kew-kew-kew" chuckle and for goofy catchphrases such as "cuff \'em and stuff \'em!" upon making an arrest. Among the most popular shows on TV in the early \'80s, "The Dukes of Hazzard" ran until 1985 and spawned TV movies, an animated series and video games. Several of Best\'s "Hazzard" co-stars paid tribute to the late actor on social media. "I laughed and learned more from Jimmie in one hour than from anyone else in a whole year," co-star John Schneider, who played Bo Duke, said on Twitter. "Give Uncle Jesse my love when you see him dear friend." "Jimmy Best was the most constantly creative person I have ever known," said Ben Jones, who played mechanic Cooter on the show, in a Facebook post. "Every minute of his long life was spent acting, writing, producing, painting, teaching, fishing, or involved in another of his life\'s many passions." Born Jewel Guy on July 26, 1926, in Powderly, Kentucky, Best was orphaned at 3 and adopted by Armen and Essa Best, who renamed him James and raised him in rural Indiana. Best served in the Army during World War II before launching his acting career. In the 1950s and 1960s, he accumulated scores of credits, playing a range of colorful supporting characters in such TV shows as "The Twilight Zone," "Bonanza," "The Andy Griffith Show" and "Gunsmoke." He later appeared in a handful of Burt Reynolds\' movies, including "Hooper" and "The End." But Best will always be best known for his "Hazzard" role, which lives on in reruns. "Jimmie was my teacher, mentor, close friend and collaborator for 26 years," Latshaw said. "I directed two of his feature films, including the recent \'Return of the Killer Shrews,\' a sequel he co-wrote and was quite proud of as he had made the first one more than 50 years earlier." People we\'ve lost in 2015 . CNN\'s Stella Chan contributed to this story.'] +[10/28/2024-16:46:22] [TRT-LLM] [I] + Reference : ['James Best, who played the sheriff on "The Dukes of Hazzard," died Monday at 88 .\n"Hazzard" ran from 1979 to 1985 and was among the most popular shows on TV .'] +[10/28/2024-16:46:22] [TRT-LLM] [I] + Output : [[' James Best, best known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," died Monday after a brief illness. He was 88.']] +[10/28/2024-16:46:22] [TRT-LLM] [I] --------------------------------------------------------- +[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT-LLM (total latency: 32.02327513694763 sec) +[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT-LLM (total output tokens: 1394) +[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT-LLM (tokens per second: 43.53083793080361) +[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT-LLM beam 0 result +[10/28/2024-16:49:33] [TRT-LLM] [I] rouge1 : 17.85755990133811 +[10/28/2024-16:49:33] [TRT-LLM] [I] rouge2 : 6.273032755727469 +[10/28/2024-16:49:33] [TRT-LLM] [I] rougeL : 14.768323033457317 +[10/28/2024-16:49:33] [TRT-LLM] [I] rougeLsum : 15.700915348496391 +[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (total latency: 189.76398921012878 sec) +[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (total output tokens: 1376) +[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (tokens per second: 7.2511123197159) +[10/28/2024-16:49:33] [TRT-LLM] [I] HF beam 0 result +[10/28/2024-16:49:33] [TRT-LLM] [I] rouge1 : 18.542590123197257 +[10/28/2024-16:49:33] [TRT-LLM] [I] rouge2 : 6.345777100488389 +[10/28/2024-16:49:33] [TRT-LLM] [I] rougeL : 15.235695878419156 +[10/28/2024-16:49:33] [TRT-LLM] [I] rougeLsum : 16.64935135356226 +``` + +At last, we can evaluate the model with [mmlu.py](../mmlu.py) script: + +```bash +# Download MMLU dataset +mkdir mmlu_data && cd mmlu_data +wget https://people.eecs.berkeley.edu/~hendrycks/data.tar && tar -xf data.tar +# Run MMLU evaluation +mpirun --allow-run-as-root -n 32 python ../mmlu.py --engine_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ + --hf_model_dir ./DeepSeek-V3 \ + --data_type bfloat16 \ + --batch_size 1 \ + --test_trt_llm \ + --test_hf \ + --data_dir ./mmlu_data/data/ +``` + + + +and the output will be like: + +``` +Average accuracy 0.480 - abstract_algebra +Average accuracy 0.741 - anatomy +Average accuracy 0.888 - astronomy +Average accuracy 0.790 - business_ethics +Average accuracy 0.845 - clinical_knowledge +Average accuracy 0.924 - college_biology +Average accuracy 0.600 - college_chemistry +Average accuracy 0.720 - college_computer_science +Average accuracy 0.510 - college_mathematics +Average accuracy 0.751 - college_medicine +Average accuracy 0.618 - college_physics +Average accuracy 0.860 - computer_security +Average accuracy 0.796 - conceptual_physics +Average accuracy 0.675 - econometrics +Average accuracy 0.800 - electrical_engineering +Average accuracy 0.741 - elementary_mathematics +Average accuracy 0.643 - formal_logic +Average accuracy 0.570 - global_facts +Average accuracy 0.910 - high_school_biology +Average accuracy 0.714 - high_school_chemistry +Average accuracy 0.890 - high_school_computer_science +Average accuracy 0.891 - high_school_european_history +Average accuracy 0.909 - high_school_geography +Average accuracy 0.953 - high_school_government_and_politics +Average accuracy 0.826 - high_school_macroeconomics +Average accuracy 0.522 - high_school_mathematics +Average accuracy 0.916 - high_school_microeconomics +Average accuracy 0.576 - high_school_physics +Average accuracy 0.923 - high_school_psychology +Average accuracy 0.704 - high_school_statistics +Average accuracy 0.907 - high_school_us_history +Average accuracy 0.932 - high_school_world_history +Average accuracy 0.834 - human_aging +Average accuracy 0.893 - human_sexuality +Average accuracy 0.909 - international_law +Average accuracy 0.880 - jurisprudence +Average accuracy 0.853 - logical_fallacies +Average accuracy 0.598 - machine_learning +Average accuracy 0.874 - management +Average accuracy 0.953 - marketing +Average accuracy 0.880 - medical_genetics +Average accuracy 0.920 - miscellaneous +Average accuracy 0.850 - moral_disputes +Average accuracy 0.613 - moral_scenarios +Average accuracy 0.830 - nutrition +Average accuracy 0.859 - philosophy +Average accuracy 0.883 - prehistory +Average accuracy 0.635 - professional_accounting +Average accuracy 0.625 - professional_law +Average accuracy 0.857 - professional_medicine +Average accuracy 0.833 - professional_psychology +Average accuracy 0.745 - public_relations +Average accuracy 0.869 - security_studies +Average accuracy 0.935 - sociology +Average accuracy 0.930 - us_foreign_policy +Average accuracy 0.596 - virology +Average accuracy 0.877 - world_religions +Average accuracy 0.632 - math +Average accuracy 0.801 - health +Average accuracy 0.738 - physics +Average accuracy 0.897 - business +Average accuracy 0.914 - biology +Average accuracy 0.677 - chemistry +Average accuracy 0.762 - computer science +Average accuracy 0.832 - economics +Average accuracy 0.800 - engineering +Average accuracy 0.736 - philosophy +Average accuracy 0.821 - other +Average accuracy 0.902 - history +Average accuracy 0.909 - geography +Average accuracy 0.883 - politics +Average accuracy 0.876 - psychology +Average accuracy 0.919 - culture +Average accuracy 0.660 - law +Average accuracy 0.727 - STEM +Average accuracy 0.740 - humanities +Average accuracy 0.873 - social sciences +Average accuracy 0.821 - other (business, health, misc.) +Average accuracy: 0.785 +``` \ No newline at end of file From 06148676829d7f196219f0dc5bac74837c509c1d Mon Sep 17 00:00:00 2001 From: jershi425 <83951930+jershi425@users.noreply.github.com> Date: Tue, 24 Dec 2024 22:20:45 +0800 Subject: [PATCH 02/12] Add unified converter (#1) * init v3 lite feat * fix moe topk method * fix noaux_tc logic * fix deepseek v3 normal rope * refactor * wo conversion ok debugging build * add quantize for attn.dense * add unified converter support * testing unified converter * add convert checkpoint and update docs --------- Co-authored-by: Zeyu Wang --- examples/deepseek_v2/convert_checkpoint.py | 55 +++- examples/deepseek_v3/README.md | 27 +- examples/deepseek_v3/convert_checkpoint.py | 253 +++++++++++++++++++ tensorrt_llm/layers/attention.py | 120 ++++++++- tensorrt_llm/layers/linear.py | 5 + tensorrt_llm/layers/moe.py | 49 +++- tensorrt_llm/models/__init__.py | 1 + tensorrt_llm/models/deepseek_v2/config.py | 168 +++++++++++++ tensorrt_llm/models/deepseek_v2/convert.py | 263 +++++++------------- tensorrt_llm/models/deepseek_v2/model.py | 126 +++++++--- tensorrt_llm/models/model_weights_loader.py | 5 +- tensorrt_llm/quantization/quantize.py | 1 + 12 files changed, 834 insertions(+), 239 deletions(-) create mode 100644 examples/deepseek_v3/convert_checkpoint.py create mode 100644 tensorrt_llm/models/deepseek_v2/config.py diff --git a/examples/deepseek_v2/convert_checkpoint.py b/examples/deepseek_v2/convert_checkpoint.py index e45ea15d58e8..7cbc7935d225 100755 --- a/examples/deepseek_v2/convert_checkpoint.py +++ b/examples/deepseek_v2/convert_checkpoint.py @@ -23,7 +23,8 @@ from tensorrt_llm.layers import MoeConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models import DeepseekV2ForCausalLM -from tensorrt_llm.models.deepseek_v2.convert import load_hf_deepseek +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo def parse_arguments(): @@ -58,7 +59,7 @@ def parse_arguments(): parser.add_argument('--load_model_on_cpu', default=False, action="store_true", - help='Choose to load HF cpkt into GPU') + help='Choose to load HF cpkt into CPU') parser.add_argument( '--use_parallel_embedding', action="store_true", @@ -75,6 +76,25 @@ def parse_arguments(): 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0)' 'To shard it along hidden dimension, set embedding_sharding_dim=1' 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') + + parser.add_argument( + '--use_weight_only', + default=False, + action="store_true", + help='Quantize weights for the various GEMMs to INT4/INT8.' + 'See --weight_only_precision to set the precision') + parser.add_argument( + '--weight_only_precision', + const='int8', + type=str, + nargs='?', + default='int8', + choices=['int8', 'int4'], + help= + 'Define the precision for the weights when using weight-only quantization.' + 'You must also use --use_weight_only for that argument to have an impact.' + ) + parser.add_argument('--output_dir', type=str, default='trtllm_checkpoint', @@ -125,12 +145,29 @@ def parse_arguments(): return args +def precision_to_config(precision, quant_config) -> QuantConfig: + '''update config dict for weight-only quantization + ''' + quant_config = QuantConfig() + precision_to_algo = {'int8': QuantAlgo.W8A16, 'int4': QuantAlgo.W4A16} + quant_config.quant_algo = precision_to_algo.get(precision) + return quant_config + + +def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: + '''return config dict with quantization info based on the command line args + ''' + quant_config = QuantConfig() + if args.use_weight_only: + quant_config = precision_to_config(args.weight_only_precision, + quant_config) + return quant_config + + def args_to_build_options(args): return { 'use_parallel_embedding': args.use_parallel_embedding, 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin, 'load_model_on_cpu': args.load_model_on_cpu } @@ -163,8 +200,7 @@ def convert_and_save_hf(args): override_fields = {} override_fields.update(args_to_build_options(args)) - load_model_on_cpu = args.load_model_on_cpu - hf_model = load_hf_deepseek(model_dir, load_model_on_cpu) + quant_config = args_to_quant_config(args) def convert_and_save_rank(args, rank): mapping = Mapping(world_size=world_size, @@ -175,7 +211,11 @@ def convert_and_save_rank(args, rank): moe_ep_size=args.moe_ep_size) deepseekv2 = DeepseekV2ForCausalLM.from_hugging_face( - hf_model, args.model_dir, args.dtype, mapping, **override_fields) + model_dir, + args.dtype, + mapping=mapping, + quant_config=quant_config, + **override_fields) deepseekv2.save_checkpoint(args.output_dir, save_config=(rank == 0)) del deepseekv2 @@ -187,7 +227,6 @@ def main(): print(tensorrt_llm.__version__) args = parse_arguments() - args.tp_size * args.pp_size if (args.moe_tp_size == -1 and args.moe_ep_size == -1): # moe default to tp-only args.moe_tp_size = args.tp_size diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index 34e13da5a13f..065e87b9d7bc 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -61,22 +61,35 @@ Below is the step-by-step to run DeepSeek-V3 with TensorRT-LLM. First the checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be build with TensorRT-LLM checkpoint. ```bash -# Convert Deepseek-v HF weights to TensorRT-LLM checkpoint format. +# Convert Deepseek-v HF weights to TensorRT-LLM checkpoint in BF16. python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ --dtype bfloat16 \ - --tp_size 32 + --tp_size 32 \ + --workers 8 # using multiple workers can accelerate the conversion process + -# With '--load_model_on_cpu' option if total GPU memory is insufficient +# Use Weight-Only Int8 quantization python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ - --output_dir ./trtllm_checkpoint_deepseek_v3_cpu_bf16 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ --dtype bfloat16 \ --tp_size 32 \ - --load_model_on_cpu -``` + --use_weight_only \ + --weight_only_precision int8 \ + --workers 8 +# Use Weight-Only Int4 quantization +python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --dtype bfloat16 \ + --tp_size 32 \ + --use_weight_only \ + --weight_only_precision int4 \ + --workers 8 +``` +We observe the checkpoint conversion time took hours, while use significant amount of CPU memory, please adjust the `--workers` parameter to balance your time and memory consumption. -We observe the checkpoint conversion time took exceed 10 hours, while use CPU memory required >= 1.3TB. +And **please note** that although Int8 and Int4 quantization can reduce the amount of total GPU memory needed, they may affect the accuracy to a certain degree. After the checkpoint conversion, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. diff --git a/examples/deepseek_v3/convert_checkpoint.py b/examples/deepseek_v3/convert_checkpoint.py new file mode 100644 index 000000000000..7cbc7935d225 --- /dev/null +++ b/examples/deepseek_v3/convert_checkpoint.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import os +import time +import traceback +from concurrent.futures import ThreadPoolExecutor, as_completed + +import tensorrt_llm +from tensorrt_llm._utils import release_gc +from tensorrt_llm.layers import MoeConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models import DeepseekV2ForCausalLM +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--model_dir', type=str, default=None, required=True) + parser.add_argument('--tp_size', + type=int, + default=1, + help='N-way tensor parallelism size') + parser.add_argument('--pp_size', + type=int, + default=1, + help='N-way pipeline parallelism size') + parser.add_argument( + '--moe_tp_size', + type=int, + default=-1, + help= + 'N-way tensor parallelism size for MoE, default is tp_size, which will do tp-only for MoE' + ) + parser.add_argument( + '--moe_ep_size', + type=int, + default=-1, + help= + 'N-way expert parallelism size for MoE, default is 1, which will do tp-only for MoE' + ) + parser.add_argument('--dtype', + type=str, + default='float16', + choices=['float32', 'bfloat16', 'float16']) + parser.add_argument('--load_model_on_cpu', + default=False, + action="store_true", + help='Choose to load HF cpkt into CPU') + parser.add_argument( + '--use_parallel_embedding', + action="store_true", + default=False, + help= + 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' + ) + parser.add_argument( + '--embedding_sharding_dim', + type=int, + default=0, + choices=[0, 1], + help= + 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0)' + 'To shard it along hidden dimension, set embedding_sharding_dim=1' + 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') + + parser.add_argument( + '--use_weight_only', + default=False, + action="store_true", + help='Quantize weights for the various GEMMs to INT4/INT8.' + 'See --weight_only_precision to set the precision') + parser.add_argument( + '--weight_only_precision', + const='int8', + type=str, + nargs='?', + default='int8', + choices=['int8', 'int4'], + help= + 'Define the precision for the weights when using weight-only quantization.' + 'You must also use --use_weight_only for that argument to have an impact.' + ) + + parser.add_argument('--output_dir', + type=str, + default='trtllm_checkpoint', + required=True, + help='The path to save the TensorRT-LLM checkpoint') + parser.add_argument( + '--workers', + type=int, + default=1, + help='The number of workers for converting checkpoint in parallel') + parser.add_argument( + '--moe_num_experts', + type=int, + default=0, + help='Specify the number of experts to use for MOE layers') + parser.add_argument( + '--moe_top_k', + type=int, + default=0, + help= + 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' + ) + parser.add_argument( + '--moe_renorm_mode', + type=int, + default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, + help= + 'Controls renormalization after gate logits. Check layers/moe.py for accepted values' + ) + parser.add_argument( + '--save_config_only', + action="store_true", + default=False, + help= + 'Only save the model config w/o read and converting weights, be careful, this is for debug only' + ) + parser.add_argument( + '--disable_weight_only_quant_plugin', + default=False, + action="store_true", + help= + 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' + 'You must also use --use_weight_only for that argument to have an impact' + ) + # Add quantization related feature later + args = parser.parse_args() + + return args + + +def precision_to_config(precision, quant_config) -> QuantConfig: + '''update config dict for weight-only quantization + ''' + quant_config = QuantConfig() + precision_to_algo = {'int8': QuantAlgo.W8A16, 'int4': QuantAlgo.W4A16} + quant_config.quant_algo = precision_to_algo.get(precision) + return quant_config + + +def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: + '''return config dict with quantization info based on the command line args + ''' + quant_config = QuantConfig() + if args.use_weight_only: + quant_config = precision_to_config(args.weight_only_precision, + quant_config) + return quant_config + + +def args_to_build_options(args): + return { + 'use_parallel_embedding': args.use_parallel_embedding, + 'embedding_sharding_dim': args.embedding_sharding_dim, + 'load_model_on_cpu': args.load_model_on_cpu + } + + +def execute(workers, func, args): + if workers == 1: + for rank, f in enumerate(func): + f(args, rank) + else: + with ThreadPoolExecutor(max_workers=workers) as p: + futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] + exceptions = [] + for future in as_completed(futures): + try: + future.result() + except Exception as e: + traceback.print_exc() + exceptions.append(e) + assert len( + exceptions + ) == 0, "Checkpoint conversion failed, please check error log." + + +def convert_and_save_hf(args): + model_dir = args.model_dir + world_size = args.tp_size * args.pp_size + # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. + # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, + # before the refactor is done. + override_fields = {} + override_fields.update(args_to_build_options(args)) + + quant_config = args_to_quant_config(args) + + def convert_and_save_rank(args, rank): + mapping = Mapping(world_size=world_size, + rank=rank, + tp_size=args.tp_size, + pp_size=args.pp_size, + moe_tp_size=args.moe_tp_size, + moe_ep_size=args.moe_ep_size) + + deepseekv2 = DeepseekV2ForCausalLM.from_hugging_face( + model_dir, + args.dtype, + mapping=mapping, + quant_config=quant_config, + **override_fields) + deepseekv2.save_checkpoint(args.output_dir, save_config=(rank == 0)) + del deepseekv2 + + execute(args.workers, [convert_and_save_rank] * world_size, args) + release_gc() + + +def main(): + print(tensorrt_llm.__version__) + args = parse_arguments() + + if (args.moe_tp_size == -1 and args.moe_ep_size == -1): + # moe default to tp-only + args.moe_tp_size = args.tp_size + args.moe_ep_size = 1 + elif (args.moe_tp_size == -1): + args.moe_tp_size = args.tp_size // args.moe_ep_size + elif (args.moe_ep_size == -1): + args.moe_ep_size = args.tp_size // args.moe_tp_size + assert (args.moe_tp_size * args.moe_ep_size == args.tp_size + ), "moe_tp_size * moe_ep_size must equal to tp_size" + + tik = time.time() + if not os.path.exists(args.output_dir): + os.makedirs(args.output_dir) + assert args.model_dir is not None + convert_and_save_hf(args) + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + print(f'Total time of converting checkpoints: {t}') + + +if __name__ == '__main__': + main() diff --git a/tensorrt_llm/layers/attention.py b/tensorrt_llm/layers/attention.py index 871ccbb119dd..4bfc86ca8c43 100755 --- a/tensorrt_llm/layers/attention.py +++ b/tensorrt_llm/layers/attention.py @@ -1913,11 +1913,6 @@ def __init__( max_position_embeddings=1024, rotary_embedding_base=10000.0, rotary_embedding_scaling=None, - rotary_embedding_beta_fast=32, - rotary_embedding_beta_slow=1, - rotary_embedding_mscale=1, - rotary_embedding_mscale_all_dim=0, - rotary_embedding_origin_max_position=4096, rotary_scaling=None, tp_group=None, tp_size=1, @@ -1944,6 +1939,7 @@ def __init__( enable_qkv=False) self.tp_size = tp_size + self.tp_rank = tp_rank if q_lora_rank is None: self.q_lora_rank = hidden_size @@ -1965,7 +1961,7 @@ def yarn_get_mscale(scale=1, mscale=1): return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 - assert self.rotary_scaling is not None + # assert self.rotary_scaling is not None if self.rotary_scaling is not None: mscale_all_dim = self.rotary_scaling.get("mscale_all_dim", 0) scaling_factor = self.rotary_scaling["factor"] @@ -1973,12 +1969,27 @@ def yarn_get_mscale(scale=1, mscale=1): mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) self.q_scaling = 1.0 / (mscale * mscale) - embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_for_deepseek_attention_plugin( - self.max_position_embeddings, self.qk_rope_head_dim, - self.rotary_embedding_base, self.rotary_scaling["factor"], - rotary_embedding_origin_max_position, rotary_embedding_beta_fast, - rotary_embedding_beta_slow, rotary_embedding_mscale, - rotary_embedding_mscale_all_dim) + embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_for_deepseek_attention_plugin( + self.max_position_embeddings, self.qk_rope_head_dim, + self.rotary_embedding_base, self.rotary_scaling["factor"], + rotary_scaling['original_max_position_embeddings'], + self.rotary_scaling["beta_fast"], + self.rotary_scaling["beta_slow"], self.rotary_scaling["mscale"], + mscale_all_dim) + else: + embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions( + self.max_position_embeddings, self.qk_rope_head_dim, + self.rotary_embedding_base) + embed_positions_for_gpt_attention = embed_positions_for_gpt_attention.squeeze( + 0) + sin, cos = np.split(embed_positions_for_gpt_attention, 2, 1) + cos_embed = np.expand_dims(np.concatenate((cos, cos), axis=1), + axis=2) + sin_embed = np.expand_dims(np.concatenate((sin, sin), axis=1), + axis=2) + embed_positions_for_gpt_attention = np.concatenate( + (cos_embed, sin_embed), axis=-1).reshape(1, -1) + self.register_parameter( 'embed_positions_for_gpt_attention', Parameter(embed_positions_for_gpt_attention, dtype='float32')) @@ -2205,3 +2216,88 @@ def forward(self, return (context, past_key_value) else: return context + + def postprocess(self, tllm_key, weights, **kwargs): + + def split(v, tp_size, idx, dim=0): + if tp_size == 1: + return v + if len(v.shape) == 1: + return torch.chunk(v, tp_size)[idx].contiguous() + else: + return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() + + if tllm_key.endswith("kv_b_proj"): + kv_b_proj = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.v_head_dim + ]) + splited_kv_b_proj = split(kv_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + k_nope_weight, v_weight = splited_kv_b_proj.split( + [self.qk_nope_head_dim, self.v_head_dim], + dim=1, + ) + kv_b_proj_weight = torch.concat([ + k_nope_weight.reshape( + self.num_attention_heads * self.qk_nope_head_dim, + self.kv_lora_rank), + v_weight.reshape(self.num_attention_heads * self.v_head_dim, + self.kv_lora_rank) + ], + dim=0) + return {tllm_key: kv_b_proj_weight} + elif tllm_key.endswith("q_b_proj"): + q_b_proj = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.qk_rope_head_dim + ]) + splited_q_b_proj = split(q_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + q_b_proj_weight = splited_q_b_proj.reshape( + self.num_attention_heads * + (self.qk_nope_head_dim + self.qk_rope_head_dim), + self.q_lora_rank) + return {tllm_key: q_b_proj_weight} + elif tllm_key.endswith("fused_q_proj"): + assert isinstance(weights, list) and len(weights) == 2 + q_b_proj = weights[0].unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.qk_rope_head_dim + ]) + splited_q_b_proj = split(q_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + kv_b_proj = weights[1].unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.v_head_dim + ]) + splited_kv_b_proj = split(kv_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + q_nope_weight, q_pe_weight = splited_q_b_proj.split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], + dim=1, + ) + k_nope_weight, _ = splited_kv_b_proj.split( + [self.qk_nope_head_dim, self.v_head_dim], + dim=1, + ) + fused_q_nope_weight = torch.einsum( + 'hdq,hdk->hkq', + q_nope_weight, + k_nope_weight, + ) + fused_q_weight = torch.cat( + [fused_q_nope_weight, q_pe_weight], + dim=1, + ).flatten(start_dim=0, end_dim=1) + return {tllm_key: fused_q_weight} + else: + return {tllm_key: weights} diff --git a/tensorrt_llm/layers/linear.py b/tensorrt_llm/layers/linear.py index 9224c066dacf..899ad446b8c3 100644 --- a/tensorrt_llm/layers/linear.py +++ b/tensorrt_llm/layers/linear.py @@ -426,6 +426,11 @@ def postprocess(self, tllm_key, weights, **kwargs): weights = w.reshape(-1, self.in_features) # Weight else: weights = w.reshape(-1) # Bias + elif tllm_key.endswith("fused_a.weight") and isinstance(weights, list): + weights = torch.cat( + [weights[0], weights[1]], + dim=0, + ) weights = weights.to(str_dtype_to_torch(self.dtype)) return {tllm_key: weights} diff --git a/tensorrt_llm/layers/moe.py b/tensorrt_llm/layers/moe.py index 2dfabe559c98..c1e83c933207 100755 --- a/tensorrt_llm/layers/moe.py +++ b/tensorrt_llm/layers/moe.py @@ -64,6 +64,11 @@ class ExpertScaleNormalizationMode(IntEnum): DEVICE_LIMITED = 3 DEVICE_LIMITED_RENORM = 4 + class TopKMethod(IntEnum): + GREEDY = 0, + GROUP_LIMITED_GREEDY = 1, + NOAUX_TC = 2 + num_experts: int = 0 shared_expert_intermediate_size: int = 0 @@ -71,6 +76,7 @@ class ExpertScaleNormalizationMode(IntEnum): normalization_mode: ExpertScaleNormalizationMode = ExpertScaleNormalizationMode.RENORMALIZE sparse_mixer_epsilon: float = 0.01 tp_mode: int = 0 + topk_method: TopKMethod = TopKMethod.GREEDY device_limited_n_group: int = 0 device_limited_topk_group: int = 0 @@ -499,6 +505,10 @@ def __init__(self, "router": "gate" } + if moe_config.topk_method == MoeConfig.TopKMethod.NOAUX_TC: + self.e_score_correction_bias = Parameter(shape=(self.num_experts, ), + dtype=trt.float32) + self.init_experts() self.max_low_rank = None @@ -542,6 +552,38 @@ def group_limited_greedy(self, logits): self.moe_config.device_limited_routed_scaling_factor return scores + def noaux_tc(self, logits): + n_group = self.moe_config.device_limited_n_group + scores = sigmoid(logits) + scores_with_bias = scores + self.e_score_correction_bias.value + scores_shape = [ + shape(scores_with_bias, i) for i in range(scores_with_bias.ndim()) + ] + group_scores = sum(topk(scores_with_bias.view( + concat(scores_shape[:-1] + [n_group, scores_shape[-1] // n_group])), + k=2, + dim=-1)[0], + dim=-1) + _, group_idx = topk(group_scores, + k=self.moe_config.device_limited_topk_group, + dim=-1) + group_mask = scatter(group_scores * 0, -1, group_idx, + group_idx.cast(group_scores.dtype) * 0 + 1) + score_mask = expand( + unsqueeze(group_mask, -1), + concat(scores_shape[:-1] + [n_group, scores_shape[-1] // n_group]), + ).view(concat(scores_shape)) + scores_with_bias = scores_with_bias * score_mask + + _, topk_idx = topk(scores_with_bias, k=self.moe_config.top_k, dim=-1) + new_mask = scatter(scores * 0, -1, topk_idx, + topk_idx.cast(scores.dtype) * 0 + 1) + scores = scores * new_mask + score_sum = sum(scores, dim=-1, keepdim=True) + 1e-20 + scores = scores / score_sum * \ + self.moe_config.device_limited_routed_scaling_factor + return scores + def forward(self, hidden_states, finished=None, @@ -556,11 +598,10 @@ def forward(self, 0, "moe_router") routing_input = cast(hidden_states, trt.float32) routing = self.router(routing_input, moe_router_lora_params) - if self.moe_config.normalization_mode in [ - MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED, - MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM - ]: + if self.moe_config.topk_method == MoeConfig.TopKMethod.GROUP_LIMITED_GREEDY: routing = self.group_limited_greedy(routing) + elif self.moe_config.topk_method == MoeConfig.TopKMethod.NOAUX_TC: + routing = self.noaux_tc(routing) output = self.forward_experts(hidden_states, routing, finished, lora_layer_params, side_stream_id) if side_stream_id != SideStreamIDType.disable: diff --git a/tensorrt_llm/models/__init__.py b/tensorrt_llm/models/__init__.py index 816f52354b0e..98f7994dc444 100755 --- a/tensorrt_llm/models/__init__.py +++ b/tensorrt_llm/models/__init__.py @@ -70,6 +70,7 @@ 'DeepseekForCausalLM', 'FalconConfig', 'DeepseekV2ForCausalLM', + 'DeepSeekV2Config' 'FalconForCausalLM', 'FalconModel', 'GPTConfig', diff --git a/tensorrt_llm/models/deepseek_v2/config.py b/tensorrt_llm/models/deepseek_v2/config.py new file mode 100644 index 000000000000..676418dacfc3 --- /dev/null +++ b/tensorrt_llm/models/deepseek_v2/config.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional, Union + +from transformers import AutoConfig + +from ...layers import MoeConfig +from ...mapping import Mapping +from ..modeling_utils import PretrainedConfig, QuantConfig + + +class DeepSeekV2Config(PretrainedConfig): + + def __init__(self, + *, + mlp_bias: bool = False, + attn_bias: bool = False, + rotary_base: float = 10000.0, + rotary_scaling: Optional[dict] = None, + residual_mlp: bool = False, + disable_weight_only_quant_plugin: bool = False, + moe: Optional[Union[MoeConfig, dict]] = None, + remove_duplicated_kv_heads: bool = False, + **kwargs): + self.mlp_bias = mlp_bias + self.attn_bias = attn_bias + self.rotary_base = rotary_base + self.rotary_scaling = rotary_scaling + self.residual_mlp = residual_mlp + self.disable_weight_only_quant_plugin = disable_weight_only_quant_plugin + if isinstance(moe, dict): + moe = MoeConfig.from_dict(moe) + assert isinstance(moe, MoeConfig) + self.moe = moe.validate() + self.remove_duplicated_kv_heads = remove_duplicated_kv_heads + self.fc_after_embed = False + self.use_input_layernorm_in_first_layer = True + self.use_last_layernorm = True + self.layer_idx_offset = 0 + + super().__init__(**kwargs) + + def to_dict(self): + output = super().to_dict() + # Serialize the fields added in LLaMAConfig + output['mlp_bias'] = self.mlp_bias + output['attn_bias'] = self.attn_bias + output['rotary_base'] = self.rotary_base + output['rotary_scaling'] = self.rotary_scaling + output['residual_mlp'] = self.residual_mlp + output[ + 'disable_weight_only_quant_plugin'] = self.disable_weight_only_quant_plugin + output['fc_after_embed'] = self.fc_after_embed + output[ + 'use_input_layernorm_in_first_layer'] = self.use_input_layernorm_in_first_layer + output['use_last_layernorm'] = self.use_last_layernorm + output['layer_idx_offset'] = self.layer_idx_offset + output['moe'] = self.moe.to_dict() + return output + + @classmethod + def from_hugging_face(cls, + model_dir: str, + dtype: str = 'auto', + mapping: Optional[Mapping] = None, + quant_config: Optional[QuantConfig] = None, + **kwargs): + trust_remote_code = kwargs.pop('trust_remote_code', True) + hf_config = AutoConfig.from_pretrained( + model_dir, trust_remote_code=trust_remote_code) + + moe_routed_scaling_factor = hf_config.routed_scaling_factor + moe_top_k = hf_config.num_experts_per_tok + assert moe_routed_scaling_factor > 0, 'routed_scaling_factor should be greater than 0' + if hf_config.topk_method == 'group_limited_greedy': + moe_topk_method = MoeConfig.TopKMethod.GROUP_LIMITED_GREEDY + if moe_top_k > 1 and hf_config.norm_topk_prob: + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM + else: + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED + elif hf_config.topk_method == 'greedy': + moe_topk_method = MoeConfig.TopKMethod.GREEDY + assert moe_routed_scaling_factor == 1.0, 'The combination of topk_method == greedy and routed_scaling_factor != 1.0 is not supported' + if moe_top_k > 1 and hf_config.norm_topk_prob: + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE + else: + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.NONE + elif hf_config.topk_method == 'noaux_tc': + moe_topk_method = MoeConfig.TopKMethod.NOAUX_TC + moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED + else: + raise AssertionError( + f'Unsupported topk_method in hf_config: {hf_config.topk_method}' + ) + + rotary_scaling = None + if hf_config.rope_scaling: + rotary_scaling = { + 'beta_fast': + hf_config.rope_scaling['beta_fast'], + 'beta_slow': + hf_config.rope_scaling['beta_slow'], + 'factor': + hf_config.rope_scaling['factor'], + 'mscale': + hf_config.rope_scaling['mscale'], + 'mscale_all_dim': + hf_config.rope_scaling['mscale_all_dim'], + 'original_max_position_embeddings': + hf_config.rope_scaling['original_max_position_embeddings'], + 'type': + 'yarn', + } + + moe_config = MoeConfig( + num_experts=hf_config.n_routed_experts, + shared_expert_intermediate_size=hf_config.n_shared_experts * + hf_config.moe_intermediate_size, + top_k=moe_top_k, + normalization_mode=moe_renorm_mode, + device_limited_n_group=hf_config.n_group, + device_limited_topk_group=hf_config.topk_group, + device_limited_routed_scaling_factor=moe_routed_scaling_factor, + topk_method=moe_topk_method) + moe_config.validate() + + return cls( + architecture=hf_config.architectures[0], + dtype=dtype, + num_hidden_layers=hf_config.num_hidden_layers, + num_attention_heads=hf_config.num_attention_heads, + hidden_size=hf_config.hidden_size, + intermediate_size=hf_config.intermediate_size, + num_key_value_heads=hf_config.num_key_value_heads, + vocab_size=hf_config.vocab_size, + position_embedding_type='rope_gpt_neox', + max_position_embeddings=hf_config.max_position_embeddings, + hidden_act='swiglu', + norm_epsilon=hf_config.rms_norm_eps, + rotary_base=hf_config.rope_theta, + rotary_scaling=rotary_scaling, # TODO: modify this + moe_inter_size=hf_config.moe_intermediate_size, + moe=moe_config, + mapping=mapping, + quantization=quant_config, + kv_lora_rank=hf_config.kv_lora_rank, + q_lora_rank=hf_config.q_lora_rank, + qk_nope_head_dim=hf_config.qk_nope_head_dim, + qk_rope_head_dim=hf_config.qk_rope_head_dim, + v_head_dim=hf_config.v_head_dim, + topk_method=hf_config.topk_method, + first_k_dense_replace=hf_config.first_k_dense_replace, + moe_layer_freq=hf_config.moe_layer_freq, + coring_func=hf_config.scoring_func, + fp8_format=False, + **kwargs) diff --git a/tensorrt_llm/models/deepseek_v2/convert.py b/tensorrt_llm/models/deepseek_v2/convert.py index e9222e6ac23c..158b527fc70f 100755 --- a/tensorrt_llm/models/deepseek_v2/convert.py +++ b/tensorrt_llm/models/deepseek_v2/convert.py @@ -18,141 +18,16 @@ import torch from transformers import AutoConfig, AutoModelForCausalLM -from tensorrt_llm.layers import MoeConfig - from ..._utils import pad_vocab_size, release_gc -from ...mapping import Mapping +from ...layers import MoeConfig +from ...quantization import QuantAlgo from ..convert_utils import get_tllm_linear_weight +from .config import DeepSeekV2Config # `Override num_hidden_layers` used for reduce number of hidden layers in DeepseekV2ForCausalLM for debug purpose OVERRIDE_HIDDEN_LAYERS = None # 2 -## Convert config parameters to dict -def create_trt_config_from_hf(model_dir, - dtype, - mapping: Mapping, - override_fields: dict = {}): - config = {} - assert isinstance(model_dir, str) - hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) - # Override num_hidden_layers - if OVERRIDE_HIDDEN_LAYERS is not None: - hf_config.num_hidden_layers = OVERRIDE_HIDDEN_LAYERS - print( - f'Override hidden layers to {hf_config.num_hidden_layers} for DeepseekV2ForCausalLM' - ) - dtype = dtype - n_layer = hf_config.num_hidden_layers - n_head = hf_config.num_attention_heads - n_embd = hf_config.hidden_size - inter_size = hf_config.intermediate_size - n_kv_head = hf_config.num_key_value_heads - vocab_size = hf_config.vocab_size - n_positions = hf_config.max_position_embeddings - hidden_act = 'swiglu' # TRT-LLM request make gated activation explicit for MOE implementation - rotary_base = hf_config.rope_theta - rms_norm_eps = hf_config.rms_norm_eps - rotary_scaling_beta_fast = hf_config.rope_scaling['beta_fast'] - rotary_scaling_beta_slow = hf_config.rope_scaling['beta_slow'] - rotary_scaling_factor = hf_config.rope_scaling['factor'] - rotary_scaling_mscale = hf_config.rope_scaling['mscale'] - rotary_scaling_mscale_all_dim = hf_config.rope_scaling['mscale_all_dim'] - rotary_scaling_original_max_position_embeddings = hf_config.rope_scaling[ - 'original_max_position_embeddings'] - rotary_scaling_type = 'yarn' - kv_lora_rank = hf_config.kv_lora_rank - q_lora_rank = hf_config.q_lora_rank - qk_nope_head_dim = hf_config.qk_nope_head_dim - qk_rope_head_dim = hf_config.qk_rope_head_dim - v_head_dim = hf_config.v_head_dim - moe_num_experts = hf_config.n_routed_experts - moe_inter_size = hf_config.moe_intermediate_size - moe_num_shared_experts = hf_config.n_shared_experts - moe_top_k = hf_config.num_experts_per_tok - moe_n_group = hf_config.n_group - moe_topk_group = hf_config.topk_group - moe_routed_scaling_factor = hf_config.routed_scaling_factor - assert moe_routed_scaling_factor > 0, 'routed_scaling_factor should be greater than 0' - if hf_config.topk_method == 'group_limited_greedy': - if moe_top_k > 1 and hf_config.norm_topk_prob: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM - else: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED - elif hf_config.topk_method == 'greedy': - assert moe_routed_scaling_factor == 1.0, 'The combination of topk_method == greedy and routed_scaling_factor != 1.0 is not supported' - if moe_top_k > 1 and hf_config.norm_topk_prob: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE - else: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.NONE - else: - raise AssertionError( - 'Unsupported topk_method in hf_config: {hf_config.topk_method}') - - config = { - 'architecture': 'DeepseekV2ForCausalLM', - 'dtype': dtype, - 'logits_type': 'float32', - 'num_hidden_layers': n_layer, - 'num_attention_heads': n_head, - 'hidden_size': n_embd, - 'intermediate_size': inter_size, - 'num_key_value_heads': n_kv_head, - 'vocab_size': vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': n_positions, - 'hidden_act': hidden_act, - 'rotary_base': rotary_base, - 'norm_epsilon': rms_norm_eps, - 'rotary_scaling': { - 'beta_fast': rotary_scaling_beta_fast, - 'beta_slow': rotary_scaling_beta_slow, - 'factor': rotary_scaling_factor, - 'mscale': rotary_scaling_mscale, - 'mscale_all_dim': rotary_scaling_mscale_all_dim, - 'original_max_position_embeddings': - rotary_scaling_original_max_position_embeddings, - 'type': rotary_scaling_type, - }, - 'mapping': { - 'world_size': mapping.tp_size * mapping.pp_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - 'moe_tp_size': mapping.moe_tp_size, - 'moe_ep_size': mapping.moe_ep_size, - }, - 'kv_lora_rank': kv_lora_rank, - 'q_lora_rank': q_lora_rank, - 'qk_nope_head_dim': qk_nope_head_dim, - 'qk_rope_head_dim': qk_rope_head_dim, - 'v_head_dim': v_head_dim, - 'moe_num_experts': moe_num_experts, - 'moe_inter_size': moe_inter_size, - 'moe_num_shared_experts': moe_num_shared_experts, - 'moe_top_k': moe_top_k, - 'moe_renorm_mode': moe_renorm_mode, - 'moe_n_group': moe_n_group, - 'moe_topk_group': moe_topk_group, - 'moe_routed_scaling_factor': moe_routed_scaling_factor, - } - - config.update(override_fields) - - moe_config = MoeConfig( - num_experts=config['moe_num_experts'], - shared_expert_intermediate_size=config['moe_num_shared_experts'] * - config['moe_inter_size'], - top_k=config['moe_top_k'], - normalization_mode=config['moe_renorm_mode'], - device_limited_n_group=config['moe_n_group'], - device_limited_topk_group=config['moe_topk_group'], - device_limited_routed_scaling_factor=config['moe_routed_scaling_factor'] - ) - moe_config.validate() - - return config - - ## Get HF model def load_hf_deepseek(model_dir, load_model_on_cpu=False): hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) @@ -225,42 +100,42 @@ def get_param_weight(weight, prefix): return results -def convert_deepseekv2(hf_model, - config, - mapping, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0): +def load_weights_from_hf_model(hf_model, + config: DeepSeekV2Config, + use_parallel_embedding=False, + sharding_dim=0): + quant_algo = config.quantization.quant_algo + use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] + if quant_algo == QuantAlgo.W8A16: + plugin_weight_only_quant_type = torch.int8 + elif quant_algo == QuantAlgo.W4A16: + plugin_weight_only_quant_type = torch.quint4x2 + else: + plugin_weight_only_quant_type = None + use_gemm_woq_plugin = True weights = {} tik = time.time() model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - moe_config = MoeConfig( - num_experts=config['moe_num_experts'], - shared_expert_intermediate_size=config['moe_num_shared_experts'] * - config['moe_inter_size'], - top_k=config['moe_top_k'], - normalization_mode=config['moe_renorm_mode'], - device_limited_n_group=config['moe_n_group'], - device_limited_topk_group=config['moe_topk_group'], - device_limited_routed_scaling_factor=config['moe_routed_scaling_factor'] - ) - - layers_range = mapping.pp_layers(config['num_hidden_layers']) + dtype = getattr(torch, config.dtype) + + mapping = config.mapping + moe_config = config.moe + first_k_dense_replace = config.first_k_dense_replace + layers_range = mapping.pp_layers(config.num_hidden_layers) def convert_layer(l): prefix = f'model.layers.{l}.' trtllm_prex = f'transformer.layers.{l - layers_range[0]}.' # Fuse matrices for compression # Split matrices for decompression - q_lora_rank = config['q_lora_rank'] - kv_lora_rank = config['kv_lora_rank'] - num_heads = config['num_attention_heads'] - qk_nope_head_dim = config['qk_nope_head_dim'] - qk_rope_head_dim = config['qk_rope_head_dim'] - v_head_dim = config['v_head_dim'] - hidden_size = config['hidden_size'] + q_lora_rank = config.q_lora_rank + kv_lora_rank = config.kv_lora_rank + num_heads = config.num_attention_heads + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + hidden_size = config.hidden_size if q_lora_rank is not None: q_a_proj_weight = get_weight(model_params, @@ -392,15 +267,19 @@ def convert_layer(l): get_param_weight(kv_b_proj_weight, trtllm_prex + 'attention.kv_b_proj')) weights.update( - get_tllm_linear_weight(o_proj_weight, - trtllm_prex + 'attention.dense.')) + get_tllm_linear_weight( + o_proj_weight, + trtllm_prex + 'attention.dense.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) if q_lora_rank is not None: weights.update( get_tllm_linear_weight(q_a_layernorm_weight, trtllm_prex + 'attention.q_a_layernorm.')) - if moe_config.has_moe() and l > 0: + if moe_config.has_moe() and l >= first_k_dense_replace: rank_experts = list(range(moe_config.num_experts)) if mapping.has_moe_ep(): rank_experts = mapping.ep_experts(moe_config.num_experts) @@ -439,14 +318,22 @@ def convert_layer(l): moe_experts_down_proj_weights = get_weight( model_params, prefix + 'mlp.experts.down_proj', dtype) weights.update( - get_tllm_linear_weight(moe_experts_down_proj_weights, - trtllm_prex + 'mlp.proj.')) + get_tllm_linear_weight( + moe_experts_down_proj_weights, + trtllm_prex + 'mlp.proj.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) ## mlp.experts.up_gate.weight moe_experts_up_gate_proj_weights = get_weight( model_params, prefix + 'mlp.experts.up_gate_proj', dtype) weights.update( - get_tllm_linear_weight(moe_experts_up_gate_proj_weights, - trtllm_prex + 'mlp.fc.')) + get_tllm_linear_weight( + moe_experts_up_gate_proj_weights, + trtllm_prex + 'mlp.fc.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) ## MOE hardcoded routing_input into trt.float32, please refer to moe.py line 397 moe_experts_gate_weights = get_weight(model_params, prefix + 'mlp.gate', @@ -455,6 +342,15 @@ def convert_layer(l): get_tllm_linear_weight(moe_experts_gate_weights, trtllm_prex + 'mlp.router.')) + if moe_config.topk_method == MoeConfig.TopKMethod.NOAUX_TC: + e_score_correction_bias = get_weight( + model_params, prefix + 'mlp.gate.e_score_correction_bias', + torch.float32, '') + weights.update( + get_param_weight( + e_score_correction_bias, + trtllm_prex + 'mlp.e_score_correction_bias')) + if moe_config.shared_expert_intermediate_size > 0: shared_moe_up_proj_weights = get_weight( model_params, prefix + 'mlp.shared_experts.up_proj', dtype) @@ -487,13 +383,21 @@ def convert_layer(l): weights.update( get_tllm_linear_weight( shared_moe_gate_up_proj_weights, - trtllm_prex + 'mlp.shared_expert.fc.')) + trtllm_prex + 'mlp.shared_expert.fc.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type= + plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) ## mlp.shared_experts.down_proj.weight weights.update( get_tllm_linear_weight( shared_moe_down_proj_weights, - trtllm_prex + 'mlp.shared_expert.proj.')) + trtllm_prex + 'mlp.shared_expert.proj.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type= + plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) else: ## Current MLP layer is only one, if it goes large consider to do fuse @@ -504,7 +408,12 @@ def convert_layer(l): mapping.tp_rank, dim=0) weights.update( - get_tllm_linear_weight(split_gate, trtllm_prex + 'mlp.gate.')) + get_tllm_linear_weight( + split_gate, + trtllm_prex + 'mlp.gate.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', dtype) @@ -513,7 +422,12 @@ def convert_layer(l): mapping.tp_rank, dim=0) weights.update( - get_tllm_linear_weight(split_fc, trtllm_prex + 'mlp.fc.')) + get_tllm_linear_weight( + split_fc, + trtllm_prex + 'mlp.fc.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', dtype) @@ -522,7 +436,12 @@ def convert_layer(l): mapping.tp_rank, dim=1) weights.update( - get_tllm_linear_weight(split_proj, trtllm_prex + 'mlp.proj.')) + get_tllm_linear_weight( + split_proj, + trtllm_prex + 'mlp.proj.', + use_weight_only=use_weight_only, + plugin_weight_only_quant_type=plugin_weight_only_quant_type, + use_gemm_woq_plugin=use_gemm_woq_plugin)) # Layer norms do not use tensor parallelism input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', @@ -540,11 +459,11 @@ def convert_layer(l): if hf_model.config.tie_word_embeddings: # lm_head.weight has the same weights as embedding if mapping.is_last_pp_rank(): - if config['vocab_size'] % mapping.tp_size != 0: + if config.vocab_size % mapping.tp_size != 0: # padding - vocab_size_padded = pad_vocab_size(config['vocab_size'], + vocab_size_padded = pad_vocab_size(config.vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - config['vocab_size'] + pad_width = vocab_size_padded - config.vocab_size v = torch.nn.functional.pad(v, (0, 0, 0, pad_width), 'constant', 0) weights['lm_head.weight'] = split(v, mapping.tp_size, @@ -559,11 +478,11 @@ def convert_layer(l): lm_head_weights = get_weight(model_params, 'lm_head', dtype) if mapping.is_last_pp_rank(): - if config['vocab_size'] % mapping.tp_size != 0: + if config.vocab_size % mapping.tp_size != 0: # padding - vocab_size_padded = pad_vocab_size(config['vocab_size'], + vocab_size_padded = pad_vocab_size(config.vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - config['vocab_size'] + pad_width = vocab_size_padded - config.vocab_size lm_head_weights = torch.nn.functional.pad(lm_head_weights, (0, 0, 0, pad_width), 'constant', diff --git a/tensorrt_llm/models/deepseek_v2/model.py b/tensorrt_llm/models/deepseek_v2/model.py index 16b13ad6883c..4fa7798b3263 100755 --- a/tensorrt_llm/models/deepseek_v2/model.py +++ b/tensorrt_llm/models/deepseek_v2/model.py @@ -13,26 +13,31 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os from typing import Optional import torch +from tqdm import tqdm from ..._utils import pad_vocab_size, torch_dtype_to_str from ...functional import Tensor, non_gated_version, recv, send from ...layers import (MOE, AttentionMaskType, ColumnLinear, - DeepseekV2Attention, Embedding, GatedMLP, MoeConfig, + DeepseekV2Attention, Embedding, GatedMLP, PositionEmbeddingType, RmsNorm, SharedMoE) +from ...layers.moe import MOEWeightWrapper from ...mapping import Mapping from ...module import Module from ...plugin import init_all_reduce_helper +from ..model_weights_loader import ModelWeightsLoader from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) -from .convert import convert_deepseekv2, create_trt_config_from_hf + QuantConfig) +from .config import DeepSeekV2Config +from .convert import load_hf_deepseek, load_weights_from_hf_model class DeepseekV2DecoderLayer(Module): - def __init__(self, config: PretrainedConfig, layer_idx: int): + def __init__(self, config: DeepSeekV2Config, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.config = config @@ -61,13 +66,6 @@ def __init__(self, config: PretrainedConfig, layer_idx: int): position_embedding_type=PositionEmbeddingType.learned_absolute, rotary_embedding_base=config.rotary_base, rotary_embedding_scaling=None, - rotary_embedding_beta_fast=config.rotary_scaling['beta_fast'], - rotary_embedding_beta_slow=config.rotary_scaling['beta_slow'], - rotary_embedding_mscale=config.rotary_scaling['mscale'], - rotary_embedding_mscale_all_dim=config. - rotary_scaling['mscale_all_dim'], - rotary_embedding_origin_max_position=config. - rotary_scaling['original_max_position_embeddings'], rotary_scaling=config.rotary_scaling, tp_group=config.mapping.tp_group, tp_size=config.mapping.tp_size, @@ -80,21 +78,12 @@ def __init__(self, config: PretrainedConfig, layer_idx: int): ### Distinguish dense MLP and MoE MLP # dense_config = DenseConfig(intermediate_size=config.intermediate_size) - moe_config = MoeConfig( - num_experts=config.moe_num_experts, - shared_expert_intermediate_size=config.moe_num_shared_experts * - config.moe_inter_size, - top_k=config.moe_top_k, - normalization_mode=config.moe_renorm_mode, - device_limited_n_group=config.moe_n_group, - device_limited_topk_group=config.moe_topk_group, - device_limited_routed_scaling_factor=config. - moe_routed_scaling_factor) + moe_config = config.moe # layer_config = LayerMLPConfig(config=[dense_config, moe_config], moe_layer_idx_min=0, # moe_layer_idx_max=config.num_hidden_layers, # total_num_layers=config.num_hidden_layers) - if moe_config.num_experts > 0 and layer_idx > 0: + if moe_config.num_experts > 0 and layer_idx >= config.first_k_dense_replace: hidden_act = config.hidden_act mlp_hidden_size = config.moe_inter_size mlp_kwargs = {'moe_config': moe_config, 'mapping': config.mapping} @@ -118,6 +107,7 @@ def __init__(self, config: PretrainedConfig, layer_idx: int): bias=False, tp_group=config.mapping.tp_group, tp_size=config.mapping.tp_size, + quant_mode=config.quant_mode, **mlp_kwargs) ### Pose layernorm in Deepseek v2 is same as Llama @@ -159,7 +149,7 @@ def forward(self, class DeepseekV2Model(Module): - def __init__(self, config: PretrainedConfig) -> None: + def __init__(self, config: DeepSeekV2Config) -> None: super().__init__() init_all_reduce_helper() # enable use_customer_all_reduce self.dtype = config.dtype @@ -221,8 +211,9 @@ def forward(self, class DeepseekV2ForCausalLM(DecoderModelForCausalLM): + config_class = DeepSeekV2Config - def __init__(self, config: PretrainedConfig): + def __init__(self, config: DeepSeekV2Config): transformer = DeepseekV2Model(config) vocab_size_padded = pad_vocab_size(config.vocab_size, config.mapping.tp_size) @@ -241,22 +232,89 @@ def __init__(self, config: PretrainedConfig): @classmethod def from_hugging_face(cls, - hf_model, model_dir, dtype: str = 'auto', mapping: Optional[Mapping] = None, - override_fields={}, + quant_config: Optional[QuantConfig] = None, **kwargs): - assert hf_model is not None + load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) if mapping is None: mapping = Mapping() - config = create_trt_config_from_hf(model_dir, - dtype, - mapping=mapping, - override_fields=override_fields) - print(config) - pretrained_config = PretrainedConfig.from_dict(config) - pretrained_config.set_rank(mapping.rank) # TODO:remove this hack + config = DeepSeekV2Config.from_hugging_face(model_dir, + dtype=dtype, + mapping=mapping, + quant_config=quant_config, + **kwargs) + if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: + if config.q_lora_rank is not None: # Deepseek-V2&V3 + custom_dict = { + "fused_a": ["q_a_proj", "kv_a_proj_with_mqa"], + "q_a_layernrom": "q_a_layernorm", + "kv_a_layernorm": "kv_a_layernorm", + "q_b_proj": "q_b_proj.weight", + "kv_b_proj": "kv_b_proj.weight", + "fused_q_proj": ["q_b_proj.weight", "kv_b_proj.weight"], + "shared_expert": "shared_experts", + "e_score_correction_bias": "gate.e_score_correction_bias", + } + else: # Deepseek-V2-Lite + custom_dict = { + "fused_a": "kv_a_proj_with_mqa", + "kv_a_layernorm": "kv_a_layernorm", + "q_b_proj": "q_proj.weight", + "kv_b_proj": "kv_b_proj.weight", + "fused_q_proj": ["q_proj.weight", "kv_b_proj.weight"], + "shared_expert": "shared_experts", + "e_score_correction_bias": "gate.e_score_correction_bias", + } + + loader = ModelWeightsLoader(model_dir, custom_dict) + model = cls(config) + for tllm_key, _ in model.named_parameters(): + sub_module = model + for attr in tllm_key.split(".")[:-1]: + sub_module = getattr(sub_module, attr) + if "router" in tllm_key or isinstance(sub_module, + MOEWeightWrapper): + sub_module_dic = sub_module.tllm_to_externel_key_dict + sub_module_dic["mlp"] = "mlp" + if "fc" in sub_module_dic.keys(): + sub_module_dic["fc"] = [ + hf_keyword.replace("w1", "gate_proj") + for hf_keyword in sub_module_dic["fc"] + ] + sub_module_dic["fc"] = [ + hf_keyword.replace("w3", "up_proj") + for hf_keyword in sub_module_dic["fc"] + ] + if "proj" in sub_module_dic.keys(): + sub_module_dic["proj"] = [ + hf_keyword.replace("w2", "down_proj") + for hf_keyword in sub_module_dic["proj"] + ] + sub_module.tllm_to_externel_key_dict = sub_module_dic + + def concat_gate_up_proj(weights): + return torch.cat(weights, dim=-2) + + loader.update_key_mapping(model) + tllm_weights = {} + for tllm_key, _ in tqdm(model.named_parameters()): + if tllm_key.endswith("shared_expert.fc.weight"): + updated_map = loader.tllm_to_externel_key_dict + updated_map["fc"] = ["up_proj", "gate_proj"] + loader.tllm_to_externel_key_dict = updated_map + tllm_weights.update( + loader.load(tllm_key, concat_gate_up_proj)) + else: + tllm_weights.update(loader.load(tllm_key)) + loader.fill(tllm_weights) + else: + hf_model = load_hf_deepseek(model_dir, load_model_on_cpu) + weights = load_weights_from_hf_model(hf_model, config) + model = cls(config) + model.load(weights) + return model if dtype == 'auto': dtype = getattr(config, 'torch_dtype', None) diff --git a/tensorrt_llm/models/model_weights_loader.py b/tensorrt_llm/models/model_weights_loader.py index 737dfc6edae1..2c822cb2a10a 100644 --- a/tensorrt_llm/models/model_weights_loader.py +++ b/tensorrt_llm/models/model_weights_loader.py @@ -286,7 +286,6 @@ def load(self, ] else: v = self.load_tensor(external_key, tp_size, tp_dim, tp_rank) - if preprocess is not None: v = preprocess(v) @@ -309,7 +308,7 @@ def load(self, weight_dict = {tllm_key: v} for k, v in weight_dict.items(): - if not v.is_contiguous(): + if v is not None and not v.is_contiguous(): weight_dict[k] = v.contiguous() return weight_dict @@ -353,6 +352,8 @@ def fill(self, weights): for tllm_key, param in self.model.named_parameters(): if param.is_buffer: continue + if weights[tllm_key] is None: + continue w_shape = weights[tllm_key].shape if w_shape != param.shape: logger.warning( diff --git a/tensorrt_llm/quantization/quantize.py b/tensorrt_llm/quantization/quantize.py index dc54899debbc..17efc23a9b00 100644 --- a/tensorrt_llm/quantization/quantize.py +++ b/tensorrt_llm/quantization/quantize.py @@ -35,6 +35,7 @@ def quantize_layers( '*position_embedding', '*block_embedding', '*shared_expert_gate', + '*fused_a', ] for name, module, parent in model.named_modules_with_parent(): From da62d7be24551fd997e27c7e8dd2708c2829dd95 Mon Sep 17 00:00:00 2001 From: yingcanw Date: Tue, 24 Dec 2024 23:59:41 -0800 Subject: [PATCH 03/12] update README --- examples/deepseek_v3/README.md | 47 ++++++++++++++++++++-------------- 1 file changed, 28 insertions(+), 19 deletions(-) diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index 065e87b9d7bc..9a9ea75abc90 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -1,19 +1,37 @@ # Deepseek-v3 -This document shows how to build and run `DeepSeek-v3` model in TensorRT-LLM. +This document shows how to build and run the `DeepSeek-v3` model in TensorRT-LLM. -***Current branch is a preview version only for supporting DeepSeek-V3 in bf16 precision, which will be part of the TRT-LLM official release in 0.18+ version.*** +***The current branch is a preview version only for supporting DeepSeek-V3, which will be part of the TRT-LLM official release in 0.18+ version.*** - [Deepseek-V3](#deepseek-v3) + - [Support Matrix](#support-matrix) - [Prerequisite](#prerequisite) - [Hardware](#hardware) - [Overview](#overview) - - [Support Matrix](#support-matrix) - [Usage](#usage) - [Build TensorRT engine(s)](#build-tensorrt-engines) +## Support Matrix + +| Model | FP16 | BF16 | FP8 | W8A16 | W4A16 | TP | EP | IB | +| :------------- | :---: | :---: | :---: | :-----: | :-----: | :-----: | :-----: | :-----: | +| DeepSeek-V3 | Y | Y | . | Y | Y | Y | Y | Y | + + +- W8A16: INT8 Weight-Only +- W4A16: INT4 Weight-Only +- TP: Tensor Parallel +- EP: Expert Parallel +- IB: Inflight Batching + +***Please Note:*** +- Prefer using BF16 over FP16 for DeepSeek-V3 since model original training precision is FP8 and we found direct convert FP8 -> FP16 may cause unknown accuracy issues. +- Although Int8 and Int4 weight-only quantization can reduce the amount of total GPU memory needed, they may affect the accuracy to a certain degree. + + ## Prerequisite First, please download DeepSeek-V3 weights from HF https://huggingface.co/deepseek-ai/DeepSeek-V3. @@ -25,9 +43,9 @@ git clone https://huggingface.co/deepseek-ai/DeepSeek-V3 ## Hardware -The DeepSeek-V3 model requires least 32x80G GPU memory, model contains 660B parameters roughly 1.3TB memory (with BF16 precision). +The DeepSeek-V3 model requires at least 32x80G GPU memory, model contains 660B parameters, roughly 1.3TB memory (with BF16 precision). -***Caution: Current TRT-LLM MLA kernel only support Hopper architecture (SM90). Ampere architecture (SM80 & SM86) will be supported in next release.*** +***Caution: Current TRT-LLM MLA kernel only supports Hopper architecture (SM90). Ampere architecture (SM80 & SM86) will be supported in the future release.*** ## Overview @@ -37,28 +55,20 @@ The TensorRT-LLM DeepSeek-V3 implementation can be found in [tensorrt_llm/models In addition, there are three shared files in the parent folder [`examples`](../) can be used for inference and evaluation: -* [`../run.py`](../run.py) to run the model inference output by given an input text. -* [`../summarize.py`](../summarize.py) to summarize the article from [cnn_dailmail](https://huggingface.co/datasets/cnn_dailymail) dataset, it can running the summarize from HF model and TensorRT-LLM model. +* [`../run.py`](../run.py) to run the model inference output by giving an input text. +* [`../summarize.py`](../summarize.py) to summarize the article from [cnn_dailmail](https://huggingface.co/datasets/cnn_dailymail) dataset, it can run the summarize from the HF model and TensorRT-LLM model. * [`../mmlu.py`](../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT-LLM model on the MMLU dataset. -## Support Matrix - -- [x] BF16 -- [ ] FP8 -- [x] Tensor Parallelism -- [ ] Expert Parallelism - -***Caution: prefer using BF16 over FP16 for DeepSeek-V3 since model original training precision is FP8 and we found direct convert FP8 -> FP16 may cause unknown accuracy issues.*** ## Usage -The TensorRT-LLM DeepSeek-V3 example code locates at [examples/deepseek_v3](./). It takes PyTorch weights as input, and builds corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. +The TensorRT-LLM DeepSeek-V3 example code is located at [examples/deepseek_v3](./). It takes PyTorch weights as input, and builds corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. ### Build TensorRT engine(s) Below is the step-by-step to run DeepSeek-V3 with TensorRT-LLM. -First the checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be build with TensorRT-LLM checkpoint. +First the checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with TensorRT-LLM checkpoint. ```bash # Convert Deepseek-v HF weights to TensorRT-LLM checkpoint in BF16. @@ -87,9 +97,8 @@ python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ --weight_only_precision int4 \ --workers 8 ``` -We observe the checkpoint conversion time took hours, while use significant amount of CPU memory, please adjust the `--workers` parameter to balance your time and memory consumption. +We observed the checkpoint conversion time took hours, while using a significant amount of CPU memory, please adjust the `--workers` parameter to balance your time and memory consumption. -And **please note** that although Int8 and Int4 quantization can reduce the amount of total GPU memory needed, they may affect the accuracy to a certain degree. After the checkpoint conversion, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. From 911c30e456a4c25641e39c5e3598c5a4c071bde7 Mon Sep 17 00:00:00 2001 From: yingcanw Date: Wed, 25 Dec 2024 00:05:34 -0800 Subject: [PATCH 04/12] add FP8 notes --- examples/deepseek_v3/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index 9a9ea75abc90..d6f240d10bbb 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -26,6 +26,7 @@ This document shows how to build and run the `DeepSeek-v3` model in TensorRT-LLM - TP: Tensor Parallel - EP: Expert Parallel - IB: Inflight Batching +- FP8: Support for FP8 is currently in progress and will be released soon ***Please Note:*** - Prefer using BF16 over FP16 for DeepSeek-V3 since model original training precision is FP8 and we found direct convert FP8 -> FP16 may cause unknown accuracy issues. From e3077bcd4b0ed696cc74990665d95de23fb44747 Mon Sep 17 00:00:00 2001 From: Hanyue He Date: Wed, 25 Dec 2024 17:28:51 +0800 Subject: [PATCH 05/12] Update run.py result --- examples/deepseek_v3/README.md | 40 ++++++++++++++++++++++++---------- 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index d6f240d10bbb..f0067c94e0d5 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -122,23 +122,29 @@ Test the engine with [run.py](../run.py) script: ```bash mpirun --allow-run-as-root -np 32 python ../run.py --engine_dir ./trtllm_engines/deepseek_3/bf16/tp32-sel4096-isl2048-bs4 \ --tokenizer_dir ./DeepSeek-V3 \ - --max_output_len 40 \ - --input_text "The president of the United States is person who" + --max_output_len 30 \ + --input_text "Today is a nice day." ``` and the output will be like: ``` -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31490468978882 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31163835525513 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31164216995239 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31491041183472 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3116364479065 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3118085861206 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3118691444397 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31516337394714 sec -Input [Text 0]: "<｜begin▁of▁sentence｜>The president of the United States is person who" -Output [Text 0 Beam 0]: " is elected by the people of the United States to lead the country. The president is the head of the executive branch of the government. The president is also the commander in chief of the armed forces." +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.4361846446991 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81118130683899 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.436166286468506 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.436179876327515 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.811140298843384 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.43616843223572 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81416082382202 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.436097621917725 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81114745140076 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.43618369102478 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.810874462127686 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.436410903930664 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.814189195632935 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.43643116950989 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81437277793884 sec +[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81150937080383 sec [TensorRT-LLM][INFO] Refreshed the MPI local session [TensorRT-LLM][INFO] Refreshed the MPI local session [TensorRT-LLM][INFO] Refreshed the MPI local session @@ -147,6 +153,16 @@ Output [Text 0 Beam 0]: " is elected by the people of the United States to lead [TensorRT-LLM][INFO] Refreshed the MPI local session [TensorRT-LLM][INFO] Refreshed the MPI local session [TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +[TensorRT-LLM][INFO] Refreshed the MPI local session +Input [Text 0]: "Today is a nice day." +Output [Text 0 Beam 0]: " I am going to the park with my friends. We are going to play soccer. We are going" ``` If we want to evaluate the model summarization ability, we can use [summarize.py](../summarize.py) script: From 723eb27fbf1e75560db9093ac241ce781b974008 Mon Sep 17 00:00:00 2001 From: yingcanw Date: Thu, 26 Dec 2024 02:40:00 +0000 Subject: [PATCH 06/12] Update V3 README --- examples/deepseek_v3/README.md | 187 ++++----------------------------- 1 file changed, 21 insertions(+), 166 deletions(-) diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index f0067c94e0d5..5c745fd90bed 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -35,11 +35,11 @@ This document shows how to build and run the `DeepSeek-v3` model in TensorRT-LLM ## Prerequisite -First, please download DeepSeek-V3 weights from HF https://huggingface.co/deepseek-ai/DeepSeek-V3. +First, please download DeepSeek-V3 weights from HF https://huggingface.co/deepseek-ai/DeepSeek-V3-Base. ```bash git lfs install -git clone https://huggingface.co/deepseek-ai/DeepSeek-V3 +git clone https://huggingface.co/deepseek-ai/DeepSeek-V3-Base ``` ## Hardware @@ -57,8 +57,6 @@ The TensorRT-LLM DeepSeek-V3 implementation can be found in [tensorrt_llm/models In addition, there are three shared files in the parent folder [`examples`](../) can be used for inference and evaluation: * [`../run.py`](../run.py) to run the model inference output by giving an input text. -* [`../summarize.py`](../summarize.py) to summarize the article from [cnn_dailmail](https://huggingface.co/datasets/cnn_dailymail) dataset, it can run the summarize from the HF model and TensorRT-LLM model. -* [`../mmlu.py`](../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT-LLM model on the MMLU dataset. ## Usage @@ -112,39 +110,37 @@ trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ --max_batch_size 4 \ --max_seq_len 4096 \ --max_input_len 2048 \ - --use_paged_context_fmha enable + --use_paged_context_fmha enable \ + --workers 8 ``` ***Caution: `--max_batch_size` and `--max_seq_len` are the main factors to determine how many GPU memory will be used during runtime, so later when try to run e.g., `summarize.py` or `mmlu.py` or `gptManagerBenchmark.cpp`may need adjust `--max_batch_size` and `--max_seq_len` accordingly to avoid OOM.(meaning rebuild TensorRT engine with smaller `--max_batch_size` and `--max_seq_len` if needed based on GPU memory size), there is beautiful technical log perf-best-practices.md (https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/performance/perf-best-practices.md) explained the mechanism.*** Test the engine with [run.py](../run.py) script: +``` +# run.sh +python3 ../run.py --input_text "Today is a nice day." \ + --max_output_len 30 \ + --tokenizer_dir ./DeepSeek-V3 \ + --engine_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ + --top_p 0.95 \ + --temperature 0.3 + + +``` +For multi-nodes inference, let's take Slurm as an example using above command (run.sh): ```bash -mpirun --allow-run-as-root -np 32 python ../run.py --engine_dir ./trtllm_engines/deepseek_3/bf16/tp32-sel4096-isl2048-bs4 \ - --tokenizer_dir ./DeepSeek-V3 \ - --max_output_len 30 \ - --input_text "Today is a nice day." +srun -N 4 -w node-[1-4] --gres=gpu:8 --ntasks-per-node 8 \ + --container-image tensorrt_llm/release:latest \ + --container-mounts ${PWD}:/workspace \ + sh /workspace/command/run.sh ``` and the output will be like: ``` -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.4361846446991 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81118130683899 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.436166286468506 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.436179876327515 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.811140298843384 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.43616843223572 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81416082382202 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.436097621917725 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81114745140076 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.43618369102478 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.810874462127686 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.436410903930664 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.814189195632935 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 38.43643116950989 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81437277793884 sec -[12/25/2024-13:34:46] [TRT-LLM] [I] Load engine takes: 37.81150937080383 sec +... [TensorRT-LLM][INFO] Refreshed the MPI local session [TensorRT-LLM][INFO] Refreshed the MPI local session [TensorRT-LLM][INFO] Refreshed the MPI local session @@ -165,144 +161,3 @@ Input [Text 0]: "Today is a nice day." Output [Text 0 Beam 0]: " I am going to the park with my friends. We are going to play soccer. We are going" ``` -If we want to evaluate the model summarization ability, we can use [summarize.py](../summarize.py) script: - -```bash -mpirun --allow-run-as-root -np 32 python ../summarize.py --engine_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ - --hf_model_dir ./DeepSeek-V3 \ - --data_type bfloat16 \ - --batch_size 1 \ - --test_trt_llm \ - --test_hf -``` - -and the output will be like: - - -``` -[10/28/2024-16:46:22] [TRT-LLM] [I] HF Generated : -[10/28/2024-16:46:22] [TRT-LLM] [I] Input : ['(CNN)James Best, best known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," died Monday after a brief illness. He was 88. Best died in hospice in Hickory, North Carolina, of complications from pneumonia, said Steve Latshaw, a longtime friend and Hollywood colleague. Although he\'d been a busy actor for decades in theater and in Hollywood, Best didn\'t become famous until 1979, when "The Dukes of Hazzard\'s" cornpone charms began beaming into millions of American homes almost every Friday night. For seven seasons, Best\'s Rosco P. Coltrane chased the moonshine-running Duke boys back and forth across the back roads of fictitious Hazzard County, Georgia, although his "hot pursuit" usually ended with him crashing his patrol car. Although Rosco was slow-witted and corrupt, Best gave him a childlike enthusiasm that got laughs and made him endearing. His character became known for his distinctive "kew-kew-kew" chuckle and for goofy catchphrases such as "cuff \'em and stuff \'em!" upon making an arrest. Among the most popular shows on TV in the early \'80s, "The Dukes of Hazzard" ran until 1985 and spawned TV movies, an animated series and video games. Several of Best\'s "Hazzard" co-stars paid tribute to the late actor on social media. "I laughed and learned more from Jimmie in one hour than from anyone else in a whole year," co-star John Schneider, who played Bo Duke, said on Twitter. "Give Uncle Jesse my love when you see him dear friend." "Jimmy Best was the most constantly creative person I have ever known," said Ben Jones, who played mechanic Cooter on the show, in a Facebook post. "Every minute of his long life was spent acting, writing, producing, painting, teaching, fishing, or involved in another of his life\'s many passions." Born Jewel Guy on July 26, 1926, in Powderly, Kentucky, Best was orphaned at 3 and adopted by Armen and Essa Best, who renamed him James and raised him in rural Indiana. Best served in the Army during World War II before launching his acting career. In the 1950s and 1960s, he accumulated scores of credits, playing a range of colorful supporting characters in such TV shows as "The Twilight Zone," "Bonanza," "The Andy Griffith Show" and "Gunsmoke." He later appeared in a handful of Burt Reynolds\' movies, including "Hooper" and "The End." But Best will always be best known for his "Hazzard" role, which lives on in reruns. "Jimmie was my teacher, mentor, close friend and collaborator for 26 years," Latshaw said. "I directed two of his feature films, including the recent \'Return of the Killer Shrews,\' a sequel he co-wrote and was quite proud of as he had made the first one more than 50 years earlier." People we\'ve lost in 2015 . CNN\'s Stella Chan contributed to this story.'] -[10/28/2024-16:46:22] [TRT-LLM] [I] - Reference : ['James Best, who played the sheriff on "The Dukes of Hazzard," died Monday at 88 .\n"Hazzard" ran from 1979 to 1985 and was among the most popular shows on TV .'] -[10/28/2024-16:46:22] [TRT-LLM] [I] - Output : [[' James Best, best known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," died Monday after a brief illness. He was 88.']] -[10/28/2024-16:46:22] [TRT-LLM] [I] --------------------------------------------------------- -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT-LLM (total latency: 32.02327513694763 sec) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT-LLM (total output tokens: 1394) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT-LLM (tokens per second: 43.53083793080361) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT-LLM beam 0 result -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge1 : 17.85755990133811 -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge2 : 6.273032755727469 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeL : 14.768323033457317 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeLsum : 15.700915348496391 -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (total latency: 189.76398921012878 sec) -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (total output tokens: 1376) -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (tokens per second: 7.2511123197159) -[10/28/2024-16:49:33] [TRT-LLM] [I] HF beam 0 result -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge1 : 18.542590123197257 -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge2 : 6.345777100488389 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeL : 15.235695878419156 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeLsum : 16.64935135356226 -``` - -At last, we can evaluate the model with [mmlu.py](../mmlu.py) script: - -```bash -# Download MMLU dataset -mkdir mmlu_data && cd mmlu_data -wget https://people.eecs.berkeley.edu/~hendrycks/data.tar && tar -xf data.tar -# Run MMLU evaluation -mpirun --allow-run-as-root -n 32 python ../mmlu.py --engine_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ - --hf_model_dir ./DeepSeek-V3 \ - --data_type bfloat16 \ - --batch_size 1 \ - --test_trt_llm \ - --test_hf \ - --data_dir ./mmlu_data/data/ -``` - - - -and the output will be like: - -``` -Average accuracy 0.480 - abstract_algebra -Average accuracy 0.741 - anatomy -Average accuracy 0.888 - astronomy -Average accuracy 0.790 - business_ethics -Average accuracy 0.845 - clinical_knowledge -Average accuracy 0.924 - college_biology -Average accuracy 0.600 - college_chemistry -Average accuracy 0.720 - college_computer_science -Average accuracy 0.510 - college_mathematics -Average accuracy 0.751 - college_medicine -Average accuracy 0.618 - college_physics -Average accuracy 0.860 - computer_security -Average accuracy 0.796 - conceptual_physics -Average accuracy 0.675 - econometrics -Average accuracy 0.800 - electrical_engineering -Average accuracy 0.741 - elementary_mathematics -Average accuracy 0.643 - formal_logic -Average accuracy 0.570 - global_facts -Average accuracy 0.910 - high_school_biology -Average accuracy 0.714 - high_school_chemistry -Average accuracy 0.890 - high_school_computer_science -Average accuracy 0.891 - high_school_european_history -Average accuracy 0.909 - high_school_geography -Average accuracy 0.953 - high_school_government_and_politics -Average accuracy 0.826 - high_school_macroeconomics -Average accuracy 0.522 - high_school_mathematics -Average accuracy 0.916 - high_school_microeconomics -Average accuracy 0.576 - high_school_physics -Average accuracy 0.923 - high_school_psychology -Average accuracy 0.704 - high_school_statistics -Average accuracy 0.907 - high_school_us_history -Average accuracy 0.932 - high_school_world_history -Average accuracy 0.834 - human_aging -Average accuracy 0.893 - human_sexuality -Average accuracy 0.909 - international_law -Average accuracy 0.880 - jurisprudence -Average accuracy 0.853 - logical_fallacies -Average accuracy 0.598 - machine_learning -Average accuracy 0.874 - management -Average accuracy 0.953 - marketing -Average accuracy 0.880 - medical_genetics -Average accuracy 0.920 - miscellaneous -Average accuracy 0.850 - moral_disputes -Average accuracy 0.613 - moral_scenarios -Average accuracy 0.830 - nutrition -Average accuracy 0.859 - philosophy -Average accuracy 0.883 - prehistory -Average accuracy 0.635 - professional_accounting -Average accuracy 0.625 - professional_law -Average accuracy 0.857 - professional_medicine -Average accuracy 0.833 - professional_psychology -Average accuracy 0.745 - public_relations -Average accuracy 0.869 - security_studies -Average accuracy 0.935 - sociology -Average accuracy 0.930 - us_foreign_policy -Average accuracy 0.596 - virology -Average accuracy 0.877 - world_religions -Average accuracy 0.632 - math -Average accuracy 0.801 - health -Average accuracy 0.738 - physics -Average accuracy 0.897 - business -Average accuracy 0.914 - biology -Average accuracy 0.677 - chemistry -Average accuracy 0.762 - computer science -Average accuracy 0.832 - economics -Average accuracy 0.800 - engineering -Average accuracy 0.736 - philosophy -Average accuracy 0.821 - other -Average accuracy 0.902 - history -Average accuracy 0.909 - geography -Average accuracy 0.883 - politics -Average accuracy 0.876 - psychology -Average accuracy 0.919 - culture -Average accuracy 0.660 - law -Average accuracy 0.727 - STEM -Average accuracy 0.740 - humanities -Average accuracy 0.873 - social sciences -Average accuracy 0.821 - other (business, health, misc.) -Average accuracy: 0.785 -``` \ No newline at end of file From 48844c488834a1f8c62042cbbd7bfae097fa7b7c Mon Sep 17 00:00:00 2001 From: yingcanw Date: Thu, 26 Dec 2024 01:08:47 -0800 Subject: [PATCH 07/12] Update usages of FP8 to BF16 instruction --- examples/deepseek_v3/README.md | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index 5c745fd90bed..6d6857d306e8 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -67,10 +67,17 @@ The TensorRT-LLM DeepSeek-V3 example code is located at [examples/deepseek_v3](. Below is the step-by-step to run DeepSeek-V3 with TensorRT-LLM. -First the checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with TensorRT-LLM checkpoint. +Firstly, convert FP8 weights to BF16: +```bash +git clone https://github.com/deepseek-ai/DeepSeek-V3.git +cd DeepSeek-V3/inferece/ +python fp8_cast_bf16.py --input-fp8-hf-path /path/to/DeepSeek-V3 --output-bf16-hf-path /path/to/deepseek-v3-bf16 +``` + +Secondly, the BF16 checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. ```bash -# Convert Deepseek-v HF weights to TensorRT-LLM checkpoint in BF16. +# Convert Deepseek-v3 HF weights to TensorRT-LLM checkpoint in BF16. python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ --dtype bfloat16 \ From 9522bbcc920f205793f3514ca2a2ab5ea7a14ec7 Mon Sep 17 00:00:00 2001 From: jershi425 <83951930+jershi425@users.noreply.github.com> Date: Thu, 2 Jan 2025 15:44:34 +0800 Subject: [PATCH 08/12] fix model name mapping (#2) --- tensorrt_llm/models/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorrt_llm/models/__init__.py b/tensorrt_llm/models/__init__.py index 98f7994dc444..c3b5231e6827 100755 --- a/tensorrt_llm/models/__init__.py +++ b/tensorrt_llm/models/__init__.py @@ -182,6 +182,7 @@ 'DeepseekForCausalLM': DeepseekForCausalLM, 'DeciLMForCausalLM': DeciLMForCausalLM, 'DeepseekV2ForCausalLM': DeepseekV2ForCausalLM, + 'DeepseekV3ForCausalLM': DeepseekV2ForCausalLM, 'EagleForCausalLM': EagleForCausalLM, 'CohereForCausalLM': CohereForCausalLM, 'MLLaMAModel': MLLaMAForCausalLM, # For modelopt From 281fd3aa47c486ad2d870878bdc1a3ed18fc4876 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 6 Jan 2025 15:12:14 +0800 Subject: [PATCH 09/12] Update HF ckpt BF16 conversion. --- examples/deepseek_v3/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index 6d6857d306e8..7cb200b904cc 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -70,8 +70,9 @@ Below is the step-by-step to run DeepSeek-V3 with TensorRT-LLM. Firstly, convert FP8 weights to BF16: ```bash git clone https://github.com/deepseek-ai/DeepSeek-V3.git -cd DeepSeek-V3/inferece/ +cd DeepSeek-V3/inference/ python fp8_cast_bf16.py --input-fp8-hf-path /path/to/DeepSeek-V3 --output-bf16-hf-path /path/to/deepseek-v3-bf16 +cp /path/to/DeepSeek-V3/config.json /path/to/DeepSeek-V3/configuration_deepseek.py /path/to/deepseek-v3-bf16/ ``` Secondly, the BF16 checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. From 0321f71e9559906096211de557d0db52d840615a Mon Sep 17 00:00:00 2001 From: Yingcan Wei <68979322+yingcanw@users.noreply.github.com> Date: Thu, 16 Jan 2025 11:50:23 +0800 Subject: [PATCH 10/12] fix config of deepseek kv cache --- tensorrt_llm/models/deepseek_v2/config.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/models/deepseek_v2/config.py b/tensorrt_llm/models/deepseek_v2/config.py index 676418dacfc3..5085a8c413ab 100644 --- a/tensorrt_llm/models/deepseek_v2/config.py +++ b/tensorrt_llm/models/deepseek_v2/config.py @@ -143,7 +143,8 @@ def from_hugging_face(cls, num_attention_heads=hf_config.num_attention_heads, hidden_size=hf_config.hidden_size, intermediate_size=hf_config.intermediate_size, - num_key_value_heads=hf_config.num_key_value_heads, + num_key_value_heads=1, + head_size=hf_config.kv_lora_rank + hf_config.qk_rope_head_dim, vocab_size=hf_config.vocab_size, position_embedding_type='rope_gpt_neox', max_position_embeddings=hf_config.max_position_embeddings, From ed789c6a40ea6ff1fd03da13b0b28387b53dfa5c Mon Sep 17 00:00:00 2001 From: yingcanw Date: Fri, 24 Jan 2025 07:42:46 -0800 Subject: [PATCH 11/12] Remove source code --- .../mixtureOfExpertsBackendBenchmarkFixture.h | 4 +- cpp/tensorrt_llm/CMakeLists.txt | 8 + .../kernels/cutlass_kernels/CMakeLists.txt | 2 + .../fp8_blockscale_gemm/fp8_blockscale_gemm.h | 102 ++++ .../libfp8_blockscale_gemm.a | 3 + .../moe_gemm/moe_gemm_kernels_bf16_fp8.cu | 25 + .../moe_gemm/moe_gemm_kernels_template.h | 230 +++++---- .../kernels/mixtureOfExperts/moe_kernels.cu | 153 ++++-- .../kernels/mixtureOfExperts/moe_kernels.h | 63 ++- cpp/tensorrt_llm/kernels/mlaKernels.cu | 16 +- cpp/tensorrt_llm/kernels/mlaKernels.h | 14 +- cpp/tensorrt_llm/plugins/CMakeLists.txt | 1 + cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp | 3 + .../CMakeLists.txt | 21 + .../fp8CurrentScalingGemmPlugin.cpp | 382 +++++++++++++++ .../fp8CurrentScalingGemmPlugin.h | 116 +++++ .../gptAttentionCommon/gptAttentionCommon.cpp | 457 +++++++++++++++--- .../gptAttentionCommon/gptAttentionCommon.h | 10 +- .../gptAttentionPlugin/gptAttentionPlugin.cpp | 71 ++- .../gptAttentionPlugin/gptAttentionPlugin.h | 9 +- .../mixtureOfExpertsPlugin.cpp | 109 ++++- .../mixtureOfExperts/mixtureOfExpertsPlugin.h | 36 +- cpp/tensorrt_llm/thop/moeOp.cpp | 5 +- cpp/tests/kernels/mixtureOfExpertsTest.cu | 13 +- examples/deepseek_v3/README.md | 76 +-- examples/deepseek_v3/convert_checkpoint.py | 45 +- .../plugin_nodes/gpt_attention_node.py | 18 +- tensorrt_llm/builder.py | 6 + tensorrt_llm/functional.py | 25 +- tensorrt_llm/layers/attention.py | 196 +++++--- tensorrt_llm/layers/moe.py | 85 +++- tensorrt_llm/models/deepseek_v2/convert.py | 24 +- tensorrt_llm/models/deepseek_v2/model.py | 13 +- tensorrt_llm/plugin/plugin.py | 12 + tensorrt_llm/quantization/functional.py | 50 ++ tensorrt_llm/quantization/layers.py | 184 ++++++- tensorrt_llm/quantization/mode.py | 42 +- tensorrt_llm/quantization/quantize.py | 47 +- tests/attention/test_deepseek_v2_attention.py | 2 +- 39 files changed, 2224 insertions(+), 454 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu create mode 100644 cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/CMakeLists.txt create mode 100644 cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.cpp create mode 100644 cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h diff --git a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h index ec3da705f6d9..12819f48845a 100644 --- a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h +++ b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h @@ -425,6 +425,8 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture QuantParams mQuantParams{}; bool mUseLora = false; LoraParams mLoraParams{}; + bool mUseDeepSeek = false; + BlockScaleParams mDeepseekParams{}; std::optional mSelectedConfig = std::nullopt; @@ -678,7 +680,7 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture mMoERunner.runMoe(mInputTensor, mInputProbabilities, mExpertWeight1, mExpertBias1, mActType, mExpertWeight2, mExpertBias2, mQuantParams, mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK, mWorkspace, mFinalOutput, nullptr, mTotalTokens, mScaleProbs, mSourceToExpandedMap, mSelectedExpert, 0.01, - parallelism_config, mNormMode, mUseLora, mLoraParams, stream); + parallelism_config, mNormMode, mUseLora, mLoraParams, mUseDeepSeek, mDeepseekParams, stream); } void runBenchmark(benchmark::State& state); diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index ad979fe60e0b..c165eed6a9ad 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -369,6 +369,7 @@ set(TRTLLM_LINK_LIBS selective_scan_src fpA_intB_gemm_src moe_gemm_src + fp8_blockscale_gemm_lib fb_gemm_src gemm_swiglu_sm90_src cutlass_src @@ -377,6 +378,13 @@ set(TRTLLM_LINK_LIBS userbuffers_src ${DECODER_SHARED_TARGET}) +set(BLOCKSCALEGEMM_LIB_LOC + "${CMAKE_CURRENT_SOURCE_DIR}/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a" +) +add_library(fp8_blockscale_gemm_lib STATIC IMPORTED) +set_property(TARGET fp8_blockscale_gemm_lib PROPERTY IMPORTED_LOCATION + ${BLOCKSCALEGEMM_LIB_LOC}) + if(ENABLE_MULTI_DEVICE) set(TRTLLM_LINK_LIBS ${TRTLLM_LINK_LIBS} ${MPI_C_LIBRARIES} ${NCCL_LIB}) endif() diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt index a457b77aa5a9..a805bed007ea 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt @@ -78,6 +78,7 @@ file(GLOB_RECURSE SRC_CU *.cu) set(ALL_SRCS ${SRC_CPP};${SRC_CU}) list(FILTER ALL_SRCS EXCLUDE REGEX "fpA_intB_gemm/.*") list(FILTER ALL_SRCS EXCLUDE REGEX "moe_gemm/.*") +list(FILTER ALL_SRCS EXCLUDE REGEX "fp8_blockscale_gemm/.*") list(FILTER ALL_SRCS EXCLUDE REGEX "fp8_rowwise_gemm/.*") list(REMOVE_ITEM ALL_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/fused_gated_gemm/gemm_swiglu_e4m3.cu") @@ -90,6 +91,7 @@ message( "Group srcs ${GROUPED_SRC_CU} ${GROUPED_SRC_CPP} ${GROUPED_CU_INSTANTIATIONS}" ) message(VERBOSE "Fbgemm srcs ${FBGEMM_SRC_CU} ${FBGEMM_CU_INSTANTIATIONS}") +message(VERBOSE "Blockscale srcs ${BLOCKSCALEGEMM_SRC_CU} ") message(VERBOSE "All srcs ${ALL_SRCS}") add_library(cutlass_src STATIC ${ALL_SRCS}) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h new file mode 100644 index 000000000000..810c59d613e9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once +#include +#include +#include +#include + +// non-persistent-cooperative GEMM +namespace tensorrt_llm::kernels +{ +namespace small_m_gemm +{ + +class CutlassFp8BlockScaleGemmRunnerInterface +{ +public: + CutlassFp8BlockScaleGemmRunnerInterface() {} + + virtual ~CutlassFp8BlockScaleGemmRunnerInterface() {} + + virtual void gemm(void* mat_d, void const* mat_a, void const* mat_b, int shape_m, int shape_n, int shape_k, + char* workspace_ptr, cudaStream_t stream, float const* scales_a = nullptr, float const* scales_b = nullptr) + = 0; + + virtual void gemm(__nv_fp8_e4m3 const* mat_a, int ld_a, __nv_fp8_e4m3 const* mat_b, int ld_b, __nv_bfloat16* mat_d, int ld_d, + int shape_m, int shape_n, int shape_k, float const* scales_a, float const* scales_b, cudaStream_t stream) = 0; + + virtual void moeGemm(void *mat_d, void const *mat_a, void const *mat_b, const int64_t *problem_m_offsets, size_t num_problems, size_t shape_n, + size_t shape_k, char *workspace_ptr, cudaStream_t stream, float const *scales_a = nullptr, + float const *scales_b = nullptr) = 0; + + virtual void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a, + __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n, int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) = 0; + + virtual void fp8CS1x128(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y, cudaStream_t stream) = 0; + virtual void fp8CS1x128Reshape(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_h, int shape_y, int stride_x, cudaStream_t stream) = 0; + virtual void fp8CS128x128(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y, cudaStream_t stream) = 0; + // Returns desired workspace size in bytes. + virtual size_t getWorkspaceSize(size_t max_shape_m, size_t shape_n, size_t shape_k, size_t num_problems = 1) = 0; + + virtual size_t getFP8DataSize(int shape_m, int shape_n, bool is_act) = 0; + virtual size_t getActScaleSize(int shape_m, int shape_k) = 0; + virtual size_t getWeightScaleSize(int shape_n, int shape_k) = 0; + virtual size_t getActWorkspaceSize(int shape_m, int shape_k) = 0; + virtual size_t getWeightWorkspaceSize(int shape_n, int shape_k) = 0; +}; + +template +class CutlassFp8BlockScaleGemmRunner : public virtual CutlassFp8BlockScaleGemmRunnerInterface +{ +public: + CutlassFp8BlockScaleGemmRunner(); + ~CutlassFp8BlockScaleGemmRunner(); + + void gemm(void* mat_d, void const* mat_a, void const* mat_b, int shape_m, int shape_n, int shape_k, + char* workspace_ptr, cudaStream_t stream, float const* scales_a = nullptr, + float const* scales_b = nullptr) override; + + void gemm(__nv_fp8_e4m3 const* mat_a, int ld_a, __nv_fp8_e4m3 const* mat_b, int ld_b, __nv_bfloat16* mat_d, int ld_d, + int shape_m, int shape_n, int shape_k, float const* scales_a, float const* scales_b, cudaStream_t stream) override; + + void moeGemm(void *mat_d, void const *mat_a, void const *mat_b, const int64_t *problem_m_offsets, size_t num_problems, size_t shape_n, + size_t shape_k, char *workspace_ptr, cudaStream_t stream, float const *scales_a = nullptr, + float const *scales_b = nullptr) override; + + void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a, + __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n, int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) override; + + void fp8CS1x128(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y, cudaStream_t stream) override; + void fp8CS1x128Reshape(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_h, int shape_y, int stride_x, cudaStream_t stream) override; + void fp8CS128x128(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y, cudaStream_t stream) override; + + // Returns desired workspace size in bytes. + size_t getWorkspaceSize(size_t max_shape_m, size_t shape_n, size_t shape_k, size_t num_problems = 1) override; + + size_t getFP8DataSize(int shape_m, int shape_n, bool is_act) override; + size_t getActScaleSize(int shape_m, int shape_k) override; + size_t getWeightScaleSize(int shape_n, int shape_k) override; + size_t getActWorkspaceSize(int shape_m, int shape_k) override; + size_t getWeightWorkspaceSize(int shape_n, int shape_k) override; +private: + int64_t max_shape_m_4_align_ = 0; +}; + +} // namespace small_m_gemm +} // namespace tensorrt_llm::kernels diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a new file mode 100644 index 000000000000..cd6cfc9597a1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dc1668c231ab9333eee605a224197f6b7e5a36ffd516a19b33f4259a8a35eac4 +size 2123214 diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu new file mode 100644 index 000000000000..cc7b0d0fb244 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu @@ -0,0 +1,25 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h" + +namespace tensorrt_llm +{ + +#ifdef ENABLE_BF16 +template class MoeGemmRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>; +#endif +} // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h index 2a337e6ca4ec..3dbd4c50a5e5 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h @@ -45,6 +45,7 @@ #include "cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h" #include "cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h" #include "cutlass_extensions/gemm/threadblock/default_mma.h" +#include #ifdef __GNUC__ // Restore GCC-specific diagnostics #pragma GCC diagnostic pop @@ -74,126 +75,155 @@ namespace kernels::cutlass_kernels { // ============================= Variable batched Gemm things =========================== + +template +struct isGeneralMoe +{ + + static constexpr bool value + = (std::is_same_v && std::is_same_v) == false; +}; + template -void genericMoeGemmKernelLauncher(T const* A, WeightType const* B, GemmOutputType const* weight_scales, - GemmOutputType const* biases, bool bias_is_broadcast, GemmOutputType* C, - int64_t const* total_tokens_including_expert, int64_t num_rows, int64_t gemm_n, int64_t gemm_k, int num_experts, - cutlass_extensions::CutlassGemmConfig gemm_config, int const multi_processor_count, bool use_fused_moe, - float const** alpha_scale_ptr_array, cudaStream_t stream, int* kernel_occupancy = nullptr) +struct genericMoeGemmKernelLauncher { + static void call(T const* A, WeightType const* B, GemmOutputType const* weight_scales, GemmOutputType const* biases, + bool bias_is_broadcast, GemmOutputType* C, int64_t const* total_tokens_including_expert, int64_t num_rows, + int64_t gemm_n, int64_t gemm_k, int num_experts, cutlass_extensions::CutlassGemmConfig gemm_config, + int const multi_processor_count, bool use_fused_moe, float const** alpha_scale_ptr_array, cudaStream_t stream, + int* kernel_occupancy = nullptr) + { #if defined(ENABLE_FP8) - static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value - || cutlass::platform::is_same::value - || cutlass::platform::is_same::value || cutlass::platform::is_same::value, - "Specialized for fp8, bfloat16, half, float"); + static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value + || cutlass::platform::is_same::value + || cutlass::platform::is_same::value || cutlass::platform::is_same::value, + "Specialized for fp8, bfloat16, half, float"); #elif defined(ENABLE_BF16) - static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value - || cutlass::platform::is_same::value, - "Specialized for bfloat16, half, float"); + static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value + || cutlass::platform::is_same::value, + "Specialized for bfloat16, half, float"); #else - static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value, - "Specialized for half, float"); + static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value, + "Specialized for half, float"); #endif - static_assert(cutlass::platform::is_same::value - || cutlass::platform::is_same::value - || cutlass::platform::is_same::value, - ""); + static_assert(cutlass::platform::is_same::value + || cutlass::platform::is_same::value + || cutlass::platform::is_same::value, + ""); - static_assert(!cutlass::platform::is_same::value, - "Sm90 architecture should use specialised kernels"); + static_assert(!cutlass::platform::is_same::value, + "Sm90 architecture should use specialised kernels"); - // The cutlass type for the input elements. This is needed to convert to cutlass::half_t if necessary. - using ElementType = typename TllmToCutlassTypeAdapter::type; - using CutlassGemmOutputType = typename TllmToCutlassTypeAdapter::type; - using CutlassWeightType = typename TllmToCutlassTypeAdapter::type; - if (!use_fused_moe) - { - // We need separate config for each architecture since we will target different tensorcore instructions. For - // float, we do not target TCs. - using MixedGemmArchTraits = cutlass::gemm::kernel::MixedGemmArchTraits; - using ElementAccumulator = typename MixedGemmArchTraits::AccType; + // The cutlass type for the input elements. This is needed to convert to cutlass::half_t if necessary. + using ElementType = typename TllmToCutlassTypeAdapter::type; + using CutlassGemmOutputType = typename TllmToCutlassTypeAdapter::type; + using CutlassWeightType = typename TllmToCutlassTypeAdapter::type; + if (!use_fused_moe) + { + // We need separate config for each architecture since we will target different tensorcore instructions. For + // float, we do not target TCs. + using MixedGemmArchTraits + = cutlass::gemm::kernel::MixedGemmArchTraits; + using ElementAccumulator = typename MixedGemmArchTraits::AccType; - using EpilogueOp = typename tensorrt_llm::cutlass_extensions::Epilogue::Op; + using EpilogueOp = typename tensorrt_llm::cutlass_extensions::Epilogue::Op; - typename EpilogueOp::Params epilogue_op( - ElementAccumulator(1.f), biases ? ElementAccumulator(1.f) : ElementAccumulator(0.f)); + typename EpilogueOp::Params epilogue_op( + ElementAccumulator(1.f), biases ? ElementAccumulator(1.f) : ElementAccumulator(0.f)); #if defined(ENABLE_FP8) - if constexpr ((std::is_same_v - || std::is_same_v) &&std::is_same_v) - { - TLLM_CHECK_WITH_INFO(weight_scales == nullptr && biases == nullptr && alpha_scale_ptr_array, - "weight_scales and biases should be nullptr and alpha_scale_ptr_array shouldn't be nullptr for FP8 " - "Ada"); - epilogue_op.alpha_ptr_array = alpha_scale_ptr_array; - } + if constexpr ((std::is_same_v + || std::is_same_v) &&std::is_same_v) + { + TLLM_CHECK_WITH_INFO(weight_scales == nullptr && biases == nullptr && alpha_scale_ptr_array, + "weight_scales and biases should be nullptr and alpha_scale_ptr_array shouldn't be nullptr for FP8 " + "Ada"); + epilogue_op.alpha_ptr_array = alpha_scale_ptr_array; + } #endif - // Finally, set up the kernel. - using GemmKernel_ = typename cutlass::gemm::kernel::DefaultGemmGrouped::GemmKernel; - - using GemmKernel = cutlass::gemm::kernel::MoeFCGemm; - - using GemmGrouped = cutlass::gemm::device::GemmGrouped; - - if (kernel_occupancy != nullptr) + // Finally, set up the kernel. + using GemmKernel_ = typename cutlass::gemm::kernel::DefaultGemmGrouped::GemmKernel; + + using GemmKernel = cutlass::gemm::kernel::MoeFCGemm; + + using GemmGrouped = cutlass::gemm::device::GemmGrouped; + + if (kernel_occupancy != nullptr) + { + *kernel_occupancy = tensorrt_llm::cutlass_extensions::compute_occupancy_for_kernel(); + return; + } + int occupancy = std::min(2, GemmGrouped::maximum_active_blocks()); + TLLM_CHECK_WITH_INFO(occupancy > 0, "GPU lacks the shared memory resources to run GroupedGEMM kernel"); + int const threadblock_count = multi_processor_count * occupancy; + + int const group_size = gemm_k; + typename GemmGrouped::Arguments args(num_experts, threadblock_count, group_size, epilogue_op, + reinterpret_cast(A), reinterpret_cast(B), + reinterpret_cast(weight_scales), + reinterpret_cast(biases), bias_is_broadcast, + reinterpret_cast(C), total_tokens_including_expert, gemm_n, gemm_k); + + GemmGrouped gemm; + + auto can_implement = gemm.can_implement(args); + TLLM_CHECK_WITH_INFO(can_implement == cutlass::Status::kSuccess, + "MoE FC kernel will fail for params. Error: " + std::string(cutlassGetStatusString(can_implement))); + + auto init_status = gemm.initialize(args); + TLLM_CHECK_WITH_INFO(init_status == cutlass::Status::kSuccess, + "Failed to initialize cutlass grouped gemm. Error: " + + std::string(cutlassGetStatusString(init_status))); + + auto run_status = gemm.run(stream); + TLLM_CHECK_WITH_INFO(run_status == cutlass::Status::kSuccess, + "Failed to run cutlass grouped gemm. Error: " + std::string(cutlassGetStatusString(run_status))); + } + else if constexpr (sizeof(ElementType) == 2 && sizeof(CutlassWeightType) == 2 + && (std::is_same_v + || std::is_same_v) ) // use fused moe gemm + // kernel.. (only support + // fp16 or bf16) { - *kernel_occupancy = tensorrt_llm::cutlass_extensions::compute_occupancy_for_kernel(); - return; + sm80_generic_fused_moe_gemm_kernelLauncher( + reinterpret_cast(A), reinterpret_cast(B), + reinterpret_cast(biases), bias_is_broadcast, reinterpret_cast(C), + total_tokens_including_expert, num_rows, gemm_n, gemm_k, num_experts, multi_processor_count, stream, + kernel_occupancy); } - int occupancy = std::min(2, GemmGrouped::maximum_active_blocks()); - TLLM_CHECK_WITH_INFO(occupancy > 0, "GPU lacks the shared memory resources to run GroupedGEMM kernel"); - int const threadblock_count = multi_processor_count * occupancy; - - int const group_size = gemm_k; - typename GemmGrouped::Arguments args(num_experts, threadblock_count, group_size, epilogue_op, - reinterpret_cast(A), reinterpret_cast(B), - reinterpret_cast(weight_scales), - reinterpret_cast(biases), bias_is_broadcast, - reinterpret_cast(C), total_tokens_including_expert, gemm_n, gemm_k); - - GemmGrouped gemm; - - auto can_implement = gemm.can_implement(args); - TLLM_CHECK_WITH_INFO(can_implement == cutlass::Status::kSuccess, - "MoE FC kernel will fail for params. Error: " + std::string(cutlassGetStatusString(can_implement))); - - auto init_status = gemm.initialize(args); - TLLM_CHECK_WITH_INFO(init_status == cutlass::Status::kSuccess, - "Failed to initialize cutlass grouped gemm. Error: " + std::string(cutlassGetStatusString(init_status))); - - auto run_status = gemm.run(stream); - TLLM_CHECK_WITH_INFO(run_status == cutlass::Status::kSuccess, - "Failed to run cutlass grouped gemm. Error: " + std::string(cutlassGetStatusString(run_status))); } - else if constexpr (sizeof(ElementType) == 2 && sizeof(CutlassWeightType) == 2 - && (std::is_same_v - || std::is_same_v) ) // use fused moe gemm - // kernel.. (only support - // fp16 or bf16) +}; + +template +struct genericMoeGemmKernelLauncher<__nv_bfloat16, __nv_fp8_e4m3, GemmOutputType, arch, EpilogueTag, ThreadblockShape, + WarpShape, Stages> +{ + static void call(__nv_bfloat16 const* A, __nv_fp8_e4m3 const* B, GemmOutputType const* weight_scales, + GemmOutputType const* biases, bool bias_is_broadcast, GemmOutputType* C, + int64_t const* total_tokens_including_expert, int64_t num_rows, int64_t gemm_n, int64_t gemm_k, int num_experts, + cutlass_extensions::CutlassGemmConfig gemm_config, int const multi_processor_count, bool use_fused_moe, + float const** alpha_scale_ptr_array, cudaStream_t stream, int* kernel_occupancy = nullptr) { - sm80_generic_fused_moe_gemm_kernelLauncher(reinterpret_cast(A), - reinterpret_cast(B), reinterpret_cast(biases), - bias_is_broadcast, reinterpret_cast(C), total_tokens_including_expert, num_rows, gemm_n, - gemm_k, num_experts, multi_processor_count, stream, kernel_occupancy); } -} - +}; } // namespace kernels::cutlass_kernels template ) ) { kernels::cutlass_kernels::genericMoeGemmKernelLauncher(A, B, weight_scales, biases, bias_is_broadcast, C, + ThreadblockShape, WarpShape, Stages>::call(A, B, weight_scales, biases, bias_is_broadcast, C, total_tokens_including_expert, num_rows, gemm_n, gemm_k, num_experts, gemm_config, multi_processor_count, use_fused_moe, alpha_scale_ptr_array, stream, occupancy); } diff --git a/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu b/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu index 992ac406826e..1858f4aa6758 100644 --- a/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu @@ -1429,7 +1429,7 @@ template CutlassMoeFCRunner::getWorkspaceDeviceBufferSizes( int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const num_experts_per_node, int const k, ActivationType activation_type, - MOEExpertScaleNormalizationMode norm_mode, bool use_lora) const + MOEExpertScaleNormalizationMode norm_mode, bool use_lora, bool use_deepseek) const { size_t const num_moe_inputs = k * num_rows; size_t const permuted_elems = num_moe_inputs * hidden_size; @@ -1512,14 +1512,24 @@ std::vector CutlassMoeFCRunner workspace{source_rows_size, // - permuted_rows_size, // - permuted_experts_size, // - expert_first_token_offset_size, // - sparse_mixer_out_size, // - softmax_out_size, // - permuted_scales_size, // - sorter_size, // + if (use_deepseek) + { + int factor = is_gated_activation ? 2 : 1; + int blockscale_fc1_output_size = factor * interbuf_elems * gemm_output_dtype; + int blockscale_fc2_output_size = permuted_elems * gemm_output_dtype; + overlapped_gemm1_gemm2_inputs = std::max(permuted_data_size, fc1_result_size); + overlapped_gemm1_gemm2_outputs = std::max(blockscale_fc1_output_size, blockscale_fc2_output_size); + } + + std::vector workspace{ + source_rows_size, // + permuted_rows_size, // + permuted_experts_size, // + expert_first_token_offset_size, // + sparse_mixer_out_size, // + softmax_out_size, // + permuted_scales_size, // + sorter_size, // // These pointers reuse the same memory overlapped_gemm1_gemm2_inputs, // overlapped_gemm1_gemm2_outputs, // @@ -1529,7 +1539,8 @@ std::vector CutlassMoeFCRunner::getWorkspaceSize(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const k, ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, MOEParallelismConfig parallelism_config, - bool use_lora) const + bool use_lora, bool use_deepseek) const { int const ep_size = parallelism_config.ep_size; TLLM_CHECK_WITH_INFO(num_experts % ep_size == 0, "Number of experts must be a multiple of ep size"); - auto workspace = getWorkspaceDeviceBufferSizes( - num_rows, hidden_size, inter_size, num_experts, num_experts / ep_size, k, activation_type, norm_mode, use_lora); + auto workspace = getWorkspaceDeviceBufferSizes(num_rows, hidden_size, inter_size, num_experts, + num_experts / ep_size, k, activation_type, norm_mode, use_lora, use_deepseek); auto ws_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspace.data(), workspace.size()); TLLM_LOG_DEBUG("Mixture Of Experts Plugin requires workspace of %2f MiB", ws_size / 1024.f / 1024.f); return ws_size; @@ -1552,11 +1563,10 @@ template ::configureWsPtrs(char* ws_ptr, int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const num_experts_per_node, int const k, ActivationType activation_type, - MOEExpertScaleNormalizationMode norm_mode, bool use_lora) - + MOEExpertScaleNormalizationMode norm_mode, bool use_lora, bool use_deepseek) { - auto ws_sizes = getWorkspaceDeviceBufferSizes( - num_rows, hidden_size, inter_size, num_experts, num_experts_per_node, k, activation_type, norm_mode, use_lora); + auto ws_sizes = getWorkspaceDeviceBufferSizes(num_rows, hidden_size, inter_size, num_experts, num_experts_per_node, + k, activation_type, norm_mode, use_lora, use_deepseek); std::vector ws_sliced{(int8_t*) ws_ptr}; for (auto size : ws_sizes) @@ -1628,6 +1638,14 @@ void CutlassMoeFCRunner::confi lora_add_bias_ = (ScaleBiasType*) ws_sliced[15]; lora_fc2_result_ = (ScaleBiasType*) ws_sliced[16]; } + + if (use_deepseek) + { + permuted_data_ = (T*) ws_sliced[8]; + fc1_result_ = (T*) ws_sliced[8]; + glu_inter_result_ = (T*) ws_sliced[9]; + fc2_result_ = (T*) ws_sliced[9]; + } } void sortAndScanSoftmaxOutput(int* expert_for_source_row, int* source_rows, int* permuted_experts, int* permuted_rows, @@ -1682,7 +1700,6 @@ void CutlassMoeFCRunner::gemm1 num_experts_per_node, input, fc1_expert_weights, fc1_fp8_dequant, nullptr, static_cast(gemm_output), stream); sync_check_cuda_error(); - gemm_runner.moeGemm(input, nullptr, nullptr, nullptr, total_tokens_including_expert, hopper_input, expanded_num_rows, fc1_out_size, hidden_size, num_experts_per_node, false, alpha_scale_ptr_array, stream, config); @@ -1835,6 +1852,58 @@ void CutlassMoeFCRunner::gemm2 sync_check_cuda_error(); } +template +void CutlassMoeFCRunner::BlockScaleFC1(T const* const input, + T* const output, void* const gemm_output, int64_t const* const expert_first_token_offset, + WeightType const* const fc1_expert_weights, ScaleBiasType const* const fc1_expert_biases, + float const* const fc2_fp8_quant, int64_t const expanded_num_rows, int64_t const hidden_size, + int64_t const inter_size, int const num_experts_per_node, ActivationType fc1_activation_type, + BlockScaleParams& deepseek_params, cudaStream_t stream) +{ + auto gemm_runner = deepseek_params.blockscale_gemm_iml; + TLLM_CHECK_WITH_INFO(gemm_runner, "blockscale gemm runner must be instantiated"); + bool const is_gated_activation = isGatedActivation(fc1_activation_type); + + int shape_n = is_gated_activation ? inter_size * 2 : inter_size; + int shape_k = hidden_size; + + gemm_runner->moeGemm(gemm_output, input, fc1_expert_weights, expert_first_token_offset, num_experts_per_node, + shape_n, shape_k, deepseek_params.workspace, stream, nullptr, deepseek_params.fc1_scales_ptrs); + + sync_check_cuda_error(); + constexpr bool bias_is_broadcast = true; + doActivation(output, static_cast(gemm_output), + fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, + inter_size, expanded_num_rows, fc1_activation_type, stream); + + sync_check_cuda_error(); +} + +template +void CutlassMoeFCRunner::BlockScaleFC2(T const* const input, + void* const gemm_output, OutputType* const final_output, int64_t const* const expert_first_token_offset, + WeightType const* const fc2_expert_weights, ScaleBiasType const* const fc2_expert_biases, + float const* const token_topk_unpermuted_scales, int const* const expanded_source_row_to_expanded_dest_row, + int const* const expert_for_source_row, int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, + int64_t const expanded_num_rows, int64_t const hidden_size, int64_t const inter_size, + int const num_experts_per_node, int64_t const k, BlockScaleParams& deepseek_params, cudaStream_t stream, + MOEParallelismConfig parallelism_config) +{ + int shape_n = hidden_size; + int shape_k = inter_size; + auto gemm_runner = deepseek_params.blockscale_gemm_iml; + TLLM_CHECK_WITH_INFO(gemm_runner, "blockscale gemm runner must be instantiated"); + gemm_runner->moeGemm(gemm_output, input, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, + shape_n, shape_k, deepseek_params.workspace, stream, nullptr, deepseek_params.fc2_scales_ptrs); + + sync_check_cuda_error(); + + finalizeMoeRoutingKernelLauncher( + static_cast(gemm_output), final_output, fc2_expert_biases, + token_topk_unpermuted_scales, expanded_source_row_to_expanded_dest_row, expert_for_source_row, num_rows, + hidden_size, k, num_valid_tokens_ptr, parallelism_config, MOEExpertScaleNormalizationMode::NONE, stream); +} + template bool CutlassMoeFCRunner::setupLoraWorkspace(int64_t expanded_num_rows, int64_t num_rows, int64_t inter_size, int64_t hidden_size, int start_expert, bool is_gated_activation, @@ -2037,7 +2106,8 @@ void CutlassMoeFCRunner::runMo int const num_experts, int const k, char* workspace_ptr, void* final_output_void, bool const* finished, int64_t const active_rows, void* token_topk_final_scales_void, int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, float sparse_mixer_epsilon, MOEParallelismConfig parallelism_config, - MOEExpertScaleNormalizationMode normalization_mode, bool use_lora, LoraParams& lora_params, cudaStream_t stream) + MOEExpertScaleNormalizationMode normalization_mode, bool use_lora, LoraParams& lora_params, bool use_deepseek, + BlockScaleParams& deepseek_params, cudaStream_t stream) { static constexpr bool int_scales_required = std::is_same::value || std::is_same::value; @@ -2097,7 +2167,7 @@ void CutlassMoeFCRunner::runMo TLLM_CHECK_WITH_INFO(fc1_fp8_dequant == nullptr && fc2_fp8_quant == nullptr && fc2_fp8_dequant == nullptr, "FP8 scales are provided for integer quantization"); } - else if (fp8_scales_required) + else if (fp8_scales_required && (!use_deepseek)) { TLLM_CHECK_WITH_INFO(fc1_expert_biases == nullptr, "Bias is not supported with FP8"); TLLM_CHECK_WITH_INFO(fc2_expert_biases == nullptr, "Bias is not supported with FP8"); @@ -2133,7 +2203,7 @@ void CutlassMoeFCRunner::runMo int const num_experts_per_node = num_experts / parallelism_config.ep_size; configureWsPtrs(workspace_ptr, num_rows, hidden_size, inter_size, num_experts, num_experts_per_node, k, - fc1_activation_type, normalization_mode, use_lora); + fc1_activation_type, normalization_mode, use_lora, use_deepseek); int const start_expert = num_experts_per_node * parallelism_config.ep_rank; int const end_expert = start_expert + num_experts_per_node; @@ -2192,10 +2262,19 @@ void CutlassMoeFCRunner::runMo } } - Self::gemm1(moe_gemm_runner_, permuted_data_, fc1_result_, glu_inter_result_, expert_first_token_offset_, - hopper_grouped_gemm_input_, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, fc1_int_scales, - fc1_fp8_dequant, fc2_fp8_quant, expanded_num_rows, hidden_size, inter_size, num_experts_per_node, - fc1_activation_type, alpha_scale_ptr_array_, !use_lora, stream, *gemm1_config_); + if (use_deepseek) + { + BlockScaleFC1(permuted_data_, fc1_result_, glu_inter_result_, expert_first_token_offset_, fc1_expert_weights, + fc1_expert_biases, fc2_fp8_quant, expanded_num_rows, hidden_size, inter_size, num_experts_per_node, + fc1_activation_type, deepseek_params, stream); + } + else + { + Self::gemm1(moe_gemm_runner_, permuted_data_, fc1_result_, glu_inter_result_, expert_first_token_offset_, + hopper_grouped_gemm_input_, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, fc1_int_scales, + fc1_fp8_dequant, fc2_fp8_quant, expanded_num_rows, hidden_size, inter_size, num_experts_per_node, + fc1_activation_type, alpha_scale_ptr_array_, !use_lora, stream, *gemm1_config_); + } sync_check_cuda_error(); @@ -2206,13 +2285,22 @@ void CutlassMoeFCRunner::runMo sync_check_cuda_error(); } - Self::gemm2(moe_gemm_runner_, fc1_result_, fc2_result_, final_output, expert_first_token_offset_, - hopper_grouped_gemm_input_, fc2_expert_weights, fc2_expert_biases, fc2_int_scales, fc2_fp8_dequant, - token_topk_unpermuted_scales, permuted_scales_, expanded_source_row_to_expanded_dest_row, permuted_rows_, - expert_for_source_row, num_valid_tokens_ptr, num_rows, expanded_num_rows, hidden_size, inter_size, - num_experts_per_node, k, !use_deterministic_hopper_reduce_, alpha_scale_ptr_array_, use_lora, lora_fc2_result_, - stream, parallelism_config, *gemm2_config_); - + if (use_deepseek) + { + BlockScaleFC2(fc1_result_, fc2_result_, final_output, expert_first_token_offset_, fc2_expert_weights, + fc2_expert_biases, token_topk_unpermuted_scales, expanded_source_row_to_expanded_dest_row, + expert_for_source_row, num_valid_tokens_ptr, num_rows, expanded_num_rows, hidden_size, inter_size, + num_experts_per_node, k, deepseek_params, stream, parallelism_config); + } + else + { + Self::gemm2(moe_gemm_runner_, fc1_result_, fc2_result_, final_output, expert_first_token_offset_, + hopper_grouped_gemm_input_, fc2_expert_weights, fc2_expert_biases, fc2_int_scales, fc2_fp8_dequant, + token_topk_unpermuted_scales, permuted_scales_, expanded_source_row_to_expanded_dest_row, permuted_rows_, + expert_for_source_row, num_valid_tokens_ptr, num_rows, expanded_num_rows, hidden_size, inter_size, + num_experts_per_node, k, use_deterministic_hopper_reduce_, alpha_scale_ptr_array_, use_lora, + lora_fc2_result_, stream, parallelism_config, *gemm2_config_); + } sync_check_cuda_error(); } @@ -2631,6 +2719,7 @@ template class CutlassMoeFCRunner; template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, half>; #ifdef ENABLE_BF16 template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>; #endif #endif diff --git a/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h b/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h index fdf5fe9ba5e0..8a003fb9727a 100644 --- a/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h @@ -19,6 +19,7 @@ #include "cutlass/gemm/gemm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels.h" #include "tensorrt_llm/kernels/lora/lora.h" #include @@ -208,13 +209,38 @@ struct LoraParams } }; +struct BlockScaleParams +{ + using BlockScaleGemmImplPtr = std::shared_ptr; + BlockScaleGemmImplPtr blockscale_gemm_iml; + + float const* fc1_scales_ptrs = nullptr; + float const* fc2_scales_ptrs = nullptr; + + char* workspace; + + cudaEvent_t* memcpy_event_ptr; + + BlockScaleParams() = default; + + BlockScaleParams(float const* fc1_scales_ptrs, float const* fc2_scales_ptrs, + BlockScaleGemmImplPtr blockscale_gemm_iml, char* workspace, cudaEvent_t* memcpy_event_ptr) + : fc1_scales_ptrs(fc1_scales_ptrs) + , fc2_scales_ptrs(fc2_scales_ptrs) + , blockscale_gemm_iml(blockscale_gemm_iml) + , workspace(workspace) + , memcpy_event_ptr(memcpy_event_ptr) + { + } +}; + class CutlassMoeFCRunnerInterface { public: virtual ~CutlassMoeFCRunnerInterface() = default; virtual size_t getWorkspaceSize(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const k, ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, - MOEParallelismConfig parallelism_config, bool use_lora) const + MOEParallelismConfig parallelism_config, bool use_lora, bool use_deepseek) const = 0; virtual void setTactic(std::optional gemm1_config, std::optional gemm2_config) @@ -228,7 +254,7 @@ class CutlassMoeFCRunnerInterface bool const* finished, int64_t const active_rows, void* token_topk_unpermuted_scales, int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, float sparse_mixer_epsilon, MOEParallelismConfig parallelism_config, MOEExpertScaleNormalizationMode normalization_mode, bool use_lora, - LoraParams& lora_params, cudaStream_t stream) + LoraParams& lora_params, bool use_deepseek, BlockScaleParams& deepseek_params, cudaStream_t stream) = 0; // Aliases for profiling the gemms @@ -255,7 +281,6 @@ class CutlassMoeFCRunnerInterface = 0; virtual size_t getGemmWorkspaceSize(int num_experts) const = 0; - bool is_profiler = false; bool use_deterministic_hopper_reduce_ = false; }; @@ -295,7 +320,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface size_t getWorkspaceSize(int64_t const num_rows, int64_t const hidden_size, int64_t const fc1_output_size, int const num_experts, int const k, ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, - MOEParallelismConfig parallelism_config, bool use_lora) const override; + MOEParallelismConfig parallelism_config, bool use_lora, bool use_deepseek) const override; void setTactic(std::optional gemm1_config, std::optional gemm2_config) override @@ -322,7 +347,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface bool const* finished, int64_t const active_rows, void* token_topk_unpermuted_scales, int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, float sparse_mixer_epsilon, MOEParallelismConfig parallelism_config, MOEExpertScaleNormalizationMode normalization_mode, bool use_lora, - LoraParams& lora_params, cudaStream_t stream) override; + LoraParams& lora_params, bool use_deepseek, BlockScaleParams& deepseek_params, cudaStream_t stream) override; // We make these GEMM1 & GEMM2 static because they need to be stateless for the profiler to work static void gemm1(MoeGemmRunner& gemm_runner, T const* const input, @@ -398,10 +423,11 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface cudaStream_t stream); std::vector getWorkspaceDeviceBufferSizes(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const num_experts_per_node, int const k, - ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, bool use_lora) const; + ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, bool use_lora, + bool use_deepseek) const; void configureWsPtrs(char* ws_ptr, int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const num_experts_per_node, int const k, ActivationType activation_type, - MOEExpertScaleNormalizationMode norm_mode, bool use_lora); + MOEExpertScaleNormalizationMode norm_mode, bool use_lora, bool use_deepseek); private: bool mayHaveDifferentGEMMOutputType() const @@ -429,6 +455,21 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface int64_t const* num_valid_tokens_ptr, int64_t num_tokens, LoraParams& lora_params, float const* fc2_fp8_quant, cudaStream_t stream); + void BlockScaleFC1(T const* const input, T* const output, void* const intermediate_result, + int64_t const* const expert_first_token_offset, WeightType const* const fc1_expert_weights, + ScaleBiasType const* const fc1_expert_biases, float const* const fc2_fp8_quant, int64_t const expanded_num_rows, + int64_t const hidden_size, int64_t const inter_size, int const num_experts_per_node, + ActivationType fc1_activation_type, BlockScaleParams& deepseek_params, cudaStream_t stream); + + void BlockScaleFC2(T const* const input, void* const gemm_output, OutputType* const final_output, + int64_t const* const expert_first_token_offset, WeightType const* const fc2_expert_weights, + ScaleBiasType const* const fc2_expert_biases, float const* const token_topk_unpermuted_scales, + int const* const expanded_source_row_to_expanded_dest_row, int const* const expert_for_source_row, + int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, + int64_t const expanded_num_rows, int64_t const hidden_size, int64_t const inter_size, + int const num_experts_per_node, int64_t const k, BlockScaleParams& deepseek_params, cudaStream_t stream, + MOEParallelismConfig parallelism_config); + CubKeyValueSorter sorter_; MoeGemmRunner moe_gemm_runner_; @@ -471,6 +512,14 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface }; HostLoraWorkspace host_lora_workspace_; + + struct HostDeepSeekWorkspace + { + std::vector host_expert_first_token_offset; + std::vector host_num_tokens_per_expert; + }; + + HostDeepSeekWorkspace host_deepseek_workspace_; }; void makeLoadBalancedRoutingConfiguration( diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index 6394ded30c83..929bb95e6f0b 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -255,9 +255,9 @@ __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* qkv_output, T const* } template -__global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T const* fuse_buf, KVCacheBuffer kv_cache, - float2 const* cos_sin_cache, size_t head_num, int head_size, int c_q, int c_k, int total_s_len, int* seqQOffset, - uint32_t* fmha_tile_counter, int32_t const* kv_cache_lengths, int* seqKVOffsets) +__global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_buf, T const* fuse_buf, + KVCacheBuffer kv_cache, float2 const* cos_sin_cache, size_t head_num, int head_size, int c_q, int c_k, + int total_s_len, int* seqQOffset, uint32_t* fmha_tile_counter, int32_t const* kv_cache_lengths, int* seqKVOffsets) { // Constants. @@ -322,12 +322,12 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T const* } else { - auto const src_q_global_offset = static_cast(global_token_idx) * head_num * (c_k + ROPE_DIM) - + (c_k + ROPE_DIM) * head_idx + c_k; + auto const src_q_global_offset + = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + + (head_size + ROPE_DIM) * head_idx + head_size; for (int i = 0; i < 2; ++i) { - ref[i] = *reinterpret_cast( - &qkv_output[src_q_global_offset + src_bias + i * ELTS_PER_VEC]); + ref[i] = *reinterpret_cast(&q_buf[src_q_global_offset + src_bias + i * ELTS_PER_VEC]); } } @@ -451,7 +451,7 @@ void invokeMLARopeGeneration(mlaParams& params, KVCacheBuffer kv_cache_buffer dim3 grid(int(tensorrt_llm::common::divUp(params.acc_q_len, 32)), params.head_num + 1 + 8); auto head_size = params.meta.qk_nope_head_dim; applyMLARopeAndAssignQKVKernelGeneration - <<>>(params.attention_input_buf, params.fused_a_input, kv_cache_buffer, + <<>>(params.attention_input_buf, params.q_buf, params.fused_a_input, kv_cache_buffer, params.cos_sin_cache, params.head_num, head_size, params.meta.q_lora_rank, params.meta.kv_lora_rank, params.acc_q_len, params.seqQOffset, params.fmha_tile_counter, params.cache_seq_lens, params.cu_kv_seqlens); } diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.h b/cpp/tensorrt_llm/kernels/mlaKernels.h index 74b683886dd9..94099a214b85 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -39,12 +39,16 @@ struct mlaMetaParams template struct mlaParams { - T const* fused_a_input; // [b, s, c_q + c_k + r] - T* attention_input_buf; // [b, s, 3, h, d_h + r] + T const* fused_a_input; // [b, s, c_q + c_k + r] + T* attention_input_buf; // [b, s, 3, h, d_h + r] T* context_buf; - T const* fused_q_proj; // [c_k + r, d] - T const* q_b_proj; // [(d_h + r) * h, c_q] - T const* kv_b_proj; // [h * d_h * 2, c_k] + T* q_buf; // [b, h, d_h + r] + T const* q_b_proj; // [(d_h + r) * h, c_q] + T const* kv_b_proj; // [h * d_h * 2, c_k] + T const* k_b_proj_trans; // [h * c_k, d_h] + float const* q_b_scale; + float const* kv_b_scale; + float const* k_b_trans_scale; float2 const* cos_sin_cache; // [s, rope] int32_t batch_size; int32_t acc_q_len; diff --git a/cpp/tensorrt_llm/plugins/CMakeLists.txt b/cpp/tensorrt_llm/plugins/CMakeLists.txt index 40fdff9f5cd4..c38aa095c193 100755 --- a/cpp/tensorrt_llm/plugins/CMakeLists.txt +++ b/cpp/tensorrt_llm/plugins/CMakeLists.txt @@ -40,6 +40,7 @@ set(PLUGIN_LISTS identityPlugin gemmPlugin gemmSwigluPlugin + fp8CurrentScalingGemmPlugin fp8RowwiseGemmPlugin smoothQuantGemmPlugin quantizePerTokenPlugin diff --git a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp b/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp index 8c54ba94c054..f05580ee4a2f 100644 --- a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp @@ -20,6 +20,7 @@ #include "tensorrt_llm/runtime/tllmLogger.h" #include "tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h" +#include "tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h" #include "tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h" #include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" #include "tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h" @@ -207,6 +208,7 @@ extern "C" static tensorrt_llm::plugins::GemmPluginCreator gemmPluginCreator; static tensorrt_llm::plugins::GemmSwigluPluginCreator gemmSwigluPluginCreator; static tensorrt_llm::plugins::Fp8RowwiseGemmPluginCreator fp8RowwiseGemmPluginCreator; + static tensorrt_llm::plugins::Fp8CurrentScalingGemmPluginCreator fp8CurrentScalingGemmPluginCreator; static tensorrt_llm::plugins::MixtureOfExpertsPluginCreator moePluginCreator; #if ENABLE_MULTI_DEVICE static tensorrt_llm::plugins::SendPluginCreator sendPluginCreator; @@ -245,6 +247,7 @@ extern "C" creatorPtr(gemmPluginCreator), creatorPtr(gemmSwigluPluginCreator), creatorPtr(fp8RowwiseGemmPluginCreator), + creatorPtr(fp8CurrentScalingGemmPluginCreator), creatorPtr(moePluginCreator), #if ENABLE_MULTI_DEVICE creatorPtr(sendPluginCreator), diff --git a/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/CMakeLists.txt new file mode 100644 index 000000000000..3b714a3928fb --- /dev/null +++ b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/CMakeLists.txt @@ -0,0 +1,21 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & +# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. +# +file(GLOB SRCS *.cpp *.cu) +set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) +set(PLUGIN_SOURCES + ${PLUGIN_SOURCES} + PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.cpp new file mode 100644 index 000000000000..75b80bdc0553 --- /dev/null +++ b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.cpp @@ -0,0 +1,382 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "fp8CurrentScalingGemmPlugin.h" +#include "cutlass_extensions/gemm_configs.h" + +#include +#include + +using namespace nvinfer1; +using namespace tensorrt_llm::common; +using namespace tensorrt_llm::kernels::small_m_gemm; +using tensorrt_llm::plugins::Fp8CurrentScalingGemmPluginCreator; +using tensorrt_llm::plugins::Fp8CurrentScalingGemmPlugin; + +static char const* FP8_CURRENT_SCALING_GEMM_PLUGIN_VERSION{"1"}; +static char const* FP8_CURRENT_SCALING_GEMM_PLUGIN_NAME{"Fp8CurrentScalingGemm"}; +PluginFieldCollection Fp8CurrentScalingGemmPluginCreator::mFC{}; +std::vector Fp8CurrentScalingGemmPluginCreator::mPluginAttributes; + +Fp8CurrentScalingGemmPlugin::Fp8CurrentScalingGemmPlugin( + int need_quantize_acts_on_demand, int need_quantize_weights_on_demand, nvinfer1::DataType type) +{ + init(type, need_quantize_acts_on_demand, need_quantize_weights_on_demand); +} + +// Parameterized constructor +Fp8CurrentScalingGemmPlugin::Fp8CurrentScalingGemmPlugin(void const* data, size_t length) +{ + char const *d = reinterpret_cast(data), *a = d; + nvinfer1::DataType type; + int need_quantize_acts_on_demand = 0; + int need_quantize_weights_on_demand = 0; + + read(d, need_quantize_acts_on_demand); + read(d, need_quantize_weights_on_demand); + read(d, type); + + init(type, need_quantize_acts_on_demand, need_quantize_weights_on_demand); + + TLLM_CHECK_WITH_INFO(d == a + length, + "Expected length (%d) != real length (%d). This is often " + "caused by using different TensorRT-LLM version to build " + "engine and run engine.", + (int) length, (int) (d - a)); +} + +void Fp8CurrentScalingGemmPlugin::init( + nvinfer1::DataType type, int need_quantize_acts_on_demand, int need_quantize_weights_on_demand) +{ + mArch = tensorrt_llm::common::getSMVersion(); + mType = type; + mNeedQuantizeActsOnDemand = need_quantize_acts_on_demand; + mNeedQuantizeWeightsOnDemand = need_quantize_weights_on_demand; + mInputIdx = 0; + mWeightInputIdx = mInputIdx + 1; + mWeightScalesIdx = mNeedQuantizeWeightsOnDemand ? mWeightInputIdx : mWeightInputIdx + 1; + mInputScalesIdx = mNeedQuantizeActsOnDemand ? mWeightScalesIdx : mWeightScalesIdx + 1; + + if (mType == nvinfer1::DataType::kBF16) + { + if (mNeedQuantizeActsOnDemand && !mNeedQuantizeWeightsOnDemand) + { + mGemmRunner + = std::make_shared>(); + } + else if (mNeedQuantizeActsOnDemand && mNeedQuantizeWeightsOnDemand) + { + mGemmRunner + = std::make_shared>(); + } + else + { + mGemmRunner + = std::make_shared>(); + } + } + else + { + TLLM_THROW("Fp8 current scaling Gemm plugin doesn't support this type now"); + } +} + +// IPluginV2DynamicExt Methods +nvinfer1::IPluginV2DynamicExt* Fp8CurrentScalingGemmPlugin::clone() const noexcept +{ + auto* plugin = new Fp8CurrentScalingGemmPlugin(*this); + return plugin; +} + +nvinfer1::DimsExprs Fp8CurrentScalingGemmPlugin::getOutputDimensions( + int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept +{ + // inputs + // 0 activations [M, K] + // 1 weights [K, N] + // 2 weight scales [K // 128, N // 128] + // 3 activation scales [M, K // 128] (optional) + + try + { + TLLM_CHECK(nbInputs == mInputScalesIdx + 1); + TLLM_CHECK(outputIndex == 0); + int const nbDimsA = inputs[0].nbDims; + int const nbDimsB = inputs[mWeightInputIdx].nbDims; + TLLM_CHECK(nbDimsA >= 2); + TLLM_CHECK(nbDimsB == 2); + DimsExprs ret; + ret.nbDims = nbDimsA; + for (int ii = 0; ii < nbDimsA - 1; ++ii) + { + ret.d[ii] = inputs[0].d[ii]; + } + ret.d[nbDimsA - 1] = inputs[1].d[0]; + return ret; + } + catch (std::exception const& e) + { + caughtError(e); + } + return DimsExprs{}; +} + +bool Fp8CurrentScalingGemmPlugin::supportsFormatCombination( + int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept +{ + if (pos == mInputIdx) + { + // activations + if (mNeedQuantizeActsOnDemand) + { + return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; + } + else + { + return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; + } + } + else if (pos == mWeightInputIdx) + { + // weights + if (mNeedQuantizeWeightsOnDemand) + { + return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; + } + else + { + return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; + } + } + else if (!mNeedQuantizeWeightsOnDemand && pos == mWeightScalesIdx) + { + // weight scales + return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; + } + else if (!mNeedQuantizeActsOnDemand && pos == mInputScalesIdx) + { + // input scales + return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; + } + else if (pos == mInputScalesIdx + 1) + { + // outputs + return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; + } + else + { + return false; + } +} + +void Fp8CurrentScalingGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, + nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept +{ + auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); + auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); + + int const maxK = in[0].max.d[in[0].max.nbDims - 1]; + int const maxN = in[1].max.d[0]; + int const minK = in[0].min.d[in[0].min.nbDims - 1]; + int const minN = in[1].min.d[0]; + + TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); + TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); + + mWorkspaceMaxSize = mGemmRunner->getWorkspaceSize(maxM, maxN, maxK); +} + +size_t Fp8CurrentScalingGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, + nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept +{ + return mWorkspaceMaxSize; +} + +int Fp8CurrentScalingGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, + nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, + cudaStream_t stream) noexcept +{ + // inputs + // 0 activations [M, K] + // 1 weights [K, N] + // 2 weight scales[K // 128, N // 128] + // 3 scales [M, K // 128] (optional) + // + // outputs + // mat [M, N] + int m = 1; + for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) + { + m *= inputDesc[0].dims.d[ii]; + } + int const n = inputDesc[1].dims.d[0]; + int const k = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; + // size_t const wsSize = mGemmRunner->getWorkspaceSize(m, n, k); + + mGemmRunner->gemm(outputs[0], inputs[0], inputs[1], m, n, k, reinterpret_cast(workspace), stream, + reinterpret_cast(inputs[3]), reinterpret_cast(inputs[2])); + sync_check_cuda_error(); + + return 0; +} + +// IPluginV2Ext Methods +nvinfer1::DataType Fp8CurrentScalingGemmPlugin::getOutputDataType( + int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept +{ + TLLM_CHECK(index == 0); + return mType; +} + +// IPluginV2 Methods + +char const* Fp8CurrentScalingGemmPlugin::getPluginType() const noexcept +{ + return FP8_CURRENT_SCALING_GEMM_PLUGIN_NAME; +} + +char const* Fp8CurrentScalingGemmPlugin::getPluginVersion() const noexcept +{ + return FP8_CURRENT_SCALING_GEMM_PLUGIN_VERSION; +} + +int Fp8CurrentScalingGemmPlugin::getNbOutputs() const noexcept +{ + return 1; +} + +int Fp8CurrentScalingGemmPlugin::initialize() noexcept +{ + // Modify here, maybe do nothing + configGemm(); // gemm profiler in action + return 0; +} + +void Fp8CurrentScalingGemmPlugin::terminate() noexcept {} + +size_t Fp8CurrentScalingGemmPlugin::getSerializationSize() const noexcept +{ + return sizeof(int) + // need_quantize_acts_on_demand + sizeof(int) + // need_quantize_weights_on_demand + sizeof(nvinfer1::DataType); // dtype +} + +void Fp8CurrentScalingGemmPlugin::serialize(void* buffer) const noexcept +{ + char *d = static_cast(buffer), *a = d; + write(d, mNeedQuantizeActsOnDemand); + write(d, mNeedQuantizeWeightsOnDemand); + write(d, mType); + + TLLM_CHECK(d == a + getSerializationSize()); +} + +void Fp8CurrentScalingGemmPlugin::destroy() noexcept +{ + // This gets called when the network containing plugin is destroyed + delete this; +} + +void Fp8CurrentScalingGemmPlugin::configGemm() {} + +Fp8CurrentScalingGemmPluginCreator::Fp8CurrentScalingGemmPluginCreator() +{ + // Fill PluginFieldCollection with PluginField arguments metadata + mPluginAttributes.clear(); + mPluginAttributes.emplace_back(PluginField("need_quantize_acts_on_demand", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("need_quantize_weights_on_demand", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1)); + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); +} + +char const* Fp8CurrentScalingGemmPluginCreator::getPluginName() const noexcept +{ + return FP8_CURRENT_SCALING_GEMM_PLUGIN_NAME; +} + +char const* Fp8CurrentScalingGemmPluginCreator::getPluginVersion() const noexcept +{ + return FP8_CURRENT_SCALING_GEMM_PLUGIN_VERSION; +} + +PluginFieldCollection const* Fp8CurrentScalingGemmPluginCreator::getFieldNames() noexcept +{ + return &mFC; +} + +IPluginV2* Fp8CurrentScalingGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept +{ + PluginField const* fields = fc->fields; + TLLM_CHECK(fc->nbFields == 3); + int needQuantizeActsOnDemand; + int needQuantizeWeightsOnDemand; + + nvinfer1::DataType type; + // Read configurations from each fields + for (int i = 0; i < fc->nbFields; ++i) + { + char const* attrName = fields[i].name; + if (!strcmp(attrName, "need_quantize_acts_on_demand")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + needQuantizeActsOnDemand = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "need_quantize_weights_on_demand")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + needQuantizeWeightsOnDemand = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "type_id")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + type = static_cast(*(static_cast(fields[i].data))); + } + } + try + { + // Fp8CurrentScalingGemmPluginCreator is unique and shared for an engine generation + // Create plugin profiler with shared tactics map + auto* obj = new Fp8CurrentScalingGemmPlugin(needQuantizeActsOnDemand, needQuantizeWeightsOnDemand, type); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; + } + catch (std::exception const& e) + { + caughtError(e); + } + return nullptr; +} + +IPluginV2* Fp8CurrentScalingGemmPluginCreator::deserializePlugin( + char const* name, void const* serialData, size_t serialLength) noexcept +{ + // This object will be deleted when the network is destroyed, which will + // call Fp8CurrentScalingGemmPlugin::destroy() + try + { + // Create plugin profiler with private tactics map which is read from the serialized engine + auto* obj = new Fp8CurrentScalingGemmPlugin(serialData, serialLength); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; + } + catch (std::exception const& e) + { + caughtError(e); + } + return nullptr; +} diff --git a/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h new file mode 100644 index 000000000000..f38de4795f39 --- /dev/null +++ b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h @@ -0,0 +1,116 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +// Add our real GEMM kernel +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" +#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" +#include "tensorrt_llm/plugins/common/plugin.h" +#include +#include +#include +#include + +namespace tensorrt_llm::plugins +{ + +using Fp8BlockScaleGemmRunnerPtr + = std::shared_ptr; + +class Fp8CurrentScalingGemmPlugin : public BasePlugin +{ +public: + Fp8CurrentScalingGemmPlugin() = delete; + + Fp8CurrentScalingGemmPlugin( + int need_quantize_acts_on_demand, int need_quantize_weights_on_demand, nvinfer1::DataType type); + + Fp8CurrentScalingGemmPlugin(void const* data, size_t length); + + ~Fp8CurrentScalingGemmPlugin() override = default; + + // IPluginV2DynamicExt Methods + nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; + nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, + nvinfer1::IExprBuilder& exprBuilder) noexcept override; + bool supportsFormatCombination( + int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; + void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, + nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; + size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, + nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; + int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, + void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; + + // IPluginV2Ext Methods + nvinfer1::DataType getOutputDataType( + int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; + + // IPluginV2 Methods + char const* getPluginType() const noexcept override; + char const* getPluginVersion() const noexcept override; + int getNbOutputs() const noexcept override; + int initialize() noexcept override; + void terminate() noexcept override; + size_t getSerializationSize() const noexcept override; + void serialize(void* buffer) const noexcept override; + void destroy() noexcept override; + +private: + void init(nvinfer1::DataType type, int need_quantize_acts_on_demand, int need_quantize_weights_on_demand); + + void configGemm(); + +private: + const std::string mLayerName; + + Fp8BlockScaleGemmRunnerPtr mGemmRunner; + size_t mWorkspaceMaxSize; + nvinfer1::DataType mType; + + int mArch; + int mNeedQuantizeActsOnDemand; + int mNeedQuantizeWeightsOnDemand; + + int mInputIdx; + int mWeightInputIdx; + int mWeightScalesIdx; + int mInputScalesIdx; +}; + +class Fp8CurrentScalingGemmPluginCreator : public BaseCreator +{ +public: + Fp8CurrentScalingGemmPluginCreator(); + + char const* getPluginName() const noexcept override; + + char const* getPluginVersion() const noexcept override; + + nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; + + nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; + + nvinfer1::IPluginV2* deserializePlugin( + char const* name, void const* serialData, size_t serialLength) noexcept override; + +private: + static nvinfer1::PluginFieldCollection mFC; + static std::vector mPluginAttributes; +}; + +} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp index 36298d53539b..0e947bd2902a 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp +++ b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp @@ -426,8 +426,8 @@ GPTAttentionPluginCommon::GPTAttentionPluginCommon(int layer_idx, int num_heads, bool use_paged_context_fmha, bool use_fp8_context_fmha, bool has_full_attention_mask, bool use_cache, bool is_spec_decoding_enabled, bool spec_decoding_is_generation_length_variable, int32_t spec_decoding_max_generation_length, bool is_mla_enabled, int q_lora_rank, int kv_lora_rank, - int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool skip_attn, int cp_size, int cp_rank, - std::set cp_group) + int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool is_ptp128c_enabled, bool is_fp8_model, + bool skip_attn, int cp_size, int cp_rank, std::set cp_group) : mLayerIdx(layer_idx) , mNumHeads(num_heads) , mVisionStart(vision_start) @@ -478,6 +478,8 @@ GPTAttentionPluginCommon::GPTAttentionPluginCommon(int layer_idx, int num_heads, , mSpecDecodingMaxGenerationLength(spec_decoding_max_generation_length) , mIsMLAEnabled(is_mla_enabled) , mMLAParams({q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim}) + , mIsPTP128CEnabled(is_ptp128c_enabled) + , mIsFP8Model(is_fp8_model) , mCpSize(cp_size) , mCpRank(cp_rank) , mCpGroup(move(cp_group)) @@ -616,6 +618,8 @@ GPTAttentionPluginCommon::GPTAttentionPluginCommon(void const* data, size_t leng read(d, mSpecDecodingMaxGenerationLength); read(d, mIsMLAEnabled); read(d, mMLAParams); + read(d, mIsPTP128CEnabled); + read(d, mIsFP8Model); read(d, mNbMultiBlockSemaphores); read(d, mSkipAttn); read(d, mCpSize); @@ -725,16 +729,38 @@ size_t GPTAttentionPluginCommon::getWorkspaceSizeForGeneration( // MLA use fmha instead of XQA in generation if (mIsMLAEnabled) { - size_t cu_seqlens_size = sizeof(int) * (max_num_tokens + 1); + size_t cu_seqlens_size = sizeof(int) * (max_num_seq + 1); size_t fmha_scheduler_counter = sizeof(uint32_t); - size_t o_buffer_size = size * max_num_tokens * mNumHeads * mMLAParams.kv_lora_rank; - int const NUM_BUFFERS = 5; + size_t q_buffer_size = size * max_num_seq * mNumHeads * (mMLAParams.q_lora_rank + mMLAParams.qk_rope_head_dim); + size_t o_buffer_size = size * max_num_seq * mNumHeads * (mMLAParams.kv_lora_rank); + size_t act_buffer_size = 0; + size_t weight_buffer_size = 0; + if (mIsPTP128CEnabled) + { + act_buffer_size + = std::max(mGemmRunner->getActWorkspaceSize(max_num_seq, + mMLAParams.q_lora_rank + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim), + mGemmRunner->getActWorkspaceSize( + max_num_seq, mNumHeads * (mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim))); + if (!mIsFP8Model) + { + weight_buffer_size = std::max(mGemmRunner->getWeightWorkspaceSize(mNumHeads + * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim), + mMLAParams.q_lora_rank), + mGemmRunner->getWeightWorkspaceSize( + mNumHeads * mMLAParams.qk_nope_head_dim * 2, mMLAParams.kv_lora_rank)); + } + } + int const NUM_BUFFERS = 8; size_t workspaces[NUM_BUFFERS]; workspaces[0] = CUBLAS_WORKSPACE_SIZE; workspaces[1] = cu_seqlens_size; // cu_q_len workspaces[2] = cu_seqlens_size; // cu_kv_len workspaces[3] = fmha_scheduler_counter; - workspaces[4] = o_buffer_size; + workspaces[4] = q_buffer_size; + workspaces[5] = o_buffer_size; + workspaces[6] = act_buffer_size; + workspaces[7] = weight_buffer_size; generation_workspace_size = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); return generation_workspace_size; } @@ -818,45 +844,164 @@ int GPTAttentionPluginCommon::mlaPreContext( auto c_k = mMLAParams.kv_lora_rank; auto context_head_size = mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim; auto v_head_dim = mMLAParams.v_head_dim; - - // kv = self.kv_b_proj(compressed_kv) [b*s, c_k] * [c_k, h * (d_h * 2)] -> [b*s, h * (d_h * 2)] + if (!mIsPTP128CEnabled) { - auto transa = CUBLAS_OP_T; - auto transb = CUBLAS_OP_N; - int m = params.head_num * context_head_size; - int n = params.acc_q_len; - int k = c_q; - int lda = k, ldb = c_q + c_k + rope_dim; - int ldc - = params.head_num * (2 * context_head_size + v_head_dim); // output shape: [(b * s), (3 * h * (d_h + rope)] - mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - mCublasWrapper->Gemm( - transa, transb, m, n, k, params.q_b_proj, lda, params.fused_a_input, ldb, params.attention_input_buf, ldc); - mCublasWrapper->destroyDescriptors(); - } + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = params.head_num * context_head_size; + int n = params.acc_q_len; + int k = c_q; + int lda = k, ldb = c_q + c_k + rope_dim; + int ldc = params.head_num + * (2 * context_head_size + v_head_dim); // output shape: [(b * s), (3 * h * (d_h + rope)] + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + mCublasWrapper->Gemm(transa, transb, m, n, k, params.q_b_proj, lda, params.fused_a_input, ldb, + params.attention_input_buf, ldc); + mCublasWrapper->destroyDescriptors(); + } + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = mMLAParams.qk_nope_head_dim; + int n = params.acc_q_len; + int k = c_k; + // int lda = k, ldb = k + params.rope_dim + params.c_q; + // int ldc = (params.head_size + params.rope_dim) * params.head_num; //params.head_size * params.c_k; + int lda = k, ldb = c_q + c_k + rope_dim; + int ldc = params.head_num * (2 * context_head_size + v_head_dim); + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, params.kv_b_proj, lda, + mMLAParams.qk_nope_head_dim * c_k, params.fused_a_input + c_q, ldb, 0, + params.attention_input_buf + static_cast(params.head_num) * context_head_size, + params.head_num * (context_head_size * 2 + v_head_dim), context_head_size, params.head_num, 1.0f, 0.0); + + mCublasWrapper->Gemm(transa, transb, params.head_num * m, n, k, + params.kv_b_proj + mMLAParams.qk_nope_head_dim * c_k * params.head_num, lda, params.fused_a_input + c_q, + ldb, params.attention_input_buf + 2 * static_cast(params.head_num) * context_head_size, ldc); + mCublasWrapper->destroyDescriptors(); + } + } + else { - auto transa = CUBLAS_OP_T; - auto transb = CUBLAS_OP_N; - int m = mMLAParams.qk_nope_head_dim; - int n = params.acc_q_len; - int k = c_k; - // int lda = k, ldb = k + params.rope_dim + params.c_q; - // int ldc = (params.head_size + params.rope_dim) * params.head_num; //params.head_size * params.c_k; - int lda = k, ldb = c_q + c_k + rope_dim; - int ldc = params.head_num * (context_head_size * 2 + v_head_dim); - mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - - mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, params.kv_b_proj, lda, - mMLAParams.qk_nope_head_dim * c_k, params.fused_a_input + c_q, ldb, 0, - params.attention_input_buf + static_cast(params.head_num) * context_head_size, - params.head_num * (context_head_size * 2 + v_head_dim), context_head_size, params.head_num, 1.0f, 0.0); - mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, - params.kv_b_proj + mMLAParams.qk_nope_head_dim * c_k * params.head_num, lda, v_head_dim * c_k, - params.fused_a_input + c_q, ldb, 0, - params.attention_input_buf + 2 * static_cast(params.head_num) * context_head_size, - params.head_num * (context_head_size * 2 + v_head_dim), v_head_dim, params.head_num, 1.0f, 0.0); - mCublasWrapper->destroyDescriptors(); + int8_t* workspace_byte_ptr = reinterpret_cast(params.workspace); + size_t offset = CUBLAS_WORKSPACE_SIZE; + + size_t act_size = mGemmRunner->getFP8DataSize(params.acc_q_len, c_q + c_k + rope_dim, true); + size_t act_scale_size = mGemmRunner->getActScaleSize(params.acc_q_len, c_q + c_k + rope_dim); + __nv_fp8_e4m3* act_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, act_size)); + float* act_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, act_scale_size)); + size_t act_offset = offset; + + mGemmRunner->fp8CS1x128(act_buffer, act_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.fused_a_input), c_q + c_k + rope_dim, params.acc_q_len, + stream); + sync_check_cuda_error(); + + // mat_a: [params.acc_q_len, c_q], mat_b: [c_q, params.head_num * context_head_size] + // + { + int m = params.head_num * context_head_size; + int n = params.acc_q_len; + int k = c_q; + int lda = k, ldb = c_q + c_k + rope_dim; + int ldc = params.head_num + * (2 * context_head_size + v_head_dim); // output shape: [(b * s), (3 * h * (d_h + rope)] + + __nv_fp8_e4m3 const* q_gemm_weight_ptr; + float const* q_gemm_scale_ptr; + if (mIsFP8Model) + { + q_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.q_b_proj); + q_gemm_scale_ptr = reinterpret_cast(params.q_b_scale); + } + else + { + size_t q_gemm_size = mGemmRunner->getFP8DataSize(params.head_num * context_head_size, c_q, false); + size_t q_gemm_scale_size = mGemmRunner->getWeightScaleSize(params.head_num * context_head_size, c_q); + __nv_fp8_e4m3* q_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, q_gemm_size)); + float* q_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, q_gemm_scale_size)); + mGemmRunner->fp8CS128x128(q_weight_buffer, q_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.q_b_proj), c_q, params.head_num * context_head_size, + stream); + + q_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(q_weight_buffer); + q_gemm_scale_ptr = reinterpret_cast(q_scale_buffer); + } + + // npcgemm2::fp8_gemm_run(act_buffer, ldb, const_cast<__nv_fp8_e4m3*>(q_gemm_weight_ptr), lda, + // reinterpret_cast<__nv_bfloat16*>(params.attention_input_buf), ldc, n, m, k, act_scale_buffer, + // const_cast(q_gemm_scale_ptr), stream); + mGemmRunner->gemm(act_buffer, ldb, const_cast<__nv_fp8_e4m3*>(q_gemm_weight_ptr), lda, + reinterpret_cast<__nv_bfloat16*>(params.attention_input_buf), ldc, n, m, k, act_scale_buffer, + const_cast(q_gemm_scale_ptr), stream); + sync_check_cuda_error(); + } + + { + offset = act_offset; + __nv_fp8_e4m3 const* kv_gemm_weight_ptr; + float const* kv_gemm_scale_ptr; + if (mIsFP8Model) + { + kv_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.kv_b_proj); + kv_gemm_scale_ptr = reinterpret_cast(params.kv_b_scale); + } + else + { + size_t kv_gemm_size + = mGemmRunner->getFP8DataSize(params.head_num * mMLAParams.qk_nope_head_dim * 2, c_k, false); + size_t kv_gemm_scale_size + = mGemmRunner->getWeightScaleSize(params.head_num * mMLAParams.qk_nope_head_dim * 2, c_k); + __nv_fp8_e4m3* kv_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, kv_gemm_size)); + float* kv_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, kv_gemm_scale_size)); + mGemmRunner->fp8CS128x128(kv_weight_buffer, kv_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.kv_b_proj), c_k, + params.head_num * mMLAParams.qk_nope_head_dim * 2, stream); + + kv_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(kv_weight_buffer); + kv_gemm_scale_ptr = reinterpret_cast(kv_scale_buffer); + } + + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = mMLAParams.qk_nope_head_dim; + int n = params.acc_q_len; + int k = c_k; + + int shape_n_4_align = ((params.acc_q_len - 1) / 4 + 1) * 4; + int lda = k, ldb = c_q + c_k + rope_dim; + int ldc = params.head_num * (2 * context_head_size + v_head_dim); + + mGemmRunner->strideBatchGemm(reinterpret_cast<__nv_bfloat16*>(params.attention_input_buf) + + static_cast(params.head_num) * context_head_size, + ldc, context_head_size, act_buffer + c_q, ldb, 0, const_cast<__nv_fp8_e4m3*>(kv_gemm_weight_ptr), lda, + m * k, params.head_num, shape_n_4_align, m, k, stream, act_scale_buffer + shape_n_4_align * (c_q / 128), + 0, const_cast(kv_gemm_scale_ptr)); + sync_check_cuda_error(); + // npcgemm2::fp8_gemm_run(act_buffer + c_q, ldb, + // const_cast<__nv_fp8_e4m3*>(kv_gemm_weight_ptr) + params.head_num * m * k, lda, + // reinterpret_cast<__nv_bfloat16*>( + // params.attention_input_buf + 2 * static_cast(params.head_num) * context_head_size), + // ldc, n, params.head_num * m, k, act_scale_buffer + shape_n_4_align * (c_q / 128), + // const_cast(kv_gemm_scale_ptr) + params.head_num * (k / 128) * (m / 128), stream); + + mGemmRunner->gemm(act_buffer + c_q, ldb, + const_cast<__nv_fp8_e4m3*>(kv_gemm_weight_ptr) + params.head_num * m * k, lda, + reinterpret_cast<__nv_bfloat16*>( + params.attention_input_buf + 2 * static_cast(params.head_num) * context_head_size), + ldc, n, params.head_num * m, k, act_scale_buffer + shape_n_4_align * (c_q / 128), + const_cast(kv_gemm_scale_ptr) + params.head_num * (k / 128) * (m / 128), stream); + sync_check_cuda_error(); + mCublasWrapper->destroyDescriptors(); + } } return 0; @@ -917,34 +1062,143 @@ int GPTAttentionPluginCommon::mlaGeneration( int8_t* workspace_byte_ptr = reinterpret_cast(params.workspace); size_t offset = CUBLAS_WORKSPACE_SIZE; - // output[b, s, :1, h * (d_h + rope)] = self.q_b_proj(q_buf) [b*s, c_q] * [c_q, h * (d_h + rope)] -> [b*s, h * (d_h - // + rope)] - { - auto transa = CUBLAS_OP_T; - auto transb = CUBLAS_OP_N; - int m = params.head_num * (mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim); - int n = params.acc_q_len; - int k = mMLAParams.q_lora_rank; - int lda = k, ldb = mMLAParams.q_lora_rank + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim; - int ldc = m; - mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - mCublasWrapper->Gemm(transa, transb, m, n, k, params.fused_q_proj, lda, params.fused_a_input, ldb, - params.attention_input_buf, ldc); - mCublasWrapper->destroyDescriptors(); - } - size_t const cu_seqlens_size = sizeof(int) * (params.batch_size + 1); size_t const fmha_scheduler_counter = sizeof(uint32_t); + size_t q_buffer_size = size * batch_beam * mNumHeads * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim); size_t o_buffer_size = size * batch_beam * mNumHeads * mMLAParams.kv_lora_rank; int* cu_q_seqlens = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, cu_seqlens_size)); int* cu_kv_seqlens = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, cu_seqlens_size)); uint32_t* fmha_tile_counter_ptr = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, fmha_scheduler_counter)); + T* q_buffer = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, q_buffer_size)); T* o_buffer = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, o_buffer_size)); + size_t quant_offset = offset; + + if (!mIsPTP128CEnabled) + { + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = params.head_num * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim); + int n = params.acc_q_len; + int k = mMLAParams.q_lora_rank; + int lda = k, ldb = mMLAParams.q_lora_rank + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim; + int ldc = m; + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + mCublasWrapper->Gemm( + transa, transb, m, n, k, params.q_b_proj, lda, params.fused_a_input, ldb, q_buffer, ldc); + mCublasWrapper->destroyDescriptors(); + } + + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = mMLAParams.kv_lora_rank; + int n = params.acc_q_len; + int k = mMLAParams.qk_nope_head_dim; + int lda = k, ldb = params.head_num * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim); + int ldc = params.head_num * (mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim); + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, params.k_b_proj_trans, lda, + mMLAParams.qk_nope_head_dim * mMLAParams.kv_lora_rank, q_buffer, ldb, + mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim, params.attention_input_buf, ldc, + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim, params.head_num, 1.0f, 0.0); + mCublasWrapper->destroyDescriptors(); + } + } + else + { + { + int m = params.head_num * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim); + int n = params.acc_q_len; + int k = mMLAParams.q_lora_rank; + int lda = k, ldb = mMLAParams.q_lora_rank + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim; + int ldc = m; + size_t act_size = mGemmRunner->getFP8DataSize(n, ldb, true); + size_t act_scale_size = mGemmRunner->getActScaleSize(n, ldb); + __nv_fp8_e4m3* act_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, act_size)); + float* act_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, act_scale_size)); + __nv_fp8_e4m3 const* q_gemm_weight_ptr; + float const* q_gemm_scale_ptr; + if (mIsFP8Model) + { + q_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.q_b_proj); + q_gemm_scale_ptr = reinterpret_cast(params.q_b_scale); + } + else + { + size_t q_gemm_size = mGemmRunner->getFP8DataSize(m, lda, false); + size_t q_gemm_scale_size = mGemmRunner->getWeightScaleSize(m, lda); + __nv_fp8_e4m3* q_gemm_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, q_gemm_size)); + float* q_gemm_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, q_gemm_scale_size)); + mGemmRunner->fp8CS128x128(q_gemm_weight_buffer, q_gemm_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.q_b_proj), lda, m, stream); + q_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(q_gemm_weight_buffer); + q_gemm_scale_ptr = reinterpret_cast(q_gemm_scale_buffer); + } + mGemmRunner->fp8CS1x128(act_buffer, act_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.fused_a_input), ldb, n, stream); + mGemmRunner->gemm(act_buffer, ldb, const_cast<__nv_fp8_e4m3*>(q_gemm_weight_ptr), lda, + reinterpret_cast<__nv_bfloat16*>(q_buffer), ldc, n, m, k, act_scale_buffer, + const_cast(q_gemm_scale_ptr), stream); + } + { + offset = quant_offset; + int m = mMLAParams.kv_lora_rank; + int n = params.acc_q_len; + int k = mMLAParams.qk_nope_head_dim; + int lda = k, ldb = mMLAParams.qk_nope_head_dim; + int ldc = params.head_num * (mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim); + size_t act_size = mGemmRunner->getFP8DataSize(n, params.head_num * mMLAParams.qk_nope_head_dim, true); + size_t act_scale_size = mGemmRunner->getActScaleSize(n, params.head_num * mMLAParams.qk_nope_head_dim); + __nv_fp8_e4m3* act_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, act_size)); + float* act_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, act_scale_size)); + __nv_fp8_e4m3 const* k_gemm_weight_ptr; + float const* k_gemm_scale_ptr; + if (mIsFP8Model) + { + k_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.k_b_proj_trans); + k_gemm_scale_ptr = reinterpret_cast(params.k_b_trans_scale); + } + else + { + size_t k_gemm_size + = mGemmRunner->getFP8DataSize(params.head_num * m, mMLAParams.qk_nope_head_dim, false); + size_t k_gemm_scale_size + = mGemmRunner->getWeightScaleSize(params.head_num * m, mMLAParams.qk_nope_head_dim); + __nv_fp8_e4m3* k_gemm_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, k_gemm_size)); + float* k_gemm_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, k_gemm_scale_size)); + mGemmRunner->fp8CS128x128(k_gemm_weight_buffer, k_gemm_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.k_b_proj_trans), mMLAParams.qk_nope_head_dim, + params.head_num * m, stream); + k_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(k_gemm_weight_buffer); + k_gemm_scale_ptr = reinterpret_cast(k_gemm_scale_buffer); + } + int shape_n_4_align = ((n - 1) / 4 + 1) * 4; + mGemmRunner->fp8CS1x128Reshape(act_buffer, act_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(q_buffer), mMLAParams.qk_nope_head_dim, params.head_num, n, + mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim, stream); + + mGemmRunner->strideBatchGemm(reinterpret_cast<__nv_bfloat16*>(params.attention_input_buf), ldc, + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim, act_buffer, ldb, n * k, + const_cast<__nv_fp8_e4m3*>(k_gemm_weight_ptr), lda, m * k, params.head_num, n, m, k, stream, + act_scale_buffer, shape_n_4_align * (mMLAParams.qk_nope_head_dim / 128), + const_cast(k_gemm_scale_ptr)); + } + } params.seqQOffset = cu_q_seqlens; params.cu_kv_seqlens = cu_kv_seqlens; params.fmha_tile_counter = fmha_tile_counter_ptr; + params.q_buf = q_buffer; invokeMLARopeGeneration(params, kv_cache_buffer, stream); @@ -980,22 +1234,68 @@ int GPTAttentionPluginCommon::mlaGeneration( // Run the fmha kernel mDecoderFMHARunner->run(fmhaParams); - + if (!mIsPTP128CEnabled) + { + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = mMLAParams.v_head_dim; + int n = params.batch_size; + int k = mMLAParams.kv_lora_rank; + int lda = k, ldb = k; + int ldc = m; // params.head_size * params.c_k; + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + + mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, + params.kv_b_proj + params.head_num * mMLAParams.kv_lora_rank * mMLAParams.qk_nope_head_dim, lda, + mMLAParams.v_head_dim * mMLAParams.kv_lora_rank, o_buffer, ldb * params.head_num, ldb, + params.context_buf, ldc * params.head_num, ldc, params.head_num, 1.0f, 0.0); + mCublasWrapper->destroyDescriptors(); + } + } + else { - auto transa = CUBLAS_OP_T; - auto transb = CUBLAS_OP_N; + // ??? 16 128 + offset = quant_offset; int m = mMLAParams.v_head_dim; int n = params.batch_size; int k = mMLAParams.kv_lora_rank; + int h = params.head_num; int lda = k, ldb = k; - int ldc = m; // params.head_size * params.c_k; - mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - - mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, - params.kv_b_proj + params.head_num * mMLAParams.kv_lora_rank * mMLAParams.qk_nope_head_dim, lda, - mMLAParams.v_head_dim * mMLAParams.kv_lora_rank, o_buffer, ldb * params.head_num, ldb, params.context_buf, - ldc * params.head_num, ldc, params.head_num, 1.0f, 0.0); - mCublasWrapper->destroyDescriptors(); + int ldc = h * m; + size_t act_size = mGemmRunner->getFP8DataSize(n, h * k, true); + size_t act_scale_size = mGemmRunner->getActScaleSize(n, h * k); + __nv_fp8_e4m3* act_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, act_size)); + float* act_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, act_scale_size)); + __nv_fp8_e4m3 const* v_gemm_weight_ptr; + float const* v_gemm_scale_ptr; + if (mIsFP8Model) + { + v_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.kv_b_proj) + h * k * m; + v_gemm_scale_ptr = reinterpret_cast(params.kv_b_scale) + h * (k / 128) * (m / 128); + } + else + { + size_t v_gemm_size = mGemmRunner->getFP8DataSize(h * m, k, false); + size_t v_gemm_scale_size = mGemmRunner->getWeightScaleSize(h * m, k); + __nv_fp8_e4m3* v_gemm_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, v_gemm_size)); + float* v_gemm_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, v_gemm_scale_size)); + mGemmRunner->fp8CS128x128(v_gemm_weight_buffer, v_gemm_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.kv_b_proj) + h * m * k, k, h * m, stream); + + v_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(v_gemm_weight_buffer); + v_gemm_scale_ptr = reinterpret_cast(v_gemm_scale_buffer); + } + int shape_n_4_align = ((n - 1) / 4 + 1) * 4; + mGemmRunner->fp8CS1x128Reshape( + act_buffer, act_scale_buffer, reinterpret_cast<__nv_bfloat16 const*>(o_buffer), k, h, n, k, stream); + mGemmRunner->strideBatchGemm(reinterpret_cast<__nv_bfloat16*>(params.context_buf), ldc, m, act_buffer, ldb, + n * k, const_cast<__nv_fp8_e4m3*>(v_gemm_weight_ptr), lda, m * k, params.head_num, n, m, k, stream, + act_scale_buffer, shape_n_4_align * (k / 128), const_cast(v_gemm_scale_ptr)); } sync_check_cuda_error(); @@ -2368,6 +2668,17 @@ int GPTAttentionPluginCommon::initialize() noexcept // Only deepseek must using fmha. TLLM_CHECK_WITH_INFO(mFMHARunner->isFmhaSupported() && mDecoderFMHARunner->isFmhaSupported(), "Deepseek should be supported by fmha in context and generation part."); + + if (mIsPTP128CEnabled && mIsFP8Model) + { + mGemmRunner = std::make_shared>(); + } + else if (mIsPTP128CEnabled && !mIsFP8Model) + { + mGemmRunner = std::make_shared>(); + } } // Fall back to unfused MHA kernels if not supported. @@ -2448,8 +2759,8 @@ size_t GPTAttentionPluginCommon::getCommonSerializationSize() const noexcept + sizeof(mPagedContextFMHA) + sizeof(mFP8ContextFMHA) + sizeof(mHasFullAttentionMask) + sizeof(mUseKVCache) + sizeof(mUnfuseQkvGemm) + sizeof(mUseLognScaling) + sizeof(mIsSpecDecodingEnabled) + sizeof(mUseSpecDecoding) + sizeof(mSpecDecodingIsGenerationLengthVariable) + sizeof(mSpecDecodingMaxGenerationLength) - + sizeof(mNbMultiBlockSemaphores) + sizeof(mIsMLAEnabled) + sizeof(mMLAParams) + sizeof(mSkipAttn) - + sizeof(uint32_t) // size of DecoderXQARunnerResource buffer. + + sizeof(mNbMultiBlockSemaphores) + sizeof(mIsMLAEnabled) + sizeof(mMLAParams) + sizeof(mIsPTP128CEnabled) + + sizeof(mIsFP8Model) + sizeof(mSkipAttn) + sizeof(uint32_t) // size of DecoderXQARunnerResource buffer. + sizeof(mCpSize) + sizeof(mCpRank) + sizeof(int32_t) * mCpGroup.size() + DecoderXQARunner::getResourceGlobal()->getSerializationSize(); } @@ -2507,6 +2818,8 @@ void GPTAttentionPluginCommon::serializeCommon(void* buffer) const noexcept write(d, mSpecDecodingMaxGenerationLength); write(d, mIsMLAEnabled); write(d, mMLAParams); + write(d, mIsPTP128CEnabled); + write(d, mIsFP8Model); write(d, mNbMultiBlockSemaphores); write(d, mSkipAttn); write(d, mCpSize); @@ -2617,6 +2930,8 @@ GPTAttentionPluginCreatorCommon::GPTAttentionPluginCreatorCommon() mPluginAttributes.emplace_back(PluginField("qk_nope_head_dim", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(PluginField("qk_rope_head_dim", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(PluginField("v_head_dim", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back(PluginField("is_ptp128c_enabled", nullptr, PluginFieldType::kINT8, 0)); + mPluginAttributes.emplace_back(PluginField("is_fp8_model", nullptr, PluginFieldType::kINT8, 0)); mPluginAttributes.emplace_back(PluginField("skip_attn", nullptr, PluginFieldType::kINT8, 0)); mPluginAttributes.emplace_back(PluginField("cp_size", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(PluginField("cp_rank", nullptr, PluginFieldType::kINT32, 0)); diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h index 0ae6cba4b055..da1a05e2e89f 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h +++ b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h @@ -21,6 +21,7 @@ #include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.h" #include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/kvCacheUtils.h" #include "tensorrt_llm/kernels/mlaKernels.h" @@ -33,6 +34,8 @@ namespace tensorrt_llm::plugins { +using Fp8BlockScaleGemmRunnerPtr + = std::shared_ptr; class GPTAttentionPluginCommon : public BasePlugin { public: @@ -57,8 +60,8 @@ class GPTAttentionPluginCommon : public BasePlugin bool use_cache = true, bool is_spec_decoding_enabled = false, bool spec_decoding_is_generation_length_variable = false, int32_t spec_decoding_max_generation_length = 1, bool is_mla_enabled = false, int q_lora_rank = 0, int kv_lora_rank = 0, int qk_nope_head_dim = 0, - int qk_rope_head_dim = 0, int v_head_dim = 0, bool skip_attn = false, int cp_size = 1, int cp_rank = 0, - std::set cp_group = {}); + int qk_rope_head_dim = 0, int v_head_dim = 0, bool is_ptp128c_enabled = false, bool is_fp8_model = false, + bool skip_attn = false, int cp_size = 1, int cp_rank = 0, std::set cp_group = {}); GPTAttentionPluginCommon(void const* data, size_t length); @@ -415,6 +418,8 @@ class GPTAttentionPluginCommon : public BasePlugin int32_t mSpecDecodingMaxGenerationLength = 1; bool mIsMLAEnabled = false; tensorrt_llm::kernels::mlaMetaParams mMLAParams; + bool mIsPTP128CEnabled = false; + bool mIsFP8Model = false; int mCpSize = 1; int mCpRank = 0; std::set mCpGroup = {}; @@ -438,6 +443,7 @@ class GPTAttentionPluginCommon : public BasePlugin UniqPtrWNullCopy mFMHARunner; UniqPtrWNullCopy mDecoderFMHARunner; UniqPtrWNullCopy mDecoderXQARunner; + Fp8BlockScaleGemmRunnerPtr mGemmRunner; bool mMultiBlockMode; bool mEnableXQA; diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp index 3f94474a7883..5bcdecf76da7 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/batch_manager/contextProgress.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/unfusedAttentionKernels.h" @@ -63,8 +64,8 @@ GPTAttentionPlugin::GPTAttentionPlugin(int layer_idx, int num_heads, int vision_ bool use_paged_context_fmha, bool use_fp8_context_fmha, bool has_full_attention_mask, bool use_cache, bool is_spec_decoding_enabled, bool spec_decoding_is_generation_length_variable, int spec_decoding_max_generation_length, bool is_mla_enabled, int q_lora_rank, int kv_lora_rank, - int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool skip_attn, int cp_size, int cp_rank, - std::set cp_group) + int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool is_ptp128c_enabled, bool is_fp8_model, + bool skip_attn, int cp_size, int cp_rank, std::set cp_group) : GPTAttentionPluginCommon(layer_idx, num_heads, vision_start, vision_length, num_kv_heads, layer_idx_in_cache_pool, head_size, unidirectional, q_scaling, attn_logit_softcapping_scale, position_embedding_type, rotary_embedding_dim, rotary_embedding_base, rotary_embedding_scale_type, rotary_embedding_scale, @@ -74,9 +75,8 @@ GPTAttentionPlugin::GPTAttentionPlugin(int layer_idx, int num_heads, int vision_ type, max_context_length, qkv_bias_enabled, cross_attention, max_distance, pos_shift_enabled, dense_context_fmha, use_paged_context_fmha, use_fp8_context_fmha, has_full_attention_mask, use_cache, is_spec_decoding_enabled, spec_decoding_is_generation_length_variable, spec_decoding_max_generation_length, - is_mla_enabled, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, skip_attn, cp_size, - - cp_rank, cp_group) + is_mla_enabled, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, is_ptp128c_enabled, + is_fp8_model, skip_attn, cp_size, cp_rank, cp_group) { initEntryIdx(); } @@ -128,6 +128,12 @@ std::string GPTAttentionPlugin::toString(IdxEntry const& entry) const case IdxEntry::LONG_ROPE_ROTARY_COS_SIN: return "LONG_ROPE_ROTARY_COS_SIN"; case IdxEntry::HOST_RUNTIME_PERF_KNOBS: return "HOST_RUNTIME_PERF_KNOBS"; case IdxEntry::HOST_CONTEXT_PROGRESS: return "HOST_CONTEXT_PROGRESS"; + case IdxEntry::MLA_Q_B_PROJ_TENSOR: return "MLA_Q_B_PROJ_TENSOR"; + case IdxEntry::MLA_KV_B_PROJ_TENSOR: return "MLA_KV_B_PROJ_TENSOR"; + case IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR: return "MLA_K_B_PROJ_TRANS_TENSOR"; + case IdxEntry::MLA_Q_B_SCALE_TENSOR: return "MLA_Q_B_SCALE_TENSOR"; + case IdxEntry::MLA_KV_B_SCALE_TENSOR: return "MLA_KV_B_SCALE_TENSOR"; + case IdxEntry::MLA_K_B_TRANS_SCALE_TENSOR: return "MLA_K_B_TRANS_SCALE_TENSOR"; case IdxEntry::SKIP_ATTN: return "SKIP_ATTN"; case IdxEntry::ENUM_SIZE: return "ENUM_SIZE"; } @@ -179,9 +185,12 @@ bool GPTAttentionPlugin::isEntryUsed(IdxEntry const& entry) const case IdxEntry::MROPE_POSITION_DELTAS: return isMRoPE(); case IdxEntry::HOST_RUNTIME_PERF_KNOBS: return true; case IdxEntry::HOST_CONTEXT_PROGRESS: return true; - case IdxEntry::MLA_FUSED_Q_PROJ_TENSOR: return mIsMLAEnabled; case IdxEntry::MLA_Q_B_PROJ_TENSOR: return mIsMLAEnabled; case IdxEntry::MLA_KV_B_PROJ_TENSOR: return mIsMLAEnabled; + case IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR: return mIsMLAEnabled; + case IdxEntry::MLA_Q_B_SCALE_TENSOR: return mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model; + case IdxEntry::MLA_KV_B_SCALE_TENSOR: return mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model; + case IdxEntry::MLA_K_B_TRANS_SCALE_TENSOR: return mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model; case IdxEntry::SKIP_ATTN: return mSkipAttn; default: return false; } @@ -323,6 +332,20 @@ bool GPTAttentionPlugin::supportsFormatCombination( posCaseLine = __LINE__; result = inOut[pos].type == nvinfer1::DataType::kFLOAT; } + else if (mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model + && (pos == getIdx(IdxEntry::MLA_Q_B_SCALE_TENSOR) || pos == getIdx(IdxEntry::MLA_KV_B_SCALE_TENSOR) + || pos == getIdx(IdxEntry::MLA_K_B_TRANS_SCALE_TENSOR))) + { + posCaseLine = __LINE__; + result = inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; + } + else if (mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model + && (pos == getIdx(IdxEntry::MLA_Q_B_PROJ_TENSOR) || pos == getIdx(IdxEntry::MLA_KV_B_PROJ_TENSOR) + || pos == getIdx(IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR))) + { + posCaseLine = __LINE__; + result = inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; + } else if (isLongRoPE() && (pos == getIdx(IdxEntry::LONG_ROPE_ROTARY_INV_FREQ) || pos == getIdx(IdxEntry::LONG_ROPE_ROTARY_COS_SIN))) { @@ -555,6 +578,7 @@ size_t GPTAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* in = getWorkspaceSizeForGeneration(type, max_num_seq, max_kv_cache_length, max_num_tokens); size_t attention_input_workspace_size = 0; + size_t context_mla_fp8_quant_size = 0; if (mIsMLAEnabled) { int32_t const size_per_head @@ -565,6 +589,25 @@ size_t GPTAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* in size_t workspaces[1]; workspaces[0] = attention_input_size; attention_input_workspace_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspaces, 1); + + if (mIsPTP128CEnabled) + { + size_t act_quant_size = mGemmRunner->getActWorkspaceSize( + max_num_tokens, mMLAParams.qk_nope_head_dim + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim); + size_t weight_size = 0; + if (!mIsFP8Model) + { + weight_size = std::max(mGemmRunner->getWeightWorkspaceSize( + mNumHeads * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim), + mMLAParams.q_lora_rank), + mGemmRunner->getWeightWorkspaceSize( + 2 * mNumHeads * mMLAParams.qk_nope_head_dim, mMLAParams.kv_lora_rank)); + } + size_t workspaces[2]; + workspaces[0] = act_quant_size; + workspaces[1] = weight_size; + context_mla_fp8_quant_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspaces, 2); + } } else if (mUnfuseQkvGemm) { @@ -579,7 +622,8 @@ size_t GPTAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* in attention_input_workspace_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspaces, 1); } - return std::max(context_workspace_size, generation_workspace_size) + attention_input_workspace_size; + return std::max(context_mla_fp8_quant_size, std::max(context_workspace_size, generation_workspace_size)) + + attention_input_workspace_size; } static size_t getStride(nvinfer1::Dims const& dims, int n) @@ -668,9 +712,10 @@ mlaParams GPTAttentionPlugin::enqueueMLAPreprocess(int32_t localNbSeq, int32_ { auto const* input = static_cast(inputs[getIdx(IdxEntry::QKV_TENSOR)]); - auto const* fused_q_proj = static_cast(inputs[getIdx(IdxEntry::MLA_FUSED_Q_PROJ_TENSOR)]); auto const* q_b_proj = static_cast(inputs[getIdx(IdxEntry::MLA_Q_B_PROJ_TENSOR)]); auto const* kv_b_proj = static_cast(inputs[getIdx(IdxEntry::MLA_KV_B_PROJ_TENSOR)]); + auto const* k_b_proj_trans = static_cast(inputs[getIdx(IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR)]); + float2 const* cos_sin_cache = static_cast(inputs[getIdx(IdxEntry::ROTARY_COS_SIN)]); AttentionOutT* context_buf_ = static_cast(outputs[0]); @@ -678,14 +723,20 @@ mlaParams GPTAttentionPlugin::enqueueMLAPreprocess(int32_t localNbSeq, int32_ mlaParams mla_params; mla_params.fused_a_input = input; mla_params.context_buf = reinterpret_cast(context_buf_); - mla_params.fused_q_proj = fused_q_proj; mla_params.q_b_proj = q_b_proj; mla_params.kv_b_proj = kv_b_proj; + mla_params.k_b_proj_trans = k_b_proj_trans; mla_params.cos_sin_cache = cos_sin_cache; mla_params.batch_size = localNbSeq; mla_params.acc_q_len = localNbTokens; mla_params.head_num = mNumHeads; mla_params.meta = mMLAParams; + if (mIsPTP128CEnabled && mIsFP8Model) + { + mla_params.q_b_scale = static_cast(inputs[getIdx(IdxEntry::MLA_Q_B_SCALE_TENSOR)]); + mla_params.kv_b_scale = static_cast(inputs[getIdx(IdxEntry::MLA_KV_B_SCALE_TENSOR)]); + mla_params.k_b_trans_scale = static_cast(inputs[getIdx(IdxEntry::MLA_K_B_TRANS_SCALE_TENSOR)]); + } // { // __nv_bfloat16 *h_input; @@ -1363,6 +1414,8 @@ IPluginV2* GPTAttentionPluginCreator::createPlugin(char const* name, PluginField static_cast(p.getScalar("qk_nope_head_dim").value()), static_cast(p.getScalar("qk_rope_head_dim").value()), static_cast(p.getScalar("v_head_dim").value()), + static_cast(p.getScalar("is_ptp128c_enabled").value()), + static_cast(p.getScalar("is_fp8_model").value()), static_cast(p.getScalar("skip_attn").value()), static_cast(p.getScalar("cp_size").value()), static_cast(p.getScalar("cp_rank").value()), diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h index ebbeb04fb75a..489d5bb448ec 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h +++ b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h @@ -119,8 +119,8 @@ class GPTAttentionPlugin : public GPTAttentionPluginCommon bool use_cache = true, bool is_spec_decoding_enabled = false, bool spec_decoding_is_generation_length_variable = false, int spec_decoding_max_generation_length = 1, bool is_mla_enabled = false, int q_lora_rank = 0, int kv_lora_rank = 0, int qk_nope_head_dim = 0, - int qk_rope_head_dim = 0, int v_head_dim = 0, bool skip_attn = false, int cp_size = 1, int cp_rank = 0, - std::set cp_group = {}); + int qk_rope_head_dim = 0, int v_head_dim = 0, bool is_ptp128c_enabled = false, bool is_fp8_model = false, + bool skip_attn = false, int cp_size = 1, int cp_rank = 0, std::set cp_group = {}); GPTAttentionPlugin(void const* data, size_t length); @@ -226,9 +226,12 @@ class GPTAttentionPlugin : public GPTAttentionPluginCommon MROPE_POSITION_DELTAS, HOST_RUNTIME_PERF_KNOBS, HOST_CONTEXT_PROGRESS, - MLA_FUSED_Q_PROJ_TENSOR, MLA_Q_B_PROJ_TENSOR, MLA_KV_B_PROJ_TENSOR, + MLA_K_B_PROJ_TRANS_TENSOR, + MLA_Q_B_SCALE_TENSOR, + MLA_KV_B_SCALE_TENSOR, + MLA_K_B_TRANS_SCALE_TENSOR, SKIP_ATTN, LOGN_SCALING, ENUM_SIZE, // Used to count the number of IdxEntry, must put in last diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp index 61cec4842616..c4ab883c369c 100644 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp @@ -45,7 +45,8 @@ MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(bool remove_input_padding, int nu bool use_finished, bool use_bias, int tp_size, int tp_rank, int ep_size, int ep_rank, MOEExpertScaleNormalizationMode normalization_mode, float sparse_mixer_epsilon, bool force_determinism, int side_stream_id, MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, bool use_lora, - nvinfer1::DataType lora_type, LoraPluginProfilerPtr lora_profiler, int max_low_rank) + nvinfer1::DataType lora_type, LoraPluginProfilerPtr lora_profiler, int max_low_rank, bool use_deepseek = false, + bool use_deepseek_with_native_fp8_weights = false) : mRemoveInputPadding(remove_input_padding) , mNumExperts(number_of_experts) , mK(top_k) @@ -68,6 +69,8 @@ MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(bool remove_input_padding, int nu , mLoraType(lora_type) , mLoraProfiler(std::move(lora_profiler)) , mMaxLowRank(max_low_rank) + , mUseDeepSeek(use_deepseek) + , mUseDeepSeekWithNativeFp8Weights(use_deepseek_with_native_fp8_weights) { init(); } @@ -105,6 +108,8 @@ tensorrt_llm::plugins::MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(MixtureOfE , mLoraImpl2(other.mLoraImpl2) , mLayerName(other.mLayerName) , mNamespace(other.mNamespace) + , mUseDeepSeek(other.mUseDeepSeek) + , mUseDeepSeekWithNativeFp8Weights(other.mUseDeepSeekWithNativeFp8Weights) { init(); } @@ -116,7 +121,8 @@ size_t MixtureOfExpertsPlugin::getSerializationSize() const noexcept + sizeof(QuantMode::BaseType) + sizeof(mUseFinished) + sizeof(mUseBias) + sizeof(mParallelismConfig) + sizeof(mNormalizationMode) + sizeof(mSparseMixerEpsilon) + sizeof(mDims) + sizeof(mUseDeterministicKernels) + sizeof(mSideStreamId) + mGemmProfiler->getSerializationSize(mGemmId1) - + mGemmProfiler->getSerializationSize(mGemmId2) + sizeof(mUseLora) + sizeof(mLoraType) + sizeof(mMaxLowRank); + + mGemmProfiler->getSerializationSize(mGemmId2) + sizeof(mUseLora) + sizeof(mLoraType) + sizeof(mMaxLowRank) + + sizeof(mUseDeepSeek) + sizeof(mUseDeepSeekWithNativeFp8Weights); if (hasLora()) { @@ -157,6 +163,8 @@ MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(void const* data, size_t length, read(d, mUseLora); read(d, mLoraType); read(d, mMaxLowRank); + read(d, mUseDeepSeek); + read(d, mUseDeepSeekWithNativeFp8Weights); // Call init before deserialising the profiler to initialize mGemmId init(); @@ -202,6 +210,8 @@ void MixtureOfExpertsPlugin::serialize(void* buffer) const noexcept write(d, mUseLora); write(d, mLoraType); write(d, mMaxLowRank); + write(d, mUseDeepSeek); + write(d, mUseDeepSeekWithNativeFp8Weights); mGemmProfiler->serialize(d, mGemmId1); mGemmProfiler->serialize(d, mGemmId2); @@ -250,6 +260,10 @@ void MixtureOfExpertsPlugin::init() { mMOERunner = std::make_unique>(); } + else if (mType == DataType::kBF16 && mWeightType == DataType::kFP8) + { + mMOERunner = std::make_unique>(); + } else if (mType == DataType::kBF16 && mWeightType == DataType::kINT8) { mMOERunner = std::make_unique>(); @@ -314,6 +328,35 @@ void MixtureOfExpertsPlugin::init() TLLM_CUDA_CHECK(cudaEventCreate(&mMemcpyEvent)); } + + if (useDeepSeekWithNativeFp8Weights()) + { + TLLM_CHECK_WITH_INFO( + useDeepSeek(), "the mUseDeepSeek should be true when the mUseDeepSeekWithNativeFp8Weights is true"); + } + if (useDeepSeek()) + { + +#ifdef ENABLE_BF16 + if (mType == DataType::kBF16) + { + if (mWeightType == DataType::kBF16) + { + mBlockScaleGemmImplPtr + = std::make_shared>(); + } +#ifdef ENABLE_FP8 + else if (mWeightType == DataType::kFP8) + { + mBlockScaleGemmImplPtr + = std::make_shared>(); + } + } +#endif +#endif + } mSideStreamPtr = nullptr; mDebugStallMain = tensorrt_llm::runtime::utils::stallStream("TLLM_DEBUG_MOE_STALL_MAIN"); mDebugStallSide = tensorrt_llm::runtime::utils::stallStream("TLLM_DEBUG_MOE_STALL_SIDE"); @@ -380,6 +423,11 @@ bool MixtureOfExpertsPlugin::supportsFormatCombination( { return inOut[pos].type == DataType::kFLOAT; } + else if (useDeepSeekWithNativeFp8Weights() && getExpertDeepseekScale1Index() <= pos + && pos <= getExpertDeepseekScale2Index()) + { + return inOut[pos].type == DataType::kFLOAT; + } else if (hasExpertIntQuantScales() && getExpertIntQuantScale1Index() <= pos && pos <= getExpertIntQuantScale2Index()) { @@ -470,7 +518,7 @@ auto MixtureOfExpertsPlugin::setupWorkspace(void* base_ptr, int64_t num_tokens, size_t dtype_size = tensorrt_llm::common::getDTypeSize(mType); size_t moe_workspace_size = mMOERunner->getWorkspaceSize(num_tokens, mExpertHiddenSize, mExpertInterSize, - mNumExperts, mK, mActivationType, mNormalizationMode, mParallelismConfig, hasLora()); + mNumExperts, mK, mActivationType, mNormalizationMode, mParallelismConfig, hasLora(), useDeepSeek()); // Output of post-softmax routing probabilities size_t scale_probabilities_size = num_tokens * mNumExperts * sizeof(float); @@ -489,12 +537,24 @@ auto MixtureOfExpertsPlugin::setupWorkspace(void* base_ptr, int64_t num_tokens, mLoraImpl2->getWorkspaceSize(num_tokens * mK, num_reqs_lora, mLoraType)); } + size_t deepseek_workspace_size = 0; + if (useDeepSeek()) + { + bool is_gated_actiation = isGatedActivation(mActivationType); + int factor = is_gated_actiation ? 2 : 1; + size_t deepseek_fc1_size = mBlockScaleGemmImplPtr->getWorkspaceSize( + num_tokens * mK, factor * mExpertInterSize, mExpertHiddenSize, mNumExperts); + size_t deepseek_fc2_size = mBlockScaleGemmImplPtr->getWorkspaceSize( + num_tokens * mK, mExpertHiddenSize, mExpertInterSize, mNumExperts); + deepseek_workspace_size = std::max(deepseek_fc1_size, deepseek_fc2_size); + } std::vector workspaces{ moe_workspace_size, scale_probabilities_size, src_to_dest_map_size, selected_expert_size, lora_workspace_size, + deepseek_workspace_size, }; WorkspaceInfo info{}; @@ -507,6 +567,7 @@ auto MixtureOfExpertsPlugin::setupWorkspace(void* base_ptr, int64_t num_tokens, info.src_to_dest_map = nextWorkspacePtr((int8_t*) info.scale_probs, scale_probabilities_size); info.selected_experts = nextWorkspacePtr((int8_t*) info.src_to_dest_map, src_to_dest_map_size); info.lora_workspace = nextWorkspacePtr((int8_t*) info.selected_experts, selected_expert_size); + info.deepseek_workspace = nextWorkspacePtr((int8_t*) info.lora_workspace, lora_workspace_size); } return info; @@ -669,6 +730,17 @@ LoraParams MixtureOfExpertsPlugin::getLoraParams( mLoraExpandGatedWeightPtrs.data()); } +BlockScaleParams MixtureOfExpertsPlugin::getBlockScaleParams( + nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace) +{ + TLLM_CHECK(useDeepSeek()); + auto fc1_scales_ptr = static_cast(inputs[getExpertDeepseekScale1Index()]); + auto fc2_scales_ptr = static_cast(inputs[getExpertDeepseekScale2Index()]); + + return BlockScaleParams( + fc1_scales_ptr, fc2_scales_ptr, mBlockScaleGemmImplPtr, static_cast(workspace), &mMemcpyEvent); +} + int MixtureOfExpertsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace_ptr, cudaStream_t stream) noexcept @@ -767,6 +839,20 @@ int MixtureOfExpertsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, auto gemm1 = mGemmProfiler->getBestConfig(num_tokens, mGemmId1); auto gemm2 = mGemmProfiler->getBestConfig(num_tokens, mGemmId2); + + BlockScaleParams deepseek_params{}; + if (useDeepSeek()) + { + deepseek_params = getBlockScaleParams(inputDesc, inputs, workspace.deepseek_workspace); + auto config + = cutlass_extensions::CutlassGemmConfig(cutlass_extensions::CutlassTileConfigSM90::CtaShape128x16x128B, + cutlass_extensions::MainloopScheduleType::AUTO, cutlass_extensions::EpilogueScheduleType::AUTO, + cutlass_extensions::ClusterShape::ClusterShape_1x1x1); + + gemm1 = std::make_optional(config); + gemm2 = std::make_optional(config); + } + mMOERunner->setTactic(gemm1, gemm2); mMOERunner->runMoe(inputs[getInputTensorIndex()], static_cast(inputs[getRoutingTensorIndex()]), inputs[getExpertWeights1Index()], hasBias() ? inputs[getExpertBias1Index()] : nullptr, mActivationType, @@ -777,7 +863,7 @@ int MixtureOfExpertsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, hasFinishedTensor() ? static_cast(inputs[getFinishedTensorIndex()]) : nullptr, num_not_finished, workspace.scale_probs, static_cast(workspace.src_to_dest_map), static_cast(workspace.selected_experts), mSparseMixerEpsilon, mParallelismConfig, mNormalizationMode, - hasLora(), lora_params, stream); + hasLora(), lora_params, useDeepSeek(), deepseek_params, stream); if (useSideStream()) { @@ -815,6 +901,11 @@ char const* MixtureOfExpertsPlugin::getPluginVersion() const noexcept int MixtureOfExpertsPlugin::initialize() noexcept { + if (useDeepSeek()) + { + mLoraProfiler->setSkip(true); + mGemmProfiler->setSkip(true); + } mGemmProfiler->setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile::GEMM_1); mGemmProfiler->profileTactics(this, mType, mDims, mGemmId1); mGemmProfiler->setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile::GEMM_2); @@ -910,6 +1001,9 @@ MixtureOfExpertsPluginCreator::MixtureOfExpertsPluginCreator() mPluginAttributes.emplace_back(nvinfer1::PluginField("use_lora", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(nvinfer1::PluginField("lora_type_id", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(nvinfer1::PluginField("max_low_rank", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back(nvinfer1::PluginField("use_deepseek", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back( + nvinfer1::PluginField("use_deepseek_with_native_fp8_weights", nullptr, PluginFieldType::kINT32, 0)); mFC.nbFields = mPluginAttributes.size(); mFC.fields = mPluginAttributes.data(); } @@ -940,6 +1034,8 @@ IPluginV2* MixtureOfExpertsPluginCreator::createPlugin( int mUseLora{}; int mLoraType{INT_MAX}; int mMaxLowRank{0}; + int mUseDeepSeek{0}; + int mUseDeepSeekWithNativeFp8Weights{0}; float mSparseMixerEpsilon = -INFINITY; @@ -977,6 +1073,8 @@ IPluginV2* MixtureOfExpertsPluginCreator::createPlugin( MapPair{"side_stream_id", std::ref(mSideStreamId), true}, MapPair{"lora_type_id", std::ref(mLoraType), true}, MapPair{"max_low_rank", std::ref(mMaxLowRank), true}, + MapPair{"use_deepseek", std::ref(mUseDeepSeek), true}, + MapPair{"use_deepseek_with_native_fp8_weights", std::ref(mUseDeepSeekWithNativeFp8Weights), true}, }; for (int i = 0; i < fc->nbFields; ++i) { @@ -1036,7 +1134,8 @@ IPluginV2* MixtureOfExpertsPluginCreator::createPlugin( QuantMode(mQuantMode), mUseFinished != 0, mUseBias != 0, mTPSize, mTPRank, mEPSize, mEPRank, static_cast(mNormalizationMode), mSparseMixerEpsilon, mRequiresDeterminism != 0, mSideStreamId, gemmProfiler, mUseLora != 0, - static_cast(mLoraType), loraProfiler, mMaxLowRank); + static_cast(mLoraType), loraProfiler, mMaxLowRank, mUseDeepSeek, + mUseDeepSeekWithNativeFp8Weights); obj->setPluginNamespace(mNamespace.c_str()); return obj; } diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h index 6a3608460308..7959e9077117 100644 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h +++ b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h @@ -20,6 +20,7 @@ #include "NvInferPlugin.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/lora/lora.h" #include "tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h" #include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" @@ -100,6 +101,7 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt using MOEExpertScaleNormalizationMode = tensorrt_llm::kernels::MOEExpertScaleNormalizationMode; using LoraPluginProfilerPtr = std::shared_ptr; using LoraImplPtr = std::shared_ptr; + using BlockScaleGemmImplPtr = std::shared_ptr; MixtureOfExpertsPlugin() = delete; MixtureOfExpertsPlugin(bool remove_input_padding, int number_of_experts, int top_k, int expert_hidden_size, @@ -108,7 +110,8 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt bool use_finished, bool use_bias, int tp_size, int tp_rank, int ep_size, int ep_rank, MOEExpertScaleNormalizationMode normalization_mode, float sparse_mixer_epsilon, bool force_determinism, int side_stream_id, MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, bool use_lora, - nvinfer1::DataType lora_type, LoraPluginProfilerPtr lora_profiler, int max_low_rank); + nvinfer1::DataType lora_type, LoraPluginProfilerPtr lora_profiler, int max_low_rank, bool use_deepseek, + bool use_deepseek_with_native_fp8_weights); MixtureOfExpertsPlugin(void const* data, size_t length, MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, LoraPluginProfilerPtr lora_profiler); MixtureOfExpertsPlugin(MixtureOfExpertsPlugin const&); @@ -181,6 +184,11 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt MixtureOfExpertsPluginProfilerPtr mGemmProfiler; + // deepseek related + bool mUseDeepSeek{}; + bool mUseDeepSeekWithNativeFp8Weights{}; + BlockScaleGemmImplPtr mBlockScaleGemmImplPtr; + // lora related bool mUseLora{}; nvinfer1::DataType mLoraType{}; @@ -216,6 +224,7 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt void* src_to_dest_map{}; void* selected_experts{}; void* lora_workspace{}; + void* deepseek_workspace{}; size_t size{}; }; @@ -230,6 +239,9 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt kernels::LoraParams getLoraParams( nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace); + kernels::BlockScaleParams getBlockScaleParams( + nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace); + enum class RequestType : int32_t { kCONTEXT = 0, @@ -295,6 +307,16 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt return mUseLora; } + bool useDeepSeek() const + { + return mUseDeepSeek; + } + + bool useDeepSeekWithNativeFp8Weights() const + { + return mUseDeepSeekWithNativeFp8Weights; + } + bool hasGatedLoraWeightsAndRanks() const { return mUseLora && isGatedActivation(mActivationType); @@ -315,6 +337,16 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt return getExpertBias2Index() + hasFinishedTensor(); } + IndexType getExpertDeepseekScale1Index() const + { + return getFinishedTensorIndex() + useDeepSeekWithNativeFp8Weights(); + } + + IndexType getExpertDeepseekScale2Index() const + { + return getExpertDeepseekScale1Index() + useDeepSeekWithNativeFp8Weights(); + } + IndexType getExpertIntQuantScale1Index() const { return getFinishedTensorIndex() + hasExpertIntQuantScales(); @@ -327,7 +359,7 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt IndexType getExpertFP8Dequant1Index() const { - return getExpertIntQuantScale2Index() + hasExpertFp8QuantScales(); + return std::max(getExpertIntQuantScale2Index(), getExpertDeepseekScale2Index()) + hasExpertFp8QuantScales(); } IndexType getExpertFP8Quant2Index() const diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index 0eab70cac409..06fe4fe0045f 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -294,6 +294,7 @@ class FusedMoeRunner : public torch::CustomClassHolder kernels::QuantParams quant_params{}; kernels::LoraParams lora_params{}; + kernels::BlockScaleParams deepseek_params{}; mKernelRunner->runMoe(input.const_data_ptr(), gating_output.const_data_ptr(), fc1_expert_weights.const_data_ptr(), nullptr, activation_type, fc2_expert_weights.const_data_ptr(), nullptr, @@ -301,7 +302,7 @@ class FusedMoeRunner : public torch::CustomClassHolder static_cast(workspace_info.workspace), output.data_ptr(), nullptr, output.sizes()[0], workspace_info.scale_probs, static_cast(workspace_info.src_to_dest_map), static_cast(workspace_info.selected_experts), 0, parallelism_config, norm_mode, false, lora_params, - stream); + false, deepseek_params, stream); return output; } @@ -451,7 +452,7 @@ class FusedMoeRunner : public torch::CustomClassHolder kernels::MOEExpertScaleNormalizationMode norm_mode, kernels::MOEParallelismConfig const& parallelismConfig) { size_t moe_workspace_size = mKernelRunner->getWorkspaceSize(num_rows, hidden_size, inter_size, num_experts, - top_k, activation_type, norm_mode, parallelismConfig, /* use_lora */ false); + top_k, activation_type, norm_mode, parallelismConfig, /* use_lora */ false, /* use_deepseek */ false); size_t scale_prob_size = num_rows * num_experts * sizeof(float); size_t src_to_dest_map_size = top_k * num_rows * sizeof(int); size_t selected_expert_size = top_k * num_rows * sizeof(int); diff --git a/cpp/tests/kernels/mixtureOfExpertsTest.cu b/cpp/tests/kernels/mixtureOfExpertsTest.cu index 609e7df23f21..4592a66969bb 100644 --- a/cpp/tests/kernels/mixtureOfExpertsTest.cu +++ b/cpp/tests/kernels/mixtureOfExpertsTest.cu @@ -239,6 +239,7 @@ protected: bool mUseBias = true; bool mUseLora = false; + bool mUseDeepSeek = false; bool mIsGated = false; int64_t mGatedMultiplier = 1; @@ -288,8 +289,8 @@ protected: // Expert weights size_t const weight_size = hidden_size * (hidden_size * 4) * num_experts * sizeof(WeightStorage) * num_gemms; // Workspace size - size_t const workspace_size = this->mMoERunner.getWorkspaceSize( - num_tokens, hidden_size, hidden_size * 4, num_experts, k, this->mActType, mNormMode, {}, mUseLora); + size_t const workspace_size = this->mMoERunner.getWorkspaceSize(num_tokens, hidden_size, hidden_size * 4, + num_experts, k, this->mActType, mNormMode, {}, mUseLora, mUseDeepSeek); // The input/output buffers size_t const in_out_size = 2 * num_tokens * hidden_size * sizeof(DataType); @@ -333,8 +334,8 @@ protected: mTotalTokens += num_tokens; } - size_t workspace_size = mMoERunner.getWorkspaceSize( - mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK, mActType, mNormMode, parallelism_config, mUseLora); + size_t workspace_size = mMoERunner.getWorkspaceSize(mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK, + mActType, mNormMode, parallelism_config, mUseLora, mUseDeepSeek); auto const stream = mStream->get(); @@ -776,12 +777,14 @@ protected: } LoraParams lora_params; + BlockScaleParams blockscale_params; mMoERunner.setTactic(tactic1, tactic2); mMoERunner.runMoe(mInputTensor, mInputProbabilities, weight1_ptr, bias1_ptr, mActType, weight2_ptr, bias2_ptr, quant_params, mTotalTokens, mHiddenSize, mInterSize / parallelism_config.tp_size, mNumExperts, mK, mWorkspace, mFinalOutput, mFinished, mActiveRows, mScaleProbs, mSourceToExpandedMap, mSelectedExpert, - mSparseMixerEpsilon, parallelism_config, mNormMode, mUseLora, lora_params, stream); + mSparseMixerEpsilon, parallelism_config, mNormMode, mUseLora, lora_params, mUseDeepSeek, blockscale_params, + stream); check_cuda_error(cudaStreamSynchronize(stream)); } diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index 7cb200b904cc..6ea8963c70ab 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -16,21 +16,18 @@ This document shows how to build and run the `DeepSeek-v3` model in TensorRT-LLM ## Support Matrix -| Model | FP16 | BF16 | FP8 | W8A16 | W4A16 | TP | EP | IB | -| :------------- | :---: | :---: | :---: | :-----: | :-----: | :-----: | :-----: | :-----: | -| DeepSeek-V3 | Y | Y | . | Y | Y | Y | Y | Y | +| Model | FP16 | BF16 | FP8 | TP | EP | IB | +| :------------- | :---: | :---: | :---: | :-----: | :-----: | :-----: | +| DeepSeek-V3 | Y | Y | Y | Y | Y | Y | -- W8A16: INT8 Weight-Only -- W4A16: INT4 Weight-Only - TP: Tensor Parallel - EP: Expert Parallel - IB: Inflight Batching - FP8: Support for FP8 is currently in progress and will be released soon -***Please Note:*** +***Please Note:*** - Prefer using BF16 over FP16 for DeepSeek-V3 since model original training precision is FP8 and we found direct convert FP8 -> FP16 may cause unknown accuracy issues. -- Although Int8 and Int4 weight-only quantization can reduce the amount of total GPU memory needed, they may affect the accuracy to a certain degree. ## Prerequisite @@ -41,6 +38,13 @@ First, please download DeepSeek-V3 weights from HF https://huggingface.co/deepse git lfs install git clone https://huggingface.co/deepseek-ai/DeepSeek-V3-Base ``` +**Optional**: Convert the FP8 checkpoint to BF16. This is not necessary unless you want to run the model E2E in BF16 precision. +```bash +git clone https://github.com/deepseek-ai/DeepSeek-V3.git +cd DeepSeek-V3/inference/ +python fp8_cast_bf16.py --input-fp8-hf-path /path/to/DeepSeek-V3 --output-bf16-hf-path /path/to/deepseek-v3-bf16 +cp /path/to/DeepSeek-V3/config.json /path/to/DeepSeek-V3/configuration_deepseek.py /path/to/deepseek-v3-bf16/ +``` ## Hardware @@ -67,50 +71,49 @@ The TensorRT-LLM DeepSeek-V3 example code is located at [examples/deepseek_v3](. Below is the step-by-step to run DeepSeek-V3 with TensorRT-LLM. -Firstly, convert FP8 weights to BF16: -```bash -git clone https://github.com/deepseek-ai/DeepSeek-V3.git -cd DeepSeek-V3/inference/ -python fp8_cast_bf16.py --input-fp8-hf-path /path/to/DeepSeek-V3 --output-bf16-hf-path /path/to/deepseek-v3-bf16 -cp /path/to/DeepSeek-V3/config.json /path/to/DeepSeek-V3/configuration_deepseek.py /path/to/deepseek-v3-bf16/ -``` -Secondly, the BF16 checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. +Firstly, convert the checkpoint to the TensorRT-LLM checkpoint format by running [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. +To convert FP8 checkpoint: ```bash -# Convert Deepseek-v3 HF weights to TensorRT-LLM checkpoint in BF16. +# Convert Deepseek-v3 HF Native FP8 weights to TensorRT-LLM checkpoint. python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ - --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_16gpu_fp8 \ --dtype bfloat16 \ - --tp_size 32 \ + --use_fp8_weights \ + --tp_size 16 \ --workers 8 # using multiple workers can accelerate the conversion process +``` - -# Use Weight-Only Int8 quantization -python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ - --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ - --dtype bfloat16 \ - --tp_size 32 \ - --use_weight_only \ - --weight_only_precision int8 \ - --workers 8 - -# Use Weight-Only Int4 quantization +To convert BF16 checkpoint: +```bash +# Convert Deepseek-v3 HF weights to TensorRT-LLM checkpoint in BF16. python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ --dtype bfloat16 \ --tp_size 32 \ - --use_weight_only \ - --weight_only_precision int4 \ - --workers 8 + --workers 8 # using multiple workers can accelerate the conversion process ``` We observed the checkpoint conversion time took hours, while using a significant amount of CPU memory, please adjust the `--workers` parameter to balance your time and memory consumption. After the checkpoint conversion, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. +For FP8: +```bash +# Build FP8 engine +trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_16gpu_fp8 \ + --output_dir ./trtllm_engines/deepseek_v3/fp8/tp16-sel4096-isl2048-bs4 \ + --max_batch_size 4 \ + --max_seq_len 4096 \ + --max_input_len 2048 \ + --use_paged_context_fmha enable \ + --workers 8 +``` + +For BF16: ```bash -# Build engine +# Build BF16 engine trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ --output_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ --gpt_attention_plugin bfloat16 \ @@ -124,13 +127,13 @@ trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ ***Caution: `--max_batch_size` and `--max_seq_len` are the main factors to determine how many GPU memory will be used during runtime, so later when try to run e.g., `summarize.py` or `mmlu.py` or `gptManagerBenchmark.cpp`may need adjust `--max_batch_size` and `--max_seq_len` accordingly to avoid OOM.(meaning rebuild TensorRT engine with smaller `--max_batch_size` and `--max_seq_len` if needed based on GPU memory size), there is beautiful technical log perf-best-practices.md (https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/performance/perf-best-practices.md) explained the mechanism.*** -Test the engine with [run.py](../run.py) script: +Test the FP8 engines with [run.py](../run.py) script: ``` # run.sh python3 ../run.py --input_text "Today is a nice day." \ --max_output_len 30 \ --tokenizer_dir ./DeepSeek-V3 \ - --engine_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ + --engine_dir ./trtllm_engines/deepseek_v3/fp8/tp16-sel4096-isl2048-bs4 \ --top_p 0.95 \ --temperature 0.3 @@ -139,7 +142,7 @@ python3 ../run.py --input_text "Today is a nice day." \ For multi-nodes inference, let's take Slurm as an example using above command (run.sh): ```bash -srun -N 4 -w node-[1-4] --gres=gpu:8 --ntasks-per-node 8 \ +srun -N 2 -w node-[1-2] --gres=gpu:8 --ntasks-per-node 8 \ --container-image tensorrt_llm/release:latest \ --container-mounts ${PWD}:/workspace \ sh /workspace/command/run.sh @@ -168,4 +171,3 @@ and the output will be like: Input [Text 0]: "Today is a nice day." Output [Text 0 Beam 0]: " I am going to the park with my friends. We are going to play soccer. We are going" ``` - diff --git a/examples/deepseek_v3/convert_checkpoint.py b/examples/deepseek_v3/convert_checkpoint.py index 7cbc7935d225..338bb6436565 100644 --- a/examples/deepseek_v3/convert_checkpoint.py +++ b/examples/deepseek_v3/convert_checkpoint.py @@ -77,23 +77,10 @@ def parse_arguments(): 'To shard it along hidden dimension, set embedding_sharding_dim=1' 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) + parser.add_argument('--use_fp8_weights', + default=False, + action="store_true", + help='Use native FP8 weights of DeepSeek-V3.') parser.add_argument('--output_dir', type=str, @@ -131,36 +118,18 @@ def parse_arguments(): help= 'Only save the model config w/o read and converting weights, be careful, this is for debug only' ) - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact' - ) - # Add quantization related feature later + args = parser.parse_args() return args -def precision_to_config(precision, quant_config) -> QuantConfig: - '''update config dict for weight-only quantization - ''' - quant_config = QuantConfig() - precision_to_algo = {'int8': QuantAlgo.W8A16, 'int4': QuantAlgo.W4A16} - quant_config.quant_algo = precision_to_algo.get(precision) - return quant_config - - def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: '''return config dict with quantization info based on the command line args ''' quant_config = QuantConfig() - if args.use_weight_only: - quant_config = precision_to_config(args.weight_only_precision, - quant_config) + if args.use_fp8_weights: + quant_config.quant_algo = QuantAlgo.FP8_CURRENT_SCALING return quant_config diff --git a/tensorrt_llm/auto_parallel/tensor_parallel/plugin_nodes/gpt_attention_node.py b/tensorrt_llm/auto_parallel/tensor_parallel/plugin_nodes/gpt_attention_node.py index b63a73159d93..a44f55c1deb1 100644 --- a/tensorrt_llm/auto_parallel/tensor_parallel/plugin_nodes/gpt_attention_node.py +++ b/tensorrt_llm/auto_parallel/tensor_parallel/plugin_nodes/gpt_attention_node.py @@ -47,9 +47,12 @@ class IdxEntry(Enum): MROPE_POSITION_DELTAS = auto() HOST_RUNTIME_PERF_KNOBS = auto() HOST_CONTEXT_PROGRESS = auto() - MLA_FUSED_Q_PROJ_TENSOR = auto() MLA_Q_B_PROJ_TENSOR = auto() MLA_KV_B_PROJ_TENSOR = auto() + MLA_K_B_PROJ_TRANS_TENSOR = auto() + MLA_Q_B_SCALE_TENSOR = auto() + MLA_KV_B_SCALE_TENSOR = auto() + MLA_K_B_TRANS_SCALE_TENSOR = auto() LOGN_SCALING = auto() @@ -78,6 +81,9 @@ def __init__(self, plugin_info): self.is_spec_decoding_enabled = bool( plugin_info.pfc_as_list['is_spec_decoding_enabled'][0]) self.is_mla_enabled = bool(plugin_info.pfc_as_list['is_mla_enabled'][0]) + self.is_ptp128c_enabled = bool( + plugin_info.pfc_as_list['is_ptp128c_enabled'][0]) + self.is_fp8_model = bool(plugin_info.pfc_as_list['is_fp8_model'][0]) self.use_logn_scaling = bool( plugin_info.pfc_as_list['use_logn_scaling'][0]) self.init_entry_to_index() @@ -157,12 +163,18 @@ def is_entry_used(self, entry: IdxEntry) -> bool: return True elif entry == IdxEntry.HOST_CONTEXT_PROGRESS: return True - elif entry == IdxEntry.MLA_FUSED_Q_PROJ_TENSOR: - return self.is_mla_enabled elif entry == IdxEntry.MLA_Q_B_PROJ_TENSOR: return self.is_mla_enabled elif entry == IdxEntry.MLA_KV_B_PROJ_TENSOR: return self.is_mla_enabled + elif entry == IdxEntry.MLA_K_B_PROJ_TRANS_TENSOR: + return self.is_mla_enabled and self.is_ptp128c_enabled and self.is_fp8_model + elif entry == IdxEntry.MLA_Q_B_SCALE_TENSOR: + return self.is_mla_enabled and self.is_ptp128c_enabled and self.is_fp8_model + elif entry == IdxEntry.MLA_KV_B_SCALE_TENSOR: + return self.is_mla_enabled and self.is_ptp128c_enabled and self.is_fp8_model + elif entry == IdxEntry.MLA_K_B_TRANS_SCALE_TENSOR: + return self.is_mla_enabled elif entry == IdxEntry.LOGN_SCALING: return self.use_logn_scaling else: diff --git a/tensorrt_llm/builder.py b/tensorrt_llm/builder.py index 32382b1ec8f4..d0f49bd0bfd0 100644 --- a/tensorrt_llm/builder.py +++ b/tensorrt_llm/builder.py @@ -1148,6 +1148,7 @@ def build(model: PretrainedModel, build_config: BuildConfig) -> Engine: disable_weight_only_quant_plugin = model.config.disable_weight_only_quant_plugin if hasattr( model.config, 'disable_weight_only_quant_plugin') else False use_fp8_rowwise = model.config.quant_mode.has_fp8_rowwise() + use_fp8_current_scaling = model.config.quant_mode.has_fp8_current_scaling() if build_config.plugin_config.manage_weights: if use_weight_only and disable_weight_only_quant_plugin: @@ -1167,6 +1168,11 @@ def build(model: PretrainedModel, build_config: BuildConfig) -> Engine: if use_fp8_rowwise: network.plugin_config.set_fp8_rowwise_quant_plugins(model.config.dtype) + + if use_fp8_current_scaling: + network.plugin_config.set_fp8_current_scaling_gemm_plugins( + model.config.dtype) + nccl_plugin = model.config.dtype if model.config.mapping.world_size > 1 else None network.plugin_config.set_nccl_plugin(nccl_plugin) diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index 7a20c1dd8fec..7c7085477e3b 100755 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -4837,9 +4837,14 @@ def gpt_attention( qk_nope_head_dim: int = 0, qk_rope_head_dim: int = 0, v_head_dim: int = 0, - fused_q_proj: Optional[Tensor] = None, + is_ptp128c_enabled_flag: bool = False, + is_fp8_model_flag: bool = False, q_b_proj: Optional[Tensor] = None, kv_b_proj: Optional[Tensor] = None, + k_b_proj_trans: Optional[Tensor] = None, + q_b_scale: Optional[Tensor] = None, + kv_b_scale: Optional[Tensor] = None, + k_b_trans_scale: Optional[Tensor] = None, skip_attn=None, cp_group: List[int] = [0], cp_size: int = 1, @@ -5247,6 +5252,12 @@ def gpt_attention( v_head_dim = trt.PluginField("v_head_dim", np.array(v_head_dim, dtype=np.int32), trt.PluginFieldType.INT32) + is_ptp128c_enabled = trt.PluginField( + "is_ptp128c_enabled", np.array(is_ptp128c_enabled_flag, dtype=np.int8), + trt.PluginFieldType.INT8) + is_fp8_model = trt.PluginField("is_fp8_model", + np.array(is_fp8_model_flag, dtype=np.int8), + trt.PluginFieldType.INT8) p_dtype = default_net().plugin_config.gpt_attention_plugin pf_type = trt.PluginField( "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), @@ -5366,7 +5377,8 @@ def gpt_attention( is_spec_decoding_enabled, spec_decoding_is_generation_length_variable, spec_decoding_max_generation_length, is_mla_enabled, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, - skip_attn_pf, cp_size, cp_rank, cp_group, use_logn_scaling + is_ptp128c_enabled, is_fp8_model, skip_attn_pf, cp_size, cp_rank, + cp_group, use_logn_scaling ]) attn_plug = attn_plg_creator.create_plugin("causal_attn", pfc) @@ -5463,10 +5475,15 @@ def gpt_attention( plug_inputs += [host_context_progress] if is_mla_enabled_flag: - assert fused_q_proj is not None assert q_b_proj is not None assert kv_b_proj is not None - plug_inputs += [fused_q_proj, q_b_proj, kv_b_proj] + assert k_b_proj_trans is not None + plug_inputs += [q_b_proj, kv_b_proj, k_b_proj_trans] + if is_ptp128c_enabled_flag and is_fp8_model_flag: + assert q_b_scale is not None + assert kv_b_scale is not None + assert k_b_trans_scale is not None + plug_inputs += [q_b_scale, kv_b_scale, k_b_trans_scale] if skip_attn is not None: plug_inputs += [skip_attn] diff --git a/tensorrt_llm/layers/attention.py b/tensorrt_llm/layers/attention.py index 4bfc86ca8c43..aad6bcadf738 100755 --- a/tensorrt_llm/layers/attention.py +++ b/tensorrt_llm/layers/attention.py @@ -1956,6 +1956,13 @@ def __init__( self.rotary_scaling = rotary_scaling self.shard_dim = 1 + self.is_ptp128c_enabled_flag = False + self.is_fp8_model_flag = False + if quant_mode.has_fp8_current_scaling(): + self.is_ptp128c_enabled_flag = True + if not quant_mode.has_fp8_quantize_weights_on_demand(): + self.is_fp8_model_flag = True + def yarn_get_mscale(scale=1, mscale=1): if scale <= 1: return 1.0 @@ -2015,20 +2022,46 @@ def yarn_get_mscale(scale=1, mscale=1): self.kv_a_layernorm = RmsNorm(kv_lora_rank, dtype=dtype, eps=eps) - self.fused_q_proj = Parameter( - shape=(self.num_attention_heads * - (self.kv_lora_rank + self.qk_rope_head_dim), - self.q_lora_rank), - dtype=dtype) + if quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + self.kv_b_proj_scale = Parameter( + shape=(int(self.num_attention_heads * self.qk_nope_head_dim / + 128 * 2), int(self.kv_lora_rank / 128)), + dtype='float32') + self.k_b_proj_trans_scale = Parameter( + shape=(int(self.num_attention_heads * self.kv_lora_rank / 128), + int(self.qk_nope_head_dim / 128)), + dtype='float32') + self.q_b_proj_scale = Parameter(shape=(int( + self.num_attention_heads * + (self.qk_nope_head_dim + self.qk_rope_head_dim) / 128), + int(self.q_lora_rank / 128)), + dtype='float32') + set_obj_attrs(self.q_b_proj_scale, { + "weight_loader": self.weight_loader, + }) + set_obj_attrs(self.k_b_proj_trans_scale, { + "weight_loader": self.weight_loader, + }) + set_obj_attrs(self.kv_b_proj_scale, { + "weight_loader": self.weight_loader, + }) + + mla_weight_type = dtype if quant_mode.has_fp8_quantize_weights_on_demand() else 'fp8' + self.kv_b_proj = Parameter( shape=(self.num_attention_heads * self.qk_nope_head_dim * 2, self.kv_lora_rank), - dtype=dtype) + dtype=mla_weight_type) + self.k_b_proj_trans = Parameter( + shape=(self.num_attention_heads * self.kv_lora_rank, + self.qk_nope_head_dim), + dtype=mla_weight_type) self.q_b_proj = Parameter( shape=(self.num_attention_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim), self.q_lora_rank), - dtype=dtype) + dtype=mla_weight_type) self.dense = RowLinear(tp_size * self.num_attention_heads * self.v_head_dim, hidden_size, @@ -2036,15 +2069,15 @@ def yarn_get_mscale(scale=1, mscale=1): dtype=dtype, tp_group=tp_group, tp_size=tp_size) - set_obj_attrs(self.fused_q_proj, { - "weight_loader": self.weight_loader, - }) set_obj_attrs(self.q_b_proj, { "weight_loader": self.weight_loader, }) set_obj_attrs(self.kv_b_proj, { "weight_loader": self.weight_loader, }) + set_obj_attrs(self.k_b_proj_trans, { + "weight_loader": self.weight_loader, + }) def weight_loader(self, mapping: Mapping, param: Parameter, loaded_weight: torch.Tensor): @@ -2206,9 +2239,17 @@ def forward(self, qk_nope_head_dim=self.qk_nope_head_dim, qk_rope_head_dim=self.qk_rope_head_dim, v_head_dim=self.v_head_dim, - fused_q_proj=self.fused_q_proj.value, + is_ptp128c_enabled_flag=self.is_ptp128c_enabled_flag, + is_fp8_model_flag=self.is_fp8_model_flag, + k_b_proj_trans=self.k_b_proj_trans.value, q_b_proj=self.q_b_proj.value, - kv_b_proj=self.kv_b_proj.value) + kv_b_proj=self.kv_b_proj.value, + q_b_scale=self.q_b_proj_scale.value + if self.is_fp8_model_flag else None, + kv_b_scale=self.kv_b_proj_scale.value + if self.is_fp8_model_flag else None, + k_b_trans_scale=self.k_b_proj_trans_scale.value + if self.is_fp8_model_flag else None) context = self.dense(context) @@ -2227,77 +2268,96 @@ def split(v, tp_size, idx, dim=0): else: return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() - if tllm_key.endswith("kv_b_proj"): + if tllm_key.endswith("kv_b_proj") or tllm_key.endswith( + "kv_b_proj_scale"): + if tllm_key.endswith("kv_b_proj"): + qk_nope_head_dim = self.qk_nope_head_dim + v_head_dim = self.v_head_dim + kv_lora_rank = self.kv_lora_rank + else: + qk_nope_head_dim = int(self.qk_nope_head_dim / 128) + v_head_dim = int(self.v_head_dim / 128) + kv_lora_rank = int(self.kv_lora_rank / 128) kv_b_proj = weights.unflatten(0, [ self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.v_head_dim + qk_nope_head_dim + v_head_dim ]) splited_kv_b_proj = split(kv_b_proj, self.tp_size, self.tp_rank, dim=0) k_nope_weight, v_weight = splited_kv_b_proj.split( - [self.qk_nope_head_dim, self.v_head_dim], + [qk_nope_head_dim, v_head_dim], dim=1, ) kv_b_proj_weight = torch.concat([ k_nope_weight.reshape( - self.num_attention_heads * self.qk_nope_head_dim, - self.kv_lora_rank), - v_weight.reshape(self.num_attention_heads * self.v_head_dim, - self.kv_lora_rank) + self.num_attention_heads * qk_nope_head_dim, kv_lora_rank), + v_weight.reshape(self.num_attention_heads * v_head_dim, + kv_lora_rank) ], dim=0) return {tllm_key: kv_b_proj_weight} - elif tllm_key.endswith("q_b_proj"): - q_b_proj = weights.unflatten(0, [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.qk_rope_head_dim - ]) - splited_q_b_proj = split(q_b_proj, - self.tp_size, - self.tp_rank, - dim=0) - q_b_proj_weight = splited_q_b_proj.reshape( - self.num_attention_heads * - (self.qk_nope_head_dim + self.qk_rope_head_dim), - self.q_lora_rank) + elif tllm_key.endswith("q_b_proj") or tllm_key.endswith( + "q_b_proj_scale"): + if tllm_key.endswith("q_b_proj"): + q_b_proj = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.qk_rope_head_dim + ]) + splited_q_b_proj = split(q_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + q_b_proj_weight = splited_q_b_proj.reshape( + self.num_attention_heads * + (self.qk_nope_head_dim + self.qk_rope_head_dim), + self.q_lora_rank) + else: + splited_q_b_proj = split(weights, + self.tp_size, + self.tp_rank, + dim=0) + q_b_proj_weight = splited_q_b_proj.reshape( + int(self.num_attention_heads * + (self.qk_nope_head_dim + self.qk_rope_head_dim) / 128), + int(self.q_lora_rank / 128)) return {tllm_key: q_b_proj_weight} - elif tllm_key.endswith("fused_q_proj"): - assert isinstance(weights, list) and len(weights) == 2 - q_b_proj = weights[0].unflatten(0, [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.qk_rope_head_dim - ]) - splited_q_b_proj = split(q_b_proj, - self.tp_size, - self.tp_rank, - dim=0) - kv_b_proj = weights[1].unflatten(0, [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.v_head_dim - ]) - splited_kv_b_proj = split(kv_b_proj, - self.tp_size, - self.tp_rank, - dim=0) - q_nope_weight, q_pe_weight = splited_q_b_proj.split( - [self.qk_nope_head_dim, self.qk_rope_head_dim], - dim=1, - ) - k_nope_weight, _ = splited_kv_b_proj.split( - [self.qk_nope_head_dim, self.v_head_dim], - dim=1, - ) - fused_q_nope_weight = torch.einsum( - 'hdq,hdk->hkq', - q_nope_weight, - k_nope_weight, - ) - fused_q_weight = torch.cat( - [fused_q_nope_weight, q_pe_weight], - dim=1, - ).flatten(start_dim=0, end_dim=1) - return {tllm_key: fused_q_weight} + elif tllm_key.endswith("k_b_proj_trans") or tllm_key.endswith( + "k_b_proj_trans_scale"): + if tllm_key.endswith("k_b_proj_trans"): + kv_b_proj = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.v_head_dim + ]) + splited_kv_b_proj = split(kv_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + k_nope_weight, v_weight = splited_kv_b_proj.split( + [self.qk_nope_head_dim, self.v_head_dim], + dim=1, + ) + k_nope_weight_trans = k_nope_weight.transpose(2, 1).reshape( + self.num_attention_heads * self.kv_lora_rank, + self.qk_nope_head_dim) + else: + kv_b_proj_scale = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim // 128 + self.v_head_dim // 128 + ]) + splited_kv_b_proj_scale = split(kv_b_proj_scale, + self.tp_size, + self.tp_rank, + dim=0) + k_nope_scale, v_scale = splited_kv_b_proj_scale.split( + [self.qk_nope_head_dim // 128, self.v_head_dim // 128], + dim=1, + ) + k_nope_weight_trans = k_nope_scale.transpose(2, 1).reshape( + self.num_attention_heads * self.kv_lora_rank // 128, + self.qk_nope_head_dim // 128) + + return {tllm_key: k_nope_weight_trans} else: return {tllm_key: weights} diff --git a/tensorrt_llm/layers/moe.py b/tensorrt_llm/layers/moe.py index c1e83c933207..d3273b670d47 100755 --- a/tensorrt_llm/layers/moe.py +++ b/tensorrt_llm/layers/moe.py @@ -238,6 +238,19 @@ def from_parameter(x): "max_low_rank", np.array(lora_max_low_rank, dtype=np.int32), trt.PluginFieldType.INT32) + # Customized plugin inputs for DeepSeek-V3 + use_deepseek = trt.PluginField( + "use_deepseek", + np.array([int(quant_mode.has_fp8_current_scaling())], dtype=np.int32), + trt.PluginFieldType.INT32) + + use_deepseek_with_native_fp8_weights = not quant_mode.has_fp8_quantize_weights_on_demand( + ) + use_deepseek_with_native_fp8_weights = trt.PluginField( + "use_deepseek_with_native_fp8_weights", + np.array([int(use_deepseek_with_native_fp8_weights)], dtype=np.int32), + trt.PluginFieldType.INT32) + pfc_inputs = [ p_remove_input_padding, p_num_experts, p_top_k, p_expert_hidden_size, p_expert_inter_size, p_activation_type, p_type_id, p_weight_type_id, @@ -250,6 +263,9 @@ def from_parameter(x): if use_lora: pfc_inputs += [p_lora_type_id, p_max_low_rank] + if quant_mode.has_fp8_current_scaling(): + pfc_inputs += [use_deepseek, use_deepseek_with_native_fp8_weights] + pfc = trt.PluginFieldCollection(pfc_inputs) # Create the plugin with our constant inputs to the constructor @@ -274,6 +290,12 @@ def from_parameter(x): assert expert_scale_2 plugin_inputs += [expert_scale_1, expert_scale_2] + if quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + assert expert_scale_1 + assert expert_scale_2 + plugin_inputs += [expert_scale_1, expert_scale_2] + # Add conditional inputs if quant_mode.has_fp8_qdq(): assert expert_scale_3 @@ -284,6 +306,12 @@ def from_parameter(x): assert output_dtype == trt.fp8 plugin_inputs += [expert_scale_4] + if quant_mode.has_fp8_current_scaling( + ) and hidden_states_raw.dtype == trt.fp8: + assert expert_scale_3 + assert expert_scale_4 + plugin_inputs += [expert_scale_3, expert_scale_4] + if use_lora: if quant_mode.has_fp8_qdq(): assert act_scale @@ -353,6 +381,11 @@ def __init__(self, in_features: int, out_features: int, self.tp_dim = tp_dim self.is_padded = False + if quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + self.dtype = 'fp8' + self.weight_dtype = 'fp8' + if quant_mode.is_weight_only(): bytes_per_col_scale = 2 if quant_mode.is_int4_weight_only() else 1 # We use a different shape here because the quantized weights have their own layout @@ -365,7 +398,7 @@ def __init__(self, in_features: int, out_features: int, self.register_parameter('per_channel_scale', None) self.weight = Parameter(shape=self.expert_shape, - dtype=weight_dtype, + dtype=self.weight_dtype, prefer_managed=True) if has_bias: @@ -379,12 +412,20 @@ def __init__(self, in_features: int, out_features: int, dtype=trt.float32) self.weights_scaling_factor = Parameter(shape=(experts_per_node, 1), dtype=trt.float32) + elif quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + self.register_parameter('activation_scaling_factor', None) + self.weights_scaling_factor = Parameter(shape=(experts_per_node, + out_features // 128, + in_features // 128), + dtype=trt.float32) else: self.register_parameter('activation_scaling_factor', None) self.register_parameter('weights_scaling_factor', None) def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("weight"): + if tllm_key.endswith("weight") or tllm_key.endswith( + "weights_scaling_factor"): if isinstance(weights, torch.Tensor): weights = [weights] if "fc" in tllm_key: @@ -395,9 +436,11 @@ def postprocess(self, tllm_key, weights, **kwargs): dim=-2) elif "proj" in tllm_key: weights = torch.stack(weights) - weights = weights.to(str_dtype_to_torch(self.dtype)) + if tllm_key.endswith("weight"): + weights = weights.to(str_dtype_to_torch(self.dtype)) - if not self.quant_mode.has_any_quant(): + if not self.quant_mode.has_any_quant( + ) or self.quant_mode.has_fp8_current_scaling(): return weights elif self.quant_mode.is_weight_only(): if "per_channel_scale" in tllm_key: @@ -476,14 +519,18 @@ def __init__(self, self.weight_dtype = trt.int8 elif quant_mode.has_fp8_qdq(): self.weight_dtype = trt.fp8 - + elif quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + self.weight_dtype = trt.fp8 rank_experts = self.mapping.ep_experts(self.num_experts) self.wrapper_tllm_to_externel_key_dict = { "mlp": "block_sparse_moe", "proj": [f"experts.{expert}.w2" for expert in rank_experts], "fc": [f"experts.{expert}.w3" for expert in rank_experts] + - [f"experts.{expert}.w1" for expert in rank_experts] + [f"experts.{expert}.w1" for expert in rank_experts], + "weights_scaling_factor": + "weight_scale_inv", } # Since output dimension is usually low (in the order of 10s), no TP at @@ -649,6 +696,32 @@ def forward_experts(self, hidden_states, routing, finished, raise RuntimeError( "Cannot output FP8 value without knowing quantization parameter" ) + elif self.quant_mode.has_fp8_current_scaling( + ) and not self.quant_mode.has_fp8_quantize_weights_on_demand(): + assert self.fc.weight.value.dtype == trt.fp8, ( + "mlp fc weight dtype should be fp8 if not quantize on demand.") + assert self.proj.weight.value.dtype == trt.fp8, ( + "mlp proj weight dtype should be fp8 if not quantize on demand." + ) + hidden_states_quant = hidden_states + if hidden_states_quant.dtype != trt.fp8: + pass + + dtype_quant = self.dtype + weight_dtype_quant = trt.fp8 + + fc1_weight_scales = self.fc.weights_scaling_factor.value + fc2_weight_scales = self.proj.weights_scaling_factor.value + # self.fc.activation_scaling_factor.value + # self.proj.activation_scaling_factor.value + + scale_1 = fc1_weight_scales + scale_2 = fc2_weight_scales + scale_3 = None # TODO(Jerry Shi): change to fc1_act_scales if cs is fused in RMSNorm + scale_4 = None # TODO(Jerry Shi): change to fc2_act_scales if cs is fused in Swiglu + scale_5 = None + + output_dtype_quant = self.dtype else: hidden_states_quant = hidden_states diff --git a/tensorrt_llm/models/deepseek_v2/convert.py b/tensorrt_llm/models/deepseek_v2/convert.py index 158b527fc70f..271481cbc038 100755 --- a/tensorrt_llm/models/deepseek_v2/convert.py +++ b/tensorrt_llm/models/deepseek_v2/convert.py @@ -214,15 +214,14 @@ def convert_layer(l): dim=1, ) - q_nope_weight, q_pe_weight = q_b_proj_weight.split( - [qk_nope_head_dim, qk_rope_head_dim], - dim=1, - ) k_nope_weight, v_weight = kv_b_proj_weight.split( [qk_nope_head_dim, v_head_dim], dim=1, ) + k_nope_weight_trans = k_nope_weight.transpose(2, 1).reshape( + num_heads // mapping.tp_size * kv_lora_rank, qk_nope_head_dim) + if q_lora_rank is None: q_b_proj_weight = q_b_proj_weight.reshape( num_heads * (qk_nope_head_dim + qk_rope_head_dim) // @@ -240,32 +239,21 @@ def convert_layer(l): ], dim=0) - # Fuse matrices for decompression - fused_q_nope_weight = torch.einsum( - 'hdq,hdk->hkq', - q_nope_weight, - k_nope_weight, - ) - fused_q_weight = torch.cat( - [fused_q_nope_weight, q_pe_weight], - dim=1, - ).flatten(start_dim=0, end_dim=1) - weights.update( get_tllm_linear_weight(fused_a_weight, trtllm_prex + 'attention.fused_a.')) weights.update( get_tllm_linear_weight(kv_a_layernorm_weight, trtllm_prex + 'attention.kv_a_layernorm.')) - weights.update( - get_param_weight(fused_q_weight, - trtllm_prex + 'attention.fused_q_proj')) weights.update( get_param_weight(q_b_proj_weight, trtllm_prex + 'attention.q_b_proj')) weights.update( get_param_weight(kv_b_proj_weight, trtllm_prex + 'attention.kv_b_proj')) + weights.update( + get_param_weight(k_nope_weight_trans, + trtllm_prex + 'attention.k_b_proj_trans')) weights.update( get_tllm_linear_weight( o_proj_weight, diff --git a/tensorrt_llm/models/deepseek_v2/model.py b/tensorrt_llm/models/deepseek_v2/model.py index 4fa7798b3263..1e30a1c85a33 100755 --- a/tensorrt_llm/models/deepseek_v2/model.py +++ b/tensorrt_llm/models/deepseek_v2/model.py @@ -69,7 +69,8 @@ def __init__(self, config: DeepSeekV2Config, layer_idx: int): rotary_scaling=config.rotary_scaling, tp_group=config.mapping.tp_group, tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank) + tp_rank=config.mapping.tp_rank, + quant_mode=config.quant_mode) ### Added deepseek MoE and shared_experts ### First decoder layer: MLA + dense MLP + input_layernorm(RMSNorm) + post_attention_layernorm(RMSNorm) @@ -252,8 +253,11 @@ def from_hugging_face(cls, "q_a_layernrom": "q_a_layernorm", "kv_a_layernorm": "kv_a_layernorm", "q_b_proj": "q_b_proj.weight", + "q_b_proj_scale": "q_b_proj.weight_scale_inv", "kv_b_proj": "kv_b_proj.weight", - "fused_q_proj": ["q_b_proj.weight", "kv_b_proj.weight"], + "kv_b_proj_scale": "kv_b_proj.weight_scale_inv", + "k_b_proj_trans": "kv_b_proj.weight", + "k_b_proj_trans_scale": "kv_b_proj.weight_scale_inv", "shared_expert": "shared_experts", "e_score_correction_bias": "gate.e_score_correction_bias", } @@ -262,8 +266,11 @@ def from_hugging_face(cls, "fused_a": "kv_a_proj_with_mqa", "kv_a_layernorm": "kv_a_layernorm", "q_b_proj": "q_proj.weight", + "q_b_proj_scale": "q_proj.weight_scale_inv", "kv_b_proj": "kv_b_proj.weight", - "fused_q_proj": ["q_proj.weight", "kv_b_proj.weight"], + "kv_b_proj_scale": "kv_b_proj.weight_scale_inv", + "k_b_proj_trans": "kv_b_proj.weight", + "k_b_proj_trans_scale": "kv_b_proj.weight_scale_inv", "shared_expert": "shared_experts", "e_score_correction_bias": "gate.e_score_correction_bias", } diff --git a/tensorrt_llm/plugin/plugin.py b/tensorrt_llm/plugin/plugin.py index c6f4e873455d..2bd68f92f528 100644 --- a/tensorrt_llm/plugin/plugin.py +++ b/tensorrt_llm/plugin/plugin.py @@ -191,6 +191,14 @@ class PluginConfig(metaclass=PluginConfigMeta): "activation and per channel static scales for weights." "Note: It also requires same calibration in checkpoint." }) + _fp8_current_scaling_gemm_plugin: Optional[str] = field( + default=None, + init=False, + metadata={ + "help": + "The quantized GEMM for fp8, which uses per block dynamic scales for " + "activation and per block static scales for weights." + }) _qserve_gemm_plugin: Optional[str] = field( default=None, init=False, @@ -506,6 +514,10 @@ def set_fp8_rowwise_quant_plugins(self, dtype: str = "auto"): self.quantize_tensor_plugin = True return self + def set_fp8_current_scaling_gemm_plugins(self, dtype: str = "auto"): + self.fp8_current_scaling_gemm_plugin = dtype + return self + def set_context_fmha(self, context_fmha_type=ContextFMHAType.enabled): assert type(context_fmha_type) == ContextFMHAType self.context_fmha_type = context_fmha_type diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py index 4a028ffe1077..275c750efc4b 100644 --- a/tensorrt_llm/quantization/functional.py +++ b/tensorrt_llm/quantization/functional.py @@ -212,6 +212,56 @@ def fp8_rowwise_gemm(input: Tensor, weights: Tensor, scales_a: Tensor, return _create_tensor(layer.get_output(0), layer) +def fp8_current_scaling_gemm( + input: Tensor, + weights: Tensor, + in_scales: Optional[Tensor] = None, + weight_scales: Optional[Tensor] = None, + need_quantize_acts_on_demand: bool = True, + need_quantize_weights_on_demand: bool = False) -> Tensor: + if not default_net().plugin_config.fp8_current_scaling_gemm_plugin: + raise TypeError( + "Fp8 current scaling GEMM is only supported with plugin") + else: + plg_creator = trt.get_plugin_registry().get_plugin_creator( + 'Fp8CurrentScalingGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) + assert plg_creator is not None + + need_quantize_acts_on_demand = 1 if need_quantize_acts_on_demand else 0 + need_quantize_acts_on_demand = trt.PluginField( + "need_quantize_acts_on_demand", + np.array(need_quantize_acts_on_demand, dtype=np.int32), + trt.PluginFieldType.INT32) + + need_quantize_weights_on_demand = 1 if need_quantize_weights_on_demand else 0 + need_quantize_weights_on_demand = trt.PluginField( + "need_quantize_weights_on_demand", + np.array(need_quantize_weights_on_demand, dtype=np.int32), + trt.PluginFieldType.INT32) + + p_dtype = default_net().plugin_config.fp8_current_scaling_gemm_plugin + pf_type = trt.PluginField( + "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), + trt.PluginFieldType.INT32) + + pfc = trt.PluginFieldCollection([ + need_quantize_acts_on_demand, need_quantize_weights_on_demand, + pf_type + ]) + fp8_cs_gemm_plug = plg_creator.create_plugin("fp8_current_scaling_gemm", + pfc) + plug_inputs = [input.trt_tensor, weights.trt_tensor] + if weight_scales is not None: + plug_inputs += [weight_scales.trt_tensor] + if in_scales is not None: + plug_inputs += [in_scales.trt_tensor] + layer = default_trtnet().add_plugin_v2(plug_inputs, fp8_cs_gemm_plug) + _add_plugin_info(layer, plg_creator, "fp8_current_scaling_gemm", pfc) + if not default_net().strongly_typed: + layer.get_input(1).set_dynamic_range(-448, 448) + return _create_tensor(layer.get_output(0), layer) + + def weight_only_quant_matmul(input: Tensor, weights: Tensor, scales: Tensor, diff --git a/tensorrt_llm/quantization/layers.py b/tensorrt_llm/quantization/layers.py index c91134602e39..2f4a581ca6a8 100644 --- a/tensorrt_llm/quantization/layers.py +++ b/tensorrt_llm/quantization/layers.py @@ -36,12 +36,13 @@ # isort: off from .functional import ( - dequantize, fp8_rowwise_gemm, fp8_rowwise_rms_norm, postprocess_fp8_rowwise, - postprocess_weight_only, postprocess_weight_only_groupwise, quantize, - quantize_fp8_per_token, quantize_per_token, quantize_tensor, - validate_group_size, smooth_quant_gemm, smooth_quant_layer_norm, - smooth_quant_rms_norm, weight_only_groupwise_quant_matmul, - weight_only_quant_matmul, qserve_gemm_per_group, qserve_gemm_per_channel) + dequantize, fp8_current_scaling_gemm, fp8_rowwise_gemm, + fp8_rowwise_rms_norm, postprocess_fp8_rowwise, postprocess_weight_only, + postprocess_weight_only_groupwise, quantize, quantize_fp8_per_token, + quantize_per_token, quantize_tensor, validate_group_size, smooth_quant_gemm, + smooth_quant_layer_norm, smooth_quant_rms_norm, + weight_only_groupwise_quant_matmul, weight_only_quant_matmul, + qserve_gemm_per_group, qserve_gemm_per_channel) # isort: on from .mode import QuantMode @@ -1620,6 +1621,177 @@ def postprocess(self, tllm_key, weights, **kwargs): return weights.to(str_dtype_to_torch(self.bias.dtype)) +class FP8CurrentScalingLinear(Linear): + + def __init__(self, + in_features, + out_features, + bias=False, + dtype=None, + tp_group=None, + tp_size=1, + gather_output=True, + prefer_managed_weight=True, + is_qkv=False, + quantize_weights_on_demand=False): + super().__init__(in_features, + out_features, + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + gather_output=gather_output, + prefer_managed_weight=prefer_managed_weight, + is_qkv=is_qkv) + self.quantize_weights_on_demand = quantize_weights_on_demand + if quantize_weights_on_demand: + self.weight = Parameter(shape=(self.out_features, self.in_features), + dtype=dtype, + prefer_managed=self.prefer_managed_weight) + else: + self.weight = Parameter(shape=(self.out_features, self.in_features), + dtype='fp8', + prefer_managed=self.prefer_managed_weight) + self.weights_scaling_factor = Parameter(shape=(math.ceil(self.out_features / 128), \ + math.ceil(self.in_features / 128)), dtype=trt.float32) + + # TODO: modify this for checkpoint loading + self.tllm_to_externel_key_dict = { + "weight": "weight", + } + + if not quantize_weights_on_demand: + self.tllm_to_externel_key_dict.update( + {"weights_scaling_factor": "weight_scale_inv"}) + + def forward(self, x, in_scales=None, lora_runtime_params=None): + weights_scaling_factor = None + w_in = self.weight.value + need_quantize_acts_on_demand = (x.dtype != trt.fp8) + need_quantize_weights_on_demand = self.quantize_weights_on_demand + if not need_quantize_acts_on_demand: + assert in_scales is not None, "FP8 input requires activation scales." + if not need_quantize_weights_on_demand: + weights_scaling_factor = self.weights_scaling_factor.value + + gemm_output = fp8_current_scaling_gemm(x, w_in, weights_scaling_factor, + in_scales, + need_quantize_acts_on_demand, + need_quantize_weights_on_demand) + + return self.collect_and_bias(gemm_output) + + def postprocess(self, tllm_key, weights, **kwargs): + # TODO: add checkpoint conversion logic + if tllm_key.endswith("scaling_factor"): + if isinstance(weights, list) and len(weights) == 2: + left = weights[0].to(torch.float32) + right = weights[1].to(torch.float32) + weights = torch.cat( + [left, right], + dim=0, + ) + return weights + else: + return weights.to(torch.float32) + elif tllm_key.endswith("weight"): + weight_dtype = str_dtype_to_torch( + self.dtype + ) if self.quantize_weights_on_demand else torch.float8_e4m3fn + if isinstance(weights, list) and len(weights) == 2: + left = weights[0].view(weight_dtype) + right = weights[1].view(weight_dtype) + weights = torch.cat( + [left, right], + dim=0, + ) + return weights + else: + return weights.view(weight_dtype) + elif tllm_key.endswith("bias"): + return weights.to(str_dtype_to_torch(self.bias.dtype)) + + +class FP8CurrentScalingRowLinear(RowLinear): + + def __init__(self, + in_features, + out_features, + bias=False, + dtype=None, + tp_group=None, + tp_size=1, + prefer_managed_weight=True, + is_expert=False, + quantize_weights_on_demand=False): + super().__init__(in_features, + out_features, + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + prefer_managed_weight=prefer_managed_weight, + is_expert=is_expert) + self.quantize_weights_on_demand = quantize_weights_on_demand + if quantize_weights_on_demand: + self.weight = Parameter( + shape=(self.out_features, self.in_features), + dtype=dtype, + prefer_managed=self.prefer_managed_weight, + ) + else: + self.weight = Parameter( + shape=(self.out_features, self.in_features), + dtype="fp8", + prefer_managed=self.prefer_managed_weight, + ) + self.weights_scaling_factor = Parameter(shape=(math.ceil(self.out_features / 128), \ + math.ceil(self.in_features / 128)), dtype=trt.float32) + + self.tllm_to_externel_key_dict = { + "weight": "weight", + } + + if not quantize_weights_on_demand: + self.tllm_to_externel_key_dict.update( + {"weights_scaling_factor": "weight_scale_inv"}) + + def forward(self, + x, + in_scales=None, + lora_runtime_params=None, + all_reduce_params=None): + + weights_scaling_factor = None + w_in = self.weight.value + + need_quantize_acts_on_demand = (x.dtype != trt.fp8) + need_quantize_weights_on_demand = self.quantize_weights_on_demand + if not need_quantize_acts_on_demand: + assert in_scales is not None, "FP8 input requires activation scales." + if not need_quantize_weights_on_demand: + weights_scaling_factor = self.weights_scaling_factor.value + gemm_output = fp8_current_scaling_gemm(x, w_in, weights_scaling_factor, + in_scales, + need_quantize_acts_on_demand, + need_quantize_weights_on_demand) + + return self.collect_and_bias(gemm_output, + all_reduce_params=all_reduce_params) + + def postprocess(self, tllm_key, weights, **kwargs): + # TODO: add checkpoint conversion logic + if tllm_key.endswith("scaling_factor"): + return weights.to(torch.float32) + elif tllm_key.endswith("weight"): + weight_dtype = str_dtype_to_torch( + self.dtype + ) if self.quantize_weights_on_demand else torch.float8_e4m3fn + return weights.view(weight_dtype) + elif tllm_key.endswith("bias"): + return weights.to(str_dtype_to_torch(self.bias.dtype)) + + class Fp8RowwiseMLP(Module): def __init__( diff --git a/tensorrt_llm/quantization/mode.py b/tensorrt_llm/quantization/mode.py index 07f830474125..a89d0138ae40 100644 --- a/tensorrt_llm/quantization/mode.py +++ b/tensorrt_llm/quantization/mode.py @@ -36,6 +36,7 @@ class QuantAlgo(StrEnum, metaclass=BaseEnumMeta): W4A8_QSERVE_PER_CHANNEL = auto() FP8 = auto() FP8_PER_CHANNEL_PER_TOKEN = auto() + FP8_CURRENT_SCALING = auto() INT8 = auto() MIXED_PRECISION = auto() NO_QUANT = auto() @@ -78,6 +79,10 @@ class QuantMode(IntFlag): FP8_QDQ = auto() # FP8 rowwise FP8_ROWWISE = auto() + # FP8 current scaling + FP8_1x128_128x128_CURRENT_SCALING = auto() + # FP8 current scaling + FP8_QUANTIZE_WEIGHTS_ON_DEMAND = auto() # W4A8 qserve W4A8_QSERVE = auto() @@ -130,6 +135,12 @@ def has_act_or_weight_quant(self): def has_per_token_dynamic_scaling(self): return self._any(self.PER_TOKEN) + def has_fp8_current_scaling(self): + return self._any(self.FP8_1x128_128x128_CURRENT_SCALING) + + def has_fp8_quantize_weights_on_demand(self): + return self._any(self.FP8_QUANTIZE_WEIGHTS_ON_DEMAND) + def has_act_static_scaling(self): return not self.has_per_token_dynamic_scaling( ) and not self.has_fp8_rowwise() @@ -162,7 +173,8 @@ def has_any_quant(self): return self._any(self.INT4_WEIGHTS | self.INT8_WEIGHTS | self.ACTIVATIONS | self.INT8_KV_CACHE | self.FP8_KV_CACHE - | self.FP8_QDQ | self.FP8_ROWWISE) + | self.FP8_QDQ | self.FP8_ROWWISE + | self.FP8_1x128_128x128_CURRENT_SCALING) def set_int8_kv_cache(self): return self | self.INT8_KV_CACHE @@ -186,6 +198,8 @@ def from_description(quantize_weights=False, use_int8_kv_cache=False, use_fp8_kv_cache=False, use_fp8_qdq=False, + use_fp8_current_scaling=False, + use_fp8_quantize_weight_on_demand=False, use_fp8_rowwise=False, use_w4a8_qserve=False): @@ -196,12 +210,14 @@ def raise_error(): f"{per_token=}, " f"{per_channel=}, " f"{per_group=}, " - f"{use_int4_weights=}" - f"{use_int8_kv_cache=}" - f"{use_fp8_kv_cache=}" - f"{use_fp8_qdq=}" - f"{use_fp8_rowwise=}" - f"{use_w4a8_qserve=}") + f"{use_int4_weights=}, " + f"{use_int8_kv_cache=}, " + f"{use_fp8_kv_cache=}, " + f"{use_fp8_qdq=}, " + f"{use_fp8_current_scaling=}, " + f"{use_fp8_quantize_weight_on_demand=}," + f"{use_fp8_rowwise=}, " + f"{use_w4a8_qserve=}, ") # We must quantize weights when we quantize activations. if quantize_activations and not quantize_weights: @@ -246,6 +262,12 @@ def raise_error(): if use_fp8_rowwise: mode = mode | QuantMode.FP8_ROWWISE | QuantMode.PER_TOKEN | QuantMode.PER_CHANNEL + if use_fp8_current_scaling: + mode = mode | QuantMode.FP8_1x128_128x128_CURRENT_SCALING + + if use_fp8_quantize_weight_on_demand: + mode = mode | QuantMode.FP8_QUANTIZE_WEIGHTS_ON_DEMAND + # W4A8 QServe if use_w4a8_qserve: mode = mode | QuantMode.W4A8_QSERVE @@ -319,6 +341,10 @@ def from_quant_algo( quant_mode = QuantMode.from_description(use_fp8_qdq=True) elif quant_algo == QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN: quant_mode = QuantMode.from_description(use_fp8_rowwise=True) + elif quant_algo == QuantAlgo.FP8_CURRENT_SCALING: + quant_mode = QuantMode.from_description( + use_fp8_current_scaling=True, + use_fp8_quantize_weight_on_demand=False) else: quant_mode = QuantMode(0) @@ -345,6 +371,8 @@ def to_dict(self): self.has_fp8_qdq(), 'enable_fp8_rowwise': self.has_fp8_rowwise(), + 'enable_fp8_current_scaling': + self.has_fp8_current_scaling(), 'fp8_kv_cache': self.has_fp8_kv_cache(), 'use_weight_only': diff --git a/tensorrt_llm/quantization/quantize.py b/tensorrt_llm/quantization/quantize.py index 17efc23a9b00..ba7636867f10 100644 --- a/tensorrt_llm/quantization/quantize.py +++ b/tensorrt_llm/quantization/quantize.py @@ -9,16 +9,20 @@ from ..layers.moe import MixtureOfExperts from ..models.modeling_utils import LayerQuantConfig, QuantConfig from ..parameter import Parameter -from .layers import (FP8Linear, FP8RowLinear, Fp8RowwiseAttention, - Fp8RowwiseGatedMLP, Fp8RowwiseMLP, Fp8RowwiseRmsNorm, - Int8SmoothQuantLinear, Int8SmoothQuantRowLinear, - QServeAttention, QServeGatedMLP, QServeMLP, QServeRmsNorm, - SmoothQuantAttention, SmoothQuantGatedMLP, - SmoothQuantLayerNorm, SmoothQuantMLP, SmoothQuantRmsNorm, - WeightOnlyGroupwiseQuantColumnLinear, - WeightOnlyGroupwiseQuantRowLinear, - WeightOnlyQuantColumnLinear, WeightOnlyQuantEmbedding, - WeightOnlyQuantRowLinear) + +# isort: off +from .layers import ( + FP8CurrentScalingLinear, FP8CurrentScalingRowLinear, FP8Linear, + FP8RowLinear, FP8RowLinear, Fp8RowwiseAttention, Fp8RowwiseGatedMLP, + Fp8RowwiseMLP, Fp8RowwiseRmsNorm, Int8SmoothQuantLinear, + Int8SmoothQuantRowLinear, QServeAttention, QServeGatedMLP, QServeMLP, + QServeRmsNorm, SmoothQuantAttention, SmoothQuantGatedMLP, + SmoothQuantLayerNorm, SmoothQuantMLP, SmoothQuantRmsNorm, + WeightOnlyGroupwiseQuantColumnLinear, WeightOnlyGroupwiseQuantRowLinear, + WeightOnlyQuantColumnLinear, WeightOnlyQuantEmbedding, + WeightOnlyQuantRowLinear) +# isort: on + from .mode import W8A8_SQ_PLUGIN_LIST, QuantAlgo, QuantMode @@ -35,7 +39,6 @@ def quantize_layers( '*position_embedding', '*block_embedding', '*shared_expert_gate', - '*fused_a', ] for name, module, parent in model.named_modules_with_parent(): @@ -71,8 +74,12 @@ def quantize_layers( if isinstance(module, ColumnLinear): init_params[ "out_features"] = module.out_features * module.tp_size + if quant_config.quant_mode.has_fp8_quantize_weights_on_demand(): + init_params["quantize_weights_on_demand"] = True elif isinstance(module, RowLinear): init_params["in_features"] = module.in_features * module.tp_size + if quant_config.quant_mode.has_fp8_quantize_weights_on_demand(): + init_params["quantize_weights_on_demand"] = True if preprocess_init_params is not None: preprocess_init_params(init_params, name, module) quant_layer = quant_cls(**init_params) @@ -227,6 +234,22 @@ def fp8_quantize(model, quant_config: QuantConfig): return model +def fp8_current_scaling_quantize(model, quant_config: QuantConfig): + assert quant_config.quant_mode.has_fp8_current_scaling() + + quant_map = { + ColumnLinear: FP8CurrentScalingLinear, + RowLinear: FP8CurrentScalingRowLinear, + } + + model = quantize_layers( + model, + quant_config, + quant_map, + ) + return model + + def fp8_rowwise_quantize(model, quant_config: QuantConfig): assert quant_config.quant_mode.has_fp8_rowwise() @@ -542,6 +565,8 @@ def quantize(model, quant_config: Union[QuantConfig, LayerQuantConfig]): if layer_quant_mode.has_fp8_qdq(): module = fp8_quantize(module, layer_quant_cfg) + elif layer_quant_mode.has_fp8_current_scaling(): + module = fp8_current_scaling_quantize(module, layer_quant_cfg) elif layer_quant_mode.has_fp8_rowwise(): module = fp8_rowwise_quantize(module, layer_quant_cfg) elif layer_quant_mode.is_qserve_w4a8(): diff --git a/tests/attention/test_deepseek_v2_attention.py b/tests/attention/test_deepseek_v2_attention.py index 7fc723ca9d4d..56cd2c54ec64 100644 --- a/tests/attention/test_deepseek_v2_attention.py +++ b/tests/attention/test_deepseek_v2_attention.py @@ -1138,7 +1138,7 @@ def _construct_execution( # position_ids=position_ids_tensor, # q_a_proj=q_a_proj_tensor, # q_a_layernorm=q_a_layernorm_tensor, - fused_q_proj=kv_a_proj_with_mqa_tensor, + k_b_proj_trans=kv_a_proj_with_mqa_tensor, q_b_proj=q_b_proj_tensor, # kv_a_proj_with_mqa=kv_a_proj_with_mqa_tensor, # kv_a_layernorm=kv_a_layernorm_tensor, From 8539574f450a68edfbed8907e8b39c94692d9cea Mon Sep 17 00:00:00 2001 From: yingcanw Date: Fri, 24 Jan 2025 10:00:05 -0800 Subject: [PATCH 12/12] Deepseek V3 FP8 Support --- .../libfp8_blockscale_gemm.a | 4 +- examples/deepseek_v3/README.md | 90 ++++++++++++++++++- 2 files changed, 90 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a index cd6cfc9597a1..2e99d3ca4d36 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:dc1668c231ab9333eee605a224197f6b7e5a36ffd516a19b33f4259a8a35eac4 -size 2123214 +oid sha256:3130be05e0b4f136a40b5ac2726ff90994f7abba0206ca205fc9345a273c34b7 +size 2123196 diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index 6ea8963c70ab..c8b64a86a998 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -24,7 +24,6 @@ This document shows how to build and run the `DeepSeek-v3` model in TensorRT-LLM - TP: Tensor Parallel - EP: Expert Parallel - IB: Inflight Batching -- FP8: Support for FP8 is currently in progress and will be released soon ***Please Note:*** - Prefer using BF16 over FP16 for DeepSeek-V3 since model original training precision is FP8 and we found direct convert FP8 -> FP16 may cause unknown accuracy issues. @@ -48,10 +47,13 @@ cp /path/to/DeepSeek-V3/config.json /path/to/DeepSeek-V3/configuration_deepseek. ## Hardware -The DeepSeek-V3 model requires at least 32x80G GPU memory, model contains 660B parameters, roughly 1.3TB memory (with BF16 precision). +The DeepSeek-V3 model requires at least 16x80G GPU memory, model contains 660B parameters, roughly 1.3TB memory (with BF16 precision). ***Caution: Current TRT-LLM MLA kernel only supports Hopper architecture (SM90). Ampere architecture (SM80 & SM86) will be supported in the future release.*** +Please follow the instructions [here](https://github.com/NVIDIA/TensorRT-LLM/blob/deepseek/docs/source/installation/build-from-source-linux.md#building-a-tensorrt-llm-docker-image +) to achieve a correct docker image. + ## Overview The TensorRT-LLM DeepSeek-V3 implementation can be found in [tensorrt_llm/models/deepseek_v2/model.py](../../tensorrt_llm/models/deepseek_v2/model.py). The TensorRT-LLM Deepseek-V3 example code is located in [`example/deepseek_v3`](./). There is one main file: @@ -62,6 +64,8 @@ In addition, there are three shared files in the parent folder [`examples`](../) * [`../run.py`](../run.py) to run the model inference output by giving an input text. +* [`../mmlu.py`](../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT-LLM model on the MMLU dataset. + ## Usage @@ -171,3 +175,85 @@ and the output will be like: Input [Text 0]: "Today is a nice day." Output [Text 0 Beam 0]: " I am going to the park with my friends. We are going to play soccer. We are going" ``` + +At last, we can evaluate the model with [mmlu.py](../mmlu.py) script: + + +```bash +# Download MMLU dataset +mkdir mmlu_data && cd mmlu_data +wget https://people.eecs.berkeley.edu/~hendrycks/data.tar && tar -xf data.tar +# Run MMLU evaluation +python3 mmlu.py \ + --hf_model_dir ${MODEL_DIR} \ + --engine_dir ./trtllm_engines/deepseek_v3/fp8/tp16-sel4096-isl2048-bs4 \ + --data_dir mmlu_data \ + --test_trt_llm 2>&1 | tee ${ENGINE_DIR}/test_with_mmlu.log +``` + +and the output will be like: + +``` +Average accuracy 0.926 - high_school_macroeconomics +Average accuracy 0.752 - high_school_mathematics +Average accuracy 0.954 - high_school_microeconomics +Average accuracy 0.848 - high_school_physics +Average accuracy 0.967 - high_school_psychology +Average accuracy 0.861 - high_school_statistics +Average accuracy 0.956 - high_school_us_history +Average accuracy 0.954 - high_school_world_history +Average accuracy 0.861 - human_aging +Average accuracy 0.931 - human_sexuality +Average accuracy 0.975 - international_law +Average accuracy 0.907 - jurisprudence +Average accuracy 0.920 - logical_fallacies +Average accuracy 0.848 - machine_learning +Average accuracy 0.951 - management +Average accuracy 0.957 - marketing +Average accuracy 0.950 - medical_genetics +Average accuracy 0.957 - miscellaneous +Average accuracy 0.870 - moral_disputes +Average accuracy 0.798 - moral_scenarios +Average accuracy 0.918 - nutrition +Average accuracy 0.916 - philosophy +Average accuracy 0.932 - prehistory +Average accuracy 0.869 - professional_accounting +Average accuracy 0.714 - professional_law +Average accuracy 0.956 - professional_medicine +Average accuracy 0.908 - professional_psychology +Average accuracy 0.800 - public_relations +Average accuracy 0.869 - security_studies +Average accuracy 0.960 - sociology +Average accuracy 0.950 - us_foreign_policy +Average accuracy 0.578 - virology +Average accuracy 0.930 - world_religions +Average accuracy 0.852 - math +Average accuracy 0.874 - health +Average accuracy 0.905 - physics +Average accuracy 0.936 - business +Average accuracy 0.958 - biology +Average accuracy 0.825 - chemistry +Average accuracy 0.888 - computer science +Average accuracy 0.912 - economics +Average accuracy 0.890 - engineering +Average accuracy 0.851 - philosophy +Average accuracy 0.917 - other +Average accuracy 0.932 - history +Average accuracy 0.944 - geography +Average accuracy 0.904 - politics +Average accuracy 0.936 - psychology +Average accuracy 0.949 - culture +Average accuracy 0.744 - law +Average accuracy 0.883 - STEM +Average accuracy 0.827 - humanities +Average accuracy 0.926 - social sciences +Average accuracy 0.898 - other (business, health, misc.) +Average accuracy: 0.877 +``` + +**Known Issue** + +1. The memory allocation for MoE is too large. + +This issue prevents running larger batch sizes and long sequence inputs. We will optimize and fix this issue soon. +