Fix CI

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>

Author: cjluo-nv

examples/deepseek/ptq.py

@@ -56,7 +56,7 @@
 from modelopt.torch.export.model_config import KV_CACHE_FP8
 from modelopt.torch.export.quant_utils import get_quant_config
 from modelopt.torch.quantization.nn import TensorQuantizer
-from modelopt.torch.quantization.triton.fp8_kernel import weight_dequant
+from modelopt.torch.quantization.triton import weight_dequant
 from modelopt.torch.quantization.utils import (
     is_quantized_column_parallel_linear,
     is_quantized_parallel_linear,

modelopt/torch/quantization/plugins/huggingface.py

@@ -49,7 +49,7 @@
 from ..conversion import register
 from ..nn import QuantInputBase, QuantModule, QuantModuleRegistry, TensorQuantizer
 from ..nn.modules.quant_linear import _QuantLinear
-from ..triton.fp8_kernel import weight_dequant
+from ..triton import weight_dequant
 from ..utils import replace_function
 from .attention import register_attention_for_kv_quant
 from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear, _QuantFunctionalMixin

modelopt/torch/quantization/triton/__init__.py

@@ -32,6 +32,7 @@
     ):
         # fp4_kernel works on any CUDA GPU with triton
         from .fp4_kernel import *
+        from .fp8_kernel import *
 
         # fp4_kernel_hopper requires compute >= 8.9 (uses tl.float8e4nv)
         if torch.cuda.get_device_capability() >= (8, 9):