perf: MLA 历史分块 prefill 与稀疏打分优化（默认 16K）

docs/en/configuration/runtime-parameter-semantics.md

@@ -62,6 +62,16 @@ Do not duplicate method policy decisions in benchmark scripts. Runtime reports
 should record the resolved method, policy, chunk size, context length, batch
 size, and checkpoint path.
 
+### MLA history workspace
+
+`mla_prefill_history_chunk_size` is a positive integer (default: `16384`).
+It limits the historical KV tokens gathered and expanded at once during MLA
+prefill, independently of the new-token limit `engine_prefill_chunk_size`.
+Smaller values reduce history workspace but add attention calls and merges.
+Current-token activations and sparse score state still require memory; this
+parameter is not a cap on the complete prefill peak. Sparse observation windows
+and the selected `sparse_prefill_score_mode` retain their existing meanings.
+
 ## Prefill sparsity
 
 `prefill_sparse_method` selects prefill acceleration independently from

docs/zh/configuration/runtime-parameter-semantics.md

@@ -58,6 +58,14 @@ Prefill policy 的唯一事实来源是 `src/sparsevllm/method_registry.py`：
 不要在 benchmark script 中复制 method policy。运行报告应记录解析后的
 method、policy、chunk size、context length、batch size 和 checkpoint 路径。
 
+### MLA 历史工作区
+
+`mla_prefill_history_chunk_size` 为正整数，默认 `16384`，限制 MLA prefill
+一次 gather 和展开的历史 KV token 数，与控制新 token 的
+`engine_prefill_chunk_size` 独立。较小的值降低历史工作区，但增加 attention
+调用和结果合并次数。当前 token 激活和稀疏分数状态仍需显存，因此该参数不是
+完整 prefill 峰值的上限。观察窗口和 `sparse_prefill_score_mode` 的语义保持不变。
+
 ## Prefill 稀疏
 
 `prefill_sparse_method` 独立选择 prefill 加速，不替代 `sparse_method`。当前有

scripts/profiling/kernel_bench/benchmark_mla_prefill.py

