xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-21 08:37:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

use for loops for fp8 linear layers init in tests

Browse Source 
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
This commit is contained in:
vllmellm 2025-11-07 07:39:08 +00:00
parent 56a05cd818
commit 9ff9b44e0d
2 changed files with 30 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
tests/compile/test_fusion.py
View File

@ -73,45 +73,36 @@ class TestModel(torch.nn.Module):
] ]
with override_cutlass_fp8_supported(not cuda_force_torch): with override_cutlass_fp8_supported(not cuda_force_torch):
self.fp8_linear_1 = TestFP8Layer( self.fp8_linear_layers = [
self.activation_quant_key, TestFP8Layer(
self.weight_quant_key, self.activation_quant_key,
self.w[0], self.weight_quant_key,
self.wscale[0], self.w[i],
self.scale[0], self.wscale[i],
) input_scale=self.scale[i],
self.fp8_linear_2 = TestFP8Layer( )
self.activation_quant_key, for i in range(3)
self.weight_quant_key, ]
self.w[1],
self.wscale[1],
self.scale[1],
)
self.fp8_linear_3 = TestFP8Layer(
self.activation_quant_key,
self.weight_quant_key,
self.w[2],
self.wscale[2],
self.scale[2],
)
self.enable_rms_norm_custom_op = self.norm[0].enabled() self.enable_rms_norm_custom_op = self.norm[0].enabled()
self.enable_quant_fp8_custom_op = self.fp8_linear.is_quant_fp8_enabled() self.enable_quant_fp8_custom_op = self.fp8_linear_layers[
0
].is_quant_fp8_enabled()
def forward(self, x): def forward(self, x):
# avoid having graph input be an arg to a pattern directly # avoid having graph input be an arg to a pattern directly
x = resid = torch.relu(x) x = resid = torch.relu(x)
y = self.norm[0](x) y = self.norm[0](x)
x2 = self.fp8_linear_1(y) x2 = self.fp8_linear_layers[0](y)
# make sure resid is used for replacement to work # make sure resid is used for replacement to work
y2, resid = self.norm[1](x2, resid) y2, resid = self.norm[1](x2, resid)
x3 = self.fp8_linear_2(y2) x3 = self.fp8_linear_layers[1](y2)
y3, resid = self.norm[2](x3, resid) # use resid here y3, resid = self.norm[2](x3, resid) # use resid here
x4 = self.fp8_linear_3(y3) x4 = self.fp8_linear_layers[2](y3)
y4, resid = self.norm[3](x4, resid) # use resid here y4, resid = self.norm[3](x4, resid) # use resid here
return y4 return y4

41
tests/compile/test_fusion_all_reduce.py
View File

@ -90,29 +90,16 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
for _ in range(3) for _ in range(3)
] ]
self.fp8_linear_1 = TestFP8Layer( self.fp8_linear_layers = [
self.quant_key, TestFP8Layer(
self.quant_key, self.quant_key,
self.weight[0], self.quant_key,
self.wscale[0], self.weight[i],
input_scale=self.input_scale[0], self.wscale[i],
) input_scale=self.input_scale[i],
)
self.fp8_linear_2 = TestFP8Layer( for i in range(3)
self.quant_key, ]
self.quant_key,
self.weight[1],
self.wscale[1],
input_scale=self.input_scale[1],
)
self.fp8_linear_3 = TestFP8Layer(
self.quant_key,
self.quant_key,
self.weight[2],
self.wscale[2],
input_scale=self.input_scale[2],
)
def forward(self, hidden_states): def forward(self, hidden_states):
# avoid having graph input be an arg to a pattern directly # avoid having graph input be an arg to a pattern directly
@ -120,17 +107,17 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
x = resid = tensor_model_parallel_all_reduce(z) x = resid = tensor_model_parallel_all_reduce(z)
y = self.norm[0](x) y = self.norm[0](x)
z2 = self.fp8_linear_1(y) z2 = self.fp8_linear_layers[0](y)
x2 = tensor_model_parallel_all_reduce(z2) x2 = tensor_model_parallel_all_reduce(z2)
y2, resid = self.norm[1](x2, resid) y2, resid = self.norm[1](x2, resid)
z3 = self.fp8_linear_2(y2) z3 = self.fp8_linear_layers[1](y2)
x3 = tensor_model_parallel_all_reduce(z3) x3 = tensor_model_parallel_all_reduce(z3)
y3, resid = self.norm[2](x3, resid) # use resid here y3, resid = self.norm[2](x3, resid) # use resid here
z4 = self.fp8_linear_3(y3) z4 = self.fp8_linear_layers[2](y3)
x4 = tensor_model_parallel_all_reduce(z4) x4 = tensor_model_parallel_all_reduce(z4)
y4, resid = self.norm[3](x4, resid) # use resid here y4, resid = self.norm[3](x4, resid) # use resid here
@ -143,7 +130,7 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
return [ return [
torch.ops.vllm.all_reduce.default, torch.ops.vllm.all_reduce.default,
torch.ops._C.static_scaled_fp8_quant.default torch.ops._C.static_scaled_fp8_quant.default
if self.fp8_linear.is_quant_fp8_enabled() if self.fp8_linear_layers[0].is_quant_fp8_enabled()
else torch.ops.aten.reciprocal.default, else torch.ops.aten.reciprocal.default,
] ]