EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

upgrade torch npu version (#4433)

Browse Source 
vLLM graph feature now rely on torch >=2.8. To make graph mode work, we
need upgrade torch version as well. For long term support, upgrade torch
to a newer one is good to go as well.

Related vLLM change: https://github.com/vllm-project/vllm/pull/25110


- vLLM version: v0.11.2
- vLLM main: https://github.com/vllm-project/vllm/commit/v0.11.2
This commit is contained in:
wangxiyuan
2025-12-01 19:01:55 +08:00
committed by GitHub
parent f1f6370ed9
commit 0d14f635b4
22 changed files with 63 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
vllm_ascend/quantization/w8a8.py
View File

@@ -119,8 +119,10 @@ class AscendW8A8LinearMethod:
weight=layer.weight,
start_flag=x,
)
quant_comm_config = getattr(layer, "_quant_comm_config", {})
try:
quant_comm_config = getattr(layer, "_quant_comm_config")
except AttributeError:
quant_comm_config = {}
comm_fn = quant_comm_config.get("communication_fn")
enable_flashcomm2_quant_comm = comm_fn is not None and (
"o_proj" in layer.prefix or "out_proj" in layer.prefix)
@@ -151,8 +153,12 @@ class AscendW8A8LinearMethod:
)
quant_bias = layer.quant_bias if tp_rank == 0 else None
if getattr(layer, "ascend_quant_method",
"") == COMPRESSED_TENSORS_METHOD:
try:
ascend_quant_method = getattr(layer, "ascend_quant_method")
except AttributeError:
ascend_quant_method = ""
if ascend_quant_method == COMPRESSED_TENSORS_METHOD:
quant_bias = bias
if get_ascend_device_type() == AscendDeviceType._310P:
@@ -194,8 +200,13 @@ class AscendW8A8LinearMethod:
layer.weight_scale.data = torch.flatten(layer.weight_scale.data)
layer.weight_offset.data = torch.flatten(layer.weight_offset.data)
layer.bias.data = layer.bias.data.to(layer.weight_scale.data.dtype)
if getattr(layer, "ascend_quant_method",
"") == COMPRESSED_TENSORS_METHOD:
try:
ascend_quant_method = getattr(layer, "ascend_quant_method")
except AttributeError:
ascend_quant_method = ""
if ascend_quant_method == COMPRESSED_TENSORS_METHOD:
deq_scale = layer.input_scale.data * layer.weight_scale.data
layer.deq_scale = torch.nn.Parameter(deq_scale,
requires_grad=False)

23
vllm_ascend/quantization/w8a8_dynamic.py
View File

@@ -15,7 +15,7 @@
# limitations under the License.
#
from typing import Any, Callable, Dict, Optional, Tuple, Union
from typing import Any, Callable, Dict, Optional
import torch
import torch_npu
@@ -73,33 +73,20 @@ class AscendW8A8DynamicLinearMethod:
@staticmethod
def apply(
layer: torch.nn.Module,
x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
tp_rank: Optional[int] = 0,
) -> torch.Tensor:
config = getattr(layer, "_ascend_quant_config", {})
if not isinstance(x, tuple):
output_dtype = config.get("output_dtype", x.dtype)
quantized_x, dynamic_scale = torch_npu.npu_dynamic_quant(x)
else:
assert "output_dtype" in config.keys(), (
f"DynamicLinearMethod needs explicitly specified `output_dtype`"
f"for pre-quantized input, got config [{config}]")
output_dtype = config["output_dtype"]
quantized_x, dynamic_scale = x
pertoken_scale = (dynamic_scale
if config.get("pertoken_scale", True) else None)
quantized_x, pertoken_scale = torch_npu.npu_dynamic_quant(x)
output = torch_npu.npu_quant_matmul(
quantized_x,
layer.weight,
layer.weight_scale,
pertoken_scale=pertoken_scale,
bias=bias,
output_dtype=output_dtype,
output_dtype=x.dtype,
)
return ((output, dynamic_scale)
if config.get("return_scale", False) else output)
return output
def process_weights_after_loading(self, layer):
if self.transpose_weight: