[#10519][feat] AutoDeploy: support HF fine grained FP8

tensorrt_llm/_torch/auto_deploy/config/default.yaml

@@ -63,6 +63,8 @@ transforms:
     stage: pattern_matcher
   quantize_nvfp4_linear_from_config:
     stage: pattern_matcher
+  quantize_hf_fp8_linear_from_config:
+    stage: pattern_matcher
   quantize_fp8_bmm_from_config:
     stage: pattern_matcher
   quantize_fp8_from_graph:

tensorrt_llm/_torch/auto_deploy/custom_ops/torch_quant.py

@@ -319,3 +319,61 @@ def _fake(
     N_half = weight_quantized.shape[-2]
     N = N_half * 2
     return torch.empty((*input.shape[:-1], N), dtype=input.dtype, device=input.device)
+
+
+@torch.library.custom_op("auto_deploy::torch_fake_quant_hf_fp8_linear", mutates_args=())
+def torch_fake_quant_hf_fp8_linear(
+    input: torch.Tensor,  # [..., K]
+    weight_quantized: torch.Tensor,  # [N, K] float8_e4m3fn
+    bias: Optional[torch.Tensor],  # [N] or None
+    input_scale: List[torch.Tensor],  # unused for HF FP8 (input quantized on the fly)
+    weight_scale: List[torch.Tensor],  # [weight_scale_inv]
+    input_zp: List[torch.Tensor],  # unused
+    weight_zp: List[torch.Tensor],  # unused
+) -> torch.Tensor:
+    """HuggingFace FineGrainedFP8 linear operation.
+    - weight_scale[0] = weight_scale_inv (per-block weight scale)
+    - input_scale, input_zp, weight_zp are unused
+    - block_size is inferred from weight and weight_scale_inv shapes
+    """
+    from transformers.integrations.finegrained_fp8 import act_quant, w8a8_block_fp8_matmul_triton
+
+    weight_scale_inv = weight_scale[0]
+
+    # Infer block_size from weight and weight_scale_inv shapes
+    # weight shape: [N, K], weight_scale_inv shape: [N/block_n, K/block_k]
+    N, K = weight_quantized.shape
+    scale_n, scale_k = weight_scale_inv.shape
+    block_n = N // scale_n
+    block_k = K // scale_k
+    block_size = [block_n, block_k]
+
+    qinput, scale = act_quant(input, block_size[1])
+    output = w8a8_block_fp8_matmul_triton(
+        qinput,
+        weight_quantized,
+        scale,
+        weight_scale_inv,
+        block_size,
+        output_dtype=input.dtype,
+    )
+
+    if bias is not None:
+        output = output + bias
+
+    return output.to(dtype=input.dtype)
+
+
+@torch_fake_quant_hf_fp8_linear.register_fake
+def _torch_fake_quant_hf_fp8_linear_fake(
+    input: torch.Tensor,
+    weight_quantized: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    input_scale: List[torch.Tensor],
+    weight_scale: List[torch.Tensor],
+    input_zp: List[torch.Tensor],
+    weight_zp: List[torch.Tensor],
+) -> torch.Tensor:
+    """Fake implementation for torch.export tracing."""
+    out_features = weight_quantized.shape[0]
+    return torch.empty((*input.shape[:-1], out_features), dtype=input.dtype, device=input.device)

tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py

@@ -146,6 +146,8 @@ class HFQuantConfigReader(QuantConfigReader):
     Quantization reader that process transformers.quantizers.HFQuantizer functionality
     """
 
+    _SUPPORTED_QUANT_METHODS = ("mxfp4", "fp8")
+
     def __init__(self):
         super().__init__()
         self._hf_quantizer = None
@@ -180,7 +182,7 @@ def from_file(cls, ckpt_dir: str) -> Optional[Tuple["HFQuantConfigReader", Dict[
         # TODO(Fridah-nv):this class is only verified with GPT-OSS MXFP4, other hf quantizers
         # should have similar workflow and will be added to the pipeline
         quant_method = str(qconf.get("quant_method", "")).lower()
-        if quant_method != "mxfp4":
+        if quant_method not in cls._SUPPORTED_QUANT_METHODS:
             return None
 
         reader = cls()

tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py

@@ -633,3 +633,93 @@ def _apply(
         return gm, TransformInfo(
             skipped=False, num_matches=cnt, is_clean=cnt == 0, has_valid_shapes=True
         )
+
+
+@TransformRegistry.register("quantize_hf_fp8_linear_from_config")
+class HFFineGrainedFP8LinearQuantization(Quantization):
+    """Quantization transform for HuggingFace FineGrainedFP8 (block-wise FP8) models.
+
+    This transform replaces linear ops with the HF FineGrainedFP8 quantized op.
+    The HF FP8 format uses per-block weight scales (weight_scale_inv) and
+    dynamic input quantization.
+
+    Config format (from HF config.json):
+        "quantization_config": {
+            "quant_method": "fp8",
+            "weight_block_size": [128, 128],
+            "modules_to_not_convert": ["lm_head"]
+        }
+    """
+
+    algo_name = "fp8"
+
+    def target_op(self):
+        return torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear.default
+
+    def quantize_weight(self, w: torch.Tensor) -> torch.Tensor:
+        return torch.empty_like(w, dtype=torch.float8_e4m3fn, device=w.device)
+
+    def scale_names(self) -> List[str]:
+        return ["weight_scale_inv"]
+
+    def default_scales(self, original_weight_shape: Tuple) -> Dict[str, torch.Tensor]:
+        # Default block size is 128x128 for HF FP8
+        N, K = original_weight_shape
+        block_n, block_k = 128, 128
+        scale_shape = (N // block_n, K // block_k)
+        return {"weight_scale_inv": torch.ones(scale_shape, dtype=torch.bfloat16)}
+
+    def build_custom_args_for_linear(self, scales: Dict[str, Node]) -> Tuple:
+        return ([], [scales["weight_scale_inv"]], [], [])
+
+    def load_hook(self, state_dict, prefix, *args, weight_name: str):
+        """Load hook to handle HF FineGrainedFP8 checkpoint format.
+
+        HF FP8 checkpoints store:
+        - weight: float8_e4m3fn tensor
+        - weight_scale_inv: per-block scale tensor
+        """
+        if weight_name not in state_dict:
+            return
+
+        weight = state_dict[weight_name]
+        if weight.dtype == torch.float8_e4m3fn:
+            scale_inv_name = weight_name + "_scale_inv"
+            if scale_inv_name in state_dict:
+                # Rename to match our buffer name
+                mod_prefix = weight_name.rsplit(".", 1)[0]
+                state_dict[mod_prefix + ".weight_scale_inv"] = state_dict[scale_inv_name]
+
+    def _apply(
+        self,
+        gm: GraphModule,
+        cm: CachedSequenceInterface,
+        factory: ModelFactory,
+        shared_config: SharedConfig,
+    ) -> Tuple[GraphModule, TransformInfo]:
+        qcfg = factory.get_quant_config()
+        if not qcfg:
+            return gm, TransformInfo(
+                skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True
+            )
+
+        quant_method = str(qcfg.get("quant_method", "")).lower()
+        if quant_method != self.algo_name:
+            return gm, TransformInfo(
+                skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True
+            )
+
+        excluded = qcfg.get("modules_to_not_convert", [])
+
+        cnt = 0
+        for n in gm.graph.nodes:
+            if not is_linear_op(n):
+                continue
+            if should_skip_quantization(n, excluded):
+                continue
+            self._insert_quantized_linear(gm, n, is_quantized_graph=False)
+            cnt += 1
+
+        return gm, TransformInfo(
+            skipped=False, num_matches=cnt, is_clean=False, has_valid_shapes=(cnt == 0)
+        )

tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py

@@ -436,6 +436,47 @@ def shard_load_hook(
         return
 
 
+class HFFineGrainedFP8WeightShardingInfo(QuantizationShardingMixin, WeightShardingInfo):
+    """Tensor-parallel sharding for HuggingFace FineGrainedFP8 quantized linears.
+
+    HF FP8 uses per-block weight scales (weight_scale_inv) with shape [N/block_n, K/block_k].
+    When sharding the weight along a dimension, we also need to shard the scale tensor.
+    """
+
+    def scale_names(self) -> List[str]:
+        return ["weight_scale_inv"]
+
+    def shard_scales(
+        self,
+        dim: int,
+        rank: int,
+        world_size: int,
+        weight_shape: torch.Size,
+        *,
+        weight_scale_inv: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        # weight_scale_inv has shape [N/block_n, K/block_k]
+        # When we shard weight along dim, we need to shard scale along the same dim
+        sharded_scale = torch.tensor_split(weight_scale_inv, world_size, dim=dim)[rank]
+        return {"weight_scale_inv": sharded_scale}
+
+    def shard_load_hook(
+        self,
+        state_dict,
+        prefix,
+        *args,
+        weight_name: str,
+        weight_shape: torch.Size,
+        dim: int,
+        rank: int,
+        world_size: int,
+    ) -> None:
+        scale_key = weight_name + "_scale_inv"
+        if scale_key in state_dict:
+            scale = state_dict[scale_key]
+            state_dict[scale_key] = torch.tensor_split(scale, world_size, dim=dim)[rank]
+
+
 def _shard_fp4_weight_scale(weight_scale, sharded_uint8_weight_shape, dim, rank, world_size):
     assert weight_scale.dim() == 1
     weight_shape_original = list(sharded_uint8_weight_shape)
@@ -1051,6 +1092,10 @@ def _validate_sharded_shapes(
         lambda n: is_op(n, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear),
         FP4WeightShardingInfo,
     ),
+    (
+        lambda n: is_op(n, torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear),
+        HFFineGrainedFP8WeightShardingInfo,
+    ),
 ]
 
 

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

@@ -362,6 +362,7 @@ def is_fake_quantized_linear_op(node: Node) -> bool:
     quantized_linear_op = {
         torch.ops.auto_deploy.torch_fake_quant_fp8_linear,
         torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear,
+        torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear,
     }
 
     return is_op(node, quantized_linear_op)

tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py

@@ -275,6 +275,46 @@ def forward(self, x):
         )
 
 
+class FakeHFFP8Linear(nn.Linear):
+    """Fake HuggingFace FineGrainedFP8 linear layer for testing.
+
+    Mimics the behavior of transformers.integrations.finegrained_fp8.FP8Linear
+    with per-block quantization (block_size = [128, 128] by default).
+    """
+
+    def __init__(self, in_features, out_features, bias=True, block_size=None):
+        super().__init__(in_features, out_features, bias)
+        device = self.weight.device
+
+        if block_size is None:
+            block_n = min(128, out_features)
+            block_k = min(128, in_features)
+            block_size = [block_n, block_k]
+        self.block_size = block_size
+
+        N, K = self.weight.shape
+        block_n, block_k = block_size
+
+        weight_reshaped = self.weight.detach().view(N // block_n, block_n, K // block_k, block_k)
+        amax = weight_reshaped.abs().amax(dim=(1, 3)).to(torch.float32)  # [N/block_n, K/block_k]
+
+        eps = torch.finfo(torch.float32).tiny
+        weight_scale_inv = torch.clamp(amax / FP8_MAX, min=eps).to(device)
+
+        weight_fp8 = (
+            self.weight.detach().float()
+            / weight_scale_inv.repeat_interleave(block_n, dim=0).repeat_interleave(block_k, dim=1)
+        ).to(torch.float8_e4m3fn)
+
+        self.weight = nn.Parameter(weight_fp8)
+        self.register_buffer("weight_scale_inv", weight_scale_inv)
+
+    def forward(self, x):
+        return torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear(
+            x, self.weight, self.bias, [], [self.weight_scale_inv], [], []
+        )
+
+
 def generate_dynamic_shapes(max_batch_size, max_seq_len):
     dynamic_shapes = (
         {

tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py

@@ -9,12 +9,13 @@
 import torch.nn.functional as F
 from _dist_test_utils import get_device_counts
 from _graph_test_helpers import run_sharding_pattern_detection_test, run_test_transformed_gm
-from _model_test_utils import FakeFP8Linear
+from _model_test_utils import FakeFP8Linear, FakeHFFP8Linear
 
 import tensorrt_llm._torch.auto_deploy.distributed.common as dist_common
 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm
 from tensorrt_llm._torch.auto_deploy.transform.library.sharding import (
     FP8WeightShardingInfo,
+    HFFineGrainedFP8WeightShardingInfo,
     LayerType,
     ShardingTransformConfig,
     SplitDimension,
@@ -125,6 +126,23 @@ def forward(self, x):
         return self.linear2(y)
 
 
+class HFFP8MLP(nn.Module):
+    """MLP using HuggingFace FineGrainedFP8 quantization for testing."""
+
+    def __init__(self, in_features, out_features, bias=False):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        # Use larger features divisible by block size (128)
+        hidden_features = max(4 * in_features, 128)
+        self.linear1 = FakeHFFP8Linear(in_features, hidden_features, bias=bias)
+        self.linear2 = FakeHFFP8Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x):
+        y = F.relu(self.linear1(x))
+        return self.linear2(y)
+
+
 def _run_sharding_execution_job(
     model_cls: nn.Module,
     dist_op_expected: str,
@@ -150,6 +168,9 @@ def _run_sharding_execution_job(
         ).to(device="cuda", dtype=torch.float16)
     elif model_cls == FP8MLP:
         model = model_cls(num_features, num_features, bias=bias).to("cuda")
+    elif model_cls == HFFP8MLP:
+        # HFFP8MLP needs features divisible by 128 (block size)
+        model = model_cls(128, 128, bias=bias).to("cuda")
     else:
         model = model_cls(num_features, num_features, bias=bias).to(
             device="cuda", dtype=torch.float16
@@ -344,6 +365,26 @@ def _run_pattern_detection_job(
                             min_local_shape=1,
                         )
                     )
+        elif model_cls == HFFP8MLP:
+            for node in gm.graph.nodes:
+                if is_op(node, torch.ops.auto_deploy.torch_fake_quant_hf_fp8_linear):
+                    # linear1 should be sharded on dim=0, add_dist=False, min_local_shape=1
+                    # linear2 should be sharded on dim=1, add_dist=True, min_local_shape=1
+                    if "linear1" in node.args[1].name:
+                        dim = SplitDimension.COLUMN
+                        dist_op = None
+                    else:
+                        dim = SplitDimension.ROW
+                        dist_op = "all_reduce"
+                    expected_transformations.append(
+                        HFFineGrainedFP8WeightShardingInfo(
+                            target_node=node.name,
+                            split_dim=dim,
+                            config=config,
+                            dist_op=dist_op,
+                            min_local_shape=1,
+                        )
+                    )
 
     # get detected transformations
     optimizer = InferenceOptimizer(
@@ -376,6 +417,7 @@ def _run_pattern_detection_job(
     (
         (MLP, "torch_dist_all_reduce"),
         (FP8MLP, "torch_dist_all_reduce"),
+        (HFFP8MLP, "torch_dist_all_reduce"),
         (nn.Linear, "torch_dist_all_gather"),
         (GQA_Block, "torch_dist_all_reduce"),
     ),
@@ -401,6 +443,7 @@ def test_sharding(
     (
         (MLP, "torch_dist_all_reduce"),
         (FP8MLP, "torch_dist_all_reduce"),
+        (HFFP8MLP, "torch_dist_all_reduce"),
         (nn.Linear, "torch_dist_all_gather"),
         (GQA_Block, "torch_dist_all_reduce"),
     ),