diff --git a/docs/en/configuration/runtime-parameter-semantics.md b/docs/en/configuration/runtime-parameter-semantics.md index 3081d8aa..acd5a624 100644 --- a/docs/en/configuration/runtime-parameter-semantics.md +++ b/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 diff --git a/docs/zh/configuration/runtime-parameter-semantics.md b/docs/zh/configuration/runtime-parameter-semantics.md index 8fee0783..f3f3bfeb 100644 --- a/docs/zh/configuration/runtime-parameter-semantics.md +++ b/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`。当前有 diff --git a/scripts/profiling/kernel_bench/benchmark_mla_prefill.py b/scripts/profiling/kernel_bench/benchmark_mla_prefill.py new file mode 100644 index 00000000..f76b66d6 --- /dev/null +++ b/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() diff --git a/src/sparsevllm/configs/runtime.py b/src/sparsevllm/configs/runtime.py index 67f84a0f..4ddd55e9 100644 --- a/src/sparsevllm/configs/runtime.py +++ b/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 diff --git a/src/sparsevllm/configs/scheduling.py b/src/sparsevllm/configs/scheduling.py index a2a24ac3..d5955cc0 100644 --- a/src/sparsevllm/configs/scheduling.py +++ b/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 " diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index 7d837ad1..d840d923 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/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, diff --git a/src/sparsevllm/engine/cache_manager/h2o.py b/src/sparsevllm/engine/cache_manager/h2o.py index 29e32e9d..3376e014 100644 --- a/src/sparsevllm/engine/cache_manager/h2o.py +++ b/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, diff --git a/src/sparsevllm/engine/cache_manager/snapkv.py b/src/sparsevllm/engine/cache_manager/snapkv.py index 2abb983e..2d9fb9cd 100644 --- a/src/sparsevllm/engine/cache_manager/snapkv.py +++ b/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]) diff --git a/src/sparsevllm/engine/cache_manager/standard.py b/src/sparsevllm/engine/cache_manager/standard.py index 783c0f7e..6d2dba52 100644 --- a/src/sparsevllm/engine/cache_manager/standard.py +++ b/src/sparsevllm/engine/cache_manager/standard.py @@ -37,6 +37,7 @@ ExplicitKVPayload, LayerBatchStates, PrefillComputeView, + PrefillScoreRequest, SparseSelection, ) from .prefix_cache_mixin import PrefixCacheMixin @@ -1134,6 +1135,14 @@ def finish_prefix_prune_scoring(self) -> torch.Tensor: raise RuntimeError("prefix-prune scoring forward produced no attention scores.") return state["score"] # type: ignore[return-value] + def prefill_score_request(self, layer_idx, seqs): + state = self._prefix_prune_scoring + if state is None: + return None + start, end = int(state["query_start"]), int(state["query_end"]) + return PrefillScoreRequest(((start, end),), "probability", + int(state["candidate_start"]), end - start) + @torch.no_grad() def collect_prefill_attention_score( self, @@ -1159,7 +1168,7 @@ def collect_prefill_attention_score( "prefix-prune query window length mismatch: " f"expected={query_end - query_start} actual={int(q.shape[0])}." ) - if not isinstance(view.payload, ExplicitKVPayload): + if view.token_scores is None and not isinstance(view.payload, ExplicitKVPayload): raise TypeError( "prefix-prune scoring requires explicit KV storage, got " f"{type(view.payload).__name__}." @@ -1170,25 +1179,28 @@ def collect_prefill_attention_score( "prefix-prune scoring currently requires an unpruned dense target path: " f"physical_context={context_len} logical_context={query_end}." ) - step_score = torch.zeros( - (1, context_len), dtype=torch.float32, device=q.device - ) - prefill_score_fwd( - q, - view.payload.k_cache, - step_score, - view.meta.req_indices, - b_start_loc, - view.meta.context_lens, - torch.tensor([query_start], dtype=torch.int32, device=q.device), - query_end - query_start, - view.meta.active_slots, - torch.tensor([query_start], dtype=torch.int32, device=q.device), - torch.tensor([query_end], dtype=torch.int32, device=q.device), - candidate_start=candidate_start, - recent_keep_tokens=query_end - query_start, - score_mode="probability", - ) + if view.token_scores is not None: + step_score = view.token_scores + else: + step_score = torch.zeros( + (1, context_len), dtype=torch.float32, device=q.device + ) + prefill_score_fwd( + q, + view.payload.k_cache, + step_score, + view.meta.req_indices, + b_start_loc, + view.meta.context_lens, + torch.tensor([query_start], dtype=torch.int32, device=q.device), + query_end - query_start, + view.meta.active_slots, + torch.tensor([query_start], dtype=torch.int32, device=q.device), + torch.tensor([query_end], dtype=torch.int32, device=q.device), + candidate_start=candidate_start, + recent_keep_tokens=query_end - query_start, + score_mode="probability", + ) score = step_score[0] accumulated = state.get("score") state["score"] = ( diff --git a/src/sparsevllm/engine/llm_engine.py b/src/sparsevllm/engine/llm_engine.py index d089e3b3..cee324ad 100644 --- a/src/sparsevllm/engine/llm_engine.py +++ b/src/sparsevllm/engine/llm_engine.py @@ -1264,6 +1264,8 @@ def worker_info( "max_num_batched_tokens", "prefill_schedule_policy", "engine_prefill_chunk_size", + "mla_prefill_history_chunk_size", + "mla_prefill_workspace_bytes", "long_prefill_offload_threshold", "sink_keep_tokens", "recent_keep_tokens", diff --git a/src/sparsevllm/kernels/external/sgl/fa3.py b/src/sparsevllm/kernels/external/sgl/fa3.py index 2f425c38..4c34b7fb 100644 --- a/src/sparsevllm/kernels/external/sgl/fa3.py +++ b/src/sparsevllm/kernels/external/sgl/fa3.py @@ -611,8 +611,10 @@ def run_contiguous_explicit_varlen( cu_seqlens_k: torch.Tensor, max_seqlen_q: int, max_seqlen_k: int, - ) -> torch.Tensor: - """Run causal varlen attention over packed contiguous KV.""" + causal: bool = True, + return_softmax_lse: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run varlen attention over packed contiguous KV.""" args: list[object] = [ q, @@ -639,7 +641,7 @@ def run_contiguous_explicit_varlen( None, None, self.softmax_scale, - True, + causal, -1, -1, ] @@ -648,6 +650,11 @@ def run_contiguous_explicit_varlen( result: Sequence[torch.Tensor] = self._op(*args) if not result or result[0].data_ptr() != output.data_ptr(): raise RuntimeError("sglang-kernel FA3 did not write to the supplied output") + if return_softmax_lse: + lse = result[1] + if lse.dtype != torch.float32 or lse.shape != (q.shape[1], q.shape[0]): + raise RuntimeError("FA3 returned invalid packed prefill LSE.") + return output, lse return output diff --git a/src/sparsevllm/kernels/triton/mla/prefill.py b/src/sparsevllm/kernels/triton/mla/prefill.py new file mode 100644 index 00000000..b44367d3 --- /dev/null +++ b/src/sparsevllm/kernels/triton/mla/prefill.py @@ -0,0 +1,199 @@ +"""Bounded MLA attention partials and online output merging.""" + +from functools import lru_cache + +import torch +import triton +import triton.language as tl + +from sparsevllm.platforms import device_runtime + + +@lru_cache(maxsize=None) +def _launch_config(device_index, head_dim): + from sparsevllm.kernels.triton.context_flashattention_nopad import ( + _device_max_shared_memory, + select_context_attention_launch_config, + ) + + # Reuse the existing prefill tiles and shared-memory compatibility bound. + return select_context_attention_launch_config( + head_dim, + max_shared_memory=_device_max_shared_memory(device_index), + is_tesla="Tesla" in device_runtime.optional_device_name(device_index), + ) + + +@triton.jit +def _attention( + Q, + K, + V, + O, + L, + CQ, + CK, + q0: tl.constexpr, + q1: tl.constexpr, + k0: tl.constexpr, + k1: tl.constexpr, + v0: tl.constexpr, + v1: tl.constexpr, + TOTAL_Q, + H: tl.constexpr, + D: tl.constexpr, + SCALE: tl.constexpr, + CAUSAL: tl.constexpr, + M: tl.constexpr, + N: tl.constexpr, +): + block, head, batch = tl.program_id(0), tl.program_id(1), tl.program_id(2) + qs, qe = tl.load(CQ + batch), tl.load(CQ + batch + 1) + ks, ke = tl.load(CK + batch), tl.load(CK + batch + 1) + qi = block * M + tl.arange(0, M) + di = tl.arange(0, D) + q = tl.load( + Q + (qs + qi[:, None]) * q0 + head * q1 + di[None, :], qi[:, None] < qe - qs, 0 + ) + maximum = tl.full((M,), -float("inf"), tl.float32) + denominator = tl.zeros((M,), tl.float32) + acc = tl.zeros((M, D), tl.float32) + # Match the causal traversal of context_flashattention_nopad: later keys + # cannot contribute to this query tile. Ragged-batch padding does no work. + end = ke - ks + if CAUSAL: + end = tl.minimum(end, (block + 1) * M + (ke - ks) - (qe - qs)) + end = tl.where(block * M < qe - qs, tl.maximum(end, 0), 0) + for start in range(0, end, N): + ki = start + tl.arange(0, N) + k = tl.load( + K + (ks + ki[None, :]) * k0 + head * k1 + di[:, None], + ki[None, :] < ke - ks, + 0, + ) + z = tl.dot(q, k) * (SCALE * 1.4426950408889634) + valid = ki[None, :] < ke - ks + if CAUSAL: + valid = valid & (ki[None, :] <= qi[:, None] + (ke - ks) - (qe - qs)) + z = tl.where(valid, z, -float("inf")) + updated = tl.maximum(maximum, tl.max(z, 1)) + safe = tl.where(updated == -float("inf"), 0.0, updated) + p = tl.exp2(z - safe[:, None]) + alpha = tl.exp2(maximum - safe) + v = tl.load( + V + (ks + ki[:, None]) * v0 + head * v1 + di[None, :], + ki[:, None] < ke - ks, + 0, + ) + acc = acc * alpha[:, None] + acc = tl.dot(p.to(v.dtype), v, acc) + denominator = denominator * alpha + tl.sum(p, 1) + maximum = updated + out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None] + tl.store( + O + ((qs + qi[:, None]) * H + head) * D + di[None, :], + out, + qi[:, None] < qe - qs, + ) + # The merge and score APIs consume natural-log LSE, despite exp2 internally. + lse = (maximum + tl.log2(denominator)) * 0.6931471805599453 + tl.store(L + head * TOTAL_Q + qs + qi, lse, qi < qe - qs) + + +def attention_partial(q, k, v, cu_q, cu_k, max_q, max_k, *, scale, causal): + output = torch.empty(q.shape, dtype=q.dtype, device=q.device) + lse = torch.empty((q.shape[1], q.shape[0]), device=q.device, dtype=torch.float32) + block_m, block_n, num_warps, num_stages = _launch_config(q.device.index, q.shape[2]) + # Observation-window chunks have few queries but can scan long history. + # Retain the original narrow, pipelined tile for these small batches. + if max_q < 1024: + block_m, block_n, num_warps, num_stages = 32, 64, 4, 3 + _attention[(triton.cdiv(max_q, block_m), q.shape[1], cu_q.numel() - 1)]( + q, + k, + v, + output, + lse, + cu_q, + cu_k, + *q.stride()[:2], + *k.stride()[:2], + *v.stride()[:2], + q.shape[0], + q.shape[1], + q.shape[2], + scale, + causal, + block_m, + block_n, + num_warps=num_warps, + num_stages=num_stages, + ) + return output, lse + + +@triton.jit +def _merge( + O, + L, + P, + PL, + o0: tl.constexpr, + o1: tl.constexpr, + l0, + l1, + p0: tl.constexpr, + p1: tl.constexpr, + pl0, + pl1, + QN, + H: tl.constexpr, + D: tl.constexpr, + BD: tl.constexpr, + ROWS: tl.constexpr = 8, +): + row = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + qi, head = row // H, row % H + valid = qi < QN + a = tl.load(L + head * l0 + qi * l1, valid, -float("inf")) + b = tl.load(PL + head * pl0 + qi * pl1, valid, -float("inf")) + m = tl.maximum(a, b) + safe = tl.where(m == -float("inf"), 0.0, m) + wa, wb = tl.exp(a - safe), tl.exp(b - safe) + total = wa + wb + denom = tl.where(total > 0, total, 1.0) + d = tl.arange(0, BD) + old = tl.load( + O + qi[:, None] * o0 + head[:, None] * o1 + d[None, :], + valid[:, None] & (d[None, :] < D), + 0, + ).to(tl.float32) + partial = tl.load( + P + qi[:, None] * p0 + head[:, None] * p1 + d[None, :], + valid[:, None] & (d[None, :] < D), + 0, + ).to(tl.float32) + result = old * (wa / denom)[:, None] + partial * (wb / denom)[:, None] + tl.store( + O + qi[:, None] * o0 + head[:, None] * o1 + d[None, :], + result, + valid[:, None] & (d[None, :] < D), + ) + tl.store(L + head * l0 + qi * l1, safe + tl.log(total), valid) + + +def merge_partial(output, lse, partial, partial_lse): + _merge[(triton.cdiv(output.shape[0] * output.shape[1], 8),)]( + output, + lse, + partial, + partial_lse, + *output.stride()[:2], + *lse.stride(), + *partial.stride()[:2], + *partial_lse.stride(), + output.shape[0], + output.shape[1], + output.shape[2], + triton.next_power_of_2(output.shape[2]), + ) diff --git a/src/sparsevllm/kernels/triton/mla/prefill_score.py b/src/sparsevllm/kernels/triton/mla/prefill_score.py new file mode 100644 index 00000000..ae150e1e --- /dev/null +++ b/src/sparsevllm/kernels/triton/mla/prefill_score.py @@ -0,0 +1,183 @@ +"""Sparse observation scores over one bounded expanded or latent key block.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _score( + Q, + K, + QR, + KR, + LSE, + OUT, + STATS, + q0: tl.constexpr, + q1: tl.constexpr, + k0: tl.constexpr, + k1: tl.constexpr, + qr0: tl.constexpr, + qr1: tl.constexpr, + kr0: tl.constexpr, + QN, + KN, + H: tl.constexpr, + D: tl.constexpr, + QSTART, + KSTART, + CSTART, + CEND, + SCALE: tl.constexpr, + MODE: tl.constexpr, + NK, + M: tl.constexpr = 32, + N: tl.constexpr = 64, +): + head, qb, kb = tl.program_id(0), tl.program_id(1), tl.program_id(2) + qi = qb * M + tl.arange(0, M) + ki = kb * N + tl.arange(0, N) + d = tl.arange(0, D) + q = tl.load(Q + qi[:, None] * q0 + head * q1 + d[None, :], qi[:, None] < QN, 0) + kh = head if MODE == 0 else 0 + k = tl.load(K + ki[None, :] * k0 + kh * k1 + d[:, None], ki[None, :] < KN, 0) + z = tl.dot(q, k) + if MODE != 0: + r = tl.arange(0, 64) + qr = tl.load( + QR + qi[:, None] * qr0 + head * qr1 + r[None, :], qi[:, None] < QN, 0 + ) + kr = tl.load(KR + ki[None, :] * kr0 + r[:, None], ki[None, :] < KN, 0) + z += tl.dot(qr, kr) + key_valid = (ki < KN) & (KSTART + ki >= CSTART) & (KSTART + ki < CEND) + valid = ( + (qi[:, None] < QN) + & key_valid[None, :] + & (QSTART + qi[:, None] >= KSTART + ki[None, :]) + ) + if MODE == 0: + maximum = tl.max(tl.where(valid, z, -float("inf")), 0) + tl.atomic_max(OUT + ki, maximum, key_valid) + elif MODE == 1: + z = tl.where(valid, z * SCALE, -float("inf")) + maximum = tl.max(z, 1) + safe = tl.where(maximum == -float("inf"), 0.0, maximum) + total = tl.sum(tl.exp(z - safe[:, None]), 1) + tl.store(STATS + (head * NK + kb) * QN + qi, maximum + tl.log(total), qi < QN) + else: + lse = tl.load(LSE + head * QN + qi, qi < QN, float("inf")) + p = tl.where( + valid & (lse[:, None] != -float("inf")), + tl.exp(z * SCALE - lse[:, None]), + 0.0, + ) + mass = tl.sum(p, 0) / QN + tl.atomic_add(OUT + head * KN + ki, mass, key_valid) + + +@triton.jit +def _merge_stats(STATS, LSE, QN, NK, BK: tl.constexpr, BQ: tl.constexpr = 16): + head, qb = tl.program_id(0), tl.program_id(1) + qi = qb * BQ + tl.arange(0, BQ) + ki = tl.arange(0, BK) + values = tl.load( + STATS + (head * NK + ki[None, :]) * QN + qi[:, None], + (qi[:, None] < QN) & (ki[None, :] < NK), + -float("inf"), + ) + old = tl.load(LSE + head * QN + qi, qi < QN, -float("inf")) + maximum = tl.maximum(tl.max(values, 1), old) + safe = tl.where(maximum == -float("inf"), 0.0, maximum) + total = tl.sum(tl.exp(values - safe[:, None]), 1) + tl.exp(old - safe) + tl.store(LSE + head * QN + qi, safe + tl.log(total), qi < QN) + + +@triton.jit +def _reduce_heads( + HEAD, OUT, H: tl.constexpr, KN, BH: tl.constexpr, BN: tl.constexpr = 128 +): + k = tl.program_id(0) * BN + tl.arange(0, BN) + h = tl.arange(0, BH) + values = tl.load( + HEAD + h[:, None] * KN + k[None, :], (h[:, None] < H) & (k[None, :] < KN), 0.0 + ) + tl.store(OUT + k, tl.max(values, 0), k < KN) + + +def score_block( + q, + k, + output, + lse, + *, + query_start, + key_start, + candidate_start, + candidate_end, + scale, + mode, + rope_q=None, + rope_k=None, +): + """Score expanded K logits, or latent probabilities with accumulated LSE.""" + qn, heads, dim = q.shape + kn = k.shape[0] + if not qn or not kn: + return + latent = mode != "logits" + if latent and ( + dim != 512 + or rope_q is None + or rope_k is None + or rope_q.shape[-1] != 64 + or rope_k.shape[-1] != 64 + ): + raise ValueError( + "MLA latent scoring requires 512 latent and 64 RoPE dimensions." + ) + nk = triton.cdiv(kn, 64) + stats = ( + torch.empty((heads, nk, qn), dtype=torch.float32, device=q.device) + if mode == "stats" + else output + ) + target = ( + torch.zeros((heads, kn), dtype=torch.float32, device=q.device) + if mode == "probability" + else output + ) + qr, kr = (rope_q, rope_k) if latent else (q, k) + _score[(heads, triton.cdiv(qn, 32), nk)]( + q, + k, + qr, + kr, + lse, + target, + stats, + *q.stride()[:2], + *k.stride()[:2], + *qr.stride()[:2], + kr.stride(0), + qn, + kn, + heads, + dim, + query_start, + key_start, + candidate_start, + candidate_end, + scale, + {"logits": 0, "stats": 1, "probability": 2}[mode], + nk, + num_warps=4, + ) + if mode == "stats": + _merge_stats[(heads, triton.cdiv(qn, 16))]( + stats, lse, qn, nk, triton.next_power_of_2(nk) + ) + elif mode == "probability": + _reduce_heads[(triton.cdiv(kn, 128),)]( + target, output, heads, kn, triton.next_power_of_2(heads) + ) diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index 56e2412a..9b90cb72 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -1,266 +1,35 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import Callable +from collections.abc import Callable +from dataclasses import replace import torch from sparsevllm.engine.cache_manager.base import ( AttentionKeyComputeView, - AttentionViewMeta, DecodeComputeView, - ExplicitKVPayload, - MlaLatentWrite, MlaLatentPayload, + MlaLatentWrite, PrefillComputeView, ) -from sparsevllm.layers.attention_backend import TritonAttentionBackend from sparsevllm.operators.mla_attention import ( MlaAttentionOpSpec, MlaAttentionProvider, resolve_mla_attention_provider, ) -from sparsevllm.kernels.triton.mla import ( - gather_latent_history, - validate_gather_metadata, +from sparsevllm.operators.mla_prefill import ( + ChunkedMlaPrefill, + estimate_mla_prefill_workspace_bytes, ) from sparsevllm.utils.context import get_context -@dataclass(frozen=True, slots=True) -class MlaPrefillHistory: - """Gathered full history and its packed logical coordinates.""" - - gathered_latent: torch.Tensor - gathered_rope: torch.Tensor - packed_offsets: torch.Tensor - packed_cu_seqlens: torch.Tensor - packed_slots: torch.Tensor - local_req_indices: torch.Tensor - context_lens: torch.Tensor - context_lengths: tuple[int, ...] - max_context_len: int - required_workspace_bytes: int - - @property - def visible_tokens(self) -> int: - return int(self.gathered_latent.shape[0]) - - -@dataclass(frozen=True, slots=True) -class MlaPrefillWorkset: - """Full-history MLA buffers ready for ordinary 256-wide attention.""" - - history: MlaPrefillHistory - expanded_k: torch.Tensor - expanded_v: torch.Tensor - - -@dataclass(frozen=True, slots=True) -class _MlaPrefillPlan: - """Step-local packing metadata shared by every MLA layer.""" - - validation_scope: object - source_active_slots: torch.Tensor - source_req_indices: torch.Tensor - source_context_lens: torch.Tensor - source_max_context_len: int | None - source_query_tokens: int - cache_slot_count: int - packed_offsets: torch.Tensor - packed_cu_seqlens: torch.Tensor - packed_slots: torch.Tensor - local_req_indices: torch.Tensor - context_lengths: tuple[int, ...] - total_visible_tokens: int - max_context_len: int - required_workspace_bytes: int - - def matches( - self, - validation_scope: object, - meta: AttentionViewMeta, - cache_slot_count: int, - query_tokens: int, - ) -> bool: - return ( - self.validation_scope is validation_scope - and self.source_active_slots is meta.active_slots - and self.source_req_indices is meta.req_indices - and self.source_context_lens is meta.context_lens - and self.source_max_context_len == meta.max_context_len - and self.source_query_tokens == int(query_tokens) - and self.cache_slot_count == int(cache_slot_count) - ) - - -@dataclass(frozen=True, slots=True) -class _MlaPrefillQueryPlan: - """Step-local query-packing validation shared by every MLA layer.""" - - validation_scope: object - source_context_lens: torch.Tensor - query_tokens: int - max_query_len: int - - def matches( - self, - validation_scope: object, - context_lens: torch.Tensor, - query_tokens: int, - ) -> bool: - return ( - self.validation_scope is validation_scope - and self.source_context_lens is context_lens - and self.query_tokens == int(query_tokens) - ) - - -def _host_int_values(tensor: torch.Tensor) -> tuple[int, ...]: - """Synchronize an integer tensor once inside a validation scope.""" - - return tuple(int(value) for value in tensor.tolist()) - - -def estimate_mla_prefill_workspace_bytes( - *, - total_visible_tokens: int, - query_tokens: int, - batch_size: int, - max_context_len: int, - local_q_heads: int, - kv_lora_rank: int, - rope_dim: int, - qk_head_dim: int, - value_head_dim: int, - hidden_size: int, - projection_chunk_size: int, - activation_dtype: torch.dtype, - cache_dtype: torch.dtype, -) -> int: - """Bound the peak modeled transient storage for MLA full-history prefill.""" - - values = { - "total_visible_tokens": total_visible_tokens, - "query_tokens": query_tokens, - "batch_size": batch_size, - "max_context_len": max_context_len, - "local_q_heads": local_q_heads, - "kv_lora_rank": kv_lora_rank, - "rope_dim": rope_dim, - "qk_head_dim": qk_head_dim, - "value_head_dim": value_head_dim, - "hidden_size": hidden_size, - "projection_chunk_size": projection_chunk_size, - } - for name, value in values.items(): - if int(value) < 0: - raise ValueError(f"{name} must be non-negative, got {value}.") - positive_values = ( - "total_visible_tokens", - "query_tokens", - "batch_size", - "local_q_heads", - "kv_lora_rank", - "rope_dim", - "qk_head_dim", - "value_head_dim", - "hidden_size", - "projection_chunk_size", - ) - for name in positive_values: - if int(values[name]) == 0: - raise ValueError(f"{name} must be positive, got 0.") - if int(query_tokens) > int(total_visible_tokens): - raise ValueError( - "query_tokens cannot exceed total_visible_tokens, got " - f"{query_tokens} > {total_visible_tokens}." - ) - qk_nope_head_dim = int(qk_head_dim) - int(rope_dim) - if qk_nope_head_dim <= 0: - raise ValueError( - "qk_head_dim must be larger than rope_dim, got " - f"{qk_head_dim} and {rope_dim}." - ) - - cache_element_size = torch.empty((), dtype=cache_dtype).element_size() - activation_element_size = torch.empty( - (), - dtype=activation_dtype, - ).element_size() - visible_tokens = int(total_visible_tokens) - current_tokens = int(query_tokens) - heads = int(local_q_heads) - projected_width = qk_nope_head_dim + int(value_head_dim) - gathered_bytes = ( - visible_tokens - * (int(kv_lora_rank) + int(rope_dim)) - * cache_element_size - ) - projected_bytes = ( - visible_tokens * heads * projected_width * activation_element_size - ) - projection_scratch_bytes = 0 - if visible_tokens > int(projection_chunk_size): - projection_scratch_bytes = ( - min(visible_tokens, int(projection_chunk_size)) - * heads - * projected_width - * activation_element_size - ) - expanded_k_bytes = ( - visible_tokens - * heads - * int(qk_head_dim) - * activation_element_size - ) - attention_output_bytes = ( - current_tokens - * heads - * int(value_head_dim) - * activation_element_size - ) - output_projection_scratch_bytes = ( - min(current_tokens, int(projection_chunk_size)) - * int(hidden_size) - * activation_element_size - ) - kv_projection_phase_bytes = ( - gathered_bytes + projected_bytes + projection_scratch_bytes - ) - attention_phase_bytes = ( - gathered_bytes - + projected_bytes - + expanded_k_bytes - + attention_output_bytes - ) - output_projection_phase_bytes = ( - attention_output_bytes + output_projection_scratch_bytes - ) - metadata_values = ( - int(batch_size) * int(max_context_len) - + 2 * int(batch_size) - + int(max_context_len) - ) - metadata_bytes = ( - metadata_values * torch.empty((), dtype=torch.int32).element_size() - ) - return int( - max( - kv_projection_phase_bytes, - attention_phase_bytes, - output_projection_phase_bytes, - ) - + metadata_bytes - ) - - class MLAAttention: """Semantic MLA execution over tagged cache views. Model code owns projection weights, query absorption, and V reconstruction. - This object owns provider binding, decode workspace, full-history gathering, - and reuse of the existing 256-wide prefill attention backend. + This object binds providers and coordinates bounded history attention + and sparse scoring over cache-manager-owned storage. """ def __init__( @@ -271,6 +40,7 @@ def __init__( prefill_workspace_bytes: int, hidden_size: int, projection_chunk_size: int, + history_chunk_size: int = 16384, ) -> None: self.spec = spec self.provider = provider @@ -301,15 +71,14 @@ def __init__( "The existing prefill backend requires equal QK/value widths, " f"got {self.spec.qk_head_dim}/{self.spec.value_head_dim}." ) - self.prefill_backend = TritonAttentionBackend() - self._prefill_plan: _MlaPrefillPlan | None = None - self._prefill_query_plan: _MlaPrefillQueryPlan | None = None - self._key_materializer_bindings: dict[tuple[int, int], tuple[object, Callable]] = {} + self.chunked_prefill = ChunkedMlaPrefill(spec, provider, history_chunk_size) + self._key_materializer_bindings: dict[ + tuple[int, int], tuple[object, Callable] + ] = {} def release_cache_runtime_bindings(self, cache_manager: object) -> None: """Drop layer bindings owned by a retiring cache runtime.""" - self._prefill_plan = None - self._prefill_query_plan = None + self.chunked_prefill.clear() for key, binding in tuple(self._key_materializer_bindings.items()): if binding[0] is cache_manager: del self._key_materializer_bindings[key] @@ -324,7 +93,8 @@ def bind( prefill_workspace_bytes: int, hidden_size: int, projection_chunk_size: int, - ) -> "MLAAttention": + history_chunk_size: int = 16384, + ) -> MLAAttention: provider = resolve_mla_attention_provider( spec, device=device, @@ -336,15 +106,12 @@ def bind( prefill_workspace_bytes=prefill_workspace_bytes, hidden_size=hidden_size, projection_chunk_size=projection_chunk_size, + history_chunk_size=history_chunk_size, ) @property def device(self) -> torch.device: - return torch.device(getattr(self.provider, "device")) - - @property - def supports_explicit_prefill(self) -> bool: - return bool(getattr(self.provider, "supports_explicit_prefill", False)) + return torch.device(self.provider.device) def _require_mla_payload( self, @@ -355,8 +122,7 @@ def _require_mla_payload( payload = view.payload if not isinstance(payload, MlaLatentPayload): raise TypeError( - f"{operation} requires MlaLatentPayload, got " - f"{type(payload).__name__}." + f"{operation} requires MlaLatentPayload, got {type(payload).__name__}." ) for name, tensor, width in ( ("latent_cache", payload.latent_cache, self.spec.kv_lora_rank), @@ -379,266 +145,6 @@ def _require_mla_payload( raise ValueError("MLA latent and RoPE caches must have equal slots.") return payload - def _get_prefill_plan( - self, - meta: AttentionViewMeta, - *, - cache_slot_count: int, - query_tokens: int, - ) -> _MlaPrefillPlan: - cached = self._prefill_plan - validation_scope = get_context().attention_validation_scope - if cached is not None and cached.matches( - validation_scope, - meta, - cache_slot_count, - query_tokens, - ): - return cached - - metadata = { - "active_slots": meta.active_slots, - "req_indices": meta.req_indices, - "context_lens": meta.context_lens, - } - for name, tensor in metadata.items(): - if tensor.device != self.device: - raise ValueError( - f"{name} is on {tensor.device}, expected {self.device}." - ) - if tensor.dtype != torch.int32: - raise TypeError( - f"{name} must use {torch.int32}, got {tensor.dtype}." - ) - if meta.context_lens.ndim != 1 or meta.context_lens.numel() == 0: - raise ValueError( - "MLA prefill context_lens must be a non-empty 1D tensor." - ) - batch_size = int(meta.context_lens.numel()) - if meta.active_slots.ndim != 2: - raise ValueError( - "MLA prefill active_slots must have shape " - "[rows, max_context_len]." - ) - if meta.req_indices.shape != (batch_size,): - raise ValueError( - f"MLA prefill req_indices must have shape ({batch_size},), " - f"got {tuple(meta.req_indices.shape)}." - ) - if batch_size > self.max_batch_size: - raise ValueError( - "MLA prefill batch exceeds the bound operator capacity: " - f"batch={batch_size} max_batch_size={self.max_batch_size}." - ) - - lengths = _host_int_values(meta.context_lens) - if any(length < 0 for length in lengths): - raise ValueError( - f"MLA prefill context lengths must be non-negative: {lengths}." - ) - total_visible_tokens = int(sum(lengths)) - if total_visible_tokens <= 0: - raise ValueError("MLA prefill requires at least one visible token.") - max_context_len = int(max(lengths)) - if ( - meta.max_context_len is not None - and int(meta.max_context_len) < max_context_len - ): - raise ValueError( - "MLA prefill max_context_len is smaller than an actual context: " - f"declared={meta.max_context_len} actual={max_context_len}." - ) - required_bytes = estimate_mla_prefill_workspace_bytes( - total_visible_tokens=total_visible_tokens, - query_tokens=query_tokens, - batch_size=batch_size, - max_context_len=max_context_len, - local_q_heads=self.spec.local_q_heads, - kv_lora_rank=self.spec.kv_lora_rank, - rope_dim=self.spec.rope_dim, - qk_head_dim=self.spec.qk_head_dim, - value_head_dim=self.spec.value_head_dim, - hidden_size=self.hidden_size, - projection_chunk_size=self.projection_chunk_size, - activation_dtype=self.spec.activation_dtype, - cache_dtype=self.spec.cache_dtype, - ) - if required_bytes > self.prefill_workspace_bytes: - raise MemoryError( - "MLA full-history prefill workspace exceeds its configured " - f"budget: required={required_bytes} bytes budget=" - f"{self.prefill_workspace_bytes} bytes visible_tokens=" - f"{total_visible_tokens} local_heads={self.spec.local_q_heads}." - ) - - packed_starts: list[int] = [] - cursor = 0 - for length in lengths: - packed_starts.append(cursor) - cursor += length - packed_offsets = torch.tensor( - packed_starts, - dtype=torch.int32, - device=self.device, - ) - packed_cu_seqlens = torch.tensor( - (*packed_starts, total_visible_tokens), - dtype=torch.int32, - device=self.device, - ) - local_req_indices = torch.arange( - batch_size, - dtype=torch.int32, - device=self.device, - ) - positions = torch.arange( - max_context_len, - dtype=torch.int32, - device=self.device, - ) - packed_slots = packed_offsets[:, None] + positions[None, :] - validate_gather_metadata( - meta.active_slots, - meta.req_indices, - meta.context_lens, - packed_offsets, - cache_slot_count=cache_slot_count, - output_capacity=total_visible_tokens, - max_context_len=max_context_len, - ) - plan = _MlaPrefillPlan( - validation_scope=validation_scope, - source_active_slots=meta.active_slots, - source_req_indices=meta.req_indices, - source_context_lens=meta.context_lens, - source_max_context_len=meta.max_context_len, - source_query_tokens=int(query_tokens), - cache_slot_count=int(cache_slot_count), - packed_offsets=packed_offsets, - packed_cu_seqlens=packed_cu_seqlens, - packed_slots=packed_slots, - local_req_indices=local_req_indices, - context_lengths=lengths, - total_visible_tokens=total_visible_tokens, - max_context_len=max_context_len, - required_workspace_bytes=required_bytes, - ) - self._prefill_plan = plan - return plan - - def prepare_prefill_history( - self, - view: PrefillComputeView, - *, - query_tokens: int, - ) -> MlaPrefillHistory: - if not isinstance(view, PrefillComputeView): - raise TypeError( - "MLA prefill requires PrefillComputeView, got " - f"{type(view).__name__}." - ) - payload = self._require_mla_payload(view, operation="MLA prefill") - meta = view.meta - plan = self._get_prefill_plan( - meta, - cache_slot_count=int(payload.latent_cache.shape[0]), - query_tokens=query_tokens, - ) - gathered_latent = torch.empty( - plan.total_visible_tokens, - self.spec.kv_lora_rank, - dtype=self.spec.cache_dtype, - device=self.device, - ) - gathered_rope = torch.empty( - plan.total_visible_tokens, - self.spec.rope_dim, - dtype=self.spec.cache_dtype, - device=self.device, - ) - gather_latent_history( - payload.latent_cache, - payload.rope_cache, - meta.active_slots, - meta.req_indices, - meta.context_lens, - plan.packed_offsets, - gathered_latent, - gathered_rope, - max_context_len=plan.max_context_len, - validate_metadata=False, - ) - return MlaPrefillHistory( - gathered_latent=gathered_latent, - gathered_rope=gathered_rope, - packed_offsets=plan.packed_offsets, - packed_cu_seqlens=plan.packed_cu_seqlens, - packed_slots=plan.packed_slots, - local_req_indices=plan.local_req_indices, - context_lens=meta.context_lens, - context_lengths=plan.context_lengths, - max_context_len=plan.max_context_len, - required_workspace_bytes=plan.required_workspace_bytes, - ) - - def bind_prefill_kv( - self, - history: MlaPrefillHistory, - *, - expanded_k: torch.Tensor, - expanded_v: torch.Tensor, - ) -> MlaPrefillWorkset: - if not isinstance(history, MlaPrefillHistory): - raise TypeError( - "bind_prefill_kv requires MlaPrefillHistory, got " - f"{type(history).__name__}." - ) - if ( - history.gathered_latent.device != self.device - or history.gathered_rope.device != self.device - ): - raise ValueError("MLA prefill history is on the wrong device.") - if ( - history.gathered_latent.dtype != self.spec.cache_dtype - or history.gathered_rope.dtype != self.spec.cache_dtype - ): - raise TypeError("MLA prefill history uses the wrong cache dtype.") - expected_k_shape = ( - history.visible_tokens, - self.spec.local_q_heads, - self.spec.qk_head_dim, - ) - expected_v_shape = ( - history.visible_tokens, - self.spec.local_q_heads, - self.spec.value_head_dim, - ) - for name, tensor, expected_shape in ( - ("expanded_k", expanded_k, expected_k_shape), - ("expanded_v", expanded_v, expected_v_shape), - ): - if tuple(tensor.shape) != expected_shape: - raise ValueError( - f"{name} must have shape {expected_shape}, got " - f"{tuple(tensor.shape)}." - ) - if tensor.device != self.device: - raise ValueError( - f"{name} is on {tensor.device}, expected {self.device}." - ) - if tensor.dtype != self.spec.activation_dtype: - raise TypeError( - f"{name} must use {self.spec.activation_dtype}, got " - f"{tensor.dtype}." - ) - if tensor.stride(-1) != 1: - raise ValueError(f"{name} must be contiguous in its last dimension.") - return MlaPrefillWorkset( - history=history, - expanded_k=expanded_k, - expanded_v=expanded_v, - ) - @torch.no_grad() def materialize_expanded_keys( self, @@ -715,149 +221,6 @@ def materialize_expanded_keys( self.spec.qk_head_dim, ) - def run_prefill( - self, - q: torch.Tensor, - workset: MlaPrefillWorkset, - *, - b_start_loc: torch.Tensor, - chunk_lens: torch.Tensor, - ) -> torch.Tensor: - history = workset.history - if q.ndim != 3: - raise ValueError( - "MLA prefill q must have shape [tokens, local_heads, 256], " - f"got {tuple(q.shape)}." - ) - expected_q_shape = ( - int(q.shape[0]), - self.spec.local_q_heads, - self.spec.qk_head_dim, - ) - if tuple(q.shape) != expected_q_shape: - raise ValueError( - f"MLA prefill q must have shape {expected_q_shape}, got " - f"{tuple(q.shape)}." - ) - if q.device != self.device or q.dtype != self.spec.activation_dtype: - raise TypeError( - "MLA prefill q must match the operator device/dtype: " - f"q={q.device}/{q.dtype} expected=" - f"{self.device}/{self.spec.activation_dtype}." - ) - query_tokens = int(q.shape[0]) - validation_scope = get_context().attention_validation_scope - batch_size = int(history.context_lens.numel()) - for name, tensor in ( - ("b_start_loc", b_start_loc), - ("chunk_lens", chunk_lens), - ): - if tensor.shape != (batch_size,): - raise ValueError( - f"{name} must have shape ({batch_size},), got " - f"{tuple(tensor.shape)}." - ) - if tensor.device != self.device or tensor.dtype != torch.int32: - raise TypeError( - f"{name} must be int32 on {self.device}, got " - f"{tensor.device}/{tensor.dtype}." - ) - cached_query_plan = self._prefill_query_plan - if cached_query_plan is None or not cached_query_plan.matches( - validation_scope, - history.context_lens, - query_tokens, - ): - chunks = _host_int_values(chunk_lens) - starts = _host_int_values(b_start_loc) - expected_starts: list[int] = [] - cursor = 0 - for chunk in chunks: - expected_starts.append(cursor) - cursor += chunk - if starts != tuple(expected_starts) or cursor != query_tokens: - raise ValueError( - "MLA prefill query packing is inconsistent: " - f"starts={starts} expected_starts={expected_starts} " - f"chunk_tokens={cursor} q_tokens={query_tokens}." - ) - contexts = history.context_lengths - if any( - chunk <= 0 or chunk > context - for chunk, context in zip(chunks, contexts) - ): - raise ValueError( - "MLA prefill chunk lengths must be positive and no larger " - f"than their contexts: chunks={chunks} contexts={contexts}." - ) - self._prefill_query_plan = _MlaPrefillQueryPlan( - validation_scope=validation_scope, - source_context_lens=history.context_lens, - query_tokens=query_tokens, - max_query_len=max(chunks), - ) - - explicit_view = self.build_prefill_explicit_view(workset) - if self.supports_explicit_prefill: - run_prefill = getattr(self.provider, "run_explicit_prefill", None) - if not callable(run_prefill): - raise RuntimeError( - f"MLA provider {self.provider.name!r} advertises explicit " - "prefill without a run_explicit_prefill implementation." - ) - cu_seqlens_q = get_context().cu_seqlens_q - if cu_seqlens_q is None: - raise RuntimeError("MLA explicit prefill requires cu_seqlens_q.") - output = torch.empty( - (query_tokens, self.spec.local_q_heads, self.spec.value_head_dim), - dtype=q.dtype, - device=q.device, - ) - return run_prefill( - q, - explicit_view, - output, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=self._prefill_query_plan.max_query_len, - validation_scope=validation_scope, - ) - return self.prefill_backend.run_prefill( - q, - explicit_view, - b_start_loc=b_start_loc, - chunk_lens=chunk_lens, - max_input_len=history.max_context_len, - ) - - def build_prefill_explicit_view( - self, - workset: MlaPrefillWorkset, - ) -> PrefillComputeView: - """Expose the exact expanded KV view used by prefill score kernels.""" - - if not isinstance(workset, MlaPrefillWorkset): - raise TypeError( - "build_prefill_explicit_view requires MlaPrefillWorkset, got " - f"{type(workset).__name__}." - ) - history = workset.history - return PrefillComputeView( - meta=AttentionViewMeta( - active_slots=history.packed_slots, - req_indices=history.local_req_indices, - context_lens=history.context_lens, - max_context_len=history.max_context_len, - ), - payload=ExplicitKVPayload( - k_cache=workset.expanded_k, - v_cache=workset.expanded_v, - metadata={ - "layout": "mla_packed_varlen", - "cu_seqlens_k": history.packed_cu_seqlens, - }, - ), - ) - def run_decode( self, q_nope_absorbed: torch.Tensor, @@ -866,8 +229,7 @@ def run_decode( ) -> torch.Tensor: if not isinstance(view, DecodeComputeView): raise TypeError( - "MLA decode requires DecodeComputeView, got " - f"{type(view).__name__}." + f"MLA decode requires DecodeComputeView, got {type(view).__name__}." ) self._require_mla_payload(view, operation="MLA decode") output = torch.empty( @@ -877,9 +239,7 @@ def run_decode( ) context = get_context() valid_batch_size = ( - int(q_nope_absorbed.shape[0]) - if context.seqs is None - else len(context.seqs) + int(q_nope_absorbed.shape[0]) if context.seqs is None else len(context.seqs) ) return self.provider.run( q_nope_absorbed, @@ -950,54 +310,44 @@ def run_cached_attention( temp_slots = view.meta.temp_slots if context.cu_seqlens_q is None or context.cu_seqlens_q.numel() <= 1: return torch.empty_like(q) - history = self.prepare_prefill_history( + self._require_mla_payload(view, operation="chunked prefill") + plan = self.chunked_prefill.prepare( view, - query_tokens=int(q.shape[0]), - ) - qk_nope_head_dim = self.spec.qk_head_dim - self.spec.rope_dim - projected_width = qk_nope_head_dim + self.spec.value_head_dim - expanded = project_latent(history.gathered_latent).view( - history.visible_tokens, - self.spec.local_q_heads, - projected_width, + context.cu_seqlens_q, + context.attention_validation_scope, ) - expanded_k_nope, expanded_v = expanded.split( - [qk_nope_head_dim, self.spec.value_head_dim], - dim=-1, + request = cache_manager.prefill_score_request(layer_idx, context.seqs) + required = estimate_mla_prefill_workspace_bytes( + plan=plan, + spec=self.spec, + chunk_size=self.chunked_prefill.chunk_size, + hidden_size=self.hidden_size, + projection_chunk_size=self.projection_chunk_size, + score_request=request, ) - expanded_k = torch.empty( - ( - history.visible_tokens, - self.spec.local_q_heads, - self.spec.qk_head_dim, - ), - dtype=expanded.dtype, - device=expanded.device, - ) - expanded_k[..., :qk_nope_head_dim].copy_(expanded_k_nope) - expanded_k[..., qk_nope_head_dim:].copy_( - history.gathered_rope[:, None, :] - ) - workset = self.bind_prefill_kv( - history, - expanded_k=expanded_k, - expanded_v=expanded_v, + if required > self.prefill_workspace_bytes: + raise MemoryError( + f"MLA chunked prefill workspace exceeds budget: required={required} " + f"budget={self.prefill_workspace_bytes}. Reduce the token batch or history chunk size." + ) + output, attention_lse, scores = self.chunked_prefill.run( + q, + view, + context.cu_seqlens_q, + context.attention_validation_scope, + project_latent, + absorb_query, + request, ) b_start_loc = context.cu_seqlens_q[:-1] chunk_lens = context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1] - output = self.run_prefill( - q, - workset, - b_start_loc=b_start_loc, - chunk_lens=chunk_lens, - ) - explicit_view = self.build_prefill_explicit_view(workset) cache_manager.collect_prefill_attention_score( layer_idx, q, - explicit_view, + replace(view, token_scores=scores), b_start_loc=b_start_loc, chunk_lens=chunk_lens, + attention_lse=attention_lse, ) cache_manager.record_prefill_query( layer_idx, @@ -1033,9 +383,8 @@ def run_cached_attention( if temp_slots is not None and temp_slots.numel() > 0: cache_manager.release_layer_temp_slots(layer_idx, temp_slots) + __all__ = [ "MLAAttention", - "MlaPrefillHistory", - "MlaPrefillWorkset", "estimate_mla_prefill_workspace_bytes", ] diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 04b8ab98..9f6b505f 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -80,6 +80,7 @@ def build_glm4_moe_lite_mla_attention( context_capacity: int, projection_chunk_size: int, score_output: AttentionScoreKind = AttentionScoreKind.NONE, + history_chunk_size: int = 16384, ) -> MLAAttention: """Bind the one process-local MLA operator from explicit runtime inputs.""" @@ -106,6 +107,7 @@ def build_glm4_moe_lite_mla_attention( prefill_workspace_bytes=prefill_workspace_bytes, hidden_size=int(config.hidden_size), projection_chunk_size=projection_chunk_size, + history_chunk_size=history_chunk_size, ) @@ -922,6 +924,7 @@ def build_runtime_kwargs( engine_config.max_decoding_seqs, ), prefill_workspace_bytes=engine_config.mla_prefill_workspace_bytes, + history_chunk_size=engine_config.mla_prefill_history_chunk_size, decode_graph=decode_graph, context_capacity=decode_context_capacity, projection_chunk_size=engine_config.mlp_chunk_size, diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index e4eeaa21..f3c661e6 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -8,7 +8,6 @@ from sparsevllm.platforms import device_runtime from sparsevllm.engine.cache_manager.base import ( DecodeComputeView, - ExplicitKVPayload, MlaLatentPayload, PrefillComputeView, ) @@ -127,7 +126,6 @@ def kernel_request(self) -> AttentionKernelRequest: class MlaAttentionProvider: name = "" capabilities: AttentionKernelCapabilities - supports_explicit_prefill = False def run( self, @@ -544,7 +542,6 @@ class MlaSglFa3Provider(MlaTritonProvider): """SGL FA3 decode with the score-producing Triton path kept explicit.""" name = "sgl_fa3_sm90" - supports_explicit_prefill = True supports_decode_graph = True def __init__( @@ -635,134 +632,20 @@ def run( ) @torch.no_grad() - def run_explicit_prefill( - self, - q: torch.Tensor, - view: PrefillComputeView, - output: torch.Tensor, - *, - cu_seqlens_q: torch.Tensor, - max_seqlen_q: int, - validation_scope: object | None = None, - ) -> torch.Tensor: - if not isinstance(view, PrefillComputeView): - raise TypeError( - "MlaSglFa3Provider.run_explicit_prefill requires " - "PrefillComputeView, got " - f"{type(view).__name__}." - ) - if not isinstance(view.payload, ExplicitKVPayload): - raise TypeError( - "MLA explicit prefill requires ExplicitKVPayload, got " - f"{type(view.payload).__name__}." - ) - query_tokens = int(q.shape[0]) - expected_q_shape = ( - query_tokens, - self.spec.local_q_heads, - self.spec.qk_head_dim, - ) - if tuple(q.shape) != expected_q_shape: - raise ValueError( - f"q must have shape {expected_q_shape}, got {tuple(q.shape)}." - ) - expected_output_shape = ( - query_tokens, - self.spec.local_q_heads, - self.spec.value_head_dim, - ) - if tuple(output.shape) != expected_output_shape: - raise ValueError( - f"output must have shape {expected_output_shape}, got " - f"{tuple(output.shape)}." - ) - batch_size = int(view.meta.context_lens.numel()) - if batch_size > self.max_batch_size: - raise ValueError( - "MLA prefill batch exceeds provider capacity: " - f"batch={batch_size} max_batch_size={self.max_batch_size}." - ) - if cu_seqlens_q.shape != (batch_size + 1,): - raise ValueError( - f"cu_seqlens_q must have shape ({batch_size + 1},), got " - f"{tuple(cu_seqlens_q.shape)}." - ) - if cu_seqlens_q.device != self.device or cu_seqlens_q.dtype != torch.int32: - raise TypeError( - "cu_seqlens_q must be int32 on the provider device, got " - f"{cu_seqlens_q.device}/{cu_seqlens_q.dtype}." - ) - if not 0 < int(max_seqlen_q) <= query_tokens: - raise ValueError( - "max_seqlen_q must be in [1, query_tokens], got " - f"{max_seqlen_q} for {query_tokens}." - ) - payload = view.payload - tensors = { - "q": q, - "output": output, - "k_cache": payload.k_cache, - "v_cache": payload.v_cache, - } - for name, tensor in tensors.items(): - if tensor.device != self.device: - raise ValueError( - f"{name} is on {tensor.device}, expected {self.device}." - ) - expected_dtype = ( - self.spec.activation_dtype - ) - if tensor.dtype != expected_dtype: - raise TypeError( - f"{name} must use {expected_dtype}, got {tensor.dtype}." - ) - metadata = payload.metadata or {} - if metadata.get("layout") == "mla_packed_varlen": - cu_seqlens_k = metadata.get("cu_seqlens_k") - if not isinstance(cu_seqlens_k, torch.Tensor): - raise TypeError( - "MLA packed varlen prefill requires tensor cu_seqlens_k." - ) - if cu_seqlens_k.shape != (batch_size + 1,): - raise ValueError( - f"cu_seqlens_k must have shape ({batch_size + 1},), got " - f"{tuple(cu_seqlens_k.shape)}." - ) - if ( - cu_seqlens_k.device != self.device - or cu_seqlens_k.dtype != torch.int32 - ): - raise TypeError( - "cu_seqlens_k must be int32 on the provider device, got " - f"{cu_seqlens_k.device}/{cu_seqlens_k.dtype}." - ) - if view.meta.max_context_len is None: - raise ValueError( - "MLA packed varlen prefill requires max_context_len." - ) - self._record_runtime_kernel_path("sgl_fa3_prefill_contiguous") - return self.fa3.run_contiguous_explicit_varlen( - q, - payload.k_cache, - payload.v_cache, - output, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=int(max_seqlen_q), - max_seqlen_k=int(view.meta.max_context_len), - ) - self._record_runtime_kernel_path("sgl_fa3_prefill_paged") - return self.fa3.run_explicit_varlen( + def run_prefill_chunk(self, q, k, v, cu_q, cu_k, max_q, max_k, *, causal): + output = torch.empty((*q.shape[:2], v.shape[-1]), dtype=q.dtype, device=q.device) + self._record_runtime_kernel_path("sgl_fa3_prefill_contiguous") + return self.fa3.run_contiguous_explicit_varlen( q, - payload.k_cache, - payload.v_cache, - view.meta.active_slots, - view.meta.req_indices, - view.meta.context_lens, + k, + v, output, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=int(max_seqlen_q), - validation_scope=validation_scope, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max_q, + max_seqlen_k=max_k, + causal=causal, + return_softmax_lse=True, ) diff --git a/src/sparsevllm/operators/mla_prefill.py b/src/sparsevllm/operators/mla_prefill.py new file mode 100644 index 00000000..8d5a1ae6 --- /dev/null +++ b/src/sparsevllm/operators/mla_prefill.py @@ -0,0 +1,320 @@ +"""MLA history chunking and sparse score execution over tagged cache views.""" + +from dataclasses import dataclass + +import torch + +from sparsevllm.engine.cache_manager.base import AttentionViewMeta +from sparsevllm.kernels.triton.mla.prefill import attention_partial, merge_partial +from sparsevllm.kernels.triton.mla.prefill_score import score_block +from sparsevllm.utils.profiler import profiler + + +@dataclass +class PrefillPlan: + scope: object + meta: AttentionViewMeta + cu_q: torch.Tensor + contexts: tuple[int, ...] + query_starts: tuple[int, ...] + rows: tuple[int, ...] + current_slots: torch.Tensor + history_chunks: tuple[tuple[int, int, int, torch.Tensor], ...] + request_cu_q: tuple[torch.Tensor, ...] + + +def estimate_mla_prefill_workspace_bytes( + *, plan, spec, chunk_size, hidden_size, projection_chunk_size, score_request=None +): + """Modeled live tensors; opaque provider workspaces remain startup-profiled. + + Count a conservative overlap of projection, partial output, FP32 merge, + and sparse score state. Historical KV is bounded independently of context. + """ + tokens = plan.query_starts[-1] + historical = max((n for _, _, n, _ in plan.history_chunks), default=0) + block = max(tokens, historical) + element = spec.activation_dtype.itemsize + cache_element = spec.cache_dtype.itemsize + heads = spec.local_q_heads + width = spec.qk_head_dim - spec.rope_dim + spec.value_head_dim + gathered = block * (spec.kv_lora_rank + spec.rope_dim) * cache_element + projected = block * heads * width * element + projection_scratch = ( + min(block, projection_chunk_size) * heads * width * element + if block > projection_chunk_size + else 0 + ) + keys = block * heads * spec.qk_head_dim * element + # Current output, one partial, final dtype conversion and FP32 accumulator. + outputs = tokens * heads * spec.value_head_dim * (3 * element + 4) + lse = 2 * tokens * heads * 4 + metadata = 2 * (tokens + historical) * 8 + len(plan.history_chunks) * 8 + score_bytes = 0 + if score_request is not None: + observed = sum(end - start for start, end in score_request.query_ranges) + max_observed = max(end - start for start, end in score_request.query_ranges) + # Final token scores, absorbed observation Q, per-query LSE, and the + # largest block's probability/statistics scratch (not all blocks). + score_bytes = ( + len(plan.contexts) * max(plan.contexts) * 4 + + observed * heads * (spec.kv_lora_rank * element + 4) + + heads * max(block, chunk_size) * 4 + + heads * ((max(block, chunk_size) + 63) // 64) * max_observed * 4 + ) + projection_output = min(tokens, projection_chunk_size) * hidden_size * element + return ( + gathered + + projected + + projection_scratch + + keys + + outputs + + lse + + metadata + + score_bytes + + projection_output + ) + + +class ChunkedMlaPrefill: + def __init__(self, spec, provider, chunk_size): + if int(chunk_size) <= 0: + raise ValueError("MLA history chunk size must be positive.") + self.spec = spec + self.chunk_size = int(chunk_size) + self.plan = None + self.partial = getattr(provider, "run_prefill_chunk", None) + + def clear(self): + self.plan = None + + def prepare(self, view, cu_q, scope): + old = self.plan + meta = view.meta + if ( + old is not None + and old.scope is scope + and old.meta.active_slots is meta.active_slots + and old.meta.req_indices is meta.req_indices + and old.meta.context_lens is meta.context_lens + and old.cu_q is cu_q + ): + return old + contexts = tuple(int(x) for x in meta.context_lens.tolist()) + starts = tuple(int(x) for x in cu_q.tolist()) + rows = tuple(int(x) for x in meta.req_indices.tolist()) + if ( + len(starts) != len(contexts) + 1 + or len(rows) != len(contexts) + or starts[0] != 0 + ): + raise ValueError("Invalid MLA prefill request packing.") + current_slots, history, request_cu = [], [], [] + for i, (row, context) in enumerate(zip(rows, contexts)): + qn = starts[i + 1] - starts[i] + if qn <= 0 or context < qn or context > meta.active_slots.shape[1]: + raise ValueError("Invalid MLA query/context length.") + if not 0 <= row < meta.active_slots.shape[0]: + raise ValueError("Invalid MLA request row.") + cached = context - qn + current_slots.append(meta.active_slots[row, cached:context]) + request_cu.append( + torch.tensor([0, qn], dtype=torch.int32, device=cu_q.device) + ) + for offset in range(0, cached, self.chunk_size): + length = min(self.chunk_size, cached - offset) + cu_k = torch.tensor([0, length], dtype=torch.int32, device=cu_q.device) + history.append((i, offset, length, cu_k)) + self.plan = PrefillPlan( + scope, + meta, + cu_q, + contexts, + starts, + rows, + torch.cat(current_slots), + tuple(history), + tuple(request_cu), + ) + return self.plan + + @staticmethod + def gather(payload, slots): + index = slots.long() + return ( + payload.latent_cache.index_select(0, index).squeeze(1), + payload.rope_cache.index_select(0, index).squeeze(1), + ) + + def expand(self, latent, rope, project): + nope = self.spec.qk_head_dim - self.spec.rope_dim + expanded = project(latent).view( + -1, self.spec.local_q_heads, nope + self.spec.value_head_dim + ) + kn, v = expanded.split((nope, self.spec.value_head_dim), dim=-1) + k = torch.empty( + (*kn.shape[:2], self.spec.qk_head_dim), dtype=kn.dtype, device=kn.device + ) + k[..., :nope].copy_(kn) + k[..., nope:].copy_(rope[:, None, :]) + return k, v + + def attention(self, q, k, v, cu_q, cu_k, max_q, max_k, causal): + if self.partial is not None: + return self.partial(q, k, v, cu_q, cu_k, max_q, max_k, causal=causal) + return attention_partial( + q, + k, + v, + cu_q, + cu_k, + max_q, + max_k, + scale=self.spec.softmax_scale, + causal=causal, + ) + + def run(self, q, view, cu_q, scope, project, absorb, score_request=None): + plan = self.prepare(view, cu_q, scope) + if plan.query_starts[-1] != q.shape[0]: + raise ValueError("MLA query count differs from packed query metadata.") + scorer = ( + MlaPrefillScores(self, q, plan, score_request, absorb) + if score_request + else None + ) + latent, rope = self.gather(view.payload, plan.current_slots) + k, v = self.expand(latent, rope, project) + max_q = max(b - a for a, b in zip(plan.query_starts, plan.query_starts[1:])) + output, lse = self.attention(q, k, v, cu_q, cu_q, max_q, max_q, True) + if scorer is not None and not scorer.is_probability: + with profiler.record("prefill_token_score"): + for i, context in enumerate(plan.contexts): + a, b = plan.query_starts[i : i + 2] + scorer.consume(i, context - (b - a), k[a:b], mode="logits") + del latent, rope, k, v + if plan.history_chunks: + # FP32 accumulation avoids one BF16 rounding per history block. + output = output.float() + for i, offset, length, cu_k in plan.history_chunks: + a, b = plan.query_starts[i : i + 2] + slots = view.meta.active_slots[plan.rows[i], offset : offset + length] + latent, rope = self.gather(view.payload, slots) + k, v = self.expand(latent, rope, project) + partial, partial_lse = self.attention( + q[a:b], k, v, plan.request_cu_q[i], cu_k, b - a, length, False + ) + merge_partial(output[a:b], lse[:, a:b], partial, partial_lse) + if scorer is not None and not scorer.is_probability: + with profiler.record("prefill_token_score"): + scorer.consume(i, offset, k, mode="logits") + del latent, rope, k, v, partial, partial_lse + if scorer is not None and scorer.is_probability: + with profiler.record("prefill_token_score"): + scorer.finish_probability(view, lse) + return output.to(q.dtype), lse, None if scorer is None else scorer.output + + +class MlaPrefillScores: + def __init__(self, owner, q, plan, request, absorb): + if request.mode not in {"logits", "probability"}: + raise ValueError("Unsupported MLA prefill score mode.") + if len(request.query_ranges) != len(plan.contexts): + raise ValueError("MLA score ranges must cover the prefill batch.") + self.owner, self.plan, self.request = owner, plan, request + self.is_probability = request.mode == "probability" + self.full_normalizer = ( + request.candidate_start == 0 and request.recent_keep_tokens == 0 + ) + self.output = torch.full( + (len(plan.contexts), max(plan.contexts)), + -torch.inf if request.mode == "logits" else 0.0, + dtype=torch.float32, + device=q.device, + ) + self.queries, self.rope_queries, self.lse = [], [], [] + nope = owner.spec.qk_head_dim - owner.spec.rope_dim + for i, (start, end) in enumerate(request.query_ranges): + a, b = plan.query_starts[i : i + 2] + cached = plan.contexts[i] - (b - a) + if end < start: + raise ValueError("MLA observation range is reversed.") + if end > start and not cached <= start < end <= plan.contexts[i]: + raise ValueError("MLA observation window is outside current queries.") + observed = ( + q[a + start - cached : a + end - cached] if end > start else q[:0] + ) + self.queries.append( + absorb(observed[..., :nope]) + if self.is_probability and end > start + else observed + ) + self.rope_queries.append( + observed[..., nope:] if self.is_probability else None + ) + self.lse.append( + torch.full( + (q.shape[1], end - start), + -torch.inf, + dtype=torch.float32, + device=q.device, + ) + ) + + def consume(self, i, offset, keys, *, mode, rope=None): + start, end = self.request.query_ranges[i] + if end <= start: + return + candidate_end = max( + self.request.candidate_start, + self.plan.contexts[i] - self.request.recent_keep_tokens, + ) + if ( + offset >= candidate_end + or offset + keys.shape[0] <= self.request.candidate_start + ): + return + score_block( + self.queries[i], + keys, + self.output[i, offset : offset + keys.shape[0]], + self.lse[i], + query_start=start, + key_start=offset, + candidate_start=self.request.candidate_start, + candidate_end=candidate_end, + scale=self.owner.spec.softmax_scale, + mode=mode, + rope_q=self.rope_queries[i], + rope_k=rope, + ) + + def finish_probability(self, view, attention_lse): + if self.full_normalizer: + for i, (start, end) in enumerate(self.request.query_ranges): + a, b = self.plan.query_starts[i : i + 2] + cached = self.plan.contexts[i] - (b - a) + if end > start: + self.lse[i] = attention_lse[ + :, a + start - cached : a + end - cached + ].contiguous() + # Candidate-only softmax needs its own denominator; full-key scoring + # reuses the main attention LSE. Both scans stay in latent space. + modes = ("probability",) if self.full_normalizer else ("stats", "probability") + for mode in modes: + for i, (start, end) in enumerate(self.request.query_ranges): + if end <= start: + continue + context = self.plan.contexts[i] + candidate_end = max( + self.request.candidate_start, + context - self.request.recent_keep_tokens, + ) + for offset in range( + self.request.candidate_start, candidate_end, self.owner.chunk_size + ): + stop = min(offset + self.owner.chunk_size, candidate_end) + slots = view.meta.active_slots[self.plan.rows[i], offset:stop] + latent, rope = self.owner.gather(view.payload, slots) + self.consume(i, offset, latent[:, None, :], mode=mode, rope=rope) + del latent, rope diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index 327d7d53..17f3dc2d 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -214,6 +214,7 @@ def test_glm_runtime_kwargs_bind_shared_operators( max_num_seqs_in_batch=4, max_decoding_seqs=8, mla_prefill_workspace_bytes=1024, + mla_prefill_history_chunk_size=17, mlp_chunk_size=16, tiny_random=False, sparse_method=sparse_method, @@ -246,6 +247,7 @@ def test_glm_runtime_kwargs_bind_shared_operators( device=torch.device("cuda", 1), max_batch_size=8, prefill_workspace_bytes=1024, + history_chunk_size=runtime.mla_prefill_history_chunk_size, decode_graph=True, context_capacity=32768, projection_chunk_size=16, diff --git a/tests/test_glm_mla_sparse_methods.py b/tests/test_glm_mla_sparse_methods.py index f2003e3e..4499b619 100644 --- a/tests/test_glm_mla_sparse_methods.py +++ b/tests/test_glm_mla_sparse_methods.py @@ -343,3 +343,64 @@ def test_omnikv_observation_selects_mla_latent_active_slots(): payload.rope_cache[selected_slots, 0, 0], selected_slots.to(torch.bfloat16) + 100, ) + + +def test_snapkv_latent_score_handoff_preserves_final_window_and_accumulation(): + # Precomputed MLA scores must enter the same final-chunk accumulator and + # selection lifecycle as explicit KV, without attempting another QK kernel. + from sparsevllm.engine.cache_manager.base import ( + AttentionViewMeta, + PrefillComputeView, + ) + + manager, payload = _latent_snap_family_manager(SnapKVCacheManager, row_len=8) + manager.config = SimpleNamespace( + sparse_method="snapkv", + snapkv_num_full_layers=0, + snapkv_window_size=2, + sink_keep_tokens=1, + recent_keep_tokens=1, + decode_keep_tokens=2, + sparse_prefill_score_mode="logits", + sparse_attn_score_dtype="float32", + ) + manager._prefill_attn_score_accumulators = {} + manager._prefill_context_lens_cpu_by_layer = {0: (8,)} + seq = Sequence(list(range(8))) + seq.seq_id = 0 + seq.current_chunk_size = 4 + seq.num_prefilled_tokens = 0 + assert manager.prefill_score_request(0, [seq]) is None + seq.num_prefilled_tokens = 4 + request = manager.prefill_score_request(0, [seq]) + assert request.query_ranges == ((6, 8),) + assert request.candidate_start == request.recent_keep_tokens == 1 + score = torch.tensor([[-torch.inf, 1.0, 2.0, 9.0, 3.0, 8.0, 4.0, -torch.inf]]) + view = PrefillComputeView( + AttentionViewMeta( + manager.buffer_req_to_token_slots[0], + torch.tensor([0], dtype=torch.int32), + torch.tensor([8], dtype=torch.int32), + ), + payload, + score, + ) + set_context(True, cache_manager=manager, seqs=[seq]) + try: + with patch.object( + manager, + "_run_prefill_score", + side_effect=AssertionError("unexpected explicit QK"), + ): + manager.collect_prefill_attention_score( + 0, + torch.empty(4, 20, 256), + view, + b_start_loc=torch.tensor([0], dtype=torch.int32), + chunk_lens=torch.tensor([4], dtype=torch.int32), + ) + torch.testing.assert_close( + manager._prefill_attn_score_accumulators[0, 0], score[0] + ) + finally: + reset_context() diff --git a/tests/test_h2o_cache_manager.py b/tests/test_h2o_cache_manager.py index e13a226a..c7a36a87 100644 --- a/tests/test_h2o_cache_manager.py +++ b/tests/test_h2o_cache_manager.py @@ -14,6 +14,7 @@ DecodeComputeView, ExplicitKVPayload, LayerBatchStates, + MlaLatentPayload, PrefillComputeView, ) from sparsevllm.engine.cache_manager.h2o import H2OCacheManager @@ -501,7 +502,8 @@ def test_h2o_logit_prefill_score_rejects_nan_or_all_inf(): ) -def test_h2o_prefill_score_collection_accumulates_in_physical_coordinates(): +@pytest.mark.parametrize("precomputed", [False, True]) +def test_h2o_prefill_score_collection_accumulates_in_physical_coordinates(precomputed): manager = _manager_with_rows([6]) seq = _seq(0, 20, prefilled=8, chunk=2) manager._h2o_scores[(0, 0)] = torch.tensor([1.0, 2.0, 3.0, 4.0]) @@ -517,6 +519,16 @@ def test_h2o_prefill_score_collection_accumulates_in_physical_coordinates(): v_cache=torch.empty((16, 1, 1)), ), ) + if precomputed: + # The latent operator hands off scores; the manager still owns weighted + # accumulation in physical coordinates after previous KV compaction. + view = PrefillComputeView( + view.meta, MlaLatentPayload(torch.empty(16, 1, 512), torch.empty(16, 1, 64)), + token_scores=torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]]), + ) + request = manager.prefill_score_request(0, [seq]) + assert request.query_ranges == ((4, 6),) + assert request.candidate_start == request.recent_keep_tokens == 0 set_context(is_prefill=True, cache_manager=manager, seqs=[seq]) def fake_run_prefill_score( @@ -542,7 +554,7 @@ def fake_run_prefill_score( manager, "_run_prefill_score", side_effect=fake_run_prefill_score, - ): + ) as scorer: manager.collect_prefill_attention_score( 0, torch.empty((2, 1, 1)), @@ -551,6 +563,7 @@ def fake_run_prefill_score( chunk_lens=torch.tensor([2], dtype=torch.int32), ) + assert scorer.call_count == (0 if precomputed else 1) assert manager._h2o_scores[(0, 0)].tolist() == pytest.approx( [1.2, 2.4, 3.6, 4.8, 1.0, 1.2] ) diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py index c6ba4a5d..e65122e0 100644 --- a/tests/test_mla_attention_layer.py +++ b/tests/test_mla_attention_layer.py @@ -7,7 +7,6 @@ import pytest import torch -import sparsevllm.layers.mla_attention as mla_attention_module from sparsevllm.engine.cache_manager import ( AttentionKeyComputeView, AttentionViewMeta, @@ -19,7 +18,6 @@ from sparsevllm.engine.cache_manager.base import CacheManager from sparsevllm.layers.mla_attention import ( MLAAttention, - estimate_mla_prefill_workspace_bytes, ) from sparsevllm.operators.mla_attention import ( MlaAttentionOpSpec, @@ -43,30 +41,6 @@ def __init__( self.max_batch_size = int(max_batch_size) -class _ExplicitPrefillProvider(_TestProvider): - supports_explicit_prefill = True - - def run_explicit_prefill( - self, - q, - view, - output, - *, - cu_seqlens_q, - max_seqlen_q, - validation_scope=None, - ): - self.prefill_call = { - "q": q, - "view": view, - "cu_seqlens_q": cu_seqlens_q, - "max_seqlen_q": max_seqlen_q, - "validation_scope": validation_scope, - } - output.copy_(q) - return output - - def _spec(tp_size: int = 4) -> MlaAttentionOpSpec: return MlaAttentionOpSpec( num_q_heads=20, @@ -143,7 +117,9 @@ def test_mla_restores_original_materializer_after_history_profiling() -> None: # Registration mocks cannot catch the duplicate callback rejected when # startup returns to its original manager after a synthetic-history probe. class Manager: - register_attention_key_materializer = CacheManager.register_attention_key_materializer + register_attention_key_materializer = ( + CacheManager.register_attention_key_materializer + ) def kv_layer_index(self, layer_idx): return layer_idx @@ -198,92 +174,19 @@ def test_mla_runtime_release_drops_cached_history_mapping() -> None: context_lens=torch.tensor([2], dtype=torch.int32), ) mapping = weakref.ref(meta.active_slots) - attention._get_prefill_plan(meta, cache_slot_count=2, query_tokens=1) + view = PrefillComputeView( + meta, MlaLatentPayload(torch.empty(2, 1, 512), torch.empty(2, 1, 64)) + ) + attention.chunked_prefill.prepare( + view, torch.tensor([0, 1], dtype=torch.int32), object() + ) + del view del meta assert mapping() is not None attention.release_cache_runtime_bindings(object()) assert mapping() is None -def _expand_history(attention: MLAAttention, history): - heads = attention.spec.local_q_heads - latent = history.gathered_latent - rope = history.gathered_rope[:, None, :].expand(-1, heads, -1) - k_nope = latent[:, None, :192].expand(-1, heads, -1) - expanded_k = torch.cat((k_nope, rope), dim=-1) - expanded_v = latent[:, None, 192:448].expand(-1, heads, -1).contiguous() - return attention.bind_prefill_kv( - history, - expanded_k=expanded_k, - expanded_v=expanded_v, - ) - - -def _torch_prefill( - q: torch.Tensor, - workset, - chunk_lens: torch.Tensor, -) -> torch.Tensor: - history = workset.history - outputs = [] - query_start = 0 - for batch_index, chunk_len in enumerate(chunk_lens.tolist()): - context_len = int(history.context_lens[batch_index].item()) - prefix_len = context_len - int(chunk_len) - history_start = int(history.packed_offsets[batch_index].item()) - keys = workset.expanded_k[ - history_start : history_start + context_len - ].float() - values = workset.expanded_v[ - history_start : history_start + context_len - ].float() - queries = q[query_start : query_start + chunk_len].float() - for query_offset, query in enumerate(queries): - visible = prefix_len + query_offset + 1 - logits = torch.einsum("hd,thd->ht", query, keys[:visible]) - probabilities = torch.softmax(logits * (256**-0.5), dim=-1) - outputs.append( - torch.einsum( - "ht,thd->hd", - probabilities, - values[:visible], - ).to(torch.bfloat16) - ) - query_start += int(chunk_len) - return torch.stack(outputs) - - -def test_mla_prefill_workspace_estimate_accounts_for_full_history() -> None: - actual = estimate_mla_prefill_workspace_bytes( - total_visible_tokens=11, - query_tokens=5, - batch_size=2, - max_context_len=7, - local_q_heads=5, - kv_lora_rank=512, - rope_dim=64, - qk_head_dim=256, - value_head_dim=256, - hidden_size=64, - projection_chunk_size=4, - activation_dtype=torch.bfloat16, - cache_dtype=torch.bfloat16, - ) - - gathered = 11 * (512 + 64) * 2 - projected = 11 * 5 * (192 + 256) * 2 - projection_scratch = 4 * 5 * (192 + 256) * 2 - expanded_k = 11 * 5 * 256 * 2 - attention_output = 5 * 5 * 256 * 2 - output_projection_scratch = 4 * 64 * 2 - metadata = (2 * 7 + 2 * 2 + 7) * 4 - assert actual == max( - gathered + projected + projection_scratch, - gathered + projected + expanded_k + attention_output, - attention_output + output_projection_scratch, - ) + metadata - - def test_mla_prefill_rejects_wrong_payload_before_gather() -> None: attention = _attention() view = PrefillComputeView( @@ -298,30 +201,8 @@ def test_mla_prefill_rejects_wrong_payload_before_gather() -> None: ), ) - with ( - patch("sparsevllm.layers.mla_attention.gather_latent_history") as gather, - pytest.raises(TypeError, match="MlaLatentPayload"), - ): - attention.prepare_prefill_history(view, query_tokens=1) - gather.assert_not_called() - - -def test_mla_prefill_budget_fails_before_allocation_or_gather() -> None: - attention = _attention(budget=1) - view = _view( - torch.empty(2, 1, 512, dtype=torch.bfloat16), - torch.empty(2, 1, 64, dtype=torch.bfloat16), - torch.tensor([[0, 1]], dtype=torch.int32), - torch.tensor([0], dtype=torch.int32), - torch.tensor([2], dtype=torch.int32), - ) - - with ( - patch("sparsevllm.layers.mla_attention.gather_latent_history") as gather, - pytest.raises(MemoryError, match="exceeds its configured budget"), - ): - attention.prepare_prefill_history(view, query_tokens=2) - gather.assert_not_called() + with pytest.raises(TypeError, match="MlaLatentPayload"): + attention._require_mla_payload(view, operation="chunked prefill") def test_mla_attention_bind_resolves_provider_once() -> None: @@ -413,168 +294,6 @@ def run( assert provider.valid_batch_size == 3 -def test_mla_prefill_reuses_validated_packing_across_layers() -> None: - attention = _attention() - active_slots = torch.tensor([[3, 1]], dtype=torch.int32) - request_indices = torch.tensor([0], dtype=torch.int32) - context_lens = torch.tensor([2], dtype=torch.int32) - first_view = _view( - torch.empty(4, 1, 512, dtype=torch.bfloat16), - torch.empty(4, 1, 64, dtype=torch.bfloat16), - active_slots, - request_indices, - context_lens, - ) - second_view = _view( - torch.empty(4, 1, 512, dtype=torch.bfloat16), - torch.empty(4, 1, 64, dtype=torch.bfloat16), - active_slots, - request_indices, - context_lens, - ) - - with ( - patch( - "sparsevllm.layers.mla_attention.validate_gather_metadata" - ) as validate, - patch( - "sparsevllm.layers.mla_attention.gather_latent_history" - ) as gather, - ): - first = attention.prepare_prefill_history(first_view, query_tokens=2) - second = attention.prepare_prefill_history(second_view, query_tokens=2) - reset_context() - attention.prepare_prefill_history(second_view, query_tokens=2) - - assert validate.call_count == 2 - assert gather.call_count == 3 - assert first.packed_offsets is second.packed_offsets - assert first.packed_cu_seqlens is second.packed_cu_seqlens - assert first.packed_slots is second.packed_slots - - -def test_mla_prefill_reuses_query_validation_across_layers() -> None: - attention = _attention() - view = _view( - torch.empty(2, 1, 512, dtype=torch.bfloat16), - torch.empty(2, 1, 64, dtype=torch.bfloat16), - torch.tensor([[0, 1]], dtype=torch.int32), - torch.tensor([0], dtype=torch.int32), - torch.tensor([2], dtype=torch.int32), - ) - with patch("sparsevllm.layers.mla_attention.gather_latent_history"): - history = attention.prepare_prefill_history(view, query_tokens=2) - workset = _expand_history(attention, history) - q = torch.empty(2, 5, 256, dtype=torch.bfloat16) - starts = torch.tensor([0], dtype=torch.int32) - chunks = torch.tensor([2], dtype=torch.int32) - - with ( - patch( - "sparsevllm.layers.mla_attention._host_int_values", - wraps=mla_attention_module._host_int_values, - ) as host_values, - patch.object( - attention.prefill_backend, - "run_prefill", - return_value=torch.empty_like(q), - ), - ): - attention.run_prefill( - q, - workset, - b_start_loc=starts, - chunk_lens=chunks, - ) - attention.run_prefill( - q, - workset, - b_start_loc=starts, - chunk_lens=chunks, - ) - assert host_values.call_count == 2 - - reset_context() - attention.run_prefill( - q, - workset, - b_start_loc=starts, - chunk_lens=chunks, - ) - assert host_values.call_count == 4 - - -def test_mla_prefill_exposes_expanded_explicit_score_view() -> None: - attention = _attention() - view = _view( - torch.empty(2, 1, 512, dtype=torch.bfloat16), - torch.empty(2, 1, 64, dtype=torch.bfloat16), - torch.tensor([[0, 1]], dtype=torch.int32), - torch.tensor([0], dtype=torch.int32), - torch.tensor([2], dtype=torch.int32), - ) - with patch("sparsevllm.layers.mla_attention.gather_latent_history"): - history = attention.prepare_prefill_history(view, query_tokens=2) - workset = _expand_history(attention, history) - - score_view = attention.build_prefill_explicit_view(workset) - - assert isinstance(score_view.payload, ExplicitKVPayload) - assert score_view.payload.k_cache is workset.expanded_k - assert score_view.payload.v_cache is workset.expanded_v - assert score_view.payload.metadata == { - "layout": "mla_packed_varlen", - "cu_seqlens_k": history.packed_cu_seqlens, - } - torch.testing.assert_close( - history.packed_cu_seqlens, - torch.tensor([0, 2], dtype=torch.int32), - ) - assert score_view.meta.active_slots is history.packed_slots - assert score_view.meta.req_indices is history.local_req_indices - assert score_view.meta.context_lens is history.context_lens - assert score_view.meta.max_context_len == history.max_context_len - - -def test_mla_prefill_dispatches_explicit_provider_with_shared_cu_seqlens() -> None: - spec = _spec() - provider = _ExplicitPrefillProvider(spec, device="cpu", max_batch_size=4) - attention = MLAAttention( - spec=spec, - provider=provider, - prefill_workspace_bytes=64 * 1024 * 1024, - hidden_size=64, - projection_chunk_size=8, - ) - view = _view( - torch.empty(2, 1, 512, dtype=torch.bfloat16), - torch.empty(2, 1, 64, dtype=torch.bfloat16), - torch.tensor([[0, 1]], dtype=torch.int32), - torch.tensor([0], dtype=torch.int32), - torch.tensor([2], dtype=torch.int32), - ) - with patch("sparsevllm.layers.mla_attention.gather_latent_history"): - history = attention.prepare_prefill_history(view, query_tokens=2) - workset = _expand_history(attention, history) - q = torch.randn(2, 5, 256, dtype=torch.bfloat16) - cu_seqlens_q = torch.tensor([0, 2], dtype=torch.int32) - set_context(True, cu_seqlens_q=cu_seqlens_q) - try: - output = attention.run_prefill( - q, - workset, - b_start_loc=cu_seqlens_q[:-1], - chunk_lens=cu_seqlens_q[1:] - cu_seqlens_q[:-1], - ) - finally: - reset_context() - - torch.testing.assert_close(output, q) - assert provider.prefill_call["view"].payload.k_cache is workset.expanded_k - assert provider.prefill_call["cu_seqlens_q"] is cu_seqlens_q - assert provider.prefill_call["max_seqlen_q"] == 2 - - def test_mla_materializes_actual_keys_for_permuted_slots() -> None: torch.manual_seed(47) attention = _attention(tp_size=4) @@ -596,11 +315,15 @@ def test_mla_materializes_actual_keys_for_permuted_slots() -> None: ) def project_latent(latent: torch.Tensor) -> torch.Tensor: - return torch.einsum( - "tr,hor->tho", - latent.float(), - weights.float(), - ).to(torch.bfloat16).flatten(1) + return ( + torch.einsum( + "tr,hor->tho", + latent.float(), + weights.float(), + ) + .to(torch.bfloat16) + .flatten(1) + ) actual = attention.materialize_expanded_keys( view, @@ -629,111 +352,41 @@ def project_latent(latent: torch.Tensor) -> torch.Tensor: torch.testing.assert_close(actual, expected) -CUDA_REQUIRED = pytest.mark.skipif( - not torch.cuda.is_available(), - reason="CUDA is required for MLA full-history prefill tests", -) - - -@CUDA_REQUIRED -def test_mla_prefill_matches_ragged_full_history_oracle() -> None: - torch.manual_seed(41) - attention = _attention(device="cuda", tp_size=4) - latent_cache = torch.randn(20, 1, 512, dtype=torch.bfloat16, device="cuda") - rope_cache = torch.randn(20, 1, 64, dtype=torch.bfloat16, device="cuda") - active_slots = torch.full((3, 7), -1, dtype=torch.int32, device="cuda") - active_slots[2, :5] = torch.tensor([13, 2, 17, 5, 11], device="cuda") - active_slots[0, :7] = torch.tensor( - [19, 1, 7, 15, 3, 9, 6], - device="cuda", - ) - context_lens = torch.tensor([5, 7], dtype=torch.int32, device="cuda") +def test_chunked_prefill_budget_fails_before_projection(): + # A deliberately small cap must reject before allocating expanded KV. + attention = _attention(budget=1) view = _view( - latent_cache, - rope_cache, - active_slots, - torch.tensor([2, 0], dtype=torch.int32, device="cuda"), - context_lens, - ) - history = attention.prepare_prefill_history(view, query_tokens=5) - workset = _expand_history(attention, history) - chunk_lens = torch.tensor([2, 3], dtype=torch.int32, device="cuda") - b_start_loc = torch.tensor([0, 2], dtype=torch.int32, device="cuda") - q = torch.randn(5, 5, 256, dtype=torch.bfloat16, device="cuda") - - output = attention.run_prefill( - q, - workset, - b_start_loc=b_start_loc, - chunk_lens=chunk_lens, - ) - torch.cuda.synchronize() - expected = _torch_prefill(q, workset, chunk_lens) - - assert history.visible_tokens == 12 - torch.testing.assert_close( - output.float(), - expected.float(), - rtol=3e-2, - atol=3e-2, - ) - - -@CUDA_REQUIRED -def test_mla_prefill_is_invariant_to_chunk_boundary() -> None: - torch.manual_seed(43) - attention = _attention(device="cuda", tp_size=4) - latent_cache = torch.randn(12, 1, 512, dtype=torch.bfloat16, device="cuda") - rope_cache = torch.randn(12, 1, 64, dtype=torch.bfloat16, device="cuda") - active_slots = torch.tensor( - [[9, 1, 11, 3, 7, 5]], - dtype=torch.int32, - device="cuda", - ) - request_indices = torch.tensor([0], dtype=torch.int32, device="cuda") - q = torch.randn(6, 5, 256, dtype=torch.bfloat16, device="cuda") - - full_history = attention.prepare_prefill_history( - _view( - latent_cache, - rope_cache, - active_slots, - request_indices, - torch.tensor([6], dtype=torch.int32, device="cuda"), - ), - query_tokens=6, - ) - full_workset = _expand_history(attention, full_history) - full_output = attention.run_prefill( - q, - full_workset, - b_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), - chunk_lens=torch.tensor([6], dtype=torch.int32, device="cuda"), - ) - - first_history = attention.prepare_prefill_history( - _view( - latent_cache, - rope_cache, - active_slots, - request_indices, - torch.tensor([4], dtype=torch.int32, device="cuda"), - ), - query_tokens=4, + torch.empty(2, 1, 512, dtype=torch.bfloat16), + torch.empty(2, 1, 64, dtype=torch.bfloat16), + torch.tensor([[0, 1]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([2], dtype=torch.int32), ) - first_output = attention.run_prefill( - q[:4], - _expand_history(attention, first_history), - b_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), - chunk_lens=torch.tensor([4], dtype=torch.int32, device="cuda"), + manager = SimpleNamespace( + register_attention_key_materializer=Mock(), + store_attention_payload=Mock(return_value=None), + on_kv_stored=Mock(), + before_prefill_layer_attention=Mock(), + build_prefill_compute_view=Mock(return_value=view), + prefill_score_request=Mock(return_value=None), ) - second_output = attention.run_prefill( - q[4:], - full_workset, - b_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), - chunk_lens=torch.tensor([2], dtype=torch.int32, device="cuda"), + set_context(True, torch.tensor([0, 1], dtype=torch.int32), manager, seqs=[]) + get_context().sparse_controller = SimpleNamespace( + get_prefill_selection=Mock(return_value=None) ) - torch.cuda.synchronize() - - torch.testing.assert_close(first_output, full_output[:4], rtol=3e-2, atol=3e-2) - torch.testing.assert_close(second_output, full_output[4:], rtol=3e-2, atol=3e-2) + project = Mock() + try: + with pytest.raises(MemoryError, match="workspace exceeds budget"): + attention.run_cached_attention( + torch.empty(1, 5, 256), + torch.empty(1, 5, 192), + torch.empty(1, 5, 64), + torch.empty(1, 512), + torch.empty(1, 64), + project_latent=project, + absorb_query=Mock(), + reconstruct_values=Mock(), + ) + project.assert_not_called() + finally: + reset_context() diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index a7e89cd1..9833c935 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -11,7 +11,6 @@ DecodeComputeView, ExplicitKVPayload, MlaLatentPayload, - PrefillComputeView, ) from sparsevllm.kernels.external.sgl.fa3 import sgl_fa3_device_support from sparsevllm.kernels.triton.mla import ( @@ -391,7 +390,8 @@ def test_mla_provider_rejects_explicit_kv_before_kernel() -> None: kernel.assert_not_called() -def test_sgl_provider_uses_packed_varlen_prefill_metadata() -> None: +@pytest.mark.parametrize("causal", [False, True]) +def test_sgl_provider_returns_chunk_output_and_lse(causal) -> None: spec = _spec(tp_size=4) workspace = _cpu_workspace(batch_size=1, head_count=5) fa3 = Mock() @@ -411,45 +411,31 @@ def test_sgl_provider_uses_packed_varlen_prefill_metadata() -> None: max_batch_size=1, ) q = torch.empty(2, 5, 256, dtype=torch.bfloat16) - output = torch.empty_like(q) + k = torch.empty(4, 5, 256, dtype=torch.bfloat16) + v = torch.empty_like(k) cu_seqlens_q = torch.tensor([0, 2], dtype=torch.int32) cu_seqlens_k = torch.tensor([0, 4], dtype=torch.int32) - view = PrefillComputeView( - meta=AttentionViewMeta( - active_slots=torch.arange(4, dtype=torch.int32).view(1, 4), - req_indices=torch.tensor([0], dtype=torch.int32), - context_lens=torch.tensor([4], dtype=torch.int32), - max_context_len=4, - ), - payload=ExplicitKVPayload( - k_cache=torch.empty(4, 5, 256, dtype=torch.bfloat16), - v_cache=torch.empty(4, 5, 256, dtype=torch.bfloat16), - metadata={ - "layout": "mla_packed_varlen", - "cu_seqlens_k": cu_seqlens_k, - }, - ), - ) - fa3.run_contiguous_explicit_varlen.return_value = output + lse = torch.empty(5, 2, dtype=torch.float32) + fa3.run_contiguous_explicit_varlen.side_effect = lambda q, k, v, out, **kw: (out, lse) - actual = provider.run_explicit_prefill( - q, - view, - output, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=2, + output, actual_lse = provider.run_prefill_chunk( + q, k, v, cu_seqlens_q, cu_seqlens_k, 2, 4, causal=causal ) - assert actual is output + assert output.shape == q.shape + assert output.dtype == q.dtype + assert actual_lse is lse fa3.run_contiguous_explicit_varlen.assert_called_once_with( q, - view.payload.k_cache, - view.payload.v_cache, + k, + v, output, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=2, max_seqlen_k=4, + causal=causal, + return_softmax_lse=True, ) fa3.run_explicit_varlen.assert_not_called() diff --git a/tests/test_mla_chunked_prefill.py b/tests/test_mla_chunked_prefill.py new file mode 100644 index 00000000..14e82fff --- /dev/null +++ b/tests/test_mla_chunked_prefill.py @@ -0,0 +1,280 @@ +from types import SimpleNamespace + +import pytest +import torch + +from sparsevllm.engine.cache_manager.base import ( + AttentionViewMeta, + MlaLatentPayload, + PrefillComputeView, + PrefillScoreRequest, +) +from sparsevllm.kernels.triton.mla.prefill import attention_partial +from sparsevllm.operators.mla_attention import MlaAttentionOpSpec, MlaSglFa3Provider +from sparsevllm.operators.mla_prefill import ChunkedMlaPrefill + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("causal", [True, False]) +@pytest.mark.parametrize( + "queries,keys", + [((1025, 7), (1025, 7)), ((65, 3), (193, 131)), ((1025, 7), (17, 0))], +) +def test_triton_partial_ragged_causal_bounds_and_lse(causal, queries, keys): + # Loop clipping must preserve bottom-right causal alignment, partial tiles, + # and fully masked rows. Existing chunked tests use only short query tiles + # and cannot catch a wrong bound after increasing the attention tile size. + torch.manual_seed(51) + device = "cuda" + q = ( + torch.randn(5, sum(queries), 256, device=device, dtype=torch.bfloat16) * 0.2 + ).transpose(0, 1) + k = torch.randn(sum(keys), 5, 256, device=device, dtype=torch.bfloat16) * 0.2 + # V is a strided view of the joint KV projection in the serving path. + v = torch.randn(sum(keys), 5, 448, device=device, dtype=torch.bfloat16)[..., 192:] + cu_q = torch.tensor([0, queries[0], sum(queries)], device=device, dtype=torch.int32) + cu_k = torch.tensor([0, keys[0], sum(keys)], device=device, dtype=torch.int32) + output, lse = attention_partial( + q, k, v, cu_q, cu_k, max(queries), max(keys), scale=0.0625, causal=causal + ) + qa = ka = 0 + for qn, kn in zip(queries, keys): + logits = torch.einsum( + "qhd,khd->hqk", q[qa : qa + qn].float(), k[ka : ka + kn].float() + ) * 0.0625 + if causal: + mask = torch.arange(kn, device=device)[None] <= ( + torch.arange(qn, device=device)[:, None] + kn - qn + ) + logits.masked_fill_(~mask[None], -torch.inf) + expected_lse = logits.logsumexp(-1) + probabilities = logits.softmax(-1).nan_to_num(0) + expected = torch.einsum("hqk,khd->qhd", probabilities, v[ka : ka + kn].float()) + torch.testing.assert_close( + output[qa : qa + qn].float(), expected, atol=0.006, rtol=0.03 + ) + torch.testing.assert_close( + lse[:, qa : qa + qn], expected_lse, atol=0.003, rtol=0.001 + ) + qa += qn + ka += kn + + +def make_case(contexts=(73, 29, 41), queries=(17, 29, 9), heads=5): + torch.manual_seed(129) + device, dtype = "cuda", torch.bfloat16 + spec = MlaAttentionOpSpec( + num_q_heads=heads, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + activation_dtype=dtype, + cache_dtype=dtype, + tp_size=1, + cuda_graph=False, + ) + capacity = sum(contexts) + latent = torch.randn(capacity, 1, 512, device=device, dtype=dtype) + rope = torch.randn(capacity, 1, 64, device=device, dtype=dtype) * 0.2 + weight = torch.randn(heads, 448, 512, device=device, dtype=dtype) * 0.03 + q = torch.randn(sum(queries), heads, 256, device=device, dtype=dtype) * 0.3 + slots = torch.full( + (len(contexts), max(contexts)), -1, device=device, dtype=torch.int32 + ) + rows = tuple(reversed(range(len(contexts)))) + permutation = torch.randperm(capacity, device=device).int() + start = 0 + for row, n in zip(rows, contexts): + slots[row, :n] = permutation[start : start + n] + start += n + view = PrefillComputeView( + AttentionViewMeta( + slots, + torch.tensor(rows, device=device, dtype=torch.int32), + torch.tensor(contexts, device=device, dtype=torch.int32), + ), + MlaLatentPayload(latent, rope), + ) + cu = torch.tensor( + [0, *torch.tensor(queries).cumsum(0).tolist()], device=device, dtype=torch.int32 + ) + + def project(x): + return torch.nn.functional.linear(x, weight.flatten(0, 1)) + + def absorb(x): + return torch.bmm(x.transpose(0, 1), weight[:, :192]).transpose(0, 1) + + return spec, q, view, cu, project, absorb + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize( + "mode,full_normalizer,backend", + [ + ("logits", False, "fa3"), + ("probability", False, "fa3"), + ("probability", True, "triton"), + ("probability", True, "fa3"), + ], +) +def test_chunked_attention_and_scores_match_explicit_oracle( + mode, full_normalizer, backend +): + # Independent full softmax protects masks, global normalization, physical + # slot indirection, and score reductions across uneven history/query blocks. + spec, q, view, cu, project, absorb = make_case() + provider = ( + MlaSglFa3Provider(op_spec=spec, device="cuda:0", max_batch_size=3) + if backend == "fa3" + else SimpleNamespace() + ) + runner = ChunkedMlaPrefill(spec, provider, 19) + contexts, starts = view.meta.context_lens.tolist(), cu.tolist() + ranges = tuple( + (n - min(7, b - a), n) for n, a, b in zip(contexts, starts, starts[1:]) + ) + request = PrefillScoreRequest( + ranges, mode, 0 if full_normalizer else 3, 0 if full_normalizer else 4 + ) + actual, lse, scores = runner.run(q, view, cu, object(), project, absorb, request) + for i, n in enumerate(contexts): + a, b = starts[i : i + 2] + row = int(view.meta.req_indices[i]) + indices = view.meta.active_slots[row, :n].long() + latent = view.payload.latent_cache[indices, 0] + rope = view.payload.rope_cache[indices, 0] + expanded = project(latent).view(n, spec.local_q_heads, 448) + k = torch.cat( + (expanded[..., :192], rope[:, None].expand(-1, spec.local_q_heads, -1)), -1 + ) + v = expanded[..., 192:] + raw = torch.einsum("qhd,khd->hqk", q[a:b].float(), k.float()) + qi = torch.arange(n - (b - a), n, device=q.device) + ki = torch.arange(n, device=q.device) + mask = qi[:, None] >= ki[None, :] + z = (raw * spec.softmax_scale).masked_fill(~mask[None], -torch.inf) + expected = torch.einsum("hqk,khd->qhd", z.softmax(-1), v.float()) + torch.testing.assert_close(actual[a:b].float(), expected, atol=0.006, rtol=0.03) + torch.testing.assert_close(lse[:, a:b], z.logsumexp(-1), atol=0.003, rtol=0.001) + observed = raw[:, -(ranges[i][1] - ranges[i][0]) :] + valid = ( + mask[-observed.shape[1] :] + & (ki[None] >= request.candidate_start) + & (ki[None] < n - request.recent_keep_tokens) + ) + if mode == "logits": + ref = observed.masked_fill(~valid[None], -torch.inf).amax((0, 1)) + tolerance = 0.001 + else: + ref = ( + (observed * spec.softmax_scale) + .masked_fill(~valid[None], -torch.inf) + .softmax(-1) + .mean(1) + .amax(0) + ) + tolerance = 0.0002 + torch.testing.assert_close(scores[i, :n], ref, atol=tolerance, rtol=0.015) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("mode", [None, "logits", "probability"]) +def test_chunk_size_preserves_outputs_and_bounds_history_projection(mode): + spec, q, view, cu, project, absorb = make_case(contexts=(171, 53), queries=(35, 5)) + provider = MlaSglFa3Provider(op_spec=spec, device="cuda:0", max_batch_size=2) + # Scoring must not re-project historical KV after the attention pass. + request = PrefillScoreRequest(((166, 171), (48, 53)), mode, 3, 4) if mode else None + outputs = [] + for size in (17, 64): + projected = [] + + def tracked_project(x, projected=projected): + projected.append(x.shape[0]) + return project(x) + + runner = ChunkedMlaPrefill(spec, provider, size) + outputs.append( + runner.run(q, view, cu, object(), tracked_project, absorb, request)[0] + ) + assert projected[0] == q.shape[0] + assert max(projected[1:]) <= size + assert sum(projected) == sum(view.meta.context_lens.tolist()) + torch.testing.assert_close(outputs[0], outputs[1], atol=0.006, rtol=0.03) + + +def test_history_budget_is_bounded_and_plan_released(): + # Long contexts previously grew the full-history workspace and retained the + # temporary startup cache's mapping after runtime retirement. + import weakref + + from sparsevllm.operators.mla_prefill import estimate_mla_prefill_workspace_bytes + + spec = MlaAttentionOpSpec( + num_q_heads=20, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + activation_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + tp_size=1, + cuda_graph=False, + ) + runner = ChunkedMlaPrefill(spec, SimpleNamespace(), 256) + estimates = [] + for length in (1024, 8192): + view = PrefillComputeView( + AttentionViewMeta( + torch.arange(length, dtype=torch.int32)[None], + torch.tensor([0], dtype=torch.int32), + torch.tensor([length], dtype=torch.int32), + ), + MlaLatentPayload(torch.empty(0, 1, 512), torch.empty(0, 1, 64)), + ) + cu = torch.tensor([0, 128], dtype=torch.int32) + scope = object() + plan = runner.prepare(view, cu, scope) + assert runner.prepare(view, cu, scope) is plan + assert all(n <= runner.chunk_size for _, _, n, _ in plan.history_chunks) + estimates.append( + estimate_mla_prefill_workspace_bytes( + plan=plan, + spec=spec, + chunk_size=256, + hidden_size=2048, + projection_chunk_size=128, + ) + ) + # Only tiny packing metadata grows with history; expanded KV does not. + assert estimates[1] - estimates[0] < 4096 + mapping = weakref.ref(view.meta.active_slots) + del view, plan + runner.clear() + assert mapping() is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_multi_tile_observations_and_empty_candidates(): + # H2O can observe more than one 32-query tile. Entire candidate blocks may + # be causally masked; they must contribute zero probability, never NaN. + spec, q, view, cu, project, absorb = make_case((239,), (137,), heads=20) + provider = MlaSglFa3Provider(op_spec=spec, device="cuda:0", max_batch_size=1) + runner = ChunkedMlaPrefill(spec, provider, 53) + req = PrefillScoreRequest(((102, 239),), "probability", 130, 3) + _, _, scores = runner.run(q, view, cu, object(), project, absorb, req) + slots = view.meta.active_slots[0, :239].long() + latent, rope = runner.gather(view.payload, slots) + expanded = project(latent).view(239, 20, 448) + keys = torch.cat((expanded[..., :192], rope[:, None].expand(-1, 20, -1)), -1) + z = torch.einsum("qhd,khd->hqk", q.float(), keys.float()) * spec.softmax_scale + ki = torch.arange(239, device=q.device) + valid = ( + (torch.arange(102, 239, device=q.device)[:, None] >= ki) + & (ki >= 130) + & (ki < 236) + ) + p = z.masked_fill(~valid, -torch.inf).softmax(-1).nan_to_num(0) + ref = p.mean(1).amax(0) + torch.testing.assert_close(scores[0], ref, atol=2e-4, rtol=0.015)