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

Fuse wk and weight_proj in Indexer for DeepSeekV3.2-FP4 (#12094)

Browse Source
This commit is contained in:
Trevor Morris
2025-10-29 18:42:56 -07:00
committed by GitHub
parent 7ed8ba05cb
commit 9ff9fa7f95
2 changed files with 110 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+65
View File
@@ -224,6 +224,17 @@ 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()
@@ -1143,6 +1154,9 @@ 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(
@@ -3413,6 +3427,10 @@ 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 = is_nsa_indexer_wk_and_weights_proj_fused(
self.config, self.quant_config
)
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}"
@@ -3584,6 +3602,53 @@ 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
)
futures.append(
executor.submit(weight_loader, 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