[Bugfix] Preserve FlashInfer CUTLASS MoE constant storage across weight reload

[RyanClark2k]

Purpose

Fixes the FlashInfer CUTLASS MoE row of RFC #48312 (category 1, storage identity), confirmed on H100 with gpt-oss-20b: rebuilding the experts object replaced gemm1_alpha, gemm1_beta, and gemm1_clamp_limit, 72/72 graph-visible tensors moved, and every capture-time storage expired.

FlashInferExperts.__init__ allocates per-expert constants that apply passes to the kernel on every forward call: the SwiGLU parameters gemm1_alpha, gemm1_beta, and gemm1_clamp_limit, plus fake_input_scale on the mxfp4 + mxfp8 path. The quant methods rebuild the modular kernel, including the experts object, on every process_weights_after_loading pass. After an RL weight update the new experts object owns fresh tensors while a captured CUDA graph keeps reading the capture-time addresses. Once the old storage is reclaimed, replay reads freed memory. The values are config-derived, so the failure is a silent stale read rather than a crash.

The fix follows the same pattern as #48438 (Marlin) and #48539 (Machete): register the tensors on the owning layer with replace_parameter(prefer_copy=True). The first construction registers them; every later rebuild copies into the same storage, so graph-captured addresses stay valid and the rebuilt experts object reads through the layer. The layer is threaded to FlashInferExperts through the existing extra_kwargs mechanism in the fp8, nvfp4, and mxfp4 kernel factories (the same idiom the HUMMING backend already uses), with an assert so a future call site that fails to supply the layer fails loudly instead of silently reintroducing the bug. Quark fp8, quark mxfp4, and INC mxfp4 called these factories without the layer argument and would have missed the fix; they now pass it.

The unquantized path is deliberately unchanged: UnquantizedFusedMoEMethod builds its kernel only on the first process_weights_after_loading call (guarded by is_weight_update), so the experts object is never rebuilt there and the constants cannot rebind.

Notes for reviewers

• The constants become registered Parameters, so each FlashInfer CUTLASS MoE layer gains up to four float32 entries of num_local_experts elements in named_parameters and state_dict. TrtLlmNvFp4ExpertsBase already registers parameters with the same names on its layers, so the naming follows existing precedent, and only one experts backend is ever active per layer.
• The constants are generated during weight processing and never loaded from checkpoints, so they must not count toward layerwise reload's load_numel_total. The accounting test covers this.
• replace_parameter(prefer_copy=True) falls back to re-registering when the existing parameter is incompatible (shape, dtype, or device mismatch). That cannot happen here because the values are config-derived and the config is constant across reloads. A fail-closed replace_parameter for all call sites is part of the [RFC] Fail-Closed Graph Storage Contract for Weight Reload #48478 discussion.
• FlashInferExperts constructed without a layer (direct construction in kernel tests and benchmarks) keeps the previous behavior: plain tensors, which is fine because nothing survives a rebuild in that usage.

Test Plan

Unit tests in tests/model_executor/model_loader/test_reload.py, mirroring the Marlin tests from #48438:

• test_flashinfer_cutlass_moe_constants_preserve_addresses, parametrized over three variants: mxfp4 weights with bf16 activations (the gpt-oss path), mxfp4 with mxfp8 activations (adds fake_input_scale), and fp8 with a config-provided clamp limit. Asserts stable data_ptr across an experts rebuild, Parameter registration, that the rebuilt experts object reads the preserved storage, and correct values.
• test_flashinfer_cutlass_moe_layerwise_reload_accounting: records reload metadata at construction, processes the layer (which builds the experts object and registers the constants), initializes layerwise reload, asserts load_numel_total equals the checkpoint-loadable numel, streams a new checkpoint, and asserts the layer finalizes during streaming and the constants keep their identity through the reload cycle.

I verified both tests discriminate. With registration disabled (simulating unfixed behavior) all three address variants and the accounting test fail. With metadata recorded after processing instead of before, the accounting total inflates from 768 to 792 (three constants times eight experts) and the first assert fails.

