vllm.models.muse_spark.amd.mono

FlyDSL decode monokernel for MuseSpark1x on MI355X (--monokernel).

Modules:

• decoder –

  MuseSpark1x decode monokernel: layers 1-60 of a verify step as FlyDSL launches.

• kernels –

  FlyDSL sources of the Muse Spark decode monokernel (see ../README.md).

• peer –

  Device buffers that every rank of a group maps through IPC handles.

Classes:

• MuseSparkMono –

  Runs decoder layers 1-60 of eligible verify steps through the kernels.

Functions:

• maybe_attach_mono –

  Attach the monokernel to lm.model after the weights are final.

MuseSparkMono

Runs decoder layers 1-60 of eligible verify steps through the kernels.

Methods:

• poll_health –

  Raise if a launch timed out waiting for a peer rank.

• refresh_weights –

  Re-take the views, e.g. after a weight reload replaced parameters.

• try_forward –

  Returns the residual stream after the last layer, or None for stock.

Source code in vllm/models/muse_spark/amd/mono/decoder.py

class MuseSparkMono:
    """Runs decoder layers 1-60 of eligible verify steps through the kernels."""

    def __init__(
        self,
        model: nn.Module,
        kernels: MonoKernels,
        rank: int,
        group: Any,
        device: torch.device,
        dims: MonoDims,
        buckets: Sequence[int],
        fused_attention: dict[int, str],
    ) -> None:
        self.rank = rank
        self.dims = dims
        self.refresh_weights(model)
        self.buckets: dict[int, _Bucket] = {}
        for b in buckets:
            kinds = FUSED_ATTENTION.get(b, ())
            layers = {i: t for i, t in fused_attention.items() if t in kinds}
            fa = None
            if layers:
                scales = {
                    i: float(model.layers[i].self_attn.attn.impl.scale) for i in layers
                }
                fa = _FusedAttention(kernels.attn_frag, layers, scales, b, dims, device)
            if b in K3_ROWS:
                k3_kinds = sorted(
                    {(self.weights[i].next_rope, t) for i, t in layers.items()}
                )
                self.buckets[b] = _K3Bucket(
                    kernels.k3, b, dims, device, group, fa, k3_kinds
                )
            else:
                self.buckets[b] = _K5Bucket(
                    kernels.k5,
                    kernels.k5m,
                    b,
                    dims,
                    device,
                    group,
                    fa,
                    sorted(set(layers.values())),
                )
        self._validated: set[int] = set()
        self._fallbacks: set[str] = set()
        self._engaged: set[int] = set()
        self._steps = 0
        self._health_host = {
            b: torch.zeros_like(buf.health, device="cpu").pin_memory()
            for b, buf in self.buckets.items()
        }
        self._health_event: torch.cuda.Event | None = None

    def refresh_weights(self, model: nn.Module) -> None:
        """Re-take the views, e.g. after a weight reload replaced parameters."""
        self.weights = {
            i: collect_layer_weights(model.layers, i, self.rank, self.dims)
            for i in range(FIRST_LAYER, LAST_LAYER + 1)
        }

    def _step_refusal(
        self, model: nn.Module, positions: torch.Tensor, x: Any
    ) -> str | None:
        if not isinstance(x, torch.Tensor):
            return "input"
        rows = x.shape[0]
        names = [
            model.layers[i].self_attn.attn.layer_name
            for i in range(FIRST_LAYER, LAST_LAYER + 2)
        ]
        reason = batch_refusal(rows, get_forward_context(), names)
        if reason is not None:
            return reason
        if bucket_of(rows) not in self.buckets:
            return f"rows:{rows}"
        if positions.dtype != torch.int64 or positions.shape != (rows,):
            return "positions"
        for name in names[1:]:
            _, _, kv_cache, slots = get_attention_context(name)
            reason = kv_cache_refusal(kv_cache, self.dims.local_kv_heads)
            if reason is not None:
                return reason
            if (
                slots is None
                or slots.dtype != torch.int64
                or slots.shape[0] < rows
                or slots.stride(-1) != 1
            ):
                return "slot-mapping"
        return None

    def try_forward(
        self,
        model: nn.Module,
        positions: torch.Tensor,
        x: Any,
        aux_hidden_states: list[torch.Tensor],
    ) -> Any:
        """Returns the residual stream after the last layer, or None for stock."""
        reason = self._step_refusal(model, positions, x)
        if reason is not None:
            if reason not in self._fallbacks:
                self._fallbacks.add(reason)
                logger.debug("Muse Spark monokernel: stock step (%s)", reason)
            return None
        n, rows = x.shape[0], bucket_of(x.shape[0])
        assert rows is not None
        if n not in self._engaged:
            self._engaged.add(n)
            kernel = "K3" if rows in K3_ROWS else "HA+HM+HB"
            logger.info(
                "Muse Spark monokernel engaged: %d rows, %s S=%d", n, kernel, rows
            )
        layers, aux_layers = model.layers, model.aux_hidden_state_layers
        bucket = self.buckets[rows]
        bucket.epoch.add_(1)

        y = layers[0](positions, x)
        if 1 in aux_layers:
            aux_hidden_states.append(norm_then_gated_residual(y))
        from ..model import norm_residual_then_norm  # model.py imports this package

        attn = layers[FIRST_LAYER].self_attn
        attn_input, stream = norm_residual_then_norm(
            y, attn.pre_qkv_proj_norm, out_dtype=layers[FIRST_LAYER].default_dtype
        )
        qkv, _ = attn.qkv_proj(attn_input)
        q, _, _, g = attn._qk_norm_rope_kv_insert(qkv, positions)
        o = attn.attn(q, None, None)
        self._validate(o, g, stream)
        stream, q, g = bucket.run(
            self, layers, positions, o, g, stream, aux_layers, aux_hidden_states
        )

        last = layers[LAST_LAYER + 1]
        o = last.self_attn.attn(q, None, None)
        y = last.after_attention(stream, last.self_attn.back(o, g))
        if LAST_LAYER + 2 in aux_layers:
            aux_hidden_states.append(norm_then_gated_residual(y))
        return y

    def _validate(
        self, attn_o: torch.Tensor, attn_g: torch.Tensor, stream: torch.Tensor
    ) -> None:
        """Check the stock layer-1 tensors the kernels read, once a row count."""
        rows, d = stream.shape[0], self.dims
        if rows in self._validated:
            return
        for name, t, dtype, cols in (
            ("attn_o", attn_o, torch.bfloat16, d.q_rows),
            ("o-gate", attn_g, torch.bfloat16, d.q_rows),
            ("stream", stream, torch.float32, d.hidden_size),
        ):
            # The o-gate is a column slice of the QKV + O-gate output.
            if (
                t.dtype != dtype
                or tuple(t.shape) != (rows, cols)
                or not (t.stride(1) == 1 if t is attn_g else t.is_contiguous())
            ):
                raise RuntimeError(
                    f"Muse Spark monokernel: {name} is {t.dtype} "
                    f"{tuple(t.shape)} {t.stride()}, expected {dtype} {(rows, cols)}"
                )
        self._validated.add(rows)

    def poll_health(self) -> None:
        """Raise if a launch timed out waiting for a peer rank.

        Reads the health words with an asynchronous copy every
        HEALTH_POLL_INTERVAL calls and checks the previous copy, so the step
        never waits for the device.
        """
        if torch.cuda.is_current_stream_capturing():
            return
        self._steps += 1
        if self._health_event is not None and self._health_event.query():
            self._health_event = None
            for b, host in self._health_host.items():
                if int(host[0]) != 0:
                    records = host[8:].view(self.buckets[b].blocks, 8)
                    stuck = records[records[:, 0] != 0][:8].tolist()
                    raise RuntimeError(
                        f"Muse Spark monokernel: rank {self.rank}, S={b}: a launch "
                        f"timed out waiting for a peer rank {stuck}"
                    )
        if self._health_event is None and self._steps % HEALTH_POLL_INTERVAL == 0:
            for b, buf in self.buckets.items():
                self._health_host[b].copy_(buf.health, non_blocking=True)
            self._health_event = torch.cuda.Event()
            self._health_event.record()

    def close(self) -> None:
        for buf in self.buckets.values():
            buf.peers.close()

