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

Support multi-node DP attention (#2925)

Browse Source 
Co-authored-by: dhou-xai <dhou@x.ai>
This commit is contained in:
Lianmin Zheng
2025-01-16 11:15:00 -08:00
committed by GitHub
co-authored by dhou-xai
parent 58f3f2b840
commit 8b6ce52e92
16 changed files with 287 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashinfer_backend.py
+10 -9
View File
@@ -18,6 +18,7 @@ import triton.language as tl
from sglang.global_config import global_config
from sglang.srt.layers.attention import AttentionBackend
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils import is_flashinfer_available
@@ -62,9 +63,9 @@ class FlashInferAttnBackend(AttentionBackend):
self.decode_use_tensor_cores = should_use_tensor_core(
kv_cache_dtype=model_runner.kv_cache_dtype,
num_attention_heads=model_runner.model_config.num_attention_heads
// model_runner.tp_size,
// get_attention_tp_size(),
num_kv_heads=model_runner.model_config.get_num_kv_heads(
model_runner.tp_size
get_attention_tp_size()
),
)
self.max_context_len = model_runner.model_config.context_len
@@ -147,7 +148,7 @@ class FlashInferAttnBackend(AttentionBackend):
self.prefill_cuda_graph_metadata = {}
def init_forward_metadata(self, forward_batch: ForwardBatch):
if forward_batch.forward_mode.is_decode():
if forward_batch.forward_mode.is_decode_or_idle():
self.indices_updater_decode.update(
forward_batch.req_pool_indices,
forward_batch.seq_lens,
@@ -238,7 +239,7 @@ class FlashInferAttnBackend(AttentionBackend):
forward_mode: ForwardMode,
spec_info: Optional[SpecInfo],
):
if forward_mode.is_decode():
if forward_mode.is_decode_or_idle():
decode_wrappers = []
for i in range(self.num_wrappers):
decode_wrappers.append(
@@ -307,7 +308,7 @@ class FlashInferAttnBackend(AttentionBackend):
forward_mode: ForwardMode,
spec_info: Optional[SpecInfo],
):
if forward_mode.is_decode():
if forward_mode.is_decode_or_idle():
self.indices_updater_decode.update(
req_pool_indices[:bs],
seq_lens[:bs],
@@ -453,10 +454,10 @@ class FlashInferIndicesUpdaterDecode:
def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
# Parse Constants
self.num_qo_heads = (
model_runner.model_config.num_attention_heads // model_runner.tp_size
model_runner.model_config.num_attention_heads // get_attention_tp_size()
)
self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
model_runner.tp_size
get_attention_tp_size()
)
self.head_dim = model_runner.model_config.head_dim
self.data_type = model_runner.kv_cache_dtype
@@ -625,10 +626,10 @@ class FlashInferIndicesUpdaterPrefill:
def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
# Parse Constants
self.num_qo_heads = (
model_runner.model_config.num_attention_heads // model_runner.tp_size
model_runner.model_config.num_attention_heads // get_attention_tp_size()
)
self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
model_runner.tp_size
get_attention_tp_size()
)
self.head_dim = model_runner.model_config.head_dim
self.data_type = model_runner.kv_cache_dtype
python/sglang/srt/layers/attention/triton_backend.py
+4 -6
View File
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.attention import AttentionBackend
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
if TYPE_CHECKING:
@@ -28,12 +29,9 @@ class TritonAttnBackend(AttentionBackend):
self.decode_attention_fwd = decode_attention_fwd
self.extend_attention_fwd = extend_attention_fwd
if model_runner.server_args.enable_dp_attention:
self.num_head = model_runner.model_config.num_attention_heads
else:
self.num_head = (
model_runner.model_config.num_attention_heads // model_runner.tp_size
)
self.num_head = (
model_runner.model_config.num_attention_heads // get_attention_tp_size()
)
self.num_kv_splits = model_runner.server_args.triton_attention_num_kv_splits
self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1]