From eddf193292d310b416473815cffb0c9d13985be2 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sun, 22 Feb 2026 00:20:51 -0800 Subject: [PATCH] [DSv32] [GLM5] Improve Model Quality by Avoiding FP32 Precision Loss in `weights_proj` (#19041) --- .../srt/layers/attention/nsa/nsa_indexer.py | 24 +++++++++++++------ .../layers/deep_gemm_wrapper/compile_utils.py | 15 ++++++++++++ .../layers/deep_gemm_wrapper/entrypoint.py | 14 +++++++++++ test/registered/kernels/test_nsa_indexer.py | 4 ++-- 4 files changed, 48 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index ca54a931b..73c1e865e 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -229,21 +229,31 @@ class Indexer(MultiPlatformOp): else: yield - @torch.compile(dynamic=True) if not _is_hip else lambda f: f - def _project_and_scale_head_gates(self, x: torch.Tensor): + def _weights_proj_bf16_in_fp32_out(self, x: torch.Tensor) -> torch.Tensor: + if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + weight = self.weights_proj.weight + out = torch.empty( + (x.shape[0], weight.shape[0]), + dtype=torch.float32, + device=x.device, + ) + deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, weight, out) + return out + if _is_hip: x = x.to(self.weights_proj.weight.dtype) weights, _ = self.weights_proj(x) - weights = weights.float() + return weights.float() + + @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_bf16_in_fp32_out(x) 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): - if _is_hip: - x = x.to(self.weights_proj.weight.dtype) - weights, _ = self.weights_proj(x) - weights = weights.float() + weights = self._weights_proj_bf16_in_fp32_out(x) weights = weights * self.n_heads**-0.5 weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale return weights diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 267e6cebe..6aef866f5 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -96,6 +96,7 @@ class DeepGemmKernelType(IntEnum): GROUPED_GEMM_NT_F8F8BF16_MASKED = auto() GROUPED_GEMM_NT_F8F8BF16_CONTIG = auto() GEMM_NT_F8F8BF16 = auto() + GEMM_NT_BF16BF16F32 = auto() _INITIALIZATION_DICT: Dict[Tuple[DeepGemmKernelType, int, int, int], bool] = dict() @@ -216,6 +217,7 @@ class _BaseWarmupExecutor: DeepGemmKernelType.GEMM_NT_F8F8BF16: _NormalWarmupExecutor, DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: _GroupedContWarmupExecutor, DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED: _GroupedMaskedWarmupExecutor, + DeepGemmKernelType.GEMM_NT_BF16BF16F32: _BF16F32WarmupExecutor, }[kernel_type](**kwargs) @staticmethod @@ -235,6 +237,9 @@ class _BaseWarmupExecutor: + num_groups * 4 + num_groups * max_m * n * 2 ) / _GB + elif kernel_type == DeepGemmKernelType.GEMM_NT_BF16BF16F32: + # bf16 lhs + bf16 rhs + fp32 out + return (max_m * k * 2 + n * k * 2 + max_m * n * 4) / _GB else: raise ValueError(f"Invalid kernel type: {kernel_type}") @@ -317,6 +322,16 @@ class _GroupedMaskedWarmupExecutor(_BaseWarmupExecutor): ) +class _BF16F32WarmupExecutor(_BaseWarmupExecutor): + def __init__(self, max_m: int, n: int, k: int, num_groups: int): + self.lhs = torch.empty((max_m, k), device="cuda", dtype=torch.bfloat16) + self.rhs = torch.empty((n, k), device="cuda", dtype=torch.bfloat16) + self.out = torch.empty((max_m, n), device="cuda", dtype=torch.float32) + + def execute(self, m): + deep_gemm.bf16_gemm_nt(self.lhs[:m], self.rhs, self.out[:m]) + + @contextmanager def deep_gemm_execution_hook( m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 88d0a959b..8eada5be6 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -102,6 +102,20 @@ def gemm_nt_f8f8bf16( ) +def gemm_nt_bf16bf16f32( + lhs: torch.Tensor, + rhs: torch.Tensor, + out: torch.Tensor, +): + m, k = lhs.shape + n, _ = rhs.shape + num_groups = 1 + kernel_type = compile_utils.DeepGemmKernelType.GEMM_NT_BF16BF16F32 + + with compile_utils.deep_gemm_execution_hook(m, n, k, num_groups, kernel_type): + deep_gemm.bf16_gemm_nt(lhs, rhs, out) + + def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): compile_utils.update_deep_gemm_config(gpu_id, server_args) diff --git a/test/registered/kernels/test_nsa_indexer.py b/test/registered/kernels/test_nsa_indexer.py index 834b1039c..42f5d316b 100644 --- a/test/registered/kernels/test_nsa_indexer.py +++ b/test/registered/kernels/test_nsa_indexer.py @@ -24,7 +24,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=2, suite="stage-b-test-small-1-gpu") +register_cuda_ci(est_time=2, suite="stage-b-test-large-1-gpu") # Global configuration for all indexer tests DEFAULT_CONFIG = { @@ -34,7 +34,7 @@ DEFAULT_CONFIG = { "context_len": 2048, "max_bs": 64, "hidden_size": 5120, - "index_n_heads": 1, + "index_n_heads": 32, "index_head_dim": 128, "rope_head_dim": 64, "index_topk": 64,