[ROCm] Enable Fused MLA Triton kernel for DeepSeekV3 (#3237)
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
Chaitanya Sri Krishna Lolla
committed by
GitHub
GitHub
parent
3758d209a0
commit
6ce9dbe828
@@ -0,0 +1,446 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""
|
||||
Memory-efficient attention for decoding.
|
||||
It supports page size = 1.
|
||||
"""
|
||||
|
||||
# Adapted from
|
||||
# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
|
||||
# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.layers.attention.triton_ops.decode_attention import (
|
||||
_decode_softmax_reducev_fwd,
|
||||
)
|
||||
from sglang.srt.layers.rotary_embedding import DeepseekScalingRotaryEmbedding
|
||||
|
||||
|
||||
def is_hip():
|
||||
return triton.runtime.driver.active.get_current_target().backend == "hip"
|
||||
|
||||
|
||||
is_hip_ = is_hip()
|
||||
|
||||
|
||||
@triton.jit
|
||||
def tanh(x):
|
||||
# Tanh is just a scaled sigmoid
|
||||
return 2 * tl.sigmoid(2 * x) - 1
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fwd_grouped_kernel_stage1_rope(
|
||||
Q, # Holds [Q_NOPE; Q_PE], b x h x (d+r)
|
||||
K_Buffer, # Holds [KV; K_PE], b*s x (c+r)
|
||||
V_buffer, # Holds [KV], b*s x (c)
|
||||
cos_sin_cache, # max_seq_len x (rotary_dim * 2)
|
||||
positions, # sequence positions
|
||||
sm_scale,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
Att_Out, # b x h x NUM_KV_SPLITS x (kv_lora_rank + 1)
|
||||
k_pe_t_out,
|
||||
stride_qb,
|
||||
stride_qh,
|
||||
stride_buf_kbs,
|
||||
stride_buf_vbs,
|
||||
stride_mid_ob,
|
||||
stride_mid_oh,
|
||||
stride_mid_os,
|
||||
stride_kpe_tokens_out_b,
|
||||
stride_cos_sin_cache_s,
|
||||
stride_positions_b,
|
||||
rotary_dim: tl.constexpr,
|
||||
kv_lora_rank: tl.constexpr,
|
||||
qk_rope_head_dim: tl.constexpr,
|
||||
kv_group_num: tl.constexpr,
|
||||
q_head_num: tl.constexpr,
|
||||
BLOCK_C: tl.constexpr,
|
||||
BLOCK_R: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_H: tl.constexpr,
|
||||
NUM_KV_SPLITS: tl.constexpr,
|
||||
logit_cap: tl.constexpr,
|
||||
USE_ROPE: tl.constexpr,
|
||||
IS_NEOX_STYLE: tl.constexpr,
|
||||
):
|
||||
|
||||
cur_batch = tl.program_id(0)
|
||||
cur_head_id = tl.program_id(1)
|
||||
split_kv_id = tl.program_id(2)
|
||||
|
||||
if BLOCK_H < kv_group_num:
|
||||
VALID_BLOCK_H: tl.constexpr = BLOCK_H
|
||||
else:
|
||||
VALID_BLOCK_H: tl.constexpr = kv_group_num
|
||||
cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
|
||||
mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
|
||||
mask_h = mask_h & (cur_head < q_head_num)
|
||||
|
||||
offs_c = tl.arange(0, BLOCK_C)
|
||||
offs_qk_r = tl.arange(kv_lora_rank, kv_lora_rank + BLOCK_R) # to get the k_pe
|
||||
|
||||
off_q_pe = (
|
||||
cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_qk_r[None, :]
|
||||
)
|
||||
offs_q = cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_c[None, :]
|
||||
|
||||
mask_c = offs_c < kv_lora_rank
|
||||
mask_qk_r = offs_qk_r < (kv_lora_rank + qk_rope_head_dim)
|
||||
|
||||
cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
|
||||
cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
|
||||
|
||||
q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_c[None, :]), other=0.0)
|
||||
q_pe = tl.load(
|
||||
Q + off_q_pe, mask=(mask_h[:, None]) & (mask_qk_r[None, :]), other=0.0
|
||||
)
|
||||
|
||||
kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
|
||||
split_kv_start = kv_len_per_split * split_kv_id
|
||||
split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
|
||||
|
||||
# apply rotary embedding for q_pe, and k_pe (last token per batch of K_PE)
|
||||
LAST_SPLIT = split_kv_end == cur_batch_seq_len
|
||||
k_pe_last_token = tl.zeros([BLOCK_R], dtype=q.dtype)
|
||||
|
||||
if USE_ROPE:
|
||||
if IS_NEOX_STYLE:
|
||||
# [BLOCK_ROTARY // 2, BLOCK_ROTARY // 2 + 1, BLOCK_ROTARY // 2 + 2, ..., 0, 1, 2, ..., BLOCK_ROTARY // 2 - 1, pass:]
|
||||
offs_qk_rot_r = kv_lora_rank + (
|
||||
(tl.arange(0, BLOCK_R) + (rotary_dim // 2)) % rotary_dim
|
||||
)
|
||||
# Which elements to flip
|
||||
mask_rotate = tl.arange(0, BLOCK_R) < (rotary_dim // 2)
|
||||
# [0 , 1, 2, ..., rotary_dim // 2 - 1, 0 , 1, 2, ..., rotary_dim // 2 - 1]
|
||||
offs_rotary = tl.arange(0, BLOCK_R) % (rotary_dim // 2)
|
||||
else:
|
||||
# [1, 0, 3, 2, 5, 4, ..., BLOCK_R, BLOCK_R - 1]
|
||||
offs_qk_rot_r = (
|
||||
kv_lora_rank
|
||||
+ (((tl.arange(0, BLOCK_R) + 1) % 2) * 2)
|
||||
- 1
|
||||
+ tl.arange(0, BLOCK_R)
|
||||
)
|
||||
mask_rotate = tl.arange(0, BLOCK_R) % 2 < 1
|
||||
# [0, 0, 1, 1, ..., rotary_dim // 2 - 1, rotary_dim // 2 - 1]
|
||||
offs_rotary = tl.arange(0, BLOCK_R) // 2
|
||||
|
||||
if qk_rope_head_dim > rotary_dim:
|
||||
offs_qk_rot_r = tl.where(
|
||||
tl.arange(0, BLOCK_R) < rotary_dim, offs_qk_rot_r, tl.arange(0, BLOCK_R)
|
||||
)
|
||||
offs_rotary = tl.where(
|
||||
tl.arange(0, BLOCK_R) < rotary_dim, offs_rotary, tl.arange(0, BLOCK_R)
|
||||
)
|
||||
|
||||
mask_rotary = tl.arange(0, BLOCK_R) < rotary_dim
|
||||
|
||||
pos = tl.load(positions + cur_batch * stride_positions_b)
|
||||
cos = tl.load(
|
||||
cos_sin_cache + pos * stride_cos_sin_cache_s + offs_rotary,
|
||||
mask=mask_rotary,
|
||||
other=1.0,
|
||||
)
|
||||
sin = tl.load(
|
||||
cos_sin_cache
|
||||
+ pos * stride_cos_sin_cache_s
|
||||
+ offs_rotary
|
||||
+ rotary_dim // 2,
|
||||
mask_rotary,
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
off_q_pe_rot = (
|
||||
cur_batch * stride_qb
|
||||
+ cur_head[:, None] * stride_qh
|
||||
+ offs_qk_rot_r[None, :]
|
||||
)
|
||||
mask_qk_rot_r = offs_qk_rot_r < (kv_lora_rank + qk_rope_head_dim)
|
||||
|
||||
# 0, 2, 4,.... 1, 3, 5...
|
||||
q_pe_rot = tl.load(
|
||||
Q + off_q_pe_rot,
|
||||
mask=(mask_h[:, None]) & (mask_qk_rot_r[None, :]),
|
||||
other=0.0,
|
||||
)
|
||||
q_pe_rot = tl.where(mask_rotate[None, :], -q_pe_rot, q_pe_rot)
|
||||
|
||||
q_pe = q_pe * cos + q_pe_rot * sin
|
||||
|
||||
# we only apply to the last token in the K_PE
|
||||
if LAST_SPLIT:
|
||||
# debug assert
|
||||
if (cur_batch == 0 and cur_head == 0) and split_kv_id < NUM_KV_SPLITS - 1:
|
||||
tl.device_assert(False, "Only last split should compute k_pe")
|
||||
|
||||
kv_loc = tl.load(
|
||||
kv_indices + cur_batch_kv_start_idx + cur_batch_seq_len - 1
|
||||
)
|
||||
offs_buf_k_pe_last_token = kv_loc * stride_buf_kbs + offs_qk_r
|
||||
offs_buf_k_pe_rot_last_token = kv_loc * stride_buf_kbs + offs_qk_rot_r
|
||||
k_pe_last_token = tl.load(K_Buffer + offs_buf_k_pe_last_token)
|
||||
|
||||
k_pe_rot_last_token = tl.load(K_Buffer + offs_buf_k_pe_rot_last_token)
|
||||
k_pe_rot_last_token = tl.where(
|
||||
mask_rotate, -k_pe_rot_last_token, k_pe_rot_last_token
|
||||
)
|
||||
|
||||
k_pe_last_token = k_pe_last_token * cos + k_pe_rot_last_token * sin
|
||||
|
||||
e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
|
||||
e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
|
||||
acc = tl.zeros([BLOCK_H, BLOCK_C], dtype=tl.float32)
|
||||
|
||||
if split_kv_end > split_kv_start:
|
||||
for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
|
||||
offs_n = start_n + tl.arange(0, BLOCK_N)
|
||||
kv_loc = tl.load(
|
||||
kv_indices + cur_batch_kv_start_idx + offs_n,
|
||||
mask=offs_n < split_kv_end,
|
||||
other=0,
|
||||
)
|
||||
|
||||
offs_buf_kv = kv_loc[None, :] * stride_buf_kbs + offs_c[:, None]
|
||||
offs_buf_k_pe = kv_loc[None, :] * stride_buf_kbs + offs_qk_r[:, None]
|
||||
|
||||
k_pe = tl.load(
|
||||
K_Buffer + offs_buf_k_pe,
|
||||
mask=(offs_n[None, :] < split_kv_end) & (mask_qk_r[:, None]),
|
||||
other=0.0,
|
||||
) # positional embedding part of keys
|
||||
|
||||
if USE_ROPE and start_n >= cur_batch_seq_len - BLOCK_N:
|
||||
k_pe = tl.where(
|
||||
offs_n[None, :] != (split_kv_end - 1),
|
||||
k_pe,
|
||||
k_pe_last_token[:, None],
|
||||
)
|
||||
|
||||
# (16, 64) x (64, 32)
|
||||
# dot product of rope parts
|
||||
qk = tl.dot(q_pe, k_pe.to(q_pe.dtype))
|
||||
|
||||
kv = tl.load(
|
||||
K_Buffer + offs_buf_kv,
|
||||
mask=(offs_n[None, :] < split_kv_end) & (mask_c[:, None]),
|
||||
other=0.0,
|
||||
) # the shared latent tensor for keys and values
|
||||
|
||||
# (16, 512) x (512, 32)
|
||||
# dot product of nope parts
|
||||
qk += tl.dot(q, kv)
|
||||
|
||||
qk *= sm_scale
|
||||
|
||||
if logit_cap > 0:
|
||||
qk = logit_cap * tanh(qk / logit_cap)
|
||||
|
||||
qk = tl.where(
|
||||
mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
|
||||
)
|
||||
|
||||
offs_buf_v = kv_loc[:, None] * stride_buf_vbs + offs_c[None, :]
|
||||
v = tl.load(
|
||||
V_buffer + offs_buf_v,
|
||||
mask=(offs_n[:, None] < split_kv_end) & (mask_c[None, :]),
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
n_e_max = tl.maximum(tl.max(qk, 1), e_max)
|
||||
re_scale = tl.exp(e_max - n_e_max)
|
||||
p = tl.exp(qk - n_e_max[:, None])
|
||||
acc *= re_scale[:, None]
|
||||
# (16, 32) x (32, 512)
|
||||
acc += tl.dot(p.to(v.dtype), v)
|
||||
|
||||
e_sum = e_sum * re_scale + tl.sum(p, 1)
|
||||
e_max = n_e_max
|
||||
|
||||
offs_mid_o = (
|
||||
cur_batch * stride_mid_ob
|
||||
+ cur_head[:, None] * stride_mid_oh
|
||||
+ split_kv_id * stride_mid_os
|
||||
+ offs_c[None, :]
|
||||
)
|
||||
|
||||
if USE_ROPE:
|
||||
if LAST_SPLIT:
|
||||
k_pe_last_token_ptrs = (
|
||||
k_pe_t_out
|
||||
+ cur_batch * stride_kpe_tokens_out_b
|
||||
+ tl.arange(0, BLOCK_R)
|
||||
)
|
||||
tl.store(k_pe_last_token_ptrs, k_pe_last_token, mask=mask_qk_r)
|
||||
|
||||
tl.store(
|
||||
Att_Out + offs_mid_o,
|
||||
acc / e_sum[:, None],
|
||||
mask=(mask_h[:, None]) & (mask_c[None, :]),
|
||||
)
|
||||
|
||||
offs_mid_o_1 = (
|
||||
cur_batch * stride_mid_ob
|
||||
+ cur_head * stride_mid_oh
|
||||
+ split_kv_id * stride_mid_os
|
||||
+ kv_lora_rank
|
||||
)
|
||||
|
||||
tl.store(
|
||||
Att_Out + offs_mid_o_1,
|
||||
e_max + tl.log(e_sum),
|
||||
mask=mask_h,
|
||||
)
|
||||
|
||||
|
||||
# TODO rope offset
|
||||
def _decode_grouped_att_m_fwd_rope(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
att_out,
|
||||
k_pe_tokens_out,
|
||||
kv_lora_rank, # c
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
rotary_dim,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap,
|
||||
use_rope,
|
||||
is_neox_style=True,
|
||||
):
|
||||
if use_rope:
|
||||
assert (
|
||||
k_pe_tokens_out is not None
|
||||
), "We must output the k_pe tokens with rope applied if rope fusion enabled."
|
||||
|
||||
BLOCK = 32
|
||||
|
||||
# # [TODO] work around shmem limit on MI3xx
|
||||
# if is_hip_ and kv_lora_rank >= 576:
|
||||
# BLOCK = 16
|
||||
|
||||
qk_rope_head_dim = k_buffer.shape[-1] - kv_lora_rank
|
||||
batch, head_num = kv_indptr.shape[0] - 1, q.shape[1]
|
||||
kv_group_num = q.shape[1] // k_buffer.shape[1]
|
||||
|
||||
BLOCK_C = triton.next_power_of_2(kv_lora_rank)
|
||||
BLOCK_R = triton.next_power_of_2(qk_rope_head_dim)
|
||||
|
||||
BLOCK_H = 16
|
||||
NUM_KV_SPLITS = num_kv_splits
|
||||
grid = (
|
||||
batch,
|
||||
triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
|
||||
NUM_KV_SPLITS,
|
||||
)
|
||||
|
||||
extra_kargs = {}
|
||||
num_stages = 2
|
||||
if is_hip_:
|
||||
# https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
|
||||
# https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
|
||||
extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
|
||||
num_stages = 1
|
||||
|
||||
_fwd_grouped_kernel_stage1_rope[grid](
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
sm_scale,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
att_out,
|
||||
k_pe_tokens_out,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
k_buffer.stride(0),
|
||||
v_buffer.stride(0),
|
||||
att_out.stride(0),
|
||||
att_out.stride(1),
|
||||
att_out.stride(2),
|
||||
k_pe_tokens_out.stride(0) if use_rope else 0,
|
||||
cos_sin_cache.stride(0) if use_rope else 0,
|
||||
positions.stride(0) if use_rope else 0,
|
||||
rotary_dim,
|
||||
kv_lora_rank,
|
||||
qk_rope_head_dim,
|
||||
kv_group_num=kv_group_num,
|
||||
q_head_num=head_num,
|
||||
BLOCK_C=BLOCK_C,
|
||||
BLOCK_R=BLOCK_R,
|
||||
BLOCK_N=BLOCK,
|
||||
BLOCK_H=BLOCK_H,
|
||||
NUM_KV_SPLITS=NUM_KV_SPLITS,
|
||||
logit_cap=logit_cap,
|
||||
USE_ROPE=use_rope,
|
||||
IS_NEOX_STYLE=is_neox_style,
|
||||
num_warps=4,
|
||||
num_stages=num_stages,
|
||||
**extra_kargs
|
||||
)
|
||||
|
||||
|
||||
def decode_attention_fwd_grouped_rope(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
k_pe_tokens,
|
||||
kv_lora_rank,
|
||||
rotary_dim,
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
attn_logits,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap=0.0,
|
||||
use_rope=False,
|
||||
is_neox_style=False,
|
||||
):
|
||||
_decode_grouped_att_m_fwd_rope(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
attn_logits,
|
||||
k_pe_tokens,
|
||||
kv_lora_rank,
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
rotary_dim,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap,
|
||||
use_rope,
|
||||
is_neox_style,
|
||||
)
|
||||
_decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, kv_indptr, num_kv_splits)
|
||||
160
python/sglang/srt/models/deepseek_v2.py
Normal file → Executable file
160
python/sglang/srt/models/deepseek_v2.py
Normal file → Executable file
@@ -16,6 +16,7 @@
|
||||
# https://github.com/vllm-project/vllm/blob/fb6af8bc086328ca6659e72d11ffd4309ce4de22/vllm/model_executor/models/deepseek_v2.py
|
||||
"""Inference-only DeepseekV2 model."""
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
@@ -31,6 +32,9 @@ from sglang.srt.distributed import (
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.layers.activation import SiluAndMul
|
||||
from sglang.srt.layers.attention.triton_ops.rocm_mla_decode_rope import (
|
||||
decode_attention_fwd_grouped_rope,
|
||||
)
|
||||
from sglang.srt.layers.layernorm import RMSNorm
|
||||
from sglang.srt.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
@@ -533,7 +537,18 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
if no_absorb():
|
||||
return self.forward_normal(positions, hidden_states, forward_batch)
|
||||
else:
|
||||
return self.forward_absorb(positions, hidden_states, forward_batch)
|
||||
if is_hip_:
|
||||
if (
|
||||
os.getenv("SGLANG_ROCM_FUSED_DECODE_MLA") == "1"
|
||||
and forward_batch.forward_mode.is_decode()
|
||||
):
|
||||
return self.forward_absorb_fused_mla_rope(
|
||||
positions, hidden_states, forward_batch
|
||||
)
|
||||
else:
|
||||
return self.forward_absorb(positions, hidden_states, forward_batch)
|
||||
else:
|
||||
return self.forward_absorb(positions, hidden_states, forward_batch)
|
||||
|
||||
def forward_normal(
|
||||
self,
|
||||
@@ -652,6 +667,149 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
|
||||
return output
|
||||
|
||||
def forward_absorb_fused_mla_rope(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
) -> torch.Tensor:
|
||||
enable_rope_fusion = (
|
||||
os.getenv("SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION", "1") == "1"
|
||||
)
|
||||
q_len = hidden_states.shape[0]
|
||||
q_input = hidden_states.new_empty(
|
||||
q_len, self.num_local_heads, self.kv_lora_rank + self.qk_rope_head_dim
|
||||
)
|
||||
if self.q_lora_rank is not None:
|
||||
q = self.q_a_proj(hidden_states)[0]
|
||||
q = self.q_a_layernorm(q)
|
||||
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
|
||||
)
|
||||
q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
|
||||
|
||||
if self.w_kc.dtype == torch.float8_e4m3fnuz:
|
||||
# TODO(kernel): add bmm_fp8 for torch.float8_e4m3fnuz
|
||||
q_nope_out = torch.bmm(
|
||||
q_nope.to(torch.bfloat16).transpose(0, 1),
|
||||
self.w_kc.to(torch.bfloat16) * self.w_scale,
|
||||
)
|
||||
elif self.w_kc.dtype == torch.float8_e4m3fn:
|
||||
q_nope_val, q_nope_scale = input_to_float8(
|
||||
q_nope.transpose(0, 1), torch.float8_e4m3fn
|
||||
)
|
||||
q_nope_out = bmm_fp8(
|
||||
q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16
|
||||
)
|
||||
else:
|
||||
q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)
|
||||
q_input[..., : self.kv_lora_rank] = q_nope_out.transpose(0, 1)
|
||||
|
||||
latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
|
||||
v_input = latent_cache[..., : self.kv_lora_rank]
|
||||
v_input = self.kv_a_layernorm(v_input.contiguous()).unsqueeze(1)
|
||||
k_input = latent_cache.unsqueeze(1)
|
||||
k_input[..., : self.kv_lora_rank] = v_input
|
||||
|
||||
if not enable_rope_fusion:
|
||||
k_pe = k_input[..., self.kv_lora_rank :]
|
||||
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
|
||||
q_input[..., self.kv_lora_rank :] = q_pe
|
||||
k_input[..., self.kv_lora_rank :] = k_pe
|
||||
k_pe_output = None
|
||||
else:
|
||||
k_pe_output = torch.empty_like(k_input[..., self.kv_lora_rank :])
|
||||
|
||||
q_input[..., self.kv_lora_rank :] = q_pe
|
||||
|
||||
# attn_output = self.attn_mqa(q_input, k_input, v_input, forward_batch)
|
||||
# Use Fused ROPE with use_rope=OFF.
|
||||
attn_output = torch.empty(
|
||||
(q_len, self.num_local_heads, self.kv_lora_rank),
|
||||
dtype=q.dtype,
|
||||
device=q.device,
|
||||
)
|
||||
attn_logits, _, kv_indptr, kv_indices, _, _, _ = (
|
||||
forward_batch.attn_backend.forward_metadata
|
||||
)
|
||||
cos_sin_cache = self.rotary_emb.cos_sin_cache
|
||||
num_kv_split = forward_batch.attn_backend.num_kv_splits
|
||||
sm_scale = self.attn_mqa.scaling
|
||||
if attn_logits is None:
|
||||
attn_logits = torch.empty(
|
||||
(
|
||||
forward_batch.batch_size,
|
||||
self.num_local_heads,
|
||||
num_kv_split,
|
||||
self.kv_lora_rank + 1,
|
||||
),
|
||||
dtype=torch.float32,
|
||||
device=q.device,
|
||||
)
|
||||
|
||||
# save current latent cache.
|
||||
forward_batch.token_to_kv_pool.set_kv_buffer(
|
||||
self.attn_mqa, forward_batch.out_cache_loc, k_input, None
|
||||
)
|
||||
key_cache_buf = forward_batch.token_to_kv_pool.get_key_buffer(
|
||||
self.attn_mqa.layer_id
|
||||
)
|
||||
val_cache_buf = key_cache_buf[..., : self.kv_lora_rank]
|
||||
|
||||
decode_attention_fwd_grouped_rope(
|
||||
q_input,
|
||||
key_cache_buf,
|
||||
val_cache_buf,
|
||||
attn_output,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
k_pe_output,
|
||||
self.kv_lora_rank,
|
||||
self.rotary_emb.rotary_dim,
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
attn_logits,
|
||||
num_kv_split,
|
||||
sm_scale,
|
||||
logit_cap=self.attn_mqa.logit_cap,
|
||||
use_rope=enable_rope_fusion,
|
||||
is_neox_style=self.rotary_emb.is_neox_style,
|
||||
)
|
||||
|
||||
if enable_rope_fusion:
|
||||
k_input[..., self.kv_lora_rank :] = k_pe_output
|
||||
forward_batch.token_to_kv_pool.set_kv_buffer(
|
||||
self.attn_mqa, forward_batch.out_cache_loc, k_input, None
|
||||
)
|
||||
|
||||
attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
|
||||
|
||||
if self.w_vc.dtype == torch.float8_e4m3fnuz:
|
||||
# TODO(kernel): add bmm_fp8 for torch.float8_e4m3fnuz
|
||||
attn_bmm_output = torch.bmm(
|
||||
attn_output.to(torch.bfloat16).transpose(0, 1),
|
||||
self.w_vc.to(torch.bfloat16) * self.w_scale,
|
||||
)
|
||||
elif self.w_vc.dtype == torch.float8_e4m3fn:
|
||||
attn_output_val, attn_output_scale = input_to_float8(
|
||||
attn_output.transpose(0, 1), torch.float8_e4m3fn
|
||||
)
|
||||
attn_bmm_output = bmm_fp8(
|
||||
attn_output_val,
|
||||
self.w_vc,
|
||||
attn_output_scale,
|
||||
self.w_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
else:
|
||||
attn_bmm_output = torch.bmm(attn_output.transpose(0, 1), self.w_vc)
|
||||
attn_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
|
||||
output, _ = self.o_proj(attn_output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def all_gather(
|
||||
input_tensor: torch.Tensor, forward_batch: ForwardBatch, rank, world_size, group
|
||||
|
||||
Reference in New Issue
Block a user