Skip to content

Commit 7e6ad73

Browse files
k50112113k50112113
committed
add fused rms quant
1 parent 090f38f commit 7e6ad73

8 files changed

Lines changed: 91 additions & 56 deletions

File tree

‎vllm/envs.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,8 @@
115115
VLLM_AITER_TRITON_FP8_BMM: bool = False
116116
VLLM_AITER_TRITON_FUSED_CONCAT_ZEROS: bool = False
117117
VLLM_AITER_TRITON_FUSED_ROPE_CACHE_CONCAT: bool = False
118-
VLLM_MLA_FP8_PADDING: bool = False
118+
VLLM_MLA_FP8_PADDING: bool = False,
119+
VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT: bool = False
119120

120121

121122
def get_default_cache_root():
@@ -755,6 +756,9 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:
755756
"VLLM_MLA_FP8_PADDING":
756757
lambda: bool(int(os.getenv("VLLM_MLA_FP8_PADDING", "0"))),
757758

759+
"VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT":
760+
lambda: bool(int(os.getenv("VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT", "0"))),
761+
758762
}
759763

760764
# end-env-vars-definition

‎vllm/model_executor/layers/linear.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -391,8 +391,8 @@ def forward(
391391
) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
392392
bias = self.bias if not self.skip_bias_add else None
393393
assert self.quant_method is not None
394-
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
395-
if isinstance(self.quant_method, Fp8LinearMethod):
394+
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import CompressedTensorsLinearMethod
395+
if isinstance(self.quant_method, CompressedTensorsLinearMethod):
396396
output = self.quant_method.apply(self, x, bias, x_quant_scales=x_quant_scales)
397397
else:
398398
assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"
@@ -623,8 +623,8 @@ def forward(
623623

624624
# Matrix multiply.
625625
assert self.quant_method is not None
626-
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
627-
if isinstance(self.quant_method, Fp8LinearMethod):
626+
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import CompressedTensorsLinearMethod
627+
if isinstance(self.quant_method, CompressedTensorsLinearMethod):
628628
output_parallel = self.quant_method.apply(self, input_, bias, x_quant_scales=x_quant_scales)
629629
else:
630630
assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"
@@ -1413,8 +1413,8 @@ def forward(
14131413
# Only fuse bias add into GEMM for rank 0 (this ensures that
14141414
# bias will not get added more than once in TP>1 case)
14151415
bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
1416-
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
1417-
if isinstance(self.quant_method, Fp8LinearMethod):
1416+
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import CompressedTensorsLinearMethod
1417+
if isinstance(self.quant_method, CompressedTensorsLinearMethod):
14181418
output_parallel = self.quant_method.apply(self, input_parallel, bias_, x_quant_scales=x_quant_scales)
14191419
else:
14201420
assert x_quant_scales is None, f"x_quant_scales input is not supported for {self.quant_method.__class__}"

‎vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -566,7 +566,8 @@ def create_weights(self, layer: torch.nn.Module,
566566
def apply(self,
567567
layer: torch.nn.Module,
568568
x: torch.Tensor,
569-
bias: Optional[torch.Tensor] = None):
569+
bias: Optional[torch.Tensor] = None,
570+
x_quant_scales: Optional[torch.Tensor] = None):
570571
"""
571572
Use the output of create_weights and the CompressedTensorsScheme
572573
associated with the layer to apply the forward pass with the
@@ -577,7 +578,11 @@ def apply(self,
577578
scheme = layer.scheme
578579
if scheme is None:
579580
raise ValueError("A scheme must be defined for each layer")
580-
return scheme.apply_weights(layer, x, bias=bias)
581+
elif isinstance(scheme, CompressedTensorsW8A8Fp8):
582+
return scheme.apply_weights(layer, x, bias=bias, x_quant_scales=x_quant_scales)
583+
else:
584+
assert x_quant_scales is None, f"x_quant_scales is not supported for {scheme.__class__}"
585+
return scheme.apply_weights(layer, x, bias=bias)
581586

582587

583588
class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):

‎vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,11 +156,13 @@ def create_weights(self, layer: torch.nn.Module,
156156
def apply_weights(self,
157157
layer: torch.nn.Module,
158158
x: torch.Tensor,
159-
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
159+
bias: Optional[torch.Tensor] = None,
160+
x_quant_scales: Optional[torch.Tensor] = None) -> torch.Tensor:
160161

161162
return self.fp8_linear.apply(input=x,
162163
weight=layer.weight,
163164
weight_scale=layer.weight_scale,
164165
out_dtype=self.out_dtype,
165166
input_scale=layer.input_scale,
166-
bias=bias)
167+
bias=bias,
168+
x_quant_scales=x_quant_scales)

‎vllm/model_executor/layers/quantization/fp8.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -404,8 +404,7 @@ def process_weights_after_loading(self, layer: Module) -> None:
404404
def apply(self,
405405
layer: torch.nn.Module,
406406
x: torch.Tensor,
407-
bias: Optional[torch.Tensor] = None,
408-
x_quant_scales: torch.Tensor = None) -> torch.Tensor:
407+
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
409408

410409
if self.use_marlin:
411410
return apply_fp8_marlin_linear(
@@ -428,7 +427,6 @@ def apply(self,
428427
bias=bias,
429428
cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
430429
use_aiter_and_is_supported=self.use_aiter_and_is_supported,
431-
input_quant_scale=x_quant_scales,
432430
)
433431

434432
return self.fp8_linear.apply(input=x,

‎vllm/model_executor/layers/quantization/utils/fp8_utils.py‎

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,6 @@ def apply_w8a8_block_fp8_linear(
115115
bias: Optional[torch.Tensor] = None,
116116
cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
117117
use_aiter_and_is_supported: bool = False,
118-
input_quant_scale: torch.Tensor = None,
119118
) -> torch.Tensor:
120119
assert input_scale is None
121120
# View input as 2D matrix for fp8 methods
@@ -136,12 +135,7 @@ def apply_w8a8_block_fp8_linear(
136135
scale_a=x_scale,
137136
scale_b=weight_scale.T)
138137
else:
139-
output_dtype = input.dtype
140-
if input_quant_scale is not None:
141-
q_input = input
142-
x_scale = input_quant_scale
143-
output_dtype = torch.bfloat16
144-
elif use_aiter_and_is_supported:
138+
if use_aiter_and_is_supported:
145139
# print(f"input_2d.shape: {input_2d.shape}, input_2d.stride: {input_2d.stride()}")
146140
q_input, x_scale = aiter_per1x128_quant(input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
147141
else:
@@ -156,11 +150,11 @@ def apply_w8a8_block_fp8_linear(
156150
x_scale,
157151
weight_scale,
158152
block_size,
159-
output_dtype=output_dtype)
153+
output_dtype=input.dtype)
160154

161155
if bias is not None:
162156
output = output + bias
163-
return output.to(dtype=output_dtype).view(*output_shape)
157+
return output.to(dtype=input.dtype).view(*output_shape)
164158

165159

166160
def apply_w8a8_block_fp8_linear_fake(
@@ -172,7 +166,6 @@ def apply_w8a8_block_fp8_linear_fake(
172166
bias: Optional[torch.Tensor] = None,
173167
cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
174168
use_aiter_and_is_supported: bool = False,
175-
input_quant_scale: torch.Tensor = None,
176169
) -> torch.Tensor:
177170
output_shape = [*input.shape[:-1], weight.shape[0]]
178171
return torch.empty(output_shape, dtype=input.dtype, device=input.device)
@@ -606,4 +599,4 @@ def grid(META):
606599
**config,
607600
)
608601

609-
return C
602+
return C

‎vllm/model_executor/layers/quantization/utils/w8a8_utils.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,7 @@ def apply(
345345
input_scale: Optional[torch.Tensor] = None,
346346
input_scale_ub: Optional[torch.Tensor] = None,
347347
bias: Optional[torch.Tensor] = None,
348+
x_quant_scales: Optional[torch.Tensor] = None,
348349
# TODO(luka) remove this parameter in favor of __init__
349350
use_per_token_if_dynamic: Optional[bool] = None
350351
) -> torch.Tensor:
@@ -404,6 +405,8 @@ def apply(
404405
qinput, x_scale = (
405406
torch.ops.vllm.rocm_aiter_per_tensor_quant_fp8(
406407
input_2d, scale=input_scale))
408+
elif x_quant_scales is not None:
409+
qinput, x_scale = input_2d, x_quant_scales
407410
else:
408411
qinput, x_scale = input_2d, input_scale
409412

‎vllm/model_executor/models/deepseek_v2.py‎

Lines changed: 61 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -62,11 +62,14 @@
6262
is_rocm_aiter_fusion_shared_expert_enabled,
6363
is_rocm_aiter_fuse_routed_scaling_factor,
6464
)
65+
from vllm.platforms import current_platform
6566

66-
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
67-
import aiter as rocm_aiter
68-
rocm_aiter_fp8_dtype = rocm_aiter.dtypes.fp8
69-
rocm_aiter_fp8_quant_group_size = 128
67+
import vllm.envs as envs
68+
69+
if current_platform.is_rocm():
70+
VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT=envs.VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT
71+
else:
72+
VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT=False
7073

7174
class DeepseekV2MLP(nn.Module):
7275

@@ -97,7 +100,10 @@ def __init__(
97100
self.act_fn = SiluAndMul()
98101

99102
def forward(self, x):
100-
gate_up, _ = self.gate_up_proj(x)
103+
x_quant_scales = None
104+
if isinstance(x, tuple):
105+
x, x_quant_scales = x
106+
gate_up, _ = self.gate_up_proj(x, x_quant_scales=x_quant_scales)
101107
x = self.act_fn(gate_up)
102108
x, _ = self.down_proj(x)
103109
return x
@@ -160,11 +166,16 @@ def __init__(
160166
)
161167

162168
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
169+
if isinstance(hidden_states, tuple):
170+
hidden_states_shared, hidden_states = hidden_states
171+
else:
172+
hidden_states_shared = hidden_states
173+
163174
num_tokens, hidden_dim = hidden_states.shape
164175
hidden_states = hidden_states.view(-1, hidden_dim)
165176
shared_output = None
166177
if self.n_shared_experts is not None and not is_rocm_aiter_fusion_shared_expert_enabled():
167-
shared_output = self.shared_experts(hidden_states)
178+
shared_output = self.shared_experts(hidden_states_shared)
168179
# router_logits: (num_tokens, n_experts)
169180
router_logits, _ = self.gate(hidden_states)
170181
if hidden_states.dtype != torch.float16:
@@ -501,13 +512,30 @@ def forward(
501512
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
502513
dim=-1,
503514
)
504-
hidden_states_or_q_c = self.q_a_layernorm(q_c)
515+
if VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT:
516+
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_per_token_quant
517+
fp8_dtype = current_platform.fp8_dtype()
518+
weight = self.q_a_layernorm.weight
519+
eps = self.q_a_layernorm.variance_epsilon
520+
weight2 = self.kv_a_layernorm.weight
521+
eps2 = self.kv_a_layernorm.variance_epsilon
522+
kv_c, k_pe = kv_lora.split(
523+
[self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
524+
(_, _), hidden_states_or_q_c, kv_c_normed, _ = fused_rms_fp8_per_token_quant(q_c, weight, eps,
525+
kv_c, weight2, eps2,
526+
dtype_quant=fp8_dtype,
527+
res1=None,
528+
output_unquantized_inp1=True,
529+
output_quantiezed_inp1=False)
530+
else:
531+
hidden_states_or_q_c = self.q_a_layernorm(q_c)
505532
else:
506533
hidden_states_or_q_c = hidden_states
507534
kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
508-
kv_c, k_pe = kv_lora.split(
509-
[self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
510-
kv_c_normed = self.kv_a_layernorm(kv_c)
535+
if not VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT:
536+
kv_c, k_pe = kv_lora.split(
537+
[self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
538+
kv_c_normed = self.kv_a_layernorm(kv_c)
511539
return self.mla_attn(hidden_states_or_q_c,
512540
kv_c_normed,
513541
k_pe,
@@ -585,28 +613,30 @@ def forward(
585613
residual: Optional[torch.Tensor],
586614
) -> torch.Tensor:
587615
# Self Attention
588-
weight = self.input_layernorm.weight
589-
eps = self.input_layernorm.variance_epsilon
590-
if residual is None:
591-
residual = hidden_states
592-
hidden_states, hidden_states_quant = fused_rms_fp8_group_quant(hidden_states, weight, eps,
593-
None, None, eps,
594-
group_size = rocm_aiter_fp8_quant_group_size,
595-
dtype_quant=rocm_aiter_fp8_dtype,
596-
res1=None)
616+
if VLLM_ROCM_USE_AITER_TRITON_FUSED_RMSNORM_FP8_QUANT:
617+
weight = self.input_layernorm.weight
618+
eps = self.input_layernorm.variance_epsilon
619+
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_per_token_quant
620+
fp8_dtype = current_platform.fp8_dtype()
621+
if residual is None:
622+
residual = hidden_states
623+
(hidden_states, hidden_states_quant), _, _, _ = fused_rms_fp8_per_token_quant(hidden_states, weight, eps,
624+
None, None, eps,
625+
dtype_quant=fp8_dtype,
626+
res1=None)
627+
else:
628+
(hidden_states, hidden_states_quant), _, _, residual = fused_rms_fp8_per_token_quant(hidden_states, weight, eps,
629+
None, None, eps,
630+
dtype_quant=fp8_dtype,
631+
res1=residual)
632+
hidden_states = (hidden_states, hidden_states_quant)
597633
else:
598-
(hidden_states, hidden_states_quant), residual = fused_rms_fp8_group_quant(hidden_states, weight, eps,
599-
None, None, eps,
600-
group_size = rocm_aiter_fp8_quant_group_size,
601-
dtype_quant=rocm_aiter_fp8_dtype,
602-
res1=residual)
603-
hidden_states = (hidden_states, hidden_states_quant)
604-
# if residual is None:
605-
# residual = hidden_states
606-
# hidden_states = self.input_layernorm(hidden_states)
607-
# else:
608-
# hidden_states, residual = self.input_layernorm(
609-
# hidden_states, residual)
634+
if residual is None:
635+
residual = hidden_states
636+
hidden_states = self.input_layernorm(hidden_states)
637+
else:
638+
hidden_states, residual = self.input_layernorm(
639+
hidden_states, residual)
610640
hidden_states = self.self_attn(
611641
positions=positions,
612642
hidden_states=hidden_states,

0 commit comments

Comments
 (0)