moe: add one-grid mixed K3/K4 Trellis execution - #104
Conversation
|
Warning Review limit reached
Next review available in: 30 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the 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 configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughAdds a W4A16 mixed K3/K4 Trellis MoE kernel with tier routing maps, combined rotations, capacity-checked runtime buffers, CUDA correctness tests, synthetic benchmarking, and checkpoint-backed validation. ChangesMixed Trellis execution
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant run_mixed_trellis
participant W4A16MixedTrellisKernel
participant topk_sum
Caller->>run_mixed_trellis: provide input, routes, maps, rotations, and buffers
run_mixed_trellis->>W4A16MixedTrellisKernel: validate and pack routes
W4A16MixedTrellisKernel->>topk_sum: pass routed partial outputs
topk_sum-->>run_mixed_trellis: return accumulated active-row output
🚥 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 review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (7)
sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py (2)
800-809: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winOptional: give the duck-typed tier object a
Protocol.
combine_trellis_rotationsandrun_mixed_trellisare both exported in__all__yet take untypedtier0/tier1and reach into ~10 attributes across them (intermediate_rotations,gate_suh,up_suh,down_svh,w13,w2,w13_scale,w2_scale,w13_global_scale,w2_global_scale). AProtocoldocuments the contract the serving framework must satisfy and catches attribute drift at type-check time rather than mid-launch.🤖 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 800 - 809, Define a typed duck-typed tier contract using a Protocol that declares the attributes consumed by combine_trellis_rotations and run_mixed_trellis, including the rotation tables and w13/w2 weight and scale fields. Annotate tier0 and tier1 with this Protocol in both exported functions, preserving the existing runtime behavior while enabling type checkers to detect attribute drift.
208-263: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: collapse the two tier branches into one parameterized emit.
The tier0/tier1 bodies are byte-identical apart from
(gemm, b, scales, global_scale, expert_bound). A local closure keeps the constexpr FC1/FC2 selection while removing the 26-line duplication, so future_run_tilesignature changes only need one edit.🤖 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 208 - 263, Collapse the duplicated tier0 and tier1 branches around the `_run_tile` calls into one parameterized emit, selecting the tier-specific `gemm`, weights (`t0_b_flat`/`t1_b_flat`), scales, global scale, and expert bound before invoking it. Preserve the `cutlass.const_expr(is_fc1)` FC1/FC2 selection and both tier validity checks, while ensuring `_run_tile` is defined in only one call site.tests/moe/test_w4a16_mixed_trellis.py (1)
229-252: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
skipped_map0clone is unused as a variation.
skipped_map0is an unmodified copy ofmap0; passingmap0directly makes the intent (only tier1's expert 3 is unmapped) clearer.♻️ Optional simplification
- skipped_map0 = map0.clone() skipped_map1 = map1.clone() skipped_map1[3] = -1 skipped_serial = _serial_tier( - x, tier0, topk_weights, topk_ids, skipped_map0 + x, tier0, topk_weights, topk_ids, map0 ) + _serial_tier(x, tier1, topk_weights, topk_ids, skipped_map1)🤖 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 229 - 252, Remove the unused skipped_map0 clone in the skipped-expert test and pass map0 directly to _serial_tier for tier0, while retaining skipped_map1 for the tier1 expert-3 unmapped variation.benchmarks/benchmark_mixed_trellis.py (3)
262-264: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the bound message from the constants.
The message hardcodes
[6, 64]while the check usesK4_EXPERTS; they will drift.♻️ Suggested change
- raise ValueError("materialized-experts must be in [6, 64]") + raise ValueError(f"materialized-experts must be in [6, {K4_EXPERTS}]")🤖 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 `@benchmarks/benchmark_mixed_trellis.py` around lines 262 - 264, Update the validation error in the materialized-experts handling to derive both displayed bounds from the existing minimum value and K4_EXPERTS constant instead of hardcoding “[6, 64]”, keeping the current validation behavior unchanged.
290-293: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRecord run provenance alongside the timings.
The header prints device, shapes, and topk, but not commit, worktree, GPU clock/persistence mode, or the invoking command — all required to make the reported ratios reproducible evidence.
As per coding guidelines: "Benchmark the real target path before making performance claims, recording the command, commit, worktree, GPU mode, correctness state, raw timings, and ratio direction."
🤖 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 `@benchmarks/benchmark_mixed_trellis.py` around lines 290 - 293, Update the benchmark header around the existing print statement to record provenance alongside device, shapes, and topk: include the invoking command, repository commit and worktree state, GPU clock/persistence mode, and correctness status, while preserving the raw timings and ratio direction in the benchmark output.Source: Coding guidelines
310-326: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBind loop variables explicitly in the captured closures.
Ruff flags ~30 B023 violations here. The current code is functionally correct (each closure is invoked and captured within the same iteration), but the lint gate will fail. Bind the loop-dependent state as default arguments.
♻️ Suggested binding
- def serial_fn(): - assert state0 is not None and state1 is not None - assert serial_output is not None + def serial_fn( + state0=state0, + state1=state1, + serial_output=serial_output, + tier_streams=tier_streams, + x=x, + weights=weights, + ids=ids, + ): + assert state0 is not None and state1 is not None + assert serial_output is not NoneApply the same pattern to
mixed_fnforlaunch,mixed_buffers,x,weights, andids.Also applies to: 357-372
🤖 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 `@benchmarks/benchmark_mixed_trellis.py` around lines 310 - 326, Update the loop-defined serial_fn and mixed_fn closures to bind every loop-dependent value through default arguments, including state0, state1, serial_output, and in mixed_fn launch, mixed_buffers, x, weights, and ids. Preserve the existing execution and stream-synchronization behavior while eliminating the B023 late-binding warnings.Source: Linters/SAST tools
benchmarks/validate_mixed_trellis_checkpoint.py (1)
212-215: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueHoist the 192/64/256 expert counts into named constants.
The K3/K4 counts and the 256-wide global namespace are duplicated across the validation check, both expert maps, and the route-slot computation. A single
K3_EXPERTS/K4_EXPERTS/TOTAL_EXPERTSset (matchingbenchmarks/benchmark_mixed_trellis.py) keeps them in sync and makes the failure at Line 213 self-documenting.Also applies to: 242-243, 275-275
🤖 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 `@benchmarks/validate_mixed_trellis_checkpoint.py` around lines 212 - 215, Define K3_EXPERTS, K4_EXPERTS, and TOTAL_EXPERTS constants in the validation module, matching the values used by benchmark_mixed_trellis.py. Replace the hard-coded 192, 64, and 256 values in the expert-count validation, both expert maps, and route-slot computation with these constants.
🤖 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 `@benchmarks/validate_mixed_trellis_checkpoint.py`:
- Line 34: Update DEFAULT_TILE_CONFIG to use the validated GLM tile geometry
with FC1_K at least 128, so the default path accepted by compile_mixed_trellis
succeeds. In the argument-parsing flow, validate tile_config[0] >= 128
immediately after parsing and report invalid --tile-config values as a CLI-level
error.
In `@sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py`:
- Around line 825-841: Split the active-row validation around m so non-positive
values report an invalid row count, while values above launch.size_m retain the
capacity-exceeded message. Update the x and topk_ids TypeError messages to
distinguish dtype mismatches from non-contiguous tensors, ensuring correctly
typed non-contiguous inputs explicitly report the contiguity failure.
- Around line 301-308: Update the extent computations used by the rotation_input
and topk_weights layouts to perform active_m scaling in Int64, and apply the
same widening to the active_m * top_k extent passed to fc2_emit. Preserve the
existing layout shapes and addressing behavior while ensuring all row-to-element
extent arithmetic avoids 32-bit evaluation.
- Around line 558-566: Clamp the cache-hit re-stamp in the mixed-trellis
specialization lookup so `size_m` and `max_m_blocks` cannot exceed the values
stored in the cached entry. Update the `replace(cached, ...)` path near
`cache_key` to preserve the compiled capacities while retaining the existing
behavior for compatible or smaller caller values; leave `run_mixed_trellis` and
cache-key construction unchanged.
- Line 748: Update the workspace allocation near run_mixed_trellis to derive its
size from the cooperative launch geometry, using blocks_per_sm * sms plus the
existing counter/lock padding instead of hard-coding 4. Ensure the allocation
matches every CTA-indexed access and remains consistent with the grid dimensions
used by run_mixed_trellis.
- Around line 858-859: Move the max-blocks capacity validation into
make_mixed_trellis_buffers before pack_topk_routes_by_expert writes to the
allocated buffers, using the actual request m or resulting packed block count
rather than launch.size_m/capacity_rows. Remove or adjust the later
block_experts.numel() guard so validation occurs before packing and reflects the
current request.
In `@tests/moe/test_w4a16_mixed_trellis.py`:
- Line 142: Replace the loose absolute-tolerance-only parity assertions for the
K3/K4 outputs and the skipped comparison with the relative-L2 validation style
already used by the large-M test. Update the assertions near the input setup and
the skipped comparison so the ~1e-3 signal is meaningfully gated, while
preserving the existing comparison targets and test behavior.
- Around line 262-352: Extend test_glm52_large_m_mixed_k3_k4_matches_serial with
a targeted big-pid paged-pool scenario, or add a nearby focused repro, that
places live page/row IDs beyond the 2^31/stride boundary while exercising the
same FC1/FC2 tile indexing and mixed k3/k4 path. Keep the serial comparison and
numerical assertion, and ensure the setup uses synthetic large strides or
expert-pool metadata without requiring an impractically large allocation.
- Around line 315-327: Update the compile_mixed_trellis setup to align
max_m_blocks with the emitted launch.moe_block_size: either assert
launch.moe_block_size == 8 before using the precomputed route_slots capacity, or
derive max_m_blocks from launch.moe_block_size instead of assuming 8. Preserve
the existing route capacity behavior for the validated block size.
---
Nitpick comments:
In `@benchmarks/benchmark_mixed_trellis.py`:
- Around line 262-264: Update the validation error in the materialized-experts
handling to derive both displayed bounds from the existing minimum value and
K4_EXPERTS constant instead of hardcoding “[6, 64]”, keeping the current
validation behavior unchanged.
- Around line 290-293: Update the benchmark header around the existing print
statement to record provenance alongside device, shapes, and topk: include the
invoking command, repository commit and worktree state, GPU clock/persistence
mode, and correctness status, while preserving the raw timings and ratio
direction in the benchmark output.
- Around line 310-326: Update the loop-defined serial_fn and mixed_fn closures
to bind every loop-dependent value through default arguments, including state0,
state1, serial_output, and in mixed_fn launch, mixed_buffers, x, weights, and
ids. Preserve the existing execution and stream-synchronization behavior while
eliminating the B023 late-binding warnings.
In `@benchmarks/validate_mixed_trellis_checkpoint.py`:
- Around line 212-215: Define K3_EXPERTS, K4_EXPERTS, and TOTAL_EXPERTS
constants in the validation module, matching the values used by
benchmark_mixed_trellis.py. Replace the hard-coded 192, 64, and 256 values in
the expert-count validation, both expert maps, and route-slot computation with
these constants.
In `@sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py`:
- Around line 800-809: Define a typed duck-typed tier contract using a Protocol
that declares the attributes consumed by combine_trellis_rotations and
run_mixed_trellis, including the rotation tables and w13/w2 weight and scale
fields. Annotate tier0 and tier1 with this Protocol in both exported functions,
preserving the existing runtime behavior while enabling type checkers to detect
attribute drift.
- Around line 208-263: Collapse the duplicated tier0 and tier1 branches around
the `_run_tile` calls into one parameterized emit, selecting the tier-specific
`gemm`, weights (`t0_b_flat`/`t1_b_flat`), scales, global scale, and expert
bound before invoking it. Preserve the `cutlass.const_expr(is_fc1)` FC1/FC2
selection and both tier validity checks, while ensuring `_run_tile` is defined
in only one call site.
In `@tests/moe/test_w4a16_mixed_trellis.py`:
- Around line 229-252: Remove the unused skipped_map0 clone in the
skipped-expert test and pass map0 directly to _serial_tier for tier0, while
retaining skipped_map1 for the tier1 expert-3 unmapped variation.
🪄 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: 10e7171f-66e2-4622-9c76-7d378f700c06
📒 Files selected for processing (4)
benchmarks/benchmark_mixed_trellis.pybenchmarks/validate_mixed_trellis_checkpoint.pysparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.pytests/moe/test_w4a16_mixed_trellis.py
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
benchmarks/benchmark_mixed_trellis.py (1)
393-418: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMake reported benchmark evidence reproducible and correctness-gated.
_times()samples are discarded into a median, provenance is not emitted, andserial/mixedvariants print timings without correctness status. Validate finite/nonzero/oracle outputs before timing, retain raw samples, and record argv, commit, worktree, GPU mode, and validation status with results.As per coding guidelines, “Benchmark the real target path before making performance claims, recording the command, commit, worktree, GPU mode, correctness state, raw timings, and ratio direction,” and “Validate correctness gates—including oracles, cosine/top-k equality, nonzero tensors, quantization semantics, and boundary behavior—before interpreting timings.”
🤖 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 `@benchmarks/benchmark_mixed_trellis.py` around lines 393 - 418, Update the benchmark flow around _times(), _capture(), and the serial/mixed result reporting to validate finite, nonzero oracle outputs and correctness before collecting or interpreting timings. Retain and emit raw timing samples alongside medians and clearly label ratio direction and validation status for both variants. Record argv, commit, worktree state, and GPU mode with each result so benchmark evidence is reproducible.Source: Coding guidelines
🤖 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 `@benchmarks/validate_mixed_trellis_checkpoint.py`:
- Around line 223-227: Update the checkpoint validation near the K3/K4 count
check to also require len(bitrates) == TOTAL_EXPERTS before loading the tiers;
reject any bitmap containing extra or missing bitrate entries while preserving
the existing K3/K4 validation and error behavior.
---
Outside diff comments:
In `@benchmarks/benchmark_mixed_trellis.py`:
- Around line 393-418: Update the benchmark flow around _times(), _capture(),
and the serial/mixed result reporting to validate finite, nonzero oracle outputs
and correctness before collecting or interpreting timings. Retain and emit raw
timing samples alongside medians and clearly label ratio direction and
validation status for both variants. Record argv, commit, worktree state, and
GPU mode with each result so benchmark evidence is reproducible.
🪄 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: 744c4cf5-cacd-42a9-9a56-0c34454396c1
📒 Files selected for processing (4)
benchmarks/benchmark_mixed_trellis.pybenchmarks/validate_mixed_trellis_checkpoint.pysparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.pytests/moe/test_w4a16_mixed_trellis.py
| if len(k3_ids) != K3_EXPERTS or len(k4_ids) != K4_EXPERTS: | ||
| raise ValueError( | ||
| f"expected {K3_EXPERTS} K3 and {K4_EXPERTS} K4 experts, " | ||
| f"got {len(k3_ids)} and {len(k4_ids)}" | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject checkpoints with extra bitrate entries.
The validation checks K3/K4 counts but not len(bitrates). A bitmap with exactly 192 K3 and 64 K4 entries plus an unsupported entry after ID 255 passes, and that extra expert is silently ignored by all later buffers and routes. Require len(bitrates) == TOTAL_EXPERTS before loading the tiers.
🐛 Proposed fix
- if len(k3_ids) != K3_EXPERTS or len(k4_ids) != K4_EXPERTS:
+ if (
+ len(bitrates) != TOTAL_EXPERTS
+ or len(k3_ids) != K3_EXPERTS
+ or len(k4_ids) != K4_EXPERTS
+ ):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if len(k3_ids) != K3_EXPERTS or len(k4_ids) != K4_EXPERTS: | |
| raise ValueError( | |
| f"expected {K3_EXPERTS} K3 and {K4_EXPERTS} K4 experts, " | |
| f"got {len(k3_ids)} and {len(k4_ids)}" | |
| ) | |
| if ( | |
| len(bitrates) != TOTAL_EXPERTS | |
| or len(k3_ids) != K3_EXPERTS | |
| or len(k4_ids) != K4_EXPERTS | |
| ): | |
| raise ValueError( | |
| f"expected {K3_EXPERTS} K3 and {K4_EXPERTS} K4 experts, " | |
| f"got {len(k3_ids)} and {len(k4_ids)}" | |
| ) |
🤖 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 `@benchmarks/validate_mixed_trellis_checkpoint.py` around lines 223 - 227,
Update the checkpoint validation near the K3/K4 count check to also require
len(bitrates) == TOTAL_EXPERTS before loading the tiers; reject any bitmap
containing extra or missing bitrate entries while preserving the existing K3/K4
validation and error behavior.
2 participants
Summary
Add a production mixed-bitrate Trellis path for rank-sliced EXL3 MoE checkpoints whose experts use two native packed bitrates, currently K3 and K4.
The API remains internal: the serving framework interprets checkpoint metadata and owns runtime planning, while SparkInfer owns prepared weights, buffers, and execution.
Supersedes #99. That comparison implementation duplicated the homogeneous kernel and used the
(64, 256, 64, 256)default tile. On the production GLM shape atM=3072, that geometry produced relative L2 error8.32e-3. This implementation requires an FC1 K tile of at least 128 and uses(128, 128, 32, 512)for the GLM geometry; measured relative L2 error is approximately4.1e-8at bothM=1andM=3072.Performance
Checkpoint:
willfalco/GLM-5.2-EXL3-TR3-3.25bpw, TP4/DCP4.Rank-0 Torch profiler:
82.38 -> 54.11 us/layer(-34.3%)6.764 -> 6.528 ms/layer(-3.5%)The exact-M specialization considered in #99 measured only
+0.22%E2E at CC1 and is intentionally not included.Validation
4 passedM=3072regression coveragepy_compile, andgit diff --checkpassThe PR is based directly on current
master; no release overlay or unrelated open PR is included.Summary by CodeRabbit