wehub-resource-sync
ec436095dd
Book-CI / test (macos-latest) (push) Has been cancelled
Book-CI / test (ubuntu-latest) (push) Has been cancelled
Book-CI / test (windows-latest) (push) Has been cancelled
Release Fake Tag / publish (push) Has been cancelled
Deploy / deploy (macos-latest) (push) Has been cancelled
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Deploy / deploy (windows-latest) (push) Has been cancelled
Release to PyPI / Build & publish sglang-kt (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.11) (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.12) (push) Has been cancelled
Release to PyPI / Publish kt-kernel to PyPI (push) Has been cancelled
814 lines
37 KiB
Python
814 lines
37 KiB
Python
#!/usr/bin/env python
|
|
# coding=utf-8
|
|
'''
|
|
Description :
|
|
Author : Boxin Zhang, Azure-Tang
|
|
Version : 0.1.0
|
|
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
|
'''
|
|
import re
|
|
import sys
|
|
import threading
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch import nn
|
|
import itertools
|
|
import time
|
|
import enum
|
|
from transformers import (
|
|
LogitsProcessorList,
|
|
TemperatureLogitsWarper,
|
|
TopKLogitsWarper,
|
|
TopPLogitsWarper,
|
|
MinPLogitsWarper,
|
|
TypicalLogitsWarper,
|
|
EpsilonLogitsWarper,
|
|
EtaLogitsWarper,
|
|
)
|
|
|
|
from ktransformers.util.custom_loader import ModelLoaderFactory, ModelLoader, SafeTensorLoader, translate_name_to_gguf
|
|
from ktransformers.operators import base_operator
|
|
from ktransformers.models.custom_cache import StaticCache
|
|
from ktransformers.util.cuda_graph_runner import CUDAGraphRunner
|
|
from ktransformers.util.textstream import TextStreamer
|
|
if not torch.xpu.is_available():
|
|
from ktransformers.operators.flashinfer_wrapper import MLAWrapperSingleton
|
|
# from ktransformers.operators.flashinfer_wrapper import MLAWrapperSingleton
|
|
import socket
|
|
|
|
warm_uped = False
|
|
CUR_DEVICE = None
|
|
W8A8_ENABLE = False
|
|
Q4_GGUF_LODER = None
|
|
_USE_NPU_GRAPH = False
|
|
_MAX_DECODE_PROFILE = 1
|
|
WARM_UP_SKIP_CNT = [1, 1]
|
|
_SPECULATE_STEP = 1
|
|
|
|
try:
|
|
import torch_npu
|
|
use_torch_npu = torch_npu.npu.is_available()
|
|
from ktransformers.util.ascend.ascend_utils import get_tensor_parallel_size
|
|
except:
|
|
use_torch_npu = False
|
|
|
|
def get_use_npu_graph():
|
|
assert _USE_NPU_GRAPH is not None, "use npu graph is not setting"
|
|
return _USE_NPU_GRAPH
|
|
|
|
from enum import StrEnum
|
|
|
|
class StatKey(StrEnum):
|
|
Embedding = "Embedding"
|
|
GraphCapture = "GraphCapture"
|
|
GraphReplay = "GraphReplay"
|
|
ExpertsForward1 = "ExpertsForward1"
|
|
ExpertsForward2 = "ExpertsForward2"
|
|
CPUExperts = "CPUExperts"
|
|
GraphDestroy = "GraphDestroy"
|
|
DecodeOneTokenPost = "DecodeOneTokenPost"
|
|
DecodeOneToken = "DecodeOneToken"
|
|
GraphInit = "GraphInit"
|
|
|
|
class TimeStat:
|
|
def __init__(self):
|
|
# open_status = os.environ["KT_PERF_STAT"] if "KT_PERF_STAT" in os.environ else "0"
|
|
# if open_status == "0":
|
|
# self.on = False
|
|
# else:
|
|
# self.on = True
|
|
self.on = True
|
|
self.prefill_stats = dict()
|
|
self.decode_stats = dict()
|
|
for key in StatKey:
|
|
self.prefill_stats[key] = StatItem()
|
|
self.decode_stats[key] = StatItem()
|
|
self.reset_all()
|
|
|
|
def record_start_time(self):
|
|
start_time = time.time_ns()
|
|
return start_time
|
|
|
|
def add_time_stat(self, key: StatKey, time_ns, is_prefill):
|
|
if not key:
|
|
return
|
|
# torch.cuda.synchronize()
|
|
cost = time.time_ns() - time_ns
|
|
if is_prefill:
|
|
item = self.prefill_stats[key]
|
|
else:
|
|
item = self.decode_stats[key]
|
|
item.add_item(cost)
|
|
|
|
def print_all(self):
|
|
# rank = f"[rank:{torch.distributed.get_rank()}]"
|
|
rank = f"[rank:0]"
|
|
msg = f"\n{rank} Prefill Time Stat\n"
|
|
msg += rank + " {:27}{:>15}{:>15}{:>15}{:>15}{:>15}\n".format("", "min(ms)", "max(ms)", "avg(ms)", "count", "total(ms)")
|
|
for key, value in self.prefill_stats.items():
|
|
msg += rank + f" {key.value:<25}:{value.get_stat()}\n"
|
|
msg += f"\n{rank} Decode Time Stat\n"
|
|
msg += rank + " {:27}{:>15}{:>15}{:>15}{:>15}{:>15}\n".format("", "min(ms)", "max(ms)", "avg(ms)", "count", "total(ms)")
|
|
for key, value in self.decode_stats.items():
|
|
msg += rank + f" {key.value:<25}:{value.get_stat()}\n"
|
|
print(msg)
|
|
|
|
def reset_all(self):
|
|
for _, value in self.prefill_stats.items():
|
|
value.reset()
|
|
for _, value in self.decode_stats.items():
|
|
value.reset()
|
|
|
|
|
|
class StatItem:
|
|
def __init__(self):
|
|
self.min_time = 100000000
|
|
self.max_time = 0
|
|
self.total_time_ns = 0
|
|
self.count = 0
|
|
|
|
def add_item(self, cost_time_ns):
|
|
self.count += 1
|
|
self.total_time_ns += cost_time_ns
|
|
self.min_time = min(self.min_time, cost_time_ns)
|
|
self.max_time = max(self.max_time, cost_time_ns)
|
|
|
|
def reset(self):
|
|
self.min_time = 100000000
|
|
self.max_time = 0
|
|
self.total_time_ns = 0
|
|
self.count = 0
|
|
|
|
def get_stat(self):
|
|
min_time = self.min_time / 1000 / 1000
|
|
max_time = self.max_time / 1000 / 1000
|
|
if self.count != 0:
|
|
avg_time = self.total_time_ns / self.count / 1000 / 1000
|
|
else:
|
|
avg_time = 0
|
|
total = self.total_time_ns / 1000 / 1000
|
|
return f"{min_time:15.2f}{max_time:15.2f}{avg_time:15.2f}{self.count:15}{total:15.2f}"
|
|
|
|
|
|
timeStat = TimeStat()
|
|
|
|
|
|
def get_free_ports(n: int, continue_prot: list):
|
|
sockets = []
|
|
ports = []
|
|
for _ in range(n):
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.bind(("", 0))
|
|
port = s.getsockname()[1]
|
|
if port in continue_prot:
|
|
s.close()
|
|
continue
|
|
ports.append(port)
|
|
sockets.append(s)
|
|
for s in sockets:
|
|
s.close()
|
|
return ports
|
|
|
|
def get_current_device():
|
|
if use_torch_npu:
|
|
return f"npu:{torch.npu.current_device()}"
|
|
else:
|
|
return f"cuda:{torch.npu.current_device()}"
|
|
|
|
def get_compute_capability(device:torch.device = None):
|
|
if use_torch_npu:
|
|
return 0
|
|
if torch.cuda.is_available():
|
|
if device is None:
|
|
num_gpus = torch.cuda.device_count()
|
|
min_compute_capability_major = 100
|
|
for gpu_id in range(num_gpus):
|
|
gpu_props = torch.cuda.get_device_properties(gpu_id)
|
|
min_compute_capability_major = min(min_compute_capability_major, gpu_props.major)
|
|
return min_compute_capability_major
|
|
else:
|
|
return torch.cuda.get_device_properties(device)
|
|
|
|
def set_module(model, submodule_key, module):
|
|
tokens = submodule_key.split('.')
|
|
sub_tokens = tokens[:-1]
|
|
cur_mod = model
|
|
for s in sub_tokens:
|
|
if hasattr(cur_mod, s):
|
|
cur_mod = getattr(cur_mod, s)
|
|
else: # nn.ModuleList or nn.ModuleList
|
|
cur_mod=cur_mod[int(s)]
|
|
if hasattr(cur_mod, tokens[-1]):
|
|
setattr(cur_mod, tokens[-1], module)
|
|
else: # nn.ModuleList or nn.ModuleList
|
|
cur_mod[int(tokens[-1])] = module
|
|
|
|
def set_param(module: nn.Module, name: str, weights: torch.Tensor):
|
|
|
|
param=nn.parameter.Parameter(weights, requires_grad=False)
|
|
if isinstance(module, nn.Linear) and len(weights.shape)==1:
|
|
param.unsqueeze_(0)
|
|
setattr(module, name, param)
|
|
|
|
def get_device(gguf_module_key:str, device_map:dict):
|
|
if gguf_module_key in device_map:
|
|
return device_map[gguf_module_key]["generate_device"]
|
|
else:
|
|
return "cuda"
|
|
|
|
def get_all_used_cuda_device(device_map:dict):
|
|
all_device_list = set()
|
|
for key in device_map:
|
|
all_device_list.add(device_map[key]["generate_device"]) if "generate_device" in device_map[key] else None
|
|
all_device_list.add(device_map[key]["prefill_device"]) if "prefill_device" in device_map[key] else None
|
|
if "cpu" in all_device_list:
|
|
all_device_list.remove("cpu")
|
|
if use_torch_npu:
|
|
all_device_list = set([device.replace('cuda', 'npu') for device in all_device_list])
|
|
all_device_list = list(all_device_list)
|
|
return all_device_list
|
|
|
|
def load_cur_state_dict_npu(module: nn.Module, gguf_loader: ModelLoader, prefix: str = "", device="npu"):
|
|
prefix = prefix.replace("orig_module.", "")
|
|
persistent_buffers = {k: v for k, v in module._buffers.items() if k not in module._non_persistent_buffers_set}
|
|
local_name_params = itertools.chain(module._parameters.items(), persistent_buffers.items())
|
|
local_state = {k: v for k, v in local_name_params if v is not None}
|
|
for name, param in local_state.items():
|
|
key = prefix + name
|
|
translated_key = translate_name_to_gguf(key)
|
|
# TODO: Merge all loader.
|
|
# I know this is ugly but lets do it for now.
|
|
if gguf_loader.safetensor_loader is not None:
|
|
load_dequantized_tensor = gguf_loader.safetensor_loader.load_dequantized_tensor
|
|
tensor_file_map = gguf_loader.safetensor_loader.tensor_file_map
|
|
else:
|
|
load_dequantized_tensor = gguf_loader.load_gguf_tensor
|
|
tensor_file_map = gguf_loader.tensor_file_map
|
|
|
|
if translated_key in tensor_file_map:
|
|
target_dtype = torch.get_default_dtype()
|
|
device = get_device(translated_key[:translated_key.rfind(".")], gguf_loader.tensor_device_map)
|
|
# Todo need fix
|
|
device = "cpu" if "embd" in translated_key else get_current_device()
|
|
print(f"loading layer {translated_key} to {device}")
|
|
torch.cuda.empty_cache()
|
|
weights = load_dequantized_tensor(translated_key, device=device).to(dtype=target_dtype)
|
|
set_param(module, name, weights)
|
|
del weights
|
|
else:
|
|
#print(load_config.tensor_file_map.keys())
|
|
raise Exception(f"can't find {translated_key} in GGUF file!")
|
|
|
|
def load_cur_state_dict(module: nn.Module, gguf_loader: ModelLoader, prefix: str = "", device="cuda"):
|
|
if use_torch_npu:
|
|
load_cur_state_dict_npu(module, gguf_loader, prefix, device)
|
|
return
|
|
|
|
prefix = prefix.replace("orig_module.", "")
|
|
persistent_buffers = {k: v for k, v in module._buffers.items() if k not in module._non_persistent_buffers_set}
|
|
local_name_params = itertools.chain(module._parameters.items(), persistent_buffers.items())
|
|
local_state = {k: v for k, v in local_name_params if v is not None}
|
|
for name, param in local_state.items():
|
|
key = prefix + name
|
|
translated_key = key
|
|
|
|
# TODO: Merge all loader.
|
|
# I know this is ugly but lets do it for now.
|
|
if isinstance(gguf_loader, SafeTensorLoader):
|
|
load_dequantized_tensor = gguf_loader.load_dequantized_tensor
|
|
else:
|
|
load_dequantized_tensor = gguf_loader.load_gguf_tensor
|
|
tensor_file_map = gguf_loader.tensor_file_map
|
|
|
|
if gguf_loader.has_tensor(translated_key) or "kv_b_proj" in translated_key:
|
|
target_dtype = torch.get_default_dtype()
|
|
device = get_device(translated_key[:translated_key.rfind(".")], gguf_loader.tensor_device_map)
|
|
print(f"loading {translated_key} to {device}")
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
elif torch.xpu.is_available():
|
|
torch.xpu.empty_cache()
|
|
if "kv_b_proj" in translated_key and not gguf_loader.has_tensor(translated_key):
|
|
attn_k_b = load_dequantized_tensor(translated_key.replace("self_attn.kv_b_proj", "attn_k_b"), device=device).to(dtype=target_dtype)
|
|
attn_k_b = attn_k_b.transpose(1, 2).contiguous()
|
|
attn_v_b = load_dequantized_tensor(translated_key.replace("self_attn.kv_b_proj", "attn_v_b"), device=device).to(dtype=target_dtype)
|
|
kv_b_proj = torch.cat((attn_k_b, attn_v_b), dim=1)
|
|
kv_b_proj = kv_b_proj.contiguous() if kv_b_proj.ndim == 2 else kv_b_proj.flatten(0, 1).contiguous()
|
|
set_param(module, name, kv_b_proj)
|
|
del attn_k_b
|
|
del attn_v_b
|
|
else:
|
|
weights = load_dequantized_tensor(translated_key, device=device).to(dtype=target_dtype)
|
|
set_param(module, name, weights)
|
|
del weights
|
|
else:
|
|
#print(load_config.tensor_file_map.keys())
|
|
raise Exception(f"can't find {translated_key} in GGUF file!")
|
|
|
|
|
|
def sync_all_device(all_device_list):
|
|
for device in all_device_list:
|
|
if "cuda" in device.lower():
|
|
torch.cuda.synchronize(device)
|
|
elif "xpu" in device.lower():
|
|
torch.xpu.synchronize(device)
|
|
elif use_torch_npu:
|
|
torch_npu.synchronize(device)
|
|
else:
|
|
raise RuntimeError("The device {} is not available".format(device))
|
|
|
|
torch_device_mapping ={"cuda": "cuda:0", "xpu": "xpu:0"}
|
|
|
|
def xpu_fp16_model(config):
|
|
# This function is to check if we run this model on XPU with FP16 dtype
|
|
if not torch.xpu.is_available():
|
|
return False
|
|
if config.architectures[0] == "DeepseekV3ForCausalLM":
|
|
return True
|
|
if config.architectures[0] == "Qwen3MoeForCausalLM" and config.hidden_size == 4096:
|
|
# Qwen3-30B seems have precision issue with FP16
|
|
# so we only use FP16 for Qwen3-235B now
|
|
return True
|
|
return False
|
|
|
|
def load_weights(module:nn.Module, gguf_loader:ModelLoader, prefix='', device="cuda"):
|
|
#print(f"recursively loading weights {prefix}")
|
|
if not isinstance(module, base_operator.BaseInjectedModule):
|
|
load_cur_state_dict(module, gguf_loader, prefix, device=device)
|
|
for name, child in module._modules.items():
|
|
load_weights(child, gguf_loader, prefix+name+".", device=device)
|
|
else:
|
|
module.load()
|
|
|
|
def tf_logits_warper(generation_config):
|
|
"""
|
|
This class returns a [`LogitsProcessorList`] list object that contains all relevant [`LogitsWarper`] instances
|
|
used for multinomial sampling.
|
|
"""
|
|
|
|
# instantiate warpers list
|
|
warpers = LogitsProcessorList()
|
|
|
|
# In beam methods, we need to keep at least one non-eos token to explore continuations that might have a
|
|
# better score (i.e. keep len(list(generation_config._eos_token_tensor)) + 1)
|
|
if generation_config.num_beams > 1:
|
|
if isinstance(generation_config._eos_token_tensor, list):
|
|
min_tokens_to_keep = len(generation_config._eos_token_tensor) + 1
|
|
elif isinstance(generation_config._eos_token_tensor, torch.Tensor):
|
|
min_tokens_to_keep = generation_config._eos_token_tensor.shape[0] + 1
|
|
else:
|
|
min_tokens_to_keep = 2
|
|
else:
|
|
min_tokens_to_keep = 1
|
|
|
|
# the following idea is largely copied from this PR: https://github.com/huggingface/transformers/pull/5420/files
|
|
# all samplers can be found in `generation_utils_samplers.py`
|
|
if generation_config.temperature is not None and generation_config.temperature != 1.0:
|
|
warpers.append(TemperatureLogitsWarper(generation_config.temperature))
|
|
if generation_config.top_k is not None and generation_config.top_k != 0:
|
|
warpers.append(TopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=min_tokens_to_keep))
|
|
if generation_config.top_p is not None and generation_config.top_p < 1.0:
|
|
warpers.append(TopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=min_tokens_to_keep))
|
|
if generation_config.min_p is not None:
|
|
# Applied after temperature scaling (see https://github.com/ggerganov/llama.cpp/pull/3841#issuecomment-2073826084)
|
|
warpers.append(MinPLogitsWarper(min_p=generation_config.min_p, min_tokens_to_keep=min_tokens_to_keep))
|
|
if generation_config.typical_p is not None and generation_config.typical_p < 1.0:
|
|
warpers.append(
|
|
TypicalLogitsWarper(mass=generation_config.typical_p, min_tokens_to_keep=min_tokens_to_keep)
|
|
)
|
|
if generation_config.epsilon_cutoff is not None and 0.0 < generation_config.epsilon_cutoff < 1.0:
|
|
warpers.append(
|
|
EpsilonLogitsWarper(epsilon=generation_config.epsilon_cutoff, min_tokens_to_keep=min_tokens_to_keep)
|
|
)
|
|
if generation_config.eta_cutoff is not None and 0.0 < generation_config.eta_cutoff < 1.0:
|
|
warpers.append(
|
|
EtaLogitsWarper(
|
|
epsilon=generation_config.eta_cutoff, min_tokens_to_keep=min_tokens_to_keep, device=device
|
|
)
|
|
)
|
|
# `LogitNormalization` should always be the last logit processor, when present
|
|
if generation_config.renormalize_logits is True:
|
|
warpers.append(LogitNormalization())
|
|
return warpers
|
|
def prefill_and_generate(model, tokenizer, inputs, max_new_tokens=10000, use_cuda_graph: bool = True,
|
|
mode = 'normal', force_think: bool = False, chunk_size = 16384, use_flashinfer_mla = False,
|
|
num_heads = None, head_dim_ckv = None, head_dim_kpe = None, q_head_dim = None,
|
|
static_cache = None, draft_model=None, draft_cache=None):
|
|
import os
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
torch._dynamo.config.suppress_errors = True
|
|
batch_size, seq_length = inputs.shape
|
|
device_map = model.gguf_loader.tensor_device_map
|
|
if use_torch_npu:
|
|
CUR_DEVICE = f"npu:{torch.npu.current_device()}"
|
|
vocabulary_size = model.config.vocab_size
|
|
topp = torch.tensor([[model.generation_config.top_p]], dtype=torch.float16).npu()
|
|
topk = torch.tensor([[model.generation_config.top_k]], dtype=torch.int32).npu()
|
|
temperature = torch.tensor([[model.generation_config.temperature]], dtype=torch.float16).npu()
|
|
next_token_fake = torch.tensor([[1]], dtype=torch.int32).npu()
|
|
next_token_probs = torch.tensor([[1.0]], dtype=torch.float16).npu()
|
|
torch_device = torch.npu.current_device()
|
|
else:
|
|
torch_device = get_device('model.layers.0.self_attn', device_map)
|
|
torch_device = torch_device_mapping[torch_device] if torch_device in torch_device_mapping else torch_device
|
|
inputs = inputs.to(torch_device)
|
|
all_cuda_device = get_all_used_cuda_device(device_map)
|
|
|
|
tokens = []
|
|
|
|
def decode_one_tokens_npu(cuda_graph_runner, cur_token, position_ids, cache_position, past_key_values, logits_warper, generation_config, use_cuda_graph: bool = True):
|
|
if cuda_graph_runner is None:
|
|
use_cuda_graph = False
|
|
|
|
inputs_embeds = model.model.embed_tokens(cur_token.to('cpu')).to(torch_device)
|
|
if use_cuda_graph:
|
|
if cuda_graph_runner.model_capture:
|
|
cuda_graph_runner.capture(model, cur_token, position_ids, cache_position, past_key_values, CUR_DEVICE, return_dict=False, use_cache=True)
|
|
cuda_graph_runner.model_capture = False
|
|
|
|
ret = cuda_graph_runner(inputs_embeds, position_ids, cache_position)
|
|
logits = ret[0]
|
|
next_token = torch.argmax(logits, dim=-1)
|
|
else:
|
|
torch_npu.npu.set_device(torch_device)
|
|
logits = model(inputs_embeds=inputs_embeds,
|
|
position_ids=position_ids,
|
|
cache_position=cache_position,
|
|
past_key_values=past_key_values,
|
|
return_dict=False, use_cache=True, is_prefill=False)[0]
|
|
if past_key_values != None:
|
|
past_key_values.change_seq_length(1)
|
|
|
|
if generation_config.do_sample:
|
|
logits = logits / temperature
|
|
torch.manual_seed(0)
|
|
probs = logits.view(batch_size, vocabulary_size)
|
|
sm = nn.Softmax(dim=-1)
|
|
probs = sm(probs).half().npu()
|
|
next_token = next_token_fake
|
|
torch_npu._npu_topk_topp_sampling(probs, topk, topp, next_token, next_token_probs)
|
|
next_token = next_token.squeeze(-1)
|
|
else:
|
|
next_token_scores = logits_warper(inputs, logits[:, -1, :])
|
|
next_token = torch.argmax(next_token_scores, dim=-1)
|
|
|
|
return next_token
|
|
|
|
|
|
def decode_one_tokens(cuda_graph_runner, cur_token, position_ids, cache_position, past_key_values, logits_warper, generation_config, use_cuda_graph: bool = True):
|
|
if use_torch_npu:
|
|
return decode_one_tokens_npu(cuda_graph_runner, cur_token, position_ids, cache_position, past_key_values, logits_warper, generation_config, use_cuda_graph)
|
|
if cuda_graph_runner is None:
|
|
use_cuda_graph = False
|
|
if use_cuda_graph:
|
|
logits = cuda_graph_runner(cur_token, position_ids, cache_position)
|
|
else:
|
|
# custom_stream = torch.cuda.Stream()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.set_device(torch_device)
|
|
elif torch.xpu.is_available():
|
|
torch.xpu.set_device(torch_device)
|
|
elif use_torch_npu:
|
|
torch_npu.set_device(torch_device)
|
|
else:
|
|
raise RuntimeError(f"The device: {torch_device} is not available")
|
|
inputs_embeds = model.model.embed_tokens(cur_token.to("cpu")).to(torch_device)
|
|
# with torch.cuda.stream(custom_stream):
|
|
logits=model(inputs_embeds=inputs_embeds,
|
|
position_ids=position_ids,
|
|
cache_position=cache_position,
|
|
past_key_values=past_key_values,
|
|
return_dict=False, use_cache=True)[0]
|
|
if past_key_values != None and isinstance(past_key_values, StaticCache):
|
|
past_key_values.change_seq_length(1)
|
|
sync_all_device(all_cuda_device)
|
|
next_token_scores = logits_warper(inputs, logits[:, -1, :])
|
|
if generation_config.do_sample:
|
|
probs = nn.functional.softmax(next_token_scores, dim=-1)
|
|
next_token = torch.multinomial(probs, num_samples=1).squeeze(1)
|
|
else:
|
|
next_token = torch.argmax(next_token_scores, dim=-1)
|
|
return next_token
|
|
|
|
# TODO: use CUDA Graph for chunk prefill, may get small improvement
|
|
def chunk_prefill(inputs, cache_position, past_key_values):
|
|
if mode == "long_context":
|
|
inputs_embeds = model.model.embed_tokens(inputs.to("cpu"))
|
|
else:
|
|
inputs_embeds = model.model.embed_tokens(inputs.to("cpu")).to(torch_device)
|
|
# inputs_embeds = torch_npu.npu_format_cast_(inputs_embeds, 29)
|
|
if use_flashinfer_mla:
|
|
MLAWrapperSingleton.update_buffer(past_key_values.max_pages)
|
|
MLAWrapperSingleton.need_plan_all()
|
|
|
|
ret = model(
|
|
inputs_embeds = inputs_embeds, cache_position=cache_position, past_key_values=past_key_values, return_dict=False, use_cache=True, is_prefill=True
|
|
)
|
|
logits = ret[0][:,-1,:].unsqueeze(0).clone().to(torch_device)
|
|
|
|
return logits
|
|
|
|
def decode_wrapper(next_token, position_ids, cache_position, cuda_graph_runner, past_key_values, inputs, seq_length, prof=None):
|
|
global warm_uped
|
|
global _USE_NPU_GRAPH
|
|
if use_cuda_graph:
|
|
from ktransformers.util.npu_graph_runner import get_or_create_runner
|
|
npu_graph_runner = get_or_create_runner(CUR_DEVICE)
|
|
npu_graph_runner.init(batch_size, seq_length)
|
|
|
|
with torch_npu.npu.stream(npu_graph_runner.main_stream):
|
|
gen_num_tokens = 1
|
|
while gen_num_tokens < max_new_tokens:
|
|
start_time = timeStat.record_start_time()
|
|
if use_flashinfer_mla:
|
|
MLAWrapperSingleton.plan_all(None,None,None,position_ids.squeeze(1)+1,None,
|
|
num_heads, head_dim_ckv, head_dim_kpe, past_key_values.page_size,
|
|
model.model.layers[0].self_attn.softmax_scale, torch.bfloat16, torch.bfloat16)
|
|
if gen_num_tokens == 1:
|
|
warm_uped = True
|
|
_USE_NPU_GRAPH = True
|
|
#np_graph_runner.capture(model, draft_model, next_token, torch.tensor(draft_token), position_ids, cache_position, past_key_values, draft_cache, torch_device, return_dict=False, use_cache=True)
|
|
cuda_graph_runner = npu_graph_runner
|
|
next_token = decode_one_tokens(cuda_graph_runner, next_token.unsqueeze(0), position_ids, cache_position, past_key_values, logits_warper, generation_config, use_cuda_graph)
|
|
next_token = next_token.to(torch_device)
|
|
inputs = torch.cat((inputs, next_token.unsqueeze(0)), dim=-1)
|
|
generated_ids[:, cache_position] = next_token.int()
|
|
tokens.append(int(next_token))
|
|
|
|
seq_length += 1
|
|
|
|
if next_token[0].item() == tokenizer.eos_token_id or tokenizer.decode(next_token.tolist()) == '<|im_end|>':
|
|
print(stream.end(), end="", flush=True)
|
|
break
|
|
else:
|
|
if torch.distributed.get_rank() % get_tensor_parallel_size() == 0:
|
|
print(stream.put(next_token.item()), end="", flush=True)
|
|
|
|
cache_position += 1
|
|
past_key_values.position[0] += 1
|
|
position_ids = cache_position.unsqueeze(0)
|
|
gen_num_tokens += 1
|
|
|
|
if prof is not None:
|
|
prof.step()
|
|
|
|
npu_graph_runner.destroy()
|
|
_USE_NPU_GRAPH = False
|
|
else:
|
|
gen_num_tokens = 1
|
|
while gen_num_tokens < max_new_tokens:
|
|
if use_flashinfer_mla:
|
|
MLAWrapperSingleton.plan_all(None,None,None,position_ids.squeeze(1)+1,None,
|
|
num_heads, head_dim_ckv, head_dim_kpe, past_key_values.page_size,
|
|
model.model.layers[0].self_attn.softmax_scale, torch.bfloat16, torch.bfloat16)
|
|
next_token = decode_one_tokens(cuda_graph_runner, next_token.unsqueeze(0), position_ids, cache_position, past_key_values, logits_warper, generation_config, use_cuda_graph)
|
|
next_token = next_token.to(torch_device)
|
|
inputs = torch.cat((inputs, next_token.unsqueeze(0)), dim=-1)
|
|
generated_ids[:, cache_position] = next_token.int()
|
|
tokens.append(int(next_token))
|
|
seq_length += 1
|
|
|
|
if next_token[0].item() == tokenizer.eos_token_id or tokenizer.decode(next_token.tolist()) == '<|im_end|>':
|
|
print(stream.end(), end="", flush=True)
|
|
break
|
|
else:
|
|
if torch.distributed.get_rank() % get_tensor_parallel_size() == 0:
|
|
print(stream.put(next_token.item()), end="", flush=True)
|
|
|
|
cache_position += 1
|
|
past_key_values.position[0] += 1
|
|
position_ids = cache_position.unsqueeze(0)
|
|
gen_num_tokens += 1
|
|
|
|
if prof is not None:
|
|
prof.step()
|
|
|
|
if prof is not None:
|
|
prof.stop()
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.set_device(torch_device)
|
|
elif torch.xpu.is_available():
|
|
torch.xpu.set_device(torch_device)
|
|
elif use_torch_npu:
|
|
torch_npu.set_device(torch_device)
|
|
else:
|
|
raise RuntimeError(f"The device: {torch_device} is not available")
|
|
|
|
with torch.no_grad():
|
|
|
|
stream = TextStreamer(tokenizer)
|
|
if torch.xpu.is_available():
|
|
from ipex_llm.transformers.kv import DynamicUnbalancedFp8Cache, DynamicNormalCache
|
|
if model.config.architectures[0] in ["DeepseekV3ForCausalLM", "DeepseekV2ForCausalLM"]:
|
|
past_key_values = DynamicUnbalancedFp8Cache.from_legacy_cache(None)
|
|
else:
|
|
past_key_values = DynamicNormalCache.from_legacy_cache(None)
|
|
elif use_torch_npu and static_cache:
|
|
assert isinstance(static_cache, StaticCache), '[ERROR] static_cache format not equal to StaticCache'
|
|
past_key_values = static_cache
|
|
if past_key_values.max_batch_size < batch_size or past_key_values.max_cache_len < seq_length + max_new_tokens:
|
|
print('[WARN] current staticCache size exceeded, try create new staticCache...')
|
|
past_key_values = StaticCache(
|
|
config=model.config, max_batch_size=1, max_cache_len=seq_length + max_new_tokens, device=device_map, dtype=model.dtype
|
|
)
|
|
else:
|
|
past_key_values.reset()
|
|
elif mode != 'long_context':
|
|
past_key_values = StaticCache(
|
|
config = model.config, max_batch_size = 1, max_cache_len = seq_length + max_new_tokens, device = device_map, dtype = model.dtype
|
|
)
|
|
else:
|
|
past_key_values = None
|
|
|
|
generation_config, model_kwargs = model._prepare_generation_config(
|
|
None, do_sample=False
|
|
# change this to modify generate config
|
|
#top_k=5, top_p=0.85, temperature=0.1
|
|
)
|
|
|
|
logits_warper = tf_logits_warper(generation_config)
|
|
|
|
cache_position = torch.arange(seq_length, device=torch_device, dtype=torch.int32)
|
|
if use_torch_npu:
|
|
past_key_values.position[0] = seq_length + 1
|
|
generated_ids = torch.zeros(
|
|
batch_size, seq_length + max_new_tokens + 1, dtype=torch.int, device=torch_device
|
|
)
|
|
generated_ids[:, cache_position] = inputs.to(torch_device).to(torch.int)
|
|
start_time = time.time()
|
|
logits = None
|
|
|
|
def prefill_wrapper(prof=None):
|
|
nonlocal logits
|
|
chunk_start = 0
|
|
while chunk_start < seq_length:
|
|
chunk_end = min(chunk_start + chunk_size, seq_length)
|
|
if past_key_values != None:
|
|
past_key_values.cur_idx=cache_position[chunk_start:chunk_end]
|
|
logits = chunk_prefill(inputs[:, chunk_start:chunk_end], cache_position[chunk_start:chunk_end], past_key_values)
|
|
chunk_start += chunk_size
|
|
if prof is not None:
|
|
prof.step()
|
|
if prof is not None:
|
|
prof.stop()
|
|
if logits is None:
|
|
raise ValueError('logits cannot be None')
|
|
|
|
if use_torch_npu:
|
|
global WARM_UP_SKIP_CNT
|
|
prof_prefill = os.environ["PROF_PREFILL"] if "PROF_PREFILL" in os.environ else "0"
|
|
if prof_prefill == "1" and WARM_UP_SKIP_CNT[0] <= 0:
|
|
experimental_config = torch_npu.profiler._ExperimentalConfig(
|
|
aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization,
|
|
profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False
|
|
)
|
|
with torch_npu.profiler.profile(
|
|
activities=[
|
|
torch_npu.profiler.ProfilerActivity.CPU,
|
|
torch_npu.profiler.ProfilerActivity.NPU
|
|
],
|
|
schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=8, repeat=1, skip_first=0),
|
|
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./prefill_prof"),
|
|
record_shapes=True,
|
|
profile_memory=True,
|
|
with_stack=False,
|
|
with_flops=False,
|
|
with_modules=False,
|
|
experimental_config=experimental_config) as prof:
|
|
prefill_wrapper(prof)
|
|
else:
|
|
prefill_wrapper()
|
|
WARM_UP_SKIP_CNT[0] -= 1
|
|
else:
|
|
|
|
chunk_start = 0
|
|
while chunk_start < seq_length:
|
|
chunk_end = min(chunk_start + chunk_size, seq_length)
|
|
if past_key_values != None:
|
|
past_key_values.cur_idx=cache_position[chunk_start:chunk_end]
|
|
logits = chunk_prefill(inputs[:, chunk_start:chunk_end], cache_position[chunk_start:chunk_end], past_key_values)
|
|
chunk_start += chunk_size
|
|
|
|
next_token_scores = logits_warper(inputs, logits[:, -1, :])
|
|
if generation_config.do_sample:
|
|
probs = nn.functional.softmax(next_token_scores, dim=-1)
|
|
next_token = torch.multinomial(probs, num_samples=1).squeeze(1)
|
|
else:
|
|
next_token = torch.argmax(next_token_scores, dim=-1)
|
|
|
|
first_token_time = time.time() - start_time
|
|
|
|
# print(f"------------------------------------- prefill next_token {next_token} draft_token {draft_token} ")
|
|
if use_flashinfer_mla:
|
|
MLAWrapperSingleton.reset_buffer()
|
|
|
|
prefill_count = seq_length
|
|
prefill_time = first_token_time
|
|
if use_torch_npu and torch.distributed.get_rank() % get_tensor_parallel_size() == 0:
|
|
if force_think:
|
|
print("<think>")
|
|
print(stream.put(next_token.item()), end="", flush=True)
|
|
elif not use_torch_npu:
|
|
if force_think:
|
|
print("<think>")
|
|
print(stream.put(next_token.item()), end="", flush=True)
|
|
|
|
generated_ids[:, seq_length] = next_token
|
|
tokens.append(int(next_token))
|
|
inputs = torch.cat((inputs, next_token.unsqueeze(0)), dim=-1)
|
|
cache_position = torch.tensor([seq_length], device=torch_device, dtype=torch.int32)
|
|
position_ids = cache_position.unsqueeze(0)
|
|
seq_length += 1
|
|
|
|
cuda_graph_runner = None
|
|
|
|
start_time = time.time()
|
|
|
|
if not use_torch_npu:
|
|
for i in range(1, max_new_tokens):
|
|
if use_flashinfer_mla:
|
|
MLAWrapperSingleton.plan_all(None,None,None,position_ids.squeeze(1)+1,None,
|
|
num_heads, head_dim_ckv, head_dim_kpe, past_key_values.page_size,
|
|
model.model.layers[0].self_attn.softmax_scale, torch.bfloat16, torch.bfloat16)
|
|
global warm_uped
|
|
if use_cuda_graph and ( (warm_uped == True and int(i) == 1) or (warm_uped == False and int(i) == 2) ):
|
|
warm_uped = True
|
|
cuda_graph_runner = CUDAGraphRunner()
|
|
cuda_graph_runner.capture(model, next_token.unsqueeze(0), position_ids, cache_position, past_key_values, torch_device, return_dict=False, use_cache=True)
|
|
next_token = decode_one_tokens(cuda_graph_runner, next_token.unsqueeze(0), position_ids, cache_position, past_key_values, logits_warper, generation_config, use_cuda_graph).to(torch_device)
|
|
inputs = torch.cat((inputs, next_token.unsqueeze(0)), dim=-1)
|
|
generated_ids[:, cache_position] = next_token.int()
|
|
tokens.append(int(next_token))
|
|
seq_length += 1
|
|
|
|
if next_token[0].item() == tokenizer.eos_token_id or tokenizer.decode(next_token.tolist()) == '<|im_end|>':
|
|
print(stream.end(), end="", flush=True)
|
|
break
|
|
else:
|
|
print(stream.put(next_token.item()), end="", flush=True)
|
|
cache_position += 1
|
|
position_ids = cache_position.unsqueeze(0)
|
|
else:
|
|
prof_decode = os.environ["PROF_DECODE"] if "PROF_DECODE" in os.environ else "0"
|
|
prof_ranks = os.environ["PROF_RANK"] if "PROF_RANK" in os.environ else "0"
|
|
prof_ranks = [int(r.strip()) for r in prof_ranks.split(",")]
|
|
if prof_decode == "1" and torch.distributed.get_rank() in prof_ranks and WARM_UP_SKIP_CNT[1] <= 0:
|
|
experimental_config = torch_npu.profiler._ExperimentalConfig(
|
|
aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization,
|
|
profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False
|
|
)
|
|
with torch_npu.profiler.profile(
|
|
activities=[
|
|
torch_npu.profiler.ProfilerActivity.CPU,
|
|
torch_npu.profiler.ProfilerActivity.NPU
|
|
],
|
|
schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=_MAX_DECODE_PROFILE, repeat=1, skip_first=0),
|
|
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./decode_prof"),
|
|
record_shapes=True,
|
|
profile_memory=True,
|
|
with_stack=False,
|
|
with_flops=False,
|
|
with_modules=False,
|
|
experimental_config=experimental_config) as prof:
|
|
decode_wrapper(next_token, position_ids, cache_position, cuda_graph_runner, past_key_values, inputs, seq_length, prof)
|
|
else:
|
|
decode_wrapper(next_token, position_ids, cache_position, cuda_graph_runner, past_key_values, inputs, seq_length)
|
|
WARM_UP_SKIP_CNT[1] -= 1
|
|
|
|
total_time = time.time() - start_time
|
|
tokens_generated = len(tokens)
|
|
tokens_per_second = tokens_generated / total_time
|
|
|
|
if not use_torch_npu:
|
|
print("")
|
|
|
|
print(f"prompt eval count: {prefill_count} token(s)")
|
|
print(f"prompt eval duration: {prefill_time}s")
|
|
print(f"prompt eval rate: {prefill_count/prefill_time} tokens/s")
|
|
print(f"eval count: {tokens_generated} token(s)")
|
|
print(f"eval duration: {total_time}s")
|
|
print(f"eval rate: {tokens_per_second} tokens/s")
|
|
else:
|
|
tp_size = get_tensor_parallel_size()
|
|
if torch.distributed.get_rank() % tp_size == 0:
|
|
rank = f"[rank:{torch.distributed.get_rank()}]"
|
|
msg = f"\n{rank} Eval Time\n"
|
|
msg += rank + f"prompt eval count: {prefill_count} token(s)\n"
|
|
msg += rank + f"prompt eval duration: {prefill_time:.9f}s\n"
|
|
msg += rank + f"prompt eval rate: {prefill_count/prefill_time:.9f} tokens/s\n"
|
|
msg += rank + f"eval count: {tokens_generated} token(s)\n"
|
|
msg += rank + f"eval duration: {total_time:.9f}s\n"
|
|
msg += rank + f"eval rate: {tokens_per_second:.9f} tokens/s\n"
|
|
print(msg)
|
|
|
|
return tokens
|
|
|
|
class InferenceState(enum.Enum):
|
|
UNLOAD = 0
|
|
PREFILL = 1
|
|
GENERATE = 2
|
|
RESTORE = 3
|