_validate(attn_o, attn_g, stream)

Check the stock layer-1 tensors the kernels read, once a row count.

Source code in vllm/models/muse_spark/amd/mono/decoder.py

def _validate(
    self, attn_o: torch.Tensor, attn_g: torch.Tensor, stream: torch.Tensor
) -> None:
    """Check the stock layer-1 tensors the kernels read, once a row count."""
    rows, d = stream.shape[0], self.dims
    if rows in self._validated:
        return
    for name, t, dtype, cols in (
        ("attn_o", attn_o, torch.bfloat16, d.q_rows),
        ("o-gate", attn_g, torch.bfloat16, d.q_rows),
        ("stream", stream, torch.float32, d.hidden_size),
    ):
        # The o-gate is a column slice of the QKV + O-gate output.
        if (
            t.dtype != dtype
            or tuple(t.shape) != (rows, cols)
            or not (t.stride(1) == 1 if t is attn_g else t.is_contiguous())
        ):
            raise RuntimeError(
                f"Muse Spark monokernel: {name} is {t.dtype} "
                f"{tuple(t.shape)} {t.stride()}, expected {dtype} {(rows, cols)}"
            )
    self._validated.add(rows)

poll_health()

Raise if a launch timed out waiting for a peer rank.

Reads the health words with an asynchronous copy every HEALTH_POLL_INTERVAL calls and checks the previous copy, so the step never waits for the device.

