vllm.models.minimax_m3.amd.mono

Bind native M3 tensors and batch metadata to ATOM's opt-in mono library.

Functions:

• detach_model_cache –

  All TP ranks detach together after graph teardown, before rebinding.

• prepare_model –

  Initialize on eager warmup, including temporary profiling caches.

detach_model_cache(model, replacement)

All TP ranks detach together after graph teardown, before rebinding.

Source code in vllm/models/minimax_m3/amd/mono.py

def detach_model_cache(model, replacement):
    """All TP ranks detach together after graph teardown, before rebinding."""
    if model._mono is not None:
        require(
            replacement.numel() == 0,
            "detach caches after destroying graphs before binding new storage",
        )
        model._mono.close()
        model._mono = None
        logger.info("ATOM mono library closed for detached model cache")

prepare_model(model, config)

Initialize on eager warmup, including temporary profiling caches.

Source code in vllm/models/minimax_m3/amd/mono.py

def prepare_model(model, config):
    """Initialize on eager warmup, including temporary profiling caches."""
    if not isinstance(get_forward_context().attn_metadata, dict):
        return
    require(
        not torch.cuda.is_current_stream_capturing(),
        "mono must be initialized by an eager forward before graph capture",
    )
    group = get_tp_group().cpu_group
    layer_ids = _collective_stage(
        group,
        "model support",
        lambda: _model_layers(model, config),
        allow_fallback=True,
    )
    if layer_ids is None:
        model._mono_config = None
        return
    attentions = [model.layers[i].self_attn for i in layer_ids]
    ready = all(
        attn.kv_cache.numel() > 0 and attn.indexer.index_cache.kv_cache.numel() > 0
        for attn in attentions
    )
    readiness = [False] * dist.get_world_size(group)
    dist.all_gather_object(readiness, ready, group=group)
    if not all(readiness):
        return
    bindings = _collective_stage(
        group,
        "weight support",
        lambda: layer_specs(model, layer_ids),
        allow_fallback=True,
    )
    if bindings is None:
        model._mono_config = None
        return
    model._mono = M3Mono(model, layer_ids, *bindings)
    for attn in attentions:
        attn._mono_model_ref = weakref.ref(model)