Add situ activation - #395
Conversation
| activation_beta: Optional[float] = None, | ||
| activation_linear_beta: Optional[float] = None, |
There was a problem hiding this comment.
🔴 critical: 保持 fast_math 的原有位置参数索引: 现有调用若以位置参数传入 fast_math（如 ..., 'swiglu', 10, False），这里新增的参数会让 False 被绑定为 activation_beta，而 fast_math 回落为 True；SwiGLU 路径不会拒绝该 beta，因此会静默改变数值和性能行为。应保留原参数顺序或将新参数设为仅限关键字。
🤖 v6
| mma_type: str = 'fp8xfp4', | ||
| activation: str = 'swiglu'): | ||
| assert activation == 'swiglu', f'Only `swiglu` activation is supported, got `{activation}`' | ||
| assert activation in ('swiglu', 'situ'), f'Unsupported activation `{activation}`' |
There was a problem hiding this comment.
🟡 warning: 不要为 BF16 路径放行 situ: 当 mma_type='bf16xbf16' 且 activation='situ' 时，此处及缓冲区大小接口会接受配置，但 bf16_mega_moe 仍执行 DG_HOST_ASSERT(activation == "swiglu")，因此该公开配置必然在运行时失败。应按 MMA 类型限制这里的取值，或同时实现 BF16 SiTU。
🤖 v6
| y=y, l1_weights=transformed_l1_weights, l2_weights=transformed_l2_weights, | ||
| sym_buffer=buffer, | ||
| cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused, | ||
| activation=args.activation, |
There was a problem hiding this comment.
🟡 warning: 避免用 SwiGLU 基线校验 SiTU 输出: 当 legacy deep_ep/tilelang 可用且使用 --activation situ 时，融合路径在这里切换为 SiTU，但 run_baseline 和共享专家基线仍固定调用 swiglu_apply_weight_to_fp8；默认正确性测试随后进行精确比较，因而必然失败。SiTU 应使用对应基线，或跳过该 legacy 比较并运行新的 PyTorch 参考实现。
🤖 v6
| l1 = x_ref[token] @ l1_weights_ref[expert].T | ||
| gate, up = l1.to(torch.bfloat16).float().chunk(2) | ||
| if args.activation == 'situ': |
There was a problem hiding this comment.
🟡 warning: 在参考实现中复现 activation_clamp: 在无 legacy 基线的单卡参考测试中，只要 activation_clamp 有限（默认是 10），内核会先将 gate 上界及 up 双边截断，而这里直接对未截断的 L1 输出应用激活。默认随机矩阵乘积通常远超 10，因此该断言比较的是不同函数，SwiGLU 和 SiTU 都可能稳定失败。
🤖 v6
| else: | ||
| create_inputs() | ||
|
|
||
| if (not is_legacy_loaded and not is_bf16xbf16 and num_ranks == 1 and |
There was a problem hiding this comment.
🔴 critical: New reference check ignores activation_clamp, so it fails with default arguments. The kernel clamps the L1 output before the activation (kActivationClamp, default --activation-clamp 10): the gate is clamped from above only (__hmin2), the up branch symmetrically (__hmax2 + __hmin2). The new PyTorch reference computes gate/up with no clamping at all. With the default --input-scale 1.0 --weight-scale 1.0 and hidden=7168 the L1 pre-activations have std ~ sqrt(7168) ~ 85, so virtually every element is clamped by the kernel and not by the reference, and assert_close(rtol=0.08, atol=5e-4) cannot pass. The check therefore only works with undocumented flag combinations (tiny scales, or --activation-clamp inf). Suggestion: mirror the kernel in the reference - gate = gate.clamp(max=args.activation_clamp) and up = up.clamp(-args.activation_clamp, args.activation_clamp) when the clamp is finite - and/or document/assert the required flags.
🤖 v7
| num_topk == 1 and num_shared_experts == 0): | ||
| fused_y = run_fused()[0] | ||
| ref_y = torch.empty_like(fused_y) | ||
| for token in range(num_tokens): |
There was a problem hiding this comment.
🟡 warning: The per-token Python reference loop is prohibitively slow, and masked tokens index expert -1. The loop runs once per token (up to --num-max-tokens-per-rank, default 8192) and does a 7168x6144 fp32 matmul plus an FP8 quantisation each iteration. Also, with --masked-ratio > 0, expert = topk_idx[token, 0].item() can be -1, silently selecting the last expert (it only happens to produce ~0 because topk_weights is 0 there). Suggestion: check a bounded random subset of tokens (e.g. 32-64) or batch the reference with a single grouped matmul, and continue explicitly when expert < 0 while asserting the fused output is zero.
🤖 v7
| f'{fused_y.float().abs().max().item():.3f}, ref mean/max=' | ||
| f'{ref_y.float().abs().mean().item():.3f}/{ref_y.float().abs().max().item():.3f}', | ||
| once_in_node=True) | ||
| torch.testing.assert_close(fused_y, ref_y, rtol=0.08, atol=5e-4) |
There was a problem hiding this comment.
🔵 suggestion: The tolerances are simultaneously too loose and too tight: rtol=0.08 allows an 8% per-element error (which would hide real regressions), while atol=5e-4 is far below the quantisation noise floor for near-zero elements of an output whose magnitude is O(10^2), making the check seed-dependent and flaky. Suggestion: use the repository's aggregate metric (calc_diff, already used at line 351) with a documented threshold, or keep assert_close with an atol derived from the output scale (e.g. 0.02 * ref_y.abs().mean()).
🤖 v7
| recipe: Tuple[int, int, int] = (1, 1, 32), | ||
| activation: str = 'swiglu', | ||
| activation_clamp: Optional[float] = None, | ||
| activation_beta: Optional[float] = None, |
There was a problem hiding this comment.
🔵 suggestion: The new public option is undocumented. activation='situ' and the activation_beta / activation_linear_beta parameters are new public API surface, but the Mega MoE section still describes the epilogue as 'SwiGLU' and the usage snippet does not mention the activation options. The kernel comments were only partially updated too (// Apply gated activation. at sm100_fp8_fp4_mega_moe.cuh:1045 while the surrounding 'Unified L1 epilogue: SwiGLU in-place' comment and csrc/jit_kernels/heuristics/mega_moe.hpp:153 still say SwiGLU). Suggestion: document the formula (gate = beta*tanh(x/beta)sigmoid(x), optional up = beta_lintanh(up/beta_lin)), the defaults (beta=1.0, beta_lin=-1.0 meaning disabled) and the FP8xFP4-only support, and refresh the SwiGLU-specific comments in the touched epilogue.
🤖 v7
| ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], | ||
| Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]: | ||
| assert activation == 'swiglu', f'Only `swiglu` activation is supported, got `{activation}`' | ||
| Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]: |
There was a problem hiding this comment.
🔵 suggestion: Unrelated indentation churn: the continuation line of the transform_weights_for_mega_moe return annotation was re-indented from 11 to 13 spaces, which has nothing to do with this feature and now mis-aligns with the opening bracket. Also activation stays a validate-only parameter in this function, so it reads as dead - a one-line comment saying situ reuses the same gate/up interleaved layout as swiglu would help. Suggestion: revert the whitespace change and add the clarifying comment.
🤖 v7
| @@ -168,6 +168,8 @@ static void fp8_fp4_mega_moe( | |||
| const std::tuple<int, int, int>& recipe, | |||
| const std::string& activation, | |||
| const std::optional<float>& activation_clamp_opt, | |||
There was a problem hiding this comment.
🔵 suggestion: Two positional parameters were inserted before fast_math in the pybind signature. _C.fp8_fp4_mega_moe is registered without py::arg names (register_apis), so the new activation_beta_opt / activation_linear_beta_opt are positional-only and any external caller that passed fast_math positionally now silently mis-binds a bool into a float slot. Suggestion: append the new parameters at the end of the signature, or add named py::arg(...) bindings, and mention the change in the release notes.
🤖 v7
🤖 ds-review-bot Code Reviewv6该变更会静默破坏现有位置参数调用，并暴露一个实际不受支持的 BF16 SiTU 配置。此外，两条新增或既有的正确性校验路径都没有与 SiTU 内核保持一致。 v4v7The MR adds a second gated activation ('situ': gate = beta*tanh(x/beta)sigmoid(x), with an optional soft-clamped up = beta_lintanh(up/beta_lin)) to the SM100 FP8xFP4 Mega-MoE epilogue, threads two new optional host parameters (activation_beta, activation_linear_beta) through the Python wrapper -> pybind API -> JIT codegen -> kernel template, and adds a PyTorch reference check plus new CLI knobs to tests/test_mega_moe.py. The kernel math is correct and matches the new PyTorch reference, the codegen bookkeeping lines up (26 placeholders / 26 arguments / 26 non-defaulted template params, same order), and hex-float Files reviewed: 5 |
2 participants
Test