H100 NVL validation on gpt-oss-20b with moe_backend="flashinfer_cutlass" (the mxfp4 + bf16 SM90 path), CUDA graphs enabled, vLLM wheel built at the base commit with only these files overlaid. Protocol: capture, census every constant (data_ptr, value sum, registration state), greedy baseline, reload_weights with the same checkpoint, second census, generation before allocation pressure, pressure, generation after. 24 MoE layers times 3 constants = 72 graph-visible tensors.

• Red (stock): 0/72 constants registered on layers, 0/72 graph-visible pointers stable across reload (72/72 moved). Every experts rebuild abandons the storage captured graphs read; the freed capture-time allocations retain their old bits until reclaimed, which is why this fails silently rather than crashing.
• Green (this fix): 72/72 registered at first load, 72/72 graph-visible pointers stable across reload, 72/72 values intact after allocation pressure, 72/72 still registered after reload.

Note on the reload flow: layerwise reload's copy-back is what carries the storage. The experts object rebuilt during process_weights_after_loading transiently binds the parameter registered before copy-back re-registers the original; both hold identical values, captured graphs read the preserved original, and only eager calls read the transient, so this is benign. The same-weights token comparison diverges after reload on both red and green builds equally, which is a sibling reload defect in this model outside the constants this PR fixes (tracked by the other #48312 rows); the pointer and value census above is the oracle for this change.

Test Result

• All 10 marlin and flashinfer tests in test_reload.py pass locally.
• ruff lint and format clean on all changed files.
• H100 red/green: see numbers above.

Related: RFC #48312 (FlashInfer CUTLASS MoE confirmed row), #48478 (fail-closed storage contract), #49789 (the transaction draft lists FlashInferExperts on its known-raw worklist; this fixes that entry on main).

This PR includes AI-generated code (Claude Code). I reviewed every changed line, and validated the fix end-to-end on H100 hardware with the red/green protocol described above.

[RyanClark2k]

Attaching the red/green harness and full results for the H100 validation described in the test plan.

Environment: H100 NVL (SM90), vLLM wheel built at base commit 4689c7d with only this PR's files overlaid for the green run, flashinfer-python 0.6.15.post1 with flashinfer-jit-cache, gpt-oss-20b, moe_backend="flashinfer_cutlass" (the mxfp4 + bf16 SM90 path selects FlashInferExperts), CUDA graphs enabled (default). 24 MoE layers times 3 SwiGLU constants = 72 graph-visible tensors.

Metric	red (stock main)	green (this PR)
Registered on layer at first load	0/72	72/72
Graph-visible pointers stable across reload	0/72	72/72
Registered after reload	0/72	72/72
Values intact at graph-visible address after allocation pressure	not measurable, see below	72/72

How to read the red row: the census can only reach tensors that are still referenced. On red the capture-time storage has no remaining reference after the experts rebuild, so the reported value fields there describe the new live tensors, not the freed allocations the captured graphs still read. The pointer diff is the oracle, which is consistent with the RFC's note that stale allocations can retain bit-identical contents until reclaimed.

On the token comparison: a same-weights reload_weights changes greedy outputs on both red and green builds equally, before any allocation pressure, with this PR's constants provably intact on green. That points at a separate reload defect on this model path (candidates are the rebuilt prepare/finalize object and the other open #48312 rows), which is why the pointer and value census is the oracle for this change rather than output comparison. I plan to follow up on that separately.

Raw results

result_red.json (moved_keys truncated to the first 6 by the harness):

