fix(trellis): make mixed expert counts runtime-dynamic - #117
Conversation
📝 WalkthroughWalkthroughMixed Trellis now reconstructs expert-sized views from runtime tier counts and raw device pointers. The benchmark pads complete tier storage. Compilation supports SUH and SVH broadcasts, and runtime validation checks storage and pointer requirements before execution. ChangesMixed Trellis runtime
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller as mixed_trellis runtime
participant Cache as compilation cache
participant Validator as tier storage validation
participant Kernel as MixedTrellis kernel
Caller->>Cache: request compilation with tier counts and broadcast flags
Cache-->>Caller: return compiled artifact metadata
Caller->>Validator: validate storage, maps, and rotations
Validator-->>Caller: return validated raw pointers
Caller->>Kernel: launch with pointers and runtime expert counts
Kernel->>Kernel: reconstruct views and enforce expert bounds
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 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}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/moe/test_w4a16_mixed_trellis.py`:
- Line 343: Update the pytest.raises call’s match pattern to use a raw string
literal, preserving the existing regex and error text; follow the nearby
rf-string style used in the same test module.
🪄 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: f0885793-7bf6-4366-a0f3-4d80afd48128
📒 Files selected for processing (4)
benchmarks/benchmark_mixed_trellis.pysparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.pytests/moe/test_benchmark_mixed_trellis.pytests/moe/test_w4a16_mixed_trellis.py
Independent AIBeast field validation: exact r25 / #117 head — PASSI independently qualified the released implementation on the full 3.36-bpw Immutable inputs
Appliance production shape
This intentionally extends the release's 2,048-token/greedy gate with our Runtime and memory evidence
Matched benchmark
Strict structured output with thinking passed. The release gate used an actual I also qualified the exact #117 source over the older r20 integration before Conclusion: #117 fixes the actual 206/50 versus 160/96 geometry and is ready |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py (1)
1219-1232: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate the 16-byte alignment of
xbefore binding it as an aligned pointer.The new validation block checks dtype, device, contiguity, element count, and
% 16 != 0alignment for every tier-storage buffer and every rotation table.xis checked only for dtype and contiguity, yet line 1346 binds it withassumed_align=16. A caller-supplied misaligned view, for example a narrowed or externally wrapped bf16 buffer, then violates the alignment promise given to the compiler and the kernel performs undefined vectorized loads. Contiguity does not imply 16-byte base alignment.Add the same alignment check that the rotation tables use.
🛡️ Proposed fail-closed check
if not tensor.is_contiguous(): raise ValueError(f"mixed Trellis {name} must be contiguous") + if int(x.data_ptr()) % 16 != 0: + raise ValueError( + "mixed Trellis input must have at least 16-byte alignment" + )Also applies to: 1341-1347
🤖 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/mixed_trellis.py` around lines 1219 - 1232, Update the input validation loop for the “input” tensor x to verify its data pointer is 16-byte aligned, matching the rotation-table alignment check, before x is bound with assumed_align=16. Preserve the existing dtype and contiguity checks and raise the established validation error when x is misaligned.
🧹 Nitpick comments (1)
tests/moe/test_w4a16_mixed_trellis.py (1)
540-552: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the reverse fail-closed case.
The test rejects broadcast rotations against a non-broadcast launch. The opposite direction is untested: expanded, per-expert rotations against
broadcast_launchmust also fail, because validation requires exactly1 * hidden_sizeelements whenlaunch.broadcast_suhis true. That direction is the one a caller hits after switching a layer to broadcast without shrinking the tables, so it is worth pinning.💚 Proposed additional assertion
with pytest.raises(ValueError, match=r"gate SUH.*512 elements"): run_mixed_trellis( x, tier0, tier1, topk_weights, topk_ids, global_to_combined, descriptor, broadcast_rotations, expanded_launch, expanded_buffers, ) + with pytest.raises(ValueError, match=r"gate SUH.*128 elements"): + run_mixed_trellis( + x, + tier0, + tier1, + topk_weights, + topk_ids, + global_to_combined, + descriptor, + expanded_rotations, + broadcast_launch, + broadcast_buffers, + )🤖 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 `@tests/moe/test_w4a16_mixed_trellis.py` around lines 540 - 552, Add a reverse validation test alongside the existing mixed-trellis ValueError case: invoke run_mixed_trellis with expanded per-expert rotation tables and broadcast_launch enabled, then assert it raises ValueError mentioning the required gate SUH size of 1 × hidden_size. Reuse the existing test fixtures and arguments while switching only the launch/rotation configuration needed to cover the broadcast validation path.
🤖 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.
Outside diff comments:
In `@sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py`:
- Around line 1219-1232: Update the input validation loop for the “input” tensor
x to verify its data pointer is 16-byte aligned, matching the rotation-table
alignment check, before x is bound with assumed_align=16. Preserve the existing
dtype and contiguity checks and raise the established validation error when x is
misaligned.
---
Nitpick comments:
In `@tests/moe/test_w4a16_mixed_trellis.py`:
- Around line 540-552: Add a reverse validation test alongside the existing
mixed-trellis ValueError case: invoke run_mixed_trellis with expanded per-expert
rotation tables and broadcast_launch enabled, then assert it raises ValueError
mentioning the required gate SUH size of 1 × hidden_size. Reuse the existing
test fixtures and arguments while switching only the launch/rotation
configuration needed to cover the broadcast validation path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: a8baca17-605c-4ff8-8689-0889e087eb50
📒 Files selected for processing (2)
sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.pytests/moe/test_w4a16_mixed_trellis.py
3 participants
Summary
Mixed Trellis K3/K4 expert counts and shared-H rotation flags are checkpoint
data, not compile geometry. This PR passes both through the runtime ABI while
preserving one compiled kernel for compatible 256-expert partitions.
The first commit retains Michel Belleau's original #114 change with authorship
preserved. The follow-up fixes the benchmark's storage contract, carries the
broadcast-H metadata through the mixed wrapper/cache, and adds focused
regression coverage.
Problem
The shared-H GLM-5.2 EXL3 3.42 bpw checkpoint changes its tier partition by
layer:
Legacy checkpoints also use other valid partitions, including 192/64.
Embedding these counts or the H-storage layout in a cached kernel can reuse
stale bounds, offsets, or strides when the next layer differs.
The old benchmark also materialized only 16 experts per tier while advertising
a 192/64 launch and padded W13 only. ABI 5 correctly rejected its tight W2 and
global-scale storage before binding raw pointers.
Changes
storage bounds, and CUDA graph replay.
Validation
changing partitions, shared-H flags, high expert IDs, and storage rejection.
206/50, shared-H on both tiers, zero relative error, cosine 1.0.(
1.567x).DCP1/MTP3, and DCP4/MTP3, including CUDA graphs and c8/c16 batching.
with CUDA graphs.
Relationship to #114
This supersedes #114 because its external head cannot be amended here. The
runtime-count change is retained and credited through its original commit; the
additional commits supply the shared-H contract and validation required for a
mergeable implementation.