chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Refactor attention backend (#1381)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-09-11 11:44:26 -07:00
committed by GitHub
parent c03cece42f
commit fec185ce0c
16 changed files with 568 additions and 564 deletions
Download Patch File Download Diff File Expand all files Collapse all files

383
python/sglang/srt/layers/attention_backend.py Normal file
View File

@@ -0,0 +1,383 @@
from __future__ import annotations
"""
Support different attention backends.
Now there are two backends: FlashInfer and Triton.
FlashInfer is faster and Triton is easier to customize.
Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode.
"""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
import torch
import torch.nn as nn
from flashinfer import (
BatchDecodeWithPagedKVCacheWrapper,
BatchPrefillWithPagedKVCacheWrapper,
BatchPrefillWithRaggedKVCacheWrapper,
)
from flashinfer.cascade import merge_state
from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
from sglang.global_config import global_config
from sglang.srt.layers.flashinfer_utils import update_flashinfer_indices
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.model_executor.forward_batch_info import ForwardMode, InputMetadata
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
class AttentionBackend(ABC):
"""The base class of attention backends"""
@abstractmethod
def init_forward_metadata(
self, batch: ScheduleBatch, input_metadata: InputMetadata
):
pass
def forward(self, q, k, v, layer, input_metadata: InputMetadata):
if input_metadata.forward_mode.is_decode():
return self.forward_decode(q, k, v, layer, input_metadata)
else:
return self.forward_extend(q, k, v, layer, input_metadata)
class FlashInferAttnBackend(AttentionBackend):
"""Flashinfer attention kernels."""
def __init__(self, model_runner: ModelRunner):
super().__init__()
self.model_runner = model_runner
if not _grouped_size_compiled_for_decode_kernels(
model_runner.model_config.num_attention_heads // model_runner.tp_size,
model_runner.model_config.get_num_kv_heads(model_runner.tp_size),
):
self.decode_use_tensor_cores = True
else:
self.decode_use_tensor_cores = False
self.workspace_buffer = torch.empty(
global_config.flashinfer_workspace_size,
dtype=torch.uint8,
device="cuda",
)
if model_runner.sliding_window_size is None:
self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
self.workspace_buffer, "NHD"
)
self.prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer, "NHD"
)
self.decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
use_tensor_cores=self.decode_use_tensor_cores,
)
else:
# Two wrappers: one for sliding window attention and one for full attention.
# Using two wrappers is unnecessary in the current PR, but are prepared for future PRs
self.prefill_wrapper_ragged = None
self.prefill_wrapper_paged = []
self.decode_wrapper = []
for _ in range(2):
self.prefill_wrapper_paged.append(
BatchPrefillWithPagedKVCacheWrapper(self.workspace_buffer, "NHD")
)
self.decode_wrapper.append(
BatchDecodeWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
use_tensor_cores=self.decode_use_tensor_cores,
)
)
self.forward_metadata = None
self.cuda_graph_metadata = {}
def init_forward_metadata(
self, batch: ScheduleBatch, input_metadata: InputMetadata
):
if input_metadata.forward_mode.is_decode():
prefix_lens = None
use_ragged = False
total_num_tokens = None
else:
prefix_lens = input_metadata.extend_prefix_lens
# Some heuristics to check whether to use ragged forward
use_ragged = False
if (
int(torch.sum(input_metadata.seq_lens)) > 4096
and self.model_runner.sliding_window_size is None
):
use_ragged = True
total_num_tokens = torch.sum(input_metadata.seq_lens).item()
update_flashinfer_indices(
input_metadata.forward_mode,
self.model_runner,
input_metadata.req_pool_indices,
input_metadata.seq_lens,
prefix_lens,
use_ragged=use_ragged,
)
self.forward_metadata = (use_ragged, total_num_tokens, self.decode_wrapper)
def init_cuda_graph_state(self, max_bs: int):
self.cuda_graph_kv_indptr = torch.zeros(
(max_bs + 1,), dtype=torch.int32, device="cuda"
)
self.cuda_graph_kv_indices = torch.zeros(
(max_bs * self.model_runner.model_config.context_len,),
dtype=torch.int32,
device="cuda",
)
self.cuda_graph_kv_last_page_len = torch.ones(
(max_bs,), dtype=torch.int32, device="cuda"
)
if self.model_runner.sliding_window_size is not None:
self.cuda_graph_kv_indptr = [
self.cuda_graph_kv_indptr,
self.cuda_graph_kv_indptr.clone(),
]
self.cuda_graph_kv_indices = [
self.cuda_graph_kv_indices,
self.cuda_graph_kv_indices.clone(),
]
def capture_cuda_graph_init(self, bs: int, req_pool_indices, seq_lens):
if self.model_runner.sliding_window_size is None:
decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
use_cuda_graph=True,
use_tensor_cores=self.decode_use_tensor_cores,
paged_kv_indptr_buffer=self.cuda_graph_kv_indptr[: bs + 1],
paged_kv_indices_buffer=self.cuda_graph_kv_indices,
paged_kv_last_page_len_buffer=self.cuda_graph_kv_last_page_len[:bs],
)
else:
decode_wrapper = []
for i in range(2):
decode_wrapper.append(
BatchDecodeWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
use_cuda_graph=True,
use_tensor_cores=self.decode_use_tensor_cores,
paged_kv_indptr_buffer=self.cuda_graph_kv_indptr[i][: bs + 1],
paged_kv_indices_buffer=self.cuda_graph_kv_indices[i],
paged_kv_last_page_len_buffer=self.cuda_graph_kv_last_page_len[
:bs
],
)
)
update_flashinfer_indices(
ForwardMode.DECODE,
self.model_runner,
req_pool_indices,
seq_lens,
None,
decode_wrapper,
)
self.cuda_graph_metadata[bs] = decode_wrapper
self.forward_metadata = (False, None, decode_wrapper)
def replay_cuda_graph_init(self, bs: int, req_pool_indices, seq_lens):
update_flashinfer_indices(
ForwardMode.DECODE,
self.model_runner,
req_pool_indices[:bs],
seq_lens[:bs],
None,
self.cuda_graph_metadata[bs],
)
def forward_extend(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
if not isinstance(self.prefill_wrapper_paged, list):
prefill_wrapper_paged = self.prefill_wrapper_paged
else:
if layer.sliding_window_size != -1:
prefill_wrapper_paged = self.prefill_wrapper_paged[0]
else:
prefill_wrapper_paged = self.prefill_wrapper_paged[1]
use_ragged, total_num_tokens, decode_wrapper = self.forward_metadata
if not use_ragged:
if k is not None:
assert v is not None
input_metadata.token_to_kv_pool.set_kv_buffer(
layer.layer_id, input_metadata.out_cache_loc, k, v
)
o = prefill_wrapper_paged.forward(
q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
input_metadata.token_to_kv_pool.get_kv_buffer(layer.layer_id),
causal=True,
sm_scale=layer.scaling,
window_left=layer.sliding_window_size,
logits_soft_cap=layer.logit_cap,
)
else:
o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(
q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
k.contiguous().view(-1, layer.tp_k_head_num, layer.head_dim),
v.contiguous().view(-1, layer.tp_v_head_num, layer.head_dim),
causal=True,
sm_scale=layer.scaling,
logits_soft_cap=layer.logit_cap,
)
if input_metadata.extend_no_prefix:
o = o1
else:
o2, s2 = prefill_wrapper_paged.forward_return_lse(
q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
input_metadata.token_to_kv_pool.get_kv_buffer(layer.layer_id),
causal=False,
sm_scale=layer.scaling,
logits_soft_cap=layer.logit_cap,
)
o, _ = merge_state(o1, s1, o2, s2)
input_metadata.token_to_kv_pool.set_kv_buffer(
layer.layer_id, input_metadata.out_cache_loc, k, v
)
if total_num_tokens >= global_config.layer_sync_threshold:
torch.cuda.synchronize()
return o.view(-1, layer.tp_q_head_num * layer.head_dim)
def forward_decode(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
use_ragged, total_num_tokens, decode_wrapper = self.forward_metadata
if isinstance(decode_wrapper, list):
if layer.sliding_window_size != -1:
decode_wrapper = decode_wrapper[0]
else:
decode_wrapper = decode_wrapper[1]
if k is not None:
assert v is not None
input_metadata.token_to_kv_pool.set_kv_buffer(
layer.layer_id, input_metadata.out_cache_loc, k, v
)
o = decode_wrapper.forward(
q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
input_metadata.token_to_kv_pool.get_kv_buffer(layer.layer_id),
sm_scale=layer.scaling,
logits_soft_cap=layer.logit_cap,
)
return o.view(-1, layer.tp_q_head_num * layer.head_dim)
class TritonAttnBackend(AttentionBackend):
def __init__(self, model_runner: ModelRunner):
# Lazy import to avoid the initialization of cuda context
from sglang.srt.layers.triton_attention.decode_attention import (
decode_attention_fwd,
)
from sglang.srt.layers.triton_attention.extend_attention import (
extend_attention_fwd,
)
super().__init__()
self.decode_attention_fwd = decode_attention_fwd
self.extend_attention_fwd = extend_attention_fwd
self.forward_metadata = None
def init_forward_metadata(
self, batch: ScheduleBatch, input_metadata: InputMetadata
):
"""Init auxiliary variables for triton attention backend."""
if input_metadata.forward_mode.is_decode():
max_seq_len = torch.max(input_metadata.seq_lens).item()
start_loc = torch.zeros_like(input_metadata.seq_lens, dtype=torch.int32)
start_loc[1:] = torch.cumsum(input_metadata.seq_lens[:-1], dim=0)
total_num_tokens = torch.sum(input_metadata.seq_lens).item()
max_extend_len = None
else:
start_loc = max_seq_len = total_num_tokens = None
prefix_lens = torch.tensor(batch.prefix_lens_cpu, device="cuda")
max_extend_len = torch.max(input_metadata.seq_lens - prefix_lens).item()
self.forward_metadata = start_loc, max_seq_len, max_extend_len, total_num_tokens
def forward_extend(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
if layer.qk_head_dim != layer.v_head_dim:
o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else:
o = torch.empty_like(q)
input_metadata.token_to_kv_pool.set_kv_buffer(
layer.layer_id, input_metadata.out_cache_loc, k, v
)
start_loc, max_seq_len, max_extend_len, total_num_tokens = self.forward_metadata
self.extend_attention_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k.contiguous(),
v.contiguous(),
o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
input_metadata.token_to_kv_pool.get_key_buffer(layer.layer_id),
input_metadata.token_to_kv_pool.get_value_buffer(layer.layer_id),
input_metadata.req_to_token_pool.req_to_token,
input_metadata.req_pool_indices,
input_metadata.seq_lens,
input_metadata.extend_seq_lens,
input_metadata.extend_start_loc,
max_extend_len,
layer.scaling,
layer.logit_cap,
)
return o
def forward_decode(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
if layer.qk_head_dim != layer.v_head_dim:
o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else:
o = torch.empty_like(q)
start_loc, max_seq_len, max_extend_len, total_num_tokens = self.forward_metadata
input_metadata.token_to_kv_pool.set_kv_buffer(
layer.layer_id, input_metadata.out_cache_loc, k, v
)
self.decode_attention_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
input_metadata.token_to_kv_pool.get_key_buffer(layer.layer_id),
input_metadata.token_to_kv_pool.get_value_buffer(layer.layer_id),
o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
input_metadata.req_to_token_pool.req_to_token,
input_metadata.req_pool_indices,
start_loc,
input_metadata.seq_lens,
max_seq_len,
total_num_tokens,
layer.scaling,
layer.logit_cap,
)
return o

72
python/sglang/srt/layers/flashinfer_utils.py
View File

@@ -10,8 +10,8 @@ def create_flashinfer_kv_indices_triton(
page_kernel_lens_ptr,
kv_indptr,
kv_start_idx,
max_context_len,
kv_indices_ptr,
max_context_len: tl.constexpr,
):
BLOCK_SIZE: tl.constexpr = 512
pid = tl.program_id(axis=0)
@@ -47,15 +47,15 @@ class FlashinferUpdater:
req_pool_indices,
seq_lens,
prefix_lens,
flashinfer_decode_wrapper=None,
flashinfer_use_ragged=False,
decode_wrapper=None,
use_ragged=False,
):
self.forward_mode = forward_mode
self.model_runner = model_runner
self.req_pool_indices = req_pool_indices
self.seq_lens = seq_lens
self.prefix_lens = prefix_lens
self.flashinfer_use_ragged = flashinfer_use_ragged
self.use_ragged = use_ragged
self.num_qo_heads = (
model_runner.model_config.num_attention_heads // model_runner.tp_size
@@ -71,20 +71,17 @@ class FlashinferUpdater:
)
(
self.flashinfer_decode_wrapper,
self.flashinfer_prefill_wrapper_ragged,
self.flashinfer_prefill_wrapper_paged,
self.decode_wrapper,
self.prefill_wrapper_ragged,
self.prefill_wrapper_paged,
) = (
flashinfer_decode_wrapper,
self.model_runner.flashinfer_prefill_wrapper_ragged,
self.model_runner.flashinfer_prefill_wrapper_paged,
decode_wrapper or self.model_runner.attn_backend.decode_wrapper,
self.model_runner.attn_backend.prefill_wrapper_ragged,
self.model_runner.attn_backend.prefill_wrapper_paged,
)
# CUDA graph uses different flashinfer_decode_wrapper
if self.flashinfer_decode_wrapper is None:
self.flashinfer_decode_wrapper = self.model_runner.flashinfer_decode_wrapper
def _init_indices_no_window(self):
if self.flashinfer_use_ragged:
def _init_indices_no_sliding_window(self):
if self.use_ragged:
paged_kernel_lens = self.prefix_lens
else:
paged_kernel_lens = self.seq_lens
@@ -103,13 +100,13 @@ class FlashinferUpdater:
paged_kernel_lens,
self.kv_indptr,
None,
self.model_runner.req_to_token_pool.req_to_token.size(1),
self.kv_indices,
self.model_runner.req_to_token_pool.req_to_token.size(1),
)
def _init_indices_window(self, wrapper_id):
# window attention use paged only
def _init_indices_sliding_window(self, wrapper_id):
if wrapper_id == 0:
# window attention use paged only
if self.forward_mode.is_decode():
paged_kernel_lens = torch.minimum(
self.seq_lens,
@@ -123,6 +120,7 @@ class FlashinferUpdater:
- self.prefix_lens,
)
else:
# full attention
paged_kernel_lens = self.seq_lens
kv_start_idx = self.seq_lens - paged_kernel_lens
@@ -139,8 +137,8 @@ class FlashinferUpdater:
paged_kernel_lens,
self.kv_indptr,
kv_start_idx,
self.model_runner.req_to_token_pool.req_to_token.size(1),
self.kv_indices,
self.model_runner.req_to_token_pool.req_to_token.size(1),
)
def _update_decode_indices(self, decode_wrapper):
@@ -164,7 +162,7 @@ class FlashinferUpdater:
)
qo_indptr[1:] = torch.cumsum(self.seq_lens - self.prefix_lens, dim=0)
if self.flashinfer_use_ragged:
if self.use_ragged:
ragged_wrapper.end_forward()
ragged_wrapper.begin_forward(
qo_indptr,
@@ -187,28 +185,28 @@ class FlashinferUpdater:
1,
)
def update_indices_no_window(self):
self._init_indices_no_window()
def update_indices_no_sliding_window(self):
self._init_indices_no_sliding_window()
if self.forward_mode.is_decode():
self._update_decode_indices(self.flashinfer_decode_wrapper)
self._update_decode_indices(self.decode_wrapper)
else:
self._update_extend_indices(
self.flashinfer_prefill_wrapper_ragged,
self.flashinfer_prefill_wrapper_paged,
self.prefill_wrapper_ragged,
self.prefill_wrapper_paged,
)
def update_indices_window(self):
assert self.flashinfer_use_ragged is False
def update_indices_sliding_window(self):
assert self.use_ragged is False
for wrapper_id in range(2):
self._init_indices_window(wrapper_id)
self._init_indices_sliding_window(wrapper_id)
if self.forward_mode.is_decode():
self._update_decode_indices(self.flashinfer_decode_wrapper[wrapper_id])
self._update_decode_indices(self.decode_wrapper[wrapper_id])
else:
self._update_extend_indices(
None,
self.flashinfer_prefill_wrapper_paged[wrapper_id],
self.prefill_wrapper_paged[wrapper_id],
)
@@ -218,20 +216,20 @@ def update_flashinfer_indices(
req_pool_indices,
seq_lens,
prefix_lens,
flashinfer_decode_wrapper=None,
flashinfer_use_ragged=False,
decode_wrapper=None,
use_ragged=False,
):
flashinfer_updater = FlashinferUpdater(
updater = FlashinferUpdater(
forward_mode,
model_runner,
req_pool_indices,
seq_lens,
prefix_lens,
flashinfer_decode_wrapper,
flashinfer_use_ragged,
decode_wrapper,
use_ragged,
)
if model_runner.sliding_window_size is None:
flashinfer_updater.update_indices_no_window()
updater.update_indices_no_sliding_window()
else:
flashinfer_updater.update_indices_window()
updater.update_indices_sliding_window()

175
python/sglang/srt/layers/radix_attention.py
View File

@@ -15,25 +15,14 @@ limitations under the License.
"""Radix attention."""
from typing import Optional
import torch
from flashinfer.cascade import merge_state
from torch import nn
from sglang.global_config import global_config
from sglang.srt.layers.triton_attention.decode_attention import decode_attention_fwd
from sglang.srt.layers.triton_attention.extend_attention import extend_attention_fwd
from sglang.srt.model_executor.forward_batch_info import ForwardMode, InputMetadata
from sglang.srt.model_executor.model_runner import global_server_args_dict
from sglang.srt.model_executor.forward_batch_info import InputMetadata
class RadixAttention(nn.Module):
"""
The attention layer implementation.
Now it has two backends: FlashInfer and Triton.
FlashInfer is faster and Triton is easier to customize.
It supports two operators: extend (i.e. prefill with cached prefix) and decode.
"""
def __init__(
@@ -43,8 +32,8 @@ class RadixAttention(nn.Module):
scaling: float,
num_kv_heads: int,
layer_id: int,
sliding_window_size: Optional[int] = None,
logit_cap: int = -1,
sliding_window_size: int = -1,
logit_cap: float = 0.0,
v_head_dim: int = -1,
):
super().__init__()
@@ -56,164 +45,14 @@ class RadixAttention(nn.Module):
self.v_head_dim = v_head_dim if v_head_dim != -1 else head_dim
self.scaling = scaling
self.layer_id = layer_id
self.logit_cap = logit_cap if logit_cap is not None and logit_cap > 0 else 0
self.sliding_window_size = sliding_window_size if sliding_window_size else -1
# Choose backend
if (
global_server_args_dict["attention_backend"] == "flashinfer"
and self.qk_head_dim == self.v_head_dim
):
self.extend_forward = self.extend_forward_flashinfer
self.decode_forward = self.decode_forward_flashinfer
elif global_server_args_dict["attention_backend"] == "triton":
self.extend_forward = self.extend_forward_triton
self.decode_forward = self.decode_forward_triton
else:
raise ValueError(
f"Invalid attention backend: {global_server_args_dict['attention_backend']}"
)
def extend_forward_triton(self, q, k, v, input_metadata: InputMetadata):
if self.qk_head_dim != self.v_head_dim:
o = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
else:
o = torch.empty_like(q)
self.store_kv_cache(k, v, input_metadata)
extend_attention_fwd(
q.view(-1, self.tp_q_head_num, self.qk_head_dim),
k.contiguous(),
v.contiguous(),
o.view(-1, self.tp_q_head_num, self.v_head_dim),
input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id),
input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id),
input_metadata.req_to_token_pool.req_to_token,
input_metadata.req_pool_indices,
input_metadata.triton_start_loc,
input_metadata.seq_lens,
input_metadata.triton_prefix_lens,
input_metadata.extend_start_loc,
input_metadata.extend_seq_lens,
input_metadata.triton_max_seq_len,
input_metadata.triton_max_extend_len,
sm_scale=self.scaling,
logit_cap=self.logit_cap,
)
return o
def decode_forward_triton(self, q, k, v, input_metadata: InputMetadata):
if self.qk_head_dim != self.v_head_dim:
o = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
else:
o = torch.empty_like(q)
self.store_kv_cache(k, v, input_metadata)
decode_attention_fwd(
q.view(-1, self.tp_q_head_num, self.qk_head_dim),
input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id),
input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id),
o.view(-1, self.tp_q_head_num, self.v_head_dim),
input_metadata.req_to_token_pool.req_to_token,
input_metadata.req_pool_indices,
input_metadata.triton_start_loc,
input_metadata.seq_lens,
input_metadata.triton_max_seq_len,
input_metadata.total_num_tokens,
sm_scale=self.scaling,
logit_cap=self.logit_cap,
)
return o
def extend_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
# using two wrappers is unnecessary in the current PR, but are prepared for future PRs
prefill_wrapper_paged = input_metadata.flashinfer_prefill_wrapper_paged
if self.sliding_window_size != -1:
prefill_wrapper_paged = prefill_wrapper_paged[0]
else:
if isinstance(prefill_wrapper_paged, list):
prefill_wrapper_paged = prefill_wrapper_paged[1]
if not input_metadata.flashinfer_use_ragged:
if k is not None:
assert v is not None
self.store_kv_cache(k, v, input_metadata)
o = prefill_wrapper_paged.forward(
q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
causal=True,
sm_scale=self.scaling,
window_left=self.sliding_window_size,
logits_soft_cap=self.logit_cap,
)
else:
o1, s1 = (
input_metadata.flashinfer_prefill_wrapper_ragged.forward_return_lse(
q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
k.contiguous().view(-1, self.tp_k_head_num, self.head_dim),
v.contiguous().view(-1, self.tp_v_head_num, self.head_dim),
causal=True,
sm_scale=self.scaling,
logits_soft_cap=self.logit_cap,
)
)
if input_metadata.extend_no_prefix:
o = o1
else:
o2, s2 = prefill_wrapper_paged.forward_return_lse(
q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
causal=False,
sm_scale=self.scaling,
logits_soft_cap=self.logit_cap,
)
o, _ = merge_state(o1, s1, o2, s2)
self.store_kv_cache(k, v, input_metadata)
if input_metadata.total_num_tokens >= global_config.layer_sync_threshold:
torch.cuda.synchronize()
return o.view(-1, self.tp_q_head_num * self.head_dim)
def decode_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
decode_wrapper = input_metadata.flashinfer_decode_wrapper
if self.sliding_window_size != -1:
decode_wrapper = decode_wrapper[0]
else:
if isinstance(decode_wrapper, list):
decode_wrapper = decode_wrapper[1]
if k is not None:
assert v is not None
self.store_kv_cache(k, v, input_metadata)
o = decode_wrapper.forward(
q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
sm_scale=self.scaling,
logits_soft_cap=self.logit_cap,
)
return o.view(-1, self.tp_q_head_num * self.head_dim)
self.logit_cap = logit_cap
self.sliding_window_size = sliding_window_size or -1
def forward(self, q, k, v, input_metadata: InputMetadata):
if k is not None:
# For cross-layer sharing, kv can be None
assert v is not None
k = k.view(-1, self.tp_k_head_num, self.qk_head_dim)
v = v.view(-1, self.tp_v_head_num, self.v_head_dim)
if input_metadata.forward_mode.is_extend():
return self.extend_forward(q, k, v, input_metadata)
elif input_metadata.forward_mode.is_decode():
return self.decode_forward(q, k, v, input_metadata)
def store_kv_cache(self, cache_k, cache_v, input_metadata: InputMetadata):
input_metadata.token_to_kv_pool.set_kv_buffer(
self.layer_id, input_metadata.out_cache_loc, cache_k, cache_v
)
return input_metadata.attn_backend.forward(q, k, v, self, input_metadata)

