From dd41764487447740cfe7e8141f62c627fd54ce09 Mon Sep 17 00:00:00 2001 From: Jacob0226 Date: Fri, 10 Apr 2026 16:08:32 +0800 Subject: [PATCH] [AMD][HIP] NSA: bf16 passthrough from RMSNorm to eliminate FP8 dequantization (#22258) Co-authored-by: Claude Opus 4.6 --- .../srt/layers/attention/nsa/nsa_indexer.py | 37 +++++++++--- python/sglang/srt/layers/communicator.py | 56 +++++++++++++------ 2 files changed, 68 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 4ffd13bdd..8469a8e43 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -2,7 +2,7 @@ from __future__ import annotations import contextlib from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import torch from einops import rearrange @@ -21,6 +21,7 @@ from sglang.srt.utils import ( ceil_align, get_bool_env_var, is_cuda, + is_gfx95_supported, is_hip, is_npu, ) @@ -31,6 +32,8 @@ _is_hip = is_hip() _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_fp8_fnuz = is_fp8_fnuz() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_gfx95_supported = is_gfx95_supported() if _is_cuda: try: import deep_gemm @@ -251,7 +254,14 @@ class Indexer(MultiPlatformOp): else: yield - def _weights_proj_bf16_in_fp32_out(self, x: torch.Tensor) -> torch.Tensor: + def _weights_proj_bf16_in_fp32_out( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] + ) -> torch.Tensor: + # aiter (ROCm gfx95): extract the passthrough bf16 tensor from the + # 3-tuple (fp8, scale, bf16) produced by fused_rms_fp8_group_quant, + # avoiding an expensive FP8-to-bf16 dequantization. + if _use_aiter and _is_gfx95_supported and isinstance(x, tuple) and len(x) == 3: + x = x[2] if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: weight = self.weights_proj.weight out = torch.empty( @@ -268,13 +278,17 @@ class Indexer(MultiPlatformOp): return weights.float() @torch.compile(dynamic=True) - def _project_and_scale_head_gates(self, x: torch.Tensor): + def _project_and_scale_head_gates( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] + ): weights = self._weights_proj_bf16_in_fp32_out(x) weights = weights * self.n_heads**-0.5 return weights @torch.compile(dynamic=True) - def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor): + def _get_logits_head_gate( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]], q_scale: torch.Tensor + ): 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 @@ -1133,9 +1147,18 @@ class Indexer(MultiPlatformOp): act_quant=act_quant, ) - # `_get_logits_head_gate` expects a Tensor. For tuple activations, dequantize - # to a float tensor here (callsite), keeping `_get_logits_head_gate` backend-agnostic. - if isinstance(x, tuple): + # aiter (ROCm gfx95): the 3-tuple (fp8, scale, bf16) from + # fused_rms_fp8_group_quant is passed directly to _get_logits_head_gate, + # which extracts the bf16 tensor via _weights_proj_bf16_in_fp32_out, + # completely skipping the FP8 dequantization path below. + if ( + _use_aiter + and _is_gfx95_supported + and isinstance(x, tuple) + and len(x) == 3 + ): + x_for_gate = x + elif isinstance(x, tuple): assert len(x) in ( 2, 3, diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 936eecb90..d20afc746 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -195,8 +195,10 @@ class AttnTpContext: self.allow_input_scattered = False self.input_scattered_ = False self.attn_inputs_: Optional[AttentionInputs] = None + self.is_nsa = False def init_context(self, q_lora_rank, is_nsa): + self.is_nsa = is_nsa self.allow_input_scattered = ( get_global_server_args().enable_attn_tp_input_scattered and (_is_cuda or _is_npu) @@ -483,8 +485,12 @@ class LayerCommunicator: None, ) elif _use_aiter and _is_gfx95_supported and ("fp8" in quant_format): - - hidden_states, _, _, _res = fused_rms_fp8_group_quant( + # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant. + # When NSA is active, also preserve the unquantized bf16 + # output as a 3-tuple (fp8, scale, bf16) so the NSA + # indexer can skip redundant FP8 dequantization. + _nsa_needs_bf16 = get_attn_tp_context().is_nsa + hidden_states, _unq_bf16, _, _res = fused_rms_fp8_group_quant( hidden_states, self.input_layernorm.weight, self.input_layernorm.variance_epsilon, @@ -494,8 +500,14 @@ class LayerCommunicator: group_size=128, dtype_quant=torch.float8_e4m3fn, res1=None, - output_unquantized_inp1=False, + output_unquantized_inp1=_nsa_needs_bf16, ) + if _nsa_needs_bf16: + hidden_states = ( + hidden_states[0], + hidden_states[1], + _unq_bf16, + ) else: hidden_states = self.input_layernorm(hidden_states) @@ -512,22 +524,30 @@ class LayerCommunicator: residual, ) elif _use_aiter and _is_gfx95_supported and ("fp8" in quant_format): - # RMSNorm + FP8 per-group quant - # return hidden_states： - # out_fp8 : FP8 activation → a8w8 GEMM - # out_bs : block-scale → gemm_a8w8_blockscale.x_scale - hidden_states, _, _, residual = fused_rms_fp8_group_quant( - hidden_states, - self.input_layernorm.weight, - self.input_layernorm.variance_epsilon, - inp2=None, - inp2_weight=None, - inp2_epsilon=None, - group_size=128, - dtype_quant=torch.float8_e4m3fn, - res1=residual, - output_unquantized_inp1=False, + # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant + # with residual addition. When NSA is active, pack + # the unquantized bf16 as a 3-tuple (fp8, scale, bf16). + _nsa_needs_bf16 = get_attn_tp_context().is_nsa + hidden_states, _unq_bf16, _, residual = ( + fused_rms_fp8_group_quant( + hidden_states, + self.input_layernorm.weight, + self.input_layernorm.variance_epsilon, + inp2=None, + inp2_weight=None, + inp2_epsilon=None, + group_size=128, + dtype_quant=torch.float8_e4m3fn, + res1=residual, + output_unquantized_inp1=_nsa_needs_bf16, + ) ) + if _nsa_needs_bf16: + hidden_states = ( + hidden_states[0], + hidden_states[1], + _unq_bf16, + ) else: hidden_states, residual = self.input_layernorm( hidden_states,