Source code in vllm/models/muse_spark/amd/mono/decoder.py

def poll_health(self) -> None:
    """Raise if a launch timed out waiting for a peer rank.

    Reads the health words with an asynchronous copy every
    HEALTH_POLL_INTERVAL calls and checks the previous copy, so the step
    never waits for the device.
    """
    if torch.cuda.is_current_stream_capturing():
        return
    self._steps += 1
    if self._health_event is not None and self._health_event.query():
        self._health_event = None
        for b, host in self._health_host.items():
            if int(host[0]) != 0:
                records = host[8:].view(self.buckets[b].blocks, 8)
                stuck = records[records[:, 0] != 0][:8].tolist()
                raise RuntimeError(
                    f"Muse Spark monokernel: rank {self.rank}, S={b}: a launch "
                    f"timed out waiting for a peer rank {stuck}"
                )
    if self._health_event is None and self._steps % HEALTH_POLL_INTERVAL == 0:
        for b, buf in self.buckets.items():
            self._health_host[b].copy_(buf.health, non_blocking=True)
        self._health_event = torch.cuda.Event()
        self._health_event.record()

refresh_weights(model)

Re-take the views, e.g. after a weight reload replaced parameters.

Source code in vllm/models/muse_spark/amd/mono/decoder.py

def refresh_weights(self, model: nn.Module) -> None:
    """Re-take the views, e.g. after a weight reload replaced parameters."""
    self.weights = {
        i: collect_layer_weights(model.layers, i, self.rank, self.dims)
        for i in range(FIRST_LAYER, LAST_LAYER + 1)
    }

try_forward(model, positions, x, aux_hidden_states)

Returns the residual stream after the last layer, or None for stock.

Source code in vllm/models/muse_spark/amd/mono/decoder.py

def try_forward(
    self,
    model: nn.Module,
    positions: torch.Tensor,
    x: Any,
    aux_hidden_states: list[torch.Tensor],
) -> Any:
    """Returns the residual stream after the last layer, or None for stock."""
    reason = self._step_refusal(model, positions, x)
    if reason is not None:
        if reason not in self._fallbacks:
            self._fallbacks.add(reason)
            logger.debug("Muse Spark monokernel: stock step (%s)", reason)
        return None
    n, rows = x.shape[0], bucket_of(x.shape[0])
    assert rows is not None
    if n not in self._engaged:
        self._engaged.add(n)
        kernel = "K3" if rows in K3_ROWS else "HA+HM+HB"
        logger.info(
            "Muse Spark monokernel engaged: %d rows, %s S=%d", n, kernel, rows
        )
    layers, aux_layers = model.layers, model.aux_hidden_state_layers
    bucket = self.buckets[rows]
    bucket.epoch.add_(1)

    y = layers[0](positions, x)
    if 1 in aux_layers:
        aux_hidden_states.append(norm_then_gated_residual(y))
    from ..model import norm_residual_then_norm  # model.py imports this package

    attn = layers[FIRST_LAYER].self_attn
    attn_input, stream = norm_residual_then_norm(
        y, attn.pre_qkv_proj_norm, out_dtype=layers[FIRST_LAYER].default_dtype
    )
    qkv, _ = attn.qkv_proj(attn_input)
    q, _, _, g = attn._qk_norm_rope_kv_insert(qkv, positions)
    o = attn.attn(q, None, None)
    self._validate(o, g, stream)
    stream, q, g = bucket.run(
        self, layers, positions, o, g, stream, aux_layers, aux_hidden_states
    )

    last = layers[LAST_LAYER + 1]
    o = last.self_attn.attn(q, None, None)
    y = last.after_attention(stream, last.self_attn.back(o, g))
    if LAST_LAYER + 2 in aux_layers:
        aux_hidden_states.append(norm_then_gated_residual(y))
    return y

