# SPDX-License-Identifier: Apache-2.0 # Adapted from # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py # Copyright 2024 The Qwen team. # Copyright 2023 The vLLM team. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # 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. """Inference-only Qwen2 model compatible with HuggingFace weights, with VACC modifications.""" from typing import Iterable, Optional, Set, Tuple, Union, List, Any, Dict import math import os import torch import numpy as np from torch import nn from transformers import Qwen2Config from vllm.attention import Attention, AttentionType from vllm.config import CacheConfig from vllm.distributed import (get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, get_tp_group) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import (MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.forward_context import ForwardContext, get_forward_context # Import original classes to be modified from vllm.model_executor.models.qwen2 import (Qwen2Attention as Qwen2AttentionOrig, Qwen2DecoderLayer as Qwen2DecoderLayerOrig, Qwen2MLP as Qwen2MLPOrig) from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod from vllm.model_executor.layers.quantization.awq import AWQLinearMethod from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase) from vllm_vacc.vllm.model_executor.layers.layernorm import RMSNorm_forward_vacc from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding, apply_interleaved_rope from ..ops.mrope_op import get_sin_cos_mrope import copy logger = init_logger(__name__) # --- VACC specific additions --- from vllm_vacc.vllm.model_executor.models.vars import TRANSPOSE_GPTQ_WEIGHT, USE_FUSED_QWEN_ATTENTION from vllm_vacc.vllm.model_executor.models.vars import BLOCK_GROUP_SIZE as env_blk_grp_size def save_tensor_pth(name, tensor): return import time if tensor is None: return process_id = os.getpid() timestamp = time.strftime("%Y-%m-%d_%H", time.localtime()) pathdir = f'./dump_qwen_pth/{timestamp}/process_{process_id}' if not os.path.exists(pathdir): os.makedirs(pathdir) counter = 0 filename = f"{name}_{counter}.pth" while os.path.exists(os.path.join(pathdir, filename)): filename = f"{name}_{counter}.pth" counter += 1 file_path = os.path.join(pathdir, filename) if isinstance(tensor, torch.Tensor): print(f'save {filename} {tensor.shape}') torch.save(tensor.cpu(), file_path) elif 'sin_cache' in name or 'cos_cache' in name: tensor_list = [i.cpu() for i in tensor] torch.save(tensor_list, file_path) else: torch.save(tensor, file_path) return # --- GPTQ Dequantization Logic (from gptq.py) --- def int32_to_int4(s0, axis = -2): # 要先拉平 shape[1, n] # 每个int32 拆成8个int4, 8个int32表示， 得到[8, n] # x32(int32) => 32bit => 4bit x 8 x4[8] 4bit # x32 31-28 => x4[7] # x32 27-24 => x4[6] # ... # x32 3-0 => x4[0] # x32[index=0] => x4[7,6,5,4,3,2,1,0] # 4bit转真实数字： # 不是按补码方式 # 1111 => 15 => 7 # 15-8 = 7 # 0101 => 6 =>-2 # 6-8 = -2 # 0x 6A CB 37 2B （内存中排列 2B 37 CB 6A） => B273BCA6 => (-8) => int4: 3, -6, -1, -5, 3, 4, 2, -2 # 内存中实际排布为小端模式： # int32: 2B 37 CB 6A => 2,11,3,7,12,11,6,10 => (-8) => -6,3, -5,-1, 4,3, -2,2 => 同一字节所在的两个交换得到 3, -6, -1, -5, 3, 4, 2, -2 # int4: 3, -6, -1, -5, 3, 4, 2, -2 s = s0.view(torch.uint32) all = [] for i in range(8): x = 15 << (i*4) # s2 = torch.bitwise_and(x,s) s2 = torch.from_numpy(np.bitwise_and(x, s.numpy())) s3 = s2 / (2 ** (i*4)) s4 = s3.to(torch.int32) # 补码， 结果不对 # s4[s4 > 7] = s4[s4 > 7]-16 # 直接 - 8 结果正确， 范围： -8-7 s4 = s4 - 8 all.append(s4.reshape(1,*s4.shape)) all = torch.concatenate(all, 0) if axis == -2 or axis == 0: # 8,K//8,N => K//8,8,N => K,N all = all.transpose(-2,0).reshape(-1,all.shape[-1]).contiguous() else: # 8,N,K//8 => N,K//8,8 => N,K all = all.permute(1,2,0).reshape(all.shape[-2],-1).contiguous() return all def dequant_weight(qw, scales, group_size = 128): N = qw.shape[1] int4_to_int32_axis = -2 if TRANSPOSE_GPTQ_WEIGHT: N = qw.shape[0] int4_to_int32_axis = -1 qweight = int32_to_int4(qw,int4_to_int32_axis).to(torch.float16) #int32 => 8 int4 +> fp16 if TRANSPOSE_GPTQ_WEIGHT: scales = scales.T.contiguous() qweight = qweight.T.contiguous() scales = torch.concatenate([scales] * group_size, 1).reshape(-1, N) # scale 按 group_size 扩展， 每 group_size 个数共用一个scale # print('qweight', qweight.shape, qweight.dtype) # print('scale', scales.shape, scales.dtype) dequant_weight = qweight * scales #dequant return dequant_weight def apply_gptq_linear( input_tensor: torch.Tensor, layer: Union[QKVParallelLinear, RowParallelLinear], ) -> torch.Tensor: """ Applies a GPTQ-quantized linear layer by dequantizing weights on-the-fly. """ out_shape = input_tensor.shape[:-1] + (layer.qweight.shape[-2 if TRANSPOSE_GPTQ_WEIGHT else -1], ) # M,N reshaped_x = input_tensor.reshape(-1, input_tensor.shape[-1]) # This assumes the linear_method is attached and is a GPTQ derivative linear_method = layer.quant_method quant_config = linear_method.quant_config # scale_k is a VACC-specific parameter for scale broadcasting scale_k = getattr(linear_method, "scale_k", 1) # Dequantize weight # weight = dequant_weight(layer.qweight.data, layer.scales.data, quant_config.group_size // scale_k) weight = dequant_weight(layer.qweight.cpu(), layer.scales.cpu(), quant_config.group_size // scale_k).to(layer.qweight.device) # Perform GEMM output = torch.matmul(reshaped_x, weight) # Add bias if it exists if hasattr(layer, 'bias') and layer.bias is not None: output = output + layer.bias # out_shape = input_tensor.shape[:-1] + (weight.shape[-1],) return output.reshape(out_shape) def get_gptq_group_size(layer: Union[QKVParallelLinear, RowParallelLinear]): # This assumes the linear_method is attached and is a GPTQ derivative linear_method = layer.quant_method quant_config = linear_method.quant_config # scale_k is a VACC-specific parameter for scale broadcasting scale_k = getattr(linear_method, "scale_k", 1) return quant_config.group_size // scale_k def apply_gptq_linear_( input_tensor: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, bias: Optional[torch.Tensor], qzeros: Optional[torch.Tensor], group_size_h: int, group_size_w: int, ) -> torch.Tensor: """ Applies a GPTQ-quantized linear layer by dequantizing weights on-the-fly. """ out_shape = input_tensor.shape[:-1] + (weight.shape[-2 if TRANSPOSE_GPTQ_WEIGHT else -1], ) # M,N reshaped_x = input_tensor.reshape(-1, input_tensor.shape[-1]) # This assumes the linear_method is attached and is a GPTQ derivative # linear_method = layer.quant_method # quant_config = linear_method.quant_config # # scale_k is a VACC-specific parameter for scale broadcasting # scale_k = getattr(linear_method, "scale_k", 1) # Dequantize weight # weight = dequant_weight(layer.qweight.data, layer.scales.data, quant_config.group_size // scale_k) # weight = dequant_weight(weight.cpu(), weight_scale.cpu(), group_size).to(weight.device) # Perform GEMM # output = torch.matmul(reshaped_x, weight) # print("entering apply_gptq_linear_, reshaped_x shape:", reshaped_x.shape, "reshaped_x stride", reshaped_x.stride(), "input_tensor", input_tensor.shape, "weight shape:", weight.shape, "weight_scale shape:", weight_scale.shape, "group_size:", group_size) output = torch.vacc.w4a8_block_int4_matmul( reshaped_x, weight.transpose(-1, -2), weight_scale.transpose(-1, -2), [group_size_h, group_size_w], ) # print("exiting apply_gptq_linear_, output shape:", output.shape) # Add bias if it exists if bias is not None: output = output + bias # out_shape = input_tensor.shape[:-1] + (weight.shape[-1],) return output.reshape(out_shape) def post_layernorm( x: torch.Tensor, residual: Optional[torch.Tensor] = None, weight: Optional[torch.Tensor] = None, variance_size_override: Optional[int] = None, variance_epsilon: float = 1e-6, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """PyTorch-native implementation equivalent to forward().""" orig_dtype = x.dtype x = x.to(torch.float32) if residual is not None: x = x + residual.to(torch.float32) residual = x.to(orig_dtype) hidden_size = x.shape[-1] if variance_size_override is None: x_var = x else: if hidden_size < variance_size_override: raise ValueError( "Expected hidden_size to be at least " f"{variance_size_override}, but found: {hidden_size}") x_var = x[:, :, :variance_size_override] variance = x_var.pow(2).mean(dim=-1, keepdim=True) x = x * torch.rsqrt(variance + variance_epsilon) x = x.to(orig_dtype) if weight is not None: x = x * weight if residual is None: return x else: return x, residual # --- VACC Fused Kernel for Qwen2 --- def vacc_fused_attn_qwen2_naive( hidden_states: torch.Tensor, residual: Optional[torch.Tensor], hidden_states_norm_weight: torch.Tensor, qkv_proj_weight: torch.Tensor, qkv_proj_weight_scale: torch.Tensor, qkv_proj_bias: Optional[torch.Tensor], qkv_proj_qzeros: Optional[torch.Tensor], sin_cache: List[torch.Tensor], cos_cache: List[torch.Tensor], slot_mapping: torch.Tensor, kv_cache: torch.Tensor, block_tables: torch.Tensor, block_group_size: int, o_proj_weight: torch.Tensor, o_proj_weight_scale: torch.Tensor, o_proj_bias: Optional[torch.Tensor], o_proj_qzeros: Optional[torch.Tensor], seq_lens: List[int], sm_scale: float, num_attention_heads: int, num_key_value_heads: int, flash_attention: bool, is_decode: bool, reduce_result: bool, world_size: int, rank: int, group_id: int, dev_info: List[int] | Tuple[int], block_size: int = 16 ): # qkv_proj_group_size = 128 # o_proj_group_size = 64 qkv_proj_group_size_h = qkv_proj_weight.shape[-1] * 8 // qkv_proj_weight_scale.shape[-1] qkv_proj_group_size_w = qkv_proj_weight.shape[-2] * 8 // qkv_proj_weight_scale.shape[-2] o_proj_group_size_h = o_proj_weight.shape[-1] * 8 // o_proj_weight_scale.shape[-1] o_proj_group_size_w = o_proj_weight.shape[-2] * 8 // o_proj_weight_scale.shape[-2] if residual is not None: hidden_states = hidden_states + residual residual_out = hidden_states save_tensor_pth("hidden_states", hidden_states) # 1. Fused RMSNorm hidden_states_norm = torch.vacc.rms_norm( hidden_states.unsqueeze(0), hidden_states_norm_weight, 1e-6).squeeze(0) # NOTE: for qwen3 and qwen2.5, head_dim is always 128 head_dim = 128 # 2. QKV Projection using on-the-fly dequantization save_tensor_pth("hidden_states_norm", hidden_states_norm) qkv = apply_gptq_linear_(hidden_states_norm, qkv_proj_weight, qkv_proj_weight_scale, qkv_proj_bias, qkv_proj_qzeros, qkv_proj_group_size_h, qkv_proj_group_size_w) save_tensor_pth("qkv", qkv) # Split Q, K, V num_q_heads = num_attention_heads // world_size num_kv_heads = max(1, num_key_value_heads // world_size) q_size = num_q_heads * head_dim kv_size = num_kv_heads * head_dim q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) save_tensor_pth("q", q) save_tensor_pth("k", k) save_tensor_pth("v", v) q = q.view(-1, num_q_heads, head_dim) k = k.view(-1, num_kv_heads, head_dim) v = v.view(-1, num_kv_heads, head_dim) save_tensor_pth("q_reshaped", q) save_tensor_pth("k_reshaped", k) save_tensor_pth("v_reshaped", v) # 3. RoPE, KV Caching, and Attention loop start = 0 attn_outs = [] if is_decode: # convert block_tables to 8K group index block_per_group = block_group_size // block_size block_tables_grouped = (block_tables // block_per_group).to(torch.int32) # logger.warning(f"decode block table: {block_tables}") num_blocks = kv_cache.shape[1] key_cache_split = kv_cache[0].view(num_blocks, -1, num_kv_heads, head_dim) value_cache_split = kv_cache[1].view(num_blocks, -1, num_kv_heads, head_dim) # bs loop for i, seq_len in enumerate(seq_lens): if not is_decode: # Prefill end = start + seq_len else: # Decode end = start + 1 cos = cos_cache[i].unsqueeze(-2) sin = sin_cache[i].unsqueeze(-2) q_i, k_i, v_i = q[start:end], k[start:end], v[start:end] # Apply Rotary Positional Embedding q_rot, k_rot = torch.vacc.RotaryPosEmbedding(q_i, k_i, cos, sin, 0, "neox") save_tensor_pth("q_rot", q_rot) save_tensor_pth("k_rot", k_rot) # Reshape and cache K/V torch.vacc.reshape_and_cache_attention(k_rot, key_cache_split, slot_mapping[start : end, ...]) torch.vacc.reshape_and_cache_attention(v_i, value_cache_split, slot_mapping[start : end, ...]) # Attention calculation if not is_decode: attn_out = torch.vacc.scaled_dot_product_attention( query=q_rot, key=k_rot, value=v_i.contiguous(), attn_mask = None, dropout_p = 0.0, is_causal = True, #causal_attn and not self.need_mask, is_train = False, recompute = False, flash_attention = False, sm_scale=sm_scale) else: # For decode, reconstruct past K/V from cache key_cache_grouped = key_cache_split.view(-1, block_group_size, num_kv_heads, head_dim) value_cache_grouped = value_cache_split.view(-1, block_group_size, num_kv_heads, head_dim) # block_per_group = block_group_size // block_size # block_tables_grouped = (block_tables // block_per_group).to(torch.int32) k_slices = key_cache_grouped[block_tables_grouped[i], :, :, :] k_past = torch.cat([k_slices[j].unsqueeze(0) for j in range(len(block_tables_grouped[i]))], dim=0) k_past = k_past.reshape(-1, num_kv_heads, head_dim)[:seq_len] v_slices = value_cache_grouped[block_tables_grouped[i], :, :, :] v_past = torch.cat([v_slices[j].unsqueeze(0) for j in range(len(block_tables_grouped[i]))], dim=0) v_past = v_past.reshape(-1, num_kv_heads, head_dim)[:seq_len] # k_past = key_cache_split.reshape(-1, key_cache_split.shape[-2], key_cache_split.shape[-1])[:seq_len] # v_past = value_cache_split.reshape(-1, value_cache_split.shape[-2], value_cache_split.shape[-1])[:seq_len] attn_out = torch.vacc.scaled_dot_product_attention( query=q_rot.contiguous(), key=k_past.contiguous(), value=v_past.contiguous(), attn_mask=None, dropout_p=0, is_causal=False, is_train=False, recompute=False, flash_attention=False,#flash_attention sm_scale=sm_scale,) attn_outs.append(attn_out) start = end attn_output_cat = torch.cat(attn_outs, dim=0) save_tensor_pth("attn_output", attn_output_cat) # 4. Output Projection attn_output_reshaped = attn_output_cat.view(hidden_states.shape[0], -1) # print(f"attn_output_reshaped: {attn_output_reshaped}") output = apply_gptq_linear_(attn_output_reshaped,o_proj_weight, o_proj_weight_scale,o_proj_bias, o_proj_qzeros, o_proj_group_size_h, o_proj_group_size_w) # 5. Optional All-Reduce and post layernorm if reduce_result: output = tensor_model_parallel_all_reduce(output) if residual is not None: return output, residual_out return output # --- Rewriting forward methods for Qwen2Attention and Qwen2DecoderLayer --- # uniform the params names from different quantize method def set_fused_params(fused_params: Dict[str, Any], quant_method: QuantizeMethodBase, layer: nn.Module, name: str): if isinstance(quant_method, UnquantizedLinearMethod): fused_params[name + '_weight'] = layer.weight fused_params[name + '_weight_scale'] = torch.Tensor() fused_params[name + '_bias'] = None if not hasattr(layer, 'bias') else layer.bias fused_params[name + '_qzeros'] = None elif isinstance(quant_method, Fp8LinearMethod): fused_params[name + '_weight'] = layer.weight fused_params[name + '_weight_scale'] = layer.weight_scale_inv fused_params[name + '_bias'] = None if not hasattr(layer, 'bias') else layer.bias fused_params[name + '_qzeros'] = None if not hasattr(layer, 'qzeros') else layer.qzeros elif isinstance(quant_method, GPTQLinearMethod): fused_params[name + '_weight'] = layer.qweight fused_params[name + '_weight_scale'] = layer.scales fused_params[name + '_bias'] = None if not hasattr(layer, 'bias') else layer.bias fused_params[name + '_qzeros'] = None if not hasattr(layer, 'qzeros') else layer.qzeros elif isinstance(quant_method, AWQLinearMethod): fused_params[name + '_weight'] = layer.qweight fused_params[name + '_weight_scale'] = layer.scales fused_params[name + '_bias'] = None if not hasattr(layer, 'bias') else layer.bias fused_params[name + '_qzeros'] = None if not hasattr(layer, 'qzeros') else layer.qzeros else: raise ValueError(f"Unsupported quant_method: {quant_method}") # 全局变量，用于为每次 dump 操作创建一个唯一的批次 ID # NOTE: This counter is process-specific. If you run multiple processes, # each will have its own counter starting from 0. The directory structure # already includes the process_id, so there won't be conflicts. DUMP_COUNTER = 0 def dump_qwen2_attention_params(hidden_states, residual, fused_params, attn_metadata, is_decode, sin_cache, cos_cache, kv_cache, env_blk_grp_size, scaling, total_num_heads, total_num_kv_heads, world_size, rank, group_id, dev_info): """ 收集所有需要dump的参数，并以“批次”为单位进行保存。 """ # if is_decode: # return # 1. 将所有需要保存的张量和数据收集到一个字典中 # 键是文件名（不含扩展名），值是数据本身。 # 这里的关键是：即使某个值是 None，我们也把它放进字典， # 让保存函数自己去处理，而不是在这里跳过。 params_to_dump = { 'hidden_states': hidden_states, 'residual': residual, 'hidden_states_norm_weight': fused_params.get('input_layernorm_weight'), 'qkv_proj_weight': fused_params.get('qkv_proj_weight'), 'qkv_proj_weight_scale': fused_params.get('qkv_proj_weight_scale'), 'qkv_proj_bias': fused_params.get('qkv_proj_bias'), 'qkv_proj_qzeros': fused_params.get('qkv_proj_qzeros'), 'o_proj_weight': fused_params.get('o_proj_weight'), 'o_proj_weight_scale': fused_params.get('o_proj_weight_scale'), 'o_proj_bias': fused_params.get('o_proj_bias'), 'o_proj_qzeros': fused_params.get('o_proj_qzeros'), 'sin_cache': sin_cache, 'cos_cache': cos_cache, 'kv_cache': kv_cache, 'slot_mapping': attn_metadata.slot_mapping if attn_metadata else None, 'block_tables': attn_metadata.block_tables if attn_metadata else None, 'seq_lens': attn_metadata.seq_lens if attn_metadata else None, 'block_group_size': env_blk_grp_size, 'sm_scale': scaling, 'num_attention_heads': total_num_heads, 'num_key_value_heads': total_num_kv_heads, 'world_size': world_size, 'rank': rank, 'group_id': group_id, 'dev_info': dev_info, } # 2. 一次性调用新的批次保存函数 save_dump_batch(params_to_dump) def save_dump_batch(tensors_dict): """ 将一个字典中的所有张量和数据保存到一个唯一的、代表“同一次调用”的子目录中。 """ global DUMP_COUNTER import time # 基础目录结构：./dump_data/{timestamp}/process_{pid} process_id = os.getpid() timestamp = time.strftime("%Y-%m-%d_%H", time.localtime()) base_dir = f'./dump_qwen2_prefill_pth/{timestamp}/process_{process_id}' # 为这次“批次”的dump创建一个唯一的子目录 # 这就是你想要的“魔法数”的实现 batch_dir = os.path.join(base_dir, f'dump_{DUMP_COUNTER}') # 递增全局计数器，为下一次调用做准备 DUMP_COUNTER += 1 # 创建目录（如果不存在） os.makedirs(batch_dir, exist_ok=True) print(f"--- Dumping batch to: {batch_dir} ---") # 遍历字典，保存每一个非 None 的项 for name, data in tensors_dict.items(): if data is None: # 如果数据是 None，直接跳过，不保存任何文件 continue file_path = os.path.join(batch_dir, f"{name}.pth") try: if isinstance(data, torch.Tensor): print(f' Saving tensor: {name}.pth, shape: {data.shape}, dtype: {data.dtype}') torch.save(data.cpu(), file_path) elif name in ['sin_cache', 'cos_cache'] and isinstance(data, (list, tuple)): # 对 sin/cos cache 的特殊处理 tensor_list = [i.cpu() for i in data] print(f' Saving tensor list: {name}.pth, count: {len(tensor_list)}') torch.save(tensor_list, file_path) else: # 保存其他可序列化的 Python 对象（如数字、字符串等） print(f' Saving metadata: {name}.pth, value: {data}') torch.save(data, file_path) except Exception as e: print(f" [ERROR] Failed to save {name}.pth: {e}") print(f"--- Finished dumping batch {DUMP_COUNTER - 1} ---") def has_nan(t1, t2): has_nan_1 = torch.isnan(t1).any() has_nan_2 = torch.isnan(t2).any() # 2. 根据检查结果进行判断和计算 if has_nan_1 or has_nan_2: print("计算被中止，因为检测到了NaN值。") if has_nan_1: print(" - Tensor 1 包含 NaN。") if has_nan_2: print(" - Tensor 2 包含 NaN。") class MlpSplitConfig: def __init__(self): self.k_split = 0 self.n_split = 0 self.w1_weight_addr = 0 self.w3_weight_addr = 0 self.matA_addr_0 = 0 self.matA_addr_1 = 0 self.w1_splitk_buffer = 0 self.w3_splitk_buffer = 0 self.w13_multiply = 0 self.w2_splitn_buffer_list = [] self.w13_dequant_block_size = 0 self.w2_dequant_block_size = 0 def get_dequant_block_size_int4(k, n, k_split, n_split): ret = [] w13_dequant_block_init_size = 128 k_per_split = k // k_split if k_per_split % w13_dequant_block_init_size == 0: ret.append(w13_dequant_block_init_size) else: while w13_dequant_block_init_size >= 16 and k_per_split % w13_dequant_block_init_size != 0: w13_dequant_block_init_size //= 2 if w13_dequant_block_init_size >= 16: ret.append(w13_dequant_block_init_size) else: ret.append(-1) w2_dequant_block_init_size = 128 n_per_split = n // n_split if n_per_split % w2_dequant_block_init_size == 0: ret.append(w2_dequant_block_init_size) else: while w2_dequant_block_init_size >= 16 and n_per_split % w2_dequant_block_init_size != 0: w2_dequant_block_init_size //= 2 if w2_dequant_block_init_size >= 16: ret.append(w2_dequant_block_init_size) else: ret.append(-1) return ret def get_split_config(total_split): ret = [] max_split_n = math.ceil(math.sqrt(total_split)) for i in range(1, max_split_n + 1): if total_split % i == 0: cur_config = [total_split // i, i] # [k_split, n_split] ret.append(cur_config) split_config_size = len(ret) for i in range(split_config_size - 1, -1, -1): if ret[i][0] != ret[i][1]: ret.append([ret[i][1], ret[i][0]]) return ret def validate_mlp_split_config(k, n, k_split, n_split, config): min_seqlen = 1024 # 尺寸可除性检查 if k % k_split != 0: return False if n % (n_split * 4) != 0: return False if k % (k_split * 4) != 0: return False if (n / (n_split * 4)) % 16 != 0: return False if (k / (k_split * 4)) % 16 != 0: return False # 计算权重和输入矩阵的大小 w13_split_size = (k // k_split) * (n // n_split // 4) * 2 * 2 matA_size = min_seqlen * (k // k_split // 4) * 2 * 2 # 检查SSRAM 4-7区空间是否足够 if w13_split_size + matA_size > 5 * 1024 * 1024: return False # 计算剩余空间 ssram_4_to_7_size_left = 5 * 1024 * 1024 - (w13_split_size + matA_size) w13_result_size = min_seqlen * (n // n_split // 4) * 2 # 检查结果存储空间 if w13_result_size * 3 > ssram_4_to_7_size_left + 5 * 1024 * 1024: return False ssram_0_to_3_size_left = 5 * 1024 * 1024 w13_result_save_type = -1 # 确定结果存储方案 if w13_result_size * 3 <= 5 * 1024 * 1024: w13_result_save_type = 0 ssram_0_to_3_size_left -= w13_result_size * 3 elif w13_result_size * 2 <= 5 * 1024 * 1024: if w13_result_size <= ssram_4_to_7_size_left: w13_result_save_type = 1 ssram_0_to_3_size_left -= w13_result_size * 2 ssram_4_to_7_size_left -= w13_result_size else: return False else: return False # 处理n_split > 1的情况 if n_split > 1: w2_split_k_buffer_total_size = min_seqlen * (k // k_split // 4) * 2 * k_split w2_split_k_buffer_size = min_seqlen * (k // k_split // 4) * 2 if w2_split_k_buffer_total_size > ssram_0_to_3_size_left + ssram_4_to_7_size_left: return False available_blocks = (ssram_0_to_3_size_left // w2_split_k_buffer_size) + \ (ssram_4_to_7_size_left // w2_split_k_buffer_size) if available_blocks < k_split: return False # 设置配置参数 config.k_split = k_split config.n_split = n_split config.w1_weight_addr = 0x54000000 config.w3_weight_addr = config.w1_weight_addr + (k // k_split) * (n // n_split // 4) * 2 config.matA_addr_0 = config.w3_weight_addr + (k // k_split) * (n // n_split // 4) * 2 config.matA_addr_1 = config.matA_addr_0 + min_seqlen * (k // k_split // 4) * 2 config.w1_splitk_buffer = 0x50000000 config.w3_splitk_buffer = config.w1_splitk_buffer + min_seqlen * (n // n_split // 4) * 2 if w13_result_save_type == 0: config.w13_multiply = config.w3_splitk_buffer + min_seqlen * (n // n_split // 4) * 2 elif w13_result_save_type == 1: config.w13_multiply = config.matA_addr_1 + min_seqlen * (k // k_split // 4) * 2 # 处理w2缓冲区分配 if n_split > 1: w2_split_k_buffer_size = min_seqlen * (k // k_split // 4) * 2 w2_split_k_buffer_in_ssram_0_to_3 = min(ssram_0_to_3_size_left // w2_split_k_buffer_size, k_split) w2_split_k_buffer_in_ssram_4_to_7 = k_split - w2_split_k_buffer_in_ssram_0_to_3 if w13_result_save_type == 0: ssram_0_to_3_start_addr = config.w13_multiply + min_seqlen * (n // n_split // 4) * 2 ssram_4_to_7_start_addr = config.matA_addr_1 + min_seqlen * (k // k_split // 4) * 2 else: ssram_0_to_3_start_addr = config.w3_splitk_buffer + min_seqlen * (n // n_split // 4) * 2 ssram_4_to_7_start_addr = config.w13_multiply + min_seqlen * (n // n_split // 4) * 2 for i in range(w2_split_k_buffer_in_ssram_0_to_3): addr = ssram_0_to_3_start_addr + i * w2_split_k_buffer_size config.w2_splitn_buffer_list.append(addr) for i in range(w2_split_k_buffer_in_ssram_4_to_7): addr = ssram_4_to_7_start_addr + i * w2_split_k_buffer_size config.w2_splitn_buffer_list.append(addr) else: config.w2_splitn_buffer_list.append(0) # 获取解量化块大小 block_size = get_dequant_block_size_int4(k, n, k_split, n_split) if block_size[0] == -1 or block_size[1] == -1: return False config.w13_dequant_block_size = block_size[0] config.w2_dequant_block_size = block_size[1] return True def get_mlp_split_schedule(k, n): min_split = math.ceil(k * n / 2 / 1310720) # 最大权重分割大小1.25MB current_split = min_split config = MlpSplitConfig() for _ in range(100): # 最多尝试100次 split_configs = get_split_config(current_split) for cfg in split_configs: if validate_mlp_split_config(k, n, cfg[0], cfg[1], config): return True, config current_split += 1 return False, None def Qwen2MLP__init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super(Qwen2MLPOrig, self).__init__() self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.down_proj", ) if hidden_act != "silu": raise ValueError(f"Unsupported activation: {hidden_act}. " "Only silu is supported for now.") self.act_fn = SiluAndMul() if hasattr(self.gate_up_proj, 'quant_method') and isinstance(self.gate_up_proj.quant_method, GPTQLinearMethod): intermediate_size_per_partition = intermediate_size // self.down_proj.tp_size success, config = get_mlp_split_schedule(hidden_size, intermediate_size_per_partition) if success: self.w13_dequant_block_size = config.w13_dequant_block_size self.w2_dequant_block_size = config.w2_dequant_block_size group_size = self.gate_up_proj.quant_method.quant_config.group_size scale_k_w13 = group_size // config.w13_dequant_block_size scale_k_w2 = group_size // config.w2_dequant_block_size self.gate_up_proj.quant_method.scale_k = max(scale_k_w13, self.gate_up_proj.quant_method.scale_k) self.down_proj.quant_method.scale_k = max(scale_k_w2, self.down_proj.quant_method.scale_k) if self.gate_up_proj.quant_method.scale_k > 1 and len(self.gate_up_proj.scales.data.shape) == 2 and \ self.gate_up_proj.scales.data.dtype in [torch.float16, torch.bfloat16, torch.float32]: w13_scales = self.gate_up_proj.scales.data scale_and_zero_size = hidden_size * self.gate_up_proj.quant_method.scale_k // group_size w13_scales = torch.empty( scale_and_zero_size, w13_scales.shape[-1], dtype=w13_scales.dtype, ) self.gate_up_proj.scales.data = w13_scales if self.down_proj.quant_method.scale_k > 1 and len(self.down_proj.scales.data.shape) == 2 and \ self.down_proj.scales.data.dtype in [torch.float16, torch.bfloat16, torch.float32]: w2_scales = self.down_proj.scales.data scale_and_zero_size = intermediate_size_per_partition * self.down_proj.quant_method.scale_k // group_size w2_scales = torch.empty( scale_and_zero_size, w2_scales.shape[-1], dtype=w2_scales.dtype, ) self.down_proj.scales.data = w2_scales else: raise ValueError(f"Cannot find valid MLP split schedule for hidden_size {hidden_size}, " "intermediate_size {intermediate_size}, tp_size {self.down_proj.tp_size}") def get_cos_sin_cache(rotary_emb: Union["MRotaryEmbedding", "RotaryEmbedding"], attn_metadata: Union["AttentionMetadata", dict[str, "AttentionMetadata"]], positions: Union[torch.Tensor, list], is_decode: bool): if isinstance(rotary_emb, MRotaryEmbedding): # get mrope sin/cos cos_cache, sin_cache = get_sin_cos_mrope(rotary_emb, positions) if len(attn_metadata.seq_lens) > 1: if is_decode: cos_cache = torch.chunk(cos_cache, len(attn_metadata.seq_lens)) sin_cache = torch.chunk(sin_cache, len(attn_metadata.seq_lens)) else: cos_cache = torch.split(cos_cache, attn_metadata.seq_lens) sin_cache = torch.split(sin_cache, attn_metadata.seq_lens) else: cos_cache = [cos_cache] sin_cache = [sin_cache] else: if is_decode: positions = [i - 1 for i in attn_metadata.seq_lens] cos_cache = [rotary_emb.cos_cache[i:i+1, ...] for i in positions] sin_cache = [rotary_emb.sin_cache[i:i+1, ...] for i in positions] else: cos_cache = [rotary_emb.cos_cache[:i, ...] for i in attn_metadata.seq_lens] sin_cache = [rotary_emb.sin_cache[:i, ...] for i in attn_metadata.seq_lens] return cos_cache, sin_cache class Qwen2MLP(nn.Module): def forward(self, x): #TODO for other quant_method if hasattr(self.gate_up_proj, 'quant_method') and isinstance(self.gate_up_proj.quant_method, GPTQLinearMethod): # bit = self.gate_up_proj.quant_method.quant_config.weight_bits # pack_num = torch.iinfo(self.gate_up_proj.qweight.dtype).bits // bit # w13_group_size = pack_num * self.gate_up_proj.qweight.shape[1] // self.gate_up_proj.scales.shape[1] # w2_group_size = pack_num * self.down_proj.qweight.shape[1] // self.down_proj.scales.shape[1] tp_group = get_tp_group() batch_size = x.shape[0] if batch_size > 4: y = torch.vacc.fuse_mlp_qwen_int4( x, self.gate_up_proj.qweight, self.down_proj.qweight, self.gate_up_proj.scales, self.down_proj.scales, self.gate_up_proj.qzeros, self.down_proj.qzeros, [1, self.w13_dequant_block_size], [1, self.w2_dequant_block_size] ) mlp_out = tensor_model_parallel_all_reduce(y) return mlp_out else: mlp_out = torch.vacc.fuse_mlp_qwen_int4_reduce( x, self.gate_up_proj.qweight, self.down_proj.qweight, self.gate_up_proj.scales, self.down_proj.scales, self.gate_up_proj.qzeros, self.down_proj.qzeros, [1, self.w13_dequant_block_size], [1, self.w2_dequant_block_size], world_size = tp_group.world_size, rank = tp_group.rank_in_group, group_id = tp_group.group_id, dev_info = tp_group.rank_device_infos ) return mlp_out elif hasattr(self.gate_up_proj, 'quant_method') and isinstance(self.gate_up_proj.quant_method, UnquantizedLinearMethod): batch_size = x.shape[0] if batch_size <= 4: hidden_states = torch.vacc.fuse_mlp_qwen_fp16_bf16_reduce(x.view(-1, x.shape[-1]), self.gate_up_proj.weight, self.down_proj.weight, world_size = get_tp_group().world_size, rank = get_tp_group().rank_in_group, group_id = get_tp_group().group_id, dev_info = get_tp_group().rank_device_infos ) return hidden_states gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj(x) return x class Qwen2Attention(nn.Module): def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None, # Added for fused kernel cos_cache: list[torch.Tensor] = None, sin_cache: list[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: forward_context: ForwardContext = get_forward_context() attn_metadata = forward_context.attn_metadata if isinstance(attn_metadata, dict): attn_metadata = attn_metadata[self.attn.layer_name] kv_cache = self.attn.kv_cache[forward_context.virtual_engine] is_decode = attn_metadata.prefill_metadata is None # print('attn_metadata', attn_metadata) use_gptq = hasattr(self.qkv_proj, 'quant_method') and isinstance(self.qkv_proj.quant_method, GPTQLinearMethod) use_unquan = hasattr(self.qkv_proj, 'quant_method') and isinstance(self.qkv_proj.quant_method, UnquantizedLinearMethod) # if USE_FUSED_QWEN_ATTENTION and not is_decode: if use_gptq: # Determine if we should reduce the result inside the kernel # print(f"Qwen2Attention forward, is_decode: {is_decode}, kv_cache shape: {kv_cache.shape}, attn_metadata.seq_lens: {attn_metadata.seq_lens}") kv_cache = kv_cache.view(2, -1, 16, self.total_num_kv_heads // get_tp_group().world_size, self.head_dim) reduce_result = is_decode # reduce_result = False # total_bytes = hidden_states.numel() * hidden_states.element_size() * get_tp_group().world_size # if total_bytes < 4194304: # Heuristic from qwen3_moe_vacc # reduce_result = True # Get rotary caches # cos_cache = [] # sin_cache = [] # start = 0 # for seq_len in attn_metadata.seq_lens: # if not is_decode: # end = start + seq_len # else: # end = start + 1 # # NOTE: # # in prefill stage, the value from positions[start] to positions[end] should be always contiguous # cos_cache.append(self.rotary_emb.cos_cache[positions[start]:positions[end-1] + 1, ...]) # sin_cache.append(self.rotary_emb.sin_cache[positions[start]:positions[end-1] + 1, ...]) # start = end if is_decode: positions = [i - 1 for i in attn_metadata.seq_lens] cos_cache = [self.rotary_emb.cos_cache[i:i+1, ...] for i in positions] sin_cache = [self.rotary_emb.sin_cache[i:i+1, ...] for i in positions] else: cos_cache = [self.rotary_emb.cos_cache[:i, ...] for i in attn_metadata.seq_lens] sin_cache = [self.rotary_emb.sin_cache[:i, ...] for i in attn_metadata.seq_lens] if residual is None: res_out = hidden_states if USE_FUSED_QWEN_ATTENTION: qkv_proj_qzeros = torch.Tensor() o_proj_bias = torch.Tensor() o_proj_qzeros = torch.Tensor() qkv_proj_bias = self.fused_params['qkv_proj_bias'] # if is_decode: # qkv_proj_bias = self.fused_params['qkv_proj_bias'] # if hidden_states.shape[0] > 1: # qkv_proj_bias = qkv_proj_bias.repeat(len(attn_metadata.seq_lens), 1) # else: # qkv_proj_bias = self.fused_params['qkv_proj_bias'] # qkv_proj_bias = qkv_proj_bias.repeat(sum(attn_metadata.seq_lens), 1) # print("Qwen2Attention forward, is_decode:", is_decode, # "kv_cache shape:", kv_cache.shape, # "attn_metadata.seq_lens:", attn_metadata.seq_lens, # "qkv_proj_bias shape:", qkv_proj_bias.shape, # "qkv_proj_bias:", qkv_proj_bias, # "qkv_proj_qzeros shape:", qkv_proj_qzeros.shape, # "o_proj_bias shape:", o_proj_bias.shape, # "o_proj_bias:", o_proj_bias, # "o_proj_qzeros shape:", o_proj_qzeros.shape, # "o_proj_qzeros:", o_proj_qzeros, # "reduce_result:", reduce_result, # "world_size:", get_tp_group().world_size, # "rank:", get_tp_group().rank_in_group, # "group_id:", get_tp_group().group_id, # "dev_info:", get_tp_group().rank_device_infos) # print("residual shape:", residual.shape if residual is not None else "None", # "hidden_states shape:", hidden_states.shape, # "hidden_states dtype:", hidden_states.dtype, # "hidden_states device:", hidden_states.device, # "attn_metadata.slot_mapping shape:", attn_metadata.slot_mapping.shape if attn_metadata.slot_mapping is not None else "None", # "attn_metadata.block_tables shape:", attn_metadata.block_tables.shape if attn_metadata.block_tables is not None else "None", # "attn_metadata.block_tables dtype:", attn_metadata.block_tables.dtype if attn_metadata.block_tables is not None else "None", # "attn_metadata.block_tables device:", attn_metadata.block_tables.device if attn_metadata.block_tables is not None else "None", # "attn_metadata.seq_lens:", attn_metadata.seq_lens, # "attn_metadata.seq_lens length:", len(attn_metadata.seq_lens), # "attn_metadata.seq_lens dtype:", type(attn_metadata.seq_lens), # # "attn_metadata.seq_lens device:", attn_metadata.seq_lens[0].device if attn_metadata.seq_lens else "None", # "attn_metadata.seq_lens type:", type(attn_metadata.seq_lens[0]) if attn_metadata.seq_lens else "None", # "attn_metadata.seq_lens content:", attn_metadata.seq_lens if attn_metadata.seq_lens else "None", # "attn_metadata.slot_mapping dtype:", attn_metadata.slot_mapping.dtype if attn_metadata.slot_mapping is not None else "None", # "attn_metadata.slot_mapping device:", attn_metadata.slot_mapping.device if attn_metadata.slot_mapping is not None else "None", # "attn_metadata.slot_mapping type:", type(attn_metadata.slot_mapping) if attn_metadata.slot_mapping is not None else "None", # "attn_metadata.slot_mapping content:", attn_metadata.slot_mapping if attn_metadata.slot_mapping is not None else "None", # "self.fused_params['input_layernorm_weight'] shape:", self.fused_params['input_layernorm_weight'].shape if 'input_layernorm_weight' in self.fused_params else "None", # "self.fused_params['input_layernorm_weight'] dtype:", self.fused_params['input_layernorm_weight'].dtype if 'input_layernorm_weight' in self.fused_params else "None", # "self.fused_params['input_layernorm_weight'] device:", self.fused_params['input_layernorm_weight'].device if 'input_layernorm_weight' in self.fused_params else "None", # "self.fused_params['qkv_proj_weight'] shape:", self.fused_params['qkv_proj_weight'].shape if 'qkv_proj_weight' in self.fused_params else "None", # "self.fused_params['qkv_proj_weight'] dtype:", self.fused_params['qkv_proj_weight'].dtype if 'qkv_proj_weight' in self.fused_params else "None", # "self.fused_params['qkv_proj_weight'] device:", self.fused_params['qkv_proj_weight'].device if 'qkv_proj_weight' in self.fused_params else "None", # "self.fused_params['qkv_proj_weight_scale'] shape:", self.fused_params['qkv_proj_weight_scale'].shape if 'qkv_proj_weight_scale' in self.fused_params else "None", # "self.fused_params['qkv_proj_weight_scale'] dtype:", self.fused_params['qkv_proj_weight_scale'].dtype if 'qkv_proj_weight_scale' in self.fused_params else "None", # "self.fused_params['qkv_proj_weight_scale'] device:", self.fused_params['qkv_proj_weight_scale'].device if 'qkv_proj_weight_scale' in self.fused_params else "None", # "self.fused_params['qkv_proj_bias'] shape:", self.fused_params['qkv_proj_bias'].shape if 'qkv_proj_bias' in self.fused_params else "None", # "self.fused_params['qkv_proj_bias'] dtype:", self.fused_params['qkv_proj_bias'].dtype if 'qkv_proj_bias' in self.fused_params else "None", # "self.fused_params['qkv_proj_bias'] device:", self.fused_params['qkv_proj_bias'].device if 'qkv_proj_bias' in self.fused_params else "None", # "self.fused_params['qkv_proj_qzeros'] shape:", self.fused_params['qkv_proj_qzeros'].shape if 'qkv_proj_qzeros' in self.fused_params else "None", # "self.fused_params['qkv_proj_qzeros'] dtype:", self.fused_params['qkv_proj_qzeros'].dtype if 'qkv_proj_qzeros' in self.fused_params else "None", # "self.fused_params['qkv_proj_qzeros'] device:", self.fused_params['qkv_proj_qzeros'].device if 'qkv_proj_qzeros' in self.fused_params else "None", # "self.fused_params['o_proj_weight'] shape:", self.fused_params['o_proj_weight'].shape if 'o_proj_weight' in self.fused_params else "None", # "self.fused_params['o_proj_weight'] dtype:", self.fused_params['o_proj_weight'].dtype if 'o_proj_weight' in self.fused_params else "None", # "self.fused_params['o_proj_weight'] device:", self.fused_params['o_proj_weight'].device if 'o_proj_weight' in self.fused_params else "None", # "self.fused_params['o_proj_weight_scale'] shape:", self.fused_params['o_proj_weight_scale'].shape if 'o_proj_weight_scale' in self.fused_params else "None", # "self.fused_params['o_proj_weight_scale'] dtype:", self.fused_params['o_proj_weight_scale'].dtype if 'o_proj_weight_scale' in self.fused_params else "None", # "self.fused_params['o_proj_weight_scale'] device:", self.fused_params['o_proj_weight_scale'].device if 'o_proj_weight_scale' in self.fused_params else "None", # # "self.fused_params['o_proj_bias'] shape:", self.fused_params['o_proj_bias'].shape if 'o_proj_bias' in self.fused_params else "None", # # "self.fused_params['o_proj_bias'] dtype:", self.fused_params['o_proj_bias'].dtype if 'o_proj_bias' in self.fused_params else "None", # # "self.fused_params['o_proj_bias'] device:", self.fused_params['o_proj_bias'].device if 'o_proj_bias' in self.fused_params else "None", # # "self.fused_params['o_proj_qzeros'] shape:", self.fused_params['o_proj_qzeros'].shape if 'o_proj_qzeros' in self.fused_params else "None", # # "self.fused_params['o_proj_qzeros'] dtype:", self.fused_params['o_proj_qzeros'].dtype if 'o_proj_qzeros' in self.fused_params else "None", # # "self.fused_params['o_proj_qzeros'] device:", self.fused_params['o_proj_qzeros'].device if 'o_proj_qzeros' in self.fused_params else "None", # "sin_cache length:", len(sin_cache), # "cos_cache length:", len(cos_cache), # "sin_cache[0] shape:", sin_cache[0].shape if sin_cache else "None", # "cos_cache[0] shape:", cos_cache[0].shape if cos_cache else "None", # "sin_cache[0] dtype:", sin_cache[0].dtype if sin_cache else "None", # "cos_cache[0] dtype:", cos_cache[0].dtype if cos_cache else "None", # "sin_cache[0] device:", sin_cache[0].device if sin_cache else "None", # "cos_cache[0] device:", cos_cache[0].device if cos_cache else "None", # "kv_cache shape:", kv_cache.shape, # "kv_cache dtype:", kv_cache.dtype, # "kv_cache device:", kv_cache.device, # "self.scaling:", self.scaling, # "self.total_num_heads:", self.total_num_heads, # "self.total_num_kv_heads:", self.total_num_kv_heads, # "is_decode:", is_decode, # "reduce_result:", reduce_result # ) # Call the naive fused kernel ####################################call fusion kernel###################################### # print("starting fused attention kernel......................................, is_decode:", is_decode) # print('hidden_states', hidden_states.shape) # print('attn_metadata.slot_mapping,', attn_metadata.slot_mapping,) # print('attn_metadata.block_tables,', attn_metadata.block_tables,) # print('attn_metadata.seq_lens,', attn_metadata.seq_lens,) # print('is_decode,', is_decode) attn_outs = torch.vacc.fuse_atten_qwen2( hidden_states, residual, self.fused_params['input_layernorm_weight'], self.fused_params['qkv_proj_weight'], self.fused_params['qkv_proj_weight_scale'], qkv_proj_bias, qkv_proj_qzeros, #self.fused_params['qkv_proj_qzeros'], sin_cache, cos_cache, attn_metadata.slot_mapping, kv_cache, attn_metadata.block_tables, env_blk_grp_size, self.fused_params['o_proj_weight'], self.fused_params['o_proj_weight_scale'], o_proj_bias, #self.fused_params['o_proj_bias'], o_proj_qzeros, #self.fused_params['o_proj_qzeros'], attn_metadata.seq_lens, self.scaling, self.total_num_heads, self.total_num_kv_heads, is_decode, is_decode, reduce_result, world_size=get_tp_group().world_size, rank=get_tp_group().rank_in_group, group_id=get_tp_group().group_id, dev_info=get_tp_group().rank_device_infos) # print("finished fused attention kernel......................................, is_decode:", is_decode) # attn_outs1 = vacc_fused_attn_qwen2_naive( # hidden_states=hidden_states, # residual=residual, # hidden_states_norm_weight=self.fused_params['input_layernorm_weight'], # ###ignore the group size, it is not used in the fused kernel # # qkv_proj_group_size=get_gptq_group_size(self.qkv_proj), # # o_proj_group_size=get_gptq_group_size(self.o_proj), # ### # qkv_proj_weight=self.fused_params['qkv_proj_weight'], # qkv_proj_weight_scale=self.fused_params['qkv_proj_weight_scale'], # qkv_proj_bias=self.fused_params['qkv_proj_bias'], # qkv_proj_qzeros=self.fused_params['qkv_proj_qzeros'], # sin_cache=sin_cache, # cos_cache=cos_cache, # slot_mapping=attn_metadata.slot_mapping, # kv_cache=kv_cache, # block_tables=attn_metadata.block_tables, # block_group_size=env_blk_grp_size, # o_proj_weight=self.fused_params['o_proj_weight'], # o_proj_weight_scale=self.fused_params['o_proj_weight_scale'], # o_proj_bias=self.fused_params['o_proj_bias'], # o_proj_qzeros=self.fused_params['o_proj_qzeros'], # seq_lens=attn_metadata.seq_lens, # sm_scale=self.scaling, # num_attention_heads=self.total_num_heads, # num_key_value_heads=self.total_num_kv_heads, # # head_dim=self.head_dim, # is_decode=is_decode, # flash_attention=is_decode, # reduce_result=reduce_result, # world_size=get_tp_group().world_size, # rank=get_tp_group().rank_in_group, # group_id=get_tp_group().group_id, # dev_info=get_tp_group().rank_device_infos # ) # if residual is None: # has_nan(attn_outs1, attn_outs) # ret = torch.isnan(attn_outs).any() # if torch.isinf(attn_outs).any() or torch.isnan(attn_outs).any(): # dump_qwen2_attention_params(hidden_states=hidden_states, # # output=output_, # residual=residual, # fused_params=self.fused_params, # attn_metadata=attn_metadata, # is_decode=is_decode, # sin_cache=sin_cache, # cos_cache=cos_cache, # kv_cache=kv_cache, # env_blk_grp_size=env_blk_grp_size, # scaling=self.scaling, # total_num_heads=self.total_num_heads, # total_num_kv_heads=self.total_num_kv_heads, # world_size=get_tp_group().world_size, # rank=get_tp_group().rank_in_group, # group_id=get_tp_group().group_id, # dev_info=get_tp_group().rank_device_infos) # cos_sim = torch.cosine_similarity(attn_outs1.cpu().reshape([-1]).float(), attn_outs.cpu().reshape([-1]).float(), dim=-1) # print("cos_sim witout residual:",cos_sim ) # else : # has_nan(attn_outs1[0], attn_outs[0]) # ret = torch.isnan(attn_outs[0]).any() # if ret: # dump_qwen2_attention_params(hidden_states=hidden_states, # # output=output_, # residual=residual, # fused_params=self.fused_params, # attn_metadata=attn_metadata, # is_decode=is_decode, # sin_cache=sin_cache, # cos_cache=cos_cache, # kv_cache=kv_cache, # env_blk_grp_size=env_blk_grp_size, # scaling=self.scaling, # total_num_heads=self.total_num_heads, # total_num_kv_heads=self.total_num_kv_heads, # world_size=get_tp_group().world_size, # rank=get_tp_group().rank_in_group, # group_id=get_tp_group().group_id, # dev_info=get_tp_group().rank_device_infos) # cos_sim = torch.cosine_similarity(attn_outs1[0].cpu().reshape([-1]).float(), attn_outs[0].cpu().reshape([-1]).float(), dim=-1) # print("cos_sim residual:",cos_sim ) else: ########################################################################## # dump_qwen2_attention_params(hidden_states=hidden_states, # # output=output_, # residual=residual, # fused_params=self.fused_params, # attn_metadata=attn_metadata, # is_decode=is_decode, # sin_cache=sin_cache, # cos_cache=cos_cache, # kv_cache=kv_cache, # env_blk_grp_size=env_blk_grp_size, # scaling=self.scaling, # total_num_heads=self.total_num_heads, # total_num_kv_heads=self.total_num_kv_heads, # world_size=get_tp_group().world_size, # rank=get_tp_group().rank_in_group, # group_id=get_tp_group().group_id, # dev_info=get_tp_group().rank_device_infos) attn_outs = vacc_fused_attn_qwen2_naive( hidden_states=hidden_states, residual=residual, hidden_states_norm_weight=self.fused_params['input_layernorm_weight'], ###ignore the group size, it is not used in the fused kernel # qkv_proj_group_size=get_gptq_group_size(self.qkv_proj), # o_proj_group_size=get_gptq_group_size(self.o_proj), ### qkv_proj_weight=self.fused_params['qkv_proj_weight'], qkv_proj_weight_scale=self.fused_params['qkv_proj_weight_scale'], qkv_proj_bias=self.fused_params['qkv_proj_bias'], qkv_proj_qzeros=None, sin_cache=sin_cache, cos_cache=cos_cache, slot_mapping=attn_metadata.slot_mapping, kv_cache=kv_cache, block_tables=attn_metadata.block_tables, block_group_size=env_blk_grp_size, o_proj_weight=self.fused_params['o_proj_weight'], o_proj_weight_scale=self.fused_params['o_proj_weight_scale'], o_proj_bias=None, o_proj_qzeros=None, seq_lens=attn_metadata.seq_lens, sm_scale=self.scaling, num_attention_heads=self.total_num_heads, num_key_value_heads=self.total_num_kv_heads, # head_dim=self.head_dim, is_decode=is_decode, flash_attention=is_decode, reduce_result=reduce_result, world_size=get_tp_group().world_size, rank=get_tp_group().rank_in_group, group_id=get_tp_group().group_id, dev_info=get_tp_group().rank_device_infos ) # if is_decode: # assert(1<0) # if residual is None: # output_ = attn_outs # else: # output_ = attn_outs[0] # print("starting tensor_model_parallel_all_reduce......................................, is_decode:", is_decode) if residual is None: attn_out = tensor_model_parallel_all_reduce(attn_outs) if not reduce_result else attn_outs else: res_out = attn_outs[1] attn_out = tensor_model_parallel_all_reduce(attn_outs[0]) if not reduce_result else attn_outs[0] # print("finished tensor_model_parallel_all_reduce......................................, is_decode:", is_decode) return attn_out, res_out elif use_unquan: # Determine if we should reduce the result inside the kernel # print(f"Qwen2Attention forward, is_decode: {is_decode}, kv_cache shape: {kv_cache.shape}, attn_metadata.seq_lens: {attn_metadata.seq_lens}") kv_cache = kv_cache.view(2, -1, 16, self.total_num_kv_heads // get_tp_group().world_size, self.head_dim) reduce_result = is_decode if cos_cache is None or sin_cache is None: cos_cache, sin_cache = get_cos_sin_cache(self.rotary_emb, attn_metadata, positions, is_decode) if residual is None: res_out = hidden_states if USE_FUSED_QWEN_ATTENTION: qkv_proj_qzeros = torch.Tensor() o_proj_bias = torch.Tensor() o_proj_qzeros = torch.Tensor() # Call the naive fused kernel ####################################call fusion kernel###################################### # print("starting fused attention kernel......................................, is_decode:", is_decode) # print('hidden_states', hidden_states.shape) # print('attn_metadata.slot_mapping,', attn_metadata.slot_mapping,) # print('attn_metadata.block_tables,', attn_metadata.block_tables,) # print('attn_metadata.seq_lens,', attn_metadata.seq_lens,) attn_outs = torch.vacc.fuse_atten_qwen2( hidden_states, residual, self.fused_params['input_layernorm_weight'], self.fused_params['qkv_proj_weight'], self.fused_params['qkv_proj_weight_scale'], # qkv_proj_bias, self.fused_params['qkv_proj_bias'], qkv_proj_qzeros, #self.fused_params['qkv_proj_qzeros'], sin_cache, cos_cache, attn_metadata.slot_mapping, kv_cache, attn_metadata.block_tables, env_blk_grp_size, self.fused_params['o_proj_weight'], self.fused_params['o_proj_weight_scale'], o_proj_bias, #self.fused_params['o_proj_bias'], o_proj_qzeros, #self.fused_params['o_proj_qzeros'], attn_metadata.seq_lens, self.scaling, self.total_num_heads, self.total_num_kv_heads, is_decode, is_decode, reduce_result, world_size=get_tp_group().world_size, rank=get_tp_group().rank_in_group, group_id=get_tp_group().group_id, dev_info=get_tp_group().rank_device_infos) else: attn_outs = vacc_fused_attn_qwen2_naive( hidden_states=hidden_states, residual=residual, hidden_states_norm_weight=self.fused_params['input_layernorm_weight'], ###ignore the group size, it is not used in the fused kernel # qkv_proj_group_size=get_gptq_group_size(self.qkv_proj), # o_proj_group_size=get_gptq_group_size(self.o_proj), ### qkv_proj_weight=self.fused_params['qkv_proj_weight'], qkv_proj_weight_scale=self.fused_params['qkv_proj_weight_scale'], qkv_proj_bias=self.fused_params['qkv_proj_bias'], qkv_proj_qzeros=None, sin_cache=sin_cache, cos_cache=cos_cache, slot_mapping=attn_metadata.slot_mapping, kv_cache=kv_cache, block_tables=attn_metadata.block_tables, block_group_size=env_blk_grp_size, o_proj_weight=self.fused_params['o_proj_weight'], o_proj_weight_scale=self.fused_params['o_proj_weight_scale'], o_proj_bias=None, o_proj_qzeros=None, seq_lens=attn_metadata.seq_lens, sm_scale=self.scaling, num_attention_heads=self.total_num_heads, num_key_value_heads=self.total_num_kv_heads, # head_dim=self.head_dim, is_decode=is_decode, flash_attention=is_decode, reduce_result=reduce_result, world_size=get_tp_group().world_size, rank=get_tp_group().rank_in_group, group_id=get_tp_group().group_id, dev_info=get_tp_group().rank_device_infos ) # if is_decode: # assert(1<0) # if residual is None: # output_ = attn_outs # else: # output_ = attn_outs[0] # print("starting tensor_model_parallel_all_reduce......................................, is_decode:", is_decode) if residual is None: attn_out = tensor_model_parallel_all_reduce(attn_outs) if not reduce_result else attn_outs # print(f"--------------------------attn_outs :{attn_out}") # print(f"--------------------------attn_outs shape :{attn_out.shape}") # print(f"--------------------------res_out :{res_out}") else: res_out = attn_outs[1] attn_out = tensor_model_parallel_all_reduce(attn_outs[0]) if not reduce_result else attn_outs[0] # print(f"--------------------------attn_outs :{attn_out}") # print(f"--------------------------attn_outs shape :{attn_out.shape}") # print(f"--------------------------res_out :{res_out}") # print("finished tensor_model_parallel_all_reduce......................................, is_decode:", is_decode) return attn_out, res_out else: # Fallback to original implementation # save_tensor_pth("hidden_states", hidden_states) qkv, _ = self.qkv_proj(hidden_states) # save_tensor_pth("qkv", qkv) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # save_tensor_pth("q", q) # save_tensor_pth("k", k) # save_tensor_pth("v", v) q, k = self.rotary_emb(positions, q, k) # save_tensor_pth("q_rot", q) # save_tensor_pth("k_rot", k) attn_output = self.attn(q, k, v) # save_tensor_pth("attn_output", attn_output) output, _ = self.o_proj(attn_output) # save_tensor_pth("output", output) return output class Qwen2DecoderLayer(nn.Module): def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: # forward_context: ForwardContext = get_forward_context() # attn_metadata = forward_context.attn_metadata # kv_cache = self.attn.kv_cache[forward_context.virtual_engine] # is_decode = attn_metadata.prefill_metadata is None use_gptq = hasattr(self.self_attn.qkv_proj, 'quant_method') and isinstance(self.self_attn.qkv_proj.quant_method, GPTQLinearMethod) use_unquan = hasattr(self.self_attn.qkv_proj, 'quant_method') and isinstance(self.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod) # if USE_FUSED_QWEN_ATTENTION and not is_decode: if use_gptq or use_unquan: # Self Attention with fused pre-layernorm # Pass the layernorm weight to the attention layer for the fused kernel self.self_attn.input_layernorm_weight = self.input_layernorm.weight self.self_attn.fused_params = {} self.self_attn.fused_params['input_layernorm_weight'] = self.input_layernorm.weight set_fused_params(self.self_attn.fused_params, self.self_attn.qkv_proj.quant_method, self.self_attn.qkv_proj, 'qkv_proj') set_fused_params(self.self_attn.fused_params, self.self_attn.o_proj.quant_method, self.self_attn.o_proj, 'o_proj') hidden_states, residual = self.self_attn( positions=positions, hidden_states=hidden_states, residual=residual, ) else: # Original non-fused path if residual is None: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm( hidden_states, residual) hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, ) # Fully Connected (MLP) part remains the same hidden_states, residual = self.post_attention_layernorm( hidden_states, residual) # save_tensor_pth("post_attention_layernorm_hidden_states", hidden_states) # save_tensor_pth("post_attention_layernorm_residual", residual) hidden_states = self.mlp(hidden_states) return hidden_states, residual