[Bug Fix]fix flyDSL MoE sorting graph capture break with dp attn + ep

[amd-weisun]

Fix the cuda graph capture issue when running GPT-OSS with dp attn + ep

Now it can pass

AITER_USE_FLYDSL_MOE_SORTING=1 python -m atom.entrypoints.openai_server --model openai/gpt-oss-120b --server-port 8000 --kv_cache_dtype fp8 -tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all --gpu-memory-utilization 0.5

run accuracy test:

lm_eval --model local-chat-completions --apply_chat_template \ --model_args model=openai/gpt-oss-120b,base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False \ --tasks gsm8k \ --num_fewshot 3

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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 4375 --add-label <label>

[Copilot]

Pull request overview

Fixes torch CUDA-graph capture/replay correctness for the FlyDSL moe_sorting path when the per-step local token count is provided as a device tensor (e.g., decode with DP-attn + expert-parallel), by ensuring the token count is read on-device rather than via a host .item().

Changes:

• Update the FlyDSL MoE sorting kernels to optionally load num_local_tokens from a device tensor and use it to bound mesh scans (so replay can reflect updated counts).
• Extend op_tests/test_moe_sorting.py to pass accumulate=..., add a FlyDSL CUDA-graph capture/replay regression test, and add a decode-style graph replay latency benchmark CLI.
• Add a dummy cached device tensor for the local_tokens_tensor argument when dynamic local-tokens are not provided.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
op_tests/test_moe_sorting.py	Adds accumulate knob passthrough and introduces CUDA-graph capture/replay regression + replay-latency benchmark paths.
aiter/ops/flydsl/kernels/moe_sorting_kernel.py	Makes FlyDSL MoE sorting capture-safe by loading local token count on-device and bounding scan work to the dynamic count.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.