@@ -115,14 +115,13 @@ def should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend: bool) -> bool:
|
||||
|
||||
|
||||
def is_dsa_enable_prefill_cp():
|
||||
if get_parallel().attn_cp_size <= 1:
|
||||
return False
|
||||
|
||||
if is_hip() or is_npu() or is_musa():
|
||||
return False
|
||||
|
||||
# Generic prefill CP derives activation from the runtime topology and model
|
||||
# architecture.
|
||||
if get_parallel().attn_cp_size <= 1:
|
||||
return False
|
||||
from sglang.srt.configs.model_config import is_deepseek_dsa, is_deepseek_v4
|
||||
|
||||
hf_config = process_model_config().hf_config
|
||||
|
||||
@@ -50,6 +50,11 @@ if TYPE_CHECKING:
|
||||
|
||||
from sgl_kernel import merge_state_v2
|
||||
|
||||
from sglang.kernels.ops.attention.flash_attention import (
|
||||
flash_attn_varlen_func,
|
||||
flash_attn_with_kvcache,
|
||||
)
|
||||
|
||||
|
||||
def _should_disable_scheduler_metadata_precompute() -> bool:
|
||||
return bool(get_parallel().enable_prefill_cp or get_parallel().enable_dp_attention)
|
||||
@@ -1278,13 +1283,7 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
aux_tensors=None,
|
||||
rel_bias=None,
|
||||
rel_bias_event=None,
|
||||
# Returns (output, lse) with lse in [total_q, num_heads].
|
||||
return_lse: bool = False,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
lse_out = None
|
||||
# Bound in __init__ so a subclass can substitute a different FA4 build.
|
||||
flash_attn_with_kvcache = self.flash_attn_with_kvcache
|
||||
flash_attn_varlen_func = self.flash_attn_varlen_func
|
||||
):
|
||||
if score_mod is not None and self.fa_impl_ver != 4:
|
||||
raise RuntimeError("score_mod is only supported by the FA4 backend.")
|
||||
cp_active = is_cp_active(forward_batch)
|
||||
@@ -1564,7 +1563,7 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
causal=False if use_cascade_attn else causal,
|
||||
window_size=window_size,
|
||||
softcap=layer.logit_cap,
|
||||
return_softmax_lse=use_cascade_attn or return_lse,
|
||||
return_softmax_lse=use_cascade_attn,
|
||||
num_splits=self.num_splits,
|
||||
out=_fa_out,
|
||||
ver=self.fa_impl_ver,
|
||||
@@ -1624,8 +1623,6 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
o_expand,
|
||||
softmax_lse_expand.T.contiguous(),
|
||||
)
|
||||
elif return_lse:
|
||||
o, lse_out, *_ = result
|
||||
else:
|
||||
o = result
|
||||
else:
|
||||
@@ -1826,12 +1823,7 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
else:
|
||||
o = result
|
||||
|
||||
o = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
||||
if return_lse:
|
||||
assert lse_out is not None
|
||||
# The varlen kernel emits LSE head-major [num_heads, total_q].
|
||||
return o, lse_out.transpose(0, 1).contiguous()
|
||||
return o
|
||||
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
||||
|
||||
def forward_decode(
|
||||
self,
|
||||
@@ -1852,13 +1844,7 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
aux_tensors=None,
|
||||
rel_bias=None,
|
||||
rel_bias_event=None,
|
||||
# Returns (output, lse) with lse in [total_q, num_heads].
|
||||
return_lse: bool = False,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
lse_out = None
|
||||
# Bound in __init__ so a subclass can substitute a different FA4 build.
|
||||
flash_attn_with_kvcache = self.flash_attn_with_kvcache
|
||||
flash_attn_varlen_func = self.flash_attn_varlen_func
|
||||
) -> torch.Tensor:
|
||||
if score_mod is not None and self.fa_impl_ver != 4:
|
||||
raise RuntimeError("score_mod is only supported by the FA4 backend.")
|
||||
if k is not None:
|
||||
@@ -2059,7 +2045,7 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
causal=False if use_cascade_attn else causal,
|
||||
window_size=window_size,
|
||||
softcap=layer.logit_cap,
|
||||
return_softmax_lse=use_cascade_attn or return_lse,
|
||||
return_softmax_lse=use_cascade_attn,
|
||||
num_splits=(
|
||||
self.decode_num_splits
|
||||
if not is_swa_layer
|
||||
@@ -2104,8 +2090,6 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
o_expand,
|
||||
softmax_lse_expand.T.contiguous(),
|
||||
)
|
||||
elif return_lse:
|
||||
o, lse_out, *_ = result
|
||||
else:
|
||||
o = result
|
||||
else:
|
||||
@@ -2184,12 +2168,7 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
else:
|
||||
o = result
|
||||
|
||||
o = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
||||
if return_lse:
|
||||
assert lse_out is not None
|
||||
# The varlen kernel emits LSE head-major [num_heads, total_q].
|
||||
return o, lse_out.transpose(0, 1).contiguous()
|
||||
return o
|
||||
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
||||
|
||||
def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
|
||||
"""Initialize CUDA graph state for the attention backend.
|
||||
|
||||
@@ -321,7 +321,6 @@ class QSAIndexer(MultiPlatformOp):
|
||||
self.compress_ratio, device=member_rows.device, dtype=torch.long
|
||||
)
|
||||
source_keys = token_k
|
||||
group_locs = group_locs.clamp_max(source_keys.shape[0] - 1)
|
||||
source_rope = metadata.extend_rope_matrix
|
||||
if source_rope is None:
|
||||
source_rope = build_rope_position_matrix(
|
||||
|
||||
@@ -221,17 +221,13 @@ def is_cpu() -> bool:
|
||||
return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" and is_host_cpu_supported
|
||||
|
||||
|
||||
try:
|
||||
import torchada # noqa: F401
|
||||
except ImportError:
|
||||
_IS_MUSA = False
|
||||
else:
|
||||
_IS_MUSA = hasattr(torch.version, "musa") and torch.version.musa is not None
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def is_musa() -> bool:
|
||||
return _IS_MUSA
|
||||
try:
|
||||
import torchada # noqa: F401
|
||||
except ImportError:
|
||||
return False
|
||||
return hasattr(torch.version, "musa") and torch.version.musa is not None
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
|
||||
Reference in New Issue
Block a user