@@ -0,0 +1,161 @@
+"""Compare MLA partial-attention wrappers against a fixed Git revision."""
+
+import argparse
+import hashlib
+import importlib.metadata
+import importlib.util
+import json
+import os
+from pathlib import Path
+import statistics
+import subprocess
+import sys
+import time
+
+import torch
+
+from sparsevllm.kernels.triton.mla import prefill
+from sparsevllm.kernels.triton.context_flashattention_nopad import context_attention_fwd
+
+
+def git(*args):
+    return subprocess.check_output(["git", *args], text=True).strip()
+
+
+def check_sampled_oracle(q, k, v, cu_q, cu_k, output, lse, causal):
+    # Bound oracle memory even for long contexts; cover first/middle/last Q.
+    for qa, qb, ka, kb in zip(cu_q, cu_q[1:], cu_k, cu_k[1:]):
+        qn, kn = qb - qa, kb - ka
+        rows = torch.tensor(sorted({0, qn // 2, qn - 1}), device=q.device)
+        z = torch.einsum("qhd,khd->hqk", q[qa + rows].float(), k[ka:kb].float()) * 0.0625
+        if causal:
+            z.masked_fill_(torch.arange(kn, device=q.device)[None, None] > (rows + kn - qn)[None, :, None], -torch.inf)
+        expected = torch.einsum("hqk,khd->qhd", z.softmax(-1), v[ka:kb].float())
+        torch.testing.assert_close(output[qa + rows].float(), expected, atol=0.006, rtol=0.03)
+        if lse is not None:
+            torch.testing.assert_close(lse[:, qa + rows], z.logsumexp(-1), atol=0.003, rtol=0.001)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--baseline-ref", required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--rounds", type=int, default=10)
+    parser.add_argument("--heads", type=int, choices=(5, 10, 20), default=20)
+    args = parser.parse_args()
+    if args.warmup < 1 or args.rounds < 1:
+        parser.error("warmup and rounds must be positive")
+    root = args.output_dir
+    root.mkdir(parents=True, exist_ok=True)
+    if (root / "run_manifest.json").exists() or (root / "raw_samples.jsonl").exists():
+        raise FileExistsError("Use a new output directory to preserve previous results")
+    path = "src/sparsevllm/kernels/triton/mla/prefill.py"
+    source = git("show", f"{args.baseline_ref}:{path}") + "\n"
+    baseline_path = root / "baseline_prefill.py"
+    baseline_path.write_text(source)
+    spec = importlib.util.spec_from_file_location("mla_prefill_baseline", baseline_path)
+    baseline = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(baseline)
+    props = torch.cuda.get_device_properties(0)
+    manifest = {
+        "status": "running", "command": [sys.executable, *sys.argv],
+        "repo": git("rev-parse", "--show-toplevel"), "head": git("rev-parse", "HEAD"),
+        "branch": git("branch", "--show-current"), "git_status": git("status", "--short"),
+        "baseline": git("rev-parse", args.baseline_ref),
+        "gpu": props.name, "capability": [props.major, props.minor],
+        "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
+        "triton_cache_dir": os.environ.get("TRITON_CACHE_DIR"),
+        "gpu_state": subprocess.check_output(["nvidia-smi", "--query-gpu=index,uuid,name,memory.used,utilization.gpu,clocks.sm,power.limit", "--format=csv"], text=True),
+        "versions": {p: importlib.metadata.version(p) for p in ("torch", "triton", "sglang-kernel")},
+        "cuda": torch.version.cuda, "seed": 42, "dtype": "bfloat16", "heads": args.heads,
+        "head_dim": 256, "graph": False, "warmup": args.warmup, "rounds": args.rounds,
+        "timing": "CUDA events around wrapper including output/LSE allocation; warm caches; ABBA pairs; no clock control",
+        "candidate_sha256": hashlib.sha256(Path(prefill.__file__).read_bytes()).hexdigest(),
+    }
+    (root / "candidate_prefill.py").write_text(Path(prefill.__file__).read_text())
+    (root / "benchmark_script.py").write_text(Path(__file__).read_text())
+    (root / "changes.patch").write_text(git("diff", "HEAD"))
+    manifest_path = root / "run_manifest.json"
+    manifest_path.write_text(json.dumps(manifest, indent=2))
+    results = []
+    try:
+        with (root / "raw_samples.jsonl").open("x") as raw:
+            for queries, keys, causal in (
+                ((4096,), (4096,), True), ((8192,), (8192,), True),
+                ((4096, 7), (4096, 7), True),
+                ((32,), (16384,), False), ((128,), (16384,), False),
+                ((129,), (16384,), False), ((512,), (16384,), False),
+                ((8192,), (16384,), False),
+            ):
+                torch.manual_seed(42)
+                q = torch.randn(sum(queries), args.heads, 256, device="cuda", dtype=torch.bfloat16) * 0.2
+                k = torch.randn(sum(keys), args.heads, 256, device="cuda", dtype=torch.bfloat16) * 0.2
+                # Match the V view returned by the serving joint KV projection.
+                projected = torch.randn(sum(keys), args.heads, 448, device="cuda", dtype=torch.bfloat16) * 0.2
+                v = projected[..., 192:]
+                cq = [0, *torch.tensor(queries).cumsum(0).tolist()]
+                ck = [0, *torch.tensor(keys).cumsum(0).tolist()]
+                cu_q, cu_k = [torch.tensor(x, device="cuda", dtype=torch.int32) for x in (cq, ck)]
+
+                def call(module):
+                    return module.attention_partial(q, k, v, cu_q, cu_k, max(queries), max(keys), scale=0.0625, causal=causal)
+
+                calls = {"pr_original": lambda: call(baseline), "candidate": lambda: call(prefill)}
+                if causal:
+                    slots = torch.full((len(keys), max(keys)), -1, device="cuda", dtype=torch.int32)
+                    for i, (a, b) in enumerate(zip(ck, ck[1:])):
+                        slots[i, :b - a] = torch.arange(a, b, device="cuda", dtype=torch.int32)
+                    rows = torch.arange(len(keys), device="cuda", dtype=torch.int32)
+                    lengths = cu_k[1:] - cu_k[:-1]
+                    cached = lengths - (cu_q[1:] - cu_q[:-1])
+                    old_output = torch.empty_like(q)
+
+                    def legacy():
+                        context_attention_fwd(q, k, v, old_output, rows, cu_q[:-1], lengths, cached, max(queries), slots)
+                        return old_output, None
+
+                    calls["legacy_causal"] = legacy
+                cold_ms = {}
+                for name, fn in calls.items():
+                    torch.cuda.synchronize()
+                    start = time.perf_counter()
+                    output, lse = fn()
+                    torch.cuda.synchronize()
+                    cold_ms[name] = (time.perf_counter() - start) * 1000
+                    check_sampled_oracle(q, k, v, cq, ck, output, lse, causal)
+                    for _ in range(args.warmup):
+                        fn()
+                torch.cuda.synchronize()
+                samples = {name: [] for name in calls}
+                pairs = [("pr_original", "candidate")]
+                if causal:
+                    pairs.append(("legacy_causal", "candidate"))
+                for pair in pairs:
+                    for iteration in range(args.rounds):
+                        for name in (*pair, *reversed(pair)):
+                            start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+                            start.record()
+                            calls[name]()
+                            end.record()
+                            end.synchronize()
+                            elapsed = start.elapsed_time(end)
+                            samples[name].append(elapsed)
+                            raw.write(json.dumps({"queries": queries, "keys": keys, "causal": causal, "pair": pair, "iteration": iteration, "implementation": name, "ms": elapsed}) + "\n")
+                row = {"queries": queries, "keys": keys, "causal": causal, "status": "success", "cold_first_call_ms": cold_ms,
+                       "strides": {"q": q.stride(), "k": k.stride(), "v": v.stride()},
+                       "timings": {name: {"n": len(values), "median_ms": statistics.median(values), "min_ms": min(values), "max_ms": max(values)} for name, values in samples.items()}}
+                results.append(row)
+                print(json.dumps(row), flush=True)
+                (root / "summary.json").write_text(json.dumps(results, indent=2))
+        manifest["status"] = "success"
+    except BaseException as error:
+        manifest.update(status="failed", error=f"{type(error).__name__}: {error}")
+        raise
+    finally:
+        manifest_path.write_text(json.dumps(manifest, indent=2))
+
+
+if __name__ == "__main__":
+    with torch.no_grad():
+        main()

src/sparsevllm/configs/runtime.py

@@ -72,6 +72,7 @@ class Config(
     long_prefill_offload_threshold: int = 64 * 1024
     mlp_chunk_size: int = 16384
     mla_prefill_workspace_bytes: int = 6 * 1024**3
+    mla_prefill_history_chunk_size: int = 16384
     prefill_schedule_policy: str = PREFILL_POLICY_AUTO
     gpu_memory_utilization: float = 0.9
     tensor_parallel_size: int = 1

src/sparsevllm/configs/scheduling.py

@@ -118,6 +118,9 @@ def normalize_scheduling(config) -> None:
         raise ValueError(f"mlp_chunk_size must be > 0, got {config.mlp_chunk_size}.")
     config.mlp_chunk_size = int(config.mlp_chunk_size)
     config.mla_prefill_workspace_bytes = int(config.mla_prefill_workspace_bytes)
+    config.mla_prefill_history_chunk_size = int(config.mla_prefill_history_chunk_size)
+    if config.mla_prefill_history_chunk_size <= 0:
+        raise ValueError("mla_prefill_history_chunk_size must be > 0.")
     if config.mla_prefill_workspace_bytes <= 0:
         raise ValueError(
             "mla_prefill_workspace_bytes must be > 0, got "

src/sparsevllm/engine/cache_manager/base.py

@@ -281,12 +281,27 @@ class DecodeComputeView:
     payload: AttentionPayload
 
 
+@dataclass(frozen=True)
+class PrefillScoreRequest:
+    """Method-owned score semantics in each request's physical coordinates.
+
+    Half-open query ranges select current queries; (0, 0) skips a request.
+    Probability denominators use eligible causal keys after sink/recent masks.
+    """
+
+    query_ranges: tuple[tuple[int, int], ...]
+    mode: str
+    candidate_start: int = 0
+    recent_keep_tokens: int = 0
+
+
 @dataclass(frozen=True)
 class PrefillComputeView:
     """Prefill metadata paired with exactly one physical payload layout."""
 
     meta: AttentionViewMeta
     payload: AttentionPayload
+    token_scores: torch.Tensor | None = None  # Temporary [batch, max_context] scores.
 
 
 @dataclass(frozen=True)
@@ -1119,6 +1134,10 @@ def build_prefill_compute_view(
             payload=payload,
         )
 
+    def prefill_score_request(self, layer_idx: int, seqs) -> PrefillScoreRequest | None:
+        """Describe optional scores before attention without changing score state."""
+        return None
+
     def collect_prefill_attention_score(
         self,
         layer_idx: int,

src/sparsevllm/engine/cache_manager/h2o.py

@@ -15,7 +15,7 @@
 from sparsevllm.utils.context import get_context
 from sparsevllm.utils.profiler import profiler
 
-from .base import ExplicitKVPayload, PrefillComputeView
+from .base import ExplicitKVPayload, PrefillComputeView, PrefillScoreRequest
 from .snapkv import SnapKVCacheManager
 from .storage import ExplicitKVStorage
 
@@ -790,6 +790,15 @@ def prefill_score_ranges(
             ranges.append((batch_idx, seq, prompt_cache_len, score_start, score_end))
         return ranges
 
+    def prefill_score_request(self, layer_idx, seqs):
+        rows = self.prefill_score_ranges(layer_idx, seqs)
+        if not rows:
+            return None
+        ranges = [(0, 0)] * len(seqs)
+        for batch_idx, _seq, _cached, start, end in rows:
+            ranges[batch_idx] = (start, end)
+        return PrefillScoreRequest(tuple(ranges), self.config.sparse_prefill_score_mode)
+
     @torch.no_grad()
     def collect_prefill_attention_score(
         self,
@@ -815,7 +824,7 @@ def collect_prefill_attention_score(
         ranges = self.prefill_score_ranges(layer_idx, seqs)
         if not ranges:
             return None
-        if not isinstance(view.payload, ExplicitKVPayload):
+        if view.token_scores is None and not isinstance(view.payload, ExplicitKVPayload):
             raise TypeError(
                 "H2O prefill scoring requires ExplicitKVPayload, got "
                 f"{type(view.payload).__name__}."
@@ -861,7 +870,9 @@ def collect_prefill_attention_score(
             score_ends=score_ends_cpu,
         )
         max_context_len = max(context_lens)
-        if meta.attn_score is None:
+        if view.token_scores is not None:
+            step_score = view.token_scores.to(self._prefill_score_dtype())
+        elif meta.attn_score is None:
             step_score = self._prefill_step_score_buffer(
                 batch_size=len(seqs),
                 max_context_len=max_context_len,

src/sparsevllm/engine/cache_manager/snapkv.py

@@ -31,6 +31,7 @@
     ExplicitKVPayload,
     LayerBatchStates,
     PrefillComputeView,
+    PrefillScoreRequest,
     SparseSelection,
 )
 from .raw_kv_offload import RawKVOffloadBuffer
@@ -1258,6 +1259,18 @@ def _prefill_step_score_buffer(
             buffers[key] = buffer
         return buffer[:batch_size, :max_context_len]
 
+    def prefill_score_request(self, layer_idx, seqs):
+        rows = self._prefill_score_rows(layer_idx, seqs)
+        if not rows:
+            return None
+        ranges = [(0, 0)] * len(seqs)
+        for batch_idx, _seq, start, end in rows:
+            ranges[batch_idx] = (start, end)
+        return PrefillScoreRequest(
+            tuple(ranges), self.config.sparse_prefill_score_mode,
+            int(self.config.sink_keep_tokens), int(self.config.recent_keep_tokens),
+        )
+
     @torch.no_grad()
     def collect_prefill_attention_score(
         self,
@@ -1282,7 +1295,7 @@ def collect_prefill_attention_score(
         rows = self._prefill_score_rows(layer_idx, seqs)
         if not rows:
             return None
-        if not isinstance(view.payload, ExplicitKVPayload):
+        if view.token_scores is None and not isinstance(view.payload, ExplicitKVPayload):
             raise TypeError(
                 "SnapKV prefill scoring requires ExplicitKVPayload, got "
                 f"{type(view.payload).__name__}."
@@ -1320,20 +1333,25 @@ def collect_prefill_attention_score(
             max_context_len=max_context_len,
             device=q.device,
         )
-        self._run_prefill_score(
-            q,
-            payload.k_cache,
-            step_score,
-            meta,
-            b_start_loc,
-            b_prompt_cache_len,
-            max_score_len,
-            score_starts,
-            score_ends,
-            candidate_start=int(self.config.sink_keep_tokens),
-            recent_keep_tokens=int(self.config.recent_keep_tokens),
-            batch_indices=score_batch_indices,
-        )
+        if view.token_scores is not None:
+            step_score.copy_(view.token_scores[:, :max_context_len].index_select(
+                0, score_batch_indices.long(),
+            ))
+        else:
+            self._run_prefill_score(
+                q,
+                payload.k_cache,
+                step_score,
+                meta,
+                b_start_loc,
+                b_prompt_cache_len,
+                max_score_len,
+                score_starts,
+                score_ends,
+                candidate_start=int(self.config.sink_keep_tokens),
+                recent_keep_tokens=int(self.config.recent_keep_tokens),
+                batch_indices=score_batch_indices,
+            )
 
         for score_row_idx, (b_idx, seq, _score_start, _score_end) in enumerate(rows):
             context_len = int(context_lens[b_idx])