[BBaditi]

Motivation & Summary

Builds on #33575 ([CUDA/CPU] Add INT4 and INT3 QMoE CPU/GPU offload support) with two decode-path optimizations (numerically equivalent to within ~1e-5 for the AVX2 cold-expert kernel, and bitwise identical for the zero-cold doorbell short-circuit) and one C++ test fixture fix:

1. AVX2 + FMA + F16C cold-expert INT4 kernel (qmoe_shared_layout.cc):
   On x86-64 hosts with AVX2+FMA+F16C that do not support AVX-512 (such as AMD EPYC Zen 2 on Azure NDv4 A100 VMs and Intel Haswell–Alder Lake), CpuHasAvx512() returns false and cold-expert INT4 dequant-GEMM and FP16/BF16 row conversions fall back to scalar C++ loops (46.36 ms/tok of CPU compute at 25% GPU residency on 48-layer Qwen3-30B-A3B). This PR adds runtime CPUID + XGETBV detection (CpuHasAvx2FmaF16c()) and an AVX2+FMA+F16C implementation for FP16/BF16 row conversions and Avx2Int4GemmRange (supporting both prepacked SM80 INT4 layout and raw INT4 layout, FP16 and BF16 scales/biases), plus an $O(1)$ lookup table in ReorderedPosition.

2. Zero-PCIe doorbell short-circuit (session.moe.offload.skip_zero_cold_doorbell):
   During partial-residency hybrid decode ($M \le 4$), when all $M \times k$ selected experts for a layer are already GPU-resident (cold_count == 0), the default doorbell path still copies the activation tile and expert indices to pinned host memory and wakes the CPU doorbell worker thread. When session.moe.offload.skip_zero_cold_doorbell="1" (or ORT_QMOE_SKIP_ZERO_COLD_DOORBELL=1) is enabled under static placement without cross-layer promotion, QMoEOffloadClassifyKernel counts cold selections via __syncthreads_count(is_cold) before issuing any host writes and, when cold_count == 0, advances doorbell_done_host and doorbell_req_host directly on the GPU (__threadfence_system()) so QMoEWaitAndAccumulateKernel proceeds without PCIe host staging or CPU worker wakeups.

3. INT3 CUDA test fixture update (moe_test.cc):
   Updates MoETest.QMoECudaInt3TiledLogHasOneCounterUpdate to match the SwiGLU QMoE contract introduced in [CUDA] Add INT4 and INT3 QMoE CPU/GPU offload support #33575 (activation_type = "swiglu", swiglu_fusion = 1, fc1 rows = 2 * inter_size, and unsigned INT3 zero-point code 4).

Key Changes

• onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout.{h,cc}:
  • Add CpuHasAvx2FmaF16c(), CompilerHasAvx2(), and C ABI qmoe_shared_cpu_has_avx2() / qmoe_shared_set_disable_avx512_for_testing().
  • Add Avx2HalfBitsToFloatRow, Avx2FloatToHalfBitsRow, Avx2BFloat16BitsToFloatRow, Avx2FloatToBFloat16BitsRow, Avx2RoundFp16InPlace, and Avx2Int4GemmRange.
  • Replace the linear scan in ReorderedPosition with a 32-entry constant lookup table.
• onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.{h,cu}, onnxruntime/contrib_ops/cuda/moe/qmoe_offload.{h,cc}, onnxruntime/core/framework/moe_offload_plan.h, include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h:
  • Add kOrtSessionOptionsMoeOffloadSkipZeroColdDoorbell (session.moe.offload.skip_zero_cold_doorbell, default "0").
  • Short-circuit PCIe host staging and CPU doorbell wakeups when cold_count == 0 and cross-layer promotion is inactive; retire pre-completed sequence numbers cleanly in DoorbellService.
  • Allocate optional QMoEOffloadTelemetry (ORT_QMOE_OFFLOAD_TELEMETRY=1) in device scratch memory and copy back once in ~QMoEOffload().
