From 3ce3b4969f3142dbbf1eef4e5a17ec85f4c998b3 Mon Sep 17 00:00:00 2001 From: Jensen Date: Fri, 18 Sep 2026 14:10:43 +0800 Subject: [PATCH] [NPU] Avoid repeated BF16 wo_a weight transposes in DeepSeek-V4 decode (#39919) --- python/sglang/srt/environ.py | 3 + python/sglang/srt/models/deepseek_v4.py | 22 ++- test/manual/dsv4/bench_npu_wo_a.py | 217 ++++++++++++++++++++++++ 3 files changed, 238 insertions(+), 4 deletions(-) create mode 100644 test/manual/dsv4/bench_npu_wo_a.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e2190ce77..c71c22d50 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index aafa3b150..da3ca2f10 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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, diff --git a/test/manual/dsv4/bench_npu_wo_a.py b/test/manual/dsv4/bench_npu_wo_a.py new file mode 100644 index 000000000..94062adbf --- /dev/null +++ b/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()