[Feat] QWen-1M context support[2/2]: Update block sparse attention backend (#5949)
This commit is contained in:
@@ -210,6 +210,7 @@ class Qwen2MoeAttention(nn.Module):
|
||||
max_position_embeddings: int = 8192,
|
||||
qkv_bias: int = True,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
dual_chunk_attention_config: Optional[dict[str, Any]] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
@@ -267,6 +268,7 @@ class Qwen2MoeAttention(nn.Module):
|
||||
max_position=max_position_embeddings,
|
||||
base=rope_theta,
|
||||
rope_scaling=rope_scaling,
|
||||
dual_chunk_attention_config=dual_chunk_attention_config,
|
||||
)
|
||||
self.attn = RadixAttention(
|
||||
self.num_heads,
|
||||
@@ -308,6 +310,9 @@ class Qwen2MoeDecoderLayer(nn.Module):
|
||||
rope_scaling = getattr(config, "rope_scaling", None)
|
||||
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
|
||||
qkv_bias = getattr(config, "qkv_bias", True)
|
||||
dual_chunk_attention_config = getattr(
|
||||
config, "dual_chunk_attention_config", None
|
||||
)
|
||||
self.self_attn = Qwen2MoeAttention(
|
||||
hidden_size=self.hidden_size,
|
||||
num_heads=config.num_attention_heads,
|
||||
@@ -317,6 +322,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
|
||||
rope_scaling=rope_scaling,
|
||||
max_position_embeddings=max_position_embeddings,
|
||||
quant_config=quant_config,
|
||||
dual_chunk_attention_config=dual_chunk_attention_config,
|
||||
qkv_bias=qkv_bias,
|
||||
prefix=add_prefix("self_attn", prefix),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user