feat(pt2): add torch.compile support for FlagGems kernels used by vLLM 0.24

[JiaryCoder]

PR Category

Other

Type of Change

New Feature

Description

为 FlagGems 中会被 vLLM 0.24 及 vLLM 0.20.2 编译路径直接追踪的 kernel 提供透明的 torch.compile
契约（PT2 contracts），不修改任何 kernel 数学实现。

背景：vLLM 0.24 的 custom_ops=all 会以 fullgraph AOT 追踪 CustomOp.forward_oot。
当前 FlagGems 的公开启动器在该路径上不可追踪，会在 profile_run 阶段中断：
LibEntry.run（Unsupported method call: __getitem__）、入口处的裸 logger.debug
（logging.Logger method not supported）、以及 kernel 内的 ext. 模块别名
（Inductor 对非 triton 模块的闭包 dump 会丢失该别名，子进程重编译报
NameError('ext is not defined')；torch 2.11 上可复现）。

改动：

• flag_gems/pt2/（新增）——torch.library.triton_op + wrap_triton 包裹
  原始 LibEntry.jit_function / Autotuner 对象：RMSNorm、fused-add RMSNorm、
  pointwise 族（SiLU/GELU-and-mul）、in-place RoPE、MoE 基础算子（topk_softmax、
  moe_sum、grouped-topk 路由）、MHC。附声明式 CompileOpSpec manifest；
  pointwise 的 plan 在图外物化。
• utils/pointwise_dynamic.py——新增 PointwiseKernelMaterialization 与
  materialize()；codegen 直接生成 tl.program_id(0).to(tl.int64) /
  tl.num_programs(0).to(tl.int64)，保留原 helper 的 64 位整数语义。
• utils/codegen_config_utils.py——新增 get_heuristics_for_num_warps_fn()。
• ops/rms_norm.py、fused/fused_add_rms_norm.py、fused/rotary_embedding.py——
  ext.program_id(0) 改为 tl.program_id(0).to(tl.int64)，移除对应 helper 导入；
  显式区分 Triton 原生函数和 FlagGems 原有的 64 位封装。
• modules/{activation,normalization,rotary_embedding}.py——删除入口路径上的
  裸 logger.debug（仅日志，数值/控制流不变）。

eager 行为与之前完全一致（torch.equal 验证）。设计与 pt2/ 契约内容与内部
0.20.2 交付件（flaggems_pt2_final_20260901）逐字节一致，仅按本仓 lint 要求
（black/isort）做格式整理并更新部分注释。

Issue

• Associated with RotaryEmbedding fails during torch.compile graph capture with vllm-plugin-FL vllm-plugin-FL#298

Related PR:

• feat(compilation): freeze FL dispatch before torch.compile; route hot ops via flag_gems.pt2 (vLLM 0.24) vllm-plugin-FL#448

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue (RotaryEmbedding fails during torch.compile graph capture with vllm-plugin-FL vllm-plugin-FL#298).
• Change is fully covered by a UT.

Validation for the explicit int64 follow-up

本次追加提交涉及 4 个文件、13 处调用或代码生成模板。展开原 helper 函数体后，
修改前后的 AST 一致；pointwise 的 27 种生成组合通过源码语法与 AST 等价检查。
仓库配置的 pre-commit 检查通过（Flake8、isort、black、black-jupyter、空白及文件结尾检查）。
当前机器没有 PyTorch/Triton，未重新执行 GPU 数值、torch.compile 或端到端测试；
下方性能及 parity 数据来自此前的验证。

Performance

未改动任何 kernel 数学实现；本 PR 不引入性能变化。compile 模式下 eager/compiled
parity（rms_norm、fused_add_rms_norm、silu_and_mul、gelu_and_mul、rope(in-place)、
topk_softmax、moe_sum，bf16，NVIDIA，torch 2.11 / triton 3.6）max|diff| = 0.0。
端到端（经 vllm-plugin-FL）：MiniCPM5-1B TP=2、custom_ops=all 冷编译 14.2 s，
51/51 PIECEWISE + 51/51 FULL CUDA graph 捕获；Qwen3.6-35B-A3B eager vs compile
tie-aware cross-phase 16/16 位置逐位一致。