_ensure_writable_flydsl_cache()

AITER points FlyDSL at its bundled, read-only cache; new kernels need a writable one.

Source code in vllm/models/muse_spark/amd/mono/__init__.py

def _ensure_writable_flydsl_cache() -> None:
    """AITER points FlyDSL at its bundled, read-only cache; new kernels need a
    writable one."""
    cur = os.environ.get("FLYDSL_RUNTIME_CACHE_DIR")
    if cur and os.access(cur, os.W_OK):
        return
    path = Path.home() / ".flydsl" / "cache"
    path.mkdir(parents=True, exist_ok=True)
    os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = str(path)

maybe_attach_mono(lm)

Attach the monokernel to lm.model after the weights are final.

Collective over the TP group: every rank calls it. A deployment that any rank cannot run raises on every rank, with every reason of this rank.

Source code in vllm/models/muse_spark/amd/mono/__init__.py

def maybe_attach_mono(lm: nn.Module) -> None:
    """Attach the monokernel to ``lm.model`` after the weights are final.

    Collective over the TP group: every rank calls it. A deployment that any
    rank cannot run raises on every rank, with every reason of this rank.
    """
    if not lm.vllm_config.kernel_config.monokernel:
        return
    model = lm.model
    if model.mono is not None:
        model.mono.refresh_weights(model)
        return
    group = get_tp_group()
    rank, world_size = group.rank_in_group, group.world_size
    text_config = lm.vllm_config.model_config.hf_text_config
    reasons = static_refusals(lm, rank, world_size)
    kernels = None
    if not reasons:
        kernels, missing = _load_kernels()
        if missing is not None:
            reasons.append(missing)
    verdict = torch.tensor([0 if reasons else 1], dtype=torch.int32)
    dist.all_reduce(verdict, op=dist.ReduceOp.MIN, group=group.cpu_group)
    if int(verdict.item()) == 0:
        raise ValueError(
            "--monokernel: Muse Spark cannot run the monokernel: "
            + ("; ".join(reasons) or "a peer rank refused the deployment")
        )
    assert kernels is not None
    buckets = built_buckets(lm.vllm_config.scheduler_config.max_num_seqs)
    max_model_len = lm.vllm_config.model_config.max_model_len
    fused = fused_attention_layers(
        model.layers, max_model_len, kernels.attn_frag.VLLM_MAX_KEYS
    )
    if not fused:
        logger.info(
            "Muse Spark monokernel: verify attention unfused (max_model_len %d, "
            "fused builds cover %d)",
            max_model_len,
            kernels.attn_frag.VLLM_MAX_KEYS,
        )
    model.mono = MuseSparkMono(
        model,
        kernels,
        rank=rank,
        group=group.cpu_group,
        device=torch.device("cuda", torch.accelerator.current_device_index()),
        dims=MonoDims.from_config(text_config, world_size),
        buckets=buckets,
        fused_attention=fused,
    )
    logger.info(
        "Muse Spark monokernel attached: layers %d-%d of verify steps with "
        "up to %d rows run as FlyDSL launches (builds S=%s; MoE activations "
        "in bf16)",
        FIRST_LAYER,
        LAST_LAYER,
        buckets[-1],
        ",".join(map(str, buckets)),
    )

static_refusals(lm, tp_rank, tp_size)

Every reason this deployment cannot run the monokernel; empty = supported.

Source code in vllm/models/muse_spark/amd/mono/__init__.py

def static_refusals(lm: nn.Module, tp_rank: int, tp_size: int) -> list[str]:
    """Every reason this deployment cannot run the monokernel; empty = supported."""
    if not current_platform.is_rocm():
        return ["not ROCm"]
    from vllm.platforms.rocm import on_gfx950

    reasons = [] if on_gfx950() else ["not gfx950"]
    cus = torch.cuda.get_device_properties(
        torch.accelerator.current_device_index()
    ).multi_processor_count
    if cus != NUM_CUS:
        reasons.append(f"{cus} CUs (kernel: {NUM_CUS})")
    text_config = lm.vllm_config.model_config.hf_text_config
    reasons += config_refusals(text_config)
    if reasons:
        return reasons
    experts = _routed_experts(lm.model.layers[FIRST_LAYER])
    moe_parallel = getattr(
        getattr(experts, "moe_config", None), "moe_parallel_config", None
    )
    reasons = deployment_refusals(lm.vllm_config, tp_rank, tp_size, moe_parallel)
    return reasons or layer_refusals(
        lm.model, MonoDims.from_config(text_config, tp_size)
    )