{
  "label": "red",
  "total": 72,
  "stable": 0,
  "moved": 72,
  "values_preserved": 72,
  "values_intact_post_pressure": 72,
  "registered_at_first_load": 0,
  "registered_after_reload": 0,
  "tokens_match_pre_pressure": false,
  "tokens_match_post_pressure": false,
  "replay_error": null,
  "moved_keys": [
    "model.layers.0.mlp.experts.routed_experts.gemm1_alpha",
    "model.layers.0.mlp.experts.routed_experts.gemm1_beta",
    "model.layers.0.mlp.experts.routed_experts.gemm1_clamp_limit",
    "model.layers.1.mlp.experts.routed_experts.gemm1_alpha",
    "model.layers.1.mlp.experts.routed_experts.gemm1_beta",
    "model.layers.1.mlp.experts.routed_experts.gemm1_clamp_limit"
  ]
}

result_green.json:

{
  "label": "green",
  "total": 72,
  "stable": 72,
  "moved": 0,
  "values_preserved": 72,
  "values_intact_post_pressure": 72,
  "registered_at_first_load": 72,
  "registered_after_reload": 72,
  "tokens_match_pre_pressure": false,
  "tokens_match_post_pressure": false,
  "replay_error": null,
  "moved_keys": []
}

Validation harness

Run as python validate_flashinfer_cutlass_constants.py --label red|green with VLLM_ALLOW_INSECURE_SERIALIZATION=1 (needed for apply_model to pass the census callables to the workers). Red is the stock wheel, green is the same wheel with this PR's files overlaid.

# SPDX-License-Identifier: Apache-2.0
"""H100 red/green validation for the FlashInfer CUTLASS MoE constants fix.

Protocol (RFC #48312 category 1):
  1. Load gpt-oss-20b with the flashinfer_cutlass MoE backend and CUDA graphs
     enabled (default). Capture bakes the constant addresses read by apply().
  2. Census the graph-visible constant storage. The graph reads whatever the
     experts object referenced at capture time: on green that tensor is a
     layer-registered Parameter, on red an unregistered experts attribute.
     Primary pointer = layer parameter when registered, experts attribute
     otherwise.
  3. Greedy baseline generation.
  4. reload_weights with the same checkpoint (rebuilds experts objects).
  5. Census again, diff the graph-visible pointers and values.
  6. Generate before allocation pressure, then apply pressure so freed
     capture-time storage gets reused, then generate again.

Usage:
  python validate_flashinfer_cutlass_constants.py --label red|green
"""

import argparse
import json

MODEL = "openai/gpt-oss-20b"
CONST_NAMES = ("gemm1_alpha", "gemm1_beta", "gemm1_clamp_limit", "fake_input_scale")
PROMPTS = [
    "The capital of France is",
    "In a shocking turn of events, the",
    "def fibonacci(n):",
    "The three laws of thermodynamics are",
]


def census_model(model):
    import torch

    out = {}
    for name, module in model.named_modules():
        quant_method = getattr(module, "quant_method", None)
        if quant_method is None:
            continue
        kernel = getattr(quant_method, "moe_kernel", None)
        if kernel is None:
            continue
        experts = getattr(kernel, "fused_experts", None)
        if experts is None or type(experts).__name__ != "FlashInferExperts":
            continue
        for attr in CONST_NAMES:
            tensor = getattr(experts, attr, None)
            if not isinstance(tensor, torch.Tensor):
                continue
            layer_param = getattr(module, attr, None)
            registered = isinstance(layer_param, torch.nn.Parameter)
            # The graph-visible storage: the layer Parameter when the fix is
            # active, the unregistered experts attribute otherwise.
            visible = layer_param if registered else tensor
            out[f"{name}.{attr}"] = {
                "ptr": visible.data_ptr(),
                "sum": float(visible.float().sum().item()),
                "experts_ptr": tensor.data_ptr(),
                "registered": registered,
            }
    return out


def pressure_model(model):
    import torch

    # Allocate, dirty, and free large blocks so the caching allocator reuses
    # any storage freed by the reload before graph replay.
    blocks = [
        torch.empty(128 * 1024 * 1024, dtype=torch.bfloat16, device="cuda")
        for _ in range(16)
    ]
    for block in blocks:
        block.fill_(1.0)
    del blocks
    torch.cuda.synchronize()
    return True


