[Kimi-K3] Prefill token graphs: the KDA/MLA backend takes a padded prefill

atom/model_ops/attentions/kimi_mla_gdn_attn.py

@@ -1,11 +1,18 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
+import numpy as np
 import torch
 from aiter.dist.parallel_state import get_tp_group
 
 from atom.model_engine.kv_block import STATE_SLOT_CLASS
 from atom.model_engine.scheduler import ScheduledBatch
+from atom.model_engine.sequence import (
+    Sequence,
+    SequenceStatus,
+    SequenceType,
+    new_block_table,
+)
 from atom.model_engine.state_runtime import StateTransfer
 from atom.model_ops.attention_mla import MLAAttention
 from atom.model_ops.attentions.pool_layout.v4_pool_fields import (
@@ -17,6 +24,7 @@
     speculative_kpool_history_size,
 )
 from atom.utils import envs
+from atom.utils.forward_context import Context
 
 from .aiter_mla import (
     MLA_ROWS,
@@ -47,9 +55,24 @@ def get_impl_cls() -> type["MLAAttention"]:
 
 
 class _KimiMLAGDNCommon(PageUnitGeometryMixin, GDNStateMixin):
+    # A prefill may pad to a recorded token-graph bucket: the KDA and MLA
+    # regions run its real rows (`AttentionMetaData.prefill_real_rows`), and
+    # every launch between them is sized by the token count alone.
+    supports_token_graph = True
+
     def __init__(self, model_runner):
         super().__init__(model_runner=model_runner)
 
+    def step_needs_eager(self, batch: ScheduledBatch) -> bool:
+        """A prefill carrying images runs eagerly, as on DeepSeek-V4.1: its
+        embeddings are scattered from the scheduled ids, which a token-graph
+        bucket would pad."""
+        if batch.total_tokens_num_prefill > 0 and getattr(
+            batch, "multimodal_data", None
+        ):
+            return True
+        return super().step_needs_eager(batch)
+
     def _module_kinds(self, module) -> tuple:
         """Two row spaces: the MLA pool's rows, and the KDA state slots.
 
@@ -502,6 +525,14 @@ def prepare_prefill(
         attn_metadata, positions = super().prepare_prefill(
             batch, running_bs, running_tokens
         )
+        scheduled = batch.total_tokens_num_prefill
+        if running_tokens > scheduled:
+            # Padded to a token-graph bucket. The metadata above is the real
+            # rows', which the regions cut to; the pad rows' slots are -1, so a
+            # recorded segment that writes by slot writes nothing for them.
+            attn_metadata.prefill_real_rows = scheduled
+            slots = self.model_runner.forward_vars["slot_mapping"].gpu
+            slots[scheduled:running_tokens].fill_(-1)
         if batch.block_tables == []:
             attn_metadata.gdn_metadata = None
             return attn_metadata, positions
@@ -540,6 +571,44 @@ def build_for_cudagraph_capture(self, bs: int):
         attn_metadata.gdn_metadata = self._build_gdn_capture_metadata(bs)
         return attn_metadata, context
 
+    def build_for_token_graph_capture(self, tokens: int):
+        """A prefill of ``tokens`` rows for recording that token-graph bucket.
+
+        The request-less prefill `ModelRunner.dummy_execution(prefill=True)`
+        runs: one sequence on placeholder page 0, holding no state slot. The KDA
+        layers get no metadata, so they write no state and hand back zeros; MLA
+        writes to the placeholder page only. The segments need only its shape.
+        """
+        seq = Sequence([0] * tokens, block_size=self.model_runner.block_size, id=-1)
+        seq.status = SequenceStatus.RUNNING
+        seq.type = SequenceType.PREFILL
+        seq.block_table = new_block_table([0])
+        batch = ScheduledBatch(
+            seqs={seq.id: seq},
+            num_scheduled_tokens=np.array([tokens], dtype=np.int32),
+            total_tokens_num=tokens,
+            total_tokens_num_prefill=tokens,
+            total_tokens_num_decode=0,
+            total_seqs_num=1,
+            total_seqs_num_prefill=1,
+            total_seqs_num_decode=0,
+            is_dummy_run=True,
+            num_spec_step=0,
+            scheduled_spec_decode_tokens=None,
+        )
+        attn_metadata, positions = self.build(batch, 1, tokens, tokens)
+        attn_metadata.gdn_metadata = None
+        context = Context(
+            positions=positions,
+            is_prefill=True,
+            is_dummy_run=True,
+            scheduled_bs=1,
+            scheduled_tokens=tokens,
+            running_bs=1,
+            running_tokens=tokens,
+        )
+        return attn_metadata, context
+
 
 class KimiAiterMLAGDNMetadataBuilder(_KimiMLAGDNCommon, AiterMLAMetadataBuilder):
     def _supports_dcp_index_staging(self) -> bool:

atom/model_ops/base_attention.py

@@ -20,6 +20,7 @@
 from atom.utils import envs, mark_spliting_op
 from atom.utils.region_graph import eager_region
 from atom.utils.selector import Family, get_attn_backend
+from atom.utils.token_graph_rows import padded_prefill_rows, real_rows, to_rows
 
 from .attention_mla import MLAModules, _mla_output_width
 
@@ -1060,23 +1061,32 @@ def unified_attention_with_output_base(
 ) -> torch.Tensor:
     atom_config = get_current_atom_config()
     self = atom_config.compilation_config.static_forward_context[layer_name]
+    rows, real = q.shape[0], padded_prefill_rows()
+    if real is not None and real < rows:
+        # A padded prefill: the metadata covers the real rows only, so attention
+        # runs on those, and its FP8 group scales must not see the pad rows.
+        q, k, v, positions, qkv = (
+            real_rows(t, real, rows) for t in (q, k, v, positions, qkv)
+        )
+        q_scale = real_rows(q_scale, real, rows, groups_major=True)
     if use_mla:
-        return self.impl.forward(
+        out = self.impl.forward(
             query=q,
             k_nope=k,
             k_rope=v,
             positions=positions,
             q_scale=q_scale,
         )
     else:
-        return self.impl.forward(
+        out = self.impl.forward(
             query=q,
             key=k,
             value=v,
             position=positions,
             q_scale=q_scale,
             qkv=qkv,
         )
+    return out if real is None else to_rows(out, rows)
 
 
 def linear_attention_with_output_base_fake(

atom/models/kimi_k3/model.py

@@ -74,6 +74,7 @@
 from atom.utils.decorators import mark_trace, support_torch_compile
 from atom.utils.forward_context import get_forward_context
 from atom.utils.region_graph import eager_region, token_graph
+from atom.utils.token_graph_rows import padded_prefill_rows, to_rows
 
 
 def _text_config(config):
@@ -1003,7 +1004,12 @@ def kda_attention_with_output(
     self = get_current_atom_config().compilation_config.static_forward_context[
         layer_name
     ]
-    return self._forward_impl(hidden_states, hidden_states_scale)
+    out = self._forward_impl(hidden_states, hidden_states_scale)
+    # A padded prefill: `_forward_impl` ran the real rows (`num_actual_tokens`);
+    # the next recorded segment reads the bucket's.
+    if padded_prefill_rows() is None:
+        return out
+    return to_rows(out, hidden_states.shape[0])
 
 
 class KimiKDAAttention(nn.Module):

atom/utils/envs.py

@@ -627,6 +627,12 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float:
     "ATOM_TOKEN_GRAPH_DP_MAX_TOKENS": lambda: int(
         os.getenv("ATOM_TOKEN_GRAPH_DP_MAX_TOKENS", "512")
     ),
+    # Debug / A-B: a step padded to a recorded token-graph bucket runs its
+    # forward eagerly instead of replaying the recording, so the padded math
+    # can be compared with and without the graphs. Off (0).
+    "ATOM_TOKEN_GRAPH_EAGER_PAD": lambda: (
+        os.getenv("ATOM_TOKEN_GRAPH_EAGER_PAD", "0") == "1"
+    ),
     # DeepSeek-V4.1's decoder on a bounded-replay prefill's tail, recorded as a
     # request graph (`models/deepseek_v41/decoder_graph.py`): the comma-separated
     # `running_bs` to record at. Any other runs the same tail eagerly; empty

atom/utils/forward_context.py

@@ -664,6 +664,11 @@ class AttentionMetaData:
     num_cached_tokens: torch.Tensor | None = None
     seq_starts: torch.Tensor | None = None
 
+    # A prefill padded to a token-graph bucket whose metadata covers only its
+    # real rows: how many there are, for the eager regions that cut to them
+    # (`atom/utils/token_graph_rows.py`). None on every other step.
+    prefill_real_rows: int | None = None
+
     def __init__(
         self,
         cu_seqlens_q: torch.Tensor | None = None,

atom/utils/region_graph.py

@@ -42,7 +42,7 @@
 import torch
 
 from atom.distributed.nccl_watchdog import drain_nccl_works
-from atom.utils import forward_context, weak_ref_tensor, weak_ref_tensors
+from atom.utils import envs, forward_context, weak_ref_tensor, weak_ref_tensors
 from atom.utils.graph_holders import register_graph_holder
 
 __all__ = [
@@ -580,6 +580,9 @@ def _run_token_graph(owner, forward: Callable, args, kwargs, per_phase=False):
     phase = _phase() if per_phase else None
     if graph.has(key, phase):
         if mode == "replay":
+            if envs.ATOM_TOKEN_GRAPH_EAGER_PAD:
+                # The padded step run eagerly: a debug A/B against the replay.
+                return forward(*args, **kwargs)
             return graph.replay(args, kwargs, phase=phase)
         if mode == "record":
             graph.mark_captured(key, phase)

atom/utils/token_graph_rows.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""The real rows of a prefill padded to a token-graph bucket.
+
+A backend whose eager regions plan their launches from the step's real rows
+(Kimi-K3's KDA and MLA) builds its prefill metadata for those rows only and
+marks the step with how many there are (`AttentionMetaData.prefill_real_rows`).
+A region then cuts its bucket-high inputs back to the real rows, and hands its
+output back at bucket height, the shape the next recorded segment reads
+(`atom/utils/region_graph.py`). Outside such a step each helper is the
+identity, so an unpadded prefill and every decode run as before.
+"""
+
+import torch
+
+from atom.utils import forward_context
+
+
+def padded_prefill_rows() -> int | None:
+    """The real rows of the running step if it is a marked padded prefill, else None."""
+    metadata = getattr(forward_context.get_forward_context(), "attn_metadata", None)
+    return getattr(metadata, "prefill_real_rows", None)
+
+
+def real_rows(
+    t: torch.Tensor | None, real: int, rows: int, groups_major: bool = False
+) -> torch.Tensor | None:
+    """`t` cut to its first `real` rows when it is a `rows`-high token-major tensor.
+
+    A 2-D per-group scale written group-major (`transpose_scale`: row r of group
+    g at g * rows + r in a (rows, groups) buffer) is rebuilt group-major at the
+    real height, which is what an unpadded producer lays out.
+    """
+    if t is None or t.dim() == 0 or t.shape[0] != rows:
+        return t
+    if groups_major and t.dim() == 2 and t.shape[1] > 1:
+        groups = t.shape[1]
+        return t.reshape(groups, rows)[:, :real].contiguous().reshape(real, groups)
+    return t[:real]
+
+
+def to_rows(out: torch.Tensor, rows: int) -> torch.Tensor:
+    """`out` at `rows` rows: its own rows, then zeros. Identity at `rows` already."""
+    if out.shape[0] == rows:
+        return out
+    if out.shape[0] > rows:
+        raise ValueError(f"a region returned {out.shape[0]} rows for a {rows}-row step")
+    padded = out.new_zeros((rows, *out.shape[1:]))
+    padded[: out.shape[0]] = out
+    return padded

docs/environment_variables.md

@@ -331,8 +331,9 @@ dependency graph that a single stream serializes.
 | Variable | Type | Default | Description |
 |----------|------|---------|-------------|
 | **ATOM_DSV41_SIDE_STREAMS** | int (0/1/2) | 0 | How many of a layer's branches leave the main stream. **0** — none. **1** — the compressor, on the MoE's `alt_stream`, waited at the scorer rather than at attention, because `visible` counts the index row a boundary crossed in this same forward writes, so the scorer is its first reader a whole top-k chain earlier. It borrows that stream because the MoE joins it inside its own forward, a sublayer after the compressor was joined, so the two are never live at once. **2** — the same, plus the indexer on a stream of its own; it is the only branch live beside both others, so it is the only one worth a queue. A value outside 0–2 raises rather than rounding. |
-| **ATOM_TOKEN_GRAPH_SIZES** | str | "256,512,1024,2048,3072" | Prefill token buckets to record as TR graphs (`atom/utils/region_graph.py`), comma-separated. Each bucket's forward is recorded once at startup as graph segments broken at the model's eager regions (indices, sparse attention, window write, compressor, indexer, Engram), and a prefill padded to the smallest bucket that holds it (under DP attention, every prefilling rank's to the one holding the group's largest) replays them instead of launching every kernel; a longer prefill runs eagerly. Empty records nothing. Only a backend with `supports_token_graph` (DeepSeek-V4.1) uses it, and a disaggregated decode instance records none. The default stops at 3072 because the graph pool grows with the largest bucket: on V4.1 TP4 per GPU, FULL decode alone takes 1.15 GB, the five default buckets 3.1 GB, buckets up to 16384 10.8 GB. Below 3072 a prefill step is CPU-bound, so a graph cuts its TTFT by 30–40 ms and padding to the next bucket costs about nothing; longer prefills are GPU-bound. The KV budget's graph estimate does not include these buckets: they come out of the safety margin. Under DP attention a rank with no prefill of its own runs a dummy prefill while a peer prefills, instead of decoding, and every rank of the step replays one bucket, the smallest holding the group's largest row count. Without the prefill delayer (`ATOM_ENABLE_PREFILL_DELAYER=0`) a rank may decode beside a peer's prefill; it pads its decode rows to that bucket and replays the same prefill recording. |
+| **ATOM_TOKEN_GRAPH_SIZES** | str | "256,512,1024,2048,3072" | Prefill token buckets to record as TR graphs (`atom/utils/region_graph.py`), comma-separated. Each bucket's forward is recorded once at startup as graph segments broken at the model's eager regions (indices, sparse attention, window write, compressor, indexer, Engram), and a prefill padded to the smallest bucket that holds it (under DP attention, every prefilling rank's to the one holding the group's largest) replays them instead of launching every kernel; a longer prefill runs eagerly. Empty records nothing. Only a backend with `supports_token_graph` (DeepSeek-V4.1, Kimi-K3) uses it, and a disaggregated decode instance records none. The default stops at 3072 because the graph pool grows with the largest bucket: on V4.1 TP4 per GPU, FULL decode alone takes 1.15 GB, the five default buckets 3.1 GB, buckets up to 16384 10.8 GB. Below 3072 a prefill step is CPU-bound, so a graph cuts its TTFT by 30–40 ms and padding to the next bucket costs about nothing; longer prefills are GPU-bound. The KV budget's graph estimate does not include these buckets: they come out of the safety margin. Under DP attention a rank with no prefill of its own runs a dummy prefill while a peer prefills, instead of decoding, and every rank of the step replays one bucket, the smallest holding the group's largest row count. Without the prefill delayer (`ATOM_ENABLE_PREFILL_DELAYER=0`) a rank may decode beside a peer's prefill; it pads its decode rows to that bucket and replays the same prefill recording. |
 | **ATOM_TOKEN_GRAPH_DP_MAX_TOKENS** | int | 512 | Under DP attention with a drafter, the largest decode token bucket recorded as a TR_HYBRID token graph; larger `bs × q` buckets are skipped at capture and left out of the graph memory estimate. Such buckets never run under DP, and recording them only grows the graph pool. The same skip on every DP rank keeps their capture loops in step. |
+| **ATOM_TOKEN_GRAPH_EAGER_PAD** | bool | 0 (false) | Debug / A-B. Set to `1` to run a step padded to a recorded `ATOM_TOKEN_GRAPH_SIZES` bucket eagerly instead of replaying the recording: the same padding and the same regions, without the graphs, so a disagreement can be put on the padding or on the replay. |
 | **ATOM_DECODER_GRAPH_SIZES** | str | "1,2,3,4,5,6,7,8,12,16" | The `running_bs` at which DeepSeek-V4.1's decoder (layers 20..39) on a bounded-replay prefill's tail is recorded as a request graph (`models/deepseek_v41/decoder_graph.py`), comma-separated. The tail is laid out for the step's `running_bs` requests, a window each, at one set of addresses per `running_bs`, so a recording replays the whole decoder in one launch; any other `running_bs` runs the same layout eagerly, and empty records nothing. Not on an FP8 index plane. Under DP attention each size also records the tail of padding a rank finishing no prompt runs beside its peers. On a 16k-token prefill step the decoder in one graph cuts the step's CPU time from 84 to about 45 ms; the six sizes 1,2,3,4,8,16 add about 0.33 GB to the graph pool on V4.1 TP4. `running_bs` is always a cudagraph capture size (`ForwardMode.decide` rounds up the ladder), so the default follows the usual ladder up to 16; 5, 6, 7 and 12 were added because without them steps of 5-7 and 9-12 requests ran the decoder eagerly (on V4.1 EP4 + DP attention the graph pool stayed at 8.71 GB per GPU). |
 
 **Level 0 is the default because forking measured slower, not faster.** Median