[AMD][HIP] NSA: bf16 passthrough from RMSNorm to eliminate FP8 dequantization (#22258)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

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
+            # aiter (ROCm gfx95): the 3-tuple (fp8, scale, bf16) from
-            # to a float tensor here (callsite), keeping `_get_logits_head_gate` backend-agnostic.
+            # fused_rms_fp8_group_quant is passed directly to _get_logits_head_gate,
-            if isinstance(x, tuple):
+            # 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,

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):
-
+                        # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant.
-                        hidden_states, _, _, _res = fused_rms_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,7 +500,13 @@ 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:
@@ -512,11 +524,12 @@ class LayerCommunicator:
                             residual,
                         )
                     elif _use_aiter and _is_gfx95_supported and ("fp8" in quant_format):
-                        # RMSNorm + FP8 per-group quant
+                        # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant
-                        # return hidden_states：
+                        # with residual addition. When NSA is active, pack
-                        #   out_fp8  : FP8 activation →  a8w8 GEMM
+                        # the unquantized bf16 as a 3-tuple (fp8, scale, bf16).
-                        #   out_bs   : block-scale →  gemm_a8w8_blockscale.x_scale
+                        _nsa_needs_bf16 = get_attn_tp_context().is_nsa
-                        hidden_states, _, _, residual = fused_rms_fp8_group_quant(
+                        hidden_states, _unq_bf16, _, residual = (
+                            fused_rms_fp8_group_quant(
                                 hidden_states,
                                 self.input_layernorm.weight,
                                 self.input_layernorm.variance_epsilon,
@@ -526,7 +539,14 @@ class LayerCommunicator:
                                 group_size=128,
                                 dtype_quant=torch.float8_e4m3fn,
                                 res1=residual,
-                            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, residual = self.input_layernorm(