diff --git a/pyproject.toml b/pyproject.toml index 1cbf03dac..95e95fed8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -252,3 +252,4 @@ fullset = "fullset" Vai = "Vai" tockens = "tockens" CANN = "CANN" +ND = "ND" diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 8a30b4c7a..e371c73a6 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -31,6 +31,14 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): _omni_enum = OmniPlatformEnum.NPU dist_backend: str = "hccl" + @classmethod + def set_device(cls, device: torch.device) -> None: + super().set_device(device) + # Ascend quantized weights are converted from ND to FRACTAL_NZ + # after loading. Enable internal format so the NZ storage layout + # is preserved for fused NPU kernels. + torch.npu.config.allow_internal_format = True + @classmethod def get_omni_ar_worker_cls(cls) -> str: return "vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker"