[Refactor][MoE] Share CUTLASS permute scratch across layers

[xijiaat]

Overview

Use the existing get_moe_permute_scratch helper in the FP8 and W4A8 CUTLASS experts instead of retaining a separate scratch object in every layer. This is the CUTLASS cross-layer sharing follow-up explicitly left out of #60447, not another scratch-capacity fix.

Claims

• Sequential CUTLASS layers with the same scratch configuration reuse one persistent resource within each workspace-manager ubatch/lane.
• Independent lanes retain independent scratch. The helper is called at runtime rather than caching the first lane's resource on the expert object.
• No changes to scratch sizing, MoE arithmetic, or the workspace manager. No throughput or model-quality improvement is claimed.

Validation

Validated before/after on one otherwise idle NVIDIA H20 (SM90). Python source is main 81198e97ba7eee2a22540caaa756b7fdddcb4d93 with/without this patch; compiled CUDA extensions and dependencies are from vllm/vllm-openai:v0.31.0 (PyTorch 2.13.0+cu130), not newly compiled main binaries.

Check	Before	After
Added parameterized pytest (both expert classes)	2 expected sharing-assertion failures	2 passed
Real CUDA scratch, 64 getter owners per class	64 objects; 152,027,584 storage bytes	1 object; 2,375,431 storage bytes
Real permute/unpermute roundtrips, including CUDA Graph replay with changed inputs	6 passed	6 passed
Existing upstream FP8 numerical cases, eager/graph and expert-map partitions	27 passed	27 passed

The scratch measurement is synthetic metadata storage (about 145 MiB to 2.27 MiB), not total memory saved by a loaded model. Same-lane sharing and separate-lane isolation were checked with real CUDA allocations. The 27 numerical cases call existing upstream test functions directly with an initialized workspace manager; this is not a claim that the whole upstream pytest suite passed. Expert-map partition cases run on one GPU, not distributed EP.

• CPU-isolated getter/helper/manager checks also reproduced both sharing failures before the change and passed after it.
• pre-commit run --files vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py tests/kernels/moe/test_cutlass_moe.py: passed, including mypy.
• git diff --check: passed.
• The added test was run in the container with the following command. --confcutdir excludes unrelated root-conftest dependencies missing from the runtime image; the test itself imports the real vLLM modules and mocks the scratch constructor as written in this PR:

.venv/bin/python -m pytest --confcutdir=tests/kernels/moe \
  tests/kernels/moe/test_cutlass_moe.py -q \
  -k test_cutlass_permute_scratch_shared_across_layers

No full-model evaluation, throughput benchmark, multi-GPU communication test, or full W4A8 GEMM-layer numerical validation is claimed. W4A8 GPU coverage is the changed getter, real scratch allocation, and permute/unpermute kernels. The PR remains Draft.

The human submitter confirmed line-by-line review and personal testing of this patch. AI assistance was used for analysis, implementation, and the validation above.

Details

CutlassExpertsFp8Base and CutlassExpertsW4A8Fp8 each lazily allocated self._permute_scratch. Consequently, sequential layers with identical configurations kept separate metadata and sorting buffers for their lifetimes. Both now use the existing persistent-resource helper already used elsewhere in the MoE path. The unsupported-device fallback remains None.

This is distinct from #60447: that PR changes capacity growth and scratch contents but explicitly retains one CUTLASS scratch per layer. This change leaves the capacity calculation untouched. The two changes may require a small interface rebase depending on merge order. Searches for open permute scratch PRs and references to #60447 found no existing CUTLASS cross-layer-sharing patch.

As with the existing helper's contract, a workspace manager must be initialized for sharing, and all required ubatch/lane resources must be warmed before locking. Without a manager the helper allocates per call; standalone callers no longer get the old per-layer cache. The added test is one parameterized test covering both expert classes.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[xijiaat]

@mgoin, could you please take a look when you have time? This addresses the CUTLASS cross-layer scratch sharing follow-up mentioned in #60447, using the existing persistent-resource helper. H20 before/after validation results are included in the description. Happy to rebase on #60447 depending on the preferred merge order.

[mgoin]

Hey @xijiaat thanks for the PR but based on issues encountered in CI in mine, I may need to address CUTLASS paths myself. Stay tuned for the final state of my PR

[xijiaat]

@mgoin Thanks for the heads-up! I'll hold off on further changes and wait for #60447 to settle, then check whether this follow-up is still needed.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @xijiaat.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[xijiaat]

Thanks @mgoin for incorporating the CUTLASS paths in #60447. I checked the merged implementation: both CutlassExpertsFp8Base and CutlassExpertsW4A8Fp8 now use get_moe_permute_scratch rather than per-instance scratch, retaining isolation by workspace ubatch/lane. This covers the production-code changes proposed here, so I am closing this PR as superseded by #60447 rather than rebasing a duplicate change. Thanks for the fix and review!