minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Jacob0226
2026-04-10 01:08:32 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent f5fd5ab622
commit dd41764487
2 changed files with 68 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+30 -7
View File
@@ -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,
python/sglang/srt/layers/communicator.py
+38 -18
View File
@@ -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,