moe: integrate EXL3 Trellis into the fused W4A16 API - #90
Conversation
📝 WalkthroughWalkthroughAdds EXL3 Trellis W4A16 support across dequantization intrinsics, dense GEMM, fused MoE preparation, routing, full-rotation kernels, workspace management, public APIs, and CUDA validation tests. ChangesEXL3 Trellis W4A16
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Caller
participant sparkinfer_moe_fp4
participant build_tp_moe_fp4_binding
participant run_w4a16_moe
participant W4A16FusedMoeKernel
Caller->>sparkinfer_moe_fp4: provide Trellis weights and routing inputs
sparkinfer_moe_fp4->>build_tp_moe_fp4_binding: prepare workspace and bindings
build_tp_moe_fp4_binding->>run_w4a16_moe: pass Trellis rotation buffers and maps
run_w4a16_moe->>W4A16FusedMoeKernel: launch full-rotation W4A16 execution
W4A16FusedMoeKernel-->>Caller: return FP32 routed 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 review |
✅ Action performedReview finished.
|
Co-authored-by: Brandon Music <brandon.m.music@gmail.com>
157c06e to
5a82b70
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (13)
tests/moe/test_fused_moe_trellis.py (1)
551-551: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTests reach into
plan._core_workspace_plan.Asserting on a private attribute couples the suite to internal plan structure. Consider exposing a small public accessor (e.g.
plan.full_rotation) so the invariant can be checked without touching internals.Also applies to: 705-705
🤖 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_fused_moe_trellis.py` at line 551, Expose the full-rotation state through a public accessor on the plan object, such as full_rotation, and update both assertions currently reading plan._core_workspace_plan.full_rotation to use that accessor. Keep the returned value equivalent to the existing internal invariant.tests/gemm/test_trellis_linear.py (1)
29-31: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer per-element identity assertions over tuple equality for tensor args.
seen["args"] == tensorsonly passes today because CPython's tuple comparison short-circuits on identity. If the implementation ever copies an argument, the comparison falls through toTensor.__eq__and the assert raises an ambiguous-truth-valueRuntimeErrorinstead of failing cleanly — which defeats the "without_copy" intent of the test.♻️ Suggested assertion
- assert seen["args"] == tensors + assert all(seen_arg is arg for seen_arg, arg in zip(seen["args"], tensors, strict=True))🤖 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/gemm/test_trellis_linear.py` around lines 29 - 31, Replace the tuple equality assertion for seen["args"] in the test with per-element identity assertions against tensors, ensuring each argument is the exact same tensor object. Keep the existing codebook and params_dtype assertions unchanged.sparkinfer/moe/_shared/kernels/w4a16/prepare.py (1)
1565-1566: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRuff RUF005: prefer unpacking over tuple concatenation.
♻️ Proposed fix
- expected_i16_shape = expected_prefix_shape + (16 * trellis_bits,) - expected_i32_shape = expected_prefix_shape + (8 * trellis_bits,) + expected_i16_shape = (*expected_prefix_shape, 16 * trellis_bits) + expected_i32_shape = (*expected_prefix_shape, 8 * trellis_bits)🤖 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 1565 - 1566, Update the shape construction in the code defining expected_i16_shape and expected_i32_shape to use tuple unpacking instead of concatenating expected_prefix_shape with trailing tuples, while preserving the existing dimensions and values.Source: Linters/SAST tools
sparkinfer/_lib/intrinsics.py (1)
6172-6174: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConstants are declared but the PTX hardcodes the same literals.
_TRELLIS_MCG,_TRELLIS_MASK, and_TRELLIS_ORare never referenced; both decoders inline0xCBAC1FED/0x8fff8fff/0x3b603b60directly. That leaves two sources of truth for the codebook parameters, which can silently diverge if one is updated.♻️ Interpolate the constants into the asm template
- mov.b32 M, 0xCBAC1FED; + mov.b32 M, __MCG__;and extend the existing substitution chain:
asm = ( asm.replace("__B1__", str(bits)) .replace("__B2__", str(2 * bits)) .replace("__B3__", str(3 * bits)) + .replace("__MCG__", f"{_TRELLIS_MCG:`#010x`}") )(apply the same for the mask/or operands in both decoders)
🤖 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/_lib/intrinsics.py` around lines 6172 - 6174, Update both decoder PTX assembly templates to interpolate _TRELLIS_MCG, _TRELLIS_MASK, and _TRELLIS_OR instead of embedding their hexadecimal literals, extending the existing substitution chain for each operand. Ensure all references in both decoders use these constants so the codebook parameters have a single source of truth.sparkinfer/moe/_shared/kernels/w4a16/kernel.py (4)
9256-9267: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueForced-tile fit check omits the 512-byte margin used elsewhere.
_select_tile_configcalls_candidate_tile_fitswithmax_shared_mem - 512, but the forced path passes the raw limit, so a forced tile can be accepted with no headroom for the+1536launch overhead_determine_blocks_per_smassumes. The new hard guard inW4A16GemmKernel.__init__fails closed afterwards, so this is only a consistency point.🤖 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/kernel.py` around lines 9256 - 9267, Update the forced-tile `_candidate_tile_fits` call in `_select_tile_config` to pass `max_shared_mem - 512`, matching the regular tile-selection path and preserving the required shared-memory headroom.
5552-5590: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
_had128_quadis duplicated verbatim across two kernel classes.The same warp butterfly and the
0.088388347648(= 1/sqrt(128)) scale exist independently inW4A16FusedMoeKernelandW4A16TopKSumKernel. Since the fused epilogue and the top-k sum must apply exactly the same H128, hoist it to one module-level@cute.jithelper with a named constant for the scale so the two cannot drift.Also applies to: 6543-6577
🤖 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/kernel.py` around lines 5552 - 5590, The _had128_quad implementation is duplicated between W4A16FusedMoeKernel and W4A16TopKSumKernel. Move the shared warp butterfly logic into one module-level `@cute.jit` helper, define a named constant for the 1/sqrt(128) scale, and update both class methods/call sites to reuse that helper so both paths remain identical.
10338-10446: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftThe 43-name tuple unpack is a silent-mismap hazard.
launch_common/launch_tailare built positionally, then destructured here into_rc_*/_lt_*locals and re-passed as keywords. Most entries areint/bool/str, so reordering either tuple (or inserting a field) mismaps arguments with no type error — e.g.swiglu_alpha/swiglu_betaor the four tile dims swapping would produce wrong results, not a crash. Build the launch arguments as a single dict once and splat it into both the registered op and_w4a16_fused_moe_launch_flat.🤖 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/kernel.py` around lines 10338 - 10446, Replace the positional `launch_common`/`launch_tail` tuple unpacking with one named launch-argument dictionary built at the source of those values. Reuse that dictionary via keyword expansion for both the registered operation and `_w4a16_fused_moe_launch_flat`, preserving existing values and overrides such as `launch_a`, `launch_a_up`, `rot_arg`, and `full_rotation` so field reordering cannot silently remap arguments.
5302-5421: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift
_run_activation_compactand_run_activation's rotation branch duplicate the whole rotation epilogue.The two bodies differ only in how the rotation-scale row is derived (
expert * 3Ivsrow * 3I) and in how rows are enumerated (packed-block traversal vs dense route index); the H128 → svh → silu/situ → suh_down → H128 math is copied verbatim. Extracting the per-quad math into one@cute.jithelper taking(g_base, u_base, s_base, out_base, lane)would keep the legacy and full-rotation numerics from diverging on the next change.Also applies to: 5605-5690
🤖 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/kernel.py` around lines 5302 - 5421, Extract the duplicated H128 rotation epilogue from _run_activation_compact and the rotation branch of _run_activation into a shared `@cute.jit` helper accepting g_base, u_base, s_base, out_base, and lane. Keep each caller’s existing row enumeration and rotation-scale base calculation (expert * 3I for compact, row * 3I for dense), and have the helper preserve the existing H128, SiLU/SwiGLU or SiT activation, down-scale, and output-store behavior.sparkinfer/moe/fused_moe/_impl.py (4)
801-805: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winValidate
w4a16_block_size_magainst the shared allowed-size constant.
plan_w4a16_buffersand_plan_core_workspaceboth key off_W4A16_ALLOWED_ROUTED_SIZES; hard-coding(8, 16, 32, 48, 64)here means a change to that constant silently splits caps validation from planning.♻️ Proposed fix
if self.w4a16_block_size_m is not None: + from sparkinfer.moe._shared.kernels.w4a16.host import ( + _W4A16_ALLOWED_ROUTED_SIZES, + ) + block_size_m = int(self.w4a16_block_size_m) - if block_size_m not in (8, 16, 32, 48, 64): - raise ValueError("w4a16_block_size_m must be one of 8, 16, 32, 48, 64") + if block_size_m not in _W4A16_ALLOWED_ROUTED_SIZES: + raise ValueError( + "w4a16_block_size_m must be one of " + f"{_W4A16_ALLOWED_ROUTED_SIZES}, got {block_size_m}" + ) object.__setattr__(self, "w4a16_block_size_m", block_size_m)🤖 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 801 - 805, Update the w4a16_block_size_m validation in the surrounding configuration logic to check membership in the shared _W4A16_ALLOWED_ROUTED_SIZES constant instead of hard-coding the tuple, while preserving integer normalization and the existing ValueError behavior.
2536-2537: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReuse
_TRELLIS256_BITSinstead of re-listing the bitrates.
kernel.pyalready owns_TRELLIS256_BITS = (3, 4, 5, 6)and validates against it inW4A16GemmKernel,compile_w4a16_fused_moe, andrun_w4a16_moe. A fourth hand-written copy here will drift if a bitrate is added or dropped.🤖 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 2536 - 2537, Update the trellis_bits validation near this ValueError to reuse the existing _TRELLIS256_BITS constant instead of hard-coding (3, 4, 5, 6), preserving the current integer conversion and error behavior.
6261-6261: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
topk_sum_launchesannotation no longer matches its keys.Under
full_rotationthe dict is keyed by(token_count, ids_dtype, mapped)tuples, but the local is annotateddict[int, object]. Widen it to match the workspace field (dict[object, object]) so type checkers and readers see the real contract.♻️ Proposed fix
- topk_sum_launches: dict[int, object] = {} + topk_sum_launches: dict[object, object] = {}Also applies to: 6314-6329
🤖 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` at line 6261, Update the topk_sum_launches local declaration to use the workspace field’s dict[object, object] annotation, covering both the initial declaration and the related handling around the full_rotation path, while preserving the existing tuple-key behavior.
9714-9720: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value
replace(prepared, ...)allocates a new prepared container on every launch.This runs on the per-call W4A16 launch path purely to swap one field. Caching the reconstructed object on the binding (or passing
kernel_workspacetorun_w4a16_moedirectly) removes a dataclass copy from the decode critical path.🤖 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 9714 - 9720, Update the full_rotation path around run_w4a16_moe so it does not call replace(prepared, ...) on every launch; cache the reconstructed prepared container on binding and reuse it, or pass the binding’s kernel_workspace directly to run_w4a16_moe while preserving the existing workspace value and behavior.sparkinfer/moe/_shared/kernels/w4a16/route_pack.py (1)
410-438: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeduplicate the two identical count+prefix launch sequences.
The
expert_counts is not Nonebranch repeats the eager branch's_w4a16_route_count_kernel+_w4a16_route_prefix_from_counts_kernellaunches verbatim; only the buffer acquisition differs. Resolving the buffer first and sharing one launch block keeps the two paths from drifting.♻️ Proposed consolidation
- if expert_counts is not None: - expert_counts = _workspace_slice( - expert_counts, - name="expert_counts", - elements=int(num_experts), - dtype=torch.int32, - device=topk_ids.device, - ) - expert_counts.zero_() - _w4a16_route_count_kernel[(triton.cdiv(numel, _FAST_COUNT_BLOCK_T),)]( - topk_ids, - expert_map_tensor, - expert_counts, - numel, - NUM_EXPERTS=int(num_experts), - HAS_EXPERT_MAP=expert_map is not None, - BLOCK_T=_FAST_COUNT_BLOCK_T, - num_warps=4, - ) - _w4a16_route_prefix_from_counts_kernel[(1,)]( - expert_counts, - packed_route_count, - expert_offsets, - BLOCK_SIZE=int(block_size), - NUM_EXPERTS=int(num_experts), - BLOCK_E=block_e, - num_warps=4, - ) - elif not torch.cuda.is_current_stream_capturing(): + counts_buffer = None + if expert_counts is not None: + counts_buffer = _workspace_slice( + expert_counts, + name="expert_counts", + elements=int(num_experts), + dtype=torch.int32, + device=topk_ids.device, + ) + counts_buffer.zero_() + elif not torch.cuda.is_current_stream_capturing(): # FAST (eager prefill): ... - expert_counts = torch.zeros( + counts_buffer = torch.zeros( int(num_experts), dtype=torch.int32, device=topk_ids.device ) + if counts_buffer is not None: _w4a16_route_count_kernel[(triton.cdiv(numel, _FAST_COUNT_BLOCK_T),)]( topk_ids, expert_map_tensor, - expert_counts, + counts_buffer, numel, NUM_EXPERTS=int(num_experts), HAS_EXPERT_MAP=expert_map is not None, BLOCK_T=_FAST_COUNT_BLOCK_T, num_warps=4, ) _w4a16_route_prefix_from_counts_kernel[(1,)]( - expert_counts, + counts_buffer, packed_route_count, expert_offsets, BLOCK_SIZE=int(block_size), NUM_EXPERTS=int(num_experts), BLOCK_E=block_e, num_warps=4, ) else:🤖 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/route_pack.py` around lines 410 - 438, Refactor the routing setup around expert_counts so the buffer is selected or acquired first, then execute a single shared _w4a16_route_count_kernel and _w4a16_route_prefix_from_counts_kernel launch sequence for both branches. Preserve the existing expert_counts workspace initialization and eager-path behavior, removing the duplicated launches without changing their arguments or ordering.
🤖 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/w4a16/kernel.py`:
- Around line 9461-9469: Update run_trellis256_dense and its callers to accept
optional caller-owned scratch buffers for the float16 input conversion, rotated
output, and float16 GEMM output; reuse these buffers instead of calling torch.to
or torch.empty_like per invocation. Preserve existing behavior by allocating
fallbacks only when buffers are not supplied, and ensure capture paths pass and
reuse the storage.
- Around line 1534-1535: Update route-major A indexing in
W4A16FusedMoeKernel.__call__, _stage_k_tile_async, and _drain_output_smem so
row/route IDs are widened to Int64 before multiplying by A or C row strides.
Apply the same Int64 offset arithmetic to the additional affected indexing paths
around the referenced locations, ensuring scaling by top_k cannot overflow Int32
before get_ptr_as_int64.
- Around line 4673-4677: Update the ValueError message in the
intermediate_rotation validation guard to state that SiTU is supported and that
the restriction applies to incompatible SwiGLU variants, matching the existing
conditions and _run_activation_compact handling. Leave the guard logic
unchanged.
- Around line 8407-8418: Update the prewarm flow associated with
_w4a16_fused_moe_launch_flat to call _rot_scales_dummy for every target device
before CUDA graph capture can begin, ensuring the cached tensor is materialized
outside the capture region. Preserve the existing per-device cache and avoid
allocating the dummy lazily on the serving launch path; alternatively, assert
that _rot_scales_dummy is never called while capture is active.
In `@sparkinfer/moe/fused_moe/_impl.py`:
- Around line 4845-4854: Make EXL3 preparation fail closed when the required MCG
marker is absent or invalid: in the surrounding EXL3 handling and the
corresponding path near the second `prepare_trellis256_moe_weights` call,
require a valid `trellis_mcg` marker before passing `codebook="mcg"`. Reject
marker-less or non-MCG checkpoints instead of preparing them with the wrong
codebook, preserving execution only for the validated MCG format.
- Around line 2282-2301: Reject non-None route_expert_map and output_expert_map
whenever plan.full_rotation is false, before the existing full-rotation
validation and execution branching. Add a clear ValueError identifying the
unsupported map arguments, while preserving the current validation and
propagation behavior for full-rotation plans.
- Around line 6876-6878: Zero kernel_workspace before it is used by the
workspace-pool bind path, ensuring recycled slots contain no stale barrier,
lock, or scalar state. Update the kernel_workspace allocation spec to use zero
initialization, or invoke the existing kernel_workspace.zero_() during the
binding flow around build_tp_moe_fp4_binding; preserve the full_rotation output
validation.
In `@tests/moe/test_fused_moe_trellis.py`:
- Line 219: Update the pytest.raises call using the `match=` argument in the
relevant test to express the regex pattern as a raw string, preserving the
existing `.*MCG marker` matching behavior and silencing RUF043.
---
Nitpick comments:
In `@sparkinfer/_lib/intrinsics.py`:
- Around line 6172-6174: Update both decoder PTX assembly templates to
interpolate _TRELLIS_MCG, _TRELLIS_MASK, and _TRELLIS_OR instead of embedding
their hexadecimal literals, extending the existing substitution chain for each
operand. Ensure all references in both decoders use these constants so the
codebook parameters have a single source of truth.
In `@sparkinfer/moe/_shared/kernels/w4a16/kernel.py`:
- Around line 9256-9267: Update the forced-tile `_candidate_tile_fits` call in
`_select_tile_config` to pass `max_shared_mem - 512`, matching the regular
tile-selection path and preserving the required shared-memory headroom.
- Around line 5552-5590: The _had128_quad implementation is duplicated between
W4A16FusedMoeKernel and W4A16TopKSumKernel. Move the shared warp butterfly logic
into one module-level `@cute.jit` helper, define a named constant for the
1/sqrt(128) scale, and update both class methods/call sites to reuse that helper
so both paths remain identical.
- Around line 10338-10446: Replace the positional `launch_common`/`launch_tail`
tuple unpacking with one named launch-argument dictionary built at the source of
those values. Reuse that dictionary via keyword expansion for both the
registered operation and `_w4a16_fused_moe_launch_flat`, preserving existing
values and overrides such as `launch_a`, `launch_a_up`, `rot_arg`, and
`full_rotation` so field reordering cannot silently remap arguments.
- Around line 5302-5421: Extract the duplicated H128 rotation epilogue from
_run_activation_compact and the rotation branch of _run_activation into a shared
`@cute.jit` helper accepting g_base, u_base, s_base, out_base, and lane. Keep each
caller’s existing row enumeration and rotation-scale base calculation (expert *
3I for compact, row * 3I for dense), and have the helper preserve the existing
H128, SiLU/SwiGLU or SiT activation, down-scale, and output-store behavior.
In `@sparkinfer/moe/_shared/kernels/w4a16/prepare.py`:
- Around line 1565-1566: Update the shape construction in the code defining
expected_i16_shape and expected_i32_shape to use tuple unpacking instead of
concatenating expected_prefix_shape with trailing tuples, while preserving the
existing dimensions and values.
In `@sparkinfer/moe/_shared/kernels/w4a16/route_pack.py`:
- Around line 410-438: Refactor the routing setup around expert_counts so the
buffer is selected or acquired first, then execute a single shared
_w4a16_route_count_kernel and _w4a16_route_prefix_from_counts_kernel launch
sequence for both branches. Preserve the existing expert_counts workspace
initialization and eager-path behavior, removing the duplicated launches without
changing their arguments or ordering.
In `@sparkinfer/moe/fused_moe/_impl.py`:
- Around line 801-805: Update the w4a16_block_size_m validation in the
surrounding configuration logic to check membership in the shared
_W4A16_ALLOWED_ROUTED_SIZES constant instead of hard-coding the tuple, while
preserving integer normalization and the existing ValueError behavior.
- Around line 2536-2537: Update the trellis_bits validation near this ValueError
to reuse the existing _TRELLIS256_BITS constant instead of hard-coding (3, 4, 5,
6), preserving the current integer conversion and error behavior.
- Line 6261: Update the topk_sum_launches local declaration to use the workspace
field’s dict[object, object] annotation, covering both the initial declaration
and the related handling around the full_rotation path, while preserving the
existing tuple-key behavior.
- Around line 9714-9720: Update the full_rotation path around run_w4a16_moe so
it does not call replace(prepared, ...) on every launch; cache the reconstructed
prepared container on binding and reuse it, or pass the binding’s
kernel_workspace directly to run_w4a16_moe while preserving the existing
workspace value and behavior.
In `@tests/gemm/test_trellis_linear.py`:
- Around line 29-31: Replace the tuple equality assertion for seen["args"] in
the test with per-element identity assertions against tensors, ensuring each
argument is the exact same tensor object. Keep the existing codebook and
params_dtype assertions unchanged.
In `@tests/moe/test_fused_moe_trellis.py`:
- Line 551: Expose the full-rotation state through a public accessor on the plan
object, such as full_rotation, and update both assertions currently reading
plan._core_workspace_plan.full_rotation to use that accessor. Keep the returned
value equivalent to the existing internal invariant.
🪄 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: 3f95085b-a40d-4347-8522-7422b18ef146
📒 Files selected for processing (15)
sparkinfer/__init__.pysparkinfer/_lib/intrinsics.pysparkinfer/gemm/__init__.pysparkinfer/gemm/trellis_linear/__init__.pysparkinfer/gemm/trellis_linear/api.pysparkinfer/moe/__init__.pysparkinfer/moe/_shared/execution.pysparkinfer/moe/_shared/kernels/w4a16/host.pysparkinfer/moe/_shared/kernels/w4a16/kernel.pysparkinfer/moe/_shared/kernels/w4a16/prepare.pysparkinfer/moe/_shared/kernels/w4a16/route_pack.pysparkinfer/moe/fused_moe/__init__.pysparkinfer/moe/fused_moe/_impl.pytests/gemm/test_trellis_linear.pytests/moe/test_fused_moe_trellis.py
|
Luke made a long post of changes to be made on #49 . Should I port that suggested work to this PR? Let me know or if you’d already done it! |
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)
sparkinfer/moe/_shared/kernels/w4a16/host.py (1)
364-379: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
full_rotationbuffers keep the callerdtypefor the intermediate caches while onlyoutputflips to fp32.
run_w4a16_moerejects full rotation unless the prepared/scratch element dtype is fp16, and the arena planner insparkinfer/moe/fused_moe/_impl.py(Line 2609) selectscache_dtype = torch.float16 if full_rotation else dtypeforintermediate_cache13/intermediate_cache2. Here both caches staydtype, so afull_rotation=True, dtype=torch.bfloat16call produces buffers that disagree with the fp16 kernel contract (rotation_a_* are already hard-coded fp16).🛡️ Mirror the planner's cache dtype
+ cache_dtype = torch.float16 if full_rotation else dtype return W4A16PackedBuffers( intermediate_cache13=torch.empty( (plan.intermediate_cache13_elements,), - dtype=dtype, + dtype=cache_dtype, device=device, ), intermediate_cache2=torch.empty( (plan.routed_rows, int(prepared.intermediate_size)), - dtype=dtype, + dtype=cache_dtype, device=device, ),🤖 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/host.py` around lines 364 - 379, Update the W4A16 buffer allocation around W4A16PackedBuffers so intermediate_cache13 and intermediate_cache2 use fp16 when full_rotation is enabled, matching the arena planner and kernel contract; retain the caller dtype otherwise and leave output’s existing dtype selection unchanged.
🧹 Nitpick comments (8)
tests/moe/test_fused_moe_trellis.py (1)
120-124: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe non-MCG rejection assertion is tautological.
_prepare_weightsraisesNotImplementedError("...only the MCG codebook")itself, sotest_prepare_weights_rejects_non_mcg_before_cuda_work(Line 208) only ever exercises this helper —fused_moe.prepare_weightsis never called with a non-MCG codebook. The production fail-closed behaviour (trellis_codebookgating inPreparedW4A16MoeWeights) is therefore untested. Either forwardcodebookintofused_moe.prepare_weights/plan_weightsand assert the library raises, or drop the assertion so it doesn't imply coverage it doesn't have.🤖 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_fused_moe_trellis.py` around lines 120 - 124, Update test_prepare_weights_rejects_non_mcg_before_cuda_work and _prepare_weights so the non-MCG codebook reaches fused_moe.prepare_weights or plan_weights, then assert the production API raises the expected NotImplementedError. Remove the helper’s own preemptive rejection, or otherwise avoid asserting that helper-only check as coverage of PreparedW4A16MoeWeights trellis_codebook gating.sparkinfer/moe/_shared/kernels/w4a16/prepare.py (2)
1535-1540: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCompute
words_per_bitafter the dtype whitelist.The ternary silently assigns 8 words for any non-int16 dtype before the
TypeErrorguard runs; moving it below keeps the invariant local and avoids a misleading intermediate value if the guard is ever relaxed.♻️ Proposed reorder
- words_per_bit = 16 if tensor.dtype == torch.int16 else 8 if tensor.dtype not in (torch.int16, torch.int32): raise TypeError( f"trellis3_t256 {name} must use native int16 or int32 storage, " f"got {tensor.dtype}" ) + words_per_bit = 16 if tensor.dtype == torch.int16 else 8🤖 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 1535 - 1540, Move the words_per_bit assignment below the dtype whitelist check in the trellis3_t256 validation logic, keeping the existing int16/int32 rejection behavior and assignment values unchanged.
1565-1566: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse unpacking instead of tuple concatenation (Ruff RUF005).
♻️ Proposed fix
- expected_i16_shape = expected_prefix_shape + (16 * trellis_bits,) - expected_i32_shape = expected_prefix_shape + (8 * trellis_bits,) + expected_i16_shape = (*expected_prefix_shape, 16 * trellis_bits) + expected_i32_shape = (*expected_prefix_shape, 8 * trellis_bits)🤖 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 1565 - 1566, Update the shape construction in the code defining expected_i16_shape and expected_i32_shape to use tuple unpacking rather than concatenating expected_prefix_shape with single-element tuples, while preserving both resulting dimensions and existing variable names.Source: Linters/SAST tools
sparkinfer/moe/_shared/kernels/w4a16/route_pack.py (1)
410-467: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe count + prefix launch pair is duplicated verbatim across the two branches.
The only difference between the
expert_counts is not Nonebranch and the eager branch is whereexpert_countscomes from (caller workspace vs.torch.zerosscratch). Resolving the tensor first and sharing one launch sequence removes ~25 duplicated lines and keeps the two paths from drifting.♻️ Sketch
- if expert_counts is not None: - expert_counts = _workspace_slice(...) - expert_counts.zero_() - <count kernel> - <prefix kernel> - elif not torch.cuda.is_current_stream_capturing(): - expert_counts = torch.zeros(...) - <count kernel> - <prefix kernel> - else: - <single-CTA fallback> + if expert_counts is not None: + expert_counts = _workspace_slice( + expert_counts, + name="expert_counts", + elements=int(num_experts), + dtype=torch.int32, + device=topk_ids.device, + ) + expert_counts.zero_() + elif not torch.cuda.is_current_stream_capturing(): + expert_counts = torch.zeros( + int(num_experts), dtype=torch.int32, device=topk_ids.device + ) + else: + expert_counts = None + if expert_counts is not None: + <count kernel> + <prefix kernel> + else: + <single-CTA fallback>🤖 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/route_pack.py` around lines 410 - 467, Refactor the branching around expert_counts so it first selects or allocates the appropriate tensor—reusing _workspace_slice for the caller-provided path and torch.zeros for the eager path—then executes the shared count and prefix kernel launches once. Preserve the existing CUDA graph capture fallback and expert_counts initialization behavior while removing the duplicated launch sequence.sparkinfer/moe/fused_moe/_impl.py (2)
6261-6261: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotation no longer matches the keys stored.
topk_sum_launches: dict[int, object]now receives(token_count, ids_dtype, mapped)tuples on the full-rotation branch (Line 6317), matchingTPW4A16Workspace.planned_topk_sum_launches: dict[object, object]. Widen the local annotation todict[object, object]so the two agree.🤖 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` at line 6261, Update the topk_sum_launches annotation in the relevant initialization to dict[object, object] so it supports the tuple keys stored by the full-rotation branch and matches TPW4A16Workspace.planned_topk_sum_launches.
801-806: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated block-size policy set; import the authoritative constant.
(8, 16, 32, 48, 64)restates_W4A16_ALLOWED_ROUTED_SIZESfromsparkinfer/moe/_shared/kernels/w4a16/host.py, which is also whatplan_w4a16_buffersvalidates against._plan_core_workspacealready imports that constant (Line 2505), so caps validation can reuse it rather than keeping a second copy that can drift.♻️ Proposed fix
if self.w4a16_block_size_m is not None: + from sparkinfer.moe._shared.kernels.w4a16.host import ( + _W4A16_ALLOWED_ROUTED_SIZES, + ) + block_size_m = int(self.w4a16_block_size_m) - if block_size_m not in (8, 16, 32, 48, 64): - raise ValueError("w4a16_block_size_m must be one of 8, 16, 32, 48, 64") + if block_size_m not in _W4A16_ALLOWED_ROUTED_SIZES: + raise ValueError( + "w4a16_block_size_m must be one of " + f"{_W4A16_ALLOWED_ROUTED_SIZES}" + )🤖 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 801 - 806, Update the w4a16 block-size validation in the relevant configuration initializer to reuse the authoritative _W4A16_ALLOWED_ROUTED_SIZES constant imported from the shared host module, instead of duplicating (8, 16, 32, 48, 64). Preserve the existing integer normalization, validation error, and attribute assignment behavior.sparkinfer/_lib/intrinsics.py (1)
6171-6174: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConstants are declared but the PTX re-hardcodes the same values.
_TRELLIS_MCG,_TRELLIS_MASK, and_TRELLIS_ORare never read; both decoders inline0xCBAC1FED/0x8fff8fff/0x3b603b60literally (Lines 6200, 6209-6216, 6259, 6268-6275). Two sources of truth for the codebook constants can drift silently. Either interpolate them into the asm template or drop the constants.♻️ Interpolate instead of duplicating
- mov.b32 M, 0xCBAC1FED; + mov.b32 M, __MCG__;plus a
.replace("__MCG__", f"{_TRELLIS_MCG:#010x}")(and equivalents for mask/or) alongside the existing__B*__substitutions.🤖 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/_lib/intrinsics.py` around lines 6171 - 6174, Use the declared _TRELLIS_MCG, _TRELLIS_MASK, and _TRELLIS_OR constants in both TRELLIS decoder PTX templates instead of hard-coded literals. Add substitutions alongside the existing __B*__ replacements, formatting each value as a zero-padded hexadecimal literal, and replace the corresponding PTX placeholders so these constants remain the single source of truth.sparkinfer/moe/_shared/kernels/w4a16/kernel.py (1)
5552-5590: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the duplicated
_had128_quadinto one shared helper.Both kernels carry byte-identical butterfly +
1/sqrt(128)code, and they must stay bit-identical for the FC1/activation rotation and the top-k-sum inverse rotation to compose. A single module-level@cute.jitfunction (called from both) removes the drift risk and keeps the explanatory comment in one place.Also applies to: 6543-6577
🤖 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/kernel.py` around lines 5552 - 5590, Extract the byte-identical butterfly and scaling logic from both `_had128_quad` implementations into one module-level `@cute.jit` helper, preserving its four-value return signature and explanatory comment. Update each kernel’s `_had128_quad` call site to invoke the shared helper, including the corresponding implementation around the second referenced region, so FC1/activation and inverse rotations remain bit-identical.
🤖 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_fused_moe_trellis.py`:
- Around line 751-758: Correct the test explanation around fused_moe.bind and
the torch.cuda.graph capture block: bind is currently executed outside capture,
while binding.run() is captured. Update the docstring to describe only the
operations and values actually captured, or move fused_moe.bind into the capture
block if the intended contract is to capture binding; keep the test’s stated
failure coverage consistent with the chosen behavior.
---
Outside diff comments:
In `@sparkinfer/moe/_shared/kernels/w4a16/host.py`:
- Around line 364-379: Update the W4A16 buffer allocation around
W4A16PackedBuffers so intermediate_cache13 and intermediate_cache2 use fp16 when
full_rotation is enabled, matching the arena planner and kernel contract; retain
the caller dtype otherwise and leave output’s existing dtype selection
unchanged.
---
Nitpick comments:
In `@sparkinfer/_lib/intrinsics.py`:
- Around line 6171-6174: Use the declared _TRELLIS_MCG, _TRELLIS_MASK, and
_TRELLIS_OR constants in both TRELLIS decoder PTX templates instead of
hard-coded literals. Add substitutions alongside the existing __B*__
replacements, formatting each value as a zero-padded hexadecimal literal, and
replace the corresponding PTX placeholders so these constants remain the single
source of truth.
In `@sparkinfer/moe/_shared/kernels/w4a16/kernel.py`:
- Around line 5552-5590: Extract the byte-identical butterfly and scaling logic
from both `_had128_quad` implementations into one module-level `@cute.jit`
helper, preserving its four-value return signature and explanatory comment.
Update each kernel’s `_had128_quad` call site to invoke the shared helper,
including the corresponding implementation around the second referenced region,
so FC1/activation and inverse rotations remain bit-identical.
In `@sparkinfer/moe/_shared/kernels/w4a16/prepare.py`:
- Around line 1535-1540: Move the words_per_bit assignment below the dtype
whitelist check in the trellis3_t256 validation logic, keeping the existing
int16/int32 rejection behavior and assignment values unchanged.
- Around line 1565-1566: Update the shape construction in the code defining
expected_i16_shape and expected_i32_shape to use tuple unpacking rather than
concatenating expected_prefix_shape with single-element tuples, while preserving
both resulting dimensions and existing variable names.
In `@sparkinfer/moe/_shared/kernels/w4a16/route_pack.py`:
- Around line 410-467: Refactor the branching around expert_counts so it first
selects or allocates the appropriate tensor—reusing _workspace_slice for the
caller-provided path and torch.zeros for the eager path—then executes the shared
count and prefix kernel launches once. Preserve the existing CUDA graph capture
fallback and expert_counts initialization behavior while removing the duplicated
launch sequence.
In `@sparkinfer/moe/fused_moe/_impl.py`:
- Line 6261: Update the topk_sum_launches annotation in the relevant
initialization to dict[object, object] so it supports the tuple keys stored by
the full-rotation branch and matches TPW4A16Workspace.planned_topk_sum_launches.
- Around line 801-806: Update the w4a16 block-size validation in the relevant
configuration initializer to reuse the authoritative _W4A16_ALLOWED_ROUTED_SIZES
constant imported from the shared host module, instead of duplicating (8, 16,
32, 48, 64). Preserve the existing integer normalization, validation error, and
attribute assignment behavior.
In `@tests/moe/test_fused_moe_trellis.py`:
- Around line 120-124: Update
test_prepare_weights_rejects_non_mcg_before_cuda_work and _prepare_weights so
the non-MCG codebook reaches fused_moe.prepare_weights or plan_weights, then
assert the production API raises the expected NotImplementedError. Remove the
helper’s own preemptive rejection, or otherwise avoid asserting that helper-only
check as coverage of PreparedW4A16MoeWeights trellis_codebook gating.
🪄 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: 059a4174-33de-4984-af8a-881dd4496d8f
📒 Files selected for processing (15)
sparkinfer/__init__.pysparkinfer/_lib/intrinsics.pysparkinfer/gemm/__init__.pysparkinfer/gemm/trellis_linear/__init__.pysparkinfer/gemm/trellis_linear/api.pysparkinfer/moe/__init__.pysparkinfer/moe/_shared/execution.pysparkinfer/moe/_shared/kernels/w4a16/host.pysparkinfer/moe/_shared/kernels/w4a16/kernel.pysparkinfer/moe/_shared/kernels/w4a16/prepare.pysparkinfer/moe/_shared/kernels/w4a16/route_pack.pysparkinfer/moe/fused_moe/__init__.pysparkinfer/moe/fused_moe/_impl.pytests/gemm/test_trellis_linear.pytests/moe/test_fused_moe_trellis.py
🚧 Files skipped from review as they are similar to previous changes (2)
- sparkinfer/moe/init.py
- sparkinfer/init.py
|
The requested #49 redesign is already ported here: Trellis now uses the existing fused_moe plan/prepare/scratch/bind/run lifecycle rather than a parallel public API or arena. I preserved your authorship on the implementation commit. Please do not duplicate the port on #49; review and any follow-up should target #90. The latest commit also closes the current CodeRabbit correctness/capture findings and passes the SM120 layer and CUDA-graph tests. |
|
Final full-checkpoint validation is complete.
The E2E run found one real planning-contract bug: the scheduler planned Paired vLLM consumer: local-inference-lab/vllm#190 |
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/fused_moe/_impl.py (1)
6724-6975: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMake
build_tp_moe_fp4_bindingreject non-Trellis expert maps too.The public
.bind(...)path already raises whenroute_expert_map/output_expert_mapare supplied withoutplan.full_rotation, butbuild_tp_moe_fp4_binding()still accepts those kwargs forTPMoEWorkspacePool,TPDynamicWorkspace,TPMicroWorkspace, and non-full-rotationTPW4A16Workspace, then stores them onTPMoEFP4Binding. Add the same guard before buildingcommon_kwargs, or restrict these kwargs to the full-rotation branches.🤖 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 6724 - 6975, Update build_tp_moe_fp4_binding to reject route_expert_map or output_expert_map unless the workspace is a full-rotation TPW4A16Workspace (Trellis) before constructing common_kwargs. Preserve the existing full-rotation validation and prevent non-Trellis workspaces—TPMoEWorkspacePool, TPDynamicWorkspace, TPMicroWorkspace, and non-full-rotation TPW4A16Workspace—from storing these maps on TPMoEFP4Binding.
🤖 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/fused_moe/_impl.py`:
- Around line 6724-6975: Update build_tp_moe_fp4_binding to reject
route_expert_map or output_expert_map unless the workspace is a full-rotation
TPW4A16Workspace (Trellis) before constructing common_kwargs. Preserve the
existing full-rotation validation and prevent non-Trellis
workspaces—TPMoEWorkspacePool, TPDynamicWorkspace, TPMicroWorkspace, and
non-full-rotation TPW4A16Workspace—from storing these maps on TPMoEFP4Binding.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: f259bd2d-d43a-4ba9-84de-d94cec4683f3
📒 Files selected for processing (10)
sparkinfer/_lib/intrinsics.pysparkinfer/gemm/trellis_linear/api.pysparkinfer/moe/_shared/kernels/w4a16/kernel.pysparkinfer/moe/_shared/kernels/w4a16/prepare.pysparkinfer/moe/_shared/kernels/w4a16/route_pack.pysparkinfer/moe/fused_moe/_impl.pytests/gemm/test_trellis_linear.pytests/moe/test_fused_moe_trellis.pytests/moe/test_tp_moe_scratch_bindings.pytests/moe/test_w4a16_route_pack.py
🚧 Files skipped from review as they are similar to previous changes (1)
- sparkinfer/moe/_shared/kernels/w4a16/prepare.py
3 participants
Summary
Integrate EXL3 Trellis/MCG as a W4A16 weight format under the existing planned
sparkinfer.moe.fused_moeAPI.This supersedes #49's separate public
sparkinfer.moe.trellis_moesurface while retaining its Trellis intrinsics, native weight preparation, full-rotation kernel path, graph-safe route packing, and dense kernel support.Why this shape
trellis3_t256,nf3_2p1, packed FP4, and ModelOpt NVFP4 all execute through the same fused-TP W4A16 strategy. The public API should therefore vary the weight recipe, not create one MoE operator per storage format.The integration now uses one lifecycle:
plan_weights(source_format="exl3_trellis_mcg", quant_modes="w4a16", ...)prepare_weights(...)Caps(..., quant_mode="w4a16")plan(...)scratch_specs()/bind(...)run(...)Trellis and NF3 share
PreparedW4A16MoeWeights,TPMoEScratchPlan, the caller-owned scratch arena, cache clearing, and activation dispatch. No second arena or parallel bind/run implementation remains.Functional changes
trellis3_t256full-rotation W4A16 execution with FP32 route accumulation.sparkinfer.gemm.trellis_linear; keep_shared.kernels.w4a16private.block_m; it must not silently invalidate the planned route and expert-block capacities.Review decomposition
Independent changes from #49 were split out:
This PR contains only the Trellis compute and unified MoE/dense API work.
Validation
On the pinned CUDA 13.2 / SM120 v20-r11 runtime:
sparkinfer.moe.fused_moe; notrellis_moecompatibility surface is required.H=6144,I=512, executed the unified plan with zero-copy weights, and passed eager plus CUDA graph replay. Independent unified runs were bitwise identical. Against the r11 legacytrellis_moeruntime on the same deterministic input, relative L2 error was2.01e-4, cosine similarity0.999999881, and maximum absolute difference5.65e-5(the comparison spans intervening shared-W4A16 changes onmaster, not only the API wrapper).brandonmusic/GLM-5.2-EXL3-TR3-3.0bpwcheckpoint passed TP4/DCP4/MTP3 load, memory profiling, piecewise/full CUDA graph capture, MTP capture, and a real HTTP generation (EXL3_OK). AtGPU_MEMORY_UTILIZATION=0.90, the run reported 735,232 KV-cache tokens and a 414.7 MiB planned EXL3 prefill arena per rank.block_m=64, while the livem=1024, top-k 8, 256-expert heuristic preferred 48. The binding now carries the planned value into route packing and GEMM; disagreement with a supplied launch fails before execution. A direct GPU check also verified that 48 and 64 compile to distinct kernels, ruling out a compile-cache collision.git diff --checkpass for every changed file.Consumer migration
The paired vLLM #190 update changes its EXL3 backend from
trellis_moeto the unified weight-plan and fused-MoE API.trellis_moewas unreleased, so no deprecated compatibility shim is added.