moe: stabilize Trellis arena memory profiling - #92
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughTrellis W4A16 scratch planning now disables source-incompatible token bucketing, precompiles full-rotation fused and top-k-sum launches, retains them on ChangesTrellis W4A16 planning
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant plan_tp_moe_scratch
participant _plan_full_rotation_w4a16_launches
participant TrellisKernels
participant _build_tp_moe_fp4_binding_from_views
plan_tp_moe_scratch->>_plan_full_rotation_w4a16_launches: plan fixed-capacity Trellis launches
_plan_full_rotation_w4a16_launches->>TrellisKernels: compile fused and top-k-sum launches
TrellisKernels-->>plan_tp_moe_scratch: return launch objects
plan_tp_moe_scratch->>_build_tp_moe_fp4_binding_from_views: build binding with fused_launch=None and topk_sum_launch=None
Possibly related PRs
Suggested reviewers: 🚥 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 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: 2
🧹 Nitpick comments (3)
tests/moe/test_tp_moe_scratch_bindings.py (2)
557-561: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBound the arena size from both sides.
core_workspace_nbytes < 1060 MiBis a one-sided magic threshold: an accidental under-sizing of the Trellis arena also satisfies it, so the assertion can't distinguish "bucketing correctly disabled" from "capacity regressed". Add a lower bound (or assert equality against the size a 3072-token plan should produce) so the exact-capacity contract is actually pinned.💚 Suggested tightening
assert plan.layout.route_workspace_nbytes == 0 - assert plan.layout.core_workspace_nbytes < 1060 * (1 << 20) + # Must land on the exact 3072-token capacity, not the 4096 compile bucket + # (~1060 MiB) and not an under-sized arena. + assert 1000 * (1 << 20) < plan.layout.core_workspace_nbytes < 1060 * (1 << 20) assert plan.layout.total_nbytes == plan.layout.core_workspace_nbytes🤖 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_tp_moe_scratch_bindings.py` around lines 557 - 561, Add a lower-bound assertion for plan.layout.core_workspace_nbytes in this test, using the expected minimum capacity for the 3072-token plan while retaining the existing upper bound. This should verify both that Trellis bucketing is disabled and that the arena is not accidentally undersized.
574-578: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe real launch planner is left untested.
Stubbing
_plan_full_rotation_w4a16_launchesis right for isolating the retention/dispatch contract, but it means the ~190 lines of new compile/prewarm logic have no direct coverage — only its CPU early-return is touched (via the sibling test). A prewarm miss fails silently as a memory-profiling regression, not an exception, so a GPU-gated test asserting the compiled fused token-count set (range(1, capacity+1)for capacity ≤ 32 vs. a single capacity entry above it) and the four(ids_dtype, mapped)top-k-sum variants would be worth adding.Want me to draft that GPU-gated test?
🤖 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_tp_moe_scratch_bindings.py` around lines 574 - 578, Add a GPU-gated test that exercises the real _plan_full_rotation_w4a16_launches planner instead of monkeypatching it, and verify the compiled fused token-count set is range(1, capacity+1) for capacities up to 32 or only the capacity entry above 32. Also assert compilation/prewarm covers all four (ids_dtype, mapped) top-k-sum variants, while preserving the existing isolated retention/dispatch test.sparkinfer/moe/fused_moe/_impl.py (1)
6198-6210: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winPrefer failing loudly over re-deriving
block_size_mwhen the plan omits it.
core_plan.route_block_size_m or select_route_block_size_m(...)re-derives a value the arena was already sized with. If the plan's block size is ever absent/zero, the fallback can choose a different block size than the one used to sizeblock_expert_ids/packed_route_indices, and that mismatchedcapacity_m_blocksis then baked into the compiled launch'smax_m_blocksgrid contract — a silent geometry divergence rather than an error.♻️ Suggested change
- block_size_m = core_plan.route_block_size_m or select_route_block_size_m( - capacity_tokens, - core_plan.num_topk, - core_plan.route_E, - ) + if core_plan.route_block_size_m: + block_size_m = int(core_plan.route_block_size_m) + else: + # No planned route geometry: derive it once, identically to the arena. + block_size_m = select_route_block_size_m( + capacity_tokens, + core_plan.num_topk, + core_plan.route_E, + )🤖 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/fused_moe/_impl.py` around lines 6198 - 6210, Update the block-size selection near capacity_m_blocks to require a valid nonzero core_plan.route_block_size_m and fail immediately when it is absent or zero; remove the select_route_block_size_m fallback. Continue using the validated plan block size for max_packed_route_slots and capacity_m_blocks so the compiled launch geometry matches the arena sizing.
🤖 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/fused_moe/_impl.py`:
- Around line 6287-6288: Move the torch.cuda.is_current_stream_capturing() guard
to the beginning of the route-pack prewarm/planning function, before any
compilation or other CUDA-visible work. Preserve the existing RuntimeError
message and ensure the function exits immediately whenever capture is active.
- Around line 6247-6261: The Trellis prewarm configuration must use the same
normalized W4a16 scale format, weight layout, and w13 layout as the runtime
path. Update the prewarm code around the visible trellis settings to derive
these values from the plan or prepared_payload through
_normalize_w4a16_scale_format and _normalize_w4a16_weight_layout, or explicitly
validate the required Trellis contract before prewarming so mismatches fail
during planning rather than producing divergent cache keys.
---
Nitpick comments:
In `@sparkinfer/moe/fused_moe/_impl.py`:
- Around line 6198-6210: Update the block-size selection near capacity_m_blocks
to require a valid nonzero core_plan.route_block_size_m and fail immediately
when it is absent or zero; remove the select_route_block_size_m fallback.
Continue using the validated plan block size for max_packed_route_slots and
capacity_m_blocks so the compiled launch geometry matches the arena sizing.
In `@tests/moe/test_tp_moe_scratch_bindings.py`:
- Around line 557-561: Add a lower-bound assertion for
plan.layout.core_workspace_nbytes in this test, using the expected minimum
capacity for the 3072-token plan while retaining the existing upper bound. This
should verify both that Trellis bucketing is disabled and that the arena is not
accidentally undersized.
- Around line 574-578: Add a GPU-gated test that exercises the real
_plan_full_rotation_w4a16_launches planner instead of monkeypatching it, and
verify the compiled fused token-count set is range(1, capacity+1) for capacities
up to 32 or only the capacity entry above 32. Also assert compilation/prewarm
covers all four (ids_dtype, mapped) top-k-sum variants, while preserving the
existing isolated retention/dispatch test.
🪄 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: 2ba3e8ac-003b-45e7-969a-5a86c0c528f3
📒 Files selected for processing (2)
sparkinfer/moe/fused_moe/_impl.pytests/moe/test_tp_moe_scratch_bindings.py
|
@coderabbitai review |
✅ Action performedReview finished.
|
2 participants
Problem
The unified Trellis
fused_moepath used the generic W4A16 compile bucket for its caller-owned arena and left several fixed launch variants to first use. For a 3,072-token EXL3 prefill plan this rounded the arena to 4,096 rows, reserving about 320 MiB/rank that execution could not use. First-use compilation during vLLM profiling also produced non-repeatable rank-local activation peaks and therefore non-repeatable distributed KV capacity.Changes
exl3_trellis_mcgscratch plans; generic W4A16 plans retain compile bucketing.The original Trellis implementation from closed #90 is already present on
master; this PR contains only the follow-up profiling and capacity fix.Validation
git diff --check: cleanPaired vLLM PRs: local-inference-lab/vllm#190 and local-inference-lab/vllm#198.
Summary by CodeRabbit
Performance Improvements
Bug Fixes
Tests