• onnxruntime/test/contrib_ops/moe_test.cc, onnxruntime/test/contrib_ops/qmoe_shared_layout_test.cc, onnxruntime/test/python/transformers/test_qmoe_shared_layout.py, onnxruntime/test/python/test_qmoe_offload_v4.py:
  • Fix MoETest.QMoECudaInt3TiledLogHasOneCounterUpdate.
  • Add C++ and Python unit tests covering AVX2+FMA+F16C vs. scalar numerical parity (max_abs_diff < 1e-5, including all 65,536 FP16 half patterns and partial-range vector interiors) and zero-cold doorbell bitwise parity, sequence counter transitions, and cross-layer promotion fallback.
• docs/contrib_ops/cuda/moe_qmoe.md:
  • Document AVX-512 -> AVX2+FMA+F16C -> scalar cold-expert dispatch, session.moe.offload.skip_zero_cold_doorbell, and environment variables (ORT_QMOE_SKIP_ZERO_COLD_DOORBELL, ORT_QMOE_OFFLOAD_TELEMETRY, ORT_QMOE_CPU_DISABLE_AVX512).

Verification & Benchmarks

Unit & Provider Tests

• onnxruntime_test_all --gtest_filter="*MoETest*:*QMoE*": 6 / 6 passed
• onnxruntime_provider_test --gtest_filter="QMoESharedLayoutTest.*:MoETest.*": 91 / 91 passed
• pytest onnxruntime/test/python/test_qmoe_offload_v4.py onnxruntime/test/python/transformers/test_qmoe_shared_layout.py onnxruntime/test/python/transformers/test_qmoe_cuda.py: 153 / 153 passed

End-to-End 48-Layer Qwen3-30B-A3B (QMoE INT4, block_size=32, weights_prepacked=1) Decode Latency (ORT_QMOE_OFFLOAD_TELEMETRY=0, 5 Trials mean ± stdev)

Measured on 1x NVIDIA A100-SXM4-80GB + AMD EPYC 7V12 (Zen 2, AVX2+FMA+F16C, no AVX-512), 5 independent trials of 10 warmup + 30 decode steps ($M = 1$, $k = 8$, 128 experts/layer) with ORT_QMOE_OFFLOAD_TELEMETRY=0 (and a separate ORT_QMOE_OFFLOAD_TELEMETRY=1 pass to record wakeup/short-circuit counts):

GPU Residency (Experts / Layer)	#33575 Baseline (Scalar Cold CPU, ms/tok)	AVX2 Cold Kernel (skip_zero_cold=0, ms/tok)	AVX2 + Zero-Cold Doorbell (skip_zero_cold=1, ms/tok)	Speedup vs. #33575 (skip=0 / skip=1)	Zero-Cold Short-Circuits / Token (skip=1)	CPU Doorbell Wakes / Token (skip=0 $\to$ skip=1)	Min Logit Cosine vs. Scalar	Top-1 Token Parity
32 / 128 (25%)	54.46 ms (18.4 tok/s)	21.60 ± 0.37 ms (46.3 tok/s)	21.89 ± 0.55 ms (45.7 tok/s)	2.52x / 2.49x	7.78 / 48	48.00 $\to$ 40.23 (-16.2%)	0.999561	30 / 30 (100%)
64 / 128 (50%)	15.52 ms (64.4 tok/s)	10.80 ± 0.10 ms (92.6 tok/s)	10.66 ± 0.23 ms (93.8 tok/s)	1.44x / 1.46x	36.65 / 48	48.00 $\to$ 11.35 (-76.4%)	0.999818	30 / 30 (100%)
96 / 128 (75%)	9.04 ms (110.6 tok/s)	8.30 ± 0.19 ms (120.4 tok/s)	8.56 ± 0.41 ms (116.8 tok/s)	1.09x / 1.06x	46.98 / 48	48.00 $\to$ 1.02 (-97.9%)	0.999883	30 / 30 (100%)
128 / 128 (100%)	8.31 ms (120.4 tok/s)	7.55 ± 0.06 ms (132.5 tok/s)	7.61 ± 0.26 ms (131.4 tok/s)	1.10x / 1.09x	—	0.00 $\to$ 0.00	1.000000	30 / 30 (100%)

