[AMD] [MORI-EP] Skip LocalExpertCount kernel in decode graph when not recording

[rbrugaro-amd]

Motivation

_should_record_expert_distribution() in the mori EP token dispatcher returned True
whenever the CUDA stream was capturing. This baked mori's LocalExpertCountKernel (and
its buffer memsets) into the decode CUDA graph for every run — including normal
serving with no expert-distribution recorder configured. The kernel then replayed on
every decode step, producing a local_expert_count that nothing consumes when recording
is off (the recorder hook is inactive), i.e. dead work in the decode hot path.

This affects any MoE model served via the mori EP all-to-all backend with CUDA graphs
enabled.

Modifications

Gate the capture-time path on whether an expert-distribution recorder is actually
configured (non-Noop) instead of on is_current_stream_capturing() alone:

• No recorder configured (common serving path) + capturing → skip the count kernel.
• A configured (non-Noop) recorder, even before start_record() → still bake the
  machinery in at capture time, so start_record() works when called after graph capture
  (the recorder's own _on_hook continues to gate the actual gather on
  recording/capturing).
• Actively recording → unchanged.

Single-file change: python/sglang/srt/layers/moe/token_dispatcher/moriep.py (+19/−2).

Because the two recording paths (configured-but-not-yet-recording, and actively-recording)
are left untouched, there is no expert-distribution recording regression: an EPLB
stat_approx record run (expert_distribution_recorder_mode=stat_approx +
/start_expert_distribution_record → traffic → /dump_expert_distribution_record) still
produces non-empty per-layer expert counts, and its gsm8k accuracy is unchanged (0.975).

Accuracy Tests

8×MI355X, DeepSeek-R1 MXFP4, mori EP (tp8/ep8/dp8, DP attention, aiter, full decode CUDA
graph). Post-patch, no recorder:

python -m sglang.launch_server --model-path <DeepSeek-R1> \
  --tp-size 8 --ep-size 8 --dp-size 8 --enable-dp-attention --moe-a2a-backend mori \
  --trust-remote-code --load-balance-method round_robin --moe-dense-tp-size 1 \
  --enable-dp-lm-head --mem-fraction-static 0.7 --chunked-prefill-size 1024 \
  --max-running-requests 128 --context-length 4096 --max-total-tokens 32768 \
  --attention-backend aiter --cuda-graph-max-bs-decode 32 --deepep-mode normal
# (env: SGLANG_USE_AITER=1 SGLANG_MORI_DISPATCH_DTYPE=bf16
#       SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128 MORI_SHMEM_MODE=ISOLATION)

python -m sglang.test.few_shot_gsm8k --num-shots 5 --num-questions 200 --parallel 128 \
  --host http://127.0.0.1 --port 30000

Accuracy: 0.965
Invalid:  0.000

Speed Tests and Profiling

Decode-dominant bench_serving A/B (baseline vs patched), same server config, same
settings, back-to-back, random input 256 / output 1024, seed 42, 8 warmup,
num-prompts = 8 × concurrency:

python -m sglang.bench_serving --backend sglang --host 127.0.0.1 --port 30000 \
  --dataset-name random --random-input-len 256 --random-output-len 1024 \
  --random-range-ratio 1.0 --num-prompts $((C*8)) --max-concurrency $C \
  --warmup-requests 8 --seed 42          # C in {16, 32, 64}

concurrency	Median TPOT (ms) baseline → patched	Output tput (tok/s) baseline → patched
16	19.78 → 19.56 (−1.1%)	793.79 → 802.86 (+1.1%)
32	21.20 → 20.98 (−1.0%)	1473.77 → 1490.12 (+1.1%)
64	23.25 → 22.85 (−1.7%)	2631.25 → 2683.34 (+2.0%)

Median ITL moves the same way (−1.1%…−1.2%). ~1–2% TPOT / output-throughput improvement,
increasing with decode concurrency — consistent in direction at every point, as expected
from removing one small per-MoE-layer-per-decode-step kernel from the decode graph. Decode
batch was capped at 32 (--cuda-graph-max-bs-decode 32); the uplift ceiling is likely a
bit higher at larger decode batches, where the kernel would have replayed more per step.

---

CI States

Latest PR Test (Base): ❌ Run #28825048244
Latest PR Test (Extra): ❌ Run #28825048081

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!