[Perf] Add a switch to enable NZ layout in weights

vllm_ascend/ascend_config.py

@@ -41,6 +41,8 @@ def __init__(self, vllm_config):
         self.expert_map_path = additional_config.get("expert_map_path", None)
         self.chunked_prefill_for_mla = additional_config.get(
             "chunked_prefill_for_mla", False)
+        self.enable_weight_nz_layout = additional_config.get(
+            "enable_weight_nz_layout", False)
 
 
 class TorchairGraphConfig:

vllm_ascend/attention/mla_v1.py

@@ -20,7 +20,8 @@
 from vllm_ascend.multistream.context import get_multistream_comm_context
 from vllm_ascend.multistream.ms_split import model_input_split_v1_mla_attn
 from vllm_ascend.ops.attention import vanilla_chunked_prefill_mla
-from vllm_ascend.utils import npu_stream_switch, npu_wait_tensor
+from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, npu_stream_switch,
+                               npu_wait_tensor)
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
@@ -651,10 +652,12 @@ def get_and_maybe_dequant_weights(layer: LinearBase):
         self.W_UV = W_UV.transpose(0, 1).contiguous()
         # Convert from (L, N, P) to (N, P, L)
         self.W_UK_T = W_UK.permute(1, 2, 0).contiguous()
-
-        # Waiting for BMM NZ support
-        # self.W_UV.data = torch_npu.npu_format_cast(self.W_UV.data, 29)
-        # self.W_UK_T.data = torch_npu.npu_format_cast(self.W_UK_T.data, 29)
+        if get_ascend_config().enable_weight_nz_layout:
+            # cast quantized weight tensors in NZ layout for higher inference speed
+            self.W_UV.data = torch_npu.npu_format_cast(self.W_UV.data,
+                                                       ACL_FORMAT_FRACTAL_NZ)
+            self.W_UK_T.data = torch_npu.npu_format_cast(
+                self.W_UK_T.data, ACL_FORMAT_FRACTAL_NZ)
 
     def _compute_prefill_context(
         self,

vllm_ascend/quantization/w8a8.py

@@ -20,6 +20,9 @@
 import torch
 import torch_npu
 
+from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ
+
 
 def quant_per_tensor(in_tensor: torch.Tensor, input_scale: torch.Tensor,
                      input_offset: torch.Tensor):
@@ -37,6 +40,8 @@ class AscendW8A8LinearMethod:
     def __init__(self) -> None:
         # aclnn quant matmul requires to transpose matrix B, set to true by default.
         self.transpose_weight = True
+        ascend_config = get_ascend_config()
+        self.enable_weight_nz_layout = ascend_config.enable_weight_nz_layout
 
     @staticmethod
     def get_weight(
@@ -110,6 +115,9 @@ def process_weights_after_loading(self, layer):
             requires_grad=False).to(layer.aclnn_input_scale.dtype)
         if self.transpose_weight:
             layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
-        layer.weight.data = torch_npu.npu_format_cast(layer.weight.data, 29)
+        if self.enable_weight_nz_layout:
+            # cast quantized weight tensors in NZ layout for higher inference speed
+            layer.weight.data = torch_npu.npu_format_cast(
+                layer.weight.data, ACL_FORMAT_FRACTAL_NZ)
         layer.weight_scale.data = torch.flatten(layer.weight_scale.data)
         layer.weight_offset.data = torch.flatten(layer.weight_offset.data)

vllm_ascend/quantization/w8a8_dynamic.py

@@ -556,6 +556,8 @@ class AscendW8A8DynamicLinearMethod:
 
     def __init__(self):
         self.transpose_weight = True
+        ascend_config = get_ascend_config()
+        self.enable_weight_nz_layout = ascend_config.enable_weight_nz_layout
 
     @staticmethod
     def get_weight(input_size: int, output_size: int,
@@ -617,8 +619,10 @@ def apply(
     def process_weights_after_loading(self, layer):
         if self.transpose_weight:
             layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
-        # cast quantized weight tensors in NZ format (29) for higher inference speed
-        layer.weight.data = torch_npu.npu_format_cast(layer.weight.data, 29)
+        if self.enable_weight_nz_layout:
+            # cast quantized weight tensors in NZ layout for higher inference speed
+            layer.weight.data = torch_npu.npu_format_cast(
+                layer.weight.data, ACL_FORMAT_FRACTAL_NZ)
         layer.weight_scale.data = layer.weight_scale.data.flatten()
         layer.weight_scale_fp32 = layer.weight_scale.data.to(torch.float32)
         layer.weight_offset.data = layer.weight_offset.data.flatten()
@@ -635,6 +639,7 @@ def __init__(self):
 
         ascend_config = get_ascend_config()
         self.torchair_graph_enabled = ascend_config.torchair_graph_config.enabled
+        self.enable_weight_nz_layout = ascend_config.enable_weight_nz_layout
 
         try:
             device_group = self.ep_group.device_group
@@ -816,6 +821,12 @@ def process_weights_after_loading(self, layer):
                 1, 2).contiguous()
             layer.w2_weight.data = layer.w2_weight.data.transpose(
                 1, 2).contiguous()
+        if self.enable_weight_nz_layout:
+            # cast quantized weight tensors in NZ layout for higher inference speed
+            layer.w13_weight.data = torch_npu.npu_format_cast(
+                layer.w13_weight.data, ACL_FORMAT_FRACTAL_NZ)
+            layer.w2_weight.data = torch_npu.npu_format_cast(
+                layer.w2_weight.data, ACL_FORMAT_FRACTAL_NZ)
         if envs.VLLM_ENABLE_FUSED_EXPERTS_ALLGATHER_EP:
             torch_npu.npu_format_cast_(layer.w2_weight, ACL_FORMAT_FRACTAL_NZ)
         layer.w13_weight_scale.data = layer.w13_weight_scale.data.view(