(Note: skip_zero_cold_doorbell="1" is 100% bitwise identical (max_abs_diff = 0.0, cosine = 1.000000) to skip_zero_cold_doorbell="0" across all residencies. On an uncontended host, single-stream decode latency for skip=0 vs. skip=1 is within noise because GPU resident GEMVs overlap the zero-cold doorbell wakeup, while skip=1 eliminates the 4 KB/layer PCIe host staging writes and reduces CPU doorbell wakeups by 76.4% at 50% residency and 97.9% at 75% residency.)

[Copilot]

🟡 Changes recommended

The tests need reliable AVX2 selection and direct coverage of vectorized partial-range writes.

2 open issues

• Partial Int4 range tests miss AVX2 vectorized interiors · New
• AVX2 tests may run AVX-512 kernels instead · New

What changed in this PR

This PR extends QMoE CPU/GPU offload with faster INT4 CPU execution on AVX2 hosts and an opt-in shortcut for decode calls whose selected experts are all GPU-resident.

Changes:

• Add AVX2 INT4 kernels, FP16/BF16 conversions, and runtime capability detection.
• Add configurable zero-cold doorbell bypass, sequence retirement, and telemetry.
• Expand numerical and decode regression tests and correct the INT3 SwiGLU fixture.

File	Description
onnxruntime/​test/​python/​transformers/​test_qmoe_shared_layout.py	Add scalar/SIMD INT4 parity coverage.
onnxruntime/​test/​python/​test_qmoe_offload_v4.py	Test bypass parity, replay, transitions, and policy handling.
onnxruntime/​test/​contrib_ops/​qmoe_shared_layout_test.cc	Expand conversion, range, and expert parity tests.
onnxruntime/​test/​contrib_ops/​moe_test.cc	Correct the INT3 fixture for SwiGLU.
onnxruntime/​core/​framework/​moe_offload_plan.h	Parse and validate the bypass policy.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_offload.h	Add telemetry configuration state.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_offload.cc	Integrate bypass handling, sequence retirement, and telemetry.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.h	Declare telemetry and extended classification parameters.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.cu	Complete zero-cold requests without host payload staging.
onnxruntime/​contrib_ops/​cpu/​moe/​qmoe_shared_layout.h	Expose AVX2 capability queries.
onnxruntime/​contrib_ops/​cpu/​moe/​qmoe_shared_layout.cc	Implement AVX2 kernels, conversions, and dispatch.
include/​onnxruntime/​core/​session/​onnxruntime_session_options_config_keys.h	Declare the bypass session option.
docs/​contrib_ops/​cuda/​moe_qmoe.md	Document the option and its applicability.

_{Review effort: Balanced}

[Tianlei Wu (tianleiwu)]

