权重预取指南

权重预取指南#

权重预取通过在需要之前将权重预加载到缓存中来优化内存使用,从而最小化模型执行期间因内存访问造成的延迟。线性层有时表现出相对较高的MTE利用率。为了解决这个问题,我们创建了一个专门用于权重预取的独立流水线,该流水线与原始向量计算流水线(如量化、MoE门控top_k、RMSNorm和SwiGlu)并行运行。这种方法允许权重提前预加载到L2缓存中,减少线性层计算期间的MTE利用率,并通过最小化资源争用和优化数据流间接提高Cube计算效率。

由于我们使用向量计算来隐藏权重预取流水线,这会对计算产生影响。如果您优先考虑低延迟而非高吞吐量,最好不要启用预取。

快速开始#

使用 --additional-config '{"weight_prefetch_config": {"enabled": true}}' 来开启权重预取。

微调预取比例#

由于权重预取使用向量计算来隐藏权重预取流水线,预取大小的设置至关重要。如果大小太小,则无法充分发挥优化优势;而较大的大小可能导致资源争用,从而导致性能下降。为了适应不同的场景,我们添加了prefetch_ratio,允许根据具体工作负载灵活配置大小,详情如下:

使用"weight_prefetch_config"中的prefetch_ratio来为特定的线性层自定义权重预取比例。

“attn”和“moe”配置选项用于MoE模型,详情如下:

"attn": { "qkv": 1.0,  "o": 1.0},  "moe": {"gate_up": 0.8}

“mlp”配置选项用于优化Dense模型的性能,详情如下:

"mlp": {"gate_up": 1.0, "down": 1.0}

以上值为默认配置,当--max-num-seqs为144时,该默认值对Qwen3-235B-A22B-W8A8有良好性能;当--max-num-seqs为72时,对Qwen3-32B-W8A8有良好性能。

然而,这可能不是您场景下的最优配置。对于更高的并发度,可以尝试增加预取大小。对于较低的并发度,预取可能不会带来任何优势,因此可以减少大小或禁用预取。通过收集性能分析数据来确定预取大小是否合适。具体来说,检查预取操作(例如,MLP Down Proj权重预取)所需的时间是否与并行向量计算算子(例如,SwiGlu计算)所需的时间重叠,以及预取操作是否不晚于向量计算算子的完成时间。在性能分析时间线中,预取操作显示为单个流上的CMO操作;此CMO操作即为预取操作。

注意:

  1. MLP down投影的权重预取依赖于序列并行,如果您想为mlp down开启预取,请同时启用序列并行。

  2. 由于当前L2缓存的大小,最大预取量不能超过18MB。如果prefetch_ratio * linear_layer_weight_size >= 18 * 1024 * 1024字节,后端将只预取18MB。

示例#

  1. 对于MoE模型:

        --additional-config \
        '{
            "weight_prefetch_config": {
                "enabled": true,
                "prefetch_ratio": {
                    "attn": {
                        "qkv": 1.0,
                        "o": 1.0
                    },
                    "moe": {
                        "gate_up": 0.8
                    }
                }
            }
        }'
    
  2. 对于Dense模型:

    以下是默认配置,当--max-num-seqs为72时,该配置可为Qwen3-32B-W8A8带来良好性能

        --additional-config \
        '{
            "weight_prefetch_config": {
                "enabled": true,
                "prefetch_ratio": {
                    "mlp": {
                        "gate_up": 1.0,
                        "down": 1.0
                    }
                }
            }
        }'