Purpose

Co-authored by: ErinYin liang.yin@daocloud.io

Humming grouped MoE currently retains permutation scratch in every layer, so its persistent memory grows with the number of layers. Marlin lock buffers are also owned by individual layers or kernels, tying their lifetime to weight preparation and reload.

This PR gives both resources a shared owner in WorkspaceManager:

• Add get_persistent_resource(key, factory) for objects with stable storage and get_persistent(...) for fixed-shape tensors. Resources are cached separately for each (ubatch, lane) and must exist before workspace locking.
• Share compatible Humming MoEPermuteScratch objects across sequential layers. Compatibility includes capacity, routing dimensions, device, hidden size, and dtype.
• Acquire Marlin locks at execution time, with separate buffers for concurrent CUDA streams as well as ubatches and lanes. Zero-initialize once; successful Marlin calls restore their counters to zero.
• Resolve dense Marlin workspace lookup inside a runtime custom op so torch.compile does not trace execution-context selection.
• Preserve standalone allocation when no workspace manager exists and support explicit Marlin workspaces.

Relationship to existing work

This combines the direction of #50626 and #56764 with a general resource API. #50626 is the earlier tensor-only Marlin implementation; #56764 adds a Humming-specific cache and accessor to WorkspaceManager. Here, the manager handles resource lifetime and execution-slot isolation, while each backend defines its own resource and compatibility key. The intent is to consolidate both changes around that shared API.

Addresses the memory issue in #56759 and relates to the runtime-storage ownership discussion in #48312. Searches for those issues and for marlin workspace, workspace persistent, and humming scratch found the related PRs above; the distinction is the combined, backend-independent resource API.

#55105 remains complementary: its dispatch-capacity and runtime-dtype fixes can feed into the shared scratch accessor, while its activation-buffer fixes stay in Humming. Those fixes are not included here.

Test Plan

Tests cover resource reuse and locking, standalone fallback, isolation across execution slots and CUDA streams, weight-preparation behavior, and compiled Marlin execution. Permutation tests compare outputs with and without persistent scratch; CUDA graph tests change inputs, routing, and graph sizes to detect stale metadata or aliasing.

Latest targeted runs on NVIDIA B300:

# 488 passed
chg run -g=1 -- .venv/bin/python -m pytest \
  tests/kernels/moe/test_moe_permute_unpermute.py \
  tests/kernels/moe/test_moe.py tests/v1/worker/test_workspace.py \
  -k 'permute or humming or workspace or persistent' --timeout=180 -q

# 20 passed
chg run -g=1 -- .venv/bin/python -m pytest \
  tests/kernels/moe/test_moe_permute_unpermute.py::test_moe_permute_scratch_isolated_across_execution_slots \
  tests/quantization/test_humming_mxfp4_block_fp8.py \
  tests/quantization/test_humming_ignore.py --timeout=180 -q

# 1 passed; run with FlashInfer 0.6.17 pinned via PYTHONPATH
chg run -g=1 -- .venv/bin/python -m pytest \
  tests/kernels/moe/test_moe_permute_unpermute.py::test_moe_permute_scratch_reused_across_graph_sizes \
  --timeout=180 -q

# 112 passed, 13 skipped
chg run -g=1 -- .venv/bin/python -m pytest \
  tests/kernels/quantization/test_marlin_tile_padding.py \
  tests/v1/worker/test_workspace.py \
  tests/model_executor/model_loader/test_reload.py \
  -k 'marlin or workspace' --timeout=180 -q

# 1 passed
chg run -g=1 -- .venv/bin/python -m pytest \
  tests/kernels/quantization/test_marlin_gemm.py -k persistent_locks -q

.venv/bin/pre-commit run
git diff --check

All applicable pre-commit hooks passed, including Ruff, formatting, mypy, and repository checks.

Before publication, the commit was rebased onto f3aa88d230. The only conflict was an obsolete Marlin workspace docstring; the implementation was preserved. Hooks passed again, followed by 60 passed in this focused run (FlashInfer 0.6.17 pinned via PYTHONPATH):

chg run -g=1 -- .venv/bin/python -m pytest \
  tests/v1/worker/test_workspace.py \
  tests/model_executor/model_loader/test_reload.py \
  tests/kernels/moe/test_moe.py \
  tests/kernels/moe/test_moe_permute_unpermute.py \
  tests/kernels/quantization/test_marlin_gemm.py \
  -k 'persistent or humming or scratch or marlin_prepare_layer or marlin_post_load or marlin_workspace' \
  --timeout=180 -q

Test Result

Humming memory

Before the rebase, compared clean base 35dc273072ff133743bc5fcfa306bf030fa8730b with this change using the same compiled extensions and dependency versions. Model: RedHatAI/Qwen3.6-35B-A3B-NVFP4, revision a1ff8c8967426709e5528307c14c52c17616b008 (40 MoE layers, hidden size 2048, 256 experts, top-k 8).

Configuration: one B300, TP=1/DP=1, BF16, Humming for linear and MoE, VLLM_HUMMING_MOE_GEMM_TYPE=grouped. Worker probes confirmed all 40 layers used HummingGroupedExperts. Fixed 8 GiB KV cache, max model length 8192, max sequences 32, async scheduling disabled, seed 0, V2 runner, compilation and CUDA graphs enabled with capture sizes [1, 2, 4, 8, 16, 32]. PyTorch 2.13.0+cu130, CUDA 13.0, FlashInfer 0.6.17.

Live PyTorch GPU allocation immediately after initialization, in GiB (2^30 bytes):

The number of scratch objects falls from 40 to 1. At 8192 tokens, their storage falls from 10330.684 MiB to 258.267 MiB: 97.5% less scratch memory. Each distinct configuration or concurrently active execution slot still needs separate storage. Fixing the KV-cache budget keeps vLLM from consuming the saving as a larger KV cache.

The 8192-token pair ran sequentially on the same reserved GPU, with fresh processes. The 2048-token baseline and repeated full baseline used another B300 of the same model. These measurements do not establish a latency improvement.

Model evaluation

Full GSM8K with the grouped Humming configuration above, 1319 questions, five shots, greedy sampling, and a 256-output-token limit, using the repository's GSM8K prompt/scoring helpers:

The candidate falls within the observed baseline accuracy range. This does not establish bitwise equivalence or rule out a causal accuracy effect: 690/1319 token sequences match between before/after, while only 669/1319 match between the two unmodified baseline runs. Repeated greedy requests also vary on both versions, including after prefix-cache resets; the cause remains unresolved.

Marlin additionally passed the repository GSM8K serving evaluation on 128 questions with five shots: 119/128 (92.97%), zero invalid answers, using the same model with both backends set to Marlin, max sequences 16, max batched tokens 2048, and graph sizes [1, 2, 4, 8, 16].

Remaining validation

No H200 or end-to-end distributed DeepEP v2/HT/LL, DBO, or DSpark serving run was performed. Expert-map and execution-slot tests cover local behavior and do not substitute for those runs. The model evaluation uses BF16/SILU and does not cover the FP8/SITU cases discussed in #55105.

AI assistance

AI assistance was used for investigation, implementation, testing, and drafting this PR.