7
python/sglang/srt/layers/triton_attention/decode_attention.py
View File

@@ -15,6 +15,7 @@ limitations under the License.
"""
Memory-efficient attention for decoding.
It supports page size = 1.
"""
# Adapted from
@@ -197,7 +198,6 @@ def _decode_att_m_fwd(
logit_cap,
):
BLOCK = 32
# shape constraints
Lq, Lk = q.shape[-1], k_buffer.shape[-1]
batch, head_num = B_req_idx.shape[0], q.shape[1]
@@ -478,7 +478,6 @@ def _decode_grouped_att_m_fwd(
logit_cap,
):
BLOCK = 32
# shape constraints
Lq, Lk = q.shape[-1], k_buffer.shape[-1]
if Lk == 576:
@@ -570,9 +569,9 @@ def _decode_grouped_softmax_reducev_fwd(
BLOCK_DMODEL=BLOCK_DMODEL,
BLOCK_N=BLOCK,
BLOCK_H=BLOCK_H,
Lv=Lv,
num_warps=num_warps,
num_stages=1,
Lv=Lv,
)
@@ -588,7 +587,7 @@ def decode_attention_fwd(
max_len_in_batch,
total_num_tokens,
sm_scale,
logit_cap=-1,
logit_cap=0.0,
att_m=None,
):
if att_m is None:

