Skip to content

esimd: fix INT4 router residual race for all shapes - #675

Merged
gc-fu merged 1 commit into
mainfrom
fix/issue-671-int4-router-race
Sep 4, 2026
Merged

esimd: fix INT4 router residual race for all shapes#675
gc-fu merged 1 commit into
mainfrom
fix/issue-671-int4-router-race

Conversation

@gc-fu

@gc-fu gc-fu commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

Summary

  • run residual-add/RMSNorm once before launching the INT4 router GEMV
  • reuse the existing optimized INT4 GEMV dispatcher, with ordered device submissions and no host synchronization
  • add a deterministic large-grid regression that compares every router row with standalone GEMV
  • make the reused GEMV input pointers const-correct

Fixes #671.

Root cause

The fused kernel launches one work-group per router output row. Every work-group reads residual, while only the n == 0 work-group updates it in place. SYCL does not provide a global barrier across those work-groups, so later-scheduled groups can read the updated residual and effectively add hidden twice. This perturbs router logits and eventually causes the repeated-token/NaN failure seen with sym_int4 and aggressive but valid sampling parameters.

PR #399 addresses the same race only for N > 512. That threshold is not a correctness guarantee: issue #671 reproduces in Qwen3.6-35B-A3B with a 256-expert router under the real TP=2 workload. This change removes the cross-work-group dependency for every shape.

Performance

The fix adds one device kernel submission but no host synchronization. RMSNorm is computed once, then the existing tuned INT4 GEMV implementation is used.

On card 2, the existing 1,000-iteration microbenchmark at N=128, K=2048 measured 11.8 us before and 13.1-13.6 us after: +1.3 to +1.8 us (about 11-15%). The absolute per-layer increase corresponds to well below 1% of observed end-to-end decode time. The final model run sustained about 62 tokens/s after warmup; no end-to-end regression was observed.

Validation

Built and reinstalled the wheel using the workspace oneAPI flow:

source /opt/intel/oneapi/setvars.sh --force
TORCH_XPU_ARCH_LIST=bmg-g21 MAX_JOBS=1 python setup.py bdist_wheel
uv pip install --python /opt/venv/bin/python --no-deps --force-reinstall dist/custom_esimd_kernels_vllm-0.1.0-cp312-cp312-linux_x86_64.whl
  • ZE_AFFINITY_MASK=2 pytest -q tests/test_resadd_norm_gemv_int4.py: 11 passed
  • regression before fix: 2,036 / 4,096 router rows mismatched; max absolute error 48.40625
  • Qwen3.6-35B-A3B, sym_int4, FP16, TP=2 on cards 2/3, exact issue sampling parameters: seeds 0-4 each generated 1,024 tokens with zero repeated ! runs
  • longer seed-0 run: stopped normally at 3,189 tokens with zero !
  • final rebuilt wheel seed-0 rerun: 1,024 tokens, zero !, 19.9 seconds

Disclosure

This change was developed with OpenAI Codex assistance. The submitter is responsible for reviewing every changed line and the validation evidence before merge.

Split residual-add/RMSNorm from router GEMV so every work-group consumes the same normalized vector. Reuse the existing optimized INT4 GEMV dispatcher to limit the performance cost and add a high-N regression test.

Co-authored-by: OpenAI Codex <noreply@openai.com>
@gc-fu
gc-fu merged commit ede4320 into main Sep 4, 2026
@gc-fu
gc-fu deleted the fix/issue-671-int4-router-race branch September 4, 2026 05:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Qwen3.6-35B-A3B sym_int4 on two b70 output ! with 0.26.0-b1 version

1 participant