huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixing moe import.

Browse Source
This commit is contained in:
Nicolas Patry 2025-01-30 11:05:21 +01:00
parent 351f3c6ee5
commit 7539881054
No known key found for this signature in database
GPG Key ID: D2920555C90F704C
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
server/text_generation_server/layers/fp8.py
View File

@ -5,7 +5,6 @@ from typing import Optional, Tuple, Type, Union, List
import torch import torch
from loguru import logger from loguru import logger
from moe_kernels.fp8_utils import w8a8_block_fp8_matmul, per_token_group_quant_fp8
from text_generation_server.utils.import_utils import SYSTEM from text_generation_server.utils.import_utils import SYSTEM
from text_generation_server.utils.weights import ( from text_generation_server.utils.weights import (
Weight, Weight,
@ -20,6 +19,12 @@ try:
except ImportError: except ImportError:
marlin_kernels = None marlin_kernels = None
try:
from moe_kernels.fp8_utils import w8a8_block_fp8_matmul, per_token_group_quant_fp8
except ImportError:
w8a8_block_fp8_matmul = None
per_token_group_quant_fp8 = None
quant_dtype: torch.dtype = ( quant_dtype: torch.dtype = (
torch.float8_e4m3fnuz if SYSTEM == "rocm" else torch.float8_e4m3fn torch.float8_e4m3fnuz if SYSTEM == "rocm" else torch.float8_e4m3fn
) )