[zhiding512]

Summary

_gen_opus_moe_stage1_a8w4_fake_tensors and _gen_opus_moe_stage2_a8w4_decode_fake_tensors were missing the scalar parameters of the ops they are registered for (gen_fake= of _opus_moe_stage1_a8w4_fwd_raw and _opus_moe_stage2_a8w4_decode_fwd_raw). torch_compile_guard calls the fake with every argument the op takes, so any call under FakeTensorMode raised TypeError. This adds the missing parameters; the fakes still return out.

Motivation

Found while replaying fused_moe on FakeTensors for the FlyDSL MoE AOT build (#6362, stacked on this PR): every tuned row whose stage 1 or stage 2 resolves to these opus kernels failed there (43 jobs).

Changes

• aiter/ops/opus/moe_stage1_a8w4.py: add topk, block_m, kernelName, inter_dim_pad, activation, swiglu_limit, situ_beta, situ_linear_beta.
• aiter/ops/opus/moe_stage2_a8w4.py: add token_num, topk, block_m, kernel_id, inter_dim_pad.

Eager execution is unaffected; only FakeTensor / torch.compile tracing reaches these functions.

Testing

Both ops called under FakeTensorMode with every argument they take (no GPU work):

	stage 1	stage 2
main	TypeError: ... takes 10 positional arguments but 18 were given	TypeError: ... takes 9 positional arguments but 14 were given
this PR	returns out	returns out

Check script

import torch
from torch._subclasses.fake_tensor import FakeTensorMode

from aiter.ops.opus.moe_stage1_a8w4 import _opus_moe_stage1_a8w4_fwd_raw
from aiter.ops.opus.moe_stage2_a8w4 import _opus_moe_stage2_a8w4_decode_fwd_raw


def t(*shape, dtype=torch.float32):
    return torch.empty(*shape, dtype=dtype, device="cuda")


with FakeTensorMode():
    fp8, i32, u8 = torch.float8_e4m3fn, torch.int32, torch.uint8
    ids, experts, valid = t(64, dtype=i32), t(4, dtype=i32), t(2, dtype=i32)
    _opus_moe_stage1_a8w4_fwd_raw(
        t(8, 256, dtype=fp8), t(4, 512, 128, dtype=u8), t(8, 8, dtype=u8),
        t(4 * 512, 8, dtype=u8), None, ids, experts, valid,
        t(64, 256, dtype=fp8), t(64, 8, dtype=u8),
        2, 16, "opus_moe1_a8w4", 0, 0, 7.0, 1.0, 1.0,
    )
    _opus_moe_stage2_a8w4_decode_fwd_raw(
        t(64, 256, dtype=fp8), t(4, 256, 128, dtype=u8), t(64, 8, dtype=u8),
        t(4 * 256, 8, dtype=u8), None, ids, experts, valid,
        t(8, 256, dtype=torch.bfloat16), 8, 2, 16, 0, 0,
    )

The FlyDSL MoE AOT replay in #6362 runs every tuned row through these ops: 0 failures with this fix, 43 without.

Dependencies

• No new third-party dependencies

Breaking Changes

None.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6310 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[Copilot]

🟢 Approval recommended

Both fake signatures now exactly match their corresponding registered raw operations without altering runtime computation.

0 open findings

What changed in this PR

Aligns Opus A8W4 MoE fake implementations with their registered raw-op signatures, fixing FakeTensor and torch.compile tracing.

Changes:

• Adds missing scalar parameters to stage 1 and stage 2 fake functions.
• Preserves existing fake outputs and eager behavior.

File	Description
aiter/​ops/​opus/​moe_stage1_a8w4.py	Completes the stage 1 fake signature.
aiter/​ops/​opus/​moe_stage2_a8w4.py	Completes the stage 2 decode fake signature.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[zhiding512]

Folded into #6362 at the author's request: that PR now targets main and carries this commit unchanged.