Skip to content

moe: harden W4A16 planning, tails, and capture resolution - #110

Open
voipmonitor wants to merge 22 commits into
masterfrom
fix/w4a16-planning-capture-current-20260801
Open

moe: harden W4A16 planning, tails, and capture resolution#110
voipmonitor wants to merge 22 commits into
masterfrom
fix/w4a16-planning-capture-current-20260801

Conversation

@voipmonitor

Copy link
Copy Markdown
Contributor

Summary

Combine two compatible W4A16 hardening series on current master:

  1. bound unpinned scratch planning to reachable launch geometries and repair generic native tail cases;
  2. reject unresolved launch resolution before compilation or persistent-cache/module loading during CUDA graph capture.

This supersedes #108 and #109. Those PRs were independently mergeable against master, but either merge order conflicted in their shared E2E test file. This PR preserves both commit series, resolves the single mechanical import conflict by retaining both imports, and adds one final repo-version Ruff formatting commit. Runtime changes merged without manual conflict resolution.

Planner and tail correctness

  • Size unpinned fused and EP route/scratch buffers only for block sizes reachable within max_tokens.
  • Account for packed-route and direct-top-k/TC-decode geometry independently.
  • Preserve explicit block pins and mapped full-rotation Trellis behavior.
  • Expose fused_moe.required_nbytes(Caps) without compiling a launch.
  • Mask stale direct-FC2 lanes and logical activation/weight/scale tails.
  • Preserve native ModelOpt E8M0 scale strides and valid non-64 shapes.
  • Choose FC1 chunk counts that partition logical intermediates into whole 16-value blocks.

For the documented 64-token planning case, the non-full W4A16 arena falls from 94.94 MiB to 30.03 MiB per rank. Aligned production shapes retain their established launch geometry.

Capture contract

  • Existing in-memory launch hits remain legal during capture.
  • Unresolved launches fail before disk-cache lookup, filesystem/module load, or compilation.
  • Capture is detected on the explicitly selected stream, including a non-current stream.
  • Planner prewarm, route packing, slicing, and runtime resolution share one capacity/key derivation.
  • Eager planning and prewarm remain legal outside runtime launch resolution.

Verification

Combined focused suite on the r16 dependency image:

89 passed, 259 GPU-gated skips

The suite covers compiler cache behavior, fused/EP planning, Trellis, route packing, replay, non-aligned native shapes, and W4A16 E2E cases. Ruff 0.15.12 lint and format, plus git diff --check, pass across all 16 changed Python files.

No merge is performed by this PR creation.

@coderabbitai

coderabbitai Bot commented Aug 1, 2026

Copy link
Copy Markdown
Contributor

Warning

Review limit reached

@voipmonitor, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 43 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: e3824d39-f579-437a-ae0b-5a95e6e96552

📥 Commits

Reviewing files that changed from the base of the PR and between b0976b7 and 0191725.

📒 Files selected for processing (16)
  • sparkinfer/_lib/compiler.py
  • sparkinfer/moe/_shared/kernels/dynamic.py
  • sparkinfer/moe/_shared/kernels/micro.py
  • sparkinfer/moe/_shared/kernels/w4a16/host.py
  • sparkinfer/moe/_shared/kernels/w4a16/kernel.py
  • sparkinfer/moe/_shared/kernels/w4a16/prepare.py
  • sparkinfer/moe/_shared/kernels/w4a16/route_pack.py
  • sparkinfer/moe/ep_moe/_impl.py
  • sparkinfer/moe/fused_moe/__init__.py
  • sparkinfer/moe/fused_moe/_impl.py
  • sparkinfer/moe/fused_moe/api.py
  • tests/_lib/test_compile_cache.py
  • tests/moe/test_ep_moe_api.py
  • tests/moe/test_fused_moe_planning.py
  • tests/moe/test_w4a16_e2e.py
  • tests/moe/test_w4a16_route_pack.py

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

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.

2 participants