Score a DSA sparse indexer in FP4

atom/model_ops/attention_mla.py

@@ -95,8 +95,27 @@ def indexer_qk_rope_quant_and_cache(
     weights_scale: float,
     preshuffle: bool = False,
     is_neox: bool = True,
+    q_scale_out: torch.Tensor | None = None,
+    kv_cache_scale: torch.Tensor | None = None,
 ) -> None:
-    """Run the fused indexer cache op with ATOM's DCP query semantics."""
+    """Run the fused indexer cache op with ATOM's DCP query semantics.
+
+    The two scale buffers together switch the op to packed E2M1 + e8m0 outputs,
+    and are forwarded only when given so that an aiter predating them still
+    accepts the FP8 call -- the same version tolerance the seg-variant import
+    above keeps. One without the other would fall back to FP8 and write that
+    layout into FP4-shaped buffers, so it is refused here rather than passed on.
+    """
+    if (q_scale_out is None) != (kv_cache_scale is None):
+        raise ValueError(
+            "FP4 output needs both q_scale_out and kv_cache_scale, got only "
+            + ("q_scale_out" if q_scale_out is not None else "kv_cache_scale")
+        )
+    fp4_out = (
+        {"q_scale_out": q_scale_out, "kv_cache_scale": kv_cache_scale}
+        if q_scale_out is not None
+        else {}
+    )
     _indexer_qk_rope_quant_and_cache(
         q,
         q_out,
@@ -117,6 +136,7 @@ def indexer_qk_rope_quant_and_cache(
         preshuffle=preshuffle,
         is_neox=is_neox,
         compute_all_q_rope=get_dcp_world_size() > 1,
+        **fp4_out,
     )
 
 
@@ -2993,6 +3013,9 @@ def _convert_req_index_to_global_index_dsa_prefill_kernel(
     MAX_NUM_BLOCKS_PER_REQ: tl.constexpr,
     PAGE_SIZE: tl.constexpr,
     BLOCK_N: tl.constexpr,  # tile width along columns
+    # The FP8 indexer scores one concatenated KV plane, so a request's own
+    # position is `indice - cu_seqlens_q[req]`; the paged FP4 scorer emits it.
+    SEQ_LOCAL: tl.constexpr,
     # strides (in elements)
     ti_stride0: tl.int64,  # topk_indices stride 0
     ti_stride1: tl.constexpr,  # topk_indices stride 1
@@ -3019,7 +3042,7 @@ def _convert_req_index_to_global_index_dsa_prefill_kernel(
     req_kv_end = tl.load(cu_seqlens_q + req_id + 1, mask=valid_req, other=0)
     req_kv_len = req_kv_end - pre_seqlens_q
 
-    seq_token_idx = indice - pre_seqlens_q
+    seq_token_idx = indice if SEQ_LOCAL else indice - pre_seqlens_q
     block_id = seq_token_idx // PAGE_SIZE
     inblock_offset = seq_token_idx % PAGE_SIZE
 
@@ -3069,6 +3092,7 @@ def triton_convert_req_index_to_global_index_dsa_prefill(
     NUM_TOPK_TOKENS: int = 2048,
     BLOCK_N: int = 1024,  # tile width along columns
     out: torch.Tensor | None = None,
+    seq_local: bool = False,
 ):
 
     assert topk_indices.shape[1] == NUM_TOPK_TOKENS
@@ -3120,6 +3144,7 @@ def triton_convert_req_index_to_global_index_dsa_prefill(
         max_num_blocks_per_req,
         PAGE_SIZE,
         BLOCK_N,
+        seq_local,
         # strides
         ti_stride0,
         ti_stride1,

atom/model_ops/attentions/aiter_mla.py

@@ -48,6 +48,14 @@
     effective_kpool_size,
     topk_output_width,
 )
+from atom.model_ops.sparse_indexer_fp4 import (
+    FP4_MQA_BLOCK_K,
+    FP4_MQA_PARALLEL_UNIT_NUM,
+    fp4_decode_parallel_units,
+    fp4_decode_schedule,
+    fp4_prefill_schedule,
+    sparse_indexer_fp4_enabled,
+)
 from atom.utils import CpuGpuBuffer, envs, upload_numpy
 from atom.utils.block_convert import (
     kv_indices_generate_triton,
@@ -257,6 +265,10 @@ class AiterMLAMetadataBuilder(CommonAttentionBuilder):
     # backend). The fused kernel handles both sparse and dense MLA.
     fuse_mtp_decode_position_update = True
 
+    # `__init__` decides this; the default serves builders assembled field by
+    # field, as tests do.
+    _indexer_fp4 = False
+
     def _global_num_draft_layers(self) -> int:
         """Return draft layers in the target MLA pool across all PP stages."""
         runner = self.model_runner
@@ -381,6 +393,9 @@ def __init__(self, model_runner):
         )
         self.index_kpool = effective_kpool_size(configured_kpool)
         self.index_topk_out = topk_output_width(self.index_topk, configured_kpool)
+        self._indexer_fp4 = self.is_sparse and sparse_indexer_fp4_enabled(
+            config.index_cache_dtype, hf_config, warn=True
+        )
         self.dtype_kv = dtypes.d_dtypes[config.kv_cache_dtype]
         self.dtype_q = self.dtype_kv
 
@@ -522,6 +537,23 @@ def __init__(self, model_runner):
                 dtype=torch.int32,
                 device=self.device,
             )
+            # `[P, 4]` int32 read by the mqa-logits kernel during a CUDAGraph
+            # replay, so it has to be a fixed address refreshed in place, and
+            # one per ubatch since TBO has two in flight.
+            self._indexer_fp4_cta_info = (
+                [
+                    torch.zeros(
+                        fp4_decode_parallel_units(self.max_bs, max_seqlen_qo),
+                        4,
+                        **i32_kwargs,
+                    )
+                    for _ in range(
+                        self._NUM_TBO_UBATCHES + 1 if config.enable_tbo else 1
+                    )
+                ]
+                if self._indexer_fp4
+                else []
+            )
             # One block-table row per query token; only MTP verify needs a
             # copy. Built once per step, not in the indexer, where every
             # full-index layer would allocate one into the CUDAGraph pool.
@@ -1181,6 +1213,8 @@ def _declare_kv_pool(self) -> MlaKvPool:
                 "index_rows_per_block": self._index_rows_per_block(),
                 "index_dim": aligned_index_cache_dim(hf_config),
                 "index_dtype": dtypes.fp8,
+                "index_head_dim": hf_config.index_head_dim,
+                "index_fp4": self._indexer_fp4,
             }
             if runner.has_mla_indexer
             else {}
@@ -1264,10 +1298,27 @@ def build_kv_cache_tensor(self, module):
                 )
             index_cache_layer_id = runner.index_cache_layer_map[global_layer_id]
             index_cache = self.kv_pool.layer("index", index_cache_layer_id)
