[NPU] Avoid repeated BF16 wo_a weight transposes in DeepSeek-V4 decode (#39919)

python/sglang/srt/environ.py

@@ -950,6 +950,9 @@ class Envs:
     SGLANG_NPU_USE_MULTI_STREAM = EnvBool(False)
     SGLANG_NPU_USE_MLAPO = EnvBool(False)
     SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD = EnvBool(False)
+    # BF16 wo_a: use F.linear for single-local-group decode (Flash TP8),
+    # retaining the original weight layout. Opt-in for A/B.
+    SGLANG_OPT_NPU_BF16_WO_A_GEMM = EnvBool(False)
     # Forward native implementation for activation gelu tanh for model Skywork-Reward-Gemma-2-27B-v0.2
     SGLANG_NPU_FORWARD_NATIVE_GELUTANH = EnvBool(False)
     # Forward native implementation for gemma rms norm for model Skywork-Reward-Gemma-2-27B-v0.2

python/sglang/srt/models/deepseek_v4.py

@@ -233,6 +233,7 @@ def _get_mhc_ops() -> MhcOps:
 logger = logging.getLogger(__name__)
 
 _FP8_WO_A_GEMM = envs.SGLANG_OPT_FP8_WO_A_GEMM.get()
+_NPU_BF16_WO_A_GEMM = _is_npu and envs.SGLANG_OPT_NPU_BF16_WO_A_GEMM.get()
 _MHC_POST_MULT_VALUE = 2.0
 _HC_PRENORM_DEEPGEMM_MIN_TOKENS = 1024
 
@@ -2055,10 +2056,23 @@ class MQALayer(MqaAttentionBase):
             else:
                 wo_a_weight = getattr(self.wo_a, "weight", None)
                 if wo_a_weight is not None:
-                    wo_a = wo_a_weight.view(self.n_local_groups, self.o_lora_rank, -1)
-                    o = _apply_wo_a_bf16_matmul(
-                        o, wo_a, is_decode=forward_batch.forward_mode.is_decode()
-                    )
+                    if (
+                        _NPU_BF16_WO_A_GEMM
+                        and forward_batch.forward_mode.is_decode()
+                        and self.n_local_groups == 1
+                        and o.dtype == wo_a_weight.dtype == torch.bfloat16
+                        and wo_a_weight.is_contiguous()
+                    ):
+                        # One local group needs no grouped contraction; linear
+                        # avoids materializing a transpose of the BF16 weight.
+                        o = F.linear(o, wo_a_weight)
+                    else:
+                        wo_a = wo_a_weight.view(
+                            self.n_local_groups, self.o_lora_rank, -1
+                        )
+                        o = _apply_wo_a_bf16_matmul(
+                            o, wo_a, is_decode=forward_batch.forward_mode.is_decode()
+                        )
                 else:
                     o = _apply_gguf_grouped_wo_a(
                         o,