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

[Deepseek V3.2] Change indexer weights_proj to fp32 (#13459)

Browse Source
This commit is contained in:
hlu1
2025-11-20 12:24:10 -08:00
committed by GitHub
parent 6bc3062894
commit 7291c72e57
3 changed files with 92 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -109,7 +109,6 @@ class Indexer(CustomOp):
prefix: str = "",
quant_config: Optional[QuantizationConfig] = None,
alt_stream: Optional[torch.cuda.Stream] = None,
fuse_wk_and_weights_proj: bool = False,
):
super().__init__()
self.hidden_size = hidden_size
@@ -120,7 +119,6 @@ class Indexer(CustomOp):
self.q_lora_rank = q_lora_rank
self.layer_id = layer_id
self.alt_stream = alt_stream
self.fuse_wk_and_weights_proj = fuse_wk_and_weights_proj
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
self.cp_size = get_attention_tp_size()
@@ -139,28 +137,22 @@ class Indexer(CustomOp):
quant_config=quant_config,
prefix=add_prefix("wq_b", prefix),
)
if self.fuse_wk_and_weights_proj:
self.fused_wk_and_weights_proj = ReplicatedLinear(
self.hidden_size,
self.head_dim + self.n_heads,
bias=False,
prefix=add_prefix("fused_wk_and_weights_proj", prefix),
)
else:
self.wk = ReplicatedLinear(
self.hidden_size,
self.head_dim,
bias=False,
quant_config=quant_config,
prefix=add_prefix("wk", prefix),
)
# NOTE: weight_proj is not quantized
self.weights_proj = ReplicatedLinear(
self.hidden_size,
self.n_heads,
bias=False,
prefix=add_prefix("weights_proj", prefix),
)
self.wk = ReplicatedLinear(
self.hidden_size,
self.head_dim,
bias=False,
quant_config=quant_config,
prefix=add_prefix("wk", prefix),
)
# NOTE: weights_proj in the checkpoint is stored in bf16, while the parameters here are stored in fp32 for convenience
self.weights_proj = ReplicatedLinear(
self.hidden_size,
self.n_heads,
bias=False,
params_dtype=torch.float32,
prefix=add_prefix("weights_proj", prefix),
)
self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)
self.rotary_emb = get_rope_wrapper(
rope_head_dim,
@@ -176,7 +168,8 @@ class Indexer(CustomOp):
self.softmax_scale = self.head_dim**-0.5
@torch.compile(dynamic=True)
def _get_logits_head_gate(self, weights: torch.Tensor, q_scale: torch.Tensor):
def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
weights, _ = self.weights_proj(x.float())
weights = weights * self.n_heads**-0.5
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
return weights
@@ -189,7 +182,6 @@ class Indexer(CustomOp):
enable_dual_stream: bool,
forward_batch: ForwardBatch,
):
weights = None
if enable_dual_stream:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
@@ -206,12 +198,7 @@ class Indexer(CustomOp):
)
with torch.cuda.stream(self.alt_stream):
# TODO we should also put DeepGEMM half SM here?
if self.fuse_wk_and_weights_proj:
key, weights = self.fused_wk_and_weights_proj(x)[0].split(
[self.head_dim, self.n_heads], dim=-1
)
else:
key, _ = self.wk(x)
key, _ = self.wk(x)
key = self.k_norm(key)
k_rope, _ = torch.split(
@@ -224,17 +211,10 @@ class Indexer(CustomOp):
else:
query, _ = self.wq_b(q_lora)
query = rearrange(query, "l (h d) -> l h d", d=self.head_dim)
q_rope, _ = torch.split(
query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1
)
if self.fuse_wk_and_weights_proj:
key, weights = self.fused_wk_and_weights_proj(x)[0].split(
[self.head_dim, self.n_heads], dim=-1
)
else:
key, _ = self.wk(x)
key, _ = self.wk(x)
key = self.k_norm(key)
k_rope, _ = torch.split(
key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1
@@ -266,7 +246,7 @@ class Indexer(CustomOp):
query = rotate_activation(query)
key = rotate_activation(key)
return query, key, weights
return query, key
def _get_k_bf16(
self,
@@ -274,13 +254,8 @@ class Indexer(CustomOp):
positions: torch.Tensor,
enable_dual_stream: bool,
):
# Compute only key, skip query and weights (weights is discarded if fused)
if self.fuse_wk_and_weights_proj:
key, _ = self.fused_wk_and_weights_proj(x)[0].split(
[self.head_dim, self.n_heads], dim=-1
)
else:
key, _ = self.wk(x)
# Compute only key, skip query
key, _ = self.wk(x)
key = self.k_norm(key)
k_rope, _ = torch.split(
key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1
@@ -779,7 +754,7 @@ class Indexer(CustomOp):
return_indices,
)
query, key, weights = self._get_q_k_bf16(
query, key = self._get_q_k_bf16(
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
)
@@ -808,9 +783,7 @@ class Indexer(CustomOp):
index_k_scale=k_scale,
)
if not self.fuse_wk_and_weights_proj:
weights, _ = self.weights_proj(x)
weights = self._get_logits_head_gate(weights, q_scale)
weights = self._get_logits_head_gate(x, q_scale)
if is_cuda():
assert forward_batch.seq_lens_cpu is not None
@@ -1037,7 +1010,7 @@ class Indexer(CustomOp):
past_key_states = forward_batch.token_to_kv_pool.get_index_k_buffer(layer_id)
x = x.view(-1, self.hidden_size)
weights = self.weights_proj(x)[0]
weights = self.weights_proj(x.float())[0]
block_table = (
block_table[: actual_seq_lengths_q.size()[0]] if is_prefill else block_table
)

71
python/sglang/srt/models/deepseek_v2.py
View File

@@ -243,17 +243,6 @@ def add_forward_absorb_core_attention_backend(backend_name):
logger.info(f"Added {backend_name} to FORWARD_ABSORB_CORE_ATTENTION_BACKENDS.")
def is_nsa_indexer_wk_and_weights_proj_fused(config, quant_config):
"""
NSA Indexer wk and weights_proj can be fused in FP4 model because they are both in BF16
"""
return (
is_deepseek_nsa(config)
and quant_config is not None
and quant_config.get_name() == "modelopt_fp4"
)
class AttnForwardMethod(IntEnum):
# Use multi-head attention
MHA = auto()
@@ -1235,9 +1224,6 @@ class DeepseekV2AttentionMLA(nn.Module):
quant_config=quant_config,
layer_id=layer_id,
alt_stream=alt_stream,
fuse_wk_and_weights_proj=is_nsa_indexer_wk_and_weights_proj_fused(
config, quant_config
),
)
self.kv_b_proj = ColumnParallelLinear(
@@ -3849,12 +3835,6 @@ class DeepseekV2ForCausalLM(nn.Module):
self.config.q_lora_rank is not None
)
cached_a_proj = {} if fuse_qkv_a_proj else None
# Fuse wk and weights_proj when NSA Indexer is enabled and quant_config is FP4. For nextn, fp4 is disabled so we cannot fuse.
fuse_wk_and_weights_proj = (
is_nsa_indexer_wk_and_weights_proj_fused(self.config, self.quant_config)
and not is_nextn
)
cached_wk_and_weights_proj = {} if fuse_wk_and_weights_proj else None
if is_nextn:
nextn_layer_prefix = f"model.layers.{nextn_layer_id}"
@@ -4040,57 +4020,6 @@ class DeepseekV2ForCausalLM(nn.Module):
)
cached_a_proj.pop(q_a_proj_name)
cached_a_proj.pop(kv_a_proj_name)
elif fuse_wk_and_weights_proj and (
"wk" in name or "weights_proj" in name
):
cached_wk_and_weights_proj[name] = loaded_weight
wk_name = (
name
if "wk" in name
else name.replace("weights_proj", "wk")
)
weights_proj_name = (
name
if "weights_proj" in name
else name.replace("wk", "weights_proj")
)
# When both wk and weights_proj has been cached, load the fused weight to parameter
if (
wk_name in cached_wk_and_weights_proj
and weights_proj_name in cached_wk_and_weights_proj
):
wk_weight = cached_wk_and_weights_proj[wk_name]
weights_proj_weight = cached_wk_and_weights_proj[
weights_proj_name
]
# todo dequantize wk for fp8
assert wk_weight.dtype == weights_proj_weight.dtype
fused_weight = torch.cat(
[wk_weight, weights_proj_weight], dim=0
)
param_name = (
name.replace("wk", "fused_wk_and_weights_proj")
if "wk" in name
else name.replace(
"weights_proj",
"fused_wk_and_weights_proj",
)
)
param = params_dict[param_name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, fused_weight),
)
cached_wk_and_weights_proj.pop(wk_name)
cached_wk_and_weights_proj.pop(weights_proj_name)
else:
if (
"k_scale" in name or "v_scale" in name