Thanks for the PR and for the detailed benchmark. The AVX2 cold-expert path is a clear win for AVX2-only hosts (25% residency: 54.46 to 22.01 ms/token), and I did not find a blocking correctness issue in the CPU or CUDA changes. Verdict: COMMENT (draft, stacked on #33575). The main open item is justifying the opt-in zero-cold bypass with repeatable measurements.

What I checked at 3cf5336 (AMD EPYC 7V12, Zen 2: AVX2+FMA+F16C, no AVX-512)

• AVX2/F16C float to half, half to float and float to bf16 converters vs the scalar ones on all 2^32 float bit patterns and all 65,536 half values: bit-exact for every non-NaN input (NaN payloads differ, see the inline comment on the test).
• Avx2Int4GemmRange vs scalar: randomized fuzz (prepacked/raw weights, block-major/row-major scales, fp16/bf16 scales and bias, partial ranges, strides), about 19M outputs under ASAN+UBSAN: no out-of-bounds access or padding clobber, max |diff| 7.6e-6 (summation order only).
• ReorderedPosition table vs the old loop for k in [-200, 200): identical.
• python3 -m unittest test_qmoe_shared_layout on this AVX2-only host: 21 run, 5 skipped (AVX-512-only), OK.
• The CUDA classify/bypass/doorbell changes were reviewed by reading only; I did not run the GPU tests or the C++ gtests.

Suggestions

1. Benchmark: the bypass shows no measurable latency gain (details below).
2. qmoe_kernels.cu: telemetry read-modify-writes on pinned host memory sit in front of the done/req stores (inline).
3. qmoe_shared_layout.cc: getenv on every CpuHasAvx512() call (inline).
4. moe_qmoe.md: the AVX2 fallback and the new env vars are undocumented, and AVX-512-only wording is stale (inline).
5. qmoe_shared_layout_test.cc: the fp16 converter test covers 10 of 18 patterns (inline).

Nitpicks / questions
6. qmoe_offload.cc: the CPU-side zero-cold shortcut in ServiceDoorbell looks unreachable (inline).
7. moe_offload_plan.h: getenv is called twice (inline).
8. PR description accuracy (see the last section).

Benchmark (skip=0 vs skip=1, ms/token, from the PR description)

Residency	skip=0	skip=1	diff	Doorbell wakes/token with skip=1
25%	22.01	22.07	+0.06	40.23 / 48
50%	10.64	10.67	+0.03	11.35 / 48 (-76%)
75%	8.76	9.04	+0.28	1.02 / 48 (-98%)
100%	7.65	7.55	-0.10	0 (no doorbell traffic)

• Wakeups drop a lot, but latency does not move. The 100% row runs the same path in both columns, so its 0.10 ms (1.3%) gap is a rough single-sample noise floor: the 25% and 50% diffs are inside it, and the 75% diff (+0.28 ms) is larger than it, in the wrong direction.
• The 75% skip=1 cell (9.04 ms, 110.6 tok/s) is identical to the #33575 baseline cell (9.04 ms), so it may be a copy slip.
• A plausible reason for no gain: the CPU doorbell service already overlaps the GPU resident GEMVs, so a zero-cold wake is hidden. I have not verified that.
• Please add repeated runs (mean +/- stdev) and a telemetry-off A/B (see the inline comment on qmoe_kernels.cu), or reframe the bypass as a PCIe / CPU-wakeup reduction (for example under CPU contention) if latency is not the goal. With the extra GPU to host handshake path and the RetirePending accounting, the justification matters.

PR description accuracy

• The speedup column uses skip=0 for the 25-75% rows (2.47x = 54.46 / 22.01) but skip=1 for the 100% row (8.31 / 7.55 = 1.10x; with skip=0 it is 1.09x).
• 'Exact' optimizations: the AVX2 path is not bit-identical to scalar (different FP32 summation order; the table shows a min logit cosine of 0.99956 at 25% residency), so 'numerically equivalent to about 1e-5' is more accurate. The bypass itself is bitwise identical per the table.
• onnxruntime/core/framework/session_state.cc is listed as modified but is not in the diff.

The two open bot-review threads on qmoe_shared_layout_test.cc (AVX2 vector interior not exercised; AVX2-named tests may run AVX-512 code) look addressed at this head, so I did not repost them.

[BBaditi]

Thanks for the PR and for the detailed benchmark. The AVX2 cold-expert path is a clear win for AVX2-only hosts (25% residency: 54.46 to 22.01 ms/token), and I did not find a blocking correctness issue in the CPU or CUDA changes. Verdict: COMMENT (draft, stacked on #33575). The main open item is justifying the opt-in zero-cold bypass with repeatable measurements.

What I checked at 3cf5336 (AMD EPYC 7V12, Zen 2: AVX2+FMA+F16C, no AVX-512)

• AVX2/F16C float to half, half to float and float to bf16 converters vs the scalar ones on all 2^32 float bit patterns and all 65,536 half values: bit-exact for every non-NaN input (NaN payloads differ, see the inline comment on the test).
• Avx2Int4GemmRange vs scalar: randomized fuzz (prepacked/raw weights, block-major/row-major scales, fp16/bf16 scales and bias, partial ranges, strides), about 19M outputs under ASAN+UBSAN: no out-of-bounds access or padding clobber, max |diff| 7.6e-6 (summation order only).
• ReorderedPosition table vs the old loop for k in [-200, 200): identical.
• python3 -m unittest test_qmoe_shared_layout on this AVX2-only host: 21 run, 5 skipped (AVX-512-only), OK.
• The CUDA classify/bypass/doorbell changes were reviewed by reading only; I did not run the GPU tests or the C++ gtests.

Suggestions

1. Benchmark: the bypass shows no measurable latency gain (details below).
2. qmoe_kernels.cu: telemetry read-modify-writes on pinned host memory sit in front of the done/req stores (inline).
3. qmoe_shared_layout.cc: getenv on every CpuHasAvx512() call (inline).
4. moe_qmoe.md: the AVX2 fallback and the new env vars are undocumented, and AVX-512-only wording is stale (inline).
5. qmoe_shared_layout_test.cc: the fp16 converter test covers 10 of 18 patterns (inline).

Nitpicks / questions 6. qmoe_offload.cc: the CPU-side zero-cold shortcut in ServiceDoorbell looks unreachable (inline). 7. moe_offload_plan.h: getenv is called twice (inline). 8. PR description accuracy (see the last section).

Benchmark (skip=0 vs skip=1, ms/token, from the PR description)

Residency skip=0 skip=1 diff Doorbell wakes/token with skip=1
25% 22.01 22.07 +0.06 40.23 / 48
50% 10.64 10.67 +0.03 11.35 / 48 (-76%)
75% 8.76 9.04 +0.28 1.02 / 48 (-98%)
100% 7.65 7.55 -0.10 0 (no doorbell traffic)

• Wakeups drop a lot, but latency does not move. The 100% row runs the same path in both columns, so its 0.10 ms (1.3%) gap is a rough single-sample noise floor: the 25% and 50% diffs are inside it, and the 75% diff (+0.28 ms) is larger than it, in the wrong direction.
• The 75% skip=1 cell (9.04 ms, 110.6 tok/s) is identical to the [CUDA] Add INT4 and INT3 QMoE CPU/GPU offload support #33575 baseline cell (9.04 ms), so it may be a copy slip.
• A plausible reason for no gain: the CPU doorbell service already overlaps the GPU resident GEMVs, so a zero-cold wake is hidden. I have not verified that.
• Please add repeated runs (mean +/- stdev) and a telemetry-off A/B (see the inline comment on qmoe_kernels.cu), or reframe the bypass as a PCIe / CPU-wakeup reduction (for example under CPU contention) if latency is not the goal. With the extra GPU to host handshake path and the RetirePending accounting, the justification matters.

PR description accuracy

• The speedup column uses skip=0 for the 25-75% rows (2.47x = 54.46 / 22.01) but skip=1 for the 100% row (8.31 / 7.55 = 1.10x; with skip=0 it is 1.09x).
• 'Exact' optimizations: the AVX2 path is not bit-identical to scalar (different FP32 summation order; the table shows a min logit cosine of 0.99956 at 25% residency), so 'numerically equivalent to about 1e-5' is more accurate. The bypass itself is bitwise identical per the table.
• onnxruntime/core/framework/session_state.cc is listed as modified but is not in the diff.

The two open bot-review threads on qmoe_shared_layout_test.cc (AVX2 vector interior not exercised; AVX2-named tests may run AVX-512 code) look addressed at this head, so I did not repost them.

Thanks for the review and fuzzing tests Tianlei Wu (@tianleiwu)! Addressed all 6 inline comments in e7e20f7 and updated the PR description with 5-trial numbers (mean ± stdev) with ORT_QMOE_OFFLOAD_TELEMETRY=0. With telemetry off on an idle GPU, single-stream latency for skip=0 and skip=1 is within noise (10.80 ± 0.10 ms vs 10.66 ± 0.23 ms at 50% residency), while skip=1 still cuts CPU doorbell wakeups by 76.4% at 50% residency and 97.9% at 75% residency.

[Tianlei Wu (tianleiwu)]

Hi BBaditi — we just updated the base QMoE offload PR stack (#33450 $\to$ #33451 $\to$ #33452 $\to$ #33575 $\to$ #33622 $\to$ #33629) on top of the latest main to align with the shared-layout CPU kernel refinement plan:

1. Per-ISA SIMD translation units in onnxruntime_mlas:
   • SIMD kernels are now split into dedicated translation units (onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout_int4_avx512.cc, qmoe_shared_layout_int4_avx2.cc, qmoe_shared_layout_int3_avx512.cc, qmoe_shared_layout_int3_avx2.cc, qmoe_shared_layout_fp4_avx512.cc) compiled with explicit ISA flags (MLAS_QMOE_AVX2_FLAGS = -mavx2 -mfma -mf16c / /arch:AVX2) in cmake/onnxruntime_mlas.cmake.
   • We included your AVX2 INT4 kernel in qmoe_shared_layout_int4_avx2.cc (in #33450) and the skip_zero_cold_doorbell short-circuit (#33629) with Co-authored-by: Bhuvaneswar Baditi <bbaditi@microsoft.com>.
2. Direct CPU execution & 100% byte-preserving FileExpertStore:
   • QMoECPU now executes directly on shared-layout weights without online prepacking or UnpackSm80Int4Weights, and FileExpertStore no longer transposes [N, K/32] scales on load (scales_are_block_major = false on CPU).

When you rebase this branch onto the updated tlwu/20261024/qmoe-offload-int4-int3 (#33575), any remaining changes from this PR can be scoped cleanly on top of the modular qmoe_shared_layout_int4_avx2.cc file and byte-preserving FileExpertStore.

[BBaditi]

Hi BBaditi — we just updated the base QMoE offload PR stack (#33450 → #33451 → #33452 → #33575 → #33622 → #33629) on top of the latest main to align with the shared-layout CPU kernel refinement plan:

1. Per-ISA SIMD translation units in onnxruntime_mlas:

   • SIMD kernels are now split into dedicated translation units (onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout_int4_avx512.cc, qmoe_shared_layout_int4_avx2.cc, qmoe_shared_layout_int3_avx512.cc, qmoe_shared_layout_int3_avx2.cc, qmoe_shared_layout_fp4_avx512.cc) compiled with explicit ISA flags (MLAS_QMOE_AVX2_FLAGS = -mavx2 -mfma -mf16c / /arch:AVX2) in cmake/onnxruntime_mlas.cmake.
   • We included your AVX2 INT4 kernel in qmoe_shared_layout_int4_avx2.cc (in #33450) and the skip_zero_cold_doorbell short-circuit (#33629) with Co-authored-by: Bhuvaneswar Baditi <bbaditi@microsoft.com>.

2. Direct CPU execution & 100% byte-preserving FileExpertStore:

   • QMoECPU now executes directly on shared-layout weights without online prepacking or UnpackSm80Int4Weights, and FileExpertStore no longer transposes [N, K/32] scales on load (scales_are_block_major = false on CPU).

When you rebase this branch onto the updated tlwu/20261024/qmoe-offload-int4-int3 (#33575), any remaining changes from this PR can be scoped cleanly on top of the modular qmoe_shared_layout_int4_avx2.cc file and byte-preserving FileExpertStore.

Rebased onto  tlwu/20261026/qmoe-offload-avx2-int3-s8-opt  (#33629) and scoped this branch down to the remaining  ORT_QMOE_CPU_DISABLE_AVX512  test toggle, single-source  skip_zero_cold_doorbell  cleanup, and unit test/doc updates.

[mirounga]

Thanks for the AVX2 path. The kernel work looks solid, but I have two questions about the offload change, and a few things in the tests and dispatch I'd like to tighten before this goes in.

Correctness

1. Zero-cold invocations no longer mark the CPU stage done (qmoe_offload.cc, around L4396, ServiceDoorbell).
   The removed zero-cold short-circuit was the only path that called MarkCpuDone() for invocations where every selected expert is GPU-resident. Those invocations now retire in the service loop without it, so the persistent invocation_rec stays cpu_done == false. Is anything downstream (completion wait, end-of-run validation, the next invocation's reuse of the record) still checking that flag? If so, this looks like a hang or a spurious failure. If the flag is now dead for this case, could you remove or document it?

2. The bypass depends only on the host flag (qmoe_kernels.cu, around L176).
   On the zero-cold bypass the kernel still writes doorbell_req_host, and only the host-side v4_policy_.skip_zero_cold_doorbell check stops the CPU from treating it as a real request. If the value seen at launch and the host's value ever differ (CUDA graph replay keeps the capture-time value, for example), the host processes stale routing and inputs. Could the kernel skip the doorbell write on the bypass, or send a "zero cold" marker the host checks, so the two sides can't disagree?

3. The BF16 half of the tail-shape parity test checks nothing (test_qmoe_shared_layout.py, around L373, test_int4_avx2_matches_scalar_across_tail_shapes).
   With is_bf16=1 the test still passes FP16 bit patterns for scales and biases. Read as BF16, those are about 1e-26 to 1e-16, so both the AVX2 and scalar outputs are near zero and the comparison passes whatever the BF16 decode does. Please generate the scale and bias bits with the matching dtype (torch.bfloat16(...).view(torch.int16) or similar).

4. ORT_QMOE_CPU_DISABLE_AVX512 drops NVFP4/MXFP4 to scalar, not AVX2 (qmoe_shared_layout.cc, around L528).
   Nvfp4GemmRange and Mxfp4GemmRange have AVX-512 and scalar branches only, so with the switch set (or on an AVX2-only host) these formats run the scalar loop. Either add the AVX2 branch or narrow the documentation of the switch to INT4.

Performance and design

5. CpuHasAvx512() re-dispatches on every call (qmoe_shared_layout.cc, around L524).
   It now does an atomic load of the test override inside per-row and per-GEMM hot paths, and the chosen path can change mid-run if the override flips. Selecting the kernels once at init (a small function-pointer table that the test override rebuilds) would avoid both.

6. The test-only AVX-512 switch ships in production (qmoe_shared_layout.h, around L446).
   ScopedDisableAvx512ForTesting and the exported C symbol are part of the production header and the library ABI, as a process-global switch: any caller can downgrade every QMoE session in the process. Could this move to a test-only header, or be compiled only for test builds?

Tests

7. No column tail in the "tail shapes" test (test_qmoe_shared_layout.py, around L367).
   hidden=128 and inter=64 make every N a multiple of 8, so only the row count varies. Please add shapes with N % 8 != 0 (and ideally N % 16 != 0) so the AVX2 column tail is exercised.

8. The AVX2 re-run duplicates work (qmoe_shared_layout_test.cc, around L1243, CanonicalInt4ExpertViewMatchesReference).
   On hosts without AVX-512 the main run already takes the AVX2 path, so the extra block runs it twice. A loop over the available dispatch levels would cover both cases without the duplication.

9. Copy-pasted 18-argument calls (test_qmoe_shared_layout.py, around L409).
   The two qmoe_int4_expert_contribution calls differ in two arguments. A small helper would keep them from drifting apart.

Items 1-4 are the ones I'd like resolved before approving. 5-9 are suggestions.

[BBaditi]

Thanks for the AVX2 path. The kernel work looks solid, but I have two questions about the offload change, and a few things in the tests and dispatch I'd like to tighten before this goes in.

Correctness

1. Zero-cold invocations no longer mark the CPU stage done (qmoe_offload.cc, around L4396, ServiceDoorbell).
   The removed zero-cold short-circuit was the only path that called MarkCpuDone() for invocations where every selected expert is GPU-resident. Those invocations now retire in the service loop without it, so the persistent invocation_rec stays cpu_done == false. Is anything downstream (completion wait, end-of-run validation, the next invocation's reuse of the record) still checking that flag? If so, this looks like a hang or a spurious failure. If the flag is now dead for this case, could you remove or document it?
2. The bypass depends only on the host flag (qmoe_kernels.cu, around L176).
   On the zero-cold bypass the kernel still writes doorbell_req_host, and only the host-side v4_policy_.skip_zero_cold_doorbell check stops the CPU from treating it as a real request. If the value seen at launch and the host's value ever differ (CUDA graph replay keeps the capture-time value, for example), the host processes stale routing and inputs. Could the kernel skip the doorbell write on the bypass, or send a "zero cold" marker the host checks, so the two sides can't disagree?
3. The BF16 half of the tail-shape parity test checks nothing (test_qmoe_shared_layout.py, around L373, test_int4_avx2_matches_scalar_across_tail_shapes).
   With is_bf16=1 the test still passes FP16 bit patterns for scales and biases. Read as BF16, those are about 1e-26 to 1e-16, so both the AVX2 and scalar outputs are near zero and the comparison passes whatever the BF16 decode does. Please generate the scale and bias bits with the matching dtype (torch.bfloat16(...).view(torch.int16) or similar).
4. ORT_QMOE_CPU_DISABLE_AVX512 drops NVFP4/MXFP4 to scalar, not AVX2 (qmoe_shared_layout.cc, around L528).
   Nvfp4GemmRange and Mxfp4GemmRange have AVX-512 and scalar branches only, so with the switch set (or on an AVX2-only host) these formats run the scalar loop. Either add the AVX2 branch or narrow the documentation of the switch to INT4.

Performance and design

5. CpuHasAvx512() re-dispatches on every call (qmoe_shared_layout.cc, around L524).
   It now does an atomic load of the test override inside per-row and per-GEMM hot paths, and the chosen path can change mid-run if the override flips. Selecting the kernels once at init (a small function-pointer table that the test override rebuilds) would avoid both.
6. The test-only AVX-512 switch ships in production (qmoe_shared_layout.h, around L446).
   ScopedDisableAvx512ForTesting and the exported C symbol are part of the production header and the library ABI, as a process-global switch: any caller can downgrade every QMoE session in the process. Could this move to a test-only header, or be compiled only for test builds?

Tests

7. No column tail in the "tail shapes" test (test_qmoe_shared_layout.py, around L367).
   hidden=128 and inter=64 make every N a multiple of 8, so only the row count varies. Please add shapes with N % 8 != 0 (and ideally N % 16 != 0) so the AVX2 column tail is exercised.
8. The AVX2 re-run duplicates work (qmoe_shared_layout_test.cc, around L1243, CanonicalInt4ExpertViewMatchesReference).
   On hosts without AVX-512 the main run already takes the AVX2 path, so the extra block runs it twice. A loop over the available dispatch levels would cover both cases without the duplication.
9. Copy-pasted 18-argument calls (test_qmoe_shared_layout.py, around L409).
   The two qmoe_int4_expert_contribution calls differ in two arguments. A small helper would keep them from drifting apart.

Items 1-4 are the ones I'd like resolved before approving. 5-9 are suggestions.

Thanks mirounga, I pushed a fix for all of these.

In qmoe_offload.cc, DoorbellService and ServiceDoorbell now check doorbell_done_host >= req directly without re-checking skip_zero_cold_doorbell on the host, and call invocation_rec->MarkCpuDone() when retiring zero-cold requests.

In qmoe_shared_layout.cc and moe_qmoe.md, CpuHasAvx512 now caches ORT_QMOE_CPU_DISABLE_AVX512 and CPUID on the first call before checking the test toggle, and I updated the comments and docs to note that NVFP4/MXFP4 fall back to scalar while INT4/INT3 and FP16/BF16 use AVX2. I also added a short note in qmoe_shared_layout.h on why the test toggle is exposed for the C++ and ctypes Python tests.

For the tests, I updated test_qmoe_shared_layout.py to pass real BF16 bit patterns when is_bf16=1, check that the output is non-zero, and use a small helper for the two contribution calls. Since the full-expert helper requires hidden and inter to be multiples of 32 along K, I expanded Avx2PartialInt4RangesPreserveOtherColumns in qmoe_shared_layout_test.cc up to 25 columns to cover the 16 + 8 + tail path, and guarded the AVX-512 disable block with CpuHasAvx512.

[Tianlei Wu (tianleiwu)]

Re-reviewed the current head and resolved my six addressed threads: cached environment selection/test toggle, device-side telemetry, redundant CPU shortcut removal, single environment lookup, SIMD documentation, and expanded conversion coverage. The AVX2 tests now force AVX2 and include vectorized interior plus scalar-tail checks; the completion paths retire pending requests and mark CPU completion consistently. No additional actionable findings.

Current-head CI is green across the complete rollup (one neutral check); I did not repeat runtime validation locally. This looks ready for manual approval. No approval submitted.