31
python/sglang/srt/layers/triton_attention/extend_attention.py
View File

@@ -61,14 +61,14 @@ def _fwd_kernel(
stride_buf_vbs,
stride_buf_vh,
stride_req_to_tokens_b,
logit_cap: tl.constexpr,
Lq: tl.constexpr,
Lv: tl.constexpr,
BLOCK_DMODEL: tl.constexpr,
BLOCK_DPE: tl.constexpr,
BLOCK_DV: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
logit_cap: tl.constexpr,
Lq: tl.constexpr,
Lv: tl.constexpr,
):
cur_seq = tl.program_id(0)
cur_head = tl.program_id(1)
@@ -111,7 +111,7 @@ def _fwd_kernel(
)
qpe = tl.load(Q_Extend + offs_qpe, mask=mask_m[:, None], other=0.0)
# stage1: compute scores with prefix
# stage 1: compute scores with prefix
offs_n = tl.arange(0, BLOCK_N)
acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
@@ -174,7 +174,7 @@ def _fwd_kernel(
e_max = n_e_max
# stage2: compute the trianlge part
# stage 2: compute the trianlge part
cur_block_m_end = tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M)
for start_n in range(0, cur_block_m_end, BLOCK_N):
@@ -255,26 +255,22 @@ def extend_attention_fwd(
v_buffer,
req_to_tokens,
b_req_idx,
b_start_loc,
b_seq_len,
b_seq_len_prefix,
b_start_loc_extend,
b_seq_len_extend,
max_len_in_batch,
b_start_loc_extend,
max_len_extend,
sm_scale=None,
logit_cap=-1,
logit_cap=0.0,
):
"""
q_extend, k_extend, v_extend, o_extend: contiguous tensors
k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
"""
Lq, Lk, Lv, Lo = (
Lq, Lk, Lv = (
q_extend.shape[-1],
k_extend.shape[-1],
v_extend.shape[-1],
o_extend.shape[-1],
)
if Lq == 576:
@@ -303,7 +299,7 @@ def extend_attention_fwd(
else:
BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)
sm_scale = 1.0 / (Lq**0.5) if sm_scale is None else sm_scale
sm_scale = sm_scale or 1.0 / (Lq**0.5)
batch_size, head_num = b_seq_len.shape[0], q_extend.shape[1]
kv_group_num = q_extend.shape[1] // k_extend.shape[1]
@@ -338,27 +334,24 @@ def extend_attention_fwd(
v_buffer.stride(0),
v_buffer.stride(1),
req_to_tokens.stride(0),
logit_cap=logit_cap,
BLOCK_DMODEL=BLOCK_DMODEL,
BLOCK_DPE=BLOCK_DPE,
BLOCK_DV=BLOCK_DV,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
num_warps=num_warps,
num_stages=num_stages,
logit_cap=logit_cap,
Lq=Lq,
Lv=Lv,
num_warps=num_warps,
num_stages=num_stages,
)
def redundant_attention(
q_extend,
k_extend,
v_extend,
o_extend,
k_buffer,
v_buffer,
req_to_tokens,
b_req_idx,
b_start_loc,
b_seq_len,