with _set_default_torch_dtype(torch.float16):
    with torch.device("cuda"):
        hf_quant_config = getattr(
            self.model_config.hf_config, "quantization_config", None
        )
        # 当然这里暂时只支持awq，实际应该要支持更多做更多ifelse判断量化方法
        if hf_quant_config is not None:
            quant_config = AWQConfig.from_config(hf_quant_config)
            linear_method = quant_config.get_linear_method()
        model = model_class(
            config=self.model_config.hf_config, linear_method=linear_method
        )
    model.load_weights(
        self.model_config.path,
        cache_dir=None,
        load_format=self.load_format,
        revision=None,
    )
self.model = model

class AWQLinearMethod(LinearMethodBase):
    def create_weights(
        self, input_size: int, output_size: int, params_dtype: torch.dtype
    ) -> Dict[str, torch.Tensor]:
        pack_factor = 32 // self.weight_bits  # 4-bit → 8

        # 1. 创建qweight: [input_size, output_size // pack_factor]
        #    形状: [2048, 512] (tp=2, 7B模型)
        qweight = Parameter(
            torch.empty(input_size, output_size // pack_factor,
                       device="cuda", dtype=torch.int32),
            requires_grad=False,
        )
        # 关键：设置属性，用于后续权重加载
        set_weight_attrs(qweight, {
            "input_dim": 0,
            "output_dim": 1,
            "packed_dim": 1,        # 标记维度1是打包的
            "pack_factor": 8,       # 8个4-bit值打包成1个int32
        })

        # 2. 创建qzeros: [input_size // group_size, output_size // pack_factor]
        #    形状: [16, 512] (group_size=128, 2048/128=16)
        qzeros = Parameter(
            torch.empty(input_size // self.quant_config.group_size,
                       output_size // pack_factor,
                       device="cuda", dtype=torch.int32),
            requires_grad=False,
        )
        set_weight_attrs(qzeros, {
            "input_dim": 0,
            "output_dim": 1,
            "packed_dim": 1,
            "pack_factor": 8,
        })

        # 3. 创建scales: [input_size // group_size, output_size]
        #    形状: [16, 4096]
        scales = Parameter(
            torch.empty(input_size // self.quant_config.group_size,
                       output_size,
                       device="cuda", dtype=params_dtype),  # float16
            requires_grad=False,
        )
        set_weight_attrs(scales, {
            "input_dim": 0,
            "output_dim": 1,
        })

        return {"qweight": qweight, "qzeros": qzeros, "scales": scales}

class RowParallelLinear(torch.nn.Module):
    def forward(self, input_):
        # input_.shape: [batch_size, seq_len, hidden_size]
        #            或: [total_tokens, hidden_size]

        # 1. 处理输入的切分
        if self.input_is_parallel:
            input_parallel = input_
        else:
            tp_rank = get_tensor_model_parallel_rank()
            splitted_input = split_tensor_along_last_dim(
                input_, num_partitions=self.tp_size
            )
            input_parallel = splitted_input[tp_rank]

        # 2. 关键：调用linear_method.apply_weights()
        #    对于AWQ：self.linear_method = AWQLinearMethod
        #    对于非量化：self.linear_method = UnquantizedLinearMethod
        output_parallel = self.linear_method.apply_weights(
            self.linear_weights,  # {qweight, qzeros, scales}
            input_parallel,
            bias=None
        )

        # 3. All-Reduce（TP场景）
        if self.reduce_results and self.tp_size > 1:
            output_ = tensor_model_parallel_all_reduce(output_parallel)
        else:
            output_ = output_parallel

        return output_, output_bias

class AWQLinearMethod(LinearMethodBase):
    def apply_weights(
        self,
        weights: Dict[str, torch.Tensor],  # {qweight, qzeros, scales}
        x: torch.Tensor,                   # [M, K]
        bias: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        qweight = weights["qweight"]    # [K, N//8]
        qzeros = weights["qzeros"]      # [K//G, N//8]
        scales = weights["scales"]      # [K//G, N]
        pack_factor = self.quant_config.pack_factor  # 8

        # 计算输出形状
        out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)  # [M, N]
        reshaped_x = x.reshape(-1, x.shape[-1])  # [M, K]

        # 关键：调用Triton内核进行量化GEMM
        # 内部会：解包4-bit → 应用zero/scale → 矩阵乘法
        out = awq_gemm_triton(
            reshaped_x,  # [M, K]
            qweight,     # [K, N//8]
            scales,      # [K//G, N]
            qzeros,      # [K//G, N//8]
            pack_factor  # 8
        )

        if bias is not None:
            out = out + bias

        return out.reshape(out_shape)  # [M, N]

class UnquantizedLinearMethod(LinearMethodBase):
    def apply_weights(
        self,
        weights: Dict[str, torch.Tensor],  # {weight}
        x: torch.Tensor,
        bias: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        weight = weights["weight"]  # [output_size, input_size]

        # 直接使用PyTorch的F.linear
        if self.separate_bias_add:
            if bias:
                return F.linear(x, weight) + bias
            return F.linear(x, weight)
        return F.linear(x, weight, bias)

# 抽象接口
class LinearMethodBase:
    def create_weights(...) -> Dict[str, Tensor]: ...
    def apply_weights(...) -> Tensor: ...

# 具体策略
class UnquantizedLinearMethod(LinearMethodBase): ...
class AWQLinearMethod(LinearMethodBase): ...

# 使用方
class RowParallelLinear:
    def __init__(self, ..., linear_method: LinearMethodBase):
        self.linear_method = linear_method  # 运行时决定使用哪个策略

# 同样的调用代码，不同的实现
output = self.linear_method.apply_weights(weights, x, bias)

# 非量化: F.linear(x, weight)
# AWQ: awq_gemm_triton(x, qweight, scales, qzeros)