[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)
+                    if (
-                    o = _apply_wo_a_bf16_matmul(
+                        _NPU_BF16_WO_A_GEMM
-                        o, wo_a, is_decode=forward_batch.forward_mode.is_decode()
+                        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,

test/manual/dsv4/bench_npu_wo_a.py

@@ -0,0 +1,217 @@
+"""Compare BF16 wo_a einsum and F.linear on one Ascend NPU.
+
+python test/manual/dsv4/bench_npu_wo_a.py --batch-sizes 1 8
+
+Simulates one DSV4 Flash TP8 rank with one local output group. Both variants
+share the same inputs and 43 distinct layer weights by default. One iteration
+replays all --layers projections; ms_per_iteration reports their total time,
+and us_per_layer reports the average. Model loading and serving are not timed.
+F.linear is called directly, as in the model, so no optimization flags are needed.
+Optional --profile-dir exports a trace for checking weight transposes.
+"""
+
+import argparse
+import json
+import math
+import statistics
+
+import torch
+import torch.nn.functional as F
+
+
+def _einsum(o, weight):
+    return torch.einsum("tgd,grd->tgr", o, weight.unsqueeze(0))
+
+
+def _capture(fn, inputs, weights):
+    def run():
+        return [fn(o, w) for o, w in zip(inputs, weights)]
+
+    stream = torch.npu.Stream()
+    stream.wait_stream(torch.npu.current_stream())
+    with torch.npu.stream(stream):
+        for _ in range(3):
+            run()
+    torch.npu.current_stream().wait_stream(stream)
+    torch.npu.synchronize()
+    graph = torch.npu.NPUGraph()
+    with torch.npu.graph(graph, stream=stream, auto_dispatch_capture=True):
+        outputs = run()
+    return graph, outputs
+
+
+def _time_graph(graph, iterations):
+    start = torch.npu.Event(enable_timing=True)
+    end = torch.npu.Event(enable_timing=True)
+    start.record()
+    for _ in range(iterations):
+        graph.replay()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) / iterations
+
+
+@torch.inference_mode()
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 8])
+    parser.add_argument(
+        "--layers",
+        type=int,
+        default=43,
+        help="Distinct wo_a projections per graph replay.",
+    )
+    parser.add_argument(
+        "--dim", type=int, default=4096, help="Input feature width of the local group."
+    )
+    parser.add_argument(
+        "--o-lora-rank",
+        type=int,
+        default=1024,
+        help="Output projection width, not a distributed rank ID.",
+    )
+    parser.add_argument(
+        "--input-stride-factor",
+        type=int,
+        default=8,
+        help="Token stride / D; 8 models the 64-head output sliced to 8 TP-local heads.",
+    )
+    parser.add_argument(
+        "--iterations", type=int, default=100, help="Graph replays per timing round."
+    )
+    parser.add_argument(
+        "--rounds", type=int, default=5, help="Timing rounds; results use the median."
+    )
+    parser.add_argument("--device", default="npu:0")
+    parser.add_argument("--profile-dir")
+    args = parser.parse_args()
+    if (
+        min(
+            *args.batch_sizes,
+            args.layers,
+            args.dim,
+            args.o_lora_rank,
+            args.input_stride_factor,
+            args.iterations,
+            args.rounds,
+        )
+        < 1
+    ):
+        parser.error("dimensions and iteration counts must be positive")
+
+    import torch_npu
+
+    torch.npu.set_device(args.device)
+    torch.manual_seed(42)
+    print(
+        json.dumps(
+            {
+                "torch": torch.__version__,
+                "torch_npu": torch_npu.__version__,
+                "device": torch.npu.get_device_name(),
+                "args": vars(args),
+            }
+        ),
+        flush=True,
+    )
+    weights = [
+        torch.randn(
+            args.o_lora_rank, args.dim, device=args.device, dtype=torch.bfloat16
+        )
+        / math.sqrt(args.dim)
+        for _ in range(args.layers)
+    ]
+    variants = {
+        "einsum_original": _einsum,
+        "linear_original": F.linear,
+    }
+    for tokens in args.batch_sizes:
+        inputs = [
+            torch.randn(
+                tokens,
+                args.input_stride_factor,
+                args.dim,
+                device=args.device,
+                dtype=torch.bfloat16,
+            )[:, :1, :]
+            for _ in weights
+        ]
+        graphs = {}
+        outputs = {}
+        for name, fn in variants.items():
+            graphs[name], outputs[name] = _capture(fn, inputs, weights)
+            graphs[name].replay()
+        torch.npu.synchronize()
+        for actual, reference in zip(
+            outputs["linear_original"], outputs["einsum_original"]
+        ):
+            torch.testing.assert_close(actual, reference, rtol=0.016, atol=0.016)
+
+        # A captured graph must see an in-place online weight update even when
+        # the update bypasses model.post_load_weights(), as direct updates do.
+        pointer = weights[0].data_ptr()
+        saved = weights[0].clone()
+        before_update = outputs["einsum_original"][0].clone()
+        weights[0].neg_()
+        for graph in graphs.values():
+            graph.replay()
+        torch.npu.synchronize()
+        assert not torch.equal(outputs["einsum_original"][0], before_update)
+        # No bias: negating W must negate the result. This also catches a stale
+        # weight in the baseline graph rather than only comparing two graphs.
+        expected_after_update = -before_update
+        for tensors in outputs.values():
+            torch.testing.assert_close(
+                tensors[0], expected_after_update, rtol=0.016, atol=0.016
+            )
+        assert weights[0].data_ptr() == pointer
+        weights[0].copy_(saved)
+        torch.npu.synchronize()
+
+        samples = {name: [] for name in graphs}
+        names = list(graphs)
+        for round_id in range(args.rounds):
+            # Rotate measurement order to reduce a fixed-order clock bias.
+            offset = round_id % len(names)
+            for name in names[offset:] + names[:offset]:
+                samples[name].append(_time_graph(graphs[name], args.iterations))
+        baseline = statistics.median(samples["einsum_original"])
+        for name, times in samples.items():
+            elapsed = statistics.median(times)
+            print(
+                json.dumps(
+                    {
+                        "batch_size": tokens,
+                        "variant": name,
+                        "layers": args.layers,
+                        "input_stride": list(inputs[0].stride()),
+                        "ms_per_iteration": elapsed,
+                        "us_per_layer": elapsed * 1000 / args.layers,
+                        "speedup": baseline / elapsed,
+                        "samples_ms": times,
+                    }
+                ),
+                flush=True,
+            )
+
+        if args.profile_dir:
+            with torch_npu.profiler.profile(
+                activities=[
+                    torch_npu.profiler.ProfilerActivity.CPU,
+                    torch_npu.profiler.ProfilerActivity.NPU,
+                ],
+                record_shapes=True,
+                with_stack=False,
+                on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
+                    f"{args.profile_dir}/bs{tokens}"
+                ),
+            ) as profiler:
+                for name, graph in graphs.items():
+                    with torch.autograd.profiler.record_function(name):
+                        graph.replay()
+                        torch.npu.synchronize()
+                    profiler.step()
+
+
+if __name__ == "__main__":
+    main()