[NPU][1/N] NPU basic functions refactor and new modelslim quant type (#13359)
This commit is contained in:
Even Zhou
@@ -0,0 +1,184 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.allocator import PagedTokenToKVPoolAllocator
|
||||
from sglang.srt.utils import get_num_new_pages, next_power_of_2
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.mem_cache.memory_pool import KVCache
|
||||
|
||||
|
||||
def _alloc_extend_naive(
|
||||
prefix_lens,
|
||||
seq_lens,
|
||||
last_loc,
|
||||
free_pages,
|
||||
out_indices,
|
||||
page_size,
|
||||
device,
|
||||
):
|
||||
extend_lens = seq_lens - prefix_lens
|
||||
end_pos = torch.cumsum(extend_lens, 0)
|
||||
start_pos = end_pos - extend_lens
|
||||
num_new_pages = (seq_lens + page_size - 1) // page_size - (
|
||||
prefix_lens + page_size - 1
|
||||
) // page_size
|
||||
num_full_new_pages = (seq_lens) // page_size - (
|
||||
prefix_lens + page_size - 1
|
||||
) // page_size
|
||||
need_page = num_new_pages - num_full_new_pages
|
||||
end_new_pages = torch.cumsum(num_new_pages, 0)
|
||||
start_new_pages = end_new_pages - num_new_pages
|
||||
pos_in_page = torch.arange(page_size, device=device, dtype=torch.int32)
|
||||
for i in range(len(prefix_lens)):
|
||||
num1 = (
|
||||
min(
|
||||
seq_lens[i],
|
||||
(prefix_lens[i] + page_size - 1) // page_size * page_size,
|
||||
)
|
||||
- prefix_lens[i]
|
||||
)
|
||||
if num1:
|
||||
out_indices[start_pos[i] : start_pos[i] + num1] = (
|
||||
last_loc[i] + 1 + pos_in_page[:num1].view(-1)
|
||||
)
|
||||
|
||||
num2 = (
|
||||
seq_lens[i] // page_size - (prefix_lens[i] + page_size - 1) // page_size
|
||||
) * page_size
|
||||
if num2:
|
||||
pages = (
|
||||
free_pages[start_new_pages[i] : end_new_pages[i] - need_page[i]]
|
||||
* page_size
|
||||
)
|
||||
out_indices[start_pos[i] + num1 : start_pos[i] + num1 + num2] = (
|
||||
pages.view(-1, 1) + pos_in_page.view(1, -1)
|
||||
).view(-1)
|
||||
|
||||
num3 = seq_lens[i] - seq_lens[i] // page_size * page_size
|
||||
if num3:
|
||||
out_indices[end_pos[i] - num3 : end_pos[i]] = (
|
||||
free_pages[end_new_pages[i] - 1] * page_size + pos_in_page[:num3]
|
||||
).view(-1)
|
||||
|
||||
|
||||
class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
def __init__(
|
||||
self,
|
||||
size: int,
|
||||
page_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
kvcache: "KVCache",
|
||||
need_sort: bool,
|
||||
):
|
||||
super().__init__(size, page_size, dtype, device, kvcache, need_sort)
|
||||
self.roundup = page_size - 1
|
||||
|
||||
def alloc_extend(
|
||||
self,
|
||||
prefix_lens: torch.Tensor,
|
||||
prefix_lens_cpu: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
last_loc: torch.Tensor,
|
||||
extend_num_tokens: int,
|
||||
):
|
||||
if self.debug_mode:
|
||||
assert torch.all(
|
||||
(last_loc + 1) % self.page_size == prefix_lens % self.page_size
|
||||
)
|
||||
|
||||
num_new_pages = (
|
||||
(seq_lens + self.roundup) // self.page_size
|
||||
- (prefix_lens + self.roundup) // self.page_size
|
||||
).sum()
|
||||
num_new_pages_item = num_new_pages.item()
|
||||
if self.need_sort and num_new_pages_item > len(self.free_pages):
|
||||
self.merge_and_sort_free()
|
||||
|
||||
if num_new_pages_item > len(self.free_pages):
|
||||
return None
|
||||
|
||||
if num_new_pages_item < 200:
|
||||
from sgl_kernel_npu.mem_cache.allocator import alloc_extend_kernel
|
||||
|
||||
out_indices = torch.empty(
|
||||
(extend_num_tokens,),
|
||||
dtype=torch.int64,
|
||||
device=self.device,
|
||||
)
|
||||
max_num_extend_tokens = next_power_of_2(extend_num_tokens)
|
||||
bs = prefix_lens.shape[0]
|
||||
alloc_extend_kernel[(bs,)](
|
||||
prefix_lens,
|
||||
seq_lens,
|
||||
last_loc,
|
||||
self.free_pages,
|
||||
out_indices,
|
||||
next_power_of_2(bs),
|
||||
self.page_size,
|
||||
max_num_extend_tokens,
|
||||
)
|
||||
|
||||
else:
|
||||
out_indices = torch.empty(
|
||||
(extend_num_tokens,),
|
||||
dtype=torch.int32,
|
||||
device=self.device,
|
||||
)
|
||||
_alloc_extend_naive(
|
||||
prefix_lens,
|
||||
seq_lens,
|
||||
last_loc,
|
||||
self.free_pages,
|
||||
out_indices,
|
||||
self.page_size,
|
||||
self.device,
|
||||
)
|
||||
|
||||
if self.debug_mode:
|
||||
assert len(torch.unique(out_indices)) == len(out_indices)
|
||||
|
||||
self.free_pages = self.free_pages[num_new_pages_item:]
|
||||
return out_indices.int()
|
||||
|
||||
def alloc_decode(
|
||||
self,
|
||||
seq_lens: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
last_loc: torch.Tensor,
|
||||
):
|
||||
if self.debug_mode:
|
||||
assert torch.all(
|
||||
(last_loc + 2) % self.page_size == seq_lens % self.page_size
|
||||
)
|
||||
|
||||
num_new_pages = get_num_new_pages(
|
||||
seq_lens=seq_lens_cpu,
|
||||
page_size=self.page_size,
|
||||
decode=True,
|
||||
)
|
||||
|
||||
if num_new_pages > len(self.free_pages):
|
||||
self.merge_and_sort_free()
|
||||
|
||||
if num_new_pages > len(self.free_pages):
|
||||
return None
|
||||
|
||||
need_new_pages = (seq_lens % self.page_size == 1).int()
|
||||
end_new_pages = torch.cumsum(need_new_pages, 0)
|
||||
start_new_pages = end_new_pages - need_new_pages
|
||||
if num_new_pages == 0:
|
||||
out_indices = last_loc + 1
|
||||
else:
|
||||
out_indices = (last_loc + 1) * (1 - need_new_pages) + self.free_pages[
|
||||
start_new_pages
|
||||
] * self.page_size * need_new_pages
|
||||
|
||||
if self.debug_mode:
|
||||
assert len(torch.unique(out_indices)) == len(out_indices)
|
||||
|
||||
self.free_pages = self.free_pages[num_new_pages:]
|
||||
return out_indices.int()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,388 @@
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
|
||||
from sglang.srt.utils import get_bool_env_var
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def is_mla_preprocess_enabled() -> bool:
|
||||
return get_bool_env_var("SGLANG_NPU_USE_MLAPO")
|
||||
|
||||
|
||||
def round_up(val: int, align: int) -> int:
|
||||
if align == 0:
|
||||
return 0
|
||||
return -(val // -align) * align
|
||||
|
||||
|
||||
def transdata(nd_mat, block_size: tuple = (16, 16)):
|
||||
r = round_up(nd_mat.shape[0], block_size[0])
|
||||
c = round_up(nd_mat.shape[1], block_size[1])
|
||||
r_pad = r - nd_mat.shape[0]
|
||||
c_pad = c - nd_mat.shape[1]
|
||||
nd_mat = F.pad(nd_mat, ((0, r_pad, 0, c_pad)))
|
||||
nz_mat = torch.permute(
|
||||
torch.reshape(
|
||||
nd_mat,
|
||||
(r // block_size[0], block_size[0], c // block_size[1], block_size[1]),
|
||||
),
|
||||
[2, 0, 1, 3],
|
||||
)
|
||||
nz_mat = torch.reshape(
|
||||
nz_mat, (nz_mat.shape[0], nz_mat.shape[1] * nz_mat.shape[2], nz_mat.shape[3])
|
||||
)
|
||||
return nz_mat
|
||||
|
||||
|
||||
def trans_rope_weight(weight, rope_dim):
|
||||
weight_1 = weight[..., -rope_dim::2, :].contiguous()
|
||||
weight_2 = weight[..., -rope_dim + 1 :: 2, :].contiguous()
|
||||
weight[..., -rope_dim:, :] = torch.cat([weight_1, weight_2], dim=-2)
|
||||
|
||||
return weight.contiguous()
|
||||
|
||||
|
||||
class NPUFusedMLAPreprocess(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
fused_qkv_a_proj_with_mqa,
|
||||
q_a_layernorm,
|
||||
kv_a_layernorm,
|
||||
q_b_proj,
|
||||
w_kc,
|
||||
rotary_emb,
|
||||
layer_id,
|
||||
num_local_heads,
|
||||
qk_nope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
quant_config: Optional["QuantizationConfig"] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.qkv_a_proj = fused_qkv_a_proj_with_mqa
|
||||
self.q_a_layernorm = q_a_layernorm
|
||||
self.kv_a_layernorm = kv_a_layernorm
|
||||
self.q_b_proj = q_b_proj
|
||||
self.w_kc = w_kc.contiguous()
|
||||
self.rotary_emb = rotary_emb
|
||||
self.layer_id = layer_id
|
||||
self.quant_config = quant_config
|
||||
self.has_preprocess_weights = False
|
||||
self.dtype = None
|
||||
|
||||
self.q_lora_rank = self.q_b_proj.input_size # 1536
|
||||
self.kv_lora_rank = self.kv_a_layernorm.hidden_size # 512
|
||||
self.num_local_heads = num_local_heads # tp
|
||||
self.qk_nope_head_dim = qk_nope_head_dim # 128
|
||||
self.qk_rope_head_dim = qk_rope_head_dim # 64
|
||||
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
||||
|
||||
def preprocess_weights(self, hidden_states):
|
||||
self.dummy = torch.empty(
|
||||
(hidden_states.shape[-1]),
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
self.qkv_a_proj_input_offset = self.qkv_a_proj.input_offset.to(dtype=torch.int8)
|
||||
self.q_b_proj_input_offset = self.q_b_proj.input_offset.to(dtype=torch.int8)
|
||||
|
||||
# matmul_0 weight [7168, 2112]
|
||||
fused_qkv_a_proj_with_mqa_weight_q = self.qkv_a_proj.weight.data[
|
||||
:, : self.q_lora_rank
|
||||
].clone() # [7168, 1536]
|
||||
fused_qkv_a_proj_with_mqa_weight_kv = self.qkv_a_proj.weight.data[
|
||||
:, self.q_lora_rank :
|
||||
].clone() # [7168, 576]
|
||||
# rope fit
|
||||
fused_qkv_a_proj_with_mqa_weight_kv_t = (
|
||||
fused_qkv_a_proj_with_mqa_weight_kv.t().contiguous()
|
||||
)
|
||||
fused_qkv_a_proj_with_mqa_weight_kv_t = trans_rope_weight(
|
||||
fused_qkv_a_proj_with_mqa_weight_kv_t, self.qk_rope_head_dim
|
||||
)
|
||||
fused_qkv_a_proj_with_mqa_weight_kv = (
|
||||
fused_qkv_a_proj_with_mqa_weight_kv_t.t().contiguous()
|
||||
)
|
||||
# cat nz
|
||||
fused_qkv_a_proj_with_mqa_weight_new = torch.cat(
|
||||
(fused_qkv_a_proj_with_mqa_weight_kv, fused_qkv_a_proj_with_mqa_weight_q),
|
||||
dim=-1,
|
||||
)
|
||||
fused_qkv_a_proj_with_mqa_weight = (
|
||||
fused_qkv_a_proj_with_mqa_weight_new.t().contiguous()
|
||||
)
|
||||
fused_qkv_a_proj_with_mqa_weight_nz = (
|
||||
transdata(fused_qkv_a_proj_with_mqa_weight, block_size=(16, 32))
|
||||
.unsqueeze(0)
|
||||
.contiguous()
|
||||
)
|
||||
self.qkv_a_proj_weight_nz = npu_format_cast(fused_qkv_a_proj_with_mqa_weight_nz)
|
||||
|
||||
# matmul_0 deq_scale [2112]
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_q = self.qkv_a_proj.deq_scale.data[
|
||||
: self.q_lora_rank
|
||||
].clone() # [7168, 1536]
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_kv = self.qkv_a_proj.deq_scale.data[
|
||||
self.q_lora_rank :
|
||||
].clone() # [7168, 576]
|
||||
# rope fit
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_kv = (
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_kv.reshape(
|
||||
self.kv_lora_rank + self.qk_rope_head_dim, -1
|
||||
).contiguous()
|
||||
)
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_kv = trans_rope_weight(
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_kv, self.qk_rope_head_dim
|
||||
)
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_kv = (
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_kv.view(
|
||||
self.kv_lora_rank + self.qk_rope_head_dim
|
||||
).contiguous()
|
||||
)
|
||||
self.qkv_a_proj_deq_scale_kvq = torch.cat(
|
||||
(
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_kv,
|
||||
fused_qkv_a_proj_with_mqa_deq_scale_q,
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
# matmul_0 quant_bias [2112]
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_q = self.qkv_a_proj.quant_bias.data[
|
||||
: self.q_lora_rank
|
||||
].clone() # [7168, 1536]
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_kv = self.qkv_a_proj.quant_bias.data[
|
||||
self.q_lora_rank :
|
||||
].clone() # [7168, 576]
|
||||
# rope fit
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_kv = (
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_kv.reshape(
|
||||
self.kv_lora_rank + self.qk_rope_head_dim, -1
|
||||
).contiguous()
|
||||
)
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_kv = trans_rope_weight(
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_kv, self.qk_rope_head_dim
|
||||
)
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_kv = (
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_kv.view(
|
||||
self.kv_lora_rank + self.qk_rope_head_dim
|
||||
).contiguous()
|
||||
)
|
||||
self.qkv_a_proj_quant_bias_kvq = torch.cat(
|
||||
(
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_kv,
|
||||
fused_qkv_a_proj_with_mqa_quant_bias_q,
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
# matmul_1 weight [1536, num_head * 192]
|
||||
q_b_proj_weight = self.q_b_proj.weight.data.clone()
|
||||
q_b_proj_weight = q_b_proj_weight.t().reshape(
|
||||
self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
|
||||
)
|
||||
q_b_proj_weight = trans_rope_weight(q_b_proj_weight, self.qk_rope_head_dim)
|
||||
q_b_proj_weight = q_b_proj_weight.reshape(
|
||||
self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim), -1
|
||||
)
|
||||
q_b_proj_weight_nz = (
|
||||
transdata(q_b_proj_weight, block_size=(16, 32)).unsqueeze(0).contiguous()
|
||||
)
|
||||
self.q_b_proj_weight_nz = npu_format_cast(q_b_proj_weight_nz)
|
||||
|
||||
# matmul_1 deq_scale [num_head * 192]
|
||||
q_b_proj_deq_scale = self.q_b_proj.deq_scale.data.clone()
|
||||
q_b_proj_deq_scale = q_b_proj_deq_scale.reshape(
|
||||
self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
|
||||
)
|
||||
q_b_proj_deq_scale = trans_rope_weight(
|
||||
q_b_proj_deq_scale, self.qk_rope_head_dim
|
||||
)
|
||||
self.q_b_proj_deq_scale = q_b_proj_deq_scale.reshape(
|
||||
self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
|
||||
)
|
||||
|
||||
# matmul_1 quant_bias [num_head * 192]
|
||||
q_b_proj_quant_bias = self.q_b_proj.quant_bias.data.clone()
|
||||
q_b_proj_quant_bias = q_b_proj_quant_bias.reshape(
|
||||
self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
|
||||
)
|
||||
q_b_proj_quant_bias = trans_rope_weight(
|
||||
q_b_proj_quant_bias, self.qk_rope_head_dim
|
||||
)
|
||||
self.q_b_proj_quant_bias = q_b_proj_quant_bias.reshape(
|
||||
self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
|
||||
)
|
||||
|
||||
def get_sin_cos(self, positions):
|
||||
cos_sin = self.rotary_emb.cos_sin_cache[positions]
|
||||
cos, sin = cos_sin.chunk(2, dim=-1)
|
||||
cos = cos.repeat(1, 2)
|
||||
sin = sin.repeat(1, 2)
|
||||
return cos, sin
|
||||
|
||||
def get_kv_cache_and_cache_idx(self, forward_batch):
|
||||
k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(self.layer_id)
|
||||
slot_mapping = forward_batch.out_cache_loc.to(dtype=torch.int32)
|
||||
return k_cache, v_cache, slot_mapping
|
||||
|
||||
def forward_absorb_prepare_npu_rms_norm_cache(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch,
|
||||
zero_allocator,
|
||||
):
|
||||
bsz, _ = hidden_states.view(-1, hidden_states.shape[-1]).shape
|
||||
self.dtype = hidden_states.dtype
|
||||
self.cos, self.sin = self.get_sin_cos(positions)
|
||||
self.kvCache, self.kvCacheRope, self.slotmapping = (
|
||||
self.get_kv_cache_and_cache_idx(forward_batch)
|
||||
)
|
||||
|
||||
if not self.has_preprocess_weights:
|
||||
self.has_preprocess_weights = True
|
||||
|
||||
cos, sin = self.cos, self.sin
|
||||
|
||||
if self.q_lora_rank is not None:
|
||||
fused_qkv_a_proj_out = self.qkv_a_proj(hidden_states)[0]
|
||||
q_lowrank, latent_cache = fused_qkv_a_proj_out.split(
|
||||
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
|
||||
)
|
||||
q = self.q_a_layernorm(q_lowrank)
|
||||
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
|
||||
else:
|
||||
q = self.q_proj(hidden_states)[0].view(
|
||||
-1, self.num_local_heads, self.qk_head_dim
|
||||
)
|
||||
latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
|
||||
|
||||
q_nope, q_pe = torch.split(
|
||||
q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
|
||||
) # b*s,n,d
|
||||
|
||||
q_nope = q_nope.view(-1, self.num_local_heads, self.qk_nope_head_dim)
|
||||
q_nope = torch.matmul(q_nope.transpose(0, 1), self.w_kc).transpose(0, 1)
|
||||
|
||||
q_pe = q_pe.view(-1, self.num_local_heads, 1, self.qk_rope_head_dim)
|
||||
cos = cos.view(-1, 1, 1, self.qk_rope_head_dim)
|
||||
sin = sin.view(-1, 1, 1, self.qk_rope_head_dim)
|
||||
q_pe = torch.ops.npu.npu_interleave_rope(q_pe, cos, sin) # (B,N,S,D)
|
||||
q_pe = q_pe.view(cos.shape[0], self.num_local_heads, self.qk_rope_head_dim)
|
||||
|
||||
latent_cache = latent_cache.view(
|
||||
-1, 1, 1, self.kv_lora_rank + self.qk_rope_head_dim
|
||||
) # (B*S,N,1,D)
|
||||
|
||||
cache_mode = "PA_BNSD"
|
||||
self.kvCache = self.kvCache.view(
|
||||
-1,
|
||||
forward_batch.attn_backend.page_size,
|
||||
1,
|
||||
forward_batch.attn_backend.kv_lora_rank,
|
||||
)
|
||||
self.kvCacheRope = self.kvCacheRope.view(
|
||||
-1,
|
||||
forward_batch.attn_backend.page_size,
|
||||
1,
|
||||
forward_batch.attn_backend.qk_rope_head_dim,
|
||||
)
|
||||
k_rope, k_nope, _, _ = torch.ops.npu.npu_kv_rmsnorm_rope_cache(
|
||||
latent_cache,
|
||||
self.kv_a_layernorm.weight,
|
||||
cos,
|
||||
sin,
|
||||
self.slotmapping.to(torch.int64),
|
||||
self.kvCacheRope,
|
||||
self.kvCache,
|
||||
epsilon=self.kv_a_layernorm.variance_epsilon,
|
||||
cache_mode=cache_mode,
|
||||
)
|
||||
|
||||
return (q_pe, k_rope, q_nope, k_nope, forward_batch, zero_allocator, positions)
|
||||
|
||||
def forward_mlapo(self, positions, hidden_states, forward_batch, zero_allocator):
|
||||
input_dtype = hidden_states.dtype
|
||||
if not self.has_preprocess_weights:
|
||||
self.preprocess_weights(hidden_states)
|
||||
self.has_preprocess_weights = True
|
||||
self.dtype = hidden_states.dtype
|
||||
|
||||
cos, sin = self.get_sin_cos(positions)
|
||||
k_cache, v_cache, slot_mapping = self.get_kv_cache_and_cache_idx(forward_batch)
|
||||
|
||||
q_nope_out = torch.empty(
|
||||
(hidden_states.shape[0], self.w_kc.shape[0], k_cache.shape[-1]),
|
||||
dtype=input_dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
q_rope_out = torch.empty(
|
||||
(hidden_states.shape[0], self.w_kc.shape[0], v_cache.shape[-1]),
|
||||
dtype=input_dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
# TODO: dummy inputs to be removed
|
||||
# https://github.com/sgl-project/sgl-kernel-npu/issues/78
|
||||
torch.ops.npu.mla_preprocess(
|
||||
hidden_states,
|
||||
self.dummy,
|
||||
self.dummy,
|
||||
self.qkv_a_proj_weight_nz,
|
||||
self.qkv_a_proj_deq_scale_kvq,
|
||||
self.q_a_layernorm.weight,
|
||||
self.q_a_layernorm.bias,
|
||||
self.q_b_proj_weight_nz,
|
||||
self.q_b_proj_deq_scale,
|
||||
self.kv_a_layernorm.weight,
|
||||
cos,
|
||||
sin,
|
||||
self.w_kc,
|
||||
k_cache,
|
||||
v_cache,
|
||||
slot_mapping,
|
||||
quant_scale0=self.qkv_a_proj.input_scale,
|
||||
quant_offset0=self.qkv_a_proj_input_offset,
|
||||
bias0=self.qkv_a_proj_quant_bias_kvq,
|
||||
quant_scale1=self.q_b_proj.input_scale,
|
||||
quant_offset1=self.q_b_proj_input_offset,
|
||||
bias1=self.q_b_proj_quant_bias,
|
||||
cache_mode="krope_ctkv",
|
||||
quant_mode="per_tensor_quant_asymm",
|
||||
q_out0=q_nope_out,
|
||||
kv_cache_out0=k_cache,
|
||||
q_out1=q_rope_out,
|
||||
kv_cache_out1=v_cache,
|
||||
)
|
||||
return (
|
||||
q_rope_out,
|
||||
v_cache,
|
||||
q_nope_out,
|
||||
k_cache,
|
||||
forward_batch,
|
||||
zero_allocator,
|
||||
positions,
|
||||
)
|
||||
|
||||
def forward(self, positions, hidden_states, forward_batch, zero_allocator):
|
||||
assert self.quant_config and self.quant_config.get_name() == "modelslim"
|
||||
# route by `qkv_a_proj` quant type as MTP layers can be unquantized
|
||||
_is_w8a8 = (
|
||||
hasattr(self.qkv_a_proj.quant_method, "quant_config")
|
||||
and self.qkv_a_proj.quant_method.quant_config.get_name() == "modelslim"
|
||||
)
|
||||
if _is_w8a8:
|
||||
return self.forward_mlapo(
|
||||
positions, hidden_states, forward_batch, zero_allocator
|
||||
)
|
||||
else:
|
||||
return self.forward_absorb_prepare_npu_rms_norm_cache(
|
||||
positions, hidden_states, forward_batch, zero_allocator
|
||||
)
|
||||
@@ -0,0 +1,54 @@
|
||||
import torch
|
||||
|
||||
cmo_stream = None
|
||||
|
||||
|
||||
def get_cmo_stream():
|
||||
"""
|
||||
Cache Management Operation(CMO).
|
||||
Launch a new stream to prefetch the weight of matmul when running other
|
||||
AIV or communication kernels, aiming to overlap the memory access time.
|
||||
"""
|
||||
global cmo_stream
|
||||
return cmo_stream
|
||||
|
||||
|
||||
def set_cmo_stream(stream):
|
||||
global cmo_stream
|
||||
cmo_stream = stream
|
||||
|
||||
|
||||
def prepare_weight_cache(handle, cache, PREFETCH_MAX_SIZE=1000000000):
|
||||
"""
|
||||
PREFETCH_MAX_SIZE: maximum size (bytes) for each prefetch operation.
|
||||
This affects the time spent in prefetch:
|
||||
time ≈ PREFETCH_MAX_SIZE / system_bandwidth
|
||||
"""
|
||||
import torch_npu
|
||||
|
||||
stream = get_cmo_stream()
|
||||
if stream is None:
|
||||
stream = torch.npu.Stream()
|
||||
set_cmo_stream(stream)
|
||||
stream.wait_stream(torch.npu.current_stream())
|
||||
with torch.npu.stream(stream):
|
||||
if isinstance(cache, list):
|
||||
for weight in cache:
|
||||
torch_npu.npu_prefetch(
|
||||
weight,
|
||||
handle,
|
||||
PREFETCH_MAX_SIZE,
|
||||
)
|
||||
else:
|
||||
torch_npu.npu_prefetch(
|
||||
cache,
|
||||
handle,
|
||||
PREFETCH_MAX_SIZE,
|
||||
)
|
||||
|
||||
|
||||
def wait_cmo_stream():
|
||||
stream = get_cmo_stream()
|
||||
if stream is not None:
|
||||
cur_stream = torch.npu.current_stream()
|
||||
cur_stream.wait_stream(stream)
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
# Copyright 2024-2025 SGLang Team
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Run the model with npu graph and torch.compile."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.model_config import is_deepseek_nsa
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
|
||||
EAGLEDraftExtendCudaGraphRunner,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
|
||||
|
||||
class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
|
||||
def __init__(self, eagle_worker: EAGLEWorker):
|
||||
super().__init__(eagle_worker)
|
||||
|
||||
def _create_graph(self):
|
||||
return torch.npu.NPUGraph()
|
||||
|
||||
def _capture_init(self, run_once_fn):
|
||||
for _ in range(2):
|
||||
torch.npu.synchronize()
|
||||
self.model_runner.tp_group.barrier()
|
||||
run_once_fn()
|
||||
|
||||
def _capture_graph(self, graph, pool, stream, run_once_fn):
|
||||
with torch.npu.graph(
|
||||
graph, pool=pool, stream=stream, auto_dispatch_capture=True
|
||||
):
|
||||
out = run_once_fn()
|
||||
return out
|
||||
|
||||
def _replay_update(self, seq_lens):
|
||||
self.graphs[self.bs].update(
|
||||
cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}]
|
||||
)
|
||||
|
||||
def _replay(self, forward_batch: ForwardBatch):
|
||||
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
|
||||
seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
|
||||
self.bs - self.raw_bs
|
||||
)
|
||||
thread = threading.Thread(target=self._replay_update, args=(seq_lens,))
|
||||
thread.start()
|
||||
self.graphs[self.bs].replay()
|
||||
thread.join()
|
||||
else:
|
||||
self.graphs[self.bs].replay()
|
||||
@@ -0,0 +1,81 @@
|
||||
# Copyright 2025 SGLang Team
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
""" Run the model with npu graph and torch.compile """
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.model_config import is_deepseek_nsa
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
|
||||
EAGLEDraftCudaGraphRunner,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
|
||||
from sglang.srt.utils import is_npu
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if is_npu():
|
||||
torch.cuda.CUDAGraph = torch.npu.NPUGraph
|
||||
torch.cuda.synchronize = torch.npu.synchronize
|
||||
torch.cuda.graph = torch.npu.graph
|
||||
torch.cuda.stream = torch.npu.stream
|
||||
torch.cuda.Stream = torch.npu.Stream
|
||||
torch.cuda.current_stream = torch.npu.current_stream
|
||||
|
||||
|
||||
class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
|
||||
def __init__(self, eagle_worker: EAGLEWorker):
|
||||
super().__init__(eagle_worker)
|
||||
|
||||
def _create_graph(self):
|
||||
return torch.npu.NPUGraph()
|
||||
|
||||
def _capture_init(self, run_once_fn):
|
||||
for _ in range(2):
|
||||
torch.npu.synchronize()
|
||||
self.model_runner.tp_group.barrier()
|
||||
run_once_fn()
|
||||
|
||||
def _capture_graph(self, graph, pool, stream, run_once_fn):
|
||||
with torch.npu.graph(
|
||||
graph, pool=pool, stream=stream, auto_dispatch_capture=True
|
||||
):
|
||||
out = run_once_fn()
|
||||
return out
|
||||
|
||||
def _replay_update(self, seq_lens):
|
||||
self.graphs[self.bs].update(
|
||||
cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}]
|
||||
)
|
||||
|
||||
def _replay(self, forward_batch: ForwardBatch):
|
||||
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
|
||||
seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
|
||||
self.bs - self.raw_bs
|
||||
)
|
||||
thread = threading.Thread(target=self._replay_update, args=(seq_lens,))
|
||||
thread.start()
|
||||
self.graphs[self.bs].replay()
|
||||
thread.join()
|
||||
else:
|
||||
self.graphs[self.bs].replay()
|
||||
@@ -0,0 +1,172 @@
|
||||
# Copyright 2023-2024 SGLang Team
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Run the model with npu graph and torch.compile."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Dict, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
|
||||
from sglang.srt.layers.dp_attention import get_attention_tp_size
|
||||
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
|
||||
from sglang.srt.utils import is_npu
|
||||
|
||||
is_npu = is_npu()
|
||||
|
||||
if is_npu:
|
||||
import torch_npu
|
||||
from torch_npu.profiler import ProfilerActivity, profile
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
|
||||
|
||||
|
||||
class NPUGraphRunner(CudaGraphRunner):
|
||||
"""A NPUGraphRunner runs the forward pass of a model with npu graph and torch.compile."""
|
||||
|
||||
def __init__(self, model_runner: ModelRunner):
|
||||
super().__init__(model_runner)
|
||||
self.update_attr_name = None
|
||||
self.update_attr_type = None
|
||||
self.model_runner = model_runner
|
||||
self._init_arch_map()
|
||||
|
||||
def _init_arch_map(self):
|
||||
self.attr_name: Dict[str, str] = {
|
||||
AttentionArch.MLA: "actual_seq_lengths_kv",
|
||||
AttentionArch.MHA: "context_lens",
|
||||
}
|
||||
self.attr_type: Dict[str, Union[list, torch.Tensor]] = {
|
||||
AttentionArch.MLA: [],
|
||||
AttentionArch.MHA: torch.Tensor(),
|
||||
}
|
||||
|
||||
def _create_device_graph(self):
|
||||
return torch.npu.NPUGraph()
|
||||
|
||||
def _capture_graph(self, graph, pool, stream, run_once_fn):
|
||||
with torch.npu.graph(
|
||||
graph,
|
||||
pool=pool,
|
||||
stream=stream,
|
||||
auto_dispatch_capture=True,
|
||||
):
|
||||
out = run_once_fn()
|
||||
return out
|
||||
|
||||
def _get_update_attr_name(self, model_runner):
|
||||
if self.bs < get_attention_tp_size():
|
||||
return self.attr_name[AttentionArch.MLA]
|
||||
return self.attr_name[model_runner.model_config.attention_arch]
|
||||
|
||||
def _get_update_attr_type(self, model_runner):
|
||||
if self.bs < get_attention_tp_size():
|
||||
return self.attr_type[AttentionArch.MLA]
|
||||
return self.attr_type[model_runner.model_config.attention_arch]
|
||||
|
||||
def _update_inputs(self, seq_lens):
|
||||
if isinstance(self.update_attr_type, torch.Tensor):
|
||||
seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
|
||||
|
||||
self.graphs[self.bs].update(
|
||||
cpu_update_input=[{self.update_attr_name: seq_lens}]
|
||||
)
|
||||
|
||||
def _cache_loc_dtype(self):
|
||||
return torch.int32
|
||||
|
||||
def _init_profile_context_and_memory_record(self):
|
||||
output_dir = os.path.join(
|
||||
os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp"), "graph_capture_profile"
|
||||
)
|
||||
if not Path(output_dir).exists():
|
||||
Path(output_dir).mkdir(parents=True, exist_ok=True)
|
||||
logger.info(
|
||||
f"Profiling starts for graph capture for NPU. Traces will be saved to: {output_dir}"
|
||||
)
|
||||
experimental_config = torch_npu.profiler._ExperimentalConfig(
|
||||
export_type=[torch_npu.profiler.ExportType.Text],
|
||||
profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
|
||||
)
|
||||
profile_context = profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
|
||||
record_shapes=True,
|
||||
profile_memory=True,
|
||||
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
|
||||
output_dir, async_mode=True
|
||||
),
|
||||
experimental_config=experimental_config,
|
||||
)
|
||||
return profile_context
|
||||
|
||||
def _post_process_after_profile(self, prof_context):
|
||||
# for NPU, profile data will be saved to disk for further analysis.
|
||||
pass
|
||||
|
||||
def replay(
|
||||
self,
|
||||
forward_batch: ForwardBatch,
|
||||
skip_attn_backend_init: bool = False,
|
||||
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
|
||||
if not skip_attn_backend_init:
|
||||
self.replay_prepare(forward_batch, pp_proxy_tensors)
|
||||
else:
|
||||
# In speculative decoding, these two fields are still needed.
|
||||
self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
|
||||
self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
|
||||
|
||||
self.update_attr_name = self._get_update_attr_name(self.model_runner)
|
||||
self.update_attr_type = self._get_update_attr_type(self.model_runner)
|
||||
# Replay
|
||||
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
|
||||
if forward_batch.forward_mode.is_target_verify():
|
||||
seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
|
||||
seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
|
||||
else:
|
||||
seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
|
||||
self.bs - self.raw_bs
|
||||
)
|
||||
thread = threading.Thread(target=self._update_inputs, args=(seq_lens,))
|
||||
thread.start()
|
||||
self.graphs[self.bs].replay()
|
||||
thread.join()
|
||||
else:
|
||||
self.graphs[self.bs].replay()
|
||||
|
||||
output = self.output_buffers[self.bs]
|
||||
if isinstance(output, LogitsProcessorOutput):
|
||||
return LogitsProcessorOutput(
|
||||
next_token_logits=output.next_token_logits[: self.raw_num_token],
|
||||
hidden_states=(
|
||||
output.hidden_states[: self.raw_num_token]
|
||||
if output.hidden_states is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
else:
|
||||
assert isinstance(output, PPProxyTensors)
|
||||
return PPProxyTensors({k: v[: self.bs] for k, v in output.tensors.items()})
|
||||
@@ -0,0 +1,300 @@
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
import torch_npu
|
||||
|
||||
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
|
||||
from sglang.srt.mem_cache.memory_pool import (
|
||||
MHATokenToKVPool,
|
||||
MLATokenToKVPool,
|
||||
get_tensor_size_bytes,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.radix_attention import RadixAttention
|
||||
|
||||
|
||||
class NPUMHATokenToKVPool(MHATokenToKVPool):
|
||||
|
||||
def _create_buffers(self):
|
||||
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
|
||||
# [size, head_num, head_dim] for each layer
|
||||
# The padded slot 0 is used for writing dummy outputs from padded tokens.
|
||||
# Continuous memory improves the efficiency of Ascend`s transmission backend,
|
||||
# while other backends remain unchanged.
|
||||
self.kv_buffer = torch.zeros(
|
||||
(
|
||||
2,
|
||||
self.layer_num,
|
||||
self.size // self.page_size + 1,
|
||||
self.page_size,
|
||||
self.head_num,
|
||||
self.head_dim,
|
||||
),
|
||||
dtype=self.store_dtype,
|
||||
device=self.device,
|
||||
)
|
||||
self.k_buffer = self.kv_buffer[0]
|
||||
self.v_buffer = self.kv_buffer[1]
|
||||
|
||||
# for disagg
|
||||
def get_contiguous_buf_infos(self):
|
||||
# layer_num x [seq_len, head_num, head_dim]
|
||||
# layer_num x [page_num, page_size, head_num, head_dim]
|
||||
kv_data_ptrs = [
|
||||
self.get_key_buffer(i).data_ptr()
|
||||
for i in range(self.start_layer, self.start_layer + self.layer_num)
|
||||
] + [
|
||||
self.get_value_buffer(i).data_ptr()
|
||||
for i in range(self.start_layer, self.start_layer + self.layer_num)
|
||||
]
|
||||
kv_data_lens = [
|
||||
self.get_key_buffer(i).nbytes
|
||||
for i in range(self.start_layer, self.start_layer + self.layer_num)
|
||||
] + [
|
||||
self.get_value_buffer(i).nbytes
|
||||
for i in range(self.start_layer, self.start_layer + self.layer_num)
|
||||
]
|
||||
kv_item_lens = [
|
||||
self.get_key_buffer(i)[0].nbytes
|
||||
for i in range(self.start_layer, self.start_layer + self.layer_num)
|
||||
] + [
|
||||
self.get_value_buffer(i)[0].nbytes
|
||||
for i in range(self.start_layer, self.start_layer + self.layer_num)
|
||||
]
|
||||
return kv_data_ptrs, kv_data_lens, kv_item_lens
|
||||
|
||||
def set_kv_buffer(
|
||||
self,
|
||||
layer: "RadixAttention",
|
||||
loc: torch.Tensor,
|
||||
cache_k: torch.Tensor,
|
||||
cache_v: torch.Tensor,
|
||||
k_scale: Optional[float] = None,
|
||||
v_scale: Optional[float] = None,
|
||||
layer_id_override: Optional[int] = None,
|
||||
):
|
||||
if layer_id_override is not None:
|
||||
layer_id = layer_id_override
|
||||
else:
|
||||
layer_id = layer.layer_id
|
||||
if cache_k.dtype != self.dtype:
|
||||
if k_scale is not None:
|
||||
cache_k.div_(k_scale)
|
||||
if v_scale is not None:
|
||||
cache_v.div_(v_scale)
|
||||
cache_k = cache_k.to(self.dtype)
|
||||
cache_v = cache_v.to(self.dtype)
|
||||
|
||||
if self.store_dtype != self.dtype:
|
||||
cache_k = cache_k.view(self.store_dtype)
|
||||
cache_v = cache_v.view(self.store_dtype)
|
||||
|
||||
torch_npu._npu_reshape_and_cache(
|
||||
key=cache_k,
|
||||
value=cache_v,
|
||||
key_cache=self.k_buffer[layer_id - self.start_layer].view(
|
||||
-1, self.page_size, self.head_num, self.head_dim
|
||||
),
|
||||
value_cache=self.v_buffer[layer_id - self.start_layer].view(
|
||||
-1, self.page_size, self.head_num, self.head_dim
|
||||
),
|
||||
slot_indices=loc,
|
||||
)
|
||||
|
||||
|
||||
class NPUMLATokenToKVPool(MLATokenToKVPool):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size: int,
|
||||
page_size: int,
|
||||
dtype: torch.dtype,
|
||||
kv_lora_rank: int,
|
||||
qk_rope_head_dim: int,
|
||||
index_head_dim: Optional[int],
|
||||
layer_num: int,
|
||||
device: str,
|
||||
enable_memory_saver: bool,
|
||||
start_layer: Optional[int] = None,
|
||||
end_layer: Optional[int] = None,
|
||||
):
|
||||
super(MLATokenToKVPool, self).__init__(
|
||||
size,
|
||||
page_size,
|
||||
dtype,
|
||||
layer_num,
|
||||
device,
|
||||
enable_memory_saver,
|
||||
start_layer,
|
||||
end_layer,
|
||||
)
|
||||
|
||||
self.kv_lora_rank = kv_lora_rank
|
||||
self.qk_rope_head_dim = qk_rope_head_dim
|
||||
self.index_head_dim = index_head_dim
|
||||
|
||||
self.custom_mem_pool = None
|
||||
|
||||
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
|
||||
# The padded slot 0 is used for writing dummy outputs from padded tokens.
|
||||
self.k_buffer = torch.zeros(
|
||||
(
|
||||
layer_num,
|
||||
self.size // self.page_size + 1,
|
||||
self.page_size,
|
||||
1,
|
||||
self.kv_lora_rank,
|
||||
),
|
||||
dtype=self.store_dtype,
|
||||
device=self.device,
|
||||
)
|
||||
self.v_buffer = torch.zeros(
|
||||
(
|
||||
layer_num,
|
||||
self.size // self.page_size + 1,
|
||||
self.page_size,
|
||||
1,
|
||||
self.qk_rope_head_dim,
|
||||
),
|
||||
dtype=self.store_dtype,
|
||||
device=self.device,
|
||||
)
|
||||
if self.index_head_dim is not None:
|
||||
self.index_k_buffer = torch.zeros(
|
||||
(
|
||||
layer_num,
|
||||
self.size // self.page_size + 1,
|
||||
self.page_size,
|
||||
1,
|
||||
self.index_head_dim,
|
||||
),
|
||||
dtype=self.store_dtype,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
self._finalize_allocation_log(size)
|
||||
|
||||
def get_kv_size_bytes(self):
|
||||
assert hasattr(self, "k_buffer")
|
||||
assert hasattr(self, "v_buffer")
|
||||
kv_size_bytes = 0
|
||||
for k_cache in self.k_buffer:
|
||||
kv_size_bytes += get_tensor_size_bytes(k_cache)
|
||||
for v_cache in self.v_buffer:
|
||||
kv_size_bytes += get_tensor_size_bytes(v_cache)
|
||||
if self.index_head_dim is not None:
|
||||
assert hasattr(self, "index_k_buffer")
|
||||
for index_k_cache in self.index_k_buffer:
|
||||
kv_size_bytes += get_tensor_size_bytes(index_k_cache)
|
||||
return kv_size_bytes
|
||||
|
||||
def get_kv_buffer(self, layer_id: int):
|
||||
if self.layer_transfer_counter is not None:
|
||||
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
|
||||
return (
|
||||
self.k_buffer[layer_id - self.start_layer],
|
||||
self.v_buffer[layer_id - self.start_layer],
|
||||
)
|
||||
|
||||
def get_key_buffer(self, layer_id: int):
|
||||
if self.layer_transfer_counter is not None:
|
||||
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
|
||||
|
||||
if self.store_dtype != self.dtype:
|
||||
return self.k_buffer[layer_id - self.start_layer].view(self.dtype)
|
||||
return self.k_buffer[layer_id - self.start_layer]
|
||||
|
||||
def get_value_buffer(self, layer_id: int):
|
||||
if self.layer_transfer_counter is not None:
|
||||
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
|
||||
|
||||
if self.store_dtype != self.dtype:
|
||||
return self.v_buffer[layer_id - self.start_layer].view(self.dtype)
|
||||
return self.v_buffer[layer_id - self.start_layer]
|
||||
|
||||
def get_index_k_buffer(self, layer_id: int):
|
||||
if self.layer_transfer_counter is not None:
|
||||
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
|
||||
|
||||
if self.store_dtype != self.dtype:
|
||||
return self.index_k_buffer[layer_id - self.start_layer].view(self.dtype)
|
||||
return self.index_k_buffer[layer_id - self.start_layer]
|
||||
|
||||
# for disagg
|
||||
def get_contiguous_buf_infos(self):
|
||||
# MLA has only one kv_buffer, so only the information of this buffer needs to be returned.
|
||||
kv_data_ptrs = [self.k_buffer[i].data_ptr() for i in range(self.layer_num)] + [
|
||||
self.v_buffer[i].data_ptr() for i in range(self.layer_num)
|
||||
]
|
||||
kv_data_lens = [self.k_buffer[i].nbytes for i in range(self.layer_num)] + [
|
||||
self.v_buffer[i].nbytes for i in range(self.layer_num)
|
||||
]
|
||||
kv_item_lens = [self.k_buffer[i][0].nbytes for i in range(self.layer_num)] + [
|
||||
self.v_buffer[i][0].nbytes for i in range(self.layer_num)
|
||||
]
|
||||
if self.index_head_dim is not None:
|
||||
kv_data_ptrs += [
|
||||
self.index_k_buffer[i].data_ptr() for i in range(self.layer_num)
|
||||
]
|
||||
kv_data_lens += [
|
||||
self.index_k_buffer[i].nbytes for i in range(self.layer_num)
|
||||
]
|
||||
kv_item_lens += [
|
||||
self.index_k_buffer[i][0].nbytes for i in range(self.layer_num)
|
||||
]
|
||||
return kv_data_ptrs, kv_data_lens, kv_item_lens
|
||||
|
||||
def set_kv_buffer(
|
||||
self,
|
||||
layer: "RadixAttention",
|
||||
loc: torch.Tensor,
|
||||
cache_k: torch.Tensor,
|
||||
cache_v: torch.Tensor,
|
||||
):
|
||||
layer_id = layer.layer_id
|
||||
if cache_k.dtype != self.dtype:
|
||||
cache_k = cache_k.to(self.dtype)
|
||||
cache_v = cache_v.to(self.dtype)
|
||||
|
||||
if self.store_dtype != self.dtype:
|
||||
cache_k = cache_k.view(self.store_dtype)
|
||||
cache_v = cache_v.view(self.store_dtype)
|
||||
|
||||
if cache_v is None:
|
||||
cache_k, cache_v = cache_k.split(
|
||||
[self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
|
||||
)
|
||||
|
||||
torch_npu.npu_scatter_nd_update_(
|
||||
self.k_buffer[layer_id - self.start_layer].view(-1, 1, self.kv_lora_rank),
|
||||
loc.view(-1, 1),
|
||||
cache_k.view(-1, 1, self.kv_lora_rank),
|
||||
)
|
||||
torch_npu.npu_scatter_nd_update_(
|
||||
self.v_buffer[layer_id - self.start_layer].view(
|
||||
-1, 1, self.qk_rope_head_dim
|
||||
),
|
||||
loc.view(-1, 1),
|
||||
cache_v.view(-1, 1, self.qk_rope_head_dim),
|
||||
)
|
||||
|
||||
def set_index_k_buffer(
|
||||
self,
|
||||
layer_id: int,
|
||||
loc: torch.Tensor,
|
||||
index_k: torch.Tensor,
|
||||
):
|
||||
if index_k.dtype != self.dtype:
|
||||
index_k = index_k.to(self.dtype)
|
||||
|
||||
if self.store_dtype != self.dtype:
|
||||
index_k = index_k.view(self.store_dtype)
|
||||
|
||||
torch_npu.npu_scatter_nd_update_(
|
||||
self.index_k_buffer[layer_id - self.start_layer].view(
|
||||
-1, 1, self.index_head_dim
|
||||
),
|
||||
loc.view(-1, 1),
|
||||
index_k.view(-1, 1, self.index_head_dim),
|
||||
)
|
||||
@@ -0,0 +1,380 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
|
||||
NPUFusedMLAPreprocess,
|
||||
is_mla_preprocess_enabled,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.utils import (
|
||||
cp_split_and_rebuild_position,
|
||||
enable_prefill_cp,
|
||||
)
|
||||
from sglang.srt.layers.communicator import get_attn_tp_context
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
|
||||
from sglang.srt.utils import BumpAllocator
|
||||
|
||||
|
||||
# region MHA
|
||||
def forward_mha_prepare_npu(
|
||||
m: "DeepseekV2AttentionMLA",
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: "ForwardBatch",
|
||||
zero_allocator: "BumpAllocator",
|
||||
):
|
||||
if m.q_lora_rank is not None:
|
||||
q, latent_cache = (
|
||||
get_attn_tp_context()
|
||||
.fetch_qkv_latent()
|
||||
.split(
|
||||
[m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
|
||||
dim=-1,
|
||||
)
|
||||
)
|
||||
|
||||
# NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk
|
||||
|
||||
if m.use_nsa:
|
||||
q_lora = m.q_a_layernorm(q)
|
||||
q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
|
||||
_ = m.indexer(
|
||||
x=hidden_states,
|
||||
q_lora=q_lora,
|
||||
positions=positions,
|
||||
forward_batch=forward_batch,
|
||||
layer_id=m.layer_id,
|
||||
return_indices=False,
|
||||
)
|
||||
|
||||
else:
|
||||
q = m.q_a_layernorm(q)
|
||||
q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
|
||||
|
||||
else:
|
||||
q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim)
|
||||
latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0]
|
||||
|
||||
_, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
|
||||
kv_a, _ = latent_cache.split([m.kv_lora_rank, m.qk_rope_head_dim], dim=-1)
|
||||
latent_cache = latent_cache.unsqueeze(1)
|
||||
|
||||
kv_a = m.kv_a_layernorm(kv_a)
|
||||
kv = m.kv_b_proj(kv_a)[0]
|
||||
|
||||
k_pe = latent_cache[:, :, m.kv_lora_rank :]
|
||||
if m.rotary_emb is not None:
|
||||
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
|
||||
q[..., m.qk_nope_head_dim :] = q_pe
|
||||
|
||||
m._set_mla_kv_buffer(latent_cache, kv_a, k_pe, forward_batch)
|
||||
if forward_batch.mha_one_shot and sum(forward_batch.extend_prefix_lens_cpu) != 0:
|
||||
if m.use_nsa and m.kv_cache_dtype == "fp8_e4m3":
|
||||
# FP8 path: dequantize NSA-specific FP8 format to BF16
|
||||
kv_a, k_pe = m._get_mla_kv_buffer_from_fp8(forward_batch)
|
||||
else:
|
||||
# BF16/FP16 path: directly fetch from cache
|
||||
kv_a, k_pe = m._get_mla_kv_buffer(
|
||||
forward_batch.fetch_mha_one_shot_kv_indices(),
|
||||
q.dtype,
|
||||
forward_batch,
|
||||
)
|
||||
kv = m.kv_b_proj(kv_a)[0]
|
||||
kv = kv.view(-1, m.num_local_heads, m.qk_nope_head_dim + m.v_head_dim)
|
||||
k_nope = kv[..., : m.qk_nope_head_dim]
|
||||
v = kv[..., m.qk_nope_head_dim :]
|
||||
|
||||
k = m._concat_and_cast_mha_k(k_nope, k_pe, forward_batch)
|
||||
return q, k, v, forward_batch
|
||||
|
||||
|
||||
def forward_mha_core_npu(
|
||||
m: "DeepseekV2AttentionMLA",
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
forward_batch: "ForwardBatch",
|
||||
) -> torch.Tensor:
|
||||
attn_output = m.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
|
||||
attn_output = attn_output.reshape(-1, m.num_local_heads * m.v_head_dim)
|
||||
output, _ = m.o_proj(attn_output)
|
||||
return output
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region MLA
|
||||
def forward_mla_prepare_npu(
|
||||
m: "DeepseekV2AttentionMLA",
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: "ForwardBatch",
|
||||
zero_allocator: "BumpAllocator",
|
||||
):
|
||||
if is_mla_preprocess_enabled():
|
||||
if not hasattr(m, "mla_preprocess"):
|
||||
m.mla_preprocess = NPUFusedMLAPreprocess(
|
||||
m.fused_qkv_a_proj_with_mqa,
|
||||
m.q_a_layernorm,
|
||||
m.kv_a_layernorm,
|
||||
m.q_b_proj,
|
||||
m.w_kc,
|
||||
m.rotary_emb,
|
||||
m.layer_id,
|
||||
m.num_local_heads,
|
||||
m.qk_nope_head_dim,
|
||||
m.qk_rope_head_dim,
|
||||
m.quant_config,
|
||||
)
|
||||
(
|
||||
q_pe,
|
||||
k_pe,
|
||||
q_nope_out,
|
||||
k_nope,
|
||||
forward_batch,
|
||||
zero_allocator,
|
||||
positions,
|
||||
) = m.mla_preprocess.forward(
|
||||
positions, hidden_states, forward_batch, zero_allocator
|
||||
)
|
||||
topk_indices = None
|
||||
else:
|
||||
q_lora = None
|
||||
if m.q_lora_rank is not None:
|
||||
q, latent_cache = (
|
||||
get_attn_tp_context()
|
||||
.fetch_qkv_latent()
|
||||
.split(
|
||||
[m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
|
||||
dim=-1,
|
||||
)
|
||||
)
|
||||
k_nope = latent_cache[..., : m.kv_lora_rank]
|
||||
|
||||
q = m.q_a_layernorm(q)
|
||||
k_nope = m.kv_a_layernorm(k_nope)
|
||||
|
||||
# q_lora needed by indexer
|
||||
if m.use_nsa:
|
||||
q_lora = q
|
||||
|
||||
k_nope = k_nope.unsqueeze(1)
|
||||
q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
|
||||
else:
|
||||
q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim)
|
||||
latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0]
|
||||
k_nope = latent_cache[..., : m.kv_lora_rank]
|
||||
k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
|
||||
|
||||
q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
|
||||
k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
|
||||
|
||||
q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
|
||||
|
||||
q_nope_out = q_nope_out.transpose(0, 1)
|
||||
|
||||
if enable_prefill_cp(forward_batch, m.nsa_enable_prefill_cp):
|
||||
positions = cp_split_and_rebuild_position(forward_batch, positions)
|
||||
|
||||
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
|
||||
|
||||
if enable_prefill_cp(forward_batch, m.nsa_enable_prefill_cp):
|
||||
# support allgather+rerrange
|
||||
k_nope, k_pe = m.rebuild_cp_kv_cache(
|
||||
latent_cache, forward_batch, k_nope, k_pe
|
||||
)
|
||||
topk_indices = None
|
||||
if q_lora is not None:
|
||||
topk_indices = m.indexer(
|
||||
x=hidden_states,
|
||||
q_lora=q_lora,
|
||||
positions=positions,
|
||||
forward_batch=forward_batch,
|
||||
layer_id=m.layer_id,
|
||||
)
|
||||
|
||||
return (
|
||||
q_pe,
|
||||
k_pe,
|
||||
q_nope_out,
|
||||
k_nope,
|
||||
forward_batch,
|
||||
zero_allocator,
|
||||
positions,
|
||||
topk_indices,
|
||||
)
|
||||
|
||||
|
||||
def forward_mla_core_npu(
|
||||
m: "DeepseekV2AttentionMLA",
|
||||
q_pe: torch.Tensor,
|
||||
k_pe: torch.Tensor,
|
||||
q_nope_out: torch.Tensor,
|
||||
k_nope: torch.Tensor,
|
||||
forward_batch: "ForwardBatch",
|
||||
zero_allocator: "BumpAllocator",
|
||||
positions: torch.Tensor,
|
||||
topk_indices: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
attn_output = m.attn_mqa(
|
||||
q_nope_out,
|
||||
k_nope,
|
||||
k_nope,
|
||||
forward_batch,
|
||||
q_rope=q_pe,
|
||||
k_rope=k_pe,
|
||||
**(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
|
||||
)
|
||||
|
||||
attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
|
||||
|
||||
attn_bmm_output = torch.empty(
|
||||
(attn_output.shape[0], m.num_local_heads * m.v_head_dim),
|
||||
dtype=attn_output.dtype,
|
||||
device=attn_output.device,
|
||||
)
|
||||
torch.bmm(
|
||||
attn_output.transpose(0, 1),
|
||||
m.w_vc,
|
||||
out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose(0, 1),
|
||||
)
|
||||
output, _ = m.o_proj(attn_bmm_output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region DSA
|
||||
def forward_dsa_prepare_npu(
|
||||
m: "DeepseekV2AttentionMLA",
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: "ForwardBatch",
|
||||
zero_allocator: "BumpAllocator",
|
||||
):
|
||||
if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode():
|
||||
if not hasattr(m, "mla_preprocess"):
|
||||
m.mla_preprocess = NPUFusedMLAPreprocess(
|
||||
m.fused_qkv_a_proj_with_mqa,
|
||||
m.q_a_layernorm,
|
||||
m.kv_a_layernorm,
|
||||
m.q_b_proj,
|
||||
m.w_kc,
|
||||
m.rotary_emb,
|
||||
m.layer_id,
|
||||
m.num_local_heads,
|
||||
m.qk_nope_head_dim,
|
||||
m.qk_rope_head_dim,
|
||||
m.quant_config,
|
||||
)
|
||||
(
|
||||
q_pe,
|
||||
k_pe,
|
||||
q_nope_out,
|
||||
k_nope,
|
||||
forward_batch,
|
||||
zero_allocator,
|
||||
positions,
|
||||
) = m.mla_preprocess.forward(
|
||||
positions, hidden_states, forward_batch, zero_allocator
|
||||
)
|
||||
|
||||
fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
|
||||
q, _ = fused_qkv_a_proj_out.split(
|
||||
[m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
|
||||
)
|
||||
q_lora = m.q_a_layernorm(q)
|
||||
else:
|
||||
fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
|
||||
q, latent_cache = fused_qkv_a_proj_out.split(
|
||||
[m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
|
||||
)
|
||||
k_nope = latent_cache[..., : m.kv_lora_rank]
|
||||
|
||||
q = m.q_a_layernorm(q)
|
||||
k_nope = m.kv_a_layernorm(k_nope)
|
||||
|
||||
q_lora = q.clone() # required for topk_indices
|
||||
k_nope = k_nope.unsqueeze(1)
|
||||
q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
|
||||
|
||||
q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
|
||||
k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
|
||||
|
||||
q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
|
||||
|
||||
q_nope_out = q_nope_out.transpose(0, 1)
|
||||
|
||||
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
|
||||
|
||||
topk_indices = m.indexer(
|
||||
hidden_states, q_lora, positions, forward_batch, m.layer_id
|
||||
)
|
||||
|
||||
return (
|
||||
q_pe,
|
||||
k_pe,
|
||||
q_nope_out,
|
||||
k_nope,
|
||||
topk_indices,
|
||||
forward_batch,
|
||||
zero_allocator,
|
||||
positions,
|
||||
)
|
||||
|
||||
|
||||
def forward_dsa_core_npu(
|
||||
m: "DeepseekV2AttentionMLA",
|
||||
q_pe: torch.Tensor,
|
||||
k_pe: torch.Tensor,
|
||||
q_nope_out: torch.Tensor,
|
||||
k_nope: torch.Tensor,
|
||||
topk_indices: torch.Tensor,
|
||||
forward_batch: "ForwardBatch",
|
||||
zero_allocator: "BumpAllocator",
|
||||
positions: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
attn_output = m.attn_mqa(
|
||||
q_nope_out.contiguous(),
|
||||
k_nope.contiguous(),
|
||||
k_nope.contiguous(),
|
||||
forward_batch,
|
||||
save_kv_cache=True, # False if forward_batch.forward_mode.is_extend() else True,
|
||||
q_rope=q_pe.contiguous(),
|
||||
k_rope=k_pe.contiguous(),
|
||||
topk_indices=topk_indices,
|
||||
)
|
||||
attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
|
||||
|
||||
attn_bmm_output = torch.empty(
|
||||
(attn_output.shape[0], m.num_local_heads, m.v_head_dim),
|
||||
dtype=attn_output.dtype,
|
||||
device=attn_output.device,
|
||||
)
|
||||
|
||||
if not forward_batch.forward_mode.is_decode():
|
||||
attn_output = attn_output.transpose(0, 1)
|
||||
torch.bmm(
|
||||
attn_output,
|
||||
m.w_vc,
|
||||
out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose(
|
||||
0, 1
|
||||
),
|
||||
)
|
||||
else:
|
||||
attn_output = attn_output.contiguous()
|
||||
torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output)
|
||||
|
||||
attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
|
||||
|
||||
output, _ = m.o_proj(attn_bmm_output)
|
||||
return output
|
||||
|
||||
|
||||
# endregion
|
||||
@@ -0,0 +1,79 @@
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
from sgl_kernel_npu.norm.l1_norm import l1_norm
|
||||
|
||||
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
||||
from sglang.srt.eplb.expert_location_dispatch import topk_ids_logical_to_physical
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput, select_experts
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
|
||||
from sglang.srt.layers.moe.topk import TopKConfig, TopKOutput
|
||||
|
||||
|
||||
def fused_topk_npu(
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
topk_config: "TopKConfig",
|
||||
num_token_non_padded: Optional[torch.Tensor] = None,
|
||||
expert_location_dispatch_info: Optional["ExpertLocationDispatchInfo"] = None,
|
||||
) -> "TopKOutput":
|
||||
|
||||
use_grouped_topk = topk_config.use_grouped_topk
|
||||
renormalize = topk_config.renormalize
|
||||
correction_bias = topk_config.correction_bias
|
||||
|
||||
if not use_grouped_topk:
|
||||
topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k_softmax(
|
||||
router_logits,
|
||||
k=topk_config.top_k,
|
||||
)
|
||||
|
||||
if renormalize:
|
||||
topk_weights = l1_norm(
|
||||
topk_weights
|
||||
if topk_config.num_fused_shared_experts == 0
|
||||
else topk_weights[:, :-1]
|
||||
)
|
||||
topk_weights = topk_weights.to(torch.float32)
|
||||
|
||||
elif use_grouped_topk and correction_bias is not None:
|
||||
routed_scaling_factor = topk_config.routed_scaling_factor or 1
|
||||
|
||||
topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k(
|
||||
router_logits.to(torch.float32),
|
||||
k=topk_config.top_k,
|
||||
bias=correction_bias.to(torch.float32),
|
||||
k_group=topk_config.topk_group,
|
||||
group_count=topk_config.num_expert_group,
|
||||
group_select_mode=1,
|
||||
renorm=0,
|
||||
norm_type=1,
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
eps=float(1e-20),
|
||||
)
|
||||
|
||||
if renormalize:
|
||||
topk_weights_sum = (
|
||||
topk_weights.sum(dim=-1, keepdim=True)
|
||||
if topk_config.num_fused_shared_experts == 0
|
||||
else topk_weights[:, :-1].sum(dim=-1, keepdim=True)
|
||||
)
|
||||
topk_weights = topk_weights / topk_weights_sum
|
||||
|
||||
else:
|
||||
topk_config.torch_native = True
|
||||
return select_experts(
|
||||
hidden_states=hidden_states,
|
||||
router_logits=router_logits,
|
||||
topk_config=topk_config,
|
||||
num_token_non_padded=num_token_non_padded,
|
||||
expert_location_dispatch_info=expert_location_dispatch_info,
|
||||
)
|
||||
|
||||
if expert_location_dispatch_info is not None:
|
||||
topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
|
||||
get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
|
||||
|
||||
return StandardTopKOutput(topk_weights, topk_ids, router_logits)
|
||||
@@ -0,0 +1,916 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
|
||||
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
|
||||
from sglang.srt.utils import set_weight_attrs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.token_dispatcher import (
|
||||
CombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
|
||||
|
||||
def npu_fused_experts(
|
||||
hidden_states: torch.Tensor,
|
||||
w13: torch.Tensor,
|
||||
w13_scale: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
top_k: int,
|
||||
**kwargs,
|
||||
):
|
||||
w13_offset = kwargs.get("w13_offset", None)
|
||||
w2_offset = kwargs.get("w2_offset", None)
|
||||
use_wna16 = kwargs.get("use_wna16", False)
|
||||
|
||||
original_shape = hidden_states.shape
|
||||
original_dtype = hidden_states.dtype
|
||||
scale_dtype = original_dtype if original_dtype == torch.bfloat16 else torch.float32
|
||||
if len(original_shape) == 3:
|
||||
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
|
||||
num_tokens = hidden_states.shape[0]
|
||||
num_experts = w13.shape[0]
|
||||
row_idx_len = num_tokens * top_k
|
||||
row_idx = (
|
||||
torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device)
|
||||
.view(top_k, -1)
|
||||
.permute(1, 0)
|
||||
.contiguous()
|
||||
)
|
||||
hidden_states, expanded_row_idx, expanded_expert_idx = (
|
||||
torch.ops.npu.npu_moe_init_routing(
|
||||
hidden_states, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens
|
||||
)
|
||||
)
|
||||
expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens(
|
||||
expanded_expert_idx, num_experts
|
||||
)
|
||||
expert_tokens = expert_tokens.to(torch.int64)
|
||||
# gmm1: gate_up_proj
|
||||
if not use_wna16:
|
||||
hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
|
||||
scale_args13 = {
|
||||
"scale": [w13_scale.to(scale_dtype)],
|
||||
"per_token_scale": [pertoken_scale],
|
||||
}
|
||||
else:
|
||||
scale_args13 = {
|
||||
"antiquant_scale": [w13_scale],
|
||||
"antiquant_offset": [w13_offset],
|
||||
}
|
||||
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[w13],
|
||||
**scale_args13,
|
||||
split_item=2,
|
||||
group_list_type=0,
|
||||
group_type=0,
|
||||
group_list=expert_tokens,
|
||||
output_dtype=original_dtype,
|
||||
)[0]
|
||||
# act_fn: swiglu
|
||||
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
|
||||
if not use_wna16:
|
||||
hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
|
||||
|
||||
scale_args2 = {
|
||||
"scale": [w2_scale.to(scale_dtype)],
|
||||
"per_token_scale": [pertoken_scale],
|
||||
}
|
||||
else:
|
||||
scale_args2 = {"antiquant_scale": [w2_scale], "antiquant_offset": [w2_offset]}
|
||||
# gmm2: down_proj
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[w2],
|
||||
**scale_args2,
|
||||
split_item=2,
|
||||
group_list_type=0,
|
||||
group_type=0,
|
||||
group_list=expert_tokens,
|
||||
output_dtype=original_dtype,
|
||||
)[0]
|
||||
|
||||
final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
|
||||
hidden_states,
|
||||
skip1=None,
|
||||
skip2=None,
|
||||
bias=None,
|
||||
scales=topk_weights,
|
||||
expanded_src_to_dst_row=expanded_row_idx,
|
||||
export_for_source_row=topk_ids,
|
||||
)
|
||||
if len(original_shape) == 3:
|
||||
final_hidden_states = final_hidden_states.view(original_shape)
|
||||
return final_hidden_states
|
||||
|
||||
|
||||
def npu_fused_moe_without_routing_weights_bf16(
|
||||
layer, hidden_states, group_list_type, group_list, output_dtype
|
||||
):
|
||||
# gmm1: gate_up_proj
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[layer.w13_weight.permute(0, 2, 1)],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
group_list=group_list,
|
||||
output_dtype=output_dtype,
|
||||
)[0]
|
||||
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
|
||||
# gmm2: down_proj
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[layer.w2_weight.permute(0, 2, 1)],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
group_list=group_list,
|
||||
output_dtype=output_dtype,
|
||||
)[0]
|
||||
return hidden_states
|
||||
|
||||
|
||||
class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase):
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
) -> None:
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
|
||||
|
||||
self.num_experts = num_experts
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
|
||||
)
|
||||
|
||||
# weight
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size,
|
||||
dtype=torch.int8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
# scale
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
# offset
|
||||
w13_weight_offset = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_offset", w13_weight_offset)
|
||||
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
|
||||
w2_weight_offset = torch.nn.Parameter(
|
||||
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_offset", w2_weight_offset)
|
||||
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
|
||||
|
||||
def release_weight_cache(self, weight: torch.Tensor):
|
||||
# .contiguous() introduces additional memory overhead and needs to be released using resize_(0)
|
||||
origin_weight = weight.data.transpose(1, 2)
|
||||
new_weight = origin_weight.contiguous()
|
||||
origin_weight.untyped_storage().resize_(0)
|
||||
return new_weight
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
weight_data = self.release_weight_cache(layer.w13_weight.data)
|
||||
layer.w13_weight = torch.nn.Parameter(weight_data, requires_grad=False)
|
||||
|
||||
weight_data = self.release_weight_cache(layer.w2_weight.data)
|
||||
layer.w2_weight = torch.nn.Parameter(weight_data, requires_grad=False)
|
||||
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
layer.w13_weight_scale.data.squeeze(-1).contiguous().to(torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
layer.w2_weight_scale.data.squeeze(-1).contiguous(), requires_grad=False
|
||||
)
|
||||
layer.w13_weight_offset = torch.nn.Parameter(
|
||||
layer.w13_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_offset = torch.nn.Parameter(
|
||||
layer.w2_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False
|
||||
)
|
||||
|
||||
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data)
|
||||
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data)
|
||||
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
|
||||
):
|
||||
self.moe_runner_config = moe_runner_config
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer,
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
) -> "CombineInput":
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
|
||||
x = dispatch_output.hidden_states
|
||||
topk_output = dispatch_output.topk_output
|
||||
|
||||
topk_weights, topk_ids, _ = topk_output
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
topk_weights = topk_weights.to(x.dtype)
|
||||
output = npu_fused_experts(
|
||||
hidden_states=x,
|
||||
w13=layer.w13_weight,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
w2=layer.w2_weight,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
top_k=topk_ids.shape[1],
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
def apply_without_routing_weights(
|
||||
self,
|
||||
layer,
|
||||
hidden_states,
|
||||
hidden_states_scale,
|
||||
group_list_type,
|
||||
group_list,
|
||||
output_dtype,
|
||||
):
|
||||
# gmm1: gate_up_proj
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[layer.w13_weight],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
group_list=group_list,
|
||||
output_dtype=torch.int32,
|
||||
)[0]
|
||||
|
||||
# act_fn: swiglu
|
||||
hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant(
|
||||
x=hidden_states,
|
||||
weight_scale=layer.w13_weight_scale,
|
||||
activation_scale=hidden_states_scale,
|
||||
bias=None,
|
||||
quant_scale=None,
|
||||
quant_offset=None,
|
||||
group_index=group_list,
|
||||
activate_left=True,
|
||||
quant_mode=1,
|
||||
)
|
||||
|
||||
# gmm2: down_proj
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[layer.w2_weight],
|
||||
scale=[layer.w2_weight_scale.to(output_dtype)],
|
||||
per_token_scale=[swiglu_out_scale],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
group_list=group_list,
|
||||
output_dtype=output_dtype,
|
||||
)[0]
|
||||
return hidden_states
|
||||
|
||||
|
||||
class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.group_size = 256
|
||||
self.tp_size = 1
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
) -> None:
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
|
||||
|
||||
self.num_experts = num_experts
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
|
||||
)
|
||||
|
||||
# >> weight
|
||||
w13_output_size = intermediate_size_per_partition
|
||||
w2_output_size = hidden_size // 2
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(num_experts, w13_output_size, hidden_size, dtype=torch.int8),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
w2_output_size,
|
||||
intermediate_size_per_partition,
|
||||
dtype=torch.int8,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# >> scale
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
# >> offset
|
||||
w13_weight_offset = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_offset", w13_weight_offset)
|
||||
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
|
||||
|
||||
w2_weight_offset = torch.nn.Parameter(
|
||||
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_offset", w2_weight_offset)
|
||||
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
|
||||
|
||||
# >>> special param for w4a8
|
||||
w13_weight_scale_second = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // self.group_size,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second)
|
||||
set_weight_attrs(w13_weight_scale_second, extra_weight_attrs)
|
||||
w13_weight_offset_second = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // self.group_size,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_offset_second", w13_weight_offset_second)
|
||||
set_weight_attrs(w13_weight_offset_second, extra_weight_attrs)
|
||||
|
||||
w2_weight_scale_second = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // self.group_size,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second)
|
||||
set_weight_attrs(w2_weight_scale_second, extra_weight_attrs)
|
||||
|
||||
w2_weight_offset_second = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // self.group_size,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second)
|
||||
set_weight_attrs(w2_weight_offset_second, extra_weight_attrs)
|
||||
|
||||
w13_scale_bias = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_scale_bias", w13_scale_bias)
|
||||
set_weight_attrs(w13_scale_bias, extra_weight_attrs)
|
||||
|
||||
w2_scale_bias = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts, hidden_size, 16 // self.tp_size, dtype=torch.float32
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_scale_bias", w2_scale_bias)
|
||||
set_weight_attrs(w2_scale_bias, extra_weight_attrs)
|
||||
|
||||
def process_scale(self, weight: torch.Tensor, scale, per_group_scale):
|
||||
scale = scale.transpose(1, 2).contiguous()
|
||||
per_group_scale = per_group_scale.transpose(1, 2).contiguous()
|
||||
group_num, k, n = weight.shape
|
||||
# the weight of the new version is reduced by half by pack n, so it needs to be restored
|
||||
n = n * 2
|
||||
per_group_scale = per_group_scale.reshape(group_num, -1, n)
|
||||
group_num, quantgroup_num, n = per_group_scale.shape
|
||||
bias = None
|
||||
|
||||
scale_fp32 = (scale * per_group_scale).to(torch.float16).to(torch.float32)
|
||||
scale_fp32_np = scale_fp32.cpu().numpy()
|
||||
scale_fp32_np.dtype = np.uint32
|
||||
sscale_uint64 = np.zeros((group_num, quantgroup_num, n * 2), dtype=np.uint32)
|
||||
|
||||
sscale_uint64[..., ::2] = scale_fp32_np
|
||||
|
||||
sscale_uint64_buffer = np.frombuffer(
|
||||
sscale_uint64.tobytes(), dtype=np.int64
|
||||
).copy()
|
||||
sscale_uint64_tensor = torch.from_numpy(sscale_uint64_buffer).reshape(
|
||||
group_num, quantgroup_num, n
|
||||
)
|
||||
sscale_uint64_tensor = sscale_uint64_tensor.npu()
|
||||
return sscale_uint64_tensor, bias
|
||||
|
||||
def update_bias(self, layer, w13_bias, w2_bias):
|
||||
layer.w13_scale_bias.data = (
|
||||
layer.w13_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1)
|
||||
)
|
||||
layer.w2_scale_bias.data = (
|
||||
layer.w2_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1)
|
||||
)
|
||||
|
||||
def pack_to_int32(self, weight: torch.Tensor):
|
||||
# pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4
|
||||
assert (
|
||||
weight.shape[-1] % 4 == 0
|
||||
), "the last dim of weight needs to be divided by 4"
|
||||
return weight.view(torch.int32).contiguous()
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
layer.w13_weight.data.transpose(1, 2).contiguous(), requires_grad=False
|
||||
)
|
||||
layer.w2_weight = torch.nn.Parameter(
|
||||
layer.w2_weight.data.transpose(1, 2).contiguous(), requires_grad=False
|
||||
)
|
||||
|
||||
w13_weight_scale_second = (
|
||||
layer.w13_weight_scale_second.data
|
||||
if hasattr(layer, "w13_weight_scale_second")
|
||||
else None
|
||||
)
|
||||
w2_weight_scale_second = (
|
||||
layer.w2_weight_scale_second.data
|
||||
if hasattr(layer, "w2_weight_scale_second")
|
||||
else None
|
||||
)
|
||||
layer.w13_weight_scale.data, w13_bias = self.process_scale(
|
||||
layer.w13_weight, layer.w13_weight_scale.data, w13_weight_scale_second
|
||||
)
|
||||
layer.w2_weight_scale.data, w2_bias = self.process_scale(
|
||||
layer.w2_weight, layer.w2_weight_scale.data, w2_weight_scale_second
|
||||
)
|
||||
if hasattr(layer, "w13_weight_scale_second"):
|
||||
# scale_second is no longer used, release this part of the memory
|
||||
del layer.w13_weight_scale_second
|
||||
del layer.w2_weight_scale_second
|
||||
del layer.w13_weight_offset_second
|
||||
del layer.w2_weight_offset_second
|
||||
|
||||
self.update_bias(layer, w13_bias, w2_bias)
|
||||
|
||||
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data)
|
||||
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data)
|
||||
layer.w13_weight.data = self.pack_to_int32(layer.w13_weight.data)
|
||||
layer.w2_weight.data = self.pack_to_int32(layer.w2_weight.data)
|
||||
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
|
||||
):
|
||||
self.moe_runner_config = moe_runner_config
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer,
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
) -> "CombineInput":
|
||||
# FIXME W4A8 only support with deepep
|
||||
raise NotImplementedError(
|
||||
f"W4A8 only support with deepep for now, please enable --moe-a2a-backend deepep"
|
||||
)
|
||||
|
||||
def apply_without_routing_weights(
|
||||
self,
|
||||
layer,
|
||||
hidden_states,
|
||||
hidden_states_scale,
|
||||
group_list_type,
|
||||
group_list,
|
||||
output_dtype,
|
||||
):
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[self.w13_weight],
|
||||
scale=[self.w13_weight_scale],
|
||||
bias=[self.w13_scale_bias],
|
||||
per_token_scale=[hidden_states_scale],
|
||||
group_list=group_list,
|
||||
split_item=2,
|
||||
group_type=0,
|
||||
group_list_type=group_list_type,
|
||||
output_dtype=output_dtype,
|
||||
)[0]
|
||||
|
||||
# act_fn: swiglu
|
||||
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
|
||||
hidden_states, swiglu_out_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
|
||||
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[self.w2_weight],
|
||||
scale=[self.w2_weight_scale],
|
||||
bias=[self.w2_scale_bias],
|
||||
per_token_scale=[swiglu_out_scale],
|
||||
group_list=group_list,
|
||||
split_item=2,
|
||||
group_type=0,
|
||||
group_list_type=group_list_type,
|
||||
output_dtype=output_dtype,
|
||||
)[0]
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class NPUW4A16Int4DynamicMoEMethod(FusedMoEMethodBase):
|
||||
|
||||
def __init__(self, quantization_config) -> None:
|
||||
self.pack_factor = 8 # weight dtype is int4, but use int32 to create
|
||||
target = (
|
||||
"MoEGMM" if "MoEGMM" in quantization_config.target_scheme_map else "Linear"
|
||||
)
|
||||
if target in quantization_config.target_scheme_map:
|
||||
self.group_size = quantization_config.target_scheme_map[target][
|
||||
"weights"
|
||||
].group_size
|
||||
else:
|
||||
self.group_size = 128
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
) -> None:
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
|
||||
|
||||
self.num_experts = num_experts
|
||||
if (
|
||||
extra_weight_attrs.get(
|
||||
"intermediate_size_full", intermediate_size_per_partition
|
||||
)
|
||||
// intermediate_size_per_partition
|
||||
> 1
|
||||
):
|
||||
quant_method = FusedMoeWeightScaleSupported.GROUP.value
|
||||
else:
|
||||
quant_method = FusedMoeWeightScaleSupported.CHANNEL.value
|
||||
extra_weight_attrs.update({"quant_method": quant_method})
|
||||
# weight
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // self.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // self.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# scale
|
||||
weight_scale_dtype = torch.bfloat16
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // self.group_size,
|
||||
dtype=weight_scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale", w13_weight_scale)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // self.group_size,
|
||||
dtype=weight_scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_scale", w2_weight_scale)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
# offset
|
||||
w13_weight_offset = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // self.group_size,
|
||||
dtype=weight_scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight_offset", w13_weight_offset)
|
||||
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
|
||||
|
||||
w2_weight_offset = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition // self.group_size,
|
||||
dtype=weight_scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight_offset", w2_weight_offset)
|
||||
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
|
||||
|
||||
def pack_to_int32(self, weight: torch.Tensor):
|
||||
assert weight.dim() == 3
|
||||
if weight.dtype == torch.int32:
|
||||
# pack 8 int4 to int32, we use a int32 to represent a int4
|
||||
assert (
|
||||
weight.shape[-1] % 8 == 0
|
||||
), "the last dim of weight needs to be divided by 8"
|
||||
new_weight = torch.ops.npu.npu_convert_weight_to_int4pack(
|
||||
weight.flatten(0, 1)
|
||||
)
|
||||
new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1)
|
||||
elif weight.dtype == torch.int8:
|
||||
# pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4
|
||||
assert (
|
||||
weight.shape[-1] % 4 == 0
|
||||
), "the last dim of weight needs to be divided by 4"
|
||||
new_weight = weight.view(torch.int32).contiguous()
|
||||
else:
|
||||
raise ValueError(f"{weight.dtype=} is not supported !")
|
||||
return new_weight
|
||||
|
||||
def unpack_from_int32(
|
||||
self,
|
||||
value: torch.Tensor,
|
||||
num_bits: int,
|
||||
shape: torch.Size = None,
|
||||
packed_dim=1,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Unpacks a tensor of packed int32 weights into individual int8s, maintaining the
|
||||
original bit range.
|
||||
|
||||
Return tensors in int8
|
||||
|
||||
:param value: tensor to unpack
|
||||
:param num_bits: number of bits to unpack each data point into
|
||||
:param shape: shape to unpack into, used to remove padding
|
||||
:returns: unpacked int8 tensor
|
||||
"""
|
||||
if value.dtype is not torch.int32:
|
||||
raise ValueError(
|
||||
f"Expected {torch.int32} but got {value.dtype}, Aborting unpack."
|
||||
)
|
||||
|
||||
if num_bits > 8:
|
||||
raise ValueError("Unpacking is only supported for less than 8 bits")
|
||||
|
||||
pack_factor = 32 // num_bits
|
||||
|
||||
# unpack
|
||||
mask = (1 << num_bits) - 1
|
||||
|
||||
if packed_dim == 1:
|
||||
unpacked = torch.zeros(
|
||||
(value.shape[0], value.shape[1] * pack_factor),
|
||||
device=value.device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
for i in range(pack_factor):
|
||||
unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask
|
||||
|
||||
# remove padding
|
||||
if shape is not None:
|
||||
original_row_size = int(shape[1])
|
||||
unpacked = unpacked[:, :original_row_size]
|
||||
else:
|
||||
unpacked = torch.zeros(
|
||||
(value.shape[0] * pack_factor, value.shape[1]),
|
||||
device=value.device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
for i in range(pack_factor):
|
||||
unpacked[i::pack_factor, :] = (value >> (num_bits * i)) & mask
|
||||
|
||||
# remove padding
|
||||
original_row_size = int(shape[0])
|
||||
unpacked = unpacked[:original_row_size, :]
|
||||
|
||||
# bits are packed in unsigned format, reformat to signed
|
||||
# update the value range from unsigned to signed
|
||||
offset = pow(2, num_bits) // 2
|
||||
unpacked = (unpacked - offset).to(torch.int8)
|
||||
|
||||
return unpacked
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
w13_weight_scale = layer.w13_weight_scale.data.transpose(-1, -2).contiguous()
|
||||
w2_weight_scale = layer.w2_weight_scale.data.transpose(-1, -2).contiguous()
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
w13_weight_scale, requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False)
|
||||
|
||||
layer.w13_weight_offset = torch.nn.Parameter(
|
||||
layer.w13_weight_offset.data.transpose(-1, -2).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_offset = torch.nn.Parameter(
|
||||
layer.w2_weight_offset.data.transpose(-1, -2).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
# w = [n, k // 8] --> [k, n // 8]
|
||||
# w13_weight = layer.w13_weight.data.transpose(1, 2).contiguous()
|
||||
# w2_weight = layer.w2_weight.data.transpose(1, 2).contiguous()
|
||||
unpacked_w13_weight = (
|
||||
self.unpack_from_int32(layer.w13_weight.data.flatten(0, 1), 4)
|
||||
.view(layer.w13_weight.data.shape[0], layer.w13_weight.data.shape[1], -1)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
.int()
|
||||
)
|
||||
unpacked_w2_weight = (
|
||||
self.unpack_from_int32(layer.w2_weight.data.flatten(0, 1), 4)
|
||||
.view(layer.w2_weight.data.shape[0], layer.w2_weight.data.shape[1], -1)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
.int()
|
||||
)
|
||||
|
||||
w13_weight = self.pack_to_int32(unpacked_w13_weight)
|
||||
w2_weight = self.pack_to_int32(unpacked_w2_weight)
|
||||
|
||||
layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
|
||||
layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
|
||||
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
|
||||
):
|
||||
self.moe_runner_config = moe_runner_config
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer,
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
) -> "CombineInput":
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
|
||||
x = dispatch_output.hidden_states
|
||||
topk_output = dispatch_output.topk_output
|
||||
|
||||
topk_weights, topk_ids, _ = topk_output
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
topk_weights = topk_weights.to(x.dtype)
|
||||
output = npu_fused_experts(
|
||||
hidden_states=x,
|
||||
w13=layer.w13_weight,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
w13_offset=layer.w13_weight_offset,
|
||||
w2=layer.w2_weight,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
w2_offset=layer.w2_weight_offset,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
top_k=topk_ids.shape[1],
|
||||
use_wna16=True,
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
def apply_without_routing_weights(
|
||||
self,
|
||||
layer,
|
||||
hidden_states,
|
||||
hidden_states_scale,
|
||||
group_list_type,
|
||||
group_list,
|
||||
output_dtype,
|
||||
):
|
||||
if hidden_states_scale is None:
|
||||
# gmm1: gate_up_proj
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[layer.w13_weight],
|
||||
antiquant_scale=[layer.w13_weight_scale],
|
||||
antiquant_offset=[layer.w13_weight_offset],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
group_list=group_list,
|
||||
output_dtype=output_dtype,
|
||||
)[0]
|
||||
|
||||
# act_fn: swiglu
|
||||
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
|
||||
|
||||
# gmm2: down_proj
|
||||
out_hidden = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[layer.w2_weight],
|
||||
antiquant_scale=[layer.w2_weight_scale],
|
||||
antiquant_offset=[layer.w2_weight_offset],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
group_list=group_list,
|
||||
output_dtype=output_dtype,
|
||||
)[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
"when weight is int4, hidden_states only supports non-quant dtype!"
|
||||
)
|
||||
|
||||
return out_hidden
|
||||
@@ -0,0 +1,215 @@
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
|
||||
from sglang.srt.layers.parameter import (
|
||||
ChannelQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
from sglang.srt.layers.quantization.base_config import LinearMethodBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
|
||||
class _NPULinearMethodBase(LinearMethodBase):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
quant_config: Optional["QuantizationConfig"] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.quant_config = quant_config
|
||||
|
||||
|
||||
class NPUW8A8Int8LinearMethod(_NPULinearMethodBase):
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: List[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
(output_size_per_partition, input_size_per_partition), dtype=torch.int8
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
weight_offset = ChannelQuantScaleParameter(
|
||||
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_offset", weight_offset)
|
||||
|
||||
input_scale = PerTensorScaleParameter(
|
||||
data=torch.empty(1, dtype=params_dtype),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
input_scale.ignore_warning = True
|
||||
layer.register_parameter("input_scale", input_scale)
|
||||
|
||||
input_offset = PerTensorScaleParameter(
|
||||
data=torch.empty(1, dtype=params_dtype),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
input_offset.ignore_warning = True
|
||||
layer.register_parameter("input_offset", input_offset)
|
||||
|
||||
quant_bias = ChannelQuantScaleParameter(
|
||||
data=torch.empty(output_size_per_partition, dtype=torch.int32),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("quant_bias", quant_bias)
|
||||
|
||||
if params_dtype == torch.bfloat16:
|
||||
deq_scale_dtype = torch.float32
|
||||
elif params_dtype == torch.float16:
|
||||
deq_scale_dtype = torch.int64
|
||||
else:
|
||||
raise ValueError(f"Unsupported params_dtype: {params_dtype}")
|
||||
deq_scale = ChannelQuantScaleParameter(
|
||||
data=torch.empty(output_size_per_partition, dtype=deq_scale_dtype),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("deq_scale", deq_scale)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
from sglang.srt.layers.linear import RowParallelLinear
|
||||
|
||||
original_dtype = x.dtype
|
||||
if original_dtype != torch.int8:
|
||||
x = torch.ops.npu.npu_quantize(
|
||||
x,
|
||||
layer.aclnn_input_scale_reciprocal,
|
||||
layer.aclnn_input_offset,
|
||||
torch.qint8,
|
||||
-1,
|
||||
False,
|
||||
)
|
||||
# Only fuse bias add into GEMM for rank 0 (this ensures that
|
||||
# bias will not get added more than once in Attention TP>1 case)
|
||||
if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0:
|
||||
quant_bias = None
|
||||
else:
|
||||
quant_bias = layer.quant_bias
|
||||
return torch.ops.npu.npu_quant_matmul(
|
||||
x,
|
||||
layer.weight,
|
||||
layer.deq_scale,
|
||||
bias=quant_bias,
|
||||
output_dtype=original_dtype,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module):
|
||||
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
|
||||
layer.weight.data = npu_format_cast(layer.weight.data)
|
||||
|
||||
layer.weight_scale.data = torch.flatten(layer.weight_scale.data)
|
||||
layer.weight_offset.data = torch.flatten(layer.weight_offset.data)
|
||||
|
||||
expanding_factor = layer.weight.data.shape[0]
|
||||
layer.aclnn_input_scale = torch.nn.Parameter(
|
||||
layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter(
|
||||
layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.aclnn_input_offset = torch.nn.Parameter(
|
||||
layer.input_offset.data.repeat(expanding_factor).to(device="npu"),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
|
||||
class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: List[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
(output_size_per_partition, input_size_per_partition), dtype=torch.int8
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
weight_offset = ChannelQuantScaleParameter(
|
||||
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_offset", weight_offset)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
original_dtype = x.dtype
|
||||
quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(x)
|
||||
return torch.ops.npu.npu_quant_matmul(
|
||||
quant_out,
|
||||
layer.weight,
|
||||
layer.weight_scale,
|
||||
pertoken_scale=dynamic_scale,
|
||||
bias=bias,
|
||||
output_dtype=original_dtype,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module):
|
||||
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
|
||||
layer.weight.data = npu_format_cast(layer.weight.data)
|
||||
|
||||
layer.weight_scale.data = layer.weight_scale.data.flatten()
|
||||
layer.weight_offset.data = layer.weight_offset.data.flatten()
|
||||
@@ -0,0 +1,241 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast
|
||||
|
||||
import torch
|
||||
from compressed_tensors.quantization import QuantizationStrategy
|
||||
|
||||
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
|
||||
NPUW4A8Int4DynamicMoEMethod,
|
||||
NPUW4A16Int4DynamicMoEMethod,
|
||||
NPUW8A8Int8DynamicMoEMethod,
|
||||
)
|
||||
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
|
||||
NPUW8A8Int8DynamicLinearMethod,
|
||||
NPUW8A8Int8LinearMethod,
|
||||
)
|
||||
from sglang.srt.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
|
||||
CompressedTensorsConfig,
|
||||
)
|
||||
from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer
|
||||
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
|
||||
from sglang.srt.utils import apply_module_patch
|
||||
|
||||
|
||||
# func refers to RMSNorm.__init__
|
||||
def npu_wrapper_rmsnorm_init(func):
|
||||
def init(self, hidden_size: int, **extra_args) -> None:
|
||||
func(self, hidden_size, **extra_args)
|
||||
self.ignore_anti = True
|
||||
# The Ascend w8a8_int8 quantization requires adding a bias in rmsnorm
|
||||
self.bias = torch.nn.Parameter(torch.zeros(hidden_size), requires_grad=False)
|
||||
|
||||
return init
|
||||
|
||||
|
||||
# func refers to RMSNorm.forward_oot
|
||||
def npu_wrapper_rmsnorm_forward(func):
|
||||
def _rmsnorm_forward_oot(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
residual: Optional[torch.Tensor] = None,
|
||||
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
|
||||
from sgl_kernel_npu.norm.add_rmsnorm_bias import add_rmsnorm_bias
|
||||
|
||||
if not x.is_contiguous():
|
||||
x = x.contiguous()
|
||||
if residual is not None:
|
||||
out, residual_out = add_rmsnorm_bias(
|
||||
x,
|
||||
residual,
|
||||
self.weight.data,
|
||||
self.bias,
|
||||
self.variance_epsilon,
|
||||
)
|
||||
return out.to(x.dtype), residual_out
|
||||
|
||||
out = torch.ops.npu.npu_rms_norm(x, self.weight.data, self.variance_epsilon)[0]
|
||||
out = out + self.bias
|
||||
return out.to(x.dtype)
|
||||
|
||||
return _rmsnorm_forward_oot
|
||||
|
||||
|
||||
class ModelSlimConfig(QuantizationConfig):
|
||||
"""
|
||||
Config class for ModelSlim Quantization, a NPU-specific quantization type.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: Dict[str, Any] = {}):
|
||||
super().__init__()
|
||||
self.quant_description = quant_config
|
||||
self.is_dynamic = quant_config.get("is_dynamic", False)
|
||||
self.is_moe_w4_dynamic = False
|
||||
ignore = cast(List[str], quant_config.get("ignore", []))
|
||||
self.ignore = ignore if ignore is not None else []
|
||||
packed_modules_mapping = quant_config.get("packed_modules_mapping", {})
|
||||
self.packed_modules_mapping = (
|
||||
packed_modules_mapping if packed_modules_mapping is not None else {}
|
||||
)
|
||||
self.target_scheme_map = (
|
||||
CompressedTensorsConfig._quantization_scheme_map_from_config(
|
||||
config=quant_config
|
||||
)
|
||||
)
|
||||
target = "MoEGMM" if "MoEGMM" in self.target_scheme_map else "Linear"
|
||||
target_scheme = self.target_scheme_map.get(target, None)
|
||||
if target_scheme is None:
|
||||
self.is_moe_w4_dynamic = False
|
||||
else:
|
||||
weight_quant = target_scheme.get("weights")
|
||||
input_quant = target_scheme.get("input_activations")
|
||||
self.is_moe_w4_dynamic = self.is_dynamic_token_w4(weight_quant, input_quant)
|
||||
self.is_moe_input_quant = input_quant
|
||||
|
||||
for name in self.quant_description.keys():
|
||||
if "norm.bias" in name:
|
||||
apply_module_patch(
|
||||
"sglang.srt.layers.layernorm.RMSNorm",
|
||||
"__init__",
|
||||
[npu_wrapper_rmsnorm_init],
|
||||
)
|
||||
apply_module_patch(
|
||||
"sglang.srt.layers.layernorm.RMSNorm",
|
||||
"forward_npu",
|
||||
[npu_wrapper_rmsnorm_forward],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
|
||||
return [torch.int8, torch.float16, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 0
|
||||
|
||||
@classmethod
|
||||
def get_name(self) -> str:
|
||||
return "modelslim"
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> List[str]:
|
||||
filenames = ["quant_model_description.json"]
|
||||
return filenames
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> ModelSlimConfig:
|
||||
return cls(config)
|
||||
|
||||
def get_quant_method(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
prefix: str,
|
||||
) -> Optional[QuantizeMethodBase]:
|
||||
from sglang.srt.layers.linear import LinearBase
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||
|
||||
if isinstance(layer, LinearBase):
|
||||
if should_ignore_layer(
|
||||
prefix,
|
||||
ignore=self.ignore,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
key = "model"
|
||||
if "vision_model" in prefix:
|
||||
key = "vision_model"
|
||||
elif "visual" in prefix:
|
||||
key = "visual"
|
||||
packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {})
|
||||
prefix_in_quant_config = prefix
|
||||
proj_name = prefix.split(".")[-1]
|
||||
if proj_name in packed_modules_mapping_subset:
|
||||
prefix_in_quant_config = prefix.replace(
|
||||
proj_name, packed_modules_mapping_subset[proj_name][0]
|
||||
)
|
||||
self.is_dynamic = (
|
||||
self.quant_description[prefix_in_quant_config + ".weight"]
|
||||
== "W8A8_DYNAMIC"
|
||||
)
|
||||
if self.is_layer_skipped(prefix, packed_modules_mapping_subset):
|
||||
return UnquantizedLinearMethod()
|
||||
return (
|
||||
NPUW8A8Int8DynamicLinearMethod(self)
|
||||
if self.is_dynamic
|
||||
else NPUW8A8Int8LinearMethod(self)
|
||||
)
|
||||
elif isinstance(layer, FusedMoE):
|
||||
prefix_in_quant_config = prefix + ".0.down_proj.weight"
|
||||
is_moe_w4a8_dynamic = (
|
||||
self.quant_description.get(prefix_in_quant_config, "STATIC")
|
||||
== "W4A8_DYNAMIC"
|
||||
)
|
||||
if (
|
||||
self.is_moe_w4_dynamic and self.is_moe_input_quant is not None
|
||||
) or is_moe_w4a8_dynamic:
|
||||
return NPUW4A8Int4DynamicMoEMethod()
|
||||
elif self.is_moe_w4_dynamic and self.is_moe_input_quant is None:
|
||||
return NPUW4A16Int4DynamicMoEMethod(self)
|
||||
else:
|
||||
return NPUW8A8Int8DynamicMoEMethod()
|
||||
return None
|
||||
|
||||
def is_layer_skipped(
|
||||
self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({})
|
||||
):
|
||||
# adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped
|
||||
proj_name = prefix.split(".")[-1]
|
||||
if proj_name in fused_mapping:
|
||||
shard_prefixes = [
|
||||
prefix.replace(proj_name, shard_proj_name)
|
||||
for shard_proj_name in fused_mapping[proj_name]
|
||||
]
|
||||
|
||||
is_skipped = None
|
||||
for shard_prefix in shard_prefixes:
|
||||
is_shard_skipped = (
|
||||
self.quant_description[shard_prefix + ".weight"] == "FLOAT"
|
||||
)
|
||||
|
||||
if is_skipped is None:
|
||||
is_skipped = is_shard_skipped
|
||||
elif is_shard_skipped != is_skipped:
|
||||
raise ValueError(
|
||||
f"Detected some but not all shards of {prefix} "
|
||||
"are quantized. All shards of fused layers "
|
||||
"to have the same precision."
|
||||
)
|
||||
else:
|
||||
is_skipped = self.quant_description[prefix + ".weight"] == "FLOAT"
|
||||
|
||||
assert is_skipped is not None
|
||||
return is_skipped
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
def is_dynamic_token_w4(self, weight_quant, input_quant) -> bool:
|
||||
is_w4 = weight_quant.num_bits == 4
|
||||
weight_strategy = (
|
||||
weight_quant.strategy == QuantizationStrategy.TENSOR.value
|
||||
or weight_quant.strategy == QuantizationStrategy.CHANNEL.value
|
||||
or weight_quant.strategy == QuantizationStrategy.GROUP.value
|
||||
)
|
||||
if input_quant is not None:
|
||||
is_token = (
|
||||
weight_strategy
|
||||
and input_quant.strategy == QuantizationStrategy.TOKEN.value
|
||||
)
|
||||
is_dynamic = not weight_quant.dynamic and input_quant.dynamic
|
||||
else:
|
||||
is_token = weight_strategy
|
||||
is_dynamic = not weight_quant.dynamic
|
||||
|
||||
# Both symmetric and asymmetric input quantization supported.
|
||||
# Only symmetric weight quantization supported.
|
||||
return is_w4 and weight_quant.symmetric and is_token and is_dynamic
|
||||
@@ -0,0 +1,104 @@
|
||||
import functools
|
||||
import logging
|
||||
from enum import IntEnum
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import is_npu
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_is_npu = is_npu()
|
||||
|
||||
|
||||
class NPUACLFormat(IntEnum):
|
||||
ACL_FORMAT_UNDEFINED = -1
|
||||
ACL_FORMAT_ND = 2
|
||||
ACL_FORMAT_FRACTAL_NZ = 29
|
||||
|
||||
|
||||
def _call_once(fn: Callable):
|
||||
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
if getattr(fn, "_has_been_called", False):
|
||||
logger.debug("Function {} has already been called.", fn.__name__)
|
||||
return
|
||||
|
||||
fn._has_been_called = True
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def set_default_server_args(args: "ServerArgs"):
|
||||
"""
|
||||
Set default server arguments for NPU backend.
|
||||
"""
|
||||
|
||||
# NPU only works with "ascend" attention backend for now
|
||||
args.attention_backend = "ascend"
|
||||
args.prefill_attention_backend = "ascend"
|
||||
args.decode_attention_backend = "ascend"
|
||||
if args.page_size is None:
|
||||
args.page_size = 128
|
||||
|
||||
# NPU does not support CustomAllReduce
|
||||
args.disable_custom_all_reduce = True
|
||||
|
||||
# handles hierarchical cache configs
|
||||
if args.enable_hierarchical_cache:
|
||||
args.hicache_io_backend = "kernel_ascend"
|
||||
if args.use_mla_backend():
|
||||
args.hicache_mem_layout = "page_first_kv_split"
|
||||
else:
|
||||
args.hicache_mem_layout = "page_first_direct"
|
||||
|
||||
|
||||
@_call_once
|
||||
def init_npu_backend():
|
||||
"""
|
||||
Initialize NPU backend. This function should be called only once.
|
||||
"""
|
||||
|
||||
assert _is_npu, "NPU backend initialization called on non-NPU device."
|
||||
|
||||
import sgl_kernel_npu # noqa: F401
|
||||
import torch_npu
|
||||
from torch_npu.contrib import transfer_to_npu # noqa: F401
|
||||
|
||||
# Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it True
|
||||
torch.cuda.is_available = lambda: False
|
||||
|
||||
torch_npu.npu.config.allow_internal_format = True
|
||||
torch_npu.npu.set_compile_mode(jit_compile=False)
|
||||
|
||||
|
||||
def npu_format_cast(
|
||||
tensor: torch.Tensor,
|
||||
acl_format: NPUACLFormat = NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Cast a tensor to a specific NPU ACL format.
|
||||
|
||||
Args:
|
||||
tensor (torch.Tensor): The input tensor.
|
||||
acl_format (NPUACLFormat): The target NPU ACL format.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The tensor cast to the specified NPU ACL format.
|
||||
"""
|
||||
|
||||
if not _is_npu:
|
||||
return tensor
|
||||
|
||||
if envs.SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT.get():
|
||||
return tensor
|
||||
|
||||
import torch_npu
|
||||
|
||||
return torch_npu.npu_format_cast(tensor, acl_format.value)
|
||||
Reference in New Issue
Block a user