Instructions to use yuhuanstudio/Yunmo-Next-1B-Instruct with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use yuhuanstudio/Yunmo-Next-1B-Instruct with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="yuhuanstudio/Yunmo-Next-1B-Instruct", trust_remote_code=True)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("yuhuanstudio/Yunmo-Next-1B-Instruct", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use yuhuanstudio/Yunmo-Next-1B-Instruct with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "yuhuanstudio/Yunmo-Next-1B-Instruct"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "yuhuanstudio/Yunmo-Next-1B-Instruct",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/yuhuanstudio/Yunmo-Next-1B-Instruct

SGLang

How to use yuhuanstudio/Yunmo-Next-1B-Instruct with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "yuhuanstudio/Yunmo-Next-1B-Instruct" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "yuhuanstudio/Yunmo-Next-1B-Instruct",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "yuhuanstudio/Yunmo-Next-1B-Instruct" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "yuhuanstudio/Yunmo-Next-1B-Instruct",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use yuhuanstudio/Yunmo-Next-1B-Instruct with Docker Model Runner:
```
docker model run hf.co/yuhuanstudio/Yunmo-Next-1B-Instruct
```

Yunmo-Next-1B-Instruct / code /yunmo_next /static_decode.py

yuhuanowo

Yunmo-Next-1B-Instruct: weights (BF16 + FP32), runtime, transformers remote code, model card

b901794 verified 10 days ago

Raw History Blame Contribute Delete

14.4 kB

"""CUDA-graph single-token decoding for Yunmo-Next.
Cached decoding through the regular modules is launch-bound: a token costs
about 5,200 kernel launches across 60 layers while the GPU is busy for roughly
a tenth of the step.  This module replays the whole step as one CUDA graph.
What changes relative to ``YunmoNextForCausalLM.forward_cached``:
- GQA keys/values live in preallocated buffers.  Global layers use a buffer of
  the bucket length, local layers a 2,048-slot ring (keys are stored after
  RoPE, so slot order does not matter); a mask computed from the position
  tensor selects valid slots.
- KDA convolution caches are FLA's own ``[N, D, W]`` buffers, updated in place
  by ``ShortConvolution``; the recurrent state is updated in place by FLA's
  fused recurrent kernel (``inplace_final_state=True``).
- Depth routing is the model's own cached Attention Residual code, unchanged:
  only each layer's mixer call is redirected to the static versions while the
  graph is captured.
Prefill still runs through ``forward_cached``; ``load`` copies its caches into
the static buffers.  One graph is captured per length bucket and reused.
With ``batch > 1`` each row is an independent sequence with its own position
(continuous batching): ``load(cache, row)`` fills a row from a batch-1
prefill, ``release(row)`` parks it, and ``step(tokens)`` advances every
active row.  A 1B model's decode step is bound by reading the weights and by
per-kernel gaps, both nearly independent of the batch, so rows are almost
free until the GPU saturates.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn.functional as F

from .attention import GatedGroupedQueryAttention
from .protocol import CONTROL_BASE

# Per-row greedy summary columns: lexical max value, its index, any-NaN flag,
# then the logits of the 256 protocol IDs.
SUMMARY_COLUMNS = 3 + 256

BUCKETS = (4096, 16384, 65536, 262144)


@dataclass
class _GQABuffers:
    keys: torch.Tensor
    values: torch.Tensor
    is_global: bool


@dataclass
class _KDABuffers:
    state: torch.Tensor
    conv: tuple[torch.Tensor, torch.Tensor, torch.Tensor]


class StaticDecoder:
    """Captured one-token decode step for one length bucket and batch size."""

    def __init__(self, model, length: int, batch: int = 1) -> None:
        if model.training:
            raise RuntimeError("static decoding requires model.eval()")
        self.model = model
        self.length = length
        self.batch = batch
        config = model.config
        self.window = config.attention_window
        param = next(model.parameters())
        device, dtype = param.device, param.dtype
        self.device = device
        self.token = torch.zeros((batch, 1), dtype=torch.long, device=device)
        self.position = torch.zeros((batch, 1), dtype=torch.long, device=device)
        # Parked rows keep rewriting one slot instead of advancing.
        self.active = torch.zeros((batch, 1), dtype=torch.long, device=device)
        self.buffers: dict[int, _GQABuffers | _KDABuffers] = {}
        for index, layer in enumerate(model.layers):
            mixer = layer.mixer
            if isinstance(mixer, GatedGroupedQueryAttention):
                slots = length if mixer.is_global else self.window
                shape = (batch, config.num_key_value_heads, slots, config.attention_head_dim)
                self.buffers[index] = _GQABuffers(
                    torch.zeros(shape, dtype=dtype, device=device),
                    torch.zeros(shape, dtype=dtype, device=device),
                    bool(mixer.is_global),
                )
            else:
                width = mixer.q_conv1d.weight.shape[-1]
                self.buffers[index] = _KDABuffers(
                    torch.zeros((batch, mixer.num_heads, mixer.head_v_dim, mixer.head_k_dim),
                                dtype=torch.float32, device=device),
                    tuple(
                        torch.zeros((batch, conv.weight.shape[0], width), dtype=dtype, device=device)
                        for conv in (mixer.q_conv1d, mixer.k_conv1d, mixer.v_conv1d)
                    ),
                )
        self.graph: torch.cuda.CUDAGraph | None = None
        self._host_position = [0] * batch
        self._host_active = [False] * batch
        self.logits: torch.Tensor | None = None
        self.summary: torch.Tensor | None = None

    # ---- per-layer static mixers -------------------------------------------------

    def _gqa(self, layer, buffers: _GQABuffers, hidden: torch.Tensor) -> torch.Tensor:
        mixer = layer.mixer
        normalized = layer.input_norm(hidden)
        raw_q, raw_k, raw_v, gate = mixer._linear_project(normalized)
        query, key, value = mixer._shape_projected(raw_q, raw_k, raw_v, self.position)
        batch, heads, slots, width = buffers.keys.shape
        position = self.position.view(batch, 1)
        slot = position if buffers.is_global else torch.remainder(position, slots)
        index = slot.view(batch, 1, 1, 1).expand(batch, heads, 1, width)
        buffers.keys.scatter_(2, index, key)
        buffers.values.scatter_(2, index, value)
        # Ring slots fill in order from 0, so the first min(position + 1, slots)
        # slots hold exactly the keys this query may see.
        valid = torch.arange(slots, device=hidden.device).view(1, slots) <= position
        # One query token: fold each KV head's query heads (SDPA's GQA maps
        # query head h to KV head h // groups) into the query-length axis and
        # attend with two batched matmuls, reading each KV buffer once.  SDPA
        # with ``enable_gqa`` expands KV to every query head (~2 ms per layer
        # at batch 16) and SDPA on the folded shape picks a slow FP32 kernel;
        # measured per layer at 4,096 slots: 0.17-0.32 ms here vs 0.27-3.9 ms.
        # Same scale as the model (SDPA default 1/sqrt(head_dim)).
        query_heads = query.shape[1]
        grouped = query.reshape(batch, heads, query_heads // heads, width)
        scores = torch.matmul(grouped, buffers.keys.transpose(-1, -2)) * width**-0.5
        scores = scores.masked_fill(~valid.view(batch, 1, 1, slots), float("-inf"))
        output = torch.matmul(torch.softmax(scores, dim=-1), buffers.values)
        output = output.reshape(batch, query_heads, 1, width)
        return mixer._apply_output_gate(output, gate)

    def _kda(self, layer, buffers: _KDABuffers, hidden: torch.Tensor) -> torch.Tensor:
        from einops import rearrange
        from fla.ops.kda.fused_recurrent import fused_recurrent_kda_fwd

        mixer = layer.mixer
        x = layer.input_norm(hidden)
        conv_q, conv_k, conv_v = buffers.conv
        q, _ = mixer.q_conv1d(x=mixer.q_proj(x), cache=conv_q, output_final_state=True)
        k, _ = mixer.k_conv1d(x=mixer.k_proj(x), cache=conv_k, output_final_state=True)
        v, _ = mixer.v_conv1d(x=mixer.v_proj(x), cache=conv_v, output_final_state=True)
        g = rearrange(mixer.f_proj(x), "... (h d) -> ... h d", d=mixer.head_k_dim)
        beta = mixer.b_proj(x)
        q, k = (rearrange(t, "... (h d) -> ... h d", d=mixer.head_k_dim) for t in (q, k))
        v = rearrange(v, "... (h d) -> ... h d", d=mixer.head_v_dim)
        output, _ = fused_recurrent_kda_fwd(
            q=q, k=k, v=v, g=g, beta=beta, A_log=mixer.A_log, dt_bias=mixer.dt_bias,
            scale=k.shape[-1] ** -0.5, initial_state=buffers.state, inplace_final_state=True,
            output_final_state=True, use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True,
            use_beta_sigmoid_in_kernel=True, allow_neg_eigval=mixer.allow_neg_eigval,
            lower_bound=mixer.lower_bound, cu_seqlens=None, state_v_first=True,
        )
        gate = rearrange(mixer.g_proj(x), "... (h d) -> ... h d", d=mixer.head_v_dim)
        output = rearrange(mixer.o_norm(output, gate), "b t h d -> b t (h d)")
        return mixer.o_proj(output)

    # ---- step and capture --------------------------------------------------------

    def _step(self) -> torch.Tensor:
        model = self.model
        originals = {}
        for index, layer in enumerate(model.layers):
            buffers = self.buffers[index]
            static = self._gqa if isinstance(buffers, _GQABuffers) else self._kda

            def mixer_output_cached(hidden, position_ids, *, fla_cache, attention_cache,
                                    _layer=layer, _buffers=buffers, _static=static):
                return _static(_layer, _buffers, hidden), None

            originals[index] = layer.__dict__.get("mixer_output_cached")
            layer.mixer_output_cached = mixer_output_cached
        try:
            state = _NullState()
            hidden = model._forward_hidden_two_phase_attention_residual_cached(
                model.embed_tokens(self.token), self.position, state, None, None,
            )
            return F.linear(hidden, model.embed_tokens.weight)[:, -1]
        finally:
            for index, layer in enumerate(model.layers):
                if originals[index] is None:
                    del layer.mixer_output_cached
                else:
                    layer.mixer_output_cached = originals[index]

    def _tensors(self):
        for buffers in self.buffers.values():
            if isinstance(buffers, _GQABuffers):
                yield buffers.keys
                yield buffers.values
            else:
                yield buffers.state
                yield from buffers.conv

    @torch.inference_mode()
    def capture(self, warmup: int = 3) -> None:
        stream = torch.cuda.Stream()
        stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(stream):
            for _ in range(warmup):
                _greedy_summary(self._step())
        torch.cuda.current_stream().wait_stream(stream)
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            self.logits = self._step()
            self.summary = _greedy_summary(self.logits)
        # Warm-up and capture ran real steps on the buffers.  ``load`` rewrites
        # every buffer of the row it fills, so clearing (instead of keeping a
        # snapshot copy, which doubled peak memory) is enough.
        for tensor in self._tensors():
            tensor.zero_()
        self.position.zero_()
        self.active.zero_()
        self.graph = graph

    # ---- runtime -----------------------------------------------------------------

    @torch.inference_mode()
    def load(self, cache, row: int = 0) -> int:
        """Copy a batch-1 ``forward_cached`` prefill cache into ``row``."""

        position = int(cache.seen_tokens)
        if position >= self.length:
            raise ValueError("prompt does not fit this decode bucket")
        for index, buffers in self.buffers.items():
            if isinstance(buffers, _GQABuffers):
                buffers.keys[row].zero_()
                buffers.values[row].zero_()
                entry = cache.attention.get(index)
                if entry is None:
                    continue
                count = entry.keys.shape[2]
                positions = torch.arange(position - count, position, device=self.device)
                slots = positions if buffers.is_global else torch.remainder(positions, self.window)
                buffers.keys[row].index_copy_(1, slots, entry.keys[0])
                buffers.values[row].index_copy_(1, slots, entry.values[0])
            else:
                layer_state = cache.fla_cache[index]
                buffers.state[row].copy_(layer_state["recurrent_state"][0])
                for target, source in zip(buffers.conv, layer_state["conv_state"], strict=True):
                    target[row].copy_(source[0])
        self.position[row].fill_(position)
        self.active[row].fill_(1)
        self._host_position[row] = position
        self._host_active[row] = True
        return position

    def release(self, row: int) -> None:
        """Park ``row``; its outputs are ignored until the next ``load``."""

        self.active[row].fill_(0)
        self.position[row].fill_(0)
        self._host_active[row] = False

    @torch.inference_mode()
    def step(self, tokens) -> torch.Tensor:
        """``advance`` and return a copy of the ``[batch, vocab]`` logits."""

        self.advance(tokens)
        return self.logits.clone()

    @torch.inference_mode()
    def advance(self, tokens) -> None:
        """Advance every active row by one token.
        ``tokens`` is an int (batch 1) or one int per row; parked rows' entries
        are ignored.  Afterwards ``logits`` and ``summary`` hold the new rows
        until the next call.
        """

        if self.graph is None:
            raise RuntimeError("capture() must run before step()")
        if isinstance(tokens, int):
            tokens = [tokens]
        if len(tokens) != self.batch:
            raise ValueError("step needs one token per row")
        # A global-layer write past the buffer would be an out-of-bounds
        # scatter inside the graph; refuse on the host instead.
        for row in range(self.batch):
            if self._host_active[row]:
                if self._host_position[row] >= self.length:
                    raise RuntimeError("decode exceeded the static decoder's length bucket")
                self._host_position[row] += 1
        self.token.copy_(torch.tensor(tokens, dtype=torch.long).view(self.batch, 1), non_blocking=False)
        self.graph.replay()
        self.position.add_(self.active)


def _greedy_summary(logits: torch.Tensor) -> torch.Tensor:
    # Float32 throughout: the index (< 2**24) is exact in float32 but not in
    # float16, and float16 logits widen exactly.
    lexical = logits[:, :CONTROL_BASE]
    index = lexical.argmax(dim=-1, keepdim=True)
    return torch.cat(
        (
            lexical.gather(-1, index).float(),
            index.float(),
            torch.isnan(logits).any(dim=-1, keepdim=True).float(),
            logits[:, CONTROL_BASE:].float(),
        ),
        dim=-1,
    )


class _NullState:
    """Stands in for ``YunmoNextCUDAState``: static mixers ignore it."""

    fla_cache = None
    attention: dict = {}


def bucket_for(length: int) -> int:
    for bucket in BUCKETS:
        if length <= bucket:
            return bucket
    raise ValueError("requested length exceeds the largest decode bucket")

	"""CUDA-graph single-token decoding for Yunmo-Next.

	Cached decoding through the regular modules is launch-bound: a token costs
	about 5,200 kernel launches across 60 layers while the GPU is busy for roughly
	a tenth of the step. This module replays the whole step as one CUDA graph.

	What changes relative to ``YunmoNextForCausalLM.forward_cached``:

	- GQA keys/values live in preallocated buffers. Global layers use a buffer of
	the bucket length, local layers a 2,048-slot ring (keys are stored after
	RoPE, so slot order does not matter); a mask computed from the position
	tensor selects valid slots.
	- KDA convolution caches are FLA's own ``[N, D, W]`` buffers, updated in place
	by ``ShortConvolution``; the recurrent state is updated in place by FLA's
	fused recurrent kernel (``inplace_final_state=True``).
	- Depth routing is the model's own cached Attention Residual code, unchanged:
	only each layer's mixer call is redirected to the static versions while the
	graph is captured.

	Prefill still runs through ``forward_cached``; ``load`` copies its caches into
	the static buffers. One graph is captured per length bucket and reused.

	With ``batch > 1`` each row is an independent sequence with its own position
	(continuous batching): ``load(cache, row)`` fills a row from a batch-1
	prefill, ``release(row)`` parks it, and ``step(tokens)`` advances every
	active row. A 1B model's decode step is bound by reading the weights and by
	per-kernel gaps, both nearly independent of the batch, so rows are almost
	free until the GPU saturates.
	"""

	from __future__ import annotations

	from dataclasses import dataclass

	import torch
	import torch.nn.functional as F

	from .attention import GatedGroupedQueryAttention
	from .protocol import CONTROL_BASE

	# Per-row greedy summary columns: lexical max value, its index, any-NaN flag,
	# then the logits of the 256 protocol IDs.
	SUMMARY_COLUMNS = 3 + 256

	BUCKETS = (4096, 16384, 65536, 262144)


	@dataclass
	class _GQABuffers:
	keys: torch.Tensor
	values: torch.Tensor
	is_global: bool


	@dataclass
	class _KDABuffers:
	state: torch.Tensor
	conv: tuple[torch.Tensor, torch.Tensor, torch.Tensor]


	class StaticDecoder:
	"""Captured one-token decode step for one length bucket and batch size."""

	def __init__(self, model, length: int, batch: int = 1) -> None:
	if model.training:
	raise RuntimeError("static decoding requires model.eval()")
	self.model = model
	self.length = length
	self.batch = batch
	config = model.config
	self.window = config.attention_window
	param = next(model.parameters())
	device, dtype = param.device, param.dtype
	self.device = device
	self.token = torch.zeros((batch, 1), dtype=torch.long, device=device)
	self.position = torch.zeros((batch, 1), dtype=torch.long, device=device)
	# Parked rows keep rewriting one slot instead of advancing.
	self.active = torch.zeros((batch, 1), dtype=torch.long, device=device)
	self.buffers: dict[int, _GQABuffers \| _KDABuffers] = {}
	for index, layer in enumerate(model.layers):
	mixer = layer.mixer
	if isinstance(mixer, GatedGroupedQueryAttention):
	slots = length if mixer.is_global else self.window
	shape = (batch, config.num_key_value_heads, slots, config.attention_head_dim)
	self.buffers[index] = _GQABuffers(
	torch.zeros(shape, dtype=dtype, device=device),
	torch.zeros(shape, dtype=dtype, device=device),
	bool(mixer.is_global),
	)
	else:
	width = mixer.q_conv1d.weight.shape[-1]
	self.buffers[index] = _KDABuffers(
	torch.zeros((batch, mixer.num_heads, mixer.head_v_dim, mixer.head_k_dim),
	dtype=torch.float32, device=device),
	tuple(
	torch.zeros((batch, conv.weight.shape[0], width), dtype=dtype, device=device)
	for conv in (mixer.q_conv1d, mixer.k_conv1d, mixer.v_conv1d)
	),
	)
	self.graph: torch.cuda.CUDAGraph \| None = None
	self._host_position = [0] * batch
	self._host_active = [False] * batch
	self.logits: torch.Tensor \| None = None
	self.summary: torch.Tensor \| None = None

	# ---- per-layer static mixers -------------------------------------------------

	def _gqa(self, layer, buffers: _GQABuffers, hidden: torch.Tensor) -> torch.Tensor:
	mixer = layer.mixer
	normalized = layer.input_norm(hidden)
	raw_q, raw_k, raw_v, gate = mixer._linear_project(normalized)
	query, key, value = mixer._shape_projected(raw_q, raw_k, raw_v, self.position)
	batch, heads, slots, width = buffers.keys.shape
	position = self.position.view(batch, 1)
	slot = position if buffers.is_global else torch.remainder(position, slots)
	index = slot.view(batch, 1, 1, 1).expand(batch, heads, 1, width)
	buffers.keys.scatter_(2, index, key)
	buffers.values.scatter_(2, index, value)
	# Ring slots fill in order from 0, so the first min(position + 1, slots)
	# slots hold exactly the keys this query may see.
	valid = torch.arange(slots, device=hidden.device).view(1, slots) <= position
	# One query token: fold each KV head's query heads (SDPA's GQA maps
	# query head h to KV head h // groups) into the query-length axis and
	# attend with two batched matmuls, reading each KV buffer once. SDPA
	# with ``enable_gqa`` expands KV to every query head (~2 ms per layer
	# at batch 16) and SDPA on the folded shape picks a slow FP32 kernel;
	# measured per layer at 4,096 slots: 0.17-0.32 ms here vs 0.27-3.9 ms.
	# Same scale as the model (SDPA default 1/sqrt(head_dim)).
	query_heads = query.shape[1]
	grouped = query.reshape(batch, heads, query_heads // heads, width)
	scores = torch.matmul(grouped, buffers.keys.transpose(-1, -2)) * width**-0.5
	scores = scores.masked_fill(~valid.view(batch, 1, 1, slots), float("-inf"))
	output = torch.matmul(torch.softmax(scores, dim=-1), buffers.values)
	output = output.reshape(batch, query_heads, 1, width)
	return mixer._apply_output_gate(output, gate)

	def _kda(self, layer, buffers: _KDABuffers, hidden: torch.Tensor) -> torch.Tensor:
	from einops import rearrange
	from fla.ops.kda.fused_recurrent import fused_recurrent_kda_fwd

	mixer = layer.mixer
	x = layer.input_norm(hidden)
	conv_q, conv_k, conv_v = buffers.conv
	q, _ = mixer.q_conv1d(x=mixer.q_proj(x), cache=conv_q, output_final_state=True)
	k, _ = mixer.k_conv1d(x=mixer.k_proj(x), cache=conv_k, output_final_state=True)
	v, _ = mixer.v_conv1d(x=mixer.v_proj(x), cache=conv_v, output_final_state=True)
	g = rearrange(mixer.f_proj(x), "... (h d) -> ... h d", d=mixer.head_k_dim)
	beta = mixer.b_proj(x)
	q, k = (rearrange(t, "... (h d) -> ... h d", d=mixer.head_k_dim) for t in (q, k))
	v = rearrange(v, "... (h d) -> ... h d", d=mixer.head_v_dim)
	output, _ = fused_recurrent_kda_fwd(
	q=q, k=k, v=v, g=g, beta=beta, A_log=mixer.A_log, dt_bias=mixer.dt_bias,
	scale=k.shape[-1] ** -0.5, initial_state=buffers.state, inplace_final_state=True,
	output_final_state=True, use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True,
	use_beta_sigmoid_in_kernel=True, allow_neg_eigval=mixer.allow_neg_eigval,
	lower_bound=mixer.lower_bound, cu_seqlens=None, state_v_first=True,
	)
	gate = rearrange(mixer.g_proj(x), "... (h d) -> ... h d", d=mixer.head_v_dim)
	output = rearrange(mixer.o_norm(output, gate), "b t h d -> b t (h d)")
	return mixer.o_proj(output)

	# ---- step and capture --------------------------------------------------------

	def _step(self) -> torch.Tensor:
	model = self.model
	originals = {}
	for index, layer in enumerate(model.layers):
	buffers = self.buffers[index]
	static = self._gqa if isinstance(buffers, _GQABuffers) else self._kda

	def mixer_output_cached(hidden, position_ids, *, fla_cache, attention_cache,
	_layer=layer, _buffers=buffers, _static=static):
	return _static(_layer, _buffers, hidden), None

	originals[index] = layer.__dict__.get("mixer_output_cached")
	layer.mixer_output_cached = mixer_output_cached
	try:
	state = _NullState()
	hidden = model._forward_hidden_two_phase_attention_residual_cached(
	model.embed_tokens(self.token), self.position, state, None, None,
	)
	return F.linear(hidden, model.embed_tokens.weight)[:, -1]
	finally:
	for index, layer in enumerate(model.layers):
	if originals[index] is None:
	del layer.mixer_output_cached
	else:
	layer.mixer_output_cached = originals[index]

	def _tensors(self):
	for buffers in self.buffers.values():
	if isinstance(buffers, _GQABuffers):
	yield buffers.keys
	yield buffers.values
	else:
	yield buffers.state
	yield from buffers.conv

	@torch.inference_mode()
	def capture(self, warmup: int = 3) -> None:
	stream = torch.cuda.Stream()
	stream.wait_stream(torch.cuda.current_stream())
	with torch.cuda.stream(stream):
	for _ in range(warmup):
	_greedy_summary(self._step())
	torch.cuda.current_stream().wait_stream(stream)
	graph = torch.cuda.CUDAGraph()
	with torch.cuda.graph(graph):
	self.logits = self._step()
	self.summary = _greedy_summary(self.logits)
	# Warm-up and capture ran real steps on the buffers. ``load`` rewrites
	# every buffer of the row it fills, so clearing (instead of keeping a
	# snapshot copy, which doubled peak memory) is enough.
	for tensor in self._tensors():
	tensor.zero_()
	self.position.zero_()
	self.active.zero_()
	self.graph = graph

	# ---- runtime -----------------------------------------------------------------

	@torch.inference_mode()
	def load(self, cache, row: int = 0) -> int:
	"""Copy a batch-1 ``forward_cached`` prefill cache into ``row``."""

	position = int(cache.seen_tokens)
	if position >= self.length:
	raise ValueError("prompt does not fit this decode bucket")
	for index, buffers in self.buffers.items():
	if isinstance(buffers, _GQABuffers):
	buffers.keys[row].zero_()
	buffers.values[row].zero_()
	entry = cache.attention.get(index)
	if entry is None:
	continue
	count = entry.keys.shape[2]
	positions = torch.arange(position - count, position, device=self.device)
	slots = positions if buffers.is_global else torch.remainder(positions, self.window)
	buffers.keys[row].index_copy_(1, slots, entry.keys[0])
	buffers.values[row].index_copy_(1, slots, entry.values[0])
	else:
	layer_state = cache.fla_cache[index]
	buffers.state[row].copy_(layer_state["recurrent_state"][0])
	for target, source in zip(buffers.conv, layer_state["conv_state"], strict=True):
	target[row].copy_(source[0])
	self.position[row].fill_(position)
	self.active[row].fill_(1)
	self._host_position[row] = position
	self._host_active[row] = True
	return position

	def release(self, row: int) -> None:
	"""Park ``row``; its outputs are ignored until the next ``load``."""

	self.active[row].fill_(0)
	self.position[row].fill_(0)
	self._host_active[row] = False

	@torch.inference_mode()
	def step(self, tokens) -> torch.Tensor:
	"""``advance`` and return a copy of the ``[batch, vocab]`` logits."""

	self.advance(tokens)
	return self.logits.clone()

	@torch.inference_mode()
	def advance(self, tokens) -> None:
	"""Advance every active row by one token.

	``tokens`` is an int (batch 1) or one int per row; parked rows' entries
	are ignored. Afterwards ``logits`` and ``summary`` hold the new rows
	until the next call.
	"""

	if self.graph is None:
	raise RuntimeError("capture() must run before step()")
	if isinstance(tokens, int):
	tokens = [tokens]
	if len(tokens) != self.batch:
	raise ValueError("step needs one token per row")
	# A global-layer write past the buffer would be an out-of-bounds
	# scatter inside the graph; refuse on the host instead.
	for row in range(self.batch):
	if self._host_active[row]:
	if self._host_position[row] >= self.length:
	raise RuntimeError("decode exceeded the static decoder's length bucket")
	self._host_position[row] += 1
	self.token.copy_(torch.tensor(tokens, dtype=torch.long).view(self.batch, 1), non_blocking=False)
	self.graph.replay()
	self.position.add_(self.active)


	def _greedy_summary(logits: torch.Tensor) -> torch.Tensor:
	# Float32 throughout: the index (< 2**24) is exact in float32 but not in
	# float16, and float16 logits widen exactly.
	lexical = logits[:, :CONTROL_BASE]
	index = lexical.argmax(dim=-1, keepdim=True)
	return torch.cat(
	(
	lexical.gather(-1, index).float(),
	index.float(),
	torch.isnan(logits).any(dim=-1, keepdim=True).float(),
	logits[:, CONTROL_BASE:].float(),
	),
	dim=-1,
	)


	class _NullState:
	"""Stands in for ``YunmoNextCUDAState``: static mixers ignore it."""

	fla_cache = None
	attention: dict = {}


	def bucket_for(length: int) -> int:
	for bucket in BUCKETS:
	if length <= bucket:
	return bucket
	raise ValueError("requested length exceeds the largest decode bucket")