在边缘设备或离线场景下部署超大语言模型一直是工程难题。120B 参数的模型即使采用 INT8 量化也需要约 150GB 存储空间，完整加载到单设备显存几乎不可能。tinygrad 作为轻量级深度学习框架，通过 JIT 编译、内存分层和混合精度等技术，为 Tinybox 这类离线设备提供了可行的部署路径。本文将从工程落地的角度，给出关键优化参数与监控要点。
1 TinyJit 图执行：消除内核启动开销
120B 模型推理的核心瓶颈在于内核启动频率。逐 token 生成时，每个前向传播若调用数百个小内核，整体延迟将迅速累积。tinygrad 的 TinyJit 装饰器可以将整个 forward 过程编译为融合后的稳定内核，这是离线部署的首要优化。
实践中应将模型的完整解码路径包装在 TinyJit 函数内，确保输入输出张量均为已实现状态。预热阶段使用与实际推理相同的 batch size、序列长度和 KV cache 形状执行 3 到 5 次虚拟调用，使 JIT 编译器完成形状推导与内核融合。生产环境中应避免在 jitted 函数内部创建新张量或进行动态分支判断，所有形状信息须在预热阶段确定。对于 Tinybox 这类资源受限设备，推荐 batch size 设为 1 到 4，序列长度根据实际需求固定在 2048 或 4096。
2 形状固定与控制流静态化
JIT 编译器在形状不确定时会产生保守代码，导致内核融合效果大打折扣。针对 120B 模型，必须在初始化阶段明确所有维度参数：embedding dimension、attention head 数量、head dimension、层数以及 KV cache 的内存布局。
建议将 KV cache 预分配为固定形状的张量，例如 [n_layers, batch, n_heads, max_seq, head_dim]，解码时通过切片而非重新分配方式访问。控制流应全部移至 jitted 函数外部：序列长度判断、batch 大小分支、采样策略选择等逻辑均在 Python 层处理，仅将计算密集的核心前向传播送入 JIT。这种做法能让 tinygrad 生成稳定且充分融合的内核，对 attention 块和 MLP 块的优化尤为关键。
3 Attention 与 KV Cache 的带宽优化
120B 模型推理中，attention 机制和 KV cache 的内存带宽通常占据 60% 以上的运行时间。tinygrad 支持类似 Flash Attention 的融合模式：将 QK^T、softmax 和 V 乘法合并为单一内核，避免完整注意力矩阵物化。
工程实现上，KV cache 应保持每层每头连续存储，避免在每个解码步进行转置操作。解码模式下每次仅处理一个新 token，直接 attend 到现有 cache 而非重新计算历史 token 的 attention。若设备支持自定义内核，可使用 tinygrad 的底层算子将解量化、矩阵乘和注意力计算融合为一到两个内核，显著降低内存带宽压力。
4 量化策略与混合精度配置
全精度 FP32 或 BF16 下的 120B 模型需要约 240GB 到 480GB 内存，远超大多数边缘设备的承载能力。量化是实现离线部署的必要手段。
推荐策略如下：模型权重存储为低精度格式（如 INT4 或自定义 4-bit 布局），在送往计算核前实时解量化为 FP16；激活值和 KV cache 保留在 FP16 以维持数值稳定性；矩阵乘法操作使用设备原生快速路径，例如 NVIDIA 或 AMD GPU 上的 FP16 计算单元。若使用外部量化后的 checkpoint（如 GPT-OSS 120B 的 FP4/INT4 版本），应确保解量化和矩阵乘法在同一内核内完成，避免将大尺寸解量化张量加载至显存。
量化精度选择需要权衡模型质量与设备内存：INT8 量化通常能保持 95% 以上的模型能力，而 INT4 可能出现明显的能力下降，建议在部署前使用 perplexity 或下游任务指标评估量化后的实际效果。
5 设备放置与模型并行策略
单 GPU 无法容纳 120B 模型全部权重与 KV cache，必须采用分布式策略。tinygrad 支持按张量指定设备，可实现以下几种并行方案：
模型并行将大权重矩阵按 hidden dimension 拆分到 2 到 4 块 GPU，通过 all-reduce 汇总输出；流水线并行将模型层划分为多个阶段，每个阶段部署在不同设备上，形成微批处理流水线；内存受限场景下，可将冷门层或 MoE 专家卸载至 CPU RAM，仅将热点层和 KV cache 保留在 GPU。
Tinybox 设备通常配备多块消费级 GPU，合理的层间分区能显著提升吞吐量。配置时应记录各阶段的内存占用，确保每块设备的峰值使用率不超过设备显存的 80%，留出余量应对激活值波动。
6 Python 层开销控制与批处理策略
解码阶段逐 token 采样时，Python 解释器开销不可忽视。优化方向包括：尽可能扩大 jitted 函数的覆盖范围，将循环内的计算全部吸入图中；批量处理多个 prompt 以摊薄启动开销；采样逻辑仅在必要时拉取最终 logits 至主机端，避免中间张量的 host-device 传输。
对于交互式离线场景，建议使用 1 到 4 的小 batch size；对于离线批处理任务，可将 batch size 提升至 8 或 16 以提升吞吐量。上下文长度应根据实际需求严格限制，KV cache 内存占用与序列长度呈线性关系，减半 max context 可释放约一半的 cache 显存，为更深层的模型或更大 batch 提供空间。
7 监控指标与性能调优
部署后应持续追踪以下关键指标：端到端 tokens per second（包括 prompt 阶段和解码阶段）、各层或各模块耗时分布（特别是 attention 与 MLP 的时间占比）、GPU 利用率与显存带宽利用率、每次前向传播的内核数量（稳态下应保持低值且不变）。
建议从缩小版模型（如 7B 参数）开始验证上述优化策略的效果，验证通过后再扩展至 120B 规模。tinygrad 提供了简洁的性能分析接口，可快速定位瓶颈所在。硬件层面，Tinybox 的散热设计需关注长时间推理下的温控策略，避免因热降频导致性能波动。
总结
在 Tinybox 等离线设备上部署 120B 模型，本质上是在极端内存约束下寻求计算效率与模型能力的平衡。tinygrad 提供的 TinyJit 编译、形状静态化、attention 融合、量化与混合精度、以及多设备并行等能力，构成了完整的优化工具链。工程落地的关键在于：预先固定所有维度信息以释放 JIT 融合潜力、通过量化与内存分层突破显存瓶颈、使用合理的模型并行策略分配计算资源、并建立持续的性能监控体系。掌握这些参数与策略，即使面对百亿级参数模型，离线部署也不再是遥不可及的目标。
资料来源：tinygrad 官方文档（docs.tinygrad.org/quickstart/）、Reddit r/LocalLLaMA 关于 GPT-OSS 120B 本地推理优化讨论。