-            # Use aligned dimension to avoid memory copy in torch inductor
-            module.indexer.k_cache.kv_cache[0] = index_cache.view(
-                -1, 1, runner.aligned_index_dim
-            )
+            # `is_sparse` is this branch's own condition, so both sides
+            # reduce to the same `sparse_indexer_fp4_enabled` call on the same
+            # two inputs and cannot legitimately disagree. Compared rather than
+            # assigned: the Indexer built `k_cache` from its own answer before
+            # this runs, so writing the builder's over it would leave the flag
+            # describing an object it no longer matches.
+            assert module.indexer._indexer_fp4 == self._indexer_fp4, (
+                "FP4 sparse indexer verdict diverged: metadata builder says "
+                f"{self._indexer_fp4}, layer {global_layer_id}'s Indexer says "
+                f"{module.indexer._indexer_fp4}"
+            )
+            if self._indexer_fp4:
+                module.indexer.k_cache.kv_cache[0] = index_cache
+                module.indexer.k_cache.kv_cache_scale = self.kv_pool.layer(
+                    "index_scale", index_cache_layer_id
+                )
+            else:
+                # Use aligned dimension to avoid memory copy in torch inductor
+                module.indexer.k_cache.kv_cache[0] = index_cache.view(
+                    -1, 1, runner.aligned_index_dim
+                )
         module.kv_cache = kv_cache
         return KVCacheTensor(
             layer_num=module.layer_num,
@@ -1289,6 +1340,17 @@ def get_kv_transfer_tensors(self):
         runner = self.model_runner
         if self.kv_pool is None:
             return None
+        if self._indexer_fp4:
+            # A connector is handed one `INDEX_CACHE_ROLE` region per layer --
+            # the whole vocabulary it and its DCP shard plan have for the
+            # indexer -- and the FP4 cache is two planes neither parses.
+            if runner.config.kv_transfer_config:
+                raise NotImplementedError(
+                    "KV transfer with the FP4 sparse indexer is unsupported: "
+                    "the region map cannot describe its separate e8m0 scale "
+                    "plane. Pass --index_cache_dtype fp8 to use a connector."
+                )
+            return None
         # What the pool was built with, not what the hook would recompute: a
         # hybrid caches for fewer layers than the model has, and the consumer
         # indices below are positions in the allocated rows.
@@ -1556,6 +1618,43 @@ def _build_dcp_indexer_prefill_meta(self, attn_metadata, bs: int, counts, var):
         attn_metadata.dcp_indexer_gather_index = torch.from_numpy(
             src.astype(np.int32)
         ).to(dev, non_blocking=True)
+        if self._indexer_fp4:
+            self._build_dcp_indexer_fp4_prefill_meta(
+                attn_metadata, bs, lpad, cu_pad, total_kv, var
+            )
+
+    def _build_dcp_indexer_fp4_prefill_meta(
+        self, attn_metadata, bs: int, lpad, cu_pad, total_kv: int, var
+    ):
+        """Publish the local slot list and identity page table FP4 staging reads.
+
+        The slot formula must track `cp_gather_indexer_k_quant_cache`, the FP8
+        plane's gather, which the FP4 planes have no equivalent of. Its padded
+        indices stay inside the sequence's own blocks, so the staging read
+        cannot leave the allocation.
+        """
+        block = self.model_runner.block_size
+        dev = self.device
+        seq_of = np.repeat(np.arange(bs, dtype=np.int64), lpad)
+        j = np.arange(int(cu_pad[bs]), dtype=np.int64) - np.repeat(cu_pad[:bs], lpad)
+        table = var["block_tables"].np[:bs].astype(np.int64)
+        slots = table[seq_of, j // block] * block + (j % block)
+        attn_metadata.dcp_indexer_fp4_local_slots = torch.from_numpy(
+            slots.astype(np.int32)
+        ).to(dev, non_blocking=True)
+
+        # Fixed width, not `pages`: the scorer specializes on this table's
+        # stride, so a per-batch width recompiles it. Sized at a whole batch's
+        # summed context rather than `block_tables`' one-sequence allowance --
+        # co-scheduled prefills run past that allowance whenever prefix caching
+        # keeps their cached tokens off `max_num_batched_tokens`, which is a
+        # legal schedule, and this table is the only thing that would have
+        # bounded it. Buys one scorer variant against the non-DCP width.
+        pages = -(-total_kv // block)
+        cols = self.max_bs * self.block_table_cols
+        staged_tables = torch.zeros(bs, cols, dtype=torch.int32, device=dev)
+        staged_tables[:, :pages] = torch.arange(pages, dtype=torch.int32, device=dev)
+        attn_metadata.dcp_indexer_fp4_block_tables = staged_tables
 
     def _sparse_selected_counts(self, seq_lens):
         """How many KV entries the indexer actually selects for each row.
@@ -1662,6 +1761,9 @@ def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
             attn_metadata.sparse_kv_indptr = var["sparse_kv_indptr"].copy_to_gpu(
                 scheduled_tokens + 1
             )
+            self._publish_indexer_fp4_prefill_schedule(
+                attn_metadata, sparse_counts, int(full_seq_lens.sum())
+            )
             if self.dcp_world_size > 1:
                 self._build_dcp_indexer_prefill_meta(attn_metadata, bs, counts, var)
             get_mla_metadata_v1(
@@ -2131,6 +2233,105 @@ def _publish_dcp_token_block_tables(
         rows.view(running_bs, max_seqlen_q, -1).copy_(block_tables.unsqueeze(1))
         attn_metadata.dcp_token_block_tables = rows
 
+    def _publish_indexer_fp4_decode_schedule(
+        self, attn_metadata: AttentionMetaData, bs: int, next_n: int, ubatch: int = 0
+    ) -> None:
+        """Refresh the FP4 decode CTA schedule; the captured kernel replays off
+        it, and every indexer layer of a step shares the one answer.
+
+        A replay reads the buffer's CONTENTS, so every decode forward has to
+        call this -- including the MTP draft's, whose rows are its own. The
+        lengths come from the published buffer rather than the metadata's view
+        of it, because a draft following a prefill carries prefill metadata.
+        """
+        if not self._indexer_fp4:
+            return
+        parallel_units = fp4_decode_parallel_units(bs, next_n)
+        cta_info = self._indexer_fp4_cta_info[ubatch][:parallel_units]
+        # `[:n]` past the end truncates rather than raising, while
+        # `indexer_fp4_n_ctas` still hands the kernel `n`. The rows past the end
+        # would then go unscheduled, and the logits buffer they should have
+        # written is `torch.empty` -- aiter only sentinels it when it builds the
+        # schedule itself, which it never does here.
+        assert cta_info.shape[0] == parallel_units
+        if self.dcp_world_size > 1:
+            # DCP scores this rank's shard. Its rows are query tokens carrying
+            # their own local window, so next_n is already flattened out of the
+            # row axis and the width is the sharded one -- `parallel_units`
+            # stays the non-DCP count, which keeps the captured grid identical.
+            from atom.model_ops.dcp_ops import dcp_local_logits_width
+            from atom.utils.forward_context import (
+                get_published_dcp_local_context_lens,
+            )
+
+            # A TBO ubatch owns rows from its own offset, not the batch's first
+            # `bs`, and it already publishes that slice on its metadata -- the
+            # same tensor the scorer reads. The global buffer stays the fallback
+            # so callers that publish nothing keep the prefix they had.
+            context_lens = get_published_dcp_local_context_lens(
+                attn_metadata, bs * next_n
+            )
+            if context_lens is None:
+                context_lens = self.model_runner.forward_vars[
+                    "dcp_local_context_lens"
+                ].gpu[: bs * next_n]
+            schedule_next_n = 1
+            width = dcp_local_logits_width(
+                self.model_runner.config.max_model_len, self.dcp_world_size
+            )
+        else:
+            context_lens = attn_metadata.context_lens[:bs]
+            schedule_next_n = next_n
+            width = self.model_runner.config.max_model_len
+        fp4_decode_schedule(
+            context_lens,
+            FP4_MQA_BLOCK_K,
+            parallel_units,
+            width,
+            schedule_next_n,
+            cta_info,
+        )
+        attn_metadata.indexer_fp4_cta_info = cta_info
+        attn_metadata.indexer_fp4_n_ctas = parallel_units
+
+    def _publish_indexer_fp4_prefill_schedule(
+        self, attn_metadata: AttentionMetaData, local_ends_np: np.ndarray, total_kv: int
+    ) -> None:
+        """Build the FP4 ragged-prefill schedule and the width it scores in.
+
+        The width comes off `sparse_counts`, which the host already holds, so
+        right-sizing the logits buffer to this batch rather than max_model_len
+        costs no device sync.
+
+        Under DCP the scorer reads a staged copy of every sequence's keys rather
+        than the in-place shard, so its columns are the flat concatenated ones
+        the FP8 path ranks in and the width is the whole key set.
+        """
+        if not self._indexer_fp4:
+            return
+        if self.dcp_world_size > 1:
+            local_starts = attn_metadata.cu_seqlen_ks
+            local_ends = attn_metadata.cu_seqlen_ke
+            max_seq_len = max(int(total_kv), 1)
+        else:
+            local_starts = None
+            local_ends = attn_metadata.cu_seqlen_ke - attn_metadata.cu_seqlen_ks
+            max_seq_len = max(int(local_ends_np.max(initial=0)), 1)
+        (
+            attn_metadata.indexer_fp4_cta_info,
+            attn_metadata.indexer_fp4_n_ctas,
+            attn_metadata.indexer_fp4_local_starts,
+        ) = fp4_prefill_schedule(
+            attn_metadata.batch_id_per_q_token,
+            local_ends,
+            FP4_MQA_BLOCK_K,
+            FP4_MQA_PARALLEL_UNIT_NUM,
+            max_seq_len,
+            local_starts,
+        )
+        attn_metadata.indexer_fp4_local_ends = local_ends
+        attn_metadata.indexer_fp4_max_seq_len = max_seq_len
+
     def prepare_decode(
         self,
         batch: ScheduledBatch,
@@ -2437,6 +2638,9 @@ def prepare_decode(
                 "sparse_kv_last_page_lens"
             ].gpu[:running_bs]
         self._publish_dcp_token_block_tables(attn_metadata, running_bs, max_seqlen_q)
+        self._publish_indexer_fp4_decode_schedule(
+            attn_metadata, running_bs, max_seqlen_q
+        )
 
         # running_bs, not scheduled_bs: the padded rows have to be split into the
         # ubatches too, or accuracy drifts.
@@ -2742,6 +2946,7 @@ def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData:
                 "sparse_kv_last_page_lens"
             ].gpu[:bs]
         self._publish_dcp_token_block_tables(attn_matadata, bs, max_q_len)
+        self._publish_indexer_fp4_decode_schedule(attn_matadata, bs, max_q_len)
         positions = var["positions"].copy_to_gpu(scheduled_tokens)
         context = Context(
             positions=positions,
@@ -2819,6 +3024,11 @@ def build_ubatch_metadata(
             # DCP, so a ubatch always runs one query per sequence.
             assert max_q_len == 1
             attn.dcp_token_block_tables = attn.block_tables
+        # Slot 0 stays the full-batch schedule, so a ubatch never overwrites
+        # one the other's kernels still read.
+        self._publish_indexer_fp4_decode_schedule(
+            attn, running_bs, max_q_len, ubatch=ubatch_idx + 1
+        )
         return attn
 
     def build_ubatch_prefill_metadata(

atom/model_ops/attentions/backends.py

@@ -479,6 +479,17 @@ def __init__(self, model_runner):
         self.model_runner.forward_vars.update(attn_metadata)
         self.has_sliding_window = hasattr(hf_config, "sliding_window")
 
+    def _publish_indexer_fp4_decode_schedule(
+        self, attn_metadata, bs: int, next_n: int, ubatch: int = 0
+    ) -> None:
+        """Nothing to refresh: this backend has no FP4 sparse indexer.
+
+        `EagleProposer` publishes on whatever builder the target uses, so every
+        backend a draft can run against has to answer. Only the MLA one
+        overrides. Inert by contract, not just by accident -- the draft reuses
+        the target's metadata, so a write here would reach the verify step.
+        """
+
     def prepare_block_tables(self, batch: ScheduledBatch):
         """Marshal the batch's block tables into `forward_vars["block_tables"]`.