moe: bound W4A16 scratch and repair tail correctness - #109
Closed
voipmonitor wants to merge 16 commits into
Closed
Conversation
Contributor
📝 WalkthroughWalkthroughThe PR updates ModelOpt W4A16 and E8M0 scale layouts, adds logical-width safety checks, replaces fixed route-size planning with capacity-aware sizing, and exposes ChangesMoE kernel and scale layout safety
Capacity-aware route and scratch planning
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sparkinfer/moe/_shared/kernels/micro.py`:
- Around line 2485-2489: Update the w1_scales layout declarations around the
E8M0/W4A16 setup to use the same padded column extent as the kernel’s scale
stride, `_align_up(cfg.two_n, 64)`, rather than raw `cfg.two_n`. Keep the layout
and preparation allocation consistent for non-64-aligned widths, while
preserving the existing E8M0 and W4A16 dimension-specific behavior.
In `@sparkinfer/moe/_shared/kernels/w4a16/prepare.py`:
- Around line 1186-1191: Align padded_w13_scale_n in the preparation path with
W4A16GemmKernel.scale_size_n: apply the 64-row E8M0 padding only when
scale_format_e8m0_k32 and has_n_tile_tail, matching the kernel’s scale_n_groups
and stride calculation. Preserve the unpadded row count for tile-aligned
dimensions without a W2 compact K tail, including multi-chunk FC1.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 115b57fb-6662-461c-af6d-8bfe2c4630cb
📒 Files selected for processing (11)
sparkinfer/moe/_shared/kernels/micro.pysparkinfer/moe/_shared/kernels/w4a16/host.pysparkinfer/moe/_shared/kernels/w4a16/prepare.pysparkinfer/moe/ep_moe/_impl.pysparkinfer/moe/fused_moe/__init__.pysparkinfer/moe/fused_moe/_impl.pysparkinfer/moe/fused_moe/api.pytests/moe/test_ep_moe_api.pytests/moe/test_fused_moe_planning.pytests/moe/test_w4a16_e2e.pytests/moe/test_w4a16_route_pack.py
Comment on lines
+2485
to
+2489
| e8m0_w1_n_cols = _align_up(cfg.two_n, 64) | ||
| ebase_sf_packed = Int64(eid) * Int64((cfg.k_dim // 32) * e8m0_w1_n_cols) | ||
| ebase_sf_packed_e4m3 = Int64(eid) * Int64( | ||
| (cfg.k_dim // 16) * _align_up(cfg.two_n, 64) | ||
| ) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Compare declared micro scale tensor extents against the padded FC1 scale stride.
set -euo pipefail
# Locate the padded-stride helper and every micro scale tensor declaration.
rg -n -C4 '_align_up\(\s*cfg\.two_n\s*,\s*64\s*\)' sparkinfer/moe/_shared/kernels/micro.py
rg -n -C6 'w1_scales\s*=\s*cute\.make_tensor' sparkinfer/moe/_shared/kernels/micro.py
# Confirm the preparation side actually allocates the padded width.
rg -n -C6 '_e8m0_logical_tail_scale_n|_PACKED_TILE_N_SIZE' sparkinfer/moe/_shared/kernels/w4a16/prepare.pyRepository: local-inference-lab/sparkinfer
Length of output: 7842
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the full FC1 padded stride definition and surrounding preparation validation/production.
sed -n '2460,2505p;5020,5050p' sparkinfer/moe/_shared/kernels/micro.py
sed -n '1,80p;1140,1205p;1205,1255p' sparkinfer/moe/_shared/kernels/w4a16/prepare.py
# Check whether declared scale extents are used to create layouts/tensors accessed later by the patched code.
python3 - <<'PY'
from pathlib import Path
p=Path('sparkinfer/moe/_shared/kernels/micro.py')
s=p.read_text()
for needle in ['w1_sf_rows * cfg.w1_sf_cols', 'cfg.two_n * cfg.k_half', 'cfg.weight_E * (cfg.k_dim // 32) * cfg.two_n', 'cfg.weight_E * (cfg.k_dim // 16) * cfg.two_n']:
print(needle, s.count(needle))
PYRepository: local-inference-lab/sparkinfer
Length of output: 11012
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the full FC1 padded stride definition and surrounding preparation validation/production.
sed -n '2460,2505p;5020,5050p' sparkinfer/moe/_shared/kernels/micro.py
sed -n '1,80p;1140,1205p;1205,1255p' sparkinfer/moe/_shared/kernels/w4a16/prepare.py
# Check whether declared scale extents are used to create layouts/tensors accessed later by the patched code.
python3 - <<'PY'
from pathlib import Path
p=Path('sparkinfer/moe/_shared/kernels/micro.py')
s=p.read_text()
for needle in ['w1_sf_rows * cfg.w1_sf_cols', 'cfg.two_n * cfg.k_half', 'cfg.weight_E * (cfg.k_dim // 32) * cfg.two_n', 'cfg.weight_E * ( cfg.k_dim // 16) * cfg.two_n']:
print(needle, s.count(needle))
PYRepository: local-inference-lab/sparkinfer
Length of output: 11013
Align w1_scales layout extents with the padded scale stride.
The kernel advances the scale base by _align_up(cfg.two_n, 64), but w1_scales layouts at lines 5036-5043 still use unpadded cfg.two_n. Preparation allocates the padded width for E8M0/W4A16, so the raw pointer reads are in bounds. The layout extent only matches if cfg.two_n is a multiple of 64; either align the declaration for non-64-aligned sizes or keep the layout size equal to the prepared padded column count.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sparkinfer/moe/_shared/kernels/micro.py` around lines 2485 - 2489, Update the
w1_scales layout declarations around the E8M0/W4A16 setup to use the same padded
column extent as the kernel’s scale stride, `_align_up(cfg.two_n, 64)`, rather
than raw `cfg.two_n`. Keep the layout and preparation allocation consistent for
non-64-aligned widths, while preserving the existing E8M0 and W4A16
dimension-specific behavior.
Comment on lines
+1186
to
+1191
| # Packed E8M0 scale rows are permuted in 64-row blocks independently of | ||
| # whether W2 has a compact K tail. Ungated 32-aligned intermediates such | ||
| # as I=224 therefore still need W13's 224 scale rows padded to 256; tying | ||
| # this padding to ``allow_w2_k_tail`` advertised the direct shape and then | ||
| # rejected it during preparation. | ||
| padded_w13_scale_n = _e8m0_logical_tail_scale_n(w13_rows) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Trace the main W4A16 GEMM's W13 scale stride derivation.
set -euo pipefail
rg -n -C6 '_e8m0_logical_tail_scale_n' sparkinfer/moe/_shared/kernels/w4a16/
# Inspect how the main GEMM sizes and strides the packed W13 E8M0 scale grid.
rg -n -C8 'w13_scales_i32_flat|w13_scale.*stride|two_n.*32|scale_format.*e8m0' sparkinfer/moe/_shared/kernels/w4a16/kernel.py | head -120Repository: local-inference-lab/sparkinfer
Length of output: 11911
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "## prepare.py relevant section"
sed -n '1140,1215p' sparkinfer/moe/_shared/kernels/w4a16/prepare.py
echo "## kernel.py W13 layout/type construction"
sed -n '650,720p' sparkinfer/moe/_shared/kernels/w4a16/kernel.py
echo "## kernel.py data preparation around scale tensors"
sed -n '778,855p' sparkinfer/moe/_shared/kernels/w4a16/kernel.py
echo "## stride/array accesses mentioning scale_size_n/scale_n_groups"
rg -n -C4 'scale_size_n|scale_n_groups|covered_size_n|self\.n_tiles|w1_scale|w13_scales_i32_flat|packed_w13_scale' sparkinfer/moe/_shared/kernels/w4a16/kernel.pyRepository: local-inference-lab/sparkinfer
Length of output: 20241
Align prepare.py’s padded_w13_scale_n with W4A16GemmKernel.scale_size_n.
prepare.py now always rounds W13 E8M0 scale rows up to 64, but W4A16GemmKernel only applies that padding when scale_format_e8m0_k32 and has_n_tile_tail. For I equal to a multiple of tile_k without a W2 compact K tail, the main kernel may use a smaller scale_n_groups/stride than the padded packed_w13_scale grid, especially with multi-chunk FC1. Use the same padding/stride calculation for both paths.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sparkinfer/moe/_shared/kernels/w4a16/prepare.py` around lines 1186 - 1191,
Align padded_w13_scale_n in the preparation path with
W4A16GemmKernel.scale_size_n: apply the 64-row E8M0 padding only when
scale_format_e8m0_k32 and has_n_tile_tail, matching the kernel’s scale_n_groups
and stride calculation. Preserve the unpadded row count for tile-aligned
dimensions without a W2 compact K tail, including multi-chunk FC1.
This was referenced Aug 1, 2026
Contributor
Author
|
Superseded by #110. Although #108 and #109 were individually mergeable against |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Bound unpinned W4A16 scratch to launch geometries reachable within the declared token capacity, expose a cheap arena-size query, and repair the tail cases that the tighter planner made reproducible.
This is the current-
mastersuccessor to #100. It contains only the commits after that PR's obsolete stacked base; no already-merged Trellis history is replayed.Planner changes
1 <= m <= max_tokens.c_tmpgeometry.fused_moe.required_nbytes(Caps)without compiling or retaining a launch.For the documented
max_tokens=64,topk=8,E=160,H=6144,I/rank=512case, the unpinned non-full W4A16 arena falls from 94.94 MiB to 30.03 MiB per rank. This does not alter kernel selection or numerical math.Correctness repairs
The bounded arena exposed pre-existing generic W4A16 tail defects. This PR also:
Aligned production shapes retain their existing geometry.
Verification
Focused suites cover fused/EP planning, route packing, replay, non-aligned native shapes, and W4A16 E2E cases. Changed-file Ruff checks and
git diff --checkpass.No merge is performed by this PR creation.
Summary by CodeRabbit
New Features
required_nbytesto estimate fused MoE workspace requirements without launching computation.Bug Fixes
Tests