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

[OPT] DeepSeekV3.2: optimize indexer weight_proj-mma performance (#17205)

Browse Source
This commit is contained in:
BingjiaWang
2026-01-20 23:09:44 +08:00
committed by GitHub
parent 55c616427d
commit 612026ad2c
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -182,12 +182,11 @@ class Indexer(MultiPlatformOp):
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,
params_dtype=torch.bfloat16,
prefix=add_prefix("weights_proj", prefix),
)
self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)
@@ -221,13 +220,15 @@ class Indexer(MultiPlatformOp):
@torch.compile(dynamic=True) if not _is_hip else lambda f: f
def _project_and_scale_head_gates(self, x: torch.Tensor):
weights, _ = self.weights_proj(x.float())
weights, _ = self.weights_proj(x)
weights = weights.float()
weights = weights * self.n_heads**-0.5
return weights
@torch.compile(dynamic=True) if not _is_hip else lambda f: f
def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
weights, _ = self.weights_proj(x.float())
weights, _ = self.weights_proj(x)
weights = weights.float()
weights = weights * self.n_heads**-0.5
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
return weights