diff --git a/.agents/skills/review-operator-organization/SKILL.md b/.agents/skills/review-operator-organization/SKILL.md index a06f6d91..267cd9b7 100644 --- a/.agents/skills/review-operator-organization/SKILL.md +++ b/.agents/skills/review-operator-organization/SKILL.md @@ -1,6 +1,6 @@ --- name: review-operator-organization -description: Review Sparse-vLLM operator architecture, provider selection, platform capability boundaries, kernel ownership, dependency compatibility, weight layouts, fallback semantics, and validation. Use for diffs touching src/sparsevllm/operators, src/sparsevllm/platforms, Triton or external kernels, model-to-operator call sites, quantized weight loading, CUDA Graph constraints, optional kernel dependencies, or backend removal and migration. +description: Review Sparse-vLLM operator architecture, provider selection, platform capability boundaries, kernel ownership, dependency compatibility, weight layouts, batch-only CUDA Graph adaptation, fallback semantics, and validation. Use for diffs touching src/sparsevllm/operators, src/sparsevllm/platforms, Triton or external kernels, model-to-operator call sites, quantized weight loading, CUDA Graph constraints, optional kernel dependencies, or backend removal and migration. --- # Review Operator Organization @@ -98,6 +98,18 @@ fallback path from model construction through execution. - Ensure one provider's build or JIT failure does not disable unrelated operators. +### Decode CUDA Graph + +Batch-only is the only maintained decode CUDA Graph shape policy. For any +review touching captured decode, graph input preparation, provider graph state, +or sparse topology paths, read and enforce +[references/batch-only-decode-graph.md](references/batch-only-decode-graph.md). +That reference defines graph identity, static versus dynamic metadata, unified +input ownership, participant lifecycles, external wrappers, validation, and +finding severity. Eager may remain as a separate correctness path or for +operators that do not support graph capture; do not preserve a second bucketed +graph architecture. + ### Kernel Portfolio - Treat standard operations as upstream-first. Prefer a mature upstream public diff --git a/.agents/skills/review-operator-organization/references/batch-only-decode-graph.md b/.agents/skills/review-operator-organization/references/batch-only-decode-graph.md new file mode 100644 index 00000000..a54bc6f3 --- /dev/null +++ b/.agents/skills/review-operator-organization/references/batch-only-decode-graph.md @@ -0,0 +1,207 @@ +# Batch-Only Decode CUDA Graph Review + +Read this reference when a review touches an operator reachable from captured +decode, decode graph input preparation, provider graph state, or sparse +short/long topology paths. + +## Scope and Vocabulary + +Batch-only is the only maintained decode CUDA Graph shape policy. Do not add or +preserve bucketed-only graph implementations, context-bucket routing, parallel +provider families, or configuration surfaces merely to keep a second graph +architecture alive. Eager may remain as an independent correctness or +unsupported-graph path; it must not leak context-dependent dispatch into +captured decode. + +Use these terms consistently: + +- **batch-only graph**: graph identity depends on batch capacity but not actual + per-step context lengths; +- **strict batch-only**: one forward graph per batch and sampling topology; +- **path-scoped batch-only**: one forward graph per batch, sampling topology, + and finite semantic topology path when kernel chains genuinely differ; +- **context capacity**: a capture-time storage and launch upper bound, not a + replay-time graph bucket; +- **static launch plan**: capture-time tile, warp, stage, split envelope, grid + envelope, compiled variant, and workspace capacity; +- **replay-before metadata**: dynamic state prepared outside the captured graph + before replay, also called graph-out preparation; +- **graph-in preparation**: fixed device work captured before operator forward; +- **stable graph state**: typed inputs, provider state, workspaces, wrappers, + and keepalive owners whose addresses and capacities remain fixed. + +Reading `context_lens` or exposing `plan()` does not by itself violate +batch-only. The violation is allowing actual context to change graph identity, +captured topology, static launch plan, workspace shape, tensor addresses, or +provider binding. + +## Operator and Provider Adaptation + +- Define graph identity from batch capacity, finite semantic topology path, + sampling topology, and capture-time tensor/layout contract. Actual + `context_lens` must not enter graph keys or cause runtime capture. +- Resolve model/hardware tuning tables and compile-time choices before capture. + A table selected for a fixed model architecture and hardware combination is + valid static configuration. Tile, warp, stage, compiled variant, grid + envelope, and workspace shape are not replay metadata. +- Flag replay-time host thresholds that switch kernel chains, launch variants, + split envelopes, or workspaces. Replace them with a fixed envelope plus + device-side effective scheduling, bind another batch-only provider, or reject + the unsupported contract during resolution/preparation. +- Dynamic lengths may drive device masking, effective split/range metadata, or + an explicit replay-before provider plan when those updates write only stable + graph state and leave the captured launch contract unchanged. +- Permit separate startup-captured short/long paths only when the semantic + kernel chain truly differs. Merge methods or length regimes with identical + topology. Seal the startup plan; transitions among declared paths must not + JIT, reselect a provider or variant, grow workspace, or recapture. +- Require `supports(spec, caps)` and preparation to validate dtype, shape, + layout, capacity, padding, workspace, and batch-only compatibility before + forward. Do not treat a few fixed-shape experiments as production support. +- When a standard upstream provider already exposes a graph-stable lifecycle, + adapt that lifecycle instead of cloning its kernel. Use a repository-owned + fixed-grid provider for missing Sparse-vLLM semantics, portable fallback, or + an exact measured override—not as an automatic replacement for upstream. +- Fail unsupported capacity or layout before cache mutation. Once bound, do not + switch provider, allocate a larger workspace, or fall back after execution + begins. + +## Unified Inputs and Participant Lifecycle + +The unified input mechanism standardizes public replay inputs and update order; +it does not combine every tensor into one allocation or expose provider and +sparse-algorithm internals to the graph runner. + +### Common input contract + +Keep shared replay inputs in typed, fixed-address runner-owned state. At minimum +distinguish token ids, positions, context lengths, request indices, KV +write-slot mappings, and valid-row state. Every registered slot declares: + +- shape, dtype, and device; +- batch axis and capacity; +- padding policy; +- semantic/value source and copy policy; +- stable-address requirement. + +Prefer explicit `DecodeGraphInputs`-style fields. Flag an indefinitely growing +`dict[str, Tensor]`, an untyped memory blob, or a positional runner API carrying +method- and provider-private tensors. + +### Ownership + +For every field distinguish storage owner, semantic owner, and per-step value +producer: + +- graph runner: common decode input storage, padding, capture/replay, and graph + identity; +- cache manager: physical KV storage, page/slot metadata, and physical cache + views; +- `SparseController`: logical sparse selection, cross-layer observation, and + attention-score coordination; +- provider: static kernel plan, schedule buffers, private graph state, + workspace, external wrapper, and physical weight/layout; +- model/attention layer: stable operator semantics only. + +Do not move provider workspace into the common registry or physical cache +metadata into `SparseController`. The runner coordinates lifecycle and copy +order without taking ownership of private algorithms or layouts. + +### Participant lifecycle + +Use a typed lifecycle equivalent to: + +```text +init_graph_state(contract, topology_path) +prepare_out_graph(step, state) +prepare_in_graph(state) +graph_keepalive_tensors(state) +``` + +- initialization allocates stable private buffers/workspaces, resolves the + static plan, initializes wrappers/JIT once, and records capacity; +- graph-out preparation updates dynamic host metadata or executes a documented + provider plan, writing only stable state; +- graph-in preparation contains fixed device work captured before forward; +- keepalive ownership prevents captured tensors, workspaces, wrappers, or + outputs from being released or replaced. + +Coordinate provider preparation once before each model replay, outside +per-layer attention forward. Model and attention code consume prepared state and +must not contain sparse-method branches, provider names, external-wrapper +access, or graph lifecycle calls. + +### Padding + +Pad real batches to their capture bucket with an explicit active-row contract. +Padding rows use safe token, position, slot, page, and score metadata. Prove +they cannot access or mutate a live request's KV cache, sparse score, or +controller state. Do not rely on an incidental sentinel that a kernel still +dereferences before masking. + +## External Graph Wrappers + +For FlashInfer paged decode and comparable external providers with a public +CUDA Graph lifecycle: + +- Use the upstream graph-enabled wrapper instead of an ordinary eager wrapper, + raw internal kernel, or repository reimplementation of its planner. Bind one + wrapper to each captured batch/topology state that needs distinct storage. +- Provider state owns fixed-capacity page indptr, page indices, last-page + lengths, integer/float workspaces, output owners, and the wrapper. The runner + invokes the participant lifecycle but never reads wrapper-private fields or + constructs provider-specific page metadata. +- Run context-dependent `plan()` or the documented fast-plan path during + replay-before preparation. Planning may change contents, not wrapper/workspace + identity, input/output addresses, launch contract, or captured `run()` + topology. +- Captured forward calls only the already-bound wrapper `run()`. Flag planning + in forward, wrapper recreation, real-length-driven `masked_select`, `cat`, or + allocation, workspace replacement, and runtime backend switching. +- Reuse persistent host/GPU staging. If the public API requires D2H or host + planning, keep the synchronization boundary explicit and report its p50/p95 + cost separately; it must not alter captured addresses. +- If the minimum supported upstream version has no public wrapper contract that + satisfies these invariants, reject the provider for batch-only during binding. + Do not reach through private APIs or silently fall back after replay starts. +- Validate constructor, plan/fast-plan, and run with a real installation at the + declared minimum version. Mocks do not prove lifecycle compatibility. + +## Required Review Evidence + +For every claimed model/method/provider topology path require: + +- one startup-captured graph per batch/topology/sampling state and no actual + context bucket in graph keys; +- repeated replay across representative, historical-threshold boundary, + ragged, padded, and maximum-capacity contexts; +- unchanged graph count, `recapture_count == 0`, and stable registered input, + workspace, output, and wrapper-owner addresses; +- no replay-time JIT, static plan/variant reselection, context-sized allocation, + workspace growth, provider switch, or per-layer host planning; +- independent numerical comparison for output and every required score, LSE, + cache mutation, or other side effect; +- padding and maximum-capacity memory-safety tests; +- real-model fixed and churn coverage, selected-provider/binding evidence, and + matched performance results; +- isolated timing for CPU metadata preparation, H2D/D2H, provider planning, + waits, and graph replay when replay-before work is nontrivial. + +For semantic short/long paths, test the algorithm threshold below, at, and +above the boundary and transition between the already captured paths without +state loss or new capture. Do not preserve a historical kernel-tuning threshold +as a semantic topology path. + +## Finding Severity Additions + +- P1: a claimed batch-only path performs runtime recapture, changes captured + topology/variant from actual context, replaces captured addresses, switches + provider after binding, or lets padding access/mutate live request state. +- P2: a changed provider rejects batch-only cleanly but leaves a required + model/method without a production batch-only provider; replay preparation has + avoidable per-layer or allocation overhead; or new code extends only the + retired bucketed graph path without an explicit migration purpose. +- P3: terminology, binding-report, or ownership documentation is unclear while + behavior remains correct and observable. + +Use the main skill's P0-P3 definitions for all other findings. diff --git a/README.md b/README.md index ae380f9d..bff21b3d 100644 --- a/README.md +++ b/README.md @@ -137,8 +137,10 @@ uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl Use `cu129` instead of `cu130` for CUDA 12.9. -`einops`, `sglang-kernel`, and the training, benchmark, and test packages are all -part of the main installation; no workflow-specific extras are required. +`einops`, `sglang-kernel==0.4.5`, and the training, benchmark, and test packages +are all part of the main installation; no workflow-specific extras are required. +The SGL kernel package is pinned because its compiled operators must match the +validated PyTorch/CUDA ABI; other versions are rejected during provider setup. Sparse-vLLM supports Qwen3.5/Qwen3.6/Qwen3.8 checkpoints in unquantized BF16 and block-scaled FP8 formats. These releases share the `qwen3_5` runtime diff --git a/README_zh.md b/README_zh.md index 510e038b..67c619c0 100644 --- a/README_zh.md +++ b/README_zh.md @@ -119,6 +119,9 @@ uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl CUDA 12.9 环境将 `cu130` 换成 `cu129`。 +主依赖固定使用 `sglang-kernel==0.4.5`，因为其编译算子必须匹配已经验证的 +PyTorch/CUDA ABI；其他版本会在 Provider 准备阶段明确失败，不会静默 fallback。 + Sparse-vLLM 支持未量化 BF16 和 block-scaled FP8 格式的 Qwen3.5/Qwen3.6/Qwen3.8 checkpoint。三者共享 `qwen3_5` 运行时架构，以及 相同的精度、并行方式、稀疏方法和多模态支持。其 prefill causal Conv1D 和 diff --git a/benchmark/efficiency/bench_probe.py b/benchmark/efficiency/bench_probe.py index a5deae2b..5ff58f43 100644 --- a/benchmark/efficiency/bench_probe.py +++ b/benchmark/efficiency/bench_probe.py @@ -163,6 +163,26 @@ def _percentile(values: list[float], quantile: float) -> float: return ordered[lower] * (1.0 - weight) + ordered[upper] * weight +_DECODE_GRAPH_COUNTERS = ( + "capture_count", + "replay_count", + "eager_static_count", + "force_eager_count", + "eviction_count", + "recapture_count", +) + + +def _decode_graph_counter_delta( + before: dict[str, Any], + after: dict[str, Any], +) -> dict[str, int]: + return { + name: int(after.get(name, 0)) - int(before.get(name, 0)) + for name in _DECODE_GRAPH_COUNTERS + } + + def _monitor_gpu_ids(explicit: str | None) -> list[int]: value = explicit or os.environ.get("CUDA_VISIBLE_DEVICES", "") if not value: @@ -877,7 +897,12 @@ def run_sparsevllm_churn( "[Sparse-vLLM Churn] Initializing " f"method={args.sparse_method}, max_concurrency={concurrency}..." ) + engine_init_started = time.perf_counter() llm = LLM(args.model_path, **engine_kwargs) + engine_init_s = time.perf_counter() - engine_init_started + startup_graph_summary = llm.debug_sparse_state_summaries()[0][ + "decode_graph" + ] try: request_count = concurrency * args.churn_request_multiplier for p_len in args.prompt_lens: @@ -917,6 +942,9 @@ def run_sparsevllm_churn( try: for iteration in range(args.num_iters): profiler.reset() + graph_before = llm.debug_sparse_state_summaries()[0][ + "decode_graph" + ] trace = _trace_for_iteration( args, model_specs, @@ -979,6 +1007,9 @@ def run_sparsevllm_churn( finished_times[seq_id] = now generated_counts[seq_id] = len(token_ids) elapsed_s = time.perf_counter() - started + graph_after = llm.debug_sparse_state_summaries()[0][ + "decode_graph" + ] expected_seq_ids = set(seq_to_request) for name, observed in ( @@ -1034,6 +1065,13 @@ def run_sparsevllm_churn( "status": "success", "elapsed_s": elapsed_s, "step_count": step_count, + "engine_init_s": engine_init_s, + "startup_decode_cuda_graph": startup_graph_summary, + "decode_cuda_graph_before": graph_before, + "decode_cuda_graph_after": graph_after, + "decode_cuda_graph_counter_delta": ( + _decode_graph_counter_delta(graph_before, graph_after) + ), "request_throughput_rps": request_count / elapsed_s, "input_token_throughput_tps": total_input / elapsed_s, "output_token_throughput_tps": total_output / elapsed_s, diff --git a/benchmark/efficiency/metrics_calculator.py b/benchmark/efficiency/metrics_calculator.py index 2768f63f..b961fb53 100644 --- a/benchmark/efficiency/metrics_calculator.py +++ b/benchmark/efficiency/metrics_calculator.py @@ -154,22 +154,22 @@ def from_config_dict(cls, cfg: dict[str, Any], bytes_per_param: int = 2) -> Mode "Model config must have positive num_attention_heads, got " f"num_attention_heads={num_attention_heads}." ) - configured_head_dim = cfg.get("head_dim") - if configured_head_dim is None: + explicit_head_dim = cfg.get("head_dim") + if explicit_head_dim is None and cfg.get("qk_nope_head_dim") is not None: + explicit_head_dim = int(cfg["qk_nope_head_dim"]) + int( + cfg.get("qk_rope_head_dim", 0) + ) + if explicit_head_dim is None: if hidden_size % num_attention_heads != 0: raise ValueError( - "Model config without an explicit head_dim requires " - "num_attention_heads to divide hidden_size, got " - f"hidden_size={hidden_size}, " - f"num_attention_heads={num_attention_heads}." - ) - head_dim = hidden_size // num_attention_heads - else: - head_dim = int(configured_head_dim) - if head_dim <= 0: - raise ValueError( - f"Model config head_dim must be positive, got {head_dim}." + "Model config must define head_dim when hidden_size is not " + "divisible by num_attention_heads, got " + f"hidden_size={hidden_size}, num_attention_heads={num_attention_heads}." ) + explicit_head_dim = hidden_size // num_attention_heads + head_dim = int(explicit_head_dim) + if head_dim <= 0: + raise ValueError(f"Model config must have positive head_dim, got {head_dim}.") vocab_size = int(cfg["vocab_size"]) # MoE parameters diff --git a/configs/debug/minimax_m2_tiny_random.json b/configs/debug/minimax_m2_tiny_random.json new file mode 100644 index 00000000..677d18d8 --- /dev/null +++ b/configs/debug/minimax_m2_tiny_random.json @@ -0,0 +1,6 @@ +{ + "num_hidden_layers": 1, + "hidden_size": 3072, + "intermediate_size": 1536, + "max_position_embeddings": 4096 +} diff --git a/docs/en/getting_started/README.md b/docs/en/getting_started/README.md index 46b738af..7885c0d6 100644 --- a/docs/en/getting_started/README.md +++ b/docs/en/getting_started/README.md @@ -34,8 +34,10 @@ MAX_JOBS=8 uv pip install flash-attn --no-build-isolation Use `cu129` instead of `cu130` for CUDA 12.9. -`einops`, `sglang-kernel`, and the training, benchmark, and test packages are all -runtime dependencies, so workflow-specific extras are not required. +`einops`, `sglang-kernel==0.4.5`, and the training, benchmark, and test packages +are runtime dependencies, so workflow-specific extras are not required. The SGL +kernel package is pinned to the validated PyTorch/CUDA ABI; other versions fail +provider setup instead of falling back silently. Sparse-vLLM supports Qwen3.5/Qwen3.6/Qwen3.8 checkpoints in unquantized BF16 and block-scaled FP8 formats. All three share the `qwen3_5` runtime architecture diff --git a/docs/zh/getting_started/README.md b/docs/zh/getting_started/README.md index 19596d3c..f5a3b0cf 100644 --- a/docs/zh/getting_started/README.md +++ b/docs/zh/getting_started/README.md @@ -31,8 +31,9 @@ MAX_JOBS=8 uv pip install flash-attn --no-build-isolation CUDA 12.9 环境将 `cu130` 换成 `cu129`。 -`einops`、`sglang-kernel` 以及训练、benchmark 和测试包均已是主依赖， -不再需要工作流专用 extra。 +`einops`、`sglang-kernel==0.4.5` 以及训练、benchmark 和测试包均已是主依赖， +不再需要工作流专用 extra。SGL kernel package 固定到已经验证的 PyTorch/CUDA ABI； +其他版本会在 Provider 准备阶段明确失败，不会静默 fallback。 Sparse-vLLM 当前支持未量化 BF16 和 block-scaled FP8 格式的 Qwen3.5/Qwen3.6/Qwen3.8 checkpoint。三者共享 `qwen3_5` 运行时架构和支持矩阵。 diff --git a/pyproject.toml b/pyproject.toml index 5b05b93c..f744348a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ dependencies = [ "pillow", "torchvision", "einops", - "sglang-kernel>=0.4.5,<0.5", + "sglang-kernel==0.4.5", "tqdm", "loguru", "fastapi>=0.100", diff --git a/scripts/debug/compare_decode_graph_eager_logits.py b/scripts/debug/compare_decode_graph_eager_logits.py index f99b85e0..af94c87e 100644 --- a/scripts/debug/compare_decode_graph_eager_logits.py +++ b/scripts/debug/compare_decode_graph_eager_logits.py @@ -570,11 +570,16 @@ def _run_decode_logits( max_tokens: int, hyper_params: dict[str, Any], use_graph: bool, + same_provider_eager: bool = False, trace_selection: bool = False, ) -> tuple[torch.Tensor, list[dict[str, Any]], dict[str, Any]]: from sparsevllm import LLM, SamplingParams - if os.getenv("SPARSEVLLM_DEBUG_SKIP_ENGINE_WARMUP", "0") == "1": + construct_with_graph = bool(use_graph or same_provider_eager) + if ( + os.getenv("SPARSEVLLM_DEBUG_SKIP_ENGINE_WARMUP", "0") == "1" + or (same_provider_eager and not use_graph) + ): LLM._warmup = lambda self: None engine_kwargs = { @@ -583,11 +588,15 @@ def _run_decode_logits( "max_model_len": max(prompt_lens) + max_tokens + 100, "max_num_seqs_in_batch": batch_size, "max_decoding_seqs": batch_size, - "decode_graph": bool(use_graph), + "decode_graph": construct_with_graph, "decode_graph_capture_sampling": False, "throughput_log_interval_s": 0.0, } llm = LLM(model_path, **engine_kwargs) + if same_provider_eager and not use_graph: + # Keep the graph-selected provider, but execute every decode step through + # DecodeCudaGraphRunner.run_eager_static() as the graph-independent oracle. + llm.config.decode_graph = False graph_counters_before = _start_graph_measurement(llm) method_calls = _install_method_instrumentation(llm) captured: list[torch.Tensor] = [] @@ -779,6 +788,14 @@ def _build_parser() -> argparse.ArgumentParser: parser.add_argument("--atol", type=float, default=0.05) parser.add_argument("--rtol", type=float, default=0.05) parser.add_argument("--trace_selection", action="store_true") + parser.add_argument( + "--same_provider_eager", + action="store_true", + help=( + "Construct the eager control with decode_graph enabled so it binds " + "the same provider, skip startup capture, then execute eager-static." + ), + ) return parser @@ -853,6 +870,7 @@ def main(argv: list[str] | None = None): max_tokens=args.max_tokens, hyper_params=hyper_params, use_graph=False, + same_provider_eager=args.same_provider_eager, trace_selection=args.trace_selection, ) graph_logits, graph_trace, graph_runtime = _run_decode_logits_isolated( @@ -863,6 +881,7 @@ def main(argv: list[str] | None = None): max_tokens=args.max_tokens, hyper_params=hyper_params, use_graph=True, + same_provider_eager=args.same_provider_eager, trace_selection=args.trace_selection, ) @@ -936,6 +955,7 @@ def main(argv: list[str] | None = None): "prompt_lens": prompt_lens, "batch_size": args.batch_size, "max_tokens": args.max_tokens, + "same_provider_eager": bool(args.same_provider_eager), "hyper_params": hyper_params, "comparison": comparison, "generated_token_ids": { @@ -980,6 +1000,7 @@ def main(argv: list[str] | None = None): { "status": "failed", "method": args.method, + "same_provider_eager": bool(args.same_provider_eager), "error": traceback.format_exc(), }, indent=2, diff --git a/src/sparsevllm/configs/cuda_graph.py b/src/sparsevllm/configs/cuda_graph.py index ddc23a75..5114404a 100644 --- a/src/sparsevllm/configs/cuda_graph.py +++ b/src/sparsevllm/configs/cuda_graph.py @@ -6,23 +6,34 @@ from sparsevllm.configs.common import _coerce_bool_config from sparsevllm.method_registry import ( DECODE_CUDA_GRAPH_SUPPORTED_METHODS, + decode_graph_path_id, + decode_sparse_long_text_threshold, is_decode_cuda_graph_supported, is_tp_decode_cuda_graph_supported, ) from sparsevllm.utils.log import log_once + def _default_decode_cuda_graph_capture_sizes(max_decoding_seqs: int) -> list[int]: + """Return at most 32 batch buckets, dense where padding hurts most.""" max_decoding_seqs = int(max_decoding_seqs) if max_decoding_seqs <= 0: raise ValueError(f"max_decoding_seqs must be > 0, got {max_decoding_seqs}.") - sizes: list[int] = [] - size = 1 - while size < max_decoding_seqs: - sizes.append(size) - size *= 2 - if not sizes or sizes[-1] != max_decoding_seqs: - sizes.append(max_decoding_seqs) + dense_limit = min(8, max_decoding_seqs) + sizes = list(range(1, dense_limit + 1)) + if max_decoding_seqs <= dense_limit: + return sizes + + # Keep small decode batches exact, then use aligned, bounded-width buckets. + # The adaptive stride caps the auto plan at 32 batch families even for a + # very large scheduler limit; explicit capture sizes remain unrestricted. + remaining_bucket_budget = 32 - dense_limit + span = max_decoding_seqs - dense_limit + stride = max(4, (span + remaining_bucket_budget - 1) // remaining_bucket_budget) + stride = ((stride + 3) // 4) * 4 + sizes.extend(range(dense_limit + stride, max_decoding_seqs, stride)) + sizes.append(max_decoding_seqs) return sizes @@ -100,7 +111,7 @@ def _select_decode_cuda_graph_batch_size( if size >= real_batch_size: return size raise ValueError( - "decode_graph capture sizes do not cover current decode batch: " + "decode_cuda_graph capture sizes do not cover current decode batch: " f"batch_size={real_batch_size}, capture_sizes={sizes}." ) @@ -163,7 +174,316 @@ def _normalize_decode_cuda_graph_context_policy(value: str | None) -> str: return policy +def _normalize_decode_graph_shape_policy(value: str | None) -> str: + policy = ( + str(value or "batch_only") + .strip() + .lower() + .replace("-", "_") + ) + policy = { + "context_bucketed": "bucketed", + "batch": "batch_only", + "bs_only": "batch_only", + }.get(policy, policy) + if policy not in {"bucketed", "batch_only"}: + raise ValueError( + "decode_graph_shape_policy must be 'bucketed' or 'batch_only', " + f"got {policy!r}." + ) + return policy + + +def _select_evenly_spaced_sizes( + sizes: list[int] | tuple[int, ...], limit: int +) -> list[int]: + candidates = sorted(set(int(size) for size in sizes)) + limit = int(limit) + if limit <= 0: + raise ValueError(f"batch-only capture limit must be positive, got {limit}.") + if len(candidates) <= limit: + return candidates + dense = candidates[: min(8, limit)] + remaining = limit - len(dense) + if remaining <= 0: + dense[-1] = candidates[-1] + return sorted(set(dense)) + tail = candidates[len(dense) :] + indices = ( + { + round(index * (len(tail) - 1) / (remaining - 1)) + for index in range(remaining) + } + if remaining > 1 + else {len(tail) - 1} + ) + return sorted(set(dense + [tail[index] for index in sorted(indices)])) + + +def _decode_cuda_graph_reachable_families(config) -> list[tuple[bool, int]]: + method = str(config.sparse_method or "") + max_model_len = int(config.max_model_len) + if not method: + return [(False, max_model_len)] + threshold = decode_sparse_long_text_threshold( + method, + num_sink_tokens=config.sink_keep_tokens, + decode_keep_tokens=config.decode_keep_tokens, + num_recent_tokens=config.recent_keep_tokens, + ) + families: list[tuple[bool, int]] = [] + if threshold >= 2: + families.append((False, min(threshold, max_model_len))) + if threshold + 2 <= max_model_len: + families.append((True, max_model_len)) + deduplicated: dict[str, tuple[bool, int]] = {} + for is_long_text, capacity in families: + deduplicated[decode_graph_path_id(method, is_long_text)] = ( + is_long_text, + capacity, + ) + if not deduplicated: + raise ValueError( + "No reachable sparse decode CUDA Graph family for batch-only capture." + ) + return list(deduplicated.values()) + + +def build_decode_cuda_graph_batch_only_startup_plan( + config, +) -> list[tuple[int, int, bool]]: + batches = sorted(set(int(size) for size in config.decode_graph_capture_sizes)) + families = _decode_cuda_graph_reachable_families(config) + limit = min( + int(config.decode_graph_startup_capture_limit), + int(config.decode_graph_max_cached_graphs), + ) + required = len(batches) * len(families) + if required > limit: + raise ValueError( + "batch-only decode CUDA Graph startup capture must cover every " + f"batch/topology lane: required={required}, limit={limit}." + ) + return sorted( + ( + (batch_size, context_capacity, is_long_text) + for batch_size in batches + for is_long_text, context_capacity in families + ), + reverse=True, + ) + + +def build_decode_cuda_graph_startup_plan( + capture_sizes: list[int] | tuple[int, ...], + context_sizes: list[int] | tuple[int, ...], + limit: int, + *, + mandatory: tuple[int, int] | None = None, +) -> list[tuple[int, int]]: + """Select dense batch coverage and coarse context coverage within ``limit``.""" + batches = sorted(set(int(size) for size in capture_sizes)) + contexts = sorted(set(int(size) for size in context_sizes)) + limit = int(limit) + if limit <= 0: + raise ValueError(f"decode_graph_startup_capture_limit must be positive, got {limit}.") + if not batches or not contexts: + return [] + if any(batch <= 0 for batch in batches) or any(context <= 0 for context in contexts): + raise ValueError( + "decode CUDA Graph startup buckets must be positive: " + f"batch_sizes={batches}, context_sizes={contexts}." + ) + if limit < len(batches): + raise ValueError( + "decode CUDA Graph startup capture limit must cover every batch bucket: " + f"limit={limit}, batch_buckets={len(batches)}." + ) + + full_plan = [(batch, context) for batch in batches for context in contexts] + if len(full_plan) <= limit: + return full_plan + + # Every batch family gets its largest context first. Remaining quota is + # biased toward smaller batches and spread over the existing power-of-two + # context buckets. A missing exact context can still reuse that batch's + # next larger captured graph, at the context-padding cost measured by the + # benchmark rather than triggering a runtime capture. + selected: list[tuple[int, int]] = [] + quotas = [limit // len(batches)] * len(batches) + for idx in range(limit % len(batches)): + quotas[idx] += 1 + for batch, quota in zip(batches, quotas): + if quota <= 0: + continue + if quota >= len(contexts): + chosen = contexts + elif quota == 1: + chosen = [contexts[-1]] + else: + indices = { + round(idx * (len(contexts) - 1) / (quota - 1)) + for idx in range(quota) + } + chosen = [contexts[idx] for idx in sorted(indices)] + selected.extend((batch, context) for context in chosen) + + if mandatory is not None: + mandatory = (int(mandatory[0]), int(mandatory[1])) + if mandatory is not None and mandatory in full_plan and mandatory not in selected: + mandatory_batch = int(mandatory[0]) + replace_idx = next( + ( + idx for idx, pair in enumerate(selected) + if pair[0] == mandatory_batch + and pair[1] != contexts[-1] + and sum(selected_pair[0] == mandatory_batch for selected_pair in selected) > 1 + ), + -1, + ) + if replace_idx >= 0: + selected[replace_idx] = mandatory + + plan = sorted(set(selected)) + if len(plan) != min(limit, len(full_plan)): + raise RuntimeError( + "decode CUDA Graph startup planner produced an incomplete plan: " + f"expected={min(limit, len(full_plan))}, actual={len(plan)}." + ) + missing_max_batches = [ + batch for batch in batches if (batch, contexts[-1]) not in plan + ] + if missing_max_batches: + raise RuntimeError( + "decode CUDA Graph startup plan must retain the largest context for " + f"every batch bucket, missing={missing_max_batches}." + ) + return plan + + +def build_decode_cuda_graph_startup_family_plan(config) -> list[tuple[int, int, bool]]: + """Build graph keys largest-first so captures reuse the shared graph pool.""" + if ( + str( + getattr( + config, + "decode_graph_shape_policy", + "batch_only", + ) + ) + == "batch_only" + ): + return build_decode_cuda_graph_batch_only_startup_plan(config) + batches = sorted(set(int(size) for size in config.decode_graph_capture_sizes)) + contexts = sorted(set(int(size) for size in config.decode_graph_context_sizes)) + limit = min( + int(config.decode_graph_startup_capture_limit), + int(config.decode_graph_max_cached_graphs), + ) + method = str(config.sparse_method or "") + if not method: + return sorted( + ( + (batch, context, False) + for batch, context in build_decode_cuda_graph_startup_plan( + batches, + contexts, + limit, + ) + ), + reverse=True, + ) + + threshold = decode_sparse_long_text_threshold( + method, + num_sink_tokens=config.sink_keep_tokens, + decode_keep_tokens=config.decode_keep_tokens, + num_recent_tokens=config.recent_keep_tokens, + ) + family_contexts: list[tuple[bool, list[int]]] = [] + if threshold >= 2: + family_contexts.append((False, contexts)) + if threshold + 2 <= int(config.max_model_len): + long_contexts = [context for context in contexts if context > threshold] + if long_contexts: + family_contexts.append((True, long_contexts)) + if not family_contexts: + raise ValueError( + "No reachable sparse decode CUDA Graph family for startup capture: " + f"method={method!r}, threshold={threshold}, max_model_len={config.max_model_len}." + ) + + lanes = [ + (batch, is_long_text, lane_contexts) + for batch in batches + for is_long_text, lane_contexts in family_contexts + ] + if limit < len(lanes): + raise ValueError( + "decode CUDA Graph sparse startup capture limit must cover every " + "batch/family lane: " + f"limit={limit}, required={len(lanes)}, batch_buckets={len(batches)}, " + f"families={len(family_contexts)}." + ) + + full_plan = [ + (batch, context, is_long_text) + for batch, is_long_text, lane_contexts in lanes + for context in lane_contexts + ] + if len(full_plan) <= limit: + return sorted(full_plan, reverse=True) + + target_size = min(limit, len(full_plan)) + quotas = [1] * len(lanes) + remaining = target_size - len(lanes) + while remaining > 0: + progressed = False + for lane_idx, (_, _, lane_contexts) in enumerate(lanes): + if quotas[lane_idx] >= len(lane_contexts): + continue + quotas[lane_idx] += 1 + remaining -= 1 + progressed = True + if remaining == 0: + break + if not progressed: + raise RuntimeError( + "decode CUDA Graph sparse startup planner could not allocate " + f"remaining budget={remaining}." + ) + selected: list[tuple[int, int, bool]] = [] + for (batch, is_long_text, lane_contexts), quota in zip(lanes, quotas): + if quota >= len(lane_contexts): + chosen = lane_contexts + elif quota == 1: + chosen = [lane_contexts[-1]] + else: + indices = { + round(idx * (len(lane_contexts) - 1) / (quota - 1)) + for idx in range(quota) + } + chosen = [lane_contexts[idx] for idx in sorted(indices)] + selected.extend( + (batch, context, is_long_text) for context in chosen + ) + plan = sorted(set(selected), reverse=True) + if len(plan) != target_size: + raise RuntimeError( + "decode CUDA Graph sparse startup planner produced an incomplete plan: " + f"expected={target_size}, actual={len(plan)}." + ) + return plan + + def normalize_decode_cuda_graph(config) -> None: + config.decode_graph_shape_policy = _normalize_decode_graph_shape_policy( + getattr( + config, + "decode_graph_shape_policy", + "batch_only", + ) + ) if config.decode_graph_max_cached_graphs is not None: config.decode_graph_max_cached_graphs = int(config.decode_graph_max_cached_graphs) if config.decode_graph_max_cached_graphs <= 0: @@ -171,6 +491,43 @@ def normalize_decode_cuda_graph(config) -> None: "decode_graph_max_cached_graphs must be a positive integer or None, " f"got {config.decode_graph_max_cached_graphs}." ) + startup_capture_setting = config.decode_graph_startup_capture + startup_capture_auto = startup_capture_setting is None + if startup_capture_auto: + config.decode_graph_startup_capture = bool(config.decode_graph) + else: + config.decode_graph_startup_capture = _coerce_bool_config( + "decode_graph_startup_capture", + startup_capture_setting, + ) + if config.decode_graph_startup_capture_limit is None: + config.decode_graph_startup_capture_limit = ( + 48 if config.sparse_method else 32 + ) + config.decode_graph_startup_capture_limit = int( + config.decode_graph_startup_capture_limit + ) + if config.decode_graph_startup_capture_limit <= 0: + raise ValueError( + "decode_graph_startup_capture_limit must be a positive integer, " + f"got {config.decode_graph_startup_capture_limit}." + ) + if config.decode_graph_startup_capture: + if not config.decode_graph: + raise ValueError("decode_graph_startup_capture requires decode_graph=True.") + if config.decode_graph_max_cached_graphs is None: + config.decode_graph_max_cached_graphs = ( + config.decode_graph_startup_capture_limit + ) + if ( + config.decode_graph + and config.decode_graph_shape_policy == "batch_only" + and not config.decode_graph_startup_capture + ): + raise ValueError( + "decode_graph_shape_policy='batch_only' requires " + "decode_graph_startup_capture=True." + ) if config.decode_graph_capture_sampling and not config.decode_graph: raise ValueError("decode_graph_capture_sampling requires decode_graph=True.") config.decode_graph_context_policy = _normalize_decode_cuda_graph_context_policy( @@ -215,11 +572,38 @@ def normalize_decode_cuda_graph(config) -> None: repr(method) for method in sorted(DECODE_CUDA_GRAPH_SUPPORTED_METHODS) if method ) raise ValueError(f"decode_graph supports these methods only: '', {supported}.") + capture_sizes_setting = config.decode_graph_capture_sizes + capture_sizes_auto = capture_sizes_setting is None or ( + isinstance(capture_sizes_setting, str) + and capture_sizes_setting.strip().lower() in {"", "auto"} + ) config.decode_graph_capture_sizes = _resolve_decode_cuda_graph_capture_sizes( - config.decode_graph_capture_sizes, + capture_sizes_setting, config.max_decoding_seqs, ) config.decode_graph_context_sizes = _resolve_decode_cuda_graph_context_sizes( config.decode_graph_context_sizes, config.max_model_len, ) + if config.decode_graph_shape_policy == "batch_only" and capture_sizes_auto: + path_count = len(_decode_cuda_graph_reachable_families(config)) + graph_budget = min( + int(config.decode_graph_startup_capture_limit), + int(config.decode_graph_max_cached_graphs), + ) + config.decode_graph_capture_sizes = _select_evenly_spaced_sizes( + config.decode_graph_capture_sizes, + graph_budget // path_count, + ) + if config.decode_graph_startup_capture: + startup_plan = build_decode_cuda_graph_startup_family_plan(config) + log_once( + "Decode CUDA Graph startup precapture enabled " + f"({'default' if startup_capture_auto else 'explicit'}): " + f"shape_policy={config.decode_graph_shape_policy}, " + f"budget={config.decode_graph_startup_capture_limit}, " + f"cache_limit={config.decode_graph_max_cached_graphs}, " + f"planned_graphs={len(startup_plan)}, " + f"batch_buckets={config.decode_graph_capture_sizes}, " + f"context_buckets={config.decode_graph_context_sizes}." + ) diff --git a/src/sparsevllm/configs/groups.py b/src/sparsevllm/configs/groups.py index ef997cc2..72f85626 100644 --- a/src/sparsevllm/configs/groups.py +++ b/src/sparsevllm/configs/groups.py @@ -25,12 +25,15 @@ class DecodeCudaGraphConfig: """Decode CUDA Graph capture and compatibility settings.""" decode_graph: bool = False + decode_graph_shape_policy: str = "batch_only" decode_graph_capture_sampling: bool = False decode_graph_capture_sizes: str | int | list[int] | tuple[int, ...] | None = "auto" decode_graph_context_sizes: str | int | list[int] | tuple[int, ...] | None = "auto" decode_graph_context_sizes_auto: bool = field(default=False, init=False) decode_graph_context_policy: str = "current" decode_graph_max_cached_graphs: int | None = None + decode_graph_startup_capture: bool | None = None + decode_graph_startup_capture_limit: int | None = None sparse_attn_score_dtype: str = "float32" @dataclass(kw_only=True) diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index 566f9c7b..cf9d6f24 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -153,7 +153,7 @@ def load_and_validate_model(config) -> None: config.hf_config, config.tiny_random_config, validate_standard_head_shape=( - model_spec.attention_cache_layout == "explicit_kv" + model_spec.tiny_random_requires_standard_head_shape ), ) log_once( @@ -178,7 +178,11 @@ def load_and_validate_model(config) -> None: or config_get(config.outer_hf_config, "torch_dtype", "bfloat16") ), ) - if config.tiny_random and config.quantization_config.enabled: + if ( + config.tiny_random + and config.quantization_config.enabled + and not model_spec.supports_quantized_tiny_random + ): raise NotImplementedError( "Tiny random mode does not support quantized model weights." ) diff --git a/src/sparsevllm/debug/tiny_random.py b/src/sparsevllm/debug/tiny_random.py index bbde1853..972142c7 100644 --- a/src/sparsevllm/debug/tiny_random.py +++ b/src/sparsevllm/debug/tiny_random.py @@ -183,12 +183,21 @@ def initialize_sparse_model( hf_config: Any, *, seed: int, + quantized: bool = False, ) -> None: from sparsevllm.utils.loader import ( _target_weight_name_for_model, default_weight_loader, ) + if quantized: + _initialize_quantized_sparse_model(model, seed=seed) + print( + "Initialized quantized model weights from deterministic tiny random " + f"seed={int(seed)} without reading checkpoint tensors" + ) + return + reference = build_tiny_random_hf_model(hf_config, seed=seed) packed_modules_mapping = getattr(model, "packed_modules_mapping", {}) loaded_count = 0 @@ -241,3 +250,43 @@ def initialize_sparse_model( f"Initialized {loaded_count} model weights from deterministic tiny random " f"seed={int(seed)} without reading checkpoint tensors" ) + + +@torch.inference_mode() +def _initialize_quantized_sparse_model(model: nn.Module, *, seed: int) -> None: + generators: dict[torch.device, torch.Generator] = {} + + def generator_for(device: torch.device) -> torch.Generator: + generator = generators.get(device) + if generator is None: + generator = torch.Generator(device=device) + generator.manual_seed(int(seed)) + generators[device] = generator + return generator + + initialized = 0 + for parameter in model.parameters(): + if not parameter.dtype.is_floating_point: + parameter.zero_() + initialized += parameter.numel() + continue + values = torch.empty( + parameter.shape, + dtype=torch.float32, + device=parameter.device, + ) + values.normal_(mean=0.0, std=0.02, generator=generator_for(parameter.device)) + parameter.copy_(values.to(parameter.dtype)) + initialized += parameter.numel() + + for name, buffer in model.named_buffers(): + if name.endswith("weight_scale_inv"): + buffer.fill_(1.0) + + for module in model.modules(): + if hasattr(module, "_quantized_weight_loaded"): + module._quantized_weight_loaded = True + module._quantized_loaded_ranges = [(0, int(module.weight.shape[0]))] + + if initialized <= 0: + raise RuntimeError("Quantized tiny random initialization found no parameters.") diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index 8810b519..e31ff5b5 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -12,12 +12,22 @@ from sparsevllm.config import Config from sparsevllm.distributed import ParallelContext -from sparsevllm.engine.sequence import Sequence +from sparsevllm.engine.decode_graph_contract import ( + CacheDecodeGraphState, + DecodeGraphContract, + DecodeGraphInputs, +) from sparsevllm.engine.prefill import ( PREFILL_EXECUTION_CHUNKED, PREFILL_EXECUTION_RAW_OFFLOAD, ) -from sparsevllm.method_registry import SUPPORTED_SPARSE_METHODS, normalize_sparse_method +from sparsevllm.engine.sequence import Sequence +from sparsevllm.method_registry import ( + SUPPORTED_SPARSE_METHODS, + decode_graph_path_id, + decode_sparse_long_text_threshold, + normalize_sparse_method, +) from sparsevllm.kernels.triton.store_kvcache import store_kvcache import sparsevllm.platforms as platforms from sparsevllm.models.layout import resolve_attention_qk_head_dim @@ -626,6 +636,46 @@ def prepare_step(self, seqs: list[Sequence], is_prefill: bool): return self._prepare_prefill(seqs) return self._prepare_decode(seqs) + def init_decode_graph_state( + self, + contract: DecodeGraphContract, + inputs: DecodeGraphInputs, + ) -> CacheDecodeGraphState: + """Bind cache-owned metadata to one graph's stable public inputs.""" + inputs.validate(contract) + return CacheDecodeGraphState(contract=contract, inputs=inputs) + + def prepare_decode_graph_step( + self, + seqs: list[Sequence], + state: CacheDecodeGraphState, + ): + """Compatibility adapter while method-specific managers migrate.""" + inputs = state.inputs + result = self.prepare_decode_static( + seqs, + inputs.input_ids, + inputs.positions, + inputs.write_slot_mapping, + inputs.context_lens, + inputs.request_indices, + ) + real_batch_size = len(seqs) + inputs.active_mask[:real_batch_size].fill_(True) + inputs.active_mask[real_batch_size:].fill_(state.contract.padding.active) + return result + + def prepare_decode_graph_in(self, state: CacheDecodeGraphState) -> None: + """Run fixed device-side cache metadata preparation during capture/replay.""" + del state + + def decode_graph_state_keepalive_tensors( + self, + state: CacheDecodeGraphState, + ) -> list[torch.Tensor]: + del state + return self.decode_graph_keepalive_tensors() + @abstractmethod def allocate_kv_cache(self): """自动计算并物理分配 KV Cache 张量""" @@ -1111,6 +1161,42 @@ def decode_graph_context_capacity( del seqs, requested_context_capacity, current_context_capacity return None + def decode_graph_path_id(self, is_long_text: bool) -> str: + return decode_graph_path_id( + str(getattr(self.config, "sparse_method", "") or ""), + bool(is_long_text), + ) + + def decode_graph_batch_only_capacity( + self, is_long_text: bool + ) -> int: + method = str(getattr(self.config, "sparse_method", "") or "") + max_model_len = int(self.config.max_model_len) + if not method or is_long_text: + return max_model_len + threshold = decode_sparse_long_text_threshold( + method, + num_sink_tokens=self.config.sink_keep_tokens, + decode_keep_tokens=self.config.decode_keep_tokens, + num_recent_tokens=self.config.recent_keep_tokens, + ) + return min(max_model_len, int(threshold)) + + def validate_decode_graph_batch_only_capacity( + self, + seqs: list[Sequence], + *, + capacity: int, + is_long_text: bool, + ) -> None: + actual = max(int(seq.num_tokens) for seq in seqs) + if int(capacity) < actual: + raise RuntimeError( + "batch-only decode CUDA Graph path capacity does not cover the " + f"request: capacity={capacity}, actual={actual}, " + f"is_long_text={is_long_text}." + ) + def decode_graph_force_eager(self) -> bool: """Whether this method should bypass graph replay for diagnostics.""" return False @@ -1506,6 +1592,13 @@ def free_slot_stats(self) -> dict[str, int]: """Return a small set of free-slot stats for logging/debugging.""" return {"free_slots": int(self.num_free_slots)} + def _debug_token_slots_for_mapping( + self, + layer_idx: int | None, + ) -> torch.Tensor: + token_slots = getattr(self, "buffer_req_to_token_slots") + return token_slots if layer_idx is None else token_slots[layer_idx] + def debug_state_summary(self) -> dict[str, Any]: """Return a synchronized-test snapshot without touching the inference hot path.""" live_rows = {} @@ -1519,10 +1612,9 @@ def debug_state_summary(self) -> dict[str, Any]: if not isinstance(mapping, dict) or not mapping: continue row_seq_lens = getattr(self, "row_seq_lens") - token_slots = getattr(self, "buffer_req_to_token_slots") + token_slots = self._debug_token_slots_for_mapping(layer_idx) if layer_idx is not None: row_seq_lens = row_seq_lens[layer_idx] - token_slots = token_slots[layer_idx] records = [] for seq_id, row_idx in sorted(mapping.items()): row_len = int(row_seq_lens[row_idx]) diff --git a/src/sparsevllm/engine/cache_manager/deltakv_base.py b/src/sparsevllm/engine/cache_manager/deltakv_base.py index 40e01ef7..f07c1dc4 100644 --- a/src/sparsevllm/engine/cache_manager/deltakv_base.py +++ b/src/sparsevllm/engine/cache_manager/deltakv_base.py @@ -1044,6 +1044,13 @@ def get_layer_buffer_req_to_token_slots(self, layer_idx: int) -> torch.Tensor: # most historical tokens are either compressed or reconstructed on-the-fly. raise NotImplementedError("DeltaKV sparse layers should use build_*_compute_view().") + def _debug_token_slots_for_mapping( + self, + layer_idx: int | None, + ) -> torch.Tensor: + del layer_idx + return self.full_layer_slots_map + @property def num_free_slots(self) -> int: # Scheduling should be conservative: we must be able to allocate both diff --git a/src/sparsevllm/engine/cache_manager/standard.py b/src/sparsevllm/engine/cache_manager/standard.py index 19ca328c..2f09aaa8 100644 --- a/src/sparsevllm/engine/cache_manager/standard.py +++ b/src/sparsevllm/engine/cache_manager/standard.py @@ -10,7 +10,10 @@ from sparsevllm.config import Config from sparsevllm.distributed import ParallelContext -from sparsevllm.engine.sequence import Sequence +from sparsevllm.engine.decode_graph_contract import ( + CacheDecodeGraphState, + DecodeGraphHostInputs, +) from sparsevllm.engine.prefix_cache import ( PrefixCacheBlock, PrefixTransferKind, @@ -19,6 +22,7 @@ select_write_through_candidates, usable_prefix_cache_tokens, ) +from sparsevllm.engine.sequence import Sequence from sparsevllm.utils.log import logger, log_level from sparsevllm.utils.profiler import profiler from sparsevllm.platforms import device_runtime @@ -1520,6 +1524,54 @@ def prepare_decode_static( Used by CUDA Graph decode replay: tensor addresses must stay stable, so this avoids the ordinary per-step metadata tensor allocation path. """ + return self._prepare_decode_graph_buffers( + seqs, + input_ids=input_ids, + positions=positions, + slot_mapping=slot_mapping, + context_lens=context_lens, + req_indices=req_indices, + ) + + def prepare_decode_graph_step( + self, + seqs: list[Sequence], + state: CacheDecodeGraphState, + ): + inputs = state.inputs + return self._prepare_decode_graph_buffers( + seqs, + input_ids=inputs.input_ids, + positions=inputs.positions, + slot_mapping=inputs.write_slot_mapping, + context_lens=inputs.context_lens, + req_indices=inputs.request_indices, + active_mask=inputs.active_mask, + host_inputs=inputs.host, + padding_write_slot=int(state.contract.padding.write_slot), + padding_active=bool(state.contract.padding.active), + mirror_first_real_row_for_reads=bool( + state.contract.padding.mirror_first_real_row_for_reads + ), + context_capacity=int(state.contract.context_capacity), + ) + + def _prepare_decode_graph_buffers( + self, + seqs: list[Sequence], + *, + input_ids: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + context_lens: torch.Tensor, + req_indices: torch.Tensor, + active_mask: torch.Tensor | None = None, + host_inputs: DecodeGraphHostInputs | None = None, + padding_write_slot: int = -1, + padding_active: bool = False, + mirror_first_real_row_for_reads: bool = True, + context_capacity: int | None = None, + ): with profiler.record("cache_prepare_decode"): self._poll_prefix_offload() self._prefix_offload_step_h2d_operations = [] @@ -1540,11 +1592,35 @@ def prepare_decode_static( "Static decode graph batch is smaller than the real decode batch: " f"graph={graph_batch_size}, real={real_batch_size}." ) + if active_mask is not None and active_mask.numel() != graph_batch_size: + raise ValueError( + "Static decode active_mask must match the graph batch size." + ) + if not mirror_first_real_row_for_reads: + raise ValueError( + "StandardCacheManager requires padded read rows to mirror the " + "first real request." + ) input_ids_list = [seq.decode_input_token for seq in seqs] positions_list = [seq.decode_input_position for seq in seqs] seq_ids = [seq.seq_id for seq in seqs] + if context_capacity is not None: + prospective_rows = np.asarray( + [self._get_free_row(seq_id) for seq_id in seq_ids], + dtype=np.int64, + ) + max_requested_context_len = int( + (self.row_seq_lens[prospective_rows] + 1).max() + ) + if max_requested_context_len > context_capacity: + raise ValueError( + "Decode request exceeded the captured graph context capacity: " + f"requested={max_requested_context_len} " + f"captured={context_capacity}." + ) + new_slots_batch, real_context_lens, row_indices = self._allocate_decode_batch_static( seq_ids, graph_batch_size, @@ -1552,27 +1628,80 @@ def prepare_decode_static( for seq, slot in zip(seqs, new_slots_batch): self._record_prefix_materialization(seq, [seq.decode_input_token], slot.reshape(1)) - input_ids[:real_batch_size].copy_(torch.tensor(input_ids_list, dtype=torch.int64)) - positions[:real_batch_size].copy_(torch.tensor(positions_list, dtype=torch.int64)) slot_mapping[:real_batch_size].copy_(new_slots_batch) - context_lens[:real_batch_size].copy_( - torch.from_numpy(real_context_lens.astype(np.int32, copy=False)) - ) - req_indices[:real_batch_size].copy_( - torch.from_numpy(row_indices.astype(np.int32, copy=False)) - ) + if host_inputs is None: + input_ids[:real_batch_size].copy_( + torch.tensor(input_ids_list, dtype=torch.int64) + ) + positions[:real_batch_size].copy_( + torch.tensor(positions_list, dtype=torch.int64) + ) + context_lens[:real_batch_size].copy_( + torch.from_numpy(real_context_lens.astype(np.int32, copy=False)) + ) + req_indices[:real_batch_size].copy_( + torch.from_numpy(row_indices.astype(np.int32, copy=False)) + ) + else: + host_inputs.input_ids.numpy()[:real_batch_size] = input_ids_list + host_inputs.positions.numpy()[:real_batch_size] = positions_list + host_inputs.context_lens.numpy()[:real_batch_size] = ( + real_context_lens.astype(np.int32, copy=False) + ) + host_inputs.request_indices.numpy()[:real_batch_size] = ( + row_indices.astype(np.int32, copy=False) + ) + host_inputs.active_mask[:real_batch_size].fill_(True) + non_blocking = bool(host_inputs.input_ids.is_pinned()) + input_ids[:real_batch_size].copy_( + host_inputs.input_ids[:real_batch_size], + non_blocking=non_blocking, + ) + positions[:real_batch_size].copy_( + host_inputs.positions[:real_batch_size], + non_blocking=non_blocking, + ) + context_lens[:real_batch_size].copy_( + host_inputs.context_lens[:real_batch_size], + non_blocking=non_blocking, + ) + req_indices[:real_batch_size].copy_( + host_inputs.request_indices[:real_batch_size], + non_blocking=non_blocking, + ) + assert active_mask is not None + active_mask[:real_batch_size].copy_( + host_inputs.active_mask[:real_batch_size], + non_blocking=non_blocking, + ) if graph_batch_size > real_batch_size: # CUDA Graph replay is shape-static. Padded rows mirror the first - # real request for read-only attention work, but use slot -1 so - # they never write KV or consume persistent cache capacity. + # real request for read-only work, but use the contract's safe + # write sentinel so they never consume persistent cache capacity. first_context_len = int(real_context_lens[0]) first_row_idx = int(row_indices[0]) input_ids[real_batch_size:].fill_(int(input_ids_list[0])) positions[real_batch_size:].fill_(int(positions_list[0])) - slot_mapping[real_batch_size:].fill_(-1) + slot_mapping[real_batch_size:].fill_(padding_write_slot) context_lens[real_batch_size:].fill_(first_context_len) req_indices[real_batch_size:].fill_(first_row_idx) + if active_mask is not None: + active_mask[real_batch_size:].fill_(padding_active) + if host_inputs is not None: + host_inputs.input_ids[real_batch_size:].fill_( + int(input_ids_list[0]) + ) + host_inputs.positions[real_batch_size:].fill_( + int(positions_list[0]) + ) + host_inputs.context_lens[real_batch_size:].fill_( + first_context_len + ) + host_inputs.request_indices[real_batch_size:].fill_( + first_row_idx + ) + host_inputs.active_mask[real_batch_size:].fill_(padding_active) self.layer_batch_state.slot_mapping = slot_mapping self.layer_batch_state.context_lens = context_lens diff --git a/src/sparsevllm/engine/decode_cuda_graph.py b/src/sparsevllm/engine/decode_cuda_graph.py index 951ec3c3..cf3736d3 100644 --- a/src/sparsevllm/engine/decode_cuda_graph.py +++ b/src/sparsevllm/engine/decode_cuda_graph.py @@ -6,9 +6,14 @@ import torch -from sparsevllm.engine.sequence import Sequence -from sparsevllm.configs.cuda_graph import _select_decode_cuda_graph_batch_size import sparsevllm.platforms as platforms +from sparsevllm.configs.cuda_graph import _select_decode_cuda_graph_batch_size +from sparsevllm.engine.decode_graph_contract import ( + DecodeGraphContract, + DecodeGraphInputs, + DecodeGraphState, +) +from sparsevllm.engine.sequence import Sequence from sparsevllm.utils.context import get_context, set_context from sparsevllm.utils.profiler import profiler @@ -51,22 +56,25 @@ class DecodeCudaGraphKey: context_capacity: int is_long_text: bool capture_sampling: bool + graph_path_id: str = "" + shape_policy: str = "batch_only" @dataclass class DecodeCudaGraphState: key: DecodeCudaGraphKey + capture_context_capacity: int = 0 + decode_state: DecodeGraphState | None = None graph: torch.cuda.CUDAGraph | None = None - input_ids: torch.Tensor | None = None - positions: torch.Tensor | None = None - slot_mapping: torch.Tensor | None = None - context_lens: torch.Tensor | None = None - req_indices: torch.Tensor | None = None logits: torch.Tensor | None = None token_ids: torch.Tensor | None = None keepalive: list[object] = field(default_factory=list) sparse_state_refs: dict[int, dict[str, object]] = field(default_factory=dict) + def __post_init__(self) -> None: + if self.capture_context_capacity <= 0 and self.key.context_capacity > 0: + self.capture_context_capacity = int(self.key.context_capacity) + class DecodeCudaGraphRunner: """Fixed-shape decode runner, optionally backed by CUDA Graph replay. @@ -90,6 +98,7 @@ def __init__( method: str, capture_sizes: list[int], context_sizes: list[int] | tuple[int, ...] | str | int | None = None, + shape_policy: str = "batch_only", graph_pool=None, ): self.runtime_state = runtime_state @@ -104,6 +113,9 @@ def __init__( if not self.capture_sizes or any(size <= 0 for size in self.capture_sizes): raise ValueError(f"decode_graph capture_sizes must be positive, got {capture_sizes}.") self.context_sizes = _normalize_context_buckets(context_sizes) + self.shape_policy = str(shape_policy).strip().lower() + if self.shape_policy not in {"bucketed", "batch_only"}: + raise ValueError(f"Unsupported decode graph shape policy {self.shape_policy!r}.") self.max_context_len_override: int | None = None self._graphs: OrderedDict[DecodeCudaGraphKey, DecodeCudaGraphState] = OrderedDict() self.max_cached_graphs = self._resolve_max_cached_graphs() @@ -114,6 +126,11 @@ def __init__( self.replay_count = 0 self.eager_static_count = 0 self.force_eager_count = 0 + self.eviction_count = 0 + self.recapture_count = 0 + self._captured_keys: set[DecodeCudaGraphKey] = set() + self.reuse_larger_context_graphs = False + self.startup_plan_sealed = False def _resolve_max_cached_graphs(self) -> int | None: resolver = getattr(self.cache_manager, "decode_graph_max_cached_graphs", None) @@ -130,6 +147,13 @@ def _resolve_max_cached_graphs(self) -> int | None: def set_max_context_len_override(self, max_context_len: int | None): self.max_context_len_override = None if max_context_len is None else int(max_context_len) + def set_reuse_larger_context_graphs(self, enabled: bool): + self.reuse_larger_context_graphs = bool(enabled) + + def seal_startup_plan(self): + if self.shape_policy == "batch_only": + self.startup_plan_sealed = True + def clear_captured_graphs(self): for state in list(self._graphs.values()): self._release_graph_state(state) @@ -140,11 +164,9 @@ def clear_captured_graphs(self): @staticmethod def _release_graph_state(state: DecodeCudaGraphState): state.graph = None - state.input_ids = None - state.positions = None - state.slot_mapping = None - state.context_lens = None - state.req_indices = None + if state.decode_state is not None: + state.decode_state.close() + state.decode_state = None state.logits = None state.token_ids = None state.keepalive.clear() @@ -165,6 +187,7 @@ def _evict_cached_graphs(self, protected_key: DecodeCudaGraphKey): continue state = self._graphs.pop(key) self._release_graph_state(state) + self.eviction_count = int(getattr(self, "eviction_count", 0)) + 1 break else: break @@ -222,43 +245,86 @@ def _select_state( context_capacity: int, is_long_text: bool, capture_sampling: bool, + graph_path_id: str = "", allow_larger_context_capacity: bool = True, ) -> DecodeCudaGraphState: + shape_policy = self.shape_policy + graph_path_id = str(graph_path_id) or ( + "dense" if not method else ("long" if is_long_text else "short") + ) candidates = [ state for key, state in self._graphs.items() if key.method == method and key.batch_size == batch_size - and key.is_long_text == is_long_text and key.capture_sampling == capture_sampling + and key.graph_path_id == graph_path_id + and key.shape_policy == shape_policy + and (shape_policy == "batch_only" or key.is_long_text == is_long_text) and ( - key.context_capacity == context_capacity + shape_policy == "batch_only" + or key.context_capacity == context_capacity or (allow_larger_context_capacity and key.context_capacity >= context_capacity) ) ] if candidates: - state = min(candidates, key=lambda state: state.key.context_capacity) + state = min(candidates, key=lambda state: state.capture_context_capacity) + if ( + shape_policy == "batch_only" + and context_capacity > state.capture_context_capacity + ): + raise RuntimeError( + "batch-only decode CUDA Graph request exceeded captured path " + f"capacity: requested={context_capacity}, " + f"captured={state.capture_context_capacity}." + ) self._touch_graph_state(state.key) return state + if shape_policy == "batch_only" and getattr(self, "startup_plan_sealed", False): + raise RuntimeError( + "batch-only decode CUDA Graph has no startup-captured graph for " + f"batch_size={batch_size}, path={graph_path_id!r}." + ) + key = DecodeCudaGraphKey( method=method, batch_size=batch_size, - context_capacity=context_capacity, + context_capacity=0 if shape_policy == "batch_only" else context_capacity, is_long_text=bool(is_long_text), capture_sampling=capture_sampling, + graph_path_id=graph_path_id, + shape_policy=shape_policy, + ) + state = DecodeCudaGraphState( + key=key, + capture_context_capacity=int(context_capacity), ) - state = DecodeCudaGraphState(key=key) device = getattr( self.cache_manager, "device", torch.device("cuda" if torch.cuda.is_available() else "cpu"), ) - state.input_ids = torch.empty((batch_size,), dtype=torch.int64, device=device) - state.positions = torch.empty((batch_size,), dtype=torch.int64, device=device) - state.slot_mapping = torch.empty((batch_size,), dtype=torch.int32, device=device) - state.context_lens = torch.empty((batch_size,), dtype=torch.int32, device=device) - state.req_indices = torch.empty((batch_size,), dtype=torch.int32, device=device) + contract = DecodeGraphContract( + method=str(method), + shape_policy=shape_policy, + topology_path_id=graph_path_id, + batch_capacity=int(batch_size), + context_capacity=int(context_capacity), + capture_sampling=bool(capture_sampling), + ) + platform = getattr(self, "platform", None) + pin_memory = bool( + device.type != "cpu" + and platform is not None + and platform.supports_pin_memory() + ) + inputs = DecodeGraphInputs.allocate( + contract, + device=device, + pin_memory=pin_memory, + ) + state.decode_state = DecodeGraphState(contract=contract, inputs=inputs) self._graphs[key] = state self._evict_cached_graphs(key) return state @@ -269,24 +335,26 @@ def _prepare_static_step( seqs: list[Sequence], is_long_text: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - prepare_decode_static = getattr(self.runtime_state, "prepare_decode_static", None) - if prepare_decode_static is None: - raise TypeError("decode_graph requires runtime_state.prepare_decode_static().") - - assert state.input_ids is not None - assert state.positions is not None - assert state.slot_mapping is not None - assert state.context_lens is not None - assert state.req_indices is not None - - self.cache_manager.set_decode_static_max_context_len(int(state.key.context_capacity)) - input_ids, positions, _ = prepare_decode_static( + prepare_decode_graph_step = getattr( + self.runtime_state, + "prepare_decode_graph_step", + None, + ) + if prepare_decode_graph_step is None: + raise TypeError( + "decode_graph requires runtime_state.prepare_decode_graph_step()." + ) + graph_state = state.decode_state + if graph_state is None: + raise RuntimeError("Decode graph state was released before preparation.") + graph_state.inputs.validate(graph_state.contract) + + self.cache_manager.set_decode_static_max_context_len( + int(state.capture_context_capacity) + ) + input_ids, positions, _ = prepare_decode_graph_step( seqs, - state.input_ids, - state.positions, - state.slot_mapping, - state.context_lens, - state.req_indices, + graph_state, ) set_context( @@ -297,7 +365,9 @@ def _prepare_static_step( seqs=seqs, recurrent_state_manager=self.recurrent_state_manager, ) - self.cache_manager.set_decode_static_max_context_len(int(state.key.context_capacity)) + self.cache_manager.set_decode_static_max_context_len( + int(state.capture_context_capacity) + ) return input_ids, positions @@ -318,16 +388,21 @@ def _graph_context_capacity_policy(self, seqs: list[Sequence]) -> tuple[int, boo or "current" ).strip().lower() if policy in {"requested", "request", "final"}: - return self._requested_context_capacity(seqs), False + return self._requested_context_capacity(seqs), bool( + getattr(self, "reuse_larger_context_graphs", False) + ) if policy not in {"current", "cur", "now"}: raise ValueError( "decode_graph_context_policy must be 'current' or 'requested', " f"got {policy!r}." ) - return self._current_context_capacity(seqs), False + return self._current_context_capacity(seqs), bool( + getattr(self, "reuse_larger_context_graphs", False) + ) def bucket_plan(self) -> dict[str, object]: return { + "shape_policy": self.shape_policy, "batch_sizes": list(self.capture_sizes), "context_sizes": list(self.context_sizes), "context_policy": str( @@ -335,8 +410,53 @@ def bucket_plan(self) -> dict[str, object]: or "current" ), "max_cached_graphs": self.max_cached_graphs, + "cached_graph_keys": [ + { + "method": key.method, + "batch_size": key.batch_size, + "context_capacity": key.context_capacity, + "capture_context_capacity": state.capture_context_capacity, + "is_long_text": key.is_long_text, + "graph_path_id": key.graph_path_id, + "capture_sampling": key.capture_sampling, + } + for key, state in self._graphs.items() + if state.graph is not None + ], } + def _graph_path_id(self, is_long_text: bool) -> str: + resolver = getattr(self.cache_manager, "decode_graph_path_id", None) + if callable(resolver): + return str(resolver(bool(is_long_text))) + return "dense" if not self.method else ("long" if is_long_text else "short") + + def _batch_only_context_capacity( + self, seqs: list[Sequence], *, is_long_text: bool + ) -> int: + if self.max_context_len_override is not None: + capacity = int(self.max_context_len_override) + else: + resolver = getattr( + self.cache_manager, + "decode_graph_batch_only_capacity", + None, + ) + if not callable(resolver): + raise TypeError( + "batch-only decode CUDA Graph requires a context-stable " + "capacity resolver." + ) + capacity = int(resolver(bool(is_long_text))) + validator = getattr( + self.cache_manager, + "validate_decode_graph_batch_only_capacity", + None, + ) + if callable(validator): + validator(seqs, capacity=capacity, is_long_text=bool(is_long_text)) + return capacity + def _cache_manager_graph_context_capacity(self, seqs: list[Sequence]) -> tuple[int, bool] | None: resolver = getattr(self.cache_manager, "decode_graph_context_capacity", None) if resolver is None: @@ -409,11 +529,17 @@ def _capture( ) -> DecodeCudaGraphState: if not self.platform.supports_graph_capture(): raise RuntimeError(f"Platform {self.platform.name!r} does not support decode CUDA graph capture.") + graph_state = state.decode_state + if graph_state is None: + raise RuntimeError("Decode graph state was released before capture.") ctx = get_context() ctx.sparse_controller = self.sparse_controller with profiler.record("decode_graph_warmup"): self.sparse_controller.prepare_forward(seqs, is_prefill=False) + participant = graph_state.runtime_state + if participant is not None: + participant.prepare_in_graph() logits = self.run_model(input_ids, positions, is_prefill=False) if state.key.capture_sampling: if logits is None: @@ -434,6 +560,9 @@ def _capture( graph = torch.cuda.CUDAGraph() try: with torch.cuda.graph(graph, pool=self.graph_pool): + participant = graph_state.runtime_state + if participant is not None: + participant.prepare_in_graph() self._reset_graph_input_attn_scores(graph_input_sparse_state_refs) logits = self.run_model(input_ids, positions, is_prefill=False) if state.key.capture_sampling: @@ -455,12 +584,8 @@ def _capture( logits, ctx.decode_mid_o, ctx.decode_mid_o_logexpsum, - state.input_ids, - state.positions, - state.slot_mapping, - state.context_lens, - state.req_indices, ] + keepalive.extend(graph_state.keepalive_tensors()) if token_ids is not None: keepalive.append(token_ids) for sparse_refs_by_layer in (graph_input_sparse_state_refs, state.sparse_state_refs): @@ -468,11 +593,18 @@ def _capture( for value in refs.values(): if isinstance(value, torch.Tensor): keepalive.append(value) - keepalive.extend(self.cache_manager.decode_graph_keepalive_tensors()) sparse_keepalive = getattr(self.sparse_controller, "decode_graph_keepalive_tensors", None) if sparse_keepalive is not None: keepalive.extend(sparse_keepalive()) state.keepalive = keepalive + captured_keys = getattr(self, "_captured_keys", None) + if captured_keys is None: + captured_keys = set() + self._captured_keys = captured_keys + if state.key in captured_keys: + self.recapture_count = int(getattr(self, "recapture_count", 0)) + 1 + else: + captured_keys.add(state.key) self.capture_count += 1 return state @@ -503,13 +635,21 @@ def run( graph_batch_size = self._select_graph_batch_size(real_batch_size) is_long_text = self.is_long_text_batch(seqs, False) - context_capacity, allow_larger_context_capacity = self._graph_context_capacity_policy(seqs) + graph_path_id = self._graph_path_id(is_long_text) + if self.shape_policy == "batch_only": + context_capacity = self._batch_only_context_capacity( + seqs, is_long_text=is_long_text + ) + allow_larger_context_capacity = False + else: + context_capacity, allow_larger_context_capacity = self._graph_context_capacity_policy(seqs) state = self._select_state( method=self.method, batch_size=graph_batch_size, context_capacity=context_capacity, is_long_text=is_long_text, capture_sampling=bool(capture_sampling), + graph_path_id=graph_path_id, allow_larger_context_capacity=allow_larger_context_capacity, ) self.last_state_key = state.key @@ -543,23 +683,37 @@ def run_eager_static(self, seqs: list[Sequence]) -> torch.Tensor | None: real_batch_size = len(seqs) graph_batch_size = self._select_graph_batch_size(real_batch_size) is_long_text = self.is_long_text_batch(seqs, False) - context_capacity, allow_larger_context_capacity = self._static_context_capacity_policy(seqs) + graph_path_id = self._graph_path_id(is_long_text) + if self.shape_policy == "batch_only": + context_capacity = self._batch_only_context_capacity( + seqs, is_long_text=is_long_text + ) + allow_larger_context_capacity = False + else: + context_capacity, allow_larger_context_capacity = self._static_context_capacity_policy(seqs) state = self._select_state( method=self.method, batch_size=graph_batch_size, context_capacity=context_capacity, is_long_text=is_long_text, capture_sampling=False, + graph_path_id=graph_path_id, allow_larger_context_capacity=allow_larger_context_capacity, ) self.last_state_key = state.key self.last_real_batch_size = real_batch_size input_ids, positions = self._prepare_static_step(state, seqs, is_long_text) + graph_state = state.decode_state + if graph_state is None: + raise RuntimeError("Decode graph state was released before static execution.") ctx = get_context() ctx.sparse_controller = self.sparse_controller with profiler.record("model_sparse_prepare"): self.sparse_controller.prepare_forward(seqs, is_prefill=False) + participant = graph_state.runtime_state + if participant is not None: + participant.prepare_in_graph() logits = self.run_model(input_ids, positions, is_prefill=False) if logits is None: return None diff --git a/src/sparsevllm/engine/decode_graph_contract.py b/src/sparsevllm/engine/decode_graph_contract.py new file mode 100644 index 00000000..765053db --- /dev/null +++ b/src/sparsevllm/engine/decode_graph_contract.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Protocol, runtime_checkable + +import torch + + +@dataclass(frozen=True) +class DecodeGraphPaddingContract: + """Safe values used for inactive rows in a fixed-capacity decode graph.""" + + write_slot: int = -1 + active: bool = False + mirror_first_real_row_for_reads: bool = True + + +@dataclass(frozen=True) +class DecodeGraphContract: + """Capture-time facts that define one decode graph family.""" + + method: str + shape_policy: str + topology_path_id: str + batch_capacity: int + context_capacity: int + capture_sampling: bool = False + dynamic_context_lens: bool = True + padding: DecodeGraphPaddingContract = field( + default_factory=DecodeGraphPaddingContract + ) + + def __post_init__(self) -> None: + if self.shape_policy not in {"bucketed", "batch_only"}: + raise ValueError( + f"Unsupported decode graph shape policy {self.shape_policy!r}." + ) + if self.batch_capacity <= 0 or self.context_capacity <= 0: + raise ValueError( + "Decode graph batch and context capacities must be positive, got " + f"batch={self.batch_capacity} context={self.context_capacity}." + ) + if not self.topology_path_id: + raise ValueError("Decode graph topology_path_id must be non-empty.") + if self.shape_policy == "batch_only" and not self.dynamic_context_lens: + raise ValueError( + "batch-only decode graphs require device-resident dynamic context lengths." + ) + + @property + def capability_level(self) -> str: + return ( + "strict" + if self.topology_path_id in {"dense", "unified"} + else "path_scoped" + ) + + +@dataclass +class DecodeGraphHostInputs: + """Persistent host mirrors for metadata copied before graph replay.""" + + input_ids: torch.Tensor + positions: torch.Tensor + context_lens: torch.Tensor + request_indices: torch.Tensor + active_mask: torch.Tensor + + def tensors(self) -> tuple[torch.Tensor, ...]: + return ( + self.input_ids, + self.positions, + self.context_lens, + self.request_indices, + self.active_mask, + ) + + +@dataclass +class DecodeGraphInputs: + """Typed, address-stable public inputs shared by decode participants.""" + + input_ids: torch.Tensor + positions: torch.Tensor + context_lens: torch.Tensor + request_indices: torch.Tensor + write_slot_mapping: torch.Tensor + active_mask: torch.Tensor + host: DecodeGraphHostInputs + + @classmethod + def allocate( + cls, + contract: DecodeGraphContract, + *, + device: torch.device, + pin_memory: bool, + ) -> DecodeGraphInputs: + batch = int(contract.batch_capacity) + + def device_buffer(dtype: torch.dtype) -> torch.Tensor: + return torch.empty(batch, dtype=dtype, device=device) + + def host_buffer(dtype: torch.dtype) -> torch.Tensor: + if pin_memory: + return torch.empty( + batch, + dtype=dtype, + device="cpu", + pin_memory=True, + ) + return torch.empty(batch, dtype=dtype, device="cpu") + + inputs = cls( + input_ids=device_buffer(torch.int64), + positions=device_buffer(torch.int64), + context_lens=device_buffer(torch.int32), + request_indices=device_buffer(torch.int32), + write_slot_mapping=device_buffer(torch.int32), + active_mask=device_buffer(torch.bool), + host=DecodeGraphHostInputs( + input_ids=host_buffer(torch.int64), + positions=host_buffer(torch.int64), + context_lens=host_buffer(torch.int32), + request_indices=host_buffer(torch.int32), + active_mask=host_buffer(torch.bool), + ), + ) + inputs.validate(contract) + return inputs + + @property + def batch_capacity(self) -> int: + return int(self.input_ids.numel()) + + def device_tensors(self) -> tuple[torch.Tensor, ...]: + return ( + self.input_ids, + self.positions, + self.context_lens, + self.request_indices, + self.write_slot_mapping, + self.active_mask, + ) + + def keepalive_tensors(self) -> tuple[torch.Tensor, ...]: + return self.device_tensors() + self.host.tensors() + + def data_ptrs(self) -> tuple[int, ...]: + return tuple(int(tensor.data_ptr()) for tensor in self.device_tensors()) + + def validate(self, contract: DecodeGraphContract) -> None: + expected = int(contract.batch_capacity) + tensors = self.device_tensors() + if any(tensor.ndim != 1 or tensor.numel() != expected for tensor in tensors): + raise ValueError( + "Decode graph public inputs must be one-dimensional and match " + f"batch_capacity={expected}." + ) + device = self.input_ids.device + if any(tensor.device != device for tensor in tensors): + raise ValueError("Decode graph public inputs must share one device.") + expected_dtypes = ( + torch.int64, + torch.int64, + torch.int32, + torch.int32, + torch.int32, + torch.bool, + ) + actual_dtypes = tuple(tensor.dtype for tensor in tensors) + if actual_dtypes != expected_dtypes: + raise TypeError( + "Decode graph public input dtypes do not match the contract: " + f"expected={expected_dtypes} actual={actual_dtypes}." + ) + host_tensors = self.host.tensors() + if any(tensor.device.type != "cpu" for tensor in host_tensors): + raise ValueError("Decode graph host mirrors must reside on CPU.") + if tuple(tensor.dtype for tensor in host_tensors) != ( + torch.int64, + torch.int64, + torch.int32, + torch.int32, + torch.bool, + ): + raise TypeError("Decode graph host mirror dtypes do not match public inputs.") + if any(tensor.ndim != 1 or tensor.numel() != expected for tensor in host_tensors): + raise ValueError( + "Decode graph host mirrors must match the graph batch capacity." + ) + + +@dataclass +class CacheDecodeGraphState: + """Per-graph cache participant state owned by one cache manager.""" + + contract: DecodeGraphContract + inputs: DecodeGraphInputs + + +@dataclass +class DecodeGraphState: + """Typed public graph state plus participant-owned private state.""" + + contract: DecodeGraphContract + inputs: DecodeGraphInputs + runtime_state: object | None = None + + def keepalive_tensors(self) -> list[torch.Tensor]: + tensors = list(self.inputs.keepalive_tensors()) + participant = self.runtime_state + keepalive = getattr(participant, "graph_keepalive_tensors", None) + if callable(keepalive): + tensors.extend(keepalive()) + return tensors + + def close(self) -> None: + participant = self.runtime_state + close = getattr(participant, "close", None) + if callable(close): + close() + self.runtime_state = None + + +@runtime_checkable +class DecodeGraphParticipant(Protocol): + """Minimal lifecycle implemented by graph metadata owners.""" + + def prepare_out_graph(self, seqs: list[object]) -> None: ... + + def prepare_in_graph(self) -> None: ... + + def graph_keepalive_tensors(self) -> Iterable[torch.Tensor]: ... diff --git a/src/sparsevllm/engine/llm_engine.py b/src/sparsevllm/engine/llm_engine.py index bcc9f964..60f6b678 100644 --- a/src/sparsevllm/engine/llm_engine.py +++ b/src/sparsevllm/engine/llm_engine.py @@ -14,6 +14,9 @@ from sparsevllm.utils.log import logger import sys +from sparsevllm.configs.cuda_graph import ( + build_decode_cuda_graph_startup_family_plan, +) from sparsevllm.config import Config from sparsevllm.sampling_params import SamplingParams @@ -352,12 +355,16 @@ def _warmup(self): graph_sized_batch = warmup_profile in ("graph", "big_prefill_only") decode_warmup = warmup_profile in ("graph", "decode_1seq") num_seqs = int(self.config.max_decoding_seqs) if graph_sized_batch else 1 + startup_capture = bool( + getattr(self.config, "decode_graph_startup_capture", False) + ) - # 预热 1 个 Token 的生成（包含 Prefill 和 Decode） + # Startup precapture owns decode warmup when enabled. Keep this first + # pass prefill-only so it cannot create unplanned short/long graph keys. sampling_params = SamplingParams( - max_tokens=2 if decode_warmup else 1, + max_tokens=2 if decode_warmup and not startup_capture else 1, temperature=0.0, - ignore_eos=decode_warmup, + ignore_eos=decode_warmup and not startup_capture, ) max_prompt_len = max(1, int(self.config.max_model_len) - int(sampling_params.max_tokens)) warmup_len = min(int(self.config.engine_prefill_chunk_size), max_prompt_len) @@ -394,9 +401,23 @@ def _warmup(self): max_warmup_len, ) warmup_len = max_warmup_len + startup_plan = ( + build_decode_cuda_graph_startup_family_plan(self.config) + if startup_capture + else [] + ) + capture_groups: dict[tuple[int, bool], list[int]] = {} + for batch_size, context_capacity, is_long_text in startup_plan: + capture_groups.setdefault((batch_size, is_long_text), []).append( + context_capacity + ) + num_warmup_rounds = 2 if warmup_profile == "graph" else 1 vocab_size = int(self.config.hf_config.vocab_size) - num_dummy_prompts = num_seqs * num_warmup_rounds + num_dummy_prompts = ( + num_seqs * num_warmup_rounds + + sum(batch_size for batch_size, _ in capture_groups) + ) if num_dummy_prompts > vocab_size: raise ValueError( "Warmup requires one distinct leading token per dummy prompt: " @@ -408,25 +429,150 @@ def _warmup(self): f"ignore_eos={sampling_params.ignore_eos})." ) - def run_warmup(params: SamplingParams, prompt_offset: int) -> None: - for request_idx in range(num_seqs): + def run_warmup( + params: SamplingParams, + prompt_offset: int, + *, + batch_size: int = num_seqs, + first_prompt_len: int = warmup_len, + ) -> int: + for request_idx in range(batch_size): # Distinct leading tokens prevent prefix-cache reuse within or # across warmup rounds. - prompt_len = warmup_len if request_idx == 0 else 1 + prompt_len = first_prompt_len if request_idx == 0 else 1 dummy_prompt = [prompt_offset + request_idx] + [0] * (prompt_len - 1) self.add_request(dummy_prompt, params) while not self.is_finished(): self.step() + return prompt_offset + batch_size + + def prepare_capture_batch( + params: SamplingParams, + prompt_offset: int, + *, + batch_size: int, + prompt_len: int, + ) -> tuple[list[Sequence], int]: + seq_ids = [] + for request_idx in range(batch_size): + dummy_prompt = [prompt_offset + request_idx] + [0] * (prompt_len - 1) + seq_ids.append(self.add_request(dummy_prompt, params)) - run_warmup(sampling_params, prompt_offset=0) + parked: list[Sequence] = [] + while self.scheduler.waiting: + self.step() + while self.scheduler.decoding: + parked.append(self.scheduler.decoding.popleft()) + while self.scheduler.decoding: + parked.append(self.scheduler.decoding.popleft()) + if len(parked) != batch_size: + raise RuntimeError( + "Startup decode CUDA Graph prefill did not park the requested " + f"batch: expected={batch_size}, actual={len(parked)}." + ) + if {int(seq.seq_id) for seq in parked} != set(seq_ids): + raise RuntimeError("Startup decode CUDA Graph prefill parked unexpected sequences.") + return parked, prompt_offset + batch_size + + prompt_offset = run_warmup(sampling_params, prompt_offset=0) + + if startup_plan: + short_graphs = sum(not is_long for _, _, is_long in startup_plan) + long_graphs = len(startup_plan) - short_graphs + logger.info( + "Startup decode CUDA Graph capture: {} coarse graphs " + "(limit={}, short={}, long={}, plan={}).", + len(startup_plan), + self.config.decode_graph_max_cached_graphs, + short_graphs, + long_graphs, + startup_plan, + ) + capture_params = SamplingParams( + max_tokens=2, + temperature=0.0, + ignore_eos=True, + ) + threshold = self.scheduler._long_text_threshold(is_prefill=False) + for (batch_size, is_long_text), context_capacities in capture_groups.items(): + prompt_len = int(threshold) if is_long_text else 1 + parked, prompt_offset = prepare_capture_batch( + capture_params, + prompt_offset, + batch_size=batch_size, + prompt_len=prompt_len, + ) + try: + observed_long = self.scheduler._is_long_text( + parked[0], + is_prefill=False, + ) + if bool(observed_long) != bool(is_long_text): + raise RuntimeError( + "Startup decode CUDA Graph family prefill crossed the " + "wrong long-text boundary: " + f"expected={is_long_text}, observed={observed_long}, " + f"threshold={threshold}, num_tokens={parked[0].num_tokens}." + ) + for context_capacity in context_capacities: + self.model_runner.call( + "set_decode_cuda_graph_max_context_len_override", + context_capacity, + ) + self.model_runner.call( + "capture_decode_cuda_graph_warmup", + parked, + ) + finally: + self.model_runner.call( + "set_decode_cuda_graph_max_context_len_override", + None, + ) + self.scheduler.decoding.extend(parked) + for seq in parked: + self.abort_request(int(seq.seq_id)) + self.model_runner.call( + "set_decode_cuda_graph_reuse_larger_context_graphs", + True, + ) + graph_runner = self.model_runner.decode_graph_runner + captured = { + ( + int(key.batch_size), + int( + state.capture_context_capacity + if key.shape_policy == "batch_only" + else key.context_capacity + ), + bool(key.is_long_text), + ) + for key, state in graph_runner._graphs.items() + if state.graph is not None + and key.method == str(self.config.sparse_method or "") + and not key.capture_sampling + } + missing = sorted(set(startup_plan) - captured) + if missing: + raise RuntimeError( + "Startup decode CUDA Graph capture did not materialize its plan: " + f"missing={missing}." + ) + self.model_runner.call("seal_decode_cuda_graph_startup_plan") + logger.info( + "Startup decode CUDA Graph capture finished: cached={} " + "capture_count={} replay_count={}.", + len(captured), + graph_runner.capture_count, + graph_runner.replay_count, + ) if warmup_profile == "graph": # CUDA Graph capture establishes its private allocator pool. Warm # prefill once more against the final allocator layout. logger.info(f"Post-capture prefill warmup (num_seqs={num_seqs}).") - run_warmup( + prompt_offset = run_warmup( SamplingParams(max_tokens=1, temperature=0.0), - prompt_offset=num_seqs, + prompt_offset=prompt_offset, ) self._warmup_moe_workspaces() @@ -1011,11 +1157,14 @@ def worker_info( "deltakv_latent_quant_bits", "deltakv_latent_quant_group_size", "decode_graph", + "decode_graph_shape_policy", "decode_graph_capture_sampling", "decode_graph_capture_sizes", "decode_graph_context_sizes", "decode_graph_context_policy", "decode_graph_max_cached_graphs", + "decode_graph_startup_capture", + "decode_graph_startup_capture_limit", "enable_prefix_caching", "prefix_cache_mode", "resolved_prefix_cache_mode", diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index a83243ce..fdbf3727 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -20,7 +20,12 @@ from sparsevllm.models.qwen2 import Qwen2ForCausalLM from sparsevllm.models.llama import LlamaForCausalLM from sparsevllm.layers.sampler import Sampler +from sparsevllm.method_registry import decode_sparse_long_text_threshold from sparsevllm.operators import registry as operator_registry +from sparsevllm.operators.decode_attention import ( + collect_decode_graph_participants, + validate_batch_only_decode_graph_model, +) from sparsevllm.utils.context import set_context, get_context, reset_context from sparsevllm.utils.loader import load_model, sync_deltakv_config_from_checkpoint @@ -246,6 +251,17 @@ def __init__( "decode_graph", bool(getattr(config, "decode_graph", False)), ) + setattr( + hf_config, + "decode_graph_shape_policy", + str( + getattr( + config, + "decode_graph_shape_policy", + "batch_only", + ) + ), + ) decode_static_capture_sizes = _resolve_decode_cuda_graph_capture_sizes( config.decode_graph_capture_sizes, config.max_decoding_seqs, @@ -262,6 +278,11 @@ def __init__( max_decoding_seqs=config.max_decoding_seqs, ), ) + if ( + self.config.decode_graph + and self.config.decode_graph_shape_policy == "batch_only" + ): + validate_batch_only_decode_graph_model(self.model) if config.tiny_random: from sparsevllm.debug.tiny_random import initialize_sparse_model @@ -269,6 +290,7 @@ def __init__( self.model, hf_config, seed=config.tiny_random_seed, + quantized=config.quantization_config.enabled, ) else: load_model( @@ -349,6 +371,7 @@ def __init__( self.recurrent_state_manager, self.prefix_cache_coordinator, self.chain_cache_coordinator, + decode_graph_participants=collect_decode_graph_participants(self.model), ) # 初始化稀疏控制器 @@ -384,6 +407,7 @@ def __init__( method=self.config.sparse_method, capture_sizes=decode_static_capture_sizes, context_sizes=decode_static_context_sizes, + shape_policy=self.config.decode_graph_shape_policy, graph_pool=self.cuda_graph_pool, ) torch.set_default_device("cpu") @@ -923,9 +947,20 @@ def parallel_group_summary(group) -> dict[str, object] | None: "force_eager_count": int( getattr(graph_runner, "force_eager_count", 0) ), + "eviction_count": int( + getattr(graph_runner, "eviction_count", 0) + ), + "recapture_count": int( + getattr(graph_runner, "recapture_count", 0) + ), "cached_graph_count": len( getattr(graph_runner, "_graphs", {}) ), + "bucket_plan": ( + graph_runner.bucket_plan() + if callable(getattr(graph_runner, "bucket_plan", None)) + else None + ), "last_state_key": ( { "method": str(graph_key.method or ""), @@ -1152,15 +1187,12 @@ def debug_sparse_state_summaries(self) -> list[dict[str, object]] | None: def _long_text_threshold(self, is_prefill: bool) -> int: del is_prefill - if self.config.sparse_method in ("streamingllm", "attention-sink", "attention_sink"): - base = self.config.sink_keep_tokens + self.config.recent_keep_tokens - else: - base = ( - self.config.sink_keep_tokens - + self.config.recent_keep_tokens - + self.config.decode_keep_tokens - ) - return base + return decode_sparse_long_text_threshold( + self.config.sparse_method, + num_sink_tokens=self.config.sink_keep_tokens, + decode_keep_tokens=self.config.decode_keep_tokens, + num_recent_tokens=self.config.recent_keep_tokens, + ) def _is_long_text_batch(self, seqs: list[Sequence], is_prefill: bool) -> bool: # Prefill execution is per-sequence and cache-manager owned. This @@ -1220,14 +1252,12 @@ def prepare_sample(self, seqs: list[Sequence]): def _auto_capture_greedy_sampling(self, seqs: list[Sequence]) -> bool: if any(self._has_sampling_penalty(seq) for seq in seqs): return False - if self.config.decode_graph_capture_sampling: - return all(bool(getattr(seq, "should_publish_sample", True)) for seq in seqs) + if not self.config.decode_graph_capture_sampling: + return False if self.config.tensor_parallel_size != 1: return False if self.config.enable_prefix_caching: return False - if str(self.config.sparse_method or "") not in {"", "omnikv"}: - return False return all( bool(getattr(seq, "should_publish_sample", True)) and seq.temperature <= 1e-10 @@ -1346,6 +1376,19 @@ def _mask_recompute_logprobs( def set_decode_cuda_graph_max_context_len_override(self, max_context_len: int | None): self.decode_graph_runner.set_max_context_len_override(max_context_len) + def set_decode_cuda_graph_reuse_larger_context_graphs(self, enabled: bool): + self.decode_graph_runner.set_reuse_larger_context_graphs(enabled) + + def seal_decode_cuda_graph_startup_plan(self): + self.decode_graph_runner.seal_startup_plan() + + def capture_decode_cuda_graph_warmup(self, seqs: list[Sequence]) -> None: + """Capture one planned graph without advancing scheduler sequence state.""" + try: + self.decode_graph_runner.run(seqs, capture_sampling=False) + finally: + reset_context() + def set_omnikv_decode_graph_max_context_len_override(self, max_context_len: int | None): self.set_decode_cuda_graph_max_context_len_override(max_context_len) diff --git a/src/sparsevllm/engine/runtime_state.py b/src/sparsevllm/engine/runtime_state.py index 09623e37..7cccc11e 100644 --- a/src/sparsevllm/engine/runtime_state.py +++ b/src/sparsevllm/engine/runtime_state.py @@ -2,18 +2,23 @@ from collections import deque from contextlib import nullcontext +from dataclasses import dataclass from typing import ContextManager from typing import Protocol import torch from sparsevllm.config import Config -from sparsevllm.engine.prefix_cache_coordinator import PrefixCacheCoordinator from sparsevllm.engine.chain_cache import ( ChainAdmissionPlan, ChainCacheCoordinator, ChainOwnerMismatchError, ) +from sparsevllm.engine.decode_graph_contract import ( + CacheDecodeGraphState, + DecodeGraphState, +) +from sparsevllm.engine.prefix_cache_coordinator import PrefixCacheCoordinator from sparsevllm.engine.recurrent_state_manager import RecurrentStateManager from sparsevllm.engine.sequence import Sequence @@ -54,6 +59,45 @@ def free_slot_stats(self) -> dict[str, int]: ... def debug_live_seq_slots(self) -> dict[int, int]: ... +@dataclass +class RuntimeDecodeGraphState: + """Per-graph runtime participant that delegates to semantic owners.""" + + owner: RuntimeState + cache: CacheDecodeGraphState + operator_states: tuple[tuple[object, object], ...] = () + + def prepare_out_graph(self, seqs: list[Sequence]) -> None: + self.owner._evict_mixed_prefix_for_step(seqs, is_prefill=False) + self.owner.cache_manager.prepare_decode_graph_step(seqs, self.cache) + for participant, state in self.operator_states: + participant.prepare_decode_graph_out(state) + if self.owner.recurrent_state_manager is not None: + inputs = self.cache.inputs + self.owner.recurrent_state_manager.prepare_decode_static( + seqs, + token_batch=inputs.batch_capacity, + device=inputs.input_ids.device, + ) + + def prepare_in_graph(self) -> None: + self.owner.cache_manager.prepare_decode_graph_in(self.cache) + for participant, state in self.operator_states: + participant.prepare_decode_graph_in(state) + + def graph_keepalive_tensors(self) -> list[torch.Tensor]: + tensors = self.owner.cache_manager.decode_graph_state_keepalive_tensors( + self.cache + ) + for participant, state in self.operator_states: + tensors.extend(participant.decode_graph_keepalive_tensors(state)) + return tensors + + def close(self) -> None: + for participant, state in reversed(self.operator_states): + participant.close_decode_graph_state(state) + + class RuntimeState: """Single lifecycle entrypoint for KV, recurrent state, and mixed prefix cache.""" @@ -64,12 +108,14 @@ def __init__( recurrent_state_manager: RecurrentStateManager | None = None, prefix_cache_coordinator: PrefixCacheCoordinator | None = None, chain_cache_coordinator: ChainCacheCoordinator | None = None, + decode_graph_participants: tuple[object, ...] = (), ): self.config = config self.cache_manager = cache_manager self.recurrent_state_manager = recurrent_state_manager self.prefix_cache_coordinator = prefix_cache_coordinator self.chain_cache_coordinator = chain_cache_coordinator + self.decode_graph_participants = tuple(decode_graph_participants) self._resident_seq_ids: set[int] = set() @property @@ -141,6 +187,56 @@ def prepare_decode_static(self, seqs: list[Sequence], *args): ) return result + def init_decode_graph_state( + self, + graph_state: DecodeGraphState, + ) -> RuntimeDecodeGraphState: + if graph_state.runtime_state is not None: + raise RuntimeError("Decode graph runtime state was initialized twice.") + cache_state = self.cache_manager.init_decode_graph_state( + graph_state.contract, + graph_state.inputs, + ) + operator_states: list[tuple[object, object]] = [] + try: + for participant in self.decode_graph_participants: + operator_states.append( + ( + participant, + participant.init_decode_graph_state( + graph_state.contract, + graph_state.inputs, + ), + ) + ) + except BaseException: + for participant, state in reversed(operator_states): + participant.close_decode_graph_state(state) + raise + state = RuntimeDecodeGraphState( + owner=self, + cache=cache_state, + operator_states=tuple(operator_states), + ) + graph_state.runtime_state = state + return state + + def prepare_decode_graph_step( + self, + seqs: list[Sequence], + graph_state: DecodeGraphState, + ): + participant = graph_state.runtime_state + if participant is None: + participant = self.init_decode_graph_state(graph_state) + if not isinstance(participant, RuntimeDecodeGraphState): + raise TypeError( + "Decode graph runtime participant has an unexpected type: " + f"{type(participant).__name__}." + ) + participant.prepare_out_graph(seqs) + return graph_state.inputs.input_ids, graph_state.inputs.positions, None + def on_forward_end(self, seqs: list[Sequence], is_prefill: bool) -> None: self.cache_manager.on_forward_end(seqs, is_prefill) if self.recurrent_state_manager is not None: diff --git a/src/sparsevllm/engine/scheduler.py b/src/sparsevllm/engine/scheduler.py index 0a22d2a9..aece3374 100644 --- a/src/sparsevllm/engine/scheduler.py +++ b/src/sparsevllm/engine/scheduler.py @@ -11,6 +11,7 @@ ) from sparsevllm.engine.sequence import Sequence, SequenceStatus from sparsevllm.engine.runtime_state import MemoryOracle +from sparsevllm.method_registry import decode_sparse_long_text_threshold from sparsevllm.sampling_params import resolve_eos_token_ids from sparsevllm.utils.log import logger @@ -67,11 +68,13 @@ def __init__( def _long_text_threshold(self, is_prefill: bool) -> int: """Long-text boundary retained only for decode batch partitioning.""" - if self.config.sparse_method in ("streamingllm", "attention-sink", "attention_sink"): - base = self.sink_keep_tokens + self.recent_keep_tokens - else: - base = self.sink_keep_tokens + self.decode_keep_tokens + self.recent_keep_tokens - return base + del is_prefill + return decode_sparse_long_text_threshold( + self.config.sparse_method, + num_sink_tokens=self.sink_keep_tokens, + decode_keep_tokens=self.decode_keep_tokens, + num_recent_tokens=self.recent_keep_tokens, + ) def _is_long_text(self, seq: Sequence, is_prefill: bool) -> bool: if not self.config.sparse_method: diff --git a/src/sparsevllm/kernels/external/flashinfer/decode.py b/src/sparsevllm/kernels/external/flashinfer/decode.py index 083f6150..79f6292b 100644 --- a/src/sparsevllm/kernels/external/flashinfer/decode.py +++ b/src/sparsevllm/kernels/external/flashinfer/decode.py @@ -67,7 +67,17 @@ def _paged_decode_wrapper_type(): ) _require_parameters( wrapper_type, - frozenset({"float_workspace_buffer", "kv_layout", "backend"}), + frozenset( + { + "float_workspace_buffer", + "kv_layout", + "use_cuda_graph", + "paged_kv_indptr_buffer", + "paged_kv_indices_buffer", + "paged_kv_last_page_len_buffer", + "backend", + } + ), feature=feature, entrypoint="BatchDecodeWithPagedKVCacheWrapper", ) @@ -110,11 +120,22 @@ def flashinfer_paged_decode_support() -> tuple[bool, str]: return True, reason -def make_flashinfer_paged_decode_wrapper(workspace: torch.Tensor): +def make_flashinfer_paged_decode_wrapper( + workspace: torch.Tensor, + *, + use_cuda_graph: bool = False, + paged_kv_indptr_buffer: torch.Tensor | None = None, + paged_kv_indices_buffer: torch.Tensor | None = None, + paged_kv_last_page_len_buffer: torch.Tensor | None = None, +): wrapper_type, _ = _paged_decode_wrapper_type() return wrapper_type( workspace, kv_layout="NHD", + use_cuda_graph=use_cuda_graph, + paged_kv_indptr_buffer=paged_kv_indptr_buffer, + paged_kv_indices_buffer=paged_kv_indices_buffer, + paged_kv_last_page_len_buffer=paged_kv_last_page_len_buffer, backend="auto", ) diff --git a/src/sparsevllm/kernels/external/sgl/support.py b/src/sparsevllm/kernels/external/sgl/support.py index 2f8c66e6..c392cde9 100644 --- a/src/sparsevllm/kernels/external/sgl/support.py +++ b/src/sparsevllm/kernels/external/sgl/support.py @@ -3,7 +3,6 @@ import importlib import importlib.metadata import importlib.util -import re from sparsevllm.kernels.external.support import ( ExternalKernelFamilyError, @@ -11,9 +10,8 @@ KernelFamilyState, ) -_MIN_VERSION = (0, 4, 5) -_MAX_VERSION = (0, 5, 0) _DISTRIBUTION = "sglang-kernel" +_REQUIRED_VERSION = "0.4.5" def sgl_kernel_health() -> KernelFamilyHealth: @@ -43,14 +41,12 @@ def sgl_kernel_health() -> KernelFamilyHealth: None, f"{_DISTRIBUTION} package metadata is unavailable", ) - match = re.match(r"^(\d+)\.(\d+)\.(\d+)", version) - parsed = tuple(map(int, match.groups())) if match else None - if parsed is None or not _MIN_VERSION <= parsed < _MAX_VERSION: + if version != _REQUIRED_VERSION: return KernelFamilyHealth( _DISTRIBUTION, KernelFamilyState.BROKEN, version, - f"requires {_DISTRIBUTION}>=0.4.5,<0.5, got {version}", + f"requires {_DISTRIBUTION}=={_REQUIRED_VERSION}, got {version}", ) try: importlib.import_module("sgl_kernel") diff --git a/src/sparsevllm/kernels/tilelang/mla/decode.py b/src/sparsevllm/kernels/tilelang/mla/decode.py index d2738cb2..8a48d974 100644 --- a/src/sparsevllm/kernels/tilelang/mla/decode.py +++ b/src/sparsevllm/kernels/tilelang/mla/decode.py @@ -102,11 +102,17 @@ def build_glm_mla_decode_kernel( VALID_BLOCK_H = min(block_H, kv_group_num) VALID_OUTPUT_HEADS = valid_output_heads HEAD_TILE_COUNT = h_q // VALID_BLOCK_H - SCORE_TILE_COUNT = HEAD_TILE_COUNT if score_mode == "partial" else 1 + SCORE_TILE_COUNT = ( + VALID_OUTPUT_HEADS + if score_mode == "per_head" + else HEAD_TILE_COUNT + if score_mode == "partial" + else 1 + ) assert h_kv == 1, "h_kv must be 1" assert h_q % VALID_BLOCK_H == 0, "h_q must use complete head tiles" assert 0 < VALID_OUTPUT_HEADS <= h_q, "valid output heads must fit h_q" - assert score_mode in ("direct", "atomic", "partial") + assert score_mode in ("direct", "atomic", "partial", "per_head") assert not need_score or score_mode != "direct" or HEAD_TILE_COUNT == 1 assert block_size >= block_N and block_size % block_N == 0, ( "block_size must be at least block_N and a multiple of block_N" @@ -190,15 +196,20 @@ def main_split( for i, j in T.Parallel(block_H, block_N): acc_s[i, j] = T.if_then_else(by * VALID_BLOCK_H + i >= VALID_OUTPUT_HEADS, -T.infinity(accum_dtype), acc_s[i, j]) if need_score: - T.reduce_max(acc_s, token_scores, dim=0) + if score_mode == "per_head": + for i, j in T.Parallel(block_H, block_N): + score_index = start + k * block_N + j + global_head = by * VALID_BLOCK_H + i + if score_index < cache_seqlens[bx]: + if global_head < VALID_OUTPUT_HEADS: + AttnScore[bx, global_head, score_index] = acc_s[i, j] + else: + T.reduce_max(acc_s, token_scores, dim=0) if score_mode == "direct": for j in T.Parallel(block_N): score_index = start + k * block_N + j - AttnScore[bx, 0, score_index] = T.if_then_else( - score_index < cache_seqlens[bx], - token_scores[j], - AttnScore[bx, 0, score_index], - ) + if score_index < cache_seqlens[bx]: + AttnScore[bx, 0, score_index] = token_scores[j] elif score_mode == "atomic": for j in T.Parallel(block_N): score_index = start + k * block_N + j @@ -207,7 +218,7 @@ def main_split( AttnScore[bx, 0, score_index], token_scores[j], ) - else: + elif score_mode == "partial": for j in T.Parallel(block_N): score_index = start + k * block_N + j AttnScore[bx, by, score_index] = T.if_then_else( @@ -351,15 +362,20 @@ def main_no_split( for i, j in T.Parallel(block_H, block_N): acc_s[i, j] = T.if_then_else(by * VALID_BLOCK_H + i >= VALID_OUTPUT_HEADS, -T.infinity(accum_dtype), acc_s[i, j]) if need_score: - T.reduce_max(acc_s, token_scores, dim=0) + if score_mode == "per_head": + for i, j in T.Parallel(block_H, block_N): + score_index = k * block_N + j + global_head = by * VALID_BLOCK_H + i + if score_index < cache_seqlens[bx]: + if global_head < VALID_OUTPUT_HEADS: + AttnScore[bx, global_head, score_index] = acc_s[i, j] + else: + T.reduce_max(acc_s, token_scores, dim=0) if score_mode == "direct": for j in T.Parallel(block_N): score_index = k * block_N + j - AttnScore[bx, 0, score_index] = T.if_then_else( - score_index < cache_seqlens[bx], - token_scores[j], - AttnScore[bx, 0, score_index], - ) + if score_index < cache_seqlens[bx]: + AttnScore[bx, 0, score_index] = token_scores[j] elif score_mode == "atomic": for j in T.Parallel(block_N): score_index = k * block_N + j @@ -368,7 +384,7 @@ def main_no_split( AttnScore[bx, 0, score_index], token_scores[j], ) - else: + elif score_mode == "partial": for j in T.Parallel(block_N): score_index = k * block_N + j AttnScore[bx, by, score_index] = T.if_then_else( diff --git a/src/sparsevllm/kernels/tilelang/mla/runtime.py b/src/sparsevllm/kernels/tilelang/mla/runtime.py index ed11a5f2..4c4d21a8 100644 --- a/src/sparsevllm/kernels/tilelang/mla/runtime.py +++ b/src/sparsevllm/kernels/tilelang/mla/runtime.py @@ -17,7 +17,7 @@ _VALID_SPLITS = (1, 2, 4, 8, 16, 32) _SUPPORTED_VALID_HEADS = (5, 10, 20) -_SCORE_MODES = ("direct", "atomic", "partial") +_SCORE_MODES = ("direct", "atomic", "partial", "per_head") _HEAD_TILE_SIZE = 16 _LATENT_DIM = 512 _ROPE_DIM = 64 @@ -99,6 +99,88 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True, slots=True) +class TileMlaLaunchPlan: + """Capture-time TileLang variants for one model/device envelope.""" + + context_capacity: int + local_q_heads: int + max_batch_size: int + need_score: bool + configs: tuple[TileMlaLaunchConfig, ...] + + @classmethod + def build( + cls, + *, + context_capacity: int, + local_q_heads: int, + max_batch_size: int, + need_score: bool, + score_mode: str | None = None, + ) -> TileMlaLaunchPlan: + if min(context_capacity, max_batch_size) <= 0: + raise ValueError( + "TileLang MLA launch plan requires positive context and batch " + f"capacities, got context={context_capacity} " + f"batch={max_batch_size}." + ) + configs = [] + for batch_size in range(1, int(max_batch_size) + 1): + config = select_tile_mla_config( + batch_size=batch_size, + context_capacity=int(context_capacity), + need_score=bool(need_score), + local_q_heads=int(local_q_heads), + ) + if score_mode is not None: + config = TileMlaLaunchConfig( + num_split=config.num_split, + block_n=config.block_n, + block_h=config.block_h, + score_mode=score_mode, + ) + configs.append(config) + return cls( + context_capacity=int(context_capacity), + local_q_heads=int(local_q_heads), + max_batch_size=int(max_batch_size), + need_score=bool(need_score), + configs=tuple(configs), + ) + + def config_for(self, batch_size: int, *, need_score: bool) -> TileMlaLaunchConfig: + if bool(need_score) != self.need_score: + raise ValueError( + "TileLang MLA launch plan score contract changed after binding: " + f"planned={self.need_score} requested={bool(need_score)}." + ) + if not 0 < int(batch_size) <= self.max_batch_size: + raise ValueError( + "TileLang MLA batch exceeds the static launch plan: " + f"batch={batch_size} max={self.max_batch_size}." + ) + return self.configs[int(batch_size) - 1] + + def metadata(self) -> dict[str, object]: + return { + "context_capacity": self.context_capacity, + "local_q_heads": self.local_q_heads, + "max_batch_size": self.max_batch_size, + "need_score": self.need_score, + "batch_configs": [ + { + "batch_size": batch_size, + "num_split": config.num_split, + "block_n": config.block_n, + "block_h": config.block_h, + "score_mode": config.score_mode, + } + for batch_size, config in enumerate(self.configs, start=1) + ], + } + + def select_tile_mla_config( *, batch_size: int, @@ -182,11 +264,21 @@ def __init__( softmax_scale: float, valid_heads: int = 10, fixed_config: TileMlaLaunchConfig | None = None, + launch_plan: TileMlaLaunchPlan | None = None, ) -> None: self.device = torch.device(device) self.softmax_scale = float(softmax_scale) self.valid_heads = int(valid_heads) self.padded_heads = _padded_head_count(self.valid_heads) + if fixed_config is not None and launch_plan is not None: + raise ValueError( + "TileLang MLA accepts either fixed_config or launch_plan, not both." + ) + if launch_plan is not None and launch_plan.local_q_heads != self.valid_heads: + raise ValueError( + "TileLang MLA launch plan head count does not match the runner: " + f"plan={launch_plan.local_q_heads} runner={self.valid_heads}." + ) if fixed_config is not None: if self.padded_heads % fixed_config.block_h: raise ValueError( @@ -195,8 +287,44 @@ def __init__( f"block_h={fixed_config.block_h}." ) self.fixed_config = fixed_config + self.launch_plan = launch_plan self._kernels: dict[_KernelKey, _BoundKernel] = {} + def runtime_metadata(self) -> dict[str, object]: + variants = [] + for key, bound in self._kernels.items(): + workspace_tensors = ( + bound.workspace.padded_latent, + bound.workspace.padded_rope, + bound.workspace.glse, + bound.workspace.partial_output, + bound.workspace.score, + ) + variants.append( + { + "batch_size": key.batch_size, + "cache_slot_count": key.cache_slot_count, + "active_slot_rows": key.active_slot_rows, + "active_slot_width": key.active_slot_width, + "score_capacity": key.score_capacity, + "num_split": key.num_split, + "block_h": key.block_h, + "score_mode": key.score_mode, + "need_score": key.need_score, + "workspace_bytes": sum( + tensor.numel() * tensor.element_size() + for tensor in workspace_tensors + ), + "workspace_data_ptrs": [ + tensor.data_ptr() for tensor in workspace_tensors + ], + } + ) + return { + "compiled_variant_count": len(variants), + "compiled_variants": variants, + } + def _config_for( self, *, @@ -204,6 +332,17 @@ def _config_for( context_capacity: int, need_score: bool, ) -> TileMlaLaunchConfig: + if self.launch_plan is not None: + if context_capacity > self.launch_plan.context_capacity: + raise ValueError( + "TileLang MLA runtime context exceeds the static launch plan: " + f"runtime={context_capacity} " + f"plan={self.launch_plan.context_capacity}." + ) + return self.launch_plan.config_for( + batch_size, + need_score=need_score, + ) if self.fixed_config is not None: return self.fixed_config return select_tile_mla_config( @@ -285,9 +424,13 @@ def _bind(self, key: _KernelKey) -> _BoundKernel: ), score=torch.empty( key.batch_size, - self.padded_heads // config.block_h - if config.score_mode == "partial" - else 1, + ( + self.valid_heads + if config.score_mode == "per_head" + else self.padded_heads // config.block_h + if config.score_mode == "partial" + else 1 + ), key.score_capacity, dtype=torch.float32, device=self.device, @@ -307,6 +450,7 @@ def _validate( output: torch.Tensor, attn_score: torch.Tensor | None, max_context_len: int, + config: TileMlaLaunchConfig, ) -> tuple[int, int]: batch_size = int(q_latent.shape[0]) expected = { @@ -361,10 +505,16 @@ def _validate( ) score_capacity = int(active_slots.shape[1]) if attn_score is not None: - if attn_score.ndim != 2 or int(attn_score.shape[0]) != batch_size: + expected_prefix = ( + (batch_size, self.valid_heads) + if config.score_mode == "per_head" + else (batch_size,) + ) + if tuple(attn_score.shape[:-1]) != expected_prefix: raise ValueError( - "TileLang MLA reduced attn_score must have shape " - f"[batch, capacity], got {tuple(attn_score.shape)}." + "TileLang MLA attn_score shape does not match the static " + f"{config.score_mode!r} contract: expected prefix " + f"{expected_prefix}, got {tuple(attn_score.shape)}." ) if ( attn_score.dtype != torch.float32 @@ -374,21 +524,24 @@ def _validate( "TileLang MLA attn_score must be FP32 on the query device, " f"got {attn_score.dtype} on {attn_score.device}." ) - if not attn_score.is_contiguous(): + if ( + not attn_score.is_contiguous() + and config.score_mode != "per_head" + ): raise ValueError( "TileLang MLA attn_score must be contiguous, got stride " f"{tuple(attn_score.stride())}." ) - score_capacity = int(attn_score.shape[1]) - if score_capacity <= 0 or score_capacity % _BLOCK_N: + score_capacity = int(attn_score.shape[-1]) + if score_capacity <= 0: raise ValueError( - "TileLang MLA context/score capacity must be a positive " - f"multiple of {_BLOCK_N}, got {score_capacity}." + "TileLang MLA context/score capacity must be positive, got " + f"{score_capacity}." ) - if score_capacity > int(active_slots.shape[1]): + if int(max_context_len) > int(active_slots.shape[1]): raise ValueError( - "TileLang MLA score capacity exceeds active slot width: " - f"score={score_capacity} slots={active_slots.shape[1]}." + "TileLang MLA active slot width does not cover max_context_len: " + f"max={max_context_len} slots={active_slots.shape[1]}." ) if not 0 < int(max_context_len) <= score_capacity: raise ValueError( @@ -412,6 +565,13 @@ def __call__( attn_score: torch.Tensor | None, max_context_len: int, ) -> torch.Tensor: + batch_size = int(q_latent.shape[0]) + need_score = attn_score is not None + config = self._config_for( + batch_size=batch_size, + context_capacity=int(max_context_len), + need_score=need_score, + ) batch_size, score_capacity = self._validate( q_latent, q_rope, @@ -423,12 +583,7 @@ def __call__( output, attn_score, max_context_len, - ) - need_score = attn_score is not None - config = self._config_for( - batch_size=batch_size, - context_capacity=int(max_context_len), - need_score=need_score, + config, ) key = _KernelKey( batch_size=batch_size, @@ -475,6 +630,11 @@ def __call__( if attn_score is not None: if config.score_mode == "partial": score_output.fill_(-1e20) + elif config.score_mode == "per_head": + if attn_score.is_contiguous(): + score_output = attn_score + else: + score_output.fill_(-1e20) else: score_output = attn_score.unsqueeze(1) bound.call( @@ -492,12 +652,19 @@ def __call__( ) if attn_score is not None and config.score_mode == "partial": torch.amax(score_output, dim=1, out=attn_score) + elif ( + attn_score is not None + and config.score_mode == "per_head" + and score_output is not attn_score + ): + attn_score.copy_(score_output) return output __all__ = [ "TileMlaDecodeKernel", "TileMlaLaunchConfig", + "TileMlaLaunchPlan", "TileMlaWorkspace", "select_tile_mla_config", "tilelang_mla_support", diff --git a/src/sparsevllm/kernels/triton/deltakv_kernels.py b/src/sparsevllm/kernels/triton/deltakv_kernels.py index df8e236b..52a32ee8 100644 --- a/src/sparsevllm/kernels/triton/deltakv_kernels.py +++ b/src/sparsevllm/kernels/triton/deltakv_kernels.py @@ -739,6 +739,9 @@ def _full_layer_kivi_flash_decode_stage1_kernel( FEAT_PER_INT: tl.constexpr, QUANT_MASK: tl.constexpr, STORE_SCORE: tl.constexpr, + FIXED_GRID: tl.constexpr, + MAX_EFFECTIVE_SPLITS: tl.constexpr, + TARGET_TOKENS_PER_SPLIT: tl.constexpr, ): cur_batch = tl.program_id(0) cur_kv_head = tl.program_id(1) @@ -750,8 +753,28 @@ def _full_layer_kivi_flash_decode_stage1_kernel( cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) cur_row = tl.load(Req_Indices + cur_batch).to(tl.int32) - cur_batch_start_index = seq_start_block * BLOCK_SEQ - cur_batch_end_index = tl.minimum(cur_batch_seq_len, cur_batch_start_index + BLOCK_SEQ) + if FIXED_GRID: + requested_splits = tl.cdiv(cur_batch_seq_len, TARGET_TOKENS_PER_SPLIT) + num_splits = tl.maximum( + 1, + tl.minimum(requested_splits, MAX_EFFECTIVE_SPLITS), + ) + split_tokens = tl.where( + requested_splits <= MAX_EFFECTIVE_SPLITS, + TARGET_TOKENS_PER_SPLIT, + tl.cdiv(cur_batch_seq_len, num_splits), + ) + cur_batch_start_index = seq_start_block * split_tokens + cur_batch_end_index = tl.minimum( + cur_batch_seq_len, + cur_batch_start_index + split_tokens, + ) + else: + cur_batch_start_index = seq_start_block * BLOCK_SEQ + cur_batch_end_index = tl.minimum( + cur_batch_seq_len, + cur_batch_start_index + BLOCK_SEQ, + ) off_q = cur_batch * stride_qbs + cur_q_head_range[:, None] * stride_qh + offs_d[None, :] * stride_qd q = tl.load( @@ -995,6 +1018,8 @@ def full_layer_kivi_flash_decode_stage1( num_warps: int = 2, num_stages: int = 3, attn_score: torch.Tensor | None = None, + max_kv_splits: int | None = None, + target_tokens_per_split: int | None = None, ): assert q.is_cuda and raw_k.is_cuda and raw_v.is_cuda assert raw_slots_map.is_cuda and kivi_block_slots_map.is_cuda and kivi_block_start_pos.is_cuda @@ -1064,7 +1089,36 @@ def full_layer_kivi_flash_decode_stage1( if int(q.shape[1]) % num_kv_heads != 0: raise ValueError(f"Q heads must be divisible by KV heads, got {q.shape[1]}/{num_kv_heads}.") - grid = (batch, num_kv_heads, triton.cdiv(max_len_in_batch, block_seq)) + fixed_grid = max_kv_splits is not None + if fixed_grid: + max_kv_splits = int(max_kv_splits) + target_tokens_per_split = int(target_tokens_per_split or 0) + if max_kv_splits <= 0 or target_tokens_per_split <= 0: + raise ValueError( + "Full-layer KIVI fixed-grid decode requires positive split capacity " + "and target tokens per split." + ) + if int(mid_out.shape[2]) != max_kv_splits or int(mid_out_logsumexp.shape[2]) != max_kv_splits: + raise ValueError( + "Full-layer KIVI fixed-grid workspace does not match the split envelope: " + f"mid_out={tuple(mid_out.shape)} mid_lse={tuple(mid_out_logsumexp.shape)} " + f"max_kv_splits={max_kv_splits}." + ) + if req_indices.dtype != torch.int32 or req_indices.stride(0) != 1: + raise TypeError("Full-layer KIVI fixed-grid req_indices must be contiguous int32.") + if context_lens.dtype != torch.int32 or context_lens.stride(0) != 1: + raise TypeError("Full-layer KIVI fixed-grid context_lens must be contiguous int32.") + req_indices_i32 = req_indices + context_lens_i32 = context_lens + grid_splits = max_kv_splits + else: + max_kv_splits = 1 + target_tokens_per_split = 1 + req_indices_i32 = req_indices.to(torch.int32).contiguous() + context_lens_i32 = context_lens.to(torch.int32).contiguous() + grid_splits = triton.cdiv(max_len_in_batch, block_seq) + + grid = (batch, num_kv_heads, grid_splits) gqa_group_size = int(q.shape[1]) // num_kv_heads score_arg = attn_score if attn_score is not None else mid_out_logsumexp score_stride_b = score_arg.stride(0) if attn_score is not None else 0 @@ -1083,8 +1137,8 @@ def full_layer_kivi_flash_decode_stage1( value_packed, value_scales, value_mins, - req_indices.to(torch.int32).contiguous(), - context_lens.to(torch.int32).contiguous(), + req_indices_i32, + context_lens_i32, mid_out, mid_out_logsumexp, score_arg, @@ -1136,6 +1190,9 @@ def full_layer_kivi_flash_decode_stage1( FEAT_PER_INT=8, QUANT_MASK=15, STORE_SCORE=attn_score is not None, + FIXED_GRID=fixed_grid, + MAX_EFFECTIVE_SPLITS=max_kv_splits, + TARGET_TOKENS_PER_SPLIT=target_tokens_per_split, num_warps=num_warps, num_stages=num_stages, ) diff --git a/src/sparsevllm/kernels/triton/flashinfer_decode_metadata.py b/src/sparsevllm/kernels/triton/flashinfer_decode_metadata.py new file mode 100644 index 00000000..21f1b20b --- /dev/null +++ b/src/sparsevllm/kernels/triton/flashinfer_decode_metadata.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _pack_page_indices_kernel( + active_slots, + request_indices, + context_lens, + packed_indices, + active_slots_stride_0: tl.constexpr, + active_slots_stride_1: tl.constexpr, + BATCH_SIZE: tl.constexpr, + BATCH_BLOCK: tl.constexpr, + CONTEXT_CAPACITY: tl.constexpr, + TOKEN_BLOCK: tl.constexpr, +): + batch_idx = tl.program_id(0) + token_block_idx = tl.program_id(1) + + batch_offsets = tl.arange(0, BATCH_BLOCK) + lengths = tl.load(context_lens + batch_offsets, mask=batch_offsets < BATCH_SIZE) + packed_start = tl.sum(tl.where(batch_offsets < batch_idx, lengths, 0)) + + token_offsets = token_block_idx * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK) + context_len = tl.load(context_lens + batch_idx) + request_idx = tl.load(request_indices + batch_idx) + valid = (token_offsets < context_len) & (token_offsets < CONTEXT_CAPACITY) + slots = tl.load( + active_slots + + request_idx * active_slots_stride_0 + + token_offsets * active_slots_stride_1, + mask=valid, + ) + tl.store(packed_indices + packed_start + token_offsets, slots, mask=valid) + + +def pack_flashinfer_page_indices( + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + packed_indices: torch.Tensor, + *, + context_capacity: int, +) -> None: + """Pack a layer's page-size-one slot table into graph-stable storage.""" + + if active_slots.ndim != 2 or active_slots.dtype != torch.int32: + raise TypeError("FlashInfer graph decode requires a rank-2 int32 slot table.") + if request_indices.ndim != 1 or request_indices.dtype != torch.int32: + raise TypeError("FlashInfer graph decode requires int32 request indices.") + if context_lens.ndim != 1 or context_lens.dtype != torch.int32: + raise TypeError("FlashInfer graph decode requires int32 context lengths.") + if request_indices.shape != context_lens.shape: + raise ValueError("FlashInfer graph request indices and context lengths must match.") + + batch_size = int(context_lens.numel()) + context_capacity = int(context_capacity) + if context_capacity <= 0 or context_capacity > int(active_slots.shape[1]): + raise ValueError( + "FlashInfer graph context capacity is outside the slot table: " + f"capacity={context_capacity} width={int(active_slots.shape[1])}." + ) + if packed_indices.ndim != 1 or packed_indices.dtype != torch.int32: + raise TypeError("FlashInfer graph packed indices must be a 1D int32 tensor.") + required = batch_size * context_capacity + if int(packed_indices.numel()) < required: + raise ValueError( + "FlashInfer graph packed-index buffer is too small: " + f"required={required} actual={int(packed_indices.numel())}." + ) + + token_block = 128 + _pack_page_indices_kernel[ + (batch_size, triton.cdiv(context_capacity, token_block)) + ]( + active_slots, + request_indices, + context_lens, + packed_indices, + active_slots.stride(0), + active_slots.stride(1), + BATCH_SIZE=batch_size, + BATCH_BLOCK=triton.next_power_of_2(batch_size), + CONTEXT_CAPACITY=context_capacity, + TOKEN_BLOCK=token_block, + ) + + +__all__ = ["pack_flashinfer_page_indices"] diff --git a/src/sparsevllm/kernels/triton/gemma4_decode_attention.py b/src/sparsevllm/kernels/triton/gemma4_decode_attention.py deleted file mode 100644 index bbf9e29a..00000000 --- a/src/sparsevllm/kernels/triton/gemma4_decode_attention.py +++ /dev/null @@ -1,320 +0,0 @@ -from __future__ import annotations - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _gemma4_decode_stage1_kernel( - q, - k, - v, - active_slots, - req_indices, - context_lens, - mid_output, - mid_lse, - attn_score, - stride_qb, - stride_qh, - stride_kt, - stride_kh, - stride_vt, - stride_vh, - stride_sb, - stride_ss, - stride_mob, - stride_moh, - stride_mos, - stride_mlb, - stride_mlh, - stride_mls, - stride_asb, - stride_ash, - stride_asl, - group_size, - HEAD_DIM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - BLOCK_N: tl.constexpr, - WINDOW: tl.constexpr, - SCORE_MODE: tl.constexpr, -): - batch = tl.program_id(0) - query_head = tl.program_id(1) - sequence_block = tl.program_id(2) - kv_head = query_head // group_size - dims = tl.arange(0, HEAD_DIM) - sequence_len = tl.load(context_lens + batch) - request = tl.load(req_indices + batch) - block_start = sequence_block * BLOCK_SEQ - mid_offset = ( - batch * stride_mob + query_head * stride_moh + sequence_block * stride_mos - ) - if block_start >= sequence_len: - tl.store(mid_output + mid_offset + dims, 0.0) - tl.store( - mid_lse - + batch * stride_mlb - + query_head * stride_mlh - + sequence_block * stride_mls, - -float("inf"), - ) - return - if WINDOW > 0: - block_start = tl.maximum(block_start, sequence_len - WINDOW) - block_end = tl.minimum(sequence_len, (sequence_block + 1) * BLOCK_SEQ) - query = tl.load(q + batch * stride_qb + query_head * stride_qh + dims) - max_logit = tl.full((), -float("inf"), tl.float32) - denominator = tl.zeros((), tl.float32) - accumulator = tl.zeros((HEAD_DIM,), tl.float32) - for offset in range(0, BLOCK_SEQ, BLOCK_N): - positions = sequence_block * BLOCK_SEQ + offset + tl.arange(0, BLOCK_N) - visible = (positions >= block_start) & (positions < block_end) - slots = tl.load( - active_slots + request * stride_sb + positions * stride_ss, - mask=visible, - other=0, - ) - key = tl.load( - k + slots[None, :] * stride_kt + kv_head * stride_kh + dims[:, None], - mask=visible[None, :], - other=0.0, - ) - logits = ( - tl.reshape(tl.dot(query[None, :], key), (BLOCK_N,)) * 1.4426950408889634 - ) - if SCORE_MODE == 3: - tl.store( - attn_score - + batch * stride_asb - + query_head * stride_ash - + positions * stride_asl, - logits * 0.6931471805599453, - mask=visible, - ) - elif SCORE_MODE == 2: - tl.atomic_max( - attn_score + batch * stride_asb + positions * stride_asl, - logits * 0.6931471805599453, - mask=visible, - ) - logits = tl.where(visible, logits, -float("inf")) - has_visible_key = tl.max(visible.to(tl.int32), axis=0) > 0 - block_max = tl.max(logits, axis=0) - new_max = tl.where(has_visible_key, tl.maximum(max_logit, block_max), max_logit) - probabilities = tl.where(visible, tl.exp2(logits - new_max), 0.0) - correction = tl.where(has_visible_key, tl.exp2(max_logit - new_max), 1.0) - denominator = denominator * correction + tl.sum(probabilities, axis=0) - accumulator *= correction - value = tl.load( - v + slots[:, None] * stride_vt + kv_head * stride_vh + dims[None, :], - mask=visible[:, None], - other=0.0, - ) - accumulator += tl.reshape( - tl.dot(probabilities[None, :].to(value.dtype), value), (HEAD_DIM,) - ) - max_logit = new_max - valid_block = block_end > block_start - tl.store( - mid_output + mid_offset + dims, - tl.where(valid_block, accumulator / denominator, 0.0), - ) - tl.store( - mid_lse - + batch * stride_mlb - + query_head * stride_mlh - + sequence_block * stride_mls, - tl.where( - valid_block, - max_logit * 0.6931471805599453 + tl.log(denominator), - -float("inf"), - ), - ) - - -@torch.no_grad() -def gemma4_decode_stage1( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - active_slots: torch.Tensor, - req_indices: torch.Tensor, - context_lens: torch.Tensor, - mid_output: torch.Tensor, - mid_lse: torch.Tensor, - *, - block_seq: int, - sliding_window: int | None, - attn_score: torch.Tensor | None = None, -) -> None: - head_dim = int(q.shape[-1]) - if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape: - raise ValueError("Gemma 4 decode requires matching rank-3 Q/K/V.") - if head_dim not in {256, 512} or int(k.shape[-1]) != head_dim: - raise ValueError( - f"Gemma 4 decode requires head_dim 256 or 512, got {head_dim}." - ) - if not all(t.is_cuda for t in (q, k, v, mid_output, mid_lse)): - raise TypeError("Gemma 4 decode requires CUDA tensors.") - if q.dtype not in {torch.float16, torch.bfloat16} or any( - t.dtype != q.dtype for t in (k, v) - ): - raise TypeError("Gemma 4 decode requires matching FP16 or BF16 Q/K/V.") - if mid_output.dtype != torch.float32 or mid_lse.dtype != torch.float32: - raise TypeError( - "Gemma 4 decode workspace must use FP32 output and LSE tensors." - ) - if int(q.shape[1]) % int(k.shape[1]): - raise ValueError("Gemma 4 decode requires divisible Q and KV heads.") - if int(block_seq) <= 0: - raise ValueError(f"Gemma 4 decode requires block_seq > 0, got {block_seq}.") - block_n = 32 if head_dim == 256 else 16 - if attn_score is not None and attn_score.dim() not in {2, 3}: - raise ValueError( - "Gemma 4 decode attention scores must be [B, L] or [B, H, L], " - f"got {tuple(attn_score.shape)}." - ) - score = mid_lse if attn_score is None else attn_score - score_head_stride = score.stride(1) if score.dim() == 3 else 0 - score_length_stride = score.stride(-1) - _gemma4_decode_stage1_kernel[ - (int(q.shape[0]), int(q.shape[1]), int(mid_output.shape[2])) - ]( - q, - k, - v, - active_slots, - req_indices, - context_lens, - mid_output, - mid_lse, - score, - q.stride(0), - q.stride(1), - k.stride(0), - k.stride(1), - v.stride(0), - v.stride(1), - active_slots.stride(0), - active_slots.stride(1), - mid_output.stride(0), - mid_output.stride(1), - mid_output.stride(2), - mid_lse.stride(0), - mid_lse.stride(1), - mid_lse.stride(2), - score.stride(0), - score_head_stride, - score_length_stride, - int(q.shape[1]) // int(k.shape[1]), - HEAD_DIM=head_dim, - BLOCK_SEQ=int(block_seq), - BLOCK_N=block_n, - WINDOW=int(sliding_window or 0), - SCORE_MODE=0 if attn_score is None else attn_score.dim(), - num_warps=8, - num_stages=1, - ) - - -@triton.jit -def _gemma4_decode_stage2_kernel( - context_lens, - mid_output, - mid_lse, - output, - stride_mob, - stride_moh, - stride_mos, - stride_mlb, - stride_mlh, - stride_mls, - stride_ob, - stride_oh, - HEAD_DIM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - WINDOW: tl.constexpr, -): - batch = tl.program_id(0) - head = tl.program_id(1) - dims = tl.arange(0, HEAD_DIM) - sequence_len = tl.load(context_lens + batch) - first_block = 0 - if WINDOW > 0: - first_block = tl.maximum(0, sequence_len - WINDOW) // BLOCK_SEQ - block_count = (sequence_len + BLOCK_SEQ - 1) // BLOCK_SEQ - max_lse = tl.full((), -float("inf"), tl.float32) - denominator = tl.zeros((), tl.float32) - accumulator = tl.zeros((HEAD_DIM,), tl.float32) - for block in range(first_block, block_count): - lse = tl.load( - mid_lse + batch * stride_mlb + head * stride_mlh + block * stride_mls - ) - value = tl.load( - mid_output - + batch * stride_mob - + head * stride_moh - + block * stride_mos - + dims - ) - new_max = tl.maximum(max_lse, lse) - old_scale = tl.exp(max_lse - new_max) - new_scale = tl.exp(lse - new_max) - accumulator = accumulator * old_scale + value * new_scale - denominator = denominator * old_scale + new_scale - max_lse = new_max - tl.store( - output + batch * stride_ob + head * stride_oh + dims, accumulator / denominator - ) - - -@torch.no_grad() -def gemma4_decode_stage2( - mid_output: torch.Tensor, - mid_lse: torch.Tensor, - context_lens: torch.Tensor, - output: torch.Tensor, - *, - block_seq: int, - sliding_window: int | None, -) -> None: - head_dim = int(mid_output.shape[-1]) - if head_dim not in {256, 512}: - raise ValueError( - f"Gemma 4 decode stage 2 requires head_dim 256 or 512, got {head_dim}." - ) - if not all(t.is_cuda for t in (mid_output, mid_lse, output)): - raise TypeError("Gemma 4 decode stage 2 requires CUDA tensors.") - if mid_output.dtype != torch.float32 or mid_lse.dtype != torch.float32: - raise TypeError("Gemma 4 decode stage 2 workspace must use FP32 tensors.") - if output.dtype not in {torch.float16, torch.bfloat16}: - raise TypeError("Gemma 4 decode stage 2 output must use FP16 or BF16.") - if output.shape[:2] != mid_output.shape[:2] or output.shape[-1] != head_dim: - raise ValueError( - "Gemma 4 decode stage 2 requires matching batch/head/output shape." - ) - if int(block_seq) <= 0: - raise ValueError( - f"Gemma 4 decode stage 2 requires block_seq > 0, got {block_seq}." - ) - _gemma4_decode_stage2_kernel[(int(output.shape[0]), int(output.shape[1]))]( - context_lens, - mid_output, - mid_lse, - output, - mid_output.stride(0), - mid_output.stride(1), - mid_output.stride(2), - mid_lse.stride(0), - mid_lse.stride(1), - mid_lse.stride(2), - output.stride(0), - output.stride(1), - HEAD_DIM=head_dim, - BLOCK_SEQ=int(block_seq), - WINDOW=int(sliding_window or 0), - num_warps=8, - num_stages=2, - ) diff --git a/src/sparsevllm/kernels/triton/gemma4_global_decode_attention.py b/src/sparsevllm/kernels/triton/gemma4_global_decode_attention.py deleted file mode 100644 index a5b33a88..00000000 --- a/src/sparsevllm/kernels/triton/gemma4_global_decode_attention.py +++ /dev/null @@ -1,185 +0,0 @@ -from __future__ import annotations - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _gemma4_global_decode_stage1_kernel( - q, - k, - v, - active_slots, - req_indices, - context_lens, - mid_output, - mid_lse, - stride_qb, - stride_qh, - stride_kt, - stride_kh, - stride_vt, - stride_vh, - stride_sb, - stride_ss, - stride_mob, - stride_moh, - stride_mos, - stride_mlb, - stride_mlh, - stride_mls, - GROUP_SIZE: tl.constexpr, - HEADS_PER_PROGRAM: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - BLOCK_N: tl.constexpr, -): - batch = tl.program_id(0) - head_group = tl.program_id(1) - sequence_block = tl.program_id(2) - heads = head_group * HEADS_PER_PROGRAM + tl.arange(0, HEADS_PER_PROGRAM) - kv_head = head_group * HEADS_PER_PROGRAM // GROUP_SIZE - dims = tl.arange(0, HEAD_DIM) - sequence_len = tl.load(context_lens + batch) - block_start = sequence_block * BLOCK_SEQ - mid_offset = ( - batch * stride_mob - + heads[:, None] * stride_moh - + sequence_block * stride_mos - ) - lse_offset = ( - batch * stride_mlb + heads * stride_mlh + sequence_block * stride_mls - ) - if block_start >= sequence_len: - tl.store(mid_output + mid_offset + dims[None, :], 0.0) - tl.store(mid_lse + lse_offset, -float("inf")) - return - - query = tl.load(q + batch * stride_qb + heads[:, None] * stride_qh + dims) - max_logit = tl.full((HEADS_PER_PROGRAM,), -float("inf"), tl.float32) - denominator = tl.zeros((HEADS_PER_PROGRAM,), tl.float32) - accumulator = tl.zeros((HEADS_PER_PROGRAM, HEAD_DIM), tl.float32) - block_end = tl.minimum(sequence_len, block_start + BLOCK_SEQ) - request = tl.load(req_indices + batch) - for offset in range(0, BLOCK_SEQ, BLOCK_N): - positions = block_start + offset + tl.arange(0, BLOCK_N) - visible = positions < block_end - slots = tl.load( - active_slots + request * stride_sb + positions * stride_ss, - mask=visible, - other=0, - ) - key = tl.load( - k + slots[None, :] * stride_kt + kv_head * stride_kh + dims[:, None], - mask=visible[None, :], - other=0.0, - ) - logits = tl.dot(query, key) * 1.4426950408889634 - logits = tl.where(visible[None, :], logits, -float("inf")) - block_max = tl.max(logits, axis=1) - new_max = tl.maximum(max_logit, block_max) - probabilities = tl.exp2(logits - new_max[:, None]) - correction = tl.exp2(max_logit - new_max) - denominator = denominator * correction + tl.sum(probabilities, axis=1) - accumulator *= correction[:, None] - value = tl.load( - v + slots[:, None] * stride_vt + kv_head * stride_vh + dims[None, :], - mask=visible[:, None], - other=0.0, - ) - accumulator += tl.dot(probabilities.to(value.dtype), value) - max_logit = new_max - tl.store(mid_output + mid_offset + dims[None, :], accumulator / denominator[:, None]) - tl.store( - mid_lse + lse_offset, - max_logit * 0.6931471805599453 + tl.log(denominator), - ) - - -def gemma4_global_decode_stage1( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - active_slots: torch.Tensor, - req_indices: torch.Tensor, - context_lens: torch.Tensor, - mid_output: torch.Tensor, - mid_lse: torch.Tensor, - *, - block_seq: int, - heads_per_program: int = 4, -) -> None: - if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape: - raise ValueError("Gemma 4 global decode requires matching rank-3 Q/K/V.") - if not all(t.is_cuda for t in (q, k, v, mid_output, mid_lse)): - raise TypeError("Gemma 4 global decode requires CUDA tensors.") - if q.dtype not in {torch.float16, torch.bfloat16} or any( - tensor.dtype != q.dtype for tensor in (k, v) - ): - raise TypeError("Gemma 4 global decode requires matching FP16 or BF16 Q/K/V.") - if int(q.shape[1]) % int(k.shape[1]): - raise ValueError("Gemma 4 global decode requires divisible Q and KV heads.") - head_dim = int(q.shape[-1]) - group_size = int(q.shape[1]) // int(k.shape[1]) - heads_per_program = int(heads_per_program) - if ( - head_dim != 512 - or int(k.shape[-1]) != head_dim - or group_size % heads_per_program - or heads_per_program not in {2, 4} - ): - raise ValueError( - "Gemma 4 global decode requires head_dim=512 and GQA groups divisible " - f"by 2 or 4, got head_dim={head_dim}, group_size={group_size}, " - f"heads_per_program={heads_per_program}." - ) - if mid_output.dtype != torch.float32 or mid_lse.dtype != torch.float32: - raise TypeError("Gemma 4 global decode workspace must use FP32 tensors.") - expected_mid = (q.shape[0], q.shape[1], mid_output.shape[2], head_dim) - expected_lse = expected_mid[:-1] - if mid_output.shape != expected_mid or mid_lse.shape != expected_lse: - raise ValueError( - f"Gemma 4 global decode workspace must have shapes {expected_mid} and " - f"{expected_lse}, got {tuple(mid_output.shape)} and {tuple(mid_lse.shape)}." - ) - _gemma4_global_decode_stage1_kernel[ - ( - int(q.shape[0]), - int(q.shape[1]) // heads_per_program, - int(mid_output.shape[2]), - ) - ]( - q, - k, - v, - active_slots, - req_indices, - context_lens, - mid_output, - mid_lse, - q.stride(0), - q.stride(1), - k.stride(0), - k.stride(1), - v.stride(0), - v.stride(1), - active_slots.stride(0), - active_slots.stride(1), - mid_output.stride(0), - mid_output.stride(1), - mid_output.stride(2), - mid_lse.stride(0), - mid_lse.stride(1), - mid_lse.stride(2), - GROUP_SIZE=group_size, - HEADS_PER_PROGRAM=heads_per_program, - HEAD_DIM=head_dim, - BLOCK_SEQ=int(block_seq), - BLOCK_N=16, - num_warps=8, - num_stages=1, - ) - - -__all__ = ["gemma4_global_decode_stage1"] diff --git a/src/sparsevllm/kernels/triton/gemma4_single_block_decode_attention.py b/src/sparsevllm/kernels/triton/gemma4_single_block_decode_attention.py deleted file mode 100644 index 2106e7ae..00000000 --- a/src/sparsevllm/kernels/triton/gemma4_single_block_decode_attention.py +++ /dev/null @@ -1,154 +0,0 @@ -from __future__ import annotations - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _gemma4_single_block_decode_kernel( - q, - k, - v, - active_slots, - req_indices, - context_lens, - output, - stride_qb, - stride_qh, - stride_kt, - stride_kh, - stride_vt, - stride_vh, - stride_sb, - stride_ss, - stride_ob, - stride_oh, - GROUP_SIZE: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - BLOCK_N: tl.constexpr, - WINDOW: tl.constexpr, -): - batch = tl.program_id(0) - kv_head = tl.program_id(1) - groups = tl.arange(0, GROUP_SIZE) - dims = tl.arange(0, HEAD_DIM) - sequence_len = tl.load(context_lens + batch) - start = tl.maximum(0, sequence_len - WINDOW) if WINDOW > 0 else 0 - query_head = kv_head * GROUP_SIZE + groups - query = tl.load( - q + batch * stride_qb + query_head[:, None] * stride_qh + dims[None, :] - ) - request = tl.load(req_indices + batch) - max_logit = tl.full((GROUP_SIZE,), -float("inf"), tl.float32) - denominator = tl.zeros((GROUP_SIZE,), tl.float32) - accumulator = tl.zeros((GROUP_SIZE, HEAD_DIM), tl.float32) - for offset in range(0, BLOCK_SEQ, BLOCK_N): - positions = offset + tl.arange(0, BLOCK_N) - visible = (positions >= start) & (positions < sequence_len) - slots = tl.load( - active_slots + request * stride_sb + positions * stride_ss, - mask=visible, - other=0, - ) - key = tl.load( - k + slots[None, :] * stride_kt + kv_head * stride_kh + dims[:, None], - mask=visible[None, :], - other=0.0, - ) - logits = tl.dot(query, key) * 1.4426950408889634 - logits = tl.where(visible[None, :], logits, -float("inf")) - block_max = tl.max(logits, axis=1) - new_max = tl.maximum(max_logit, block_max) - probabilities = tl.exp2(logits - new_max[:, None]) - correction = tl.exp2(max_logit - new_max) - denominator = denominator * correction + tl.sum(probabilities, axis=1) - accumulator *= correction[:, None] - value = tl.load( - v + slots[:, None] * stride_vt + kv_head * stride_vh + dims[None, :], - mask=visible[:, None], - other=0.0, - ) - accumulator += tl.dot(probabilities.to(value.dtype), value) - max_logit = new_max - offsets = batch * stride_ob + query_head[:, None] * stride_oh + dims[None, :] - tl.store(output + offsets, accumulator / denominator[:, None]) - - -def gemma4_single_block_decode( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - active_slots: torch.Tensor, - req_indices: torch.Tensor, - context_lens: torch.Tensor, - output: torch.Tensor, - *, - block_seq: int, - sliding_window: int | None, -) -> None: - if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape or output.shape != q.shape: - raise ValueError( - "Gemma 4 single-block decode requires matching rank-3 Q/K/V/output." - ) - if not all(t.is_cuda for t in (q, k, v, output)): - raise TypeError("Gemma 4 single-block decode requires CUDA tensors.") - if q.dtype not in {torch.float16, torch.bfloat16} or any( - t.dtype != q.dtype for t in (k, v, output) - ): - raise TypeError( - "Gemma 4 single-block decode requires matching FP16 or BF16 tensors." - ) - if any(t.stride(-1) != 1 for t in (q, k, v, output)): - raise ValueError( - "Gemma 4 single-block decode requires contiguous head dimensions." - ) - if int(q.shape[1]) % int(k.shape[1]): - raise ValueError( - "Gemma 4 single-block decode requires divisible Q and KV heads." - ) - group_size = int(q.shape[1]) // int(k.shape[1]) - if group_size not in {2, 4, 8}: - raise ValueError( - f"Gemma 4 single-block decode requires GQA group 2, 4, or 8, got {group_size}." - ) - head_dim = int(q.shape[-1]) - if head_dim not in {256, 512} or int(k.shape[-1]) != head_dim: - raise ValueError( - f"Gemma 4 single-block decode requires head_dim 256 or 512, got {head_dim}." - ) - if int(block_seq) <= 0: - raise ValueError( - f"Gemma 4 single-block decode requires block_seq > 0, got {block_seq}." - ) - block_n = 32 if head_dim == 256 else 16 - _gemma4_single_block_decode_kernel[(int(q.shape[0]), int(k.shape[1]))]( - q, - k, - v, - active_slots, - req_indices, - context_lens, - output, - q.stride(0), - q.stride(1), - k.stride(0), - k.stride(1), - v.stride(0), - v.stride(1), - active_slots.stride(0), - active_slots.stride(1), - output.stride(0), - output.stride(1), - GROUP_SIZE=group_size, - HEAD_DIM=head_dim, - BLOCK_SEQ=int(block_seq), - BLOCK_N=block_n, - WINDOW=int(sliding_window or 0), - num_warps=8, - num_stages=1, - ) - - -__all__ = ["gemma4_single_block_decode"] diff --git a/src/sparsevllm/kernels/triton/gemma4_window_decode_attention.py b/src/sparsevllm/kernels/triton/gemma4_window_decode_attention.py deleted file mode 100644 index 118fb2d5..00000000 --- a/src/sparsevllm/kernels/triton/gemma4_window_decode_attention.py +++ /dev/null @@ -1,266 +0,0 @@ -from __future__ import annotations - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _gemma4_window_decode_stage1_kernel( - q, - k, - v, - active_slots, - req_indices, - context_lens, - mid_output, - mid_lse, - stride_qb, - stride_qh, - stride_kt, - stride_kh, - stride_vt, - stride_vh, - stride_sb, - stride_ss, - stride_mob, - stride_moh, - stride_mos, - stride_mlb, - stride_mlh, - stride_mls, - GROUP_SIZE: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - BLOCK_N: tl.constexpr, - WINDOW: tl.constexpr, -): - batch = tl.program_id(0) - kv_head = tl.program_id(1) - sequence_block = tl.program_id(2) - groups = tl.arange(0, GROUP_SIZE) - dims = tl.arange(0, HEAD_DIM) - sequence_len = tl.load(context_lens + batch) - window_start = tl.maximum(0, sequence_len - WINDOW) - block_start = window_start + sequence_block * BLOCK_SEQ - query_head = kv_head * GROUP_SIZE + groups - mid_offset = ( - batch * stride_mob - + query_head[:, None] * stride_moh - + sequence_block * stride_mos - ) - lse_offset = ( - batch * stride_mlb - + query_head * stride_mlh - + sequence_block * stride_mls - ) - if block_start >= sequence_len: - tl.store(mid_output + mid_offset + dims[None, :], 0.0) - tl.store(mid_lse + lse_offset, -float("inf")) - return - - block_end = tl.minimum(sequence_len, block_start + BLOCK_SEQ) - query = tl.load( - q + batch * stride_qb + query_head[:, None] * stride_qh + dims[None, :] - ) - max_logit = tl.full((GROUP_SIZE,), -float("inf"), tl.float32) - denominator = tl.zeros((GROUP_SIZE,), tl.float32) - accumulator = tl.zeros((GROUP_SIZE, HEAD_DIM), tl.float32) - for offset in range(0, BLOCK_SEQ, BLOCK_N): - positions = block_start + offset + tl.arange(0, BLOCK_N) - visible = positions < block_end - request = tl.load(req_indices + batch) - slots = tl.load( - active_slots + request * stride_sb + positions * stride_ss, - mask=visible, - other=0, - ) - key = tl.load( - k + slots[None, :] * stride_kt + kv_head * stride_kh + dims[:, None], - mask=visible[None, :], - other=0.0, - ) - logits = tl.dot(query, key) * 1.4426950408889634 - logits = tl.where(visible[None, :], logits, -float("inf")) - block_max = tl.max(logits, axis=1) - new_max = tl.maximum(max_logit, block_max) - probabilities = tl.exp2(logits - new_max[:, None]) - correction = tl.exp2(max_logit - new_max) - denominator = denominator * correction + tl.sum(probabilities, axis=1) - accumulator *= correction[:, None] - value = tl.load( - v + slots[:, None] * stride_vt + kv_head * stride_vh + dims[None, :], - mask=visible[:, None], - other=0.0, - ) - accumulator += tl.dot(probabilities.to(value.dtype), value) - max_logit = new_max - tl.store(mid_output + mid_offset + dims[None, :], accumulator / denominator[:, None]) - tl.store( - mid_lse + lse_offset, - max_logit * 0.6931471805599453 + tl.log(denominator), - ) - - -@triton.jit -def _gemma4_window_decode_stage2_kernel( - context_lens, - mid_output, - mid_lse, - output, - stride_mob, - stride_moh, - stride_mos, - stride_mlb, - stride_mlh, - stride_mls, - stride_ob, - stride_oh, - GROUP_SIZE: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - NUM_BLOCKS: tl.constexpr, - WINDOW: tl.constexpr, -): - batch = tl.program_id(0) - kv_head = tl.program_id(1) - groups = tl.arange(0, GROUP_SIZE) - dims = tl.arange(0, HEAD_DIM) - query_head = kv_head * GROUP_SIZE + groups - sequence_len = tl.load(context_lens + batch) - block_count = (tl.minimum(sequence_len, WINDOW) + BLOCK_SEQ - 1) // BLOCK_SEQ - max_lse = tl.full((GROUP_SIZE,), -float("inf"), tl.float32) - denominator = tl.zeros((GROUP_SIZE,), tl.float32) - accumulator = tl.zeros((GROUP_SIZE, HEAD_DIM), tl.float32) - for block in range(0, NUM_BLOCKS): - valid = block < block_count - lse = tl.load( - mid_lse - + batch * stride_mlb - + query_head * stride_mlh - + block * stride_mls - ) - lse = tl.where(valid, lse, -float("inf")) - value = tl.load( - mid_output - + batch * stride_mob - + query_head[:, None] * stride_moh - + block * stride_mos - + dims[None, :] - ) - new_max = tl.maximum(max_lse, lse) - old_scale = tl.exp(max_lse - new_max) - new_scale = tl.exp(lse - new_max) - accumulator = accumulator * old_scale[:, None] + value * new_scale[:, None] - denominator = denominator * old_scale + new_scale - max_lse = new_max - tl.store( - output - + batch * stride_ob - + query_head[:, None] * stride_oh - + dims[None, :], - accumulator / denominator[:, None], - ) - - -def gemma4_window_decode( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - active_slots: torch.Tensor, - req_indices: torch.Tensor, - context_lens: torch.Tensor, - mid_output: torch.Tensor, - mid_lse: torch.Tensor, - output: torch.Tensor, - *, - block_seq: int, - sliding_window: int, -) -> None: - if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape or output.shape != q.shape: - raise ValueError("Gemma 4 window decode requires matching rank-3 Q/K/V/output.") - if not all(t.is_cuda for t in (q, k, v, mid_output, mid_lse, output)): - raise TypeError("Gemma 4 window decode requires CUDA tensors.") - if q.dtype not in {torch.float16, torch.bfloat16} or any( - tensor.dtype != q.dtype for tensor in (k, v, output) - ): - raise TypeError("Gemma 4 window decode requires matching FP16 or BF16 Q/K/V.") - if mid_output.dtype != torch.float32 or mid_lse.dtype != torch.float32: - raise TypeError("Gemma 4 window decode workspace must use FP32 tensors.") - if int(q.shape[1]) % int(k.shape[1]): - raise ValueError("Gemma 4 window decode requires divisible Q and KV heads.") - group_size = int(q.shape[1]) // int(k.shape[1]) - head_dim = int(q.shape[-1]) - if group_size not in {2, 4} or head_dim != 256 or int(k.shape[-1]) != head_dim: - raise ValueError( - "Gemma 4 window decode requires head_dim=256 and GQA group 2 or 4, " - f"got head_dim={head_dim}, group_size={group_size}." - ) - block_seq, sliding_window = int(block_seq), int(sliding_window) - if block_seq <= 0 or sliding_window <= 0: - raise ValueError("Gemma 4 window decode requires positive block and window sizes.") - num_blocks = triton.cdiv(sliding_window, block_seq) - if mid_output.shape[2] < num_blocks or mid_lse.shape[2] < num_blocks: - raise ValueError( - f"Gemma 4 window workspace needs {num_blocks} blocks, got " - f"{mid_output.shape[2]}/{mid_lse.shape[2]}." - ) - mid_output = mid_output[:, :, :num_blocks] - mid_lse = mid_lse[:, :, :num_blocks] - _gemma4_window_decode_stage1_kernel[ - (int(q.shape[0]), int(k.shape[1]), num_blocks) - ]( - q, - k, - v, - active_slots, - req_indices, - context_lens, - mid_output, - mid_lse, - q.stride(0), - q.stride(1), - k.stride(0), - k.stride(1), - v.stride(0), - v.stride(1), - active_slots.stride(0), - active_slots.stride(1), - mid_output.stride(0), - mid_output.stride(1), - mid_output.stride(2), - mid_lse.stride(0), - mid_lse.stride(1), - mid_lse.stride(2), - GROUP_SIZE=group_size, - HEAD_DIM=head_dim, - BLOCK_SEQ=block_seq, - BLOCK_N=32, - WINDOW=sliding_window, - num_warps=8, - num_stages=1, - ) - _gemma4_window_decode_stage2_kernel[(int(q.shape[0]), int(k.shape[1]))]( - context_lens, - mid_output, - mid_lse, - output, - mid_output.stride(0), - mid_output.stride(1), - mid_output.stride(2), - mid_lse.stride(0), - mid_lse.stride(1), - mid_lse.stride(2), - output.stride(0), - output.stride(1), - GROUP_SIZE=group_size, - HEAD_DIM=head_dim, - BLOCK_SEQ=block_seq, - NUM_BLOCKS=num_blocks, - WINDOW=sliding_window, - num_warps=8, - num_stages=2, - ) - - -__all__ = ["gemma4_window_decode"] diff --git a/src/sparsevllm/kernels/triton/mla/decode_schedule.py b/src/sparsevllm/kernels/triton/mla/decode_schedule.py index 56074dbb..4a5b57a2 100644 --- a/src/sparsevllm/kernels/triton/mla/decode_schedule.py +++ b/src/sparsevllm/kernels/triton/mla/decode_schedule.py @@ -114,15 +114,15 @@ def __post_init__(self) -> None: def select_glm_mla_decode_config( *, batch_size: int, - max_context_len: int, + context_capacity: int, local_q_heads: int, ) -> MlaDecodeLaunchConfig: - """Select a graph-stable launch config from static decode dimensions.""" + """Select a launch config from a capture-time context capacity.""" if batch_size <= 0: raise ValueError("batch_size must be positive") - if max_context_len <= 0: - raise ValueError("max_context_len must be positive") + if context_capacity <= 0: + raise ValueError("context_capacity must be positive") if local_q_heads <= 0: raise ValueError("local_q_heads must be positive") if local_q_heads != 10: @@ -131,7 +131,7 @@ def select_glm_mla_decode_config( return _GLM_MLA_TP2_SMALL_BATCH_CONFIG if batch_size <= 8: return _GLM_MLA_TP2_MEDIUM_BATCH_CONFIG - if max_context_len <= 1024: + if context_capacity <= 1024: return _GLM_MLA_TP2_SHORT_CONTEXT_CONFIG return _GLM_MLA_TP2_LARGE_BATCH_CONFIG diff --git a/src/sparsevllm/kernels/triton/paged_flash_decoding.py b/src/sparsevllm/kernels/triton/paged_flash_decoding.py new file mode 100644 index 00000000..3302936d --- /dev/null +++ b/src/sparsevllm/kernels/triton/paged_flash_decoding.py @@ -0,0 +1,579 @@ +"""Graph-stable split-KV paged decode attention. + +The stable decode kernels intentionally remain unchanged. This variant fixes +the CUDA launch grid and workspace split dimension while deriving the effective +split ranges from the device-resident context lengths. The split scheduling +follows the fixed-upper-bound design used by SGLang's Triton decode attention +(reference revision ed0a62e4), adapted to Sparse-vLLM's slot-table layout. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _paged_decode_stage1( + Q, + K, + V, + sm_scale, + Req_to_tokens, + B_req_idx, + B_Seqlen, + Mid_O, + Mid_Lse, + Attn_Score, + stride_req_b, + stride_req_s, + stride_qb, + stride_qh, + stride_kb, + stride_kh, + stride_vb, + stride_vh, + stride_mid_b, + stride_mid_h, + stride_mid_s, + stride_lse_b, + stride_lse_h, + stride_lse_s, + stride_score_b, + stride_score_h, + stride_score_s, + GQA_GROUP_SIZE: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + MAX_KV_SPLITS: tl.constexpr, + MAX_EFFECTIVE_SPLITS: tl.constexpr, + TARGET_TOKENS_PER_SPLIT: tl.constexpr, + SCORE_MODE: tl.constexpr, +): + batch_id = tl.program_id(0) + head_id = tl.program_id(1) + split_id = tl.program_id(2) + kv_head_id = head_id // GQA_GROUP_SIZE + + seq_len = tl.load(B_Seqlen + batch_id) + requested_splits = tl.cdiv(seq_len, TARGET_TOKENS_PER_SPLIT) + num_splits = tl.maximum(1, tl.minimum(requested_splits, MAX_EFFECTIVE_SPLITS)) + split_tokens = tl.where( + requested_splits <= MAX_EFFECTIVE_SPLITS, + TARGET_TOKENS_PER_SPLIT, + tl.cdiv(seq_len, num_splits), + ) + split_start = split_id * split_tokens + split_end = tl.minimum(split_start + split_tokens, seq_len) + split_valid = (split_id < num_splits) & (split_start < split_end) + if not split_valid: + return + + offs_d = tl.arange(0, HEAD_DIM) + q = tl.load(Q + batch_id * stride_qb + head_id * stride_qh + offs_d) + req_id = tl.load(B_req_idx + batch_id) + + max_logit = -float("inf") + exp_sum = 0.0 + acc = tl.zeros([HEAD_DIM], dtype=tl.float32) + block_count = tl.where(split_valid, tl.cdiv(split_end - split_start, BLOCK_N), 0) + + for block_id in range(0, block_count): + positions = split_start + block_id * BLOCK_N + tl.arange(0, BLOCK_N) + position_mask = positions < split_end + slots = tl.load( + Req_to_tokens + req_id * stride_req_b + positions * stride_req_s, + mask=position_mask, + other=0, + ).to(tl.int64) + k_offsets = slots[:, None] * stride_kb + kv_head_id * stride_kh + offs_d[None, :] + v_offsets = slots[:, None] * stride_vb + kv_head_id * stride_vh + offs_d[None, :] + k = tl.load(K + k_offsets, mask=position_mask[:, None], other=0.0) + v = tl.load(V + v_offsets, mask=position_mask[:, None], other=0.0) + logits = tl.sum(q[None, :].to(tl.float32) * k.to(tl.float32), axis=1) + logits = tl.where(position_mask, logits, -float("inf")) + if SCORE_MODE == 3: + score_offsets = ( + batch_id * stride_score_b + + head_id * stride_score_h + + positions * stride_score_s + ) + tl.store(Attn_Score + score_offsets, logits, mask=position_mask) + elif SCORE_MODE == 2: + score_offsets = batch_id * stride_score_b + positions * stride_score_s + tl.atomic_max(Attn_Score + score_offsets, logits, mask=position_mask) + logits *= sm_scale + + block_max = tl.max(logits, axis=0) + next_max = tl.maximum(max_logit, block_max) + old_scale = tl.exp(max_logit - next_max) + probs = tl.exp(logits - next_max) + acc = acc * old_scale + tl.sum(probs[:, None] * v, axis=0) + exp_sum = exp_sum * old_scale + tl.sum(probs, axis=0) + max_logit = next_max + + mid_offset = ( + batch_id * stride_mid_b + + head_id * stride_mid_h + + split_id * stride_mid_s + + offs_d + ) + lse_offset = ( + batch_id * stride_lse_b + head_id * stride_lse_h + split_id * stride_lse_s + ) + safe_sum = tl.where(split_valid, exp_sum, 1.0) + tl.store(Mid_O + mid_offset, tl.where(split_valid, acc / safe_sum, 0.0)) + tl.store( + Mid_Lse + lse_offset, + tl.where(split_valid, max_logit + tl.log(safe_sum), -float("inf")), + ) + + +@triton.jit +def _paged_grouped_decode_stage1( + Q, + K, + V, + sm_scale, + Req_to_tokens, + B_req_idx, + B_Seqlen, + Mid_O, + Mid_Lse, + Attn_Score, + stride_req_b, + stride_req_s, + stride_qb, + stride_qh, + stride_kb, + stride_kh, + stride_vb, + stride_vh, + stride_mid_b, + stride_mid_h, + stride_mid_s, + stride_lse_b, + stride_lse_h, + stride_lse_s, + stride_score_b, + stride_score_h, + stride_score_s, + GQA_GROUP_SIZE: tl.constexpr, + QUERY_HEAD_BLOCK: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + MAX_KV_SPLITS: tl.constexpr, + MAX_EFFECTIVE_SPLITS: tl.constexpr, + TARGET_TOKENS_PER_SPLIT: tl.constexpr, + SCORE_MODE: tl.constexpr, +): + batch_id = tl.program_id(0) + kv_head_id = tl.program_id(1) + split_id = tl.program_id(2) + head_offsets = tl.arange(0, QUERY_HEAD_BLOCK) + query_heads = kv_head_id * GQA_GROUP_SIZE + head_offsets + head_mask = head_offsets < GQA_GROUP_SIZE + + seq_len = tl.load(B_Seqlen + batch_id) + requested_splits = tl.cdiv(seq_len, TARGET_TOKENS_PER_SPLIT) + num_splits = tl.maximum(1, tl.minimum(requested_splits, MAX_EFFECTIVE_SPLITS)) + split_tokens = tl.where( + requested_splits <= MAX_EFFECTIVE_SPLITS, + TARGET_TOKENS_PER_SPLIT, + tl.cdiv(seq_len, num_splits), + ) + split_start = split_id * split_tokens + split_end = tl.minimum(split_start + split_tokens, seq_len) + split_valid = (split_id < num_splits) & (split_start < split_end) + if not split_valid: + return + + offs_d = tl.arange(0, HEAD_DIM) + q_offsets = batch_id * stride_qb + query_heads[:, None] * stride_qh + offs_d[None, :] + q = tl.load(Q + q_offsets, mask=head_mask[:, None], other=0.0) + req_id = tl.load(B_req_idx + batch_id) + + max_logit = tl.zeros([QUERY_HEAD_BLOCK], dtype=tl.float32) - float("inf") + exp_sum = tl.zeros([QUERY_HEAD_BLOCK], dtype=tl.float32) + acc = tl.zeros([QUERY_HEAD_BLOCK, HEAD_DIM], dtype=tl.float32) + block_count = tl.where(split_valid, tl.cdiv(split_end - split_start, BLOCK_N), 0) + + for block_id in range(0, block_count): + positions = split_start + block_id * BLOCK_N + tl.arange(0, BLOCK_N) + position_mask = positions < split_end + slots = tl.load( + Req_to_tokens + req_id * stride_req_b + positions * stride_req_s, + mask=position_mask, + other=0, + ).to(tl.int64) + k_offsets = slots[None, :] * stride_kb + kv_head_id * stride_kh + offs_d[:, None] + v_offsets = slots[:, None] * stride_vb + kv_head_id * stride_vh + offs_d[None, :] + k = tl.load(K + k_offsets, mask=position_mask[None, :], other=0.0) + v = tl.load(V + v_offsets, mask=position_mask[:, None], other=0.0) + logits = tl.dot(q, k) + logits = tl.where(position_mask[None, :], logits, -float("inf")) + if SCORE_MODE == 3: + score_offsets = ( + batch_id * stride_score_b + + query_heads[:, None] * stride_score_h + + positions[None, :] * stride_score_s + ) + tl.store( + Attn_Score + score_offsets, + logits, + mask=head_mask[:, None] & position_mask[None, :], + ) + elif SCORE_MODE == 2: + score_offsets = batch_id * stride_score_b + positions * stride_score_s + reduced_logits = tl.max( + tl.where(head_mask[:, None], logits, -float("inf")), + axis=0, + ) + tl.atomic_max( + Attn_Score + score_offsets, + reduced_logits, + mask=position_mask, + ) + logits *= sm_scale + + block_max = tl.max(logits, axis=1) + next_max = tl.maximum(max_logit, block_max) + old_scale = tl.exp(max_logit - next_max) + probs = tl.exp(logits - next_max[:, None]) + acc *= old_scale[:, None] + acc += tl.dot(probs.to(v.dtype), v) + exp_sum = exp_sum * old_scale + tl.sum(probs, axis=1) + max_logit = next_max + + safe_sum = tl.where(split_valid, exp_sum, 1.0) + mid_offsets = ( + batch_id * stride_mid_b + + query_heads[:, None] * stride_mid_h + + split_id * stride_mid_s + + offs_d[None, :] + ) + lse_offsets = ( + batch_id * stride_lse_b + + query_heads * stride_lse_h + + split_id * stride_lse_s + ) + tl.store( + Mid_O + mid_offsets, + tl.where(split_valid, acc / safe_sum[:, None], 0.0), + mask=head_mask[:, None], + ) + tl.store( + Mid_Lse + lse_offsets, + tl.where(split_valid, max_logit + tl.log(safe_sum), -float("inf")), + mask=head_mask, + ) + + +@triton.jit +def _paged_decode_stage2( + B_Seqlen, + Mid_O, + Mid_Lse, + O, + Out_Lse, + stride_mid_b, + stride_mid_h, + stride_mid_s, + stride_lse_b, + stride_lse_h, + stride_lse_s, + stride_ob, + stride_oh, + stride_out_lse_h, + stride_out_lse_b, + HEAD_DIM: tl.constexpr, + MAX_KV_SPLITS: tl.constexpr, + MAX_EFFECTIVE_SPLITS: tl.constexpr, + TARGET_TOKENS_PER_SPLIT: tl.constexpr, +): + batch_id = tl.program_id(0) + head_id = tl.program_id(1) + seq_len = tl.load(B_Seqlen + batch_id) + num_splits = tl.maximum( + 1, + tl.minimum( + tl.cdiv(seq_len, TARGET_TOKENS_PER_SPLIT), + MAX_EFFECTIVE_SPLITS, + ), + ) + + offs_d = tl.arange(0, HEAD_DIM) + mid_base = batch_id * stride_mid_b + head_id * stride_mid_h + offs_d + lse_base = batch_id * stride_lse_b + head_id * stride_lse_h + max_lse = -float("inf") + exp_sum = 0.0 + acc = tl.zeros([HEAD_DIM], dtype=tl.float32) + for split_id in range(0, num_splits): + split_lse = tl.load(Mid_Lse + lse_base + split_id * stride_lse_s) + split_o = tl.load(Mid_O + mid_base + split_id * stride_mid_s) + next_max = tl.maximum(max_lse, split_lse) + old_scale = tl.exp(max_lse - next_max) + split_scale = tl.exp(split_lse - next_max) + acc = acc * old_scale + split_scale * split_o + exp_sum = exp_sum * old_scale + split_scale + max_lse = next_max + + tl.store(O + batch_id * stride_ob + head_id * stride_oh + offs_d, acc / exp_sum) + tl.store( + Out_Lse + head_id * stride_out_lse_h + batch_id * stride_out_lse_b, + max_lse + tl.log(exp_sum), + ) + + +def _check_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + mid_o: torch.Tensor, + mid_lse: torch.Tensor, + attn_score: torch.Tensor | None, +) -> None: + head_dim = int(q.shape[-1]) + if head_dim not in {16, 32, 64, 128, 256}: + raise ValueError(f"unsupported context-stable decode head_dim={head_dim}") + if q.dtype != k.dtype or k.dtype != v.dtype: + raise TypeError("query, key, and value tensors must have the same dtype") + if int(q.shape[1]) % int(k.shape[1]): + raise ValueError("query head count must be divisible by KV head count") + if q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1: + raise ValueError("query, key, and value head dimensions must be contiguous") + if k.stride() != v.stride(): + raise ValueError("key and value cache layouts must match") + if active_slots.dim() != 2 or active_slots.stride(-1) != 1: + raise ValueError("active_slots must be a contiguous 2D slot table") + if req_indices.stride(0) != 1 or context_lens.stride(0) != 1: + raise ValueError("request indices and context lengths must be contiguous") + expected_workspace = (int(q.shape[0]), int(q.shape[1])) + if tuple(mid_o.shape[:2]) != expected_workspace or tuple(mid_lse.shape[:2]) != expected_workspace: + raise ValueError( + "workspace batch/head dimensions do not match query: " + f"q={tuple(q.shape)} mid_o={tuple(mid_o.shape)} mid_lse={tuple(mid_lse.shape)}" + ) + if int(mid_o.shape[2]) != int(mid_lse.shape[2]): + raise ValueError("workspace split dimensions must match") + if attn_score is not None and attn_score.dim() not in {2, 3}: + raise ValueError("attention score output must be 2D or 3D") + + +@torch.no_grad() +def fixed_grid_flash_decode_stage2( + mid_o: torch.Tensor, + mid_lse: torch.Tensor, + context_lens: torch.Tensor, + output: torch.Tensor, + output_lse: torch.Tensor, + *, + target_tokens_per_split: int, + num_warps: int | None = None, + num_stages: int = 2, +) -> None: + """Reduce a fixed split envelope using device-resident effective lengths.""" + if mid_o.dim() != 4 or mid_lse.dim() != 3: + raise ValueError("fixed-grid decode workspaces must be rank 4 and rank 3") + if tuple(mid_o.shape[:3]) != tuple(mid_lse.shape): + raise ValueError( + "fixed-grid decode workspace shapes do not match: " + f"mid_o={tuple(mid_o.shape)} mid_lse={tuple(mid_lse.shape)}" + ) + batch, num_heads, max_kv_splits, head_dim = map(int, mid_o.shape) + if tuple(output.shape) != (batch, num_heads, head_dim): + raise ValueError( + "fixed-grid decode output shape does not match its workspace: " + f"output={tuple(output.shape)} expected={(batch, num_heads, head_dim)}" + ) + if tuple(output_lse.shape) != (num_heads, batch): + raise ValueError( + "fixed-grid decode LSE output must be [heads, batch], got " + f"{tuple(output_lse.shape)}." + ) + if context_lens.dtype != torch.int32 or context_lens.stride(0) != 1: + raise TypeError("fixed-grid decode context_lens must be contiguous int32") + if int(context_lens.numel()) != batch: + raise ValueError("fixed-grid decode expects one context length per batch row") + if max_kv_splits <= 0 or int(target_tokens_per_split) <= 0: + raise ValueError("fixed-grid decode split envelope must be positive") + if head_dim not in {16, 32, 64, 128, 256}: + raise ValueError(f"unsupported fixed-grid decode head_dim={head_dim}") + if output_lse.dtype != torch.float32 or output_lse.device != output.device: + raise TypeError("fixed-grid decode LSE output must be FP32 on the output device") + if num_warps is None: + num_warps = 8 if head_dim == 256 else 4 + if int(num_warps) <= 0 or int(num_stages) <= 0: + raise ValueError("fixed-grid decode stage2 warps/stages must be positive") + + _paged_decode_stage2[(batch, num_heads)]( + context_lens, + mid_o, + mid_lse, + output, + output_lse, + mid_o.stride(0), + mid_o.stride(1), + mid_o.stride(2), + mid_lse.stride(0), + mid_lse.stride(1), + mid_lse.stride(2), + output.stride(0), + output.stride(1), + output_lse.stride(0), + output_lse.stride(1), + HEAD_DIM=head_dim, + MAX_KV_SPLITS=max_kv_splits, + MAX_EFFECTIVE_SPLITS=max_kv_splits, + TARGET_TOKENS_PER_SPLIT=int(target_tokens_per_split), + num_warps=int(num_warps), + num_stages=int(num_stages), + ) + + +@torch.no_grad() +def paged_flash_decode( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + mid_o: torch.Tensor, + mid_lse: torch.Tensor, + *, + attn_score: torch.Tensor | None = None, + softmax_scale: float | None = None, + target_tokens_per_split: int, + block_n: int = 32, + num_warps: int = 4, + num_stages: int = 2, + stage2_num_warps: int | None = None, + stage2_num_stages: int = 2, + return_softmax_lse: bool = False, + output_lse: torch.Tensor | None = None, + output: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run fixed-grid split-KV decode for MHA or GQA.""" + _check_inputs( + q, + k, + v, + active_slots, + req_indices, + context_lens, + mid_o, + mid_lse, + attn_score, + ) + max_kv_splits = int(mid_o.shape[2]) + if max_kv_splits <= 0 or target_tokens_per_split <= 0: + raise ValueError("split count and target tokens per split must be positive") + if block_n not in {16, 32, 64, 128}: + raise ValueError(f"unsupported BLOCK_N={block_n}") + if num_warps <= 0 or num_stages <= 0 or stage2_num_stages <= 0: + raise ValueError("Triton launch warps/stages must be positive") + + batch, num_heads, head_dim = map(int, q.shape) + if softmax_scale is None: + softmax_scale = 1.0 / (head_dim**0.5) + if softmax_scale <= 0: + raise ValueError("softmax_scale must be positive") + if stage2_num_warps is None: + stage2_num_warps = 8 if head_dim == 256 else 4 + group_size = num_heads // int(k.shape[1]) + max_effective_splits = max_kv_splits + score = mid_lse if attn_score is None else attn_score + if attn_score is None: + score_strides = (0, 0, 0) + elif attn_score.dim() == 3: + score_strides = tuple(int(stride) for stride in attn_score.stride()) + else: + score_strides = (int(attn_score.stride(0)), 0, int(attn_score.stride(1))) + stage1_args = ( + q, + k, + v, + softmax_scale, + active_slots, + req_indices, + context_lens, + mid_o, + mid_lse, + score, + active_slots.stride(0), + active_slots.stride(1), + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + mid_o.stride(0), + mid_o.stride(1), + mid_o.stride(2), + mid_lse.stride(0), + mid_lse.stride(1), + mid_lse.stride(2), + *score_strides, + ) + stage1_meta = dict( + GQA_GROUP_SIZE=group_size, + HEAD_DIM=head_dim, + BLOCK_N=block_n, + MAX_KV_SPLITS=max_kv_splits, + MAX_EFFECTIVE_SPLITS=max_effective_splits, + TARGET_TOKENS_PER_SPLIT=target_tokens_per_split, + SCORE_MODE=0 if attn_score is None else attn_score.dim(), + num_warps=num_warps, + num_stages=num_stages, + ) + if group_size > 1: + _paged_grouped_decode_stage1[ + (batch, int(k.shape[1]), max_kv_splits) + ]( + *stage1_args, + QUERY_HEAD_BLOCK=max(16, triton.next_power_of_2(group_size)), + **stage1_meta, + ) + else: + _paged_decode_stage1[(batch, num_heads, max_kv_splits)]( + *stage1_args, + **stage1_meta, + ) + + if output is None: + output = torch.empty_like(q) + elif tuple(output.shape) != tuple(q.shape): + raise ValueError( + "decode output workspace must match Q shape, got " + f"output={tuple(output.shape)} q={tuple(q.shape)}" + ) + elif output.dtype != q.dtype or output.device != q.device: + raise TypeError("decode output workspace must match Q dtype and device") + if output_lse is None: + output_lse = torch.empty( + (num_heads, batch), dtype=torch.float32, device=q.device + ) + elif tuple(output_lse.shape) != (num_heads, batch): + raise ValueError( + "softmax LSE workspace must be [query_heads, batch], got " + f"{tuple(output_lse.shape)}." + ) + if output_lse.dtype != torch.float32 or output_lse.device != q.device: + raise TypeError("softmax LSE workspace must be FP32 on the query device") + fixed_grid_flash_decode_stage2( + mid_o, + mid_lse, + context_lens, + output, + output_lse, + target_tokens_per_split=target_tokens_per_split, + num_warps=stage2_num_warps, + num_stages=stage2_num_stages, + ) + return (output, output_lse) if return_softmax_lse else output diff --git a/src/sparsevllm/kernels/triton/sglang_gemma4_decode_attention.py b/src/sparsevllm/kernels/triton/sglang_gemma4_decode_attention.py new file mode 100644 index 00000000..1ed4091f --- /dev/null +++ b/src/sparsevllm/kernels/triton/sglang_gemma4_decode_attention.py @@ -0,0 +1,625 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""SGLang fixed-grid Triton decode adapted for Sparse-vLLM Gemma 4. + +Source: sglang/srt/layers/attention/triton_ops/decode_attention.py at +ed0a62e4dd006132a2c6434378962528f010c906. + +The kernel topology and split scheduling follow SGLang. The local changes are +limited to Sparse-vLLM's two-dimensional slot table, Gemma 4 sliding-window +coordinates, and the optional raw-QK score output used by sparse methods. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +_MIN_BLOCK_KV = tl.constexpr(32) + + +@triton.jit +def _get_num_kv_splits( + num_kv_splits, + context_lens, + num_seq, + num_heads: tl.constexpr, + num_kv_heads: tl.constexpr, + max_kv_splits: tl.constexpr, + multi_processor_count: tl.constexpr, + window: tl.constexpr, + max_num_seq: tl.constexpr, +): + offsets = tl.arange(0, max_num_seq) + mask = offsets < num_seq + seq_lens = tl.load(context_lens + offsets, mask=mask, other=0) + if window > 0: + seq_lens = tl.minimum(seq_lens, window) + max_seq_len = tl.max(seq_lens) + seq_lens_for_min = tl.load( + context_lens + offsets, mask=mask, other=max_seq_len + ) + if window > 0: + seq_lens_for_min = tl.minimum(seq_lens_for_min, window) + min_seq_len = tl.min(seq_lens_for_min) + if max_seq_len * 8 < min_seq_len * 10: + min_seq_len = max_seq_len + + split_cap_by_lengths = tl.minimum( + tl.cdiv(max_seq_len, min_seq_len), max_kv_splits + ) + chunk_by_lengths = tl.cdiv(max_seq_len, split_cap_by_lengths) + + extended_len = tl.cast(max_seq_len, tl.float32) / 64.0 + extended_cores = tl.cast( + multi_processor_count * tl.maximum(tl.log2(extended_len), 1.0), tl.int32 + ) + group_size: tl.constexpr = num_heads // num_kv_heads + if group_size == 1: + token_grid = num_seq * num_heads + else: + block_h: tl.constexpr = min(16, group_size) + token_grid = num_seq * tl.cdiv(num_heads, block_h) + split_cap_by_cores = tl.minimum( + tl.cdiv(extended_cores, token_grid), max_kv_splits + ) + chunk_by_cores = tl.cdiv(max_seq_len, split_cap_by_cores) + splits = tl.maximum( + tl.cdiv(seq_lens, chunk_by_lengths), tl.cdiv(seq_lens, chunk_by_cores) + ) + # Every split consumed by stage2 must have at least one 32-token block. + # Otherwise stage1 leaves that split's workspace uninitialized. + splits = tl.maximum( + 1, + tl.minimum(splits, tl.cdiv(seq_lens, _MIN_BLOCK_KV)), + ) + tl.store(num_kv_splits + offsets, splits, mask=mask) + + +@triton.jit +def _decode_stage1_normal( + q, + k, + v, + active_slots, + req_indices, + context_lens, + mid_output, + mid_lse, + num_kv_splits, + attn_score, + stride_qb, + stride_qh, + stride_kt, + stride_kh, + stride_vt, + stride_vh, + stride_sb, + stride_ss, + stride_mob, + stride_moh, + stride_mos, + stride_mlb, + stride_mlh, + stride_mls, + stride_asb, + stride_ash, + stride_asl, + head_dim: tl.constexpr, + block_dim: tl.constexpr, + block_n: tl.constexpr, + window: tl.constexpr, + score_mode: tl.constexpr, +): + batch = tl.program_id(0) + head = tl.program_id(1) + split = tl.program_id(2) + dims = tl.arange(0, block_dim) + dim_mask = dims < head_dim + sequence_len = tl.load(context_lens + batch) + visible_len = sequence_len + visible_start = 0 + if window > 0: + visible_len = tl.minimum(sequence_len, window) + visible_start = sequence_len - visible_len + splits = tl.load(num_kv_splits + batch) + tokens_per_split = ( + tl.cdiv(tl.cdiv(visible_len, splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV + ) + split_start = split * tokens_per_split + split_end = tl.minimum(split_start + tokens_per_split, visible_len) + + max_logit = -float("inf") + exp_sum = 0.0 + accumulator = tl.zeros((block_dim,), tl.float32) + if split_end > split_start: + query = tl.load( + q + batch * stride_qb + head * stride_qh + dims, + mask=dim_mask, + other=0.0, + ) + request = tl.load(req_indices + batch) + for start in range(split_start, split_end, block_n): + local_positions = start + tl.arange(0, block_n) + positions = visible_start + local_positions + position_mask = local_positions < split_end + slots = tl.load( + active_slots + request * stride_sb + positions * stride_ss, + mask=position_mask, + other=0, + ) + keys = tl.load( + k + slots[:, None] * stride_kt + head * stride_kh + dims[None, :], + mask=position_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + logits = tl.sum(query[None, :] * keys, axis=1) + if score_mode == 3: + tl.store( + attn_score + + batch * stride_asb + + head * stride_ash + + positions * stride_asl, + logits, + mask=position_mask, + ) + elif score_mode == 2: + tl.atomic_max( + attn_score + batch * stride_asb + positions * stride_asl, + logits, + mask=position_mask, + ) + logits = tl.where(position_mask, logits, -float("inf")) + values = tl.load( + v + slots[:, None] * stride_vt + head * stride_vh + dims[None, :], + mask=position_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + next_max = tl.maximum(tl.max(logits, axis=0), max_logit) + old_scale = tl.exp(max_logit - next_max) + probabilities = tl.exp(logits - next_max) + accumulator *= old_scale + accumulator += tl.sum(probabilities[:, None] * values, axis=0) + exp_sum = exp_sum * old_scale + tl.sum(probabilities, axis=0) + max_logit = next_max + + mid_offset = batch * stride_mob + head * stride_moh + split * stride_mos + tl.store( + mid_output + mid_offset + dims, + accumulator / exp_sum, + mask=dim_mask, + ) + tl.store( + mid_lse + batch * stride_mlb + head * stride_mlh + split * stride_mls, + max_logit + tl.log(exp_sum), + ) + + +@triton.jit +def _decode_stage1_grouped( + q, + k, + v, + active_slots, + req_indices, + context_lens, + mid_output, + mid_lse, + num_kv_splits, + attn_score, + stride_qb, + stride_qh, + stride_kt, + stride_kh, + stride_vt, + stride_vh, + stride_sb, + stride_ss, + stride_mob, + stride_moh, + stride_mos, + stride_mlb, + stride_mlh, + stride_mls, + stride_asb, + stride_ash, + stride_asl, + group_size: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + block_dim: tl.constexpr, + block_n: tl.constexpr, + block_h: tl.constexpr, + window: tl.constexpr, + score_mode: tl.constexpr, +): + batch = tl.program_id(0) + head_group = tl.program_id(1) + split = tl.program_id(2) + valid_block_h: tl.constexpr = min(block_h, group_size) + kv_head = head_group // tl.cdiv(group_size, block_h) + heads = head_group * valid_block_h + tl.arange(0, block_h) + head_mask = (heads < (head_group + 1) * valid_block_h) & (heads < num_heads) + dims = tl.arange(0, block_dim) + dim_mask = dims < head_dim + + sequence_len = tl.load(context_lens + batch) + visible_len = sequence_len + visible_start = 0 + if window > 0: + visible_len = tl.minimum(sequence_len, window) + visible_start = sequence_len - visible_len + splits = tl.load(num_kv_splits + batch) + tokens_per_split = ( + tl.cdiv(tl.cdiv(visible_len, splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV + ) + split_start = split * tokens_per_split + split_end = tl.minimum(split_start + tokens_per_split, visible_len) + + max_logit = tl.full((block_h,), -float("inf"), tl.float32) + exp_sum = tl.zeros((block_h,), tl.float32) + accumulator = tl.zeros((block_h, block_dim), tl.float32) + if split_end > split_start: + query = tl.load( + q + batch * stride_qb + heads[:, None] * stride_qh + dims[None, :], + mask=head_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + request = tl.load(req_indices + batch) + key_base = kv_head * stride_kh + dims[:, None] + value_base = kv_head * stride_vh + dims[None, :] + for start in tl.range(split_start, split_end, block_n): + local_positions = start + tl.arange(0, block_n) + positions = visible_start + local_positions + position_mask = local_positions < split_end + slots = tl.load( + active_slots + request * stride_sb + positions * stride_ss, + mask=position_mask, + other=0, + ) + keys = tl.load( + k + slots[None, :] * stride_kt + key_base, + mask=dim_mask[:, None] & position_mask[None, :], + other=0.0, + ) + logits = tl.dot(query.to(k.dtype.element_ty), keys) + if score_mode == 3: + tl.store( + attn_score + + batch * stride_asb + + heads[:, None] * stride_ash + + positions[None, :] * stride_asl, + logits, + mask=head_mask[:, None] & position_mask[None, :], + ) + elif score_mode == 2: + reduced_logits = tl.max( + tl.where(head_mask[:, None], logits, -float("inf")), axis=0 + ) + tl.atomic_max( + attn_score + batch * stride_asb + positions * stride_asl, + reduced_logits, + mask=position_mask, + ) + logits = tl.where( + head_mask[:, None] & position_mask[None, :], + logits, + -float("inf"), + ) + values = tl.load( + v + slots[:, None] * stride_vt + value_base, + mask=position_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + next_max = tl.maximum(tl.max(logits, axis=1), max_logit) + old_scale = tl.exp(max_logit - next_max) + probabilities = tl.exp(logits - next_max[:, None]) + accumulator *= old_scale[:, None] + accumulator += tl.dot(probabilities.to(values.dtype), values) + exp_sum = exp_sum * old_scale + tl.sum(probabilities, axis=1) + max_logit = next_max + + mid_offsets = ( + batch * stride_mob + + heads[:, None] * stride_moh + + split * stride_mos + + dims[None, :] + ) + lse_offsets = batch * stride_mlb + heads * stride_mlh + split * stride_mls + tl.store( + mid_output + mid_offsets, + accumulator / exp_sum[:, None], + mask=head_mask[:, None] & dim_mask[None, :], + ) + tl.store(mid_lse + lse_offsets, max_logit + tl.log(exp_sum), mask=head_mask) + + +@triton.jit +def _decode_stage2( + mid_output, + mid_lse, + output, + context_lens, + num_kv_splits, + stride_mob, + stride_moh, + stride_mos, + stride_mlb, + stride_mlh, + stride_mls, + stride_ob, + stride_oh, + head_dim: tl.constexpr, + block_dim: tl.constexpr, + max_kv_splits: tl.constexpr, + window: tl.constexpr, +): + batch = tl.program_id(0) + head = tl.program_id(1) + dims = tl.arange(0, block_dim) + dim_mask = dims < head_dim + sequence_len = tl.load(context_lens + batch) + visible_len = sequence_len + if window > 0: + visible_len = tl.minimum(sequence_len, window) + splits = tl.load(num_kv_splits + batch) + tokens_per_split = ( + tl.cdiv(tl.cdiv(visible_len, splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV + ) + + max_lse = -float("inf") + exp_sum = 0.0 + accumulator = tl.zeros((block_dim,), tl.float32) + value_offset = batch * stride_mob + head * stride_moh + dims + lse_offset = batch * stride_mlb + head * stride_mlh + for split in tl.range(0, max_kv_splits, num_stages=2): + split_start = tokens_per_split * split + split_end = tl.minimum(split_start + tokens_per_split, visible_len) + if split_end > split_start: + value = tl.load( + mid_output + value_offset + split * stride_mos, + mask=dim_mask, + other=0.0, + ) + lse = tl.load(mid_lse + lse_offset + split * stride_mls) + next_max = tl.maximum(lse, max_lse) + old_scale = tl.exp(max_lse - next_max) + split_scale = tl.exp(lse - next_max) + accumulator = accumulator * old_scale + value * split_scale + exp_sum = exp_sum * old_scale + split_scale + max_lse = next_max + tl.store( + output + batch * stride_ob + head * stride_oh + dims, + accumulator / exp_sum, + mask=dim_mask, + ) + + +def _check_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + mid_output: torch.Tensor, + mid_lse: torch.Tensor, + num_kv_splits: torch.Tensor, + attn_score: torch.Tensor | None, +) -> None: + tensors = ( + q, + k, + v, + active_slots, + req_indices, + context_lens, + mid_output, + mid_lse, + num_kv_splits, + ) + if not all(tensor.is_cuda and tensor.device == q.device for tensor in tensors): + raise TypeError("Gemma 4 decode tensors must share one CUDA device.") + if attn_score is not None and ( + not attn_score.is_cuda or attn_score.device != q.device + ): + raise TypeError("Gemma 4 attention scores must share the Q/K/V CUDA device.") + if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape: + raise ValueError("Gemma 4 decode requires matching rank-3 Q/K/V tensors.") + head_dim = int(q.shape[-1]) + if head_dim not in {256, 512} or int(k.shape[-1]) != head_dim: + raise ValueError(f"Gemma 4 decode requires head_dim 256 or 512, got {head_dim}.") + if q.dtype not in {torch.float16, torch.bfloat16} or any( + tensor.dtype != q.dtype for tensor in (k, v) + ): + raise TypeError("Gemma 4 decode requires matching FP16 or BF16 Q/K/V.") + if any(tensor.stride(-1) != 1 for tensor in (q, k, v)): + raise ValueError("Gemma 4 Q/K/V head dimensions must be contiguous.") + if int(q.shape[1]) % int(k.shape[1]): + raise ValueError("Gemma 4 query heads must be divisible by KV heads.") + if active_slots.ndim != 2 or active_slots.stride(-1) != 1: + raise ValueError("Gemma 4 active_slots must be a contiguous 2D slot table.") + if active_slots.dtype not in {torch.int32, torch.int64}: + raise TypeError("Gemma 4 active_slots must use int32 or int64 indices.") + batch, heads = int(q.shape[0]), int(q.shape[1]) + if req_indices.shape != (batch,) or context_lens.shape != (batch,): + raise ValueError("Gemma 4 request indices and context lengths must match batch size.") + if req_indices.dtype not in {torch.int32, torch.int64}: + raise TypeError("Gemma 4 request indices must use int32 or int64.") + if context_lens.dtype not in {torch.int32, torch.int64}: + raise TypeError("Gemma 4 context lengths must use int32 or int64.") + if mid_output.shape[:2] != (batch, heads) or mid_lse.shape[:2] != (batch, heads): + raise ValueError("Gemma 4 workspace batch/head dimensions must match query.") + if mid_output.dtype != torch.float32 or mid_lse.dtype != torch.float32: + raise TypeError("Gemma 4 decode workspace must use FP32 tensors.") + if mid_output.shape[2] != mid_lse.shape[2] or mid_output.shape[-1] != head_dim: + raise ValueError("Gemma 4 workspace split/head dimensions do not match.") + if num_kv_splits.dtype != torch.int32 or num_kv_splits.shape != (batch,): + raise ValueError("Gemma 4 num_kv_splits must be a batch-sized int32 tensor.") + if attn_score is not None and attn_score.dim() not in {2, 3}: + raise ValueError("Gemma 4 attention scores must be rank 2 or 3.") + if attn_score is not None: + expected_prefix = (batch, heads) if attn_score.dim() == 3 else (batch,) + if attn_score.shape[:-1] != expected_prefix: + raise ValueError("Gemma 4 attention score batch/head dimensions do not match.") + if attn_score.dtype != torch.float32: + raise TypeError("Gemma 4 attention scores must use FP32.") + + +@torch.no_grad() +def sglang_gemma4_decode( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + mid_output: torch.Tensor, + mid_lse: torch.Tensor, + num_kv_splits: torch.Tensor, + *, + sliding_window: int | None, + multi_processor_count: int, + attn_score: torch.Tensor | None = None, +) -> torch.Tensor: + """Run SGLang's context-stable fixed-grid Gemma 4 decode.""" + _check_inputs( + q, + k, + v, + active_slots, + req_indices, + context_lens, + mid_output, + mid_lse, + num_kv_splits, + attn_score, + ) + batch, num_heads, head_dim = map(int, q.shape) + num_kv_heads = int(k.shape[1]) + max_kv_splits = int(mid_output.shape[2]) + if max_kv_splits <= 0 or int(multi_processor_count) <= 0: + raise ValueError( + "Gemma 4 split count and multi-processor count must be positive." + ) + max_num_seq = 256 if batch < 256 else triton.next_power_of_2(batch) + window = int(sliding_window or 0) + _get_num_kv_splits[(1,)]( + num_kv_splits, + context_lens, + batch, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + max_kv_splits=max_kv_splits, + multi_processor_count=int(multi_processor_count), + window=window, + max_num_seq=max_num_seq, + ) + + if attn_score is None: + score = mid_lse + score_mode = 0 + score_strides = (0, 0, 0) + else: + score = attn_score + score_mode = attn_score.dim() + score_strides = ( + int(attn_score.stride(0)), + int(attn_score.stride(1)) if score_mode == 3 else 0, + int(attn_score.stride(-1)), + ) + block_dim = triton.next_power_of_2(head_dim) + common_args = ( + q, + k, + v, + active_slots, + req_indices, + context_lens, + mid_output, + mid_lse, + num_kv_splits, + score, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + active_slots.stride(0), + active_slots.stride(1), + mid_output.stride(0), + mid_output.stride(1), + mid_output.stride(2), + mid_lse.stride(0), + mid_lse.stride(1), + mid_lse.stride(2), + *score_strides, + ) + group_size = num_heads // num_kv_heads + if group_size == 1: + _decode_stage1_normal[(batch, num_heads, max_kv_splits)]( + *common_args, + head_dim=head_dim, + block_dim=block_dim, + block_n=64, + window=window, + score_mode=score_mode, + num_warps=4, + num_stages=2, + ) + else: + block_h = 16 + _decode_stage1_grouped[ + (batch, triton.cdiv(num_heads, min(block_h, group_size)), max_kv_splits) + ]( + *common_args, + group_size=group_size, + num_heads=num_heads, + head_dim=head_dim, + block_dim=block_dim, + block_n=32, + block_h=block_h, + window=window, + score_mode=score_mode, + num_warps=4, + num_stages=2, + ) + + output = torch.empty_like(q) + _decode_stage2[(batch, num_heads)]( + mid_output, + mid_lse, + output, + context_lens, + num_kv_splits, + mid_output.stride(0), + mid_output.stride(1), + mid_output.stride(2), + mid_lse.stride(0), + mid_lse.stride(1), + mid_lse.stride(2), + output.stride(0), + output.stride(1), + head_dim=head_dim, + block_dim=block_dim, + max_kv_splits=max_kv_splits, + window=window, + num_warps=4, + num_stages=2, + ) + return output + + +__all__ = ["sglang_gemma4_decode"] diff --git a/src/sparsevllm/layers/attention.py b/src/sparsevllm/layers/attention.py index a817010d..f9454bee 100644 --- a/src/sparsevllm/layers/attention.py +++ b/src/sparsevllm/layers/attention.py @@ -230,14 +230,27 @@ def forward( num_seq_blocks = ( max_len_in_batch + block_seq - 1 ) // block_seq - mid_o, mid_o_logexpsum = get_decode_workspace( - context, - batch_size, - self.num_heads, - num_seq_blocks, - self.head_dim, - q.device, + workspace_provider = getattr( + self.attention_backend, + "get_decode_workspace", + None, ) + if callable(workspace_provider): + mid_o, mid_o_logexpsum = workspace_provider( + batch_size=batch_size, + num_heads=self.num_heads, + head_dim=self.head_dim, + device=q.device, + ) + else: + mid_o, mid_o_logexpsum = get_decode_workspace( + context, + batch_size, + self.num_heads, + num_seq_blocks, + self.head_dim, + q.device, + ) o = self.attention_backend.run_decode( q, decode_view, diff --git a/src/sparsevllm/method_registry.py b/src/sparsevllm/method_registry.py index 925e3e92..335f5fae 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -96,13 +96,14 @@ class SparsePrefillAttentionContract: {"snapkv", "pyramidkv", "h2o", "rkv"} ) -# These methods can request a score-producing decode launch on at least one -# layer or decode step. The answer is deliberately static so Provider -# selection happens before CUDA Graph capture and never changes in run(). -_DECODE_ATTENTION_SCORE_METHODS = frozenset( - {"pyramidkv", "omnikv", "skipkv", "deltakv"} -) - +# Static method score contracts let providers bind before CUDA Graph capture +# instead of changing the score-producing implementation during replay. +_DECODE_ATTENTION_SCORE_KINDS = { + "pyramidkv": AttentionScoreKind.RAW_QK_REDUCED, + "omnikv": AttentionScoreKind.RAW_QK_PER_HEAD, + "skipkv": AttentionScoreKind.RAW_QK_PER_HEAD, + "deltakv": AttentionScoreKind.RAW_QK_PER_HEAD, +} def sparse_prefill_attention_contract( method: str | None, @@ -146,10 +147,26 @@ def h2o_uses_fused_prefill_score(config) -> bool: def sparse_decode_attention_requires_scores(method: str | None) -> bool: """Return whether a prepared decode implementation must support scores.""" + return sparse_decode_attention_score_kind(method) is not AttentionScoreKind.NONE + + +def sparse_decode_attention_score_kind( + method: str | None, +) -> AttentionScoreKind: + """Return the score representation consumed by sparse decode logic. + + OmniKV, SkipKV, and DeltaKV normalize each head in ``SparseController`` + before reducing across heads, so providers must preserve raw per-head QK. + PyramidKV consumes the existing fused head-reduced raw-QK representation. + """ + normalized = normalize_sparse_method(method) if normalized not in CANONICAL_SPARSE_METHODS: raise ValueError(f"Unknown sparse method {normalized!r}.") - return normalized in _DECODE_ATTENTION_SCORE_METHODS + return _DECODE_ATTENTION_SCORE_KINDS.get( + normalized, + AttentionScoreKind.NONE, + ) _MOE_SPARSE_METHODS = frozenset( @@ -249,6 +266,35 @@ def sparse_decode_attention_requires_scores(method: str | None) -> bool: "skipkv", } + +def decode_sparse_long_text_threshold( + method: str, + *, + num_sink_tokens: int, + decode_keep_tokens: int, + num_recent_tokens: int, +) -> int: + """Return the shared decode boundary between short and sparse graph families.""" + method = str(method or "") + if not method: + return 0 + if method in {"streamingllm", "attention-sink", "attention_sink"}: + return int(num_sink_tokens) + int(num_recent_tokens) + return ( + int(num_sink_tokens) + + int(decode_keep_tokens) + + int(num_recent_tokens) + ) + + +def decode_graph_path_id(method: str, is_long_text: bool) -> str: + """Identify one graph-stable decode topology family.""" + method = str(method or "") + if not method: + return "dense" + return "long" if is_long_text else "short" + + _DEFAULT_PREFILL_POLICY_BY_METHOD = { "": PREFILL_POLICY_ALL_CHUNKED, "streamingllm": PREFILL_POLICY_ALL_CHUNKED, diff --git a/src/sparsevllm/models/attention_runtime.py b/src/sparsevllm/models/attention_runtime.py index 226dd996..d426748d 100644 --- a/src/sparsevllm/models/attention_runtime.py +++ b/src/sparsevllm/models/attention_runtime.py @@ -8,7 +8,6 @@ sparse_decode_attention_requires_scores, sparse_prefill_attention_contract, ) -from sparsevllm.operators.moe import model_activation_dtype from sparsevllm.operators.decode_attention import ( DecodeAttentionOpSpec, PreparedDecodeAttentionOp, @@ -19,6 +18,7 @@ FullAttentionProvider, prepare_full_attention_provider, ) +from sparsevllm.operators.moe import model_activation_dtype from sparsevllm.operators.prefill_attention import ( PreparedPrefillAttentionOp, PrefillAttentionOpSpec, @@ -164,6 +164,42 @@ def build_mha_decode_attention_spec( h2o_layerwise_probability_scores=( normalized_method == "h2o" and requires_decode_scores ), + batch_only_cuda_graph=( + bool(cuda_graph) + and str( + getattr( + runtime_config, + "decode_graph_shape_policy", + "batch_only", + ) + ) + == "batch_only" + ), + context_capacity=int(getattr(runtime_config, "max_model_len", 0) or 0) + or None, + may_use_full_layer_kivi_int4=( + normalized_method == "deltakv" + and int( + getattr(runtime_config, "full_layer_kv_quant_bits", 0) or 0 + ) + == 4 + and bool( + getattr(runtime_config, "enable_full_layer_kivi_quant", True) + ) + ), + full_layer_kivi_decode_block_seq=int( + getattr(runtime_config, "full_layer_kivi_decode_block_seq", 256) + or 256 + ), + full_layer_kivi_decode_block_n=int( + getattr(runtime_config, "full_layer_kivi_decode_block_n", 16) or 16 + ), + full_layer_kivi_decode_num_warps=int( + getattr(runtime_config, "full_layer_kivi_decode_num_warps", 2) or 2 + ), + full_layer_kivi_decode_num_stages=int( + getattr(runtime_config, "full_layer_kivi_decode_num_stages", 3) or 3 + ), ) diff --git a/src/sparsevllm/models/gdn_runtime.py b/src/sparsevllm/models/gdn_runtime.py index be2e28c5..4b1d1d7d 100644 --- a/src/sparsevllm/models/gdn_runtime.py +++ b/src/sparsevllm/models/gdn_runtime.py @@ -46,6 +46,17 @@ def build_gated_delta_rule_op( activation_dtype=model_activation_dtype(config), recurrent_state_dtype=recurrent_state_dtype, cuda_graph_decode=bool(cuda_graph), + batch_only_cuda_graph=( + bool(cuda_graph) + and str( + getattr( + config, + "decode_graph_shape_policy", + "batch_only", + ) + ) + == "batch_only" + ), ), device_index=int(device.index or 0), ) diff --git a/src/sparsevllm/models/gemma4.py b/src/sparsevllm/models/gemma4.py index 031b918c..cbf086c2 100644 --- a/src/sparsevllm/models/gemma4.py +++ b/src/sparsevllm/models/gemma4.py @@ -696,6 +696,20 @@ def build_runtime_kwargs( head_dims=head_dims, cuda_graph=bool(engine_config.decode_graph), attention_contracts=attention_contracts, + max_batch_size=int(getattr(engine_config, "max_decoding_seqs", 1)), + batch_only_cuda_graph=( + bool(engine_config.decode_graph) + and str( + getattr( + engine_config, + "decode_graph_shape_policy", + "batch_only", + ) + ) + == "batch_only" + ), + context_capacity=int(getattr(engine_config, "max_model_len", 0) or 0) + or None, ), device_index=device.index, ) diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index b0bac9b3..c9400fe7 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -28,7 +28,8 @@ from sparsevllm.layers.mla_attention import MLAAttention from sparsevllm.layers.packed_moe import PackedMoeExperts from sparsevllm.layers.rotary_embedding import RotaryEmbedding, get_rope -from sparsevllm.method_registry import sparse_decode_attention_requires_scores +from sparsevllm.method_registry import sparse_decode_attention_score_kind +from sparsevllm.operators.attention_capabilities import AttentionScoreKind from sparsevllm.models.qwen3 import Qwen3MLP from sparsevllm.operators.mla_attention import MlaAttentionOpSpec from sparsevllm.operators.all_reduce import ( @@ -120,8 +121,9 @@ def build_glm4_moe_lite_mla_attention( max_batch_size: int, prefill_workspace_bytes: int, decode_graph: bool, + context_capacity: int, projection_chunk_size: int, - may_require_attention_scores: bool = False, + score_output: AttentionScoreKind = AttentionScoreKind.NONE, ) -> MLAAttention: """Bind the one process-local MLA operator from explicit runtime inputs.""" @@ -137,7 +139,19 @@ def build_glm4_moe_lite_mla_attention( cache_dtype=activation_dtype, tp_size=int(parallel_context.attention_tp_size), cuda_graph=bool(decode_graph), - may_require_attention_scores=bool(may_require_attention_scores), + score_output=score_output, + batch_only_cuda_graph=( + bool(decode_graph) + and str( + getattr( + config, + "decode_graph_shape_policy", + "batch_only", + ) + ) + == "batch_only" + ), + context_capacity=int(context_capacity), ) return MLAAttention.bind( spec=spec, @@ -892,11 +906,10 @@ def build_runtime_kwargs( ), prefill_workspace_bytes=engine_config.mla_prefill_workspace_bytes, decode_graph=decode_graph, + context_capacity=int(engine_config.max_model_len), projection_chunk_size=engine_config.mlp_chunk_size, - may_require_attention_scores=( - sparse_decode_attention_requires_scores( - engine_config.sparse_method - ) + score_output=sparse_decode_attention_score_kind( + engine_config.sparse_method ), ), "mlp_chunk_size": engine_config.mlp_chunk_size, diff --git a/src/sparsevllm/models/spec.py b/src/sparsevllm/models/spec.py index 981efc28..12c89a74 100644 --- a/src/sparsevllm/models/spec.py +++ b/src/sparsevllm/models/spec.py @@ -15,6 +15,8 @@ class ModelSpec: mixed_attention: bool = False allow_raw_config: bool = False supports_tiny_random: bool = True + supports_quantized_tiny_random: bool = False + tiny_random_requires_standard_head_shape: bool = True supports_expert_parallel: bool = False supports_outer_tp_moe: bool = False outer_tp_moe_config_field: str | None = None @@ -163,6 +165,8 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: "minimax_m2": ModelSpec( "MiniMax M2.7", requires_fp8=True, + supports_quantized_tiny_random=True, + tiny_random_requires_standard_head_shape=False, supports_expert_parallel=True, supports_outer_tp_moe=True, runtime_class_name="MiniMaxM2ForCausalLM", @@ -173,6 +177,7 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: ), "glm4_moe_lite": ModelSpec( "GLM-4.7-Flash", + tiny_random_requires_standard_head_shape=False, supports_expert_parallel=True, supports_outer_tp_moe=True, runtime_class_name="Glm4MoeLiteForCausalLM", diff --git a/src/sparsevllm/operators/decode_attention.py b/src/sparsevllm/operators/decode_attention.py index 2b5f3612..f3c7aa1c 100644 --- a/src/sparsevllm/operators/decode_attention.py +++ b/src/sparsevllm/operators/decode_attention.py @@ -1,7 +1,7 @@ from __future__ import annotations import math -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any import torch @@ -87,6 +87,13 @@ class DecodeAttentionOpSpec: layer_varying_page_table: bool = False cuda_graph: bool = True h2o_layerwise_probability_scores: bool = False + batch_only_cuda_graph: bool = False + context_capacity: int | None = None + may_use_full_layer_kivi_int4: bool = False + full_layer_kivi_decode_block_seq: int = 256 + full_layer_kivi_decode_block_n: int = 16 + full_layer_kivi_decode_num_warps: int = 2 + full_layer_kivi_decode_num_stages: int = 3 def __post_init__(self) -> None: if self.num_query_heads <= 0 or self.num_kv_heads <= 0: @@ -104,6 +111,51 @@ def __post_init__(self) -> None: raise ValueError( "H2O layer-wise probability scoring requires decode score output." ) + if self.context_capacity is not None and self.context_capacity <= 0: + raise ValueError("Decode attention context_capacity must be positive.") + if self.may_use_full_layer_kivi_int4 and not self.layer_varying_page_table: + raise ValueError( + "Full-layer KIVI decode requires a layer-varying KV view contract." + ) + if ( + self.may_use_full_layer_kivi_int4 + and ( + self.full_layer_kivi_decode_block_seq <= 0 + or self.full_layer_kivi_decode_block_seq % 16 + ) + ): + raise ValueError( + "Full-layer KIVI decode block_seq must be a positive multiple " + f"of 16, got {self.full_layer_kivi_decode_block_seq}." + ) + if self.may_use_full_layer_kivi_int4 and ( + self.full_layer_kivi_decode_block_n <= 0 + or self.full_layer_kivi_decode_block_n % 16 + or self.full_layer_kivi_decode_block_seq + % self.full_layer_kivi_decode_block_n + ): + raise ValueError( + "Full-layer KIVI decode block_n must be a positive multiple of " + "16 and divide block_seq, got " + f"block_n={self.full_layer_kivi_decode_block_n}, " + f"block_seq={self.full_layer_kivi_decode_block_seq}." + ) + if ( + self.may_use_full_layer_kivi_int4 + and self.full_layer_kivi_decode_num_warps not in {1, 2, 4, 8} + ): + raise ValueError( + "Full-layer KIVI decode num_warps must be one of 1, 2, 4, " + f"or 8, got {self.full_layer_kivi_decode_num_warps}." + ) + if ( + self.may_use_full_layer_kivi_int4 + and self.full_layer_kivi_decode_num_stages <= 0 + ): + raise ValueError( + "Full-layer KIVI decode num_stages must be positive, got " + f"{self.full_layer_kivi_decode_num_stages}." + ) @property def kernel_request(self) -> AttentionKernelRequest: @@ -127,6 +179,7 @@ def kernel_request(self) -> AttentionKernelRequest: class DecodeAttentionProvider: name = "" capabilities: AttentionKernelCapabilities + decode_graph_lifecycle = False def prepare( self, @@ -155,6 +208,109 @@ class DecodeAttentionRunResult: softmax_lse: torch.Tensor +@dataclass(frozen=True) +class GraphStableDecodeLaunchPlan: + """Capture-time launch envelope for context-stable MHA/GQA decode.""" + + plan_id: str + context_capacity: int + max_kv_splits: int + target_tokens_per_split: int + block_n: int + stage1_num_warps: int + stage1_num_stages: int + stage2_num_warps: int + stage2_num_stages: int + + def __post_init__(self) -> None: + positive = ( + self.context_capacity, + self.max_kv_splits, + self.target_tokens_per_split, + self.block_n, + self.stage1_num_warps, + self.stage1_num_stages, + self.stage2_num_warps, + self.stage2_num_stages, + ) + if any(value <= 0 for value in positive): + raise ValueError(f"Decode launch plan values must be positive: {self}.") + + def as_dict(self) -> dict[str, int | str]: + return { + "plan_id": self.plan_id, + "context_capacity": self.context_capacity, + "max_kv_splits": self.max_kv_splits, + "target_tokens_per_split": self.target_tokens_per_split, + "block_n": self.block_n, + "stage1_num_warps": self.stage1_num_warps, + "stage1_num_stages": self.stage1_num_stages, + "stage2_num_warps": self.stage2_num_warps, + "stage2_num_stages": self.stage2_num_stages, + } + + +def build_graph_stable_decode_launch_plan( + spec: DecodeAttentionOpSpec, + caps: DeviceCaps, +) -> GraphStableDecodeLaunchPlan: + """Resolve one context-invariant portable plan before provider preparation.""" + del caps + if spec.context_capacity is None: + raise ValueError( + "context-stable decode requires a static context_capacity." + ) + if spec.head_dim == 256: + block_n, stage1_warps, stage2_warps = 128, 4, 8 + elif spec.head_dim in {64, 128}: + block_n, stage1_warps, stage2_warps = 64, 2, 4 + else: + raise ValueError( + f"No context-stable decode launch plan for head_dim={spec.head_dim}." + ) + + # The grid is derived from the configured capacity, never the current + # request length. Capping the envelope bounds workspace and empty programs; + # each replay derives its effective split count from device context_lens. + max_kv_splits = min( + 64, + max(16, math.ceil(int(spec.context_capacity) / 4096)), + ) + return GraphStableDecodeLaunchPlan( + plan_id="portable_fixed_grid_v1", + context_capacity=int(spec.context_capacity), + max_kv_splits=max_kv_splits, + target_tokens_per_split=256, + block_n=block_n, + stage1_num_warps=stage1_warps, + stage1_num_stages=2, + stage2_num_warps=stage2_warps, + stage2_num_stages=2, + ) + + +def build_deltakv_kivi_decode_launch_plan( + spec: DecodeAttentionOpSpec, + caps: DeviceCaps, +) -> GraphStableDecodeLaunchPlan: + """Resolve the fixed split envelope for packed full-layer KIVI decode.""" + base = build_graph_stable_decode_launch_plan(spec, caps) + target_tokens_per_split = int(spec.full_layer_kivi_decode_block_seq) + max_kv_splits = min( + 64, + max(4, math.ceil(base.context_capacity / target_tokens_per_split)), + ) + return replace( + base, + plan_id="deltakv_kivi_fixed_grid_v1", + max_kv_splits=max_kv_splits, + target_tokens_per_split=target_tokens_per_split, + block_n=int(spec.full_layer_kivi_decode_block_n), + stage1_num_warps=int(spec.full_layer_kivi_decode_num_warps), + stage1_num_stages=int(spec.full_layer_kivi_decode_num_stages), + ) + + DECODE_ATTENTION_REGISTRY: OpRegistry[ DecodeAttentionOpSpec, DecodeAttentionProvider ] = OpRegistry( @@ -164,7 +320,11 @@ class DecodeAttentionRunResult: "sgl_fa3_paged_decode_sm90", "flashinfer_paged_decode", ), - repo_portable=("triton_paged_decode",), + repo_portable=( + "triton_paged_decode", + "triton_fixed_grid_paged_decode", + ), + repo_nonstandard=("triton_deltakv_fixed_grid_decode",), ), ) @@ -172,6 +332,7 @@ class DecodeAttentionRunResult: @DECODE_ATTENTION_REGISTRY.register_atomic(ProviderRole.UPSTREAM_STANDARD) class SglFa3PagedDecodeAttentionProvider(DecodeAttentionProvider): name = "sgl_fa3_paged_decode_sm90" + supports_batch_only_cuda_graph = True capabilities = AttentionKernelCapabilities( platforms=frozenset({PlatformEnum.CUDA}), compute_capabilities=frozenset({(9, 0)}), @@ -195,6 +356,10 @@ def supports( spec: DecodeAttentionOpSpec, caps: DeviceCaps, ) -> SupportResult: + if spec.may_use_full_layer_kivi_int4: + return SupportResult.unsupported( + "does not support mixed dense and full-layer KIVI int4 storage" + ) common = match_attention_capabilities( spec.kernel_request, caps, @@ -313,6 +478,8 @@ def _run_sgl( @DECODE_ATTENTION_REGISTRY.register_atomic(ProviderRole.UPSTREAM_STANDARD) class FlashInferPagedDecodeAttentionProvider(DecodeAttentionProvider): name = "flashinfer_paged_decode" + decode_graph_lifecycle = True + supports_batch_only_cuda_graph = True capabilities = AttentionKernelCapabilities( platforms=frozenset({PlatformEnum.CUDA}), activation_dtypes=frozenset({torch.bfloat16, torch.float16}), @@ -321,7 +488,7 @@ class FlashInferPagedDecodeAttentionProvider(DecodeAttentionProvider): returns_softmax_lse=True, layer_varying_page_table=True, varlen=True, - cuda_graph=False, + cuda_graph=True, ) def __init__(self) -> None: @@ -333,6 +500,10 @@ def supports( spec: DecodeAttentionOpSpec, caps: DeviceCaps, ) -> SupportResult: + if spec.may_use_full_layer_kivi_int4: + return SupportResult.unsupported( + "does not support mixed dense and full-layer KIVI int4 storage" + ) common = match_attention_capabilities( spec.kernel_request, caps, @@ -351,7 +522,6 @@ def prepare( *, device_index: int | None = None, ) -> None: - del spec if self._state is not None: return current_device = torch.cuda.current_device() @@ -362,12 +532,14 @@ def prepare( "FlashInfer decode must be prepared on the selected CUDA device: " f"selected={device_index} current={current_device}." ) - self._state = _FlashInferPagedDecodeState( - torch.device("cuda", int(device_index)) - ) + if not spec.cuda_graph: + self._state = _FlashInferPagedDecodeState( + torch.device("cuda", int(device_index)) + ) def close(self) -> None: self._state = None + self._active_graph_state = None def binding_metadata(self) -> dict[str, object]: return { @@ -376,9 +548,55 @@ def binding_metadata(self) -> dict[str, object]: "kernel_path": "flashinfer.BatchDecodeWithPagedKVCacheWrapper", "kv_layout": "NHD", "page_size": 1, - "cuda_graph": False, + "cuda_graph": True, + "graph_metadata": "fixed buffers + graph-out plan + graph-in page packing", } + def init_decode_graph_state( + self, + spec: DecodeAttentionOpSpec, + contract, + inputs, + ) -> _FlashInferPagedDecodeGraphState: + if not spec.cuda_graph: + raise RuntimeError("FlashInfer graph state requires a CUDA Graph spec.") + if contract.batch_capacity > spec.max_batch_size: + raise ValueError( + "FlashInfer graph batch exceeds the prepared operator capacity: " + f"graph={contract.batch_capacity} operator={spec.max_batch_size}." + ) + return _FlashInferPagedDecodeGraphState( + spec, + contract=contract, + inputs=inputs, + ) + + def prepare_decode_graph_out( + self, + state: _FlashInferPagedDecodeGraphState, + ) -> None: + self._active_graph_state = state + state.prepare_out_graph() + + def prepare_decode_graph_in( + self, + state: _FlashInferPagedDecodeGraphState, + ) -> None: + self._active_graph_state = state + + def decode_graph_keepalive_tensors( + self, + state: _FlashInferPagedDecodeGraphState, + ) -> list[torch.Tensor]: + return state.keepalive_tensors() + + def close_decode_graph_state( + self, + state: _FlashInferPagedDecodeGraphState, + ) -> None: + if getattr(self, "_active_graph_state", None) is state: + self._active_graph_state = None + def run( self, spec: DecodeAttentionOpSpec, @@ -392,8 +610,15 @@ def run( "FlashInfer decode received unsupported runtime arguments: " f"{sorted(kwargs)}." ) - if self._state is None: - raise RuntimeError("FlashInfer decode provider was not prepared.") + graph_state = getattr(self, "_active_graph_state", None) + if spec.cuda_graph: + if not isinstance(graph_state, _FlashInferPagedDecodeGraphState): + raise RuntimeError("FlashInfer graph decode state is not active.") + state = graph_state + else: + if self._state is None: + raise RuntimeError("FlashInfer decode provider was not prepared.") + state = self._state payload = view.payload meta = view.meta if q.dtype != spec.activation_dtype: @@ -405,31 +630,38 @@ def run( "FlashInfer decode requires Q/K/V with the same dtype, got " f"{q.dtype}/{payload.k_cache.dtype}/{payload.v_cache.dtype}." ) - max_context_len = getattr(meta, "max_context_len", None) - if max_context_len is None: - raise RuntimeError( - "FlashInfer decode requires host-side max_context_len metadata." - ) - context = get_context() - plan_key = ( - context.attention_validation_scope, - meta.active_slots.data_ptr(), - meta.req_indices.data_ptr(), - meta.context_lens.data_ptr(), - int(max_context_len), - ) - if getattr(self._state, "plan_key", None) != plan_key: - self._state.plan( - spec, + if spec.cuda_graph: + state.pack_page_indices( active_slots=meta.active_slots, req_indices=meta.req_indices, context_lens=meta.context_lens, - max_context_len=int(max_context_len), ) - self._state.plan_key = plan_key + else: + max_context_len = getattr(meta, "max_context_len", None) + if max_context_len is None: + raise RuntimeError( + "FlashInfer decode requires host-side max_context_len metadata." + ) + context = get_context() + plan_key = ( + context.attention_validation_scope, + meta.active_slots.data_ptr(), + meta.req_indices.data_ptr(), + meta.context_lens.data_ptr(), + int(max_context_len), + ) + if getattr(state, "plan_key", None) != plan_key: + state.plan( + spec, + active_slots=meta.active_slots, + req_indices=meta.req_indices, + context_lens=meta.context_lens, + max_context_len=int(max_context_len), + ) + state.plan_key = plan_key output = torch.empty_like(q) return_softmax_lse = spec.kernel_request.requires_softmax_lse - result = self._state.wrapper.run( + result = state.wrapper.run( q, ( payload.k_cache.unsqueeze(1), @@ -539,6 +771,114 @@ def plan( ) +class _FlashInferPagedDecodeGraphState: + def __init__(self, spec, *, contract, inputs) -> None: + self.spec = spec + self.contract = contract + self.inputs = inputs + device = inputs.context_lens.device + batch_size = int(contract.batch_capacity) + context_capacity = int(contract.context_capacity) + self.workspace = torch.empty( + 128 * 1024 * 1024, + dtype=torch.uint8, + device=device, + ) + self.indptr = torch.empty( + batch_size + 1, + dtype=torch.int32, + device=device, + ) + self.indices = torch.empty( + batch_size * context_capacity, + dtype=torch.int32, + device=device, + ) + self.last_page_len = torch.ones( + batch_size, + dtype=torch.int32, + device=device, + ) + pin_memory = bool(inputs.host.context_lens.is_pinned()) + self.host_indptr = torch.empty( + batch_size + 1, + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) + self.host_last_page_len = torch.ones( + batch_size, + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) + self.wrapper = make_flashinfer_paged_decode_wrapper( + self.workspace, + use_cuda_graph=True, + paged_kv_indptr_buffer=self.indptr, + paged_kv_indices_buffer=self.indices, + paged_kv_last_page_len_buffer=self.last_page_len, + ) + self.planned = False + + def prepare_out_graph(self) -> None: + context_lens = self.inputs.host.context_lens + if torch.any(context_lens <= 0): + raise ValueError("FlashInfer graph decode requires positive context lengths.") + if torch.any(context_lens > self.contract.context_capacity): + raise ValueError( + "FlashInfer graph decode context exceeds its captured capacity." + ) + self.host_indptr[0] = 0 + torch.cumsum(context_lens, dim=0, dtype=torch.int32, out=self.host_indptr[1:]) + total_pages = int(self.host_indptr[-1]) + self.wrapper.plan( + self.host_indptr, + self.indices[:total_pages], + self.host_last_page_len, + num_qo_heads=self.spec.num_query_heads, + num_kv_heads=self.spec.num_kv_heads, + head_dim=self.spec.head_dim, + page_size=self.spec.page_size, + sm_scale=self.spec.softmax_scale, + q_data_type=self.spec.activation_dtype, + kv_data_type=self.spec.activation_dtype, + non_blocking=True, + ) + self.planned = True + + def pack_page_indices( + self, + *, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + ) -> None: + if not self.planned: + raise RuntimeError("FlashInfer graph decode was not planned before forward.") + from sparsevllm.kernels.triton.flashinfer_decode_metadata import ( + pack_flashinfer_page_indices, + ) + + pack_flashinfer_page_indices( + active_slots, + req_indices, + context_lens, + self.indices, + context_capacity=int(self.contract.context_capacity), + ) + + def keepalive_tensors(self) -> list[torch.Tensor]: + return [ + self.workspace, + self.indptr, + self.indices, + self.last_page_len, + self.host_indptr, + self.host_last_page_len, + ] + + @DECODE_ATTENTION_REGISTRY.register_atomic(ProviderRole.REPO_PORTABLE) class TritonPagedDecodeAttentionProvider(DecodeAttentionProvider): name = "triton_paged_decode" @@ -570,6 +910,8 @@ def supports( spec: DecodeAttentionOpSpec, caps: DeviceCaps, ) -> SupportResult: + if spec.batch_only_cuda_graph: + return SupportResult.unsupported("split count depends on context length") return match_attention_capabilities( spec.kernel_request, caps, @@ -674,6 +1016,561 @@ def run( ) +@DECODE_ATTENTION_REGISTRY.register_atomic(ProviderRole.REPO_PORTABLE) +class FixedGridTritonPagedDecodeAttentionProvider(DecodeAttentionProvider): + """Fixed-grid Triton MHA/GQA decode provider for batch-only graphs.""" + + name = "triton_fixed_grid_paged_decode" + supports_batch_only_cuda_graph = True + capabilities = replace( + TritonPagedDecodeAttentionProvider.capabilities, + activation_dtypes=frozenset({torch.bfloat16, torch.float16}), + head_dims=frozenset({64, 128, 256}), + returns_softmax_lse=True, + ) + + def __init__( + self, + *, + launch_plan: GraphStableDecodeLaunchPlan, + ) -> None: + self.launch_plan = launch_plan + self._mid_o: torch.Tensor | None = None + self._mid_lse: torch.Tensor | None = None + self._softmax_lse: torch.Tensor | None = None + + @classmethod + def bind( + cls, + spec: DecodeAttentionOpSpec, + caps: DeviceCaps, + **provider_kwargs, + ) -> FixedGridTritonPagedDecodeAttentionProvider: + if provider_kwargs: + raise TypeError( + "Fixed-grid Triton decode does not accept provider " + f"arguments: {sorted(provider_kwargs)}." + ) + return cls(launch_plan=build_graph_stable_decode_launch_plan(spec, caps)) + + @classmethod + def supports( + cls, spec: DecodeAttentionOpSpec, caps: DeviceCaps + ) -> SupportResult: + if not spec.batch_only_cuda_graph: + return SupportResult.unsupported("reserved for batch-only CUDA Graph") + if spec.may_use_full_layer_kivi_int4: + return SupportResult.unsupported( + "full-layer KIVI int4 requires the DeltaKV fixed-grid provider" + ) + if spec.context_capacity is None: + return SupportResult.unsupported("requires a static context capacity") + return match_attention_capabilities( + spec.kernel_request, + caps, + cls.capabilities, + ) + + def prepare( + self, + spec: DecodeAttentionOpSpec, + *, + device_index: int | None = None, + ) -> None: + if self.launch_plan.context_capacity != spec.context_capacity: + raise RuntimeError( + "Fixed-grid decode launch plan does not match the operator " + f"capacity: plan={self.launch_plan.context_capacity} " + f"spec={spec.context_capacity}." + ) + if device_index is None: + device_index = torch.cuda.current_device() + device = torch.device("cuda", int(device_index)) + self._mid_o = torch.empty( + ( + spec.max_batch_size, + spec.num_query_heads, + self.launch_plan.max_kv_splits, + spec.head_dim, + ), + dtype=torch.float32, + device=device, + ) + self._mid_lse = torch.empty( + ( + spec.max_batch_size, + spec.num_query_heads, + self.launch_plan.max_kv_splits, + ), + dtype=torch.float32, + device=device, + ) + self._softmax_lse = torch.empty( + (spec.num_query_heads, spec.max_batch_size), + dtype=torch.float32, + device=device, + ) + + def close(self) -> None: + self._mid_o = None + self._mid_lse = None + self._softmax_lse = None + + def binding_metadata(self) -> dict[str, object]: + return { + "implementation_kind": "atomic_provider", + "implementation_source": "repo_triton", + "kernel_path": "paged_flash_decode", + "cuda_graph_shape_policy": "batch_only", + "launch_plan": self.launch_plan.as_dict(), + "workspace_owner": "provider", + } + + def run( + self, + spec: DecodeAttentionOpSpec, + q: torch.Tensor, + view: Any, + **kwargs, + ) -> torch.Tensor: + kwargs.pop("decode_launch_op", None) + if kwargs: + raise TypeError( + "Fixed-grid decode received unsupported arguments: " + f"{sorted(kwargs)}." + ) + if ( + self._mid_o is None + or self._mid_lse is None + or self._softmax_lse is None + ): + raise RuntimeError("Fixed-grid decode provider was not prepared.") + payload = view.payload + if getattr(payload, "backend", None) != "dense": + raise RuntimeError( + "Fixed-grid decode requires dense explicit KV storage." + ) + batch_size = int(q.shape[0]) + from sparsevllm.kernels.triton.paged_flash_decoding import ( + paged_flash_decode, + ) + + result = paged_flash_decode( + q, + payload.k_cache, + payload.v_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + self._mid_o[:batch_size], + self._mid_lse[:batch_size], + attn_score=( + None + if spec.h2o_layerwise_probability_scores + else view.meta.attn_score + ), + softmax_scale=spec.softmax_scale, + target_tokens_per_split=self.launch_plan.target_tokens_per_split, + block_n=self.launch_plan.block_n, + num_warps=self.launch_plan.stage1_num_warps, + num_stages=self.launch_plan.stage1_num_stages, + stage2_num_warps=self.launch_plan.stage2_num_warps, + stage2_num_stages=self.launch_plan.stage2_num_stages, + return_softmax_lse=spec.h2o_layerwise_probability_scores, + output_lse=self._softmax_lse[:, :batch_size], + ) + if not spec.h2o_layerwise_probability_scores: + return result + if not isinstance(result, tuple): + raise RuntimeError("Fixed-grid decode did not return softmax LSE.") + return DecodeAttentionRunResult(output=result[0], softmax_lse=result[1]) + + +@dataclass +class _DeltaKVFixedGridDecodeState: + batch_capacity: int + launch_plan: GraphStableDecodeLaunchPlan + kivi_launch_plan: GraphStableDecodeLaunchPlan + mid_o: torch.Tensor + mid_lse: torch.Tensor + kivi_mid_o: torch.Tensor + kivi_mid_lse: torch.Tensor + output: torch.Tensor + output_lse: torch.Tensor + + @classmethod + def allocate( + cls, + spec: DecodeAttentionOpSpec, + caps: DeviceCaps, + *, + batch_capacity: int, + context_capacity: int, + device: torch.device, + ) -> _DeltaKVFixedGridDecodeState: + graph_spec = replace( + spec, + max_batch_size=int(batch_capacity), + context_capacity=int(context_capacity), + ) + launch_plan = build_graph_stable_decode_launch_plan(graph_spec, caps) + kivi_launch_plan = build_deltakv_kivi_decode_launch_plan(graph_spec, caps) + return cls( + batch_capacity=int(batch_capacity), + launch_plan=launch_plan, + kivi_launch_plan=kivi_launch_plan, + mid_o=torch.empty( + ( + batch_capacity, + spec.num_query_heads, + launch_plan.max_kv_splits, + spec.head_dim, + ), + dtype=torch.float32, + device=device, + ), + mid_lse=torch.empty( + ( + batch_capacity, + spec.num_query_heads, + launch_plan.max_kv_splits, + ), + dtype=torch.float32, + device=device, + ), + kivi_mid_o=torch.empty( + ( + batch_capacity, + spec.num_query_heads, + kivi_launch_plan.max_kv_splits, + spec.head_dim, + ), + dtype=torch.float32, + device=device, + ), + kivi_mid_lse=torch.empty( + ( + batch_capacity, + spec.num_query_heads, + kivi_launch_plan.max_kv_splits, + ), + dtype=torch.float32, + device=device, + ), + output=torch.empty( + (batch_capacity, spec.num_query_heads, spec.head_dim), + dtype=spec.activation_dtype, + device=device, + ), + output_lse=torch.empty( + (spec.num_query_heads, batch_capacity), + dtype=torch.float32, + device=device, + ), + ) + + def keepalive_tensors(self) -> list[torch.Tensor]: + return [ + self.mid_o, + self.mid_lse, + self.kivi_mid_o, + self.kivi_mid_lse, + self.output, + self.output_lse, + ] + + +@DECODE_ATTENTION_REGISTRY.register_atomic(ProviderRole.REPO_NONSTANDARD) +class DeltaKVFixedGridDecodeAttentionProvider(DecodeAttentionProvider): + """Batch-only provider for DeltaKV's dense and full-layer KIVI views.""" + + name = "triton_deltakv_fixed_grid_decode" + supports_batch_only_cuda_graph = True + decode_graph_lifecycle = True + capabilities = replace( + FixedGridTritonPagedDecodeAttentionProvider.capabilities, + head_dims=frozenset({64, 128}), + ) + + def __init__( + self, + *, + caps: DeviceCaps, + launch_plan: GraphStableDecodeLaunchPlan, + kivi_launch_plan: GraphStableDecodeLaunchPlan, + ) -> None: + self._caps = caps + self.launch_plan = launch_plan + self.kivi_launch_plan = kivi_launch_plan + self._active_graph_state: _DeltaKVFixedGridDecodeState | None = None + + @classmethod + def bind( + cls, + spec: DecodeAttentionOpSpec, + caps: DeviceCaps, + **provider_kwargs, + ) -> DeltaKVFixedGridDecodeAttentionProvider: + if provider_kwargs: + raise TypeError( + "DeltaKV fixed-grid decode does not accept provider arguments: " + f"{sorted(provider_kwargs)}." + ) + return cls( + caps=caps, + launch_plan=build_graph_stable_decode_launch_plan(spec, caps), + kivi_launch_plan=build_deltakv_kivi_decode_launch_plan(spec, caps), + ) + + @classmethod + def supports( + cls, + spec: DecodeAttentionOpSpec, + caps: DeviceCaps, + ) -> SupportResult: + if not spec.batch_only_cuda_graph: + return SupportResult.unsupported("reserved for batch-only CUDA Graph") + if not spec.may_use_full_layer_kivi_int4: + return SupportResult.unsupported( + "reserved for mixed dense and full-layer KIVI int4 storage" + ) + if spec.context_capacity is None: + return SupportResult.unsupported("requires a static context capacity") + return match_attention_capabilities( + spec.kernel_request, + caps, + cls.capabilities, + ) + + def prepare( + self, + spec: DecodeAttentionOpSpec, + *, + device_index: int | None = None, + ) -> None: + del device_index + plans = (self.launch_plan, self.kivi_launch_plan) + if any(plan.context_capacity != spec.context_capacity for plan in plans): + raise RuntimeError( + "DeltaKV fixed-grid launch plans do not match the operator " + f"capacity: plans={[plan.context_capacity for plan in plans]} " + f"spec={spec.context_capacity}." + ) + + def close(self) -> None: + self._active_graph_state = None + + def binding_metadata(self) -> dict[str, object]: + return { + "implementation_kind": "atomic_provider", + "implementation_source": "repo_triton", + "kernel_path": "paged_flash_decode + full_layer_kivi_flash_decode", + "cuda_graph_shape_policy": "batch_only", + "launch_plan": self.launch_plan.as_dict(), + "kivi_launch_plan": self.kivi_launch_plan.as_dict(), + "workspace_owner": "per_graph_provider_state", + "payload_routes": ["dense", "full_layer_kivi"], + } + + def init_decode_graph_state( + self, + spec: DecodeAttentionOpSpec, + contract, + inputs, + ) -> _DeltaKVFixedGridDecodeState: + if contract.shape_policy != "batch_only": + raise ValueError("DeltaKV fixed-grid state requires a batch-only contract.") + if int(contract.batch_capacity) > int(spec.max_batch_size): + raise ValueError( + "DeltaKV graph batch exceeds the prepared operator capacity: " + f"graph={contract.batch_capacity} operator={spec.max_batch_size}." + ) + if spec.context_capacity is None or int(contract.context_capacity) > int( + spec.context_capacity + ): + raise ValueError( + "DeltaKV graph context exceeds the prepared operator capacity: " + f"graph={contract.context_capacity} operator={spec.context_capacity}." + ) + return _DeltaKVFixedGridDecodeState.allocate( + spec, + self._caps, + batch_capacity=int(contract.batch_capacity), + context_capacity=int(contract.context_capacity), + device=inputs.context_lens.device, + ) + + def prepare_decode_graph_out( + self, + state: _DeltaKVFixedGridDecodeState, + ) -> None: + self._active_graph_state = state + + def prepare_decode_graph_in( + self, + state: _DeltaKVFixedGridDecodeState, + ) -> None: + self._active_graph_state = state + + def decode_graph_keepalive_tensors( + self, + state: _DeltaKVFixedGridDecodeState, + ) -> list[torch.Tensor]: + return state.keepalive_tensors() + + def close_decode_graph_state( + self, + state: _DeltaKVFixedGridDecodeState, + ) -> None: + if self._active_graph_state is state: + self._active_graph_state = None + + def run( + self, + spec: DecodeAttentionOpSpec, + q: torch.Tensor, + view: Any, + **kwargs, + ) -> torch.Tensor: + kwargs.pop("decode_launch_op", None) + if kwargs: + raise TypeError( + "DeltaKV fixed-grid decode received unsupported arguments: " + f"{sorted(kwargs)}." + ) + state = self._active_graph_state + if state is None: + raise RuntimeError( + "DeltaKV fixed-grid decode has no active graph participant state." + ) + batch_size = int(q.shape[0]) + if batch_size > state.batch_capacity: + raise RuntimeError( + "DeltaKV fixed-grid decode batch exceeds active state capacity: " + f"batch={batch_size} capacity={state.batch_capacity}." + ) + payload = view.payload + backend = getattr(payload, "backend", None) + if backend == "dense": + return self._run_dense(spec, q, view, state, batch_size) + if backend == "full_layer_kivi": + return self._run_full_layer_kivi(q, view, state, batch_size) + raise RuntimeError( + "DeltaKV fixed-grid decode requires dense or full-layer KIVI storage, " + f"got {backend!r}." + ) + + @staticmethod + def _run_dense( + spec: DecodeAttentionOpSpec, + q: torch.Tensor, + view: Any, + state: _DeltaKVFixedGridDecodeState, + batch_size: int, + ) -> torch.Tensor: + from sparsevllm.kernels.triton.paged_flash_decoding import ( + paged_flash_decode, + ) + + return paged_flash_decode( + q, + view.payload.k_cache, + view.payload.v_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + state.mid_o[:batch_size], + state.mid_lse[:batch_size], + attn_score=view.meta.attn_score, + softmax_scale=spec.softmax_scale, + target_tokens_per_split=state.launch_plan.target_tokens_per_split, + block_n=state.launch_plan.block_n, + num_warps=state.launch_plan.stage1_num_warps, + num_stages=state.launch_plan.stage1_num_stages, + stage2_num_warps=state.launch_plan.stage2_num_warps, + stage2_num_stages=state.launch_plan.stage2_num_stages, + output_lse=state.output_lse[:, :batch_size], + output=state.output[:batch_size], + ) + + @staticmethod + def _run_full_layer_kivi( + q: torch.Tensor, + view: Any, + state: _DeltaKVFixedGridDecodeState, + batch_size: int, + ) -> torch.Tensor: + metadata = getattr(view.payload, "metadata", None) + if metadata is None: + raise RuntimeError("Full-layer KIVI decode view is missing metadata.") + required = ( + "kivi_block_slots_map", + "kivi_block_start_pos", + "key_packed", + "key_scales", + "key_mins", + "value_packed", + "value_scales", + "value_mins", + "group_size", + ) + missing = [name for name in required if name not in metadata] + if missing: + raise RuntimeError( + f"Full-layer KIVI decode view is missing metadata: {missing}." + ) + + from sparsevllm.kernels.triton.deltakv_kernels import ( + full_layer_kivi_flash_decode_stage1, + ) + from sparsevllm.kernels.triton.paged_flash_decoding import ( + fixed_grid_flash_decode_stage2, + ) + + plan = state.kivi_launch_plan + mid_o = state.kivi_mid_o[:batch_size] + mid_lse = state.kivi_mid_lse[:batch_size] + full_layer_kivi_flash_decode_stage1( + q=q, + raw_k=view.payload.k_cache, + raw_v=view.payload.v_cache, + raw_slots_map=view.meta.active_slots, + kivi_block_slots_map=metadata["kivi_block_slots_map"], + kivi_block_start_pos=metadata["kivi_block_start_pos"], + key_packed=metadata["key_packed"], + key_scales=metadata["key_scales"], + key_mins=metadata["key_mins"], + value_packed=metadata["value_packed"], + value_scales=metadata["value_scales"], + value_mins=metadata["value_mins"], + req_indices=view.meta.req_indices, + context_lens=view.meta.context_lens, + max_len_in_batch=plan.context_capacity, + mid_out=mid_o, + mid_out_logsumexp=mid_lse, + group_size=int(metadata["group_size"]), + block_seq=plan.target_tokens_per_split, + block_n=plan.block_n, + num_warps=plan.stage1_num_warps, + num_stages=plan.stage1_num_stages, + attn_score=view.meta.attn_score, + max_kv_splits=plan.max_kv_splits, + target_tokens_per_split=plan.target_tokens_per_split, + ) + output = state.output[:batch_size] + fixed_grid_flash_decode_stage2( + mid_o, + mid_lse, + view.meta.context_lens, + output, + state.output_lse[:, :batch_size], + target_tokens_per_split=plan.target_tokens_per_split, + num_warps=plan.stage2_num_warps, + num_stages=plan.stage2_num_stages, + ) + return output + + class PreparedDecodeAttentionOp: """One prepared decode provider shared by all compatible MHA layers.""" @@ -690,6 +1587,12 @@ def __init__( def name(self) -> str: return self.provider.name + @property + def supports_batch_only_cuda_graph(self) -> bool: + return bool( + getattr(self.provider, "supports_batch_only_cuda_graph", False) + ) + def run(self, q: torch.Tensor, view: Any, **kwargs) -> torch.Tensor: if self._closed: raise RuntimeError("Decode attention operator is closed.") @@ -730,6 +1633,30 @@ def run(self, q: torch.Tensor, view: Any, **kwargs) -> torch.Tensor: ) return result.output + @property + def decode_graph_lifecycle(self) -> bool: + return bool(getattr(self.provider, "decode_graph_lifecycle", False)) + + def init_decode_graph_state(self, contract, inputs): + initializer = getattr(self.provider, "init_decode_graph_state", None) + if not callable(initializer): + raise TypeError( + f"Decode provider {self.provider.name!r} has no graph-state initializer." + ) + return initializer(self.spec, contract, inputs) + + def prepare_decode_graph_out(self, state) -> None: + self.provider.prepare_decode_graph_out(state) + + def prepare_decode_graph_in(self, state) -> None: + self.provider.prepare_decode_graph_in(state) + + def decode_graph_keepalive_tensors(self, state) -> list[torch.Tensor]: + return list(self.provider.decode_graph_keepalive_tensors(state)) + + def close_decode_graph_state(self, state) -> None: + self.provider.close_decode_graph_state(state) + def close(self) -> None: if self._closed: return @@ -746,7 +1673,10 @@ def prepare_decode_attention_op( if device_index is None: device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 caps = platform.get_device_caps(int(device_index)) - resolved = OpResolver(DECODE_ATTENTION_REGISTRY).resolve(spec, caps) + resolved = OpResolver(DECODE_ATTENTION_REGISTRY).resolve( + spec, + caps, + ) logger.info( "Resolved MHA decode provider={} rejected={}", resolved.provider.name, @@ -756,6 +1686,77 @@ def prepare_decode_attention_op( return PreparedDecodeAttentionOp(spec, resolved.provider) +def collect_decode_graph_participants(model: torch.nn.Module) -> tuple[object, ...]: + """Collect unique prepared decode operators with graph-out lifecycle state.""" + + from sparsevllm.layers.attention import Attention + + participants: list[object] = [] + seen: set[int] = set() + for module in model.modules(): + if not isinstance(module, Attention): + continue + participant = getattr(module, "decode_op", None) + if participant is None or not bool( + getattr(participant, "decode_graph_lifecycle", False) + ): + continue + identity = id(participant) + if identity not in seen: + seen.add(identity) + participants.append(participant) + return tuple(participants) + + +def validate_batch_only_decode_graph_model(model: torch.nn.Module) -> int: + """Audit every semantic decode path after construction-time binding.""" + from sparsevllm.layers.attention import Attention + + validated = 0 + for module in model.modules(): + if isinstance(module, Attention): + decode_op = getattr(module, "decode_op", None) + implementation = ( + decode_op + if decode_op is not None + else getattr(module, "attention_backend", None) + ) + if not bool( + getattr(implementation, "supports_batch_only_cuda_graph", False) + ): + raise RuntimeError( + "batch-only decode CUDA Graph requires a graph-stable " + f"attention provider, got {type(implementation).__name__}." + ) + validated += 1 + if getattr(module, "is_gated_delta_rule_layer", False): + op = getattr(module, "gated_delta_rule_op", None) + if not bool(getattr(op, "supports_batch_only_cuda_graph", False)): + raise RuntimeError( + "batch-only decode CUDA Graph requires a graph-stable " + "GDN provider." + ) + validated += 1 + + model_body = getattr(model, "model", None) + mla_attention = getattr(model_body, "mla_attention", None) + if mla_attention is not None: + provider = getattr(mla_attention, "provider", None) + if not bool( + getattr(provider, "supports_batch_only_cuda_graph", False) + ): + raise RuntimeError( + "batch-only decode CUDA Graph requires a graph-stable " + "MLA provider." + ) + validated += 1 + if validated == 0: + raise RuntimeError( + "batch-only decode CUDA Graph found no validated decode operator." + ) + return validated + + @dataclass(frozen=True) class DecodeAttentionLaunchSpec: num_query_heads: int diff --git a/src/sparsevllm/operators/gated_delta_rule.py b/src/sparsevllm/operators/gated_delta_rule.py index 54758dd7..43064ed4 100644 --- a/src/sparsevllm/operators/gated_delta_rule.py +++ b/src/sparsevllm/operators/gated_delta_rule.py @@ -37,6 +37,7 @@ class GatedDeltaRuleOpSpec: state_layout_id: str = "k_major_hkv" varlen_prefill: bool = True cuda_graph_decode: bool = True + batch_only_cuda_graph: bool = False def __post_init__(self) -> None: if self.num_key_heads <= 0 or self.num_value_heads <= 0: @@ -379,6 +380,10 @@ def __init__( def name(self) -> str: return self.provider.name + @property + def supports_batch_only_cuda_graph(self) -> bool: + return bool(self.spec.batch_only_cuda_graph) + def run_prefill(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: if self._closed: raise RuntimeError("GDN operator is closed.") diff --git a/src/sparsevllm/operators/gemma4.py b/src/sparsevllm/operators/gemma4.py index 572a6194..9dfbb140 100644 --- a/src/sparsevllm/operators/gemma4.py +++ b/src/sparsevllm/operators/gemma4.py @@ -28,10 +28,19 @@ class Gemma4OpSpec: head_dims: tuple[int, ...] cuda_graph: bool attention_contracts: tuple[tuple[int, int, int, int], ...] = () + max_batch_size: int = 1 + batch_only_cuda_graph: bool = False + context_capacity: int | None = None def __post_init__(self) -> None: if not self.head_dims or any(int(value) <= 0 for value in self.head_dims): raise ValueError("Gemma 4 head dimensions must be positive.") + if self.max_batch_size <= 0: + raise ValueError("Gemma 4 max_batch_size must be positive.") + if self.context_capacity is not None and self.context_capacity <= 0: + raise ValueError("Gemma 4 context_capacity must be positive.") + if self.batch_only_cuda_graph and self.context_capacity is None: + raise ValueError("Gemma 4 batch-only decode requires context_capacity.") class Gemma4OperatorProvider: @@ -54,8 +63,7 @@ def binding_metadata(self) -> dict[str, object]: "triton_context", ], "decode_routes": [ - "triton_single_block", - "triton_two_stage", + "sglang_fixed_grid", ], }, } @@ -134,7 +142,9 @@ def rmsnorm_residual( GEMMA4_REGISTRY: OpRegistry[Gemma4OpSpec, Gemma4OperatorProvider] = OpRegistry( "Gemma 4 model operations", - portfolio=PortfolioPolicy(repo_nonstandard=("triton",)), + portfolio=PortfolioPolicy( + repo_nonstandard=("triton",) + ), profile_order=("gemma4_h20_profile",), ) @@ -143,6 +153,22 @@ def rmsnorm_residual( class TritonGemma4OperatorProvider(Gemma4OperatorProvider): name = "triton" + def __init__( + self, + *, + spec: Gemma4OpSpec | None = None, + caps: DeviceCaps | None = None, + ) -> None: + super().__init__() + self.spec = spec + self.device = None if caps is None else torch.device("cuda", caps.device_index) + self.multi_processor_count = ( + None if caps is None else int(caps.multi_processor_count or 0) + ) + if caps is not None and self.multi_processor_count <= 0: + raise ValueError("Gemma 4 requires a positive multi-processor count.") + self._decode_workspaces: dict[tuple[int, int, int], object] = {} + @classmethod def supports(cls, spec: Gemma4OpSpec, caps: DeviceCaps) -> SupportResult: if caps.platform != PlatformEnum.CUDA or not caps.supports_triton: @@ -153,15 +179,92 @@ def supports(cls, spec: Gemma4OpSpec, caps: DeviceCaps) -> SupportResult: return SupportResult.unsupported("requires BF16 or FP16 activations") if any(head_dim not in {256, 512} for head_dim in spec.head_dims): return SupportResult.unsupported("requires attention head dimensions 256 or 512") + if ( + caps.multi_processor_count is None + or int(caps.multi_processor_count) <= 0 + ): + return SupportResult.unsupported( + "requires a positive multi-processor count" + ) return SupportResult.yes() + @classmethod + def bind( + cls, + spec: Gemma4OpSpec, + caps: DeviceCaps, + **kwargs, + ) -> TritonGemma4OperatorProvider: + if kwargs: + raise TypeError(f"Unexpected Gemma 4 bind arguments: {sorted(kwargs)}.") + return cls(spec=spec, caps=caps) + def attention_backend(self, *, sliding_window: int | None): - from sparsevllm.operators.gemma4_attention import Gemma4AttentionBackend + from sparsevllm.operators.gemma4_attention import ( + Gemma4AttentionBackend, + Gemma4DecodeWorkspace, + ) + + if ( + self.spec is None + or self.device is None + or self.multi_processor_count is None + ): + raise RuntimeError( + "Gemma 4 attention requires a provider bound from Gemma4OpSpec." + ) + window_left = -1 if sliding_window is None else int(sliding_window) - 1 + matching = [ + contract + for contract in self.spec.attention_contracts + if int(contract[3]) == window_left + ] + if len(matching) != 1: + raise RuntimeError( + "Gemma 4 provider requires one attention contract for " + f"window_left={window_left}, got {matching}." + ) + query_heads, _, head_dim, _ = matching[0] + max_kv_splits = 8 + signature = (int(query_heads), int(head_dim), max_kv_splits) + workspace = self._decode_workspaces.get(signature) + if workspace is None: + workspace = Gemma4DecodeWorkspace( + mid_output=torch.empty( + ( + self.spec.max_batch_size, + signature[0], + signature[2], + signature[1], + ), + dtype=torch.float32, + device=self.device, + ), + mid_lse=torch.empty( + (self.spec.max_batch_size, signature[0], signature[2]), + dtype=torch.float32, + device=self.device, + ), + num_kv_splits=torch.empty( + (self.spec.max_batch_size,), + dtype=torch.int32, + device=self.device, + ), + ) + self._decode_workspaces[signature] = workspace return self._register_attention_backend( - Gemma4AttentionBackend(sliding_window=sliding_window) + Gemma4AttentionBackend( + sliding_window=sliding_window, + decode_workspace=workspace, + multi_processor_count=self.multi_processor_count, + ) ) + def close(self) -> None: + super().close() + self._decode_workspaces.clear() + def rmsnorm(self, x, weight, eps): from sparsevllm.kernels.triton.gemma4_rmsnorm import gemma4_rmsnorm @@ -241,6 +344,8 @@ def bind( f"{cls.name} does not accept provider arguments: {sorted(kwargs)}" ) return cls( + spec=spec, + caps=caps, device_index=caps.device_index, max_prefill_contracts=( len(spec.attention_contracts) or len(spec.head_dims) @@ -250,10 +355,12 @@ def bind( def __init__( self, *, + spec: Gemma4OpSpec, + caps: DeviceCaps, device_index: int | None = None, max_prefill_contracts: int = 2, ) -> None: - super().__init__() + super().__init__(spec=spec, caps=caps) from sparsevllm.operators.gemma4_attention import Gemma4FlashInferPrefill self._prefill = Gemma4FlashInferPrefill() @@ -273,10 +380,7 @@ def binding_metadata(self) -> dict[str, object]: "triton_context", ], "decode_routes": [ - "triton_window", - "triton_single_block", - "triton_global", - "triton_two_stage", + "sglang_fixed_grid", ], }, "flashinfer_backend": "fa2", @@ -290,16 +394,9 @@ def close(self) -> None: super().close() def attention_backend(self, *, sliding_window: int | None): - from sparsevllm.operators.gemma4_attention import Gemma4AttentionBackend - - return self._register_attention_backend( - Gemma4AttentionBackend( - sliding_window=sliding_window, - flashinfer_prefill=self._prefill, - use_window_decode=True, - global_decode_heads_per_program=4, - ) - ) + backend = super().attention_backend(sliding_window=sliding_window) + backend.flashinfer_prefill = self._prefill + return backend @GEMMA4_REGISTRY.register_profile diff --git a/src/sparsevllm/operators/gemma4_attention.py b/src/sparsevllm/operators/gemma4_attention.py index 829a7ebb..c383aa40 100644 --- a/src/sparsevllm/operators/gemma4_attention.py +++ b/src/sparsevllm/operators/gemma4_attention.py @@ -21,6 +21,13 @@ class _FlashInferState: plan_key: tuple[object, ...] | None = None +@dataclass +class Gemma4DecodeWorkspace: + mid_output: torch.Tensor + mid_lse: torch.Tensor + num_kv_splits: torch.Tensor + + class Gemma4FlashInferPrefill: """Shared FlashInfer plans for Gemma 4 text-prefill head shapes.""" @@ -166,35 +173,67 @@ class Gemma4AttentionBackend(TritonAttentionBackend): """Gemma 4 attention semantics isolated from the tuned generic kernels.""" name = "triton_gemma4" + supports_batch_only_cuda_graph = True def __init__( self, *, sliding_window: int | None, flashinfer_prefill: Gemma4FlashInferPrefill | None = None, - use_window_decode: bool = False, - global_decode_heads_per_program: int | None = None, + decode_workspace: Gemma4DecodeWorkspace | None = None, + multi_processor_count: int | None = None, ) -> None: super().__init__() self.sliding_window = None if sliding_window is None else int(sliding_window) self.flashinfer_prefill = flashinfer_prefill - self.use_window_decode = bool(use_window_decode) - self.global_decode_heads_per_program = global_decode_heads_per_program + self.decode_workspace = decode_workspace + self.multi_processor_count = ( + None + if multi_processor_count is None + else int(multi_processor_count) + ) self._runtime_kernel_path_counts: dict[str, dict[str, int]] = {} + def get_decode_workspace( + self, + *, + batch_size: int, + num_heads: int, + head_dim: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + workspace = self.decode_workspace + if workspace is None: + raise RuntimeError("Gemma 4 decode backend has no prepared workspace.") + if ( + self.multi_processor_count is None + or self.multi_processor_count <= 0 + ): + raise RuntimeError( + "Gemma 4 decode backend requires a positive multi-processor count." + ) + if ( + batch_size > workspace.mid_output.shape[0] + or num_heads != workspace.mid_output.shape[1] + or head_dim != workspace.mid_output.shape[3] + or device != workspace.mid_output.device + ): + raise RuntimeError( + "Gemma 4 fixed-grid workspace does not match the decode contract: " + f"actual={(batch_size, num_heads, head_dim, device)} " + f"workspace={tuple(workspace.mid_output.shape)}/" + f"{workspace.mid_output.device}." + ) + return workspace.mid_output[:batch_size], workspace.mid_lse[:batch_size] + def binding_metadata(self) -> dict[str, object]: prefill_routes = ["triton_multimodal_context", "triton_context"] if self.flashinfer_prefill is not None: prefill_routes.insert(1, "flashinfer_paged_prefill_fa2") - decode_routes = ["triton_single_block", "triton_two_stage"] - if self.use_window_decode: - decode_routes.insert(0, "triton_window") - if self.global_decode_heads_per_program is not None: - decode_routes.insert(-1, "triton_global") return { "implementation_kind": "dispatch_plan", "prefill_routes": prefill_routes, - "decode_routes": decode_routes, + "decode_routes": ["sglang_fixed_grid"], "sliding_window": self.sliding_window, } @@ -229,41 +268,6 @@ def _prefill_route(self, view) -> str: return "flashinfer_paged_prefill_fa2" return "triton_context" - def _decode_route( - self, - q: torch.Tensor, - view, - *, - mid_o: torch.Tensor, - block_seq: int, - group_size: int, - ) -> str: - if ( - self.use_window_decode - and self.sliding_window is not None - and view.meta.attn_score is None - and int(q.shape[-1]) == 256 - and group_size in {2, 4} - and mid_o.shape[2] - >= (self.sliding_window + block_seq - 1) // block_seq - ): - return "triton_window" - if ( - mid_o.shape[2] == 1 - and view.meta.attn_score is None - and group_size in {2, 4, 8} - ): - return "triton_single_block" - if ( - self.sliding_window is None - and view.meta.attn_score is None - and int(q.shape[-1]) == 512 - and self.global_decode_heads_per_program is not None - and group_size % self.global_decode_heads_per_program == 0 - ): - return "triton_global" - return "triton_two_stage" - def run_prefill( self, q: torch.Tensor, @@ -345,99 +349,49 @@ def run_decode( gqa_block_n: int = 16, gqa_num_warps: int = 2, ) -> torch.Tensor: - del max_len_in_batch, num_heads, num_kv_heads, gqa_block_n, gqa_num_warps + del ( + max_len_in_batch, + block_seq, + num_heads, + num_kv_heads, + gqa_block_n, + gqa_num_warps, + ) + workspace = self.decode_workspace + if workspace is None: + raise RuntimeError("Gemma 4 decode backend has no prepared workspace.") + multi_processor_count = self.multi_processor_count + if multi_processor_count is None or multi_processor_count <= 0: + raise RuntimeError( + "Gemma 4 decode backend requires a positive multi-processor count." + ) payload = _require_explicit_payload(view, operation="Gemma 4 decode") - from sparsevllm.kernels.triton.gemma4_decode_attention import ( - gemma4_decode_stage1, - gemma4_decode_stage2, + if payload.backend != "dense": + raise RuntimeError("Gemma 4 fixed-grid decode requires dense explicit KV.") + from sparsevllm.kernels.triton.sglang_gemma4_decode_attention import ( + sglang_gemma4_decode, ) - group_size = int(q.shape[1]) // int(payload.k_cache.shape[1]) - route = self._decode_route( + batch_size = int(q.shape[0]) + self._record_kernel_path("sglang_fixed_grid") + return sglang_gemma4_decode( q, - view, - mid_o=mid_o, - block_seq=block_seq, - group_size=group_size, - ) - self._record_kernel_path(route) - if route == "triton_window": - from sparsevllm.kernels.triton.gemma4_window_decode_attention import ( - gemma4_window_decode, - ) - - output = torch.empty_like(q) - window_blocks = (self.sliding_window + block_seq - 1) // block_seq - gemma4_window_decode( - q, - payload.k_cache, - payload.v_cache, - view.meta.active_slots, - view.meta.req_indices, - view.meta.context_lens, - mid_o[:, :, :window_blocks], - mid_o_logexpsum[:, :, :window_blocks], - output, - block_seq=block_seq, - sliding_window=self.sliding_window, - ) - return output - if route == "triton_single_block": - from sparsevllm.kernels.triton.gemma4_single_block_decode_attention import ( - gemma4_single_block_decode, - ) - - output = torch.empty_like(q) - gemma4_single_block_decode( - q, payload.k_cache, payload.v_cache, view.meta.active_slots, - view.meta.req_indices, view.meta.context_lens, output, - block_seq=block_seq, sliding_window=self.sliding_window, - ) - return output - if route == "triton_global": - from sparsevllm.kernels.triton.gemma4_global_decode_attention import ( - gemma4_global_decode_stage1, - ) - - gemma4_global_decode_stage1( - q, - payload.k_cache, - payload.v_cache, - view.meta.active_slots, - view.meta.req_indices, - view.meta.context_lens, - mid_o, - mid_o_logexpsum, - block_seq=block_seq, - heads_per_program=self.global_decode_heads_per_program, - ) - output = torch.empty_like(q) - gemma4_decode_stage2( - mid_o, - mid_o_logexpsum, - view.meta.context_lens, - output, - block_seq=block_seq, - sliding_window=None, - ) - return output - gemma4_decode_stage1( - q, payload.k_cache, payload.v_cache, view.meta.active_slots, - view.meta.req_indices, view.meta.context_lens, mid_o, - mid_o_logexpsum, block_seq=block_seq, - sliding_window=self.sliding_window, - attn_score=view.meta.attn_score, - ) - output = torch.empty_like(q) - gemma4_decode_stage2( - mid_o, - mid_o_logexpsum, + payload.k_cache, + payload.v_cache, + view.meta.active_slots, + view.meta.req_indices, view.meta.context_lens, - output, - block_seq=block_seq, + workspace.mid_output[:batch_size], + workspace.mid_lse[:batch_size], + workspace.num_kv_splits[:batch_size], sliding_window=self.sliding_window, + multi_processor_count=multi_processor_count, + attn_score=view.meta.attn_score, ) - return output -__all__ = ["Gemma4AttentionBackend", "Gemma4FlashInferPrefill"] +__all__ = [ + "Gemma4AttentionBackend", + "Gemma4DecodeWorkspace", + "Gemma4FlashInferPrefill", +] diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index fc5312d7..2f695d6a 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -18,6 +18,7 @@ ) from sparsevllm.kernels.tilelang.mla.runtime import ( TileMlaDecodeKernel, + TileMlaLaunchPlan, tilelang_mla_support, ) from sparsevllm.kernels.triton.mla import ( @@ -66,7 +67,9 @@ class MlaAttentionOpSpec: cache_dtype: torch.dtype tp_size: int cuda_graph: bool - may_require_attention_scores: bool = False + score_output: AttentionScoreKind = AttentionScoreKind.NONE + batch_only_cuda_graph: bool = False + context_capacity: int | None = None def __post_init__(self) -> None: dimensions = { @@ -85,6 +88,18 @@ def __post_init__(self) -> None: "MLA query heads must be divisible by tensor parallel size: " f"heads={self.num_q_heads} tp_size={self.tp_size}." ) + if self.context_capacity is not None and self.context_capacity <= 0: + raise ValueError("MLA context_capacity must be positive.") + if self.score_output not in { + AttentionScoreKind.NONE, + AttentionScoreKind.RAW_QK_PER_HEAD, + AttentionScoreKind.RAW_QK_REDUCED, + }: + raise ValueError( + "MLA decode currently supports NONE, RAW_QK_PER_HEAD, or " + "RAW_QK_REDUCED score " + f"contracts, got {self.score_output.name}." + ) @property def local_q_heads(self) -> int: @@ -99,11 +114,7 @@ def kernel_request(self) -> AttentionKernelRequest: return AttentionKernelRequest( activation_dtype=self.activation_dtype, head_dim=self.qk_head_dim, - score_output=( - AttentionScoreKind.RAW_QK_REDUCED - if self.may_require_attention_scores - else AttentionScoreKind.NONE - ), + score_output=self.score_output, layer_varying_page_table=True, varlen=True, cuda_graph=self.cuda_graph, @@ -146,6 +157,7 @@ class MlaTritonProvider(MlaAttentionProvider): """Portable SM90 provider with caller-independent decode workspace.""" name = "triton_sm90" + supports_batch_only_cuda_graph = True capabilities = AttentionKernelCapabilities( platforms=frozenset({PlatformEnum.CUDA}), compute_capabilities=frozenset({(9, 0)}), @@ -200,11 +212,19 @@ def __init__( self._runtime_fallback_reasons: dict[str, int] = {} def binding_metadata(self) -> dict[str, object]: - return { + metadata = { "implementation_kind": "atomic_provider", "implementation_source": "repo_triton", "decode_kernel_path": "triton_mla_stage1_stage2", } + if not self.spec.batch_only_cuda_graph: + return metadata + return { + **metadata, + "cuda_graph_shape_policy": "batch_only", + "context_capacity": self.spec.context_capacity, + "launch_plan_source": "batch_tp_heads_context_capacity", + } def _record_runtime_kernel_path(self, path: str) -> None: counts = getattr(self, "_runtime_kernel_path_counts", None) @@ -243,7 +263,7 @@ def runtime_kernel_stats(self) -> dict[str, object]: } @classmethod - def _contract_support( + def _common_contract_support( cls, spec: MlaAttentionOpSpec, caps: DeviceCaps, @@ -291,7 +311,9 @@ def supports( spec: MlaAttentionOpSpec, caps: DeviceCaps, ) -> SupportResult: - return cls._contract_support(spec, caps) + if spec.batch_only_cuda_graph and spec.context_capacity is None: + return SupportResult.unsupported("requires a static context capacity") + return cls._common_contract_support(spec, caps) def _validate_run_inputs( self, @@ -428,14 +450,21 @@ def _launch_config_for( ) -> MlaDecodeLaunchConfig: if self._fixed_launch_config is not None: return self._fixed_launch_config - context_capacity = ( - active_slot_width - if max_context_len is None - else int(max_context_len) - ) + if self.spec.batch_only_cuda_graph: + if self.spec.context_capacity is None: + raise RuntimeError( + "Batch-only MLA requires a static context capacity." + ) + context_capacity = self.spec.context_capacity + else: + context_capacity = ( + active_slot_width + if max_context_len is None + else int(max_context_len) + ) return select_glm_mla_decode_config( batch_size=batch_size, - max_context_len=context_capacity, + context_capacity=context_capacity, local_q_heads=self.spec.local_q_heads, ) @@ -494,6 +523,7 @@ class MlaSglFa3Provider(MlaTritonProvider): name = "sgl_fa3_sm90" supports_explicit_prefill = True + supports_batch_only_cuda_graph = True def __init__( self, @@ -521,10 +551,10 @@ def supports( spec: MlaAttentionOpSpec, caps: DeviceCaps, ) -> SupportResult: - base = cls._contract_support(spec, caps) + base = cls._common_contract_support(spec, caps) if not base.supported: return base - if spec.may_require_attention_scores: + if spec.score_output is not AttentionScoreKind.NONE: return SupportResult.unsupported( "does not satisfy the prepared score-output contract" ) @@ -719,9 +749,10 @@ def run_explicit_prefill( profile_only=True, ) class MlaTileLangScoreProvider(MlaSglFa3Provider): - """Explicit score-aware Composite over FA3, TileLang, and Triton.""" + """Score-aware Composite over FA3 and statically planned TileLang.""" name = "tilelang_score_sgl_fa3_h100" + supports_batch_only_cuda_graph = True def __init__( self, @@ -737,10 +768,22 @@ def __init__( max_batch_size=max_batch_size, launch_config=launch_config, ) + if self.spec.context_capacity is None: + raise ValueError( + "TileLang MLA requires a capture-time context capacity." + ) + self.tilelang_launch_plan = TileMlaLaunchPlan.build( + context_capacity=self.spec.context_capacity, + local_q_heads=self.spec.local_q_heads, + max_batch_size=self.max_batch_size, + need_score=True, + score_mode="per_head", + ) self.tilelang_score = TileMlaDecodeKernel( device=self.device, softmax_scale=self.spec.softmax_scale, valid_heads=self.spec.local_q_heads, + launch_plan=self.tilelang_launch_plan, ) @classmethod @@ -749,12 +792,16 @@ def supports( spec: MlaAttentionOpSpec, caps: DeviceCaps, ) -> SupportResult: - base = cls._contract_support(spec, caps) + base = cls._common_contract_support(spec, caps) if not base.supported: return base - if not spec.may_require_attention_scores: + if spec.score_output is not AttentionScoreKind.RAW_QK_PER_HEAD: + return SupportResult.unsupported( + "requires the RAW_QK_PER_HEAD decode score contract" + ) + if spec.context_capacity is None: return SupportResult.unsupported( - "score-capable Composite is not required by this operation" + "requires a capture-time context capacity" ) supported, reason = sgl_fa3_device_support(caps.device_index) if not supported: @@ -765,29 +812,49 @@ def supports( def binding_metadata(self) -> dict[str, object]: return { "implementation_kind": "composite_provider", - "implementation_source": "sglang-kernel+tilelang+repo_triton", + "implementation_source": "sglang-kernel+tilelang", "routes": { "score_free": "sgl_kernel.fa3.fwd", - "reduced_score": "tilelang_mla_decode", - "unsupported_score_contract": "triton_mla_stage1_stage2", + "raw_qk_per_head": "tilelang_mla_decode", }, + "tilelang_launch_plan": self.tilelang_launch_plan.metadata(), } - @staticmethod - def _tilelang_score_shape_supported( + def runtime_kernel_stats(self) -> dict[str, object]: + return { + **super().runtime_kernel_stats(), + "tilelang": self.tilelang_score.runtime_metadata(), + } + + def _validate_tilelang_score_contract( + self, attn_score: torch.Tensor, *, max_context_len: int | None, - ) -> bool: - score_capacity = int(attn_score.shape[1]) if attn_score.ndim >= 2 else 0 - return ( - attn_score.ndim == 2 - and attn_score.dtype == torch.float32 - and score_capacity > 0 - and score_capacity % 64 == 0 - and max_context_len is not None - and int(max_context_len) <= score_capacity - ) + ) -> None: + if attn_score.ndim != 3: + raise ValueError( + "TileLang MLA RAW_QK_PER_HEAD score must have shape " + f"[batch, heads, capacity], got {tuple(attn_score.shape)}." + ) + if int(attn_score.shape[1]) != self.spec.local_q_heads: + raise ValueError( + "TileLang MLA score head count does not match the bound TP " + f"shape: expected={self.spec.local_q_heads} " + f"got={attn_score.shape[1]}." + ) + if attn_score.dtype != torch.float32: + raise TypeError( + "TileLang MLA RAW_QK_PER_HEAD score must use FP32, got " + f"{attn_score.dtype}." + ) + if max_context_len is None or not 0 < int(max_context_len) <= int( + attn_score.shape[2] + ): + raise ValueError( + "TileLang MLA score capacity must cover max_context_len: " + f"max={max_context_len} capacity={attn_score.shape[2]}." + ) @staticmethod def _tilelang_layout_rejection_reason( @@ -797,19 +864,12 @@ def _tilelang_layout_rejection_reason( if not isinstance(view.payload, MlaLatentPayload): return "payload_type" attn_score = view.meta.attn_score - if ( - attn_score is not None - and attn_score.ndim == 2 - and int(attn_score.shape[1]) > int(view.meta.active_slots.shape[1]) - ): - return "score_capacity_exceeds_active_slots" tensors = { "latent_cache": view.payload.latent_cache, "rope_cache": view.payload.rope_cache, "active_slots": view.meta.active_slots, "request_indices": view.meta.req_indices, "context_lens": view.meta.context_lens, - "attn_score": view.meta.attn_score, "output": output, } rejected = [ @@ -840,34 +900,15 @@ def run( validation_scope=validation_scope, valid_batch_size=valid_batch_size, ) - # Per-head or non-tile-aligned score buffers remain on the existing - # Triton implementation. This is a static shape dispatch before any - # TileLang kernel launch, not an exception-driven runtime fallback. - if not self._tilelang_score_shape_supported( + self._validate_tilelang_score_contract( attn_score, max_context_len=view.meta.max_context_len, - ): - self._record_runtime_fallback("unsupported_score_shape") - return MlaTritonProvider.run( - self, - q_nope_absorbed, - q_rope, - view, - output, - validation_scope=validation_scope, - valid_batch_size=valid_batch_size, - ) + ) layout_rejection = self._tilelang_layout_rejection_reason(view, output) if layout_rejection is not None: - self._record_runtime_fallback(layout_rejection) - return MlaTritonProvider.run( - self, - q_nope_absorbed, - q_rope, - view, - output, - validation_scope=validation_scope, - valid_batch_size=valid_batch_size, + raise ValueError( + "TileLang MLA runtime view violates the bound layout contract: " + f"{layout_rejection}." ) payload = self._validate_run_inputs( q_nope_absorbed, diff --git a/src/sparsevllm/platforms/cuda.py b/src/sparsevllm/platforms/cuda.py index e2b347b9..ac4b3a9b 100644 --- a/src/sparsevllm/platforms/cuda.py +++ b/src/sparsevllm/platforms/cuda.py @@ -65,6 +65,15 @@ def barrier_device_ids(self, rank: int) -> list[int] | None: def get_device_caps(self, device_index: int = 0) -> DeviceCaps: device_index = int(device_index) major, minor = torch.cuda.get_device_capability(device_index) + try: + multi_processor_count = int( + torch.cuda.get_device_properties(device_index).multi_processor_count + ) + except (AssertionError, RuntimeError): + # Capability-only unit tests may stub the public capability probes + # without initializing a CUDA driver. This optional performance + # fact is resolved on real devices and may remain unknown otherwise. + multi_processor_count = None return DeviceCaps( platform=self.enum, device_type=self.device_type, @@ -79,6 +88,7 @@ def get_device_caps(self, device_index: int = 0) -> DeviceCaps: supports_bfloat16=(int(major), int(minor)) >= (8, 0), # Ada (SM89), Hopper and Blackwell provide native FP8 tensor cores. supports_native_fp8=(int(major), int(minor)) >= (8, 9), + multi_processor_count=multi_processor_count, ) def get_default_attention_backend(self) -> str: diff --git a/src/sparsevllm/platforms/interface.py b/src/sparsevllm/platforms/interface.py index e382aac6..fd8a2860 100644 --- a/src/sparsevllm/platforms/interface.py +++ b/src/sparsevllm/platforms/interface.py @@ -37,6 +37,7 @@ class DeviceCaps: supports_pin_memory: bool = False supports_bfloat16: bool = False supports_native_fp8: bool = False + multi_processor_count: int | None = None class Platform: diff --git a/tests/test_batch_only_decode_graph.py b/tests/test_batch_only_decode_graph.py new file mode 100644 index 00000000..76859714 --- /dev/null +++ b/tests/test_batch_only_decode_graph.py @@ -0,0 +1,659 @@ +from types import SimpleNamespace + +import pytest +import torch + +from sparsevllm.configs.cuda_graph import ( + _normalize_decode_graph_shape_policy, + build_decode_cuda_graph_batch_only_startup_plan, +) +from sparsevllm.engine.decode_cuda_graph import DecodeCudaGraphRunner +from sparsevllm.engine.decode_graph_contract import ( + DecodeGraphContract, + DecodeGraphInputs, + DecodeGraphState, +) +from sparsevllm.engine.runtime_state import RuntimeState +from sparsevllm.kernels.triton.paged_flash_decoding import ( + paged_flash_decode, +) +from sparsevllm.kernels.triton.sglang_gemma4_decode_attention import ( + sglang_gemma4_decode, +) +from sparsevllm.operators.decode_attention import ( + FixedGridTritonPagedDecodeAttentionProvider, + DECODE_ATTENTION_REGISTRY, + DecodeAttentionOpSpec, + TritonPagedDecodeAttentionProvider, + build_graph_stable_decode_launch_plan, +) +from sparsevllm.operators.registry import OpResolver +from sparsevllm.operators.gemma4 import Gemma4OpSpec, TritonGemma4OperatorProvider +from sparsevllm.operators.mla_attention import ( + MlaAttentionOpSpec, + MlaTritonProvider, +) +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum + + +def _cuda_caps() -> DeviceCaps: + return DeviceCaps( + platform=PlatformEnum.CUDA, + device_type="cuda", + device_index=0, + device_name="test sm90", + compute_capability=(9, 0), + runtime_version="12.8", + supports_graph_capture=True, + supports_triton=True, + supports_bfloat16=True, + multi_processor_count=120, + ) + + +def test_batch_only_policy_aliases_and_rejects_unknown_values() -> None: + assert _normalize_decode_graph_shape_policy("batch") == "batch_only" + assert _normalize_decode_graph_shape_policy(None) == "batch_only" + with pytest.raises(ValueError, match="shape_policy"): + _normalize_decode_graph_shape_policy("sequence_only") + + +def test_batch_only_startup_plan_has_one_graph_per_batch_and_path() -> None: + config = SimpleNamespace( + decode_graph_capture_sizes=[1, 4], + decode_graph_startup_capture_limit=8, + decode_graph_max_cached_graphs=8, + sparse_method="quest", + max_model_len=32768, + sink_keep_tokens=64, + decode_keep_tokens=4096, + recent_keep_tokens=512, + ) + assert build_decode_cuda_graph_batch_only_startup_plan(config) == [ + (4, 32768, True), + (4, 4672, False), + (1, 32768, True), + (1, 4672, False), + ] + + +def test_batch_only_state_identity_omits_context_capacity() -> None: + runner = object.__new__(DecodeCudaGraphRunner) + runner.shape_policy = "batch_only" + runner.startup_plan_sealed = False + runner._graphs = {} + runner.max_cached_graphs = None + runner.cache_manager = SimpleNamespace(device=torch.device("cpu")) + runner.eviction_count = 0 + + state = runner._select_state( + method="quest", + batch_size=4, + context_capacity=32768, + is_long_text=True, + capture_sampling=False, + graph_path_id="long", + ) + reused = runner._select_state( + method="quest", + batch_size=4, + context_capacity=8192, + is_long_text=True, + capture_sampling=False, + graph_path_id="long", + ) + assert reused is state + assert state.key.context_capacity == 0 + assert state.capture_context_capacity == 32768 + assert state.decode_state is not None + assert state.decode_state.contract == DecodeGraphContract( + method="quest", + shape_policy="batch_only", + topology_path_id="long", + batch_capacity=4, + context_capacity=32768, + ) + assert state.decode_state.inputs.batch_capacity == 4 + assert state.decode_state.contract.capability_level == "path_scoped" + + reused_with_default_path = runner._select_state( + method="quest", + batch_size=4, + context_capacity=16384, + is_long_text=True, + capture_sampling=False, + ) + assert reused_with_default_path is state + + with pytest.raises(RuntimeError, match="exceeded captured path capacity"): + runner._select_state( + method="quest", + batch_size=4, + context_capacity=65536, + is_long_text=True, + capture_sampling=False, + graph_path_id="long", + ) + + +def test_typed_decode_graph_participant_delegates_to_cache_owner() -> None: + calls = [] + private_keepalive = torch.empty(1) + operator_keepalive = torch.empty(1) + + class CacheOwner: + num_free_slots = 16 + + def init_decode_graph_state(self, contract, inputs): + calls.append(("init", contract.topology_path_id)) + return SimpleNamespace(contract=contract, inputs=inputs) + + def prepare_decode_graph_step(self, seqs, state): + calls.append(("prepare_out", len(seqs))) + state.inputs.input_ids.fill_(7) + + def prepare_decode_graph_in(self, state): + calls.append(("prepare_in", state.contract.topology_path_id)) + + def decode_graph_state_keepalive_tensors(self, state): + calls.append(("keepalive", state.contract.topology_path_id)) + return [private_keepalive] + + class OperatorOwner: + def init_decode_graph_state(self, contract, inputs): + calls.append(("operator_init", contract.batch_capacity)) + return SimpleNamespace(contract=contract, inputs=inputs) + + def prepare_decode_graph_out(self, state): + calls.append(("operator_out", state.contract.context_capacity)) + + def prepare_decode_graph_in(self, state): + calls.append(("operator_in", state.contract.topology_path_id)) + + def decode_graph_keepalive_tensors(self, state): + calls.append(("operator_keepalive", state.contract.topology_path_id)) + return [operator_keepalive] + + def close_decode_graph_state(self, state): + calls.append(("operator_close", state.contract.batch_capacity)) + + contract = DecodeGraphContract( + method="", + shape_policy="batch_only", + topology_path_id="dense", + batch_capacity=2, + context_capacity=32, + ) + graph_state = DecodeGraphState( + contract=contract, + inputs=DecodeGraphInputs.allocate( + contract, + device=torch.device("cpu"), + pin_memory=False, + ), + ) + runtime = RuntimeState( + SimpleNamespace(), + CacheOwner(), + decode_graph_participants=(OperatorOwner(),), + ) + + participant = runtime.init_decode_graph_state(graph_state) + runtime.prepare_decode_graph_step([object()], graph_state) + participant.prepare_in_graph() + keepalive = graph_state.keepalive_tensors() + + assert graph_state.runtime_state is participant + assert graph_state.inputs.input_ids.tolist() == [7, 7] + assert any(tensor is private_keepalive for tensor in keepalive) + assert any(tensor is operator_keepalive for tensor in keepalive) + assert calls == [ + ("init", "dense"), + ("operator_init", 2), + ("prepare_out", 1), + ("operator_out", 32), + ("prepare_in", "dense"), + ("operator_in", "dense"), + ("keepalive", "dense"), + ("operator_keepalive", "dense"), + ] + graph_state.close() + assert calls[-1] == ("operator_close", 2) + + +def test_mha_resolver_prefers_sgl_fa3_for_batch_only_on_supported_sm90() -> None: + spec = DecodeAttentionOpSpec( + num_query_heads=8, + num_kv_heads=2, + head_dim=128, + activation_dtype=torch.bfloat16, + softmax_scale=128**-0.5, + max_batch_size=8, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + caps = _cuda_caps() + assert FixedGridTritonPagedDecodeAttentionProvider.supports(spec, caps).supported + assert not TritonPagedDecodeAttentionProvider.supports(spec, caps).supported + + h2o_spec = DecodeAttentionOpSpec( + num_query_heads=8, + num_kv_heads=2, + head_dim=128, + activation_dtype=torch.bfloat16, + softmax_scale=128**-0.5, + max_batch_size=8, + may_require_attention_scores=True, + h2o_layerwise_probability_scores=True, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + assert FixedGridTritonPagedDecodeAttentionProvider.supports( + h2o_spec, caps + ).supported + + unsupported = DecodeAttentionOpSpec( + **{ + **spec.__dict__, + "activation_dtype": torch.float32, + } + ) + assert not FixedGridTritonPagedDecodeAttentionProvider.supports( + unsupported, + caps, + ).supported + + plan = build_graph_stable_decode_launch_plan(spec, caps) + assert plan.context_capacity == spec.context_capacity + assert plan.max_kv_splits > 0 + assert plan.target_tokens_per_split > 0 + assert plan.block_n > 0 + from unittest.mock import patch + + with patch( + "sparsevllm.operators.decode_attention.sgl_fa3_device_support", + return_value=(True, "available"), + ): + resolved = OpResolver(DECODE_ATTENTION_REGISTRY).resolve(spec, caps) + assert resolved.provider.name == "sgl_fa3_paged_decode_sm90" + assert resolved.report.selection_basis == "upstream_default" + + +def test_mha_resolver_falls_back_to_fixed_grid_when_upstream_is_ineligible() -> None: + spec = DecodeAttentionOpSpec( + num_query_heads=8, + num_kv_heads=2, + head_dim=128, + activation_dtype=torch.bfloat16, + softmax_scale=128**-0.5, + max_batch_size=8, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + caps = DeviceCaps( + **{ + **_cuda_caps().__dict__, + "compute_capability": (8, 0), + } + ) + from unittest.mock import patch + + with patch( + "sparsevllm.operators.decode_attention.flashinfer_paged_decode_support", + return_value=(False, "unavailable"), + ): + resolved = OpResolver(DECODE_ATTENTION_REGISTRY).resolve(spec, caps) + assert isinstance( + resolved.provider, + FixedGridTritonPagedDecodeAttentionProvider, + ) + metadata = resolved.report.as_dict()["provider_metadata"] + assert metadata["launch_plan"]["plan_id"] == "portable_fixed_grid_v1" + + +def test_mla_resolver_contract_selects_fixed_launch_provider() -> None: + 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=2, + cuda_graph=True, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + caps = _cuda_caps() + assert MlaTritonProvider.supports(spec, caps).supported + + +def test_gemma4_resolver_contract_selects_fixed_grid_provider() -> None: + spec = Gemma4OpSpec( + activation_dtype=torch.bfloat16, + head_dims=(256, 512), + cuda_graph=True, + attention_contracts=((8, 2, 256, 1023), (8, 1, 512, -1)), + max_batch_size=8, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + caps = _cuda_caps() + assert TritonGemma4OperatorProvider.supports(spec, caps).supported + provider = TritonGemma4OperatorProvider.bind(spec, caps) + assert provider.name == "triton" + assert provider.binding_metadata()["attention_dispatch"]["decode_routes"] == [ + "sglang_fixed_grid" + ] + + +@pytest.mark.parametrize("multi_processor_count", [None, 0, -1]) +def test_gemma4_provider_rejects_missing_multi_processor_count( + multi_processor_count, +) -> None: + spec = Gemma4OpSpec( + activation_dtype=torch.bfloat16, + head_dims=(256,), + cuda_graph=True, + attention_contracts=((8, 2, 256, -1),), + max_batch_size=8, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + caps = DeviceCaps( + **{ + **_cuda_caps().__dict__, + "multi_processor_count": multi_processor_count, + } + ) + result = TritonGemma4OperatorProvider.supports(spec, caps) + assert not result.supported + assert "multi-processor count" in result.reason + with pytest.raises(ValueError, match="multi-processor count"): + TritonGemma4OperatorProvider.bind(spec, caps) + + +def _decode_reference( + q, k, v, slots, req_indices, lengths, window=None, *, scale=True +): + output = torch.empty_like(q) + lse = torch.empty( + q.shape[1], q.shape[0], dtype=torch.float32, device=q.device + ) + group_size = q.shape[1] // k.shape[1] + for batch, length in enumerate(lengths.tolist()): + start = max(0, length - int(window or length)) + indices = slots[req_indices[batch], start:length].long() + keys = k[indices].repeat_interleave(group_size, dim=1) + values = v[indices].repeat_interleave(group_size, dim=1) + logits = torch.einsum("hd,lhd->hl", q[batch].float(), keys.float()) + if scale: + logits = logits / q.shape[-1] ** 0.5 + probabilities = logits.softmax(-1) + output[batch] = torch.einsum( + "hl,lhd->hd", probabilities, values.float() + ).to(q.dtype) + lse[:, batch] = logits.logsumexp(-1) + return output, lse + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires idle CUDA GPU") +@pytest.mark.parametrize( + ("dtype", "heads", "kv_heads", "head_dim"), + [ + (torch.bfloat16, 4, 4, 64), + (torch.float16, 8, 2, 64), + (torch.bfloat16, 8, 2, 128), + (torch.float16, 4, 4, 128), + (torch.bfloat16, 8, 2, 256), + (torch.float16, 4, 4, 256), + ], +) +def test_batch_only_mha_matches_reference_and_replays_new_lengths( + dtype, + heads, + kv_heads, + head_dim, +) -> None: + torch.manual_seed(11) + batch, capacity = 2, 257 + device = torch.device("cuda") + q = torch.randn(batch, heads, head_dim, dtype=dtype, device=device) + k = torch.randn( + batch * capacity, kv_heads, head_dim, dtype=dtype, device=device + ) + v = torch.randn_like(k) + slots = torch.arange( + batch * capacity, dtype=torch.int32, device=device + ).view(batch, capacity) + req_indices = torch.arange(batch, dtype=torch.int32, device=device) + lengths = torch.tensor([129, 257], dtype=torch.int32, device=device) + mid_o = torch.empty( + batch, heads, 8, head_dim, dtype=torch.float32, device=device + ) + mid_lse = torch.empty(batch, heads, 8, dtype=torch.float32, device=device) + output_lse = torch.empty(heads, batch, dtype=torch.float32, device=device) + + def run(): + return paged_flash_decode( + q, + k, + v, + slots, + req_indices, + lengths, + mid_o, + mid_lse, + target_tokens_per_split=64, + return_softmax_lse=True, + output_lse=output_lse, + ) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output, graph_lse = run() + lengths.copy_(torch.tensor([17, 201], dtype=torch.int32, device=device)) + q.copy_(torch.randn_like(q)) + graph.replay() + expected_output, expected_lse = _decode_reference( + q, k, v, slots, req_indices, lengths + ) + torch.cuda.synchronize() + torch.testing.assert_close(graph_output, expected_output, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(graph_lse, expected_lse, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires idle CUDA GPU") +def test_batch_only_gqa_replays_exact_context_capacity() -> None: + torch.manual_seed(29) + batch, heads, kv_heads, head_dim, capacity = 1, 8, 2, 128, 8352 + device = torch.device("cuda") + q = torch.randn(batch, heads, head_dim, dtype=torch.bfloat16, device=device) + k = torch.randn( + capacity, + kv_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + v = torch.randn_like(k) + slots = torch.arange(capacity, dtype=torch.int32, device=device).view( + batch, + capacity, + ) + req_indices = torch.zeros(batch, dtype=torch.int32, device=device) + lengths = torch.tensor([4097], dtype=torch.int32, device=device) + mid_o = torch.empty( + batch, + heads, + 16, + head_dim, + dtype=torch.float32, + device=device, + ) + mid_lse = torch.empty(batch, heads, 16, dtype=torch.float32, device=device) + output_lse = torch.empty(heads, batch, dtype=torch.float32, device=device) + + def run(): + return paged_flash_decode( + q, + k, + v, + slots, + req_indices, + lengths, + mid_o, + mid_lse, + target_tokens_per_split=256, + return_softmax_lse=True, + output_lse=output_lse, + ) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output, graph_lse = run() + lengths.fill_(capacity) + q.copy_(torch.randn_like(q)) + graph.replay() + expected_output, expected_lse = _decode_reference( + q, + k, + v, + slots, + req_indices, + lengths, + ) + torch.cuda.synchronize() + torch.testing.assert_close(graph_output, expected_output, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(graph_lse, expected_lse, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires idle CUDA GPU") +def test_batch_only_gqa_produces_raw_per_head_scores() -> None: + torch.manual_seed(23) + batch, heads, kv_heads, head_dim, capacity = 2, 4, 2, 64, 33 + device = torch.device("cuda") + q = torch.randn(batch, heads, head_dim, dtype=torch.bfloat16, device=device) + k = torch.randn( + batch * capacity, + kv_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + v = torch.randn_like(k) + slots = torch.arange( + batch * capacity, + dtype=torch.int32, + device=device, + ).view(batch, capacity) + req_indices = torch.arange(batch, dtype=torch.int32, device=device) + lengths = torch.tensor([17, 29], dtype=torch.int32, device=device) + mid_o = torch.empty( + batch, + heads, + 8, + head_dim, + dtype=torch.float32, + device=device, + ) + mid_lse = torch.empty(batch, heads, 8, dtype=torch.float32, device=device) + scores = torch.full( + (batch, heads, capacity), + -torch.inf, + dtype=torch.float32, + device=device, + ) + + output = paged_flash_decode( + q, + k, + v, + slots, + req_indices, + lengths, + mid_o, + mid_lse, + attn_score=scores, + target_tokens_per_split=8, + ) + expected, _ = _decode_reference(q, k, v, slots, req_indices, lengths) + torch.testing.assert_close(output, expected, rtol=2e-2, atol=2e-2) + + group_size = heads // kv_heads + for batch_idx, length in enumerate(lengths.tolist()): + keys = k[slots[batch_idx, :length].long()].repeat_interleave( + group_size, + dim=1, + ) + expected_scores = torch.einsum( + "hd,lhd->hl", + q[batch_idx].float(), + keys.float(), + ) + torch.testing.assert_close( + scores[batch_idx, :, :length], + expected_scores, + rtol=2e-2, + atol=2e-2, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires idle CUDA GPU") +@pytest.mark.parametrize("window", [None, 8]) +def test_gemma4_fixed_grid_matches_reference_and_graph(window) -> None: + torch.manual_seed(19) + device = torch.device("cuda") + batch, heads, kv_heads, head_dim, capacity = 2, 4, 2, 256, 33 + q = torch.randn(batch, heads, head_dim, dtype=torch.bfloat16, device=device) + k = torch.randn( + batch * capacity, kv_heads, head_dim, dtype=torch.bfloat16, device=device + ) + v = torch.randn_like(k) + slots = torch.arange( + batch * capacity, dtype=torch.int32, device=device + ).view(batch, capacity) + req_indices = torch.arange(batch, dtype=torch.int32, device=device) + lengths = torch.tensor([33, 21], dtype=torch.int32, device=device) + mid_o = torch.empty( + batch, heads, 8, head_dim, dtype=torch.float32, device=device + ) + mid_lse = torch.empty(batch, heads, 8, dtype=torch.float32, device=device) + splits = torch.empty(batch, dtype=torch.int32, device=device) + + def run(): + return sglang_gemma4_decode( + q, + k, + v, + slots, + req_indices, + lengths, + mid_o, + mid_lse, + splits, + sliding_window=window, + multi_processor_count=120, + ) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = run() + lengths.copy_(torch.tensor([17, 29], dtype=torch.int32, device=device)) + q.copy_(torch.randn_like(q)) + graph.replay() + expected, _ = _decode_reference( + q, k, v, slots, req_indices, lengths, window, scale=False + ) + torch.cuda.synchronize() + cosine = torch.nn.functional.cosine_similarity( + graph_output.float().flatten(), expected.float().flatten(), dim=0 + ) + assert cosine > 0.999 diff --git a/tests/test_decode_attention_provider.py b/tests/test_decode_attention_provider.py index d0fae0d4..05579717 100644 --- a/tests/test_decode_attention_provider.py +++ b/tests/test_decode_attention_provider.py @@ -5,6 +5,10 @@ import pytest import torch +from sparsevllm.engine.decode_graph_contract import ( + DecodeGraphContract, + DecodeGraphInputs, +) from sparsevllm.kernels.external.flashinfer.decode import ( flashinfer_paged_decode_support, ) @@ -12,13 +16,17 @@ from sparsevllm.method_registry import sparse_decode_attention_requires_scores from sparsevllm.models.attention_runtime import build_mha_decode_attention_spec from sparsevllm.operators.decode_attention import ( + DECODE_ATTENTION_REGISTRY, + DeltaKVFixedGridDecodeAttentionProvider, DecodeAttentionRunResult, DecodeAttentionOpSpec, FlashInferPagedDecodeAttentionProvider, + FixedGridTritonPagedDecodeAttentionProvider, PreparedDecodeAttentionOp, SglFa3PagedDecodeAttentionProvider, TritonPagedDecodeAttentionProvider, ) +from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum @@ -66,6 +74,134 @@ def test_h2o_runtime_decode_spec_is_score_free_while_eviction_is_disabled(): assert not spec.kernel_request.requires_softmax_lse +def test_batch_only_decode_spec_carries_static_context_capacity(): + config = SimpleNamespace( + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + torch_dtype=torch.bfloat16, + ) + runtime_config = SimpleNamespace( + decode_graph_shape_policy="batch_only", + max_model_len=40960, + ) + + spec = build_mha_decode_attention_spec( + config, + sparse_method="vanilla", + attention_tp_size=1, + max_batch_size=8, + cuda_graph=True, + runtime_config=runtime_config, + ) + + assert spec.batch_only_cuda_graph + assert spec.context_capacity == runtime_config.max_model_len + + +def test_deltakv_kivi_decode_spec_carries_mixed_storage_contract(): + config = SimpleNamespace( + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + torch_dtype=torch.bfloat16, + ) + runtime_config = SimpleNamespace( + decode_graph_shape_policy="batch_only", + max_model_len=131072, + full_layer_kv_quant_bits=4, + enable_full_layer_kivi_quant=True, + full_layer_kivi_decode_block_seq=512, + full_layer_kivi_decode_block_n=32, + full_layer_kivi_decode_num_warps=4, + full_layer_kivi_decode_num_stages=2, + ) + + spec = build_mha_decode_attention_spec( + config, + sparse_method="deltakv", + attention_tp_size=1, + max_batch_size=32, + cuda_graph=True, + runtime_config=runtime_config, + ) + + assert spec.may_use_full_layer_kivi_int4 + assert spec.full_layer_kivi_decode_block_seq == 512 + assert spec.full_layer_kivi_decode_block_n == 32 + assert spec.full_layer_kivi_decode_num_warps == 4 + assert spec.full_layer_kivi_decode_num_stages == 2 + assert spec.may_require_attention_scores + assert spec.layer_varying_page_table + + +def test_deltakv_kivi_batch_only_resolves_nonstandard_fixed_grid_provider(): + spec = _spec( + may_require_attention_scores=True, + layer_varying_page_table=True, + batch_only_cuda_graph=True, + context_capacity=131072, + may_use_full_layer_kivi_int4=True, + full_layer_kivi_decode_block_seq=512, + full_layer_kivi_decode_block_n=32, + full_layer_kivi_decode_num_warps=4, + full_layer_kivi_decode_num_stages=2, + ) + caps = _cuda_caps() + + resolved = OpResolver(DECODE_ATTENTION_REGISTRY).resolve(spec, caps) + + assert isinstance( + resolved.provider, + DeltaKVFixedGridDecodeAttentionProvider, + ) + assert resolved.report.selection_basis == "semantic_fallback" + assert not FixedGridTritonPagedDecodeAttentionProvider.supports( + spec, + caps, + ).supported + assert not SglFa3PagedDecodeAttentionProvider.supports(spec, caps).supported + assert not FlashInferPagedDecodeAttentionProvider.supports(spec, caps).supported + + +def test_deltakv_kivi_bucketed_keeps_legacy_triton_baseline(): + spec = _spec( + may_require_attention_scores=True, + layer_varying_page_table=True, + batch_only_cuda_graph=False, + context_capacity=131072, + may_use_full_layer_kivi_int4=True, + ) + + assert TritonPagedDecodeAttentionProvider.supports(spec, _cuda_caps()).supported + assert not DeltaKVFixedGridDecodeAttentionProvider.supports( + spec, + _cuda_caps(), + ).supported + + +def test_deltakv_kivi_fixed_grid_rejects_unsupported_head_dim(): + spec = _spec( + head_dim=256, + may_require_attention_scores=True, + layer_varying_page_table=True, + batch_only_cuda_graph=True, + context_capacity=131072, + may_use_full_layer_kivi_int4=True, + full_layer_kivi_decode_block_seq=512, + full_layer_kivi_decode_block_n=32, + full_layer_kivi_decode_num_warps=4, + full_layer_kivi_decode_num_stages=2, + ) + + support = DeltaKVFixedGridDecodeAttentionProvider.supports( + spec, + _cuda_caps(), + ) + + assert not support.supported + + def _cuda_caps( *, device_name: str = "NVIDIA H100 80GB HBM3", @@ -104,7 +240,79 @@ def _spec(**overrides) -> DecodeAttentionOpSpec: return DecodeAttentionOpSpec(**values) -def test_flashinfer_lse_decode_rejects_cuda_graph_before_dependency_probe(): +def test_deltakv_fixed_grid_graph_state_owns_per_graph_workspace(): + spec = _spec( + max_batch_size=8, + may_require_attention_scores=True, + layer_varying_page_table=True, + batch_only_cuda_graph=True, + context_capacity=131072, + may_use_full_layer_kivi_int4=True, + full_layer_kivi_decode_block_seq=512, + full_layer_kivi_decode_block_n=32, + full_layer_kivi_decode_num_warps=4, + full_layer_kivi_decode_num_stages=2, + ) + provider = DeltaKVFixedGridDecodeAttentionProvider.bind(spec, _cuda_caps()) + long_contract = DecodeGraphContract( + method="deltakv", + shape_policy="batch_only", + topology_path_id="long", + batch_capacity=4, + context_capacity=131072, + ) + short_contract = DecodeGraphContract( + method="deltakv", + shape_policy="batch_only", + topology_path_id="short", + batch_capacity=4, + context_capacity=8192, + ) + long_inputs = DecodeGraphInputs.allocate( + long_contract, + device=torch.device("cpu"), + pin_memory=False, + ) + short_inputs = DecodeGraphInputs.allocate( + short_contract, + device=torch.device("cpu"), + pin_memory=False, + ) + + long_state = provider.init_decode_graph_state(spec, long_contract, long_inputs) + short_state = provider.init_decode_graph_state(spec, short_contract, short_inputs) + + assert long_state.launch_plan.context_capacity == 131072 + assert short_state.launch_plan.context_capacity == 8192 + assert long_state.kivi_launch_plan.target_tokens_per_split == 512 + assert long_state.kivi_launch_plan.block_n == 32 + assert long_state.kivi_launch_plan.stage1_num_warps == 4 + assert long_state.kivi_launch_plan.stage1_num_stages == 2 + assert long_state.launch_plan.max_kv_splits > short_state.launch_plan.max_kv_splits + assert ( + long_state.kivi_launch_plan.max_kv_splits + >= short_state.kivi_launch_plan.max_kv_splits + ) + assert ( + long_state.kivi_mid_o.shape[2] + == long_state.kivi_launch_plan.max_kv_splits + ) + assert ( + short_state.kivi_mid_o.shape[2] + == short_state.kivi_launch_plan.max_kv_splits + ) + assert long_state.mid_o.data_ptr() != short_state.mid_o.data_ptr() + assert long_state.kivi_mid_o.data_ptr() != long_state.mid_o.data_ptr() + assert len(provider.decode_graph_keepalive_tensors(long_state)) == 6 + provider.prepare_decode_graph_out(long_state) + assert provider._active_graph_state is long_state + provider.prepare_decode_graph_in(short_state) + assert provider._active_graph_state is short_state + provider.close_decode_graph_state(short_state) + assert provider._active_graph_state is None + + +def test_flashinfer_lse_decode_accepts_cuda_graph_contract(): spec = _spec( may_require_attention_scores=True, layer_varying_page_table=True, @@ -116,13 +324,13 @@ def test_flashinfer_lse_decode_rejects_cuda_graph_before_dependency_probe(): ) with patch( - "sparsevllm.operators.decode_attention.flashinfer_paged_decode_support" + "sparsevllm.operators.decode_attention.flashinfer_paged_decode_support", + return_value=(True, "available"), ) as support: result = FlashInferPagedDecodeAttentionProvider.supports(spec, caps) - assert not result.supported - assert "CUDA Graph" in result.reason - support.assert_not_called() + assert result.supported + support.assert_called_once_with() def test_prepared_h2o_decode_applies_fixed_probability_scorer(): @@ -437,6 +645,123 @@ def test_sgl_decode_provider_uses_prepared_explicit_kv_adapter(): decode_launch_op.launch_config.assert_not_called() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_flashinfer_graph_decode_replans_and_replays_new_metadata(): + flashinfer_paged_decode_support() + torch.manual_seed(20260825) + device = torch.device("cuda") + batch, query_heads, kv_heads, head_dim, capacity = 2, 8, 2, 128, 17 + spec = _spec( + num_query_heads=query_heads, + num_kv_heads=kv_heads, + head_dim=head_dim, + activation_dtype=torch.bfloat16, + softmax_scale=head_dim**-0.5, + max_batch_size=batch, + cuda_graph=True, + batch_only_cuda_graph=True, + context_capacity=capacity, + ) + contract = DecodeGraphContract( + method="", + shape_policy="batch_only", + topology_path_id="dense", + batch_capacity=batch, + context_capacity=capacity, + ) + inputs = DecodeGraphInputs.allocate( + contract, + device=device, + pin_memory=False, + ) + q = torch.randn( + batch, + query_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + slots = 3 * capacity + k_cache = torch.randn( + slots, + kv_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + v_cache = torch.randn_like(k_cache) + page_table = torch.arange( + slots, + dtype=torch.int32, + device=device, + ).view(3, capacity) + view = SimpleNamespace( + payload=SimpleNamespace(k_cache=k_cache, v_cache=v_cache), + meta=SimpleNamespace( + active_slots=page_table, + req_indices=inputs.request_indices, + context_lens=inputs.context_lens, + attn_score=None, + ), + ) + provider = FlashInferPagedDecodeAttentionProvider() + provider.prepare(spec, device_index=torch.cuda.current_device()) + state = provider.init_decode_graph_state(spec, contract, inputs) + + def update_metadata(lengths, rows) -> None: + inputs.host.context_lens.copy_(torch.tensor(lengths, dtype=torch.int32)) + inputs.host.request_indices.copy_(torch.tensor(rows, dtype=torch.int32)) + inputs.context_lens.copy_(inputs.host.context_lens) + inputs.request_indices.copy_(inputs.host.request_indices) + provider.prepare_decode_graph_out(state) + provider.prepare_decode_graph_in(state) + + def reference() -> torch.Tensor: + expected = [] + group_size = query_heads // kv_heads + for batch_idx in range(batch): + length = int(inputs.host.context_lens[batch_idx]) + row = int(inputs.host.request_indices[batch_idx]) + active = page_table[row, :length].long() + keys = k_cache[active].repeat_interleave(group_size, dim=1) + values = v_cache[active].repeat_interleave(group_size, dim=1) + logits = torch.einsum( + "hd,lhd->hl", + q[batch_idx].float(), + keys.float(), + ) * spec.softmax_scale + expected.append( + torch.einsum( + "hl,lhd->hd", + torch.softmax(logits, dim=-1), + values.float(), + ).to(q.dtype) + ) + return torch.stack(expected) + + try: + update_metadata([17, 13], [2, 0]) + provider.run(spec, q, view) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = provider.run(spec, q, view) + + update_metadata([9, 16], [1, 2]) + expected = reference() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + graph_output, + expected, + rtol=3e-2, + atol=3e-2, + ) + finally: + provider.close_decode_graph_state(state) + provider.close() + + def test_triton_provider_owns_launch_config_and_workspace_preparation(): provider = TritonPagedDecodeAttentionProvider() provider._backend = Mock(name="triton_backend") diff --git a/tests/test_deltakv_less_memory_kernel.py b/tests/test_deltakv_less_memory_kernel.py index 6c311c17..e2eb8b09 100644 --- a/tests/test_deltakv_less_memory_kernel.py +++ b/tests/test_deltakv_less_memory_kernel.py @@ -21,6 +21,9 @@ from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import ( flash_decode_stage1_with_score as gqa_flash_decode_stage1_with_score, ) +from sparsevllm.kernels.triton.paged_flash_decoding import ( + fixed_grid_flash_decode_stage2, +) from sparsevllm.kernels.triton.quant import ( triton_dequantize_2d_int4_grouped, triton_quantize_and_pack_2d_int4_grouped, @@ -1268,6 +1271,107 @@ def test_full_layer_kivi_flash_decode_stage1_matches_dense_stage1(self): self.assertTrue(torch.allclose(lse_token_group, lse_ref, atol=3e-2, rtol=3e-2)) self.assertTrue(torch.allclose(score_token_group, score_ref, atol=3e-2, rtol=3e-2)) + max_kv_splits = 4 + target_tokens_per_split = 16 + mid_fixed = torch.empty( + (batch, num_heads, max_kv_splits, head_dim), + device=device, + dtype=torch.float32, + ) + lse_fixed = torch.empty( + (batch, num_heads, max_kv_splits), + device=device, + dtype=torch.float32, + ) + out_fixed = torch.empty_like(q) + out_lse_fixed = torch.empty( + (num_heads, batch), + device=device, + dtype=torch.float32, + ) + score_fixed = torch.full_like(score_ref, -1e20) + + def run_fixed_grid(): + full_layer_kivi_flash_decode_stage1( + q=q, + raw_k=raw_k, + raw_v=raw_v, + raw_slots_map=raw_slots_map, + kivi_block_slots_map=kivi_block_slots_map, + kivi_block_start_pos=kivi_block_start_pos, + key_packed=key_packed, + key_scales=key_scales, + key_mins=key_mins, + value_packed=value_packed, + value_scales=value_scales, + value_mins=value_mins, + req_indices=req_indices, + context_lens=context_lens, + max_len_in_batch=seq_len, + mid_out=mid_fixed, + mid_out_logsumexp=lse_fixed, + group_size=group_size, + block_seq=block_seq, + attn_score=score_fixed, + max_kv_splits=max_kv_splits, + target_tokens_per_split=target_tokens_per_split, + ) + fixed_grid_flash_decode_stage2( + mid_fixed, + lse_fixed, + context_lens, + out_fixed, + out_lse_fixed, + target_tokens_per_split=target_tokens_per_split, + ) + + run_fixed_grid() + torch.cuda.synchronize() + output_ptr = out_fixed.data_ptr() + workspace_ptrs = (mid_fixed.data_ptr(), lse_fixed.data_ptr()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_fixed_grid() + + context_lens.copy_( + torch.tensor([17, seq_len], device=device, dtype=torch.int32) + ) + q.copy_(torch.randn_like(q)) + graph.replay() + torch.cuda.synchronize() + + expected = torch.empty_like(q) + for batch_idx, length in enumerate(context_lens.tolist()): + keys = dense_k[ + batch_idx * seq_len : batch_idx * seq_len + length + ].repeat_interleave(num_heads // num_kv_heads, dim=1) + values = dense_v[ + batch_idx * seq_len : batch_idx * seq_len + length + ].repeat_interleave(num_heads // num_kv_heads, dim=1) + logits = torch.einsum( + "hd,lhd->hl", + q[batch_idx].float(), + keys.float(), + ) + expected[batch_idx] = torch.einsum( + "hl,lhd->hd", + (logits / head_dim**0.5).softmax(-1), + values.float(), + ).to(dtype) + torch.testing.assert_close( + score_fixed[batch_idx, :, :length], + logits, + atol=3e-2, + rtol=3e-2, + ) + + torch.testing.assert_close(out_fixed, expected, atol=3e-2, rtol=3e-2) + self.assertEqual(out_fixed.data_ptr(), output_ptr) + self.assertEqual( + (mid_fixed.data_ptr(), lse_fixed.data_ptr()), + workspace_ptrs, + ) + def test_full_layer_kivi_token_map_flash_decode_stage1_matches_dense_stage1(self): torch.manual_seed(8) device = "cuda" diff --git a/tests/test_efficiency_benchmark_contracts.py b/tests/test_efficiency_benchmark_contracts.py index 47f83004..54b9d3cb 100644 --- a/tests/test_efficiency_benchmark_contracts.py +++ b/tests/test_efficiency_benchmark_contracts.py @@ -13,6 +13,7 @@ _attach_churn_comparisons, _attach_saturation_metrics, _physical_gpu_metadata, + _decode_graph_counter_delta, _record_batch_first_tokens, _resolve_sparse_probe_protocol, _vllm_phase_metrics, @@ -77,6 +78,23 @@ def test_physical_gpu_metadata_uses_nvidia_smi_without_cuda_init(monkeypatch): ] +def test_decode_graph_counter_delta_reports_runtime_capture_churn(): + delta = _decode_graph_counter_delta( + {"capture_count": 28, "replay_count": 10, "eviction_count": 0}, + { + "capture_count": 30, + "replay_count": 110, + "eviction_count": 2, + "recapture_count": 1, + }, + ) + + assert delta["capture_count"] == 2 + assert delta["replay_count"] == 100 + assert delta["eviction_count"] == 2 + assert delta["recapture_count"] == 1 + + def test_unknown_hardware_does_not_fall_back_to_h100(): with pytest.raises(ValueError, match="Unknown GPU hardware"): detect_gpu_hardware("Mystery Accelerator") @@ -116,7 +134,7 @@ def test_model_specs_accept_nested_explicit_non_factorized_head_dim(): def test_model_specs_require_factorized_head_dim_when_not_explicit(): - with pytest.raises(ValueError, match="without an explicit head_dim"): + with pytest.raises(ValueError, match="must define head_dim"): ModelArchitectureSpecs.from_config_dict( { "hidden_size": 5120, @@ -128,6 +146,26 @@ def test_model_specs_require_factorized_head_dim_when_not_explicit(): ) +def test_model_specs_resolve_mla_qk_head_dim(): + specs = ModelArchitectureSpecs.from_config_dict( + { + "hidden_size": 2048, + "num_hidden_layers": 47, + "num_attention_heads": 20, + "num_key_value_heads": 20, + "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, + "vocab_size": 154880, + "n_routed_experts": 64, + "num_experts_per_tok": 4, + "moe_intermediate_size": 1536, + "intermediate_size": 10240, + } + ) + + assert specs.head_dim == 256 + + def test_probe_writes_metric_failed_when_model_discovery_fails(tmp_path, monkeypatch): model_dir = tmp_path / "model" model_dir.mkdir() diff --git a/tests/test_flashinfer_decode.py b/tests/test_flashinfer_decode.py index 715eee31..d362e5ec 100644 --- a/tests/test_flashinfer_decode.py +++ b/tests/test_flashinfer_decode.py @@ -17,6 +17,10 @@ def __init__( self, float_workspace_buffer, kv_layout="NHD", + use_cuda_graph=False, + paged_kv_indptr_buffer=None, + paged_kv_indices_buffer=None, + paged_kv_last_page_len_buffer=None, backend="auto", ): pass diff --git a/tests/test_gemma4_attention_kernels.py b/tests/test_gemma4_attention_kernels.py index 70160dce..50beee48 100644 --- a/tests/test_gemma4_attention_kernels.py +++ b/tests/test_gemma4_attention_kernels.py @@ -6,23 +6,13 @@ import torch from sparsevllm.engine.cache_manager.base import ExplicitKVPayload -from sparsevllm.kernels.triton.gemma4_context_attention import gemma4_context_attention -from sparsevllm.kernels.triton.gemma4_decode_attention import ( - gemma4_decode_stage1, - gemma4_decode_stage2, -) -from sparsevllm.kernels.triton.gemma4_global_decode_attention import ( - gemma4_global_decode_stage1, -) -from sparsevllm.kernels.triton.gemma4_single_block_decode_attention import ( - gemma4_single_block_decode, -) -from sparsevllm.kernels.triton.gemma4_window_decode_attention import ( - gemma4_window_decode, +from sparsevllm.kernels.triton.gemma4_context_attention import ( + gemma4_context_attention, ) from sparsevllm.operators.gemma4_attention import Gemma4FlashInferPrefill from sparsevllm.utils.context import reset_context, set_context + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -65,126 +55,31 @@ def test_gemma4_flashinfer_prefill_matches_torch( max_context_len=length, sliding_window=sliding_window, ) - second_slots = slots.flip(0).contiguous() - view.meta.active_slots = second_slots.view(1, -1) - second_output = prefill.run( - query, - view, - q_start=torch.zeros(1, device="cuda", dtype=torch.int32), - chunk_lens=torch.tensor([chunk], device="cuda", dtype=torch.int32), - max_context_len=length, - sliding_window=sliding_window, - ) - set_context( - True, - cu_seqlens_q=torch.tensor([0, chunk], device="cuda", dtype=torch.int32), - ) - second_slots.copy_(slots) - reused_output = prefill.run( - query, - view, - q_start=torch.zeros(1, device="cuda", dtype=torch.int32), - chunk_lens=torch.tensor([chunk], device="cuda", dtype=torch.int32), - max_context_len=length, - sliding_window=sliding_window, - ) finally: prefill.close() reset_context() + kv_head_ids = torch.arange(q_heads, device="cuda") // (q_heads // kv_heads) + logical_key, logical_value = key[slots.long()], value[slots.long()] + logits = torch.einsum( + "qhd,khd->hqk", query, logical_key[:, kv_head_ids] + ).float() query_positions = prefix + torch.arange(chunk, device="cuda") key_positions = torch.arange(length, device="cuda") visible = key_positions[None] <= query_positions[:, None] if sliding_window is not None: visible &= key_positions[None] > query_positions[:, None] - sliding_window - for actual, slot_ids in ( - (output, slots), - (second_output, slots.flip(0)), - (reused_output, slots), - ): - logical_key, logical_value = key[slot_ids.long()], value[slot_ids.long()] - logits = torch.einsum( - "qhd,khd->hqk", query, logical_key[:, kv_head_ids] - ).float() - probabilities = logits.masked_fill(~visible[None], -torch.inf).softmax(-1) - reference = torch.einsum( - "hqk,khd->qhd", - probabilities.to(value.dtype), - logical_value[:, kv_head_ids], - ) - cosine = torch.nn.functional.cosine_similarity( - actual.float().flatten(), reference.float().flatten(), dim=0 - ) - assert torch.isfinite(actual).all() - assert cosine > 0.999 - - -def _slots_and_lengths(): - lengths = torch.tensor([21, 13], dtype=torch.int32, device="cuda") - slots = torch.zeros((2, 21), dtype=torch.int32, device="cuda") - slots[0, :21] = torch.arange(21, dtype=torch.int32, device="cuda") - slots[1, :13] = torch.arange(21, 34, dtype=torch.int32, device="cuda") - return slots, lengths - - -def _decode_reference(q, k, v, slots, lengths, window): - output = torch.empty_like(q) - for batch, length in enumerate(lengths.tolist()): - start = max(0, length - (window or length)) - indices = slots[batch, start:length].long() - for head in range(q.shape[1]): - logits = q[batch, head] @ k[indices, head // (q.shape[1] // k.shape[1])].T - output[batch, head] = logits.softmax(-1) @ v[ - indices, head // (q.shape[1] // v.shape[1]) - ] - return output - - -@pytest.mark.parametrize("group_size", [8, 16]) -@pytest.mark.parametrize("length", [513, 8193]) -def test_gemma4_global_decode_matches_torch_and_graph(group_size, length): - torch.manual_seed(20260813) - block_seq = 256 - slots = torch.arange(length, device="cuda", dtype=torch.int32).view(1, -1) - lengths = torch.tensor([length], device="cuda", dtype=torch.int32) - key = torch.randn(length, 1, 512, device="cuda", dtype=torch.bfloat16) - value = torch.randn_like(key) - query = torch.randn(1, group_size, 512, device="cuda", dtype=torch.bfloat16) - blocks = (length + block_seq - 1) // block_seq - mid = torch.empty(1, group_size, blocks, 512, device="cuda", dtype=torch.float32) - lse = torch.empty(1, group_size, blocks, device="cuda", dtype=torch.float32) - output = torch.empty_like(query) - - def run(): - gemma4_global_decode_stage1( - query, - key, - value, - slots, - torch.zeros(1, device="cuda", dtype=torch.int32), - lengths, - mid, - lse, - block_seq=block_seq, - ) - gemma4_decode_stage2( - mid, lse, lengths, output, block_seq=block_seq, sliding_window=None - ) - - run() - reference = _decode_reference(query, key, value, slots, lengths, None) + probabilities = logits.masked_fill(~visible[None], -torch.inf).softmax(-1) + reference = torch.einsum( + "hqk,khd->qhd", + probabilities.to(value.dtype), + logical_value[:, kv_head_ids], + ) cosine = torch.nn.functional.cosine_similarity( output.float().flatten(), reference.float().flatten(), dim=0 ) + assert torch.isfinite(output).all() assert cosine > 0.999 - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() - query.copy_(torch.randn_like(query)) - graph.replay() - replay = output.clone() - graph.replay() - torch.testing.assert_close(output, replay, rtol=0, atol=0) @pytest.mark.parametrize("head_dim", [256, 512]) @@ -226,7 +121,9 @@ def test_gemma4_prefill_matches_torch(head_dim, sliding_window): indices = slots[batch, begin:end].long() for head in range(query.shape[1]): logits = query[start + offset, head] @ key[indices, head // 2].T - reference[start + offset, head] = logits.softmax(-1) @ value[indices, head // 2] + reference[start + offset, head] = logits.softmax(-1) @ value[ + indices, head // 2 + ] cosine = torch.nn.functional.cosine_similarity( output.float().flatten(), reference.float().flatten(), dim=0 ) @@ -264,12 +161,12 @@ def test_gemma4_long_window_prefill_matches_torch(): ).float() query_positions = 256 + torch.arange(1088, device="cuda") key_positions = torch.arange(1344, device="cuda") - visible = (key_positions[None, :] <= query_positions[:, None]) & ( - key_positions[None, :] > query_positions[:, None] - 1024 + visible = (key_positions[None] <= query_positions[:, None]) & ( + key_positions[None] > query_positions[:, None] - 1024 ) - probabilities = logits.masked_fill(~visible, -float("inf")).softmax(-1) reference = torch.bmm( - probabilities.to(value.dtype), value[:, kv_heads].permute(1, 0, 2) + logits.masked_fill(~visible, -torch.inf).softmax(-1).to(value.dtype), + value[:, kv_heads].permute(1, 0, 2), ).permute(1, 0, 2) cosine = torch.nn.functional.cosine_similarity( output.float().flatten(), reference.float().flatten(), dim=0 @@ -278,303 +175,18 @@ def test_gemma4_long_window_prefill_matches_torch(): assert cosine > 0.999 -@pytest.mark.parametrize("head_dim", [256, 512]) -@pytest.mark.parametrize("sliding_window", [None, 4]) -def test_gemma4_decode_matches_torch(head_dim, sliding_window): - torch.manual_seed(7) - slots, lengths = _slots_and_lengths() - key = torch.randn(34, 2, head_dim, dtype=torch.bfloat16, device="cuda") - value = torch.randn_like(key) - query = torch.randn(2, 4, head_dim, dtype=torch.bfloat16, device="cuda") - block_seq = 8 - blocks = (int(lengths.max()) + block_seq - 1) // block_seq - mid = torch.empty(2, 4, blocks, head_dim, dtype=torch.float32, device="cuda") - lse = torch.empty(2, 4, blocks, dtype=torch.float32, device="cuda") - output = torch.empty_like(query) - - gemma4_decode_stage1( - query, - key, - value, - slots, - torch.tensor([0, 1], dtype=torch.int32, device="cuda"), - lengths, - mid, - lse, - block_seq=block_seq, - sliding_window=sliding_window, - ) - gemma4_decode_stage2( - mid, - lse, - lengths, - output, - block_seq=block_seq, - sliding_window=sliding_window, - ) - - reference = _decode_reference(query, key, value, slots, lengths, sliding_window) - cosine = torch.nn.functional.cosine_similarity( - output.float().flatten(), reference.float().flatten(), dim=0 - ) - assert torch.isfinite(output).all() - assert cosine > 0.999 - - -def test_gemma4_long_window_decode_matches_torch(): - torch.manual_seed(13) - length, window, block_seq = 1486, 1024, 256 - slots = torch.arange(length, dtype=torch.int32, device="cuda").view(1, -1) - lengths = torch.tensor([length], dtype=torch.int32, device="cuda") - key = torch.randn(length, 2, 256, dtype=torch.bfloat16, device="cuda") - value = torch.randn_like(key) - query = torch.randn(1, 4, 256, dtype=torch.bfloat16, device="cuda") - blocks = (length + block_seq - 1) // block_seq - mid = torch.empty(1, 4, blocks, 256, dtype=torch.float32, device="cuda") - lse = torch.empty(1, 4, blocks, dtype=torch.float32, device="cuda") - output = torch.empty_like(query) - - gemma4_decode_stage1( - query, - key, - value, - slots, - torch.zeros(1, dtype=torch.int32, device="cuda"), - lengths, - mid, - lse, - block_seq=block_seq, - sliding_window=window, - ) - gemma4_decode_stage2( - mid, - lse, - lengths, - output, - block_seq=block_seq, - sliding_window=window, - ) - - reference = _decode_reference(query, key, value, slots, lengths, window) - cosine = torch.nn.functional.cosine_similarity( - output.float().flatten(), reference.float().flatten(), dim=0 - ) - assert torch.isfinite(output).all() - assert cosine > 0.999 - - -@pytest.mark.parametrize("group_size", [2, 4, 8]) -@pytest.mark.parametrize("head_dim", [256, 512]) -def test_gemma4_single_block_decode_matches_torch(group_size, head_dim): - torch.manual_seed(11) - slots, lengths = _slots_and_lengths() - key = torch.randn(34, 2, head_dim, dtype=torch.bfloat16, device="cuda") - value = torch.randn_like(key) - query = torch.randn(2, 2 * group_size, head_dim, dtype=torch.bfloat16, device="cuda") - output = torch.empty_like(query) - gemma4_single_block_decode( - query, - key, - value, - slots, - torch.tensor([0, 1], dtype=torch.int32, device="cuda"), - lengths, - output, - block_seq=256, - sliding_window=None, - ) - reference = _decode_reference(query, key, value, slots, lengths, None) - cosine = torch.nn.functional.cosine_similarity( - output.float().flatten(), reference.float().flatten(), dim=0 - ) - assert torch.isfinite(output).all() - assert cosine > 0.999 - - -@pytest.mark.parametrize("group_size", [2, 4]) -@pytest.mark.parametrize("block_seq", [250, 256]) -def test_gemma4_window_decode_matches_torch(group_size, block_seq): - torch.manual_seed(17) - lengths = torch.tensor([1301, 1177], dtype=torch.int32, device="cuda") - slots = torch.zeros((2, 1301), dtype=torch.int32, device="cuda") - slots[0, :1301] = torch.arange(1301, dtype=torch.int32, device="cuda") - slots[1, :1177] = torch.arange(1301, 2478, dtype=torch.int32, device="cuda") - key = torch.randn(2478, 2, 256, dtype=torch.bfloat16, device="cuda") - value = torch.randn_like(key) - query = torch.randn(2, 2 * group_size, 256, dtype=torch.bfloat16, device="cuda") - blocks = (1024 + block_seq - 1) // block_seq - mid = torch.empty( - 2, 2 * group_size, blocks, 256, dtype=torch.float32, device="cuda" - ) - lse = torch.empty( - 2, 2 * group_size, blocks, dtype=torch.float32, device="cuda" - ) - output = torch.empty_like(query) - gemma4_window_decode( - query, - key, - value, - slots, - torch.tensor([0, 1], dtype=torch.int32, device="cuda"), - lengths, - mid, - lse, - output, - block_seq=block_seq, - sliding_window=1024, - ) - reference = _decode_reference(query, key, value, slots, lengths, 1024) - cosine = torch.nn.functional.cosine_similarity( - output.float().flatten(), reference.float().flatten(), dim=0 - ) - assert torch.isfinite(output).all() - assert cosine > 0.999 - - -def test_gemma4_window_decode_supports_cuda_graph(): - torch.manual_seed(23) - slots, lengths = _slots_and_lengths() - key = torch.randn(34, 2, 256, dtype=torch.bfloat16, device="cuda") - value = torch.randn_like(key) - query = torch.randn(2, 4, 256, dtype=torch.bfloat16, device="cuda") - mid = torch.empty(2, 4, 2, 256, dtype=torch.float32, device="cuda") - lse = torch.empty(2, 4, 2, dtype=torch.float32, device="cuda") - output = torch.empty_like(query) - request_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - - def run(): - gemma4_window_decode( - query, - key, - value, - slots, - request_indices, - lengths, - mid, - lse, - output, - block_seq=8, - sliding_window=16, - ) - - for _ in range(3): - run() - reference = _decode_reference(query, key, value, slots, lengths, 16) - cosine = torch.nn.functional.cosine_similarity( - output.float().flatten(), reference.float().flatten(), dim=0 - ) - assert cosine > 0.999 - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() - query.copy_(torch.randn_like(query)) - graph.replay() - first = output.clone() - graph.replay() - assert torch.equal(first, output) - - -@pytest.mark.parametrize("head_dim", [256, 512]) -def test_gemma4_decode_supports_cuda_graph(head_dim): - slots, lengths = _slots_and_lengths() - query = torch.randn(2, 4, head_dim, dtype=torch.bfloat16, device="cuda") - key = torch.randn(34, 2, head_dim, dtype=torch.bfloat16, device="cuda") - value = torch.randn_like(key) - mid = torch.empty(2, 4, 1, head_dim, dtype=torch.float32, device="cuda") - lse = torch.empty(2, 4, 1, dtype=torch.float32, device="cuda") - output = torch.empty_like(query) - request_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - - def run(): - gemma4_decode_stage1( - query, - key, - value, - slots, - request_indices, - lengths, - mid, - lse, - block_seq=256, - sliding_window=1024, - ) - gemma4_decode_stage2( - mid, - lse, - lengths, - output, - block_seq=256, - sliding_window=1024, - ) - - for _ in range(3): - run() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() - query.copy_(torch.randn_like(query)) - graph.replay() - first = output.clone() - graph.replay() - assert torch.equal(first, output) - - -@pytest.mark.parametrize("score_dims", [2, 3]) -def test_gemma4_decode_collects_raw_qk_scores(score_dims): - torch.manual_seed(13) - slots, lengths = _slots_and_lengths() - head_dim = 256 - query = torch.randn(2, 4, head_dim, dtype=torch.bfloat16, device="cuda") - key = torch.randn(34, 2, head_dim, dtype=torch.bfloat16, device="cuda") - value = torch.randn_like(key) - mid = torch.empty(2, 4, 3, head_dim, dtype=torch.float32, device="cuda") - lse = torch.empty(2, 4, 3, dtype=torch.float32, device="cuda") - score = torch.full( - (2, 4, 21) if score_dims == 3 else (2, 21), - -1e20, - dtype=torch.float32, - device="cuda", - ) - gemma4_decode_stage1( - query, - key, - value, - slots, - torch.tensor([0, 1], dtype=torch.int32, device="cuda"), - lengths, - mid, - lse, - block_seq=8, - sliding_window=None, - attn_score=score, - ) - expected = torch.empty(2, 4, 21, dtype=torch.float32, device="cuda") - expected.fill_(-1e20) - for batch, length in enumerate(lengths.tolist()): - for head in range(4): - indices = slots[batch, :length].long() - expected[batch, head, :length] = ( - query[batch, head].float() - @ key[indices, head // 2].float().T - ) - expected = expected if score_dims == 3 else expected.max(1).values - torch.testing.assert_close(score, expected, rtol=2e-2, atol=1.0) - - -@pytest.mark.parametrize("score_dims", [2, 3]) -def test_gemma4_prefill_collects_raw_qk_scores(score_dims): +@pytest.mark.parametrize("score_rank", [2, 3]) +def test_gemma4_prefill_collects_raw_qk_scores(score_rank): torch.manual_seed(17) head_dim = 256 - prefix = torch.tensor([0], dtype=torch.int32, device="cuda") lengths = torch.tensor([4], dtype=torch.int32, device="cuda") - starts = torch.tensor([0], dtype=torch.int32, device="cuda") slots = torch.arange(4, dtype=torch.int32, device="cuda").unsqueeze(0) key = torch.randn(4, 2, head_dim, dtype=torch.bfloat16, device="cuda") value = torch.randn_like(key) query = torch.randn(4, 4, head_dim, dtype=torch.bfloat16, device="cuda") output = torch.empty_like(query) score = torch.zeros( - (1, 4, 4) if score_dims == 3 else (1, 4), + (1, 4, 4) if score_rank == 3 else (1, 4), dtype=torch.float32, device="cuda", ) @@ -584,9 +196,9 @@ def test_gemma4_prefill_collects_raw_qk_scores(score_dims): value, output, torch.tensor([0], dtype=torch.int32, device="cuda"), - starts, + torch.tensor([0], dtype=torch.int32, device="cuda"), lengths, - prefix, + torch.zeros(1, dtype=torch.int32, device="cuda"), 4, slots, sliding_window=None, @@ -594,11 +206,9 @@ def test_gemma4_prefill_collects_raw_qk_scores(score_dims): ) expected = torch.zeros(1, 4, 4, dtype=torch.float32, device="cuda") for head in range(4): - logits = query[:, head].float() @ key[:, head // 2].float().T - expected[0, head] = logits.tril().sum(0) - expected = ( - expected - if score_dims == 3 - else (expected / 4).max(1).values.clamp_min_(0) - ) + expected[0, head] = ( + query[:, head].float() @ key[:, head // 2].float().T + ).tril().sum(0) + if score_rank == 2: + expected = (expected / 4).max(1).values.clamp_min_(0) torch.testing.assert_close(score, expected, rtol=2e-2, atol=1.0) diff --git a/tests/test_gemma4_fixed_grid_decode.py b/tests/test_gemma4_fixed_grid_decode.py new file mode 100644 index 00000000..41aafdfc --- /dev/null +++ b/tests/test_gemma4_fixed_grid_decode.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +import pytest +import torch + +from sparsevllm.kernels.triton.sglang_gemma4_decode_attention import ( + sglang_gemma4_decode, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _reference(q, k, v, slots, request_indices, lengths, window): + output = torch.empty_like(q) + group_size = q.shape[1] // k.shape[1] + for batch_idx, length in enumerate(lengths.tolist()): + start = max(0, length - int(window or length)) + active = slots[request_indices[batch_idx], start:length].long() + keys = k[active].repeat_interleave(group_size, dim=1) + values = v[active].repeat_interleave(group_size, dim=1) + logits = torch.einsum("hd,lhd->hl", q[batch_idx].float(), keys.float()) + output[batch_idx] = torch.einsum( + "hl,lhd->hd", + torch.softmax(logits, dim=-1), + values.float(), + ).to(q.dtype) + return output + + +def _case(*, dtype, head_dim, query_heads, kv_heads, capacity, lengths, window): + batch = len(lengths) + torch.manual_seed(20260825 + head_dim + query_heads + capacity) + q = torch.randn( + batch, + query_heads, + head_dim, + dtype=dtype, + device="cuda", + ).mul_(0.25) + slots = torch.arange( + batch * capacity, + dtype=torch.int32, + device="cuda", + ).view(batch, capacity) + k = torch.randn( + batch * capacity, + kv_heads, + head_dim, + dtype=dtype, + device="cuda", + ).mul_(0.25) + v = torch.randn_like(k) + request_indices = torch.arange(batch, dtype=torch.int32, device="cuda") + context_lens = torch.tensor(lengths, dtype=torch.int32, device="cuda") + mid_output = torch.empty( + batch, + query_heads, + 8, + head_dim, + dtype=torch.float32, + device="cuda", + ) + mid_lse = torch.empty( + batch, + query_heads, + 8, + dtype=torch.float32, + device="cuda", + ) + num_kv_splits = torch.empty(batch, dtype=torch.int32, device="cuda") + return ( + q, + k, + v, + slots, + request_indices, + context_lens, + mid_output, + mid_lse, + num_kv_splits, + window, + ) + + +def _run(case, *, score=None): + q, k, v, slots, request_indices, lengths, mid, lse, splits, window = case + return sglang_gemma4_decode( + q, + k, + v, + slots, + request_indices, + lengths, + mid, + lse, + splits, + sliding_window=window, + multi_processor_count=torch.cuda.get_device_properties(0).multi_processor_count, + attn_score=score, + ) + + +@pytest.mark.parametrize( + "case_kwargs", + [ + dict( + dtype=torch.bfloat16, + head_dim=256, + query_heads=2, + kv_heads=2, + capacity=21, + lengths=[21, 13], + window=None, + ), + dict( + dtype=torch.float16, + head_dim=256, + query_heads=8, + kv_heads=2, + capacity=1301, + lengths=[1301, 1177], + window=1024, + ), + dict( + dtype=torch.bfloat16, + head_dim=512, + query_heads=8, + kv_heads=1, + capacity=513, + lengths=[513, 377], + window=None, + ), + dict( + dtype=torch.bfloat16, + head_dim=512, + query_heads=16, + kv_heads=1, + capacity=8193, + lengths=[8193], + window=None, + ), + ], +) +def test_gemma4_fixed_grid_decode_matches_independent_oracle(case_kwargs): + case = _case(**case_kwargs) + actual = _run(case) + q, k, v, slots, request_indices, lengths, *_, window = case + expected = _reference(q, k, v, slots, request_indices, lengths, window) + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=3e-2) + + +@pytest.mark.parametrize("score_rank", [2, 3]) +def test_gemma4_fixed_grid_decode_produces_raw_qk_scores(score_rank): + case = _case( + dtype=torch.bfloat16, + head_dim=256, + query_heads=4, + kv_heads=2, + capacity=33, + lengths=[33, 21], + window=None, + ) + q, k, _, slots, request_indices, lengths, *_ = case + score = torch.full( + (2, 4, 33) if score_rank == 3 else (2, 33), + -1e20, + dtype=torch.float32, + device="cuda", + ) + _run(case, score=score) + expected = torch.full( + (2, 4, 33), + -1e20, + dtype=torch.float32, + device="cuda", + ) + group_size = q.shape[1] // k.shape[1] + for batch_idx, length in enumerate(lengths.tolist()): + active = slots[request_indices[batch_idx], :length].long() + keys = k[active].repeat_interleave(group_size, dim=1) + expected[batch_idx, :, :length] = torch.einsum( + "hd,lhd->hl", + q[batch_idx].float(), + keys.float(), + ) + if score_rank == 2: + expected = expected.max(dim=1).values + torch.testing.assert_close(score, expected, rtol=2e-2, atol=1.0) + + +@pytest.mark.parametrize( + ("head_dim", "window"), + [(256, None), (256, 16), (512, None)], +) +def test_gemma4_fixed_grid_decode_replays_new_lengths_and_rows(head_dim, window): + case = _case( + dtype=torch.bfloat16, + head_dim=head_dim, + query_heads=4, + kv_heads=2, + capacity=33, + lengths=[33, 21], + window=window, + ) + _run(case) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = _run(case) + + q, k, v, slots, request_indices, lengths, *_, window = case + q.copy_(torch.randn_like(q)) + request_indices.copy_(torch.tensor([1, 0], dtype=torch.int32, device="cuda")) + lengths.copy_(torch.tensor([17, 29], dtype=torch.int32, device="cuda")) + expected = _reference(q, k, v, slots, request_indices, lengths, window) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(graph_output, expected, rtol=3e-2, atol=3e-2) diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index a706682f..2b093725 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -29,8 +29,10 @@ Glm4MoeLiteForCausalLM, Glm4MoeLiteRouter, Glm4MoeLiteSparseMoeBlock, + build_glm4_moe_lite_mla_attention, ) from sparsevllm.models.qwen3 import Qwen3MLP +from sparsevllm.operators.attention_capabilities import AttentionScoreKind from sparsevllm.operators.mla_attention import MlaAttentionOpSpec from sparsevllm.operators.moe import TritonMoeProvider from sparsevllm.operators.moe_router import GlmBiasedSigmoidRouterProvider @@ -194,6 +196,8 @@ def test_glm_runtime_kwargs_bind_model_owned_operators( context = _tp_context(tp_size=2) runtime = SimpleNamespace( decode_graph=True, + decode_graph_shape_policy="batch_only", + max_model_len=32768, max_num_seqs_in_batch=4, max_decoding_seqs=8, mla_prefill_workspace_bytes=1024, @@ -233,8 +237,13 @@ def test_glm_runtime_kwargs_bind_model_owned_operators( max_batch_size=8, prefill_workspace_bytes=1024, decode_graph=True, + context_capacity=32768, projection_chunk_size=16, - may_require_attention_scores=requires_scores, + score_output=( + AttentionScoreKind.RAW_QK_PER_HEAD + if requires_scores + else AttentionScoreKind.NONE + ), ) build_all_reduce.assert_called_once_with( config, @@ -245,6 +254,36 @@ def test_glm_runtime_kwargs_bind_model_owned_operators( ) +def test_glm_batch_only_mla_spec_owns_context_capacity() -> None: + config = _config() + config.decode_graph_shape_policy = "batch_only" + bound = object() + with ( + patch( + "sparsevllm.models.glm4_moe_lite.get_parallel_context", + return_value=_tp_context(tp_size=2), + ), + patch( + "sparsevllm.models.glm4_moe_lite.MLAAttention.bind", + return_value=bound, + ) as bind, + ): + actual = build_glm4_moe_lite_mla_attention( + config, + device="cpu", + max_batch_size=8, + prefill_workspace_bytes=1024, + decode_graph=True, + context_capacity=32768, + projection_chunk_size=16, + ) + + assert actual is bound + spec = bind.call_args.kwargs["spec"] + assert spec.batch_only_cuda_graph + assert spec.context_capacity == 32768 + + def test_glm_interleaved_rope_matches_transformers() -> None: torch.manual_seed(13) q = torch.randn(1, 3, 5, 64) diff --git a/tests/test_glm_cuda_graph.py b/tests/test_glm_cuda_graph.py index 2546a270..399f3307 100644 --- a/tests/test_glm_cuda_graph.py +++ b/tests/test_glm_cuda_graph.py @@ -12,7 +12,13 @@ from torch import nn from sparsevllm.config import RuntimeLayout +from sparsevllm.configs.cuda_graph import ( + _default_decode_cuda_graph_capture_sizes, + build_decode_cuda_graph_startup_family_plan, + build_decode_cuda_graph_startup_plan, +) from sparsevllm.models.layout import resolve_attention_qk_head_dim +from sparsevllm.method_registry import sparse_decode_attention_score_kind from sparsevllm.distributed import ParallelContext from sparsevllm.engine.cache_manager import LayerBatchStates from sparsevllm.engine.cache_manager.h2o import H2OCacheManager @@ -35,6 +41,7 @@ Glm4MoeLiteForCausalLM, Glm4MoeLiteSparseMoeBlock, ) +from sparsevllm.operators.attention_capabilities import AttentionScoreKind from sparsevllm.operators.mla_attention import MlaAttentionOpSpec from sparsevllm.utils.context import get_context @@ -45,6 +52,154 @@ ) +@pytest.mark.parametrize( + ("method", "expected"), + [ + ("pyramidkv", AttentionScoreKind.RAW_QK_REDUCED), + ("omnikv", AttentionScoreKind.RAW_QK_PER_HEAD), + ("skipkv", AttentionScoreKind.RAW_QK_PER_HEAD), + ("deltakv", AttentionScoreKind.RAW_QK_PER_HEAD), + ("vanilla", AttentionScoreKind.NONE), + ], +) +def test_glm_sparse_method_declares_exact_decode_score_contract( + method, + expected, +): + assert sparse_decode_attention_score_kind(method) is expected + + +def test_startup_graph_plan_captures_complete_coarse_grid_when_it_fits(): + plan = build_decode_cuda_graph_startup_plan( + [1, 2, 4, 8], + [1024, 2048, 4096, 8192, 16384, 32768, 33280], + 32, + ) + + assert len(plan) == 28 + assert plan[0] == (1, 1024) + assert plan[-1] == (8, 33280) + + +def test_startup_graph_plan_spreads_contexts_and_preserves_mandatory_graph(): + plan = build_decode_cuda_graph_startup_plan( + [1, 2, 4, 8], + [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144], + 12, + mandatory=(8, 8192), + ) + + assert len(plan) == 12 + assert {batch for batch, _ in plan} == {1, 2, 4, 8} + assert (8, 8192) in plan + assert all(any(context == 262144 for b, context in plan if b == batch) for batch in (1, 2, 4)) + + +def test_startup_graph_plan_prioritizes_dense_batch_coverage(): + batches = list(range(1, 9)) + contexts = [1024, 2048, 4096, 8192, 16384, 32768, 65536] + + plan = build_decode_cuda_graph_startup_plan(batches, contexts, 32) + + assert len(plan) == 32 + assert {batch for batch, _ in plan} == set(batches) + assert all((batch, 65536) in plan for batch in batches) + assert all(len([pair for pair in plan if pair[0] == batch]) == 4 for batch in batches) + + +def test_startup_graph_plan_keeps_max_context_when_mandatory_cannot_fit(): + plan = build_decode_cuda_graph_startup_plan( + [1, 2, 3], + [1024, 2048, 4096], + 3, + mandatory=(3, 1024), + ) + + assert plan == [(1, 4096), (2, 4096), (3, 4096)] + + +def test_sparse_startup_graph_plan_covers_short_and_long_families(): + config = SimpleNamespace( + decode_graph_capture_sizes=list(range(1, 9)), + decode_graph_context_sizes=[1024, 2048, 4096, 8192, 16384, 32768], + decode_graph_startup_capture_limit=48, + decode_graph_max_cached_graphs=48, + sparse_method="snapkv", + sink_keep_tokens=64, + decode_keep_tokens=4096, + recent_keep_tokens=512, + max_model_len=32768, + ) + + plan = build_decode_cuda_graph_startup_family_plan(config) + + assert len(plan) == 48 + assert {(batch, is_long) for batch, _, is_long in plan} == { + (batch, is_long) + for batch in range(1, 9) + for is_long in (False, True) + } + assert all(context > 4672 for _, context, is_long in plan if is_long) + assert all( + len([key for key in plan if key[0] == batch and key[2] == is_long]) == 3 + for batch in range(1, 9) + for is_long in (False, True) + ) + + +def test_h2o_startup_graph_plan_uses_normal_context_buckets(): + config = SimpleNamespace( + decode_graph_capture_sizes=[1, 2, 4], + decode_graph_context_sizes=[1024, 2048, 4096, 8192, 16384], + decode_graph_startup_capture_limit=48, + decode_graph_max_cached_graphs=48, + sparse_method="h2o", + sink_keep_tokens=64, + decode_keep_tokens=4096, + recent_keep_tokens=512, + max_model_len=16384, + ) + + plan = build_decode_cuda_graph_startup_family_plan(config) + + assert plan == sorted( + [ + (batch, context, False) + for batch in (1, 2, 4) + for context in (1024, 2048, 4096, 8192, 16384) + ] + + [ + (batch, context, True) + for batch in (1, 2, 4) + for context in (8192, 16384) + ], + reverse=True, + ) + + +def test_sparse_startup_graph_plan_covers_default_64_sequence_limit(): + batches = _default_decode_cuda_graph_capture_sizes(64) + config = SimpleNamespace( + decode_graph_capture_sizes=batches, + decode_graph_context_sizes=[1024, 2048, 4096, 8192, 16384, 32768, 65536], + decode_graph_startup_capture_limit=48, + decode_graph_max_cached_graphs=48, + sparse_method="snapkv", + sink_keep_tokens=64, + decode_keep_tokens=4096, + recent_keep_tokens=512, + max_model_len=65536, + ) + + plan = build_decode_cuda_graph_startup_family_plan(config) + + assert len(batches) == 22 + assert len(plan) == 48 + assert {(batch, is_long) for batch, _, is_long in plan} == { + (batch, is_long) for batch in batches for is_long in (False, True) + } + + def _make_glm_graph_lane( *, device: torch.device, @@ -978,6 +1133,7 @@ def _make_glm_method_graph_lane( cache_dtype=torch.bfloat16, tp_size=1, cuda_graph=True, + score_output=sparse_decode_attention_score_kind(method), ) mla_attention = MLAAttention.bind( spec=spec, diff --git a/tests/test_glm_runtime_compatibility.py b/tests/test_glm_runtime_compatibility.py index 1d8893c5..7e6049ac 100644 --- a/tests/test_glm_runtime_compatibility.py +++ b/tests/test_glm_runtime_compatibility.py @@ -107,3 +107,59 @@ def test_glm_config_rejects_nondivisible_outer_tp_moe_ep_layout(): def test_glm_config_rejects_data_parallelism(): with pytest.raises(ValueError, match="does not support data parallelism"): _glm_config(data_parallel_size=2) + + +def test_glm_config_defaults_to_bounded_vanilla_startup_graph_capture(): + config = _glm_config(decode_graph=True) + + assert config.decode_graph_startup_capture is True + assert config.decode_graph_startup_capture_limit == 32 + assert config.decode_graph_max_cached_graphs == 32 + + +def test_glm_config_allows_disabling_default_startup_graph_capture(): + config = _glm_config( + decode_graph=True, + decode_graph_startup_capture=False, + ) + + assert config.decode_graph_startup_capture is False + assert config.decode_graph_max_cached_graphs is None + + +def test_glm_config_defaults_to_larger_sparse_startup_capture_budget(): + config = _glm_config( + decode_graph=True, + sparse_method="snapkv", + ) + + assert config.decode_graph_startup_capture is True + assert config.decode_graph_startup_capture_limit == 48 + assert config.decode_graph_max_cached_graphs == 48 + + +def test_glm_config_rejects_startup_capture_without_cuda_graph(): + with pytest.raises(ValueError, match="requires decode_graph=True"): + _glm_config(decode_graph_startup_capture=True) + + +def test_glm_config_allows_disabling_sparse_startup_capture(): + config = _glm_config( + decode_graph=True, + decode_graph_startup_capture=False, + sparse_method="snapkv", + ) + + assert config.decode_graph_startup_capture is False + assert config.decode_graph_max_cached_graphs is None + + +def test_glm_config_rejects_startup_budget_smaller_than_batch_plan(): + with pytest.raises(ValueError, match="must cover every batch bucket"): + _glm_config( + decode_graph=True, + decode_graph_startup_capture=True, + decode_graph_capture_sizes=[1, 2, 3, 4, 5], + decode_graph_max_cached_graphs=4, + max_decoding_seqs=5, + ) diff --git a/tests/test_minimax_m2_attention_graph.py b/tests/test_minimax_m2_attention_graph.py index b41cb20e..67e8a238 100644 --- a/tests/test_minimax_m2_attention_graph.py +++ b/tests/test_minimax_m2_attention_graph.py @@ -1,9 +1,23 @@ +from types import SimpleNamespace +from unittest.mock import Mock, patch + import pytest import torch +from sparsevllm.engine.cache_manager import ( + AttentionViewMeta, + DecodeComputeView, + ExplicitKVPayload, +) +from sparsevllm.kernels.external.sgl.fa3 import sgl_fa3_support from sparsevllm.kernels.triton.flash_decoding_stage2 import flash_decode_stage2 from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import flash_decode_stage1 from sparsevllm.kernels.triton.store_kvcache import store_kvcache +from sparsevllm.operators.decode_attention import ( + DecodeAttentionOpSpec, + SglFa3PagedDecodeAttentionProvider, + prepare_decode_attention_op, +) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -130,3 +144,98 @@ def run_decode(): run_decode() torch.cuda.synchronize() assert torch.equal(graph_output, output) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_fa3_support()[0], + reason="CUDA and a validated sglang-kernel are required", +) +def test_minimax_m2_production_provider_replays_across_32k_boundary(): + torch.manual_seed(20260825) + device = torch.device("cuda") + query_heads, kv_heads, head_dim = 12, 2, 128 + capacity = 32769 + spec = DecodeAttentionOpSpec( + num_query_heads=query_heads, + num_kv_heads=kv_heads, + head_dim=head_dim, + activation_dtype=torch.bfloat16, + softmax_scale=head_dim**-0.5, + max_batch_size=1, + batch_only_cuda_graph=True, + context_capacity=capacity, + ) + prepared = prepare_decode_attention_op(spec, device_index=device.index or 0) + assert isinstance(prepared.provider, SglFa3PagedDecodeAttentionProvider) + + q = 0.25 * torch.randn( + 1, + query_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + k_cache = 0.25 * torch.randn( + capacity, + kv_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + v_cache = torch.randn_like(k_cache) + active_slots = torch.arange( + capacity, + dtype=torch.int32, + device=device, + ).unsqueeze(0) + context_lens = torch.tensor([32767], dtype=torch.int32, device=device) + view = DecodeComputeView( + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=torch.zeros(1, dtype=torch.int32, device=device), + context_lens=context_lens, + max_context_len=capacity, + ), + payload=ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), + ) + + launch_profile = Mock(name="context_dependent_launch_profile") + validation_scope = object() + with patch( + "sparsevllm.operators.decode_attention.get_context", + return_value=SimpleNamespace(attention_validation_scope=validation_scope), + ): + prepared.run(q, view, decode_launch_op=launch_profile) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = prepared.run(q, view, decode_launch_op=launch_profile) + + for context_len in (32767, 32768, 32769): + context_lens.fill_(context_len) + graph.replay() + torch.cuda.synchronize() + + active = active_slots[0, :context_len].long() + group_size = query_heads // kv_heads + expanded_k = k_cache[active].repeat_interleave(group_size, dim=1) + expanded_v = v_cache[active].repeat_interleave(group_size, dim=1) + logits = torch.einsum( + "hd,lhd->hl", + q[0].float(), + expanded_k.float(), + ) + probabilities = torch.softmax(logits * spec.softmax_scale, dim=-1) + expected = torch.einsum( + "hl,lhd->hd", + probabilities, + expanded_v.float(), + ).to(torch.bfloat16) + torch.testing.assert_close( + graph_output[0], + expected, + rtol=3e-2, + atol=3e-2, + ) + + launch_profile.launch_config.assert_not_called() + prepared.close() diff --git a/tests/test_minimax_m2_config.py b/tests/test_minimax_m2_config.py index 4b368e85..ae9c8771 100644 --- a/tests/test_minimax_m2_config.py +++ b/tests/test_minimax_m2_config.py @@ -1,3 +1,4 @@ +import json from types import SimpleNamespace from unittest.mock import patch @@ -93,6 +94,36 @@ def test_minimax_config_requires_all_fp8_exclusions(tmp_path): _make_config(tmp_path, hf_config=hf_config) +def test_minimax_config_supports_quantized_tiny_random(tmp_path): + tiny_config = tmp_path / "tiny.json" + tiny_config.write_text( + json.dumps( + { + "num_hidden_layers": 1, + "hidden_size": 3072, + "intermediate_size": 1536, + "num_attention_heads": 48, + "num_key_value_heads": 8, + "head_dim": 128, + "vocab_size": 256, + "max_position_embeddings": 512, + } + ), + encoding="utf-8", + ) + + config = _make_config( + tmp_path, + tiny_random=True, + tiny_random_config=str(tiny_config), + max_model_len=512, + ) + + assert config.tiny_random + assert config.quantization_config.enabled + assert config.hf_config.hidden_size == 3072 + + @pytest.mark.parametrize( "parallel_kwargs", [ diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index 38b96b10..4869a2f5 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -13,6 +13,10 @@ MlaLatentPayload, PrefillComputeView, ) +from sparsevllm.kernels.external.sgl.fa3 import sgl_fa3_support +from sparsevllm.kernels.triton.mla import ( + MlaDecodeWorkspace, +) from sparsevllm.operators.mla_attention import ( MLA_ATTENTION_REGISTRY, MlaAttentionOpSpec, @@ -21,9 +25,6 @@ ) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum -from sparsevllm.kernels.triton.mla import ( - MlaDecodeWorkspace, -) def _spec(**overrides) -> MlaAttentionOpSpec: @@ -37,6 +38,7 @@ def _spec(**overrides) -> MlaAttentionOpSpec: "cache_dtype": torch.bfloat16, "tp_size": 4, "cuda_graph": False, + "context_capacity": 65536, } values.update(overrides) return MlaAttentionOpSpec(**values) @@ -80,6 +82,7 @@ def _cpu_workspace(batch_size: int, head_count: int) -> MlaDecodeWorkspace: {"value_head_dim": 0}, {"tp_size": 0}, {"num_q_heads": 20, "tp_size": 3}, + {"context_capacity": 0}, ], ) def test_mla_attention_spec_rejects_invalid_dimensions(overrides) -> None: @@ -103,6 +106,108 @@ def test_mla_triton_atomic_support_is_not_narrowed_by_device_name() -> None: assert result.supported +def test_batch_only_mla_requires_static_capacity() -> None: + spec = _spec( + cuda_graph=True, + batch_only_cuda_graph=True, + context_capacity=None, + ) + + result = MlaTritonProvider.supports(spec, _h100_caps()) + + assert not result.supported + assert "static context capacity" in result.reason + + +def test_batch_only_mla_launch_config_ignores_runtime_context() -> None: + spec = _spec( + tp_size=2, + cuda_graph=True, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + workspace = _cpu_workspace(batch_size=32, head_count=10) + with patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ): + provider = MlaTritonProvider( + op_spec=spec, + device="cpu", + max_batch_size=32, + ) + launch_config = object() + with patch( + "sparsevllm.operators.mla_attention.select_glm_mla_decode_config", + return_value=launch_config, + ) as select: + first = provider._launch_config_for( + batch_size=32, + max_context_len=1, + active_slot_width=64, + ) + second = provider._launch_config_for( + batch_size=32, + max_context_len=32000, + active_slot_width=65536, + ) + + assert first is launch_config + assert second is launch_config + assert select.call_count == 2 + select.assert_called_with( + batch_size=32, + context_capacity=32768, + local_q_heads=10, + ) + + +def test_sgl_mla_accepts_batch_only_score_free_contract() -> None: + spec = _spec( + cuda_graph=True, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + with patch( + "sparsevllm.operators.mla_attention.sgl_fa3_device_support", + return_value=(True, "sgl test"), + ): + result = MlaSglFa3Provider.supports(spec, _h100_caps()) + + assert result.supported + assert MlaSglFa3Provider.supports_batch_only_cuda_graph + + +def test_batch_only_mla_resolver_prefers_sgl_fa3() -> None: + spec = _spec( + cuda_graph=True, + batch_only_cuda_graph=True, + context_capacity=32768, + ) + workspace = _cpu_workspace(batch_size=8, head_count=5) + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_device_support", + return_value=(True, "sgl test"), + ), + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ), + patch("sparsevllm.operators.mla_attention.SglFa3DecodeKernel"), + ): + resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( + spec, + _h100_caps(), + op_spec=spec, + device="cpu", + max_batch_size=8, + ) + + assert type(resolved.provider) is MlaSglFa3Provider + assert resolved.report.selection_basis == "upstream_default" + + @pytest.mark.parametrize( ("spec_overrides", "caps_overrides", "reason"), [ @@ -467,3 +572,126 @@ def test_mla_provider_runs_static_padded_batch() -> None: torch.testing.assert_close(output[1], torch.zeros_like(output[1])) assert bool(torch.isfinite(output).all().item()) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_fa3_support()[0], + reason="CUDA and a validated sglang-kernel are required", +) +@torch.inference_mode() +def test_glm_production_provider_replays_across_1k_boundary() -> None: + torch.manual_seed(20260825) + device = torch.device("cuda") + capacity = 1025 + spec = _spec( + tp_size=1, + cuda_graph=True, + batch_only_cuda_graph=True, + context_capacity=capacity, + ) + provider = MlaSglFa3Provider( + op_spec=spec, + device=device, + max_batch_size=1, + ) + + q_nope_absorbed = 0.125 * torch.randn( + 1, + spec.local_q_heads, + spec.kv_lora_rank, + dtype=torch.bfloat16, + device=device, + ) + q_rope = 0.125 * torch.randn( + 1, + spec.local_q_heads, + spec.rope_dim, + dtype=torch.bfloat16, + device=device, + ) + payload = MlaLatentPayload( + latent_cache=0.125 + * torch.randn( + capacity, + 1, + spec.kv_lora_rank, + dtype=torch.bfloat16, + device=device, + ), + rope_cache=0.125 + * torch.randn( + capacity, + 1, + spec.rope_dim, + dtype=torch.bfloat16, + device=device, + ), + ) + active_slots = torch.arange( + capacity, + dtype=torch.int32, + device=device, + ).unsqueeze(0) + context_lens = torch.tensor([1023], dtype=torch.int32, device=device) + view = DecodeComputeView( + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=torch.zeros(1, dtype=torch.int32, device=device), + context_lens=context_lens, + max_context_len=capacity, + ), + payload=payload, + ) + output = torch.empty_like(q_nope_absorbed) + validation_scope = object() + + provider.run( + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=validation_scope, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + provider.run( + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=validation_scope, + ) + + static_ptrs = { + "active_slots": active_slots.data_ptr(), + "context_lens": context_lens.data_ptr(), + "output": output.data_ptr(), + } + captured_plans = provider.fa3._captured_scheduler_plans + assert len(captured_plans) == 1 + scheduler_metadata_ptr = captured_plans[0].metadata.data_ptr() + + for context_len in (1023, 1024, 1025): + context_lens.fill_(context_len) + graph.replay() + torch.cuda.synchronize() + + latent = payload.latent_cache[:context_len, 0].float() + rope = payload.rope_cache[:context_len, 0].float() + logits = torch.einsum( + "hd,ld->hl", + q_nope_absorbed[0].float(), + latent, + ) + torch.einsum("hd,ld->hl", q_rope[0].float(), rope) + probabilities = torch.softmax(logits * spec.softmax_scale, dim=-1) + expected = torch.einsum("hl,ld->hd", probabilities, latent).to(torch.bfloat16) + torch.testing.assert_close(output[0], expected, rtol=3e-2, atol=3e-2) + + assert active_slots.data_ptr() == static_ptrs["active_slots"] + assert context_lens.data_ptr() == static_ptrs["context_lens"] + assert output.data_ptr() == static_ptrs["output"] + assert len(provider.fa3._captured_scheduler_plans) == 1 + assert ( + provider.fa3._captured_scheduler_plans[0].metadata.data_ptr() + == scheduler_metadata_ptr + ) diff --git a/tests/test_platforms.py b/tests/test_platforms.py index 80422fe0..e680ad72 100644 --- a/tests/test_platforms.py +++ b/tests/test_platforms.py @@ -1,4 +1,5 @@ import importlib +from types import SimpleNamespace import pytest import torch @@ -32,6 +33,11 @@ def test_cuda_device_caps_are_the_capability_source(monkeypatch): monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (9, 0)) monkeypatch.setattr(torch.cuda, "get_device_name", lambda _: "Test H100") + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _: SimpleNamespace(multi_processor_count=120), + ) monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) platform = CudaPlatform() @@ -40,6 +46,7 @@ def test_cuda_device_caps_are_the_capability_source(monkeypatch): assert caps.device_index == 7 assert caps.device_name == "Test H100" assert caps.compute_capability == (9, 0) + assert caps.multi_processor_count == 120 assert caps.supports_native_fp8 assert platform.supports_fp8() diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index 97a6dcc0..ce58fab5 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -11,6 +11,10 @@ import torch from sparsevllm.config import Config +from sparsevllm.configs.cuda_graph import ( + _default_decode_cuda_graph_capture_sizes, + _resolve_decode_static_batch_capacity, +) from sparsevllm.engine.cache_manager.standard import StandardCacheManager from sparsevllm.engine.cache_manager.deltakv import DeltaKVCacheManager from sparsevllm.engine.cache_manager.deltakv_less_memory import DeltaKVLessMemoryCacheManager @@ -714,8 +718,13 @@ def test_decode_cuda_graph_capture_sampling_requires_graph(self): decode_graph_capture_sampling=True, ) - def test_decode_cuda_graph_auto_capture_sizes_cover_decode_limit(self): - for max_decoding_seqs in (1, 6, 8, 24): + def test_decode_cuda_graph_auto_capture_sizes_end_at_decode_limit(self): + for max_decoding_seqs, expected_sizes in ( + (1, [1]), + (6, [1, 2, 3, 4, 5, 6]), + (8, [1, 2, 3, 4, 5, 6, 7, 8]), + (24, [1, 2, 3, 4, 5, 6, 7, 8, 12, 16, 20, 24]), + ): with self.subTest(max_decoding_seqs=max_decoding_seqs): cfg = self.make_config( sparse_method="omnikv", @@ -727,7 +736,37 @@ def test_decode_cuda_graph_auto_capture_sizes_cover_decode_limit(self): self.assertEqual(capture_sizes[-1], max_decoding_seqs) self.assertTrue(all(0 < size <= max_decoding_seqs for size in capture_sizes)) self.assertTrue(cfg.decode_graph) - self.assertEqual(cfg.decode_graph_capture_sizes, capture_sizes) + self.assertEqual(cfg.decode_graph_capture_sizes, expected_sizes) + + def test_decode_cuda_graph_auto_capture_sizes_are_bounded_for_large_limits(self): + for max_decoding_seqs in (64, 80, 128, 256, 1024): + with self.subTest(max_decoding_seqs=max_decoding_seqs): + sizes = _default_decode_cuda_graph_capture_sizes(max_decoding_seqs) + self.assertLessEqual(len(sizes), 32) + self.assertEqual(sizes[:8], list(range(1, 9))) + self.assertEqual(sizes[-1], max_decoding_seqs) + self.assertEqual(sizes, sorted(set(sizes))) + + def test_decode_static_batch_capacity_uses_reachable_padding_bucket(self): + cases = ( + ([1, 2, 4, 8, 16, 32, 64], 32, 64, 32), + ([1, 4, 8, 64], 32, 64, 64), + ([1, 2, 4, 8, 16, 32, 64], 80, 64, 64), + ) + for capture_sizes, max_batch, max_decode, expected in cases: + with self.subTest( + capture_sizes=capture_sizes, + max_batch=max_batch, + max_decode=max_decode, + ): + self.assertEqual( + _resolve_decode_static_batch_capacity( + capture_sizes, + max_num_seqs_in_batch=max_batch, + max_decoding_seqs=max_decode, + ), + expected, + ) def test_legacy_platform_aliases_are_not_config_fields(self): fields = Config.__dataclass_fields__ @@ -767,10 +806,10 @@ def test_auto_capture_greedy_sampling_scope(self): enable_prefix_caching=False, sparse_method="", ) - self.assertTrue(runner._auto_capture_greedy_sampling(seqs)) + self.assertFalse(runner._auto_capture_greedy_sampling(seqs)) runner.config.sparse_method = "omnikv" - self.assertTrue(runner._auto_capture_greedy_sampling(seqs)) + self.assertFalse(runner._auto_capture_greedy_sampling(seqs)) runner.config.sparse_method = "quest" self.assertFalse(runner._auto_capture_greedy_sampling(seqs)) @@ -788,8 +827,14 @@ def test_auto_capture_greedy_sampling_scope(self): self.assertFalse(runner._auto_capture_greedy_sampling(seqs)) runner.config.decode_graph_capture_sampling = True + self.assertFalse(runner._auto_capture_greedy_sampling(seqs)) + runner.config.tensor_parallel_size = 1 self.assertTrue(runner._auto_capture_greedy_sampling(seqs)) + seqs[0].temperature = 0.7 + self.assertFalse(runner._auto_capture_greedy_sampling(seqs)) + seqs[0].temperature = 0.0 + seqs[0].presence_penalty = 0.1 self.assertFalse(runner._auto_capture_greedy_sampling(seqs)) seqs[0].presence_penalty = 0.0 @@ -905,6 +950,7 @@ def make_graph_manager(self, *, context_policy="current", max_cached_graphs=None def make_runner(self, method="quest", cache_manager=None): runner = object.__new__(DecodeCudaGraphRunner) runner.method = method + runner.shape_policy = "bucketed" runner.cache_manager = cache_manager if cache_manager is not None else SimpleNamespace() runner.runtime_state = runner.cache_manager runner.recurrent_state_manager = None @@ -973,7 +1019,14 @@ def test_deltakv_graph_eager_static_uses_current_capacity_policy(self): ), ) runner.runtime_state = SimpleNamespace( - prepare_decode_static=runner.cache_manager.prepare_decode_static, + prepare_decode_graph_step=lambda seqs, state: runner.cache_manager.prepare_decode_static( + seqs, + state.inputs.input_ids, + state.inputs.positions, + state.inputs.write_slot_mapping, + state.inputs.context_lens, + state.inputs.request_indices, + ), ) runner.sparse_controller = SimpleNamespace(prepare_forward=lambda seqs, is_prefill: None) runner.is_long_text_batch = lambda seqs, is_prefill: False @@ -1009,7 +1062,14 @@ def test_eager_static_allows_tp_worker_without_logits(self): ), ) runner.runtime_state = SimpleNamespace( - prepare_decode_static=runner.cache_manager.prepare_decode_static, + prepare_decode_graph_step=lambda seqs, state: runner.cache_manager.prepare_decode_static( + seqs, + state.inputs.input_ids, + state.inputs.positions, + state.inputs.write_slot_mapping, + state.inputs.context_lens, + state.inputs.request_indices, + ), ) runner.sparse_controller = SimpleNamespace( prepare_forward=lambda seqs, is_prefill: calls.append(f"prepare:{is_prefill}") @@ -1033,7 +1093,14 @@ def empty_on_cpu(shape, *, dtype=None, device=None): def test_exact_current_policy_does_not_reuse_larger_warmup_state(self): runner = self.make_runner("quest") - warmup_key = DecodeCudaGraphKey("quest", 1, 16384, False, False) + warmup_key = DecodeCudaGraphKey( + "quest", + 1, + 16384, + False, + False, + shape_policy="bucketed", + ) warmup_state = DecodeCudaGraphState(key=warmup_key) runner._graphs[warmup_key] = warmup_state real_empty = torch.empty @@ -1070,8 +1137,22 @@ def test_evict_cached_graphs_releases_oldest_unprotected_state(self): runner = self.make_runner("deltakv") runner.max_cached_graphs = 1 runner._graphs = OrderedDict() - old_key = DecodeCudaGraphKey("deltakv", 1, 1024, False, False) - new_key = DecodeCudaGraphKey("deltakv", 1, 2048, False, False) + old_key = DecodeCudaGraphKey( + "deltakv", + 1, + 1024, + False, + False, + shape_policy="bucketed", + ) + new_key = DecodeCudaGraphKey( + "deltakv", + 1, + 2048, + False, + False, + shape_policy="bucketed", + ) old_state = DecodeCudaGraphState(key=old_key) old_state.keepalive.append(object()) old_state.sparse_state_refs[0] = {"attn_score": object()} @@ -1554,6 +1635,65 @@ def test_vanilla_model_runner_does_not_partition_long_and_short(self): ModelRunner._is_long_text_batch(runner, seqs, is_prefill=False) ) + def test_sparse_model_runner_rejects_mixed_decode_topology_batch(self): + runner = object.__new__(ModelRunner) + runner.config = SimpleNamespace( + sparse_method="quest", + sink_keep_tokens=1, + recent_keep_tokens=1, + decode_keep_tokens=4, + ) + + with self.assertRaisesRegex(ValueError, "Mixed long/short batch"): + ModelRunner._is_long_text_batch( + runner, + [seq_with_len(4), seq_with_len(20)], + is_prefill=False, + ) + + def test_sparse_decode_transition_and_prefix_restore_select_long_path(self): + runner = object.__new__(ModelRunner) + runner.config = SimpleNamespace( + sparse_method="omnikv", + sink_keep_tokens=1, + recent_keep_tokens=1, + decode_keep_tokens=4, + ) + threshold = ModelRunner._long_text_threshold( + runner, + is_prefill=False, + ) + sequence = seq_with_len(threshold) + + self.assertFalse( + ModelRunner._is_long_text_batch( + runner, + [sequence], + is_prefill=False, + ) + ) + sequence.append_token(0) + self.assertTrue( + ModelRunner._is_long_text_batch( + runner, + [sequence], + is_prefill=False, + ) + ) + + restored = seq_with_len(threshold + 1) + restored.prefix_cache_enabled = True + restored.prefix_cache_hit_len = threshold + restored.prefix_cache_hit_block_count = 1 + restored.prefix_cache_hit_last_block_id = b"prefix" + self.assertTrue( + ModelRunner._is_long_text_batch( + runner, + [restored], + is_prefill=False, + ) + ) + def test_all_chunked_batches_sparse_mixed_lengths(self): scheduler = make_scheduler( PREFILL_POLICY_ALL_CHUNKED, @@ -1587,6 +1727,29 @@ def test_vanilla_decode_batches_across_sparse_long_text_boundary(self): self.assertFalse(is_prefill) self.assertEqual(scheduled, [short_seq, long_seq]) + def test_sparse_decode_schedules_short_and_long_topologies_separately(self): + scheduler = make_scheduler( + PREFILL_POLICY_ALL_CHUNKED, + method="quest", + ) + short_seq = seq_with_len(4) + long_seq = seq_with_len(20) + short_seq.num_prefilled_tokens = short_seq.num_prompt_tokens + long_seq.num_prefilled_tokens = long_seq.num_prompt_tokens + scheduler.decoding.extend((short_seq, long_seq)) + + short_batch, is_prefill, _ = scheduler.schedule() + + self.assertFalse(is_prefill) + self.assertEqual(short_batch, [short_seq]) + self.assertIn(long_seq, scheduler.decoding) + + scheduler.decoding.remove(short_seq) + long_batch, is_prefill_long, _ = scheduler.schedule() + + self.assertFalse(is_prefill_long) + self.assertEqual(long_batch, [long_seq]) + def test_all_chunked_caps_each_prefill_by_chunk_size(self): scheduler = make_scheduler(PREFILL_POLICY_ALL_CHUNKED, method="", chunk=5, max_tokens=20) seq_a = seq_with_len(20) diff --git a/tests/test_prefix_cache.py b/tests/test_prefix_cache.py index 0242d6e6..3a70653b 100644 --- a/tests/test_prefix_cache.py +++ b/tests/test_prefix_cache.py @@ -21,7 +21,10 @@ QuestPrefixOffloadController, StandardPrefixOffloadController, ) -from sparsevllm.engine.sequence import Sequence +from sparsevllm.engine.decode_graph_contract import ( + DecodeGraphContract, + DecodeGraphInputs, +) from sparsevllm.engine.prefix_cache import ( PrefixBlockResidency, PrefixCacheBlock, @@ -32,6 +35,7 @@ resolve_prefix_cache_block_size, usable_prefix_cache_tokens, ) +from sparsevllm.engine.sequence import Sequence from sparsevllm.platforms import device_runtime @@ -2255,6 +2259,72 @@ def test_standard_static_decode_padding_does_not_materialize_padded_rows(): assert manager._num_free_slots == 86 +def test_standard_decode_graph_state_updates_stable_typed_inputs(): + manager = _make_standard_manager_for_prefix(block_size=4) + seq = Sequence([1, 2, 3]) + prompt_slots = manager._allocate(seq.seq_id, 3) + manager._record_prefix_materialization(seq, [1, 2, 3], prompt_slots) + seq.num_prefilled_tokens = seq.num_prompt_tokens + seq.append_token(4) + + contract = DecodeGraphContract( + method="", + shape_policy="batch_only", + topology_path_id="dense", + batch_capacity=4, + context_capacity=16, + ) + inputs = DecodeGraphInputs.allocate( + contract, + device=torch.device("cpu"), + pin_memory=False, + ) + state = manager.init_decode_graph_state(contract, inputs) + pointers = inputs.data_ptrs() + assert contract.capability_level == "strict" + + manager.prepare_decode_graph_step([seq], state) + + assert inputs.data_ptrs() == pointers + assert inputs.input_ids.tolist() == [4, 4, 4, 4] + assert inputs.positions.tolist() == [3, 3, 3, 3] + assert inputs.write_slot_mapping.tolist()[1:] == [-1, -1, -1] + assert inputs.context_lens.tolist() == [4, 4, 4, 4] + assert inputs.request_indices.tolist() == [0, 0, 0, 0] + assert inputs.active_mask.tolist() == [True, False, False, False] + assert all(tensor.data_ptr() for tensor in inputs.keepalive_tensors()) + + +def test_standard_decode_graph_rejects_capacity_before_cache_mutation(): + manager = _make_standard_manager_for_prefix(block_size=4) + seq = Sequence([1, 2, 3]) + prompt_slots = manager._allocate(seq.seq_id, 3) + manager._record_prefix_materialization(seq, [1, 2, 3], prompt_slots) + seq.num_prefilled_tokens = seq.num_prompt_tokens + seq.append_token(4) + contract = DecodeGraphContract( + method="", + shape_policy="batch_only", + topology_path_id="dense", + batch_capacity=1, + context_capacity=3, + ) + inputs = DecodeGraphInputs.allocate( + contract, + device=torch.device("cpu"), + pin_memory=False, + ) + state = manager.init_decode_graph_state(contract, inputs) + free_slots_before = manager._num_free_slots + row_len_before = int(manager.row_seq_lens[manager.seq_id_to_row[seq.seq_id]]) + + with pytest.raises(ValueError, match="exceeded the captured graph context"): + manager.prepare_decode_graph_step([seq], state) + + assert manager._num_free_slots == free_slots_before + assert int(manager.row_seq_lens[manager.seq_id_to_row[seq.seq_id]]) == row_len_before + + def test_standard_decode_materialized_block_can_seed_later_prefix_hit(): manager = _make_standard_manager_for_prefix(block_size=4) first = Sequence([1, 2, 3]) diff --git a/tests/test_sgl_fa3.py b/tests/test_sgl_fa3.py index 0aa9c422..2661b6d4 100644 --- a/tests/test_sgl_fa3.py +++ b/tests/test_sgl_fa3.py @@ -28,8 +28,8 @@ def test_sgl_fa3_support_rejects_missing_package() -> None: assert "sglang-kernel is not installed" in str(exc_info.value) -@pytest.mark.parametrize("version", ["0.4.4", "0.5.0"]) -def test_sgl_fa3_support_rejects_outside_declared_range(version: str) -> None: +@pytest.mark.parametrize("version", ["0.4.4", "0.4.5.post1", "0.4.6.post1"]) +def test_sgl_fa3_support_rejects_unpinned_version(version: str) -> None: with ( patch("importlib.util.find_spec", return_value=object()), patch("importlib.metadata.version", return_value=version), @@ -38,11 +38,11 @@ def test_sgl_fa3_support_rejects_outside_declared_range(version: str) -> None: sgl_fa3_support() assert exc_info.value.health.state is KernelFamilyState.BROKEN - assert "sglang-kernel>=0.4.5,<0.5" in str(exc_info.value) + assert "sglang-kernel==0.4.5" in str(exc_info.value) -@pytest.mark.parametrize("version", ["0.4.5", "0.4.6.post1"]) -def test_sgl_fa3_support_accepts_declared_range(version: str) -> None: +def test_sgl_fa3_support_accepts_pinned_version() -> None: + version = "0.4.5" op = SimpleNamespace( _schema=SimpleNamespace( arguments=[ @@ -429,6 +429,40 @@ def counted_scheduler_op(*args, **kwargs): atol=3e-2, ) + request_indices.copy_( + torch.tensor([1, 3, 0], device=device, dtype=torch.int32) + ) + context_lens.copy_( + torch.tensor([3, 8, 6], device=device, dtype=torch.int32) + ) + replay_rows = [] + for batch_index in range(batch_size): + length = int(context_lens[batch_index].item()) + row = int(request_indices[batch_index].item()) + active = page_table[row, :length].long() + logits = q_rope[batch_index].float() @ rope_cache[active, 0].float().T + logits += q_latent[batch_index].float() @ latent_cache[active, 0].float().T + probs = torch.softmax(logits * (256**-0.5), dim=-1) + replay_rows.append( + (probs @ latent_cache[active, 0].float()).to(torch.bfloat16) + ) + replay_expected = torch.stack(replay_rows) + graph.replay() + second_graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + graph_output, + replay_expected, + rtol=3e-2, + atol=3e-2, + ) + torch.testing.assert_close( + second_graph_output, + replay_expected, + rtol=3e-2, + atol=3e-2, + ) + @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_fa3_support()[0], diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py index f43a4f20..baa9e1c9 100644 --- a/tests/test_sgl_moe.py +++ b/tests/test_sgl_moe.py @@ -276,8 +276,8 @@ def test_sgl_moe_support_rejects_missing_package() -> None: assert "sglang-kernel is not installed" in str(exc_info.value) -@pytest.mark.parametrize("version", ["0.4.4", "0.5.0"]) -def test_sgl_moe_support_rejects_outside_declared_range(version: str) -> None: +@pytest.mark.parametrize("version", ["0.4.4", "0.4.5.post1", "0.4.6.post1"]) +def test_sgl_moe_support_rejects_unpinned_version(version: str) -> None: with ( patch("importlib.util.find_spec", return_value=object()), patch("importlib.metadata.version", return_value=version), @@ -286,11 +286,11 @@ def test_sgl_moe_support_rejects_outside_declared_range(version: str) -> None: sgl_moe_alignment_support() assert exc_info.value.health.state is KernelFamilyState.BROKEN - assert "sglang-kernel>=0.4.5,<0.5" in str(exc_info.value) + assert "sglang-kernel==0.4.5" in str(exc_info.value) -@pytest.mark.parametrize("version", ["0.4.5", "0.4.6.post1"]) -def test_sgl_moe_support_accepts_declared_range(version: str) -> None: +def test_sgl_moe_support_accepts_pinned_version() -> None: + version = "0.4.5" module = SimpleNamespace(moe_align_block_size=lambda *_args: None) with ( patch("importlib.util.find_spec", return_value=object()), diff --git a/tests/test_sparse_state_summary.py b/tests/test_sparse_state_summary.py index a9543a80..5b01d9b4 100644 --- a/tests/test_sparse_state_summary.py +++ b/tests/test_sparse_state_summary.py @@ -120,7 +120,10 @@ def gather(output, local, group): "replay_count": 3, "eager_static_count": 0, "force_eager_count": 0, + "eviction_count": 0, + "recapture_count": 0, "cached_graph_count": 1, + "bucket_plan": None, "last_state_key": { "method": "snapkv", "batch_size": 2, diff --git a/tests/test_tilelang_mla_kernel.py b/tests/test_tilelang_mla_kernel.py index 27d706cc..41556060 100644 --- a/tests/test_tilelang_mla_kernel.py +++ b/tests/test_tilelang_mla_kernel.py @@ -6,6 +6,7 @@ from sparsevllm.kernels.tilelang.mla.runtime import ( TileMlaDecodeKernel, TileMlaLaunchConfig, + TileMlaLaunchPlan, ) CUDA_REQUIRED = pytest.mark.skipif( @@ -20,7 +21,7 @@ def _torch_oracle( latent_cache: torch.Tensor, rope_cache: torch.Tensor, slots: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: latent_keys = latent_cache[slots.long(), 0].float() rope_keys = rope_cache[slots.long(), 0].float() raw = torch.matmul(q_latent.float(), latent_keys.T) + torch.matmul( @@ -29,7 +30,7 @@ def _torch_oracle( score = raw.max(dim=0).values probability = torch.softmax(raw * (256**-0.5), dim=-1) output = torch.matmul(probability, latent_keys) - return output.to(torch.bfloat16), score + return output.to(torch.bfloat16), raw, score @CUDA_REQUIRED @@ -43,6 +44,9 @@ def _torch_oracle( (20, 1, 16, "atomic"), (20, 4, 16, "partial"), (20, 4, 32, "direct"), + (5, 4, 16, "per_head"), + (10, 16, 16, "per_head"), + (20, 4, 32, "per_head"), ], ) def test_tilelang_mla_score_matches_torch_with_indirect_slots_and_graph( @@ -83,7 +87,11 @@ def test_tilelang_mla_score_matches_torch_with_indirect_slots_and_graph( q_latent.shape, dtype=q_latent.dtype, device=q_latent.device ) score = torch.full( - (batch_size, capacity), + ( + (batch_size, valid_heads, capacity) + if score_mode == "per_head" + else (batch_size, capacity) + ), -1e20, dtype=torch.float32, device=device, @@ -119,7 +127,7 @@ def run() -> None: run() torch.cuda.synchronize() - expected_output, expected_score = _torch_oracle( + expected_output, expected_per_head_score, expected_reduced_score = _torch_oracle( q_latent[0], q_rope[0], latent_cache, @@ -129,11 +137,20 @@ def run() -> None: torch.testing.assert_close( output[0], expected_output, rtol=3e-2, atol=3e-2 ) - torch.testing.assert_close( - score[0, :33], expected_score, rtol=3e-2, atol=3e-2 - ) + if score_mode == "per_head": + torch.testing.assert_close( + score[0, :, :33], + expected_per_head_score, + rtol=3e-2, + atol=3e-2, + ) + assert torch.all(score[0, :, 33:] == -1e20) + else: + torch.testing.assert_close( + score[0, :33], expected_reduced_score, rtol=3e-2, atol=3e-2 + ) + assert torch.all(score[0, 33:] == -1e20) torch.testing.assert_close(output[1], torch.zeros_like(output[1])) - assert torch.all(score[0, 33:] == -1e20) assert torch.all(score[1] == -1e20) run() @@ -207,3 +224,194 @@ def test_tilelang_score_ignores_zero_padded_heads(valid_heads: int) -> None: rtol=0, atol=0, ) + + +@CUDA_REQUIRED +def test_static_plan_replays_across_contexts_with_unaligned_capacity() -> None: + torch.manual_seed(20260825) + device = torch.device("cuda") + valid_heads = 10 + capacity = 127 + cache_slots = 160 + q_latent = torch.randn( + 1, valid_heads, 512, dtype=torch.bfloat16, device=device + ) + q_rope = torch.randn( + 1, valid_heads, 64, dtype=torch.bfloat16, device=device + ) + latent_cache = torch.randn( + cache_slots, 1, 512, dtype=torch.bfloat16, device=device + ) + rope_cache = torch.randn( + cache_slots, 1, 64, dtype=torch.bfloat16, device=device + ) + active_slots = torch.randperm( + cache_slots, dtype=torch.int64, device=device + )[:capacity].to(torch.int32).unsqueeze(0) + request_indices = torch.zeros(1, dtype=torch.int32, device=device) + context_lens = torch.full((1,), 31, dtype=torch.int32, device=device) + output = torch.empty_like(q_latent) + score_storage = torch.empty( + 1, valid_heads, capacity + 1, dtype=torch.float32, device=device + ) + score = score_storage[:, :, :capacity] + assert not score.is_contiguous() + plan = TileMlaLaunchPlan.build( + context_capacity=8192, + local_q_heads=valid_heads, + max_batch_size=1, + need_score=True, + score_mode="per_head", + ) + runner = TileMlaDecodeKernel( + device=device, + softmax_scale=256**-0.5, + valid_heads=valid_heads, + launch_plan=plan, + ) + + def run() -> None: + score.fill_(-1e20) + runner( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + attn_score=score, + max_context_len=capacity, + ) + + run() + torch.cuda.synchronize() + metadata = runner.runtime_metadata() + assert metadata["compiled_variant_count"] == 1 + workspace_ptrs = metadata["compiled_variants"][0]["workspace_data_ptrs"] + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for context_len in (31, 64, 65, capacity, 31): + context_lens.fill_(context_len) + graph.replay() + torch.cuda.synchronize() + expected_output, expected_score, _ = _torch_oracle( + q_latent[0], + q_rope[0], + latent_cache, + rope_cache, + active_slots[0, :context_len], + ) + torch.testing.assert_close( + output[0], expected_output, rtol=3e-2, atol=3e-2 + ) + torch.testing.assert_close( + score[0, :, :context_len], + expected_score, + rtol=3e-2, + atol=3e-2, + ) + assert torch.all(score[0, :, context_len:] == -1e20) + metadata = runner.runtime_metadata() + assert metadata["compiled_variant_count"] == 1 + assert ( + metadata["compiled_variants"][0]["workspace_data_ptrs"] + == workspace_ptrs + ) + + +@CUDA_REQUIRED +def test_static_plan_replays_representative_contexts_through_64k() -> None: + torch.manual_seed(20260825) + device = torch.device("cuda") + valid_heads = 5 + capacity = 65536 + q_latent = torch.randn( + 1, valid_heads, 512, dtype=torch.bfloat16, device=device + ) + q_rope = torch.randn( + 1, valid_heads, 64, dtype=torch.bfloat16, device=device + ) + latent_cache = torch.randn( + capacity, 1, 512, dtype=torch.bfloat16, device=device + ) + rope_cache = torch.randn( + capacity, 1, 64, dtype=torch.bfloat16, device=device + ) + active_slots = torch.arange( + capacity, dtype=torch.int32, device=device + ).unsqueeze(0) + request_indices = torch.zeros(1, dtype=torch.int32, device=device) + context_lens = torch.full((1,), 1024, dtype=torch.int32, device=device) + output = torch.empty_like(q_latent) + score = torch.empty( + 1, valid_heads, capacity, dtype=torch.float32, device=device + ) + plan = TileMlaLaunchPlan.build( + context_capacity=capacity, + local_q_heads=valid_heads, + max_batch_size=1, + need_score=True, + score_mode="per_head", + ) + runner = TileMlaDecodeKernel( + device=device, + softmax_scale=256**-0.5, + valid_heads=valid_heads, + launch_plan=plan, + ) + + def run() -> None: + score.fill_(-1e20) + runner( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + attn_score=score, + max_context_len=capacity, + ) + + run() + torch.cuda.synchronize() + metadata = runner.runtime_metadata() + assert metadata["compiled_variant_count"] == 1 + workspace_ptrs = metadata["compiled_variants"][0]["workspace_data_ptrs"] + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for context_len in (1024, 4096, 8192, 16384, 32768, capacity): + context_lens.fill_(context_len) + graph.replay() + torch.cuda.synchronize() + expected_output, expected_score, _ = _torch_oracle( + q_latent[0], + q_rope[0], + latent_cache, + rope_cache, + active_slots[0, :context_len], + ) + torch.testing.assert_close( + output[0], expected_output, rtol=3e-2, atol=3e-2 + ) + torch.testing.assert_close( + score[0, :, :context_len], + expected_score, + rtol=3e-2, + atol=3e-2, + ) + assert torch.all(score[0, :, context_len:] == -1e20) + metadata = runner.runtime_metadata() + assert metadata["compiled_variant_count"] == 1 + assert ( + metadata["compiled_variants"][0]["workspace_data_ptrs"] + == workspace_ptrs + ) diff --git a/tests/test_tilelang_mla_operator.py b/tests/test_tilelang_mla_operator.py index f523d24d..bc229c51 100644 --- a/tests/test_tilelang_mla_operator.py +++ b/tests/test_tilelang_mla_operator.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dataclasses import replace import subprocess import sys from importlib import metadata @@ -16,13 +17,15 @@ ) from sparsevllm.kernels.tilelang.mla.runtime import ( TileMlaDecodeKernel, + TileMlaLaunchConfig, + TileMlaLaunchPlan, tilelang_mla_support, ) from sparsevllm.kernels.triton.mla import MlaDecodeWorkspace +from sparsevllm.operators.attention_capabilities import AttentionScoreKind from sparsevllm.operators.mla_attention import ( MLA_ATTENTION_REGISTRY, MlaAttentionOpSpec, - MlaSglFa3Provider, MlaTileLangScoreProvider, MlaTritonProvider, ) @@ -41,7 +44,8 @@ def _spec(*, tp_size: int = 2) -> MlaAttentionOpSpec: cache_dtype=torch.bfloat16, tp_size=tp_size, cuda_graph=True, - may_require_attention_scores=True, + score_output=AttentionScoreKind.RAW_QK_PER_HEAD, + context_capacity=65536, ) @@ -206,11 +210,16 @@ def test_tilelang_provider_binds_rank_local_head_count( max_batch_size=2, ) - tilelang_cls.assert_called_once_with( - device=torch.device("cpu"), - softmax_scale=256**-0.5, - valid_heads=local_heads, - ) + tilelang_cls.assert_called_once() + kwargs = tilelang_cls.call_args.kwargs + assert kwargs["device"] == torch.device("cpu") + assert kwargs["softmax_scale"] == 256**-0.5 + assert kwargs["valid_heads"] == local_heads + plan = kwargs["launch_plan"] + assert isinstance(plan, TileMlaLaunchPlan) + assert plan.context_capacity == 65536 + assert plan.local_q_heads == local_heads + assert all(config.score_mode == "per_head" for config in plan.configs) def test_missing_tilelang_binds_score_capable_triton_provider() -> None: @@ -311,9 +320,83 @@ def test_tilelang_mla_exact_h100_profile_overrides_default_portfolio() -> None: assert resolved.report.selection_basis == "profile_override" +def test_batch_only_score_contract_binds_static_tilelang_plan() -> None: + spec = replace( + _spec(), + batch_only_cuda_graph=True, + ) + workspace = _cpu_workspace() + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_device_support", + return_value=(True, "sgl test"), + ), + patch("sparsevllm.operators.mla_attention.SglFa3DecodeKernel"), + patch("sparsevllm.operators.mla_attention.TileMlaDecodeKernel"), + patch( + "sparsevllm.operators.mla_attention.tilelang_mla_support", + return_value=(True, "tilelang test"), + ), + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ), + ): + resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( + spec, + _h100_caps(), + op_spec=spec, + device="cpu", + max_batch_size=2, + ) + + assert type(resolved.provider) is MlaTileLangScoreProvider + assert resolved.provider.supports_batch_only_cuda_graph + assert resolved.provider.tilelang_launch_plan.context_capacity == 65536 + + +def test_batch_only_reduced_score_contract_binds_static_triton_provider() -> None: + spec = replace( + _spec(), + score_output=AttentionScoreKind.RAW_QK_REDUCED, + batch_only_cuda_graph=True, + ) + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_device_support", + return_value=(True, "sgl test"), + ), + patch( + "sparsevllm.operators.mla_attention.tilelang_mla_support", + return_value=(True, "tilelang test"), + ), + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=_cpu_workspace(), + ), + ): + resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( + spec, + _h100_caps(), + op_spec=spec, + device="cpu", + max_batch_size=2, + ) + + assert type(resolved.provider) is MlaTritonProvider + assert ( + "tilelang_score_sgl_fa3_h100", + "requires the RAW_QK_PER_HEAD decode score contract", + ) in resolved.rejected + + def _provider_with_mocks() -> tuple[MlaTileLangScoreProvider, Mock, Mock]: fa3 = Mock(return_value=torch.empty(2, 10, 512, dtype=torch.bfloat16)) tilelang = Mock(return_value=torch.empty(2, 10, 512, dtype=torch.bfloat16)) + tilelang.runtime_metadata.return_value = { + "compiled_variant_count": 1, + "compiled_variants": [], + } with ( patch( "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", @@ -336,9 +419,9 @@ def _provider_with_mocks() -> tuple[MlaTileLangScoreProvider, Mock, Mock]: return provider, fa3, tilelang -def test_score_path_routes_to_tilelang_with_caller_owned_score() -> None: +def test_per_head_score_path_routes_to_tilelang_with_caller_owned_score() -> None: provider, fa3, tilelang = _provider_with_mocks() - score = torch.full((2, 64), -1e20, dtype=torch.float32) + score = torch.full((2, 10, 64), -1e20, dtype=torch.float32) view = _view(score=score) q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) @@ -370,12 +453,16 @@ def test_score_path_routes_to_tilelang_with_caller_owned_score() -> None: } }, "fallback_reasons": {}, + "tilelang": { + "compiled_variant_count": 1, + "compiled_variants": [], + }, } def test_noncontiguous_glm_queries_route_to_tilelang() -> None: provider, fa3, tilelang = _provider_with_mocks() - score = torch.full((2, 64), -1e20, dtype=torch.float32) + score = torch.full((2, 10, 64), -1e20, dtype=torch.float32) view = _view(score=score) q_latent = torch.empty(10, 2, 512, dtype=torch.bfloat16).transpose(0, 1) q_rope = torch.empty(10, 2, 64, dtype=torch.bfloat16).transpose(0, 1) @@ -394,7 +481,7 @@ def test_noncontiguous_glm_queries_route_to_tilelang() -> None: def test_runtime_kernel_stats_distinguish_cuda_graph_capture() -> None: provider, _, tilelang = _provider_with_mocks() - view = _view(score=torch.empty(2, 64, dtype=torch.float32)) + view = _view(score=torch.empty(2, 10, 64, dtype=torch.float32)) q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) output = torch.empty_like(q_latent) @@ -435,51 +522,50 @@ def test_no_score_path_remains_fa3() -> None: @pytest.mark.parametrize( - "score", + ("score", "message"), [ - torch.empty(2, 10, 64, dtype=torch.float32), - torch.empty(2, 64, dtype=torch.bfloat16), - torch.empty(2, 63, dtype=torch.float32), + (torch.empty(2, 64, dtype=torch.float32), "RAW_QK_PER_HEAD"), + (torch.empty(2, 10, 64, dtype=torch.bfloat16), "must use FP32"), + (torch.empty(2, 9, 64, dtype=torch.float32), "head count"), ], ) -def test_unsupported_score_contract_uses_explicit_triton_path(score) -> None: +def test_unsupported_score_contract_fails_instead_of_falling_back( + score, + message: str, +) -> None: provider, fa3, tilelang = _provider_with_mocks() view = _view(score=score) q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) output = torch.empty_like(q_latent) - with patch.object( - MlaSglFa3Provider.__mro__[1], "run", return_value=output - ) as triton: + with pytest.raises((TypeError, ValueError), match=message): provider.run(q_latent, q_rope, view, output) fa3.assert_not_called() tilelang.assert_not_called() - triton.assert_called_once() + assert provider.runtime_kernel_stats()["fallback_reasons"] == {} -def test_score_capacity_smaller_than_declared_context_uses_triton() -> None: +def test_score_capacity_smaller_than_declared_context_fails() -> None: provider, fa3, tilelang = _provider_with_mocks() - view = _view(score=torch.empty(2, 64, dtype=torch.float32)) + view = _view(score=torch.empty(2, 10, 64, dtype=torch.float32)) object.__setattr__(view.meta, "max_context_len", 128) q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) output = torch.empty_like(q_latent) - with patch.object( - MlaSglFa3Provider.__mro__[1], "run", return_value=output - ) as triton: + with pytest.raises(ValueError, match="must cover max_context_len"): provider.run(q_latent, q_rope, view, output) fa3.assert_not_called() tilelang.assert_not_called() - triton.assert_called_once() + assert provider.runtime_kernel_stats()["fallback_reasons"] == {} -def test_score_capacity_larger_than_active_slots_uses_triton() -> None: +def test_score_capacity_may_include_padding_beyond_active_slots() -> None: provider, fa3, tilelang = _provider_with_mocks() - view = _view(score=torch.empty(2, 64, dtype=torch.float32)) + view = _view(score=torch.empty(2, 10, 64, dtype=torch.float32)) object.__setattr__( view.meta, "active_slots", @@ -490,52 +576,59 @@ def test_score_capacity_larger_than_active_slots_uses_triton() -> None: q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) output = torch.empty_like(q_latent) - with patch.object( - MlaSglFa3Provider.__mro__[1], "run", return_value=output - ) as triton: + with patch( + "sparsevllm.operators.mla_attention.validate_mla_decode_metadata" + ): provider.run(q_latent, q_rope, view, output) fa3.assert_not_called() - tilelang.assert_not_called() - triton.assert_called_once() + tilelang.assert_called_once() -@pytest.mark.parametrize("noncontiguous", ["active_slots", "attn_score"]) -def test_noncontiguous_tilelang_inputs_use_triton(noncontiguous: str) -> None: +def test_noncontiguous_active_slots_fail_instead_of_falling_back() -> None: provider, fa3, tilelang = _provider_with_mocks() - score = torch.empty(2, 128, dtype=torch.float32)[:, ::2] - view = _view( - score=( - score - if noncontiguous == "attn_score" - else torch.empty(2, 64, dtype=torch.float32) - ) - ) - if noncontiguous == "active_slots": - backing = torch.full((3, 128), -1, dtype=torch.int32) - backing[2, :6:2] = torch.tensor([5, 2, 7], dtype=torch.int32) - object.__setattr__(view.meta, "active_slots", backing[:, ::2]) + view = _view(score=torch.empty(2, 10, 64, dtype=torch.float32)) + backing = torch.full((3, 128), -1, dtype=torch.int32) + backing[2, :6:2] = torch.tensor([5, 2, 7], dtype=torch.int32) + object.__setattr__(view.meta, "active_slots", backing[:, ::2]) q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) output = torch.empty_like(q_latent) - with patch.object( - MlaSglFa3Provider.__mro__[1], "run", return_value=output - ) as triton: + with pytest.raises(ValueError, match="noncontiguous:active_slots"): provider.run(q_latent, q_rope, view, output) fa3.assert_not_called() tilelang.assert_not_called() - triton.assert_called_once() - assert provider.runtime_kernel_stats()["fallback_reasons"] == { - f"noncontiguous:{noncontiguous}": 1 - } + assert provider.runtime_kernel_stats()["fallback_reasons"] == {} + + +def test_noncontiguous_per_head_score_routes_to_tilelang_staging() -> None: + provider, fa3, tilelang = _provider_with_mocks() + score = torch.empty(2, 10, 128, dtype=torch.float32)[:, :, ::2] + view = _view(score=score) + q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) + q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) + output = torch.empty_like(q_latent) + with patch( + "sparsevllm.operators.mla_attention.validate_mla_decode_metadata" + ): + provider.run(q_latent, q_rope, view, output) -def test_tilelang_runner_rejects_unaligned_score_capacity_before_import() -> None: - runner = TileMlaDecodeKernel(device="cpu", softmax_scale=0.0625) - view = _view(score=torch.empty(2, 63, dtype=torch.float32)) - with pytest.raises(ValueError, match="multiple of 64"): + fa3.assert_not_called() + tilelang.assert_called_once() + assert tilelang.call_args.kwargs["attn_score"] is score + + +def test_tilelang_runner_rejects_score_capacity_smaller_than_context() -> None: + runner = TileMlaDecodeKernel( + device="cpu", + softmax_scale=0.0625, + fixed_config=TileMlaLaunchConfig(1, score_mode="per_head"), + ) + view = _view(score=torch.empty(2, 10, 63, dtype=torch.float32)) + with pytest.raises(ValueError, match="must fit"): runner( torch.empty(2, 10, 512, dtype=torch.bfloat16), torch.empty(2, 10, 64, dtype=torch.bfloat16), @@ -546,7 +639,7 @@ def test_tilelang_runner_rejects_unaligned_score_capacity_before_import() -> Non view.meta.context_lens, torch.empty(2, 10, 512, dtype=torch.bfloat16), attn_score=view.meta.attn_score, - max_context_len=63, + max_context_len=64, ) diff --git a/tests/test_tiny_random.py b/tests/test_tiny_random.py index aea5398b..a7e8cf31 100644 --- a/tests/test_tiny_random.py +++ b/tests/test_tiny_random.py @@ -10,6 +10,7 @@ from sparsevllm.debug.tiny_random import ( apply_tiny_random_overrides, build_tiny_random_hf_model, + initialize_sparse_model, load_tiny_random_overrides, resolve_tiny_random_settings, ) @@ -110,3 +111,25 @@ def test_normal_config_import_does_not_import_tiny_random_module(): text=True, ) assert completed.returncode == 0, completed.stderr + + +def test_quantized_tiny_random_initializes_fp8_weights_and_scales(): + class QuantizedModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.empty(128, 128, dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + self.register_buffer( + "weight_scale_inv", + torch.empty(1, 1, dtype=torch.float32), + ) + + first = QuantizedModel() + second = QuantizedModel() + initialize_sparse_model(first, None, seed=23, quantized=True) + initialize_sparse_model(second, None, seed=23, quantized=True) + + assert torch.equal(first.weight, second.weight) + assert torch.equal(first.weight_scale_inv, torch.ones_like(first.weight_scale_inv)) diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index d5a03e7a..5ac271a3 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -493,6 +493,36 @@ def test_model_runner_reset_after_warmup_resets_local_runtime_state(): assert calls == ["runtime"] +def test_model_runner_decode_graph_startup_controls_use_live_runner(): + calls = [] + runner = object.__new__(ModelRunner) + runner.decode_graph_runner = SimpleNamespace( + set_reuse_larger_context_graphs=lambda enabled: calls.append( + ("reuse", enabled) + ), + seal_startup_plan=lambda: calls.append(("seal",)), + run=lambda seqs, capture_sampling: calls.append( + ("capture", seqs, capture_sampling) + ), + ) + seqs = [object()] + + with patch( + "sparsevllm.engine.model_runner.reset_context", + side_effect=lambda: calls.append(("reset",)), + ): + ModelRunner.set_decode_cuda_graph_reuse_larger_context_graphs(runner, True) + ModelRunner.seal_decode_cuda_graph_startup_plan(runner) + ModelRunner.capture_decode_cuda_graph_warmup(runner, seqs) + + assert calls == [ + ("reuse", True), + ("seal",), + ("capture", seqs, False), + ("reset",), + ] + + def test_model_runner_exit_drains_graphs_before_barrier(): calls = [] runner = object.__new__(ModelRunner)