def generate(llm):
    from vllm import SamplingParams

    params = SamplingParams(temperature=0.0, max_tokens=32)
    outputs = llm.generate(PROMPTS, params)
    return [list(o.outputs[0].token_ids) for o in outputs]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--label", required=True, choices=["red", "green"])
    args = parser.parse_args()
    label = args.label

    from vllm import LLM

    llm = LLM(
        model=MODEL,
        moe_backend="flashinfer_cutlass",
        max_model_len=2048,
        enable_prefix_caching=False,
        gpu_memory_utilization=0.85,
    )

    census1 = llm.apply_model(census_model)[0]
    assert census1, "no FlashInferExperts constants found; wrong MoE backend?"
    registered1 = [k for k in census1 if census1[k]["registered"]]
    print(f"[{label}] census1: {len(census1)} graph-visible constant tensors")
    print(f"[{label}] census1 layer-registered: {len(registered1)}/{len(census1)}")

    tokens_before = generate(llm)
    print(f"[{label}] baseline generation done")

    llm.collective_rpc("reload_weights", kwargs={"weights_path": MODEL})
    print(f"[{label}] reload_weights done")

    census2 = llm.apply_model(census_model)[0]
    assert census2.keys() == census1.keys(), (
        f"census key mismatch: {census1.keys() ^ census2.keys()}"
    )

    stable = [k for k in census1 if census2[k]["ptr"] == census1[k]["ptr"]]
    moved = [k for k in census1 if census2[k]["ptr"] != census1[k]["ptr"]]
    value_ok = [k for k in census1 if census2[k]["sum"] == census1[k]["sum"]]
    registered2 = [k for k in census2 if census2[k]["registered"]]

    tokens_after_reload = generate(llm)
    match_pre_pressure = tokens_after_reload == tokens_before

    llm.apply_model(pressure_model)
    replay_error = None
    tokens_after_pressure = None
    try:
        tokens_after_pressure = generate(llm)
    except Exception as exc:  # noqa: BLE001
        replay_error = repr(exc)
    match_post_pressure = tokens_after_pressure == tokens_before

    # Values at the graph-visible addresses after pressure: stale storage that
    # was reclaimed shows up here as changed sums.
    census3 = llm.apply_model(census_model)[0]
    value_ok_post = [k for k in census1 if census3[k]["sum"] == census1[k]["sum"]]

    total = len(census1)
    print(f"\n===== RESULT [{label}] =====")
    print(f"tensors:                        {total}")
    print(f"graph-visible pointers stable:  {len(stable)}/{total}")
    print(f"graph-visible pointers moved:   {len(moved)}/{total}")
    print(f"values preserved after reload:  {len(value_ok)}/{total}")
    print(f"values intact after pressure:   {len(value_ok_post)}/{total}")
    print(f"layer-registered before/after:  {len(registered1)}/{total} -> "
          f"{len(registered2)}/{total}")
    print(f"tokens match (pre-pressure):    {match_pre_pressure}")
    print(f"tokens match (post-pressure):   {match_post_pressure}")
    if replay_error:
        print(f"replay error: {replay_error}")
    if moved:
        print(f"sample moved: {moved[:3]}")

    with open(f"result_{label}.json", "w") as f:
        json.dump(
            {
                "label": label,
                "total": total,
                "stable": len(stable),
                "moved": len(moved),
                "values_preserved": len(value_ok),
                "values_intact_post_pressure": len(value_ok_post),
                "registered_at_first_load": len(registered1),
                "registered_after_reload": len(registered2),
                "tokens_match_pre_pressure": match_pre_pressure,
                "tokens_match_post_pressure": match_post_pressure,
                "replay_error": replay_error,
                "moved_keys": moved[:6],
            },
            f,
            indent=2,
        )
    print(f"wrote result_{label}.json")


if __name__ == "__main__":
    main()

[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.

[mergify]

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

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